[OMNIML-6181] [11/11] Share one Megatron worker pool across the gpu_megatron modules

tests/_test_utils/torch/distributed/utils.py

@@ -15,6 +15,7 @@
 
 import contextlib
 import faulthandler
+import gc
 import os
 import queue
 import signal
@@ -96,6 +97,36 @@ def get_live_pool(pools, key, make_pool):
     return pool
 
 
+_BASELINE_TORCH_FLAGS = {}
+
+
+def _capture_baseline_torch_flags():
+    """Remember the torch settings a freshly spawned worker starts with."""
+    _BASELINE_TORCH_FLAGS.update(
+        cudnn_deterministic=torch.backends.cudnn.deterministic,
+        cudnn_benchmark=torch.backends.cudnn.benchmark,
+        cudnn_tf32=torch.backends.cudnn.allow_tf32,
+        matmul_tf32=torch.backends.cuda.matmul.allow_tf32,
+        deterministic_algorithms=torch.are_deterministic_algorithms_enabled(),
+    )
+
+
+def reset_worker_state(rank, world_size):
+    """Return a persistent worker to the state it had right after spawning.
+
+    ``set_seed`` and similar helpers change process-global torch settings; run this job at the
+    boundary between test modules that share workers.
+    """
+    torch.backends.cudnn.deterministic = _BASELINE_TORCH_FLAGS["cudnn_deterministic"]
+    torch.backends.cudnn.benchmark = _BASELINE_TORCH_FLAGS["cudnn_benchmark"]
+    torch.backends.cudnn.allow_tf32 = _BASELINE_TORCH_FLAGS["cudnn_tf32"]
+    torch.backends.cuda.matmul.allow_tf32 = _BASELINE_TORCH_FLAGS["matmul_tf32"]
+    torch.use_deterministic_algorithms(_BASELINE_TORCH_FLAGS["deterministic_algorithms"])
+    torch.manual_seed(1234)
+    gc.collect()
+    torch.cuda.empty_cache()
+
+
 def default_worker_teardown(rank, world_size):
     """Minimal cleanup between tests in persistent workers."""
     try:
