File size: 1,970 Bytes
68442cd | 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 | #!/usr/bin/env python
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 ( # noqa: E402
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())
|