feat: forward rollout session IDs to Dynamo

nemo_rl/environments/nemo_gym.py

@@ -47,6 +47,7 @@
     RolloutDataFailure,
     http_status_is_infra,
 )
+from nemo_rl.models.generation.dynamo.token_wrapper import DYNAMO_SESSION_ID_HEADER
 from nemo_rl.models.generation.interfaces import should_use_async_rollouts
 from nemo_rl.models.policy import PolicyConfig, TokenizerConfig
 from nemo_rl.utils.routed_experts_codec import decode_routed_experts
@@ -1120,6 +1121,15 @@ def setup_nemo_gym_config(config, tokenizer) -> None:
     generation_config["stop_strings"] = None
     generation_config["stop_token_ids"] = None
 
+    if generation_config["backend"] == "dynamo":
+        model_config = (
+            config.env.setdefault("nemo_gym", {})
+            .setdefault("policy_model", {})
+            .setdefault("responses_api_models", {})
+            .setdefault("vllm_model", {})
+        )
+        model_config["session_id_header"] = DYNAMO_SESSION_ID_HEADER
+
     # For VLM runs, plumb the tokenizer config into the gym env config so the
     # NemoGym actor can reconstruct the processor inside itself (needed for
     # multi-turn multimodal postprocessing).

nemo_rl/experience/rollout_manager.py

