zenml-nightly 0.80.2.dev20250415__py3-none-any.whl → 0.81.0.dev20250417__py3-none-any.whl

zenml/VERSION

@@ -1 +1 @@
-0.80.2.dev20250415
+0.81.0.dev20250417

zenml/cli/utils.py

@@ -308,13 +308,13 @@ def print_pydantic_models(
             if isinstance(model, BaseIdentifiedResponse):
                 include_columns = ["id"]
 
-                if "name" in model.model_fields:
+                if "name" in type(model).model_fields:
                     include_columns.append("name")
 
                 include_columns.extend(
                     [
                         k
-                        for k in model.get_body().model_fields.keys()
+                        for k in type(model.get_body()).model_fields.keys()
                         if k not in exclude_columns
                     ]
                 )
@@ -323,7 +323,9 @@ def print_pydantic_models(
                     include_columns.extend(
                         [
                             k
-                            for k in model.get_metadata().model_fields.keys()
+                            for k in type(
+                                model.get_metadata()
+                            ).model_fields.keys()
                             if k not in exclude_columns
                         ]
                     )
@@ -347,7 +349,7 @@ def print_pydantic_models(
             #  we want to attempt to represent them by name, if they contain
             #  such a field, else the id is used
             if isinstance(value, BaseIdentifiedResponse):
-                if "name" in value.model_fields:
+                if "name" in type(value).model_fields:
                     items[k] = str(getattr(value, "name"))
                 else:
                     items[k] = str(value.id)
@@ -357,7 +359,7 @@ def print_pydantic_models(
             elif isinstance(value, list):
                 for v in value:
                     if isinstance(v, BaseIdentifiedResponse):
-                        if "name" in v.model_fields:
+                        if "name" in type(v).model_fields:
                             items.setdefault(k, []).append(
                                 str(getattr(v, "name"))
                             )
@@ -448,13 +450,13 @@ def print_pydantic_model(
         if isinstance(model, BaseIdentifiedResponse):
             include_columns = ["id"]
 
-            if "name" in model.model_fields:
+            if "name" in type(model).model_fields:
                 include_columns.append("name")
 
             include_columns.extend(
                 [
                     k
-                    for k in model.get_body().model_fields.keys()
+                    for k in type(model.get_body()).model_fields.keys()
                     if k not in exclude_columns
                 ]
             )
@@ -463,7 +465,7 @@ def print_pydantic_model(
                 include_columns.extend(
                     [
                         k
-                        for k in model.get_metadata().model_fields.keys()
+                        for k in type(model.get_metadata()).model_fields.keys()
                         if k not in exclude_columns
                     ]
                 )
@@ -482,7 +484,7 @@ def print_pydantic_model(
     for k in include_columns:
         value = getattr(model, k)
         if isinstance(value, BaseIdentifiedResponse):
-            if "name" in value.model_fields:
+            if "name" in type(value).model_fields:
                 items[k] = str(getattr(value, "name"))
             else:
                 items[k] = str(value.id)
@@ -492,7 +494,7 @@ def print_pydantic_model(
         elif isinstance(value, list):
             for v in value:
                 if isinstance(v, BaseIdentifiedResponse):
-                    if "name" in v.model_fields:
+                    if "name" in type(v).model_fields:
                         items.setdefault(k, []).append(str(getattr(v, "name")))
                     else:
                         items.setdefault(k, []).append(str(v.id))
@@ -2138,7 +2140,7 @@ def _scrub_secret(config: StackComponentConfig) -> Dict[str, Any]:
         A configuration with secret values removed.
     """
     config_dict = {}
-    config_fields = config.__class__.model_fields
+    config_fields = type(config).model_fields
     for key, value in config_fields.items():
         if getattr(config, key):
             if secret_utils.is_secret_field(value):

zenml/config/compiler.py

@@ -210,6 +210,7 @@ class Compiler:
                 enable_artifact_metadata=config.enable_artifact_metadata,
                 enable_artifact_visualization=config.enable_artifact_visualization,
                 enable_step_logs=config.enable_step_logs,
+                enable_pipeline_logs=config.enable_pipeline_logs,
                 settings=config.settings,
                 tags=config.tags,
                 extra=config.extra,

zenml/config/global_config.py

@@ -447,7 +447,7 @@ class GlobalConfiguration(BaseModel, metaclass=GlobalConfigMetaClass):
         """
         environment_vars = {}
 
-        for key in self.model_fields.keys():
+        for key in type(self).model_fields.keys():
             if key == "store":
                 # The store configuration uses its own environment variable
                 # naming scheme

zenml/config/pipeline_configurations.py

@@ -41,6 +41,7 @@ class PipelineConfigurationUpdate(StrictBaseModel):
     enable_artifact_metadata: Optional[bool] = None
     enable_artifact_visualization: Optional[bool] = None
     enable_step_logs: Optional[bool] = None
+    enable_pipeline_logs: Optional[bool] = None
     settings: Dict[str, SerializeAsAny[BaseSettings]] = {}
     tags: Optional[List[Union[str, "Tag"]]] = None
     extra: Dict[str, Any] = {}

zenml/config/pipeline_run_configuration.py

@@ -40,6 +40,7 @@ class PipelineRunConfiguration(
     enable_artifact_metadata: Optional[bool] = None
     enable_artifact_visualization: Optional[bool] = None
     enable_step_logs: Optional[bool] = None
+    enable_pipeline_logs: Optional[bool] = None
     schedule: Optional[Schedule] = None
     build: Union[PipelineBuildBase, UUID, None] = Field(
         default=None, union_mode="left_to_right"

zenml/constants.py

@@ -168,6 +168,7 @@ ENV_ZENML_SKIP_STACK_VALIDATION = "ZENML_SKIP_STACK_VALIDATION"
 ENV_ZENML_SERVER = "ZENML_SERVER"
 ENV_ZENML_ENFORCE_TYPE_ANNOTATIONS = "ZENML_ENFORCE_TYPE_ANNOTATIONS"
 ENV_ZENML_ENABLE_IMPLICIT_AUTH_METHODS = "ZENML_ENABLE_IMPLICIT_AUTH_METHODS"
+ENV_ZENML_DISABLE_PIPELINE_LOGS_STORAGE = "ZENML_DISABLE_PIPELINE_LOGS_STORAGE"
 ENV_ZENML_DISABLE_STEP_LOGS_STORAGE = "ZENML_DISABLE_STEP_LOGS_STORAGE"
 ENV_ZENML_DISABLE_STEP_NAMES_IN_LOGS = "ZENML_DISABLE_STEP_NAMES_IN_LOGS"
 ENV_ZENML_IGNORE_FAILURE_HOOK = "ZENML_IGNORE_FAILURE_HOOK"
@@ -199,7 +200,9 @@ ENV_ZENML_RUN_SINGLE_STEPS_WITHOUT_STACK = (
 ENV_ZENML_PREVENT_CLIENT_SIDE_CACHING = "ZENML_PREVENT_CLIENT_SIDE_CACHING"
 ENV_ZENML_DISABLE_CREDENTIALS_DISK_CACHING = "DISABLE_CREDENTIALS_DISK_CACHING"
 ENV_ZENML_RUNNER_IMAGE_DISABLE_UV = "ZENML_RUNNER_IMAGE_DISABLE_UV"
-
+ENV_ZENML_WORKLOAD_TOKEN_EXPIRATION_LEEWAY = (
+    "ZENML_WORKLOAD_TOKEN_EXPIRATION_LEEWAY"
+)
 # Logging variables
 IS_DEBUG_ENV: bool = handle_bool_env_var(ENV_ZENML_DEBUG, default=False)
 

zenml/integrations/gcp/orchestrators/vertex_orchestrator.py

@@ -341,13 +341,55 @@ class VertexOrchestrator(ContainerizedOrchestrator, GoogleCredentialsMixin):
                 self.config.workload_service_account
             )
 
+        # Create a dictionary of explicit parameters
+        params = custom_job_parameters.model_dump(
+            exclude_none=True, exclude={"additional_training_job_args"}
+        )
+
+        # Remove None values to let defaults be set by the function
+        params = {k: v for k, v in params.items() if v is not None}
+
+        # Add environment variables
+        params["env"] = [
+            {"name": key, "value": value} for key, value in environment.items()
+        ]
+
+        # Check if any advanced parameters will override explicit parameters
+        if custom_job_parameters.additional_training_job_args:
+            overridden_params = set(params.keys()) & set(
+                custom_job_parameters.additional_training_job_args.keys()
+            )
+            if overridden_params:
+                logger.warning(
+                    f"The following explicit parameters are being overridden by values in "
+                    f"additional_training_job_args: {', '.join(overridden_params)}. "
+                    f"This may lead to unexpected behavior. Consider using either explicit "
+                    f"parameters or additional_training_job_args, but not both for the same parameters."
+                )
+
+        # Add any advanced parameters - these will override explicit parameters if provided
+        params.update(custom_job_parameters.additional_training_job_args)
+
+        # Add other parameters from orchestrator config if not already in params
+        if self.config.network and "network" not in params:
+            params["network"] = self.config.network
+
+        if (
+            self.config.encryption_spec_key_name
+            and "encryption_spec_key_name" not in params
+        ):
+            params["encryption_spec_key_name"] = (
+                self.config.encryption_spec_key_name
+            )
+        if (
+            self.config.workload_service_account
+            and "service_account" not in params
+        ):
+            params["service_account"] = self.config.workload_service_account
+
         custom_job_component = create_custom_training_job_from_component(
             component_spec=component,
-            env=[
-                {"name": key, "value": value}
-                for key, value in environment.items()
-            ],
-            **custom_job_parameters.model_dump(),
+            **params,
         )
 
         return custom_job_component

zenml/integrations/gcp/vertex_custom_job_parameters.py

@@ -13,7 +13,7 @@
 #  permissions and limitations under the License.
 """Vertex custom job parameter model."""
 
-from typing import Optional
+from typing import Any, Dict, Optional
 
 from pydantic import BaseModel
 
@@ -37,8 +37,21 @@ class VertexCustomJobParameters(BaseModel):
         boot_disk_type: Type of the boot disk. (Default: pd-ssd)
             https://cloud.google.com/vertex-ai/docs/training/configure-compute#boot_disk_options
         persistent_resource_id: The ID of the persistent resource to use for the job.
+            If empty (default), the job will not use a persistent resource.
+            When using a persistent resource, you must also specify a service_account.
+            Conversely, when explicitly setting this to an empty string, you
+            should not specify a service_account (ZenML will handle this automatically).
             https://cloud.google.com/vertex-ai/docs/training/persistent-resource-overview
         service_account: Specifies the service account to be used.
+            This is required when using a persistent_resource_id, and
+            should not be set when persistent_resource_id="".
+        additional_training_job_args: Additional arguments to pass to the create_custom_training_job_from_component
+            function. This allows passing any additional parameters supported by the Google
+            Cloud Pipeline Components library without requiring ZenML to update its API.
+            Note: If you specify parameters in this dictionary that are also defined as explicit
+            attributes (like machine_type or boot_disk_size_gb), the values in this dictionary
+            will override the explicit values.
+            See: https://google-cloud-pipeline-components.readthedocs.io/en/google-cloud-pipeline-components-2.19.0/api/v1/custom_job.html
     """
 
     accelerator_type: Optional[str] = None
@@ -48,3 +61,4 @@ class VertexCustomJobParameters(BaseModel):
     boot_disk_type: str = "pd-ssd"
     persistent_resource_id: Optional[str] = None
     service_account: Optional[str] = None
+    additional_training_job_args: Dict[str, Any] = {}

zenml/integrations/kubernetes/orchestrators/kubernetes_orchestrator.py

@@ -543,7 +543,6 @@ class KubernetesOrchestrator(ContainerizedOrchestrator):
                 mount_local_stores=self.config.is_local,
             )
 
-            logger.info("Waiting for Kubernetes orchestrator pod to start...")
             kube_utils.create_and_wait_for_pod_to_start(
                 core_api=self._k8s_core_api,
                 pod_display_name="Kubernetes orchestrator pod",

zenml/integrations/kubernetes/orchestrators/kubernetes_orchestrator_entrypoint.py

@@ -173,7 +173,6 @@ def main() -> None:
             mount_local_stores=mount_local_stores,
         )
 
-        logger.info(f"Waiting for pod of step `{step_name}` to start...")
         kube_utils.create_and_wait_for_pod_to_start(
             core_api=core_api,
             pod_display_name=f"pod for step `{step_name}`",

zenml/integrations/kubernetes/step_operators/kubernetes_step_operator.py

@@ -218,9 +218,6 @@ class KubernetesStepOperator(BaseStepOperator):
             mount_local_stores=False,
         )
 
-        logger.info(
-            "Waiting for pod of step `%s` to start...", info.pipeline_step_name
-        )
         kube_utils.create_and_wait_for_pod_to_start(
             core_api=self._k8s_core_api,
             pod_display_name=f"pod of step `{info.pipeline_step_name}`",

zenml/logging/step_logging.py

@@ -13,6 +13,7 @@
 #  permissions and limitations under the License.
 """ZenML logging handler."""
 
+import logging
 import os
 import re
 import sys
@@ -48,6 +49,7 @@ logger = get_logger(__name__)
 redirected: ContextVar[bool] = ContextVar("redirected", default=False)
 
 LOGS_EXTENSION = ".log"
+PIPELINE_RUN_LOGS_FOLDER = "pipeline_runs"
 
 
 def remove_ansi_escape_codes(text: str) -> str:
@@ -65,14 +67,14 @@ def remove_ansi_escape_codes(text: str) -> str:
 
 def prepare_logs_uri(
     artifact_store: "BaseArtifactStore",
-    step_name: str,
+    step_name: Optional[str] = None,
     log_key: Optional[str] = None,
 ) -> str:
     """Generates and prepares a URI for the log file or folder for a step.
 
     Args:
         artifact_store: The artifact store on which the artifact will be stored.
-        step_name: Name of the step.
+        step_name: Name of the step. Skipped for global pipeline run logs.
         log_key: The unique identification key of the log file.
 
     Returns:
@@ -81,11 +83,8 @@ def prepare_logs_uri(
     if log_key is None:
         log_key = str(uuid4())
 
-    logs_base_uri = os.path.join(
-        artifact_store.path,
-        step_name,
-        "logs",
-    )
+    subfolder = step_name or PIPELINE_RUN_LOGS_FOLDER
+    logs_base_uri = os.path.join(artifact_store.path, subfolder, "logs")
 
     # Create the dir
     if not artifact_store.exists(logs_base_uri):
@@ -210,7 +209,7 @@ def fetch_logs(
         artifact_store.cleanup()
 
 
-class StepLogsStorage:
+class PipelineLogsStorage:
     """Helper class which buffers and stores logs to a given URI."""
 
     def __init__(
@@ -324,6 +323,18 @@ class StepLogsStorage:
             self.disabled = True
 
             try:
+                # The configured logging handler uses a lock to ensure that
+                # logs generated by different threads are not interleaved.
+                # Given that most artifact stores are based on fsspec, which
+                # use a separate thread for async operations, it may happen that
+                # the fsspec library itself will log something, which will end
+                # up in a deadlock.
+                # To avoid this, we temporarily disable the lock in the logging
+                # handler while writing to the file.
+                logging_handler = logging.getLogger().handlers[0]
+                logging_lock = logging_handler.lock
+                logging_handler.lock = None
+
                 if self.buffer:
                     if self.artifact_store.config.IS_IMMUTABLE_FILESYSTEM:
                         _logs_uri = self._get_timestamped_filename()
@@ -353,6 +364,9 @@ class StepLogsStorage:
                 # I/O errors.
                 logger.error(f"Error while trying to write logs: {e}")
             finally:
+                # Restore the original logging handler lock
+                logging_handler.lock = logging_lock
+
                 self.buffer = []
                 self.last_save_time = time.time()
 
@@ -418,27 +432,32 @@ class StepLogsStorage:
                         )
 
 
-class StepLogsStorageContext:
-    """Context manager which patches stdout and stderr during step execution."""
+class PipelineLogsStorageContext:
+    """Context manager which patches stdout and stderr during pipeline run execution."""
 
     def __init__(
-        self, logs_uri: str, artifact_store: "BaseArtifactStore"
+        self,
+        logs_uri: str,
+        artifact_store: "BaseArtifactStore",
+        prepend_step_name: bool = True,
     ) -> None:
         """Initializes and prepares a storage object.
 
         Args:
             logs_uri: the URI of the logs file.
-            artifact_store: Artifact Store from the current step context.
+            artifact_store: Artifact Store from the current pipeline run context.
+            prepend_step_name: Whether to prepend the step name to the logs.
         """
-        self.storage = StepLogsStorage(
+        self.storage = PipelineLogsStorage(
             logs_uri=logs_uri, artifact_store=artifact_store
         )
+        self.prepend_step_name = prepend_step_name
 
-    def __enter__(self) -> "StepLogsStorageContext":
+    def __enter__(self) -> "PipelineLogsStorageContext":
         """Enter condition of the context manager.
 
         Wraps the `write` method of both stderr and stdout, so each incoming
-        message gets stored in the step logs storage.
+        message gets stored in the pipeline logs storage.
 
         Returns:
             self
@@ -499,14 +518,17 @@ class StepLogsStorageContext:
         """
 
         def wrapped_write(*args: Any, **kwargs: Any) -> Any:
-            # Check if step names in logs are disabled via env var
-            step_names_disabled = handle_bool_env_var(
-                ENV_ZENML_DISABLE_STEP_NAMES_IN_LOGS, default=False
+            step_names_disabled = (
+                handle_bool_env_var(
+                    ENV_ZENML_DISABLE_STEP_NAMES_IN_LOGS, default=False
+                )
+                or not self.prepend_step_name
             )
 
             if step_names_disabled:
                 output = method(*args, **kwargs)
             else:
+                message = args[0]
                 # Try to get step context if not available yet
                 step_context = None
                 try:
@@ -515,9 +537,7 @@ class StepLogsStorageContext:
                     pass
 
                 if step_context and args[0] != "\n":
-                    message = f"[{step_context.step_name}] " + args[0]
-                else:
-                    message = args[0]
+                    message = f"[{step_context.step_name}] " + message
 
                 output = method(message, *args[1:], **kwargs)
 

zenml/login/credentials_store.py

@@ -25,6 +25,7 @@ from zenml.constants import (
 from zenml.io import fileio
 from zenml.logger import get_logger
 from zenml.login.credentials import APIToken, ServerCredentials, ServerType
+from zenml.login.pro.constants import ZENML_PRO_API_URL
 from zenml.login.pro.workspace.models import WorkspaceRead
 from zenml.models import OAuthTokenResponse, ServerModel
 from zenml.utils import yaml_utils
@@ -396,6 +397,36 @@ class CredentialsStore(metaclass=SingletonMetaClass):
         """
         return self.get_pro_token(pro_api_url) is not None
 
+    def can_login(self, server_url: str) -> bool:
+        """Check if credentials to login to the given server exist.
+
+        Args:
+            server_url: The server URL for which to check the authentication.
+
+        Returns:
+            True if the credentials store contains credentials that can be used
+            to login to the given server URL, False otherwise.
+        """
+        self.check_and_reload_from_file()
+        credentials = self.get_credentials(server_url)
+        if not credentials:
+            return False
+
+        if credentials.api_key is not None:
+            return True
+        elif (
+            credentials.username is not None
+            and credentials.password is not None
+        ):
+            return True
+        elif credentials.type == ServerType.PRO:
+            pro_api_url = credentials.pro_api_url or ZENML_PRO_API_URL
+            pro_token = self.get_pro_token(pro_api_url, allow_expired=False)
+            if pro_token:
+                return True
+
+        return False
+
     def set_api_key(
         self,
         server_url: str,

zenml/models/v2/base/base.py

@@ -134,7 +134,7 @@ class BaseResponse(BaseZenModel, Generic[AnyBody, AnyMetadata, AnyResources]):
             )
 
         # Check if the name has changed
-        if "name" in self.model_fields:
+        if "name" in type(self).model_fields:
             original_name = getattr(self, "name")
             hydrated_name = getattr(hydrated_model, "name")
 
@@ -172,7 +172,7 @@ class BaseResponse(BaseZenModel, Generic[AnyBody, AnyMetadata, AnyResources]):
                     )
 
         # Check all the fields in the body
-        for field in self.get_body().model_fields:
+        for field in type(self.get_body()).model_fields:
             original_value = getattr(self.get_body(), field)
             hydrated_value = getattr(hydrated_model.get_body(), field)
 
@@ -255,7 +255,9 @@ class BaseResponse(BaseZenModel, Generic[AnyBody, AnyMetadata, AnyResources]):
         """
         if self.metadata is None:
             # If the metadata is not there, check the class first.
-            metadata_annotation = self.model_fields["metadata"].annotation
+            metadata_annotation = (
+                type(self).model_fields["metadata"].annotation
+            )
             assert metadata_annotation is not None, (
                 "For each response model, an annotated metadata"
                 "field should exist."
@@ -293,7 +295,9 @@ class BaseResponse(BaseZenModel, Generic[AnyBody, AnyMetadata, AnyResources]):
         """
         if self.resources is None:
             # If the resources are not there, check the class first.
-            resources_annotation = self.model_fields["resources"].annotation
+            resources_annotation = (
+                type(self).model_fields["resources"].annotation
+            )
             assert resources_annotation is not None, (
                 "For each response model, an annotated resources"
                 "field should exist."

zenml/models/v2/base/filter.py

@@ -665,7 +665,7 @@ class BaseFilter(BaseModel):
             A list of Filter models.
         """
         return self._generate_filter_list(
-            {key: getattr(self, key) for key in self.model_fields}
+            {key: getattr(self, key) for key in type(self).model_fields}
         )
 
     @property

zenml/models/v2/core/pipeline_run.py

@@ -45,6 +45,7 @@ from zenml.models.v2.base.scoped import (
     RunMetadataFilterMixin,
     TaggableFilter,
 )
+from zenml.models.v2.core.logs import LogsRequest
 from zenml.models.v2.core.model_version import ModelVersionResponse
 from zenml.models.v2.core.tag import TagResponse
 from zenml.utils.tag_utils import Tag
@@ -55,6 +56,7 @@ if TYPE_CHECKING:
     from zenml.models import TriggerExecutionResponse
     from zenml.models.v2.core.artifact_version import ArtifactVersionResponse
     from zenml.models.v2.core.code_reference import CodeReferenceResponse
+    from zenml.models.v2.core.logs import LogsResponse
     from zenml.models.v2.core.pipeline import PipelineResponse
     from zenml.models.v2.core.pipeline_build import (
         PipelineBuildResponse,
@@ -124,6 +126,10 @@ class PipelineRunRequest(ProjectScopedRequest):
         default=None,
         title="Tags of the pipeline run.",
     )
+    logs: Optional[LogsRequest] = Field(
+        default=None,
+        title="Logs of the pipeline run.",
+    )
 
     model_config = ConfigDict(protected_namespaces=())
 
@@ -252,6 +258,10 @@ class PipelineRunResponseResources(ProjectScopedResponseResources):
     tags: List[TagResponse] = Field(
         title="Tags associated with the pipeline run.",
     )
+    logs: Optional["LogsResponse"] = Field(
+        title="Logs associated with this pipeline run.",
+        default=None,
+    )
 
     # TODO: In Pydantic v2, the `model_` is a protected namespaces for all
     #  fields defined under base models. If not handled, this raises a warning.
@@ -579,6 +589,15 @@ class PipelineRunResponse(
         """
         return self.get_resources().tags
 
+    @property
+    def logs(self) -> Optional["LogsResponse"]:
+        """The `logs` property.
+
+        Returns:
+            the value of the property.
+        """
+        return self.get_resources().logs
+
 
 # ------------------ Filter Model ------------------
 

zenml/orchestrators/step_launcher.py

@@ -158,7 +158,7 @@ class StepLauncher:
                 step_name=self._step_name,
             )
 
-            logs_context = step_logging.StepLogsStorageContext(
+            logs_context = step_logging.PipelineLogsStorageContext(
                 logs_uri=logs_uri, artifact_store=self._stack.artifact_store
             )  # type: ignore[assignment]
 
@@ -240,7 +240,7 @@ class StepLauncher:
                             # the external jobs in step operators
                             if isinstance(
                                 logs_context,
-                                step_logging.StepLogsStorageContext,
+                                step_logging.PipelineLogsStorageContext,
                             ):
                                 force_write_logs = partial(
                                     logs_context.storage.save_to_file,
@@ -421,7 +421,6 @@ class StepLauncher:
         )
         environment = orchestrator_utils.get_config_environment_vars(
             pipeline_run_id=step_run_info.run_id,
-            step_run_id=step_run_info.step_run_id,
         )
         if last_retry:
             environment[ENV_ZENML_IGNORE_FAILURE_HOOK] = str(False)

zenml/orchestrators/step_runner.py

@@ -40,7 +40,7 @@ from zenml.constants import (
 from zenml.enums import ArtifactSaveType
 from zenml.exceptions import StepInterfaceError
 from zenml.logger import get_logger
-from zenml.logging.step_logging import StepLogsStorageContext, redirected
+from zenml.logging.step_logging import PipelineLogsStorageContext, redirected
 from zenml.materializers.base_materializer import BaseMaterializer
 from zenml.models.v2.core.step_run import StepRunInputResponse
 from zenml.orchestrators.publish_utils import (
@@ -136,7 +136,7 @@ class StepRunner:
         logs_context = nullcontext()
         if step_logging_enabled and not redirected.get():
             if step_run.logs:
-                logs_context = StepLogsStorageContext(  # type: ignore[assignment]
+                logs_context = PipelineLogsStorageContext(  # type: ignore[assignment]
                     logs_uri=step_run.logs.uri,
                     artifact_store=self._stack.artifact_store,
                 )