Update http_storage.py

http_storage.py

@@ -158,85 +158,71 @@ class HTTPGPUStorage:
             return False
 
     def load_tensor(self, tensor_id: str) -> Optional[np.ndarray]:
-        """Load tensor data via HTTP API"""
+        """Load tensor data from local storage"""
         try:
-            # Check tensor registry first
-            if tensor_id not in self.tensor_registry:
-                logging.warning(f"Tensor {tensor_id} not registered in VRAM")
-                # Still try to load it in case it exists on server
-            
-            response = self._make_request('GET', f'/vram/blocks/{tensor_id}')
+            tensor_path = self.vram_path / f"{tensor_id}.npy"
+            metadata_path = self.vram_path / f"{tensor_id}_meta.json"
             
-            if response and response.get('status') == 'success':
-                data = response.get('data')
-                metadata = response.get('metadata', {})
-                
-                if data is None:
-                    logging.error(f"No data found for tensor {tensor_id}")
-                    return None
-                
-                try:
-                    # Convert to numpy array with correct dtype
-                    expected_dtype = metadata.get('dtype', 'float32')
-                    expected_shape = metadata.get('shape')
-                    
-                    arr = np.array(data, dtype=np.dtype(expected_dtype))
-                    if expected_shape and arr.shape != tuple(expected_shape):
-                        arr = arr.reshape(expected_shape)
-                    
-                    # Update registry if not present
-                    if tensor_id not in self.tensor_registry:
-                        with self.lock:
-                            self.tensor_registry[tensor_id] = metadata
-                    
-                    return arr
-                    
-                except Exception as e:
-                    logging.error(f"Error converting tensor data: {str(e)}")
-                    return None
-            else:
-                logging.error(f"Failed to load tensor {tensor_id}: {response.get('message', 'Unknown error')}")
+            # Check if tensor files exist
+            if not tensor_path.exists() or not metadata_path.exists():
+                logging.warning(f"Tensor {tensor_id} not found in local storage")
                 return None
+            
+            # Load metadata
+            with open(metadata_path, 'r') as f:
+                metadata = json.load(f)
+            
+            # Load tensor data
+            arr = np.load(str(tensor_path))
+            
+            # Update registry if not present
+            if tensor_id not in self.tensor_registry:
+                with self.lock:
+                    self.tensor_registry[tensor_id] = metadata
+            
+            return arr
                 
         except Exception as e:
             logging.error(f"Error loading tensor {tensor_id}: {str(e)}")
             return None
 
     def store_state(self, component: str, state_id: str, state_data: Dict[str, Any]) -> bool:
-        """Store component state via HTTP API"""
+        """Store component state in local storage"""
         try:
