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