nvidia-gpu-tests: optionally fail the play on a failed DCGM diagnostic

roles/nvidia-gpu-tests/README.md

@@ -9,7 +9,7 @@ Running this will perform the following:
 * Install the CUDA toolkit
 * Download and build cuda-samples
 * Run the Peer2Peer and MatrixMultiply tests
-* Run the DCGM diagnostics
+* Run the DCGM diagnostics (set `gpu_test_dcgm_fail_on_error: yes` to fail the play when `dcgmi diag` exits non-zero; by default failures are ignored)
 * Run a basic Tensorflow DL job
 
 

roles/nvidia-gpu-tests/defaults/main.yml

@@ -14,6 +14,10 @@ gpu_test_build: yes
 # 1 -> quick test (seconds), 2 -> regular test (minute), 3 -> long test (minutes)
 gpu_test_dcgm_level: 3
 
+# Whether a failed DCGM diagnostic (dcgmi exits non-zero) fails the play. Set to
+# yes to stop on a failing GPU; by default the failure is ignored and only reported.
+gpu_test_dcgm_fail_on_error: no
+
 # Batch size to use for Tensorflow tests
 # DGX-2 -> (512, 50); DGX-1 -> (256,50); DGX-Station -> (128, 50)
 gpu_test_tf_batch_size: 128 # (choose from 32, 64, 128, 256, 512)

roles/nvidia-gpu-tests/tasks/main.yml

@@ -44,7 +44,7 @@
 - name: run dcgm diag for up to 15 minutes
   shell: "dcgmi diag -r {{ gpu_test_dcgm_level }}"
   register: dcgm
-  ignore_errors: yes
+  ignore_errors: "{{ not (gpu_test_dcgm_fail_on_error | bool) }}"
 
 # Note, this may fail to run depending on Ansible settings
 # We only look at the last few lines with relevant output

tests/nvidia-gpu-tests/test_dcgm_failure.py

@@ -0,0 +1,72 @@
+#!/usr/bin/env python3
+"""Exercise the real DCGM task locally with a fake dcgmi; no GPU required.
+
+Run with: python3 tests/nvidia-gpu-tests/test_dcgm_failure.py
+Requires PyYAML and ansible-playbook on PATH.
+"""
+import os
+from pathlib import Path
+import subprocess
+import tempfile
+import unittest
+
+import yaml
+
+
+ROOT = Path(__file__).resolve().parents[2]
+ROLE = ROOT / "roles/nvidia-gpu-tests"
+
+
+class DcgmFailureTest(unittest.TestCase):
+    def test_exit_policy(self):
+        tasks = yaml.safe_load((ROLE / "tasks/main.yml").read_text())
+        task = next(task for task in tasks if task.get("register") == "dcgm")
+        defaults = yaml.safe_load((ROLE / "defaults/main.yml").read_text())
+        # Removing boolean normalization must break the string-false cases.
+        cases = [
+            ("default", [], False),
+            ("boolean false", ["-e", '{"gpu_test_dcgm_fail_on_error": false}'], False),
+            ("boolean true", ["-e", '{"gpu_test_dcgm_fail_on_error": true}'], True),
+            ("string false", ["-e", "gpu_test_dcgm_fail_on_error=false"], False),
+            ("string no", ["-e", "gpu_test_dcgm_fail_on_error=no"], False),
+            ("string true", ["-e", "gpu_test_dcgm_fail_on_error=true"], True),
+            ("string yes", ["-e", "gpu_test_dcgm_fail_on_error=yes"], True),
+        ]
+        with tempfile.TemporaryDirectory(prefix="dcgm-test-") as directory:
+            temp = Path(directory)
+            stub = temp / "dcgmi"
+            stub.write_text('#!/bin/sh\n[ "$*" = "diag -r 3" ] || exit 99\nprintf "diagnostic fixture\\n"\nexit "$DIAG_RC"\n')
+            stub.chmod(0o755)
+            config = temp / "ansible.cfg"
+            config.write_text("[defaults]\nretry_files_enabled = False\nhost_key_checking = False\n")
+            play = [{
+                "hosts": "localhost", "connection": "local", "gather_facts": False,
+                "vars": defaults,
+                "environment": {"PATH": str(temp) + os.pathsep + os.environ["PATH"]},
+                "tasks": [task, {"name": "Verify diagnostic result", "assert": {"that": [
+                    "dcgm.rc == expected_rc | int", "'diagnostic fixture' in dcgm.stdout"
+                ]}}],
+            }]
+            playbook = temp / "test.yml"
+            playbook.write_text(yaml.safe_dump(play))
+            env = dict(os.environ, ANSIBLE_CONFIG=str(config), ANSIBLE_HOME=str(temp / "home"),
+                       ANSIBLE_LOCAL_TEMP=str(temp / "local"), ANSIBLE_REMOTE_TEMP=str(temp / "remote"),
+                       ANSIBLE_NOCOLOR="1")
+            for name, extra, fail_on_error in cases:
+                for rc in (0, 1):
+                    with self.subTest(setting=name, diagnostic_rc=rc):
+                        result = subprocess.run(
+                            ["ansible-playbook", "-i", "localhost,", str(playbook),
+                             "-e", f"expected_rc={rc}", *extra],
+                            cwd=temp, env=dict(env, DIAG_RC=str(rc)),
+                            text=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, timeout=60,
+                        )
+                        if rc:
+                            self.assertIn("diagnostic fixture", result.stdout, result.stdout)
+                        self.assertEqual(result.returncode, 2 if rc and fail_on_error else 0, result.stdout)
+                        if not (rc and fail_on_error):
+                            self.assertIn("All assertions passed", result.stdout, result.stdout)
+
+
+if __name__ == "__main__":
+    unittest.main()