Fix RLM command timeout + consistent hashing

src/prime_rl/configs/inference.py

@@ -147,21 +147,64 @@ class DisaggregatedInferenceDeploymentConfig(BaseInferenceDeploymentConfig):
 
     Each inference replica is split into separate prefill and decode node groups.
     Requires NIXL for KV transfer and a vllm-router for request routing.
+
+    Multi-replica support: set ``num_prefill_replicas`` / ``num_decode_replicas``
+    to run multiple independent vLLM instances within the prefill / decode node
+    groups.  For example, ``num_prefill_nodes=4, num_prefill_replicas=2`` creates
+    two prefill vLLM instances each spanning 2 nodes (EP16 with 8 GPUs/node).
     """
 
     type: Literal["disaggregated"] = "disaggregated"
 
-    num_prefill_nodes: Annotated[int, Field(ge=1, description="Number of prefill nodes per replica.")] = 1
-    num_decode_nodes: Annotated[int, Field(ge=1, description="Number of decode nodes per replica.")] = 1
+    num_prefill_nodes: Annotated[int, Field(ge=1, description="Total number of prefill nodes.")] = 1
+    num_decode_nodes: Annotated[int, Field(ge=1, description="Total number of decode nodes.")] = 1
+
+    num_prefill_replicas: Annotated[
+        int,
+        Field(
+            ge=1,
+            description="Number of independent prefill vLLM instances. Must evenly divide num_prefill_nodes.",
+        ),
+    ] = 1
+    num_decode_replicas: Annotated[
+        int,
+        Field(
+            ge=1,
+            description="Number of independent decode vLLM instances. Must evenly divide num_decode_nodes.",
+        ),
+    ] = 1
 
     router_port: Annotated[int, Field(description="Port for the vllm-router on each replica.")] = 8000
     prefill_port: Annotated[int, Field(description="Port for prefill vLLM instances.")] = 8100
     decode_port: Annotated[int, Field(description="Port for decode vLLM instances.")] = 8200
 
+    prefill_env_overrides: Annotated[
+        dict[str, str],
+        Field(description="Extra environment variables exported only on prefill nodes."),
+    ] = {}
+    decode_env_overrides: Annotated[
+        dict[str, str],
+        Field(description="Extra environment variables exported only on decode nodes."),
+    ] = {}
+
     @property
     def num_nodes(self) -> int:
         return self.num_prefill_nodes + self.num_decode_nodes
 
+    @model_validator(mode="after")
+    def validate_replicas_divide_nodes(self):
+        if self.num_prefill_nodes % self.num_prefill_replicas != 0:
+            raise ValueError(
+                f"num_prefill_replicas ({self.num_prefill_replicas}) must evenly divide "
+                f"num_prefill_nodes ({self.num_prefill_nodes})"
+            )
+        if self.num_decode_nodes % self.num_decode_replicas != 0:
+            raise ValueError(
+                f"num_decode_replicas ({self.num_decode_replicas}) must evenly divide "
+                f"num_decode_nodes ({self.num_decode_nodes})"
+            )
+        return self
+
 
 InferenceDeploymentConfig: TypeAlias = Annotated[
     SingleNodeInferenceDeploymentConfig | MultiNodeInferenceDeploymentConfig | DisaggregatedInferenceDeploymentConfig,

src/prime_rl/entrypoints/inference.py

@@ -50,11 +50,15 @@ def write_slurm_script(config: InferenceConfig, config_path: Path, script_path:
         template_vars.update(
             num_prefill_nodes=config.deployment.num_prefill_nodes,
             num_decode_nodes=config.deployment.num_decode_nodes,
+            num_prefill_replicas=config.deployment.num_prefill_replicas,
+            num_decode_replicas=config.deployment.num_decode_replicas,
             prefill_port=config.deployment.prefill_port,
             decode_port=config.deployment.decode_port,
             router_port=config.deployment.router_port,
             data_parallel_rpc_port=config.data_parallel_rpc_port,
             use_deep_gemm=config.use_deep_gemm,
+            prefill_env_overrides=config.deployment.prefill_env_overrides,
+            decode_env_overrides=config.deployment.decode_env_overrides,
         )
     elif is_multi_node:
         template_vars.update(

src/prime_rl/entrypoints/rl.py

@@ -419,13 +419,17 @@ def write_slurm_script(config: RLConfig, config_dir: Path, script_path: Path) ->
             num_infer_replicas=config.deployment.num_infer_replicas,
             num_prefill_nodes=infer_deploy.num_prefill_nodes,
             num_decode_nodes=infer_deploy.num_decode_nodes,
+            num_prefill_replicas=infer_deploy.num_prefill_replicas,
+            num_decode_replicas=infer_deploy.num_decode_replicas,
             gpus_per_node=config.deployment.gpus_per_node,
             router_port=infer_deploy.router_port,
             prefill_port=infer_deploy.prefill_port,
             decode_port=infer_deploy.decode_port,
             inference_tp=config.inference.parallel.tp,
             inference_data_parallel_rpc_port=config.inference.data_parallel_rpc_port,
             use_deep_gemm=config.inference.use_deep_gemm,
+            prefill_env_overrides=infer_deploy.prefill_env_overrides,
+            decode_env_overrides=infer_deploy.decode_env_overrides,
             use_nccl_broadcast=config.weight_broadcast is not None and config.weight_broadcast.type == "nccl",
             wandb_shared=config.wandb is not None and config.wandb.shared,
         )

src/prime_rl/inference/patches.py

@@ -716,32 +716,83 @@ def _apply_with_saved_kernel(self, layer, x, topk_weights, topk_ids, shared_expe
 
 
 def monkey_patch_dp_engine_core_pause_resume_deadlock():
-    """Fix deadlock with pause/resume and collective_rpc in DP engine core.
+    """Fix DP pause/resume deadlocks around weight updates.
 
-    When a request arrives for an already-completed wave while the scheduler is
-    paused, the unpatched code sends a start_wave notification that triggers
-    collective_rpc on other DP engines. But the paused engine can't participate
-    in the collective, causing a deadlock.
+    Bug 1 (job 3756): while paused, START_DP_WAVE can wake idle ranks into the
+    DP loop. Those ranks then run dummy batches and hit DP collectives while
+    other ranks are still in NCCL weight transfer.
 
-    Fix: only send the start_wave notification when the scheduler is unpaused,
-    and explicitly set engines_running=True before notifying.
+    Bug 2 (jobs 3769/3771): resume ties the DP running state to local
+    unfinished requests, but the DP wave state is global. Ranks with no local
+    work still need to re-enter the loop so they can participate in the same
+    DP collectives as ranks that are resuming remote-KV or decode work.
 
-    Upstream: https://github.com/vllm-project/vllm/pull/37024
+    Fix:
+    - ignore START_DP_WAVE wakeups while paused
+    - on resume, wake every DP rank and force an immediate global unfinished
+      sync instead of waiting for the normal 32-step cadence
+
+    This keeps the upstream pause-side fix from
+    https://github.com/vllm-project/vllm/pull/37024 and extends it with the
+    resume-side wave-state fix.
     """
