verify_setup.py

#!/usr/bin/env python3
"""Verify that the training environment is correctly set up.

Run this BEFORE training to catch missing packages, GPU issues,
and version mismatches early.

Usage:
    python verify_setup.py
"""

import sys
import importlib


# ── Required packages and minimum versions ──────────────────────────────────

REQUIRED_PACKAGES = [
    # (import_name, pip_name, min_version_tuple, required_for)
    ("torch",           "torch",           (2, 0, 0),   "training"),
    ("transformers",    "transformers",    (4, 40, 0),  "training"),
    ("trl",             "trl",             (0, 18, 0),  "training (hackathon guideline §10)"),
    ("peft",            "peft",            (0, 9, 0),   "training"),
    ("accelerate",      "accelerate",      (0, 27, 0),  "training"),
    ("datasets",        "datasets",        (3, 0, 0),   "TRL GRPOTrainer"),
    ("bitsandbytes",    "bitsandbytes",    (0, 43, 0),  "4-bit quantization"),
    ("numpy",           "numpy",           (1, 24, 0),  "training"),
    ("matplotlib",      "matplotlib",      (3, 5, 0),   "reward curve plotting"),
]

OPTIONAL_PACKAGES = [
    ("unsloth",         "unsloth",         None,        "2x speed boost (hackathon guideline §10)"),
]


def parse_version(version_str: str) -> tuple:
    """Parse a version string like '2.11.0' into a tuple of ints."""
    parts = []
    for part in version_str.split(".")[:3]:
        # Handle versions like '0.45.0.dev0'
        num = ""
        for ch in part:
            if ch.isdigit():
                num += ch
            else:
                break
        parts.append(int(num) if num else 0)
    while len(parts) < 3:
        parts.append(0)
    return tuple(parts)


def check_package(import_name, pip_name, min_version, required_for, optional=False):
    """Check if a package is installed and meets minimum version."""
    prefix = "  💡" if optional else "  ❌"
    ok_prefix = "  ✅"

    try:
        mod = importlib.import_module(import_name)
    except ImportError:
        if optional:
            print(f"{prefix} {pip_name}: NOT INSTALLED (optional — {required_for})")
            return True  # optional, so not a failure
        print(f"{prefix} {pip_name}: NOT INSTALLED — pip install {pip_name}")
        print(f"      Required for: {required_for}")
        return False

    version_str = getattr(mod, "__version__", "unknown")
    if min_version and version_str != "unknown":
        installed = parse_version(version_str)
        if installed < min_version:
            min_str = ".".join(str(x) for x in min_version)
            print(f"  ⚠️  {pip_name}: {version_str} (minimum: {min_str})")
            print(f"      Upgrade: pip install -U {pip_name}")
            return False

    print(f"{ok_prefix} {pip_name}: {version_str}")
    return True


def check_gpu():
    """Check GPU availability and CUDA setup."""
    print("\n─── GPU / CUDA ───────────────────────────────────")
    try:
        import torch
    except ImportError:
        print("  ❌ Cannot check GPU — torch not installed")
        return False

    if not torch.cuda.is_available():
        print("  ⚠️  CUDA: NOT AVAILABLE")
        print("      Training will run on CPU (very slow).")
        print("      For GPU training, install CUDA-enabled PyTorch:")
        print("      pip install torch --index-url https://download.pytorch.org/whl/cu126")
        return True  # not a failure, just a warning

    device_name = torch.cuda.get_device_name(0)
    device_count = torch.cuda.device_count()
    vram_gb = torch.cuda.get_device_properties(0).total_mem / (1024**3)
    cuda_version = torch.version.cuda or "unknown"

    print(f"  ✅ CUDA: {cuda_version}")
    print(f"  ✅ GPU: {device_name} ({device_count} device(s))")
    print(f"  ✅ VRAM: {vram_gb:.1f} GB")

    if vram_gb < 6:
        print("  ⚠️  Low VRAM — 4-bit quantization required, may still OOM")
    elif vram_gb < 16:
        print("  💡 Moderate VRAM — 4-bit quantization recommended")
    else:
        print("  💡 Ample VRAM — can use FP16 or 4-bit")

    return True


