Files
mule-image/backend/app/services/vision/bootstrap_models.py
root f090a809a9 fix: harden pipeline — retries, acks_late, time limits, session safety
Addresses 16 robustness, transparency, and performance issues across
the Celery media processing pipeline:

Critical:
- Singleton DB engine in vision tasks (was leaking one per task call)
- acks_late + task_reject_on_worker_lost so crashed workers don't lose tasks
- Global soft/hard time limits (5/10 min) to prevent hung worker slots
- Thumbnail copy-before-resize (in-place mutation degraded larger sizes)
- backfill_vision now checks each task type independently (OCR, faces, etc.)
- Parameterized LIMIT in backfill_vision (was f-string SQL injection)

High:
- try/except + retry(max=3) on all vision inference tasks
- extract_metadata writes processing_error on exiftool failure
- PIL Image handles closed in _load_thumb/_load_original
- Scan progress Redis keys auto-expire after 1 hour
- Watcher lock renewal is wall-clock based (30s) not event-count based
- worker_process_init signal warms up vision models on startup

Medium:
- Explicit task_routes for every task name (wildcards never matched)
- app.services.metadata added to Celery include list
- POST /maintenance/recover-stuck endpoint for photos stuck in processing
- Docker healthchecks for worker-light, worker-vision, and Redis
- Task ID in vision log lines for distributed tracing
- Bare except:pass narrowed to specific exceptions

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-04-13 08:57:10 +02:00

120 lines
3.9 KiB
Python

"""
Download vision model weights on first worker boot.
Run as: python -m app.services.vision.bootstrap_models
Or called from the vision worker entrypoint before Celery starts.
Downloads are idempotent — existing files are skipped.
For models that require export (OpenCLIP, YOLOv8n), see export_models.py.
Those must be exported once on any machine with pip, then placed in
the models volume before the worker starts.
"""
import logging
import os
from pathlib import Path
from urllib.request import urlretrieve
from app.config import settings
logger = logging.getLogger(__name__)
# (relative_path, url, description)
# Models with url=None must be pre-exported via export_models.py.
# InsightFace (RetinaFace + ArcFace) auto-downloads via the insightface
# package on first use — no manual download entries needed.
DOWNLOADS = []
# Models that need manual export via export_models.py
EXPORTS = [
("embed/visual.onnx", "OpenCLIP ViT-B/32 visual encoder"),
("embed/textual.onnx", "OpenCLIP ViT-B/32 textual encoder"),
("detect/yolov8n.onnx", "YOLOv8n object detector"),
]
def bootstrap(models_dir: str | None = None):
"""Ensure all model files are present. Download what we can, warn about
files that need manual export."""
base = Path(models_dir or settings.vision.models_dir)
base.mkdir(parents=True, exist_ok=True)
# Download auto-downloadable models
for rel_path, url, desc in DOWNLOADS:
dest = base / rel_path
dest.parent.mkdir(parents=True, exist_ok=True)
if dest.exists():
logger.debug("Already exists: %s (%s)", dest, desc)
continue
logger.info("Downloading %s%s", desc, dest)
try:
urlretrieve(url, str(dest))
size_kb = dest.stat().st_size / 1024
logger.info("Downloaded %s (%.0f KB)", desc, size_kb)
except Exception as e:
logger.error("Failed to download %s: %s", desc, e)
if dest.exists():
dest.unlink()
# Check for manually-exported models
missing = []
for rel_path, desc in EXPORTS:
dest = base / rel_path
if not dest.exists():
missing.append((rel_path, desc))
if missing:
logger.warning(
"Missing %d model file(s); attempting automatic export:",
len(missing),
)
for rel_path, desc in missing:
logger.warning(" %s%s", base / rel_path, desc)
try:
from app.services.vision import export_models
export_models.export_openclip(base)
export_models.export_siglip2(base)
export_models.export_yolov8n(base)
except Exception as e:
logger.error(
"Automatic export failed: %s. "
"Run `python -m app.services.vision.export_models "
"--models-dir %s` manually before starting the worker.",
e,
base,
)
return
# Re-check what's still missing after the export pass.
still_missing = [
(rel_path, desc)
for rel_path, desc in EXPORTS
if not (base / rel_path).exists()
]
if still_missing:
for rel_path, desc in still_missing:
logger.error(" still missing: %s%s", base / rel_path, desc)
else:
logger.info("All model files present in %s", base)
else:
logger.info("All model files present in %s", base)
# Signal readiness via Redis so the scan pipeline knows the vision
# worker can accept tasks.
try:
import redis as _redis
r = _redis.from_url(settings.redis_url)
r.set("mulita:vision:ready", "1")
logger.info("Set mulita:vision:ready in Redis")
except Exception as e:
logger.warning("Could not set vision readiness flag in Redis: %s", e)
if __name__ == "__main__":
logging.basicConfig(level=logging.INFO)
bootstrap()