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"