HansAschauer-Privat/mcp-forge
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Implement ZMQ port management and kernel readiness checks

Browse source 
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
Download patch file Download diff file Expand all files Collapse all files

39
docs/JUPYTER_IMPLEMENTATION_STATUS.md
View file

@ -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):

102
src/mcp_forge/execution/jupyter/kernel.py
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,

21
src/mcp_forge/podman/containers.py
View file

@ -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