virtual_gpu_setup/virtual_gpu/virtual_ram.py

"""
Virtual RAM Module - 128GB System Memory Abstraction

This module implements a symbolic representation of 128GB system RAM using
efficient data structures and lazy allocation strategies. It avoids allocating
real memory and uses dictionaries or sparse mappings to simulate blocks.
"""

import time
from typing import Dict, Any, Optional, Union
from dataclasses import dataclass
import numpy as np


@dataclass
class RAMBlock:
    """Represents a block of memory in the symbolic RAM."""
    name: str
    size_bytes: int
    allocated_time: float
    last_accessed: float
    access_count: int = 0
    # We use a symbolic representation instead of actual data
    # The data field will be None for large blocks to avoid memory allocation
    data: Optional[Union[np.ndarray, bytes]] = None
    is_symbolic: bool = True  # True if this is a symbolic block (no real data)


class VirtualRAM:
    """
    Virtual RAM class that simulates 128GB of system memory symbolically.
    
    This class provides block allocation, tracking, and transfer capabilities
    without actually allocating large amounts of physical memory.
    """
    
    def __init__(self, capacity_gb: int = 128):
        self.capacity_bytes = capacity_gb * 1024 * 1024 * 1024  # Convert GB to bytes
        self.capacity_gb = capacity_gb
        
        # Block registry - stores metadata about allocated blocks
        self.blocks: Dict[str, RAMBlock] = {}
        
        # Memory usage tracking
        self.allocated_bytes = 0
        self.allocation_counter = 0
        
        # Access simulation parameters
        self.access_delay_ms = 0.1  # Simulated RAM access delay
        self.transfer_bandwidth_gbps = 51.2  # DDR5-6400 bandwidth
        
        # Statistics
        self.total_allocations = 0
        self.total_deallocations = 0
        self.total_accesses = 0
        self.total_transfers = 0
        
        print(f"VirtualRAM initialized with {capacity_gb}GB capacity")
        
    def allocate_block(self, name: str, size_bytes: int, 
                      store_data: bool = False) -> bool:
        """
        Allocate a block of memory symbolically.
        
        Args:
            name: Unique name for the block
            size_bytes: Size of the block in bytes
            store_data: If True, actually allocate small amounts of real data for testing
                       If False (default), only store metadata symbolically
        
        Returns:
            True if allocation successful, False if not enough space or name exists
        """
        # Check if name already exists
        if name in self.blocks:
            print(f"Block '{name}' already exists")
            return False
            
        # Check if we have enough capacity
        if self.allocated_bytes + size_bytes > self.capacity_bytes:
            print(f"Not enough capacity: requested {size_bytes:,} bytes, "
                  f"available {self.capacity_bytes - self.allocated_bytes:,} bytes")
            return False
            
        # Create the block
        current_time = time.time()
              # For all blocks, we store only metadata to avoid memory issues
        actual_data = None
        is_symbolic = True
        
        # If store_data is explicitly requested and size is very small, we can store actual data
        if store_data and size_bytes <= 1024 * 1024 * 10: # Up to 10MB for actual data
            actual_data = np.zeros(size_bytes, dtype=np.uint8)
            is_symbolic = False
            print(f"Allocated real data for block \'{name}\' ({size_bytes:,} bytes)")
        else:
            print(f"Created symbolic block \'{name}\' of {size_bytes:,} bytes")            
        block = RAMBlock(
            name=name,
            size_bytes=size_bytes,
            allocated_time=current_time,
            last_accessed=current_time,
            data=actual_data,
            is_symbolic=is_symbolic
        )
        
        self.blocks[name] = block
        self.allocated_bytes += size_bytes
        self.total_allocations += 1
        self.allocation_counter += 1
        
        print(f"Allocated block '{name}': {size_bytes:,} bytes "
              f"({'symbolic' if is_symbolic else 'real data'})")
        
        return True
        
    def get_block(self, name: str) -> Optional[RAMBlock]:
        """
        Retrieve a block by name and simulate access delay.
        
        Args:
            name: Name of the block to retrieve
            
        Returns:
            RAMBlock if found, None otherwise
        """
        if name not in self.blocks:
            return None
            
        # Simulate access delay
        time.sleep(self.access_delay_ms / 1000.0)
        
        # Update access statistics
        block = self.blocks[name]
        block.last_accessed = time.time()
        block.access_count += 1
        self.total_accesses += 1
        
        return block
        
    def release_block(self, name: str) -> bool:
        """
        Deallocate a block of memory.
        
        Args:
            name: Name of the block to deallocate
            
        Returns:
            True if deallocation successful, False if block not found
        """
        if name not in self.blocks:
            print(f"Block '{name}' not found")
            return False
            
        block = self.blocks[name]
        self.allocated_bytes -= block.size_bytes
        self.total_deallocations += 1
        
        del self.blocks[name]
        
        print(f"Released block '{name}': {block.size_bytes:,} bytes")
        return True
        
    def transfer_to_vram(self, block_name: str, vram_instance, 
                        vram_name: Optional[str] = None) -> Optional[str]:
        """
        Transfer a RAM block to VRAM with delay simulation.
        
        Args:
            block_name: Name of the RAM block to transfer
            vram_instance: Instance of VRAM to transfer to
            vram_name: Optional name for the block in VRAM
            
        Returns:
            VRAM block ID if successful, None otherwise
        """
        # Get the block from RAM
        block = self.get_block(block_name)
        if block is None:
            print(f"Block '{block_name}' not found in RAM")
            return None
            
        # Calculate transfer time based on bandwidth
        transfer_time_ms = (block.size_bytes / (self.transfer_bandwidth_gbps * 1e9)) * 1000
        
        print(f"Transferring '{block_name}' ({block.size_bytes:,} bytes) "
              f"from RAM to VRAM (estimated {transfer_time_ms:.2f}ms)")
        
        # Prepare data for transfer
        if block.is_symbolic:
            # For symbolic blocks, create a small representative data sample
            sample_size = min(1024, block.size_bytes)  # 1KB sample
            transfer_data = np.random.randint(0, 256, sample_size, dtype=np.uint8)
            print(f"Using {sample_size} byte sample for symbolic block transfer")
        else:
            # Use actual data
            transfer_data = block.data
            
        # Perform the transfer to VRAM
        if vram_name is None:
            vram_name = f"ram_transfer_{block_name}"
            
        vram_id = vram_instance.transfer_from_ram(vram_name, transfer_data, 
                                                 delay_ms=transfer_time_ms)
        
        if vram_id:
            self.total_transfers += 1
            print(f"Successfully transferred '{block_name}' to VRAM as '{vram_id}'")
        else:
            print(f"Failed to transfer '{block_name}' to VRAM")
            
        return vram_id
        
    def create_tensor_block(self, name: str, shape: tuple, dtype=np.float32) -> bool:
        """
        Create a tensor block with specified shape and data type.
        
        Args:
            name: Name for the tensor block
            shape: Shape of the tensor (e.g., (1024, 1024, 3))
            dtype: Data type of the tensor
            
        Returns:
            True if creation successful, False otherwise
        """
        # Calculate size in bytes
        element_size = np.dtype(dtype).itemsize
        total_elements = np.prod(shape)
        size_bytes = total_elements * element_size
        
        # Allocate the block symbolically
        success = self.allocate_block(name, size_bytes, store_data=False)
        
        if success:
            # Store tensor metadata
            block = self.blocks[name]
            block.tensor_shape = shape
            block.tensor_dtype = dtype
            print(f"Created tensor block '{name}' with shape {shape} and dtype {dtype}")
            
        return success
        
    def info(self) -> Dict[str, Any]:
        """
        Get comprehensive information about the Virtual RAM state.
        
        Returns:
            Dictionary containing RAM usage statistics and metadata
        """
        used_bytes = self.allocated_bytes
        free_bytes = self.capacity_bytes - used_bytes
        utilization_percent = (used_bytes / self.capacity_bytes) * 100
        
        # Calculate average block size
        avg_block_size = used_bytes / len(self.blocks) if self.blocks else 0
        
        # Find largest and smallest blocks
        largest_block = max(self.blocks.values(), key=lambda b: b.size_bytes) if self.blocks else None
        smallest_block = min(self.blocks.values(), key=lambda b: b.size_bytes) if self.blocks else None
        
        # Count symbolic vs real blocks
        symbolic_blocks = sum(1 for b in self.blocks.values() if b.is_symbolic)
        real_blocks = len(self.blocks) - symbolic_blocks
        
        info_dict = {
            "capacity_gb": self.capacity_gb,
            "capacity_bytes": self.capacity_bytes,
            "used_bytes": used_bytes,
            "free_bytes": free_bytes,
            "utilization_percent": utilization_percent,
            "total_blocks": len(self.blocks),
            "symbolic_blocks": symbolic_blocks,
            "real_data_blocks": real_blocks,
            "avg_block_size_bytes": avg_block_size,
            "largest_block_name": largest_block.name if largest_block else None,
            "largest_block_size": largest_block.size_bytes if largest_block else 0,
            "smallest_block_name": smallest_block.name if smallest_block else None,
            "smallest_block_size": smallest_block.size_bytes if smallest_block else 0,
            "total_allocations": self.total_allocations,
            "total_deallocations": self.total_deallocations,
            "total_accesses": self.total_accesses,
            "total_transfers": self.total_transfers,
            "block_names": list(self.blocks.keys())
        }
        
        return info_dict
        
    def print_info(self) -> None:
        """Print a formatted summary of Virtual RAM information."""
        info = self.info()
        
        print("\n" + "="*50)
        print("VIRTUAL RAM INFORMATION")
        print("="*50)
        print(f"Capacity: {info['capacity_gb']} GB ({info['capacity_bytes']:,} bytes)")
        print(f"Used: {info['used_bytes']:,} bytes ({info['utilization_percent']:.2f}%)")
        print(f"Free: {info['free_bytes']:,} bytes")
        print(f"Total Blocks: {info['total_blocks']}")
        print(f"  - Symbolic blocks: {info['symbolic_blocks']}")
        print(f"  - Real data blocks: {info['real_data_blocks']}")
        
        if info['total_blocks'] > 0:
            print(f"Average block size: {info['avg_block_size_bytes']:,.0f} bytes")
            print(f"Largest block: '{info['largest_block_name']}' ({info['largest_block_size']:,} bytes)")
            print(f"Smallest block: '{info['smallest_block_name']}' ({info['smallest_block_size']:,} bytes)")
            
        print(f"\nStatistics:")
        print(f"  - Total allocations: {info['total_allocations']}")
        print(f"  - Total deallocations: {info['total_deallocations']}")
        print(f"  - Total accesses: {info['total_accesses']}")
        print(f"  - Total transfers: {info['total_transfers']}")
        
        if info['block_names']:
            print(f"\nBlock names: {', '.join(info['block_names'])}")
            
        print("="*50)
        
    def simulate_workload(self, num_operations: int = 100) -> None:
        """
        Simulate a typical workload with allocations, accesses, and deallocations.
        
        Args:
            num_operations: Number of operations to simulate
        """
        print(f"\nSimulating workload with {num_operations} operations...")
        
        import random
        
        for i in range(num_operations):
            operation = random.choice(['allocate', 'access', 'deallocate'])
            
            if operation == 'allocate' and len(self.blocks) < 50:  # Limit to 50 blocks
                size = random.randint(1024, 100 * 1024 * 1024)  # 1KB to 100MB
                name = f"workload_block_{i}"
                self.allocate_block(name, size)
                
            elif operation == 'access' and self.blocks:
                block_name = random.choice(list(self.blocks.keys()))
                self.get_block(block_name)
                
            elif operation == 'deallocate' and self.blocks:
                block_name = random.choice(list(self.blocks.keys()))
                self.release_block(block_name)
                
        print(f"Workload simulation completed.")


