Avoid starting a Slurm step for every SWE-bench tool call (#522)

We currently launch a new srun step and start Enroot for every tool command, even though we already have the node allocation and container. With several SWE evaluations running together, this creates thousands of Slurm steps and puts a lot of pressure on the scheduler. Simple commands end up taking seconds, and some runs hit the outer watchdog timeout and fail with RunnerError

This change keeps one worker running inside each SWE instance’s container and reuses it for tool execution. Persistent execution prevents calling of srun for every tool call. Setup, grading, and cleanup retain their Slurm steps, while command exit codes, timeouts, and process cleanup are preserved.

I've tested this, no impact to accuracy evaluation.

Resolves this issue: #521

Author: leopck

AGENTS.md

@@ -140,6 +140,19 @@ The aggregator is a separate process (`python -m inference_endpoint.async_utils.
 - **Post-run service wait**: `drain_and_build_report` waits for service subprocesses **unbounded** (`wait_for_exit(None)`). `wait_for_exit` SIGKILLs on expiry, and the aggregator writes `final_snapshot.json` — the Report's primary source — as the last thing it does, so a parent-side deadline here trades a wedged-service hang for a lost snapshot. The abort path keeps its own ceiling (`interrupted_teardown_grace_s`, 30s), and the whole-run watchdog stays armed throughout.
 - **Histogram bucket edges are dynamic per snapshot**: log-spaced over the observed `[min, max]`. Bucket count is fixed at construction; consumers MUST re-render from the snapshot's `(lo, hi, count)` triples each frame and MUST NOT track bucket-by-index across snapshots.
 
+### SWE-bench Pyxis command transport
+
+`evaluation/swebench_service/swebench_service/pyxis_persistent.py` owns one
+long-lived Slurm step per agent environment. The container runs the packaged
+`pyxis_command_worker.sh`, staged in a private control mount separate from tool
+`/tmp`. Worker startup retries only confirmed prelaunch Slurm failures before any
+request is published. A lock serializes callers through one atomically published
+request directory and completion marker.
+Commands run in fresh PID namespaces; accepted requests are never replayed after an
+uncertain failure. `pyxis_environment.py` owns container creation, command result
+mapping, and worker-before-container cleanup. All tool commands use the persistent
+worker. Evaluation remains in `pyxis_worker.py`.
+
 ### CLI Modes
 
 CLI is auto-generated from `config/schema.py` Pydantic models via cyclopts. Fields annotated with `cyclopts.Parameter(alias="--flag")` get flat shorthands; all other fields get auto-generated dotted flags (kebab-case).
@@ -283,7 +296,7 @@ src/inference_endpoint/
 │   ├── types.py               # Pydantic: VideoPathRequest, VideoPathResponse, VideoPayloadResponse
 │   └── adapter.py             # VideoGenAdapter (HttpRequestAdapter) + VideoGenAccumulator (no-op)
 ├── evaluation/                # Accuracy evaluation (extractor, scoring, livecodebench)
-│   └── swebench_service/      # Isolated uv service for Docker-backed SWE-bench runs
+│   └── swebench_service/      # Isolated uv service for Docker/Pyxis SWE-bench runs
 ├── compliance/                # Submission compliance checks (config-lock, accuracy gate, run validity)
 │   ├── __init__.py
 │   └── checker.py             # check_submission() + Check/ComplianceReport (Edge-Agentic ruleset)

examples/10_Agentic_Inference/README.md

@@ -260,7 +260,7 @@ Reference mean values are shown in parentheses.
 
 | Metric             | Kimi K3                  | Qwen3.6-35B-A3B          | DeepSeek-V4.1-Flash      |
 | ------------------ | ------------------------ | ------------------------ | ------------------------ |
-| Inline accuracy    | `>= 58.32%` (`58.9%`)    | `>= 55.86%` (`56.43%`)   | `>= 52.36%` (`53.16%`)     |
+| Inline accuracy    | `>= 58.32%` (`58.9%`)    | `>= 55.86%` (`56.43%`)   | `>= 52.36%` (`53.16%`)   |
 | OSL per-turn mean¹ | `425-520` tokens (`472`) | `344-422` tokens (`383`) | `793-970` tokens (`882`) |
 | SWE-bench accuracy | `>= 93.5%` (`94.83%`)    | `>= 69%` (`71.7%`)       | `>= 96.4%` (`97.5%`)     |
 

pyproject.toml

@@ -96,7 +96,7 @@ dependencies = [
     "colorama==0.4.6",
     # Fix pytz-2024 import warning
     "pytz==2026.1.post1",
-    "urllib3==2.7.0",
+    "urllib3==2.8.0",
     "pyyaml==6.0.3",
     # anyio is pulled in transitively (openai -> httpx -> anyio); pinned here to
     # force past CVE-2026-63374 / CVE-2026-64847 (both fixed in 4.14.2), which
@@ -130,7 +130,7 @@ dev = [
     # patched here, pinned only inside the bfcl fork. Closes CVE-2025-68146 /
     # CVE-2026-22701 (filelock) and CVE-2026-22702 (virtualenv).
     "filelock>=3.20.3",
-    "virtualenv>=20.36.1",
+    "virtualenv==21.7.13",
     # pip is pulled in transitively (pip-audit -> pip-api -> pip); force past
     # PYSEC-2026-3721 (fixed in 26.2), which otherwise fails `uv run pip-audit`.
     "pip==26.2",

src/inference_endpoint/evaluation/swebench_service/README.md

@@ -66,10 +66,32 @@ never forwarded.
 
 During generation, the service still uses mini-swe-agent for the agent loop and
 model requests, but replaces its Docker environment with `PyxisEnvironment`. Every
-trajectory receives a named, writable Pyxis container. Each tool call becomes an
-overlapping `srun` step in that container, preserving filesystem changes across
-turns. Tool commands run in private PID namespaces so one trajectory cannot signal
-processes belonging to another trajectory.
+trajectory receives a named, writable Pyxis container and one long-lived overlapping
+`srun` command worker. Tool calls use atomic request and response files in the
+container's private `/tmp` mount, avoiding Slurm step creation and Enroot startup
+on every turn. Each worker handles one request at a time; it publishes a completion
+marker only after the command exits and its output is closed.
+Filesystem changes persist, while each command runs in a fresh shell
+and private PID namespace. Commands cannot signal the worker or other trajectories;
+remaining child processes are removed when the command's PID namespace exits.
+The service must stay on the allocated node, and node-local `TMPDIR` is recommended
+for the request files.
+
+The command worker is the packaged `swebench_service/pyxis_command_worker.sh`.
+The Python service copies it into the private `/tmp` mount and starts it with Bash
+inside the task container; no service Python installation is needed in task images.
+
+Command failures preserve their exit status and merged stdout/stderr. A command
+timeout terminates its process group with a five-second kill grace; loss of the
+worker, invalid responses, and driver deadlines fail the run as infrastructure
+errors. An accepted request is never automatically replayed because its execution
+may already have changed the repository. The worker is stopped and reaped before
+its container and temporary files are removed.
+
+All tool commands use the persistent worker.
+Container initialization, worker startup, evaluation, and cleanup still use Slurm
+steps. This reduces per-command scheduler traffic; it does not bypass allocation
+limits or guarantee any particular end-to-end evaluation time.
 
 After generation, the Pyxis worker evaluates each prediction in a fresh `srun`
 container step because the Docker-based SWE-bench evaluator cannot run on the

src/inference_endpoint/evaluation/swebench_service/swebench_service/pyxis_command_worker.sh

@@ -0,0 +1,49 @@
+#!/usr/bin/env bash
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+set -uo pipefail
+root=$1
+shift
+interpreter=("$@")
+request="$root/running"
+touch "$root/ready" || exit 70
+
+while [ ! -e "$root/stop" ]; do
+    if [ ! -d "$root/request" ]; then
+        sleep 0.05
+        continue
+    fi
+    # The host serializes callers and removes the previous result before publishing.
+    [ ! -e "$request" ] || exit 70
+    mv -- "$root/request" "$request" || exit 70
+    timeout_s=$(cat "$request/timeout") || exit 70
+    case "$timeout_s" in ''|*[!0-9]*|0) exit 70 ;; esac
+    # A sentinel preserves trailing newlines in the command and working directory.
+    cwd=$(cat "$request/cwd" && printf x) || exit 70
+    # Keep tool text out of supervisor argv so pkill -f cannot match it there.
+    unshare --pid --fork --mount-proc \
+        timeout -k 5 "$timeout_s" bash -c '
+            status=$1; cwd=$2; shift 2
+            command=$(cat <&3 && printf x) || exit 70
+            exec 3<&-
+            if cd -- "$cwd"; then
+                "$@" "${command%x}"
+                rc=$?
+            else
+                rc=125
+            fi
+            printf "%s\n" "$rc" > "$status" || exit 70
+            exit "$rc"
+        ' command-status "$request/command_status" "${cwd%x}" \
+        "${interpreter[@]}" 3< "$request/command" > "$request/output" 2>&1
+    returncode=$?
+    timed_out=0
+    # Explicit exits 124/137 are command results, not timeout notifications.
+    if [ "$(cat "$request/command_status" 2>/dev/null)" != "$returncode" ]; then
+        case "$returncode" in 124|137) timed_out=1 ;; *) exit 70 ;; esac
+    fi
+    size=$(wc -c < "$request/output") || exit 70
+    printf '%s %s %s\n' "$returncode" "$timed_out" "$((size))" > "$request/complete.tmp" || exit 70
+    mv -- "$request/complete.tmp" "$request/complete" || exit 70
+done

src/inference_endpoint/evaluation/swebench_service/swebench_service/pyxis_environment.py

@@ -6,18 +6,23 @@
 import logging
 import os
 import platform
-import random
 import re
+import shutil
 import subprocess
 import tempfile
 import threading
-import time
 import uuid
 from pathlib import Path
 from typing import Any
 
 from pydantic import AliasChoices, BaseModel, Field
 
+from .pyxis_persistent import PersistentExecChannel
+from .pyxis_slurm import SRUN_MAX_ATTEMPTS as _SRUN_MAX_ATTEMPTS
+from .pyxis_slurm import (
+    is_retryable_prelaunch_failure as _is_retryable_prelaunch_failure,
+)
+from .pyxis_slurm import wait_for_prelaunch_retry
 from .runner import RunnerError
 
 logger = logging.getLogger(__name__)
@@ -45,18 +50,8 @@
     "SLURM_CONF",
 )
 _STEP_STATUS = "/tmp/.mlperf_srun_status"
-_SRUN_MAX_ATTEMPTS = 5
-_RETRYABLE_PRELAUNCH_ERRORS = (
-    "spank_sybil: rpc request error",
-    "required plugin spank_sybil.so",
-    "failed to connect to any sack sockets",
-    "failed to create token",
-    "curl: (56) connect tunnel failed",
-    "unable to confirm allocation for job",
-)
-_IGNORABLE_SRUN_PREAMBLE_LINES = {
-    "srun: lua: Checking requeue policy with options:",
-}
+_PERSISTENT_ROOT = "/.mlperf_persistent_exec"
+_COMMAND_WORKER = Path(__file__).with_name("pyxis_command_worker.sh")
 _STEP_SCRIPT = r"""set +e
 status_path=$1
 timeout_s=$2
@@ -73,14 +68,6 @@ def safe_srun_env() -> dict[str, str]:
     return {name: os.environ[name] for name in _SAFE_SRUN_ENV if name in os.environ}
 
 
-def _is_retryable_prelaunch_failure(status: str, output: str) -> bool:
-    """Return whether Slurm rejected the step before its command started."""
-    if status != "pending":
-        return False
-    lowered = output.lower()
-    return any(marker in lowered for marker in _RETRYABLE_PRELAUNCH_ERRORS)
-
-
 def build_srun_command(
     *,
     argv: list[str],
@@ -187,15 +174,7 @@ def run_srun_step(
         ).strip()
         retryable = _is_retryable_prelaunch_failure(status, output)
         if retryable and attempt < _SRUN_MAX_ATTEMPTS:
-            backoff_s = min(2**attempt, 16)
-            delay_s = backoff_s + random.uniform(0.0, backoff_s)
-            logger.warning(
-                "Retrying Pyxis pre-launch failure in %.1fs (attempt %d/%d)",
-                delay_s,
-                attempt,
-                _SRUN_MAX_ATTEMPTS,
-            )
-            time.sleep(delay_s)
+            wait_for_prelaunch_retry(attempt)
             continue
 
         if failure_path is not None:
@@ -232,6 +211,7 @@ class PyxisEnvironmentConfig(BaseModel):
     env: dict[str, str] = Field(default_factory=dict)
     timeout_s: int = Field(
         default=30,
+        gt=0,
         validation_alias=AliasChoices("timeout_s", "timeout"),
         serialization_alias="timeout",
     )
@@ -246,70 +226,89 @@ def __init__(self, **kwargs: Any):
         self.name = f"mswe_{safe_run_id}_{uuid.uuid4().hex[:8]}"
         self._tmp = tempfile.TemporaryDirectory(prefix=f"pyxis_{self.name}_")
         self._tmp_dir = Path(self._tmp.name)
-        self._tmp_dir.chmod(0o1777)
         self._lock = threading.Lock()
         self._cleaned = False
         try:
+            (self._tmp_dir / "tmp").mkdir(mode=0o700)
+            shutil.copyfile(_COMMAND_WORKER, self._tmp_dir / _COMMAND_WORKER.name)
             # A no-op initializes and validates the named persistent container.
             run_srun_step(
                 image=self.config.image,
                 name=self.name,
-                mounts=[(self._tmp_dir, "/tmp")],
+                mounts=[(self._tmp_dir / "tmp", "/tmp")],
                 workdir=self.config.cwd,
                 argv=["true"],
-                status_path=self._tmp_dir / Path(_STEP_STATUS).name,
+                status_path=self._tmp_dir / "tmp" / Path(_STEP_STATUS).name,
                 timeout_s=self.config.timeout_s,
                 failure_path=self.config.infrastructure_failure_path,
             )
-        except RunnerError as exc:
+            self._persistent_channel = PersistentExecChannel(
+                self._tmp_dir / "channel",
+                self._persistent_server_command(),
+                safe_srun_env(),
+                failure_path=self.config.infrastructure_failure_path,
+                launch_timeout_s=self.config.timeout_s + 30,
+            )
+            self._persistent_channel.start()
+        except (RunnerError, OSError) as exc:
             self.cleanup()
             raise RunnerError(
                 f"failed to start Pyxis container for {self.config.image}"
             ) from exc
 
+    def _persistent_server_command(self) -> list[str]:
+        return build_srun_command(
+            name=self.name,
+            # Tool cleanup of /tmp must not remove the worker or its protocol files.
+            mounts=[
+                (self._tmp_dir / "tmp", "/tmp"),
+                (self._tmp_dir, _PERSISTENT_ROOT),
+            ],
+            workdir=self.config.cwd,
+            argv=[
+                "env",
+                *(f"{key}={value}" for key, value in self.config.env.items()),
+                "unshare",
+                "--pid",
+                "--fork",
+                "--mount-proc",
+                "--kill-child",
+                "bash",
+                f"{_PERSISTENT_ROOT}/{_COMMAND_WORKER.name}",
+                f"{_PERSISTENT_ROOT}/channel",
+                *self.config.interpreter,
+            ],
+        )
+
     def execute(
         self, action: dict[str, Any], cwd: str = "", *, timeout: int | None = None
     ) -> dict[str, Any]:
         command = action.get("command", "")
         logger.debug("Executing Pyxis command: %s", command)
-        argv = ["env"]
-        argv.extend(f"{key}={value}" for key, value in self.config.env.items())
-        argv.extend([*self.config.interpreter, command])
-        result = run_srun_step(
-            argv=argv,
-            status_path=self._tmp_dir / Path(_STEP_STATUS).name,
-            timeout_s=timeout or self.config.timeout_s,
-            failure_path=self.config.infrastructure_failure_path,
-            name=self.name,
-            mounts=[(self._tmp_dir, "/tmp")],
-            workdir=cwd or self.config.cwd,
+        timeout_s = self.config.timeout_s if timeout is None else timeout
+        result = self._persistent_channel.execute(
+            command=command,
+            cwd=cwd or self.config.cwd,
+            timeout_s=timeout_s,
         )
         output: dict[str, Any]
-        if result.returncode == 124:
+        if result.timed_out:
             output = {
-                "output": result.stdout,
+                "output": result.output,
                 "returncode": -1,
                 "exception_info": "The command timed out",
                 "extra": {
                     "exception_type": "TimeoutExpired",
-                    "exception": (
-                        f"command timed out after {timeout or self.config.timeout_s}s"
-                    ),
+                    "exception": f"command timed out after {timeout_s}s",
                 },
             }
         else:
             output = {
-                "output": result.stdout,
+                "output": result.output,
                 "returncode": result.returncode,
                 "exception_info": "",
             }
         lines = output.get("output", "").lstrip().splitlines(keepends=True)
-        # Some Slurm cli_filter plugins write informational messages to stderr.
-        # run_srun_step merges stderr into stdout so command errors remain visible,
-        # which can place this cluster-generated preamble before mini-swe-agent's
-        # otherwise first-line submission marker.
-        while lines and lines[0].strip() in _IGNORABLE_SRUN_PREAMBLE_LINES:
-            lines.pop(0)
         if (
             lines
             and lines[0].strip() == "COMPLETE_TASK_AND_SUBMIT_FINAL_OUTPUT"
@@ -353,6 +352,14 @@ def cleanup(self) -> None:
                 return
             self._cleaned = True
         try:
+            channel = getattr(self, "_persistent_channel", None)
+            if channel is not None:
+                try:
+                    channel.close()
+                except (OSError, subprocess.SubprocessError):
+                    logger.warning(
+                        "Could not stop Pyxis worker %s", self.name, exc_info=True
+                    )
             if os.environ.get("SLURM_JOB_ID", "").strip():
                 try:
                     subprocess.run(