Implement real Jupyter backend with jupyter-client
- 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
parent
372af75b90
commit
244a3e5574
6 changed files with 940 additions and 253 deletions
|
|
@ -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
|
||||
|
|
|
|||
418
src/mcp_forge/execution/jupyter/kernel_old.py
Normal file
418
src/mcp_forge/execution/jupyter/kernel_old.py
Normal file
|
|
@ -0,0 +1,418 @@
|
|||
"""
|
||||
Jupyter kernel management for stateful execution.
|
||||
|
||||
This module implements a real Jupyter kernel manager that:
|
||||
- Uses jupyter-client (runs on host) to connect to kernels
|
||||
- Runs ipykernel processes inside Podman containers
|
||||
- Communicates via ZMQ protocol
|
||||
- Maintains 1:1 mapping of sessions to containers/kernels
|
||||
"""
|
||||
|
||||
from typing import Dict, Optional, List, Any
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime, timedelta
|
||||
import uuid
|
||||
import json
|
||||
import tempfile
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
from jupyter_client import BlockingKernelClient
|
||||
from jupyter_client.manager import KernelManager
|
||||
import zmq
|
||||
|
||||
from mcp_forge.podman.containers import SecureContainerManager, ContainerConfig
|
||||
from mcp_forge.security.resource_limits import ResourceLimits
|
||||
from mcp_forge.execution.simple.executor import ExecutionResult
|
||||
|
||||
|
||||
class KernelError(Exception):
|
||||
"""Raised when kernel operations fail."""
|
||||
pass
|
||||
|
||||
|
||||
@dataclass
|
||||
class KernelInfo:
|
||||
"""Information about a running kernel."""
|
||||
kernel_id: str
|
||||
container_id: str
|
||||
session_id: str
|
||||
connection_info: Dict[str, Any] # ZMQ ports and keys
|
||||
started_at: datetime
|
||||
last_activity: datetime
|
||||
client: Optional[BlockingKernelClient] = None
|
||||
|
||||
def to_dict(self) -> dict:
|
||||
"""Convert to dictionary for JSON serialization."""
|
||||
return {
|
||||
"kernel_id": self.kernel_id,
|
||||
"container_id": self.container_id,
|
||||
"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.
|
||||
|
||||
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.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
container_manager: SecureContainerManager,
|
||||
image: str,
|
||||
resource_limits: Optional[ResourceLimits]
|
||||
):
|
||||
"""
|
||||
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)
|
||||
"""
|
||||
self.container_manager = container_manager
|
||||
self.image = image
|
||||
self.resource_limits = resource_limits
|
||||
self.kernels: Dict[str, KernelInfo] = {}
|
||||
|
||||
def start_kernel(
|
||||
self,
|
||||
session_id: str,
|
||||
volumes: Optional[Dict[str, dict]] = None
|
||||
) -> str:
|
||||
"""
|
||||
Start a new kernel in a container.
|
||||
|
||||
Creates a long-running container with Python that will accept
|
||||
and execute code, maintaining namespace state between executions.
|
||||
|
||||
Args:
|
||||
session_id: Session ID this kernel belongs to
|
||||
volumes: Optional volume mounts
|
||||
|
||||
Returns:
|
||||
kernel_id: Unique identifier for the kernel
|
||||
"""
|
||||
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 {}
|
||||
)
|
||||
|
||||
# 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)
|
||||
|
||||
# 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
|
||||
|
||||
def execute_code(
|
||||
self,
|
||||
kernel_id: str,
|
||||
code: str,
|
||||
timeout: int = 300
|
||||
) -> ExecutionResult:
|
||||
"""
|
||||
Execute code in the kernel.
|
||||
|
||||
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
|
||||
|
||||
Args:
|
||||
kernel_id: ID of kernel to execute in
|
||||
code: Python code to execute
|
||||
timeout: Maximum execution time
|
||||
|
||||
Returns:
|
||||
ExecutionResult with output and status
|
||||
|
||||
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.kernels[kernel_id]
|
||||
|
||||
# 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
|
||||
|
||||
result_value = None
|
||||
error = None
|
||||
|
||||
try:
|
||||
# Redirect stdout/stderr
|
||||
sys.stdout = stdout_capture
|
||||
sys.stderr = stderr_capture
|
||||
|
||||
# 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()
|
||||
|
||||
execution_time = time.time() - start_time
|
||||
|
||||
return ExecutionResult(
|
||||
success=(error is None),
|
||||
stdout=stdout,
|
||||
stderr=stderr,
|
||||
result=result_value,
|
||||
execution_time=execution_time,
|
||||
exit_code=0 if error is None else 1,
|
||||
error=error
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
execution_time = time.time() - start_time
|
||||
return ExecutionResult(
|
||||
success=False,
|
||||
stdout="",
|
||||
stderr="",
|
||||
result=None,
|
||||
execution_time=execution_time,
|
||||
exit_code=1,
|
||||
error=f"Execution failed: {str(e)}"
|
||||
)
|
||||
|
||||
def shutdown_kernel(self, kernel_id: str) -> None:
|
||||
"""
|
||||
Shutdown kernel and cleanup container.
|
||||
|
||||
Args:
|
||||
kernel_id: ID of kernel to shutdown
|
||||
|
||||
Raises:
|
||||
KernelError: If kernel not found
|
||||
"""
|
||||
if kernel_id not in self.kernels:
|
||||
raise KernelError(f"Kernel {kernel_id} not found")
|
||||
|
||||
kernel_info = self.kernels[kernel_id]
|
||||
|
||||
# Stop and remove container
|
||||
try:
|
||||
self.container_manager.stop_container(kernel_info.container_id, timeout=10)
|
||||
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]
|
||||
|
||||
def inspect_namespace(self, kernel_id: str) -> List[str]:
|
||||
"""
|
||||
Get list of variables in kernel namespace.
|
||||
|
||||
Args:
|
||||
kernel_id: ID of kernel to inspect
|
||||
|
||||
Returns:
|
||||
List of variable names (excluding private vars)
|
||||
|
||||
Raises:
|
||||
KernelError: If kernel not found
|
||||
"""
|
||||
if kernel_id not in self.kernels:
|
||||
raise KernelError(f"Kernel {kernel_id} not found")
|
||||
|
||||
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
|
||||
|
||||
def get_variable_info(
|
||||
self,
|
||||
kernel_id: str,
|
||||
variable_name: str
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
Get information about a variable.
|
||||
|
||||
Args:
|
||||
kernel_id: ID of kernel
|
||||
variable_name: Name of variable to inspect
|
||||
|
||||
Returns:
|
||||
Dictionary with type, size, and repr info
|
||||
|
||||
Raises:
|
||||
KernelError: If kernel or variable not found
|
||||
"""
|
||||
if kernel_id not in self.kernels:
|
||||
raise KernelError(f"Kernel {kernel_id} not found")
|
||||
|
||||
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__'):
|
||||
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
|
||||
|
||||
def restart_kernel(self, kernel_id: str) -> None:
|
||||
"""
|
||||
Restart kernel (reset namespace).
|
||||
|
||||
Args:
|
||||
kernel_id: ID of kernel to restart
|
||||
|
||||
Raises:
|
||||
KernelError: If kernel not found
|
||||
"""
|
||||
if kernel_id not in self.kernels:
|
||||
raise KernelError(f"Kernel {kernel_id} not found")
|
||||
|
||||
# Clear namespace to reset state
|
||||
kernel_info = self.kernels[kernel_id]
|
||||
kernel_info.namespace.clear()
|
||||
kernel_info.last_activity = datetime.utcnow()
|
||||
|
||||
def cleanup_idle_kernels(self, idle_timeout: timedelta) -> int:
|
||||
"""
|
||||
Cleanup kernels idle longer than timeout.
|
||||
|
||||
Args:
|
||||
idle_timeout: Maximum idle time before cleanup
|
||||
|
||||
Returns:
|
||||
Number of kernels cleaned up
|
||||
"""
|
||||
now = datetime.utcnow()
|
||||
kernels_to_remove = []
|
||||
|
||||
for kernel_id, kernel_info in self.kernels.items():
|
||||
idle_time = now - kernel_info.last_activity
|
||||
if idle_time > idle_timeout:
|
||||
kernels_to_remove.append(kernel_id)
|
||||
|
||||
# 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)
|
||||
|
||||
def _wrap_code_with_namespace(self, code: str, namespace: Dict[str, Any]) -> str:
|
||||
"""
|
||||
Wrap code to load/save namespace.
|
||||
|
||||
This is a helper for the real implementation where code would be
|
||||
executed in a container with namespace persistence.
|
||||
|
||||
Args:
|
||||
code: User code to wrap
|
||||
namespace: Current namespace state
|
||||
|
||||
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
|
||||
Loading…
Add table
Add a link
Reference in a new issue