`nemo_rl.models.generation.vllm.vllm_worker_async`#

Module Contents#

Classes#

`CapturedRequest`	Request-local ownership of tokens and processed-media snapshots.
`_AsyncLLMHTTPClient`	Keep HTTP generation on the loop that owns AsyncLLM request state.
`VllmAsyncGenerationWorkerImpl`
`VllmAsyncGenerationWorker`

Functions#

install_http_server_api_key_check

Refuse OpenAI-route requests that do not carry Authorization: Bearer <api_key>.

Data#

`LOGGER`
`HTTP_SERVER_PROTECTED_PATH_PREFIXES`

API#

nemo_rl.models.generation.vllm.vllm_worker_async.LOGGER#: ‘getLogger(…)’

class nemo_rl.models.generation.vllm.vllm_worker_async.CapturedRequest#

Request-local ownership of tokens and processed-media snapshots.

call: nemo_gym.token_id_capture.staging.capture.ActiveCall#: None

prompt_token_ids: list[int]#: None

media: nemo_rl.data.captured_media.CapturedMedia | None#: None

nemo_rl.models.generation.vllm.vllm_worker_async.HTTP_SERVER_PROTECTED_PATH_PREFIXES#: (‘/v1/’, ‘/tokenize’)

nemo_rl.models.generation.vllm.vllm_worker_async.install_http_server_api_key_check( app: fastapi.FastAPI, api_key: str, ) → None#

Refuse OpenAI-route requests that do not carry Authorization: Bearer <api_key>.

Registers an HTTP middleware on app. For a request whose path starts with one of HTTP_SERVER_PROTECTED_PATH_PREFIXES, the Authorization header is compared in constant time with Bearer <api_key>; a missing or different value is answered with HTTP 401 in OpenAI’s error envelope and never reaches the route. Requests to other paths pass through unchanged.

class nemo_rl.models.generation.vllm.vllm_worker_async._AsyncLLMHTTPClient( engine_client: Any, engine_loop: asyncio.AbstractEventLoop, )#

Keep HTTP generation on the loop that owns AsyncLLM request state.

The engine-client surface is explicit. Do not add a __getattr__ fallback. Add each new member here and decide whether it must run on the engine loop.

Initialization

async _run_on_engine_loop( operation: collections.abc.Callable[[], collections.abc.Awaitable[Any]], ) → Any#

generate(

prompt: Any,

sampling_params: Any,

request_id: str,

**kwargs: Any,

) → AsyncGenerator[Any, None]#

async _generate( prompt: Any, sampling_params: Any, request_id: str, kwargs: dict[str, Any], ) → AsyncGenerator[Any, None]#

property errored: bool#

property dead_error: BaseException#

check_admission(n: int = 1, request_id: str | None = None) → None#

Queue-limit admission check vLLM >= 0.29 runs before every response.

