Fix: drain the host-log writer before scene tests read captured spans

tests/st/a2a3/host_build_graph/concurrent_prepare_stress/test_concurrent_prepare_stress.py

@@ -202,7 +202,7 @@ def _drive_pipeline(self, chip_worker, callable_id, orch_sig, config, iters, inf
         while inflight:
             self._retire(inflight.pop(0))
 
-    def test_concurrent_prepare_overlap(self, st_platform, st_worker, capfd):
+    def test_concurrent_prepare_overlap(self, st_platform, st_worker, capfd, drain_host_log):
         """Golden-check a 2-deep overlapping pipeline over both arena banks.
 
         Each submission prepares (fully binds) its run against one bank while the
@@ -219,9 +219,9 @@ def test_concurrent_prepare_overlap(self, st_platform, st_worker, capfd):
         callable_id = _ARM_CALLABLE_ID
 
         with self._registered_callable(chip_worker, callable_id, callable_obj):
-            capfd.readouterr()  # only this arm's spans reach the verdict
+            drain_host_log(capfd)  # only this arm's spans reach the verdict
             self._drive_pipeline(chip_worker, callable_id, orch_sig, config, _ITERS, inflight_limit=2)
-            captured = capfd.readouterr().err
+            captured = drain_host_log(capfd)
 
         # Each pair must prove prepare(N+1) ran concurrently with device(N) — the
         # property this pipeline exists for. Not that prepare *finished* inside
@@ -230,7 +230,7 @@ def test_concurrent_prepare_overlap(self, st_platform, st_worker, capfd):
         checks = assert_native_overlap(parse_spans(captured.splitlines()))
         assert len(checks) == _ITERS - 1
 
-    def test_serial_submission_is_rejected_as_not_overlapping(self, st_platform, st_worker, capfd):
+    def test_serial_submission_is_rejected_as_not_overlapping(self, st_platform, st_worker, capfd, drain_host_log):
         """One run at a time must fail the same assertion the overlapping arm passes.
 
         The single-variable control for the positive arm: same lane, same
@@ -254,14 +254,14 @@ def test_serial_submission_is_rejected_as_not_overlapping(self, st_platform, st_
         callable_id = _ARM_CALLABLE_ID
 
         with self._registered_callable(chip_worker, callable_id, callable_obj):
-            capfd.readouterr()
+            drain_host_log(capfd)
             self._drive_pipeline(chip_worker, callable_id, orch_sig, config, _CONTROL_ITERS, inflight_limit=1)
-            captured = capfd.readouterr().err
+            captured = drain_host_log(capfd)
 
         with pytest.raises(NativeOverlapError, match="did not overlap"):
             assert_native_overlap(parse_spans(captured.splitlines()))
 
-    def test_diagnostics_config_serializes_the_native_lane(self, st_platform, st_worker, capfd):
+    def test_diagnostics_config_serializes_the_native_lane(self, st_platform, st_worker, capfd, drain_host_log):
         """A diagnostic flag turns staging off, and the log must then be rejected.
 
         ``allow_prepared_successor`` folds in ``CallConfig::diagnostics_any()`` —
@@ -303,9 +303,9 @@ def test_diagnostics_config_serializes_the_native_lane(self, st_platform, st_wor
             # is required by CallConfig::validate() whenever one of them is.
             config = self._build_config({}, enable_scope_stats=True, output_prefix=output_dir)
             with self._registered_callable(chip_worker, callable_id, callable_obj):
-                capfd.readouterr()
+                drain_host_log(capfd)
                 self._drive_pipeline(chip_worker, callable_id, orch_sig, config, _CONTROL_ITERS, inflight_limit=2)
-                captured = capfd.readouterr().err
+                captured = drain_host_log(capfd)
 
         with pytest.raises(NativeOverlapError, match="did not overlap"):
             assert_native_overlap(parse_spans(captured.splitlines()))

tests/st/a2a3/host_build_graph/native_run_lifecycle/test_native_run_lifecycle.py

@@ -68,10 +68,10 @@ def generate_args(self, params):
     def compute_golden(self, args, params):
         args.out[:] = args.a + args.b + _CHAIN_LENGTH
 
-    def test_run(self, st_platform, st_worker, request, capfd):
+    def test_run(self, st_platform, st_worker, request, capfd, drain_host_log):
         super().test_run(st_platform, st_worker, request)
 
-        spans = list(parse_spans(capfd.readouterr().err.splitlines()))
+        spans = list(parse_spans(drain_host_log(capfd).splitlines()))
         invocations = [inv for inv in group_invocations(spans) if "chip.run" in inv.by_name()]
         # Two of these are the abandoned diagnostic prepares below, which record a
         # chip.run invocation without ever reaching chip.run.runner_run.

tests/st/conftest.py

@@ -10,3 +10,52 @@
 
 sys.path is handled by pyproject.toml [tool.pytest.ini_options] pythonpath.
 """
+
+import time
+
+import pytest
+
+
+@pytest.fixture
+def drain_host_log():
+    """Read captured output with the host-log writer drained first.
+
+    A `[STRACE]` record reaches captured stderr through the process writer
+    thread, so a bare `capfd.readouterr()` races it: the reader can return before
+    the last records are written, and the failure then looks like a missing span
+    rather than a timing problem. Every test that counts or matches spans in
+    captured output needs this.
+
+    The wait is bounded but it is not the verdict. Producers are quiescent by the
+    time a test reads — the run has completed — so `pending_record_count`
+    reaching zero is a real drain rather than a deadline standing in for
+    correctness, and the caller's own assertion stays the thing that decides.
+    Exhausting the bound means the writer is genuinely stuck, and the message
+    reports the drop counter so a queue loss is not mistaken for a slow drain.
+    """
+    from simpler.task_interface import (  # noqa: PLC0415
+        _flush_host_log,
+        _host_log_dropped_records,
+        _host_log_pending_records,
+    )
+
+    def _drain(capfd, timeout_s: float = 5.0) -> str:
+        dropped_before = _host_log_dropped_records()
+        chunks: list[str] = []
+        deadline = time.monotonic() + timeout_s
+        while True:
+            flushed = _flush_host_log(100)
+            captured = capfd.readouterr()
+            chunks.extend((captured.err, captured.out))
+            if flushed and _host_log_pending_records() == 0:
+                return "".join(chunks)
+            if time.monotonic() >= deadline:
+                pending = _host_log_pending_records()
+                dropped = _host_log_dropped_records() - dropped_before
+                raise AssertionError(
+                    f"host-log writer did not drain within {timeout_s:.1f}s: "
+                    f"pending={pending}, dropped_delta={dropped}, last_flush={flushed}"
+                )
+            time.sleep(0.01)
+
+    return _drain

tests/st/task_timing/task_timing_slots/test_task_timing_e2e.py

@@ -155,10 +155,10 @@ def _drive(
 @pytest.mark.runtime("tensormap_and_ringbuffer")
 @pytest.mark.device_count(1)
 @pytest.mark.manual(["a2a3sim", "a5sim"])
-def test_distinct_slots_emit_markers(st_platform, st_device_ids, capfd):
+def test_distinct_slots_emit_markers(st_platform, st_device_ids, capfd, drain_host_log):
     # Two-task chain: t0 -> slot 0, t1 -> slot 1. out = a + 2b.
     _drive(st_platform, int(st_device_ids[0]), "task_timing_orchestration", 2)
-    err = capfd.readouterr().err
+    err = drain_host_log(capfd)
 
     slot0 = _slot_spans(err, 0)
     slot1 = _slot_spans(err, 1)
@@ -176,15 +176,15 @@ def test_distinct_slots_emit_markers(st_platform, st_device_ids, capfd):
 @pytest.mark.runtime("tensormap_and_ringbuffer")
 @pytest.mark.device_count(1)
 @pytest.mark.manual(["a2a3sim", "a5sim"])
-def test_duplicate_slot_merges_window(st_platform, st_device_ids, capfd):
+def test_duplicate_slot_merges_window(st_platform, st_device_ids, capfd, drain_host_log):
     dev = int(st_device_ids[0])
 
     # A 3-task chain (t0 -> t1 -> t2) all tagged slot 0. out = a + 3b. The three
     # tagged tasks must fold their dispatch/finish into a SINGLE slot-0 window
     # (min dispatch .. max finish), not three separate markers, and no other slot
     # is touched.
     _drive(st_platform, dev, "task_timing_dup_orchestration", 3)
-    err = capfd.readouterr().err
+    err = drain_host_log(capfd)
 
     slot0 = _slot_spans(err, 0)
     assert len(slot0) == 1, f"expected exactly ONE merged task_slot_0 marker, got {len(slot0)}: {slot0}"
@@ -209,10 +209,10 @@ def test_duplicate_slot_merges_window(st_platform, st_device_ids, capfd):
 @pytest.mark.runtime("host_build_graph")
 @pytest.mark.device_count(1)
 @pytest.mark.manual(["a2a3sim", "a5sim"])
-def test_hbg_distinct_slots_emit_markers(st_platform, st_device_ids, capfd):
+def test_hbg_distinct_slots_emit_markers(st_platform, st_device_ids, capfd, drain_host_log):
     # Same two-task chain as test_distinct_slots_emit_markers, on the hbg path.
     _drive(st_platform, int(st_device_ids[0]), "task_timing_orchestration", 2, runtime="host_build_graph")
-    err = capfd.readouterr().err
+    err = drain_host_log(capfd)
 
     slot0 = _slot_spans(err, 0)
     slot1 = _slot_spans(err, 1)
@@ -230,11 +230,11 @@ def test_hbg_distinct_slots_emit_markers(st_platform, st_device_ids, capfd):
 @pytest.mark.runtime("host_build_graph")
 @pytest.mark.device_count(1)
 @pytest.mark.manual(["a2a3sim", "a5sim"])
-def test_hbg_duplicate_slot_merges_window(st_platform, st_device_ids, capfd):
+def test_hbg_duplicate_slot_merges_window(st_platform, st_device_ids, capfd, drain_host_log):
     # Same 3-task same-slot merge as test_duplicate_slot_merges_window, on the hbg
     # path: min(dispatch)/max(finish) must fold into a single slot-0 window.
     _drive(st_platform, int(st_device_ids[0]), "task_timing_dup_orchestration", 3, runtime="host_build_graph")
-    err = capfd.readouterr().err
+    err = drain_host_log(capfd)
 
     slot0 = _slot_spans(err, 0)
     assert len(slot0) == 1, f"expected exactly ONE merged task_slot_0 marker, got {len(slot0)}: {slot0}"
@@ -293,7 +293,7 @@ def incore(rel, core_type):
 @pytest.mark.runtime("tensormap_and_ringbuffer")
 @pytest.mark.device_count(1)
 @pytest.mark.manual(["a2a3sim"])
-def test_mix_task_aggregates_across_subtasks(st_platform, st_device_ids, capfd):
+def test_mix_task_aggregates_across_subtasks(st_platform, st_device_ids, capfd, drain_host_log):
     # One MIX task (AIC matmul + AIV0 add + AIV1 mul) tagged slot 0. All three
     # subtasks fold their dispatch/finish into slot 0 -> one complete window.
     import torch  # noqa: PLC0415
@@ -351,7 +351,7 @@ def test_mix_task_aggregates_across_subtasks(st_platform, st_device_ids, capfd):
     finally:
         worker.close()
 
-    err = capfd.readouterr().err
+    err = drain_host_log(capfd)
     slot0 = _slot_spans(err, 0)
     assert len(slot0) == 1, f"expected exactly one task_slot_0 marker for the MIX task, got {slot0}"
     assert slot0[0][1] > 0, f"MIX task_slot_0 must be a complete window across its subtasks, got {slot0}"
@@ -361,11 +361,11 @@ def test_mix_task_aggregates_across_subtasks(st_platform, st_device_ids, capfd):
 @pytest.mark.runtime("tensormap_and_ringbuffer")
 @pytest.mark.device_count(1)
 @pytest.mark.manual(["a2a3sim", "a5sim"])
-def test_spmd_task_aggregates_across_threads(st_platform, st_device_ids, capfd):
+def test_spmd_task_aggregates_across_threads(st_platform, st_device_ids, capfd, drain_host_log):
     # One SPMD task (block_num=8) tagged slot 0; blocks dispatch across multiple
     # scheduler threads and must reduce to one complete slot. out = a + b.
     _drive(st_platform, int(st_device_ids[0]), "task_timing_spmd_orchestration", 1)
-    err = capfd.readouterr().err
+    err = drain_host_log(capfd)
 
     slot0 = _slot_spans(err, 0)
     assert len(slot0) == 1, f"expected exactly one task_slot_0 marker for the SPMD task, got {slot0}"