scripts/omni/plan_finetune_sample_budget.py

#!/usr/bin/env python3
"""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())