Merge pull request #25 from RobGeada/DB

feat: Add MariaDB storage_interface

Author: RobGeada

.github/workflows/python-tests.yaml

@@ -22,34 +22,45 @@ jobs:
         with:
           python-version: ${{ matrix.python-version }}
           cache: "pip"
-
+      - name: Install MariaDB and MariaDB Connector
+        run: |
+            sudo apt-get install
+            sudo apt-get install -y libmariadb3 libmariadb-dev mariadb-server
       - name: Install dependencies
         run: |
           python -m pip install --upgrade pip
           python -m pip install uv
           python -m venv .venv
           source .venv/bin/activate
           if [ -f pyproject.toml ]; then
-            uv pip install -e ".[dev,protobuf]"
+            uv pip install -e ".[dev,protobuf,mariadb]"
           else
             uv pip install pytest
           fi
 
       - name: Install protobuf compiler
         run: |
-          sudo apt-get update
           sudo apt-get install -y protobuf-compiler
-
       - name: Generate protobuf stubs
         run: |
           source .venv/bin/activate
           bash scripts/generate_protos.sh
-
+      - name: Shutdown Ubuntu MySQL
+        run: sudo service mysql stop
+      - name: Set up MariaDB
+        uses: getong/[email protected]
+        with:
+          mysql user: 'trustyai'
+          mysql password: 'trustyai'
+          mysql root password: ''
+          mysql database: 'trustyai-database'
+      - name: Populate Database
+        run: |
+          mariadb -u trustyai --password=trustyai -D trustyai-database -h 127.0.0.1 < tests/resources/legacy_database_dump.sql
       - name: Run tests with pytest
         run: |
           source .venv/bin/activate
           pytest tests/ -v --cov=src --cov-report=xml
-
       - name: Upload coverage to Codecov
         uses: codecov/codecov-action@v4
         with:

Dockerfile

@@ -10,10 +10,18 @@ WORKDIR /app
 COPY pyproject.toml poetry.lock* README.md ./
 
 
-
 USER root
+
+# install mariadb connector from MariaDB Community Service package repository
+RUN if [[ "$EXTRAS" == *"mariadb"* ]]; then  \
+		curl -LsSO https://r.mariadb.com/downloads/mariadb_repo_setup && \
+    	echo "c4a0f3dade02c51a6a28ca3609a13d7a0f8910cccbb90935a2f218454d3a914a  mariadb_repo_setup" | sha256sum -c - && \
+    	chmod +x mariadb_repo_setup && \
+    	./mariadb_repo_setup --mariadb-server-version="mariadb-10.6" && \
+    	dnf install -y MariaDB-shared MariaDB-devel;  \
+    fi
 RUN pip install uv==0.6.16 && \
-    uv pip install .[$EXTRAS]
+    uv pip install ".[$EXTRAS]"
 COPY . .
 USER 1001
 EXPOSE 4443

README.md

@@ -32,31 +32,27 @@ environment, a Jupyter Notebook, or in Kubernetes.
 
 ---
 ## 📦 Building 📦
-### Locally (Without Eval Support)
-```bash
-uv pip install .
-````
-
-### Locally (With Eval Support)
+### Locally
 ```bash
-uv pip install .[eval]
-````
+uv pip install ".[$EXTRAS]"
+```
 
-### Container (Without Eval Support)
+### Container
 ```bash
