Add LightningStore interface and implementation (#118)

agentlightning/store/base.py

@@ -0,0 +1,258 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+from typing import Any, Dict, List, Literal, Optional, Sequence
+
+from opentelemetry.sdk.trace import ReadableSpan
+
+from agentlightning.tracer import Span
+from agentlightning.types import (
+    Attempt,
+    AttemptedRollout,
+    AttemptStatus,
+    NamedResources,
+    ResourcesUpdate,
+    RolloutConfig,
+    RolloutStatus,
+    RolloutV2,
+    TaskInput,
+)
+
+
+def is_queuing(rollout: RolloutV2) -> bool:
+    return rollout.status == "queuing" or rollout.status == "requeuing"
+
+
+def is_running(rollout: RolloutV2) -> bool:
+    return rollout.status == "preparing" or rollout.status == "running"
+
+
+def is_finished(rollout: RolloutV2) -> bool:
+    return rollout.status == "failed" or rollout.status == "succeeded" or rollout.status == "cancelled"
+
+
+class _UnsetType:
+    """A sentinel type to indicate an unset value."""
+
+    __slots__ = ()
+
+    def __repr__(self) -> str:
+        return "UNSET"
+
+    def __reduce__(self):
+        return (_get_unset, ())
+
+
+def _get_unset() -> _UnsetType:
+    return UNSET
+
+
+UNSET = _UnsetType()
+Unset = _UnsetType  # Alias for convenience
+
+
+class LightningStore:
+    """
+    A centralized, thread-safe, async, data store for the lightning's state.
+    This holds the task queue, versioned resources, and completed rollouts.
+
+    The store has a built-in clock and it should be responsible for tracking the times.
+    All the time-based operations like retry, timeout, etc. should be handled by the store.
+    """
+
+    async def start_rollout(
+        self,
+        input: TaskInput,
+        mode: Literal["train", "val", "test"] | None = None,
+        resources_id: str | None = None,
+        metadata: Dict[str, Any] | None = None,
+    ) -> AttemptedRollout:
+        """
+        Add one incomplete rollout to the store, and get an attempt created for it.
+        This will immediately sets the rollout to a preparing state, and should be
+        used by whoever is going to execute the rollout.
+
+        Return a special rollout with attempt object. Do not update it directly.
+
+        But if the rollout fails or timeouts, it's still possible that the watchdog
+        sends it back to the queue for retry.
+
+        To enqueue a rollout to the task queue, use `enqueue_rollout` instead.
+        """
+        raise NotImplementedError()
+
+    async def enqueue_rollout(
+        self,
+        input: TaskInput,
+        mode: Literal["train", "val", "test"] | None = None,
+        resources_id: str | None = None,
+        metadata: Dict[str, Any] | None = None,
+    ) -> RolloutV2:
+        """
+        Adds a new task to the queue with specific metadata and
+        returns the rollout object with its unique ID.
+        """
+        raise NotImplementedError()
+
+    async def dequeue_rollout(self) -> Optional[AttemptedRollout]:
+        """
+        Retrieves the next task from the queue without blocking.
+        Returns None if the queue is empty.
+
+        Will set the rollout status to preparing.
+        """
+        raise NotImplementedError()
+
+    async def start_attempt(self, rollout_id: str) -> AttemptedRollout:
+        """
+        Create a new attempt for a given rollout ID and return the attempt details.
+        """
+        raise NotImplementedError()
+
+    async def add_span(self, span: Span) -> Span:
+        """
+        Add a span to the store.
+
+        This method is responsible for updating the rollout/attempt status to "running" if needed.
+        """
+        raise NotImplementedError()
+
+    async def add_otel_span(
+        self,
+        rollout_id: str,
+        attempt_id: str,
+        readable_span: ReadableSpan,
+        sequence_id: int | None = None,
+    ) -> Span:
+        """
+        Add an opentelemetry span to the store.
+
+        If sequence_id is not provided, it will be fetched from `get_next_span_sequence_id` and assigned automatically.
+        """
+        raise NotImplementedError()
+
+    async def query_rollouts(
+        self, *, status: Optional[Sequence[RolloutStatus]] = None, rollout_ids: Optional[Sequence[str]] = None
+    ) -> List[RolloutV2]:
+        """
+        Query and retrieve rollouts filtered by their status.
+        If no status is provided, returns all rollouts.
+        """
+        raise NotImplementedError()
+
+    async def query_attempts(self, rollout_id: str) -> List[Attempt]:
+        """
+        Query and retrieve all attempts associated with a specific rollout ID.
+        Returns an empty list if no attempts are found.
+        """
+        raise NotImplementedError()
+
+    async def get_rollout_by_id(self, rollout_id: str) -> Optional[RolloutV2]:
+        """
+        Safely retrieves a specific rollout by its ID.
+        """
+        raise NotImplementedError()
+
+    async def get_latest_attempt(self, rollout_id: str) -> Optional[Attempt]:
+        """
+        Safely retrieves the latest attempt for a given rollout ID.
+        """
+        raise NotImplementedError()
+
+    async def get_resources_by_id(self, resources_id: str) -> Optional[ResourcesUpdate]:
+        """
+        Safely retrieves a specific version of named resources by its ID.
+        """
+        raise NotImplementedError()
+
+    async def get_latest_resources(self) -> Optional[ResourcesUpdate]:
+        """
+        Safely retrieves the latest version of named resources.
+        """
+        raise NotImplementedError()
+
+    async def get_next_span_sequence_id(self, rollout_id: str, attempt_id: str) -> int:
+        """
+        Get the next span sequence ID for a given rollout and attempt.
+        This should be used to assign a unique sequence ID to each span within an attempt.
+
+        Recommend getting the ID before the operation even begins to avoid racing conditions.
+        """
+        raise NotImplementedError()
+
+    async def wait_for_rollouts(self, *, rollout_ids: List[str], timeout: Optional[float] = None) -> List[RolloutV2]:
+        """
+        Wait for specified rollouts to complete with a timeout.
+        Returns the completed rollouts, potentially incomplete if timeout is reached.
+
+        TODO: Add support for waiting for 20 new rollouts, or wait until 80% of the pending ids are completed.
+        """
+        raise NotImplementedError()
+
+    async def query_spans(self, rollout_id: str, attempt_id: str | Literal["latest"] | None = None) -> List[Span]:
+        """
+        Query and retrieve all spans associated with a specific rollout ID.
+        Returns an empty list if no spans are found.
+        """
+        raise NotImplementedError()
+
+    async def update_resources(self, resources_id: str, resources: NamedResources) -> ResourcesUpdate:
+        """
+        Safely stores a new version of named resources and sets it as the latest.
+        """
+        raise NotImplementedError()
+
+    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,
+    ) -> RolloutV2:
+        """
+        Update the rollout status and related metadata.
+
+        Not-listed fields here either cannot be updated, or should be auto-updated (e.g., end_time).
+
+        When status is updated to a finished / problematic state, other states like task
+        queues will be updated accordingly.
+
+        Args:
+            rollout_id: Unique identifier for the rollout to update
+            input: New input data for the rollout. If set, will be updated. Can be updated to None
+            mode: New mode for the rollout. If set, will be updated. Can be updated to None
+            resources_id: New resources ID for the rollout. If set, will be updated. Can be updated to None
+            status: New status for the rollout. If set, will be updated
+            config: New config for the rollout. If set, will be updated
+            metadata: Dictionary of additional metadata to update. If set, will replace the existing metadata
+        """
+        raise NotImplementedError()
+
+    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:
+        """
+        Update a specific or latest attempt for a given rollout.
+
+        Update the latest attempt will NOT affect the corresponding rollout status.
+
+
+        Args:
+            rollout_id: Unique identifier for the rollout
+            attempt_id: Unique identifier for the attempt
+            status: Status to set for the attempt, update if provided
+            worker_id: Worker identifier, update if provided
+            last_heartbeat_time: Timestamp of the last heartbeat from the worker
+            metadata: Dictionary of additional metadata to update, will replace the existing metadata
+        """
+        raise NotImplementedError()

agentlightning/store/client_server.py