@@ -45,8 +45,10 @@
 from nemo_rl.experience.metric_utils import calculate_single_metric, pct
 from nemo_rl.experience.rollouts import (
     EffortLevelsConfig,
+    _add_dynamo_session_id,
     _apply_effort_shaping,
     _attach_routed_experts_to_message_log_prefix,
+    _create_dynamo_session_id,
     _dummy_routed_experts_for_tokens,
     _effort_shaping_metrics,
     _find_routed_experts_template,
@@ -439,6 +441,7 @@ async def _run_single_rollout(
         current_extra_env_info = copy.deepcopy(input_sample["extra_env_info"])
         current_stop_strings = input_sample.get("stop_strings", None)
         task_name = input_sample["task_name"]
+        session_id = _create_dynamo_session_id(self._policy_generation)
 
         total_reward = 0.0
         turn_count = 0
@@ -476,6 +479,7 @@ async def _run_single_rollout(
                 ) = await self._generate_response(
                     current_message_log,
                     current_stop_strings,
+                    session_id=session_id,
                 )
             except Exception as e:
                 raise _classify_generation_failure(
@@ -602,6 +606,8 @@ async def _generate_response(
         self,
         message_log: list[dict],
         stop_strings: list[str] | None,
+        *,
+        session_id: str | None = None,
     ) -> tuple[dict, torch.Tensor, dict[str, Any]]:
         """Generate a single-turn response for one sample.
 
@@ -618,6 +624,11 @@ async def _generate_response(
                 "stop_strings": [stop_strings],
             }
         )
+        _add_dynamo_session_id(
+            generation_input_data,
+            self._policy_generation,
+            session_id,
+        )
 
         # Generate response
         # TODO: update generate_async to return a single item directly

nemo_rl/experience/rollouts.py

@@ -24,6 +24,7 @@
 from collections.abc import AsyncGenerator, Mapping, Sequence
 from dataclasses import dataclass
 from typing import Any, Optional
+from uuid import uuid4
 
 import ray
 import torch
@@ -77,6 +78,25 @@
 TokenizerType = PreTrainedTokenizerBase
 
 
+def _create_dynamo_session_id(
+    policy_generation: GenerationInterface,
+) -> str | None:
+    generation_config = getattr(policy_generation, "cfg", {})
+    if generation_config.get("backend") == "dynamo":
+        return str(uuid4())
+    return None
+
+
+def _add_dynamo_session_id(
+    generation_input_data: BatchedDataDict[GenerationDatumSpec],
+    policy_generation: GenerationInterface,
+    session_id: str | None,
+) -> None:
+    generation_config = getattr(policy_generation, "cfg", {})
+    if session_id is not None and generation_config.get("backend") == "dynamo":
+        generation_input_data["session_ids"] = [session_id]
+
+
 def attach_initial_nemo_gym_image_payloads(
     batch: BatchedDataDict[DatumSpec],
     processor: Any,
@@ -1127,6 +1147,7 @@ async def async_generate_response_for_sample_turn(
     max_seq_len: int,
     greedy: bool = False,
     *,
+    session_id: str | None = None,
     sample_multimodal_data: dict[str, Any] | None = None,
     deduplicate_multimodal_data: bool = False,
 ) -> tuple[list[dict], torch.Tensor, torch.Tensor, dict[str, float]]:
@@ -1139,6 +1160,8 @@ async def async_generate_response_for_sample_turn(
         tokenizer: Tokenizer to use
         max_seq_len: Maximum sequence length
         greedy: Whether to use greedy decoding
+        session_id: Stable Dynamo session ID reused across every turn of one
+            trajectory attempt. Ignored unless the generation backend is Dynamo.
         sample_multimodal_data: Native vLLM media fields for this sample.
         deduplicate_multimodal_data: Avoid sending both native and policy-ready
             media through the async generation boundary.
@@ -1165,6 +1188,11 @@ async def async_generate_response_for_sample_turn(
             "stop_strings": [sample_stop_strings],
         }
     )
+    _add_dynamo_session_id(
+        generation_input_data,
+        policy_generation,
+        session_id,
+    )
 
     # Create a dummy batch for generate_responses_async
     dummy_batch = BatchedDataDict[DatumSpec](
@@ -1237,6 +1265,7 @@ async def run_sample_multi_turn_rollout(
     current_extra_env_info = copy.deepcopy(initial_sample_state["extra_env_info"])
     current_stop_strings = initial_sample_state.get("stop_strings", None)
     task_name = initial_sample_state["task_name"]
+    session_id = _create_dynamo_session_id(policy_generation)
     sample_multimodal_data = {
         key: initial_sample_state[key]
         for key in NATIVE_MULTIMODAL_KEYS
@@ -1287,6 +1316,7 @@ async def run_sample_multi_turn_rollout(
                 tokenizer,
                 max_seq_len,
                 greedy=greedy,
+                session_id=session_id,
                 sample_multimodal_data=turn_multimodal_data,
                 deduplicate_multimodal_data=deduplicate_multimodal_data,
             )

nemo_rl/models/generation/dynamo/dynamo_generation.py

@@ -30,7 +30,10 @@
 from nemo_rl.models.generation.dynamo.managed_runtime import ManagedDynamoRuntime
 from nemo_rl.models.generation.dynamo.metrics import DynamoMetricsSampler
 from nemo_rl.models.generation.dynamo.refit import DynamoRefitChannel
-from nemo_rl.models.generation.dynamo.token_wrapper import DynamoTokenWrapperServer
+from nemo_rl.models.generation.dynamo.token_wrapper import (
+    DYNAMO_SESSION_ID_HEADER,
+    DynamoTokenWrapperServer,
+)
 from nemo_rl.models.generation.interfaces import (
     CollectiveSenderSpec,
     GenerationDatumSpec,
@@ -441,6 +444,7 @@ async def _post_completion_request(
         greedy: bool,
         stop_strings: Optional[list[str]],
         max_new_tokens: int,
+        session_id: Optional[str] = None,
     ) -> tuple[list[int], list[float], bool]:
         request_url = self._completion_url()
         payload = self._build_completion_request(
@@ -449,12 +453,16 @@ async def _post_completion_request(
             stop_strings=stop_strings,
             max_new_tokens=max_new_tokens,
         )
+        request_headers = (
+            {DYNAMO_SESSION_ID_HEADER: session_id} if session_id is not None else None
+        )
         response: dict[str, Any] = {}
         for attempt in range(1, _HTTP_MAX_ATTEMPTS + 1):
             response = await async_http_post_json(
                 request_url,
                 payload,
                 self._request_timeout_s(),
+                headers=request_headers,
             )
             if not _is_retryable_http_response(response):
                 break
@@ -569,6 +577,16 @@ async def generate_async(
             "outside this method."
         )
         sample_idx = 0
+        session_ids = data.get("session_ids")
+        session_id = None
+        if session_ids is not None:
+            if len(session_ids) != batch_size:
+                raise ValueError(
+                    "Dynamo session_ids must contain one value for each input sample."
+                )
+            session_id = session_ids[sample_idx]
+            if not isinstance(session_id, str) or not session_id.strip():
+                raise ValueError("Dynamo session IDs must be non-empty strings.")
         input_length = int(input_lengths_batch[sample_idx].item())
         batch_stop_strings = data.get("stop_strings", [[] for _ in range(batch_size)])
         per_sample_stop_strings = None
@@ -589,6 +607,7 @@ async def generate_async(
             greedy=greedy,
             stop_strings=final_stop_strings,
             max_new_tokens=allowed_new_tokens,
+            session_id=session_id,
         )
 
         yield (

nemo_rl/models/generation/dynamo/http_client.py

@@ -17,6 +17,7 @@
 import json
 import urllib.error
 import urllib.request
+from collections.abc import Mapping
 from typing import Any
 
 import aiohttp
@@ -63,13 +64,17 @@ def http_post_json(
 
 
 async def async_http_post_json(
-    url: str, payload: dict[str, Any], timeout_s: float
+    url: str,
+    payload: dict[str, Any],
+    timeout_s: float,
+    *,
+    headers: Mapping[str, str] | None = None,
 ) -> dict[str, Any]:
     """POST JSON without blocking the rollout actor event loop."""
     timeout = aiohttp.ClientTimeout(total=timeout_s)
     try:
         async with aiohttp.ClientSession(timeout=timeout) as session:
-            async with session.post(url, json=payload) as response:
+            async with session.post(url, json=payload, headers=headers) as response:
                 body = await response.read()
                 if response.status >= 400:
                     return {

nemo_rl/models/generation/dynamo/token_wrapper.py

@@ -35,6 +35,7 @@
     "generation_log_probs",
 )
 _TOOL_ARGUMENT_MAPPING_ERROR = "Can only get item pairs from a mapping."
+DYNAMO_SESSION_ID_HEADER = "X-Dynamo-Session-ID"
 
 
 def _coerce_token_id_list(value: Any, field_name: str) -> list[int]:
@@ -474,6 +475,7 @@ async def chat_completions(request: Request) -> JSONResponse:
             status_code, response_body = await self._forward_chat_completion(
                 prepared_body,
                 authorization=request.headers.get("authorization"),
+                session_id=request.headers.get(DYNAMO_SESSION_ID_HEADER),
             )
             if 200 <= status_code < 300:
                 try:
@@ -512,13 +514,16 @@ async def _forward_chat_completion(
         request_body: dict[str, Any],
         *,
         authorization: Optional[str],
+        session_id: Optional[str] = None,
     ) -> tuple[int, dict[str, Any]]:
         import aiohttp
 
         url = f"{self.dynamo_frontend_base_url.rstrip('/')}/chat/completions"
         headers = {"Content-Type": "application/json"}
         if authorization:
             headers["Authorization"] = authorization
+        if session_id:
+            headers[DYNAMO_SESSION_ID_HEADER] = session_id
 
         session = self._client_session
         if session is None:

nemo_rl/models/generation/interfaces.py

@@ -294,6 +294,7 @@ class GenerationDatumSpec(TypedDict):
     - input_ids: Tensor of token IDs representing the input sequences (right padded)
     - input_lengths: Tensor containing the actual length of each sequence (without padding)
     - stop_strings: Optional list of strings to stop generation (per sample)
+    - session_ids: Optional per-sample stable session IDs; honored only by the Dynamo backend
     - __extra__: Additional model-specific data fields
 
     Example of a batch with 4 entries with different sequence lengths:
@@ -319,6 +320,7 @@ class GenerationDatumSpec(TypedDict):
     input_ids: torch.Tensor
     input_lengths: torch.Tensor
     stop_strings: NotRequired[list[str]]
+    session_ids: NotRequired[list[str]]
     __extra__: Any
 
 

tests/unit/environments/test_nemo_gym.py

@@ -60,6 +60,7 @@
     _reattach_original_multimodal_payloads,
     attach_static_multimodal_payload,
 )
+from nemo_rl.models.generation.dynamo.token_wrapper import DYNAMO_SESSION_ID_HEADER
 from nemo_rl.models.generation.vllm import VllmGeneration
 
 # cluster and tokenizer are fixture imports
@@ -953,6 +954,45 @@ def test_validate_reward_components_match_scalar():
         )
 
 
+def test_setup_nemo_gym_config_merges_session_header_for_dynamo() -> None:
+    config = SimpleNamespace(
+        policy={"generation": {"backend": "dynamo", "vllm_cfg": {}}},
+        env={
+            "nemo_gym": {
+                "num_gpu_nodes": 2,
+                "policy_model": {
+                    "responses_api_models": {
+                        "vllm_model": {"model_name": "keep-me"},
+                        "other_model": {"model_name": "untouched"},
+                    }
+                },
+            }
+        },
+    )
+
+    setup_nemo_gym_config(config, tokenizer=object())
+
+    nemo_gym = config.env["nemo_gym"]
+    responses_api_models = nemo_gym["policy_model"]["responses_api_models"]
+    assert responses_api_models["vllm_model"] == {
+        "model_name": "keep-me",
+        "session_id_header": DYNAMO_SESSION_ID_HEADER,
+    }
+    assert responses_api_models["other_model"] == {"model_name": "untouched"}
+    assert nemo_gym["num_gpu_nodes"] == 2
+
+
+def test_setup_nemo_gym_config_does_not_set_session_header_for_vllm() -> None:
+    config = SimpleNamespace(
+        policy={"generation": {"backend": "vllm", "vllm_cfg": {}}},
+        env={},
+    )
+
+    setup_nemo_gym_config(config, tokenizer=object())
+
+    assert config.env == {}
+
+
 @pytest.mark.nemo_gym
 def test_nemo_gym_stub_module():
     from nemo_gym import config_types