Add gpu_devices policy for device-level GPU isolation

Signed-off-by: Cong Wang <cwang@multikernel.io>

Author: congwang-mk

src/sandlock/_context.py

@@ -633,6 +633,25 @@ def __enter__(self) -> "SandboxContext":
                 writable = list(self._policy.fs_writable)
                 readable = list(self._policy.fs_readable)
                 denied = list(self._policy.fs_denied)
+
+                # GPU device access
+                if self._policy.gpu_devices is not None:
+                    _gpu_rw = [
+                        "/dev/nvidia*", "/dev/nvidiactl",
+                        "/dev/nvidia-uvm", "/dev/nvidia-uvm-tools",
+                        "/dev/dri",
+                    ]
+                    _gpu_ro = [
+                        "/proc/driver/nvidia",
+                        "/sys/bus/pci/devices",
+                        "/sys/module/nvidia",
+                    ]
+                    for p in _gpu_rw:
+                        if p not in writable:
+                            writable.append(p)
+                    for p in _gpu_ro:
+                        if p not in readable:
+                            readable.append(p)
                 bind_ports = self._policy.bind_ports() or None
                 connect_ports = self._policy.connect_ports() or None
                 if (writable or readable or bind_ports or connect_ports
@@ -747,6 +766,14 @@ def __enter__(self) -> "SandboxContext":
                 if self._policy.env:
                     os.environ.update(self._policy.env)
 
+                # 8b. GPU device visibility
+                if self._policy.gpu_devices is not None:
+                    devs = self._policy.gpu_devices
+                    if len(devs) > 0:
+                        vis = ",".join(str(d) for d in devs)
+                        os.environ["CUDA_VISIBLE_DEVICES"] = vis
+                        os.environ["ROCR_VISIBLE_DEVICES"] = vis
+
                 # 9b. Disable vDSO for time virtualization
                 if (self._notif_policy is not None
                         and self._notif_policy.time_start is not None):

src/sandlock/policy.py

@@ -195,6 +195,14 @@ class Policy:
     pages, which causes nondeterministic memory layout, RSS measurements,
     and page fault timing.  Applied via prctl(PR_SET_THP_DISABLE)."""
 
+    # GPU access
+    gpu_devices: Sequence[int] | None = None
+    """GPU device indices visible to the sandbox.  When set, Landlock
+    rules are added for GPU device files (/dev/nvidia*, /dev/dri/*) and
+    driver paths (/proc/driver/nvidia, /sys/bus/pci/devices), and
+    ``CUDA_VISIBLE_DEVICES`` / ``ROCR_VISIBLE_DEVICES`` are set.
+    ``None`` = no GPU access.  ``[]`` (empty list) = all GPUs visible."""
+
     # Optional chroot
     chroot: str | None = None
     """Path to chroot into before applying other confinement."""

tests/test_policy.py

@@ -153,6 +153,20 @@ def test_env_set(self):
         assert p.env == {"FOO": "bar", "BAZ": "qux"}
 
 
+class TestGpuDevices:
+    def test_default_none(self):
+        p = Policy()
+        assert p.gpu_devices is None
+
+    def test_specific_devices(self):
+        p = Policy(gpu_devices=[0, 2])
+        assert p.gpu_devices == [0, 2]
+
+    def test_all_gpus(self):
+        p = Policy(gpu_devices=[])
+        assert p.gpu_devices == []
+
+
 class TestIpcScoping:
     def test_defaults_to_off(self):
         p = Policy()

tests/test_sandbox.py

@@ -482,3 +482,53 @@ def test_throttle_result_correct(self):
         result = Sandbox(Policy(max_cpu=50)).run(["python3", "-c", self._BURN_CODE])
         assert result.success
         assert result.stdout.strip() == b"20000000"
+
+
+class TestGpuDevices:
+    """Test GPU device isolation via CUDA_VISIBLE_DEVICES."""
+
+    _GPU_POLICY_READABLE = _PYTHON_READABLE
+
+    def test_gpu_devices_sets_env(self):
+        """gpu_devices=[0,2] sets CUDA_VISIBLE_DEVICES=0,2."""
+        code = (
+            "import os; "
+            "print(os.environ.get('CUDA_VISIBLE_DEVICES', 'UNSET'))"
+        )
+        policy = Policy(gpu_devices=[0, 2], fs_readable=self._GPU_POLICY_READABLE)
+        result = Sandbox(policy).run(["python3", "-c", code])
+        assert result.success, f"failed: {result.stderr}"
+        assert result.stdout.strip() == b"0,2"
+
+    def test_gpu_devices_sets_rocr(self):
+        """gpu_devices also sets ROCR_VISIBLE_DEVICES for AMD GPUs."""
+        code = (
+            "import os; "
+            "print(os.environ.get('ROCR_VISIBLE_DEVICES', 'UNSET'))"
+        )
+        policy = Policy(gpu_devices=[1], fs_readable=self._GPU_POLICY_READABLE)
+        result = Sandbox(policy).run(["python3", "-c", code])
+        assert result.success, f"failed: {result.stderr}"
+        assert result.stdout.strip() == b"1"
+
+    def test_gpu_devices_empty_no_env(self):
+        """gpu_devices=[] (all GPUs) does not set CUDA_VISIBLE_DEVICES."""
+        code = (
+            "import os; "
+            "print(os.environ.get('CUDA_VISIBLE_DEVICES', 'UNSET'))"
+        )
+        policy = Policy(gpu_devices=[], fs_readable=self._GPU_POLICY_READABLE)
+        result = Sandbox(policy).run(["python3", "-c", code])
+        assert result.success, f"failed: {result.stderr}"
+        assert result.stdout.strip() == b"UNSET"
+
+    def test_no_gpu_no_env(self):
+        """gpu_devices=None (default) does not set CUDA_VISIBLE_DEVICES."""
+        code = (
+            "import os; "
+            "print(os.environ.get('CUDA_VISIBLE_DEVICES', 'UNSET'))"
+        )
+        policy = Policy()
+        result = Sandbox(policy).run(["python3", "-c", code])
+        assert result.success
+        assert result.stdout.strip() == b"UNSET"