Make docker_build_env.sh editable, add payload override and build stage detection

src/datasmith/agents/sandbox.py

@@ -11,6 +11,7 @@
 
 import hashlib
 import json
+import os
 import shutil
 import subprocess
 import sys
@@ -35,7 +36,6 @@
 _IMMUTABLE_FILES = (
     "Dockerfile.pr",
     "docker_build_base.sh",
-    "docker_build_env.sh",
     "docker_build_final.sh",
     "profile.sh",
     "run-tests.sh",
@@ -54,14 +54,34 @@ def _compute_immutable_hashes(task_dir: Path) -> dict[str, str]:
     return hashes
 
 
+def _read_env_payload_override(task_dir: Path) -> str | None:
+    """Read and validate ``env_payload_override.json`` from *task_dir*.
+
+    Returns the raw JSON string if the file exists and contains a valid
+    JSON list, otherwise ``None``.
+    """
+    override_file = task_dir / "env_payload_override.json"
+    if not override_file.exists():
+        return None
+    try:
+        raw = override_file.read_text()
+        parsed = json.loads(raw)
+        if isinstance(parsed, list):
+            return raw
+        logger.warning("env_payload_override.json is not a JSON list, ignoring")
+    except (json.JSONDecodeError, Exception):
+        logger.warning("Failed to parse env_payload_override.json, ignoring")
+    return None
+
+
 @dataclass
 class SandboxConfig:
     """Configuration for the Codex sandbox runner."""
 
-    timeout_s: int = 3600
+    timeout_s: int = int(os.environ.get("SYNTHESIS_TIMEOUT_S", "14400"))
     """Total wall-clock timeout for the codex session (seconds)."""
 
-    codex_timeout_s: int = 3600
+    codex_timeout_s: int = int(os.environ.get("SYNTHESIS_TIMEOUT_S", "14400"))
     """Timeout passed to subprocess.run for the codex process (seconds)."""
 
 
@@ -78,6 +98,7 @@ class SandboxResult:
     agent_name: str = ""
     files_changed: list[str] = field(default_factory=list)
     resource_metrics: dict = field(default_factory=dict)
+    env_payload_override: str | None = None
 
 
 class SandboxRunner:
@@ -296,7 +317,7 @@ def _extract_results(self, workspace: Path, codex_result: AgentResult, agent_nam
 
         success = success_file.exists()
 
-        # Read back only the two agent-editable scripts (the rest are templates)
+        # Read back the agent-editable scripts (the rest are templates)
         docker_context: DockerContext | None = None
         try:
             pkg_sh = (
@@ -305,10 +326,16 @@ def _extract_results(self, workspace: Path, codex_result: AgentResult, agent_nam
             run_sh = (
                 (task_dir / "docker_build_run.sh").read_text() if (task_dir / "docker_build_run.sh").exists() else ""
             )
-            docker_context = DockerContext(build_pkg_sh=pkg_sh, build_run_sh=run_sh)
+            env_sh = (
+                (task_dir / "docker_build_env.sh").read_text() if (task_dir / "docker_build_env.sh").exists() else ""
+            )
+            docker_context = DockerContext(build_pkg_sh=pkg_sh, build_run_sh=run_sh, build_env_sh=env_sh)
         except Exception:
             logger.warning("Failed to read Docker context from workspace")
 
+        # Read env_payload override if the agent wrote one
+        env_payload_override = _read_env_payload_override(task_dir)
+
         # Read failure.json if present
         failure_json: dict | None = None
         if failure_file.exists():
@@ -347,6 +374,7 @@ def _extract_results(self, workspace: Path, codex_result: AgentResult, agent_nam
             agent_name=agent_name,
             files_changed=codex_result.files_changed,
             resource_metrics=resource_metrics,
+            env_payload_override=env_payload_override if success else None,
         )
 
 

src/datasmith/agents/synthesizer.py

@@ -343,11 +343,13 @@ def _save_context(
         issue_number: int,
         ctx: DockerContext,
         resource_metrics: dict | None = None,
+        env_payload_override: str | None = None,
     ) -> None:
         """Persist the agent-edited scripts to the ``candidate_containers`` table.
 
-        Only ``build_pkg_sh`` and ``build_run_sh`` are saved — the other
-        fields come from templates and don't need to be persisted.
+        Saves ``build_pkg_sh``, ``build_run_sh``, and ``build_env_sh``.
+        When the agent also modified the env payload, ``env_payload_override``
+        is persisted to the ``env_payload`` column.
         """
         if not sha:
             return
@@ -360,9 +362,12 @@ def _save_context(
                 "issue_number": issue_number,
                 "build_pkg_sh": ctx.build_pkg_sh,
                 "build_run_sh": ctx.build_run_sh,
+                "build_env_sh": ctx.build_env_sh,
             }
             if resource_metrics:
                 row["resource_metrics"] = resource_metrics
+            if env_payload_override:
+                row["env_payload"] = env_payload_override
             client.table("candidate_containers").upsert(row).execute()
             logger.info("Saved context for %s/%s@%s", owner, repo, sha[:12])
         except Exception:

src/datasmith/agents/templates/AGENTS.md.j2

@@ -35,24 +35,31 @@ Follow this iterative cycle:
 
 1. Run `python3 sandbox_verify.py` (A good timeout is 60 minutes, but adjust as needed)
 2. If it fails, read `task/failure.json` for error details
-3. Edit `task/docker_build_pkg.sh` and/or `task/docker_build_run.sh`
+3. Edit `task/docker_build_pkg.sh`, `task/docker_build_run.sh`, and/or `task/docker_build_env.sh`
 4. Re-run `python3 sandbox_verify.py`
 5. Repeat until `task/verification_success.json` is created
 
-**IMPORTANT**: Only modify `task/docker_build_pkg.sh` and `task/docker_build_run.sh`.
+**IMPORTANT**: Only modify `task/docker_build_pkg.sh`, `task/docker_build_run.sh`, and `task/docker_build_env.sh`.
 Do **not** modify the Dockerfile, `task.txt`, or any other scripts.
 
+You may also write `task/env_payload_override.json` — a JSON list of pinned package strings
+(e.g. `["numpy==1.21.0", "scipy==1.7.0"]`) — to replace the default env payload when the
+original package versions fail to build.
+
 ## File Constraints
 
 ### Files you CAN edit
 - `task/docker_build_pkg.sh` — Primary build script. Install the package, add build deps.
 - `task/docker_build_run.sh` — Runtime setup. Add test/benchmark deps, repo-specific config.
+- `task/docker_build_env.sh` — Python environment setup. Add build prerequisites needed before env payload installation (e.g. Cython, distutils).
+
+### Files you CAN create
+- `task/env_payload_override.json` — JSON list of pinned package strings to replace the default env payload. Use this when the original package versions are fundamentally broken (e.g. no wheels available, incompatible with the Python version).
 
 ### Files you MUST NOT edit
 - `task/Dockerfile.pr` — Multi-stage build definition (base and repo stages are pre-built)
 - `task/task.txt` — Task configuration (owner, repo, sha, deps)
 - `task/docker_build_base.sh` — Base system setup (template)
-- `task/docker_build_env.sh` — Python environment creation (template)
 - `task/docker_build_final.sh` — Final image setup (template)
 - `task/profile.sh` — ASV benchmark runner
 - `task/run-tests.sh` — Test runner
@@ -64,16 +71,19 @@ When verification fails, `task/failure.json` contains:
 
 ```json
 {
-  "stage": "build|tests",
+  "stage": "env|pkg|run|tests|build",
   "return_code": 1,
   "stderr": "Full error output from the failed stage",
   "stdout": "Full standard output from the failed stage",
   "error_message": "Human-readable error description"
 }
 ```
 
-- **build** failures: Usually missing dependencies or compilation errors → fix in `docker_build_pkg.sh`
+- **env** failures: Env payload packages fail to install (missing build tools, broken wheels) → fix in `docker_build_env.sh` or write `env_payload_override.json`
+- **pkg** failures: Package installation errors or missing build deps → fix in `docker_build_pkg.sh`
+- **run** failures: Runtime setup issues → fix in `docker_build_run.sh`
 - **tests** failures: Missing test deps, import errors, or benchmark setup issues → fix in `docker_build_run.sh` or `docker_build_pkg.sh`
+- **build** (generic): Stage could not be determined from logs — check stdout/stderr for details
 
 ## Common Fixes for docker_build_pkg.sh
 
@@ -132,13 +142,15 @@ These are set by earlier build stages and available in both scripts:
 ## Docker Build Stages
 
 The Dockerfile.pr has 4 stages: env → pkg → run → final (base and repo images are pre-built).
-Only the `pkg` and `run` stages use your editable scripts.
+The `env`, `pkg`, and `run` stages use your editable scripts.
 
+- **env** stage: Runs `docker_build_env.sh` to set up the Python environment and install pinned deps
 - **pkg** stage: Runs `docker_build_pkg.sh` to install the package
 - **run** stage: Runs `docker_build_run.sh` to prepare for tests/benchmarks
 
+Changes to `docker_build_env.sh` rebuild from the env stage onward (slowest).
 Changes to `docker_build_pkg.sh` rebuild from the pkg stage onward (fast iteration).
-Changes to `docker_build_run.sh` rebuild only the run stage.
+Changes to `docker_build_run.sh` rebuild only the run stage (fastest).
 
 ## Tips
 
@@ -148,3 +160,6 @@ Changes to `docker_build_run.sh` rebuild only the run stage.
 - If the build times out, look for ways to simplify or skip expensive steps
 - Docker layer caching means earlier stages are cached — only your changed stage rebuilds
 - A lot of bookkeeping code is used in downstream stages; avoid removing code unless you understand its purpose.
+- **Fixing env payload packages**: If packages in the env payload fail to build (e.g. no wheels, missing build deps), you have two options:
+  1. **Edit `docker_build_env.sh`** to install build prerequisites (Cython, compiler toolchains, etc.) *before* the payload is installed. This is preferred when the package versions are correct but just need build tools.
+  2. **Write `task/env_payload_override.json`** with a corrected JSON list of pinned packages (e.g. `["numpy==1.21.0", "h5py==3.1.0"]`). `sandbox_verify.py` will use this instead of the original payload. Use this when a package version is fundamentally broken (no source dist, incompatible with the Python version, yanked). Keep changes minimal — stay as close as possible to the originals.

src/datasmith/agents/templates/sandbox_verify.py

@@ -34,7 +34,6 @@
 _IMMUTABLE_FILES = (
     "Dockerfile.pr",
     "docker_build_base.sh",
-    "docker_build_env.sh",
     "docker_build_final.sh",
     "profile.sh",
     "run-tests.sh",
@@ -173,6 +172,31 @@ def __init__(self, message: str, stdout: str, stderr: str, rc: int) -> None:
         self.rc = rc
 
 
+def _parse_failed_stage(build_log: str) -> str:
+    """Identify which Dockerfile stage (env/pkg/run) failed from Docker build output.
+
+    Scans for BuildKit stage markers like ``[env 2/2]`` or legacy markers like
+    ``FROM env AS pkg``.  Returns the name of the last stage seen before the
+    error, or ``"build"`` as a fallback.
+    """
+    # BuildKit format: #N [stage_name step/total] ...
+    buildkit_re = re.compile(r"\[(\w+)\s+\d+/\d+\]")
+    # Legacy format: Step N/M : FROM x AS stage
+    legacy_re = re.compile(r"Step \d+/\d+\s*:\s*FROM\s+\S+\s+AS\s+(\w+)", re.IGNORECASE)
+
+    last_stage = ""
+    for line in build_log.splitlines():
+        m = buildkit_re.search(line)
+        if m:
+            last_stage = m.group(1)
+            continue
+        m = legacy_re.search(line)
+        if m:
+            last_stage = m.group(1)
+
+    return last_stage if last_stage else "build"
+
+
 _MEM_UNITS = {"B": 1, "KIB": 1024, "MIB": 1024**2, "GIB": 1024**3, "TIB": 1024**4}
 _MEM_RE = re.compile(r"([\d.]+)\s*((?:[KMGT]i)?B)", re.IGNORECASE)
 
@@ -380,13 +404,38 @@ def verify(task_dir: Path) -> bool:
         _write_failure(task_dir, "parse", stderr="Task.sha is None", metrics=metrics)
         return False
 
+    # Check for env_payload override written by the agent
+    override_file = task_dir / "env_payload_override.json"
+    if override_file.exists():
+        try:
+            raw = override_file.read_text()
+            parsed = json.loads(raw)
+            if isinstance(parsed, list):
+                print(f"Using env_payload_override.json ({len(parsed)} packages)")
+                task = Task(
+                    owner=task.owner,
+                    repo=task.repo,
+                    sha=task.sha,
+                    commit_date=task.commit_date,
+                    env_payload=json.dumps(parsed),
+                    python_version=task.python_version,
+                    tag=task.tag,
+                    benchmarks=task.benchmarks,
+                    repo_image=task.repo_image,
+                )
+            else:
+                print("WARNING: env_payload_override.json is not a JSON list, ignoring")
+        except (json.JSONDecodeError, Exception) as e:
+            print(f"WARNING: Failed to parse env_payload_override.json: {e}")
+
     # Build
     try:
         tag = build_image(docker, task_dir, task, target="run", metrics=metrics)
         print(f"Build succeeded: {tag}")
     except BuildError as e:
-        print(f"Build failed: {e}")
-        _write_failure(task_dir, "build", stdout=e.stdout, stderr=e.stderr, rc=e.rc, metrics=metrics)
+        stage = _parse_failed_stage(e.stdout)
+        print(f"Build failed at stage '{stage}': {e}")
+        _write_failure(task_dir, stage, stdout=e.stdout, stderr=e.stderr, rc=e.rc, metrics=metrics)
         return False
     except Exception as e:
         print(f"Build failed: {str(e)[:200]}")