From 78c6e2970870374dd8ccbef27cc6a455b807bb17 Mon Sep 17 00:00:00 2001 From: gwarmstrong Date: Wed, 1 Jul 2026 11:43:19 -0700 Subject: [PATCH 1/5] inference: high-throughput async path + OpenAI request-routing keys Scale `ns generate` cleanly to high concurrency (5k+ in-flight) against OpenAI-compatible endpoints. At high concurrency the litellm/httpx dispatch was the GIL-bound bottleneck (~95% of one core in the async loop); this moves the hot path onto a C-extension-backed aiohttp transport and removes the per-completion write/postprocess serialization in the generate loop. Native AsyncOpenAI + aiohttp fast-path (model/base.py): - Routes non-streaming AND streaming chat through a native AsyncOpenAI client using the DefaultAioHttpClient (aiohttp) transport, bypassing litellm's httpx dispatch. Default-on for OpenAI-compatible providers, gated by a new SUPPORTS_NATIVE_OPENAI class attribute (openai/vllm/sglang); azure, gemini, megatron, and multimodal stay on litellm. Disable with NEMO_SKILLS_OPENAI_AIOHTTP=0. - Reasoning content is read from both .reasoning (newer vLLM) and .reasoning_content (older) -- litellm normalized this for us, but the native path sees vLLM's raw field. - Transient socket-error retries with backoff; extended httpx client TTL patch to survive long high-concurrency runs. Scalable async loop (inference/generate.py): - Bounded worker pool (memory O(in_flight)) with a rate-limited connection ramp (ramp_rate_per_sec) to avoid overflowing the kernel accept queue on prime. - Dedicated batched writer task instead of a per-completion fout lock. - Separate postprocess/eval processor pool (postprocess_workers) so a lane frees its slot the instant the HTTP response returns. - Installs uvloop when available (NEMO_SKILLS_DISABLE_UVLOOP=1 to opt out). Request-routing keys (model/base.py): - X-Request-Id: a per-request id for idempotency/tracing, kept stable across the OpenAI SDK retry loop; distinct calls get distinct ids. - prompt_cache_key: the standard OpenAI/vLLM prefix-cache hint, opt-in via cache_key=. Requests sharing a value hint a shared prompt prefix. Kept separate from X-Request-Id (shared vs per-request). A plain endpoint ignores both. Misc: - Dumps the full response.usage dict + reasoning-token breakdown. - Auto-passthrough of NEMO_SKILLS_* env vars to the inner python launched by the pipeline (pipeline/utils/cluster.py). - core/requirements.txt: openai -> openai[aiohttp] so the transport is installed. Tests: tests/test_prompt_cache_routing.py (routing-key idempotency/affinity decoupling) and tests/test_native_openai_path.py (reasoning extraction from both fields, stream/non-stream parsing, empty-choices guard, per-provider native eligibility, env-var default/opt-out). Signed-off-by: gwarmstrong --- core/requirements.txt | 2 +- nemo_skills/inference/generate.py | 273 ++++++++++++++++-- nemo_skills/inference/model/azure.py | 3 + nemo_skills/inference/model/base.py | 255 +++++++++++++++- nemo_skills/inference/model/openai.py | 10 + nemo_skills/inference/model/vllm.py | 5 + .../inference/model/vllm_multimodal.py | 5 + nemo_skills/pipeline/utils/cluster.py | 10 + tests/test_native_openai_path.py | 183 ++++++++++++ tests/test_prompt_cache_routing.py | 56 ++++ 10 files changed, 767 insertions(+), 35 deletions(-) create mode 100644 tests/test_native_openai_path.py create mode 100644 tests/test_prompt_cache_routing.py diff --git a/core/requirements.txt b/core/requirements.txt index 0d0fdd2c58..f88715ac30 100644 --- a/core/requirements.txt +++ b/core/requirements.txt @@ -27,7 +27,7 @@ litellm[caching]==1.83.14 math-verify[antlr4_9_3] mcp numpy -openai +openai[aiohttp] openpyxl>=3.1.0 pandas>=2.0.0 pyxlsb>=1.0.10 diff --git a/nemo_skills/inference/generate.py b/nemo_skills/inference/generate.py index 232be7d0ed..739f7d499d 100644 --- a/nemo_skills/inference/generate.py +++ b/nemo_skills/inference/generate.py @@ -15,11 +15,30 @@ import asyncio import json import logging +import os import random import shutil import subprocess import sys import time + +# uvloop is a drop-in replacement for asyncio's default event loop +# that's ~2-4x faster on CPU-bound asyncio workloads. ns generate's +# async_loop at high concurrency (5k+ lanes) becomes CPU-bound on +# the event loop's per-task scheduling overhead -- not on actual +# Python work in user code. Installing uvloop closes that gap so +# the bounded pool can refill at the rate the upstream can sustain +# instead of the rate one Python thread can scheduler-switch. +# +# Set NEMO_SKILLS_DISABLE_UVLOOP=1 to opt out (useful if a tool in +# the stack assumes selector loop semantics; rare). +if not os.environ.get("NEMO_SKILLS_DISABLE_UVLOOP"): + try: + import uvloop + + uvloop.install() + except ImportError: + pass from copy import deepcopy from dataclasses import asdict, field, is_dataclass from pathlib import Path @@ -125,6 +144,47 @@ class GenerationTaskConfig: # maximum number of concurrent requests to the server for the async loop # if sync loop is used, this is the batch size max_concurrent_requests: int = 512 + # Rate at which the async loop opens NEW HTTP connections during + # initial prime (target connections per second). The bounded + # worker pool would otherwise spawn `max_concurrent_requests` + # tasks in a tight loop, each opening a TCP connection nearly + # simultaneously -- the kernel accept queue on the receiving end + # (typically capped at net.core.somaxconn=4096) can briefly + # overflow and SYN-drop the excess. Limiting the ramp rate so + # each ~ramp_rate/4-connection burst is followed by a short pause + # keeps the accept queue draining cleanly. + # + # Default 4000/sec accounts for ns generate's per-task pre-work + # (format_prompt, build litellm request, semaphore acquires) -- + # the time from asyncio.create_task to "TCP request fired" is + # significant, so the configured task-creation rate has to be + # higher than the SYN rate you actually want at the kernel. + # Each burst is still ramp_rate/4 SYNs (1000 at default), well + # under net.core.somaxconn=4096. + # + # Set to 0 (or negative) to disable ramping and prime + # instantaneously. The transient-retry logic in + # BaseModel.generate_async would still cover any resulting SYN + # drops; the ramp just makes them not happen in the first place. + ramp_rate_per_sec: int = 4000 + # Number of background coroutines that postprocess + (optionally) + # evaluate completed datapoints. The async loop's bounded worker + # pool used to keep each lane "alive" (occupying a slot) until + # postprocess + evaluate + queue.put all completed -- which + # serialized those steps with the next HTTP refill. For + # reasoning-heavy workloads the per-lane post-HTTP work is small + # in absolute terms (~ms per record) but multiplied across + # thousands of lanes it can keep the bounded-pool slot occupied + # while real HTTP capacity sits idle. Moving postprocess + eval + # to a separate processor pool lets the lane exit the moment + # generate_async returns, so the bounded pool refills with the + # next datapoint immediately. + # + # Default 16 is enough headroom for synchronous postprocess (no + # eval) at >3k records/sec. If you enable an LLM-judge evaluator + # inside generate, raise this to match the eval concurrency you + # want. + postprocess_workers: int = 16 # chunk the dataset into equal sized parts and index into them num_chunks: int | None = None # if specified, will split the data into chunks and only generate for one chunk chunk_id: int | None = None # if specified, will index the specified chunk only @@ -409,6 +469,14 @@ def __init__(self, cfg: GenerationTaskConfig): # output_lock will be initialized when async_loop is called self.output_lock = None + # Three-stage pipeline queues (set by async_loop): + # _raw_queue: lanes -> processors (un-postprocessed outputs) + # _output_queue: processors -> writer (ready-to-write records) + # None outside of async_loop; _generate_and_save_datapoint + # falls back to inline postprocess + locked write in that + # case (used by entrypoints like prover.py). + self._raw_queue = None + self._output_queue = None def setup_prompt(self): if self.cfg.prompt_config is None: @@ -849,16 +917,23 @@ async def _generate_and_save_datapoint(self, data_point, all_data, fout, pbar): output["generation_end_time"] = end_time output["generation_time"] = end_time - start_time - await self.postprocess_single_output(output, data_point) - - # evaluate single-data point if requested and evaluator supports that - if self.should_run_evaluation and self.evaluator: - output = await self.evaluate_single_datapoint({**data_point, **output}) - - # Thread-safe output writing - async with self.output_lock: - self.dump_outputs([output], [data_point], fout) - pbar.update(1) + # Hand the raw (un-postprocessed) completion off to the + # processor pool. The bounded-pool slot in async_loop frees + # the moment this coroutine returns, so the next datapoint + # can begin its HTTP request immediately. Postprocess + + # evaluate run in the processor pool; the writer task drains + # processed records to disk in batches. + if self._raw_queue is not None: + await self._raw_queue.put((output, data_point)) + else: + # Fallback path (async_loop not entered through the + # processor/writer split). Run synchronously inline. + await self.postprocess_single_output(output, data_point) + if self.should_run_evaluation and self.evaluator: + output = await self.evaluate_single_datapoint({**data_point, **output}) + async with self.output_lock: + self.dump_outputs([output], [data_point], fout) + pbar.update(1) async def async_loop(self, data): """Async loop to generate generations using asyncio.""" @@ -881,7 +956,15 @@ async def async_loop(self, data): pbar = tqdm(total=len(remaining_data_points), desc="Remaining generations") - with open(self.cfg.output_file + "-async", "at", encoding="utf-8", buffering=1) as fout: + # Default Python buffering (8KB on most platforms). The + # previous `buffering=1` forced a kernel flush on every + # newline -- combined with the per-completion output_lock, + # that flush became the throughput bottleneck (~60-80/sec + # cap regardless of client concurrency). The writer task + # below issues an explicit fout.flush() after every batch + # drained from the queue, so durability is preserved without + # paying for a syscall per record. + with open(self.cfg.output_file + "-async", "at", encoding="utf-8") as fout: # Dump prefilled data first if len(prefilled_data_points) > 0: for output, data_point in zip(prefilled_outputs, prefilled_data_points): @@ -893,15 +976,165 @@ async def async_loop(self, data): async with self.output_lock: self.dump_outputs(prefilled_outputs, prefilled_data_points, fout) - # Create tasks for all remaining data points - tasks = [] - for data_point in remaining_data_points: - task = asyncio.create_task(self._generate_and_save_datapoint(data_point, data, fout, pbar)) - tasks.append(task) - - # Wait for all tasks to complete - if tasks: - await asyncio.gather(*tasks) + # Three-stage pipeline: + # + # lanes ── raw_queue ──> processors ── write_queue ──> writer + # + # `lanes` are the bounded-pool tasks below; each one does + # one HTTP call and exits the instant the response is + # parsed. `processors` is a small pool of background + # coroutines that run postprocess_single_output and + # (optionally) evaluate_single_datapoint -- both of which + # used to run INSIDE the lane, holding its bounded-pool + # slot for ~5-10ms past the HTTP return for postprocess + # alone, and seconds or more if an LLM-judge evaluator + # was attached. Decoupling them lets the lane refill + # instantly. `writer` batches processed records and + # flushes once per batch. + # + # Queue sizes are picked so producer puts never block in + # steady state. A `None` sentinel propagates through each + # stage on shutdown. + queue_cap = max(self.cfg.max_concurrent_requests * 2, 1024) + self._raw_queue: asyncio.Queue = asyncio.Queue(maxsize=queue_cap) + self._output_queue = asyncio.Queue(maxsize=queue_cap) + num_processors = max(1, int(self.cfg.postprocess_workers)) + + async def _processor(): + # Drain raw_queue, postprocess (+ evaluate), forward + # to output_queue. A `None` sentinel terminates this + # worker; the dispatcher below pushes one sentinel + # per worker. + while True: + item = await self._raw_queue.get() + if item is None: + return + output, data_point = item + try: + await self.postprocess_single_output(output, data_point) + if self.should_run_evaluation and self.evaluator: + output = await self.evaluate_single_datapoint({**data_point, **output}) + except Exception: + pos = data_point.get(self.cfg.async_position_key, "?") + LOG.exception("postprocess/evaluate failed for async_pos=%s", pos) + await self._output_queue.put((output, data_point)) + + async def _writer(): + # Drain in batches: pull one item (blocking), then + # opportunistically drain any others already in the + # queue without blocking, write them all, and flush + # once at the end. Throughput scales with completion + # rate, not per-record syscall cost. + while True: + item = await self._output_queue.get() + if item is None: + return + batch = [item] + # Drain anything that's already enqueued. + while not self._output_queue.empty(): + nxt = self._output_queue.get_nowait() + if nxt is None: + for o, _dp in batch: + fout.write(json.dumps(o) + "\n") + pbar.update(len(batch)) + fout.flush() + return + batch.append(nxt) + for o, _dp in batch: + fout.write(json.dumps(o) + "\n") + pbar.update(len(batch)) + fout.flush() + + processor_tasks = [asyncio.create_task(_processor()) for _ in range(num_processors)] + writer_task = asyncio.create_task(_writer()) + + # Bounded worker pool. At most `max_in_flight` tasks + # exist at any moment; as each one completes we spawn the + # next datapoint's task. Memory is O(in_flight), not + # O(total inputs), so the loop runs cleanly on 100k/1M+ + # input files without preallocating a coroutine frame + # for each. + # + # Was: tasks = [asyncio.create_task(...) for dp in data] + # That created ALL coroutines upfront. At a 100k input + # file and ~1-5 KB per coroutine frame (closure + Task + # bookkeeping) the upfront allocation was 100-500 MB -- + # avoidable memory pressure when co-located with other + # services. The outer semaphore inside each task bounded + # in-flight requests but didn't help that upfront cost. + # + # Each task is wrapped in a try/except so a single + # exception doesn't crash the whole gather (the old + # default behavior was first-exception-wins, killing + # all other in-flight datapoints). + max_in_flight = self.cfg.max_concurrent_requests + + async def _safe_one(dp): + try: + await self._generate_and_save_datapoint(dp, data, fout, pbar) + except Exception: + pos = dp.get(self.cfg.async_position_key, "?") + LOG.exception("datapoint async_pos=%s crashed", pos) + + data_iter = iter(remaining_data_points) + in_flight: set = set() + + # Prime the pool up to the concurrency cap, with rate- + # limited ramp so we don't open all max_in_flight TCP + # connections in a microsecond burst. + ramp_rate = self.cfg.ramp_rate_per_sec + if ramp_rate > 0: + # Spawn in batches sized so the implied pause is + # ~250ms -- short enough not to extend wall-time + # noticeably (target_depth=20k at 1000/s = 20s + # ramp), long enough for the kernel accept queue + # to drain between bursts. + ramp_batch = max(50, ramp_rate // 4) + ramp_interval = ramp_batch / ramp_rate + else: + ramp_batch = max_in_flight # instant prime + ramp_interval = 0.0 + + primed = 0 + for dp in data_iter: + in_flight.add(asyncio.create_task(_safe_one(dp))) + primed += 1 + if primed >= max_in_flight: + break + if ramp_interval > 0 and primed % ramp_batch == 0: + await asyncio.sleep(ramp_interval) + + # Refill as tasks complete. asyncio.wait with + # FIRST_COMPLETED returns as soon as ANY task finishes, + # which lets us start a fresh task immediately rather + # than waiting for the whole batch. + while in_flight: + done, in_flight = await asyncio.wait( + in_flight, + return_when=asyncio.FIRST_COMPLETED, + ) + for _ in done: + try: + dp = next(data_iter) + except StopIteration: + continue + in_flight.add(asyncio.create_task(_safe_one(dp))) + + # Shutdown sequence: cascade sentinels through the + # pipeline so every in-flight record is flushed. + # 1. Put one None on raw_queue per processor -- each + # processor exits after seeing its sentinel. + # 2. Wait for all processors. Any record put on + # output_queue by a processor is already there. + # 3. Put one None on output_queue -- writer drains + # whatever's queued ahead of it, then exits. + for _ in range(num_processors): + await self._raw_queue.put(None) + await asyncio.gather(*processor_tasks) + await self._output_queue.put(None) + await writer_task + self._raw_queue = None + self._output_queue = None pbar.close() diff --git a/nemo_skills/inference/model/azure.py b/nemo_skills/inference/model/azure.py index e097e4f80a..f3459db9ab 100644 --- a/nemo_skills/inference/model/azure.py +++ b/nemo_skills/inference/model/azure.py @@ -19,6 +19,9 @@ class AzureOpenAIModel(OpenAIModel): MODEL_PROVIDER = "azure" + # Azure uses a different endpoint layout + api-version/api-key auth that a + # plain AsyncOpenAI client can't target; keep it on litellm. + SUPPORTS_NATIVE_OPENAI = False def __init__( self, diff --git a/nemo_skills/inference/model/base.py b/nemo_skills/inference/model/base.py index 6079521f2f..3a7f3bccdf 100644 --- a/nemo_skills/inference/model/base.py +++ b/nemo_skills/inference/model/base.py @@ -15,6 +15,7 @@ import asyncio import logging import os +import uuid from enum import Enum from typing import Union @@ -73,6 +74,14 @@ class BaseModel: # Litellm provider name MODEL_PROVIDER = "openai" + # Whether this model may use the native AsyncOpenAI + aiohttp fast-path + # (see __init__). Only OpenAI-compatible endpoints that speak the plain + # /v1/chat/completions schema at self.base_url can opt in. Subclasses that + # talk to a non-OpenAI server (megatron), a differently-authed endpoint + # (azure), a non-OpenAI provider (gemini), or need litellm's request + # rewriting (multimodal audio) leave this False and stay on litellm. + SUPPORTS_NATIVE_OPENAI = False + def __init__( self, model: str, @@ -95,6 +104,17 @@ def __init__( output_dir: str | None = None, # Request tokenizer initialization independent of soft_fail require_tokenizer: bool = False, + # Max in-flight HTTP requests gated by an internal semaphore. + # The semaphore was originally added because earlier + # combinations of httpx + litellm would hang under truly + # unbounded concurrency. Modern httpx (>=0.27) handles this + # fine, so the default is now high enough to effectively be + # "no internal limit" -- the script-level concurrency knob + # (generate.py `max_concurrent_requests`) is the single + # user-facing cap. Override via $NEMO_SKILLS_MODEL_CONCURRENCY + # only if you have a specific reason to throttle BELOW the + # script-level cap. + concurrent_requests_limit: int = int(os.environ.get("NEMO_SKILLS_MODEL_CONCURRENCY", "65536")), ): self._tunnel = None self.model_name_or_path = model @@ -169,7 +189,69 @@ def __init__( litellm.aclient_session = httpx.AsyncClient(limits=httpx_limits) # Controlling concurrent requests using semaphore since large # concurrent requests result into httpx hanging - self.concurrent_semaphore = asyncio.Semaphore(2048) + self.concurrent_semaphore = asyncio.Semaphore(concurrent_requests_limit) + + # Native aiohttp fast-path (DEFAULT for OpenAI-compatible + # providers). The litellm.acompletion path goes through httpx, + # a pure-Python HTTP implementation -- at 5k+ concurrent + # connections it becomes the GIL-bound bottleneck (~95% of + # single-core CPU in the async loop, per py-spy). The OpenAI + # SDK natively supports an aiohttp transport via + # DefaultAioHttpClient. aiohttp's request/response parsing is + # C-extension backed, so the same workload runs at <5% CPU on + # the same hardware. + # + # Enabled by default for providers with SUPPORTS_NATIVE_OPENAI + # (openai, vllm, sglang). Both the non-streaming and streaming + # chat endpoints go through this path; text completions, the + # responses API, and non-OpenAI-shaped providers (gemini, + # azure, megatron, multimodal) still use litellm. Set + # NEMO_SKILLS_OPENAI_AIOHTTP=0 to force everything back onto + # litellm/httpx. + self._async_openai_client = None + native_enabled = os.environ.get("NEMO_SKILLS_OPENAI_AIOHTTP", "1") != "0" + if native_enabled and self.SUPPORTS_NATIVE_OPENAI: + try: + from openai import AsyncOpenAI, DefaultAioHttpClient + + # openai SDK defaults the aiohttp connector to + # max_connections=1000, max_keepalive_connections=100, + # which silently caps the lane count at high + # concurrency (8k lanes see only ~1000 open TCP + # connections; the rest stall in the aiohttp connector + # queue). Also, httpx-aiohttp's transport translates + # httpx.Limits(max_connections=None) to "omit the + # `limit` kwarg," which makes aiohttp fall back to ITS + # default of 100 -- worse than the openai SDK + # default. We therefore pass an explicit large + # integer (NEMO_SKILLS_OPENAI_AIOHTTP_LIMIT, default + # 65536) instead of None. + aiohttp_limit = int(os.environ.get("NEMO_SKILLS_OPENAI_AIOHTTP_LIMIT", "65536")) + self._async_openai_client = AsyncOpenAI( + api_key=api_key, + base_url=self.base_url, + http_client=DefaultAioHttpClient( + limits=httpx.Limits( + max_connections=aiohttp_limit, + max_keepalive_connections=aiohttp_limit, + ), + ), + max_retries=0, # we own retries; SDK's retry would + # rebuild the request and double-count + # concurrent_semaphore holds. + ) + except Exception as e: + # ImportError if openai[aiohttp] isn't installed, but the + # DefaultAioHttpClient constructor can also raise other errors + # when the aiohttp extra is only partially present. Degrade to + # litellm/httpx rather than failing model construction. + LOG.warning( + "Native AsyncOpenAI aiohttp fast-path unavailable (%s: %s); " + "falling back to litellm/httpx. Install openai[aiohttp] to enable it.", + type(e).__name__, + e, + ) + self._async_openai_client = None def _get_api_key(self, api_key: str | None, api_key_env_var: str | None, base_url: str) -> str | None: if api_key: # explicit cmd argument always takes precedence @@ -185,8 +267,11 @@ def _get_api_key(self, api_key: str | None, api_key_env_var: str | None, base_ur return api_key def __del__(self): - if self._tunnel: - self._tunnel.stop() + # getattr guard: __del__ can fire on an instance whose __init__ raised + # (or never ran), before _tunnel was set. + tunnel = getattr(self, "_tunnel", None) + if tunnel: + tunnel.stop() def _maybe_apply_stop_phrase_removal( self, result: dict, remove_stop_phrases: bool, stop_phrases: list[str] | None @@ -239,6 +324,27 @@ def _build_completion_request_params(self, **kwargs) -> dict: def _build_responses_request_params(self, **kwargs) -> dict: raise NotImplementedError("Responses completion is not not supported or implemented for this model.") + @staticmethod + def _apply_routing_keys(request_params: dict, cache_key: str | None) -> None: + """Attach optional, provider-agnostic request-routing keys. Plain + OpenAI/vLLM/gemini endpoints ignore (or echo) them. + + * `X-Request-Id` (extra_headers) = per-request id for idempotency / + tracing. Kept stable across the OpenAI SDK's internal retry loop so + a retry is identifiable as the same logical call; distinct calls + (incl. parallel samples) get distinct ids. + * `prompt_cache_key` (opt-in via `cache_key`) = the standard + OpenAI/vLLM prefix-cache hint. Requests that share a value are + hinting they share a common prompt prefix and can reuse a warm + prefix cache. Omitted when cache_key is None. + + Kept separate on purpose: prompt_cache_key is meant to be SHARED across + distinct requests, so it must not double as the per-request idempotency id. + """ + request_params.setdefault("extra_headers", {}).setdefault("X-Request-Id", str(uuid.uuid4())) + if cache_key is not None: + request_params["prompt_cache_key"] = cache_key + @with_context_retry async def generate_async( self, @@ -261,6 +367,7 @@ async def generate_async( include_response: bool = False, extra_body: dict = None, response_format=None, + cache_key: str | None = None, ) -> dict: if endpoint_type is None: # Infering completion type from prompt @@ -289,9 +396,23 @@ async def generate_async( "response_format": response_format, } - # TODO: remove this after we no longer use gpt-oss or it's fixed in vllm + # Two separate retry budgets: + # - bad_request_count is for the gpt-oss "output messages" + # bug; small, app-level. Kept as-is until we drop that + # model (or vLLM fixes it upstream). + # - transient_count handles socket-level failures that the + # raw HTTP path can hit under burst load (kernel SYN-drops + # when the accept queue overflows, brief upstream + # unavailability during worker turnover, etc). litellm's + # own max_retries covers HTTP 429/5xx but NOT these + # connection-class exceptions, so a single ramp-up burst + # into the upstream could otherwise hard-fail individual + # datapoints. Retries here turn those into invisible + # ~50-200ms re-tries. max_retries = 2 retry_count = 0 + max_transient_retries = int(os.environ.get("NEMO_SKILLS_TRANSIENT_RETRIES", "3")) + transient_count = 0 async with self.concurrent_semaphore: while retry_count <= max_retries: @@ -299,7 +420,38 @@ async def generate_async( if endpoint_type == EndpointType.chat: assert isinstance(prompt, list), "Chat completion requests must be a list of messages." request_params = self._build_chat_request_params(messages=prompt, stream=stream, **kwargs) - response = await litellm.acompletion(**request_params, **self.litellm_kwargs) + self._apply_routing_keys(request_params, cache_key) + # Fast path: native AsyncOpenAI + aiohttp transport, + # set up in __init__ for OpenAI-compatible providers + # (disable with NEMO_SKILLS_OPENAI_AIOHTTP=0). Skips + # litellm's httpx-based dispatch which is the dominant + # single-core CPU consumer at high concurrency. Both + # streaming and non-streaming chat are routed here; the + # responses API, text completions, and non-OpenAI + # providers still fall through to litellm. + use_native = self._async_openai_client is not None + if use_native: + native_params = { + k: v + for k, v in request_params.items() + # litellm-only knobs we strip before + # handing the body to AsyncOpenAI. The + # SDK validates against the OpenAI + # schema and rejects unknown top-level + # fields. + if k not in ("allowed_openai_params",) + } + # AsyncOpenAI needs `model` as a top-level + # kwarg; litellm consumes it via + # self.litellm_kwargs (with the `openai/` + # provider prefix). Use the raw model name + # here. + native_params.setdefault("model", self.model_name_or_path) + response = await self._async_openai_client.chat.completions.create( + **native_params, + ) + else: + response = await litellm.acompletion(**request_params, **self.litellm_kwargs) if stream: result = self._stream_chat_chunks_async(response) else: @@ -309,6 +461,7 @@ async def generate_async( elif endpoint_type == EndpointType.text: assert isinstance(prompt, str), "Text completion requests must be a string." request_params = self._build_completion_request_params(prompt=prompt, stream=stream, **kwargs) + self._apply_routing_keys(request_params, cache_key) response = await litellm.atext_completion(**request_params, **self.litellm_kwargs) if stream: result = self._stream_completion_chunks_async(response) @@ -319,6 +472,7 @@ async def generate_async( elif endpoint_type == EndpointType.responses: assert isinstance(prompt, list), "Responses completion requests must be a list." request_params = self._build_responses_request_params(input=prompt, stream=stream, **kwargs) + self._apply_routing_keys(request_params, cache_key) response = await litellm.aresponses(**request_params, **self.litellm_kwargs) if stream: raise NotImplementedError("Streaming responses is not supported yet.") @@ -348,6 +502,41 @@ async def generate_async( } else: raise e + except ( + httpx.ConnectError, + httpx.ReadError, + httpx.RemoteProtocolError, + httpx.PoolTimeout, + httpx.NetworkError, + openai.APIConnectionError, + ConnectionError, + ConnectionResetError, + ) as e: + # Transient socket-level failure. Common cause at + # high concurrency: kernel SYN-drop during ramp-up + # when the upstream accept queue briefly overflows. + # Backoff doubles each retry (50ms, 100ms, 200ms) + # to let the accept queue drain. + if transient_count < max_transient_retries: + transient_count += 1 + backoff = 0.05 * (2 ** (transient_count - 1)) + LOG.warning( + "transient connect error, retry %d/%d after %.0fms: %s: %s", + transient_count, + max_transient_retries, + backoff * 1000, + type(e).__name__, + e, + ) + await asyncio.sleep(backoff) + continue + LOG.error( + "transient connect error after %d retries, propagating: %s: %s", + max_transient_retries, + type(e).__name__, + e, + ) + raise return result @@ -384,6 +573,23 @@ def _parse_completion_response( return result + @staticmethod + def _extract_reasoning(message_or_delta) -> str | None: + """Return reasoning text from an OpenAI chat message or streaming delta, + tolerating both field names. + + vLLM historically exposed reasoning as `reasoning_content`; newer vLLM + versions expose it as `reasoning`. litellm normalizes these to + `reasoning_content` for us, but the native AsyncOpenAI fast-path sees + vLLM's raw field unchanged -- so we accept either here. Prefers + `reasoning_content` when both are present. + """ + for attr in ("reasoning_content", "reasoning"): + value = getattr(message_or_delta, attr, None) + if value: + return value + return None + def _parse_chat_completion_response(self, response, include_response: bool = False, **kwargs) -> dict: choice = response.choices[0] output = choice.message.content @@ -395,9 +601,28 @@ def _parse_chat_completion_response(self, response, include_response: bool = Fal elif getattr(response.usage, "input_tokens", None) is not None: result["num_input_tokens"] = response.usage.input_tokens - # Add reasoning_content if available - if hasattr(choice.message, "reasoning_content") and choice.message.reasoning_content: - result["reasoning_content"] = choice.message.reasoning_content + # Add reasoning_content if available (accepts .reasoning or .reasoning_content) + reasoning = self._extract_reasoning(choice.message) + if reasoning: + result["reasoning_content"] = reasoning + + # Dump the full usage object verbatim into the output. Keeping + # None-valued fields (rather than excluding them) preserves the + # distinction between "server explicitly sent null" and "field + # was absent" — both can be meaningful (e.g. some providers + # send prompt_tokens_details=null when there's no cache hit). + try: + result["usage"] = response.usage.model_dump() + except AttributeError: + # Older or non-pydantic usage objects: fall back to dict(). + try: + result["usage"] = dict(response.usage) + except Exception: + result["usage"] = { + "prompt_tokens": getattr(response.usage, "prompt_tokens", None), + "completion_tokens": getattr(response.usage, "completion_tokens", None), + "total_tokens": getattr(response.usage, "total_tokens", None), + } # Extract detailed token breakdown for reasoning models if available if hasattr(response.usage, "completion_tokens_details") and response.usage.completion_tokens_details: @@ -463,12 +688,8 @@ def _process_chat_chunk(self, chunk): """Process a single chat chunk and return data to yield.""" if hasattr(chunk.choices[0], "delta"): cur_delta = chunk.choices[0].delta.content - # Check for reasoning_content in delta - reasoning_delta = ( - getattr(chunk.choices[0].delta, "reasoning_content", None) - if hasattr(chunk.choices[0].delta, "reasoning_content") - else None - ) + # Check for reasoning in delta (accepts .reasoning or .reasoning_content) + reasoning_delta = self._extract_reasoning(chunk.choices[0].delta) tool_calls_delta = getattr(chunk.choices[0].delta, "tool_calls", None) else: cur_delta = chunk.choices[0].text @@ -561,6 +782,12 @@ def _serialize_output(self, response): async def _stream_chat_chunks_async(self, response): async for chunk in response: + # Some servers emit chunks with no choices (e.g. a trailing + # usage-only chunk when stream_options.include_usage is set, or + # keepalive chunks). They carry no delta -- skip them so + # _process_chat_chunk's choices[0] access stays safe. + if not getattr(chunk, "choices", None): + continue results = self._process_chat_chunk(chunk) for result in results: yield result diff --git a/nemo_skills/inference/model/openai.py b/nemo_skills/inference/model/openai.py index f8d4dc7bde..68f0eb4070 100644 --- a/nemo_skills/inference/model/openai.py +++ b/nemo_skills/inference/model/openai.py @@ -20,6 +20,10 @@ class OpenAIModel(BaseModel): + # OpenAI-compatible /v1/chat/completions endpoint -- eligible for the + # native AsyncOpenAI + aiohttp fast-path (see BaseModel.__init__). + SUPPORTS_NATIVE_OPENAI = True + def __init__( self, host: str = "127.0.0.1", @@ -129,6 +133,12 @@ def _build_chat_request_params( "stream": stream, "tools": tools, "response_format": response_format, + # NOTE: idempotency (X-Request-Id) and prefix-cache affinity + # (prompt_cache_key, opt-in via generate_async(cache_key=...)) are + # attached centrally in BaseModel._apply_routing_keys. We deliberately + # do NOT set a per-request-UUID prompt_cache_key here: that would give + # every request a unique affinity key and defeat cross-request prefix + # caching. } if self._is_reasoning_model(self.model): diff --git a/nemo_skills/inference/model/vllm.py b/nemo_skills/inference/model/vllm.py index e2c8942d29..aa666540ef 100644 --- a/nemo_skills/inference/model/vllm.py +++ b/nemo_skills/inference/model/vllm.py @@ -91,6 +91,11 @@ def process_image_content(content: list | str | None, data_dir: str = "") -> lis class VLLMModel(BaseModel): + # vLLM (and sglang, which subclasses this) serve an OpenAI-compatible + # /v1/chat/completions API -- eligible for the native AsyncOpenAI + + # aiohttp fast-path (see BaseModel.__init__). + SUPPORTS_NATIVE_OPENAI = True + def __init__(self, **kwargs): super().__init__(**kwargs) diff --git a/nemo_skills/inference/model/vllm_multimodal.py b/nemo_skills/inference/model/vllm_multimodal.py index 4b8006abe5..2924d393de 100644 --- a/nemo_skills/inference/model/vllm_multimodal.py +++ b/nemo_skills/inference/model/vllm_multimodal.py @@ -69,6 +69,11 @@ class VLLMMultimodalModel(VLLMModel): ) """ + # Audio/VLM request rewriting and response post-processing go through + # litellm; keep this off the native fast-path until it's verified for + # multimodal payloads. Overrides VLLMModel's True. + SUPPORTS_NATIVE_OPENAI = False + def __init__( self, model: str | None = None, diff --git a/nemo_skills/pipeline/utils/cluster.py b/nemo_skills/pipeline/utils/cluster.py index d525746f3e..f57fee58ba 100644 --- a/nemo_skills/pipeline/utils/cluster.py +++ b/nemo_skills/pipeline/utils/cluster.py @@ -215,6 +215,16 @@ def get_env_variables(cluster_config): "HF_TOKEN", "NGC_API_KEY", } + # Auto-passthrough for NEMO_SKILLS_* runtime tuning vars set in + # the user's shell. Without this, opt-in env knobs like + # NEMO_SKILLS_OPENAI_AIOHTTP, NEMO_SKILLS_DISABLE_UVLOOP, and + # NEMO_SKILLS_TRANSIENT_RETRIES are stripped when the launcher + # spawns the inner python -- they only live in the wrapper + # shell. Scoping by prefix keeps the surface tight (no risk of + # accidentally exporting unrelated user env). + for name in os.environ: + if name.startswith("NEMO_SKILLS_") and name not in optional_env_vars_to_add: + optional_env_vars_to_add.add(name) default_factories = { "HF_TOKEN": lambda: str(token) if (token := get_token()) else "", } diff --git a/tests/test_native_openai_path.py b/tests/test_native_openai_path.py new file mode 100644 index 0000000000..f711f1d858 --- /dev/null +++ b/tests/test_native_openai_path.py @@ -0,0 +1,183 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# 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. +"""Unit tests for the native AsyncOpenAI + aiohttp fast-path in BaseModel. + +Covers the three feature-completion items: + 1. Native path is the DEFAULT for OpenAI-compatible providers and can be + disabled with NEMO_SKILLS_OPENAI_AIOHTTP=0; non-OpenAI providers never + use it. + 2. Reasoning content is read from BOTH `.reasoning` and `.reasoning_content` + (litellm normalizes this for us; the native path sees vLLM's raw field). + 3. Streaming is routed through the native path and tolerates chunks with no + choices. + +These are pure unit tests -- no server or network access. +""" + +import asyncio +import types + +import openai +import pytest + +from nemo_skills.inference.model.azure import AzureOpenAIModel +from nemo_skills.inference.model.base import BaseModel +from nemo_skills.inference.model.gemini import GeminiModel +from nemo_skills.inference.model.megatron import MegatronModel +from nemo_skills.inference.model.openai import OpenAIModel +from nemo_skills.inference.model.sglang import SGLangModel +from nemo_skills.inference.model.vllm import VLLMModel +from nemo_skills.inference.model.vllm_multimodal import VLLMMultimodalModel + + +# --------------------------------------------------------------------------- +# Item 2: reasoning extraction accepts both field names +# --------------------------------------------------------------------------- +def test_extract_reasoning_from_reasoning_content(): + obj = types.SimpleNamespace(reasoning_content="thinking") + assert BaseModel._extract_reasoning(obj) == "thinking" + + +def test_extract_reasoning_from_reasoning(): + # New vLLM field name; only `.reasoning` present. + obj = types.SimpleNamespace(reasoning="thinking") + assert BaseModel._extract_reasoning(obj) == "thinking" + + +def test_extract_reasoning_prefers_reasoning_content(): + obj = types.SimpleNamespace(reasoning_content="primary", reasoning="secondary") + assert BaseModel._extract_reasoning(obj) == "primary" + + +def test_extract_reasoning_absent_or_empty(): + assert BaseModel._extract_reasoning(types.SimpleNamespace()) is None + assert BaseModel._extract_reasoning(types.SimpleNamespace(reasoning="", reasoning_content="")) is None + + +class _FakeUsage: + completion_tokens = 5 + prompt_tokens = 3 + + def model_dump(self): + return {"completion_tokens": 5, "prompt_tokens": 3, "total_tokens": 8} + + +class _FakeChoice: + def __init__(self, message): + self.message = message + self.finish_reason = "stop" + + def model_dump(self): + return {"message": {"role": "assistant", "content": self.message.content}} + + +class _FakeResponse: + def __init__(self, message): + self.choices = [_FakeChoice(message)] + self.usage = _FakeUsage() + + +def test_parse_chat_completion_reads_reasoning_field(): + # A response carrying ONLY `.reasoning` (as newer vLLM emits over the raw + # OpenAI schema) must still populate reasoning_content in the result. + inst = BaseModel.__new__(BaseModel) + message = types.SimpleNamespace(content="the answer", reasoning="the thinking") + result = inst._parse_chat_completion_response(_FakeResponse(message)) + assert result["generation"] == "the answer" + assert result["reasoning_content"] == "the thinking" + assert result["num_generated_tokens"] == 5 + + +# --------------------------------------------------------------------------- +# Item 3: streaming chunk parsing (reasoning + empty-choices guard) +# --------------------------------------------------------------------------- +def test_process_chat_chunk_reads_reasoning_field(): + inst = BaseModel.__new__(BaseModel) + delta = types.SimpleNamespace(content="hi", reasoning="step") + chunk = types.SimpleNamespace(choices=[types.SimpleNamespace(delta=delta, finish_reason=None)]) + [result] = inst._process_chat_chunk(chunk) + assert result["generation"] == "hi" + assert result["reasoning_content"] == "step" + + +def test_stream_skips_empty_choices_chunk(): + inst = BaseModel.__new__(BaseModel) + + async def fake_stream(): + # usage-only / keepalive chunk with no choices -- must be skipped + yield types.SimpleNamespace(choices=[]) + yield types.SimpleNamespace( + choices=[types.SimpleNamespace(delta=types.SimpleNamespace(content="hello"), finish_reason=None)] + ) + + async def collect(): + return [r async for r in inst._stream_chat_chunks_async(fake_stream())] + + results = asyncio.run(collect()) + assert results == [{"generation": "hello"}] + + +# --------------------------------------------------------------------------- +# Item 1: which providers are eligible for the native fast-path +# --------------------------------------------------------------------------- +def test_native_eligibility_by_provider(): + assert BaseModel.SUPPORTS_NATIVE_OPENAI is False + # OpenAI-compatible endpoints opt in. + assert OpenAIModel.SUPPORTS_NATIVE_OPENAI is True + assert VLLMModel.SUPPORTS_NATIVE_OPENAI is True + assert SGLangModel.SUPPORTS_NATIVE_OPENAI is True + # Non-OpenAI-shaped endpoints stay on litellm. + assert AzureOpenAIModel.SUPPORTS_NATIVE_OPENAI is False + assert GeminiModel.SUPPORTS_NATIVE_OPENAI is False + assert MegatronModel.SUPPORTS_NATIVE_OPENAI is False + assert VLLMMultimodalModel.SUPPORTS_NATIVE_OPENAI is False + + +class _DummyAioHttpClient: + def __init__(self, *args, **kwargs): + pass + + +class _DummyAsyncOpenAI: + def __init__(self, *args, **kwargs): + self.kwargs = kwargs + + +@pytest.fixture +def _mock_openai_client(monkeypatch): + # Avoid needing the real openai[aiohttp] transport / a live endpoint. + monkeypatch.setattr(openai, "AsyncOpenAI", _DummyAsyncOpenAI, raising=False) + monkeypatch.setattr(openai, "DefaultAioHttpClient", _DummyAioHttpClient, raising=False) + + +def test_native_client_on_by_default(monkeypatch, _mock_openai_client): + monkeypatch.delenv("NEMO_SKILLS_OPENAI_AIOHTTP", raising=False) + model = VLLMModel(model="test-model", base_url="http://localhost:9999/v1") + assert isinstance(model._async_openai_client, _DummyAsyncOpenAI) + + +def test_native_client_opt_out(monkeypatch, _mock_openai_client): + monkeypatch.setenv("NEMO_SKILLS_OPENAI_AIOHTTP", "0") + model = VLLMModel(model="test-model", base_url="http://localhost:9999/v1") + assert model._async_openai_client is None + + +def test_native_client_off_for_ineligible_provider(monkeypatch, _mock_openai_client): + monkeypatch.delenv("NEMO_SKILLS_OPENAI_AIOHTTP", raising=False) + + class _NoNative(VLLMModel): + SUPPORTS_NATIVE_OPENAI = False + + model = _NoNative(model="test-model", base_url="http://localhost:9999/v1") + assert model._async_openai_client is None diff --git a/tests/test_prompt_cache_routing.py b/tests/test_prompt_cache_routing.py new file mode 100644 index 0000000000..d5af83b12d --- /dev/null +++ b/tests/test_prompt_cache_routing.py @@ -0,0 +1,56 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# 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. +"""Unit tests for the request-routing keys attached by BaseModel. + +Idempotency (X-Request-Id) and prefix-cache affinity (prompt_cache_key) are +DECOUPLED: X-Request-Id is per-request (unique per call, safe for retries), +prompt_cache_key is opt-in and SHARED across requests with a common prefix. +""" + +from nemo_skills.inference.model.base import BaseModel + + +def test_idempotency_header_always_set_affinity_opt_in(): + # No cache_key → X-Request-Id present (idempotency), prompt_cache_key OMITTED. + params = {} + BaseModel._apply_routing_keys(params, None) + assert "prompt_cache_key" not in params + xid = params["extra_headers"]["X-Request-Id"] + assert isinstance(xid, str) and len(xid) >= 16 + + +def test_cache_key_sets_prompt_cache_key(): + params = {} + BaseModel._apply_routing_keys(params, "conv-1") + assert params["prompt_cache_key"] == "conv-1" + assert params["extra_headers"]["X-Request-Id"] + + +def test_shared_affinity_distinct_request_ids(): + # The branching/multi-sample invariant: same cache_key (shared prefix), but + # each call gets a DISTINCT X-Request-Id (unique per request). + p1, p2 = {}, {} + BaseModel._apply_routing_keys(p1, "conv-1") + BaseModel._apply_routing_keys(p2, "conv-1") + assert p1["prompt_cache_key"] == p2["prompt_cache_key"] == "conv-1" + assert p1["extra_headers"]["X-Request-Id"] != p2["extra_headers"]["X-Request-Id"] + + +def test_preserves_existing_headers_and_does_not_override_request_id(): + # setdefault semantics: a caller-set X-Request-Id (e.g. for app-level retry + # idempotency) and other headers survive. + params = {"extra_headers": {"X-Custom": "v", "X-Request-Id": "preset"}} + BaseModel._apply_routing_keys(params, None) + assert params["extra_headers"]["X-Custom"] == "v" + assert params["extra_headers"]["X-Request-Id"] == "preset" From c4de6d34ba023c13a32fb7d3967526d2353e5a16 Mon Sep 17 00:00:00 2001 From: gwarmstrong Date: Wed, 1 Jul 2026 13:32:11 -0700 Subject: [PATCH 2/5] inference: address async-loop data-integrity and native-path review findings Follow-up fixes from code review of the high-throughput async path. async_loop data integrity (generate.py): - On a datapoint that fails generation, emit a placeholder record (routed through the postprocess queue so it is position-stamped) instead of dropping it silently. A dropped position made restore_async_order either IndexError at finalization (a non-last position missing) or silently drop/misalign rows. - On a postprocess/eval failure, stamp async_position (and a postprocess_error marker) onto the record so it stays placeable even if postprocess raised before merging the original datapoint. - restore_async_order sizes by max(position)+1 and tolerates a missing position (pop-with-default + skip) instead of assuming exactly-one-line-per-position. - The writer flushes through self.dump_outputs (the overridable hook the prefilled path uses) via a single helper, so subclass serialization applies to async records and the duplicated write path is gone. Native fast-path (model/base.py): - Give the native AsyncOpenAI client the same max_retries as the litellm path (was 0). The SDK retries HTTP 429/5xx (which our transient-error except does not catch); leaving it at 0 dropped those under load on the default path. - Drop both litellm-only keys (named _LITELLM_ONLY_CHAT_PARAMS) and None-valued params from the native request body: litellm omits None fields, but the SDK serializes an explicit None as JSON null, which some servers reject. Env passthrough (pipeline/utils/cluster.py): - Keep forwarding the NEMO_SKILLS_ prefix, but skip the SSH-tunnel vars (NEMO_SKILLS_SSH_SERVER / NEMO_SKILLS_SSH_KEY_PATH): forwarding a launcher's tunnel config into a co-located cluster job makes model init open a tunnel it shouldn't and fail on a host key path that doesn't exist in the container. Tests: tests/test_async_loop_integrity.py (reorder tolerance), plus native None/litellm-only param stripping and the env-passthrough allow/deny behavior. Signed-off-by: gwarmstrong --- nemo_skills/inference/generate.py | 61 +++++++++++++++----- nemo_skills/inference/model/base.py | 34 +++++++++--- nemo_skills/pipeline/utils/cluster.py | 33 ++++++++--- tests/test_async_loop_integrity.py | 80 +++++++++++++++++++++++++++ tests/test_native_openai_path.py | 52 +++++++++++++++++ tests/test_pipeline_utils.py | 34 ++++++++++++ 6 files changed, 264 insertions(+), 30 deletions(-) create mode 100644 tests/test_async_loop_integrity.py diff --git a/nemo_skills/inference/generate.py b/nemo_skills/inference/generate.py index 739f7d499d..68e1027c23 100644 --- a/nemo_skills/inference/generate.py +++ b/nemo_skills/inference/generate.py @@ -1014,11 +1014,28 @@ async def _processor(): await self.postprocess_single_output(output, data_point) if self.should_run_evaluation and self.evaluator: output = await self.evaluate_single_datapoint({**data_point, **output}) - except Exception: + except Exception as e: pos = data_point.get(self.cfg.async_position_key, "?") LOG.exception("postprocess/evaluate failed for async_pos=%s", pos) + # postprocess_single_output merges data_point (which carries + # async_position_key) into output only partway through, so a + # failure before that merge would write a record with no + # position key -- and restore_async_order would then KeyError + # on the whole file. Stamp the position and an error marker so + # the record is still placeable and the failure is visible. + if self.cfg.async_position_key in data_point: + output[self.cfg.async_position_key] = data_point[self.cfg.async_position_key] + output.setdefault("postprocess_error", repr(e)) await self._output_queue.put((output, data_point)) + def _flush_batch(batch): + # Route through self.dump_outputs (the overridable hook the + # prefilled path also uses) rather than an inline json.dumps, so + # a subclass customizing serialization applies to async records too. + self.dump_outputs([o for o, _dp in batch], [dp for _o, dp in batch], fout) + pbar.update(len(batch)) + fout.flush() + async def _writer(): # Drain in batches: pull one item (blocking), then # opportunistically drain any others already in the @@ -1030,20 +1047,17 @@ async def _writer(): if item is None: return batch = [item] + stop = False # Drain anything that's already enqueued. while not self._output_queue.empty(): nxt = self._output_queue.get_nowait() if nxt is None: - for o, _dp in batch: - fout.write(json.dumps(o) + "\n") - pbar.update(len(batch)) - fout.flush() - return + stop = True + break batch.append(nxt) - for o, _dp in batch: - fout.write(json.dumps(o) + "\n") - pbar.update(len(batch)) - fout.flush() + _flush_batch(batch) + if stop: + return processor_tasks = [asyncio.create_task(_processor()) for _ in range(num_processors)] writer_task = asyncio.create_task(_writer()) @@ -1072,9 +1086,17 @@ async def _writer(): async def _safe_one(dp): try: await self._generate_and_save_datapoint(dp, data, fout, pbar) - except Exception: + except Exception as e: pos = dp.get(self.cfg.async_position_key, "?") LOG.exception("datapoint async_pos=%s crashed", pos) + # Emit a placeholder so this datapoint's position still + # produces exactly one output line. Without it, + # restore_async_order sees a hole and either IndexErrors (a + # non-last position missing) or silently drops/misaligns rows. + # Routed through raw_queue so postprocess stamps the position + # and merges the original fields, same as a normal record. + if self._raw_queue is not None: + await self._raw_queue.put(({"generation": "", "generation_error": repr(e)}, dp)) data_iter = iter(remaining_data_points) in_flight: set = set() @@ -1145,13 +1167,26 @@ def restore_async_order(self): with open(self.cfg.output_file + "-async", "rt", encoding="utf-8") as fin: generations = [json.loads(line) for line in fin] - ordered_generations = [None] * len(generations) + # Size by the max position, not the line count: those differ if any + # record was dropped, and indexing a count-sized list by original + # position would then IndexError. pop(default) + skip means a + # missing/position-less record can never crash finalization. The + # pipeline emits a placeholder per datapoint so holes should not occur; + # this is a safety net. + key = self.cfg.async_position_key + positions = [gen[key] for gen in generations if key in gen] + ordered_generations = [None] * ((max(positions) + 1) if positions else 0) for gen_dict in generations: - async_pos = gen_dict.pop(self.cfg.async_position_key) + async_pos = gen_dict.pop(key, None) + if async_pos is None: + LOG.warning("async record missing %s, dropping from reorder", key) + continue ordered_generations[async_pos] = gen_dict with open(self.cfg.output_file, "wt", encoding="utf-8") as fout: for gen_dict in ordered_generations: + if gen_dict is None: + continue fout.write(json.dumps(gen_dict) + "\n") Path(self.cfg.output_file + "-async").unlink() diff --git a/nemo_skills/inference/model/base.py b/nemo_skills/inference/model/base.py index 3a7f3bccdf..f2e6144833 100644 --- a/nemo_skills/inference/model/base.py +++ b/nemo_skills/inference/model/base.py @@ -82,6 +82,12 @@ class BaseModel: # rewriting (multimodal audio) leave this False and stay on litellm. SUPPORTS_NATIVE_OPENAI = False + # Keys that litellm understands but AsyncOpenAI's strict schema rejects as + # unknown top-level fields; stripped before handing the body to the native + # client. Keep this in sync with any litellm-only knob a _build_chat_request_params + # implementation may emit. + _LITELLM_ONLY_CHAT_PARAMS = frozenset({"allowed_openai_params"}) + def __init__( self, model: str, @@ -236,9 +242,16 @@ def __init__( max_keepalive_connections=aiohttp_limit, ), ), - max_retries=0, # we own retries; SDK's retry would - # rebuild the request and double-count - # concurrent_semaphore holds. + # Mirror the litellm path's retry budget (self.litellm_kwargs + # also passes max_retries). The SDK's own retry loop honors + # Retry-After / backoff for HTTP 429 and 5xx -- error classes + # that our app-level transient-retry except clause does NOT + # catch (it only handles connection-class errors). Leaving this + # at 0 silently dropped 429/5xx under load on the native path, + # which is the DEFAULT for openai/vllm/sglang. SDK retries run + # inside our concurrent_semaphore hold, so one logical request + # still counts as one in-flight slot. + max_retries=max_retries, ) except Exception as e: # ImportError if openai[aiohttp] isn't installed, but the @@ -431,15 +444,18 @@ async def generate_async( # providers still fall through to litellm. use_native = self._async_openai_client is not None if use_native: + # Strip (a) litellm-only knobs the strict SDK rejects + # and (b) None-valued params. litellm OMITS None-valued + # fields from the wire body, but AsyncOpenAI serializes + # an explicit None as JSON `null` -- e.g. the vllm builder + # emits response_format/tools/stop/seed=None, and some + # OpenAI-compatible servers reject `"response_format": null`. + # Dropping None here keeps the native body byte-compatible + # with what the litellm path sent. native_params = { k: v for k, v in request_params.items() - # litellm-only knobs we strip before - # handing the body to AsyncOpenAI. The - # SDK validates against the OpenAI - # schema and rejects unknown top-level - # fields. - if k not in ("allowed_openai_params",) + if k not in self._LITELLM_ONLY_CHAT_PARAMS and v is not None } # AsyncOpenAI needs `model` as a top-level # kwarg; litellm consumes it via diff --git a/nemo_skills/pipeline/utils/cluster.py b/nemo_skills/pipeline/utils/cluster.py index f57fee58ba..4237da23e7 100644 --- a/nemo_skills/pipeline/utils/cluster.py +++ b/nemo_skills/pipeline/utils/cluster.py @@ -160,6 +160,20 @@ def parse_kwargs(kwargs: str | dict | None, **extra_kwargs) -> dict | None: return full_kwargs +# NEMO_SKILLS_* vars that must NOT be auto-forwarded from the launcher shell +# into the job/container. Scoped tightly to the SSH-tunnel controls: when BOTH +# are set, model/base.py opens an SSH tunnel at init -- so forwarding a laptop's +# tunnel config into a cluster job (where the client and inference server are +# co-located and the host key path doesn't exist) makes model init fail or +# mis-route. Everything else with the NEMO_SKILLS_ prefix -- runtime knobs, +# model/base-url selection, sandbox settings, host paths, and future knobs -- is +# forwarded automatically (worst case a stale value is redundant, not fatal). +_NON_FORWARDED_ENV_VARS = { + "NEMO_SKILLS_SSH_SERVER", + "NEMO_SKILLS_SSH_KEY_PATH", +} + + def get_env_variables(cluster_config): """ Will get the environment variables from the cluster config and the user environment. @@ -215,15 +229,18 @@ def get_env_variables(cluster_config): "HF_TOKEN", "NGC_API_KEY", } - # Auto-passthrough for NEMO_SKILLS_* runtime tuning vars set in - # the user's shell. Without this, opt-in env knobs like - # NEMO_SKILLS_OPENAI_AIOHTTP, NEMO_SKILLS_DISABLE_UVLOOP, and - # NEMO_SKILLS_TRANSIENT_RETRIES are stripped when the launcher - # spawns the inner python -- they only live in the wrapper - # shell. Scoping by prefix keeps the surface tight (no risk of - # accidentally exporting unrelated user env). + # Auto-passthrough for NEMO_SKILLS_* runtime knobs set in the user's shell. + # Without this, opt-in env knobs (e.g. NEMO_SKILLS_OPENAI_AIOHTTP, + # NEMO_SKILLS_DISABLE_UVLOOP) are stripped when the launcher spawns the inner + # python -- they only live in the wrapper shell. Forward the whole prefix so + # future knobs work without maintaining a list, except the SSH-tunnel vars in + # _NON_FORWARDED_ENV_VARS (see there). for name in os.environ: - if name.startswith("NEMO_SKILLS_") and name not in optional_env_vars_to_add: + if ( + name.startswith("NEMO_SKILLS_") + and name not in _NON_FORWARDED_ENV_VARS + and name not in optional_env_vars_to_add + ): optional_env_vars_to_add.add(name) default_factories = { "HF_TOKEN": lambda: str(token) if (token := get_token()) else "", diff --git a/tests/test_async_loop_integrity.py b/tests/test_async_loop_integrity.py new file mode 100644 index 0000000000..1ffa70445d --- /dev/null +++ b/tests/test_async_loop_integrity.py @@ -0,0 +1,80 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# 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. +"""restore_async_order must tolerate a dropped datapoint position. + +The async loop is fault-tolerant (a failed datapoint is logged and skipped, and +a placeholder is emitted), so restore_async_order can encounter an -async file +whose positions are non-contiguous. It must not IndexError (which would destroy +the whole run's output at the final reorder step) nor mis-align rows. +""" + +import json +from types import SimpleNamespace + +from nemo_skills.inference.generate import GenerationTask + + +def _run_restore(tmp_path, records): + task = GenerationTask.__new__(GenerationTask) + out = tmp_path / "out.jsonl" + task.cfg = SimpleNamespace( + output_file=str(out), + async_position_key="_async_position", + enable_litellm_cache=False, + ) + with open(str(out) + "-async", "w", encoding="utf-8") as f: + for rec in records: + f.write(json.dumps(rec) + "\n") + task.restore_async_order() + with open(str(out), encoding="utf-8") as f: + rows = [json.loads(line) for line in f] + return rows, out + + +def test_restore_async_order_tolerates_missing_middle_position(tmp_path): + # positions 0 and 2 present, 1 dropped. Old code did [None]*len(lines)=2 + # and indexed at 2 -> IndexError. Must not crash; must keep order. + rows, out = _run_restore( + tmp_path, + [ + {"_async_position": 2, "generation": "c"}, + {"_async_position": 0, "generation": "a"}, + ], + ) + assert rows == [{"generation": "a"}, {"generation": "c"}] + assert not (out.parent / "out.jsonl-async").exists() + + +def test_restore_async_order_complete_is_ordered(tmp_path): + rows, _ = _run_restore( + tmp_path, + [ + {"_async_position": 1, "generation": "b"}, + {"_async_position": 0, "generation": "a"}, + {"_async_position": 2, "generation": "c"}, + ], + ) + assert rows == [{"generation": "a"}, {"generation": "b"}, {"generation": "c"}] + + +def test_restore_async_order_skips_positionless_record(tmp_path): + # A record with no position key must be skipped (logged), not KeyError. + rows, _ = _run_restore( + tmp_path, + [ + {"_async_position": 0, "generation": "a"}, + {"generation": "orphan"}, # no _async_position + ], + ) + assert rows == [{"generation": "a"}] diff --git a/tests/test_native_openai_path.py b/tests/test_native_openai_path.py index f711f1d858..17907b656f 100644 --- a/tests/test_native_openai_path.py +++ b/tests/test_native_openai_path.py @@ -181,3 +181,55 @@ class _NoNative(VLLMModel): model = _NoNative(model="test-model", base_url="http://localhost:9999/v1") assert model._async_openai_client is None + + +class _CapturingAsyncOpenAI: + """Captures the kwargs passed to chat.completions.create.""" + + def __init__(self, *args, **kwargs): + self.captured_kwargs = None + outer = self + + class _Completions: + async def create(self, **kw): + outer.captured_kwargs = kw + return _FakeResponse(types.SimpleNamespace(content="ok")) + + class _Chat: + completions = _Completions() + + self.chat = _Chat() + + +def test_native_params_drop_none_and_litellm_only(monkeypatch): + # The native path must not forward litellm-only keys (allowed_openai_params) + # nor explicit None values (which the SDK would serialize as JSON null, + # unlike litellm which omits them). + monkeypatch.delenv("NEMO_SKILLS_OPENAI_AIOHTTP", raising=False) + monkeypatch.setattr(openai, "AsyncOpenAI", _CapturingAsyncOpenAI, raising=False) + monkeypatch.setattr(openai, "DefaultAioHttpClient", _DummyAioHttpClient, raising=False) + + model = VLLMModel(model="test-model", base_url="http://localhost:9999/v1") + assert isinstance(model._async_openai_client, _CapturingAsyncOpenAI) + + asyncio.run( + model.generate_async( + prompt=[{"role": "user", "content": "hi"}], + tokens_to_generate=8, + temperature=0.0, + top_p=0.95, + random_seed=None, # -> seed=None in the builder, must be dropped + stop_phrases=None, # -> stop=None, must be dropped + tools=None, # -> tools=None, must be dropped + response_format=None, # -> response_format=None, must be dropped + ) + ) + + captured = model._async_openai_client.captured_kwargs + assert captured is not None + assert "allowed_openai_params" not in captured + assert all(v is not None for v in captured.values()), f"None leaked: {captured}" + # explicit-None builder fields were dropped, not sent as null + for dropped in ("seed", "stop", "tools", "response_format", "top_logprobs"): + assert dropped not in captured + assert captured["model"] == "test-model" diff --git a/tests/test_pipeline_utils.py b/tests/test_pipeline_utils.py index bb2827d48f..43c6bf6ac7 100644 --- a/tests/test_pipeline_utils.py +++ b/tests/test_pipeline_utils.py @@ -457,3 +457,37 @@ def test_add_task_sandbox_mounts_override_keep_mounts_true(mock_port, mock_get_e ) assert mock_get_executor.call_args_list[-1].kwargs["mounts"] == ["/host/data:/sandbox/data:ro"] + + +def test_get_env_variables_forwards_prefix_except_ssh(monkeypatch): + """NEMO_SKILLS_* passthrough forwards the whole prefix (so future knobs work + without a maintained list), except the SSH-tunnel vars -- forwarding a + laptop's tunnel config into a co-located cluster job would break model init. + """ + from nemo_skills.pipeline.utils.cluster import get_env_variables + + # Forwarded: runtime knobs, model/base-url selection, sandbox, host paths, + # and future knobs (worst case a stale value is redundant, not fatal). + monkeypatch.setenv("NEMO_SKILLS_OPENAI_AIOHTTP", "1") + monkeypatch.setenv("NEMO_SKILLS_DISABLE_UVLOOP", "1") + monkeypatch.setenv("NEMO_SKILLS_OPENAI_BASE_URL", "http://model:8000/v1") + monkeypatch.setenv("NEMO_SKILLS_SANDBOX_HOST", "sandbox") + monkeypatch.setenv("NEMO_SKILLS_CONFIG_DIR", "/home/user/cluster_configs") + monkeypatch.setenv("NEMO_SKILLS_SOME_FUTURE_KNOB", "x") + # Not forwarded: SSH-tunnel controls. + monkeypatch.setenv("NEMO_SKILLS_SSH_SERVER", "login-node") + monkeypatch.setenv("NEMO_SKILLS_SSH_KEY_PATH", "/home/user/.ssh/id_rsa") + + env = get_env_variables({}) + + for forwarded in ( + "NEMO_SKILLS_OPENAI_AIOHTTP", + "NEMO_SKILLS_DISABLE_UVLOOP", + "NEMO_SKILLS_OPENAI_BASE_URL", + "NEMO_SKILLS_SANDBOX_HOST", + "NEMO_SKILLS_CONFIG_DIR", + "NEMO_SKILLS_SOME_FUTURE_KNOB", + ): + assert forwarded in env, f"{forwarded} should be forwarded" + for denied in ("NEMO_SKILLS_SSH_SERVER", "NEMO_SKILLS_SSH_KEY_PATH"): + assert denied not in env, f"{denied} (SSH tunnel) must not be forwarded into the job env" From cb1516735bcd21801e93b1c91c986b3861a0bb0c Mon Sep 17 00:00:00 2001 From: gwarmstrong Date: Wed, 1 Jul 2026 13:51:13 -0700 Subject: [PATCH 3/5] model/base: drop redundant routing-key note, clarify why reasoning is normalized - Remove the redundant "kept separate on purpose" sentence from the _apply_routing_keys docstring (the X-Request-Id / prompt_cache_key bullets already cover it). - Lead the _extract_reasoning docstring with why the normalization is manual: litellm normally maps vLLM's reasoning field to reasoning_content, but the native AsyncOpenAI fast-path bypasses litellm and sees the raw field. Signed-off-by: gwarmstrong --- nemo_skills/inference/model/base.py | 14 ++++++-------- 1 file changed, 6 insertions(+), 8 deletions(-) diff --git a/nemo_skills/inference/model/base.py b/nemo_skills/inference/model/base.py index f2e6144833..dfee6b608f 100644 --- a/nemo_skills/inference/model/base.py +++ b/nemo_skills/inference/model/base.py @@ -350,9 +350,6 @@ def _apply_routing_keys(request_params: dict, cache_key: str | None) -> None: OpenAI/vLLM prefix-cache hint. Requests that share a value are hinting they share a common prompt prefix and can reuse a warm prefix cache. Omitted when cache_key is None. - - Kept separate on purpose: prompt_cache_key is meant to be SHARED across - distinct requests, so it must not double as the per-request idempotency id. """ request_params.setdefault("extra_headers", {}).setdefault("X-Request-Id", str(uuid.uuid4())) if cache_key is not None: @@ -594,11 +591,12 @@ def _extract_reasoning(message_or_delta) -> str | None: """Return reasoning text from an OpenAI chat message or streaming delta, tolerating both field names. - vLLM historically exposed reasoning as `reasoning_content`; newer vLLM - versions expose it as `reasoning`. litellm normalizes these to - `reasoning_content` for us, but the native AsyncOpenAI fast-path sees - vLLM's raw field unchanged -- so we accept either here. Prefers - `reasoning_content` when both are present. + We normalize the field name here manually because litellm typically + takes care of it for us (it maps vLLM's field to `reasoning_content`), + but the native AsyncOpenAI fast-path bypasses litellm and sees vLLM's + raw field. vLLM historically exposed it as `reasoning_content`; newer + versions use `reasoning`. We accept either, preferring `reasoning_content` + when both are present. """ for attr in ("reasoning_content", "reasoning"): value = getattr(message_or_delta, attr, None) From de62c929abee120371aed4af9555277372e58eee Mon Sep 17 00:00:00 2001 From: gwarmstrong Date: Wed, 1 Jul 2026 15:56:03 -0700 Subject: [PATCH 4/5] model/base: route pydantic response_format through litellm, not native path CI failure: tests/test_generation.py::test_judge_generations_with_structured_output produced an empty judgement -> JSONDecodeError. Root cause: the judge uses ++structured_output (a pydantic BaseModel response_format) on server_type=openai, which now defaults to the native AsyncOpenAI path. The SDK's chat.completions.create() rejects a BaseModel class: TypeError: You tried to pass a `BaseModel` class to chat.completions.create(); You must use chat.completions.parse() instead litellm accepts a BaseModel and converts it to a JSON schema, so route such requests back through litellm: when response_format is a pydantic BaseModel subclass, skip the native fast-path. dict/json_schema response_formats stay on the native path. (pydantic import is aliased -- this module defines its own class named BaseModel.) Verified live against an OpenAI-compatible endpoint: native default + pydantic response_format now returns valid structured JSON. Adds a unit test asserting the fallback. Signed-off-by: gwarmstrong --- nemo_skills/inference/model/base.py | 13 +++++++++- tests/test_native_openai_path.py | 40 +++++++++++++++++++++++++++++ 2 files changed, 52 insertions(+), 1 deletion(-) diff --git a/nemo_skills/inference/model/base.py b/nemo_skills/inference/model/base.py index dfee6b608f..0b1ac0eaf0 100644 --- a/nemo_skills/inference/model/base.py +++ b/nemo_skills/inference/model/base.py @@ -13,6 +13,7 @@ # limitations under the License. import abc import asyncio +import inspect import logging import os import uuid @@ -25,6 +26,7 @@ import litellm.llms.custom_httpx.http_handler import litellm.llms.openai.common_utils import openai +from pydantic import BaseModel as PydanticBaseModel from nemo_skills.inference.patch_litellm_logging import patch_litellm_logging_worker from nemo_skills.utils import get_logger_name @@ -439,7 +441,16 @@ async def generate_async( # streaming and non-streaming chat are routed here; the # responses API, text completions, and non-OpenAI # providers still fall through to litellm. - use_native = self._async_openai_client is not None + # + # Exception: a pydantic BaseModel response_format (structured + # output) must go through litellm. The native SDK's + # chat.completions.create() rejects a BaseModel class + # ("must use .parse() instead"); litellm accepts it and + # converts it to a JSON schema. (dict/json_schema + # response_formats are fine on the native path.) + rf = request_params.get("response_format") + needs_litellm_response_format = inspect.isclass(rf) and issubclass(rf, PydanticBaseModel) + use_native = self._async_openai_client is not None and not needs_litellm_response_format if use_native: # Strip (a) litellm-only knobs the strict SDK rejects # and (b) None-valued params. litellm OMITS None-valued diff --git a/tests/test_native_openai_path.py b/tests/test_native_openai_path.py index 17907b656f..c688de6c4a 100644 --- a/tests/test_native_openai_path.py +++ b/tests/test_native_openai_path.py @@ -233,3 +233,43 @@ def test_native_params_drop_none_and_litellm_only(monkeypatch): for dropped in ("seed", "stop", "tools", "response_format", "top_logprobs"): assert dropped not in captured assert captured["model"] == "test-model" + + +def test_pydantic_response_format_falls_back_to_litellm(monkeypatch): + # A pydantic BaseModel response_format (structured output) must NOT go through + # the native client -- chat.completions.create() rejects a BaseModel class + # ("must use .parse() instead"). litellm accepts it, so we fall back. + import litellm + from pydantic import BaseModel + + monkeypatch.delenv("NEMO_SKILLS_OPENAI_AIOHTTP", raising=False) + monkeypatch.setattr(openai, "AsyncOpenAI", _CapturingAsyncOpenAI, raising=False) + monkeypatch.setattr(openai, "DefaultAioHttpClient", _DummyAioHttpClient, raising=False) + + litellm_called = {"hit": False} + + async def fake_acompletion(**kwargs): + litellm_called["hit"] = True + return _FakeResponse(types.SimpleNamespace(content='{"ok": true}')) + + monkeypatch.setattr(litellm, "acompletion", fake_acompletion) + + class _RF(BaseModel): + answer: str + + model = VLLMModel(model="test-model", base_url="http://localhost:9999/v1") + assert isinstance(model._async_openai_client, _CapturingAsyncOpenAI) # native is available... + + asyncio.run( + model.generate_async( + prompt=[{"role": "user", "content": "hi"}], + tokens_to_generate=8, + temperature=0.0, + top_p=0.95, + response_format=_RF, + ) + ) + + # ...but it was NOT used for the structured-output request; litellm handled it. + assert model._async_openai_client.captured_kwargs is None + assert litellm_called["hit"] is True From 3665bb3b01997f751eded8f9471f1f3b5718448f Mon Sep 17 00:00:00 2001 From: gwarmstrong Date: Wed, 1 Jul 2026 17:19:10 -0700 Subject: [PATCH 5/5] Clean up high-throughput inference comments Signed-off-by: gwarmstrong --- nemo_skills/inference/generate.py | 184 +++--------------- nemo_skills/inference/model/azure.py | 4 +- nemo_skills/inference/model/base.py | 169 +++++----------- nemo_skills/inference/model/openai.py | 6 - .../inference/model/vllm_multimodal.py | 5 +- nemo_skills/pipeline/utils/cluster.py | 16 +- tests/test_async_loop_integrity.py | 13 +- tests/test_native_openai_path.py | 38 +--- tests/test_pipeline_utils.py | 10 +- tests/test_prompt_cache_routing.py | 15 +- 10 files changed, 95 insertions(+), 365 deletions(-) diff --git a/nemo_skills/inference/generate.py b/nemo_skills/inference/generate.py index 68e1027c23..47a4916313 100644 --- a/nemo_skills/inference/generate.py +++ b/nemo_skills/inference/generate.py @@ -22,16 +22,8 @@ import sys import time -# uvloop is a drop-in replacement for asyncio's default event loop -# that's ~2-4x faster on CPU-bound asyncio workloads. ns generate's -# async_loop at high concurrency (5k+ lanes) becomes CPU-bound on -# the event loop's per-task scheduling overhead -- not on actual -# Python work in user code. Installing uvloop closes that gap so -# the bounded pool can refill at the rate the upstream can sustain -# instead of the rate one Python thread can scheduler-switch. -# -# Set NEMO_SKILLS_DISABLE_UVLOOP=1 to opt out (useful if a tool in -# the stack assumes selector loop semantics; rare). +# Prefer uvloop when available to reduce event-loop scheduling overhead. +# Set NEMO_SKILLS_DISABLE_UVLOOP to use the default asyncio event loop. if not os.environ.get("NEMO_SKILLS_DISABLE_UVLOOP"): try: import uvloop @@ -144,46 +136,11 @@ class GenerationTaskConfig: # maximum number of concurrent requests to the server for the async loop # if sync loop is used, this is the batch size max_concurrent_requests: int = 512 - # Rate at which the async loop opens NEW HTTP connections during - # initial prime (target connections per second). The bounded - # worker pool would otherwise spawn `max_concurrent_requests` - # tasks in a tight loop, each opening a TCP connection nearly - # simultaneously -- the kernel accept queue on the receiving end - # (typically capped at net.core.somaxconn=4096) can briefly - # overflow and SYN-drop the excess. Limiting the ramp rate so - # each ~ramp_rate/4-connection burst is followed by a short pause - # keeps the accept queue draining cleanly. - # - # Default 4000/sec accounts for ns generate's per-task pre-work - # (format_prompt, build litellm request, semaphore acquires) -- - # the time from asyncio.create_task to "TCP request fired" is - # significant, so the configured task-creation rate has to be - # higher than the SYN rate you actually want at the kernel. - # Each burst is still ramp_rate/4 SYNs (1000 at default), well - # under net.core.somaxconn=4096. - # - # Set to 0 (or negative) to disable ramping and prime - # instantaneously. The transient-retry logic in - # BaseModel.generate_async would still cover any resulting SYN - # drops; the ramp just makes them not happen in the first place. + # Target task-creation rate while initially filling the async worker pool. + # Tasks are created in short batches to limit connection bursts. Values <= 0 disable ramping. ramp_rate_per_sec: int = 4000 - # Number of background coroutines that postprocess + (optionally) - # evaluate completed datapoints. The async loop's bounded worker - # pool used to keep each lane "alive" (occupying a slot) until - # postprocess + evaluate + queue.put all completed -- which - # serialized those steps with the next HTTP refill. For - # reasoning-heavy workloads the per-lane post-HTTP work is small - # in absolute terms (~ms per record) but multiplied across - # thousands of lanes it can keep the bounded-pool slot occupied - # while real HTTP capacity sits idle. Moving postprocess + eval - # to a separate processor pool lets the lane exit the moment - # generate_async returns, so the bounded pool refills with the - # next datapoint immediately. - # - # Default 16 is enough headroom for synchronous postprocess (no - # eval) at >3k records/sec. If you enable an LLM-judge evaluator - # inside generate, raise this to match the eval concurrency you - # want. + # Number of background coroutines that postprocess and optionally evaluate completed datapoints. + # Increase this when per-datapoint evaluation benefits from additional concurrency. postprocess_workers: int = 16 # chunk the dataset into equal sized parts and index into them num_chunks: int | None = None # if specified, will split the data into chunks and only generate for one chunk @@ -469,12 +426,7 @@ def __init__(self, cfg: GenerationTaskConfig): # output_lock will be initialized when async_loop is called self.output_lock = None - # Three-stage pipeline queues (set by async_loop): - # _raw_queue: lanes -> processors (un-postprocessed outputs) - # _output_queue: processors -> writer (ready-to-write records) - # None outside of async_loop; _generate_and_save_datapoint - # falls back to inline postprocess + locked write in that - # case (used by entrypoints like prover.py). + # Pipeline queues are set by async_loop; None selects inline processing for other callers. self._raw_queue = None self._output_queue = None @@ -917,17 +869,11 @@ async def _generate_and_save_datapoint(self, data_point, all_data, fout, pbar): output["generation_end_time"] = end_time output["generation_time"] = end_time - start_time - # Hand the raw (un-postprocessed) completion off to the - # processor pool. The bounded-pool slot in async_loop frees - # the moment this coroutine returns, so the next datapoint - # can begin its HTTP request immediately. Postprocess + - # evaluate run in the processor pool; the writer task drains - # processed records to disk in batches. + # Async-loop callers enqueue completed generations for postprocessing and batched output. if self._raw_queue is not None: await self._raw_queue.put((output, data_point)) else: - # Fallback path (async_loop not entered through the - # processor/writer split). Run synchronously inline. + # Other callers postprocess and write inline. await self.postprocess_single_output(output, data_point) if self.should_run_evaluation and self.evaluator: output = await self.evaluate_single_datapoint({**data_point, **output}) @@ -956,14 +902,7 @@ async def async_loop(self, data): pbar = tqdm(total=len(remaining_data_points), desc="Remaining generations") - # Default Python buffering (8KB on most platforms). The - # previous `buffering=1` forced a kernel flush on every - # newline -- combined with the per-completion output_lock, - # that flush became the throughput bottleneck (~60-80/sec - # cap regardless of client concurrency). The writer task - # below issues an explicit fout.flush() after every batch - # drained from the queue, so durability is preserved without - # paying for a syscall per record. + # The writer flushes after each batch, so default buffering avoids per-record flushes. with open(self.cfg.output_file + "-async", "at", encoding="utf-8") as fout: # Dump prefilled data first if len(prefilled_data_points) > 0: @@ -976,35 +915,15 @@ async def async_loop(self, data): async with self.output_lock: self.dump_outputs(prefilled_outputs, prefilled_data_points, fout) - # Three-stage pipeline: - # - # lanes ── raw_queue ──> processors ── write_queue ──> writer - # - # `lanes` are the bounded-pool tasks below; each one does - # one HTTP call and exits the instant the response is - # parsed. `processors` is a small pool of background - # coroutines that run postprocess_single_output and - # (optionally) evaluate_single_datapoint -- both of which - # used to run INSIDE the lane, holding its bounded-pool - # slot for ~5-10ms past the HTTP return for postprocess - # alone, and seconds or more if an LLM-judge evaluator - # was attached. Decoupling them lets the lane refill - # instantly. `writer` batches processed records and - # flushes once per batch. - # - # Queue sizes are picked so producer puts never block in - # steady state. A `None` sentinel propagates through each - # stage on shutdown. + # Generation tasks -> raw queue -> postprocessors -> output queue -> batched writer. + # Bounded queues provide backpressure; None sentinels stop each processing stage. queue_cap = max(self.cfg.max_concurrent_requests * 2, 1024) self._raw_queue: asyncio.Queue = asyncio.Queue(maxsize=queue_cap) self._output_queue = asyncio.Queue(maxsize=queue_cap) num_processors = max(1, int(self.cfg.postprocess_workers)) async def _processor(): - # Drain raw_queue, postprocess (+ evaluate), forward - # to output_queue. A `None` sentinel terminates this - # worker; the dispatcher below pushes one sentinel - # per worker. + # Each processor receives its own sentinel after all raw outputs are queued. while True: item = await self._raw_queue.get() if item is None: @@ -1017,31 +936,20 @@ async def _processor(): except Exception as e: pos = data_point.get(self.cfg.async_position_key, "?") LOG.exception("postprocess/evaluate failed for async_pos=%s", pos) - # postprocess_single_output merges data_point (which carries - # async_position_key) into output only partway through, so a - # failure before that merge would write a record with no - # position key -- and restore_async_order would then KeyError - # on the whole file. Stamp the position and an error marker so - # the record is still placeable and the failure is visible. + # Preserve ordering metadata and failure details for the output record. if self.cfg.async_position_key in data_point: output[self.cfg.async_position_key] = data_point[self.cfg.async_position_key] output.setdefault("postprocess_error", repr(e)) await self._output_queue.put((output, data_point)) def _flush_batch(batch): - # Route through self.dump_outputs (the overridable hook the - # prefilled path also uses) rather than an inline json.dumps, so - # a subclass customizing serialization applies to async records too. + # Use the overridable hook so subclasses retain custom serialization. self.dump_outputs([o for o, _dp in batch], [dp for _o, dp in batch], fout) pbar.update(len(batch)) fout.flush() async def _writer(): - # Drain in batches: pull one item (blocking), then - # opportunistically drain any others already in the - # queue without blocking, write them all, and flush - # once at the end. Throughput scales with completion - # rate, not per-record syscall cost. + # Block for the first item, then drain and flush all currently queued items as one batch. while True: item = await self._output_queue.get() if item is None: @@ -1062,25 +970,8 @@ async def _writer(): processor_tasks = [asyncio.create_task(_processor()) for _ in range(num_processors)] writer_task = asyncio.create_task(_writer()) - # Bounded worker pool. At most `max_in_flight` tasks - # exist at any moment; as each one completes we spawn the - # next datapoint's task. Memory is O(in_flight), not - # O(total inputs), so the loop runs cleanly on 100k/1M+ - # input files without preallocating a coroutine frame - # for each. - # - # Was: tasks = [asyncio.create_task(...) for dp in data] - # That created ALL coroutines upfront. At a 100k input - # file and ~1-5 KB per coroutine frame (closure + Task - # bookkeeping) the upfront allocation was 100-500 MB -- - # avoidable memory pressure when co-located with other - # services. The outer semaphore inside each task bounded - # in-flight requests but didn't help that upfront cost. - # - # Each task is wrapped in a try/except so a single - # exception doesn't crash the whole gather (the old - # default behavior was first-exception-wins, killing - # all other in-flight datapoints). + # Bound both request concurrency and allocated task state to max_in_flight. + # Each task handles its own error so other datapoints continue processing. max_in_flight = self.cfg.max_concurrent_requests async def _safe_one(dp): @@ -1089,28 +980,18 @@ async def _safe_one(dp): except Exception as e: pos = dp.get(self.cfg.async_position_key, "?") LOG.exception("datapoint async_pos=%s crashed", pos) - # Emit a placeholder so this datapoint's position still - # produces exactly one output line. Without it, - # restore_async_order sees a hole and either IndexErrors (a - # non-last position missing) or silently drops/misaligns rows. - # Routed through raw_queue so postprocess stamps the position - # and merges the original fields, same as a normal record. + # Emit one placeholder per failed datapoint to preserve positional alignment. + # Route it through the raw queue so it follows normal postprocessing. if self._raw_queue is not None: await self._raw_queue.put(({"generation": "", "generation_error": repr(e)}, dp)) data_iter = iter(remaining_data_points) in_flight: set = set() - # Prime the pool up to the concurrency cap, with rate- - # limited ramp so we don't open all max_in_flight TCP - # connections in a microsecond burst. + # Fill the concurrency pool at the configured rate to limit initial connection bursts. ramp_rate = self.cfg.ramp_rate_per_sec if ramp_rate > 0: - # Spawn in batches sized so the implied pause is - # ~250ms -- short enough not to extend wall-time - # noticeably (target_depth=20k at 1000/s = 20s - # ramp), long enough for the kernel accept queue - # to drain between bursts. + # Quarter-second batches keep the ramp smooth while filling the pool promptly. ramp_batch = max(50, ramp_rate // 4) ramp_interval = ramp_batch / ramp_rate else: @@ -1126,10 +1007,7 @@ async def _safe_one(dp): if ramp_interval > 0 and primed % ramp_batch == 0: await asyncio.sleep(ramp_interval) - # Refill as tasks complete. asyncio.wait with - # FIRST_COMPLETED returns as soon as ANY task finishes, - # which lets us start a fresh task immediately rather - # than waiting for the whole batch. + # Refill one worker slot for each completed task. while in_flight: done, in_flight = await asyncio.wait( in_flight, @@ -1142,14 +1020,7 @@ async def _safe_one(dp): continue in_flight.add(asyncio.create_task(_safe_one(dp))) - # Shutdown sequence: cascade sentinels through the - # pipeline so every in-flight record is flushed. - # 1. Put one None on raw_queue per processor -- each - # processor exits after seeing its sentinel. - # 2. Wait for all processors. Any record put on - # output_queue by a processor is already there. - # 3. Put one None on output_queue -- writer drains - # whatever's queued ahead of it, then exits. + # Stop processors after all raw outputs are queued, then stop the writer after processors finish. for _ in range(num_processors): await self._raw_queue.put(None) await asyncio.gather(*processor_tasks) @@ -1167,12 +1038,7 @@ def restore_async_order(self): with open(self.cfg.output_file + "-async", "rt", encoding="utf-8") as fin: generations = [json.loads(line) for line in fin] - # Size by the max position, not the line count: those differ if any - # record was dropped, and indexing a count-sized list by original - # position would then IndexError. pop(default) + skip means a - # missing/position-less record can never crash finalization. The - # pipeline emits a placeholder per datapoint so holes should not occur; - # this is a safety net. + # Original positions can be sparse, so size from the maximum and skip missing entries. key = self.cfg.async_position_key positions = [gen[key] for gen in generations if key in gen] ordered_generations = [None] * ((max(positions) + 1) if positions else 0) diff --git a/nemo_skills/inference/model/azure.py b/nemo_skills/inference/model/azure.py index f3459db9ab..951fa8325d 100644 --- a/nemo_skills/inference/model/azure.py +++ b/nemo_skills/inference/model/azure.py @@ -19,8 +19,8 @@ class AzureOpenAIModel(OpenAIModel): MODEL_PROVIDER = "azure" - # Azure uses a different endpoint layout + api-version/api-key auth that a - # plain AsyncOpenAI client can't target; keep it on litellm. + # Azure requires provider-specific endpoint and authentication handling + # from litellm. SUPPORTS_NATIVE_OPENAI = False def __init__( diff --git a/nemo_skills/inference/model/base.py b/nemo_skills/inference/model/base.py index 0b1ac0eaf0..d3fe428cb1 100644 --- a/nemo_skills/inference/model/base.py +++ b/nemo_skills/inference/model/base.py @@ -76,18 +76,15 @@ class BaseModel: # Litellm provider name MODEL_PROVIDER = "openai" - # Whether this model may use the native AsyncOpenAI + aiohttp fast-path - # (see __init__). Only OpenAI-compatible endpoints that speak the plain - # /v1/chat/completions schema at self.base_url can opt in. Subclasses that - # talk to a non-OpenAI server (megatron), a differently-authed endpoint - # (azure), a non-OpenAI provider (gemini), or need litellm's request - # rewriting (multimodal audio) leave this False and stay on litellm. + # Whether chat completions may use the native AsyncOpenAI + aiohttp path. + # Subclasses should opt in only when their endpoint accepts the standard + # OpenAI chat-completions request without litellm transformations. SUPPORTS_NATIVE_OPENAI = False # Keys that litellm understands but AsyncOpenAI's strict schema rejects as # unknown top-level fields; stripped before handing the body to the native - # client. Keep this in sync with any litellm-only knob a _build_chat_request_params - # implementation may emit. + # client. Keep this in sync with any litellm-only parameter emitted by a + # _build_chat_request_params implementation. _LITELLM_ONLY_CHAT_PARAMS = frozenset({"allowed_openai_params"}) def __init__( @@ -112,16 +109,9 @@ def __init__( output_dir: str | None = None, # Request tokenizer initialization independent of soft_fail require_tokenizer: bool = False, - # Max in-flight HTTP requests gated by an internal semaphore. - # The semaphore was originally added because earlier - # combinations of httpx + litellm would hang under truly - # unbounded concurrency. Modern httpx (>=0.27) handles this - # fine, so the default is now high enough to effectively be - # "no internal limit" -- the script-level concurrency knob - # (generate.py `max_concurrent_requests`) is the single - # user-facing cap. Override via $NEMO_SKILLS_MODEL_CONCURRENCY - # only if you have a specific reason to throttle BELOW the - # script-level cap. + # Upper bound on requests admitted by this model instance. The high + # default leaves normal concurrency control to the caller; use + # NEMO_SKILLS_MODEL_CONCURRENCY to impose a lower model-level limit. concurrent_requests_limit: int = int(os.environ.get("NEMO_SKILLS_MODEL_CONCURRENCY", "65536")), ): self._tunnel = None @@ -195,45 +185,23 @@ def __init__( httpx_limits = httpx.Limits(max_keepalive_connections=None, max_connections=None) litellm.client_session = httpx.Client(limits=httpx_limits) litellm.aclient_session = httpx.AsyncClient(limits=httpx_limits) - # Controlling concurrent requests using semaphore since large - # concurrent requests result into httpx hanging + # Gate in-flight requests at the configured model-level limit. self.concurrent_semaphore = asyncio.Semaphore(concurrent_requests_limit) - # Native aiohttp fast-path (DEFAULT for OpenAI-compatible - # providers). The litellm.acompletion path goes through httpx, - # a pure-Python HTTP implementation -- at 5k+ concurrent - # connections it becomes the GIL-bound bottleneck (~95% of - # single-core CPU in the async loop, per py-spy). The OpenAI - # SDK natively supports an aiohttp transport via - # DefaultAioHttpClient. aiohttp's request/response parsing is - # C-extension backed, so the same workload runs at <5% CPU on - # the same hardware. - # - # Enabled by default for providers with SUPPORTS_NATIVE_OPENAI - # (openai, vllm, sglang). Both the non-streaming and streaming - # chat endpoints go through this path; text completions, the - # responses API, and non-OpenAI-shaped providers (gemini, - # azure, megatron, multimodal) still use litellm. Set - # NEMO_SKILLS_OPENAI_AIOHTTP=0 to force everything back onto - # litellm/httpx. + # Use the lower-overhead AsyncOpenAI + aiohttp transport for providers + # that opt in. Text completions, the Responses API, and providers that + # require litellm transformations continue to use litellm. Set + # NEMO_SKILLS_OPENAI_AIOHTTP=0 to disable the native path. self._async_openai_client = None native_enabled = os.environ.get("NEMO_SKILLS_OPENAI_AIOHTTP", "1") != "0" if native_enabled and self.SUPPORTS_NATIVE_OPENAI: try: from openai import AsyncOpenAI, DefaultAioHttpClient - # openai SDK defaults the aiohttp connector to - # max_connections=1000, max_keepalive_connections=100, - # which silently caps the lane count at high - # concurrency (8k lanes see only ~1000 open TCP - # connections; the rest stall in the aiohttp connector - # queue). Also, httpx-aiohttp's transport translates - # httpx.Limits(max_connections=None) to "omit the - # `limit` kwarg," which makes aiohttp fall back to ITS - # default of 100 -- worse than the openai SDK - # default. We therefore pass an explicit large - # integer (NEMO_SKILLS_OPENAI_AIOHTTP_LIMIT, default - # 65536) instead of None. + # Set an explicit connector limit: passing None through the + # httpx-aiohttp adapter omits aiohttp's `limit` setting and + # restores its lower default. The environment variable allows + # deployments to choose an appropriate transport-level cap. aiohttp_limit = int(os.environ.get("NEMO_SKILLS_OPENAI_AIOHTTP_LIMIT", "65536")) self._async_openai_client = AsyncOpenAI( api_key=api_key, @@ -244,15 +212,10 @@ def __init__( max_keepalive_connections=aiohttp_limit, ), ), - # Mirror the litellm path's retry budget (self.litellm_kwargs - # also passes max_retries). The SDK's own retry loop honors - # Retry-After / backoff for HTTP 429 and 5xx -- error classes - # that our app-level transient-retry except clause does NOT - # catch (it only handles connection-class errors). Leaving this - # at 0 silently dropped 429/5xx under load on the native path, - # which is the DEFAULT for openai/vllm/sglang. SDK retries run - # inside our concurrent_semaphore hold, so one logical request - # still counts as one in-flight slot. + # Match the litellm retry budget. The SDK handles HTTP 429 + # and 5xx retries; the application-level handler below is + # limited to connection errors. SDK retries occur within the + # semaphore, so a logical request retains one in-flight slot. max_retries=max_retries, ) except Exception as e: @@ -341,18 +304,7 @@ def _build_responses_request_params(self, **kwargs) -> dict: @staticmethod def _apply_routing_keys(request_params: dict, cache_key: str | None) -> None: - """Attach optional, provider-agnostic request-routing keys. Plain - OpenAI/vLLM/gemini endpoints ignore (or echo) them. - - * `X-Request-Id` (extra_headers) = per-request id for idempotency / - tracing. Kept stable across the OpenAI SDK's internal retry loop so - a retry is identifiable as the same logical call; distinct calls - (incl. parallel samples) get distinct ids. - * `prompt_cache_key` (opt-in via `cache_key`) = the standard - OpenAI/vLLM prefix-cache hint. Requests that share a value are - hinting they share a common prompt prefix and can reuse a warm - prefix cache. Omitted when cache_key is None. - """ + """Attach a request ID and an optional prompt-cache key.""" request_params.setdefault("extra_headers", {}).setdefault("X-Request-Id", str(uuid.uuid4())) if cache_key is not None: request_params["prompt_cache_key"] = cache_key @@ -408,19 +360,8 @@ async def generate_async( "response_format": response_format, } - # Two separate retry budgets: - # - bad_request_count is for the gpt-oss "output messages" - # bug; small, app-level. Kept as-is until we drop that - # model (or vLLM fixes it upstream). - # - transient_count handles socket-level failures that the - # raw HTTP path can hit under burst load (kernel SYN-drops - # when the accept queue overflows, brief upstream - # unavailability during worker turnover, etc). litellm's - # own max_retries covers HTTP 429/5xx but NOT these - # connection-class exceptions, so a single ramp-up burst - # into the upstream could otherwise hard-fail individual - # datapoints. Retries here turn those into invisible - # ~50-200ms re-tries. + # Keep the targeted bad-request retry budget separate from retries for + # transient connection failures. max_retries = 2 retry_count = 0 max_transient_retries = int(os.environ.get("NEMO_SKILLS_TRANSIENT_RETRIES", "3")) @@ -433,14 +374,9 @@ async def generate_async( assert isinstance(prompt, list), "Chat completion requests must be a list of messages." request_params = self._build_chat_request_params(messages=prompt, stream=stream, **kwargs) self._apply_routing_keys(request_params, cache_key) - # Fast path: native AsyncOpenAI + aiohttp transport, - # set up in __init__ for OpenAI-compatible providers - # (disable with NEMO_SKILLS_OPENAI_AIOHTTP=0). Skips - # litellm's httpx-based dispatch which is the dominant - # single-core CPU consumer at high concurrency. Both - # streaming and non-streaming chat are routed here; the - # responses API, text completions, and non-OpenAI - # providers still fall through to litellm. + # Use the native client when this provider and request + # support it. Other endpoint types continue through + # litellm. # # Exception: a pydantic BaseModel response_format (structured # output) must go through litellm. The native SDK's @@ -452,14 +388,10 @@ async def generate_async( needs_litellm_response_format = inspect.isclass(rf) and issubclass(rf, PydanticBaseModel) use_native = self._async_openai_client is not None and not needs_litellm_response_format if use_native: - # Strip (a) litellm-only knobs the strict SDK rejects - # and (b) None-valued params. litellm OMITS None-valued - # fields from the wire body, but AsyncOpenAI serializes - # an explicit None as JSON `null` -- e.g. the vllm builder - # emits response_format/tools/stop/seed=None, and some - # OpenAI-compatible servers reject `"response_format": null`. - # Dropping None here keeps the native body byte-compatible - # with what the litellm path sent. + # Strip litellm-only parameters and None-valued fields. + # AsyncOpenAI rejects unknown keywords, and some + # compatible endpoints reject explicit JSON nulls that + # litellm omits from the request body. native_params = { k: v for k, v in request_params.items() @@ -536,11 +468,8 @@ async def generate_async( ConnectionError, ConnectionResetError, ) as e: - # Transient socket-level failure. Common cause at - # high concurrency: kernel SYN-drop during ramp-up - # when the upstream accept queue briefly overflows. - # Backoff doubles each retry (50ms, 100ms, 200ms) - # to let the accept queue drain. + # Retry transient connection failures with a short + # exponential backoff. if transient_count < max_transient_retries: transient_count += 1 backoff = 0.05 * (2 ** (transient_count - 1)) @@ -599,15 +528,11 @@ def _parse_completion_response( @staticmethod def _extract_reasoning(message_or_delta) -> str | None: - """Return reasoning text from an OpenAI chat message or streaming delta, - tolerating both field names. - - We normalize the field name here manually because litellm typically - takes care of it for us (it maps vLLM's field to `reasoning_content`), - but the native AsyncOpenAI fast-path bypasses litellm and sees vLLM's - raw field. vLLM historically exposed it as `reasoning_content`; newer - versions use `reasoning`. We accept either, preferring `reasoning_content` - when both are present. + """Return reasoning text from either supported response field. + + litellm uses `reasoning_content`, while direct OpenAI-compatible + responses may use `reasoning`. Prefer `reasoning_content` when both are + present. """ for attr in ("reasoning_content", "reasoning"): value = getattr(message_or_delta, attr, None) @@ -626,20 +551,17 @@ def _parse_chat_completion_response(self, response, include_response: bool = Fal elif getattr(response.usage, "input_tokens", None) is not None: result["num_input_tokens"] = response.usage.input_tokens - # Add reasoning_content if available (accepts .reasoning or .reasoning_content) + # Normalize either supported reasoning field in the result. reasoning = self._extract_reasoning(choice.message) if reasoning: result["reasoning_content"] = reasoning - # Dump the full usage object verbatim into the output. Keeping - # None-valued fields (rather than excluding them) preserves the - # distinction between "server explicitly sent null" and "field - # was absent" — both can be meaningful (e.g. some providers - # send prompt_tokens_details=null when there's no cache hit). + # Preserve the complete usage payload, including explicit nulls, so + # provider-specific usage metadata remains available. try: result["usage"] = response.usage.model_dump() except AttributeError: - # Older or non-pydantic usage objects: fall back to dict(). + # Support non-pydantic usage objects. try: result["usage"] = dict(response.usage) except Exception: @@ -713,7 +635,7 @@ def _process_chat_chunk(self, chunk): """Process a single chat chunk and return data to yield.""" if hasattr(chunk.choices[0], "delta"): cur_delta = chunk.choices[0].delta.content - # Check for reasoning in delta (accepts .reasoning or .reasoning_content) + # Normalize either supported reasoning field in the delta. reasoning_delta = self._extract_reasoning(chunk.choices[0].delta) tool_calls_delta = getattr(chunk.choices[0].delta, "tool_calls", None) else: @@ -807,10 +729,7 @@ def _serialize_output(self, response): async def _stream_chat_chunks_async(self, response): async for chunk in response: - # Some servers emit chunks with no choices (e.g. a trailing - # usage-only chunk when stream_options.include_usage is set, or - # keepalive chunks). They carry no delta -- skip them so - # _process_chat_chunk's choices[0] access stays safe. + # Usage-only and keepalive chunks have no choices to process. if not getattr(chunk, "choices", None): continue results = self._process_chat_chunk(chunk) diff --git a/nemo_skills/inference/model/openai.py b/nemo_skills/inference/model/openai.py index 68f0eb4070..4854f5fb0b 100644 --- a/nemo_skills/inference/model/openai.py +++ b/nemo_skills/inference/model/openai.py @@ -133,12 +133,6 @@ def _build_chat_request_params( "stream": stream, "tools": tools, "response_format": response_format, - # NOTE: idempotency (X-Request-Id) and prefix-cache affinity - # (prompt_cache_key, opt-in via generate_async(cache_key=...)) are - # attached centrally in BaseModel._apply_routing_keys. We deliberately - # do NOT set a per-request-UUID prompt_cache_key here: that would give - # every request a unique affinity key and defeat cross-request prefix - # caching. } if self._is_reasoning_model(self.model): diff --git a/nemo_skills/inference/model/vllm_multimodal.py b/nemo_skills/inference/model/vllm_multimodal.py index 2924d393de..98bd743d45 100644 --- a/nemo_skills/inference/model/vllm_multimodal.py +++ b/nemo_skills/inference/model/vllm_multimodal.py @@ -69,9 +69,8 @@ class VLLMMultimodalModel(VLLMModel): ) """ - # Audio/VLM request rewriting and response post-processing go through - # litellm; keep this off the native fast-path until it's verified for - # multimodal payloads. Overrides VLLMModel's True. + # Multimodal request and response transformations require litellm, so this + # overrides VLLMModel's native-client support. SUPPORTS_NATIVE_OPENAI = False def __init__( diff --git a/nemo_skills/pipeline/utils/cluster.py b/nemo_skills/pipeline/utils/cluster.py index 4237da23e7..6d5dabd5a1 100644 --- a/nemo_skills/pipeline/utils/cluster.py +++ b/nemo_skills/pipeline/utils/cluster.py @@ -160,14 +160,7 @@ def parse_kwargs(kwargs: str | dict | None, **extra_kwargs) -> dict | None: return full_kwargs -# NEMO_SKILLS_* vars that must NOT be auto-forwarded from the launcher shell -# into the job/container. Scoped tightly to the SSH-tunnel controls: when BOTH -# are set, model/base.py opens an SSH tunnel at init -- so forwarding a laptop's -# tunnel config into a cluster job (where the client and inference server are -# co-located and the host key path doesn't exist) makes model init fail or -# mis-route. Everything else with the NEMO_SKILLS_ prefix -- runtime knobs, -# model/base-url selection, sandbox settings, host paths, and future knobs -- is -# forwarded automatically (worst case a stale value is redundant, not fatal). +# SSH tunnel settings are launcher-local and are not forwarded automatically. _NON_FORWARDED_ENV_VARS = { "NEMO_SKILLS_SSH_SERVER", "NEMO_SKILLS_SSH_KEY_PATH", @@ -229,12 +222,7 @@ def get_env_variables(cluster_config): "HF_TOKEN", "NGC_API_KEY", } - # Auto-passthrough for NEMO_SKILLS_* runtime knobs set in the user's shell. - # Without this, opt-in env knobs (e.g. NEMO_SKILLS_OPENAI_AIOHTTP, - # NEMO_SKILLS_DISABLE_UVLOOP) are stripped when the launcher spawns the inner - # python -- they only live in the wrapper shell. Forward the whole prefix so - # future knobs work without maintaining a list, except the SSH-tunnel vars in - # _NON_FORWARDED_ENV_VARS (see there). + # Forward NEMO_SKILLS_* runtime configuration except launcher-local settings. for name in os.environ: if ( name.startswith("NEMO_SKILLS_") diff --git a/tests/test_async_loop_integrity.py b/tests/test_async_loop_integrity.py index 1ffa70445d..54a3ed9d21 100644 --- a/tests/test_async_loop_integrity.py +++ b/tests/test_async_loop_integrity.py @@ -11,13 +11,7 @@ # 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. -"""restore_async_order must tolerate a dropped datapoint position. - -The async loop is fault-tolerant (a failed datapoint is logged and skipped, and -a placeholder is emitted), so restore_async_order can encounter an -async file -whose positions are non-contiguous. It must not IndexError (which would destroy -the whole run's output at the final reorder step) nor mis-align rows. -""" +"""Tests for restoring ordered output from partial async result files.""" import json from types import SimpleNamespace @@ -43,8 +37,7 @@ def _run_restore(tmp_path, records): def test_restore_async_order_tolerates_missing_middle_position(tmp_path): - # positions 0 and 2 present, 1 dropped. Old code did [None]*len(lines)=2 - # and indexed at 2 -> IndexError. Must not crash; must keep order. + # Position 1 is absent; the remaining records retain their relative order. rows, out = _run_restore( tmp_path, [ @@ -69,7 +62,7 @@ def test_restore_async_order_complete_is_ordered(tmp_path): def test_restore_async_order_skips_positionless_record(tmp_path): - # A record with no position key must be skipped (logged), not KeyError. + # Positionless records are ignored. rows, _ = _run_restore( tmp_path, [ diff --git a/tests/test_native_openai_path.py b/tests/test_native_openai_path.py index c688de6c4a..25c2f439ed 100644 --- a/tests/test_native_openai_path.py +++ b/tests/test_native_openai_path.py @@ -11,19 +11,7 @@ # 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. -"""Unit tests for the native AsyncOpenAI + aiohttp fast-path in BaseModel. - -Covers the three feature-completion items: - 1. Native path is the DEFAULT for OpenAI-compatible providers and can be - disabled with NEMO_SKILLS_OPENAI_AIOHTTP=0; non-OpenAI providers never - use it. - 2. Reasoning content is read from BOTH `.reasoning` and `.reasoning_content` - (litellm normalizes this for us; the native path sees vLLM's raw field). - 3. Streaming is routed through the native path and tolerates chunks with no - choices. - -These are pure unit tests -- no server or network access. -""" +"""Unit tests for BaseModel's native AsyncOpenAI aiohttp path.""" import asyncio import types @@ -42,7 +30,7 @@ # --------------------------------------------------------------------------- -# Item 2: reasoning extraction accepts both field names +# Reasoning extraction # --------------------------------------------------------------------------- def test_extract_reasoning_from_reasoning_content(): obj = types.SimpleNamespace(reasoning_content="thinking") @@ -50,7 +38,6 @@ def test_extract_reasoning_from_reasoning_content(): def test_extract_reasoning_from_reasoning(): - # New vLLM field name; only `.reasoning` present. obj = types.SimpleNamespace(reasoning="thinking") assert BaseModel._extract_reasoning(obj) == "thinking" @@ -89,8 +76,7 @@ def __init__(self, message): def test_parse_chat_completion_reads_reasoning_field(): - # A response carrying ONLY `.reasoning` (as newer vLLM emits over the raw - # OpenAI schema) must still populate reasoning_content in the result. + # The `.reasoning` field is supported when `.reasoning_content` is absent. inst = BaseModel.__new__(BaseModel) message = types.SimpleNamespace(content="the answer", reasoning="the thinking") result = inst._parse_chat_completion_response(_FakeResponse(message)) @@ -100,7 +86,7 @@ def test_parse_chat_completion_reads_reasoning_field(): # --------------------------------------------------------------------------- -# Item 3: streaming chunk parsing (reasoning + empty-choices guard) +# Streaming chunk parsing # --------------------------------------------------------------------------- def test_process_chat_chunk_reads_reasoning_field(): inst = BaseModel.__new__(BaseModel) @@ -129,7 +115,7 @@ async def collect(): # --------------------------------------------------------------------------- -# Item 1: which providers are eligible for the native fast-path +# Native-client eligibility # --------------------------------------------------------------------------- def test_native_eligibility_by_provider(): assert BaseModel.SUPPORTS_NATIVE_OPENAI is False @@ -202,9 +188,7 @@ class _Chat: def test_native_params_drop_none_and_litellm_only(monkeypatch): - # The native path must not forward litellm-only keys (allowed_openai_params) - # nor explicit None values (which the SDK would serialize as JSON null, - # unlike litellm which omits them). + # The native client receives only supported parameters with concrete values. monkeypatch.delenv("NEMO_SKILLS_OPENAI_AIOHTTP", raising=False) monkeypatch.setattr(openai, "AsyncOpenAI", _CapturingAsyncOpenAI, raising=False) monkeypatch.setattr(openai, "DefaultAioHttpClient", _DummyAioHttpClient, raising=False) @@ -229,16 +213,15 @@ def test_native_params_drop_none_and_litellm_only(monkeypatch): assert captured is not None assert "allowed_openai_params" not in captured assert all(v is not None for v in captured.values()), f"None leaked: {captured}" - # explicit-None builder fields were dropped, not sent as null + # Builder fields with None values are omitted. for dropped in ("seed", "stop", "tools", "response_format", "top_logprobs"): assert dropped not in captured assert captured["model"] == "test-model" def test_pydantic_response_format_falls_back_to_litellm(monkeypatch): - # A pydantic BaseModel response_format (structured output) must NOT go through - # the native client -- chat.completions.create() rejects a BaseModel class - # ("must use .parse() instead"). litellm accepts it, so we fall back. + # Pydantic response formats require LiteLLM because the native completion + # method does not accept model classes. import litellm from pydantic import BaseModel @@ -258,7 +241,7 @@ class _RF(BaseModel): answer: str model = VLLMModel(model="test-model", base_url="http://localhost:9999/v1") - assert isinstance(model._async_openai_client, _CapturingAsyncOpenAI) # native is available... + assert isinstance(model._async_openai_client, _CapturingAsyncOpenAI) asyncio.run( model.generate_async( @@ -270,6 +253,5 @@ class _RF(BaseModel): ) ) - # ...but it was NOT used for the structured-output request; litellm handled it. assert model._async_openai_client.captured_kwargs is None assert litellm_called["hit"] is True diff --git a/tests/test_pipeline_utils.py b/tests/test_pipeline_utils.py index 43c6bf6ac7..03f084c500 100644 --- a/tests/test_pipeline_utils.py +++ b/tests/test_pipeline_utils.py @@ -460,21 +460,17 @@ def test_add_task_sandbox_mounts_override_keep_mounts_true(mock_port, mock_get_e def test_get_env_variables_forwards_prefix_except_ssh(monkeypatch): - """NEMO_SKILLS_* passthrough forwards the whole prefix (so future knobs work - without a maintained list), except the SSH-tunnel vars -- forwarding a - laptop's tunnel config into a co-located cluster job would break model init. - """ + """NEMO_SKILLS_* values are forwarded except launcher-local SSH settings.""" from nemo_skills.pipeline.utils.cluster import get_env_variables - # Forwarded: runtime knobs, model/base-url selection, sandbox, host paths, - # and future knobs (worst case a stale value is redundant, not fatal). + # Representative runtime settings, including an unknown future setting. monkeypatch.setenv("NEMO_SKILLS_OPENAI_AIOHTTP", "1") monkeypatch.setenv("NEMO_SKILLS_DISABLE_UVLOOP", "1") monkeypatch.setenv("NEMO_SKILLS_OPENAI_BASE_URL", "http://model:8000/v1") monkeypatch.setenv("NEMO_SKILLS_SANDBOX_HOST", "sandbox") monkeypatch.setenv("NEMO_SKILLS_CONFIG_DIR", "/home/user/cluster_configs") monkeypatch.setenv("NEMO_SKILLS_SOME_FUTURE_KNOB", "x") - # Not forwarded: SSH-tunnel controls. + # SSH tunnel settings remain launcher-local. monkeypatch.setenv("NEMO_SKILLS_SSH_SERVER", "login-node") monkeypatch.setenv("NEMO_SKILLS_SSH_KEY_PATH", "/home/user/.ssh/id_rsa") diff --git a/tests/test_prompt_cache_routing.py b/tests/test_prompt_cache_routing.py index d5af83b12d..6bceba4138 100644 --- a/tests/test_prompt_cache_routing.py +++ b/tests/test_prompt_cache_routing.py @@ -11,18 +11,13 @@ # 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. -"""Unit tests for the request-routing keys attached by BaseModel. - -Idempotency (X-Request-Id) and prefix-cache affinity (prompt_cache_key) are -DECOUPLED: X-Request-Id is per-request (unique per call, safe for retries), -prompt_cache_key is opt-in and SHARED across requests with a common prefix. -""" +"""Tests for BaseModel request IDs and optional prompt-cache affinity keys.""" from nemo_skills.inference.model.base import BaseModel def test_idempotency_header_always_set_affinity_opt_in(): - # No cache_key → X-Request-Id present (idempotency), prompt_cache_key OMITTED. + # Request IDs do not depend on cache affinity. params = {} BaseModel._apply_routing_keys(params, None) assert "prompt_cache_key" not in params @@ -38,8 +33,7 @@ def test_cache_key_sets_prompt_cache_key(): def test_shared_affinity_distinct_request_ids(): - # The branching/multi-sample invariant: same cache_key (shared prefix), but - # each call gets a DISTINCT X-Request-Id (unique per request). + # Requests may share cache affinity without sharing request IDs. p1, p2 = {}, {} BaseModel._apply_routing_keys(p1, "conv-1") BaseModel._apply_routing_keys(p2, "conv-1") @@ -48,8 +42,7 @@ def test_shared_affinity_distinct_request_ids(): def test_preserves_existing_headers_and_does_not_override_request_id(): - # setdefault semantics: a caller-set X-Request-Id (e.g. for app-level retry - # idempotency) and other headers survive. + # Existing request and custom headers are preserved. params = {"extra_headers": {"X-Custom": "v", "X-Request-Id": "preset"}} BaseModel._apply_routing_keys(params, None) assert params["extra_headers"]["X-Custom"] == "v"