Files
computer-playbook/tests/unit/python/roles/sys-svc-container/test_gpu_detection.py
T
kevinveenbirkenbachandClaude Opus 5.5 11308259b3 test(sys-svc-container): read the role vars through ROLE_FILE_VARS_MAIN
Gate pass 5 flagged test_gpu_detection.py:31 in test_file_spot: the
`ROLE / "vars" / "main.yml"` chain hard-codes a role-file path the
utils.roles.mapping SPOT already names.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
2026-10-30 13:43:02 +01:00

127 lines
4.7 KiB
Python

"""Which hosts the role treats as GPU hosts.
A privileged container, such as a swarm lab node on a workstation with an
NVIDIA card, sees the host's ``/dev/nvidia0`` but carries no driver, so
``nvidia-smi`` does not exist there. The device node alone must therefore not
make a GPU host: the GPU include and the ``gpu`` lookup both have to say no
there, and a bare host in that state has to fail instead of silently running
its GPU services on the CPU.
"""
import unittest
from pathlib import Path
from jinja2 import Environment
from plugins.lookup.gpu import _device_present
from utils import PROJECT_ROOT
from utils.cache.yaml import load_yaml_any
from utils.roles.mapping import ROLE_FILE_VARS_MAIN
ROLE = Path(PROJECT_ROOT) / "roles" / "sys-svc-container"
CORE = ROLE / "tasks" / "00_core.yml"
GPU_INCLUDE = "00_core/gpu.yml"
DEVICE = "/dev/nvidia0"
NVIDIA_REGISTER = "sys_svc_container_nvidia_"
_ENV = Environment(autoescape=True)
_ENV.filters["bool"] = bool
def _role_vars() -> dict:
return load_yaml_any(str(ROLE / ROLE_FILE_VARS_MAIN))
def _render(value: str, context: dict) -> str:
return _ENV.from_string(value).render(**context)
def _all_hold(expressions, context: dict) -> bool:
guards = expressions if isinstance(expressions, list) else [expressions]
return all(bool(_ENV.compile_expression(expr)(**context)) for expr in guards)
def _walk(*, driver: bool, device: bool, in_container: bool) -> tuple[str, dict]:
"""Run the GPU detection of 00_core.yml over a fake filesystem.
Args:
driver: whether the NVIDIA driver binary exists on the host.
device: whether the host shows an NVIDIA device node.
in_container: whether the host is itself a container.
Returns:
"gpu" when the GPU include runs, "cpu" when it does not, "refused"
when an assert stops the play; and the registers the tasks left.
"""
context = dict(_role_vars())
context["DOCKER_IN_CONTAINER"] = in_container
existing = {
_render(context["SYS_SVC_CONTAINER_GPU_DRIVER"], context): driver,
DEVICE: device,
}
outcome = "cpu"
for task in load_yaml_any(str(CORE)):
stat = task.get("ansible.builtin.stat")
assertion = task.get("ansible.builtin.assert")
if assertion is not None and NVIDIA_REGISTER not in str(assertion["that"]):
assertion = None
gpu_include = task.get("include_tasks") == GPU_INCLUDE
if stat is None and assertion is None and not gpu_include:
continue
if not _all_hold(task.get("when", []), context):
continue
if stat is not None:
path = _render(stat["path"], context)
context[task["register"]] = {"stat": {"exists": existing.get(path, False)}}
elif assertion is not None and not _all_hold(assertion["that"], context):
return "refused", context
elif gpu_include:
outcome = "gpu"
return outcome, context
class TestGpuDetection(unittest.TestCase):
def test_a_container_seeing_a_driverless_device_runs_on_the_cpu(self):
outcome, registers = _walk(driver=False, device=True, in_container=True)
self.assertEqual(
outcome,
"cpu",
"a privileged container sees the host's /dev/nvidia0 without "
"nvidia-smi; the include would fail on the missing binary",
)
self.assertFalse(_device_present(registers))
def test_a_bare_host_with_a_driverless_device_is_refused(self):
outcome, _ = _walk(driver=False, device=True, in_container=False)
self.assertEqual(outcome, "refused")
def test_driver_and_device_make_a_gpu_host(self):
for in_container in (False, True):
with self.subTest(in_container=in_container):
outcome, registers = _walk(
driver=True, device=True, in_container=in_container
)
self.assertEqual(outcome, "gpu")
self.assertTrue(_device_present(registers))
def test_no_device_is_no_gpu_host(self):
for driver in (False, True):
for in_container in (False, True):
with self.subTest(driver=driver, in_container=in_container):
outcome, registers = _walk(
driver=driver, device=False, in_container=in_container
)
self.assertEqual(outcome, "cpu")
self.assertFalse(_device_present(registers))
def test_the_query_runs_the_binary_the_detection_probed(self):
self.assertTrue(
_role_vars()["SYS_SVC_CONTAINER_GPU_QUERY"].startswith(
"{{ SYS_SVC_CONTAINER_GPU_DRIVER }} "
)
)
if __name__ == "__main__":
unittest.main()