diff --git a/cvs/runners/aorta.py b/cvs/runners/aorta.py index a0b3afea9..0d9bba229 100644 --- a/cvs/runners/aorta.py +++ b/cvs/runners/aorta.py @@ -11,6 +11,7 @@ from __future__ import annotations import logging +import shlex import subprocess import time from concurrent.futures import ThreadPoolExecutor, as_completed @@ -272,6 +273,43 @@ def _get_remote_uid_gid(self, node: str) -> Optional[Tuple[int, int]]: log.debug(f"Could not get remote UID/GID for {node}: {e}") return None + def _resolve_group_add(self, node: str) -> List[str]: + """ + Resolve the container's group_add list for GPU device access on `node`. + + "video" exists on essentially every distro. "render" does not (e.g. some + minimal/older images), so probe for it over SSH before requesting it. Docker + resolves a group_add name against the *container image's* /etc/group, not the + host's, so a name found on the host (e.g. "render") can fail to resolve inside + the image even though the host device node needs that host GID. Pass the + numeric GID instead so it's applied directly, independent of the image's group + database. + """ + groups = ["video"] + try: + result = subprocess.run( + [ + "ssh", + "-o", + "BatchMode=yes", + "-o", + "ConnectTimeout=10", + f"{self.config.username}@{node}", + "getent group render", + ], + capture_output=True, + text=True, + timeout=15, + ) + if result.returncode == 0 and result.stdout.strip(): + gid = result.stdout.strip().split(":")[2] + groups.append(gid) + else: + log.debug(f"No 'render' group on {node}; launching with group_add={groups}") + except (subprocess.TimeoutExpired, FileNotFoundError, IndexError) as e: + log.debug(f"Could not check for 'render' group on {node}: {e}") + return groups + def _launch_container(self, client: docker.DockerClient, node: str) -> Container: """ Launch Aorta container on a node. @@ -308,7 +346,8 @@ def _launch_container(self, client: docker.DockerClient, node: str) -> Container volumes=volumes, devices=devices, working_dir=self.config.container_mount_path, - group_add=["video"], + user="root", + group_add=self._resolve_group_add(node), cap_add=["SYS_PTRACE"], security_opt=["seccomp=unconfined"], ulimits=[ @@ -553,6 +592,44 @@ def setup(self) -> bool: log.info(f"All {num_nodes} node(s) set up successfully") return True + def _build_base_env(self) -> Dict[str, str]: + """Build the environment dict exported into the container before launch.""" + env = self.config.environment.to_dict() + + rccl_path = self.config.rccl.build_path + env["LD_LIBRARY_PATH"] = ( + f"{rccl_path}/build/release/:/opt/rocm/lib:/opt/rocm/lib64:" + f"/opt/openmpi/lib:/opt/rccl-tests/build:$LD_LIBRARY_PATH" + ) + env["rccl_path"] = rccl_path + + if self.config.training_overrides: + # aorta train.py exposes `--override` with `nargs="*"`; multiple + # `--override` groups collapse to the last group's values. Emit a + # single `--override` followed by all key=value tokens so that + # downstream legacy launch scripts also forward them correctly. + # No shell re-parses this string, so embedded quote characters + # would reach argparse literally (e.g. training.max_steps="15" + # instead of training.max_steps=15) -- emit bare key=value tokens. + tokens = " ".join(f"{key}={value}" for key, value in self.config.training_overrides.items()) + env["AORTA_OVERRIDE_ARGS"] = f"--override {tokens}" + + return env + + def _build_experiment_command(self) -> str: + """Build the ``bash experiment_script config_path [--override ...]`` command.""" + config_path = f"{self.config.container_mount_path}/{self.config.base_config}" + exp_cmd = f"bash {self.config.container_mount_path}/{self.config.experiment_script} {config_path}" + if self.config.training_overrides: + # AORTA_OVERRIDE_ARGS (see _build_base_env) only reaches the process + # environment; nothing consumes it unless it's also appended to the + # argv the experiment script receives. + override_parts = ["--override"] + for key, value in self.config.training_overrides.items(): + override_parts.append(f"{key}={shlex.quote(str(value))}") + exp_cmd += " " + " ".join(override_parts) + return exp_cmd + def run(self, **kwargs) -> RunResult: """ Execute the Aorta benchmark. @@ -579,33 +656,13 @@ def run(self, **kwargs) -> RunResult: error_message=f"No container found for {node}", ) - # Build environment with computed values - env = self.config.environment.to_dict() - - # Add RCCL library path - rccl_path = self.config.rccl.build_path - env["LD_LIBRARY_PATH"] = ( - f"{rccl_path}/build/release/:/opt/rocm/lib:/opt/rocm/lib64:" - f"/opt/openmpi/lib:/opt/rccl-tests/build:$LD_LIBRARY_PATH" - ) - env["rccl_path"] = rccl_path - - # Build override arguments if any - override_args = "" - if self.config.training_overrides: - for key, value in self.config.training_overrides.items(): - override_args += f' --override {key}="{value}"' - - # Execute experiment script with streaming output for real-time feedback - # Note: override_args is passed via environment if the script supports it - if override_args: - env["AORTA_OVERRIDE_ARGS"] = override_args.strip() - log.info(f"Training overrides: {override_args.strip()}") + env = self._build_base_env() + if env.get("AORTA_OVERRIDE_ARGS"): + log.info(f"Training overrides: {env['AORTA_OVERRIDE_ARGS']}") # Pass the base config file to the experiment script # launch_rocm.sh expects: CONFIG=${1:-default.yaml} - config_path = f"{self.config.container_mount_path}/{self.config.base_config}" - exp_cmd = f"bash {self.config.container_mount_path}/{self.config.experiment_script} {config_path}" + exp_cmd = self._build_experiment_command() log.info(f"Running experiment: {exp_cmd}") log.info("Streaming output (this may take several minutes)...") diff --git a/cvs/runners/unittests/test_aorta.py b/cvs/runners/unittests/test_aorta.py new file mode 100644 index 000000000..cedd3b250 --- /dev/null +++ b/cvs/runners/unittests/test_aorta.py @@ -0,0 +1,156 @@ +""" +Unit tests for ``AortaRunner`` helpers that do not require a live cluster. + +The networked container/SSH paths are not exercised here; see +``cvs/tests/benchmark/test_aorta.py`` for the end-to-end pytest suite that runs +against a real cluster. + +Copyright 2025 Advanced Micro Devices, Inc. +All rights reserved. +""" + +import unittest +from pathlib import Path +from unittest.mock import Mock, patch + +import cvs.runners.aorta as aorta_mod +from cvs.runners.aorta import ( + AortaConfig, + AortaDockerConfig, + AortaEnvironment, + AortaRunner, + RcclConfig, +) + + +def _make_runner( + *, + nodes, + aorta_path, + base_config="config/distributed.yaml", + experiment_script="scripts/launch_rocm.sh", + **config_overrides, +): + cfg = AortaConfig( + nodes=list(nodes), + username="testuser", + pkey="/home/testuser/.ssh/id_rsa", + aorta_path=Path(aorta_path), + base_config=base_config, + docker=AortaDockerConfig(), + rccl=RcclConfig(), + environment=AortaEnvironment(), + build_script="scripts/launch_rocm.sh", + experiment_script=experiment_script, + gpus_per_node=8, + **config_overrides, + ) + # The runner's __init__ aborts when the docker SDK is unavailable. None of + # the helpers under test actually call into docker, so flip the module flag + # for the duration of this call. This keeps the unit tests runnable in + # minimal CI environments without the docker package. + with patch.object(aorta_mod, "DOCKER_SDK_AVAILABLE", True): + return AortaRunner(cfg) + + +class TestBuildBaseEnv(unittest.TestCase): + def test_rccl_paths_are_exported(self): + runner = _make_runner(nodes=["a"], aorta_path="/tmp/aorta") + env = runner._build_base_env() + self.assertIn("LD_LIBRARY_PATH", env) + self.assertEqual(env["rccl_path"], runner.config.rccl.build_path) + # Existing NCCL knobs should still be there. + self.assertEqual(env["NCCL_MAX_NCHANNELS"], "112") + + def test_no_override_var_without_training_overrides(self): + runner = _make_runner(nodes=["a"], aorta_path="/tmp/aorta") + self.assertNotIn("AORTA_OVERRIDE_ARGS", runner._build_base_env()) + + def test_training_overrides_become_env_var(self): + runner = _make_runner(nodes=["a"], aorta_path="/tmp/aorta") + runner.config.training_overrides = {"training.max_steps": 5} + env = runner._build_base_env() + # No shell re-parses this string, so it must not contain embedded quote + # characters -- those would reach argparse literally as part of the value. + self.assertEqual(env["AORTA_OVERRIDE_ARGS"], "--override training.max_steps=5") + + def test_multi_key_overrides_share_one_override_group(self): + # Aorta train.py uses argparse(--override, nargs="*"); multiple + # `--override` groups would silently keep only the last group's values. + # Guarantee a single group regardless of how many keys are configured. + runner = _make_runner(nodes=["a"], aorta_path="/tmp/aorta") + runner.config.training_overrides = { + "training.max_steps": 5, + "training.batch_size": 8, + "profiling.active": 3, + } + env = runner._build_base_env() + self.assertEqual( + env["AORTA_OVERRIDE_ARGS"], + "--override training.max_steps=5 training.batch_size=8 profiling.active=3", + ) + + +class TestBuildExperimentCommand(unittest.TestCase): + def test_no_override_args_without_training_overrides(self): + runner = _make_runner(nodes=["a"], aorta_path="/tmp/aorta") + cmd = runner._build_experiment_command() + self.assertEqual(cmd, "bash /mnt/scripts/launch_rocm.sh /mnt/config/distributed.yaml") + + def test_training_overrides_are_appended_to_command(self): + # Computing AORTA_OVERRIDE_ARGS is not enough -- launch_rocm.sh only + # reads argv, so the override tokens must also land on the command + # line actually executed, not just in the environment. + runner = _make_runner(nodes=["a"], aorta_path="/tmp/aorta") + runner.config.training_overrides = {"training.max_steps": 5} + cmd = runner._build_experiment_command() + self.assertTrue(cmd.endswith("--override training.max_steps=5"), cmd) + + def test_multi_key_overrides_share_one_override_group(self): + runner = _make_runner(nodes=["a"], aorta_path="/tmp/aorta") + runner.config.training_overrides = { + "training.max_steps": 5, + "training.batch_size": 8, + } + cmd = runner._build_experiment_command() + self.assertTrue(cmd.endswith("--override training.max_steps=5 training.batch_size=8"), cmd) + + def test_override_values_needing_quoting_are_shell_safe(self): + runner = _make_runner(nodes=["a"], aorta_path="/tmp/aorta") + runner.config.training_overrides = {"training.tag": "a b"} + cmd = runner._build_experiment_command() + self.assertTrue(cmd.endswith("--override training.tag='a b'"), cmd) + + +class TestLaunchContainerGpuAccess(unittest.TestCase): + def _launch(self, ssh_returncode=0, ssh_stdout="render:x:104:testuser\n"): + runner = _make_runner(nodes=["a"], aorta_path="/tmp/aorta") + client = Mock() + client.containers.run.return_value.status = "running" + ssh_result = Mock(returncode=ssh_returncode, stdout=ssh_stdout) + with ( + patch.object(aorta_mod, "docker", Mock()), + patch.object(aorta_mod.subprocess, "run", return_value=ssh_result), + ): + runner._launch_container(client, "a") + return client.containers.run.call_args.kwargs + + def test_container_runs_as_root(self): + # Without this the container inherits the image's default UID, which on + # the validation cluster could not open /dev/kfd even with --privileged. + self.assertEqual(self._launch()["user"], "root") + + def test_render_group_is_added_when_present_on_host(self): + # Docker resolves group_add names against the *container image's* /etc/group, + # not the host's, so the host's numeric GID must be passed instead of "render". + self.assertEqual(self._launch()["group_add"], ["video", "104"]) + + def test_render_group_is_skipped_when_absent_on_host(self): + # containers.run() fails outright if a requested group is missing on the + # host, so a host without a "render" group must not request it. + kwargs = self._launch(ssh_returncode=2, ssh_stdout="") + self.assertEqual(kwargs["group_add"], ["video"]) + + +if __name__ == "__main__": + unittest.main()