fix(aorta): container /dev/kfd access and --override argument packing

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)...")
 

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()