# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import asyncio import base64 import binascii import json import logging import os from io import BytesIO from queue import Queue from typing import AsyncIterator, Optional from urllib.parse import urlparse import av import connect import httpx import numpy as np import torch import torch.nn.functional as F from utils.protocol import EncodeRequest from utils.vllm import parse_vllm_args from dynamo.sdk import async_on_start, endpoint, service logger = logging.getLogger(__name__) CACHE_SIZE_MAXIMUM = 8 @service( dynamo={ "namespace": "dynamo", }, resources={"gpu": 1, "cpu": "10", "memory": "20Gi"}, workers=1, ) class VllmEncodeWorker: def __init__(self) -> None: class_name = self.__class__.__name__ self.engine_args = parse_vllm_args(class_name, "") self.num_frames_to_sample = getattr(self.engine_args, "num_sampled_frames", 8) self.frame_height = getattr(self.engine_args, "frame_height", 336) self.frame_width = getattr(self.engine_args, "frame_width", 336) self.frame_channels = getattr(self.engine_args, "frame_channels", 3) self.dummy_token_id = getattr(self.engine_args, "dummy_token_id", 0) self.video_token_id = getattr(self.engine_args, "video_token_id", 32000) self.dummy_tokens_per_frame = getattr( self.engine_args, "dummy_tokens_per_frame", 144 ) self._video_content_cache: dict[str, BytesIO] = {} self._cache_queue: Queue[str] = Queue(maxsize=CACHE_SIZE_MAXIMUM) self._http_client: Optional[httpx.AsyncClient] = None self._http_timeout = 60.0 async def _read_video_pyav( self, container: av.container.InputContainer, indices: np.ndarray ) -> np.ndarray: """ Decode the video with PyAV decoder. Async wrapper. """ def blocking_decode(): container.seek(0) # Reset container for decoding processed_indices = set(indices) # Determine min/max index to optimize decoding loop slightly min_idx = 0 max_idx = -1 if len(indices) > 0: min_idx = np.min(indices) max_idx = np.max(indices) if ( not processed_indices and container.streams.video and container.streams.video[0].frames > 0 ): logger.warning( "_read_video_pyav called with empty indices for a non-empty video, attempting to read first frame." ) try: frame = next(container.decode(video=0)) return np.stack([frame.to_ndarray(format="rgb24")]) except StopIteration: logger.error( "Failed to read even the first frame despite non-empty indices check." ) return np.array([]) decoded_frames_list = [] for i, frame in enumerate(container.decode(video=0)): if i > max_idx and max_idx != -1: # max_idx is -1 if indices is empty break if i >= min_idx and i in processed_indices: decoded_frames_list.append(frame) if not decoded_frames_list and len(processed_indices) > 0: actual_decoded_count = 0 try: container.seek(0) # Reset for counting for _ in container.decode(video=0): actual_decoded_count += 1 except Exception: # Handle cases where re-decoding/counting fails pass # Keep original error message raise ValueError( f"Could not decode any frames for the given indices: {indices.tolist()}. " f"Video might be shorter than expected or indices out of bounds. " f"Actual decodable frames in container (approx): {actual_decoded_count}." ) return ( np.stack([x.to_ndarray(format="rgb24") for x in decoded_frames_list]) if decoded_frames_list else np.array([]) ) return await asyncio.to_thread(blocking_decode) async def _load_video_content(self, video_url: str) -> BytesIO: parsed_url = urlparse(video_url) video_url_lower = video_url.lower() if parsed_url.scheme in ("http", "https"): if video_url_lower in self._video_content_cache: logger.info(f"Video content found in cache for URL: {video_url}") cached_content = self._video_content_cache[video_url_lower] cached_content.seek(0) return cached_content try: video_data: BytesIO if parsed_url.scheme == "data": if not parsed_url.path.startswith( ("video/", "application/octet-stream") ): raise ValueError("Data URL must be a video type or octet-stream") media_type_and_data = parsed_url.path.split(",", 1) if len(media_type_and_data) != 2: raise ValueError("Invalid Data URL format: missing comma separator") media_type, data_segment = media_type_and_data if ";base64" not in media_type: raise ValueError("Video Data URL currently must be base64 encoded") try: video_bytes = base64.b64decode(data_segment) video_data = BytesIO(video_bytes) except binascii.Error as e: raise ValueError( f"Invalid base64 encoding for video data: {e}" ) from e elif parsed_url.scheme in ("http", "https"): if not self._http_client: await self._init_http_client() if not self._http_client: # Double check after initialization raise RuntimeError("Failed to initialize HTTP client") logger.info(f"Downloading video from URL: {video_url}") response = await self._http_client.get( video_url, timeout=self._http_timeout ) response.raise_for_status() if not response.content: raise ValueError( f"Empty response content from video URL: {video_url}" ) video_data = BytesIO(response.content) video_data.seek(0) logger.info( f"Video downloaded from {video_url}, size: {len(response.content)} bytes." ) elif parsed_url.scheme == "file" or not parsed_url.scheme: file_path = parsed_url.path if parsed_url.scheme else video_url # Ensure path is absolute or resolve relative to a known base if necessary # For simplicity, assuming it's an accessible path. if not os.path.exists(file_path): raise FileNotFoundError(f"Error reading file: {file_path}") with open(file_path, "rb") as f: video_bytes = f.read() video_data = BytesIO(video_bytes) else: raise ValueError( f"Unsupported video source scheme: {parsed_url.scheme} for URL {video_url}" ) if parsed_url.scheme in ( "http", "https", ): # Cache successfully downloaded content if self._cache_queue.full(): oldest_url = self._cache_queue.get_nowait() if oldest_url in self._video_content_cache: del self._video_content_cache[oldest_url] # Store the BytesIO object directly; it will be seek(0)'d when retrieved self._video_content_cache[video_url_lower] = video_data self._cache_queue.put(video_url_lower) return video_data except httpx.HTTPStatusError as e: logger.error( f"HTTP error {e.response.status_code} loading video {video_url}: {e.response.text[:200]}" ) raise ValueError( f"Failed to download video {video_url}: HTTP {e.response.status_code}" ) from e except httpx.RequestError as e: logger.error(f"Request error loading video {video_url}: {e}") raise ValueError(f"Network request failed for video {video_url}") from e except FileNotFoundError as e: logger.error(f"File error loading video {video_url}: {e}") raise except Exception as e: logger.error( f"Error loading video content from {video_url}: {type(e).__name__} - {e}" ) raise ValueError(f"Failed to load video content: {e}") from e @endpoint() async def encode(self, request: EncodeRequest) -> AsyncIterator[str]: request_id = request.request_id video_url = getattr(request, "video_url", None) if not video_url: logger.error(f"Request {request_id}: 'video_url' not provided.") raise ValueError("'video_url' is required for encoding.") if request.serialized_request is None: logger.error( f"Request serialized_request is None for request: {{ id: {request_id} }}." ) raise ValueError("'serialized_request' is required for encoding.") logger.info( f"Received encode request: {{ id: {request_id}, video_url: '{video_url[:100]}...' }}" ) container: Optional[av.container.InputContainer] = None try: video_content_stream = await self._load_video_content(video_url) def open_video_container_sync(): try: return av.open(video_content_stream, mode="r") except av.FFmpegError as ave: logger.error( f"PyAV error opening video stream from {video_url}: {ave}" ) raise ValueError( f"Invalid video format or corrupted data from {video_url}." ) from ave except Exception as e: logger.error( f"Unexpected error opening video stream from {video_url} with PyAV: {e}" ) raise ValueError( f"Unexpected error opening video from {video_url}." ) from e container = await asyncio.to_thread(open_video_container_sync) if not container or not container.streams.video: logger.error(f"No video stream found in {video_url}.") raise ValueError(f"No video stream in {video_url}.") stream_info = container.streams.video[0] total_frames = stream_info.frames # Duration can be useful for streams where total_frames is 0 if stream_info.duration and stream_info.time_base: duration_sec = float(stream_info.duration * stream_info.time_base) else: duration_sec = 0 if total_frames == 0 and duration_sec == 0: logger.error(f"Video file '{video_url}' has 0 frames and 0 duration.") raise ValueError(f"Video {video_url} has 0 frames and 0 duration.") if total_frames == 0 and duration_sec > 0: logger.warning( f"Video {video_url} reports 0 frames but has duration {duration_sec:.2f}s. Frame sampling may be based on requested count directly." ) logger.debug( f"Video {video_url} has {total_frames} frames (duration: {duration_sec:.2f}s). Sampling {self.num_frames_to_sample} frames." ) indices: np.ndarray if total_frames > 0: if total_frames < self.num_frames_to_sample: logger.warning( f"Video frames ({total_frames}) < samples ({self.num_frames_to_sample}). Using all {total_frames} available frames." ) indices = np.arange(0, total_frames).astype(int) else: indices = np.linspace( 0, total_frames - 1, self.num_frames_to_sample, dtype=int ) else: # total_frames is 0 (likely a stream), sample by count. logger.warning( f"Video {video_url} frame count is 0. Attempting to sample {self.num_frames_to_sample} frames by index. This might fail if stream is too short." ) indices = np.arange(0, self.num_frames_to_sample).astype(int) # Ensure indices are unique, especially after linspace for small numbers. indices = np.unique(indices) if ( len(indices) == 0 and total_frames > 0 ): # Safety for linspace oddities with few frames # If unique resulted in empty but there are frames, sample at least one or up to num_frames_to_sample actual_samples = min(self.num_frames_to_sample, total_frames) indices = np.arange(0, actual_samples).astype(int) elif len(indices) == 0 and total_frames == 0: # If indices is empty and total_frames is 0, this means num_frames_to_sample might be 0 or indices logic failed. # This case implies we might not be able to sample any frames. # _read_video_pyav handles empty indices with non-empty video by trying to read the first frame. # If indices is empty here due to num_frames_to_sample=0, _read_video_pyav will return empty. pass # Let _read_video_pyav handle this. logger.info(f"Selected frame indices for {video_url}: {indices.tolist()}") if not container: raise ValueError(f"Container is None for {video_url}") clip_np: np.ndarray = await self._read_video_pyav(container, indices) if clip_np.size == 0: raise ValueError( f"Failed to extract any video frames from {video_url} for indices {indices.tolist()}. Clip is empty." ) logger.info( f"Successfully extracted {len(clip_np) if clip_np.ndim > 1 and clip_np.shape[0] > 0 else 0} frames for {video_url} with original shape {clip_np.shape}." ) # Convert the NumPy array from the video decoder into a PyTorch tensor. # This is a required step to use PyTorch functions for GPU-accelerated image processing. frames_tensor_orig_res = torch.from_numpy(clip_np) # Shape: (T, H, W, C) # Permute to (T, C, H, W) for interpolate frames_tensor_chw = frames_tensor_orig_res.permute( 0, 3, 1, 2 ).float() # Ensure float for interpolate # Resize resized_frames_tensor_chw = F.interpolate( frames_tensor_chw, size=(self.frame_height, self.frame_width), mode="bilinear", align_corners=False, ) # Permute back to (T, H_new, W_new, C) resized_frames_tensor_hwc = resized_frames_tensor_chw.permute(0, 2, 3, 1) logger.debug(f"Resized frames to shape: {resized_frames_tensor_hwc.shape}") # Ensure the tensor is contiguous, on CUDA and uint8 for the NIXL buffer. tensor_for_descriptor: torch.Tensor = resized_frames_tensor_hwc.to( device="cpu", dtype=torch.uint8 ).contiguous() logger.info( f"Req {request_id}: Preparing raw frames tensor (shape: {tensor_for_descriptor.shape}, " f"dtype: {tensor_for_descriptor.dtype}, device: {tensor_for_descriptor.device}, " f"contiguous: {tensor_for_descriptor.is_contiguous()}) for RDMA." ) # Create a descriptor for the tensor to be sent via the connector. descriptor = connect.Descriptor(tensor_for_descriptor) logger.info(f"Req {request_id}: Beginning connector write operation.") # Pass the remote worker's SerializedRequest (representing its WritableOperation) to begin_write. # This initiates the data transfer to the memory buffer on the other worker. write_op = await self._connector.begin_write( descriptor, request.serialized_request ) # Wait for the RDMA/transfer operation to complete. await write_op.wait_for_completion() logger.info(f"Req {request_id}: Connector write operation completed.") # Yield a simplified response, assuming EncodeResponse in protocol is adapted final_response_data = { "request_id": request.request_id, } yield json.dumps(final_response_data) logger.info(f"Encode request {request_id} processed successfully.") except ( FileNotFoundError, av.FFmpegError, ValueError, ) as e: logger.error( f"Error processing request {request_id} ({video_url[:100]}...): {type(e).__name__} - {e}" ) raise # Re-raise to be handled by the service framework except Exception as e: logger.exception( f"Unexpected error processing request {request_id} ({video_url[:100]}...): {e}" ) raise finally: if container: await asyncio.to_thread(container.close) async def _init_http_client(self): if ( not self._http_client or self._http_client.is_closed ): # Check if closed as well self._http_client = httpx.AsyncClient(timeout=self._http_timeout) logger.info("HTTP client (re)initialized.") @async_on_start async def async_init(self): logger.info(f"{self.__class__.__name__} async_init started.") # Initialize the connector for RDMA transfers. self._connector = connect.Connector() await self._connector.initialize() logger.info("Dynamo connector initialized.") await self._init_http_client() logger.info( f"{self.__class__.__name__} async_init completed. Ready to encode video frames." )