Add Pyxis runtime for SWE-bench accuracy - #438
Conversation
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
|
MLCommons CLA bot All contributors have signed the MLCommons CLA ✍️ ✅ |
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #438 +/- ##
=======================================
Coverage ? 81.76%
=======================================
Files ? 148
Lines ? 19615
Branches ? 0
=======================================
Hits ? 16039
Misses ? 3576
Partials ? 0 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
nv-alicheng
left a comment
nv-alicheng
left a comment
There was a problem hiding this comment.
Review Council — Multi-AI Code Review
Reviewed by: Claude + Code-Quality | Depth: thorough
codex was unavailable in this environment — Claude + Code-Quality review. See the summary comment for the tiered breakdown.
| futures = [ | ||
| executor.submit(_evaluate_instance, **payload) for payload in payloads | ||
| ] | ||
| for future in concurrent.futures.as_completed(futures): |
There was a problem hiding this comment.
[Both] high (data-integrity): _run_eval swallows every per-instance failure (except Exception: ... print(...)) and main() then unconditionally calls make_run_report, which writes a result file even if zero instances were evaluated. The parent PyxisSweBenchRunner._run_eval (runner.py:747) only checks the file exists and returns success. So an infrastructure fault — srun scheduling failure, missing SLURM_JOB_ID, node down, registry auth failure (all raise/exit non-zero per instance) — is indistinguishable from a genuinely unresolved instance: the service reports a successful run with a silently degraded 0%/partial score. For a compliance-grade accuracy tool this turns a broken cluster run into a passing-shaped artifact with an undercounted score. There is no worker→runner propagation path for "all instances failed for infra reasons." This is the root of the two medium findings below (pyxis_worker.py:138, pyxis_environment.py:162) — fix the propagation here (distinguish infra failure from unresolved; fail the run when instances error for non-eval reasons) and both symptom-patches downstream become unnecessary.
| ) | ||
| with _PRINT_LOCK: | ||
| print(f"[{instance_id}]\n{result.stdout}{result.stderr}", flush=True) | ||
| if result.returncode != 0: |
There was a problem hiding this comment.
[Both] medium (error-handling): _evaluate_instance drops the instance on ANY non-zero container exit with no report and no signal (if result.returncode != 0: return). report_path.unlink(missing_ok=True) ran at the top, so a missing report makes make_run_report count the instance unresolved. But the eval script returns non-zero for a failed patch apply (1), test timeout (124), AND whenever the srun/enroot layer itself fails to run the step. A transient container/enroot/node error therefore permanently marks a possibly-correct patch unresolved, with the only trace being printed stdout. No retry, no infra-vs-patch-fail distinction. Raise (the caller has a try/except) or write an explicit error report so the two cases stay distinguishable.
| "returncode": result.returncode, | ||
| "exception_info": "", | ||
| } | ||
| except Exception as exc: |
There was a problem hiding this comment.
[Both] medium (error-handling): execute converts any subprocess exception — including subprocess.TimeoutExpired from a stuck/queued srun — into a synthetic command result with returncode: -1. Reasonable for surfacing a shell failure to the agent, but a host-side srun scheduling/timeout failure is then presented to the agent as an ordinary failed command rather than aborting the instance; combined with the swallow above, cluster-contention failures are laundered into "the model's command failed," silently biasing the trajectory. At minimum make a TimeoutExpired on the container-start/step path distinguishable from a real command exit. (The broad except Exception also masks genuine command-assembly bugs as -1 — narrow it or log unexpected types.)
| ) | ||
| super().__init__(*agent_args, instance_id=instance_id, **kwargs) | ||
|
|
||
| swebench.get_sb_environment = get_pyxis_environment |
There was a problem hiding this comment.
[Both] medium (design): _run_agent injects the Pyxis env by reassigning upstream module globals: swebench.get_sb_environment = get_pyxis_environment / swebench.ProgressTrackingAgent = LiveTrajectoryAgent. It works because the worker is a one-shot subprocess, but it's a fragile abstraction leak / hidden global mutable state: any upstream rename or pre-capture of these names breaks silently, with no assertion the attributes existed before overwrite. Add a guard (assert hasattr(swebench, "get_sb_environment")) so the coupling fails loudly on upstream drift, or isolate it behind a documented shim.
| secret_values=secret_values, | ||
| cancel_token=cancel_token, | ||
| ) | ||
| result_path = ( |
There was a problem hiding this comment.
[Claude] medium (bug): The Pyxis eval result-path resolution is stricter than Docker's and can spuriously fail a valid run. Docker _run_eval (619-629) falls back to rglob(f"*{run_id}*.json") when the exact {safe_model}.{run_id}.json isn't at top level; the Pyxis override only checks the exact path and raises RunnerError("...result file not found...") otherwise. If upstream ever changes the report filename or writes into a subdir, Pyxis raises while Docker recovers. Share the resolution logic or document why Pyxis can assume the exact name.
| 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) |
There was a problem hiding this comment.
[Claude] low (security): The per-run staging tmp dir is made world-writable: self._tmp_dir.chmod(0o1777), then bind-mounted into the container as /tmp with --container-remap-root. On a shared SLURM node this dir (under the default tempfile root) is writable by every user on the node for the run's lifetime; the sticky bit stops deletion of others' files but not read/write of predictable filenames the eval steps use (e.g. swebench_patch.diff). Prefer a per-user path ($XDG_RUNTIME_DIR/$TMPDIR) with 0o700, or justify why 0o1777 is required for the remapped-root container to write.
| return command | ||
|
|
||
|
|
||
| def resolve_image(image_registry: str, instance_id: str) -> str: |
There was a problem hiding this comment.
[Claude] low (bug): resolve_image does not lowercase instance_id, while the Docker cleanup path does (runner.py:511, instance_id.lower()). Registry image path components are case-sensitive and SWE-bench image tags are conventionally lowercased, so any uppercase instance ID (or a future subset) yields sweb.eval.arm64.{instance_id}:v4.1.0-arm64 that won't match the pushed lowercase image → image-pull failure. Latent today (Verified/Lite IDs are lowercase) but inconsistent with Docker — normalize.
| ) from exc | ||
|
|
||
|
|
||
| def create_runner( |
There was a problem hiding this comment.
[Quality] low (code-quality): create_runner(runtime: str, ...) dispatches on a stringly-typed flag with a runtime raise for the illegal value, though __main__.py already constrains it via choices=("docker","pyxis"). A Literal["docker","pyxis"] (or small enum) on the parameter makes the illegal state unrepresentable at the type boundary and lets mypy catch a bad caller.
| secret_values: set[str], | ||
| cancel_token: CancellationToken | None = None, | ||
| ) -> Path: | ||
| run_id = f"endpoints_{uuid.uuid4().hex[:8]}" |
There was a problem hiding this comment.
[Quality] low (code-quality): PyxisSweBenchRunner._run_eval copy-pastes non-trivial logic from base SweBenchRunner._run_eval (582-589): the run_id = f"endpoints_{uuid.uuid4().hex[:8]}" + persist step and the identical {"verified": "princeton-nlp/SWE-bench_Verified", "lite": "princeton-nlp/SWE-bench_Lite"}.get(request.subset) map. DRY: factor run-id generation and the subset→dataset map into a shared helper/constant instead of duplicating across both overrides.
| raise RunnerError(f"SWE-bench result file not found for run_id={run_id}") | ||
| return result_path | ||
|
|
||
| def _cleanup_containers( |
There was a problem hiding this comment.
[Claude] low (concurrency): PyxisSweBenchRunner._cleanup_containers sweeps leaked agent containers by listing enroot on a single node (build_srun_command pins --nodelist=$SLURMD_NODENAME). Correct for the README's advertised one-node allocation, but there's no assertion or comment tying the single-node cleanup to that invariant. If the allocation is ever >1 node (agent workers srun --overlap'd onto siblings), named pyxis_mswe_* containers on other nodes leak silently. Add a comment (or guard) documenting the single-node cleanup invariant. Relatedly, the README (line 20) should state the service must run inside a live SLURM allocation (SLURM_JOB_ID/SLURMD_NODENAME set) and that running outside one yields degraded — not errored — results, given failures are swallowed.
Review Council — Multi-AI Code ReviewReviewed by: Claude + Code-Quality | Depth: thorough New Pyxis/SLURM SWE-bench subsystem. Subprocess calls use list-argv form with a sanitized env ( 🔴 Must Fix (high)
🟡 Should Fix (medium)
🔵 Consider (low)
Also raised, dropped as minor at this cap: magic image arch/version constants (
|
| environment_class: docker | ||
| pull_timeout: 3600 | ||
| container_timeout: 10h | ||
| container_timeout: 6h |
There was a problem hiding this comment.
why is the container_timeout lowered? I guess we don't really need 10h but just wondering if there's an issue that requires a lowered container timeout
|
@hvagadia looking through this code, for the pyxis environment and worker support, as we are adding support into mini-swe agent itself, I would think it would make sense to upstream support for pyxis into mini-swe agent and open a PR in their repo instead for the pyxis support. https://github.com/SWE-agent/mini-swe-agent/tree/main/src/minisweagent/environments seems like the correct place. This would be great for the wider audience at Nvidia as well since others at Nvidia would be able to reuse your components for the pyxis support unblocking others. Right now my biggest concern is that there is some monkey patched codes: swebench.get_sb_environment = get_pyxis_environment
swebench.ProgressTrackingAgent = LiveTrajectoryAgentit will cause this module to break if the API changes. If we want to adopt this into this repo, I think we must at the minimum do an assert to pin this version of mini-swe of this library support to ensure that users don't try to use other versions of mini-swe with this library e.g. ( Upstreaming this module would be the path of less maintenance. In the meantime we could sustain a fork while we wait for the upstream to accept this PR and only if the upstream maintainers do not accept, then I believe we should sustain this. I'll let @nv-alicheng to top up any thoughts on this as well. Let me know what you think about this @hvagadia. Btw, I'm only referring to the agent part, but the eval part of this code look justified and it does look like something that we should own. ( |
|
Agreed with @leopck - the monkeypatching seems a little hacky to me (I believe this is the Medium - 4 in my agent review post), which I think should be maybe bumped up. Is it possible to make some child class to override and use that instead, or add it directly to swebench / make a fork? |
4 participants
Summary
srunand Enroot on a retained one-node Slurm allocationsweb.eval.arm64.<instance_id>:v4.1.0-arm64conventionWhy
ARM64 Slurm nodes provide Pyxis and Enroot but may not provide Docker. This allows the existing external SWE-bench accuracy service to run agent and evaluator containers natively on those nodes without changing its HTTP API or the endpoint client YAML.
Validation
uv run pytest -q tests/unit/evaluation/swebench_service— 72 passedpre-commitchecks for all five changed files — passed