| 1 | """Local Echo 1.5 inference scheduler and media assembly API. |
| 2 | |
| 3 | The control plane schedules work from in-memory queues and mirrors job state to |
| 4 | SQLite for restart recovery. GPU workers admit work from live GPU/RAM telemetry, |
| 5 | stage conditioning and generation weights, execute the shared local pipeline, |
| 6 | index finished artifacts, and notify callbacks. |
| 7 | """ |
| 8 | |
| 9 | from __future__ import annotations |
| 10 | |
| 11 | import argparse |
| 12 | import asyncio |
| 13 | import base64 |
| 14 | import binascii |
| 15 | import json |
| 16 | import os |
| 17 | import shutil |
| 18 | import sqlite3 |
| 19 | import subprocess |
| 20 | import tempfile |
| 21 | import threading |
| 22 | import time |
| 23 | import traceback |
| 24 | import uuid |
| 25 | from collections import OrderedDict |
| 26 | from contextlib import asynccontextmanager |
| 27 | from dataclasses import dataclass |
| 28 | from datetime import UTC, datetime, timedelta |
| 29 | from hashlib import sha256 |
| 30 | from pathlib import Path |
| 31 | from typing import Any |
| 32 | from urllib.parse import quote, unquote, urlparse |
| 33 | |
| 34 | import httpx |
| 35 | import yaml |
| 36 | from fastapi import FastAPI, Header, HTTPException, Request |
| 37 | from fastapi.responses import FileResponse |
| 38 | from pydantic import ( |
| 39 | AliasChoices, |
| 40 | BaseModel, |
| 41 | Field, |
| 42 | ValidationError, |
| 43 | field_validator, |
| 44 | model_validator, |
| 45 | ) |
| 46 | |
| 47 | from r2v_schema import MAX_MEMORY_SLOTS |
| 48 | from .state import JobJournal |
| 49 | |
| 50 | REPO_ROOT = Path(__file__).resolve().parent.parent |
| 51 | MAX_IMAGE_BYTES = 20 * 1024 * 1024 |
| 52 | MAX_AUDIO_BYTES = 100 * 1024 * 1024 |
| 53 | _DATA_MIME_EXTENSIONS = { |
| 54 | "image/jpeg": ".jpg", |
| 55 | "image/png": ".png", |
| 56 | "image/webp": ".webp", |
| 57 | "image/gif": ".gif", |
| 58 | "image/bmp": ".bmp", |
| 59 | "audio/wav": ".wav", |
| 60 | "audio/x-wav": ".wav", |
| 61 | "audio/flac": ".flac", |
| 62 | "audio/mpeg": ".mp3", |
| 63 | "audio/mp4": ".m4a", |
| 64 | "audio/ogg": ".ogg", |
| 65 | } |
| 66 | |
| 67 | |
| 68 | def _now() -> str: |
| 69 | return datetime.now(UTC).isoformat() |
| 70 | |
| 71 | |
| 72 | def _json_loads(value: str | None, default: Any) -> Any: |
| 73 | if not value: |
| 74 | return default |
| 75 | try: |
| 76 | return json.loads(value) |
| 77 | except (TypeError, json.JSONDecodeError): |
| 78 | return default |
| 79 | |
| 80 | |
| 81 | def _canonical_json(value: Any) -> str: |
| 82 | return json.dumps(value, ensure_ascii=False, sort_keys=True, separators=(",", ":")) |
| 83 | |
| 84 | |
| 85 | def _file_sha256(path: Path) -> str: |
| 86 | digest = sha256() |
| 87 | with path.open("rb") as source: |
| 88 | for chunk in iter(lambda: source.read(1024 * 1024), b""): |
| 89 | digest.update(chunk) |
| 90 | return digest.hexdigest() |
| 91 | |
| 92 | |
| 93 | def _callback_retry_at(attempts_before_failure: int) -> str: |
| 94 | delay_seconds = min(2 ** max(attempts_before_failure, 0), 30) |
| 95 | return (datetime.now(UTC) + timedelta(seconds=delay_seconds)).isoformat() |
| 96 | |
| 97 | |
| 98 | def _repo_path(value: str) -> Path: |
| 99 | path = Path(value).expanduser() |
| 100 | if not path.is_absolute(): |
| 101 | path = REPO_ROOT / path |
| 102 | return path.resolve() |
| 103 | |
| 104 | |
| 105 | def _dit_residency_policy(value: str) -> str: |
| 106 | policy = value.strip().lower() |
| 107 | if policy not in {"auto", "resident", "swap"}: |
| 108 | raise ValueError("ECHO_DIT_RESIDENCY must be auto, resident, or swap") |
| 109 | return policy |
| 110 | |
| 111 | |
| 112 | def _decode_data_url(value: str, *, kind: str, max_bytes: int) -> tuple[bytes, str]: |
| 113 | """Decode a bounded base64 data URL and return its bytes and extension.""" |
| 114 | |
| 115 | header, separator, encoded = value.partition(",") |
| 116 | if not separator or not header.startswith("data:") or not header.endswith(";base64"): |
| 117 | raise ValueError(f"{kind} data URL must use base64 encoding") |
| 118 | mime_type = header[5:-7].strip().lower() |
| 119 | expected_prefix = "image/" if kind == "image" else "audio/" |
| 120 | if not mime_type.startswith(expected_prefix): |
| 121 | raise ValueError(f"{kind} data URL has incompatible MIME type: {mime_type}") |
| 122 | extension = _DATA_MIME_EXTENSIONS.get(mime_type) |
| 123 | if not extension: |
| 124 | raise ValueError(f"unsupported {kind} data URL MIME type: {mime_type}") |
| 125 | if len(encoded) > ((max_bytes + 2) // 3) * 4 + 4: |
| 126 | raise ValueError(f"{kind} data URL exceeds {max_bytes} decoded bytes") |
| 127 | try: |
| 128 | payload = base64.b64decode(encoded, validate=True) |
| 129 | except (ValueError, binascii.Error) as error: |
| 130 | raise ValueError(f"invalid base64 in {kind} data URL") from error |
| 131 | if not payload or len(payload) > max_bytes: |
| 132 | raise ValueError(f"{kind} data URL must contain 1..{max_bytes} bytes") |
| 133 | return payload, extension |
| 134 | |
| 135 | |
| 136 | def _config_section(config: dict[str, Any], name: str) -> dict[str, Any]: |
| 137 | value = config.get(name, {}) |
| 138 | if value is None: |
| 139 | return {} |
| 140 | if not isinstance(value, dict): |
| 141 | raise ValueError(f"server config section {name!r} must be a mapping") |
| 142 | return value |
| 143 | |
| 144 | |
| 145 | def load_server_config(path: str | Path) -> dict[str, Any]: |
| 146 | config_path = _repo_path(str(path)) |
| 147 | with config_path.open("r", encoding="utf-8") as source: |
| 148 | config = yaml.safe_load(source) or {} |
| 149 | if not isinstance(config, dict): |
| 150 | raise ValueError("server config root must be a mapping") |
| 151 | for section in ("inference", "runtime", "server"): |
| 152 | _config_section(config, section) |
| 153 | return config |
| 154 | |
| 155 | |
| 156 | def _environment_value(value: Any) -> str: |
| 157 | if isinstance(value, bool): |
| 158 | return "1" if value else "0" |
| 159 | if isinstance(value, list): |
| 160 | return ",".join(str(item) for item in value) |
| 161 | return str(value) |
| 162 | |
| 163 | |
| 164 | def apply_server_config(path: str | Path) -> dict[str, Any]: |
| 165 | """Load deployment settings while preserving explicit environment overrides.""" |
| 166 | |
| 167 | config = load_server_config(path) |
| 168 | inference = _config_section(config, "inference") |
| 169 | runtime = _config_section(config, "runtime") |
| 170 | server = _config_section(config, "server") |
| 171 | mappings = ( |
| 172 | ("ECHO_INFERENCE_CONFIG", inference, "config"), |
| 173 | ("ECHO_CHECKPOINT", inference, "checkpoint"), |
| 174 | ("ECHO_CONDITIONING_CACHE_DIR", inference, "conditioning_cache_dir"), |
| 175 | ("ECHO_GPU_IDS", runtime, "gpu_ids"), |
| 176 | ("ECHO_DISABLE_INFERENCE_WORKERS", runtime, "disable_inference_workers"), |
| 177 | ("ECHO_DIT_RESIDENCY", runtime, "dit_residency"), |
| 178 | ("ECHO_GPU_HEADROOM_FRACTION", runtime, "gpu_headroom_fraction"), |
| 179 | ("ECHO_RAM_HEADROOM_FRACTION", runtime, "ram_headroom_fraction"), |
| 180 | ("ECHO_MODEL_IDLE_SECONDS", runtime, "model_idle_seconds"), |
| 181 | ("ECHO_ARTIFACT_DB_PATH", server, "artifact_db_path"), |
| 182 | ("ECHO_R2V_TIMEOUT_SECONDS", server, "request_timeout_seconds"), |
| 183 | ("ECHO_R2V_POLL_SECONDS", server, "poll_seconds"), |
| 184 | ("ECHO_PUBLIC_BASE_URL", server, "public_base_url"), |
| 185 | ("ECHO_R2V_QUEUE_CAPACITY", server, "queue_capacity"), |
| 186 | ("ECHO_CALLBACK_MAX_ATTEMPTS", server, "callback_max_attempts"), |
| 187 | ("ECHO_MEDIA_ROOT", server, "media_root"), |
| 188 | ("ECHO_FFMPEG_PATH", server, "ffmpeg_path"), |
| 189 | ("ECHO_MERGE_TIMEOUT_SECONDS", server, "merge_timeout_seconds"), |
| 190 | ("ECHO_MERGE_MAX_INPUT_BYTES", server, "merge_max_input_bytes"), |
| 191 | ) |
| 192 | for environment_name, section, key in mappings: |
| 193 | value = section.get(key) |
| 194 | if value is not None: |
| 195 | os.environ.setdefault(environment_name, _environment_value(value)) |
| 196 | return config |
| 197 | |
| 198 | |
| 199 | class MemorySlot(BaseModel): |
| 200 | """One ordered memory slot accepted by the production R2V endpoint.""" |
| 201 | |
| 202 | shot_id: str | None = None |
| 203 | image_url: str | None = None |
| 204 | image_mode: str | None = None |
| 205 | audio_url: str | None = None |
| 206 | audio_mode: str | None = None |
| 207 | metadata: dict[str, Any] | None = None |
| 208 | |
| 209 | @model_validator(mode="after") |
| 210 | def validate_source(self) -> "MemorySlot": |
| 211 | self.shot_id = (self.shot_id or "").strip() or None |
| 212 | self.image_url = (self.image_url or "").strip() or None |
| 213 | self.image_mode = (self.image_mode or "").strip().lower() or None |
| 214 | self.audio_url = (self.audio_url or "").strip() or None |
| 215 | audio_mode = (self.audio_mode or "").strip().lower() |
| 216 | |
| 217 | if self.image_mode: |
| 218 | raise ValueError("image_mode is not supported") |
| 219 | if bool(self.shot_id) == bool(self.image_url): |
| 220 | raise ValueError("exactly one of shot_id or image_url is required") |
| 221 | if audio_mode not in {"", "empty"}: |
| 222 | raise ValueError("audio_mode currently supports only 'empty'") |
| 223 | self.audio_mode = audio_mode or None |
| 224 | if self.shot_id and (self.audio_url or self.audio_mode): |
| 225 | raise ValueError("shot_id cannot be combined with audio_url or audio_mode") |
| 226 | if self.audio_url and self.audio_mode == "empty": |
| 227 | raise ValueError("audio_url conflicts with audio_mode='empty'") |
| 228 | if self.image_url and not self.audio_url and self.audio_mode is None: |
| 229 | self.audio_mode = "empty" |
| 230 | return self |
| 231 | |
| 232 | |
| 233 | class R2VGenerateRequest(BaseModel): |
| 234 | work_id: str = Field(min_length=1) |
| 235 | shot_id: str = Field(min_length=1) |
| 236 | job_id: str | None = None |
| 237 | prompt: str = Field(min_length=1) |
| 238 | memory_slots: list[MemorySlot] = Field(max_length=MAX_MEMORY_SLOTS) |
| 239 | condition_img: str | None = None |
| 240 | callback_url: str | None = None |
| 241 | callback_context: dict[str, Any] = Field(default_factory=dict) |
| 242 | num_frames: int | None = Field(default=None, gt=0) |
| 243 | duration_sec: float | None = Field(default=None, gt=0) |
| 244 | width: int | None = Field(default=None, gt=0) |
| 245 | height: int | None = Field(default=None, gt=0) |
| 246 | seed: int | None = None |
| 247 | |
| 248 | @field_validator("work_id", "shot_id", "prompt") |
| 249 | @classmethod |
| 250 | def strip_required_text(cls, value: str) -> str: |
| 251 | value = value.strip() |
| 252 | if not value: |
| 253 | raise ValueError("must not be blank") |
| 254 | return value |
| 255 | |
| 256 | |
| 257 | class MergeShotInput(BaseModel): |
| 258 | """One ordered video input for a merge job.""" |
| 259 | |
| 260 | version_id: str | None = None |
| 261 | video_url: str | None = Field( |
| 262 | default=None, |
| 263 | validation_alias=AliasChoices("video_url", "artifact_url", "result_url"), |
| 264 | ) |
| 265 | |
| 266 | @model_validator(mode="after") |
| 267 | def validate_source(self) -> "MergeShotInput": |
| 268 | self.version_id = (self.version_id or "").strip() or None |
| 269 | self.video_url = (self.video_url or "").strip() or None |
| 270 | if not self.version_id and not self.video_url: |
| 271 | raise ValueError("version_id or video_url is required") |
| 272 | return self |
| 273 | |
| 274 | |
| 275 | class MergeRequest(BaseModel): |
| 276 | """Ordered shot list assembled into one server-owned MP4 artifact.""" |
| 277 | |
| 278 | work_id: str = Field(min_length=1) |
| 279 | job_id: str | None = None |
| 280 | shots: list[MergeShotInput] = Field(min_length=1) |
| 281 | callback_url: str | None = None |
| 282 | callback_context: dict[str, Any] = Field(default_factory=dict) |
| 283 | |
| 284 | @field_validator("work_id") |
| 285 | @classmethod |
| 286 | def strip_work_id(cls, value: str) -> str: |
| 287 | value = value.strip() |
| 288 | if not value: |
| 289 | raise ValueError("must not be blank") |
| 290 | return value |
| 291 | |
| 292 | |
| 293 | @dataclass(frozen=True) |
| 294 | class Settings: |
| 295 | artifact_db_path: Path |
| 296 | request_timeout: float |
| 297 | poll_interval: float |
| 298 | public_base_url: str |
| 299 | queue_capacity: int |
| 300 | callback_attempts: int |
| 301 | media_root: Path |
| 302 | ffmpeg_binary: str |
| 303 | merge_timeout: float |
| 304 | merge_max_input_bytes: int |
| 305 | inference_config: Path |
| 306 | checkpoint: str | None |
| 307 | conditioning_cache_dir: Path |
| 308 | gpu_ids: str |
| 309 | disable_inference_workers: bool |
| 310 | dit_residency: str |
| 311 | gpu_headroom_fraction: float |
| 312 | ram_headroom_fraction: float |
| 313 | model_idle_seconds: float |
| 314 | |
| 315 | @classmethod |
| 316 | def from_env(cls) -> "Settings": |
| 317 | return cls( |
| 318 | artifact_db_path=_repo_path( |
| 319 | os.environ.get("ECHO_ARTIFACT_DB_PATH", "data/artifacts.sqlite3") |
| 320 | ), |
| 321 | request_timeout=float(os.environ.get("ECHO_R2V_TIMEOUT_SECONDS", "30")), |
| 322 | poll_interval=max(float(os.environ.get("ECHO_R2V_POLL_SECONDS", "5")), 0.1), |
| 323 | public_base_url=os.environ.get( |
| 324 | "ECHO_PUBLIC_BASE_URL", "http://127.0.0.1:8221" |
| 325 | ).strip().rstrip("/"), |
| 326 | queue_capacity=max(int(os.environ.get("ECHO_R2V_QUEUE_CAPACITY", "1000")), 1), |
| 327 | callback_attempts=max(int(os.environ.get("ECHO_CALLBACK_MAX_ATTEMPTS", "0")), 0), |
| 328 | media_root=_repo_path(os.environ.get("ECHO_MEDIA_ROOT", "data/media")), |
| 329 | ffmpeg_binary=os.environ.get("ECHO_FFMPEG_PATH", "ffmpeg").strip() or "ffmpeg", |
| 330 | merge_timeout=max(float(os.environ.get("ECHO_MERGE_TIMEOUT_SECONDS", "3600")), 1), |
| 331 | merge_max_input_bytes=max( |
| 332 | int(os.environ.get("ECHO_MERGE_MAX_INPUT_BYTES", str(2 * 1024**3))), |
| 333 | 1, |
| 334 | ), |
| 335 | inference_config=_repo_path( |
| 336 | os.environ.get("ECHO_INFERENCE_CONFIG", "configs/inference.bf16.yaml") |
| 337 | ), |
| 338 | checkpoint=os.environ.get("ECHO_CHECKPOINT", "").strip() or None, |
| 339 | conditioning_cache_dir=_repo_path( |
| 340 | os.environ.get("ECHO_CONDITIONING_CACHE_DIR", "data/conditioning_cache") |
| 341 | ), |
| 342 | gpu_ids=os.environ.get("ECHO_GPU_IDS", "0").strip(), |
| 343 | disable_inference_workers=os.environ.get( |
| 344 | "ECHO_DISABLE_INFERENCE_WORKERS", "0" |
| 345 | ).strip().lower() |
| 346 | in {"1", "true", "yes"}, |
| 347 | dit_residency=_dit_residency_policy( |
| 348 | os.environ.get("ECHO_DIT_RESIDENCY", "auto") |
| 349 | ), |
| 350 | gpu_headroom_fraction=min( |
| 351 | max(float(os.environ.get("ECHO_GPU_HEADROOM_FRACTION", "0.05")), 0), |
| 352 | 0.5, |
| 353 | ), |
| 354 | ram_headroom_fraction=min( |
| 355 | max(float(os.environ.get("ECHO_RAM_HEADROOM_FRACTION", "0.10")), 0), |
| 356 | 0.5, |
| 357 | ), |
| 358 | model_idle_seconds=max( |
| 359 | float(os.environ.get("ECHO_MODEL_IDLE_SECONDS", "0")), 0 |
| 360 | ), |
| 361 | ) |
| 362 | |
| 363 | |
| 364 | class JobStore: |
| 365 | """In-memory FIFO mirrored to a durable restart journal.""" |
| 366 | |
| 367 | def __init__( |
| 368 | self, |
| 369 | capacity: int, |
| 370 | *, |
| 371 | journal: JobJournal | None = None, |
| 372 | journal_kind: str = "r2v", |
| 373 | ) -> None: |
| 374 | self.capacity = capacity |
| 375 | self.journal = journal |
| 376 | self.journal_kind = journal_kind |
| 377 | self._jobs: OrderedDict[str, dict[str, Any]] = OrderedDict() |
| 378 | self._lock = threading.RLock() |
| 379 | |
| 380 | def initialize(self) -> None: |
| 381 | if self.journal is None: |
| 382 | return |
| 383 | self.journal.initialize() |
| 384 | restored = self.journal.load(self.journal_kind) |
| 385 | with self._lock: |
| 386 | self._jobs.clear() |
| 387 | for job in restored: |
| 388 | if job.get("status") == "running": |
| 389 | job.update( |
| 390 | status="queued", |
| 391 | stage="recovered", |
| 392 | gpu_id=None, |
| 393 | resource_json=None, |
| 394 | updated_at=_now(), |
| 395 | ) |
| 396 | self.journal.save(self.journal_kind, job) |
| 397 | self._jobs[str(job["version_id"])] = job |
| 398 | self._prune_terminal() |
| 399 | |
| 400 | def _persist(self, job: dict[str, Any]) -> None: |
| 401 | if self.journal is not None: |
| 402 | self.journal.save(self.journal_kind, job) |
| 403 | |
| 404 | def _prune_terminal(self) -> None: |
| 405 | while len(self._jobs) > self.capacity * 2: |
| 406 | terminal_id = next( |
| 407 | ( |
| 408 | version_id |
| 409 | for version_id, job in self._jobs.items() |
| 410 | if job["status"] in {"succeeded", "failed"} |
| 411 | and ( |
| 412 | not job.get("callback_url") |
| 413 | or job.get("callback_status") == "delivered" |
| 414 | ) |
| 415 | ), |
| 416 | None, |
| 417 | ) |
| 418 | if terminal_id is None: |
| 419 | return |
| 420 | self._jobs.pop(terminal_id, None) |
| 421 | |
| 422 | def enqueue( |
| 423 | self, |
| 424 | version_id: str, |
| 425 | request: R2VGenerateRequest, |
| 426 | payload: dict[str, Any], |
| 427 | callback_url: str | None, |
| 428 | ) -> str: |
| 429 | now = _now() |
| 430 | request_json = _canonical_json(payload) |
| 431 | with self._lock: |
| 432 | if request.job_id: |
| 433 | existing = next( |
| 434 | ( |
| 435 | job |
| 436 | for job in self._jobs.values() |
| 437 | if job.get("agent_job_id") == request.job_id |
| 438 | ), |
| 439 | None, |
| 440 | ) |
| 441 | if existing is None and self.journal is not None: |
| 442 | existing = self.journal.get_by_agent_job_id( |
| 443 | self.journal_kind, request.job_id |
| 444 | ) |
| 445 | if existing is not None: |
| 446 | if existing["request_json"] != request_json: |
| 447 | raise ValueError( |
| 448 | "job_id is already associated with a different R2V request" |
| 449 | ) |
| 450 | existing.update( |
| 451 | callback_url=callback_url or existing.get("callback_url"), |
| 452 | callback_context_json=json.dumps( |
| 453 | request.callback_context, ensure_ascii=False |
| 454 | ), |
| 455 | updated_at=now, |
| 456 | ) |
| 457 | self._persist(existing) |
| 458 | return str(existing["version_id"]) |
| 459 | pending = sum( |
| 460 | job["status"] in {"queued", "running"} for job in self._jobs.values() |
| 461 | ) |
| 462 | if pending >= self.capacity: |
| 463 | raise OverflowError("R2V submission queue is full") |
| 464 | self._jobs[version_id] = { |
| 465 | "version_id": version_id, |
| 466 | "work_id": request.work_id, |
| 467 | "shot_id": request.shot_id, |
| 468 | "agent_job_id": request.job_id, |
| 469 | "request_json": request_json, |
| 470 | "callback_url": callback_url, |
| 471 | "callback_context_json": json.dumps( |
| 472 | request.callback_context, ensure_ascii=False |
| 473 | ), |
| 474 | "callback_status": None, |
| 475 | "callback_response": None, |
| 476 | "callback_error": None, |
| 477 | "callback_attempts": 0, |
| 478 | "callback_next_at": None, |
| 479 | "status": "queued", |
| 480 | "stage": "queued", |
| 481 | "result_json": None, |
| 482 | "gpu_id": None, |
| 483 | "resource_json": None, |
| 484 | "error": None, |
| 485 | "attempts": 0, |
| 486 | "created_at": now, |
| 487 | "updated_at": now, |
| 488 | "started_at": None, |
| 489 | "completed_at": None, |
| 490 | } |
| 491 | self._persist(self._jobs[version_id]) |
| 492 | self._prune_terminal() |
| 493 | return version_id |
| 494 | |
| 495 | def claim_batch(self, limit: int) -> list[dict[str, Any]]: |
| 496 | with self._lock: |
| 497 | rows = [ |
| 498 | job for job in self._jobs.values() if job["status"] == "queued" |
| 499 | ][: max(int(limit), 1)] |
| 500 | if not rows: |
| 501 | return [] |
| 502 | now = _now() |
| 503 | for job in rows: |
| 504 | job.update( |
| 505 | status="running", |
| 506 | stage="claimed", |
| 507 | attempts=int(job["attempts"]) + 1, |
| 508 | started_at=job["started_at"] or now, |
| 509 | updated_at=now, |
| 510 | error=None, |
| 511 | ) |
| 512 | self._persist(job) |
| 513 | return [dict(job) for job in rows] |
| 514 | |
| 515 | def claim_next(self) -> dict[str, Any] | None: |
| 516 | jobs = self.claim_batch(1) |
| 517 | return jobs[0] if jobs else None |
| 518 | |
| 519 | def update_stage( |
| 520 | self, |
| 521 | version_id: str, |
| 522 | stage: str, |
| 523 | *, |
| 524 | gpu_id: int | None = None, |
| 525 | resources: dict[str, Any] | None = None, |
| 526 | ) -> None: |
| 527 | with self._lock: |
| 528 | job = self._jobs[version_id] |
| 529 | job.update( |
| 530 | status="running", |
| 531 | stage=stage, |
| 532 | gpu_id=gpu_id, |
| 533 | resource_json=json.dumps(resources, ensure_ascii=False) |
| 534 | if resources |
| 535 | else None, |
| 536 | updated_at=_now(), |
| 537 | ) |
| 538 | self._persist(job) |
| 539 | |
| 540 | def complete(self, version_id: str, result: dict[str, Any]) -> None: |
| 541 | now = _now() |
| 542 | with self._lock: |
| 543 | self._jobs[version_id].update( |
| 544 | status="succeeded", |
| 545 | stage="succeeded", |
| 546 | result_json=json.dumps(result, ensure_ascii=False), |
| 547 | updated_at=now, |
| 548 | completed_at=now, |
| 549 | error=None, |
| 550 | ) |
| 551 | self._persist(self._jobs[version_id]) |
| 552 | self._prune_terminal() |
| 553 | |
| 554 | def fail(self, version_id: str, error: str) -> None: |
| 555 | now = _now() |
| 556 | with self._lock: |
| 557 | self._jobs[version_id].update( |
| 558 | status="failed", |
| 559 | stage="failed", |
| 560 | error=error[:4000], |
| 561 | updated_at=now, |
| 562 | completed_at=now, |
| 563 | ) |
| 564 | self._persist(self._jobs[version_id]) |
| 565 | self._prune_terminal() |
| 566 | |
| 567 | def callback_candidates(self, max_attempts: int) -> list[dict[str, Any]]: |
| 568 | now = _now() |
| 569 | with self._lock: |
| 570 | return [ |
| 571 | dict(job) |
| 572 | for job in self._jobs.values() |
| 573 | if job["status"] in {"succeeded", "failed"} |
| 574 | and job["callback_url"] |
| 575 | and job["callback_status"] != "delivered" |
| 576 | and (max_attempts == 0 or job["callback_attempts"] < max_attempts) |
| 577 | and (job["callback_next_at"] is None or job["callback_next_at"] <= now) |
| 578 | ] |
| 579 | |
| 580 | def record_callback( |
| 581 | self, |
| 582 | version_id: str, |
| 583 | *, |
| 584 | delivered: bool, |
| 585 | response: str | None = None, |
| 586 | error: str | None = None, |
| 587 | ) -> None: |
| 588 | with self._lock: |
| 589 | job = self._jobs[version_id] |
| 590 | attempts = int(job["callback_attempts"]) |
| 591 | job.update( |
| 592 | callback_status="delivered" if delivered else "error", |
| 593 | callback_response=response[:4000] if response else None, |
| 594 | callback_error=error[:1000] if error else None, |
| 595 | callback_attempts=attempts + 1, |
| 596 | callback_next_at=None |
| 597 | if delivered |
| 598 | else _callback_retry_at(attempts), |
| 599 | updated_at=_now(), |
| 600 | ) |
| 601 | self._persist(job) |
| 602 | |
| 603 | def get(self, version_id: str) -> dict[str, Any] | None: |
| 604 | with self._lock: |
| 605 | job = self._jobs.get(version_id) |
| 606 | if job is not None: |
| 607 | return dict(job) |
| 608 | if self.journal is not None: |
| 609 | return self.journal.get(self.journal_kind, version_id) |
| 610 | return None |
| 611 | |
| 612 | def queue_position(self, version_id: str) -> int | None: |
| 613 | job = self.get(version_id) |
| 614 | if job is None or job["status"] != "queued": |
| 615 | return None |
| 616 | with self._lock: |
| 617 | queued = [ |
| 618 | candidate["version_id"] |
| 619 | for candidate in self._jobs.values() |
| 620 | if candidate["status"] == "queued" |
| 621 | ] |
| 622 | return queued.index(version_id) + 1 |
| 623 | |
| 624 | def counts(self) -> dict[str, int]: |
| 625 | result = {name: 0 for name in ("queued", "running", "succeeded", "failed")} |
| 626 | with self._lock: |
| 627 | for job in self._jobs.values(): |
| 628 | result[str(job["status"])] += 1 |
| 629 | return result |
| 630 | |
| 631 | |
| 632 | class ArtifactStore: |
| 633 | """Small metadata index for generated and merged video artifacts.""" |
| 634 | |
| 635 | def __init__(self, path: Path) -> None: |
| 636 | self.path = path.expanduser().resolve() |
| 637 | |
| 638 | def connect(self) -> sqlite3.Connection: |
| 639 | connection = sqlite3.connect(self.path, timeout=30) |
| 640 | connection.row_factory = sqlite3.Row |
| 641 | connection.execute("PRAGMA busy_timeout = 30000") |
| 642 | return connection |
| 643 | |
| 644 | def initialize(self) -> None: |
| 645 | self.path.parent.mkdir(parents=True, exist_ok=True) |
| 646 | with self.connect() as connection: |
| 647 | connection.execute( |
| 648 | """ |
| 649 | CREATE TABLE IF NOT EXISTS artifacts ( |
| 650 | artifact_id TEXT PRIMARY KEY, |
| 651 | version_id TEXT NOT NULL, |
| 652 | work_id TEXT NOT NULL, |
| 653 | shot_id TEXT, |
| 654 | kind TEXT NOT NULL, |
| 655 | role TEXT NOT NULL, |
| 656 | url TEXT NOT NULL, |
| 657 | local_path TEXT, |
| 658 | size_bytes INTEGER, |
| 659 | sha256 TEXT, |
| 660 | metadata_json TEXT NOT NULL, |
| 661 | created_at TEXT NOT NULL, |
| 662 | updated_at TEXT NOT NULL, |
| 663 | UNIQUE(version_id, role) |
| 664 | ) |
| 665 | """ |
| 666 | ) |
| 667 | columns = { |
| 668 | str(row["name"]) |
| 669 | for row in connection.execute("PRAGMA table_info(artifacts)").fetchall() |
| 670 | } |
| 671 | if "shot_id" not in columns: |
| 672 | connection.execute("ALTER TABLE artifacts ADD COLUMN shot_id TEXT") |
| 673 | connection.execute( |
| 674 | "CREATE INDEX IF NOT EXISTS idx_artifacts_work_id " |
| 675 | "ON artifacts(work_id, created_at)" |
| 676 | ) |
| 677 | connection.execute( |
| 678 | "CREATE INDEX IF NOT EXISTS idx_artifacts_shot " |
| 679 | "ON artifacts(work_id, shot_id, created_at)" |
| 680 | ) |
| 681 | for row in connection.execute( |
| 682 | "SELECT artifact_id, metadata_json FROM artifacts WHERE shot_id IS NULL" |
| 683 | ).fetchall(): |
| 684 | metadata = _json_loads(row["metadata_json"], {}) |
| 685 | request = metadata.get("request") if isinstance(metadata, dict) else None |
| 686 | shot_id = ( |
| 687 | request.get("shot_id") |
| 688 | if isinstance(request, dict) |
| 689 | else metadata.get("shot_id") if isinstance(metadata, dict) else None |
| 690 | ) |
| 691 | if shot_id: |
| 692 | connection.execute( |
| 693 | "UPDATE artifacts SET shot_id = ? WHERE artifact_id = ?", |
| 694 | (str(shot_id), row["artifact_id"]), |
| 695 | ) |
| 696 | |
| 697 | def upsert( |
| 698 | self, |
| 699 | *, |
| 700 | version_id: str, |
| 701 | work_id: str, |
| 702 | shot_id: str | None = None, |
| 703 | kind: str, |
| 704 | role: str, |
| 705 | url: str, |
| 706 | local_path: str | None = None, |
| 707 | size_bytes: int | None = None, |
| 708 | sha256: str | None = None, |
| 709 | metadata: dict[str, Any] | None = None, |
| 710 | ) -> dict[str, Any]: |
| 711 | now = _now() |
| 712 | artifact_id = str(uuid.uuid5(uuid.NAMESPACE_URL, f"echo:{version_id}:{role}")) |
| 713 | with self.connect() as connection: |
| 714 | connection.execute( |
| 715 | """ |
| 716 | INSERT INTO artifacts ( |
| 717 | artifact_id, version_id, work_id, shot_id, kind, role, url, local_path, |
| 718 | size_bytes, sha256, metadata_json, created_at, updated_at |
| 719 | ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) |
| 720 | ON CONFLICT(version_id, role) DO UPDATE SET |
| 721 | work_id = excluded.work_id, |
| 722 | shot_id = excluded.shot_id, |
| 723 | kind = excluded.kind, |
| 724 | url = excluded.url, |
| 725 | local_path = excluded.local_path, |
| 726 | size_bytes = excluded.size_bytes, |
| 727 | sha256 = excluded.sha256, |
| 728 | metadata_json = excluded.metadata_json, |
| 729 | updated_at = excluded.updated_at |
| 730 | """, |
| 731 | ( |
| 732 | artifact_id, |
| 733 | version_id, |
| 734 | work_id, |
| 735 | shot_id, |
| 736 | kind, |
| 737 | role, |
| 738 | url, |
| 739 | local_path, |
| 740 | size_bytes, |
| 741 | sha256, |
| 742 | json.dumps(metadata or {}, ensure_ascii=False), |
| 743 | now, |
| 744 | now, |
| 745 | ), |
| 746 | ) |
| 747 | return self.get(artifact_id) or {} |
| 748 | |
| 749 | def get(self, artifact_id: str) -> dict[str, Any] | None: |
| 750 | with self.connect() as connection: |
| 751 | row = connection.execute( |
| 752 | "SELECT * FROM artifacts WHERE artifact_id = ?", (artifact_id,) |
| 753 | ).fetchone() |
| 754 | return self._public(dict(row)) if row else None |
| 755 | |
| 756 | def for_version(self, version_id: str) -> list[dict[str, Any]]: |
| 757 | with self.connect() as connection: |
| 758 | rows = connection.execute( |
| 759 | "SELECT * FROM artifacts WHERE version_id = ? ORDER BY created_at, role", |
| 760 | (version_id,), |
| 761 | ).fetchall() |
| 762 | return [self._public(dict(row)) for row in rows] |
| 763 | |
| 764 | def latest_local_for_shot(self, work_id: str, shot_id: str) -> dict[str, Any] | None: |
| 765 | """Return the newest generated primary artifact for one logical shot.""" |
| 766 | |
| 767 | with self.connect() as connection: |
| 768 | row = connection.execute( |
| 769 | """ |
| 770 | SELECT * FROM artifacts |
| 771 | WHERE work_id = ? AND shot_id = ? AND kind = 'r2v' AND role = 'primary' |
| 772 | ORDER BY created_at DESC LIMIT 1 |
| 773 | """, |
| 774 | (work_id, shot_id), |
| 775 | ).fetchone() |
| 776 | return dict(row) if row else None |
| 777 | |
| 778 | @staticmethod |
| 779 | def _public(row: dict[str, Any]) -> dict[str, Any]: |
| 780 | row["metadata"] = _json_loads(row.pop("metadata_json", None), {}) |
| 781 | # Filesystem layout is an implementation detail and is never returned. |
| 782 | row.pop("local_path", None) |
| 783 | return row |
| 784 | |
| 785 | |
| 786 | class RequestAssetResolver: |
| 787 | """Materialize inline assets and prior-shot references for local inference.""" |
| 788 | |
| 789 | def __init__(self, settings: Settings, artifacts: ArtifactStore) -> None: |
| 790 | self.settings = settings |
| 791 | self.artifacts = artifacts |
| 792 | self.root = settings.media_root / "request_assets" |
| 793 | self._lock = threading.Lock() |
| 794 | |
| 795 | def _write_inline(self, value: str, *, kind: str) -> str: |
| 796 | limit = MAX_IMAGE_BYTES if kind == "image" else MAX_AUDIO_BYTES |
| 797 | payload, extension = _decode_data_url(value, kind=kind, max_bytes=limit) |
| 798 | digest = sha256(payload).hexdigest() |
| 799 | target = self.root / "inline" / f"{digest}{extension}" |
| 800 | if not target.is_file(): |
| 801 | target.parent.mkdir(parents=True, exist_ok=True) |
| 802 | temporary = target.with_name(f".{target.name}.{uuid.uuid4().hex}.tmp") |
| 803 | temporary.write_bytes(payload) |
| 804 | try: |
| 805 | temporary.replace(target) |
| 806 | finally: |
| 807 | temporary.unlink(missing_ok=True) |
| 808 | return str(target.resolve()) |
| 809 | |
| 810 | def _extract_reference_assets( |
| 811 | self, |
| 812 | *, |
| 813 | work_id: str, |
| 814 | shot_id: str, |
| 815 | ) -> tuple[str, str]: |
| 816 | artifact = self.artifacts.latest_local_for_shot(work_id, shot_id) |
| 817 | if artifact is None: |
| 818 | raise LookupError( |
| 819 | f"memory shot is not available: work_id={work_id} shot_id={shot_id}" |
| 820 | ) |
| 821 | source = Path(str(artifact.get("local_path") or "")).expanduser().resolve() |
| 822 | if not source.is_file(): |
| 823 | raise FileNotFoundError( |
| 824 | f"memory shot artifact is missing: work_id={work_id} shot_id={shot_id}" |
| 825 | ) |
| 826 | binary = shutil.which(self.settings.ffmpeg_binary) |
| 827 | if binary is None: |
| 828 | raise RuntimeError( |
| 829 | f"FFmpeg executable was not found: {self.settings.ffmpeg_binary}" |
| 830 | ) |
| 831 | |
| 832 | version_id = str(artifact["version_id"]) |
| 833 | output_dir = self.root / "memory" / version_id |
| 834 | frame_path = output_dir / "representative.png" |
| 835 | audio_path = output_dir / "audio.wav" |
| 836 | with self._lock: |
| 837 | output_dir.mkdir(parents=True, exist_ok=True) |
| 838 | if not frame_path.is_file(): |
| 839 | temporary = output_dir / f".representative.{uuid.uuid4().hex}.png" |
| 840 | try: |
| 841 | subprocess.run( |
| 842 | [ |
| 843 | binary, |
| 844 | "-hide_banner", |
| 845 | "-loglevel", |
| 846 | "error", |
| 847 | "-y", |
| 848 | "-i", |
| 849 | str(source), |
| 850 | "-vf", |
| 851 | "thumbnail", |
| 852 | "-frames:v", |
| 853 | "1", |
| 854 | str(temporary), |
| 855 | ], |
| 856 | check=True, |
| 857 | capture_output=True, |
| 858 | text=True, |
| 859 | timeout=self.settings.merge_timeout, |
| 860 | ) |
| 861 | temporary.replace(frame_path) |
| 862 | finally: |
| 863 | temporary.unlink(missing_ok=True) |
| 864 | if not audio_path.is_file(): |
| 865 | temporary = output_dir / f".audio.{uuid.uuid4().hex}.wav" |
| 866 | try: |
| 867 | subprocess.run( |
| 868 | [ |
| 869 | binary, |
| 870 | "-hide_banner", |
| 871 | "-loglevel", |
| 872 | "error", |
| 873 | "-y", |
| 874 | "-i", |
| 875 | str(source), |
| 876 | "-vn", |
| 877 | "-ac", |
| 878 | "2", |
| 879 | "-ar", |
| 880 | "48000", |
| 881 | "-c:a", |
| 882 | "pcm_s16le", |
| 883 | str(temporary), |
| 884 | ], |
| 885 | check=True, |
| 886 | capture_output=True, |
| 887 | text=True, |
| 888 | timeout=self.settings.merge_timeout, |
| 889 | ) |
| 890 | temporary.replace(audio_path) |
| 891 | finally: |
| 892 | temporary.unlink(missing_ok=True) |
| 893 | return str(frame_path.resolve()), str(audio_path.resolve()) |
| 894 | |
| 895 | def materialize_r2v_payload(self, payload: dict[str, Any]) -> dict[str, Any]: |
| 896 | """Return a durable payload understood by the local inference loader.""" |
| 897 | |
| 898 | materialized = json.loads(_canonical_json(payload)) |
| 899 | condition_img = materialized.get("condition_img") |
| 900 | if isinstance(condition_img, str) and condition_img.startswith("data:"): |
| 901 | materialized["condition_img"] = self._write_inline( |
| 902 | condition_img, kind="image" |
| 903 | ) |
| 904 | |
| 905 | slots: list[dict[str, Any]] = [] |
| 906 | for raw_slot in materialized.get("memory_slots", []): |
| 907 | slot = dict(raw_slot) |
| 908 | reference_shot_id = str(slot.get("shot_id") or "").strip() |
| 909 | if reference_shot_id: |
| 910 | image_path, audio_path = self._extract_reference_assets( |
| 911 | work_id=str(materialized["work_id"]), |
| 912 | shot_id=reference_shot_id, |
| 913 | ) |
| 914 | metadata = dict(slot.get("metadata") or {}) |
| 915 | metadata.update( |
| 916 | { |
| 917 | "source_shot_id": reference_shot_id, |
| 918 | "resolved_by": "local_server", |
| 919 | } |
| 920 | ) |
| 921 | slot = { |
| 922 | "image_url": image_path, |
| 923 | "audio_url": audio_path, |
| 924 | "metadata": metadata, |
| 925 | } |
| 926 | else: |
| 927 | image_url = slot.get("image_url") |
| 928 | audio_url = slot.get("audio_url") |
| 929 | if isinstance(image_url, str) and image_url.startswith("data:"): |
| 930 | slot["image_url"] = self._write_inline(image_url, kind="image") |
| 931 | if isinstance(audio_url, str) and audio_url.startswith("data:"): |
| 932 | slot["audio_url"] = self._write_inline(audio_url, kind="audio") |
| 933 | slots.append(slot) |
| 934 | materialized["memory_slots"] = slots |
| 935 | return materialized |
| 936 | |
| 937 | |
| 938 | class MergeStore: |
| 939 | """In-memory merge FIFO mirrored to the durable restart journal.""" |
| 940 | |
| 941 | def __init__( |
| 942 | self, |
| 943 | capacity: int, |
| 944 | *, |
| 945 | journal: JobJournal | None = None, |
| 946 | journal_kind: str = "merge", |
| 947 | ) -> None: |
| 948 | self.capacity = capacity |
| 949 | self.journal = journal |
| 950 | self.journal_kind = journal_kind |
| 951 | self._jobs: OrderedDict[str, dict[str, Any]] = OrderedDict() |
| 952 | self._lock = threading.RLock() |
| 953 | |
| 954 | def initialize(self) -> None: |
| 955 | if self.journal is None: |
| 956 | return |
| 957 | self.journal.initialize() |
| 958 | restored = self.journal.load(self.journal_kind) |
| 959 | with self._lock: |
| 960 | self._jobs.clear() |
| 961 | for job in restored: |
| 962 | if job.get("status") == "running": |
| 963 | job.update(status="queued", updated_at=_now(), error=None) |
| 964 | self.journal.save(self.journal_kind, job) |
| 965 | self._jobs[str(job["version_id"])] = job |
| 966 | self._prune_terminal() |
| 967 | |
| 968 | def _persist(self, job: dict[str, Any]) -> None: |
| 969 | if self.journal is not None: |
| 970 | self.journal.save(self.journal_kind, job) |
| 971 | |
| 972 | def _prune_terminal(self) -> None: |
| 973 | while len(self._jobs) > self.capacity * 2: |
| 974 | terminal_id = next( |
| 975 | ( |
| 976 | version_id |
| 977 | for version_id, job in self._jobs.items() |
| 978 | if job["status"] in {"succeeded", "failed"} |
| 979 | and ( |
| 980 | not job.get("callback_url") |
| 981 | or job.get("callback_status") == "delivered" |
| 982 | ) |
| 983 | ), |
| 984 | None, |
| 985 | ) |
| 986 | if terminal_id is None: |
| 987 | return |
| 988 | self._jobs.pop(terminal_id, None) |
| 989 | |
| 990 | def enqueue( |
| 991 | self, |
| 992 | version_id: str, |
| 993 | request: MergeRequest, |
| 994 | *, |
| 995 | callback_url: str | None, |
| 996 | public_base_url: str, |
| 997 | ) -> str: |
| 998 | now = _now() |
| 999 | payload = request.model_dump( |
| 1000 | exclude_none=True, |
| 1001 | exclude={"callback_url", "callback_context", "job_id"}, |
| 1002 | ) |
| 1003 | payload["public_base_url"] = public_base_url |
| 1004 | request_json = _canonical_json(payload) |
| 1005 | with self._lock: |
| 1006 | if request.job_id: |
| 1007 | existing = next( |
| 1008 | ( |
| 1009 | job |
| 1010 | for job in self._jobs.values() |
| 1011 | if job.get("agent_job_id") == request.job_id |
| 1012 | ), |
| 1013 | None, |
| 1014 | ) |
| 1015 | if existing is None and self.journal is not None: |
| 1016 | existing = self.journal.get_by_agent_job_id( |
| 1017 | self.journal_kind, request.job_id |
| 1018 | ) |
| 1019 | if existing is not None: |
| 1020 | if existing["request_json"] != request_json: |
| 1021 | raise ValueError( |
| 1022 | "job_id is already associated with a different merge request" |
| 1023 | ) |
| 1024 | existing.update( |
| 1025 | callback_url=callback_url or existing.get("callback_url"), |
| 1026 | callback_context_json=json.dumps( |
| 1027 | request.callback_context, ensure_ascii=False |
| 1028 | ), |
| 1029 | updated_at=now, |
| 1030 | ) |
| 1031 | self._persist(existing) |
| 1032 | return str(existing["version_id"]) |
| 1033 | pending = sum( |
| 1034 | job["status"] in {"queued", "running"} for job in self._jobs.values() |
| 1035 | ) |
| 1036 | if pending >= self.capacity: |
| 1037 | raise OverflowError("merge queue is full") |
| 1038 | self._jobs[version_id] = { |
| 1039 | "version_id": version_id, |
| 1040 | "work_id": request.work_id, |
| 1041 | "agent_job_id": request.job_id, |
| 1042 | "request_json": request_json, |
| 1043 | "callback_url": callback_url, |
| 1044 | "callback_context_json": json.dumps( |
| 1045 | request.callback_context, ensure_ascii=False |
| 1046 | ), |
| 1047 | "callback_status": None, |
| 1048 | "callback_error": None, |
| 1049 | "callback_attempts": 0, |
| 1050 | "callback_next_at": None, |
| 1051 | "status": "queued", |
| 1052 | "video_url": None, |
| 1053 | "output_path": None, |
| 1054 | "error": None, |
| 1055 | "created_at": now, |
| 1056 | "updated_at": now, |
| 1057 | "started_at": None, |
| 1058 | "completed_at": None, |
| 1059 | } |
| 1060 | self._persist(self._jobs[version_id]) |
| 1061 | self._prune_terminal() |
| 1062 | return version_id |
| 1063 | |
| 1064 | def claim_next(self) -> dict[str, Any] | None: |
| 1065 | with self._lock: |
| 1066 | row = next( |
| 1067 | (job for job in self._jobs.values() if job["status"] == "queued"), |
| 1068 | None, |
| 1069 | ) |
| 1070 | if row is None: |
| 1071 | return None |
| 1072 | now = _now() |
| 1073 | row.update( |
| 1074 | status="running", |
| 1075 | started_at=row["started_at"] or now, |
| 1076 | updated_at=now, |
| 1077 | error=None, |
| 1078 | ) |
| 1079 | self._persist(row) |
| 1080 | return dict(row) |
| 1081 | |
| 1082 | def complete(self, version_id: str, *, video_url: str, output_path: Path) -> None: |
| 1083 | now = _now() |
| 1084 | with self._lock: |
| 1085 | self._jobs[version_id].update( |
| 1086 | status="succeeded", |
| 1087 | video_url=video_url, |
| 1088 | output_path=str(output_path), |
| 1089 | updated_at=now, |
| 1090 | completed_at=now, |
| 1091 | error=None, |
| 1092 | ) |
| 1093 | self._persist(self._jobs[version_id]) |
| 1094 | self._prune_terminal() |
| 1095 | |
| 1096 | def fail(self, version_id: str, error: str) -> None: |
| 1097 | now = _now() |
| 1098 | with self._lock: |
| 1099 | self._jobs[version_id].update( |
| 1100 | status="failed", |
| 1101 | error=error[:4000], |
| 1102 | updated_at=now, |
| 1103 | completed_at=now, |
| 1104 | ) |
| 1105 | self._persist(self._jobs[version_id]) |
| 1106 | self._prune_terminal() |
| 1107 | |
| 1108 | def get(self, version_id: str) -> dict[str, Any] | None: |
| 1109 | with self._lock: |
| 1110 | job = self._jobs.get(version_id) |
| 1111 | if job is not None: |
| 1112 | return dict(job) |
| 1113 | if self.journal is not None: |
| 1114 | return self.journal.get(self.journal_kind, version_id) |
| 1115 | return None |
| 1116 | |
| 1117 | def callback_candidates(self, max_attempts: int) -> list[dict[str, Any]]: |
| 1118 | now = _now() |
| 1119 | with self._lock: |
| 1120 | return [ |
| 1121 | dict(job) |
| 1122 | for job in self._jobs.values() |
| 1123 | if job["status"] in {"succeeded", "failed"} |
| 1124 | and job["callback_url"] |
| 1125 | and job["callback_status"] != "delivered" |
| 1126 | and (max_attempts == 0 or job["callback_attempts"] < max_attempts) |
| 1127 | and (job["callback_next_at"] is None or job["callback_next_at"] <= now) |
| 1128 | ] |
| 1129 | |
| 1130 | def record_callback(self, version_id: str, *, delivered: bool, error: str = "") -> None: |
| 1131 | with self._lock: |
| 1132 | job = self._jobs[version_id] |
| 1133 | attempts = int(job["callback_attempts"]) |
| 1134 | job.update( |
| 1135 | callback_status="delivered" if delivered else "error", |
| 1136 | callback_error=None if delivered else error[:1000] or None, |
| 1137 | callback_attempts=attempts + 1, |
| 1138 | callback_next_at=None |
| 1139 | if delivered |
| 1140 | else _callback_retry_at(attempts), |
| 1141 | updated_at=_now(), |
| 1142 | ) |
| 1143 | self._persist(job) |
| 1144 | |
| 1145 | def counts(self) -> dict[str, int]: |
| 1146 | result = {name: 0 for name in ("queued", "running", "succeeded", "failed")} |
| 1147 | with self._lock: |
| 1148 | for job in self._jobs.values(): |
| 1149 | result[str(job["status"])] += 1 |
| 1150 | return result |
| 1151 | |
| 1152 | |
| 1153 | class R2VQueueService: |
| 1154 | """In-memory local inference queue with one staged runtime per GPU.""" |
| 1155 | |
| 1156 | def __init__(self, settings: Settings) -> None: |
| 1157 | self.settings = settings |
| 1158 | self.journal = JobJournal(settings.artifact_db_path) |
| 1159 | self.store = JobStore(settings.queue_capacity, journal=self.journal) |
| 1160 | self.artifacts = ArtifactStore(settings.artifact_db_path) |
| 1161 | self.assets = RequestAssetResolver(settings, self.artifacts) |
| 1162 | self.stop_event = threading.Event() |
| 1163 | self.wake_event = threading.Event() |
| 1164 | self.threads: list[threading.Thread] = [] |
| 1165 | self.runtimes: dict[int, Any] = {} |
| 1166 | self.worker_states: dict[int, dict[str, Any]] = {} |
| 1167 | self._state_lock = threading.Lock() |
| 1168 | self._conditioning_lock = threading.Lock() |
| 1169 | |
| 1170 | def start(self) -> None: |
| 1171 | self.artifacts.initialize() |
| 1172 | self.store.initialize() |
| 1173 | self.threads = [ |
| 1174 | threading.Thread(target=self._callback_loop, name="r2v-callback", daemon=True) |
| 1175 | ] |
| 1176 | if not self.settings.disable_inference_workers: |
| 1177 | from .runtime import LocalModelRuntime, resolve_gpu_ids |
| 1178 | |
| 1179 | if not self.settings.inference_config.is_file(): |
| 1180 | raise FileNotFoundError( |
| 1181 | f"inference config not found: {self.settings.inference_config}" |
| 1182 | ) |
| 1183 | requests_root = self.settings.media_root / "requests" |
| 1184 | output_root = self.settings.media_root / "r2v" |
| 1185 | for gpu_id in resolve_gpu_ids(self.settings.gpu_ids): |
| 1186 | runtime = LocalModelRuntime( |
| 1187 | gpu_id=gpu_id, |
| 1188 | config_path=self.settings.inference_config, |
| 1189 | checkpoint=self.settings.checkpoint, |
| 1190 | conditioning_cache_dir=self.settings.conditioning_cache_dir, |
| 1191 | requests_root=requests_root, |
| 1192 | output_root=output_root, |
| 1193 | dit_residency=self.settings.dit_residency, |
| 1194 | gpu_headroom_fraction=self.settings.gpu_headroom_fraction, |
| 1195 | ram_headroom_fraction=self.settings.ram_headroom_fraction, |
| 1196 | ) |
| 1197 | self.runtimes[gpu_id] = runtime |
| 1198 | self.worker_states[gpu_id] = { |
| 1199 | "gpu_id": gpu_id, |
| 1200 | "state": "starting", |
| 1201 | "current_task": None, |
| 1202 | "model_loaded": False, |
| 1203 | "weights_loaded": False, |
| 1204 | "model_location": "unloaded", |
| 1205 | "resources": None, |
| 1206 | "error": None, |
| 1207 | } |
| 1208 | self.threads.append( |
| 1209 | threading.Thread( |
| 1210 | target=self._worker_loop, |
| 1211 | args=(gpu_id,), |
| 1212 | name=f"r2v-gpu-{gpu_id}", |
| 1213 | daemon=True, |
| 1214 | ) |
| 1215 | ) |
| 1216 | for thread in self.threads: |
| 1217 | thread.start() |
| 1218 | |
| 1219 | def stop(self) -> None: |
| 1220 | self.stop_event.set() |
| 1221 | self.wake_event.set() |
| 1222 | for thread in self.threads: |
| 1223 | thread.join(timeout=5) |
| 1224 | |
| 1225 | def enqueue( |
| 1226 | self, |
| 1227 | request: R2VGenerateRequest, |
| 1228 | payload: dict[str, Any], |
| 1229 | callback_url: str | None, |
| 1230 | ) -> str: |
| 1231 | version_id = str(uuid.uuid4()) |
| 1232 | version_id = self.store.enqueue(version_id, request, payload, callback_url) |
| 1233 | self.wake_event.set() |
| 1234 | return version_id |
| 1235 | |
| 1236 | def materialize_request(self, payload: dict[str, Any]) -> dict[str, Any]: |
| 1237 | return self.assets.materialize_r2v_payload(payload) |
| 1238 | |
| 1239 | def status(self) -> dict[str, Any]: |
| 1240 | with self._state_lock: |
| 1241 | workers = [dict(value) for _, value in sorted(self.worker_states.items())] |
| 1242 | return { |
| 1243 | "enabled": not self.settings.disable_inference_workers, |
| 1244 | "workers": workers, |
| 1245 | "queue_backend": "memory", |
| 1246 | "recovery_backend": "sqlite", |
| 1247 | "precision_config": str(self.settings.inference_config), |
| 1248 | "dit_residency": self.settings.dit_residency, |
| 1249 | "model_idle_seconds": self.settings.model_idle_seconds, |
| 1250 | "gpu_headroom_fraction": self.settings.gpu_headroom_fraction, |
| 1251 | "ram_headroom_fraction": self.settings.ram_headroom_fraction, |
| 1252 | } |
| 1253 | |
| 1254 | def _set_worker_state(self, gpu_id: int, **updates: Any) -> None: |
| 1255 | with self._state_lock: |
| 1256 | state = self.worker_states.setdefault(gpu_id, {"gpu_id": gpu_id}) |
| 1257 | state.update(updates) |
| 1258 | |
| 1259 | @staticmethod |
| 1260 | def _weights_loaded(runtime: Any) -> bool: |
| 1261 | return bool(getattr(runtime, "weights_loaded", runtime.model_loaded)) |
| 1262 | |
| 1263 | def _resource_ready(self, gpu_id: int) -> tuple[bool, dict[str, Any]]: |
| 1264 | from .runtime import probe_resources |
| 1265 | |
| 1266 | snapshot = probe_resources(gpu_id) |
| 1267 | payload = snapshot.as_dict() |
| 1268 | gpu_headroom = int(snapshot.gpu_total_bytes * self.settings.gpu_headroom_fraction) |
| 1269 | ram_headroom = int(snapshot.ram_total_bytes * self.settings.ram_headroom_fraction) |
| 1270 | runtime = self.runtimes[gpu_id] |
| 1271 | requirements = getattr(runtime, "admission_requirements", None) |
| 1272 | if requirements is not None: |
| 1273 | required_gpu, required_ram, plan = requirements(snapshot) |
| 1274 | else: |
| 1275 | required_gpu, required_ram, plan = gpu_headroom, ram_headroom, { |
| 1276 | "mode": "headroom_only" |
| 1277 | } |
| 1278 | ready = ( |
| 1279 | snapshot.gpu_free_bytes >= required_gpu |
| 1280 | and snapshot.ram_available_bytes >= required_ram |
| 1281 | ) |
| 1282 | payload["gpu_headroom_gib"] = round(gpu_headroom / 2**30, 3) |
| 1283 | payload["ram_headroom_gib"] = round(ram_headroom / 2**30, 3) |
| 1284 | payload["required_gpu_free_gib"] = round(required_gpu / 2**30, 3) |
| 1285 | payload["required_ram_available_gib"] = round(required_ram / 2**30, 3) |
| 1286 | payload["admission_plan"] = plan |
| 1287 | payload["ready"] = ready |
| 1288 | return ready, payload |
| 1289 | |
| 1290 | def _worker_loop(self, gpu_id: int) -> None: |
| 1291 | runtime = self.runtimes[gpu_id] |
| 1292 | while not self.stop_event.is_set(): |
| 1293 | if self.store.counts()["queued"] == 0: |
| 1294 | if ( |
| 1295 | self._weights_loaded(runtime) |
| 1296 | and self.settings.model_idle_seconds > 0 |
| 1297 | and time.monotonic() - runtime.last_used_at >= self.settings.model_idle_seconds |
| 1298 | ): |
| 1299 | runtime.unload() |
| 1300 | self._set_worker_state( |
| 1301 | gpu_id, |
| 1302 | state=runtime.state if self._weights_loaded(runtime) else "idle", |
| 1303 | current_task=None, |
| 1304 | model_loaded=runtime.model_loaded, |
| 1305 | weights_loaded=self._weights_loaded(runtime), |
| 1306 | model_location=getattr(runtime, "model_location", "unloaded"), |
| 1307 | error=None, |
| 1308 | ) |
| 1309 | self.wake_event.wait(timeout=0.5) |
| 1310 | self.wake_event.clear() |
| 1311 | continue |
| 1312 | try: |
| 1313 | ready, resources = self._resource_ready(gpu_id) |
| 1314 | except Exception as exc: # noqa: BLE001 |
| 1315 | self._set_worker_state(gpu_id, state="resource_error", error=str(exc)) |
| 1316 | self.stop_event.wait(self.settings.poll_interval) |
| 1317 | continue |
| 1318 | if not ready: |
| 1319 | if self._weights_loaded(runtime) and not runtime.model_loaded: |
| 1320 | # Decoder-only leftovers cannot make forward progress when |
| 1321 | # a cold generator load does not fit. Drop them and re-probe. |
| 1322 | runtime.unload() |
| 1323 | elif self._weights_loaded(runtime) and ( |
| 1324 | resources["ram_available_gib"] < resources["ram_headroom_gib"] |
| 1325 | ): |
| 1326 | runtime.unload() |
| 1327 | self._set_worker_state( |
| 1328 | gpu_id, |
| 1329 | state="waiting_resources", |
| 1330 | current_task=None, |
| 1331 | model_loaded=runtime.model_loaded, |
| 1332 | weights_loaded=self._weights_loaded(runtime), |
| 1333 | model_location=getattr(runtime, "model_location", "unloaded"), |
| 1334 | resources=resources, |
| 1335 | error=None, |
| 1336 | ) |
| 1337 | self.stop_event.wait(self.settings.poll_interval) |
| 1338 | continue |
| 1339 | job = self.store.claim_next() |
| 1340 | if job is None: |
| 1341 | continue |
| 1342 | self._run_job(gpu_id, job, resources) |
| 1343 | runtime.unload() |
| 1344 | |
| 1345 | def _request_file(self, payload: dict[str, Any]) -> Path: |
| 1346 | canonical = json.dumps( |
| 1347 | payload, ensure_ascii=False, sort_keys=True, separators=(",", ":") |
| 1348 | ) |
| 1349 | digest = sha256(canonical.encode("utf-8")).hexdigest() |
| 1350 | path = self.settings.media_root / "requests" / f"{digest}.json" |
| 1351 | if not path.is_file(): |
| 1352 | path.parent.mkdir(parents=True, exist_ok=True) |
| 1353 | temporary = path.with_suffix(f".{uuid.uuid4().hex}.tmp") |
| 1354 | temporary.write_text(canonical + "\n", encoding="utf-8") |
| 1355 | try: |
| 1356 | temporary.replace(path) |
| 1357 | finally: |
| 1358 | temporary.unlink(missing_ok=True) |
| 1359 | return path |
| 1360 | |
| 1361 | def _fail_job(self, gpu_id: int, job: dict[str, Any], error: str) -> None: |
| 1362 | runtime = self.runtimes[gpu_id] |
| 1363 | version_id = str(job["version_id"]) |
| 1364 | self.store.fail(version_id, error) |
| 1365 | self._set_worker_state( |
| 1366 | gpu_id, |
| 1367 | state="failed", |
| 1368 | current_task=version_id, |
| 1369 | model_loaded=runtime.model_loaded, |
| 1370 | weights_loaded=self._weights_loaded(runtime), |
| 1371 | model_location=getattr(runtime, "model_location", "unloaded"), |
| 1372 | error=error, |
| 1373 | ) |
| 1374 | failed = self.store.get(version_id) |
| 1375 | if failed is not None: |
| 1376 | self.deliver_callback(failed) |
| 1377 | |
| 1378 | def _run_job( |
| 1379 | self, |
| 1380 | gpu_id: int, |
| 1381 | job: dict[str, Any], |
| 1382 | resources: dict[str, Any], |
| 1383 | ) -> None: |
| 1384 | runtime = self.runtimes[gpu_id] |
| 1385 | version_id = str(job["version_id"]) |
| 1386 | try: |
| 1387 | self.store.update_stage( |
| 1388 | version_id, "validating", gpu_id=gpu_id, resources=resources |
| 1389 | ) |
| 1390 | payload = _json_loads(job["request_json"], {}) |
| 1391 | request_file = self._request_file(payload) |
| 1392 | request = runtime.load_request(request_file) |
| 1393 | except Exception as exc: # noqa: BLE001 |
| 1394 | traceback.print_exc() |
| 1395 | self._fail_job(gpu_id, job, f"{type(exc).__name__}: {exc}") |
| 1396 | return |
| 1397 | |
| 1398 | def update_condition_stage(stage: str) -> None: |
| 1399 | try: |
| 1400 | _, current_resources = self._resource_ready(gpu_id) |
| 1401 | except Exception: # noqa: BLE001 |
| 1402 | current_resources = resources |
| 1403 | self.store.update_stage( |
| 1404 | version_id, |
| 1405 | stage, |
| 1406 | gpu_id=gpu_id, |
| 1407 | resources=current_resources, |
| 1408 | ) |
| 1409 | self._set_worker_state( |
| 1410 | gpu_id, |
| 1411 | state=stage, |
| 1412 | current_task=version_id, |
| 1413 | model_loaded=runtime.model_loaded, |
| 1414 | weights_loaded=self._weights_loaded(runtime), |
| 1415 | model_location=getattr(runtime, "model_location", "unloaded"), |
| 1416 | resources=current_resources, |
| 1417 | error=None, |
| 1418 | ) |
| 1419 | |
| 1420 | try: |
| 1421 | cache_loader = getattr(runtime, "load_cached_conditions", None) |
| 1422 | bundles = ( |
| 1423 | cache_loader([request_file], [request], update_condition_stage) |
| 1424 | if cache_loader |
| 1425 | else None |
| 1426 | ) |
| 1427 | if bundles is None: |
| 1428 | # Gemma is the largest transient conditioning component. Only |
| 1429 | # cache misses are serialized; cache hits remain parallel. |
| 1430 | with self._conditioning_lock: |
| 1431 | from .runtime import probe_resources |
| 1432 | |
| 1433 | snapshot = probe_resources(gpu_id) |
| 1434 | policy_selector = getattr( |
| 1435 | runtime, "conditioning_generator_policy", None |
| 1436 | ) |
| 1437 | generator_policy = ( |
| 1438 | str(policy_selector(snapshot)) |
| 1439 | if policy_selector is not None |
| 1440 | else "release" |
| 1441 | ) |
| 1442 | bundles = runtime.prepare_conditions( |
| 1443 | [request_file], |
| 1444 | [request], |
| 1445 | update_condition_stage, |
| 1446 | generator_policy=generator_policy, |
| 1447 | ) |
| 1448 | bundle = bundles[request_file] |
| 1449 | except Exception as exc: # noqa: BLE001 |
| 1450 | error = f"{type(exc).__name__}: {exc}" |
| 1451 | if "out of memory" in error.lower(): |
| 1452 | runtime.unload() |
| 1453 | traceback.print_exc() |
| 1454 | self._fail_job(gpu_id, job, error) |
| 1455 | return |
| 1456 | |
| 1457 | self.store.update_stage( |
| 1458 | version_id, "ready", gpu_id=gpu_id, resources=resources |
| 1459 | ) |
| 1460 | self._run_prepared_job( |
| 1461 | gpu_id, |
| 1462 | job, |
| 1463 | request_file, |
| 1464 | request, |
| 1465 | bundle, |
| 1466 | resources, |
| 1467 | ) |
| 1468 | |
| 1469 | def _run_prepared_job( |
| 1470 | self, |
| 1471 | gpu_id: int, |
| 1472 | job: dict[str, Any], |
| 1473 | request_file: Path, |
| 1474 | request: Any, |
| 1475 | bundle: Any, |
| 1476 | resources: dict[str, Any], |
| 1477 | ) -> None: |
| 1478 | runtime = self.runtimes[gpu_id] |
| 1479 | version_id = str(job["version_id"]) |
| 1480 | |
| 1481 | def update_stage(stage: str) -> None: |
| 1482 | try: |
| 1483 | _, current_resources = self._resource_ready(gpu_id) |
| 1484 | except Exception: # noqa: BLE001 |
| 1485 | current_resources = resources |
| 1486 | self.store.update_stage( |
| 1487 | version_id, stage, gpu_id=gpu_id, resources=current_resources |
| 1488 | ) |
| 1489 | self._set_worker_state( |
| 1490 | gpu_id, |
| 1491 | state=stage, |
| 1492 | current_task=version_id, |
| 1493 | model_loaded=runtime.model_loaded, |
| 1494 | weights_loaded=self._weights_loaded(runtime), |
| 1495 | model_location=getattr(runtime, "model_location", "unloaded"), |
| 1496 | resources=current_resources, |
| 1497 | error=None, |
| 1498 | ) |
| 1499 | |
| 1500 | try: |
| 1501 | update_stage("loading_generator") |
| 1502 | output_dir = ( |
| 1503 | self.settings.media_root |
| 1504 | / "r2v" |
| 1505 | / request.work_id |
| 1506 | / request.shot_id |
| 1507 | / version_id |
| 1508 | ) |
| 1509 | result = runtime.run(request_file, request, output_dir, bundle, update_stage) |
| 1510 | output_path = Path(result["output_path"]).resolve() |
| 1511 | relative = output_path.relative_to(self.settings.media_root.expanduser().resolve()) |
| 1512 | video_url = f"{self.settings.public_base_url}/media/{quote(relative.as_posix())}" |
| 1513 | artifact = self.artifacts.upsert( |
| 1514 | version_id=version_id, |
| 1515 | work_id=str(job["work_id"]), |
| 1516 | shot_id=request.shot_id, |
| 1517 | kind="r2v", |
| 1518 | role="primary", |
| 1519 | url=video_url, |
| 1520 | local_path=str(output_path), |
| 1521 | size_bytes=output_path.stat().st_size, |
| 1522 | sha256=_file_sha256(output_path), |
| 1523 | metadata=result["metadata"], |
| 1524 | ) |
| 1525 | local_result = { |
| 1526 | "video_id": artifact.get("artifact_id"), |
| 1527 | "video_url": video_url, |
| 1528 | "artifact_url": video_url, |
| 1529 | "output_path": str(output_path), |
| 1530 | "metadata": result["metadata"], |
| 1531 | } |
| 1532 | self.store.complete(version_id, local_result) |
| 1533 | completed = self.store.get(version_id) |
| 1534 | if completed is not None: |
| 1535 | self.deliver_callback(completed) |
| 1536 | except Exception as exc: # noqa: BLE001 |
| 1537 | error = f"{type(exc).__name__}: {exc}" |
| 1538 | if "out of memory" in error.lower(): |
| 1539 | runtime.unload() |
| 1540 | elif runtime.model_loaded: |
| 1541 | runtime.state = "ready" |
| 1542 | traceback.print_exc() |
| 1543 | self._fail_job(gpu_id, job, error) |
| 1544 | finally: |
| 1545 | self._set_worker_state( |
| 1546 | gpu_id, |
| 1547 | state=runtime.state if self._weights_loaded(runtime) else "idle", |
| 1548 | current_task=None, |
| 1549 | model_loaded=runtime.model_loaded, |
| 1550 | weights_loaded=self._weights_loaded(runtime), |
| 1551 | model_location=getattr(runtime, "model_location", "unloaded"), |
| 1552 | ) |
| 1553 | |
| 1554 | def deliver_callback(self, job: dict[str, Any]) -> None: |
| 1555 | callback_url = job.get("callback_url") |
| 1556 | if not callback_url or job.get("callback_status") == "delivered": |
| 1557 | return |
| 1558 | result = _json_loads(job.get("result_json"), {}) |
| 1559 | completed = job["status"] == "succeeded" |
| 1560 | context = _json_loads(job.get("callback_context_json"), {}) |
| 1561 | body: dict[str, Any] = { |
| 1562 | key: value |
| 1563 | for key, value in context.items() |
| 1564 | if key in {"session_key", "channel", "chat_id"} and value is not None |
| 1565 | } |
| 1566 | body.update( |
| 1567 | { |
| 1568 | "work_id": job["work_id"], |
| 1569 | "job_id": job.get("agent_job_id"), |
| 1570 | "status": "completed" if completed else "failed", |
| 1571 | "shot_id": job["shot_id"], |
| 1572 | "remote_task_id": job["version_id"], |
| 1573 | } |
| 1574 | ) |
| 1575 | if completed: |
| 1576 | body.update( |
| 1577 | { |
| 1578 | "video_id": result.get("video_id"), |
| 1579 | "result_url": result.get("video_url"), |
| 1580 | "artifact_url": result.get("video_url"), |
| 1581 | } |
| 1582 | ) |
| 1583 | body = {key: value for key, value in body.items() if value is not None} |
| 1584 | else: |
| 1585 | body["error"] = job.get("error") or "generation failed" |
| 1586 | try: |
| 1587 | with httpx.Client(timeout=self.settings.request_timeout) as client: |
| 1588 | response = client.post( |
| 1589 | callback_url, |
| 1590 | json=body, |
| 1591 | headers={"Content-Type": "application/json", "Accept": "application/json"}, |
| 1592 | ) |
| 1593 | response.raise_for_status() |
| 1594 | self.store.record_callback( |
| 1595 | job["version_id"], delivered=True, response=getattr(response, "text", "") |
| 1596 | ) |
| 1597 | except httpx.HTTPError as exc: |
| 1598 | self.store.record_callback(job["version_id"], delivered=False, error=str(exc)) |
| 1599 | |
| 1600 | def _callback_loop(self) -> None: |
| 1601 | while not self.stop_event.wait(self.settings.poll_interval): |
| 1602 | for job in self.store.callback_candidates(self.settings.callback_attempts): |
| 1603 | if self.stop_event.is_set(): |
| 1604 | return |
| 1605 | self.deliver_callback(job) |
| 1606 | |
| 1607 | |
| 1608 | class MergeQueueService: |
| 1609 | """Download ordered server-side artifacts and assemble them with FFmpeg.""" |
| 1610 | |
| 1611 | def __init__(self, settings: Settings, r2v_service: R2VQueueService) -> None: |
| 1612 | self.settings = settings |
| 1613 | self.r2v_service = r2v_service |
| 1614 | self.store = MergeStore( |
| 1615 | settings.queue_capacity, |
| 1616 | journal=r2v_service.journal, |
| 1617 | ) |
| 1618 | self.artifacts = r2v_service.artifacts |
| 1619 | self.stop_event = threading.Event() |
| 1620 | self.wake_event = threading.Event() |
| 1621 | self.thread: threading.Thread | None = None |
| 1622 | |
| 1623 | def start(self) -> None: |
| 1624 | self.settings.media_root.expanduser().resolve().mkdir(parents=True, exist_ok=True) |
| 1625 | self.store.initialize() |
| 1626 | self.thread = threading.Thread(target=self._run_loop, name="video-merge", daemon=True) |
| 1627 | self.thread.start() |
| 1628 | |
| 1629 | def stop(self) -> None: |
| 1630 | self.stop_event.set() |
| 1631 | self.wake_event.set() |
| 1632 | if self.thread is not None: |
| 1633 | self.thread.join(timeout=5) |
| 1634 | |
| 1635 | def enqueue( |
| 1636 | self, |
| 1637 | request: MergeRequest, |
| 1638 | *, |
| 1639 | callback_url: str | None, |
| 1640 | public_base_url: str, |
| 1641 | ) -> str: |
| 1642 | version_id = str(uuid.uuid4()) |
| 1643 | version_id = self.store.enqueue( |
| 1644 | version_id, |
| 1645 | request, |
| 1646 | callback_url=callback_url, |
| 1647 | public_base_url=public_base_url, |
| 1648 | ) |
| 1649 | self.wake_event.set() |
| 1650 | return version_id |
| 1651 | |
| 1652 | def _source_url(self, shot: dict[str, Any]) -> str: |
| 1653 | version_id = str(shot.get("version_id") or "").strip() |
| 1654 | if version_id: |
| 1655 | source_job = self.r2v_service.store.get(version_id) |
| 1656 | if source_job is None: |
| 1657 | raise RuntimeError(f"R2V version not found: {version_id}") |
| 1658 | if source_job.get("status") != "succeeded": |
| 1659 | raise RuntimeError(f"R2V version is not ready: {version_id}") |
| 1660 | result = _json_loads(source_job.get("result_json"), {}) |
| 1661 | url = result.get("video_url") |
| 1662 | else: |
| 1663 | url = shot.get("video_url") |
| 1664 | if not isinstance(url, str) or not url.strip(): |
| 1665 | raise RuntimeError("merge input does not resolve to a video URL") |
| 1666 | parsed = urlparse(url.strip()) |
| 1667 | if parsed.scheme not in {"http", "https"} or not parsed.netloc: |
| 1668 | raise RuntimeError("merge video_url must be an absolute HTTP(S) URL") |
| 1669 | return url.strip() |
| 1670 | |
| 1671 | def _download(self, url: str, target: Path) -> None: |
| 1672 | total = 0 |
| 1673 | with httpx.stream( |
| 1674 | "GET", |
| 1675 | url, |
| 1676 | timeout=self.settings.request_timeout, |
| 1677 | follow_redirects=True, |
| 1678 | ) as response: |
| 1679 | response.raise_for_status() |
| 1680 | with target.open("wb") as output: |
| 1681 | for chunk in response.iter_bytes(chunk_size=1024 * 1024): |
| 1682 | total += len(chunk) |
| 1683 | if total > self.settings.merge_max_input_bytes: |
| 1684 | raise RuntimeError( |
| 1685 | f"merge input exceeds {self.settings.merge_max_input_bytes} bytes" |
| 1686 | ) |
| 1687 | output.write(chunk) |
| 1688 | if total == 0: |
| 1689 | raise RuntimeError(f"merge input is empty: {url}") |
| 1690 | |
| 1691 | def _run_ffmpeg(self, inputs: list[Path], output: Path, work_dir: Path) -> None: |
| 1692 | binary = shutil.which(self.settings.ffmpeg_binary) |
| 1693 | if binary is None: |
| 1694 | raise RuntimeError( |
| 1695 | f"FFmpeg executable was not found: {self.settings.ffmpeg_binary}" |
| 1696 | ) |
| 1697 | concat_file = work_dir / "inputs.txt" |
| 1698 | concat_file.write_text( |
| 1699 | "".join(f"file '{path.as_posix()}'\n" for path in inputs), |
| 1700 | encoding="utf-8", |
| 1701 | ) |
| 1702 | common = [ |
| 1703 | binary, |
| 1704 | "-nostdin", |
| 1705 | "-hide_banner", |
| 1706 | "-loglevel", |
| 1707 | "error", |
| 1708 | "-y", |
| 1709 | "-f", |
| 1710 | "concat", |
| 1711 | "-safe", |
| 1712 | "0", |
| 1713 | "-i", |
| 1714 | str(concat_file), |
| 1715 | ] |
| 1716 | copy_result = subprocess.run( |
| 1717 | [*common, "-c", "copy", "-movflags", "+faststart", str(output)], |
| 1718 | capture_output=True, |
| 1719 | text=True, |
| 1720 | timeout=self.settings.merge_timeout, |
| 1721 | check=False, |
| 1722 | ) |
| 1723 | if copy_result.returncode == 0 and output.is_file() and output.stat().st_size > 0: |
| 1724 | return |
| 1725 | output.unlink(missing_ok=True) |
| 1726 | encode_result = subprocess.run( |
| 1727 | [ |
| 1728 | *common, |
| 1729 | "-c:v", |
| 1730 | "libx264", |
| 1731 | "-preset", |
| 1732 | "medium", |
| 1733 | "-crf", |
| 1734 | "18", |
| 1735 | "-c:a", |
| 1736 | "aac", |
| 1737 | "-b:a", |
| 1738 | "192k", |
| 1739 | "-movflags", |
| 1740 | "+faststart", |
| 1741 | str(output), |
| 1742 | ], |
| 1743 | capture_output=True, |
| 1744 | text=True, |
| 1745 | timeout=self.settings.merge_timeout, |
| 1746 | check=False, |
| 1747 | ) |
| 1748 | if encode_result.returncode != 0 or not output.is_file() or output.stat().st_size == 0: |
| 1749 | detail = (encode_result.stderr or copy_result.stderr or "unknown FFmpeg error")[-2000:] |
| 1750 | raise RuntimeError(f"FFmpeg merge failed: {detail}") |
| 1751 | |
| 1752 | def _process(self, job: dict[str, Any]) -> None: |
| 1753 | payload = _json_loads(job.get("request_json"), {}) |
| 1754 | shots = payload.get("shots") |
| 1755 | if not isinstance(shots, list) or not shots: |
| 1756 | raise RuntimeError("merge request has no shots") |
| 1757 | media_root = self.settings.media_root.expanduser().resolve() |
| 1758 | output_dir = media_root / "merges" |
| 1759 | output_dir.mkdir(parents=True, exist_ok=True) |
| 1760 | output = output_dir / f"{job['version_id']}.mp4" |
| 1761 | output.unlink(missing_ok=True) |
| 1762 | work_dir = Path(tempfile.mkdtemp(prefix="merge-", dir=media_root)) |
| 1763 | try: |
| 1764 | inputs: list[Path] = [] |
| 1765 | for index, shot in enumerate(shots, start=1): |
| 1766 | if not isinstance(shot, dict): |
| 1767 | raise RuntimeError(f"merge shot {index} is invalid") |
| 1768 | target = work_dir / f"shot-{index:04d}.mp4" |
| 1769 | self._download(self._source_url(shot), target) |
| 1770 | inputs.append(target) |
| 1771 | self._run_ffmpeg(inputs, output, work_dir) |
| 1772 | finally: |
| 1773 | shutil.rmtree(work_dir, ignore_errors=True) |
| 1774 | |
| 1775 | base = str(payload.get("public_base_url") or "").rstrip("/") |
| 1776 | video_url = f"{base}/media/merges/{quote(output.name)}" |
| 1777 | digest = _file_sha256(output) |
| 1778 | self.artifacts.upsert( |
| 1779 | version_id=str(job["version_id"]), |
| 1780 | work_id=str(job["work_id"]), |
| 1781 | kind="merge", |
| 1782 | role="merged", |
| 1783 | url=video_url, |
| 1784 | local_path=str(output), |
| 1785 | size_bytes=output.stat().st_size, |
| 1786 | sha256=digest, |
| 1787 | metadata={"input_count": len(inputs)}, |
| 1788 | ) |
| 1789 | self.store.complete(str(job["version_id"]), video_url=video_url, output_path=output) |
| 1790 | |
| 1791 | def deliver_callback(self, job: dict[str, Any]) -> None: |
| 1792 | callback_url = str(job.get("callback_url") or "").strip() |
| 1793 | if not callback_url or job.get("callback_status") == "delivered": |
| 1794 | return |
| 1795 | context = _json_loads(job.get("callback_context_json"), {}) |
| 1796 | body: dict[str, Any] = { |
| 1797 | key: value |
| 1798 | for key, value in context.items() |
| 1799 | if key in {"session_key", "channel", "chat_id"} and value is not None |
| 1800 | } |
| 1801 | completed = job.get("status") == "succeeded" |
| 1802 | body.update( |
| 1803 | { |
| 1804 | "work_id": job["work_id"], |
| 1805 | "job_id": job.get("agent_job_id"), |
| 1806 | "remote_task_id": job["version_id"], |
| 1807 | "status": "completed" if completed else "failed", |
| 1808 | } |
| 1809 | ) |
| 1810 | if completed: |
| 1811 | body["result"] = { |
| 1812 | "artifact_url": job.get("video_url"), |
| 1813 | "result_url": job.get("video_url"), |
| 1814 | } |
| 1815 | else: |
| 1816 | body["error"] = job.get("error") or "merge failed" |
| 1817 | body = {key: value for key, value in body.items() if value is not None} |
| 1818 | try: |
| 1819 | with httpx.Client(timeout=self.settings.request_timeout) as client: |
| 1820 | response = client.post(callback_url, json=body) |
| 1821 | response.raise_for_status() |
| 1822 | self.store.record_callback(str(job["version_id"]), delivered=True) |
| 1823 | except httpx.HTTPError as exc: |
| 1824 | self.store.record_callback( |
| 1825 | str(job["version_id"]), delivered=False, error=str(exc) |
| 1826 | ) |
| 1827 | |
| 1828 | def _run_loop(self) -> None: |
| 1829 | while not self.stop_event.is_set(): |
| 1830 | job = self.store.claim_next() |
| 1831 | if job is not None: |
| 1832 | try: |
| 1833 | self._process(job) |
| 1834 | except (OSError, RuntimeError, subprocess.SubprocessError, httpx.HTTPError) as exc: |
| 1835 | self.store.fail(str(job["version_id"]), str(exc)) |
| 1836 | completed = self.store.get(str(job["version_id"])) |
| 1837 | if completed is not None: |
| 1838 | self.deliver_callback(completed) |
| 1839 | continue |
| 1840 | for candidate in self.store.callback_candidates(self.settings.callback_attempts): |
| 1841 | if self.stop_event.is_set(): |
| 1842 | return |
| 1843 | self.deliver_callback(candidate) |
| 1844 | self.wake_event.wait(timeout=0.5) |
| 1845 | self.wake_event.clear() |
| 1846 | |
| 1847 | |
| 1848 | def _validate_callback_url(value: str | None) -> str | None: |
| 1849 | if value is None: |
| 1850 | return None |
| 1851 | value = value.strip() |
| 1852 | if not value: |
| 1853 | return None |
| 1854 | parsed = urlparse(value) |
| 1855 | if parsed.scheme not in {"http", "https"} or not parsed.netloc: |
| 1856 | raise HTTPException(status_code=422, detail="callback URL must be absolute HTTP(S)") |
| 1857 | return value |
| 1858 | |
| 1859 | |
| 1860 | def _validate_local_resource(value: str, field_name: str, *, kind: str) -> None: |
| 1861 | """Reject ambiguous paths before they enter the local GPU queue.""" |
| 1862 | |
| 1863 | if value.startswith("data:"): |
| 1864 | limit = MAX_IMAGE_BYTES if kind == "image" else MAX_AUDIO_BYTES |
| 1865 | try: |
| 1866 | _decode_data_url(value, kind=kind, max_bytes=limit) |
| 1867 | except ValueError as exc: |
| 1868 | raise HTTPException(status_code=422, detail=f"{field_name}: {exc}") from exc |
| 1869 | return |
| 1870 | parsed = urlparse(value) |
| 1871 | if parsed.scheme in {"http", "https"}: |
| 1872 | if not parsed.netloc: |
| 1873 | raise HTTPException( |
| 1874 | status_code=422, detail=f"{field_name} must be an absolute HTTP(S) URL" |
| 1875 | ) |
| 1876 | return |
| 1877 | if parsed.scheme == "file": |
| 1878 | if parsed.netloc not in {"", "localhost"}: |
| 1879 | raise HTTPException( |
| 1880 | status_code=422, detail=f"{field_name} does not support remote file URLs" |
| 1881 | ) |
| 1882 | local_path = Path(unquote(parsed.path)) |
| 1883 | elif parsed.scheme: |
| 1884 | raise HTTPException( |
| 1885 | status_code=422, |
| 1886 | detail=( |
| 1887 | f"{field_name} must use a data URL, HTTP(S), file://, " |
| 1888 | "or an absolute local path" |
| 1889 | ), |
| 1890 | ) |
| 1891 | else: |
| 1892 | local_path = Path(value).expanduser() |
| 1893 | if not local_path.is_absolute(): |
| 1894 | raise HTTPException( |
| 1895 | status_code=422, |
| 1896 | detail=f"{field_name} must be absolute when submitted through server.py", |
| 1897 | ) |
| 1898 | if not local_path.is_file(): |
| 1899 | raise HTTPException(status_code=422, detail=f"{field_name} does not exist") |
| 1900 | |
| 1901 | |
| 1902 | def _validate_local_r2v_request(body: R2VGenerateRequest) -> None: |
| 1903 | if body.condition_img: |
| 1904 | _validate_local_resource(body.condition_img, "condition_img", kind="image") |
| 1905 | for index, slot in enumerate(body.memory_slots): |
| 1906 | if slot.shot_id: |
| 1907 | continue |
| 1908 | assert slot.image_url is not None |
| 1909 | _validate_local_resource( |
| 1910 | slot.image_url, |
| 1911 | f"memory_slots[{index}].image_url", |
| 1912 | kind="image", |
| 1913 | ) |
| 1914 | if slot.audio_url: |
| 1915 | _validate_local_resource( |
| 1916 | slot.audio_url, |
| 1917 | f"memory_slots[{index}].audio_url", |
| 1918 | kind="audio", |
| 1919 | ) |
| 1920 | |
| 1921 | |
| 1922 | def _service(request: Request) -> R2VQueueService: |
| 1923 | service = getattr(request.app.state, "r2v_service", None) |
| 1924 | if service is None: |
| 1925 | raise HTTPException(status_code=503, detail="R2V queue is not ready") |
| 1926 | return service |
| 1927 | |
| 1928 | |
| 1929 | def _merge_service(request: Request) -> MergeQueueService: |
| 1930 | service = getattr(request.app.state, "merge_service", None) |
| 1931 | if service is None: |
| 1932 | raise HTTPException(status_code=503, detail="merge queue is not ready") |
| 1933 | return service |
| 1934 | |
| 1935 | |
| 1936 | def _status_url(request: Request, settings: Settings, version_id: str) -> str: |
| 1937 | base = settings.public_base_url or str(request.base_url).rstrip("/") |
| 1938 | return f"{base}/version/{quote(version_id, safe='')}" |
| 1939 | |
| 1940 | |
| 1941 | def _public_job(request: Request, service: R2VQueueService, job: dict[str, Any]) -> dict[str, Any]: |
| 1942 | payload = _json_loads(job.get("request_json"), {}) |
| 1943 | local_result = _json_loads(job.get("result_json"), {}) |
| 1944 | resources = _json_loads(job.get("resource_json"), {}) |
| 1945 | result = { |
| 1946 | "accepted": True, |
| 1947 | "kind": "r2v", |
| 1948 | "task_id": job["version_id"], |
| 1949 | "version_id": job["version_id"], |
| 1950 | "remote_task_id": job["version_id"], |
| 1951 | "work_id": job["work_id"], |
| 1952 | "job_id": job.get("agent_job_id"), |
| 1953 | "shot_id": job["shot_id"], |
| 1954 | "status": job["status"], |
| 1955 | "stage": job.get("stage") or job["status"], |
| 1956 | "queue_position": service.store.queue_position(job["version_id"]), |
| 1957 | "status_url": _status_url(request, service.settings, job["version_id"]), |
| 1958 | "gpu_id": job.get("gpu_id"), |
| 1959 | "resources": resources or None, |
| 1960 | "video_id": local_result.get("video_id"), |
| 1961 | "video_url": local_result.get("video_url"), |
| 1962 | "width": payload.get("width"), |
| 1963 | "height": payload.get("height"), |
| 1964 | "memory_slots": payload.get("memory_slots", []), |
| 1965 | "error": job.get("error"), |
| 1966 | "callback_status": job.get("callback_status"), |
| 1967 | "callback_error": job.get("callback_error"), |
| 1968 | "created_at": job["created_at"], |
| 1969 | "updated_at": job["updated_at"], |
| 1970 | "started_at": job.get("started_at"), |
| 1971 | "completed_at": job.get("completed_at"), |
| 1972 | "artifacts": service.artifacts.for_version(str(job["version_id"])), |
| 1973 | } |
| 1974 | if local_result: |
| 1975 | result["result"] = local_result |
| 1976 | return result |
| 1977 | |
| 1978 | |
| 1979 | def _public_merge_job( |
| 1980 | request: Request, |
| 1981 | service: MergeQueueService, |
| 1982 | job: dict[str, Any], |
| 1983 | ) -> dict[str, Any]: |
| 1984 | return { |
| 1985 | "accepted": True, |
| 1986 | "kind": "merge", |
| 1987 | "task_id": job["version_id"], |
| 1988 | "version_id": job["version_id"], |
| 1989 | "work_id": job["work_id"], |
| 1990 | "job_id": job.get("agent_job_id"), |
| 1991 | "status": job["status"], |
| 1992 | "status_url": _status_url(request, service.settings, str(job["version_id"])), |
| 1993 | "video_url": job.get("video_url"), |
| 1994 | "error": job.get("error"), |
| 1995 | "callback_status": job.get("callback_status"), |
| 1996 | "callback_error": job.get("callback_error"), |
| 1997 | "created_at": job["created_at"], |
| 1998 | "updated_at": job["updated_at"], |
| 1999 | "started_at": job.get("started_at"), |
| 2000 | "completed_at": job.get("completed_at"), |
| 2001 | "artifacts": service.artifacts.for_version(str(job["version_id"])), |
| 2002 | } |
| 2003 | |
| 2004 | |
| 2005 | def _parse_merge_request( |
| 2006 | body: dict[str, Any], |
| 2007 | callback_header: str | None, |
| 2008 | ) -> tuple[MergeRequest, str | None]: |
| 2009 | payload_source = body.get("payload") |
| 2010 | payload = dict(payload_source) if isinstance(payload_source, dict) else dict(body) |
| 2011 | job = body.get("job") |
| 2012 | if isinstance(job, dict) and not payload.get("job_id"): |
| 2013 | payload["job_id"] = job.get("job_id") |
| 2014 | callback = body.get("callback") |
| 2015 | if isinstance(callback, dict): |
| 2016 | payload["callback_context"] = { |
| 2017 | key: callback.get(key) for key in ("session_key", "channel", "chat_id") |
| 2018 | } |
| 2019 | payload["callback_url"] = callback_header or callback.get("url") |
| 2020 | elif callback_header: |
| 2021 | payload["callback_url"] = callback_header |
| 2022 | try: |
| 2023 | parsed = MergeRequest.model_validate(payload) |
| 2024 | except ValidationError as exc: |
| 2025 | raise HTTPException(status_code=422, detail=exc.errors()) from exc |
| 2026 | return parsed, _validate_callback_url(parsed.callback_url) |
| 2027 | |
| 2028 | |
| 2029 | @asynccontextmanager |
| 2030 | async def lifespan(app: FastAPI): |
| 2031 | loop = asyncio.get_running_loop() |
| 2032 | previous_exception_handler = loop.get_exception_handler() |
| 2033 | |
| 2034 | def handle_loop_exception( |
| 2035 | current_loop: asyncio.AbstractEventLoop, |
| 2036 | context: dict[str, Any], |
| 2037 | ) -> None: |
| 2038 | exc = context.get("exception") |
| 2039 | message = str(context.get("message") or "") |
| 2040 | if ( |
| 2041 | isinstance(exc, ConnectionResetError) |
| 2042 | and getattr(exc, "winerror", None) == 10054 |
| 2043 | and "_call_connection_lost" in message |
| 2044 | ): |
| 2045 | return |
| 2046 | if previous_exception_handler is not None: |
| 2047 | previous_exception_handler(current_loop, context) |
| 2048 | else: |
| 2049 | current_loop.default_exception_handler(context) |
| 2050 | |
| 2051 | loop.set_exception_handler(handle_loop_exception) |
| 2052 | settings = Settings.from_env() |
| 2053 | r2v_service = R2VQueueService(settings) |
| 2054 | merge_service = MergeQueueService(settings, r2v_service) |
| 2055 | r2v_service.start() |
| 2056 | merge_service.start() |
| 2057 | app.state.r2v_service = r2v_service |
| 2058 | app.state.merge_service = merge_service |
| 2059 | try: |
| 2060 | yield |
| 2061 | finally: |
| 2062 | merge_service.stop() |
| 2063 | r2v_service.stop() |
| 2064 | app.state.r2v_service = None |
| 2065 | app.state.merge_service = None |
| 2066 | loop.set_exception_handler(previous_exception_handler) |
| 2067 | |
| 2068 | |
| 2069 | app = FastAPI( |
| 2070 | title="Echo 1.5 Server", |
| 2071 | version="1.5", |
| 2072 | description="R2V generation, video merge, status, and health APIs for Echo Director.", |
| 2073 | lifespan=lifespan, |
| 2074 | ) |
| 2075 | |
| 2076 | |
| 2077 | @app.get("/health") |
| 2078 | async def health(request: Request) -> dict[str, Any]: |
| 2079 | service = _service(request) |
| 2080 | merge_service = _merge_service(request) |
| 2081 | inference = service.status() |
| 2082 | r2v_counts = service.store.counts() |
| 2083 | workers = inference["workers"] |
| 2084 | worker_count = len(workers) |
| 2085 | gpu_busy = sum(worker.get("current_task") is not None for worker in workers) |
| 2086 | return { |
| 2087 | "status": "ok", |
| 2088 | "inference": inference, |
| 2089 | "queue": r2v_counts, |
| 2090 | "queues": { |
| 2091 | "r2v": r2v_counts, |
| 2092 | "merge": merge_service.store.counts(), |
| 2093 | }, |
| 2094 | "scheduler": { |
| 2095 | "queues": { |
| 2096 | "inference": { |
| 2097 | "workers": worker_count, |
| 2098 | "busy": gpu_busy, |
| 2099 | "idle": max(worker_count - gpu_busy, 0), |
| 2100 | "pending": r2v_counts["queued"], |
| 2101 | } |
| 2102 | }, |
| 2103 | "gpu_workers": worker_count, |
| 2104 | "gpu_busy": gpu_busy, |
| 2105 | "gpu_idle": max(worker_count - gpu_busy, 0), |
| 2106 | "active_tasks": r2v_counts["queued"] + r2v_counts["running"], |
| 2107 | "total_queued": r2v_counts["queued"], |
| 2108 | }, |
| 2109 | "ffmpeg_available": shutil.which(service.settings.ffmpeg_binary) is not None, |
| 2110 | } |
| 2111 | |
| 2112 | |
| 2113 | @app.post("/r2v") |
| 2114 | async def generate_r2v( |
| 2115 | body: R2VGenerateRequest, |
| 2116 | request: Request, |
| 2117 | callback_header: str | None = Header( |
| 2118 | default=None, |
| 2119 | alias="X-Nanobot-Director-Callback-Url", |
| 2120 | ), |
| 2121 | ) -> dict[str, Any]: |
| 2122 | service = _service(request) |
| 2123 | _validate_local_r2v_request(body) |
| 2124 | callback_url = _validate_callback_url(callback_header or body.callback_url) |
| 2125 | payload = body.model_dump( |
| 2126 | exclude_none=True, |
| 2127 | exclude={"callback_url", "callback_context", "job_id"}, |
| 2128 | ) |
| 2129 | try: |
| 2130 | payload = service.materialize_request(payload) |
| 2131 | version_id = service.enqueue(body, payload, callback_url) |
| 2132 | except LookupError as exc: |
| 2133 | raise HTTPException(status_code=409, detail=str(exc)) from exc |
| 2134 | except (OSError, RuntimeError, subprocess.SubprocessError, ValueError) as exc: |
| 2135 | status_code = 409 if "job_id is already associated" in str(exc) else 422 |
| 2136 | raise HTTPException(status_code=status_code, detail=str(exc)) from exc |
| 2137 | except OverflowError as exc: |
| 2138 | raise HTTPException(status_code=503, detail=str(exc)) from exc |
| 2139 | job = service.store.get(version_id) |
| 2140 | assert job is not None |
| 2141 | return _public_job(request, service, job) |
| 2142 | |
| 2143 | |
| 2144 | @app.post("/merge") |
| 2145 | async def merge_videos( |
| 2146 | body: dict[str, Any], |
| 2147 | request: Request, |
| 2148 | callback_header: str | None = Header( |
| 2149 | default=None, |
| 2150 | alias="X-Nanobot-Director-Callback-Url", |
| 2151 | ), |
| 2152 | ) -> dict[str, Any]: |
| 2153 | service = _merge_service(request) |
| 2154 | merge_request, callback_url = _parse_merge_request(body, callback_header) |
| 2155 | public_base_url = service.settings.public_base_url or str(request.base_url).rstrip("/") |
| 2156 | try: |
| 2157 | version_id = service.enqueue( |
| 2158 | merge_request, |
| 2159 | callback_url=callback_url, |
| 2160 | public_base_url=public_base_url, |
| 2161 | ) |
| 2162 | except ValueError as exc: |
| 2163 | raise HTTPException(status_code=409, detail=str(exc)) from exc |
| 2164 | except OverflowError as exc: |
| 2165 | raise HTTPException(status_code=503, detail=str(exc)) from exc |
| 2166 | job = service.store.get(version_id) |
| 2167 | assert job is not None |
| 2168 | return _public_merge_job(request, service, job) |
| 2169 | |
| 2170 | |
| 2171 | @app.get("/version/{version_id}") |
| 2172 | async def get_version( |
| 2173 | version_id: str, |
| 2174 | request: Request, |
| 2175 | ) -> dict[str, Any]: |
| 2176 | service = _service(request) |
| 2177 | job = service.store.get(version_id) |
| 2178 | if job is not None: |
| 2179 | return _public_job(request, service, job) |
| 2180 | merge_service = _merge_service(request) |
| 2181 | merge_job = merge_service.store.get(version_id) |
| 2182 | if merge_job is not None: |
| 2183 | return _public_merge_job(request, merge_service, merge_job) |
| 2184 | raise HTTPException(status_code=404, detail="version not found") |
| 2185 | |
| 2186 | |
| 2187 | @app.get("/artifact/{artifact_id}") |
| 2188 | async def get_artifact_metadata( |
| 2189 | artifact_id: str, |
| 2190 | request: Request, |
| 2191 | ) -> dict[str, Any]: |
| 2192 | service = _service(request) |
| 2193 | artifact = service.artifacts.get(artifact_id) |
| 2194 | if artifact is None: |
| 2195 | raise HTTPException(status_code=404, detail="artifact not found") |
| 2196 | return artifact |
| 2197 | |
| 2198 | |
| 2199 | @app.get("/media/{asset_path:path}") |
| 2200 | async def get_media(asset_path: str, request: Request) -> FileResponse: |
| 2201 | service = _service(request) |
| 2202 | root = service.settings.media_root.expanduser().resolve() |
| 2203 | candidate = (root / asset_path).resolve() |
| 2204 | try: |
| 2205 | candidate.relative_to(root) |
| 2206 | except ValueError as exc: |
| 2207 | raise HTTPException(status_code=404, detail="media not found") from exc |
| 2208 | if not candidate.is_file(): |
| 2209 | raise HTTPException(status_code=404, detail="media not found") |
| 2210 | return FileResponse(candidate, media_type="video/mp4") |
| 2211 | |
| 2212 | |
| 2213 | def main() -> None: |
| 2214 | parser = argparse.ArgumentParser(description="Run the Echo 1.5 local server") |
| 2215 | parser.add_argument( |
| 2216 | "--config", |
| 2217 | required=True, |
| 2218 | help="Server YAML that references an inference YAML", |
| 2219 | ) |
| 2220 | parser.add_argument("--host", help="Override server.host") |
| 2221 | parser.add_argument("--port", type=int, help="Override server.port") |
| 2222 | args = parser.parse_args() |
| 2223 | |
| 2224 | config_path = _repo_path(args.config) |
| 2225 | config = apply_server_config(config_path) |
| 2226 | server_config = _config_section(config, "server") |
| 2227 | host = args.host or str(server_config.get("host", "127.0.0.1")) |
| 2228 | port = args.port or int(server_config.get("port", 8221)) |
| 2229 | workers = int(server_config.get("workers", 1)) |
| 2230 | if workers != 1: |
| 2231 | parser.error("server.workers must be 1 because GPU workers live in process memory") |
| 2232 | |
| 2233 | os.environ["ECHO_SERVER_CONFIG"] = str(config_path) |
| 2234 | import uvicorn |
| 2235 | |
| 2236 | uvicorn.run("server.app:app", host=host, port=port, workers=workers) |
| 2237 | |
| 2238 | |
| 2239 | if __name__ == "__main__": |
| 2240 | main() |
| 2241 |