diff --git a/tests/st/a2a3/host_build_graph/concurrent_prepare_stress/test_concurrent_prepare_stress.py b/tests/st/a2a3/host_build_graph/concurrent_prepare_stress/test_concurrent_prepare_stress.py index 412f78320a..f82e439923 100644 --- a/tests/st/a2a3/host_build_graph/concurrent_prepare_stress/test_concurrent_prepare_stress.py +++ b/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())) diff --git a/tests/st/a2a3/host_build_graph/native_run_lifecycle/test_native_run_lifecycle.py b/tests/st/a2a3/host_build_graph/native_run_lifecycle/test_native_run_lifecycle.py index 8561aa061c..b3b7c1fdab 100644 --- a/tests/st/a2a3/host_build_graph/native_run_lifecycle/test_native_run_lifecycle.py +++ b/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. diff --git a/tests/st/conftest.py b/tests/st/conftest.py index a83bfa018c..8db1b6e9bc 100644 --- a/tests/st/conftest.py +++ b/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 diff --git a/tests/st/task_timing/task_timing_slots/test_task_timing_e2e.py b/tests/st/task_timing/task_timing_slots/test_task_timing_e2e.py index 917089f775..bdce896d73 100644 --- a/tests/st/task_timing/task_timing_slots/test_task_timing_e2e.py +++ b/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,7 +176,7 @@ 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 @@ -184,7 +184,7 @@ def test_duplicate_slot_merges_window(st_platform, st_device_ids, capfd): # (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}"