Add automated ML train/save/load test scripts (#404)

Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Edoardo Zoni <59625522+EZoni@users.noreply.github.com>

Author: 4 people

dashboard/app.py

@@ -6,7 +6,7 @@
 from trame.ui.vuetify3 import SinglePageWithDrawerLayout
 from trame.widgets import plotly, router, vuetify3 as vuetify, html
 
-from model_manager import ModelManager
+from model_manager import ModelManager, model_type_tag_dict
 from outputs_manager import OutputManager
 from optimization_manager import OptimizationManager
 from parameters_manager import ParametersManager
@@ -16,6 +16,7 @@
 from error_manager import error_panel, add_error
 from utils import (
     data_depth_panel,
+    load_config_dict,
     load_experiments,
     load_database,
     load_data,
@@ -64,14 +65,18 @@ def update(
         state.experiment
     )
     # load data
-    db = load_database(state.experiment)
-    exp_data, sim_data = load_data(db)
+    config_dict = load_config_dict(state.experiment)
+    db = load_database(config_dict)
+    exp_data, sim_data = load_data(db, state.experiment, state.experiment_date_range)
     # reset output
     if reset_output:
         out_manager = OutputManager(output_variables)
     # reset model
     if reset_model:
-        mod_manager = ModelManager()
+        mod_manager = ModelManager(
+            config_dict=config_dict,
+            model_type_tag=model_type_tag_dict[state.model_type],
+        )
         opt_manager = OptimizationManager(mod_manager)
     # reset parameters
     if reset_parameters:
@@ -257,7 +262,8 @@ def find_simulation(event, db):
 
 
 def open_simulation_dialog(event):
-    db = load_database(state.experiment)
+    config_dict = load_config_dict(state.experiment)
+    db = load_database(config_dict)
     try:
         data_directory, file_path = find_simulation(event, db)
         state.simulation_video = file_path.endswith(".mp4")

dashboard/error_manager.py

@@ -1,8 +1,13 @@
 from trame.widgets import vuetify3 as vuetify, html
-from state_manager import state
+from state_manager import state, server
 
 
 def add_error(title, msg):
+    if not server.running:
+        # Outside of a Trame app (e.g. check_model.py), raise a Python error
+        # to surface the error to the caller.
+        raise RuntimeError(f"{title}: {msg}")
+    # Otherwise: Inside a Trame app, add the error to the state.
     state.errors.append(
         {
             "id": state.error_counter,

dashboard/model_manager.py

@@ -65,29 +65,17 @@ def patched(host_creds, endpoint, method, *args, **kwargs):
 
 
 class ModelManager:
-    def __init__(self):
+    def __init__(self, config_dict, model_type_tag):
         print("Initializing model manager...")
-        # Set initial default values
         self.__model = None
         self.__is_neural_network = False
         self.__is_gaussian_process = False
         self.__is_neural_network_ensemble = False
-
-        model_type_tag = model_type_tag_dict[state.model_type]
-        try:
-            config_dict = load_config_dict(state.experiment)
-        except Exception as e:
-            title = "Unable to load experiment configuration"
-            msg = (
-                f"Error occurred when loading configuration for {state.experiment}: {e}"
-            )
-            add_error(title, msg)
-            print(msg)
-            return
+        self.__model_type_tag = model_type_tag
 
         if "mlflow" not in config_dict or not config_dict["mlflow"].get("tracking_uri"):
             print(
-                f"No mlflow.tracking_uri in configuration file for {state.experiment}; cannot load model from MLflow."
+                f"No mlflow.tracking_uri in configuration file for {config_dict['experiment']}; cannot load model from MLflow."
             )
             return
 
@@ -99,7 +87,9 @@ def __init__(self):
             == "https://mlflow.american-science-cloud.org"
         ):
             enable_amsc_x_api_key(config_dict)
-        model_name = f"{state.experiment}_{model_type_tag}"
+
+        experiment = config_dict["experiment"]
+        model_name = f"{experiment}_{model_type_tag}"
 
         try:
             # Download model from MLflow server
@@ -108,16 +98,16 @@ def __init__(self):
                 .unwrap_python_model()
                 .model
             )
-            if state.model_type == "Neural Network (single)":
+            if model_type_tag == "NN":
                 self.__is_neural_network = True
-            elif state.model_type == "Neural Network (ensemble)":
+            elif model_type_tag == "ensemble_NN":
                 self.__is_neural_network_ensemble = True
-            elif state.model_type == "Gaussian Process":
+            elif model_type_tag == "GP":
                 self.__is_gaussian_process = True
             else:
-                raise ValueError(f"Unsupported model type: {state.model_type}")
+                raise ValueError(f"Unsupported model type: {model_type_tag}")
         except Exception as e:
-            title = f"Unable to load model {state.model_type}"
+            title = f"Unable to load model {model_type_tag}"
             msg = f"Error occurred when loading model from MLflow: {e}"
             add_error(title, msg)
             print(msg)
@@ -156,7 +146,7 @@ def evaluate(self, parameters, output):
                 std_dev = output_dict[output].variance.sqrt().detach()
                 mean_error = 2.0 * std_dev
             else:
-                raise ValueError(f"Unsupported model type: {state.model_type}")
+                raise ValueError(f"Unsupported model type: {self.__model_type_tag}")
             # compute lower/upper bounds for error range
             lower = mean - mean_error
             upper = mean + mean_error

dashboard/utils.py

@@ -103,15 +103,13 @@ def create_date_filter(experiment_date_range):
 
 
 @timer
-def load_data(db):
+def load_data(db, experiment, date_range=None):
     print("Loading data from database...")
     # create date filter if date range is set
-    date_filter = create_date_filter(state.experiment_date_range)
+    date_filter = create_date_filter(date_range)
     # load experiment and simulation data points in dataframes
-    exp_data = pd.DataFrame(
-        db[state.experiment].find({"experiment_flag": 1, **date_filter})
-    )
-    sim_data = pd.DataFrame(db[state.experiment].find({"experiment_flag": 0}))
+    exp_data = pd.DataFrame(db[experiment].find({"experiment_flag": 1, **date_filter}))
+    sim_data = pd.DataFrame(db[experiment].find({"experiment_flag": 0}))
     # Store '_id', 'date' as string
     for key in ["_id", "date"]:
         if key in exp_data.columns:
@@ -122,10 +120,8 @@ def load_data(db):
 
 
 @timer
-def load_database(experiment):
+def load_database(config_dict):
     print("Loading database...")
-    # load configuration dictionary
-    config_dict = load_config_dict(experiment)
     # read database information from configuration dictionary
     db_host = config_dict["database"]["host"]
     db_port = config_dict["database"]["port"]

ml/README.md

@@ -72,6 +72,30 @@ This section describes how to train ML models locally.
    python train_model.py --test --model <your_model> --config_file <your_config_file>
    ```
 
+### Test the full train/save/load cycle: `test_ml_pipeline.py`
+
+[`tests/test_ml_pipeline.py`](../tests/test_ml_pipeline.py) exercises the full ML lifecycle: training → upload to MLflow → download → accuracy check. It requires a local, empty MLflow server to avoid touching any production server.
+
+1. Start a local MLflow server, e.g. with Docker:
+   ```bash
+   docker run -p 127.0.0.1:5000:5000 ghcr.io/mlflow/mlflow mlflow server --host 0.0.0.0
+   ```
+
+2. Run the test script from the root of the repository (by default this expects the MLflow server to run on `localhost:5000`):
+   ```bash
+   python tests/test_ml_pipeline.py
+   ```
+
+   Optionally, restrict to a specific model type or config file:
+   ```bash
+   python tests/test_ml_pipeline.py --model NN --config_file experiments/synapse-bella-ip2
+   ```
+
+   If your MLflow server is running on a different port (e.g. 5001 instead of 5000), pass it explicitly:
+   ```bash
+   python tests/test_ml_pipeline.py --test-mlflow-uri http://localhost:5001
+   ```
+
 ## With Docker
 
 Coming soon.

tests/check_model.py

@@ -0,0 +1,115 @@
+#!/usr/bin/env python
+"""
+Check that a model stored in MLflow loads and evaluates correctly,
+using the same logic as the dashboard.
+
+Usage:
+    python check_model.py --config_file <path/to/config.yaml> --model <GP|NN|ensemble_NN>
+"""
+
+import argparse
+import os
+import sys
+import torch
+import yaml
+
+DASHBOARD_DIR = os.path.join(
+    os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "dashboard"
+)  # similar to "cd ../dashboard"
+sys.path.insert(0, DASHBOARD_DIR)
+from model_manager import ModelManager  # noqa: E402
+from utils import load_database, load_data  # noqa: E402
+
+
+MODEL_TYPES = ["GP", "NN", "ensemble_NN"]
+ACCURACY_TOLERANCE = 0.80
+
+
+def load_experimental_data(config_dict):
+    """Fetch all experimental points from the database."""
+    input_names = [v["name"] for v in config_dict["inputs"].values()]
+    output_names = [v["name"] for v in config_dict["outputs"].values()]
+
+    db = load_database(config_dict)
+    exp_data, _ = load_data(db, config_dict["experiment"])
+
+    return exp_data, input_names, output_names
+
+
+def check_evaluate(config_dict, model_type):
+    """Load model and evaluate with experimental data; verify accuracy (relative RMSE <= threshold)."""
+    # Load model
+    mm = ModelManager(config_dict=config_dict, model_type_tag=model_type)
+    # Load experimental data
+    df_exp, input_names, output_names = load_experimental_data(config_dict)
+
+    # Skip accuracy check if no experimental data available
+    if len(df_exp) == 0:
+        print(
+            f"[SKIP] No experimental data available for {config_dict['experiment']}; skipping accuracy check."
+        )
+        return
+
+    # Convert input to the format expected by the model manager
+    inputs = {n: torch.tensor(df_exp[n].values) for n in input_names}
+
+    # Check accuracy
+    all_passed = True
+    for output_name in output_names:
+        actual = torch.tensor(df_exp[output_name].values)
+        if actual.isnan().all():
+            print(
+                f"  [SKIP] Output '{output_name}': all actual values are NaN; skipping."
+            )
+            continue
+        prediction, _, _ = mm.evaluate(inputs, output_name)
+        rel_errors = (prediction - actual) / torch.max(
+            torch.abs(actual), torch.abs(prediction)
+        )
+        rmse = torch.sqrt(torch.nanmean(rel_errors**2)).item()
+        if rmse <= ACCURACY_TOLERANCE:
+            status = "PASS"
+        else:
+            status = "FAIL"
+            all_passed = False
+        print(
+            f"  [{status}] Output '{output_name}': relative RMSE = {rmse:.1%} (tolerance {ACCURACY_TOLERANCE:.0%})"
+        )
+
+    if not all_passed:
+        raise RuntimeError(
+            f"Accuracy check failed: relative RMSE exceeded {ACCURACY_TOLERANCE:.0%} for one or more outputs."
+        )
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description="Verify that an MLflow model loads and evaluates correctly."
+    )
+    parser.add_argument(
+        "--config_file",
+        help="Path to the configuration file",
+        type=str,
+        required=True,
+    )
+    parser.add_argument(
+        "--model",
+        help="Model type: GP, NN, or ensemble_NN",
+        choices=MODEL_TYPES,
+        required=True,
+    )
+    args = parser.parse_args()
+
+    # Load configuration
+    with open(args.config_file) as f:
+        config_dict = yaml.safe_load(f)
+    print(f"Experiment: {config_dict['experiment']}")
+
+    # Load model and evaluate with experimental data
+    try:
+        check_evaluate(config_dict, args.model)
+    except Exception as e:
+        print(f"[FAIL] {e}")
+        sys.exit(1)
+
+    print("[PASS] Model loaded and evaluated successfully.")