321 lines
12 KiB
Python
321 lines
12 KiB
Python
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()
|