zenml-nightly 0.83.1.dev20250710__py3-none-any.whl → 0.84.0.dev20250712__py3-none-any.whl

zenml/integrations/kubernetes/flavors/kubernetes_orchestrator_flavor.py

@@ -15,7 +15,7 @@
 
 from typing import TYPE_CHECKING, Any, Dict, Optional, Type
 
-from pydantic import NonNegativeInt, PositiveInt, field_validator
+from pydantic import Field, NonNegativeInt, PositiveInt, field_validator
 
 from zenml.config.base_settings import BaseSettings
 from zenml.constants import KUBERNETES_CLUSTER_RESOURCE_TYPE
@@ -33,92 +33,122 @@ if TYPE_CHECKING:
 class KubernetesOrchestratorSettings(BaseSettings):
     """Settings for the Kubernetes orchestrator.
 
-    Attributes:
-        synchronous: If `True`, the client running a pipeline using this
-            orchestrator waits until all steps finish running. If `False`,
-            the client returns immediately and the pipeline is executed
-            asynchronously. Defaults to `True`.
-        timeout: How many seconds to wait for synchronous runs. `0` means
-            to wait for an unlimited duration.
-        stream_step_logs: If `True`, the orchestrator pod will stream the logs
-            of the step pods. This only has an effect if specified on the
-            pipeline, not on individual steps.
-        service_account_name: Name of the service account to use for the
-            orchestrator pod. If not provided, a new service account with "edit"
-            permissions will be created.
-        step_pod_service_account_name: Name of the service account to use for the
-            step pods. If not provided, the default service account will be used.
-        privileged: If the container should be run in privileged mode.
-        pod_settings: Pod settings to apply to pods executing the steps.
-        orchestrator_pod_settings: Pod settings to apply to the pod which is
-            launching the actual steps.
-        pod_name_prefix: Prefix to use for the pod name.
-        pod_startup_timeout: The maximum time to wait for a pending step pod to
-            start (in seconds).
-        pod_failure_max_retries: The maximum number of times to retry a step
-            pod if the step Kubernetes pod fails to start
-        pod_failure_retry_delay: The delay in seconds between pod
-            failure retries and pod startup retries (in seconds)
-        pod_failure_backoff: The backoff factor for pod failure retries and
-            pod startup retries.
-        max_parallelism: Maximum number of steps to run in parallel.
-        successful_jobs_history_limit: The number of successful jobs
-            to retain. This only applies to jobs created when scheduling a
-            pipeline.
-        failed_jobs_history_limit: The number of failed jobs to retain.
-            This only applies to jobs created when scheduling a pipeline.
-        ttl_seconds_after_finished: The amount of seconds to keep finished jobs
-            before deleting them. **Note**: This does not clean up the
-            orchestrator pod for non-scheduled runs.
-        active_deadline_seconds: The active deadline seconds for the job that is
-            executing the step.
-        backoff_limit_margin: The value to add to the backoff limit in addition
-            to the step retries. The retry configuration defined on the step
-            defines the maximum number of retries that the server will accept
-            for a step. For this orchestrator, this controls how often the
-            job running the step will try to start the step pod. There are some
-            circumstances however where the job will start the pod, but the pod
-            doesn't actually get to the point of running the step. That means
-            the server will not receive the maximum amount of retry requests,
-            which in turn causes other inconsistencies like wrong step statuses.
-            To mitigate this, this attribute allows to add a margin to the
-            backoff limit. This means that the job will retry the pod startup
-            for the configured amount of times plus the margin, which increases
-            the chance of the server receiving the maximum amount of retry
-            requests.
-        pod_failure_policy: The pod failure policy to use for the job that is
-            executing the step.
-        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.
+    Configuration options for how pipelines are executed on Kubernetes clusters.
+    Field descriptions are defined inline using Field() descriptors.
     """
 
