diff --git a/docs/algo/rollout_server_routing.md b/docs/algo/rollout_server_routing.md new file mode 100644 index 00000000..b7af25d5 --- /dev/null +++ b/docs/algo/rollout_server_routing.md @@ -0,0 +1,206 @@ +(rollout_server_routing)= +# Uid-Affinity Server Routing for Diffusion Rollouts + +Last updated: 06/14/2026 + +Diffusion async rollouts in VeRL-Omni already run across multiple `OmniLLMServer` +replicas via verl's server-manager stack. This feature adds **uid-affinity server +routing**: configurable replica selection keyed on per-prompt `uid`, so all +`rollout.n` copies of the same prompt can co-locate on one replica. It is aimed at +online diffusion RL (FlowGRPO, MixGRPO, DiffusionNFT) when +`actor_rollout_ref.rollout.nnodes > 0` or multiple rollout servers are active. + +## Motivation + +Upstream `verl` uses `GlobalRequestLoadBalancer` with a single behavior: +**least-inflight routing**, sticky only on the per-request `request_id`. That is a +good default for LLM serving fairness and multi-turn prefix caching, but it is a +poor fit for diffusion RL batching patterns. + +In a typical FlowGRPO step with `train_batch_size=32` and `rollout.n=16`, the +trainer issues **512 independent HTTP rollout requests**. Each request has a +unique `request_id`, but only **32 prompt groups** exist because 16 requests +share the same per-prompt `uid`. + +With least-inflight routing keyed on `request_id`, those 16 copies scatter across +replicas. Each GPU sees a fragmented subset of the workload instead of a dense +group of homogeneous prompts. + +`prompt_uid_affinity` routes by the prompt-group key (`uid` by default) so all +`rollout.n` copies for one prompt co-locate on one replica. That improves local +batch formation when combined with vLLM-Omni request-level batching or wide +`max_num_seqs` scheduling. + +| Goal | `least_inflight` | `prompt_uid_affinity` | +| --- | --- | --- | +| Optimize for | Fair spread / low tail latency | Co-locate related rollout copies | +| Sticky key | `request_id` (unique per HTTP call) | `routing_key` (default: batch `uid`) | +| Typical effect with `rollout.n=16` | ~1–4 requests per replica at a time | Up to 16 requests per replica per prompt | + +## What uid-affinity server routing means + +Uid-affinity server routing in VeRL-Omni is **verl-side replica selection policy** +in `OmniLLMServerManager` / `OmniLLMServerClient`. It replaces verl's +`GlobalRequestLoadBalancer` on the diffusion trainer path with +`OmniRequestLoadBalancer`, which can stick requests by `uid` instead of only by +`request_id`. It is distinct from: + +- **`scheduling_policy`** (`fcfs`, etc.) — internal vLLM-Omni diffusion scheduler + policy on a single replica. +- **Request-level batching** — fusing concurrent requests into one GPU forward + inside a replica (vLLM-Omni feature; orthogonal to routing). + +The routing layer sits between the diffusion agent loop and the rollout HTTP +servers: + +1. The trainer assigns one `uid` per prompt, then repeats the batch + `rollout.n` times (interleaved). +2. Each agent-loop worker calls `server_manager.generate(..., routing_key=uid)`. +3. `OmniRequestLoadBalancer` picks a replica according to `policy`. +4. The chosen `OmniLLMServer` runs generation as usual. + +Training semantics stay on-policy: routing only affects **where** a request +executes, not **what** policy generates it. + +## Quickstart + +Diffusion trainers default to `prompt_uid_affinity` via +`verl_omni/trainer/config/diffusion/rollout/diffusion_rollout.yaml`. For +multi-replica FlowGRPO, set standalone rollout nodes and keep the default policy: + +```bash +bash examples/flowgrpo_trainer/run_qwen_image_ocr_lora.sh \ + actor_rollout_ref.rollout.nnodes=1 \ + actor_rollout_ref.rollout.n_gpus_per_node=4 \ + actor_rollout_ref.rollout.tensor_model_parallel_size=1 \ + actor_rollout_ref.rollout.n=16 +``` + +To restore verl-style fair spreading across replicas: + +```bash +actor_rollout_ref.rollout.server_routing.policy=least_inflight +``` + +Explicit affinity override (equivalent to the diffusion default): + +```bash +actor_rollout_ref.rollout.server_routing.policy=prompt_uid_affinity +``` + +## Config reference + +Settings live under `actor_rollout_ref.rollout.server_routing`: + +| Config | Meaning | +| --- | --- | +| `server_routing.policy` | Replica routing policy (see table below). Diffusion defaults to `prompt_uid_affinity`; shared `server_routing.yaml` keeps `least_inflight` for omni LLM rollouts. | +| `server_routing.routing_key_field` | Per-sample kwargs field used as the routing key. Defaults to `uid` in code; usually omitted from yaml. | + +Supported policies: + +| Policy | Behavior | +| --- | --- | +| `prompt_uid_affinity` | Sticky route by `routing_key`. New keys pick the least-loaded replica; subsequent requests with the same key follow that replica. Recommended for diffusion RL with `rollout.n > 1`. | +| `least_inflight` | Sticky route by `request_id` with least-loaded assignment for new keys. Spreads concurrent requests for fairness. | +| `prompt_hash_sharding` | Deterministic shard `hash(routing_key) % num_replicas`. | +| `round_robin` | Rotate replicas regardless of load. | + +Base schema: `verl_omni/trainer/config/rollout/server_routing.yaml`. +Diffusion override: `verl_omni/trainer/config/diffusion/rollout/diffusion_rollout.yaml`. + +## What is `uid`? + +`uid` is a **prompt-group identifier** assigned by the diffusion trainer each +step (one UUID per prompt before `rollout.n` expansion). After interleaved +repeat, all `rollout.n` copies of the same prompt share the same `uid`. The same +field is used for GRPO advantage grouping. + +You rarely need to set `routing_key_field`; keep the default `uid` unless you +introduce a custom grouping field in the agent-loop kwargs. + +## How it plugs in + +The trainer creates per-prompt `uid` values, then expands for rollout: + +```python +batch.non_tensor_batch["uid"] = np.array( + [str(uuid.uuid4()) for _ in range(len(batch.batch))], dtype=object +) +gen_batch_output = gen_batch.repeat( + repeat_times=self.config.actor_rollout_ref.rollout.n, interleave=True +) +``` + +The single-turn diffusion agent loop extracts the routing key and forwards it to +the server manager: + +```python +routing_key_field = OmegaConf.select( + self.rollout_config, "server_routing.routing_key_field", default="uid" +) +value = kwargs.get(routing_key_field) +routing_key = str(value) if value is not None else None + +output = await self.server_manager.generate( + request_id=uuid4().hex, + prompt_ids=prompt_ids, + sampling_params=sampling_params, + routing_key=routing_key, +) +``` + +`OmniLLMServerManager` installs `OmniRequestLoadBalancer` and reads the +policy from config: + +```python +policy = OmegaConf.select( + self.rollout_config, "server_routing.policy", default="least_inflight" +) +self.global_load_balancer = OmniRequestLoadBalancer.remote( + servers=dict(zip(self.server_addresses, self.server_handles, strict=True)), + policy=policy, +) +``` + +For `prompt_uid_affinity`, the load balancer sticks on `routing_key` (falling +back to `request_id` when the key is missing). For `least_inflight`, it sticks on +`request_id` only. + +Use `OmniLLMServerManager` (not upstream `LLMServerManager`) in diffusion agent-loop +tests and trainers so `routing_key` is accepted end-to-end. + +## Interaction with request-level batching + +Replica routing and vLLM-Omni request-level batching are **orthogonal**: + +- Routing decides **which replica** receives a request. +- Batching decides **how requests on one replica** are fused into GPU forwards. + +They compose well: `prompt_uid_affinity` sends all `rollout.n` copies of a prompt +to one replica, giving local admission windows more homogeneous concurrent traffic. +In multi-replica FlowGRPO experiments, using both reduced step time compared with +`least_inflight` alone when replicas could form wider fused batches. + +Routing does not require request-level batching to be enabled, and request-level +batching does not require affinity routing. Enable both when training throughput +is limited by fragmented per-replica load. + +## Tests + +| Test | Purpose | +| --- | --- | +| `tests/workers/rollout/test_rollout_server_routing.py` | CPU tests for policies, Hydra defaults, and `uid` clustering. | +| `tests/workers/rollout/rollout_vllm/test_rollout_server_routing_perf.py` | Multi-replica GPU benchmark comparing policies under FlowGRPO-like bursts. | + +Run CPU coverage: + +```bash +pytest tests/workers/rollout/test_rollout_server_routing.py -q +``` + +## References + +- [Flow-GRPO: Training Flow Matching Models via Online RL](https://arxiv.org/abs/2505.05470) +- [FlowGRPO quickstart](../start/flowgrpo_quickstart.md) +- [Async reward](async_reward.md) — another advanced feature for hiding reward latency behind rollout diff --git a/docs/index.md b/docs/index.md index 26c8d6da..efebd169 100644 --- a/docs/index.md +++ b/docs/index.md @@ -38,6 +38,7 @@ start/metrics.md :caption: Advanced Features algo/async_reward.md +algo/rollout_server_routing.md algo/rollout_correction.md start/http_scorer.md ``` diff --git a/tests/agent_loop/test_diffusion_agent_loop.py b/tests/agent_loop/test_diffusion_agent_loop.py index 799a0919..4b5e402a 100644 --- a/tests/agent_loop/test_diffusion_agent_loop.py +++ b/tests/agent_loop/test_diffusion_agent_loop.py @@ -21,9 +21,9 @@ from omegaconf import DictConfig from verl.experimental.agent_loop.agent_loop import AgentLoopManager from verl.protocol import DataProto -from verl.workers.rollout.llm_server import LLMServerManager from verl_omni.agent_loop import DiffusionAgentLoopWorker +from verl_omni.workers.rollout.omni_llm_server import OmniLLMServerManager from ..utils.gpu_test_topology import resolve_diffusion_agent_loop_gpu_topology @@ -114,7 +114,7 @@ def test_single_turn(init_config): ) try: AgentLoopManager.agent_loop_workers_class = ray.remote(DiffusionAgentLoopWorker) - llm_server_manager = LLMServerManager.create(config=init_config) + llm_server_manager = OmniLLMServerManager.create(config=init_config) agent_loop_manager = AgentLoopManager.create( config=init_config, llm_client=llm_server_manager.get_client(), diff --git a/tests/agent_loop/test_diffusion_rollout_seed_gpu.py b/tests/agent_loop/test_diffusion_rollout_seed_gpu.py index 05c86933..d1d40591 100644 --- a/tests/agent_loop/test_diffusion_rollout_seed_gpu.py +++ b/tests/agent_loop/test_diffusion_rollout_seed_gpu.py @@ -25,9 +25,9 @@ from omegaconf import DictConfig from verl.experimental.agent_loop.agent_loop import AgentLoopManager from verl.protocol import DataProto -from verl.workers.rollout.llm_server import LLMServerManager from verl_omni.agent_loop import DiffusionAgentLoopWorker +from verl_omni.workers.rollout.omni_llm_server import OmniLLMServerManager from ..utils.gpu_test_topology import resolve_diffusion_agent_loop_gpu_topology @@ -163,7 +163,7 @@ def test_rollout_without_seed_produces_different_initial_latents(multi_worker_se ) try: AgentLoopManager.agent_loop_workers_class = ray.remote(DiffusionAgentLoopWorker) - llm_server_manager = LLMServerManager.create(config=multi_worker_seed_rollout_config) + llm_server_manager = OmniLLMServerManager.create(config=multi_worker_seed_rollout_config) agent_loop_manager = AgentLoopManager.create( config=multi_worker_seed_rollout_config, llm_client=llm_server_manager.get_client(), @@ -211,7 +211,7 @@ def test_rollout_seeds_unique_across_agent_loop_workers(multi_worker_seed_rollou ) try: AgentLoopManager.agent_loop_workers_class = ray.remote(DiffusionAgentLoopWorker) - llm_server_manager = LLMServerManager.create(config=multi_worker_seed_rollout_config) + llm_server_manager = OmniLLMServerManager.create(config=multi_worker_seed_rollout_config) agent_loop_manager = AgentLoopManager.create( config=multi_worker_seed_rollout_config, llm_client=llm_server_manager.get_client(), diff --git a/tests/workers/rollout/rollout_vllm/test_rollout_server_routing_perf.py b/tests/workers/rollout/rollout_vllm/test_rollout_server_routing_perf.py new file mode 100644 index 00000000..3f826b9b --- /dev/null +++ b/tests/workers/rollout/rollout_vllm/test_rollout_server_routing_perf.py @@ -0,0 +1,327 @@ +# Copyright 2026 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Multi-server FlowGRPO routing benchmark. + +Compares rollout replica routing policies under the same OCR burst pattern used +in training (train_batch_size x rollout.n concurrent generates). + +Usage: + pytest tests/workers/rollout/rollout_vllm/test_rollout_server_routing_perf.py -v -s + +Env overrides: + FLOWGRPO_ROUTING_NUM_REPLICAS (default: 4) + FLOWGRPO_ROUTING_NUM_PROMPTS (default: 32) + FLOWGRPO_ROUTING_N (default: 16) + FLOWGRPO_ROUTING_POLICY (default: prompt_uid_affinity) + FLOWGRPO_ROUTING_COMPARE_POLICIES (default: 1, run least_inflight then uid affinity) +""" + +from __future__ import annotations + +import asyncio +import importlib.util +import os +import re +import time +from pathlib import Path +from typing import Any +from uuid import uuid4 + +import pytest +import ray +import torch +from omegaconf import OmegaConf +from verl.workers.rollout.replica import RolloutMode + +from verl_omni.workers.rollout.omni_llm_server import OmniLLMServerClient +from verl_omni.workers.rollout.replica import DiffusionOutput +from verl_omni.workers.rollout.request_routing import OmniRequestLoadBalancer +from verl_omni.workers.rollout.vllm_rollout.vllm_omni_async_server import vLLMOmniHttpServer + +_flowgrpo_spec = importlib.util.spec_from_file_location( + "flowgrpo_batch_helpers", + Path(__file__).with_name("test_vllm_omni_request_batch_flowgrpo.py"), +) +_flowgrpo_helpers = importlib.util.module_from_spec(_flowgrpo_spec) +assert _flowgrpo_spec.loader is not None +_flowgrpo_spec.loader.exec_module(_flowgrpo_helpers) + +DEFAULT_MODEL_PATH = _flowgrpo_helpers.DEFAULT_MODEL_PATH +DEFAULT_OCR_PARQUET = _flowgrpo_helpers.DEFAULT_OCR_PARQUET +REQUEST_BATCH_MAX_WAIT_MS = _flowgrpo_helpers.REQUEST_BATCH_MAX_WAIT_MS +REQUEST_BATCH_MIN_SIZE = _flowgrpo_helpers.REQUEST_BATCH_MIN_SIZE +flowgrpo_training_sampling_params = _flowgrpo_helpers.flowgrpo_training_sampling_params +load_flowgrpo_ocr_samples = _flowgrpo_helpers.load_flowgrpo_ocr_samples +parse_batch_verify_log = _flowgrpo_helpers.parse_batch_verify_log +summarize_batch_histogram = _flowgrpo_helpers.summarize_batch_histogram +tokenize_chat_messages = _flowgrpo_helpers.tokenize_chat_messages + +NUM_REPLICAS = int(os.environ.get("FLOWGRPO_ROUTING_NUM_REPLICAS", "4")) +NUM_PROMPTS = int(os.environ.get("FLOWGRPO_ROUTING_NUM_PROMPTS", "32")) +ROLLOUT_N = int(os.environ.get("FLOWGRPO_ROUTING_N", "16")) +ROUTING_POLICY = os.environ.get("FLOWGRPO_ROUTING_POLICY", "prompt_uid_affinity") +COMPARE_POLICIES = os.environ.get("FLOWGRPO_ROUTING_COMPARE_POLICIES", "1") == "1" +MAX_NUM_SEQS = int(os.environ.get("FLOWGRPO_REQUEST_BATCH_MAX_NUM_SEQS", "32")) +STEP_EXECUTION = os.environ.get("FLOWGRPO_ROUTING_STEP_EXECUTION", "0") == "1" +NUM_INFERENCE_STEPS = int(os.environ.get("FLOWGRPO_ROUTING_NUM_INFERENCE_STEPS", "10")) +_MIN_PROMPT_TOKENS = 35 + +_NUM_REQS_RE = re.compile(r"num_reqs=(\d+)") + + +def _build_rollout_cfg() -> Any: + return OmegaConf.create( + { + "_target_": "verl_omni.workers.config.diffusion.DiffusionRolloutConfig", + "name": "vllm_omni", + "mode": "async", + "tensor_model_parallel_size": 1, + "data_parallel_size": 1, + "pipeline_model_parallel_size": 1, + "gpu_memory_utilization": float(os.environ.get("FLOWGRPO_ROUTING_GPU_MEM_UTIL", "0.45")), + "max_num_batched_tokens": 8192, + "max_num_seqs": MAX_NUM_SEQS, + "dtype": "bfloat16", + "load_format": "safetensors", + "enforce_eager": True, + "enable_chunked_prefill": False, + "enable_prefix_caching": False, + "enable_sleep_mode": False, + "free_cache_engine": False, + "disable_log_stats": True, + "calculate_log_probs": True, + "step_execution": STEP_EXECUTION, + "engine_kwargs": { + "vllm_omni": { + "step_execution": STEP_EXECUTION, + } + }, + "pipeline": { + "_target_": "verl_omni.workers.config.diffusion.DiffusionPipelineConfig", + "height": 512, + "width": 512, + "num_inference_steps": NUM_INFERENCE_STEPS, + "true_cfg_scale": 4.0, + "max_sequence_length": 256, + }, + "algo": { + "_target_": "verl_omni.workers.config.diffusion.DiffusionRolloutAlgoConfig", + "noise_level": 1.2, + "sde_type": "sde", + "sde_window_size": 2, + "sde_window_range": [0, 5], + }, + "server_routing": { + "_target_": "verl_omni.workers.config.RolloutServerRoutingConfig", + "policy": ROUTING_POLICY, + }, + } + ) + + +def _build_model_cfg() -> Any: + return OmegaConf.create( + { + "_target_": "verl_omni.workers.config.diffusion.DiffusionModelConfig", + "path": str(DEFAULT_MODEL_PATH), + "tokenizer_path": str(DEFAULT_MODEL_PATH / "tokenizer"), + "trust_remote_code": True, + "load_tokenizer": True, + "algorithm": "flow_grpo", + } + ) + + +def _build_trainer_cfg(policy: str) -> Any: + rollout_cfg = _build_rollout_cfg() + rollout_cfg.server_routing.policy = policy + return OmegaConf.create( + { + "actor_rollout_ref": { + "rollout": rollout_cfg, + "model": _build_model_cfg(), + } + } + ) + + +async def _run_routed_burst( + *, + client: OmniLLMServerClient, + tokenized: list[tuple[list[int], list[int]]], + num_prompts: int, + rollout_n: int, +) -> float: + async def _one_generate(prompt_idx: int, rollout_idx: int, prompt_ids: list[int], neg_ids: list[int]): + uid = f"flowgrpo-uid-{prompt_idx}" + return await client.generate( + request_id=uuid4().hex, + prompt_ids=prompt_ids, + negative_prompt_ids=neg_ids, + sampling_params=flowgrpo_training_sampling_params(seed=1000 + prompt_idx * rollout_n + rollout_idx), + routing_key=uid, + ) + + tasks = [] + for prompt_idx in range(num_prompts): + prompt_ids, neg_ids = tokenized[prompt_idx % len(tokenized)] + for rollout_idx in range(rollout_n): + tasks.append(_one_generate(prompt_idx, rollout_idx, prompt_ids, neg_ids)) + + start = time.perf_counter() + results = await asyncio.gather(*tasks) + elapsed = time.perf_counter() - start + + for output in results: + assert isinstance(output, DiffusionOutput) + assert output.diffusion_output is not None + return elapsed + + +@pytest.fixture(scope="module") +def multi_server_cluster(tmp_path_factory): + if not torch.cuda.is_available() or torch.cuda.device_count() < NUM_REPLICAS: + pytest.skip(f"requires >= {NUM_REPLICAS} CUDA devices") + if not DEFAULT_MODEL_PATH.is_dir() or not DEFAULT_OCR_PARQUET.is_file(): + pytest.skip("model or OCR parquet missing") + + verify_log = tmp_path_factory.mktemp("routing_perf") / "flowgrpo_batch_verify.log" + os.environ["FLOWGRPO_BATCH_VERIFY_LOG"] = str(verify_log) + + ray.init( + runtime_env={ + "env_vars": { + "TOKENIZERS_PARALLELISM": "true", + "NCCL_DEBUG": "WARN", + "VLLM_LOGGING_LEVEL": "INFO", + "FLOWGRPO_BATCH_VERIFY_LOG": str(verify_log), + } + }, + ignore_reinit_error=True, + ) + + rollout_cfg = _build_rollout_cfg() + model_cfg = _build_model_cfg() + ServerCls = ray.remote(vLLMOmniHttpServer) + servers = [] + server_map: dict[str, ray.actor.ActorHandle] = {} + + for replica_rank in range(NUM_REPLICAS): + server = ServerCls.options( + runtime_env={ + "env_vars": { + "RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES": "1", + "RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES": "1", + "NCCL_CUMEM_ENABLE": "0", + "FLOWGRPO_BATCH_VERIFY_LOG": str(verify_log), + } + }, + max_concurrency=NUM_PROMPTS * ROLLOUT_N + 8, + ).remote( + config=rollout_cfg, + model_config=model_cfg, + rollout_mode=RolloutMode.STANDALONE, + workers=[], + replica_rank=replica_rank, + node_rank=0, + gpus_per_node=NUM_REPLICAS, + nnodes=1, + cuda_visible_devices=str(replica_rank), + ) + servers.append(server) + server_map[f"replica-{replica_rank}"] = server + + ray.get([server.launch_server.remote() for server in servers]) + + yield servers, server_map, verify_log + ray.shutdown() + + +@pytest.mark.parametrize("policy", ["prompt_uid_affinity"]) +def test_flowgrpo_routing_policy_performance(multi_server_cluster, policy: str): + servers, server_map, verify_log = multi_server_cluster + del servers # accessed via LB + + samples = load_flowgrpo_ocr_samples(DEFAULT_OCR_PARQUET, num_samples=NUM_PROMPTS) + tokenized = [ + ( + tokenize_chat_messages(DEFAULT_MODEL_PATH, sample["raw_prompt"]), + tokenize_chat_messages(DEFAULT_MODEL_PATH, sample["raw_negative_prompt"]), + ) + for sample in samples + ] + + policies = ["least_inflight", "prompt_uid_affinity"] if COMPARE_POLICIES else [policy] + results: list[dict[str, Any]] = [] + + for routing_policy in policies: + verify_log.write_text("") + lb = OmniRequestLoadBalancer.remote(servers=server_map, policy=routing_policy) + trainer_cfg = _build_trainer_cfg(routing_policy) + client = OmniLLMServerClient(config=trainer_cfg, load_balancer_handle=lb) + + elapsed = asyncio.run( + _run_routed_burst( + client=client, + tokenized=tokenized, + num_prompts=NUM_PROMPTS, + rollout_n=ROLLOUT_N, + ) + ) + + hist = parse_batch_verify_log(verify_log) + num_reqs_hist = hist["num_reqs"] + max_num_reqs = max(num_reqs_hist) if num_reqs_hist else 0 + lb_status = ray.get(lb.get_status.remote()) + + row = { + "policy": routing_policy, + "elapsed_s": elapsed, + "max_num_reqs": max_num_reqs, + "total_requests": NUM_PROMPTS * ROLLOUT_N, + "lb_status": lb_status, + "num_reqs_hist": num_reqs_hist, + } + results.append(row) + + print( + f"\n=== Routing benchmark ({routing_policy}) ===\n" + f"replicas={NUM_REPLICAS} prompts={NUM_PROMPTS} rollout_n={ROLLOUT_N}\n" + f"wall_time_s={elapsed:.2f}\n" + f"max_num_reqs={max_num_reqs}\n" + f"lb_status={lb_status}\n" + f"num_reqs histogram:\n{summarize_batch_histogram(num_reqs_hist)}\n" + ) + + if COMPARE_POLICIES and len(results) == 2: + least = next(r for r in results if r["policy"] == "least_inflight") + affinity = next(r for r in results if r["policy"] == "prompt_uid_affinity") + speedup = least["elapsed_s"] / affinity["elapsed_s"] if affinity["elapsed_s"] > 0 else 0.0 + print( + "\n=== Policy comparison ===\n" + f"least_inflight: {least['elapsed_s']:.2f}s, max_num_reqs={least['max_num_reqs']}\n" + f"prompt_uid_affinity: {affinity['elapsed_s']:.2f}s, max_num_reqs={affinity['max_num_reqs']}\n" + f"affinity_speedup={speedup:.2f}x\n" + ) + assert affinity["max_num_reqs"] >= least["max_num_reqs"], ( + "prompt_uid_affinity should co-locate rollout.n copies and improve fused batch sizes" + ) + + final = results[-1] + if not STEP_EXECUTION: + assert final["max_num_reqs"] >= 8, ( + f"Expected fused batches under {final['policy']}, got max_num_reqs={final['max_num_reqs']}" + ) diff --git a/tests/workers/rollout/rollout_vllm/test_vllm_omni_request_batch_flowgrpo.py b/tests/workers/rollout/rollout_vllm/test_vllm_omni_request_batch_flowgrpo.py new file mode 100644 index 00000000..4acfc443 --- /dev/null +++ b/tests/workers/rollout/rollout_vllm/test_vllm_omni_request_batch_flowgrpo.py @@ -0,0 +1,331 @@ +# Copyright 2026 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Debug/regression test for vllm-omni request-wise batching under FlowGRPO OCR. + +Replays OCR train-parquet prompts, negative prompts, and FlowGRPO sampling +params, then fires a burst of concurrent ``vLLMOmniHttpServer.generate()`` +calls — the same ingress path as ``DiffusionSingleTurnAgentLoop``. + +Instrumentation in ``QwenImagePipelineWithLogProb`` writes per-forward batch +sizes to ``FLOWGRPO_BATCH_VERIFY_LOG``. This test parses that log and fails +with a histogram when fused batch sizes stay too small. + +Usage: + FLOWGRPO_BATCH_VERIFY_LOG=/tmp/flowgrpo_batch_verify_test.log \\ + pytest tests/workers/rollout/rollout_vllm/test_vllm_omni_request_batch_flowgrpo.py -v -s + +Optional env overrides: + FLOWGRPO_REQUEST_BATCH_MODEL_PATH + FLOWGRPO_REQUEST_BATCH_OCR_PARQUET + FLOWGRPO_REQUEST_BATCH_NUM_CONCURRENT (default: 32, train_batch_size) + FLOWGRPO_REQUEST_BATCH_MAX_NUM_SEQS (default: 32) + FLOWGRPO_MIN_EXPECTED_MAX_BATCH (default: 8; smoke: 1) + FLOWGRPO_REQUEST_BATCH_MAX_WAIT_MS (default: 250) + FLOWGRPO_REQUEST_BATCH_MIN_SIZE (default: 0, auto max_num_seqs//2) +""" + +from __future__ import annotations + +import os +import re +from collections import Counter +from pathlib import Path +from typing import Any +from uuid import uuid4 + +import pandas as pd +import pytest +import ray +import torch +from omegaconf import OmegaConf +from transformers import AutoTokenizer +from verl.utils.tokenizer import normalize_token_ids +from verl.workers.rollout.replica import RolloutMode + +from verl_omni.workers.rollout.replica import DiffusionOutput +from verl_omni.workers.rollout.vllm_rollout.vllm_omni_async_server import vLLMOmniHttpServer + +DEFAULT_MODEL_PATH = Path(os.environ.get("FLOWGRPO_REQUEST_BATCH_MODEL_PATH", "/home/public/yx/models/Qwen/Qwen-Image")) +DEFAULT_OCR_PARQUET = Path( + os.environ.get( + "FLOWGRPO_REQUEST_BATCH_OCR_PARQUET", + "/home/public/yx/data/ocr/qwen_image/train.parquet", + ) +) +NUM_CONCURRENT = int(os.environ.get("FLOWGRPO_REQUEST_BATCH_NUM_CONCURRENT", "32")) +MAX_NUM_SEQS = int(os.environ.get("FLOWGRPO_REQUEST_BATCH_MAX_NUM_SEQS", "32")) +MIN_EXPECTED_MAX_BATCH = int(os.environ.get("FLOWGRPO_MIN_EXPECTED_MAX_BATCH", "8")) +REQUEST_BATCH_MAX_WAIT_MS = float(os.environ.get("FLOWGRPO_REQUEST_BATCH_MAX_WAIT_MS", "250")) +REQUEST_BATCH_MIN_SIZE = int(os.environ.get("FLOWGRPO_REQUEST_BATCH_MIN_SIZE", "0")) +_MIN_PROMPT_TOKENS = 35 + +_NUM_REQS_RE = re.compile(r"num_reqs=(\d+)") +_FUSED_B_RE = re.compile(r"fused_B=(\d+)") + + +def _messages_from_parquet_cell(cell: Any) -> list[dict[str, str]]: + return [{"role": item["role"], "content": item["content"]} for item in cell] + + +def load_flowgrpo_ocr_samples(parquet_path: Path, *, num_samples: int) -> list[dict[str, Any]]: + """Load OCR parquet rows in the same message shape used by RLHFDataset.""" + df = pd.read_parquet(parquet_path) + samples: list[dict[str, Any]] = [] + for idx in range(min(num_samples, len(df))): + row = df.iloc[idx] + samples.append( + { + "raw_prompt": _messages_from_parquet_cell(row["prompt"]), + "raw_negative_prompt": _messages_from_parquet_cell(row["negative_prompt"]), + "ground_truth": row["reward_model"]["ground_truth"], + } + ) + return samples + + +def tokenize_chat_messages(model_path: Path, messages: list[dict[str, str]]) -> list[int]: + tokenizer = AutoTokenizer.from_pretrained( + model_path / "tokenizer", + trust_remote_code=True, + ) + token_ids = normalize_token_ids(tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=False)) + assert len(token_ids) > _MIN_PROMPT_TOKENS, ( + f"Prompt too short ({len(token_ids)} tokens, need >{_MIN_PROMPT_TOKENS}). " + "Qwen-Image drops the first 34 chat-template prefix tokens." + ) + return token_ids + + +def flowgrpo_training_sampling_params(*, seed: int) -> dict[str, Any]: + """Sampling params aligned with ``run_qwen_image_ocr_lora.sh`` rollout.""" + return { + "height": 512, + "width": 512, + "num_inference_steps": 10, + "true_cfg_scale": 4.0, + "max_sequence_length": 256, + "noise_level": 1.2, + "sde_type": "sde", + "sde_window_size": 2, + "sde_window_range": (0, 5), + "logprobs": True, + "seed": seed, + } + + +def parse_batch_verify_log(log_path: Path) -> dict[str, Counter[int]]: + """Parse ``[flowgrpo_reqbatch]`` lines into histograms.""" + num_reqs_hist: Counter[int] = Counter() + fused_b_hist: Counter[int] = Counter() + if not log_path.exists(): + return {"num_reqs": num_reqs_hist, "fused_B": fused_b_hist} + + with log_path.open(encoding="utf-8") as fh: + for line in fh: + if "num_reqs=" in line: + for match in _NUM_REQS_RE.finditer(line): + num_reqs_hist[int(match.group(1))] += 1 + if "fused_B=" in line: + for match in _FUSED_B_RE.finditer(line): + fused_b_hist[int(match.group(1))] += 1 + return {"num_reqs": num_reqs_hist, "fused_B": fused_b_hist} + + +def summarize_batch_histogram(hist: Counter[int]) -> str: + if not hist: + return "(empty)" + lines = [f" size={size}: {count}" for size, count in sorted(hist.items())] + return "\n".join(lines) + + +pytestmark = [ + pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA"), + pytest.mark.skipif(not DEFAULT_MODEL_PATH.is_dir(), reason=f"model missing at {DEFAULT_MODEL_PATH}"), + pytest.mark.skipif(not DEFAULT_OCR_PARQUET.is_file(), reason=f"OCR parquet missing at {DEFAULT_OCR_PARQUET}"), +] + + +@pytest.fixture +def flowgrpo_request_batch_server(tmp_path): + """Launch one vLLMOmniHttpServer with FlowGRPO OCR rollout settings.""" + verify_log = tmp_path / "flowgrpo_batch_verify.log" + os.environ["FLOWGRPO_BATCH_VERIFY_LOG"] = str(verify_log) + + ray.init( + runtime_env={ + "env_vars": { + "TOKENIZERS_PARALLELISM": "true", + "NCCL_DEBUG": "WARN", + "VLLM_LOGGING_LEVEL": "INFO", + "FLOWGRPO_BATCH_VERIFY_LOG": str(verify_log), + } + }, + ignore_reinit_error=True, + ) + + rollout_cfg = OmegaConf.create( + { + "_target_": "verl_omni.workers.config.diffusion.DiffusionRolloutConfig", + "name": "vllm_omni", + "mode": "async", + "tensor_model_parallel_size": 1, + "data_parallel_size": 1, + "pipeline_model_parallel_size": 1, + "gpu_memory_utilization": 0.5, + "max_num_batched_tokens": 8192, + "max_num_seqs": MAX_NUM_SEQS, + "dtype": "bfloat16", + "load_format": "safetensors", + "enforce_eager": True, + "enable_chunked_prefill": False, + "enable_prefix_caching": False, + "enable_sleep_mode": False, + "free_cache_engine": False, + "disable_log_stats": True, + "calculate_log_probs": True, + "engine_kwargs": { + "vllm_omni": { + "step_execution": False, + "request_batch_max_wait_ms": REQUEST_BATCH_MAX_WAIT_MS, + "request_batch_min_size": REQUEST_BATCH_MIN_SIZE, + } + }, + "pipeline": { + "_target_": "verl_omni.workers.config.diffusion.DiffusionPipelineConfig", + "height": 512, + "width": 512, + "num_inference_steps": 10, + "true_cfg_scale": 4.0, + "max_sequence_length": 256, + }, + "algo": { + "_target_": "verl_omni.workers.config.diffusion.DiffusionRolloutAlgoConfig", + "noise_level": 1.2, + "sde_type": "sde", + "sde_window_size": 2, + "sde_window_range": [0, 5], + }, + } + ) + + model_cfg = OmegaConf.create( + { + "_target_": "verl_omni.workers.config.diffusion.DiffusionModelConfig", + "path": str(DEFAULT_MODEL_PATH), + "tokenizer_path": str(DEFAULT_MODEL_PATH / "tokenizer"), + "trust_remote_code": True, + "load_tokenizer": True, + "algorithm": "flow_grpo", + } + ) + + ServerCls = ray.remote(vLLMOmniHttpServer) + server = ServerCls.options( + runtime_env={ + "env_vars": { + "RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES": "1", + "RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES": "1", + "NCCL_CUMEM_ENABLE": "0", + "FLOWGRPO_BATCH_VERIFY_LOG": str(verify_log), + } + }, + max_concurrency=NUM_CONCURRENT + 4, + ).remote( + config=rollout_cfg, + model_config=model_cfg, + rollout_mode=RolloutMode.STANDALONE, + workers=[], + replica_rank=0, + node_rank=0, + gpus_per_node=1, + nnodes=1, + cuda_visible_devices="0", + ) + + ray.get(server.launch_server.remote()) + + yield server, verify_log + + ray.shutdown() + + +def test_flowgrpo_concurrent_request_batch_sizes(flowgrpo_request_batch_server): + """Burst concurrent FlowGRPO OCR requests and report fused scheduler batch sizes.""" + server, verify_log = flowgrpo_request_batch_server + samples = load_flowgrpo_ocr_samples(DEFAULT_OCR_PARQUET, num_samples=NUM_CONCURRENT) + tokenized = [ + ( + tokenize_chat_messages(DEFAULT_MODEL_PATH, sample["raw_prompt"]), + tokenize_chat_messages(DEFAULT_MODEL_PATH, sample["raw_negative_prompt"]), + ) + for sample in samples + ] + + refs = [] + for i, (prompt_ids, negative_prompt_ids) in enumerate(tokenized): + rid = f"flowgrpo_batch_{i}_{uuid4().hex[:8]}" + refs.append( + server.generate.remote( + prompt_ids=prompt_ids, + negative_prompt_ids=negative_prompt_ids, + sampling_params=flowgrpo_training_sampling_params(seed=1000 + i), + request_id=rid, + ) + ) + + results = ray.get(refs, timeout=3600) + for i, output in enumerate(results): + assert isinstance(output, DiffusionOutput), f"request {i}: expected DiffusionOutput" + assert output.diffusion_output is not None + + hist = parse_batch_verify_log(verify_log) + num_reqs_hist = hist["num_reqs"] + fused_b_hist = hist["fused_B"] + max_num_reqs = max(num_reqs_hist) if num_reqs_hist else 0 + max_fused_b = max(fused_b_hist) if fused_b_hist else 0 + + print( + "\n=== FlowGRPO request-batch debug summary ===\n" + f"concurrent_requests={NUM_CONCURRENT}\n" + f"max_num_seqs={MAX_NUM_SEQS}\n" + f"request_batch_max_wait_ms={REQUEST_BATCH_MAX_WAIT_MS}\n" + f"request_batch_min_size={REQUEST_BATCH_MIN_SIZE}\n" + f"verify_log={verify_log}\n" + f"max_num_reqs_seen={max_num_reqs}\n" + f"max_fused_B_seen={max_fused_b}\n" + "num_reqs histogram:\n" + f"{summarize_batch_histogram(num_reqs_hist)}\n" + "fused_B histogram:\n" + f"{summarize_batch_histogram(fused_b_hist)}\n" + ) + + assert num_reqs_hist, ( + f"No [flowgrpo_reqbatch] lines in {verify_log}. " + "Ensure QwenImagePipelineWithLogProb batch instrumentation is active." + ) + + assert max_num_reqs >= MIN_EXPECTED_MAX_BATCH, ( + "vllm-omni request-wise batching did not reach the expected fused batch size " + f"under FlowGRPO OCR ingress (max_num_reqs={max_num_reqs}, " + f"max_fused_B={max_fused_b}, max_num_seqs={MAX_NUM_SEQS}, " + f"concurrent={NUM_CONCURRENT}).\n" + "num_reqs histogram:\n" + f"{summarize_batch_histogram(num_reqs_hist)}\n" + "fused_B histogram:\n" + f"{summarize_batch_histogram(fused_b_hist)}\n" + "This reproduces training ingress: OCR parquet prompts, true_cfg_scale=4.0, " + "negative prompts, per-request seeds, and burst concurrent HTTP generates." + ) diff --git a/tests/workers/rollout/test_rollout_server_routing.py b/tests/workers/rollout/test_rollout_server_routing.py new file mode 100644 index 00000000..04ed5197 --- /dev/null +++ b/tests/workers/rollout/test_rollout_server_routing.py @@ -0,0 +1,132 @@ +# Copyright 2026 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +from collections import Counter + +import pytest +import ray +from omegaconf import OmegaConf + +from verl_omni.workers.config.rollout_routing import RolloutServerRoutingConfig +from verl_omni.workers.rollout.request_routing import ( + OmniRequestLoadBalancer, + stable_shard_index, +) + + +@ray.remote +class _DummyServer: + pass + + +def _make_lb(policy: str, num_servers: int = 4) -> ray.actor.ActorHandle: + servers = {f"s{i}": _DummyServer.remote() for i in range(num_servers)} + return OmniRequestLoadBalancer.remote(servers=servers, policy=policy) + + +@pytest.fixture(scope="module", autouse=True) +def _ray_init(): + if not ray.is_initialized(): + ray.init(ignore_reinit_error=True, num_cpus=4) + yield + ray.shutdown() + + +def test_stable_shard_index_is_deterministic(): + assert stable_shard_index("uid-42", 4) == stable_shard_index("uid-42", 4) + assert 0 <= stable_shard_index("uid-42", 4) < 4 + + +def test_rollout_server_routing_config_from_rollout_yaml(): + import os + + from hydra import compose, initialize_config_dir + + with initialize_config_dir(config_dir=os.path.abspath("verl_omni/trainer/config"), version_base=None): + cfg = compose(config_name="diffusion_trainer") + assert OmegaConf.select(cfg.actor_rollout_ref.rollout, "server_routing.policy") == "prompt_uid_affinity" + assert OmegaConf.select(cfg.actor_rollout_ref.rollout, "server_routing.routing_key_field", default="uid") == "uid" + + +def test_rollout_server_routing_config_override(): + cfg = OmegaConf.create( + { + "actor_rollout_ref": { + "rollout": { + "server_routing": { + "_target_": "verl_omni.workers.config.RolloutServerRoutingConfig", + "policy": "prompt_uid_affinity", + } + } + } + } + ) + assert OmegaConf.select(cfg.actor_rollout_ref.rollout, "server_routing.policy") == "prompt_uid_affinity" + + +def test_rollout_server_routing_config_rejects_invalid_policy(): + with pytest.raises(ValueError, match="Invalid rollout server_routing.policy"): + RolloutServerRoutingConfig(policy="random_shuffle") + + +@pytest.mark.parametrize("policy", ["prompt_hash_sharding", "round_robin"]) +def test_hash_and_round_robin_cover_all_servers(policy: str): + lb = _make_lb(policy) + counts: Counter[str] = Counter() + for i in range(64): + server_id, _ = ray.get(lb.acquire_server.remote(request_id=f"req-{i}", routing_key=f"uid-{i % 8}")) + counts[server_id] += 1 + ray.get(lb.release_server.remote(server_id)) + assert len(counts) == 4 + assert all(count > 0 for count in counts.values()) + + +def test_prompt_uid_affinity_clusters_rollout_copies(): + """FlowGRPO pattern: rollout.n copies of the same uid must share one replica.""" + lb = _make_lb("prompt_uid_affinity") + per_uid_server: dict[str, str] = {} + + num_prompts = 32 + rollout_n = 16 + for uid in range(num_prompts): + uid_key = f"prompt-{uid}" + acquired_servers: list[str] = [] + for _ in range(rollout_n): + server_id, _ = ray.get(lb.acquire_server.remote(request_id=f"req-{uid}-{_}", routing_key=uid_key)) + per_uid_server.setdefault(uid_key, server_id) + assert per_uid_server[uid_key] == server_id + acquired_servers.append(server_id) + + for server_id in acquired_servers: + ray.get(lb.release_server.remote(server_id)) + + assert len(per_uid_server) == num_prompts + + +def test_least_inflight_spreads_unique_request_ids(): + lb = _make_lb("least_inflight") + server_load: Counter[str] = Counter() + acquired_servers: list[str] = [] + for i in range(64): + server_id, _ = ray.get(lb.acquire_server.remote(request_id=f"unique-req-{i}")) + server_load[server_id] += 1 + acquired_servers.append(server_id) + + for server_id in acquired_servers: + ray.get(lb.release_server.remote(server_id)) + + assert len(server_load) == 4 + assert max(server_load.values()) - min(server_load.values()) <= 8 diff --git a/verl_omni/agent_loop/single_turn_agent_loop.py b/verl_omni/agent_loop/single_turn_agent_loop.py index 9d154482..b59349f1 100644 --- a/verl_omni/agent_loop/single_turn_agent_loop.py +++ b/verl_omni/agent_loop/single_turn_agent_loop.py @@ -16,6 +16,7 @@ from typing import Any from uuid import uuid4 +from omegaconf import OmegaConf from verl.experimental.agent_loop.agent_loop import AgentLoopBase, register from verl.utils.profiler import simple_timer @@ -58,6 +59,10 @@ async def run(self, sampling_params: dict[str, Any], **kwargs) -> DiffusionAgent negative_prompt_ids = None # 3. generate sequences + routing_key_field = OmegaConf.select(self.rollout_config, "server_routing.routing_key_field", default="uid") + value = kwargs.get(routing_key_field) if routing_key_field else None + routing_key = str(value) if value is not None else None + metrics = {} with simple_timer("generate_sequences", metrics): output = await self.server_manager.generate( @@ -67,6 +72,7 @@ async def run(self, sampling_params: dict[str, Any], **kwargs) -> DiffusionAgent image_data=images, video_data=videos, negative_prompt_ids=negative_prompt_ids, + routing_key=routing_key, ) if metrics.get("num_preempted") is None: metrics["num_preempted"] = output.num_preempted if output.num_preempted is not None else -1 diff --git a/verl_omni/trainer/config/_generated_diffusion_trainer.yaml b/verl_omni/trainer/config/_generated_diffusion_trainer.yaml index 69e001b3..2f7437bf 100644 --- a/verl_omni/trainer/config/_generated_diffusion_trainer.yaml +++ b/verl_omni/trainer/config/_generated_diffusion_trainer.yaml @@ -172,6 +172,10 @@ actor_rollout_ref: trace_alloc_max_entries: ${oc.select:global_profiler.global_tool_config.torch_memory.trace_alloc_max_entries,100000} stack_depth: ${oc.select:global_profiler.global_tool_config.torch_memory.stack_depth,32} name: torch_memory + server_routing: + _target_: verl_omni.workers.config.RolloutServerRoutingConfig + policy: prompt_uid_affinity + max_imbalance: null _target_: verl_omni.workers.config.diffusion.DiffusionRolloutConfig name: ??? mode: async diff --git a/verl_omni/trainer/config/_generated_diffusion_veomni_trainer.yaml b/verl_omni/trainer/config/_generated_diffusion_veomni_trainer.yaml index eed93cf0..8be68a2d 100644 --- a/verl_omni/trainer/config/_generated_diffusion_veomni_trainer.yaml +++ b/verl_omni/trainer/config/_generated_diffusion_veomni_trainer.yaml @@ -213,6 +213,10 @@ actor_rollout_ref: trace_alloc_max_entries: ${oc.select:global_profiler.global_tool_config.torch_memory.trace_alloc_max_entries,100000} stack_depth: ${oc.select:global_profiler.global_tool_config.torch_memory.stack_depth,32} name: torch_memory + server_routing: + _target_: verl_omni.workers.config.RolloutServerRoutingConfig + policy: prompt_uid_affinity + max_imbalance: null _target_: verl_omni.workers.config.diffusion.DiffusionRolloutConfig name: ??? mode: async diff --git a/verl_omni/trainer/config/diffusion/rollout/diffusion_rollout.yaml b/verl_omni/trainer/config/diffusion/rollout/diffusion_rollout.yaml index 91b436d6..2072a3c7 100644 --- a/verl_omni/trainer/config/diffusion/rollout/diffusion_rollout.yaml +++ b/verl_omni/trainer/config/diffusion/rollout/diffusion_rollout.yaml @@ -4,6 +4,9 @@ defaults: # per-role profiler config, inheriting from trainer/config/profiler/profiler.yaml - ../../profiler@profiler: profiler + # shared rollout replica routing (diffusion + omni LLM) + - /rollout/server_routing@server_routing + # load the reference default config, then apply the fields in the current yaml - _self_ @@ -87,6 +90,13 @@ logprobs_mode: processed_logprobs # scheduling policy for vllm-omni rollout scheduling_policy: fcfs +# Replica routing across rollout HTTP servers (verl-side load balancer). +# Distinct from vllm-omni internal diffusion scheduler batching. +# Base schema from trainer/config/rollout/server_routing.yaml; diffusion defaults +# to prompt_uid_affinity so rollout.n copies for one prompt co-locate on a replica. +server_routing: + policy: prompt_uid_affinity + # Which loader to use for rollout model weights: dummy, hf, megatron, etc. # safetensors (for huge model, and set use_shm=True); dummy: randomly init model weight load_format: dummy diff --git a/verl_omni/trainer/config/rollout/server_routing.yaml b/verl_omni/trainer/config/rollout/server_routing.yaml new file mode 100644 index 00000000..dffcbb10 --- /dev/null +++ b/verl_omni/trainer/config/rollout/server_routing.yaml @@ -0,0 +1,15 @@ +# Shared rollout replica routing defaults (diffusion + omni LLM async rollouts). +# Distinct from vllm-omni internal diffusion scheduler ``scheduling_policy``. +defaults: + - _self_ + +_target_: verl_omni.workers.config.RolloutServerRoutingConfig + +# least_inflight | prompt_uid_affinity | prompt_hash_sharding | round_robin +policy: least_inflight + +# For soft affinity (prompt_uid_affinity), maximum allowed imbalance in queue length +# before re-routing to an underloaded replica. Set to null or <= 0 for strict affinity. +max_imbalance: null + +# routing_key_field defaults to uid (FlowGRPO prompt group id in per-sample kwargs). diff --git a/verl_omni/trainer/diffusion/ray_diffusion_trainer.py b/verl_omni/trainer/diffusion/ray_diffusion_trainer.py index 60a67af2..119d8bef 100644 --- a/verl_omni/trainer/diffusion/ray_diffusion_trainer.py +++ b/verl_omni/trainer/diffusion/ray_diffusion_trainer.py @@ -49,7 +49,6 @@ from verl.utils.metric import reduce_metrics from verl.utils.py_functional import rename_dict from verl.utils.tracking import ValidationGenerationsLogger -from verl.workers.rollout.llm_server import LLMServerManager from verl_omni.trainer.config import DiffusionAlgoConfig from verl_omni.trainer.diffusion.diffusion_algos import ( @@ -70,6 +69,7 @@ apply_rollout_correction_to_diffusion_batch, rollout_correction_enabled, ) +from verl_omni.workers.rollout.omni_llm_server import OmniLLMServerManager from verl_omni.workers.utils.padding import embeds_padding_2_no_padding sys_logger = logging.getLogger(__name__) @@ -673,7 +673,7 @@ def _init_online_rollout_stack(self, actor_rollout_resource_pool): # to stream reward computation with actor rollout reward_loop_worker_handles = self.reward_loop_manager.reward_loop_workers if enable_agent_reward_loop else None - self.llm_server_manager = LLMServerManager.create( + self.llm_server_manager = OmniLLMServerManager.create( config=self.config, worker_group=self.actor_rollout_wg, rollout_resource_pool=actor_rollout_resource_pool, diff --git a/verl_omni/workers/config/diffusion/rollout.py b/verl_omni/workers/config/diffusion/rollout.py index 830169d1..5f71b40c 100644 --- a/verl_omni/workers/config/diffusion/rollout.py +++ b/verl_omni/workers/config/diffusion/rollout.py @@ -26,11 +26,14 @@ PrometheusConfig, ) +from verl_omni.workers.config.rollout_routing import RolloutServerRoutingConfig + __all__ = [ "DiffusionRolloutAlgoConfig", "DiffusionPipelineConfig", "DiffusionSamplingConfig", "DiffusionRolloutConfig", + "RolloutServerRoutingConfig", ] @@ -166,6 +169,9 @@ class DiffusionRolloutConfig(BaseConfig): disaggregation: DisaggregationConfig = field(default_factory=DisaggregationConfig) + # Replica routing for async rollout HTTP servers (diffusion + omni LLM). + server_routing: RolloutServerRoutingConfig = field(default_factory=RolloutServerRoutingConfig) + external_lib: Optional[str] = None def __post_init__(self): diff --git a/verl_omni/workers/config/rollout_routing.py b/verl_omni/workers/config/rollout_routing.py new file mode 100644 index 00000000..abf0e61a --- /dev/null +++ b/verl_omni/workers/config/rollout_routing.py @@ -0,0 +1,61 @@ +# Copyright 2026 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from dataclasses import dataclass +from typing import Optional + +from verl.base_config import BaseConfig + +VALID_ROLLOUT_SERVER_ROUTING_POLICIES = frozenset( + { + "least_inflight", + "prompt_uid_affinity", + "prompt_hash_sharding", + "round_robin", + } +) + + +@dataclass +class RolloutServerRoutingConfig(BaseConfig): + """How agent loops route HTTP rollout requests across omni server replicas. + + This applies to any async rollout using ``OmniLLMServerClient`` and + ``OmniRequestLoadBalancer`` (diffusion and omni LLM). It controls + **verl-side replica selection**, not the internal vllm-omni diffusion + scheduler policy (``scheduling_policy``). + """ + + # Replica routing policy: + # - least_inflight: spread requests for fairness (verl default) + # - prompt_uid_affinity: sticky route by routing_key (e.g. batch uid) + # - prompt_hash_sharding: deterministic shard by hash(routing_key) + # - round_robin: rotate replicas regardless of load + policy: str = "least_inflight" + + # Field on per-sample agent kwargs used as routing_key (e.g. ``uid`` for + # FlowGRPO rollout.n copies). Ignored when null. + routing_key_field: Optional[str] = "uid" + + # For soft affinity (prompt_uid_affinity), maximum allowed imbalance in queue length + # (inflight requests) between the assigned replica and the least loaded replica + # before re-routing. If null or <= 0, strict/hard affinity is used. + max_imbalance: Optional[int] = None + + def __post_init__(self) -> None: + if self.policy not in VALID_ROLLOUT_SERVER_ROUTING_POLICIES: + raise ValueError( + f"Invalid rollout server_routing.policy={self.policy!r}. " + f"Must be one of {sorted(VALID_ROLLOUT_SERVER_ROUTING_POLICIES)}." + ) diff --git a/verl_omni/workers/rollout/omni_llm_server.py b/verl_omni/workers/rollout/omni_llm_server.py new file mode 100644 index 00000000..77d213c9 --- /dev/null +++ b/verl_omni/workers/rollout/omni_llm_server.py @@ -0,0 +1,109 @@ +# Copyright 2026 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import logging +import os +from typing import Any, Optional +from uuid import uuid4 + +import ray +from omegaconf import DictConfig, OmegaConf +from verl.utils.rollout_trace import rollout_trace_op +from verl.workers.rollout.llm_server import LLMServerClient, LLMServerManager +from verl.workers.rollout.replica import TokenOutput + +from verl_omni.workers.rollout.request_routing import ( + DEFAULT_ROUTING_CACHE_SIZE, + OmniRequestLoadBalancer, +) + +logger = logging.getLogger(__name__) +logger.setLevel(os.getenv("VERL_LOGGING_LEVEL", "WARN")) + + +class OmniLLMServerClient(LLMServerClient): + """LLM server client that routes requests via ``OmniRequestLoadBalancer``.""" + + def __init__(self, config: DictConfig, load_balancer_handle: ray.actor.ActorHandle | None = None, **kwargs): + super().__init__(config=config, load_balancer_handle=load_balancer_handle, **kwargs) + + async def _acquire_server( + self, + request_id: str, + routing_key: str | None = None, + ) -> tuple[str, ray.actor.ActorHandle]: + return await self._load_balancer.acquire_server.remote( + request_id=request_id, + routing_key=routing_key, + ) + + @rollout_trace_op + async def generate( + self, + request_id, + *, + prompt_ids: list[int], + sampling_params: dict[str, Any], + image_data: Optional[list[Any]] = None, + video_data: Optional[list[Any]] = None, + audio_data: Optional[list[Any]] = None, + mm_processor_kwargs: Optional[dict[str, Any]] = None, + routing_key: str | None = None, + **kwargs: Any, + ) -> TokenOutput: + server_id, server = await self._acquire_server(request_id, routing_key=routing_key) + try: + multimodal_kwargs = {} + if audio_data is not None: + multimodal_kwargs["audio_data"] = audio_data + if mm_processor_kwargs: + multimodal_kwargs["mm_processor_kwargs"] = mm_processor_kwargs + output: TokenOutput = await server.generate.remote( + request_id=uuid4().hex, + prompt_ids=prompt_ids, + sampling_params=sampling_params, + image_data=image_data, + video_data=video_data, + **multimodal_kwargs, + **kwargs, + ) + return output + finally: + self._release_server(server_id) + + +class OmniLLMServerManager(LLMServerManager): + """Launch rollout replicas with ``OmniRequestLoadBalancer``.""" + + async def _init_global_load_balancer(self) -> None: + policy = OmegaConf.select(self.rollout_config, "server_routing.policy", default="least_inflight") + routing_key_field = OmegaConf.select(self.rollout_config, "server_routing.routing_key_field", default="uid") + max_imbalance = OmegaConf.select(self.rollout_config, "server_routing.max_imbalance", default=8) + self.global_load_balancer = OmniRequestLoadBalancer.remote( + servers=dict(zip(self.server_addresses, self.server_handles, strict=True)), + policy=policy, + max_cache_size=DEFAULT_ROUTING_CACHE_SIZE, + max_imbalance=max_imbalance, + ) + logger.info( + "[OmniLLMServerManager] OmniRequestLoadBalancer policy=%s routing_key_field=%s max_imbalance=%s", + policy, + routing_key_field, + max_imbalance, + ) + + def get_client(self, client_cls=OmniLLMServerClient, **kwargs) -> OmniLLMServerClient: + return super().get_client(client_cls=client_cls, **kwargs) diff --git a/verl_omni/workers/rollout/request_routing.py b/verl_omni/workers/rollout/request_routing.py new file mode 100644 index 00000000..1b105de9 --- /dev/null +++ b/verl_omni/workers/rollout/request_routing.py @@ -0,0 +1,154 @@ +# Copyright 2026 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import hashlib +import logging +import os + +import ray +from cachetools import LRUCache + +from verl_omni.workers.config.rollout_routing import ( + VALID_ROLLOUT_SERVER_ROUTING_POLICIES, +) + +logger = logging.getLogger(__name__) +logger.setLevel(os.getenv("VERL_LOGGING_LEVEL", "WARN")) + +DEFAULT_ROUTING_CACHE_SIZE = 10000 + + +def stable_shard_index(key: str, num_shards: int) -> int: + if num_shards <= 0: + raise ValueError(f"num_shards must be positive, got {num_shards}") + digest = hashlib.sha256(key.encode("utf-8")).digest() + return int.from_bytes(digest[:8], "big") % num_shards + + +@ray.remote +class OmniRequestLoadBalancer: + """Multi-replica request load balancer for verl-omni rollout servers. + + Replaces verl's ``GlobalRequestLoadBalancer`` with pluggable routing policies + (for example ``prompt_uid_affinity`` and ``least_inflight``) for diffusion and + omni async rollouts. + """ + + def __init__( + self, + servers: dict[str, ray.actor.ActorHandle], + policy: str = "least_inflight", + max_cache_size: int = DEFAULT_ROUTING_CACHE_SIZE, + max_imbalance: int | None = None, + ) -> None: + if not servers: + raise ValueError("servers must be non-empty") + if policy not in VALID_ROLLOUT_SERVER_ROUTING_POLICIES: + raise ValueError(f"Unsupported routing policy: {policy!r}") + + self._servers: dict[str, ray.actor.ActorHandle] = dict(servers) + self._policy = policy + self._max_imbalance = max_imbalance + self._inflight_requests: dict[str, int] = {sid: 0 for sid in servers} + self._request_id_to_server: LRUCache = LRUCache(maxsize=max_cache_size) + self._round_robin_idx = 0 + + def acquire_server( + self, + request_id: str, + routing_key: str | None = None, + ) -> tuple[str, ray.actor.ActorHandle]: + sticky_key = routing_key or request_id + + if self._policy == "prompt_hash_sharding": + server_id = self._pick_sharded_server(sticky_key) + self._inflight_requests[server_id] += 1 + return server_id, self._servers[server_id] + + if self._policy == "round_robin": + server_id = self._pick_round_robin_server() + self._inflight_requests[server_id] += 1 + return server_id, self._servers[server_id] + + if self._policy == "prompt_uid_affinity": + return self._acquire_sticky(sticky_key) + + return self._acquire_sticky(request_id) + + def _pick_sharded_server(self, sticky_key: str) -> str: + server_ids = sorted(self._servers.keys()) + return server_ids[stable_shard_index(sticky_key, len(server_ids))] + + def _pick_round_robin_server(self) -> str: + server_ids = sorted(self._servers.keys()) + server_id = server_ids[self._round_robin_idx % len(server_ids)] + self._round_robin_idx += 1 + return server_id + + def _acquire_sticky(self, sticky_key: str) -> tuple[str, ray.actor.ActorHandle]: + if sticky_key in self._request_id_to_server: + server_id = self._request_id_to_server[sticky_key] + if server_id in self._inflight_requests: + # Soft Affinity load check + if ( + self._policy == "prompt_uid_affinity" + and self._max_imbalance is not None + and self._max_imbalance > 0 + ): + min_server_id = min(self._inflight_requests, key=self._inflight_requests.get) + if ( + self._inflight_requests[server_id] - self._inflight_requests[min_server_id] + > self._max_imbalance + ): + # Divert request to least loaded replica and update stickiness mapping + server_id = min_server_id + self._request_id_to_server[sticky_key] = server_id + + self._inflight_requests[server_id] += 1 + return server_id, self._servers[server_id] + del self._request_id_to_server[sticky_key] + + if not self._inflight_requests: + raise RuntimeError("No available servers in load balancer") + + server_id = min(self._inflight_requests, key=self._inflight_requests.get) + self._request_id_to_server[sticky_key] = server_id + self._inflight_requests[server_id] += 1 + return server_id, self._servers[server_id] + + def release_server(self, server_id: str) -> None: + if server_id not in self._inflight_requests: + return + if self._inflight_requests[server_id] > 0: + self._inflight_requests[server_id] -= 1 + + def add_servers(self, servers: dict[str, ray.actor.ActorHandle]) -> None: + for sid, handle in servers.items(): + self._inflight_requests[sid] = 0 + self._servers[sid] = handle + + def remove_servers(self, server_ids: list[str]) -> None: + for sid in server_ids: + self._inflight_requests.pop(sid, None) + self._servers.pop(sid, None) + + def get_status(self) -> dict: + return { + "policy": self._policy, + "servers": dict(self._inflight_requests), + "total_inflight": sum(self._inflight_requests.values()), + "active_servers": len(self._inflight_requests), + }