Spaces:
Sleeping
Sleeping
| from websocket_storage import WebSocketGPUStorage | |
| import numpy as np | |
| from typing import Dict, Any, Optional | |
| import time | |
| class VirtualVRAM: | |
| def __init__(self, size_gb: int = None): | |
| """Initialize virtual VRAM with unlimited storage capability""" | |
| self.storage = WebSocketGPUStorage() | |
| if not self.storage.wait_for_connection(): | |
| raise RuntimeError("Could not connect to GPU storage server") | |
| # Initialize VRAM state with unlimited capacity | |
| self.vram_state = { | |
| "total_size": float('inf'), # Unlimited size | |
| "allocated": 0, | |
| "blocks": {}, | |
| "memory_map": {}, | |
| "is_unlimited": True | |
| } | |
| self.store_vram_state() | |
| def store_vram_state(self): | |
| """Store VRAM state in WebSocket storage""" | |
| self.storage.store_state("vram", "state", self.vram_state) | |
| def allocate_block(self, size: int, block_id: Optional[str] = None) -> str: | |
| """Allocate a block of VRAM""" | |
| if self.vram_state["allocated"] + size > self.vram_state["total_size"]: | |
| raise MemoryError("Not enough VRAM available") | |
| if block_id is None: | |
| block_id = f"block_{time.time_ns()}" | |
| self.vram_state["blocks"][block_id] = { | |
| "size": size, | |
| "allocated_at": time.time_ns(), | |
| "last_accessed": time.time_ns() | |
| } | |
| self.vram_state["allocated"] += size | |
| # Store updated state | |
| self.store_vram_state() | |
| return block_id | |
| def free_block(self, block_id: str): | |
| """Free a block of VRAM""" | |
| if block_id in self.vram_state["blocks"]: | |
| self.vram_state["allocated"] -= self.vram_state["blocks"][block_id]["size"] | |
| del self.vram_state["blocks"][block_id] | |
| self.store_vram_state() | |
| # Remove block data from storage | |
| self.storage.store_tensor(block_id, None) | |
| def write_block(self, block_id: str, data: np.ndarray): | |
| """Write data to a VRAM block""" | |
| if block_id not in self.vram_state["blocks"]: | |
| raise ValueError(f"Block {block_id} not allocated") | |
| self.vram_state["blocks"][block_id]["last_accessed"] = time.time_ns() | |
| self.store_vram_state() | |
| return self.storage.store_tensor(block_id, data) | |
| def read_block(self, block_id: str) -> Optional[np.ndarray]: | |
| """Read data from a VRAM block""" | |
| if block_id not in self.vram_state["blocks"]: | |
| raise ValueError(f"Block {block_id} not allocated") | |
| self.vram_state["blocks"][block_id]["last_accessed"] = time.time_ns() | |
| self.store_vram_state() | |
| return self.storage.load_tensor(block_id) | |
| def map_address(self, virtual_addr: str, block_id: str): | |
| """Map virtual address to VRAM block""" | |
| self.vram_state["memory_map"][virtual_addr] = block_id | |
| self.store_vram_state() | |
| def get_block_from_address(self, virtual_addr: str) -> Optional[str]: | |
| """Get block ID from virtual address""" | |
| return self.vram_state["memory_map"].get(virtual_addr) | |
| def get_stats(self) -> Dict[str, Any]: | |
| """Get VRAM statistics""" | |
| return { | |
| "total_gb": self.size_gb, | |
| "used_gb": self.vram_state["allocated"] / (1024 * 1024 * 1024), | |
| "free_gb": (self.vram_state["total_size"] - self.vram_state["allocated"]) / (1024 * 1024 * 1024), | |
| "num_blocks": len(self.vram_state["blocks"]), | |
| "mappings": len(self.vram_state["memory_map"]) | |
| } | |