scripts/runpod_setup_check.py

#!/usr/bin/env python3
"""RunPod environment check before LoRA training."""

from __future__ import annotations

import json
import platform
import sys
from pathlib import Path

LFS_POINTER_PREFIX = "version https://git-lfs.github.com/spec/v1"


def project_root() -> Path:
    return Path(__file__).resolve().parent.parent


def main() -> int:
    root = project_root()
    train_file = root / "data" / "train" / "mythos_sft_messages_clean.jsonl"
    train_script = root / "scripts" / "train_lora_sft.py"

    print("=== Mythos-Coder RunPod Setup Check ===")
    print(f"Python version:     {platform.python_version()}")
    print(f"Platform:           {platform.platform()}")

    try:
        import torch

        print(f"torch version:      {torch.__version__}")
        cuda = torch.cuda.is_available()
        print(f"CUDA available:     {cuda}")
        if cuda:
            print(f"GPU name:           {torch.cuda.get_device_name(0)}")
            props = torch.cuda.get_device_properties(0)
            total_gb = props.total_memory / (1024**3)
            print(f"GPU memory:         {total_gb:.2f} GB")
        else:
            print("GPU name:           (none)")
            print("GPU memory:         (none)")
    except ImportError:
        print("torch version:      NOT INSTALLED")
        print("CUDA available:     unknown")
        print("GPU name:           unknown")
        print("GPU memory:         unknown")

    try:
        import transformers

        print(f"transformers version: {transformers.__version__}")
    except ImportError:
        print("transformers version: NOT INSTALLED")

    print(f"train file exists:  {train_file.exists()} ({train_file})")
    print(f"train script exists: {train_script.exists()} ({train_script})")

    train_valid = False
    if train_file.exists():
        rows = 0
        with train_file.open(encoding="utf-8") as handle:
            for line_num, line in enumerate(handle, 1):
                stripped = line.strip()
                if not stripped:
                    continue
                rows += 1
                if line_num == 1 and stripped.startswith(LFS_POINTER_PREFIX):
                    print("train file format:  Git LFS pointer (run: git lfs install && git lfs pull)")
                    break
                try:
                    json.loads(stripped)
                except json.JSONDecodeError as exc:
                    print(f"train file format:  invalid JSON on line {line_num}: {exc}")
                    break
            else:
                train_valid = rows > 0
                print(f"train file rows:    {rows}")
                print(f"train file format:  valid JSONL")

    ok = train_file.exists() and train_script.exists() and train_valid
    print("=== Status:", "READY" if ok else "NOT READY", "===")
    return 0 if ok else 1


if __name__ == "__main__":
    raise SystemExit(main())