fix(prime-runs): stop blocking and dropping on upload bursts

packages/prime-runs/README.md

@@ -93,8 +93,9 @@ run.finish()
   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
+  sink (a training run pauses it for five minutes instead). Logging never waits
+  for the upload: a queue over its budget (about 4 GiB of encoded records) drops
+  new records 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

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

@@ -1,11 +1,14 @@
-"""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.
+"""Background uploader: one daemon thread draining a budgeted queue into sinks.
+
+Backpressure: submitting never waits. A producer that outruns the uploader
+queues up to a record budget and an estimated byte budget, then drops (counted).
+One write takes about one request's worth of records, so a slow service costs a
+long drain rather than a stalled producer. 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
@@ -15,14 +18,23 @@
 from dataclasses import dataclass, field
 from typing import Any, Callable, Optional, Sequence
 
+from prime_traces.batching import DEFAULT_TARGET_BATCH_BYTES
+
 from . import _fork
+from ._http import encode_json
 from .exceptions import is_record_rejection, is_transient
-from .sinks.base import Sink, SinkWriteError, is_episode
+from .sinks.base import Sink, SinkWriteError, is_episode, to_mapping
 
 logger = logging.getLogger(__name__)
 
-DEFAULT_QUEUE_SIZE = 256
-DEFAULT_PUT_TIMEOUT = 5.0
+DEFAULT_MAX_QUEUED_RECORDS = 65_536
+#: Estimated encoded bytes, a proxy for the memory the queued records hold.
+DEFAULT_MAX_QUEUED_BYTES = 4 * 1024**3
+DEFAULT_MAX_WRITE_BYTES = DEFAULT_TARGET_BATCH_BYTES
+#: Share of the newest sample in the running estimate of a record's size.
+RECORD_BYTES_SMOOTHING = 0.25
+#: Seconds between warnings about dropped records; a burst can drop hundreds.
+DROP_WARNING_INTERVAL = 10.0
 #: Consecutive transient failures before a sink is retired.
 TRANSIENT_FAILURE_LIMIT = 3
 
@@ -47,25 +59,34 @@ def time_left(deadline: Optional[float]) -> Optional[float]:
 
 
 class UploadWorker:
-    """Drains a bounded queue of record batches into a list of sinks."""
+    """Drains a budgeted queue of record batches 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,
+        max_queued_bytes: int = DEFAULT_MAX_QUEUED_BYTES,
+        max_write_bytes: int = DEFAULT_MAX_WRITE_BYTES,
         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.max_queued_bytes = max_queued_bytes
+        self.max_write_bytes = max_write_bytes
         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 taken into a write.
+        self._queued_records = 0
+        self._queued_lock = threading.Lock()
+        #: Running estimate of one record's encoded size, sampled by the uploader
+        #: so that submitting never pays for encoding. ``None`` until the first sample.
+        self._record_bytes: Optional[float] = None
+        self._last_drop_warning: Optional[float] = None
         self._thread: Optional[threading.Thread] = None
         self._stopping = threading.Event()
         self._lock = threading.Lock()
@@ -109,21 +130,28 @@ def _run(self) -> None:
                     item.event.set()
                 else:
                     batch, pending = self._coalesce(item)
+                    with self._queued_lock:
+                        self._queued_records -= len(batch)
                     self._dispatch(batch)
             except Exception as exc:  # noqa: BLE001 - the thread must outlive one bad batch
                 logger.debug("Uploader iteration failed: %s", exc)
-            finally:
-                self._queue.task_done()
 
     def _coalesce(self, first: Sequence[Any]) -> "tuple[list, Any]":
-        """Merge everything already queued behind ``first`` into one batch, so
-        the request rate tracks upload latency rather than rollout throughput.
+        """Merge what is already queued behind ``first`` into one batch of about
+        ``max_write_bytes``, so the request rate tracks upload latency rather
+        than rollout throughput, and one slow or failed write stays small.
         Stops at a control item (a flush barrier must not be reordered past the
         records before it) and at a batch of the other record kind (the traces
-        sink infers the line format from the first record)."""
+        sink infers the line format from the first record). Records past the
+        limit are processed next, ahead of anything queued after them."""
         batch = list(first)
         kind = is_episode(batch[0]) if batch else None
+        if batch:
+            self._sample_record_bytes(batch[0])
+        limit = self._write_limit()
         while True:
+            if len(batch) >= limit:
+                return batch[:limit], (batch[limit:] or _NOTHING)
             try:
                 item = self._queue.get_nowait()
             except queue.Empty:
@@ -133,7 +161,23 @@ def _coalesce(self, first: Sequence[Any]) -> "tuple[list, Any]":
             if item and (kind is None or is_episode(item[0]) != kind):
                 return batch, item
             batch.extend(item)
-            self._queue.task_done()
+
+    def _sample_record_bytes(self, record: Any) -> None:
+        try:
+            size = len(encode_json(to_mapping(record)))
+        except Exception as exc:  # noqa: BLE001 - the sink decides what a bad record means
+            logger.debug("Could not size a record: %s", exc)
+            return
+        if self._record_bytes is None:
+            self._record_bytes = float(size)
+        else:
+            self._record_bytes += RECORD_BYTES_SMOOTHING * (size - self._record_bytes)
+
+    def _write_limit(self) -> int:
+        """Records per write. Without a size estimate, as many as may be queued."""
+        if not self._record_bytes:
+            return max(1, self.max_queued_records)
+        return max(1, int(self.max_write_bytes / self._record_bytes))
 
     def _dispatch(self, records: Sequence[Any]) -> None:
         for sink in self.sinks:
@@ -239,63 +283,75 @@ 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 waiting. ``False`` means it was dropped: the
+        queue is over its budget, or the worker is closed."""
         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)