@@ -0,0 +1,575 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+import asyncio
+import logging
+from typing import Any, Dict, List, Literal, Optional, Sequence, Union
+
+import aiohttp
+import uvicorn
+from fastapi import FastAPI
+from opentelemetry.sdk.trace import ReadableSpan
+from pydantic import BaseModel, Field
+
+from agentlightning.tracer import Span
+from agentlightning.types import (
+    Attempt,
+    AttemptedRollout,
+    AttemptStatus,
+    NamedResources,
+    ResourcesUpdate,
+    RolloutConfig,
+    RolloutStatus,
+    RolloutV2,
+    TaskInput,
+)
+
+from .base import UNSET, LightningStore, Unset
+
+logger = logging.getLogger(__name__)
+
+
+class PydanticUnset(BaseModel):
+    _type: Literal["UNSET"] = "UNSET"
+
+
+class RolloutRequest(BaseModel):
+    input: TaskInput
+    mode: Optional[Literal["train", "val", "test"]] = None
+    resources_id: Optional[str] = None
+    metadata: Optional[Dict[str, Any]] = None
+
+
+class QueryRolloutsRequest(BaseModel):
+    status: Optional[List[RolloutStatus]] = None
+    rollout_ids: Optional[List[str]] = None
+
+
+class WaitForRolloutsRequest(BaseModel):
+    rollout_ids: List[str]
+    timeout: Optional[float] = None
+
+
+class RolloutId(BaseModel):
+    rollout_id: str
+
+
+class UpdateRolloutRequest(BaseModel):
+    rollout_id: str
+    input: Union[TaskInput, PydanticUnset] = Field(default_factory=PydanticUnset)
+    mode: Union[Optional[Literal["train", "val", "test"]], PydanticUnset] = Field(default_factory=PydanticUnset)
+    resources_id: Union[Optional[str], PydanticUnset] = Field(default_factory=PydanticUnset)
+    status: Union[RolloutStatus, PydanticUnset] = Field(default_factory=PydanticUnset)
+    config: Union[RolloutConfig, PydanticUnset] = Field(default_factory=PydanticUnset)
+    metadata: Union[Dict[str, Any], PydanticUnset] = Field(default_factory=PydanticUnset)
+
+
+class UpdateAttemptRequest(BaseModel):
+    rollout_id: str
+    attempt_id: Union[str, Literal["latest"]]
+    status: Union[AttemptStatus, PydanticUnset] = Field(default_factory=PydanticUnset)
+    worker_id: Union[str, PydanticUnset] = Field(default_factory=PydanticUnset)
+    last_heartbeat_time: Union[float, PydanticUnset] = Field(default_factory=PydanticUnset)
+    metadata: Union[Dict[str, Any], PydanticUnset] = Field(default_factory=PydanticUnset)
+
+
+class LightningStoreServer(LightningStore):
+    """
+    Server wrapper that exposes a LightningStore via HTTP API.
+    Delegates all operations to an underlying store implementation.
+
+    Healthcheck and watchdog relies on the underlying store.
+    """
+
+    def __init__(self, store: LightningStore, host: str, port: int):
+        super().__init__()
+        self.store = store
+        self.host = host
+        self.port = port
+        self.app = FastAPI(title="LightningStore Server")
+        self._setup_routes()
+        self._uvicorn_config = uvicorn.Config(self.app, host=self.host, port=self.port, log_level="info")
+        self._uvicorn_server = uvicorn.Server(self._uvicorn_config)
+
+    @property
+    def endpoint(self) -> str:
+        return f"http://{self.host}:{self.port}"
+
+    async def start(self):
+        """Starts the FastAPI server in the background."""
+        logger.info(f"Starting server at {self.endpoint}")
+        asyncio.create_task(self._uvicorn_server.serve())
+        await asyncio.sleep(1)  # Allow time for server to start up.
+
+    async def stop(self):
+        """Gracefully stops the running FastAPI server."""
+        if self._uvicorn_server.started:
+            logger.info("Stopping server...")
+            self._uvicorn_server.should_exit = True
+            await asyncio.sleep(1)  # Allow time for graceful shutdown.
+            logger.info("Server stopped.")
+
+    def _setup_routes(self):
+        """Set up FastAPI routes for all store operations."""
+
+        @self.app.post("/start_rollout", response_model=AttemptedRollout)
+        async def start_rollout(request: RolloutRequest):  # pyright: ignore[reportUnusedFunction]
+            return await self.store.start_rollout(
+                input=request.input,
+                mode=request.mode,
+                resources_id=request.resources_id,
+                metadata=request.metadata,
+            )
+
+        @self.app.post("/enqueue_rollout", response_model=RolloutV2)
+        async def enqueue_rollout(request: RolloutRequest):  # pyright: ignore[reportUnusedFunction]
+            return await self.store.enqueue_rollout(
+                input=request.input,
+                mode=request.mode,
+                resources_id=request.resources_id,
+                metadata=request.metadata,
+            )
+
+        @self.app.get("/dequeue_rollout", response_model=Optional[AttemptedRollout])
+        async def dequeue_rollout():  # pyright: ignore[reportUnusedFunction]
+            return await self.store.dequeue_rollout()
+
+        @self.app.post("/start_attempt", response_model=AttemptedRollout)
+        async def start_attempt(request: RolloutId):  # pyright: ignore[reportUnusedFunction]
+            return await self.store.start_attempt(request.rollout_id)
+
+        @self.app.post("/query_rollouts", response_model=List[RolloutV2])
+        async def query_rollouts(request: QueryRolloutsRequest):  # pyright: ignore[reportUnusedFunction]
+            return await self.store.query_rollouts(status=request.status)
+
+        @self.app.get("/query_attempts/{rollout_id}", response_model=List[Attempt])
+        async def query_attempts(rollout_id: str):  # pyright: ignore[reportUnusedFunction]
+            return await self.store.query_attempts(rollout_id)
+
+        @self.app.get("/get_latest_attempt/{rollout_id}", response_model=Optional[Attempt])
+        async def get_latest_attempt(rollout_id: str):  # pyright: ignore[reportUnusedFunction]
+            return await self.store.get_latest_attempt(rollout_id)
+
+        @self.app.get("/get_rollout_by_id/{rollout_id}", response_model=Optional[RolloutV2])
+        async def get_rollout_by_id(rollout_id: str):  # pyright: ignore[reportUnusedFunction]
+            return await self.store.get_rollout_by_id(rollout_id)
+
+        @self.app.post("/update_resources", response_model=ResourcesUpdate)
+        async def update_resources(update: ResourcesUpdate):  # pyright: ignore[reportUnusedFunction]
+            return await self.store.update_resources(update.resources_id, update.resources)
+
+        @self.app.get("/get_resources_by_id/{resources_id}", response_model=Optional[ResourcesUpdate])
+        async def get_resources_by_id(resources_id: str):  # pyright: ignore[reportUnusedFunction]
+            return await self.store.get_resources_by_id(resources_id)
+
+        @self.app.get("/get_latest_resources", response_model=Optional[ResourcesUpdate])
+        async def get_latest_resources():  # pyright: ignore[reportUnusedFunction]
+            return await self.store.get_latest_resources()
+
+        @self.app.post("/add_span", response_model=Span)
+        async def add_span(span: Span):  # pyright: ignore[reportUnusedFunction]
+            return await self.store.add_span(span)
+
+        @self.app.get("/get_next_span_sequence_id/{rollout_id}/{attempt_id}", response_model=int)
+        async def get_next_span_sequence_id(rollout_id: str, attempt_id: str):  # pyright: ignore[reportUnusedFunction]
+            return await self.store.get_next_span_sequence_id(rollout_id, attempt_id)
+
+        @self.app.post("/wait_for_rollouts", response_model=List[RolloutV2])
+        async def wait_for_rollouts(request: WaitForRolloutsRequest):  # pyright: ignore[reportUnusedFunction]
+            return await self.store.wait_for_rollouts(rollout_ids=request.rollout_ids, timeout=request.timeout)
+
+        @self.app.get("/query_spans/{rollout_id}", response_model=List[Span])
+        async def query_spans(  # pyright: ignore[reportUnusedFunction]
+            rollout_id: str, attempt_id: Optional[str] = None
+        ):
+            return await self.store.query_spans(rollout_id, attempt_id)
+
+        @self.app.post("/update_rollout", response_model=RolloutV2)
+        async def update_rollout(request: UpdateRolloutRequest):  # pyright: ignore[reportUnusedFunction]
+            return await self.store.update_rollout(
+                rollout_id=request.rollout_id,
+                input=request.input if not isinstance(request.input, PydanticUnset) else UNSET,
+                mode=request.mode if not isinstance(request.mode, PydanticUnset) else UNSET,
+                resources_id=request.resources_id if not isinstance(request.resources_id, PydanticUnset) else UNSET,
+                status=request.status if not isinstance(request.status, PydanticUnset) else UNSET,
+                config=request.config if not isinstance(request.config, PydanticUnset) else UNSET,
+                metadata=request.metadata if not isinstance(request.metadata, PydanticUnset) else UNSET,
+            )
+
+        @self.app.post("/update_attempt", response_model=Attempt)
+        async def update_attempt(request: UpdateAttemptRequest):  # pyright: ignore[reportUnusedFunction]
+            return await self.store.update_attempt(
+                rollout_id=request.rollout_id,
+                attempt_id=request.attempt_id,
+                status=request.status if not isinstance(request.status, PydanticUnset) else UNSET,
+                worker_id=request.worker_id if not isinstance(request.worker_id, PydanticUnset) else UNSET,
+                last_heartbeat_time=(
+                    request.last_heartbeat_time if not isinstance(request.last_heartbeat_time, PydanticUnset) else UNSET
+                ),
+                metadata=request.metadata if not isinstance(request.metadata, PydanticUnset) else UNSET,
+            )
+
+    # Delegate methods -------------------------------------------------
+    async def start_rollout(
+        self,
+        input: TaskInput,
+        mode: Literal["train", "val", "test"] | None = None,
+        resources_id: str | None = None,
+        metadata: Dict[str, Any] | None = None,
+    ) -> AttemptedRollout:
+        return await self.store.start_rollout(input, mode, resources_id, metadata)
+
+    async def enqueue_rollout(
+        self,
+        input: TaskInput,
+        mode: Literal["train", "val", "test"] | None = None,
+        resources_id: str | None = None,
+        metadata: Dict[str, Any] | None = None,
+    ) -> RolloutV2:
+        return await self.store.enqueue_rollout(input, mode, resources_id, metadata)
+
+    async def dequeue_rollout(self) -> Optional[AttemptedRollout]:
+        return await self.store.dequeue_rollout()
+
+    async def start_attempt(self, rollout_id: str) -> AttemptedRollout:
+        return await self.store.start_attempt(rollout_id)
+
+    async def query_rollouts(
+        self, *, status: Optional[Sequence[RolloutStatus]] = None, rollout_ids: Optional[Sequence[str]] = None
+    ) -> List[RolloutV2]:
+        return await self.store.query_rollouts(status=status, rollout_ids=rollout_ids)
+
+    async def query_attempts(self, rollout_id: str) -> List[Attempt]:
+        return await self.store.query_attempts(rollout_id)
+
+    async def get_latest_attempt(self, rollout_id: str) -> Optional[Attempt]:
+        return await self.store.get_latest_attempt(rollout_id)
+
+    async def get_rollout_by_id(self, rollout_id: str) -> Optional[RolloutV2]:
+        return await self.store.get_rollout_by_id(rollout_id)
+
+    async def update_resources(self, resources_id: str, resources: NamedResources) -> ResourcesUpdate:
+        return await self.store.update_resources(resources_id, resources)
+
+    async def get_resources_by_id(self, resources_id: str) -> Optional[ResourcesUpdate]:
+        return await self.store.get_resources_by_id(resources_id)
+
+    async def get_latest_resources(self) -> Optional[ResourcesUpdate]:
+        return await self.store.get_latest_resources()
+
+    async def add_span(self, span: Span) -> Span:
+        return await self.store.add_span(span)
+
+    async def get_next_span_sequence_id(self, rollout_id: str, attempt_id: str) -> int:
+        return await self.store.get_next_span_sequence_id(rollout_id, attempt_id)
+
+    async def add_otel_span(
+        self,
+        rollout_id: str,
+        attempt_id: str,
+        readable_span: ReadableSpan,
+        sequence_id: int | None = None,
+    ) -> Span:
+        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[RolloutV2]:
+        return await self.store.wait_for_rollouts(rollout_ids=rollout_ids, timeout=timeout)
+
+    async def query_spans(
+        self,
+        rollout_id: str,
+        attempt_id: str | Literal["latest"] | None = None,
+    ) -> List[Span]:
+        return await self.store.query_spans(rollout_id, attempt_id)
+
+    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,
+    ) -> RolloutV2:
+        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:
+        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,
+        )
+
+
+class LightningStoreClient(LightningStore):
+    """HTTP client that talks to a remote LightningStoreServer."""
+
+    def __init__(self, server_address: str):
+        self.server_address = server_address.rstrip("/")
+        self.session: Optional[aiohttp.ClientSession] = None
+
+    async def _get_session(self) -> aiohttp.ClientSession:
+        """Get or create aiohttp session."""
+        if self.session is None or self.session.closed:
+            self.session = aiohttp.ClientSession()
+        return self.session
+
+    async def close(self):
+        """Close the HTTP session."""
+        if self.session and not self.session.closed:
+            await self.session.close()
+
+    async def start_rollout(
+        self,
+        input: TaskInput,
+        mode: Literal["train", "val", "test"] | None = None,
+        resources_id: str | None = None,
+        metadata: Dict[str, Any] | None = None,
+    ) -> AttemptedRollout:
+        session = await self._get_session()
+        request_data = RolloutRequest(input=input, mode=mode, resources_id=resources_id, metadata=metadata)
+
+        async with session.post(f"{self.server_address}/start_rollout", json=request_data.model_dump()) as response:
+            response.raise_for_status()
+            data = await response.json()
+            return AttemptedRollout.model_validate(data)
+
+    async def enqueue_rollout(
+        self,
+        input: TaskInput,
+        mode: Literal["train", "val", "test"] | None = None,
+        resources_id: str | None = None,
+        metadata: Dict[str, Any] | None = None,
+    ) -> RolloutV2:
+        session = await self._get_session()
+        request_data = RolloutRequest(input=input, mode=mode, resources_id=resources_id, metadata=metadata)
+
+        async with session.post(f"{self.server_address}/enqueue_rollout", json=request_data.model_dump()) as response:
+            response.raise_for_status()
+            data = await response.json()
+            return RolloutV2.model_validate(data)
+
+    async def dequeue_rollout(self) -> Optional[AttemptedRollout]:
+        session = await self._get_session()
+
+        async with session.get(f"{self.server_address}/dequeue_rollout") as response:
+            response.raise_for_status()
+            data = await response.json()
+            return AttemptedRollout.model_validate(data) if data else None
+
+    async def start_attempt(self, rollout_id: str) -> AttemptedRollout:
+        session = await self._get_session()
+        request_data = RolloutId(rollout_id=rollout_id)
+
+        async with session.post(f"{self.server_address}/start_attempt", json=request_data.model_dump()) as response:
+            response.raise_for_status()
+            data = await response.json()
+            return AttemptedRollout.model_validate(data)
+
+    async def query_rollouts(
+        self, *, status: Optional[Sequence[RolloutStatus]] = None, rollout_ids: Optional[Sequence[str]] = None
+    ) -> List[RolloutV2]:
+        session = await self._get_session()
+        request_data = QueryRolloutsRequest(
+            status=list(status) if status else None, rollout_ids=list(rollout_ids) if rollout_ids else None
+        )
+
+        async with session.post(f"{self.server_address}/query_rollouts", json=request_data.model_dump()) as response:
+            response.raise_for_status()
+            data = await response.json()
+            return [RolloutV2.model_validate(item) for item in data]
+
+    async def query_attempts(self, rollout_id: str) -> List[Attempt]:
+        session = await self._get_session()
+
+        async with session.get(f"{self.server_address}/query_attempts/{rollout_id}") as response:
+            response.raise_for_status()
+            data = await response.json()
+            return [Attempt.model_validate(item) for item in data]
+
+    async def get_latest_attempt(self, rollout_id: str) -> Optional[Attempt]:
+        session = await self._get_session()
+
+        async with session.get(f"{self.server_address}/get_latest_attempt/{rollout_id}") as response:
+            response.raise_for_status()
+            data = await response.json()
+            return Attempt.model_validate(data) if data else None
+
+    async def get_rollout_by_id(self, rollout_id: str) -> Optional[RolloutV2]:
+        session = await self._get_session()
+
+        async with session.get(f"{self.server_address}/get_rollout_by_id/{rollout_id}") as response:
+            response.raise_for_status()
+            data = await response.json()
+            return RolloutV2.model_validate(data) if data else None
+
+    async def update_resources(self, resources_id: str, resources: NamedResources) -> ResourcesUpdate:
+        session = await self._get_session()
+        update = ResourcesUpdate(resources_id=resources_id, resources=resources)
+
+        async with session.post(f"{self.server_address}/update_resources", json=update.model_dump()) as response:
+            response.raise_for_status()
+            data = await response.json()
+            return ResourcesUpdate.model_validate(data)
+
+    async def get_resources_by_id(self, resources_id: str) -> Optional[ResourcesUpdate]:
+        session = await self._get_session()
+
+        async with session.get(f"{self.server_address}/get_resources_by_id/{resources_id}") as response:
+            response.raise_for_status()
+            data = await response.json()
+            return ResourcesUpdate.model_validate(data) if data else None
+
+    async def get_latest_resources(self) -> Optional[ResourcesUpdate]:
+        session = await self._get_session()
+
+        async with session.get(f"{self.server_address}/get_latest_resources") as response:
+            response.raise_for_status()
+            data = await response.json()
+            return ResourcesUpdate.model_validate(data) if data else None
+
+    async def add_span(self, span: Span) -> Span:
+        session = await self._get_session()
+
+        async with session.post(f"{self.server_address}/add_span", json=span.model_dump(mode="json")) as response:
+            response.raise_for_status()
+            data = await response.json()
+            return Span.model_validate(data)
+
+    async def get_next_span_sequence_id(self, rollout_id: str, attempt_id: str) -> int:
+        session = await self._get_session()
+
+        async with session.get(
+            f"{self.server_address}/get_next_span_sequence_id/{rollout_id}/{attempt_id}"
+        ) as response:
+            response.raise_for_status()
+            data = await response.json()
+            return data
+
+    async def add_otel_span(
+        self,
+        rollout_id: str,
+        attempt_id: str,
+        readable_span: ReadableSpan,
+        sequence_id: int | None = None,
+    ) -> Span:
+        if sequence_id is None:
+            sequence_id = await self.get_next_span_sequence_id(rollout_id, attempt_id)
+        span = Span.from_opentelemetry(
+            readable_span,
+            rollout_id=rollout_id,
+            attempt_id=attempt_id,
+            sequence_id=sequence_id,
+        )
+        await self.add_span(span)
+        return span
+
+    async def wait_for_rollouts(self, *, rollout_ids: List[str], timeout: Optional[float] = None) -> List[RolloutV2]:
+        session = await self._get_session()
+        if timeout is not None and timeout > 0.1:
+            raise ValueError(
+                "Timeout must be less than 0.1 seconds in LightningStoreClient to avoid blocking the event loop"
+            )
+        request_data = WaitForRolloutsRequest(rollout_ids=rollout_ids, timeout=timeout)
+
+        async with session.post(f"{self.server_address}/wait_for_rollouts", json=request_data.model_dump()) as response:
+            response.raise_for_status()
+            data = await response.json()
+            return [RolloutV2.model_validate(item) for item in data]
+
+    async def query_spans(
+        self,
+        rollout_id: str,
+        attempt_id: str | Literal["latest"] | None = None,
+    ) -> List[Span]:
+        session = await self._get_session()
+
+        url = f"{self.server_address}/query_spans/{rollout_id}"
+        if attempt_id is not None:
+            url += f"?attempt_id={attempt_id}"
+
+        async with session.get(url) as response:
+            response.raise_for_status()
+            data = await response.json()
+            return [Span.model_validate(item) for item in data]
+
+    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,
+    ) -> RolloutV2:
+        session = await self._get_session()
+
+        payload: Dict[str, Any] = {"rollout_id": rollout_id}
+        if not isinstance(input, Unset):
+            payload["input"] = input
+        if not isinstance(mode, Unset):
+            payload["mode"] = mode
+        if not isinstance(resources_id, Unset):
+            payload["resources_id"] = resources_id
+        if not isinstance(status, Unset):
+            payload["status"] = status
+        if not isinstance(config, Unset):
+            payload["config"] = config.model_dump()
+        if not isinstance(metadata, Unset):
+            payload["metadata"] = metadata
+
+        async with session.post(f"{self.server_address}/update_rollout", json=payload) as response:
+            response.raise_for_status()
+            data = await response.json()
+            return RolloutV2.model_validate(data)
+
+    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:
+        session = await self._get_session()
+
+        payload: Dict[str, Any] = {
+            "rollout_id": rollout_id,
+            "attempt_id": attempt_id,
+        }
+        if not isinstance(status, Unset):
+            payload["status"] = status
+        if not isinstance(worker_id, Unset):
+            payload["worker_id"] = worker_id
+        if not isinstance(last_heartbeat_time, Unset):
+            payload["last_heartbeat_time"] = last_heartbeat_time
+        if not isinstance(metadata, Unset):
+            payload["metadata"] = metadata
+
+        async with session.post(f"{self.server_address}/update_attempt", json=payload) as response:
+            response.raise_for_status()
+            data = await response.json()
+            return Attempt.model_validate(data)

