memory_optimization/test.py

"""
Memory Optimization Module Test
Quick test to verify all components work correctly.
"""
import torch
import sys
import os

# Add parent directory to path
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from memory_optimization import (
    UnifiedMemoryManager,
    MemoryOptimizationConfig,
    get_shared_model,
    get_tensor,
    return_tensor,
    get_memory_stats,
    clear_memory
)

def test_memory_optimization():
    """Test memory optimization module"""
    print("=" * 70)
    print("Testing Memory Optimization Module")
    print("=" * 70)
    
    # Initialize config
    config = MemoryOptimizationConfig(
        use_shared_model=True,
        device="cuda" if torch.cuda.is_available() else "cpu"
    )
    
    print(f"\n[CONFIG] Device: {config.device}")
    print(f"[CONFIG] Shared Model: {config.use_shared_model}")
    
    # Initialize manager
    manager = UnifiedMemoryManager(config)
    print("\n[OK] UnifiedMemoryManager initialized")
    
    # Test tensor pooling
    print("\n[TEST] Tensor Pooling...")
    tensor1 = manager.get_tensor((10, 1024), dtype=torch.float32)
    print(f"  [OK] Created tensor: {tensor1.shape}, device: {tensor1.device}")
    
    manager.return_tensor(tensor1)
    print("  [OK] Returned tensor to pool")
    
    tensor2 = manager.get_tensor((10, 1024), dtype=torch.float32)
    print(f"  [OK] Retrieved tensor from pool: {tensor2.shape}")
    
    # Test shared model (if available)
    print("\n[TEST] Shared Model...")
    try:
        # This will use shared Qwen model if available
        model = manager.get_shared_model("Qwen/Qwen3-0.6B", "transformer")
        print(f"  [OK] Got shared model: {type(model).__name__}")
    except Exception as e:
        print(f"  [WARN] Could not get shared model: {e}")
    
    # Test memory stats
    print("\n[TEST] Memory Stats...")
    stats = manager.get_memory_stats()
    print(f"  [OK] Got memory stats: {len(stats)} categories")
    
    # Test cleanup
    print("\n[TEST] Memory Cleanup...")
    manager.clear_all_memory()
    print("  [OK] Memory cleared")
    
    print("\n" + "=" * 70)
    print("[SUCCESS] All tests passed!")
    print("=" * 70)

if __name__ == "__main__":
    test_memory_optimization()