+            count = len(records)
+            with self._queued_lock:
+                accepted = self._has_room(count)
+                if accepted:
+                    self._queued_records += count
+            if accepted:
+                self._queue.put_nowait(records)
                 return True
-            except queue.Full:
-                count = len(records)
-                self.dropped += count
-                logger.warning(
-                    "Upload queue full after %.1fs; dropped %d item(s) (%d total). "
-                    "The producer is outrunning the uploader.",
-                    self.put_timeout,
-                    count,
-                    self.dropped,
-                )
-                return False
+            self.dropped += count
+            self._warn_dropped(count)
+            return False
+
+    def _has_room(self, count: int) -> bool:
+        """An empty queue takes any batch, so an oversized one is not dropped outright."""
+        queued = self._queued_records
+        if not queued:
+            return True
+        total = queued + count
+        if total > self.max_queued_records:
+            return False
+        return self._record_bytes is None or total * self._record_bytes <= self.max_queued_bytes
+
+    def _warn_dropped(self, count: int) -> None:
+        now = time.monotonic()
+        last = self._last_drop_warning
+        if last is not None and now - last < DROP_WARNING_INTERVAL:
+            return
+        self._last_drop_warning = now
+        logger.warning(
+            "Upload queue over budget (%d record(s) queued); dropped %d item(s) (%d total). "
+            "The producer is outrunning the uploader.",
+            self._queued_records,
+            count,
+            self.dropped,
+        )
 
     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_nowait(barrier)
+        return barrier.event.wait(timeout if timeout is None else max(0.0, 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_nowait(None)
+            thread.join(timeout if timeout is None else max(0.0, timeout))
             if thread.is_alive():
                 # Closing the sinks now would pull a client out from under a request.
                 logger.warning(
-                    "Uploader still running after %ss; leaving it and its sinks open. "
-                    "Records still in flight may not finish before the process exits.",
+                    "Uploader still running after %ss with %d record(s) queued; leaving it "
+                    "and its sinks open. They may not finish before the process exits.",
                     timeout,
+                    self._queued_records,
                 )
                 return
         self._thread = None
@@ -306,7 +362,9 @@ 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._queued_lock = threading.Lock()
         self._thread = None
         self._stopping = threading.Event()
         self._lock = threading.Lock()