| |
| """Freeze the exact FrameSkip cache coverage required by the 20K->20.5K run.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| from pathlib import Path |
|
|
| from starVLA.dataloader.infoflow_sampler import LogicalDraw, ValidAnchorManifest |
|
|
|
|
| def main() -> None: |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--manifest", required=True) |
| parser.add_argument("--output", required=True) |
| parser.add_argument("--seed", type=int, default=42) |
| parser.add_argument("--start-step", type=int, default=20_000) |
| parser.add_argument("--updates", type=int, default=500) |
| parser.add_argument("--global-batch-size", type=int, default=64) |
| parser.add_argument("--global-microbatch-size", type=int, default=8) |
| parser.add_argument("--pruned-batches", type=int, default=5) |
| parser.add_argument("--full-batches", type=int, default=1) |
| args = parser.parse_args() |
|
|
| manifest = ValidAnchorManifest.load(args.manifest) |
| start_draw = args.start_step * args.global_batch_size |
| end_draw = (args.start_step + args.updates) * args.global_batch_size |
| period = args.pruned_batches + args.full_batches |
| required: dict[str, set[int | str]] = { |
| dataset.name: set() for dataset in manifest.datasets |
| } |
| pruned_draws = 0 |
| full_draws = 0 |
| for global_draw_id in range(start_draw, end_draw): |
| microbatch = (global_draw_id - start_draw) // args.global_microbatch_size |
| is_pruned = microbatch % period < args.pruned_batches |
| trace = manifest.sample(args.seed, LogicalDraw(0, global_draw_id)) |
| if is_pruned: |
| required[trace.dataset_name].add(trace.trajectory_id) |
| pruned_draws += 1 |
| else: |
| full_draws += 1 |
|
|
| payload = { |
| "schema_version": 1, |
| "method": "FrameSkip-adapted (Gaussian-TPI fallback)", |
| "manifest_path": str(Path(args.manifest).resolve()), |
| "sampling_seed": args.seed, |
| "start_step": args.start_step, |
| "updates": args.updates, |
| "global_batch_size": args.global_batch_size, |
| "global_microbatch_size": args.global_microbatch_size, |
| "start_global_draw_id": start_draw, |
| "end_global_draw_id_exclusive": end_draw, |
| "pruned_batches": args.pruned_batches, |
| "full_batches": args.full_batches, |
| "pruned_draws": pruned_draws, |
| "full_draws": full_draws, |
| "datasets": { |
| name: sorted(ids, key=lambda value: (type(value).__name__, str(value))) |
| for name, ids in required.items() |
| }, |
| } |
| payload["required_trajectory_count"] = sum( |
| len(values) for values in payload["datasets"].values() |
| ) |
| destination = Path(args.output) |
| destination.parent.mkdir(parents=True, exist_ok=True) |
| destination.write_text(json.dumps(payload, indent=2) + "\n", encoding="utf-8") |
| print(json.dumps({ |
| "output": str(destination), |
| "pruned_draws": pruned_draws, |
| "full_draws": full_draws, |
| "required_trajectory_count": payload["required_trajectory_count"], |
| "per_dataset": {k: len(v) for k, v in payload["datasets"].items()}, |
| }, indent=2)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
