diff --git a/deps/verifiers b/deps/verifiers index bccc2c52293..6d5edc7c7f7 160000 --- a/deps/verifiers +++ b/deps/verifiers @@ -1 +1 @@ -Subproject commit bccc2c52293fe2435651de0c9f70436ca7332b9b +Subproject commit 6d5edc7c7f7d178e3511bed06212d911c4eaf9b8 diff --git a/src/prime_rl/monitors/file/monitor.py b/src/prime_rl/monitors/file/monitor.py index 02b560cdf48..3741b89f10e 100644 --- a/src/prime_rl/monitors/file/monitor.py +++ b/src/prime_rl/monitors/file/monitor.py @@ -15,7 +15,7 @@ from prime_rl.monitors.base import Kind, Monitor, Subset from prime_rl.monitors.file.traces import get_annotations_dir, get_index_path, get_trace_stream from prime_rl.monitors.file.traces.chunks import ChunkedJsonl -from prime_rl.monitors.file.traces.index import index_row +from prime_rl.monitors.file.traces.index import index_record, index_row from prime_rl.monitors.file.traces.live import get_live_dir, get_pending_dir from prime_rl.monitors.file.traces.update import update_index_row from prime_rl.utils.pathing import get_eval_plan_path, get_file_monitor_dir @@ -159,14 +159,17 @@ async def log_episodes(self, episodes: list[vf.Episode], step: int, kind: Kind, def write() -> None: stream, index = self._stream(get_trace_stream(self.output_dir)) - # An index row goes out with every episode: summarising the record here, - # while it is already in hand, saves every reader from parsing a stream + # An index row goes out with every episode: summarising the episode here, + # while it is still in hand, saves every reader from parsing a stream # that outgrows memory long before the run does. for episode in episodes: - record = episode.to_record(float_decimals=self.config.float_decimals) - chunk, offset = stream.append(orjson.dumps(record, default=str, option=OPTS)) + # `to_json` encodes in short steps, so a long episode never holds the GIL + # (and the event loop) for the whole encode. + line = episode.to_json(float_decimals=self.config.float_decimals) + chunk, offset = stream.append(line, b"\n") self._logged += 1 - index.write(orjson.dumps(index_row(self._logged, record, chunk, offset), default=str, option=OPTS)) + row = index_row(self._logged, index_record(episode), chunk, offset) + index.write(orjson.dumps(row, default=str, option=OPTS)) stream.flush() index.flush() diff --git a/src/prime_rl/monitors/file/traces/chunks.py b/src/prime_rl/monitors/file/traces/chunks.py index f2ace2036aa..74f6a7f2c92 100644 --- a/src/prime_rl/monitors/file/traces/chunks.py +++ b/src/prime_rl/monitors/file/traces/chunks.py @@ -78,13 +78,16 @@ def __init__(self, directory: Path, max_bytes: int, compress: bool) -> None: self.file = open(chunk_path(directory, self.number), "ab") self.size = self.file.tell() - def append(self, line: bytes) -> tuple[int, int]: - """Write one line; returns ``(chunk, offset)``, the address a reader seeks to.""" - if self.size and self.size + len(line) > self.max_bytes: + def append(self, *parts: bytes) -> tuple[int, int]: + """Write one line, given in parts so a long line is never copied to add its + newline; returns ``(chunk, offset)``, the address a reader seeks to.""" + size = sum(map(len, parts)) + if self.size and self.size + size > self.max_bytes: self._roll() offset = self.size - self.file.write(line) - self.size += len(line) + for part in parts: + self.file.write(part) + self.size += size return self.number, offset def flush(self) -> None: diff --git a/src/prime_rl/monitors/file/traces/index.py b/src/prime_rl/monitors/file/traces/index.py index 96ed720f598..b5d4f65e8e8 100644 --- a/src/prime_rl/monitors/file/traces/index.py +++ b/src/prime_rl/monitors/file/traces/index.py @@ -1,12 +1,35 @@ """The episode index - one compact row per episode, so a reader never has to parse the stream to browse it. -The file monitor writes a row as each episode lands, when the record is already in -hand and summarising it is nearly free, and records the chunk and byte offset it wrote -the episode at so a reader can seek straight to one. A stream another producer wrote -has no index, so its reader derives the same rows itself. +The file monitor writes a row as each episode lands, while the episode is still in +hand, and records the chunk and byte offset it wrote the episode at so a reader can +seek straight to one. A stream another producer wrote has no index, so its reader +derives the same rows itself. """ +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + import verifiers.v1 as vf + + +def index_record(episode: vf.Episode) -> dict[str, Any]: + """The part of the episode's record (`to_record`) that `summarize_episode` reads: + everything but the nodes' payloads. Built one trace at a time, so no single call + holds the GIL for a whole long episode.""" + record = episode.model_dump(mode="json", exclude={"traces"}, exclude_none=True) + record["traces"] = [] + for trace in episode.traces: + head = trace.model_dump(mode="json", exclude={"nodes"}, exclude_none=True) + head["nodes"] = [ + {"sampled": node.sampled} if node.parent is None else {"parent": node.parent, "sampled": node.sampled} + for node in trace.nodes + ] + record["traces"].append(head) + return record + def walk_timing(obj: dict, prefix: str, out: dict[str, float]) -> None: """Flatten a trace timing tree to phase -> seconds (same walk as the viewer).""" diff --git a/src/prime_rl/orchestrator/orchestrator.py b/src/prime_rl/orchestrator/orchestrator.py index 50f72d21472..4892f34cbc7 100644 --- a/src/prime_rl/orchestrator/orchestrator.py +++ b/src/prime_rl/orchestrator/orchestrator.py @@ -235,8 +235,9 @@ async def setup(self) -> None: train_env_names=[env.resolved_name for env in config.train.source], eval_env_names=[source.resolved_name for source in config.eval.source] if config.eval is not None else [], ) - # The launcher-set $PRL_RUN_ID is the run identity; standalone runs mint a local one. - self.run_id = os.environ.get("PRL_RUN_ID") or uuid.uuid4().hex + # The launcher-set $PRL_RUN_ID is the run identity; a hosted launch sets only the + # platform's $RUN_ID, and standalone runs mint a local one. + self.run_id = os.environ.get("PRL_RUN_ID") or os.environ.get("RUN_ID") or uuid.uuid4().hex # Base labels for sandboxes created in this process; env-server processes read # the same launcher-set env var themselves. self.run_name = os.environ.get("PRL_RUN_NAME") diff --git a/tests/unit/monitors/test_monitor_load.py b/tests/unit/monitors/test_monitor_load.py index 9c91b1fc3fc..f56616c0be3 100644 --- a/tests/unit/monitors/test_monitor_load.py +++ b/tests/unit/monitors/test_monitor_load.py @@ -49,8 +49,8 @@ METRIC_KEYS = 200 # The orchestrator's own work runs between monitor calls; a stall this long delays # dispatch and batch assembly visibly. -MAX_LAG_S = 0.25 -P99_LAG_S = 0.1 +MAX_LAG_S = 0.1 +P99_LAG_S = 0.025 # A log() call may await the monitor's own work, such as the file monitor's write on its # writer thread; the loop keeps serving other tasks meanwhile, which the lag limits cover. MAX_LOG_S = 1.0 diff --git a/tests/unit/orchestrator/test_episode_records.py b/tests/unit/orchestrator/test_episode_records.py index 924731ba296..25e6fbc0304 100644 --- a/tests/unit/orchestrator/test_episode_records.py +++ b/tests/unit/orchestrator/test_episode_records.py @@ -9,13 +9,16 @@ import verifiers.v1 as vf from prime_runs._http import encode_json from prime_runs.projection import build_samples, episodes_to_parquet_bytes +from prime_runs.sinks.base import stamp_run, to_mapping +from prime_runs.sinks.traces import names_run from prime_traces.traces import _encode_record from verifiers.v1.serve.delta import DeltaStreamer, EpisodeAssembly, TraceSummary, dump, pack, unpack from prime_rl.configs.monitors import FileMonitorConfig from prime_rl.monitors.file.monitor import FileMonitor -from prime_rl.monitors.file.traces import get_annotations_dir, get_trace_stream +from prime_rl.monitors.file.traces import get_annotations_dir, get_index_path, get_trace_stream from prime_rl.monitors.file.traces.chunks import chunk_path +from prime_rl.monitors.file.traces.index import index_row from prime_rl.monitors.file.traces.live import get_live_dir, read_live from prime_rl.orchestrator.algo.routing import assign_advantages from prime_rl.orchestrator.annotations import stamp_batch @@ -92,3 +95,29 @@ def test_received_episode_uploads_plain_json(): # The samples table and the training viewer's parquet rows assert encode_json(build_samples([received])) assert episodes_to_parquet_bytes([received], "run", 1) + + +def test_received_episode_encodes_as_its_record_did(tmp_path): + """The file monitor and the Prime SDK encode episodes with `to_json`, one node at a + time: the stream line and its index row match what encoding the whole record gave, + and the upload carries the same record.""" + received, _ = _over_the_wire(episode([trace(), trace()])) + record = received.to_record() + + monitor = FileMonitor(FileMonitorConfig()) + asyncio.run(monitor.init(tmp_path)) + asyncio.run(monitor.log_episodes([received], step=1, kind="train", subset="all")) + + stream = get_trace_stream(tmp_path) + (line,) = chunk_path(stream, 0).read_bytes().splitlines() + assert line == orjson.dumps(record, option=orjson.OPT_SERIALIZE_NUMPY) + (row,) = get_index_path(stream).read_bytes().splitlines() + assert orjson.loads(row) == orjson.loads(orjson.dumps(index_row(1, record, 0, 0), default=str)) + + # With the run recorded on the episode and its traces, as the dispatcher does, the + # SDK uploads `to_json` as is; it carries what the SDK stamped on the record before. + received.record_run(vf.TrainRunInfo(id="prime-run", work=vf.TrainWorkInfo(step=1))) + assert names_run(received, "prime-run", "train") + assert orjson.loads(received.to_json()) == orjson.loads( + _encode_record(stamp_run(to_mapping(received), "prime-run", "train"), 1) + ) diff --git a/uv.lock b/uv.lock index d86fb8f01d5..87381b4f2c1 100644 --- a/uv.lock +++ b/uv.lock @@ -7949,7 +7949,6 @@ dependencies = [ { name = "openai" }, { name = "orjson" }, { name = "prime-pydantic-config", extra = ["toml"] }, - { name = "prime-runs" }, { name = "prime-sandboxes" }, { name = "prime-tunnel" }, { name = "pydantic" }, @@ -7990,8 +7989,7 @@ requires-dist = [ { name = "openenv", marker = "extra == 'openenv'", specifier = ">=0.4.1" }, { name = "orjson", specifier = ">=3.10.0" }, { name = "prime-pydantic-config", extras = ["toml"], specifier = ">=0.4.3" }, - { name = "prime-runs", specifier = ">=0.1.5" }, - { name = "prime-sandboxes", specifier = ">=0.5.3" }, + { name = "prime-sandboxes", specifier = ">=0.5.4" }, { name = "prime-tunnel", specifier = ">=0.1.14" }, { name = "pydantic", specifier = ">=2.12.3" }, { name = "pyzmq", specifier = ">=27.1.0" },