Merge pull request #6512 from hpcaitech/fix/ci-container-gpu-access

[ci] Preserve GPU access and isolate container scratch files

Author: yuxuan-lou

.github/scripts/ci/container_batch.py

@@ -9,8 +9,10 @@
 import os
 import re
 import signal
+import stat
 import subprocess
 import uuid
+import xml.etree.ElementTree as ET
 from pathlib import Path
 
 OWNER_LABEL = "org.colossalai.ci.owner"
@@ -63,6 +65,44 @@ def validate_scope(scope):
     return scope
 
 
+def gpu_device_paths(gpu_uuids):
+    """Resolve allocated UUIDs to character devices, including shared CUDA controls."""
+    inventory = subprocess.run(
+        # minor_number is present in the XML inventory but is not a supported
+        # --query-gpu field on all drivers (including the H20 runner's driver).
+        ["nvidia-smi", "-q", "-x"],
+        check=True,
+        capture_output=True,
+        text=True,
+        timeout=15,
+    ).stdout
+    minors = {}
+    try:
+        entries = ET.fromstring(inventory).findall("gpu")
+    except ET.ParseError as error:
+        raise RuntimeError("Invalid GPU XML inventory from nvidia-smi") from error
+    for entry in entries:
+        gpu = entry.findtext("uuid", "").strip()
+        minor = entry.findtext("minor_number", "").strip()
+        if not gpu or gpu in minors or not minor.isdecimal():
+            raise RuntimeError("Invalid GPU device inventory from nvidia-smi")
+        minors[gpu] = int(minor)
+    if any(gpu not in minors for gpu in gpu_uuids):
+        raise RuntimeError("An allocated GPU UUID is missing from the device inventory")
+    selected_minors = [minors[gpu] for gpu in gpu_uuids]
+    if len(set(selected_minors)) != len(selected_minors):
+        raise RuntimeError("Allocated GPU UUIDs resolve to duplicate device minors")
+    devices = [Path(f"/dev/nvidia{minor}") for minor in selected_minors]
+    devices += [Path("/dev/nvidiactl"), Path("/dev/nvidia-uvm"), Path("/dev/nvidia-uvm-tools")]
+    for index, device in enumerate(devices):
+        metadata = device.stat()
+        if not stat.S_ISCHR(metadata.st_mode):
+            raise RuntimeError(f"GPU device is not a character device: {device}")
+        if index < len(selected_minors) and os.minor(metadata.st_rdev) != selected_minors[index]:
+            raise RuntimeError(f"GPU device minor does not match the allocation: {device}")
+    return devices
+
+
 def create_command(name, scope, image_id, source, output, gpu_uuids, batch, env):
     validate_scope(scope)
     if not IMAGE_ID.fullmatch(image_id):
@@ -114,6 +154,10 @@ def create_command(name, scope, image_id, source, output, gpu_uuids, batch, env)
         shm,
         "--pids-limit=4096",
         "--ulimit=memlock=-1:-1",
+        # Keep disposable source, HOME and compiler caches out of overlay2.
+        # tmpfs pages count towards the existing container memory limit.
+        "--tmpfs",
+        f"/tmp:rw,exec,nosuid,nodev,size={memory},mode=1777",
         "--gpus",
         '"device=' + ",".join(gpu_uuids) + '"',
         "--mount",
@@ -123,6 +167,11 @@ def create_command(name, scope, image_id, source, output, gpu_uuids, batch, env)
         "--workdir",
         "/tmp",
     ]
