Speed up local vLLM ASR client requests

[DongjiGao]

Summary

• Replace httpcore's async connection-pool assignment hot path (httpcore 1.0.9 is O(N²+N·M)) with a single-pass O(N+M) rewrite ported from Fix 2 bugs and improve performance from O(N²+NxM) to O(N+M) in _assign_requests_to_connections encode/httpcore#1035; HTTP/2 pools fall back to the original implementation. This also fixes the two logic bugs that PR documents: the max_keepalive_connections idle-count miscount, and the HTTP/1.1 double-assign that can raise ConnectionNotAvailable.
• Send local vLLM multimodal audio as file:// audio URLs by default for local audio_url requests, avoiding inline base64 payloads (base64 retained for external APIs and as an explicit fallback).
• Add a shared make_audio_url_content_block helper and expand vLLM audio tests (file:// vs base64 content, plus direct VLLMMultimodalModel.__init__ tests for audio_as_path defaulting/validation).

Test plan

• python3 -m py_compile nemo_skills/inference/generate.py nemo_skills/inference/model/vllm_multimodal.py nemo_skills/inference/model/audio_utils.py
• pytest tests/test_vllm_audio.py → 16 passed; ruff check + ruff format --check clean.
• httpcore rewrite validated for assignment-decision equivalence + liveness across 20k randomized pool states (vs the original 1.0.9 implementation).
• Full 8-chunk ASR leaderboard end-to-end: WER (%) 6.21, no throughput regression; controlled traced runs show the patch applied (max_queued=1, max_conn=512). Residual wall-clock variance traces to server-startup / chunk-start skew, not the request path.

Summary by CodeRabbit

• New Features

  • Added audio_as_path option to multimodal models to represent local audio as file:// URLs instead of inline base64.
  • Added support for producing audio URL content blocks for URL-based audio inputs.
  • Improved async connection pooling to make request assignment more efficient.

• Tests

  • Expanded tests to cover audio_as_path defaults, fallback to base64, validation for file:// and data:audio values, and related edge cases.

[DongjiGao]

@pzelasko

[coderabbitai]

📝 Walkthrough

Walkthrough

Optional file-path (file://) audio inputs added to VLLMMultimodalModel via audio_as_path, a shared make_audio_url_content_block helper, expanded tests for both path and base64 modes, and an idempotent httpcore AsyncConnectionPool assignment monkey-patch applied at generation startup.

Changes

Audio Path Representation and Async Optimizations

Layer / File(s)	Summary
Audio URL content block utility nemo_skills/inference/model/audio_utils.py, nemo_skills/inference/model/__init__.py	New helper make_audio_url_content_block() constructs vLLM-style audio_url blocks from URL strings and is exported through the model package.
VLLMMultimodalModel audio_as_path feature nemo_skills/inference/model/vllm_multimodal.py	Adds `audio_as_path: bool
Audio representation test coverage tests/test_vllm_audio.py	Updates validation to accept file:// audio_url, adjusts fixtures to control audio_as_path, adds tests for path and base64 fallback, defaults, overrides, and invalid configurations, and introduces a VLLM init stubbing helper.
httpcore connection pool optimization nemo_skills/inference/generate.py	Adds _patch_httpcore_connection_pool_assignment() that idempotently monkey-patches AsyncConnectionPool._assign_requests_to_connections with a single-pass O(N+M) assigner and is invoked at GenerationTask.async_loop startup; preserves HTTP/2 fallback.

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~45 minutes

Suggested labels

run GPU tests

Suggested reviewers

• melllinia

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title 'Speed up local vLLM ASR client requests' directly and specifically describes the main objective of the changeset: optimizing performance for local vLLM audio requests through httpcore patching and file:// URL optimization.
Docstring Coverage	✅ Passed	Docstring coverage is 92.00% which is sufficient. The required threshold is 80.00%.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

Actionable comments posted: 1

🧹 Nitpick comments (1)

tests/test_vllm_audio.py (1)

55-87: ⚡ Quick win

Add direct constructor tests for audio_as_path defaulting/validation.

These fixtures bypass VLLMMultimodalModel.__init__, so the new init behavior isn’t covered (defaulting from mode + rejection of invalid combinations). A small set of direct init tests would lock this down.

🤖 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_vllm_audio.py` around lines 55 - 87, The fixtures patch
VLLMMultimodalModel.__init__ so constructor logic (defaulting audio_as_path
based on audio_format and rejecting invalid combinations) isn’t tested; add
direct unit tests that call VLLMMultimodalModel.__init__ (real constructor) to
verify defaulting and validation: instantiate with audio_format values
("audio_url", "input_audio") and assert audio_as_path defaults correctly, and
assert that invalid combos (e.g., audio_format="audio_url" with
audio_as_path=False or other disallowed combinations) raise the expected error;
reference VLLMMultimodalModel.__init__, the audio_as_path and audio_format
parameters, and any validation behavior around
enable_audio_chunking/audio_chunk_task_types/chunk_audio_threshold_sec to ensure
constructor enforces invariants.

🤖 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 298-309: The cleanup loop is comparing len(self._connections)
(total connections) to self._max_keepalive_connections and thus may close too
many idle connections; change the surplus-idle check to count idle connections
only (e.g. compute idle_count = sum(1 for c in self._connections if c.is_idle())
or an equivalent generator/scalar) and use idle_count >
self._max_keepalive_connections in the condition that evaluates
connection.is_idle(); keep using the same loop over self._connections and
connection.is_idle()/is_closed()/has_expired() methods but base the surplus
decision on the idle-only count.

---

Nitpick comments:
In `@tests/test_vllm_audio.py`:
- Around line 55-87: The fixtures patch VLLMMultimodalModel.__init__ so
constructor logic (defaulting audio_as_path based on audio_format and rejecting
invalid combinations) isn’t tested; add direct unit tests that call
VLLMMultimodalModel.__init__ (real constructor) to verify defaulting and
validation: instantiate with audio_format values ("audio_url", "input_audio")
and assert audio_as_path defaults correctly, and assert that invalid combos
(e.g., audio_format="audio_url" with audio_as_path=False or other disallowed
combinations) raise the expected error; reference VLLMMultimodalModel.__init__,
the audio_as_path and audio_format parameters, and any validation behavior
around enable_audio_chunking/audio_chunk_task_types/chunk_audio_threshold_sec to
ensure constructor enforces invariants.

🪄 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: 9f87cae9-790a-44e7-9b06-42f9498ecafe

📥 Commits

Reviewing files that changed from the base of the PR and between 69fbc9a and 47f27fb.

📒 Files selected for processing (5)

• nemo_skills/inference/generate.py
• nemo_skills/inference/model/__init__.py
• nemo_skills/inference/model/audio_utils.py
• nemo_skills/inference/model/vllm_multimodal.py
• tests/test_vllm_audio.py

[DongjiGao]

Addressed in eca6b2b (patch) and 701cbf0 (tests).

Rather than only fix the max_keepalive count, I replaced the assignment routine with a single-pass O(N+M) port of encode/httpcore#1035. This resolves both issues that PR documents:

• the max_keepalive_connections miscount — it now enforces the limit on the count of idle connections, not len(self._connections); and
• the HTTP/1.1 double-assign (ConnectionNotAvailable) bug — each connection is popped from the available bucket as it is assigned, so one connection is never handed to two requests in a single pass.

HTTP/2 pools fall back to the original implementation, since the single-pass path intentionally omits the per-connection stream-capacity tracking #1035 adds for multiplexing.

Validated for assignment-decision equivalence + liveness across 20k randomized pool states, and end-to-end on the 8-chunk ASR leaderboard (WER 6.21, no throughput regression).

For the test nitpick: added direct VLLMMultimodalModel.__init__ tests covering audio_as_path defaulting (by mode/format) and the validation that rejects path mode for external APIs / non-audio_url formats.

[coderabbitai]

Tip

For best results, initiate chat on the files or code changes.

@DongjiGao Thanks for the detailed write-up! The single-pass O(N+M) port of encode/httpcore#1035 is a much more principled fix than patching only the keepalive count — resolving both the miscount and the double-assign race in one pass is clearly the right call. The rationale for keeping HTTP/2 pools on the original path (stream-capacity multiplexing requires per-connection tracking that the simplified path skips) is sound and worth having documented in the commit message / code comment as-is.

The VLLMMultimodalModel.__init__ unit tests covering the defaulting logic and the path-mode-for-external-API rejection case close the gap that was flagged — appreciate that those are direct constructor tests rather than just exercising the behaviour through content_text_to_list.

🐇✨

[pzelasko]

LGTM, very welcome speedup, great work!

[DongjiGao]

Recreated as in-repo branch PR #1487 so CI can run (fork PRs don't receive CI secrets). Closing this fork PR in favor of #1487.