mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-29 06:33:28 +00:00
feat(cookbook): profiling live runner, preflight and cleanup guards
Co-Authored-By: bot_apk <apk@cognition.ai>
This commit is contained in:
Devin AI
parent
0cf05cc15b
commit
ef1cb8cd1b
7 changed files with 1140 additions and 178 deletions
|
|
@ -14,7 +14,7 @@ The deterministic selection seed is `litellm-capability-profiling-v1`. Candidate
|
|||
|
||||
The effective API ceiling is $185, consisting of a $5 labeling allowance and a $180 solver allowance. The solver limit is 180 runs, at most 60 tasks across three solver models. A future run uses $1 per run, at most 50 steps, and a 600-second timeout, with three workers. The runner accepts `--max-batches N` to stop after a bounded number of three-task batches. For task index `i` in manifest order, model order rotates by `i % 3`: `haiku-4.5, sonnet-5, opus-5`, then `sonnet-5, opus-5, haiku-4.5`, then `opus-5, haiku-4.5, sonnet-5`. Each run record stores that task's model order
|
||||
|
||||
Keep at least 50 GiB free and abort before any image pull if free disk is below 60 GiB. Process images in batches of three, validate with gold and noop before solving, and never replace a task that fails preflight. Clean only task-owned image tags or IDs after a batch, then verify disk; do not run `docker system prune` or `docker image prune`. Write raw results and a `final=false` map checkpoint after each batch, with a progress file for resume. The current runner entry point only builds dry-run schedules and checkpoints; it launches no pulls, gold/noop checks, or solves
|
||||
Keep at least 50 GiB free and abort before any image pull if free disk is below 60 GiB. Process images in batches of three, validate with gold and noop before solving, and never replace a task that fails preflight. Clean only task-owned image tags or IDs after a batch, then verify disk; do not run `docker system prune` or `docker image prune`. The live runner appends attempt and preflight records to `profiling_runs.jsonl`, compiles a map after each batch, and uses progress only as advisory output. Dry-run prints the schedule without writing progress or live artifacts. The live path is implemented but was not launched
|
||||
|
||||
## Profile records and routing
|
||||
|
||||
|
|
|
|||
|
|
@ -203,7 +203,24 @@ def attempt_from_mapping(value: object) -> TaskAttempt:
|
|||
|
||||
|
||||
def read_attempts_jsonl(path: Path) -> tuple[TaskAttempt, ...]:
|
||||
return tuple(attempt_from_mapping(json.loads(line)) for line in path.read_text().splitlines() if line.strip())
|
||||
if not path.exists():
|
||||
return ()
|
||||
rows: Final = tuple(
|
||||
_record_mapping(json.loads(line)) for line in path.read_text().splitlines() if line.strip()
|
||||
)
|
||||
return tuple(
|
||||
attempt_from_mapping(row)
|
||||
for row in rows
|
||||
if row.get("record_type", "attempt") == "attempt"
|
||||
)
|
||||
|
||||
|
||||
def _run_rows(path: Path) -> tuple[Mapping[str, object], ...]:
|
||||
if not path.exists():
|
||||
return ()
|
||||
return tuple(
|
||||
_record_mapping(json.loads(line)) for line in path.read_text().splitlines() if line.strip()
|
||||
)
|
||||
|
||||
|
||||
def _manifest_category_ids(value: object) -> tuple[str, ...]:
|
||||
|
|
@ -319,7 +336,6 @@ def aggregate_profiles(
|
|||
attempt.key.category,
|
||||
)
|
||||
for attempt in attempts
|
||||
if not attempt.is_infra
|
||||
)
|
||||
cohort_groups: Final = groupby(sorted(cohorts), key=lambda row: row[:-1])
|
||||
if any(len({row[-1] for row in rows}) > 1 for _, rows in cohort_groups):
|
||||
|
|
@ -332,14 +348,20 @@ def aggregate_profiles(
|
|||
|
||||
|
||||
def _versions_match(profile: Profile, versions: ProfileVersions) -> bool:
|
||||
expected_model_version: Final = dict(versions.model_versions).get(profile.key.model_alias)
|
||||
expected_deployment_id: Final = dict(versions.model_deployments).get(profile.key.model_alias)
|
||||
model_versions_match: Final = (
|
||||
not versions.model_versions
|
||||
or (profile.key.model_alias, profile.key.returned_model_version) in versions.model_versions
|
||||
)
|
||||
deployments_match: Final = (
|
||||
not versions.model_deployments
|
||||
or (profile.key.model_alias, profile.key.deployment_id) in versions.model_deployments
|
||||
)
|
||||
return (
|
||||
profile.key.harness_version == versions.harness_version
|
||||
and profile.key.budget_version == versions.budget_version
|
||||
and profile.key.harness_config_hash == versions.harness_config_hash
|
||||
and profile.key.returned_model_version == expected_model_version
|
||||
and (expected_deployment_id is None or profile.key.deployment_id == expected_deployment_id)
|
||||
and model_versions_match
|
||||
and deployments_match
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -398,6 +420,7 @@ def compile_map(
|
|||
tau: float = 0.80,
|
||||
min_evidence: int = 13,
|
||||
final: bool = False,
|
||||
eligible: frozenset[tuple[str, str | None]] | None = None,
|
||||
) -> dict[str, object]:
|
||||
_validate_thresholds(tau, min_evidence)
|
||||
categories: Final = (
|
||||
|
|
@ -428,6 +451,7 @@ def compile_map(
|
|||
category_profiles[category],
|
||||
tau,
|
||||
min_evidence=min_evidence,
|
||||
eligible=eligible,
|
||||
)
|
||||
)
|
||||
for category in categories
|
||||
|
|
@ -446,8 +470,47 @@ def compile_map(
|
|||
"harness_version": versions.harness_version,
|
||||
"budget_version": versions.budget_version,
|
||||
"harness_config_hash": versions.harness_config_hash,
|
||||
"model_versions": dict(versions.model_versions),
|
||||
"model_deployments": dict(versions.model_deployments),
|
||||
"model_versions": {
|
||||
alias: tuple(sorted({profile.key.returned_model_version for profile in matching
|
||||
if profile.key.model_alias == alias}))
|
||||
for alias in sorted({profile.key.model_alias for profile in matching})
|
||||
},
|
||||
"model_deployments": {
|
||||
alias: tuple(sorted({profile.key.deployment_id for profile in matching
|
||||
if profile.key.model_alias == alias}))
|
||||
for alias in sorted({profile.key.model_alias for profile in matching})
|
||||
},
|
||||
"observed_model_cohorts": tuple(
|
||||
{
|
||||
"model_alias": alias,
|
||||
"returned_model_version": returned_model,
|
||||
"deployment_id": deployment_id or None,
|
||||
}
|
||||
for alias, returned_model, deployment_id in sorted(
|
||||
{
|
||||
(profile.key.model_alias, profile.key.returned_model_version, profile.key.deployment_id)
|
||||
for profile in matching
|
||||
}
|
||||
)
|
||||
),
|
||||
"version_changes": tuple(
|
||||
{
|
||||
"model_alias": alias,
|
||||
"returned_model_versions": versions_for_alias,
|
||||
"deployment_ids": deployments_for_alias,
|
||||
}
|
||||
for alias, versions_for_alias, deployments_for_alias in (
|
||||
(
|
||||
alias,
|
||||
tuple(sorted({profile.key.returned_model_version for profile in matching
|
||||
if profile.key.model_alias == alias})),
|
||||
tuple(sorted({profile.key.deployment_id for profile in matching
|
||||
if profile.key.model_alias == alias})),
|
||||
)
|
||||
for alias in sorted({profile.key.model_alias for profile in matching})
|
||||
)
|
||||
if len(versions_for_alias) > 1 or len(deployments_for_alias) > 1
|
||||
),
|
||||
},
|
||||
"profiles": [
|
||||
asdict(profile)
|
||||
|
|
@ -468,16 +531,147 @@ def compile_map_from_profiling_runs(
|
|||
) -> dict[str, object]:
|
||||
_validate_thresholds(tau, min_evidence)
|
||||
selected_task_ids, actual_manifest_sha256 = manifest_selection(manifest_path, expected_manifest_sha256)
|
||||
rows: Final = _run_rows(path)
|
||||
for row in rows:
|
||||
row_task_id: Final = row.get("task_id")
|
||||
row_manifest_sha256: Final = row.get("manifest_sha256")
|
||||
if row.get("record_type", "attempt") in {
|
||||
"attempt",
|
||||
"preflight_passed",
|
||||
"preflight_failed",
|
||||
"quarantined_attempt",
|
||||
}:
|
||||
if not isinstance(row_task_id, str) or not isinstance(row_manifest_sha256, str):
|
||||
raise ValueError("profiling run records must include task_id and manifest_sha256")
|
||||
validate_manifest_attempt(
|
||||
row_task_id,
|
||||
row_manifest_sha256,
|
||||
selected_task_ids,
|
||||
actual_manifest_sha256,
|
||||
)
|
||||
attempts: Final = read_attempts_jsonl(path)
|
||||
for attempt in attempts:
|
||||
validate_manifest_attempt(
|
||||
attempt.task_id,
|
||||
attempt.manifest_sha256,
|
||||
selected_task_ids,
|
||||
actual_manifest_sha256,
|
||||
observed_versions: Final = tuple(
|
||||
sorted({(attempt.key.model_alias, attempt.key.returned_model_version) for attempt in attempts})
|
||||
)
|
||||
observed_deployments: Final = tuple(
|
||||
sorted(
|
||||
{
|
||||
(attempt.key.model_alias, attempt.key.deployment_id)
|
||||
for attempt in attempts
|
||||
}
|
||||
)
|
||||
)
|
||||
effective_versions: Final = ProfileVersions(
|
||||
versions.harness_version,
|
||||
versions.budget_version,
|
||||
versions.harness_config_hash,
|
||||
observed_versions,
|
||||
observed_deployments,
|
||||
)
|
||||
profiles: Final = aggregate_profiles(attempts, tau, min_evidence)
|
||||
return compile_map(tenant_id, versions, profiles, tau, min_evidence, final=False)
|
||||
preflight: Final = tuple(row for row in rows if row.get("record_type") in {"preflight_passed", "preflight_failed"})
|
||||
preflight_by_task: Final = {
|
||||
str(row["task_id"]): row
|
||||
for row in preflight
|
||||
if isinstance(row.get("task_id"), str)
|
||||
}
|
||||
task_order: Final = tuple(selected_task_ids)
|
||||
expected_models: Final = {
|
||||
task_id: (
|
||||
("haiku-4.5", "sonnet-5", "opus-5")[index % 3: ]
|
||||
+ ("haiku-4.5", "sonnet-5", "opus-5")[: index % 3]
|
||||
)
|
||||
for index, task_id in enumerate(task_order)
|
||||
}
|
||||
attempt_rows: Final = tuple(row for row in rows if row.get("record_type", "attempt") == "attempt")
|
||||
completion: Final = _map_completeness(task_order, expected_models, preflight_by_task, attempt_rows)
|
||||
quarantined: Final = tuple(
|
||||
dict(row)
|
||||
for row in rows
|
||||
if row.get("record_type") == "quarantined_attempt"
|
||||
)
|
||||
failed_preflights: Final = tuple(
|
||||
dict(row)
|
||||
for row in preflight
|
||||
if row.get("record_type") == "preflight_failed"
|
||||
)
|
||||
return (
|
||||
compile_map(
|
||||
tenant_id,
|
||||
effective_versions,
|
||||
profiles,
|
||||
tau,
|
||||
min_evidence,
|
||||
final=completion["final"] is True,
|
||||
eligible=frozenset(
|
||||
(attempt.key.model_alias, attempt.key.deployment_id or None)
|
||||
for attempt in attempts
|
||||
if attempt.key.model_alias in {"haiku-4.5", "sonnet-5", "opus-5"}
|
||||
),
|
||||
)
|
||||
| {
|
||||
"infra_deficits": completion["infra_deficits"],
|
||||
"preflight_incomplete": completion["preflight_incomplete"],
|
||||
"preflight_failed": failed_preflights,
|
||||
"quarantined_attempts": quarantined,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _map_completeness(
|
||||
task_order: tuple[str, ...],
|
||||
expected_models: Mapping[str, tuple[str, ...]],
|
||||
preflight_by_task: Mapping[str, Mapping[str, object]],
|
||||
attempt_rows: Sequence[Mapping[str, object]],
|
||||
) -> dict[str, object]:
|
||||
preflight_incomplete: Final = tuple(task_id for task_id in task_order if task_id not in preflight_by_task)
|
||||
passed_task_ids: Final = tuple(
|
||||
task_id
|
||||
for task_id in task_order
|
||||
if preflight_by_task.get(task_id, {}).get("record_type") == "preflight_passed"
|
||||
)
|
||||
grouped_rows: Final = {
|
||||
(str(row.get("task_id")), str(row.get("model_alias", row.get("model")))): tuple(
|
||||
candidate
|
||||
for candidate in attempt_rows
|
||||
if candidate.get("task_id") == row.get("task_id")
|
||||
and candidate.get("model_alias", candidate.get("model")) == row.get("model_alias", row.get("model"))
|
||||
)
|
||||
for row in attempt_rows
|
||||
}
|
||||
infra_deficits: Final = tuple(
|
||||
{
|
||||
"task_id": task_id,
|
||||
"model_alias": model_alias,
|
||||
"error_class": latest.get("error_class"),
|
||||
"cost_usd": latest.get("cost_usd"),
|
||||
}
|
||||
for task_id in passed_task_ids
|
||||
for model_alias in expected_models[task_id]
|
||||
for rows in (grouped_rows.get((task_id, model_alias), ()),)
|
||||
if rows
|
||||
for latest in (rows[-1],)
|
||||
if latest.get("is_infra") is True
|
||||
and latest.get("retry_eligible") is False
|
||||
and not any(row.get("is_infra") is False for row in rows)
|
||||
)
|
||||
incomplete_pairs: Final = tuple(
|
||||
(task_id, model_alias)
|
||||
for task_id in passed_task_ids
|
||||
for model_alias in expected_models[task_id]
|
||||
if not any(
|
||||
row.get("is_infra") is False
|
||||
for row in grouped_rows.get((task_id, model_alias), ())
|
||||
)
|
||||
and (task_id, model_alias) not in {
|
||||
(str(row["task_id"]), str(row["model_alias"])) for row in infra_deficits
|
||||
}
|
||||
)
|
||||
return {
|
||||
"final": not preflight_incomplete and not incomplete_pairs,
|
||||
"preflight_incomplete": preflight_incomplete,
|
||||
"infra_deficits": infra_deficits,
|
||||
}
|
||||
|
||||
|
||||
def _validate_thresholds(tau: float, min_evidence: int) -> None:
|
||||
|
|
|
|||
|
|
@ -59,6 +59,7 @@ class AttemptPaths:
|
|||
image_digests: Mapping[str, str]
|
||||
image_overlay_ids: Mapping[str, str]
|
||||
dataset_metadata: Mapping[str, str]
|
||||
root_dir: Path | None = None
|
||||
TOKEN_FIELDS = {
|
||||
"prompt": "prompt_tokens",
|
||||
"completion": "completion_tokens",
|
||||
|
|
@ -88,8 +89,12 @@ def read_jsonl_locked(path: Path) -> list[dict[str, object]]:
|
|||
|
||||
|
||||
def committed_spend_usd(runs_path: Path = RUNS_PATH, aux_path: Path = AUX_COSTS_PATH) -> float:
|
||||
run_cost = sum(_cost(row.get("cost_usd")) for row in read_jsonl_locked(runs_path))
|
||||
aux_cost = sum(_cost(row.get("cost_usd")) for row in read_jsonl_locked(aux_path))
|
||||
run_rows = read_jsonl_locked(runs_path)
|
||||
aux_rows = read_jsonl_locked(aux_path)
|
||||
if any(not valid_cost(row.get("cost_usd")) for row in run_rows + aux_rows):
|
||||
raise ValueError("committed spend contains missing or invalid cost")
|
||||
run_cost = sum(_cost(row.get("cost_usd")) for row in run_rows)
|
||||
aux_cost = sum(_cost(row.get("cost_usd")) for row in aux_rows)
|
||||
return run_cost + aux_cost
|
||||
|
||||
|
||||
|
|
@ -184,11 +189,11 @@ def _utc_now() -> str:
|
|||
return datetime.now(UTC).isoformat()
|
||||
|
||||
|
||||
def _read_pilot_key() -> str:
|
||||
path = ROOT / ".private" / "pilot_key"
|
||||
if stat.S_IMODE(path.stat().st_mode) != 0o600:
|
||||
def _read_pilot_key(path: Path | None = None) -> str:
|
||||
key_path = path if path is not None else ROOT / ".private" / "pilot_key"
|
||||
if stat.S_IMODE(key_path.stat().st_mode) != 0o600:
|
||||
raise PermissionError("pilot key must have mode 0600")
|
||||
key = path.read_text().strip()
|
||||
key = key_path.read_text().strip()
|
||||
if not key:
|
||||
raise ValueError("pilot key is empty")
|
||||
return key
|
||||
|
|
@ -283,6 +288,7 @@ def _grade(
|
|||
run_dir: Path,
|
||||
env: dict[str, str],
|
||||
dataset_name: str = "SWE-bench/SWE-bench",
|
||||
root_dir: Path = ROOT,
|
||||
) -> tuple[str, float]:
|
||||
if not submission.strip():
|
||||
return "skipped_empty_patch", 0.0
|
||||
|
|
@ -303,7 +309,7 @@ def _grade(
|
|||
+ "\n"
|
||||
)
|
||||
command = [
|
||||
str(ROOT / ".venv" / "bin" / "python"),
|
||||
str(root_dir / ".venv" / "bin" / "python"),
|
||||
"-m",
|
||||
"swebench.harness.run_evaluation",
|
||||
"-d",
|
||||
|
|
@ -343,10 +349,10 @@ def _grade(
|
|||
return status, time.monotonic() - started
|
||||
|
||||
|
||||
def _cleanup_containers(run_id: str, env: dict[str, str]) -> int:
|
||||
def _cleanup_containers(run_id: str, env: dict[str, str], root_dir: Path = ROOT) -> int:
|
||||
result = subprocess.run(
|
||||
["docker", "ps", "-aq", "--filter", f"label=routing_poc_run={run_id}"],
|
||||
cwd=ROOT,
|
||||
cwd=root_dir,
|
||||
env=_utility_env(env),
|
||||
capture_output=True,
|
||||
text=True,
|
||||
|
|
@ -357,7 +363,7 @@ def _cleanup_containers(run_id: str, env: dict[str, str]) -> int:
|
|||
return 0
|
||||
removed = subprocess.run(
|
||||
["docker", "rm", "-f", *container_ids],
|
||||
cwd=ROOT,
|
||||
cwd=root_dir,
|
||||
env=_utility_env(env),
|
||||
capture_output=True,
|
||||
)
|
||||
|
|
@ -391,6 +397,7 @@ def _attempt(
|
|||
started_at = _utc_now()
|
||||
started = time.monotonic()
|
||||
run_id = f"{task_id}__{model}__a{attempt}__{_timestamp()}"
|
||||
root_dir: Final = paths.root_dir if paths is not None and paths.root_dir is not None else ROOT
|
||||
run_root: Final = paths.run_root if paths is not None else RUNS_ROOT
|
||||
dataset_root: Final = paths.dataset_root if paths is not None else ROOT / "agent_dataset"
|
||||
run_dir = run_root / run_id
|
||||
|
|
@ -403,7 +410,7 @@ def _attempt(
|
|||
command = [
|
||||
"timeout",
|
||||
"600",
|
||||
str(ROOT / ".venv" / "bin" / "mini-extra"),
|
||||
str(root_dir / ".venv" / "bin" / "mini-extra"),
|
||||
"swebench",
|
||||
"--subset",
|
||||
str(dataset_root),
|
||||
|
|
@ -426,7 +433,7 @@ def _attempt(
|
|||
with solve_log.open("wb") as log:
|
||||
process = subprocess.Popen(
|
||||
command,
|
||||
cwd=ROOT,
|
||||
cwd=root_dir,
|
||||
env=env,
|
||||
stdout=log,
|
||||
stderr=subprocess.STDOUT,
|
||||
|
|
@ -444,7 +451,7 @@ def _attempt(
|
|||
rc = process.wait()
|
||||
_bounded_log(solve_log)
|
||||
solve_latency_s = time.monotonic() - started
|
||||
containers_removed = _cleanup_containers(run_id, env)
|
||||
containers_removed = _cleanup_containers(run_id, env, root_dir)
|
||||
trajectory_path = solve_dir / task_id / f"{task_id}.traj.json"
|
||||
info, trajectory_found = _trajectory(trajectory_path)
|
||||
exit_status = info.get("exit_status")
|
||||
|
|
@ -461,6 +468,7 @@ def _attempt(
|
|||
run_dir,
|
||||
env,
|
||||
paths.eval_dataset if paths is not None else "SWE-bench/SWE-bench",
|
||||
root_dir,
|
||||
)
|
||||
if eval_status == "skipped_empty_patch":
|
||||
success = False
|
||||
|
|
@ -590,6 +598,13 @@ class DispatchGuard:
|
|||
{"task_id": task_id, "reason": "cost_accounting_incomplete", "ts": _utc_now()},
|
||||
)
|
||||
|
||||
def stop(self, task_id: str, reason: str) -> None:
|
||||
with self._lock:
|
||||
if self._stopped:
|
||||
return
|
||||
self._stopped = True
|
||||
append_jsonl(self._stop_path, {"task_id": task_id, "reason": reason, "ts": _utc_now()})
|
||||
|
||||
def release(self) -> None:
|
||||
with self._lock:
|
||||
self._in_flight -= 1
|
||||
|
|
|
|||
|
|
@ -1,13 +1,22 @@
|
|||
import argparse
|
||||
import hashlib
|
||||
import json
|
||||
import math
|
||||
import os
|
||||
import shutil
|
||||
from collections.abc import Callable, Mapping
|
||||
import stat
|
||||
import subprocess
|
||||
import time
|
||||
from collections.abc import Callable, Mapping, Sequence
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
from datasets import load_dataset
|
||||
|
||||
from poc import run_pilot
|
||||
from poc.prepare import DATASET_NAME, DATASET_REVISION, DATASET_SPLIT, agent_rows
|
||||
from poc.profiles import (
|
||||
PUBLIC_DEV_TENANT,
|
||||
HarnessConfig,
|
||||
|
|
@ -25,6 +34,16 @@ BATCH_SIZE: Final = 3
|
|||
WORKERS: Final = 3
|
||||
MIN_FREE_BYTES: Final = 60 * 1024**3
|
||||
MANIFEST_SHA256: Final = "15720dee12d40b62780837ed12a403f3951e9c5bb1ed18ec0d5d5d1f4776bff4"
|
||||
EVAL_DATASET_REVISION: Final = "c6fe717fd7a4c3ac1daa4055a4fd082c6a1d28a2"
|
||||
NOOP_PATCH: Final = (
|
||||
"diff --git a/routing_poc_noop.txt b/routing_poc_noop.txt\n"
|
||||
"new file mode 100644\n"
|
||||
"index 0000000..d0951a7\n"
|
||||
"--- /dev/null\n"
|
||||
"+++ b/routing_poc_noop.txt\n"
|
||||
"@@ -0,0 +1 @@\n"
|
||||
"+noop\n"
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
|
|
@ -49,6 +68,8 @@ class ProfileTask:
|
|||
task_id: str
|
||||
category: str
|
||||
model_order: tuple[str, ...]
|
||||
image_ref: str | None = None
|
||||
image_digest: str | None = None
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
|
|
@ -57,34 +78,47 @@ class ProfileBatch:
|
|||
tasks: tuple[ProfileTask, ...]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class AttemptPaths:
|
||||
pilot: run_pilot.AttemptPaths
|
||||
task_rows: Mapping[str, dict[str, object]]
|
||||
gold_patches: Mapping[str, str]
|
||||
solver_configs: Mapping[str, dict[str, object]]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class RuntimeOps:
|
||||
disk_stats: Callable[[], Mapping[str, int]]
|
||||
inspect_image: Callable[[str], str | None]
|
||||
pull_image: Callable[[str], None]
|
||||
evaluate: Callable[[str, str, str, run_pilot.AttemptPaths], tuple[str, float]]
|
||||
remove_image: Callable[[str], bool]
|
||||
image_exists: Callable[[str], bool]
|
||||
attempt: Callable[
|
||||
[str, str, int, dict[str, object], dict[str, object], object, object, dict[str, str], str,
|
||||
run_pilot.AttemptPaths],
|
||||
dict[str, object],
|
||||
]
|
||||
|
||||
|
||||
def _mapping(value: object, field: str) -> Mapping[str, object]:
|
||||
if not isinstance(value, Mapping):
|
||||
raise ValueError(f"{field} must be a JSON object")
|
||||
return value
|
||||
|
||||
|
||||
def model_order_for_task(index: int) -> tuple[str, ...]:
|
||||
def _model_order(index: int) -> tuple[str, ...]:
|
||||
offset: Final = index % len(MODEL_ORDER)
|
||||
return MODEL_ORDER[offset:] + MODEL_ORDER[:offset]
|
||||
|
||||
|
||||
def _category_tasks(
|
||||
category: str,
|
||||
manifest_category: Mapping[str, object],
|
||||
index_offset: int,
|
||||
) -> tuple[ProfileTask, ...]:
|
||||
selected_ids: Final = manifest_category.get("selected_ids")
|
||||
if not isinstance(selected_ids, list) or not all(isinstance(task_id, str) for task_id in selected_ids):
|
||||
raise ValueError(f"{category} selected_ids must be strings")
|
||||
return tuple(
|
||||
ProfileTask(
|
||||
index=index_offset + index,
|
||||
task_id=task_id,
|
||||
category=category,
|
||||
model_order=model_order_for_task(index_offset + index),
|
||||
)
|
||||
for index, task_id in enumerate(selected_ids)
|
||||
)
|
||||
model_order_for_task = _model_order
|
||||
|
||||
|
||||
class MixedVersionAttemptError(ValueError):
|
||||
def __init__(self, task_id: str):
|
||||
super().__init__("mixed_version attempt quarantined; stopping profiling run")
|
||||
self.task_id = task_id
|
||||
|
||||
|
||||
def profile_batches(config: ProfileConfig, max_batches: int | None = None) -> tuple[ProfileBatch, ...]:
|
||||
|
|
@ -93,17 +127,37 @@ def profile_batches(config: ProfileConfig, max_batches: int | None = None) -> tu
|
|||
selected_ids, _ = manifest_selection(config.manifest_path, config.manifest_sha256)
|
||||
manifest: Final = _mapping(json.loads(config.manifest_path.read_text()), "manifest")
|
||||
categories: Final = _mapping(manifest.get("categories"), "manifest categories")
|
||||
repository_tasks: Final = _category_tasks(
|
||||
"repository_debugging",
|
||||
_mapping(categories.get("repository_debugging"), "repository_debugging"),
|
||||
0,
|
||||
candidate_rows: Final = manifest.get("candidates", ())
|
||||
candidate_by_id: Final = {
|
||||
row["task_id"]: row
|
||||
for row in candidate_rows
|
||||
if isinstance(row, Mapping) and isinstance(row.get("task_id"), str)
|
||||
}
|
||||
tasks: Final = tuple(
|
||||
ProfileTask(
|
||||
index=index,
|
||||
task_id=task_id,
|
||||
category=category,
|
||||
model_order=_model_order(index),
|
||||
image_ref=(
|
||||
f"{candidate_by_id[task_id]['image']}@{candidate_by_id[task_id]['image_digest']}"
|
||||
if task_id in candidate_by_id
|
||||
and isinstance(candidate_by_id[task_id].get("image"), str)
|
||||
and isinstance(candidate_by_id[task_id].get("image_digest"), str)
|
||||
else None
|
||||
),
|
||||
image_digest=(
|
||||
str(candidate_by_id[task_id]["image_digest"])
|
||||
if task_id in candidate_by_id and isinstance(candidate_by_id[task_id].get("image_digest"), str)
|
||||
else None
|
||||
),
|
||||
)
|
||||
for category in ("repository_debugging", "feature_implementation")
|
||||
for ids in (_mapping(categories.get(category), category).get("selected_ids"),)
|
||||
if isinstance(ids, list)
|
||||
for task_id in ids
|
||||
for index in (selected_ids.index(task_id),)
|
||||
)
|
||||
feature_tasks: Final = _category_tasks(
|
||||
"feature_implementation",
|
||||
_mapping(categories.get("feature_implementation"), "feature_implementation"),
|
||||
len(repository_tasks),
|
||||
)
|
||||
tasks: Final = repository_tasks + feature_tasks
|
||||
if tuple(task.task_id for task in tasks) != selected_ids:
|
||||
raise ValueError("profile categories differ from manifest selection")
|
||||
batches: Final = tuple(
|
||||
|
|
@ -118,16 +172,6 @@ def require_disk_for_pull(free_bytes: int) -> None:
|
|||
raise OSError("abort profiling batch before pull: less than 60 GiB free")
|
||||
|
||||
|
||||
def _single_result_value(result: Mapping[str, object], field: str) -> str | None:
|
||||
value: Final = result.get(field, ())
|
||||
values: Final = (value,) if isinstance(value, str) else value
|
||||
if not isinstance(values, (list, tuple)) or not all(isinstance(item, str) for item in values):
|
||||
raise ValueError(f"{field} must contain strings")
|
||||
if len(values) > 1:
|
||||
raise ValueError(f"{field} must contain at most one value")
|
||||
return values[0] if values else None
|
||||
|
||||
|
||||
def profile_guard(
|
||||
config: ProfileConfig,
|
||||
free_bytes: Callable[[], int] = lambda: shutil.disk_usage("/").free,
|
||||
|
|
@ -147,6 +191,21 @@ def profile_guard(
|
|||
)
|
||||
|
||||
|
||||
def _result_values(result: Mapping[str, object], field: str) -> tuple[str, ...]:
|
||||
value: Final = result.get(field, ())
|
||||
values: Final = (value,) if isinstance(value, str) else value
|
||||
if not isinstance(values, (list, tuple)) or not all(isinstance(item, str) for item in values):
|
||||
raise ValueError(f"{field} must contain strings")
|
||||
return tuple(values)
|
||||
|
||||
|
||||
def _required_metric(result: Mapping[str, object], field: str) -> float:
|
||||
value: Final = result.get(field)
|
||||
if isinstance(value, bool) or not isinstance(value, (int, float)) or not math.isfinite(value) or value < 0:
|
||||
raise ValueError(f"{field} must be finite and non-negative")
|
||||
return float(value)
|
||||
|
||||
|
||||
def record_task_result(
|
||||
config: ProfileConfig,
|
||||
task: ProfileTask,
|
||||
|
|
@ -159,84 +218,498 @@ def record_task_result(
|
|||
raise ValueError("task result must match the scheduled task and include manifest_sha256")
|
||||
selected_ids, actual_sha256 = manifest_selection(config.manifest_path, config.manifest_sha256)
|
||||
validate_manifest_attempt(task.task_id, manifest_sha256, selected_ids, actual_sha256)
|
||||
run_id: Final = result.get("run_id")
|
||||
if not isinstance(run_id, str) or not run_id:
|
||||
raise ValueError("task result must include run_id")
|
||||
prior_run_ids: Final = {
|
||||
str(existing["run_id"])
|
||||
for existing in run_pilot.read_jsonl_locked(config.runs_path)
|
||||
if isinstance(existing.get("run_id"), str)
|
||||
}
|
||||
if run_id in prior_run_ids:
|
||||
return
|
||||
if result.get("cost_incomplete") is True:
|
||||
run_pilot.append_jsonl(
|
||||
config.runs_path,
|
||||
{
|
||||
**dict(result),
|
||||
"record_type": "quarantined_attempt",
|
||||
"task_id": task.task_id,
|
||||
"tenant_id": config.tenant_id,
|
||||
"category": task.category,
|
||||
"manifest_sha256": manifest_sha256,
|
||||
"cost_usd": None,
|
||||
"is_infra": True,
|
||||
"success": None,
|
||||
"error_class": "accounting",
|
||||
},
|
||||
)
|
||||
raise ValueError("attempt cost accounting is incomplete")
|
||||
model_alias: Final = result.get("model")
|
||||
if not isinstance(model_alias, str) or model_alias not in task.model_order:
|
||||
raise ValueError("task result model must match the scheduled model order")
|
||||
returned_model: Final = _single_result_value(result, "returned_models") or "unreturned"
|
||||
deployment_id: Final = _single_result_value(result, "deployment_ids")
|
||||
cost_usd: Final = _required_metric(result, "cost_usd")
|
||||
solve_latency: Final = _required_metric(result, "solve_latency_s")
|
||||
eval_latency: Final = _required_metric(result, "eval_latency_s")
|
||||
returned_models: Final = _result_values(result, "returned_models")
|
||||
deployment_ids: Final = _result_values(result, "deployment_ids")
|
||||
latency: Final = solve_latency + eval_latency
|
||||
error_class: Final = result.get("error_class")
|
||||
is_infra: Final = error_class in ("infra", "accounting")
|
||||
success: Final = result.get("success")
|
||||
harness_sha256: Final = hashlib.sha256(Path(run_pilot.__file__).read_bytes()).hexdigest()
|
||||
row: Final = {
|
||||
**result,
|
||||
"record_type": "attempt",
|
||||
"task_id": task.task_id,
|
||||
"tenant_id": config.tenant_id,
|
||||
"category": task.category,
|
||||
"model_alias": model_alias,
|
||||
"returned_model_version": returned_models[0] if len(returned_models) == 1 else "unreturned",
|
||||
"deployment_id": deployment_ids[0] if len(deployment_ids) == 1 else None,
|
||||
"harness_version": versions.harness_version,
|
||||
"budget_version": versions.budget_version,
|
||||
"harness_config_hash": versions.harness_config_hash,
|
||||
"is_infra": is_infra,
|
||||
"success": None if is_infra else success,
|
||||
"cost_usd": cost_usd,
|
||||
"latency_s": latency,
|
||||
"primary_label": task.category,
|
||||
"secondary_label": None,
|
||||
"manifest_sha256": manifest_sha256,
|
||||
"harness_sha256": harness_sha256,
|
||||
"model_order": task.model_order,
|
||||
"returned_model": returned_models[0] if len(returned_models) == 1 else "unreturned",
|
||||
"x_litellm_model_id": deployment_ids[0] if len(deployment_ids) == 1 else None,
|
||||
"error_class": error_class,
|
||||
"raw_result": dict(result),
|
||||
}
|
||||
if not is_infra and not isinstance(success, bool):
|
||||
raise ValueError("graded task results must include a boolean success")
|
||||
solve_latency: Final = result.get("solve_latency_s", 0.0)
|
||||
eval_latency: Final = result.get("eval_latency_s", 0.0)
|
||||
latency: Final = float(solve_latency) + float(eval_latency)
|
||||
harness_sha256: Final = hashlib.sha256(Path(run_pilot.__file__).read_bytes()).hexdigest()
|
||||
run_pilot.append_jsonl(
|
||||
config.runs_path,
|
||||
{
|
||||
**result,
|
||||
"task_id": task.task_id,
|
||||
"tenant_id": config.tenant_id,
|
||||
"category": task.category,
|
||||
"model_alias": model_alias,
|
||||
"returned_model_version": returned_model,
|
||||
"deployment_id": deployment_id,
|
||||
"harness_version": versions.harness_version,
|
||||
"budget_version": versions.budget_version,
|
||||
"harness_config_hash": versions.harness_config_hash,
|
||||
"is_infra": is_infra,
|
||||
"success": None if is_infra else success,
|
||||
"cost_usd": result.get("cost_usd", 0.0),
|
||||
"latency_s": latency,
|
||||
"primary_label": task.category,
|
||||
"secondary_label": None,
|
||||
"manifest_sha256": manifest_sha256,
|
||||
"harness_sha256": harness_sha256,
|
||||
"model_order": task.model_order,
|
||||
"returned_model": returned_model,
|
||||
"x_litellm_model_id": deployment_id,
|
||||
"error_class": error_class,
|
||||
"raw_result": dict(result),
|
||||
},
|
||||
if len(returned_models) > 1 or len(deployment_ids) > 1:
|
||||
run_pilot.append_jsonl(
|
||||
config.runs_path,
|
||||
row
|
||||
| {
|
||||
"record_type": "quarantined_attempt",
|
||||
"error_class": "mixed_version",
|
||||
"is_infra": True,
|
||||
"success": None,
|
||||
"returned_models": returned_models,
|
||||
"deployment_ids": deployment_ids,
|
||||
},
|
||||
)
|
||||
raise MixedVersionAttemptError(task.task_id)
|
||||
run_pilot.append_jsonl(config.runs_path, row)
|
||||
|
||||
|
||||
def _guard_sha256(root_dir: Path) -> str:
|
||||
source_dir: Final = Path(__file__).resolve().parent
|
||||
paths: Final = (
|
||||
source_dir / "guarded_model.py",
|
||||
source_dir / "run_profiling.py",
|
||||
source_dir / "run_pilot.py",
|
||||
root_dir / "poc" / "pilot_models.json",
|
||||
root_dir / "poc" / "pilot.yaml",
|
||||
)
|
||||
digest = hashlib.sha256()
|
||||
for path in paths:
|
||||
digest.update(path.name.encode())
|
||||
digest.update(b"\0")
|
||||
digest.update(path.read_bytes())
|
||||
return digest.hexdigest()
|
||||
|
||||
|
||||
def _profile_versions(root_dir: Path) -> ProfileVersions:
|
||||
model_manifest: Final = _mapping(
|
||||
def _profile_versions(root_dir: Path, runs_path: Path | None = None) -> ProfileVersions:
|
||||
pilot_config: Final = _mapping(
|
||||
json.loads((root_dir / "poc" / "pilot_models.json").read_text()),
|
||||
"pilot model manifest",
|
||||
)
|
||||
solver_data: Final = _mapping(model_manifest.get("solvers"), "pilot solvers")
|
||||
model_versions: Final = tuple(
|
||||
(alias, str(_mapping(solver_data.get(alias), alias)["gateway_model"])) for alias in MODEL_ORDER
|
||||
)
|
||||
model_deployments: Final = tuple(
|
||||
(alias, str(_mapping(solver_data.get(alias), alias)["deployment_id"])) for alias in MODEL_ORDER
|
||||
)
|
||||
run_pilot_sha256: Final = hashlib.sha256((root_dir / "poc" / "run_pilot.py").read_bytes()).hexdigest()
|
||||
config: Final = HarnessConfig(
|
||||
if not isinstance(pilot_config.get("solvers"), Mapping):
|
||||
raise ValueError("pilot model manifest must contain solvers")
|
||||
harness_config: Final = HarnessConfig(
|
||||
harness_version="run_pilot_attempt_v1",
|
||||
budget_version="budget_1usd_50_600_v1",
|
||||
per_run_budget_usd=1.0,
|
||||
step_limit=50,
|
||||
wall_timeout_s=600,
|
||||
mini_swe_agent_version=run_pilot._package_version("mini-swe-agent"),
|
||||
guard_sha256=run_pilot_sha256,
|
||||
guard_sha256=_guard_sha256(root_dir),
|
||||
)
|
||||
rows: Final = run_pilot.read_jsonl_locked(runs_path) if runs_path is not None else []
|
||||
model_versions: Final = tuple(
|
||||
sorted(
|
||||
{
|
||||
(str(row["model_alias"]), str(row["returned_model_version"]))
|
||||
for row in rows
|
||||
if isinstance(row.get("model_alias"), str) and isinstance(row.get("returned_model_version"), str)
|
||||
}
|
||||
)
|
||||
)
|
||||
deployments: Final = tuple(
|
||||
sorted(
|
||||
{
|
||||
(str(row["model_alias"]), str(row["deployment_id"]))
|
||||
for row in rows
|
||||
if isinstance(row.get("model_alias"), str)
|
||||
}
|
||||
)
|
||||
)
|
||||
return ProfileVersions(
|
||||
harness_version=config.harness_version,
|
||||
budget_version=config.budget_version,
|
||||
harness_config_hash=harness_config_sha256(config),
|
||||
model_versions=model_versions,
|
||||
model_deployments=model_deployments,
|
||||
harness_config.harness_version,
|
||||
harness_config.budget_version,
|
||||
harness_config_sha256(harness_config),
|
||||
model_versions,
|
||||
deployments,
|
||||
)
|
||||
|
||||
|
||||
def _load_dataset(name: str, revision: str, split: str) -> tuple[Mapping[str, object], ...]:
|
||||
return tuple(load_dataset(name, revision=revision, split=split))
|
||||
|
||||
|
||||
def _task_data(
|
||||
config: ProfileConfig,
|
||||
tasks: Sequence[ProfileTask],
|
||||
agent_source: Sequence[Mapping[str, object]],
|
||||
eval_source: Sequence[Mapping[str, object]],
|
||||
) -> AttemptPaths:
|
||||
selected_ids: Final = tuple(task.task_id for task in tasks)
|
||||
safe_rows: Final = agent_rows(agent_source, selected_ids)
|
||||
eval_by_id: Final = {
|
||||
row["instance_id"]: row
|
||||
for row in eval_source
|
||||
if isinstance(row.get("instance_id"), str) and row["instance_id"] in selected_ids
|
||||
}
|
||||
if set(eval_by_id) != set(selected_ids):
|
||||
raise ValueError("selected IDs are missing from the pinned evaluator dataset")
|
||||
if any(not isinstance(eval_by_id[task_id].get("patch"), str) for task_id in selected_ids):
|
||||
raise ValueError("pinned evaluator rows must contain gold patches")
|
||||
candidate_by_id: Final = {
|
||||
task.task_id: task
|
||||
for task in tasks
|
||||
if task.image_ref is not None and task.image_digest is not None
|
||||
}
|
||||
if set(candidate_by_id) != set(selected_ids):
|
||||
raise ValueError("selected manifest tasks require digest-pinned image references")
|
||||
task_rows: Final = {
|
||||
str(row["instance_id"]): dict(row)
|
||||
| {"image": candidate_by_id[str(row["instance_id"])].image_ref}
|
||||
for row in safe_rows
|
||||
}
|
||||
private_root: Final = config.root_dir / ".private"
|
||||
private_root.mkdir(parents=True, exist_ok=True, mode=0o700)
|
||||
if stat.S_IMODE(private_root.stat().st_mode) != 0o700:
|
||||
raise PermissionError("private profiling data directory must have mode 0700")
|
||||
data_dir: Final = private_root / "profiling_data"
|
||||
data_dir.mkdir(exist_ok=True, mode=0o700)
|
||||
if stat.S_IMODE(data_dir.stat().st_mode) != 0o700:
|
||||
raise PermissionError("profiling data directory must have mode 0700")
|
||||
agent_path: Final = data_dir / "agent_dataset"
|
||||
agent_path.mkdir(exist_ok=True, mode=0o700)
|
||||
(agent_path / "train.jsonl").write_text(
|
||||
"".join(json.dumps(task_rows[task_id]) + "\n" for task_id in selected_ids)
|
||||
)
|
||||
eval_path: Final = data_dir / "swebench_dev.json"
|
||||
eval_path.write_text(json.dumps([eval_by_id[task_id] for task_id in selected_ids]) + "\n")
|
||||
base_commits: Final = {
|
||||
task_id: str(task_rows[task_id]["base_commit"])
|
||||
for task_id in selected_ids
|
||||
if isinstance(task_rows[task_id].get("base_commit"), str)
|
||||
}
|
||||
image_digests: Final = {
|
||||
str(candidate.image_ref): str(candidate.image_digest)
|
||||
for candidate in candidate_by_id.values()
|
||||
}
|
||||
solver_payload: Final = _mapping(
|
||||
json.loads((config.root_dir / "poc" / "pilot_models.json").read_text()),
|
||||
"pilot model manifest",
|
||||
)
|
||||
solvers: Final = _mapping(solver_payload.get("solvers"), "pilot solvers")
|
||||
solver_configs: Final = {
|
||||
alias: dict(_mapping(solvers.get(alias), alias))
|
||||
for alias in MODEL_ORDER
|
||||
}
|
||||
pilot_paths: Final = run_pilot.AttemptPaths(
|
||||
run_root=private_root / "profiling_runs",
|
||||
dataset_root=agent_path,
|
||||
eval_dataset=str(eval_path),
|
||||
base_commits=base_commits,
|
||||
image_digests=image_digests,
|
||||
image_overlay_ids={},
|
||||
dataset_metadata={
|
||||
"agent_dataset": f"{DATASET_NAME}@{DATASET_REVISION} {DATASET_SPLIT} (sanitized)",
|
||||
"eval_dataset": f"SWE-bench/SWE-bench@{EVAL_DATASET_REVISION} dev",
|
||||
},
|
||||
root_dir=config.root_dir,
|
||||
)
|
||||
return AttemptPaths(
|
||||
pilot_paths,
|
||||
task_rows,
|
||||
{task_id: str(eval_by_id[task_id]["patch"]) for task_id in selected_ids},
|
||||
solver_configs,
|
||||
)
|
||||
|
||||
|
||||
def _default_disk_stats() -> Mapping[str, int]:
|
||||
usage: Final = shutil.disk_usage("/")
|
||||
return {"total_bytes": usage.total, "used_bytes": usage.used, "free_bytes": usage.free}
|
||||
|
||||
|
||||
def _inspect_image(reference: str) -> str | None:
|
||||
result: Final = subprocess.run(
|
||||
["docker", "image", "inspect", "--format", "{{.Id}}", reference],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
check=False,
|
||||
)
|
||||
return result.stdout.strip() if result.returncode == 0 else None
|
||||
|
||||
|
||||
def _pull_image(reference: str) -> None:
|
||||
subprocess.run(["docker", "pull", reference], check=True, capture_output=True, text=True)
|
||||
|
||||
|
||||
def _remove_image(reference: str) -> bool:
|
||||
return subprocess.run(["docker", "rmi", reference], capture_output=True, check=False).returncode == 0
|
||||
|
||||
|
||||
def _evaluate(task_id: str, patch: str, label: str, paths: run_pilot.AttemptPaths,
|
||||
env: dict[str, str]) -> tuple[str, float]:
|
||||
run_id: Final = f"{task_id}__preflight_{label}_{time.time_ns()}"
|
||||
run_dir: Final = paths.run_root / "preflight" / task_id / label
|
||||
return run_pilot._grade(task_id, f"preflight-{label}", patch, run_id, run_dir, env, paths.eval_dataset,
|
||||
paths.root_dir or ROOT)
|
||||
|
||||
|
||||
def _append_preflight(
|
||||
config: ProfileConfig,
|
||||
task: ProfileTask,
|
||||
status: str,
|
||||
error_class: str | None,
|
||||
gold_status: str | None,
|
||||
noop_status: str | None,
|
||||
latency_s: float,
|
||||
) -> None:
|
||||
run_pilot.append_jsonl(
|
||||
config.runs_path,
|
||||
{
|
||||
"record_type": "preflight_failed" if error_class is not None else "preflight_passed",
|
||||
"task_id": task.task_id,
|
||||
"tenant_id": config.tenant_id,
|
||||
"category": task.category,
|
||||
"manifest_sha256": config.manifest_sha256,
|
||||
"status": status,
|
||||
"error_class": error_class,
|
||||
"gold_status": gold_status,
|
||||
"noop_status": noop_status,
|
||||
"latency_s": latency_s,
|
||||
"cost_usd": 0.0,
|
||||
"model_order": task.model_order,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def execute_batch(
|
||||
config: ProfileConfig,
|
||||
batch: ProfileBatch,
|
||||
data: AttemptPaths,
|
||||
versions: ProfileVersions,
|
||||
ops: RuntimeOps,
|
||||
env: dict[str, str],
|
||||
base_url: str,
|
||||
guard: run_pilot.DispatchGuard,
|
||||
) -> None:
|
||||
existing: Final = run_pilot.read_jsonl_locked(config.runs_path)
|
||||
preflight_rows: Final = {
|
||||
str(row["task_id"]): row
|
||||
for row in existing
|
||||
if row.get("record_type") in {"preflight_passed", "preflight_failed"}
|
||||
and isinstance(row.get("task_id"), str)
|
||||
}
|
||||
pulled: Final = _pull_batch(config, batch, ops, preflight_rows)
|
||||
current_rows: Final = run_pilot.read_jsonl_locked(config.runs_path)
|
||||
current_preflight: Final = {
|
||||
str(row["task_id"]): row
|
||||
for row in current_rows
|
||||
if row.get("record_type") in {"preflight_passed", "preflight_failed"}
|
||||
and isinstance(row.get("task_id"), str)
|
||||
}
|
||||
runnable: Final = tuple(
|
||||
task
|
||||
for task in batch.tasks
|
||||
if task.task_id in current_preflight
|
||||
and current_preflight[task.task_id].get("record_type") == "preflight_passed"
|
||||
and task.image_ref is not None
|
||||
and ops.inspect_image(task.image_ref) is not None
|
||||
)
|
||||
new_tasks: Final = tuple(task for task in batch.tasks if task.task_id not in current_preflight)
|
||||
passed_new: Final = tuple(
|
||||
task
|
||||
for task in new_tasks
|
||||
if _preflight_task(config, batch, task, data, ops, env)
|
||||
)
|
||||
solve_tasks: Final = tuple(task for task in runnable + passed_new if task.image_ref is not None)
|
||||
|
||||
def solve_model(task: ProfileTask, model: str) -> None:
|
||||
rows: Final = run_pilot.read_jsonl_locked(config.runs_path)
|
||||
pair: Final = tuple(
|
||||
row
|
||||
for row in rows
|
||||
if row.get("record_type") in {"attempt", "quarantined_attempt"}
|
||||
and row.get("task_id") == task.task_id
|
||||
and row.get("model_alias", row.get("model")) == model
|
||||
)
|
||||
if any(row.get("record_type") == "quarantined_attempt" for row in pair):
|
||||
return
|
||||
if any(row.get("is_infra") is False for row in pair):
|
||||
return
|
||||
if pair and (
|
||||
pair[-1].get("is_infra") is not True or pair[-1].get("retry_eligible") is not True
|
||||
):
|
||||
return
|
||||
attempt_number: Final = 1 + len(pair)
|
||||
if attempt_number > 2 or not guard.acquire(task.task_id):
|
||||
return
|
||||
try:
|
||||
raw: Final = ops.attempt(
|
||||
task.task_id,
|
||||
model,
|
||||
attempt_number,
|
||||
data.task_rows[task.task_id],
|
||||
data.solver_configs[model],
|
||||
None,
|
||||
None,
|
||||
env,
|
||||
base_url,
|
||||
data.pilot,
|
||||
)
|
||||
result: Final = dict(raw) | {
|
||||
"task_id": task.task_id,
|
||||
"manifest_sha256": config.manifest_sha256,
|
||||
"model": model,
|
||||
}
|
||||
record_task_result(config, task, result, versions)
|
||||
if result.get("error_class") == "accounting":
|
||||
guard.stop_for_accounting(task.task_id)
|
||||
finally:
|
||||
guard.release()
|
||||
if result.get("error_class") == "infra" and result.get("retry_eligible") is True:
|
||||
solve_model(task, model)
|
||||
|
||||
def solve_task(task: ProfileTask) -> None:
|
||||
tuple(solve_model(task, model) for model in task.model_order)
|
||||
|
||||
try:
|
||||
with ThreadPoolExecutor(max_workers=WORKERS) as executor:
|
||||
tuple(executor.map(solve_task, solve_tasks))
|
||||
except MixedVersionAttemptError as error:
|
||||
guard.stop(error.task_id, "mixed_version_attempt")
|
||||
finally:
|
||||
tuple(ops.remove_image(reference) for reference in pulled)
|
||||
tuple(
|
||||
(_ for _ in ()).throw(RuntimeError(f"pulled image remains after cleanup: {reference}"))
|
||||
for reference in pulled
|
||||
if ops.image_exists(reference)
|
||||
)
|
||||
run_pilot.append_jsonl(
|
||||
config.profile_dir / "disk_events.jsonl",
|
||||
{"batch": batch.index, "event": "cleanup", **dict(ops.disk_stats())},
|
||||
)
|
||||
_write_checkpoint(config, batch, versions)
|
||||
|
||||
|
||||
def _pull_batch(
|
||||
config: ProfileConfig,
|
||||
batch: ProfileBatch,
|
||||
ops: RuntimeOps,
|
||||
preflight_rows: Mapping[str, Mapping[str, object]],
|
||||
) -> tuple[str, ...]:
|
||||
if not batch.tasks:
|
||||
return ()
|
||||
task: Final = batch.tasks[0]
|
||||
remaining: Final = ProfileBatch(batch.index, batch.tasks[1:])
|
||||
if preflight_rows.get(task.task_id, {}).get("record_type") == "preflight_failed":
|
||||
return _pull_batch(config, remaining, ops, preflight_rows)
|
||||
if task.image_ref is None:
|
||||
if task.task_id not in preflight_rows:
|
||||
_append_preflight(config, task, "preflight_failed", "env", None, None, 0.0)
|
||||
return _pull_batch(config, remaining, ops, preflight_rows)
|
||||
try:
|
||||
existing_image_id: Final = ops.inspect_image(task.image_ref)
|
||||
except OSError:
|
||||
if task.task_id not in preflight_rows:
|
||||
_append_preflight(config, task, "preflight_failed", "infra", None, None, 0.0)
|
||||
return _pull_batch(config, remaining, ops, preflight_rows)
|
||||
if existing_image_id is not None:
|
||||
return _pull_batch(config, remaining, ops, preflight_rows)
|
||||
try:
|
||||
require_disk_for_pull(int(ops.disk_stats().get("free_bytes", 0)))
|
||||
except OSError:
|
||||
tuple(
|
||||
_append_preflight(config, later, "preflight_failed", "infra", None, None, 0.0)
|
||||
for later in batch.tasks
|
||||
if later.task_id not in preflight_rows
|
||||
)
|
||||
return ()
|
||||
try:
|
||||
ops.pull_image(task.image_ref)
|
||||
except Exception:
|
||||
if task.task_id not in preflight_rows:
|
||||
_append_preflight(config, task, "preflight_failed", "infra", None, None, 0.0)
|
||||
return (
|
||||
(task.image_ref,) if ops.inspect_image(task.image_ref) is not None else ()
|
||||
) + _pull_batch(config, remaining, ops, preflight_rows)
|
||||
image_id: Final = ops.inspect_image(task.image_ref)
|
||||
run_pilot.append_jsonl(
|
||||
config.profile_dir / "disk_events.jsonl",
|
||||
{
|
||||
"batch": batch.index,
|
||||
"task_id": task.task_id,
|
||||
"event": "pull",
|
||||
"image_id": image_id,
|
||||
**dict(ops.disk_stats()),
|
||||
},
|
||||
)
|
||||
return (task.image_ref,) + _pull_batch(config, remaining, ops, preflight_rows)
|
||||
|
||||
|
||||
def _preflight_task(
|
||||
config: ProfileConfig,
|
||||
batch: ProfileBatch,
|
||||
task: ProfileTask,
|
||||
data: AttemptPaths,
|
||||
ops: RuntimeOps,
|
||||
env: dict[str, str],
|
||||
) -> bool:
|
||||
started: Final = time.monotonic()
|
||||
try:
|
||||
if task.image_ref is None or ops.inspect_image(task.image_ref) is None:
|
||||
raise OSError("manifest image is not present after pull")
|
||||
gold_status, gold_latency = ops.evaluate(task.task_id, data.gold_patches[task.task_id], "gold",
|
||||
data.pilot)
|
||||
noop_status, noop_latency = ops.evaluate(task.task_id, NOOP_PATCH, "noop", data.pilot)
|
||||
except (OSError, subprocess.SubprocessError):
|
||||
_append_preflight(config, task, "preflight_failed", "infra", None, None, time.monotonic() - started)
|
||||
return False
|
||||
except Exception:
|
||||
_append_preflight(config, task, "preflight_failed", "evaluator", None, None, time.monotonic() - started)
|
||||
return False
|
||||
passed: Final = gold_status == "resolved" and noop_status != "resolved" and noop_status != "error"
|
||||
error_class: Final = None if passed else ("evaluator" if noop_status in {"resolved", "error"} else "env")
|
||||
_append_preflight(
|
||||
config,
|
||||
task,
|
||||
"preflight_passed" if passed else "preflight_failed",
|
||||
error_class,
|
||||
gold_status,
|
||||
noop_status,
|
||||
gold_latency + noop_latency,
|
||||
)
|
||||
return passed
|
||||
|
||||
|
||||
def _write_checkpoint(config: ProfileConfig, batch: ProfileBatch, versions: ProfileVersions) -> None:
|
||||
config.profile_dir.mkdir(parents=True, exist_ok=True)
|
||||
config.runs_path.touch(exist_ok=True)
|
||||
capability_map: Final = compile_map_from_profiling_runs(
|
||||
config.runs_path,
|
||||
config.manifest_path,
|
||||
|
|
@ -244,79 +717,83 @@ def _write_checkpoint(config: ProfileConfig, batch: ProfileBatch, versions: Prof
|
|||
config.tenant_id,
|
||||
versions,
|
||||
)
|
||||
serialized_map: Final = json.dumps(capability_map, indent=2, sort_keys=True) + "\n"
|
||||
(config.profile_dir / f"capability_map_batch_{batch.index:03d}.json").write_text(serialized_map)
|
||||
(config.profile_dir / "capability_map_intermediate.json").write_text(serialized_map)
|
||||
progress: Final = {
|
||||
"dry_run": True,
|
||||
"manifest_sha256": config.manifest_sha256,
|
||||
"tenant_id": config.tenant_id,
|
||||
"completed_batches": batch.index,
|
||||
"next_batch": batch.index + 1,
|
||||
"tasks": [
|
||||
{
|
||||
"task_id": task.task_id,
|
||||
"category": task.category,
|
||||
"model_order": task.model_order,
|
||||
"status": "planned",
|
||||
"preflight_status": "not_run_dry_run",
|
||||
}
|
||||
for task in batch.tasks
|
||||
],
|
||||
}
|
||||
(config.profile_dir / "progress.json").write_text(json.dumps(progress, indent=2) + "\n")
|
||||
config.profile_dir.mkdir(parents=True, exist_ok=True)
|
||||
serialized: Final = json.dumps(capability_map, indent=2, sort_keys=True) + "\n"
|
||||
(config.profile_dir / f"capability_map_batch_{batch.index:03d}.json").write_text(serialized)
|
||||
(config.profile_dir / "capability_map_intermediate.json").write_text(serialized)
|
||||
run_pilot.append_jsonl(
|
||||
config.profile_dir / "progress.jsonl",
|
||||
{"batch": batch.index, "manifest_sha256": config.manifest_sha256, "ts": run_pilot._utc_now()},
|
||||
)
|
||||
|
||||
|
||||
def _completed_batches(config: ProfileConfig) -> int:
|
||||
path: Final = config.profile_dir / "progress.json"
|
||||
if not path.exists():
|
||||
return 0
|
||||
progress: Final = _mapping(json.loads(path.read_text()), "progress")
|
||||
if progress.get("manifest_sha256") != config.manifest_sha256:
|
||||
raise ValueError("progress manifest sha256 mismatch")
|
||||
if progress.get("tenant_id") != config.tenant_id:
|
||||
raise ValueError("progress tenant mismatch")
|
||||
completed: Final = progress.get("completed_batches")
|
||||
if not isinstance(completed, int) or isinstance(completed, bool) or completed < 0:
|
||||
raise ValueError("progress completed_batches must be a non-negative integer")
|
||||
return completed
|
||||
def _default_ops(
|
||||
config: ProfileConfig,
|
||||
data: AttemptPaths,
|
||||
env: dict[str, str],
|
||||
) -> RuntimeOps:
|
||||
return RuntimeOps(
|
||||
disk_stats=_default_disk_stats,
|
||||
inspect_image=_inspect_image,
|
||||
pull_image=_pull_image,
|
||||
evaluate=lambda task_id, patch, label, paths: _evaluate(task_id, patch, label, paths, env),
|
||||
remove_image=_remove_image,
|
||||
image_exists=lambda reference: _inspect_image(reference) is not None,
|
||||
attempt=run_pilot._attempt,
|
||||
)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description="Validate and checkpoint profiling batches without launching solves")
|
||||
parser = argparse.ArgumentParser(description="Run capability profiling in guarded, resumable batches")
|
||||
parser.add_argument("--root", type=Path, default=ROOT)
|
||||
parser.add_argument("--manifest", type=Path)
|
||||
parser.add_argument("--manifest-sha256", default=MANIFEST_SHA256)
|
||||
parser.add_argument("--tenant-id", default=PUBLIC_DEV_TENANT)
|
||||
parser.add_argument("--max-batches", type=int)
|
||||
parser.add_argument("--execute", action="store_true")
|
||||
arguments = parser.parse_args()
|
||||
default_manifest: Final = arguments.root / "profiling" / "profiling_manifest.json"
|
||||
manifest_path: Final = (
|
||||
arguments.manifest
|
||||
if arguments.manifest is not None
|
||||
else default_manifest
|
||||
if default_manifest.exists()
|
||||
else arguments.root / "profiling_manifest.json"
|
||||
manifest_path: Final = arguments.manifest or (
|
||||
default_manifest if default_manifest.exists() else arguments.root / "profiling_manifest.json"
|
||||
)
|
||||
config: Final = ProfileConfig(arguments.root, manifest_path, arguments.manifest_sha256, arguments.tenant_id)
|
||||
completed_batches: Final = _completed_batches(config)
|
||||
pending_batches: Final = profile_batches(config)[completed_batches:]
|
||||
batches: Final = pending_batches if arguments.max_batches is None else pending_batches[: arguments.max_batches]
|
||||
if arguments.max_batches is not None and arguments.max_batches < 1:
|
||||
raise ValueError("max_batches must be positive")
|
||||
versions: Final = _profile_versions(config.root_dir)
|
||||
for batch in batches:
|
||||
_write_checkpoint(config, batch, versions)
|
||||
print(
|
||||
json.dumps(
|
||||
{
|
||||
"dry_run": True,
|
||||
"batches": len(batches),
|
||||
"planned_runs": sum(len(batch.tasks) * len(MODEL_ORDER) for batch in batches),
|
||||
"workers": WORKERS,
|
||||
}
|
||||
batches: Final = profile_batches(config, arguments.max_batches)
|
||||
if not arguments.execute:
|
||||
print(
|
||||
json.dumps(
|
||||
{
|
||||
"dry_run": True,
|
||||
"guard_path": "poc/run_pilot.py:DispatchGuard.acquire -> global_budget_allows",
|
||||
"batches": [
|
||||
{
|
||||
"batch": batch.index,
|
||||
"tasks": [
|
||||
{"task_id": task.task_id, "model_order": task.model_order} for task in batch.tasks
|
||||
],
|
||||
}
|
||||
for batch in batches
|
||||
],
|
||||
"workers": WORKERS,
|
||||
},
|
||||
indent=2,
|
||||
)
|
||||
)
|
||||
)
|
||||
return
|
||||
versions: Final = _profile_versions(config.root_dir, config.runs_path)
|
||||
agent_source: Final = tuple(load_dataset(DATASET_NAME, revision=DATASET_REVISION, split=DATASET_SPLIT))
|
||||
eval_source: Final = _load_dataset("SWE-bench/SWE-bench", EVAL_DATASET_REVISION, "dev")
|
||||
data: Final = _task_data(config, tuple(task for batch in batches for task in batch.tasks), agent_source, eval_source)
|
||||
key_path: Final = config.root_dir / ".private" / "profiling_key"
|
||||
env: Final = os.environ.copy()
|
||||
env.pop("LITELLM_GATEWAY_ADMIN_KEY", None)
|
||||
env["LITELLM_PROXY_API_KEY"] = run_pilot._read_pilot_key(key_path)
|
||||
env["PYTHONPATH"] = str(config.root_dir)
|
||||
env["MSWEA_MODEL_RETRY_STOP_AFTER_ATTEMPT"] = "1"
|
||||
base_url: Final = os.environ.get("LITELLM_GATEWAY_BASE_URL", "https://gateway.litellm-sandbox.ai").rstrip("/")
|
||||
guard: Final = profile_guard(config, free_bytes=lambda: _default_disk_stats()["free_bytes"])
|
||||
ops: Final = _default_ops(config, data, env)
|
||||
for batch in batches:
|
||||
execute_batch(config, batch, data, versions, ops, env, base_url, guard)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
|
|
|||
2
cookbook/capability_routing_pilot/pytest.ini
Normal file
2
cookbook/capability_routing_pilot/pytest.ini
Normal file
|
|
@ -0,0 +1,2 @@
|
|||
[pytest]
|
||||
pythonpath = .
|
||||
|
|
@ -479,6 +479,81 @@ def test_map_builder_rejects_attempt_membership_or_digest_mismatch(
|
|||
)
|
||||
|
||||
|
||||
def test_final_map_requires_every_selected_attempt_or_accounted_infra_deficit(tmp_path: Path) -> None:
|
||||
manifest = tmp_path / "manifest.json"
|
||||
manifest_bytes = json.dumps(
|
||||
{
|
||||
"categories": {
|
||||
"repository_debugging": {"selected_ids": ["task-0"]},
|
||||
"feature_implementation": {"selected_ids": []},
|
||||
"shell_build": {"selected_ids": []},
|
||||
"algorithms_data": {"selected_ids": []},
|
||||
"unknown": {"selected_ids": []},
|
||||
}
|
||||
}
|
||||
).encode()
|
||||
manifest.write_bytes(manifest_bytes)
|
||||
manifest_sha256 = hashlib.sha256(manifest_bytes).hexdigest()
|
||||
versions = _versions(
|
||||
models=(
|
||||
("haiku-4.5", MODEL_VERSION),
|
||||
("sonnet-5", "anthropic/claude-sonnet-5"),
|
||||
("opus-5", "anthropic/claude-opus-5"),
|
||||
),
|
||||
deployments=(
|
||||
("haiku-4.5", "deployment-haiku"),
|
||||
("sonnet-5", "deployment-sonnet"),
|
||||
("opus-5", "deployment-opus"),
|
||||
),
|
||||
)
|
||||
preflight = {
|
||||
"record_type": "preflight_passed",
|
||||
"task_id": "task-0",
|
||||
"manifest_sha256": manifest_sha256,
|
||||
}
|
||||
attempt_rows = tuple(
|
||||
{
|
||||
**asdict(_attempts(alias, 1, 1, 0.1, returned_model_version=version, deployment_id=deployment)[0]),
|
||||
**asdict(_attempts(alias, 1, 1, 0.1, returned_model_version=version, deployment_id=deployment)[0].key),
|
||||
"record_type": "attempt",
|
||||
"model_alias": alias,
|
||||
"model": alias,
|
||||
"run_id": f"run-{alias}",
|
||||
"manifest_sha256": manifest_sha256,
|
||||
}
|
||||
for alias, version, deployment in (
|
||||
("haiku-4.5", MODEL_VERSION, "deployment-haiku"),
|
||||
("sonnet-5", "anthropic/claude-sonnet-5", "deployment-sonnet"),
|
||||
("opus-5", "anthropic/claude-opus-5", "deployment-opus"),
|
||||
)
|
||||
)
|
||||
profiling_runs = tmp_path / "profiling_runs.jsonl"
|
||||
profiling_runs.write_text(
|
||||
"\n".join(json.dumps(row) for row in (preflight,) + attempt_rows) + "\n"
|
||||
)
|
||||
|
||||
complete_map = compile_map_from_profiling_runs(
|
||||
profiling_runs,
|
||||
manifest,
|
||||
manifest_sha256,
|
||||
PUBLIC_DEV_TENANT,
|
||||
versions,
|
||||
)
|
||||
profiling_runs.write_text(
|
||||
"\n".join(json.dumps(row) for row in (preflight,) + attempt_rows[:-1]) + "\n"
|
||||
)
|
||||
incomplete_map = compile_map_from_profiling_runs(
|
||||
profiling_runs,
|
||||
manifest,
|
||||
manifest_sha256,
|
||||
PUBLIC_DEV_TENANT,
|
||||
versions,
|
||||
)
|
||||
|
||||
assert complete_map["final"] is True
|
||||
assert incomplete_map["final"] is False
|
||||
|
||||
|
||||
def test_harness_hash_is_canonical_json_over_the_full_budget_config() -> None:
|
||||
config = HarnessConfig(
|
||||
harness_version="run_pilot_attempt_v1",
|
||||
|
|
|
|||
|
|
@ -4,10 +4,16 @@ from pathlib import Path
|
|||
|
||||
import pytest
|
||||
|
||||
from poc import run_pilot
|
||||
from poc.run_profiling import (
|
||||
PHASE_BUDGET_USD,
|
||||
AttemptPaths,
|
||||
ProfileConfig,
|
||||
ProfileBatch,
|
||||
ProfileTask,
|
||||
RuntimeOps,
|
||||
MixedVersionAttemptError,
|
||||
execute_batch,
|
||||
model_order_for_task,
|
||||
profile_batches,
|
||||
profile_guard,
|
||||
|
|
@ -86,6 +92,7 @@ def test_runner_record_boundary_checks_membership_and_manifest_sha(tmp_path: Pat
|
|||
task,
|
||||
{
|
||||
"task_id": "task-0",
|
||||
"run_id": "run-0",
|
||||
"manifest_sha256": config.manifest_sha256,
|
||||
"model": "haiku-4.5",
|
||||
"returned_models": ["anthropic/claude-haiku-4-5-20251001"],
|
||||
|
|
@ -107,6 +114,189 @@ def test_runner_record_boundary_checks_membership_and_manifest_sha(tmp_path: Pat
|
|||
assert result["raw_result"]["model"] == "haiku-4.5"
|
||||
|
||||
|
||||
def test_attempt_cost_and_latency_are_required(tmp_path: Path) -> None:
|
||||
config = _config(tmp_path, ("task-0",))
|
||||
task = ProfileTask(0, "task-0", "repository_debugging", model_order_for_task(0))
|
||||
versions = ProfileVersions("harness-v1", "budget-v1", "config-hash", ())
|
||||
result = {
|
||||
"task_id": "task-0",
|
||||
"run_id": "run-0",
|
||||
"manifest_sha256": config.manifest_sha256,
|
||||
"model": "haiku-4.5",
|
||||
"returned_models": ["concrete-haiku"],
|
||||
"deployment_ids": ["deployment-1"],
|
||||
"error_class": "solver",
|
||||
"success": True,
|
||||
"solve_latency_s": 3.0,
|
||||
"eval_latency_s": 2.0,
|
||||
}
|
||||
|
||||
with pytest.raises(ValueError, match="cost_usd must be finite"):
|
||||
record_task_result(config, task, result, versions)
|
||||
assert not config.runs_path.exists()
|
||||
|
||||
|
||||
def test_mixed_versions_are_quarantined_with_cost_and_stop_signal(tmp_path: Path) -> None:
|
||||
config = _config(tmp_path, ("task-0",))
|
||||
task = ProfileTask(0, "task-0", "repository_debugging", model_order_for_task(0))
|
||||
versions = ProfileVersions("harness-v1", "budget-v1", "config-hash", ())
|
||||
result = {
|
||||
"task_id": "task-0",
|
||||
"run_id": "run-mixed",
|
||||
"manifest_sha256": config.manifest_sha256,
|
||||
"model": "haiku-4.5",
|
||||
"returned_models": ["concrete-haiku-a", "concrete-haiku-b"],
|
||||
"deployment_ids": ["deployment-1"],
|
||||
"error_class": "solver",
|
||||
"success": True,
|
||||
"cost_usd": 0.37,
|
||||
"solve_latency_s": 3.0,
|
||||
"eval_latency_s": 2.0,
|
||||
}
|
||||
|
||||
with pytest.raises(MixedVersionAttemptError):
|
||||
record_task_result(config, task, result, versions)
|
||||
row = json.loads(config.runs_path.read_text())
|
||||
assert row["record_type"] == "quarantined_attempt"
|
||||
assert row["error_class"] == "mixed_version"
|
||||
assert row["cost_usd"] == 0.37
|
||||
assert row["success"] is None
|
||||
|
||||
|
||||
def test_injected_batch_preflight_cleanup_and_attempts(tmp_path: Path) -> None:
|
||||
task_id = "task-0"
|
||||
image_ref = "swebench/example@sha256:abc"
|
||||
manifest = tmp_path / "profiling_manifest.json"
|
||||
manifest_row = {
|
||||
"categories": {
|
||||
"repository_debugging": {"selected_ids": [task_id]},
|
||||
"feature_implementation": {"selected_ids": []},
|
||||
"shell_build": {"selected_ids": []},
|
||||
"algorithms_data": {"selected_ids": []},
|
||||
"unknown": {"selected_ids": []},
|
||||
},
|
||||
"candidates": [
|
||||
{"task_id": task_id, "image": "swebench/example:latest", "image_digest": "sha256:abc"}
|
||||
],
|
||||
}
|
||||
manifest_bytes = json.dumps(manifest_row).encode()
|
||||
manifest.write_bytes(manifest_bytes)
|
||||
config = ProfileConfig(tmp_path, manifest, hashlib.sha256(manifest_bytes).hexdigest())
|
||||
task = ProfileTask(0, task_id, "repository_debugging", model_order_for_task(0), image_ref, "sha256:abc")
|
||||
batch = ProfileBatch(1, (task,))
|
||||
pilot_paths = run_pilot.AttemptPaths(
|
||||
run_root=tmp_path / "raw",
|
||||
dataset_root=tmp_path / "agent_dataset",
|
||||
eval_dataset=str(tmp_path / "eval.json"),
|
||||
base_commits={task_id: "commit"},
|
||||
image_digests={image_ref: "sha256:abc"},
|
||||
image_overlay_ids={},
|
||||
dataset_metadata={},
|
||||
root_dir=tmp_path,
|
||||
)
|
||||
data = AttemptPaths(
|
||||
pilot_paths,
|
||||
{task_id: {"instance_id": task_id, "image": image_ref}},
|
||||
{task_id: "gold patch"},
|
||||
{model: {} for model in model_order_for_task(0)},
|
||||
)
|
||||
images: list[str] = []
|
||||
evaluated: list[str] = []
|
||||
removed: list[str] = []
|
||||
attempted: list[tuple[str, int]] = []
|
||||
|
||||
def inspect(reference: str) -> str | None:
|
||||
return "sha256:image-id" if reference in images else None
|
||||
|
||||
def pull(reference: str) -> None:
|
||||
images.append(reference)
|
||||
|
||||
def evaluate(_task: str, _patch: str, label: str, _paths: object) -> tuple[str, float]:
|
||||
evaluated.append(label)
|
||||
return ("resolved" if label == "gold" else "unresolved", 1.0)
|
||||
|
||||
def remove(reference: str) -> bool:
|
||||
images.remove(reference)
|
||||
removed.append(reference)
|
||||
return True
|
||||
|
||||
def attempt(
|
||||
task_id_arg: str,
|
||||
model: str,
|
||||
number: int,
|
||||
_task_row: dict[str, object],
|
||||
_model_data: dict[str, object],
|
||||
_offline: object,
|
||||
_jev: object,
|
||||
_env: dict[str, str],
|
||||
_base_url: str,
|
||||
_paths: object,
|
||||
) -> dict[str, object]:
|
||||
attempted.append((model, number))
|
||||
return {
|
||||
"task_id": task_id_arg,
|
||||
"run_id": f"{model}-{number}",
|
||||
"model": model,
|
||||
"returned_models": [f"concrete-{model}"],
|
||||
"deployment_ids": [f"deployment-{model}"],
|
||||
"error_class": "solver",
|
||||
"success": True,
|
||||
"cost_usd": 0.1,
|
||||
"solve_latency_s": 1.0,
|
||||
"eval_latency_s": 1.0,
|
||||
"retry_eligible": False,
|
||||
}
|
||||
|
||||
operations = RuntimeOps(
|
||||
disk_stats=lambda: {"free_bytes": 100 * 1024**3},
|
||||
inspect_image=inspect,
|
||||
pull_image=pull,
|
||||
evaluate=evaluate,
|
||||
remove_image=remove,
|
||||
image_exists=lambda reference: reference in images,
|
||||
attempt=attempt,
|
||||
)
|
||||
versions = ProfileVersions("harness-v1", "budget-v1", "config-hash", ())
|
||||
guard = profile_guard(
|
||||
config,
|
||||
free_bytes=lambda: 100 * 1024**3,
|
||||
committed_spend=lambda: 0.0,
|
||||
)
|
||||
for model, run_id, error_class, success, retry_eligible in (
|
||||
("haiku-4.5", "seed-graded", "solver", True, False),
|
||||
("sonnet-5", "seed-infra", "infra", None, True),
|
||||
):
|
||||
record_task_result(
|
||||
config,
|
||||
task,
|
||||
{
|
||||
"task_id": task_id,
|
||||
"run_id": run_id,
|
||||
"manifest_sha256": config.manifest_sha256,
|
||||
"model": model,
|
||||
"returned_models": [f"concrete-{model}"],
|
||||
"deployment_ids": [f"deployment-{model}"],
|
||||
"error_class": error_class,
|
||||
"success": success,
|
||||
"cost_usd": 0.1,
|
||||
"solve_latency_s": 1.0,
|
||||
"eval_latency_s": 1.0,
|
||||
"retry_eligible": retry_eligible,
|
||||
},
|
||||
versions,
|
||||
)
|
||||
|
||||
execute_batch(config, batch, data, versions, operations, {}, "https://gateway.invalid", guard)
|
||||
|
||||
assert evaluated == ["gold", "noop"]
|
||||
assert removed == [image_ref]
|
||||
assert not images
|
||||
records = [json.loads(line) for line in config.runs_path.read_text().splitlines()]
|
||||
assert sum(row["record_type"] == "attempt" for row in records) == 4
|
||||
assert attempted == [("sonnet-5", 2), ("opus-5", 1)]
|
||||
assert (config.profile_dir / "capability_map_intermediate.json").exists()
|
||||
|
||||
|
||||
def test_profile_guard_enforces_phase_cap_and_pre_pull_disk_reserve(tmp_path: Path) -> None:
|
||||
config = _config(tmp_path, ("task-0",))
|
||||
guard = profile_guard(
|
||||
|
|
@ -119,3 +309,12 @@ def test_profile_guard_enforces_phase_cap_and_pre_pull_disk_reserve(tmp_path: Pa
|
|||
require_disk_for_pull(60 * 1024**3)
|
||||
with pytest.raises(OSError, match="less than 60 GiB"):
|
||||
require_disk_for_pull(60 * 1024**3 - 1)
|
||||
|
||||
|
||||
def test_phase_guard_counts_failed_retry_spend(tmp_path: Path) -> None:
|
||||
config = _config(tmp_path, ("task-0",))
|
||||
run_pilot.append_jsonl(config.runs_path, {"cost_usd": 179.5, "error_class": "infra"})
|
||||
run_pilot.append_jsonl(config.runs_path, {"cost_usd": 0.5, "error_class": "solver"})
|
||||
guard = profile_guard(config, free_bytes=lambda: 100 * 1024**3)
|
||||
|
||||
assert not guard.acquire("task-after-retries")
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue