From 14aba0a048280a80ea69dcc7bf5a983aad6ca243 Mon Sep 17 00:00:00 2001 From: Hans Aschauer Date: Sat, 7 Feb 2026 08:16:51 +0100 Subject: [PATCH] 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. --- docs/JUPYTER_IMPLEMENTATION_STATUS.md | 39 +++++---- src/mcp_forge/execution/jupyter/kernel.py | 102 +++++++++++++++++++--- src/mcp_forge/podman/containers.py | 21 ++++- 3 files changed, 128 insertions(+), 34 deletions(-) diff --git a/docs/JUPYTER_IMPLEMENTATION_STATUS.md b/docs/JUPYTER_IMPLEMENTATION_STATUS.md index b57dc0d..6e218f2 100644 --- a/docs/JUPYTER_IMPLEMENTATION_STATUS.md +++ b/docs/JUPYTER_IMPLEMENTATION_STATUS.md @@ -55,7 +55,20 @@ - Error handling for injection failures - 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 - `ipykernel` removed from server (will be in container image) - `pyzmq>=27.1.0` for ZMQ support @@ -68,17 +81,7 @@ ### 🚧 In Progress / TODO -1. **Container Configuration** - - [ ] 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** +1. **Restart Implementation** - [ ] Use `KernelManager` instead of just `BlockingKernelClient` - [ ] Proper restart via `KernelManager.restart_kernel()` - [ ] Handle restart failures gracefully @@ -145,16 +148,16 @@ Current test results: **20 failed, 2 passed** 1. **Immediate** (Core functionality): ``` - 1. Implement port mapping in ContainerConfig - 2. Mount connection file into container - 3. Test kernel startup with real container + 1. Update container image with ipykernel + 2. Test kernel startup with real container + 3. Mock tests for unit testing ``` 2. **Short-term** (Stability): ``` - 4. Implement proper readiness check - 5. Update container image with ipykernel - 6. Mock tests for unit testing + 4. Implement proper kernel restart + 5. Error handling improvements + 6. Port conflict handling ``` 3. **Medium-term** (Production-ready): diff --git a/src/mcp_forge/execution/jupyter/kernel.py b/src/mcp_forge/execution/jupyter/kernel.py index 9aa913a..3cfdf77 100644 --- a/src/mcp_forge/execution/jupyter/kernel.py +++ b/src/mcp_forge/execution/jupyter/kernel.py @@ -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, diff --git a/src/mcp_forge/podman/containers.py b/src/mcp_forge/podman/containers.py index d3d097f..e84f3f7 100644 --- a/src/mcp_forge/podman/containers.py +++ b/src/mcp_forge/podman/containers.py @@ -27,7 +27,9 @@ class ContainerConfig: volumes: Optional[Dict[str, dict]] = None, resource_limits: Optional[ResourceLimits] = 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. @@ -40,6 +42,8 @@ class ContainerConfig: resource_limits: Resource limits to apply working_dir: Working directory in container (None to use image default) 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.command = command or [] @@ -48,16 +52,19 @@ class ContainerConfig: self.resource_limits = resource_limits self.working_dir = working_dir self.user = user + self.network_mode = network_mode + self.port_bindings = port_bindings or {} def to_podman_params(self) -> dict: """ Convert to Podman container create parameters. Ensures all security requirements are included: - - network_mode: none + - network_mode: configurable (default 'none' for security) - read_only: True - security_opt: ["no-new-privileges"] - resource limits + - port_bindings: for host networking mode Returns: Dictionary of parameters for Podman containers.create() @@ -68,7 +75,7 @@ class ContainerConfig: "environment": self.environment, "user": self.user, # Security requirements - "network_mode": "none", + "network_mode": self.network_mode, "read_only": True, "security_opt": ["no-new-privileges"], } @@ -77,6 +84,14 @@ class ContainerConfig: if self.working_dir is not None: 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 if self.volumes: params["volumes"] = self.volumes