agentlightning/store/memory.py

@@ -0,0 +1,634 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+import asyncio
+import functools
+import logging
+import time
+import uuid
+from collections import deque
+from typing import Any, Callable, Counter, Dict, List, Literal, Optional, Sequence, TypeVar, cast
+
+from opentelemetry.sdk.trace import ReadableSpan
+
+from agentlightning.tracer import Span
+from agentlightning.types import (
+    Attempt,
+    AttemptedRollout,
+    AttemptStatus,
+    NamedResources,
+    ResourcesUpdate,
+    RolloutConfig,
+    RolloutStatus,
+    RolloutV2,
+    TaskInput,
+)
+
+from .base import UNSET, LightningStore, Unset, is_finished, is_queuing
+from .utils import healthcheck, propagate_status
+
+T_callable = TypeVar("T_callable", bound=Callable[..., Any])
+
+logger = logging.getLogger(__name__)
+
+
+def _healthcheck_wrapper(func: T_callable) -> T_callable:
+    """
+    Decorator to run the watchdog healthcheck **before** executing the decorated method.
+    Only runs if the store has a watchdog configured.
+    Prevents recursive healthcheck execution using a flag on the store instance.
+    """
+
+    @functools.wraps(func)
+    async def wrapper(self: InMemoryLightningStore, *args: Any, **kwargs: Any) -> Any:
+        # Check if healthcheck is already running to prevent recursion
+        if getattr(self, "_healthcheck_running", False):
+            # Skip healthcheck if already running
+            return await func(self, *args, **kwargs)
+
+        # Set flag to prevent recursive healthcheck calls
+        # This flag is not asyncio/thread-safe, but it doesn't matter
+        self._healthcheck_running = True  # type: ignore
+        try:
+            # The following methods should live inside one lock.
+            await self._healthcheck()  # pyright: ignore[reportPrivateUsage]
+        finally:
+            # Always clear the flag, even if healthcheck fails
+            self._healthcheck_running = False  # type: ignore
+
+        # Execute the original method
+        # This should be outside the lock.
+        return await func(self, *args, **kwargs)
+
+    return cast(T_callable, wrapper)
+
+
+class InMemoryLightningStore(LightningStore):
+    """
+    In-memory implementation of LightningStore using Python data structures.
+    Thread-safe and async-compatible but data is not persistent.
+
+    The methods in this class should generally not call each other,
+    especially those that are locked.
+    """
+
+    def __init__(self):
+        self._lock = asyncio.Lock()
+
+        # Task queue and rollouts storage
+        self._task_queue: deque[RolloutV2] = deque()
+        self._rollouts: Dict[str, RolloutV2] = {}
+
+        # Resources storage (similar to legacy server.py)
+        self._resources: Dict[str, ResourcesUpdate] = {}
+        self._latest_resources_id: Optional[str] = None
+
+        # Spans storage
+        self._spans: Dict[str, List[Span]] = {}  # rollout_id -> list of spans
+        self._span_sequence_ids: Dict[str, int] = Counter()  # rollout_id -> sequence_id
+
+        # Attempt tracking
+        self._attempts: Dict[str, List[Attempt]] = {}  # rollout_id -> list of attempts
+
+        # Completion tracking for wait_for_rollouts
+        self._completion_events: Dict[str, asyncio.Event] = {}
+
+    @_healthcheck_wrapper
+    async def start_rollout(
+        self,
+        input: TaskInput,
+        mode: Literal["train", "val", "test"] | None = None,
+        resources_id: str | None = None,
+        metadata: Dict[str, Any] | None = None,
+    ) -> AttemptedRollout:
+        """
+        Notify the store that I'm about to run a rollout.
+        """
+        async with self._lock:
+            rollout_id = f"rollout-{uuid.uuid4()}"
+            current_time = time.time()
+
+            rollout = RolloutV2(
+                rollout_id=rollout_id,
+                input=input,
+                mode=mode,
+                resources_id=resources_id or self._latest_resources_id,
+                start_time=current_time,
+                status="preparing",
+                metadata=metadata or {},
+            )
+
+            # Create the initial attempt
+            attempt_id = f"attempt-{uuid.uuid4()}"
+            attempt = Attempt(
+                rollout_id=rollout.rollout_id,
+                attempt_id=attempt_id,
+                sequence_id=1,
+                start_time=current_time,
+                status="preparing",
+            )
+
+            self._attempts[rollout.rollout_id] = [attempt]
+            self._rollouts[rollout.rollout_id] = rollout
+
+            # Manully added rollout is not added to task queue. It's already preparing
+            self._completion_events.setdefault(rollout.rollout_id, asyncio.Event())
+
+            return AttemptedRollout(**rollout.model_dump(), attempt=attempt)
+
+    @_healthcheck_wrapper
+    async def enqueue_rollout(
+        self,
+        input: TaskInput,
+        mode: Literal["train", "val", "test"] | None = None,
+        resources_id: str | None = None,
+        metadata: Dict[str, Any] | None = None,
+    ) -> RolloutV2:
+        """
+        Adds a new task to the queue with specific metadata and returns its unique ID.
+        """
+        async with self._lock:
+            rollout_id = f"rollout-{uuid.uuid4()}"
+            current_time = time.time()
+
+            rollout = RolloutV2(
+                rollout_id=rollout_id,
+                input=input,
+                mode=mode,
+                resources_id=resources_id or self._latest_resources_id,
+                start_time=current_time,
+                status="queuing",  # should be queuing
+                metadata=metadata or {},
+            )
+
+            self._rollouts[rollout.rollout_id] = rollout
+            self._task_queue.append(rollout)  # add it to the end of the queue
+            self._completion_events.setdefault(rollout.rollout_id, asyncio.Event())
+
+            return rollout
+
+    @_healthcheck_wrapper
+    async def dequeue_rollout(self) -> Optional[AttemptedRollout]:
+        """
+        Retrieves the next task from the queue without blocking.
+        Returns None if the queue is empty.
+
+        Will set the rollout status to preparing and create a new attempt.
+        """
+        async with self._lock:
+            # Keep looking until we find a rollout that's still in queuing status
+            # or the queue is empty
+            while self._task_queue:
+                rollout = self._task_queue.popleft()
+
+                # Check if rollout is still in a queuing state
+                # (it might have been updated to a different status while in queue)
+                if is_queuing(rollout):
+                    # Update status to preparing
+                    rollout.status = "preparing"
+
+                    # Create a new attempt (could be first attempt or retry)
+                    attempt_id = f"attempt-{uuid.uuid4()}"
+                    current_time = time.time()
+
+                    # Get existing attempts to determine sequence number
+                    existing_attempts = self._attempts.get(rollout.rollout_id, [])
+                    sequence_id = len(existing_attempts) + 1
+
+                    attempt = Attempt(
+                        rollout_id=rollout.rollout_id,
+                        attempt_id=attempt_id,
+                        sequence_id=sequence_id,
+                        start_time=current_time,
+                        status="preparing",
+                    )
+
+                    if rollout.rollout_id not in self._attempts:
+                        self._attempts[rollout.rollout_id] = []
+                    self._attempts[rollout.rollout_id].append(attempt)
+
+                    return AttemptedRollout(**rollout.model_dump(), attempt=attempt)
+
+                # If not in queuing state, skip this rollout and continue
+                # (it was updated externally and should not be processed)
+
+            # No valid rollouts found
+            return None
+
+    @_healthcheck_wrapper
+    async def start_attempt(self, rollout_id: str) -> AttemptedRollout:
+        """
+        Create a new attempt for a given rollout ID and return the attempt details.
+        """
+        async with self._lock:
+            # Get the rollout
+            rollout = self._rollouts.get(rollout_id)
+            if not rollout:
+                raise ValueError(f"Rollout {rollout_id} not found")
+
+            # Get existing attempts to determine sequence number
+            existing_attempts = self._attempts.get(rollout_id, [])
+            sequence_id = len(existing_attempts) + 1
+
+            # We don't care whether the max attempts have reached or not
+            # This attempt is from user trigger
+
+            # Create new attempt
+            attempt_id = f"attempt-{uuid.uuid4()}"
+            current_time = time.time()
+
+            attempt = Attempt(
+                rollout_id=rollout_id,
+                attempt_id=attempt_id,
+                sequence_id=sequence_id,
+                start_time=current_time,
+                status="preparing",
+            )
+
+            # Add attempt to storage
+            if rollout_id not in self._attempts:
+                self._attempts[rollout_id] = []
+            self._attempts[rollout_id].append(attempt)
+
+            self._completion_events.setdefault(rollout.rollout_id, asyncio.Event())
+
+            return AttemptedRollout(**rollout.model_dump(), attempt=attempt)
+
+    @_healthcheck_wrapper
+    async def query_rollouts(
+        self, *, status: Optional[Sequence[RolloutStatus]] = None, rollout_ids: Optional[Sequence[str]] = None
+    ) -> List[RolloutV2]:
+        """
+        Query and retrieve rollouts filtered by their status and rollout ids.
+        If no status is provided, returns all rollouts.
+        """
+        async with self._lock:
+            rollouts = list(self._rollouts.values())
+
+            # Filter by rollout_ids if provided
+            if rollout_ids is not None:
+                rollout_ids_set = set(rollout_ids)
+                rollouts = [rollout for rollout in rollouts if rollout.rollout_id in rollout_ids_set]
+
+            # Filter by status if provided
+            if status is not None:
+                status_set = set(status)
+                rollouts = [rollout for rollout in rollouts if rollout.status in status_set]
+
+            return rollouts
+
+    @_healthcheck_wrapper
+    async def get_rollout_by_id(self, rollout_id: str) -> Optional[RolloutV2]:
+        """
+        Safely retrieves a specific rollout by its ID.
+        """
+        async with self._lock:
+            return self._rollouts.get(rollout_id)
+
+    @_healthcheck_wrapper
+    async def query_attempts(self, rollout_id: str) -> List[Attempt]:
+        """
+        Query and retrieve all attempts associated with a specific rollout ID.
+        Returns an empty list if no attempts are found.
+        """
+        async with self._lock:
+            return self._attempts.get(rollout_id, [])
+
+    @_healthcheck_wrapper
+    async def get_latest_attempt(self, rollout_id: str) -> Optional[Attempt]:
+        """
+        Safely retrieves the latest attempt for a given rollout ID.
+        """
+        async with self._lock:
+            attempts = self._attempts.get(rollout_id, [])
+            if not attempts:
+                return None
+            return max(attempts, key=lambda a: a.sequence_id)
+
+    @_healthcheck_wrapper
+    async def update_resources(self, resources_id: str, resources: NamedResources) -> ResourcesUpdate:
+        """
+        Safely stores a new version of named resources and sets it as the latest.
+        """
+        async with self._lock:
+            update = ResourcesUpdate(resources_id=resources_id, resources=resources)
+            self._resources[resources_id] = update
+            self._latest_resources_id = resources_id
+            return update
+
+    @_healthcheck_wrapper
+    async def get_resources_by_id(self, resources_id: str) -> Optional[ResourcesUpdate]:
+        """
+        Safely retrieves a specific version of named resources by its ID.
+        """
+        async with self._lock:
+            return self._resources.get(resources_id)
+
+    @_healthcheck_wrapper
+    async def get_latest_resources(self) -> Optional[ResourcesUpdate]:
+        """
+        Safely retrieves the latest version of named resources.
+        """
+        async with self._lock:
+            if self._latest_resources_id:
+                return self._resources.get(self._latest_resources_id)
+            return None
+
+    async def get_next_span_sequence_id(self, rollout_id: str, attempt_id: str) -> int:
+        """
+        Get the next span sequence ID for a given rollout and attempt.
+        The number is strictly increasing for each rollout.
+        The store will not issue the same sequence ID twice.
+        """
+        async with self._lock:
+            self._span_sequence_ids[rollout_id] += 1
+            return self._span_sequence_ids[rollout_id]
+
+    async def add_span(self, span: Span) -> Span:
+        """Persist a pre-converted span."""
+        async with self._lock:
+            self._span_sequence_ids[span.rollout_id] = max(self._span_sequence_ids[span.rollout_id], span.sequence_id)
+            return await self._add_span_unlocked(span)
+
+    async def add_otel_span(
+        self, rollout_id: str, attempt_id: str, readable_span: ReadableSpan, sequence_id: int | None = None
+    ) -> Span:
+        """Add an opentelemetry span to the store."""
+        async with self._lock:
+            if sequence_id is None:
+                # Issue a new sequence ID for the rollout
+                self._span_sequence_ids[rollout_id] += 1
+                sequence_id = self._span_sequence_ids[rollout_id]
+            else:
+                # Comes from a provided sequence ID
+                # Make sure our counter is strictly increasing
+                self._span_sequence_ids[rollout_id] = max(self._span_sequence_ids[rollout_id], sequence_id)
+
+            span = Span.from_opentelemetry(
+                readable_span, rollout_id=rollout_id, attempt_id=attempt_id, sequence_id=sequence_id
+            )
+            await self._add_span_unlocked(span)
+            return span
+
+    async def _add_span_unlocked(self, span: Span) -> Span:
+        rollout = self._rollouts.get(span.rollout_id)
+        if not rollout:
+            raise ValueError(f"Rollout {span.rollout_id} not found")
+        attempts = self._attempts.get(span.rollout_id, [])
+        current_attempt = next((a for a in attempts if a.attempt_id == span.attempt_id), None)
+        latest_attempt = max(attempts, key=lambda a: a.sequence_id) if attempts else None
+        if not current_attempt:
+            raise ValueError(f"Attempt {span.attempt_id} not found for rollout {span.rollout_id}")
+        if not latest_attempt:
+            raise ValueError(f"No attempts found for rollout {span.rollout_id}")
+
+        if span.rollout_id not in self._spans:
+            self._spans[span.rollout_id] = []
+        self._spans[span.rollout_id].append(span)
+
+        # Update attempt heartbeat
+        current_attempt.last_heartbeat_time = time.time()
+        if current_attempt.status == "preparing":
+            current_attempt.status = "running"
+
+        # If the status has already timed out or failed, do not change it
+
+        # Update rollout status if it's the latest attempt
+        if rollout.status == "preparing" and current_attempt == latest_attempt:
+            rollout.status = "running"
+
+        return span
+
+    @_healthcheck_wrapper
+    async def wait_for_rollouts(self, *, rollout_ids: List[str], timeout: Optional[float] = None) -> List[RolloutV2]:
+        """
+        Wait for specified rollouts to complete with a timeout.
+        Returns the completed rollouts, potentially incomplete if timeout is reached.
+
+        This method does not change the state of the store.
+        """
+        completed_rollouts: List[RolloutV2] = []
+
+        async def wait_for_rollout(rollout_id: str):
+            # First check if already completed
+            async with self._lock:
+                rollout = self._rollouts.get(rollout_id)
+                if rollout and is_finished(rollout):
+                    completed_rollouts.append(rollout)
+                    return
+
+            # If not completed and we have an event, wait for completion
+            if rollout_id in self._completion_events:
+                try:
+                    await asyncio.wait_for(self._completion_events[rollout_id].wait(), timeout=timeout)
+                    async with self._lock:
+                        rollout = self._rollouts.get(rollout_id)
+                        if rollout and is_finished(rollout):
+                            completed_rollouts.append(rollout)
+                except asyncio.TimeoutError:
+                    pass
+
+        # Wait for all rollouts concurrently
+        await asyncio.gather(*[wait_for_rollout(rid) for rid in rollout_ids], return_exceptions=True)
+
+        return completed_rollouts
+
+    @_healthcheck_wrapper
+    async def query_spans(self, rollout_id: str, attempt_id: str | Literal["latest"] | None = None) -> List[Span]:
+        """
+        Query and retrieve all spans associated with a specific rollout ID.
+        Returns an empty list if no spans are found.
+        """
+        async with self._lock:
+            spans = self._spans.get(rollout_id, [])
+            if attempt_id is None:
+                return spans
+            elif attempt_id == "latest":
+                # Find the latest attempt_id
+                if not spans:
+                    return []
+                latest_attempt = max(spans, key=lambda s: s.sequence_id if s.attempt_id else "").attempt_id
+                return [s for s in spans if s.attempt_id == latest_attempt]
+            else:
+                return [s for s in spans if s.attempt_id == attempt_id]
+
+    @_healthcheck_wrapper
+    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,
+    ) -> RolloutV2:
+        """
+        Update the rollout status and related metadata.
+        """
+        async with self._lock:
+            return await self._update_rollout_unlocked(
+                rollout_id=rollout_id,
+                input=input,
+                mode=mode,
+                resources_id=resources_id,
+                status=status,
+                config=config,
+                metadata=metadata,
+            )
+
+    @_healthcheck_wrapper
+    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:
+        """
+        Update a specific or latest attempt for a given rollout.
+        """
+        async with self._lock:
+            attempt = await self._update_attempt_unlocked(
+                rollout_id=rollout_id,
+                attempt_id=attempt_id,
+                status=status,
+                worker_id=worker_id,
+                last_heartbeat_time=last_heartbeat_time,
+                metadata=metadata,
+            )
+
+        return attempt
+
+    async def _update_rollout_unlocked(
+        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,
+    ) -> RolloutV2:
+        # No lock inside this one.
+        rollout = self._rollouts.get(rollout_id)
+        if not rollout:
+            raise ValueError(f"Rollout {rollout_id} not found")
+
+        # Update fields if they are not UNSET
+        if not isinstance(input, Unset):
+            rollout.input = input
+        if not isinstance(mode, Unset):
+            rollout.mode = mode
+        if not isinstance(resources_id, Unset):
+            rollout.resources_id = resources_id
+        if not isinstance(status, Unset):
+            rollout.status = status
+        if not isinstance(config, Unset):
+            rollout.config = config
+        if not isinstance(metadata, Unset):
+            rollout.metadata = metadata
+
+        # Set end time for finished rollouts
+        # Rollout is only finished when it succeeded or fail with no more retries.
+        if status is not UNSET and is_finished(rollout):
+            rollout.end_time = time.time()
+            # Signal completion
+            if rollout_id in self._completion_events:
+                self._completion_events[rollout_id].set()
+
+        # If requeuing, add back to queue
+        elif is_queuing(rollout) and rollout not in self._task_queue:
+            self._task_queue.append(rollout)
+
+        # Re-validate the rollout to ensure legality
+        RolloutV2.model_validate(rollout.model_dump())
+
+        return rollout
+
+    async def _update_attempt_unlocked(
+        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:
+        # No lock, but with status propagation.
+        rollout = self._rollouts.get(rollout_id)
+        if not rollout:
+            raise ValueError(f"Rollout {rollout_id} not found")
+
+        attempts = self._attempts.get(rollout_id, [])
+        if not attempts:
+            raise ValueError(f"No attempts found for rollout {rollout_id}")
+
+        latest_attempt = max(attempts, key=lambda a: a.sequence_id)
+
+        # Find the attempt to update
+        if attempt_id == "latest":
+            attempt = latest_attempt
+        else:
+            attempt = next((a for a in attempts if a.attempt_id == attempt_id), None)
+            if not attempt:
+                raise ValueError(f"Attempt {attempt_id} not found for rollout {rollout_id}")
+
+        # Update fields if they are not UNSET
+        if not isinstance(status, Unset):
+            attempt.status = status
+            # Also update end_time if the status indicates completion
+            if status in ["failed", "succeeded"]:
+                attempt.end_time = time.time()
+        if not isinstance(worker_id, Unset):
+            attempt.worker_id = worker_id
+        if not isinstance(last_heartbeat_time, Unset):
+            attempt.last_heartbeat_time = last_heartbeat_time
+        if not isinstance(metadata, Unset):
+            attempt.metadata = metadata
+
+        # Re-validate the attempt to ensure legality
+        Attempt.model_validate(attempt.model_dump())
+
+        if attempt == latest_attempt:
+
+            async def _update_status(rollout_id: str, status: RolloutStatus) -> RolloutV2:
+                return await self._update_rollout_unlocked(rollout_id, status=status)
+
+            # Propagate the status to the rollout
+            await propagate_status(
+                _update_status,
+                attempt,
+                rollout.config,
+            )
+
+        return attempt
+
+    async def _healthcheck(self) -> None:
+        """Perform healthcheck against all running rollouts in the store."""
+        async with self._lock:
+            running_rollouts: List[AttemptedRollout] = []
+            for rollout in self._rollouts.values():
+                if rollout.status in ["preparing", "running"]:
+                    all_attempts = self._attempts.get(rollout.rollout_id, [])
+                    if not all_attempts:
+                        # The rollout is running but has no attempts, this should not happen
+                        logger.error(f"Rollout {rollout.rollout_id} is running but has no attempts")
+                        continue
+                    latest_attempt = max(all_attempts, key=lambda a: a.sequence_id)
+                    running_rollouts.append(AttemptedRollout(**rollout.model_dump(), attempt=latest_attempt))
+
+            async def _update_attempt_status(rollout_id: str, attempt_id: str, status: AttemptStatus) -> Attempt:
+                return await self._update_attempt_unlocked(rollout_id, attempt_id, status=status)
+
+            async def _update_rollout_status(rollout_id: str, status: RolloutStatus) -> RolloutV2:
+                return await self._update_rollout_unlocked(rollout_id, status=status)
+
+            await healthcheck(
+                running_rollouts,
+                _update_rollout_status,
+                _update_attempt_status,
+            )

agentlightning/store/sqlite.py

@@ -0,0 +1 @@
+# Copyright (c) Microsoft. All rights reserved.

agentlightning/store/threading.py

@@ -0,0 +1,167 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+import threading
+from typing import Any, Dict, List, Literal, Optional, Sequence
+
+from opentelemetry.sdk.trace import ReadableSpan
+
+from agentlightning.tracer import Span
+from agentlightning.types import (
+    Attempt,
+    AttemptedRollout,
+    AttemptStatus,
+    NamedResources,
+    ResourcesUpdate,
+    RolloutConfig,
+    RolloutStatus,
+    RolloutV2,
+    TaskInput,
+)
+
+from .base import UNSET, LightningStore, 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()
+
+    async def start_rollout(
+        self,
+        input: TaskInput,
+        mode: Literal["train", "val", "test"] | None = None,
+        resources_id: str | None = None,
+        metadata: Dict[str, Any] | None = None,
+    ) -> AttemptedRollout:
+        with self._lock:
+            return await self.store.start_rollout(input, mode, resources_id, metadata)
+
+    async def enqueue_rollout(
+        self,
+        input: TaskInput,
+        mode: Literal["train", "val", "test"] | None = None,
+        resources_id: str | None = None,
+        metadata: Dict[str, Any] | None = None,
+    ) -> RolloutV2:
+        with self._lock:
+            return await self.store.enqueue_rollout(input, mode, resources_id, metadata)
+
+    async def dequeue_rollout(self) -> Optional[AttemptedRollout]:
+        with self._lock:
+            return await self.store.dequeue_rollout()
+
+    async def start_attempt(self, rollout_id: str) -> AttemptedRollout:
+        with self._lock:
+            return await self.store.start_attempt(rollout_id)
+
+    async def query_rollouts(
+        self,
+        *,
+        status: Optional[Sequence[RolloutStatus]] = None,
+        rollout_ids: Optional[Sequence[str]] = None,
+    ) -> List[RolloutV2]:
+        with self._lock:
+            return await self.store.query_rollouts(status=status, rollout_ids=rollout_ids)
+
+    async def query_attempts(self, rollout_id: str) -> List[Attempt]:
+        with self._lock:
+            return await self.store.query_attempts(rollout_id)
+
+    async def get_rollout_by_id(self, rollout_id: str) -> Optional[RolloutV2]:
+        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 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_span(self, span: Span) -> 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,
+    ) -> 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[RolloutV2]:
+        # 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 query_spans(
+        self,
+        rollout_id: str,
+        attempt_id: str | Literal["latest"] | None = None,
+    ) -> List[Span]:
+        with self._lock:
+            return await self.store.query_spans(rollout_id, attempt_id)
+
+    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,
+    ) -> RolloutV2:
+        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,
+            )

