From dd2a6cc3ff9034ae047983cc62febf75608f3586 Mon Sep 17 00:00:00 2001 From: yueming-yuan Date: Thu, 16 Jul 2026 22:01:09 -0700 Subject: [PATCH 1/3] rollout: support consistent_hashing/manual routing in the inference_rollout stack - rename Sample.session_id to Sample.routing_key to disambiguate from the session server's and p2p transfer engine's session ids - unify all routing-key sites behind policy_uses_routing_key, covering the manual policy (#1690) as well: both stacks' group assignment and eval, single_turn/multi_turn/legacy generate headers, prefill recompute - drop the MILES_EXPERIMENTAL_ROLLOUT_REFACTOR=1 restriction on --sglang-router-policy Keyless-request enforcement lives router-side (radixark/sgl-router-for-miles#8), which covers every client including proxy layers that bypass miles' http utils. --- miles/backends/sglang_utils/arguments.py | 7 ------- miles/rollout/generate_hub/multi_turn.py | 3 ++- miles/rollout/generate_hub/single_turn.py | 3 ++- .../generate_utils/generate_endpoint_utils.py | 15 +++++++++++++ .../generate_utils/prefill_logprobs.py | 8 +++---- miles/rollout/generate_utils/sample_utils.py | 2 +- .../inference_rollout_common.py | 7 +++++++ .../inference_rollout_eval.py | 4 ++++ miles/rollout/sglang_rollout.py | 21 +++++++++++-------- miles/utils/types.py | 7 +++---- 10 files changed, 49 insertions(+), 28 deletions(-) diff --git a/miles/backends/sglang_utils/arguments.py b/miles/backends/sglang_utils/arguments.py index 71e3f48fb4a..bb9c1871e04 100644 --- a/miles/backends/sglang_utils/arguments.py +++ b/miles/backends/sglang_utils/arguments.py @@ -147,12 +147,5 @@ def validate_args(args): if args.sglang_dp_size > 1: assert args.sglang_enable_dp_attention - if args.sglang_router_policy: - from miles.utils.environ import enable_experimental_rollout_refactor - - assert ( - not enable_experimental_rollout_refactor() - ), "--sglang-router-policy is not supported with MILES_EXPERIMENTAL_ROLLOUT_REFACTOR=1" - if getattr(args, "sglang_router_ip", None): args.sglang_router_ip = _wrap_ipv6(args.sglang_router_ip) diff --git a/miles/rollout/generate_hub/multi_turn.py b/miles/rollout/generate_hub/multi_turn.py index 97814ecb3d1..99bec2b1a9d 100644 --- a/miles/rollout/generate_hub/multi_turn.py +++ b/miles/rollout/generate_hub/multi_turn.py @@ -9,6 +9,7 @@ from miles.rollout.generate_utils.generate_endpoint_utils import ( compute_prompt_ids_from_sample, compute_request_payload, + compute_routing_headers, update_sample_from_response, ) from miles.rollout.generate_utils.tool_call_utils import ( @@ -56,7 +57,7 @@ async def generate(input: GenerateFnInput) -> GenerateFnOutput: if args.generate_multi_samples: sample = deepcopy(input.sample) - output = await post(url, payload) + output = await post(url, payload, headers=compute_routing_headers(args, sample)) await update_sample_from_response(args, sample, payload=payload, output=output, update_loss_mask=True) if args.generate_multi_samples: diff --git a/miles/rollout/generate_hub/single_turn.py b/miles/rollout/generate_hub/single_turn.py index 5c0a15b5b4b..bc744572de3 100644 --- a/miles/rollout/generate_hub/single_turn.py +++ b/miles/rollout/generate_hub/single_turn.py @@ -6,6 +6,7 @@ from miles.rollout.generate_utils.generate_endpoint_utils import ( compute_prompt_ids_from_sample, compute_request_payload, + compute_routing_headers, update_sample_from_response, ) from miles.utils.http_utils import post @@ -40,7 +41,7 @@ async def generate(input: GenerateFnInput) -> GenerateFnOutput: sample.status = halt_status return GenerateFnOutput(samples=sample) - output = await post(url, payload) + output = await post(url, payload, headers=compute_routing_headers(args, sample)) await update_sample_from_response(args, sample, payload=payload, output=output) return GenerateFnOutput(samples=sample) diff --git a/miles/rollout/generate_utils/generate_endpoint_utils.py b/miles/rollout/generate_utils/generate_endpoint_utils.py index d50098e686b..a3d2e775ec6 100644 --- a/miles/rollout/generate_utils/generate_endpoint_utils.py +++ b/miles/rollout/generate_utils/generate_endpoint_utils.py @@ -36,6 +36,21 @@ def compute_prompt_ids_from_sample(state, sample, tools=None): return state.tokenizer.encode(prompt, add_special_tokens=False) +def policy_uses_routing_key(args) -> bool: + return args.sglang_router_policy in ("consistent_hashing", "manual") + + +def compute_routing_headers(args, sample: Sample) -> dict[str, str] | None: + if policy_uses_routing_key(args): + assert sample.routing_key, ( + f"router policy {args.sglang_router_policy} routes by X-SMG-Routing-Key, " + f"but sample (index={sample.index}) has no routing_key set" + ) + if sample.routing_key: + return {"X-SMG-Routing-Key": sample.routing_key} + return None + + def compute_request_payload( args, input_ids: list[int], diff --git a/miles/rollout/generate_utils/prefill_logprobs.py b/miles/rollout/generate_utils/prefill_logprobs.py index d11752fd1f3..8512be8732b 100644 --- a/miles/rollout/generate_utils/prefill_logprobs.py +++ b/miles/rollout/generate_utils/prefill_logprobs.py @@ -4,6 +4,7 @@ from collections.abc import Mapping from typing import Any +from miles.rollout.generate_utils.generate_endpoint_utils import compute_routing_headers, policy_uses_routing_key from miles.utils.http_utils import post from miles.utils.lora import LORA_ADAPTER_NAME, is_lora_enabled from miles.utils.processing_utils import encode_image_for_rollout_engine @@ -48,7 +49,7 @@ def _build_prefill_scoring_payload( def _can_batch_prefill_score(args: Any, samples: list[Sample]) -> bool: - if getattr(args, "sglang_router_policy", None) == "consistent_hashing": + if policy_uses_routing_key(args): return False return not any(sample.multimodal_inputs and sample.multimodal_inputs.get("images") for sample in samples) @@ -161,10 +162,7 @@ async def recompute_samples_rollout_logprobs_via_prefill( return for sample in samples_to_score: - headers = None - uses_consistent_hashing = getattr(args, "sglang_router_policy", None) == "consistent_hashing" - if uses_consistent_hashing and sample.session_id: - headers = {"X-SMG-Routing-Key": sample.session_id} + headers = compute_routing_headers(args, sample) await post(flush_url, {}, headers=headers) await recompute_rollout_logprobs_via_prefill( diff --git a/miles/rollout/generate_utils/sample_utils.py b/miles/rollout/generate_utils/sample_utils.py index 1594dac28e9..effbce562b1 100644 --- a/miles/rollout/generate_utils/sample_utils.py +++ b/miles/rollout/generate_utils/sample_utils.py @@ -142,7 +142,7 @@ def _merge_metadata(): metadata=_merge_metadata(), generate_function_path=_merge_equal_value("generate_function_path"), train_metadata=_merge_equal_value("train_metadata"), - session_id=_merge_equal_value("session_id"), + routing_key=_merge_equal_value("routing_key"), non_generation_time=_merge_equal_value("non_generation_time"), spec_info=_merge_spec_info(a.spec_info, b.spec_info), prefix_cache_info=_merge_prefix_cache_info(a.prefix_cache_info, b.prefix_cache_info), diff --git a/miles/rollout/inference_rollout/inference_rollout_common.py b/miles/rollout/inference_rollout/inference_rollout_common.py index 9f1cc603b01..1b835104409 100644 --- a/miles/rollout/inference_rollout/inference_rollout_common.py +++ b/miles/rollout/inference_rollout/inference_rollout_common.py @@ -1,5 +1,6 @@ import asyncio import logging +import uuid from argparse import Namespace from copy import deepcopy from typing import Any @@ -15,6 +16,7 @@ RolloutFnTrainOutput, ) from miles.rollout.generate_hub.single_turn import generate +from miles.rollout.generate_utils.generate_endpoint_utils import policy_uses_routing_key from miles.rollout.inference_rollout.compatibility import load_generate_function from miles.rollout.rm_hub import async_rm, batched_async_rm from miles.utils.processing_utils import load_processor, load_tokenizer @@ -125,6 +127,11 @@ async def generate_and_rm_group( if state.aborted: return group + if policy_uses_routing_key(args): + for sample in group: + if sample.routing_key is None: + sample.routing_key = str(uuid.uuid4()) + log_prefix = f"[group indices={[getattr(s, 'index', '?') for s in group]}]" logger.debug(f"{log_prefix} Starting group with {len(group)} samples") tasks = [] diff --git a/miles/rollout/inference_rollout/inference_rollout_eval.py b/miles/rollout/inference_rollout/inference_rollout_eval.py index 2747776791e..1826f9f38de 100644 --- a/miles/rollout/inference_rollout/inference_rollout_eval.py +++ b/miles/rollout/inference_rollout/inference_rollout_eval.py @@ -1,10 +1,12 @@ import asyncio import copy import logging +import uuid from typing import Any from tqdm import tqdm +from miles.rollout.generate_utils.generate_endpoint_utils import policy_uses_routing_key from miles.rollout.inference_rollout.inference_rollout_common import ( GenerateState, compute_sampling_params, @@ -66,6 +68,8 @@ async def eval_rollout_single_dataset( sample.index = sample_index sample_index += 1 sample.metadata = dataset_cfg.inject_metadata(getattr(sample, "metadata", None)) + if policy_uses_routing_key(args): + sample.routing_key = str(uuid.uuid4()) sampling_params = base_sampling_params if getattr(args, "sglang_enable_deterministic_inference", False): sampling_params = base_sampling_params.copy() diff --git a/miles/rollout/sglang_rollout.py b/miles/rollout/sglang_rollout.py index e9bf9aa6a73..2d79dbf236c 100644 --- a/miles/rollout/sglang_rollout.py +++ b/miles/rollout/sglang_rollout.py @@ -31,7 +31,11 @@ ) from miles.utils.types import Sample -from .generate_utils.generate_endpoint_utils import get_indexer_topk_from_response +from .generate_utils.generate_endpoint_utils import ( + compute_routing_headers, + get_indexer_topk_from_response, + policy_uses_routing_key, +) from .generate_utils.prefill_logprobs import recompute_samples_rollout_logprobs_via_prefill from .rm_hub import async_rm, batched_async_rm @@ -192,10 +196,7 @@ async def generate(args: Namespace, sample: Sample, sampling_params: dict[str, A if not sample.tokens: # Initialize sample.tokens for the first turn sample.tokens = prompt_ids - # Use session_id for consistent hashing routing if router uses consistent_hashing policy - headers = None - if args.sglang_router_policy == "consistent_hashing" and sample.session_id: - headers = {"X-SMG-Routing-Key": sample.session_id} + headers = compute_routing_headers(args, sample) output = await post(url, payload, headers=headers) if getattr(args, "use_opd", False) and opd_top_k > 0 and opd_top_k_strategy != "only-teacher": @@ -313,11 +314,11 @@ async def generate_and_rm_group( if state.aborted: return group - # Generate a unique session_id for each sample in the group (consistent hashing only) - if args.sglang_router_policy == "consistent_hashing": + # Generate a unique routing_key for each sample in the group (routing-key policies only) + if policy_uses_routing_key(args): for sample in group: - if sample.session_id is None: - sample.session_id = str(uuid.uuid4()) + if sample.routing_key is None: + sample.routing_key = str(uuid.uuid4()) tasks = [] for idx, sample in enumerate(group): @@ -565,6 +566,8 @@ async def eval_rollout_single_dataset( sample_index += 1 sample.metadata = dataset_cfg.inject_metadata(getattr(sample, "metadata", None)) sample.generate_function_path = getattr(dataset_cfg, "custom_generate_function_path", None) + if policy_uses_routing_key(args): + sample.routing_key = str(uuid.uuid4()) sampling_params = base_sampling_params if getattr(args, "sglang_enable_deterministic_inference", False): sampling_params = base_sampling_params.copy() diff --git a/miles/utils/types.py b/miles/utils/types.py index cd7637d4c44..6b501456501 100644 --- a/miles/utils/types.py +++ b/miles/utils/types.py @@ -52,9 +52,8 @@ class Status(Enum): # metadata used during training, e.g., what loss to use for this sample. train_metadata: dict | None = None - # Session ID for consistent hashing routing (used when router policy is consistent_hashing) - # TODO: Its definition needs to merge with the session server's session id in the new rollout function. - session_id: str | None = None + # Per-sample routing key for the router's consistent_hashing policy (sent as X-SMG-Routing-Key) + routing_key: str | None = None non_generation_time: float = 0.0 # time spent in non-generation steps @@ -217,7 +216,7 @@ def reset_for_retry(self) -> None: """Reset generated outputs so the original prompt can be re-sampled. Keeps identity / prompt fields (group_index, index, prompt, label, - multimodal_inputs, metadata, generate_function_path, session_id) and + multimodal_inputs, metadata, generate_function_path, routing_key) and restores everything else to dataclass defaults. """ self.tokens = [] From a1719748f3eed79113b122216dc20cb213144f7e Mon Sep 17 00:00:00 2001 From: Yueming Yuan Date: Mon, 20 Jul 2026 11:51:48 -0700 Subject: [PATCH 2/3] [router] set manual policy (sticky + min_load) as default agentic routing policy (#1690) --- miles/backends/sglang_utils/arguments.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/miles/backends/sglang_utils/arguments.py b/miles/backends/sglang_utils/arguments.py index bb9c1871e04..02939d4f63a 100644 --- a/miles/backends/sglang_utils/arguments.py +++ b/miles/backends/sglang_utils/arguments.py @@ -147,5 +147,10 @@ def validate_args(args): if args.sglang_dp_size > 1: assert args.sglang_enable_dp_attention + if args.sglang_router_policy is None and args.use_session_server: + args.sglang_router_policy = "manual" + if args.router_assignment_mode == "random": + args.router_assignment_mode = "min_load" + if getattr(args, "sglang_router_ip", None): args.sglang_router_ip = _wrap_ipv6(args.sglang_router_ip) From 341dcf6443628be1d251de759efa2d86515aad9f Mon Sep 17 00:00:00 2001 From: Zhichenzzz Date: Mon, 20 Jul 2026 14:17:08 -0700 Subject: [PATCH 3/3] rollout: raise ValueError instead of assert for missing routing_key assert is stripped entirely under python -O / PYTHONOPTIMIZE=1, which would silently skip this check and reintroduce the keyless-request degradation this PR fixes. Enforcement now lives router-side too (sgl-router-for-miles#8), so this is defense-in-depth, but it should still fail loudly rather than be optimizable away. --- miles/rollout/generate_utils/generate_endpoint_utils.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/miles/rollout/generate_utils/generate_endpoint_utils.py b/miles/rollout/generate_utils/generate_endpoint_utils.py index a3d2e775ec6..1de0e66d079 100644 --- a/miles/rollout/generate_utils/generate_endpoint_utils.py +++ b/miles/rollout/generate_utils/generate_endpoint_utils.py @@ -41,8 +41,8 @@ def policy_uses_routing_key(args) -> bool: def compute_routing_headers(args, sample: Sample) -> dict[str, str] | None: - if policy_uses_routing_key(args): - assert sample.routing_key, ( + if policy_uses_routing_key(args) and not sample.routing_key: + raise ValueError( f"router policy {args.sglang_router_policy} routes by X-SMG-Routing-Key, " f"but sample (index={sample.index}) has no routing_key set" )