wehub-resource-sync
358 lines
12 KiB
Python
358 lines
12 KiB
Python
"""
|
|
Ray Serve application: Video Embedding → Multi-Decoder.
|
|
|
|
Processes entire videos by chunking into segments.
|
|
Videos are downloaded from S3 to temp file, then processed locally (faster than streaming).
|
|
|
|
Encoder refs are passed directly to decoder; Ray Serve resolves dependencies automatically.
|
|
|
|
Usage:
|
|
serve run app:app
|
|
|
|
# With custom bucket:
|
|
S3_BUCKET=my-bucket serve run app:app
|
|
"""
|
|
|
|
import logging
|
|
import os
|
|
import tempfile
|
|
import time
|
|
from collections import defaultdict
|
|
from pathlib import Path
|
|
from urllib.parse import urlparse
|
|
|
|
import aioboto3
|
|
import numpy as np
|
|
from fastapi import FastAPI, HTTPException
|
|
from pydantic import BaseModel
|
|
from ray import serve
|
|
from ray.serve.handle import DeploymentResponse
|
|
|
|
from deployments.encoder import VideoEncoder
|
|
from deployments.decoder import MultiDecoder
|
|
from utils.video import chunk_video_async
|
|
from constants import DEFAULT_NUM_FRAMES, DEFAULT_CHUNK_DURATION, FFMPEG_THREADS, NUM_WORKERS
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def parse_s3_uri(s3_uri: str) -> tuple[str, str]:
|
|
"""Parse s3://bucket/key into (bucket, key)."""
|
|
parsed = urlparse(s3_uri)
|
|
if parsed.scheme != "s3":
|
|
raise ValueError(f"Invalid S3 URI: {s3_uri}")
|
|
bucket = parsed.netloc
|
|
key = parsed.path.lstrip("/")
|
|
return bucket, key
|
|
|
|
|
|
class AnalyzeRequest(BaseModel):
|
|
"""Request schema for /analyze endpoint."""
|
|
stream_id: str
|
|
video_path: str # S3 URI: s3://bucket/key
|
|
num_frames: int = DEFAULT_NUM_FRAMES
|
|
chunk_duration: float = DEFAULT_CHUNK_DURATION
|
|
use_batching: bool = False # Set False to compare unbatched performance
|
|
|
|
|
|
class TagResult(BaseModel):
|
|
text: str
|
|
score: float
|
|
|
|
|
|
class CaptionResult(BaseModel):
|
|
text: str
|
|
score: float
|
|
|
|
|
|
class TimingResult(BaseModel):
|
|
s3_download_ms: float
|
|
decode_video_ms: float
|
|
encode_ms: float
|
|
decode_ms: float
|
|
total_ms: float
|
|
|
|
|
|
class SceneChange(BaseModel):
|
|
"""Detected scene change event."""
|
|
timestamp: float # Seconds from video start
|
|
score: float # Scene change score (higher = bigger change)
|
|
chunk_index: int
|
|
frame_index: int # Frame index within chunk
|
|
|
|
|
|
class ChunkResult(BaseModel):
|
|
"""Result for a single chunk."""
|
|
chunk_index: int
|
|
start_time: float
|
|
duration: float
|
|
tags: list[TagResult]
|
|
retrieval_caption: CaptionResult
|
|
# Detected scene changes in this chunk
|
|
scene_changes: list[SceneChange]
|
|
|
|
|
|
class AnalyzeResponse(BaseModel):
|
|
"""Response schema for /analyze endpoint."""
|
|
stream_id: str
|
|
# Aggregated results (across all chunks)
|
|
tags: list[TagResult]
|
|
retrieval_caption: CaptionResult
|
|
# Scene change detection
|
|
scene_changes: list[SceneChange] # All detected scene changes
|
|
num_scene_changes: int
|
|
# Per-chunk results
|
|
chunks: list[ChunkResult]
|
|
num_chunks: int
|
|
video_duration: float
|
|
timing_ms: TimingResult
|
|
|
|
|
|
# FastAPI app
|
|
fastapi_app = FastAPI(
|
|
title="Video Embedding API",
|
|
description="GPU encoder → CPU multi-decoder using SigLIP embeddings",
|
|
)
|
|
|
|
|
|
@serve.deployment(
|
|
# setting this to twice that of the encoder. So that requests can complete the
|
|
# upfront CPU work and be queued for GPU processing.
|
|
num_replicas="auto",
|
|
ray_actor_options={"num_cpus": FFMPEG_THREADS},
|
|
max_ongoing_requests=4,
|
|
autoscaling_config={
|
|
"min_replicas": 2,
|
|
"max_replicas": 20,
|
|
"target_num_ongoing_requests": 2,
|
|
},
|
|
)
|
|
@serve.ingress(fastapi_app)
|
|
class VideoAnalyzer:
|
|
"""
|
|
Main ingress deployment that orchestrates VideoEncoder and MultiDecoder.
|
|
|
|
Encoder refs are passed directly to decoder; Ray Serve resolves dependencies.
|
|
Downloads video from S3 to temp file for fast local processing.
|
|
"""
|
|
|
|
def __init__(self, encoder: VideoEncoder, decoder: MultiDecoder):
|
|
self.encoder = encoder
|
|
self.decoder = decoder
|
|
self._s3_session = aioboto3.Session()
|
|
self._s3_client = None # Cached client for reuse across requests
|
|
logger.info("VideoAnalyzer ready")
|
|
|
|
async def _get_s3_client(self):
|
|
"""Get or create a reusable S3 client."""
|
|
if self._s3_client is None:
|
|
self._s3_client = await self._s3_session.client("s3").__aenter__()
|
|
return self._s3_client
|
|
|
|
async def _download_video(self, s3_uri: str) -> Path:
|
|
"""Download video from S3 to temp file. Returns local path."""
|
|
bucket, key = parse_s3_uri(s3_uri)
|
|
|
|
# Create temp file with video extension
|
|
suffix = Path(key).suffix or ".mp4"
|
|
temp_file = tempfile.NamedTemporaryFile(suffix=suffix, delete=False)
|
|
temp_path = Path(temp_file.name)
|
|
temp_file.close()
|
|
|
|
try:
|
|
s3 = await self._get_s3_client()
|
|
await s3.download_file(bucket, key, str(temp_path))
|
|
except Exception:
|
|
# Clean up temp file if download fails
|
|
temp_path.unlink(missing_ok=True)
|
|
raise
|
|
|
|
return temp_path
|
|
|
|
def _aggregate_results(
|
|
self,
|
|
chunk_results: list[dict],
|
|
top_k_tags: int = 5,
|
|
) -> dict:
|
|
"""
|
|
Aggregate results from multiple chunks.
|
|
|
|
Strategy:
|
|
- Tags: Average scores across chunks, return top-k
|
|
- Caption: Return the one with highest score across all chunks
|
|
"""
|
|
# Aggregate tag scores
|
|
tag_scores = defaultdict(list)
|
|
for result in chunk_results:
|
|
for tag in result["tags"]:
|
|
tag_scores[tag["text"]].append(tag["score"])
|
|
|
|
# Average tag scores and sort
|
|
aggregated_tags = [
|
|
{"text": text, "score": np.mean(scores)}
|
|
for text, scores in tag_scores.items()
|
|
]
|
|
aggregated_tags.sort(key=lambda x: x["score"], reverse=True)
|
|
top_tags = aggregated_tags[:top_k_tags]
|
|
|
|
# Best caption across all chunks
|
|
best_caption = max(
|
|
(r["retrieval_caption"] for r in chunk_results),
|
|
key=lambda x: x["score"],
|
|
)
|
|
|
|
return {
|
|
"tags": top_tags,
|
|
"retrieval_caption": best_caption,
|
|
}
|
|
|
|
def _encode_chunk(self, frames: np.ndarray, use_batching: bool = False) -> DeploymentResponse:
|
|
"""Encode a single chunk's frames to embeddings. Returns DeploymentResponse ref."""
|
|
return self.encoder.remote(frames, use_batching=use_batching)
|
|
|
|
async def _decode_chunk(
|
|
self,
|
|
encoder_output: dict,
|
|
chunk_index: int,
|
|
chunk_start_time: float,
|
|
chunk_duration: float,
|
|
ema_state=None,
|
|
) -> dict:
|
|
"""Decode embeddings to tags, caption, scene changes."""
|
|
return await self.decoder.remote(
|
|
encoder_output=encoder_output,
|
|
chunk_index=chunk_index,
|
|
chunk_start_time=chunk_start_time,
|
|
chunk_duration=chunk_duration,
|
|
ema_state=ema_state,
|
|
)
|
|
|
|
@fastapi_app.post("/analyze", response_model=AnalyzeResponse)
|
|
async def analyze(self, request: AnalyzeRequest) -> AnalyzeResponse:
|
|
"""
|
|
Analyze a video from S3 and return tags, caption, and scene changes.
|
|
|
|
Downloads video to temp file for fast local processing.
|
|
Chunks the entire video and aggregates results.
|
|
Encoder refs are passed directly to decoder for dependency resolution.
|
|
"""
|
|
total_start = time.perf_counter()
|
|
temp_path = None
|
|
|
|
try:
|
|
# Download video from S3 to temp file
|
|
download_start = time.perf_counter()
|
|
try:
|
|
temp_path = await self._download_video(request.video_path)
|
|
except Exception as e:
|
|
raise HTTPException(status_code=400, detail=f"Cannot download S3 video: {e}")
|
|
s3_download_ms = (time.perf_counter() - download_start) * 1000
|
|
|
|
# Chunk video with PARALLEL frame extraction from local file
|
|
decode_start = time.perf_counter()
|
|
try:
|
|
chunks = await chunk_video_async(
|
|
str(temp_path),
|
|
chunk_duration=request.chunk_duration,
|
|
num_frames_per_chunk=request.num_frames,
|
|
ffmpeg_threads=FFMPEG_THREADS,
|
|
use_single_ffmpeg=True,
|
|
)
|
|
except Exception as e:
|
|
raise HTTPException(status_code=400, detail=f"Cannot process video: {e}")
|
|
|
|
decode_video_ms = (time.perf_counter() - decode_start) * 1000
|
|
|
|
if not chunks:
|
|
raise HTTPException(status_code=400, detail="No chunks extracted from video")
|
|
|
|
# Calculate video duration from chunks
|
|
video_duration = chunks[-1].start_time + chunks[-1].duration
|
|
|
|
# Fire off all encoder calls (returns refs, not awaited)
|
|
encode_start = time.perf_counter()
|
|
encode_refs = [
|
|
self._encode_chunk(chunk.frames, use_batching=request.use_batching)
|
|
for chunk in chunks
|
|
]
|
|
encode_ms = (time.perf_counter() - encode_start) * 1000
|
|
|
|
# Decode chunks SERIALLY, passing encoder refs directly.
|
|
# Ray Serve resolves the encoder result when decoder needs it.
|
|
# EMA state is tracked here (not in decoder) to ensure continuity
|
|
# even when autoscaling routes requests to different replicas.
|
|
decode_start = time.perf_counter()
|
|
decode_results = []
|
|
ema_state = None # Will be initialized from first chunk's first frame
|
|
for chunk, enc_ref in zip(chunks, encode_refs):
|
|
dec_result = await self._decode_chunk(
|
|
encoder_output=enc_ref,
|
|
chunk_index=chunk.index,
|
|
chunk_start_time=chunk.start_time,
|
|
chunk_duration=chunk.duration,
|
|
ema_state=ema_state,
|
|
)
|
|
decode_results.append(dec_result)
|
|
ema_state = dec_result["ema_state"] # Carry forward for next chunk
|
|
decode_ms = (time.perf_counter() - decode_start) * 1000
|
|
|
|
# Collect results
|
|
chunk_results = []
|
|
per_chunk_results = []
|
|
all_scene_changes = []
|
|
|
|
for chunk, decoder_result in zip(chunks, decode_results):
|
|
chunk_results.append(decoder_result)
|
|
|
|
# Scene changes come directly from decoder
|
|
chunk_scene_changes = [
|
|
SceneChange(**sc) for sc in decoder_result["scene_changes"]
|
|
]
|
|
all_scene_changes.extend(chunk_scene_changes)
|
|
|
|
per_chunk_results.append(ChunkResult(
|
|
chunk_index=chunk.index,
|
|
start_time=chunk.start_time,
|
|
duration=chunk.duration,
|
|
tags=[TagResult(**t) for t in decoder_result["tags"]],
|
|
retrieval_caption=CaptionResult(**decoder_result["retrieval_caption"]),
|
|
scene_changes=chunk_scene_changes,
|
|
))
|
|
|
|
# Aggregate results
|
|
aggregated = self._aggregate_results(chunk_results)
|
|
|
|
total_ms = (time.perf_counter() - total_start) * 1000
|
|
|
|
return AnalyzeResponse(
|
|
stream_id=request.stream_id,
|
|
tags=[TagResult(**t) for t in aggregated["tags"]],
|
|
retrieval_caption=CaptionResult(**aggregated["retrieval_caption"]),
|
|
scene_changes=all_scene_changes,
|
|
num_scene_changes=len(all_scene_changes),
|
|
chunks=per_chunk_results,
|
|
num_chunks=len(chunks),
|
|
video_duration=video_duration,
|
|
timing_ms=TimingResult(
|
|
s3_download_ms=round(s3_download_ms, 2),
|
|
decode_video_ms=round(decode_video_ms, 2),
|
|
encode_ms=round(encode_ms, 2),
|
|
decode_ms=round(decode_ms, 2),
|
|
total_ms=round(total_ms, 2),
|
|
),
|
|
)
|
|
finally:
|
|
# Clean up temp file
|
|
if temp_path and temp_path.exists():
|
|
temp_path.unlink(missing_ok=True)
|
|
|
|
@fastapi_app.get("/health")
|
|
async def health(self):
|
|
"""Health check endpoint."""
|
|
return {"status": "healthy"}
|
|
|
|
|
|
encoder = VideoEncoder.bind()
|
|
decoder = MultiDecoder.bind(bucket=os.environ.get("S3_BUCKET"))
|
|
app = VideoAnalyzer.bind(encoder=encoder, decoder=decoder)
|