🚧 fixed storage interface issue by adding fresh storage per request

m-misiura · m-misiura · commit 8a0eb289c586 · 2025-06-03T11:52:24.000+01:00
diff --git a/src/endpoints/data/data_upload.py b/src/endpoints/data/data_upload.py
@@ -7,7 +7,7 @@
 from fastapi import APIRouter, HTTPException
 from pydantic import BaseModel
 
-from src.service.constants import METADATA_SUFFIX, OUTPUT_SUFFIX
+from src.service.constants import INPUT_SUFFIX, METADATA_SUFFIX, OUTPUT_SUFFIX
 from src.service.data.modelmesh_parser import ModelMeshPayloadParser
 from src.service.data.storage import get_storage_interface
 from src.service.utils.upload import (
@@ -21,8 +21,6 @@
 
 router = APIRouter()
 logger = logging.getLogger(__name__)
-storage = get_storage_interface()
-
 
 class UploadPayload(BaseModel):
     model_name: str
@@ -35,69 +33,81 @@ class UploadPayload(BaseModel):
 @router.post("/data/upload")
 async def upload(payload: UploadPayload) -> Dict[str, str]:
     """Upload model data - regular or ground truth."""
-    model_name = ModelMeshPayloadParser.standardize_model_id(payload.model_name)
-    if payload.data_tag and (error := validate_data_tag(payload.data_tag)):
-        raise HTTPException(400, error)
-    inputs = payload.request.get("inputs", [])
-    outputs = payload.response.get("outputs", [])
-    if not inputs or not outputs:
-        raise HTTPException(400, "Missing input or output tensors")
-    input_arrays, input_names, _, execution_ids = process_tensors(inputs)
-    output_arrays, output_names, _, _ = process_tensors(outputs)
-    if error := validate_input_shapes(input_arrays, input_names):
-        raise HTTPException(400, f"One or more errors in input tensors: {error}")
-    if payload.is_ground_truth:
-        if not execution_ids:
-            raise HTTPException(400, "Ground truth requires execution IDs")
-        result = await handle_ground_truths(
-            model_name,
-            input_arrays,
-            input_names,
-            output_arrays,
-            output_names,
-            [sanitize_id(id) for id in execution_ids],
-        )
-        if not result.success:
-            raise HTTPException(400, result.message)
-        result_data = result.data
-        if result_data is None:
-            raise HTTPException(500, "Ground truth processing failed")
-        gt_name = f"{model_name}_ground_truth"
-        await storage.write_data(gt_name + OUTPUT_SUFFIX, result_data["outputs"], result_data["output_names"])
-        await storage.write_data(
-            gt_name + METADATA_SUFFIX,
-            result_data["metadata"],
-            result_data["metadata_names"],
-        )
-        return {"message": result.message}
-    else:
-        n_rows = input_arrays[0].shape[0]
-        exec_ids = execution_ids or [str(uuid.uuid4()) for _ in range(n_rows)]
+    try:
+        # Get fresh storage interface for each request
+        storage = get_storage_interface()
+        
+        model_name = ModelMeshPayloadParser.standardize_model_id(payload.model_name)
+        if payload.data_tag and (error := validate_data_tag(payload.data_tag)):
+            raise HTTPException(400, error)
+        inputs = payload.request.get("inputs", [])
+        outputs = payload.response.get("outputs", [])
+        if not inputs or not outputs:
+            raise HTTPException(400, "Missing input or output tensors")
+        input_arrays, input_names, _, execution_ids = process_tensors(inputs)
+        output_arrays, output_names, _, _ = process_tensors(outputs)
+        if error := validate_input_shapes(input_arrays, input_names):
+            raise HTTPException(400, f"One or more errors in input tensors: {error}")
+        
+        if payload.is_ground_truth:
+            if not execution_ids:
+                raise HTTPException(400, "Ground truth requires execution IDs")
+            result = await handle_ground_truths(
+                model_name,
+                input_arrays,
+                input_names,
+                output_arrays,
+                output_names,
+                [sanitize_id(id) for id in execution_ids],
+            )
+            if not result.success:
+                raise HTTPException(400, result.message)
+            result_data = result.data
+            if result_data is None:
+                raise HTTPException(500, "Ground truth processing failed")
+            gt_name = f"{model_name}_ground_truth"
+            await storage.write_data(gt_name + OUTPUT_SUFFIX, result_data["outputs"], result_data["output_names"])
+            await storage.write_data(
+                gt_name + METADATA_SUFFIX,
+                result_data["metadata"],
+                result_data["metadata_names"],
+            )
+            return {"message": result.message}
+        else:
+            n_rows = input_arrays[0].shape[0]
+            exec_ids = execution_ids or [str(uuid.uuid4()) for _ in range(n_rows)]
 
-        def flatten(arrays: List[np.ndarray], row: int) -> List[Any]:
-            return [x for arr in arrays for x in (arr[row].flatten() if arr.ndim > 1 else [arr[row]])]
+            def flatten(arrays: List[np.ndarray], row: int) -> List[Any]:
+                return [x for arr in arrays for x in (arr[row].flatten() if arr.ndim > 1 else [arr[row]])]
 
-        input_data = [flatten(input_arrays, i) for i in range(n_rows)]
-        output_data = [flatten(output_arrays, i) for i in range(n_rows)]
-        cols = ["id", "model_id", "timestamp", "tag"]
-        current_timestamp = datetime.now().isoformat()
-        metadata_rows = [
-            [
-                str(eid),
-                str(model_name),
-                str(current_timestamp),
-                str(payload.data_tag or ""),
+            input_data = [flatten(input_arrays, i) for i in range(n_rows)]
+            output_data = [flatten(output_arrays, i) for i in range(n_rows)]
+            cols = ["id", "model_id", "timestamp", "tag"]
+            current_timestamp = datetime.now().isoformat()
+            metadata_rows = [
+                [
+                    str(eid),
+                    str(model_name),
+                    str(current_timestamp),
+                    str(payload.data_tag or ""),
+                ]
+                for eid in exec_ids
             ]
-            for eid in exec_ids
-        ]
-        metadata = np.array(metadata_rows, dtype="<U100")
-        await save_model_data(
-            model_name,
-            np.array(input_data),
-            input_names,
-            np.array(output_data),
-            output_names,
-            metadata,
-            cols,
-        )
-        return {"message": f"{n_rows} datapoints added to {model_name}"}
+            metadata = np.array(metadata_rows, dtype="<U100")
+            await save_model_data(
+                model_name,
+                np.array(input_data),
+                input_names,
+                np.array(output_data),
+                output_names,
+                metadata,
+                cols,
+            )
+            return {"message": f"{n_rows} datapoints added to {model_name}"}
+            
+    except HTTPException:
+        # Re-raise HTTP exceptions as-is
+        raise
+    except Exception as e:
+        logger.error(f"Unexpected error in upload endpoint for model {payload.model_name}: {str(e)}", exc_info=True)
+        raise HTTPException(500, f"Internal server error: {str(e)}")
diff --git a/src/service/utils/upload.py b/src/service/utils/upload.py
@@ -170,6 +170,8 @@ def __init__(self, model_name: str):
 
     async def initialize(self) -> None:
         """Load existing data."""
+        # Get fresh storage interface for each call
+        storage_interface = get_storage_interface()
         self.inputs, _ = await storage_interface.read_data(self.model_name + INPUT_SUFFIX)
         self.outputs, _ = await storage_interface.read_data(self.model_name + OUTPUT_SUFFIX)
         self.metadata, _ = await storage_interface.read_data(self.model_name + METADATA_SUFFIX)
@@ -216,6 +218,8 @@ async def validate_data(
                 return f"ID={exec_id} output type mismatch at position {i + 1}: Class={existing_type} != Class={uploaded_type}"
         if output_names:
             try:
+                # Get fresh storage interface for each call
+                storage_interface = get_storage_interface()
                 stored_output_names = await storage_interface.get_original_column_names(self.model_name + OUTPUT_SUFFIX)
                 if len(stored_output_names) != len(output_names):
                     return (
@@ -233,6 +237,8 @@ async def validate_data(
                 logger.warning(f"Could not validate output names for {exec_id}: {e}")
         if input_names:
             try:
+                # Get fresh storage interface for each call
+                storage_interface = get_storage_interface()
                 stored_input_names = await storage_interface.get_original_column_names(self.model_name + INPUT_SUFFIX)
                 if len(stored_input_names) != len(input_names):
                     return (
@@ -262,6 +268,8 @@ async def handle_ground_truths(
     """Handle ground truth validation."""
     if not execution_ids:
         return GroundTruthValidationResult(success=False, message="No execution IDs provided.")
+    # Get fresh storage interface for each call
+    storage_interface = get_storage_interface()
     if not await storage_interface.dataset_exists(model_name + INPUT_SUFFIX):
         return GroundTruthValidationResult(success=False, message=f"Model {model_name} not found.")
     validator = GroundTruthValidator(model_name)
@@ -316,11 +324,13 @@ async def save_model_data(
     metadata_data: np.ndarray,
     metadata_names: List[str],
 ) -> Dict[str, Any]:
+    # Get fresh storage interface for each call
+    storage_interface = get_storage_interface()
     """Save model data to storage."""
     await storage_interface.write_data(model_name + INPUT_SUFFIX, input_data, input_names)
     await storage_interface.write_data(model_name + OUTPUT_SUFFIX, output_data, output_names)
     await storage_interface.write_data(model_name + METADATA_SUFFIX, metadata_data, metadata_names)
     return {
         "model_name": model_name,
         "rows": len(input_data),
-    }
+    }
diff --git a/tests/endpoints/test_upload_endpoint.py b/tests/endpoints/test_upload_endpoint.py
@@ -6,7 +6,6 @@
 import sys
 import tempfile
 import uuid
-import asyncio
 
 import h5py
 import numpy as np
@@ -25,13 +24,7 @@
 )
 from src.service.data.storage import get_storage_interface
 
-@pytest.fixture(autouse=True)
-def reset_storage():
-    """Reset storage interface for each test."""
-    import src.service.data.storage
-    src.service.data.storage._storage_interface = None
-    yield
-    
+
 def pytest_sessionfinish(session, exitstatus):
     """Clean up the temporary directory after all tests are done."""
     if os.path.exists(TEMP_DIR):
@@ -136,43 +129,46 @@ def generate_mismatched_shape_no_unique_name_multi_input_payload(n_rows, n_input
 
 def get_data_from_storage(model_name, suffix):
     """Get data from storage file."""
-    async def _get_data():
-        storage = get_storage_interface()
-        try:
-            data, column_names = await storage.read_data(model_name + suffix)
-            return {"data": data, "column_names": column_names}
-        except Exception as e:
-            print(f"Error reading {model_name + suffix}: {e}")
-            return None
-    
-    try:
-        return asyncio.run(_get_data())
-    except Exception as e:
-        print(f"Async error for {model_name + suffix}: {e}")
+    storage = get_storage_interface()
+    filename = storage._get_filename(model_name + suffix)
+    if not os.path.exists(filename):
         return None
+    with h5py.File(filename, "r") as f:
+        if model_name + suffix in f:
+            data = f[model_name + suffix][:]
+            column_names = f[model_name + suffix].attrs.get("column_names", [])
+            return {"data": data, "column_names": column_names}
 
 
 def get_metadata_ids(model_name):
     """Extract actual IDs from metadata storage."""
-    async def _get_ids():
-        storage = get_storage_interface()
-        try:
-            metadata, column_names = await storage.read_data(model_name + METADATA_SUFFIX)
+    storage = get_storage_interface()
+    filename = storage._get_filename(model_name + METADATA_SUFFIX)
+    if not os.path.exists(filename):
+        return []
+    ids = []
+    with h5py.File(filename, "r") as f:
+        if model_name + METADATA_SUFFIX in f:
+            metadata = f[model_name + METADATA_SUFFIX][:]
+            column_names = f[model_name + METADATA_SUFFIX].attrs.get("column_names", [])
             id_idx = next((i for i, name in enumerate(column_names) if name.lower() == "id"), None)
-            if id_idx is not None and metadata is not None:
-                ids = []
+            if id_idx is not None:
                 for row in metadata:
-                    if hasattr(row, "__len__") and len(row) > id_idx:
-                        ids.append(str(row[id_idx]))
-                return ids
-        except Exception as e:
-            print(f"Error getting metadata: {e}")
-        return []
-    
-    try:
-        return asyncio.run(_get_ids())
-    except Exception:
-        return []
+                    try:
+                        if hasattr(row, "__getitem__") and len(row) > id_idx:
+                            id_val = row[id_idx]
+                        else:
+                            row_data = pickle.loads(row.tobytes())
+                            id_val = row_data[id_idx]
+                        if isinstance(id_val, np.ndarray):
+                            ids.append(str(id_val))
+                        else:
+                            ids.append(str(id_val))
+                    except Exception as e:
+                        print(f"Error processing ID from row {len(ids)}: {e}")
+                        continue
+    print(f"Successfully extracted {len(ids)} IDs: {ids}")
+    return ids
 
 
 def get_metadata_from_storage(model_name):
@@ -443,4 +439,4 @@ def test_upload_gaussian_data():
             ],
         },
     }
-    post_test(payload, 200, ["2 datapoints"])
+    post_test(payload, 200, ["2 datapoints"])
