Implement AI Host functionality with model caching and job management

- Added ai_host.py to handle AI tasks with FastAPI integration. - Introduced HostState and ModelCache classes for managing client connections and model loading. - Implemented job submission and cancellation endpoints in the AI host. - Enhanced tests for AI host functionality, including client tracking and job cancellation. - Updated existing tests to accommodate changes in the operator and command structure. - Refactored runtime commands to use request-based architecture instead of command-line arguments. - Improved model loading logic to handle CUDA memory management and reuse models efficiently.
2026-08-04 07:35:40 +08:00
parent d694980da7
commit 838af0ee96
16 changed files with 2112 additions and 225 deletions
--- a/src/snapform/runtime.py
+++ b/src/snapform/runtime.py
@@ -1,10 +1,16 @@
 import json
 import os
+import secrets
 import shutil
+import socket
 import subprocess
+import time
 import uuid
+from dataclasses import dataclass
 from datetime import datetime
 from pathlib import Path
+from urllib.error import HTTPError, URLError
+from urllib.request import ProxyHandler, Request, build_opener

 import bpy

@@ -13,10 +19,40 @@ from .models import BEN2_MODEL, get_downloadable_model, selected_model_key

 DA3_COMMIT = "3d835ec1a5802d64a8b8b15f817a1ab54809bfe4"
 BEN2_COMMIT = "2c99a5da477b5523585bfa5c893888a6e818a8f6"
-RUNTIME_SCHEMA = 3
+RUNTIME_SCHEMA = 4
+HOST_SCHEMA = 2
+HOST_ADDRESS = "127.0.0.1"
 DEFAULT_STORAGE_ROOT = Path.home() / ".snapform"


+@dataclass(frozen=True)
+class HostJobRequest:
+    payload: dict
+
+
+@dataclass
+class HostConnection:
+    port: int
+    token: str
+    pid: int
+    instance_id: str
+
+
+@dataclass
+class LaunchedHost:
+    process: subprocess.Popen
+    log_file: object
+
+
+_host_connection = None
+_launched_host = None
+_registered_instance_id = None
+_client_id = uuid.uuid4().hex
+_host_auto_start = False
+_host_snapshot = {"state": "STOPPED", "message": "Stopped"}
+_local_opener = build_opener(ProxyHandler({}))
+
+
 def preferences():
    addon = bpy.context.preferences.addons.get(__package__)
    return addon.preferences if addon is not None else None
@@ -124,6 +160,8 @@ def install_runtime_command(status_path):
        executable,
        "--runtime-dir",
        str(runtime_directory()),
+        "--models-dir",
+        str(models_directory()),
        "--backend",
        backend.lower(),
        "--da3-commit",
@@ -135,6 +173,419 @@ def install_runtime_command(status_path):
    ]


