feat: non-blocking, concurrent run uploads in prime-runs and prime-traces

packages/prime-runs/README.md

@@ -87,15 +87,16 @@ run.finish()
 ## How it behaves
 
 - **Streams.** Records go out on a background thread as they are logged;
-  whatever queues up during one request goes out as the next.
+  whatever queues up during one upload goes out as the next, with up to eight
+  requests in flight. Logging never blocks the caller.
 - **Contains its errors.** With the default `on_error="warn"` nothing the
   platform raises escapes into your loop; `on_error="raise"` surfaces the first
   failure from `flush()` or `finish()`, for tests and CI. Platform errors are
   the `prime_traces` exception family.
 - **Degrades.** A transient failure costs its batch, three in a row retire the
   sink (a training run pauses it for five minutes instead), and a full queue
-  drops records rather than stalling the run. Losses are counted in
-  `run.dropped_records` and `run.failed_records`.
+  (10,000 records waiting) drops a new batch rather than stalling the run.
+  Losses are counted in `run.dropped_records` and `run.failed_records`.
 - **Drains on exit.** `finish()` gives queued uploads up to `finish_timeout`
   (300 s) before closing the run out; an abort path can pass
   `finish(timeout=...)`.

packages/prime-runs/src/prime_runs/sinks/traces.py

@@ -93,18 +93,22 @@ def write(self, records: Sequence[Any]) -> None:
         self._ensure_client()
         # The line format is inferred from the first record; batches are homogeneous.
         line_format = LineFormat.EPISODE if is_episode(records[0]) else LineFormat.TRACE
-        payload = [self._prepare(record) for record in records]
-        offset = 0
+        # Prepared one at a time as the client encodes them, so a backlog is never
+        # held as wire mappings all at once.
+        has_traces: list = []
+
+        def prepared() -> Any:
+            for record in records:
+                mapping = self._prepare(record)
+                if line_format is LineFormat.EPISODE:
+                    has_traces.append(bool(mapping.get("traces")))
+                yield mapping
 
         def committed(batch: Any, receipt: Any) -> None:
-            nonlocal offset
             if line_format is LineFormat.EPISODE:
                 # Match hosted eval confirmation: empty episodes produce no samples.
-                self.episodes_written += sum(
-                    bool(record.get("traces"))
-                    for record in payload[offset : offset + batch.num_lines]
-                )
-            offset += batch.num_lines
+                start = batch.first_line_number - 1
+                self.episodes_written += sum(has_traces[start : start + batch.num_lines])
             self.receipts_received += 1
             if self._receipt_history_size:
                 self.receipts.append(receipt)
