Add R2E-Gym patch execution reward

docs/r2e_patch_execution.md

@@ -0,0 +1,24 @@
+# R2E-Gym patch execution
+
+The default R2E-Gym workflow generates an issue and scores its text. It does not execute a patch. Set `R2E_PATCH_EXECUTION=1` to switch the same workflow to unified-diff generation and container test reward. The model must return a `git diff` style patch. Each turn runs tests before and after applying the patch in a fresh task checkout. Reward is 1 only when the baseline fails and the patched tests pass; failed execution is returned as tool feedback for a possible revision.
+
+The public dataset index contains `repo_name`, `commit_hash`, and `docker_image`, but no repository checkout or test command. Supply these explicitly before enabling patch mode:
+
+```sh
+export R2E_PATCH_EXECUTION=1
+export R2E_REPO_ROOT=/absolute/path/to/task-repositories
+export R2E_TEST_COMMAND_JSON='["pytest", "-q"]'
+export R2E_PATCH_TIMEOUT=300
+```
+
+`R2E_REPO_ROOT/<repo_name>` must be a local Git repository containing the requested commit. `R2E_TEST_COMMAND_JSON` is an argument array run inside the dataset's Docker image with the isolated checkout mounted at `/workspace`. Use an image and test command appropriate to your R2E-Gym task; the dataset's `expected_output_json` is not itself an executable test specification. Docker must be available on each rollout worker. The harness disables container networking and drops Linux capabilities. Never point patch mode at a privileged Docker daemon when processing untrusted model output.
+
+For R2E-Gym `.sif` task images containing `/testbed` and `/r2e_tests`, use Singularity:
+
+```sh
+export R2E_PATCH_BACKEND=singularity
+export R2E_SIF_ROOT=/absolute/path/to/sif-images
+export R2E_TEST_COMMAND_JSON='["bash", "/testbed/run_tests.sh"]'
+```
+
+The entrypoint resolves each image as `<repo_name>_<first-eight-commit-chars>.sif` under `R2E_SIF_ROOT`. Every rollout worker must have Singularity and access to these images.

env/r2e_patch_harness.py

