# SPDX-FileCopyrightText: Copyright (c) 2025-2026 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 argparse import asyncio import logging import os import signal import sys from typing import AsyncIterator, Tuple import torch import uvloop from transformers import AutoProcessor, Qwen2AudioForConditionalGeneration from vllm.engine.arg_utils import AsyncEngineArgs from vllm.utils.argparse_utils import FlexibleArgumentParser import dynamo.nixl_connect as connect from dynamo.runtime import Client, DistributedRuntime, dynamo_worker from dynamo.runtime.logging import configure_dynamo_logging sys.path.append(os.path.join(os.path.dirname(os.path.abspath(__file__)), "..")) from utils.args import Config, base_parse_args, parse_endpoint from utils.audio_loader import AudioLoader from utils.protocol import MyRequestOutput, vLLMMultimodalRequest configure_dynamo_logging() logger = logging.getLogger(__name__) try: import cupy as array_module if not array_module.cuda.is_available(): raise ImportError("CUDA is not available.") DEVICE = "cuda" logger.info("Using cupy for array operations (GPU mode).") except ImportError as e: logger.warning(f"Failed to import cupy, falling back to numpy: {e}.") import numpy as array_module DEVICE = "cpu" CACHE_SIZE_MAXIMUM = 8 class VllmEncodeWorker: def __init__( self, args: argparse.Namespace, engine_args: AsyncEngineArgs, pd_worker_client: Client, ) -> None: self.pd_worker_client = pd_worker_client self.engine_args = engine_args self.model = self.engine_args.model self.audio_loader = AudioLoader(cache_size=CACHE_SIZE_MAXIMUM) self.audio_processor = AutoProcessor.from_pretrained( self.model, trust_remote_code=True ) self.audio_model = Qwen2AudioForConditionalGeneration.from_pretrained( self.model, device_map="auto", dtype=torch.float16 ).eval() def get_audio_embeddings(self, audio_features): input_features, feature_attention_mask = ( audio_features.input_features, audio_features.feature_attention_mask, ) with torch.no_grad(): ( audio_feat_lengths, audio_output_lengths, ) = self.audio_model.audio_tower._get_feat_extract_output_lengths( feature_attention_mask.sum(-1) ) batch_size, _, max_mel_seq_len = input_features.shape max_seq_len = (max_mel_seq_len - 2) // 2 + 1 # Create a sequence tensor of shape (batch_size, max_seq_len) seq_range = ( torch.arange( 0, max_seq_len, dtype=audio_feat_lengths.dtype, device=audio_feat_lengths.device, ) .unsqueeze(0) .expand(batch_size, max_seq_len) ) lengths_expand = audio_feat_lengths.unsqueeze(1).expand( batch_size, max_seq_len ) # Create mask padding_mask = seq_range >= lengths_expand audio_attention_mask_ = padding_mask.view( batch_size, 1, 1, max_seq_len ).expand(batch_size, 1, max_seq_len, max_seq_len) audio_attention_mask = audio_attention_mask_.to( dtype=self.audio_model.audio_tower.conv1.weight.dtype, device=self.audio_model.audio_tower.conv1.weight.device, ) audio_attention_mask[audio_attention_mask_] = float("-inf") audio_outputs = self.audio_model.audio_tower( input_features, attention_mask=audio_attention_mask ) selected_audio_feature = audio_outputs.last_hidden_state audio_features = self.audio_model.multi_modal_projector( selected_audio_feature ) num_audios, max_audio_tokens, embed_dim = audio_features.shape audio_features_mask = torch.arange( max_audio_tokens, device=audio_output_lengths.device )[None, :] audio_features_mask = audio_features_mask < audio_output_lengths[:, None] audio_features = audio_features[audio_features_mask] return audio_features def cleanup(self): pass async def generate( self, request: vLLMMultimodalRequest ) -> AsyncIterator[MyRequestOutput]: logger.debug(f"Got raw request: {request}") if not isinstance(request, vLLMMultimodalRequest): if isinstance(request, str): request = vLLMMultimodalRequest.model_validate_json(request) else: request = vLLMMultimodalRequest.model_validate(request) logger.debug(f"Received encode request: {{ id: {request.request_id} }}.") request_id = request.request_id # The following steps encode the requested audio and provided useful embeddings. # 1. Open the audio from the provided URL. # 2. Process the audio using the audio processor. # 3. Run the audio through the audio model's audio tower. # 4. Run the results of the audio tower through the multi-modal projector. # 5. Create a descriptor for the embeddings. # 6. Create a write operation using the serialized request and the descriptor. # 7. Await for the write operation to complete. # 8. Yield the encode response. try: audio, sr = await self.audio_loader.load_audio( request.multimodal_input.audio_url ) audio_features = self.audio_processor( text="test<|AUDIO|>", audio=audio, return_tensors="pt", padding=False ) with torch.no_grad(): audio_embeddings = self.get_audio_embeddings(audio_features) descriptor = connect.Descriptor(audio_embeddings) with await self._connector.create_readable(descriptor) as readable: request.serialized_request = readable.metadata() # Clear the audio URL as hint that the audio is passed as embeddings. request.multimodal_input.audio_url = None request.embeddings_shape = tuple(audio_embeddings.shape) logger.debug(f"Request: {request.model_dump_json()}") response_generator = await self.pd_worker_client.round_robin( request.model_dump_json() ) await readable.wait_for_completion() async for response in response_generator: output = MyRequestOutput.model_validate_json(response.data()) yield MyRequestOutput( request_id=output.request_id, prompt=output.prompt, prompt_token_ids=output.prompt_token_ids, prompt_logprobs=output.prompt_logprobs, outputs=output.outputs, finished=output.finished, ).model_dump_json() except Exception as e: logger.error(f"Error processing request {request_id}: {e}") raise async def async_init(self, runtime: DistributedRuntime): logger.info("Startup started.") # Create and initialize a dynamo connector for this worker. # We'll needs this to move data between this worker and remote workers efficiently. self._connector = connect.Connector() logger.info("Startup completed.") @classmethod def parse_args(cls) -> Tuple[argparse.Namespace, Config]: DYN_NAMESPACE = os.environ.get("DYN_NAMESPACE", "dynamo") DEFAULT_ENDPOINT = f"dyn://{DYN_NAMESPACE}.encoder.generate" DEFAULT_DOWNSTREAM_ENDPOINT = f"dyn://{DYN_NAMESPACE}.llm.generate" parser = FlexibleArgumentParser( description="vLLM based encoder for Dynamo LLM." ) parser.add_argument( "--endpoint", type=str, default=DEFAULT_ENDPOINT, help=f"Dynamo endpoint string in 'dyn://namespace.component.endpoint' format. Default: '{DEFAULT_ENDPOINT}'", ) parser.add_argument( "--downstream-endpoint", type=str, default=DEFAULT_DOWNSTREAM_ENDPOINT, help=f"The endpoint string of the downstream LLM in 'dyn://namespace.component.endpoint' format. Default: '{DEFAULT_DOWNSTREAM_ENDPOINT}'", ) args, config = base_parse_args(parser) return args, config async def graceful_shutdown(runtime): """ By calling `runtime.shutdown()`, the endpoints will immediately be unavailable. However, in-flight requests will still be processed until they are finished. After all in-flight requests are finished, the `serve_endpoint` functions will return and the engine will be shutdown by Python's garbage collector. """ logging.info("Received shutdown signal, shutting down DistributedRuntime") runtime.shutdown() logging.info("DistributedRuntime shutdown complete") @dynamo_worker() async def worker(runtime: DistributedRuntime): # Runtime setup # Set up signal handler for graceful shutdown loop = asyncio.get_running_loop() def signal_handler(): asyncio.create_task(graceful_shutdown(runtime)) for sig in (signal.SIGTERM, signal.SIGINT): loop.add_signal_handler(sig, signal_handler) logging.info("Signal handlers set up for graceful shutdown") # worker setup args, config = VllmEncodeWorker.parse_args() await init(runtime, args, config) async def init(runtime: DistributedRuntime, args: argparse.Namespace, config: Config): """ Instantiate and serve """ component = runtime.namespace(config.namespace).component(config.component) generate_endpoint = component.endpoint(config.endpoint) parsed_namespace, parsed_component_name, parsed_endpoint_name = parse_endpoint( args.downstream_endpoint ) pd_worker_client = ( await runtime.namespace(parsed_namespace) .component(parsed_component_name) .endpoint(parsed_endpoint_name) .client() ) handler = VllmEncodeWorker(args, config.engine_args, pd_worker_client) await handler.async_init(runtime) logger.info("Waiting for PD Worker Instances ...") await pd_worker_client.wait_for_instances() logger.info(f"Starting to serve the {args.endpoint} endpoint...") try: await asyncio.gather( generate_endpoint.serve_endpoint( handler.generate, metrics_labels=[("model", config.model)] ), ) except Exception as e: logger.error(f"Failed to serve endpoints: {e}") raise finally: handler.cleanup() if __name__ == "__main__": uvloop.install() asyncio.run(worker())