nemo_rl.models.generation.trtllm.trtllm_worker_async

Ray actor wrapping tensorrt_llm._torch.async_llm.AsyncLLM.

The sole TRT-LLM generation worker (the synchronous engine path was removed; see :class:TrtllmGeneration, which asserts trtllm_cfg.async_engine=true). Every method that calls into AsyncLLM is exposed as async def with the _async suffix, so Ray’s actor runtime runs them on the actor’s own asyncio loop; process-lifecycle / helper methods (e.g. shutdown, configure_worker) stay sync.

Weight updates flow through NcclExtension inside TRT-LLM’s internal RayGPUWorker, invoked via llm.collective_rpc().

Module Contents

Classes

TrtllmAsyncGenerationWorkerImpl	Plain (non-actor) implementation of the async TRT-LLM generation worker.
TrtllmAsyncGenerationWorker	Ray actor wrapper around :class:TrtllmAsyncGenerationWorkerImpl.

API

class nemo_rl.models.generation.trtllm.trtllm_worker_async.TrtllmAsyncGenerationWorkerImpl(

config: nemo_rl.models.generation.trtllm.config.TrtllmConfig,

bundle_indices: Optional[list[int]] = None,

fraction_of_gpus: float = 1.0,

seed: Optional[int] = None,

)

Plain (non-actor) implementation of the async TRT-LLM generation worker.

Held separately from the @ray.remote-wrapped

Class:

TrtllmAsyncGenerationWorker so it can be exercised without Ray.

Initialization

static configure_worker(

num_gpus: int | float,

bundle_indices: Optional[tuple[int, list[int]]] = None,

num_gpus_per_node: Optional[int] = None,

) → tuple[dict[str, Any], dict[str, str], dict[str, Any], dict[str, Any]]

__repr__() → str

is_alive() → bool

async post_init_async() → None

Finish async engine setup on the Ray actor’s asyncio loop and (optionally) start HTTP server.

shutdown() → bool

start_http_server(port: int = 0) → str

Start an OpenAI-compatible HTTP server backed by self.llm.

stop_http_server() → None

async report_dp_openai_server_base_url() → Optional[str]

async init_collective_async(

rank_prefix: int,

ip: str,

port: int,

world_size: int,

train_world_size: int,

) → None

async prepare_refit_info_async(

state_dict_info: dict[str, Any],

) → None

async update_weights_from_collective_async(

*,

drain: bool = True,

recompute_kv: bool = False,

) → bool

Async version of update_weights_from_collective.

Parameters:

• drain – If False, run the refit at a scheduler step boundary without draining in-flight requests (in-flight weight update). Default True preserves the original drain-first behavior.

• recompute_kv – If True (and drain=False), preempt all in-flight requests after the refit so the scheduler re-prefills them under the new weights.

async update_weights_via_ipc_zmq_async() → bool

async report_device_id_async() → list[str]

classmethod _weights_tags() → list[str]

classmethod _all_sleep_tags() → list[str]

_resolve_wake_tags(tags: Optional[list[str]]) → list[str]

async sleep_async(**kwargs: Any) → bool

async wake_up_async(**kwargs: Any) → bool

async reset_prefix_cache_async(**kwargs: Any) → bool

async start_gpu_profiling_async() → None

async stop_gpu_profiling_async() → None

async generate_async(

data: nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.models.generation.interfaces.GenerationDatumSpec],

greedy: bool = False,

) → nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.models.generation.interfaces.GenerationOutputSpec]

_build_sampling_params(*, greedy: bool)

class nemo_rl.models.generation.trtllm.trtllm_worker_async.TrtllmAsyncGenerationWorker(

config: nemo_rl.models.generation.trtllm.config.TrtllmConfig,

bundle_indices: Optional[list[int]] = None,

fraction_of_gpus: float = 1.0,

seed: Optional[int] = None,

)

Bases: nemo_rl.models.generation.trtllm.trtllm_worker_async.TrtllmAsyncGenerationWorkerImpl

Ray actor wrapper around :class:TrtllmAsyncGenerationWorkerImpl.

Initialization