+    # The legacy NVIDIA hook's cgroup permissions can disappear after a
+    # container/systemd update. Explicit devices are retained by runc. Resolve
+    # UUIDs to actual minors rather than assuming host indices equal minors.
+    for device in gpu_device_paths(gpu_uuids):
+        command.extend(["--device", f"{device}:{device}:rw"])
     container_env = {
         "NVIDIA_VISIBLE_DEVICES": ",".join(gpu_uuids),
         "NVIDIA_DRIVER_CAPABILITIES": "compute,utility",
@@ -185,6 +234,7 @@ def interrupted(signum, frame):
         raise RuntimeError(f"Container batch interrupted by signal {signum}")
 
     old_handlers = {sig: signal.signal(sig, interrupted) for sig in (signal.SIGTERM, signal.SIGINT)}
+    execution_error = None
     try:
         try:
             subprocess.run(command, check=True, capture_output=True, text=True, timeout=60)
@@ -194,13 +244,18 @@ def interrupted(signum, frame):
         client_code = stream(["docker", "start", "--attach", name], output / "pytest.log", env)
         state = json.loads(docker("inspect", "--format", "{{json .State}}", name))
         result["container_state"] = state
+        result["container_exit_code"] = state.get("ExitCode")
         if state.get("Running") or state.get("Status") != "exited":
             raise RuntimeError("Docker client exited before the container finished")
         if state.get("OOMKilled"):
             raise RuntimeError("Container exceeded its configured memory limit")
         if state.get("Error"):
             raise RuntimeError(f"Container runtime error: {state['Error']}")
         return state["ExitCode"] or client_code
+    except (OSError, RuntimeError, ValueError, subprocess.SubprocessError) as error:
+        execution_error = error
+        result["container_execution_error"] = str(error)
+        raise
     finally:
         # Also runs when create times out: the daemon might have created the
         # container despite losing the client's response. The job's always()
@@ -212,7 +267,12 @@ def interrupted(signum, frame):
             result["container_cleanup"] = "completed"
         except (OSError, RuntimeError, ValueError, subprocess.SubprocessError) as error:
             result["container_cleanup"] = f"failed: {error}"
-            raise
+            # Retain the original execution exception when both paths fail;
+            # cleanup errors are still recorded and the job cannot pass.
+            if execution_error is None:
+                raise RuntimeError(
+                    f"Container exited with code {result.get('container_exit_code', 'unknown')}; cleanup failed: {error}"
+                ) from error
         finally:
             for sig, handler in old_handlers.items():
                 signal.signal(sig, handler)

.github/scripts/ci/cuda_probe.py

@@ -0,0 +1,24 @@
+"""Check actual computation on every container-visible CUDA device."""
+
+
+def probe_devices(torch, expected):
+    if expected < 1:
+        raise RuntimeError("The CUDA probe requires at least one allocated GPU")
+    if not torch.cuda.is_available():
+        raise RuntimeError("CUDA is not available inside the container")
+    count = torch.cuda.device_count()
+    if count != expected:
+        raise RuntimeError(f"Expected {expected} CUDA device(s), found {count}")
+    names = []
+    for index in range(expected):
+        try:
+            # Enumeration can succeed before device access is lost. Allocate,
+            # launch a kernel and synchronize so that failures reach this gate.
+            value = torch.ones(4, dtype=torch.float32, device=f"cuda:{index}").sum().item()
+            torch.cuda.synchronize(index)
+            if value != 4:
+                raise RuntimeError(f"Unexpected CUDA probe result: {value}")
+            names.append(torch.cuda.get_device_name(index))
+        except Exception as error:
+            raise RuntimeError(f"CUDA compute probe failed on cuda:{index}: {error}") from error
+    return names

.github/scripts/ci/run_container_batch.sh

@@ -50,17 +50,20 @@ from pathlib import Path
 
 import torch
 
+sys.path.insert(0, str(Path(os.environ["COLOSSAL_REPO"]) / ".github/scripts/ci"))
+from cuda_probe import probe_devices
+
 assert torch.__version__.startswith("2.13."), torch.__version__
 assert torch.version.cuda == "13.0", torch.version.cuda
-assert torch.cuda.is_available(), "CUDA is not available inside the container"
-assert torch.cuda.device_count() == len(os.environ["CUDA_VISIBLE_DEVICES"].split(","))
+gpu_names = probe_devices(torch, len(os.environ["CUDA_VISIBLE_DEVICES"].split(",")))
 profile = {
     "source_sha": os.environ["CI_SOURCE_SHA"],
     "image_id": os.environ["CI_IMAGE_ID"],
     "python": sys.version,
     "torch": torch.__version__,
     "cuda": torch.version.cuda,
-    "gpu_names": [torch.cuda.get_device_name(i) for i in range(torch.cuda.device_count())],
+    "gpu_names": gpu_names,
+    "cuda_compute_probe": "passed",
 }
 print(json.dumps(profile, indent=2), flush=True)
 Path("/ci-results/runtime.json").write_text(json.dumps(profile, indent=2) + "\n")

.github/scripts/ci/test_container_batch.py

@@ -2,11 +2,13 @@
 import json
 import os
 import signal
+import stat
 import subprocess
 import tempfile
 import unittest
 from contextlib import redirect_stdout
 from pathlib import Path
+from types import SimpleNamespace
 from unittest.mock import Mock, patch
 
 import container_batch as container
@@ -29,6 +31,19 @@ def setUp(self):
         script.touch()
         self.output = self.root / "reports"
         self.output.mkdir()
+        devices = patch.object(
+            container,
+            "gpu_device_paths",
+            return_value=[
+                Path("/dev/nvidia4"),
+                Path("/dev/nvidia7"),
+                Path("/dev/nvidiactl"),
+                Path("/dev/nvidia-uvm"),
+                Path("/dev/nvidia-uvm-tools"),
+            ],
+        )
+        devices.start()
+        self.addCleanup(devices.stop)
 
     def command(self, **env):
         return container.create_command(
@@ -48,6 +63,17 @@ def test_uuid_filter_and_logical_device_mapping(self):
         self.assertIn("CUDA_VISIBLE_DEVICES=0,1", command)
         self.assertIn(f"NVIDIA_VISIBLE_DEVICES={GPU_A},{GPU_B}", command)
         self.assertEqual(command[-2:], ["5", "0,1"])
+        devices = [command[index + 1] for index, value in enumerate(command) if value == "--device"]
+        self.assertEqual(
+            devices,
+            [
+                "/dev/nvidia4:/dev/nvidia4:rw",
+                "/dev/nvidia7:/dev/nvidia7:rw",
+                "/dev/nvidiactl:/dev/nvidiactl:rw",
+                "/dev/nvidia-uvm:/dev/nvidia-uvm:rw",
+                "/dev/nvidia-uvm-tools:/dev/nvidia-uvm-tools:rw",
+            ],
+        )
 
     def test_only_source_and_reports_are_mounted(self):
         command = self.command(COLOSSAL_VENV="/host/private/venv", SECRET_TOKEN="not-for-container")
@@ -75,6 +101,7 @@ def test_resource_limits(self):
             ("--shm-size", "4g"),
         ):
             self.assertEqual(command[command.index(flag) + 1], expected)
+        self.assertEqual(command[command.index("--tmpfs") + 1], "/tmp:rw,exec,nosuid,nodev,size=32g,mode=1777")
         for env in (
             {"COLOSSAL_CI_CPUS": "0"},
             {"COLOSSAL_CI_MEMORY": "unlimited"},
@@ -208,6 +235,92 @@ def test_cleanup_failure_cannot_report_success(self):
             with self.assertRaisesRegex(RuntimeError, "daemon unavailable"):
                 self.run_batch(Mock(return_value=0))
 
+    def test_cleanup_failure_preserves_execution_error(self):
+        result = {}
+        with (
+            patch.object(container.subprocess, "run"),
+            patch.object(container, "remove_owned_container", side_effect=RuntimeError("overlay2 removal failed")),
+        ):
+            with self.assertRaisesRegex(RuntimeError, "original execution failed"):
+                container.run_container(
+                    "5",
+                    [(4, GPU_A), (7, GPU_B)],
+                    self.root,
+                    self.output,
+                    IMAGE,
+                    {"CI_CONTAINER_SCOPE": SCOPE},
+                    Mock(side_effect=RuntimeError("original execution failed")),
+                    result,
+                )
+        self.assertEqual(result["container_execution_error"], "original execution failed")
+        self.assertIn("overlay2 removal failed", result["container_cleanup"])
+
+
+class GPUDeviceTests(unittest.TestCase):
+    @staticmethod
+    def inventory(*pairs):
+        return (
+            "<nvidia_smi_log>"
+            + "".join(f"<gpu><uuid>{gpu}</uuid><minor_number>{minor}</minor_number></gpu>" for gpu, minor in pairs)
+            + "</nvidia_smi_log>"
+        )
+
+    def test_devices_are_resolved_by_uuid_and_minor_not_inventory_order(self):
+        metadata = [
+            SimpleNamespace(st_mode=stat.S_IFCHR, st_rdev=os.makedev(195, 7)),
+            SimpleNamespace(st_mode=stat.S_IFCHR, st_rdev=os.makedev(195, 4)),
+            *[SimpleNamespace(st_mode=stat.S_IFCHR, st_rdev=0) for _ in range(3)],
+        ]
+        with (
+            patch.object(container.subprocess, "run", return_value=Mock(stdout=self.inventory((GPU_B, 4), (GPU_A, 7)))),
+            patch.object(Path, "stat", side_effect=metadata),
+        ):
+            devices = container.gpu_device_paths([GPU_A, GPU_B])
+        self.assertEqual([str(path) for path in devices[:2]], ["/dev/nvidia7", "/dev/nvidia4"])
+        self.assertEqual(len(devices), 5)
+
+    def test_missing_or_ambiguous_allocations_fail_before_device_access(self):
+        for inventory in (
+            self.inventory((GPU_A, 4)),
+            self.inventory((GPU_A, 4), (GPU_A, 7)),
+            self.inventory((GPU_A, 4), (GPU_B, 4)),
+            self.inventory((GPU_A, "N/A")),
+            "malformed XML",
+        ):
+            with (
+                self.subTest(inventory=inventory),
+                patch.object(container.subprocess, "run", return_value=Mock(stdout=inventory)),
+                patch.object(Path, "stat") as metadata,
+            ):
+                with self.assertRaises(RuntimeError):
+                    container.gpu_device_paths([GPU_A, GPU_B])
+                metadata.assert_not_called()
+
+    def test_non_device_and_mismatched_minor_are_rejected(self):
+        for metadata in (
+            SimpleNamespace(st_mode=stat.S_IFREG, st_rdev=0),
+            SimpleNamespace(st_mode=stat.S_IFCHR, st_rdev=os.makedev(195, 9)),
+        ):
+            with (
+                self.subTest(metadata=metadata),
+                patch.object(container.subprocess, "run", return_value=Mock(stdout=self.inventory((GPU_A, 4)))),
+                patch.object(Path, "stat", return_value=metadata),
+            ):
+                with self.assertRaises(RuntimeError):
+                    container.gpu_device_paths([GPU_A])
+
+    def test_missing_control_device_fails_closed(self):
+        with (
+            patch.object(container.subprocess, "run", return_value=Mock(stdout=self.inventory((GPU_A, 4)))),
+            patch.object(
+                Path,
+                "stat",
+                side_effect=[SimpleNamespace(st_mode=stat.S_IFCHR, st_rdev=os.makedev(195, 4)), FileNotFoundError()],
+            ),
+        ):
+            with self.assertRaises(FileNotFoundError):
+                container.gpu_device_paths([GPU_A])
+
 
 class ContainerOwnershipTests(unittest.TestCase):
     def test_absent_container_needs_no_removal(self):

.github/scripts/ci/test_cuda_probe.py

@@ -0,0 +1,50 @@
+import unittest
+from unittest.mock import Mock, call
+
+from cuda_probe import probe_devices
+
+
+class CUDAProbeTests(unittest.TestCase):
+    def setUp(self):
+        self.torch = Mock()
+        self.torch.cuda.is_available.return_value = True
+        self.torch.cuda.device_count.return_value = 2
+        self.torch.cuda.get_device_name.side_effect = lambda index: f"GPU {index}"
+        self.torch.ones.return_value.sum.return_value.item.return_value = 4
+
+    def test_computes_and_synchronizes_every_allocated_device(self):
+        self.assertEqual(probe_devices(self.torch, 2), ["GPU 0", "GPU 1"])
+        self.assertEqual(
+            self.torch.ones.call_args_list,
+            [
+                call(4, dtype=self.torch.float32, device="cuda:0"),
+                call(4, dtype=self.torch.float32, device="cuda:1"),
+            ],
+        )
+        self.assertEqual(self.torch.cuda.synchronize.call_args_list, [call(0), call(1)])
+
+    def test_enumeration_success_cannot_hide_allocation_failure(self):
+        self.torch.ones.side_effect = [self.torch.ones.return_value, RuntimeError("device access denied")]
+        with self.assertRaisesRegex(RuntimeError, "cuda:1: device access denied"):
+            probe_devices(self.torch, 2)
+        self.torch.cuda.synchronize.assert_called_once_with(0)
+
+    def test_asynchronous_cuda_failure_is_not_a_pass(self):
+        self.torch.cuda.synchronize.side_effect = RuntimeError("CUDA device unavailable")
+        with self.assertRaisesRegex(RuntimeError, "cuda:0: CUDA device unavailable"):
+            probe_devices(self.torch, 2)
+
+    def test_unavailable_or_wrong_device_count_fails_before_computation(self):
+        self.torch.cuda.is_available.return_value = False
+        with self.assertRaisesRegex(RuntimeError, "CUDA is not available"):
+            probe_devices(self.torch, 2)
+        self.torch.cuda.is_available.return_value = True
+        with self.assertRaisesRegex(RuntimeError, "Expected 1 CUDA device"):
+            probe_devices(self.torch, 1)
+        with self.assertRaisesRegex(RuntimeError, "at least one"):
+            probe_devices(self.torch, 0)
+        self.torch.ones.assert_not_called()
+
+
+if __name__ == "__main__":
+    unittest.main()