playbooks: run nvidia-smi tasks outside the ssh cgroup so they survive the login GPU guard

playbooks/nvidia-software/nvidia-cuda.yml

@@ -31,8 +31,12 @@
       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']

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']

playbooks/nvidia-software/nvidia-mig.yml

@@ -6,29 +6,73 @@
 
 # Check if MIG capabilities and software on nodes
 - hosts: all
+  become: true
   vars:
     mig_manager_reboot_timeout: 900
   tasks:
+    # 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: nvidia-smi --query-gpu=mig.mode.current --format=csv,noheader | grep -v 'N/A'
+      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
       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:
         src: "../../config/nvidia-mig-config.yml"
         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
@@ -49,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

playbooks/utilities/gpu-clocks.yml

@@ -6,8 +6,14 @@
       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']

playbooks/utilities/nvidia-set-gpu-clocks.yml

@@ -4,10 +4,16 @@
   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

roles/nvidia-mig-manager/tasks/main.yml

@@ -1,15 +1,42 @@
 ---
 # 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. 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 && nvidia-smi --query-gpu=mig.mode.current --format=csv,noheader | grep -v 'N/A'
+  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"
   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
@@ -29,5 +56,6 @@
     name: "{{ mig_manager_url_rpm }}"
     state: present
   when:
+    - has_mig.rc == 0
     - has_mig_parted.rc != 0
     - ansible_os_family == "RedHat"