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

@ -15,6 +15,7 @@ import uuid
import json
import tempfile
import time
import socket
from pathlib import Path
from jupyter_client.blocking.client import BlockingKernelClient
@ -138,7 +139,7 @@ class JupyterKernelManager:
connection_file = self._create_connection_file(kernel_id, connection_info)
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 {}
if bridge_socket_path:
container_volumes[bridge_socket_path] = {
@ -146,17 +147,30 @@ class JupyterKernelManager:
"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(
image=self.image,
command=[
"python", "-m", "ipykernel_launcher",
"-f", f"/tmp/kernel-{kernel_id}.json"
"-f", container_connection_path
],
resource_limits=self.resource_limits,
volumes=container_volumes,
# TODO: Port mappings for ZMQ
# TODO: Mount connection file into container
network_mode="host", # Use host network for ZMQ communication
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(
@ -168,8 +182,9 @@ class JupyterKernelManager:
# Start container
self.container_manager.start_container(container_id)
# Wait for kernel to be ready
time.sleep(2) # TODO: Better readiness check
# Wait for kernel to be ready with polling
if not self._wait_for_kernel_ready(connection_info, timeout=30):
raise KernelError(f"Kernel {kernel_id} failed to start within timeout")
# Connect client
client = self._connect_client(connection_info)
@ -451,15 +466,18 @@ _info
return self.kernels[kernel_id]
def _generate_connection_info(self) -> Dict[str, Any]:
"""Generate ZMQ connection information."""
"""Generate ZMQ connection information with allocated ports."""
import secrets
# Allocate 5 ports for ZMQ channels
ports = self._allocate_ports(5)
return {
"shell_port": 0, # Let ZMQ assign
"iopub_port": 0,
"stdin_port": 0,
"control_port": 0,
"hb_port": 0,
"shell_port": ports[0],
"iopub_port": ports[1],
"stdin_port": ports[2],
"control_port": ports[3],
"hb_port": ports[4],
"ip": "127.0.0.1",
"key": secrets.token_hex(32),
"transport": "tcp",
@ -467,6 +485,26 @@ _info
"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(
self,
kernel_id: str,
@ -497,6 +535,44 @@ _info
except Exception:
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(
self,
client: BlockingKernelClient,