nemo_rl.models.generation.trtllm.trtllm_disagg_server

One replica’s disaggregation front-end, backed by TRT-LLM’s OpenAIDisaggServer.

Started like :mod:trtllm_http_server: a uvicorn app in a daemon thread, returning the URL NeMo-Gym talks to. NeMo RL hands it the two address pools and the router policies; everything below that – engine choice per leg, the KV handshake – is the disagg server’s.

Module Contents

Classes

_DropGymOnlyRequestFields	ASGI middleware stripping the vLLM-only fields from request bodies.
DisaggServerActorImpl	Hosts one replica’s disagg server.
DisaggServerActor	Ray actor wrapper around :class:DisaggServerActorImpl.

Functions

build_config	Assemble the DisaggServerConfig for one replica.
generation_token_ids	Generated token ids, from whichever declared field carries them.
attach_rollout_fields	Re-attach, in place, the fields NeMo-Gym reads off choices[].message.
_build_adaptor_class	Build OpenAIDisaggServerAdaptor lazily.
_build_frontend_tokenizer	Async payload -> b64 prompt ids using the adapters’ exact pipeline.
start_server	Start the disagg server in a daemon thread and return (thread, base_url, server).
wait_ready	Block until the disagg server answers /health.

Data

logger
__all__
_GYM_ONLY_REQUEST_FIELDS
_TOKEN_ID_PREFIX
_ADAPTED_PATHS

API

nemo_rl.models.generation.trtllm.trtllm_disagg_server.logger

‘getLogger(…)’

nemo_rl.models.generation.trtllm.trtllm_disagg_server.__all__

