github-featured/microsoft--agent-lightning
2
0
Fork 0
Code Issues Pull Requests Deployments Models Agent Notes Packages Projects Releases Wiki Activity

refactor: PATCH semantics for update_rollout — drop expected_version

Browse Source 
UpdateRolloutRequest → PatchRolloutRequest:
- All fields optional (true partial update)
- Uses model_dump(exclude_unset=True) to apply only sent fields
- Absent field = untouched, explicit null = clear field
- No expected_version — no GET-before-PATCH round trip needed

Controller can now just send:
  PATCH /api/rollouts/{rid} {"status": "succeeded", "succeeded_attempt_id": "pod-uid"}

State transition validation still enforced server-side.
Version field kept on Rollout (informational, bumped internally).

Updated architecture doc, controller pseudocode, todo.md.
New tests: empty_patch_is_noop, update_without_status, explicit_null_clears_field.
114 tests passing, ruff clean.
This commit is contained in:
Yuqing
2026-03-19 23:21:44 +08:00
parent 894e17b58f
commit 95d2bc188f
8 changed files with 145 additions and 137 deletions
Download Patch File Download Diff File Expand all files Collapse all files
agl_lite/schemas/api.py
+7 -5
View File
@@ -6,7 +6,7 @@ from typing import Any
from pydantic import BaseModel, Field
from agl_lite.schemas.rollout import RolloutConfig
from agl_lite.schemas.rollout import RolloutConfig, RolloutStatus
# --- Rollout API ---
@@ -27,11 +27,13 @@ class EnqueueBatchRequest(BaseModel):
rollouts: list[EnqueueRolloutRequest]
class UpdateRolloutRequest(BaseModel):
"""Controller updates rollout status (with optimistic locking)."""
class PatchRolloutRequest(BaseModel):
"""Partial update for a rollout. Only fields present in the request body are applied.
status: str
expected_version: int
Use `model_dump(exclude_unset=True)` to get only the fields the caller explicitly set.
"""
status: RolloutStatus | None = None
job_name: str | None = None
succeeded_attempt_id: str | None = None
error_message: str | None = None
agl_lite/store/memory.py
+16 -16
View File
@@ -12,8 +12,8 @@ import uuid
from pathlib import Path
from typing import Any
from agl_lite.schemas.api import ArchiveBackend, ArchiveResult, EnqueueRolloutRequest, UpdateRolloutRequest
from agl_lite.schemas.errors import ConflictError, InvalidTransitionError, NotFoundError
from agl_lite.schemas.api import ArchiveBackend, ArchiveResult, EnqueueRolloutRequest, PatchRolloutRequest
from agl_lite.schemas.errors import InvalidTransitionError, NotFoundError
from agl_lite.schemas.event import Event
from agl_lite.schemas.model_server import ModelServer
from agl_lite.schemas.resources import ResourcesUpdate
@@ -72,35 +72,35 @@ class InMemoryStore:
"""Check if a rollout exists. Used by gateway for fast validation (~100ns)."""
return rollout_id in self._rollouts
def update_rollout(self, rollout_id: str, req: UpdateRolloutRequest) -> Rollout:
"""Update rollout status with optimistic locking and transition validation.
def update_rollout(self, rollout_id: str, req: PatchRolloutRequest) -> Rollout:
"""Partial update of a rollout. Only fields explicitly set in req are applied.
If `status` is being changed, validates the state transition.
Bumps `version` and `updated_at` on every successful update.
Raises:
NotFoundError: rollout doesn't exist
ConflictError: version mismatch
InvalidTransitionError: illegal state transition
"""
rollout = self.get_rollout(rollout_id)
updates = req.model_dump(exclude_unset=True)
# Optimistic locking.
if rollout.version != req.expected_version:
raise ConflictError("Rollout", rollout_id, req.expected_version, rollout.version)
if not updates:
return rollout # no-op
# State transition validation.
status = RolloutStatus(req.status)
if status not in VALID_TRANSITIONS[rollout.status]:
raise InvalidTransitionError(rollout_id, rollout.status, status)
# Validate state transition if status is changing.
if "status" in updates:
new_status = updates["status"]
if new_status not in VALID_TRANSITIONS[rollout.status]:
raise InvalidTransitionError(rollout_id, rollout.status, new_status)
# Apply update.
now = time.time()
updated = rollout.model_copy(
update={
"status": status,
**updates,
"version": rollout.version + 1,
"updated_at": now,
**({"job_name": req.job_name} if req.job_name is not None else {}),
**({"succeeded_attempt_id": req.succeeded_attempt_id} if req.succeeded_attempt_id is not None else {}),
**({"error_message": req.error_message} if req.error_message is not None else {}),
}
)
self._rollouts[rollout_id] = updated
dev/todo.md
+2 -2
View File
@@ -16,7 +16,7 @@
## Completed
- [x] **Phase 0: Schemas and Project Skeleton** — frozen data models, project structure, dev tooling, 36 schema tests
- [x] **Phase 0.4: State transition rules** — valid transitions table, cancel_requested rules, optimistic locking, exhaustive tests
- [x] **Phase 0.4: State transition rules** — valid transitions table, cancel_requested rules, exhaustive tests
- [x] **Phase 1: In-Memory Store** — `InMemoryStore` with all operations, 74 store tests (110 total)
- Rollout: enqueue, update (transition + version check), cancel, query
- Events: add, query (smart attempt_id resolution), list_attempts
@@ -177,7 +177,7 @@ These are settled and should not be revisited during implementation:
| `job_defaults` schema | Typed `JobDefaults` model, validated at POST time. Known fields validated; `overrides` dict for unknown K8s fields. |
| Auth | API keys, 3 roles, `OPENAI_API_KEY` trick for agents |
| Health endpoint | `GET /healthz`, no auth |
| Error codes | 401 missing/invalid key, 403 wrong role, 404 rollout not found, 409 version conflict / invalid transition |
| Error codes | 401 missing/invalid key, 403 wrong role, 404 rollout not found, 409 invalid transition |
| Archive format | JSONL, user-specified file path (`*.jsonl`). Append if file exists, create if not. Includes rollout + events + resources per archive call. |
| Gateway config | Static YAML at startup (param adjustment), no runtime changes |
| Rollout existence check | On both LLM proxy and event ingestion (in-process, ~100ns) |
docs/design/0_architecture.md
+25 -39
View File
@@ -356,7 +356,7 @@ class Rollout:
error_message: Optional[str] # error info (set on terminal_failed)
# Concurrency control
version: int # optimistic locking — incremented on every update
version: int # incremented on every update (informational)
created_at: float
updated_at: float
@@ -437,11 +437,10 @@ queuing ──[controller creates Job]──────────→ running
class Store:
# Rollout management
async def enqueue_rollout(input, config, resources_id=None) -> Rollout
async def update_rollout(rollout_id, status, expected_version,
job_name=None, succeeded_attempt_id=None,
error_message=None) -> Rollout
# Enforces: valid transition + optimistic locking (version check)
# Raises: ConflictError (version mismatch), InvalidTransitionError
async def update_rollout(rollout_id, **updates) -> Rollout
# Partial update — only provided fields are changed.
# Validates state transition if `status` is included.
# Raises: InvalidTransitionError
async def cancel_rollout(rollout_id) -> Rollout
# Sets cancel_requested=True. Rejects if already terminal.
async def query_rollouts(ids=None, status_in=None, cancel_requested=None,
@@ -495,7 +494,7 @@ The Gateway (LLM proxy) and Store (data management) are combined into a **single
| `POST` | `/rollout/{rid}/attempt/{aid}/v1/...` | **LLM reverse proxy** — forwards to model server, auto-captures `model_request` events. Supports both OpenAI (`/v1/chat/completions`) and Anthropic (`/v1/messages`) formats. | Agent pods |
| `POST` | `/rollout/{rid}/attempt/{aid}/events` | **Event ingestion** — accepts reward and user-defined events | Agent pods, runner, environment |
| `POST` `GET` | `/api/rollouts` | **Rollout management** — enqueue, query (with batch ID support) | Algorithm, K8s controller |
| `GET` `PATCH` | `/api/rollouts/{rid}` | **Single rollout** — get, update (with optimistic locking) | K8s controller |
| `GET` `PATCH` | `/api/rollouts/{rid}` | **Single rollout** — get, partial update | K8s controller |
| `POST` | `/api/rollouts/{rid}/cancel` | **Cancel rollout** — set cancel_requested flag | Algorithm, user |
| `POST` | `/api/rollouts/archive` | **Data lifecycle** — archive and purge consumed rollouts (optional JSONL persistence) | Algorithm |
| `POST` `GET` `DELETE` | `/api/models` | **Model server management** — register, list, remove inference servers | Algorithm / Compute Backend |
@@ -586,7 +585,7 @@ The gateway is a single Python async process. The concurrency profile is excelle
|----------|---------------|------------|
| Model server registry | Read every request, write once per weight update | Near zero (read-heavy, write-rare) |
| Event store per `(rid, aid)` | Append per event, naturally partitioned by pod | Near zero (different agents never touch the same partition) |
| Rollout records | Controller updates status, Algorithm enqueues | Low (not on hot path, optimistic locking) |
| Rollout records | Controller updates status, Algorithm enqueues | Low (not on hot path) |
No locks needed on the hot path. Single-threaded asyncio serializes all writes naturally. The partition key `(rollout_id, attempt_id)` eliminates cross-agent contention entirely.
@@ -650,7 +649,7 @@ The archive feature thus serves dual purpose: **data lifecycle** (purge hot stor
| `POST` | `/api/rollouts` | Enqueue rollout(s). Body: single `{input, config?, resources_id?}` or batch `{config, resources_id?, rollouts: [{input, config?}, ...]}`. Batch-level `config` and `resources_id` apply to all rollouts; per-rollout fields override. Returns `Rollout` or `List[Rollout]` with status `queuing`. |
| `GET` | `/api/rollouts` | Query rollouts. Params: `ids` (comma-separated for batch fetch), `status_in`, `cancel_requested`, `limit`, `offset`. Returns `List[Rollout]`. |
| `GET` | `/api/rollouts/{rollout_id}` | Get a single rollout by ID. Returns `Rollout`. |
| `PATCH` | `/api/rollouts/{rollout_id}` | Update rollout status. Body: `{status, expected_version, job_name?, succeeded_attempt_id?, error_message?}`. Enforces valid transitions + optimistic locking. Used by K8s controller. |
| `PATCH` | `/api/rollouts/{rollout_id}` | Partial update. Body: any subset of `{status, job_name, succeeded_attempt_id, error_message}`. Only fields present in the body are applied. Validates state transitions when `status` is included. Used by K8s controller. |
| `POST` | `/api/rollouts/{rollout_id}/cancel` | Set `cancel_requested=true`. Rejects if already terminal. Used by Algorithm or user. |
**Event / trajectory access:**
@@ -963,15 +962,13 @@ def handle_cancel(rollout):
if rollout.status == QUEUING and rollout.job_name is None:
# No Job exists. Straight to cancelled.
store.update_rollout(rollout.rollout_id,
status=CANCELLED,
expected_version=rollout.version)
status=CANCELLED)
return
if rollout.job_name is None:
# Running but no job_name? Shouldn't happen, but be safe.
store.update_rollout(rollout.rollout_id,
status=CANCELLED,
expected_version=rollout.version)
status=CANCELLED)
return
job = k8s.get_job(rollout.job_name)
@@ -979,8 +976,7 @@ def handle_cancel(rollout):
if job is None:
# Job already gone. Mark cancelled.
store.update_rollout(rollout.rollout_id,
status=CANCELLED,
expected_version=rollout.version)
status=CANCELLED)
return
# Job exists. Check if it already succeeded before we delete.
@@ -989,8 +985,7 @@ def handle_cancel(rollout):
succeeded_pod_uid = find_succeeded_pod_uid(job)
store.update_rollout(rollout.rollout_id,
status=SUCCEEDED,
succeeded_attempt_id=succeeded_pod_uid,
expected_version=rollout.version)
succeeded_attempt_id=succeeded_pod_uid)
k8s.delete_job(rollout.job_name)
return
@@ -998,8 +993,7 @@ def handle_cancel(rollout):
# Job already failed on its own. User wanted cancel — mark cancelled,
# not terminal_failed. The intent was cancellation.
store.update_rollout(rollout.rollout_id,
status=CANCELLED,
expected_version=rollout.version)
status=CANCELLED)
k8s.delete_job(rollout.job_name)
return
@@ -1020,32 +1014,27 @@ def handle_queuing(rollout):
job = k8s.get_job(rollout.job_name)
if job is not None:
store.update_rollout(rollout.rollout_id,
status=RUNNING, job_name=rollout.job_name,
expected_version=rollout.version)
status=RUNNING, job_name=rollout.job_name)
return
# Job name set but Job gone? Something went wrong.
store.update_rollout(rollout.rollout_id,
status=TERMINAL_FAILED,
error_message="Job not found during recovery",
expected_version=rollout.version)
error_message="Job not found during recovery")
return
job_name = f"agl-rollout-{rollout.rollout_id}"
try:
k8s.create_job(make_job_spec(rollout, job_name))
store.update_rollout(rollout.rollout_id,
status=RUNNING, job_name=job_name,
expected_version=rollout.version)
status=RUNNING, job_name=job_name)
except K8sAlreadyExistsError:
# Job exists (duplicate from previous attempt). Fetch and proceed.
store.update_rollout(rollout.rollout_id,
status=RUNNING, job_name=job_name,
expected_version=rollout.version)
status=RUNNING, job_name=job_name)
except K8sError as e:
store.update_rollout(rollout.rollout_id,
status=TERMINAL_FAILED,
error_message=f"Job creation failed: {e}",
expected_version=rollout.version)
error_message=f"Job creation failed: {e}")
def handle_running(rollout):
@@ -1056,8 +1045,7 @@ def handle_running(rollout):
# Job disappeared (manually deleted, namespace cleanup).
store.update_rollout(rollout.rollout_id,
status=TERMINAL_FAILED,
error_message="K8s Job not found",
expected_version=rollout.version)
error_message="K8s Job not found")
return
conditions = {c.type: c for c in (job.status.conditions or [])}
@@ -1066,19 +1054,17 @@ def handle_running(rollout):
succeeded_pod_uid = find_succeeded_pod_uid(job)
store.update_rollout(rollout.rollout_id,
status=SUCCEEDED,
succeeded_attempt_id=succeeded_pod_uid,
expected_version=rollout.version)
succeeded_attempt_id=succeeded_pod_uid)
elif "Failed" in conditions:
reason = conditions["Failed"].reason # BackoffLimitExceeded, DeadlineExceeded
message = conditions["Failed"].message
store.update_rollout(rollout.rollout_id,
status=TERMINAL_FAILED,
error_message=f"{reason}: {message}",
expected_version=rollout.version)
error_message=f"{reason}: {message}")
# else: Job still active (running or between retries). No update needed.
```
All `update_rollout` calls use optimistic locking (`expected_version`). On `ConflictError`, the controller re-fetches the rollout and re-evaluates — another instance may have already handled it.
All `update_rollout` calls are partial updates — only the fields included in the PATCH body are changed. State transition validation is enforced server-side. The controller is the sole writer for status transitions, so no locking is needed.
#### Edge cases
@@ -1086,7 +1072,7 @@ All `update_rollout` calls use optimistic locking (`expected_version`). On `Conf
On restart, periodic full reconciliation scans all non-terminal rollouts and syncs them with K8s Job status. This is idempotent — if a Job exists, its status is checked; if it's gone, the rollout is marked `terminal_failed`. The deterministic Job name (`agl-rollout-{rollout_id}`) ensures the controller can always find the Job for a rollout.
**Two controller instances (leader election gap):**
Both read `version=N`, both try to update the same rollout. One succeeds (`version→N+1`), the other gets `ConflictError`, re-fetches, sees the update was already done. Optimistic locking is the single serialization point.
Both read the same rollout, both try to update. State transition validation prevents invalid updates — e.g., if one controller already moved it to `succeeded`, the other's `PATCH {status: succeeded}` is a no-op (same state) or rejected (if trying a different transition). Single-threaded asyncio in the Store ensures no interleaving.
**Store unavailable:**
The K8s controller pattern naturally handles this: if `update_rollout` fails due to Store being unreachable, the event is requeued with exponential backoff. The Job keeps running regardless of Store availability. When the Store comes back, the controller retries. No data loss — K8s Job status is the durable record.
@@ -1263,7 +1249,7 @@ A single HTTP service with **~18 endpoints** across 6 domains:
| `POST /api/rollouts` | `enqueue_rollout`, `enqueue_many_rollouts` (batch via JSON array body) |
| `GET /api/rollouts` | `query_rollouts` + `wait_for_rollouts` (params: `ids` for batch fetch, `status_in`, `cancel_requested`, `limit`, `offset`). Waiting is client-side polling. |
| `GET /api/rollouts/{rid}` | `get_rollout_by_id` |
| `PATCH /api/rollouts/{rid}` | `update_rollout` (with optimistic locking via `expected_version`) |
| `PATCH /api/rollouts/{rid}` | `update_rollout` — partial update, only provided fields changed. Validates state transitions. |
| `POST /api/rollouts/{rid}/cancel` | *New.* Sets `cancel_requested` flag. Original used `update_rollout(status="cancelled")`. |
**Model server management (4 endpoints):**
@@ -1328,7 +1314,7 @@ Attempt listing is folded into `GET /api/rollouts/{rid}` response (includes `att
| `POST /rollout/{rid}/attempt/{aid}/events` | Explicit event ingestion (reward, user-defined types). Original had no direct event API — everything went through OTEL spans. |
| `POST /api/rollouts/{rid}/cancel` | Explicit cancel with `cancel_requested` flag. Cleaner than overloading `update_rollout(status="cancelled")`. |
| `cancel_requested` flag on Rollout | Separate intent from execution. Original used status directly. |
| `expected_version` on `PATCH /api/rollouts/{rid}` | Optimistic locking for safe concurrent updates. Original relied on in-process thread locks or single-writer patterns. |
| `version` on `Rollout` | Monotonically incrementing counter, bumped on every update. Informational — included in responses for debugging/observability. Not required in PATCH requests. Original relied on in-process thread locks or single-writer patterns. |
| `succeeded_attempt_id` on Rollout | Directly links successful rollout to its trajectory data. Original required querying attempts to find the successful one. |
| Open event types (`event_type: str`) | Extensible without schema changes. Original was locked to OTEL span format. |
| Unified service (proxy + store) | Single deployment, in-process event capture on hot path. Original had separate LLM Proxy and Store Server. |
tests/schemas/test_api.py
+20 -12
View File
@@ -5,11 +5,11 @@ from agl_lite.schemas.api import (
ArchiveRequest,
EnqueueBatchRequest,
EnqueueRolloutRequest,
PatchRolloutRequest,
PostEventRequest,
RegisterModelRequest,
UpdateRolloutRequest,
)
from agl_lite.schemas.rollout import RolloutConfig
from agl_lite.schemas.rollout import RolloutConfig, RolloutStatus
class TestEnqueueRolloutRequest:
@@ -41,18 +41,26 @@ class TestEnqueueBatchRequest:
assert b.config.image == "agent:v1"
class TestUpdateRolloutRequest:
def test_status_update(self):
u = UpdateRolloutRequest(status="running", expected_version=1)
assert u.job_name is None
class TestPatchRolloutRequest:
def test_status_only(self):
p = PatchRolloutRequest(status=RolloutStatus.RUNNING)
dumped = p.model_dump(exclude_unset=True)
assert dumped == {"status": RolloutStatus.RUNNING}
def test_with_optional_fields(self):
u = UpdateRolloutRequest(
status="succeeded",
expected_version=2,
succeeded_attempt_id="pod-uid-1",
)
assert u.succeeded_attempt_id == "pod-uid-1"
p = PatchRolloutRequest(status=RolloutStatus.SUCCEEDED, succeeded_attempt_id="pod-uid-1")
dumped = p.model_dump(exclude_unset=True)
assert dumped == {"status": RolloutStatus.SUCCEEDED, "succeeded_attempt_id": "pod-uid-1"}
def test_empty_patch(self):
p = PatchRolloutRequest()
dumped = p.model_dump(exclude_unset=True)
assert dumped == {}
def test_non_status_field_only(self):
p = PatchRolloutRequest(job_name="agl-rollout-abc")
dumped = p.model_dump(exclude_unset=True)
assert dumped == {"job_name": "agl-rollout-abc"}
class TestPostEventRequest:
tests/store/test_archive.py
+10 -12
View File
@@ -5,7 +5,7 @@ from pathlib import Path
import pytest
from agl_lite.schemas.api import ArchiveBackend, EnqueueRolloutRequest, UpdateRolloutRequest
from agl_lite.schemas.api import ArchiveBackend, EnqueueRolloutRequest, PatchRolloutRequest
from agl_lite.schemas.errors import NotFoundError
from agl_lite.schemas.rollout import RolloutConfig, RolloutStatus
from agl_lite.store.memory import InMemoryStore
@@ -22,10 +22,8 @@ def _enqueue(store: InMemoryStore, **kwargs):
return store.enqueue_rollout(EnqueueRolloutRequest(**kwargs))
def _update(store: InMemoryStore, rollout_id: str, status: str, expected_version: int, **kwargs):
return store.update_rollout(
rollout_id, UpdateRolloutRequest(status=status, expected_version=expected_version, **kwargs)
)
def _patch(store: InMemoryStore, rollout_id: str, **kwargs):
return store.update_rollout(rollout_id, PatchRolloutRequest(**kwargs))
def _make_terminal_rollout(
@@ -34,10 +32,10 @@ def _make_terminal_rollout(
"""Helper: create a rollout and move it to a terminal state. Returns rollout_id."""
r = _enqueue(store, **enqueue_kwargs)
if status == RolloutStatus.SUCCEEDED:
_update(store, r.rollout_id, "running", 1)
_update(store, r.rollout_id, "succeeded", 2, succeeded_attempt_id="pod-1")
_patch(store, r.rollout_id, status=RolloutStatus.RUNNING)
_patch(store, r.rollout_id, status=RolloutStatus.SUCCEEDED, succeeded_attempt_id="pod-1")
else:
_update(store, r.rollout_id, status.value, 1)
_patch(store, r.rollout_id, status=status)
return r.rollout_id
@@ -87,8 +85,8 @@ class TestArchivePurge:
def test_jsonl_includes_resources(self, store: InMemoryStore, tmp_path: Path):
res = store.add_resources({"system_prompt": "Be helpful"})
r = _enqueue(store, resources_id=res.resources_id)
_update(store, r.rollout_id, "running", 1)
_update(store, r.rollout_id, "succeeded", 2, succeeded_attempt_id="pod-1")
_patch(store, r.rollout_id, status=RolloutStatus.RUNNING)
_patch(store, r.rollout_id, status=RolloutStatus.SUCCEEDED, succeeded_attempt_id="pod-1")
archive_path = tmp_path / "archive.jsonl"
store.archive_rollouts([r.rollout_id], backend=ArchiveBackend(path=str(archive_path)))
@@ -132,8 +130,8 @@ class TestArchivePurge:
r = _enqueue(store)
store.add_event(r.rollout_id, "pod-1", "model_request", {"attempt": 1})
store.add_event(r.rollout_id, "pod-2", "model_request", {"attempt": 2})
_update(store, r.rollout_id, "running", 1)
_update(store, r.rollout_id, "succeeded", 2, succeeded_attempt_id="pod-2")
_patch(store, r.rollout_id, status=RolloutStatus.RUNNING)
_patch(store, r.rollout_id, status=RolloutStatus.SUCCEEDED, succeeded_attempt_id="pod-2")
archive_path = tmp_path / "archive.jsonl"
store.archive_rollouts([r.rollout_id], backend=ArchiveBackend(path=str(archive_path)))
tests/store/test_events.py
+6 -8
View File
@@ -4,9 +4,9 @@ import time
import pytest
from agl_lite.schemas.api import EnqueueRolloutRequest, UpdateRolloutRequest
from agl_lite.schemas.api import EnqueueRolloutRequest, PatchRolloutRequest
from agl_lite.schemas.errors import NotFoundError
from agl_lite.schemas.rollout import RolloutConfig
from agl_lite.schemas.rollout import RolloutConfig, RolloutStatus
from agl_lite.store.memory import InMemoryStore
@@ -21,10 +21,8 @@ def _enqueue(store: InMemoryStore, **kwargs):
return store.enqueue_rollout(EnqueueRolloutRequest(**kwargs))
def _update(store: InMemoryStore, rollout_id: str, status: str, expected_version: int, **kwargs):
return store.update_rollout(
rollout_id, UpdateRolloutRequest(status=status, expected_version=expected_version, **kwargs)
)
def _patch(store: InMemoryStore, rollout_id: str, **kwargs):
return store.update_rollout(rollout_id, PatchRolloutRequest(**kwargs))
class TestAddEvent:
@@ -116,8 +114,8 @@ class TestSmartAttemptResolution:
r = _enqueue(store)
store.add_event(r.rollout_id, "pod-1", "model_request", {"attempt": "failed"})
store.add_event(r.rollout_id, "pod-2", "model_request", {"attempt": "succeeded"})
_update(store, r.rollout_id, "running", 1)
_update(store, r.rollout_id, "succeeded", 2, succeeded_attempt_id="pod-2")
_patch(store, r.rollout_id, status=RolloutStatus.RUNNING)
_patch(store, r.rollout_id, status=RolloutStatus.SUCCEEDED, succeeded_attempt_id="pod-2")
events = store.query_events(r.rollout_id) # no attempt_id
assert len(events) == 1
assert events[0].data["attempt"] == "succeeded"
tests/store/test_rollouts.py
+59 -43
View File
@@ -2,8 +2,8 @@
import pytest
from agl_lite.schemas.api import EnqueueRolloutRequest, UpdateRolloutRequest
from agl_lite.schemas.errors import ConflictError, InvalidTransitionError, NotFoundError
from agl_lite.schemas.api import EnqueueRolloutRequest, PatchRolloutRequest
from agl_lite.schemas.errors import InvalidTransitionError, NotFoundError
from agl_lite.schemas.rollout import TERMINAL_STATUSES, RolloutConfig, RolloutStatus
from agl_lite.store.memory import InMemoryStore
@@ -20,11 +20,9 @@ def _enqueue(store: InMemoryStore, **kwargs):
return store.enqueue_rollout(EnqueueRolloutRequest(**kwargs))
def _update(store: InMemoryStore, rollout_id: str, status: str, expected_version: int, **kwargs):
"""Helper: update rollout status."""
return store.update_rollout(
rollout_id, UpdateRolloutRequest(status=status, expected_version=expected_version, **kwargs)
)
def _patch(store: InMemoryStore, rollout_id: str, **kwargs):
"""Helper: partial update."""
return store.update_rollout(rollout_id, PatchRolloutRequest(**kwargs))
class TestEnqueueRollout:
@@ -78,40 +76,42 @@ class TestRolloutExists:
class TestUpdateRollout:
def test_queuing_to_running(self, store: InMemoryStore):
r = _enqueue(store)
updated = _update(store, r.rollout_id, "running", 1, job_name="agl-rollout-x")
updated = _patch(store, r.rollout_id, status=RolloutStatus.RUNNING, job_name="agl-rollout-x")
assert updated.status == RolloutStatus.RUNNING
assert updated.version == 2
assert updated.job_name == "agl-rollout-x"
def test_running_to_succeeded(self, store: InMemoryStore):
r = _enqueue(store)
_update(store, r.rollout_id, "running", 1)
updated = _update(store, r.rollout_id, "succeeded", 2, succeeded_attempt_id="pod-uid-1")
_patch(store, r.rollout_id, status=RolloutStatus.RUNNING)
updated = _patch(store, r.rollout_id, status=RolloutStatus.SUCCEEDED, succeeded_attempt_id="pod-uid-1")
assert updated.status == RolloutStatus.SUCCEEDED
assert updated.succeeded_attempt_id == "pod-uid-1"
assert updated.version == 3
def test_running_to_terminal_failed(self, store: InMemoryStore):
r = _enqueue(store)
_update(store, r.rollout_id, "running", 1)
updated = _update(store, r.rollout_id, "terminal_failed", 2, error_message="BackoffLimitExceeded")
_patch(store, r.rollout_id, status=RolloutStatus.RUNNING)
updated = _patch(
store, r.rollout_id, status=RolloutStatus.TERMINAL_FAILED, error_message="BackoffLimitExceeded"
)
assert updated.status == RolloutStatus.TERMINAL_FAILED
assert updated.error_message == "BackoffLimitExceeded"
def test_running_to_cancelled(self, store: InMemoryStore):
r = _enqueue(store)
_update(store, r.rollout_id, "running", 1)
updated = _update(store, r.rollout_id, "cancelled", 2)
_patch(store, r.rollout_id, status=RolloutStatus.RUNNING)
updated = _patch(store, r.rollout_id, status=RolloutStatus.CANCELLED)
assert updated.status == RolloutStatus.CANCELLED
def test_queuing_to_terminal_failed(self, store: InMemoryStore):
r = _enqueue(store)
updated = _update(store, r.rollout_id, "terminal_failed", 1, error_message="Job creation failed")
updated = _patch(store, r.rollout_id, status=RolloutStatus.TERMINAL_FAILED, error_message="Job creation failed")
assert updated.status == RolloutStatus.TERMINAL_FAILED
def test_queuing_to_cancelled(self, store: InMemoryStore):
r = _enqueue(store)
updated = _update(store, r.rollout_id, "cancelled", 1)
updated = _patch(store, r.rollout_id, status=RolloutStatus.CANCELLED)
assert updated.status == RolloutStatus.CANCELLED
# --- Invalid transitions ---
@@ -119,52 +119,68 @@ class TestUpdateRollout:
def test_queuing_to_succeeded_rejected(self, store: InMemoryStore):
r = _enqueue(store)
with pytest.raises(InvalidTransitionError):
_update(store, r.rollout_id, "succeeded", 1)
_patch(store, r.rollout_id, status=RolloutStatus.SUCCEEDED)
def test_running_to_queuing_rejected(self, store: InMemoryStore):
r = _enqueue(store)
_update(store, r.rollout_id, "running", 1)
_patch(store, r.rollout_id, status=RolloutStatus.RUNNING)
with pytest.raises(InvalidTransitionError):
_update(store, r.rollout_id, "queuing", 2)
_patch(store, r.rollout_id, status=RolloutStatus.QUEUING)
def test_terminal_to_anything_rejected(self, store: InMemoryStore):
for terminal in TERMINAL_STATUSES:
s = InMemoryStore()
r = _enqueue(s)
if terminal == RolloutStatus.SUCCEEDED:
_update(s, r.rollout_id, "running", 1)
_update(s, r.rollout_id, terminal.value, 2)
version = 3
_patch(s, r.rollout_id, status=RolloutStatus.RUNNING)
_patch(s, r.rollout_id, status=terminal)
else:
_update(s, r.rollout_id, terminal.value, 1)
version = 2
_patch(s, r.rollout_id, status=terminal)
for target in RolloutStatus:
with pytest.raises(InvalidTransitionError):
_update(s, r.rollout_id, target.value, version)
# --- Optimistic locking ---
def test_version_mismatch(self, store: InMemoryStore):
r = _enqueue(store)
with pytest.raises(ConflictError, match="expected version 99"):
_update(store, r.rollout_id, "running", 99)
_patch(s, r.rollout_id, status=target)
def test_not_found(self, store: InMemoryStore):
with pytest.raises(NotFoundError):
_update(store, "nonexistent", "running", 1)
_patch(store, "nonexistent", status=RolloutStatus.RUNNING)
# --- Optional fields preserved ---
# --- Partial update semantics ---
def test_optional_fields_not_overwritten(self, store: InMemoryStore):
"""Passing None for optional fields should not overwrite existing values."""
def test_only_set_fields_applied(self, store: InMemoryStore):
"""Fields not in request body are untouched."""
r = _enqueue(store)
_update(store, r.rollout_id, "running", 1, job_name="my-job")
# Update to succeeded without re-specifying job_name.
updated = _update(store, r.rollout_id, "succeeded", 2, succeeded_attempt_id="pod-1")
_patch(store, r.rollout_id, status=RolloutStatus.RUNNING, job_name="my-job")
# Patch status only — job_name should be preserved.
updated = _patch(store, r.rollout_id, status=RolloutStatus.SUCCEEDED, succeeded_attempt_id="pod-1")
assert updated.job_name == "my-job" # preserved
assert updated.succeeded_attempt_id == "pod-1"
def test_update_without_status(self, store: InMemoryStore):
"""Can update non-status fields without changing status."""
r = _enqueue(store)
_patch(store, r.rollout_id, status=RolloutStatus.RUNNING)
updated = _patch(store, r.rollout_id, job_name="agl-rollout-xyz")
assert updated.status == RolloutStatus.RUNNING # unchanged
assert updated.job_name == "agl-rollout-xyz"
assert updated.version == 3
def test_empty_patch_is_noop(self, store: InMemoryStore):
"""Empty body = no changes, no version bump."""
r = _enqueue(store)
updated = store.update_rollout(r.rollout_id, PatchRolloutRequest())
assert updated.version == 1 # no bump
def test_explicit_null_clears_field(self, store: InMemoryStore):
"""Explicitly sending null should set the field to None."""
r = _enqueue(store)
_patch(store, r.rollout_id, status=RolloutStatus.RUNNING, job_name="my-job")
# Explicitly clear job_name by constructing request with job_name=None set.
req = PatchRolloutRequest.model_validate({"job_name": None})
updated = store.update_rollout(r.rollout_id, req)
assert updated.job_name is None
assert updated.version == 3
class TestCancelRollout:
def test_cancel_queuing(self, store: InMemoryStore):
@@ -175,7 +191,7 @@ class TestCancelRollout:
def test_cancel_running(self, store: InMemoryStore):
r = _enqueue(store)
_update(store, r.rollout_id, "running", 1)
_patch(store, r.rollout_id, status=RolloutStatus.RUNNING)
updated = store.cancel_rollout(r.rollout_id)
assert updated.cancel_requested is True
@@ -187,7 +203,7 @@ class TestCancelRollout:
def test_cancel_terminal_rejected(self, store: InMemoryStore):
r = _enqueue(store)
_update(store, r.rollout_id, "terminal_failed", 1)
_patch(store, r.rollout_id, status=RolloutStatus.TERMINAL_FAILED)
with pytest.raises(InvalidTransitionError, match="cancel_requested"):
store.cancel_rollout(r.rollout_id)
@@ -219,7 +235,7 @@ class TestQueryRollouts:
def test_filter_status(self, store: InMemoryStore):
r1 = _enqueue(store)
_enqueue(store)
_update(store, r1.rollout_id, "running", 1)
_patch(store, r1.rollout_id, status=RolloutStatus.RUNNING)
results = store.query_rollouts(status_in=[RolloutStatus.RUNNING])
assert len(results) == 1
assert results[0].rollout_id == r1.rollout_id
@@ -237,7 +253,7 @@ class TestQueryRollouts:
r2 = _enqueue(store)
store.cancel_rollout(r1.rollout_id)
store.cancel_rollout(r2.rollout_id)
_update(store, r1.rollout_id, "running", 2)
_patch(store, r1.rollout_id, status=RolloutStatus.RUNNING)
results = store.query_rollouts(status_in=[RolloutStatus.RUNNING], cancel_requested=True)
assert len(results) == 1
assert results[0].rollout_id == r1.rollout_id