diff --git a/packages/prime-runs/README.md b/packages/prime-runs/README.md index 8f953da2f..56ec0a75b 100644 --- a/packages/prime-runs/README.md +++ b/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 diff --git a/packages/prime-runs/src/prime_runs/worker.py b/packages/prime-runs/src/prime_runs/worker.py index 6e33c96d3..17bfe83f2 100644 --- a/packages/prime-runs/src/prime_runs/worker.py +++ b/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 record in the running estimate of a record's size. +RECORD_BYTES_SMOOTHING = 0.1 +#: 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, measured by the uploader + #: so that submitting never pays for encoding. ``None`` until it has one. + 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,32 @@ 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 at most + ``max_write_bytes`` (or one record, if it is bigger), 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).""" - batch = list(first) - kind = is_episode(batch[0]) if batch else None + sink infers the line format from the first record). Records past the + limit are processed next, ahead of anything queued after them.""" + batch: list = [] + batch_bytes = 0 + item = first + kind = is_episode(item[0]) if item else None while True: + for position, record in enumerate(item): + size = self._measure(record) + if batch and batch_bytes + size > self.max_write_bytes: + return batch, list(item[position:]) + batch.append(record) + batch_bytes += size try: item = self._queue.get_nowait() except queue.Empty: @@ -132,8 +164,20 @@ def _coalesce(self, first: Sequence[Any]) -> "tuple[list, Any]": return batch, item if item and (kind is None or is_episode(item[0]) != kind): return batch, item - batch.extend(item) - self._queue.task_done() + + def _measure(self, record: Any) -> int: + """The record's encoded size, folded into the queue's running estimate. A + record that cannot be encoded counts as nothing: the sink decides what it means.""" + try: + size = len(encode_json(to_mapping(record))) + except Exception as exc: # noqa: BLE001 + logger.debug("Could not size a record: %s", exc) + return 0 + if self._record_bytes is None: + self._record_bytes = float(size) + else: + self._record_bytes += RECORD_BYTES_SMOOTHING * (size - self._record_bytes) + return size def _dispatch(self, records: Sequence[Any]) -> None: for sink in self.sinks: @@ -239,8 +283,8 @@ 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. @@ -248,35 +292,50 @@ def submit(self, records: Sequence[Any]) -> bool: 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.""" @@ -284,18 +343,15 @@ def close(self, timeout: Optional[float] = 30.0) -> None: 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() diff --git a/packages/prime-runs/tests/test_worker.py b/packages/prime-runs/tests/test_worker.py index f429b809d..f21da9dff 100644 --- a/packages/prime-runs/tests/test_worker.py +++ b/packages/prime-runs/tests/test_worker.py @@ -1,7 +1,9 @@ """The background uploader: backpressure, containment, fork safety.""" +import logging import queue import threading +import time import pytest from conftest import FakeSink @@ -82,15 +84,17 @@ def write(self, records) -> None: def test_a_full_queue_drops_rather_than_blocking_the_producer(): sink = BlockingSink() - worker = UploadWorker([sink], max_queue_size=1, put_timeout=0.05) + worker = UploadWorker([sink], max_queued_records=1) # First item is picked up and wedges the uploader inside sink.write(). 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}]) + started = time.monotonic() accepted = worker.submit([{"id": 2}, {"id": 3}]) + assert time.monotonic() - started < 0.5 assert accepted is False assert worker.dropped == 2 @@ -98,9 +102,151 @@ def test_a_full_queue_drops_rather_than_blocking_the_producer(): worker.close() -def test_flush_uses_the_drain_budget_to_get_behind_a_full_queue(): +def _record(index: int, size: int = 100) -> dict: + return {"id": index, "pad": "x" * size} + + +def test_the_byte_budget_drops_once_the_estimated_queue_is_over_it(): + sink = BlockingSink() + # Each record encodes to a little over 100 bytes: three fit the budget, four do not. + worker = UploadWorker([sink], max_queued_bytes=400) + assert worker.submit([_record(0)]) # sampled, then wedges the uploader + assert sink.entered.wait(5.0) + + accepted = [worker.submit([_record(index)]) for index in range(1, 6)] + + assert accepted == [True, True, True, False, False] + assert worker.dropped == 2 + + sink.released.set() + worker.close() + + +def test_one_write_takes_about_one_request_of_what_is_queued(): + sink = BlockingSink() + worker = UploadWorker([sink], max_write_bytes=250) # two records a write + assert worker.submit([_record(0)]) + assert sink.entered.wait(5.0) + for index in range(1, 6): + assert worker.submit([_record(index)]) + + sink.released.set() + drain(worker) + + assert [[record["id"] for record in batch] for batch in sink.batches] == [ + [0], + [1, 2], + [3, 4], + [5], + ] + assert worker._queued_records == 0 + worker.close() + + +def test_a_small_record_first_does_not_let_large_ones_through_uncapped(): + sink = BlockingSink() + worker = UploadWorker([sink], max_write_bytes=2_500) + assert worker.submit([_record(0)]) + assert sink.entered.wait(5.0) + assert worker.submit([_record(1)]) # small, and first in the next window + for index in range(2, 7): + assert worker.submit([_record(index, size=1_000)]) + + sink.released.set() + drain(worker) + + assert [[record["id"] for record in batch] for batch in sink.batches] == [ + [0], + [1, 2, 3], + [4, 5], + [6], + ] + assert worker._queued_records == 0 + worker.close() + + +def test_a_record_over_the_write_limit_goes_out_alone(): + sink = BlockingSink() + worker = UploadWorker([sink], max_write_bytes=250) + assert worker.submit([_record(0)]) + assert sink.entered.wait(5.0) + assert worker.submit([_record(1, size=1_000)]) + assert worker.submit([_record(2)]) + + sink.released.set() + drain(worker) + + assert [[record["id"] for record in batch] for batch in sink.batches] == [[0], [1], [2]] + worker.close() + + +def test_an_oversized_batch_goes_out_in_order_across_writes_ahead_of_later_records(): sink = BlockingSink() - worker = UploadWorker([sink], max_queue_size=1, put_timeout=0.01) + worker = UploadWorker([sink], max_write_bytes=250) + assert worker.submit([_record(0)]) + assert sink.entered.wait(5.0) + assert worker.submit([_record(index) for index in range(1, 6)]) + assert worker.submit([_record(6)]) + + sink.released.set() + drain(worker) + + assert [[record["id"] for record in batch] for batch in sink.batches] == [ + [0], + [1, 2], + [3, 4], + [5, 6], + ] + assert worker._queued_records == 0 + worker.close() + + +def test_a_batch_over_the_budgets_is_still_taken_by_an_empty_queue(): + sink = FakeSink() + worker = UploadWorker([sink], max_queued_records=1, max_queued_bytes=1) + + assert worker.submit([_record(0), _record(1), _record(2)]) + drain(worker) + + assert worker.dropped == 0 + assert sum(len(batch) for batch in sink.batches) == 3 + worker.close() + + +def test_a_record_that_cannot_be_sized_still_reaches_the_sink(): + sink = FakeSink() + worker = UploadWorker([sink]) + unsizable = object() + + assert worker.submit([unsizable]) + drain(worker) + + assert sink.batches == [[unsizable]] + assert worker._record_bytes is None + worker.close() + + +def test_a_burst_of_drops_warns_once(caplog): + sink = BlockingSink() + worker = UploadWorker([sink], max_queued_records=1) + assert worker.submit([{"id": 0}]) + assert sink.entered.wait(5.0) + assert worker.submit([{"id": 1}]) + + with caplog.at_level(logging.WARNING, logger="prime_runs.worker"): + for index in range(2, 12): + assert worker.submit([{"id": index}]) is False + + assert worker.dropped == 10 + assert len([r for r in caplog.records if "over budget" in r.getMessage()]) == 1 + + sink.released.set() + worker.close() + + +def test_flush_waits_for_the_records_queued_ahead_of_it(): + sink = BlockingSink() + worker = UploadWorker([sink], max_queued_records=1) assert worker.submit([{"id": 0}]) assert sink.entered.wait(1.0) assert worker.submit([{"id": 1}]) @@ -117,9 +263,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_drains_the_records_queued_ahead_of_the_stop(): sink = BlockingSink() - worker = UploadWorker([sink], max_queue_size=1, put_timeout=0.01) + worker = UploadWorker([sink], max_queued_records=1) assert worker.submit([{"id": 0}]) assert sink.entered.wait(1.0) assert worker.submit([{"id": 1}]) @@ -160,14 +306,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