[‘DisaggServerActor’, ‘DisaggServerActorImpl’, ‘attach_rollout_fields’, ‘build_config’, ‘generation_…

nemo_rl.models.generation.trtllm.trtllm_disagg_server.build_config(

ctx_addrs: list[tuple[str, int]],

gen_addrs: list[tuple[str, int]],

*,

node_id: int,

ctx_router: str,

gen_router: str,

gen_tokids_ctxbytes: bool = False,

gen_strip_message_history: bool = False,

) → Any

Assemble the DisaggServerConfig for one replica.

Parameters:

• gen_addrs (ctx_addrs /) – (hostname, port) of every engine in this replica. An address is the only thing the disagg server is given about an engine.

• node_id – Must be distinct per replica. TRT-LLM’s default is uuid.getnode() % 256, documented as assuming a single disagg server per machine, and we run one per replica.

• ctx_router – Stateful, so a trajectory’s turns keep reaching the prefill engine holding its prefix.

• gen_router – Stateless, so placement stays local to this replica and no coordinator process is needed.

nemo_rl.models.generation.trtllm.trtllm_disagg_server._GYM_ONLY_REQUEST_FIELDS

(‘return_tokens_as_token_ids’, ‘return_token_ids’, ‘required_prefix_token_ids’, ‘parallel_tool_calls…

nemo_rl.models.generation.trtllm.trtllm_disagg_server._TOKEN_ID_PREFIX

‘token_id:’

nemo_rl.models.generation.trtllm.trtllm_disagg_server._ADAPTED_PATHS

‘frozenset(…)’

nemo_rl.models.generation.trtllm.trtllm_disagg_server.generation_token_ids(

choice: dict[str, Any],

) → Optional[list[int]]

Generated token ids, from whichever declared field carries them.

ChatCompletionResponseChoice.token_ids is declared since TRT-LLM f94eeb37 and is the live path. On builds without it they ride in logprobs.content[].token using vLLM’s token_id:N encoding, which is a declared string field and therefore survives the hop.

nemo_rl.models.generation.trtllm.trtllm_disagg_server.attach_rollout_fields(

payload: dict[str, Any],

) → dict[str, Any]

Re-attach, in place, the fields NeMo-Gym reads off choices[].message.

The counterpart of the aggregated path’s msg_dict assignment in

Mod:

trtllm_http_server: there the engine answers Gym directly and writes the fields onto the message, here the engine had to route them through declared response fields to survive the disagg server’s extra="forbid" re-validation, so they are moved back. The two must agree field for field – a message missing generation_token_ids is skipped outright by nemo_gym.py, silently dropping that turn’s training data – which is why this lives at module level with no tensorrt_llm import in its way.

class nemo_rl.models.generation.trtllm.trtllm_disagg_server._DropGymOnlyRequestFields(

app: Any,

tokenize_fn: Any = None,

)

ASGI middleware stripping the vLLM-only fields from request bodies.

Deliberately raw ASGI rather than a Starlette BaseHTTPMiddleware: that class passes the downstream app its own captured receive channel, so reassigning request._receive never reaches FastAPI’s validation and the request is rejected anyway. Replacing receive here is what the route actually reads.

Initialization

async __call__(

scope: Any,

receive: Any,

send: Any,

) → None

nemo_rl.models.generation.trtllm.trtllm_disagg_server._build_adaptor_class() → type

Build OpenAIDisaggServerAdaptor lazily.

Deferred so importing this module does not require tensorrt_llm.

nemo_rl.models.generation.trtllm.trtllm_disagg_server._build_frontend_tokenizer(

model_name: str,

server_template_kwargs: dict[str, Any],

) → Any

Async payload -> b64 prompt ids using the adapters’ exact pipeline.

Loads the tokenizer and HF config the same way the engine adapters do (AutoTokenizer/AutoConfig.from_pretrained(model_name, trust_remote_code=True)) and renders through the shared build_spliced_prompt_ids – the ids must be bit-identical to what the ctx adapter would compute, since a divergence silently trains on off-policy tokens (the adapter shadow-validates a sample under NRL_TRTLLM_TOKENIZE_SHADOW_RATE).

nemo_rl.models.generation.trtllm.trtllm_disagg_server.start_server(

config: Any,

host: str = '0.0.0.0',

port: int = 0,

req_timeout_secs: int = 1800,

tokenize_fn: Any = None,

) → tuple[threading.Thread, str, Any]

Start the disagg server in a daemon thread and return (thread, base_url, server).

class nemo_rl.models.generation.trtllm.trtllm_disagg_server.DisaggServerActorImpl(

replica_idx: int,

frontend_idx: int = 0,

num_frontends: int = 1,

port: int = 0,

serve_args: Optional[dict[str, Any]] = None,

)

Hosts one replica’s disagg server.

Its own process rather than an engine worker’s: this is the request hot path for the whole replica, and sharing a process with an engine would couple the replica’s routing latency to that one engine’s load – two uvicorn loops plus the engine’s generate thread pool contending for a single GIL. It holds no GPUs, so the isolation is cheap.

Held separately from the @ray.remote-wrapped :class:DisaggServerActor so it can be exercised without Ray.

Initialization

_start_serving(

ctx_addrs: list[tuple[str, int]],

gen_addrs: list[tuple[str, int]],

*,

ctx_router: str,

gen_router: str,

gen_tokids_ctxbytes: bool = False,

gen_strip_message_history: bool = False,

frontend_tokenize: bool = False,

model_name: Optional[str] = None,

default_chat_template_kwargs: Optional[dict[str, Any]] = None,

) → None

start() → str

Wait for the server to answer and return the URL NeMo-Gym will talk to.

Serving is the constructor’s job, not this method’s – that is what makes Ray actor restart replay it. A second entry point that could also start would be a silently less configured one (no frontend_tokenize, no model name, hardcoded routers), so the invariant is asserted instead of being papered over with a fallback.

base_url() → Optional[str]

shutdown() → bool

nemo_rl.models.generation.trtllm.trtllm_disagg_server.wait_ready(base_url: str, timeout_s: float = 300.0) → None

Block until the disagg server answers /health.

It reaches out to every engine in its pools on startup, so readiness lags the thread start by more than a socket bind.

class nemo_rl.models.generation.trtllm.trtllm_disagg_server.DisaggServerActor(

replica_idx: int,

frontend_idx: int = 0,

num_frontends: int = 1,

port: int = 0,

serve_args: Optional[dict[str, Any]] = None,

)

Bases: nemo_rl.models.generation.trtllm.trtllm_disagg_server.DisaggServerActorImpl

Ray actor wrapper around :class:DisaggServerActorImpl.

Initialization