Initial public ModelForge release
This commit is contained in:
Jens
@@ -0,0 +1,52 @@
|
||||
FROM python:3.12-slim-trixie@sha256:09f7da3bc104798d0afb40bc08d23ab2da20a76130cec1f2ef170848f5d85217 AS builder
|
||||
|
||||
ENV PYTHONDONTWRITEBYTECODE=1 \
|
||||
PYTHONUNBUFFERED=1
|
||||
|
||||
WORKDIR /build
|
||||
COPY backend ./backend
|
||||
COPY node-agent ./node-agent
|
||||
RUN python -m venv /opt/venv \
|
||||
&& /opt/venv/bin/pip install --no-cache-dir --upgrade pip "setuptools>=78.1.1" "msgpack>=1.2.1" \
|
||||
&& /opt/venv/bin/pip install --no-cache-dir ./backend ./node-agent \
|
||||
&& /opt/venv/bin/pip check \
|
||||
&& /opt/venv/bin/python -m pip uninstall --yes pip setuptools
|
||||
|
||||
FROM python:3.12-slim-trixie@sha256:09f7da3bc104798d0afb40bc08d23ab2da20a76130cec1f2ef170848f5d85217
|
||||
|
||||
ENV PYTHONDONTWRITEBYTECODE=1 \
|
||||
PYTHONUNBUFFERED=1 \
|
||||
PATH="/opt/venv/bin:${PATH}"
|
||||
|
||||
RUN apt-get update \
|
||||
&& DEBIAN_FRONTEND=noninteractive apt-get upgrade --yes \
|
||||
&& rm -rf /var/lib/apt/lists/* \
|
||||
&& /usr/local/bin/python -m pip uninstall --yes pip setuptools
|
||||
|
||||
# Build identity. These are stamped in at build time so a running container can say exactly
|
||||
# where it came from; an argument that is never passed stays empty and is reported as null
|
||||
# rather than becoming a claimed commit.
|
||||
ARG MODELFORGE_VERSION=0.0.0
|
||||
ARG MODELFORGE_COMMIT=""
|
||||
ARG MODELFORGE_BUILT_AT=""
|
||||
ENV MODELFORGE_BUILD_COMMIT=${MODELFORGE_COMMIT}
|
||||
ENV MODELFORGE_BUILD_TIMESTAMP=${MODELFORGE_BUILT_AT}
|
||||
LABEL org.opencontainers.image.title="ITWorx ModelForge node agent"
|
||||
LABEL org.opencontainers.image.description="Outbound-only compute node agent for ITWorx ModelForge"
|
||||
LABEL org.opencontainers.image.version="${MODELFORGE_VERSION}"
|
||||
LABEL org.opencontainers.image.revision="${MODELFORGE_COMMIT}"
|
||||
LABEL org.opencontainers.image.created="${MODELFORGE_BUILT_AT}"
|
||||
LABEL org.opencontainers.image.source="https://git.example.com/example/modelforge.git"
|
||||
LABEL org.opencontainers.image.vendor="ITWorx"
|
||||
LABEL org.opencontainers.image.licenses="AGPL-3.0-or-later"
|
||||
|
||||
COPY --from=builder /opt/venv /opt/venv
|
||||
|
||||
RUN groupadd --system --gid 101 modelforge-agent \
|
||||
&& useradd --system --uid 100 --gid 101 --home-dir /app --create-home modelforge-agent \
|
||||
&& mkdir -p /data/state /data/hf-cache /data/artifacts /data/quarantine \
|
||||
&& chown -R modelforge-agent:modelforge-agent /data
|
||||
WORKDIR /app
|
||||
USER modelforge-agent
|
||||
|
||||
CMD ["modelforge-node-agent"]
|
||||
@@ -0,0 +1,54 @@
|
||||
[build-system]
|
||||
requires = ["setuptools>=80"]
|
||||
build-backend = "setuptools.build_meta"
|
||||
|
||||
[project]
|
||||
name = "modelforge-node-agent"
|
||||
version = "1.2.2"
|
||||
license = "AGPL-3.0-or-later"
|
||||
requires-python = ">=3.12"
|
||||
dependencies = [
|
||||
"certifi==2026.7.22",
|
||||
"httpx==0.28.1",
|
||||
# Ships as one product with the control plane, so this tracks VERSION exactly. A
|
||||
# stale pin here made the Node Agent image unbuildable the moment the product
|
||||
# version moved, which is a release failure that only appears at build time.
|
||||
"modelforge-api==1.2.2",
|
||||
"pydantic-settings==2.15.0",
|
||||
]
|
||||
|
||||
[project.optional-dependencies]
|
||||
dev = [
|
||||
"pytest==8.4.2",
|
||||
"ruff==0.16.5",
|
||||
"mypy==1.20.2",
|
||||
]
|
||||
|
||||
[project.scripts]
|
||||
modelforge-node-agent = "modelforge_node_agent.main:main"
|
||||
|
||||
[tool.setuptools.packages.find]
|
||||
where = ["src"]
|
||||
|
||||
[tool.pytest.ini_options]
|
||||
pythonpath = ["src", "../backend/src"]
|
||||
testpaths = ["tests"]
|
||||
|
||||
[tool.mypy]
|
||||
# The images run on Linux, so that is the platform the type check must analyse. Checking on the
|
||||
# developer's platform meant Windows-only branches were analysed and Linux-only ones were not —
|
||||
# a winreg block type-checked cleanly here and failed in CI on identical code.
|
||||
platform = "linux"
|
||||
python_version = "3.12"
|
||||
strict = true
|
||||
mypy_path = ["../backend/src"]
|
||||
|
||||
[tool.ruff]
|
||||
target-version = "py312"
|
||||
line-length = 100
|
||||
|
||||
[tool.ruff.lint]
|
||||
select = ["E", "F", "I", "B", "UP", "S"]
|
||||
|
||||
[tool.ruff.lint.per-file-ignores]
|
||||
"tests/*" = ["S101"]
|
||||
@@ -0,0 +1 @@
|
||||
__version__ = "1.2.2"
|
||||
@@ -0,0 +1,372 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
import shutil
|
||||
import ssl
|
||||
import struct
|
||||
from pathlib import Path, PurePosixPath
|
||||
from typing import Any, Literal
|
||||
from urllib.parse import quote
|
||||
|
||||
import certifi
|
||||
import httpx
|
||||
from modelforge_api.domain.acquisition import (
|
||||
AgentArtifactJobLease,
|
||||
AgentJobComplete,
|
||||
AgentJobFailure,
|
||||
AgentJobProgress,
|
||||
CompletedFile,
|
||||
)
|
||||
|
||||
from modelforge_node_agent.settings import AgentSettings
|
||||
|
||||
|
||||
class AcquisitionFailure(RuntimeError):
|
||||
def __init__(self, code: str, message: str, *, retryable: bool = False) -> None:
|
||||
super().__init__(message)
|
||||
self.code = code
|
||||
self.retryable = retryable
|
||||
|
||||
|
||||
class AcquisitionCancelled(AcquisitionFailure):
|
||||
def __init__(self) -> None:
|
||||
super().__init__("cancelled", "artifact acquisition was cancelled")
|
||||
|
||||
|
||||
def _confined(base: Path, candidate: Path) -> Path:
|
||||
resolved_base = base.resolve()
|
||||
resolved = candidate.resolve()
|
||||
if resolved != resolved_base and resolved_base not in resolved.parents:
|
||||
raise AcquisitionFailure("path_escape", "artifact path escaped its approved storage root")
|
||||
return resolved
|
||||
|
||||
|
||||
def _relative_path(value: str) -> PurePosixPath:
|
||||
path = PurePosixPath(value.replace("\\", "/"))
|
||||
if path.is_absolute() or not path.parts or ".." in path.parts:
|
||||
raise AcquisitionFailure("invalid_path", "upstream file path is not confined")
|
||||
return path
|
||||
|
||||
|
||||
def _hash_file(path: Path) -> tuple[str, int]:
|
||||
digest = hashlib.sha256()
|
||||
size = 0
|
||||
with path.open("rb") as handle:
|
||||
for chunk in iter(lambda: handle.read(1024 * 1024), b""):
|
||||
digest.update(chunk)
|
||||
size += len(chunk)
|
||||
return digest.hexdigest(), size
|
||||
|
||||
|
||||
def _inspect_file(path: Path, plan_file: Any) -> list[dict[str, Any]]:
|
||||
findings: list[dict[str, Any]] = []
|
||||
if path.is_symlink():
|
||||
return [{"type": "symlink", "status": "blocked", "severity": "block", "evidence": {}}]
|
||||
risks = set(plan_file.risk_flags)
|
||||
if "pickle_or_executable_serialization" in risks:
|
||||
findings.append(
|
||||
{
|
||||
"type": "serialization",
|
||||
"status": "blocked",
|
||||
"severity": "block",
|
||||
"evidence": {"format": plan_file.file_format},
|
||||
}
|
||||
)
|
||||
if plan_file.file_format == "python" or "remote_code" in risks:
|
||||
findings.append(
|
||||
{
|
||||
"type": "remote_code",
|
||||
"status": "blocked",
|
||||
"severity": "block",
|
||||
"evidence": {"trust_remote_code": False},
|
||||
}
|
||||
)
|
||||
if plan_file.file_format == "safetensors":
|
||||
with path.open("rb") as handle:
|
||||
raw = handle.read(8)
|
||||
if len(raw) != 8:
|
||||
raise AcquisitionFailure(
|
||||
"invalid_safetensors", f"truncated safetensors header: {plan_file.path}"
|
||||
)
|
||||
header_size = struct.unpack("<Q", raw)[0]
|
||||
if header_size > 100 * 1024 * 1024 or header_size > path.stat().st_size - 8:
|
||||
raise AcquisitionFailure(
|
||||
"invalid_safetensors", f"unsafe safetensors header: {plan_file.path}"
|
||||
)
|
||||
try:
|
||||
header = json.loads(handle.read(header_size))
|
||||
except (UnicodeDecodeError, json.JSONDecodeError) as exc:
|
||||
raise AcquisitionFailure(
|
||||
"invalid_safetensors", f"invalid safetensors JSON header: {plan_file.path}"
|
||||
) from exc
|
||||
if not isinstance(header, dict):
|
||||
raise AcquisitionFailure(
|
||||
"invalid_safetensors", f"invalid safetensors structure: {plan_file.path}"
|
||||
)
|
||||
findings.append(
|
||||
{
|
||||
"type": "safetensors_header",
|
||||
"status": "passed",
|
||||
"severity": "info",
|
||||
"evidence": {"header_bytes": header_size, "tensor_entries": len(header)},
|
||||
}
|
||||
)
|
||||
if plan_file.file_format == "json" and path.stat().st_size <= 16 * 1024 * 1024:
|
||||
try:
|
||||
payload = json.loads(path.read_text(encoding="utf-8"))
|
||||
except (UnicodeDecodeError, json.JSONDecodeError):
|
||||
payload = None
|
||||
if isinstance(payload, dict) and payload.get("auto_map"):
|
||||
findings.append(
|
||||
{
|
||||
"type": "remote_code_configuration",
|
||||
"status": "blocked",
|
||||
"severity": "block",
|
||||
"evidence": {"auto_map_present": True, "trust_remote_code": False},
|
||||
}
|
||||
)
|
||||
if not findings:
|
||||
findings.append(
|
||||
{
|
||||
"type": "file_policy",
|
||||
"status": "passed",
|
||||
"severity": "info",
|
||||
"evidence": {"format": plan_file.file_format},
|
||||
}
|
||||
)
|
||||
return findings
|
||||
|
||||
|
||||
class ArtifactAcquirer:
|
||||
def __init__(self, settings: AgentSettings, transport: Any, credential: str) -> None:
|
||||
self.settings = settings
|
||||
self.transport = transport
|
||||
self.credential = credential
|
||||
|
||||
def _progress(
|
||||
self,
|
||||
lease: AgentArtifactJobLease,
|
||||
status: Literal["claimed", "downloading", "verifying", "promoting"],
|
||||
progress_bytes: int,
|
||||
current_file: str | None,
|
||||
quarantine_relative_path: str | None,
|
||||
) -> None:
|
||||
response = self.transport.artifact_progress(
|
||||
str(lease.job_id),
|
||||
AgentJobProgress(
|
||||
lease_token=lease.lease_token,
|
||||
status=status,
|
||||
progress_bytes=progress_bytes,
|
||||
current_file=current_file,
|
||||
quarantine_relative_path=quarantine_relative_path,
|
||||
).model_dump(mode="json"),
|
||||
self.credential,
|
||||
)
|
||||
if response.get("cancel_requested"):
|
||||
raise AcquisitionCancelled()
|
||||
|
||||
def _download(
|
||||
self,
|
||||
client: httpx.Client,
|
||||
lease: AgentArtifactJobLease,
|
||||
plan_file: Any,
|
||||
destination: Path,
|
||||
progress_before: int,
|
||||
quarantine_relative_path: str,
|
||||
) -> None:
|
||||
destination.parent.mkdir(parents=True, exist_ok=True)
|
||||
partial = destination.with_suffix(destination.suffix + ".part")
|
||||
offset = partial.stat().st_size if partial.is_file() else 0
|
||||
if offset > plan_file.size_bytes:
|
||||
partial.unlink()
|
||||
offset = 0
|
||||
elif offset == plan_file.size_bytes:
|
||||
# A transport may fail while closing a response after every expected byte was
|
||||
# durably written. Treat the complete partial as resumable input and let the
|
||||
# normal verification phase decide whether its digest is acceptable.
|
||||
os.replace(partial, destination)
|
||||
return
|
||||
headers: dict[str, str] = {}
|
||||
if offset:
|
||||
headers["Range"] = f"bytes={offset}-"
|
||||
token = self.settings.hf_token
|
||||
if token and token.get_secret_value():
|
||||
headers["Authorization"] = f"Bearer {token.get_secret_value()}"
|
||||
url = (
|
||||
"https://huggingface.co/"
|
||||
f"{quote(lease.repository_id, safe='/')}/resolve/"
|
||||
f"{quote(lease.resolved_commit_sha, safe='')}/{quote(plan_file.path, safe='/')}"
|
||||
)
|
||||
with client.stream("GET", url, headers=headers) as response:
|
||||
response.raise_for_status()
|
||||
append = offset > 0 and response.status_code == 206
|
||||
if not append:
|
||||
offset = 0
|
||||
mode = "ab" if append else "wb"
|
||||
last_report = offset
|
||||
with partial.open(mode) as handle:
|
||||
for chunk in response.iter_bytes(1024 * 1024):
|
||||
if offset + len(chunk) > plan_file.size_bytes:
|
||||
raise AcquisitionFailure(
|
||||
"download_size_mismatch",
|
||||
f"download exceeded declared size for {plan_file.path}",
|
||||
retryable=False,
|
||||
)
|
||||
handle.write(chunk)
|
||||
offset += len(chunk)
|
||||
if offset - last_report >= 8 * 1024 * 1024:
|
||||
self._progress(
|
||||
lease,
|
||||
"downloading",
|
||||
progress_before + offset,
|
||||
plan_file.path,
|
||||
quarantine_relative_path,
|
||||
)
|
||||
last_report = offset
|
||||
handle.flush()
|
||||
os.fsync(handle.fileno())
|
||||
if offset != plan_file.size_bytes:
|
||||
raise AcquisitionFailure(
|
||||
"download_size_mismatch",
|
||||
f"downloaded size mismatch for {plan_file.path}",
|
||||
retryable=True,
|
||||
)
|
||||
os.replace(partial, destination)
|
||||
|
||||
def execute(self, raw_lease: dict[str, Any]) -> None:
|
||||
lease = AgentArtifactJobLease.model_validate(raw_lease)
|
||||
configured_root = self.settings.artifact_path.resolve()
|
||||
root = _confined(configured_root, Path(lease.target_root))
|
||||
root.mkdir(parents=True, exist_ok=True)
|
||||
if root.is_symlink():
|
||||
raise AcquisitionFailure("symlink_root", "approved artifact root is a symlink")
|
||||
usage = shutil.disk_usage(root)
|
||||
reserve = max(lease.reserve_bytes, usage.total * lease.reserve_percent // 100)
|
||||
if lease.total_size_bytes > max(0, usage.free - reserve):
|
||||
raise AcquisitionFailure("capacity_changed", "execution capacity preflight failed")
|
||||
safe_repo = lease.repository_id.replace("/", "--").replace("\\", "--")
|
||||
final = _confined(root, root / "repositories" / safe_repo / lease.resolved_commit_sha)
|
||||
stage = _confined(root, root / ".quarantine" / str(lease.job_id))
|
||||
reconcile_promoted = final.exists()
|
||||
if reconcile_promoted:
|
||||
existing_manifest = final / ".modelforge-manifest.json"
|
||||
if not existing_manifest.is_file():
|
||||
raise AcquisitionFailure(
|
||||
"promotion_conflict", "existing target has no ModelForge manifest"
|
||||
)
|
||||
existing = json.loads(existing_manifest.read_text(encoding="utf-8"))
|
||||
if (
|
||||
existing.get("repository_id") != lease.repository_id
|
||||
or existing.get("resolved_commit_sha") != lease.resolved_commit_sha
|
||||
):
|
||||
raise AcquisitionFailure("promotion_conflict", "existing target provenance differs")
|
||||
else:
|
||||
stage.mkdir(parents=True, exist_ok=True)
|
||||
worktree = final if reconcile_promoted else stage
|
||||
quarantine_relative = stage.relative_to(root).as_posix()
|
||||
self._progress(lease, "claimed", 0, None, quarantine_relative)
|
||||
completed: list[CompletedFile] = []
|
||||
progress = 0
|
||||
with httpx.Client(
|
||||
timeout=self.settings.artifact_download_timeout_seconds,
|
||||
follow_redirects=True,
|
||||
verify=ssl.create_default_context(cafile=certifi.where()),
|
||||
) as client:
|
||||
for plan_file in lease.files:
|
||||
relative = _relative_path(plan_file.path)
|
||||
destination = _confined(worktree, worktree.joinpath(*relative.parts))
|
||||
if not destination.is_file():
|
||||
if reconcile_promoted:
|
||||
raise AcquisitionFailure(
|
||||
"promotion_conflict",
|
||||
f"promoted set is incomplete: {plan_file.path}",
|
||||
)
|
||||
self._download(
|
||||
client,
|
||||
lease,
|
||||
plan_file,
|
||||
destination,
|
||||
progress,
|
||||
quarantine_relative,
|
||||
)
|
||||
self._progress(
|
||||
lease,
|
||||
"verifying",
|
||||
progress + plan_file.size_bytes,
|
||||
plan_file.path,
|
||||
quarantine_relative,
|
||||
)
|
||||
digest, size = _hash_file(destination)
|
||||
if size != plan_file.size_bytes:
|
||||
raise AcquisitionFailure(
|
||||
"verification_size_mismatch", f"size mismatch for {plan_file.path}"
|
||||
)
|
||||
if plan_file.upstream_sha256 and digest != plan_file.upstream_sha256:
|
||||
raise AcquisitionFailure(
|
||||
"verification_digest_mismatch", f"digest mismatch for {plan_file.path}"
|
||||
)
|
||||
inspections = _inspect_file(destination, plan_file)
|
||||
if any(item["severity"] == "block" for item in inspections):
|
||||
raise AcquisitionFailure(
|
||||
"static_inspection_blocked", f"static inspection blocked {plan_file.path}"
|
||||
)
|
||||
progress += size
|
||||
completed.append(
|
||||
CompletedFile(
|
||||
path=plan_file.path,
|
||||
relative_path="pending",
|
||||
size_bytes=size,
|
||||
sha256=digest,
|
||||
inspections=inspections,
|
||||
)
|
||||
)
|
||||
self._progress(lease, "promoting", progress, None, quarantine_relative)
|
||||
if not reconcile_promoted:
|
||||
final.parent.mkdir(parents=True, exist_ok=True)
|
||||
manifest = {
|
||||
"job_id": str(lease.job_id),
|
||||
"repository_id": lease.repository_id,
|
||||
"resolved_commit_sha": lease.resolved_commit_sha,
|
||||
"files": [item.model_dump(mode="json") for item in completed],
|
||||
}
|
||||
(stage / ".modelforge-manifest.json").write_text(
|
||||
json.dumps(manifest, sort_keys=True, indent=2), encoding="utf-8"
|
||||
)
|
||||
os.replace(stage, final)
|
||||
promoted_base = final.relative_to(root).as_posix()
|
||||
promoted_files = [
|
||||
item.model_copy(update={"relative_path": f"{promoted_base}/{item.path}"})
|
||||
for item in completed
|
||||
]
|
||||
after = shutil.disk_usage(root)
|
||||
self.transport.artifact_complete(
|
||||
str(lease.job_id),
|
||||
AgentJobComplete(
|
||||
lease_token=lease.lease_token,
|
||||
promoted_relative_path=promoted_base,
|
||||
capacity_observation={
|
||||
"capacity_bytes": after.total,
|
||||
"free_bytes": after.free,
|
||||
"reserve_bytes": reserve,
|
||||
"agent_path": str(root),
|
||||
},
|
||||
files=promoted_files,
|
||||
).model_dump(mode="json"),
|
||||
self.credential,
|
||||
)
|
||||
|
||||
def fail(self, raw_lease: dict[str, Any], error: AcquisitionFailure) -> None:
|
||||
lease = AgentArtifactJobLease.model_validate(raw_lease)
|
||||
self.transport.artifact_fail(
|
||||
str(lease.job_id),
|
||||
AgentJobFailure(
|
||||
lease_token=lease.lease_token,
|
||||
error_code=error.code,
|
||||
error_message=str(error),
|
||||
retryable=error.retryable,
|
||||
details={"exception_type": type(error).__name__},
|
||||
).model_dump(mode="json"),
|
||||
self.credential,
|
||||
)
|
||||
@@ -0,0 +1,309 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import threading
|
||||
import time
|
||||
from datetime import UTC, datetime
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
from modelforge_api.domain.agent_protocol import (
|
||||
AGENT_PROTOCOL_CAPABILITIES,
|
||||
AGENT_PROTOCOL_VERSION,
|
||||
AgentMetadata,
|
||||
EnrollmentRequest,
|
||||
HeartbeatRequest,
|
||||
InventoryNvidiaPayload,
|
||||
InventoryReport,
|
||||
TelemetryReport,
|
||||
)
|
||||
from modelforge_api.domain.hardware import NvidiaCollection
|
||||
from modelforge_api.hardware.collectors import (
|
||||
NodeIdentityProvider,
|
||||
SystemHostCollector,
|
||||
build_nvml_collector,
|
||||
)
|
||||
|
||||
from modelforge_node_agent import __version__
|
||||
from modelforge_node_agent.acquisition import AcquisitionFailure, ArtifactAcquirer
|
||||
from modelforge_node_agent.preflight import (
|
||||
AgentObservationError,
|
||||
AgentStartupError,
|
||||
nvidia_required,
|
||||
validate_nvidia_collection,
|
||||
)
|
||||
from modelforge_node_agent.settings import AgentSettings
|
||||
from modelforge_node_agent.state import AgentPersistentState, AgentStateStore
|
||||
from modelforge_node_agent.transport import AgentTransport
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def describe_failure(error: BaseException) -> str:
|
||||
"""A bounded, non-secret description an operator can act on.
|
||||
|
||||
`type(error).__name__` alone cannot distinguish a revoked credential from a DNS failure, which
|
||||
is exactly the question an operator asks when an agent stops publishing. The status code, the
|
||||
method and the path answer it. The response body and the credential never appear: the body can
|
||||
echo request content and the credential is the secret itself.
|
||||
"""
|
||||
|
||||
if isinstance(error, AgentObservationError):
|
||||
return str(error.problem)
|
||||
if isinstance(error, httpx.HTTPStatusError):
|
||||
request = error.request
|
||||
return (
|
||||
f"{type(error).__name__} {error.response.status_code} "
|
||||
f"on {request.method} {request.url.path}"
|
||||
)
|
||||
if isinstance(error, httpx.RequestError):
|
||||
# A transport error raised before the request was attached has no request to describe.
|
||||
try:
|
||||
request = error.request
|
||||
except RuntimeError:
|
||||
return type(error).__name__
|
||||
return f"{type(error).__name__} on {request.method} {request.url.path}"
|
||||
return type(error).__name__
|
||||
|
||||
|
||||
class NodeAgent:
|
||||
def __init__(
|
||||
self,
|
||||
settings: AgentSettings,
|
||||
transport: Any = None,
|
||||
host_collector: Any = None,
|
||||
accelerator_collector: Any = None,
|
||||
state_store: AgentStateStore | None = None,
|
||||
) -> None:
|
||||
self.settings = settings
|
||||
self.transport = transport or AgentTransport(
|
||||
settings.control_plane_url,
|
||||
settings.request_timeout_seconds,
|
||||
settings.tls_verify,
|
||||
)
|
||||
identity = NodeIdentityProvider(
|
||||
settings.identity_file,
|
||||
settings.identity,
|
||||
force_persisted=settings.identity_mode == "persisted",
|
||||
)
|
||||
self.host_collector = host_collector or SystemHostCollector(
|
||||
identity,
|
||||
{
|
||||
"model_cache": settings.model_cache_path,
|
||||
"artifacts": settings.artifact_path,
|
||||
"quarantine": settings.quarantine_path,
|
||||
},
|
||||
)
|
||||
self.accelerator_collector = accelerator_collector or build_nvml_collector()
|
||||
self.state_store = state_store or AgentStateStore(
|
||||
settings.state_file, settings.credential_file
|
||||
)
|
||||
self.state: AgentPersistentState = self.state_store.load()
|
||||
self.credential = self.state_store.credential()
|
||||
self.started_at = datetime.now(UTC)
|
||||
self._stop = asyncio.Event()
|
||||
# The publication loop and the artifact loop run in separate worker threads and both
|
||||
# need a credential. Without this lock they can enrol concurrently and burn one
|
||||
# single-use token on two node identities for the same hardware.
|
||||
self._enrollment_lock = threading.Lock()
|
||||
self._preflight_passed = False
|
||||
|
||||
def collect_accelerators(self, *, startup: bool = False) -> NvidiaCollection:
|
||||
collection: NvidiaCollection = self.accelerator_collector.collect()
|
||||
validate_nvidia_collection(
|
||||
collection,
|
||||
required=nvidia_required(self.settings.accelerator_mode),
|
||||
startup=startup,
|
||||
)
|
||||
return collection
|
||||
|
||||
def preflight(self) -> NvidiaCollection:
|
||||
"""Validate the accelerator contract before enrollment or publication."""
|
||||
|
||||
collection = self.collect_accelerators(startup=True)
|
||||
self._preflight_passed = True
|
||||
return collection
|
||||
|
||||
def metadata(self) -> AgentMetadata:
|
||||
return AgentMetadata(
|
||||
agent_version=__version__,
|
||||
protocol_version=AGENT_PROTOCOL_VERSION,
|
||||
supported_capabilities=AGENT_PROTOCOL_CAPABILITIES,
|
||||
started_at=self.started_at,
|
||||
)
|
||||
|
||||
def ensure_enrolled(self) -> None:
|
||||
if not self._preflight_passed:
|
||||
self.preflight()
|
||||
if self.credential:
|
||||
return
|
||||
with self._enrollment_lock:
|
||||
if self.credential:
|
||||
return
|
||||
token = self.settings.enrollment_token
|
||||
if token is None or not token.get_secret_value():
|
||||
raise RuntimeError("agent is not enrolled and no enrollment token was provided")
|
||||
host = self.host_collector.collect()
|
||||
response = self.transport.enroll(
|
||||
EnrollmentRequest(
|
||||
enrollment_token=token.get_secret_value(),
|
||||
identity_key=host.identity_key,
|
||||
identity_source=host.identity_source,
|
||||
hostname=host.hostname,
|
||||
display_name=self.settings.display_name or host.display_name,
|
||||
metadata=self.metadata(),
|
||||
).model_dump(mode="json")
|
||||
)
|
||||
credential = str(response["node_credential"])
|
||||
self.state_store.save_credential(credential)
|
||||
self.credential = credential
|
||||
|
||||
def enrolled_credential(self) -> str:
|
||||
self.ensure_enrolled()
|
||||
if self.credential is None: # Defensive invariant for custom state stores.
|
||||
raise RuntimeError("enrollment did not produce a node credential")
|
||||
return self.credential
|
||||
|
||||
def heartbeat_once(self, last_error: str | None = None) -> None:
|
||||
credential = self.enrolled_credential()
|
||||
host = self.host_collector.collect()
|
||||
request = HeartbeatRequest(
|
||||
identity_key=host.identity_key,
|
||||
metadata=self.metadata(),
|
||||
observed_at=datetime.now(UTC),
|
||||
last_error=last_error,
|
||||
)
|
||||
self.transport.heartbeat(request.model_dump(mode="json"), credential)
|
||||
|
||||
def inventory_once(self) -> None:
|
||||
nvidia = self.collect_accelerators() if self._preflight_passed else self.preflight()
|
||||
credential = self.enrolled_credential()
|
||||
host = self.host_collector.collect()
|
||||
self.state.inventory_sequence += 1
|
||||
self.state_store.save(self.state)
|
||||
request = InventoryReport(
|
||||
identity_key=host.identity_key,
|
||||
protocol_version=AGENT_PROTOCOL_VERSION,
|
||||
sequence=self.state.inventory_sequence,
|
||||
observed_at=max(host.inventory_at, nvidia.observed_at),
|
||||
host=host,
|
||||
nvidia=InventoryNvidiaPayload(
|
||||
availability=nvidia.availability,
|
||||
reason=nvidia.reason,
|
||||
inventory=nvidia.inventory,
|
||||
),
|
||||
)
|
||||
self.transport.inventory(request.model_dump(mode="json"), credential)
|
||||
|
||||
def telemetry_once(self) -> None:
|
||||
nvidia = self.collect_accelerators() if self._preflight_passed else self.preflight()
|
||||
credential = self.enrolled_credential()
|
||||
host = self.host_collector.collect()
|
||||
self.state.telemetry_sequence += 1
|
||||
self.state_store.save(self.state)
|
||||
request = TelemetryReport(
|
||||
identity_key=host.identity_key,
|
||||
protocol_version=AGENT_PROTOCOL_VERSION,
|
||||
sequence=self.state.telemetry_sequence,
|
||||
observed_at=max(host.inventory_at, nvidia.observed_at),
|
||||
available_ram_bytes=host.available_ram_bytes,
|
||||
storage=host.storage,
|
||||
accelerators=nvidia.telemetry,
|
||||
)
|
||||
self.transport.telemetry(request.model_dump(mode="json"), credential)
|
||||
|
||||
async def run(self) -> None:
|
||||
if not self._preflight_passed:
|
||||
self.preflight()
|
||||
artifact_task = asyncio.create_task(self._artifact_loop())
|
||||
next_heartbeat = next_inventory = next_telemetry = 0.0
|
||||
backoff = 1
|
||||
last_error: str | None = None
|
||||
try:
|
||||
while not self._stop.is_set():
|
||||
now = time.monotonic()
|
||||
try:
|
||||
if now >= next_heartbeat:
|
||||
await asyncio.to_thread(self.heartbeat_once, last_error)
|
||||
next_heartbeat = now + self.settings.heartbeat_interval_seconds
|
||||
if now >= next_inventory:
|
||||
await asyncio.to_thread(self.inventory_once)
|
||||
next_inventory = now + self.settings.inventory_interval_seconds
|
||||
if now >= next_telemetry:
|
||||
await asyncio.to_thread(self.telemetry_once)
|
||||
next_telemetry = now + self.settings.telemetry_interval_seconds
|
||||
backoff = 1
|
||||
last_error = None
|
||||
delay = (
|
||||
min(
|
||||
next_heartbeat,
|
||||
next_inventory,
|
||||
next_telemetry,
|
||||
)
|
||||
- time.monotonic()
|
||||
)
|
||||
delay = max(0.1, delay)
|
||||
except AgentStartupError:
|
||||
raise
|
||||
except (httpx.HTTPError, OSError, RuntimeError, ValueError, KeyError) as exc:
|
||||
last_error = describe_failure(exc)
|
||||
logger.warning(
|
||||
"control-plane publication failed; retrying: %s (backoff %ss)",
|
||||
last_error,
|
||||
backoff,
|
||||
)
|
||||
delay = backoff
|
||||
backoff = min(backoff * 2, self.settings.max_backoff_seconds)
|
||||
try:
|
||||
await asyncio.wait_for(self._stop.wait(), delay)
|
||||
except TimeoutError:
|
||||
continue
|
||||
finally:
|
||||
self._stop.set()
|
||||
await artifact_task
|
||||
|
||||
def artifact_once(self) -> bool:
|
||||
credential = self.enrolled_credential()
|
||||
if not hasattr(self.transport, "next_artifact_job"):
|
||||
return False
|
||||
lease = self.transport.next_artifact_job(credential)
|
||||
if not lease:
|
||||
return False
|
||||
acquirer = ArtifactAcquirer(self.settings, self.transport, credential)
|
||||
try:
|
||||
acquirer.execute(lease)
|
||||
except AcquisitionFailure as exc:
|
||||
acquirer.fail(lease, exc)
|
||||
except httpx.HTTPError as exc:
|
||||
acquirer.fail(
|
||||
lease,
|
||||
AcquisitionFailure(
|
||||
"download_transport_error", type(exc).__name__, retryable=True
|
||||
),
|
||||
)
|
||||
except (OSError, RuntimeError, ValueError, KeyError) as exc:
|
||||
acquirer.fail(
|
||||
lease,
|
||||
AcquisitionFailure("agent_execution_error", type(exc).__name__),
|
||||
)
|
||||
return True
|
||||
|
||||
async def _artifact_loop(self) -> None:
|
||||
while not self._stop.is_set():
|
||||
try:
|
||||
worked = await asyncio.to_thread(self.artifact_once)
|
||||
delay = 0.1 if worked else self.settings.artifact_job_poll_interval_seconds
|
||||
except (httpx.HTTPError, OSError, RuntimeError, ValueError, KeyError) as exc:
|
||||
logger.warning("artifact-job poll failed; retrying: %s", describe_failure(exc))
|
||||
delay = self.settings.max_backoff_seconds
|
||||
try:
|
||||
await asyncio.wait_for(self._stop.wait(), delay)
|
||||
except TimeoutError:
|
||||
continue
|
||||
|
||||
def stop(self) -> None:
|
||||
self._stop.set()
|
||||
|
||||
def close(self) -> None:
|
||||
self.transport.close()
|
||||
@@ -0,0 +1,43 @@
|
||||
import argparse
|
||||
import asyncio
|
||||
import logging
|
||||
import signal
|
||||
|
||||
from modelforge_node_agent.agent import NodeAgent
|
||||
from modelforge_node_agent.preflight import AgentStartupError
|
||||
from modelforge_node_agent.settings import AgentSettings
|
||||
|
||||
|
||||
async def run_agent(*, preflight_only: bool = False) -> None:
|
||||
agent = NodeAgent(AgentSettings())
|
||||
loop = asyncio.get_running_loop()
|
||||
for signum in (signal.SIGINT, signal.SIGTERM):
|
||||
loop.add_signal_handler(signum, agent.stop)
|
||||
try:
|
||||
if preflight_only:
|
||||
agent.preflight()
|
||||
logging.getLogger(__name__).info("Node Agent accelerator preflight passed")
|
||||
return
|
||||
await agent.run()
|
||||
finally:
|
||||
agent.close()
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description="ITWorx ModelForge Node Agent")
|
||||
parser.add_argument(
|
||||
"--preflight",
|
||||
action="store_true",
|
||||
help="validate the accelerator runtime without enrollment or publication",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
try:
|
||||
asyncio.run(run_agent(preflight_only=args.preflight))
|
||||
except AgentStartupError as exc:
|
||||
logging.getLogger(__name__).critical("%s", exc.problem)
|
||||
raise SystemExit(3) from None
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,129 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import platform
|
||||
from dataclasses import dataclass
|
||||
from enum import StrEnum
|
||||
from pathlib import Path
|
||||
|
||||
from modelforge_api.domain.enums import Availability
|
||||
from modelforge_api.domain.hardware import NvidiaCollection
|
||||
|
||||
|
||||
class AgentStartupFailureCode(StrEnum):
|
||||
NVIDIA_NVML_UNAVAILABLE = "NVIDIA_NVML_UNAVAILABLE"
|
||||
NVIDIA_DEVICE_NOT_FOUND = "NVIDIA_DEVICE_NOT_FOUND"
|
||||
NVIDIA_TELEMETRY_UNAVAILABLE = "NVIDIA_TELEMETRY_UNAVAILABLE"
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class AgentStartupProblem:
|
||||
code: AgentStartupFailureCode
|
||||
setting: str
|
||||
message: str
|
||||
|
||||
def __str__(self) -> str:
|
||||
return f"[{self.code}] {self.setting}: {self.message}"
|
||||
|
||||
|
||||
class AgentStartupError(RuntimeError):
|
||||
def __init__(self, problem: AgentStartupProblem) -> None:
|
||||
super().__init__(str(problem))
|
||||
self.problem = problem
|
||||
|
||||
|
||||
class AgentObservationError(RuntimeError):
|
||||
def __init__(self, problem: AgentStartupProblem) -> None:
|
||||
super().__init__(str(problem))
|
||||
self.problem = problem
|
||||
|
||||
|
||||
def nvidia_runtime_observed(
|
||||
*,
|
||||
environ: dict[str, str] | None = None,
|
||||
device_root: Path = Path("/dev"),
|
||||
) -> bool:
|
||||
"""Return whether the container has affirmative NVIDIA runtime evidence."""
|
||||
|
||||
environment = os.environ if environ is None else environ
|
||||
visible_devices = environment.get("NVIDIA_VISIBLE_DEVICES", "").strip().lower()
|
||||
if visible_devices not in {"", "none", "void"}:
|
||||
return True
|
||||
return any((device_root / name).exists() for name in ("nvidiactl", "nvidia-uvm", "nvidia0"))
|
||||
|
||||
|
||||
def nvidia_required(
|
||||
mode: str,
|
||||
*,
|
||||
environ: dict[str, str] | None = None,
|
||||
device_root: Path = Path("/dev"),
|
||||
) -> bool:
|
||||
if mode == "nvidia":
|
||||
return True
|
||||
if mode == "cpu":
|
||||
return False
|
||||
return nvidia_runtime_observed(environ=environ, device_root=device_root)
|
||||
|
||||
|
||||
def _runtime_description() -> str:
|
||||
libc_name, libc_version = platform.libc_ver()
|
||||
if libc_name:
|
||||
return f"{libc_name} {libc_version}".strip()
|
||||
if Path("/lib/ld-musl-x86_64.so.1").exists():
|
||||
return "musl"
|
||||
return "unknown libc"
|
||||
|
||||
|
||||
def validate_nvidia_collection(
|
||||
collection: NvidiaCollection,
|
||||
*,
|
||||
required: bool,
|
||||
startup: bool = True,
|
||||
) -> None:
|
||||
"""Enforce the GPU-node startup/publication contract without fabricating observations."""
|
||||
|
||||
if not required:
|
||||
return
|
||||
setting = "MODELFORGE_AGENT_ACCELERATOR_MODE=nvidia"
|
||||
if collection.availability is Availability.TEMPORARILY_FAILED:
|
||||
reason = collection.reason or "NVML telemetry collection temporarily failed"
|
||||
problem = AgentStartupProblem(
|
||||
code=AgentStartupFailureCode.NVIDIA_TELEMETRY_UNAVAILABLE,
|
||||
setting=setting,
|
||||
message=f"NVML returned a temporary incomplete observation ({reason})",
|
||||
)
|
||||
if startup:
|
||||
raise AgentStartupError(problem)
|
||||
raise AgentObservationError(problem)
|
||||
if collection.availability is not Availability.KNOWN:
|
||||
reason = collection.reason or "NVML did not provide a reason"
|
||||
raise AgentStartupError(
|
||||
AgentStartupProblem(
|
||||
code=AgentStartupFailureCode.NVIDIA_NVML_UNAVAILABLE,
|
||||
setting=setting,
|
||||
message=(
|
||||
f"NVML is unavailable ({reason}) under {_runtime_description()}; "
|
||||
"verify the NVIDIA Container Toolkit, injected driver libraries, and "
|
||||
"/dev/nvidia* device mapping"
|
||||
),
|
||||
)
|
||||
)
|
||||
if not collection.inventory:
|
||||
raise AgentStartupError(
|
||||
AgentStartupProblem(
|
||||
code=AgentStartupFailureCode.NVIDIA_DEVICE_NOT_FOUND,
|
||||
setting=setting,
|
||||
message="NVML initialized but returned no NVIDIA accelerator inventory",
|
||||
)
|
||||
)
|
||||
inventory_uuids = {item.device_uuid for item in collection.inventory}
|
||||
telemetry_uuids = {item.device_uuid for item in collection.telemetry}
|
||||
missing = sorted(inventory_uuids - telemetry_uuids)
|
||||
if missing:
|
||||
raise AgentStartupError(
|
||||
AgentStartupProblem(
|
||||
code=AgentStartupFailureCode.NVIDIA_TELEMETRY_UNAVAILABLE,
|
||||
setting=setting,
|
||||
message="NVML returned no telemetry for accelerator(s): " + ", ".join(missing),
|
||||
)
|
||||
)
|
||||
@@ -0,0 +1,33 @@
|
||||
from pathlib import Path
|
||||
from typing import Literal
|
||||
|
||||
from pydantic import Field, SecretStr
|
||||
from pydantic_settings import BaseSettings, SettingsConfigDict
|
||||
|
||||
|
||||
class AgentSettings(BaseSettings):
|
||||
model_config = SettingsConfigDict(
|
||||
env_prefix="MODELFORGE_AGENT_", env_file=".env", extra="ignore"
|
||||
)
|
||||
|
||||
control_plane_url: str = "http://api:8000"
|
||||
enrollment_token: SecretStr | None = None
|
||||
identity: str | None = None
|
||||
identity_mode: str = "auto"
|
||||
accelerator_mode: Literal["auto", "nvidia", "cpu"] = "auto"
|
||||
identity_file: Path = Path("/data/state/node-id")
|
||||
credential_file: Path = Path("/data/state/node-credential")
|
||||
state_file: Path = Path("/data/state/agent-state.json")
|
||||
display_name: str | None = None
|
||||
heartbeat_interval_seconds: int = Field(default=10, ge=1)
|
||||
inventory_interval_seconds: int = Field(default=300, ge=5)
|
||||
telemetry_interval_seconds: int = Field(default=30, ge=5)
|
||||
request_timeout_seconds: float = Field(default=10, gt=0)
|
||||
max_backoff_seconds: int = Field(default=60, ge=1)
|
||||
tls_verify: bool = True
|
||||
model_cache_path: Path = Path("/data/hf-cache")
|
||||
artifact_path: Path = Path("/data/artifacts")
|
||||
quarantine_path: Path = Path("/data/quarantine")
|
||||
hf_token: SecretStr | None = None
|
||||
artifact_job_poll_interval_seconds: int = Field(default=5, ge=1, le=60)
|
||||
artifact_download_timeout_seconds: float = Field(default=120, gt=0, le=600)
|
||||
@@ -0,0 +1,46 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
from dataclasses import asdict, dataclass
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
@dataclass
|
||||
class AgentPersistentState:
|
||||
inventory_sequence: int = 0
|
||||
telemetry_sequence: int = 0
|
||||
|
||||
|
||||
class AgentStateStore:
|
||||
def __init__(self, state_file: Path, credential_file: Path) -> None:
|
||||
self.state_file = state_file
|
||||
self.credential_file = credential_file
|
||||
|
||||
def load(self) -> AgentPersistentState:
|
||||
try:
|
||||
payload = json.loads(self.state_file.read_text(encoding="utf-8"))
|
||||
return AgentPersistentState(
|
||||
inventory_sequence=int(payload.get("inventory_sequence", 0)),
|
||||
telemetry_sequence=int(payload.get("telemetry_sequence", 0)),
|
||||
)
|
||||
except (OSError, ValueError, TypeError, json.JSONDecodeError):
|
||||
return AgentPersistentState()
|
||||
|
||||
def save(self, state: AgentPersistentState) -> None:
|
||||
self.state_file.parent.mkdir(parents=True, exist_ok=True)
|
||||
temporary = self.state_file.with_suffix(".tmp")
|
||||
temporary.write_text(json.dumps(asdict(state), sort_keys=True), encoding="utf-8")
|
||||
os.replace(temporary, self.state_file)
|
||||
|
||||
def credential(self) -> str | None:
|
||||
try:
|
||||
value = self.credential_file.read_text(encoding="utf-8").strip()
|
||||
return value or None
|
||||
except OSError:
|
||||
return None
|
||||
|
||||
def save_credential(self, credential: str) -> None:
|
||||
self.credential_file.parent.mkdir(parents=True, exist_ok=True)
|
||||
self.credential_file.write_text(credential, encoding="utf-8")
|
||||
os.chmod(self.credential_file, 0o600)
|
||||
@@ -0,0 +1,63 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
|
||||
|
||||
class AgentTransport:
|
||||
def __init__(self, base_url: str, timeout: float, verify: bool) -> None:
|
||||
self.client = httpx.Client(base_url=base_url.rstrip("/"), timeout=timeout, verify=verify)
|
||||
|
||||
def _request(
|
||||
self,
|
||||
method: str,
|
||||
path: str,
|
||||
payload: dict[str, Any] | None = None,
|
||||
credential: str | None = None,
|
||||
) -> dict[str, Any]:
|
||||
headers = {"Authorization": f"Bearer {credential}"} if credential else {}
|
||||
response = self.client.request(method, path, json=payload, headers=headers)
|
||||
response.raise_for_status()
|
||||
decoded = response.json()
|
||||
return dict(decoded) if isinstance(decoded, dict) else {}
|
||||
|
||||
def enroll(self, payload: dict[str, Any]) -> dict[str, Any]:
|
||||
return self._request("POST", "/api/v1/agent/enroll", payload)
|
||||
|
||||
def heartbeat(self, payload: dict[str, Any], credential: str) -> dict[str, Any]:
|
||||
return self._request("POST", "/api/v1/agent/heartbeat", payload, credential)
|
||||
|
||||
def inventory(self, payload: dict[str, Any], credential: str) -> dict[str, Any]:
|
||||
return self._request("PUT", "/api/v1/agent/inventory", payload, credential)
|
||||
|
||||
def telemetry(self, payload: dict[str, Any], credential: str) -> dict[str, Any]:
|
||||
return self._request("PUT", "/api/v1/agent/telemetry", payload, credential)
|
||||
|
||||
def next_artifact_job(self, credential: str) -> dict[str, Any] | None:
|
||||
response = self._request("GET", "/api/v1/agent/artifact-jobs/next", credential=credential)
|
||||
return response or None
|
||||
|
||||
def artifact_progress(
|
||||
self, job_id: str, payload: dict[str, Any], credential: str
|
||||
) -> dict[str, Any]:
|
||||
return self._request(
|
||||
"POST", f"/api/v1/agent/artifact-jobs/{job_id}/progress", payload, credential
|
||||
)
|
||||
|
||||
def artifact_complete(
|
||||
self, job_id: str, payload: dict[str, Any], credential: str
|
||||
) -> dict[str, Any]:
|
||||
return self._request(
|
||||
"POST", f"/api/v1/agent/artifact-jobs/{job_id}/complete", payload, credential
|
||||
)
|
||||
|
||||
def artifact_fail(
|
||||
self, job_id: str, payload: dict[str, Any], credential: str
|
||||
) -> dict[str, Any]:
|
||||
return self._request(
|
||||
"POST", f"/api/v1/agent/artifact-jobs/{job_id}/fail", payload, credential
|
||||
)
|
||||
|
||||
def close(self) -> None:
|
||||
self.client.close()
|
||||
@@ -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