Publish curated ITWorx-ModelForge source from 990401737f96c8d46a9f7b0d84994516cadcb6a0

runtime-worker/Dockerfile

@@ -18,20 +18,28 @@ LABEL org.opencontainers.image.description="Private typed model runtime for ITWo
 LABEL org.opencontainers.image.version="${MODELFORGE_VERSION}"
 LABEL org.opencontainers.image.revision="${MODELFORGE_COMMIT}"
 LABEL org.opencontainers.image.created="${MODELFORGE_BUILT_AT}"
-LABEL org.opencontainers.image.source="https://git.example.com/example/modelforge.git"
+LABEL org.opencontainers.image.source="https://git.example.com/example/modelforge"
 LABEL org.opencontainers.image.vendor="ITWorx"
 LABEL org.opencontainers.image.licenses="AGPL-3.0-or-later"
 
 WORKDIR /build
+COPY config/security/build-toolchain.lock /tmp/modelforge-build-toolchain.lock
 COPY runtime-worker ./runtime-worker
-RUN apt-get update \
+RUN python -m pip install --no-cache-dir --require-hashes --only-binary=:all: \
+        --index-url https://pypi.org/simple -r /tmp/modelforge-build-toolchain.lock \
+    && apt-get update \
     && DEBIAN_FRONTEND=noninteractive apt-get upgrade --yes \
