Test notebooks (#31)

* Cleanup unnecessary env when bringing up mlflow.

* USE_JUPYTER setting and startup script.

* Refactor MLflow startup scripts to conditionally disable Jupyter support

* Refactor GitHub Actions workflow for Jupyter testing and add notebook execution script

* Update MLflow startup script to support detached mode and improve SSH handling

* f

* Fix Docker command in Jupyter test setup and update MLflow service start command

* Only build mlflow when no jupyter.

* f

Author: superdosh

.env

@@ -17,7 +17,7 @@ JUPYTER_TOKEN=changeme
 GIT_PYTHON_REFRESH=quiet
 
 # container uri for mlflow -- adjust this if you have a remote tracking server
-MLFLOW_TRACKING_URI=http://mlflow:8080
+# MLFLOW_TRACKING_URI=
 
 MLFLOW_ARTIFACT_DESTINATION=./mlruns
 # To use cloud storage for artifacts, uncomment below and provide the necessary locations for credentials.

.github/workflows/tests.yml

@@ -7,7 +7,7 @@ on:
   pull_request:
 
 jobs:
-  integration-test:
+  cli-test:
     runs-on: ubuntu-latest
 
     steps:
@@ -19,16 +19,9 @@ jobs:
         with:
           python-version: "3.12"
 
-      - name: Start MLflow server
-        run: |
-          MLFLOW_TRACKING_URI=http://localhost:8080 docker compose up -d mlflow
-
-      - name: Wait for MLflow server to be ready
+      - name: Start MLflow server (no jupyter)
         run: |
-          until curl -f http://localhost:8080/health; do
-            echo "Waiting for MLflow server..."
-            sleep 5
-          done
+          ./start_services.sh no-jupyter -d
 
       - name: Install poetry
         run: pipx install "poetry == 1.8.5"
@@ -41,3 +34,31 @@ jobs:
 
       - name: Stop MLflow server
         run: docker compose down
+
+  jupyter-test:
+    runs-on: ubuntu-latest
+
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v3
+
+      - name: Set up Python
+        uses: actions/setup-python@v4
+        with:
+          python-version: "3.12"
+
+      - name: Start MLflow server
+        run: |
+          ./start_services.sh -d
+
+      - name: Copy test script to Jupyter container
+        run: |
+          docker cp tests/notebooks/run.py modelplane-jupyter-1:/app/test_notebooks.py
+
+      - name: Execute all notebooks
+        run: |
+          docker exec modelplane-jupyter-1 poetry run python /app/test_notebooks.py
+
+          
+      - name: Stop MLflow server
+        run: docker compose down

docker-compose.yaml

@@ -22,7 +22,6 @@ services:
       POSTGRES_HOST: ${POSTGRES_HOST}
       POSTGRES_PORT: ${POSTGRES_PORT}
       MLFLOW_TRACKING_URI: ${MLFLOW_TRACKING_URI}
-      # MLFLOW_ARTIFACT_TRACKING_URI: ${MLFLOW_ARTIFACT_TRACKING_URI}
       MLFLOW_ARTIFACT_DESTINATION: ${MLFLOW_ARTIFACT_DESTINATION}
       # if not provided via volume below, GS will not work as artifact store
       GOOGLE_APPLICATION_CREDENTIALS: /creds/gcp-key.json
@@ -40,6 +39,11 @@ services:
         --port 8080
     ports:
       - "8080:8080"
+    healthcheck:
+      test: ["CMD", "curl", "-f", "http://localhost:8080/health"]
+      interval: 10s
+      timeout: 5s
+      retries: 5
     volumes:
       # Volume only needed for local storage of artifacts
       - ./mlruns:/mlruns

start_services.sh

@@ -0,0 +1,37 @@
+#!/bin/bash
+
+# load .env
+set -a
+source .env
+set +a
+
+# Check if SSH agent is running
+if ssh-add -l >/dev/null 2>&1; then
+  SSH_FLAG="--ssh default"
+else
+  echo "SSH agent not running. Proceeding without SSH."
+  SSH_FLAG=""
+fi
+
+# Default values
+USE_JUPYTER=true
+DETACHED=""
+
+# Parse arguments
+for arg in "$@"; do
+  case $arg in
+    no-jupyter)
+      USE_JUPYTER=false
+      ;;
+    -d)
+      DETACHED="-d"
+      ;;
+  esac
+done
+
+# Start services based on the options
+if [ "$USE_JUPYTER" = "true" ]; then
+  docker compose down && docker compose build $SSH_FLAG && MLFLOW_TRACKING_URI="http://mlflow:8080" docker compose up $DETACHED
+else
+  docker compose down && docker compose build $SSH_FLAG mlflow && MLFLOW_TRACKING_URI="http://localhost:8080" docker compose up $DETACHED mlflow
+fi

tests/notebooks/run.py

@@ -0,0 +1,31 @@
+# This is tested via GitHub Actions (within the Jupyter container), but we
+# don't want pytest to run it directly, hence the strange name.
+
+import os
+import nbformat
+from nbconvert.preprocessors import ExecutePreprocessor
+
+
+def test_notebooks(notebooks_dir):
+    notebooks = [f for f in os.listdir(notebooks_dir) if f.endswith(".ipynb")]
+    if not notebooks:
+        print(f"No notebooks found in {notebooks_dir}")
+        return
+
+    for notebook in notebooks:
+        notebook_path = os.path.join(notebooks_dir, notebook)
+        print(f"Testing notebook: {notebook_path}")
+        with open(notebook_path) as f:
+            nb = nbformat.read(f, as_version=4)
+            ep = ExecutePreprocessor(timeout=600, kernel_name="python3")
+            try:
+                ep.preprocess(nb, {"metadata": {"path": notebooks_dir}})
+                print(f"Notebook {notebook} executed successfully.")
+            except Exception as e:
+                print(f"Error executing notebook {notebook}: {e}")
+                raise
+
+
+if __name__ == "__main__":
+    notebooks_dir = "/app/flightpaths"
+    test_notebooks(notebooks_dir)