386 lines
14 KiB
Python
386 lines
14 KiB
Python
import asyncio
|
|
import logging
|
|
from pathlib import Path
|
|
|
|
import httpx
|
|
import pytest
|
|
from modelforge_api.domain.enums import Availability
|
|
from modelforge_api.domain.hardware import (
|
|
AcceleratorInventory,
|
|
AcceleratorTelemetry,
|
|
HostInventory,
|
|
NvidiaCollection,
|
|
ObservedValue,
|
|
StorageObservation,
|
|
)
|
|
from pydantic import SecretStr
|
|
|
|
from modelforge_node_agent.agent import NodeAgent, describe_failure
|
|
from modelforge_node_agent.preflight import (
|
|
AgentObservationError,
|
|
AgentStartupError,
|
|
AgentStartupFailureCode,
|
|
)
|
|
from modelforge_node_agent.settings import AgentSettings
|
|
|
|
|
|
class HostCollector:
|
|
def collect(self) -> HostInventory:
|
|
return HostInventory(
|
|
identity_key="stable-agent-node",
|
|
identity_source="persisted_uuid",
|
|
hostname="server01",
|
|
display_name="Server 01",
|
|
os_name="Linux",
|
|
os_version=ObservedValue.known("1"),
|
|
architecture="x86_64",
|
|
kernel_version=ObservedValue.known("6.6"),
|
|
cpu_model=ObservedValue.known("CPU"),
|
|
logical_cpu_count=ObservedValue.known(16),
|
|
physical_core_count=ObservedValue.known(8),
|
|
total_ram_bytes=ObservedValue.known(64 * 1024**3),
|
|
available_ram_bytes=ObservedValue.known(32 * 1024**3),
|
|
agent_version="0.1.0",
|
|
storage=[
|
|
StorageObservation(
|
|
purpose="artifacts",
|
|
path="/data/artifacts",
|
|
total_bytes=ObservedValue.known(1000),
|
|
used_bytes=ObservedValue.known(100),
|
|
free_bytes=ObservedValue.known(900),
|
|
)
|
|
],
|
|
)
|
|
|
|
|
|
class AcceleratorCollector:
|
|
def collect(self) -> NvidiaCollection:
|
|
return NvidiaCollection(
|
|
availability=Availability.UNAVAILABLE,
|
|
reason="NVMLError_LibraryNotFound",
|
|
)
|
|
|
|
|
|
class CollectionCollector:
|
|
def __init__(self, collection: NvidiaCollection) -> None:
|
|
self.collection = collection
|
|
|
|
def collect(self) -> NvidiaCollection:
|
|
return self.collection
|
|
|
|
|
|
def gpu_collection() -> NvidiaCollection:
|
|
unavailable_int = ObservedValue[int].absent(Availability.UNSUPPORTED, "not exposed")
|
|
return NvidiaCollection(
|
|
availability=Availability.KNOWN,
|
|
inventory=[
|
|
AcceleratorInventory(
|
|
device_index=0,
|
|
device_uuid="GPU-MOCK-4080",
|
|
pci_bus_id=ObservedValue.known("0000:01:00.0"),
|
|
name="NVIDIA GeForce RTX 4080 SUPER",
|
|
architecture=ObservedValue.known("ada"),
|
|
compute_capability_major=ObservedValue.known(8),
|
|
compute_capability_minor=ObservedValue.known(9),
|
|
total_vram_bytes=ObservedValue.known(16 * 1024**3),
|
|
driver_version=ObservedValue.known("mock-driver"),
|
|
cuda_driver_version=ObservedValue.known("mock-cuda"),
|
|
mig_mode_current=ObservedValue.absent(Availability.UNSUPPORTED),
|
|
)
|
|
],
|
|
telemetry=[
|
|
AcceleratorTelemetry(
|
|
device_uuid="GPU-MOCK-4080",
|
|
used_vram_bytes=ObservedValue.known(1024**3),
|
|
free_vram_bytes=ObservedValue.known(15 * 1024**3),
|
|
gpu_utilization_percent=ObservedValue.known(12),
|
|
memory_utilization_percent=ObservedValue.known(7),
|
|
temperature_c=ObservedValue.known(42),
|
|
power_draw_w=ObservedValue.known(80.0),
|
|
power_limit_w=ObservedValue.known(320.0),
|
|
graphics_clock_mhz=unavailable_int,
|
|
memory_clock_mhz=unavailable_int,
|
|
fan_speed_percent=unavailable_int,
|
|
performance_state=ObservedValue.absent(Availability.UNSUPPORTED),
|
|
)
|
|
],
|
|
)
|
|
|
|
|
|
class FakeTransport:
|
|
def __init__(self, fail_heartbeat: bool = False) -> None:
|
|
self.fail_heartbeat = fail_heartbeat
|
|
self.enrollments: list[dict] = []
|
|
self.heartbeats: list[dict] = []
|
|
self.inventories: list[dict] = []
|
|
self.telemetry_samples: list[dict] = []
|
|
self.closed = False
|
|
|
|
def enroll(self, payload: dict) -> dict:
|
|
self.enrollments.append(payload)
|
|
return {
|
|
"node_id": "node-id",
|
|
"credential_id": "credential-id",
|
|
"node_credential": "mfnode_00000000-0000-0000-0000-000000000001_secret",
|
|
}
|
|
|
|
def heartbeat(self, payload: dict, credential: str) -> dict:
|
|
if self.fail_heartbeat:
|
|
raise httpx.ConnectError("temporary outage")
|
|
self.heartbeats.append(payload)
|
|
return {"accepted": True}
|
|
|
|
def inventory(self, payload: dict, credential: str) -> dict:
|
|
self.inventories.append(payload)
|
|
return {"accepted": True}
|
|
|
|
def telemetry(self, payload: dict, credential: str) -> dict:
|
|
self.telemetry_samples.append(payload)
|
|
return {"accepted": True}
|
|
|
|
def close(self) -> None:
|
|
self.closed = True
|
|
|
|
|
|
def agent_settings(tmp_path: Path, token: str | None = "e" * 40) -> AgentSettings:
|
|
return AgentSettings(
|
|
_env_file=None,
|
|
enrollment_token=SecretStr(token) if token else None,
|
|
identity="stable-agent-node",
|
|
accelerator_mode="cpu",
|
|
identity_file=tmp_path / "node-id",
|
|
credential_file=tmp_path / "credential",
|
|
state_file=tmp_path / "state.json",
|
|
heartbeat_interval_seconds=1,
|
|
inventory_interval_seconds=5,
|
|
telemetry_interval_seconds=5,
|
|
max_backoff_seconds=2,
|
|
)
|
|
|
|
|
|
def build_agent(
|
|
tmp_path: Path, transport: FakeTransport, token: str | None = "e" * 40
|
|
) -> NodeAgent:
|
|
return NodeAgent(
|
|
agent_settings(tmp_path, token),
|
|
transport=transport,
|
|
host_collector=HostCollector(),
|
|
accelerator_collector=AcceleratorCollector(),
|
|
)
|
|
|
|
|
|
def test_first_enrollment_and_restart_reuse_persisted_credential(tmp_path: Path) -> None:
|
|
first_transport = FakeTransport()
|
|
first = build_agent(tmp_path, first_transport)
|
|
first.ensure_enrolled()
|
|
assert len(first_transport.enrollments) == 1
|
|
assert first.settings.credential_file.read_text(encoding="utf-8").startswith("mfnode_")
|
|
second_transport = FakeTransport()
|
|
restarted = build_agent(tmp_path, second_transport, token=None)
|
|
restarted.ensure_enrolled()
|
|
assert second_transport.enrollments == []
|
|
assert restarted.credential == first.credential
|
|
|
|
|
|
def test_agent_advertises_typed_runtime_worker_capabilities(tmp_path: Path) -> None:
|
|
capabilities = build_agent(tmp_path, FakeTransport()).metadata().supported_capabilities
|
|
assert "runtime.probe.v1" in capabilities
|
|
assert "runtime.health.v1" in capabilities
|
|
assert "runtime.unload.v1" in capabilities
|
|
assert not any("shell" in capability or "command" in capability for capability in capabilities)
|
|
|
|
|
|
def test_heartbeat_inventory_telemetry_and_sequences_survive_restart(tmp_path: Path) -> None:
|
|
transport = FakeTransport()
|
|
agent = build_agent(tmp_path, transport)
|
|
agent.heartbeat_once()
|
|
agent.inventory_once()
|
|
agent.telemetry_once()
|
|
assert transport.heartbeats[0]["identity_key"] == "stable-agent-node"
|
|
assert transport.inventories[0]["nvidia"]["availability"] == "unavailable"
|
|
assert transport.telemetry_samples[0]["sequence"] == 1
|
|
restarted = build_agent(tmp_path, FakeTransport(), token=None)
|
|
restarted.inventory_once()
|
|
restarted.telemetry_once()
|
|
assert restarted.state.inventory_sequence == 2
|
|
assert restarted.state.telemetry_sequence == 2
|
|
|
|
|
|
def test_temporary_outage_uses_bounded_retry_and_clean_shutdown(tmp_path: Path, caplog) -> None:
|
|
transport = FakeTransport(fail_heartbeat=True)
|
|
agent = build_agent(tmp_path, transport)
|
|
|
|
async def exercise() -> None:
|
|
task = asyncio.create_task(agent.run())
|
|
await asyncio.sleep(0.05)
|
|
agent.stop()
|
|
await asyncio.wait_for(task, timeout=1)
|
|
|
|
with caplog.at_level(logging.WARNING):
|
|
asyncio.run(exercise())
|
|
assert "retrying" in caplog.text
|
|
assert "e" * 40 not in caplog.text
|
|
assert "mfnode_" not in caplog.text
|
|
agent.close()
|
|
assert transport.closed
|
|
|
|
|
|
def test_zero_gpu_inventory_is_a_successful_empty_observation(tmp_path: Path) -> None:
|
|
transport = FakeTransport()
|
|
agent = NodeAgent(
|
|
agent_settings(tmp_path),
|
|
transport=transport,
|
|
host_collector=HostCollector(),
|
|
accelerator_collector=CollectionCollector(
|
|
NvidiaCollection(availability=Availability.KNOWN)
|
|
),
|
|
)
|
|
agent.inventory_once()
|
|
assert transport.inventories[0]["nvidia"]["availability"] == "known"
|
|
assert transport.inventories[0]["nvidia"]["inventory"] == []
|
|
|
|
|
|
def test_mock_nvml_gpu_inventory_and_telemetry_are_published(tmp_path: Path) -> None:
|
|
transport = FakeTransport()
|
|
agent = NodeAgent(
|
|
agent_settings(tmp_path),
|
|
transport=transport,
|
|
host_collector=HostCollector(),
|
|
accelerator_collector=CollectionCollector(gpu_collection()),
|
|
)
|
|
agent.inventory_once()
|
|
agent.telemetry_once()
|
|
assert transport.inventories[0]["nvidia"]["inventory"][0]["device_uuid"] == "GPU-MOCK-4080"
|
|
assert transport.telemetry_samples[0]["accelerators"][0]["temperature_c"]["value"] == 42
|
|
|
|
|
|
def test_nvidia_failure_refuses_before_enrollment_or_publication(tmp_path: Path) -> None:
|
|
transport = FakeTransport()
|
|
settings = agent_settings(tmp_path)
|
|
settings.accelerator_mode = "nvidia"
|
|
agent = NodeAgent(
|
|
settings,
|
|
transport=transport,
|
|
host_collector=HostCollector(),
|
|
accelerator_collector=AcceleratorCollector(),
|
|
)
|
|
|
|
try:
|
|
agent.inventory_once()
|
|
except AgentStartupError as exc:
|
|
assert exc.problem.code is AgentStartupFailureCode.NVIDIA_NVML_UNAVAILABLE
|
|
else:
|
|
raise AssertionError("an NVIDIA node must refuse unavailable NVML")
|
|
assert transport.enrollments == []
|
|
assert transport.inventories == []
|
|
assert not settings.identity_file.exists()
|
|
|
|
|
|
def test_nvidia_collection_passes_preflight_and_keeps_protocol_one(tmp_path: Path) -> None:
|
|
settings = agent_settings(tmp_path)
|
|
settings.accelerator_mode = "nvidia"
|
|
agent = NodeAgent(
|
|
settings,
|
|
transport=FakeTransport(),
|
|
host_collector=HostCollector(),
|
|
accelerator_collector=CollectionCollector(gpu_collection()),
|
|
)
|
|
|
|
collection = agent.preflight()
|
|
|
|
assert collection.inventory[0].device_uuid == "GPU-MOCK-4080"
|
|
assert collection.telemetry[0].used_vram_bytes.value == 1024**3
|
|
assert agent.metadata().protocol_version == 1
|
|
|
|
|
|
def test_direct_run_preflights_before_enrollment(tmp_path: Path) -> None:
|
|
transport = FakeTransport()
|
|
settings = agent_settings(tmp_path)
|
|
settings.accelerator_mode = "nvidia"
|
|
agent = NodeAgent(
|
|
settings,
|
|
transport=transport,
|
|
host_collector=HostCollector(),
|
|
accelerator_collector=AcceleratorCollector(),
|
|
)
|
|
|
|
with pytest.raises(AgentStartupError, match="NVIDIA_NVML_UNAVAILABLE"):
|
|
asyncio.run(agent.run())
|
|
|
|
assert transport.enrollments == []
|
|
assert transport.heartbeats == []
|
|
|
|
|
|
def test_temporary_nvml_failure_is_not_published_after_preflight(tmp_path: Path) -> None:
|
|
transport = FakeTransport()
|
|
settings = agent_settings(tmp_path)
|
|
settings.accelerator_mode = "nvidia"
|
|
collector = CollectionCollector(gpu_collection())
|
|
agent = NodeAgent(
|
|
settings,
|
|
transport=transport,
|
|
host_collector=HostCollector(),
|
|
accelerator_collector=collector,
|
|
)
|
|
agent.preflight()
|
|
collector.collection = NvidiaCollection(
|
|
availability=Availability.TEMPORARILY_FAILED,
|
|
reason="device 0: NVMLError_Unknown",
|
|
)
|
|
|
|
with pytest.raises(AgentObservationError, match="NVIDIA_TELEMETRY_UNAVAILABLE"):
|
|
agent.telemetry_once()
|
|
|
|
assert transport.enrollments == []
|
|
assert transport.telemetry_samples == []
|
|
|
|
|
|
def test_a_failed_publication_names_the_status_and_path_not_just_the_exception() -> None:
|
|
"""M16 operability regression.
|
|
|
|
The agent logged a bare 'control-plane publication failed; retrying' for two days while a
|
|
credential was rejected. `type(exc).__name__` cannot tell an operator whether the cause is a
|
|
revoked credential or a DNS failure; the status code, method and path can.
|
|
"""
|
|
|
|
request = httpx.Request("POST", "http://api:8000/api/v1/agent/heartbeat")
|
|
response = httpx.Response(401, request=request)
|
|
described = describe_failure(
|
|
httpx.HTTPStatusError("unauthorised", request=request, response=response)
|
|
)
|
|
assert "401" in described
|
|
assert "POST" in described
|
|
assert "/api/v1/agent/heartbeat" in described
|
|
|
|
|
|
def test_a_transport_failure_names_the_endpoint_it_could_not_reach() -> None:
|
|
request = httpx.Request("PUT", "http://api:8000/api/v1/agent/inventory")
|
|
described = describe_failure(httpx.ConnectError("no route", request=request))
|
|
assert "ConnectError" in described
|
|
assert "/api/v1/agent/inventory" in described
|
|
|
|
|
|
def test_a_non_http_failure_still_produces_a_bounded_description() -> None:
|
|
assert describe_failure(ValueError("boom")) == "ValueError"
|
|
assert describe_failure(OSError("disk gone")) == "OSError"
|
|
assert describe_failure(httpx.ConnectError("no request attached")) == "ConnectError"
|
|
|
|
|
|
def test_a_failure_description_never_carries_the_response_body_or_a_credential() -> None:
|
|
"""A diagnostic that echoes the body can leak exactly what the agent was publishing."""
|
|
|
|
request = httpx.Request(
|
|
"POST",
|
|
"http://api:8000/api/v1/agent/heartbeat",
|
|
headers={"Authorization": "Bearer mfnode_super_secret_value"},
|
|
)
|
|
response = httpx.Response(
|
|
403, request=request, text="forbidden: credential mfnode_super_secret_value"
|
|
)
|
|
described = describe_failure(
|
|
httpx.HTTPStatusError("forbidden", request=request, response=response)
|
|
)
|
|
assert "mfnode_super_secret_value" not in described
|
|
assert "forbidden" not in described
|
|
assert described == "HTTPStatusError 403 on POST /api/v1/agent/heartbeat"
|