[BUG]: filter Mlflow models by stage

[whoknowsb]

Contact Details [Optional]

No response

System Information

ZENML_LOCAL_VERSION: 0.58.2.dev20240620
ZENML_SERVER_VERSION: 0.58.2.dev20240620
ZENML_SERVER_DATABASE: sqlite
ZENML_SERVER_DEPLOYMENT_TYPE: other
ZENML_CONFIG_DIR: /Users/sballoch/Library/Application Support/zenml
ZENML_LOCAL_STORE_DIR: /Users/sballoch/Library/Application Support/zenml/local_stores
ZENML_SERVER_URL: sqlite:////Users/sballoch/Library/Application Support/zenml/local_stores/default_zen_store/zenml.db
ZENML_ACTIVE_REPOSITORY_ROOT: None
PYTHON_VERSION: 3.10.14
ENVIRONMENT: native
SYSTEM_INFO: {'os': 'mac', 'mac_version': '14.4.1'}
ACTIVE_WORKSPACE: default
ACTIVE_STACK: deros_forecasting_dev_stack
ACTIVE_USER: default
TELEMETRY_STATUS: enabled
ANALYTICS_CLIENT_ID: bf2620ad-d9d2-48f9-9b81-3e0e8e492776
ANALYTICS_USER_ID: 80ceaf7b-8fa9-4fdd-abf3-5f45b0a3d265
ANALYTICS_SERVER_ID: bf2620ad-d9d2-48f9-9b81-3e0e8e492776
INTEGRATIONS: ['airflow', 'bitbucket', 'kaniko', 'pigeon', 'pillow', 'scipy', 'sklearn']
PACKAGES: {'babel': '2.15.0', 'cython': '3.0.10', 'deprecated': '1.2.14', 'gputil': '1.4.0', 'gitpython': '3.1.43', 'mako': '1.3.5', 'markdown': '3.6', 'markupsafe': '2.1.5', 'pyjwt': '2.7.0', 'pymysql':
'1.1.1', 'pyyaml': '6.0.1', 'sqlalchemy': '2.0.31', 'sqlalchemy-jsonfield': '1.0.2', 'sqlalchemy-utils': '0.41.2', 'absl-py': '2.1.0', 'aiobotocore': '2.5.4', 'aiohttp': '3.9.5', 'aioitertools': '0.11.0',
'aiosignal': '1.3.1', 'alabaster': '0.7.16', 'alembic': '1.8.1', 'aniso8601': '9.0.1', 'annotated-types': '0.7.0', 'antlr4-python3-runtime': '4.9.3', 'anyio': '4.4.0', 'asttokens': '2.4.1', 'astunparse':
'1.6.3', 'async-timeout': '4.0.3', 'attrs': '23.2.0', 'azure-common': '1.1.28', 'azure-core': '1.30.2', 'azure-mgmt-core': '1.4.0', 'azure-mgmt-resource': '23.1.1', 'backports.tarfile': '1.2.0', 'bcrypt':
'4.0.1', 'black': '23.3.0', 'blinker': '1.8.2', 'boto3': '1.28.17', 'botocore': '1.31.17', 'build': '1.2.1', 'cachetools': '5.3.3', 'certifi': '2024.6.2', 'cffi': '1.16.0', 'cfgv': '3.4.0',
'charset-normalizer': '2.1.1', 'clarabel': '0.9.0', 'click': '8.1.3', 'click-params': '0.3.0', 'cloudpickle': '2.2.1', 'comm': '0.2.2', 'contourpy': '1.2.1', 'coverage': '7.5.3', 'cryptography': '42.0.8',
'cvxopt': '1.3.2', 'cvxpy': '1.5.2', 'cycler': '0.12.1', 'decorator': '5.1.1', 'deros-base-strategy': '0.0.6', 'deros-clients': '0.3.2', 'deros-cloud-client-core': '2.0.0', 'deros-forecasting-pipeline':
'0.0.12', 'deros-ms-config-client': '1.0.0', 'deros-stats-client': '2.0.0', 'deros-topology-client': '2.0.1', 'deros-weather-service-client': '0.0.1', 'dill': '0.3.8', 'distlib': '0.3.8', 'distro': '1.9.0',
'docformatter': '1.5.1', 'docker': '6.1.3', 'docstr-coverage': '2.2.0', 'docutils': '0.19', 'ecos': '2.0.14', 'enelx-logging': '0.1.7', 'entrypoints': '0.4', 'exceptiongroup': '1.2.1', 'execnet': '2.1.1',
'executing': '2.0.1', 'fancyimpute': '0.7.0', 'fastapi': '0.110.0', 'fastapi-utils': '0.2.1', 'filelock': '3.15.3', 'flask': '3.0.3', 'flatbuffers': '24.3.25', 'fonttools': '4.53.0',
'forecasting-premo-strategy': '0.0.1', 'frozenlist': '1.4.1', 'fsspec': '2023.4.0', 'gast': '0.5.4', 'gitdb': '4.0.11', 'google-pasta': '0.2.0', 'graphene': '3.3', 'graphql-core': '3.2.3', 'graphql-relay':
'3.2.0', 'grpcio': '1.64.1', 'gunicorn': '22.0.0', 'h11': '0.14.0', 'h5py': '3.11.0', 'holidays': '0.51', 'httplib2': '0.19.1', 'httptools': '0.6.1', 'hydra-core': '1.3.2', 'identify': '2.5.36', 'idna':
'3.7', 'imagesize': '1.4.1', 'immutabledict': '3.0.0', 'importlib-metadata': '7.1.0', 'iniconfig': '2.0.0', 'ipinfo': '5.0.1', 'ipython': '8.25.0', 'ipython-genutils': '0.2.0', 'ipywidgets': '8.1.3',
'isodate': '0.6.1', 'itsdangerous': '2.2.0', 'jaraco.classes': '3.4.0', 'jaraco.context': '5.3.0', 'jaraco.functools': '4.0.1', 'jedi': '0.19.1', 'jinja2': '3.1.4', 'jmespath': '1.0.1', 'joblib': '1.4.2',
'jsonschema': '4.22.0', 'jsonschema-specifications': '2023.12.1', 'jupyter-core': '5.7.2', 'jupyterlab-widgets': '3.0.11', 'keras': '3.3.3', 'keyring': '25.2.1', 'kiwisolver': '1.4.5', 'knnimpute': '0.1.0',
'libclang': '18.1.1', 'm2r2': '0.3.3.post2', 'markdown-it-py': '3.0.0', 'marshmallow': '3.21.3', 'matplotlib': '3.9.0', 'matplotlib-inline': '0.1.7', 'mdurl': '0.1.2', 'mistune': '0.8.4', 'ml-dtypes':
'0.3.2', 'mlflow': '2.14.1', 'more-itertools': '10.3.0', 'multidict': '6.0.5', 'mypy': '1.2.0', 'mypy-extensions': '1.0.0', 'namex': '0.0.8', 'nbformat': '4.4.0', 'nh3': '0.2.17', 'nodeenv': '1.9.1',
'nose': '1.3.7', 'numpy': '1.26.4', 'numpydoc': '1.5.0', 'omegaconf': '2.3.0', 'opentelemetry-api': '1.25.0', 'opentelemetry-sdk': '1.25.0', 'opentelemetry-semantic-conventions': '0.46b0', 'opt-einsum':
'3.3.0', 'optree': '0.11.0', 'orjson': '3.10.5', 'osqp': '0.6.7', 'overrides': '7.7.0', 'packaging': '24.1', 'pandas': '2.2.2', 'parso': '0.8.4', 'passlib': '1.7.4', 'pathspec': '0.12.1', 'patsy': '0.5.6',
'pexpect': '4.9.0', 'pillow': '10.3.0', 'pip': '24.0', 'pkginfo': '1.11.1', 'platformdirs': '4.2.2', 'plotly': '5.22.0', 'pluggy': '1.5.0', 'pmdarima': '2.0.4', 'pre-commit': '3.3.1', 'premo': '1.0.0',
'prompt-toolkit': '3.0.47', 'protobuf': '4.25.3', 'psutil': '6.0.0', 'ptyprocess': '0.7.0', 'pure-eval': '0.2.2', 'pyarrow': '15.0.2', 'pycparser': '2.22', 'pydantic': '2.7.4', 'pydantic-core': '2.18.4',
'pydantic-settings': '2.2.1', 'pygments': '2.18.0', 'pyparsing': '2.4.7', 'pyproject-hooks': '1.1.0', 'pytest': '7.3.1', 'pytest-cov': '4.0.0', 'pytest-mock': '3.14.0', 'pytest-sugar': '0.9.7',
'pytest-xdist': '3.2.1', 'python-dateutil': '2.9.0.post0', 'python-dotenv': '1.0.1', 'python-multipart': '0.0.9', 'pytz': '2024.1', 'qdldl': '0.1.7.post3', 'querystring-parser': '1.2.4', 'readme-renderer':
'43.0', 'referencing': '0.35.1', 'requests': '2.31.0', 'requests-toolbelt': '1.0.0', 'responses': '0.24.1', 'returns': '0.20.1', 'rfc3986': '2.0.0', 'rich': '13.7.1', 'rpds-py': '0.18.1', 'rstr': '3.2.2',
'ruff': '0.0.264', 's3fs': '2023.4.0', 's3transfer': '0.6.2', 'scikit-learn': '1.5.0', 'scipy': '1.13.1', 'scs': '3.2.4.post3', 'secure': '0.3.0', 'setuptools': '67.8.0', 'six': '1.16.0', 'smmap': '5.0.1',
'sniffio': '1.3.1', 'snowballstemmer': '2.2.0', 'sphinx': '6.2.1', 'sphinx-autodoc-typehints': '1.23.0', 'sphinx-rtd-theme': '0.5.1', 'sphinxcontrib-applehelp': '1.0.8', 'sphinxcontrib-confluencebuilder':
'2.0.0', 'sphinxcontrib-devhelp': '1.0.6', 'sphinxcontrib-htmlhelp': '2.0.5', 'sphinxcontrib-jsmath': '1.0.1', 'sphinxcontrib-qthelp': '1.0.7', 'sphinxcontrib-serializinghtml': '1.1.10',
'spring-config-client': '0.2', 'sqlalchemy2-stubs': '0.0.2a38', 'sqlmodel': '0.0.18', 'sqlparse': '0.5.0', 'stack-data': '0.6.3', 'starlette': '0.36.3', 'statsmodels': '0.14.2', 'tailer': '0.4.1',
'tenacity': '8.4.1', 'tensorboard': '2.16.2', 'tensorboard-data-server': '0.7.2', 'tensorflow': '2.16.1', 'tensorflow-io-gcs-filesystem': '0.37.0', 'termcolor': '2.4.0', 'threadpoolctl': '3.5.0',
'tokenize-rt': '5.2.0', 'tomli': '2.0.1', 'tqdm': '4.66.4', 'traitlets': '5.14.3', 'twine': '4.0.2', 'types-requests': '2.31.0.6', 'types-urllib3': '1.26.25.14', 'typing-extensions': '4.12.2', 'tzdata':
'2024.1', 'ufcore': '0.3.1', 'untokenize': '0.1.1', 'urllib3': '1.26.19', 'uvicorn': '0.30.1', 'uvloop': '0.19.0', 'validators': '0.18.2', 'vcrpy': '5.1.0', 'virtualenv': '20.26.2', 'watchfiles': '0.22.0',
'wcwidth': '0.2.13', 'websocket-client': '1.8.0', 'websockets': '12.0', 'werkzeug': '3.0.3', 'wheel': '0.43.0', 'widgetsnbextension': '4.0.11', 'wrapt': '1.16.0', 'yarl': '1.9.4', 'zenml-nightly':
'0.58.2.dev20240620', 'zipp': '3.19.2'}