-    synchronous: bool = True
-    timeout: int = 0
-    stream_step_logs: bool = True
-    service_account_name: Optional[str] = None
-    step_pod_service_account_name: Optional[str] = None
-    privileged: bool = False
-    pod_settings: Optional[KubernetesPodSettings] = None
-    orchestrator_pod_settings: Optional[KubernetesPodSettings] = None
-    pod_name_prefix: Optional[str] = None
-    pod_startup_timeout: int = 60 * 10  # Default 10 minutes
-    pod_failure_max_retries: int = 3
-    pod_failure_retry_delay: int = 10
-    pod_failure_backoff: float = 1.0
-    max_parallelism: Optional[PositiveInt] = None
-    successful_jobs_history_limit: Optional[NonNegativeInt] = None
-    failed_jobs_history_limit: Optional[NonNegativeInt] = None
-    ttl_seconds_after_finished: Optional[NonNegativeInt] = None
-    active_deadline_seconds: Optional[NonNegativeInt] = None
-    backoff_limit_margin: NonNegativeInt = 0
-    pod_failure_policy: Optional[Dict[str, Any]] = None
-    prevent_orchestrator_pod_caching: bool = False
-    always_build_pipeline_image: bool = False
-    pod_stop_grace_period: PositiveInt = 30
+    synchronous: bool = Field(
+        default=True,
+        description="Whether to wait for all pipeline steps to complete. "
+        "When `False`, the client returns immediately and execution continues asynchronously.",
+    )
+    timeout: int = Field(
+        default=0,
+        description="Maximum seconds to wait for synchronous runs. Set to `0` for unlimited duration.",
+    )
+    stream_step_logs: bool = Field(
+        default=True,
+        description="If `True`, the orchestrator pod will stream the logs "
+        "of the step pods. This only has an effect if specified on the "
+        "pipeline, not on individual steps.",
+    )
+    service_account_name: Optional[str] = Field(
+        default=None,
+        description="Kubernetes service account for the orchestrator pod. "
+        "If not specified, creates a new account with 'edit' permissions.",
+    )
+    step_pod_service_account_name: Optional[str] = Field(
+        default=None,
+        description="Kubernetes service account for step execution pods. "
+        "Uses the default service account if not specified.",
+    )
+    privileged: bool = Field(
+        default=False,
+        description="Whether to run containers in privileged mode with extended permissions.",
+    )
+    pod_settings: Optional[KubernetesPodSettings] = Field(
+        default=None,
+        description="Pod configuration for step execution containers.",
+    )
+    orchestrator_pod_settings: Optional[KubernetesPodSettings] = Field(
+        default=None,
+        description="Pod configuration for the orchestrator container that launches step pods.",
+    )
+    pod_name_prefix: Optional[str] = Field(
+        default=None,
+        description="Custom prefix for generated pod names. Helps identify pods in the cluster.",
+    )
+    pod_startup_timeout: int = Field(
+        default=600,
+        description="Maximum seconds to wait for step pods to start. Default is 10 minutes.",
+    )
+    pod_failure_max_retries: int = Field(
+        default=3,
+        description="Maximum retry attempts when step pods fail to start.",
+    )
+    pod_failure_retry_delay: int = Field(
+        default=10,
+        description="Delay in seconds between pod failure retry attempts.",
+    )
+    pod_failure_backoff: float = Field(
+        default=1.0,
+        description="Exponential backoff factor for retry delays. Values > 1.0 increase delay with each retry.",
+    )
+    max_parallelism: Optional[PositiveInt] = Field(
+        default=None,
+        description="Maximum number of step pods to run concurrently. No limit if not specified.",
+    )
+    successful_jobs_history_limit: Optional[NonNegativeInt] = Field(
+        default=None,
+        description="Number of successful scheduled jobs to retain in cluster history.",
+    )
+    failed_jobs_history_limit: Optional[NonNegativeInt] = Field(
+        default=None,
+        description="Number of failed scheduled jobs to retain in cluster history.",
+    )
+    ttl_seconds_after_finished: Optional[NonNegativeInt] = Field(
+        default=None,
+        description="Seconds to keep finished scheduled jobs before automatic cleanup.",
+    )
+    active_deadline_seconds: Optional[NonNegativeInt] = Field(
+        default=None,
+        description="Deadline in seconds for the active pod. If the pod is inactive for this many seconds, it will be terminated.",
+    )
+    backoff_limit_margin: NonNegativeInt = Field(
+        default=0,
+        description="The value to add to the backoff limit in addition "
+        "to the step retries. The retry configuration defined on the step "
+        "defines the maximum number of retries that the server will accept "
+        "for a step. For this orchestrator, this controls how often the "
+        "job running the step will try to start the step pod. There are some "
+        "circumstances however where the job will start the pod, but the pod "
+        "doesn't actually get to the point of running the step. That means "
+        "the server will not receive the maximum amount of retry requests, "
+        "which in turn causes other inconsistencies like wrong step statuses. "
+        "To mitigate this, this attribute allows to add a margin to the "
+        "backoff limit. This means that the job will retry the pod startup "
+        "for the configured amount of times plus the margin, which increases "
+        "the chance of the server receiving the maximum amount of retry "
+        "requests.",
+    )
+    pod_failure_policy: Optional[Dict[str, Any]] = Field(
+        default=None,
+        description="The pod failure policy to use for the job that is "
+        "executing the step.",
+    )
+    prevent_orchestrator_pod_caching: bool = Field(
+        default=False,
+        description="Whether to disable caching optimization in the orchestrator pod.",
+    )
+    always_build_pipeline_image: bool = Field(
+        default=False,
+        description="If `True`, the orchestrator will always build the pipeline image, "
+        "even if all steps have a custom build.",
+    )
+    pod_stop_grace_period: PositiveInt = Field(
+        default=30,
+        description="When stopping a pipeline run, the amount of seconds to wait for a step pod to shutdown gracefully.",
+    )
 
     @field_validator("pod_failure_policy", mode="before")
     @classmethod
