From 7ca5ae38a7a18e8eb7942f9defbada60272faa17 Mon Sep 17 00:00:00 2001 From: Shi Dong Date: Sat, 11 Jul 2026 22:50:53 -0700 Subject: [PATCH 1/3] rollout: flush Harbor agent server on oversampling abort When oversampling collects enough samples, abort() aborts the in-flight SGLang generations but leaves the Harbor agent loops running. Those loops keep issuing fresh completion requests to SGLang until they hit their own max_seq_len / timeout, wasting inference and holding containers. Flush the agent server's in-flight /run trials (agent-agnostic, so it works for mini-swe-agent which bypasses harbor's lite_llm) right after the SGLang abort so they stop hitting SGLang and release their containers. No-op unless --agent-server-url / AGENT_SERVER_URL and session_server_instance_id are set. --- miles/rollout/sglang_rollout.py | 35 +++++++++++++++++++++++++++++++++ miles/utils/arguments.py | 9 +++++++++ 2 files changed, 44 insertions(+) diff --git a/miles/rollout/sglang_rollout.py b/miles/rollout/sglang_rollout.py index 64662ed522..e03a60481b 100644 --- a/miles/rollout/sglang_rollout.py +++ b/miles/rollout/sglang_rollout.py @@ -1,6 +1,7 @@ import asyncio import copy import logging +import os import uuid from argparse import Namespace from collections.abc import Callable @@ -340,6 +341,36 @@ async def generate_and_rm_group( return group +async def flush_agent_server(args: Namespace) -> None: + """Cancel in-flight Harbor agent-server trials for this run. + + Aborting SGLang only stops the in-flight generations; the Harbor agent loops + keep running and keep issuing fresh completion requests to SGLang until they + hit their own max_seq_len / timeout. Flushing the agent server cancels the + in-flight ``/run`` tasks (agent-agnostic) and releases their containers. + """ + agent_server_url = getattr(args, "agent_server_url", None) or os.environ.get("AGENT_SERVER_URL") + instance_id = getattr(args, "session_server_instance_id", None) + if not agent_server_url or not instance_id: + return + + headers = None + admin_secret = os.environ.get("HARBOR_ADMIN_SECRET") + if admin_secret: + headers = {"Authorization": f"Bearer {admin_secret}"} + + try: + result = await post( + f"{agent_server_url.rstrip('/')}/flush", + {"session_server_instance_id": instance_id}, + max_retries=3, + headers=headers, + ) + logger.info(f"Flushed agent server {agent_server_url}: {result}") + except Exception as e: + logger.warning(f"Failed to flush agent server {agent_server_url}: {e}") + + async def abort(args: Namespace, rollout_id: int) -> list[list[Sample]]: aborted_samples = [] @@ -361,6 +392,10 @@ async def abort(args: Namespace, rollout_id: int) -> list[list[Sample]]: if isinstance(result, Exception): logger.warning(f"Failed to abort worker at {url}: {result}") + # Cancel in-flight Harbor trials so they stop hitting SGLang and release their + # containers instead of running on until their own max_seq_len / timeout. + await flush_agent_server(args) + # make sure all the pending tasks are finished count = 0 while state.pendings: diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index 53dcc34f8e..55a0b4b847 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -2033,6 +2033,15 @@ def add_session_arguments(parser): default=None, help="Port of the standalone session server. Auto-allocated if not set.", ) + parser.add_argument( + "--agent-server-url", + type=str, + default=os.environ.get("AGENT_SERVER_URL"), + help="Base URL of the Harbor agent server driving agentic rollouts. " + "When set, oversampling abort also flushes the server's in-flight " + "trials so they stop hitting SGLang. Defaults to the AGENT_SERVER_URL " + "env var.", + ) parser.add_argument( "--tito-model", type=str, From b893ba00788fdb40671b4a233e00acf036ead091 Mon Sep 17 00:00:00 2001 From: Shi Dong Date: Sat, 11 Jul 2026 23:03:39 -0700 Subject: [PATCH 2/3] rollout: decouple abort teardown into a pluggable agent hook Replace the hard-coded Harbor /flush in core sglang_rollout.abort() with a generic, optional hook: it looks for a sibling abort() in the module named by --custom-agent-function-path and calls it if present. Core Miles no longer assumes the agent host speaks /flush; non-agentic and non-Harbor backends just drain as before. Move the Harbor /flush teardown into the swe_agent_function plugin (where the agent-server URL and /run already live) as its abort() hook. Drop the core --agent-server-url arg added earlier. --- .../swe-agent-v2/swe_agent_function.py | 31 ++++++++++++ miles/rollout/sglang_rollout.py | 47 +++++++++---------- miles/utils/arguments.py | 9 ---- 3 files changed, 53 insertions(+), 34 deletions(-) diff --git a/examples/experimental/swe-agent-v2/swe_agent_function.py b/examples/experimental/swe-agent-v2/swe_agent_function.py index d1460fbe06..b347240b08 100644 --- a/examples/experimental/swe-agent-v2/swe_agent_function.py +++ b/examples/experimental/swe-agent-v2/swe_agent_function.py @@ -92,3 +92,34 @@ async def run( "eval_report": response.get("eval_report", {}), "agent_metrics": response.get("agent_metrics", {}), } + + +async def abort(args) -> None: + """Teardown hook for oversampling abort (called by sglang_rollout.abort). + + When Miles has enough samples and aborts SGLang, the in-flight Harbor trials + keep looping and hitting SGLang until they hit their own max_seq_len/timeout. + Flush the agent server so it cancels those ``/run`` tasks and releases their + containers. No-op unless AGENT_SERVER_URL and session_server_instance_id are + available. + """ + agent_server_url = os.getenv("AGENT_SERVER_URL", os.getenv("SWE_AGENT_URL")) + instance_id = getattr(args, "session_server_instance_id", None) + if not agent_server_url or not instance_id: + return + + headers = None + admin_secret = os.getenv("HARBOR_ADMIN_SECRET") + if admin_secret: + headers = {"Authorization": f"Bearer {admin_secret}"} + + try: + result = await post( + f"{agent_server_url.rstrip('/')}/flush", + {"session_server_instance_id": instance_id}, + max_retries=3, + headers=headers, + ) + logger.info(f"Flushed agent server {agent_server_url}: {result}") + except Exception as e: + logger.warning(f"Failed to flush agent server {agent_server_url}: {e}") diff --git a/miles/rollout/sglang_rollout.py b/miles/rollout/sglang_rollout.py index e03a60481b..75a9015df2 100644 --- a/miles/rollout/sglang_rollout.py +++ b/miles/rollout/sglang_rollout.py @@ -1,7 +1,6 @@ import asyncio import copy import logging -import os import uuid from argparse import Namespace from collections.abc import Callable @@ -341,34 +340,32 @@ async def generate_and_rm_group( return group -async def flush_agent_server(args: Namespace) -> None: - """Cancel in-flight Harbor agent-server trials for this run. +async def _call_agent_abort_hook(args: Namespace) -> None: + """Invoke the agent plugin's optional abort hook, if it defines one. - Aborting SGLang only stops the in-flight generations; the Harbor agent loops - keep running and keep issuing fresh completion requests to SGLang until they - hit their own max_seq_len / timeout. Flushing the agent server cancels the - in-flight ``/run`` tasks (agent-agnostic) and releases their containers. + Aborting SGLang only stops the in-flight generations; an external agent loop + (driven by ``--custom-agent-function-path``) keeps running and keeps issuing + fresh completion requests until it hits its own limit. The agent integration + knows how to tell its backend to stop, so we look for a sibling ``abort`` + callable in the same module as the configured agent function and call it. + Backends that don't expose one are simply left to drain as before. """ - agent_server_url = getattr(args, "agent_server_url", None) or os.environ.get("AGENT_SERVER_URL") - instance_id = getattr(args, "session_server_instance_id", None) - if not agent_server_url or not instance_id: + agent_function_path = getattr(args, "custom_agent_function_path", None) + if not agent_function_path: return - headers = None - admin_secret = os.environ.get("HARBOR_ADMIN_SECRET") - if admin_secret: - headers = {"Authorization": f"Bearer {admin_secret}"} + module_path, _, _ = agent_function_path.rpartition(".") + if not module_path: + return + try: + abort_hook = load_function(f"{module_path}.abort") + except (AttributeError, ModuleNotFoundError): + return # plugin doesn't expose an abort hook; nothing to tear down try: - result = await post( - f"{agent_server_url.rstrip('/')}/flush", - {"session_server_instance_id": instance_id}, - max_retries=3, - headers=headers, - ) - logger.info(f"Flushed agent server {agent_server_url}: {result}") + await abort_hook(args) except Exception as e: - logger.warning(f"Failed to flush agent server {agent_server_url}: {e}") + logger.warning(f"Agent abort hook {module_path}.abort failed: {e}") async def abort(args: Namespace, rollout_id: int) -> list[list[Sample]]: @@ -392,9 +389,9 @@ async def abort(args: Namespace, rollout_id: int) -> list[list[Sample]]: if isinstance(result, Exception): logger.warning(f"Failed to abort worker at {url}: {result}") - # Cancel in-flight Harbor trials so they stop hitting SGLang and release their - # containers instead of running on until their own max_seq_len / timeout. - await flush_agent_server(args) + # 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. + await _call_agent_abort_hook(args) # make sure all the pending tasks are finished count = 0 diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index 55a0b4b847..53dcc34f8e 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -2033,15 +2033,6 @@ def add_session_arguments(parser): default=None, help="Port of the standalone session server. Auto-allocated if not set.", ) - parser.add_argument( - "--agent-server-url", - type=str, - default=os.environ.get("AGENT_SERVER_URL"), - help="Base URL of the Harbor agent server driving agentic rollouts. " - "When set, oversampling abort also flushes the server's in-flight " - "trials so they stop hitting SGLang. Defaults to the AGENT_SERVER_URL " - "env var.", - ) parser.add_argument( "--tito-model", type=str, From 052df28951331b5b1089b45cfc6413f458d0c291 Mon Sep 17 00:00:00 2001 From: Shi Dong Date: Sat, 11 Jul 2026 23:13:07 -0700 Subject: [PATCH 3/3] rollout: also fire the agent abort hook on the async (inference_rollout) path Async training (train_async.py + MILES_EXPERIMENTAL_ROLLOUT_REFACTOR=1) uses InferenceRolloutFn, whose oversampling abort lives in inference_rollout_train.abort, a separate implementation from sglang_rollout.abort. The previous commit only wired the legacy sync path, so async agentic runs still leaked in-flight trials. Hoist the hook resolver into miles.utils.misc.call_agent_abort_hook (shared, no new cross-module dep) and call it from both abort paths. --- .../inference_rollout_train.py | 6 +++- miles/rollout/sglang_rollout.py | 32 ++----------------- miles/utils/misc.py | 28 ++++++++++++++++ 3 files changed, 35 insertions(+), 31 deletions(-) diff --git a/miles/rollout/inference_rollout/inference_rollout_train.py b/miles/rollout/inference_rollout/inference_rollout_train.py index eed15557bd..ac7b9a742b 100644 --- a/miles/rollout/inference_rollout/inference_rollout_train.py +++ b/miles/rollout/inference_rollout/inference_rollout_train.py @@ -13,7 +13,7 @@ from miles.rollout.inference_rollout.inference_rollout_common import GenerateState, generate_and_rm_group from miles.utils import dumper_utils from miles.utils.http_utils import get, post -from miles.utils.misc import as_completed_async, load_function +from miles.utils.misc import as_completed_async, call_agent_abort_hook, load_function from miles.utils.types import Sample logger = logging.getLogger(__name__) @@ -29,6 +29,10 @@ async def abort(state: GenerateState, pendings: set, rollout_id: int) -> list[li 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. + await call_agent_abort_hook(args) + # make sure all the pending tasks are finished aborted_samples = [] async for group in as_completed_async(pendings): diff --git a/miles/rollout/sglang_rollout.py b/miles/rollout/sglang_rollout.py index 75a9015df2..e9bf9aa6a7 100644 --- a/miles/rollout/sglang_rollout.py +++ b/miles/rollout/sglang_rollout.py @@ -22,7 +22,7 @@ from miles.utils.eval_config import EvalDatasetConfig from miles.utils.http_utils import get, post from miles.utils.lora import LORA_ADAPTER_NAME, is_lora_enabled -from miles.utils.misc import SingletonMeta, load_function +from miles.utils.misc import SingletonMeta, call_agent_abort_hook, load_function from miles.utils.processing_utils import ( call_processor, encode_image_for_rollout_engine, @@ -340,34 +340,6 @@ async def generate_and_rm_group( return group -async def _call_agent_abort_hook(args: Namespace) -> None: - """Invoke the agent plugin's optional abort hook, if it defines one. - - Aborting SGLang only stops the in-flight generations; an external agent loop - (driven by ``--custom-agent-function-path``) keeps running and keeps issuing - fresh completion requests until it hits its own limit. The agent integration - knows how to tell its backend to stop, so we look for a sibling ``abort`` - callable in the same module as the configured agent function and call it. - Backends that don't expose one are simply left to drain as before. - """ - agent_function_path = getattr(args, "custom_agent_function_path", None) - if not agent_function_path: - return - - module_path, _, _ = agent_function_path.rpartition(".") - if not module_path: - return - try: - abort_hook = load_function(f"{module_path}.abort") - except (AttributeError, ModuleNotFoundError): - return # plugin doesn't expose an abort hook; nothing to tear down - - try: - await abort_hook(args) - except Exception as e: - logger.warning(f"Agent abort hook {module_path}.abort failed: {e}") - - async def abort(args: Namespace, rollout_id: int) -> list[list[Sample]]: aborted_samples = [] @@ -391,7 +363,7 @@ async def abort(args: Namespace, rollout_id: int) -> list[list[Sample]]: # 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. - await _call_agent_abort_hook(args) + await call_agent_abort_hook(args) # make sure all the pending tasks are finished count = 0 diff --git a/miles/utils/misc.py b/miles/utils/misc.py index 3c160cc5cb..d99baf69c4 100644 --- a/miles/utils/misc.py +++ b/miles/utils/misc.py @@ -62,6 +62,34 @@ def load_function(path): return getattr(module, attr) +async def call_agent_abort_hook(args) -> None: + """Invoke the agent plugin's optional abort hook, if it defines one. + + When oversampling collects enough samples, the rollout aborts SGLang, but an + external agent loop (driven by ``--custom-agent-function-path``) keeps running + and keeps issuing fresh completion requests until it hits its own limit. The + agent integration knows how to tell its backend to stop, so we look for a + sibling ``abort`` callable in the same module as the configured agent function + and call it. Backends that don't expose one are left to drain as before. + """ + agent_function_path = getattr(args, "custom_agent_function_path", None) + if not agent_function_path: + return + + module_path, _, _ = agent_function_path.rpartition(".") + if not module_path: + return + try: + abort_hook = load_function(f"{module_path}.abort") + except (AttributeError, ModuleNotFoundError): + return # plugin doesn't expose an abort hook; nothing to tear down + + try: + await abort_hook(args) + except Exception as e: + logger.warning(f"Agent abort hook {module_path}.abort failed: {e}") + + class SingletonMeta(type): """ A metaclass for creating singleton classes.