@@ -0,0 +1,140 @@
+"""Execute a unified diff against an R2E-Gym repository in Docker."""
+from __future__ import annotations
+
+import subprocess
+import tempfile
+import shlex
+import shutil
+from pathlib import Path
+from typing import Any, Sequence
+
+
+class R2EPatchExecutionHarness:
+    requires_repo_path = True
+    def __init__(self, timeout: int = 300):
+        self.timeout = max(1, int(timeout))
+
+    def execute(
+        self,
+        patch: str,
+        workspace: str | Path,
+        test_command: Sequence[str],
+        docker_image: str,
+        commit_hash: str,
+    ) -> dict[str, Any]:
+        """Apply a patch to an isolated commit checkout, then run tests in Docker.
+
+        The caller must supply a trusted image and command. Model output is only
+        used as patch data; it never becomes a shell command.
+        """
+        source = Path(workspace).resolve()
+        if not source.is_dir() or not (source / ".git").exists():
+            raise ValueError(f"workspace is not a git checkout: {source}")
+        if not patch.strip() or not docker_image or not commit_hash or not test_command:
+            raise ValueError("patch, docker_image, commit_hash and test_command are required")
+        with tempfile.TemporaryDirectory(prefix="r2e-patch-") as td:
+            checkout = Path(td) / "repo"
+            clone = subprocess.run(["git", "clone", "--quiet", "--no-hardlinks", "--", str(source), str(checkout)],
+                                   capture_output=True, text=True, timeout=self.timeout)
+            if clone.returncode:
+                raise RuntimeError(f"cannot clone task repository: {clone.stderr}")
+            reset = subprocess.run(["git", "checkout", "--detach", commit_hash], cwd=checkout,
+                                   capture_output=True, text=True, timeout=self.timeout)
+            if reset.returncode:
+                raise ValueError(f"commit is unavailable in workspace: {commit_hash}")
+            patch_file = Path(td) / "candidate.patch"
+            patch_file.write_text(patch, encoding="utf-8")
+            command = ["docker", "run", "--rm", "--network", "none", "--cap-drop", "ALL",
+                       "-v", f"{checkout}:/workspace", "-w", "/workspace", docker_image,
+                       *test_command]
+            try:
+                baseline = subprocess.run(command, capture_output=True, text=True, timeout=self.timeout)
+            except subprocess.TimeoutExpired:
+                raise RuntimeError("R2E baseline test timed out")
+            check = subprocess.run(["git", "apply", "--check", "--", str(patch_file)], cwd=checkout,
+                                   capture_output=True, text=True, timeout=self.timeout)
+            if check.returncode:
+                return {"patch_applied": False, "tests_passed": False,
+                        "returncode": check.returncode, "stdout": check.stdout, "stderr": check.stderr}
+            apply = subprocess.run(["git", "apply", "--", str(patch_file)], cwd=checkout,
+                                   capture_output=True, text=True, timeout=self.timeout)
+            if apply.returncode:
+                return {"patch_applied": False, "tests_passed": False,
+                        "returncode": apply.returncode, "stdout": apply.stdout, "stderr": apply.stderr}
+            try:
+                result = subprocess.run(command, capture_output=True, text=True, timeout=self.timeout)
+            except subprocess.TimeoutExpired:
+                return {"patch_applied": True, "tests_passed": False, "returncode": None,
+                        "stdout": "", "stderr": "test timeout"}
+            return {"patch_applied": True, "baseline_failed": baseline.returncode != 0,
+                    "tests_passed": baseline.returncode != 0 and result.returncode == 0,
+                    "returncode": result.returncode, "stdout": result.stdout[-10000:],
+                    "stderr": result.stderr[-10000:]}
+
+
+class R2ESingularityPatchHarness:
+    """Execute patches in R2E-Gym SIF images containing /testbed."""
+
+    requires_repo_path = False
+
+    def __init__(self, timeout: int = 300):
+        self.timeout = max(1, int(timeout))
+
+    def execute(
+        self,
+        patch: str,
+        workspace: str | Path | None,
+        test_command: Sequence[str],
+        docker_image: str,
+        commit_hash: str,
+    ) -> dict[str, Any]:
+        del workspace, commit_hash
+        image = Path(docker_image).resolve()
+        if not image.is_file() or image.suffix != ".sif":
+            raise ValueError(f"R2E SIF image is unavailable: {image}")
+        if not patch.strip() or not test_command:
+            raise ValueError("patch and test_command are required")
+        if shutil.which("singularity") is None:
+            raise RuntimeError("Singularity is unavailable on this worker")
+        with tempfile.TemporaryDirectory(prefix="r2e-sif-") as td:
+            sandbox = Path(td) / "sandbox"
+            build = subprocess.run(
+                ["singularity", "build", "--sandbox", str(sandbox), str(image)],
+                capture_output=True, text=True, timeout=self.timeout,
+            )
+            if build.returncode:
+                raise RuntimeError(f"cannot build R2E sandbox: {build.stderr[-3000:]}")
+            testbed = sandbox / "testbed"
+            if not (testbed / ".git").exists():
+                raise ValueError("R2E image does not contain a /testbed git checkout")
+            patch_path = testbed / ".r2e_candidate.patch"
+            patch_path.write_text(patch, encoding="utf-8")
+            tests_link = testbed / "r2e_tests"
+            if not tests_link.exists() and (sandbox / "r2e_tests").exists():
+                tests_link.symlink_to("/r2e_tests")
+            prefix = ["singularity", "exec", "--writable", "--no-home", str(sandbox),
+                      "bash", "-lc"]
+            setup = "export GIT_CONFIG_GLOBAL=/tmp/r2e_gitconfig; git config --global --add safe.directory /testbed; cd /testbed && "
+            command = setup + shlex.join(test_command)
+            try:
+                baseline = subprocess.run(prefix + [command], capture_output=True,
+                                          text=True, timeout=self.timeout)
+            except subprocess.TimeoutExpired:
+                raise RuntimeError("R2E baseline test timed out")
+            for action in ("git apply --check .r2e_candidate.patch", "git apply .r2e_candidate.patch"):
+                result = subprocess.run(prefix + [setup + action], capture_output=True,
+                                        text=True, timeout=self.timeout)
+                if result.returncode:
+                    return {"patch_applied": False, "tests_passed": False,
+                            "returncode": result.returncode, "stdout": result.stdout[-10000:],
+                            "stderr": result.stderr[-10000:]}
+            try:
+                result = subprocess.run(prefix + [command], capture_output=True,
+                                        text=True, timeout=self.timeout)
+            except subprocess.TimeoutExpired:
+                return {"patch_applied": True, "tests_passed": False,
+                        "returncode": None, "stdout": "", "stderr": "test timeout"}
+            return {"patch_applied": True, "baseline_failed": baseline.returncode != 0,
+                    "tests_passed": baseline.returncode != 0 and result.returncode == 0,
+                    "returncode": result.returncode, "stdout": result.stdout[-10000:],
+                    "stderr": result.stderr[-10000:]}