CURRENT STACK

Name: deros_forecasting_dev_stack
ID: 7f9c7762-9538-412f-8700-f4a565e695c7
User: default / 80ceaf7b-8fa9-4fdd-abf3-5f45b0a3d265
Workspace: default / 52d9629d-af58-4921-964e-161ad286c0f9

ORCHESTRATOR: default

Name: default
ID: 9561ed2b-9506-4e6f-9aeb-82a692327988
Type: orchestrator
Flavor: local
Configuration: {}
Workspace: default / 52d9629d-af58-4921-964e-161ad286c0f9

ARTIFACT_STORE: default

Name: default
ID: bd8263b8-27f5-4ef7-a847-75530b58f73f
Type: artifact_store
Flavor: local
Configuration: {'path': ''}
Workspace: default / 52d9629d-af58-4921-964e-161ad286c0f9

EXPERIMENT_TRACKER: mlflow_experiment_tracker

Name: mlflow_experiment_tracker
ID: 56214e94-602a-4258-9145-dc6b6200607c
Type: experiment_tracker
Flavor: mlflow
Configuration: {'experiment_name': None, 'nested': False, 'tags': {}, 'tracking_uri': None, 'tracking_username': '', 'tracking_password': '', 'tracking_token': '********',
'tracking_insecure_tls': False, 'databricks_host': None}
User: default / 80ceaf7b-8fa9-4fdd-abf3-5f45b0a3d265
Workspace: default / 52d9629d-af58-4921-964e-161ad286c0f9

