Update server.py

server.py

@@ -104,51 +104,51 @@ class VirtualGPUServer:
         """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')
-            data = operation.get('data')
+            if not block_id:
+                raise ValueError("Missing block_id")
             
-            # Handle non-serializable objects
+            data = operation.get('data')
             if data and isinstance(data, (dict, list)):
                 data = self._make_json_serializable(data)
 
             if op_type == 'write':
-                # Additional error handling and logging
-                if not block_id:
-                    raise ValueError("Missing block_id for VRAM write operation")
                 if data is None:
-                    raise ValueError("Missing data for VRAM write operation")
-
-        if op_type == 'write':
-            # Generate unique file path for this block
-            file_path = self.vram_path / f"{block_id}.npy"
-            # Save data to file
-            np.save(file_path, np.array(data))
-            # Update cache
-            self.vram_cache[block_id] = np.array(data)
-            return {'status': 'success', 'message': f'Block {block_id} written'}
+                    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'}
             
-        elif op_type == 'read':
-            # Try cache first
-            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'
-                }
+            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'}
             
-            # Load from file
-            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': 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)}'}
                 
-            return {'status': 'error', 'message': 'Block not found'}
-
     async def handle_state_operation(self, operation: dict) -> dict:
         """Handle GPU state operations"""
         op_type = operation.get('type')