h3-blackwell-runtime/tools/serve_hot_runtime.py

"""Serve a resident prompt-only H3 runtime over a small JSON HTTP API."""

from __future__ import annotations

import argparse
import json
import threading
import time
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path
from urllib.parse import urlparse

from h3_blackwell_runtime.attention import AVAILABLE_BACKENDS
from h3_blackwell_runtime.runtime import H3HotRuntime, RuntimeConfig


parser = argparse.ArgumentParser()
parser.add_argument("--host", default="0.0.0.0")
parser.add_argument("--port", type=int, default=8000)
parser.add_argument("--attention", choices=AVAILABLE_BACKENDS, default="sage2", help="Initial attention backend. Requests can switch with the JSON attention field.")
parser.add_argument("--vae-dtype", choices=("float32", "float16", "bfloat16"), default="float16")
parser.add_argument("--vae-tile-size", type=int, default=256)
parser.add_argument("--mlp-chunks", type=int, default=1)
parser.add_argument("--mlp-chunk-threshold", type=int, default=4096)
parser.add_argument("--warmup", action="store_true", help="Run a tiny generation before accepting traffic.")
parser.add_argument("--warmup-output", type=Path, default=Path("/output/h3-blackwell-runtime/hot-runtime-warmup.mp4"))
args = parser.parse_args()

runtime = H3HotRuntime(RuntimeConfig(attention=args.attention, vae_dtype=args.vae_dtype, vae_tile_size=args.vae_tile_size, mlp_chunks=args.mlp_chunks, mlp_chunk_threshold=args.mlp_chunk_threshold))
runtime_lock = threading.Lock()
warmup_result = None
if args.warmup:
    warmup_result = runtime.generate(
        prompt="A small warmup cat blinks in soft light.",
        output=args.warmup_output,
        width=320,
        height=192,
        frames=22,
        steps=2,
        seed=440501,
        mux_audio=True,
    )


def service_status() -> dict:
    return {
        "ready": True,
        "attention_backends": list(AVAILABLE_BACKENDS),
        "runtime": runtime.status(),
    }


def write_json(handler: BaseHTTPRequestHandler, status: int, payload: dict) -> None:
    body = json.dumps(payload, indent=2).encode("utf-8")
    handler.send_response(status)
    handler.send_header("Content-Type", "application/json")
    handler.send_header("Content-Length", str(len(body)))
    handler.end_headers()
    handler.wfile.write(body)


class Handler(BaseHTTPRequestHandler):
    def log_message(self, format: str, *args) -> None:
        return

    def do_GET(self) -> None:
        path = urlparse(self.path).path
        if path in {"/health", "/ready"}:
            write_json(self, 200, {**service_status(), "warmup_result": warmup_result})
            return
        write_json(self, 404, {"error": "not found"})

    def do_POST(self) -> None:
        path = urlparse(self.path).path
        if path != "/generate":
            write_json(self, 404, {"error": "not found"})
            return
        try:
            length = int(self.headers.get("Content-Length", "0"))
            payload = json.loads(self.rfile.read(length).decode("utf-8")) if length else {}
            prompt = payload["prompt"]
            output = payload["output"]
            width = int(payload.get("width", 960))
            height = int(payload.get("height", 544))
            frames = int(payload.get("frames", 124))
            steps = int(payload.get("steps", 12))
            seed = int(payload.get("seed", 440407))
            attention = payload.get("attention")
            if attention is not None and attention not in AVAILABLE_BACKENDS:
                write_json(self, 400, {"error": "unsupported attention", "attention": attention, "available": list(AVAILABLE_BACKENDS)})
                return
            mux_audio = bool(payload.get("mux_audio", True))
            ffmpeg_loglevel = payload.get("ffmpeg_loglevel", "error")
            save_latent = payload.get("save_latent")
            started = time.perf_counter()
            with runtime_lock:
                result = runtime.generate(
                    prompt=prompt,
                    output=output,
                    width=width,
                    height=height,
                    frames=frames,
                    steps=steps,
                    seed=seed,
                    attention=attention,
                    mux_audio=mux_audio,
                    ffmpeg_loglevel=ffmpeg_loglevel,
                    save_latent=save_latent,
                )
            result["wall_seconds"] = time.perf_counter() - started
            write_json(self, 200, result)
        except Exception as exc:
            write_json(self, 500, {"error": type(exc).__name__, "message": str(exc)})


server = ThreadingHTTPServer((args.host, args.port), Handler)
print(json.dumps({"serving": True, "host": args.host, "port": args.port, **service_status(), "warmup_result": warmup_result}, indent=2), flush=True)
server.serve_forever()