ai-toolkit/extensions_built_in/inference_engine/InferenceEngine.py

"""The inference engine as an AI Toolkit job: a resident generation server
that stays up until the job is stopped.

    job: extension
    config:
      name: inference_engine
      process:
        - type: InferenceEngine
          sqlite_db_path: ./aitk_db.db
          device: cuda
          engine:
            job_folder: output/inference_engine   # engine.json + outputs/ + assets/
            default_model: {arch: zimage}         # optional warm load

The server binds 127.0.0.1 on an ephemeral port and publishes it (with a
per-launch token) to <job_folder>/engine.json; the UI proxies to it.
"""

import json
import os
import secrets
import signal
import time
from collections import OrderedDict

import torch

from jobs.process import BaseExtensionProcess
from toolkit.paths import TOOLKIT_ROOT
from toolkit.ui_job_status import UIJobStatus


class InferenceEngineConfig:
    def __init__(self, **kwargs):
        self.job_folder = kwargs.get("job_folder", None)
        self.output_folder = kwargs.get("output_folder", None)
        self.default_model = kwargs.get("default_model", None)
        self.dtype = kwargs.get("dtype", "bf16")
        self.host = kwargs.get("host", "127.0.0.1")
        # 0 = ephemeral (recommended); a fixed port is for external clients
        self.port = int(kwargs.get("port", 0))
        self.token = kwargs.get("token", None)


class InferenceEngine(BaseExtensionProcess):
    def __init__(self, process_id: int, job, config: OrderedDict, **kwargs):
        super().__init__(process_id, job, config, **kwargs)
        self.ui = UIJobStatus(self.config.get("sqlite_db_path", "./aitk_db.db"))
        self.engine_config = InferenceEngineConfig(**self.get_conf("engine", {}))
        self.device = self.get_conf("device", "cuda")
        job_folder = self.engine_config.job_folder or os.path.join(TOOLKIT_ROOT, "output", self.name)
        self.job_folder = os.path.abspath(job_folder)
        self.output_folder = os.path.abspath(
            self.engine_config.output_folder or os.path.join(self.job_folder, "outputs")
        )
        self.endpoint_path = os.path.join(self.job_folder, "engine.json")
        self.engine = None
        self.server = None
        self.server_thread = None
        self.endpoint = None

    def run(self):
        super().run()
        from .engine import Engine
        from .server import create_app

        self.ui.update_status("running", "Starting server")
        os.makedirs(self.job_folder, exist_ok=True)
        self.engine = Engine(
            device=self.device,
            output_folder=self.output_folder,
            assets_folder=os.path.join(self.job_folder, "assets"),
            status_cb=self.ui.update_info,
            dtype=self.engine_config.dtype,
        )
        token = self.engine_config.token or secrets.token_urlsafe(24)
        self.start_server(create_app(self.engine, token), token)

        def _on_signal(signum, frame):
            # first signal: stop after the current request; if a request (or a
            # model load) is in flight, interrupt it right here on the main
            # thread so a long load cannot hold the shutdown hostage
            self.engine.stop("stopped")
            if self.engine.current is not None or self.engine.loading:
                raise KeyboardInterrupt()

        for sig in (signal.SIGINT, signal.SIGTERM):
            try:
                signal.signal(sig, _on_signal)
            except Exception:
                pass

        def _on_db_stop(reason):
            self.engine.stop(reason)
            # deliver a real SIGINT to the main thread (never os.kill: on
            # Windows that is TerminateProcess)
            signal.raise_signal(signal.SIGINT)

        self.ui.start_stop_watcher(_on_db_stop)

        if self.engine_config.default_model:
            with torch.no_grad():
                self.engine.preload(dict(self.engine_config.default_model))

        self.engine._idle_status()
        try:
            with torch.no_grad():
                self.engine.run_forever()
        finally:
            self.shutdown()

    def start_server(self, app, token: str):
        from .server import serve_in_thread

        self.server, self.server_thread, host, port = serve_in_thread(
            app, self.engine_config.host, self.engine_config.port
        )
        self.endpoint = {
            "host": host,
            "port": port,
            "token": token,
            "pid": os.getpid(),
            "job_id": self.ui.job_id,
            "device": self.device,
            "job_folder": self.job_folder,
            "output_folder": self.output_folder,
            "started_at": time.time(),
        }
        tmp = self.endpoint_path + ".tmp"
        with open(tmp, "w") as f:
            json.dump(self.endpoint, f, indent=2)
        os.replace(tmp, self.endpoint_path)
        print(f"[AITK] Inference engine listening on http://{host}:{port} (endpoint file: {self.endpoint_path})")

    def shutdown(self):
        reason = self.engine.stop_reason if self.engine else "stopped"
        print(f"[AITK] Inference engine shutting down ({reason})")
        if self.server is not None:
            self.server.should_exit = True
            if self.server_thread is not None:
                self.server_thread.join(timeout=10)
        try:
            os.remove(self.endpoint_path)
        except FileNotFoundError:
            pass
        if self.engine is not None:
            try:
                self.engine.unload()
            except Exception:
                pass
        self.ui.update_status("stopped", "Engine stopped")

    def on_error(self, e: Exception):
        super().on_error(e)
        if isinstance(e, KeyboardInterrupt):
            self.ui.update_status("stopped", "Engine stopped")
        else:
            self.ui.update_status("error", str(e))
        try:
            os.remove(self.endpoint_path)
        except Exception:
            pass