scripts/verify_component1_setup.py

"""
This script verifies Component 1 setup in plain English.
It checks Python, key packages, and GPU visibility for PyTorch.
"""

import importlib
import importlib.util
import platform
import sys
from typing import List


def check_imports(packages: List[str]) -> List[str]:
    """
    Tries importing required packages.
    Returns a list of package names that failed to import.
    """
    failed = []
    for package in packages:
        try:
            importlib.import_module(package)
        except Exception:
            failed.append(package)
    return failed


def check_optional_installed(packages: List[str]) -> List[str]:
    """
    Checks whether optional packages exist without importing them.
    Returns packages that are missing.
    """
    missing = []
    for package in packages:
        if importlib.util.find_spec(package) is None:
            missing.append(package)
    return missing


def main() -> None:
    print("=== Component 1 Verification ===")
    print(f"Python version: {sys.version.split()[0]}")
    print(f"Operating system: {platform.system()} {platform.release()}")

    # Python 3.10/3.11 is the target for best compatibility on Windows.
    if sys.version_info.major != 3 or sys.version_info.minor not in (10, 11):
        print("")
        print("Verification failed.")
        print("This project currently requires Python 3.10 or 3.11 on Windows.")
        print("Fix suggestion: install Python 3.11, recreate .venv, and reinstall requirements.")
        raise SystemExit(1)

    # These are required for Component 1 success.
    required = [
        "torch",
        "transformers",
        "tokenizers",
        "datasets",
        "accelerate",
        "gradio",
        "tree_sitter",
    ]

    failed = check_imports(required)
    if failed:
        print("")
        print("Verification failed.")
        print("The following packages could not be imported:")
        for package in failed:
            print(f"- {package}")
        print("")
        print("Fix suggestion: activate .venv and run 'pip install -r requirements.txt' again.")
        raise SystemExit(1)

    # Optional imports should not fail Component 1.
    optional = ["bitsandbytes"]
    optional_failed = check_optional_installed(optional)

    import torch

    print(f"PyTorch version: {torch.__version__}")
    cuda_available = torch.cuda.is_available()
    print(f"CUDA available: {cuda_available}")

    if cuda_available:
        gpu_name = torch.cuda.get_device_name(0)
        total_vram_gb = torch.cuda.get_device_properties(0).total_memory / (1024**3)
        print(f"Detected GPU: {gpu_name}")
        print(f"Total VRAM: {total_vram_gb:.2f} GB")
    else:
        print("No CUDA GPU was detected by PyTorch.")
        print("You can still continue, but training speed will be much slower.")

    if optional_failed:
        print("")
        print("Optional package warning:")
        print("- bitsandbytes is not available in this environment.")
        print("This does not block Component 1.")
        print("For Component 5 on native Windows, we will use an automatic fallback optimizer if needed.")

    print("")
    print("Component 1 is verified successfully.")


if __name__ == "__main__":
    try:
        main()
    except Exception as exc:
        print("")
        print("Verification script crashed.")
        print(f"What went wrong: {exc}")
        print("Fix suggestion: ensure .venv is active and dependencies are installed.")
        raise SystemExit(1)