@@ -113,7 +117,7 @@ def committed(batch: Any, receipt: Any) -> None:
         try:
             list(
                 self._client.upload_records(
-                    payload,
+                    prepared(),
                     line_format=line_format,
                     context=dict(self._context) or None,
                     compress=self._compress,

packages/prime-runs/src/prime_runs/worker.py

@@ -1,11 +1,12 @@
 """Background uploader: one daemon thread draining a bounded queue into sinks.
 
-Backpressure: a producer that outruns the uploader blocks briefly, then drops
-(counted). Fork safety: a child starts over with an empty queue; the queued
-records belong to the parent. Containment: record-specific failures drop only
-their batch, transient failures get a few consecutive strikes before the sink
-is retired (for good, or until ``retire_cooldown`` has passed), and any other
-failure retires the sink at once. The thread never dies on one bad batch.
+Backpressure: submitting never blocks the producer. The queue holds up to
+``max_queued_records`` records; a batch that does not fit is dropped (counted).
+Fork safety: a child starts over with an empty queue; the queued records belong
+to the parent. Containment: record-specific failures drop only their batch,
+transient failures get a few consecutive strikes before the sink is retired
+(for good, or until ``retire_cooldown`` has passed), and any other failure
+retires the sink at once. The thread never dies on one bad batch.
 """
 
 import logging
@@ -21,8 +22,9 @@
 
 logger = logging.getLogger(__name__)
 
-DEFAULT_QUEUE_SIZE = 256
-DEFAULT_PUT_TIMEOUT = 5.0
+#: Records queued or being written before new batches are dropped. A burst of a
+#: few thousand episodes (one large training step) fits.
+DEFAULT_MAX_QUEUED_RECORDS = 10_000
 #: Consecutive transient failures before a sink is retired.
 TRANSIENT_FAILURE_LIMIT = 3
 
@@ -47,25 +49,25 @@ def time_left(deadline: Optional[float]) -> Optional[float]:
 
 
 class UploadWorker:
-    """Drains a bounded queue of record batches into a list of sinks."""
+    """Drains a queue of record batches, bounded by record count, into a list of sinks."""
 
     def __init__(
         self,
         sinks: Sequence[Sink],
         *,
-        max_queue_size: int = DEFAULT_QUEUE_SIZE,
-        put_timeout: float = DEFAULT_PUT_TIMEOUT,
+        max_queued_records: int = DEFAULT_MAX_QUEUED_RECORDS,
         on_error: Optional[Callable[[str, Exception], None]] = None,
         retire_cooldown: Optional[float] = None,
     ) -> None:
         self.sinks = sinks
-        self.max_queue_size = max_queue_size
-        self.put_timeout = put_timeout
+        self.max_queued_records = max_queued_records
         self._on_error = on_error
         #: Seconds after which a sink retired on transient failures is tried
         #: again; ``None`` retires it for the rest of the run.
         self._retire_cooldown = retire_cooldown
-        self._queue: "queue.Queue[Any]" = queue.Queue(maxsize=max_queue_size)
+        self._queue: "queue.Queue[Any]" = queue.Queue()
+        #: Records submitted and not yet written; guarded by ``_lock``.
+        self._queued_records = 0
         self._thread: Optional[threading.Thread] = None
         self._stopping = threading.Event()
         self._lock = threading.Lock()
@@ -109,7 +111,11 @@ def _run(self) -> None:
                     item.event.set()
                 else:
                     batch, pending = self._coalesce(item)
-                    self._dispatch(batch)
+                    try:
+                        self._dispatch(batch)
+                    finally:
+                        with self._lock:
+                            self._queued_records -= len(batch)
             except Exception as exc:  # noqa: BLE001 - the thread must outlive one bad batch
                 logger.debug("Uploader iteration failed: %s", exc)
             finally:
@@ -239,57 +245,47 @@ def _notify(self, name: str, exc: Exception) -> None:
             logger.debug("Error handler raised while reporting a sink failure", exc_info=True)
 
     def submit(self, records: Sequence[Any]) -> bool:
-        """Queue a batch. ``False`` means it was dropped: the queue stayed full
-        for the whole put timeout."""
+        """Queue a batch without blocking. ``False`` means it was dropped: it
+        does not fit in the queue. A batch always fits an empty queue."""
         if not self.sinks:  # a disabled run: nowhere for the records to go
             return True
         # Serialize with close(): an accepted batch is queued before the stop sentinel.
         with self._lock:
             if self._stopping.is_set():
                 return False
             self._start_locked()
-            try:
-                self._queue.put(records, timeout=self.put_timeout)
-                return True
-            except queue.Full:
-                count = len(records)
+            count = len(records)
+            if self._queued_records and self._queued_records + count > self.max_queued_records:
                 self.dropped += count
                 logger.warning(
-                    "Upload queue full after %.1fs; dropped %d item(s) (%d total). "
+                    "Upload queue full (%d records); dropped %d item(s) (%d total). "
                     "The producer is outrunning the uploader.",
-                    self.put_timeout,
+                    self._queued_records,
                     count,
                     self.dropped,
                 )
                 return False
+            self._queued_records += count
+            self._queue.put(records)
+            return True
 
     def flush(self, timeout: Optional[float] = None) -> bool:
         """Block until everything queued so far has been written."""
         if self._thread is None or not self._thread.is_alive():
             self._flush_sinks()
             return True
-        deadline = deadline_after(timeout)
         barrier = _Flush()
-        try:
-            # The barrier gets the caller's drain budget, not the producer put timeout.
-            self._queue.put(barrier, timeout=time_left(deadline))
-        except queue.Full:
-            logger.warning("Could not enqueue a flush barrier before the drain timeout")
-            return False
-        return barrier.event.wait(time_left(deadline))
+        self._queue.put(barrier)
+        return barrier.event.wait(timeout)
 
     def close(self, timeout: Optional[float] = 30.0) -> None:
         """Drain, stop the thread, and close every sink."""
         with self._lock:
             self._stopping.set()
             thread = self._thread
         if thread is not None and thread.is_alive():
-            deadline = deadline_after(timeout)
-            try:
-                self._queue.put(None, timeout=time_left(deadline))
-            except queue.Full:
-                logger.warning("Upload queue remained saturated through the close timeout")
-            thread.join(time_left(deadline))
+            self._queue.put(None)
+            thread.join(timeout)
             if thread.is_alive():
                 # Closing the sinks now would pull a client out from under a request.
                 logger.warning(
@@ -306,7 +302,8 @@ def close(self, timeout: Optional[float] = 30.0) -> None:
                 logger.debug("Error closing sink %s: %s", sink.name, exc)
 
     def reset_after_fork(self) -> None:
-        self._queue = queue.Queue(maxsize=self.max_queue_size)
+        self._queue = queue.Queue()
+        self._queued_records = 0
         self._thread = None
         self._stopping = threading.Event()
         self._lock = threading.Lock()

packages/prime-runs/tests/test_traces_sink.py

@@ -17,10 +17,11 @@ def __init__(self, raises: Exception = None) -> None:
         self.closed = False
 
     def upload_records(self, records, **kwargs):
-        self.calls.append((list(records), kwargs))
+        records = list(records)
+        self.calls.append((records, kwargs))
         if self.raises is not None:
             raise self.raises
-        kwargs["on_batch"](SimpleNamespace(num_lines=len(records)), "receipt")
+        kwargs["on_batch"](SimpleNamespace(num_lines=len(records), first_line_number=1), "receipt")
         return ["receipt"]
 
     def close(self) -> None:

packages/prime-runs/tests/test_worker.py

@@ -2,6 +2,7 @@
 
 import queue
 import threading
+import time
 
 import pytest
 from conftest import FakeSink
@@ -80,27 +81,44 @@ def write(self, records) -> None:
     worker.close()
 
 
-def test_a_full_queue_drops_rather_than_blocking_the_producer():
+def test_a_full_queue_drops_at_once_rather_than_blocking_the_producer():
     sink = BlockingSink()
-    worker = UploadWorker([sink], max_queue_size=1, put_timeout=0.05)
+    worker = UploadWorker([sink], max_queued_records=2)
 
-    # First item is picked up and wedges the uploader inside sink.write().
+    # The first record is picked up and wedges the uploader inside sink.write();
+    # it counts against the bound until it is written.
     assert worker.submit([{"id": 0}])
     assert sink.entered.wait(5.0), "the uploader never reached the sink"
-    # Second fills the one-slot queue; third has nowhere to go.
-    worker.submit([{"id": 1}])
+    assert worker.submit([{"id": 1}])
+    start = time.monotonic()
     accepted = worker.submit([{"id": 2}, {"id": 3}])
 
+    assert time.monotonic() - start < 0.05
     assert accepted is False
     assert worker.dropped == 2
 
     sink.released.set()
+    drain(worker)
+    # Written records free their room.
+    assert worker.submit([{"id": 4}, {"id": 5}])
     worker.close()
+    assert [record["id"] for batch in sink.batches for record in batch] == [0, 1, 4, 5]
+
 
+def test_a_batch_larger_than_the_bound_fits_an_empty_queue():
+    sink = FakeSink()
+    worker = UploadWorker([sink], max_queued_records=2)
+
+    assert worker.submit([{"id": i} for i in range(5)])
+    drain(worker)
 
-def test_flush_uses_the_drain_budget_to_get_behind_a_full_queue():
+    assert len(sink.batches[0]) == 5
+    worker.close()
+
+
+def test_flush_waits_for_the_records_queued_before_it():
     sink = BlockingSink()
-    worker = UploadWorker([sink], max_queue_size=1, put_timeout=0.01)
+    worker = UploadWorker([sink])
     assert worker.submit([{"id": 0}])
     assert sink.entered.wait(1.0)
     assert worker.submit([{"id": 1}])
@@ -117,9 +135,9 @@ def test_flush_uses_the_drain_budget_to_get_behind_a_full_queue():
     assert [batch[0]["id"] for batch in sink.batches] == [0, 1]
 
 
-def test_close_uses_its_budget_to_queue_the_stop_behind_pending_records():
+def test_close_queues_the_stop_behind_pending_records():
     sink = BlockingSink()
-    worker = UploadWorker([sink], max_queue_size=1, put_timeout=0.01)
+    worker = UploadWorker([sink])
     assert worker.submit([{"id": 0}])
     assert sink.entered.wait(1.0)
     assert worker.submit([{"id": 1}])
@@ -160,14 +178,16 @@ def test_submitting_after_close_is_refused():
 
 def test_a_forked_child_starts_over_instead_of_re_uploading_the_parents_queue():
     sink = FakeSink()
-    worker = UploadWorker([sink], max_queue_size=4)
+    worker = UploadWorker([sink])
     worker._queue.put([{"id": "parents"}])
+    worker._queued_records = 1
     old_queue = worker._queue
 
     worker.reset_after_fork()
 
     assert worker._queue is not old_queue
     assert worker._queue.empty()
+    assert worker._queued_records == 0
     assert worker._thread is None
     assert isinstance(worker._queue, queue.Queue)
     assert isinstance(worker._lock, type(threading.Lock())) or worker._lock is not None

packages/prime-traces/src/prime_traces/core/async_client.py

@@ -11,6 +11,7 @@
 from .client import (
     AMBIGUOUS_TRANSPORT_ERRORS,
     DEFAULT_TIMEOUT,
+    GZIP_LEVEL,
     IDEMPOTENT_RETRY_ATTEMPTS,
     BaseTracesAPIClient,
     is_ambiguous_delete_failure,
@@ -95,9 +96,9 @@ async def upload_batch(
         self._check_auth()
 
         # mtime=0 keeps the compressed bytes reproducible across retries.
-        compressed = (
-            await asyncio.to_thread(gzip_module.compress, body, mtime=0) if compress else None
-        )
+        compressed = None
+        if compress:
+            compressed = await asyncio.to_thread(gzip_module.compress, body, GZIP_LEVEL, mtime=0)
         files = self._multipart_files(
             body, schema_version=schema_version, context=context, compressed=compressed
         )

packages/prime-traces/src/prime_traces/core/client.py

@@ -48,6 +48,11 @@
 #: budget — see ``TracesClient._send_with_retry``.
 IDEMPOTENT_RETRY_ATTEMPTS = 3
 
+#: Gzip level for upload bodies. Level 1 compresses a multi-MB episode about 9x
+#: faster than the default 9 for ~13% more bytes; the upload thread is otherwise
+#: CPU-bound on compression and falls behind a fast producer.
+GZIP_LEVEL = 1
+
 #: Retryable statuses that leave it unknown whether the request reached the
 #: service. 502/504 come from a gateway, which may already have forwarded the
 #: request upstream and then lost or timed out the response.
@@ -379,7 +384,9 @@ def upload_batch(
         self._check_auth()
 
         # mtime=0 keeps the compressed bytes reproducible across retries.
-        compressed = gzip_module.compress(body, mtime=0) if compress else None
+        compressed = None
+        if compress:
+            compressed = gzip_module.compress(body, compresslevel=GZIP_LEVEL, mtime=0)
         files = self._multipart_files(
             body, schema_version=schema_version, context=context, compressed=compressed
         )