mantisdk 0.1.0__py3-none-any.whl

mantisdk/types/core.py

@@ -0,0 +1,553 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Core data models shared across Mantisdk components."""
+
+from __future__ import annotations
+
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Dict,
+    Generic,
+    Iterator,
+    List,
+    Literal,
+    Mapping,
+    Optional,
+    Protocol,
+    Sequence,
+    SupportsIndex,
+    TypedDict,
+    TypeVar,
+    Union,
+    cast,
+    overload,
+)
+
+from opentelemetry.sdk.trace import ReadableSpan
+from pydantic import BaseModel, Field, model_validator
+
+from .tracer import Span, SpanCoreFields
+
+if TYPE_CHECKING:
+    from mantisdk.litagent import LitAgent
+    from mantisdk.runner.base import Runner
+    from mantisdk.tracer.base import Tracer
+
+__all__ = [
+    "Triplet",
+    "RolloutLegacy",
+    "Task",
+    "TaskInput",
+    "TaskIfAny",
+    "RolloutRawResultLegacy",
+    "RolloutRawResult",
+    "RolloutMode",
+    "GenericResponse",
+    "ParallelWorkerBase",
+    "Dataset",
+    "AttemptStatus",
+    "RolloutStatus",
+    "RolloutConfig",
+    "Rollout",
+    "Attempt",
+    "AttemptedRollout",
+    "EnqueueRolloutRequest",
+    "Hook",
+    "Worker",
+    "WorkerStatus",
+    "PaginatedResult",
+    "FilterOptions",
+    "SortOptions",
+    "FilterField",
+]
+
+T_co = TypeVar("T_co", covariant=True)
+
+
+class Triplet(BaseModel):
+    """Single interaction turn captured during reinforcement learning."""
+
+    prompt: Any
+    response: Any
+    reward: Optional[float] = None
+    metadata: Dict[str, Any] = Field(default_factory=dict)
+
+
+class RolloutLegacy(BaseModel):
+    """Legacy reporting payload exchanged with the deprecated HTTP server.
+
+    !!! warning "Deprecated"
+        Use [`Rollout`][mantisdk.Rollout] instead.
+    """
+
+    rollout_id: str
+
+    # Echoing the input task
+    task: Optional[Task] = None
+
+    # Primary, high-level feedback
+    final_reward: Optional[float] = None
+
+    # Structured, sequential feedback for RL-style optimization
+    triplets: Optional[List[Triplet]] = None
+
+    # Optional, rich-context data for deep analysis
+    trace: Optional[List[Dict[str, Any]]] = Field(
+        default=None,
+        description="A list of spans that conform to the OpenTelemetry JSON format. "
+        "Users of the opentelemetry-sdk can generate this by calling "
+        "json.loads(readable_span.to_json()).",
+    )
+    logs: Optional[List[str]] = None
+
+    # A bucket for any other relevant information
+    metadata: Dict[str, Any] = Field(default_factory=dict)
+
+
+RolloutStatus = Literal[
+    "queuing",  # initial status
+    "preparing",  # after the trace is claimed
+    "running",  # after receiving the first trace
+    "failed",  # crashed
+    "succeeded",  # status OK
+    "cancelled",  # cancelled by user (or watchdog)
+    "requeuing",  # retrying
+]
+"""The status of a rollout."""
+
+AttemptStatus = Literal[
+    # A status is essentially a process.
+    # It should not have scheduling/management statuses like "queuing" or "cancelled".
+    "preparing",
+    "running",
+    "failed",
+    "succeeded",
+    "unresponsive",  # the worker has not reported results for a while
+    "timeout",  # the worker has been emitting new logs, but have been working on the task for too long
+]
+"""The status of an attempt."""
+
+RolloutMode = Literal["train", "val", "test"]
+"""Possible rollout modes."""
+
+
+class Attempt(BaseModel):
+    """Execution attempt for a rollout, including metadata for retries."""
+
+    rollout_id: str
+    """The rollout which this attempt belongs to."""
+    attempt_id: str
+    """The universal id for current attempt."""
+    sequence_id: int
+    """The sequence number of the attempt, starting from 1."""
+    start_time: float
+    """The time when the attempt has started."""
+    end_time: Optional[float] = None
+    """The time when the attempt has ended."""
+    status: AttemptStatus = "preparing"
+    """The status of the attempt."""
+    worker_id: Optional[str] = None
+    """The rollout worker which is executing this attempt."""
+
+    last_heartbeat_time: Optional[float] = None
+    """The last time when the worker has reported progress (i.e., a span)."""
+
+    metadata: Optional[Dict[str, Any]] = None
+    """A bucket for any other relevant information."""
+
+
+class RolloutConfig(BaseModel):
+    """Configuration controlling rollout retries and timeouts."""
+
+    timeout_seconds: Optional[float] = None
+    """The timeout for the rollout, in seconds. None indicates no timeout."""
+    unresponsive_seconds: Optional[float] = None
+    """The unresponsive timeout for the rollout, in seconds. None indicates no unresponsive timeout."""
+    max_attempts: int = Field(default=1, ge=1)
+    """The maximum number of attempts for the rollout, including the first attempt."""
+    retry_condition: List[AttemptStatus] = Field(default_factory=cast(Callable[[], List[AttemptStatus]], list))
+    """The list of statuses that should trigger a retry."""
+
+
+class Rollout(BaseModel):
+    rollout_id: str
+    """Unique identifier for the rollout."""
+
+    input: TaskInput
+    """Task input used to generate the rollout."""
+
+    # Time to track the lifecycle of the rollout
+    start_time: float
+    """Timestamp when the rollout started."""
+    end_time: Optional[float] = None
+    """Timestamp when the rollout ended."""
+
+    mode: Optional[RolloutMode] = None
+    """Execution mode such as `"train"`, `"val"` or `"test"`. See [`RolloutMode`][mantisdk.RolloutMode]."""
+    resources_id: Optional[str] = None
+    """Identifier of the resources required to execute the rollout."""
+
+    status: RolloutStatus = "queuing"
+    """Latest status emitted by the controller."""
+
+    config: RolloutConfig = Field(default_factory=RolloutConfig)
+    """Retry and timeout configuration associated with the rollout."""
+
+    metadata: Optional[Dict[str, Any]] = None
+    """Additional metadata attached to the rollout."""
+
+
+class AttemptedRollout(Rollout):
+    """Rollout paired with the currently active attempt."""
+
+    attempt: Attempt
+    """The attempt that is currently processing the rollout."""
+
+    @model_validator(mode="after")
+    def check_consistency(self) -> AttemptedRollout:
+        if self.attempt.rollout_id != self.rollout_id:
+            raise ValueError("Inconsistent rollout_id between Rollout and Attempt")
+        return self
+
+
+class EnqueueRolloutRequest(BaseModel):
+    """Payload describing a rollout to be queued via [`enqueue_rollout`][mantisdk.LightningStore.enqueue_rollout].
+
+    A subset of fields from [`Rollout`][mantisdk.Rollout] used for queuing new rollouts.
+    """
+
+    input: TaskInput
+    """Task input used to generate the rollout."""
+    mode: Optional[RolloutMode] = None
+    """Execution mode such as `"train"`, `"val"` or `"test"`. See [`RolloutMode`][mantisdk.RolloutMode]."""
+    resources_id: Optional[str] = None
+    """Identifier of the resources required to execute the rollout."""
+    config: Optional[RolloutConfig] = None
+    """Retry and timeout configuration associated with the rollout."""
+    metadata: Optional[Dict[str, Any]] = None
+    """Additional metadata attached to the rollout."""
+
+
+WorkerStatus = Literal["idle", "busy", "unknown"]
+
+
+class Worker(BaseModel):
+    """Worker information. This is actually the same as Runner info."""
+
+    worker_id: str
+    """The ID of the worker."""
+    status: WorkerStatus = "unknown"
+    """The status of the worker."""
+    heartbeat_stats: Optional[Dict[str, Any]] = None
+    """Statistics about the worker's heartbeat."""
+    last_heartbeat_time: Optional[float] = None
+    """The last time when the worker has reported the stats."""
+    last_dequeue_time: Optional[float] = None
+    """The last time when the worker has tried to dequeue a rollout."""
+    last_busy_time: Optional[float] = None
+    """The last time when the worker has started an attempt and became busy."""
+    last_idle_time: Optional[float] = None
+    """The last time when the worker has triggered the end of an attempt and became idle."""
+    current_rollout_id: Optional[str] = None
+    """The ID of the current rollout that the worker is processing."""
+    current_attempt_id: Optional[str] = None
+    """The ID of the current attempt that the worker is processing."""
+
+
+TaskInput = Any
+"""Task input type. Accepts arbitrary payloads."""
+
+
+class Task(BaseModel):
+    """Rollout request served to client agents.
+
+    !!! warning "Deprecated"
+        The legacy HTTP client/server stack still uses this model. Prefer
+        [`LightningStore`][mantisdk.LightningStore] APIs for new workflows.
+    """
+
+    rollout_id: str
+    input: TaskInput
+
+    mode: Optional[RolloutMode] = None
+    resources_id: Optional[str] = None
+
+    # Optional fields for tracking task lifecycle
+    create_time: Optional[float] = None
+    last_claim_time: Optional[float] = None
+    num_claims: Optional[int] = None
+
+    # Allow additional metadata fields
+    metadata: Dict[str, Any] = Field(default_factory=dict)
+
+
+class TaskIfAny(BaseModel):
+    """A task or indication that no task is available.
+
+    !!! warning "Deprecated"
+        Use [`LightningStore`][mantisdk.LightningStore] APIs for new workflows.
+    """
+
+    is_available: bool
+    """Indication that a task is available."""
+    task: Optional[Task] = None
+
+
+RolloutRawResultLegacy = Union[None, float, List[Triplet], List[Dict[str, Any]], List[ReadableSpan], RolloutLegacy]
+"""Legacy rollout result type.
+
+!!! warning "Deprecated"
+    Use [`RolloutRawResult`][mantisdk.RolloutRawResult] instead.
+"""
+
+RolloutRawResult = Union[
+    None,  # nothing (relies on tracer)
+    float,  # only final reward
+    List[ReadableSpan],  # constructed OTEL spans by user
+    List[Span],  # constructed Span objects by user
+    List[SpanCoreFields],  # constructed SpanCoreFields objects by user
+]
+"""Rollout result type.
+
+Possible return values of [`rollout`][mantisdk.LitAgent.rollout].
+"""
+
+
+class GenericResponse(BaseModel):
+    """Generic server response used by compatibility endpoints.
+
+    !!! warning "Deprecated"
+        This response is no longer used by the new
+        [`LightningStore`][mantisdk.LightningStore] APIs.
+
+    Attributes:
+        status: Status string describing the result of the request.
+        message: Optional human readable explanation.
+        data: Arbitrary payload serialized as JSON.
+    """
+
+    status: str = "success"
+    message: Optional[str] = None
+    data: Optional[Dict[str, Any]] = None
+
+
+class ParallelWorkerBase:
+    """Base class for workloads executed across multiple worker processes.
+
+    The lifecycle is orchestrated by the main process:
+
+    * [`init()`][mantisdk.ParallelWorkerBase.init] prepares shared state.
+    * Each worker calls [`init_worker()`][mantisdk.ParallelWorkerBase.init_worker] during start-up.
+    * [`run()`][mantisdk.ParallelWorkerBase.run] performs the parallel workload.
+    * Workers call [`teardown_worker()`][mantisdk.ParallelWorkerBase.teardown_worker] before exiting.
+    * The main process finalizes through [`teardown()`][mantisdk.ParallelWorkerBase.teardown].
+
+    Subclasses must implement [`run()`][mantisdk.ParallelWorkerBase.run]
+    and can override other lifecycle hooks.
+    """
+
+    def __init__(self) -> None:
+        """Initialize the base class. This method can be overridden by subclasses."""
+        self.worker_id: Optional[int] = None
+
+    def init(self, *args: Any, **kwargs: Any) -> None:
+        """Initialize before spawning the workers. This method can be overridden by subclasses."""
+        pass
+
+    def init_worker(self, worker_id: int, *args: Any, **kwargs: Any) -> None:
+        """Initialize the worker. This method can be overridden by subclasses."""
+        self.worker_id = worker_id
+
+    def run(self, *args: Any, **kwargs: Any) -> Any:
+        """Run the workload. This method can be overridden by subclasses."""
+        pass
+
+    def teardown_worker(self, worker_id: int, *args: Any, **kwargs: Any) -> None:
+        """Teardown the worker. This method can be overridden by subclasses."""
+        pass
+
+    def teardown(self, *args: Any, **kwargs: Any) -> None:
+        """Teardown after the workers have exited. This method can be overridden by subclasses."""
+        pass
+
+
+class Dataset(Protocol, Generic[T_co]):
+    """The general interface for a dataset.
+
+    It's currently implemented as a protocol, having a similar interface to `torch.utils.data.Dataset`.
+    You don't have to inherit from this class; you can use a simple list if you want to.
+    """
+
+    def __getitem__(self, index: SupportsIndex, /) -> T_co: ...
+
+    def __len__(self) -> int: ...
+
+
+class Hook(ParallelWorkerBase):
+    """Base class for defining hooks in the agent runner's lifecycle."""
+
+    async def on_trace_start(
+        self, *, agent: LitAgent[Any], runner: Runner[Any], tracer: Tracer, rollout: Rollout
+    ) -> None:
+        """Hook called immediately after the tracer enters the trace context but before the rollout begins.
+
+        Args:
+            agent: The [`LitAgent`][mantisdk.LitAgent] instance associated with the runner.
+            runner: The [`Runner`][mantisdk.Runner] managing the rollout.
+            tracer: The [`Tracer`][mantisdk.Tracer] instance associated with the runner.
+            rollout: The [`Rollout`][mantisdk.Rollout] object that will be processed.
+
+        Subclasses can override this method to implement custom logic such as logging,
+        metric collection, or resource setup. By default, this is a no-op.
+        """
+
+    async def on_trace_end(
+        self, *, agent: LitAgent[Any], runner: Runner[Any], tracer: Tracer, rollout: Rollout
+    ) -> None:
+        """Hook called immediately after the rollout completes but before the tracer exits the trace context.
+
+        Args:
+            agent: The [`LitAgent`][mantisdk.LitAgent] instance associated with the runner.
+            runner: The [`Runner`][mantisdk.Runner] managing the rollout.
+            tracer: The [`Tracer`][mantisdk.Tracer] instance associated with the runner.
+            rollout: The [`Rollout`][mantisdk.Rollout] object that has been processed.
+
+        Subclasses can override this method to implement custom logic such as logging,
+        metric collection, or resource cleanup. By default, this is a no-op.
+        """
+
+    async def on_rollout_start(self, *, agent: LitAgent[Any], runner: Runner[Any], rollout: Rollout) -> None:
+        """Hook called immediately before a rollout *attempt* begins.
+
+        Args:
+            agent: The [`LitAgent`][mantisdk.LitAgent] instance associated with the runner.
+            runner: The [`Runner`][mantisdk.Runner] managing the rollout.
+            rollout: The [`Rollout`][mantisdk.Rollout] object that will be processed.
+
+        Subclasses can override this method to implement custom logic such as
+        logging, metric collection, or resource setup. By default, this is a
+        no-op.
+        """
+
+    async def on_rollout_end(
+        self,
+        *,
+        agent: LitAgent[Any],
+        runner: Runner[Any],
+        rollout: Rollout,
+        spans: Union[List[ReadableSpan], List[Span]],
+    ) -> None:
+        """Hook called after a rollout *attempt* completes.
+
+        Args:
+            agent: The [`LitAgent`][mantisdk.LitAgent] instance associated with the runner.
+            runner: The [`Runner`][mantisdk.Runner] managing the rollout.
+            rollout: The [`Rollout`][mantisdk.Rollout] object that has been processed.
+            spans: The spans that have been added to the store.
+
+        Subclasses can override this method for cleanup or additional
+        logging. By default, this is a no-op.
+        """
+
+
+class FilterField(TypedDict, total=False):
+    """An operator dict for a single field."""
+
+    exact: Any
+    within: Sequence[Any]
+    contains: str
+
+
+FilterOptions = Mapping[
+    Union[str, Literal["_aggregate", "_must"]],
+    Union[FilterField, Literal["and", "or"], Mapping[str, FilterField]],
+]
+"""A mapping of field name -> operator dict.
+
+Each operator dict can contain:
+
+- "exact": value for exact equality.
+- "within": iterable of allowed values.
+- "contains": substring to search for in string fields.
+
+The filter can also have a special field called "_aggregate" that can be used to specify the logic
+to combine the results of the filters:
+
+- "and": all conditions must match. This is the default value if not specified.
+- "or": at least one condition must match.
+
+All conditions within a field and between different fields are
+stored in a unified pool and combined using `_aggregate`.
+
+The filter can also have a special group called "_must", which is a mapping of filters that must all match,
+no matter whether the aggregate logic is "and" or "or".
+
+Example:
+
+```json
+{
+    "_aggregate": "or",
+    "_must": {
+        "city": {"exact": "New York"},
+        "timezone": {"within": ["America/New_York", "America/Los_Angeles"]},
+    },
+    "status": {"exact": "active"},
+    "id": {"within": [1, 2, 3]},
+    "name": {"contains": "foo"},
+}
+```
+"""
+
+
+class SortOptions(TypedDict):
+    """Options for sorting the collection."""
+
+    name: str
+    """The name of the field to sort by."""
+    order: Literal["asc", "desc"]
+    """The order to sort by."""
+
+
+T_item = TypeVar("T_item")
+
+
+class PaginatedResult(BaseModel, Sequence[T_item]):
+    """Result of a paginated query.
+
+    Behaves like a sequence, but also carries pagination metadata (limit, offset, total).
+    """
+
+    items: Sequence[T_item]
+    """Items in the result."""
+    limit: int
+    """Limit of the result."""
+    offset: int
+    """Offset of the result."""
+    total: int
+    """Total number of items in the collection."""
+
+    def __len__(self) -> int:
+        return len(self.items)
+
+    @overload
+    def __getitem__(self, index: int) -> T_item: ...
+
+    @overload
+    def __getitem__(self, index: slice) -> Sequence[T_item]: ...
+
+    def __getitem__(self, index: Union[int, slice]) -> Union[T_item, Sequence[T_item]]:
+        return self.items[index]
+
+    # Overriding __iter__ enables list(paginated_result) to work as expected,
+    # but changes Pydantic's default dict iteration behavior (which would otherwise
+    # iterate over field names).
+    def __iter__(self) -> Iterator[T_item]:  # type: ignore
+        return iter(self.items)
+
+    def __repr__(self) -> str:
+        first_item_repr = repr(self.items[0]) if self.items else "empty"
+        items_repr = f"[{first_item_repr}, ...]" if len(self.items) > 1 else first_item_repr
+        slice_repr = f"{self.offset}:" if self.limit == -1 else f"{self.offset}:{self.offset + self.limit}"
+        return f"<PaginatedResult ({slice_repr} of {self.total}) {items_repr}>"