agentlightning/store/utils.py

@@ -0,0 +1,126 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import time
+from typing import Awaitable, Callable, List, cast
+
+from agentlightning.types import Attempt, AttemptedRollout, AttemptStatus, RolloutConfig, RolloutStatus, RolloutV2
+
+UpdateRolloutStatus = Callable[[str, RolloutStatus], Awaitable[RolloutV2]]
+UpdateAttemptStatus = Callable[[str, str, AttemptStatus], Awaitable[Attempt]]
+
+
+async def propagate_status(
+    update_rollout_status: UpdateRolloutStatus,  # this should be unlocked
+    attempt: Attempt,
+    config: RolloutConfig,
+) -> RolloutV2:
+    """
+    Propagate the status of an attempt to the rollout.
+
+    The rollout should be made sure in a state to be outdated.
+    Requeue the rollout if it should be retried.
+
+    This operation is completely unlocked. The caller is responsible for locking the store.
+    """
+    # Propagate the status directly to the rollout
+    if attempt.status == "preparing" or attempt.status == "running" or attempt.status == "succeeded":
+        return await update_rollout_status(
+            attempt.rollout_id,
+            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 await update_rollout_status(
+                    attempt.rollout_id,
+                    "requeuing",
+                )
+
+        # If we can't retry or shouldn't retry, mark as failed
+        return await update_rollout_status(
+            attempt.rollout_id,
+            "failed",
+        )
+
+    raise ValueError(f"Invalid attempt status: {attempt.status}")
+
+
+async def healthcheck(
+    rollouts: List[AttemptedRollout],
+    update_rollout_status: UpdateRolloutStatus,
+    update_attempt_status: UpdateAttemptStatus,
+) -> None:
+    """
+    Perform health check on all running rollouts in the store.
+
+    This method should be called periodically to:
+    1. Update rollout status to failed to succeeded when the attempt is done
+    2. Check for unresponsive attempts (no heartbeat or spans for a while)
+    3. Check for timed-out rollouts (running too long since start_time)
+    4. Update attempt/rollout status accordingly
+
+    This operation is completely unlocked. The caller is responsible for locking the store.
+
+    Args:
+        store: The LightningStore instance to check rollouts from
+    """
+    current_time = time.time()
+
+    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:
+            continue
+
+        # Check if the attempt has already failed or succeeded
+        if latest_attempt.status == "failed" or latest_attempt.status == "succeeded":
+            await propagate_status(update_rollout_status, latest_attempt, config)
+            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:
+            await update_attempt_status(
+                latest_attempt.rollout_id,
+                latest_attempt.attempt_id,
+                "timeout",
+            )
+            continue
+
+        # Check for unresponsive condition (based on last heartbeat)
+        if latest_attempt.last_heartbeat_time:
+            if latest_attempt.status == "preparing":
+                # If still preparing, mark it as running
+                latest_attempt = await update_attempt_status(
+                    latest_attempt.rollout_id,
+                    latest_attempt.attempt_id,
+                    "running",
+                )
+
+            # Haven't received heartbeat for a while
+            if (
+                config.unresponsive_seconds is not None
+                and current_time - cast(float, latest_attempt.last_heartbeat_time) > config.unresponsive_seconds
+            ):
+                await update_attempt_status(
+                    latest_attempt.rollout_id,
+                    latest_attempt.attempt_id,
+                    "unresponsive",
+                )
+                continue
+
+        # 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
+        ):
+            await update_attempt_status(
+                latest_attempt.rollout_id,
+                latest_attempt.attempt_id,
+                "unresponsive",
+            )

