Model Registry: Add ability to download artifacts in MLflow
What does this MR do and why?
This MR adds support for downloading a model with MLflow. It adds new endpoints for managing and downloading MLflow models and artifacts, allowing users to easily integrate MLflow with GitLab.
MR acceptance checklist
Please evaluate this MR against the MR acceptance checklist. It helps you analyze changes to reduce risks in quality, performance, reliability, security, and maintainability.
Screenshots or screen recordings
Screenshots are required for UI changes, and strongly recommended for all other merge requests.
| Before | After |
|---|---|
How to set up and validate locally
- Open GDK: http://127.0.0.1:3000/flightjs/Flight/-/ml/models
- Create model if one doesn't exist, give it a model name and version, click
Create - Add a filed called
conda.yamlto the version with the following content
channels:
- conda-forge
dependencies:
- python=3.12.5
- pip<=23.2.1
- pip:
- mlflow==2.16.1
- cloudpickle==3.0.0
- numpy==1.26.4
- scikit-learn==1.5.0
- scipy==1.13.1
name: mlflow-env
- Run the following cURL
curl -X GET --header "Authorization: Bearer glpat-<redacted>" "http://127.0.0.1:3000/api/v4/projects/7/ml/mlflow/api/2.0/mlflow-artifacts/artifacts/17/conda.yaml"Output should look like the following
channels:
- conda-forge
dependencies:
- python=3.12.5
- pip<=23.2.1
- pip:
- mlflow==2.16.1
- cloudpickle==3.0.0
- numpy==1.26.4
- scikit-learn==1.5.0
- scipy==1.13.1
name: mlflow-env```
4. Run the following Python script to ensure it works with MLflow
```python
# mlflow==2.16.1
import os
import mlflow.sklearn
from mlflow import MlflowClient
os.environ['MLFLOW_TRACKING_URI'] = "http://127.0.0.1:3000/api/v4/projects/7/ml/mlflow/"
os.environ['MLFLOW_TRACKING_TOKEN'] = 'glpat-<redcacted>'
mlflow.set_tracking_uri(os.environ['MLFLOW_TRACKING_URI'])
client = MlflowClient()
model_name = '<your-model-name>'
version = '<your-version>'
latest_version = mlflow.search_registered_models(filter_string=f"name='{model_name}'")[0].latest_versions[0].version
print(latest_version)
loaded_model = mlflow.pyfunc.load_model(f"models:/{model_name}/{latest_version}", dst_path="models")
# Alternatively you can load the model without specifying the version and get the latest version
loaded_model = mlflow.pyfunc.load_model(f"models:/{model_name}/latest")
# Make predictions with the loaded model
predictions = loaded_model.predict(X_test)
print(f"Predictions: {predictions}")
After running this script, you should see the files downloaded to the models directory:

Numbered steps to set up and validate the change are strongly suggested.
Related to #455671 (closed)
Edited by Fred de Gier