GitLab mlflow Integration - Loading models does not work

Summary

Hi,

I've tried to implement a simple end-to-end process of training (tensorflow) models and using them in project. With some workarounds (like storing models only under registered versions (not in experiments), there still seems to be a problem with loading the models that makes it impossible (or I make some mistake).

Appreciate the help :)

Steps to reproduce

1. Register Model Version

from mlflow import MlflowClient

os.environ['MLFLOW_TRACKING_TOKEN'] = ''
os.environ['MLFLOW_TRACKING_URI'] = ''


client = MlflowClient()

model_name = 'MODEL_EXAMPLE'
description = 'Exemplary model'

model = client.create_registered_model(model_name, description=description)
model_version = client.create_model_version(model_name, source="", description=description)

I run it and get model into registry: image

2. Train and log simple tensorflow model under version that you just created

# Create simple regression model with tensorflow
my_model = create_tf_model()
fake_data = pd.DataFrame({'var1':[3,2,1], 'var2':[1,2,3],'target':[15,25,35]})


# Get the run_id of registered model version to log the model there
version = '1.0.0'
model_version = client.get_model_version(model_name, version)
run_id = model_version.run_id

# Train and log the model
with mlflow.start_run(run_id=run_id) as run:

    my_model.fit(fake_data[['var1', 'var2']], fake_data['target'], epochs=5)
    mlflow.tensorflow.log_model(my_model, 'tfmodel')

and the model appears to be logged: image

3. Try to load the model

Either with model version

loaded_model = mlflow.tensorflow.load_model(model_uri=f"models:/{model_name}/{version}")

returns error: image

or by the run_id

loaded_model = mlflow.tensorflow.load_model(f'runs:/{run_id}/tfmodel')

returns error: image

Tried also to download the artifacts with:

mlflow.artifacts.download_artifacts(run_id=run_id)

but also failed with the same error as above

What is the current bug behavior?

Errors listed above

What is the expected correct behavior?

I should be able to load the model

Environment info

GitLab v16.9
python 3.11.5
tensorflow 2.12.1
mlflow 2.11.3

Edited by Piotr Grzywacz