Implement ZMQ port management and kernel readiness checks

Container Configuration:
- Added network_mode and port_bindings to ContainerConfig
- Support for 'none', 'host', and 'bridge' network modes
- Default remains 'none' for security

Jupyter Kernel Manager:
- Dynamic port allocation for 5 ZMQ channels using socket.socket()
- _allocate_ports() finds available ports via OS binding
- Host networking mode for Jupyter kernels (network_mode='host')
- Connection file properly mounted into container
- Port bindings tracked for documentation

Kernel Readiness:
- _wait_for_kernel_ready() polls shell port until kernel responds
- Configurable timeout (30s) and poll interval (0.5s)
- Replaced time.sleep(2) with proper connectivity check
- Early return when kernel is ready

This completes the core ZMQ communication infrastructure needed
for real Jupyter kernel operation.
This commit is contained in:
Hans Aschauer 2026-02-07 08:16:51 +01:00
parent 9cc43e4166
commit 14aba0a048
3 changed files with 128 additions and 34 deletions

View file

@ -55,7 +55,20 @@
- Error handling for injection failures - Error handling for injection failures
- Passed through SessionManager and JupyterBackend - Passed through SessionManager and JupyterBackend
7. **Dependencies** 7. **Network & Port Management**
- Host networking mode for ZMQ communication (`network_mode="host"`)
- Dynamic port allocation for 5 ZMQ channels (shell, iopub, stdin, control, hb)
- Port availability checking before allocation
- Port bindings tracked in ContainerConfig
- Connection file mounted into container at `/tmp/kernel-{id}.json`
8. **Kernel Readiness**
- `_wait_for_kernel_ready()`: Polls shell port until kernel responds
- Configurable timeout (default 30s) and poll interval (0.5s)
- Replaces simple sleep with proper port connectivity check
- Returns early when kernel is ready
9. **Dependencies**
- `jupyter-client>=8.8.0` added to server dependencies - `jupyter-client>=8.8.0` added to server dependencies
- `ipykernel` removed from server (will be in container image) - `ipykernel` removed from server (will be in container image)
- `pyzmq>=27.1.0` for ZMQ support - `pyzmq>=27.1.0` for ZMQ support
@ -68,17 +81,7 @@
### 🚧 In Progress / TODO ### 🚧 In Progress / TODO
1. **Container Configuration** 1. **Restart Implementation**
- [ ] Port mappings for ZMQ sockets (shell, iopub, stdin, control, hb)
- [ ] Mount connection file into container at `/tmp/kernel-{id}.json`
- [ ] Proper network configuration for host-container ZMQ communication
2. **Kernel Readiness**
- [ ] Replace `time.sleep(2)` with proper readiness check
- [ ] Poll kernel status messages
- [ ] Implement retry logic with timeout
3. **Restart Implementation**
- [ ] Use `KernelManager` instead of just `BlockingKernelClient` - [ ] Use `KernelManager` instead of just `BlockingKernelClient`
- [ ] Proper restart via `KernelManager.restart_kernel()` - [ ] Proper restart via `KernelManager.restart_kernel()`
- [ ] Handle restart failures gracefully - [ ] Handle restart failures gracefully
@ -145,16 +148,16 @@ Current test results: **20 failed, 2 passed**
1. **Immediate** (Core functionality): 1. **Immediate** (Core functionality):
``` ```
1. Implement port mapping in ContainerConfig 1. Update container image with ipykernel
2. Mount connection file into container 2. Test kernel startup with real container
3. Test kernel startup with real container 3. Mock tests for unit testing
``` ```
2. **Short-term** (Stability): 2. **Short-term** (Stability):
``` ```
4. Implement proper readiness check 4. Implement proper kernel restart
5. Update container image with ipykernel 5. Error handling improvements
6. Mock tests for unit testing 6. Port conflict handling
``` ```
3. **Medium-term** (Production-ready): 3. **Medium-term** (Production-ready):

View file

