Initial public ModelForge release

node-agent/tests/test_preflight.py

@@ -0,0 +1,80 @@
+from pathlib import Path
+
+import pytest
+from modelforge_api.domain.enums import Availability
+from modelforge_api.domain.hardware import NvidiaCollection
+from test_agent import gpu_collection
+
+from modelforge_node_agent.preflight import (
+    AgentStartupError,
+    nvidia_required,
+    nvidia_runtime_observed,
+    validate_nvidia_collection,
+)
+
+
+def test_auto_observes_injected_nvidia_device(tmp_path: Path) -> None:
+    assert not nvidia_runtime_observed(environ={}, device_root=tmp_path)
+    (tmp_path / "nvidiactl").touch()
+    assert nvidia_runtime_observed(environ={}, device_root=tmp_path)
+
+
+@pytest.mark.parametrize("value", ["all", "0", "GPU-3aaf512c"])
+def test_auto_observes_nvidia_visible_devices(value: str, tmp_path: Path) -> None:
+    assert nvidia_runtime_observed(
+        environ={"NVIDIA_VISIBLE_DEVICES": value}, device_root=tmp_path
+    )
+
+
+@pytest.mark.parametrize("value", ["", "none", "void", " NONE "])
+def test_auto_rejects_non_device_environment_values(value: str, tmp_path: Path) -> None:
+    assert not nvidia_runtime_observed(
+        environ={"NVIDIA_VISIBLE_DEVICES": value}, device_root=tmp_path
+    )
+
+
+def test_explicit_modes_do_not_depend_on_host_detection(tmp_path: Path) -> None:
+    (tmp_path / "nvidia0").touch()
+    assert nvidia_required("nvidia", environ={}, device_root=tmp_path)
+    assert not nvidia_required("cpu", environ={}, device_root=tmp_path)
+
+
+def test_nvidia_node_with_nvml_available_succeeds() -> None:
+    validate_nvidia_collection(gpu_collection(), required=True)
+
+
+def test_nvidia_node_with_nvml_unavailable_fails_closed() -> None:
+    with pytest.raises(AgentStartupError, match="NVIDIA_NVML_UNAVAILABLE") as caught:
+        validate_nvidia_collection(
+            NvidiaCollection(
+                availability=Availability.UNAVAILABLE,
+                reason="NVMLError_LibraryNotFound",
+            ),
+            required=True,
+        )
+    assert "NVMLError_LibraryNotFound" in str(caught.value)
+    assert "MODELFORGE_AGENT_ACCELERATOR_MODE=nvidia" in str(caught.value)
+
+
+def test_cpu_node_with_nvml_unavailable_is_allowed() -> None:
+    validate_nvidia_collection(
+        NvidiaCollection(
+            availability=Availability.UNAVAILABLE,
+            reason="NVMLError_LibraryNotFound",
+        ),
+        required=False,
+    )
+
+
+def test_nvidia_node_cannot_publish_empty_success_inventory() -> None:
+    with pytest.raises(AgentStartupError, match="NVIDIA_DEVICE_NOT_FOUND"):
+        validate_nvidia_collection(
+            NvidiaCollection(availability=Availability.KNOWN),
+            required=True,
+        )
+
+
+def test_nvidia_inventory_requires_matching_telemetry() -> None:
+    collection = gpu_collection().model_copy(update={"telemetry": []})
+    with pytest.raises(AgentStartupError, match="NVIDIA_TELEMETRY_UNAVAILABLE"):
+        validate_nvidia_collection(collection, required=True)