diff --git a/src/agents/tracing/processors.py b/src/agents/tracing/processors.py index b61f3e7976..547aed9fa8 100644 --- a/src/agents/tracing/processors.py +++ b/src/agents/tracing/processors.py @@ -568,6 +568,9 @@ def __init__( self._max_batch_size = max_batch_size self._schedule_delay = schedule_delay self._shutdown_event = threading.Event() + # Set when the worker has something to do before its next scheduled export: the queue + # reached the trigger size, or we are shutting down. + self._wake_event = threading.Event() # The queue size threshold at which we export immediately. self._export_trigger_size = max(1, int(max_queue_size * export_trigger_ratio)) @@ -602,6 +605,8 @@ def on_trace_start(self, trace: Trace) -> None: self._queue.put_nowait(trace) except queue.Full: logger.warning("Queue is full, dropping trace.") + return + self._wake_worker_if_queue_is_full_enough() def on_trace_end(self, trace: Trace) -> None: # We send traces via on_trace_start, so we don't need to do anything here. @@ -619,12 +624,20 @@ def on_span_end(self, span: Span[Any]) -> None: self._queue.put_nowait(span) except queue.Full: logger.warning("Queue is full, dropping span.") + return + self._wake_worker_if_queue_is_full_enough() + + def _wake_worker_if_queue_is_full_enough(self) -> None: + """Wake the worker early when the queue reaches the export trigger size.""" + if self._queue.qsize() >= self._export_trigger_size: + self._wake_event.set() def shutdown(self, timeout: float | None = None): """ Called when the application stops. We signal our thread to stop, then join it. """ self._shutdown_event.set() + self._wake_event.set() if timeout is not None: request_exporter_shutdown = getattr(self._exporter, "_request_shutdown", None) if callable(request_exporter_shutdown): @@ -660,9 +673,14 @@ def _run(self): self._export_batches() # Reset the next scheduled flush time self._next_export_time = time.monotonic() + self._schedule_delay - else: - # Sleep a short interval so we don't busy-wait. - time.sleep(0.2) + continue + + # Block until the next scheduled flush rather than polling. Producers set the + # event once the queue reaches the trigger size and shutdown() sets it to break + # out immediately, so an idle process wakes once per schedule_delay instead of + # five times a second. + self._wake_event.wait(max(0.0, self._next_export_time - time.monotonic())) + self._wake_event.clear() # Final drain after shutdown self._export_batches(deadline=self._shutdown_deadline) diff --git a/tests/test_trace_processor.py b/tests/test_trace_processor.py index 07e975ccb9..3d25e71146 100644 --- a/tests/test_trace_processor.py +++ b/tests/test_trace_processor.py @@ -282,6 +282,52 @@ def export(self, items: list[Trace | Span[Any]]) -> None: assert exporter.call_count >= 3 +def test_batch_trace_processor_queue_trigger_wakes_the_idle_worker() -> None: + """The worker parks until its next scheduled export, so producers have to wake it.""" + exported = threading.Event() + + class SignalingExporter(TracingExporter): + def export(self, items: list[Trace | Span[Any]]) -> None: + exported.set() + + processor = BatchTraceProcessor( + exporter=SignalingExporter(), + max_queue_size=4, + schedule_delay=60.0, + export_trigger_ratio=0.5, + ) + + processor.on_span_end(get_span(processor)) + # One span is below the two-item trigger, so the worker has no reason to wake yet. + assert not exported.wait(timeout=0.2) + + processor.on_span_end(get_span(processor)) + assert exported.wait(timeout=2.0), "reaching the trigger size must wake the worker" + + processor.shutdown(timeout=2.0) + + +def test_batch_trace_processor_shutdown_wakes_the_idle_worker() -> None: + """shutdown() must not wait out the schedule delay before the final drain.""" + exported: list[Trace | Span[Any]] = [] + + class RecordingExporter(TracingExporter): + def export(self, items: list[Trace | Span[Any]]) -> None: + exported.extend(items) + + processor = BatchTraceProcessor(exporter=RecordingExporter(), schedule_delay=60.0) + processor.on_span_end(get_span(processor)) + + start = time.monotonic() + processor.shutdown(timeout=2.0) + elapsed = time.monotonic() - start + + assert elapsed < 1.0 + assert processor._worker_thread is not None + assert not processor._worker_thread.is_alive() + assert len(exported) == 1 + + @pytest.mark.parametrize( ("adjusted_wall_time", "adjusted_monotonic_time", "expected_scheduled_exports"), [ @@ -300,7 +346,6 @@ class ControlledTime: def __init__(self) -> None: self.wall_time = 1000.0 self.monotonic_time = 100.0 - self.sleep_calls = 0 def time(self) -> float: return self.wall_time @@ -309,12 +354,25 @@ def monotonic(self) -> float: return self.monotonic_time def sleep(self, _seconds: float) -> None: - self.sleep_calls += 1 - if self.sleep_calls == 1: - self.wall_time = adjusted_wall_time - self.monotonic_time = adjusted_monotonic_time + raise AssertionError("the worker must wait on its event, not poll with sleep") + + class ControlledWakeEvent: + """Stands in for the worker's wake event and drives the loop from its waits.""" + + def __init__(self) -> None: + self.wait_calls = 0 + + def wait(self, _timeout: float | None = None) -> bool: + self.wait_calls += 1 + if self.wait_calls == 1: + controlled_time.wall_time = adjusted_wall_time + controlled_time.monotonic_time = adjusted_monotonic_time else: processor._shutdown_event.set() + return False + + def clear(self) -> None: + pass controlled_time = ControlledTime() monkeypatch.setattr("agents.tracing.processors.time", controlled_time) @@ -324,6 +382,7 @@ def sleep(self, _seconds: float) -> None: schedule_delay=1.0, export_trigger_ratio=1.0, ) + monkeypatch.setattr(processor, "_wake_event", ControlledWakeEvent()) processor._queue.put_nowait(get_span(processor)) scheduled_export = object() export_deadlines: list[float | None | object] = []