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:
Hans Aschauer
parent
8b6b237be9
commit
db75b822f4
14 changed files with 2842 additions and 0 deletions
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"
|
||||
)
|
||||
Loading…
Add table
Add a link
Reference in a new issue