| from __future__ import annotations |
|
|
| import asyncio |
| from pathlib import Path |
| from typing import Any |
|
|
| from autocad_bench.orchestration.trace_stream import IncrementalTracePublisher |
|
|
|
|
| class FakeS3: |
| def __init__(self) -> None: |
| self.objects: dict[tuple[str, str], bytes] = {} |
| self.puts: list[str] = [] |
|
|
| def put_object(self, **kwargs: Any) -> None: |
| key = str(kwargs["Key"]) |
| self.puts.append(key) |
| self.objects[(str(kwargs["Bucket"]), key)] = bytes(kwargs["Body"]) |
|
|
| def upload_file( |
| self, |
| filename: str, |
| bucket: str, |
| key: str, |
| **_kwargs: Any, |
| ) -> None: |
| self.puts.append(key) |
| self.objects[(bucket, key)] = Path(filename).read_bytes() |
|
|
|
|
| def test_incremental_trace_publisher_streams_only_changed_live_artifacts( |
| tmp_path: Path, |
| ) -> None: |
| root = tmp_path / "run" |
| frame = root / "rollouts" / "task-001" / "frames" / "frame_0000.png" |
| state = root / "rollouts" / "task-001" / "run-state.json" |
| raw_turns = root / "rollouts" / "task-001" / "model_turns.jsonl" |
| frame.parent.mkdir(parents=True) |
| frame.write_bytes(b"png") |
| state.write_text('{"state":"running"}\n', encoding="utf-8") |
| raw_turns.write_text('{"turn":1}\n', encoding="utf-8") |
| s3 = FakeS3() |
| publisher = IncrementalTracePublisher( |
| root=root, |
| bucket="bucket", |
| prefix="controller/runs/benchmark", |
| s3=s3, |
| interval_s=0.01, |
| ) |
|
|
| asyncio.run(publisher.publish_once()) |
|
|
| assert ( |
| s3.objects[ |
| ( |
| "bucket", |
| "controller/runs/benchmark/rollouts/task-001/frames/frame_0000.png", |
| ) |
| ] |
| == b"png" |
| ) |
| assert not any(key.endswith("model_turns.jsonl") for key in s3.puts) |
| first_put_count = len(s3.puts) |
|
|
| asyncio.run(publisher.publish_once()) |
| assert len(s3.puts) == first_put_count + 1 |
|
|
| state.write_text('{"state":"completed"}\n', encoding="utf-8") |
| asyncio.run(publisher.publish_once(final=True)) |
|
|
| assert any(key.endswith("model_turns.jsonl") for key in s3.puts) |
| assert ( |
| s3.objects[ |
| ( |
| "bucket", |
| "controller/runs/benchmark/rollouts/task-001/run-state.json", |
| ) |
| ] |
| == b'{"state":"completed"}\n' |
| ) |
|
|