-    && pip install --no-cache-dir './runtime-worker[inference]' \
+    && python -m pip install --no-cache-dir --no-build-isolation './runtime-worker[inference]' \
     && DEBIAN_FRONTEND=noninteractive apt-get purge --yes linux-libc-dev \
     && rm -rf /var/lib/apt/lists/* \
     && pip check \
+    && python -m pip uninstall --yes pip setuptools wheel editables hatchling pathspec pluggy \
+        tomlkit trove-classifiers \
+    # This is the only shared group: it grants the worker read-only access to the Node Agent's
+    # 0640 node credential in the project-scoped, read-only state volume.
+    && groupadd --system --gid 101 modelforge-agent \
     && groupadd --system modelforge-runtime \
-    && useradd --system --gid modelforge-runtime --home /app modelforge-runtime \
+    && useradd --system --gid modelforge-runtime --groups modelforge-agent --home /app modelforge-runtime \
     && mkdir -p /app /tmp/huggingface \
     && chown -R modelforge-runtime:modelforge-runtime /app /tmp/huggingface
 

runtime-worker/pyproject.toml

@@ -1,10 +1,12 @@
 [build-system]
-requires = ["setuptools==80.9.0"]
+requires = [
+  "setuptools==84.0.0",
+]
 build-backend = "setuptools.build_meta"
 
 [project]
 name = "modelforge-runtime-worker"
-version = "1.2.2"
+version = "1.3.0-rc.1"
 license = "AGPL-3.0-or-later"
 requires-python = ">=3.11"
 dependencies = [
@@ -27,8 +29,14 @@ inference = [
   # makes the final runtime environment reproducible and prevents an old base package from
   # satisfying an unconstrained transitive dependency.
   "Brotli==1.2.0",
+  "filelock==3.20.3",
+  "h2==4.4.1",
+  "idna==3.15",
   "jaraco.context==6.1.2",
+  "lief==0.17.2",
   "msgpack==1.2.2",
+  "Pygments==2.20.0",
+  "requests==2.33.0",
   "sentencepiece==0.2.2",
   "sentence-transformers==6.0.1",
   "setuptools==84.0.0",
@@ -38,7 +46,7 @@ inference = [
   "wheel==0.48.0",
 ]
 dev = [
-  "pytest==8.4.2",
+  "pytest==9.0.3",
   "ruff==0.16.5",
   "mypy==1.20.2",
 ]

runtime-worker/src/modelforge_runtime_worker/__init__.py

@@ -1 +1 @@
-__version__ = "1.2.2"
+__version__ = "1.3.0-rc.1"

runtime-worker/src/modelforge_runtime_worker/contracts.py

@@ -81,6 +81,7 @@ class AgentServingJobLease(WireModel):
     top_n: int | None = Field(default=None, ge=1, le=10)
     modality_payload: dict[str, Any] | None = None
     request_id: uuid.UUID | None = None
+    lab_run_id: uuid.UUID | None = None
 
     @model_validator(mode="after")
     def validate_capability_payload(self) -> AgentServingJobLease:

runtime-worker/src/modelforge_runtime_worker/serving.py

@@ -278,7 +278,7 @@ class ServingEngine:
             "capability": "healthy",
         }
         timings = result.get("timings", {})
-        return result, {
+        metrics = {
             "inference_time_ms": result["inference_time_ms"],
             "preprocess_time_ms": float(timings.get("preprocess_ms", 0.0)),
             "serialization_time_ms": float(timings.get("serialization_ms", 0.0)),
@@ -289,6 +289,16 @@ class ServingEngine:
             "temperature_c": adapter.during["temperature_c"],
             "power_draw_w": adapter.during["power_draw_w"],
         }
+        if lease.capability == "rag.embedding":
+            metrics.update(
+                {
+                    "gpu_uuid": adapter.during["gpu_uuid"],
+                    "process_used_vram_bytes": adapter.during["process_used_vram_bytes"],
+                    "free_vram_bytes": adapter.during["free_vram_bytes"],
+                    "gpu_utilization_percent": adapter.during["gpu_utilization_percent"],
+                }
+            )
+        return result, metrics
 
     def _health(self, lease: AgentServingJobLease) -> tuple[dict[str, Any], dict[str, Any]]:
         if lease.capability not in {"rag.embedding", "rag.reranking"}:

runtime-worker/src/modelforge_runtime_worker/transport.py

@@ -82,6 +82,32 @@ class RuntimeTransport:
             or {}
         )
 
+    def check_serving_job_cancellation(
+        self, job_id: str, credential: str, payload: dict[str, Any]
+    ) -> dict[str, Any]:
+        return (
+            self._request(
+                "POST",
+                f"/api/v1/agent/serving-jobs/{job_id}/cancellation",
+                credential,
+                payload,
+            )
+            or {}
+        )
+
+    def acknowledge_serving_job_cancellation(
+        self, job_id: str, credential: str, payload: dict[str, Any]
+    ) -> dict[str, Any]:
+        return (
+            self._request(
+                "POST",
+                f"/api/v1/agent/serving-jobs/{job_id}/cancelled",
+                credential,
+                payload,
+            )
+            or {}
+        )
+
     def report_serving_state(self, credential: str, payload: dict[str, Any]) -> dict[str, Any]:
         return self._request("POST", "/api/v1/agent/serving-state", credential, payload) or {}
 

runtime-worker/src/modelforge_runtime_worker/worker.py

@@ -1,6 +1,8 @@
 from __future__ import annotations
 
+import logging
 import multiprocessing
+import threading
 import time
 from pathlib import Path
 from queue import Empty
@@ -11,6 +13,7 @@ from modelforge_runtime_worker.contracts import (
     AgentRuntimeProbeFailure,
     AgentRuntimeProbeLease,
     AgentRuntimeProbeProgress,
+    AgentServingJobLease,
 )
 from modelforge_runtime_worker.probe import (
     ProbeFailure,
@@ -23,6 +26,8 @@ from modelforge_runtime_worker.serving import ServingEngine
 from modelforge_runtime_worker.settings import RuntimeWorkerSettings
 from modelforge_runtime_worker.transport import RuntimeTransport
 
+logger = logging.getLogger(__name__)
+
 
 class ProbeCancelled(ProbeFailure):
     def __init__(self) -> None:
@@ -186,30 +191,83 @@ class RuntimeWorker:
                 serving_lease = self.transport.next_serving_job(credential)
             if serving_lease:
                 job_id = str(serving_lease["job_id"])
+                parsed_lease = AgentServingJobLease.model_validate(serving_lease)
+                stop_monitor = threading.Event()
+                cancellation_requested = threading.Event()
+                monitor: threading.Thread | None = None
+                if parsed_lease.lab_run_id and hasattr(
+                    self.transport, "check_serving_job_cancellation"
+                ):
+                    monitor = threading.Thread(
+                        target=self._monitor_lab_cancellation,
+                        args=(
+                            job_id,
+                            parsed_lease.lease_token,
+                            credential,
+                            stop_monitor,
+                            cancellation_requested,
+                        ),
+                        name="lab-run-cancellation-heartbeat",
+                        daemon=True,
+                    )
+                    monitor.start()
+
+                def stop_heartbeat() -> None:
+                    stop_monitor.set()
+                    if monitor is not None:
+                        monitor.join(timeout=1.0)
+
+                def acknowledge_if_requested(worker_state: str) -> bool:
+                    if (
+                        parsed_lease.operation != "invoke"
+                        or not cancellation_requested.is_set()
+                        or not hasattr(self.transport, "acknowledge_serving_job_cancellation")
+                    ):
+                        return False
+                    self.transport.acknowledge_serving_job_cancellation(
+                        job_id,
+                        credential,
+                        {
+                            "lease_token": parsed_lease.lease_token,
+                            "worker_instance_id": self.serving.worker_instance_id,
+                            "state": worker_state,
+                        },
+                    )
+                    return True
+
                 try:
                     completed = self.serving.execute(serving_lease)
-                    self.transport.complete_serving_job(
-                        job_id,
-                        credential,
-                        completed.model_dump(mode="json"),
-                    )
                 except ProbeFailure as exc:
                     failure = self.serving.failure(serving_lease, exc)
-                    self.transport.fail_serving_job(
-                        job_id,
-                        credential,
-                        failure.model_dump(mode="json"),
-                    )
+                    stop_heartbeat()
+                    if not acknowledge_if_requested(failure.state):
+                        self.transport.fail_serving_job(
+                            job_id,
+                            credential,
+                            failure.model_dump(mode="json"),
+                        )
                 except (OSError, RuntimeError, ValueError, KeyError) as exc:
                     failure = self.serving.failure(
                         serving_lease,
                         ProbeFailure("RUNTIME_CRASH", type(exc).__name__),
                     )
-                    self.transport.fail_serving_job(
-                        job_id,
-                        credential,
-                        failure.model_dump(mode="json"),
-                    )
+                    stop_heartbeat()
+                    if not acknowledge_if_requested(failure.state):
+                        self.transport.fail_serving_job(
+                            job_id,
+                            credential,
+                            failure.model_dump(mode="json"),
+                        )
+                else:
+                    stop_heartbeat()
+                    if not acknowledge_if_requested(completed.state):
+                        self.transport.complete_serving_job(
+                            job_id,
+                            credential,
+                            completed.model_dump(mode="json"),
+                        )
+                finally:
+                    stop_heartbeat()
                 return True
         lease = self.transport.next_probe(credential)
         if not lease:
@@ -230,5 +288,29 @@ class RuntimeWorker:
             )
         return True
 
+    def _monitor_lab_cancellation(
+        self,
+        job_id: str,
+        lease_token: str,
+        credential: str,
+        stop: threading.Event,
+        cancellation_requested: threading.Event,
+    ) -> None:
+        while not stop.is_set():
+            try:
+                response = self.transport.check_serving_job_cancellation(
+                    job_id,
+                    credential,
+                    {"lease_token": lease_token},
+                )
+                if response.get("cancel_requested") is True:
+                    cancellation_requested.set()
+            except Exception as exc:
+                # A missed poll does not authorize retry or completion: the control plane
+                # rechecks cancellation atomically when the worker reports its result.
+                logger.debug("LAB cancellation poll failed (%s)", type(exc).__name__)
+            if stop.wait(1.0):
+                return
+
     def close(self) -> None:
         self.transport.close()

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)