✨ added upload and download endpoints together with tests akin to the Java ones

Signed-off-by: m-misiura <[email protected]>

Author: m-misiura

src/endpoints/data/data_download.py

@@ -1,32 +1,57 @@
-from fastapi import APIRouter, HTTPException
-from pydantic import BaseModel
-from typing import List, Any, Optional
 import logging
 
-router = APIRouter()
-logger = logging.getLogger(__name__)
-
-
-class RowMatcher(BaseModel):
-    columnName: str
-    operation: str
-    values: List[Any]
+import pandas as pd
+from fastapi import APIRouter, HTTPException
 
+from src.service.utils.download import (
+    DataRequestPayload,
+    DataResponsePayload,
+    apply_matcher,
+    load_model_dataframe,
+)
 
-class DataRequestPayload(BaseModel):
-    modelId: str
-    matchAny: Optional[List[RowMatcher]] = None
-    matchAll: Optional[List[RowMatcher]] = None
-    matchNone: Optional[List[RowMatcher]] = None
+router = APIRouter()
+logger = logging.getLogger(__name__)
 
 
 @router.post("/data/download")
-async def download_data(payload: DataRequestPayload):
-    """Download model data."""
+async def download_data(payload: DataRequestPayload) -> DataResponsePayload:
+    """Download model data with filtering."""
     try:
         logger.info(f"Received data download request for model: {payload.modelId}")
-        # TODO: Implement
-        return {"status": "success", "data": []}
+
+        # Load the dataframe
+        df = await load_model_dataframe(payload.modelId)
+
+        if df.empty:
+            return DataResponsePayload(dataCSV="")
+        # Apply matchAll filters (AND logic)
+        if payload.matchAll:
+            for matcher in payload.matchAll:
+                df = apply_matcher(df, matcher, negate=False)
+        # Apply matchNone filters (NOT logic)
+        if payload.matchNone:
+            for matcher in payload.matchNone:
+                df = apply_matcher(df, matcher, negate=True)
+        base_df = df.copy()
+        # Apply matchAny filters (OR logic)
+        if payload.matchAny:
+            matching_dfs = []
+            for matcher in payload.matchAny:
+                matched_df = apply_matcher(base_df, matcher, negate=False)
+                if not matched_df.empty:
+                    matching_dfs.append(matched_df)
+            # Union all results
+            if matching_dfs:
+                df = pd.concat(matching_dfs, ignore_index=True).drop_duplicates()
+            else:
+                # No matches found, return empty dataframe with same columns
+                df = pd.DataFrame(columns=df.columns)
+        # Convert to CSV
+        csv_data = df.to_csv(index=False)
+        return DataResponsePayload(dataCSV=csv_data)
+    except HTTPException:
+        raise
     except Exception as e:
         logger.error(f"Error downloading data: {str(e)}")
         raise HTTPException(status_code=500, detail=f"Error downloading data: {str(e)}")

src/endpoints/data/data_upload.py

@@ -1,27 +1,103 @@
+import logging
+import uuid
+from datetime import datetime
+from typing import Any, Dict, List, Optional
+
+import numpy as np
 from fastapi import APIRouter, HTTPException
 from pydantic import BaseModel
-from typing import Dict, Any
-import logging
+
+from src.service.constants import 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 (
+    handle_ground_truths,
+    process_tensors,
+    sanitize_id,
+    save_model_data,
+    validate_data_tag,
+    validate_input_shapes,
+)
 
 router = APIRouter()
 logger = logging.getLogger(__name__)
+storage = get_storage_interface()
 
 
-class ModelInferJointPayload(BaseModel):
+class UploadPayload(BaseModel):
     model_name: str
-    data_tag: str = None
+    data_tag: Optional[str] = None
     is_ground_truth: bool = False
     request: Dict[str, Any]
     response: Dict[str, Any]
 
 
 @router.post("/data/upload")
-async def upload_data(payload: ModelInferJointPayload):
-    """Upload a batch of model data to TrustyAI."""
-    try:
-        logger.info(f"Received data upload for model: {payload.model_name}")
-        # TODO: Implement
-        return {"status": "success", "message": "Data uploaded successfully"}
-    except Exception as e:
-        logger.error(f"Error uploading data: {str(e)}")
-        raise HTTPException(status_code=500, detail=f"Error uploading data: {str(e)}")
+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)]
+
+        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 ""),
+            ]
+            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}"}

src/service/utils/download.py

