@mcoda/mswarm 0.1.101 → 0.1.102

package/scripts/stable-audio-3-openai-server.py

@@ -0,0 +1,738 @@
+#!/usr/bin/env python3
+"""Loopback-only OpenAI-shaped HTTP wrapper for Stable Audio 3 TensorRT."""
+
+from __future__ import annotations
+
+import argparse
+import base64
+import io
+import ipaddress
+import json
+import logging
+import math
+import os
+import secrets
+import signal
+import subprocess
+import sys
+import tempfile
+import threading
+import time
+import wave
+from dataclasses import dataclass, field
+from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
+from pathlib import Path
+from socketserver import TCPServer
+from typing import Any, Protocol
+
+
+DEFAULT_HOST = "127.0.0.1"
+DEFAULT_PORT = 11446
+DEFAULT_MODEL_ID = "stable-audio-3-local"
+DEFAULT_RUNTIME_ROOT = (
+    "/mnt/githubActions/piriatlas/tools/stable-audio-3/optimized/tensorRT"
+)
+SAMPLE_RATE = 44_100
+CHANNELS = 2
+SAMPLE_WIDTH_BYTES = 2
+ALLOWED_FIELDS = {
+    "model",
+    "prompt",
+    "duration_seconds",
+    "response_format",
+    "n",
+    "seed",
+    "steps",
+}
+
+
+class AudioBackend(Protocol):
+    def generate(
+        self,
+        prompt: str,
+        *,
+        seconds: float,
+        steps: int,
+        seed: int | None,
+    ) -> tuple[bytes, dict[str, Any]]: ...
+
+
+class ApiError(Exception):
+    def __init__(
+        self,
+        status: int,
+        message: str,
+        *,
+        error_type: str = "invalid_request_error",
+        param: str | None = None,
+        code: str | None = None,
+    ) -> None:
+        super().__init__(message)
+        self.status = status
+        self.message = message
+        self.error_type = error_type
+        self.param = param
+        self.code = code
+
+
+class MockBackend:
+    """Dependency-free backend used only by the HTTP contract test."""
+
+    def generate(
+        self,
+        prompt: str,
+        *,
+        seconds: float,
+        steps: int,
+        seed: int | None,
+    ) -> tuple[bytes, dict[str, Any]]:
+        del prompt, steps
+        sample_count = int(round(seconds * SAMPLE_RATE))
+        output = io.BytesIO()
+        with wave.open(output, "wb") as wav_file:
+            wav_file.setnchannels(CHANNELS)
+            wav_file.setsampwidth(SAMPLE_WIDTH_BYTES)
+            wav_file.setframerate(SAMPLE_RATE)
+            wav_file.writeframes(b"\x00" * (sample_count * CHANNELS * SAMPLE_WIDTH_BYTES))
+        return output.getvalue(), {"seed": 0 if seed is None else seed}
+
+
+class SubprocessBackend:
+    """Runs one canonical Stable Audio CLI process per request.
+
+    Upstream documents that the TensorRT decoder workspace evolves after the
+    first inference in a process. A fresh process per request preserves the
+    canonical CLI's deterministic behavior while the HTTP listener remains
+    resident and bounded.
+    """
+
+    def __init__(
+        self,
+        runtime_root: Path,
+        *,
+        dit: str,
+        decoder: str,
+        precision: str | None,
+        models_dir: Path | None,
+        engine_arch: str,
+        max_output_bytes: int,
+        timeout_seconds: float,
+    ) -> None:
+        self.runtime_root = runtime_root
+        self.python = runtime_root / ".venv" / "bin" / "python"
+        self.entrypoint = runtime_root / "scripts" / "sa3_trt.py"
+        if not self.python.is_file() or not os.access(self.python, os.X_OK):
+            raise RuntimeError(f"Stable Audio Python is missing or not executable: {self.python}")
+        if not self.entrypoint.is_file():
+            raise RuntimeError(f"Stable Audio TensorRT entrypoint not found at {self.entrypoint}")
+        self.dit = dit
+        self.decoder = decoder
+        self.precision = precision
+        self.models_dir = models_dir
+        self.max_output_bytes = max_output_bytes
+        self.timeout_seconds = timeout_seconds
+        precision_name = precision or "fp16mixed"
+        dit_directory = {
+            "sm-music": "sa3-sm-music",
+            "sm-sfx": "sa3-sm-sfx",
+            "medium": "sa3-m",
+        }[dit]
+        decoder_filename = (
+            "dec_dynamic_fp32.trt"
+            if precision_name == "fp32"
+            else {
+                "same-s": "dec_dynamic_bf16.trt",
+                "same-l": "dec_dynamic_triton_swa.trt",
+            }[decoder]
+        )
+        models_root = models_dir or runtime_root / "models"
+        required_files = (
+            runtime_root / "scripts" / "tokenizer.json",
+            models_root / engine_arch / "t5gemma" / "t5gemma_fp16mixed.trt",
+            models_root / engine_arch / dit_directory / f"dit_{precision_name}.trt",
+            models_root / engine_arch / decoder / decoder_filename,
+        )
+        missing = [path for path in required_files if not path.is_file() or path.stat().st_size == 0]
+        if missing:
+            raise RuntimeError(f"Stable Audio runtime file is missing or empty: {missing[0]}")
+
+    def generate(
+        self,
+        prompt: str,
+        *,
+        seconds: float,
+        steps: int,
+        seed: int | None,
+    ) -> tuple[bytes, dict[str, Any]]:
+        effective_seed = seed if seed is not None else secrets.randbelow(2**31)
+        with tempfile.TemporaryDirectory(prefix="stable-audio-3-") as temp_dir:
+            output_path = Path(temp_dir) / "generation.wav"
+            command = [
+                str(self.python),
+                str(self.entrypoint),
+                f"--prompt={prompt}",
+                "--dit",
+                self.dit,
+                "--decoder",
+                self.decoder,
+                "--seconds",
+                str(seconds),
+                "--steps",
+                str(steps),
+                "--seed",
+                str(effective_seed),
+                "--quiet",
+                "--out",
+                str(output_path),
+            ]
+            if self.precision:
+                command.extend(("--precision", self.precision))
+            if self.models_dir:
+                command.extend(("--models-dir", str(self.models_dir)))
+            process = subprocess.Popen(
+                    command,
+                    cwd=self.runtime_root,
+                    stdin=subprocess.DEVNULL,
+                    stdout=subprocess.DEVNULL,
+                    stderr=subprocess.DEVNULL,
+                    start_new_session=True,
+            )
+            try:
+                return_code = process.wait(timeout=self.timeout_seconds)
+            except subprocess.TimeoutExpired as error:
+                try:
+                    os.killpg(process.pid, signal.SIGTERM)
+                    process.wait(timeout=5)
+                except subprocess.TimeoutExpired:
+                    os.killpg(process.pid, signal.SIGKILL)
+                    process.wait()
+                except ProcessLookupError:
+                    pass
+                raise RuntimeError("Stable Audio CLI timed out") from error
+            if return_code != 0:
+                raise RuntimeError(
+                    f"Stable Audio CLI exited with status {return_code}"
+                )
+            try:
+                if output_path.stat().st_size > self.max_output_bytes:
+                    raise RuntimeError("Stable Audio CLI output exceeds the configured limit")
+                wav_data = output_path.read_bytes()
+            except OSError as error:
+                raise RuntimeError("Stable Audio CLI did not produce a WAV file") from error
+        _validate_wav(wav_data, seconds)
+        return wav_data, {"seed": effective_seed}
+
+
+@dataclass
+class ServerState:
+    backend: AudioBackend
+    backend_name: str
+    model_id: str
+    dit: str
+    decoder: str
+    max_request_bytes: int
+    max_output_bytes: int
+    max_prompt_chars: int
+    min_duration_seconds: float
+    max_duration_seconds: float
+    max_steps: int
+    default_steps: int
+    fatal_error: threading.Event = field(default_factory=threading.Event)
+    inference_slots: threading.BoundedSemaphore = field(
+        default_factory=lambda: threading.BoundedSemaphore(1)
+    )
+
+
+class AudioHttpServer(ThreadingHTTPServer):
+    daemon_threads = True
+    request_queue_size = 8
+
+    def server_bind(self) -> None:
+        # HTTPServer.server_bind performs a reverse-DNS lookup via getfqdn().
+        # That lookup is unnecessary for this loopback-only service and can
+        # stall startup on hosts whose resolver does not map loopback quickly.
+        TCPServer.server_bind(self)
+        host, port = self.server_address[:2]
+        self.server_name = str(host)
+        self.server_port = int(port)
+
+
+def _is_int(value: Any) -> bool:
+    return isinstance(value, int) and not isinstance(value, bool)
+
+
+def _parse_generation_request(payload: Any, state: ServerState) -> dict[str, Any]:
+    if not isinstance(payload, dict):
+        raise ApiError(400, "request body must be a JSON object", code="invalid_json")
+
+    unknown = sorted(set(payload) - ALLOWED_FIELDS)
+    if unknown:
+        raise ApiError(
+            400,
+            f"unsupported request parameter: {unknown[0]}",
+            param=unknown[0],
+            code="unsupported_parameter",
+        )
+
+    model = payload.get("model")
+    if model != state.model_id:
+        raise ApiError(
+            400,
+            f"model must be {state.model_id}",
+            param="model",
+            code="model_not_found",
+        )
+
+    prompt = payload.get("prompt")
+    if not isinstance(prompt, str) or not prompt.strip():
+        raise ApiError(
+            400,
+            "prompt must be a non-empty string",
+            param="prompt",
+            code="invalid_prompt",
+        )
+    if len(prompt) > state.max_prompt_chars:
+        raise ApiError(
+            400,
+            "prompt exceeds the configured limit",
+            param="prompt",
+            code="prompt_limit_exceeded",
+        )
+    if "\x00" in prompt:
+        raise ApiError(
+            400,
+            "prompt must not contain a null character",
+            param="prompt",
+            code="invalid_prompt",
+        )
+    try:
+        prompt.encode("utf-8", errors="strict")
+    except UnicodeEncodeError as error:
+        raise ApiError(
+            400,
+            "prompt must be valid UTF-8 text",
+            param="prompt",
+            code="invalid_prompt",
+        ) from error
+
+    duration = payload.get("duration_seconds")
+    if isinstance(duration, bool) or not isinstance(duration, (int, float)):
+        raise ApiError(
+            400,
+            "duration_seconds must be a number",
+            param="duration_seconds",
+            code="invalid_duration",
+        )
+    duration = float(duration)
+    if not math.isfinite(duration):
+        raise ApiError(
+            400,
+            "duration_seconds must be finite",
+            param="duration_seconds",
+            code="invalid_duration",
+        )
+    if not state.min_duration_seconds <= duration <= state.max_duration_seconds:
+        raise ApiError(
+            400,
+            "duration_seconds is outside the configured range",
+            param="duration_seconds",
+            code="duration_limit_exceeded",
+        )
+
+    response_format = payload.get("response_format", "wav")
+    if response_format != "wav":
+        raise ApiError(
+            400,
+            "response_format must be wav",
+            param="response_format",
+            code="unsupported_response_format",
+        )
+
+    n = payload.get("n", 1)
+    if not _is_int(n) or n != 1:
+        raise ApiError(400, "n must be 1", param="n", code="invalid_n")
+
+    seed = payload.get("seed")
+    if seed is not None and (not _is_int(seed) or not 0 <= seed <= 2**31 - 1):
+        raise ApiError(
+            400,
+            "seed must be an integer between 0 and 2147483647",
+            param="seed",
+            code="invalid_seed",
+        )
+
+    steps = payload.get("steps", state.default_steps)
+    if not _is_int(steps) or not 1 <= steps <= state.max_steps:
+        raise ApiError(
+            400,
+            f"steps must be an integer between 1 and {state.max_steps}",
+            param="steps",
+            code="invalid_steps",
+        )
+
+    return {
+        "prompt": prompt,
+        "duration_seconds": duration,
+        "steps": steps,
+        "seed": seed,
+    }
+
+
+def _validate_wav(wav_data: bytes, duration_seconds: float) -> None:
+    expected_samples = int(round(duration_seconds * SAMPLE_RATE))
+    expected_pcm_bytes = expected_samples * CHANNELS * SAMPLE_WIDTH_BYTES
+    try:
+        with wave.open(io.BytesIO(wav_data), "rb") as wav_file:
+            if wav_file.getnchannels() != CHANNELS:
+                raise RuntimeError("backend WAV must be stereo")
+            if wav_file.getsampwidth() != SAMPLE_WIDTH_BYTES:
+                raise RuntimeError("backend WAV must use signed 16-bit samples")
+            if wav_file.getframerate() != SAMPLE_RATE:
+                raise RuntimeError(f"backend WAV must use {SAMPLE_RATE} Hz")
+            if wav_file.getnframes() != expected_samples:
+                raise RuntimeError(
+                    f"backend WAV has {wav_file.getnframes()} frames; expected {expected_samples}"
+                )
+            pcm_data = wav_file.readframes(expected_samples + 1)
+            if len(pcm_data) != expected_pcm_bytes:
+                raise RuntimeError(
+                    f"backend WAV has {len(pcm_data)} PCM bytes; expected {expected_pcm_bytes}"
+                )
+    except (EOFError, wave.Error) as error:
+        raise RuntimeError("backend returned an invalid WAV file") from error
+    if len(wav_data) != 44 + expected_pcm_bytes:
+        raise RuntimeError("backend WAV contains truncated or trailing data")
+
+
+def _handler_for(state: ServerState) -> type[BaseHTTPRequestHandler]:
+    class Handler(BaseHTTPRequestHandler):
+        server_version = "stable-audio-3-openai/1"
+
+        def log_message(self, format_string: str, *args: Any) -> None:
+            del format_string, args
+
+        def _send_json(self, status: int, payload: dict[str, Any]) -> None:
+            encoded = json.dumps(payload, separators=(",", ":")).encode("utf-8")
+            self.send_response(status)
+            self.send_header("content-type", "application/json")
+            self.send_header("content-length", str(len(encoded)))
+            self.send_header("cache-control", "no-store")
+            self.end_headers()
+            try:
+                self.wfile.write(encoded)
+            except BrokenPipeError:
+                pass
+
+        def _send_error(self, error: ApiError) -> None:
+            self._send_json(
+                error.status,
+                {
+                    "error": {
+                        "message": error.message,
+                        "type": error.error_type,
+                        "param": error.param,
+                        "code": error.code,
+                    }
+                },
+            )
+
+        def do_GET(self) -> None:
+            if self.path == "/healthz":
+                if state.fatal_error.is_set():
+                    self._send_json(503, {"status": "unhealthy", "model": state.model_id})
+                    return
+                self._send_json(
+                    200,
+                    {
+                        "status": "ok",
+                        "model": state.model_id,
+                        "backend": state.backend_name,
+                        "dit": state.dit,
+                        "decoder": state.decoder,
+                    },
+                )
+                return
+            if self.path == "/v1/models":
+                self._send_json(
+                    200,
+                    {
+                        "object": "list",
+                        "data": [
+                            {
+                                "id": state.model_id,
+                                "object": "model",
+                                "owned_by": "local",
+                            }
+                        ],
+                    },
+                )
+                return
+            self._send_error(ApiError(404, "not found", error_type="not_found_error"))
+
+        def do_POST(self) -> None:
+            if self.path != "/v1/audio/generations":
+                self._send_error(ApiError(404, "not found", error_type="not_found_error"))
+                return
+
+            try:
+                if self.headers.get("transfer-encoding"):
+                    raise ApiError(
+                        400,
+                        "chunked request bodies are not supported",
+                        code="unsupported_transfer_encoding",
+                    )
+                content_type = self.headers.get("content-type", "")
+                if content_type.split(";", 1)[0].strip().lower() != "application/json":
+                    raise ApiError(
+                        400,
+                        "content-type must be application/json",
+                        code="invalid_content_type",
+                    )
+                content_length_header = self.headers.get("content-length")
+                if content_length_header is None:
+                    raise ApiError(411, "content-length is required", code="length_required")
+                try:
+                    content_length = int(content_length_header)
+                except ValueError as error:
+                    raise ApiError(
+                        400, "content-length is invalid", code="invalid_content_length"
+                    ) from error
+                if content_length < 1:
+                    raise ApiError(400, "request body is empty", code="invalid_json")
+                if content_length > state.max_request_bytes:
+                    raise ApiError(
+                        413,
+                        "request body exceeds the configured limit",
+                        code="request_too_large",
+                    )
+                raw_body = self.rfile.read(content_length)
+                if len(raw_body) != content_length:
+                    raise ApiError(400, "request body is incomplete", code="invalid_json")
+                try:
+                    payload = json.loads(raw_body)
+                except (UnicodeDecodeError, json.JSONDecodeError) as error:
+                    raise ApiError(400, "request body is invalid JSON", code="invalid_json") from error
+                request = _parse_generation_request(payload, state)
+            except ApiError as error:
+                self._send_error(error)
+                return
+
+            if not state.inference_slots.acquire(blocking=False):
+                self._send_error(
+                    ApiError(
+                        429,
+                        "the local audio generator is busy",
+                        error_type="rate_limit_error",
+                        code="local_concurrency_saturated",
+                    )
+                )
+                return
+
+            started_at = time.monotonic()
+            try:
+                wav_data, timing = state.backend.generate(
+                    request["prompt"],
+                    seconds=request["duration_seconds"],
+                    steps=request["steps"],
+                    seed=request["seed"],
+                )
+                _validate_wav(wav_data, request["duration_seconds"])
+                if len(wav_data) > state.max_output_bytes:
+                    raise ApiError(
+                        413,
+                        "generated audio exceeds the configured output limit",
+                        code="output_too_large",
+                    )
+                effective_seed = timing.get("seed") if isinstance(timing, dict) else None
+                if not _is_int(effective_seed):
+                    raise RuntimeError("backend did not report an effective seed")
+                response = {
+                    "created": int(time.time()),
+                    "model": state.model_id,
+                    "data": [
+                        {
+                            "b64_audio": base64.b64encode(wav_data).decode("ascii"),
+                            "mime_type": "audio/wav",
+                            "duration_seconds": request["duration_seconds"],
+                            "sample_rate": SAMPLE_RATE,
+                            "seed": effective_seed,
+                        }
+                    ],
+                }
+                self._send_json(200, response)
+                logging.info(
+                    "generation completed duration_seconds=%.3f steps=%d elapsed_ms=%d",
+                    request["duration_seconds"],
+                    request["steps"],
+                    int((time.monotonic() - started_at) * 1000),
+                )
+            except ApiError as error:
+                self._send_error(error)
+            except Exception as error:
+                logging.error("generation failed error_type=%s", type(error).__name__)
+                self._send_error(
+                    ApiError(
+                        500,
+                        "audio generation failed",
+                        error_type="server_error",
+                        code="generation_failed",
+                    )
+                )
+                state.fatal_error.set()
+                threading.Thread(target=self.server.shutdown, daemon=True).start()
+            finally:
+                state.inference_slots.release()
+
+    return Handler
+
+
+def _load_subprocess_backend(args: argparse.Namespace) -> AudioBackend:
+    runtime_root = Path(args.runtime_root).expanduser().resolve()
+    models_dir = (
+        Path(args.models_dir).expanduser().resolve() if args.models_dir else None
+    )
+    return SubprocessBackend(
+        runtime_root,
+        dit=args.dit,
+        decoder=args.decoder,
+        precision=args.precision,
+        models_dir=models_dir,
+        engine_arch=args.engine_arch,
+        max_output_bytes=args.max_output_bytes,
+        timeout_seconds=args.inference_timeout_seconds,
+    )
+
+
+def _parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--host", default=DEFAULT_HOST)
+    parser.add_argument("--port", type=int, default=DEFAULT_PORT)
+    parser.add_argument("--backend", choices=("subprocess", "mock"), default="subprocess")
+    parser.add_argument("--runtime-root", default=os.environ.get("STABLE_AUDIO_3_TRT_ROOT", DEFAULT_RUNTIME_ROOT))
+    parser.add_argument("--models-dir")
+    parser.add_argument(
+        "--engine-arch",
+        default=os.environ.get("STABLE_AUDIO_3_ENGINE_ARCH", "sm_86"),
+    )
+    parser.add_argument("--dit", choices=("sm-music", "sm-sfx", "medium"), default="sm-music")
+    parser.add_argument("--decoder", choices=("same-s", "same-l"), default="same-s")
+    parser.add_argument("--precision", choices=("fp16mixed", "bf16", "fp8", "fp32"))
+    parser.add_argument("--model-id", default=DEFAULT_MODEL_ID)
+    parser.add_argument("--max-request-bytes", type=int, default=131_072)
+    parser.add_argument("--max-output-bytes", type=int, default=67_108_864)
+    parser.add_argument("--max-prompt-chars", type=int, default=8_192)
+    parser.add_argument("--min-duration-seconds", type=float, default=3.0)
+    parser.add_argument("--max-duration-seconds", type=float, default=120.0)
+    parser.add_argument("--default-steps", type=int, default=8)
+    parser.add_argument("--max-steps", type=int, default=32)
+    parser.add_argument("--inference-timeout-seconds", type=float, default=900.0)
+    parser.add_argument(
+        "--startup-self-test",
+        action=argparse.BooleanOptionalAction,
+        default=True,
+    )
+    return parser.parse_args()
+
+
+def main() -> int:
+    args = _parse_args()
+    try:
+        address = ipaddress.ip_address(args.host)
+    except ValueError:
+        if args.host != "localhost":
+            raise SystemExit("--host must be a loopback IP address or localhost")
+    else:
+        if not address.is_loopback:
+            raise SystemExit("--host must resolve to a loopback address")
+    if not 1 <= args.port <= 65_535 and args.port != 0:
+        raise SystemExit("--port must be 0 or between 1 and 65535")
+    if args.decoder != ("same-l" if args.dit == "medium" else "same-s"):
+        raise SystemExit("medium requires same-l; sm-music and sm-sfx require same-s")
+    if args.precision in ("bf16", "fp8") and args.dit != "medium":
+        raise SystemExit("bf16 and fp8 precision are available only for medium")
+    if not 1 <= args.default_steps <= args.max_steps:
+        raise SystemExit("--default-steps must be within the configured steps range")
+    if (
+        not math.isfinite(args.min_duration_seconds)
+        or not math.isfinite(args.max_duration_seconds)
+        or args.min_duration_seconds < 3
+        or args.max_duration_seconds < args.min_duration_seconds
+    ):
+        raise SystemExit("duration limits are invalid")
+    if not math.isfinite(args.inference_timeout_seconds) or args.inference_timeout_seconds <= 0:
+        raise SystemExit("--inference-timeout-seconds must be positive")
+    if args.max_request_bytes <= 0 or args.max_output_bytes <= 0 or args.max_prompt_chars <= 0:
+        raise SystemExit("request, output, and prompt limits must be positive")
+    if not args.model_id.strip():
+        raise SystemExit("--model-id must be non-empty")
+    if not args.startup_self_test and args.backend != "mock":
+        raise SystemExit("--no-startup-self-test is permitted only with --backend mock")
+
+    logging.basicConfig(
+        level=logging.INFO,
+        format="%(asctime)s %(levelname)s %(message)s",
+    )
+    logging.info("loading backend=%s dit=%s decoder=%s", args.backend, args.dit, args.decoder)
+    backend: AudioBackend = (
+        MockBackend() if args.backend == "mock" else _load_subprocess_backend(args)
+    )
+    if args.startup_self_test:
+        logging.info("running startup inference self-test")
+        probe_seconds = max(3.0, args.min_duration_seconds)
+        probe_wav, probe_timing = backend.generate(
+            "readiness check: a clean warm analog synthesizer tone",
+            seconds=probe_seconds,
+            steps=args.default_steps,
+            seed=1_001,
+        )
+        _validate_wav(probe_wav, probe_seconds)
+        if len(probe_wav) > args.max_output_bytes:
+            raise RuntimeError("startup self-test output exceeds the configured limit")
+        if not _is_int(probe_timing.get("seed") if isinstance(probe_timing, dict) else None):
+            raise RuntimeError("startup self-test did not report an effective seed")
+        logging.info("startup inference self-test passed")
+    state = ServerState(
+        backend=backend,
+        backend_name=args.backend,
+        model_id=args.model_id,
+        dit=args.dit,
+        decoder=args.decoder,
+        max_request_bytes=args.max_request_bytes,
+        max_output_bytes=args.max_output_bytes,
+        max_prompt_chars=args.max_prompt_chars,
+        min_duration_seconds=args.min_duration_seconds,
+        max_duration_seconds=args.max_duration_seconds,
+        max_steps=args.max_steps,
+        default_steps=args.default_steps,
+    )
+    server = AudioHttpServer((args.host, args.port), _handler_for(state))
+
+    def request_shutdown(signum: int, frame: Any) -> None:
+        del signum, frame
+        threading.Thread(target=server.shutdown, daemon=True).start()
+
+    signal.signal(signal.SIGTERM, request_shutdown)
+    signal.signal(signal.SIGINT, request_shutdown)
+    host, port = server.server_address[:2]
+    print(
+        json.dumps(
+            {
+                "event": "listening",
+                "host": host,
+                "port": port,
+                "model": state.model_id,
+                "backend": state.backend_name,
+            },
+            separators=(",", ":"),
+        ),
+        flush=True,
+    )
+    try:
+        server.serve_forever(poll_interval=0.5)
+    finally:
+        server.server_close()
+    return 1 if state.fatal_error.is_set() else 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())