Initial public ModelForge release

2026-09-01 21:30:16 +02:00
commit 7082ab955a
490 changed files with 104252 additions and 0 deletions
@@ -0,0 +1,385 @@
+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"