@@ -144,42 +174,50 @@ class KubernetesOrchestratorSettings(BaseSettings):
 class KubernetesOrchestratorConfig(
     BaseOrchestratorConfig, KubernetesOrchestratorSettings
 ):
-    """Configuration for the Kubernetes orchestrator.
-
-    Attributes:
-        incluster: If `True`, the orchestrator will run the pipeline inside the
-            same cluster in which it itself is running. This requires the client
-            to run in a Kubernetes pod itself. If set, the `kubernetes_context`
-            config option is ignored. If the stack component is linked to a
-            Kubernetes service connector, this field is ignored.
-        kubernetes_context: Name of a Kubernetes context to run pipelines in.
-            If the stack component is linked to a Kubernetes service connector,
-            this field is ignored. Otherwise, it is mandatory.
-        kubernetes_namespace: Name of the Kubernetes namespace to be used.
-            If not provided, `zenml` namespace will be used.
-        local: If `True`, the orchestrator will assume it is connected to a
-            local kubernetes cluster and will perform additional validations and
-            operations to allow using the orchestrator in combination with other
-            local stack components that store data in the local filesystem
-            (i.e. it will mount the local stores directory into the pipeline
-            containers).
-        skip_local_validations: If `True`, the local validations will be
-            skipped.
-        parallel_step_startup_waiting_period: How long to wait in between
-            starting parallel steps. This can be used to distribute server
-            load when running pipelines with a huge amount of parallel steps.
-        pass_zenml_token_as_secret: If `True`, the ZenML token will be passed
-            as a Kubernetes secret to the pods. For this to work, the Kubernetes
-            client must have permissions to create secrets in the namespace.
-    """
-
-    incluster: bool = False
-    kubernetes_context: Optional[str] = None
-    kubernetes_namespace: str = "zenml"
-    local: bool = False
-    skip_local_validations: bool = False
-    parallel_step_startup_waiting_period: Optional[float] = None
-    pass_zenml_token_as_secret: bool = False
+    """Configuration for the Kubernetes orchestrator."""
+
+    incluster: bool = Field(
+        False,
+        description="If `True`, the orchestrator will run the pipeline inside the "
+        "same cluster in which it itself is running. This requires the client "
+        "to run in a Kubernetes pod itself. If set, the `kubernetes_context` "
+        "config option is ignored. If the stack component is linked to a "
+        "Kubernetes service connector, this field is ignored.",
+    )
+    kubernetes_context: Optional[str] = Field(
+        None,
+        description="Name of a Kubernetes context to run pipelines in. "
+        "If the stack component is linked to a Kubernetes service connector, "
+        "this field is ignored. Otherwise, it is mandatory.",
+    )
+    kubernetes_namespace: str = Field(
+        "zenml",
+        description="Name of the Kubernetes namespace to be used. "
+        "If not provided, `zenml` namespace will be used.",
+    )
+    local: bool = Field(
+        False,
+        description="If `True`, the orchestrator will assume it is connected to a "
+        "local kubernetes cluster and will perform additional validations and "
+        "operations to allow using the orchestrator in combination with other "
+        "local stack components that store data in the local filesystem "
+        "(i.e. it will mount the local stores directory into the pipeline containers).",
+    )
+    skip_local_validations: bool = Field(
+        False, description="If `True`, the local validations will be skipped."
+    )
+    parallel_step_startup_waiting_period: Optional[float] = Field(
+        None,
+        description="How long to wait in between starting parallel steps. "
+        "This can be used to distribute server load when running pipelines "
+        "with a huge amount of parallel steps.",
+    )
+    pass_zenml_token_as_secret: bool = Field(
+        False,
+        description="If `True`, the ZenML token will be passed as a Kubernetes secret "
+        "to the pods. For this to work, the Kubernetes client must have permissions "
+        "to create secrets in the namespace.",
+    )
 
     @property
     def is_remote(self) -> bool:

