CP Tests batching using subprocess worker pool

[sudhakarsingh27]

Design

Problem

8×H100, test_essential=True, 38 runnable CP attention configs:

• ~554 s wall when each config runs in its own torchrun (baseline measured in [PyTorch] Batch CP attention tests in single torchrun to amortize NCC… #2965).
• Of that, only ~216 s is the test work itself (CP fwd/bwd + assertions, ~5.7 s/config).
• The remaining ~338 s is per-spawn overhead — Python imports + NCCL global init/teardown — paid 38 times at ~8.9 s each.

We need that overhead amortised, without changing how tests are written or how skips report. #2965 takes one approach (dry-run + per-batch torchrun). This PR proposes a simpler alternative.

Approach

One long-lived torchrun per world_size, fed by JSON-over-stdio. No dry-run, no batch chunking, no two-pass dispatch.

A session-scoped fixture spawns at most two workers (one for world_size=2, one for world_size=4, both lazily). Each test body calls pool.submit(kwargs) which writes one JSON request line to rank-0's stdin and reads one JSON response line from rank-0's stdout. NCCL initialises once per pool; every subsequent case reuses the same process group.

pytest (session)
    ├─ fixture cp_pool spawns torchrun --standalone (world_size=2)
    │    ├─ rank 0 ── stdin/stdout ──┐
    │    ├─ rank 1                   │
    │                                │
    ├─ test_cp_with_flash_attention[case0] ── pool.submit(kwargs) ──┘
    ├─ test_cp_with_flash_attention[case1] ── pool.submit(kwargs)
    ├─ ...
    └─ fixture teardown sends {"op":"shutdown"} to each pool

How the pool works

Worker (run_attention_with_cp_pool.py) is launched once per world_size and loops:

def main():
    rank = int(os.environ["RANK"])
    world_size = int(os.environ["WORLD_SIZE"])
    torch.cuda.set_device(rank % torch.cuda.device_count())
    dist.init_process_group(backend="nccl", rank=rank, world_size=world_size)
    os.environ["NVTE_CP_POOL_PG"] = "1"  # tell run_dpa_with_cp to reuse our PG

    while True:
        req = _recv_request(rank)              # rank 0 reads stdin, broadcasts
        if req.get("op") == "shutdown":
            break
        ok, msg = _run_one(req, rank)          # _reset_between_cases + run_dpa_with_cp
        gathered = [None] * world_size
        dist.gather_object((ok, msg), gathered if rank == 0 else None, dst=0)
        if rank == 0:
            _send_response(rank, {"ok": all_ok, "error": first_failing_traceback})
    dist.destroy_process_group()

run_dpa_with_cp checks NVTE_CP_POOL_PG and skips its own dist.init_process_group call; on teardown it destroys only the per-case CP sub-groups, leaving the main PG intact for the next case.

Pytest side (PoolWorker) is a thin wrapper around subprocess.Popen of python -m torch.distributed.run --standalone. submit() writes one request line, selects on stdout with a configurable timeout, reads one sentinel-prefixed JSON response line. On timeout or pipe error it terminates the pool; the next submit() lazily respawns it.

A daemon thread drains the worker's stderr into a bounded ring buffer (200 lines / ~40 KB) and echoes each line live to sys.stderr. On a crash/timeout, the last 4 KB of that buffer is attached to the AssertionError raised on the failing test — so CI JUnit XML carries the actual cause (NCCL error, Python traceback, OOM) inline rather than just "pool worker died". Equivalent in spirit to PR #2965's run_distributed() stderr capture.

Bug found while validating: stream race in AttnFuncWithCPAndKVAllGather.forward

A latent TE bug surfaced once the persistent pool ran the full test_essential=False matrix. In the all-gather CP forward, max_logit_per_step[i] is written inside with torch.cuda.stream(flash_attn_streams[i]): — and for i=1 that's cp_stream, not the default stream. Later, at loop iteration i=2, the code reads max_logit_per_step[1] via torch.maximum(...) on the default stream without a wait_stream in between. The post-loop current_stream().wait_stream(cp_stream) is too late — the race has already fired.

The race is latent: outcome depends on stream scheduling. In a fresh-process subprocess (one-torchrun-per-test path) streams are cleanly initialised and timing happens to put the write before the read. In a long-running persistent-worker process — exactly what the pool exposes — prior workloads leave stream state in a different shape and the read can fire before the write completes. The result was max_logit stale on 3 of 12 head entries (~0.3 abs diff), surfacing in 5 specific configs (all cp_comm_type=all_gather, cp_1_0/cp_1_1, bshd or fp16).

Fix is one line in context_parallel.py: current_stream().wait_stream(flash_attn_streams[i-1]) before the torch.maximum read. No-op when streams are identical (i=1), only fires when reading from cp_stream (i=2). Independently useful for anyone running CP attention in a long-lived process, not just the pool.

Performance

Measured back-to-back on the same 8×H100 box.

test_essential=True (38 runnable configs: 34 × 2-GPU + 4 × 4-GPU)

Run	Approach	torchrun spawns	Wall	Pass / Fail	Speedup
Unbatched (baseline)	one torchrun per config	38	~554 s	n/a (PR #2965)	1.0×
#2965 B=16	dry-run + per-batch torchrun	4	263 s	38 / 0	2.1×
This PR (pool)	persistent NCCL pool per world_size	2	248 s	38 / 0	2.23×

test_essential=False (full matrix, 50 976 collected, 348 runnable)

Run	Approach	Wall	Pass / Fail	Notes
#2965 (B=16)	dry-run + per-batch torchrun	28 m 10 s (1 690 s)	328 / 0	clean
This PR (pool) — before the all-gather race fix	persistent pool	26 m 25 s (1 585 s)	323 / 5	5 latent-TE-bug failures; same configs pass on #2965 because per-batch fresh-process re-init hides the race
This PR (pool) — current	persistent pool + race fix	26 m 23 s (1 583 s)	348 / 0	clean

Pool is ~7 % faster on the full matrix. Note the runnable count rises from 328 to 348 once the race fix lands — the 20 extra cases were silently dropped earlier by the same numerical-corruption path.

Knobs

Env var	Effect
NVTE_CP_POOL_TIMEOUT_SEC=N	Max seconds the fixture waits for one case's response. Default 90. Slowest observed case on H100 (test_essential=True) is ~15 s; 90 s gives ~6× headroom. Override for slower machines or heavier matrices.
NVTE_CP_POOL_TIMING=1	Emit [POOL-TIMING] case_idx=N world_size=W wall_s=X.XXX ok=B on rank-0 stderr per case. Off by default. Used to recalibrate the timeout against a new matrix.
NVTE_CP_POOL_PG=1	Set by the pool worker itself before calling run_dpa_with_cp; tells that function not to call dist.init_process_group and to leave the main PG alone on teardown. Not for end-user use.

There is intentionally no batch-size knob — there's no concept of a batch to size.

Adding a pooled test

1. Write the test as usual: @pytest.mark.parametrize stack + inline pytest.skip(...) checks.
2. Add cp_pool to the function signature.
3. At the end, do pool = cp_pool(num_gpus); _submit(pool, **kwargs) where kwargs becomes run_dpa_with_cp(**kwargs).

That's it. No two-pass logic, no fixture stubs.

Failure semantics

Outcome	What pytest sees
Inline pytest.skip(...) fires	Standard SKIP — body returns before pool.submit, no pool work.
@pytest.mark.skip(if) marker fires	Standard SKIP via pytest's normal path.
Config ran, assertion failed on any rank	FAIL with the first failing rank's full traceback (via dist.gather_object). Other ranks' tracebacks visible in subprocess stderr. No retry — real test failures propagate.
Pool worker timed out	Retried once on a fresh pool. If the retry passes, test PASSES with a [POOL-RETRY] line in <system-err>. If the retry also dies, FAIL with "pool worker (world_size=N) timed out after 90s; ..." plus the last 4 KB of the worker's stderr and a [POOL-RETRY-FAIL] line summarising both attempts.
Pool worker died mid-request	Retried once on a fresh pool, same as above. Final FAIL message: "pool worker died mid-request".
Pool worker died before request	Retried once on a fresh pool, same as above. Final FAIL message: "pool worker died before request could be sent".

Cross-rank failure detail is strictly better than all_reduce(ok, op=MIN) (which only tells you some rank failed): gather_object brings back each rank's (ok, traceback) tuple so the reported error is the actual non-zero-rank stack trace, not "see subprocess stderr."

What happens when a pool worker stalls

Three terminal states for any submit(): response arrives → normal handling; no response after POOL_SUBMIT_TIMEOUT_SEC → timeout path; process died → mid-request-death path. Any stall (application hang, NCCL deadlock, GPU wedge, even a stdout-pipe-full self-deadlock) eventually resolves to one of the latter two. The pool is SIGTERM'd (5 s grace) then SIGKILL'd, and submit() lazily respawns a fresh pool of the same world_size (~6–9 s NCCL re-init) for the retry attempt. The other pool (different world_size) is unaffected. Blast radius: one respawn per stall; the test only FAILs if the retry also dies.

Single-retry on pool-infrastructure failures

The pool worker subprocess can die mid-case due to async NCCL aborts or flaky 4-GPU collective state that doesn't reproduce on a fresh pool — we observed this once per ~350-case full-matrix run on cp_comm_type=a2a+p2p (4-GPU pool, fused-attention). Each submit() retries once if the failure mode matches "pool worker died", "timed out", or "before request could be sent". Test-assertion failures (the worker returned {"ok": false, "error": ...}) do not retry — those are real bugs and propagate.

Every retry leaves a [POOL-RETRY] line in stderr; if both attempts die, a [POOL-RETRY-FAIL] line summarising the first attempt's error is also written. pytest captures per-test stderr and writes it into the JUnit <testcase>/<system-err> field, so a test that PASSED on retry is visible to anyone querying CI artifacts for POOL-RETRY — a 50 %-flaky test would show up as ~50 % of its CI runs having a [POOL-RETRY] line in <system-err>, even though the case row is green.

Mitigations for shared-process state

All cases share one Python process and one NCCL world per world_size, so anything that needs a clean per-test starting point is reset before each case (in _run_one, not in a finally block, so the first case is also clean):

• torch.manual_seed(1234) + torch.cuda.manual_seed(1234) — RNG reseeded so input tensors are reproducible per case.
• FP8GlobalStateManager.reset() — drops FP8 amax history etc. that would otherwise leak across cases.
• torch.cuda.empty_cache().
• copy.deepcopy(model_configs_*[model]) inside run_dpa_with_cp — the THD branch rewrites attn_mask_type in place; without deepcopy the change leaks into the module-level dict.

run_dpa_with_cp itself re-sets NVTE_FUSED_ATTN/NVTE_FLASH_ATTN unconditionally at the top of every call and pops the FP8-related transient env vars, so no explicit env-key reset is needed in the pool worker.

The single-shot run_dpa_with_cp does some of these inherently (it's a fresh process). For the pooled path we replicate them explicitly so the two execution modes produce identical per-case state.

Edge cases

1. Cross-pytest-session port collisions. torch.distributed.run --standalone picks a free rendezvous port at bind time. No MASTER_PORT plumbing needed; parallel pytest sessions (e.g., L1_pytorch_distributed_unittest's det vs non-det concurrent runs on disjoint GPU sets) cannot collide.
2. Pool wedges after N cases. Per-case timeout is 90 s by default. On expiry, _kill() terminates the pool; the retry attempt lazily respawns a fresh worker. Blast radius: one respawn (~6–9 s) plus, if the retry also dies, one FAIL.
3. Pool worker crashes mid-case. BrokenPipeError / empty read → AssertionError. The wrapper retries once on a fresh pool; the failing test only surfaces as FAILED if the retry also dies.
4. Two pools (2-GPU and 4-GPU) co-resident on overlapping low-index GPUs. torch.cuda.set_device(rank % device_count) means both pools claim GPUs 0–N starting at 0. They never run CUDA concurrently (pytest serialises tests), so the idle pool only holds ~1 GB of CUDA context per shared GPU — well within H100's 80 GB. NCCL worlds are independent. No collision. This is intentional: setting an explicit CUDA_VISIBLE_DEVICES per pool would prevent the parallel-pytest-session pattern used by qa/L1_pytorch_distributed_unittest/test.sh (det / non-det on disjoint GPU sets); see Follow-ups.
5. NVTE_CP_POOL_PG collision. The env var is set by the pool worker after dist.init_process_group and only read by run_dpa_with_cp — no other consumer. If end-user code accidentally sets it without an init'd PG, run_dpa_with_cp will fail when it tries to use the (non-existent) PG. Harmless; same failure as setting any other internal flag incorrectly.

Validation

Run	Result
8×H100, test_essential=True	38 passed / 10234 skipped / 0 failed in 248 s
8×H100, test_essential=False	348 passed / 50628 skipped / 0 failed in 26 m 23 s

Per-case wall-time distribution on H100 (test_essential=True, with NVTE_CP_POOL_TIMING=1): min 1.87 s, p50 4.77 s, p95 12.43 s, max 15.39 s. Drove the POOL_SUBMIT_TIMEOUT_SEC default of 90 s (~6× max).

Deterministic mode

qa/L1_pytorch_distributed_unittest/test.sh on this branch runs test_attention_with_cp.py twice — once with default NVTE_ALLOW_NONDETERMINISTIC_ALGO=1, once with =0 — in parallel on disjoint GPU sets when ≥8 GPUs are available (sequentially otherwise). Both modes pass cleanly.

An inline skip at test_attention_with_cp.py:605-609 handles a known cuDNN sm90 issue:

if qkv_format == "thd" and config.num_heads >= 20 and get_device_compute_capability() == (9, 0):
    pytest.skip(
        "Deterministic FusedAttention backward with THD format OOMs on sm90"
        " (cuDNN reserves memory proportional to bHSS)."
    )

Without this skip, 5 fused-attention THD configs OOM under NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 with Tried to allocate 512.01 GiB inside tex.fused_attn_bwd (workspace-size pathology in cuDNN's deterministic backward path on Hopper, not actual memory exhaustion). The skip is independent of the pool design but lives in the same test file; documented here so deterministic-mode skips don't get mistaken for batching-infra bugs.

Follow-ups (not in this PR)

Items intentionally deferred to keep this PR scoped:

• L8 — drop the try/except Exception: pass around dist.destroy_process_group calls. Real semantic change: would surface errors that the swallow currently hides. Needs its own validation.
• M1 — replace the NVTE_CP_POOL_PG env-var contract between the pool worker and run_dpa_with_cp with a module-level flag (run_attention_with_cp._POOL_MANAGED_PG = True). Cleaner, type-checkable; same effect.
• M2 — redirect non-rank-0 stdout to /dev/null at worker startup instead of sentinel-prefix scanning. Closes the stdout-pollution class at the source rather than papering over it.
• M3 — intentionally not done. Setting explicit CUDA_VISIBLE_DEVICES per pool (so 2-GPU and 4-GPU pools claim disjoint GPUs) would break the parallel-pytest-session pattern used by qa/L1_pytorch_distributed_unittest/test.sh. Each top-level pytest session sets CUDA_VISIBLE_DEVICES itself (e.g. det on GPUs 0-3, non-det on 4-7); the pool inherits that and uses rank % device_count to map within the session's slice. Adding per-pool device pinning would override the session's slice and produce overlap across sessions. The current "overlap within a session, idle context only" behaviour (edge case 4) is the right trade-off.

Files

• transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py — wait_stream fix in AttnFuncWithCPAndKVAllGather.forward (the latent stream race the pool exposed).
• tests/pytorch/attention/test_attention_with_cp.py — cp_pool session fixture + PoolWorker (lazy spawn, --standalone, 90 s timeout, kill-and-respawn, stderr drainer thread + tail attached to crash AssertionErrors, single retry on pool-infrastructure failures with [POOL-RETRY] stderr tag).
• tests/pytorch/attention/run_attention_with_cp_pool.py — pool worker: init NCCL once, dispatch loop, _reset_between_cases (seed/FP8/cache), gather_object for cross-rank failure detail, optional NVTE_CP_POOL_TIMING=1 per-case timing log.
• tests/pytorch/attention/run_attention_with_cp.py — run_dpa_with_cp honours NVTE_CP_POOL_PG (skips its own PG init/destroy), deep-copies model configs.
• qa/L1_pytorch_distributed_unittest/test.sh — runs test_attention_with_cp.py twice (non-det + det) in parallel on disjoint GPU sets when ≥8 GPUs available; sequential fallback otherwise.

Comparison to #2965

Both PRs solve the same problem; this one is structurally smaller and has fewer concepts.

	#2965 (dry-run + batched torchrun)	This PR (persistent pool)
Batching-only diff vs origin/main	355 + / 69 − (2 files)	593 + / 111 − (3 files)
Knobs	CP_TEST_BATCH_SIZE, CP_TEST_BATCH_RETRY	NVTE_CP_POOL_TIMEOUT_SEC (default 90 s), NVTE_CP_POOL_TIMING
Required concepts	dry-run, _COLLECT_MODE, _DummyRequest, _item_static_skip, _BACKEND_CACHE, batch chunking, atomic JSON flush, singleton retry	long-lived worker, JSON-over-stdio, single retry on infra failure
Cross-rank failure	all_reduce(ok, MIN) — boolean only	gather_object — full traceback per rank
Crash-path stderr in AssertionError	✓ (run_distributed() attaches last 4 KB)	✓ (drainer thread, last 4 KB attached)
Master-port handling	needs explicit MASTER_PORT env per parallel pytest session	--standalone, automatic
Per-case timeout	none	90 s default
Spawns (test_essential=True)	4 @ B=16, 2 @ B=50	2 always
Pre-pytest overhead	~14 s dry-run + collection	none
Full-matrix wall (H100, 348 runnable)	n/a (#2965 measured at 328 runnable, before the race fix)	26 m 23 s
Full-matrix correctness	clean (race didn't surface, fresh process per batch hides it)	clean (race fixed in TE)

Type of change

• Code refactoring (test infrastructure)
• Bug fix (latent CP stream race exposed by the pool — fix in TE source)

Checklist

• Contributing guidelines followed
• Functionality complete
• Code commented where non-obvious
• Documentation (n/a — internal test infra)
• No new warnings
• Existing test suite serves as input + validation
• Existing tests pass locally (8×H100, test_essential=False: 348 / 0, both det and non-det)

[greptile-apps]

Greptile Summary

This PR replaces per-test torchrun spawning with a persistent worker-pool design: one long-lived torchrun --standalone process per world_size (2-GPU and 4-GPU), fed JSON requests over rank-0 stdin/stdout with a sentinel prefix to guard the protocol against interleaved output. It also patches a latent CUDA stream race in AttnFuncWithCPAndKVAllGather.forward that the persistent pool exposed.

• Pool infrastructure (test_attention_with_cp.py, run_attention_with_cp_pool.py): Session-scoped cp_pool fixture lazily spawns a PoolWorker per world_size; each worker initialises NCCL once, loops on requests via gather_object for cross-rank failure capture, and a daemon thread drains its stderr into a ring buffer that is attached to crash AssertionErrors.
• run_attention_with_cp.py gains NVTE_CP_POOL_PG awareness to skip redundant init_process_group/destroy_process_group calls, plus copy.deepcopy for model configs to prevent attn_mask_type mutations from leaking across pool cases.
• context_parallel.py stream-race fix: adds current_stream().wait_stream(flash_attn_streams[i-1]) before the torch.maximum(max_logit, max_logit_per_step[i-1]) read at loop iteration i=2, where the write happened on cp_stream and the read was previously racing on the default stream.

Confidence Score: 5/5

Safe to merge. The stream-race fix is correct, the pool protocol is well-guarded with sentinel prefixes and a 90s per-case timeout, and test_essential is left as True.

The stream-race fix in context_parallel.py is a precise, well-commented one-liner that inserts the missing wait_stream at exactly the right point — no-op for the default-stream case (i=1) and fires only when reading across stream boundaries (i=2). The pool infrastructure handles all three terminal outcomes (response, timeout, worker death) with appropriate cleanup and respawn. Cross-rank failure detail via gather_object is strictly better than the all_reduce-min approach. No new uncovered failure modes were found beyond those already discussed in earlier review threads.

No files require special attention.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py	One-line stream-race fix in AttnFuncWithCPAndKVAllGather.forward: adds wait_stream(cp_stream) before torch.maximum reads max_logit_per_step[1], which was written on cp_stream but read on the default stream without synchronization.
tests/pytorch/attention/run_attention_with_cp_pool.py	New pool-worker entrypoint: NCCL init once, dispatch loop with gather_object for cross-rank traceback collection, optional POOL_TIMING instrumentation, and NCCL teardown in finally block.
tests/pytorch/attention/test_attention_with_cp.py	Adds PoolWorker class and cp_pool session fixture; test functions now call _submit(pool, **kwargs) instead of launching per-case subprocesses; test_essential is correctly left as True.
tests/pytorch/attention/run_attention_with_cp.py	run_dpa_with_cp gains NVTE_CP_POOL_PG awareness, deep-copies model configs to prevent cross-case mutation, and reuses pool-shared CP comm groups when available.

_{Reviews (9): Last reviewed commit: "[pre-commit.ci] auto fixes from pre-comm..." | Re-trigger Greptile}

[greptile-apps]

stdout pollution can silently corrupt the JSON protocol

torchrun (and the worker processes it spawns for ranks 1–N) all inherit the same stdout file descriptor as rank 0. If torchrun writes any status line to stdout, or if any non-rank-0 worker accidentally prints (e.g. via a print call in a library, NCCL debug output, or a Python warning), those bytes are interleaved with rank 0's JSON responses. The parent's readline() in PoolWorker.submit would then receive a non-JSON line and raise a json.JSONDecodeError, killing the pool and failing the test with a misleading error message.

Consider redirecting torchrun's own output or adding a sentinel prefix to every response line so the reader can skip unrecognised lines.

[sudhakarsingh27]

Implemented the stronger fix in 36b65fb3: redirect non-rank-0 stdout to /dev/null at the fd level via dup2 at worker startup. Catches both Python print and C-level (NCCL, libc, etc.) writes — the sentinel-prefix approach could only mitigate the former by scanning + skipping.

With non-rank-0 stdout silenced, rank 0's JSON line is the only thing that reaches the parent's pipe, so the _RESP_PREFIX machinery + the sentinel-scanning while loop in PoolWorker._submit_once are gone. The reader collapses to one select + one readline + one json.loads.

Validated on 8×H100 (test_essential=True, flash-attn): 9 passed / 55 skipped / 0 failed in 56 s; no JSONDecodeError, no protocol corruption. Closes the M2 follow-up listed in the PR description.

[sudhakarsingh27]

Thanks for the review @greptile-apps — all three findings addressed in e162a9ec:

P1 — NCCL communicator leak on exception (run_attention_with_cp.py)

Wrapped the body of run_dpa_with_cp in try/finally. cp_comm_sub_groups = [] is now initialised unconditionally before the a2a+p2p branch so the finally block is safe even when (a) cp_comm_type != "a2a+p2p" or (b) the inline assert at the top of the a2a+p2p block fires before the populate loop. Each dist.destroy_process_group() is itself wrapped in try/except Exception: pass so a destroy failure on one group can't leak the others.

P2 — stdout protocol can be corrupted by interleaved chatter (run_attention_with_cp_pool.py + test_attention_with_cp.py)

Worker now prefixes every response with [CP_POOL_RESP] (constant _RESP_PREFIX). PoolWorker.submit() reads lines in a loop and skips anything that doesn't start with the sentinel — echoing the non-protocol lines to stderr so they're still visible in CI logs. Bounded by _MAX_NOISE_LINES = 1000 so a chatty worker can't keep us spinning past the deadline. The timeout budget is preserved across iterations via a monotonic deadline.

P2 — select.select on a pipe fd is Linux/macOS only (test_attention_with_cp.py)

Added a short comment noting the Windows portability gap. Tests target Linux GPU hosts only; this is a doc issue, not a real bug, but flagged for future readers.

PR-body note

The "default 600 s" inconsistency you flagged on the env-var table was already reconciled in the previous push — POOL_SUBMIT_TIMEOUT_SEC defaults to 90 s in both code and description. (You reviewed at 385e9662, which was pre-polish.)

Validation

Re-ran test_essential=True on H100 after the three fixes: 38 passed / 0 failed in 270 s (was 248 s pre-P2 — the +22 s is the new sentinel-scan loop's per-line overhead, within noise).

The 5 state-cascade failures on test_essential=False documented under Known issues are unchanged by these fixes (different root cause — process-shared state outside the _reset_between_cases hook). Tracking separately.

[greptile-apps]

All three fixes look correct. Let me note one small remaining issue in the deadline logic before signing off.

In submit(), the deadline is initialized after the first select.select returns, not before the loop:

deadline = None
while scanned < self._MAX_NOISE_LINES:
    remaining = timeout if deadline is None else max(0.0, deadline - time.monotonic())
    ready, _, _ = select.select([self.proc.stdout], [], [], remaining)
    if deadline is None:
        deadline = time.monotonic() + timeout   # ← set after first select completes

If the first line (noise or response) arrives after, say, 50 s, the deadline is then set to now + 90 s, making the total possible wall time for a chatty worker ~140 s instead of 90 s. The comment "The timeout budget is preserved across iterations via a monotonic deadline" overstates what's implemented — the deadline only bounds iterations 2..N, not the first.

The fix is to set the deadline once, before the loop:

deadline = time.monotonic() + timeout
while scanned < self._MAX_NOISE_LINES:
    remaining = max(0.0, deadline - time.monotonic())
    ready, _, _ = select.select([self.proc.stdout], [], [], remaining)
    if not ready:
        ...

In practice this won't matter much — torchrun status chatter arrives nearly instantaneously before the response — but the current code doesn't actually enforce the stated timeout bound end-to-end.

Everything else looks good:

• P1: cp_comm_sub_groups = [] is now unconditionally initialized before the a2a+p2p branch, the test body is wrapped in try/finally, and each destroy_process_group() is individually guarded with try/except Exception: pass. Communicators will be reclaimed even if the body raises mid-assertion.
• P2 sentinel: _RESP_PREFIX = "[CP_POOL_RESP] " is written by the worker and the parent scans for it correctly, echoing non-sentinel lines to stderr. The _MAX_NOISE_LINES = 1000 cap prevents an infinite spin.
• P2 doc: The select.select Windows portability note is in place.

The deadline fix above is the only remaining item; everything else is addressed cleanly.

Tip: You can customize Greptile's behavior for this repo with .greptile/rules.md and .greptile/config.json.

[greptile-apps]

Group creation sits outside the try/finally guard

cp_comm_group (line 250) and the sub-group creation loop (lines 260–263) all run before the try block that starts at line 268. If the sub-group loop raises — e.g., NCCL rejects a new_group call on a partially recovered communicator after a previous case failed — the finally block is never entered, so cp_comm_group and any sub-groups already appended to cp_comm_sub_groups are never destroyed. In a long-running pool this accumulates a communicator leak on every such failure, eventually hitting NCCL's internal communicator limit and breaking unrelated subsequent cases. Moving all three dist.new_group sites inside the try block (or at minimum wrapping the sub-group loop and the unconditional cp_comm_group creation in a nested try/finally) would close the gap.

[sudhakarsingh27]

Same fix as #3243445773 — the pool worker now owns the group lifecycle, so the dist.new_group calls live in the single-shot branch only (run_attention_with_cp.py:275-285, inside if not _reusing_pool_groups:). The original accumulation concern doesn't apply because the pool never creates per-case groups to leak; the single-shot path's subprocess exits at function return and the OS reclaims communicators. (01db9f5 + fc0d929)

[cyanguwa]

I think the structure is a lot simpler/cleaner than #2965 - thanks!

For the 5 remaining failing tests, I wonder if there's something else we need to reset for FP8 (please check with @ksivaman).

Also, please compare the number of tests before/after this PR and make sure we're still running the same number of tests! If the new, reduced runtime allows now, we can turn on test_essential=False, but I'll leave that to you.

If no major issues, I approve! Thanks!

[cyanguwa]

Do we want to destroy (and create earlier in the file) cp_comm_group for every config? I feel if the pool is the same, the world_size would be the same, and so is cp_comm_group?

[sudhakarsingh27]

Yes — exactly that, fixed in this direction. The pool worker creates cp_comm_group once per world_size at startup (run_attention_with_cp_pool.py:154); run_dpa_with_cp in pool mode reads the cached pointer instead of dist.new_group'ing per call (run_attention_with_cp.py:271-273). Destruction is once-per-worker at shutdown (run_attention_with_cp_pool.py:182-193). The a2a+p2p sub-groups are pre-created the same way. This dropped per-case NCCL overhead and also closed the leak class greptile raised in #3243445773 / #3244111437. Thanks for the catch — clean simplification. (01db9f5 + fc0d929)

[greptile-apps]

test_essential left as False — will break CI

The default was True (38 essential configs, all passing) and has been changed to False (full matrix, 328 runnable, 5 known failures). The PR description explicitly calls these out as "state-cascade failures" under Known Issues. Merging with test_essential = False will cause CI to report 5 failures by default, regressing from the current baseline.

[sudhakarsingh27]

Reverted in a018a53 — test_essential = True at line 47 to match the shipping default. The full-matrix run (test_essential=False, 348/0) in the PR description is a one-off validation pass, not the CI default.

[sudhakarsingh27]

/te-ci pytorch L3