diff --git a/plugins/lookup/gpu.py b/plugins/lookup/gpu.py new file mode 100644 index 0000000000..d72e80422e --- /dev/null +++ b/plugins/lookup/gpu.py @@ -0,0 +1,82 @@ +"""Whether a service gets a GPU: it asked for one AND the host carries one. + +Two facts have to agree and used to be restated at every call site. The +declaration lives in ``applications`` under ``services..gpu``; the +hardware is ``sys-svc-container``'s stat of ``/dev/nvidia0`` on the target +host, registered as ``sys_svc_container_nvidia_device``. Asking only the +declaration selects a ``-cuda`` image on a host that has no device, and that +tag has no arm64 manifest at all, so the build fails with ``no match for +platform in manifest``. + +The stat is a target-host fact, so this reads the register out of the play's +variables rather than inspecting the controller it runs on. +""" + +from __future__ import annotations + +from typing import Any + +from ansible.errors import AnsibleError +from ansible.module_utils.parsing.convert_bool import boolean +from ansible.plugins.loader import lookup_loader +from ansible.plugins.lookup import LookupBase + +from plugins.filter.resource_filter import resource_filter + +DEVICE_FACT = "sys_svc_container_nvidia_device" + + +def _device_present(variables: dict) -> bool: + """Whether the target host registered an NVIDIA device node. + + Args: + variables: the play's variables, carrying the register. + + Returns: + False while the register is absent, which is also the answer before + ``sys-svc-container`` has run. + """ + register = variables.get(DEVICE_FACT) or {} + if not isinstance(register, dict): + return False + return bool((register.get("stat") or {}).get("exists")) + + +class LookupModule(LookupBase): + def run( + self, + terms: list[Any], + variables: dict | None = None, + **kwargs: Any, + ) -> list[bool]: + """Return ``[bool]``: the service is GPU-accelerated on this host. + + Args: + terms: ``application_id`` and optionally the service name; the + service defaults to the application's own entity. + variables: the play's variables. + + Raises: + AnsibleError: wrong number of terms. + """ + if not 1 <= len(terms) <= 2: + raise AnsibleError( + "gpu: expected 1 or 2 terms: lookup('gpu', application_id[, service])" + ) + application_id = str(terms[0]) + service_name = str(terms[1]) if len(terms) == 2 else "" + + templar = getattr(self, "_templar", None) + variables = variables or getattr(templar, "available_variables", {}) or {} + + applications = lookup_loader.get( + "applications", loader=self._loader, templar=templar + ).run([], variables=variables)[0] + + wanted = resource_filter(applications, application_id, "gpu", service_name, False) + # Exception: boolean(), not bool(). The call sites this replaces read + # the flag through Jinja's `| bool`, which answers False for the + # strings "false", "False", "0" and "no" where bool() answers True. A + # templated declaration renders to such a string, and the difference + # would attach a GPU runtime nothing asked for. + return [boolean(wanted, strict=False) and _device_present(variables)] diff --git a/roles/svc-ai-libretranslate-engine/tasks/00_core.yml b/roles/svc-ai-libretranslate-engine/tasks/00_core.yml index 5bd00fca59..ae597c1159 100644 --- a/roles/svc-ai-libretranslate-engine/tasks/00_core.yml +++ b/roles/svc-ai-libretranslate-engine/tasks/00_core.yml @@ -56,5 +56,3 @@ when: - IS_STACK_HOST | bool - LIBRETRANSLATE_ENGINE_GPU | bool - - sys_svc_container_nvidia_device is defined - - sys_svc_container_nvidia_device.stat.exists diff --git a/roles/svc-ai-libretranslate-engine/vars/main.yml b/roles/svc-ai-libretranslate-engine/vars/main.yml index 72e30a8353..e842a32357 100644 --- a/roles/svc-ai-libretranslate-engine/vars/main.yml +++ b/roles/svc-ai-libretranslate-engine/vars/main.yml @@ -5,7 +5,7 @@ entity_name: "{{ application_id | entity_name }}" LIBRETRANSLATE_ENGINE_SERVICE: "{{ entity_name }}" LIBRETRANSLATE_ENGINE_IMAGE: "{{ lookup('config', application_id, 'services.libretranslate-engine.image') }}" -LIBRETRANSLATE_ENGINE_GPU: "{{ lookup('applications') | resource_filter(application_id, 'gpu', 'libretranslate-engine', false) }}" +LIBRETRANSLATE_ENGINE_GPU: "{{ lookup('gpu', application_id, 'libretranslate-engine') }}" LIBRETRANSLATE_ENGINE_VERSION: "{{ lookup('config', application_id, 'services.libretranslate-engine.version') ~ ('-cuda' if LIBRETRANSLATE_ENGINE_GPU | bool else '') }}" LIBRETRANSLATE_ENGINE_LOAD_ONLY: "{{ lookup('config', application_id, 'services.libretranslate-engine.load_only') | join(',') }}" diff --git a/roles/sys-svc-container/templates/resource.yml.j2 b/roles/sys-svc-container/templates/resource.yml.j2 index 395c1952af..845b0fa4e0 100644 --- a/roles/sys-svc-container/templates/resource.yml.j2 +++ b/roles/sys-svc-container/templates/resource.yml.j2 @@ -2,6 +2,6 @@ cpus: {{ lookup('applications') | resource_filter(application_id, mem_reservation: {{ lookup('applications') | resource_filter(application_id, 'mem_reservation', service_name | default(''), RESOURCE_MEM_RESERVATION) }} mem_limit: {{ lookup('applications') | resource_filter(application_id, 'mem_limit', service_name | default(''), RESOURCE_MEM_LIMIT) }} pids_limit: {{ lookup('applications') | resource_filter(application_id, 'pids_limit', service_name | default(''), RESOURCE_PIDS_LIMIT) }} -{% if (lookup('applications') | resource_filter(application_id, 'gpu', service_name | default(''), false)) | bool and (sys_svc_container_nvidia_device is defined and sys_svc_container_nvidia_device.stat.exists) %} +{% if lookup('gpu', application_id, service_name | default('')) %} runtime: nvidia {% endif %} diff --git a/roles/web-svc-libretranslate/tasks/00_core.yml b/roles/web-svc-libretranslate/tasks/00_core.yml index 5a748a5a1a..3975e6a20b 100644 --- a/roles/web-svc-libretranslate/tasks/00_core.yml +++ b/roles/web-svc-libretranslate/tasks/00_core.yml @@ -49,5 +49,3 @@ when: - IS_STACK_HOST | bool - LIBRETRANSLATE_GPU | bool - - sys_svc_container_nvidia_device is defined - - sys_svc_container_nvidia_device.stat.exists diff --git a/roles/web-svc-libretranslate/vars/main.yml b/roles/web-svc-libretranslate/vars/main.yml index 2ce3f36145..396203f5f7 100644 --- a/roles/web-svc-libretranslate/vars/main.yml +++ b/roles/web-svc-libretranslate/vars/main.yml @@ -5,7 +5,7 @@ entity_name: "{{ application_id | entity_name }}" LIBRETRANSLATE_SERVICE: "{{ entity_name }}" LIBRETRANSLATE_IMAGE: "{{ lookup('config', application_id, 'services.libretranslate.image') }}" -LIBRETRANSLATE_GPU: "{{ lookup('applications') | resource_filter(application_id, 'gpu', 'libretranslate', false) }}" +LIBRETRANSLATE_GPU: "{{ lookup('gpu', application_id, 'libretranslate') }}" LIBRETRANSLATE_VERSION: "{{ lookup('config', application_id, 'services.libretranslate.version') ~ ('-cuda' if LIBRETRANSLATE_GPU | bool else '') }}" LIBRETRANSLATE_LOAD_ONLY: "{{ lookup('config', application_id, 'services.libretranslate.load_only') | join(',') }}" LIBRETRANSLATE_DIRECTIONS: "{{ lookup('config', application_id, 'services.libretranslate.directions') }}" diff --git a/tests/unit/python/plugins/lookup/test_gpu.py b/tests/unit/python/plugins/lookup/test_gpu.py new file mode 100644 index 0000000000..8d123526c0 --- /dev/null +++ b/tests/unit/python/plugins/lookup/test_gpu.py @@ -0,0 +1,101 @@ +"""Unit tests for plugins/lookup/gpu.py. + +Mocks the Ansible-side machinery (loader, lookup_loader, templar) and pins +that the lookup answers True only when the declaration AND the host fact +agree. The declaration resolution itself belongs to resource_filter and is +covered by tests.unit.python.plugins.filter.test_resource_filter. +""" + +from __future__ import annotations + +import unittest +import unittest.mock as mock +from typing import ClassVar + +from ansible.errors import AnsibleError + +from plugins.lookup.gpu import LookupModule + +APPLICATIONS: ClassVar[dict] = { + "web-svc-libretranslate": {"services": {"libretranslate": {"gpu": True}}}, + "web-app-plain": {"services": {"plain": {"gpu": False}}}, + "web-app-stringy": {"services": {"stringy": {"gpu": "false"}}}, + "web-app-truthy": {"services": {"truthy": {"gpu": "yes"}}}, +} + +PRESENT: ClassVar[dict] = {"sys_svc_container_nvidia_device": {"stat": {"exists": True}}} +ABSENT: ClassVar[dict] = {"sys_svc_container_nvidia_device": {"stat": {"exists": False}}} + + +class _Templar: + def __init__(self, vars_=None): + self.available_variables = vars_ or {} + + def template(self, value, **_): + return value + + +def _make_lookup(vars_=None): + lm = LookupModule() + lm._templar = _Templar(vars_ or {}) + lm._loader = mock.MagicMock() + return lm + + +def _run(lm, terms): + loaded = mock.MagicMock() + loaded.run.return_value = [APPLICATIONS] + with mock.patch("plugins.lookup.gpu.lookup_loader") as loader: + loader.get.return_value = loaded + return lm.run(terms, variables=lm._templar.available_variables) + + +class TestGpuLookup(unittest.TestCase): + def test_declared_and_device_present_is_true(self): + lm = _make_lookup(PRESENT) + self.assertEqual(_run(lm, ["web-svc-libretranslate", "libretranslate"]), [True]) + + def test_declared_without_a_device_is_false(self): + """The case that selected a -cuda tag with no manifest for the platform.""" + lm = _make_lookup(ABSENT) + self.assertEqual(_run(lm, ["web-svc-libretranslate", "libretranslate"]), [False]) + + def test_undeclared_with_a_device_is_false(self): + lm = _make_lookup(PRESENT) + self.assertEqual(_run(lm, ["web-app-plain", "plain"]), [False]) + + def test_missing_register_is_false(self): + """Before sys-svc-container runs the register is absent, not falsy.""" + lm = _make_lookup({}) + self.assertEqual(_run(lm, ["web-svc-libretranslate", "libretranslate"]), [False]) + + def test_non_mapping_register_is_false(self): + lm = _make_lookup({"sys_svc_container_nvidia_device": "skipped"}) + self.assertEqual(_run(lm, ["web-svc-libretranslate", "libretranslate"]), [False]) + + def test_a_false_string_declaration_is_false(self): + """Jinja's `| bool` reads "false" as False where bool() reads it as True. + + The call sites this lookup replaces used `| bool`, so answering True + here would attach a GPU runtime nothing asked for. + """ + lm = _make_lookup(PRESENT) + self.assertEqual(_run(lm, ["web-app-stringy", "stringy"]), [False]) + + def test_a_yes_string_declaration_is_true(self): + lm = _make_lookup(PRESENT) + self.assertEqual(_run(lm, ["web-app-truthy", "truthy"]), [True]) + + def test_rejects_no_terms(self): + lm = _make_lookup(PRESENT) + with self.assertRaises(AnsibleError): + _run(lm, []) + + def test_rejects_three_terms(self): + lm = _make_lookup(PRESENT) + with self.assertRaises(AnsibleError): + _run(lm, ["a", "b", "c"]) + + +if __name__ == "__main__": + unittest.main()