File size: 2,347 Bytes
1866d85 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 | 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 # publisher status advanced
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'
)
|