dovla_cil/generation/distributed.py

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",
]