examples/r2e_gym_async_rl.py

@@ -4,6 +4,7 @@
 
 import json
 import os
+import shutil
 import sys
 from collections import Counter
 from datetime import timedelta
@@ -20,6 +21,10 @@
 from RL_Framework import AsyncRLTrainer, parse_args_and_load_config
 from RL_Framework.engine.device_utils import distributed_backend, set_device
 from RL_Framework.env.r2e_gym_reward import r2e_gym_reward_fn
+from RL_Framework.env.r2e_patch_harness import (
+    R2EPatchExecutionHarness,
+    R2ESingularityPatchHarness,
+)
 from RL_Framework.workflow.r2e_gym import R2EGymWorkflow
 
 
@@ -83,6 +88,21 @@ def main():
         print(f"Scheduler: {config.heterogeneous_rollout.scheduling.scheduler_type}")
 
     dataset = load_dataset("json", data_files=data_path, split="train")
+    patch_mode = os.environ.get("R2E_PATCH_EXECUTION", "0") == "1"
+    patch_backend = os.environ.get("R2E_PATCH_BACKEND", "docker").lower()
+    if patch_mode:
+        if patch_backend not in {"docker", "singularity"}:
+            raise ValueError("R2E_PATCH_BACKEND must be docker or singularity")
+        if shutil.which(patch_backend) is None:
+            raise RuntimeError(f"R2E patch mode requires {patch_backend} on this worker")
+        repo_root = os.environ.get("R2E_REPO_ROOT")
+        sif_root = os.environ.get("R2E_SIF_ROOT")
+        test_command_json = os.environ.get("R2E_TEST_COMMAND_JSON")
+        if (patch_backend == "docker" and not repo_root) or (patch_backend == "singularity" and not sif_root) or not test_command_json:
+            raise ValueError("Patch mode requires R2E_TEST_COMMAND_JSON and the selected backend's repo/image root")
+        test_command = json.loads(test_command_json)
+        if not isinstance(test_command, list) or not test_command or not all(isinstance(part, str) for part in test_command):
+            raise ValueError("R2E_TEST_COMMAND_JSON must be a nonempty JSON string array")
 
     def preprocess(example):
         prompt = (example.get("prompt") or "").strip()
@@ -92,7 +112,7 @@ def preprocess(example):
             or ""
         ).strip()
         prompt_id = f"{example.get('repo_name', 'repo')}:{example.get('commit_hash', '')}"
