332 lines
12 KiB
Python
332 lines
12 KiB
Python
from __future__ import annotations
|
|
|
|
from hashlib import sha256
|
|
import json
|
|
import subprocess
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
from fastapi.testclient import TestClient
|
|
|
|
from app.core.config import Settings
|
|
from app.main import app
|
|
from app.services.runtime_model_provenance_service import RuntimeModelProvenanceService
|
|
from app.services.yolo_preflight_service import YoloPreflightService
|
|
|
|
|
|
ROOT = Path(__file__).resolve().parents[2]
|
|
|
|
|
|
class AvailableAdapter:
|
|
@staticmethod
|
|
def dependencies_available() -> bool:
|
|
return True
|
|
|
|
def __init__(self, settings: Settings) -> None:
|
|
self.settings = settings
|
|
|
|
def load_model(self, model_path: Path) -> object:
|
|
return {"model_path": str(model_path)}
|
|
|
|
|
|
class MissingDependencyAdapter:
|
|
@staticmethod
|
|
def dependencies_available() -> bool:
|
|
return False
|
|
|
|
|
|
class FailingLoadAdapter(AvailableAdapter):
|
|
def load_model(self, model_path: Path) -> object:
|
|
raise RuntimeError(f"cannot load {model_path}")
|
|
|
|
|
|
def _manifest(tmp_path: Path, tile_count: int = 1) -> Path:
|
|
tiles = []
|
|
for index in range(tile_count):
|
|
tile_path = tmp_path / f"tile_{index:04d}.tif"
|
|
tile_path.write_bytes(b"tile")
|
|
tiles.append(
|
|
{
|
|
"path": str(tile_path),
|
|
"pixel_window": [0, 0, 100, 100],
|
|
"bounds": [4.0, 51.0, 5.0, 52.0],
|
|
"transform": [4.0, 0.01, 0.0, 52.0, 0.0, -0.01],
|
|
"index": index,
|
|
}
|
|
)
|
|
manifest_path = tmp_path / "manifest.json"
|
|
manifest_path.write_text(json.dumps({"tiles": tiles, "count": tile_count}), encoding="utf-8")
|
|
return manifest_path
|
|
|
|
|
|
def _write_model_sidecar(model_path: Path, settings: Settings) -> None:
|
|
model_sha256 = sha256(model_path.read_bytes()).hexdigest()
|
|
payload = {
|
|
"schema_version": RuntimeModelProvenanceService.MANIFEST_SCHEMA_VERSION,
|
|
"data_contract": {"key": "geointel.model.pytorch", "version": "1.0.0"},
|
|
"model": {
|
|
"model_id": settings.yolo_model_id,
|
|
"task_type": "object_detection",
|
|
"sha256": model_sha256,
|
|
"model_format": "pytorch",
|
|
"framework": "ultralytics/pytorch",
|
|
"class_mapping": {"0": "building"},
|
|
"source_version": settings.yolo_model_version or "test-v1",
|
|
},
|
|
"source": {
|
|
"source_registry_id": "11111111-1111-4111-8111-111111111111",
|
|
"source_snapshot_id": "22222222-2222-4222-8222-222222222222",
|
|
"source_registry_key": "model",
|
|
"source_snapshot_checksum_sha256": model_sha256,
|
|
},
|
|
"lineage": {
|
|
"upstream_asset_ids": ["test-training-corpus"],
|
|
"upstream_checksums_sha256": ["a" * 64],
|
|
"transformations": [
|
|
{"name": "test-training", "version": "1.0.0", "checksum_sha256": "b" * 64}
|
|
],
|
|
},
|
|
"metadata": {"training_manifest_sha256": "c" * 64},
|
|
"imported_at": "2026-08-01T10:00:00+00:00",
|
|
}
|
|
payload["metadata"]["runtime_manifest_sha256"] = RuntimeModelProvenanceService.manifest_self_checksum(payload)
|
|
RuntimeModelProvenanceService.manifest_path_for_model(model_path).write_text(
|
|
json.dumps(payload, sort_keys=True),
|
|
encoding="utf-8",
|
|
)
|
|
|
|
|
|
def test_yolo_preflight_reports_disabled_without_loading_model(tmp_path: Path, monkeypatch) -> None:
|
|
monkeypatch.setenv("YOLO_CONFIG_DIR", str(tmp_path / "ultralytics"))
|
|
|
|
result = YoloPreflightService.run(
|
|
settings=Settings(yolo_enabled=False, yolo_model_path=str(tmp_path / "missing.pt")),
|
|
tile_manifest_path=str(tmp_path / "missing-manifest.json"),
|
|
yolo_adapter_class=AvailableAdapter,
|
|
)
|
|
|
|
assert result["status"] == "not_configured"
|
|
assert result["checks"]["enabled"] is False
|
|
assert result["checks"]["dependencies_available"] is None
|
|
assert result["checks"]["model_file_exists"] is None
|
|
assert result["checks"]["manifest_valid"] is None
|
|
assert result["runtime"]["dependencies_assumed"] is False
|
|
assert result["runtime"]["model_directory"] == str(tmp_path)
|
|
assert result["runtime"]["yolo_config_dir"] == str(tmp_path / "ultralytics")
|
|
assert "torch_version" in result["runtime"]
|
|
assert "ultralytics_version" in result["runtime"]
|
|
assert "cuda_available" in result["runtime"]
|
|
|
|
|
|
def test_yolo_preflight_distinguishes_missing_dependencies_from_missing_model(tmp_path: Path) -> None:
|
|
model_path = tmp_path / "model.pt"
|
|
model_path.write_bytes(b"weights")
|
|
|
|
result = YoloPreflightService.run(
|
|
settings=Settings(yolo_enabled=True, yolo_model_path=str(model_path)),
|
|
tile_manifest_path=str(_manifest(tmp_path)),
|
|
yolo_adapter_class=MissingDependencyAdapter,
|
|
)
|
|
|
|
assert result["status"] == "dependency_unavailable"
|
|
assert result["checks"]["dependencies_available"] is False
|
|
assert result["checks"]["model_file_exists"] is None
|
|
assert result["checks"]["manifest_valid"] is None
|
|
|
|
|
|
def test_yolo_preflight_validates_model_and_manifest_without_importing_yolo(tmp_path: Path) -> None:
|
|
model_path = tmp_path / "model.pt"
|
|
model_path.write_bytes(b"weights")
|
|
manifest_path = _manifest(tmp_path, tile_count=2)
|
|
settings = Settings(yolo_enabled=True, yolo_model_path=str(model_path), yolo_max_tiles=4)
|
|
_write_model_sidecar(model_path, settings)
|
|
|
|
result = YoloPreflightService.run(
|
|
settings=settings,
|
|
tile_manifest_path=str(manifest_path),
|
|
yolo_adapter_class=AvailableAdapter,
|
|
)
|
|
|
|
assert result["status"] == "ready"
|
|
assert result["checks"]["dependencies_available"] is True
|
|
assert result["checks"]["model_file_exists"] is True
|
|
assert result["checks"]["model_provenance_valid"] is True
|
|
assert result["checks"]["manifest_valid"] is True
|
|
assert result["tile_count"] == 2
|
|
assert result["will_download_models"] is False
|
|
assert result["will_run_inference"] is False
|
|
assert result["runtime"]["dependencies_assumed"] is False
|
|
|
|
|
|
def test_yolo_preflight_marks_assumed_dependencies_in_runtime_details(tmp_path: Path) -> None:
|
|
model_path = tmp_path / "model.pt"
|
|
model_path.write_bytes(b"weights")
|
|
manifest_path = _manifest(tmp_path, tile_count=1)
|
|
settings = Settings(yolo_enabled=True, yolo_model_path=str(model_path), yolo_max_tiles=4)
|
|
_write_model_sidecar(model_path, settings)
|
|
|
|
result = YoloPreflightService.run(
|
|
settings=settings,
|
|
tile_manifest_path=str(manifest_path),
|
|
yolo_adapter_class=MissingDependencyAdapter,
|
|
assume_dependencies=True,
|
|
)
|
|
|
|
assert result["status"] == "ready"
|
|
assert result["checks"]["dependencies_available"] is True
|
|
assert result["runtime"]["dependencies_assumed"] is True
|
|
assert result["runtime"]["model_directory"] == str(tmp_path)
|
|
|
|
|
|
def test_yolo_preflight_can_explicitly_smoke_load_local_model(tmp_path: Path) -> None:
|
|
model_path = tmp_path / "model.pt"
|
|
model_path.write_bytes(b"weights")
|
|
manifest_path = _manifest(tmp_path)
|
|
settings = Settings(yolo_enabled=True, yolo_model_path=str(model_path), yolo_max_tiles=4)
|
|
_write_model_sidecar(model_path, settings)
|
|
|
|
result = YoloPreflightService.run(
|
|
settings=settings,
|
|
tile_manifest_path=str(manifest_path),
|
|
yolo_adapter_class=AvailableAdapter,
|
|
check_model_load=True,
|
|
)
|
|
|
|
assert result["status"] == "ready"
|
|
assert result["checks"]["model_load_requested"] is True
|
|
assert result["checks"]["model_load_ok"] is True
|
|
assert result["will_download_models"] is False
|
|
assert result["will_run_inference"] is False
|
|
|
|
|
|
def test_yolo_preflight_reports_explicit_model_load_failure(tmp_path: Path) -> None:
|
|
model_path = tmp_path / "model.pt"
|
|
model_path.write_bytes(b"weights")
|
|
settings = Settings(yolo_enabled=True, yolo_model_path=str(model_path), yolo_max_tiles=4)
|
|
_write_model_sidecar(model_path, settings)
|
|
|
|
result = YoloPreflightService.run(
|
|
settings=settings,
|
|
tile_manifest_path=str(_manifest(tmp_path)),
|
|
yolo_adapter_class=FailingLoadAdapter,
|
|
check_model_load=True,
|
|
)
|
|
|
|
assert result["status"] == "model_load_failed"
|
|
assert result["error_code"] == "DETECTION_MODEL_LOAD_FAILED"
|
|
assert result["checks"]["model_load_ok"] is False
|
|
|
|
|
|
def test_yolo_preflight_script_outputs_json(tmp_path: Path) -> None:
|
|
model_path = tmp_path / "model.pt"
|
|
model_path.write_bytes(b"weights")
|
|
manifest_path = _manifest(tmp_path)
|
|
_write_model_sidecar(model_path, Settings(yolo_enabled=True, yolo_model_path=str(model_path)))
|
|
|
|
result = subprocess.run(
|
|
[
|
|
sys.executable,
|
|
str(ROOT / "scripts" / "yolo_preflight.py"),
|
|
"--model-path",
|
|
str(model_path),
|
|
"--tile-manifest-path",
|
|
str(manifest_path),
|
|
"--assume-dependencies",
|
|
"--json",
|
|
],
|
|
cwd=ROOT,
|
|
check=True,
|
|
capture_output=True,
|
|
text=True,
|
|
)
|
|
payload = json.loads(result.stdout)
|
|
|
|
assert payload["status"] == "ready"
|
|
assert payload["model_path"] == str(model_path)
|
|
assert payload["tile_manifest_path"] == str(manifest_path)
|
|
|
|
|
|
def test_yolo_preflight_script_uses_environment_configuration(tmp_path: Path, monkeypatch) -> None:
|
|
model_path = tmp_path / "model.pt"
|
|
model_path.write_bytes(b"weights")
|
|
manifest_path = _manifest(tmp_path)
|
|
_write_model_sidecar(model_path, Settings(yolo_enabled=True, yolo_model_path=str(model_path), yolo_max_tiles=4))
|
|
monkeypatch.setenv("YOLO_ENABLED", "true")
|
|
monkeypatch.setenv("YOLO_MODEL_PATH", str(model_path))
|
|
monkeypatch.setenv("YOLO_MAX_TILES", "4")
|
|
|
|
result = subprocess.run(
|
|
[
|
|
sys.executable,
|
|
str(ROOT / "scripts" / "yolo_preflight.py"),
|
|
"--tile-manifest-path",
|
|
str(manifest_path),
|
|
"--assume-dependencies",
|
|
"--json",
|
|
],
|
|
cwd=ROOT,
|
|
check=True,
|
|
capture_output=True,
|
|
text=True,
|
|
)
|
|
payload = json.loads(result.stdout)
|
|
|
|
assert payload["status"] == "ready"
|
|
assert payload["checks"]["enabled"] is True
|
|
assert payload["model_path"] == str(model_path)
|
|
assert payload["max_tiles"] == 4
|
|
|
|
|
|
def test_yolo_preflight_refuses_unmanifested_local_weights(tmp_path: Path) -> None:
|
|
model_path = tmp_path / "model.pt"
|
|
model_path.write_bytes(b"unmanifested weights")
|
|
|
|
result = YoloPreflightService.run(
|
|
settings=Settings(yolo_enabled=True, yolo_model_path=str(model_path)),
|
|
tile_manifest_path=str(_manifest(tmp_path)),
|
|
yolo_adapter_class=AvailableAdapter,
|
|
)
|
|
|
|
assert result["status"] == "contract_incomplete"
|
|
assert result["checks"]["model_provenance_valid"] is False
|
|
assert result["error_code"] == "MODEL_PROVENANCE_MANIFEST_MISSING"
|
|
|
|
|
|
def test_yolo_preflight_script_rejects_assumed_dependencies_for_model_load(tmp_path: Path) -> None:
|
|
result = subprocess.run(
|
|
[
|
|
sys.executable,
|
|
str(ROOT / "scripts" / "yolo_preflight.py"),
|
|
"--model-path",
|
|
str(tmp_path / "model.pt"),
|
|
"--assume-dependencies",
|
|
"--check-model-load",
|
|
"--json",
|
|
],
|
|
cwd=ROOT,
|
|
check=False,
|
|
capture_output=True,
|
|
text=True,
|
|
)
|
|
|
|
assert result.returncode != 0
|
|
assert "--check-model-load cannot be combined with --assume-dependencies" in result.stderr
|
|
|
|
|
|
def test_yolo_preflight_api_returns_canonical_envelope(monkeypatch, tmp_path: Path) -> None:
|
|
monkeypatch.setenv("YOLO_ENABLED", "false")
|
|
monkeypatch.setenv("YOLO_MODEL_PATH", str(tmp_path / "missing.pt"))
|
|
monkeypatch.setenv("YOLO_CONFIG_DIR", str(tmp_path / "ultralytics"))
|
|
|
|
response = TestClient(app).get("/api/v1/detection/yolo/preflight")
|
|
|
|
assert response.status_code == 200
|
|
payload = response.json()
|
|
assert set(payload) == {"data"}
|
|
assert payload["data"]["status"] == "not_configured"
|
|
assert payload["data"]["checks"]["enabled"] is False
|
|
assert payload["data"]["runtime"]["model_directory"] == str(tmp_path)
|
|
assert payload["data"]["runtime"]["yolo_config_dir"] == str(tmp_path / "ultralytics")
|
|
assert payload["data"]["will_download_models"] is False
|
|
assert payload["data"]["will_run_inference"] is False
|