mlflow: [BUG] load_model always loads latest version from registry

Issues Policy acknowledgement

I have read and agree to submit bug reports in accordance with the issues policy

Willingness to contribute

Yes. I would be willing to contribute a fix for this bug with guidance from the MLflow community.

MLflow version

Client: 2.2.2
Tracking server: 2.2.2

System information

Rocky Linux 8.7 (Green Obsidian)
Python 3.9.13

Describe the problem

I’m currently using an mlflow server in scenario 2. I am following the documentation to fetch a registered model in a specific version / stage. When I run:

import mlflow.pyfunc
model = mlflow.pyfunc.load_model(model_uri="models:/token-classification-template_albert-base-v2/4")

I will always get the latest model, not version 4. This also applies when replacing 4 with Staging (I only have one staging model at a time), the latest model is loaded, even if it’s not the staging one.

Tracking information

System information: Linux #1 SMP Tue Sep 13 18:09:48 UTC 2022 Python version: 3.9.16 MLflow version: 2.2.2 MLflow module location: /udir/tschilla/anaconda3/envs/text_classify/lib/python3.9/site-packages/mlflow/init.py Tracking URI: http://127.0.0.1:5002 Registry URI: http://127.0.0.1:5002 MLflow environment variables: MLFLOW_TRACKING_URI: http://127.0.0.1:5002 MLflow dependencies: Flask: 2.2.3 Jinja2: 3.1.2 alembic: 1.9.4 click: 8.1.3 cloudpickle: 2.2.1 databricks-cli: 0.17.4 docker: 6.0.1 entrypoints: 0.4 gitpython: 3.1.30 gunicorn: 20.1.0 importlib-metadata: 5.2.0 markdown: 3.4.1 matplotlib: 3.7.0 numpy: 1.23.5 packaging: 22.0 pandas: 1.5.3 protobuf: 3.20.2 pyarrow: 10.0.1 pytz: 2022.7.1 pyyaml: 6.0 querystring-parser: 1.2.4 requests: 2.28.2 scikit-learn: 1.1.3 scipy: 1.9.3 shap: 0.41.0 sqlalchemy: 1.4.46 sqlparse: 0.4.3

Code to reproduce issue

Register a model which will raise an error when called.
Fetch it and call it to make sure it’s faulty.
Register a working model.
Fetch the faulty version registered in step 1 and call it -> no error is raised.

Stack trace

N/A

Other info / logs

N/A

What component(s) does this bug affect?

area/artifacts: Artifact stores and artifact logging
area/build: Build and test infrastructure for MLflow
area/docs: MLflow documentation pages
area/examples: Example code
area/model-registry: Model Registry service, APIs, and the fluent client calls for Model Registry
area/models: MLmodel format, model serialization/deserialization, flavors
area/recipes: Recipes, Recipe APIs, Recipe configs, Recipe Templates
area/projects: MLproject format, project running backends
area/scoring: MLflow Model server, model deployment tools, Spark UDFs
area/server-infra: MLflow Tracking server backend
area/tracking: Tracking Service, tracking client APIs, autologging

What interface(s) does this bug affect?

area/uiux: Front-end, user experience, plotting, JavaScript, JavaScript dev server
area/docker: Docker use across MLflow’s components, such as MLflow Projects and MLflow Models
area/sqlalchemy: Use of SQLAlchemy in the Tracking Service or Model Registry
area/windows: Windows support

What language(s) does this bug affect?

language/r: R APIs and clients
language/java: Java APIs and clients
language/new: Proposals for new client languages

What integration(s) does this bug affect?

integrations/azure: Azure and Azure ML integrations
integrations/sagemaker: SageMaker integrations
integrations/databricks: Databricks integrations

About this issue

Original URL
State: closed
Created a year ago
Comments: 24 (13 by maintainers)

Most upvoted comments

Glad we indentified the root cause, thanks for using mlflow!

harupy on Mar 24, 2023

I get it now, thank you for the answer! I’m going to have to review the way I use mlflow as I have a multitude of trainers which can be invoked from my main.py, and I need to be able to modify the class definitions.

thomas-schillaci on Mar 24, 2023

I’m expecting this class to be pickled

The class is not pickled.

# a.py

import pickle


class Foo:
    def foo(self):
        return "foo"


foo = Foo()
with open("foo.pkl", "wb") as f:
    pickle.dump(foo, f)

# b.py

import pickle


class Foo:
    def foo(self):
        return "bar"


f = "foo.pkl"
with open(f, "rb") as f:
    foo = pickle.load(f)
    print(foo.foo())

> python a.py
> python b.py
bar

harupy on Mar 24, 2023