Use the standard router abort path for Omni

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 "

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")
 

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.

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.

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)."""