Merge branch 'main' into prometheus-publishing

Author: saichandrapandraju

.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:

.github/workflows/security-scan.yaml

@@ -0,0 +1,65 @@
+name: Security Scan
+
+on:
+  push:
+    branches: [main]
+  pull_request:
+    branches: [main]
+  workflow_dispatch:
+
+jobs:
+  bandit-scan:
+    runs-on: ubuntu-latest
+    permissions:
+      security-events: write
+      actions: read
+      contents: read
+      checks: write
+
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.11"
+          cache: "pip"
+
+      - name: Create virtual environment
+        run: |
+          python -m pip install --upgrade pip
+          python -m pip install uv
+          python -m venv .venv
+
+      - name: Install dependencies
+        run: |
+          source .venv/bin/activate
+          if [ -f pyproject.toml ]; then
+            uv pip install -e ".[dev]"
+          fi
+
+      - name: Install Bandit
+        run: |
+          source .venv/bin/activate
+          python -m pip install bandit[toml]
+
+      - name: Run Bandit Security Scan
+        uses: PyCQA/bandit-action@v1
+        with:
+          targets: "src/"
+
+      - name: Upload SARIF results to GitHub Security tab
+        if: github.ref == 'refs/heads/main'
+        uses: github/codeql-action/upload-sarif@v3
+        with:
+          sarif_file: results.sarif
+          category: bandit-security-scan
+        continue-on-error: true
+
+      - name: Upload SARIF as artifact
+        uses: actions/upload-artifact@v4
+        with:
+          name: bandit-sarif-results
+          path: results.sarif
+          retention-days: 30
+        continue-on-error: true

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

@@ -8,27 +8,29 @@ readme = "README.md"
 dependencies = [
     "fastapi>=0.115.9,<0.116",
     "pandas>=2.2.3,<3",
-    "prometheus-client>=0.21.1,<0.22",
+    "prometheus-client>=0.21.1,<0.23",
     "pydantic>=2.4.2,<3",
     "uvicorn>=0.34.0,<0.35",
     "protobuf>=4.24.4,<7",
     "requests>=2.31.0,<3",
-    "cryptography>=44.0.2,<45",
+    "cryptography>=44.0.2,<46",
     "h5py>=3.13.0,<4",
 ]
 
 [project.optional-dependencies]
 dev = [
     "pytest>=7.4.2,<9",
     "pytest-asyncio>=0.26.0,<2",
-    "isort>=5.12.0,<6",
-    "flake8>=6.1.0,<7",
+    "isort>=5.12.0,<7",
+    "flake8>=6.1.0,<8",
     "mypy>=1.5.1,<2",
-    "pytest-cov>=4.1.0,<5",
-    "httpx>=0.25.0,<0.26",
+    "pytest-cov>=4.1.0,<7",
+    "httpx>=0.25.0,<0.29",
 ]
 eval = ["lm-eval[api]==0.4.4", "fastapi-utils>=0.8.0", "typing-inspect==0.9.0"]
-protobuf = ["numpy>=1.24.0,<2", "grpcio>=1.62.1,<2", "grpcio-tools>=1.62.1,<2"]
+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.")