Create standalone pod_executor package
Created a standalone code execution package independent of MCP-Forge: Structure: - pod_executor/security/ - Resource limits, audit protocols, validation - pod_executor/containers/ - Podman client and container management - pod_executor/simple/ - Stateless code executor - pod_executor/jupyter/ - Stateful Jupyter backend with sessions Key changes: - Removed ForgeConfig dependency - all parameters explicit - Audit logger now a protocol with NullAuditLogger/SimpleFileAuditLogger - Validator now a protocol with NoOpValidator/BasicValidator - All imports updated to pod_executor namespace - Audit calls use simple strings instead of enums Benefits: - Standalone package usable without MCP-Forge - Clear separation between execution engine and MCP protocol - Easier testing and development - Reusable in other projects
This commit is contained in:
parent
8b6b237be9
commit
db75b822f4
14 changed files with 2842 additions and 0 deletions
82
src/pod_executor/__init__.py
Normal file
82
src/pod_executor/__init__.py
Normal file
|
|
@ -0,0 +1,82 @@
|
||||||
|
"""
|
||||||
|
Pod Executor - Standalone Python code execution in Podman containers.
|
||||||
|
|
||||||
|
Provides stateless and stateful (Jupyter) code execution backends with
|
||||||
|
security isolation via Podman containers.
|
||||||
|
|
||||||
|
Examples:
|
||||||
|
# Simple stateless execution
|
||||||
|
from pod_executor import CodeExecutor, ResourceLimits
|
||||||
|
from pod_executor.containers import SecureContainerManager, PodmanClient
|
||||||
|
from pod_executor.security import NoOpValidator, NullAuditLogger
|
||||||
|
|
||||||
|
client = PodmanClient(socket_path="/run/podman/podman.sock",
|
||||||
|
validator=NoOpValidator(),
|
||||||
|
audit_logger=NullAuditLogger())
|
||||||
|
container_manager = SecureContainerManager(client, NoOpValidator(), NullAuditLogger())
|
||||||
|
limits = ResourceLimits(memory="512m", cpu_quota=100000, storage="1g", timeout=30)
|
||||||
|
executor = CodeExecutor(container_manager, "python:3.12", limits)
|
||||||
|
|
||||||
|
result = executor.execute("print('Hello World')")
|
||||||
|
print(result.stdout)
|
||||||
|
|
||||||
|
# Stateful Jupyter execution
|
||||||
|
from pod_executor import JupyterBackend
|
||||||
|
|
||||||
|
backend = JupyterBackend(
|
||||||
|
container_manager=container_manager,
|
||||||
|
image="jupyter/base-notebook",
|
||||||
|
default_timeout=300
|
||||||
|
)
|
||||||
|
|
||||||
|
result = backend.execute("x = 42", session_id="my-session")
|
||||||
|
result = backend.execute("print(x * 2)", session_id="my-session")
|
||||||
|
"""
|
||||||
|
|
||||||
|
from pod_executor.simple.executor import CodeExecutor, ExecutionResult
|
||||||
|
from pod_executor.jupyter.backend import JupyterBackend
|
||||||
|
from pod_executor.jupyter.kernel import JupyterKernelManager, KernelError
|
||||||
|
from pod_executor.jupyter.sessions import SessionManager, SessionState, SessionError
|
||||||
|
from pod_executor.containers.manager import SecureContainerManager, ContainerConfig
|
||||||
|
from pod_executor.containers.client import PodmanClient, PodmanConnectionError
|
||||||
|
from pod_executor.security.resource_limits import ResourceLimits
|
||||||
|
from pod_executor.security.audit import (
|
||||||
|
AuditLoggerProtocol,
|
||||||
|
NullAuditLogger,
|
||||||
|
SimpleFileAuditLogger,
|
||||||
|
)
|
||||||
|
from pod_executor.security.validation import (
|
||||||
|
SecurityError,
|
||||||
|
OperationValidatorProtocol,
|
||||||
|
NoOpValidator,
|
||||||
|
BasicValidator,
|
||||||
|
)
|
||||||
|
|
||||||
|
__version__ = "0.1.0"
|
||||||
|
|
||||||
|
__all__ = [
|
||||||
|
# Simple execution
|
||||||
|
"CodeExecutor",
|
||||||
|
"ExecutionResult",
|
||||||
|
# Jupyter execution
|
||||||
|
"JupyterBackend",
|
||||||
|
"JupyterKernelManager",
|
||||||
|
"KernelError",
|
||||||
|
"SessionManager",
|
||||||
|
"SessionState",
|
||||||
|
"SessionError",
|
||||||
|
# Container management
|
||||||
|
"SecureContainerManager",
|
||||||
|
"ContainerConfig",
|
||||||
|
"PodmanClient",
|
||||||
|
"PodmanConnectionError",
|
||||||
|
# Security
|
||||||
|
"ResourceLimits",
|
||||||
|
"AuditLoggerProtocol",
|
||||||
|
"NullAuditLogger",
|
||||||
|
"SimpleFileAuditLogger",
|
||||||
|
"SecurityError",
|
||||||
|
"OperationValidatorProtocol",
|
||||||
|
"NoOpValidator",
|
||||||
|
"BasicValidator",
|
||||||
|
]
|
||||||
14
src/pod_executor/containers/__init__.py
Normal file
14
src/pod_executor/containers/__init__.py
Normal file
|
|
@ -0,0 +1,14 @@
|
||||||
|
"""Container management for pod_executor."""
|
||||||
|
|
||||||
|
from pod_executor.containers.client import PodmanClient, PodmanConnectionError
|
||||||
|
from pod_executor.containers.manager import (
|
||||||
|
ContainerConfig,
|
||||||
|
SecureContainerManager,
|
||||||
|
)
|
||||||
|
|
||||||
|
__all__ = [
|
||||||
|
"PodmanClient",
|
||||||
|
"PodmanConnectionError",
|
||||||
|
"ContainerConfig",
|
||||||
|
"SecureContainerManager",
|
||||||
|
]
|
||||||
157
src/pod_executor/containers/client.py
Normal file
157
src/pod_executor/containers/client.py
Normal file
|
|
@ -0,0 +1,157 @@
|
||||||
|
"""
|
||||||
|
Podman client wrapper with security validation.
|
||||||
|
|
||||||
|
Wraps Podman API with security validation and error handling.
|
||||||
|
All container operations are validated against security policy.
|
||||||
|
"""
|
||||||
|
|
||||||
|
import os
|
||||||
|
from pathlib import Path
|
||||||
|
from typing import Optional, Union
|
||||||
|
from podman import PodmanClient as BasePodmanClient
|
||||||
|
|
||||||
|
from pod_executor.security.validation import OperationValidatorProtocol
|
||||||
|
from pod_executor.security.audit import AuditLoggerProtocol, NullAuditLogger
|
||||||
|
|
||||||
|
|
||||||
|
class PodmanConnectionError(Exception):
|
||||||
|
"""Raised when connection to Podman fails."""
|
||||||
|
pass
|
||||||
|
|
||||||
|
|
||||||
|
class PodmanClient:
|
||||||
|
"""
|
||||||
|
Wrapper around Podman API with security validation.
|
||||||
|
|
||||||
|
All container operations are validated against security policy
|
||||||
|
before being sent to Podman. Provides lazy connection and
|
||||||
|
proper error handling.
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
socket_path: Union[str, Path],
|
||||||
|
validator: OperationValidatorProtocol,
|
||||||
|
audit_logger: Optional[AuditLoggerProtocol] = None
|
||||||
|
):
|
||||||
|
"""
|
||||||
|
Initialize Podman client wrapper.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
socket_path: Path to Podman socket
|
||||||
|
validator: Operation validator for security checks
|
||||||
|
audit_logger: Optional audit logger (uses NullAuditLogger if None)
|
||||||
|
"""
|
||||||
|
self.socket_path = Path(socket_path)
|
||||||
|
self.validator = validator
|
||||||
|
self.audit_logger = audit_logger or NullAuditLogger()
|
||||||
|
self._client: Optional[BasePodmanClient] = None
|
||||||
|
|
||||||
|
def connect(self) -> None:
|
||||||
|
"""
|
||||||
|
Connect to Podman via socket.
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
PodmanConnectionError: If connection fails
|
||||||
|
"""
|
||||||
|
try:
|
||||||
|
# Verify socket exists and is accessible
|
||||||
|
self.verify_socket_access()
|
||||||
|
|
||||||
|
# Create Podman client with Unix socket
|
||||||
|
base_url = f"unix://{self.socket_path}"
|
||||||
|
self._client = BasePodmanClient(base_url=base_url)
|
||||||
|
|
||||||
|
# Test connection with ping (only if client supports it)
|
||||||
|
if hasattr(self._client, 'ping'):
|
||||||
|
self._client.ping()
|
||||||
|
|
||||||
|
except PodmanConnectionError:
|
||||||
|
# Re-raise our own exceptions
|
||||||
|
raise
|
||||||
|
except Exception as e:
|
||||||
|
raise PodmanConnectionError(
|
||||||
|
f"Failed to connect to Podman at {self.socket_path}: {e}"
|
||||||
|
) from e
|
||||||
|
|
||||||
|
def ping(self) -> bool:
|
||||||
|
"""
|
||||||
|
Test connection to Podman.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
True if connection is healthy
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
PodmanConnectionError: If not connected or ping fails
|
||||||
|
"""
|
||||||
|
if self._client is None:
|
||||||
|
raise PodmanConnectionError("Not connected to Podman")
|
||||||
|
|
||||||
|
try:
|
||||||
|
result = self._client.ping()
|
||||||
|
return result == "OK" or result is True
|
||||||
|
except Exception as e:
|
||||||
|
raise PodmanConnectionError(f"Ping failed: {e}") from e
|
||||||
|
|
||||||
|
def disconnect(self) -> None:
|
||||||
|
"""Disconnect from Podman and cleanup."""
|
||||||
|
if self._client is not None:
|
||||||
|
try:
|
||||||
|
self._client.close()
|
||||||
|
except Exception:
|
||||||
|
pass # Ignore errors during cleanup
|
||||||
|
finally:
|
||||||
|
self._client = None
|
||||||
|
|
||||||
|
def verify_socket_access(self) -> None:
|
||||||
|
"""
|
||||||
|
Verify that socket exists and is accessible.
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
PodmanConnectionError: If socket is not accessible
|
||||||
|
"""
|
||||||
|
if not self.socket_path.exists():
|
||||||
|
raise PodmanConnectionError(
|
||||||
|
f"Podman socket not found: {self.socket_path}"
|
||||||
|
)
|
||||||
|
|
||||||
|
if not os.access(self.socket_path, os.R_OK):
|
||||||
|
raise PodmanConnectionError(
|
||||||
|
f"Podman socket is not readable: {self.socket_path}"
|
||||||
|
)
|
||||||
|
|
||||||
|
def check_api_version(self) -> dict:
|
||||||
|
"""
|
||||||
|
Get Podman API version information.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Dictionary with version information
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
PodmanConnectionError: If not connected
|
||||||
|
"""
|
||||||
|
if self._client is None:
|
||||||
|
raise PodmanConnectionError("Not connected to Podman")
|
||||||
|
|
||||||
|
try:
|
||||||
|
return self._client.version()
|
||||||
|
except Exception as e:
|
||||||
|
raise PodmanConnectionError(
|
||||||
|
f"Failed to get API version: {e}"
|
||||||
|
) from e
|
||||||
|
|
||||||
|
@property
|
||||||
|
def client(self) -> BasePodmanClient:
|
||||||
|
"""
|
||||||
|
Get underlying Podman client (lazy connection).
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Connected Podman client
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
PodmanConnectionError: If connection fails
|
||||||
|
"""
|
||||||
|
if self._client is None:
|
||||||
|
self.connect()
|
||||||
|
assert self._client is not None # Type narrowing for mypy
|
||||||
|
return self._client
|
||||||
501
src/pod_executor/containers/manager.py
Normal file
501
src/pod_executor/containers/manager.py
Normal file
|
|
@ -0,0 +1,501 @@
|
||||||
|
""
|
||||||
|
Secure container management with security enforcement.
|
||||||
|
|
||||||
|
All container operations are validated against security policy
|
||||||
|
before being sent to Podman. Provides lifecycle management
|
||||||
|
with comprehensive audit logging.
|
||||||
|
""
|
||||||
|
|
||||||
|
from typing import Optional, Dict, List, Union
|
||||||
|
from datetime import datetime, timedelta
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
from pod_executor.containers.client import PodmanClient
|
||||||
|
from pod_executor.security.validation import OperationValidatorProtocol, SecurityError
|
||||||
|
from pod_executor.security.audit import AuditLoggerProtocol, NullAuditLogger
|
||||||
|
from pod_executor.security.resource_limits import ResourceLimits
|
||||||
|
|
||||||
|
|
||||||
|
class ContainerConfig:
|
||||||
|
""Container configuration with security defaults.""
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self",
|
||||||
|
image: str",
|
||||||
|
command: Optional[List[str]] = None",
|
||||||
|
environment: Optional[Dict[str, str]] = None",
|
||||||
|
volumes: Optional[Dict[str, dict]] = None",
|
||||||
|
resource_limits: Optional[ResourceLimits] = None",
|
||||||
|
working_dir: Optional[str] = None",
|
||||||
|
user: str = "1000:1000",
|
||||||
|
network_mode: str = "none",
|
||||||
|
port_bindings: Optional[Dict[str, int]] = None
|
||||||
|
):
|
||||||
|
""
|
||||||
|
Initialize container configuration.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
image: Container image to use
|
||||||
|
command: Command to run in container
|
||||||
|
environment: Environment variables
|
||||||
|
volumes: Volume mounts (host_path -> {bind, mode})
|
||||||
|
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 []
|
||||||
|
self.environment = environment or {}
|
||||||
|
self.volumes = volumes or {}
|
||||||
|
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: 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()
|
||||||
|
""
|
||||||
|
params = {
|
||||||
|
"image": self.image",
|
||||||
|
"command": self.command if self.command else None",
|
||||||
|
"environment": self.environment",
|
||||||
|
"user": self.user",
|
||||||
|
# Security requirements
|
||||||
|
"network_mode": self.network_mode",
|
||||||
|
"read_only": True",
|
||||||
|
"security_opt": ["no-new-privileges"]",
|
||||||
|
}
|
||||||
|
|
||||||
|
# Add working_dir only if explicitly set
|
||||||
|
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
|
||||||
|
|
||||||
|
# Add resource limits if present
|
||||||
|
if self.resource_limits:
|
||||||
|
limit_params = self.resource_limits.to_podman_params()
|
||||||
|
params.update(limit_params)
|
||||||
|
# Disable swap to avoid cgroup swap.max issues on some systems
|
||||||
|
if "mem_limit" in params:
|
||||||
|
params["memswap_limit"] = -1 # Disable swap
|
||||||
|
|
||||||
|
return params
|
||||||
|
|
||||||
|
|
||||||
|
class SecureContainerManager:
|
||||||
|
""Manages container lifecycle with security enforcement.""
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self",
|
||||||
|
podman_client: PodmanClient",
|
||||||
|
validator: OperationValidatorProtocol",
|
||||||
|
audit_logger: Optional[AuditLoggerProtocol] = None
|
||||||
|
):
|
||||||
|
""
|
||||||
|
Initialize secure container manager.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
podman_client: Podman client wrapper
|
||||||
|
validator: Operation validator for security checks
|
||||||
|
audit_logger: Optional audit logger (uses NullAuditLogger if None)
|
||||||
|
""
|
||||||
|
self.podman = podman_client
|
||||||
|
self.validator = validator
|
||||||
|
self.audit_logger = audit_logger or NullAuditLogger()
|
||||||
|
|
||||||
|
def create_container(
|
||||||
|
self",
|
||||||
|
config: ContainerConfig",
|
||||||
|
session_id: Optional[str] = None",
|
||||||
|
name: Optional[str] = None",
|
||||||
|
**extra_params
|
||||||
|
) -> str:
|
||||||
|
""
|
||||||
|
Create a container with security validation.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
config: Container configuration
|
||||||
|
session_id: Session ID for tracking
|
||||||
|
name: Optional container name
|
||||||
|
**extra_params: Additional parameters (checked for forbidden values)
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Container ID
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
SecurityError: If configuration violates security policy
|
||||||
|
""
|
||||||
|
# Convert config to Podman parameters
|
||||||
|
params = config.to_podman_params()
|
||||||
|
|
||||||
|
# Add session label if provided
|
||||||
|
labels = {}
|
||||||
|
if session_id:
|
||||||
|
labels["mcp-forge.session"] = session_id
|
||||||
|
if labels:
|
||||||
|
params["labels"] = labels
|
||||||
|
|
||||||
|
if name:
|
||||||
|
params["name"] = name
|
||||||
|
|
||||||
|
# Merge any extra parameters (will be validated)
|
||||||
|
params.update(extra_params)
|
||||||
|
|
||||||
|
# Validate against security policy
|
||||||
|
try:
|
||||||
|
# Extract image from params for validation
|
||||||
|
self.validator.validate_container_create(
|
||||||
|
image=config.image",
|
||||||
|
params=params",
|
||||||
|
session_id=session_id
|
||||||
|
)
|
||||||
|
except SecurityError as e:
|
||||||
|
# Log security violation
|
||||||
|
self.audit_logger.log(event_type="security.violation", severity="critical", message=f"Security violation: container_create", details={"operation": "container_create", "reason": str(str(e)")}, session_id=session_id
|
||||||
|
)
|
||||||
|
raise
|
||||||
|
|
||||||
|
# Create container
|
||||||
|
try:
|
||||||
|
container = self.podman.client.containers.create(**params)
|
||||||
|
container_id = container.id
|
||||||
|
|
||||||
|
# Register with validator
|
||||||
|
if session_id:
|
||||||
|
self.validator.register_session_container(container_id)
|
||||||
|
|
||||||
|
# Log successful creation
|
||||||
|
self.audit_logger.log(event_type="container.operation", severity="info", operation="create",
|
||||||
|
container_id=container_id",
|
||||||
|
image=config.image",
|
||||||
|
session_id=session_id",
|
||||||
|
details={
|
||||||
|
"name": name",
|
||||||
|
"command": config.command
|
||||||
|
}
|
||||||
|
)
|
||||||
|
|
||||||
|
return container_id
|
||||||
|
|
||||||
|
except Exception as e:
|
||||||
|
self.audit_logger.log(
|
||||||
|
event_type="CONTAINER_CREATE",
|
||||||
|
severity="ERROR",
|
||||||
|
message=f"Container creation failed: {e}",
|
||||||
|
details={
|
||||||
|
"image": config.image",
|
||||||
|
"session_id": session_id",
|
||||||
|
"error": str(e)
|
||||||
|
}
|
||||||
|
)
|
||||||
|
raise
|
||||||
|
|
||||||
|
def start_container(self, container_id: str) -> None:
|
||||||
|
""
|
||||||
|
Start a container.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
container_id: Container ID to start
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
SecurityError: If container is not a session container
|
||||||
|
""
|
||||||
|
# Verify container is registered (security check)
|
||||||
|
if container_id not in self.validator.session_containers:
|
||||||
|
self.audit_logger.log_security_violation(
|
||||||
|
operation="container_start",
|
||||||
|
reason=f"Attempted to start unregistered container: {container_id}"
|
||||||
|
)
|
||||||
|
raise SecurityError(
|
||||||
|
f"Container {container_id} is not a registered session container"
|
||||||
|
)
|
||||||
|
|
||||||
|
try:
|
||||||
|
container = self.podman.client.containers.get(container_id)
|
||||||
|
container.start()
|
||||||
|
|
||||||
|
self.audit_logger.log(event_type="container.operation", severity="info", operation="start",
|
||||||
|
container_id=container_id",
|
||||||
|
image=" # Not available without extra lookup
|
||||||
|
)
|
||||||
|
|
||||||
|
except Exception as e:
|
||||||
|
self.audit_logger.log(
|
||||||
|
event_type="CONTAINER_START",
|
||||||
|
severity="ERROR",
|
||||||
|
message=f"Container start failed: {e}",
|
||||||
|
details={"container_id": container_id, "error": str(e)}
|
||||||
|
)
|
||||||
|
raise
|
||||||
|
|
||||||
|
def stop_container(
|
||||||
|
self",
|
||||||
|
container_id: str",
|
||||||
|
timeout: int = 10
|
||||||
|
) -> None:
|
||||||
|
""
|
||||||
|
Stop a container.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
container_id: Container ID to stop
|
||||||
|
timeout: Timeout in seconds
|
||||||
|
""
|
||||||
|
# Verify container is registered
|
||||||
|
if container_id not in self.validator.session_containers:
|
||||||
|
self.audit_logger.log_security_violation(
|
||||||
|
operation="container_stop",
|
||||||
|
reason=f"Attempted to stop unregistered container: {container_id}"
|
||||||
|
)
|
||||||
|
raise SecurityError(
|
||||||
|
f"Container {container_id} is not a registered session container"
|
||||||
|
)
|
||||||
|
|
||||||
|
try:
|
||||||
|
container = self.podman.client.containers.get(container_id)
|
||||||
|
container.stop(timeout=timeout)
|
||||||
|
|
||||||
|
self.audit_logger.log(event_type="container.operation", severity="info", operation="stop",
|
||||||
|
container_id=container_id",
|
||||||
|
image="",
|
||||||
|
details={"timeout": timeout}
|
||||||
|
)
|
||||||
|
|
||||||
|
except Exception as e:
|
||||||
|
self.audit_logger.log(
|
||||||
|
event_type="CONTAINER_STOP",
|
||||||
|
severity="ERROR",
|
||||||
|
message=f"Container stop failed: {e}",
|
||||||
|
details={"container_id": container_id, "error": str(e)}
|
||||||
|
)
|
||||||
|
raise
|
||||||
|
|
||||||
|
def remove_container(
|
||||||
|
self",
|
||||||
|
container_id: str",
|
||||||
|
force: bool = False
|
||||||
|
) -> None:
|
||||||
|
""
|
||||||
|
Remove a container.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
container_id: Container ID to remove
|
||||||
|
force: Force removal even if running
|
||||||
|
""
|
||||||
|
# Verify container is registered
|
||||||
|
if container_id not in self.validator.session_containers:
|
||||||
|
self.audit_logger.log_security_violation(
|
||||||
|
operation="container_remove",
|
||||||
|
reason=f"Attempted to remove unregistered container: {container_id}"
|
||||||
|
)
|
||||||
|
raise SecurityError(
|
||||||
|
f"Container {container_id} is not a registered session container"
|
||||||
|
)
|
||||||
|
|
||||||
|
try:
|
||||||
|
container = self.podman.client.containers.get(container_id)
|
||||||
|
container.remove(force=force)
|
||||||
|
|
||||||
|
# Unregister from validator
|
||||||
|
self.validator.unregister_session_container(container_id)
|
||||||
|
|
||||||
|
self.audit_logger.log(event_type="container.operation", severity="info", operation="remove",
|
||||||
|
container_id=container_id",
|
||||||
|
image="",
|
||||||
|
details={"force": force}
|
||||||
|
)
|
||||||
|
|
||||||
|
except Exception as e:
|
||||||
|
self.audit_logger.log(
|
||||||
|
event_type="CONTAINER_REMOVE",
|
||||||
|
severity="ERROR",
|
||||||
|
message=f"Container removal failed: {e}",
|
||||||
|
details={"container_id": container_id, "error": str(e)}
|
||||||
|
)
|
||||||
|
raise
|
||||||
|
|
||||||
|
def get_container_logs(
|
||||||
|
self",
|
||||||
|
container_id: str",
|
||||||
|
tail: int = 100
|
||||||
|
) -> tuple[str, str]:
|
||||||
|
""
|
||||||
|
Get container stdout and stderr logs.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
container_id: Container ID
|
||||||
|
tail: Number of lines to retrieve
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
(stdout, stderr) as strings
|
||||||
|
""
|
||||||
|
if container_id not in self.validator.session_containers:
|
||||||
|
raise SecurityError(
|
||||||
|
f"Container {container_id} is not a registered session container"
|
||||||
|
)
|
||||||
|
|
||||||
|
try:
|
||||||
|
container = self.podman.client.containers.get(container_id)
|
||||||
|
logs = container.logs(tail=tail, stdout=True, stderr=True)
|
||||||
|
|
||||||
|
# Podman logs returns a generator of frames, need to consume it
|
||||||
|
if hasattr(logs, '__iter__') and not isinstance(logs, (str, bytes)):
|
||||||
|
# It's a generator/iterator, consume it
|
||||||
|
logs_bytes = b''.join(logs)
|
||||||
|
logs_str = logs_bytes.decode('utf-8', errors='replace')
|
||||||
|
elif isinstance(logs, bytes):
|
||||||
|
logs_str = logs.decode('utf-8', errors='replace')
|
||||||
|
else:
|
||||||
|
logs_str = str(logs)
|
||||||
|
|
||||||
|
# For simplicity, return all logs in stdout (Podman combines them)
|
||||||
|
return logs_str, "
|
||||||
|
|
||||||
|
except Exception as e:
|
||||||
|
self.audit_logger.log(
|
||||||
|
event_type="EXECUTION_REQUEST",
|
||||||
|
severity="ERROR",
|
||||||
|
message=f"Failed to get container logs: {e}",
|
||||||
|
details={"container_id": container_id, "error": str(e)}
|
||||||
|
)
|
||||||
|
raise
|
||||||
|
|
||||||
|
def wait_for_container(
|
||||||
|
self",
|
||||||
|
container_id: str",
|
||||||
|
timeout: int = 300
|
||||||
|
) -> int:
|
||||||
|
""
|
||||||
|
Wait for container to exit.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
container_id: Container ID
|
||||||
|
timeout: Timeout in seconds
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Exit code
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
TimeoutError: If container doesn't exit within timeout
|
||||||
|
""
|
||||||
|
if container_id not in self.validator.session_containers:
|
||||||
|
raise SecurityError(
|
||||||
|
f"Container {container_id} is not a registered session container"
|
||||||
|
)
|
||||||
|
|
||||||
|
try:
|
||||||
|
container = self.podman.client.containers.get(container_id)
|
||||||
|
result = container.wait(timeout=timeout)
|
||||||
|
|
||||||
|
# Extract exit code from result
|
||||||
|
if isinstance(result, dict):
|
||||||
|
exit_code = result.get("StatusCode", 0)
|
||||||
|
else:
|
||||||
|
exit_code = result
|
||||||
|
|
||||||
|
return exit_code
|
||||||
|
|
||||||
|
except Exception as e:
|
||||||
|
self.audit_logger.log(
|
||||||
|
event_type="EXECUTION_REQUEST",
|
||||||
|
severity="ERROR",
|
||||||
|
message=f"Failed to wait for container: {e}",
|
||||||
|
details={"container_id": container_id, "error": str(e)}
|
||||||
|
)
|
||||||
|
raise
|
||||||
|
|
||||||
|
def cleanup_old_containers(
|
||||||
|
self",
|
||||||
|
max_age: timedelta = timedelta(hours=24)
|
||||||
|
) -> int:
|
||||||
|
""
|
||||||
|
Cleanup containers older than max_age.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
max_age: Maximum age for containers
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Number of containers removed
|
||||||
|
""
|
||||||
|
try:
|
||||||
|
# Get all containers with mcp-forge.session label
|
||||||
|
containers = self.podman.client.containers.list(
|
||||||
|
all=True",
|
||||||
|
filters={"label": ["mcp-forge.session"]}
|
||||||
|
)
|
||||||
|
|
||||||
|
removed_count = 0
|
||||||
|
now = datetime.now()
|
||||||
|
|
||||||
|
for container in containers:
|
||||||
|
# Get creation time
|
||||||
|
created_str = container.attrs.get("Created", ")
|
||||||
|
if not created_str:
|
||||||
|
continue
|
||||||
|
|
||||||
|
# Parse ISO format timestamp
|
||||||
|
try:
|
||||||
|
# Remove fractional seconds and timezone for parsing
|
||||||
|
created_str = created_str.split('.')[0]
|
||||||
|
created = datetime.fromisoformat(created_str.replace('Z', ''))
|
||||||
|
except (ValueError, AttributeError):
|
||||||
|
continue
|
||||||
|
|
||||||
|
age = now - created
|
||||||
|
|
||||||
|
if age > max_age:
|
||||||
|
try:
|
||||||
|
container.remove(force=True)
|
||||||
|
removed_count += 1
|
||||||
|
|
||||||
|
self.audit_logger.log(
|
||||||
|
event_type="CONTAINER_REMOVE",
|
||||||
|
severity="INFO",
|
||||||
|
message=f"Cleaned up old container: {container.id}",
|
||||||
|
details={
|
||||||
|
"container_id": container.id",
|
||||||
|
"age_hours": age.total_seconds() / 3600
|
||||||
|
}
|
||||||
|
)
|
||||||
|
except Exception as e:
|
||||||
|
self.audit_logger.log(
|
||||||
|
event_type="CONTAINER_REMOVE",
|
||||||
|
severity="WARNING",
|
||||||
|
message=f"Failed to remove old container: {e}",
|
||||||
|
details={"container_id": container.id, "error": str(e)}
|
||||||
|
)
|
||||||
|
|
||||||
|
return removed_count
|
||||||
|
|
||||||
|
except Exception as e:
|
||||||
|
self.audit_logger.log(
|
||||||
|
event_type="CONTAINER_REMOVE",
|
||||||
|
severity="ERROR",
|
||||||
|
message=f"Cleanup failed: {e}",
|
||||||
|
details={"error": str(e)}
|
||||||
|
)
|
||||||
|
raise
|
||||||
14
src/pod_executor/jupyter/__init__.py
Normal file
14
src/pod_executor/jupyter/__init__.py
Normal file
|
|
@ -0,0 +1,14 @@
|
||||||
|
"""Jupyter stateful code executor."""
|
||||||
|
|
||||||
|
from pod_executor.jupyter.backend import JupyterBackend
|
||||||
|
from pod_executor.jupyter.kernel import JupyterKernelManager, KernelError
|
||||||
|
from pod_executor.jupyter.sessions import SessionManager, SessionState, SessionError
|
||||||
|
|
||||||
|
__all__ = [
|
||||||
|
"JupyterBackend",
|
||||||
|
"JupyterKernelManager",
|
||||||
|
"KernelError",
|
||||||
|
"SessionManager",
|
||||||
|
"SessionState",
|
||||||
|
"SessionError",
|
||||||
|
]
|
||||||
229
src/pod_executor/jupyter/backend.py
Normal file
229
src/pod_executor/jupyter/backend.py
Normal file
|
|
@ -0,0 +1,229 @@
|
||||||
|
"""Jupyter backend for stateful code execution."""
|
||||||
|
|
||||||
|
from typing import Optional, Dict, List
|
||||||
|
import hashlib
|
||||||
|
|
||||||
|
from pod_executor.containers.manager import SecureContainerManager
|
||||||
|
from pod_executor.security.resource_limits import ResourceLimits, parse_memory_string
|
||||||
|
from pod_executor.security.audit import AuditLoggerProtocol, NullAuditLogger
|
||||||
|
from pod_executor.simple.executor import ExecutionResult
|
||||||
|
from pod_executor.jupyter.kernel import JupyterKernelManager
|
||||||
|
from pod_executor.jupyter.sessions import SessionManager, SessionState, SessionError
|
||||||
|
|
||||||
|
|
||||||
|
class JupyterBackend:
|
||||||
|
"""Stateful code execution backend using Jupyter kernels."""
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
container_manager: SecureContainerManager,
|
||||||
|
image: str,
|
||||||
|
default_timeout: int = 300,
|
||||||
|
default_memory: str = "512m",
|
||||||
|
default_cpu_quota: int = 100000,
|
||||||
|
max_timeout: int = 3600,
|
||||||
|
max_memory: str = "2g",
|
||||||
|
max_cpu_quota: int = 200000,
|
||||||
|
idle_timeout: int = 3600,
|
||||||
|
max_sessions: int = 10,
|
||||||
|
resource_limits: Optional[ResourceLimits] = None,
|
||||||
|
audit_logger: Optional[AuditLoggerProtocol] = None
|
||||||
|
):
|
||||||
|
"""
|
||||||
|
Initialize Jupyter backend.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
container_manager: Container lifecycle manager
|
||||||
|
image: Docker/Podman image with ipykernel installed
|
||||||
|
default_timeout: Default execution timeout in seconds
|
||||||
|
default_memory: Default memory limit (e.g., "512m", "1g")
|
||||||
|
default_cpu_quota: Default CPU quota (100000 = 1 CPU)
|
||||||
|
max_timeout: Maximum allowed timeout
|
||||||
|
max_memory: Maximum allowed memory
|
||||||
|
max_cpu_quota: Maximum allowed CPU quota
|
||||||
|
idle_timeout: Session idle timeout in seconds
|
||||||
|
max_sessions: Maximum concurrent sessions
|
||||||
|
resource_limits: Default resource limits for kernels
|
||||||
|
audit_logger: Optional audit logger (uses NullAuditLogger if None)
|
||||||
|
"""
|
||||||
|
self.container_manager = container_manager
|
||||||
|
self.image = image
|
||||||
|
self.default_timeout = default_timeout
|
||||||
|
self.default_memory = default_memory
|
||||||
|
self.default_cpu_quota = default_cpu_quota
|
||||||
|
self.max_timeout = max_timeout
|
||||||
|
self.max_memory = max_memory
|
||||||
|
self.max_cpu_quota = max_cpu_quota
|
||||||
|
self.idle_timeout = idle_timeout
|
||||||
|
self.max_sessions = max_sessions
|
||||||
|
self.audit_logger = audit_logger or NullAuditLogger()
|
||||||
|
|
||||||
|
# Initialize kernel manager
|
||||||
|
kernel_manager = JupyterKernelManager(
|
||||||
|
container_manager=container_manager,
|
||||||
|
image=image,
|
||||||
|
resource_limits=resource_limits
|
||||||
|
)
|
||||||
|
|
||||||
|
# Initialize session manager
|
||||||
|
self.session_manager = SessionManager(
|
||||||
|
kernel_manager=kernel_manager,
|
||||||
|
idle_timeout=idle_timeout,
|
||||||
|
max_sessions=max_sessions,
|
||||||
|
audit_logger=self.audit_logger
|
||||||
|
)
|
||||||
|
|
||||||
|
def execute(
|
||||||
|
self,
|
||||||
|
code: str,
|
||||||
|
session_id: str,
|
||||||
|
timeout: Optional[int] = None,
|
||||||
|
memory: Optional[str] = None,
|
||||||
|
cpu_quota: Optional[int] = None,
|
||||||
|
volumes: Optional[Dict[str, dict]] = None,
|
||||||
|
injection_code: Optional[str] = None,
|
||||||
|
bridge_socket_path: Optional[str] = None
|
||||||
|
) -> ExecutionResult:
|
||||||
|
"""
|
||||||
|
Execute code in stateful session.
|
||||||
|
|
||||||
|
Creates session if it doesn't exist, reuses existing session otherwise.
|
||||||
|
Session maintains namespace state across multiple executions.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
code: Python code to execute
|
||||||
|
session_id: Unique session identifier
|
||||||
|
timeout: Max execution time in seconds (uses default if None)
|
||||||
|
memory: Memory limit string (uses default if None)
|
||||||
|
cpu_quota: CPU quota (uses default if None)
|
||||||
|
volumes: Volume mounts dict
|
||||||
|
injection_code: Optional code executed once at session start
|
||||||
|
bridge_socket_path: Optional path to socket for mounting
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
ExecutionResult with execution output and metadata
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
ValueError: If limits exceed configured maximums
|
||||||
|
SessionError: If session operation fails
|
||||||
|
"""
|
||||||
|
# Use defaults if not specified
|
||||||
|
timeout = timeout if timeout is not None else self.default_timeout
|
||||||
|
memory = memory if memory is not None else self.default_memory
|
||||||
|
cpu_quota = cpu_quota if cpu_quota is not None else self.default_cpu_quota
|
||||||
|
|
||||||
|
# Validate limits against maximums
|
||||||
|
self._validate_limits(timeout, memory, cpu_quota)
|
||||||
|
|
||||||
|
# Log execution (hash code, don't log actual content)
|
||||||
|
code_hash = hashlib.sha256(code.encode()).hexdigest()
|
||||||
|
self.audit_logger.log(
|
||||||
|
event_type="execution.request",
|
||||||
|
severity="info",
|
||||||
|
message="Stateful code execution requested",
|
||||||
|
session_id=session_id,
|
||||||
|
details={
|
||||||
|
"code_hash": code_hash,
|
||||||
|
"timeout": timeout,
|
||||||
|
"memory": memory,
|
||||||
|
"cpu_quota": cpu_quota
|
||||||
|
}
|
||||||
|
)
|
||||||
|
|
||||||
|
# Check if session exists, create if needed
|
||||||
|
try:
|
||||||
|
self.session_manager.get_session(session_id)
|
||||||
|
except SessionError:
|
||||||
|
# Session doesn't exist, create it
|
||||||
|
resource_limits = ResourceLimits(
|
||||||
|
memory=memory,
|
||||||
|
cpu_quota=cpu_quota,
|
||||||
|
storage="1g",
|
||||||
|
timeout=timeout
|
||||||
|
)
|
||||||
|
|
||||||
|
self.session_manager.create_session(
|
||||||
|
session_id=session_id,
|
||||||
|
resource_limits=resource_limits,
|
||||||
|
volumes=volumes,
|
||||||
|
injection_code=injection_code,
|
||||||
|
bridge_socket_path=bridge_socket_path
|
||||||
|
)
|
||||||
|
|
||||||
|
# Execute in session
|
||||||
|
result = self.session_manager.execute_in_session(
|
||||||
|
session_id=session_id,
|
||||||
|
code=code,
|
||||||
|
timeout=timeout
|
||||||
|
)
|
||||||
|
|
||||||
|
return result
|
||||||
|
|
||||||
|
def document_state(
|
||||||
|
self,
|
||||||
|
session_id: str,
|
||||||
|
variables: Dict[str, str],
|
||||||
|
note: str = "",
|
||||||
|
clear: bool = False
|
||||||
|
) -> dict:
|
||||||
|
"""
|
||||||
|
Document important variables in session.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
session_id: Session to document
|
||||||
|
variables: Dictionary of variable_name -> description
|
||||||
|
note: Optional note about session state
|
||||||
|
clear: If True, replace all documented variables; if False, merge
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Dictionary with updated state info
|
||||||
|
"""
|
||||||
|
self.session_manager.document_state(
|
||||||
|
session_id=session_id,
|
||||||
|
variables=variables,
|
||||||
|
note=note,
|
||||||
|
clear=clear
|
||||||
|
)
|
||||||
|
|
||||||
|
state = self.session_manager.get_session_state(session_id)
|
||||||
|
return state.to_dict()
|
||||||
|
|
||||||
|
def get_session_state(self, session_id: str) -> SessionState:
|
||||||
|
"""Get documented state for session."""
|
||||||
|
return self.session_manager.get_session_state(session_id)
|
||||||
|
|
||||||
|
def destroy_session(self, session_id: str) -> None:
|
||||||
|
"""Destroy session and cleanup kernel."""
|
||||||
|
self.session_manager.destroy_session(session_id)
|
||||||
|
|
||||||
|
def list_sessions(self) -> List[dict]:
|
||||||
|
"""List all active sessions with metadata."""
|
||||||
|
return self.session_manager.list_sessions()
|
||||||
|
|
||||||
|
def cleanup_idle_sessions(self) -> int:
|
||||||
|
"""Cleanup sessions idle beyond configured timeout."""
|
||||||
|
return self.session_manager.cleanup_idle_sessions()
|
||||||
|
|
||||||
|
def _validate_limits(self, timeout: int, memory: str, cpu_quota: int) -> None:
|
||||||
|
"""
|
||||||
|
Validate resource limits against configured maximums.
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
ValueError: If any limit exceeds maximum
|
||||||
|
"""
|
||||||
|
if timeout > self.max_timeout:
|
||||||
|
raise ValueError(
|
||||||
|
f"Timeout {timeout} exceeds maximum {self.max_timeout}"
|
||||||
|
)
|
||||||
|
|
||||||
|
memory_bytes = parse_memory_string(memory)
|
||||||
|
max_memory_bytes = parse_memory_string(self.max_memory)
|
||||||
|
if memory_bytes > max_memory_bytes:
|
||||||
|
raise ValueError(
|
||||||
|
f"Memory {memory} exceeds maximum {self.max_memory}"
|
||||||
|
)
|
||||||
|
|
||||||
|
if cpu_quota > self.max_cpu_quota:
|
||||||
|
raise ValueError(
|
||||||
|
f"CPU quota {cpu_quota} exceeds maximum {self.max_cpu_quota}"
|
||||||
|
)
|
||||||
631
src/pod_executor/jupyter/kernel.py
Normal file
631
src/pod_executor/jupyter/kernel.py
Normal file
|
|
@ -0,0 +1,631 @@
|
||||||
|
"""
|
||||||
|
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
|
||||||
|
from datetime import datetime, timedelta
|
||||||
|
import uuid
|
||||||
|
import json
|
||||||
|
import tempfile
|
||||||
|
import time
|
||||||
|
import socket
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
from jupyter_client.blocking.client import BlockingKernelClient
|
||||||
|
import zmq
|
||||||
|
|
||||||
|
from pod_executor.containers.manager import SecureContainerManager, ContainerConfig
|
||||||
|
from pod_executor.security.resource_limits import ResourceLimits
|
||||||
|
from pod_executor.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_file: Path
|
||||||
|
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 via jupyter-client.
|
||||||
|
|
||||||
|
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] = None
|
||||||
|
):
|
||||||
|
"""
|
||||||
|
Initialize kernel manager.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
container_manager: Container lifecycle manager
|
||||||
|
image: Docker/Podman image with ipykernel installed
|
||||||
|
resource_limits: Default resource limits for kernels
|
||||||
|
"""
|
||||||
|
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,
|
||||||
|
injection_code: Optional[str] = None,
|
||||||
|
bridge_socket_path: Optional[str] = None
|
||||||
|
) -> str:
|
||||||
|
"""
|
||||||
|
Start IPython kernel in dedicated container.
|
||||||
|
|
||||||
|
Process:
|
||||||
|
1. Generate ZMQ connection info (ports, keys)
|
||||||
|
2. Create connection file
|
||||||
|
3. Create container with ipykernel command
|
||||||
|
4. Mount bridge socket if provided (for MCP tools)
|
||||||
|
5. Start container
|
||||||
|
6. Wait for kernel to be ready
|
||||||
|
7. Connect jupyter-client to kernel via ZMQ
|
||||||
|
8. Execute injection code (MCP tools setup) if provided
|
||||||
|
9. Verify kernel is responsive
|
||||||
|
|
||||||
|
Args:
|
||||||
|
session_id: Session ID this kernel belongs to
|
||||||
|
volumes: Optional volume mounts
|
||||||
|
injection_code: Optional MCP tool injection code to execute at startup
|
||||||
|
bridge_socket_path: Optional path to MCP bridge socket for mounting
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
kernel_id: Unique identifier for this kernel
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
KernelError: If kernel startup fails
|
||||||
|
"""
|
||||||
|
kernel_id = f"kernel-{uuid.uuid4().hex[:16]}"
|
||||||
|
|
||||||
|
# Generate connection info
|
||||||
|
connection_info = self._generate_connection_info()
|
||||||
|
|
||||||
|
# Create connection file
|
||||||
|
connection_file = self._create_connection_file(kernel_id, connection_info)
|
||||||
|
|
||||||
|
try:
|
||||||
|
# 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] = {
|
||||||
|
"bind": bridge_socket_path,
|
||||||
|
"mode": "rw"
|
||||||
|
}
|
||||||
|
|
||||||
|
# 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", container_connection_path
|
||||||
|
],
|
||||||
|
resource_limits=self.resource_limits,
|
||||||
|
volumes=container_volumes,
|
||||||
|
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(
|
||||||
|
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 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)
|
||||||
|
|
||||||
|
# Verify kernel is responsive
|
||||||
|
if not self._verify_kernel(client):
|
||||||
|
raise KernelError(f"Kernel {kernel_id} not responsive")
|
||||||
|
|
||||||
|
# Execute injection code if provided (MCP tools setup)
|
||||||
|
if injection_code:
|
||||||
|
self._execute_injection_code(client, injection_code, kernel_id)
|
||||||
|
|
||||||
|
# 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,
|
||||||
|
kernel_id: str,
|
||||||
|
code: str,
|
||||||
|
timeout: int = 300
|
||||||
|
) -> ExecutionResult:
|
||||||
|
"""
|
||||||
|
Execute code in kernel via ZMQ.
|
||||||
|
|
||||||
|
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: Kernel to execute in
|
||||||
|
code: Python code to execute
|
||||||
|
timeout: Maximum execution time in seconds
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
ExecutionResult with stdout, stderr, result
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
KernelError: If kernel not found or execution fails
|
||||||
|
"""
|
||||||
|
kernel_info = self._get_kernel(kernel_id)
|
||||||
|
client = kernel_info.client
|
||||||
|
|
||||||
|
if not client:
|
||||||
|
raise KernelError(f"Kernel {kernel_id} has no connected client")
|
||||||
|
|
||||||
|
start_time = time.time()
|
||||||
|
|
||||||
|
try:
|
||||||
|
# Execute code
|
||||||
|
_msg_id = client.execute(code, silent=False, store_history=True)
|
||||||
|
|
||||||
|
# Collect output
|
||||||
|
stdout_parts = []
|
||||||
|
stderr_parts = []
|
||||||
|
result = None
|
||||||
|
has_error = False
|
||||||
|
|
||||||
|
# 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':
|
||||||
|
has_error = True
|
||||||
|
stderr_parts.append('\n'.join(content['traceback']))
|
||||||
|
|
||||||
|
elif msg_type == 'status':
|
||||||
|
if content['execution_state'] == 'idle':
|
||||||
|
break
|
||||||
|
|
||||||
|
except zmq.error.Again:
|
||||||
|
break
|
||||||
|
|
||||||
|
execution_time = time.time() - start_time
|
||||||
|
|
||||||
|
# Update last activity
|
||||||
|
kernel_info.last_activity = datetime.utcnow()
|
||||||
|
|
||||||
|
stderr_text = ''.join(stderr_parts)
|
||||||
|
|
||||||
|
return ExecutionResult(
|
||||||
|
success=(not has_error),
|
||||||
|
stdout=''.join(stdout_parts),
|
||||||
|
stderr=stderr_text,
|
||||||
|
result=result,
|
||||||
|
execution_time=execution_time,
|
||||||
|
exit_code=1 if has_error else 0,
|
||||||
|
error=stderr_text if has_error else None
|
||||||
|
)
|
||||||
|
|
||||||
|
except Exception as e:
|
||||||
|
execution_time = time.time() - start_time
|
||||||
|
return ExecutionResult(
|
||||||
|
success=False,
|
||||||
|
stdout='',
|
||||||
|
stderr=str(e),
|
||||||
|
result=None,
|
||||||
|
execution_time=execution_time,
|
||||||
|
exit_code=1,
|
||||||
|
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: Kernel to shutdown
|
||||||
|
"""
|
||||||
|
kernel_info = self._get_kernel(kernel_id)
|
||||||
|
|
||||||
|
try:
|
||||||
|
# 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)
|
||||||
|
|
||||||
|
# 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: Kernel to inspect
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
List of variable names
|
||||||
|
"""
|
||||||
|
code = "[var for var in dir() if not var.startswith('_')]"
|
||||||
|
result = self.execute_code(kernel_id, code, timeout=5)
|
||||||
|
|
||||||
|
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,
|
||||||
|
kernel_id: str,
|
||||||
|
variable_name: str
|
||||||
|
) -> Dict[str, Any]:
|
||||||
|
"""
|
||||||
|
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: Kernel to inspect
|
||||||
|
variable_name: Name of variable to inspect
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Dict with type, size, shape, repr
|
||||||
|
"""
|
||||||
|
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)
|
||||||
|
|
||||||
|
if result.success and result.result:
|
||||||
|
try:
|
||||||
|
return eval(result.result) # nosec - controlled code
|
||||||
|
except Exception:
|
||||||
|
return {}
|
||||||
|
return {}
|
||||||
|
|
||||||
|
def restart_kernel(self, kernel_id: str) -> None:
|
||||||
|
"""
|
||||||
|
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: Kernel to restart
|
||||||
|
"""
|
||||||
|
kernel_info = self._get_kernel(kernel_id)
|
||||||
|
|
||||||
|
# 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:
|
||||||
|
"""
|
||||||
|
Cleanup kernels idle longer than timeout.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
idle_timeout: Maximum idle time before cleanup
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Number of kernels cleaned up
|
||||||
|
"""
|
||||||
|
now = datetime.utcnow()
|
||||||
|
cleaned_up = 0
|
||||||
|
|
||||||
|
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:
|
||||||
|
try:
|
||||||
|
self.shutdown_kernel(kernel_id)
|
||||||
|
cleaned_up += 1
|
||||||
|
except Exception:
|
||||||
|
pass # Continue cleanup even if one fails
|
||||||
|
|
||||||
|
return cleaned_up
|
||||||
|
|
||||||
|
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 with allocated ports."""
|
||||||
|
import secrets
|
||||||
|
|
||||||
|
# Allocate 5 ports for ZMQ channels
|
||||||
|
ports = self._allocate_ports(5)
|
||||||
|
|
||||||
|
return {
|
||||||
|
"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",
|
||||||
|
"signature_scheme": "hmac-sha256",
|
||||||
|
"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,
|
||||||
|
connection_info: Dict[str, Any]
|
||||||
|
) -> Path:
|
||||||
|
"""Create connection file for kernel."""
|
||||||
|
# Create temp file
|
||||||
|
fd, path = tempfile.mkstemp(suffix=f"-kernel-{kernel_id}.json")
|
||||||
|
|
||||||
|
# Write connection info
|
||||||
|
with open(fd, 'w') as f:
|
||||||
|
json.dump(connection_info, f)
|
||||||
|
|
||||||
|
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
|
||||||
|
|
||||||
|
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,
|
||||||
|
injection_code: str,
|
||||||
|
kernel_id: str
|
||||||
|
) -> None:
|
||||||
|
"""
|
||||||
|
Execute MCP tool injection code on kernel startup.
|
||||||
|
|
||||||
|
This runs once when the kernel starts to set up MCP tools.
|
||||||
|
Unlike regular code execution, we don't capture output.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
client: Connected kernel client
|
||||||
|
injection_code: Python code to inject (MCP tools setup)
|
||||||
|
kernel_id: Kernel ID for error messages
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
KernelError: If injection code fails to execute
|
||||||
|
"""
|
||||||
|
try:
|
||||||
|
# Execute injection code silently
|
||||||
|
_msg_id = client.execute(injection_code, silent=True, store_history=False)
|
||||||
|
|
||||||
|
# Wait for execution to complete
|
||||||
|
timeout = 10 # Injection should be fast
|
||||||
|
while True:
|
||||||
|
try:
|
||||||
|
msg = client.get_iopub_msg(timeout=timeout)
|
||||||
|
msg_type = msg['header']['msg_type']
|
||||||
|
|
||||||
|
if msg_type == 'error':
|
||||||
|
content = msg['content']
|
||||||
|
error_msg = '\n'.join(content.get('traceback', [str(content)]))
|
||||||
|
raise KernelError(
|
||||||
|
f"MCP injection failed in kernel {kernel_id}: {error_msg}"
|
||||||
|
)
|
||||||
|
|
||||||
|
elif msg_type == 'status':
|
||||||
|
if msg['content']['execution_state'] == 'idle':
|
||||||
|
break # Injection complete
|
||||||
|
|
||||||
|
except zmq.error.Again:
|
||||||
|
break # Timeout, assume success
|
||||||
|
|
||||||
|
except KernelError:
|
||||||
|
raise
|
||||||
|
except Exception as e:
|
||||||
|
raise KernelError(
|
||||||
|
f"Failed to execute MCP injection code in kernel {kernel_id}: {e}"
|
||||||
|
) from e
|
||||||
437
src/pod_executor/jupyter/sessions.py
Normal file
437
src/pod_executor/jupyter/sessions.py
Normal file
|
|
@ -0,0 +1,437 @@
|
||||||
|
"""Session management for stateful execution."""
|
||||||
|
|
||||||
|
from typing import Dict, Optional, List, Any
|
||||||
|
from dataclasses import dataclass, field
|
||||||
|
from datetime import datetime, timedelta
|
||||||
|
|
||||||
|
from pod_executor.jupyter.kernel import JupyterKernelManager
|
||||||
|
from pod_executor.security.audit import AuditLoggerProtocol, NullAuditLogger
|
||||||
|
from pod_executor.security.resource_limits import ResourceLimits
|
||||||
|
from pod_executor.simple.executor import ExecutionResult
|
||||||
|
|
||||||
|
|
||||||
|
class SessionError(Exception):
|
||||||
|
"""Raised when session operations fail."""
|
||||||
|
pass
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class SessionState:
|
||||||
|
"""Documented state for a session."""
|
||||||
|
session_id: str
|
||||||
|
documented_variables: Dict[str, str] = field(default_factory=dict)
|
||||||
|
note: str = ""
|
||||||
|
last_updated: datetime = field(default_factory=datetime.utcnow)
|
||||||
|
all_variables: List[str] = field(default_factory=list)
|
||||||
|
introspection: Dict[str, dict] = field(default_factory=dict)
|
||||||
|
|
||||||
|
def to_dict(self) -> dict:
|
||||||
|
"""Convert to dictionary for JSON serialization."""
|
||||||
|
return {
|
||||||
|
"session_id": self.session_id,
|
||||||
|
"documented_variables": self.documented_variables,
|
||||||
|
"note": self.note,
|
||||||
|
"last_updated": self.last_updated.isoformat(),
|
||||||
|
"all_variables": self.all_variables,
|
||||||
|
"introspection": self.introspection
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
class Session:
|
||||||
|
"""Stateful execution session."""
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
session_id: str,
|
||||||
|
kernel_id: str,
|
||||||
|
created_at: datetime,
|
||||||
|
resource_limits: ResourceLimits
|
||||||
|
):
|
||||||
|
"""
|
||||||
|
Initialize session.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
session_id: Unique session identifier
|
||||||
|
kernel_id: ID of associated kernel
|
||||||
|
created_at: Session creation timestamp
|
||||||
|
resource_limits: Resource limits for this session
|
||||||
|
"""
|
||||||
|
self.session_id = session_id
|
||||||
|
self.kernel_id = kernel_id
|
||||||
|
self.created_at = created_at
|
||||||
|
self.last_activity = created_at
|
||||||
|
self.resource_limits = resource_limits
|
||||||
|
self.state = SessionState(session_id=session_id)
|
||||||
|
self.documented_variables: Dict[str, str] = {}
|
||||||
|
self.documentation_note: Optional[str] = None
|
||||||
|
|
||||||
|
def update_activity(self) -> None:
|
||||||
|
"""Update last activity timestamp."""
|
||||||
|
self.last_activity = datetime.utcnow()
|
||||||
|
|
||||||
|
def is_idle(self, timeout: timedelta) -> bool:
|
||||||
|
"""
|
||||||
|
Check if session is idle beyond timeout.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
timeout: Maximum idle time
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
True if session has been idle longer than timeout
|
||||||
|
"""
|
||||||
|
now = datetime.utcnow()
|
||||||
|
idle_time = now - self.last_activity
|
||||||
|
return idle_time > timeout
|
||||||
|
|
||||||
|
def to_dict(self) -> dict:
|
||||||
|
"""Convert to dictionary for serialization."""
|
||||||
|
return {
|
||||||
|
"session_id": self.session_id,
|
||||||
|
"kernel_id": self.kernel_id,
|
||||||
|
"created_at": self.created_at.isoformat(),
|
||||||
|
"last_activity": self.last_activity.isoformat(),
|
||||||
|
"state": self.state.to_dict()
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
class SessionManager:
|
||||||
|
"""Manages stateful execution sessions."""
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
kernel_manager: JupyterKernelManager,
|
||||||
|
idle_timeout: int = 3600,
|
||||||
|
max_sessions: int = 10,
|
||||||
|
audit_logger: Optional[AuditLoggerProtocol] = None
|
||||||
|
):
|
||||||
|
"""
|
||||||
|
Initialize session manager.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
kernel_manager: Kernel lifecycle manager
|
||||||
|
idle_timeout: Session idle timeout in seconds (default: 3600)
|
||||||
|
max_sessions: Maximum concurrent sessions (default: 10)
|
||||||
|
audit_logger: Optional audit logger (uses NullAuditLogger if None)
|
||||||
|
"""
|
||||||
|
self.kernel_manager = kernel_manager
|
||||||
|
self.idle_timeout = idle_timeout
|
||||||
|
self.max_sessions = max_sessions
|
||||||
|
self.audit_logger = audit_logger or NullAuditLogger()
|
||||||
|
self.sessions: Dict[str, Session] = {}
|
||||||
|
|
||||||
|
def create_session(
|
||||||
|
self,
|
||||||
|
session_id: str,
|
||||||
|
resource_limits: ResourceLimits,
|
||||||
|
volumes: Optional[Dict[str, dict]] = None,
|
||||||
|
injection_code: Optional[str] = None,
|
||||||
|
bridge_socket_path: Optional[str] = None
|
||||||
|
) -> Session:
|
||||||
|
"""
|
||||||
|
Create new stateful session.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
session_id: Unique identifier for session
|
||||||
|
resource_limits: Resource limits for session
|
||||||
|
volumes: Optional volume mounts
|
||||||
|
injection_code: Optional MCP tool injection code to execute at startup
|
||||||
|
bridge_socket_path: Optional path to MCP bridge socket for mounting
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Created Session object
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
SessionError: If session_id already exists
|
||||||
|
SessionError: If max concurrent sessions exceeded
|
||||||
|
"""
|
||||||
|
if session_id in self.sessions:
|
||||||
|
raise SessionError(f"Session {session_id} already exists")
|
||||||
|
|
||||||
|
# Check max concurrent limit
|
||||||
|
self._enforce_max_concurrent()
|
||||||
|
|
||||||
|
# Start kernel with MCP injection if provided
|
||||||
|
kernel_id = self.kernel_manager.start_kernel(
|
||||||
|
session_id,
|
||||||
|
volumes=volumes,
|
||||||
|
injection_code=injection_code,
|
||||||
|
bridge_socket_path=bridge_socket_path
|
||||||
|
)
|
||||||
|
|
||||||
|
# Create session
|
||||||
|
now = datetime.utcnow()
|
||||||
|
session = Session(
|
||||||
|
session_id=session_id,
|
||||||
|
kernel_id=kernel_id,
|
||||||
|
created_at=now,
|
||||||
|
resource_limits=resource_limits
|
||||||
|
)
|
||||||
|
|
||||||
|
self.sessions[session_id] = session
|
||||||
|
|
||||||
|
# Log session creation
|
||||||
|
self.audit_logger.log(
|
||||||
|
event_type="SESSION_CREATE,
|
||||||
|
severity="INFO,
|
||||||
|
message=f"Session created: {session_id}",
|
||||||
|
session_id=session_id,
|
||||||
|
details={
|
||||||
|
"kernel_id": kernel_id,
|
||||||
|
"memory": resource_limits.memory_bytes,
|
||||||
|
"cpu_quota": resource_limits.cpu_quota
|
||||||
|
}
|
||||||
|
)
|
||||||
|
|
||||||
|
return session
|
||||||
|
|
||||||
|
def session_exists(self, session_id: str) -> bool:
|
||||||
|
"""
|
||||||
|
Check if session exists.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
session_id: Session identifier
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
True if session exists, False otherwise
|
||||||
|
"""
|
||||||
|
return session_id in self.sessions
|
||||||
|
|
||||||
|
def get_session(self, session_id: str) -> Session:
|
||||||
|
"""
|
||||||
|
Get session by ID.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
session_id: Session identifier
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Session object
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
SessionError: If session doesn't exist
|
||||||
|
"""
|
||||||
|
if session_id not in self.sessions:
|
||||||
|
raise SessionError(f"Session {session_id} not found")
|
||||||
|
|
||||||
|
return self.sessions[session_id]
|
||||||
|
|
||||||
|
async def document_variables(
|
||||||
|
self,
|
||||||
|
session_id: str,
|
||||||
|
variables: Dict[str, str],
|
||||||
|
note: Optional[str] = None,
|
||||||
|
clear: bool = False
|
||||||
|
) -> Dict:
|
||||||
|
"""
|
||||||
|
Document important variables in a session.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
session_id: Session identifier
|
||||||
|
variables: Dict mapping variable names to descriptions
|
||||||
|
note: Optional general note about session state
|
||||||
|
clear: Whether to clear existing documentation first
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Result dict with success status and documented count
|
||||||
|
"""
|
||||||
|
session = self.get_session(session_id)
|
||||||
|
|
||||||
|
if clear:
|
||||||
|
session.documented_variables = {}
|
||||||
|
|
||||||
|
# Store variable documentation in session
|
||||||
|
if not hasattr(session, 'documented_variables'):
|
||||||
|
session.documented_variables = {}
|
||||||
|
|
||||||
|
session.documented_variables.update(variables)
|
||||||
|
|
||||||
|
if note:
|
||||||
|
session.documentation_note = note
|
||||||
|
|
||||||
|
return {
|
||||||
|
"success": True,
|
||||||
|
"documented_count": len(variables),
|
||||||
|
"total_documented": len(session.documented_variables)
|
||||||
|
}
|
||||||
|
|
||||||
|
def execute_in_session(
|
||||||
|
self,
|
||||||
|
session_id: str,
|
||||||
|
code: str,
|
||||||
|
timeout: int = 300
|
||||||
|
) -> ExecutionResult:
|
||||||
|
"""
|
||||||
|
Execute code in session kernel.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
session_id: Session to execute in
|
||||||
|
code: Python code to execute
|
||||||
|
timeout: Maximum execution time
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
ExecutionResult with output
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
SessionError: If session doesn't exist
|
||||||
|
"""
|
||||||
|
session = self.get_session(session_id)
|
||||||
|
|
||||||
|
# Update activity
|
||||||
|
session.update_activity()
|
||||||
|
|
||||||
|
# Execute in kernel
|
||||||
|
result = self.kernel_manager.execute_code(
|
||||||
|
session.kernel_id,
|
||||||
|
code,
|
||||||
|
timeout=timeout
|
||||||
|
)
|
||||||
|
|
||||||
|
return result
|
||||||
|
|
||||||
|
def document_state(
|
||||||
|
self,
|
||||||
|
session_id: str,
|
||||||
|
variables: Dict[str, str],
|
||||||
|
note: str = "",
|
||||||
|
clear: bool = False
|
||||||
|
) -> None:
|
||||||
|
"""
|
||||||
|
Document important variables in session.
|
||||||
|
|
||||||
|
Updates session.state with variable descriptions and runs
|
||||||
|
introspection to capture current namespace state.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
session_id: Session to document
|
||||||
|
variables: Dictionary of variable_name -> description
|
||||||
|
note: Optional note about session state
|
||||||
|
clear: If True, replace all documented variables; if False, merge
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
SessionError: If session doesn't exist
|
||||||
|
"""
|
||||||
|
session = self.get_session(session_id)
|
||||||
|
|
||||||
|
# Update documented variables
|
||||||
|
if clear:
|
||||||
|
session.state.documented_variables = variables.copy()
|
||||||
|
else:
|
||||||
|
session.state.documented_variables.update(variables)
|
||||||
|
|
||||||
|
# Update note if provided
|
||||||
|
if note:
|
||||||
|
session.state.note = note
|
||||||
|
|
||||||
|
# Run introspection to get current namespace state
|
||||||
|
session.state.all_variables = self.kernel_manager.inspect_namespace(session.kernel_id)
|
||||||
|
|
||||||
|
# Get variable info for documented variables
|
||||||
|
session.state.introspection = {}
|
||||||
|
for var_name in variables.keys():
|
||||||
|
if var_name in session.state.all_variables:
|
||||||
|
try:
|
||||||
|
info = self.kernel_manager.get_variable_info(session.kernel_id, var_name)
|
||||||
|
session.state.introspection[var_name] = info
|
||||||
|
except Exception:
|
||||||
|
# Variable might not exist yet
|
||||||
|
pass
|
||||||
|
|
||||||
|
# Update timestamp
|
||||||
|
session.state.last_updated = datetime.utcnow()
|
||||||
|
session.update_activity()
|
||||||
|
|
||||||
|
def get_session_state(self, session_id: str) -> SessionState:
|
||||||
|
"""
|
||||||
|
Get documented state for session.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
session_id: Session identifier
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
SessionState object
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
SessionError: If session doesn't exist
|
||||||
|
"""
|
||||||
|
session = self.get_session(session_id)
|
||||||
|
return session.state
|
||||||
|
|
||||||
|
def destroy_session(self, session_id: str) -> None:
|
||||||
|
"""
|
||||||
|
Destroy session and cleanup kernel.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
session_id: Session to destroy
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
SessionError: If session doesn't exist
|
||||||
|
"""
|
||||||
|
session = self.get_session(session_id)
|
||||||
|
|
||||||
|
# Shutdown kernel
|
||||||
|
try:
|
||||||
|
self.kernel_manager.shutdown_kernel(session.kernel_id)
|
||||||
|
except Exception as e:
|
||||||
|
# Log but continue with cleanup
|
||||||
|
self.audit_logger.log(
|
||||||
|
event_type="SESSION_DESTROY,
|
||||||
|
severity="WARNING,
|
||||||
|
message=f"Error shutting down kernel for session {session_id}",
|
||||||
|
session_id=session_id,
|
||||||
|
error=str(e)
|
||||||
|
)
|
||||||
|
|
||||||
|
# Remove session
|
||||||
|
del self.sessions[session_id]
|
||||||
|
|
||||||
|
# Log destruction
|
||||||
|
self.audit_logger.log(
|
||||||
|
event_type="SESSION_DESTROY,
|
||||||
|
severity="INFO,
|
||||||
|
message=f"Session destroyed: {session_id}",
|
||||||
|
session_id=session_id
|
||||||
|
)
|
||||||
|
|
||||||
|
def cleanup_idle_sessions(self) -> int:
|
||||||
|
"""
|
||||||
|
Cleanup sessions idle beyond configured timeout.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Number of sessions cleaned up
|
||||||
|
"""
|
||||||
|
timeout = timedelta(seconds=self.idle_timeout)
|
||||||
|
sessions_to_remove = []
|
||||||
|
|
||||||
|
for session_id, session in self.sessions.items():
|
||||||
|
if session.is_idle(timeout):
|
||||||
|
sessions_to_remove.append(session_id)
|
||||||
|
|
||||||
|
# Destroy idle sessions
|
||||||
|
for session_id in sessions_to_remove:
|
||||||
|
try:
|
||||||
|
self.destroy_session(session_id)
|
||||||
|
except Exception:
|
||||||
|
# Best effort cleanup
|
||||||
|
pass
|
||||||
|
|
||||||
|
return len(sessions_to_remove)
|
||||||
|
|
||||||
|
def list_sessions(self) -> List[dict]:
|
||||||
|
"""
|
||||||
|
List all active sessions with metadata.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
List of session dictionaries
|
||||||
|
"""
|
||||||
|
return [session.to_dict() for session in self.sessions.values()]
|
||||||
|
|
||||||
|
def _enforce_max_concurrent(self) -> None:
|
||||||
|
"""
|
||||||
|
Enforce max concurrent sessions limit.
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
SessionError: If at max concurrent sessions
|
||||||
|
"""
|
||||||
|
if len(self.sessions) >= self.max_sessions:
|
||||||
|
raise SessionError(
|
||||||
|
f"Maximum concurrent sessions ({self.max_sessions}) reached"
|
||||||
|
)
|
||||||
33
src/pod_executor/security/__init__.py
Normal file
33
src/pod_executor/security/__init__.py
Normal file
|
|
@ -0,0 +1,33 @@
|
||||||
|
"""Security components for pod_executor."""
|
||||||
|
|
||||||
|
from pod_executor.security.resource_limits import (
|
||||||
|
ResourceLimits,
|
||||||
|
parse_memory_string,
|
||||||
|
parse_cpu_quota,
|
||||||
|
parse_storage_string,
|
||||||
|
)
|
||||||
|
from pod_executor.security.audit import (
|
||||||
|
AuditLoggerProtocol,
|
||||||
|
NullAuditLogger,
|
||||||
|
SimpleFileAuditLogger,
|
||||||
|
)
|
||||||
|
from pod_executor.security.validation import (
|
||||||
|
SecurityError,
|
||||||
|
OperationValidatorProtocol,
|
||||||
|
NoOpValidator,
|
||||||
|
BasicValidator,
|
||||||
|
)
|
||||||
|
|
||||||
|
__all__ = [
|
||||||
|
"ResourceLimits",
|
||||||
|
"parse_memory_string",
|
||||||
|
"parse_cpu_quota",
|
||||||
|
"parse_storage_string",
|
||||||
|
"AuditLoggerProtocol",
|
||||||
|
"NullAuditLogger",
|
||||||
|
"SimpleFileAuditLogger",
|
||||||
|
"SecurityError",
|
||||||
|
"OperationValidatorProtocol",
|
||||||
|
"NoOpValidator",
|
||||||
|
"BasicValidator",
|
||||||
|
]
|
||||||
132
src/pod_executor/security/audit.py
Normal file
132
src/pod_executor/security/audit.py
Normal file
|
|
@ -0,0 +1,132 @@
|
||||||
|
"""
|
||||||
|
Audit logger protocol for pod_executor.
|
||||||
|
|
||||||
|
Provides a protocol (interface) for audit logging that can be implemented
|
||||||
|
by consuming applications. A default no-op implementation is provided.
|
||||||
|
"""
|
||||||
|
|
||||||
|
from typing import Protocol, Any, Optional, Dict
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
|
||||||
|
class AuditLoggerProtocol(Protocol):
|
||||||
|
"""Protocol for audit logging (optional dependency)."""
|
||||||
|
|
||||||
|
def log(
|
||||||
|
self,
|
||||||
|
event_type: str,
|
||||||
|
severity: str,
|
||||||
|
message: str,
|
||||||
|
session_id: Optional[str] = None,
|
||||||
|
user_id: Optional[str] = None,
|
||||||
|
details: Optional[Dict[str, Any]] = None,
|
||||||
|
error: Optional[str] = None,
|
||||||
|
**kwargs: Any
|
||||||
|
) -> None:
|
||||||
|
"""
|
||||||
|
Log an audit event.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
event_type: Type of event (e.g., "container.create", "execution.request")
|
||||||
|
severity: Severity level ("info", "warning", "error", "critical")
|
||||||
|
message: Human-readable message describing the event
|
||||||
|
session_id: Optional session ID associated with event
|
||||||
|
user_id: Optional user ID associated with event
|
||||||
|
details: Optional dictionary of additional details
|
||||||
|
error: Optional error message if event represents an error
|
||||||
|
**kwargs: Additional keyword arguments for extensibility
|
||||||
|
"""
|
||||||
|
...
|
||||||
|
|
||||||
|
|
||||||
|
class NullAuditLogger:
|
||||||
|
"""No-op audit logger for standalone usage without audit requirements."""
|
||||||
|
|
||||||
|
def log(
|
||||||
|
self,
|
||||||
|
event_type: str = "",
|
||||||
|
severity: str = "info",
|
||||||
|
message: str = "",
|
||||||
|
session_id: Optional[str] = None,
|
||||||
|
user_id: Optional[str] = None,
|
||||||
|
details: Optional[Dict[str, Any]] = None,
|
||||||
|
error: Optional[str] = None,
|
||||||
|
**kwargs: Any
|
||||||
|
) -> None:
|
||||||
|
"""Do nothing - audit logging disabled."""
|
||||||
|
pass
|
||||||
|
|
||||||
|
|
||||||
|
class SimpleFileAuditLogger:
|
||||||
|
"""
|
||||||
|
Simple file-based audit logger for basic use cases.
|
||||||
|
|
||||||
|
Logs events to a JSON Lines file (one JSON object per line).
|
||||||
|
Thread-safe via file locking.
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(self, log_path: Path):
|
||||||
|
"""
|
||||||
|
Initialize file audit logger.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
log_path: Path to audit log file
|
||||||
|
"""
|
||||||
|
self.log_path = Path(log_path)
|
||||||
|
self._ensure_log_file()
|
||||||
|
|
||||||
|
def _ensure_log_file(self) -> None:
|
||||||
|
"""Ensure log file and directory exist."""
|
||||||
|
self.log_path.parent.mkdir(parents=True, exist_ok=True)
|
||||||
|
if not self.log_path.exists():
|
||||||
|
self.log_path.touch()
|
||||||
|
|
||||||
|
def log(
|
||||||
|
self,
|
||||||
|
event_type: str = "",
|
||||||
|
severity: str = "info",
|
||||||
|
message: str = "",
|
||||||
|
session_id: Optional[str] = None,
|
||||||
|
user_id: Optional[str] = None,
|
||||||
|
details: Optional[Dict[str, Any]] = None,
|
||||||
|
error: Optional[str] = None,
|
||||||
|
**kwargs: Any
|
||||||
|
) -> None:
|
||||||
|
"""
|
||||||
|
Log an audit event to JSON Lines file.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
event_type: Type of event
|
||||||
|
severity: Severity level
|
||||||
|
message: Human-readable message
|
||||||
|
session_id: Optional session ID
|
||||||
|
user_id: Optional user ID
|
||||||
|
details: Optional details dictionary
|
||||||
|
error: Optional error message
|
||||||
|
**kwargs: Additional fields
|
||||||
|
"""
|
||||||
|
import json
|
||||||
|
from datetime import datetime, timezone
|
||||||
|
|
||||||
|
entry = {
|
||||||
|
"timestamp": datetime.now(timezone.utc).isoformat(),
|
||||||
|
"event_type": event_type,
|
||||||
|
"severity": severity,
|
||||||
|
"message": message,
|
||||||
|
}
|
||||||
|
|
||||||
|
if session_id is not None:
|
||||||
|
entry["session_id"] = session_id
|
||||||
|
if user_id is not None:
|
||||||
|
entry["user_id"] = user_id
|
||||||
|
if details is not None:
|
||||||
|
entry["details"] = details
|
||||||
|
if error is not None:
|
||||||
|
entry["error"] = error
|
||||||
|
|
||||||
|
# Add any additional kwargs
|
||||||
|
entry.update(kwargs)
|
||||||
|
|
||||||
|
# Write to file (append mode, file locking via 'a' mode)
|
||||||
|
with open(self.log_path, 'a') as f:
|
||||||
|
f.write(json.dumps(entry) + '\n')
|
||||||
150
src/pod_executor/security/resource_limits.py
Normal file
150
src/pod_executor/security/resource_limits.py
Normal file
|
|
@ -0,0 +1,150 @@
|
||||||
|
"""
|
||||||
|
Resource limit parser and validator.
|
||||||
|
|
||||||
|
Parses and validates resource limit strings (memory, CPU, storage).
|
||||||
|
All values must be positive and within reasonable limits.
|
||||||
|
"""
|
||||||
|
|
||||||
|
import re
|
||||||
|
from typing import Dict, Any
|
||||||
|
|
||||||
|
|
||||||
|
def parse_memory_string(memory: str) -> int:
|
||||||
|
"""
|
||||||
|
Parse memory string to bytes.
|
||||||
|
|
||||||
|
Supports: k, m, g suffixes (case-insensitive)
|
||||||
|
Examples: "512m" → 536870912, "2g" → 2147483648
|
||||||
|
|
||||||
|
Args:
|
||||||
|
memory: Memory string with suffix (e.g., "512m", "2g", "1024k")
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Memory in bytes
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
ValueError: If format is invalid or value is <= 0
|
||||||
|
"""
|
||||||
|
memory = memory.strip()
|
||||||
|
|
||||||
|
# Pattern: optional sign, number (int or float), suffix (k/m/g)
|
||||||
|
pattern = r'^(-?\d+(?:\.\d+)?)\s*([kmgKMG])$'
|
||||||
|
match = re.match(pattern, memory)
|
||||||
|
|
||||||
|
if not match:
|
||||||
|
raise ValueError(
|
||||||
|
f"Invalid memory format: '{memory}'. "
|
||||||
|
f"Expected format: <number><k|m|g> (e.g., '512m', '2g')"
|
||||||
|
)
|
||||||
|
|
||||||
|
value_str, suffix = match.groups()
|
||||||
|
value = float(value_str)
|
||||||
|
|
||||||
|
if value <= 0:
|
||||||
|
raise ValueError(
|
||||||
|
f"Memory value must be positive, got: {value}"
|
||||||
|
)
|
||||||
|
|
||||||
|
# Convert to bytes
|
||||||
|
suffix_lower = suffix.lower()
|
||||||
|
multipliers = {
|
||||||
|
'k': 1024,
|
||||||
|
'm': 1024 * 1024,
|
||||||
|
'g': 1024 * 1024 * 1024,
|
||||||
|
}
|
||||||
|
|
||||||
|
bytes_value = int(value * multipliers[suffix_lower])
|
||||||
|
|
||||||
|
return bytes_value
|
||||||
|
|
||||||
|
|
||||||
|
def parse_cpu_quota(cpu_quota: int) -> int:
|
||||||
|
"""
|
||||||
|
Validate CPU quota value.
|
||||||
|
|
||||||
|
CPU quota is in microseconds per 100ms period.
|
||||||
|
100000 = 100% of one CPU core
|
||||||
|
|
||||||
|
Args:
|
||||||
|
cpu_quota: CPU quota in microseconds (e.g., 50000 for 50% of one core)
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Validated CPU quota value
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
ValueError: If quota <= 0 or > 1000000 (10 cores max)
|
||||||
|
"""
|
||||||
|
if cpu_quota <= 0:
|
||||||
|
raise ValueError(
|
||||||
|
f"CPU quota must be positive, got: {cpu_quota}"
|
||||||
|
)
|
||||||
|
|
||||||
|
# Maximum of 10 cores (1000000 microseconds)
|
||||||
|
if cpu_quota > 1000000:
|
||||||
|
raise ValueError(
|
||||||
|
f"CPU quota exceeds maximum of 1000000 (10 cores), got: {cpu_quota}"
|
||||||
|
)
|
||||||
|
|
||||||
|
return cpu_quota
|
||||||
|
|
||||||
|
|
||||||
|
def parse_storage_string(storage: str) -> int:
|
||||||
|
"""
|
||||||
|
Parse storage string to bytes (same as memory).
|
||||||
|
|
||||||
|
Args:
|
||||||
|
storage: Storage string with suffix (e.g., "1g", "512m")
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Storage in bytes
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
ValueError: If format is invalid or value is <= 0
|
||||||
|
"""
|
||||||
|
return parse_memory_string(storage)
|
||||||
|
|
||||||
|
|
||||||
|
class ResourceLimits:
|
||||||
|
"""
|
||||||
|
Resource limits with validation.
|
||||||
|
|
||||||
|
Encapsulates memory, storage, CPU, and timeout limits with validation.
|
||||||
|
Provides conversion to Podman container parameters.
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
memory: str,
|
||||||
|
storage: str,
|
||||||
|
cpu_quota: int,
|
||||||
|
timeout: int = 300
|
||||||
|
):
|
||||||
|
"""
|
||||||
|
Initialize resource limits with validation.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
memory: Memory limit string (e.g., "512m", "2g")
|
||||||
|
storage: Storage limit string (e.g., "1g", "10g")
|
||||||
|
cpu_quota: CPU quota in microseconds per 100ms period
|
||||||
|
timeout: Execution timeout in seconds (default: 300)
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
ValueError: If any limit is invalid
|
||||||
|
"""
|
||||||
|
self.memory_bytes = parse_memory_string(memory)
|
||||||
|
self.storage_bytes = parse_storage_string(storage)
|
||||||
|
self.cpu_quota = parse_cpu_quota(cpu_quota)
|
||||||
|
self.timeout = timeout
|
||||||
|
|
||||||
|
def to_podman_params(self) -> Dict[str, Any]:
|
||||||
|
"""
|
||||||
|
Convert to Podman container create parameters.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Dictionary of parameters suitable for Podman container creation
|
||||||
|
"""
|
||||||
|
return {
|
||||||
|
"mem_limit": str(self.memory_bytes), # Podman expects string
|
||||||
|
"cpu_quota": self.cpu_quota
|
||||||
|
# Note: storage_bytes tracked internally but not passed to Podman (not supported)
|
||||||
|
}
|
||||||
232
src/pod_executor/security/validation.py
Normal file
232
src/pod_executor/security/validation.py
Normal file
|
|
@ -0,0 +1,232 @@
|
||||||
|
"""
|
||||||
|
Security validation protocol for pod_executor.
|
||||||
|
|
||||||
|
Provides protocols (interfaces) for security validation that can be implemented
|
||||||
|
by consuming applications. Default implementations are provided.
|
||||||
|
"""
|
||||||
|
|
||||||
|
from typing import Protocol, Set, Optional, Dict, Any
|
||||||
|
|
||||||
|
|
||||||
|
class SecurityError(Exception):
|
||||||
|
"""Raised when security policy is violated."""
|
||||||
|
pass
|
||||||
|
|
||||||
|
|
||||||
|
class OperationValidatorProtocol(Protocol):
|
||||||
|
"""Protocol for validating container operations."""
|
||||||
|
|
||||||
|
session_containers: Set[str]
|
||||||
|
|
||||||
|
def validate_container_create(
|
||||||
|
self,
|
||||||
|
image: str,
|
||||||
|
params: Dict[str, Any],
|
||||||
|
session_id: Optional[str] = None
|
||||||
|
) -> None:
|
||||||
|
"""
|
||||||
|
Validate container creation parameters against security policy.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
image: Container image name
|
||||||
|
params: Container creation parameters
|
||||||
|
session_id: Optional session ID for volume validation
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
SecurityError: If any security policy is violated
|
||||||
|
"""
|
||||||
|
...
|
||||||
|
|
||||||
|
def validate_container_start(self, container_id: str) -> None:
|
||||||
|
"""
|
||||||
|
Validate container start.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
container_id: Container ID to start
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
SecurityError: If operation is not allowed
|
||||||
|
"""
|
||||||
|
...
|
||||||
|
|
||||||
|
def validate_container_stop(self, container_id: str) -> None:
|
||||||
|
"""
|
||||||
|
Validate container stop.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
container_id: Container ID to stop
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
SecurityError: If operation is not allowed
|
||||||
|
"""
|
||||||
|
...
|
||||||
|
|
||||||
|
def validate_operation(self, operation: str, target: str) -> tuple[bool, Optional[str]]:
|
||||||
|
"""
|
||||||
|
Validate generic operation.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
operation: Operation type
|
||||||
|
target: Operation target
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
(allowed, reason) tuple
|
||||||
|
"""
|
||||||
|
...
|
||||||
|
|
||||||
|
def register_session_container(self, container_id: str) -> None:
|
||||||
|
"""Register a container as belonging to a session."""
|
||||||
|
...
|
||||||
|
|
||||||
|
def unregister_session_container(self, container_id: str) -> None:
|
||||||
|
"""Unregister a session container."""
|
||||||
|
...
|
||||||
|
|
||||||
|
|
||||||
|
class NoOpValidator:
|
||||||
|
"""
|
||||||
|
No-op validator that allows all operations.
|
||||||
|
|
||||||
|
WARNING: This validator provides NO SECURITY. Only use for testing
|
||||||
|
or in fully trusted environments.
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(self):
|
||||||
|
"""Initialize with empty session container set."""
|
||||||
|
self.session_containers: Set[str] = set()
|
||||||
|
|
||||||
|
def validate_container_create(
|
||||||
|
self,
|
||||||
|
image: str,
|
||||||
|
params: Dict[str, Any],
|
||||||
|
session_id: Optional[str] = None
|
||||||
|
) -> None:
|
||||||
|
"""Allow all container creations."""
|
||||||
|
pass
|
||||||
|
|
||||||
|
def validate_container_start(self, container_id: str) -> None:
|
||||||
|
"""Allow all container starts."""
|
||||||
|
pass
|
||||||
|
|
||||||
|
def validate_container_stop(self, container_id: str) -> None:
|
||||||
|
"""Allow all container stops."""
|
||||||
|
pass
|
||||||
|
|
||||||
|
def validate_operation(self, operation: str, target: str) -> tuple[bool, Optional[str]]:
|
||||||
|
"""Allow all operations."""
|
||||||
|
return (True, None)
|
||||||
|
|
||||||
|
def register_session_container(self, container_id: str) -> None:
|
||||||
|
"""Track session container."""
|
||||||
|
self.session_containers.add(container_id)
|
||||||
|
|
||||||
|
def unregister_session_container(self, container_id: str) -> None:
|
||||||
|
"""Untrack session container."""
|
||||||
|
self.session_containers.discard(container_id)
|
||||||
|
|
||||||
|
|
||||||
|
class BasicValidator:
|
||||||
|
"""
|
||||||
|
Basic validator with minimal security checks.
|
||||||
|
|
||||||
|
Enforces:
|
||||||
|
- Allowed image patterns
|
||||||
|
- Required security parameters
|
||||||
|
- Forbidden dangerous parameters
|
||||||
|
- Session container tracking
|
||||||
|
"""
|
||||||
|
|
||||||
|
# Allowed container images with wildcard support
|
||||||
|
DEFAULT_ALLOWED_IMAGES = [
|
||||||
|
"python:3.11*",
|
||||||
|
"python:3.12*",
|
||||||
|
"jupyter/*",
|
||||||
|
"mcp-forge/*",
|
||||||
|
]
|
||||||
|
|
||||||
|
# Parameters that are forbidden
|
||||||
|
FORBIDDEN_PARAMS = [
|
||||||
|
"privileged",
|
||||||
|
"cap_add",
|
||||||
|
"devices",
|
||||||
|
"pid_mode",
|
||||||
|
]
|
||||||
|
|
||||||
|
def __init__(self, allowed_images: Optional[list[str]] = None):
|
||||||
|
"""
|
||||||
|
Initialize basic validator.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
allowed_images: List of allowed image patterns (supports wildcards)
|
||||||
|
"""
|
||||||
|
self.allowed_images = allowed_images or self.DEFAULT_ALLOWED_IMAGES
|
||||||
|
self.session_containers: Set[str] = set()
|
||||||
|
|
||||||
|
def validate_container_create(
|
||||||
|
self,
|
||||||
|
image: str,
|
||||||
|
params: Dict[str, Any],
|
||||||
|
session_id: Optional[str] = None
|
||||||
|
) -> None:
|
||||||
|
"""
|
||||||
|
Validate container creation.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
image: Container image name
|
||||||
|
params: Container creation parameters
|
||||||
|
session_id: Optional session ID
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
SecurityError: If validation fails
|
||||||
|
"""
|
||||||
|
# Validate image is allowed
|
||||||
|
if not self._is_image_allowed(image):
|
||||||
|
raise SecurityError(
|
||||||
|
f"Image '{image}' not in allowlist. "
|
||||||
|
f"Allowed patterns: {self.allowed_images}"
|
||||||
|
)
|
||||||
|
|
||||||
|
# Check for forbidden parameters
|
||||||
|
for forbidden in self.FORBIDDEN_PARAMS:
|
||||||
|
if forbidden in params:
|
||||||
|
raise SecurityError(
|
||||||
|
f"Forbidden parameter '{forbidden}' in container creation"
|
||||||
|
)
|
||||||
|
|
||||||
|
# Ensure security_opt includes no-new-privileges
|
||||||
|
security_opts = params.get("security_opt", [])
|
||||||
|
if "no-new-privileges" not in security_opts:
|
||||||
|
raise SecurityError(
|
||||||
|
"Container must include security_opt=['no-new-privileges']"
|
||||||
|
)
|
||||||
|
|
||||||
|
def validate_container_start(self, container_id: str) -> None:
|
||||||
|
"""Validate container start."""
|
||||||
|
if container_id not in self.session_containers:
|
||||||
|
raise SecurityError(
|
||||||
|
f"Container {container_id} is not a registered session container"
|
||||||
|
)
|
||||||
|
|
||||||
|
def validate_container_stop(self, container_id: str) -> None:
|
||||||
|
"""Validate container stop."""
|
||||||
|
if container_id not in self.session_containers:
|
||||||
|
raise SecurityError(
|
||||||
|
f"Container {container_id} is not a registered session container"
|
||||||
|
)
|
||||||
|
|
||||||
|
def validate_operation(self, operation: str, target: str) -> tuple[bool, Optional[str]]:
|
||||||
|
"""Validate generic operation."""
|
||||||
|
return (True, None) # Allow by default
|
||||||
|
|
||||||
|
def register_session_container(self, container_id: str) -> None:
|
||||||
|
"""Register session container."""
|
||||||
|
self.session_containers.add(container_id)
|
||||||
|
|
||||||
|
def unregister_session_container(self, container_id: str) -> None:
|
||||||
|
"""Unregister session container."""
|
||||||
|
self.session_containers.discard(container_id)
|
||||||
|
|
||||||
|
def _is_image_allowed(self, image: str) -> bool:
|
||||||
|
"""Check if image matches any allowed pattern."""
|
||||||
|
import fnmatch
|
||||||
|
return any(fnmatch.fnmatch(image, pattern) for pattern in self.allowed_images)
|
||||||
8
src/pod_executor/simple/__init__.py
Normal file
8
src/pod_executor/simple/__init__.py
Normal file
|
|
@ -0,0 +1,8 @@
|
||||||
|
"""Simple stateless code executor."""
|
||||||
|
|
||||||
|
from pod_executor.simple.executor import CodeExecutor, ExecutionResult
|
||||||
|
|
||||||
|
__all__ = [
|
||||||
|
"CodeExecutor",
|
||||||
|
"ExecutionResult",
|
||||||
|
]
|
||||||
222
src/pod_executor/simple/executor.py
Normal file
222
src/pod_executor/simple/executor.py
Normal file
|
|
@ -0,0 +1,222 @@
|
||||||
|
"""Code execution in isolated containers."""
|
||||||
|
|
||||||
|
from typing import Any, Optional
|
||||||
|
from dataclasses import dataclass, asdict
|
||||||
|
import json
|
||||||
|
import time
|
||||||
|
import textwrap
|
||||||
|
import uuid
|
||||||
|
|
||||||
|
from pod_executor.containers.manager import SecureContainerManager, ContainerConfig
|
||||||
|
from pod_executor.security.resource_limits import ResourceLimits
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class ExecutionResult:
|
||||||
|
"""Result of code execution."""
|
||||||
|
success: bool
|
||||||
|
stdout: str
|
||||||
|
stderr: str
|
||||||
|
result: Optional[Any]
|
||||||
|
execution_time: float
|
||||||
|
exit_code: int
|
||||||
|
error: Optional[str] = None
|
||||||
|
|
||||||
|
def to_dict(self) -> dict:
|
||||||
|
"""Convert to dictionary for JSON serialization."""
|
||||||
|
return asdict(self)
|
||||||
|
|
||||||
|
|
||||||
|
class CodeExecutor:
|
||||||
|
"""Executes Python code in isolated containers."""
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
container_manager: SecureContainerManager,
|
||||||
|
image: str,
|
||||||
|
resource_limits: ResourceLimits
|
||||||
|
):
|
||||||
|
self.container_manager = container_manager
|
||||||
|
self.image = image
|
||||||
|
self.resource_limits = resource_limits
|
||||||
|
|
||||||
|
def execute(
|
||||||
|
self,
|
||||||
|
code: str,
|
||||||
|
timeout: Optional[int] = None,
|
||||||
|
injection_code: Optional[str] = None,
|
||||||
|
bridge_socket_path: Optional[str] = None
|
||||||
|
) -> ExecutionResult:
|
||||||
|
"""
|
||||||
|
Execute Python code in a fresh container.
|
||||||
|
|
||||||
|
Process:
|
||||||
|
1. Create container with code
|
||||||
|
2. Start container
|
||||||
|
3. Wait for completion (with timeout)
|
||||||
|
4. Capture stdout/stderr
|
||||||
|
5. Extract result from last expression
|
||||||
|
6. Cleanup container
|
||||||
|
|
||||||
|
Args:
|
||||||
|
code: Python code to execute
|
||||||
|
timeout: Maximum execution time in seconds
|
||||||
|
injection_code: Optional MCP tool injection code to prepend
|
||||||
|
bridge_socket_path: Optional path to MCP bridge socket for mounting
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
ExecutionResult with stdout, stderr, result, and timing
|
||||||
|
"""
|
||||||
|
start_time = time.time()
|
||||||
|
container_id = None
|
||||||
|
|
||||||
|
# Use provided timeout or default from resource limits
|
||||||
|
exec_timeout = timeout if timeout is not None else self.resource_limits.timeout
|
||||||
|
|
||||||
|
try:
|
||||||
|
# Prepare code wrapper (with injection if provided)
|
||||||
|
wrapped_code = self._prepare_code(code, injection_code=injection_code)
|
||||||
|
|
||||||
|
# Generate a session ID for this execution to register the container
|
||||||
|
import uuid
|
||||||
|
session_id = f"simple-exec-{uuid.uuid4().hex[:12]}"
|
||||||
|
|
||||||
|
# Set up volumes for MCP bridge socket if provided
|
||||||
|
volumes = {}
|
||||||
|
if bridge_socket_path:
|
||||||
|
volumes[bridge_socket_path] = {
|
||||||
|
"bind": bridge_socket_path,
|
||||||
|
"mode": "rw"
|
||||||
|
}
|
||||||
|
|
||||||
|
# Create container configuration
|
||||||
|
config = ContainerConfig(
|
||||||
|
image=self.image,
|
||||||
|
command=["python", "-c", wrapped_code],
|
||||||
|
resource_limits=self.resource_limits,
|
||||||
|
volumes=volumes if volumes else None
|
||||||
|
)
|
||||||
|
|
||||||
|
# Create and start container with session_id for proper registration
|
||||||
|
container_id = self.container_manager.create_container(config, session_id=session_id)
|
||||||
|
self.container_manager.start_container(container_id)
|
||||||
|
|
||||||
|
# Wait for completion
|
||||||
|
exit_code = self.container_manager.wait_for_container(
|
||||||
|
container_id,
|
||||||
|
timeout=exec_timeout
|
||||||
|
)
|
||||||
|
|
||||||
|
# Get logs
|
||||||
|
stdout, stderr = self.container_manager.get_container_logs(container_id)
|
||||||
|
|
||||||
|
# Parse output to extract result
|
||||||
|
result, error = self._parse_output(stdout)
|
||||||
|
|
||||||
|
execution_time = time.time() - start_time
|
||||||
|
|
||||||
|
return ExecutionResult(
|
||||||
|
success=(exit_code == 0 and error is None),
|
||||||
|
stdout=stdout,
|
||||||
|
stderr=stderr,
|
||||||
|
result=result,
|
||||||
|
execution_time=execution_time,
|
||||||
|
exit_code=exit_code,
|
||||||
|
error=error
|
||||||
|
)
|
||||||
|
|
||||||
|
except TimeoutError 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 timeout: {str(e)}"
|
||||||
|
)
|
||||||
|
|
||||||
|
finally:
|
||||||
|
# Cleanup container
|
||||||
|
if container_id is not None:
|
||||||
|
try:
|
||||||
|
self.container_manager.remove_container(container_id)
|
||||||
|
except Exception:
|
||||||
|
pass # Best effort cleanup
|
||||||
|
|
||||||
|
def _prepare_code(self, code: str, injection_code: Optional[str] = None) -> str:
|
||||||
|
"""
|
||||||
|
Wrap code to capture result and handle errors.
|
||||||
|
|
||||||
|
Wraps code in try/except and captures:
|
||||||
|
- Last expression result
|
||||||
|
- Exceptions with traceback
|
||||||
|
- Execution metadata
|
||||||
|
|
||||||
|
Args:
|
||||||
|
code: User code to execute
|
||||||
|
injection_code: Optional MCP tool injection code to prepend
|
||||||
|
|
||||||
|
Returns wrapped code that outputs JSON to stdout.
|
||||||
|
"""
|
||||||
|
# Prepend injection code if provided
|
||||||
|
if injection_code:
|
||||||
|
full_code = injection_code + "\n\n" + code
|
||||||
|
else:
|
||||||
|
full_code = code
|
||||||
|
|
||||||
|
# Escape the code for embedding in exec string
|
||||||
|
escaped_code = full_code.replace('\\', '\\\\').replace('"', '\\"').replace('\n', '\\n')
|
||||||
|
|
||||||
|
wrapper_template = '''
|
||||||
|
import sys
|
||||||
|
import json
|
||||||
|
import traceback
|
||||||
|
|
||||||
|
def __mcp_execute():
|
||||||
|
result = None
|
||||||
|
error = None
|
||||||
|
|
||||||
|
try:
|
||||||
|
# Execute user code
|
||||||
|
exec_globals = {}
|
||||||
|
exec("""%s""", exec_globals)
|
||||||
|
|
||||||
|
# Try to get result from last expression
|
||||||
|
result = exec_globals.get('_', None)
|
||||||
|
|
||||||
|
except SyntaxError as e:
|
||||||
|
error = f"SyntaxError: {e.msg} (line {e.lineno})"
|
||||||
|
except Exception as e:
|
||||||
|
error = f"{type(e).__name__}: {str(e)}"
|
||||||
|
|
||||||
|
# Output result as JSON
|
||||||
|
print(json.dumps({"result": result, "error": error}))
|
||||||
|
|
||||||
|
__mcp_execute()
|
||||||
|
'''
|
||||||
|
|
||||||
|
return wrapper_template % escaped_code
|
||||||
|
|
||||||
|
def _parse_output(self, stdout: str) -> tuple[Optional[Any], Optional[str]]:
|
||||||
|
"""
|
||||||
|
Parse execution output to extract result and error.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
(result, error_message)
|
||||||
|
"""
|
||||||
|
if not stdout:
|
||||||
|
return None, None
|
||||||
|
|
||||||
|
try:
|
||||||
|
# First line should be JSON output
|
||||||
|
lines = stdout.split('\n', 1)
|
||||||
|
json_line = lines[0]
|
||||||
|
|
||||||
|
data = json.loads(json_line)
|
||||||
|
return data.get("result"), data.get("error")
|
||||||
|
|
||||||
|
except (json.JSONDecodeError, ValueError):
|
||||||
|
# If can't parse JSON, treat entire output as result
|
||||||
|
return None, None
|
||||||
Loading…
Add table
Add a link
Reference in a new issue