Upload 12 files

Dockerfile

@@ -0,0 +1,40 @@
+# Use an official Python runtime as a parent image
+FROM python:3.9-slim
+
+# Set environment variables
+ENV PYTHONUNBUFFERED=1
+
+# Set working directory in container
+WORKDIR /app
+
+# Install system dependencies
+RUN apt-get update && apt-get install -y \
+    gcc \
+    python3-dev \
+    && rm -rf /var/lib/apt/lists/*
+
+# Copy only the requirements file first
+COPY virtual_gpu_server/requirements.txt .
+
+# Install Python dependencies
+RUN pip install --no-cache-dir -r requirements.txt
+
+# Copy the server code
+COPY virtual_gpu_server/ .
+
+# Create storage directories
+RUN mkdir -p /app/storage/vram_blocks/active \
+    /app/storage/vram_blocks/archived \
+    /app/storage/vram_blocks/temp \
+    /app/storage/gpu_state \
+    /app/storage/cache
+
+# Expose ports
+EXPOSE 7860
+EXPOSE 8765
+
+# Set environment variable for storage path
+ENV STORAGE_PATH=/app/storage
+
+# Run the server using Uvicorn
+CMD ["uvicorn", "server:app", "--host", "0.0.0.0", "--port", "7860", "--ws", "websockets"]

config.json

@@ -0,0 +1,27 @@
+{
+    "server": {
+        "host": "localhost",
+        "port": 8765,
+        "max_connections": 1000
+    },
+    "storage": {
+        "vram": {
+            "block_size": 1048576,  // 1MB blocks
+            "max_cache_size": 8589934592,  // 8GB cache
+            "persistence": true
+        },
+        "state": {
+            "cache_ttl": 3600,  // 1 hour cache lifetime
+            "compression": true
+        },
+        "cache": {
+            "max_memory_size": 4294967296,  // 4GB memory cache
+            "eviction_policy": "lru"
+        }
+    },
+    "performance": {
+        "prefetch_enabled": true,
+        "compression_level": 1,
+        "cache_warmup": true
+    }
+}

requirements.txt

@@ -0,0 +1,19 @@
+# Virtual GPU Server Core Dependencies
+numpy==1.24.3         # Array operations for GPU data handling
+websockets==11.0.3    # WebSocket server implementation
+aiohttp==3.8.5       # HTTP server and async web framework
+uvicorn==0.23.2      # ASGI server implementation
+
+# WebSocket/HTTP Server Dependencies
+aiosignal==1.3.1     # Async signals
+async-timeout==4.0.3 # Timeouts for async operations
+attrs==23.1.0        # Class builders
+charset-normalizer==3.2.0  # Unicode normalization
+frozenlist==1.4.0    # Immutable lists
+multidict==6.0.4     # Dict with multiple values per key
+yarl==1.9.2          # URL handling
+
+# Performance & Utilities
+ujson==5.8.0         # Fast JSON processing
+psutil==5.9.5        # System and process monitoring
+pathlib==1.0.1       # Path manipulation

requirements_new.txt

@@ -0,0 +1,19 @@
+# Virtual GPU Server Core Dependencies
+numpy==1.24.3         # Array operations for GPU data handling
+fastapi==0.103.0     # FastAPI framework
+uvicorn==0.23.2      # ASGI server
+
+# WebSocket Dependencies
+websockets==11.0.3    # WebSocket support
+websocket-client==1.6.1  # WebSocket client
+
+# Performance & Utilities
+ujson==5.8.0         # Fast JSON processing
+psutil==5.9.5        # System and process monitoring
+pathlib==1.0.1       # Path manipulation
+
+# FastAPI Dependencies
+starlette==0.27.0    # Web framework required by FastAPI
+pydantic==2.3.0      # Data validation
+python-multipart==0.0.6  # Form/File uploads
+typing-extensions==4.7.1  # Enhanced typing support

server.py

@@ -0,0 +1,365 @@
+import asyncio
+import json
+import os
+from pathlib import Path
+import uuid
+import time
+from typing import Dict, Any, Optional
+import numpy as np
+from fastapi import FastAPI, WebSocket
+from fastapi.responses import HTMLResponse
+from datetime import datetime
+
+# Create FastAPI instance
+app = FastAPI()
+
+class VirtualGPUServer:
+    def __init__(self):
+        self.base_path = Path(__file__).parent / "storage"
+        self.vram_path = self.base_path / "vram_blocks"
+        self.state_path = self.base_path / "gpu_state"
+        self.cache_path = self.base_path / "cache"
+        
+        # Ensure all storage directories exist
+        self.vram_path.mkdir(parents=True, exist_ok=True)
+        self.state_path.mkdir(parents=True, exist_ok=True)
+        self.cache_path.mkdir(parents=True, exist_ok=True)
+        
+        # In-memory caches for faster access
+        self.vram_cache: Dict[str, Any] = {}
+        self.state_cache: Dict[str, Any] = {}
+        self.memory_cache: Dict[str, Any] = {}
+        
+        # Active connections and sessions
+        self.active_connections: Dict[str, WebSocket] = {}
+        self.active_sessions: Dict[str, Dict[str, Any]] = {}
+        self.heartbeat_interval = 5  # seconds
+        self.connection_timeout = 30  # seconds
+        
+        # Performance monitoring
+        self.ops_counter = 0
+        self.start_time = time.time()
+
+    def _make_json_serializable(self, obj):
+        """Convert non-JSON-serializable objects to serializable format"""
+        if isinstance(obj, dict):
+            return {k: self._make_json_serializable(v) for k, v in obj.items()}
+        elif isinstance(obj, list):
+            return [self._make_json_serializable(i) for i in obj]
+        elif isinstance(obj, tuple):
+            return list(obj)
+        elif isinstance(obj, (np.ndarray, np.generic)):
+            return obj.tolist()
+        elif isinstance(obj, (Path, uuid.UUID)):
+            return str(obj)
+        elif hasattr(obj, '__dict__'):
+            # Handle custom objects by converting their __dict__ to serializable format
+            return self._make_json_serializable(obj.__dict__)
+        elif isinstance(obj, (int, float, str, bool, type(None))):
+            return obj
+        else:
+            # Convert any other types to string representation
+            return str(obj)
+
+    async def monitor_connection(self, websocket: WebSocket, session_id: str):
+        """Monitor connection health and handle reconnection"""
+        try:
+            while session_id in self.active_connections:
+                try:
+                    await asyncio.wait_for(websocket.receive_text(), timeout=self.heartbeat_interval)
+                except asyncio.TimeoutError:
+                    try:
+                        await websocket.send_json({"type": "ping"})
+                    except:
+                        print(f"Connection lost for session {session_id}")
+                        break
+                except Exception:
+                    break
+                await asyncio.sleep(self.heartbeat_interval)
+        finally:
+            await self.handle_disconnect(session_id)
+
+    async def handle_disconnect(self, session_id: str):
+        """Clean up resources when a client disconnects"""
+        if session_id in self.active_connections:
+            try:
+                await self.active_connections[session_id].close()
+            except:
+                pass
+            del self.active_connections[session_id]
+        if session_id in self.active_sessions:
+            # Save any pending state before removing session
+            session_data = self.active_sessions[session_id]
+            if session_data.get('pending_state'):
+                await self.handle_state_operation({
+                    'type': 'save',
+                    'component': 'session',
+                    'state_id': session_id,
+                    'data': session_data['pending_state']
+                })
+            del self.active_sessions[session_id]
+
+    async def handle_vram_operation(self, operation: dict) -> dict:
+        """Handle VRAM read/write operations"""
+        try:
+            op_type = operation.get('type')
+            if not op_type:
+                raise ValueError("Missing operation type")
+            
+            block_id = operation.get('block_id')
+            if not block_id:
+                raise ValueError("Missing block_id")
+            
+            data = operation.get('data')
+            if data and isinstance(data, (dict, list)):
+                data = self._make_json_serializable(data)
+
+            if op_type == 'write':
+                if data is None:
+                    raise ValueError("Missing data for write operation")
+                file_path = self.vram_path / f"{block_id}.npy"
+                np.save(file_path, np.array(data))
+                self.vram_cache[block_id] = np.array(data)
+                return {'status': 'success', 'message': f'Block {block_id} written'}
+            
+            if op_type == 'read':
+                if block_id in self.vram_cache:
+                    return {
+                        'status': 'success',
+                        'data': self.vram_cache[block_id] if isinstance(self.vram_cache[block_id], list) else self.vram_cache[block_id].tolist(),
+                        'source': 'cache'
+                    }
+                
+                file_path = self.vram_path / f"{block_id}.npy"
+                if file_path.exists():
+                    data = np.load(file_path)
+                    self.vram_cache[block_id] = np.array(data)
+                    return {
+                        'status': 'success',
+                        'data': data.tolist(),
+                        'source': 'disk'
+                    }
+                return {'status': 'error', 'message': 'Block not found'}
+            
+            return {'status': 'error', 'message': f'Unknown operation type: {op_type}'}
+            
+        except ValueError as e:
+            return {'status': 'error', 'message': str(e)}
+        except Exception as e:
+            return {'status': 'error', 'message': f'Operation failed: {str(e)}'}
+                
+    async def handle_state_operation(self, operation: dict) -> dict:
+        """Handle GPU state operations"""
+        op_type = operation.get('type')
+        component = operation.get('component')
+        state_id = operation.get('state_id')
+        state_data = operation.get('data')
+
+        file_path = self.state_path / component / f"{state_id}.json"
+        
+        if op_type == 'save':
+            file_path.parent.mkdir(exist_ok=True)
+            with open(file_path, 'w') as f:
+                json.dump(state_data, f)
+            self.state_cache[f"{component}:{state_id}"] = state_data
+            return {'status': 'success', 'message': f'State {state_id} saved'}
+            
+        elif op_type == 'load':
+            cache_key = f"{component}:{state_id}"
+            if cache_key in self.state_cache:
+                return {
+                    'status': 'success',
+                    'data': self.state_cache[cache_key],
+                    'source': 'cache'
+                }
+                
+            if file_path.exists():
+                with open(file_path) as f:
+                    state_data = json.load(f)
+                self.state_cache[cache_key] = state_data
+                return {
+                    'status': 'success',
+                    'data': state_data,
+                    'source': 'disk'
+                }
+                
+            return {'status': 'error', 'message': 'State not found'}
+
+    async def handle_cache_operation(self, operation: dict) -> dict:
+        """Handle cache operations"""
+        op_type = operation.get('type')
+        key = operation.get('key')
+        data = operation.get('data')
+        
+        if op_type == 'set':
+            self.memory_cache[key] = data
+            # Also persist to disk for recovery
+            file_path = self.cache_path / f"{key}.json"
+            with open(file_path, 'w') as f:
+                json.dump(data, f)
+            return {'status': 'success', 'message': f'Cache key {key} set'}
+            
+        elif op_type == 'get':
+            if key in self.memory_cache:
+                return {
+                    'status': 'success',
+                    'data': self.memory_cache[key],
+                    'source': 'memory'
+                }
+                
+            file_path = self.cache_path / f"{key}.json"
+            if file_path.exists():
+                with open(file_path) as f:
+                    data = json.load(f)
+                self.memory_cache[key] = data
+                return {
+                    'status': 'success',
+                    'data': data,
+                    'source': 'disk'
+                }
+                
+            return {'status': 'error', 'message': 'Cache key not found'}
+
+
+    async def handle_connection(self, websocket: WebSocket, connection_type: str):
+        """Handle incoming WebSocket connections"""
+        # Generate unique session ID
+        session_id = str(uuid.uuid4())
+        self.active_connections[session_id] = websocket
+        self.active_sessions[session_id] = {
+            'start_time': time.time(),
+            'ops_count': 0,
+            'connection_type': connection_type
+        }
+        
+        try:
+            while True:
+                message = await websocket.receive_json()
+                
+                # Route operation to appropriate handler
+                operation_type = message.get('operation')
+                if operation_type == 'vram':
+                    response = await self.handle_vram_operation(message)
+                elif operation_type == 'state':
+                    response = await self.handle_state_operation(message)
+                elif operation_type == 'cache':
+                    response = await self.handle_cache_operation(message)
+                else:
+                    response = {
+                        'status': 'error',
+                        'message': 'Unknown operation type'
+                    }
+                
+                # Update statistics
+                self.ops_counter += 1
+                self.active_sessions[session_id]['ops_count'] += 1
+                
+                # Send response
+                await websocket.send_json(response)
+                
+        except Exception as e:
+            print(f"WebSocket error: {e}")
+        finally:
+            # Cleanup on disconnect
+            if session_id in self.active_connections:
+                del self.active_connections[session_id]
+            if session_id in self.active_sessions:
+                del self.active_sessions[session_id]
+
+    def get_stats(self) -> dict:
+        uptime = current_time - self.start_time
+        ops_per_second = self.ops_counter / uptime if uptime > 0 else 0
+        
+        return {
+            'uptime': uptime,
+            'total_operations': self.ops_counter,
+            'ops_per_second': ops_per_second,
+            'active_connections': len(self.active_connections),
+            'vram_cache_size': len(self.vram_cache),
+            'state_cache_size': len(self.state_cache),
+            'memory_cache_size': len(self.memory_cache)
+        }
+
+server = VirtualGPUServer()
+
+@app.get("/", response_class=HTMLResponse)
+async def handle_index():
+    """Handle HTTP index request"""
+    stats = server.get_stats()
+    html = f"""
+    <!DOCTYPE html>
+    <html>
+    <head>
+        <title>Virtual GPU Server</title>
+        <style>
+            body {{ font-family: Arial, sans-serif; margin: 40px; }}
+            table {{ border-collapse: collapse; width: 100%; margin-top: 20px; }}
+            th, td {{ padding: 12px; text-align: left; border-bottom: 1px solid #ddd; }}
+            th {{ background-color: #f2f2f2; }}
+            .stats {{ background-color: #f9f9f9; padding: 20px; border-radius: 5px; }}
+        </style>
+    </head>
+    <body>
+        <h1>Virtual GPU Server Status</h1>
+        <div class="stats">
+            <h2>Server Statistics</h2>
+            <ul>
+                <li>Uptime: {stats['uptime']:.2f} seconds</li>
+                <li>Total Operations: {stats['total_operations']}</li>
+                <li>Operations per Second: {stats['ops_per_second']:.2f}</li>
+                <li>Active Connections: {stats['active_connections']}</li>
+                <li>VRAM Cache Size: {stats['vram_cache_size']}</li>
+                <li>State Cache Size: {stats['state_cache_size']}</li>
+                <li>Memory Cache Size: {stats['memory_cache_size']}</li>
+            </ul>
+        </div>
+        <h2>Server Files</h2>
+        <iframe src="/files" style="width: 100%; height: 500px; border: none;"></iframe>
+    </body>
+    </html>
+    """
+    return HTMLResponse(content=html)
+
+@app.get("/files", response_class=HTMLResponse)
+async def handle_files():
+    """Handle HTTP files listing request"""
+    def format_size(size):
+        for unit in ['B', 'KB', 'MB', 'GB']:
+            if size < 1024:
+                return f"{size:.2f} {unit}"
+            size /= 1024
+        return f"{size:.2f} TB"
+
+    html = ['<!DOCTYPE html><html><head>',
+            '<style>',
+            'body { font-family: Arial, sans-serif; margin: 20px; }',
+            'table { border-collapse: collapse; width: 100%; }',
+            'th, td { padding: 12px; text-align: left; border-bottom: 1px solid #ddd; }',
+            'th { background-color: #f2f2f2; }',
+            '</style></head><body>',
+            '<h2>Server Files</h2>',
+            '<table><tr><th>Path</th><th>Size</th><th>Last Modified</th></tr>']
+
+    for root, _, files in os.walk(server.base_path):
+        for file in files:
+            full_path = Path(root) / file
+            rel_path = full_path.relative_to(server.base_path)
+            size = format_size(os.path.getsize(full_path))
+            mtime = datetime.fromtimestamp(os.path.getmtime(full_path))
+            html.append(f'<tr><td>{rel_path}</td><td>{size}</td><td>{mtime}</td></tr>')
+
+    html.extend(['</table></body></html>'])
+    return HTMLResponse(content='\n'.join(html))
+
+# WebSocket endpoint
+@app.websocket("/ws")
+async def websocket_endpoint(websocket: WebSocket):
+    await websocket.accept()
+    await server.handle_connection(websocket, "gpu")
+
+@app.websocket("/ws/model")
+async def websocket_model_endpoint(websocket: WebSocket):
+    await websocket.accept()
+    await server.handle_connection(websocket, "model")
+
+

storage/cache/model_info.json

@@ -0,0 +1 @@
+{"model_version": "v1.2", "accuracy": 0.98}

storage/gpu_state/gpu_core/current_state.json

@@ -0,0 +1 @@
+{"gpu_temp": 65, "fan_speed": 2000}

storage/vram_blocks/test_model_123_11.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6b77c2d6a8693eece477e2ac5cd425c1a2e1bc2d8bb1fbb684d3a6804588e01e
+size 8388736

storage/vram_blocks/test_model_123_14.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:39c2d4f10e76b70763505974c095aba40a2127a7153ad741390ad8e6f796c7b1
+size 8388736

storage_manager.py

@@ -0,0 +1,222 @@
+from pathlib import Path
+import json
+import shutil
+import os
+from typing import Dict, Any, Optional
+import numpy as np
+from datetime import datetime
+
+class StorageManager:
+    def __init__(self, base_path: Path):
+        self.base_path = base_path
+        
+        # Storage paths
+        self.vram_path = base_path / "vram_blocks"
+        self.state_path = base_path / "gpu_state"
+        self.cache_path = base_path / "cache"
+        
+        # Create directory structure
+        self._init_storage()
+        
+        # Component paths
+        self.component_paths = {
+            'cores': self.state_path / "cores",
+            'tensor_cores': self.state_path / "tensor_cores",
+            'warps': self.state_path / "warps",
+            'threads': self.state_path / "threads",
+            'memory': self.state_path / "memory"
+        }
+        
+    def _init_storage(self):
+        """Initialize storage directory structure with unlimited capacity support"""
+        # Create main directories with unlimited storage capacity
+        self.vram_path.mkdir(parents=True, exist_ok=True)
+        self.state_path.mkdir(parents=True, exist_ok=True)
+        self.cache_path.mkdir(parents=True, exist_ok=True)
+        
+        # Create subdirectories for different types of state
+        for component in ['cores', 'tensor_cores', 'warps', 'threads', 'memory']:
+            (self.state_path / component).mkdir(exist_ok=True)
+            
+        # Create VRAM block subdirectories with unlimited capacity
+        for subdir in ['active', 'archived', 'temp']:
+            (self.vram_path / subdir).mkdir(exist_ok=True)
+            
+        # Initialize storage configuration for unlimited capacity
+        self.storage_config = {
+            "unlimited_storage": True,
+            "dynamic_scaling": True,
+            "auto_compress": False  # Disabled for maximum performance
+        }
+
+    def store_vram_block(self, block_id: str, data: np.ndarray, temp: bool = False) -> Path:
+        """Store a VRAM block with metadata"""
+        # Determine storage location
+        if temp:
+            target_dir = self.vram_path / "temp"
+        else:
+            target_dir = self.vram_path / "active"
+            
+        file_path = target_dir / f"{block_id}.npz"
+        metadata_path = target_dir / f"{block_id}.meta.json"
+        
+        # Save data and metadata
+        np.savez_compressed(file_path, data=data)
+        metadata = {
+            'created': datetime.now().isoformat(),
+            'shape': data.shape,
+            'dtype': str(data.dtype),
+            'size_bytes': data.nbytes
+        }
+        with open(metadata_path, 'w') as f:
+            json.dump(metadata, f)
+            
+        return file_path
+
+    def load_vram_block(self, block_id: str) -> Optional[np.ndarray]:
+        """Load a VRAM block with fallback to archived"""
+        # Check active blocks first
+        active_path = self.vram_path / "active" / f"{block_id}.npz"
+        if active_path.exists():
+            return np.load(active_path)['data']
+            
+        # Check archived blocks
+        archived_path = self.vram_path / "archived" / f"{block_id}.npz"
+        if archived_path.exists():
+            data = np.load(archived_path)['data']
+            # Move back to active
+            self.store_vram_block(block_id, data)
+            return data
+            
+        return None
+
+    def store_component_state(self, component: str, state_id: str, state: Dict) -> Path:
+        """Store component state with versioning"""
+        component_path = self.component_paths[component]
+        state_dir = component_path / state_id
+        state_dir.mkdir(exist_ok=True)
+        
+        # Create versioned state file
+        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+        state_file = state_dir / f"state_{timestamp}.json"
+        
+        with open(state_file, 'w') as f:
+            json.dump(state, f)
+            
+        # Update latest symlink
+        latest_link = state_dir / "latest.json"
+        if latest_link.exists():
+            latest_link.unlink()
+        latest_link.symlink_to(state_file)
+        
+        return state_file
+
+    def load_component_state(self, component: str, state_id: str, version: str = "latest") -> Optional[Dict]:
+        """Load component state with version support"""
+        component_path = self.component_paths[component]
+        state_dir = component_path / state_id
+        
+        if version == "latest":
+            state_file = state_dir / "latest.json"
+        else:
+            state_file = state_dir / f"state_{version}.json"
+            
+        if state_file.exists():
+            with open(state_file) as f:
+                return json.load(f)
+                
+        return None
+
+    def archive_vram_block(self, block_id: str):
+        """Move a VRAM block to archive"""
+        source_path = self.vram_path / "active" / f"{block_id}.npz"
+        target_path = self.vram_path / "archived" / f"{block_id}.npz"
+        
+        if source_path.exists():
+            shutil.move(source_path, target_path)
+            
+        # Also move metadata
+        source_meta = self.vram_path / "active" / f"{block_id}.meta.json"
+        target_meta = self.vram_path / "archived" / f"{block_id}.meta.json"
+        if source_meta.exists():
+            shutil.move(source_meta, target_meta)
+
+    def cleanup_temp_blocks(self, max_age_hours: int = 24):
+        """Clean up old temporary VRAM blocks"""
+        temp_dir = self.vram_path / "temp"
+        current_time = datetime.now()
+        
+        for file_path in temp_dir.glob("*.npz"):
+            # Check corresponding metadata
+            meta_path = file_path.with_suffix('.meta.json')
+            if meta_path.exists():
+                with open(meta_path) as f:
+                    metadata = json.load(f)
+                created = datetime.fromisoformat(metadata['created'])
+                
+                # Remove if too old
+                if (current_time - created).total_seconds() > max_age_hours * 3600:
+                    file_path.unlink()
+                    meta_path.unlink()
+
+    def get_storage_stats(self) -> Dict[str, Any]:
+        """Get storage statistics"""
+        stats = {
+            'vram': {
+                'active_blocks': len(list(self.vram_path.glob("active/*.npz"))),
+                'archived_blocks': len(list(self.vram_path.glob("archived/*.npz"))),
+                'temp_blocks': len(list(self.vram_path.glob("temp/*.npz"))),
+                'total_size_bytes': self._get_dir_size(self.vram_path)
+            },
+            'state': {
+                'components': {},
+                'total_size_bytes': self._get_dir_size(self.state_path)
+            },
+            'cache': {
+                'size_bytes': self._get_dir_size(self.cache_path)
+            }
+        }
+        
+        # Add component-specific stats
+        for component in self.component_paths:
+            stats['state']['components'][component] = {
+                'states': len(list(self.component_paths[component].glob("**/state_*.json"))),
+                'size_bytes': self._get_dir_size(self.component_paths[component])
+            }
+            
+        return stats
+
+    def _get_dir_size(self, path: Path) -> int:
+        """Calculate total size of a directory"""
+        return sum(f.stat().st_size for f in path.glob('**/*') if f.is_file())
+
+    def create_snapshot(self, snapshot_id: str):
+        """Create a snapshot of the entire storage"""
+        snapshot_dir = self.base_path / "snapshots" / snapshot_id
+        snapshot_dir.mkdir(parents=True)
+        
+        # Copy current state
+        shutil.copytree(self.state_path, snapshot_dir / "gpu_state")
+        shutil.copytree(self.vram_path / "active", snapshot_dir / "vram_blocks")
+        
+        # Create snapshot metadata
+        metadata = {
+            'created': datetime.now().isoformat(),
+            'stats': self.get_storage_stats()
+        }
+        with open(snapshot_dir / "snapshot.json", 'w') as f:
+            json.dump(metadata, f)
+
+    def restore_snapshot(self, snapshot_id: str):
+        """Restore from a snapshot"""
+        snapshot_dir = self.base_path / "snapshots" / snapshot_id
+        if not snapshot_dir.exists():
+            raise ValueError(f"Snapshot {snapshot_id} not found")
+            
+        # Clear current state
+        shutil.rmtree(self.state_path)
+        shutil.rmtree(self.vram_path / "active")
+        
+        # Restore from snapshot
+        shutil.copytree(snapshot_dir / "gpu_state", self.state_path)
+        shutil.copytree(snapshot_dir / "vram_blocks", self.vram_path / "active")

