Skip to content

Model Registry: MLflow versions compatibility

Fred de Gier requested to merge 455671-gitlab-mlflow-integration-loading-models-does-not-work-2 into master

What does this MR do and why?

A few Python methods in MLflow are not compatible with GitLab Model Registry due to a different approach in versioning.

E.g. below call will fail

mlflow.pyfunc.load_model(f"models:/{model_name}/{latest_version}", dst_path="models")

Fails:

    return max(int(x.version) for x in latest)
               ^^^^^^^^^^^^^^
ValueError: invalid literal for int() with base 10: '14.0.0'

This is resolved by returning the incremental ID instead of the semantic version. Initially this looked like a "hack" but in other calls such as calls to download artifacts, the incremental ID is also used instead of the semantic version. E.g. http://127.0.0.1:3000/api/v4/projects/7/packages/ml_models/17/files/conda.yaml where 17 is the id of the version.

Ultimately it is a breaking change, or a bug fix, as the endpoint /projects/#{project_id}/ml/mlflow/api/2.0/mlflow/registered-models/get-latest-versions" used to return the semantic version but not returns the id, which in the context of MLflow is IMHO a bug fix. But want to make sure we document and release it properly.

Furthermore, loading the latest model also works and corresponds to the actual latest model. The only consideration is that we currently allow older SemVers, which makes sense as it might be that a patch is needed for an older major version. But this will cause confusion on what is actually the "latests" version. Technically, a patch is the latest version but it's not the the highest SemVer.

model = mlflow.pyfunc.load_model(f"models:/{model_name}/latest")
model.metadata.run_id
'1b95bddd746a4da3b901b237f7008f1d'

Furthermore, versions are not exposed so the following call will return an empty list

mlflow.search_registered_models(filter_string="name='my-model2'")[0].latest_versions
[]

By exposing the version expose :versions, as: :latest_versions a user is able to retrieve the version

mlflow.search_registered_models(filter_string="name='my-model2'")[0].latest_versions[0].version
`14.0.0`

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

  1. Open GDK: http://127.0.0.1:3000/flightjs/Flight/-/ml/models
  2. Create model if one doesn't exist, give it a model name and version, click Create
  3. 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/model-versions/get?name=my-model2&version=14.0.0"↵ Output should look like the following, optionally verify in Rails console that it's the correct model_version ```↵ {"model_version":{"name":"my-model2","version":"17","creation_timestamp":1726563129,"last_updated_timestamp":1726563129,"user_id":null,"current_stage":"development","description":"","source":"http://127.0.0.1:3000/flightjs/Flight/-/ml/models/4/versions/17","run_id":"beb67dce-a0ca-4e6d-b171-c34351505aa9","status":"READY","status_message":"","tags":[],"run_link":"","aliases":[]}}-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
  1. Run the following Python script to ensure it works with MLflow
# 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)

Numbered steps to set up and validate the change are strongly suggested.

Related to #455671 (closed)

Edited by Fred de Gier

Merge request reports