zenml-nightly 0.83.1.dev20250702__py3-none-any.whl → 0.83.1.dev20250704__py3-none-any.whl

zenml/VERSION

@@ -1 +1 @@
-0.83.1.dev20250702
+0.83.1.dev20250704

zenml/cli/pipeline.py

@@ -34,7 +34,7 @@ from zenml.models import (
     ScheduleFilter,
 )
 from zenml.pipelines.pipeline_definition import Pipeline
-from zenml.utils import source_utils, uuid_utils
+from zenml.utils import run_utils, source_utils, uuid_utils
 from zenml.utils.yaml_utils import write_yaml
 
 logger = get_logger(__name__)
@@ -511,6 +511,59 @@ def list_pipeline_runs(**kwargs: Any) -> None:
         cli_utils.print_page_info(pipeline_runs)
 
 
+@runs.command("stop")
+@click.argument("run_name_or_id", type=str, required=True)
+@click.option(
+    "--graceful",
+    "-g",
+    is_flag=True,
+    default=False,
+    help="Use graceful shutdown (default is False).",
+)
+@click.option(
+    "--yes",
+    "-y",
+    is_flag=True,
+    default=False,
+    help="Don't ask for confirmation.",
+)
+def stop_pipeline_run(
+    run_name_or_id: str,
+    graceful: bool = False,
+    yes: bool = False,
+) -> None:
+    """Stop a running pipeline.
+
+    Args:
+        run_name_or_id: The name or ID of the pipeline run to stop.
+        graceful: If True, uses graceful shutdown. If False, forces immediate termination.
+        yes: If set, don't ask for confirmation.
+    """
+    # Ask for confirmation to stop run.
+    if not yes:
+        action = "gracefully stop" if graceful else "force stop"
+        confirmation = cli_utils.confirmation(
+            f"Are you sure you want to {action} pipeline run `{run_name_or_id}`?"
+        )
+        if not confirmation:
+            cli_utils.declare("Not stopping the pipeline run.")
+            return
+
+    # Stop run.
+    try:
+        run = Client().get_pipeline_run(name_id_or_prefix=run_name_or_id)
+        run_utils.stop_run(run=run, graceful=graceful)
+        action = "Gracefully stopped" if graceful else "Force stopped"
+        cli_utils.declare(f"{action} pipeline run '{run.name}'.")
+    except NotImplementedError:
+        cli_utils.error(
+            "The orchestrator used for this pipeline run does not support "
+            f"{'gracefully' if graceful else 'forcefully'} stopping runs."
+        )
+    except Exception as e:
+        cli_utils.error(f"Failed to stop pipeline run: {e}")
+
+
 @runs.command("delete")
 @click.argument("run_name_or_id", type=str, required=True)
 @click.option(

zenml/cli/utils.py

@@ -2214,6 +2214,8 @@ def get_execution_status_emoji(status: "ExecutionStatus") -> str:
         return ":white_check_mark:"
     if status == ExecutionStatus.CACHED:
         return ":package:"
+    if status == ExecutionStatus.STOPPED or status == ExecutionStatus.STOPPING:
+        return ":stop_sign:"
     raise RuntimeError(f"Unknown status: {status}")
 
 

zenml/config/compiler.py

@@ -556,7 +556,16 @@ class Compiler:
 
         for name, step in steps.items():
             step_operator = step.config.step_operator
-            if step_operator and step_operator not in available_step_operators:
+            if step_operator is True:
+                if not available_step_operators:
+                    raise StackValidationError(
+                        f"Step `{name}` requires a step operator, but no step "
+                        f"operators are configured in the stack '{stack.name}'."
+                    )
+            elif (
+                isinstance(step_operator, str)
+                and step_operator not in available_step_operators
+            ):
                 raise StackValidationError(
                     f"Step `{name}` requires step operator "
                     f"'{step_operator}' which is not configured in "
@@ -565,8 +574,15 @@ class Compiler:
                 )
 
             experiment_tracker = step.config.experiment_tracker
-            if (
-                experiment_tracker
+            if experiment_tracker is True:
+                if not available_experiment_trackers:
+                    raise StackValidationError(
+                        f"Step `{name}` requires an experiment tracker, but no "
+                        f"experiment trackers are configured in the stack "
+                        f"'{stack.name}'."
+                    )
+            elif (
+                isinstance(experiment_tracker, str)
                 and experiment_tracker not in available_experiment_trackers
             ):
                 raise StackValidationError(

zenml/config/step_configurations.py

@@ -145,8 +145,8 @@ class StepConfigurationUpdate(StrictBaseModel):
     enable_artifact_metadata: Optional[bool] = None
     enable_artifact_visualization: Optional[bool] = None
     enable_step_logs: Optional[bool] = None
-    step_operator: Optional[str] = None
-    experiment_tracker: Optional[str] = None
+    step_operator: Optional[Union[bool, str]] = None
+    experiment_tracker: Optional[Union[bool, str]] = None
     parameters: Dict[str, Any] = {}
     settings: Dict[str, SerializeAsAny[BaseSettings]] = {}
     extra: Dict[str, Any] = {}
@@ -158,6 +158,38 @@ class StepConfigurationUpdate(StrictBaseModel):
 
     outputs: Mapping[str, PartialArtifactConfiguration] = {}
 
+    def uses_step_operator(self, name: str) -> bool:
+        """Checks if the step configuration uses the given step operator.
+
+        Args:
+            name: The name of the step operator.
+
+        Returns:
+            If the step configuration uses the given step operator.
+        """
+        if self.step_operator is True:
+            return True
+        elif isinstance(self.step_operator, str):
+            return self.step_operator == name
+        else:
+            return False
+
+    def uses_experiment_tracker(self, name: str) -> bool:
+        """Checks if the step configuration uses the given experiment tracker.
+
+        Args:
+            name: The name of the experiment tracker.
+
+        Returns:
+            If the step configuration uses the given experiment tracker.
+        """
+        if self.experiment_tracker is True:
+            return True
+        elif isinstance(self.experiment_tracker, str):
+            return self.experiment_tracker == name
+        else:
+            return False
+
 
 class PartialStepConfiguration(StepConfigurationUpdate):
     """Class representing a partial step configuration."""

zenml/constants.py

@@ -416,6 +416,7 @@ STATISTICS = "/statistics"
 STATUS = "/status"
 STEP_CONFIGURATION = "/step-configuration"
 STEPS = "/steps"
+STOP = "/stop"
 TAGS = "/tags"
 TAG_RESOURCES = "/tag_resources"
 TRIGGERS = "/triggers"

zenml/enums.py

@@ -71,25 +71,28 @@ class ZenMLServiceType(StrEnum):
 
 
 class ExecutionStatus(StrEnum):
-    """Enum that represents the current status of a step or pipeline run."""
+    """Enum that represents the execution status of a step or pipeline run."""
 
     INITIALIZING = "initializing"
     FAILED = "failed"
     COMPLETED = "completed"
     RUNNING = "running"
     CACHED = "cached"
+    STOPPED = "stopped"
+    STOPPING = "stopping"
 
     @property
     def is_finished(self) -> bool:
-        """Whether the execution status refers to a finished execution.
+        """Returns whether the execution status is in a finished state.
 
         Returns:
-            Whether the execution status refers to a finished execution.
+            Whether the execution status is finished.
         """
         return self in {
             ExecutionStatus.FAILED,
             ExecutionStatus.COMPLETED,
             ExecutionStatus.CACHED,
+            ExecutionStatus.STOPPED,
         }
 
 

zenml/exceptions.py

@@ -122,6 +122,14 @@ class IllegalOperationError(ZenMLBaseException):
     """Raised when an illegal operation is attempted."""
 
 
+class RunStoppedException(ZenMLBaseException):
+    """Raised when a ZenML pipeline run gets stopped by the user."""
+
+
+class RunInterruptedException(ZenMLBaseException):
+    """Raised when a ZenML step gets interrupted for an unknown reason."""
+
+
 class MethodNotAllowedError(ZenMLBaseException):
     """Raised when the server does not allow a request method."""
 

zenml/integrations/aws/orchestrators/sagemaker_orchestrator.py

@@ -853,12 +853,16 @@ class SagemakerOrchestrator(ContainerizedOrchestrator):
         )["PipelineExecutionStatus"]
 
         # Map the potential outputs to ZenML ExecutionStatus. Potential values:
-        # https://cloud.google.com/vertex-ai/docs/reference/rest/v1beta1/PipelineState
-        if status in ["Executing", "Stopping"]:
+        # https://docs.aws.amazon.com/sagemaker/latest/APIReference/API_DescribePipelineExecution.html
+        if status == "Executing":
             return ExecutionStatus.RUNNING
-        elif status in ["Stopped", "Failed"]:
+        elif status == "Stopping":
+            return ExecutionStatus.STOPPING
+        elif status == "Stopped":
+            return ExecutionStatus.STOPPED
+        elif status == "Failed":
             return ExecutionStatus.FAILED
-        elif status in ["Succeeded"]:
+        elif status == "Succeeded":
             return ExecutionStatus.COMPLETED
         else:
             raise ValueError("Unknown status for the pipeline execution.")

zenml/integrations/aws/step_operators/sagemaker_step_operator.py

@@ -152,7 +152,7 @@ class SagemakerStepOperator(BaseStepOperator):
         """
         builds = []
         for step_name, step in deployment.step_configurations.items():
-            if step.config.step_operator == self.name:
+            if step.config.uses_step_operator(self.name):
                 build = BuildConfiguration(
                     key=SAGEMAKER_DOCKER_IMAGE_KEY,
                     settings=step.config.docker_settings,

zenml/integrations/azure/orchestrators/azureml_orchestrator.py

@@ -515,14 +515,16 @@ class AzureMLOrchestrator(ContainerizedOrchestrator):
             return ExecutionStatus.INITIALIZING
         elif status in ["Running", "Finalizing"]:
             return ExecutionStatus.RUNNING
+        elif status == "CancelRequested":
+            return ExecutionStatus.STOPPING
+        elif status == "Canceled":
+            return ExecutionStatus.STOPPED
         elif status in [
-            "CancelRequested",
             "Failed",
-            "Canceled",
             "NotResponding",
         ]:
             return ExecutionStatus.FAILED
-        elif status in ["Completed"]:
+        elif status == "Completed":
             return ExecutionStatus.COMPLETED
         else:
             raise ValueError("Unknown status for the pipeline job.")

zenml/integrations/azure/step_operators/azureml_step_operator.py

@@ -149,7 +149,7 @@ class AzureMLStepOperator(BaseStepOperator):
         """
         builds = []
         for step_name, step in deployment.step_configurations.items():
-            if step.config.step_operator == self.name:
+            if step.config.uses_step_operator(self.name):
                 build = BuildConfiguration(
                     key=AZUREML_STEP_OPERATOR_DOCKER_IMAGE_KEY,
                     settings=step.config.docker_settings,

zenml/integrations/gcp/orchestrators/vertex_orchestrator.py

@@ -942,7 +942,7 @@ class VertexOrchestrator(ContainerizedOrchestrator, GoogleCredentialsMixin):
 
         # Map the potential outputs to ZenML ExecutionStatus. Potential values:
         # https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sagemaker/client/describe_pipeline_execution.html#
-        if status in [PipelineState.PIPELINE_STATE_UNSPECIFIED]:
+        if status == PipelineState.PIPELINE_STATE_UNSPECIFIED:
             return run.status
         elif status in [
             PipelineState.PIPELINE_STATE_QUEUED,
@@ -954,14 +954,13 @@ class VertexOrchestrator(ContainerizedOrchestrator, GoogleCredentialsMixin):
             PipelineState.PIPELINE_STATE_PAUSED,
         ]:
             return ExecutionStatus.RUNNING
-        elif status in [PipelineState.PIPELINE_STATE_SUCCEEDED]:
+        elif status == PipelineState.PIPELINE_STATE_SUCCEEDED:
             return ExecutionStatus.COMPLETED
-
-        elif status in [
-            PipelineState.PIPELINE_STATE_FAILED,
-            PipelineState.PIPELINE_STATE_CANCELLING,
-            PipelineState.PIPELINE_STATE_CANCELLED,
-        ]:
+        elif status == PipelineState.PIPELINE_STATE_CANCELLING:
+            return ExecutionStatus.STOPPING
+        elif status == PipelineState.PIPELINE_STATE_CANCELLED:
+            return ExecutionStatus.STOPPED
+        elif status == PipelineState.PIPELINE_STATE_FAILED:
             return ExecutionStatus.FAILED
         else:
             raise ValueError("Unknown status for the pipeline job.")

zenml/integrations/gcp/step_operators/vertex_step_operator.py

@@ -161,7 +161,7 @@ class VertexStepOperator(BaseStepOperator, GoogleCredentialsMixin):
         """
         builds = []
         for step_name, step in deployment.step_configurations.items():
-            if step.config.step_operator == self.name:
+            if step.config.uses_step_operator(self.name):
                 build = BuildConfiguration(
                     key=VERTEX_DOCKER_IMAGE_KEY,
                     settings=step.config.docker_settings,

zenml/integrations/kubernetes/flavors/kubernetes_orchestrator_flavor.py

@@ -69,6 +69,10 @@ class KubernetesOrchestratorSettings(BaseSettings):
             scheduling a pipeline.
         prevent_orchestrator_pod_caching: If `True`, the orchestrator pod will
             not try to compute cached steps before starting the step pods.
+        always_build_pipeline_image: If `True`, the orchestrator will always
+            build the pipeline image, even if all steps have a custom build.
+        pod_stop_grace_period: When stopping a pipeline run, the amount of
+            seconds to wait for a step pod to shutdown gracefully.
     """
 
     synchronous: bool = True
@@ -88,6 +92,8 @@ class KubernetesOrchestratorSettings(BaseSettings):
     failed_jobs_history_limit: Optional[NonNegativeInt] = None
     ttl_seconds_after_finished: Optional[NonNegativeInt] = None
     prevent_orchestrator_pod_caching: bool = False
+    always_build_pipeline_image: bool = False
+    pod_stop_grace_period: PositiveInt = 30
 
 
 class KubernetesOrchestratorConfig(

zenml/integrations/kubernetes/orchestrators/kubernetes_orchestrator.py

@@ -70,7 +70,11 @@ from zenml.orchestrators.utils import get_orchestrator_run_name
 from zenml.stack import StackValidator
 
 if TYPE_CHECKING:
-    from zenml.models import PipelineDeploymentResponse, PipelineRunResponse
+    from zenml.models import (
+        PipelineDeploymentBase,
+        PipelineDeploymentResponse,
+        PipelineRunResponse,
+    )
     from zenml.stack import Stack
 
 logger = get_logger(__name__)
@@ -84,6 +88,22 @@ class KubernetesOrchestrator(ContainerizedOrchestrator):
 
     _k8s_client: Optional[k8s_client.ApiClient] = None
 
+    def should_build_pipeline_image(
+        self, deployment: "PipelineDeploymentBase"
+    ) -> bool:
+        """Whether to always build the pipeline image.
+
+        Args:
+            deployment: The pipeline deployment.
+
+        Returns:
+            Whether to always build the pipeline image.
+        """
+        settings = cast(
+            KubernetesOrchestratorSettings, self.get_settings(deployment)
+        )
+        return settings.always_build_pipeline_image
+
     def get_kube_client(
         self, incluster: Optional[bool] = None
     ) -> k8s_client.ApiClient:
@@ -545,6 +565,7 @@ class KubernetesOrchestrator(ContainerizedOrchestrator):
                 successful_jobs_history_limit=settings.successful_jobs_history_limit,
                 failed_jobs_history_limit=settings.failed_jobs_history_limit,
                 ttl_seconds_after_finished=settings.ttl_seconds_after_finished,
+                termination_grace_period_seconds=settings.pod_stop_grace_period,
                 labels=orchestrator_pod_labels,
             )
 
@@ -570,6 +591,7 @@ class KubernetesOrchestrator(ContainerizedOrchestrator):
                 env=environment,
                 labels=orchestrator_pod_labels,
                 mount_local_stores=self.config.is_local,
+                termination_grace_period_seconds=settings.pod_stop_grace_period,
             )
 
             kube_utils.create_and_wait_for_pod_to_start(
@@ -663,6 +685,92 @@ class KubernetesOrchestrator(ContainerizedOrchestrator):
                 f"{ENV_ZENML_KUBERNETES_RUN_ID}."
             )
 
+    def _stop_run(
+        self, run: "PipelineRunResponse", graceful: bool = True
+    ) -> None:
+        """Stops a specific pipeline run by terminating step pods.
+
+        Args:
+            run: The run that was executed by this orchestrator.
+            graceful: If True, does nothing (lets the orchestrator and steps finish naturally).
+                If False, stops all running step pods.
+
+        Raises:
+            RuntimeError: If we fail to stop the run.
+        """
+        # If graceful, do nothing and let the orchestrator handle the stop naturally
+        if graceful:
+            logger.info(
+                "Graceful stop requested - the orchestrator pod will handle "
+                "stopping naturally"
+            )
+            return
+
+        pods_stopped = []
+        errors = []
+
+        # Find all pods with the orchestrator run ID label
+        label_selector = f"run_id={kube_utils.sanitize_label(str(run.id))}"
+        try:
+            pods = self._k8s_core_api.list_namespaced_pod(
+                namespace=self.config.kubernetes_namespace,
+                label_selector=label_selector,
+            )
+        except Exception as e:
+            raise RuntimeError(
+                f"Failed to list step pods with run ID {run.id}: {e}"
+            )
+
+        # Filter to only include running or pending pods
+        for pod in pods.items:
+            if pod.status.phase not in ["Running", "Pending"]:
+                logger.debug(
+                    f"Skipping pod {pod.metadata.name} with status {pod.status.phase}"
+                )
+                continue
+
+            try:
+                self._k8s_core_api.delete_namespaced_pod(
+                    name=pod.metadata.name,
+                    namespace=self.config.kubernetes_namespace,
+                )
+                pods_stopped.append(f"step pod: {pod.metadata.name}")
+                logger.debug(
+                    f"Successfully initiated graceful stop of step pod: {pod.metadata.name}"
+                )
+            except Exception as e:
+                error_msg = f"Failed to stop step pod {pod.metadata.name}: {e}"
+                logger.warning(error_msg)
+                errors.append(error_msg)
+
+        # Summary logging
+        settings = cast(KubernetesOrchestratorSettings, self.get_settings(run))
+        grace_period_seconds = settings.pod_stop_grace_period
+        if pods_stopped:
+            logger.debug(
+                f"Successfully initiated graceful termination of: {', '.join(pods_stopped)}. "
+                f"Pods will terminate within {grace_period_seconds} seconds."
+            )
+
+        if errors:
+            error_summary = "; ".join(errors)
+            if not pods_stopped:
+                # If nothing was stopped successfully, raise an error
+                raise RuntimeError(
+                    f"Failed to stop pipeline run: {error_summary}"
+                )
+            else:
+                # If some things were stopped but others failed, raise an error
+                raise RuntimeError(
+                    f"Partial stop operation completed with errors: {error_summary}"
+                )
+
+        # If no step pods were found and no errors occurred
+        if not pods_stopped and not errors:
+            logger.info(
+                f"No running step pods found for pipeline run with ID: {run.id}"
+            )
+
     def get_pipeline_run_metadata(
         self, run_id: UUID
     ) -> Dict[str, "MetadataType"]:

zenml/integrations/kubernetes/orchestrators/kubernetes_orchestrator_entrypoint.py

@@ -18,6 +18,7 @@ import socket
 from typing import Callable, Dict, Optional, cast
 
 from kubernetes import client as k8s_client
+from kubernetes.client.rest import ApiException
 
 from zenml.client import Client
 from zenml.entrypoints.step_entrypoint_configuration import (
@@ -248,6 +249,7 @@ def main() -> None:
             or settings.service_account_name,
             mount_local_stores=mount_local_stores,
             owner_references=owner_references,
+            termination_grace_period_seconds=settings.pod_stop_grace_period,
             labels=step_pod_labels,
         )
 
@@ -330,6 +332,38 @@ def main() -> None:
             # as the pipeline run status will already have been published.
             pass
 
+    def check_pipeline_cancellation() -> bool:
+        """Check if the pipeline should continue execution.
+
+        Returns:
+            True if execution should continue, False if it should stop.
+        """
+        try:
+            run = client.get_pipeline_run(
+                name_id_or_prefix=pipeline_run.id,
+                project=pipeline_run.project_id,
+                hydrate=False,  # We only need status, not full hydration
+            )
+
+            # If the run is STOPPING or STOPPED, we should stop the execution
+            if run.status in [
+                ExecutionStatus.STOPPING,
+                ExecutionStatus.STOPPED,
+            ]:
+                logger.info(
+                    f"Pipeline run is in {run.status} state, stopping execution"
+                )
+                return False
+
+            return True
+
+        except Exception as e:
+            # If we can't check the status, assume we should continue
+            logger.warning(
+                f"Failed to check pipeline cancellation status: {e}"
+            )
+            return True
+
     parallel_node_startup_waiting_period = (
         orchestrator.config.parallel_step_startup_waiting_period or 0.0
     )
@@ -344,6 +378,7 @@ def main() -> None:
             run_fn=run_step_on_kubernetes,
             preparation_fn=pre_step_run,
             finalize_fn=finalize_run,
+            continue_fn=check_pipeline_cancellation,
             parallel_node_startup_waiting_period=parallel_node_startup_waiting_period,
             max_parallelism=pipeline_settings.max_parallelism,
         ).run()
@@ -360,7 +395,7 @@ def main() -> None:
                     namespace=namespace,
                     secret_name=secret_name,
                 )
-            except k8s_client.rest.ApiException as e:
+            except ApiException as e:
                 logger.error(f"Error cleaning up secret {secret_name}: {e}")
 
 

zenml/integrations/kubernetes/orchestrators/manifest_utils.py

@@ -106,6 +106,7 @@ def build_pod_manifest(
     labels: Optional[Dict[str, str]] = None,
     mount_local_stores: bool = False,
     owner_references: Optional[List[k8s_client.V1OwnerReference]] = None,
+    termination_grace_period_seconds: Optional[int] = 30,
 ) -> k8s_client.V1Pod:
     """Build a Kubernetes pod manifest for a ZenML run or step.
 
@@ -124,6 +125,8 @@ def build_pod_manifest(
         mount_local_stores: Whether to mount the local stores path inside the
             pod.
         owner_references: List of owner references for the pod.
+        termination_grace_period_seconds: The amount of seconds to wait for a
+            pod to shutdown gracefully.
 
     Returns:
         Pod manifest.
@@ -154,19 +157,20 @@ def build_pod_manifest(
         containers=[container_spec],
         restart_policy="Never",
         image_pull_secrets=image_pull_secrets,
+        termination_grace_period_seconds=termination_grace_period_seconds,
     )
 
     if service_account_name is not None:
         pod_spec.service_account_name = service_account_name
 
+    # Apply pod settings if provided
     labels = labels or {}
 
     if pod_settings:
         add_pod_settings(pod_spec, pod_settings)
 
-        # Add pod_settings.labels to the labels
-        if pod_settings.labels:
-            labels.update(pod_settings.labels)
+    if pod_settings and pod_settings.labels:
+        labels.update(pod_settings.labels)
 
     pod_metadata = k8s_client.V1ObjectMeta(
         name=pod_name,
@@ -273,6 +277,7 @@ def build_cron_job_manifest(
     successful_jobs_history_limit: Optional[int] = None,
     failed_jobs_history_limit: Optional[int] = None,
     ttl_seconds_after_finished: Optional[int] = None,
+    termination_grace_period_seconds: Optional[int] = 30,
 ) -> k8s_client.V1CronJob:
     """Create a manifest for launching a pod as scheduled CRON job.
 
@@ -295,6 +300,8 @@ def build_cron_job_manifest(
         failed_jobs_history_limit: The number of failed jobs to retain.
         ttl_seconds_after_finished: The amount of seconds to keep finished jobs
             before deleting them.
+        termination_grace_period_seconds: The amount of seconds to wait for a
+            pod to shutdown gracefully.
 
     Returns:
         CRON job manifest.
@@ -310,6 +317,7 @@ def build_cron_job_manifest(
         env=env,
         labels=labels,
         mount_local_stores=mount_local_stores,
+        termination_grace_period_seconds=termination_grace_period_seconds,
     )
 
     job_spec = k8s_client.V1CronJobSpec(

zenml/integrations/kubernetes/step_operators/kubernetes_step_operator.py

@@ -120,7 +120,7 @@ class KubernetesStepOperator(BaseStepOperator):
         """
         builds = []
         for step_name, step in deployment.step_configurations.items():
-            if step.config.step_operator == self.name:
+            if step.config.uses_step_operator(self.name):
                 build = BuildConfiguration(
                     key=KUBERNETES_STEP_OPERATOR_DOCKER_IMAGE_KEY,
                     settings=step.config.docker_settings,

zenml/integrations/modal/step_operators/modal_step_operator.py

@@ -139,7 +139,7 @@ class ModalStepOperator(BaseStepOperator):
         """
         builds = []
         for step_name, step in deployment.step_configurations.items():
-            if step.config.step_operator == self.name:
+            if step.config.uses_step_operator(self.name):
                 build = BuildConfiguration(
                     key=MODAL_STEP_OPERATOR_DOCKER_IMAGE_KEY,
                     settings=step.config.docker_settings,

zenml/integrations/spark/step_operators/kubernetes_step_operator.py

@@ -124,7 +124,7 @@ class KubernetesSparkStepOperator(SparkStepOperator):
         builds = []
         extra_files = {ENTRYPOINT_NAME: LOCAL_ENTRYPOINT}
         for step_name, step in deployment.step_configurations.items():
-            if step.config.step_operator == self.name:
+            if step.config.uses_step_operator(self.name):
                 build = BuildConfiguration(
                     key=SPARK_DOCKER_IMAGE_KEY,
                     settings=step.config.docker_settings,