server.py

"""Local FastAPI service: STEP-file-in, classification-out. One process, models loaded once.

Designed to be launched by a C# host app as a child process. The service binds to
127.0.0.1 only; never exposes a network-facing port.

Startup protocol:
  - Bind a free port (default: OS-assigned via --port 0).
  - Once models are loaded and the server is accepting requests, write
    `READY port=<PORT>` to stdout (single line). The host reads that line to learn
    where to connect.
  - On `POST /shutdown` (or SIGTERM), drain in-flight requests and exit.

Endpoints:
  GET  /health                 -> service / model state
  POST /classify {step_path}   -> single STEP classification
  POST /classify_batch {paths} -> multiple in one call (sequential, single-process)
  POST /shutdown               -> graceful exit
"""
from __future__ import annotations
import argparse
import atexit
import os
import shutil
import sys
import tempfile
import threading
import time
from pathlib import Path
from typing import List, Optional

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import uvicorn

from heg_brep import (
    DEFAULT_PASS1_MODEL, DEFAULT_ELBOW_MODEL, DEFAULT_TEE_MODEL,
)
from heg_brep.inference import LoadedModel, TwoPassClassifier
from heg_brep.extraction import extract_step_to_npz

# Eagerly import the BRepExtractor pipeline at startup. Without this, the first
# /classify call pays ~20-30s of OCC + occwl + igl cold-import cost — terrible
# UX for an interactive viewer.
import pipeline.extract_brep_extractor_data_from_step  # noqa: F401


class State:
    classifier: Optional[TwoPassClassifier] = None
    started_at: float = 0.0
    device: str = "cpu"
    pass2_min_conf: float = 0.85
    pass2_tau: float = 0.0
    temp_artifact_baseline: set[str] = set()
    temp_cleanup_done: bool = False
    temp_cleanup_lock = threading.Lock()


app = FastAPI(title="HEG BRep component identification")


class ClassifyRequest(BaseModel):
    step_path: str
    # Optional override: persist the NPZ to disk for inspection.
    npz_keep_dir: Optional[str] = None


class ClassifyBatchRequest(BaseModel):
    step_paths: List[str]
    npz_keep_dir: Optional[str] = None


@app.get("/health")
def health():
    return {
        "status": "ok",
        "models_loaded": State.classifier is not None,
        "device": State.device,
        "uptime_sec": round(time.time() - State.started_at, 2),
    }


def _classify_one(step_path: str, npz_keep_dir: Optional[str]) -> dict:
    if State.classifier is None:
        raise HTTPException(status_code=503, detail="models not loaded yet")
    sp = Path(step_path).expanduser()
    if not sp.exists():
        return {"step_path": str(sp), "status": "error", "error": "step_path not found"}
    temp_npz_dir = npz_keep_dir is None
    out_dir = Path(npz_keep_dir).expanduser().resolve() if npz_keep_dir \
              else Path(tempfile.mkdtemp(prefix="heg_brep_npz_"))
    try:
        try:
            npz = extract_step_to_npz(sp, out_dir)
        except Exception as exc:
            return {"step_path": str(sp), "status": "extraction_failed", "error": str(exc)[:500]}
        try:
            result = State.classifier.classify_npz(npz)
        except Exception as exc:
            return {"step_path": str(sp), "status": "inference_failed", "error": str(exc)[:500],
                    "npz_path": str(npz)}
        result.update({"step_path": str(sp), "status": "ok", "npz_path": str(npz)})
        return result
    finally:
        if temp_npz_dir:
            shutil.rmtree(out_dir, ignore_errors=True)


def _snapshot_service_temp_artifacts() -> set[str]:
    temp_root = Path(tempfile.gettempdir())
    out: set[str] = set()
    for p in temp_root.glob("heg_brep_npz_*"):
        out.add(str(p.resolve()))
    for p in temp_root.glob("torch_geometric.nn.conv.sage_conv_SAGEConv_propagate_*.py"):
        out.add(str(p.resolve()))
    pycache = temp_root / "__pycache__"
    if pycache.exists():
        for p in pycache.glob("torch_geometric.nn.conv.sage_conv_SAGEConv_propagate_*.pyc"):
            out.add(str(p.resolve()))
    return out


