File size: 2,032 Bytes
da08b7d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 | 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()
|