398 lines
12 KiB
Python
398 lines
12 KiB
Python
|
"""Jupyter kernel management for stateful execution."""
|
||
|
|
|
||
|
|
from typing import Dict, Optional, List, Any
|
||
|
|
from dataclasses import dataclass, field
|
||
|
|
from datetime import datetime, timedelta
|
||
|
|
import uuid
|
||
|
|
import json
|
||
|
|
import sys
|
||
|
|
import io
|
||
|
|
|
||
|
|
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
|
||
|
|
started_at: datetime
|
||
|
|
last_activity: datetime
|
||
|
|
namespace: Dict[str, Any] = field(default_factory=dict)
|
||
|
|
|
||
|
|
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(),
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
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
|