tests/test_distributed_generation.py

from __future__ import annotations

import importlib.util
from pathlib import Path

import pytest

from dovla_cil.generation.distributed import (
    DistributedCILConfig,
    plan_generation_jobs,
    require_ray,
    run_distributed_cil_generation,
)
from dovla_cil.tasks.library import built_in_toy_tasks


def test_ray_absent_gives_clear_error() -> None:
    if importlib.util.find_spec("ray") is not None:
        pytest.skip("Ray is installed in this environment")

    with pytest.raises(ImportError) as exc_info:
        require_ray()

    assert "Ray is optional" in str(exc_info.value)
    assert "pip install ray" in str(exc_info.value)


def test_deterministic_group_ids_under_seed() -> None:
    tasks = built_in_toy_tasks()[:2]

    first = plan_generation_jobs(tasks, backend="toy", num_states_per_task=2, seed=7)
    second = plan_generation_jobs(tasks, backend="toy", num_states_per_task=2, seed=7)
    skipped = plan_generation_jobs(
        tasks,
        backend="toy",
        num_states_per_task=2,
        seed=7,
        completed_group_ids={first[0].group_id},
    )

    assert [job.group_id for job in first] == [job.group_id for job in second]
    assert len(first) == 4
    assert len(skipped) == 3
    assert first[0].group_id not in {job.group_id for job in skipped}


def test_tiny_distributed_generation_if_ray_present(tmp_path: Path) -> None:
    ray = pytest.importorskip("ray")
    tasks = built_in_toy_tasks()[:1]
    config = DistributedCILConfig(
        backend="toy",
        output_dir=tmp_path,
        num_workers=1,
        num_states_per_task=1,
        k=2,
        seed=3,
        shard_size=10,
    )

    try:
        summary = run_distributed_cil_generation(config, tasks)
    finally:
        if ray.is_initialized():
            ray.shutdown()

    assert summary.num_groups == 1
    assert summary.num_records == 2
    assert (tmp_path / "manifest.json").exists()
    assert (tmp_path / "distributed_manifest.json").exists()
    assert (tmp_path / "group_index.jsonl").exists()