zenml/integrations/kubernetes/flavors/kubernetes_step_operator_flavor.py

@@ -15,6 +15,8 @@
 
 from typing import TYPE_CHECKING, Optional, Type
 
+from pydantic import Field
+
 from zenml.config.base_settings import BaseSettings
 from zenml.constants import KUBERNETES_CLUSTER_RESOURCE_TYPE
 from zenml.integrations.kubernetes import KUBERNETES_STEP_OPERATOR_FLAVOR
@@ -31,27 +33,38 @@ if TYPE_CHECKING:
 class KubernetesStepOperatorSettings(BaseSettings):
     """Settings for the Kubernetes step operator.
 
-    Attributes:
-        pod_settings: Pod settings to apply to pods executing the steps.
-        service_account_name: Name of the service account to use for the pod.
-        privileged: If the container should be run in privileged mode.
-        pod_startup_timeout: The maximum time to wait for a pending step pod to
-            start (in seconds).
-        pod_failure_max_retries: The maximum number of times to retry a step
-            pod if the step Kubernetes pod fails to start
-        pod_failure_retry_delay: The delay in seconds between pod
-            failure retries and pod startup retries (in seconds)
-        pod_failure_backoff: The backoff factor for pod failure retries and
-            pod startup retries.
+    Configuration options for individual step execution on Kubernetes.
+    Field descriptions are defined inline using Field() descriptors.
     """
 
