Publish curated ITWorx-ModelForge source from 990401737f96c8d46a9f7b0d84994516cadcb6a0

runtime-worker/tests/test_worker.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 import hashlib
 import json
 import socket
+import threading
 import uuid
 from pathlib import Path
 from typing import Any
@@ -18,7 +19,11 @@ from modelforge_runtime_worker.adapters import (
     WhisperAdapter,
     adapter_for,
 )
-from modelforge_runtime_worker.contracts import AgentRuntimeProbeLease, AgentServingJobLease
+from modelforge_runtime_worker.contracts import (
+    AgentRuntimeProbeLease,
+    AgentServingJobComplete,
+    AgentServingJobLease,
+)
 from modelforge_runtime_worker.probe import (
     ProbeFailure,
     child_entry,
@@ -62,6 +67,41 @@ class FakeTransport:
         pass
 
 
+class FakeLabCancellationTransport(FakeTransport):
+    def __init__(self, lease: dict[str, Any]) -> None:
+        super().__init__()
+        self.serving_lease = lease
+        self.cancel_polled = threading.Event()
+        self.cancellation_acks: list[dict[str, Any]] = []
+        self.serving_completions: list[dict[str, Any]] = []
+        self.serving_failures: list[dict[str, Any]] = []
+
+    def next_serving_job(self, _credential: str, *, wait_seconds: float = 1.0):
+        del wait_seconds
+        value, self.serving_lease = self.serving_lease, None
+        return value
+
+    def check_serving_job_cancellation(
+        self, _job_id: str, _credential: str, _payload: dict[str, Any]
+    ) -> dict[str, bool]:
+        self.cancel_polled.set()
+        return {"cancel_requested": True}
+
+    def acknowledge_serving_job_cancellation(
+        self, _job_id: str, _credential: str, payload: dict[str, Any]
+    ) -> dict[str, str]:
+        self.cancellation_acks.append(payload)
+        return {"status": "accepted"}
+
+    def complete_serving_job(self, _job_id: str, _credential: str, payload: dict[str, Any]):
+        self.serving_completions.append(payload)
+        return {"status": "accepted"}
+
+    def fail_serving_job(self, _job_id: str, _credential: str, payload: dict[str, Any]):
+        self.serving_failures.append(payload)
+        return {"status": "accepted"}
+
+
 def lease_payload() -> dict[str, Any]:
     return {
         "probe_id": "00000000-0000-0000-0000-000000000001",
@@ -116,6 +156,37 @@ def serving_lease_payload(operation: str = "load") -> dict[str, Any]:
     }
 
 
+def test_lab_worker_polls_while_inference_runs_and_acks_at_safe_boundary(tmp_path: Path) -> None:
+    lease = {
+        **serving_lease_payload("invoke"),
+        "lab_run_id": "00000000-0000-0000-0000-000000000099",
+    }
+    transport = FakeLabCancellationTransport(lease)
+    worker = RuntimeWorker(make_settings(tmp_path), transport)
+    worker.serving.worker_instance_id = "worker-test"
+
+    def slow_safe_inference(_lease: dict[str, Any]) -> AgentServingJobComplete:
+        assert transport.cancel_polled.wait(timeout=2)
+        return AgentServingJobComplete(
+            lease_token=lease["lease_token"],
+            worker_instance_id="worker-test",
+            state="warm",
+            result={"vectors": [[0.0] * 1024], "count": 1},
+            metrics={},
+            health={},
+            runtime_facts={},
+        )
+
+    worker.serving.execute = slow_safe_inference  # type: ignore[method-assign]
+    assert worker.run_once()
+
+    assert len(transport.cancellation_acks) == 1
+    assert transport.cancellation_acks[0]["lease_token"] == lease["lease_token"]
+    assert transport.cancellation_acks[0]["state"] == "warm"
+    assert transport.serving_completions == []
+    assert transport.serving_failures == []
+
+
 def make_settings(tmp_path: Path) -> RuntimeWorkerSettings:
     credential = tmp_path / "credential"
     credential.write_text("credential", encoding="utf-8")
@@ -363,6 +434,8 @@ class FakeServingAdapter(SentenceTransformersAdapter):
         self.during = {
             "used_vram_bytes": 320,
             "process_used_vram_bytes": 220,
+            "free_vram_bytes": 280,
+            "gpu_utilization_percent": 42,
             "temperature_c": 50,
             "power_draw_w": 40.0,
             "gpu_uuid": "GPU-test",
@@ -491,6 +564,19 @@ def test_serving_engine_load_invoke_warm_reuse_and_unload(
     result = engine.execute(invoke)
     assert result.result["count"] == 1
     assert len(result.result["vectors"][0]) == 1024
+    assert result.metrics == {
+        "inference_time_ms": 25.0,
+        "preprocess_time_ms": 0.0,
+        "serialization_time_ms": 0.0,
+        "worker_total_time_ms": 25.0,
+        "gpu_uuid": "GPU-test",
+        "used_vram_bytes": 320,
+        "process_used_vram_bytes": 220,
+        "free_vram_bytes": 280,
+        "gpu_utilization_percent": 42,
+        "temperature_c": 50,
+        "power_draw_w": 40.0,
+    }
     unload = serving_lease_payload("unload")
     unload["artifact_root"] = str(tmp_path)
     unloaded = engine.execute(unload)