from __future__ import annotations from collections import Counter from collections.abc import Iterable from dataclasses import asdict, dataclass from pathlib import Path from typing import Any from dovla_cil.data.schema import CILRecord, compute_regret_and_ranks, compute_state_hash from dovla_cil.data.sharding import ShardReader, ShardWriter, group_records from dovla_cil.generation.pipeline import ( GenerationSummary, _execute_group, _make_group_id, _reward_distribution, _stable_seed, load_task_specs, plan_expert_actions, sample_toy_scene, ) from dovla_cil.interventions.samplers import InterventionSampler from dovla_cil.sim.registry import get_simulator_backend from dovla_cil.tasks.schema import TaskSpec from dovla_cil.tasks.validators import validate_task from dovla_cil.utils.io import ensure_dir, write_json from dovla_cil.utils.seeding import seed_everything RAY_INSTALL_HINT = ( "Ray is optional for DoVLA-CIL distributed generation. Install it with " "`pip install ray` or `pip install -e .[ray]` before using distributed generation." ) @dataclass(frozen=True) class DistributedCILConfig: backend: str = "toy" output_dir: str | Path = "data/cil_large" num_workers: int = 4 num_states_per_task: int = 1000 k: int = 32 seed: int = 0 shard_size: int = 10000 resume: bool = False ray_address: str | None = None inline_observations: bool = True def __post_init__(self) -> None: if self.backend != "toy": raise NotImplementedError("Distributed generation currently supports toy backend only.") if self.num_workers <= 0: raise ValueError("num_workers must be positive") if self.num_states_per_task <= 0: raise ValueError("num_states_per_task must be positive") if self.k <= 0: raise ValueError("k must be positive") if self.shard_size <= 0: raise ValueError("shard_size must be positive") if not self.inline_observations: raise NotImplementedError( "Distributed generation currently writes inline observations only." ) @dataclass(frozen=True) class GenerationJob: task_index: int task_payload: dict[str, Any] state_index: int group_seed: int group_id: str state_hash: str @property def task_id(self) -> str: return str(self.task_payload["task_id"]) def to_dict(self) -> dict[str, Any]: return asdict(self) @classmethod def from_dict(cls, payload: dict[str, Any]) -> GenerationJob: return cls( task_index=int(payload["task_index"]), task_payload=dict(payload["task_payload"]), state_index=int(payload["state_index"]), group_seed=int(payload["group_seed"]), group_id=str(payload["group_id"]), state_hash=str(payload["state_hash"]), ) def require_ray(): try: import ray except ImportError as exc: raise ImportError(RAY_INSTALL_HINT) from exc return ray def load_completed_group_ids(output_dir: str | Path) -> set[str]: dataset_dir = Path(output_dir) if not ((dataset_dir / "metadata.json").exists() or (dataset_dir / "manifest.json").exists()): return set() try: return {str(entry["group_id"]) for entry in ShardReader(dataset_dir).index.load_group_index()} except Exception: return set() def load_existing_records(output_dir: str | Path) -> list[CILRecord]: dataset_dir = Path(output_dir) if not ((dataset_dir / "metadata.json").exists() or (dataset_dir / "manifest.json").exists()): return [] try: return list(ShardReader(dataset_dir).iterate_records()) except Exception: return [] def plan_generation_jobs( tasks: list[TaskSpec], *, backend: str, num_states_per_task: int, seed: int, completed_group_ids: Iterable[str] = (), ) -> list[GenerationJob]: if backend != "toy": raise NotImplementedError("Distributed job planning currently supports toy backend only.") completed = set(completed_group_ids) jobs: list[GenerationJob] = [] for task_index, task in enumerate(tasks): validate_task(task) for state_index in range(num_states_per_task): identity = compute_group_identity( task, backend=backend, state_index=state_index, seed=seed, ) if identity["group_id"] in completed: continue jobs.append( GenerationJob( task_index=task_index, task_payload=task.to_dict(), state_index=state_index, group_seed=int(identity["group_seed"]), group_id=str(identity["group_id"]), state_hash=str(identity["state_hash"]), ) ) return jobs def compute_group_identity( task: TaskSpec, *, backend: str, state_index: int, seed: int, ) -> dict[str, Any]: if backend != "toy": raise NotImplementedError("Group identity computation currently supports toy backend only.") group_seed = _stable_seed(seed, task.task_id, state_index) scene = sample_toy_scene(task, state_index=state_index, seed=group_seed) simulator = get_simulator_backend(backend) try: simulator.seed(group_seed) simulator.reset_task(task, scene) state_blob = simulator.serialize_state() state_hash = compute_state_hash(state_blob) group_id = _make_group_id(task.task_id, state_index, state_hash) return {"group_seed": group_seed, "state_hash": state_hash, "group_id": group_id} finally: simulator.close() def generate_group_records_for_job( job_payload: dict[str, Any], *, backend: str, k: int, ) -> dict[str, Any]: job = GenerationJob.from_dict(job_payload) task = TaskSpec.from_dict(job.task_payload) validate_task(task) simulator = get_simulator_backend(backend) try: scene = sample_toy_scene(task, state_index=job.state_index, seed=job.group_seed) simulator.seed(job.group_seed) simulator.reset_task(task, scene) state_blob = simulator.serialize_state() state_hash = compute_state_hash(state_blob) group_id = _make_group_id(task.task_id, job.state_index, state_hash) state_blob_ref = f"states/{group_id}.pkl" observation0 = simulator.render_observation() symbolic0 = simulator.get_symbolic_state() expert_actions = plan_expert_actions( backend=backend, task=task, symbolic_state=symbolic0, ) candidates = InterventionSampler(seed=job.group_seed).sample( task=task, observation=observation0, symbolic_state=symbolic0, expert_actions=expert_actions, k=k, ) records = _execute_group( backend=backend, simulator=simulator, task=task, scene=scene, observation0=observation0, state_blob=state_blob, state_hash=state_hash, state_blob_ref=state_blob_ref, group_id=group_id, group_seed=job.group_seed, state_index=job.state_index, task_index=job.task_index, candidates=candidates, output_dir=Path("."), inline_observations=True, ) ranked = compute_regret_and_ranks(records) return { "status": "generated", "group_id": group_id, "state_hash": state_hash, "state_blob_ref": state_blob_ref, "state_blob": state_blob, "records": [record.to_dict() for record in ranked], } finally: simulator.close() def run_distributed_cil_generation( config: DistributedCILConfig, tasks: list[TaskSpec], ) -> GenerationSummary: ray = require_ray() seed_everything(config.seed) output_dir = ensure_dir(config.output_dir) ensure_dir(output_dir / "states") completed_group_ids = load_completed_group_ids(output_dir) if config.resume else set() existing_records = load_existing_records(output_dir) if config.resume else [] jobs = plan_generation_jobs( tasks, backend=config.backend, num_states_per_task=config.num_states_per_task, seed=config.seed, completed_group_ids=completed_group_ids, ) _write_distributed_manifest( output_dir, config=config, status="running", planned_groups=len(tasks) * config.num_states_per_task, skipped_groups=len(completed_group_ids), completed_groups=len(completed_group_ids), generated_groups=0, ) if not ray.is_initialized(): ray.init(address=config.ray_address, ignore_reinit_error=True, include_dashboard=False) @ray.remote class TaskSceneSamplerActor: def __init__(self, job_rows: list[dict[str, Any]]) -> None: self._jobs = list(job_rows) def jobs(self) -> list[dict[str, Any]]: return list(self._jobs) @ray.remote class SimulatorWorkerActor: def __init__(self, backend: str, k: int) -> None: self.backend = backend self.k = int(k) def generate(self, job_row: dict[str, Any]) -> dict[str, Any]: return generate_group_records_for_job(job_row, backend=self.backend, k=self.k) @ray.remote class ShardWriterActor: def __init__( self, output_dir: str, backend: str, k: int, task_count: int, seed: int, shard_size: int, existing_rows: list[dict[str, Any]], ) -> None: self.output_dir = Path(output_dir) ensure_dir(self.output_dir / "states") self.writer = ShardWriter( self.output_dir, dataset_name=f"cil_{backend}", backend=backend, k=k, task_count=task_count, seed=seed, shard_size=shard_size, overwrite=True, ) self.completed_group_ids: set[str] = set() self.generated_groups = 0 for records in group_records(CILRecord.from_dict(row) for row in existing_rows).values(): for record in records: self.writer.write(record) self.completed_group_ids.add(records[0].group_id) def write_group(self, result: dict[str, Any]) -> dict[str, Any]: if result.get("status") != "generated": return {"status": "skipped", "group_id": result.get("group_id")} group_id = str(result["group_id"]) if group_id in self.completed_group_ids: return {"status": "duplicate_skipped", "group_id": group_id} state_blob_ref = str(result["state_blob_ref"]) (self.output_dir / state_blob_ref).parent.mkdir(parents=True, exist_ok=True) (self.output_dir / state_blob_ref).write_bytes(result["state_blob"]) for row in result["records"]: self.writer.write(CILRecord.from_dict(row)) self.completed_group_ids.add(group_id) self.generated_groups += 1 return {"status": "written", "group_id": group_id, "records": len(result["records"])} def close(self) -> dict[str, Any]: metadata = self.writer.close() return { "metadata": metadata, "completed_group_ids": sorted(self.completed_group_ids), "generated_groups": self.generated_groups, } sampler = TaskSceneSamplerActor.remote([job.to_dict() for job in jobs]) job_rows = ray.get(sampler.jobs.remote()) writer = ShardWriterActor.remote( str(output_dir), config.backend, config.k, len(tasks), config.seed, config.shard_size, [record.to_dict() for record in existing_records], ) workers = [ SimulatorWorkerActor.remote(config.backend, config.k) for _index in range(config.num_workers) ] pending = [] for index, job_row in enumerate(job_rows): worker = workers[index % len(workers)] pending.append(worker.generate.remote(job_row)) written = [] while pending: ready, pending = ray.wait(pending, num_returns=1) result = ray.get(ready[0]) written.append(ray.get(writer.write_group.remote(result))) writer_result = ray.get(writer.close.remote()) generated_groups = int(writer_result.get("generated_groups", 0)) completed_after = len(writer_result.get("completed_group_ids", [])) _write_distributed_manifest( output_dir, config=config, status="complete", planned_groups=len(tasks) * config.num_states_per_task, skipped_groups=len(completed_group_ids), completed_groups=completed_after, generated_groups=generated_groups, written_results=written, ) records = list(ShardReader(output_dir).iterate_records()) return _summary_from_records(records, output_dir=output_dir) def run_distributed_from_task_file( *, task_path: str | Path, config: DistributedCILConfig, ) -> GenerationSummary: return run_distributed_cil_generation(config, load_task_specs(task_path)) def _summary_from_records(records: list[CILRecord], *, output_dir: Path) -> GenerationSummary: rewards = [record.reward.progress for record in records] successes = [record.reward.terminal_success for record in records] candidate_counts = Counter(record.candidate_type for record in records) return GenerationSummary( output_dir=output_dir, manifest_path=output_dir / "manifest.json", group_index_path=output_dir / "group_index.jsonl", num_groups=len({record.group_id for record in records}), num_records=len(records), success_rate=sum(1 for value in successes if value) / len(successes) if successes else 0.0, reward_distribution=_reward_distribution(rewards), candidate_type_distribution=dict(sorted(candidate_counts.items())), ) def _write_distributed_manifest( output_dir: Path, *, config: DistributedCILConfig, status: str, planned_groups: int, skipped_groups: int, completed_groups: int, generated_groups: int, written_results: list[dict[str, Any]] | None = None, ) -> None: write_json( { "status": status, "backend": config.backend, "num_workers": config.num_workers, "num_states_per_task": config.num_states_per_task, "k": config.k, "seed": config.seed, "shard_size": config.shard_size, "resume": config.resume, "planned_groups": planned_groups, "skipped_groups": skipped_groups, "completed_groups": completed_groups, "generated_groups": generated_groups, "written_results": written_results or [], }, output_dir / "distributed_manifest.json", ) __all__ = [ "DistributedCILConfig", "GenerationJob", "RAY_INSTALL_HINT", "compute_group_identity", "generate_group_records_for_job", "load_completed_group_ids", "plan_generation_jobs", "require_ray", "run_distributed_cil_generation", "run_distributed_from_task_file", ]