-    pod_settings: Optional[KubernetesPodSettings] = None
-    service_account_name: Optional[str] = None
-    privileged: bool = False
-    pod_startup_timeout: int = 60 * 10  # Default 10 minutes
-    pod_failure_max_retries: int = 3
-    pod_failure_retry_delay: int = 10
-    pod_failure_backoff: float = 1.0
+    pod_settings: Optional[KubernetesPodSettings] = Field(
+        default=None,
+        description="Pod configuration for step execution containers.",
+    )
+    service_account_name: Optional[str] = Field(
+        default=None,
+        description="Kubernetes service account for step pods. Uses default account if not specified.",
+    )
+    privileged: bool = Field(
+        default=False,
+        description="Whether to run step containers in privileged mode with extended permissions.",
+    )
+    pod_startup_timeout: int = Field(
+        default=600,
+        description="Maximum seconds to wait for step pods to start. Default is 10 minutes.",
+    )
+    pod_failure_max_retries: int = Field(
+        default=3,
+        description="Maximum retry attempts when step pods fail to start.",
+    )
+    pod_failure_retry_delay: int = Field(
+        default=10,
+        description="Delay in seconds between pod failure retry attempts.",
+    )
+    pod_failure_backoff: float = Field(
+        default=1.0,
+        description="Exponential backoff factor for retry delays. Values > 1.0 increase delay with each retry.",
+    )
 
 
 class KubernetesStepOperatorConfig(
@@ -59,22 +72,24 @@ class KubernetesStepOperatorConfig(
 ):
     """Configuration for the Kubernetes step operator.
 
-    Attributes:
-        kubernetes_namespace: Name of the Kubernetes namespace to be used.
-        incluster: If `True`, the step operator will run the pipeline inside the
-            same cluster in which the orchestrator is running. For this to work,
-            the pod running the orchestrator needs permissions to create new
-            pods. If set, the `kubernetes_context` config option is ignored. If
-            the stack component is linked to a Kubernetes service connector,
-            this field is ignored.
-        kubernetes_context: Name of a Kubernetes context to run pipelines in.
-            If the stack component is linked to a Kubernetes service connector,
-            this field is ignored. Otherwise, it is mandatory.
+    Defines cluster connection and execution settings.
+    Field descriptions are defined inline using Field() descriptors.
     """
 
-    kubernetes_namespace: str = "zenml"
-    incluster: bool = False
-    kubernetes_context: Optional[str] = None
+    kubernetes_namespace: str = Field(
+        default="zenml",
+        description="Kubernetes namespace for step execution. Must be a valid namespace name.",
+    )
+    incluster: bool = Field(
+        default=False,
+        description="Whether to execute within the same cluster as the orchestrator. "
+        "Requires appropriate pod creation permissions.",
+    )
+    kubernetes_context: Optional[str] = Field(
+        default=None,
+        description="Kubernetes context name for cluster connection. "
+        "Ignored when using service connectors or in-cluster execution.",
+    )
 
     @property
     def is_remote(self) -> bool:

zenml/integrations/kubernetes/orchestrators/kubernetes_orchestrator.py

@@ -49,7 +49,7 @@ from zenml.config.base_settings import BaseSettings
 from zenml.constants import (
     METADATA_ORCHESTRATOR_RUN_ID,
 )
-from zenml.enums import StackComponentType
+from zenml.enums import ExecutionStatus, StackComponentType
 from zenml.integrations.kubernetes.flavors.kubernetes_orchestrator_flavor import (
     KubernetesOrchestratorConfig,
     KubernetesOrchestratorSettings,
@@ -785,6 +785,187 @@ class KubernetesOrchestrator(ContainerizedOrchestrator):
                 f"No running step jobs found for pipeline run with ID: {run.id}"
             )
 
+    def fetch_status(
+        self, run: "PipelineRunResponse", include_steps: bool = False
+    ) -> Tuple[
+        Optional[ExecutionStatus], Optional[Dict[str, ExecutionStatus]]
+    ]:
+        """Refreshes the status of a specific pipeline run.
+
+        Args:
+            run: The run that was executed by this orchestrator.
+            include_steps: If True, also fetch the status of individual steps.
+
+        Returns:
+            A tuple of (pipeline_status, step_statuses).
+            If include_steps is False, step_statuses will be None.
+            If include_steps is True, step_statuses will be a dict (possibly empty).
+
+        Raises:
+            ValueError: If the orchestrator run ID cannot be found or if the
+                stack components are not accessible.
+        """
+        # Get the orchestrator run ID which corresponds to the orchestrator pod name
+        orchestrator_run_id = run.orchestrator_run_id
+        if not orchestrator_run_id:
+            raise ValueError(
+                "Cannot determine orchestrator run ID for the run. "
+                "Unable to fetch the status."
+            )
+
+        # Check the orchestrator pod status (only if run is not finished)
+        if not run.status.is_finished:
+            orchestrator_pod_phase = self._check_pod_status(
+                pod_name=orchestrator_run_id,
+            )
+            pipeline_status = self._map_pod_phase_to_execution_status(
+                orchestrator_pod_phase
+            )
+        else:
+            # Run is already finished, don't change status
+            pipeline_status = None
+
+        step_statuses = None
+        if include_steps:
+            step_statuses = self._fetch_step_statuses(run)
+
+        return pipeline_status, step_statuses
+
+    def _check_pod_status(
+        self,
+        pod_name: str,
+    ) -> kube_utils.PodPhase:
+        """Check pod status and handle deletion scenarios for both orchestrator and step pods.
+
+        This method should only be called for non-finished pipeline runs/steps.
+
+        Args:
+            pod_name: The name of the pod to check.
+
+        Returns:
+            The pod phase if the pod exists, or PodPhase.FAILED if pod was deleted.
+        """
+        pod = kube_utils.get_pod(
+            core_api=self._k8s_core_api,
+            pod_name=pod_name,
+            namespace=self.config.kubernetes_namespace,
+        )
+
+        if pod and pod.status and pod.status.phase:
+            try:
+                return kube_utils.PodPhase(pod.status.phase)
+            except ValueError:
+                # Handle unknown pod phases
+                logger.warning(
+                    f"Unknown pod phase for pod {pod_name}: {pod.status.phase}"
+                )
+                return kube_utils.PodPhase.UNKNOWN
+        else:
+            logger.warning(
+                f"Can't fetch the status of pod {pod_name} "
+                f"in namespace {self.config.kubernetes_namespace}."
+            )
+            return kube_utils.PodPhase.UNKNOWN
+
+    def _map_pod_phase_to_execution_status(
+        self, pod_phase: kube_utils.PodPhase
+    ) -> Optional[ExecutionStatus]:
+        """Map Kubernetes pod phase to ZenML execution status.
+
+        Args:
+            pod_phase: The Kubernetes pod phase.
+
+        Returns:
+            The corresponding ZenML execution status.
+        """
+        if pod_phase == kube_utils.PodPhase.PENDING:
+            return ExecutionStatus.INITIALIZING
+        elif pod_phase == kube_utils.PodPhase.RUNNING:
+            return ExecutionStatus.RUNNING
+        elif pod_phase == kube_utils.PodPhase.SUCCEEDED:
+            return ExecutionStatus.COMPLETED
+        elif pod_phase == kube_utils.PodPhase.FAILED:
+            return ExecutionStatus.FAILED
+        else:  # UNKNOWN - no update
+            return None
+
+    def _map_job_status_to_execution_status(
+        self, job: k8s_client.V1Job
+    ) -> Optional[ExecutionStatus]:
+        """Map Kubernetes job status to ZenML execution status.
+
+        Args:
+            job: The Kubernetes job.
+
+        Returns:
+            The corresponding ZenML execution status, or None if no clear status.
+        """
+        # Check job conditions first
+        if job.status and job.status.conditions:
+            for condition in job.status.conditions:
+                if condition.type == "Complete" and condition.status == "True":
+                    return ExecutionStatus.COMPLETED
+                elif condition.type == "Failed" and condition.status == "True":
+                    return ExecutionStatus.FAILED
+
+        # Return None if no clear status - don't update
+        return None
+
+    def _fetch_step_statuses(
+        self, run: "PipelineRunResponse"
+    ) -> Dict[str, ExecutionStatus]:
+        """Fetch the statuses of individual pipeline steps.
+
+        Args:
+            run: The pipeline run response.
+
+        Returns:
+            A dictionary mapping step names to their execution statuses.
+        """
+        step_statuses = {}
+
+        # Query all jobs for this run and match them to steps
+        label_selector = f"run_id={kube_utils.sanitize_label(str(run.id))}"
+
+        try:
+            jobs = self._k8s_batch_api.list_namespaced_job(
+                namespace=self.config.kubernetes_namespace,
+                label_selector=label_selector,
+            )
+        except Exception as e:
+            logger.warning(f"Failed to list jobs for run {run.id}: {e}")
+            return {}
+
+        # Fetch the steps from the run response
+        steps_dict = run.steps
+
+        for job in jobs.items:
+            # Extract step name from job labels
+            if not job.metadata or not job.metadata.labels:
+                continue
+
+            step_name = job.metadata.labels.get("step_name")
+            if not step_name:
+                continue
+
+            # Check if this step is already finished
+            step_response = steps_dict.get(step_name, None)
+
+            # If the step is not in the run response yet, skip, we can't update
+            if step_response is None:
+                continue
+
+            # If the step is already in a finished state, skip
+            if step_response and step_response.status.is_finished:
+                continue
+
+            # Check job status and map to execution status
+            execution_status = self._map_job_status_to_execution_status(job)
+            if execution_status is not None:
+                step_statuses[step_name] = execution_status
+
+        return step_statuses
+
     def get_pipeline_run_metadata(
         self, run_id: UUID
     ) -> Dict[str, "MetadataType"]:

zenml/integrations/kubernetes/orchestrators/kubernetes_orchestrator_entrypoint.py

@@ -187,7 +187,7 @@ def main() -> None:
 
                 return True
 
-        step_pod_labels = {
+        base_labels = {
             "run_id": kube_utils.sanitize_label(str(pipeline_run.id)),
             "run_name": kube_utils.sanitize_label(str(pipeline_run.name)),
             "pipeline": kube_utils.sanitize_label(
@@ -234,6 +234,10 @@ def main() -> None:
                 pod_name, namespace=namespace
             )
 
+            # Add step name to labels so both pod and job have consistent labeling
+            step_labels = base_labels.copy()
+            step_labels["step_name"] = kube_utils.sanitize_label(step_name)
+
             image = KubernetesOrchestrator.get_image(
                 deployment=deployment, step_name=step_name
             )
@@ -281,7 +285,7 @@ def main() -> None:
                 or settings.service_account_name,
                 mount_local_stores=mount_local_stores,
                 termination_grace_period_seconds=settings.pod_stop_grace_period,
-                labels=step_pod_labels,
+                labels=step_labels,
             )
 
             retry_config = step_config.retry
@@ -347,7 +351,7 @@ def main() -> None:
                 active_deadline_seconds=settings.active_deadline_seconds,
                 pod_failure_policy=pod_failure_policy,
                 owner_references=owner_references,
-                labels=step_pod_labels,
+                labels=step_labels,
             )
 
             kube_utils.create_job(