feat: encode episodes for the file and prime monitors one node at a time

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()
 

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:

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)."""

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")

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

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)
+    )