"""
Handle Manager for sktime MCP.
Manages references to instantiated estimator objects.
"""
import logging
import threading
import uuid
from collections import deque
from dataclasses import dataclass, field
from datetime import datetime
from typing import Any
logger = logging.getLogger(__name__)
[docs]
@dataclass
class HandleInfo:
"""Information about a managed handle."""
handle_id: str
estimator_name: str
instance: Any
params: dict[str, Any]
created_at: datetime
fitted: bool = False
metadata: dict[str, Any] = field(default_factory=dict)
[docs]
def to_dict(self) -> dict[str, Any]:
return {
"handle_id": self.handle_id,
"estimator_name": self.estimator_name,
"params": self.params,
"created_at": self.created_at.isoformat(),
"fitted": self.fitted,
"metadata": self.metadata,
}
[docs]
class HandleManager:
"""Manager for estimator instance handles."""
[docs]
def __init__(self, max_handles: int = 100):
self._handles: dict[str, HandleInfo] = {}
self._max_handles = max_handles
# Tombstones: ids evicted to stay under the cap, so a later lookup can
# say "evicted" instead of an indistinguishable "not found".
self._evicted: deque[str] = deque(maxlen=1024)
# Handles are read and written from worker threads (fit_async runs fit
# in a thread pool) while the asyncio thread serves other tool calls.
# Reentrant because create_handle evicts while already holding it.
self._lock = threading.RLock()
[docs]
def describe_missing(self, handle_id: str) -> str:
"""Message for a handle that isn't present — distinguishes evicted from unknown."""
with self._lock:
if handle_id in self._evicted:
return (
f"Estimator handle '{handle_id}' was evicted (handle limit "
f"{self._max_handles} reached); re-create it with instantiate."
)
return f"Handle not found: {handle_id}"
[docs]
def was_evicted(self, handle_id: str) -> bool:
with self._lock:
return handle_id in self._evicted
[docs]
def create_handle(
self,
estimator_name: str,
instance: Any,
params: dict[str, Any] | None = None,
metadata: dict[str, Any] | None = None,
) -> str:
with self._lock:
if len(self._handles) >= self._max_handles:
self._cleanup_oldest()
handle_id = f"est_{uuid.uuid4().hex[:12]}"
handle_info = HandleInfo(
handle_id=handle_id,
estimator_name=estimator_name,
instance=instance,
params=params or {},
created_at=datetime.now(),
metadata=metadata or {},
)
self._handles[handle_id] = handle_info
return handle_id
[docs]
def get_instance(self, handle_id: str) -> Any:
with self._lock:
if handle_id not in self._handles:
raise KeyError(self.describe_missing(handle_id))
return self._handles[handle_id].instance
[docs]
def get_info(self, handle_id: str) -> HandleInfo:
with self._lock:
if handle_id not in self._handles:
raise KeyError(self.describe_missing(handle_id))
return self._handles[handle_id]
[docs]
def exists(self, handle_id: str) -> bool:
with self._lock:
return handle_id in self._handles
[docs]
def replace_instance(self, handle_id: str, instance: Any) -> None:
"""Swap the live instance behind a handle (e.g. rollback after a failed update)."""
with self._lock:
if handle_id in self._handles:
self._handles[handle_id].instance = instance
[docs]
def mark_fitted(self, handle_id: str) -> None:
with self._lock:
if handle_id in self._handles:
self._handles[handle_id].fitted = True
[docs]
def is_fitted(self, handle_id: str) -> bool:
with self._lock:
if handle_id not in self._handles:
return False
return self._handles[handle_id].fitted
[docs]
def release_handle(self, handle_id: str) -> bool:
with self._lock:
if handle_id in self._handles:
del self._handles[handle_id]
return True
return False
[docs]
def list_handles(self) -> list[dict[str, Any]]:
with self._lock:
return [info.to_dict() for info in self._handles.values()]
[docs]
def clear_all(self) -> int:
with self._lock:
count = len(self._handles)
self._handles.clear()
return count
def _cleanup_oldest(self, count: int = 10) -> None:
"""Evict the oldest handles. Caller must hold ``self._lock``."""
with self._lock:
sorted_handles = sorted(
self._handles.items(),
key=lambda x: x[1].created_at,
)
for handle_id, _ in sorted_handles[:count]:
del self._handles[handle_id]
self._evicted.append(handle_id)
logger.info(
"Evicted estimator handle %s (limit %d reached)", handle_id, self._max_handles
)
_handle_manager_instance: HandleManager | None = None
[docs]
def get_handle_manager() -> HandleManager:
global _handle_manager_instance
if _handle_manager_instance is None:
_handle_manager_instance = HandleManager()
return _handle_manager_instance