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

Implement real Jupyter backend with jupyter-client

Browse source 
- Replaced mock exec() implementation with real Jupyter protocol
- Uses jupyter-client (host) to connect to ipykernel (container) via ZMQ
- 1:1 mapping: one container per session, one kernel per container
- Proper Jupyter message protocol for code execution
- Kernel lifecycle management (start, execute, shutdown, restart)
- Namespace inspection via introspection code
- Idle kernel cleanup
- Connection file management
- Backed up old implementation as kernel_old.py

TODO:
- Container image needs ipykernel installed
- Need to implement proper port mapping for ZMQ
- Need to mount connection file into container
- Add better kernel readiness check
- Implement restart_kernel properly with KernelManager
- Write tests
This commit is contained in:
Hans Aschauer 2026-02-07 08:06:44 +01:00
parent 372af75b90
commit 244a3e5574
6 changed files with 940 additions and 253 deletions
Download patch file Download diff file Expand all files Collapse all files

521
src/mcp_forge/execution/jupyter/kernel.py
View file

@ -1,12 +1,24 @@
"""Jupyter kernel management for stateful execution."""
"""
Real Jupyter kernel management for stateful execution.
This module implements proper Jupyter kernel management:
- jupyter-client runs on host (MCP-Forge server)
- ipykernel runs inside Podman containers
- Communication via ZMQ protocol
- 1:1 mapping: one container per session, one kernel per container
"""
from typing import Dict, Optional, List, Any
from dataclasses import dataclass, field
from dataclasses import dataclass
from datetime import datetime, timedelta
import uuid
import json
import sys
import io
import tempfile
import time
from pathlib import Path
from jupyter_client.blocking.client import BlockingKernelClient
import zmq
from mcp_forge.podman.containers import SecureContainerManager, ContainerConfig
from mcp_forge.security.resource_limits import ResourceLimits
@ -24,9 +36,11 @@ class KernelInfo:
kernel_id: str
container_id: str
session_id: str
connection_file: Path
connection_info: Dict[str, Any] # ZMQ ports and keys
started_at: datetime
last_activity: datetime
namespace: Dict[str, Any] = field(default_factory=dict)
client: Optional[BlockingKernelClient] = None
def to_dict(self) -> dict:
"""Convert to dictionary for JSON serialization."""
@ -36,31 +50,46 @@ class KernelInfo:
"session_id": self.session_id,
"started_at": self.started_at.isoformat(),
"last_activity": self.last_activity.isoformat(),
"connection_info": {
"shell_port": self.connection_info.get("shell_port"),
"iopub_port": self.connection_info.get("iopub_port"),
"stdin_port": self.connection_info.get("stdin_port"),
"control_port": self.connection_info.get("control_port"),
"hb_port": self.connection_info.get("hb_port"),
}
}
class JupyterKernelManager:
"""
Manages IPython kernels in containers for stateful execution.
Manages IPython kernels in containers via jupyter-client.
This is a simplified implementation that uses containers to maintain
state between executions. Each kernel runs in its own container and
maintains a Python namespace that persists across execute calls.
Architecture:
- This class runs on host (MCP-Forge server process)
- Creates one container per session with ipykernel running inside
- Connects to kernel via ZMQ protocol (jupyter-client)
- Communicates using Jupyter message protocol
Each session gets:
- Dedicated container
- Dedicated kernel process
- Isolated Python namespace
- Independent resource limits
"""
def __init__(
self,
container_manager: SecureContainerManager,
image: str,
resource_limits: Optional[ResourceLimits]
resource_limits: Optional[ResourceLimits] = None
):
"""
Initialize kernel manager.
Args:
container_manager: Container lifecycle manager
image: Docker/Podman image with Python/IPython
resource_limits: Default resource limits for kernels (None to disable) (None to disable)
image: Docker/Podman image with ipykernel installed
resource_limits: Default resource limits for kernels
"""
self.container_manager = container_manager
self.image = image
@ -73,49 +102,88 @@ class JupyterKernelManager:
volumes: Optional[Dict[str, dict]] = None
) -> str:
"""
Start a new kernel in a container.
Start IPython kernel in dedicated container.
Creates a long-running container with Python that will accept
and execute code, maintaining namespace state between executions.
Process:
1. Generate ZMQ connection info (ports, keys)
2. Create connection file
3. Create container with ipykernel command
4. Start container
5. Wait for kernel to be ready
6. Connect jupyter-client to kernel via ZMQ
7. Verify kernel is responsive
Args:
session_id: Session ID this kernel belongs to
volumes: Optional volume mounts
Returns:
kernel_id: Unique identifier for the kernel
kernel_id: Unique identifier for this kernel
Raises:
KernelError: If kernel startup fails
"""
kernel_id = f"kernel-{uuid.uuid4().hex[:16]}"
# Create container configuration for long-running kernel
# We use a shell that stays running so we can exec into it
config = ContainerConfig(
image=self.image,
command=["sleep", "infinity"], # Keep container running
resource_limits=self.resource_limits,
volumes=volumes or {}
)
# Generate connection info
connection_info = self._generate_connection_info()
# Create and start container
container_id = self.container_manager.create_container(
config,
session_id=session_id,
name=f"kernel-{kernel_id}"
)
self.container_manager.start_container(container_id)
# Create connection file
connection_file = self._create_connection_file(kernel_id, connection_info)
# Register kernel
now = datetime.utcnow()
kernel_info = KernelInfo(
kernel_id=kernel_id,
container_id=container_id,
session_id=session_id,
started_at=now,
last_activity=now
)
self.kernels[kernel_id] = kernel_info
return kernel_id
try:
# Create container with ipykernel
config = ContainerConfig(
image=self.image,
command=[
"python", "-m", "ipykernel_launcher",
"-f", f"/tmp/kernel-{kernel_id}.json"
],
resource_limits=self.resource_limits,
volumes=volumes or {},
# TODO: Port mappings for ZMQ
# TODO: Mount connection file into container
)
container_id = self.container_manager.create_container(
config,
session_id=session_id,
name=f"jupyter-{kernel_id}"
)
# Start container
self.container_manager.start_container(container_id)
# Wait for kernel to be ready
time.sleep(2) # TODO: Better readiness check
# Connect client
client = self._connect_client(connection_info)
# Verify kernel is responsive
if not self._verify_kernel(client):
raise KernelError(f"Kernel {kernel_id} not responsive")
# Register kernel
now = datetime.utcnow()
kernel_info = KernelInfo(
kernel_id=kernel_id,
container_id=container_id,
session_id=session_id,
connection_file=connection_file,
connection_info=connection_info,
started_at=now,
last_activity=now,
client=client
)
self.kernels[kernel_id] = kernel_info
return kernel_id
except Exception as e:
# Cleanup on failure
connection_file.unlink(missing_ok=True)
raise KernelError(f"Failed to start kernel: {e}") from e
def execute_code(
self,
@ -124,161 +192,148 @@ class JupyterKernelManager:
timeout: int = 300
) -> ExecutionResult:
"""
Execute code in the kernel.
Execute code in kernel via ZMQ.
This is a simplified implementation that:
1. Validates kernel exists
2. Wraps code to capture output and maintain namespace
3. Executes in the kernel's container
4. Returns results
Uses jupyter-client to:
1. Send execute_request message
2. Receive stream (stdout/stderr) messages
3. Receive execute_result/display_data messages
4. Collect and parse all output
Args:
kernel_id: ID of kernel to execute in
kernel_id: Kernel to execute in
code: Python code to execute
timeout: Maximum execution time
timeout: Maximum execution time in seconds
Returns:
ExecutionResult with output and status
ExecutionResult with stdout, stderr, result
Raises:
KernelError: If kernel not found or execution fails
"""
if kernel_id not in self.kernels:
raise KernelError(f"Kernel {kernel_id} not found")
kernel_info = self._get_kernel(kernel_id)
client = kernel_info.client
kernel_info = self.kernels[kernel_id]
if not client:
raise KernelError(f"Kernel {kernel_id} has no connected client")
# Update activity
kernel_info.last_activity = datetime.utcnow()
# For simplified implementation, we execute code by creating
# a Python script that:
# 1. Loads namespace from kernel_info
# 2. Executes user code
# 3. Saves namespace back
# 4. Returns result as JSON
# Execute in container using Python
# In real implementation, this would use docker exec or similar
# For now, we simulate execution with proper stdout/stderr capture
import time
start_time = time.time()
try:
# Capture stdout and stderr
stdout_capture = io.StringIO()
stderr_capture = io.StringIO()
old_stdout = sys.stdout
old_stderr = sys.stderr
# Execute code
_msg_id = client.execute(code, silent=False, store_history=True)
result_value = None
error = None
# Collect output
stdout_parts = []
stderr_parts = []
result = None
try:
# Redirect stdout/stderr
sys.stdout = stdout_capture
sys.stderr = stderr_capture
# Wait for execution to complete
while True:
try:
msg = client.get_iopub_msg(timeout=timeout)
msg_type = msg['header']['msg_type']
content = msg['content']
if msg_type == 'stream':
if content['name'] == 'stdout':
stdout_parts.append(content['text'])
elif content['name'] == 'stderr':
stderr_parts.append(content['text'])
elif msg_type == 'execute_result':
result = content.get('data', {}).get('text/plain', '')
elif msg_type == 'error':
stderr_parts.append('\n'.join(content['traceback']))
elif msg_type == 'status':
if content['execution_state'] == 'idle':
break
# Execute and update namespace
exec_globals = kernel_info.namespace.copy()
exec(code, exec_globals)
# Update kernel namespace
kernel_info.namespace.update(exec_globals)
# Try to get result from last expression
result_value = exec_globals.get('_', None)
except SyntaxError as e:
error = f"SyntaxError: {e.msg}"
stderr_capture.write(f"{error}\n")
except Exception as e:
error = f"{type(e).__name__}: {str(e)}"
stderr_capture.write(f"{error}\n")
finally:
# Restore stdout/stderr
sys.stdout = old_stdout
sys.stderr = old_stderr
# Get captured output
stdout = stdout_capture.getvalue()
stderr = stderr_capture.getvalue()
except zmq.error.Again:
break
execution_time = time.time() - start_time
# Update last activity
kernel_info.last_activity = datetime.utcnow()
return ExecutionResult(
success=(error is None),
stdout=stdout,
stderr=stderr,
result=result_value,
success=True,
stdout=''.join(stdout_parts),
stderr=''.join(stderr_parts),
result=result,
execution_time=execution_time,
exit_code=0 if error is None else 1,
error=error
exit_code=0
)
except Exception as e:
execution_time = time.time() - start_time
return ExecutionResult(
success=False,
stdout="",
stderr="",
stdout='',
stderr=str(e),
result=None,
execution_time=execution_time,
exit_code=1,
error=f"Execution failed: {str(e)}"
error=str(e)
)
def shutdown_kernel(self, kernel_id: str) -> None:
"""
Shutdown kernel and cleanup container.
1. Send shutdown_request via ZMQ
2. Wait for kernel shutdown
3. Stop and remove container
4. Cleanup connection file
Args:
kernel_id: ID of kernel to shutdown
Raises:
KernelError: If kernel not found
kernel_id: Kernel to shutdown
"""
if kernel_id not in self.kernels:
raise KernelError(f"Kernel {kernel_id} not found")
kernel_info = self._get_kernel(kernel_id)
kernel_info = self.kernels[kernel_id]
# Stop and remove container
try:
self.container_manager.stop_container(kernel_info.container_id, timeout=10)
# Shutdown kernel
if kernel_info.client:
kernel_info.client.shutdown()
kernel_info.client.stop_channels()
# Stop and remove container
self.container_manager.stop_container(kernel_info.container_id)
self.container_manager.remove_container(kernel_info.container_id)
except Exception as e:
# Log but don't fail - best effort cleanup
pass
# Remove from registry
del self.kernels[kernel_id]
# Cleanup connection file
kernel_info.connection_file.unlink(missing_ok=True)
finally:
# Remove from registry
del self.kernels[kernel_id]
def inspect_namespace(self, kernel_id: str) -> List[str]:
"""
Get list of variables in kernel namespace.
Executes introspection code:
[var for var in dir() if not var.startswith('_')]
Args:
kernel_id: ID of kernel to inspect
kernel_id: Kernel to inspect
Returns:
List of variable names (excluding private vars)
Raises:
KernelError: If kernel not found
List of variable names
"""
if kernel_id not in self.kernels:
raise KernelError(f"Kernel {kernel_id} not found")
code = "[var for var in dir() if not var.startswith('_')]"
result = self.execute_code(kernel_id, code, timeout=5)
kernel_info = self.kernels[kernel_id]
# Filter out private variables and builtins
variables = [
name for name in kernel_info.namespace.keys()
if not name.startswith('_') and name not in ['__builtins__']
]
return variables
if result.success and result.result:
# Parse result (it's a string representation of a list)
try:
return eval(result.result) # nosec - controlled code
except Exception:
return []
return []
def get_variable_info(
self,
@ -286,68 +341,66 @@ class JupyterKernelManager:
variable_name: str
) -> Dict[str, Any]:
"""
Get information about a variable.
Get detailed information about a variable.
Executes introspection code to get:
- type(var).__name__
- sys.getsizeof(var) if available
- var.shape if hasattr(var, 'shape')
- repr(var)[:100]
Args:
kernel_id: ID of kernel
kernel_id: Kernel to inspect
variable_name: Name of variable to inspect
Returns:
Dictionary with type, size, and repr info
Raises:
KernelError: If kernel or variable not found
Dict with type, size, shape, repr
"""
if kernel_id not in self.kernels:
raise KernelError(f"Kernel {kernel_id} not found")
code = f"""
import sys
_var = {variable_name}
_info = {{
'type': type(_var).__name__,
'repr': repr(_var)[:100],
}}
try:
_info['size_bytes'] = sys.getsizeof(_var)
except:
pass
if hasattr(_var, 'shape'):
_info['shape'] = _var.shape
_info
"""
result = self.execute_code(kernel_id, code, timeout=5)
kernel_info = self.kernels[kernel_id]
if variable_name not in kernel_info.namespace:
raise KernelError(f"Variable {variable_name} not found in kernel namespace")
value = kernel_info.namespace[variable_name]
info = {
"type": type(value).__name__,
"repr": repr(value)[:100], # Truncate long reprs
}
# Add size for sized objects
if hasattr(value, '__len__'):
if result.success and result.result:
try:
info["size"] = len(value)
except:
pass
# Add shape for array-like objects
if hasattr(value, 'shape'):
try:
info["shape"] = value.shape
except:
pass
return info
return eval(result.result) # nosec - controlled code
except Exception:
return {}
return {}
def restart_kernel(self, kernel_id: str) -> None:
"""
Restart kernel (reset namespace).
Restart kernel (namespace reset, container kept).
Strategy: shutdown current kernel and start new one in same container.
Note: In a full implementation, we'd use KernelManager.restart_kernel().
Args:
kernel_id: ID of kernel to restart
Raises:
KernelError: If kernel not found
kernel_id: Kernel to restart
"""
if kernel_id not in self.kernels:
raise KernelError(f"Kernel {kernel_id} not found")
kernel_info = self._get_kernel(kernel_id)
# Clear namespace to reset state
kernel_info = self.kernels[kernel_id]
kernel_info.namespace.clear()
# For now, just record activity - full restart implementation requires
# KernelManager integration (not just BlockingKernelClient)
# TODO: Implement proper kernel restart via KernelManager
kernel_info.last_activity = datetime.utcnow()
def cleanup_idle_kernels(self, idle_timeout: timedelta) -> int:
def cleanup_idle_kernels(
self,
idle_timeout: timedelta
) -> int:
"""
Cleanup kernels idle longer than timeout.
@ -358,40 +411,70 @@ class JupyterKernelManager:
Number of kernels cleaned up
"""
now = datetime.utcnow()
kernels_to_remove = []
cleaned_up = 0
for kernel_id, kernel_info in self.kernels.items():
for kernel_id in list(self.kernels.keys()):
kernel_info = self.kernels[kernel_id]
idle_time = now - kernel_info.last_activity
if idle_time > idle_timeout:
kernels_to_remove.append(kernel_id)
try:
self.shutdown_kernel(kernel_id)
cleaned_up += 1
except Exception:
pass # Continue cleanup even if one fails
# Shutdown idle kernels
for kernel_id in kernels_to_remove:
try:
self.shutdown_kernel(kernel_id)
except Exception:
# Best effort cleanup
pass
return len(kernels_to_remove)
return cleaned_up
def _wrap_code_with_namespace(self, code: str, namespace: Dict[str, Any]) -> str:
"""
Wrap code to load/save namespace.
def _get_kernel(self, kernel_id: str) -> KernelInfo:
"""Get kernel info or raise error."""
if kernel_id not in self.kernels:
raise KernelError(f"Kernel {kernel_id} not found")
return self.kernels[kernel_id]
def _generate_connection_info(self) -> Dict[str, Any]:
"""Generate ZMQ connection information."""
import secrets
This is a helper for the real implementation where code would be
executed in a container with namespace persistence.
return {
"shell_port": 0, # Let ZMQ assign
"iopub_port": 0,
"stdin_port": 0,
"control_port": 0,
"hb_port": 0,
"ip": "127.0.0.1",
"key": secrets.token_hex(32),
"transport": "tcp",
"signature_scheme": "hmac-sha256",
"kernel_name": "python3"
}
def _create_connection_file(
self,
kernel_id: str,
connection_info: Dict[str, Any]
) -> Path:
"""Create connection file for kernel."""
# Create temp file
fd, path = tempfile.mkstemp(suffix=f"-kernel-{kernel_id}.json")
Args:
code: User code to wrap
namespace: Current namespace state
# Write connection info
with open(fd, 'w') as f:
json.dump(connection_info, f)
Returns:
Wrapped code with namespace handling
"""
# In real implementation, this would serialize namespace,
# inject it into container execution, run code, and extract
# updated namespace.
# For this simplified version, we don't need the wrapping
# since we're executing directly in Python.
return code
return Path(path)
def _connect_client(self, connection_info: Dict[str, Any]) -> BlockingKernelClient:
"""Connect jupyter-client to kernel."""
client = BlockingKernelClient()
client.load_connection_info(connection_info)
client.start_channels()
return client
def _verify_kernel(self, client: BlockingKernelClient, timeout: int = 10) -> bool:
"""Verify kernel is responsive."""
try:
client.wait_for_ready(timeout=timeout)
return True
except Exception:
return False