agentlightning/tracer/__init__.py

@@ -2,5 +2,6 @@
 
 from .agentops import AgentOpsTracer
 from .base import BaseTracer
+from .types import Span
 
-__all__ = ["AgentOpsTracer", "BaseTracer"]
+__all__ = ["AgentOpsTracer", "BaseTracer", "Span"]

agentlightning/tracer/types.py

@@ -0,0 +1,231 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+import json
+from typing import Any, Dict, List, Optional, Sequence, Union
+
+from opentelemetry import trace as trace_api
+from opentelemetry.sdk.resources import Resource as OtelResource
+from opentelemetry.sdk.trace import Event as OtelEvent
+from opentelemetry.sdk.trace import ReadableSpan
+from opentelemetry.trace.status import Status as OtelStatus
+from pydantic import BaseModel
+
+
+def convert_timestamp(timestamp: Optional[int]) -> Optional[float]:
+    """Convert timestamp from nanoseconds to seconds if needed.
+
+    Auto-detects format: if > 1e12, assumes nanoseconds; otherwise seconds.
+    """
+    if not timestamp:
+        return None
+    return timestamp / 1_000_000_000 if timestamp > 1e12 else timestamp
+
+
+def extract_extra_fields(src: Any, excluded_fields: List[str]) -> Dict[str, Any]:
+    """Extract extra fields from source object, excluding specified fields and private fields."""
+    excluded_fields_set = set(excluded_fields) | set(["_" + k for k in excluded_fields])
+    # Exclude the function fields
+    excluded_fields_set |= set(src.__class__.__dict__.keys())
+    stripped_dict = {k.lstrip("_"): v for k, v in src.__dict__.items()}
+    candidates = {k: v for k, v in stripped_dict.items() if k not in excluded_fields_set and not k.startswith("_")}
+    # This should strip or flatten the unserializable fields
+    candidates_serialized = json.dumps(candidates, default=str)
+    return json.loads(candidates_serialized)
+
+
+AttributeValue = Union[
+    str,
+    bool,
+    int,
+    float,
+    Sequence[str],
+    Sequence[bool],
+    Sequence[int],
+    Sequence[float],
+]
+Attributes = Dict[str, AttributeValue]
+TraceState = Dict[str, str]
+
+
+class SpanContext(BaseModel):
+    """Corresponding to opentelemetry.trace.SpanContext"""
+
+    trace_id: str
+    span_id: str
+    is_remote: bool
+    trace_state: TraceState
+
+    class Config:
+        allow_extra = True
+
+    @classmethod
+    def from_opentelemetry(cls, src: trace_api.SpanContext) -> "SpanContext":
+        return cls(
+            trace_id=trace_api.format_trace_id(src.trace_id),
+            span_id=trace_api.format_span_id(src.span_id),
+            is_remote=src.is_remote,
+            trace_state={k: v for k, v in src.trace_state.items()} if src.trace_state else {},
+            **extract_extra_fields(src, ["trace_id", "span_id", "is_remote", "trace_state"]),
+        )
+
+
+class TraceStatus(BaseModel):
+    """Corresponding to opentelemetry.trace.Status"""
+
+    status_code: str
+    description: Optional[str] = None
+
+    class Config:
+        allow_extra = True
+
+    @classmethod
+    def from_opentelemetry(cls, src: OtelStatus) -> "TraceStatus":
+        return cls(
+            status_code=src.status_code.name,
+            description=src.description,
+            **extract_extra_fields(src, ["status_code", "description"]),
+        )
+
+
+class Event(BaseModel):
+    """Corresponding to opentelemetry.trace.Event"""
+
+    name: str
+    attributes: Attributes
+    timestamp: Optional[float] = None
+
+    class Config:
+        allow_extra = True
+
+    @classmethod
+    def from_opentelemetry(cls, src: OtelEvent) -> "Event":
+        return cls(
+            name=src.name,
+            attributes=dict(src.attributes) if src.attributes else {},
+            timestamp=convert_timestamp(src.timestamp),
+            **extract_extra_fields(src, ["name", "attributes", "timestamp"]),
+        )
+
+
+class Link(BaseModel):
+    """Corresponding to opentelemetry.trace.Link"""
+
+    context: SpanContext
+    attributes: Optional[Attributes] = None
+
+    class Config:
+        allow_extra = True
+
+    @classmethod
+    def from_opentelemetry(cls, src: trace_api.Link) -> "Link":
+        return cls(
+            context=SpanContext.from_opentelemetry(src.context),
+            attributes=dict(src.attributes) if src.attributes else None,
+            **extract_extra_fields(src, ["context", "attributes"]),
+        )
+
+
+class Resource(BaseModel):
+    """Corresponding to opentelemetry.sdk.resources.Resource"""
+
+    attributes: Attributes
+    schema_url: str
+
+    @classmethod
+    def from_opentelemetry(cls, src: OtelResource) -> "Resource":
+        return cls(
+            attributes=dict(src.attributes) if src.attributes else {},
+            schema_url=src.schema_url if src.schema_url else "",
+            **extract_extra_fields(src, ["attributes", "schema_url"]),
+        )
+
+
+class Span(BaseModel):
+
+    class Config:
+        allow_extra = True  # allow extra fields if needed
+
+    rollout_id: str
+    attempt_id: str
+    # The ID to make spans ordered within a single attempt
+    sequence_id: int
+
+    # Current ID (in hex, formatted via trace_api.format_*)
+    trace_id: str  # one rollout can have traces coming from multiple places
+    span_id: str
+    parent_id: Optional[str]
+
+    # Core ReadableSpan fields
+    name: str
+    status: TraceStatus
+    attributes: Attributes
+    events: List[Event]
+    links: List[Link]
+
+    # Timestamps
+    start_time: Optional[float]
+    end_time: Optional[float]
+
+    # Other parsable fields
+    context: Optional[SpanContext]
+    parent: Optional[SpanContext]
+    resource: Resource
+
+    # Preserve other fields in the readable span as extra fields
+    # Make sure that are json serializable (so no bytes, complex objects, ...)
+
+    @classmethod
+    def from_opentelemetry(
+        cls,
+        src: ReadableSpan,
+        rollout_id: str,
+        attempt_id: str,
+        sequence_id: int,
+    ) -> "Span":
+        context = src.get_span_context()
+        if context is None:
+            trace_id = span_id = 0
+        else:
+            trace_id = context.trace_id
+            span_id = context.span_id
+        return cls(
+            rollout_id=rollout_id,
+            attempt_id=attempt_id,
+            sequence_id=sequence_id,
+            trace_id=trace_api.format_trace_id(trace_id),
+            span_id=trace_api.format_span_id(span_id),
+            parent_id=(trace_api.format_span_id(src.parent.span_id) if src.parent else None),
+            name=src.name,
+            status=TraceStatus.from_opentelemetry(src.status),
+            attributes=dict(src.attributes) if src.attributes else {},
+            events=[Event.from_opentelemetry(event) for event in src.events] if src.events else [],
+            links=[Link.from_opentelemetry(link) for link in src.links] if src.links else [],
+            start_time=convert_timestamp(src.start_time),
+            end_time=convert_timestamp(src.end_time),
+            context=SpanContext.from_opentelemetry(context) if context else None,
+            parent=(SpanContext.from_opentelemetry(src.parent) if src.parent else None),
+            resource=Resource.from_opentelemetry(src.resource),
+            **extract_extra_fields(
+                src,
+                [
+                    "name",
+                    "context",
+                    "parent",
+                    "resource",
+                    "attributes",
+                    "events",
+                    "links",
+                    "start_time",
+                    "end_time",
+                    "status",
+                    "span_processor",
+                    "rollout_id",
+                    "attempt_id",
+                    "trace_id",
+                    "span_id",
+                    "parent_id",
+                ],
+            ),
+        )

agentlightning/types.py

@@ -2,10 +2,10 @@
 
 from __future__ import annotations
 
-from typing import Annotated, Any, Dict, Generic, List, Literal, Optional, Protocol, TypeVar, Union
+from typing import Annotated, Any, Callable, Dict, Generic, List, Literal, Optional, Protocol, TypeVar, Union, cast
 
 from opentelemetry.sdk.trace import ReadableSpan
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, Field, model_validator
 
 __all__ = [
     "Triplet",
@@ -23,6 +23,12 @@ __all__ = [
     "GenericResponse",
     "ParallelWorkerBase",
     "Dataset",
+    "AttemptStatus",
+    "RolloutStatus",
+    "RolloutConfig",
+    "RolloutV2",
+    "Attempt",
+    "AttemptedRollout",
 ]
 
 T_co = TypeVar("T_co", covariant=True)
@@ -64,6 +70,92 @@ class Rollout(BaseModel):
     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
+]
+
+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
+]
+
+
+class Attempt(BaseModel):
+    """An attempt to execute a rollout. A rollout can have multiple attempts if retries are needed."""
+
+    rollout_id: str  # the rollout 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  # time when the attempt has started
+    end_time: Optional[float] = None  # time when the attempt has ended
+
+    status: AttemptStatus = "preparing"
+    # The rollout worker which is executing this attempt
+    worker_id: Optional[str] = None
+
+    last_heartbeat_time: Optional[float] = None  # last time when the worker has reported progress
+
+    # A bucket for any other relevant information
+    metadata: Optional[Dict[str, Any]] = None
+
+
+class RolloutConfig(BaseModel):
+    """Configurations for rollout execution."""
+
+    timeout_seconds: Optional[float] = None  # none indicates no timeout
+    unresponsive_seconds: Optional[float] = None  # none indicates no unresponsive timeout
+    max_attempts: int = Field(default=1, ge=1)  # including the first attempt
+    retry_condition: List[AttemptStatus] = Field(
+        default_factory=cast(Callable[[], List[AttemptStatus]], list)
+    )  # list of statuses that should trigger a retry
+
+
+class RolloutV2(BaseModel):
+    rollout_id: str
+
+    # Inputs
+    input: TaskInput
+
+    # Time to track the lifecycle of the rollout
+    start_time: float
+    end_time: Optional[float] = None
+
+    mode: Optional[Literal["train", "val", "test"]] = None
+    resources_id: Optional[str] = None
+
+    # Overall scheduling/running information
+    status: RolloutStatus = "queuing"
+
+    config: RolloutConfig = Field(default_factory=RolloutConfig)
+
+    # A bucket for any other relevant information
+    metadata: Optional[Dict[str, Any]] = None
+
+
+class AttemptedRollout(RolloutV2):
+    """A rollout along with its active attempt."""
+
+    attempt: Attempt
+
+    @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
+
+
 TaskInput = Any
 
 
@@ -115,6 +207,7 @@ class LLM(Resource):
     resource_type: Literal["llm"] = "llm"
     endpoint: str
     model: str
+    api_key: Optional[str] = None
     sampling_parameters: Dict[str, Any] = Field(default_factory=dict)
 
 

tests/store/conftest.py

@@ -0,0 +1,50 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import time
+from unittest.mock import Mock
+
+import pytest
+from opentelemetry.sdk.trace import ReadableSpan
+
+from agentlightning.store.memory import InMemoryLightningStore
+
+__all__ = [
+    "store",
+    "mock_readable_span",
+]
+
+
+@pytest.fixture
+def store() -> InMemoryLightningStore:
+    """Create a fresh InMemoryLightningStore instance."""
+    return InMemoryLightningStore()
+
+
+@pytest.fixture
+def mock_readable_span() -> ReadableSpan:
+    """Create a mock ReadableSpan for testing."""
+    span = Mock()
+    span.name = "test_span"
+
+    # Mock context
+    context = Mock()
+    context.trace_id = 111111
+    context.span_id = 222222
+    context.is_remote = False
+    context.trace_state = {}  # Make it an empty dict instead of Mock
+    span.get_span_context = Mock(return_value=context)
+
+    # Mock other attributes
+    span.parent = None
+    # Fix mock status to return proper string values
+    status_code_mock = Mock()
+    status_code_mock.name = "OK"
+    span.status = Mock(status_code=status_code_mock, description=None)
+    span.attributes = {"test": "value"}
+    span.events = []
+    span.links = []
+    span.start_time = time.time_ns()
+    span.end_time = time.time_ns() + 1000000
+    span.resource = Mock(attributes={}, schema_url="")
+
+    return span

tests/store/test_client_server.py

@@ -0,0 +1,243 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import asyncio
+import contextlib
+import socket
+from typing import AsyncGenerator, Tuple
+from unittest.mock import patch
+
+import pytest
+import pytest_asyncio
+from opentelemetry.sdk.trace import ReadableSpan
+
+from agentlightning.store.base import UNSET
+from agentlightning.store.client_server import LightningStoreClient, LightningStoreServer
+from agentlightning.store.memory import InMemoryLightningStore
+from agentlightning.tracer import Span
+from agentlightning.tracer.types import Resource, TraceStatus
+
+
+def _get_free_port() -> int:
+    with contextlib.closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as sock:
+        sock.bind(("127.0.0.1", 0))
+        return sock.getsockname()[1]
+
+
+def _make_span(rollout_id: str, attempt_id: str, sequence_id: int, name: str) -> Span:
+    return Span(
+        rollout_id=rollout_id,
+        attempt_id=attempt_id,
+        sequence_id=sequence_id,
+        trace_id=f"{sequence_id:032x}",
+        span_id=f"{sequence_id:016x}",
+        parent_id=None,
+        name=name,
+        status=TraceStatus(status_code="OK"),
+        attributes={},
+        events=[],
+        links=[],
+        start_time=None,
+        end_time=None,
+        context=None,
+        parent=None,
+        resource=Resource(attributes={}, schema_url=""),
+    )
+
+
+@pytest_asyncio.fixture
+async def server_client() -> AsyncGenerator[Tuple[LightningStoreServer, LightningStoreClient], None]:
+    store = InMemoryLightningStore()
+    port = _get_free_port()
+    server = LightningStoreServer(store, "127.0.0.1", port)
+    await server.start()
+    client = LightningStoreClient(server.endpoint)
+    try:
+        yield server, client
+    finally:
+        await client.close()
+        await server.stop()
+
+
+@pytest.mark.asyncio
+async def test_client_server_end_to_end(
+    server_client: Tuple[LightningStoreServer, LightningStoreClient], mock_readable_span: ReadableSpan
+) -> None:
+    server, client = server_client
+
+    # Server delegate coverage -------------------------------------------------
+    await server.update_resources("server-resources", {})
+    assert await server.get_resources_by_id("server-resources") is not None
+    assert await server.get_latest_resources() is not None
+
+    await server.start_rollout(input={"origin": "server"})
+    queued_rollout = await server.enqueue_rollout(input={"origin": "server-queue"})
+    dequeued = await server.dequeue_rollout()
+    started_attempt = await server.start_attempt(queued_rollout.rollout_id)
+
+    await server.query_rollouts()
+    await server.query_attempts(queued_rollout.rollout_id)
+    assert await server.get_latest_attempt(queued_rollout.rollout_id) is not None
+    assert await server.get_rollout_by_id(queued_rollout.rollout_id) is not None
+
+    assert dequeued is not None
+
+    server_span = _make_span(dequeued.rollout_id, dequeued.attempt.attempt_id, 0, "server-span")
+    await server.add_span(server_span)
+    assert await server.get_next_span_sequence_id(dequeued.rollout_id, dequeued.attempt.attempt_id) == 1
+
+    with patch("agentlightning.store.client_server.Span.from_opentelemetry", autospec=True) as mocked:
+        mocked.side_effect = lambda readable, rollout_id, attempt_id, sequence_id: _make_span(  # pyright: ignore[reportUnknownLambdaType]
+            rollout_id,  # pyright: ignore[reportUnknownArgumentType]
+            attempt_id,  # pyright: ignore[reportUnknownArgumentType]
+            sequence_id,  # pyright: ignore[reportUnknownArgumentType]
+            f"server-otel-{sequence_id}",  # pyright: ignore[reportUnknownArgumentType]
+        )
+        await server.add_otel_span(dequeued.rollout_id, dequeued.attempt.attempt_id, mock_readable_span)
+
+    await server.query_spans(dequeued.rollout_id)
+    await server.update_rollout(queued_rollout.rollout_id, status="running")
+    await server.update_attempt(
+        queued_rollout.rollout_id,
+        started_attempt.attempt.attempt_id,
+        status="running",
+        worker_id="server-worker",
+        metadata={"phase": "warmup"},
+    )
+    await server.update_attempt(queued_rollout.rollout_id, "latest", status="succeeded")
+    completed = await server.wait_for_rollouts(rollout_ids=[queued_rollout.rollout_id], timeout=0.1)
+    assert completed and completed[0].status in {"succeeded", "failed", "cancelled"}
+
+    # Client HTTP round trip ---------------------------------------------------
+    resource_update = await client.update_resources("client-resources", {})
+    assert resource_update.resources == {}
+    assert await client.get_resources_by_id("client-resources") is not None
+    assert await client.get_latest_resources() is not None
+
+    _attempted = await client.start_rollout(input={"origin": "client"}, mode="train", metadata={"step": 0})
+    enqueued = await client.enqueue_rollout(input={"origin": "client-queue"})
+    dequeued_client = await client.dequeue_rollout()
+    assert dequeued_client is not None
+    started_client_attempt = await client.start_attempt(dequeued_client.rollout_id)
+
+    all_rollouts = await client.query_rollouts()
+    assert any(r.rollout_id == enqueued.rollout_id for r in all_rollouts)
+    assert await client.query_rollouts(rollout_ids=[enqueued.rollout_id])
+    attempts = await client.query_attempts(dequeued_client.rollout_id)
+    assert attempts
+    assert await client.get_latest_attempt(dequeued_client.rollout_id) is not None
+    assert await client.get_rollout_by_id(dequeued_client.rollout_id) is not None
+
+    client_span = _make_span(dequeued_client.rollout_id, dequeued_client.attempt.attempt_id, 101, "client-span")
+    stored_span = await client.add_span(client_span)
+    assert stored_span.name == "client-span"
+    assert await client.get_next_span_sequence_id(dequeued_client.rollout_id, dequeued_client.attempt.attempt_id) == 102
+
+    with patch("agentlightning.store.client_server.Span.from_opentelemetry", autospec=True) as mocked:
+        mocked.side_effect = lambda readable, rollout_id, attempt_id, sequence_id: _make_span(  # pyright: ignore[reportUnknownLambdaType]
+            rollout_id,  # pyright: ignore[reportUnknownArgumentType]
+            attempt_id,  # pyright: ignore[reportUnknownArgumentType]
+            sequence_id,  # pyright: ignore[reportUnknownArgumentType]
+            f"client-otel-{sequence_id}",
+        )
+        await client.add_otel_span(dequeued_client.rollout_id, dequeued_client.attempt.attempt_id, mock_readable_span)
+
+    spans = await client.query_spans(dequeued_client.rollout_id)
+    assert spans
+
+    await client.update_rollout(dequeued_client.rollout_id, mode="val", metadata={"step": 1})
+    await client.update_attempt(
+        dequeued_client.rollout_id,
+        started_client_attempt.attempt.attempt_id,
+        worker_id="client-worker",
+        metadata={"info": "started"},
+    )
+    await client.update_attempt(dequeued_client.rollout_id, "latest", status="succeeded")
+    await client.update_rollout(dequeued_client.rollout_id, status="succeeded")
+
+    wait_result = await client.wait_for_rollouts(rollout_ids=[dequeued_client.rollout_id], timeout=0.05)
+    assert wait_result and wait_result[0].status == "succeeded"
+
+
+@pytest.mark.asyncio
+async def test_update_rollout_none_vs_unset(server_client: Tuple[LightningStoreServer, LightningStoreClient]) -> None:
+    _, client = server_client
+
+    attempted = await client.start_rollout(input={"payload": True}, metadata={"keep": True})
+    rollout_id = attempted.rollout_id
+
+    await client.update_rollout(rollout_id, mode="train", metadata={"extra": 1})
+    updated = await client.get_rollout_by_id(rollout_id)
+
+    assert updated is not None
+    assert updated.mode == "train"
+    assert updated.metadata is not None
+    assert updated.metadata["extra"] == 1
+
+    await client.update_rollout(rollout_id, mode=None, metadata={"extra1": 2})
+    cleared = await client.get_rollout_by_id(rollout_id)
+    assert cleared is not None
+    assert cleared.mode is None
+    assert cleared.metadata is not None
+    assert cleared.metadata == {"extra1": 2}
+
+    await client.update_rollout(rollout_id, mode=UNSET, metadata=UNSET, status="running")
+    preserved = await client.get_rollout_by_id(rollout_id)
+    assert preserved is not None
+    assert preserved.mode is None
+    assert preserved.metadata == {"extra1": 2}
+    assert preserved.status == "running"
+
+
+@pytest.mark.asyncio
+async def test_update_attempt_none_vs_unset(server_client: Tuple[LightningStoreServer, LightningStoreClient]) -> None:
+    _, client = server_client
+
+    attempted = await client.start_rollout(input={"payload": True})
+    rollout_id = attempted.rollout_id
+    attempt_id = attempted.attempt.attempt_id
+
+    await client.update_attempt(rollout_id, attempt_id, worker_id="worker-1", metadata={"stage": "init"})
+    initial = await client.get_latest_attempt(rollout_id)
+    assert initial is not None
+    assert initial.worker_id == "worker-1"
+    assert initial.metadata is not None
+    assert initial.metadata["stage"] == "init"
+
+    await client.update_attempt(rollout_id, "latest", worker_id="", metadata={})
+    cleared = await client.get_latest_attempt(rollout_id)
+    assert cleared is not None
+    assert cleared.worker_id == ""
+    assert cleared.metadata == {}
+
+    await client.update_attempt(rollout_id, "latest", status="running", worker_id=UNSET, metadata=UNSET)
+    preserved = await client.get_latest_attempt(rollout_id)
+    assert preserved is not None
+    assert preserved.worker_id == ""
+    assert preserved.metadata == {}
+    assert preserved.status == "running"
+
+
+@pytest.mark.asyncio
+async def test_concurrent_add_otel_span_sequence_ids_unique(
+    server_client: Tuple[LightningStoreServer, LightningStoreClient], mock_readable_span: ReadableSpan
+) -> None:
+    _, client = server_client
+
+    attempted = await client.start_rollout(input={"payload": True})
+    rollout_id = attempted.rollout_id
+    attempt_id = attempted.attempt.attempt_id
+
+    def _build_concurrent_span(readable: ReadableSpan, rollout_id: str, attempt_id: str, sequence_id: int) -> Span:
+        return _make_span(rollout_id, attempt_id, sequence_id, f"concurrent-{sequence_id}")
+
+    with patch("agentlightning.store.client_server.Span.from_opentelemetry", autospec=True) as mocked:
+        mocked.side_effect = _build_concurrent_span
+        spans = await asyncio.gather(
+            *[client.add_otel_span(rollout_id, attempt_id, mock_readable_span) for _ in range(20)]
+        )
+    sequence_ids = [span.sequence_id for span in spans]
+    assert len(set(sequence_ids)) == 20
+    assert set(sequence_ids) == set(range(1, 21))
+
+    stored_spans = await client.query_spans(rollout_id, attempt_id="latest")
+    assert len(stored_spans) >= 2

