High-throughput async inference: native aiohttp fast-path, scalable async loop, and request-routing keys#1502
High-throughput async inference: native aiohttp fast-path, scalable async loop, and request-routing keys#1502gwarmstrong wants to merge 5 commits into
Conversation
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 <gwarmstrong@users.noreply.github.com>
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughThis PR adds native OpenAI-compatible request handling and prompt-cache routing, expands async generation into a bounded pipeline, updates async ordering restoration, and forwards selected ChangesNative OpenAI fast-path and request handling
Async generation pipeline and environment passthrough
Estimated code review effort: 4 (Complex) | ~60 minutes Suggested labels: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 5
🧹 Nitpick comments (4)
nemo_skills/pipeline/utils/cluster.py (1)
218-227: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winBlanket
NEMO_SKILLS_*prefix passthrough may forward unintended vars.Any environment variable prefixed
NEMO_SKILLS_in the launching shell is auto-forwarded to the remote job, not just the intended tuning knobs (NEMO_SKILLS_OPENAI_AIOHTTP,NEMO_SKILLS_DISABLE_UVLOOP, etc.). If the prefix is reused for local-only debug/test vars in the future (or by other tooling), they'd silently leak into launched jobs/containers. Consider an explicit allowlist of known tuning-knob names instead of a blanket prefix scan, to keep the passthrough surface intentional.♻️ Example of a tighter allowlist approach
- 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) + known_tuning_env_vars = { + "NEMO_SKILLS_OPENAI_AIOHTTP", + "NEMO_SKILLS_DISABLE_UVLOOP", + "NEMO_SKILLS_TRANSIENT_RETRIES", + # add other explicit knobs here + } + for name in known_tuning_env_vars: + if name in os.environ and name not in optional_env_vars_to_add: + optional_env_vars_to_add.add(name)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@nemo_skills/pipeline/utils/cluster.py` around lines 218 - 227, The passthrough logic in the cluster environment setup is too broad because it forwards every NEMO_SKILLS_* variable from os.environ into optional_env_vars_to_add. Tighten this in the env-collection block by replacing the prefix scan with an explicit allowlist of known tuning variables (for example the existing NEMO_SKILLS_OPENAI_AIOHTTP, NEMO_SKILLS_DISABLE_UVLOOP, and NEMO_SKILLS_TRANSIENT_RETRIES names), so only intended knobs are exported by the launcher.nemo_skills/inference/model/base.py (3)
213-254: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueBroad
except Exceptionaround native client construction.The
try/except Exceptionswallows any error fromAsyncOpenAI(...)/DefaultAioHttpClient(...)construction, not just the documentedImportErrorcase. The repo's own guideline is explicit about not catching exceptions that aren't expected to normally be raised. This is mitigated somewhat by the informativeLOG.warning, but it also silently degrades performance (falls back to the slower litellm/httpx path) for genuinely unexpected bugs (e.g. a typo inhttpx.Limitskwargs) instead of surfacing them.Consider narrowing to
except (ImportError, TypeError, ValueError) as e:(the realistic set for "aiohttp extra missing/partially installed") so truly unexpected errors still surface.As per coding guidelines, "Don't catch exceptions when we don't expect them to be normally raised... let the code fail when something unexpected happens." Also flagged by static analysis (Ruff BLE001, "Do not catch blind exception:
Exception").🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@nemo_skills/inference/model/base.py` around lines 213 - 254, The native AsyncOpenAI client setup in base.py is catching a broad Exception, which hides unexpected construction bugs and silently falls back to litellm/httpx. Narrow the handler around the AsyncOpenAI/DefaultAioHttpClient block in the native_enabled and self.SUPPORTS_NATIVE_OPENAI path to only the expected failures for missing or partially installed aiohttp support, keeping the existing LOG.warning and fallback behavior for those cases while letting other errors surface.Sources: Coding guidelines, Linters/SAST tools
107-117: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick win
concurrent_requests_limitdefault is frozen at import time, not read per-instantiation.
int(os.environ.get("NEMO_SKILLS_MODEL_CONCURRENCY", "65536"))as a parameter default is evaluated once, when thisdef __init__line is executed at module import — not each time a model is constructed. UnlikeNEMO_SKILLS_SSH_SERVER/NEMO_SKILLS_SSH_KEY_PATH, which are read viaos.getenv()inside the method body, any change to this env var after the module is first imported (e.g. in-process config, tests usingmonkeypatch.setenv) won't be picked up. Typically harmless for a fresh process, but inconsistent with the stated "override via env var" contract in the comment.♻️ Proposed fix: read the env var lazily like the other overrides
- concurrent_requests_limit: int = int(os.environ.get("NEMO_SKILLS_MODEL_CONCURRENCY", "65536")), + concurrent_requests_limit: int | None = None, ): + if concurrent_requests_limit is None: + concurrent_requests_limit = int(os.environ.get("NEMO_SKILLS_MODEL_CONCURRENCY", "65536"))🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@nemo_skills/inference/model/base.py` around lines 107 - 117, The default for concurrent_requests_limit in Base.__init__ is being evaluated at module import time, so later changes to NEMO_SKILLS_MODEL_CONCURRENCY are ignored. Move the environment lookup into the __init__ body and assign the default lazily, matching the pattern already used for the other env-based overrides in Base.__init__. Keep the existing fallback value and preserve the current user-facing behavior, but ensure each model construction reads the latest env var value.
604-625: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winTriple-nested fallback for serializing
usageadds complexity for a rarely-hit corner case.
model_dump()→dict()→ manual field construction is three levels of defensive fallback (with a bareexcept Exceptionat the innermost level, flagged by Ruff BLE001) to build one dict. Given essentially all supported providers here return pydanticusageobjects withmodel_dump(), the two fallback layers mostly guard against an untested/unexpected shape — which per the repo's guideline should probably just fail loudly rather than degrade to a partial dict.Consider narrowing the fallback to a single, more specific
except (AttributeError, TypeError):and dropping the innermost blindexcept Exception, or removing the fallback chain entirely if none of the supported providers realistically hit it.As per coding guidelines, "Don't catch exceptions when we don't expect them to be normally raised... it's much better to just [fail] with a clear error," and "Keep code simple and elegant... minimize conditional checks." Also flagged by static analysis (Ruff BLE001 at line 620).
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@nemo_skills/inference/model/base.py` around lines 604 - 625, The `usage` serialization in `BaseModel` is over-defensive and includes a broad inner exception handler that Ruff flags. Simplify the fallback path around `response.usage` in the `model_dump()` block: keep the primary `model_dump()` handling, and either remove the extra fallback layers entirely or collapse them into a single specific exception catch such as `AttributeError`/`TypeError`. Avoid the bare `except Exception` and the manual partial-dict fallback so unexpected `usage` shapes fail clearly instead of silently degrading.Sources: Coding guidelines, Linters/SAST tools
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@nemo_skills/inference/generate.py`:
- Around line 1013-1020: Stop swallowing datapoint processing failures in the
async output path: the exception handling around postprocess/evaluate in
generate.py currently logs and continues, which can enqueue incomplete results
and break restore_async_order() later. Update the async draining logic in the
relevant handler(s) around postprocess_single_output, evaluate_single_datapoint,
and the similar block near the other referenced section to collect any failure
state while draining the queue, keep processing until the pipeline is drained,
and then raise the captured error before restoring/removing the async file. Also
use direct access for expected keys like _async_position instead of .get where
the position must exist.
- Line 1001: The postprocess worker count is being silently clamped in
generate.py, which hides invalid user input. Update the logic around
self.cfg.postprocess_workers in the generation path to validate the configured
value explicitly and raise a ValueError when it is 0 or negative instead of
forcing num_processors to 1. Keep the check close to the num_processors
assignment so the failure is immediate and easy to trace.
- Around line 1037-1044: The async writer path in the output loop bypasses the
`dump_outputs()` hook, so custom task subclasses lose their overridden output
handling. Update the writer logic in `generate.py` to route each record through
`dump_outputs()` the same way the fallback and prefilled paths do, rather than
writing `json.dumps(...)` directly. Keep the fix localized around the async
batch flush/write section so the `dump_outputs()` behavior is preserved for all
generated records.
In `@nemo_skills/inference/model/base.py`:
- Around line 213-254: The AsyncOpenAI aiohttp client created in BaseModelModel
initialization is never closed, which can leak aiohttp sessions on repeated
model construction. Update the model shutdown path in base.py so the existing
cleanup logic also awaits self._async_openai_client.close() when that client was
created, and keep the current _tunnel cleanup intact. Use the existing
_async_openai_client field and the destructor/teardown path as the place to add
the async close handling, with a safe fallback if the client was never
initialized.
- Around line 327-347: The native OpenAI request path is passing
prompt_cache_key directly in request_params, which breaks older openai installs;
update the OpenAI-specific call site to route cache_key through extra_body
instead, or enforce a minimum openai version of 1.98.0 in requirements. Use
_apply_routing_keys to keep X-Request-Id behavior unchanged, and make the later
AsyncOpenAI.chat.completions.create path follow the same compatibility handling
so cache_key never hits unsupported SDK versions.
---
Nitpick comments:
In `@nemo_skills/inference/model/base.py`:
- Around line 213-254: The native AsyncOpenAI client setup in base.py is
catching a broad Exception, which hides unexpected construction bugs and
silently falls back to litellm/httpx. Narrow the handler around the
AsyncOpenAI/DefaultAioHttpClient block in the native_enabled and
self.SUPPORTS_NATIVE_OPENAI path to only the expected failures for missing or
partially installed aiohttp support, keeping the existing LOG.warning and
fallback behavior for those cases while letting other errors surface.
- Around line 107-117: The default for concurrent_requests_limit in
Base.__init__ is being evaluated at module import time, so later changes to
NEMO_SKILLS_MODEL_CONCURRENCY are ignored. Move the environment lookup into the
__init__ body and assign the default lazily, matching the pattern already used
for the other env-based overrides in Base.__init__. Keep the existing fallback
value and preserve the current user-facing behavior, but ensure each model
construction reads the latest env var value.
- Around line 604-625: The `usage` serialization in `BaseModel` is
over-defensive and includes a broad inner exception handler that Ruff flags.
Simplify the fallback path around `response.usage` in the `model_dump()` block:
keep the primary `model_dump()` handling, and either remove the extra fallback
layers entirely or collapse them into a single specific exception catch such as
`AttributeError`/`TypeError`. Avoid the bare `except Exception` and the manual
partial-dict fallback so unexpected `usage` shapes fail clearly instead of
silently degrading.
In `@nemo_skills/pipeline/utils/cluster.py`:
- Around line 218-227: The passthrough logic in the cluster environment setup is
too broad because it forwards every NEMO_SKILLS_* variable from os.environ into
optional_env_vars_to_add. Tighten this in the env-collection block by replacing
the prefix scan with an explicit allowlist of known tuning variables (for
example the existing NEMO_SKILLS_OPENAI_AIOHTTP, NEMO_SKILLS_DISABLE_UVLOOP, and
NEMO_SKILLS_TRANSIENT_RETRIES names), so only intended knobs are exported by the
launcher.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: bdb60ce6-6f03-4464-8a0d-e4ad125afb39
📒 Files selected for processing (10)
core/requirements.txtnemo_skills/inference/generate.pynemo_skills/inference/model/azure.pynemo_skills/inference/model/base.pynemo_skills/inference/model/openai.pynemo_skills/inference/model/vllm.pynemo_skills/inference/model/vllm_multimodal.pynemo_skills/pipeline/utils/cluster.pytests/test_native_openai_path.pytests/test_prompt_cache_routing.py
| 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)) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Reject invalid postprocess_workers instead of silently clamping.
Line 1001 turns 0/negative user input into 1, so a bad config silently runs with different behavior than requested. Prefer an explicit ValueError.
Proposed fix
- num_processors = max(1, int(self.cfg.postprocess_workers))
+ num_processors = int(self.cfg.postprocess_workers)
+ if num_processors < 1:
+ raise ValueError(f"postprocess_workers must be >= 1, got {self.cfg.postprocess_workers}")As per coding guidelines, “Avoid cases when user-passed parameters are unused” and “the code should fail” for unsupported arguments.
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| num_processors = max(1, int(self.cfg.postprocess_workers)) | |
| num_processors = int(self.cfg.postprocess_workers) | |
| if num_processors < 1: | |
| raise ValueError(f"postprocess_workers must be >= 1, got {self.cfg.postprocess_workers}") |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@nemo_skills/inference/generate.py` at line 1001, The postprocess worker count
is being silently clamped in generate.py, which hides invalid user input. Update
the logic around self.cfg.postprocess_workers in the generation path to validate
the configured value explicitly and raise a ValueError when it is 0 or negative
instead of forcing num_processors to 1. Keep the check close to the
num_processors assignment so the failure is immediate and easy to trace.
Source: Coding guidelines
| 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)) |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
Don’t swallow datapoint failures; raise after draining the pipeline.
These handlers log and continue. A postprocess/eval failure can enqueue raw output without _async_position; a generation failure can skip a position entirely. restore_async_order() then fails late or writes incomplete results. Collect failures while draining queues, then raise before restoring/removing the async file.
Proposed pattern
+ errors = []
+
async def _processor():
@@
- try:
+ 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, "?")
+ except Exception as exc:
+ errors.append(exc)
+ pos = data_point[self.cfg.async_position_key]
LOG.exception("postprocess/evaluate failed for async_pos=%s", pos)
+ continue
await self._output_queue.put((output, data_point))
@@
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, "?")
+ except Exception as exc:
+ errors.append(exc)
+ pos = dp[self.cfg.async_position_key]
LOG.exception("datapoint async_pos=%s crashed", pos)
@@
await self._output_queue.put(None)
await writer_task
+ if errors:
+ raise RuntimeError(f"{len(errors)} datapoint(s) failed during async generation")
self._raw_queue = None
self._output_queue = NoneAs per coding guidelines, “Don’t catch exceptions when we don’t expect them to be normally raised” and prefer direct key access when keys are expected.
Also applies to: 1072-1077
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@nemo_skills/inference/generate.py` around lines 1013 - 1020, Stop swallowing
datapoint processing failures in the async output path: the exception handling
around postprocess/evaluate in generate.py currently logs and continues, which
can enqueue incomplete results and break restore_async_order() later. Update the
async draining logic in the relevant handler(s) around
postprocess_single_output, evaluate_single_datapoint, and the similar block near
the other referenced section to collect any failure state while draining the
queue, keep processing until the pipeline is drained, and then raise the
captured error before restoring/removing the async file. Also use direct access
for expected keys like _async_position instead of .get where the position must
exist.
Source: Coding guidelines
There was a problem hiding this comment.
@gwarmstrong not sure if this is correct but if it is, we should probably fail on exceptions
There was a problem hiding this comment.
@Kipok, agreed — this is a real gap. Right now _processor() and _safe_one() both catch broad Exception, log it, and keep going:
- A postprocess/eval failure in
_processor()logs and then still doesawait self._output_queue.put((output, data_point))with whatever partialoutputexists — ifpostprocess_single_outputfailed before setting_async_position-related fields or before mergingoriginal_data_point, that record can be malformed or missing the position key entirely. - A crash in
_generate_and_save_datapointinside_safe_one()is caught, logged, and silently dropped — no placeholder is emitted for that position at all, so that datapoint's slot is simply missing from the async file.
Both cases mean restore_async_order() either writes an incomplete/corrupted output file or (per the current implementation) silently drops that position with just a warning log — so a large async run could finish "successfully" while quietly missing results, which is easy to miss without inspecting logs closely.
This does match the coding guideline of not swallowing unexpected exceptions — the fix should be to track these failures while draining the pipeline (or emit clearly-marked placeholder records with the position preserved) and then raise once the queues are fully drained, rather than letting the run report success with silently missing/malformed data. Happy to help draft that once you and @gwarmstrong agree on the exact approach (hard fail vs. always-emit-placeholder-then-fail).
| 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 |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Check for any explicit close/aclose handling of the native client elsewhere in the codebase.
rg -n -A3 -B3 '_async_openai_client' nemo_skills/inferenceRepository: NVIDIA-NeMo/Skills
Length of output: 3445
🏁 Script executed:
rg -n -A2 -B2 'aclose\(|close\(|__del__|async with AsyncOpenAI|DefaultAioHttpClient' nemo_skillsRepository: NVIDIA-NeMo/Skills
Length of output: 15419
🏁 Script executed:
sed -n '260,330p' nemo_skills/inference/model/base.pyRepository: NVIDIA-NeMo/Skills
Length of output: 3148
🌐 Web query:
openai python AsyncOpenAI DefaultAioHttpClient aclose close documentation
💡 Result:
In the OpenAI Python SDK, the recommended way to manage the lifecycle of an AsyncOpenAI client is through an asynchronous context manager, which automatically handles the closing of underlying HTTP connections [1][2][3]. ### Closing the Client While the SDK may automatically close connections when the client is garbage collected, relying on this is not recommended for production applications [1][3]. Instead, you should use one of the following methods: 1. Asynchronous Context Manager (Recommended) This is the cleanest approach and ensures the underlying http_client is closed correctly when the block exits [2][3]. python async with AsyncOpenAI() as client: # perform your operations pass # client is automatically closed here 2. Explicit await client.close() If you cannot use a context manager (e.g., in a long-running service where the client is stored in application state), you must manually call await client.close() during application shutdown [2][3]. python # During application startup client = AsyncOpenAI() # During application shutdown await client.close() ### DefaultAioHttpClient The DefaultAioHttpClient is a utility provided by the SDK to use aiohttp as the underlying HTTP backend instead of the default httpx client [1][4]. It is used when instantiating AsyncOpenAI: python from openai import AsyncOpenAI, DefaultAioHttpClient async with AsyncOpenAI(http_client=DefaultAioHttpClient()) as client: # perform operations pass Regardless of whether you use the default httpx transport or DefaultAioHttpClient, the AsyncOpenAI client exposes an await client.close() method that shuts down the underlying HTTP client [5][4]. Note: Avoid using aclose() directly on client instances; the SDK exposes the standard close() method, which is an async method (making it awaitable) that internally calls the appropriate closure method on the underlying transport [5][4]. You may encounter aclose() in the context of specific stream objects or underlying library components, but for the top-level AsyncOpenAI client, await client.close() is the official public API [5][2][6].
Citations:
- 1: https://github.com/openai/openai-python/blob/main/README.md
- 2: https://openai-openai-python-73.mintlify.app/concepts/async
- 3: Why docs do not use context manager for several examples openai/openai-python#2779
- 4: https://github.com/openai/openai-python/blob/main/src/openai/_base_client.py
- 5: https://github.com/openai/openai-python/blob/5e8f09c2/src/openai/_base_client.py
- 6: AsyncStream missing aclose() causes AttributeError in streaming structured output openai/openai-python#2853
Close the native OpenAI client on shutdown.
self._async_openai_client wraps an aiohttp-backed AsyncOpenAI, but this class never calls await client.close() anywhere. __del__ only stops _tunnel, so repeated model construction in one process can leave aiohttp sessions/connectors open.
🧰 Tools
🪛 Ruff (0.15.20)
[warning] 243-243: Do not catch blind exception: Exception
(BLE001)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@nemo_skills/inference/model/base.py` around lines 213 - 254, The AsyncOpenAI
aiohttp client created in BaseModelModel initialization is never closed, which
can leak aiohttp sessions on repeated model construction. Update the model
shutdown path in base.py so the existing cleanup logic also awaits
self._async_openai_client.close() when that client was created, and keep the
current _tunnel cleanup intact. Use the existing _async_openai_client field and
the destructor/teardown path as the place to add the async close handling, with
a safe fallback if the client was never initialized.
| @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 | ||
|
|
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🌐 Web query:
Does openai-python chat.completions.create() accept prompt_cache_key as a typed top-level parameter, and since which version?
💡 Result:
Yes, the openai-python library's chat.completions.create method now accepts prompt_cache_key as a typed top-level parameter [1][2]. This parameter was added to the library in version 1.98.0 [3]. Prior to this version, developers attempting to use prompt_cache_key as a parameter would encounter a TypeError, and it was recommended to use the extra_body parameter as a workaround [4]. The prompt_cache_key parameter is used to influence routing and improve cache hit rates by allowing you to associate requests with specific cache identifiers, effectively replacing the previously used user field for this purpose [1][5][6].
Citations:
- 1: https://github.com/openai/openai-python/blob/main/src/openai/resources/chat/completions/completions.py
- 2: openai/openai-python@00b49ae
- 3: openai/openai-python@v1.97.2...v1.98.0
- 4: python sdk prompt_cache_key support? openai/openai-python#2499
- 5: https://developers.openai.com/api/reference/resources/chat/subresources/completions/methods/create/
- 6: https://developers.openai.com/api/docs/guides/prompt-caching
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '\n== openai dependency pins ==\n'
git ls-files | rg '(^|/)(requirements\.txt|pyproject\.toml|uv\.lock|poetry\.lock|Pipfile|Pipfile\.lock)$|core/requirements\.txt|requirements/pipeline\.txt'
printf '\n== relevant openai constraints ==\n'
rg -n --hidden --no-heading 'openai|litellm|prompt_cache_key|extra_body|allowed_openai_params|SUPPORTS_NATIVE_OPENAI' \
core nemo_skills . 2>/dev/null | sed -n '1,220p'
printf '\n== target file excerpt ==\n'
sed -n '300,480p' nemo_skills/inference/model/base.pyRepository: NVIDIA-NeMo/Skills
Length of output: 33390
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '\n== core/requirements.txt ==\n'
sed -n '1,80p' core/requirements.txt
printf '\n== pyproject openai/litellm constraints ==\n'
sed -n '1,140p' pyproject.toml
printf '\n== nemo_skills/inference/model/openai.py ==\n'
sed -n '1,240p' nemo_skills/inference/model/openai.pyRepository: NVIDIA-NeMo/Skills
Length of output: 12682
Route prompt_cache_key through extra_body on the native OpenAI path, or pin openai>=1.98.0. core/requirements.txt doesn’t set a minimum openai version, so older installs will raise TypeError when cache_key is used with the native AsyncOpenAI.chat.completions.create() call. Same issue applies to the later native OpenAI call site.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@nemo_skills/inference/model/base.py` around lines 327 - 347, The native
OpenAI request path is passing prompt_cache_key directly in request_params,
which breaks older openai installs; update the OpenAI-specific call site to
route cache_key through extra_body instead, or enforce a minimum openai version
of 1.98.0 in requirements. Use _apply_routing_keys to keep X-Request-Id behavior
unchanged, and make the later AsyncOpenAI.chat.completions.create path follow
the same compatibility handling so cache_key never hits unsupported SDK
versions.
…indings 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 <gwarmstrong@users.noreply.github.com>
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@nemo_skills/inference/generate.py`:
- Around line 1176-1184: The reorder logic in generate.py is dropping async
records that are missing self.cfg.async_position_key, which causes silent data
loss during finalization. Update the reordering path around ordered_generations
and the gen_dict.pop(key, None) handling so positionless records are preserved
instead of skipped, and make the failure explicit if the key is truly required.
Keep the recovery flow intact in the finalization path and ensure
generate/reorder does not delete the only recoverable data when a record lacks
the async position key.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: d73f83ba-187f-4a2e-971a-3b70bca598d2
📒 Files selected for processing (6)
nemo_skills/inference/generate.pynemo_skills/inference/model/base.pynemo_skills/pipeline/utils/cluster.pytests/test_async_loop_integrity.pytests/test_native_openai_path.pytests/test_pipeline_utils.py
🚧 Files skipped from review as they are similar to previous changes (2)
- nemo_skills/pipeline/utils/cluster.py
- nemo_skills/inference/model/base.py
| 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 |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
Don't drop positionless async records during finalization.
If a record reaches reorder without self.cfg.async_position_key, this branch logs a warning, omits the row from the final JSONL, and then deletes the only recovery file. That is silent data loss.
Proposed fix
for gen_dict in generations:
- async_pos = gen_dict.pop(key, None)
- if async_pos is None:
- LOG.warning("async record missing %s, dropping from reorder", key)
- continue
+ if key not in gen_dict:
+ raise ValueError(f"async record missing required key {key!r} during reorder")
+ async_pos = gen_dict.pop(key)
ordered_generations[async_pos] = gen_dictAs per coding guidelines, “Don't use .get() for accessing dictionary keys if the code expects them to be present” and “Errors should never pass silently.”
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| 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 | |
| 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: | |
| if key not in gen_dict: | |
| raise ValueError(f"async record missing required key {key!r} during reorder") | |
| async_pos = gen_dict.pop(key) | |
| ordered_generations[async_pos] = gen_dict |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@nemo_skills/inference/generate.py` around lines 1176 - 1184, The reorder
logic in generate.py is dropping async records that are missing
self.cfg.async_position_key, which causes silent data loss during finalization.
Update the reordering path around ordered_generations and the gen_dict.pop(key,
None) handling so positionless records are preserved instead of skipped, and
make the failure explicit if the key is truly required. Keep the recovery flow
intact in the finalization path and ensure generate/reorder does not delete the
only recoverable data when a record lacks the async position key.
Source: Coding guidelines
… 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 <gwarmstrong@users.noreply.github.com>
…e 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 <gwarmstrong@users.noreply.github.com>
Signed-off-by: gwarmstrong <gwarmstrong@users.noreply.github.com>
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
tests/test_pipeline_utils.py (1)
479-487: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAssert forwarded values, not just key presence.
This test would still pass if
get_env_variables()forwarded the right names with wrong values. Use direct key access and compare the expected env values.Proposed fix
- 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" + expected_forwarded = { + "NEMO_SKILLS_OPENAI_AIOHTTP": "1", + "NEMO_SKILLS_DISABLE_UVLOOP": "1", + "NEMO_SKILLS_OPENAI_BASE_URL": "http://model:8000/v1", + "NEMO_SKILLS_SANDBOX_HOST": "sandbox", + "NEMO_SKILLS_CONFIG_DIR": "/home/user/cluster_configs", + "NEMO_SKILLS_SOME_FUTURE_KNOB": "x", + } + for forwarded, value in expected_forwarded.items(): + assert env[forwarded] == value🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/test_pipeline_utils.py` around lines 479 - 487, The test only checks that get_env_variables() forwards the expected keys, so it can still pass with incorrect values. Update the assertions in test_pipeline_utils.py to read the forwarded values directly from env and compare them against the expected source env values for each symbol, using the existing get_env_variables() result and the listed NEMO_SKILLS_* keys.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@nemo_skills/inference/generate.py`:
- Around line 992-996: The ramp scheduling in generate.py ignores very small
ramp_rate_per_sec values because ramp_batch is forced to at least 50 in the ramp
logic. Update the ramp calculation around ramp_rate, ramp_batch, and
ramp_interval so user-provided low rates are honored exactly, without a minimum
burst that exceeds the configured limit. Keep the behavior in the same ramp
control block and adjust the batching formula to derive batch size directly from
ramp_rate_per_sec while preserving the smooth quarter-second pacing intent.
---
Nitpick comments:
In `@tests/test_pipeline_utils.py`:
- Around line 479-487: The test only checks that get_env_variables() forwards
the expected keys, so it can still pass with incorrect values. Update the
assertions in test_pipeline_utils.py to read the forwarded values directly from
env and compare them against the expected source env values for each symbol,
using the existing get_env_variables() result and the listed NEMO_SKILLS_* keys.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 3009abb2-0558-4723-824f-8d00c901a929
📒 Files selected for processing (10)
nemo_skills/inference/generate.pynemo_skills/inference/model/azure.pynemo_skills/inference/model/base.pynemo_skills/inference/model/openai.pynemo_skills/inference/model/vllm_multimodal.pynemo_skills/pipeline/utils/cluster.pytests/test_async_loop_integrity.pytests/test_native_openai_path.pytests/test_pipeline_utils.pytests/test_prompt_cache_routing.py
💤 Files with no reviewable changes (1)
- nemo_skills/inference/model/openai.py
🚧 Files skipped from review as they are similar to previous changes (7)
- nemo_skills/inference/model/vllm_multimodal.py
- nemo_skills/inference/model/azure.py
- tests/test_prompt_cache_routing.py
- nemo_skills/pipeline/utils/cluster.py
- tests/test_async_loop_integrity.py
- tests/test_native_openai_path.py
- nemo_skills/inference/model/base.py
| 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 |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Honor low ramp_rate_per_sec values.
With ramp_rate_per_sec=1, Line 995 still creates an initial burst of 50 tasks before sleeping, which violates the configured ramp limit and can overload endpoints the knob is meant to protect.
Proposed fix
- ramp_batch = max(50, ramp_rate // 4)
+ ramp_batch = max(1, ramp_rate // 4)
ramp_interval = ramp_batch / ramp_rateAs per coding guidelines, user-passed parameters should not be silently ignored or changed.
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| 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 | |
| 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(1, ramp_rate // 4) | |
| ramp_interval = ramp_batch / ramp_rate |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@nemo_skills/inference/generate.py` around lines 992 - 996, The ramp
scheduling in generate.py ignores very small ramp_rate_per_sec values because
ramp_batch is forced to at least 50 in the ramp logic. Update the ramp
calculation around ramp_rate, ramp_batch, and ramp_interval so user-provided low
rates are honored exactly, without a minimum burst that exceeds the configured
limit. Keep the behavior in the same ramp control block and adjust the batching
formula to derive batch size directly from ramp_rate_per_sec while preserving
the smooth quarter-second pacing intent.
Source: Coding guidelines
| # 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 |
There was a problem hiding this comment.
could you clean this up please - I think current retry count is never incremented and while loop is basically while True and there is a separate counter for transient retries. Should either have explicity while True and remove these outdated vars or just use retry_count without adding transient_count like it was before
| # 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) |
There was a problem hiding this comment.
this looks a bit fragile as a logic for routing to litellm, but I guess we will retire all this code soon anyway, so hopefully ok for now. And gym doesn't use litellm at all
| # Set NEMO_SKILLS_DISABLE_UVLOOP to use the default asyncio event loop. | ||
| if not os.environ.get("NEMO_SKILLS_DISABLE_UVLOOP"): | ||
| try: | ||
| import uvloop |
There was a problem hiding this comment.
does this need to be added to requirements?
| 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)) |
There was a problem hiding this comment.
@gwarmstrong not sure if this is correct but if it is, we should probably fail on exceptions
| async def _safe_one(dp): | ||
| try: | ||
| await self._generate_and_save_datapoint(dp, data, fout, pbar) | ||
| except Exception as e: |
There was a problem hiding this comment.
I guess similar here, why do we want to catch all exceptions?
2 participants
Summary
Makes
ns generatescale cleanly to high concurrency (5k+ in-flight requests) against OpenAI-compatible endpoints. At high concurrency the old litellm/httpx dispatch was the GIL-bound bottleneck (~95% of a single core in the async loop, per py-spy); this branch moves the hot path onto a C-extension-backed aiohttp transport and removes the per-completion write/postprocess serialization in the generate loop. It also adds optional, provider-agnostic request-routing keys that a proxy/gateway may use (a plain endpoint ignores them).The fast path is the default and is correct for current vLLM (reasoning field rename) and streaming.
What's included
Native AsyncOpenAI + aiohttp fast-path (
model/base.py)AsyncOpenAIclient using theDefaultAioHttpClient(aiohttp) transport, bypassing litellm's httpx dispatch.SUPPORTS_NATIVE_OPENAIclass attribute → enabled foropenai/vllm/sglang;azure,gemini,megatron, and multimodal stay on litellm. Disable withNEMO_SKILLS_OPENAI_AIOHTTP=0..reasoning(newer vLLM) and.reasoning_content(older) — litellm normalized this for us, but the native path sees vLLM's raw field.Scalable async loop (
inference/generate.py)ramp_rate_per_sec, default 4000) to avoid overflowing the kernel accept queue on prime.foutlock.postprocess_workers, default 16) so a lane frees its slot the instant the HTTP response returns.uvloopwhen available (NEMO_SKILLS_DISABLE_UVLOOP=1to opt out).Request-routing keys (
model/base.py) — optional, provider-agnostic; a plain endpoint ignores them: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 viacache_key=. Requests sharing a value hint a shared prompt prefix. Kept separate fromX-Request-Id(shared vs per-request).Misc
response.usagedict + reasoning-token breakdown into the output.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 actually installed (per CONTRIBUTING, inference deps live here).Config / env knobs
NEMO_SKILLS_OPENAI_AIOHTTP10forces litellmNEMO_SKILLS_OPENAI_AIOHTTP_LIMIT65536NEMO_SKILLS_TRANSIENT_RETRIES3NEMO_SKILLS_DISABLE_UVLOOPmax_concurrent_requests(cfg)512ramp_rate_per_sec(cfg)4000postprocess_workers(cfg)16Testing
Unit (
tests/test_prompt_cache_routing.py,tests/test_native_openai_path.py— 15 tests): routing-key idempotency/affinity decoupling; reasoning extraction from both field names; stream + non-stream parsing; empty-choiceschunk guard; per-provider native eligibility; env-var default/opt-out.Live parity (real model on an OpenAI-compatible endpoint,
nvidia/nemotron-3-nano-30b-a3b):AIOHTTP=0), × non-stream × stream → identical final answer;reasoning_contentcaptured on all four; native transport active only by default and off withAIOHTTP=0.generation/reasoning_content/finish_reason).finish_reason=tool_calls, reassembledget_weather({"city":"Paris"})(the exact chunk shapeToolCallingWrapper._stream_singleconsumes).Known non-blocking notes
num_generated_tokensisNoneon both native and litellm paths (unchanged;ToolCallingWrappercounts via tokenizer). I did not enablestream_options.include_usage, but added the empty-choicesguard so it's safe to turn on later.aclient_session); harmless since the process exits after a run.Summary by CodeRabbit
NEMO_SKILLS_*environment variables into jobs (excluding SSH tunnel controls).