-        return {
+        result = {
             "prompt_id": prompt_id,
             "prompt": prompt,
             "task_text": target_issue,
@@ -103,6 +123,18 @@ def preprocess(example):
             "expected_output_json": example.get("expected_output_json", "{}"),
             "modified_files": _normalize_modified_files(example.get("modified_files")),
         }
+        if patch_mode:
+            repo_name = example.get("repo_name", "")
+            if not repo_name or os.path.basename(repo_name) != repo_name:
+                raise ValueError(f"Invalid repo_name: {repo_name!r}")
+            if patch_backend == "docker":
+                result["repo_path"] = os.path.join(repo_root, repo_name)
+            else:
+                result["docker_image"] = os.path.join(
+                    sif_root, f"{repo_name}_{str(example.get('commit_hash', ''))[:8]}.sif"
+                )
+            result["test_command"] = test_command
+        return result
 
     dataset = dataset.map(preprocess)
     if is_main_process:
@@ -134,6 +166,9 @@ def preprocess(example):
         temperature=config.temperature,
         top_p=config.top_p,
         n_samples=config.n_samples,
+        patch_harness=(R2ESingularityPatchHarness if patch_backend == "singularity" else R2EPatchExecutionHarness)(
+            timeout=int(os.environ.get("R2E_PATCH_TIMEOUT", "300"))
+        ) if patch_mode else None,
     )
 
     trainer = AsyncRLTrainer(config)

scripts/test_cpu_ci.sh

@@ -11,4 +11,6 @@ python -m pytest -q --strict-markers --junitxml=reports/cpu-tests.xml \
   tests/test_cmlfq_cost_scheduler.py \
   tests/test_hetero_cmlfq_integration.py \
   tests/test_rollout_engine.py \
-  tests/test_cpu_offload_backend.py
+  tests/test_cpu_offload_backend.py \
+  tests/test_r2e_patch_harness.py \
+  tests/test_r2e_gym_workflow.py

tests/test_r2e_patch_harness.py

@@ -0,0 +1,34 @@
+import subprocess
+from pathlib import Path
+
+from env.r2e_patch_harness import R2EPatchExecutionHarness
+
+
+def test_patch_harness_uses_isolated_checkout_and_docker(tmp_path: Path, monkeypatch):
+    repo = tmp_path / "source"
+    repo.mkdir()
+    subprocess.run(["git", "init", "-q"], cwd=repo, check=True)
+    (repo / "value.py").write_text("VALUE = 1\n")
+    subprocess.run(["git", "add", "."], cwd=repo, check=True)
+    subprocess.run(["git", "-c", "user.name=t", "-c", "user.email=t@t", "commit", "-qm", "base"], cwd=repo, check=True)
+    commit = subprocess.check_output(["git", "rev-parse", "HEAD"], cwd=repo, text=True).strip()
+    (repo / "value.py").write_text("VALUE = 2\n")
+    patch = subprocess.check_output(["git", "diff", "--", "value.py"], cwd=repo, text=True)
+    subprocess.run(["git", "checkout", "--", "value.py"], cwd=repo, check=True)
+    real_run = subprocess.run
+    calls = 0
+
+    def run(command, **kwargs):
+        nonlocal calls
+        if command[0] == "docker":
+            calls += 1
+            mount = command[command.index("-v") + 1].split(":/workspace")[0]
+            expected = "VALUE = 1\n" if calls == 1 else "VALUE = 2\n"
+            assert (Path(mount) / "value.py").read_text() == expected
+            return subprocess.CompletedProcess(command, 1 if calls == 1 else 0, "", "")
+        return real_run(command, **kwargs)
+
+    monkeypatch.setattr(subprocess, "run", run)
+    result = R2EPatchExecutionHarness().execute(patch, repo, ["pytest", "-q"], "example:test", commit)
+    assert result["patch_applied"] and result["baseline_failed"] and result["tests_passed"]
+    assert (repo / "value.py").read_text() == "VALUE = 1\n"