Spaces:

Lishika
/

crime

Sleeping

crime / verify_setup.py

Added training + fixes + logs

eb9ad07 6 months ago

8.95 kB

#!/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())

	#!/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())