f090a809a9
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>
494 lines
18 KiB
Python
494 lines
18 KiB
Python
"""
|
||
Celery tasks for thumbnail generation
|
||
"""
|
||
import os
|
||
import asyncio
|
||
from pathlib import Path
|
||
import logging
|
||
from typing import Tuple, Optional
|
||
import json
|
||
|
||
from celery import shared_task
|
||
from sqlalchemy import select
|
||
from sqlalchemy.ext.asyncio import AsyncSession
|
||
from PIL import Image
|
||
import imageio
|
||
from pillow_heif import register_heif_opener
|
||
import ffmpeg
|
||
|
||
# Try to import optional libraries
|
||
try:
|
||
import pyvips
|
||
PYVIPS_AVAILABLE = True
|
||
except ImportError:
|
||
PYVIPS_AVAILABLE = False
|
||
print("pyvips not available, using Pillow for image processing")
|
||
|
||
try:
|
||
import rawpy
|
||
RAWPY_AVAILABLE = True
|
||
except ImportError:
|
||
RAWPY_AVAILABLE = False
|
||
print("rawpy not available, using exiftool for RAW preview extraction")
|
||
|
||
from app.database import AsyncSessionLocal
|
||
from app.models import Photo
|
||
from app.config import settings
|
||
|
||
# Register HEIF opener with Pillow
|
||
register_heif_opener()
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
# Thumbnail sizes configuration
|
||
THUMB_SIZES = {
|
||
'small': settings.thumbnails.small,
|
||
'medium': settings.thumbnails.medium,
|
||
'large': settings.thumbnails.large
|
||
}
|
||
|
||
def get_thumb_path(photo_id: str, size: str, user_id: str = None) -> str:
|
||
"""Get the path for a thumbnail file.
|
||
|
||
When user_id is provided, thumbnails are stored under a user-specific
|
||
subdirectory to enforce isolation between users.
|
||
"""
|
||
if user_id:
|
||
thumb_dir = f"/data/thumbs/{user_id}/{photo_id}"
|
||
else:
|
||
thumb_dir = f"/data/thumbs/{photo_id}"
|
||
os.makedirs(thumb_dir, exist_ok=True)
|
||
return f"{thumb_dir}/{size}.{settings.thumbnails.format}"
|
||
|
||
def process_standard_image(filepath: str) -> Image.Image:
|
||
"""Process standard image formats (JPEG, PNG, etc.)"""
|
||
return Image.open(filepath)
|
||
|
||
def process_raw_image(filepath: str) -> Image.Image:
|
||
"""Process RAW image formats"""
|
||
if RAWPY_AVAILABLE:
|
||
try:
|
||
with rawpy.imread(filepath) as raw:
|
||
# Use half_size for faster processing
|
||
rgb = raw.postprocess(use_camera_wb=True, half_size=True)
|
||
# Convert numpy array to PIL Image
|
||
return Image.fromarray(rgb, 'RGB')
|
||
except Exception as e:
|
||
logger.error(f"Error processing RAW file {filepath}: {e}")
|
||
# Try to extract embedded JPEG preview
|
||
return extract_raw_preview(filepath)
|
||
else:
|
||
# Use exiftool to extract embedded preview
|
||
return extract_raw_preview(filepath)
|
||
|
||
def extract_raw_preview(filepath: str) -> Optional[Image.Image]:
|
||
"""Extract embedded JPEG preview from RAW file"""
|
||
try:
|
||
# Use exiftool to extract preview
|
||
import subprocess
|
||
import tempfile
|
||
|
||
with tempfile.NamedTemporaryFile(suffix='.jpg', delete=False) as tmp:
|
||
cmd = ['exiftool', '-b', '-PreviewImage', filepath]
|
||
result = subprocess.run(cmd, capture_output=True)
|
||
|
||
if result.returncode == 0 and result.stdout:
|
||
tmp.write(result.stdout)
|
||
tmp.flush()
|
||
return Image.open(tmp.name)
|
||
except Exception as e:
|
||
logger.error(f"Error extracting RAW preview from {filepath}: {e}")
|
||
|
||
return None
|
||
|
||
def process_heic_image(filepath: str) -> Image.Image:
|
||
"""Process HEIC/HEIF image formats"""
|
||
try:
|
||
# Use pillow-heif to open the image
|
||
img = Image.open(filepath)
|
||
# Convert to RGB if needed
|
||
if img.mode != 'RGB':
|
||
img = img.convert('RGB')
|
||
return img
|
||
except Exception as e:
|
||
logger.error(f"Error processing HEIC file {filepath}: {e}")
|
||
raise
|
||
|
||
def process_video_thumbnail(filepath: str) -> Image.Image:
|
||
"""Extract a still frame from a video file as a PIL Image."""
|
||
import tempfile
|
||
from io import BytesIO
|
||
|
||
tmp_path: Optional[str] = None
|
||
try:
|
||
# Find a usable seek timestamp. Some camera MOVs only expose
|
||
# duration at the format level, and stream 0 isn't always the
|
||
# video stream — search explicitly and fall back to the format
|
||
# duration, then to t=0 if neither is available.
|
||
probe = ffmpeg.probe(filepath)
|
||
duration: Optional[float] = None
|
||
for stream_info in probe.get('streams', []):
|
||
if stream_info.get('codec_type') != 'video':
|
||
continue
|
||
raw_duration = stream_info.get('duration')
|
||
if raw_duration is not None:
|
||
try:
|
||
duration = float(raw_duration)
|
||
break
|
||
except (TypeError, ValueError):
|
||
pass
|
||
if duration is None:
|
||
raw_duration = probe.get('format', {}).get('duration')
|
||
if raw_duration is not None:
|
||
try:
|
||
duration = float(raw_duration)
|
||
except (TypeError, ValueError):
|
||
duration = None
|
||
|
||
# Seek to 10% in for a representative frame; clamp very short
|
||
# clips to t=0 so we don't seek past the end.
|
||
timestamp = max(0.0, (duration or 0.0) * 0.1)
|
||
|
||
# NamedTemporaryFile creates the file on disk, so we MUST tell
|
||
# ffmpeg to overwrite it (otherwise it prompts on stdin and the
|
||
# call hangs/fails — which is why videos were getting the gray
|
||
# placeholder). We close the handle immediately and clean up
|
||
# in `finally` ourselves.
|
||
with tempfile.NamedTemporaryFile(suffix='.jpg', delete=False) as tmp:
|
||
tmp_path = tmp.name
|
||
|
||
stream = ffmpeg.input(filepath, ss=timestamp)
|
||
stream = ffmpeg.output(
|
||
stream,
|
||
tmp_path,
|
||
vframes=1,
|
||
format='image2',
|
||
vcodec='mjpeg',
|
||
)
|
||
ffmpeg.run(
|
||
stream,
|
||
capture_stdout=True,
|
||
capture_stderr=True,
|
||
overwrite_output=True,
|
||
)
|
||
|
||
# Load the frame fully into memory so we can delete the temp
|
||
# file immediately. Pillow's `Image.open` is lazy, which would
|
||
# otherwise leave the file dangling.
|
||
with open(tmp_path, 'rb') as fh:
|
||
data = fh.read()
|
||
if not data:
|
||
raise RuntimeError("ffmpeg produced an empty frame")
|
||
return Image.open(BytesIO(data)).copy()
|
||
except ffmpeg.Error as e:
|
||
stderr = (e.stderr or b'').decode('utf-8', errors='replace')
|
||
logger.error(
|
||
f"ffmpeg failed extracting video thumbnail from {filepath}: {stderr}"
|
||
)
|
||
return create_placeholder_thumbnail('video')
|
||
except Exception as e:
|
||
logger.error(f"Error extracting video thumbnail from {filepath}: {e}")
|
||
return create_placeholder_thumbnail('video')
|
||
finally:
|
||
if tmp_path and os.path.exists(tmp_path):
|
||
try:
|
||
os.unlink(tmp_path)
|
||
except OSError:
|
||
pass
|
||
|
||
def create_placeholder_thumbnail(media_type: str) -> Image.Image:
|
||
"""Create a placeholder thumbnail for failed processing"""
|
||
# Create a simple gray placeholder
|
||
img = Image.new('RGB', (640, 480), color=(128, 128, 128))
|
||
return img
|
||
|
||
def auto_rotate_image(image: Image.Image) -> Image.Image:
|
||
"""Auto-rotate image based on EXIF orientation"""
|
||
try:
|
||
# Get EXIF data
|
||
exif = image._getexif()
|
||
if exif:
|
||
orientation = exif.get(274) # Orientation tag
|
||
|
||
rotation_map = {
|
||
3: 180,
|
||
6: 270, # Note: PIL uses different rotation values than vips
|
||
8: 90
|
||
}
|
||
|
||
if orientation in rotation_map:
|
||
image = image.rotate(rotation_map[orientation], expand=True)
|
||
except (AttributeError, KeyError, TypeError):
|
||
pass # No orientation data available
|
||
|
||
return image
|
||
|
||
def generate_thumbnail(image: Image.Image, size: int, output_path: str):
|
||
"""Generate a thumbnail of the specified size.
|
||
|
||
Works on a copy so the caller's image is never mutated — this is
|
||
critical because the thumbnail loop iterates multiple sizes and
|
||
in-place shrinking would degrade later (larger) sizes.
|
||
"""
|
||
img = image.copy()
|
||
img.thumbnail((size, size), Image.Resampling.LANCZOS)
|
||
|
||
img.save(
|
||
output_path,
|
||
'WEBP',
|
||
quality=settings.thumbnails.quality,
|
||
method=4 # Balance between speed and compression
|
||
)
|
||
img.close()
|
||
|
||
@shared_task(bind=True, name='generate_thumbnails')
|
||
def generate_thumbnails(self, photo_id: str):
|
||
"""Generate thumbnails for a photo"""
|
||
return asyncio.run(_generate_thumbnails_async(photo_id, self))
|
||
|
||
async def _generate_thumbnails_async(photo_id: str, task):
|
||
"""Async implementation of thumbnail generation"""
|
||
async with AsyncSessionLocal() as session:
|
||
# Declared up front so the except block below can safely check it
|
||
# even if the initial SELECT raises (e.g. asyncpg transport error).
|
||
photo: Optional[Photo] = None
|
||
try:
|
||
# Get photo from database
|
||
result = await session.execute(
|
||
select(Photo).where(Photo.id == photo_id)
|
||
)
|
||
photo = result.scalar_one_or_none()
|
||
|
||
if not photo:
|
||
logger.error(f"Photo not found: {photo_id}")
|
||
return {'status': 'error', 'message': 'Photo not found'}
|
||
|
||
# Check if file exists
|
||
if not os.path.exists(photo.filepath):
|
||
logger.error(f"File not found: {photo.filepath}")
|
||
photo.processing_status = 'failed'
|
||
photo.processing_error = 'File not found'
|
||
await session.commit()
|
||
return {'status': 'error', 'message': 'File not found'}
|
||
|
||
# Update processing status
|
||
photo.processing_status = 'processing'
|
||
await session.commit()
|
||
|
||
# Load and process the image based on type
|
||
image = None
|
||
|
||
if photo.media_type == 'photo':
|
||
image = process_standard_image(photo.filepath)
|
||
elif photo.media_type == 'raw':
|
||
image = process_raw_image(photo.filepath)
|
||
elif photo.media_type == 'heic':
|
||
image = process_heic_image(photo.filepath)
|
||
elif photo.media_type == 'video':
|
||
image = process_video_thumbnail(photo.filepath)
|
||
else:
|
||
logger.error(f"Unsupported media type: {photo.media_type}")
|
||
image = create_placeholder_thumbnail(photo.media_type)
|
||
|
||
if not image:
|
||
raise Exception("Failed to process image")
|
||
|
||
# Auto-rotate based on EXIF
|
||
image = auto_rotate_image(image)
|
||
|
||
# Store original dimensions
|
||
photo.width = image.width
|
||
photo.height = image.height
|
||
|
||
# Perceptual hash from the original-resolution decoded frame.
|
||
# pHash is robust to resize/recompression but the thumbnail
|
||
# loop below mutates `image` in place, so this MUST run before
|
||
# the loop sees it. Failures are non-fatal — phash is a
|
||
# nice-to-have, not a blocker for thumbnail generation.
|
||
try:
|
||
import imagehash
|
||
photo.phash = str(imagehash.phash(image)) # 16-char hex
|
||
except Exception as e:
|
||
logger.warning(f"phash failed for {photo_id}: {e}")
|
||
photo.phash = None
|
||
|
||
# Generate thumbnails for each size
|
||
for size_name, size_value in THUMB_SIZES.items():
|
||
thumb_path = get_thumb_path(photo_id, size_name, photo.user_id)
|
||
generate_thumbnail(image, size_value, thumb_path)
|
||
|
||
# Update database with thumbnail path
|
||
setattr(photo, f'thumb_{size_name}', thumb_path)
|
||
|
||
# Update progress
|
||
task.update_state(
|
||
state='PROGRESS',
|
||
meta={'current_size': size_name, 'photo_id': photo_id}
|
||
)
|
||
|
||
# Update processing status
|
||
photo.processing_status = 'completed'
|
||
photo.processing_error = None
|
||
await session.commit()
|
||
|
||
logger.info(f"Thumbnails generated for photo {photo_id}")
|
||
|
||
# Dispatch vision pipeline only after thumbnails succeeded —
|
||
# vision tasks need the generated thumbnails to run inference.
|
||
try:
|
||
from app.tasks.vision import vision_fanout
|
||
vision_fanout.delay(photo_id)
|
||
except Exception as e:
|
||
logger.warning(f"Could not dispatch vision_fanout for {photo_id}: {e}")
|
||
|
||
return {'status': 'success', 'photo_id': photo_id}
|
||
|
||
except Exception as e:
|
||
logger.error(f"Error generating thumbnails for {photo_id}: {e}")
|
||
|
||
# Update error status. If the session is in a bad state (e.g.
|
||
# the original failure was a transport error) rollback first so
|
||
# the status write has a clean transaction to commit into.
|
||
try:
|
||
await session.rollback()
|
||
except Exception:
|
||
pass
|
||
|
||
if photo is not None:
|
||
try:
|
||
photo.processing_status = 'failed'
|
||
photo.processing_error = str(e)
|
||
await session.commit()
|
||
except Exception:
|
||
logger.exception(
|
||
f"Could not mark photo {photo_id} as failed"
|
||
)
|
||
|
||
return {'status': 'error', 'message': str(e)}
|
||
|
||
@shared_task(name='regenerate_all_thumbnails')
|
||
def regenerate_all_thumbnails():
|
||
"""Regenerate thumbnails for all photos"""
|
||
return asyncio.run(_regenerate_all_thumbnails_async())
|
||
|
||
async def _regenerate_all_thumbnails_async():
|
||
"""Async implementation of regenerating all thumbnails.
|
||
|
||
Queue order matters on first-boot and recovery runs: we dispatch
|
||
newest-first (by EXIF taken_at, fallback added_at) so the user's
|
||
most recent photos become fully-indexed before the 2012 archive even
|
||
starts. Picking up the library in pipeline order means the grid,
|
||
timeline and All Photos view populate top-down instead of the worker
|
||
chewing through random insertion-order rows while the UI still
|
||
shows grey placeholders.
|
||
"""
|
||
async with AsyncSessionLocal() as session:
|
||
# Get all photos that need thumbnails, newest first.
|
||
result = await session.execute(
|
||
select(Photo)
|
||
.where(Photo.processing_status.in_(['pending', 'failed']))
|
||
.order_by(
|
||
Photo.taken_at.desc().nullslast(),
|
||
Photo.added_at.desc().nullslast(),
|
||
)
|
||
)
|
||
photos = result.scalars().all()
|
||
|
||
logger.info(f"Regenerating thumbnails for {len(photos)} photos")
|
||
|
||
for photo in photos:
|
||
generate_thumbnails.delay(photo.id)
|
||
|
||
return {'status': 'queued', 'count': len(photos)}
|
||
|
||
|
||
# ── Perceptual hash backfill ────────────────────────────────────────────
|
||
#
|
||
# When phash was added post-launch, every existing photo has phash=NULL.
|
||
# This task fills them in by reading the existing thumb_large (the cheap
|
||
# option — pHash is robust to scale, and the thumb is already on local
|
||
# disk so we avoid re-decoding the original RAW/HEIC). Falls back to the
|
||
# original filepath if the thumb isn't available for some reason. Runs
|
||
# in batches to keep memory bounded and to give the user incremental
|
||
# progress visible in the worker logs.
|
||
|
||
@shared_task(name='backfill_phashes')
|
||
def backfill_phashes():
|
||
"""Compute and persist phash for every photo currently missing one."""
|
||
return asyncio.run(_backfill_phashes_async())
|
||
|
||
|
||
async def _backfill_phashes_async():
|
||
import imagehash
|
||
from PIL import Image as _PILImage
|
||
|
||
BATCH = 100
|
||
total_done = 0
|
||
total_failed = 0
|
||
|
||
async with AsyncSessionLocal() as session:
|
||
while True:
|
||
# Newest-first so the recent end of the library gets phashes
|
||
# (and therefore duplicate detection) ahead of the archive.
|
||
result = await session.execute(
|
||
select(Photo)
|
||
.where(Photo.phash.is_(None))
|
||
.where(Photo.processing_status == 'completed')
|
||
.order_by(
|
||
Photo.taken_at.desc().nullslast(),
|
||
Photo.added_at.desc().nullslast(),
|
||
)
|
||
.limit(BATCH)
|
||
)
|
||
batch = result.scalars().all()
|
||
if not batch:
|
||
break
|
||
|
||
for photo in batch:
|
||
source = photo.thumb_large or photo.filepath
|
||
try:
|
||
if not source or not os.path.exists(source):
|
||
photo.phash = None
|
||
total_failed += 1
|
||
continue
|
||
with _PILImage.open(source) as im:
|
||
photo.phash = str(imagehash.phash(im))
|
||
total_done += 1
|
||
except Exception as e:
|
||
logger.warning(f"phash backfill failed for {photo.id}: {e}")
|
||
total_failed += 1
|
||
|
||
await session.commit()
|
||
logger.info(
|
||
f"Backfilled phashes: {total_done} done, {total_failed} failed"
|
||
)
|
||
|
||
return {
|
||
'status': 'success',
|
||
'computed': total_done,
|
||
'failed': total_failed,
|
||
}
|
||
|
||
|
||
@shared_task(name='regroup_duplicates')
|
||
def regroup_duplicates_task():
|
||
"""Full recompute of duplicate groups (pHash + CLIP similarity).
|
||
|
||
Used by the Settings → Re-detect duplicates button."""
|
||
from app.services.duplicates import regroup_duplicates
|
||
return asyncio.run(regroup_duplicates())
|
||
|
||
|
||
@shared_task(name='incremental_regroup_duplicates')
|
||
def incremental_regroup_duplicates_task(since_iso: str | None = None):
|
||
"""Incremental duplicate detection for newly added photos.
|
||
|
||
Compares only photos added after `since_iso` against the full library
|
||
using CLIP vector similarity (O(new × log N) via HNSW) plus pHash.
|
||
Default post-scan path — much faster than a full regroup."""
|
||
from app.services.duplicates import incremental_regroup
|
||
from datetime import datetime, timezone
|
||
since = None
|
||
if since_iso:
|
||
since = datetime.fromisoformat(since_iso)
|
||
return asyncio.run(incremental_regroup(since=since)) |