[andrewwhitecdw]

This PR addresses the following issue in tests/jax/test_distributed_fused_attn.py: scan-ring env var may leak across tests on pytest.skip.

Changes

• tests/jax/test_distributed_fused_attn.py: scan-ring env var may leak across tests on pytest.skip.

Details

--- a/tests/jax/test_distributed_fused_attn.py
+++ b/tests/jax/test_distributed_fused_attn.py
@@ -1,7 +1,17 @@
-        if use_scan_ring:
-            os.environ["NVTE_FUSED_RING_ATTENTION_USE_SCAN"] = "1"
-        else:
-            os.environ["NVTE_FUSED_RING_ATTENTION_USE_SCAN"] = "0"
-...
-        runner.test_backward()
-        del os.environ["NVTE_FUSED_RING_ATTENTION_USE_SCAN"]
+@contextmanager
+def _scan_env(use_scan_ring):
+    """Set NVTE_FUSED_RING_ATTENTION_USE_SCAN and restore the prior value on exit."""
+    key = "NVTE_FUSED_RING_ATTENTION_USE_SCAN"
+    old_value = os.environ.get(key)
+    os.environ[key] = "1" if use_scan_ring else "0"
+    try:
+        yield
+    finally:
+        if old_value is None:
+            os.environ.pop(key, None)
+        else:
+            os.environ[key] = old_value
+
+...
+        with _scan_env(use_scan_ring):
+            runner.test_backward()

Tests

• tests/jax/test_distributed_fused_attn.py

--- a/tests/jax/test_distributed_fused_attn.py
+++ b/tests/jax/test_distributed_fused_attn.py
@@ -486,3 +499,25 @@ class TestReorderCausalLoadBalancing:
         reordered = reorder(tensor, reorder_strategy, cp_size, seq_dim, stripe_size)
         inversed = inverse(reordered, reorder_strategy, cp_size, seq_dim, stripe_size)
 
         assert jnp.array_equal(inversed, ref)
+
+
+def test_scan_env_restored():
+    """_scan_env restores the original env value even on exception."""
+    key = "NVTE_FUSED_RING_ATTENTION_USE_SCAN"
+    sentinel = "original"
+    original = os.environ.get(key)
+    os.environ[key] = sentinel
+    try:
+        with _scan_env(True):
+            assert os.environ.get(key) == "1"
+            raise RuntimeError("expected")
+    except RuntimeError:
+        pass
+    assert os.environ.get(key) == sentinel
+    if original is None:
+        os.environ.pop(key, None)
+    else:
+        os.environ[key] = original

---

Greptile feedback addressed

• Wrapped the outer cleanup of test_scan_env_restored in try/finally, so an unexpected assertion or exception no longer leaks the sentinel value into later tests.

Local verification: python3 -m py_compile tests/jax/test_distributed_fused_attn.py passed. Full pytest execution was not feasible because JAX is not installed in this environment.

[greptile-apps]

Greptile Summary

The PR prevents NVTE_FUSED_RING_ATTENTION_USE_SCAN from leaking between distributed JAX attention tests.

• Adds a context manager that restores the environment variable’s prior value on every exit path.
• Restricts the override to backward-test execution.
• Adds exception-path coverage and guarantees that the test itself restores its sentinel setup.

Confidence Score: 5/5

The PR appears safe to merge.

No blocking failure remains; the previously reported outer-cleanup issue is fixed by the test’s new finally block.

Important Files Changed

Filename	Overview
tests/jax/test_distributed_fused_attn.py	The context-managed environment override and regression test safely address the previously reported cleanup failure.

_{Reviews (6): Last reviewed commit: "fix: scan-ring env var may leak across t..." | Re-trigger Greptile}

[andrewwhitecdw]

@greptile-apps thanks for the review. I applied the suggested fix:

• Wrapped the outer cleanup in test_scan_env_restored with try/finally, so the original value of NVTE_FUSED_RING_ATTENTION_USE_SCAN is always restored even if the test assertion fails or an unexpected exception escapes.

The branch has been amended to a single signed-off commit (376a997) and force-pushed. python3 -m py_compile tests/jax/test_distributed_fused_attn.py passed; full pytest was not feasible because JAX is not installed in this environment.

[greptile-apps]

The fix looks correct. The test_scan_env_restored cleanup is now properly wrapped in try/finally, so NVTE_FUSED_RING_ATTENTION_USE_SCAN is always restored to its pre-test state regardless of whether the inner assertion fails or an unexpected exception escapes.

The overall implementation is clean:

• _scan_env correctly snapshots and restores the prior value (including the absent-key case via pop)
• The with _scan_env(use_scan_ring): call site in impl_test_context_parallel_attn eliminates the original leak-on-skip bug
• test_scan_env_restored is now itself exception-safe, which is the right thing to do for a test that mutates global state

LGTM.