Initial public ModelForge release
This commit is contained in:
Jens
@@ -0,0 +1,320 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import shutil
|
||||
import struct
|
||||
import uuid
|
||||
from datetime import UTC, datetime, timedelta
|
||||
from pathlib import Path
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from modelforge_api.domain.acquisition import AgentArtifactJobFile, AgentArtifactJobLease
|
||||
from pydantic import ValidationError
|
||||
|
||||
from modelforge_node_agent.acquisition import (
|
||||
AcquisitionCancelled,
|
||||
AcquisitionFailure,
|
||||
ArtifactAcquirer,
|
||||
_confined,
|
||||
_inspect_file,
|
||||
_relative_path,
|
||||
)
|
||||
from modelforge_node_agent.settings import AgentSettings
|
||||
from modelforge_node_agent.transport import AgentTransport
|
||||
|
||||
|
||||
class FakeTransport:
|
||||
def __init__(self, *, cancel: bool = False) -> None:
|
||||
self.cancel = cancel
|
||||
self.progress: list[dict] = []
|
||||
self.completed: dict | None = None
|
||||
|
||||
def artifact_progress(self, _job_id: str, payload: dict, _credential: str) -> dict:
|
||||
self.progress.append(payload)
|
||||
return {
|
||||
"accepted": True,
|
||||
"cancel_requested": self.cancel and len(self.progress) > 1,
|
||||
"lease_expires_at": (datetime.now(UTC) + timedelta(seconds=120)).isoformat(),
|
||||
}
|
||||
|
||||
def artifact_complete(self, _job_id: str, payload: dict, _credential: str) -> dict:
|
||||
self.completed = payload
|
||||
return {"status": "completed"}
|
||||
|
||||
|
||||
def lease(target_root: Path, *, size: int = 8) -> AgentArtifactJobLease:
|
||||
return AgentArtifactJobLease(
|
||||
job_id=uuid.uuid4(),
|
||||
lease_token="x" * 40,
|
||||
lease_expires_at=datetime.now(UTC) + timedelta(minutes=2),
|
||||
repository_id="org/model",
|
||||
resolved_commit_sha="a" * 40,
|
||||
storage_root_id=uuid.uuid4(),
|
||||
target_root=str(target_root),
|
||||
total_size_bytes=size,
|
||||
reserve_bytes=0,
|
||||
reserve_percent=0,
|
||||
files=[
|
||||
AgentArtifactJobFile(
|
||||
ordinal=0,
|
||||
path="config.json",
|
||||
size_bytes=size,
|
||||
upstream_sha256=None,
|
||||
file_format="json",
|
||||
role="configuration",
|
||||
risk_flags=[],
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
def test_path_confinement_rejects_escape(tmp_path: Path) -> None:
|
||||
with pytest.raises(AcquisitionFailure, match="confined"):
|
||||
_relative_path("../secret")
|
||||
with pytest.raises(AcquisitionFailure, match="escaped"):
|
||||
_confined(tmp_path, tmp_path / ".." / "outside")
|
||||
|
||||
|
||||
def test_empty_job_poll_decodes_json_null_as_no_job() -> None:
|
||||
transport = AgentTransport("https://control.invalid", 5, True)
|
||||
transport.client = httpx.Client(
|
||||
base_url="https://control.invalid",
|
||||
transport=httpx.MockTransport(
|
||||
lambda request: httpx.Response(200, content=b"null", request=request)
|
||||
),
|
||||
)
|
||||
try:
|
||||
assert transport.next_artifact_job("credential") is None
|
||||
finally:
|
||||
transport.close()
|
||||
|
||||
|
||||
def test_typed_job_rejects_arbitrary_command_data(tmp_path: Path) -> None:
|
||||
payload = lease(tmp_path).model_dump(mode="json")
|
||||
payload["command"] = "sh -c anything"
|
||||
with pytest.raises(ValidationError, match="Extra inputs are not permitted"):
|
||||
AgentArtifactJobLease.model_validate(payload)
|
||||
|
||||
|
||||
def test_static_inspection_reads_header_without_deserializing(tmp_path: Path) -> None:
|
||||
header = json.dumps({"weight": {"dtype": "F32", "shape": [1], "data_offsets": [0, 4]}}).encode()
|
||||
target = tmp_path / "model.safetensors"
|
||||
target.write_bytes(struct.pack("<Q", len(header)) + header + b"\0\0\0\0")
|
||||
plan_file = AgentArtifactJobFile(
|
||||
ordinal=0,
|
||||
path=target.name,
|
||||
size_bytes=target.stat().st_size,
|
||||
upstream_sha256=None,
|
||||
file_format="safetensors",
|
||||
role="weights",
|
||||
risk_flags=[],
|
||||
)
|
||||
findings = _inspect_file(target, plan_file)
|
||||
assert findings[0]["type"] == "safetensors_header"
|
||||
assert findings[0]["status"] == "passed"
|
||||
|
||||
|
||||
def test_pickle_and_remote_code_are_blocking(tmp_path: Path) -> None:
|
||||
target = tmp_path / "weights.bin"
|
||||
target.write_bytes(b"not executed")
|
||||
plan_file = AgentArtifactJobFile(
|
||||
ordinal=0,
|
||||
path=target.name,
|
||||
size_bytes=target.stat().st_size,
|
||||
upstream_sha256=None,
|
||||
file_format="bin",
|
||||
role="weights",
|
||||
risk_flags=["pickle_or_executable_serialization"],
|
||||
)
|
||||
assert _inspect_file(target, plan_file)[0]["severity"] == "block"
|
||||
|
||||
|
||||
def test_auto_map_configuration_is_blocking_without_loading_code(tmp_path: Path) -> None:
|
||||
target = tmp_path / "config.json"
|
||||
target.write_text('{"auto_map":{"AutoModel":"modeling.Custom"}}', encoding="utf-8")
|
||||
plan_file = AgentArtifactJobFile(
|
||||
ordinal=0,
|
||||
path=target.name,
|
||||
size_bytes=target.stat().st_size,
|
||||
upstream_sha256=None,
|
||||
file_format="json",
|
||||
role="configuration",
|
||||
risk_flags=[],
|
||||
)
|
||||
assert _inspect_file(target, plan_file)[0]["type"] == "remote_code_configuration"
|
||||
assert _inspect_file(target, plan_file)[0]["severity"] == "block"
|
||||
|
||||
|
||||
def test_mocked_download_resumes_and_promotes_atomically(
|
||||
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
root = tmp_path / "artifacts" / "model-registry"
|
||||
job = lease(root)
|
||||
stage = root / ".quarantine" / str(job.job_id)
|
||||
stage.mkdir(parents=True)
|
||||
(stage / "config.json.part").write_bytes(b"1234")
|
||||
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
assert request.headers["range"] == "bytes=4-"
|
||||
return httpx.Response(206, content=b"5678", request=request)
|
||||
|
||||
original_client = httpx.Client
|
||||
monkeypatch.setattr(
|
||||
"modelforge_node_agent.acquisition.httpx.Client",
|
||||
lambda **kwargs: original_client(transport=httpx.MockTransport(handler), **kwargs),
|
||||
)
|
||||
transport = FakeTransport()
|
||||
settings = AgentSettings(
|
||||
artifact_path=tmp_path / "artifacts",
|
||||
artifact_download_timeout_seconds=10,
|
||||
)
|
||||
ArtifactAcquirer(settings, transport, "credential").execute(job.model_dump(mode="json"))
|
||||
final = root / "repositories" / "org--model" / ("a" * 40) / "config.json"
|
||||
assert final.read_bytes() == b"12345678"
|
||||
assert not stage.exists()
|
||||
assert transport.completed is not None
|
||||
assert transport.completed["files"][0]["sha256"]
|
||||
verifying = next(item for item in transport.progress if item["status"] == "verifying")
|
||||
assert verifying["progress_bytes"] == 8
|
||||
|
||||
|
||||
def test_download_rejects_bytes_beyond_declared_size_before_writing_them(
|
||||
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
root = tmp_path / "artifacts" / "model-registry"
|
||||
job = lease(root)
|
||||
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
return httpx.Response(200, content=b"123456789", request=request)
|
||||
|
||||
original_client = httpx.Client
|
||||
monkeypatch.setattr(
|
||||
"modelforge_node_agent.acquisition.httpx.Client",
|
||||
lambda **kwargs: original_client(transport=httpx.MockTransport(handler), **kwargs),
|
||||
)
|
||||
settings = AgentSettings(artifact_path=tmp_path / "artifacts")
|
||||
with pytest.raises(AcquisitionFailure, match="exceeded declared size") as failure:
|
||||
ArtifactAcquirer(settings, FakeTransport(), "credential").execute(
|
||||
job.model_dump(mode="json")
|
||||
)
|
||||
|
||||
assert failure.value.code == "download_size_mismatch"
|
||||
assert failure.value.retryable is False
|
||||
partial = root / ".quarantine" / str(job.job_id) / "config.json.part"
|
||||
assert partial.stat().st_size == 0
|
||||
|
||||
|
||||
def test_complete_partial_is_verified_without_another_request(
|
||||
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
root = tmp_path / "artifacts" / "model-registry"
|
||||
job = lease(root)
|
||||
stage = root / ".quarantine" / str(job.job_id)
|
||||
stage.mkdir(parents=True)
|
||||
(stage / "config.json.part").write_bytes(b"12345678")
|
||||
|
||||
def fail_if_called(_request: httpx.Request) -> httpx.Response:
|
||||
raise AssertionError("a complete partial must not trigger a network request")
|
||||
|
||||
original_client = httpx.Client
|
||||
monkeypatch.setattr(
|
||||
"modelforge_node_agent.acquisition.httpx.Client",
|
||||
lambda **kwargs: original_client(
|
||||
transport=httpx.MockTransport(fail_if_called), **kwargs
|
||||
),
|
||||
)
|
||||
transport = FakeTransport()
|
||||
settings = AgentSettings(artifact_path=tmp_path / "artifacts")
|
||||
ArtifactAcquirer(settings, transport, "credential").execute(job.model_dump(mode="json"))
|
||||
final = root / "repositories" / "org--model" / ("a" * 40) / "config.json"
|
||||
assert final.read_bytes() == b"12345678"
|
||||
assert transport.completed is not None
|
||||
|
||||
|
||||
def test_promoted_manifest_reconciles_without_redownload(
|
||||
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
root = tmp_path / "artifacts" / "model-registry"
|
||||
job = lease(root)
|
||||
final = root / "repositories" / "org--model" / ("a" * 40)
|
||||
final.mkdir(parents=True)
|
||||
(final / "config.json").write_bytes(b"12345678")
|
||||
(final / ".modelforge-manifest.json").write_text(
|
||||
json.dumps(
|
||||
{
|
||||
"repository_id": "org/model",
|
||||
"resolved_commit_sha": "a" * 40,
|
||||
"files": [],
|
||||
}
|
||||
),
|
||||
encoding="utf-8",
|
||||
)
|
||||
|
||||
def fail_if_called(_request: httpx.Request) -> httpx.Response:
|
||||
raise AssertionError("a promoted exact-revision set must not be downloaded again")
|
||||
|
||||
original_client = httpx.Client
|
||||
monkeypatch.setattr(
|
||||
"modelforge_node_agent.acquisition.httpx.Client",
|
||||
lambda **kwargs: original_client(
|
||||
transport=httpx.MockTransport(fail_if_called), **kwargs
|
||||
),
|
||||
)
|
||||
transport = FakeTransport()
|
||||
settings = AgentSettings(artifact_path=tmp_path / "artifacts")
|
||||
ArtifactAcquirer(settings, transport, "credential").execute(job.model_dump(mode="json"))
|
||||
assert transport.completed is not None
|
||||
assert transport.completed["promoted_relative_path"].endswith("a" * 40)
|
||||
|
||||
|
||||
def test_cancel_is_observed_before_promotion(tmp_path: Path) -> None:
|
||||
root = tmp_path / "artifacts" / "model-registry"
|
||||
job = lease(root, size=0)
|
||||
job.files[0].size_bytes = 0
|
||||
stage = root / ".quarantine" / str(job.job_id)
|
||||
stage.mkdir(parents=True)
|
||||
(stage / "config.json").write_bytes(b"")
|
||||
transport = FakeTransport(cancel=True)
|
||||
settings = AgentSettings(artifact_path=tmp_path / "artifacts")
|
||||
with pytest.raises(AcquisitionCancelled):
|
||||
ArtifactAcquirer(settings, transport, "credential").execute(job.model_dump(mode="json"))
|
||||
assert not (root / "repositories" / "org--model" / ("a" * 40)).exists()
|
||||
|
||||
|
||||
def test_execution_capacity_recheck_fails_before_network(
|
||||
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
root = tmp_path / "artifacts" / "model-registry"
|
||||
job = lease(root, size=500)
|
||||
monkeypatch.setattr(
|
||||
"modelforge_node_agent.acquisition.shutil.disk_usage",
|
||||
lambda _path: shutil._ntuple_diskusage(1000, 950, 50),
|
||||
)
|
||||
settings = AgentSettings(artifact_path=tmp_path / "artifacts")
|
||||
with pytest.raises(AcquisitionFailure, match="capacity preflight"):
|
||||
ArtifactAcquirer(settings, FakeTransport(), "credential").execute(
|
||||
job.model_dump(mode="json")
|
||||
)
|
||||
|
||||
|
||||
def test_size_mismatch_never_promotes_partial(
|
||||
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
root = tmp_path / "artifacts" / "model-registry"
|
||||
job = lease(root, size=8)
|
||||
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
return httpx.Response(200, content=b"short", request=request)
|
||||
|
||||
original_client = httpx.Client
|
||||
monkeypatch.setattr(
|
||||
"modelforge_node_agent.acquisition.httpx.Client",
|
||||
lambda **kwargs: original_client(transport=httpx.MockTransport(handler), **kwargs),
|
||||
)
|
||||
settings = AgentSettings(artifact_path=tmp_path / "artifacts")
|
||||
with pytest.raises(AcquisitionFailure, match="size mismatch"):
|
||||
ArtifactAcquirer(settings, FakeTransport(), "credential").execute(
|
||||
job.model_dump(mode="json")
|
||||
)
|
||||
assert not (root / "repositories" / "org--model" / ("a" * 40)).exists()
|
||||
@@ -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"
|
||||
@@ -0,0 +1,30 @@
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
import modelforge_node_agent.main as main_module
|
||||
from modelforge_node_agent.preflight import (
|
||||
AgentStartupError,
|
||||
AgentStartupFailureCode,
|
||||
AgentStartupProblem,
|
||||
)
|
||||
|
||||
|
||||
def test_cli_reports_typed_preflight_failure_with_exit_three(monkeypatch) -> None:
|
||||
async def fail_preflight(*, preflight_only: bool = False) -> None:
|
||||
assert preflight_only
|
||||
raise AgentStartupError(
|
||||
AgentStartupProblem(
|
||||
AgentStartupFailureCode.NVIDIA_NVML_UNAVAILABLE,
|
||||
"MODELFORGE_AGENT_ACCELERATOR_MODE=nvidia",
|
||||
"bounded diagnostic",
|
||||
)
|
||||
)
|
||||
|
||||
monkeypatch.setattr(main_module, "run_agent", fail_preflight)
|
||||
monkeypatch.setattr(sys, "argv", ["modelforge-node-agent", "--preflight"])
|
||||
|
||||
with pytest.raises(SystemExit) as caught:
|
||||
main_module.main()
|
||||
|
||||
assert caught.value.code == 3
|
||||
@@ -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)
|
||||
Reference in New Issue
Block a user