From 476c701ef73cab491059d4e8edc12f8b612d1901 Mon Sep 17 00:00:00 2001 From: speriaswamy-amd Date: Fri, 14 Aug 2026 15:25:20 -0400 Subject: [PATCH 1/2] fix(aorta): bound setup and run by timeout_seconds on daemon threads Neither setup() nor run() put a deadline on their parallel per-node work, so a single stalled node hung the entire CVS invocation with no result. This was hit live: an NCCL collective stalled on one node of a 2-node run and the process sat there indefinitely instead of reporting TIMEOUT. ThreadPoolExecutor cannot fix this on its own. Its workers are non-daemon and CPython's atexit hook joins every one of them at interpreter shutdown regardless of shutdown(wait=False), so a node stuck in a blocking Docker or SSH call would still wedge the process on the way out. _run_bounded_parallel() runs each task on a daemon thread against a shared deadline and reports per-key results, errors, and timeouts; abandoned threads cannot block exit. Also closes a setup/teardown race that the new deadline makes reachable: _setup_single_node() takes a cancel Event and, if it finishes launching its container after setup() has given up, tears that container down itself. Otherwise it would register into self._containers after teardown()'s one-time snapshot had already run, orphaning the container on the node. Co-Authored-By: Claude --- cvs/runners/aorta.py | 177 ++++++++++++------ cvs/runners/unittests/test_aorta_multinode.py | 145 +++++++++++++- 2 files changed, 266 insertions(+), 56 deletions(-) diff --git a/cvs/runners/aorta.py b/cvs/runners/aorta.py index 4dd16912d..dbb784259 100644 --- a/cvs/runners/aorta.py +++ b/cvs/runners/aorta.py @@ -14,12 +14,11 @@ import shlex import subprocess import time -from concurrent.futures import ThreadPoolExecutor, as_completed from dataclasses import dataclass, field from functools import partial from pathlib import Path -from threading import Lock -from typing import Any, Dict, List, Optional, Tuple, TYPE_CHECKING +from threading import Event, Lock, Thread +from typing import Any, Callable, Dict, List, Optional, Tuple, TYPE_CHECKING if TYPE_CHECKING: import docker @@ -502,12 +501,59 @@ def _exec_in_container_streaming( return exit_code, '\n'.join(output_lines) - def _setup_single_node(self, node: str) -> Tuple[str, bool, Optional[str]]: + @staticmethod + def _run_bounded_parallel( + tasks: Dict[Any, Callable[[], Any]], timeout_seconds: float + ) -> Tuple[Dict[Any, Any], Dict[Any, Exception], List[Any]]: + """ + Run each zero-arg callable in ``tasks`` in its own thread, bounded by a + single deadline shared across all of them. + + ``ThreadPoolExecutor`` worker threads are non-daemon; CPython's atexit + hook joins every such thread at interpreter shutdown regardless of + ``executor.shutdown(wait=False)``, so a node genuinely stuck in a + blocking Docker/SSH call would still hang the whole process at exit, + not just this method. Daemon threads are exempt from that join, so a + stuck task here can never block CVS from exiting. + + Returns ``(results, errors, timed_out_keys)`` keyed by ``tasks``' keys. + A key lands in exactly one of ``results``, ``errors``, or + ``timed_out_keys``. + """ + results: Dict[Any, Any] = {} + errors: Dict[Any, Exception] = {} + + def _worker(key: Any, fn: Callable[[], Any]) -> None: + try: + results[key] = fn() + except Exception as e: # noqa: BLE001 - reported back per-key, not swallowed + errors[key] = e + + threads = {key: Thread(target=_worker, args=(key, fn), daemon=True) for key, fn in tasks.items()} + for t in threads.values(): + t.start() + + deadline = time.monotonic() + timeout_seconds + timed_out = [] + for key, t in threads.items(): + t.join(timeout=max(0.0, deadline - time.monotonic())) + if t.is_alive(): + timed_out.append(key) + + return results, errors, timed_out + + def _setup_single_node(self, node: str, cancel_event: Event) -> Tuple[str, bool, Optional[str]]: """ Set up a single node (thread-safe helper for parallel deployment). Args: node: Hostname or IP of the node + cancel_event: set by ``setup()`` once its overall deadline has + passed. A node that finishes launching its container after + that point tears the container down itself instead of + registering it into ``self._containers`` — ``teardown()``'s + single unsynchronized pass over that dict already ran and + won't see it otherwise, orphaning the container. Returns: Tuple of (node, success, error_message) @@ -529,6 +575,15 @@ def _setup_single_node(self, node: str) -> Tuple[str, bool, Optional[str]]: # Launch container container = self._launch_container(client, node) + if cancel_event.is_set(): + log.warning(f"Setup on {node} finished after the deadline; cleaning up its container") + try: + container.stop(timeout=30) + container.remove(force=True) + except Exception as e: + log.warning(f"Error removing late container on {node}: {e}") + return (node, False, f"Setup timed out after {self.config.timeout_seconds}s") + # Thread-safe update of shared state with self._lock: self._containers[node] = container @@ -599,26 +654,31 @@ def setup(self) -> bool: log.info(f"Setting up {num_nodes} node(s) in parallel...") - # Use ThreadPoolExecutor for parallel deployment - # Max workers = number of nodes (each node gets its own thread) - with ThreadPoolExecutor(max_workers=num_nodes) as executor: - # Submit all setup tasks - futures = {executor.submit(self._setup_single_node, node): node for node in nodes} - - # Collect results as they complete - failed_nodes = [] - for future in as_completed(futures): - node = futures[future] - try: - node_name, success, error_msg = future.result() - if not success: - log.error(f"Setup failed on {node_name}: {error_msg}") - failed_nodes.append((node_name, error_msg)) - else: - log.info(f"Setup completed on {node_name}") - except Exception as e: - log.exception(f"Unexpected error setting up {node}: {e}") - failed_nodes.append((node, str(e))) + # Bounded by timeout_seconds so a stuck image pull or RCCL build cannot + # hang setup() forever; a still-stuck node's thread is abandoned (daemon) + # rather than joined, so it can't hang CVS itself either. + cancel_event = Event() + tasks = {node: partial(self._setup_single_node, node, cancel_event) for node in nodes} + results, errors, timed_out = self._run_bounded_parallel(tasks, self.config.timeout_seconds) + # Deadline has now passed for anyone still running; tell them to clean + # up their own container instead of registering it after the fact. + cancel_event.set() + + failed_nodes = [] + for node in nodes: + if node in timed_out: + log.error(f"Setup timed out on {node} after {self.config.timeout_seconds}s") + failed_nodes.append((node, f"Timed out after {self.config.timeout_seconds}s")) + elif node in errors: + log.exception(f"Unexpected error setting up {node}: {errors[node]}") + failed_nodes.append((node, str(errors[node]))) + else: + node_name, success, error_msg = results[node] + if not success: + log.error(f"Setup failed on {node_name}: {error_msg}") + failed_nodes.append((node_name, error_msg)) + else: + log.info(f"Setup completed on {node_name}") if failed_nodes: log.error(f"Setup failed on {len(failed_nodes)}/{num_nodes} nodes:") @@ -906,43 +966,50 @@ def run(self, **kwargs) -> RunResult: f"master={master_addr}:{master_port}" ) - with ThreadPoolExecutor(max_workers=nnodes) as executor: - futures = {} - for rank, node in enumerate(nodes): - cmd = self._build_torchrun_command( - node_rank=rank, - nnodes=nnodes, - master_addr=master_addr, - master_port=master_port, - nproc_per_node=nproc_per_node, - ) - future = executor.submit( - partial( - self._run_single_node, - node=node, - node_rank=rank, - launch_cmd=cmd, - env=base_env, - ) + tasks = {} + for rank, node in enumerate(nodes): + cmd = self._build_torchrun_command( + node_rank=rank, + nnodes=nnodes, + master_addr=master_addr, + master_port=master_port, + nproc_per_node=nproc_per_node, + ) + tasks[(rank, node)] = partial( + self._run_single_node, + node=node, + node_rank=rank, + launch_cmd=cmd, + env=base_env, + ) + + # Bounded by timeout_seconds so a stalled NCCL collective (or any + # other indefinite hang inside the container) cannot hang run() + # forever with no CVS-level report; a still-stuck node's thread is + # abandoned (daemon) rather than joined, so it can't hang CVS itself. + results, errors, not_done = self._run_bounded_parallel(tasks, self.config.timeout_seconds) + + for rank, node in tasks: + if (rank, node) in not_done: + log.error(f"Node {node} (rank {rank}) timed out after {self.config.timeout_seconds}s") + stdout_dict[node] = ( + f"Timed out after {self.config.timeout_seconds}s waiting for node to complete" ) - futures[future] = (rank, node) - - for future in as_completed(futures): - rank, node = futures[future] - try: - n, ec, out = future.result() - stdout_dict[n] = out - exit_codes[n] = ec - except Exception as e: - log.exception(f"Node {node} (rank {rank}) raised: {e}") - stdout_dict[node] = str(e) - exit_codes[node] = -1 + exit_codes[node] = -1 + elif (rank, node) in errors: + log.exception(f"Node {node} (rank {rank}) raised: {errors[(rank, node)]}") + stdout_dict[node] = str(errors[(rank, node)]) + exit_codes[node] = -1 + else: + n, ec, out = results[(rank, node)] + stdout_dict[n] = out + exit_codes[n] = ec failed = {n: c for n, c in exit_codes.items() if c != 0} if failed: log.error(f"Disaggregated run failed on {len(failed)}/{nnodes} nodes: {failed}") return RunResult( - status=RunStatus.FAILED, + status=RunStatus.TIMEOUT if not_done else RunStatus.FAILED, start_time=start_time, end_time=time.time(), stdout=stdout_dict, diff --git a/cvs/runners/unittests/test_aorta_multinode.py b/cvs/runners/unittests/test_aorta_multinode.py index e8cc7e411..e267d28ef 100644 --- a/cvs/runners/unittests/test_aorta_multinode.py +++ b/cvs/runners/unittests/test_aorta_multinode.py @@ -12,11 +12,14 @@ import subprocess import tempfile +import threading +import time import unittest from pathlib import Path -from unittest.mock import patch +from unittest.mock import Mock, patch import cvs.runners.aorta as aorta_mod +from cvs.runners._base_runner import RunStatus from cvs.runners.aorta import ( AortaConfig, AortaDockerConfig, @@ -268,5 +271,145 @@ def test_script_mode_does_not_require_train_script(self): ) +class TestRunBoundedParallel(unittest.TestCase): + def test_all_tasks_succeed(self): + tasks = {"a": lambda: 1, "b": lambda: 2} + results, errors, timed_out = AortaRunner._run_bounded_parallel(tasks, timeout_seconds=5) + self.assertEqual(results, {"a": 1, "b": 2}) + self.assertEqual(errors, {}) + self.assertEqual(timed_out, []) + + def test_task_exception_is_captured_per_key(self): + def boom(): + raise ValueError("bad") + + tasks = {"good": lambda: "ok", "bad": boom} + results, errors, timed_out = AortaRunner._run_bounded_parallel(tasks, timeout_seconds=5) + self.assertEqual(results, {"good": "ok"}) + self.assertIsInstance(errors["bad"], ValueError) + self.assertEqual(timed_out, []) + + def test_hung_task_times_out_without_blocking_caller(self): + never_set = threading.Event() + + def hang(): + never_set.wait() + return "unreachable" + + tasks = {"fast": lambda: "ok", "stuck": hang} + start = time.time() + results, errors, timed_out = AortaRunner._run_bounded_parallel(tasks, timeout_seconds=0.05) + elapsed = time.time() - start + + self.assertLess(elapsed, 0.5) + self.assertEqual(results, {"fast": "ok"}) + self.assertEqual(errors, {}) + self.assertEqual(timed_out, ["stuck"]) + never_set.set() + + def test_hung_task_runs_on_daemon_thread(self): + # Proves a stuck node cannot hang the whole process at interpreter exit + # (CPython's atexit joins every non-daemon thread regardless of any + # shutdown(wait=False) call the caller might make on an executor). + captured = [] + never_set = threading.Event() + + def hang(): + captured.append(threading.current_thread()) + never_set.wait() + + AortaRunner._run_bounded_parallel({"stuck": hang}, timeout_seconds=0.05) + + self.assertEqual(len(captured), 1) + self.assertTrue(captured[0].daemon) + self.assertTrue(captured[0].is_alive()) + never_set.set() + captured[0].join(timeout=1) + + +class TestRunMultiNodeTimeout(unittest.TestCase): + def test_hung_node_times_out_without_blocking_run(self): + r = _make_runner(nodes=["10.0.0.1", "10.0.0.2"], aorta_path="/tmp/aorta") + r.config.timeout_seconds = 0.05 + + def fake_run_single_node(*, node, node_rank, launch_cmd, env): + if node == "10.0.0.2": + time.sleep(0.3) + return (node, 0, "ok") + + with ( + patch.object(r, "_run_single_node", side_effect=fake_run_single_node), + patch.object(r, "_pick_master_port", return_value=29500), + ): + start = time.time() + result = r.run() + elapsed = time.time() - start + + self.assertLess(elapsed, 0.25) + self.assertEqual(result.status, RunStatus.TIMEOUT) + self.assertEqual(result.exit_codes["10.0.0.1"], 0) + self.assertEqual(result.exit_codes["10.0.0.2"], -1) + self.assertIn("Timed out", result.stdout["10.0.0.2"]) + + +class TestSetupTimeout(unittest.TestCase): + def test_hung_node_times_out_without_blocking_setup(self): + with tempfile.TemporaryDirectory() as tmp: + r = _make_runner(nodes=["10.0.0.1", "10.0.0.2"], aorta_path=tmp) + r.config.timeout_seconds = 0.05 + + def fake_setup_single_node(node, cancel_event): + if node == "10.0.0.2": + time.sleep(0.3) + return (node, True, None) + + with patch.object(r, "_setup_single_node", side_effect=fake_setup_single_node): + start = time.time() + ok = r.setup() + elapsed = time.time() - start + + self.assertFalse(ok) + self.assertLess(elapsed, 0.25) + + +class TestSetupSingleNodeCancelledLate(unittest.TestCase): + def test_container_launched_after_cancel_is_torn_down_not_registered(self): + r = _make_runner(nodes=["10.0.0.1"], aorta_path="/tmp/aorta") + fake_container = Mock() + cancel_event = threading.Event() + cancel_event.set() + + with ( + patch.object(r, "_connect_docker", return_value=Mock()), + patch.object(r, "_cleanup_existing_containers"), + patch.object(r, "_launch_container", return_value=fake_container), + ): + node, success, error = r._setup_single_node("10.0.0.1", cancel_event) + + self.assertFalse(success) + self.assertIn("timed out", error.lower()) + self.assertNotIn("10.0.0.1", r._containers) + fake_container.stop.assert_called_once() + fake_container.remove.assert_called_once() + + def test_container_launched_before_cancel_is_registered_normally(self): + r = _make_runner(nodes=["10.0.0.1"], aorta_path="/tmp/aorta") + r.config.skip_rccl_build = True + fake_container = Mock() + cancel_event = threading.Event() + + with ( + patch.object(r, "_connect_docker", return_value=Mock()), + patch.object(r, "_cleanup_existing_containers"), + patch.object(r, "_launch_container", return_value=fake_container), + ): + node, success, error = r._setup_single_node("10.0.0.1", cancel_event) + + self.assertTrue(success) + self.assertIsNone(error) + self.assertIs(r._containers["10.0.0.1"], fake_container) + fake_container.stop.assert_not_called() + + if __name__ == "__main__": unittest.main() From 9faf30e592425df84b9c9c74f03c660f048812ea Mon Sep 17 00:00:00 2001 From: speriaswamy-amd Date: Thu, 10 Sep 2026 17:44:36 -0400 Subject: [PATCH 2/2] fix(aorta): close setup/teardown race on container registration _setup_single_node() checked cancel_event and registered the launched container into self._containers as two separate, unsynchronized steps. teardown() could take its unsynchronized snapshot of that dict in the gap between them, orphaning the container, or crash outright with "dictionary changed size during iteration" if a registration landed mid-iteration. Make the check-and-register atomic under self._lock, and have teardown() snapshot-and-clear under the same lock while setting a _teardown_started flag a straggling setup thread can observe. --- cvs/runners/aorta.py | 49 ++++++++++++++----- cvs/runners/unittests/test_aorta_multinode.py | 41 ++++++++++++++++ 2 files changed, 78 insertions(+), 12 deletions(-) diff --git a/cvs/runners/aorta.py b/cvs/runners/aorta.py index dbb784259..691c7230a 100644 --- a/cvs/runners/aorta.py +++ b/cvs/runners/aorta.py @@ -200,7 +200,11 @@ def __init__(self, config: AortaConfig): # Thread-safe storage for parallel deployment self._docker_clients: Dict[str, docker.DockerClient] = {} self._containers: Dict[str, Container] = {} - self._lock = Lock() # Protects _docker_clients and _containers + self._lock = Lock() # Protects _docker_clients, _containers, and _teardown_started + # Set under _lock once teardown() has taken its snapshot of _containers, so a + # straggling setup thread that registers after that point notices and cleans up + # after itself instead of leaking a container teardown() will never see again. + self._teardown_started = False def validate_config(self) -> List[str]: """Validate Aorta-specific configuration.""" @@ -551,9 +555,12 @@ def _setup_single_node(self, node: str, cancel_event: Event) -> Tuple[str, bool, cancel_event: set by ``setup()`` once its overall deadline has passed. A node that finishes launching its container after that point tears the container down itself instead of - registering it into ``self._containers`` — ``teardown()``'s - single unsynchronized pass over that dict already ran and - won't see it otherwise, orphaning the container. + registering it into ``self._containers`` — ``teardown()`` + may already have taken its snapshot of that dict and won't + see it otherwise, orphaning the container. The check and the + registration happen atomically under ``self._lock`` (see + ``self._teardown_started``) so there is no window between + them for ``teardown()`` to race past. Returns: Tuple of (node, success, error_message) @@ -575,7 +582,16 @@ def _setup_single_node(self, node: str, cancel_event: Event) -> Tuple[str, bool, # Launch container container = self._launch_container(client, node) - if cancel_event.is_set(): + # Check-and-register must be atomic: if this happened as two separate + # steps, teardown() could take its snapshot of self._containers in the + # gap between the check and the registration and never see this + # container, orphaning it. + with self._lock: + too_late = cancel_event.is_set() or self._teardown_started + if not too_late: + self._containers[node] = container + + if too_late: log.warning(f"Setup on {node} finished after the deadline; cleaning up its container") try: container.stop(timeout=30) @@ -584,10 +600,6 @@ def _setup_single_node(self, node: str, cancel_event: Event) -> Tuple[str, bool, log.warning(f"Error removing late container on {node}: {e}") return (node, False, f"Setup timed out after {self.config.timeout_seconds}s") - # Thread-safe update of shared state - with self._lock: - self._containers[node] = container - # Build RCCL if not skipping if not self.config.skip_rccl_build: log.info(f"Building RCCL on {node}...") @@ -641,6 +653,9 @@ def setup(self) -> bool: This significantly reduces setup time for multi-node clusters. """ + with self._lock: + self._teardown_started = False + if not self.config.aorta_path.exists() and not self._ensure_aorta_repo(): log.error("Aorta path does not exist and auto-clone failed or is disabled") return False @@ -1290,7 +1305,19 @@ def teardown(self) -> bool: success = True - for node, container in self._containers.items(): + # Snapshot-and-clear under the same lock _setup_single_node uses to register + # containers, so a straggling setup thread deterministically either lands in + # this snapshot (and gets torn down below) or observes _teardown_started and + # cleans up after itself -- never both, and never neither. Iterating + # self._containers directly here (unsynchronized) would also risk a + # "dictionary changed size during iteration" crash if a setup thread inserted + # into it concurrently. + with self._lock: + self._teardown_started = True + containers = dict(self._containers) + self._containers.clear() + + for node, container in containers.items(): try: log.info(f"Stopping container on {node}...") container.stop(timeout=30) @@ -1328,8 +1355,6 @@ def teardown(self) -> bool: log.warning(f"Error removing container on {node}: {e}") success = False - self._containers.clear() - # Close Docker clients - suppress BrokenPipeError during SSH cleanup for node, client in self._docker_clients.items(): try: diff --git a/cvs/runners/unittests/test_aorta_multinode.py b/cvs/runners/unittests/test_aorta_multinode.py index e267d28ef..f85371038 100644 --- a/cvs/runners/unittests/test_aorta_multinode.py +++ b/cvs/runners/unittests/test_aorta_multinode.py @@ -410,6 +410,47 @@ def test_container_launched_before_cancel_is_registered_normally(self): self.assertIs(r._containers["10.0.0.1"], fake_container) fake_container.stop.assert_not_called() + def test_registration_after_teardown_started_is_torn_down_not_registered(self): + # A straggling setup thread can reach the registration point after + # teardown() has already taken its snapshot of self._containers, even + # though cancel_event was never set (e.g. setup() returned True/False + # for unrelated reasons and execute() moved straight to teardown()). + # It must notice via self._teardown_started and clean up after itself + # instead of silently registering into a dict teardown() will never + # look at again. + r = _make_runner(nodes=["10.0.0.1"], aorta_path="/tmp/aorta") + fake_container = Mock() + cancel_event = threading.Event() + r._teardown_started = True + + with ( + patch.object(r, "_connect_docker", return_value=Mock()), + patch.object(r, "_cleanup_existing_containers"), + patch.object(r, "_launch_container", return_value=fake_container), + ): + node, success, error = r._setup_single_node("10.0.0.1", cancel_event) + + self.assertFalse(success) + self.assertIn("timed out", error.lower()) + self.assertNotIn("10.0.0.1", r._containers) + fake_container.stop.assert_called_once() + fake_container.remove.assert_called_once() + + def test_teardown_snapshots_containers_so_concurrent_registration_does_not_crash(self): + # teardown() must iterate a snapshot, not the live self._containers dict. + # A straggling setup thread can insert into self._containers while + # teardown() is mid-iteration; if teardown() iterated the live dict, + # that mutation would raise "dictionary changed size during iteration". + r = _make_runner(nodes=["10.0.0.1"], aorta_path="/tmp/aorta") + existing = Mock() + existing.stop.side_effect = lambda *a, **k: r._containers.__setitem__("10.0.0.2", Mock()) + r._containers["10.0.0.1"] = existing + + with patch.object(r, "_get_remote_uid_gid", return_value=None): + r.teardown() # must not raise + + self.assertTrue(r._teardown_started) + if __name__ == "__main__": unittest.main()