Merge branch 'develop' into feature/PRD-586-ochestrator-urls

zenml-io · Sep 24, 2024 · 29652c5 · 29652c5
2 parents a4647dc + 40e1990
commit 29652c5
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/zenml/integrations/mlflow/steps/mlflow_registry.py b/src/zenml/integrations/mlflow/steps/mlflow_registry.py
@@ -146,8 +146,8 @@ 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.get("mlflow_run_id", None) is None:
-        metadata.model_extra["mlflow_run_id"] = mlflow_run_id
+    if getattr(metadata, "mlflow_run_id", None) is None:
+        setattr(metadata, "mlflow_run_id", mlflow_run_id)
 
     # Register model version
     model_version = model_registry.register_model_version(