scripts/push_checkpoints.py

"""Push checkpoints to GitHub via Git LFS.

Called automatically at the end of training, or manually.
Requires GH_TOKEN environment variable.
"""

import os
import subprocess
import sys


def run(cmd, check=True):
    """Run a shell command and return output."""
    result = subprocess.run(cmd, shell=True, capture_output=True, text=True)
    if check and result.returncode != 0:
        print(f"Command failed: {cmd}")
        print(f"stderr: {result.stderr}")
        return None
    return result.stdout.strip()


def main():
    repo_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
    os.chdir(repo_dir)

    # Check for new checkpoints
    checkpoints = []
    ckpt_dir = os.path.join(repo_dir, "checkpoints")
    if os.path.exists(ckpt_dir):
        for f in os.listdir(ckpt_dir):
            if f.endswith((".pt", ".pth", ".json", ".bin", ".safetensors")):
                checkpoints.append(os.path.join("checkpoints", f))

    if not checkpoints:
        print("No checkpoints found to push.")
        return

    print(f"Found {len(checkpoints)} checkpoint files:")
    for cp in checkpoints:
        size = os.path.getsize(cp)
        print(f"  {cp}: {size / 1e9:.2f} GB" if size > 1e6 else f"  {cp}: {size / 1e3:.1f} KB")

    # Configure git
    run("git config user.name 'Hardik Sanghvi'")
    run("git config user.email 'hardik@adiance.com'")

    # Stage checkpoint files
    for cp in checkpoints:
        run(f"git add -f {cp}")

    # Also add any updated configs or results
    run("git add -f configs/ || true")
    run("git add -f checkpoints/benchmark_results.json || true")

    # Commit
    status = run("git status --porcelain")
    if not status:
        print("No changes to commit.")
        return

    run('git commit -m "feat: update checkpoints from training run"')

    # Push
    print("Pushing to GitHub (with LFS)...")
    result = run("git push origin main 2>&1", check=False)
    if result:
        print(result)
        print("Push successful!")
    else:
        print("Push failed. You may need to set GH_TOKEN or push manually.")


if __name__ == "__main__":
    main()