OpenAIServing._preflight calls this (vllm-project/vllm#49445, max_num_queued_reqs / max_num_queued_tokens); without it every chat completion 500s with '_AsyncLLMHTTPClient' object has no attribute 'check_admission'. It only reads scheduler config and unfinished-request counters, so it stays off the engine loop like the other status reads above. Raises vLLM’s HTTP-mapped overflow errors.

async is_tracing_enabled() → bool#

class nemo_rl.models.generation.vllm.vllm_worker_async.VllmAsyncGenerationWorkerImpl( config, bundle_indices=None, fraction_of_gpus: float = 1.0, seed=None, extra_env_vars: Optional[list[str]] = None, defer_model_load: bool = False, )#

Bases: nemo_rl.models.generation.vllm.checkpoint_engine.VllmAsyncCheckpointEngineRpcMixin, nemo_rl.models.generation.vllm.vllm_worker.BaseVllmGenerationWorker

async start_gpu_profiling() → None#

Start GPU profiling on the inner TP workers.

The async engine’s collective_rpc is a coroutine; the inherited sync start_gpu_profiling would drop it un-awaited, so the inner TP workers never receive cudaProfilerStart. Await it here so every inner worker starts. This outer actor deliberately does not call cudaProfilerStart itself, so its own nsys session never opens a capture and writes no report.

async stop_gpu_profiling() → None#

Stop GPU profiling on the inner TP workers.

Awaits collective_rpc(“stop_gpu_profiling”) so cudaProfilerStop reaches every inner worker and ends its capture.

_return_routed_experts_enabled() → bool#

_reserve_port() → None#

Bind and listen on a TCP socket to reserve a free port from the OS.

The socket is held open in LISTENING state and later passed directly to uvicorn via the sockets= parameter in server.serve(). The socket is never closed and re-opened, so there is zero gap where another process could steal the port.

load_model() → None#

Load the vLLM model and create the engine.

Called after a deferred init to perform the heavy model loading.

_create_engine(llm_kwargs: dict[str, Any]) → None#

_start_vllm_metrics_logger() → None#

Start a background thread that periodically collects vLLM logger metrics.

Controlled by the required vllm_metrics_logger_interval in vllm_cfg. Runs only on the model-owner actor.

get_vllm_logger_metrics() → dict[str, Any]#

drain_latest_vllm_logger_metrics() → dict[str, Any]#: Return latest samples and prune histories after a telemetry poll.

clear_vllm_logger_metrics() → None#

async post_init_async()#

async get_reserved_url() → Optional[str]#: Return the URL from the reserved socket, available before model loading.

async report_dp_openai_server_base_url() → Optional[str]#

install_token_capture(capture: Any) → None#: Gym’s install_capture seam (the CaptureHost contract).

async setup_token_capture( dp_cfg: dict[str, Any], staging_partition: str, *, capture_media: bool = False, ) → nemo_rl.data.captured_media.MediaColumnSpec | None#

Host ledger-authoritative token capture in this worker.

Fan-out target (token_capture.enabled only): builds the in-worker data-plane client and TQTokenSink, then makes the single install_capture call wiring Gym’s engine-blind capture core + vLLM adapter into this worker. Returns the media column spec this worker pinned the staging column to (None for text-only capture and on non-model-owner ranks, which serve no HTTP and install nothing).

async mooncake_checkpoint( body: dict[str, Any], ) → dict[str, Any] | None#: Run owner-local checkpoint I/O without blocking the actor event loop.

async set_rollout_weight_version(version: int) → None#: Rotate the weight version stamped on subsequent captured calls.

_capture_admission(request: Any) → Any | None#

Parse the ledger’s ng_capture context into a CaptureAdmission.

Returns None unless capture is installed and the request carries the context. The dict itself is never mutated: the admission is the typed, read-only contract that the prefix resolution and begin_call share.

_begin_request_capture( request: Any, prompt_token_ids: list[int], *, admission: Any | None = None, prefix_token_ids: list[int] | None = None, media: nemo_rl.data.captured_media.CapturedMedia | None = None, ) → None#

Admit one ledger-forwarded call into the capture layer.

Called from preprocess_chat once the exact engine prompt is known (post-splice in token-in mode, full render in text mode). No-op unless capture is installed and the request carries the ledger’s ng_capture context.

prefix_token_ids is the prefix resolved by

Meth:: _resolve_admission_prefix. Gym’s begin_call requires it for a staging_chain admission and checks its length against prev_len; violations raise CaptureError before any state is kept.

_capture_request_media( engine_prompt: dict[str, Any], *, admission: Any | None, splice: nemo_rl.models.generation.openai_server_utils.PrefixSplice | None = None, ) → nemo_rl.data.captured_media.CapturedMedia | None#: Run off-loop: resolve retained geometry and snapshot processed pixels.

_fetch_chain_prefix(staging_chain: list[str]) → list[int]#: Resolve a staging chain through the shared, cached TQ read.

_resolve_admission_prefix(admission: Any) → list[int]#: Resolve a CaptureAdmission to the flat prefix the engine prompt starts with.

_enter_request_prefix( request: Any, prefix_token_ids: list[int], ) → None#

Attach the resolved prefix to the request through the capture adapter.

VLLMCaptureAdapter.enter_prefix writes the engine-native field (required_prefix_token_ids) into a payload; the same fields are applied to the pydantic request so the existing prefix-splice branch of preprocess_chat handles staged and inline prefixes alike.

static _delta_align_routed_experts( payload: dict[str, Any], *, prev_len: int, prompt_len: int, generated_len: int, ) → None#: Normalize optional vLLM routes to the staged token delta (plus boundary row).

_finish_request_capture(request: Any, content: dict) → dict#

Stage the finished call and ride its coords on the response.

Tokens and the call’s new media tensors go to TQ in one write (the media ride as opaque attachments beside the record), so staged coords vouch for both and any failure is capture_failed at call time. Token ids, logprobs, and routes are stripped after staging, so the worker->gate hop carries the completion and coords only.

_abort_request_capture(request: Any, *, reason: str) → None#: Drop the in-flight capture state for a request that errored.

_setup_vllm_openai_api_server(app: fastapi.FastAPI) → fastapi.FastAPI#

_setup_vllm_server() → tuple[threading.Thread, str, uvicorn.Server]#

async init_collective_async( rank_prefix: int, ip: str, port: int, world_size: int, train_world_size: int, ) → None#

async generate_async( data: nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.models.generation.interfaces.GenerationDatumSpec], greedy: bool = False, ) → AsyncGenerator[tuple[int, nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.models.generation.interfaces.GenerationOutputSpec]], None]#

Generate a batch of data using vLLM’s AsyncLLMEngine, yielding results as they are ready.

Parameters:

data – BatchedDataDict with input_ids and input_lengths
greedy – Whether to use greedy decoding instead of sampling

Yields:

Tuple of (original_index, BatchedDataDict conforming to GenerationOutputSpec for the single sequence)

async generate_text_async( data: nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.models.generation.interfaces.GenerationDatumSpec], greedy: bool = False, ) → AsyncGenerator[tuple[int, nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.models.generation.interfaces.GenerationOutputSpec]], None]#

Generate text responses asynchronously, yielding results as they are ready.