From 4f24a84ff1279366a8dba6889fbd26852d8adffa Mon Sep 17 00:00:00 2001 From: Jingwen Gu Date: Tue, 28 Jul 2026 00:12:28 +0000 Subject: [PATCH] Use the standard router abort path for Omni --- .../run_qwen3_omni_thinker_grpo.py | 14 ++++---- .../generate_utils/generate_endpoint_utils.py | 18 ---------- .../inference_rollout_train.py | 11 ++---- miles/rollout/sglang_rollout.py | 29 +++++++-------- .../rollout/generate_hub/test_sglang_omni.py | 35 ------------------- 5 files changed, 23 insertions(+), 84 deletions(-) diff --git a/examples/omni_thinker/run_qwen3_omni_thinker_grpo.py b/examples/omni_thinker/run_qwen3_omni_thinker_grpo.py index b45146c5bc..f5aecff543 100644 --- a/examples/omni_thinker/run_qwen3_omni_thinker_grpo.py +++ b/examples/omni_thinker/run_qwen3_omni_thinker_grpo.py @@ -1,8 +1,9 @@ """GRPO on the Qwen3-Omni-30B-A3B thinker with audio-input AVQA (MCQ reward). -Topology: rollout runs on a standalone sglang-omni text server (external engines; miles -launches nothing locally); the trainer holds the extracted text backbone and injects -frozen-audio-tower embeddings at placeholder positions (--qwen3-omni-audio-encoder-path). +Topology: rollout runs on a standalone sglang-omni text worker registered behind an +sglang-omni router (external engines; miles launches nothing locally); the trainer holds +the extracted text backbone and injects frozen-audio-tower embeddings at placeholder +positions (--qwen3-omni-audio-encoder-path). Weight sync: `--sync-mode skip` freezes the server (off-policy debug; TIS absorbs the gap), `--sync-mode distributed` pushes thinker.* weights over NCCL each step (on-policy). @@ -38,6 +39,8 @@ class ScriptArgs(U.ExecuteTrainConfig): megatron_path: str = "/root/Megatron-LM" omni_server_ip: str = "127.0.0.1" omni_server_port: int = 30000 + omni_router_ip: str = "127.0.0.1" + omni_router_port: int = 30001 omni_server_tp: int = 4 # TP size of the external omni server (NCCL group world_size = tp + 1) avqa_max_samples: int = 5120 extra_args: str = "" @@ -96,9 +99,8 @@ def execute(args: ScriptArgs): "--rollout-temperature 1 " f"--global-batch-size {32 if debug_minimal else 256} " "--balance-data " - # the standalone omni server doubles as the router: the adapter posts straight to it - f"--sglang-router-ip {args.omni_server_ip} " - f"--sglang-router-port {args.omni_server_port} " + f"--sglang-router-ip {args.omni_router_ip} " + f"--sglang-router-port {args.omni_router_port} " "--rollout-external " f"--rollout-external-engine-addrs {args.omni_server_ip}:{args.omni_server_port} " "--rollout-external-admin-api sglang-omni " diff --git a/miles/rollout/generate_utils/generate_endpoint_utils.py b/miles/rollout/generate_utils/generate_endpoint_utils.py index a45a0631bb..37472d553b 100644 --- a/miles/rollout/generate_utils/generate_endpoint_utils.py +++ b/miles/rollout/generate_utils/generate_endpoint_utils.py @@ -9,7 +9,6 @@ import pybase64 import torch -from miles.utils.http_utils import post from miles.utils.lora import LORA_ADAPTER_NAME, is_lora_enabled from miles.utils.processing_utils import ( call_processor, @@ -63,23 +62,6 @@ def compute_prompt_ids_from_sample(state, sample, tools=None): return state.tokenizer.encode(prompt, add_special_tokens=False) -def is_omni_external_admin(args) -> bool: - return getattr(args, "rollout_external_admin_api", "sglang") == "sglang-omni" - - -async def abort_external_omni_requests(args) -> None: - """Abort in-flight generation on external sglang-omni servers. - - A bare omni server exposes neither the router /workers listing nor /abort_request; - pause_generation(mode="abort") followed by continue_generation is the omni-native - equivalent of abort_all. - """ - for addr in args.rollout_external_engine_addrs: - base_url = f"http://{addr}" - await post(f"{base_url}/pause_generation", {"mode": "abort"}) - await post(f"{base_url}/continue_generation", {}) - - def policy_uses_routing_key(args) -> bool: return args.sglang_router_policy in ("consistent_hashing", "manual") diff --git a/miles/rollout/inference_rollout/inference_rollout_train.py b/miles/rollout/inference_rollout/inference_rollout_train.py index 64626455a2..ac7b9a742b 100644 --- a/miles/rollout/inference_rollout/inference_rollout_train.py +++ b/miles/rollout/inference_rollout/inference_rollout_train.py @@ -9,7 +9,6 @@ from miles.rollout.base_types import RolloutFnTrainOutput from miles.rollout.filter_hub.base_types import MetricGatherer, call_dynamic_filter -from miles.rollout.generate_utils.generate_endpoint_utils import abort_external_omni_requests, is_omni_external_admin from miles.rollout.generate_utils.prefill_logprobs import recompute_samples_rollout_logprobs_via_prefill from miles.rollout.inference_rollout.inference_rollout_common import GenerateState, generate_and_rm_group from miles.utils import dumper_utils @@ -26,13 +25,9 @@ async def abort(state: GenerateState, pendings: set, rollout_id: int) -> list[li assert not state.aborted state.aborted = True - if is_omni_external_admin(args): - # bare omni servers have no router /workers nor /abort_request - await abort_external_omni_requests(args) - else: - urls = await get_worker_urls(args) - logger.info(f"Abort request for {urls}") - await asyncio.gather(*[post(f"{url}/abort_request", {"abort_all": True}) for url in urls]) + urls = await get_worker_urls(args) + logger.info(f"Abort request for {urls}") + await asyncio.gather(*[post(f"{url}/abort_request", {"abort_all": True}) for url in urls]) # Let the agent integration tear down its in-flight trials so they stop hitting # SGLang, instead of running on until their own max_seq_len / timeout. diff --git a/miles/rollout/sglang_rollout.py b/miles/rollout/sglang_rollout.py index bcc23e6f90..7e6f60e182 100644 --- a/miles/rollout/sglang_rollout.py +++ b/miles/rollout/sglang_rollout.py @@ -15,7 +15,6 @@ from miles.rollout.base_types import GenerateFnInput, RolloutFnEvalOutput, RolloutFnTrainOutput from miles.rollout.filter_hub.base_types import MetricGatherer, call_dynamic_filter -from miles.rollout.generate_utils.generate_endpoint_utils import abort_external_omni_requests, is_omni_external_admin from miles.rollout.inference_rollout.compatibility import load_generate_function from miles.utils import dumper_utils from miles.utils.async_utils import run @@ -378,23 +377,19 @@ async def abort(args: Namespace, rollout_id: int) -> list[list[Sample]]: assert not state.aborted state.aborted = True - if is_omni_external_admin(args): - # bare omni servers have no router /workers nor /abort_request - await abort_external_omni_requests(args) + if parse(sglang_router.__version__) <= parse("0.2.1") or args.use_miles_router: + response = await get(f"http://{args.sglang_router_ip}:{args.sglang_router_port}/list_workers") + urls = response["urls"] else: - if parse(sglang_router.__version__) <= parse("0.2.1") or args.use_miles_router: - response = await get(f"http://{args.sglang_router_ip}:{args.sglang_router_port}/list_workers") - urls = response["urls"] - else: - response = await get(f"http://{args.sglang_router_ip}:{args.sglang_router_port}/workers") - urls = [worker["url"] for worker in response["workers"]] - - logger.info(f"Abort request for {urls}") - abort_tasks = [post(f"{url}/abort_request", {"abort_all": True}) for url in urls] - abort_results = await asyncio.gather(*abort_tasks, return_exceptions=True) - for url, result in zip(urls, abort_results, strict=False): - if isinstance(result, Exception): - logger.warning(f"Failed to abort worker at {url}: {result}") + response = await get(f"http://{args.sglang_router_ip}:{args.sglang_router_port}/workers") + urls = [worker["url"] for worker in response["workers"]] + + logger.info(f"Abort request for {urls}") + abort_tasks = [post(f"{url}/abort_request", {"abort_all": True}) for url in urls] + abort_results = await asyncio.gather(*abort_tasks, return_exceptions=True) + for url, result in zip(urls, abort_results, strict=False): + if isinstance(result, Exception): + logger.warning(f"Failed to abort worker at {url}: {result}") # Let the agent integration tear down its in-flight trials so they stop hitting # SGLang, instead of running on until their own max_seq_len / timeout. diff --git a/tests/fast/rollout/generate_hub/test_sglang_omni.py b/tests/fast/rollout/generate_hub/test_sglang_omni.py index c4f426550e..daee9954ff 100644 --- a/tests/fast/rollout/generate_hub/test_sglang_omni.py +++ b/tests/fast/rollout/generate_hub/test_sglang_omni.py @@ -204,41 +204,6 @@ async def fail_post(url, payload, headers=None): assert output.samples is sample -def test_abort_external_omni_requests_uses_pause_abort_continue(monkeypatch): - # a bare omni server has neither the router /workers listing nor /abort_request; - # pause(mode=abort) + continue is the omni-native equivalent - from miles.rollout.generate_utils import generate_endpoint_utils as geu - - calls = [] - - async def fake_post(url, payload, **kwargs): - calls.append((url, payload)) - return {} - - monkeypatch.setattr(geu, "post", fake_post) - args = SimpleNamespace( - rollout_external_admin_api="sglang-omni", - rollout_external_engine_addrs=["10.0.0.1:30111", "10.0.0.2:30111"], - ) - - asyncio.run(geu.abort_external_omni_requests(args)) - - assert calls == [ - ("http://10.0.0.1:30111/pause_generation", {"mode": "abort"}), - ("http://10.0.0.1:30111/continue_generation", {}), - ("http://10.0.0.2:30111/pause_generation", {"mode": "abort"}), - ("http://10.0.0.2:30111/continue_generation", {}), - ] - - -def test_is_omni_external_admin(monkeypatch): - from miles.rollout.generate_utils import generate_endpoint_utils as geu - - assert geu.is_omni_external_admin(SimpleNamespace(rollout_external_admin_api="sglang-omni")) is True - assert geu.is_omni_external_admin(SimpleNamespace(rollout_external_admin_api="sglang")) is False - assert geu.is_omni_external_admin(SimpleNamespace()) is False - - class TestHarness: """Through the real parse_args + mock sglang server harness (generation_fixtures)."""