From e93ddb03472041d7a4a070b8de7032d583e175f3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=EB=B0=B1=EC=A7=80=EB=AA=85?= Date: Sun, 6 Sep 2026 16:15:47 +0900 Subject: [PATCH 1/3] playbooks: run nvidia-smi tasks outside the ssh cgroup so they survive the login GPU guard MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit With slurm_login_on_compute enabled, login-compute-setup restricts the ssh service's cgroup to /dev/nvidiactl so interactive users cannot reach GPUs outside a Slurm job. Ansible arrives over the same ssh service, so every task that runs nvidia-smi afterwards is restricted too (the filter is per-cgroup, become does not help): nvidia-smi -> rc=6, "No devices were found" nvidia-smi -acp UNRESTRICTED -> rc=6 nvidia-smi --query-gpu=mig.mode -> rc=0 with "No devices were found" That breaks the second run slurm-single-node.md asks for ("reboot manually when prompted and then run Ansible again"): gpu-clocks.yml and the nvidia-cuda.yml driver test fail, and nvidia-mig.yml treats the "No devices" text as a MIG-capable GPU because it only filtered 'N/A'. Run those nvidia-smi invocations through `systemd-run --wait --pipe --collect --quiet`. The command executes in a transient unit under system.slice, outside the ssh cgroup, so the check is still real instead of being skipped. Make the MIG probe accept only Enabled/Disabled. nvidia-driver.yml has the same test; it is left alone here to avoid conflicting with #1396 and can take the same wrapper after that lands. Verified on a DGX B300 (DGX OS 7.5.0, driver 580.126.20) with the guard active: plain nvidia-smi in the ssh session returns rc=6, while `sudo systemd-run --wait --pipe --collect --quiet nvidia-smi -L` lists all eight GPUs with rc=0. Signed-off-by: 백지명 --- playbooks/nvidia-software/nvidia-cuda.yml | 4 +++- playbooks/nvidia-software/nvidia-mig.yml | 14 +++++++++++++- playbooks/utilities/gpu-clocks.yml | 6 +++++- 3 files changed, 21 insertions(+), 3 deletions(-) diff --git a/playbooks/nvidia-software/nvidia-cuda.yml b/playbooks/nvidia-software/nvidia-cuda.yml index 7945f671c..b023b1e65 100644 --- a/playbooks/nvidia-software/nvidia-cuda.yml +++ b/playbooks/nvidia-software/nvidia-cuda.yml @@ -31,8 +31,10 @@ include_role: name: nvidia_cuda + # Run in a transient unit so the check also works on nodes where the + # Slurm login setup hides GPUs from ssh sessions (see login-compute-setup). - name: test nvidia-smi - command: nvidia-smi + command: systemd-run --wait --pipe --collect --quiet nvidia-smi changed_when: false when: - ansible_local['gpus']['count'] diff --git a/playbooks/nvidia-software/nvidia-mig.yml b/playbooks/nvidia-software/nvidia-mig.yml index 02bff055b..14a32089f 100644 --- a/playbooks/nvidia-software/nvidia-mig.yml +++ b/playbooks/nvidia-software/nvidia-mig.yml @@ -9,10 +9,22 @@ vars: mig_manager_reboot_timeout: 900 tasks: + # Run nvidia-smi in a transient unit so the probe also works on nodes where + # the Slurm login setup hides GPUs from ssh sessions, and only accept a real + # MIG mode: with hidden GPUs nvidia-smi prints "No devices were found", + # which the previous 'grep -v N/A' let through as a match. - name: check for MIG capable devices - shell: nvidia-smi --query-gpu=mig.mode.current --format=csv,noheader | grep -v 'N/A' + shell: >- + set -o pipefail && + systemd-run --wait --pipe --collect --quiet + nvidia-smi --query-gpu=mig.mode.current --format=csv,noheader + | grep -E '^(Enabled|Disabled)$' + args: + executable: /bin/bash + become: true register: has_mig failed_when: false + changed_when: false - name: Install MIG Manager include_role: diff --git a/playbooks/utilities/gpu-clocks.yml b/playbooks/utilities/gpu-clocks.yml index ef4d69732..ec5cafa13 100644 --- a/playbooks/utilities/gpu-clocks.yml +++ b/playbooks/utilities/gpu-clocks.yml @@ -6,8 +6,12 @@ include_role: name: facts + # Run nvidia-smi in a transient unit: on nodes where the Slurm login + # setup hides GPUs from ssh sessions (DeviceAllow on the ssh service), + # commands run over the Ansible connection cannot see the devices. + # Note: -acp is deprecated in current drivers and treated as a no-op. - name: set GPU clocks permissions - command: nvidia-smi -acp UNRESTRICTED + command: systemd-run --wait --pipe --collect --quiet nvidia-smi -acp UNRESTRICTED changed_when: false when: - ansible_local['gpus']['count'] From bc4a353ae09d854acab95c1a44e22b981c340591 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=EB=B0=B1=EC=A7=80=EB=AA=85?= Date: Sun, 6 Sep 2026 16:48:58 +0900 Subject: [PATCH 2/3] nvidia-mig-manager, nvidia-set-gpu-clocks: cover the remaining nvidia-smi call sites MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The nvidia-mig-manager role repeats the MIG probe that nvidia-mig.yml runs, with the same 'grep -v N/A' filter, and utilities/nvidia-set-gpu-clocks.yml locks/resets clocks with plain nvidia-smi. Both run over the Ansible connection and hit the same ssh-cgroup device guard on slurm_login_on_compute nodes. Apply the same systemd-run wrapper, accept only Enabled/Disabled from the MIG probe, and make that probe become root explicitly since the rest of the role already requires it. Signed-off-by: 백지명 --- playbooks/utilities/nvidia-set-gpu-clocks.yml | 6 ++++-- roles/nvidia-mig-manager/tasks/main.yml | 12 ++++++++++-- 2 files changed, 14 insertions(+), 4 deletions(-) diff --git a/playbooks/utilities/nvidia-set-gpu-clocks.yml b/playbooks/utilities/nvidia-set-gpu-clocks.yml index 4c734726d..586f2a263 100644 --- a/playbooks/utilities/nvidia-set-gpu-clocks.yml +++ b/playbooks/utilities/nvidia-set-gpu-clocks.yml @@ -4,10 +4,12 @@ gather_facts: no become: yes tasks: + # nvidia-smi runs in a transient unit so this also works on nodes where the + # Slurm login setup hides GPUs from ssh sessions (see login-compute-setup). - name: set the gpu clock to a specified amount - shell: "nvidia-smi -lgc {{ gpu_clock_lock }}" + command: "systemd-run --wait --pipe --collect --quiet nvidia-smi -lgc {{ gpu_clock_lock }}" when: not gpu_clock_reset - name: reset the gpu clock to the default - shell: "nvidia-smi -rgc" + command: systemd-run --wait --pipe --collect --quiet nvidia-smi -rgc when: gpu_clock_reset diff --git a/roles/nvidia-mig-manager/tasks/main.yml b/roles/nvidia-mig-manager/tasks/main.yml index be1b2376e..1e171737f 100644 --- a/roles/nvidia-mig-manager/tasks/main.yml +++ b/roles/nvidia-mig-manager/tasks/main.yml @@ -1,11 +1,19 @@ --- # Install the NVIDIA MIG Manager tooling on all MIG-capable nodes -# Check node state +# Check node state. nvidia-smi runs in a transient unit so the probe also works +# on nodes where the Slurm login setup hides GPUs from ssh sessions, and only a +# real MIG mode counts: with hidden GPUs nvidia-smi prints "No devices were +# found", which 'grep -v N/A' would accept as a match. - name: check for MIG capable devices - shell: set -o pipefail && nvidia-smi --query-gpu=mig.mode.current --format=csv,noheader | grep -v 'N/A' + shell: >- + set -o pipefail && + systemd-run --wait --pipe --collect --quiet + nvidia-smi --query-gpu=mig.mode.current --format=csv,noheader + | grep -E '^(Enabled|Disabled)$' args: executable: "/bin/bash" + become: true register: has_mig failed_when: false changed_when: false From 32b04b74cd753b35f5bb63824ee24bab4d92fb87 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=EB=B0=B1=EC=A7=80=EB=AA=85?= Date: Fri, 11 Sep 2026 09:21:17 +0900 Subject: [PATCH 3/3] nvidia-mig, nvidia-driver: finish the escape from the ssh GPU cgroup MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Address the review on running GPU tasks outside the login device guard. The MIG probe was the only part of the MIG path that escaped the cgroup. `nvidia-mig-parted apply` and `assert` drive the very same GPUs, so once the repaired probe succeeds they still ran in the restricted ssh context and failed there. They now run in a transient unit as well, and the play runs privileged, which the standalone `nvidia-mig-parted` invocations needed anyway. `nvidia-driver.yml`'s `test nvidia-smi` gets the same treatment: that was the remaining driver check, and #1396 closed without landing it. The probe could not tell a GPU that does not support MIG from a query that never ran. Both end up with a non-zero status, and with the query piped into the match, `pipefail` reports the match's status in either case, so an operational failure -- a driver that does not answer, a transient service that could not start -- was silently read as "no MIG here" and the node was reported as converged without ever being looked at. The query is now kept apart from the match so its own status survives: 0 a GPU reports a MIG mode -> configure the node 1 queried fine, no MIG support -> skip the node 2 the query itself failed -> fail, with the reason Exit 2 is only fatal on a node where lspci actually found NVIDIA GPUs, so a node with no NVIDIA hardware still passes through untouched. Two smaller corrections from the same reading: the match accepts only a real MIG mode, since `grep -v N/A` also matched warnings nvidia-smi prints next to the values; and the Red Hat install branch now requires a successful capability probe, which the Debian branch already did. scripts/validation/tests/test_gpu_task_contracts.py covers this. The probe is executed as written against a mocked nvidia-smi and systemd-run, for a failed query (reported on stdout, on stderr, and not at all), an unsupported device, a noisy but successful query, and a repeated run with the device guard already active. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01UiFVTVAyEMmn68Xvo9LSTf Signed-off-by: 백지명 --- playbooks/nvidia-software/nvidia-cuda.yml | 4 +- playbooks/nvidia-software/nvidia-driver.yml | 6 +- playbooks/nvidia-software/nvidia-mig.yml | 61 +++- playbooks/utilities/gpu-clocks.yml | 4 +- playbooks/utilities/nvidia-set-gpu-clocks.yml | 8 +- roles/nvidia-mig-manager/tasks/main.yml | 38 ++- .../tests/test_gpu_task_contracts.py | 314 ++++++++++++++++++ 7 files changed, 408 insertions(+), 27 deletions(-) create mode 100644 scripts/validation/tests/test_gpu_task_contracts.py diff --git a/playbooks/nvidia-software/nvidia-cuda.yml b/playbooks/nvidia-software/nvidia-cuda.yml index b023b1e65..e423a7805 100644 --- a/playbooks/nvidia-software/nvidia-cuda.yml +++ b/playbooks/nvidia-software/nvidia-cuda.yml @@ -34,7 +34,9 @@ # Run in a transient unit so the check also works on nodes where the # Slurm login setup hides GPUs from ssh sessions (see login-compute-setup). - name: test nvidia-smi - command: systemd-run --wait --pipe --collect --quiet nvidia-smi + command: >- + systemd-run --wait --pipe --collect --quiet -- + nvidia-smi changed_when: false when: - ansible_local['gpus']['count'] diff --git a/playbooks/nvidia-software/nvidia-driver.yml b/playbooks/nvidia-software/nvidia-driver.yml index c445e5574..d8921d551 100644 --- a/playbooks/nvidia-software/nvidia-driver.yml +++ b/playbooks/nvidia-software/nvidia-driver.yml @@ -23,8 +23,12 @@ name: nvidia.nvidia_driver when: (ansible_local['gpus']['count'] and is_dgx.stat.exists == False) or (nvidia_driver_force_install|default(false)) + # Run in a transient unit so the check also works on nodes where the Slurm + # login setup hides GPUs from ssh sessions (see login-compute-setup). - name: test nvidia-smi - command: nvidia-smi + command: >- + systemd-run --wait --pipe --collect --quiet -- + nvidia-smi changed_when: false when: - ansible_local['gpus']['count'] diff --git a/playbooks/nvidia-software/nvidia-mig.yml b/playbooks/nvidia-software/nvidia-mig.yml index 14a32089f..a77bda84c 100644 --- a/playbooks/nvidia-software/nvidia-mig.yml +++ b/playbooks/nvidia-software/nvidia-mig.yml @@ -6,32 +6,58 @@ # Check if MIG capabilities and software on nodes - hosts: all + become: true vars: mig_manager_reboot_timeout: 900 tasks: - # Run nvidia-smi in a transient unit so the probe also works on nodes where - # the Slurm login setup hides GPUs from ssh sessions, and only accept a real - # MIG mode: with hidden GPUs nvidia-smi prints "No devices were found", - # which the previous 'grep -v N/A' let through as a match. + # ansible_local['gpus']['count'] comes from lspci, so it reports the GPUs the + # node has even while the Slurm login setup hides them from ssh sessions. + - name: install custom facts + include_role: + name: facts + + # nvidia-smi runs in a transient unit so the probe also works on nodes where + # the Slurm login setup hides GPUs from ssh sessions. The query is kept apart + # from the match so its exit status survives: piping the two together reports + # the same status whether nvidia-smi answered 'N/A' or never answered at all. + # 0 - a GPU reports a MIG mode, configure the node + # 1 - the query succeeded and no GPU supports MIG, skip the node + # 2 - the query itself failed, which is an error rather than an answer + # Only a real mode counts: 'grep -v N/A' also matched warnings and notices + # nvidia-smi prints alongside the values. - name: check for MIG capable devices - shell: >- - set -o pipefail && - systemd-run --wait --pipe --collect --quiet - nvidia-smi --query-gpu=mig.mode.current --format=csv,noheader - | grep -E '^(Enabled|Disabled)$' + shell: | + set -o pipefail + mig_mode=$(systemd-run --wait --pipe --collect --quiet -- \ + nvidia-smi --query-gpu=mig.mode.current --format=csv,noheader) || exit 2 + printf '%s\n' "$mig_mode" | grep -E '^(Enabled|Disabled)$' args: executable: /bin/bash - become: true register: has_mig failed_when: false changed_when: false + # A node whose GPUs cannot be queried is not a node without MIG: skipping it + # silently would report a MIG layout as converged that was never looked at. + # ansible_local['gpus']['count'] comes from lspci, so a node with no NVIDIA + # hardware at all still passes through untouched. + - name: fail when the MIG capability probe could not run + fail: + msg: >- + Could not read the MIG mode of the GPUs on this node. + nvidia-smi is unavailable or did not answer, while lspci reports + {{ ansible_local['gpus']['count'] }} NVIDIA GPU(s). + Install the driver on this node, or exclude it from this play. + when: + - has_mig.rc == 2 + - ansible_local['gpus']['count'] | default(0) | int > 0 + - name: Install MIG Manager include_role: name: nvidia-mig-manager when: - has_mig.rc == 0 - + # TODO: Consider storing a custom copy of the hooks.yaml configuration alongside the config.yaml - name: copy cluster-wide mig config file copy: @@ -39,8 +65,14 @@ dest: "{{ mig_manager_config }}" when: has_mig.rc == 0 + # nvidia-mig-parted drives the same GPUs as the probe, so it needs the same + # isolation from the ssh cgroup. A failure here is left unhandled on purpose: + # a MIG layout that could not be applied must not pass as applied. - name: Apply MIG configuration - command: nvidia-mig-parted apply -f {{ mig_manager_config }} -c {{ mig_manager_profile }} -k {{ mig_manager_hooks }} + command: >- + systemd-run --wait --pipe --collect --quiet -- + nvidia-mig-parted apply -f {{ mig_manager_config }} + -c {{ mig_manager_profile }} -k {{ mig_manager_hooks }} when: has_mig.rc == 0 # Reboot nodes if necessary and poll for them to come up @@ -61,5 +93,8 @@ - has_mig.rc == 0 - name: Assert MIG configuration was applied - command: nvidia-mig-parted assert -f {{ mig_manager_config }} -c {{ mig_manager_profile }} + command: >- + systemd-run --wait --pipe --collect --quiet -- + nvidia-mig-parted assert -f {{ mig_manager_config }} + -c {{ mig_manager_profile }} when: has_mig.rc == 0 diff --git a/playbooks/utilities/gpu-clocks.yml b/playbooks/utilities/gpu-clocks.yml index ec5cafa13..ed5767427 100644 --- a/playbooks/utilities/gpu-clocks.yml +++ b/playbooks/utilities/gpu-clocks.yml @@ -11,7 +11,9 @@ # commands run over the Ansible connection cannot see the devices. # Note: -acp is deprecated in current drivers and treated as a no-op. - name: set GPU clocks permissions - command: systemd-run --wait --pipe --collect --quiet nvidia-smi -acp UNRESTRICTED + command: >- + systemd-run --wait --pipe --collect --quiet -- + nvidia-smi -acp UNRESTRICTED changed_when: false when: - ansible_local['gpus']['count'] diff --git a/playbooks/utilities/nvidia-set-gpu-clocks.yml b/playbooks/utilities/nvidia-set-gpu-clocks.yml index 586f2a263..4f98eeb5e 100644 --- a/playbooks/utilities/nvidia-set-gpu-clocks.yml +++ b/playbooks/utilities/nvidia-set-gpu-clocks.yml @@ -7,9 +7,13 @@ # nvidia-smi runs in a transient unit so this also works on nodes where the # Slurm login setup hides GPUs from ssh sessions (see login-compute-setup). - name: set the gpu clock to a specified amount - command: "systemd-run --wait --pipe --collect --quiet nvidia-smi -lgc {{ gpu_clock_lock }}" + command: >- + systemd-run --wait --pipe --collect --quiet -- + nvidia-smi -lgc {{ gpu_clock_lock }} when: not gpu_clock_reset - name: reset the gpu clock to the default - command: systemd-run --wait --pipe --collect --quiet nvidia-smi -rgc + command: >- + systemd-run --wait --pipe --collect --quiet -- + nvidia-smi -rgc when: gpu_clock_reset diff --git a/roles/nvidia-mig-manager/tasks/main.yml b/roles/nvidia-mig-manager/tasks/main.yml index 1e171737f..9d73aa804 100644 --- a/roles/nvidia-mig-manager/tasks/main.yml +++ b/roles/nvidia-mig-manager/tasks/main.yml @@ -2,22 +2,41 @@ # Install the NVIDIA MIG Manager tooling on all MIG-capable nodes # Check node state. nvidia-smi runs in a transient unit so the probe also works -# on nodes where the Slurm login setup hides GPUs from ssh sessions, and only a -# real MIG mode counts: with hidden GPUs nvidia-smi prints "No devices were -# found", which 'grep -v N/A' would accept as a match. +# on nodes where the Slurm login setup hides GPUs from ssh sessions. The query is +# kept apart from the match so its exit status survives: piping the two together +# reports the same status whether nvidia-smi answered 'N/A' or never answered. +# 0 - a GPU reports a MIG mode, install the MIG manager +# 1 - the query succeeded and no GPU supports MIG, skip the node +# 2 - the query itself failed, which is an error rather than an answer +# Only a real mode counts: 'grep -v N/A' also matched warnings and notices +# nvidia-smi prints alongside the values. - name: check for MIG capable devices - shell: >- - set -o pipefail && - systemd-run --wait --pipe --collect --quiet - nvidia-smi --query-gpu=mig.mode.current --format=csv,noheader - | grep -E '^(Enabled|Disabled)$' + become: true + shell: | + set -o pipefail + mig_mode=$(systemd-run --wait --pipe --collect --quiet -- \ + nvidia-smi --query-gpu=mig.mode.current --format=csv,noheader) || exit 2 + printf '%s\n' "$mig_mode" | grep -E '^(Enabled|Disabled)$' args: executable: "/bin/bash" - become: true register: has_mig failed_when: false changed_when: false +# A node whose GPUs cannot be queried is not a node without MIG: installing +# nothing and reporting it as converged hides a broken driver. The lspci fact +# keeps nodes with no NVIDIA hardware out of this. +- name: fail when the MIG capability probe could not run + fail: + msg: >- + Could not read the MIG mode of the GPUs on this node. + nvidia-smi is unavailable or did not answer, while lspci reports + {{ ansible_local['gpus']['count'] }} NVIDIA GPU(s). + Install the driver on this node, or exclude it from this play. + when: + - has_mig.rc == 2 + - ansible_local['gpus']['count'] | default(0) | int > 0 + - name: check for NVIDIA MIG parted shell: which nvidia-mig-parted # noqa command-instead-of-shell register: has_mig_parted @@ -37,5 +56,6 @@ name: "{{ mig_manager_url_rpm }}" state: present when: + - has_mig.rc == 0 - has_mig_parted.rc != 0 - ansible_os_family == "RedHat" diff --git a/scripts/validation/tests/test_gpu_task_contracts.py b/scripts/validation/tests/test_gpu_task_contracts.py new file mode 100644 index 000000000..3035e5400 --- /dev/null +++ b/scripts/validation/tests/test_gpu_task_contracts.py @@ -0,0 +1,314 @@ +"""Tests for the GPU tasks that have to escape the Slurm login GPU cgroup. + +The Slurm login setup puts a DeviceAllow guard on the ssh service, so anything +Ansible runs over its ssh connection sees no GPUs. Every privileged GPU command +therefore runs in a transient systemd unit. These tests pin that contract and +exercise the MIG capability probe against a mocked nvidia-smi, so the difference +between "this GPU does not support MIG" and "the query never ran" is covered. + +Run with: python3 -m unittest discover scripts/validation/tests +""" + +import os +import subprocess +import tempfile +import textwrap +import unittest +from pathlib import Path + +import yaml + + +ROOT = Path(__file__).resolve().parents[3] +TRANSIENT_SERVICE = "systemd-run --wait --pipe --collect --quiet --" +PROBE_NAME = "check for MIG capable devices" +MIG_PLAYBOOK = "playbooks/nvidia-software/nvidia-mig.yml" +MIG_ROLE = "roles/nvidia-mig-manager/tasks/main.yml" +GPU_TASK_PATHS = ( + "playbooks/nvidia-software/nvidia-driver.yml", + "playbooks/nvidia-software/nvidia-cuda.yml", + "playbooks/utilities/gpu-clocks.yml", + "playbooks/utilities/nvidia-set-gpu-clocks.yml", + MIG_PLAYBOOK, + MIG_ROLE, +) + + +def tasks(path): + """Every command/shell task in a playbook or task file, flattened.""" + found = [] + + def visit(value): + if isinstance(value, dict): + if "command" in value or "shell" in value: + found.append(value) + for child in value.values(): + visit(child) + elif isinstance(value, list): + for child in value: + visit(child) + + visit(yaml.safe_load((ROOT / path).read_text(encoding="utf-8"))) + return found + + +def gpu_commands(path): + """The (task, normalized command) pairs that drive the GPUs.""" + pairs = [] + for task in tasks(path): + command = task.get("command", task.get("shell")) + if isinstance(command, str) and ( + "nvidia-smi" in command + or "nvidia-mig-parted apply" in command + or "nvidia-mig-parted assert" in command + ): + pairs.append((task, " ".join(command.split()))) + return pairs + + +def all_tasks(path): + """Every task in a playbook or task file, whatever module it uses.""" + found = [] + + def visit(value): + if isinstance(value, dict): + # A task carries a name and a module; `include_role: {name: x}` does not. + if isinstance(value.get("name"), str) and len(value) > 1: + found.append(value) + for child in value.values(): + visit(child) + elif isinstance(value, list): + for child in value: + visit(child) + + visit(yaml.safe_load((ROOT / path).read_text(encoding="utf-8"))) + return found + + +def task_named(path, name): + matches = [task for task in all_tasks(path) if task.get("name") == name] + assert len(matches) == 1, f"{path}: {name}" + return matches[0] + + +class TransientExecutionContract(unittest.TestCase): + """Nothing that touches a GPU may run directly in the ssh context.""" + + def test_every_gpu_command_runs_in_a_transient_unit(self): + for path in GPU_TASK_PATHS: + for task, command in gpu_commands(path): + with self.subTest(path=path, task=task.get("name")): + self.assertIn(TRANSIENT_SERVICE, command) + driven = "nvidia-smi" if "nvidia-smi" in command else "nvidia-mig-parted" + self.assertLess( + command.index(TRANSIENT_SERVICE), command.index(driven) + ) + + def test_mig_administration_is_covered_too(self): + # The probe is useless on its own: once it succeeds, apply and assert + # drive the very same GPUs and need the same isolation. + commands = [command for _, command in gpu_commands(MIG_PLAYBOOK)] + for fragment in ( + "nvidia-smi --query-gpu=mig.mode.current", + "nvidia-mig-parted apply", + "nvidia-mig-parted assert", + ): + with self.subTest(fragment=fragment): + matching = [c for c in commands if fragment in c] + self.assertEqual(len(matching), 1, fragment) + self.assertIn(TRANSIENT_SERVICE, matching[0]) + + def test_mig_operations_run_privileged(self): + plays = yaml.safe_load((ROOT / MIG_PLAYBOOK).read_text(encoding="utf-8")) + self.assertEqual(len(plays), 1) + self.assertIs(plays[0]["become"], True) + self.assertIs(task_named(MIG_ROLE, PROBE_NAME)["become"], True) + + def test_gpu_clock_arguments_avoid_shell_interpolation(self): + path = "playbooks/utilities/nvidia-set-gpu-clocks.yml" + for name in ( + "set the gpu clock to a specified amount", + "reset the gpu clock to the default", + ): + with self.subTest(name=name): + task = task_named(path, name) + self.assertIn("command", task) + self.assertNotIn("shell", task) + + +class MigApplyFailureContract(unittest.TestCase): + """A MIG layout that could not be applied must not pass as applied.""" + + def test_apply_and_assert_surface_their_failures(self): + for name in ("Apply MIG configuration", "Assert MIG configuration was applied"): + with self.subTest(name=name): + task = task_named(MIG_PLAYBOOK, name) + self.assertNotIn("failed_when", task) + self.assertNotIn("ignore_errors", task) + + def test_installation_requires_a_successful_capability_probe(self): + # Without this the Red Hat branch installed the MIG manager on nodes + # whose capability was never established. + for name, family in ( + ("Install MIG Manager (apt)", "Debian"), + ("Install MIG Manager (yum)", "RedHat"), + ): + with self.subTest(name=name): + tasks_by_name = { + task["name"]: task + for task in yaml.safe_load( + (ROOT / MIG_ROLE).read_text(encoding="utf-8") + ) + } + self.assertEqual( + set(tasks_by_name[name]["when"]), + { + "has_mig.rc == 0", + "has_mig_parted.rc != 0", + f'ansible_os_family == "{family}"', + }, + ) + + def test_a_failed_probe_is_reported_rather_than_skipped(self): + # rc 2 is the probe's own "the query never ran"; the lspci fact keeps + # nodes without NVIDIA hardware out of it. + for path in (MIG_PLAYBOOK, MIG_ROLE): + with self.subTest(path=path): + guard = task_named(path, "fail when the MIG capability probe could not run") + self.assertEqual( + set(guard["when"]), + { + "has_mig.rc == 2", + "ansible_local['gpus']['count'] | default(0) | int > 0", + }, + ) + + +class MigProbeBehaviour(unittest.TestCase): + """Run the probe as written, against a mocked nvidia-smi.""" + + def probe(self, path): + probe = task_named(path, PROBE_NAME) + self.assertIs(probe["failed_when"], False) + self.assertIs(probe["changed_when"], False) + return probe["shell"] + + def run_probe(self, path, mode, ssh_device_guard=False): + """Execute the probe's shell exactly as Ansible would.""" + script = self.probe(path) + with tempfile.TemporaryDirectory() as tmp: + mocks = Path(tmp) + + # systemd-run drops the ssh cgroup: it consumes its own options up to + # the '--' separator and runs the rest outside the device guard. + (mocks / "systemd-run").write_text( + textwrap.dedent( + """\ + #!/usr/bin/env bash + while [ "$1" != "--" ]; do + case "$1" in + --wait|--pipe|--collect|--quiet) shift ;; + *) echo "systemd-run: unexpected argument $1" >&2; exit 1 ;; + esac + done + shift + SSH_DEVICE_GUARD=0 exec "$@" + """ + ), + encoding="utf-8", + ) + # nvidia-smi reports a missing device on stdout and a broken driver + # on stderr, so the probe must not depend on which stream carried it. + (mocks / "nvidia-smi").write_text( + textwrap.dedent( + """\ + #!/usr/bin/env bash + if [ "${SSH_DEVICE_GUARD:-0}" = "1" ]; then + echo "No devices were found" + exit 6 + fi + case "$MODE" in + mig) printf 'Enabled\\nDisabled\\n'; exit 0 ;; + no-mig) printf 'N/A\\nN/A\\n'; exit 0 ;; + noisy) printf 'N/A\\nWARNING: infoROM is corrupted\\n'; exit 0 ;; + hidden) echo "No devices were found"; exit 6 ;; + nodriver) echo "NVIDIA-SMI has failed to communicate with the driver" >&2 + exit 9 ;; + missing) exit 127 ;; + esac + """ + ), + encoding="utf-8", + ) + for mock in ("systemd-run", "nvidia-smi"): + (mocks / mock).chmod(0o755) + + return subprocess.run( + ["/bin/bash", "-c", script], + env={ + **os.environ, + "PATH": f"{mocks}:{os.environ['PATH']}", + "MODE": mode, + "SSH_DEVICE_GUARD": "1" if ssh_device_guard else "0", + }, + text=True, + capture_output=True, + check=False, + ) + + def assert_operational_failure(self, result): + """What the playbook reads as "the query never ran".""" + self.assertEqual(result.returncode, 2) + + def assert_unsupported_device(self, result): + """What the playbook reads as "these GPUs do not do MIG".""" + self.assertEqual(result.returncode, 1) + + def test_mig_capable_devices_are_detected(self): + for path in (MIG_PLAYBOOK, MIG_ROLE): + with self.subTest(path=path): + result = self.run_probe(path, "mig") + self.assertEqual(result.returncode, 0, result.stderr) + self.assertEqual(result.stdout.split(), ["Enabled", "Disabled"]) + + def test_devices_without_mig_support_skip_configuration(self): + for path in (MIG_PLAYBOOK, MIG_ROLE): + with self.subTest(path=path): + self.assert_unsupported_device(self.run_probe(path, "no-mig")) + + def test_only_a_real_mig_mode_counts_as_a_capability(self): + # nvidia-smi prints warnings alongside the values it was asked for, and + # 'grep -v N/A' accepted those as a MIG mode. + for path in (MIG_PLAYBOOK, MIG_ROLE): + with self.subTest(path=path): + self.assert_unsupported_device(self.run_probe(path, "noisy")) + + def test_a_failed_query_does_not_look_like_a_missing_capability(self): + # Whether the diagnostic went to stdout, to stderr or nowhere at all, + # a query that did not answer has to be distinguishable. + for path in (MIG_PLAYBOOK, MIG_ROLE): + for mode in ("hidden", "nodriver", "missing"): + with self.subTest(path=path, mode=mode): + self.assert_operational_failure(self.run_probe(path, mode)) + + def test_the_probe_escapes_an_active_ssh_device_guard(self): + # The regression this PR is about: with the guard already applied, the + # bare command sees no devices, and the transient unit still does. + for path in (MIG_PLAYBOOK, MIG_ROLE): + with self.subTest(path=path): + result = self.run_probe(path, "mig", ssh_device_guard=True) + self.assertEqual(result.returncode, 0, result.stderr) + + def test_the_probe_is_stable_across_repeated_runs(self): + # Convergence runs after the login isolation is in place have to reach + # the same answer as the run that applied it. + for path in (MIG_PLAYBOOK, MIG_ROLE): + with self.subTest(path=path): + first = self.run_probe(path, "mig", ssh_device_guard=False) + second = self.run_probe(path, "mig", ssh_device_guard=True) + self.assertEqual(first.returncode, second.returncode) + self.assertEqual(first.stdout, second.stdout) + + +if __name__ == "__main__": + unittest.main()