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/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 64662ed522..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, @@ -361,6 +361,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}") + # 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 while state.pendings: 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.