@@ -208,6 +239,7 @@ def _worker_loop(rank, world_size, backend, port, cmd_queue, result_queue, teard
         if backend == "nccl" and torch.cuda.is_available():
             torch.cuda.set_device(rank)
         torch.manual_seed(1234)
+        _capture_baseline_torch_flags()
         result_queue.put("ready")
 
         while True:
@@ -245,7 +277,9 @@ def run(self, fn, *args, **kwargs):
         not finish within ``failure_grace`` (they are typically blocked in a collective waiting for
         it), the workers are still inside the job and every later ``run`` would hang until its own
         timeout. The workers are killed instead and the pool is marked ``broken``; callers that
-        want to continue need a new pool.
+        want to continue need a new pool. A pool is also retired after any failed job: a failure
+        can leave its workers unusable (a CUDA error such as an illegal memory access is sticky,
+        and every later job in that process fails with it).
         """
         if self.broken:
             raise RuntimeError("Worker pool was killed after an earlier job hung; create a new one")
@@ -270,12 +304,15 @@ def run(self, fn, *args, **kwargs):
             raise
 
         if errors:
+            if not self.broken:
+                # Every rank replied, but the workers may be poisoned: never hand them to another job.
+                self._abort(dump_stacks=False)
             raise RuntimeError("Worker(s) failed:\n" + "\n".join(errors))
 
-    def _abort(self):
-        """Print every worker's stack (to stderr, for the CI log), then kill the workers."""
+    def _abort(self, dump_stacks=True):
+        """Kill the workers, first printing their stacks (to stderr, for the CI log) if asked."""
         self.broken = True
-        if hasattr(signal, "SIGUSR1"):
+        if dump_stacks and hasattr(signal, "SIGUSR1"):
             for p in self._processes:
                 with contextlib.suppress(OSError):
                     os.kill(p.pid, signal.SIGUSR1)

tests/gpu/torch/utils/test_distributed_worker_pool.py

@@ -22,8 +22,13 @@
 from types import SimpleNamespace
 
 import pytest
+import torch
 import torch.distributed as dist
-from _test_utils.torch.distributed.utils import DistributedWorkerPool, get_live_pool
+from _test_utils.torch.distributed.utils import (
+    DistributedWorkerPool,
+    get_live_pool,
+    reset_worker_state,
+)
 
 
 def _hang(rank, size):
@@ -36,12 +41,43 @@ def _rank_0_fails_and_rank_1_waits_for_it(rank, size):
     dist.barrier()  # never completes: its partner left
 
 
+def _rank_0_fails(rank, size):
+    if rank == 0:
+        raise ValueError("rank 0 failed")
+
+
 def _rank_0_dies_and_rank_1_waits_for_it(rank, size):
     if rank == 0:
         os._exit(13)  # what an illegal memory access or an OOM kill does: no exception, no result
     time.sleep(600)
 
 
+def _torch_settings():
+    return (
+        torch.backends.cudnn.deterministic,
+        torch.backends.cudnn.benchmark,
+        torch.backends.cudnn.allow_tf32,
+        torch.backends.cuda.matmul.allow_tf32,
+        torch.are_deterministic_algorithms_enabled(),
+    )
+
+
+def _change_the_settings_then_reset_them(rank, size):
+    spawn_time = _torch_settings()
+    torch.backends.cudnn.deterministic = not spawn_time[0]
+    torch.backends.cudnn.benchmark = not spawn_time[1]
+    torch.backends.cudnn.allow_tf32 = not spawn_time[2]
+    torch.backends.cuda.matmul.allow_tf32 = not spawn_time[3]
+    torch.use_deterministic_algorithms(not spawn_time[4])
+    torch.manual_seed(5)
+    assert _torch_settings() != spawn_time
+
+    reset_worker_state(rank, size)
+
+    assert _torch_settings() == spawn_time
+    assert torch.initial_seed() == 1234
+
+
 def test_a_hung_job_kills_the_pool_instead_of_stalling_later_runs(capfd):
     pool = DistributedWorkerPool(world_size=1, backend="gloo", teardown_fn=None)
     try:
@@ -74,6 +110,21 @@ def test_a_failed_rank_does_not_make_its_peers_wait_out_the_job_timeout():
         pool.shutdown()
 
 
+def test_a_failed_job_retires_the_pool_even_when_every_rank_replied():
+    pool = DistributedWorkerPool(world_size=2, backend="gloo", teardown_fn=None)
+    try:
+        with pytest.raises(RuntimeError, match="rank 0 failed"):
+            pool.run(_rank_0_fails)  # rank 1 returns at once, so no rank is left waiting
+
+        # A sticky CUDA error would fail every later job in these workers
+        assert pool.broken
+        assert not any(p.is_alive() for p in pool._processes)
+        with pytest.raises(RuntimeError, match="killed"):
+            pool.run(_rank_0_fails)
+    finally:
+        pool.shutdown()
+
+
 def test_a_dead_worker_fails_the_run_at_once_and_kills_its_peers():
     pool = DistributedWorkerPool(world_size=2, backend="gloo", teardown_fn=None)
     try:
@@ -104,6 +155,14 @@ def test_a_failed_start_raises_and_leaves_no_workers(kwargs, error):
     assert {p.pid for p in multiprocessing.active_children()} <= before
 
 
+def test_reset_worker_state_restores_the_settings_a_worker_starts_with():
+    pool = DistributedWorkerPool(world_size=1, backend="gloo", teardown_fn=None)
+    try:
+        pool.run(_change_the_settings_then_reset_them)
+    finally:
+        pool.shutdown()
+
+
 def test_a_killed_pool_is_replaced_and_a_live_one_is_reused():
     pools = {}
     made = []

tests/gpu_megatron/conftest.py

@@ -13,11 +13,17 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import contextlib
+import gc
+import queue
 
 import pytest
 import torch
 from _test_utils.fs_utils import assert_unmodified_tree
-from _test_utils.torch.distributed.utils import DistributedWorkerPool, get_live_pool
+from _test_utils.torch.distributed.utils import (
+    DistributedWorkerPool,
+    get_live_pool,
+    reset_worker_state,
+)
 from _test_utils.torch.transformers_models import get_tiny_tokenizer
 from megatron.core.parallel_state import destroy_model_parallel
 
@@ -68,6 +74,8 @@ def _prestart_preprocess_forkserver():
 
 def megatron_worker_teardown(rank, world_size):
     """Clean up model-parallel state between tests in persistent workers."""
+    # Surface asynchronous CUDA errors at the test that caused them, not at a later one
+    torch.cuda.synchronize()
     if dist.is_initialized():
         dist.barrier()
     try:
@@ -79,6 +87,15 @@ def megatron_worker_teardown(rank, world_size):
             apex_destroy()
         except Exception as e:
             print(f"Error destroying model parallel with Apex: {e}")
+    # destroy_model_parallel() only drops Megatron's references to its NCCL groups, so every test's
+    # communicators (and their GPU buffers) piled up in the persistent workers (~20 groups and
+    # ~0.4 GiB per test), slowing later tests until they timed out. Destroy all but the default.
+    for group in list(torch.distributed.distributed_c10d._world.pg_map):
+        if group is not torch.distributed.GroupMember.WORLD:
+            torch.distributed.destroy_process_group(group)
+    # The finished test's modules can still reference the destroyed groups. Collect them here, with
+    # no kernels in flight, rather than at an arbitrary point of the next job.
+    gc.collect()
     torch.cuda.empty_cache()
 
 
@@ -90,38 +107,50 @@ def _make_pool(world_size):
     )
 
 
