| |
| from __future__ import annotations |
|
|
| import argparse |
| import sys |
| from pathlib import Path |
|
|
| PROJECT_ROOT = Path(__file__).resolve().parents[1] |
| if str(PROJECT_ROOT) not in sys.path: |
| sys.path.insert(0, str(PROJECT_ROOT)) |
|
|
| from dovla_cil.generation.distributed import ( |
| DistributedCILConfig, |
| run_distributed_from_task_file, |
| ) |
| from dovla_cil.generation.pipeline import print_generation_summary |
|
|
|
|
| def main(argv: list[str] | None = None) -> int: |
| parser = argparse.ArgumentParser(description="Generate CIL data with optional Ray workers.") |
| parser.add_argument("--backend", choices=["toy"], default="toy") |
| parser.add_argument("--tasks", type=Path, required=True, help="TaskSpec JSONL or JSON path.") |
| parser.add_argument("--out", type=Path, required=True) |
| parser.add_argument("--num-workers", type=int, default=4) |
| parser.add_argument("--num-states-per-task", type=int, default=1000) |
| parser.add_argument("--k", type=int, default=32) |
| parser.add_argument("--seed", type=int, default=0) |
| parser.add_argument("--shard-size", type=int, default=10000) |
| parser.add_argument("--resume", action="store_true") |
| parser.add_argument("--ray-address", default=None) |
| args = parser.parse_args(argv) |
|
|
| config = DistributedCILConfig( |
| backend=args.backend, |
| output_dir=args.out, |
| num_workers=args.num_workers, |
| num_states_per_task=args.num_states_per_task, |
| k=args.k, |
| seed=args.seed, |
| shard_size=args.shard_size, |
| resume=args.resume, |
| ray_address=args.ray_address, |
| ) |
| try: |
| summary = run_distributed_from_task_file(task_path=args.tasks, config=config) |
| except ImportError as exc: |
| print(str(exc), file=sys.stderr) |
| return 2 |
| print_generation_summary(summary) |
| print(f"distributed manifest: {args.out / 'distributed_manifest.json'}") |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