websocket_test.py

@@ -0,0 +1,80 @@
+import asyncio
+import websockets
+import json
+import numpy as np
+
+async def test_websocket():
+    uri = "wss://8765-int35aifjdyxr46oxipjg-cf14c883.manusvm.computer"
+    async with websockets.connect(uri) as websocket:
+        print("Connected to websocket.")
+
+        # Test VRAM write operation
+        vram_data = np.random.rand(10, 10).tolist()
+        vram_write_message = {
+            "operation": "vram",
+            "type": "write",
+            "block_id": "test_block_1",
+            "data": vram_data
+        }
+        await websocket.send(json.dumps(vram_write_message))
+        response = await websocket.recv()
+        print(f"VRAM write response: {response}")
+
+        # Test VRAM read operation
+        vram_read_message = {
+            "operation": "vram",
+            "type": "read",
+            "block_id": "test_block_1"
+        }
+        await websocket.send(json.dumps(vram_read_message))
+        response = await websocket.recv()
+        print(f"VRAM read response: {response}")
+
+        # Test state save operation
+        state_data = {"gpu_temp": 65, "fan_speed": 2000}
+        state_save_message = {
+            "operation": "state",
+            "type": "save",
+            "component": "gpu_core",
+            "state_id": "current_state",
+            "data": state_data
+        }
+        await websocket.send(json.dumps(state_save_message))
+        response = await websocket.recv()
+        print(f"State save response: {response}")
+
+        # Test state load operation
+        state_load_message = {
+            "operation": "state",
+            "type": "load",
+            "component": "gpu_core",
+            "state_id": "current_state"
+        }
+        await websocket.send(json.dumps(state_load_message))
+        response = await websocket.recv()
+        print(f"State load response: {response}")
+
+        # Test cache set operation
+        cache_data = {"model_version": "v1.2", "accuracy": 0.98}
+        cache_set_message = {
+            "operation": "cache",
+            "type": "set",
+            "key": "model_info",
+            "data": cache_data
+        }
+        await websocket.send(json.dumps(cache_set_message))
+        response = await websocket.recv()
+        print(f"Cache set response: {response}")
+
+        # Test cache get operation
+        cache_get_message = {
+            "operation": "cache",
+            "type": "get",
+            "key": "model_info"
+        }
+        await websocket.send(json.dumps(cache_get_message))
+        response = await websocket.recv()
+        print(f"Cache get response: {response}")
+
+asyncio.run(test_websocket())
+