NODEDC_MISSION_CORE/simulation/ai-polygon/model_stack.py

"""Run-owned additive Docker stack. Caller must hold Core's GPU reservation."""

import hashlib
import http.client
import json
import subprocess
import sys
import time
from pathlib import Path

ROOT = Path(__file__).resolve().parent
sys.path.insert(0, str(ROOT.parents[1] / "src"))
NAVIGATION_NAME = "ndc-mission-core-ai-module-simulation-cmu-navigation"
NAMES = (
    "ndc-ai-polygon-ddrnet",
    "ndc-ai-polygon-segformer",
    "ndc-ai-polygon-rfdetr",
    NAVIGATION_NAME,
)


def sha256(path):
    with Path(path).open("rb") as stream:
        return hashlib.file_digest(stream, "sha256").hexdigest()


def docker(*args, check=True):
    return subprocess.run(
        ["docker", *args], check=check, capture_output=True, text=True, timeout=120
    )


class ModelStack:
    def __init__(self):
        self.profile_path = ROOT / "models.worker-006.json"
        self.profile = json.loads(self.profile_path.read_text())
        self.ids = []

    def preflight(self):
        from k1link.simulation.ai_polygon.composition import compose

        compose(ROOT)  # Portable graph admission happens before any Isaac boot.
        for model in [*self.profile["models"], self.profile["navigation"]]:
            for key in ("checkpoint", "runner"):
                if key in model and sha256(model[key]) != model[key + "_sha256"]:
                    raise RuntimeError("Simulation model identity changed: " + model["id"])
            actual = docker(
                "image", "inspect", model["image"], "--format", "{{.Id}}"
            ).stdout.strip()
            if actual != model["image"]:
                raise RuntimeError("Simulation model image is unavailable")
        if sha256(self.profile["labels"]) != self.profile["labels_sha256"]:
            raise RuntimeError("GOOSE label table changed")

    def start(self, cancelled=None, on_acquired=None, selection=None):
        from k1link.simulation.ai_polygon.composition import compose

        graph = compose(ROOT, selection)
        services = {
            "simulation-ddrnet-goose": "ddrnet",
            "simulation-segformer-ade": "segformer",
            "simulation-rf-detr": "detector",
        }
        module_ids = [node.module.module_id for node in graph.nodes]
        installed = {*services, "simulation-cmu-navigation", "simulation-waypoint-mission"}
        if set(module_ids) - installed:
            raise ValueError("Uninstalled simulation provider")
        self.preflight()
        for name in NAMES:
            if docker("container", "inspect", name, check=False).returncode == 0:
                raise RuntimeError(
                    "An existing simulation container requires reconciliation: " + name
                )
        if cancelled and cancelled():
            raise InterruptedError("Model startup cancelled")
        try:
            selected = [services[module] for module in module_ids if module in services]
            docker("compose", "-f", str(ROOT / "compose.models.yaml"), "up", "-d", *selected)
            if "simulation-cmu-navigation" in module_ids:
                docker(
                    "run",
                    "-d",
                    "--name",
                    NAVIGATION_NAME,
                    "--restart",
                    "no",
                    "--label",
                    "com.nodedc.product=mission-core",
                    "--label",
                    "com.nodedc.stack=ai-polygon",
                    "--label",
                    "com.nodedc.role=ai-module",
                    "--label",
                    "com.nodedc.managed-by=ai-polygon-worker",
                    "--label",
                    "com.nodedc.composition-sha256=" + graph.sha256,
                    "--cpus",
                    "3",
                    "--memory",
                    "2g",
                    "-p",
                    "127.0.0.1:18093:8010",
                    "--mount",
                    f"type=bind,source={ROOT},target=/adapter,readonly",
                    self.profile["navigation"]["image"],
                )
        finally:
            # Capture immutable IDs even after a partial Compose start.
            for name in NAMES:
                result = docker("container", "inspect", name, check=False)
                if result.returncode == 0:
                    row = json.loads(result.stdout)[0]
                    if row["Config"]["Labels"].get("com.nodedc.stack") == "ai-polygon":
                        self.ids.append(row["Id"])
            if on_acquired:
                on_acquired(list(self.ids))
        deadline = time.monotonic() + 120
        while time.monotonic() < deadline:
            if cancelled and cancelled():
                raise InterruptedError("Model startup cancelled")
            ready = True
            for port, path in (
                (18091, "/ready"),
                (18092, "/v2/models/rf_detr_large/versions/1/ready"),
                (18093, "/ready"),
            ):
                connection = http.client.HTTPConnection("127.0.0.1", port, timeout=2)
                try:
                    connection.request("GET", path)
                    response = connection.getresponse()
                    response.read(65536)
                    ready &= response.status == 200
                except OSError:
                    ready = False
                finally:
                    connection.close()
            if ready:
                return
            time.sleep(0.5)
        raise TimeoutError("Simulation models did not become ready")

    def stop(self):
        failures = []
        for identity in self.ids:
            result = docker("container", "rm", "-f", identity, check=False)
            if (
                result.returncode
                and docker("container", "inspect", identity, check=False).returncode == 0
            ):
                failures.append(identity)
        if failures:
            raise RuntimeError("Simulation GPU resources were not released")
        self.ids.clear()
        # Only our now-unused network; Docker refuses removal if another endpoint uses it.
        docker("network", "rm", "ndc-ai-polygon-models", check=False)