tests/unit/test_exec_executor.py - master - egorbu/harbor_fix

tests/unit/test_exec_executor.py
242 строки8 KB
Alex Shaw
Add compileable task evaluation and auto-verification support (#2150)
02 июл 2026, 10:00
02 июл 2026, 10:0018f525b
from pathlib import Path

import pytest

import harbor.exec.executor as executor_module
from harbor.exec import Executor
from harbor.models.compile import CompileConfig, CompileInstruction
from harbor.models.exec import (
    ExecConfig,
    ExecJobConfig,
    ExecMapConfig,
    ExecReduceConfig,
    ExecReduceEnvironment,
    ExecReduceTaskConfig,
)
from harbor.models.job.config import JobConfig
from harbor.models.job.result import JobResult
from harbor.models.task.paths import TaskPaths
from harbor.models.trial.config import TaskConfig
from harbor.models.trial.result import TrialResult


class FakeJob:
    def __init__(self, config: JobConfig, result: JobResult):
        self.job_dir = config.jobs_dir / config.job_name
        self.result = result

    async def run(self) -> JobResult:
        return self.result


def _patch_job_create(
    monkeypatch: pytest.MonkeyPatch,
    results: list[JobResult],
) -> list[JobConfig]:
    captured_configs: list[JobConfig] = []

    async def fake_create(config: JobConfig) -> FakeJob:
        captured_configs.append(config)
        return FakeJob(config, results.pop(0))

    monkeypatch.setattr(executor_module.Job, "create", staticmethod(fake_create))
    return captured_configs


def _trial_result(trial_name: str, trial_dir: Path) -> TrialResult:
    return TrialResult.model_construct(
        trial_name=trial_name,
        trial_uri=trial_dir.resolve().as_uri(),
    )


def _executor() -> Executor:
    return Executor(
        ExecConfig(
            map=ExecMapConfig(
                compile=CompileConfig(
                    instructions=[CompileInstruction(text="Do it.")],
                ),
                job=ExecJobConfig(),
            )
        )
    )


@pytest.mark.asyncio
async def test_executor_compiles_map_tasks_and_runs_map_job(
    tmp_path: Path,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    map_result = JobResult.model_construct(trial_results=[])
    captured_configs = _patch_job_create(monkeypatch, [map_result])

    config = ExecConfig(
        map=ExecMapConfig(
            compile=CompileConfig(
                output_dir=tmp_path / "map-tasks",
                instructions=[CompileInstruction(text="Write the answer.")],
            ),
            job=ExecJobConfig(
                job_name="map-job",
                jobs_dir=tmp_path / "jobs",
                quiet=True,
            ),
        )
    )

    result = await Executor(config).execute()

    assert result.map.job_result is map_result
    assert result.map.job_dir == tmp_path / "jobs" / "map-job"
    assert result.reduce is None
    assert len(result.map.task_dirs) == 1
    assert TaskPaths(result.map.task_dirs[0]).instruction_path.read_text() == (
        "Write the answer.\n"
    )
    assert captured_configs == [
        JobConfig(
            job_name="map-job",
            jobs_dir=tmp_path / "jobs",
            quiet=True,
            tasks=[TaskConfig(path=result.map.task_dirs[0])],
        )
    ]


@pytest.mark.asyncio
async def test_executor_stages_map_artifacts_into_reduce_task_environment(
    tmp_path: Path,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    reduce_template_dir = tmp_path / "reduce-template"
    reduce_template_environment_dir = reduce_template_dir / "environment"
    reduce_template_environment_dir.mkdir(parents=True)
    (reduce_template_environment_dir / "keep.txt").write_text("template environment")
    reduce_template_artifacts_dir = reduce_template_environment_dir / "artifacts"
    reduce_template_artifacts_dir.mkdir()
    (reduce_template_artifacts_dir / "keep.txt").write_text("template artifacts")

    trial_dir = tmp_path / "map-job" / "map-trial"
    artifacts_dir = trial_dir / "artifacts"
    artifacts_dir.mkdir(parents=True)
    (artifacts_dir / "result.txt").write_text("map output")

    map_result = JobResult.model_construct(
        trial_results=[
            TrialResult.model_construct(
                trial_name="Map Trial",
                trial_uri=trial_dir.resolve().as_uri(),
            )
        ]
    )
    reduce_result = JobResult.model_construct(trial_results=[])
    captured_configs = _patch_job_create(monkeypatch, [map_result, reduce_result])

    config = ExecConfig(
        map=ExecMapConfig(
            compile=CompileConfig(
                output_dir=tmp_path / "map-tasks",
                artifacts=["/app/result.txt"],
                instructions=[CompileInstruction(text="Write the answer.")],
            ),
            job=ExecJobConfig(
                job_name="map-job",
                jobs_dir=tmp_path / "map-jobs",
            ),
        ),
        reduce=ExecReduceConfig(
            task=ExecReduceTaskConfig(
                task_name="reduce",
                output_dir=tmp_path / "reduce-tasks",
                task_template=reduce_template_dir,
                instruction=CompileInstruction(text="Summarize the artifacts."),
                artifacts=["/app/summary.txt"],
                environment=ExecReduceEnvironment(
                    docker_image="python:3.12",
                    workdir="/app",
                ),
            ),
            job=ExecJobConfig(
                job_name="reduce-job",
                jobs_dir=tmp_path / "reduce-jobs",
            ),
        ),
    )

    result = await Executor(config).execute()

    assert result.reduce is not None
    assert result.reduce.job_result is reduce_result
    assert result.reduce.job_dir == tmp_path / "reduce-jobs" / "reduce-job"
    assert len(result.reduce.task_dirs) == 1
    reduce_task_dir = result.reduce.task_dirs[0]
    reduce_paths = TaskPaths(reduce_task_dir)
    assert (
        reduce_paths.environment_dir / "keep.txt"
    ).read_text() == "template environment"
    assert (
        reduce_paths.environment_dir
        / executor_module.REDUCE_ARTIFACTS_DIRNAME
        / "keep.txt"
    ).read_text() == "template artifacts"
    assert (
        reduce_paths.environment_dir
        / executor_module.REDUCE_ARTIFACTS_DIRNAME
        / "0001-map-trial"
        / "result.txt"
    ).read_text() == "map output"

    reduce_task_config = TaskConfig(path=reduce_task_dir)
    assert captured_configs[1] == JobConfig(
        job_name="reduce-job",
        jobs_dir=tmp_path / "reduce-jobs",
        tasks=[reduce_task_config],
    )


def test_executor_skips_map_trials_missing_artifacts(tmp_path: Path) -> None:
    present_trial_dir = tmp_path / "map-job" / "present-trial"
    present_artifacts_dir = present_trial_dir / "artifacts"
    present_artifacts_dir.mkdir(parents=True)
    (present_artifacts_dir / "result.txt").write_text("map output")

    missing_trial_dir = tmp_path / "map-job" / "missing-trial"
    missing_trial_dir.mkdir()

    destination_dir = tmp_path / "staged-artifacts"
    destination_dir.mkdir()

    _executor()._stage_map_artifacts(
        [
            _trial_result("Present Trial", present_trial_dir),
            _trial_result("Missing Trial", missing_trial_dir),
        ],
        destination_dir,
    )

    assert (destination_dir / "0001-present-trial" / "result.txt").read_text() == (
        "map output"
    )
    assert [path.name for path in destination_dir.iterdir()] == ["0001-present-trial"]


def test_executor_raises_when_all_map_trial_artifacts_are_missing(
    tmp_path: Path,
) -> None:
    first_trial_dir = tmp_path / "map-job" / "first-trial"
    second_trial_dir = tmp_path / "map-job" / "second-trial"
    first_trial_dir.mkdir(parents=True)
    second_trial_dir.mkdir()

    destination_dir = tmp_path / "staged-artifacts"
    destination_dir.mkdir()

    with pytest.raises(ValueError, match="no trial artifacts"):
        _executor()._stage_map_artifacts(
            [
                _trial_result("First Trial", first_trial_dir),
                _trial_result("Second Trial", second_trial_dir),
            ],
            destination_dir,
        )