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:
parent
9cc43e4166
commit
14aba0a048
3 changed files with 128 additions and 34 deletions
|
|
@ -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,
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue