# Copyright (c) Microsoft. All rights reserved. """ Comprehensive tests for LightningStore. Test categories: - Core CRUD operations - Queue operations (FIFO behavior) - Resource versioning - Span tracking and sequencing - Rollout lifecycle and status transitions - Concurrent access patterns - Error handling and edge cases It should work for multiple store implementations (InMemory, SQL, etc.). """ import asyncio import logging import sys import time from typing import Any, List, Optional, Protocol, Sequence, cast from unittest.mock import Mock import pytest from pydantic import BaseModel from agentlightning.store import CollectionBasedLightningStore from agentlightning.store.base import UNSET, LightningStore from agentlightning.store.memory import InMemoryLightningStore, estimate_model_size from agentlightning.types import ( LLM, Attempt, AttemptedRollout, EnqueueRolloutRequest, Event, Link, OtelResource, PaginatedResult, PromptTemplate, ResourcesUpdate, Rollout, RolloutConfig, Span, SpanContext, TraceStatus, ) pytestmark = [pytest.mark.store] class FakeTimeController(Protocol): def set(self, value: float) -> None: ... def advance(self, delta: float) -> None: ... def test_paginated_result_behaves_like_sequence() -> None: result = PaginatedResult(items=["a", "b", "c"], limit=2, offset=1, total=5) assert isinstance(result, Sequence) assert len(result) == 3 assert result[0] == "a" assert result[1:] == ["b", "c"] assert list(result) == ["a", "b", "c"] assert repr(result) == "" result2 = PaginatedResult(items=["a", "b", "c"], limit=-1, offset=1, total=5) assert repr(result2) == "" @pytest.mark.asyncio async def test_statistics_updates_counts(store_fixture: LightningStore, mock_readable_span: Mock) -> None: """Statistics should reflect entity counts across the store.""" initial = await store_fixture.statistics() rollout = await store_fixture.start_rollout(input={"origin": "stats"}) await store_fixture.add_otel_span(rollout.rollout_id, rollout.attempt.attempt_id, mock_readable_span) await store_fixture.add_resources( { "stat_llm": LLM( resource_type="llm", endpoint="http://localhost:8000/v1", model="stats-model", ) } ) await store_fixture.update_worker("stats-worker", heartbeat_stats={"cpu": 0.5}) updated = await store_fixture.statistics() assert updated["name"] == store_fixture.__class__.__name__ # type: ignore assert updated["total_rollouts"] == initial.get("total_rollouts", 0) + 1 # type: ignore assert updated["total_attempts"] == initial.get("total_attempts", 0) + 1 # type: ignore assert updated["total_spans"] == initial.get("total_spans", 0) + 1 # type: ignore assert updated["total_resources"] == initial.get("total_resources", 0) + 1 # type: ignore assert updated["total_workers"] == initial.get("total_workers", 0) + 1 # type: ignore assert updated["uptime"] > initial.get("uptime", 0.0) # type: ignore # Core CRUD Operations Tests @pytest.mark.asyncio async def test_enqueue_rollout_creates_rollout(store_fixture: LightningStore) -> None: """Test that enqueue_rollout creates a properly initialized rollout.""" sample = {"input": "test_data"} metadata = {"key": "value", "number": 42} rollout = await store_fixture.enqueue_rollout(input=sample, mode="train", resources_id="res-123", metadata=metadata) assert rollout.rollout_id.startswith("ro-") assert rollout.input == sample assert rollout.mode == "train" assert rollout.resources_id == "res-123" assert rollout.metadata == metadata assert rollout.status == "queuing" assert rollout.start_time is not None @pytest.mark.asyncio async def test_enqueue_rollout_accepts_config(store_fixture: LightningStore) -> None: """Rollout-specific configs can be provided when enqueuing tasks.""" config = RolloutConfig(timeout_seconds=12.0, max_attempts=3, retry_condition=["timeout"]) rollout = await store_fixture.enqueue_rollout(input={"sample": True}, config=config) assert rollout.config.timeout_seconds == 12.0 assert rollout.config.max_attempts == 3 assert rollout.config.retry_condition == ["timeout"] stored = await store_fixture.get_rollout_by_id(rollout.rollout_id) assert stored is not None assert stored.config.timeout_seconds == 12.0 assert stored.config.max_attempts == 3 assert stored.config.retry_condition == ["timeout"] @pytest.mark.asyncio async def test_add_rollout_initializes_attempt(store_fixture: LightningStore) -> None: """Test that add_rollout immediately tracks a preparing attempt.""" sample = {"payload": "value"} attempt_rollout = await store_fixture.start_rollout(input=sample, mode="val", resources_id="res-add") assert attempt_rollout.status == "preparing" assert attempt_rollout.rollout_id.startswith("ro-") assert attempt_rollout.attempt.attempt_id.startswith("at-") assert attempt_rollout.attempt.sequence_id == 1 assert attempt_rollout.attempt.status == "preparing" stored = await store_fixture.query_rollouts(status=["preparing"]) assert len(stored) == 1 assert stored[0].rollout_id == attempt_rollout.rollout_id assert stored[0].resources_id == "res-add" attempts = await store_fixture.query_attempts(attempt_rollout.rollout_id) assert len(attempts) == 1 assert attempts[0].attempt_id == attempt_rollout.attempt.attempt_id latest_attempt = await store_fixture.get_latest_attempt(attempt_rollout.rollout_id) assert latest_attempt is not None assert latest_attempt.attempt_id == attempt_rollout.attempt.attempt_id @pytest.mark.asyncio async def test_start_rollout_accepts_config(store_fixture: LightningStore) -> None: """Custom rollout config is preserved for started rollouts.""" config = RolloutConfig(unresponsive_seconds=5.0, max_attempts=2, retry_condition=["unresponsive"]) attempt_rollout = await store_fixture.start_rollout(input={"payload": "value"}, config=config) assert attempt_rollout.config.unresponsive_seconds == 5.0 assert attempt_rollout.config.max_attempts == 2 assert attempt_rollout.config.retry_condition == ["unresponsive"] stored = await store_fixture.get_rollout_by_id(attempt_rollout.rollout_id) assert stored is not None assert stored.config.unresponsive_seconds == 5.0 assert stored.config.max_attempts == 2 assert stored.config.retry_condition == ["unresponsive"] @pytest.mark.asyncio async def test_query_rollouts_by_status(store_fixture: LightningStore) -> None: """Test querying rollouts filtered by status.""" # Create rollouts with different statuses r1 = await store_fixture.enqueue_rollout(input={"id": 1}) r2 = await store_fixture.enqueue_rollout(input={"id": 2}) r3 = await store_fixture.enqueue_rollout(input={"id": 3}) # Modify statuses await store_fixture.dequeue_rollout() # r1 becomes "preparing" await store_fixture.update_rollout(rollout_id=r2.rollout_id, status="failed") # r3 remains "queuing" # Test various queries all_rollouts = await store_fixture.query_rollouts() assert len(all_rollouts) == 3 queuing = await store_fixture.query_rollouts(status=["queuing"]) assert len(queuing) == 1 assert queuing[0].rollout_id == r3.rollout_id preparing = await store_fixture.query_rollouts(status=["preparing"]) assert len(preparing) == 1 assert preparing[0].rollout_id == r1.rollout_id finished = await store_fixture.query_rollouts(status=["failed", "succeeded"]) assert len(finished) == 1 assert finished[0].rollout_id == r2.rollout_id # Empty status list none = await store_fixture.query_rollouts(status=[]) assert len(none) == 0 @pytest.mark.asyncio async def test_query_rollouts_returns_latest_attempt(store_fixture: LightningStore) -> None: """Querying rollouts should attach the most recent attempt when present.""" attempted = await store_fixture.start_rollout(input={"sample": "latest"}) latest_attempt = await store_fixture.start_attempt(attempted.rollout_id) results = await store_fixture.query_rollouts(rollout_ids=[attempted.rollout_id]) assert len(results) == 1 retrieved = results[0] assert type(retrieved) is AttemptedRollout assert retrieved.attempt.attempt_id == latest_attempt.attempt.attempt_id assert retrieved.attempt.sequence_id == latest_attempt.attempt.sequence_id @pytest.mark.asyncio async def test_query_rollouts_supports_new_filters(store_fixture: LightningStore) -> None: """The expanded query interface should honor filtering, sorting, and pagination.""" rollouts = [await store_fixture.enqueue_rollout(input={"idx": idx}) for idx in range(3)] await store_fixture.update_rollout(rollout_id=rollouts[2].rollout_id, status="failed") failed = await store_fixture.query_rollouts(status_in=["failed"]) assert [r.rollout_id for r in failed] == [rollouts[2].rollout_id] sorted_desc = sorted([r.rollout_id for r in rollouts], reverse=True) paged = await store_fixture.query_rollouts(sort_by="rollout_id", sort_order="desc", limit=2) assert [r.rollout_id for r in paged] == sorted_desc[:2] offset_item = await store_fixture.query_rollouts(sort_by="rollout_id", sort_order="desc", limit=1, offset=1) assert [r.rollout_id for r in offset_item] == sorted_desc[1:2] contains = await store_fixture.query_rollouts( rollout_id_contains=rollouts[0].rollout_id[-4:], status_in=["queuing"], ) assert any(r.rollout_id == rollouts[0].rollout_id for r in contains) or_filtered = await store_fixture.query_rollouts( status_in=["succeeded"], rollout_id_in=[rollouts[1].rollout_id], filter_logic="or", ) assert [r.rollout_id for r in or_filtered] == [rollouts[1].rollout_id] @pytest.mark.asyncio async def test_query_rollouts_status_in_takes_precedence(store_fixture: LightningStore) -> None: """status_in should override the legacy status parameter when both are provided.""" failed = await store_fixture.enqueue_rollout(input={"kind": "failed"}) pending = await store_fixture.enqueue_rollout(input={"kind": "pending"}) await store_fixture.update_rollout(rollout_id=failed.rollout_id, status="failed") results = await store_fixture.query_rollouts(status=["queuing"], status_in=["failed"]) assert [rollout.rollout_id for rollout in results] == [failed.rollout_id] legacy = await store_fixture.query_rollouts(status=["queuing"]) assert any(rollout.rollout_id == pending.rollout_id for rollout in legacy) @pytest.mark.asyncio async def test_query_rollouts_rollout_id_in_takes_precedence(store_fixture: LightningStore) -> None: """rollout_id_in should override the legacy rollout_ids parameter.""" keep = await store_fixture.enqueue_rollout(input={"kind": "keep"}) ignored = await store_fixture.enqueue_rollout(input={"kind": "ignored"}) results = await store_fixture.query_rollouts(rollout_ids=[ignored.rollout_id], rollout_id_in=[keep.rollout_id]) assert [rollout.rollout_id for rollout in results] == [keep.rollout_id] @pytest.mark.asyncio async def test_query_rollouts_filter_logic_controls_contains_behavior(store_fixture: LightningStore) -> None: """Changing filter_logic should alter how status and substring filters are combined.""" failed = await store_fixture.enqueue_rollout(input={"kind": "failed"}) still_queueing = await store_fixture.enqueue_rollout(input={"kind": "queue"}) await store_fixture.update_rollout(rollout_id=failed.rollout_id, status="failed") substring = still_queueing.rollout_id[-6:] and_results = await store_fixture.query_rollouts( status_in=["failed"], rollout_id_contains=substring, ) assert len(and_results) == 0 or_results = await store_fixture.query_rollouts( status_in=["failed"], rollout_id_contains=substring, filter_logic="or", sort_by="rollout_id", ) returned_ids = {rollout.rollout_id for rollout in or_results} assert returned_ids == {failed.rollout_id, still_queueing.rollout_id} @pytest.mark.asyncio async def test_query_rollouts_combined_filters_with_sort_and_pagination(store_fixture: LightningStore) -> None: """Complex combinations of filters should work with explicit sorting and pagination.""" rollouts = [await store_fixture.enqueue_rollout(input={"idx": idx}) for idx in range(4)] await store_fixture.update_rollout(rollout_id=rollouts[0].rollout_id, status="failed") await store_fixture.update_rollout(rollout_id=rollouts[3].rollout_id, status="failed") substring = rollouts[2].rollout_id[-5:] filtered = await store_fixture.query_rollouts( status_in=["failed"], rollout_id_contains=substring, filter_logic="or", sort_by="rollout_id", sort_order="asc", limit=2, offset=1, ) expected_ids = sorted({rollouts[0].rollout_id, rollouts[2].rollout_id, rollouts[3].rollout_id}) assert [item.rollout_id for item in filtered] == expected_ids[1:3] @pytest.mark.asyncio async def test_query_rollouts_reports_pagination_metadata(store_fixture: LightningStore) -> None: """Paginated rollouts should expose limit/offset/total values.""" rollouts = [await store_fixture.enqueue_rollout(input={"idx": idx}) for idx in range(3)] paginated = await store_fixture.query_rollouts(sort_by="rollout_id", sort_order="asc", limit=1, offset=1) assert isinstance(paginated, PaginatedResult) assert paginated.limit == 1 assert paginated.offset == 1 assert paginated.total == len(rollouts) @pytest.mark.asyncio async def test_query_attempts_supports_sort_and_limit(store_fixture: LightningStore) -> None: """Attempt queries should respect sorting and pagination controls.""" attempted = await store_fixture.start_rollout(input={"payload": "attempt-filters"}) await store_fixture.start_attempt(attempted.rollout_id) await store_fixture.start_attempt(attempted.rollout_id) attempts_desc = await store_fixture.query_attempts(attempted.rollout_id, sort_by="sequence_id", sort_order="desc") assert [attempt.sequence_id for attempt in attempts_desc] == [3, 2, 1] middle_attempt = await store_fixture.query_attempts( attempted.rollout_id, sort_by="sequence_id", sort_order="desc", limit=1, offset=1, ) assert [attempt.sequence_id for attempt in middle_attempt] == [2] @pytest.mark.asyncio async def test_query_attempts_offset_past_end_returns_empty(store_fixture: LightningStore) -> None: """Offsets beyond the result size should return an empty list.""" attempted = await store_fixture.start_rollout(input={"payload": "attempt-offset"}) await store_fixture.start_attempt(attempted.rollout_id) await store_fixture.start_attempt(attempted.rollout_id) results = await store_fixture.query_attempts( attempted.rollout_id, sort_by="sequence_id", sort_order="asc", limit=1, offset=10, ) assert len(results) == 0 @pytest.mark.asyncio async def test_query_attempts_zero_limit_returns_no_items(store_fixture: LightningStore) -> None: """A zero limit should be treated as 'return nothing' even when attempts exist.""" attempted = await store_fixture.start_rollout(input={"payload": "attempt-limit"}) await store_fixture.start_attempt(attempted.rollout_id) results = await store_fixture.query_attempts(attempted.rollout_id, limit=0) assert len(results) == 0 @pytest.mark.asyncio async def test_query_attempts_reports_pagination_metadata(store_fixture: LightningStore) -> None: """Attempt pagination should retain metadata.""" attempted = await store_fixture.start_rollout(input={"payload": "attempt-pagination"}) await store_fixture.start_attempt(attempted.rollout_id) await store_fixture.start_attempt(attempted.rollout_id) paginated = await store_fixture.query_attempts( attempted.rollout_id, sort_by="sequence_id", sort_order="asc", limit=1, offset=1, ) assert isinstance(paginated, PaginatedResult) assert paginated.limit == 1 assert paginated.offset == 1 assert paginated.total == 3 @pytest.mark.asyncio async def test_get_rollout_by_id_returns_latest_attempt(store_fixture: LightningStore) -> None: """Fetching a rollout by ID should include the latest attempt when available.""" attempted = await store_fixture.start_rollout(input={"foo": "bar"}) second_attempt = await store_fixture.start_attempt(attempted.rollout_id) retrieved = await store_fixture.get_rollout_by_id(attempted.rollout_id) assert retrieved is not None assert type(retrieved) is AttemptedRollout assert retrieved.attempt.attempt_id == second_attempt.attempt.attempt_id assert retrieved.attempt.sequence_id == second_attempt.attempt.sequence_id @pytest.mark.asyncio async def test_get_rollout_by_id_without_attempt_returns_rollout( store_fixture: LightningStore, ) -> None: """Rollouts with no attempts should be returned without the Attempt wrapper.""" queued = await store_fixture.enqueue_rollout(input={"foo": "bar"}) retrieved = await store_fixture.get_rollout_by_id(queued.rollout_id) assert retrieved is not None assert type(retrieved) is Rollout assert not hasattr(retrieved, "attempt") @pytest.mark.asyncio async def test_get_rollout_by_id(store_fixture: LightningStore) -> None: """Test retrieving rollouts by their ID.""" # Test getting non-existent rollout rollout = await store_fixture.get_rollout_by_id("nonexistent") assert rollout is None # Create a rollout created = await store_fixture.enqueue_rollout(input={"test": "data"}, mode="train") # Retrieve by ID retrieved = await store_fixture.get_rollout_by_id(created.rollout_id) assert retrieved is not None assert retrieved.rollout_id == created.rollout_id assert retrieved.input == created.input assert retrieved.mode == created.mode assert retrieved.status == created.status # Update rollout and verify changes are reflected await store_fixture.update_rollout(rollout_id=created.rollout_id, status="running") updated = await store_fixture.get_rollout_by_id(created.rollout_id) assert updated is not None assert updated.status == "running" @pytest.mark.asyncio async def test_store_lock_rebinds_to_new_event_loop( store_fixture: LightningStore, ) -> None: """The in-memory store can be reused after switching to a new event loop.""" rollout = await store_fixture.enqueue_rollout(input={"foo": "bar"}) def run_in_new_loop() -> Optional[Rollout]: loop = asyncio.new_event_loop() try: return loop.run_until_complete(store_fixture.get_rollout_by_id(rollout.rollout_id)) finally: loop.close() retrieved = await asyncio.to_thread(run_in_new_loop) assert retrieved is not None assert retrieved.rollout_id == rollout.rollout_id @pytest.mark.asyncio async def test_query_rollouts_by_rollout_ids(store_fixture: LightningStore) -> None: """Test querying rollouts filtered by rollout IDs.""" # Create multiple rollouts r1 = await store_fixture.enqueue_rollout(input={"id": 1}) r2 = await store_fixture.enqueue_rollout(input={"id": 2}) r3 = await store_fixture.enqueue_rollout(input={"id": 3}) # Query by specific IDs selected = await store_fixture.query_rollouts(rollout_ids=[r1.rollout_id, r3.rollout_id]) assert len(selected) == 2 selected_ids = {r.rollout_id for r in selected} assert selected_ids == {r1.rollout_id, r3.rollout_id} # Query by single ID single = await store_fixture.query_rollouts(rollout_ids=[r2.rollout_id]) assert len(single) == 1 assert single[0].rollout_id == r2.rollout_id # Query by non-existent ID none = await store_fixture.query_rollouts(rollout_ids=["nonexistent"]) assert len(none) == 0 # Combine with status filter await store_fixture.update_rollout(rollout_id=r1.rollout_id, status="succeeded") await store_fixture.update_rollout(rollout_id=r2.rollout_id, status="failed") filtered = await store_fixture.query_rollouts( rollout_ids=[r1.rollout_id, r2.rollout_id, r3.rollout_id], status=["succeeded", "queuing"] ) assert len(filtered) == 2 filtered_ids = {r.rollout_id for r in filtered} assert filtered_ids == {r1.rollout_id, r3.rollout_id} # r1 succeeded, r3 still queuing @pytest.mark.asyncio async def test_update_rollout_fields(store_fixture: LightningStore) -> None: """Test updating various rollout fields.""" rollout = await store_fixture.enqueue_rollout(input={"test": "data"}) # Update multiple fields at once including config config = RolloutConfig( timeout_seconds=60.0, unresponsive_seconds=30.0, max_attempts=3, retry_condition=["timeout", "unresponsive"] ) await store_fixture.update_rollout( rollout_id=rollout.rollout_id, status="running", mode="train", resources_id="new-resources", config=config, metadata={"custom_field": "custom_value"}, ) # Verify all updates updated_rollouts = await store_fixture.query_rollouts() updated = updated_rollouts[0] assert updated.status == "running" assert updated.mode == "train" assert updated.resources_id == "new-resources" assert updated.config.timeout_seconds == 60.0 assert updated.config.unresponsive_seconds == 30.0 assert updated.config.max_attempts == 3 assert updated.config.retry_condition == ["timeout", "unresponsive"] assert updated.metadata is not None assert updated.metadata["custom_field"] == "custom_value" @pytest.mark.asyncio async def test_rollout_config_functionality(store_fixture: LightningStore) -> None: """Test RolloutConfig controls retry and timeout behavior.""" # Create rollout with specific retry configuration config = RolloutConfig( timeout_seconds=30.0, unresponsive_seconds=15.0, max_attempts=2, retry_condition=["timeout", "unresponsive", "failed"], ) rollout = await store_fixture.enqueue_rollout(input={"test": "retry"}) await store_fixture.update_rollout(rollout_id=rollout.rollout_id, config=config) # Verify config is stored stored = await store_fixture.get_rollout_by_id(rollout.rollout_id) assert stored is not None assert stored.config.timeout_seconds == 30.0 assert stored.config.max_attempts == 2 assert "failed" in stored.config.retry_condition # Test that different rollouts can have different configs config2 = RolloutConfig(timeout_seconds=120.0, max_attempts=5, retry_condition=["timeout"]) rollout2 = await store_fixture.enqueue_rollout(input={"test": "different_config"}) await store_fixture.update_rollout(rollout_id=rollout2.rollout_id, config=config2) stored2 = await store_fixture.get_rollout_by_id(rollout2.rollout_id) assert stored2 is not None assert stored2.config.timeout_seconds == 120.0 assert stored2.config.max_attempts == 5 assert stored2.config.retry_condition == ["timeout"] # Verify first rollout config unchanged stored1_again = await store_fixture.get_rollout_by_id(rollout.rollout_id) assert stored1_again is not None assert stored1_again.config.timeout_seconds == 30.0 # Queue Operations Tests @pytest.mark.asyncio async def test_dequeue_rollout_skips_non_queuing_status(store_fixture: LightningStore) -> None: """Test that dequeue_rollout skips rollouts that have been updated to non-queuing status.""" # Add multiple rollouts to the queue r1 = await store_fixture.enqueue_rollout(input={"id": 1}) r2 = await store_fixture.enqueue_rollout(input={"id": 2}) r3 = await store_fixture.enqueue_rollout(input={"id": 3}) # Update r1 to succeeded status while it's still in the queue await store_fixture.update_rollout(rollout_id=r1.rollout_id, status="succeeded") # Update r2 to failed status await store_fixture.update_rollout(rollout_id=r2.rollout_id, status="failed") # r3 should still be in queuing status # Pop should skip r1 and r2 (both non-queuing) and return r3 popped = await store_fixture.dequeue_rollout() assert popped is not None assert popped.rollout_id == r3.rollout_id assert popped.status == "preparing" assert popped.input["id"] == 3 # Second pop should return None since no queuing rollouts remain popped2 = await store_fixture.dequeue_rollout() assert popped2 is None # Verify r1 and r2 are still in their non-queuing states all_rollouts = await store_fixture.query_rollouts() rollout_statuses = {r.rollout_id: r.status for r in all_rollouts} assert rollout_statuses[r1.rollout_id] == "succeeded" assert rollout_statuses[r2.rollout_id] == "failed" assert rollout_statuses[r3.rollout_id] == "preparing" @pytest.mark.asyncio async def test_fifo_ordering(store_fixture: LightningStore) -> None: """Test that queue maintains FIFO order.""" rollouts: List[Rollout] = [] for i in range(5): r = await store_fixture.enqueue_rollout(input={"order": i}) rollouts.append(r) # Pop all and verify order for i in range(5): popped = await store_fixture.dequeue_rollout() assert popped is not None assert popped.rollout_id == rollouts[i].rollout_id assert popped.input["order"] == i assert popped.status == "preparing" @pytest.mark.asyncio async def test_pop_empty_queue(store_fixture: LightningStore) -> None: """Test popping from empty queue returns None.""" result = await store_fixture.dequeue_rollout() assert result is None # Multiple pops should all return None for _ in range(3): assert await store_fixture.dequeue_rollout() is None @pytest.mark.asyncio async def test_requeue_mechanism(store_fixture: LightningStore) -> None: """Test requeuing puts rollout back in queue.""" rollout = await store_fixture.enqueue_rollout(input={"data": "test"}) original_id = rollout.rollout_id # Pop and verify it's not in queue popped = await store_fixture.dequeue_rollout() assert popped is not None assert await store_fixture.dequeue_rollout() is None # Requeue it await store_fixture.update_rollout(rollout_id=original_id, status="requeuing") # Should be back in queue requeued = await store_fixture.dequeue_rollout() assert requeued is not None assert requeued.rollout_id == original_id assert requeued.status == "preparing" # Changes when popped # Check that a new attempt was created attempts = await store_fixture.query_attempts(requeued.rollout_id) assert len(attempts) == 2 # First attempt plus requeued attempt latest_attempt = await store_fixture.get_latest_attempt(requeued.rollout_id) assert latest_attempt is not None assert latest_attempt.status == "preparing" assert latest_attempt.sequence_id == 2 @pytest.mark.asyncio async def test_enqueue_many_rollouts_preserves_order(store_fixture: LightningStore) -> None: """enqueue_many_rollouts should enqueue tasks in the provided order with matching metadata.""" requests = [ EnqueueRolloutRequest(input={"idx": 0}, metadata={"batch": "a"}), EnqueueRolloutRequest(input={"idx": 1}, mode="train"), EnqueueRolloutRequest(input={"idx": 2}, config=RolloutConfig(timeout_seconds=3.5)), ] rollouts = await store_fixture.enqueue_many_rollouts(requests) assert [rollout.input["idx"] for rollout in rollouts] == [0, 1, 2] assert all(rollout.status == "queuing" for rollout in rollouts) assert rollouts[0].metadata == {"batch": "a"} assert rollouts[1].mode == "train" assert rollouts[2].config.timeout_seconds == 3.5 for expected_idx in range(3): dequeued = await store_fixture.dequeue_rollout() assert dequeued is not None assert dequeued.input["idx"] == expected_idx @pytest.mark.asyncio async def test_dequeue_many_rollouts_with_limit(store_fixture: LightningStore) -> None: """dequeue_many_rollouts should honor limit and propagate worker IDs to attempts.""" requests = [EnqueueRolloutRequest(input={"idx": idx}) for idx in range(4)] await store_fixture.enqueue_many_rollouts(requests) first_batch = await store_fixture.dequeue_many_rollouts(limit=2, worker_id="bulk-worker") assert len(first_batch) == 2 assert [attempt.input["idx"] for attempt in first_batch] == [0, 1] assert all(attempt.attempt.worker_id == "bulk-worker" for attempt in first_batch) second_batch = await store_fixture.dequeue_many_rollouts(limit=5, worker_id="bulk-worker") assert len(second_batch) == 2 assert [attempt.input["idx"] for attempt in second_batch] == [2, 3] assert await store_fixture.dequeue_many_rollouts(limit=1) == [] @pytest.mark.asyncio async def test_update_and_query_workers(store_fixture: LightningStore) -> None: """Workers can be created, heartbeats recorded, and telemetry auto-updated.""" first = await store_fixture.update_worker("worker-1", heartbeat_stats={"cpu": 0.5}) assert first.worker_id == "worker-1" assert first.heartbeat_stats == {"cpu": 0.5} assert isinstance(first.last_heartbeat_time, float) assert first.status == "unknown" rollout = await store_fixture.enqueue_rollout(input={"task": "work"}) claimed = await store_fixture.dequeue_rollout(worker_id="worker-1") assert claimed is not None assert claimed.rollout_id == rollout.rollout_id await store_fixture.update_attempt(claimed.rollout_id, claimed.attempt.attempt_id, worker_id="worker-1") busy = await store_fixture.get_worker_by_id("worker-1") assert busy is not None assert busy.status == "busy" assert busy.current_rollout_id == claimed.rollout_id assert busy.current_attempt_id == claimed.attempt.attempt_id assert isinstance(busy.last_dequeue_time, float) assert isinstance(busy.last_busy_time, float) heartbeat = await store_fixture.update_worker("worker-1") assert heartbeat.last_heartbeat_time is not None assert heartbeat.last_heartbeat_time >= first.last_heartbeat_time await store_fixture.update_attempt(claimed.rollout_id, claimed.attempt.attempt_id, status="succeeded") idle = await store_fixture.get_worker_by_id("worker-1") assert idle is not None assert idle.status == "idle" assert idle.current_rollout_id is None assert idle.current_attempt_id is None assert isinstance(idle.last_idle_time, float) workers = await store_fixture.query_workers() assert any(w.worker_id == "worker-1" for w in workers) assert await store_fixture.get_worker_by_id("missing") is None with pytest.raises(TypeError): await store_fixture.update_worker("worker-1", heartbeat_stats=None) # type: ignore[arg-type] @pytest.mark.asyncio async def test_worker_sync_preserves_existing_fields(store_fixture: LightningStore) -> None: """Worker heartbeats should survive status syncs triggered by attempts.""" worker_id = "worker-sync-preserve" initial = await store_fixture.update_worker(worker_id, heartbeat_stats={"cpu": 0.42}) assert initial.last_heartbeat_time is not None attempted = await store_fixture.start_rollout(input={"task": "preserve"}, worker_id=worker_id) busy = await store_fixture.get_worker_by_id(worker_id) assert busy is not None assert busy.heartbeat_stats == {"cpu": 0.42} assert busy.last_heartbeat_time == initial.last_heartbeat_time await store_fixture.update_attempt(attempted.rollout_id, attempted.attempt.attempt_id, status="succeeded") idle = await store_fixture.get_worker_by_id(worker_id) assert idle is not None assert idle.heartbeat_stats == {"cpu": 0.42} assert idle.last_heartbeat_time == initial.last_heartbeat_time @pytest.mark.asyncio async def test_start_rollout_assigns_worker(store_fixture: LightningStore) -> None: """start_rollout should immediately associate attempts with the provided worker.""" attempted = await store_fixture.start_rollout(input={"task": "direct"}, worker_id="worker-direct") assert attempted.attempt.worker_id == "worker-direct" worker = await store_fixture.get_worker_by_id("worker-direct") assert worker is not None assert worker.status == "busy" assert worker.current_rollout_id == attempted.rollout_id assert worker.current_attempt_id == attempted.attempt.attempt_id @pytest.mark.asyncio async def test_dequeue_rollout_assigns_worker(store_fixture: LightningStore) -> None: """dequeue_rollout should stamp attempts and worker telemetry with worker_id.""" await store_fixture.enqueue_rollout(input={"task": "queued"}) dequeued = await store_fixture.dequeue_rollout(worker_id="worker-dequeue") assert dequeued is not None assert dequeued.attempt.worker_id == "worker-dequeue" worker = await store_fixture.get_worker_by_id("worker-dequeue") assert worker is not None assert worker.status == "busy" assert worker.current_rollout_id == dequeued.rollout_id assert worker.current_attempt_id == dequeued.attempt.attempt_id @pytest.mark.asyncio async def test_start_attempt_assigns_worker(store_fixture: LightningStore) -> None: """Manual retries should also update worker state when worker_id is provided.""" initial = await store_fixture.start_rollout(input={"task": "retry-seed"}) retry = await store_fixture.start_attempt(initial.rollout_id, worker_id="worker-retry") assert retry.attempt.sequence_id == 2 assert retry.attempt.worker_id == "worker-retry" worker = await store_fixture.get_worker_by_id("worker-retry") assert worker is not None assert worker.status == "busy" assert worker.current_rollout_id == retry.rollout_id assert worker.current_attempt_id == retry.attempt.attempt_id @pytest.mark.asyncio async def test_query_workers_supports_filters(store_fixture: LightningStore) -> None: """Worker queries should support filtering, sorting, and pagination.""" await store_fixture.update_worker("alpha-worker", heartbeat_stats={"cpu": 0.2}) await store_fixture.update_worker("beta-worker", heartbeat_stats={"cpu": 0.8}) busy_rollout = await store_fixture.start_rollout(input={"worker": "alpha"}) await store_fixture.update_attempt( busy_rollout.rollout_id, busy_rollout.attempt.attempt_id, worker_id="alpha-worker", status="running", ) idle_rollout = await store_fixture.start_rollout(input={"worker": "beta"}) await store_fixture.update_attempt( idle_rollout.rollout_id, idle_rollout.attempt.attempt_id, worker_id="beta-worker", status="succeeded", ) busy_workers = await store_fixture.query_workers(status_in=["busy"]) assert [worker.worker_id for worker in busy_workers] == ["alpha-worker"] contains_beta = await store_fixture.query_workers(worker_id_contains="beta") assert [worker.worker_id for worker in contains_beta] == ["beta-worker"] sorted_workers = await store_fixture.query_workers(sort_by="worker_id", sort_order="desc") assert [worker.worker_id for worker in sorted_workers] == ["beta-worker", "alpha-worker"] paged = await store_fixture.query_workers(sort_by="worker_id", sort_order="desc", limit=1, offset=1) assert [worker.worker_id for worker in paged] == ["alpha-worker"] @pytest.mark.asyncio async def test_query_workers_filter_logic_or_combines_conditions(store_fixture: LightningStore) -> None: """filter_logic should dictate whether worker filters act as AND or OR.""" await store_fixture.update_worker("cpu-worker", heartbeat_stats={"cpu": 0.1}) busy_rollout = await store_fixture.start_rollout(input={"task": "busy"}) await store_fixture.update_attempt( busy_rollout.rollout_id, busy_rollout.attempt.attempt_id, worker_id="busy-worker", status="running", ) and_results = await store_fixture.query_workers(status_in=["busy"], worker_id_contains="cpu") assert len(and_results) == 0 or_results = await store_fixture.query_workers( status_in=["busy"], worker_id_contains="cpu", filter_logic="or", sort_by="worker_id", ) assert [worker.worker_id for worker in or_results] == ["busy-worker", "cpu-worker"] @pytest.mark.asyncio async def test_query_workers_filter_logic_and_with_sort(store_fixture: LightningStore) -> None: """Explicit AND logic should combine substring filters with status filters and respect sorting.""" await store_fixture.update_worker("alpha-worker", heartbeat_stats={"cpu": 0.1}) await store_fixture.update_worker("beta-worker", heartbeat_stats={"cpu": 0.1}) busy_rollout = await store_fixture.start_rollout(input={"task": "alpha"}) await store_fixture.update_attempt( busy_rollout.rollout_id, busy_rollout.attempt.attempt_id, worker_id="alpha-worker", status="running", ) idle_rollout = await store_fixture.start_rollout(input={"task": "beta"}) await store_fixture.update_attempt( idle_rollout.rollout_id, idle_rollout.attempt.attempt_id, worker_id="beta-worker", status="succeeded", ) filtered = await store_fixture.query_workers( status_in=["busy"], worker_id_contains="alpha", filter_logic="and", sort_by="worker_id", sort_order="asc", limit=1, ) assert [worker.worker_id for worker in filtered] == ["alpha-worker"] @pytest.mark.asyncio async def test_query_workers_reports_pagination_metadata(store_fixture: LightningStore) -> None: """Worker pagination should expose metadata for callers.""" for worker_id in ["worker-a", "worker-b", "worker-c"]: await store_fixture.update_worker(worker_id, heartbeat_stats={"cpu": 0.1}) paginated = await store_fixture.query_workers(sort_by="worker_id", sort_order="asc", limit=1, offset=1) assert isinstance(paginated, PaginatedResult) assert paginated.limit == 1 assert paginated.offset == 1 assert paginated.total == 3 # Resource Management Tests @pytest.mark.asyncio async def test_add_resources_generates_id_and_stores(store_fixture: LightningStore) -> None: """Test that add_resources generates a resources_id and stores the resources.""" # Initially no resources assert await store_fixture.get_latest_resources() is None # Add resources using add_resources (auto-generates ID) llm = LLM( resource_type="llm", endpoint="http://localhost:8080/v1", model="test-model", sampling_parameters={"temperature": 0.7}, ) prompt = PromptTemplate(resource_type="prompt_template", template="Hello {name}!", engine="f-string") resources_update = await store_fixture.add_resources({"main_llm": llm, "greeting": prompt}) # Verify resources_id was auto-generated with correct prefix assert resources_update.resources_id.startswith("rs-") assert len(resources_update.resources_id) == 15 # "rs-" + 12 char hash # Verify resources were stored correctly assert isinstance(resources_update.resources["main_llm"], LLM) assert resources_update.resources["main_llm"].model == "test-model" assert isinstance(resources_update.resources["greeting"], PromptTemplate) assert resources_update.resources["greeting"].template == "Hello {name}!" # Verify it's set as latest latest = await store_fixture.get_latest_resources() assert latest is not None assert latest.resources_id == resources_update.resources_id assert latest.resources["main_llm"].model == "test-model" # type: ignore # Verify we can retrieve by ID retrieved = await store_fixture.get_resources_by_id(resources_update.resources_id) assert retrieved is not None assert retrieved.resources_id == resources_update.resources_id @pytest.mark.asyncio async def test_add_resources_multiple_times_generates_unique_ids(store_fixture: LightningStore) -> None: """Test that multiple calls to add_resources generate unique IDs.""" llm1 = LLM(resource_type="llm", endpoint="http://localhost:8080", model="model-v1") llm2 = LLM(resource_type="llm", endpoint="http://localhost:8080", model="model-v2") update1 = await store_fixture.add_resources({"llm": llm1}) update2 = await store_fixture.add_resources({"llm": llm2}) # IDs should be different assert update1.resources_id != update2.resources_id assert update1.resources_id.startswith("rs-") assert update2.resources_id.startswith("rs-") # Both should be retrievable retrieved1 = await store_fixture.get_resources_by_id(update1.resources_id) retrieved2 = await store_fixture.get_resources_by_id(update2.resources_id) assert retrieved1 is not None assert retrieved2 is not None assert retrieved1.resources["llm"].model == "model-v1" # type: ignore assert retrieved2.resources["llm"].model == "model-v2" # type: ignore # Latest should be the second one latest = await store_fixture.get_latest_resources() assert latest is not None assert latest.resources_id == update2.resources_id @pytest.mark.asyncio async def test_query_resources_returns_history(store_fixture: LightningStore) -> None: """query_resources should list snapshots in the order they were stored.""" assert len(await store_fixture.query_resources()) == 0 first = await store_fixture.add_resources( { "llm": LLM(resource_type="llm", endpoint="http://localhost:8080", model="model-v1"), } ) second = await store_fixture.update_resources( "custom-snapshot", { "prompt": PromptTemplate(resource_type="prompt_template", template="Hi {name}", engine="f-string"), }, ) history = await store_fixture.query_resources() assert set([item.resources_id for item in history]) == {first.resources_id, second.resources_id} assert isinstance(history[0], ResourcesUpdate) assert isinstance(history[1], ResourcesUpdate) @pytest.mark.asyncio async def test_query_resources_supports_filters(store_fixture: LightningStore) -> None: """Resource queries should support substring filters and pagination.""" alpha = PromptTemplate(resource_type="prompt_template", template="alpha", engine="jinja") beta = PromptTemplate(resource_type="prompt_template", template="beta", engine="jinja") await store_fixture.update_resources("manual-alpha", {"prompt": alpha}) await store_fixture.update_resources("manual-beta", {"prompt": beta}) contains_beta = await store_fixture.query_resources(resources_id_contains="beta") assert [item.resources_id for item in contains_beta] == ["manual-beta"] sorted_ids = sorted(["manual-alpha", "manual-beta"], reverse=True) paged = await store_fixture.query_resources(sort_by="resources_id", sort_order="desc", limit=1) assert [item.resources_id for item in paged] == sorted_ids[:1] offset_item = await store_fixture.query_resources(sort_by="resources_id", sort_order="desc", limit=1, offset=1) assert [item.resources_id for item in offset_item] == sorted_ids[1:2] @pytest.mark.asyncio async def test_query_resources_combines_exact_and_contains_filters(store_fixture: LightningStore) -> None: """Exact and substring filters should be usable together.""" alpha = PromptTemplate(resource_type="prompt_template", template="alpha", engine="jinja") beta = PromptTemplate(resource_type="prompt_template", template="beta", engine="jinja") await store_fixture.update_resources("manual-alpha", {"prompt": alpha}) await store_fixture.update_resources("manual-beta", {"prompt": beta}) results = await store_fixture.query_resources(resources_id="manual-alpha", resources_id_contains="manual") assert [item.resources_id for item in results] == ["manual-alpha"] @pytest.mark.asyncio async def test_query_resources_offset_beyond_range_returns_empty(store_fixture: LightningStore) -> None: """Large offsets should simply return an empty result.""" await store_fixture.update_resources( "snapshot-a", {"prompt": PromptTemplate(resource_type="prompt_template", template="a", engine="jinja")}, ) await store_fixture.update_resources( "snapshot-b", {"prompt": PromptTemplate(resource_type="prompt_template", template="b", engine="jinja")}, ) results = await store_fixture.query_resources(sort_by="resources_id", limit=1, offset=5) assert len(results) == 0 @pytest.mark.asyncio async def test_query_resources_contains_with_sort_and_pagination(store_fixture: LightningStore) -> None: """Substring filters should combine with sort order, limit, and offset.""" for suffix in ["alpha", "beta", "gamma"]: await store_fixture.update_resources( f"manual-{suffix}", {"prompt": PromptTemplate(resource_type="prompt_template", template=suffix, engine="jinja")}, ) filtered = await store_fixture.query_resources( resources_id_contains="manual-", sort_by="resources_id", sort_order="asc", limit=2, offset=1, ) expected_ids = sorted([f"manual-{suffix}" for suffix in ["alpha", "beta", "gamma"]]) assert [item.resources_id for item in filtered] == expected_ids[1:3] @pytest.mark.asyncio async def test_query_resources_reports_pagination_metadata(store_fixture: LightningStore) -> None: """Resource pagination should expose metadata fields.""" await store_fixture.update_resources( "snapshot-a", {"prompt": PromptTemplate(resource_type="prompt_template", template="a", engine="jinja")}, ) await store_fixture.update_resources( "snapshot-b", {"prompt": PromptTemplate(resource_type="prompt_template", template="b", engine="jinja")}, ) paginated = await store_fixture.query_resources(sort_by="resources_id", limit=1, offset=1) assert isinstance(paginated, PaginatedResult) assert paginated.limit == 1 assert paginated.offset == 1 assert paginated.total == 2 @pytest.mark.asyncio async def test_resource_lifecycle(store_fixture: LightningStore) -> None: """Test adding, updating, and retrieving resources.""" # Initially no resources assert await store_fixture.get_latest_resources() is None assert await store_fixture.get_resources_by_id("any-id") is None # Add first version with proper LLM resource llm_v1 = LLM( resource_type="llm", endpoint="http://localhost:8080/v1", model="test-model-v1", sampling_parameters={"temperature": 0.7}, ) update = await store_fixture.update_resources("v1", {"main_llm": llm_v1}) assert update.resources_id == "v1" latest = await store_fixture.get_latest_resources() assert latest is not None assert latest.resources_id == "v1" assert isinstance(latest.resources["main_llm"], LLM) assert latest.resources["main_llm"].model == "test-model-v1" # Add second version with different LLM llm_v2 = LLM( resource_type="llm", endpoint="http://localhost:8080/v2", model="test-model-v2", sampling_parameters={"temperature": 0.8}, ) v2 = await store_fixture.update_resources("v2", {"main_llm": llm_v2}) assert v2.resources_id == "v2" assert isinstance(v2.resources["main_llm"], LLM) assert v2.resources["main_llm"].model == "test-model-v2" # Latest should be v2 latest = await store_fixture.get_latest_resources() assert latest is not None assert latest.resources_id == "v2" # Can still retrieve v1 old = await store_fixture.get_resources_by_id("v1") assert old is not None assert isinstance(old.resources["main_llm"], LLM) assert old.resources["main_llm"].model == "test-model-v1" @pytest.mark.asyncio async def test_latest_resources_rehydrates_cache(store_fixture: LightningStore) -> None: """get_latest_resources should consult storage even if the cache is unset.""" llm = LLM( resource_type="llm", endpoint="http://localhost:8080/v1", model="cache-model", sampling_parameters={"temperature": 0.1}, ) update = await store_fixture.update_resources("cache-test", {"main_llm": llm}) # Simulate a fresh process by clearing the cache. store_fixture._latest_resources_id = UNSET # type: ignore[attr-defined] latest = await store_fixture.get_latest_resources() assert latest is not None assert latest.resources_id == update.resources_id assert latest.resources["main_llm"].model == "cache-model" # type: ignore @pytest.mark.asyncio async def test_task_inherits_latest_resources(store_fixture: LightningStore) -> None: """Test that new tasks inherit latest resources_id if not specified.""" # Set up resources with proper PromptTemplate prompt = PromptTemplate(resource_type="prompt_template", template="Hello {name}!", engine="f-string") update = ResourcesUpdate( resources_id="current", resources={"greeting": prompt}, create_time=time.time(), update_time=time.time(), version=1, ) await store_fixture.update_resources(update.resources_id, update.resources) # Task without explicit resources_id r1 = await store_fixture.enqueue_rollout(input={"id": 1}) assert r1.resources_id == "current" # Task with explicit resources_id r2 = await store_fixture.enqueue_rollout(input={"id": 2}, resources_id="override") assert r2.resources_id == "override" # Update resources new_prompt = PromptTemplate(resource_type="prompt_template", template="Hi {name}!", engine="f-string") update2 = ResourcesUpdate( resources_id="new", resources={"greeting": new_prompt}, create_time=time.time(), update_time=time.time(), version=1, ) await store_fixture.update_resources(update2.resources_id, update2.resources) # New task gets new resources r3 = await store_fixture.enqueue_rollout(input={"id": 3}) assert r3.resources_id == "new" # Span Management Tests @pytest.mark.asyncio async def test_span_sequence_generation(store_fixture: LightningStore, mock_readable_span: Mock) -> None: """Test automatic sequence ID generation for spans.""" rollout = await store_fixture.enqueue_rollout(input={"test": "data"}) # Pop to create an attempt await store_fixture.dequeue_rollout() attempts = await store_fixture.query_attempts(rollout.rollout_id) attempt_id = attempts[0].attempt_id # First span gets sequence_id 1 seq_id = await store_fixture.get_next_span_sequence_id(rollout.rollout_id, attempt_id) assert seq_id == 1 span1 = await store_fixture.add_otel_span(rollout.rollout_id, attempt_id, mock_readable_span) assert span1 is not None and span1.sequence_id == 2 # Next span gets sequence_id 3 seq_id = await store_fixture.get_next_span_sequence_id(rollout.rollout_id, attempt_id) assert seq_id == 3 span2 = await store_fixture.add_otel_span(rollout.rollout_id, attempt_id, mock_readable_span) assert span2 is not None and span2.sequence_id == 4 # Different attempt reuses the same rollout_id seq_id = await store_fixture.get_next_span_sequence_id(rollout.rollout_id, "attempt-does-not-exist") assert seq_id == 5 @pytest.mark.asyncio async def test_get_many_span_sequence_ids_handles_mixed_batches(store_fixture: LightningStore) -> None: """Bulk sequence allocation should work across mixed rollouts and attempts.""" first = await store_fixture.start_rollout(input={"origin": "bulk-seq"}) second = await store_fixture.start_rollout(input={"origin": "bulk-seq-2"}) await store_fixture.update_rollout(first.rollout_id, status="requeuing") retried = await store_fixture.start_attempt(first.rollout_id) sequence_pairs = [ (first.rollout_id, first.attempt.attempt_id), (second.rollout_id, second.attempt.attempt_id), (first.rollout_id, retried.attempt.attempt_id), (first.rollout_id, first.attempt.attempt_id), ] sequence_ids = await store_fixture.get_many_span_sequence_ids(sequence_pairs) assert sequence_ids == [1, 1, 2, 3] next_first = await store_fixture.get_next_span_sequence_id(first.rollout_id, first.attempt.attempt_id) next_second = await store_fixture.get_next_span_sequence_id(second.rollout_id, second.attempt.attempt_id) assert next_first == 4 assert next_second == 2 @pytest.mark.asyncio async def test_add_many_spans_handles_mixed_rollouts_and_attempts(store_fixture: LightningStore) -> None: """Batch span insertion should handle mixed rollouts and skip duplicates.""" first = await store_fixture.start_rollout(input={"origin": "batch-spans"}) second = await store_fixture.start_rollout(input={"origin": "batch-spans-2"}) await store_fixture.update_rollout(first.rollout_id, status="requeuing") retried = await store_fixture.start_attempt(first.rollout_id) def _build_span(seed: int, rollout_id: str, attempt_id: str) -> Span: trace_hex = f"{seed:032x}" span_hex = f"{seed:016x}" return Span( rollout_id=rollout_id, attempt_id=attempt_id, sequence_id=seed, trace_id=trace_hex, span_id=span_hex, parent_id=None, name=f"batch-span-{seed}", status=TraceStatus(status_code="OK"), attributes={}, events=[], links=[], start_time=None, end_time=None, context=SpanContext(trace_id=trace_hex, span_id=span_hex, is_remote=False, trace_state={}), parent=None, resource=OtelResource(attributes={}, schema_url=""), ) span_first = _build_span(1, first.rollout_id, first.attempt.attempt_id) span_retry = _build_span(2, retried.rollout_id, retried.attempt.attempt_id) span_second = _build_span(3, second.rollout_id, second.attempt.attempt_id) duplicate_first = _build_span(1, first.rollout_id, first.attempt.attempt_id) stored = await store_fixture.add_many_spans([span_first, span_retry, span_second, duplicate_first]) if isinstance(store_fixture, InMemoryLightningStore): assert {span.span_id for span in stored} == {span_first.span_id, span_retry.span_id, span_second.span_id} spans_first = await store_fixture.query_spans(first.rollout_id) assert {span.span_id for span in spans_first} >= {span_first.span_id, span_retry.span_id} spans_second = await store_fixture.query_spans(second.rollout_id) assert {span.span_id for span in spans_second} >= {span_second.span_id} @pytest.mark.asyncio async def test_add_span_returns_none_on_duplicate(store_fixture: LightningStore) -> None: """Duplicate span IDs should return None instead of raising.""" attempted = await store_fixture.start_rollout(input={"origin": "duplicate-span"}) attempt_id = attempted.attempt.attempt_id span = Span( rollout_id=attempted.rollout_id, attempt_id=attempt_id, sequence_id=1, trace_id="0" * 32, span_id="0" * 16, parent_id=None, name="dup-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=OtelResource(attributes={}, schema_url=""), ) assert await store_fixture.add_span(span) is span assert await store_fixture.add_span(span) is None @pytest.mark.asyncio async def test_span_updates_attempt_status(store_fixture: LightningStore, mock_readable_span: Mock) -> None: """Adding a span should persist attempt heartbeat and transition status to running.""" rollout = await store_fixture.enqueue_rollout(input={"test": "attempt-status"}) await store_fixture.dequeue_rollout() attempts = await store_fixture.query_attempts(rollout.rollout_id) assert attempts attempt_id = attempts[0].attempt_id assert attempts[0].status == "preparing" assert attempts[0].last_heartbeat_time is None await store_fixture.add_otel_span(rollout.rollout_id, attempt_id, mock_readable_span) updated_attempt = (await store_fixture.query_attempts(rollout.rollout_id))[0] assert updated_attempt.status == "running" assert updated_attempt.last_heartbeat_time is not None @pytest.mark.asyncio async def test_spans_promote_preparing_attempt_with_heartbeat( store_fixture: LightningStore, mock_readable_span: Mock, monkeypatch: pytest.MonkeyPatch ) -> None: """Spans should set heartbeat time and promote preparing attempts/rollouts to running.""" rollout = await store_fixture.enqueue_rollout(input={"test": "preparing-heartbeat"}) dequeued = await store_fixture.dequeue_rollout() assert dequeued is not None attempts_before = await store_fixture.query_attempts(rollout.rollout_id) assert attempts_before attempt_id = attempts_before[0].attempt_id assert attempts_before[0].status == "preparing" assert attempts_before[0].last_heartbeat_time is None heartbeat_time = 1234.5 monkeypatch.setattr("agentlightning.store.collection_based.time.time", lambda: heartbeat_time) await store_fixture.add_otel_span(rollout.rollout_id, attempt_id, mock_readable_span) attempt_after = (await store_fixture.query_attempts(rollout.rollout_id))[0] assert attempt_after.status == "running" assert attempt_after.last_heartbeat_time == heartbeat_time rollout_after = await store_fixture.get_rollout_by_id(rollout.rollout_id) assert rollout_after is not None assert rollout_after.status == "running" @pytest.mark.asyncio async def test_unresponsive_attempt_recovers_after_span( store_fixture: LightningStore, mock_readable_span: Mock ) -> None: """Spans arriving for an unresponsive attempt should mark it running again.""" rollout = await store_fixture.enqueue_rollout(input={"test": "unresponsive"}) dequeued = await store_fixture.dequeue_rollout() assert dequeued is not None attempt_id = dequeued.attempt.attempt_id await store_fixture.update_attempt( rollout_id=rollout.rollout_id, attempt_id=attempt_id, status="unresponsive", ) attempt_before = (await store_fixture.query_attempts(rollout.rollout_id))[0] assert attempt_before.status == "unresponsive" await store_fixture.add_otel_span(rollout.rollout_id, attempt_id, mock_readable_span) attempt_after = (await store_fixture.query_attempts(rollout.rollout_id))[0] assert attempt_after.status == "running" assert attempt_after.last_heartbeat_time is not None @pytest.mark.asyncio async def test_running_attempt_updates_heartbeat( store_fixture: LightningStore, mock_readable_span: Mock, monkeypatch: pytest.MonkeyPatch ) -> None: """Adding spans to an already running attempt should advance its heartbeat.""" rollout = await store_fixture.enqueue_rollout(input={"test": "running-heartbeat"}) await store_fixture.dequeue_rollout() attempt_id = (await store_fixture.query_attempts(rollout.rollout_id))[0].attempt_id await store_fixture.add_otel_span(rollout.rollout_id, attempt_id, mock_readable_span) attempt_after_first = (await store_fixture.query_attempts(rollout.rollout_id))[0] assert attempt_after_first.status == "running" first_heartbeat = attempt_after_first.last_heartbeat_time assert first_heartbeat is not None monkeypatch.setattr("agentlightning.store.collection_based.time.time", lambda: first_heartbeat + 100.0) await store_fixture.add_otel_span(rollout.rollout_id, attempt_id, mock_readable_span) attempt_after_second = (await store_fixture.query_attempts(rollout.rollout_id))[0] assert attempt_after_second.last_heartbeat_time == first_heartbeat + 100.0 @pytest.mark.asyncio async def test_span_post_add_preserves_concurrent_updates( store_fixture: LightningStore, mock_readable_span: Mock, monkeypatch: pytest.MonkeyPatch ) -> None: """Concurrent attempt updates should not be clobbered when spans record heartbeats.""" rollout = await store_fixture.enqueue_rollout(input={"test": "span-concurrency"}) dequeued = await store_fixture.dequeue_rollout() assert dequeued is not None attempt_id = dequeued.attempt.attempt_id original_post = store_fixture._post_add_spans # type: ignore async def patched_post(spans: List[Span], rollout_id: str, mutated_attempt_id: str) -> None: await store_fixture.update_attempt(rollout_id, mutated_attempt_id, metadata={"concurrent": True}) await original_post(spans, rollout_id, mutated_attempt_id) monkeypatch.setattr(store_fixture, "_post_add_spans", patched_post) await store_fixture.add_otel_span(rollout.rollout_id, attempt_id, mock_readable_span) attempt_after = (await store_fixture.query_attempts(rollout.rollout_id))[0] assert attempt_after.metadata == {"concurrent": True} @pytest.mark.asyncio async def test_healthcheck_marks_unresponsive_and_updates_worker( store_fixture: LightningStore, monkeypatch: pytest.MonkeyPatch ) -> None: """Healthcheck should mark attempts unresponsive and sync worker state via the helper.""" class TimeStub: def __init__(self, value: float): self.value = value def __call__(self) -> float: return self.value time_stub = TimeStub(100.0) monkeypatch.setattr("agentlightning.store.collection_based.time.time", time_stub) rollout = await store_fixture.enqueue_rollout( input={"test": "healthcheck-worker"}, config=RolloutConfig(unresponsive_seconds=1.0), ) dequeued = await store_fixture.dequeue_rollout(worker_id="worker-sync") assert dequeued is not None attempt_id = dequeued.attempt.attempt_id await store_fixture.update_attempt(rollout.rollout_id, attempt_id, worker_id="worker-sync") original_sync = store_fixture._sync_workers_with_attempts # type: ignore sync_calls: List[str] = [] async def tracking_sync(attempts: Sequence[Attempt]) -> None: for attempt in attempts: sync_calls.append(attempt.attempt_id) await original_sync(attempts) monkeypatch.setattr(store_fixture, "_sync_workers_with_attempts", tracking_sync) time_stub.value = 105.0 await store_fixture.get_rollout_by_id(rollout.rollout_id) worker = await store_fixture.get_worker_by_id("worker-sync") assert worker is not None assert worker.status == "unknown" attempt_after = (await store_fixture.query_attempts(rollout.rollout_id))[0] assert attempt_after.status == "unresponsive" assert sync_calls == [attempt_after.attempt_id] @pytest.mark.asyncio async def test_duplicate_span_id_error( store_fixture: LightningStore, mock_readable_span: Mock, caplog: pytest.LogCaptureFixture ) -> None: """Adding two spans with the same span_id should raise a ValueError.""" caplog.set_level(logging.ERROR) rollout = await store_fixture.enqueue_rollout(input={"test": "data"}) await store_fixture.dequeue_rollout() attempts = await store_fixture.query_attempts(rollout.rollout_id) attempt_id = attempts[0].attempt_id # Force the mock to reuse the same span context for every call. fixed_context = mock_readable_span.get_span_context() mock_readable_span.get_span_context = Mock(return_value=fixed_context) await store_fixture.add_otel_span(rollout.rollout_id, attempt_id, mock_readable_span) duplicate = await store_fixture.add_otel_span(rollout.rollout_id, attempt_id, mock_readable_span) assert duplicate is None assert "Duplicated span added" in caplog.text @pytest.mark.asyncio async def test_span_with_explicit_sequence_id(store_fixture: LightningStore, mock_readable_span: Mock) -> None: """Test providing explicit sequence_id to spans.""" rollout = await store_fixture.enqueue_rollout(input={"test": "data"}) # Pop to create an attempt await store_fixture.dequeue_rollout() attempts = await store_fixture.query_attempts(rollout.rollout_id) attempt_id = attempts[0].attempt_id # Add span with explicit sequence_id span = await store_fixture.add_otel_span(rollout.rollout_id, attempt_id, mock_readable_span, sequence_id=100) assert span is not None and span.sequence_id == 100 next_seq = await store_fixture.get_next_span_sequence_id(rollout.rollout_id, attempt_id) assert next_seq == 101 @pytest.mark.asyncio async def test_query_spans_by_attempt(store_fixture: LightningStore, mock_readable_span: Mock) -> None: """Test querying spans filtered by attempt_id.""" rollout = await store_fixture.enqueue_rollout(input={"test": "data"}) # Pop to create first attempt await store_fixture.dequeue_rollout() attempts = await store_fixture.query_attempts(rollout.rollout_id) attempt1_id = attempts[0].attempt_id # Add spans for first attempt for _ in range(2): await store_fixture.add_otel_span(rollout.rollout_id, attempt1_id, mock_readable_span) # Simulate requeue and create second attempt await store_fixture.update_rollout(rollout_id=rollout.rollout_id, status="requeuing") await store_fixture.dequeue_rollout() attempts = await store_fixture.query_attempts(rollout.rollout_id) attempt2_id = attempts[1].attempt_id # Add spans for second attempt for _ in range(3): await store_fixture.add_otel_span(rollout.rollout_id, attempt2_id, mock_readable_span) # Query all spans all_spans = await store_fixture.query_spans(rollout.rollout_id) assert len(all_spans) == 5 # Query specific attempt attempt1_spans = await store_fixture.query_spans(rollout.rollout_id, attempt_id=attempt1_id) assert len(attempt1_spans) == 2 assert all(s.attempt_id == attempt1_id for s in attempt1_spans) # Query latest attempt latest_spans = await store_fixture.query_spans(rollout.rollout_id, attempt_id="latest") assert len(latest_spans) == 3 assert all(s.attempt_id == attempt2_id for s in latest_spans) # Query non-existent attempt no_spans = await store_fixture.query_spans(rollout.rollout_id, attempt_id="nonexistent") assert len(no_spans) == 0 @pytest.mark.asyncio async def test_query_spans_supports_filters(store_fixture: LightningStore) -> None: """Span queries should honor filtering logic and pagination.""" attempted = await store_fixture.start_rollout(input={"payload": "span-filters"}) attempt_id = attempted.attempt.attempt_id def build_span(idx: int, *, name: str, parent: Optional[str]) -> Span: trace_hex = f"{idx:032x}" span_hex = f"{idx:016x}" return Span( rollout_id=attempted.rollout_id, attempt_id=attempt_id, sequence_id=idx, trace_id=trace_hex, span_id=span_hex, parent_id=parent, name=name, status=TraceStatus(status_code="OK"), attributes={}, events=[Event(name=f"event-{idx}", attributes={})], links=[], start_time=None, end_time=None, context=SpanContext(trace_id=trace_hex, span_id=span_hex, is_remote=False, trace_state={}), parent=None, resource=OtelResource(attributes={}, schema_url=""), ) created_spans = [ build_span(1, name="reward", parent=None), build_span(2, name="planner", parent=f"{1:016x}"), build_span(3, name="tool-call", parent=f"{2:016x}"), ] for span in created_spans: await store_fixture.add_span(span) trace_filtered = await store_fixture.query_spans(attempted.rollout_id, trace_id=created_spans[1].trace_id) assert [s.span_id for s in trace_filtered] == [created_spans[1].span_id] or_filtered = await store_fixture.query_spans( attempted.rollout_id, span_id=created_spans[0].span_id, trace_id_contains=created_spans[2].trace_id[-4:], filter_logic="or", ) assert {s.span_id for s in or_filtered} == {created_spans[0].span_id, created_spans[2].span_id} parent_filtered = await store_fixture.query_spans( attempted.rollout_id, parent_id_contains=created_spans[1].span_id[-4:], ) assert [s.span_id for s in parent_filtered] == [created_spans[2].span_id] sorted_ids = sorted([span.span_id for span in created_spans], reverse=True) paged = await store_fixture.query_spans( attempted.rollout_id, sort_by="span_id", sort_order="desc", limit=1, offset=1, ) assert [span.span_id for span in paged] == sorted_ids[1:2] @pytest.mark.asyncio async def test_query_spans_filter_logic_respects_rollout_scope(store_fixture: LightningStore) -> None: """Even with OR logic, query_spans should not leak spans from other rollouts.""" first = await store_fixture.start_rollout(input={"payload": "first-span"}) second = await store_fixture.start_rollout(input={"payload": "second-span"}) def make_span(rollout_id: str, attempt_id: str, idx: int, name: str) -> Span: span_hex = f"{idx:016x}" trace_hex = f"{idx:032x}" return Span( rollout_id=rollout_id, attempt_id=attempt_id, sequence_id=idx, trace_id=trace_hex, span_id=span_hex, parent_id=None, name=name, status=TraceStatus(status_code="OK"), attributes={}, events=[], links=[], start_time=None, end_time=None, context=SpanContext(trace_id=trace_hex, span_id=span_hex, is_remote=False, trace_state={}), parent=None, resource=OtelResource(attributes={}, schema_url=""), ) span_first = make_span(first.rollout_id, first.attempt.attempt_id, 1, "alpha") span_second = make_span(second.rollout_id, second.attempt.attempt_id, 2, "beta") await store_fixture.add_span(span_first) await store_fixture.add_span(span_second) results = await store_fixture.query_spans( first.rollout_id, attempt_id=first.attempt.attempt_id, name_contains="alpha", trace_id=span_second.trace_id, filter_logic="or", ) assert [span.span_id for span in results] == [span_first.span_id] @pytest.mark.asyncio async def test_query_spans_supports_name_contains_filter(store_fixture: LightningStore) -> None: """name_contains filtering should happen before pagination.""" attempted = await store_fixture.start_rollout(input={"payload": "span-names"}) attempt_id = attempted.attempt.attempt_id def build_span(idx: int, name: str) -> Span: span_hex = f"{idx:016x}" trace_hex = f"{idx:032x}" return Span( rollout_id=attempted.rollout_id, attempt_id=attempt_id, sequence_id=idx, trace_id=trace_hex, span_id=span_hex, parent_id=None, name=name, status=TraceStatus(status_code="OK"), attributes={}, events=[], links=[], start_time=None, end_time=None, context=SpanContext(trace_id=trace_hex, span_id=span_hex, is_remote=False, trace_state={}), parent=None, resource=OtelResource(attributes={}, schema_url=""), ) matching = build_span(1, "planner-step") non_matching = build_span(2, "tool-call") await store_fixture.add_span(matching) await store_fixture.add_span(non_matching) results = await store_fixture.query_spans( attempted.rollout_id, name_contains="plan", sort_by="sequence_id", limit=1, ) assert [span.span_id for span in results] == [matching.span_id] @pytest.mark.asyncio async def test_query_spans_multiple_filters_require_all(store_fixture: LightningStore) -> None: """Using multiple exact/substring filters together should narrow down to a single span.""" attempted = await store_fixture.start_rollout(input={"payload": "span-multi"}) attempt_id = attempted.attempt.attempt_id def build_span(seq: int, parent: Optional[str], name: str) -> Span: span_hex = f"{seq:016x}" trace_hex = f"{(seq * 10):032x}" return Span( rollout_id=attempted.rollout_id, attempt_id=attempt_id, sequence_id=seq, trace_id=trace_hex, span_id=span_hex, parent_id=parent, name=name, status=TraceStatus(status_code="OK"), attributes={}, events=[], links=[], start_time=None, end_time=None, context=SpanContext(trace_id=trace_hex, span_id=span_hex, is_remote=False, trace_state={}), parent=None, resource=OtelResource(attributes={}, schema_url=""), ) spans = [ build_span(1, None, "phase-plan"), build_span(2, f"{1:016x}", "phase-run"), build_span(3, f"{2:016x}", "tool-call"), ] for span in spans: await store_fixture.add_span(span) target = spans[2] filtered = await store_fixture.query_spans( attempted.rollout_id, trace_id_contains=target.trace_id[-4:], span_id_contains=target.span_id[-4:], parent_id=target.parent_id, name_contains="tool", ) assert [span.span_id for span in filtered] == [target.span_id] @pytest.mark.asyncio async def test_query_spans_reports_pagination_metadata(store_fixture: LightningStore) -> None: """Span pagination should return limit/offset/total values.""" attempted = await store_fixture.start_rollout(input={"payload": "span-pagination"}) attempt_id = attempted.attempt.attempt_id def build_span(idx: int) -> Span: span_hex = f"{idx:016x}" trace_hex = f"{idx:032x}" return Span( rollout_id=attempted.rollout_id, attempt_id=attempt_id, sequence_id=idx, trace_id=trace_hex, span_id=span_hex, parent_id=None, name=f"span-{idx}", status=TraceStatus(status_code="OK"), attributes={}, events=[], links=[], start_time=None, end_time=None, context=SpanContext(trace_id=trace_hex, span_id=span_hex, is_remote=False, trace_state={}), parent=None, resource=OtelResource(attributes={}, schema_url=""), ) for idx in range(1, 4): await store_fixture.add_span(build_span(idx)) paginated = await store_fixture.query_spans( attempted.rollout_id, attempt_id=attempt_id, sort_by="sequence_id", sort_order="asc", limit=1, offset=1, ) assert isinstance(paginated, PaginatedResult) assert paginated.limit == 1 assert paginated.offset == 1 assert paginated.total == 3 @pytest.mark.asyncio async def test_query_spans_or_filters_with_sort_and_offset(store_fixture: LightningStore) -> None: """OR logic combined with sort + pagination should return deterministic slices.""" attempted = await store_fixture.start_rollout(input={"payload": "span-or"}) attempt_id = attempted.attempt.attempt_id def build_span(seq: int, name: str, parent: Optional[str]) -> Span: span_hex = f"{seq:016x}" trace_hex = f"{(seq * 3):032x}" return Span( rollout_id=attempted.rollout_id, attempt_id=attempt_id, sequence_id=seq, trace_id=trace_hex, span_id=span_hex, parent_id=parent, name=name, status=TraceStatus(status_code="OK"), attributes={}, events=[], links=[], start_time=None, end_time=None, context=SpanContext(trace_id=trace_hex, span_id=span_hex, is_remote=False, trace_state={}), parent=None, resource=OtelResource(attributes={}, schema_url=""), ) span_plan = build_span(10, "phase-plan", None) span_run = build_span(20, "phase-run", span_plan.span_id) span_tool = build_span(30, "tool-call", span_run.span_id) for span in [span_plan, span_run, span_tool]: await store_fixture.add_span(span) filtered = await store_fixture.query_spans( attempted.rollout_id, name_contains="phase", parent_id_contains=span_run.span_id[-4:], span_id=span_tool.span_id, filter_logic="or", sort_by="sequence_id", sort_order="desc", limit=1, offset=1, ) assert [span.span_id for span in filtered] == [span_run.span_id] @pytest.mark.asyncio async def test_span_eviction_removes_oldest_rollouts(mock_readable_span: Mock, monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr("agentlightning.store.memory._detect_total_memory_bytes", lambda: 100) store = InMemoryLightningStore( eviction_memory_threshold=0.5, safe_memory_threshold=0.05, span_size_estimator=lambda span: 20, ) attempted_rollouts: List[AttemptedRollout] = [] for index in range(4): attempted = await store.start_rollout(input={"index": index}) attempted_rollouts.append(attempted) await store.add_otel_span(attempted.rollout_id, attempted.attempt.attempt_id, mock_readable_span) for attempted in attempted_rollouts[:3]: with pytest.raises(RuntimeError): await store.query_spans(attempted.rollout_id) remaining_spans = await store.query_spans(attempted_rollouts[3].rollout_id) assert len(remaining_spans) == 1 assert remaining_spans[0].rollout_id == attempted_rollouts[3].rollout_id def test_memory_threshold_accepts_byte_values() -> None: store = InMemoryLightningStore( eviction_memory_threshold=150, safe_memory_threshold=20, ) assert store._eviction_threshold_bytes == 150 # pyright: ignore[reportPrivateUsage] assert store._safe_threshold_bytes == 20 # pyright: ignore[reportPrivateUsage] def test_memory_threshold_accepts_ratios_with_zero_safe(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr("agentlightning.store.memory._detect_total_memory_bytes", lambda: 200) store = InMemoryLightningStore( eviction_memory_threshold=0.6, safe_memory_threshold=0.0, ) assert store._eviction_threshold_bytes == int(200 * 0.6) # pyright: ignore[reportPrivateUsage] assert store._safe_threshold_bytes == 0 # pyright: ignore[reportPrivateUsage] def test_invalid_safe_threshold_raises_value_error() -> None: with pytest.raises(ValueError): InMemoryLightningStore( eviction_memory_threshold=50, safe_memory_threshold=100, ) def test_estimate_model_size_counts_nested_models() -> None: class Inner(BaseModel): value: int data: List[int] class Outer(BaseModel): inner: Inner mapping: dict[str, str] tags: List[str] inner = Inner(value=7, data=[1, 2, 3]) outer = Outer(inner=inner, mapping={"alpha": "beta"}, tags=["x", "yz"]) inner_expected = ( sys.getsizeof(inner) + sys.getsizeof(inner.value) + sys.getsizeof(inner.data) + sum(sys.getsizeof(item) for item in inner.data) ) assert estimate_model_size(inner) == inner_expected mapping_expected = sys.getsizeof(outer.mapping) + sum(sys.getsizeof(v) for v in outer.mapping.values()) tags_expected = sys.getsizeof(outer.tags) + sum(sys.getsizeof(tag) for tag in outer.tags) outer_expected = sys.getsizeof(outer) + inner_expected + mapping_expected + tags_expected assert estimate_model_size(outer) == outer_expected def test_estimate_model_size_handles_span_objects() -> None: status = TraceStatus(status_code="OK", description="fine") context = SpanContext(trace_id="trace", span_id="parent", is_remote=False, trace_state={"foo": "bar"}) event = Event(name="step", attributes={"detail": "value"}, timestamp=1.0) link = Link(context=context, attributes=None) resource = OtelResource(attributes={"service.name": "unit"}, schema_url="schema") span = Span( rollout_id="ro-1", attempt_id="at-1", sequence_id=1, trace_id="trace", span_id="span", parent_id=None, name="operation", status=status, attributes={"foo": "bar", "answer": 42}, events=[event], links=[link], start_time=1.0, end_time=2.0, context=None, parent=None, resource=resource, ) status_expected = sys.getsizeof(status) + sys.getsizeof(status.status_code) + sys.getsizeof(status.description) trace_state_values = context.trace_state.values() context_expected = ( sys.getsizeof(context) + sys.getsizeof(context.trace_id) + sys.getsizeof(context.span_id) + sys.getsizeof(context.is_remote) + sys.getsizeof(context.trace_state) + sum(sys.getsizeof(v) for v in trace_state_values) ) event_attributes_expected = sys.getsizeof(event.attributes) + sys.getsizeof("value") event_expected = ( sys.getsizeof(event) + sys.getsizeof(event.name) + event_attributes_expected + sys.getsizeof(event.timestamp) ) events_expected = sys.getsizeof(span.events) + event_expected link_attributes = cast(Optional[dict[str, str]], link.attributes) link_attribute_values = link_attributes.values() if link_attributes is not None else () link_attributes_expected = sys.getsizeof(link_attributes if link_attributes is not None else None) + sum( sys.getsizeof(v) for v in link_attribute_values ) link_expected = sys.getsizeof(link) + context_expected + link_attributes_expected links_expected = sys.getsizeof(span.links) + link_expected attributes_expected = ( sys.getsizeof(span.attributes) + sys.getsizeof("bar") + sys.getsizeof(span.attributes["answer"]) ) resource_expected = ( sys.getsizeof(resource) + sys.getsizeof(resource.attributes) + sum(sys.getsizeof(v) for v in resource.attributes.values()) + sys.getsizeof(resource.schema_url) ) expected_size = ( sys.getsizeof(span) + sys.getsizeof(span.rollout_id) + sys.getsizeof(span.attempt_id) + sys.getsizeof(span.sequence_id) + sys.getsizeof(span.trace_id) + sys.getsizeof(span.span_id) + sys.getsizeof(span.parent_id) + sys.getsizeof(span.name) + status_expected + attributes_expected + events_expected + links_expected + sys.getsizeof(span.start_time) + sys.getsizeof(span.end_time) + sys.getsizeof(span.context) + sys.getsizeof(span.parent) + resource_expected ) assert estimate_model_size(span) == expected_size @pytest.mark.asyncio async def test_span_triggers_status_transition(store_fixture: LightningStore, mock_readable_span: Mock) -> None: """Test that adding first span transitions rollout from preparing to running.""" rollout = await store_fixture.enqueue_rollout(input={"test": "data"}) # Pop to set status to preparing and create attempt popped = await store_fixture.dequeue_rollout() assert popped is not None assert popped.status == "preparing" # Verify status in store rollouts = await store_fixture.query_rollouts(status=["preparing"]) assert len(rollouts) == 1 # Get the attempt attempts = await store_fixture.query_attempts(rollout.rollout_id) attempt_id = attempts[0].attempt_id # Add first span await store_fixture.add_otel_span(rollout.rollout_id, attempt_id, mock_readable_span) # Status should transition to running rollouts = await store_fixture.query_rollouts(status=["running"]) assert len(rollouts) == 1 assert rollouts[0].rollout_id == rollout.rollout_id # Rollout Lifecycle Tests @pytest.mark.asyncio async def test_span_does_not_reset_timeout_attempt(store_fixture: LightningStore, mock_readable_span: Mock) -> None: """Adding a span to a timed-out attempt should not mark it running again.""" rollout = await store_fixture.enqueue_rollout(input={"test": "timeout-span"}) # Create the first attempt dequeued = await store_fixture.dequeue_rollout() assert dequeued is not None attempt_id = dequeued.attempt.attempt_id # Simulate the attempt timing out await store_fixture.update_attempt( rollout_id=rollout.rollout_id, attempt_id=attempt_id, status="timeout", ) attempts_before = await store_fixture.query_attempts(rollout.rollout_id) assert attempts_before[0].status == "timeout" rollout_before = await store_fixture.get_rollout_by_id(rollout.rollout_id) assert rollout_before is not None assert rollout_before.status != "running" # Adding a new span should keep the attempt in timeout state await store_fixture.add_otel_span(rollout.rollout_id, attempt_id, mock_readable_span) attempts_after = await store_fixture.query_attempts(rollout.rollout_id) assert attempts_after[0].status == "timeout" assert attempts_after[0].last_heartbeat_time is not None rollout_after = await store_fixture.get_rollout_by_id(rollout.rollout_id) assert rollout_after is not None assert rollout_after.status == rollout_before.status @pytest.mark.asyncio async def test_completion_sets_end_time(store_fixture: LightningStore) -> None: """Test that completing a rollout sets end_time.""" rollout = await store_fixture.enqueue_rollout(input={"test": "data"}) # Initially no end_time assert rollout.end_time is None # Complete as succeeded await store_fixture.update_rollout(rollout_id=rollout.rollout_id, status="succeeded") completed_rollouts = await store_fixture.query_rollouts() completed = completed_rollouts[0] assert completed.status == "succeeded" assert completed.end_time is not None assert completed.end_time > completed.start_time @pytest.mark.asyncio async def test_wait_for_rollouts(store_fixture: LightningStore) -> None: """Test waiting for rollout completion.""" # Add multiple rollouts r1 = await store_fixture.enqueue_rollout(input={"id": 1}) r2 = await store_fixture.enqueue_rollout(input={"id": 2}) _r3 = await store_fixture.enqueue_rollout(input={"id": 3}) # Start waiting for r1 and r2 async def wait_for_completion() -> List[Rollout]: return await store_fixture.wait_for_rollouts(rollout_ids=[r1.rollout_id, r2.rollout_id], timeout=5.0) wait_task = asyncio.create_task(wait_for_completion()) await asyncio.sleep(0.01) # Let wait task start # Complete r1 await store_fixture.update_rollout(rollout_id=r1.rollout_id, status="succeeded") # Complete r2 await store_fixture.update_rollout(rollout_id=r2.rollout_id, status="failed") # Get results completed = await wait_task assert len(completed) == 2 assert {r.rollout_id for r in completed} == {r1.rollout_id, r2.rollout_id} assert {r.status for r in completed} == {"succeeded", "failed"} @pytest.mark.asyncio async def test_wait_timeout(store_fixture: LightningStore) -> None: """Test wait_for_rollouts timeout behavior.""" rollout = await store_fixture.enqueue_rollout(input={"test": "data"}) start = time.time() completed = await store_fixture.wait_for_rollouts(rollout_ids=[rollout.rollout_id], timeout=0.1) elapsed = time.time() - start assert elapsed < 0.2 # Should timeout quickly assert len(completed) == 0 # No completions @pytest.mark.asyncio async def test_wait_with_timeout_none_polling(store_fixture: LightningStore) -> None: """Test wait_for_rollouts with timeout=None uses polling and can be cancelled.""" rollout = await store_fixture.enqueue_rollout(input={"test": "indefinite"}) async def wait_indefinitely(): return await store_fixture.wait_for_rollouts(rollout_ids=[rollout.rollout_id], timeout=None) # Start waiting with timeout=None wait_task = asyncio.create_task(wait_indefinitely()) # Give it a moment to start polling await asyncio.sleep(0.1) # Complete the rollout await store_fixture.update_rollout(rollout_id=rollout.rollout_id, status="succeeded") # The wait should complete now timeout = 1.0 if isinstance(store_fixture, InMemoryLightningStore) else 11.0 completed = await asyncio.wait_for(wait_task, timeout=timeout) assert len(completed) == 1 assert completed[0].rollout_id == rollout.rollout_id assert completed[0].status == "succeeded" @pytest.mark.asyncio async def test_wait_with_timeout_none_can_be_cancelled(store_fixture: LightningStore) -> None: """Test that wait_for_rollouts with timeout=None can be cancelled cleanly.""" rollout = await store_fixture.enqueue_rollout(input={"test": "cancel"}) async def wait_indefinitely(): return await store_fixture.wait_for_rollouts(rollout_ids=[rollout.rollout_id], timeout=None) # Start waiting with timeout=None wait_task = asyncio.create_task(wait_indefinitely()) # Give it time to start polling await asyncio.sleep(0.15) # Wait for at least one poll cycle # Cancel the task wait_task.cancel() # Should raise CancelledError with pytest.raises(asyncio.CancelledError): await wait_task # Task should be cancelled, no hanging threads assert wait_task.cancelled() @pytest.mark.asyncio async def test_wait_with_timeout_zero(store_fixture: LightningStore) -> None: """Test wait_for_rollouts with timeout=0 returns immediately.""" rollout = await store_fixture.enqueue_rollout(input={"test": "zero"}) start = time.time() completed = await store_fixture.wait_for_rollouts(rollout_ids=[rollout.rollout_id], timeout=0) elapsed = time.time() - start # Should return almost immediately assert elapsed < 0.05 assert len(completed) == 0 @pytest.mark.asyncio async def test_wait_with_already_completed_rollout(store_fixture: LightningStore) -> None: """Test wait_for_rollouts returns immediately for already completed rollouts.""" rollout = await store_fixture.enqueue_rollout(input={"test": "already_done"}) # Complete it first await store_fixture.update_rollout(rollout_id=rollout.rollout_id, status="succeeded") # Wait should return immediately without blocking start = time.time() completed = await store_fixture.wait_for_rollouts(rollout_ids=[rollout.rollout_id], timeout=5.0) elapsed = time.time() - start assert elapsed < 0.1 # Should be instant assert len(completed) == 1 assert completed[0].rollout_id == rollout.rollout_id assert completed[0].status == "succeeded" @pytest.mark.asyncio async def test_wait_multiple_rollouts_different_completion_times(store_fixture: LightningStore) -> None: """Test waiting for multiple rollouts that complete at different times.""" r1 = await store_fixture.enqueue_rollout(input={"id": 1}) r2 = await store_fixture.enqueue_rollout(input={"id": 2}) r3 = await store_fixture.enqueue_rollout(input={"id": 3}) async def wait_for_all(): return await store_fixture.wait_for_rollouts( rollout_ids=[r1.rollout_id, r2.rollout_id, r3.rollout_id], timeout=2.0 ) wait_task = asyncio.create_task(wait_for_all()) # Complete them at different times await asyncio.sleep(0.05) await store_fixture.update_rollout(rollout_id=r2.rollout_id, status="succeeded") await asyncio.sleep(0.05) await store_fixture.update_rollout(rollout_id=r1.rollout_id, status="failed") await asyncio.sleep(0.05) await store_fixture.update_rollout(rollout_id=r3.rollout_id, status="succeeded") # All should be collected completed = await wait_task assert len(completed) == 3 completed_ids = {r.rollout_id for r in completed} assert completed_ids == {r1.rollout_id, r2.rollout_id, r3.rollout_id} @pytest.mark.asyncio async def test_wait_partial_completion_on_timeout(store_fixture: LightningStore) -> None: """Test that wait_for_rollouts returns partial results when timeout occurs.""" r1 = await store_fixture.enqueue_rollout(input={"id": 1}) r2 = await store_fixture.enqueue_rollout(input={"id": 2}) r3 = await store_fixture.enqueue_rollout(input={"id": 3}) async def wait_with_short_timeout(): return await store_fixture.wait_for_rollouts( rollout_ids=[r1.rollout_id, r2.rollout_id, r3.rollout_id], timeout=0.2 ) wait_task = asyncio.create_task(wait_with_short_timeout()) # Only complete one before timeout await asyncio.sleep(0.05) await store_fixture.update_rollout(rollout_id=r1.rollout_id, status="succeeded") # Wait for timeout completed = await wait_task # Should only get r1 assert len(completed) == 1 assert completed[0].rollout_id == r1.rollout_id @pytest.mark.asyncio async def test_wait_concurrent_waiters_on_same_rollout(store_fixture: LightningStore) -> None: """Test multiple concurrent waiters on the same rollout.""" rollout = await store_fixture.enqueue_rollout(input={"test": "concurrent"}) async def wait_for_completion(): return await store_fixture.wait_for_rollouts(rollout_ids=[rollout.rollout_id], timeout=2.0) # Start multiple waiters concurrently wait_tasks = [asyncio.create_task(wait_for_completion()) for _ in range(5)] await asyncio.sleep(0.05) # Complete the rollout await store_fixture.update_rollout(rollout_id=rollout.rollout_id, status="succeeded") # All waiters should complete results = await asyncio.gather(*wait_tasks) # Each waiter should get the completed rollout for completed in results: assert len(completed) == 1 assert completed[0].rollout_id == rollout.rollout_id assert completed[0].status == "succeeded" @pytest.mark.asyncio async def test_wait_nonexistent_rollout_with_finite_timeout(store_fixture: LightningStore) -> None: """Test waiting for non-existent rollout with finite timeout.""" start = time.time() completed = await store_fixture.wait_for_rollouts(rollout_ids=["nonexistent"], timeout=0.1) elapsed = time.time() - start if isinstance(store_fixture, InMemoryLightningStore): # Should timeout quickly (not wait indefinitely) assert elapsed < 0.2 else: # Should be slower, but not too slow assert elapsed < 2.0 assert len(completed) == 0 @pytest.mark.asyncio async def test_wait_mixed_existing_and_nonexistent_rollouts(store_fixture: LightningStore) -> None: """Test waiting for mix of existing and non-existent rollouts.""" r1 = await store_fixture.enqueue_rollout(input={"id": 1}) async def wait_for_mixed(): return await store_fixture.wait_for_rollouts( rollout_ids=[r1.rollout_id, "nonexistent1", "nonexistent2"], timeout=0.5 ) wait_task = asyncio.create_task(wait_for_mixed()) await asyncio.sleep(0.05) await store_fixture.update_rollout(rollout_id=r1.rollout_id, status="succeeded") completed = await wait_task # Should only get the existing, completed rollout assert len(completed) == 1 assert completed[0].rollout_id == r1.rollout_id @pytest.mark.asyncio async def test_wait_event_set_before_wait_starts(store_fixture: LightningStore) -> None: """Test that waiting on an already-set event returns immediately.""" rollout = await store_fixture.enqueue_rollout(input={"test": "early_complete"}) # Complete it before waiting await store_fixture.update_rollout(rollout_id=rollout.rollout_id, status="succeeded") # Now start waiting - should return immediately start = time.time() completed = await store_fixture.wait_for_rollouts(rollout_ids=[rollout.rollout_id], timeout=10.0) elapsed = time.time() - start assert elapsed < 0.05 # Should be instant assert len(completed) == 1 assert completed[0].status == "succeeded" @pytest.mark.asyncio async def test_wait_polling_interval_with_timeout_none(store_fixture: LightningStore) -> None: """Test that timeout=None polling doesn't busy-wait (uses reasonable intervals).""" rollout = await store_fixture.enqueue_rollout(input={"test": "polling"}) start = time.time() async def wait_and_complete(): # Start waiting with timeout=None wait_task = asyncio.create_task(store_fixture.wait_for_rollouts(rollout_ids=[rollout.rollout_id], timeout=None)) # Wait for 0.5 seconds to let polling happen await asyncio.sleep(0.5) # Complete the rollout await store_fixture.update_rollout(rollout_id=rollout.rollout_id, status="succeeded") return await wait_task completed = await wait_and_complete() elapsed = time.time() - start if isinstance(store_fixture, InMemoryLightningStore): # Should complete after ~0.5s (when we set the event) assert 0.4 < elapsed < 0.7 else: # Should be more than 5 seconds assert 5 < elapsed < 15 assert len(completed) == 1 assert completed[0].status == "succeeded" # Concurrent Access Tests @pytest.mark.asyncio async def test_concurrent_task_addition(store_fixture: LightningStore) -> None: """Test adding tasks concurrently.""" async def enqueue_rollout(index: int) -> Rollout: return await store_fixture.enqueue_rollout(input={"index": index}) # Add 50 tasks concurrently tasks = [enqueue_rollout(i) for i in range(50)] rollouts = await asyncio.gather(*tasks) # All should succeed with unique IDs assert len(rollouts) == 50 ids = [r.rollout_id for r in rollouts] assert len(set(ids)) == 50 # All should be in store all_rollouts = await store_fixture.query_rollouts() assert len(all_rollouts) == 50 @pytest.mark.asyncio async def test_concurrent_pop_operations(store_fixture: LightningStore) -> None: """Test concurrent popping ensures each rollout is popped once.""" # Add 20 tasks for i in range(20): await store_fixture.enqueue_rollout(input={"index": i}) async def pop_task() -> Rollout | None: return await store_fixture.dequeue_rollout() # Pop concurrently (more attempts than available) tasks = [pop_task() for _ in range(30)] results = await asyncio.gather(*tasks) # Should get exactly 20 rollouts and 10 None valid = [r for r in results if r is not None] none_results = [r for r in results if r is None] assert len(valid) == 20 assert len(none_results) == 10 # Each rollout popped exactly once ids = [r.rollout_id for r in valid] assert len(set(ids)) == 20 @pytest.mark.asyncio async def test_concurrent_span_additions(store_fixture: LightningStore, mock_readable_span: Mock) -> None: """Test concurrent span additions maintain consistency.""" await store_fixture.enqueue_rollout(input={"test": "data"}) rollout = await store_fixture.dequeue_rollout() # Create an attempt assert rollout is not None async def add_span(index: int) -> Span: return await store_fixture.add_otel_span(rollout.rollout_id, rollout.attempt.attempt_id, mock_readable_span) # type: ignore # Add 30 spans concurrently tasks = [add_span(i) for i in range(30)] spans = await asyncio.gather(*tasks) # All should have unique sequence IDs seq_ids = [s.sequence_id for s in spans] assert len(set(seq_ids)) == 30 assert set(seq_ids) == set(range(1, 31)) @pytest.mark.asyncio async def test_concurrent_resource_updates(store_fixture: LightningStore) -> None: """Test concurrent resource updates are atomic.""" async def update_resource(ver: int) -> None: llm = LLM( resource_type="llm", endpoint=f"http://localhost:808{ver % 10}", model=f"model-v{ver}", sampling_parameters={"temperature": 0.5 + ver * 0.01}, ) update = ResourcesUpdate( resources_id=f"v{ver}", resources={"llm": llm}, create_time=time.time(), update_time=time.time(), version=1 ) await store_fixture.update_resources(update.resources_id, update.resources) # Update concurrently tasks = [update_resource(i) for i in range(50)] await asyncio.gather(*tasks) # Latest should be one of the versions latest = await store_fixture.get_latest_resources() assert latest is not None assert latest.resources_id.startswith("v") # All versions should be stored for i in range(50): res = await store_fixture.get_resources_by_id(f"v{i}") assert res is not None assert isinstance(res.resources["llm"], LLM) assert res.resources["llm"].model == f"model-v{i}" # Error Handling Tests @pytest.mark.asyncio async def test_update_nonexistent_rollout(store_fixture: LightningStore) -> None: """Test updating non-existent rollout raises error.""" with pytest.raises(ValueError, match=r"Item.*does not exist"): await store_fixture.update_rollout(rollout_id="nonexistent", status="failed") @pytest.mark.asyncio async def test_add_span_without_rollout(store_fixture: LightningStore, mock_readable_span: Mock) -> None: """Test adding span to non-existent rollout raises error.""" with pytest.raises(ValueError, match="Rollout nonexistent not found"): await store_fixture.add_otel_span("nonexistent", "attempt-1", mock_readable_span) @pytest.mark.asyncio async def test_add_span_with_missing_attempt(store_fixture: LightningStore, mock_readable_span: Mock) -> None: """Test adding span with an unknown attempt_id raises a helpful error.""" rollout = await store_fixture.enqueue_rollout(input={"test": "data"}) # Create a valid attempt to ensure rollout exists in store await store_fixture.dequeue_rollout() invalid_span = Span.from_opentelemetry( mock_readable_span, rollout_id=rollout.rollout_id, attempt_id="attempt-missing", sequence_id=1, ) with pytest.raises(ValueError, match="Attempt attempt-missing not found"): await store_fixture.add_span(invalid_span) @pytest.mark.asyncio async def test_query_empty_spans(store_fixture: LightningStore) -> None: """Test querying spans for non-existent rollout returns empty.""" spans = await store_fixture.query_spans("nonexistent") assert len(spans) == 0 # With attempt_id spans = await store_fixture.query_spans("nonexistent", attempt_id="attempt-1") assert len(spans) == 0 # With latest spans = await store_fixture.query_spans("nonexistent", attempt_id="latest") assert len(spans) == 0 @pytest.mark.asyncio async def test_query_latest_with_no_spans(store_fixture: LightningStore) -> None: """Test querying 'latest' attempt when no spans exist.""" rollout = await store_fixture.enqueue_rollout(input={"test": "data"}) spans = await store_fixture.query_spans(rollout.rollout_id, attempt_id="latest") assert len(spans) == 0 @pytest.mark.asyncio async def test_wait_for_nonexistent_rollout(store_fixture: LightningStore) -> None: """Test waiting for non-existent rollout handles gracefully.""" completed = await store_fixture.wait_for_rollouts(rollout_ids=["nonexistent"], timeout=0.1) assert len(completed) == 0 # Attempt Management Tests @pytest.mark.asyncio async def test_query_attempts(store_fixture: LightningStore) -> None: """Test querying attempts for a rollout.""" rollout = await store_fixture.enqueue_rollout(input={"test": "data"}) # Initially no attempts attempts = await store_fixture.query_attempts(rollout.rollout_id) assert len(attempts) == 0 # Pop creates first attempt await store_fixture.dequeue_rollout() attempts = await store_fixture.query_attempts(rollout.rollout_id) assert len(attempts) == 1 assert attempts[0].sequence_id == 1 assert attempts[0].status == "preparing" # Requeue and pop creates second attempt await store_fixture.update_rollout(rollout_id=rollout.rollout_id, status="requeuing") await store_fixture.dequeue_rollout() attempts = await store_fixture.query_attempts(rollout.rollout_id) assert len(attempts) == 2 assert attempts[0].sequence_id == 1 assert attempts[1].sequence_id == 2 @pytest.mark.asyncio async def test_get_latest_attempt(store_fixture: LightningStore) -> None: """Test getting the latest attempt.""" rollout = await store_fixture.enqueue_rollout(input={"test": "data"}) # No attempts initially latest = await store_fixture.get_latest_attempt(rollout.rollout_id) assert latest is None # Create first attempt await store_fixture.dequeue_rollout() latest = await store_fixture.get_latest_attempt(rollout.rollout_id) assert latest is not None assert latest.sequence_id == 1 # Create second attempt await store_fixture.update_rollout(rollout_id=rollout.rollout_id, status="requeuing") await store_fixture.dequeue_rollout() latest = await store_fixture.get_latest_attempt(rollout.rollout_id) assert latest is not None assert latest.sequence_id == 2 @pytest.mark.asyncio async def test_update_attempt_fields(store_fixture: LightningStore) -> None: """Test updating attempt fields.""" rollout = await store_fixture.enqueue_rollout(input={"test": "data"}) await store_fixture.dequeue_rollout() attempts = await store_fixture.query_attempts(rollout.rollout_id) attempt = attempts[0] # Update various fields updated = await store_fixture.update_attempt( rollout_id=rollout.rollout_id, attempt_id=attempt.attempt_id, status="running", worker_id="worker-123", last_heartbeat_time=time.time(), metadata={"custom": "value"}, ) assert updated.status == "running" assert updated.worker_id == "worker-123" assert updated.last_heartbeat_time is not None assert updated.metadata is not None assert updated.metadata["custom"] == "value" @pytest.mark.asyncio async def test_update_latest_attempt(store_fixture: LightningStore) -> None: """Test updating latest attempt using 'latest' identifier.""" rollout = await store_fixture.enqueue_rollout(input={"test": "data"}) await store_fixture.dequeue_rollout() # Update using 'latest' updated = await store_fixture.update_attempt(rollout_id=rollout.rollout_id, attempt_id="latest", status="succeeded") assert updated.status == "succeeded" assert updated.end_time is not None # Should auto-set end_time @pytest.mark.asyncio async def test_update_attempt_sets_end_time_for_terminal_status(store_fixture: LightningStore) -> None: """Terminal attempt statuses set end_time while in-progress statuses don't.""" rollout = await store_fixture.enqueue_rollout(input={"test": "data"}) await store_fixture.dequeue_rollout() attempt = (await store_fixture.query_attempts(rollout.rollout_id))[0] assert attempt.end_time is None running = await store_fixture.update_attempt( rollout_id=rollout.rollout_id, attempt_id=attempt.attempt_id, status="running", ) assert running.status == "running" assert running.end_time is None failed = await store_fixture.update_attempt( rollout_id=rollout.rollout_id, attempt_id=attempt.attempt_id, status="failed", ) assert failed.status == "failed" assert failed.end_time is not None assert failed.end_time >= failed.start_time rollout = await store_fixture.get_rollout_by_id(rollout_id=rollout.rollout_id) assert rollout is not None assert rollout.status == "failed" assert rollout.end_time is not None assert rollout.end_time >= failed.end_time @pytest.mark.asyncio async def test_rollout_retry_lifecycle_updates_statuses( store_fixture: LightningStore, mock_readable_span: Mock ) -> None: """Rollout retry creates new attempts and updates statuses via spans and completions.""" rollout = await store_fixture.enqueue_rollout(input={"test": "data"}) first_attempted = await store_fixture.dequeue_rollout() assert first_attempted is not None assert first_attempted.status == "preparing" first_attempt = (await store_fixture.query_attempts(rollout.rollout_id))[0] await store_fixture.add_otel_span(rollout.rollout_id, first_attempt.attempt_id, mock_readable_span) # Status should reflect running state after span is recorded running_rollout = await store_fixture.query_rollouts(status=["running"]) assert running_rollout and running_rollout[0].rollout_id == rollout.rollout_id running_attempts = await store_fixture.query_attempts(rollout.rollout_id) assert running_attempts[0].status == "running" # Mark first attempt as failed and requeue rollout failed_attempt = await store_fixture.update_attempt( rollout_id=rollout.rollout_id, attempt_id=first_attempt.attempt_id, status="failed", ) assert failed_attempt.end_time is not None await store_fixture.update_rollout(rollout_id=rollout.rollout_id, status="requeuing") attempts_after_failure = await store_fixture.query_attempts(rollout.rollout_id) assert [a.status for a in attempts_after_failure] == ["failed"] retry_attempted = await store_fixture.dequeue_rollout() assert retry_attempted is not None assert retry_attempted.status == "preparing" assert retry_attempted.attempt.sequence_id == 2 latest_pre_span = await store_fixture.get_latest_attempt(rollout.rollout_id) assert latest_pre_span is not None and latest_pre_span.sequence_id == 2 assert latest_pre_span.status == "preparing" await store_fixture.add_otel_span(rollout.rollout_id, retry_attempted.attempt.attempt_id, mock_readable_span) latest_running = await store_fixture.get_latest_attempt(rollout.rollout_id) assert latest_running is not None assert latest_running.sequence_id == 2 assert latest_running.status == "running" await store_fixture.update_attempt( rollout_id=rollout.rollout_id, attempt_id=retry_attempted.attempt.attempt_id, status="succeeded", ) await store_fixture.update_rollout(rollout_id=rollout.rollout_id, status="succeeded") final_rollout = await store_fixture.query_rollouts(status=["succeeded"]) assert final_rollout and final_rollout[0].rollout_id == rollout.rollout_id final_attempts = await store_fixture.query_attempts(rollout.rollout_id) assert [a.status for a in final_attempts] == ["failed", "succeeded"] @pytest.mark.asyncio async def test_update_nonexistent_attempt(store_fixture: LightningStore) -> None: """Test updating non-existent attempt raises error.""" rollout = await store_fixture.enqueue_rollout(input={"test": "data"}) with pytest.raises(ValueError, match="No attempts found"): await store_fixture.update_attempt(rollout_id=rollout.rollout_id, attempt_id="nonexistent", status="failed") # Add Attempt Tests @pytest.mark.asyncio async def test_add_attempt_creates_new_attempt(store_fixture: LightningStore) -> None: """Test add_attempt creates a new attempt for existing rollout.""" # Create a rollout rollout = await store_fixture.enqueue_rollout(input={"test": "data"}) # Add first manual attempt attempted_rollout = await store_fixture.start_attempt(rollout.rollout_id) assert attempted_rollout.rollout_id == rollout.rollout_id assert attempted_rollout.attempt.sequence_id == 1 assert attempted_rollout.attempt.status == "preparing" assert attempted_rollout.attempt.rollout_id == rollout.rollout_id assert attempted_rollout.attempt.attempt_id.startswith("at-") # Verify attempt is stored attempts = await store_fixture.query_attempts(rollout.rollout_id) assert len(attempts) == 1 assert attempts[0].attempt_id == attempted_rollout.attempt.attempt_id @pytest.mark.asyncio async def test_add_attempt_increments_sequence_id(store_fixture: LightningStore) -> None: """Test add_attempt correctly increments sequence_id.""" # Create a rollout and dequeue to create first attempt rollout = await store_fixture.enqueue_rollout(input={"test": "data"}) await store_fixture.dequeue_rollout() # Creates attempt with sequence_id=1 # Add second attempt manually attempted_rollout2 = await store_fixture.start_attempt(rollout.rollout_id) assert attempted_rollout2.attempt.sequence_id == 2 # Add third attempt manually attempted_rollout3 = await store_fixture.start_attempt(rollout.rollout_id) assert attempted_rollout3.attempt.sequence_id == 3 # Verify all attempts exist attempts = await store_fixture.query_attempts(rollout.rollout_id) assert len(attempts) == 3 assert [a.sequence_id for a in attempts] == [1, 2, 3] @pytest.mark.asyncio async def test_add_attempt_nonexistent_rollout(store_fixture: LightningStore) -> None: """Test add_attempt raises error for nonexistent rollout.""" with pytest.raises(ValueError, match="Rollout nonexistent not found"): await store_fixture.start_attempt("nonexistent") @pytest.mark.asyncio async def test_add_attempt_ignores_max_attempts(store_fixture: LightningStore) -> None: """Test add_attempt ignores max_attempts configuration.""" # Create rollout with max_attempts=2 rollout = await store_fixture.enqueue_rollout(input={"test": "data"}) config = RolloutConfig(max_attempts=2) await store_fixture.update_rollout(rollout_id=rollout.rollout_id, config=config) # Add attempts beyond max_attempts attempt1 = await store_fixture.start_attempt(rollout.rollout_id) attempt2 = await store_fixture.start_attempt(rollout.rollout_id) attempt3 = await store_fixture.start_attempt(rollout.rollout_id) # Should succeed despite max_attempts=2 assert attempt1.attempt.sequence_id == 1 assert attempt2.attempt.sequence_id == 2 assert attempt3.attempt.sequence_id == 3 # All attempts should exist attempts = await store_fixture.query_attempts(rollout.rollout_id) assert len(attempts) == 3 # Latest Attempt Status Propagation Tests @pytest.mark.asyncio async def test_status_propagation_only_for_latest_attempt(store_fixture: LightningStore) -> None: """Test that status changes only propagate to rollout when updating latest attempt.""" rollout = await store_fixture.enqueue_rollout(input={"test": "propagation"}) # Create multiple attempts attempt1 = await store_fixture.start_attempt(rollout.rollout_id) _attempt2 = await store_fixture.start_attempt(rollout.rollout_id) attempt3 = await store_fixture.start_attempt(rollout.rollout_id) # This is the latest # Update attempt1 (not latest) to succeeded await store_fixture.update_attempt( rollout_id=rollout.rollout_id, attempt_id=attempt1.attempt.attempt_id, status="succeeded" ) # Rollout status should NOT change since attempt1 is not the latest updated_rollout = await store_fixture.get_rollout_by_id(rollout.rollout_id) assert updated_rollout is not None assert updated_rollout.status == "preparing" # Should be status of attempt 2 # Update attempt3 (latest) to succeeded await store_fixture.update_attempt( rollout_id=rollout.rollout_id, attempt_id=attempt3.attempt.attempt_id, status="succeeded" ) # Now rollout status should change since we updated the latest attempt updated_rollout = await store_fixture.get_rollout_by_id(rollout.rollout_id) assert updated_rollout is not None assert updated_rollout.status == "succeeded" @pytest.mark.asyncio async def test_status_propagation_with_retry_for_latest_attempt(store_fixture: LightningStore) -> None: """Test retry logic only applies when updating latest attempt.""" rollout = await store_fixture.enqueue_rollout(input={"test": "retry"}) config = RolloutConfig(max_attempts=3, retry_condition=["failed"]) await store_fixture.update_rollout(rollout_id=rollout.rollout_id, config=config) # Create multiple attempts attempt1 = await store_fixture.start_attempt(rollout.rollout_id) # sequence_id=1 attempt2 = await store_fixture.start_attempt(rollout.rollout_id) # sequence_id=2 (latest) # Fail attempt1 (not latest) - should NOT trigger retry await store_fixture.update_attempt( rollout_id=rollout.rollout_id, attempt_id=attempt1.attempt.attempt_id, status="failed" ) updated_rollout = await store_fixture.get_rollout_by_id(rollout.rollout_id) assert updated_rollout is not None assert updated_rollout.status == "preparing" # Should be status of attempt 2 # Fail attempt2 (latest) - should trigger retry since sequence_id=2 < max_attempts=3 await store_fixture.update_attempt( rollout_id=rollout.rollout_id, attempt_id=attempt2.attempt.attempt_id, status="failed" ) updated_rollout = await store_fixture.get_rollout_by_id(rollout.rollout_id) assert updated_rollout is not None assert updated_rollout.status == "requeuing" # Should be requeued for retry @pytest.mark.asyncio async def test_status_propagation_latest_changes_when_new_attempt_added(store_fixture: LightningStore) -> None: """Test that the 'latest attempt' changes as new attempts are added.""" rollout = await store_fixture.enqueue_rollout(input={"test": "latest_changes"}) # Create first attempt and update it to succeeded attempt1 = await store_fixture.start_attempt(rollout.rollout_id) await store_fixture.update_attempt( rollout_id=rollout.rollout_id, attempt_id=attempt1.attempt.attempt_id, status="succeeded" ) # Rollout should be succeeded since attempt1 is latest updated_rollout = await store_fixture.get_rollout_by_id(rollout.rollout_id) assert updated_rollout is not None assert updated_rollout.status == "succeeded" # Add second attempt (now this becomes latest) attempt2 = await store_fixture.start_attempt(rollout.rollout_id) # Update attempt1 to failed - should NOT affect rollout since it's no longer latest await store_fixture.update_attempt( rollout_id=rollout.rollout_id, attempt_id=attempt1.attempt.attempt_id, status="failed" ) updated_rollout = await store_fixture.get_rollout_by_id(rollout.rollout_id) assert updated_rollout is not None assert updated_rollout.status == "preparing" # Should be the status of attempt 2 # Update attempt2 (now latest) to failed await store_fixture.update_attempt( rollout_id=rollout.rollout_id, attempt_id=attempt2.attempt.attempt_id, status="failed" ) # Now rollout should change since we updated the new latest attempt updated_rollout = await store_fixture.get_rollout_by_id(rollout.rollout_id) assert updated_rollout is not None assert updated_rollout.status == "failed" @pytest.mark.asyncio async def test_status_propagation_update_latest_by_reference(store_fixture: LightningStore) -> None: """Test status propagation when updating latest attempt using 'latest' reference.""" rollout = await store_fixture.enqueue_rollout(input={"test": "latest_ref"}) # Create multiple attempts await store_fixture.start_attempt(rollout.rollout_id) await store_fixture.start_attempt(rollout.rollout_id) attempt3 = await store_fixture.start_attempt(rollout.rollout_id) # This is latest # Update using "latest" reference updated_attempt = await store_fixture.update_attempt( rollout_id=rollout.rollout_id, attempt_id="latest", status="succeeded" ) # Should have updated attempt3 assert updated_attempt.attempt_id == attempt3.attempt.attempt_id assert updated_attempt.status == "succeeded" # Rollout should be updated since we updated the latest attempt updated_rollout = await store_fixture.get_rollout_by_id(rollout.rollout_id) assert updated_rollout is not None assert updated_rollout.status == "succeeded" @pytest.mark.asyncio async def test_healthcheck_timeout_behavior(store_fixture: LightningStore, mock_readable_span: Mock) -> None: """Test that healthcheck detects and handles timeout conditions.""" # Create rollout with short timeout configuration config = RolloutConfig( timeout_seconds=0.1, max_attempts=2, retry_condition=["timeout"] # Very short timeout for testing ) rollout = await store_fixture.enqueue_rollout(input={"test": "timeout"}) await store_fixture.update_rollout(rollout_id=rollout.rollout_id, config=config) # Dequeue to create an attempt and add span to make it running attempted = await store_fixture.dequeue_rollout() assert attempted is not None await store_fixture.add_otel_span(rollout.rollout_id, attempted.attempt.attempt_id, mock_readable_span) # Verify it's running running_rollouts = await store_fixture.query_rollouts(status=["running"]) assert len(running_rollouts) == 1 # Wait for timeout to occur await asyncio.sleep(0.15) # Wait longer than timeout_seconds # Trigger healthcheck by calling any decorated method # Verify the attempt was marked as timeout and rollout was requeued attempts = await store_fixture.query_attempts(rollout.rollout_id) assert len(attempts) == 1 assert attempts[0].status == "timeout" # Since retry_condition includes "timeout" and max_attempts=2, should requeue rollout_after = await store_fixture.get_rollout_by_id(rollout.rollout_id) assert rollout_after is not None assert rollout_after.status == "requeuing" @pytest.mark.asyncio async def test_healthcheck_unresponsive_behavior(store_fixture: LightningStore, mock_readable_span: Mock) -> None: """Test that healthcheck detects and handles unresponsive conditions.""" # Create rollout with short unresponsive timeout but no retry for unresponsive config = RolloutConfig( unresponsive_seconds=0.2, # Very short unresponsive timeout max_attempts=3, retry_condition=["timeout"], # Note: "unresponsive" not in retry_condition ) rollout = await store_fixture.enqueue_rollout(input={"test": "unresponsive"}) await store_fixture.update_rollout(rollout_id=rollout.rollout_id, config=config) # Dequeue and add span to make it running (this sets last_heartbeat_time) attempted = await store_fixture.dequeue_rollout() assert attempted is not None await store_fixture.add_otel_span(rollout.rollout_id, attempted.attempt.attempt_id, mock_readable_span) # Verify it's running and has heartbeat running_attempts = await store_fixture.query_attempts(rollout.rollout_id) assert running_attempts[0].status == "running" assert running_attempts[0].last_heartbeat_time is not None # Wait for unresponsive timeout await asyncio.sleep(0.25) # Wait longer than unresponsive_seconds # Verify attempt was marked as unresponsive attempts_after = await store_fixture.query_attempts(rollout.rollout_id) assert attempts_after[0].status == "unresponsive" # Since "unresponsive" not in retry_condition, rollout should be failed rollout_after = await store_fixture.get_rollout_by_id(rollout.rollout_id) assert rollout_after is not None assert rollout_after.status == "failed" # Full Lifecycle Integration Tests @pytest.mark.asyncio async def test_full_lifecycle_success(store_fixture: LightningStore, mock_readable_span: Mock) -> None: """Test successful rollout lifecycle: queue -> prepare -> run -> succeed.""" # 1. Create task rollout = await store_fixture.enqueue_rollout(input={"test": "data"}, mode="train") assert rollout.status == "queuing" # 2. Pop to start processing (creates attempt) popped = await store_fixture.dequeue_rollout() assert popped is not None assert popped.status == "preparing" attempts = await store_fixture.query_attempts(rollout.rollout_id) assert len(attempts) == 1 attempt = attempts[0] assert attempt.status == "preparing" # 3. Add span (transitions to running) span = await store_fixture.add_otel_span(rollout.rollout_id, attempt.attempt_id, mock_readable_span) assert span is not None and span.sequence_id == 1 # Check status transitions rollouts = await store_fixture.query_rollouts(status=["running"]) assert len(rollouts) == 1 attempts = await store_fixture.query_attempts(rollout.rollout_id) assert attempts[0].status == "running" assert attempts[0].last_heartbeat_time is not None # 4. Complete successfully await store_fixture.update_attempt(rollout_id=rollout.rollout_id, attempt_id=attempt.attempt_id, status="succeeded") await store_fixture.update_rollout(rollout_id=rollout.rollout_id, status="succeeded") # Verify final state final = (await store_fixture.query_rollouts())[0] assert final.status == "succeeded" assert final.end_time is not None final_attempt = await store_fixture.get_latest_attempt(rollout.rollout_id) assert final_attempt is not None assert final_attempt.status == "succeeded" assert final_attempt.end_time is not None # Retry and requeue interactions def _retry_config() -> RolloutConfig: """Helper to create a rollout config that retries unresponsive attempts.""" return RolloutConfig(max_attempts=2, retry_condition=["unresponsive"]) @pytest.mark.asyncio async def test_requeued_attempt_recovers_before_retry(store_fixture: LightningStore, mock_readable_span: Mock) -> None: """A requeued attempt that resumes should be removed from the queue.""" attempted = await store_fixture.start_rollout(input={"foo": "bar"}) await store_fixture.update_rollout(rollout_id=attempted.rollout_id, config=_retry_config()) await store_fixture.update_attempt( rollout_id=attempted.rollout_id, attempt_id=attempted.attempt.attempt_id, status="unresponsive" ) rollout = await store_fixture.get_rollout_by_id(attempted.rollout_id) assert rollout is not None assert rollout.status == "requeuing" await store_fixture.add_otel_span(attempted.rollout_id, attempted.attempt.attempt_id, mock_readable_span) latest_attempt = await store_fixture.get_latest_attempt(attempted.rollout_id) assert latest_attempt is not None assert latest_attempt.attempt_id == attempted.attempt.attempt_id assert latest_attempt.status == "running" rollout = await store_fixture.get_rollout_by_id(attempted.rollout_id) assert rollout is not None assert rollout.status == "running" # Queue should no longer return the rollout for retry. assert await store_fixture.dequeue_rollout() is None @pytest.mark.asyncio async def test_requeued_attempt_succeeds_without_new_attempt( store_fixture: LightningStore, mock_readable_span: Mock ) -> None: """Recovered attempts can finish successfully without spawning a retry.""" attempted = await store_fixture.start_rollout(input={"foo": "bar"}) await store_fixture.update_rollout(rollout_id=attempted.rollout_id, config=_retry_config()) await store_fixture.update_attempt( rollout_id=attempted.rollout_id, attempt_id=attempted.attempt.attempt_id, status="unresponsive" ) await store_fixture.add_otel_span(attempted.rollout_id, attempted.attempt.attempt_id, mock_readable_span) await store_fixture.update_attempt( rollout_id=attempted.rollout_id, attempt_id=attempted.attempt.attempt_id, status="succeeded" ) rollout = await store_fixture.get_rollout_by_id(attempted.rollout_id) assert rollout is not None assert rollout.status == "succeeded" latest_attempt = await store_fixture.get_latest_attempt(attempted.rollout_id) assert latest_attempt is not None assert latest_attempt.status == "succeeded" assert latest_attempt.end_time is not None assert await store_fixture.dequeue_rollout() is None @pytest.mark.asyncio async def test_requeued_attempt_fails_without_new_attempt( store_fixture: LightningStore, mock_readable_span: Mock ) -> None: """Recovered attempts that fail should mark the rollout failed without retries.""" attempted = await store_fixture.start_rollout(input={"foo": "bar"}) await store_fixture.update_rollout(rollout_id=attempted.rollout_id, config=_retry_config()) await store_fixture.update_attempt( rollout_id=attempted.rollout_id, attempt_id=attempted.attempt.attempt_id, status="unresponsive" ) await store_fixture.add_otel_span(attempted.rollout_id, attempted.attempt.attempt_id, mock_readable_span) await store_fixture.update_attempt( rollout_id=attempted.rollout_id, attempt_id=attempted.attempt.attempt_id, status="failed" ) rollout = await store_fixture.get_rollout_by_id(attempted.rollout_id) assert rollout is not None assert rollout.status == "failed" latest_attempt = await store_fixture.get_latest_attempt(attempted.rollout_id) assert latest_attempt is not None assert latest_attempt.status == "failed" assert latest_attempt.end_time is not None assert await store_fixture.dequeue_rollout() is None @pytest.mark.asyncio async def test_requeued_attempt_recovers_after_retry_started( store_fixture: LightningStore, mock_readable_span: Mock ) -> None: """Data from an old attempt should not disrupt a newly started retry.""" attempted = await store_fixture.start_rollout(input={"foo": "bar"}) await store_fixture.update_rollout(rollout_id=attempted.rollout_id, config=_retry_config()) await store_fixture.update_attempt( rollout_id=attempted.rollout_id, attempt_id=attempted.attempt.attempt_id, status="unresponsive" ) # Start a new attempt by dequeuing the rollout from the queue. retried = await store_fixture.dequeue_rollout() assert retried is not None assert retried.attempt.sequence_id == 2 await store_fixture.add_otel_span(attempted.rollout_id, attempted.attempt.attempt_id, mock_readable_span) latest_attempt = await store_fixture.get_latest_attempt(attempted.rollout_id) assert latest_attempt is not None assert latest_attempt.attempt_id == retried.attempt.attempt_id assert latest_attempt.sequence_id == 2 # The old attempt is still marked running but does not change the rollout state. first_attempts = await store_fixture.query_attempts(attempted.rollout_id) assert first_attempts[0].status == "running" rollout = await store_fixture.get_rollout_by_id(attempted.rollout_id) assert rollout is not None assert rollout.status == "preparing" assert await store_fixture.dequeue_rollout() is None @pytest.mark.asyncio async def test_resources_update_tracks_create_and_update_times(store_fixture: LightningStore) -> None: """Test that ResourcesUpdate tracks create_time and update_time correctly.""" llm = LLM( resource_type="llm", endpoint="http://localhost:8080", model="test-model", sampling_parameters={"temperature": 0.7}, ) # Add initial resource start_time = time.time() update1 = await store_fixture.add_resources({"main_llm": llm}) # Verify create_time is set and reasonable assert update1.create_time >= start_time assert update1.create_time <= time.time() # Initially, update_time should equal create_time assert update1.update_time == update1.create_time assert update1.version == 1 # Wait a bit and update the same resource await asyncio.sleep(0.01) llm_v2 = LLM( resource_type="llm", endpoint="http://localhost:8080", model="test-model-v2", sampling_parameters={"temperature": 0.8}, ) update2 = await store_fixture.update_resources(update1.resources_id, {"main_llm": llm_v2}) # Verify update_time changed but create_time stayed the same assert update2.resources_id == update1.resources_id assert update2.create_time == update1.create_time # create_time should not change assert update2.update_time > update1.update_time # update_time should be newer assert update2.version == 2 # version should increment @pytest.mark.asyncio async def test_resources_update_version_increments(store_fixture: LightningStore) -> None: """Test that ResourcesUpdate version increments correctly with each update.""" llm = LLM( resource_type="llm", endpoint="http://localhost:8080", model="test-model", sampling_parameters={"temperature": 0.7}, ) # Add initial resource update1 = await store_fixture.add_resources({"main_llm": llm}) assert update1.version == 1 # Update it multiple times for i in range(2, 6): llm_updated = LLM( resource_type="llm", endpoint="http://localhost:8080", model=f"test-model-v{i}", sampling_parameters={"temperature": 0.7}, ) update = await store_fixture.update_resources(update1.resources_id, {"main_llm": llm_updated}) assert update.version == i assert update.resources_id == update1.resources_id assert update.create_time == update1.create_time @pytest.mark.asyncio async def test_resources_different_ids_have_independent_versions(store_fixture: LightningStore) -> None: """Test that different resources_ids have independent version counters.""" llm1 = LLM( resource_type="llm", endpoint="http://localhost:8080", model="model-1", sampling_parameters={"temperature": 0.7}, ) llm2 = LLM( resource_type="llm", endpoint="http://localhost:8080", model="model-2", sampling_parameters={"temperature": 0.8}, ) # Add two different resources res1 = await store_fixture.add_resources({"llm": llm1}) res2 = await store_fixture.add_resources({"llm": llm2}) # Both should start at version 1 assert res1.version == 1 assert res2.version == 1 assert res1.resources_id != res2.resources_id # Update res1 twice for i in range(2): llm_updated = LLM( resource_type="llm", endpoint="http://localhost:8080", model=f"model-1-v{i+2}", sampling_parameters={"temperature": 0.7}, ) res1 = await store_fixture.update_resources(res1.resources_id, {"llm": llm_updated}) # res1 should be at version 3, res2 should still be at version 1 assert res1.version == 3 retrieved_res2 = await store_fixture.get_resources_by_id(res2.resources_id) assert retrieved_res2 is not None assert retrieved_res2.version == 1 @pytest.mark.asyncio async def test_query_resources_returns_all_fields(store_fixture: LightningStore) -> None: """Test that query_resources returns all ResourcesUpdate fields.""" llm = LLM( resource_type="llm", endpoint="http://localhost:8080", model="test-model", sampling_parameters={"temperature": 0.7}, ) # Add multiple resources await store_fixture.add_resources({"llm": llm}) await asyncio.sleep(0.01) await store_fixture.add_resources({"llm": llm}) # Query all resources all_resources = await store_fixture.query_resources() assert len(all_resources) == 2 for res in all_resources: assert res.resources_id is not None assert res.create_time > 0 assert res.update_time > 0 assert res.version >= 1 assert res.resources is not None @pytest.mark.asyncio async def test_scan_debounce_allows_initial_scan( fake_time: FakeTimeController, debounced_store: CollectionBasedLightningStore[Any] ) -> None: """The first watchdog scan should run immediately even with debouncing enabled.""" fake_time.set(100.0) assert await debounced_store._should_scan_for_unhealthy_rollouts() # pyright: ignore[reportPrivateUsage] @pytest.mark.asyncio async def test_scan_debounce_blocks_until_interval( fake_time: FakeTimeController, debounced_store: CollectionBasedLightningStore[Any] ) -> None: """Subsequent scans wait until the debounce window elapses.""" fake_time.set(200.0) assert await debounced_store._should_scan_for_unhealthy_rollouts() # pyright: ignore[reportPrivateUsage] assert not await debounced_store._should_scan_for_unhealthy_rollouts() # pyright: ignore[reportPrivateUsage] fake_time.advance(debounced_store._scan_debounce_seconds - 1) # pyright: ignore[reportPrivateUsage] assert not await debounced_store._should_scan_for_unhealthy_rollouts() # pyright: ignore[reportPrivateUsage] fake_time.advance(1.0) assert await debounced_store._should_scan_for_unhealthy_rollouts() # pyright: ignore[reportPrivateUsage] @pytest.mark.asyncio async def test_scan_debounce_allows_single_concurrent_scan( fake_time: FakeTimeController, debounced_store: CollectionBasedLightningStore[Any] ) -> None: """When multiple coroutines race, only one should trigger the scan.""" fake_time.set(300.0) results = await asyncio.gather( debounced_store._should_scan_for_unhealthy_rollouts(), # pyright: ignore[reportPrivateUsage] debounced_store._should_scan_for_unhealthy_rollouts(), # pyright: ignore[reportPrivateUsage] debounced_store._should_scan_for_unhealthy_rollouts(), # pyright: ignore[reportPrivateUsage] ) assert results.count(True) == 1 assert results.count(False) == 2 @pytest.mark.asyncio async def test_scan_debounce_disabled_when_zero(store_fixture: CollectionBasedLightningStore[Any]) -> None: """Setting debounce to zero should run the scan every time.""" assert store_fixture._scan_debounce_seconds == 0 # pyright: ignore[reportPrivateUsage] assert await store_fixture._should_scan_for_unhealthy_rollouts() # pyright: ignore[reportPrivateUsage] assert await store_fixture._should_scan_for_unhealthy_rollouts() # pyright: ignore[reportPrivateUsage]