#!/usr/bin/env python3
"""
Test model loading and GPU distribution without training.
"""

import torch
from transformers import AutoModelForCausalLM, BitsAndBytesConfig

def test_model_loading():
    print("=" * 80)
    print("Testing model loading and GPU distribution")
    print("=" * 80)
    
    # Check GPU availability
    print(f"\n1. GPU Check:")
    print(f"   CUDA available: {torch.cuda.is_available()}")
    print(f"   GPU count: {torch.cuda.device_count()}")
    for i in range(torch.cuda.device_count()):
        props = torch.cuda.get_device_properties(i)
        print(f"   GPU {i}: {props.name} ({props.total_memory / 1e9:.2f} GB)")
    
    # Test 1: Load with device_map="auto" (distributed)
    print("\n2. Test 1: Load with device_map='auto' (should distribute across GPUs)")
    try:
        bnb_config = BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_quant_type="nf4",
            bnb_4bit_compute_dtype=torch.bfloat16,
            bnb_4bit_use_double_quant=True,
        )
        
        print("   Loading model...")
        model = AutoModelForCausalLM.from_pretrained(
            "/data/models/Ornith-1.0-35B",
            quantization_config=bnb_config,
            device_map="auto",
            trust_remote_code=True,
            low_cpu_mem_usage=True,
        )
        print("   ✓ Model loaded successfully")
        
        # Check memory usage
        print("\n3. Memory Usage:")
        for i in range(torch.cuda.device_count()):
            mem_allocated = torch.cuda.memory_allocated(i) / 1e9
            mem_reserved = torch.cuda.memory_reserved(i) / 1e9
            total = torch.cuda.get_device_properties(i).total_memory / 1e9
            print(f"   GPU {i}:")
            print(f"      Allocated: {mem_allocated:.2f} GB")
            print(f"      Reserved: {mem_reserved:.2f} GB")
            print(f"      Total: {total:.2f} GB")
            print(f"      Free: {total - mem_allocated:.2f} GB")
        
        # Check if model is distributed
        print("\n4. Distribution Check:")
        
        # Get memory on each GPU
        gpu0_mem = torch.cuda.memory_allocated(0) / 1e9
        gpu1_mem = torch.cuda.memory_allocated(1) / 1e9
        
        print(f"   GPU 0 memory: {gpu0_mem:.2f} GB")
        print(f"   GPU 1 memory: {gpu1_mem:.2f} GB")
        
        # Determine distribution pattern
        print("\n5. Distribution Pattern:")
        if abs(gpu0_mem - gpu1_mem) < 1.0:  # Within 1GB
            if gpu0_mem < 10.0:  # Less than 10GB each
                print("   ✓ DISTRIBUTED: Model split across both GPUs")
                print(f"      Each GPU has ~{gpu0_mem:.2f}GB of the model")
            else:
                print("   ⚠ DUPLICATE: Same model loaded on both GPUs")
                print(f"      Each GPU has ~{gpu0_mem:.2f}GB (wasteful but fits)")
        else:
            print("   ✗ NOT DISTRIBUTED: Model on one GPU only")
            if gpu0_mem > gpu1_mem:
                print(f"      GPU 0: {gpu0_mem:.2f}GB, GPU 1: {gpu1_mem:.2f}GB")
            else:
                print(f"      GPU 0: {gpu0_mem:.2f}GB, GPU 1: {gpu1_mem:.2f}GB")
        
        print("\n" + "=" * 80)
        print("TEST PASSED: Model loaded and distributed across GPUs")
        print("=" * 80)
        
    except Exception as e:
        print(f"\n✗ Test 1 FAILED: {e}")
        print("\nThis means the model is NOT being distributed properly!")
        print("It might be trying to fit the entire model on one GPU.")
        return False
    
    return True

if __name__ == "__main__":
    success = test_model_loading()
    exit(0 if success else 1)