| |
| """Plan Xperience-10M episode counts for a fine-tuning run. |
| |
| This is a storage and evaluation-design helper. It does not train a model and |
| does not invent results. Use it before downloading many episodes. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import math |
| import shutil |
| from pathlib import Path |
|
|
|
|
| GB = 1024 ** 3 |
|
|
|
|
| def parse_args() -> argparse.Namespace: |
| parser = argparse.ArgumentParser(description="Estimate feasible Xperience-10M fine-tuning sample counts.") |
| parser.add_argument("--storage-root", type=Path, default=Path("."), help="Disk root to inspect.") |
| parser.add_argument("--free-gb", type=float, default=None, help="Override measured free space in GiB.") |
| parser.add_argument("--target-free-after-download-gb", type=float, default=800.0) |
| parser.add_argument("--model-cache-gb", type=float, default=250.0) |
| parser.add_argument("--checkpoint-cache-gb", type=float, default=200.0) |
| parser.add_argument("--log-cache-gb", type=float, default=50.0) |
| parser.add_argument("--minimal-per-episode-gb", type=float, default=2.02) |
| parser.add_argument("--all-training-per-episode-gb", type=float, default=2.40) |
| parser.add_argument("--full-preview-per-episode-gb", type=float, default=5.10) |
| parser.add_argument("--windows-per-episode", type=int, default=1161) |
| parser.add_argument("--test-fraction", type=float, default=0.20) |
| parser.add_argument("--output", type=Path, default=Path("outputs/omni_exploration/finetune_sample_budget.json")) |
| return parser.parse_args() |
|
|
|
|
| def measured_free_gb(storage_root: Path, override: float | None) -> float: |
| if override is not None: |
| return float(override) |
| if not storage_root.exists(): |
| raise FileNotFoundError(f"storage root does not exist: {storage_root}") |
| return shutil.disk_usage(storage_root).free / GB |
|
|
|
|
| def max_episodes_for_budget(available_data_gb: float, per_episode_gb: float) -> int: |
| if available_data_gb <= 0 or per_episode_gb <= 0: |
| return 0 |
| return max(0, int(math.floor(available_data_gb / per_episode_gb))) |
|
|
|
|
| def split_windows(episodes: int, windows_per_episode: int, test_fraction: float) -> dict: |
| if episodes <= 0: |
| return {"train_episodes": 0, "test_episodes": 0, "train_windows": 0, "test_windows": 0} |
| test_episodes = max(1, int(round(episodes * test_fraction))) if episodes > 1 else 1 |
| train_episodes = max(0, episodes - test_episodes) |
| return { |
| "train_episodes": train_episodes, |
| "test_episodes": test_episodes, |
| "train_windows": train_episodes * windows_per_episode, |
| "test_windows": test_episodes * windows_per_episode, |
| } |
|
|
|
|
| def phase_rows(max_all_training: int, windows_per_episode: int, test_fraction: float) -> list[dict]: |
| phase_specs = [ |
| ("smoke", 1, "Verify loaders, alignment, and heads."), |
| ("smoke_plus", 3, "Catch obvious multi-episode path issues."), |
| ("pilot", 16, "First held-out-episode evaluation."), |
| ("recommended_next", 32, "Default next run if download layout is clean."), |
| ("useful_lora_small", 64, "Train sensor adapters plus selected LoRA layers."), |
| ("useful_lora_medium", 128, "More useful LoRA run after pilot is stable."), |
| ("storage_heavy", 256, "Only after checkpoint size and data layout are stable."), |
| ] |
| rows = [] |
| for name, episodes, purpose in phase_specs: |
| split = split_windows(episodes, windows_per_episode, test_fraction) |
| rows.append({ |
| "phase": name, |
| "episodes": episodes, |
| "feasible_under_all_training_budget": episodes <= max_all_training, |
| "approx_windows": episodes * windows_per_episode, |
| **split, |
| "purpose": purpose, |
| }) |
| return rows |
|
|
|
|
| def choose_recommendation(max_all_training: int) -> int: |
| for candidate in (32, 16, 8, 3, 1): |
| if max_all_training >= candidate: |
| return candidate |
| return 0 |
|
|
|
|
| def main() -> int: |
| args = parse_args() |
| free_gb = measured_free_gb(args.storage_root.expanduser(), args.free_gb) |
| reserved_gb = args.target_free_after_download_gb + args.model_cache_gb + args.checkpoint_cache_gb + args.log_cache_gb |
| available_data_gb = max(0.0, free_gb - reserved_gb) |
|
|
| modes = { |
| "minimal_annotation_plus_one_video": args.minimal_per_episode_gb, |
| "all_training_files_no_rrd": args.all_training_per_episode_gb, |
| "full_preview_including_rrd": args.full_preview_per_episode_gb, |
| } |
| mode_summary = { |
| name: { |
| "per_episode_gb": per_episode_gb, |
| "max_episodes": max_episodes_for_budget(available_data_gb, per_episode_gb), |
| } |
| for name, per_episode_gb in modes.items() |
| } |
| max_all_training = mode_summary["all_training_files_no_rrd"]["max_episodes"] |
| recommended = choose_recommendation(max_all_training) |
|
|
| payload = { |
| "assumptions": { |
| "storage_root": str(args.storage_root), |
| "measured_or_overridden_free_gb": round(free_gb, 3), |
| "target_free_after_download_gb": args.target_free_after_download_gb, |
| "reserved_model_cache_gb": args.model_cache_gb, |
| "reserved_checkpoint_cache_gb": args.checkpoint_cache_gb, |
| "reserved_log_cache_gb": args.log_cache_gb, |
| "available_for_episode_data_gb": round(available_data_gb, 3), |
| "windows_per_episode": args.windows_per_episode, |
| "test_fraction": args.test_fraction, |
| "note": "Episode sizes are estimates until build_episode_manifest.py scans the actual downloaded folders.", |
| }, |
| "modes": mode_summary, |
| "recommended_next_episodes": recommended, |
| "recommended_next_reason": ( |
| "Use 32 episodes first when feasible; otherwise use the largest smaller phase. " |
| "Scale to 64 or 128 only after the pilot download and held-out-episode evaluation are stable." |
| ), |
| "phases": phase_rows(max_all_training, args.windows_per_episode, args.test_fraction), |
| } |
|
|
| args.output.parent.mkdir(parents=True, exist_ok=True) |
| args.output.write_text(json.dumps(payload, indent=2), encoding="utf-8") |
| print(json.dumps(payload["assumptions"], indent=2)) |
| print(json.dumps({"recommended_next_episodes": recommended, "modes": mode_summary}, indent=2)) |
| print(f"Wrote {args.output}") |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