tests/store/test_threading.py

@@ -0,0 +1,415 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import asyncio
+from concurrent.futures import ThreadPoolExecutor
+from typing import Any, Dict, List, Literal, Optional, Sequence
+from unittest.mock import MagicMock
+
+import pytest
+from opentelemetry.sdk.trace import ReadableSpan
+
+from agentlightning.store.base import UNSET, LightningStore
+from agentlightning.store.threading import LightningStoreThreaded
+from agentlightning.tracer import Span
+from agentlightning.tracer.types import Resource, SpanContext, TraceStatus
+from agentlightning.types import (
+    Attempt,
+    AttemptedRollout,
+    AttemptStatus,
+    NamedResources,
+    ResourcesUpdate,
+    RolloutStatus,
+    RolloutV2,
+    TaskInput,
+)
+
+
+class DummyLightningStore(LightningStore):
+    def __init__(self, return_values: Dict[str, Any]) -> None:
+        super().__init__()
+        self.calls: List[tuple[str, tuple[Any, ...], Dict[str, Any]]] = []
+        self.return_values = return_values
+
+    async def start_rollout(
+        self,
+        input: TaskInput,
+        mode: Optional[str] = None,
+        resources_id: Optional[str] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+    ) -> AttemptedRollout:
+        self.calls.append(("start_rollout", (input, mode, resources_id, metadata), {}))
+        return self.return_values["start_rollout"]
+
+    async def enqueue_rollout(
+        self,
+        input: TaskInput,
+        mode: Optional[str] = None,
+        resources_id: Optional[str] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+    ) -> RolloutV2:
+        self.calls.append(("enqueue_rollout", (input, mode, resources_id, metadata), {}))
+        return self.return_values["enqueue_rollout"]
+
+    async def dequeue_rollout(self) -> Optional[AttemptedRollout]:
+        self.calls.append(("dequeue_rollout", (), {}))
+        return self.return_values["dequeue_rollout"]
+
+    async def start_attempt(self, rollout_id: str) -> AttemptedRollout:
+        self.calls.append(("start_attempt", (rollout_id,), {}))
+        return self.return_values["start_attempt"]
+
+    async def query_rollouts(
+        self, *, status: Optional[Sequence[RolloutStatus]] = None, rollout_ids: Optional[Sequence[str]] = None
+    ) -> List[RolloutV2]:
+        self.calls.append(("query_rollouts", (), {"status": status, "rollout_ids": rollout_ids}))
+        return self.return_values["query_rollouts"]
+
+    async def query_attempts(self, rollout_id: str) -> List[Attempt]:
+        self.calls.append(("query_attempts", (rollout_id,), {}))
+        return self.return_values["query_attempts"]
+
+    async def get_rollout_by_id(self, rollout_id: str) -> Optional[RolloutV2]:
+        self.calls.append(("get_rollout_by_id", (rollout_id,), {}))
+        return self.return_values["get_rollout_by_id"]
+
+    async def get_latest_attempt(self, rollout_id: str) -> Optional[Attempt]:
+        self.calls.append(("get_latest_attempt", (rollout_id,), {}))
+        return self.return_values["get_latest_attempt"]
+
+    async def update_resources(self, resources_id: str, resources: NamedResources) -> ResourcesUpdate:
+        self.calls.append(("update_resources", (resources_id, resources), {}))
+        return self.return_values["update_resources"]
+
+    async def get_resources_by_id(self, resources_id: str) -> Optional[ResourcesUpdate]:
+        self.calls.append(("get_resources_by_id", (resources_id,), {}))
+        return self.return_values["get_resources_by_id"]
+
+    async def get_latest_resources(self) -> Optional[ResourcesUpdate]:
+        self.calls.append(("get_latest_resources", (), {}))
+        return self.return_values["get_latest_resources"]
+
+    async def add_span(self, span: Span) -> Span:
+        self.calls.append(("add_span", (span,), {}))
+        return self.return_values["add_span"]
+
+    async def add_otel_span(
+        self,
+        rollout_id: str,
+        attempt_id: str,
+        readable_span: ReadableSpan,
+        sequence_id: Optional[int] = None,
+    ) -> Span:
+        self.calls.append(("add_otel_span", (rollout_id, attempt_id, readable_span, sequence_id), {}))
+        return self.return_values["add_otel_span"]
+
+    async def wait_for_rollouts(self, *, rollout_ids: List[str], timeout: Optional[float] = None) -> List[RolloutV2]:
+        self.calls.append(("wait_for_rollouts", (), {"rollout_ids": rollout_ids, "timeout": timeout}))
+        return self.return_values["wait_for_rollouts"]
+
+    async def get_next_span_sequence_id(self, rollout_id: str, attempt_id: str) -> int:
+        self.calls.append(("get_next_span_sequence_id", (rollout_id, attempt_id), {}))
+        return self.return_values["get_next_span_sequence_id"]
+
+    async def query_spans(
+        self,
+        rollout_id: str,
+        attempt_id: str | Literal["latest"] | None = None,
+    ) -> List[Span]:
+        self.calls.append(("query_spans", (rollout_id, attempt_id), {}))
+        return self.return_values["query_spans"]
+
+    async def update_rollout(
+        self,
+        rollout_id: str,
+        input: TaskInput | Any = UNSET,
+        mode: Optional[str] | Any = UNSET,
+        resources_id: Optional[str] | Any = UNSET,
+        status: RolloutStatus | Any = UNSET,
+        config: Any = UNSET,
+        metadata: Optional[Dict[str, Any]] | Any = UNSET,
+    ) -> RolloutV2:
+        self.calls.append(
+            (
+                "update_rollout",
+                (rollout_id, input, mode, resources_id, status, config, metadata),
+                {},
+            )
+        )
+        return self.return_values["update_rollout"]
+
+    async def update_attempt(
+        self,
+        rollout_id: str,
+        attempt_id: str | Literal["latest"],
+        status: AttemptStatus | Any = UNSET,
+        worker_id: str | Any = UNSET,
+        last_heartbeat_time: float | Any = UNSET,
+        metadata: Optional[Dict[str, Any]] | Any = UNSET,
+    ) -> Attempt:
+        self.calls.append(
+            (
+                "update_attempt",
+                (rollout_id, attempt_id, status, worker_id, last_heartbeat_time, metadata),
+                {},
+            )
+        )
+        return self.return_values["update_attempt"]
+
+
+class SlowAttemptStore(LightningStore):
+    def __init__(self) -> None:
+        super().__init__()
+        self.active_calls = 0
+        self.max_active_calls = 0
+        self.sequence = 0
+
+    async def update_attempt(
+        self,
+        rollout_id: str,
+        attempt_id: str,
+        status: AttemptStatus | Any = UNSET,
+        worker_id: str | Any = UNSET,
+        last_heartbeat_time: float | Any = UNSET,
+        metadata: Dict[str, Any] | Any = UNSET,
+    ) -> Attempt:
+        self.active_calls += 1
+        self.max_active_calls = max(self.max_active_calls, self.active_calls)
+        await asyncio.sleep(0.01)
+        self.active_calls -= 1
+
+        self.sequence += 1
+        status_value = status if status is not UNSET else "preparing"
+        worker_value = worker_id if worker_id is not UNSET else None
+        heartbeat_value = last_heartbeat_time if last_heartbeat_time is not UNSET else None
+        metadata_value = metadata if metadata is not UNSET else {}
+
+        return Attempt(
+            rollout_id=rollout_id,
+            attempt_id=attempt_id,
+            sequence_id=self.sequence,
+            start_time=float(self.sequence),
+            status=status_value,
+            worker_id=worker_value,
+            last_heartbeat_time=heartbeat_value,
+            metadata=metadata_value,
+        )
+
+
+class IncrementingResourceStore(LightningStore):
+    def __init__(self) -> None:
+        super().__init__()
+        self.counter = 0
+
+    async def update_resources(self, resources_id: str, resources: NamedResources) -> ResourcesUpdate:
+        snapshot = self.counter
+        await asyncio.sleep(0.01)
+        self.counter = snapshot + 1
+        return ResourcesUpdate(resources_id=f"res-{self.counter}", resources=resources)
+
+
+def make_span(rollout_id: str, attempt_id: str, sequence_id: int = 1) -> Span:
+    return Span(
+        rollout_id=rollout_id,
+        attempt_id=attempt_id,
+        sequence_id=sequence_id,
+        trace_id="0" * 32,
+        span_id="0" * 16,
+        parent_id=None,
+        name="test-span",
+        status=TraceStatus(status_code="OK"),
+        attributes={},
+        events=[],
+        links=[],
+        start_time=None,
+        end_time=None,
+        context=SpanContext(trace_id="0" * 32, span_id="0" * 16, is_remote=False, trace_state={}),
+        parent=None,
+        resource=Resource(attributes={}, schema_url=""),
+    )
+
+
+@pytest.mark.asyncio
+async def test_threaded_store_delegates_all_methods() -> None:
+    task_input: TaskInput = {"foo": "bar"}
+    rollout_id = "rollout-1"
+    attempt_id = "attempt-1"
+
+    base_attempt = Attempt(
+        rollout_id=rollout_id,
+        attempt_id=attempt_id,
+        sequence_id=1,
+        start_time=0.0,
+    )
+    base_rollout = RolloutV2(
+        rollout_id=rollout_id,
+        input=task_input,
+        start_time=0.0,
+        mode="train",
+    )
+    attempted_rollout = AttemptedRollout(
+        rollout_id=rollout_id,
+        input=task_input,
+        start_time=0.0,
+        mode="train",
+        status="preparing",
+        metadata={},
+        attempt=base_attempt,
+    )
+    resources_update = ResourcesUpdate(resources_id="resources-1", resources={})
+    span = make_span(rollout_id, attempt_id)
+    readable_span = MagicMock(spec=ReadableSpan)
+
+    updated_rollout = RolloutV2(
+        rollout_id=rollout_id,
+        input=task_input,
+        start_time=1.0,
+        mode="val",
+        status="succeeded",
+        resources_id="resources-2",
+        metadata={"note": "done"},
+    )
+    updated_attempt = Attempt(
+        rollout_id=rollout_id,
+        attempt_id=attempt_id,
+        sequence_id=2,
+        start_time=1.0,
+        status="running",
+        worker_id="worker-1",
+        last_heartbeat_time=1.5,
+        metadata={"idx": 0},
+    )
+
+    return_values = {
+        "start_rollout": attempted_rollout,
+        "enqueue_rollout": base_rollout,
+        "dequeue_rollout": attempted_rollout,
+        "start_attempt": attempted_rollout,
+        "query_rollouts": [base_rollout],
+        "query_attempts": [base_attempt],
+        "get_rollout_by_id": base_rollout,
+        "get_latest_attempt": base_attempt,
+        "update_resources": resources_update,
+        "get_resources_by_id": resources_update,
+        "get_latest_resources": resources_update,
+        "add_span": span,
+        "add_otel_span": span,
+        "wait_for_rollouts": [base_rollout],
+        "get_next_span_sequence_id": 42,
+        "query_spans": [span],
+        "update_rollout": updated_rollout,
+        "update_attempt": updated_attempt,
+    }
+
+    dummy_store = DummyLightningStore(return_values)
+    threaded_store = LightningStoreThreaded(dummy_store)
+
+    assert (
+        await threaded_store.start_rollout(task_input, mode="train", resources_id="resources-1", metadata={"a": 1})
+        == attempted_rollout
+    )
+    assert (
+        await threaded_store.enqueue_rollout(task_input, mode="train", resources_id="resources-1", metadata={"b": 2})
+        == base_rollout
+    )
+    assert await threaded_store.dequeue_rollout() == attempted_rollout
+    assert await threaded_store.start_attempt(rollout_id) == attempted_rollout
+    assert await threaded_store.query_rollouts(status=["preparing", "running"], rollout_ids=[rollout_id]) == [
+        base_rollout
+    ]
+    assert await threaded_store.query_attempts(rollout_id) == [base_attempt]
+    assert await threaded_store.get_rollout_by_id(rollout_id) == base_rollout
+    assert await threaded_store.get_latest_attempt(rollout_id) == base_attempt
+    assert await threaded_store.update_resources("resources-1", {}) == resources_update
+    assert await threaded_store.get_resources_by_id("resources-1") == resources_update
+    assert await threaded_store.get_latest_resources() == resources_update
+    assert await threaded_store.add_span(span) == span
+    assert await threaded_store.add_otel_span(rollout_id, attempt_id, readable_span, sequence_id=5) == span
+    assert await threaded_store.wait_for_rollouts(rollout_ids=[rollout_id], timeout=1.0) == [base_rollout]
+    assert await threaded_store.get_next_span_sequence_id(rollout_id, attempt_id) == 42
+    assert await threaded_store.query_spans(rollout_id, attempt_id="latest") == [span]
+    assert (
+        await threaded_store.update_rollout(
+            rollout_id,
+            input=task_input,
+            mode="val",
+            resources_id="resources-2",
+            status="succeeded",
+            metadata={"note": "done"},
+        )
+        == updated_rollout
+    )
+    assert (
+        await threaded_store.update_attempt(
+            rollout_id,
+            attempt_id,
+            status="running",
+            worker_id="worker-1",
+            last_heartbeat_time=1.5,
+            metadata={"idx": 0},
+        )
+        == updated_attempt
+    )
+
+    expected_order = [
+        "start_rollout",
+        "enqueue_rollout",
+        "dequeue_rollout",
+        "start_attempt",
+        "query_rollouts",
+        "query_attempts",
+        "get_rollout_by_id",
+        "get_latest_attempt",
+        "update_resources",
+        "get_resources_by_id",
+        "get_latest_resources",
+        "add_span",
+        "add_otel_span",
+        "wait_for_rollouts",
+        "get_next_span_sequence_id",
+        "query_spans",
+        "update_rollout",
+        "update_attempt",
+    ]
+    assert [name for name, *_ in dummy_store.calls] == expected_order
+
+
+def test_threaded_store_serializes_update_attempt_calls() -> None:
+    store = SlowAttemptStore()
+    threaded_store = LightningStoreThreaded(store)
+    rollout_id = "rollout-race"
+    num_calls = 5
+
+    def invoke(idx: int) -> Attempt:
+        return asyncio.run(
+            threaded_store.update_attempt(
+                rollout_id,
+                f"attempt-{idx}",
+                status="running",
+                worker_id=f"worker-{idx}",
+                last_heartbeat_time=float(idx),
+                metadata={"idx": idx},
+            )
+        )
+
+    with ThreadPoolExecutor(max_workers=num_calls) as executor:
+        results = list(executor.map(invoke, range(num_calls)))
+
+    assert store.max_active_calls == 1
+    assert len(results) == num_calls
+    assert sorted(result.attempt_id for result in results) == [f"attempt-{i}" for i in range(num_calls)]
+
+
+def test_threaded_store_prevents_race_conditions_on_resource_updates() -> None:
+    store = IncrementingResourceStore()
+    threaded_store = LightningStoreThreaded(store)
+    num_updates = 20
+
+    def invoke(idx: int) -> ResourcesUpdate:
+        return asyncio.run(threaded_store.update_resources(f"resources-{idx}", {}))
+
+    with ThreadPoolExecutor(max_workers=num_updates) as executor:
+        updates = list(executor.map(invoke, range(num_updates)))
+
+    assert store.counter == num_updates
+    assert len(updates) == num_updates
+    assert {update.resources_id for update in updates} == {f"res-{i + 1}" for i in range(num_updates)}

