From 7c4dc9e3c8a8fdf963ce024936e96a8662cb3b70 Mon Sep 17 00:00:00 2001 From: Michael Schuster Date: Tue, 15 Oct 2024 17:52:00 +0200 Subject: [PATCH] Refactor function --- src/zenml/artifacts/utils.py | 11 +++-------- 1 file changed, 3 insertions(+), 8 deletions(-) diff --git a/src/zenml/artifacts/utils.py b/src/zenml/artifacts/utils.py index e6a333fddcf..2ed5ea11a45 100644 --- a/src/zenml/artifacts/utils.py +++ b/src/zenml/artifacts/utils.py @@ -793,11 +793,8 @@ def _link_artifact_version_to_the_step_and_model( ), ) error_message = "model" - model_version = ( - step_context.step_run.model_version - or step_context.pipeline_run.model_version - ) - if model_version: + + if step_context.model_version: from zenml.model.utils import ( link_artifact_version_to_model_version, ) @@ -808,11 +805,9 @@ def _link_artifact_version_to_the_step_and_model( ) link_artifact_version_to_model_version( artifact_version=artifact_version, - model_version=model_version, + model_version=step_context.model_version, artifact_config=artifact_config, ) - else: - raise RuntimeError except (RuntimeError, StepContextError): logger.debug(f"Unable to link saved artifact to {error_message}.")