+def host_log_path():
+    return storage_root() / "ai-host.log"
+
+
+def host_discovery_path():
+    return storage_root() / "host.json"
+
+
+def host_lock_path():
+    return storage_root() / "host.lock"
+
+
+def _available_host_port():
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as listener:
+        listener.bind((HOST_ADDRESS, 0))
+        return listener.getsockname()[1]
+
+
+def _host_url(connection, path):
+    return f"http://{HOST_ADDRESS}:{connection.port}{path}"
+
+
+def _host_http_request(connection, method, path, payload=None, timeout=2.0):
+    data = None
+    headers = {"X-SnapForm-Token": connection.token}
+    if payload is not None:
+        data = json.dumps(payload).encode("utf-8")
+        headers["Content-Type"] = "application/json"
+    request = Request(
+        _host_url(connection, path),
+        data=data,
+        headers=headers,
+        method=method,
+    )
+    try:
+        with _local_opener.open(request, timeout=timeout) as response:
+            return json.loads(response.read().decode("utf-8"))
+    except HTTPError as error:
+        try:
+            detail = json.loads(error.read().decode("utf-8")).get("detail")
+        except (json.JSONDecodeError, UnicodeDecodeError):
+            detail = None
+        raise RuntimeError(detail or f"AI Host returned HTTP {error.code}")
+    except (OSError, URLError) as error:
+        raise RuntimeError(f"Unable to reach AI Host: {error}")
+
+
+def _read_host_discovery():
+    try:
+        payload = json.loads(host_discovery_path().read_text(encoding="utf-8"))
+        return HostConnection(
+            port=int(payload["port"]),
+            token=str(payload["token"]),
+            pid=int(payload["pid"]),
+            instance_id=str(payload["instance_id"]),
+        )
+    except (KeyError, TypeError, ValueError, OSError, json.JSONDecodeError):
+        return None
+
+
+def _check_host(connection, timeout=0.3):
+    health = _host_http_request(
+        connection,
+        "GET",
+        "/health",
+        timeout=timeout,
+    )
+    if health.get("host_schema") != HOST_SCHEMA:
+        raise RuntimeError("AI Host protocol is out of date")
+    if health.get("instance_id") != connection.instance_id:
+        raise RuntimeError("AI Host discovery is stale")
+    return health
+
+
+def _discover_live_host(timeout=0.3):
+    connection = _read_host_discovery()
+    if connection is None:
+        return None, None
+    try:
+        return connection, _check_host(connection, timeout=timeout)
+    except RuntimeError:
+        return None, None
+
+
+def start_host():
+    global _launched_host
+
+    if _launched_host is not None and _launched_host.process.poll() is None:
+        return
+    _close_launched_host()
+    port = _available_host_port()
+    token = secrets.token_urlsafe(32)
+    instance_id = uuid.uuid4().hex
+    log_file = host_log_path().open("a", encoding="utf-8")
+    command = [
+        str(runtime_python()),
+        "-u",
+        str(Path(__file__).with_name("ai_host.py")),
+        "--host",
+        HOST_ADDRESS,
+        "--port",
+        str(port),
+        "--token",
+        token,
+        "--instance-id",
+        instance_id,
+        "--discovery",
+        str(host_discovery_path()),
+        "--lock",
+        str(host_lock_path()),
+    ]
+    try:
+        process = subprocess.Popen(
+            command,
+            stdout=log_file,
+            stderr=subprocess.STDOUT,
+            cwd=storage_root(),
+            creationflags=creation_flags(),
+        )
+    except OSError:
+        log_file.close()
+        raise
+    _launched_host = LaunchedHost(process, log_file)
+
+
+def ensure_host(timeout=60.0):
+    global _host_auto_start, _host_connection, _host_snapshot
+
+    _host_auto_start = True
+    deadline = time.monotonic() + timeout
+    last_launch = 0.0
+    starting = False
+    while time.monotonic() < deadline:
+        connection, _health = _discover_live_host(timeout=0.5)
+        if connection is not None:
+            _host_connection = connection
+            _register_host_client(connection)
+            health = _check_host(connection, timeout=0.5)
+            _host_snapshot = _format_host_snapshot(connection, health)
+            _redraw_host_ui()
+            return connection
+        if not starting:
+            _host_snapshot = {"state": "STARTING", "message": "Starting"}
+            _redraw_host_ui()
+            starting = True
+        _close_finished_launched_host()
+        if (
+            _launched_host is None
+            and time.monotonic() - last_launch >= 1.0
+        ):
+            start_host()
+            last_launch = time.monotonic()
+        time.sleep(0.05)
+    raise RuntimeError(f"AI Host did not become ready; log: {host_log_path()}")
+
+
+def submit_host_job(request):
+    connection = ensure_host()
+    payload = dict(request.payload)
+    payload["client_id"] = _client_id
+    _host_http_request(
+        connection,
+        "POST",
+        "/jobs",
+        payload=payload,
+        timeout=5.0,
+    )
+    return connection
+
+
+def stop_host():
+    global _host_auto_start, _host_connection, _host_snapshot
+    global _registered_instance_id
+
+    connection = _host_connection
+    if connection is None:
+        connection, _health = _discover_live_host(timeout=0.5)
+    if connection is None:
+        _close_launched_host(terminate=True)
+        _host_auto_start = False
+        _host_snapshot = {"state": "STOPPED", "message": "Stopped"}
+        return
+    _register_host_client(connection)
+    _host_http_request(
+        connection,
+        "POST",
+        "/shutdown",
+        payload={"client_id": _client_id},
+        timeout=2.0,
+    )
+    _host_auto_start = False
+    deadline = time.monotonic() + 10.0
+    while time.monotonic() < deadline:
+        try:
+            _check_host(connection, timeout=0.2)
+        except RuntimeError:
+            break
+        time.sleep(0.05)
+    else:
+        raise RuntimeError("AI Host did not stop")
+    _host_connection = None
+    _registered_instance_id = None
+    _close_finished_launched_host()
+    _host_snapshot = {"state": "STOPPED", "message": "Stopped"}
+
+
+def restart_host():
+    global _host_auto_start, _host_connection, _registered_instance_id
+    global _host_snapshot
+
+    connection = ensure_host()
+    _host_http_request(
+        connection,
+        "POST",
+        "/shutdown",
+        payload={"client_id": _client_id},
+        timeout=2.0,
+    )
+    _host_auto_start = True
+    _host_connection = None
+    _registered_instance_id = None
+    _host_snapshot = {"state": "STARTING", "message": "Restarting"}
+
+
+def enable_host():
+    global _host_auto_start, _host_connection, _host_snapshot
+
+    _host_auto_start = True
+    connection, health = _discover_live_host(timeout=0.2)
+    if connection is not None:
+        _host_connection = connection
+        _register_host_client(connection)
+        _host_snapshot = _format_host_snapshot(connection, health)
+        _redraw_host_ui()
+        return
+    start_host()
+    _host_snapshot = {"state": "STARTING", "message": "Starting"}
+    _redraw_host_ui()
+
+
+def clear_host_models():
+    connection = ensure_host()
+    return _host_http_request(
+        connection,
+        "POST",
+        "/models/clear",
+        timeout=5.0,
+    )
+
+
+def cancel_host_job(job_id):
+    connection = ensure_host()
+    return _host_http_request(
+        connection,
+        "DELETE",
+        f"/jobs/{job_id}",
+        payload={"client_id": _client_id},
+        timeout=2.0,
+    )
+
+
+def detach_host_client():
+    global _host_connection, _registered_instance_id
+
+    connection = _host_connection
+    if connection is not None and _registered_instance_id is not None:
+        try:
+            _host_http_request(
+                connection,
+                "DELETE",
+                f"/clients/{_client_id}",
+                timeout=0.5,
+            )
+        except RuntimeError:
+            pass
+    _host_connection = None
+    _registered_instance_id = None
+    _close_launched_host(terminate=False)
+
+
+def schedule_host_start():
+    timers = getattr(bpy.app, "timers", None)
+    if timers is None:
+        return
+    if timers.is_registered(_service_host_timer):
+        return
+    timers.register(_service_host_timer, first_interval=1.0, persistent=True)
+
+
+def unschedule_host():
+    timers = getattr(bpy.app, "timers", None)
+    if timers is None:
+        return
+    if timers.is_registered(_service_host_timer):
+        timers.unregister(_service_host_timer)
+
+
+def host_status():
+    return dict(_host_snapshot)
+
+
+def host_busy():
+    return _host_snapshot.get("state") == "BUSY"
+
+
+def _service_host_timer():
+    global _host_connection, _host_snapshot, _registered_instance_id
+
+    if not runtime_ready():
+        _host_snapshot = {
+            "state": "UNAVAILABLE",
+            "message": "Runtime is not installed",
+        }
+        return 2.0
+    connection, health = _discover_live_host(timeout=0.2)
+    if connection is None:
+        _host_connection = None
+        _registered_instance_id = None
+        _close_finished_launched_host()
+        if not _host_auto_start:
+            if _host_snapshot.get("state") != "STOPPED":
+                _host_snapshot = {"state": "STOPPED", "message": "Stopped"}
+                _redraw_host_ui()
+            return 2.0
+        try:
+            start_host()
+        except OSError as error:
+            _host_snapshot = {"state": "ERROR", "message": str(error)}
+            return 2.0
+        _host_snapshot = {"state": "STARTING", "message": "Starting"}
+        _redraw_host_ui()
+        return 1.0
+    try:
+        _host_connection = connection
+        if _registered_instance_id == connection.instance_id:
+            try:
+                _host_http_request(
+                    connection,
+                    "POST",
+                    f"/clients/{_client_id}/heartbeat",
+                    timeout=0.5,
+                )
+            except RuntimeError:
+                _registered_instance_id = None
+        _register_host_client(connection)
+        health = _check_host(connection, timeout=0.5)
+        _host_snapshot = _format_host_snapshot(connection, health)
+    except RuntimeError as error:
+        _host_connection = None
+        _registered_instance_id = None
+        _host_snapshot = {"state": "ERROR", "message": str(error)}
+    _redraw_host_ui()
+    return 5.0
+
+
+def _register_host_client(connection):
+    global _registered_instance_id
+
+    if _registered_instance_id == connection.instance_id:
+        return
+    label = f"Blender {getattr(bpy.app, 'version_string', '')}".strip()
+    _host_http_request(
+        connection,
+        "POST",
+        "/clients",
+        payload={"client_id": _client_id, "label": label},
+        timeout=2.0,
+    )
+    _registered_instance_id = connection.instance_id
+
+
+def _format_host_snapshot(connection, health):
+    state = "BUSY" if health.get("busy") else "READY"
+    return {
+        **health,
+        "state": state,
+        "message": "Busy" if state == "BUSY" else "Ready",
+        "port": connection.port,
+    }
+
+
+def _redraw_host_ui():
+    try:
+        from .status_bar import redraw_ui
+
+        redraw_ui(bpy.context, force=True)
+    except (AttributeError, RuntimeError):
+        pass
+
+
+def _close_finished_launched_host():
+    if _launched_host is None or _launched_host.process.poll() is None:
+        return
+    _close_launched_host()
+
+
+def _close_launched_host(terminate=False):
+    global _launched_host
+
+    launched = _launched_host
+    _launched_host = None
+    if launched is None:
+        return
+    if terminate and launched.process.poll() is None:
+        launched.process.terminate()
+        try:
+            launched.process.wait(timeout=3.0)
+        except subprocess.TimeoutExpired:
+            launched.process.kill()
+            launched.process.wait(timeout=3.0)
+    launched.log_file.close()
+
+
 def download_model_command(model_key, status_path):
    model = get_downloadable_model(model_key)
    command = [
@@ -153,31 +604,29 @@ def download_model_command(model_key, status_path):
    return command


-def remove_background_command(
+def remove_background_request(
    settings,
    input_path,
    status_path,
    result_directory,
    model_key,
 ):
-    return [
-        str(runtime_python()),
-        str(Path(__file__).with_name("worker.py")),
-        "remove-background",
-        "--input",
-        str(input_path),
-        "--model-dir",
-        str(model_directory(model_key)),
-        "--output-dir",
-        str(result_directory),
-        "--status",
-        str(status_path),
-        "--device",
-        settings.device.lower(),
-    ]
+    return HostJobRequest(
+        {
+            "job_id": Path(status_path).stem,
+            "command": "remove-background",
+            "parameters": {
+                "input": str(input_path),
+                "model_dir": str(model_directory(model_key)),
+                "output_dir": str(result_directory),
+                "status": str(status_path),
+                "device": settings.device.lower(),
+            },
+        }
+    )


-def reconstruct_command(
+def reconstruct_request(
    settings,
    export_format,
    status_path,
@@ -201,53 +650,39 @@ def reconstruct_command(
    snapform_selected_ray_pose = (
        settings.use_ray_pose if use_ray_pose is None else use_ray_pose
    )
-    command = [
-        str(runtime_python()),
-        str(Path(__file__).with_name("worker.py")),
-        "reconstruct",
-        "--input",
-        str(
+    parameters = {
+        "input": str(
            Path(input_path)
            if input_path is not None
            else Path(bpy.path.abspath(settings.input_path))
        ),
-        "--model-dir",
-        str(model_directory(snapform_selected_key)),
-        "--output-dir",
-        str(result_directory),
-        "--export-format",
-        export_format,
-        "--status",
-        str(status_path),
-        "--process-res",
-        str(snapform_selected_process_res),
-        "--process-res-method",
-        snapform_selected_process_res_method,
-        "--conf-thresh-percentile",
-        str(settings.conf_thresh_percentile),
-        "--num-max-points",
-        str(settings.num_max_points),
-        "--show-cameras",
-        "1" if settings.show_cameras else "0",
-        "--gs-views-interval",
-        str(settings.gs_views_interval),
-        "--video-frames",
-        str(settings.video_frames),
-        "--device",
-        settings.device.lower(),
-        "--use-ray-pose",
-        "1" if snapform_selected_ray_pose else "0",
-        "--ref-view-strategy",
-        settings.ref_view_strategy,
-    ]
+        "model_dir": str(model_directory(snapform_selected_key)),
+        "output_dir": str(result_directory),
+        "export_format": export_format,
+        "status": str(status_path),
+        "process_res": snapform_selected_process_res,
+        "process_res_method": snapform_selected_process_res_method,
+        "conf_thresh_percentile": settings.conf_thresh_percentile,
+        "num_max_points": settings.num_max_points,
+        "show_cameras": 1 if settings.show_cameras else 0,
+        "gs_views_interval": settings.gs_views_interval,
+        "video_frames": settings.video_frames,
+        "device": settings.device.lower(),
+        "use_ray_pose": 1 if snapform_selected_ray_pose else 0,
+        "ref_view_strategy": settings.ref_view_strategy,
+        "remove_background_model_dir": "",
+    }
    if remove_background:
-        command.extend(
-            [
-                "--remove-background-model-dir",
-                str(model_directory(BEN2_MODEL.key)),
-            ]
+        parameters["remove_background_model_dir"] = str(
+            model_directory(BEN2_MODEL.key)
        )
-    return command
+    return HostJobRequest(
+        {
+            "job_id": Path(status_path).stem,
+            "command": "reconstruct",
+            "parameters": parameters,
+        }
+    )


 def creation_flags():