-            request_data = {
+            # Create component directory if needed
+            component_dir = self.state_path / component
+            component_dir.mkdir(parents=True, exist_ok=True)
+            
+            # Save state data with timestamp
+            state_file = component_dir / f"{state_id}.json"
+            data_to_save = {
                 "data": state_data,
                 "timestamp": time.time()
             }
             
-            response = self._make_request(
-                'POST',
-                f'/state/{component}/{state_id}',
-                json=request_data
-            )
+            with open(state_file, 'w') as f:
+                json.dump(data_to_save, f, indent=2)
             
-            if response and response.get('status') == 'success':
-                return True
-            else:
-                logging.error(f"Failed to store state for {component}/{state_id}: {response.get('message', 'Unknown error')}")
-                return False
+            return True
                 
         except Exception as e:
             logging.error(f"Error storing state for {component}/{state_id}: {str(e)}")
             return False
 
     def load_state(self, component: str, state_id: str) -> Optional[Dict[str, Any]]:
-        """Load component state via HTTP API"""
+        """Load component state from local storage"""
         try:
-            response = self._make_request("GET", f"/api/v1/state/{component}/{state_id}")
+            state_file = self.state_path / component / f"{state_id}.json"
             
-            if response and response.get('status') == 'success':
-                return response.get('data')
-            else:
-                logging.error(f"Failed to load state for {component}/{state_id}: {response.get('message', 'Unknown error')}")
+            if not state_file.exists():
+                logging.warning(f"State file not found for {component}/{state_id}")
                 return None
                 
+            with open(state_file, 'r') as f:
+                saved_data = json.load(f)
+                
+            return saved_data.get('data')
+                
         except Exception as e:
             logging.error(f"Error loading state for {component}/{state_id}: {str(e)}")
             return None
@@ -271,23 +257,6 @@ class HTTPGPUStorage:
             logging.error(f"Error getting cached data for key {key}: {str(e)}")
             return None
 
-    def is_model_loaded(self, model_name: str) -> bool:
-        """Check if a model is loaded via HTTP API"""
-        try:
-            response = self._make_request(
-                "GET", 
-                f"/models/{model_name}/status",
-                timeout=60
-            )
-            
-            if response and response.get('status') == 'loaded':
-                return True
-            return False
-            
-        except Exception as e:
-            logging.error(f"Error checking model status for {model_name}: {str(e)}")
-            return False
-
     def load_model(self, model_name: str, model_path: Optional[str] = None, model_data: Optional[Dict] = None) -> bool:
         """Load a model from local storage"""
         try:
@@ -349,92 +318,18 @@ class HTTPGPUStorage:
             logging.error(f"Error calculating model hash: {str(e)}")
             return ""
 
-    def start_inference(self, model_name: str, input_data: np.ndarray) -> Optional[Dict[str, Any]]:
-        """Start inference with a loaded model via HTTP API"""
-        try:
-            if not self.is_model_loaded(model_name):
-                logging.error(f"Model {model_name} not loaded. Please load the model first.")
-                return None
-
-            request_data = {
-                "input_data": input_data.tolist() if isinstance(input_data, np.ndarray) else input_data
-            }
-
-            response = self._make_request(
-                'POST',
-                f'/models/{model_name}/inference',
-                json=request_data
-            )
-
-            if response and response.get('status') == 'success':
-                return {
-                    'output': np.array(response['output']) if 'output' in response else None,
-                    'metrics': response.get('metrics', {}),
-                    'model_info': self.model_registry.get(model_name, {})
-                }
-            else:
-                logging.error(f"Inference failed for model {model_name}: {response.get('message', 'Unknown error')}")
-                return None
-                
-        except Exception as e:
-            logging.error(f"Error during inference for model {model_name}: {str(e)}")
-            return None
 
     def ping(self) -> bool:
-        """Ping the server to check connection status."""
+        """Check if local storage is accessible"""
         try:
-            response = self._make_request('GET', '/status')
-            return response and response.get('status') == 'ok'
+            # Check if all storage directories exist and are accessible
+            for path in [self.vram_path, self.models_path, self.cache_path, self.state_path]:
+                if not path.exists() or not os.access(str(path), os.R_OK | os.W_OK):
+                    return False
+            return True
         except Exception as e:
-            logging.error(f"Ping failed: {e}")
+            logging.error(f"Storage check failed: {e}")
             return False
-
-    def is_connected(self) -> bool:
-        """Check if the client is connected to the server."""
-        return self.ping()
-
-    def get_connection_status(self) -> Dict[str, Any]:
-        """Get detailed connection status."""
-        if self.is_connected():
-            return {"status": "connected", "session_id": self.session_id}
-        else:
-            return {"status": "disconnected", "error_count": self.error_count}
-
-    def set_keep_alive(self, interval: int):
-        """Set keep-alive interval (compatibility method)."""
-        logging.info(f"Keep-alive interval set to {interval} seconds (HTTP client does not use websockets).")
-
-    def reconnect(self):
-        """Attempt to reconnect (compatibility method)."""
-        logging.info("Attempting to reconnect HTTP client...")
-        self._create_session()
-
-    def wait_for_connection(self, timeout: float = 30.0) -> bool:
-        """Wait for HTTP connection to be established (compatibility method)"""
-        start_time = time.time()
-        while time.time() - start_time < timeout:
-            if self.is_connected():
-                logging.info("HTTP connection established.")
-                return True
-            time.sleep(1) # Wait for 1 second before retrying
-        logging.error("HTTP connection not established within timeout.")
-        return False
-
-    def close(self):
-        """Close HTTP client"""
-        self._closing = True
-        logging.info("HTTP client is closing.")
-        # Invalidate session on server side if possible
-        if self.session_token:
-            try:
-                self.http_session.post(f"{self.api_base}/sessions/invalidate",
-                                       headers={'Authorization': f'Bearer {self.session_token}'},
-                                       timeout=5)
-            except Exception as e:
-                logging.warning(f"Failed to invalidate session on server: {e}")
-        self.http_session.close()
-        HTTPGPUStorage._instance = None # Clear singleton instance
-
 # Compatibility alias for existing code
 WebSocketGPUStorage = HTTPGPUStorage