"""Prompt-scoped remote session helpers for split Modal proxy execution.""" from __future__ import annotations import logging import threading import time from dataclasses import dataclass, field from typing import Any, Mapping logger = logging.getLogger(__name__) _REMOTE_SESSION_HANDLE_MARKER = "__comfy_modal_remote_session_handle__" _REMOTE_SESSION_VALUE_REF_MARKER = "__comfy_modal_remote_session_value_ref__" class RemoteSessionStateError(RuntimeError): """Raised when remote session metadata is invalid or missing.""" @dataclass(frozen=True) class RemoteSessionHandle: """Opaque prompt-scoped handle used to group remote-only runtime state.""" session_id: str prompt_id: str | None = None owner_component_id: str | None = None def to_payload(self) -> dict[str, Any]: """Serialize this handle into a JSON-safe payload mapping.""" return { _REMOTE_SESSION_HANDLE_MARKER: True, "session_id": self.session_id, "prompt_id": self.prompt_id, "owner_component_id": self.owner_component_id, } @classmethod def from_payload(cls, payload: Mapping[str, Any]) -> "RemoteSessionHandle": """Deserialize a previously serialized handle payload.""" session_id = str(payload.get("session_id") or "").strip() if not session_id: raise RemoteSessionStateError("Remote session handles must define session_id.") prompt_id = payload.get("prompt_id") owner_component_id = payload.get("owner_component_id") return cls( session_id=session_id, prompt_id=(str(prompt_id) if prompt_id is not None else None), owner_component_id=( str(owner_component_id) if owner_component_id is not None else None ), ) @dataclass(frozen=True) class RemoteSessionValueRef: """Reference one remote-only value stored inside a prompt-scoped session.""" session_id: str node_id: str output_index: int def to_payload(self) -> dict[str, Any]: """Serialize this value reference into a JSON-safe payload mapping.""" return { _REMOTE_SESSION_VALUE_REF_MARKER: True, "session_id": self.session_id, "node_id": self.node_id, "output_index": self.output_index, } @classmethod def from_payload(cls, payload: Mapping[str, Any]) -> "RemoteSessionValueRef": """Deserialize a previously serialized value-reference payload.""" session_id = str(payload.get("session_id") or "").strip() node_id = str(payload.get("node_id") or "").strip() if not session_id or not node_id: raise RemoteSessionStateError( "Remote session value refs must define session_id and node_id." ) output_index = payload.get("output_index") if output_index is None: raise RemoteSessionStateError("Remote session value refs must define output_index.") return cls( session_id=session_id, node_id=node_id, output_index=int(output_index), ) def is_remote_session_handle_payload(payload: Any) -> bool: """Return whether one arbitrary value looks like a session-handle payload.""" return isinstance(payload, Mapping) and bool(payload.get(_REMOTE_SESSION_HANDLE_MARKER)) def is_remote_session_value_ref_payload(payload: Any) -> bool: """Return whether one arbitrary value looks like a session-value-ref payload.""" return isinstance(payload, Mapping) and bool(payload.get(_REMOTE_SESSION_VALUE_REF_MARKER)) @dataclass class _RemoteSessionBucket: """Hold the live values cached for one prompt-scoped remote session.""" handle: RemoteSessionHandle values: dict[tuple[str, int], Any] = field(default_factory=dict) created_at: float = field(default_factory=time.time) class InMemoryRemoteSessionStore: """Thread-safe in-memory store for remote-only values shared across proxy calls.""" def __init__(self) -> None: """Initialize the empty session store.""" self._lock = threading.Lock() self._sessions: dict[str, _RemoteSessionBucket] = {} def ensure_session(self, handle: RemoteSessionHandle) -> RemoteSessionHandle: """Create one session bucket when it does not yet exist.""" with self._lock: bucket = self._sessions.get(handle.session_id) if bucket is None: self._sessions[handle.session_id] = _RemoteSessionBucket(handle=handle) logger.info( "Created remote session session_id=%s prompt_id=%s owner_component_id=%s.", handle.session_id, handle.prompt_id, handle.owner_component_id, ) return handle logger.info( "Reusing remote session session_id=%s prompt_id=%s owner_component_id=%s age_seconds=%.3f value_count=%d.", bucket.handle.session_id, bucket.handle.prompt_id, bucket.handle.owner_component_id, max(0.0, time.time() - bucket.created_at), len(bucket.values), ) return bucket.handle def put_output( self, handle: RemoteSessionHandle, *, node_id: str, output_index: int, value: Any, ) -> RemoteSessionValueRef: """Store one output value under a prompt-scoped session and return its ref.""" self.ensure_session(handle) ref = RemoteSessionValueRef( session_id=handle.session_id, node_id=str(node_id), output_index=int(output_index), ) with self._lock: bucket = self._sessions[handle.session_id] value_key = (ref.node_id, ref.output_index) replacing_existing = value_key in bucket.values bucket.values[value_key] = value logger.info( "Stored remote session value session_id=%s node_id=%s output_index=%d result=%s value_type=%s total_value_count=%d.", ref.session_id, ref.node_id, ref.output_index, "replace" if replacing_existing else "create", type(value).__name__, len(bucket.values), ) return ref def get_output(self, ref: RemoteSessionValueRef) -> Any: """Return one stored output value or raise when it no longer exists.""" with self._lock: bucket = self._sessions.get(ref.session_id) if bucket is None: logger.warning( "Remote session lookup missed session_id=%s node_id=%s output_index=%d reason=session-missing.", ref.session_id, ref.node_id, ref.output_index, ) raise RemoteSessionStateError( f"Remote session {ref.session_id!r} was not found." ) value_key = (ref.node_id, ref.output_index) if value_key not in bucket.values: logger.warning( "Remote session lookup missed session_id=%s node_id=%s output_index=%d reason=value-missing age_seconds=%.3f value_count=%d.", ref.session_id, ref.node_id, ref.output_index, max(0.0, time.time() - bucket.created_at), len(bucket.values), ) raise RemoteSessionStateError( "Remote session value was not found for " f"session_id={ref.session_id!r} node_id={ref.node_id!r} " f"output_index={ref.output_index}." ) logger.info( "Resolved remote session value session_id=%s node_id=%s output_index=%d age_seconds=%.3f value_count=%d.", ref.session_id, ref.node_id, ref.output_index, max(0.0, time.time() - bucket.created_at), len(bucket.values), ) return bucket.values[value_key] def clear_session(self, handle: RemoteSessionHandle) -> None: """Drop one prompt-scoped session and every value stored inside it.""" with self._lock: removed = self._sessions.pop(handle.session_id, None) if removed is not None: logger.info( "Cleared remote session session_id=%s prompt_id=%s owner_component_id=%s age_seconds=%.3f value_count=%d.", handle.session_id, removed.handle.prompt_id, removed.handle.owner_component_id, max(0.0, time.time() - removed.created_at), len(removed.values), ) return logger.warning( "Remote session clear skipped session_id=%s prompt_id=%s owner_component_id=%s reason=session-missing.", handle.session_id, handle.prompt_id, handle.owner_component_id, ) def resolve_value(self, value: Any) -> Any: """Resolve one possible session reference back into the underlying live value.""" if not is_remote_session_value_ref_payload(value): return value ref = RemoteSessionValueRef.from_payload(value) logger.info( "Resolving remote session ref session_id=%s node_id=%s output_index=%d.", ref.session_id, ref.node_id, ref.output_index, ) return self.get_output(ref)