@@ -0,0 +1,172 @@
+import logging
+import pickle
+from datetime import datetime
+from typing import Any, List, Optional
+
+import pandas as pd
+from fastapi import HTTPException
+from pydantic import BaseModel
+
+from src.service.data.storage import get_storage_interface
+
+logger = logging.getLogger(__name__)
+
+
+class RowMatcher(BaseModel):
+    columnName: str
+    operation: str
+    values: List[Any]
+
+
+class DataRequestPayload(BaseModel):
+    modelId: str
+    matchAny: Optional[List[RowMatcher]] = []
+    matchAll: Optional[List[RowMatcher]] = []
+    matchNone: Optional[List[RowMatcher]] = []
+
+
+class DataResponsePayload(BaseModel):
+    dataCSV: str
+
+
+def get_storage() -> Any:
+    """Get storage instance"""
+    return get_storage_interface()
+
+
+def apply_matcher(df: pd.DataFrame, matcher: RowMatcher, negate: bool = False) -> pd.DataFrame:
+    """Apply a single matcher to the dataframe."""
+    if matcher.operation not in ["EQUALS", "BETWEEN"]:
+        raise HTTPException(
+            status_code=400,
+            detail="RowMatch operation must be one of [BETWEEN, EQUALS]",
+        )
+    if matcher.operation == "EQUALS":
+        return apply_equals_matcher(df, matcher, negate)
+    elif matcher.operation == "BETWEEN":
+        return apply_between_matcher(df, matcher, negate)
+
+
+def apply_equals_matcher(df: pd.DataFrame, matcher: RowMatcher, negate: bool = False) -> pd.DataFrame:
+    """Apply EQUALS matcher to dataframe."""
+    column_name = matcher.columnName
+    values = matcher.values
+    if column_name not in df.columns:
+        raise HTTPException(
+            status_code=400,
+            detail=f"No feature or output found with name={column_name}",
+        )
+    mask = df[column_name].isin(values)
+    if negate:
+        mask = ~mask
+    return df[mask]
+
+
+def apply_between_matcher(df: pd.DataFrame, matcher: RowMatcher, negate: bool = False) -> pd.DataFrame:
+    """Apply BETWEEN matcher to dataframe."""
+    column_name = matcher.columnName
+    values = matcher.values
+
+    if column_name not in df.columns:
+        raise HTTPException(
+            status_code=400,
+            detail=f"No feature or output found with name={column_name}",
+        )
+    errors = []
+    if len(values) != 2:
+        errors.append(
+            f"BETWEEN operation must contain exactly two values, describing the lower and upper bounds of the desired range. Received {len(values)} values"
+        )
+    if column_name == "trustyai.TIMESTAMP":
+        if errors:
+            combined_error = ", ".join(errors)
+            raise HTTPException(status_code=400, detail=combined_error)
+        try:
+            start_time = pd.to_datetime(str(values[0]))
+            end_time = pd.to_datetime(str(values[1]))
+            df_times = pd.to_datetime(df[column_name])
+            mask = (df_times >= start_time) & (df_times < end_time)
+        except Exception as e:
+            raise HTTPException(
+                status_code=400,
+                detail=f"Timestamp value is unparseable as an ISO_LOCAL_DATE_TIME: {str(e)}",
+            )
+    elif column_name == "trustyai.INDEX":
+        if errors:
+            combined_error = ", ".join(errors)
+            raise HTTPException(status_code=400, detail=combined_error)
+        min_val, max_val = sorted([int(v) for v in values])
+        mask = (df[column_name] >= min_val) & (df[column_name] < max_val)
+    else:
+        if not all(isinstance(v, (int, float)) for v in values):
+            errors.append(
+                "BETWEEN operation must only contain numbers, describing the lower and upper bounds of the desired range. Received non-numeric values"
+            )
+        if errors:
+            combined_error = ", ".join(errors)
+            raise HTTPException(status_code=400, detail=combined_error)
+        min_val, max_val = sorted(values)
+        try:
+            mask = (pd.to_numeric(df[column_name], errors="coerce") >= min_val) & (
+                pd.to_numeric(df[column_name], errors="coerce") < max_val
+            )
+        except:
+            mask = (df[column_name].astype(str) >= str(min_val)) & (df[column_name].astype(str) < str(max_val))
+    if negate:
+        mask = ~mask
+    return df[mask]
+
+
+async def load_model_dataframe(model_id: str) -> pd.DataFrame:
+    """Load model dataframe from storage."""
+    storage = get_storage()
+    try:
+        input_data, input_cols = await storage.read_data(f"{model_id}_inputs")
+        output_data, output_cols = await storage.read_data(f"{model_id}_outputs")
+        metadata_data, metadata_cols = await storage.read_data(f"{model_id}_metadata")
+        if input_data is None or output_data is None or metadata_data is None:
+            raise HTTPException(status_code=404, detail=f"Model {model_id} not found")
+        df = pd.DataFrame()
+        if len(input_data) > 0:
+            if input_data.ndim == 2 and len(input_cols) == 1 and input_data.shape[1] > 1:
+                col_name = input_cols[0]
+                for j in range(input_data.shape[1]):
+                    df[f"{col_name}_{j}"] = input_data[:, j]
+            else:
+                input_df = pd.DataFrame(input_data, columns=input_cols)
+                for col in input_cols:
+                    df[col] = input_df[col]
+        if len(output_data) > 0:
+            if output_data.ndim == 2 and len(output_cols) == 1 and output_data.shape[1] > 1:
+                col_name = output_cols[0]
+                for j in range(output_data.shape[1]):
+                    df[f"{col_name}_{j}"] = output_data[:, j]
+            else:
+                if output_data.ndim == 2:
+                    output_data = output_data.flatten()
+                output_df = pd.DataFrame({output_cols[0]: output_data})
+                for col in output_cols:
+                    df[col] = output_df[col]
+        if len(metadata_data) > 0 and isinstance(metadata_data[0], bytes):
+            deserialized_metadata = []
+            for row in metadata_data:
+                deserialized_row = pickle.loads(row)
+                deserialized_metadata.append(deserialized_row)
+            metadata_df = pd.DataFrame(deserialized_metadata, columns=metadata_cols)
+        else:
+            metadata_df = pd.DataFrame(metadata_data, columns=metadata_cols)
+        trusty_mapping = {
+            "id": "trustyai.ID",
+            "model_id": "trustyai.MODEL_ID",
+            "timestamp": "trustyai.TIMESTAMP",
+            "tag": "trustyai.TAG",
+        }
+        for orig_col in metadata_cols:
+            trusty_col = trusty_mapping.get(orig_col.lower(), orig_col)
+            df[trusty_col] = metadata_df[orig_col]
+        df["trustyai.INDEX"] = range(len(df))
+        return df
+    except Exception as e:
+        if "not found" in str(e).lower() or "MissingH5PYDataException" in str(type(e).__name__):
+            raise HTTPException(status_code=404, detail=f"Model {model_id} not found")
+        raise HTTPException(status_code=500, detail=f"Error loading model data: {str(e)}")