def _remove_path(p: Path) -> None:
    if p.is_dir():
        shutil.rmtree(p, ignore_errors=True)
    else:
        p.unlink(missing_ok=True)


def _cleanup_service_temp_artifacts() -> None:
    with State.temp_cleanup_lock:
        if State.temp_cleanup_done:
            return
        current = _snapshot_service_temp_artifacts()
        to_remove = sorted(current - State.temp_artifact_baseline)
        for raw in to_remove:
            _remove_path(Path(raw))
        State.temp_cleanup_done = True


@app.post("/classify")
def classify(req: ClassifyRequest):
    return _classify_one(req.step_path, req.npz_keep_dir)


@app.post("/classify_batch")
def classify_batch(req: ClassifyBatchRequest):
    return {"results": [_classify_one(p, req.npz_keep_dir) for p in req.step_paths]}


@app.post("/shutdown")
def shutdown():
    def _exit_soon():
        time.sleep(0.2)
        try:
            _cleanup_service_temp_artifacts()
        finally:
            os._exit(0)
    threading.Thread(target=_exit_soon, daemon=True).start()
    return {"status": "shutting_down"}


def _load_models(args) -> None:
    pass1 = LoadedModel(Path(args.pass1_model), device=args.device)
    elbow = LoadedModel(Path(args.elbow_model), device=args.device) if Path(args.elbow_model).exists() else None
    tee   = LoadedModel(Path(args.tee_model),   device=args.device) if Path(args.tee_model).exists()   else None
    State.classifier = TwoPassClassifier(
        pass1=pass1, elbow=elbow, tee=tee,
        pass2_min_conf=args.pass2_min_conf, pass2_tau=args.pass2_tau,
    )
    State.device = args.device
    State.pass2_min_conf = args.pass2_min_conf
    State.pass2_tau = args.pass2_tau


def parse_args() -> argparse.Namespace:
    ap = argparse.ArgumentParser(description="HEG BRep classification service")
    ap.add_argument("--host", default="127.0.0.1")
    ap.add_argument("--port", type=int, default=0,
                    help="Port to bind. 0 = let the OS pick. The chosen port is "
                         "printed to stdout as a single line `READY port=<N>`.")
    ap.add_argument("--pass1_model", default=str(DEFAULT_PASS1_MODEL))
    ap.add_argument("--elbow_model", default=str(DEFAULT_ELBOW_MODEL))
    ap.add_argument("--tee_model",   default=str(DEFAULT_TEE_MODEL))
    ap.add_argument("--device", default="cpu", choices=["cpu", "cuda"])
    ap.add_argument("--pass2_min_conf", type=float, default=0.85)
    ap.add_argument("--pass2_tau", type=float, default=0.0)
    return ap.parse_args()


def main() -> int:
    args = parse_args()
    State.temp_artifact_baseline = _snapshot_service_temp_artifacts()
    atexit.register(_cleanup_service_temp_artifacts)
    State.started_at = time.time()
    print(f"[heg_brep] loading models from:", file=sys.stderr)
    print(f"  pass1: {args.pass1_model}", file=sys.stderr)
    print(f"  elbow: {args.elbow_model}", file=sys.stderr)
    print(f"  tee  : {args.tee_model}",   file=sys.stderr)
    print(f"  device: {args.device}",     file=sys.stderr)
    t0 = time.time()
    _load_models(args)
    print(f"[heg_brep] models loaded in {time.time() - t0:.1f}s", file=sys.stderr)

    # Bind socket up front so we can tell the parent which port we got.
    import socket
    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    sock.bind((args.host, int(args.port)))
    bound_port = sock.getsockname()[1]
    print(f"READY port={bound_port}", flush=True)

    config = uvicorn.Config(app=app, host=args.host, port=bound_port, log_level="warning")
    server = uvicorn.Server(config)
    # uvicorn doesn't accept an already-bound socket in the simple Config API; close
    # ours and let uvicorn rebind. Race window is fine because we're loopback only.
    sock.close()
    try:
        server.run()
    finally:
        _cleanup_service_temp_artifacts()
    return 0


if __name__ == "__main__":
    sys.exit(main())