diff --git a/Makefile b/Makefile index 3783387c..86cd9daf 100755 --- a/Makefile +++ b/Makefile @@ -6,7 +6,7 @@ SHELL := /bin/bash -.PHONY: help docker-shell docker-check-agents docker-smoke docker-run docker-parallel-run docker-setup-flydsl \ +.PHONY: help docker-shell docker-check-agents docker-smoke docker-run docker-parallel-run docker-setup-flydsl docker-setup-geak \ check-docker-runner check-evaluator check-held-out check-visualization \ visualization-build visualization-serve visualization-run \ sync-perf-helpers check-perf-helpers materialize-perf-workspace \ @@ -17,7 +17,7 @@ help: @echo "======================================================" @echo "Docker-first workflow (the only supported path):" @echo "make docker-shell - Enter the runtime image with repo and agent auth mounted" - @echo "make docker-check-agents - Verify the first-class host CLI selected by CONFIG" + @echo "make docker-check-agents - Verify the agent stack selected by CONFIG" @echo " Use CONFIG=... for another config; AGENTS=... overrides it" @echo " AGENTS=all explicitly checks all three first-class CLIs" @echo "make docker-smoke - Verify Docker Python, ROCm tools, imports, and GPU access" @@ -27,6 +27,7 @@ help: @echo " On other GPUs, pass a matching CONFIG explicitly" @echo " Images: gfx942->mi30x, gfx950->mi35x; override with AKA_DOCKER_IMAGE=..." @echo "make docker-setup-flydsl - Install FlyDSL when absent (for flydsl2flydsl, torch2flydsl, and triton2flydsl)" + @echo "make docker-setup-geak - Install the GEAK v4 Claude Agent SDK dependency" @echo "make check-docker-runner - Check Docker runner syntax and runtime-specific arguments" @echo "make check-evaluator - Run centralized evaluator unit tests" @echo "make check-held-out - Run held-out module unit tests" @@ -55,7 +56,11 @@ VISUALIZATION_PORT ?= 8080 MATERIALIZE_FORCE_ARG := $(if $(filter 1 true yes,$(FORCE)),--force,) docker-shell: - @$(DOCKER_RUNNER) shell + @if [[ -n "$(AGENTS)" ]]; then \ + AKA_AGENTS="$(AGENTS)" $(DOCKER_RUNNER) shell; \ + else \ + $(DOCKER_RUNNER) shell; \ + fi docker-check-agents: @AKA_AGENTS="$(AGENTS)" $(DOCKER_RUNNER) check-agents --config_name $(CONFIG) @@ -69,11 +74,16 @@ docker-run: docker-parallel-run: @GPU_IDS="$(GPU_IDS)" $(DOCKER_RUNNER) parallel-run --config_name $(CONFIG) $(RUN_ARGS) -# Install FlyDSL into the container's persistent pip user-base when the selected +# Install FlyDSL into the container's persistent Python dependency target when the selected # image does not ship it. Needed by all three FlyDSL task types. docker-setup-flydsl: @$(DOCKER_RUNNER) setup-flydsl +# Install the Claude Agent SDK into the persistent container Python dependency target. +# The GEAK Workflow checkout itself is bind-mounted read-only from AKA_GEAK_ROOT. +docker-setup-geak: + @$(DOCKER_RUNNER) setup-geak + check-docker-runner: @bash tests/test_docker_benchmark.sh diff --git a/README.md b/README.md index b15d264b..a77e4bb8 100755 --- a/README.md +++ b/README.md @@ -74,6 +74,7 @@ AgentKernelArena/ │ ├── cursor/ # Cursor Agent CLI │ ├── claude_code/ # Claude Code CLI │ ├── codex/ # Codex CLI +│ ├── geak_v4/ # GEAK v4 deterministic Workflow │ ├── geak_v3/ # GEAK HIP optimization │ ├── geak_v3_triton/ # GEAK Triton optimization │ ├── mini_swe_triton/ # mini-swe-agent Triton optimization @@ -114,12 +115,13 @@ Each run selects one `agent.template`. Repeated runs can compare different agent | `cursor` | Cursor Agent CLI integration | | `claude_code` | Claude Code CLI integration | | `codex` | Codex CLI integration | +| `geak_v4` | GEAK v4 deterministic kernel Workflow through Claude Code | | `geak_v3` | GEAK optimization for HIP tasks | | `geak_v3_triton` | GEAK optimization for Triton tasks | | `mini_swe_triton` | mini-swe-agent-based Triton optimization | | `task_validator` | Task quality validation; does not optimize kernels | -Agent-specific models, effort settings, iteration guidance, timeouts, and provider configuration live under `agents//agent_config.yaml` or in the selected agent CLI. Specialized agents may require additional setup; inspect their directories and agent-specific README files where present. +Agent-specific models, effort settings, iteration guidance, timeouts, and provider configuration live under `agents//agent_config.yaml` or in the selected agent CLI. Specialized agents may require additional setup; inspect their directories and agent-specific README files where present. GEAK v4 setup and its single-source task contract are documented in [agents/geak_v4/README.md](agents/geak_v4/README.md). ## Task Environments @@ -188,12 +190,13 @@ installation. The npm path requires Node.js 22+ and npm. See the [official Claude Code setup guide](https://code.claude.com/docs/en/installation) for the current alternatives. -The repository provides three ready-to-use run configurations: +The repository provides four ready-to-use run configurations: | Configuration | Purpose | | --- | --- | | `example_configs/quickstart_claude_mi300.yaml` | One Claude Code GELU task on MI300/MI300X (`gfx942`); use this for a first run on MI300-series hardware. | | `example_configs/quickstart_claude_mi355x.yaml` | One Claude Code GELU task on MI355X (`gfx950`); use this for a first run on MI355X. | +| `example_configs/quickstart_geak_v4_mi300.yaml` | One GEAK v4 GELU task on MI300/MI300X (`gfx942`); requires the additional GEAK checkout and SDK setup. | | `example_configs/benchmark_cursor_mi355x.yaml` | Curated 60-task Cursor Agent benchmark on MI355X; use this for a longer benchmark only after installing and authenticating Cursor Agent. | Running `make docker-run` without `CONFIG` uses the MI300/MI300X Claude @@ -206,6 +209,18 @@ FlyDSL tasks require FlyDSL in the container. The pinned image may already provi make docker-setup-flydsl ``` +GEAK v4 requires a GEAK checkout beside AgentKernelArena by default, Claude +Code 2.1.177 or newer logged in on the host, and the Agent SDK: + +```bash +make docker-setup-geak +``` + +Set `AKA_GEAK_ROOT=/absolute/path/to/GEAK` when the checkout is not the default +sibling directory. This target installs only `claude-agent-sdk`; it does not +`pip install` GEAK, which is mounted read-only. See the +[GEAK v4 agent README](agents/geak_v4/README.md) before selecting its example. + Performance timing helpers are maintained in `src/tools/perf/` and materialized into run workspaces. See [src/tools/perf/README.md](src/tools/perf/README.md) before changing task timing code. For detailed installation and compatibility information, see [docs/install/install.md](docs/install/install.md) and [docs/reference/compatibility-matrix.md](docs/reference/compatibility-matrix.md). @@ -227,18 +242,19 @@ Run agent-specific settings such as `model`, `effort`, `max_iterations`, and `timeout_seconds` are configured in the selected agent's `agent_config.yaml`, not in the run configuration. -For a Cursor, Claude Code, Codex, or task-validator config, verify only the -selected first-class host CLI (the validator resolves to its configured backend): +For a Cursor, Claude Code, Codex, GEAK v4, or task-validator config, verify only +the dependencies selected by the config (the validator resolves to its +configured backend): ```bash CONFIG_PATH=my_experiment.yaml make docker-check-agents CONFIG="$CONFIG_PATH" ``` -Use `AGENTS=claude_code,codex` to check an explicit subset or `AGENTS=all` to -check Cursor, Claude Code, and Codex together. Specialized integrations such as -GEAK and mini-swe have their own dependency checks and are not handled by this -command. +Use `AGENTS=claude_code,codex` to check an explicit subset, +`AGENTS=geak_v4` to check the GEAK v4 stack, or `AGENTS=all` to check Cursor, +Claude Code, and Codex together. Legacy GEAK and mini-swe integrations retain +their own dependency checks. ### Run Serially @@ -258,9 +274,10 @@ make docker-parallel-run CONFIG="$CONFIG_PATH" ``` The Docker parallel path is verified for `cursor`, `claude_code`, `codex`, and -`task_validator`. Specialized GEAK/mini-swe integrations need their own -dependencies and GPU-ID configuration before they are used with isolated -workers. +`task_validator`. `geak_v4` maps each isolated worker to logical GPU 0, but +still requires the setup in its agent README. A paid Claude workflow invocation +is not part of the integration's offline validation. Legacy GEAK/mini-swe +integrations need their own dependencies and GPU-ID configuration. ### Resume a Run diff --git a/agents/geak_v4/README.md b/agents/geak_v4/README.md new file mode 100644 index 00000000..e0958349 --- /dev/null +++ b/agents/geak_v4/README.md @@ -0,0 +1,127 @@ +# GEAK v4 Agent + +The `geak_v4` integration runs GEAK's deterministic +`kernel_workflow/kernel_workflow.js` through Claude Code's dynamic Workflow +tool. GEAK works on a disposable copy of the task workspace; AgentKernelArena +imports only a validated patch for the one declared kernel source and then runs +its normal correctness and performance evaluation. + +## Prerequisites + +- An AMD Instinct GPU and `rocprof-compute`, as required by the Arena Docker + smoke/run contract. +- A local GEAK checkout. By default, the Docker runner looks for `GEAK` beside + the Arena checkout: + + ```text + parent/ + ├── AgentKernelArena/ + └── GEAK/ + ``` + +- Claude Code 2.1.177 or newer, installed and logged in on the host. The minimum + version is required for the dynamic Workflow feature. +- Access to the Claude model configured in + `agents/geak_v4/agent_config.yaml`. +- FlyDSL installed with `make docker-setup-flydsl` when selecting a + `flydsl2flydsl` task. + +## Setup + +Clone GEAK beside AgentKernelArena, authenticate Claude Code on the host, and +install the Python SDK into the persistent container dependency directory: + +```bash +cd /path/to/parent +git clone https://github.com/AMD-AGI/GEAK.git +cd AgentKernelArena + +claude --version +claude +claude auth status + +make docker-setup-geak +``` + +This adapter was developed against GEAK commit +`4965d5b2ccde927925c8c5501a25c1233daa52eb` +(`v4.0.0-102-g4965d5b`). For a reproducible review, check out that revision; +newer GEAK revisions may require an adapter/schema update. + +If the GEAK checkout is elsewhere, provide its absolute host path for setup, +preflight, and every run: + +```bash +export AKA_GEAK_ROOT=/absolute/path/to/GEAK +make docker-setup-geak +``` + +`make docker-setup-geak` installs only `claude-agent-sdk`. It does **not** +`pip install` GEAK. The Docker runner bind-mounts the checkout read-only at +`/opt/geak`, so the workflow code, roles, and knowledge remain unchanged by an +Arena run. The persistent Python dependency directory is writable only in the +explicit setup container and is nested-mounted read-only during agent runs. + +## Run the example + +The included MI300/MI300X example runs one HIP GELU task: + +```bash +CONFIG_PATH=example_configs/quickstart_geak_v4_mi300.yaml +make docker-check-agents CONFIG="$CONFIG_PATH" +make docker-run CONFIG="$CONFIG_PATH" +``` + +`docker-check-agents` verifies the Claude login and version, the Agent SDK, the +GEAK workflow checkout, and an available profiler without starting an +optimization. + +## V1 supported task contract + +V1 of the Arena integration supports: + +- `hip2hip` +- `triton2triton` +- `flydsl2flydsl` + +A task must declare exactly one existing source in `source_file_path`. That +source must be a normal, non-symlink file and cannot be a protected config, +test, or harness path. In particular, names such as `test_*.py`, `*_test.py`, +`test_*.cpp`, `*_test.hip`, and `*_harness.cu` are rejected. Tasks with +multiple sources or a source that also contains the test harness are +intentionally out of scope. + +Authoring, translation, repository, and image-level tasks are not supported by +V1. Use a supported task with a single standalone kernel source. + +## Isolation and artifacts + +For a task workspace named ``, GEAK artifacts are kept outside the +scored workspace under: + +```text +._geak_v4// +├── input/ # Disposable task copy seen by GEAK +├── eval/ # GEAK validation and final patch artifacts +├── runs/ # Workflow experiment artifacts +├── handoff.json +└── result.json +``` + +The hidden artifact directory is a sibling of the task workspace. GEAK is +instructed not to apply changes to the original input, and the Arena accepts a +result only when GEAK's Director validation passes and the final patch changes +the one allowlisted source without creating, deleting, renaming, or changing +the mode of a file. Arena also compares a full manifest of the scored workspace +before and after Workflow execution; a direct mutation makes the task fail +before scoring or patch import. + +## Validation boundary + +Offline tests cover handoff mapping, result parsing, patch filtering, artifact +isolation, and Docker argument construction without contacting Claude. A real +paid Claude workflow invocation is not part of that offline validation; run +the one-task example with an authorized account before relying on this +integration for a benchmark campaign. As with other Arena task workspaces, the +disposable-copy and manifest checks are fail-closed integrity controls, not an +OS security sandbox; a failed run is not automatically rolled back. diff --git a/agents/geak_v4/__init__.py b/agents/geak_v4/__init__.py new file mode 100644 index 00000000..c54dbfb2 --- /dev/null +++ b/agents/geak_v4/__init__.py @@ -0,0 +1,4 @@ +# Copyright(C) [2026] Advanced Micro Devices, Inc. All rights reserved. +from .launch_agent import launch_agent + +__all__ = ["launch_agent"] diff --git a/agents/geak_v4/agent_config.yaml b/agents/geak_v4/agent_config.yaml new file mode 100644 index 00000000..4fe19699 --- /dev/null +++ b/agents/geak_v4/agent_config.yaml @@ -0,0 +1,29 @@ +# GEAK v4 runs the deterministic kernel_workflow through Claude Code's dynamic +# Workflow tool. The GEAK checkout is mounted read-only at /opt/geak; all run +# artifacts are written beside the Arena task workspace. +workflow_dir: /opt/geak/kernel_workflow + +# Defaults from GEAK v4's supported standalone kernel workflow. +model: claude-opus-4-8 +effort: ultracode +budget: 6 +min_improve: 0.02 +deep_cost: 2 +use_expert_skills: false + +# Hard wall-clock bound for the complete multi-agent workflow. +timeout_seconds: 43200 + +# Newer Claude Code builds may finish detached validation work after the +# background task notification. Keep the SDK client alive for this bounded +# on-disk completion grace period. +done_grace_seconds: 1800 +done_poll_seconds: 5 + +# V1 deliberately targets existing, isolated kernels with one declared source +# file. Author/translation tasks need a separate frozen-baseline adapter, while +# repository/image tasks need a storage/copy-cost qualification first. +supported_task_types: + - hip2hip + - triton2triton + - flydsl2flydsl diff --git a/agents/geak_v4/launch_agent.py b/agents/geak_v4/launch_agent.py new file mode 100644 index 00000000..fe738dcb --- /dev/null +++ b/agents/geak_v4/launch_agent.py @@ -0,0 +1,927 @@ +# Copyright(C) [2026] Advanced Micro Devices, Inc. All rights reserved. +"""AgentKernelArena adapter for GEAK v4's deterministic kernel workflow.""" + +from __future__ import annotations + +import hashlib +import json +import logging +import math +import os +import re +import shutil +import signal +import stat +import subprocess +import sys +import tempfile +import threading +import time +import uuid +from datetime import datetime, timezone +from pathlib import Path, PurePosixPath +from typing import Any + +import yaml + +from agents import register_agent +from src.harness_guard import ( + is_protected_workspace_path, + snapshot_workspace_harness, + verify_workspace_harness, +) +from src.module_registration import AgentType, load_prompt_builder + + +_COPY_IGNORED_DIRS = { + ".git", + ".rocprofv3", + ".torch_ext", + "__pycache__", + "build", +} +_COPY_IGNORED_SUFFIXES = {".o", ".pyc", ".so"} +_PATCH_SIZE_LIMIT = 16 * 1024 * 1024 +_JSON_SIZE_LIMIT = 8 * 1024 * 1024 +_PROCESS_OUTPUT_LIMIT = 4 * 1024 * 1024 +_UNSAFE_SOURCE_NAME = re.compile( + r"(?:test_.*|.*_(?:test|harness))\.(?:py|c|cc|cpp|cxx|cu|hip)", + re.IGNORECASE, +) + + +def _load_agent_config() -> dict[str, Any]: + path = Path(__file__).with_name("agent_config.yaml") + with path.open("r", encoding="utf-8") as stream: + return yaml.safe_load(stream) or {} + + +def _load_task_config(task_config_dir: str) -> dict[str, Any]: + with Path(task_config_dir).open("r", encoding="utf-8") as stream: + return yaml.safe_load(stream) or {} + + +def _single_declared_source( + task_config: dict[str, Any], + workspace: Path, +) -> PurePosixPath: + raw = task_config.get("source_file_path") + values = [raw] if isinstance(raw, str) else raw + if not isinstance(values, list) or len(values) != 1: + raise ValueError( + "GEAK v4 V1 requires exactly one source_file_path; " + f"got {raw!r}" + ) + value = values[0] + if not isinstance(value, str) or not value.strip(): + raise ValueError("GEAK v4 source_file_path must be a non-empty string") + if value != value.strip(): + raise ValueError(f"source_file_path must not contain outer whitespace: {value!r}") + if "\\" in value or any(ord(char) < 32 for char in value): + raise ValueError(f"unsafe source_file_path: {value!r}") + + relative = PurePosixPath(value) + if ( + relative.is_absolute() + or relative.as_posix() != value + or any(part in ("", ".", "..") for part in relative.parts) + ): + raise ValueError(f"source_file_path must be a safe relative path: {value!r}") + if is_protected_workspace_path(Path(*relative.parts)): + raise ValueError(f"declared source is protected by the Arena harness guard: {value}") + if _UNSAFE_SOURCE_NAME.fullmatch(relative.name): + raise ValueError( + "GEAK v4 V1 does not support a source file that also looks like a " + f"co-located test/harness: {value}" + ) + + source = workspace.joinpath(*relative.parts) + try: + metadata = source.lstat() + except OSError as exc: + raise ValueError(f"declared source is not readable: {source}") from exc + if ( + not stat.S_ISREG(metadata.st_mode) + or metadata.st_nlink != 1 + ): + raise ValueError(f"declared source must be an existing regular file: {source}") + try: + source.resolve().relative_to(workspace.resolve()) + except ValueError as exc: + raise ValueError(f"declared source escapes the workspace: {source}") from exc + return relative + + +def _snapshot_ignore(directory: str, names: list[str]) -> set[str]: + ignored: set[str] = set() + for name in names: + path = Path(directory) / name + if name in _COPY_IGNORED_DIRS: + ignored.add(name) + elif path.is_file() and path.suffix in _COPY_IGNORED_SUFFIXES: + ignored.add(name) + elif name in {"task_result.yaml", "validation_report.yaml"}: + ignored.add(name) + return ignored + + +def _materialize_disposable_input(workspace: Path, destination: Path) -> None: + """Copy the task so Workflow tools never receive the scoring workspace path.""" + if destination.exists(): + raise FileExistsError(f"disposable GEAK input already exists: {destination}") + symlinks = sorted( + str(path.relative_to(workspace)) + for path in workspace.rglob("*") + if path.is_symlink() + ) + if symlinks: + raise ValueError( + "GEAK v4 does not accept symlinks in a task workspace because a " + f"disposable copy could retain references to protected data: {symlinks[:10]}" + ) + shutil.copytree( + workspace, + destination, + symlinks=False, + ignore=_snapshot_ignore, + ) + + +def _logical_gpu_ids(eval_config: dict[str, Any]) -> str: + """Return GPU IDs in the process-visible namespace. + + Arena parallel workers mask one physical GPU with ROCR_VISIBLE_DEVICES and + expose it as logical HIP/CUDA device 0. GEAK's gpu_lock wrapper rewrites + HIP_VISIBLE_DEVICES again, so forwarding the host ID would hide the GPU. + """ + if os.environ.get("AGENT_KERNEL_ARENA_HOST_GPU_ID") is not None: + return "0" + + visible = ( + os.environ.get("HIP_VISIBLE_DEVICES") + or os.environ.get("CUDA_VISIBLE_DEVICES") + ) + if visible: + count = len([part for part in visible.split(",") if part.strip()]) + if count: + return ",".join(str(index) for index in range(count)) + + override = os.environ.get("GEAK_V4_GPU_IDS") + configured = override if override is not None else eval_config.get("gpu_ids", "0") + if isinstance(configured, (list, tuple)): + configured = ",".join(str(item) for item in configured) + return str(configured) + + +def _directory_identity(path: Path) -> tuple[int, int]: + try: + metadata = path.lstat() + except OSError as exc: + raise RuntimeError(f"artifact directory is not readable: {path}") from exc + if not stat.S_ISDIR(metadata.st_mode): + raise RuntimeError( + f"artifact directory must be a real directory, not a symlink: {path}" + ) + return metadata.st_dev, metadata.st_ino + + +def _open_directory_fd( + path: Path, + expected_identity: tuple[int, int] | None = None, +) -> int: + flags = os.O_RDONLY + flags |= getattr(os, "O_CLOEXEC", 0) + flags |= getattr(os, "O_DIRECTORY", 0) + flags |= getattr(os, "O_NOFOLLOW", 0) + descriptor = os.open(path, flags) + metadata = os.fstat(descriptor) + identity = (metadata.st_dev, metadata.st_ino) + if ( + not stat.S_ISDIR(metadata.st_mode) + or ( + expected_identity is not None + and identity != expected_identity + ) + ): + os.close(descriptor) + raise RuntimeError(f"artifact directory identity changed: {path}") + return descriptor + + +def _verify_directory_identity( + path: Path, + expected_identity: tuple[int, int], +) -> None: + if _directory_identity(path) != expected_identity: + raise RuntimeError(f"artifact directory identity changed: {path}") + + +def _verify_artifact_directories( + artifact_root: Path, + artifact_root_identity: tuple[int, int], + run_dir: Path, + run_dir_identity: tuple[int, int], +) -> None: + _verify_directory_identity(artifact_root, artifact_root_identity) + _verify_directory_identity(run_dir, run_dir_identity) + + +def _new_run_paths(workspace: Path) -> dict[str, Path]: + root = workspace.parent / f".{workspace.name}_geak_v4" + run_id = ( + datetime.now(timezone.utc).strftime("%Y%m%d_%H%M%S") + + f"_{os.getpid()}_{uuid.uuid4().hex[:8]}" + ) + run_dir = root / run_id + root.mkdir(mode=0o700, exist_ok=True) + root_identity = _directory_identity(root) + root_fd = _open_directory_fd(root, root_identity) + try: + os.mkdir(run_id, mode=0o700, dir_fd=root_fd) + finally: + os.close(root_fd) + _verify_directory_identity(root, root_identity) + _directory_identity(run_dir) + return { + "artifact_root": root, + "run_dir": run_dir, + "input": run_dir / "input", + "eval": run_dir / "eval", + "exp_root": run_dir / "runs", + "handoff": run_dir / "handoff.json", + "result": run_dir / "result.json", + } + + +def _atomic_write_json( + path: Path, + value: dict[str, Any], + *, + expected_parent_identity: tuple[int, int] | None = None, +) -> None: + encoded = ( + json.dumps(value, indent=2, sort_keys=True) + "\n" + ).encode("utf-8") + directory_fd = _open_directory_fd(path.parent, expected_parent_identity) + temporary_fd = -1 + temporary_name = "" + try: + proc_directory = Path(f"/proc/self/fd/{directory_fd}") + temporary_fd, temporary_path = tempfile.mkstemp( + prefix=f".{path.name}.tmp.", + dir=proc_directory, + ) + temporary_name = Path(temporary_path).name + metadata = os.fstat(temporary_fd) + if not stat.S_ISREG(metadata.st_mode) or metadata.st_nlink != 1: + raise OSError("atomic JSON temporary is not a private regular file") + with os.fdopen(temporary_fd, "wb", closefd=True) as stream: + temporary_fd = -1 + stream.write(encoded) + stream.flush() + os.fsync(stream.fileno()) + os.replace( + temporary_name, + path.name, + src_dir_fd=directory_fd, + dst_dir_fd=directory_fd, + ) + temporary_name = "" + os.fsync(directory_fd) + finally: + if temporary_fd >= 0: + os.close(temporary_fd) + if temporary_name: + try: + os.unlink(temporary_name, dir_fd=directory_fd) + except FileNotFoundError: + pass + os.close(directory_fd) + + +def _read_bounded_text(path: Path, size_limit: int) -> str | None: + flags = os.O_RDONLY + flags |= getattr(os, "O_CLOEXEC", 0) + flags |= getattr(os, "O_NOFOLLOW", 0) + flags |= getattr(os, "O_NONBLOCK", 0) + try: + descriptor = os.open(path, flags) + except OSError: + return None + try: + metadata = os.fstat(descriptor) + if ( + not stat.S_ISREG(metadata.st_mode) + or metadata.st_size > size_limit + ): + return None + with os.fdopen(descriptor, "rb", closefd=False) as stream: + content = stream.read(size_limit + 1) + if len(content) > size_limit: + return None + return content.decode("utf-8") + except (OSError, UnicodeDecodeError): + return None + finally: + os.close(descriptor) + + +def _read_json(path: Path) -> dict[str, Any] | None: + content = _read_bounded_text(path, _JSON_SIZE_LIMIT) + if content is None: + return None + try: + value = json.loads(content) + except json.JSONDecodeError: + return None + return value if isinstance(value, dict) else None + + +def _process_group_exists(pgid: int) -> bool: + try: + os.killpg(pgid, 0) + except ProcessLookupError: + return False + except PermissionError: + return True + return True + + +def _wait_group_exit(process: subprocess.Popen[str], pgid: int, timeout: float) -> bool: + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + process.poll() + if not _process_group_exists(pgid): + return True + time.sleep(0.1) + return not _process_group_exists(pgid) + + +def _terminate_process_group( + process: subprocess.Popen[str], + logger: logging.Logger, +) -> None: + pgid = process.pid + if not _process_group_exists(pgid): + process.poll() + return + try: + os.killpg(pgid, signal.SIGTERM) + except ProcessLookupError: + return + if _wait_group_exit(process, pgid, 10): + return + logger.warning("Force killing GEAK runner process group") + try: + os.killpg(pgid, signal.SIGKILL) + except ProcessLookupError: + return + _wait_group_exit(process, pgid, 5) + + +def _stream_pipe( + stream, + prefix: str, + output: list[str], + log, +) -> None: + captured = 0 + truncated = False + try: + while True: + chunk = stream.read(4096) + if not chunk: + break + remaining = _PROCESS_OUTPUT_LIMIT - captured + if remaining > 0: + retained = chunk[:remaining] + output.append(retained) + captured += len(retained) + compact = " ".join(retained[:2000].split()) + if compact: + log(f"{prefix} {compact[:500]}") + if len(chunk) > remaining and not truncated: + truncated = True + log(f"{prefix} output truncated at {_PROCESS_OUTPUT_LIMIT} characters") + finally: + stream.close() + + +def _run_workflow_runner( + handoff_path: Path, + result_path: Path, + *, + timeout_seconds: int, + logger: logging.Logger, +) -> str: + runner = Path(__file__).with_name("workflow_runner.py") + command = [sys.executable, str(runner), str(handoff_path), str(result_path)] + process = subprocess.Popen( + command, + cwd=str(handoff_path.parent), + stdin=subprocess.DEVNULL, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + errors="replace", + bufsize=1, + start_new_session=True, + ) + assert process.stdout is not None + assert process.stderr is not None + stdout: list[str] = [] + stderr: list[str] = [] + stdout_thread = threading.Thread( + target=_stream_pipe, + args=(process.stdout, "[GEAK]", stdout, logger.info), + daemon=True, + ) + stderr_thread = threading.Thread( + target=_stream_pipe, + args=(process.stderr, "[GEAK STDERR]", stderr, logger.warning), + daemon=True, + ) + stdout_thread.start() + stderr_thread.start() + + timed_out = False + try: + process.wait(timeout=timeout_seconds + 30) + except subprocess.TimeoutExpired: + timed_out = True + logger.error("GEAK v4 runner exceeded its hard timeout") + finally: + # Also clean descendants if the runner leader exited while a background + # Workflow/Claude process remained in the session. + _terminate_process_group(process, logger) + stdout_thread.join(timeout=2) + stderr_thread.join(timeout=2) + + if timed_out: + raise TimeoutError(f"GEAK v4 timed out after {timeout_seconds} seconds") + + result = _read_json(result_path) + stderr_text = "".join(stderr) + if process.returncode != 0: + detail = result.get("error") if result else stderr_text[-4000:] + raise RuntimeError( + f"GEAK v4 runner failed with exit {process.returncode}: {detail}" + ) + if result is None: + raise RuntimeError(f"GEAK v4 runner did not write a valid result: {result_path}") + return "\n".join( + part for part in ("".join(stdout), stderr_text) if part + ) + + +def _run_git_apply_inspection( + patch: Path, + *options: str, + cwd: Path | None = None, + binary: bool = False, + isolate_from_parent_repo: bool = False, +) -> subprocess.CompletedProcess: + environment = None + if isolate_from_parent_repo: + if cwd is None: + raise ValueError("isolated git apply inspection requires cwd") + environment = dict(os.environ) + environment.pop("GIT_DIR", None) + environment.pop("GIT_WORK_TREE", None) + environment["GIT_CEILING_DIRECTORIES"] = str(cwd.parent.resolve()) + return subprocess.run( + ["git", "apply", *options, str(patch)], + cwd=str(cwd) if cwd else None, + env=environment, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=not binary, + check=False, + ) + + +def _patch_paths(patch: Path) -> list[PurePosixPath]: + result = _run_git_apply_inspection(patch, "--numstat", "-z", binary=True) + if result.returncode != 0: + stderr = os.fsdecode(result.stderr) + raise RuntimeError(f"cannot parse GEAK patch: {stderr.strip()}") + + paths: list[PurePosixPath] = [] + for record in result.stdout.split(b"\0"): + if not record: + continue + fields = record.split(b"\t", 2) + if len(fields) != 3: + raise RuntimeError("GEAK patch has malformed numstat output") + added, deleted, raw_path = fields + if not added.isdigit() or not deleted.isdigit(): + raise RuntimeError("binary GEAK patches are not accepted") + value = os.fsdecode(raw_path) + if ( + not value + or "\\" in value + or any(ord(char) < 32 for char in value) + ): + raise RuntimeError(f"GEAK patch contains an unsafe path: {value!r}") + path = PurePosixPath(value) + if path.is_absolute() or any(part in ("", ".", "..") for part in path.parts): + raise RuntimeError(f"GEAK patch path escapes the workspace: {value!r}") + paths.append(path) + return paths + + +def _sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as stream: + for chunk in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def _workspace_manifest(workspace: Path) -> dict[str, tuple[Any, ...]]: + """Capture all workspace entries so Workflow cannot bypass patch import.""" + manifest: dict[str, tuple[Any, ...]] = {} + for path in sorted(workspace.rglob("*")): + relative = str(path.relative_to(workspace)) + metadata = path.lstat() + mode = stat.S_IMODE(metadata.st_mode) + if stat.S_ISLNK(metadata.st_mode): + manifest[relative] = ("symlink", mode, os.readlink(path)) + elif stat.S_ISREG(metadata.st_mode): + manifest[relative] = ("file", mode, metadata.st_size, _sha256(path)) + elif stat.S_ISDIR(metadata.st_mode): + manifest[relative] = ("directory", mode) + else: + manifest[relative] = ("special", stat.S_IFMT(metadata.st_mode), mode) + return manifest + + +def _verify_workspace_manifest( + expected: dict[str, tuple[Any, ...]], + workspace: Path, +) -> None: + """Reject any direct mutation of the Arena scoring workspace.""" + current = _workspace_manifest(workspace) + if current == expected: + return + + expected_paths = set(expected) + current_paths = set(current) + added = sorted(current_paths - expected_paths) + deleted = sorted(expected_paths - current_paths) + changed = sorted( + path + for path in expected_paths & current_paths + if expected[path] != current[path] + ) + raise RuntimeError( + "GEAK v4 detected a direct mutation of the Arena scoring workspace; " + "only the validated single-file patch import is allowed " + f"(added={added[:10]}, deleted={deleted[:10]}, changed={changed[:10]})" + ) + + +def _apply_validated_patch( + *, + result: dict[str, Any], + expected_eval_dir: Path, + workspace: Path, + source_path: PurePosixPath, + min_improve: float, + run_dir: Path, + expected_workspace_manifest: dict[str, tuple[Any, ...]] | None = None, + artifact_root: Path | None = None, + artifact_root_identity: tuple[int, int] | None = None, + run_dir_identity: tuple[int, int] | None = None, +) -> bool: + if result.get("schema_version") != 1: + raise RuntimeError( + "GEAK result schema is not supported: " + f"{result.get('schema_version')!r}" + ) + status = str(result.get("status") or "") + if status in {"no_gain", "rejected"}: + return False + if status != "ok": + raise RuntimeError(f"GEAK result is not importable: status={status!r}") + if result.get("validation_status") != "accepted": + raise RuntimeError("GEAK result was not accepted by its Director") + if result.get("correctness") != "pass": + raise RuntimeError("GEAK Director did not report correctness=pass") + if str(result.get("applied_to_original", "unknown")).lower() != "false": + raise RuntimeError("GEAK unexpectedly wrote directly to its input snapshot") + if not math.isfinite(min_improve) or min_improve < 0: + raise RuntimeError("GEAK min_improve policy must be finite and non-negative") + + try: + speedup = float(result["final_speedup"]) + except (KeyError, TypeError, ValueError) as exc: + raise RuntimeError("GEAK result has no numeric final_speedup") from exc + if not math.isfinite(speedup): + raise RuntimeError("GEAK final_speedup must be finite") + if speedup < 1.0 + min_improve: + return False + + raw_eval_dir = result.get("eval_dir") + if not isinstance(raw_eval_dir, str): + raise RuntimeError("GEAK result is missing eval_dir") + if raw_eval_dir != str(expected_eval_dir): + raise RuntimeError("GEAK result eval_dir does not match the pinned directory") + if expected_eval_dir.is_symlink() or not expected_eval_dir.is_dir(): + raise RuntimeError( + f"GEAK eval_dir is missing, not a directory, or a symlink: {expected_eval_dir}" + ) + + # The patch location is fixed by the workflow contract. Never trust an + # arbitrary absolute path returned by an agent. + patch = expected_eval_dir / "final_patch.diff" + for field in ("final_patch", "director_final_patch"): + reported = result.get(field) + if not isinstance(reported, str) or reported != str(patch): + raise RuntimeError( + f"GEAK {field} does not match the Director-validated pinned patch" + ) + workflow_patch = result.get("workflow_final_patch") + if workflow_patch is not None and workflow_patch != str(patch): + raise RuntimeError( + "GEAK workflow_final_patch does not match the pinned patch" + ) + if patch.is_symlink() or not patch.is_file(): + raise RuntimeError(f"GEAK final patch is missing or not a regular file: {patch}") + size = patch.stat().st_size + if size <= 0 or size > _PATCH_SIZE_LIMIT: + raise RuntimeError(f"GEAK final patch has an invalid size: {size} bytes") + + touched = _patch_paths(patch) + if touched != [source_path]: + raise RuntimeError( + "GEAK patch must modify exactly the declared source file; " + f"declared={source_path}, touched={touched}" + ) + if is_protected_workspace_path(Path(*source_path.parts)): + raise RuntimeError(f"GEAK patch targets a protected harness path: {source_path}") + + if ( + artifact_root is not None + and artifact_root_identity is not None + and run_dir_identity is not None + ): + _verify_artifact_directories( + artifact_root, + artifact_root_identity, + run_dir, + run_dir_identity, + ) + + baseline_manifest = ( + dict(expected_workspace_manifest) + if expected_workspace_manifest is not None + else _workspace_manifest(workspace) + ) + _verify_workspace_manifest(baseline_manifest, workspace) + + summary = _run_git_apply_inspection(patch, "--summary") + if summary.returncode != 0: + raise RuntimeError(f"cannot summarize GEAK patch: {summary.stderr.strip()}") + if summary.stdout.strip(): + raise RuntimeError( + "GEAK patch contains a create/delete/rename/mode operation: " + + summary.stdout.strip() + ) + + destination = workspace.joinpath(*source_path.parts) + destination_metadata = destination.lstat() + if ( + not stat.S_ISREG(destination_metadata.st_mode) + or destination_metadata.st_nlink != 1 + ): + raise RuntimeError(f"Arena source changed type before patch import: {destination}") + original_mode = destination_metadata.st_mode + original_digest = _sha256(destination) + harness_snapshot = snapshot_workspace_harness(workspace) + + with tempfile.TemporaryDirectory(prefix="patch_stage_", dir=run_dir) as temporary: + staging = Path(temporary) + staged_source = staging.joinpath(*source_path.parts) + staged_source.parent.mkdir(parents=True, exist_ok=True) + shutil.copy2(destination, staged_source) + + check = _run_git_apply_inspection( + patch, + "--check", + cwd=staging, + isolate_from_parent_repo=True, + ) + if check.returncode != 0: + raise RuntimeError(f"GEAK patch does not apply cleanly: {check.stderr.strip()}") + apply_result = _run_git_apply_inspection( + patch, + cwd=staging, + isolate_from_parent_repo=True, + ) + if apply_result.returncode != 0: + raise RuntimeError(f"failed to apply GEAK patch: {apply_result.stderr.strip()}") + staged_metadata = staged_source.lstat() + if ( + not stat.S_ISREG(staged_metadata.st_mode) + or staged_metadata.st_nlink != 1 + ): + raise RuntimeError("GEAK patch changed the source file type") + if staged_metadata.st_mode != original_mode: + raise RuntimeError("GEAK patch changed the source file mode") + if _sha256(staged_source) == original_digest: + raise RuntimeError("GEAK patch reported success but did not change the source") + + current_metadata = destination.lstat() + if ( + not stat.S_ISREG(current_metadata.st_mode) + or current_metadata.st_nlink != 1 + or current_metadata.st_mode != original_mode + or _sha256(destination) != original_digest + ): + raise RuntimeError("Arena source changed while GEAK patch was staged") + # TemporaryDirectory lives beside the workspace, so os.replace is an + # atomic same-filesystem mutation of the one approved source file. + os.replace(staged_source, destination) + + verify_workspace_harness(harness_snapshot) + updated_metadata = destination.lstat() + expected_after = dict(baseline_manifest) + expected_after[str(source_path)] = ( + "file", + stat.S_IMODE(updated_metadata.st_mode), + updated_metadata.st_size, + _sha256(destination), + ) + _verify_workspace_manifest(expected_after, workspace) + + if ( + artifact_root is not None + and artifact_root_identity is not None + and run_dir_identity is not None + ): + _verify_artifact_directories( + artifact_root, + artifact_root_identity, + run_dir, + run_dir_identity, + ) + audit_parent_identity = ( + run_dir_identity + if run_dir_identity is not None + else _directory_identity(run_dir) + ) + _atomic_write_json( + run_dir / "applied_patch.json", + { + "source_file": str(source_path), + "patch": str(patch), + "patch_sha256": _sha256(patch), + "source_sha256": _sha256(destination), + "director_speedup": speedup, + }, + expected_parent_identity=audit_parent_identity, + ) + _verify_workspace_manifest(expected_after, workspace) + return True + + +@register_agent("geak_v4") +def launch_agent( + eval_config: dict[str, Any], + task_config_dir: str, + workspace: str, +) -> str: + """Run GEAK v4 against a disposable copy, then import one validated patch.""" + logger = logging.getLogger(__name__) + agent_config = _load_agent_config() + task_config = _load_task_config(task_config_dir) + workspace_path = Path(workspace).resolve() + if not workspace_path.is_dir(): + raise FileNotFoundError(f"Arena workspace does not exist: {workspace_path}") + + task_type = str(task_config.get("task_type") or "") + supported = {str(value) for value in agent_config.get("supported_task_types", [])} + if task_type not in supported: + raise ValueError( + f"GEAK v4 V1 does not support task_type={task_type!r}; " + f"supported task types: {sorted(supported)}" + ) + source_path = _single_declared_source(task_config, workspace_path) + workspace_manifest = _workspace_manifest(workspace_path) + + workflow_dir = Path( + os.environ.get("GEAK_V4_WORKFLOW_DIR") + or agent_config.get("workflow_dir") + or "/opt/geak/kernel_workflow" + ).resolve() + workflow_script = workflow_dir / "kernel_workflow.js" + if not workflow_script.is_file(): + raise FileNotFoundError( + f"GEAK v4 workflow not found: {workflow_script}. " + "Set AKA_GEAK_ROOT on the host and run make docker-setup-geak." + ) + claude_binary = shutil.which("claude") + if not claude_binary: + raise RuntimeError("Claude Code CLI not found on PATH") + + paths = _new_run_paths(workspace_path) + artifact_root_identity = _directory_identity(paths["artifact_root"]) + run_dir_identity = _directory_identity(paths["run_dir"]) + _materialize_disposable_input(workspace_path, paths["input"]) + + prompt_builder = load_prompt_builder(AgentType.GEAK_V4, logger) + task_prompt = prompt_builder( + task_config_dir, + str(paths["input"]), + eval_config, + logger, + ) + task_prompt += ( + "\n\n### GEAK/Arena Integration Contract\n" + "The config.yaml compile, correctness, and performance commands are the " + "measurement source of truth. Do not create, modify, or replace any test, " + "harness, config, eval_tools, reference, or timing file. Optimize only " + f"`{source_path}`. The caller will accept a patch only when it modifies " + "that one existing file." + ) + + timeout_seconds = int(agent_config.get("timeout_seconds", 43200)) + handoff = { + "schema_version": 1, + "kernel_path": str(paths["input"]), + "workflow_dir": str(workflow_dir), + "eval_dir": str(paths["eval"]), + "exp_root": str(paths["exp_root"]), + "gpu_ids": _logical_gpu_ids(eval_config), + "budget": int(agent_config.get("budget", 6)), + "min_improve": float(agent_config.get("min_improve", 0.02)), + "deep_cost": int(agent_config.get("deep_cost", 2)), + "use_expert_skills": bool(agent_config.get("use_expert_skills", False)), + "task": task_prompt, + "model": str(agent_config.get("model", "claude-opus-4-8")), + "effort": str(agent_config.get("effort", "ultracode")), + "claude_cli_path": claude_binary, + "timeout_seconds": timeout_seconds, + "done_grace_seconds": float(agent_config.get("done_grace_seconds", 1800)), + "done_poll_seconds": float(agent_config.get("done_poll_seconds", 5)), + } + _atomic_write_json( + paths["handoff"], + handoff, + expected_parent_identity=run_dir_identity, + ) + + logger.info("GEAK v4 preflight") + logger.info(" workflow: %s", workflow_script) + logger.info(" disposable input: %s", paths["input"]) + logger.info(" eval dir: %s", paths["eval"]) + logger.info(" source allowlist: %s", source_path) + logger.info(" logical GPU IDs: %s", handoff["gpu_ids"]) + logger.info(" budget: %s", handoff["budget"]) + logger.info(" timeout: %ss", timeout_seconds) + + try: + output = _run_workflow_runner( + paths["handoff"], + paths["result"], + timeout_seconds=timeout_seconds, + logger=logger, + ) + finally: + try: + _verify_artifact_directories( + paths["artifact_root"], + artifact_root_identity, + paths["run_dir"], + run_dir_identity, + ) + finally: + _verify_workspace_manifest(workspace_manifest, workspace_path) + result = _read_json(paths["result"]) + if result is None: + raise RuntimeError("GEAK v4 result disappeared after runner completion") + + applied = _apply_validated_patch( + result=result, + expected_eval_dir=paths["eval"], + workspace=workspace_path, + source_path=source_path, + min_improve=float(agent_config.get("min_improve", 0.02)), + run_dir=paths["run_dir"], + expected_workspace_manifest=workspace_manifest, + artifact_root=paths["artifact_root"], + artifact_root_identity=artifact_root_identity, + run_dir_identity=run_dir_identity, + ) + _verify_artifact_directories( + paths["artifact_root"], + artifact_root_identity, + paths["run_dir"], + run_dir_identity, + ) + if applied: + logger.info("Imported GEAK v4 Director-validated patch into Arena workspace") + else: + logger.info("GEAK v4 produced no accepted gain; Arena workspace is unchanged") + return output + "\n" + json.dumps(result, sort_keys=True) diff --git a/agents/geak_v4/workflow_runner.py b/agents/geak_v4/workflow_runner.py new file mode 100644 index 00000000..85eaa316 --- /dev/null +++ b/agents/geak_v4/workflow_runner.py @@ -0,0 +1,857 @@ +#!/usr/bin/env python3 +"""Stable GEAK v4 kernel-workflow runner for external orchestrators. + +The GEAK JavaScript workflow can only execute inside Claude Code's dynamic +``Workflow`` runtime. This module keeps that volatile SDK/tool lifecycle out of +the Arena launcher: + +* map a versioned handoff onto ``kernel_workflow.js`` arguments; +* pin a known evaluation directory for completion/recovery; +* keep the SDK client alive when Workflow runs as a background task; +* recover the authoritative result from on-disk GEAK artifacts; and +* never let GEAK write directly to the Arena workspace. + +The command is intentionally usable in ``--dry-run`` mode without importing the +Claude Agent SDK or contacting a model. +""" + +from __future__ import annotations + +import argparse +import json +import math +import os +import shutil +import stat +import sys +import tempfile +import time +from pathlib import Path +from typing import Any, Iterable + + +SCHEMA_VERSION = 1 +ALLOWED_TOOLS = ["Workflow", "Bash", "Read", "Write"] +VALID_EFFORTS = {"low", "medium", "high", "xhigh", "max"} +DEFAULT_SETTINGS = {"enableWorkflows": True, "ultracode": True} +_JSON_SIZE_LIMIT = 8 * 1024 * 1024 +_SDK_OUTPUT_FILE_LIMIT = 8 * 1024 * 1024 +_TRANSCRIPT_SIZE_LIMIT = 8 * 1024 * 1024 +_TRANSCRIPT_JSON_LINE_LIMIT = 64 + + +class HandoffError(ValueError): + """The caller supplied an invalid GEAK handoff.""" + + +def _open_directory_fd( + path: Path | str, + *, + parent_fd: int | None = None, +) -> int: + flags = os.O_RDONLY + flags |= getattr(os, "O_CLOEXEC", 0) + flags |= getattr(os, "O_DIRECTORY", 0) + flags |= getattr(os, "O_NOFOLLOW", 0) + descriptor = os.open(path, flags, dir_fd=parent_fd) + metadata = os.fstat(descriptor) + if not stat.S_ISDIR(metadata.st_mode): + os.close(descriptor) + raise OSError(f"not a directory: {path}") + return descriptor + + +def _atomic_write_json( + path: Path, + payload: dict[str, Any], + *, + directory_fd: int | None = None, +) -> None: + """Write JSON without following attacker-created file symlinks. + + When ``directory_fd`` is supplied, the write stays pinned to that already + opened directory even if Workflow renames or replaces its pathname. + """ + encoded = ( + json.dumps(payload, indent=2, sort_keys=True) + "\n" + ).encode("utf-8") + owned_directory_fd = directory_fd is None + if directory_fd is None: + directory_fd = _open_directory_fd(path.parent) + + temporary_fd = -1 + temporary_name = "" + try: + proc_directory = Path(f"/proc/self/fd/{directory_fd}") + temporary_fd, temporary_path = tempfile.mkstemp( + prefix=f".{path.name}.tmp.", + dir=proc_directory, + ) + temporary_name = Path(temporary_path).name + metadata = os.fstat(temporary_fd) + if not stat.S_ISREG(metadata.st_mode) or metadata.st_nlink != 1: + raise OSError("atomic JSON temporary is not a private regular file") + with os.fdopen(temporary_fd, "wb", closefd=True) as stream: + temporary_fd = -1 + stream.write(encoded) + stream.flush() + os.fsync(stream.fileno()) + os.replace( + temporary_name, + path.name, + src_dir_fd=directory_fd, + dst_dir_fd=directory_fd, + ) + temporary_name = "" + os.fsync(directory_fd) + finally: + if temporary_fd >= 0: + os.close(temporary_fd) + if temporary_name: + try: + os.unlink(temporary_name, dir_fd=directory_fd) + except FileNotFoundError: + pass + if owned_directory_fd: + os.close(directory_fd) + + +def _read_bounded_text(path: Path, size_limit: int) -> str | None: + flags = os.O_RDONLY + flags |= getattr(os, "O_CLOEXEC", 0) + flags |= getattr(os, "O_NOFOLLOW", 0) + flags |= getattr(os, "O_NONBLOCK", 0) + try: + descriptor = os.open(path, flags) + except OSError: + return None + try: + metadata = os.fstat(descriptor) + if ( + not stat.S_ISREG(metadata.st_mode) + or metadata.st_size > size_limit + ): + return None + with os.fdopen(descriptor, "rb", closefd=False) as stream: + content = stream.read(size_limit + 1) + if len(content) > size_limit: + return None + return content.decode("utf-8") + except (OSError, UnicodeDecodeError): + return None + finally: + os.close(descriptor) + + +def _read_json(path: Path) -> dict[str, Any] | None: + content = _read_bounded_text(path, _JSON_SIZE_LIMIT) + if content is None: + return None + try: + value = json.loads(content) + except json.JSONDecodeError: + return None + return value if isinstance(value, dict) else None + + +def load_handoff(path: Path) -> dict[str, Any]: + value = _read_json(path) + if value is None: + raise HandoffError(f"handoff is not a readable JSON object: {path}") + if value.get("schema_version") != SCHEMA_VERSION: + raise HandoffError( + f"unsupported handoff schema_version={value.get('schema_version')!r}; " + f"expected {SCHEMA_VERSION}" + ) + return value + + +def _absolute_path(value: Any, field: str) -> Path: + if not isinstance(value, str) or not value.strip(): + raise HandoffError(f"{field} must be a non-empty absolute path") + path = Path(value) + if not path.is_absolute(): + raise HandoffError(f"{field} must be absolute: {path}") + return path.resolve() + + +def _positive_int(value: Any, field: str) -> int: + try: + parsed = int(value) + except (TypeError, ValueError) as exc: + raise HandoffError(f"{field} must be an integer") from exc + if parsed <= 0: + raise HandoffError(f"{field} must be positive") + return parsed + + +def _nonnegative_float(value: Any, field: str) -> float: + try: + parsed = float(value) + except (TypeError, ValueError) as exc: + raise HandoffError(f"{field} must be numeric") from exc + if not math.isfinite(parsed) or parsed < 0: + raise HandoffError(f"{field} must be finite and non-negative") + return parsed + + +def _gpu_ids(value: Any) -> str: + if isinstance(value, (list, tuple)): + value = ",".join(str(item) for item in value) + text = str(value if value is not None else "0") + parts = [part.strip() for part in text.split(",") if part.strip()] + if not parts or any(not part.isdigit() for part in parts): + raise HandoffError(f"gpu_ids must be comma-separated non-negative integers: {text!r}") + return ",".join(parts) + + +def map_workflow_args(handoff: dict[str, Any]) -> tuple[Path, dict[str, Any]]: + """Validate a handoff and return ``(script_path, workflow_args)``.""" + kernel_path = _absolute_path(handoff.get("kernel_path"), "kernel_path") + workflow_dir = _absolute_path(handoff.get("workflow_dir"), "workflow_dir") + eval_dir = _absolute_path(handoff.get("eval_dir"), "eval_dir") + exp_root = _absolute_path( + handoff.get("exp_root") or str(eval_dir.parent), + "exp_root", + ) + script_path = workflow_dir / "kernel_workflow.js" + + if not kernel_path.is_dir(): + raise HandoffError(f"kernel_path is not a directory: {kernel_path}") + if not script_path.is_file(): + raise HandoffError(f"GEAK kernel workflow not found: {script_path}") + if eval_dir.exists() and ( + not eval_dir.is_dir() or any(eval_dir.iterdir()) + ): + raise HandoffError(f"eval_dir must be absent or an empty directory: {eval_dir}") + for path, field in ((eval_dir, "eval_dir"), (exp_root, "exp_root")): + try: + path.relative_to(kernel_path) + except ValueError: + pass + else: + raise HandoffError( + f"{field} must not be inside kernel_path; GEAK copies kernel_path " + f"and would recursively copy its own outputs: {path}" + ) + + args: dict[str, Any] = { + "kernel_path": str(kernel_path), + "workflow_dir": str(workflow_dir), + "eval_dir": str(eval_dir), + "exp_root": str(exp_root), + "gpu_ids": _gpu_ids(handoff.get("gpu_ids", "0")), + "budget": _positive_int(handoff.get("budget", 6), "budget"), + "min_improve": _nonnegative_float( + handoff.get("min_improve", 0.02), + "min_improve", + ), + "deep_cost": _positive_int(handoff.get("deep_cost", 2), "deep_cost"), + "mode": "optimize", + # Arena owns the task workspace and applies only a policy-checked patch. + "apply_to_original": "false", + } + task = handoff.get("task") + if task: + args["task"] = str(task) + if bool(handoff.get("use_expert_skills", False)): + args["use_expert_skills"] = "true" + return script_path, args + + +def build_prompt(script_path: Path, workflow_args: dict[str, Any]) -> str: + eval_dir = workflow_args["eval_dir"] + return ( + "Invoke the Workflow tool exactly once with:\n" + f' scriptPath: "{script_path}"\n' + f" args: {json.dumps(workflow_args, ensure_ascii=False)}\n" + "Run the complete GEAK kernel pipeline through independent Director " + "validation. Do not edit the original kernel_path directly; " + "apply_to_original is false and the caller owns patch import. When the " + "Workflow finishes, write its exact full return object as compact JSON to " + f'"{eval_dir}/workflow_return.json", then print exactly that compact JSON ' + "as the final line and print nothing after it." + ) + + +def _iter_message_text(message: Any) -> Iterable[str]: + """Yield text fragments from SDK objects across supported SDK shapes.""" + if message is None: + return + if isinstance(message, str): + if message.strip(): + yield message + return + if isinstance(message, dict): + for key in ("result", "text", "summary"): + value = message.get(key) + if isinstance(value, str) and value.strip(): + yield value + content = message.get("content") + if isinstance(content, str): + if content.strip(): + yield content + elif isinstance(content, (list, tuple)): + for item in content: + yield from _iter_message_text(item) + return + + for attribute in ("result", "text", "summary"): + value = getattr(message, attribute, None) + if isinstance(value, str) and value.strip(): + yield value + content = getattr(message, "content", None) + if isinstance(content, str): + if content.strip(): + yield content + elif isinstance(content, (list, tuple)): + for item in content: + yield from _iter_message_text(item) + + +def _finite_number(value: Any) -> bool: + try: + return math.isfinite(float(value)) + except (TypeError, ValueError): + return False + + +def _matches_pinned_patch(value: Any, eval_dir: Path) -> bool: + return isinstance(value, str) and value == str(eval_dir / "final_patch.diff") + + +def _valid_workflow_return( + value: Any, + eval_dir: Path, + *, + require_pinned_patch: bool = False, +) -> bool: + if not isinstance(value, dict): + return False + raw_eval_dir = value.get("eval_dir") + if not isinstance(raw_eval_dir, str): + return False + matches = raw_eval_dir == str(eval_dir) + status = value.get("validation_status") + patch_contract_ok = ( + not require_pinned_patch + or status not in {"accepted", "flagged"} + or _matches_pinned_patch(value.get("final_patch"), eval_dir) + ) + return ( + matches + and isinstance(status, str) + and _finite_number(value.get("final_geomean")) + and isinstance(value.get("final_patch"), str) + and patch_contract_ok + and ( + "workload_aligned" not in value + or isinstance(value.get("workload_aligned"), bool) + ) + ) + + +def _valid_director_validation( + value: Any, + eval_dir: Path | None = None, +) -> bool: + valid = ( + isinstance(value, dict) + and value.get("validation_status") in {"accepted", "flagged"} + and value.get("correctness") in {"pass", "fail"} + and _finite_number(value.get("director_verified_speedup_geomean")) + and value.get("applied_to_original") in {"true", "false"} + and isinstance(value.get("final_patch"), str) + ) + return valid and ( + eval_dir is None + or _matches_pinned_patch(value.get("final_patch"), eval_dir) + ) + + +def _terminal_artifact_exists(eval_dir: Path) -> bool: + workflow_return = _read_json(eval_dir / "workflow_return.json") + director_validation = _read_json(eval_dir / "director_validation.json") + return _valid_workflow_return( + workflow_return, + eval_dir, + require_pinned_patch=True, + ) or _valid_director_validation(director_validation, eval_dir) + + +def _extract_workflow_return(transcript: str, expected_eval_dir: Path) -> dict[str, Any] | None: + """Read the final compact JSON line without quadratic brace scanning.""" + lines = transcript.splitlines() + for raw_line in reversed(lines[-_TRANSCRIPT_JSON_LINE_LIMIT:]): + line = raw_line.strip() + if not line or len(line.encode("utf-8")) > _JSON_SIZE_LIMIT: + continue + try: + value = json.loads(line) + except json.JSONDecodeError: + continue + if _valid_workflow_return( + value, + expected_eval_dir, + require_pinned_patch=True, + ): + return value + return None + + +def _completed_producer_error( + state: dict[str, bool], + pending: set[str], + producer_error: list[BaseException], +) -> BaseException | None: + if not state["producer_done"]: + return None + if producer_error: + return producer_error[0] + if pending: + return RuntimeError( + "Claude SDK message stream ended with unfinished GEAK tasks: " + f"{sorted(pending)}" + ) + if not state["result_seen"]: + return RuntimeError( + "Claude SDK message stream ended without a ResultMessage or " + "a valid GEAK terminal artifact" + ) + return None + + +def invoke_via_sdk( + prompt: str, + *, + workflow_dir: Path, + eval_dir: Path, + model: str, + effort: str, + settings: str, + cli_path: str, + timeout_seconds: int, + done_grace_seconds: float, + done_poll_seconds: float, +) -> str: + """Invoke Claude Code while surviving synchronous and background Workflows.""" + try: + import anyio + from claude_agent_sdk import ClaudeAgentOptions, ClaudeSDKClient + except ImportError as exc: + raise RuntimeError( + "claude_agent_sdk is required for reliable GEAK Workflow lifecycle " + "handling; run `make docker-setup-geak` first" + ) from exc + + option_extras: dict[str, Any] = {} + if effort in VALID_EFFORTS: + option_extras["effort"] = effort + if hasattr(os, "geteuid") and os.geteuid() == 0: + raise RuntimeError( + "GEAK v4 refuses to run Claude as root without a real OS sandbox; " + "use the Docker runner's non-root host UID mapping" + ) + sdk_env = { + "CLAUDE_CODE_DISABLE_AUTO_MEMORY": "1", + } + + options = ClaudeAgentOptions( + model=model, + allowed_tools=ALLOWED_TOOLS, + permission_mode="bypassPermissions", + settings=settings, + extra_args=option_extras, + cwd=str(workflow_dir), + env=sdk_env, + **({"cli_path": cli_path} if cli_path else {}), + ) + + async def _run() -> str: + chunks: list[str] = [] + captured_chars = 0 + pending: set[str] = set() + state = { + "background_started": False, + "terminal_task_seen": False, + "result_seen": False, + "producer_done": False, + } + + with anyio.fail_after(timeout_seconds): + async with ClaudeSDKClient(options=options) as client: + await client.query(prompt) + producer_error: list[BaseException] = [] + + async def _receive() -> None: + nonlocal captured_chars + try: + async for message in client.receive_messages(): + for text in _iter_message_text(message): + remaining = _TRANSCRIPT_SIZE_LIMIT - captured_chars + if remaining > 0: + retained = text[:remaining] + chunks.append(retained) + captured_chars += len(retained) + compact = " ".join(text[:2000].split()) + if compact: + print( + f"[GEAK SDK] {compact[:500]}", + file=sys.stderr, + flush=True, + ) + + name = type(message).__name__ + if name == "TaskStartedMessage": + task_id = getattr(message, "task_id", None) + if task_id: + pending.add(str(task_id)) + state["background_started"] = True + elif name == "TaskNotificationMessage": + state["terminal_task_seen"] = True + task_id = getattr(message, "task_id", None) + if task_id: + pending.discard(str(task_id)) + output_file = getattr(message, "output_file", None) + if output_file: + output = _read_bounded_text( + Path(output_file), + _SDK_OUTPUT_FILE_LIMIT, + ) + remaining = _TRANSCRIPT_SIZE_LIMIT - captured_chars + if output and remaining > 0: + retained = output[:remaining] + chunks.append(retained) + captured_chars += len(retained) + elif name == "ResultMessage": + state["result_seen"] = True + except BaseException as exc: + producer_error.append(exc) + finally: + state["producer_done"] = True + + async with anyio.create_task_group() as task_group: + task_group.start_soon(_receive) + weak_deadline: float | None = None + while True: + # Match GEAK's lifecycle contract: a task notification + # is authoritative over an on-disk marker. The Director + # writes its JSON before StructuredOutput returns and + # the JS Workflow assembles its final result. + if pending and not state["producer_done"]: + await anyio.sleep(max(0.1, done_poll_seconds)) + continue + if _terminal_artifact_exists(eval_dir): + break + if state["result_seen"] and not state["background_started"]: + break + completion_error = _completed_producer_error( + state, + pending, + producer_error, + ) + if completion_error is not None: + raise completion_error + + weak_terminal = ( + state["background_started"] + and state["result_seen"] + and not pending + and ( + state["terminal_task_seen"] + or state["producer_done"] + ) + ) + if weak_terminal and weak_deadline is None: + weak_deadline = ( + time.monotonic() + max(0.0, done_grace_seconds) + ) + if weak_deadline is not None and time.monotonic() >= weak_deadline: + break + if ( + state["producer_done"] + and not state["background_started"] + and not state["result_seen"] + ): + break + await anyio.sleep(max(0.1, done_poll_seconds)) + task_group.cancel_scope.cancel() + return "\n".join(chunks)[:_TRANSCRIPT_SIZE_LIMIT] + + return anyio.run(_run) + + +def _number(value: Any) -> float | None: + try: + parsed = float(value) + except (TypeError, ValueError): + return None + return parsed if math.isfinite(parsed) else None + + +def normalize_result( + eval_dir: Path, + workflow_return: dict[str, Any] | None = None, +) -> dict[str, Any]: + """Build the stable runner result from GEAK's authoritative artifacts.""" + disk_return_path = eval_dir / "workflow_return.json" + disk_return = _read_json(disk_return_path) + disk_return_present = disk_return_path.exists() or disk_return_path.is_symlink() + disk_return_valid = _valid_workflow_return( + disk_return, + eval_dir, + require_pinned_patch=True, + ) + workflow_contract_invalid = disk_return_present and not disk_return_valid + if disk_return_valid: + workflow_return = disk_return + elif not _valid_workflow_return( + workflow_return, + eval_dir, + require_pinned_patch=True, + ): + workflow_return = {} + validation = _read_json(eval_dir / "director_validation.json") or {} + + validation_status = str( + validation.get("validation_status") + or workflow_return.get("validation_status") + or "unknown" + ).lower() + correctness = str(validation.get("correctness") or "unknown").lower() + workload_aligned = workflow_return.get("workload_aligned") is True + weighted = _number(validation.get("director_verified_speedup_weighted")) + geomean = _number(validation.get("director_verified_speedup_geomean")) + workflow_speedup = _number(workflow_return.get("final_speedup")) + speedup = ( + weighted + if workload_aligned and weighted is not None + else (geomean if geomean is not None else workflow_speedup) + ) + + patch_path = eval_dir / "final_patch.diff" + patch_exists = patch_path.is_file() and patch_path.stat().st_size > 0 + + accepted = validation_status == "accepted" and correctness == "pass" + gained = speedup is not None and speedup > 1.0 + director_valid = _valid_director_validation(validation, eval_dir) + primary_metric_valid = speedup is not None + if ( + accepted + and gained + and patch_exists + and director_valid + and not workflow_contract_invalid + ): + status = "ok" + elif accepted and director_valid and workflow_contract_invalid: + status = "error" + elif accepted and director_valid and not primary_metric_valid: + status = "error" + elif accepted and director_valid and not gained: + status = "no_gain" + elif validation_status == "flagged" or correctness == "fail": + status = "rejected" + else: + status = "error" + + if not director_valid and validation_status in {"accepted", "flagged"}: + reason = "GEAK Director artifact is missing or invalid" + elif workflow_contract_invalid: + reason = "GEAK workflow return artifact is present but invalid" + elif not accepted: + reason = ( + f"GEAK validation did not accept the candidate " + f"(status={validation_status}, correctness={correctness})" + ) + elif not primary_metric_valid: + metric = "weighted" if workload_aligned else "geomean" + reason = f"GEAK Director artifact has no finite {metric} speedup" + elif not gained: + reason = f"GEAK did not verify a speedup above 1.0x (speedup={speedup})" + elif not patch_exists: + reason = f"GEAK accepted a gain but produced no non-empty patch at {patch_path}" + else: + reason = "" + + return { + "schema_version": SCHEMA_VERSION, + "status": status, + "eval_dir": str(eval_dir), + "validation_status": validation_status, + "correctness": correctness, + "workload_aligned": workload_aligned, + "final_speedup": speedup, + "final_geomean": geomean, + "final_weighted": weighted, + "final_patch": str(patch_path), + "director_final_patch": validation.get("final_patch"), + "workflow_final_patch": workflow_return.get("final_patch"), + "report_path": str( + workflow_return.get("report_path") + or eval_dir / "tech_lead_report.md" + ), + "budget_used": workflow_return.get("budget_used"), + "budget_total": workflow_return.get("budget_total"), + "applied_to_original": validation.get("applied_to_original", "unknown"), + "reason": reason, + } + + +def run_handoff(handoff: dict[str, Any]) -> dict[str, Any]: + script_path, workflow_args = map_workflow_args(handoff) + eval_dir = Path(workflow_args["eval_dir"]) + eval_dir.parent.mkdir(parents=True, exist_ok=True) + prompt = build_prompt(script_path, workflow_args) + + timeout_seconds = _positive_int( + handoff.get("timeout_seconds", 43200), + "timeout_seconds", + ) + model = str(handoff.get("model") or "claude-opus-4-8") + effort = str(handoff.get("effort") or "ultracode") + settings_value = handoff.get("settings", DEFAULT_SETTINGS) + settings = ( + settings_value + if isinstance(settings_value, str) + else json.dumps(settings_value) + ) + cli_path = str( + handoff.get("claude_cli_path") + or os.environ.get("GEAK_CLAUDE_BIN") + or shutil.which("claude") + or "" + ).strip() + if not cli_path: + raise RuntimeError("Claude Code CLI not found; cannot run GEAK Workflow") + done_grace = _nonnegative_float( + handoff.get("done_grace_seconds", 1800), + "done_grace_seconds", + ) + done_poll = _nonnegative_float( + handoff.get("done_poll_seconds", 5), + "done_poll_seconds", + ) + + run_directory_fd = _open_directory_fd(eval_dir.parent) + try: + transcript = "" + invocation_error: Exception | None = None + try: + transcript = invoke_via_sdk( + prompt, + workflow_dir=script_path.parent, + eval_dir=eval_dir, + model=model, + effort=effort, + settings=settings, + cli_path=cli_path, + timeout_seconds=timeout_seconds, + done_grace_seconds=done_grace, + done_poll_seconds=done_poll, + ) + except Exception as exc: # disk recovery below may still prove completion + invocation_error = exc + + parsed_return = _extract_workflow_return(transcript, eval_dir) + if parsed_return: + eval_directory_fd = _open_directory_fd( + eval_dir.name, + parent_fd=run_directory_fd, + ) + try: + try: + os.stat( + "workflow_return.json", + dir_fd=eval_directory_fd, + follow_symlinks=False, + ) + except FileNotFoundError: + _atomic_write_json( + eval_dir / "workflow_return.json", + parsed_return, + directory_fd=eval_directory_fd, + ) + finally: + os.close(eval_directory_fd) + + if _terminal_artifact_exists(eval_dir): + result = normalize_result(eval_dir, parsed_return) + if invocation_error: + result["recovered_after_error"] = type(invocation_error).__name__ + return result + if invocation_error: + raise invocation_error + raise RuntimeError( + "GEAK Workflow exited without workflow_return.json or " + f"director_validation.json under {eval_dir}" + ) + finally: + os.close(run_directory_fd) + + +def _dry_run_result( + handoff: dict[str, Any], + script_path: Path, + workflow_args: dict[str, Any], +) -> dict[str, Any]: + return { + "schema_version": SCHEMA_VERSION, + "status": "dry_run", + "script_path": str(script_path), + "workflow_args": workflow_args, + "prompt": build_prompt(script_path, workflow_args), + "model": str(handoff.get("model") or "claude-opus-4-8"), + "effort": str(handoff.get("effort") or "ultracode"), + } + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser(description="Run GEAK v4 kernel_workflow") + parser.add_argument("handoff", type=Path) + parser.add_argument("result", type=Path) + parser.add_argument( + "--dry-run", + action="store_true", + help="validate/map the handoff without importing the SDK or invoking Claude", + ) + namespace = parser.parse_args(argv) + + result_directory_fd = _open_directory_fd(namespace.result.parent) + try: + result: dict[str, Any] + try: + handoff = load_handoff(namespace.handoff) + if namespace.dry_run: + script_path, workflow_args = map_workflow_args(handoff) + result = _dry_run_result(handoff, script_path, workflow_args) + else: + result = run_handoff(handoff) + except Exception as exc: + result = { + "schema_version": SCHEMA_VERSION, + "status": "error", + "error_type": type(exc).__name__, + "error": str(exc), + } + _atomic_write_json( + namespace.result, + result, + directory_fd=result_directory_fd, + ) + print(json.dumps(result, ensure_ascii=False), flush=True) + return 1 + + _atomic_write_json( + namespace.result, + result, + directory_fd=result_directory_fd, + ) + print(json.dumps(result, ensure_ascii=False), flush=True) + return 1 if result.get("status") == "error" else 0 + finally: + os.close(result_directory_fd) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/docs/how-to/agents.md b/docs/how-to/agents.md index 0527c164..36f4b5d2 100644 --- a/docs/how-to/agents.md +++ b/docs/how-to/agents.md @@ -21,6 +21,7 @@ The following agents are available. | `cursor` | Cursor Agent CLI | | `claude_code` | Anthropic Claude Code CLI | | `codex` | OpenAI Codex CLI | +| `geak_v4` | GEAK v4 deterministic kernel Workflow through Claude Code (see [GEAK v4 setup](#geak-v4)) | | `geak_v3` | Specialized GEAK integration for HIP optimization | | `geak_v3_triton` | Specialized GEAK integration for Triton optimization | | `mini_swe_triton` | mini-swe-agent-based Triton optimization | @@ -37,8 +38,10 @@ Each agent lives under `agents//` and is registered into a shared registry, so the framework loads only the agent you select. The Cursor, Claude Code, and Codex integrations reuse their host CLI login -state. Specialized integrations have additional setup and configuration under -their respective `agents//` directories. +state. GEAK v4 also reuses the Claude Code login and adds a read-only GEAK +checkout plus the Claude Agent SDK. Specialized integrations have additional +setup and configuration under their respective `agents//` +directories. ## Models, providers, and agent settings @@ -48,19 +51,44 @@ effort, timeout, and iteration settings through its CLI and `agents//agent_config.yaml`. For Cursor, Claude Code, and Codex, authenticate with the host CLI. A normal run -preflights only the selected CLI. When the config selects one of these -first-class integrations (or `task_validator`), select the run configuration -first and check its CLI/backend: +preflights only the selected dependencies. When the config selects one of these +first-class integrations, GEAK v4, or `task_validator`, select the run +configuration first and check its CLI/backend: ```bash CONFIG_PATH=example_configs/quickstart_claude_mi300.yaml make docker-check-agents CONFIG="$CONFIG_PATH" ``` -Use `AGENTS=` for an explicit subset or `AGENTS=all` for -all three first-class CLIs and login states. Specialized integrations are not -handled by this command; their README files document their own dependencies, -API keys, and endpoint configuration. +Use `AGENTS=` for an explicit subset, +`AGENTS=geak_v4` for GEAK v4, or `AGENTS=all` for all three first-class CLIs +and login states. Legacy specialized integrations retain their own dependency +checks; their README files document their dependencies and endpoint +configuration. + +## GEAK v4 + +V1 of the GEAK v4 Arena integration supports `hip2hip`, `triton2triton`, and +`flydsl2flydsl` tasks that declare exactly one ordinary kernel source. A test or +harness source, a protected path, or multiple source files is rejected. + +Place the GEAK checkout beside AgentKernelArena, or set +`AKA_GEAK_ROOT=/absolute/path/to/GEAK`. Install and log in to Claude Code +2.1.177 or newer on the host, then run: + +```bash +make docker-setup-geak +CONFIG_PATH=example_configs/quickstart_geak_v4_mi300.yaml +make docker-check-agents CONFIG="$CONFIG_PATH" +make docker-run CONFIG="$CONFIG_PATH" +``` + +The setup target installs only `claude-agent-sdk`; it does not `pip install` +GEAK. The checkout is mounted read-only, while workflow artifacts are written +to a hidden directory beside the scored task workspace. Offline validation +does not invoke a real paid Claude workflow. See the +[repository agent README](https://github.com/AMD-AGI/AgentKernelArena/blob/main/agents/geak_v4/README.md) +for the full task and patch contract. `make vllm` starts an OpenAI-compatible local endpoint on port `30001`, but it does not automatically reconfigure an agent. Point the selected integration at diff --git a/docs/how-to/run-evaluation.md b/docs/how-to/run-evaluation.md index ddf0f36f..fe607631 100644 --- a/docs/how-to/run-evaluation.md +++ b/docs/how-to/run-evaluation.md @@ -15,12 +15,13 @@ resume, and inspect a run. ## Choose or create a run configuration A run configuration selects the agent, tasks, and target GPU. The repository -ships three examples: +ships four examples: | Configuration | Purpose | | --- | --- | | `example_configs/quickstart_claude_mi300.yaml` | One Claude Code GELU task on MI300/MI300X (`gfx942`). | | `example_configs/quickstart_claude_mi355x.yaml` | One Claude Code GELU task on MI355X (`gfx950`). | +| `example_configs/quickstart_geak_v4_mi300.yaml` | One GEAK v4 GELU task on MI300/MI300X (`gfx942`); complete its agent-specific setup first. | | `example_configs/benchmark_cursor_mi355x.yaml` | Curated 60-task Cursor Agent benchmark on MI355X; use only after installing and authenticating Cursor Agent. | For a first run, select the quickstart that matches the physical GPU: @@ -102,10 +103,14 @@ For debugging, enter the same Docker runtime used by the experiment: ```bash make docker-shell + +# Mount the composite GEAK v4 stack when debugging that integration: +make docker-shell AGENTS=geak_v4 ``` -The Docker runner currently supports Codex, Claude Code, and Cursor Agent login -reuse from the host. It preflights the selected config before starting the run. +The Docker runner supports Codex, Claude Code, and Cursor Agent login reuse from +the host. It also provisions the composite GEAK v4 stack when selected. It +preflights the selected config before starting the run. ## Run across multiple GPUs @@ -134,7 +139,8 @@ make docker-parallel-run \ ``` The Docker parallel path is verified for `cursor`, `claude_code`, `codex`, and -`task_validator`. Specialized GEAK/mini-swe templates require their own +`task_validator`. `geak_v4` maps each isolated worker to logical GPU 0 after +its agent-specific setup. Legacy GEAK/mini-swe templates require their own dependencies and worker-visible GPU configuration. See [Run tasks in parallel across multiple GPUs](parallel-run.md) for scheduling, GPU isolation, resume behavior, and failure handling. diff --git a/docs/install/install.md b/docs/install/install.md index 92748162..e6012845 100644 --- a/docs/install/install.md +++ b/docs/install/install.md @@ -87,19 +87,50 @@ installation and its alternative npm installation. See the for current installation alternatives. The `geak_v3`, `geak_v3_triton`, and `mini_swe_triton` integrations require -their own runtime dependencies. Review the corresponding directory under -`agents/` before selecting one. +their own runtime dependencies. GEAK v4 has a dedicated setup path described +below. Review the corresponding directory under `agents/` before selecting a +specialized integration. + +## Set up GEAK v4 (optional) + +GEAK v4 requires Claude Code 2.1.177 or newer, authenticated on the host, and a +local GEAK checkout. The Docker runner expects this sibling layout by default: + +```text +parent/ +├── AgentKernelArena/ +└── GEAK/ +``` + +Clone `https://github.com/AMD-AGI/GEAK.git` into that location, or export +`AKA_GEAK_ROOT=/absolute/path/to/GEAK`. Then install the Agent SDK into the +persistent container Python dependency directory: + +```bash +make docker-setup-geak +``` + +This target installs only `claude-agent-sdk`; it does not `pip install` GEAK. +The Docker runner mounts the GEAK checkout read-only at `/opt/geak`. V1 of the +Arena integration supports only `hip2hip`, `triton2triton`, and +`flydsl2flydsl` tasks with exactly one non-test, non-harness source file. +Workflow artifacts are written in a hidden sibling directory outside the +scored task workspace. + +See the [GEAK v4 agent guidance](../how-to/agents.md#geak-v4) for the complete +setup flow. Its offline tests do not make a real paid Claude workflow call. ## Choose an example configuration Choose the configuration that matches the physical GPU and installed agent. -The two quickstart configurations each run one GELU task; the benchmark +The three quickstart configurations each run one GELU task; the benchmark configuration is a longer 60-task Cursor Agent run. | Configuration | Purpose | | --- | --- | | `example_configs/quickstart_claude_mi300.yaml` | First Claude Code run on MI300/MI300X (`gfx942`). | | `example_configs/quickstart_claude_mi355x.yaml` | First Claude Code run on MI355X (`gfx950`). | +| `example_configs/quickstart_geak_v4_mi300.yaml` | First GEAK v4 run on MI300/MI300X (`gfx942`); requires the GEAK v4 setup above. | | `example_configs/benchmark_cursor_mi355x.yaml` | Curated 60-task Cursor Agent benchmark on MI355X; requires an installed and authenticated Cursor Agent CLI. | The default `make docker-run` configuration is the MI300/MI300X quickstart. @@ -135,7 +166,7 @@ cp "$CONFIG_PATH" my_experiment.yaml `flydsl2flydsl`, `torch2flydsl`, and `triton2flydsl` tasks need the `flydsl` package inside the container. The selected image may already ship it (`make docker-smoke` prints `flydsl=ok ` when present). If yours does -not, install it once into the container's persistent pip user-base: +not, install it once into the container's persistent dependency directory: ```bash make docker-setup-flydsl @@ -146,23 +177,25 @@ This is a no-op when the image already provides FlyDSL. ## Configure authentication and providers Cursor, Claude Code, and Codex reuse their host CLI authentication. A normal run -preflights only its selected CLI. For another config that selects one of these -CLIs—or `task_validator`, which resolves to its configured backend—check the -same CLI without starting a task by passing that run config: +preflights only its selected dependencies. For another config that selects one +of these CLIs, GEAK v4, or `task_validator`, which resolves to its configured +backend, check the same stack without starting a task by passing that run +config: ```bash make docker-check-agents CONFIG=my_experiment.yaml # Optional overrides: make docker-check-agents AGENTS=claude_code,codex +make docker-check-agents AGENTS=geak_v4 make docker-check-agents AGENTS=all ``` `AGENTS=all` is the explicit strict check for Cursor, Claude Code, and Codex. -Specialized integrations such as GEAK and mini-swe use their own dependency and -authentication checks. They read credentials and provider endpoints from their -own environment/configuration; there is no shared provider field in the root -run configuration. +GEAK v4 has its own composite check for Claude, the Agent SDK, the read-only +workflow checkout, and a profiler. Legacy GEAK and mini-swe integrations use +their own dependency and authentication checks. There is no shared provider +field in the root run configuration. To run against a self-hosted model instead of a hosted provider, start a local vLLM server: diff --git a/docs/reference/api-reference.md b/docs/reference/api-reference.md index 7b430567..87635ead 100644 --- a/docs/reference/api-reference.md +++ b/docs/reference/api-reference.md @@ -90,7 +90,8 @@ The following Make targets are available for running experiments. | `make docker-run CONFIG=example_configs/quickstart_claude_mi300.yaml` | Run tasks serially in one Docker container | | `make docker-parallel-run CONFIG=example_configs/benchmark_cursor_mi355x.yaml GPU_IDS=0,1` | Run one Docker worker per listed GPU, using a shared dynamic task queue | | `make docker-smoke` | Verify Docker, ROCm runtime visibility, Python imports, and GPU access | -| `make docker-check-agents CONFIG=example_configs/quickstart_claude_mi300.yaml` | Verify the first-class host CLI selected by the config inside Docker (`task_validator` resolves to its backend). Override with `AGENTS=claude_code,codex`; use `AGENTS=all` for all three. Specialized integrations use their own checks | +| `make docker-check-agents CONFIG=example_configs/quickstart_claude_mi300.yaml` | Verify the agent stack selected by the config inside Docker (`task_validator` resolves to its backend). Override with `AGENTS=claude_code,codex` or `AGENTS=geak_v4`; `AGENTS=all` checks the three host CLIs. | +| `make docker-setup-geak` | Install only `claude-agent-sdk` into the persistent container dependency directory; GEAK itself is mounted read-only from `AKA_GEAK_ROOT`. | | `make docker-shell` | Open an interactive shell in the experiment runtime | `docker-parallel-run` accepts these environment variables: diff --git a/docs/reference/compatibility-matrix.md b/docs/reference/compatibility-matrix.md index c5225bcd..d9946c00 100644 --- a/docs/reference/compatibility-matrix.md +++ b/docs/reference/compatibility-matrix.md @@ -36,6 +36,7 @@ The following software versions are required or verified. | Triton | Bundled with the image's ROCm PyTorch | Required for Triton task categories. | | AITER | `0.1.17.dev110+g9127c94a1` in the verified `gfx950` image | Required by AITER-backed task oracles and kernels. | | FlyDSL | `0.2.2` in the verified `gfx950` image (or `make docker-setup-flydsl` when absent) | Required for `flydsl2flydsl`, `torch2flydsl`, and `triton2flydsl` tasks. | +| Claude Code for GEAK v4 | 2.1.177 or newer | Must be authenticated on the host; GEAK's dynamic Workflow also requires `claude-agent-sdk`. | ## Agents @@ -48,6 +49,7 @@ The following templates are selectable in the current `AgentType` registry. See | `cursor` | Cursor Agent CLI and host login state. | | `claude_code` | Native/local or npm-installed Claude Code CLI and host login state. | | `codex` | Codex CLI and host login state. | +| `geak_v4` | Claude Code 2.1.177+, host login state, `claude-agent-sdk`, a read-only local GEAK checkout, and a supported profiler. | | `geak_v3` | GEAK CLI; HIP-oriented integration. | | `geak_v3_triton` | GEAK CLI; Triton-oriented integration. | | `mini_swe_triton` | mini-swe-agent/GEAK dependencies. | diff --git a/docs/reference/release-notes.md b/docs/reference/release-notes.md index 6d229c38..8d2683e4 100644 --- a/docs/reference/release-notes.md +++ b/docs/reference/release-notes.md @@ -71,12 +71,18 @@ The supported agent templates are now: - `claude_code` - `codex` - `cursor` +- `geak_v4` - `geak_v3` - `geak_v3_triton` - `mini_swe_triton` - `task_validator` -The task validator now includes Codex backend support, repository-task validation, improved Python-environment propagation, stronger source and target checks, starter-stub detection, and standardized validation reports. +The GEAK v4 integration runs the deterministic kernel Workflow through Claude +Code against a disposable task copy and imports only a Director-validated, +single-source patch. The task validator now includes Codex backend support, +repository-task validation, improved Python-environment propagation, stronger +source and target checks, starter-stub detection, and standardized validation +reports. #### Documentation and onboarding diff --git a/example_configs/quickstart_geak_v4_mi300.yaml b/example_configs/quickstart_geak_v4_mi300.yaml new file mode 100644 index 00000000..fcc622c6 --- /dev/null +++ b/example_configs/quickstart_geak_v4_mi300.yaml @@ -0,0 +1,11 @@ +# Minimal first-run example for GEAK v4 on MI300/MI300X (gfx942). +# Complete agents/geak_v4/README.md setup before running it. +agent: + template: geak_v4 + +tasks: + - hip2hip/gpumode/GELU + +target_gpu_model: MI300 +log_directory: logs +workspace_directory_prefix: workspace diff --git a/src/harness_guard.py b/src/harness_guard.py index da475af0..7485e953 100644 --- a/src/harness_guard.py +++ b/src/harness_guard.py @@ -53,6 +53,16 @@ def _is_protected_path(rel: Path) -> bool: return name.endswith(_HARNESS_FILE_SUFFIXES) +def is_protected_workspace_path(path: str | Path) -> bool: + """Return whether a workspace-relative path belongs to the task harness. + + Agent integrations that import patches can use this public predicate before + touching the workspace. The final digest check remains authoritative, but + rejecting protected paths up front avoids partially applying an unsafe patch. + """ + return _is_protected_path(Path(path)) + + def _iter_protected_files(root: Path) -> Iterable[Path]: for path in root.rglob("*"): if not path.is_file(): diff --git a/src/module_registration.py b/src/module_registration.py index d6000820..01eb3b14 100755 --- a/src/module_registration.py +++ b/src/module_registration.py @@ -11,6 +11,7 @@ class AgentType(Enum): CLAUDE_CODE = "claude_code" CODEX = "codex" TASK_VALIDATOR = "task_validator" + GEAK_V4 = "geak_v4" GEAK_V3 = "geak_v3" GEAK_V3_TRITON = "geak_v3_triton" MINI_SWE_TRITON = "mini_swe_triton" @@ -66,6 +67,8 @@ def load_agent_launcher(agent_type: AgentType, logger: logging.Logger) -> Callab from agents.codex import launch_agent # noqa: F401 elif agent_type == AgentType.TASK_VALIDATOR: from agents.task_validator import launch_agent # noqa: F401 + elif agent_type == AgentType.GEAK_V4: + from agents.geak_v4 import launch_agent # noqa: F401 elif agent_type == AgentType.GEAK_V3: from agents.geak_v3 import launch_agent # noqa: F401 elif agent_type == AgentType.GEAK_V3_TRITON: @@ -109,7 +112,7 @@ def load_post_processing_handler(agent_type: AgentType, logger: logging.Logger) from agents.task_validator.validation_postprocessing import validation_post_processing logger.info(f"Using validation_post_processing for agent: {agent_name}") return validation_post_processing - elif agent_type in [AgentType.CURSOR, AgentType.CLAUDE_CODE, AgentType.CODEX, AgentType.GEAK_V3, AgentType.GEAK_V3_TRITON, AgentType.MINI_SWE_TRITON, AgentType.FORGE]: + elif agent_type in [AgentType.CURSOR, AgentType.CLAUDE_CODE, AgentType.CODEX, AgentType.GEAK_V4, AgentType.GEAK_V3, AgentType.GEAK_V3_TRITON, AgentType.MINI_SWE_TRITON, AgentType.FORGE]: logger.info(f"Using general_post_processing for agent: {agent_name}") return general_post_processing else: @@ -135,8 +138,8 @@ def load_prompt_builder(agent_type: AgentType, logger: logging.Logger) -> Callab agent_name = agent_type.value # Map agents to their prompt builder functions - if agent_type in [AgentType.CURSOR, AgentType.CLAUDE_CODE, AgentType.CODEX]: + if agent_type in [AgentType.CURSOR, AgentType.CLAUDE_CODE, AgentType.CODEX, AgentType.GEAK_V4]: logger.info(f"Using standard prompt_builder for agent: {agent_name}") return prompt_builder else: - raise NotImplementedError(f"Prompt builder not implemented for agent: {agent_name}") \ No newline at end of file + raise NotImplementedError(f"Prompt builder not implemented for agent: {agent_name}") diff --git a/src/scripts/docker_benchmark.sh b/src/scripts/docker_benchmark.sh index fe601cb3..95308e6e 100755 --- a/src/scripts/docker_benchmark.sh +++ b/src/scripts/docker_benchmark.sh @@ -7,6 +7,9 @@ DEFAULT_DOCKER_IMAGE_GFX950="${AKA_DOCKER_IMAGE_GFX950:-$GFX950_V0514_DOCKER_IMA CONTAINER_WORKDIR="${AKA_DOCKER_WORKDIR:-/workspace}" HOST_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)" HOST_HOME="${HOME:?HOME must be set}" +HOST_GEAK_ROOT="${AKA_GEAK_ROOT:-$(dirname "$HOST_ROOT")/GEAK}" +CONTAINER_GEAK_ROOT="/opt/geak" +CONTAINER_PYTHON_DEPS="${CONTAINER_WORKDIR}/.aka-pyuserbase/site-packages" HOST_UID="$(id -u)" HOST_GID="$(id -g)" SELECTED_GPU_ARCH="" @@ -28,6 +31,7 @@ Usage: src/scripts/docker_benchmark.sh preflight [--config_name ] src/scripts/docker_benchmark.sh shell src/scripts/docker_benchmark.sh check-agents [--config_name ] + src/scripts/docker_benchmark.sh setup-geak src/scripts/docker_benchmark.sh smoke Default run config: @@ -43,7 +47,8 @@ Environment overrides: AKA_DOCKER_IMAGE_GFX942 Default image for gfx942. AKA_DOCKER_IMAGE_GFX950 Default image for gfx950. AKA_NODE_PREFIX Host Node prefix containing bin/node and npm-installed agent CLI(s). - AKA_AGENTS Agent CLI(s) to check, comma/space separated; use all for all three. + AKA_GEAK_ROOT Host GEAK checkout (default: sibling GEAK directory). + AKA_AGENTS Agent stack(s) to check, comma/space separated; `all` checks the three host CLIs. EOF } @@ -309,6 +314,7 @@ resolve_required_agents() { [[ "$tmpl" == "task_validator" ]] && tmpl="$(read_validator_backend)" case "$tmpl" in claude|claude_code) printf 'claude_code\n' ;; + geak_v4) printf 'geak_v4\n' ;; cursor|cursor-agent) printf 'cursor\n' ;; codex) printf 'codex\n' ;; *) printf '%s\n' "$tmpl" ;; @@ -330,6 +336,9 @@ normalize_check_agents() { claude|claude_code) normalized+=(claude_code) ;; + geak_v4) + normalized+=(geak_v4) + ;; cursor|cursor-agent) normalized+=(cursor) ;; @@ -337,7 +346,7 @@ normalize_check_agents() { normalized+=(codex) ;; *) - die "docker-check-agents only supports codex, claude_code, cursor, or all; got '$agent'" + die "docker-check-agents only supports codex, claude_code, cursor, geak_v4, or all; got '$agent'" ;; esac done @@ -396,6 +405,29 @@ mount_agent() { add_mount "$HOST_HOME/.claude.json" "$HOST_HOME/.claude.json" fi ;; + geak_v4) + # GEAK v4 is a composite integration: it needs the normal mounted + # Claude CLI/login plus a read-only GEAK Workflow checkout. + mount_agent claude_code "$strict" + need_path \ + "$HOST_GEAK_ROOT/kernel_workflow/kernel_workflow.js" \ + "GEAK v4 kernel workflow (set AKA_GEAK_ROOT)" \ + "$strict" || return 0 + need_path \ + "$HOST_GEAK_ROOT/kernel_workflow/roles" \ + "GEAK v4 workflow roles" \ + "$strict" || return 0 + need_path \ + "$HOST_GEAK_ROOT/kernel_workflow/knowledge" \ + "GEAK v4 workflow knowledge" \ + "$strict" || return 0 + need_path \ + "$HOST_GEAK_ROOT/kernel_workflow/scripts/gpu_lock.sh" \ + "GEAK v4 GPU lock helper" \ + "$strict" || return 0 + add_mount "$HOST_GEAK_ROOT" "$CONTAINER_GEAK_ROOT" ro + docker_args+=(-e "GEAK_V4_ROOT=${CONTAINER_GEAK_ROOT}") + ;; cursor) need_path "$HOST_HOME/.local/bin" "host local bin directory" "$strict" || return 0 need_path "$HOST_HOME/.local/share/cursor-agent" "Cursor Agent local install" "$strict" || return 0 @@ -469,6 +501,7 @@ build_docker_args() { -e "TORCH_EXTENSIONS_DIR=/tmp/torch-extensions${cache_postfix}" -e "TRITON_CACHE_DIR=/tmp/triton-cache${cache_postfix}" -e "PYTHONUSERBASE=${CONTAINER_WORKDIR}/.aka-pyuserbase" + -e "PYTHONPATH=${CONTAINER_PYTHON_DEPS}" -e "MIOPEN_USER_DB_PATH=/tmp/miopen-cache${cache_postfix}" -e "MIOPEN_CACHE_DIR=/tmp/miopen-cache${cache_postfix}" -e "MIOPEN_CUSTOM_CACHE_DIR=/tmp/miopen-cache${cache_postfix}" @@ -527,11 +560,24 @@ build_docker_args() { add_device_if_present /dev/mem add_mount "$HOST_ROOT" "$CONTAINER_WORKDIR" - # Persistent pip user-base (PYTHONUSERBASE) so `make docker-setup-flydsl` survives - # across runs. It lives INSIDE the repo dir, which is already bind-mounted above and - # is owned by the host user — this avoids a separate mount whose source the docker - # daemon would have to create (which fails on NFS/root-squashed homes). - mkdir -p "$HOST_ROOT/.aka-pyuserbase" 2>/dev/null || true + # Persistent dependency target so optional Python packages survive across + # runs. The image's /opt/venv disables user-site packages, so setup targets + # use pip --target and this explicit PYTHONPATH instead of pip --user. + # It lives inside the repo's existing bind mount and remains host-user-owned. + mkdir -p "$HOST_ROOT/.aka-pyuserbase/site-packages" 2>/dev/null || true + if [[ "${PYTHON_DEPS_WRITABLE:-0}" == "1" ]]; then + add_mount \ + "$HOST_ROOT/.aka-pyuserbase" \ + "${CONTAINER_WORKDIR}/.aka-pyuserbase" + else + # Override the parent repo's RW bind at this nested path. Agent code is + # imported from PYTHONPATH, so allowing Workflow to write here would + # enable sitecustomize-based benchmark/harness injection. + add_mount \ + "$HOST_ROOT/.aka-pyuserbase" \ + "${CONTAINER_WORKDIR}/.aka-pyuserbase" \ + ro + fi local _agent for _agent in $agents; do mount_agent "$_agent" "$strict" @@ -657,7 +703,7 @@ if "codex" in agents: codex_line = next((line for line in codex_status.splitlines() if "Logged in" in line), codex_status.splitlines()[-1]) print(f"codex_status={codex_line}") -if "claude_code" in agents: +if "claude_code" in agents or "geak_v4" in agents: require_cmd("claude") claude_version = run_checked(["claude", "--version"]).splitlines()[-1] claude_status_raw = run_checked(["claude", "auth", "status"]) @@ -671,6 +717,52 @@ if "claude_code" in agents: f"version={claude_version}" ) +if "geak_v4" in agents: + import inspect + import pathlib + import re + + version_match = re.search(r"(\d+)\.(\d+)\.(\d+)", claude_version) + if not version_match or tuple(map(int, version_match.groups())) < (2, 1, 177): + raise SystemExit( + f"GEAK v4 requires Claude Code >=2.1.177; found {claude_version}" + ) + try: + from claude_agent_sdk import ClaudeAgentOptions, ClaudeSDKClient + except (ImportError, AttributeError) as exc: + raise SystemExit( + "claude_agent_sdk lacks the persistent-client API; " + "run make docker-setup-geak" + ) from exc + if "cli_path" not in inspect.signature(ClaudeAgentOptions).parameters: + raise SystemExit( + "claude_agent_sdk lacks ClaudeAgentOptions.cli_path; " + "run make docker-setup-geak" + ) + + geak_root = pathlib.Path(os.environ.get("GEAK_V4_ROOT", "/opt/geak")) + required = [ + (geak_root / "kernel_workflow" / "kernel_workflow.js", "file"), + (geak_root / "kernel_workflow" / "roles", "directory"), + (geak_root / "kernel_workflow" / "knowledge", "directory"), + (geak_root / "kernel_workflow" / "scripts" / "gpu_lock.sh", "file"), + ] + invalid = [ + f"{path} (expected {kind})" + for path, kind in required + if ( + path.is_symlink() + or (kind == "file" and not path.is_file()) + or (kind == "directory" and not path.is_dir()) + ) + ] + if invalid: + raise SystemExit("GEAK v4 checkout is incomplete: " + ", ".join(invalid)) + # Arena's common smoke/run contract requires rocprof-compute, so the + # agent-only preflight must enforce the same dependency. + require_cmd("rocprof-compute") + print(f"geak_v4=ok root={geak_root} sdk=claude_agent_sdk") + if "cursor" in agents: require_cmd("cursor-agent") cursor_version = run_checked(["cursor-agent", "--version"]).splitlines()[-1] @@ -710,19 +802,32 @@ PY } container_setup_flydsl() { - # If the image already provides FlyDSL, do nothing — installing a --user copy + # If the image already provides FlyDSL, do nothing — installing a target copy # could shadow the image version with an incompatible one. if python -c 'import flydsl' 2>/dev/null; then python -c 'import flydsl; print("flydsl already provided by image: " + str(getattr(flydsl, "__version__", "unknown")) + "; nothing to install")' return 0 fi - # Otherwise install into the persistent pip user-base (PYTHONUSERBASE), a - # host-mounted dir, so it survives the --rm container and is importable in later runs. - echo "flydsl not found in image; installing into persistent pip user-base..." - python -m pip install --user --upgrade flydsl + # /opt/venv disables user-site packages. Install into the explicit target + # already placed on PYTHONPATH so it survives the --rm container. + local target="${PYTHONUSERBASE:?PYTHONUSERBASE must be set}/site-packages" + mkdir -p "$target" + echo "flydsl not found in image; installing into persistent dependency target..." + python -m pip install --upgrade --target "$target" flydsl python -c 'import flydsl; print("flydsl=" + str(getattr(flydsl, "__version__", "unknown")) + " setup OK")' } +container_setup_geak() { + # The Workflow sources are used directly from the read-only /opt/geak + # checkout. Only install the SDK lifecycle dependency; never pip-install + # GEAK itself (its package bootstrap performs unrelated host-side setup). + local target="${PYTHONUSERBASE:?PYTHONUSERBASE must be set}/site-packages" + mkdir -p "$target" + echo "Installing/upgrading claude-agent-sdk in the persistent dependency target..." + python -m pip install --upgrade --target "$target" claude-agent-sdk + python -c 'import inspect; from claude_agent_sdk import ClaudeAgentOptions, ClaudeSDKClient; assert "cli_path" in inspect.signature(ClaudeAgentOptions).parameters; print("claude-agent-sdk persistent client setup OK")' +} + container_prepare_worker_home() { local state_root="${AGENT_STATE_MOUNT_ROOT:-/opt/aka-agent-state}" mkdir -p "$HOME" @@ -1013,11 +1118,22 @@ case "${1:-}" in # FlyDSL install needs no agent CLIs. REQUIRED_AGENTS="" AGENTS_STRICT=0 + PYTHON_DEPS_WRITABLE=1 docker_exec 0 bash src/scripts/docker_benchmark.sh _container_setup_flydsl ;; + setup-geak) + select_runtime_for_host + REQUIRED_AGENTS="geak_v4" + AGENTS_STRICT=1 + PYTHON_DEPS_WRITABLE=1 + docker_exec 0 bash src/scripts/docker_benchmark.sh _container_setup_geak + ;; _container_setup_flydsl) container_setup_flydsl ;; + _container_setup_geak) + container_setup_geak + ;; _container_smoke) container_smoke ;; diff --git a/tests/test_docker_benchmark.sh b/tests/test_docker_benchmark.sh index d68e2ffd..f8611749 100755 --- a/tests/test_docker_benchmark.sh +++ b/tests/test_docker_benchmark.sh @@ -57,6 +57,16 @@ run_check_args() { bash "$RUNNER" check-agents --config_name "$config" 2>/dev/null } +run_setup_geak_args() { + local home="$1" + local geak_root="$2" + env \ + HOME="$home" \ + AKA_GPU_ARCH=gfx950 \ + AKA_GEAK_ROOT="$geak_root" \ + bash "$RUNNER" setup-geak 2>/dev/null +} + assert_cache_args_present() { local suffix="$1" shift @@ -156,6 +166,76 @@ assert_has "$NATIVE_CLAUDE_HOME/.claude:$NATIVE_CLAUDE_HOME/.claude" "${args[@]} assert_has "$NATIVE_CLAUDE_HOME/.claude.json:$NATIVE_CLAUDE_HOME/.claude.json" "${args[@]}" assert_has "claude_code" "${args[@]}" +# GEAK v4 composes the native Claude mount with a read-only workflow checkout. +# The fake Docker captures argv only; no daemon, network, Claude, or GEAK code runs. +GEAK_ROOT="$TEST_HOME/geak" +GEAK_CONFIG="$TEST_HOME/geak-config.yaml" +mkdir -p \ + "$GEAK_ROOT/kernel_workflow/roles" \ + "$GEAK_ROOT/kernel_workflow/knowledge" \ + "$GEAK_ROOT/kernel_workflow/scripts" +touch \ + "$GEAK_ROOT/kernel_workflow/kernel_workflow.js" \ + "$GEAK_ROOT/kernel_workflow/scripts/gpu_lock.sh" +printf 'agent:\n template: geak_v4\n' > "$GEAK_CONFIG" + +mapfile -t args < <(run_check_args \ + "$NATIVE_CLAUDE_HOME" \ + "$GEAK_CONFIG" \ + AKA_GEAK_ROOT="$GEAK_ROOT") +assert_has "$NATIVE_CLAUDE_HOME/.local/bin:$NATIVE_CLAUDE_HOME/.local/bin:ro" "${args[@]}" +assert_has "$NATIVE_CLAUDE_HOME/.local/share/claude:$NATIVE_CLAUDE_HOME/.local/share/claude:ro" "${args[@]}" +assert_has "$NATIVE_CLAUDE_HOME/.claude:$NATIVE_CLAUDE_HOME/.claude" "${args[@]}" +assert_has "$NATIVE_CLAUDE_HOME/.claude.json:$NATIVE_CLAUDE_HOME/.claude.json" "${args[@]}" +assert_has "$GEAK_ROOT:/opt/geak:ro" "${args[@]}" +assert_has "GEAK_V4_ROOT=/opt/geak" "${args[@]}" +assert_has "PYTHONPATH=/workspace/.aka-pyuserbase/site-packages" "${args[@]}" +assert_has "$ROOT/.aka-pyuserbase:/workspace/.aka-pyuserbase:ro" "${args[@]}" +assert_has "_container_check_agents" "${args[@]}" +assert_has "geak_v4" "${args[@]}" +assert_not_has "claude_code" "${args[@]}" + +# An explicitly selected GEAK agent is strict about the host checkout. +if run_check_args \ + "$NATIVE_CLAUDE_HOME" \ + "$GEAK_CONFIG" \ + AKA_GEAK_ROOT="$TEST_HOME/missing-geak" >/dev/null; then + fail "GEAK agent check unexpectedly accepted a missing AKA_GEAK_ROOT" +fi + +# setup-geak uses the same mounts and routes into the SDK-only container setup. +mapfile -t args < <(run_setup_geak_args "$NATIVE_CLAUDE_HOME" "$GEAK_ROOT") +assert_has "$GEAK_ROOT:/opt/geak:ro" "${args[@]}" +assert_has "GEAK_V4_ROOT=/opt/geak" "${args[@]}" +assert_has "$ROOT/.aka-pyuserbase:/workspace/.aka-pyuserbase" "${args[@]}" +assert_not_has "$ROOT/.aka-pyuserbase:/workspace/.aka-pyuserbase:ro" "${args[@]}" +assert_has "_container_setup_geak" "${args[@]}" +assert_not_has "_container_check_agents" "${args[@]}" + +# The runtime venv disables user-site packages. Verify setup uses the persistent +# PYTHONPATH target instead of pip --user, without contacting a package index. +SETUP_FAKE_BIN="$TEST_HOME/setup-fake-bin" +SETUP_CALLS="$TEST_HOME/setup-python-calls" +mkdir -p "$SETUP_FAKE_BIN" +printf '%s\n' \ + '#!/usr/bin/env bash' \ + 'if [[ "$1" == "-c" && "$2" == "import claude_agent_sdk" ]]; then exit 1; fi' \ + 'printf "%s\n" "$*" >> "$AKA_SETUP_LOG"' \ + 'exit 0' \ + > "$SETUP_FAKE_BIN/python" +chmod +x "$SETUP_FAKE_BIN/python" +env \ + PATH="$SETUP_FAKE_BIN:$PATH" \ + AKA_SETUP_LOG="$SETUP_CALLS" \ + PYTHONUSERBASE="$TEST_HOME/python-deps" \ + bash "$RUNNER" _container_setup_geak >/dev/null +mapfile -t setup_calls < "$SETUP_CALLS" +assert_has \ + "-m pip install --upgrade --target $TEST_HOME/python-deps/site-packages claude-agent-sdk" \ + "${setup_calls[@]}" +[[ "$(tr '\n' ' ' < "$SETUP_CALLS")" != *"--user"* ]] \ + || fail "GEAK setup unexpectedly used pip --user" + # Omitting --config_name uses the one-task MI300/MI300X Claude quickstart. mapfile -t args < <( env \ diff --git a/tests/test_geak_v4.py b/tests/test_geak_v4.py new file mode 100644 index 00000000..d6d40ea4 --- /dev/null +++ b/tests/test_geak_v4.py @@ -0,0 +1,835 @@ +"""Offline tests for the GEAK v4 Arena adapter. + +These tests exercise handoff validation, result recovery, GPU namespace +mapping, and the single-file patch import boundary. They intentionally do not +invoke Claude, the Claude Agent SDK, GEAK, a container, or a GPU. +""" + +from __future__ import annotations + +import importlib +import json +import logging +import subprocess +from pathlib import Path, PurePosixPath +from typing import Callable + +import pytest + +from agents.geak_v4 import workflow_runner +from src.module_registration import AgentType, load_agent_launcher + + +geak_launcher = importlib.import_module("agents.geak_v4.launch_agent") + + +SOURCE = PurePosixPath("src/kernel.py") +ORIGINAL_SOURCE = "value = 1\nsecond = 2\nthird = 3\n" +OPTIMIZED_SOURCE = "value = 4\nsecond = 5\nthird = 6\n" + + +def _write_json(path: Path, value: object) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(value) + "\n", encoding="utf-8") + + +def _handoff(tmp_path: Path) -> tuple[dict[str, object], Path, Path]: + kernel = tmp_path / "kernel" + workflow_dir = tmp_path / "geak" / "kernel_workflow" + eval_dir = tmp_path / "artifacts" / "eval" + kernel.mkdir(parents=True) + workflow_dir.mkdir(parents=True) + (workflow_dir / "kernel_workflow.js").write_text( + "// offline fixture\n", + encoding="utf-8", + ) + handoff: dict[str, object] = { + "schema_version": workflow_runner.SCHEMA_VERSION, + "kernel_path": str(kernel), + "workflow_dir": str(workflow_dir), + "eval_dir": str(eval_dir), + "exp_root": str(tmp_path / "artifacts" / "runs"), + "gpu_ids": "7", + "budget": 3, + "min_improve": 0.03, + "deep_cost": 1, + # These untrusted values must never override the Arena policy. + "mode": "author", + "apply_to_original": True, + } + return handoff, kernel, eval_dir + + +def _workflow_return( + eval_dir: Path, + *, + workload_aligned: bool = False, +) -> dict[str, object]: + return { + "eval_dir": str(eval_dir), + "validation_status": "accepted", + "final_geomean": 1.20, + "final_speedup": 1.19, + "final_patch": str(eval_dir / "final_patch.diff"), + "workload_aligned": workload_aligned, + } + + +def _director_validation( + eval_dir: Path, + *, + validation_status: str = "accepted", + correctness: str = "pass", + geomean: object = 1.20, + weighted: object = 1.50, +) -> dict[str, object]: + return { + "validation_status": validation_status, + "correctness": correctness, + "director_verified_speedup_geomean": geomean, + "director_verified_speedup_weighted": weighted, + "final_patch": str(eval_dir / "final_patch.diff"), + "applied_to_original": "false", + } + + +def _prepare_normalized_result( + tmp_path: Path, + *, + workload_aligned: bool = False, + validation_status: str = "accepted", + correctness: str = "pass", + geomean: object = 1.20, + weighted: object = 1.50, +) -> tuple[Path, dict[str, object]]: + eval_dir = tmp_path / "eval" + eval_dir.mkdir() + _write_json( + eval_dir / "workflow_return.json", + _workflow_return(eval_dir, workload_aligned=workload_aligned), + ) + _write_json( + eval_dir / "director_validation.json", + _director_validation( + eval_dir, + validation_status=validation_status, + correctness=correctness, + geomean=geomean, + weighted=weighted, + ), + ) + (eval_dir / "final_patch.diff").write_text( + "non-empty offline fixture\n", + encoding="utf-8", + ) + return eval_dir, workflow_runner.normalize_result(eval_dir) + + +def _workspace(tmp_path: Path) -> Path: + workspace = tmp_path / "workspace" + source = workspace.joinpath(*SOURCE.parts) + source.parent.mkdir(parents=True) + source.write_text(ORIGINAL_SOURCE, encoding="utf-8") + (workspace / "config.yaml").write_text( + "task_type: triton2triton\n", + encoding="utf-8", + ) + scripts = workspace / "scripts" + scripts.mkdir() + (scripts / "test_kernel.py").write_text( + "def test_kernel(): pass\n", + encoding="utf-8", + ) + return workspace + + +def _git(repo: Path, *args: str) -> subprocess.CompletedProcess[str]: + return subprocess.run( + ["git", *args], + cwd=repo, + check=True, + text=True, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + ) + + +def _write_repo_file(repo: Path, relative: str, content: str | bytes) -> None: + path = repo / relative + path.parent.mkdir(parents=True, exist_ok=True) + if isinstance(content, bytes): + path.write_bytes(content) + else: + path.write_text(content, encoding="utf-8") + + +def _make_git_patch( + tmp_path: Path, + baseline: dict[str, str | bytes], + mutate: Callable[[Path], None], + *, + find_renames: bool = False, +) -> bytes: + repo = tmp_path / "patch_repo" + repo.mkdir() + _git(repo, "init", "-q") + for relative, content in baseline.items(): + _write_repo_file(repo, relative, content) + _git(repo, "add", "-A") + _git( + repo, + "-c", + "user.name=Offline Test", + "-c", + "user.email=offline@example.invalid", + "commit", + "--allow-empty", + "-qm", + "baseline", + ) + + mutate(repo) + _git(repo, "add", "-A") + command = ["git", "diff", "--cached", "--binary", "--no-ext-diff"] + if find_renames: + command.append("--find-renames") + command.append("HEAD") + diff = subprocess.run( + command, + cwd=repo, + check=True, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + ).stdout + assert diff + return diff + + +def _import_result(eval_dir: Path, *, status: str = "ok") -> dict[str, object]: + patch = str(eval_dir / "final_patch.diff") + return { + "schema_version": workflow_runner.SCHEMA_VERSION, + "status": status, + "validation_status": "accepted", + "correctness": "pass", + "applied_to_original": "false", + "final_speedup": 1.10, + "eval_dir": str(eval_dir), + "final_patch": patch, + "director_final_patch": patch, + "workflow_final_patch": patch, + } + + +def _apply_patch( + tmp_path: Path, + patch: bytes, + *, + status: str = "ok", +) -> tuple[bool, Path, Path]: + workspace = _workspace(tmp_path) + eval_dir = tmp_path / "eval" + run_dir = tmp_path / "run" + eval_dir.mkdir() + run_dir.mkdir() + (eval_dir / "final_patch.diff").write_bytes(patch) + applied = geak_launcher._apply_validated_patch( + result=_import_result(eval_dir, status=status), + expected_eval_dir=eval_dir, + workspace=workspace, + source_path=SOURCE, + min_improve=0.02, + run_dir=run_dir, + ) + return applied, workspace, run_dir + + +def _workspace_contents(workspace: Path) -> dict[str, bytes]: + return { + str(path.relative_to(workspace)): path.read_bytes() + for path in sorted(workspace.rglob("*")) + if path.is_file() + } + + +def test_dry_run_forces_optimize_without_importing_sdk(tmp_path, monkeypatch): + handoff, _, _ = _handoff(tmp_path) + handoff_path = tmp_path / "handoff.json" + result_path = tmp_path / "result.json" + _write_json(handoff_path, handoff) + + # A dry run must not enter the only function that imports claude_agent_sdk. + monkeypatch.setattr( + workflow_runner, + "invoke_via_sdk", + lambda *args, **kwargs: pytest.fail("dry-run invoked Claude SDK"), + ) + assert workflow_runner.main( + [str(handoff_path), str(result_path), "--dry-run"] + ) == 0 + + result = json.loads(result_path.read_text(encoding="utf-8")) + assert result["status"] == "dry_run" + assert result["workflow_args"]["mode"] == "optimize" + assert result["workflow_args"]["apply_to_original"] == "false" + assert result["workflow_args"]["gpu_ids"] == "7" + assert "Invoke the Workflow tool exactly once" in result["prompt"] + assert '"apply_to_original": "false"' in result["prompt"] + + +def test_agent_registry_loads_geak_v4(): + assert AgentType.from_string("geak_v4") is AgentType.GEAK_V4 + assert ( + load_agent_launcher(AgentType.GEAK_V4, logging.getLogger(__name__)) + is geak_launcher.launch_agent + ) + + +@pytest.mark.parametrize("isolated_field", ["eval_dir", "exp_root"]) +def test_handoff_rejects_artifacts_inside_kernel(tmp_path, isolated_field): + handoff, kernel, _ = _handoff(tmp_path) + handoff[isolated_field] = str(kernel / "recursive-output") + + with pytest.raises( + workflow_runner.HandoffError, + match=rf"{isolated_field} must not be inside kernel_path", + ): + workflow_runner.map_workflow_args(handoff) + + +def test_disposable_input_is_independent_and_omits_run_artifacts(tmp_path): + workspace = _workspace(tmp_path) + (workspace / "task_result.yaml").write_text("score: 1\n", encoding="utf-8") + (workspace / "__pycache__").mkdir() + (workspace / "__pycache__" / "kernel.pyc").write_bytes(b"cache") + disposable = tmp_path / "disposable" + + geak_launcher._materialize_disposable_input(workspace, disposable) + disposable.joinpath(*SOURCE.parts).write_text( + OPTIMIZED_SOURCE, + encoding="utf-8", + ) + + assert workspace.joinpath(*SOURCE.parts).read_text() == ORIGINAL_SOURCE + assert not (disposable / "task_result.yaml").exists() + assert not (disposable / "__pycache__").exists() + + +def test_disposable_input_rejects_workspace_symlinks(tmp_path): + workspace = _workspace(tmp_path) + (workspace / "src" / "kernel_alias.py").symlink_to("kernel.py") + + with pytest.raises((ValueError, RuntimeError), match="[Ss]ymlink"): + geak_launcher._materialize_disposable_input( + workspace, + tmp_path / "disposable", + ) + + +def test_artifact_root_must_not_be_a_symlink(tmp_path): + workspace = tmp_path / "workspace" + workspace.mkdir() + redirected = tmp_path / "redirected" + redirected.mkdir() + (tmp_path / ".workspace_geak_v4").symlink_to(redirected, target_is_directory=True) + + with pytest.raises(RuntimeError, match="real directory"): + geak_launcher._new_run_paths(workspace) + + +def test_json_readers_reject_symlinks_and_oversized_files(tmp_path, monkeypatch): + target = tmp_path / "target.json" + target.write_text('{"status": "ok"}\n', encoding="utf-8") + alias = tmp_path / "alias.json" + alias.symlink_to(target) + + assert workflow_runner._read_json(alias) is None + assert geak_launcher._read_json(alias) is None + + monkeypatch.setattr(workflow_runner, "_JSON_SIZE_LIMIT", 4) + monkeypatch.setattr(geak_launcher, "_JSON_SIZE_LIMIT", 4) + assert workflow_runner._read_json(target) is None + assert geak_launcher._read_json(target) is None + + +def test_atomic_json_writers_replace_destination_symlink(tmp_path): + sentinel = tmp_path / "sentinel.txt" + sentinel.write_text("SAFE\n", encoding="utf-8") + + launcher_dir = tmp_path / "launcher" + launcher_dir.mkdir() + launcher_result = launcher_dir / "result.json" + launcher_result.symlink_to(sentinel) + geak_launcher._atomic_write_json( + launcher_result, + {"status": "ok"}, + expected_parent_identity=geak_launcher._directory_identity(launcher_dir), + ) + + runner_dir = tmp_path / "runner" + runner_dir.mkdir() + runner_result = runner_dir / "result.json" + runner_result.symlink_to(sentinel) + workflow_runner._atomic_write_json(runner_result, {"status": "ok"}) + + assert sentinel.read_text(encoding="utf-8") == "SAFE\n" + for result in (launcher_result, runner_result): + assert not result.is_symlink() + assert json.loads(result.read_text(encoding="utf-8")) == {"status": "ok"} + + +def test_workspace_manifest_detects_direct_source_modification(tmp_path): + workspace = _workspace(tmp_path) + manifest = geak_launcher._workspace_manifest(workspace) + workspace.joinpath(*SOURCE.parts).write_text( + OPTIMIZED_SOURCE, + encoding="utf-8", + ) + + with pytest.raises(RuntimeError, match="direct mutation.*changed="): + geak_launcher._verify_workspace_manifest(manifest, workspace) + + +def test_workspace_manifest_detects_added_file(tmp_path): + workspace = _workspace(tmp_path) + manifest = geak_launcher._workspace_manifest(workspace) + (workspace / "rogue-output.txt").write_text( + "written outside the disposable input\n", + encoding="utf-8", + ) + + with pytest.raises(RuntimeError, match="direct mutation.*added="): + geak_launcher._verify_workspace_manifest(manifest, workspace) + + +def test_terminal_artifacts_require_complete_schema(tmp_path): + eval_dir = tmp_path / "eval" + eval_dir.mkdir() + _write_json( + eval_dir / "workflow_return.json", + { + "eval_dir": str(eval_dir), + "validation_status": "accepted", + "final_patch": "final_patch.diff", + }, + ) + _write_json( + eval_dir / "director_validation.json", + { + "validation_status": "accepted", + "correctness": "pass", + "final_patch": "final_patch.diff", + }, + ) + assert not workflow_runner._terminal_artifact_exists(eval_dir) + + _write_json( + eval_dir / "director_validation.json", + _director_validation(eval_dir), + ) + assert workflow_runner._terminal_artifact_exists(eval_dir) + + +def test_full_workflow_return_is_terminal_and_transcript_parser_ignores_noise( + tmp_path, +): + eval_dir = tmp_path / "eval" + eval_dir.mkdir() + complete = _workflow_return(eval_dir) + _write_json(eval_dir / "workflow_return.json", complete) + assert workflow_runner._terminal_artifact_exists(eval_dir) + + wrong_eval = dict(complete, eval_dir=str(tmp_path / "other")) + transcript = ( + 'setup={"enableWorkflows": true}\n' + + json.dumps(wrong_eval) + + '\npartial={"eval_dir": ' + + json.dumps(str(eval_dir)) + + "}\n" + + json.dumps(complete) + ) + assert workflow_runner._extract_workflow_return(transcript, eval_dir) == complete + + +def test_completed_background_stream_without_result_fails_immediately(): + state = { + "background_started": True, + "terminal_task_seen": True, + "result_seen": False, + "producer_done": True, + } + + error = workflow_runner._completed_producer_error(state, set(), []) + + assert isinstance(error, RuntimeError) + assert "without a ResultMessage" in str(error) + + +def test_normalize_non_workload_uses_director_geomean(tmp_path): + _, result = _prepare_normalized_result( + tmp_path, + workload_aligned=False, + geomean=1.20, + weighted=9.0, + ) + + assert result["status"] == "ok" + assert result["final_speedup"] == pytest.approx(1.20) + assert result["final_geomean"] == pytest.approx(1.20) + assert result["final_weighted"] == pytest.approx(9.0) + + +def test_normalize_workload_aligned_uses_director_weighted_speedup(tmp_path): + _, result = _prepare_normalized_result( + tmp_path, + workload_aligned=True, + geomean=1.20, + weighted=1.50, + ) + + assert result["status"] == "ok" + assert result["final_speedup"] == pytest.approx(1.50) + + +def test_normalize_rejects_non_finite_director_geomean(tmp_path): + _, result = _prepare_normalized_result( + tmp_path, + geomean=float("nan"), + weighted=1.50, + ) + + assert result["status"] == "error" + assert result["final_geomean"] is None + assert "missing or invalid" in result["reason"] + + +def test_normalize_workload_falls_back_when_weighted_is_non_finite(tmp_path): + _, result = _prepare_normalized_result( + tmp_path, + workload_aligned=True, + geomean=1.20, + weighted=float("nan"), + ) + + assert result["status"] == "ok" + assert result["final_speedup"] == pytest.approx(1.20) + assert result["final_weighted"] is None + + +def test_normalize_flagged_candidate_is_rejected(tmp_path): + _, result = _prepare_normalized_result( + tmp_path, + validation_status="flagged", + correctness="pass", + ) + + assert result["status"] == "rejected" + assert "did not accept" in result["reason"] + + +def test_parallel_worker_maps_host_gpu_to_logical_zero(monkeypatch): + monkeypatch.setenv("AGENT_KERNEL_ARENA_HOST_GPU_ID", "7") + monkeypatch.setenv("ROCR_VISIBLE_DEVICES", "7") + monkeypatch.setenv("HIP_VISIBLE_DEVICES", "7") + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "7") + monkeypatch.setenv("GEAK_V4_GPU_IDS", "7") + + assert geak_launcher._logical_gpu_ids({"gpu_ids": "7"}) == "0" + + +def test_single_declared_source_accepts_one_regular_kernel(tmp_path): + workspace = _workspace(tmp_path) + + assert geak_launcher._single_declared_source( + {"source_file_path": [str(SOURCE)]}, + workspace, + ) == SOURCE + + +@pytest.mark.parametrize( + ("source_value", "message"), + [ + (["src/kernel.py", "src/other.py"], "exactly one"), + ("../kernel.py", "safe relative path"), + ("config.yaml", "protected"), + ("scripts/test_kernel.py", "protected"), + ("test_kernel.py", "co-located test/harness"), + ("test_kernel.cpp", "co-located test/harness"), + ("test_kernel.hip", "co-located test/harness"), + ("src/kernel_harness.cu", "co-located test/harness"), + ("src/kernel_harness.py", "protected"), + ], +) +def test_single_declared_source_rejects_unsafe_allowlists( + tmp_path, + source_value, + message, +): + workspace = _workspace(tmp_path) + (workspace / "test_kernel.py").write_text("kernel = 1\n", encoding="utf-8") + (workspace / "src" / "kernel_harness.py").write_text( + "kernel = 1\n", + encoding="utf-8", + ) + + with pytest.raises(ValueError, match=message): + geak_launcher._single_declared_source( + {"source_file_path": source_value}, + workspace, + ) + + +def test_valid_single_file_patch_is_imported_atomically(tmp_path): + patch = _make_git_patch( + tmp_path, + {str(SOURCE): ORIGINAL_SOURCE}, + lambda repo: _write_repo_file(repo, str(SOURCE), OPTIMIZED_SOURCE), + ) + + applied, workspace, run_dir = _apply_patch(tmp_path, patch) + + assert applied is True + assert workspace.joinpath(*SOURCE.parts).read_text() == OPTIMIZED_SOURCE + audit = json.loads( + (run_dir / "applied_patch.json").read_text(encoding="utf-8") + ) + assert audit["source_file"] == str(SOURCE) + assert audit["director_speedup"] == pytest.approx(1.10) + + +def test_patch_audit_replaces_symlink_without_overwriting_target(tmp_path): + patch = _make_git_patch( + tmp_path, + {str(SOURCE): ORIGINAL_SOURCE}, + lambda repo: _write_repo_file(repo, str(SOURCE), OPTIMIZED_SOURCE), + ) + workspace = _workspace(tmp_path) + eval_dir = tmp_path / "eval" + run_dir = tmp_path / "run" + eval_dir.mkdir() + run_dir.mkdir() + (eval_dir / "final_patch.diff").write_bytes(patch) + sentinel = tmp_path / "sentinel.txt" + sentinel.write_text("SAFE\n", encoding="utf-8") + audit = run_dir / "applied_patch.json" + audit.symlink_to(sentinel) + + assert geak_launcher._apply_validated_patch( + result=_import_result(eval_dir), + expected_eval_dir=eval_dir, + workspace=workspace, + source_path=SOURCE, + min_improve=0.02, + run_dir=run_dir, + ) + + assert sentinel.read_text(encoding="utf-8") == "SAFE\n" + assert not audit.is_symlink() + assert json.loads(audit.read_text(encoding="utf-8"))["source_file"] == str(SOURCE) + + +def test_valid_patch_applies_inside_an_outer_git_worktree(tmp_path): + patch = _make_git_patch( + tmp_path, + {str(SOURCE): ORIGINAL_SOURCE}, + lambda repo: _write_repo_file(repo, str(SOURCE), OPTIMIZED_SOURCE), + ) + outer = tmp_path / "outer_repo" + outer.mkdir() + _git(outer, "init", "-q") + workspace = _workspace(outer) + eval_dir = outer / "eval" + run_dir = outer / "run" + eval_dir.mkdir() + run_dir.mkdir() + (eval_dir / "final_patch.diff").write_bytes(patch) + + assert geak_launcher._apply_validated_patch( + result=_import_result(eval_dir), + expected_eval_dir=eval_dir, + workspace=workspace, + source_path=SOURCE, + min_improve=0.02, + run_dir=run_dir, + ) + assert workspace.joinpath(*SOURCE.parts).read_text() == OPTIMIZED_SOURCE + + +def test_normalize_rejects_patch_not_named_by_director(tmp_path): + eval_dir = tmp_path / "eval" + eval_dir.mkdir() + _write_json(eval_dir / "workflow_return.json", _workflow_return(eval_dir)) + validation = _director_validation(eval_dir) + validation["final_patch"] = str(eval_dir / "validated_elsewhere.diff") + _write_json(eval_dir / "director_validation.json", validation) + (eval_dir / "final_patch.diff").write_text( + "unvalidated patch\n", + encoding="utf-8", + ) + + result = workflow_runner.normalize_result(eval_dir) + + assert result["status"] == "error" + assert "Director artifact is missing or invalid" in result["reason"] + + +def test_patch_import_rechecks_director_patch_provenance(tmp_path): + patch = _make_git_patch( + tmp_path, + {str(SOURCE): ORIGINAL_SOURCE}, + lambda repo: _write_repo_file(repo, str(SOURCE), OPTIMIZED_SOURCE), + ) + workspace = _workspace(tmp_path) + eval_dir = tmp_path / "eval" + run_dir = tmp_path / "run" + eval_dir.mkdir() + run_dir.mkdir() + (eval_dir / "final_patch.diff").write_bytes(patch) + result = _import_result(eval_dir) + result["director_final_patch"] = str(eval_dir / "validated_elsewhere.diff") + + with pytest.raises(RuntimeError, match="director_final_patch"): + geak_launcher._apply_validated_patch( + result=result, + expected_eval_dir=eval_dir, + workspace=workspace, + source_path=SOURCE, + min_improve=0.02, + run_dir=run_dir, + ) + + assert workspace.joinpath(*SOURCE.parts).read_text() == ORIGINAL_SOURCE + + +def test_patch_import_rejects_unknown_result_schema(tmp_path): + patch = _make_git_patch( + tmp_path, + {str(SOURCE): ORIGINAL_SOURCE}, + lambda repo: _write_repo_file(repo, str(SOURCE), OPTIMIZED_SOURCE), + ) + workspace = _workspace(tmp_path) + eval_dir = tmp_path / "eval" + run_dir = tmp_path / "run" + eval_dir.mkdir() + run_dir.mkdir() + (eval_dir / "final_patch.diff").write_bytes(patch) + result = _import_result(eval_dir) + result["schema_version"] = 999 + + with pytest.raises(RuntimeError, match="schema"): + geak_launcher._apply_validated_patch( + result=result, + expected_eval_dir=eval_dir, + workspace=workspace, + source_path=SOURCE, + min_improve=0.02, + run_dir=run_dir, + ) + + assert workspace.joinpath(*SOURCE.parts).read_text() == ORIGINAL_SOURCE + + +@pytest.mark.parametrize("status", ["no_gain", "rejected"]) +def test_non_accepted_status_never_parses_or_applies_patch(tmp_path, status): + traversal = ( + b"diff --git a/../escape.py b/../escape.py\n" + b"--- a/../escape.py\n" + b"+++ b/../escape.py\n" + b"@@ -1 +1 @@\n" + b"-unsafe\n" + b"+escaped\n" + ) + + applied, workspace, _ = _apply_patch(tmp_path, traversal, status=status) + + assert applied is False + assert workspace.joinpath(*SOURCE.parts).read_text() == ORIGINAL_SOURCE + assert not (tmp_path / "escape.py").exists() + + +def _malicious_patch(tmp_path: Path, kind: str) -> bytes: + if kind == "multi_file": + def mutate(repo: Path) -> None: + _write_repo_file(repo, str(SOURCE), OPTIMIZED_SOURCE) + _write_repo_file(repo, "config.yaml", "task_type: fake\n") + + return _make_git_patch( + tmp_path, + { + str(SOURCE): ORIGINAL_SOURCE, + "config.yaml": "task_type: triton2triton\n", + }, + mutate, + ) + if kind == "new_file": + return _make_git_patch( + tmp_path, + {}, + lambda repo: _write_repo_file(repo, str(SOURCE), OPTIMIZED_SOURCE), + ) + if kind == "delete": + return _make_git_patch( + tmp_path, + {str(SOURCE): ORIGINAL_SOURCE}, + lambda repo: (repo / str(SOURCE)).unlink(), + ) + if kind == "rename": + def rename(repo: Path) -> None: + destination = repo / "src" / "renamed.py" + (repo / str(SOURCE)).rename(destination) + + return _make_git_patch( + tmp_path, + {str(SOURCE): ORIGINAL_SOURCE}, + rename, + find_renames=True, + ) + if kind == "binary": + return _make_git_patch( + tmp_path, + {str(SOURCE): b"old\x00binary\n"}, + lambda repo: _write_repo_file( + repo, + str(SOURCE), + b"new\x00binary\n", + ), + ) + if kind == "traversal": + return ( + b"diff --git a/../escape.py b/../escape.py\n" + b"index 1234567..7654321 100644\n" + b"--- a/../escape.py\n" + b"+++ b/../escape.py\n" + b"@@ -1 +1 @@\n" + b"-unsafe\n" + b"+escaped\n" + ) + raise AssertionError(f"unhandled malicious patch kind: {kind}") + + +@pytest.mark.parametrize( + "kind", + ["multi_file", "new_file", "delete", "rename", "binary", "traversal"], +) +def test_malicious_patch_is_rejected_without_workspace_mutation(tmp_path, kind): + patch = _malicious_patch(tmp_path, kind) + workspace = _workspace(tmp_path) + before = _workspace_contents(workspace) + eval_dir = tmp_path / "eval" + run_dir = tmp_path / "run" + eval_dir.mkdir() + run_dir.mkdir() + (eval_dir / "final_patch.diff").write_bytes(patch) + + with pytest.raises(RuntimeError): + geak_launcher._apply_validated_patch( + result=_import_result(eval_dir), + expected_eval_dir=eval_dir, + workspace=workspace, + source_path=SOURCE, + min_improve=0.02, + run_dir=run_dir, + ) + + assert _workspace_contents(workspace) == before + assert not (tmp_path / "escape.py").exists()