+    from vllm.config import ParallelConfig
     from vllm.v1.core.sched.interface import PauseState
-    from vllm.v1.engine import EngineCoreOutputs
-    from vllm.v1.engine.core import DPEngineCoreProc, EngineCore
+    from vllm.v1.engine import EngineCoreOutputs, EngineCoreRequestType
+    from vllm.v1.engine.core import DPEngineCoreProc, EngineCore, EngineCoreProc
     from vllm.v1.request import Request
 
     _base_add_request = EngineCore.add_request
+    _base_handle_client_request = EngineCoreProc._handle_client_request
+    _base_resume_scheduler = DPEngineCoreProc.resume_scheduler
 
     def _patched_add_request(self, request: Request, request_wave: int = 0):
         _base_add_request(self, request, request_wave)
         if self.has_coordinator and request_wave != self.current_wave:
             if request_wave > self.current_wave:
                 self.current_wave = request_wave
-            elif not self.engines_running and self.scheduler.pause_state == PauseState.UNPAUSED:
+            elif (
+                not self.engines_running
+                and self.scheduler.pause_state == PauseState.UNPAUSED
+            ):
                 self.engines_running = True
-                self.output_queue.put_nowait((-1, EngineCoreOutputs(start_wave=self.current_wave)))
+                self.output_queue.put_nowait(
+                    (-1, EngineCoreOutputs(start_wave=self.current_wave))
+                )
+
+    def _patched_handle_client_request(self, request_type, request):
+        if request_type == EngineCoreRequestType.START_DP_WAVE:
+            new_wave, exclude_eng_index = request
+            if (
+                exclude_eng_index != self.engine_index
+                and new_wave >= self.current_wave
+            ):
+                self.current_wave = new_wave
+                if (
+                    not self.engines_running
+                    and self.scheduler.pause_state == PauseState.UNPAUSED
+                ):
+                    self.engines_running = True
+        else:
+            _base_handle_client_request(self, request_type, request)
+
+    def _patched_resume_scheduler(self):
+        was_paused = self.scheduler.pause_state != PauseState.UNPAUSED
+        _base_resume_scheduler(self)
+        if was_paused:
+            self.engines_running = True
+            self._force_dp_running_state_sync = True
+
+    def _patched_has_global_unfinished_reqs(self, local_unfinished: bool) -> bool:
+        self.step_counter += 1
+        if getattr(self, "_force_dp_running_state_sync", False):
+            self._force_dp_running_state_sync = False
+            return ParallelConfig.has_unfinished_dp(self.dp_group, local_unfinished)
+        if self.step_counter % 32 != 0:
+            return True
+        return ParallelConfig.has_unfinished_dp(self.dp_group, local_unfinished)
 
     DPEngineCoreProc.add_request = _patched_add_request
