81 lines
2.8 KiB
Python
81 lines
2.8 KiB
Python
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)
|