zenml-nightly 0.84.0.dev20250728__py3-none-any.whl → 0.84.0.dev20250730__py3-none-any.whl

zenml/VERSION

@@ -1 +1 @@
-0.84.0.dev20250728
+0.84.0.dev20250730

zenml/client.py

@@ -7339,19 +7339,24 @@ class Client(metaclass=ClientMetaClass):
     def create_service_account(
         self,
         name: str,
+        full_name: Optional[str] = None,
         description: str = "",
     ) -> ServiceAccountResponse:
         """Create a new service account.
 
         Args:
             name: The name of the service account.
+            full_name: The display name of the service account.
             description: The description of the service account.
 
         Returns:
             The created service account.
         """
         service_account = ServiceAccountRequest(
-            name=name, description=description, active=True
+            name=name,
+            full_name=full_name or "",
+            description=description,
+            active=True,
         )
         created_service_account = self.zen_store.create_service_account(
             service_account=service_account

zenml/code_repositories/git/local_git_repository_context.py

@@ -13,7 +13,7 @@
 #  permissions and limitations under the License.
 """Implementation of the Local git repository context."""
 
-from typing import TYPE_CHECKING, Callable, Optional, cast
+from typing import TYPE_CHECKING, Callable, Optional
 
 from zenml.code_repositories import (
     LocalRepositoryContext,
@@ -25,7 +25,6 @@ from zenml.constants import (
 from zenml.logger import get_logger
 
 if TYPE_CHECKING:
-    from git.objects import Commit
     from git.remote import Remote
     from git.repo.base import Repo
 
@@ -180,7 +179,7 @@ class LocalGitRepositoryContext(LocalRepositoryContext):
             # Branch doesn't exist on remote
             return True
 
-        return cast("Commit", remote_commit_object) != local_commit_object
+        return remote_commit_object != local_commit_object
 
     @property
     def current_commit(self) -> str:
@@ -189,4 +188,4 @@ class LocalGitRepositoryContext(LocalRepositoryContext):
         Returns:
             The current commit sha.
         """
-        return cast(str, self.git_repo.head.object.hexsha)
+        return self.git_repo.head.object.hexsha

zenml/constants.py

@@ -174,6 +174,12 @@ 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_LOGS_STORAGE_MAX_QUEUE_SIZE = "ZENML_LOGS_STORAGE_MAX_QUEUE_SIZE"
+ENV_ZENML_LOGS_STORAGE_QUEUE_TIMEOUT = "ZENML_LOGS_STORAGE_QUEUE_TIMEOUT"
+
+ENV_ZENML_LOGS_WRITE_INTERVAL_SECONDS = "ZENML_LOGS_WRITE_INTERVAL_SECONDS"
+ENV_ZENML_LOGS_MERGE_INTERVAL_SECONDS = "ZENML_LOGS_MERGE_INTERVAL_SECONDS"
+
 ENV_ZENML_CUSTOM_SOURCE_ROOT = "ZENML_CUSTOM_SOURCE_ROOT"
 ENV_ZENML_PIPELINE_RUN_API_TOKEN_EXPIRATION = (
     "ZENML_PIPELINE_API_TOKEN_EXPIRATION"
@@ -181,6 +187,9 @@ ENV_ZENML_PIPELINE_RUN_API_TOKEN_EXPIRATION = (
 ENV_ZENML_CODE_REPOSITORY_IGNORE_UNTRACKED_FILES = (
     "ZENML_CODE_REPOSITORY_IGNORE_UNTRACKED_FILES"
 )
+# Environment variable that indicates whether the current environment is running
+# a step operator.
+ENV_ZENML_STEP_OPERATOR = "ZENML_STEP_OPERATOR"
 
 # Materializer environment variables
 ENV_ZENML_MATERIALIZER_ALLOW_NON_ASCII_JSON_DUMPS = (
@@ -501,3 +510,22 @@ STACK_DEPLOYMENT_API_TOKEN_EXPIRATION = 60 * 6  # 6 hours
 ZENML_PIPELINE_RUN_API_TOKEN_EXPIRATION = handle_int_env_var(
     ENV_ZENML_PIPELINE_RUN_API_TOKEN_EXPIRATION, default=0
 )
+
+# Logs storage constants
+# Maximum number of log batches to queue.
+LOGS_STORAGE_MAX_QUEUE_SIZE = handle_int_env_var(
+    ENV_ZENML_LOGS_STORAGE_MAX_QUEUE_SIZE, default=100000
+)
+# Queue timeout controls log dropping vs application blocking:
+# - Positive value (e.g., 30): Wait N seconds, then drop logs if queue full
+# - Negative value (e.g., -1): Block indefinitely until queue has space (never drop logs)
+LOGS_STORAGE_QUEUE_TIMEOUT = handle_int_env_var(
+    ENV_ZENML_LOGS_STORAGE_QUEUE_TIMEOUT, default=-1
+)
+# Log batch-write and merge windows in seconds.
+LOGS_WRITE_INTERVAL_SECONDS = handle_int_env_var(
+    ENV_ZENML_LOGS_WRITE_INTERVAL_SECONDS, default=2
+)
+LOGS_MERGE_INTERVAL_SECONDS = handle_int_env_var(
+    ENV_ZENML_LOGS_MERGE_INTERVAL_SECONDS, default=10 * 60
+)

zenml/integrations/azure/service_connectors/azure_service_connector.py

@@ -1490,7 +1490,8 @@ class AzureServiceConnector(ServiceConnector):
                 aks_clusters = [
                     cluster
                     for cluster in aks_clusters
-                    if self._get_resource_group(cluster.id)
+                    if cluster.id
+                    and self._get_resource_group(cluster.id)
                     == self.config.resource_group
                 ]
 
@@ -1507,7 +1508,8 @@ class AzureServiceConnector(ServiceConnector):
             clusters = [
                 (cluster.name, self._get_resource_group(cluster.id))
                 for cluster in aks_clusters
-                if cluster.name
+                if cluster.id
+                and cluster.name
                 and (not cluster_name or cluster.name == cluster_name)
             ]
 
@@ -1809,16 +1811,25 @@ class AzureServiceConnector(ServiceConnector):
             cluster_name, resource_group = clusters[0]
 
             try:
-                client = ContainerServiceClient(credential, subscription_id)
+                cs_client = ContainerServiceClient(credential, subscription_id)
 
-                creds = client.managed_clusters.list_cluster_admin_credentials(
-                    resource_group_name=resource_group,
-                    resource_name=cluster_name,
+                creds = (
+                    cs_client.managed_clusters.list_cluster_admin_credentials(
+                        resource_group_name=resource_group,
+                        resource_name=cluster_name,
+                    )
                 )
 
-                kubeconfig_yaml = creds.kubeconfigs[0].value.decode(
-                    encoding="UTF-8"
-                )
+                if creds.kubeconfigs and creds.kubeconfigs[0].value:
+                    kubeconfig_yaml = creds.kubeconfigs[0].value.decode(
+                        encoding="UTF-8"
+                    )
+                else:
+                    raise AuthorizationException(
+                        f"failed to list credentials for Azure Kubernetes "
+                        f"Service cluster '{cluster_name}' in resource group "
+                        f"'{resource_group}': no kubeconfig found"
+                    )
             except AzureError as e:
                 raise AuthorizationException(
                     f"failed to list credentials for Azure Kubernetes "

zenml/integrations/gcp/__init__.py

@@ -43,9 +43,12 @@ class GcpIntegration(Integration):
     """Definition of Google Cloud Platform integration for ZenML."""
 
     NAME = GCP
+    # Adding the gcsfs<=2024.12.0 for now to solve the issue with the 
+    # increased number of list calls. This is a temporary fix and should be
+    # removed once the issue is resolved.
     REQUIREMENTS = [
         "kfp>=2.6.0",
-        "gcsfs!=2025.5.0,!=2025.5.0.post1",
+        "gcsfs!=2025.5.0,!=2025.5.0.post1,<=2024.12.0",
         "google-cloud-secret-manager",
         "google-cloud-container>=2.21.0",
         "google-cloud-artifact-registry>=1.11.3",

zenml/integrations/kubernetes/flavors/kubernetes_orchestrator_flavor.py

@@ -13,7 +13,7 @@
 #  permissions and limitations under the License.
 """Kubernetes orchestrator flavor."""
 
-from typing import TYPE_CHECKING, Any, Dict, Optional, Type
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Type
 
 from pydantic import Field, NonNegativeInt, PositiveInt, field_validator
 
@@ -131,6 +131,25 @@ class KubernetesOrchestratorSettings(BaseSettings):
         "the chance of the server receiving the maximum amount of retry "
         "requests.",
     )
+    fail_on_container_waiting_reasons: Optional[List[str]] = Field(
+        default=[
+            "InvalidImageName",
+            "ErrImagePull",
+            "ImagePullBackOff",
+            "CreateContainerConfigError",
+        ],
+        description="List of container waiting reasons that should cause the "
+        "job to fail immediately. This should be set to a list of "
+        "nonrecoverable reasons, which if found in any "
+        "`pod.status.containerStatuses[*].state.waiting.reason` of a job pod, "
+        "should cause the job to fail immediately.",
+    )
+    job_monitoring_interval: int = Field(
+        default=3,
+        description="The interval in seconds to monitor the job. Each interval "
+        "is used to check for container issues and streaming logs for the "
+        "job pods.",
+    )
     pod_failure_policy: Optional[Dict[str, Any]] = Field(
         default=None,
         description="The pod failure policy to use for the job that is "

zenml/integrations/kubernetes/orchestrators/kube_utils.py

@@ -37,7 +37,7 @@ import json
 import re
 import time
 from collections import defaultdict
-from typing import Any, Callable, Dict, List, Optional, TypeVar, cast
+from typing import Any, Callable, Dict, List, Optional, Tuple, TypeVar, cast
 
 from kubernetes import client as k8s_client
 from kubernetes import config as k8s_config
@@ -657,6 +657,50 @@ def create_job(
     )
 
 
+def get_container_status(
+    pod: k8s_client.V1Pod, container_name: str
+) -> Optional[k8s_client.V1ContainerState]:
+    """Get the status of a container.
+
+    Args:
+        pod: The pod to get the container status for.
+        container_name: The container name.
+
+    Returns:
+        The container status.
+    """
+    if not pod.status or not pod.status.container_statuses:
+        return None
+
+    for container_status in pod.status.container_statuses:
+        if container_status.name == container_name:
+            return container_status.state
+
+    return None
+
+
+def get_container_termination_reason(
+    pod: k8s_client.V1Pod, container_name: str
+) -> Optional[Tuple[int, str]]:
+    """Get the termination reason for a container.
+
+    Args:
+        pod: The pod to get the termination reason for.
+        container_name: The container name.
+
+    Returns:
+        The exit code and termination reason for the container.
+    """
+    container_state = get_container_status(pod, container_name)
+    if not container_state or not container_state.terminated:
+        return None
+
+    return (
+        container_state.terminated.exit_code,
+        container_state.terminated.reason or "Unknown",
+    )
+
+
 def wait_for_job_to_finish(
     batch_api: k8s_client.BatchV1Api,
     core_api: k8s_client.CoreV1Api,
@@ -665,8 +709,9 @@ def wait_for_job_to_finish(
     backoff_interval: float = 1,
     maximum_backoff: float = 32,
     exponential_backoff: bool = False,
-    container_name: Optional[str] = None,
+    fail_on_container_waiting_reasons: Optional[List[str]] = None,
     stream_logs: bool = True,
+    container_name: Optional[str] = None,
 ) -> None:
     """Wait for a job to finish.
 
@@ -679,6 +724,8 @@ def wait_for_job_to_finish(
         maximum_backoff: The maximum interval to wait between polling the job
             status.
         exponential_backoff: Whether to use exponential backoff.
+        fail_on_container_waiting_reasons: List of container waiting reasons
+            that will cause the job to fail.
         stream_logs: Whether to stream the job logs.
         container_name: Name of the container to stream logs from.
 
@@ -703,9 +750,39 @@ def wait_for_job_to_finish(
                         f"{condition.message}"
                     )
 
+        if fail_on_container_waiting_reasons:
+            pod_list: k8s_client.V1PodList = retry_on_api_exception(
+                core_api.list_namespaced_pod
+            )(
+                namespace=namespace,
+                label_selector=f"job-name={job_name}",
+                field_selector="status.phase=Pending",
+            )
+            for pod in pod_list.items:
+                container_state = get_container_status(
+                    pod, container_name or "main"
+                )
+
+                if (
+                    container_state
+                    and (waiting_state := container_state.waiting)
+                    and waiting_state.reason
+                    in fail_on_container_waiting_reasons
+                ):
+                    retry_on_api_exception(batch_api.delete_namespaced_job)(
+                        name=job_name,
+                        namespace=namespace,
+                        propagation_policy="Foreground",
+                    )
+                    raise RuntimeError(
+                        f"Job `{namespace}:{job_name}` failed: "
+                        f"Detected container in state "
+                        f"{waiting_state.reason}"
+                    )
+
         if stream_logs:
             try:
-                pod_list: k8s_client.V1PodList = core_api.list_namespaced_pod(
+                pod_list = core_api.list_namespaced_pod(
                     namespace=namespace,
                     label_selector=f"job-name={job_name}",
                 )

zenml/integrations/kubernetes/orchestrators/kubernetes_orchestrator_entrypoint.py

@@ -367,33 +367,37 @@ def main() -> None:
                     core_api=core_api,
                     namespace=namespace,
                     job_name=job_name,
+                    fail_on_container_waiting_reasons=settings.fail_on_container_waiting_reasons,
                     stream_logs=pipeline_settings.stream_step_logs,
+                    backoff_interval=settings.job_monitoring_interval,
                 )
 
                 logger.info(f"Job for step `{step_name}` completed.")
             except Exception:
-                reason = ""
+                reason = "Unknown"
                 try:
-                    pods_by_job = core_api.list_namespaced_pod(
+                    pods = core_api.list_namespaced_pod(
                         label_selector=f"job-name={job_name}",
                         namespace=namespace,
                     ).items
-                    first_pod = pods_by_job[0]
-                    if status := first_pod.status:
-                        if container_statuses := status.container_statuses:
-                            for cs in container_statuses:
-                                if cs.name == "main":
-                                    if state := cs.state:
-                                        if terminated := state.terminated:
-                                            if terminated.exit_code != 0:
-                                                reason = f"{terminated.reason}(exit_code={terminated.exit_code})"
-                                    break
+                    # Sort pods by creation timestamp, oldest first
+                    pods.sort(
+                        key=lambda pod: pod.metadata.creation_timestamp,
+                    )
+                    if pods:
+                        if (
+                            termination_reason
+                            := kube_utils.get_container_termination_reason(
+                                pods[-1], "main"
+                            )
+                        ):
+                            exit_code, reason = termination_reason
+                            if exit_code != 0:
+                                reason = f"{reason} (exit_code={exit_code})"
                 except Exception:
                     pass
                 logger.error(
                     f"Job for step `{step_name}` failed. Reason: {reason}"
-                    if reason
-                    else ""
                 )
 
                 raise

zenml/logging/__init__.py

@@ -11,14 +11,3 @@
 #  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
 #  or implied. See the License for the specific language governing
 #  permissions and limitations under the License.
-
-"""Logging utilities."""
-
-# How many seconds to wait before uploading logs to the artifact store
-STEP_LOGS_STORAGE_INTERVAL_SECONDS: int = 15
-
-# How many messages to buffer before uploading logs to the artifact store
-STEP_LOGS_STORAGE_MAX_MESSAGES: int = 100
-
-# How often to merge logs into a single file
-STEP_LOGS_STORAGE_MERGE_INTERVAL_SECONDS: int = 10 * 60