Linking registered models in MLflow with the corresponding MLflow run #3020
Conversation
|
Important Review skippedAuto reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the You can disable this status message by setting the Thank you for using CodeRabbit. We offer it for free to the OSS community and would appreciate your support in helping us grow. If you find it useful, would you consider giving us a shout-out on your favorite social media? TipsChatThere are 3 ways to chat with CodeRabbit:
Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
Documentation and Community
|
| @@ -146,6 +146,11 @@ def mlflow_register_model_step( | |||
| metadata.zenml_pipeline_run_uuid = pipeline_run_uuid | |||
| if metadata.zenml_workspace is None: | |||
| metadata.zenml_workspace = zenml_workspace | |||
| if metadata.model_extra and "mlflow_run_id" in metadata.model_extra: | |||
There was a problem hiding this comment.
I'm a little confused by this change, can you explain the intended behaviour of this code?
To me it seems like:
- We first check whether the
mlflow_run_idexists in themodel_extradict - If it doesn't exist, we do nothing. Why don't we add the key in this case?
- If it does exist, we have another
ifcondition, but both theifandelsebranch do exactly the same thing: They overwrite the key with the mlflow run id.
It seems like both the if and else condition do the same thing: Overwrite mlflow_run_id in the model_extra dict, no matter if it's None or not.
And why do we need the actual check ahead if the key even exists and not just always set it?
There was a problem hiding this comment.
Yes, you are right, there is a mistake. I was probably too focused on getting rid of the mypy warning 😄
Now it should work.
I now use the get method to directly retrieve the value for the key mlflow_run_id and check for None. If mlflow_run_id does not exist, None is returned. If None is returned, the value mlflow_run_id is set.
As long as the value does not already exist in the metadata and != None, the determined run ID is added to the metadata.
aiakide
force-pushed
the
feature/mlflow-register-link-run
branch
from
c9f5588 to
c2c54d7
Compare
…l_version` The `mlflow_run_id` is required to assign a registered model to an MLflow run. The MLflow run can be traced back to a ZenML pipeline run. This is not a necessary change. However, it is more convenient to have a reference to the MLflow run for the registered models as well, as this makes it possible to view parameters or metrics, for example.
aiakide
force-pushed
the
feature/mlflow-register-link-run
branch
from
c2c54d7 to
7f5e3fb
Compare
Describe changes
To link a registered model in MLflow with the corresponding MLflow run, the parameter
mlflow_run_idmust be included in the metadata argument of themlflow_register_model_step.The MLflow run can be traced back to a ZenML pipeline run. This is not a necessary change. However, it is more convenient to have a reference to the MLflow run for the registered models as well, as this makes it possible to view parameters or metrics, for example.
I have adjusted the
mlflow_register_model_stepso that themlflow_run_idis added by default if it is not already present in the metadata.The following screenshots show the effects in the MLFlow user interface.
Pre-requisites
Please ensure you have done the following:
developand the open PR is targetingdevelop. If your branch wasn't based on develop read Contribution guide on rebasing branch to develop.Types of changes