-podman build -t $IMAGE_NAME .
-````
+podman build -t $IMAGE_NAME --build-arg EXTRAS="$EXTRAS" .
+```
 
-### Container (With Eval Support)
-```bash
-podman build -t $IMAGE_NAME --build-arg EXTRAS=eval .
-````
+### Available Extras
+Pass these extras as a comma separated list, e.g., `"mariadb,protobuf"`
+* `protobuf`: To process model inference data from ModelMesh models, you can install with `protobuf` support. Otherwise, only KServe models will be supported.
+* `eval`: To enable the Language Model Evaluation servers, install with `eval` support.
+* `mariadb` (If installing locally, install the [MariaDB Connector/C](https://mariadb.com/docs/server/connect/programming-languages/c/install/) first.)
 
-### Locally (With ModelMesh/Protobuf Support)
+### Examples
 ```bash
-uv pip install .[protobuf]
-````
-
+uv pip install ".[mariadb,protobuf,eval]"
+podman build -t $IMAGE_NAME --build-arg EXTRAS="mariadb,protobuf,eval" .
+```
 
 ## 🏃Running 🏃‍♀️
 ### Locally
@@ -91,12 +87,6 @@ python -m pytest --cov=src
 ## 🔄 Protobuf Support 🔄
 To process model inference data from ModelMesh models, you can install protobuf support. Otherwise, only KServe models will be supported.
 
-### Installing Dependencies
-Install the required dependencies for protobuf support:
-```bash
-uv pip install -e ".[protobuf]"
-```
-
 ### Generating Protobuf Code
 After installing dependencies, generate Python code from the protobuf definitions:
 

pyproject.toml

@@ -28,6 +28,8 @@ dev = [
 ]
 eval = ["lm-eval[api]==0.4.4", "fastapi-utils>=0.8.0", "typing-inspect==0.9.0"]
 protobuf = ["numpy>=1.24.0,<3", "grpcio>=1.62.1,<2", "grpcio-tools>=1.62.1,<2"]
+mariadb = ["mariadb>=1.1.12", "javaobj-py3==0.4.4"]
+
 
 [tool.hatch.build.targets.sdist]
 include = ["src"]

src/endpoints/consumer/consumer_endpoint.py

@@ -27,7 +27,7 @@
 logger = logging.getLogger(__name__)
 
 PartialKind = Literal["request", "response"]
-storage_inferface = get_storage_interface()
+storage_interface = get_storage_interface()
 unreconciled_inputs = {}
 unreconciled_outputs = {}
 
@@ -160,11 +160,11 @@ async def consume_inference_payload(
                 ) from e
 
             # Store the input payload
-            await storage_inferface.persist_modelmesh_payload(
+            await storage_interface.persist_modelmesh_payload(
                 partial_payload, payload_id, is_input
             )
 
-            output_payload = await storage_inferface.get_modelmesh_payload(
+            output_payload = await storage_interface.get_modelmesh_payload(
                 payload_id, False
             )
 
@@ -188,11 +188,11 @@ async def consume_inference_payload(
                 ) from e
 
             # Store the output payload
-            await storage_inferface.persist_modelmesh_payload(
+            await storage_interface.persist_modelmesh_payload(
                 partial_payload, payload_id, is_input
             )
 
-            input_payload = await storage_inferface.get_modelmesh_payload(
+            input_payload = await storage_interface.get_modelmesh_payload(
                 payload_id, True
             )
 
@@ -253,11 +253,11 @@ async def reconcile_modelmesh_payloads(
     metadata_cols = ["iso_time", "unix_timestamp", "tags"]
 
     await asyncio.gather(
-        storage_inferface.write_data(input_dataset, df[input_cols].values, input_cols),
-        storage_inferface.write_data(
+        storage_interface.write_data(input_dataset, df[input_cols].values, input_cols),
+        storage_interface.write_data(
             output_dataset, df[output_cols].values, output_cols
         ),
-        storage_inferface.write_data(metadata_dataset, metadata, metadata_cols),
+        storage_interface.write_data(metadata_dataset, metadata, metadata_cols),
     )
 
     shapes = await ModelData(model_id).shapes()
@@ -273,8 +273,8 @@ async def reconcile_modelmesh_payloads(
     )
 
     # Clean up
-    await storage_inferface.delete_modelmesh_payload(request_id, True)
-    await storage_inferface.delete_modelmesh_payload(request_id, False)
+    await storage_interface.delete_modelmesh_payload(request_id, True)
+    await storage_interface.delete_modelmesh_payload(request_id, False)
 
 
 def reconcile_mismatching_shape_error(shape_tuples, payload_type, payload_id):
@@ -377,9 +377,9 @@ async def reconcile(
     metadata_dataset = output_payload.model_name + METADATA_SUFFIX
 
     await asyncio.gather(
-        storage_inferface.write_data(input_dataset, input_array, input_names),
-        storage_inferface.write_data(output_dataset, output_array, output_names),
-        storage_inferface.write_data(metadata_dataset, metadata, metadata_names),
+        storage_interface.write_data(input_dataset, input_array, input_names),
+        storage_interface.write_data(output_dataset, output_array, output_names),
+        storage_interface.write_data(metadata_dataset, metadata, metadata_names),
     )
 
     shapes = await ModelData(output_payload.model_name).shapes()
@@ -411,13 +411,13 @@ async def consume_cloud_event(
         else:
             logger.info(f"KServe Inference Input {payload.id} received.")
             # if a match is found, the payload is auto-deleted from data
-            partial_output = await storage_inferface.get_partial_payload(
+            partial_output = await storage_interface.get_partial_payload(
                 payload.id, is_input=False
             )
             if partial_output is not None:
                 await reconcile(payload, partial_output)
             else:
-                await storage_inferface.persist_partial_payload(payload, is_input=True)
+                await storage_interface.persist_partial_payload(payload, is_input=True)
             return {
                 "status": "success",
                 "message": f"Input payload {payload.id} processed successfully",
@@ -435,13 +435,13 @@ async def consume_cloud_event(
             logger.info(
                 f"KServe Inference Output {payload.id} received from model={payload.model_name}."
             )
-            partial_input = await storage_inferface.get_partial_payload(
+            partial_input = await storage_interface.get_partial_payload(
                 payload.id, is_input=True
             )
             if partial_input is not None:
                 await reconcile(partial_input, payload)
             else:
-                await storage_inferface.persist_partial_payload(payload, is_input=False)
+                await storage_interface.persist_partial_payload(payload, is_input=False)
 
         return {
             "status": "success",

src/service/data/storage/__init__.py

@@ -1,9 +1,21 @@
 import os
+
+from src.service.data.storage.maria.legacy_maria_reader import LegacyMariaDBStorageReader
+from src.service.data.storage.maria.maria import MariaDBStorage
 from src.service.data.storage.pvc import PVCStorage
 
 def get_storage_interface():
     storage_format = os.environ.get("SERVICE_STORAGE_FORMAT", "PVC")
     if storage_format == "PVC":
         return PVCStorage(data_directory=os.environ.get("STORAGE_DATA_FOLDER", "/tmp"), data_file=os.environ.get("STORAGE_DATA_FILENAME", "trustyai.hdf5"))
+    elif storage_format == "MARIA":
+        return MariaDBStorage(
+            user=os.environ.get("DATABASE_USERNAME"),
+            password=os.environ.get("DATABASE_PASSWORD"),
+            host=os.environ.get("DATABASE_HOST"),
+            port=int(os.environ.get("DATABASE_PORT")),
+            database=os.environ.get("DATABASE_DATABASE"),
+            attempt_migration=True
+        )
     else:
         raise ValueError(f"Storage format={storage_format} not yet supported by the Python implementation of the service.")