rollout: decouple oversampling-abort teardown into a pluggable agent hook

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

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

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:

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.