| from __future__ import annotations |
|
|
| import asyncio |
| import json |
| from pathlib import Path |
|
|
| from autocad_bench.orchestration import resume as resume_batch |
|
|
|
|
| class _FakeEc2: |
| def __init__(self, instances: list[dict]) -> None: |
| self.instances = instances |
|
|
| def describe_instances(self, **_kwargs): |
| return {"Reservations": [{"Instances": self.instances}]} |
|
|
|
|
| def _interrupted_rollout(batch: Path, task_id: str, instance_id: str) -> None: |
| rollout = batch / "rollouts" / task_id |
| rollout.mkdir(parents=True) |
| (rollout / "run-state.json").write_text( |
| json.dumps( |
| { |
| "task_id": task_id, |
| "provider": "openai", |
| "model_id": "gpt-test", |
| } |
| ) |
| ) |
| logs = batch / "logs" |
| logs.mkdir() |
| (logs / f"{task_id}.log").write_text( |
| json.dumps( |
| { |
| "event": "sandbox_instance_launched", |
| "instance_id": instance_id, |
| } |
| ) |
| + "\n" |
| ) |
|
|
|
|
| def test_discovery_accepts_parent_of_multiple_benchmark_runs(tmp_path: Path) -> None: |
| benchmarks = tmp_path / "benchmarks" |
| _interrupted_rollout(benchmarks / "run-a", "task-001", "i-a") |
| _interrupted_rollout(benchmarks / "run-b", "task-002", "i-b") |
|
|
| plan = resume_batch._discover_plan([benchmarks]) |
|
|
| assert {entry["rollout_id"] for entry in plan["entries"]} == { |
| "run-a/task-001", |
| "run-b/task-002", |
| } |
|
|
|
|
| def test_empty_reconciliation_scan_exits_successfully(tmp_path: Path, capsys) -> None: |
| benchmarks = tmp_path / "benchmarks" |
| (benchmarks / "finished" / "rollouts").mkdir(parents=True) |
|
|
| exit_code = resume_batch.main(["--run-root", str(benchmarks)]) |
|
|
| assert exit_code == 0 |
| assert json.loads(capsys.readouterr().out) == { |
| "entries": 0, |
| "message": "no interrupted rollout directories were discovered", |
| "state": "idle", |
| } |
|
|
|
|
| def test_completed_execution_with_incomplete_evaluation_is_cleanup_only( |
| tmp_path: Path, |
| ) -> None: |
| batch = tmp_path / "benchmarks" / "run-a" |
| _interrupted_rollout(batch, "task-001", "i-a") |
| rollout = batch / "rollouts" / "task-001" |
| (rollout / "result.json").write_text( |
| json.dumps( |
| { |
| "task_id": "task-001", |
| "completed": True, |
| "stop_reason": "done", |
| "artifact_bytes": 100, |
| } |
| ) |
| ) |
|
|
| plan = resume_batch._discover_plan([tmp_path / "benchmarks"]) |
|
|
| assert [entry["rollout_id"] for entry in plan["entries"]] == [ |
| "run-a/task-001" |
| ] |
| assert plan["entries"][0]["cleanup_only"] is True |
|
|
|
|
| def test_deferred_vision_only_retry_does_not_hold_a_windows_session( |
| tmp_path: Path, |
| ) -> None: |
| batch = tmp_path / "benchmarks" / "run-a" |
| _interrupted_rollout(batch, "task-001", "i-a") |
| rollout = batch / "rollouts" / "task-001" |
| (rollout / "result.json").write_text( |
| json.dumps( |
| { |
| "task_id": "task-001", |
| "completed": True, |
| "stop_reason": "done", |
| "artifact_bytes": 100, |
| "evaluation": {"status": "retry_pending"}, |
| } |
| ) |
| ) |
|
|
| try: |
| resume_batch._discover_plan([tmp_path / "benchmarks"]) |
| except resume_batch.NoInterruptedRollouts: |
| pass |
| else: |
| raise AssertionError("vision-only retry unexpectedly retained the worker") |
|
|
|
|
| def test_budget_exhaustion_does_not_restart_the_model(tmp_path: Path) -> None: |
| batch = tmp_path / "benchmarks" / "run-a" |
| _interrupted_rollout(batch, "task-001", "i-a") |
| rollout = batch / "rollouts" / "task-001" |
| (batch / "batch-state.json").write_text( |
| json.dumps({"batch_id": "batch-a"}) |
| ) |
| (rollout / "result.json").write_text( |
| json.dumps( |
| { |
| "task_id": "task-001", |
| "execution_status": "budget_exhausted", |
| "stop_reason": "max_actions", |
| "artifact_bytes": 100, |
| "evaluation": {"status": "not_started"}, |
| } |
| ) |
| ) |
| ec2 = _FakeEc2( |
| [ |
| { |
| "InstanceId": "i-a", |
| "Tags": [ |
| {"Key": "BatchId", "Value": "batch-a"}, |
| {"Key": "RolloutId", "Value": "task-001"}, |
| ], |
| } |
| ] |
| ) |
|
|
| plan = resume_batch._discover_plan([batch.parent], ec2=ec2) |
|
|
| assert plan["entries"][0]["cleanup_only"] is True |
|
|
|
|
| def test_malformed_model_failure_mislabeled_as_infrastructure_is_cleanup_only( |
| tmp_path: Path, |
| ) -> None: |
| batch = tmp_path / "benchmarks" / "run-a" |
| _interrupted_rollout(batch, "task-001", "i-a") |
| rollout = batch / "rollouts" / "task-001" |
| (rollout / "result.json").write_text( |
| json.dumps( |
| { |
| "execution_status": "infrastructure_failed", |
| "stop_reason": "invalid_action_limit", |
| "failure": "Model returned invalid computer-action arguments", |
| "artifact_bytes": 0, |
| } |
| ) |
| ) |
|
|
| plan = resume_batch._discover_plan([batch.parent]) |
|
|
| assert plan["entries"][0]["cleanup_only"] is True |
| assert plan["entries"][0]["recovery_exhausted"] is False |
|
|
|
|
| def test_recovery_attempt_limit_reaps_worker_without_another_model_turn( |
| tmp_path: Path, |
| ) -> None: |
| batch = tmp_path / "benchmarks" / "run-a" |
| _interrupted_rollout(batch, "task-001", "i-a") |
| rollout = batch / "rollouts" / "task-001" |
| result = { |
| "stop_reason": "infrastructure_error", |
| "failure": "BrokerError: transport unavailable", |
| "artifact_bytes": 0, |
| } |
| (rollout / "result.json").write_text(json.dumps(result)) |
| (batch / "cloud-resume-state.json").write_text( |
| json.dumps( |
| { |
| "entries": { |
| "run-a/task-001": { |
| "state": "infrastructure_failed", |
| "recovery_attempts": 3, |
| "last_result_fingerprint": ( |
| resume_batch._result_recovery_fingerprint(result) |
| ), |
| } |
| } |
| } |
| ) |
| ) |
|
|
| plan = resume_batch._discover_plan([batch.parent]) |
|
|
| assert plan["entries"][0]["cleanup_only"] is True |
| assert plan["entries"][0]["recovery_exhausted"] is True |
|
|
|
|
| def test_discovery_caps_subprocess_recovery_concurrency(tmp_path: Path) -> None: |
| benchmarks = tmp_path / "benchmarks" |
| for index in range(6): |
| _interrupted_rollout( |
| benchmarks / f"run-{index}", |
| f"basic-00{index + 1}", |
| f"i-{index}", |
| ) |
|
|
| plan = resume_batch._discover_plan([benchmarks]) |
|
|
| assert len(plan["entries"]) == 6 |
| assert plan["max_concurrency"] == 1 |
|
|
|
|
| def test_discovery_uses_ec2_tags_when_launch_log_was_not_persisted( |
| tmp_path: Path, |
| ) -> None: |
| batch = tmp_path / "benchmarks" / "run-a" |
| _interrupted_rollout(batch, "task-001", "i-from-log") |
| (batch / "logs" / "task-001.log").write_text("") |
| (batch / "batch-state.json").write_text( |
| json.dumps({"batch_id": "batch-a"}) |
| ) |
| ec2 = _FakeEc2( |
| [ |
| { |
| "InstanceId": "i-from-tags", |
| "Tags": [ |
| {"Key": "BatchId", "Value": "batch-a"}, |
| {"Key": "RolloutId", "Value": "task-001"}, |
| ], |
| } |
| ] |
| ) |
|
|
| plan = resume_batch._discover_plan([batch.parent], ec2=ec2) |
|
|
| assert plan["entries"][0]["instance_id"] == "i-from-tags" |
|
|
|
|
| def test_discovery_skips_rollout_owned_by_live_primary_process( |
| tmp_path: Path, monkeypatch |
| ) -> None: |
| batch = tmp_path / "benchmarks" / "run-a" |
| _interrupted_rollout(batch, "task-001", "i-a") |
| rollout = batch / "rollouts" / "task-001" |
| (batch / "batch-state.json").write_text( |
| json.dumps( |
| { |
| "batch_id": "batch-a", |
| "rollouts": { |
| "task-001": { |
| "state": "running", |
| "pid": 1234, |
| } |
| }, |
| } |
| ) |
| ) |
| monkeypatch.setattr(resume_batch, "_pid_owns_rollout", lambda *_args, **_kwargs: True) |
|
|
| try: |
| resume_batch._discover_plan([batch.parent]) |
| except resume_batch.NoInterruptedRollouts: |
| pass |
| else: |
| raise AssertionError("live primary rollout was scheduled for duplicate resume") |
|
|
|
|
| def test_shared_primary_pid_is_bound_to_batch_root( |
| tmp_path: Path, monkeypatch |
| ) -> None: |
| batch = tmp_path / "benchmarks" / "run-a" |
| rollout = batch / "rollouts" / "task-001" |
| rollout.mkdir(parents=True) |
| checked_paths: list[Path] = [] |
|
|
| def owns(_pid: object, path: Path) -> bool: |
| checked_paths.append(path) |
| return True |
|
|
| monkeypatch.setattr(resume_batch, "_pid_owns_rollout", owns) |
|
|
| assert resume_batch._primary_rollout_is_active( |
| { |
| "rollouts": { |
| "task-001": { |
| "state": "running", |
| "pid": 1234, |
| "controller_execution_mode": "shared", |
| } |
| } |
| }, |
| rollout_id="task-001", |
| output_dir=rollout, |
| ) |
| assert checked_paths == [batch] |
|
|
|
|
| def test_discovery_recovers_active_row_after_primary_process_dies( |
| tmp_path: Path, monkeypatch |
| ) -> None: |
| batch = tmp_path / "benchmarks" / "run-a" |
| _interrupted_rollout(batch, "task-001", "i-a") |
| (batch / "batch-state.json").write_text( |
| json.dumps( |
| { |
| "batch_id": "batch-a", |
| "rollouts": { |
| "task-001": { |
| "state": "running", |
| "pid": 1234, |
| } |
| }, |
| } |
| ) |
| ) |
| monkeypatch.setattr( |
| resume_batch, |
| "_pid_owns_rollout", |
| lambda *_args, **_kwargs: False, |
| ) |
| ec2 = _FakeEc2( |
| [ |
| { |
| "InstanceId": "i-a", |
| "Tags": [ |
| {"Key": "BatchId", "Value": "batch-a"}, |
| {"Key": "RolloutId", "Value": "task-001"}, |
| ], |
| } |
| ] |
| ) |
|
|
| plan = resume_batch._discover_plan([batch.parent], ec2=ec2) |
|
|
| assert [entry["rollout_id"] for entry in plan["entries"]] == [ |
| "run-a/task-001" |
| ] |
|
|
|
|
| def test_discovery_skips_rollout_owned_by_live_resume_process( |
| tmp_path: Path, monkeypatch |
| ) -> None: |
| batch = tmp_path / "benchmarks" / "run-a" |
| _interrupted_rollout(batch, "task-001", "i-a") |
| (batch / "cloud-resume-state.json").write_text( |
| json.dumps( |
| { |
| "entries": { |
| "run-a/task-001": { |
| "state": "running", |
| "pid": 5678, |
| } |
| } |
| } |
| ) |
| ) |
| monkeypatch.setattr(resume_batch, "_pid_owns_rollout", lambda *_args, **_kwargs: True) |
|
|
| try: |
| resume_batch._discover_plan([batch.parent]) |
| except resume_batch.NoInterruptedRollouts: |
| pass |
| else: |
| raise AssertionError("live recovery rollout was scheduled for duplicate resume") |
|
|
|
|
| def test_discovery_does_not_reopen_terminated_tagged_worker(tmp_path: Path) -> None: |
| batch = tmp_path / "benchmarks" / "run-a" |
| _interrupted_rollout(batch, "task-001", "i-terminated") |
| (batch / "batch-state.json").write_text( |
| json.dumps({"batch_id": "batch-a"}) |
| ) |
|
|
| try: |
| resume_batch._discover_plan([batch.parent], ec2=_FakeEc2([])) |
| except resume_batch.NoInterruptedRollouts: |
| pass |
| else: |
| raise AssertionError("terminated tagged worker was scheduled for resume") |
|
|
|
|
| def test_discovery_reaps_a_live_worker_only_after_evaluation_is_durable( |
| tmp_path: Path, |
| ) -> None: |
| batch = tmp_path / "benchmarks" / "run-a" |
| _interrupted_rollout(batch, "task-001", "i-a") |
| (batch / "batch-state.json").write_text( |
| json.dumps({"batch_id": "batch-a"}) |
| ) |
| rollout = batch / "rollouts" / "task-001" |
| (rollout / "result.json").write_text( |
| json.dumps( |
| { |
| "task_id": "task-001", |
| "completed": True, |
| "artifact_bytes": 100, |
| "evaluation": {"status": "completed"}, |
| } |
| ) |
| ) |
| ec2 = _FakeEc2( |
| [ |
| { |
| "InstanceId": "i-a", |
| "Tags": [{"Key": "RolloutId", "Value": "task-001"}], |
| } |
| ] |
| ) |
|
|
| plan = resume_batch._discover_plan([batch.parent], ec2=ec2) |
|
|
| assert plan["entries"][0]["cleanup_only"] is True |
|
|
|
|
| def test_released_without_checkpoint_can_restart_fresh_on_same_worker( |
| tmp_path: Path, |
| ) -> None: |
| log = tmp_path / "resume.log" |
| log.write_text("BrokerError: checkpoint is required\n") |
| plan = {"aws_profile": "test-profile", "aws_region": "us-east-1"} |
| entry = { |
| "task_id": "task-001", |
| "instance_id": "i-a", |
| "output_dir": str(tmp_path / "replacement"), |
| "provider": "openai", |
| "model_id": "gpt-test", |
| "resume": False, |
| } |
|
|
| command = resume_batch._command(plan, entry) |
|
|
| assert "--resume" not in command |
| assert resume_batch._needs_fresh_same_worker_recovery( |
| exit_code=1, |
| result=None, |
| log_path=log, |
| ) |
|
|
|
|
| def test_ready_session_before_first_action_restarts_fresh_on_same_worker( |
| tmp_path: Path, |
| ) -> None: |
| log = tmp_path / "resume.log" |
| log.write_text("resume session is not running: ready\n") |
|
|
| assert resume_batch._needs_fresh_same_worker_recovery( |
| exit_code=1, |
| result=None, |
| log_path=log, |
| ) |
|
|
|
|
| def test_openrouter_recovery_preserves_required_tool_choice(tmp_path: Path) -> None: |
| plan = {"aws_profile": "test-profile", "aws_region": "us-east-1"} |
| entry = { |
| "task_id": "task-001", |
| "instance_id": "i-a", |
| "output_dir": str(tmp_path / "replacement"), |
| "provider": "openrouter", |
| "model_id": "moonshotai/kimi-k3", |
| "resume": False, |
| } |
|
|
| command = resume_batch._command(plan, entry) |
|
|
| index = command.index("--tool-choice-mode") |
| assert command[index + 1] == "required" |
|
|
|
|
| def test_legacy_aws_plan_checks_and_cleans_worker_through_plugin( |
| tmp_path: Path, |
| ) -> None: |
| class Ec2: |
| def __init__(self) -> None: |
| self.terminated: list[str] = [] |
|
|
| def describe_instances(self, **_kwargs): |
| return { |
| "Reservations": [ |
| { |
| "Instances": [ |
| { |
| "InstanceId": "i-a", |
| "State": {"Name": "running"}, |
| } |
| ] |
| } |
| ] |
| } |
|
|
| def terminate_instances(self, *, InstanceIds): |
| self.terminated.extend(InstanceIds) |
| return {} |
|
|
| output_dir = tmp_path / "rollout" |
| output_dir.mkdir() |
| (output_dir / "result.json").write_text( |
| json.dumps( |
| { |
| "completed": True, |
| "artifact_bytes": 1, |
| "evaluation": {"status": "completed"}, |
| } |
| ), |
| encoding="utf-8", |
| ) |
| state_path = tmp_path / "resume-state.json" |
| plan_path = tmp_path / "resume-plan.json" |
| plan = { |
| "schema_version": 1, |
| "aws_profile": "test-profile", |
| "aws_region": "us-east-1", |
| "state_path": str(state_path), |
| "max_concurrency": 1, |
| "entries": [ |
| { |
| "rollout_id": "run-a/task-001", |
| "task_id": "task-001", |
| "instance_id": "i-a", |
| "output_dir": str(output_dir), |
| "provider": "openai", |
| "model_id": "gpt-test", |
| "cleanup_only": True, |
| } |
| ], |
| } |
| ec2 = Ec2() |
|
|
| exit_code = asyncio.run( |
| resume_batch.execute_plan( |
| plan, |
| plan_path=plan_path, |
| sessions={"aws": ec2}, |
| ) |
| ) |
| state = json.loads(state_path.read_text(encoding="utf-8")) |
|
|
| assert exit_code == 0 |
| assert ec2.terminated == ["i-a"] |
| assert state["entries"]["run-a/task-001"]["worker_id"] == "i-a" |
| assert state["entries"]["run-a/task-001"]["instance_id"] == "i-a" |
|
|
