chore: import upstream snapshot with attribution

doc/source/serve/tutorials/video-analysis/app.py

@@ -0,0 +1,357 @@
+"""
+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)