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..47a4916313 100644 --- a/nemo_skills/inference/generate.py +++ b/nemo_skills/inference/generate.py @@ -15,11 +15,22 @@ import asyncio import json import logging +import os import random import shutil import subprocess import sys import time + +# 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 + + uvloop.install() + except ImportError: + pass from copy import deepcopy from dataclasses import asdict, field, is_dataclass from pathlib import Path @@ -125,6 +136,12 @@ 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 + # 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 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 chunk_id: int | None = None # if specified, will index the specified chunk only @@ -409,6 +426,9 @@ def __init__(self, cfg: GenerationTaskConfig): # output_lock will be initialized when async_loop is called self.output_lock = None + # Pipeline queues are set by async_loop; None selects inline processing for other callers. + self._raw_queue = None + self._output_queue = None def setup_prompt(self): if self.cfg.prompt_config is None: @@ -849,16 +869,17 @@ 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) + # 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: + # 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}) + 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 +902,8 @@ 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: + # 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: for output, data_point in zip(prefilled_outputs, prefilled_data_points): @@ -893,15 +915,119 @@ 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) + # 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(): + # Each processor receives its own sentinel after all raw outputs are queued. + 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 as e: + pos = data_point.get(self.cfg.async_position_key, "?") + LOG.exception("postprocess/evaluate failed for async_pos=%s", pos) + # 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): + # 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(): + # 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: + 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: + stop = True + break + batch.append(nxt) + _flush_batch(batch) + if stop: + return + + processor_tasks = [asyncio.create_task(_processor()) for _ in range(num_processors)] + writer_task = asyncio.create_task(_writer()) + + # 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): + try: + await self._generate_and_save_datapoint(dp, data, fout, pbar) + except Exception as e: + pos = dp.get(self.cfg.async_position_key, "?") + LOG.exception("datapoint async_pos=%s crashed", pos) + # 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() + + # 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: + # 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: + 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 one worker slot for each completed task. + 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))) + + # 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) + await self._output_queue.put(None) + await writer_task + self._raw_queue = None + self._output_queue = None pbar.close() @@ -912,13 +1038,21 @@ 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) + # 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) 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/azure.py b/nemo_skills/inference/model/azure.py index e097e4f80a..951fa8325d 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 requires provider-specific endpoint and authentication handling + # from 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..d3fe428cb1 100644 --- a/nemo_skills/inference/model/base.py +++ b/nemo_skills/inference/model/base.py @@ -13,8 +13,10 @@ # limitations under the License. import abc import asyncio +import inspect import logging import os +import uuid from enum import Enum from typing import Union @@ -24,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 @@ -73,6 +76,17 @@ class BaseModel: # Litellm provider name MODEL_PROVIDER = "openai" + # 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 parameter emitted by a + # _build_chat_request_params implementation. + _LITELLM_ONLY_CHAT_PARAMS = frozenset({"allowed_openai_params"}) + def __init__( self, model: str, @@ -95,6 +109,10 @@ def __init__( output_dir: str | None = None, # Request tokenizer initialization independent of soft_fail require_tokenizer: bool = False, + # 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 self.model_name_or_path = model @@ -167,9 +185,51 @@ 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 - self.concurrent_semaphore = asyncio.Semaphore(2048) + # Gate in-flight requests at the configured model-level limit. + self.concurrent_semaphore = asyncio.Semaphore(concurrent_requests_limit) + + # 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 + + # 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, + base_url=self.base_url, + http_client=DefaultAioHttpClient( + limits=httpx.Limits( + max_connections=aiohttp_limit, + max_keepalive_connections=aiohttp_limit, + ), + ), + # 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: + # 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 +245,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 +302,13 @@ 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 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 + @with_context_retry async def generate_async( self, @@ -261,6 +331,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 +360,12 @@ async def generate_async( "response_format": response_format, } - # TODO: remove this after we no longer use gpt-oss or it's fixed in vllm + # 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")) + transient_count = 0 async with self.concurrent_semaphore: while retry_count <= max_retries: @@ -299,7 +373,41 @@ 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) + # 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 + # 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 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() + 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 + # 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 +417,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 +428,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 +458,38 @@ async def generate_async( } else: raise e + except ( + httpx.ConnectError, + httpx.ReadError, + httpx.RemoteProtocolError, + httpx.PoolTimeout, + httpx.NetworkError, + openai.APIConnectionError, + ConnectionError, + ConnectionResetError, + ) as e: + # 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)) + 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 +526,20 @@ def _parse_completion_response( return result + @staticmethod + def _extract_reasoning(message_or_delta) -> str | None: + """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) + 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 +551,25 @@ 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 + # Normalize either supported reasoning field in the result. + reasoning = self._extract_reasoning(choice.message) + if reasoning: + result["reasoning_content"] = reasoning + + # Preserve the complete usage payload, including explicit nulls, so + # provider-specific usage metadata remains available. + try: + result["usage"] = response.usage.model_dump() + except AttributeError: + # Support non-pydantic usage objects. + 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 +635,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 - ) + # 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: cur_delta = chunk.choices[0].text @@ -561,6 +729,9 @@ def _serialize_output(self, response): async def _stream_chat_chunks_async(self, response): async for chunk in response: + # Usage-only and keepalive chunks have no choices to process. + 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..4854f5fb0b 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", 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..98bd743d45 100644 --- a/nemo_skills/inference/model/vllm_multimodal.py +++ b/nemo_skills/inference/model/vllm_multimodal.py @@ -69,6 +69,10 @@ class VLLMMultimodalModel(VLLMModel): ) """ + # Multimodal request and response transformations require litellm, so this + # overrides VLLMModel's native-client support. + 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..6d5dabd5a1 100644 --- a/nemo_skills/pipeline/utils/cluster.py +++ b/nemo_skills/pipeline/utils/cluster.py @@ -160,6 +160,13 @@ def parse_kwargs(kwargs: str | dict | None, **extra_kwargs) -> dict | None: return full_kwargs +# SSH tunnel settings are launcher-local and are not forwarded automatically. +_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,6 +222,14 @@ def get_env_variables(cluster_config): "HF_TOKEN", "NGC_API_KEY", } + # Forward NEMO_SKILLS_* runtime configuration except launcher-local settings. + for name in os.environ: + 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..54a3ed9d21 --- /dev/null +++ b/tests/test_async_loop_integrity.py @@ -0,0 +1,73 @@ +# 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. +"""Tests for restoring ordered output from partial async result files.""" + +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): + # Position 1 is absent; the remaining records retain their relative 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): + # Positionless records are ignored. + 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 new file mode 100644 index 0000000000..25c2f439ed --- /dev/null +++ b/tests/test_native_openai_path.py @@ -0,0 +1,257 @@ +# 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 BaseModel's native AsyncOpenAI aiohttp path.""" + +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 + + +# --------------------------------------------------------------------------- +# Reasoning extraction +# --------------------------------------------------------------------------- +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(): + 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(): + # 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)) + assert result["generation"] == "the answer" + assert result["reasoning_content"] == "the thinking" + assert result["num_generated_tokens"] == 5 + + +# --------------------------------------------------------------------------- +# Streaming chunk parsing +# --------------------------------------------------------------------------- +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"}] + + +# --------------------------------------------------------------------------- +# Native-client eligibility +# --------------------------------------------------------------------------- +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 + + +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 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) + + 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}" + # 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): + # Pydantic response formats require LiteLLM because the native completion + # method does not accept model classes. + 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) + + asyncio.run( + model.generate_async( + prompt=[{"role": "user", "content": "hi"}], + tokens_to_generate=8, + temperature=0.0, + top_p=0.95, + response_format=_RF, + ) + ) + + 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 bb2827d48f..03f084c500 100644 --- a/tests/test_pipeline_utils.py +++ b/tests/test_pipeline_utils.py @@ -457,3 +457,33 @@ 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_* values are forwarded except launcher-local SSH settings.""" + from nemo_skills.pipeline.utils.cluster import get_env_variables + + # 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") + # 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") + + 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" diff --git a/tests/test_prompt_cache_routing.py b/tests/test_prompt_cache_routing.py new file mode 100644 index 0000000000..6bceba4138 --- /dev/null +++ b/tests/test_prompt_cache_routing.py @@ -0,0 +1,49 @@ +# 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. +"""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(): + # Request IDs do not depend on cache affinity. + 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(): + # 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") + 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(): + # 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" + assert params["extra_headers"]["X-Request-Id"] == "preset"