tests/store/test_utils.py

@@ -0,0 +1,325 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import time
+from typing import List, Optional, cast
+from unittest.mock import AsyncMock, patch
+
+import pytest
+
+from agentlightning.store.utils import healthcheck, propagate_status
+from agentlightning.types import (
+    Attempt,
+    AttemptedRollout,
+    AttemptStatus,
+    RolloutConfig,
+)
+
+# Tests for propagate_status function
+
+
+@pytest.mark.parametrize(
+    "status,expected_call",
+    [
+        ("preparing", "preparing"),
+        ("running", "running"),
+        ("succeeded", "succeeded"),
+    ],
+)
+@pytest.mark.asyncio
+async def test_propagate_status_direct_statuses(status: AttemptStatus, expected_call: AttemptStatus) -> None:
+    """Test propagate_status directly propagates preparing/running/succeeded statuses."""
+    attempt = Attempt(
+        rollout_id="test-rollout", attempt_id="test-attempt", sequence_id=1, start_time=time.time(), status=status
+    )
+    config = RolloutConfig()
+    update_rollout_mock = AsyncMock()
+
+    await propagate_status(update_rollout_mock, attempt, config)
+
+    update_rollout_mock.assert_called_once_with("test-rollout", expected_call)
+
+
+@pytest.mark.parametrize(
+    "status,in_retry_condition,sequence_id,max_attempts,expected_call",
+    [
+        ("failed", True, 1, 3, "requeuing"),  # Should retry
+        ("failed", True, 3, 3, "failed"),  # Max attempts reached
+        ("failed", False, 1, 3, "failed"),  # Not in retry condition
+        ("timeout", True, 2, 3, "requeuing"),  # Should retry
+        ("timeout", False, 1, 3, "failed"),  # Not in retry condition
+        ("unresponsive", True, 1, 2, "requeuing"),  # Should retry
+    ],
+)
+@pytest.mark.asyncio
+async def test_propagate_status_retry_logic(
+    status: AttemptStatus, in_retry_condition: bool, sequence_id: int, max_attempts: int, expected_call: AttemptStatus
+) -> None:
+    """Test propagate_status retry logic for different combinations."""
+    attempt = Attempt(
+        rollout_id="test-rollout",
+        attempt_id="test-attempt",
+        sequence_id=sequence_id,
+        start_time=time.time(),
+        status=status,
+    )
+
+    retry_condition: List[AttemptStatus] = [status] if in_retry_condition else []
+    config = RolloutConfig(max_attempts=max_attempts, retry_condition=retry_condition)
+    update_rollout_mock = AsyncMock()
+
+    await propagate_status(update_rollout_mock, attempt, config)
+
+    update_rollout_mock.assert_called_once_with("test-rollout", expected_call)
+
+
+@pytest.mark.asyncio
+async def test_propagate_status_invalid_status() -> None:
+    """Test propagate_status raises error for invalid status."""
+    # Create a valid attempt first, then modify its status
+    attempt = Attempt(
+        rollout_id="test-rollout", attempt_id="test-attempt", sequence_id=1, start_time=time.time(), status="failed"
+    )
+    # Bypass Pydantic validation by directly setting the attribute
+    attempt.status = cast(AttemptStatus, "invalid_status")  # Invalid status
+
+    config = RolloutConfig()
+    update_rollout_mock = AsyncMock()
+
+    with pytest.raises(ValueError, match="Invalid attempt status: invalid_status"):
+        await propagate_status(update_rollout_mock, attempt, config)
+
+
+# Tests for healthcheck function
+
+
+@pytest.mark.asyncio
+async def test_healthcheck_empty_rollouts_list() -> None:
+    """Test healthcheck handles empty rollouts list gracefully."""
+    update_rollout_mock = AsyncMock()
+    update_attempt_mock = AsyncMock()
+
+    await healthcheck([], update_rollout_mock, update_attempt_mock)
+
+    # Should not call any updates
+    update_rollout_mock.assert_not_called()
+    update_attempt_mock.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_healthcheck_multiple_rollouts_different_timeouts() -> None:
+    """Test healthcheck handles multiple rollouts with different timeout configs."""
+    current_time = time.time()
+
+    # Rollout 1: Short timeout, should timeout
+    config1 = RolloutConfig(timeout_seconds=1.0)
+    attempt1 = Attempt(
+        rollout_id="rollout-1",
+        attempt_id="attempt-1",
+        sequence_id=1,
+        start_time=current_time - 2.0,
+        status="running",  # 2 seconds ago
+    )
+    rollout1 = AttemptedRollout(
+        rollout_id="rollout-1",
+        input={"test": 1},
+        status="running",
+        start_time=current_time,
+        config=config1,
+        attempt=attempt1,
+    )
+
+    # Rollout 2: Long timeout, should not timeout
+    config2 = RolloutConfig(timeout_seconds=10.0)
+    attempt2 = Attempt(
+        rollout_id="rollout-2",
+        attempt_id="attempt-2",
+        sequence_id=1,
+        start_time=current_time - 2.0,
+        status="running",  # 2 seconds ago
+    )
+    rollout2 = AttemptedRollout(
+        rollout_id="rollout-2",
+        input={"test": 2},
+        status="running",
+        start_time=current_time,
+        config=config2,
+        attempt=attempt2,
+    )
+
+    update_rollout_mock = AsyncMock()
+    update_attempt_mock = AsyncMock()
+
+    with patch("time.time", return_value=current_time):
+        await healthcheck([rollout1, rollout2], update_rollout_mock, update_attempt_mock)
+
+    # Only rollout1 should be marked as timeout
+    update_attempt_mock.assert_called_once_with("rollout-1", "attempt-1", "timeout")
+
+
+@pytest.mark.parametrize(
+    "timeout_seconds,unresponsive_seconds,should_timeout,should_unresponsive",
+    [
+        (None, None, False, False),  # No timeouts configured
+        (None, 1.0, False, True),  # Only unresponsive timeout
+        (1.0, None, True, False),  # Only regular timeout
+        (0.5, 1.0, True, False),  # Timeout triggers first
+        (2.0, 0.5, False, True),  # Unresponsive triggers first
+    ],
+)
+@pytest.mark.asyncio
+async def test_healthcheck_timeout_configurations(
+    timeout_seconds: Optional[float],
+    unresponsive_seconds: Optional[float],
+    should_timeout: bool,
+    should_unresponsive: bool,
+) -> None:
+    """Test healthcheck with various timeout configurations."""
+    current_time = time.time()
+
+    config = RolloutConfig(timeout_seconds=timeout_seconds, unresponsive_seconds=unresponsive_seconds)
+
+    attempt = Attempt(
+        rollout_id="test-rollout",
+        attempt_id="test-attempt",
+        sequence_id=1,
+        start_time=current_time - 1.5,
+        status="running",  # 1.5 seconds ago
+        last_heartbeat_time=None,  # No heartbeat for unresponsive detection
+    )
+
+    rollout = AttemptedRollout(
+        rollout_id="test-rollout",
+        input={"test": 1},
+        status="running",
+        start_time=current_time,
+        config=config,
+        attempt=attempt,
+    )
+
+    update_rollout_mock = AsyncMock()
+    update_attempt_mock = AsyncMock()
+
+    with patch("time.time", return_value=current_time):
+        await healthcheck([rollout], update_rollout_mock, update_attempt_mock)
+
+    if should_timeout:
+        update_attempt_mock.assert_called_once_with("test-rollout", "test-attempt", "timeout")
+    elif should_unresponsive:
+        update_attempt_mock.assert_called_once_with("test-rollout", "test-attempt", "unresponsive")
+    else:
+        update_attempt_mock.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_healthcheck_unresponsive_with_heartbeat_timing() -> None:
+    """Test unresponsive detection considers heartbeat timing correctly."""
+    current_time = time.time()
+    config = RolloutConfig(unresponsive_seconds=1.0)
+
+    # Case 1: Recent heartbeat - should not be unresponsive
+    attempt_recent = Attempt(
+        rollout_id="rollout-recent",
+        attempt_id="attempt-recent",
+        sequence_id=1,
+        start_time=current_time - 5.0,
+        status="running",
+        last_heartbeat_time=current_time - 0.5,  # Recent heartbeat
+    )
+    rollout_recent = AttemptedRollout(
+        rollout_id="rollout-recent",
+        input={"test": 1},
+        status="running",
+        start_time=current_time,
+        config=config,
+        attempt=attempt_recent,
+    )
+
+    # Case 2: Old heartbeat - should be unresponsive
+    attempt_old = Attempt(
+        rollout_id="rollout-old",
+        attempt_id="attempt-old",
+        sequence_id=1,
+        start_time=current_time - 5.0,
+        status="running",
+        last_heartbeat_time=current_time - 2.0,  # Old heartbeat
+    )
+    rollout_old = AttemptedRollout(
+        rollout_id="rollout-old",
+        input={"test": 2},
+        status="running",
+        start_time=current_time,
+        config=config,
+        attempt=attempt_old,
+    )
+
+    update_rollout_mock = AsyncMock()
+    update_attempt_mock = AsyncMock()
+
+    with patch("time.time", return_value=current_time):
+        await healthcheck([rollout_recent, rollout_old], update_rollout_mock, update_attempt_mock)
+
+    # Only the old heartbeat should trigger unresponsive
+    update_attempt_mock.assert_called_once_with("rollout-old", "attempt-old", "unresponsive")
+
+
+@pytest.mark.asyncio
+async def test_healthcheck_preparing_with_heartbeat_promotion() -> None:
+    """Test healthcheck promotes preparing attempts with heartbeat to running."""
+    current_time = time.time()
+
+    config = RolloutConfig()
+    attempt = Attempt(
+        rollout_id="test-rollout",
+        attempt_id="test-attempt",
+        sequence_id=1,
+        start_time=current_time,
+        status="preparing",
+        last_heartbeat_time=current_time,  # Has heartbeat
+    )
+    rollout = AttemptedRollout(
+        rollout_id="test-rollout",
+        input={"test": 1},
+        status="preparing",
+        start_time=current_time,
+        config=config,
+        attempt=attempt,
+    )
+
+    update_rollout_mock = AsyncMock()
+    update_attempt_mock = AsyncMock()
+
+    await healthcheck([rollout], update_rollout_mock, update_attempt_mock)
+
+    # Should promote to running
+    update_attempt_mock.assert_called_once_with("test-rollout", "test-attempt", "running")
+
+
+@pytest.mark.asyncio
+async def test_healthcheck_skips_rollouts_without_attempts() -> None:
+    """Test healthcheck gracefully skips rollouts with no attempts."""
+    config = RolloutConfig()
+
+    # Create a valid attempt first, then set it to None
+    attempt = Attempt(
+        rollout_id="test-rollout", attempt_id="test-attempt", sequence_id=1, start_time=time.time(), status="running"
+    )
+    rollout = AttemptedRollout(
+        rollout_id="test-rollout",
+        input={"test": 1},
+        status="running",
+        start_time=time.time(),
+        config=config,
+        attempt=attempt,
+    )
+
+    # Bypass Pydantic validation by directly setting the attribute
+    rollout.attempt = cast(Attempt, None)  # No attempt
+
+    update_rollout_mock = AsyncMock()
+    update_attempt_mock = AsyncMock()
+
+    await healthcheck([rollout], update_rollout_mock, update_attempt_mock)
+
+    # Should not call any updates
+    update_rollout_mock.assert_not_called()
+    update_attempt_mock.assert_not_called()