-@pytest.fixture(scope="module")
+@pytest.fixture(scope="session")
 def _pool_cache():
-    """Module-scoped cache of worker pools keyed by world_size.
+    """Session-scoped cache of worker pools keyed by world_size.
 
     Spinning up a pool cold-imports the full torch/megatron/modelopt stack per worker
-    (~tens of seconds), so fixtures that request the same world_size share one pool
-    instead of spawning duplicates — e.g. on a 2-GPU runner ``dist_workers`` and
-    ``dist_workers_size_2`` are both size 2. The cache is module-scoped and torn down at
-    module end, so workers are never reused across modules (avoids cross-test
-    state contamination).
+    (75-100 s in CI), so every module that requests the same world_size shares one pool,
+    e.g. on a 2-GPU runner ``dist_workers`` and ``dist_workers_size_2`` are both size 2.
+    Isolation between tests comes from ``megatron_worker_teardown`` (model-parallel state and
+    process groups after every job), from ``_reset_shared_pools`` (process-global torch settings
+    at each module boundary) and from ``get_live_pool`` replacing a pool that a failed, hung,
+    timed-out or crashed job retired (a sticky CUDA error would otherwise fail every later test).
     """
     pools: dict[int, DistributedWorkerPool] = {}
     yield pools
     for pool in pools.values():
         pool.shutdown()
 
 
+@pytest.fixture(scope="module", autouse=True)
+def _reset_shared_pools(_pool_cache):
+    """Workers outlive test modules: undo what the previous module did to process-global torch state."""
+    for pool in _pool_cache.values():
+        if not pool.broken:
+            # A pool whose reset fails is retired by run(); the next test that needs one gets a new
+            # pool, and the module's tests, which may not use a pool at all, are not errored.
+            with contextlib.suppress(RuntimeError, queue.Empty):
+                pool.run(reset_worker_state)
+
+
 def _get_pool(cache, world_size):
     # A hung test kills its pool; a fresh one keeps a single hang from timing out every later
-    # test in the module (that cascade cost the nightly 20+ min).
+    # test (that cascade cost the nightly 20+ min).
     return get_live_pool(cache, world_size, lambda: _make_pool(world_size))
 
 
 @pytest.fixture
 def dist_workers(_pool_cache):
-    """Pool (cached per module) with world_size=torch.cuda.device_count()."""
+    """Shared pool with world_size=torch.cuda.device_count()."""
     return _get_pool(_pool_cache, torch.cuda.device_count())
 
 
 @pytest.fixture
 def dist_workers_size_1(_pool_cache):
-    """Pool (cached per module) with world_size=1 for tests that require a single process."""
+    """Shared pool with world_size=1 for tests that require a single process."""
     return _get_pool(_pool_cache, 1)