@ -15,6 +15,7 @@ import uuid
import json import json
import tempfile import tempfile
import time import time
import socket
from pathlib import Path from pathlib import Path
from jupyter_client.blocking.client import BlockingKernelClient from jupyter_client.blocking.client import BlockingKernelClient
@ -138,7 +139,7 @@ class JupyterKernelManager:
connection_file = self._create_connection_file(kernel_id, connection_info) connection_file = self._create_connection_file(kernel_id, connection_info)
try: try:
# Set up volumes (user volumes + bridge socket if provided) # Set up volumes (user volumes + bridge socket + connection file)
container_volumes = volumes.copy() if volumes else {} container_volumes = volumes.copy() if volumes else {}
if bridge_socket_path: if bridge_socket_path:
container_volumes[bridge_socket_path] = { container_volumes[bridge_socket_path] = {
@ -146,17 +147,30 @@ class JupyterKernelManager:
"mode": "rw" "mode": "rw"
} }
# Create container with ipykernel # Mount connection file into container
container_connection_path = f"/tmp/kernel-{kernel_id}.json"
container_volumes[str(connection_file)] = {
"bind": container_connection_path,
"mode": "ro"
}
# Create container with ipykernel using host networking
config = ContainerConfig( config = ContainerConfig(
image=self.image, image=self.image,
command=[ command=[
"python", "-m", "ipykernel_launcher", "python", "-m", "ipykernel_launcher",
"-f", f"/tmp/kernel-{kernel_id}.json" "-f", container_connection_path
], ],
resource_limits=self.resource_limits, resource_limits=self.resource_limits,
volumes=container_volumes, volumes=container_volumes,
# TODO: Port mappings for ZMQ network_mode="host", # Use host network for ZMQ communication
# TODO: Mount connection file into container port_bindings={
connection_info["shell_port"]: connection_info["shell_port"],
connection_info["iopub_port"]: connection_info["iopub_port"],
connection_info["stdin_port"]: connection_info["stdin_port"],
connection_info["control_port"]: connection_info["control_port"],
connection_info["hb_port"]: connection_info["hb_port"],
}
) )
container_id = self.container_manager.create_container( container_id = self.container_manager.create_container(
@ -168,8 +182,9 @@ class JupyterKernelManager:
# Start container # Start container
self.container_manager.start_container(container_id) self.container_manager.start_container(container_id)
# Wait for kernel to be ready # Wait for kernel to be ready with polling
time.sleep(2) # TODO: Better readiness check if not self._wait_for_kernel_ready(connection_info, timeout=30):
raise KernelError(f"Kernel {kernel_id} failed to start within timeout")
# Connect client # Connect client
client = self._connect_client(connection_info) client = self._connect_client(connection_info)
@ -451,15 +466,18 @@ _info
return self.kernels[kernel_id] return self.kernels[kernel_id]
def _generate_connection_info(self) -> Dict[str, Any]: def _generate_connection_info(self) -> Dict[str, Any]:
"""Generate ZMQ connection information.""" """Generate ZMQ connection information with allocated ports."""
import secrets import secrets
# Allocate 5 ports for ZMQ channels
ports = self._allocate_ports(5)
return { return {
"shell_port": 0, # Let ZMQ assign "shell_port": ports[0],
"iopub_port": 0, "iopub_port": ports[1],
"stdin_port": 0, "stdin_port": ports[2],
"control_port": 0, "control_port": ports[3],
"hb_port": 0, "hb_port": ports[4],
"ip": "127.0.0.1", "ip": "127.0.0.1",
"key": secrets.token_hex(32), "key": secrets.token_hex(32),
"transport": "tcp", "transport": "tcp",
@ -467,6 +485,26 @@ _info
"kernel_name": "python3" "kernel_name": "python3"
} }
def _allocate_ports(self, count: int) -> List[int]:
"""
Allocate available ports for ZMQ.
Args:
count: Number of ports to allocate
Returns:
List of allocated port numbers
"""
ports = []
for _ in range(count):
# Let OS assign available port
sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
sock.bind(('127.0.0.1', 0)) # Bind to any available port
port = sock.getsockname()[1]
sock.close()
ports.append(port)
return ports
def _create_connection_file( def _create_connection_file(
self, self,
kernel_id: str, kernel_id: str,
@ -497,6 +535,44 @@ _info
except Exception: except Exception:
return False return False
def _wait_for_kernel_ready(
self,
connection_info: Dict[str, Any],
timeout: int = 30,
poll_interval: float = 0.5
) -> bool:
"""
Wait for kernel to be ready by polling ports.
Args:
connection_info: Kernel connection information
timeout: Maximum time to wait in seconds
poll_interval: Time between polls in seconds
Returns:
True if kernel is ready, False if timeout
"""
start_time = time.time()
shell_port = connection_info["shell_port"]
while time.time() - start_time < timeout:
try:
# Try to connect to shell port
sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
sock.settimeout(1)
result = sock.connect_ex(('127.0.0.1', shell_port))
sock.close()
if result == 0:
# Port is open, kernel is ready
return True
except Exception:
pass
time.sleep(poll_interval)
return False
def _execute_injection_code( def _execute_injection_code(
self, self,
client: BlockingKernelClient, client: BlockingKernelClient,

View file

@ -27,7 +27,9 @@ class ContainerConfig:
volumes: Optional[Dict[str, dict]] = None, volumes: Optional[Dict[str, dict]] = None,
resource_limits: Optional[ResourceLimits] = None, resource_limits: Optional[ResourceLimits] = None,
working_dir: Optional[str] = None, working_dir: Optional[str] = None,
user: str = "1000:1000" user: str = "1000:1000",
network_mode: str = "none",
port_bindings: Optional[Dict[str, int]] = None
): ):
""" """
Initialize container configuration. Initialize container configuration.
@ -40,6 +42,8 @@ class ContainerConfig:
resource_limits: Resource limits to apply resource_limits: Resource limits to apply
working_dir: Working directory in container (None to use image default) working_dir: Working directory in container (None to use image default)
user: User to run as (UID:GID) user: User to run as (UID:GID)
network_mode: Network mode (none, host, bridge). Default is 'none' for security.
port_bindings: Port mappings for network_mode=host (container_port -> host_port)
""" """
self.image = image self.image = image
self.command = command or [] self.command = command or []
@ -48,16 +52,19 @@ class ContainerConfig:
self.resource_limits = resource_limits self.resource_limits = resource_limits
self.working_dir = working_dir self.working_dir = working_dir
self.user = user self.user = user
self.network_mode = network_mode
self.port_bindings = port_bindings or {}
def to_podman_params(self) -> dict: def to_podman_params(self) -> dict:
""" """
Convert to Podman container create parameters. Convert to Podman container create parameters.
Ensures all security requirements are included: Ensures all security requirements are included:
- network_mode: none - network_mode: configurable (default 'none' for security)
- read_only: True - read_only: True
- security_opt: ["no-new-privileges"] - security_opt: ["no-new-privileges"]
- resource limits - resource limits
- port_bindings: for host networking mode
Returns: Returns:
Dictionary of parameters for Podman containers.create() Dictionary of parameters for Podman containers.create()
@ -68,7 +75,7 @@ class ContainerConfig:
"environment": self.environment, "environment": self.environment,
"user": self.user, "user": self.user,
# Security requirements # Security requirements
"network_mode": "none", "network_mode": self.network_mode,
"read_only": True, "read_only": True,
"security_opt": ["no-new-privileges"], "security_opt": ["no-new-privileges"],
} }
@ -77,6 +84,14 @@ class ContainerConfig:
if self.working_dir is not None: if self.working_dir is not None:
params["working_dir"] = self.working_dir params["working_dir"] = self.working_dir
# Add port bindings if using host network mode
# Note: In host mode, ports are directly accessible
# port_bindings are informational for tracking
if self.network_mode == "host" and self.port_bindings:
# With host networking, container uses host's network stack directly
# No explicit port mapping needed, but we track for documentation
pass
# Add volumes if present # Add volumes if present
if self.volumes: if self.volumes:
params["volumes"] = self.volumes params["volumes"] = self.volumes