From db952b717c34910014bcfe2f01d81e8c207ef2ac Mon Sep 17 00:00:00 2001 From: Jie Gong Date: Fri, 26 Jun 2026 13:07:39 -0600 Subject: [PATCH 1/3] Add Parakeet TDT 0.6B V3 dynamic-batching example NeMo-based ASR truss for nvidia/parakeet-tdt-0.6b-v3 with server-side dynamic batching (collector + inference pipeline, adaptive collect window, CUDA graphs) for high-throughput HTTP transcription. Batching counterpart to the existing nvidia/parakeet-tdt-0_6b-v2 example. Co-authored-by: Cursor --- .../parakeet-tdt-0_6b-v3-batching/README.md | 60 +++ .../parakeet-tdt-0_6b-v3-batching/config.yaml | 47 ++ .../model/__init__.py | 0 .../model/model.py | 483 ++++++++++++++++++ 4 files changed, 590 insertions(+) create mode 100644 nvidia/parakeet-tdt-0_6b-v3-batching/README.md create mode 100644 nvidia/parakeet-tdt-0_6b-v3-batching/config.yaml create mode 100644 nvidia/parakeet-tdt-0_6b-v3-batching/model/__init__.py create mode 100644 nvidia/parakeet-tdt-0_6b-v3-batching/model/model.py diff --git a/nvidia/parakeet-tdt-0_6b-v3-batching/README.md b/nvidia/parakeet-tdt-0_6b-v3-batching/README.md new file mode 100644 index 000000000..5fb05d5fb --- /dev/null +++ b/nvidia/parakeet-tdt-0_6b-v3-batching/README.md @@ -0,0 +1,60 @@ +# Parakeet TDT 0.6B V3 — dynamic batching + +A [Truss](https://truss.baseten.co/) for NVIDIA's +[`parakeet-tdt-0.6b-v3`](https://huggingface.co/nvidia/parakeet-tdt-0.6b-v3) English +ASR model, with **server-side dynamic batching** for high-throughput HTTP +transcription. It runs the NeMo model directly (fp16 on Turing/T4, bf16 on Ampere+) +and batches concurrent requests on the GPU. + +This is the batching counterpart to [`../parakeet-tdt-0_6b-v2`](../parakeet-tdt-0_6b-v2) +(single-request). + +## How batching works + +Incoming requests are decoded to audio on the request thread, then handed to a small +pipeline: + +- a **collector** thread drains the request queue (with an adaptive collect window), + pads the batch, and runs the mel preprocessor on the GPU; +- an **inference** thread runs the encoder + TDT decoder (with CUDA graphs) on the + previous batch while the collector prepares the next one. + +The adaptive collect window only waits for stragglers when 2+ requests are already +queued, so single-arrival tail latency is unchanged while concurrent load grows +batches above 1. + +## Configuration (`config.yaml` → `environment_variables`) + +| var | default | meaning | +| --- | --- | --- | +| `BATCH_INFERENCE` | `true` | enable the batched pipeline (`false` = serial) | +| `MAX_BATCH_SIZE` | `64` | max requests per GPU batch | +| `BATCH_COLLECT_MS` | `20` | adaptive collect window (ms) for stragglers | +| `CUDA_GRAPHS` | `true` | CUDA graphs for the TDT decoder loop | +| `TORCH_COMPILE` | `false` | `torch.compile` the encoder | + +The model weights are pulled from Hugging Face into a Baseten volume cache +(`model_cache` in `config.yaml`) for fast cold starts; `nvidia/parakeet-tdt-0.6b-v3` +is public, so no token is required. Default accelerator is `L4`. + +## Deploy + +```bash +truss push +``` + +## Call + +```bash +curl -X POST https://model-.api.baseten.co/environments/production/predict \ + -H "Authorization: Api-Key $BASETEN_API_KEY" \ + -d '{"audio_url": "https://dldata-public.s3.us-east-2.amazonaws.com/2086-149220-0033.wav", "timestamps": false}' +``` + +Input fields: + +- `audio_url` — URL fetched and decoded server-side, **or** +- `audio_base64` — base64-encoded audio bytes (wav/mp3/opus/… via ffmpeg) +- `timestamps` — `true` returns word-level timestamps (serial path; not batched) + +Response: `{"transcript": "..."}`. diff --git a/nvidia/parakeet-tdt-0_6b-v3-batching/config.yaml b/nvidia/parakeet-tdt-0_6b-v3-batching/config.yaml new file mode 100644 index 000000000..fcc36ccea --- /dev/null +++ b/nvidia/parakeet-tdt-0_6b-v3-batching/config.yaml @@ -0,0 +1,47 @@ +model_name: Parakeet TDT 0.6B V3 (Batching) +description: > + NVIDIA Parakeet TDT 0.6B V3 English ASR with server-side dynamic batching + for high-throughput HTTP transcription. +python_version: py312 +model_metadata: + repo_id: nvidia/parakeet-tdt-0.6b-v3 + avatar_url: https://cdn-avatars.huggingface.co/v1/production/uploads/1613114437487-60262a8e0703121c822a80b6.png + example_model_input: + { + "audio_url": "https://dldata-public.s3.us-east-2.amazonaws.com/2086-149220-0033.wav", + "timestamps": false, + } +system_packages: + - ffmpeg +resources: + accelerator: L4 + use_gpu: true +runtime: + predict_concurrency: 128 +secrets: + hf_access_token: null +model_cache: + - repo_id: nvidia/parakeet-tdt-0.6b-v3 + revision: main + use_volume: true + volume_folder: parakeet-tdt-0.6b-v3 + kind: hf +requirements: + - nemo_toolkit[asr] + - torch>=2.6.0,<2.8.0 + - torchaudio>=2.6.0,<2.8.0 + - requests + - soundfile + - numpy + - numpy<=2.3 + - pyarrow==20.0.0 + - cuda-python +environment_variables: + BATCH_INFERENCE: "true" + MAX_BATCH_SIZE: "64" + # Adaptive collect window: only waits when 2+ items already drained, + # so single-arrival tail latency is unchanged. 20ms picks up + # stragglers under concurrent load to grow batches above 1. + BATCH_COLLECT_MS: "20" + CUDA_GRAPHS: "true" + TORCH_COMPILE: "false" diff --git a/nvidia/parakeet-tdt-0_6b-v3-batching/model/__init__.py b/nvidia/parakeet-tdt-0_6b-v3-batching/model/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/nvidia/parakeet-tdt-0_6b-v3-batching/model/model.py b/nvidia/parakeet-tdt-0_6b-v3-batching/model/model.py new file mode 100644 index 000000000..a1024e9ba --- /dev/null +++ b/nvidia/parakeet-tdt-0_6b-v3-batching/model/model.py @@ -0,0 +1,483 @@ +""" +Parakeet TDT 0.6B batching endpoint (NeMo PyTorch fp16 on T4). + +Architecture: + + HTTP request → predict() → audio decode (CPU) → enqueue(_BatchRequest) + │ + ▼ + ┌──────────────────────────────────────────────────────────────────┐ + │ _batch_queue (FIFO) │ + └──────────────────────────────────────────────────────────────────┘ + │ │ + drain + adaptive wait (blocks for first item) + │ ▼ + ┌────────┴──────────────────────────────────┐ + │ Collector thread │ + │ - drain queue + adaptive collect window │ + │ - pad audio (CPU) to max batch length │ + │ - run mel preprocessor on GPU │ + │ - put → _ready_queue (maxsize=1) │ + └────────┬──────────────────────────────────┘ + │ + ┌────────▼──────────────────────────────────┐ + │ Inference thread │ + │ - take preprocessed batch │ + │ - encoder + TDT decoder (with CUDA │ + │ graphs for the per-step decoder loop) │ + │ - dispatch text → request Futures │ + └───────────────────────────────────────────┘ + +Key design choices for performance on T4: + + 1. Auto-detect dtype: fp16 on Turing (T4, sm75), bf16 on Ampere+ + (sm>=80). Hardcoding bf16 on T4 silently falls back to slow paths + because T4 has no bf16 tensor cores. + + 2. Decompose model.transcribe(): bypass NeMo's high-level wrapper + (which has per-call freeze/unfreeze + mode switching overhead) + and call preprocessor → encoder → rnnt_decoder_predictions_tensor + directly. + + 3. Pipelined two-thread design: a collector thread prepares the next + batch (drain + pad + preprocessor) while an inference thread runs + encoder + decoder on the previous batch. The single-slot + `_ready_queue` between them keeps exactly one batch pre-staged. + + 4. CUDA graphs ON for the TDT decoder: captures the per-token kernel + sequence and replays without per-step Python overhead. Big win at + small/medium batch sizes. + + 5. Adaptive collect window: the collector drains the queue + immediately (no fixed sleep), then waits up to BATCH_COLLECT_MS + for stragglers ONLY if the initial drain produced ≥2 items. At + low load (single arrivals) requests fire instantly; at high load + the wait window grows batches above 1. + + 6. Input audio is sorted by length within a batch: minor + padding-efficiency win for variable-length clips. + +Environment variables: + BATCH_INFERENCE "true" | "false" (default: true) + MAX_BATCH_SIZE int (default: 64) + BATCH_COLLECT_MS float, ms (default: 20) + CUDA_GRAPHS "true" | "false" (default: true) + TORCH_COMPILE "true" | "false" (default: false) + COMPILE_MODE reduce-overhead|default (default: reduce-overhead) +""" + +import base64 +import io +import logging +import os +import queue +import subprocess +import tempfile +import threading +import time +from concurrent.futures import Future + +import numpy as np +import requests +import soundfile as sf +import torch + +logger = logging.getLogger(__name__) + +MODEL_CACHE_DIR = "/app/model_cache/parakeet-tdt-0.6b-v3" +MODEL_NEMO_FILE = "parakeet-tdt-0.6b-v3.nemo" +SAMPLE_RATE = 16000 + + +# ── Helpers ─────────────────────────────────────────────────────────── + + +def json_serialize_recursive(obj): + if isinstance(obj, dict): + return {k: json_serialize_recursive(v) for k, v in obj.items()} + elif isinstance(obj, list): + return [json_serialize_recursive(v) for v in obj] + elif isinstance(obj, tuple): + return tuple(json_serialize_recursive(v) for v in obj) + elif isinstance(obj, (int, float, str, bool, type(None))): + return obj + elif isinstance(obj, torch.Tensor): + return obj.tolist() + elif hasattr(obj, "__dict__"): + return json_serialize_recursive(obj.__dict__) + else: + return str(obj) + + +def download_and_decode_audio(audio_url: str) -> np.ndarray: + """Fetch audio over HTTP and decode to 16kHz mono float32 numpy array.""" + response = requests.get(audio_url, timeout=60) + response.raise_for_status() + audio_bytes = response.content + + try: + audio_data, sr = sf.read(io.BytesIO(audio_bytes), dtype="float32") + if len(audio_data.shape) > 1: + audio_data = audio_data.mean(axis=1) + if sr == SAMPLE_RATE: + return audio_data + except Exception: + sr = None + + # Fall back to ffmpeg for codecs soundfile can't handle (mp3, opus, etc.) + url_clean = audio_url.split("?")[0].lower() + ext = os.path.splitext(url_clean)[1] or ".audio" + with tempfile.NamedTemporaryFile(suffix=ext, delete=True) as tmp: + tmp.write(audio_bytes) + tmp.flush() + result = subprocess.run( + [ + "ffmpeg", "-i", tmp.name, "-ar", str(SAMPLE_RATE), "-ac", "1", + "-f", "f32le", "-loglevel", "error", "-", + ], + capture_output=True, check=True, + ) + return np.frombuffer(result.stdout, dtype=np.float32) + + +def decode_base64_audio(audio_b64: str) -> np.ndarray: + """Decode base64-encoded audio bytes to 16kHz mono float32 numpy array.""" + raw = base64.b64decode(audio_b64) + audio_data, sr = sf.read(io.BytesIO(raw), dtype="float32") + if len(audio_data.shape) > 1: + audio_data = audio_data.mean(axis=1) + if sr != SAMPLE_RATE: + with tempfile.NamedTemporaryFile(suffix=".wav", delete=True) as tmp: + tmp.write(raw) + tmp.flush() + result = subprocess.run( + [ + "ffmpeg", "-i", tmp.name, "-ar", str(SAMPLE_RATE), "-ac", "1", + "-f", "f32le", "-loglevel", "error", "-", + ], + capture_output=True, check=True, + ) + return np.frombuffer(result.stdout, dtype=np.float32) + return audio_data + + +def _detect_dtype() -> torch.dtype: + """bfloat16 on Ampere+ (sm>=80); float16 on Turing/T4 (sm75); fp32 on CPU. + + T4 has no bf16 tensor cores — bf16 there silently falls back to + slow non-tensor-core paths. fp16 is the right choice on T4. + """ + if not torch.cuda.is_available(): + return torch.float32 + major, _ = torch.cuda.get_device_capability() + return torch.bfloat16 if major >= 8 else torch.float16 + + +# ── Batch request ───────────────────────────────────────────────────── + + +class _BatchRequest: + __slots__ = ("audio_array", "future") + + def __init__(self, audio_array: np.ndarray): + self.audio_array = audio_array + self.future: Future = Future() + + +# ── Model loading ───────────────────────────────────────────────────── + + +def _load_nemo_model( + model_path: str, dtype: torch.dtype, use_compile: bool, + compile_mode: str, enable_cuda_graphs: bool = True, +): + """Load a NeMo .nemo file, configure it for fast inference, and warm up. + + Configuration: + - CUDA graphs for the TDT decoder (enabled by default; captures the + per-token decoder kernel sequence so it replays without per-step + Python overhead). + - Local windowed self-attention (window=256) instead of full + attention. Lower memory + faster on long audio without quality + cost on typical short-utterance traffic. + - Subsampling conv chunking factor = 1 (process the whole input + in one chunk; chunking>1 is for OOM mitigation on very long + audio that we don't need here). + - Cast weights to the target dtype (fp16 on T4, bf16 on Ampere+). + """ + import nemo.collections.asr as nemo_asr + + model = nemo_asr.models.ASRModel.restore_from(restore_path=model_path) + + if not enable_cuda_graphs: + model.decoding.decoding.decoding_computer.disable_cuda_graphs() + logger.info("CUDA graphs DISABLED for TDT decoder") + else: + logger.info("CUDA graphs ENABLED for TDT decoder") + + model.change_attention_model("rel_pos_local_attn", [256, 256]) + model.change_subsampling_conv_chunking_factor(1) + model.to(dtype) + model.eval() + + if use_compile: + try: + model.encoder = torch.compile(model.encoder, mode=compile_mode) + except Exception as exc: + logger.warning("torch.compile failed, using eager: %s", exc) + + # Warm up so the first real request doesn't pay JIT / kernel-autotune cost. + dummy = np.random.randn(SAMPLE_RATE * 5).astype(np.float32) + with torch.inference_mode(): + model.transcribe([dummy], timestamps=False) + + return model + + +# ── Model class ─────────────────────────────────────────────────────── + + +class Model: + def __init__(self, lazy_data_resolver, **kwargs) -> None: + self._lazy_data_resolver = lazy_data_resolver + self._hf_access_token = kwargs["secrets"].get("hf_access_token") + self._transcribe_lock = threading.Lock() + + self._batch_enabled = os.getenv("BATCH_INFERENCE", "true").lower() == "true" + self._max_batch_size = int(os.getenv("MAX_BATCH_SIZE", "64")) + self._batch_collect_s = float(os.getenv("BATCH_COLLECT_MS", "20")) / 1000.0 + self._use_compile = os.getenv("TORCH_COMPILE", "false").lower() == "true" + self._compile_mode = os.getenv("COMPILE_MODE", "reduce-overhead") + self._cuda_graphs = os.getenv("CUDA_GRAPHS", "true").lower() == "true" + + self.model = None + self._device = None + # Inbound request queue (one entry per HTTP request). + self._batch_queue: queue.Queue[_BatchRequest] = queue.Queue() + # Single-slot handoff between collector and inference threads. + # maxsize=1 ensures the collector pre-stages exactly one batch + # ahead — enough to overlap CPU prep with previous-batch GPU + # work, but no more (extra slots would just inflate latency). + self._ready_queue: queue.Queue = queue.Queue(maxsize=1) + + # ── Load ────────────────────────────────────────────────────────── + + def load(self): + self._lazy_data_resolver.block_until_download_complete() + model_path = os.path.join(MODEL_CACHE_DIR, MODEL_NEMO_FILE) + + dtype = _detect_dtype() + gpu = (torch.cuda.get_device_name() + if torch.cuda.is_available() else "cpu") + logger.info( + "Config: max_batch=%d collect_ms=%.0f dtype=%s " + "compile=%s cuda_graphs=%s gpu=%s", + self._max_batch_size, self._batch_collect_s * 1000, dtype, + self._use_compile, self._cuda_graphs, gpu, + ) + + t0 = time.time() + self.model = _load_nemo_model( + model_path, dtype, self._use_compile, self._compile_mode, + enable_cuda_graphs=self._cuda_graphs, + ) + self._device = next(self.model.parameters()).device + logger.info("Model loaded + warmed up in %.1fs", time.time() - t0) + + if self._batch_enabled: + threading.Thread( + target=self._batch_collector, daemon=True, name="collector", + ).start() + threading.Thread( + target=self._batch_inference, daemon=True, name="inference", + ).start() + logger.info("Pipelined batch worker started (collector + inference)") + + logger.info("Load complete — ready to serve") + + # ── Pipelined batch worker ─────────────────────────────────────── + + def _drain_batch(self) -> list[_BatchRequest]: + """Collect a batch: block for first item, drain queue, adaptive wait. + + Adaptive collect window: only wait for stragglers if the initial + drain already produced 2+ items (evidence of concurrent traffic). + When a request arrives alone, fire it immediately to preserve + tail latency at low load. At medium-to-high load the queue is + rarely empty after the first drain, so the wait window kicks in + and grows the batch above 1. + """ + first = self._batch_queue.get() + batch = [first] + + while len(batch) < self._max_batch_size: + try: + batch.append(self._batch_queue.get_nowait()) + except queue.Empty: + break + + if ( + len(batch) >= 2 + and len(batch) < self._max_batch_size + and self._batch_collect_s > 0 + ): + deadline = time.monotonic() + self._batch_collect_s + while len(batch) < self._max_batch_size: + remaining = deadline - time.monotonic() + if remaining <= 0: + break + try: + batch.append(self._batch_queue.get(timeout=remaining)) + except queue.Empty: + break + + # Sort by audio length so padding within the batch is minimized + # (the longest item drives the padded length anyway, but keeping + # similar lengths together improves cache locality slightly). + batch.sort(key=lambda r: len(r.audio_array)) + return batch + + def _batch_collector(self): + """Thread 1: collect batches, pad audio, run mel preprocessor. + + Runs concurrently with _batch_inference. While the GPU is busy + with batch N's encoder + decoder, this thread builds batch N+1: + - drain inbound queue (with adaptive collect window) + - pad audio arrays to the batch's max length + - move to GPU + - run NeMo's mel preprocessor (on the default CUDA stream) + - hand off to _ready_queue for the inference thread + """ + while True: + batch = [] + try: + batch = self._drain_batch() + audios = [r.audio_array for r in batch] + + max_len = max(len(a) for a in audios) + padded = np.zeros((len(audios), max_len), dtype=np.float32) + lengths = np.zeros(len(audios), dtype=np.int64) + for i, a in enumerate(audios): + padded[i, :len(a)] = a + lengths[i] = len(a) + + input_signal = torch.from_numpy(padded).to( + device=self._device, dtype=torch.float32, + ) + input_lengths = torch.from_numpy(lengths).to(device=self._device) + + with torch.inference_mode(): + processed, proc_len = self.model.preprocessor( + input_signal=input_signal, length=input_lengths, + ) + + self._ready_queue.put((processed, proc_len, batch)) + + except Exception as exc: + logger.error("Collector error: %s", exc) + for req in batch: + if not req.future.done(): + req.future.set_exception(exc) + + def _batch_inference(self): + """Thread 2: run encoder + TDT decoder on preprocessed batches. + + Decomposed pipeline: we call NeMo's preprocessor / encoder / + decoder modules directly instead of going through the high-level + `model.transcribe()` API, which has per-call setup/teardown + overhead (freeze/unfreeze, mode switching, hypothesis post- + processing) that adds up at high request rates. + + The TDT decoder's internal token loop calls .item() per step, + which implicitly synchronizes with the GPU — so by the time + rnnt_decoder_predictions_tensor returns, all GPU work for this + batch is complete and we can safely dispatch text to futures. + """ + while True: + batch = [] + try: + processed, proc_len, batch = self._ready_queue.get() + + with torch.inference_mode(): + encoded, encoded_len = self.model.encoder( + audio_signal=processed, length=proc_len, + ) + dec_result = self.model.decoding.rnnt_decoder_predictions_tensor( + encoder_output=encoded, + encoded_lengths=encoded_len, + return_hypotheses=False, + ) + + raw_hyps = ( + dec_result[0] if isinstance(dec_result, tuple) + else dec_result + ) + for req, hyp in zip(batch, raw_hyps): + text = hyp.text if hasattr(hyp, "text") else str(hyp) + req.future.set_result(text) + + except Exception as exc: + logger.error("Inference error: %s", exc) + for req in batch: + if not req.future.done(): + req.future.set_exception(exc) + + # ── Predict ─────────────────────────────────────────────────────── + + def predict(self, request: dict): + """HTTP entrypoint. Accepts: + audio_url: URL to fetch and decode + audio_base64: base64-encoded audio bytes (any soundfile/ffmpeg-readable) + timestamps: bool — if true, returns word-level timestamps + (uses serial path; not batched). + """ + audio_url = request.get("audio_url") + audio_b64 = request.get("audio_base64") + is_timestamps = request.get("timestamps", False) + + try: + if audio_b64: + audio_array = decode_base64_audio(audio_b64) + else: + audio_array = download_and_decode_audio(audio_url) + + if self._batch_enabled and not is_timestamps: + return self._predict_batched(audio_array) + return self._predict_serial(audio_array, is_timestamps) + + except Exception as exc: + logger.error("Predict error: %s", exc) + return {"error": str(exc)} + + def _predict_batched(self, audio_array: np.ndarray) -> dict: + """Batched path (no timestamps): enqueue and wait for batch worker.""" + req = _BatchRequest(audio_array) + self._batch_queue.put(req) + text = req.future.result(timeout=120) + return {"transcript": text} + + def _predict_serial(self, audio_array: np.ndarray, is_timestamps: bool) -> dict: + """Serial path: used when timestamps are requested (the batched + decoder doesn't produce per-token timestamps in this code path). + Single-threaded under self._transcribe_lock so concurrent serial + requests don't trip on NeMo's shared transcribe() state. + """ + with self._transcribe_lock: + with torch.inference_mode(): + results = self.model.transcribe( + [audio_array], timestamps=is_timestamps, + ) + hyp = results[0] + if isinstance(hyp, (list, tuple)): + hyp = hyp[0] + + if not is_timestamps: + text = hyp.text if hasattr(hyp, "text") else str(hyp) + return {"transcript": text} + return { + "transcript": json_serialize_recursive({ + "text": hyp.text, + "score": hyp.score, + "timestep": hyp.timestep, + }) + } From 3681081be8f56b614ef82af44aaf28d64ab22c8c Mon Sep 17 00:00:00 2001 From: Jie Gong Date: Tue, 30 Jun 2026 12:07:24 -0600 Subject: [PATCH 2/3] Apply ruff-format to parakeet batching model.py Co-authored-by: Cursor --- .../model/model.py | 102 ++++++++++++------ 1 file changed, 71 insertions(+), 31 deletions(-) diff --git a/nvidia/parakeet-tdt-0_6b-v3-batching/model/model.py b/nvidia/parakeet-tdt-0_6b-v3-batching/model/model.py index a1024e9ba..506e0c363 100644 --- a/nvidia/parakeet-tdt-0_6b-v3-batching/model/model.py +++ b/nvidia/parakeet-tdt-0_6b-v3-batching/model/model.py @@ -132,10 +132,21 @@ def download_and_decode_audio(audio_url: str) -> np.ndarray: tmp.flush() result = subprocess.run( [ - "ffmpeg", "-i", tmp.name, "-ar", str(SAMPLE_RATE), "-ac", "1", - "-f", "f32le", "-loglevel", "error", "-", + "ffmpeg", + "-i", + tmp.name, + "-ar", + str(SAMPLE_RATE), + "-ac", + "1", + "-f", + "f32le", + "-loglevel", + "error", + "-", ], - capture_output=True, check=True, + capture_output=True, + check=True, ) return np.frombuffer(result.stdout, dtype=np.float32) @@ -152,10 +163,21 @@ def decode_base64_audio(audio_b64: str) -> np.ndarray: tmp.flush() result = subprocess.run( [ - "ffmpeg", "-i", tmp.name, "-ar", str(SAMPLE_RATE), "-ac", "1", - "-f", "f32le", "-loglevel", "error", "-", + "ffmpeg", + "-i", + tmp.name, + "-ar", + str(SAMPLE_RATE), + "-ac", + "1", + "-f", + "f32le", + "-loglevel", + "error", + "-", ], - capture_output=True, check=True, + capture_output=True, + check=True, ) return np.frombuffer(result.stdout, dtype=np.float32) return audio_data @@ -188,8 +210,11 @@ def __init__(self, audio_array: np.ndarray): def _load_nemo_model( - model_path: str, dtype: torch.dtype, use_compile: bool, - compile_mode: str, enable_cuda_graphs: bool = True, + model_path: str, + dtype: torch.dtype, + use_compile: bool, + compile_mode: str, + enable_cuda_graphs: bool = True, ): """Load a NeMo .nemo file, configure it for fast inference, and warm up. @@ -267,18 +292,24 @@ def load(self): model_path = os.path.join(MODEL_CACHE_DIR, MODEL_NEMO_FILE) dtype = _detect_dtype() - gpu = (torch.cuda.get_device_name() - if torch.cuda.is_available() else "cpu") + gpu = torch.cuda.get_device_name() if torch.cuda.is_available() else "cpu" logger.info( "Config: max_batch=%d collect_ms=%.0f dtype=%s " "compile=%s cuda_graphs=%s gpu=%s", - self._max_batch_size, self._batch_collect_s * 1000, dtype, - self._use_compile, self._cuda_graphs, gpu, + self._max_batch_size, + self._batch_collect_s * 1000, + dtype, + self._use_compile, + self._cuda_graphs, + gpu, ) t0 = time.time() self.model = _load_nemo_model( - model_path, dtype, self._use_compile, self._compile_mode, + model_path, + dtype, + self._use_compile, + self._compile_mode, enable_cuda_graphs=self._cuda_graphs, ) self._device = next(self.model.parameters()).device @@ -286,10 +317,14 @@ def load(self): if self._batch_enabled: threading.Thread( - target=self._batch_collector, daemon=True, name="collector", + target=self._batch_collector, + daemon=True, + name="collector", ).start() threading.Thread( - target=self._batch_inference, daemon=True, name="inference", + target=self._batch_inference, + daemon=True, + name="inference", ).start() logger.info("Pipelined batch worker started (collector + inference)") @@ -358,17 +393,19 @@ def _batch_collector(self): padded = np.zeros((len(audios), max_len), dtype=np.float32) lengths = np.zeros(len(audios), dtype=np.int64) for i, a in enumerate(audios): - padded[i, :len(a)] = a + padded[i, : len(a)] = a lengths[i] = len(a) input_signal = torch.from_numpy(padded).to( - device=self._device, dtype=torch.float32, + device=self._device, + dtype=torch.float32, ) input_lengths = torch.from_numpy(lengths).to(device=self._device) with torch.inference_mode(): processed, proc_len = self.model.preprocessor( - input_signal=input_signal, length=input_lengths, + input_signal=input_signal, + length=input_lengths, ) self._ready_queue.put((processed, proc_len, batch)) @@ -400,7 +437,8 @@ def _batch_inference(self): with torch.inference_mode(): encoded, encoded_len = self.model.encoder( - audio_signal=processed, length=proc_len, + audio_signal=processed, + length=proc_len, ) dec_result = self.model.decoding.rnnt_decoder_predictions_tensor( encoder_output=encoded, @@ -409,8 +447,7 @@ def _batch_inference(self): ) raw_hyps = ( - dec_result[0] if isinstance(dec_result, tuple) - else dec_result + dec_result[0] if isinstance(dec_result, tuple) else dec_result ) for req, hyp in zip(batch, raw_hyps): text = hyp.text if hasattr(hyp, "text") else str(hyp) @@ -426,10 +463,10 @@ def _batch_inference(self): def predict(self, request: dict): """HTTP entrypoint. Accepts: - audio_url: URL to fetch and decode - audio_base64: base64-encoded audio bytes (any soundfile/ffmpeg-readable) - timestamps: bool — if true, returns word-level timestamps - (uses serial path; not batched). + audio_url: URL to fetch and decode + audio_base64: base64-encoded audio bytes (any soundfile/ffmpeg-readable) + timestamps: bool — if true, returns word-level timestamps + (uses serial path; not batched). """ audio_url = request.get("audio_url") audio_b64 = request.get("audio_base64") @@ -465,7 +502,8 @@ def _predict_serial(self, audio_array: np.ndarray, is_timestamps: bool) -> dict: with self._transcribe_lock: with torch.inference_mode(): results = self.model.transcribe( - [audio_array], timestamps=is_timestamps, + [audio_array], + timestamps=is_timestamps, ) hyp = results[0] if isinstance(hyp, (list, tuple)): @@ -475,9 +513,11 @@ def _predict_serial(self, audio_array: np.ndarray, is_timestamps: bool) -> dict: text = hyp.text if hasattr(hyp, "text") else str(hyp) return {"transcript": text} return { - "transcript": json_serialize_recursive({ - "text": hyp.text, - "score": hyp.score, - "timestep": hyp.timestep, - }) + "transcript": json_serialize_recursive( + { + "text": hyp.text, + "score": hyp.score, + "timestep": hyp.timestep, + } + ) } From b134039e27818eb6801974b8800692f1ee88c190 Mon Sep 17 00:00:00 2001 From: Jie Gong Date: Tue, 30 Jun 2026 12:07:24 -0600 Subject: [PATCH 3/3] chore: ruff lint+format fixes for triton example (from #577) Unblocks the `pre-commit run --all-files` lint job, which currently fails on every PR because these files (added in #577) have one ruff F-class error and are not ruff-formatted. Fixes are the safe auto-fixes pre-commit itself applies. Co-authored-by: Cursor --- .../tensorrt_llm/1/helpers.py | 96 +++-- .../model_repository/tensorrt_llm/1/model.py | 354 +++++++++++------- 2 files changed, 264 insertions(+), 186 deletions(-) diff --git a/triton-inference-server/tensorrtllm-backend/data/model_repository/tensorrt_llm/1/helpers.py b/triton-inference-server/tensorrtllm-backend/data/model_repository/tensorrt_llm/1/helpers.py index 0b18b5d74..f61d6d2c2 100644 --- a/triton-inference-server/tensorrtllm-backend/data/model_repository/tensorrt_llm/1/helpers.py +++ b/triton-inference-server/tensorrtllm-backend/data/model_repository/tensorrt_llm/1/helpers.py @@ -4,12 +4,12 @@ from torch.utils.dlpack import from_dlpack -def convert_request_input_to_dict(request, param_mappings, default_values, - batch_size, batch_index): +def convert_request_input_to_dict( + request, param_mappings, default_values, batch_size, batch_index +): kwargs = {} for source_name, target_name in param_mappings.items(): - value = get_input_scalar_by_name(request, source_name, batch_size, - batch_index) + value = get_input_scalar_by_name(request, source_name, batch_size, batch_index) if value is None and source_name in default_values: kwargs[target_name] = default_values[source_name] elif value is not None: @@ -24,16 +24,16 @@ def get_sampling_params_from_request(request, batch_size=1, batch_index=0): Used in llmapi/tensorrt_llm """ sampling_params_args = [ - 'best_of', - 'temperature', - 'top_k', - 'top_p', - 'frequency_penalty', - 'presence_penalty', - 'max_tokens', - 'seed', - 'exclude_input_from_output', - 'return_perf_metrics', + "best_of", + "temperature", + "top_k", + "top_p", + "frequency_penalty", + "presence_penalty", + "max_tokens", + "seed", + "exclude_input_from_output", + "return_perf_metrics", ] param_mappings = {} for arg in sampling_params_args: @@ -45,13 +45,13 @@ def get_sampling_params_from_request(request, batch_size=1, batch_index=0): param_mappings.setdefault(arg, arg) default_values = { - 'sampling_param_best_of': 1, - 'sampling_param_exclude_input_from_output': False, - 'sampling_param_return_perf_metrics': False, + "sampling_param_best_of": 1, + "sampling_param_exclude_input_from_output": False, + "sampling_param_return_perf_metrics": False, } - kwargs = convert_request_input_to_dict(request, param_mappings, - default_values, batch_size, - batch_index) + kwargs = convert_request_input_to_dict( + request, param_mappings, default_values, batch_size, batch_index + ) # Parse stop as a list of strings not scalar kwargs["stop"] = get_input_tensor_by_name(request, "sampling_param_stop") if kwargs["stop"] is None: @@ -65,21 +65,22 @@ def get_output_config_from_request(request, batch_size=1, batch_index=0): Used in llmapi/tensorrt_llm """ output_config_args = [ - 'return_finish_reason', 'return_stop_reason', - 'return_cumulative_logprob' + "return_finish_reason", + "return_stop_reason", + "return_cumulative_logprob", ] param_mappings = {} for arg in output_config_args: param_mappings[arg] = arg default_values = { - 'return_finish_reason': False, - 'return_stop_reason': False, - 'return_cumulative_logprob': False + "return_finish_reason": False, + "return_stop_reason": False, + "return_cumulative_logprob": False, } - kwargs = convert_request_input_to_dict(request, param_mappings, - default_values, batch_size, - batch_index) + kwargs = convert_request_input_to_dict( + request, param_mappings, default_values, batch_size, batch_index + ) return kwargs @@ -88,20 +89,13 @@ def get_streaming_from_request(request, batch_size=1, batch_index=0): Helper function to get streaming from request Used in llmapi/tensorrt_llm """ - streaming = get_input_scalar_by_name( - request, "streaming", batch_size, batch_index - ) + streaming = get_input_scalar_by_name(request, "streaming", batch_size, batch_index) if streaming is None: - streaming = get_input_scalar_by_name( - request, "stream", batch_size, batch_index - ) + streaming = get_input_scalar_by_name(request, "stream", batch_size, batch_index) return streaming or False -def get_input_scalar_by_name(request, - name, - expected_batch_size=1, - batch_index=0): +def get_input_scalar_by_name(request, name, expected_batch_size=1, batch_index=0): tensor = pb_utils.get_input_tensor_by_name(request, name) if tensor is None: return None @@ -109,16 +103,15 @@ def get_input_scalar_by_name(request, if tensor.size != expected_batch_size: raise pb_utils.TritonModelException( - f"Expected a scalar tensor for tensor {name}") + f"Expected a scalar tensor for tensor {name}" + ) return tensor.reshape(-1)[batch_index] -def get_input_tensor_by_name(request, - name, - expected_batch_size=None, - batch_index=None, - force_on_torch=False): +def get_input_tensor_by_name( + request, name, expected_batch_size=None, batch_index=None, force_on_torch=False +): tensor = pb_utils.get_input_tensor_by_name(request, name) if tensor is None: return None @@ -128,15 +121,19 @@ def get_input_tensor_by_name(request, else: tensor = from_dlpack(tensor.to_dlpack()) - if expected_batch_size is not None and tensor.shape[ - 0] != expected_batch_size: + if expected_batch_size is not None and tensor.shape[0] != expected_batch_size: raise pb_utils.TritonModelException( f"Expected batch size doesn't match batch size for tensor {name}. Expected {expected_batch_size} got {tensor.shape[0]}" ) - if batch_index is not None and expected_batch_size is not None and batch_index >= expected_batch_size: + if ( + batch_index is not None + and expected_batch_size is not None + and batch_index >= expected_batch_size + ): raise pb_utils.TritonModelException( - f"Invalid batch index in get_input_tensor_by_name for {name}") + f"Invalid batch index in get_input_tensor_by_name for {name}" + ) if batch_index is not None: # Add leading 1 batch dimension @@ -169,4 +166,5 @@ def get_parameter(model_config, name, pytype=str): if name not in model_config["parameters"]: return None return read_parameter_as_type( - model_config["parameters"][name]['string_value'], name, pytype) + model_config["parameters"][name]["string_value"], name, pytype + ) diff --git a/triton-inference-server/tensorrtllm-backend/data/model_repository/tensorrt_llm/1/model.py b/triton-inference-server/tensorrtllm-backend/data/model_repository/tensorrt_llm/1/model.py index 1c0f261e2..603e74a8d 100644 --- a/triton-inference-server/tensorrtllm-backend/data/model_repository/tensorrt_llm/1/model.py +++ b/triton-inference-server/tensorrtllm-backend/data/model_repository/tensorrt_llm/1/model.py @@ -40,9 +40,12 @@ import pandas as pd import triton_python_backend_utils as pb_utils import yaml -from helpers import (get_input_tensor_by_name, get_output_config_from_request, - get_sampling_params_from_request, - get_streaming_from_request) +from helpers import ( + get_input_tensor_by_name, + get_output_config_from_request, + get_sampling_params_from_request, + get_streaming_from_request, +) # NOTE: tensorrt_llm imports are deferred to initialize() to support multi-instance deployments. # Root cause: CUDA is not fork-safe. Importing tensorrt_llm imports torch, which initializes @@ -67,27 +70,23 @@ def get_model_config(filename, include_keys=None, exclude_keys=None): engine_config = yaml.safe_load(file) except Exception as e: raise pb_utils.TritonModelException( - f"Failed to parse YAML engine config: {e}") + f"Failed to parse YAML engine config: {e}" + ) - assert engine_config is not None, f"'{filename}' containing TRT-LLM engine args not found in '{pb_utils.get_model_dir()}'" + assert engine_config is not None, ( + f"'{filename}' containing TRT-LLM engine args not found in '{pb_utils.get_model_dir()}'" + ) if include_keys: - engine_config = { - k: v - for k, v in engine_config.items() if k in include_keys - } + engine_config = {k: v for k, v in engine_config.items() if k in include_keys} if exclude_keys: engine_config = { - k: v - for k, v in engine_config.items() if k not in exclude_keys + k: v for k, v in engine_config.items() if k not in exclude_keys } return engine_config -def get_input_scalar_by_name(request, - name, - expected_batch_size=1, - batch_index=0): +def get_input_scalar_by_name(request, name, expected_batch_size=1, batch_index=0): tensor = pb_utils.get_input_tensor_by_name(request, name) if tensor is None: return None @@ -95,13 +94,13 @@ def get_input_scalar_by_name(request, if tensor.size != expected_batch_size: raise pb_utils.TritonModelException( - f"Expected a scalar tensor for tensor {name}") + f"Expected a scalar tensor for tensor {name}" + ) return tensor.item(batch_index) class TritonPythonModel: - @classmethod def auto_complete_config(cls, auto_complete_model_config): """ @@ -117,14 +116,16 @@ def auto_complete_config(cls, auto_complete_model_config): - This function is called when Triton server starts. - It combines the default configurations in config.pbtxt with the triton_config values in model.yaml """ - triton_config = get_model_config(os.environ.get('LLM_CONFIG_PATH', - 'model.yaml'), - include_keys=["triton_config" - ])["triton_config"] + triton_config = get_model_config( + os.environ.get("LLM_CONFIG_PATH", "model.yaml"), + include_keys=["triton_config"], + )["triton_config"] auto_complete_model_config.set_model_transaction_policy( - dict(decoupled=bool(triton_config["decoupled"]))) + dict(decoupled=bool(triton_config["decoupled"])) + ) auto_complete_model_config.set_max_batch_size( - int(triton_config["max_batch_size"])) + int(triton_config["max_batch_size"]) + ) return auto_complete_model_config @@ -139,37 +140,41 @@ def initialize(self, args): - Implementing `initialize` function is optional. """ self.model_config = json.loads(args["model_config"]) - triton_config = get_model_config(os.environ.get('LLM_CONFIG_PATH', - 'model.yaml'), - include_keys=["triton_config" - ])["triton_config"] + triton_config = get_model_config( + os.environ.get("LLM_CONFIG_PATH", "model.yaml"), + include_keys=["triton_config"], + )["triton_config"] self.decoupled = bool(triton_config["decoupled"]) - self.params = self.model_config['parameters'] + self.params = self.model_config["parameters"] self.logger = pb_utils.Logger # Check instance count from config.pbtxt instance_count = int( - self.model_config.get("instance_group", [{}])[0].get("count", 1)) + self.model_config.get("instance_group", [{}])[0].get("count", 1) + ) self._use_multi_instance = instance_count > 1 self._instance_idx = 0 # Get GPU device IDs from config.pbtxt parameters # Format: "0,1;2,3" means instance 0 gets GPUs 0,1, instance 1 gets GPUs 2,3 if self._use_multi_instance: - gpu_device_ids_param = self.params.get("gpu_device_ids", - {}).get("string_value", "") + gpu_device_ids_param = self.params.get("gpu_device_ids", {}).get( + "string_value", "" + ) if not gpu_device_ids_param: self.logger.log_error( "[trtllm] gpu_device_ids parameter required for multi-instance mode" ) raise ValueError( - "gpu_device_ids parameter required when instance count > 1") + "gpu_device_ids parameter required when instance count > 1" + ) # Get instance index from instance name (e.g., "tensorrt_llm_0" -> 0) import re + instance_name = args.get("model_instance_name", "") if instance_name: - match = re.search(r'_(\d+)$', instance_name) + match = re.search(r"_(\d+)$", instance_name) if match: self._instance_idx = int(match.group(1)) @@ -195,24 +200,27 @@ def initialize(self, args): ) # Import tensorrt_llm AFTER setting CUDA_VISIBLE_DEVICES - from tensorrt_llm.llmapi.llm_utils import \ - update_llm_args_with_extra_dict + from tensorrt_llm.llmapi.llm_utils import update_llm_args_with_extra_dict text_output_config = pb_utils.get_output_config_by_name( - self.model_config, "text_output") + self.model_config, "text_output" + ) self.output_dtype = pb_utils.triton_string_to_numpy( - text_output_config["data_type"]) + text_output_config["data_type"] + ) from tensorrt_llm._utils import global_mpi_rank + is_leader = global_mpi_rank() == 0 if is_leader: # Initialize engine arguments self.llm_engine_args = update_llm_args_with_extra_dict( {}, - get_model_config(os.environ.get('LLM_CONFIG_PATH', - 'model.yaml'), - exclude_keys=["triton_config"]), + get_model_config( + os.environ.get("LLM_CONFIG_PATH", "model.yaml"), + exclude_keys=["triton_config"], + ), ) # For multi-instance mode, set gpus_per_node to match the assigned GPUs @@ -224,19 +232,25 @@ def initialize(self, args): ) triton_config = get_model_config( - os.environ.get('LLM_CONFIG_PATH', 'model.yaml'), - include_keys=["triton_config"])["triton_config"] - self.cancellation_check_period_ms = int( - triton_config["cancellation_check_period_ms"] - ) if "cancellation_check_period_ms" in triton_config else 100 + os.environ.get("LLM_CONFIG_PATH", "model.yaml"), + include_keys=["triton_config"], + )["triton_config"] + self.cancellation_check_period_ms = ( + int(triton_config["cancellation_check_period_ms"]) + if "cancellation_check_period_ms" in triton_config + else 100 + ) from tensorrt_llm._utils import global_mpi_size + if global_mpi_size() > 1: from mpi4py.MPI import COMM_WORLD from tensorrt_llm.llmapi import MpiCommSession - mpi_session = MpiCommSession(comm=COMM_WORLD, - n_workers=COMM_WORLD.Get_size()) + + mpi_session = MpiCommSession( + comm=COMM_WORLD, n_workers=COMM_WORLD.Get_size() + ) self.llm_engine_args["_mpi_session"] = mpi_session # Starting the TRT-LLM engine with LLM API and its event thread running the AsyncIO event loop. @@ -253,13 +267,13 @@ def initialize(self, args): self.req_id_to_request_data = {} self.triton_user_id_to_req_ids = {} self.lock = threading.Lock() - self.cancellation_thread = threading.Thread( - target=self.cancellation_loop) + self.cancellation_thread = threading.Thread(target=self.cancellation_loop) self.running = True self.cancellation_thread.start() else: from mpi4py.futures import MPICommExecutor from mpi4py.MPI import COMM_WORLD + self.logger.log_info( f"[trtllm] rank{global_mpi_rank()} is waiting for the leader node..." ) @@ -277,8 +291,9 @@ def _init_engine(self): self._llm_engine = None self._llm_engine_start_cv = threading.Condition() self._llm_engine_shutdown_event = asyncio.Event() - self._event_thread = threading.Thread(target=asyncio.run, - args=(self._run_llm_engine(), )) + self._event_thread = threading.Thread( + target=asyncio.run, args=(self._run_llm_engine(),) + ) self._event_thread.start() with self._llm_engine_start_cv: while self._llm_engine is None: @@ -309,10 +324,11 @@ async def async_llm_wrapper(): loop = asyncio.get_running_loop() try: llm = await loop.run_in_executor( - None, lambda: LLM(**self.llm_engine_args)) + None, lambda: LLM(**self.llm_engine_args) + ) yield llm finally: - if 'llm' in locals(): + if "llm" in locals(): # Run shutdown in a thread to avoid blocking await loop.run_in_executor(None, llm.shutdown) @@ -333,7 +349,9 @@ async def async_llm_wrapper(): while self._ongoing_request_count > 0: self.logger.log_info( "[trtllm] Awaiting remaining {} requests".format( - self._ongoing_request_count)) + self._ongoing_request_count + ) + ) await asyncio.sleep(1) # Cancel all tasks in the event loop. @@ -369,11 +387,11 @@ def _response_loop(self): response_sender.send(response, response_flag) # Stop checking for cancellation if the last response is generated. if not response_state["last_response_generated"]: - response_state[ - "is_cancelled"] = response_sender.is_cancelled() + response_state["is_cancelled"] = response_sender.is_cancelled() except Exception as e: self.logger.log_error( - f"An error occurred while sending a response: {e}") + f"An error occurred while sending a response: {e}" + ) finally: if response_flag == pb_utils.TRITONSERVER_RESPONSE_COMPLETE_FINAL: self._ongoing_request_count -= 1 @@ -382,6 +400,7 @@ def cancellation_loop(self): """Checks if any pending requests have been cancelled.""" while self.running: import time + time.sleep(self.cancellation_check_period_ms / 1000.0) with self.lock: cancelled_req_ids = [] @@ -389,14 +408,18 @@ def cancellation_loop(self): if request_data.triton_request.is_cancelled(): request_data.response_iterator.abort() - response_sender = request_data.triton_request.get_response_sender( + response_sender = ( + request_data.triton_request.get_response_sender() ) response_sender.send( pb_utils.InferenceResponse( error=pb_utils.TritonError( "Request cancelled by client", - pb_utils.TritonError.CANCELLED)), - flags=pb_utils.TRITONSERVER_RESPONSE_COMPLETE_FINAL) + pb_utils.TritonError.CANCELLED, + ) + ), + flags=pb_utils.TRITONSERVER_RESPONSE_COMPLETE_FINAL, + ) cancelled_req_ids.append(req_id) for req_id in cancelled_req_ids: del self.req_id_to_request_data[req_id] @@ -404,9 +427,13 @@ def cancellation_loop(self): def handle_stop_request(self, triton_user_id, response_sender): if triton_user_id is None or triton_user_id == "": response_sender.send( - pb_utils.InferenceResponse(error=pb_utils.TritonError( - "A request id must be provided for request cancellation")), - flags=pb_utils.TRITONSERVER_RESPONSE_COMPLETE_FINAL) + pb_utils.InferenceResponse( + error=pb_utils.TritonError( + "A request id must be provided for request cancellation" + ) + ), + flags=pb_utils.TRITONSERVER_RESPONSE_COMPLETE_FINAL, + ) return with self.lock: @@ -418,9 +445,13 @@ def handle_stop_request(self, triton_user_id, response_sender): del self.req_id_to_request_data[req_id] response_sender.send( - pb_utils.InferenceResponse(error=pb_utils.TritonError( - "Request cancelled by client", pb_utils.TritonError.CANCELLED)), - flags=pb_utils.TRITONSERVER_RESPONSE_COMPLETE_FINAL) + pb_utils.InferenceResponse( + error=pb_utils.TritonError( + "Request cancelled by client", pb_utils.TritonError.CANCELLED + ) + ), + flags=pb_utils.TRITONSERVER_RESPONSE_COMPLETE_FINAL, + ) def execute(self, requests): """ @@ -436,9 +467,9 @@ def execute(self, requests): for request in requests: # TODO : [JIRA-4040] Verify Lora if request is not None: - assert ( - self._llm_engine_shutdown_event.is_set() is False - ), "Cannot create tasks after shutdown has been requested" + assert self._llm_engine_shutdown_event.is_set() is False, ( + "Cannot create tasks after shutdown has been requested" + ) coro = self._execute_single_request(request) asyncio.run_coroutine_threadsafe(coro, self._event_loop) @@ -451,7 +482,7 @@ async def _execute_single_request(self, request): response_sender = request.get_response_sender() triton_user_id = request.request_id() - stop = get_input_scalar_by_name(request, 'stop') + stop = get_input_scalar_by_name(request, "stop") if stop: self.handle_stop_request(triton_user_id, response_sender) return @@ -459,8 +490,7 @@ async def _execute_single_request(self, request): response_state = { "response_sender": response_sender, "is_cancelled": False, - "last_response_generated": - False, # last response ready but not yet sent + "last_response_generated": False, # last response ready but not yet sent } self._ongoing_request_count += 1 decrement_ongoing_request_count = True @@ -472,14 +502,17 @@ async def _execute_single_request(self, request): from tensorrt_llm import SamplingParams # TODO: [JIRA-4496] Implement when request contains batched prompts - (prompt, sampling_params, streaming, - output_config) = self._convert_request(request) + (prompt, sampling_params, streaming, output_config) = self._convert_request( + request + ) if streaming and not self.decoupled: raise pb_utils.TritonModelException( - "Streaming is only supported in decoupled mode.") + "Streaming is only supported in decoupled mode." + ) # Generate the response. response_iterator = self._llm_engine.generate_async( - prompt, SamplingParams(**sampling_params), streaming) + prompt, SamplingParams(**sampling_params), streaming + ) with self.lock: self.req_id_to_request_data[triton_req_id] = RequestData( @@ -488,26 +521,28 @@ async def _execute_single_request(self, request): triton_request=request, response_iterator=response_iterator, ) - if triton_user_id is not None and triton_user_id != "" and triton_user_id: + if ( + triton_user_id is not None + and triton_user_id != "" + and triton_user_id + ): self.triton_user_id_to_req_ids[triton_user_id] = set() # TODO: [JIRA-4496] Add all batched request ids to the set - self.triton_user_id_to_req_ids[triton_user_id].add( - triton_req_id) + self.triton_user_id_to_req_ids[triton_user_id].add(triton_req_id) async for request_output in response_iterator: # Send each response if streaming. if streaming: response = self._create_response( - request_output=request_output, - output_config=output_config) + request_output=request_output, output_config=output_config + ) flags = 0 if request_output.finished: response_state["last_response_generated"] = True flags = pb_utils.TRITONSERVER_RESPONSE_COMPLETE_FINAL # with streaming, self._response_loop will decrement self._ongoing_request_count decrement_ongoing_request_count = False - self._response_queue.put_nowait( - (response_state, response, flags)) + self._response_queue.put_nowait((response_state, response, flags)) # Send the last response which contains all the outputs if not streaming. if not streaming: @@ -521,19 +556,24 @@ async def _execute_single_request(self, request): if not was_cancelled: response_sender.send( - self._create_response(request_output=request_output, - output_config=output_config), - flags=pb_utils.TRITONSERVER_RESPONSE_COMPLETE_FINAL) + self._create_response( + request_output=request_output, output_config=output_config + ), + flags=pb_utils.TRITONSERVER_RESPONSE_COMPLETE_FINAL, + ) except Exception as e: self.logger.log_error(f"[trtllm] Error generating request: {e}") error = pb_utils.TritonError(f"Error generating request: {e}") text_output_tensor = pb_utils.Tensor( - "text_output", np.asarray(["N/A"], dtype=self.output_dtype)) + "text_output", np.asarray(["N/A"], dtype=self.output_dtype) + ) response = pb_utils.InferenceResponse( - output_tensors=[text_output_tensor], error=error) + output_tensors=[text_output_tensor], error=error + ) response_sender.send( - response, flags=pb_utils.TRITONSERVER_RESPONSE_COMPLETE_FINAL) + response, flags=pb_utils.TRITONSERVER_RESPONSE_COMPLETE_FINAL + ) raise e finally: @@ -542,7 +582,11 @@ async def _execute_single_request(self, request): with self.lock: if triton_req_id in self.req_id_to_request_data: del self.req_id_to_request_data[triton_req_id] - if triton_user_id is not None and triton_user_id != "" and triton_user_id in self.triton_user_id_to_req_ids: + if ( + triton_user_id is not None + and triton_user_id != "" + and triton_user_id in self.triton_user_id_to_req_ids + ): del self.triton_user_id_to_req_ids[triton_user_id] def _convert_request(self, request): @@ -557,15 +601,17 @@ def _convert_request(self, request): Notes: - Supports a single prompt per request (1D or [1,1] from the OpenAI frontend). """ - text_input = get_input_tensor_by_name(request, 'text_input') + text_input = get_input_tensor_by_name(request, "text_input") if text_input is None: raise pb_utils.TritonModelException( - f"text_input is missing from the request") + "text_input is missing from the request" + ) flat = np.squeeze(text_input).reshape(-1) if flat.size == 0: raise pb_utils.TritonModelException( - "text_input must contain at least one prompt") + "text_input must contain at least one prompt" + ) prompt = flat[0] if isinstance(prompt, bytes): @@ -606,23 +652,26 @@ def _create_response(self, request_output, output_config): """ # TODO: [JIRA-4040] Check if request_output has_error and handle it response = [] - text_output = [ - output.text.encode("utf-8") for output in request_output.outputs - ] + text_output = [output.text.encode("utf-8") for output in request_output.outputs] response.append( - pb_utils.Tensor("text_output", - np.asarray(text_output, dtype=self.output_dtype))) + pb_utils.Tensor( + "text_output", np.asarray(text_output, dtype=self.output_dtype) + ) + ) # Extract and add configurable output fields # The output_config loads related input from request output_fields = { - "return_finish_reason": - ("finish_reason", lambda output: output.finish_reason), - "return_stop_reason": - ("stop_reason", lambda output: output.stop_reason), - "return_cumulative_logprob": - ("cumulative_logprob", lambda output: output.cumulative_logprob) + "return_finish_reason": ( + "finish_reason", + lambda output: output.finish_reason, + ), + "return_stop_reason": ("stop_reason", lambda output: output.stop_reason), + "return_cumulative_logprob": ( + "cumulative_logprob", + lambda output: output.cumulative_logprob, + ), } for config_key, (output_name, extractor) in output_fields.items(): @@ -631,12 +680,15 @@ def _create_response(self, request_output, output_config): str(extractor(output)) for output in request_output.outputs ] response.append( - pb_utils.Tensor(output_name, - np.asarray(tensor_data, dtype=np.object_))) - - if hasattr(request_output.outputs[0], 'request_perf_metrics' - ) and request_output.outputs[0].request_perf_metrics: + pb_utils.Tensor( + output_name, np.asarray(tensor_data, dtype=np.object_) + ) + ) + if ( + hasattr(request_output.outputs[0], "request_perf_metrics") + and request_output.outputs[0].request_perf_metrics + ): perf_metrics = request_output.outputs[0].request_perf_metrics # kv cache perf metrics per request @@ -645,28 +697,37 @@ def _create_response(self, request_output, output_config): response.append( pb_utils.Tensor( "kv_cache_reused_block", - np.asarray([kv_metrics.num_reused_blocks], - dtype=self.output_dtype))) + np.asarray([kv_metrics.num_reused_blocks], dtype=self.output_dtype), + ) + ) response.append( pb_utils.Tensor( "kv_cache_hit_rate", - np.asarray([kv_metrics.kv_cache_hit_rate], - dtype=self.output_dtype))) + np.asarray([kv_metrics.kv_cache_hit_rate], dtype=self.output_dtype), + ) + ) response.append( pb_utils.Tensor( "kv_cache_alloc_new_blocks", - np.asarray([kv_metrics.num_new_allocated_blocks], - dtype=self.output_dtype))) + np.asarray( + [kv_metrics.num_new_allocated_blocks], dtype=self.output_dtype + ), + ) + ) response.append( pb_utils.Tensor( "kv_cache_alloc_total_blocks", - np.asarray([kv_metrics.num_total_allocated_blocks], - dtype=self.output_dtype))) + np.asarray( + [kv_metrics.num_total_allocated_blocks], dtype=self.output_dtype + ), + ) + ) response.append( pb_utils.Tensor( "kv_cache_missed_block", - np.asarray([kv_metrics.num_missed_blocks], - dtype=self.output_dtype))) + np.asarray([kv_metrics.num_missed_blocks], dtype=self.output_dtype), + ) + ) # timing perf metrics per request timing_metrics = perf_metrics.timing_metrics @@ -675,50 +736,71 @@ def _create_response(self, request_output, output_config): "arrival_time_ns", np.asarray( [pd.Timedelta(timing_metrics.arrival_time).value], - dtype=self.output_dtype))) + dtype=self.output_dtype, + ), + ) + ) response.append( pb_utils.Tensor( "first_scheduled_time_ns", - np.asarray([ - pd.Timedelta(timing_metrics.first_scheduled_time).value - ], - dtype=self.output_dtype))) + np.asarray( + [pd.Timedelta(timing_metrics.first_scheduled_time).value], + dtype=self.output_dtype, + ), + ) + ) response.append( pb_utils.Tensor( "first_token_time_ns", np.asarray( [pd.Timedelta(timing_metrics.first_token_time).value], - dtype=self.output_dtype))) + dtype=self.output_dtype, + ), + ) + ) response.append( pb_utils.Tensor( "last_token_time_ns", np.asarray( [pd.Timedelta(timing_metrics.last_token_time).value], - dtype=self.output_dtype))) + dtype=self.output_dtype, + ), + ) + ) - #spec dec perf metrics per request + # spec dec perf metrics per request spec_dec_metrics = perf_metrics.speculative_decoding response.append( pb_utils.Tensor( "acceptance_rate", - np.asarray([spec_dec_metrics.acceptance_rate], - dtype=self.output_dtype))) + np.asarray( + [spec_dec_metrics.acceptance_rate], dtype=self.output_dtype + ), + ) + ) response.append( pb_utils.Tensor( "total_accepted_draft_tokens", - np.asarray([spec_dec_metrics.total_accepted_draft_tokens], - dtype=self.output_dtype))) + np.asarray( + [spec_dec_metrics.total_accepted_draft_tokens], + dtype=self.output_dtype, + ), + ) + ) response.append( pb_utils.Tensor( "total_draft_tokens", - np.asarray([spec_dec_metrics.total_draft_tokens], - dtype=self.output_dtype))) + np.asarray( + [spec_dec_metrics.total_draft_tokens], dtype=self.output_dtype + ), + ) + ) return pb_utils.InferenceResponse(output_tensors=response) @@ -731,8 +813,7 @@ def finalize(self): - Implementing `finalize` function is optional. """ self.logger.log_info("[trtllm] Issuing finalize to trtllm backend") - self._event_loop.call_soon_threadsafe( - self._llm_engine_shutdown_event.set) + self._event_loop.call_soon_threadsafe(self._llm_engine_shutdown_event.set) # Shutdown the event thread. if self._event_thread is not None: @@ -752,7 +833,6 @@ def finalize(self): # When using parallel tensors, the stub process may not shutdown due to # unreleased references, so manually run the garbage collector once. - self.logger.log_info( - "[trtllm] Running Garbage Collector on finalize...") + self.logger.log_info("[trtllm] Running Garbage Collector on finalize...") gc.collect() self.logger.log_info("[trtllm] Garbage Collector on finalize... done")