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()
|
||||
Reference in New Issue
Block a user