if __name__ == "__main__":
    # Test the VirtualRAM module
    print("Testing VirtualRAM module...")
    
    # Create a VirtualRAM instance with 128GB capacity
    vram = VirtualRAM(capacity_gb=128)
    
    # Test basic allocation
    print("\n1. Testing basic allocation...")
    vram.allocate_block("small_buffer", 1024 * 1024, store_data=True)  # 1MB with real data
    vram.allocate_block("medium_buffer", 50 * 1024 * 1024)  # 50MB symbolic
    vram.allocate_block("large_tensor", 16 * 1024 * 1024 * 1024)  # 16GB symbolic
    
    # Test tensor creation
    print("\n2. Testing tensor creation...")
    vram.create_tensor_block("ai_weights", (1000, 1000, 512), np.float32)
    vram.create_tensor_block("image_batch", (32, 224, 224, 3), np.uint8)
    
    # Test block access
    print("\n3. Testing block access...")
    block = vram.get_block("small_buffer")
    if block:
        print(f"Accessed block: {block.name}, size: {block.size_bytes:,} bytes")
        
    # Test info display
    print("\n4. Testing info display...")
    vram.print_info()
    
    # Test workload simulation
    print("\n5. Testing workload simulation...")
    vram.simulate_workload(20)
    
    # Final info
    print("\n6. Final state...")
    vram.print_info()
    
    print("\nVirtualRAM test completed!")