mantisdk 0.1.0__py3-none-any.whl

mantisdk/store/threading.py

@@ -0,0 +1,357 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+import threading
+from typing import Any, Dict, List, Literal, Optional, Sequence, Tuple
+
+from opentelemetry.sdk.trace import ReadableSpan
+
+from mantisdk.types import (
+    Attempt,
+    AttemptedRollout,
+    AttemptStatus,
+    EnqueueRolloutRequest,
+    NamedResources,
+    ResourcesUpdate,
+    Rollout,
+    RolloutConfig,
+    RolloutStatus,
+    Span,
+    TaskInput,
+    Worker,
+    WorkerStatus,
+)
+
+from .base import UNSET, LightningStore, LightningStoreCapabilities, LightningStoreStatistics, Unset
+
+
+class LightningStoreThreaded(LightningStore):
+    """Facade that delegates all store operations to a underlying store instance.
+
+    The operations are guaranteed to be thread-safe.
+    Make sure the threaded stores are instantiated before initializing the threads.
+    """
+
+    def __init__(self, store: LightningStore) -> None:
+        super().__init__()  # watchdog relies on the underlying store
+        self.store = store
+        self._lock = threading.Lock()
+
+    @property
+    def capabilities(self) -> LightningStoreCapabilities:
+        """Return the capabilities of the store."""
+        capabilities = self.store.capabilities
+        return {
+            **capabilities,
+            "async_safe": True,
+            "thread_safe": True,
+        }
+
+    async def statistics(self) -> LightningStoreStatistics:
+        """Return the statistics of the store."""
+        with self._lock:
+            return await self.store.statistics()
+
+    def otlp_traces_endpoint(self) -> str:
+        """Return the OTLP/HTTP traces endpoint of the underlying store."""
+        return self.store.otlp_traces_endpoint()
+
+    def get_otlp_headers(self) -> Dict[str, str]:
+        """Return the OTLP authentication headers from the underlying store."""
+        if hasattr(self.store, "get_otlp_headers"):
+            return self.store.get_otlp_headers()
+        return {}
+
+    async def start_rollout(
+        self,
+        input: TaskInput,
+        mode: Literal["train", "val", "test"] | None = None,
+        resources_id: str | None = None,
+        config: RolloutConfig | None = None,
+        metadata: Dict[str, Any] | None = None,
+        worker_id: Optional[str] = None,
+    ) -> AttemptedRollout:
+        with self._lock:
+            return await self.store.start_rollout(
+                input,
+                mode,
+                resources_id,
+                config,
+                metadata,
+                worker_id,
+            )
+
+    async def enqueue_rollout(
+        self,
+        input: TaskInput,
+        mode: Literal["train", "val", "test"] | None = None,
+        resources_id: str | None = None,
+        config: RolloutConfig | None = None,
+        metadata: Dict[str, Any] | None = None,
+    ) -> Rollout:
+        with self._lock:
+            return await self.store.enqueue_rollout(input, mode, resources_id, config, metadata)
+
+    async def enqueue_many_rollouts(self, rollouts: Sequence[EnqueueRolloutRequest]) -> Sequence[Rollout]:
+        with self._lock:
+            return await self.store.enqueue_many_rollouts(rollouts)
+
+    async def dequeue_rollout(self, worker_id: Optional[str] = None) -> Optional[AttemptedRollout]:
+        with self._lock:
+            return await self.store.dequeue_rollout(worker_id=worker_id)
+
+    async def dequeue_many_rollouts(
+        self,
+        *,
+        limit: int = 1,
+        worker_id: Optional[str] = None,
+    ) -> Sequence[AttemptedRollout]:
+        with self._lock:
+            return await self.store.dequeue_many_rollouts(limit=limit, worker_id=worker_id)
+
+    async def start_attempt(self, rollout_id: str, worker_id: Optional[str] = None) -> AttemptedRollout:
+        with self._lock:
+            return await self.store.start_attempt(rollout_id, worker_id)
+
+    async def query_rollouts(
+        self,
+        *,
+        status_in: Optional[Sequence[RolloutStatus]] = None,
+        rollout_id_in: Optional[Sequence[str]] = None,
+        rollout_id_contains: Optional[str] = None,
+        filter_logic: Literal["and", "or"] = "and",
+        sort_by: Optional[str] = None,
+        sort_order: Literal["asc", "desc"] = "asc",
+        limit: int = -1,
+        offset: int = 0,
+        status: Optional[Sequence[RolloutStatus]] = None,
+        rollout_ids: Optional[Sequence[str]] = None,
+    ) -> Sequence[Rollout]:
+        with self._lock:
+            return await self.store.query_rollouts(
+                status_in=status_in,
+                rollout_id_in=rollout_id_in,
+                rollout_id_contains=rollout_id_contains,
+                filter_logic=filter_logic,
+                sort_by=sort_by,
+                sort_order=sort_order,
+                limit=limit,
+                offset=offset,
+                status=status,
+                rollout_ids=rollout_ids,
+            )
+
+    async def query_attempts(
+        self,
+        rollout_id: str,
+        *,
+        sort_by: Optional[str] = "sequence_id",
+        sort_order: Literal["asc", "desc"] = "asc",
+        limit: int = -1,
+        offset: int = 0,
+    ) -> Sequence[Attempt]:
+        with self._lock:
+            return await self.store.query_attempts(
+                rollout_id,
+                sort_by=sort_by,
+                sort_order=sort_order,
+                limit=limit,
+                offset=offset,
+            )
+
+    async def get_rollout_by_id(self, rollout_id: str) -> Optional[Rollout]:
+        with self._lock:
+            return await self.store.get_rollout_by_id(rollout_id)
+
+    async def get_latest_attempt(self, rollout_id: str) -> Optional[Attempt]:
+        with self._lock:
+            return await self.store.get_latest_attempt(rollout_id)
+
+    async def query_resources(
+        self,
+        *,
+        resources_id: Optional[str] = None,
+        resources_id_contains: Optional[str] = None,
+        sort_by: Optional[str] = None,
+        sort_order: Literal["asc", "desc"] = "asc",
+        limit: int = -1,
+        offset: int = 0,
+    ) -> Sequence[ResourcesUpdate]:
+        with self._lock:
+            return await self.store.query_resources(
+                resources_id=resources_id,
+                resources_id_contains=resources_id_contains,
+                sort_by=sort_by,
+                sort_order=sort_order,
+                limit=limit,
+                offset=offset,
+            )
+
+    async def add_resources(self, resources: NamedResources) -> ResourcesUpdate:
+        with self._lock:
+            return await self.store.add_resources(resources)
+
+    async def update_resources(self, resources_id: str, resources: NamedResources) -> ResourcesUpdate:
+        with self._lock:
+            return await self.store.update_resources(resources_id, resources)
+
+    async def get_resources_by_id(self, resources_id: str) -> Optional[ResourcesUpdate]:
+        with self._lock:
+            return await self.store.get_resources_by_id(resources_id)
+
+    async def get_latest_resources(self) -> Optional[ResourcesUpdate]:
+        with self._lock:
+            return await self.store.get_latest_resources()
+
+    async def add_many_spans(self, spans: Sequence[Span]) -> Sequence[Span]:
+        with self._lock:
+            return await self.store.add_many_spans(spans)
+
+    async def add_span(self, span: Span) -> Optional[Span]:
+        with self._lock:
+            return await self.store.add_span(span)
+
+    async def add_otel_span(
+        self,
+        rollout_id: str,
+        attempt_id: str,
+        readable_span: ReadableSpan,
+        sequence_id: int | None = None,
+    ) -> Optional[Span]:
+        with self._lock:
+            return await self.store.add_otel_span(rollout_id, attempt_id, readable_span, sequence_id)
+
+    async def wait_for_rollouts(self, *, rollout_ids: List[str], timeout: Optional[float] = None) -> List[Rollout]:
+        # This method does not change the state of the store, and it's not thread-safe.
+        return await self.store.wait_for_rollouts(rollout_ids=rollout_ids, timeout=timeout)
+
+    async def get_next_span_sequence_id(self, rollout_id: str, attempt_id: str) -> int:
+        with self._lock:
+            return await self.store.get_next_span_sequence_id(rollout_id, attempt_id)
+
+    async def get_many_span_sequence_ids(self, rollout_attempt_ids: Sequence[Tuple[str, str]]) -> Sequence[int]:
+        with self._lock:
+            return await self.store.get_many_span_sequence_ids(rollout_attempt_ids)
+
+    async def query_spans(
+        self,
+        rollout_id: str,
+        attempt_id: str | Literal["latest"] | None = None,
+        *,
+        trace_id: Optional[str] = None,
+        trace_id_contains: Optional[str] = None,
+        span_id: Optional[str] = None,
+        span_id_contains: Optional[str] = None,
+        parent_id: Optional[str] = None,
+        parent_id_contains: Optional[str] = None,
+        name: Optional[str] = None,
+        name_contains: Optional[str] = None,
+        filter_logic: Literal["and", "or"] = "and",
+        limit: int = -1,
+        offset: int = 0,
+        sort_by: Optional[str] = "sequence_id",
+        sort_order: Literal["asc", "desc"] = "asc",
+    ) -> Sequence[Span]:
+        with self._lock:
+            return await self.store.query_spans(
+                rollout_id,
+                attempt_id,
+                trace_id=trace_id,
+                trace_id_contains=trace_id_contains,
+                span_id=span_id,
+                span_id_contains=span_id_contains,
+                parent_id=parent_id,
+                parent_id_contains=parent_id_contains,
+                name=name,
+                name_contains=name_contains,
+                filter_logic=filter_logic,
+                limit=limit,
+                offset=offset,
+                sort_by=sort_by,
+                sort_order=sort_order,
+            )
+
+    async def update_rollout(
+        self,
+        rollout_id: str,
+        input: TaskInput | Unset = UNSET,
+        mode: Optional[Literal["train", "val", "test"]] | Unset = UNSET,
+        resources_id: Optional[str] | Unset = UNSET,
+        status: RolloutStatus | Unset = UNSET,
+        config: RolloutConfig | Unset = UNSET,
+        metadata: Optional[Dict[str, Any]] | Unset = UNSET,
+    ) -> Rollout:
+        with self._lock:
+            return await self.store.update_rollout(
+                rollout_id=rollout_id,
+                input=input,
+                mode=mode,
+                resources_id=resources_id,
+                status=status,
+                config=config,
+                metadata=metadata,
+            )
+
+    async def update_attempt(
+        self,
+        rollout_id: str,
+        attempt_id: str | Literal["latest"],
+        status: AttemptStatus | Unset = UNSET,
+        worker_id: str | Unset = UNSET,
+        last_heartbeat_time: float | Unset = UNSET,
+        metadata: Optional[Dict[str, Any]] | Unset = UNSET,
+    ) -> Attempt:
+        with self._lock:
+            return await self.store.update_attempt(
+                rollout_id=rollout_id,
+                attempt_id=attempt_id,
+                status=status,
+                worker_id=worker_id,
+                last_heartbeat_time=last_heartbeat_time,
+                metadata=metadata,
+            )
+
+    async def query_workers(
+        self,
+        *,
+        status_in: Optional[Sequence[WorkerStatus]] = None,
+        worker_id_contains: Optional[str] = None,
+        filter_logic: Literal["and", "or"] = "and",
+        sort_by: Optional[str] = None,
+        sort_order: Literal["asc", "desc"] = "asc",
+        limit: int = -1,
+        offset: int = 0,
+    ) -> Sequence[Worker]:
+        with self._lock:
+            return await self.store.query_workers(
+                status_in=status_in,
+                worker_id_contains=worker_id_contains,
+                sort_by=sort_by,
+                sort_order=sort_order,
+                limit=limit,
+                offset=offset,
+            )
+
+    async def get_worker_by_id(self, worker_id: str) -> Optional[Worker]:
+        with self._lock:
+            return await self.store.get_worker_by_id(worker_id)
+
+    async def update_worker(
+        self,
+        worker_id: str,
+        heartbeat_stats: Dict[str, Any] | Unset = UNSET,
+    ) -> Worker:
+        with self._lock:
+            return await self.store.update_worker(
+                worker_id=worker_id,
+                heartbeat_stats=heartbeat_stats,
+            )
+
+    def complete_job(self, summary: Optional[Dict[str, Any]] = None) -> None:
+        """Delegate job completion to the underlying store."""
+        with self._lock:
+            # We check if the underlying store has the method to be safe,
+            # though base LightningStore now has it.
+            if hasattr(self.store, "complete_job"):
+                self.store.complete_job(summary)

