diff --git a/cvs/runners/aorta.py b/cvs/runners/aorta.py index 4dd16912d..691c7230a 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 @@ -201,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.""" @@ -502,12 +505,62 @@ 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()`` + 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) @@ -529,9 +582,23 @@ def _setup_single_node(self, node: str) -> Tuple[str, bool, Optional[str]]: # Launch container container = self._launch_container(client, node) - # Thread-safe update of shared state + # 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: - self._containers[node] = container + 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) + 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") # Build RCCL if not skipping if not self.config.skip_rccl_build: @@ -586,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 @@ -599,26 +669,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 +981,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, @@ -1223,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) @@ -1261,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 e8cc7e411..f85371038 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,186 @@ 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() + + 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()