Update virtual_gpu_server_http.py

virtual_gpu_server_http.py

@@ -649,24 +649,54 @@ async def load_model(
         # Get safe filename for storage
         safe_name = sanitize_filename(model_name)
         
-        # Store model information
+        if not request.model_data:
+            raise HTTPException(
+                status_code=400,
+                detail="model_data is required and must include architecture configuration"
+            )
+            
+        # Validate required model configuration
+        required_fields = ['num_sms', 'tensor_cores_per_sm', 'cuda_cores_per_sm']
+        missing_fields = [field for field in required_fields if field not in request.model_data]
+        if missing_fields:
+            raise HTTPException(
+                status_code=400,
+                detail=f"Missing required model configuration fields: {missing_fields}"
+            )
+
+        # Store model information with full configuration
         model_info = {
             'model_name': model_name,
             'model_data': request.model_data,
             'model_path': request.model_path,
             'model_hash': request.model_hash,
             'loaded_at': time.time(),
-            'session_id': session['session_id']
+            'session_id': session['session_id'],
+            'architecture': {
+                'num_sms': request.model_data['num_sms'],
+                'tensor_cores_per_sm': request.model_data['tensor_cores_per_sm'],
+                'cuda_cores_per_sm': request.model_data['cuda_cores_per_sm'],
+                'vram_allocation': request.model_data.get('vram_allocation', 'dynamic'),
+                'compute_capability': request.model_data.get('compute_capability', '8.0')
+            }
         }
         
         server.model_cache[model_name] = model_info
         
         # Store in persistent storage
         model_file = server.models_path / f"{safe_name}.json"
+        model_data_file = server.models_path / f"{safe_name}.data"
         logging.info(f"Storing model info at: {model_file}")
         
+        # Store metadata and configuration
         with open(model_file, 'w') as f:
             json.dump(model_info, f)
+            
+        # Store actual model data separately
+        if request.model_data.get('weights') or request.model_data.get('parameters'):
+            logging.info(f"Storing model data at: {model_data_file}")
+            with open(model_data_file, 'w') as f:
+                json.dump(request.model_data, f)
         
         server.ops_counter += 1
         return {
@@ -674,6 +704,7 @@ async def load_model(
             "message": f"Model {model_name} loaded successfully",
             "model_info": {
                 "name": model_name,
+                "architecture": model_info['architecture'],
                 "loaded_at": model_info['loaded_at']
             }
         }