def check_bitsandbytes_cuda():
    """Verify bitsandbytes CUDA integration."""
    print("\n─── bitsandbytes CUDA ────────────────────────────")
    try:
        import bitsandbytes as bnb
        import torch
    except ImportError:
        print("  ⚠️  Skipping — bitsandbytes or torch not installed")
        return True

    if not torch.cuda.is_available():
        print("  ⚠️  CUDA not available — 4-bit quantization will be skipped")
        return True

    try:
        # Quick smoke test: create a quantized linear layer
        layer = bnb.nn.Linear4bit(64, 64, bias=False)
        print("  ✅ bitsandbytes 4-bit: working")
    except Exception as e:
        print(f"  ❌ bitsandbytes 4-bit test failed: {e}")
        return False

    return True


def check_environment():
    """Check that the crime_env package is importable."""
    print("\n─── Project Environment ──────────────────────────")
    try:
        from crime_env.environment import CrimeInvestigationEnv
        from crime_env.case_generator import generate_case
        from crime_env.reward_calculator import RewardCalculator
        from crime_env.consistency_tracker import ConsistencyTracker
        from crime_env.agent_prompts import build_system_prompt
        from crime_env.constants import SUSPECT_A, SUSPECT_B

        # Quick smoke test
        case = generate_case()
        env = CrimeInvestigationEnv()
        obs = env.reset(case_data=case)
        assert obs["role"] == "detective"
        print("  ✅ crime_env: all modules importable")
        print("  ✅ Environment: reset() works correctly")
    except ImportError as e:
        print(f"  ❌ crime_env import failed: {e}")
        print("      Make sure you're running from the project root directory")
        return False
    except Exception as e:
        print(f"  ❌ Environment smoke test failed: {e}")
        return False

    return True


def check_hackathon_stack():
    """Check the hackathon-required stack: TRL + Unsloth + OpenEnv."""
    print("\n─── Hackathon Stack (§10 Guidelines) ─────────────")
    all_ok = True

    # TRL
    try:
        from trl import GRPOTrainer, GRPOConfig
        print("  ✅ TRL GRPOTrainer: available")
    except ImportError:
        print("  ❌ TRL GRPOTrainer: NOT AVAILABLE")
        print("      pip install trl")
        all_ok = False

    # Unsloth
    try:
        from unsloth import FastLanguageModel
        print("  ✅ Unsloth FastLanguageModel: available")
        try:
            from unsloth import PatchFastRL
            print("  ✅ Unsloth PatchFastRL: available")
        except ImportError:
            print("  ⚠️  Unsloth PatchFastRL: not found (may need newer version)")
    except ImportError:
        print("  ⚠️  Unsloth: NOT INSTALLED (optional but recommended)")
        print("      pip install unsloth")

    # OpenEnv
    try:
        from openenv.core.env_server.http_server import create_app
        print("  ✅ OpenEnv: available")
    except ImportError:
        print("  ⚠️  OpenEnv: NOT INSTALLED (needed for server only)")
        print("      pip install openenv-core")

    return all_ok


def main():
    print("=" * 55)
    print("  AI Crime Investigation World — Setup Verification")
    print("=" * 55)
    print(f"\n  Python: {sys.version}")
    print(f"  Platform: {sys.platform}")

    all_ok = True

    # Required packages
    print("\n─── Required Packages ────────────────────────────")
    for import_name, pip_name, min_ver, desc in REQUIRED_PACKAGES:
        if not check_package(import_name, pip_name, min_ver, desc):
            all_ok = False

    # Optional packages
    print("\n─── Optional Packages ────────────────────────────")
    for import_name, pip_name, min_ver, desc in OPTIONAL_PACKAGES:
        check_package(import_name, pip_name, min_ver, desc, optional=True)

    # GPU
    check_gpu()

    # bitsandbytes CUDA
    check_bitsandbytes_cuda()

    # Project environment
    if not check_environment():
        all_ok = False

    # Hackathon stack
    if not check_hackathon_stack():
        all_ok = False

    # Summary
    print("\n" + "=" * 55)
    if all_ok:
        print("  ✅ ALL CHECKS PASSED — Ready to train!")
    else:
        print("  ❌ SOME CHECKS FAILED — Fix the issues above")
    print("=" * 55)

    return 0 if all_ok else 1


if __name__ == "__main__":
    sys.exit(main())