+    DPEngineCoreProc._handle_client_request = _patched_handle_client_request
+    DPEngineCoreProc.resume_scheduler = _patched_resume_scheduler
+    DPEngineCoreProc._has_global_unfinished_reqs = _patched_has_global_unfinished_reqs

src/prime_rl/orchestrator/eval_utils.py

@@ -103,6 +103,7 @@ async def evaluate_env(
     logger = get_logger()
     logger.info(f"Evaluating {env_name} ({num_examples=}, {rollouts_per_example=})")
     eval_start_time = time.perf_counter()
+    total_inputs = len(env._get_eval_inputs(num_examples, rollouts_per_example))
     outputs = await evaluate(
         env=env,
         model_name=model_name,
@@ -113,9 +114,15 @@ async def evaluate_env(
         max_retries=max_retries,
     )
     eval_time = time.perf_counter() - eval_start_time
+    failed_rollouts = total_inputs - len(outputs)
 
     if not outputs:
-        logger.warning(f"All rollouts failed for {env_name}, skipping metrics")
+        logger.warning(f"All rollouts failed for {env_name} ({failed_rollouts} failed), skipping metrics")
+        monitor = get_monitor()
+        monitor.log(
+            {f"eval/{env_name}/failed_rollouts": failed_rollouts, "progress/ckpt_step": ckpt_step, "step": step},
+            step=step,
+        )
         return
 
     rows = []
@@ -166,6 +173,7 @@ async def evaluate_env(
         "completion_len/max": results_df.completion_len.max().item(),
         "completion_len/min": results_df.completion_len.min().item(),
         "is_truncated/mean": results_df.is_truncated.mean().item(),
+        "failed_rollouts": failed_rollouts,
         "time": eval_time,
     }
     if could_be_binary:
@@ -175,3 +183,4 @@ async def evaluate_env(
     eval_metrics.update({"progress/ckpt_step": ckpt_step, "step": step})
     monitor = get_monitor()
     monitor.log(eval_metrics, step=step)
+    monitor.log_eval_samples(outputs, env_name=env_name, step=step)

src/prime_rl/orchestrator/vf_utils.py

@@ -207,6 +207,10 @@ async def run_group_with_progress(example) -> list[vf.RolloutOutput] | None:
     finally:
         pbar.close()
 
+    failed_groups = sum(1 for g in group_outputs_list if g is None)
+    if failed_groups:
+        get_logger().warning(f"{failed_groups}/{len(group_outputs_list)} groups failed")
+
     return [output for group_outputs in group_outputs_list if group_outputs is not None for output in group_outputs]
 
 
@@ -229,7 +233,7 @@ async def evaluate(
 
     """
     inputs = env._get_eval_inputs(num_examples, rollouts_per_example)
-    outputs = await generate(
+    return await generate(
         env=env,
         clients=clients,
         get_client=get_client,
@@ -244,7 +248,6 @@ async def evaluate(
         max_retries=max_retries,
         state_columns=state_columns,
     )
-    return outputs
 
 
 # TODO: remove once usage is tracked by verifiers

src/prime_rl/templates/inference.sbatch.j2

@@ -27,7 +27,11 @@ export GPUS_PER_NODE={{ gpus_per_node }}
 {% if disaggregated -%}
 export NUM_PREFILL_NODES={{ num_prefill_nodes }}
 export NUM_DECODE_NODES={{ num_decode_nodes }}
+export NUM_PREFILL_REPLICAS={{ num_prefill_replicas }}
+export NUM_DECODE_REPLICAS={{ num_decode_replicas }}
 export NODES_PER_REPLICA=$((NUM_PREFILL_NODES + NUM_DECODE_NODES))
+export NODES_PER_PREFILL_REPLICA=$((NUM_PREFILL_NODES / NUM_PREFILL_REPLICAS))
+export NODES_PER_DECODE_REPLICA=$((NUM_DECODE_NODES / NUM_DECODE_REPLICAS))
 export PREFILL_PORT={{ prefill_port }}
 export DECODE_PORT={{ decode_port }}
 export ROUTER_PORT={{ router_port }}
@@ -142,19 +146,24 @@ srun bash -c '
     KV_CFG='"'"'{"kv_connector":"NixlConnector","kv_role":"kv_both","kv_connector_extra_config":{"num_threads":1}}'"'"'
     DECODE_COMPILE_CFG='"'"'{"cudagraph_mode":"FULL_DECODE_ONLY"}'"'"'
 
-    # Determine replica and role
+    # Determine outer replica (P/D pair) and role
     REPLICA_IDX=$((INFER_NODE_RANK / NODES_PER_REPLICA))
     RANK_IN_REPLICA=$((INFER_NODE_RANK % NODES_PER_REPLICA))
+    REPLICA_BASE=$((REPLICA_IDX * NODES_PER_REPLICA))
 
     if [ "$RANK_IN_REPLICA" -lt "$NUM_PREFILL_NODES" ]; then
         # ── Prefill node ──
         ROLE="prefill"
-        ROLE_RANK=$RANK_IN_REPLICA
-        EP=$((NUM_PREFILL_NODES * GPUS_PER_NODE))
+        SUB_REPLICA=$((RANK_IN_REPLICA / NODES_PER_PREFILL_REPLICA))
+        ROLE_RANK=$((RANK_IN_REPLICA % NODES_PER_PREFILL_REPLICA))
+        EP=$((NODES_PER_PREFILL_REPLICA * GPUS_PER_NODE))
         PORT=$PREFILL_PORT
         START_RANK=$((ROLE_RANK * GPUS_PER_NODE))
-        HEAD_HOST="${HOSTNAMES[$((REPLICA_IDX * NODES_PER_REPLICA))]}"
+        HEAD_HOST="${HOSTNAMES[$((REPLICA_BASE + SUB_REPLICA * NODES_PER_PREFILL_REPLICA))]}"
 
+{% for key, value in prefill_env_overrides.items() %}
+        export {{ key }}="{{ value }}"
+{% endfor %}
         ROLE_EXTRA="\"all2all_backend\": \"deepep_high_throughput\""
 
         if [ "$ROLE_RANK" -eq 0 ]; then
@@ -166,12 +175,17 @@ srun bash -c '
         # ── Decode node ──
         export UCX_NET_DEVICES=mlx5_0:1
         ROLE="decode"
-        ROLE_RANK=$((RANK_IN_REPLICA - NUM_PREFILL_NODES))
-        EP=$((NUM_DECODE_NODES * GPUS_PER_NODE))
+        RANK_IN_DECODE=$((RANK_IN_REPLICA - NUM_PREFILL_NODES))
+        SUB_REPLICA=$((RANK_IN_DECODE / NODES_PER_DECODE_REPLICA))
+        ROLE_RANK=$((RANK_IN_DECODE % NODES_PER_DECODE_REPLICA))
+        EP=$((NODES_PER_DECODE_REPLICA * GPUS_PER_NODE))
         PORT=$DECODE_PORT
         START_RANK=$((ROLE_RANK * GPUS_PER_NODE))
-        HEAD_HOST="${HOSTNAMES[$((REPLICA_IDX * NODES_PER_REPLICA + NUM_PREFILL_NODES))]}"
+        HEAD_HOST="${HOSTNAMES[$((REPLICA_BASE + NUM_PREFILL_NODES + SUB_REPLICA * NODES_PER_DECODE_REPLICA))]}"
 
+{% for key, value in decode_env_overrides.items() %}
+        export {{ key }}="{{ value }}"
+{% endfor %}
         ROLE_EXTRA="\"all2all_backend\": \"deepep_low_latency\", \"compilation_config\": $DECODE_COMPILE_CFG"
 
         if [ "$ROLE_RANK" -eq 0 ]; then
@@ -181,7 +195,7 @@ srun bash -c '
         fi
     fi
 
-    echo "NODE_RANK=$INFER_NODE_RANK ROLE=$ROLE REPLICA=$REPLICA_IDX ROLE_RANK=$ROLE_RANK EP=$EP PORT=$PORT LOCAL_IP=$LOCAL_IP" \
+    echo "NODE_RANK=$INFER_NODE_RANK ROLE=$ROLE REPLICA=$REPLICA_IDX SUB_REPLICA=$SUB_REPLICA ROLE_RANK=$ROLE_RANK EP=$EP PORT=$PORT LOCAL_IP=$LOCAL_IP" \
         | tee $OUTPUT_DIR/slurm/latest_infer_node_rank_${INFER_NODE_RANK}.log \
               $OUTPUT_DIR/slurm/job_${SLURM_JOB_ID}_infer_node_rank_${INFER_NODE_RANK}.log
 
@@ -193,7 +207,7 @@ srun bash -c '
         REPLICA_ROUTER_ARGS=$(echo "$ALL_ROUTER_ARGS" | cut -d"|" -f$((REPLICA_IDX + 1)))
 
         vllm-router \
-            --policy round_robin \
+            --policy consistent_hash \
             --vllm-pd-disaggregation \
             $REPLICA_ROUTER_ARGS \
             --host 0.0.0.0 \
@@ -226,7 +240,7 @@ srun bash -c '
         echo "Starting vllm-router on $LOCAL_IP:$ROUTER_PORT" | tee $ROUTER_LOG
 
         vllm-router \
-            --policy round_robin \
+            --policy consistent_hash \
             --worker-urls $ROUTER_ARGS \
             --host 0.0.0.0 \
             --port $ROUTER_PORT \