mantisdk/store/utils.py

@@ -0,0 +1,142 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import time
+from typing import Awaitable, Callable, Dict, List, Tuple
+
+from mantisdk.types import Attempt, AttemptedRollout, AttemptStatus, Rollout, RolloutConfig, RolloutStatus
+
+UpdateRolloutStatus = Callable[[str, RolloutStatus], Awaitable[Rollout]]
+UpdateAttemptStatus = Callable[[str, str, AttemptStatus], Awaitable[Attempt]]
+
+
+LATENCY_BUCKETS = [
+    0.000001,
+    0.000002,
+    0.000005,
+    0.00001,
+    0.00002,
+    0.00005,
+    0.0001,
+    0.0002,
+    0.0005,
+    0.001,
+    0.002,
+    0.003,
+    0.005,
+    0.007,
+    0.01,
+    0.015,
+    0.02,
+    0.03,
+    0.05,
+    0.07,
+    0.1,
+    0.2,
+    0.3,
+    0.5,
+    0.7,
+    1.0,
+    2.0,
+    3.0,
+    5.0,
+    7.0,
+    10.0,
+    12.0,
+    15.0,
+    20.0,
+    25.0,
+    30.0,
+    40.0,
+    50.0,
+    60.0,
+    90.0,
+    120.0,
+    180.0,
+    240.0,
+    300.0,
+]
+
+
+async def rollout_status_from_attempt(
+    attempt: Attempt,
+    config: RolloutConfig,
+) -> RolloutStatus:
+    """
+    Propagate the status of an attempt to the rollout.
+
+    Returns:
+        The status of the rollout from the perspective of the attempt.
+    """
+    # Propagate the status directly to the rollout
+    if attempt.status == "preparing" or attempt.status == "running" or attempt.status == "succeeded":
+        return attempt.status
+
+    if attempt.status == "failed" or attempt.status == "timeout" or attempt.status == "unresponsive":
+        # Check if this status should trigger a retry
+        if attempt.status in config.retry_condition:
+            # If we haven't exceeded max attempts, retry
+            if attempt.sequence_id < config.max_attempts:
+                return "requeuing"
+
+        # If we can't retry or shouldn't retry, mark as failed
+        return "failed"
+
+    raise ValueError(f"Invalid attempt status: {attempt.status}")
+
+
+async def scan_unhealthy_rollouts(
+    rollouts: List[AttemptedRollout],
+) -> Dict[Tuple[str, str], AttemptStatus]:
+    """
+    Perform health check on all running rollouts in the store.
+
+    This method should be called periodically to:
+
+    1. Check for unresponsive attempts (no heartbeat or spans for a while)
+    2. Check for timed-out rollouts (running too long since start_time)
+
+    This operation is completely unlocked. The caller is responsible for locking the store.
+
+    Args:
+        rollouts: The list of running rollouts to check.
+
+    Returns:
+        A dictionary of updates to the rollouts.
+    """
+    current_time = time.time()
+    updates: Dict[Tuple[str, str], AttemptStatus] = {}
+
+    for rollout in rollouts:
+        config = rollout.config  # policy for retry and timeout
+
+        # Get the latest attempt for this rollout
+        latest_attempt = rollout.attempt
+        if not latest_attempt:
+            # This should not happen
+            continue
+
+        # Check for timeout condition (based on attempt start_time, instead of rollout start_time)
+        if config.timeout_seconds is not None and current_time - latest_attempt.start_time > config.timeout_seconds:
+            updates[(latest_attempt.rollout_id, latest_attempt.attempt_id)] = "timeout"
+            continue
+
+        # Check for unresponsive condition (based on last heartbeat)
+        # (1) Haven't received heartbeat for a while
+        if (
+            latest_attempt.last_heartbeat_time
+            and config.unresponsive_seconds is not None
+            and current_time - latest_attempt.last_heartbeat_time > config.unresponsive_seconds
+        ):
+            updates[(latest_attempt.rollout_id, latest_attempt.attempt_id)] = "unresponsive"
+            continue
+
+        # (2) Check if there's no last heartbeat (no spans) at all
+        if (
+            latest_attempt.last_heartbeat_time is None
+            and config.unresponsive_seconds is not None
+            and current_time - latest_attempt.start_time > config.unresponsive_seconds
+        ):
+            updates[(latest_attempt.rollout_id, latest_attempt.attempt_id)] = "unresponsive"
+            continue
+
+    return updates

mantisdk/tracer/__init__.py

@@ -0,0 +1,16 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from .agentops import AgentOpsTracer
+from .base import Tracer, clear_active_tracer, get_active_tracer, set_active_tracer
+from .dummy import DummyTracer
+from .otel import OtelTracer
+
+__all__ = [
+    "AgentOpsTracer",
+    "Tracer",
+    "OtelTracer",
+    "DummyTracer",
+    "get_active_tracer",
+    "set_active_tracer",
+    "clear_active_tracer",
+]