MODEL_REGISTRY: mlflow_model_registry

Name: mlflow_model_registry
ID: 88bd3589-2db0-426c-8e96-e724728734ed
Type: model_registry
Flavor: mlflow
Configuration: {}
User: default / 80ceaf7b-8fa9-4fdd-abf3-5f45b0a3d265
Workspace: default / 52d9629d-af58-4921-964e-161ad286c0f9
deros-forecasting-pipeline-py3.10sballoch@Riccardos-MacBook-Pro-2 forecasting-pipeline %

What happened?

filtering by stages on Mlflow does not work. always evaluate

mlflow_model_version.current_stage == str(stage)
False

https://github.com/zenml-io/zenml/blob/51192ccdf1cc9693e43d173d82f2a51bc6dbe461/src/zenml/integrations/mlflow/model_registries/mlflow_model_registry.py#L618C13-L618C79

Fix ->

ModelVersionStage(mlflow_model_version.current_stage) == stage
True

Fix proposed in: #2798

Reproduction steps

...

Relevant log output

No response

Code of Conduct

• I agree to follow this project's Code of Conduct

[whoknowsb]

Is there anyone who can review this? 😅

[htahir1]

@whoknowsb so sorry for the delay, @bcdurak will be on this ASAP

[htahir1]

Fixed by https://github.com/zenml-io/zenml/pull/2798/files