zenml-nightly 0.81.0.dev20250428__py3-none-any.whl → 0.81.0.dev20250430__py3-none-any.whl

zenml/VERSION

@@ -1 +1 @@
-0.81.0.dev20250428
+0.81.0.dev20250430

zenml/cli/model.py

@@ -398,18 +398,13 @@ def version() -> None:
 
 @cli_utils.list_options(ModelVersionFilter)
 @version.command("list", help="List model versions with filter.")
-@click.argument("model_name_or_id")
-def list_model_versions(model_name_or_id: str, **kwargs: Any) -> None:
+def list_model_versions(**kwargs: Any) -> None:
     """List model versions with filter in the Model Control Plane.
 
     Args:
-        model_name_or_id: The name or ID of the parent model.
         **kwargs: Keyword arguments to filter models.
     """
-    model_versions = Client().list_model_versions(
-        model_name_or_id=model_name_or_id,
-        **kwargs,
-    )
+    model_versions = Client().list_model_versions(**kwargs)
 
     if not model_versions:
         cli_utils.declare("No model versions found.")

zenml/client.py

@@ -6504,7 +6504,8 @@ class Client(metaclass=ClientMetaClass):
 
     def list_model_versions(
         self,
-        model_name_or_id: Union[str, UUID],
+        model: Optional[Union[str, UUID]] = None,
+        model_name_or_id: Optional[Union[str, UUID]] = None,
         sort_by: str = "number",
         page: int = PAGINATION_STARTING_PAGE,
         size: int = PAGE_SIZE_DEFAULT,
@@ -6525,6 +6526,7 @@ class Client(metaclass=ClientMetaClass):
         """Get model versions by filter from Model Control Plane.
 
         Args:
+            model: The model to filter by.
             model_name_or_id: name or id of the model containing the model
                 version.
             sort_by: The column to sort by
@@ -6548,6 +6550,19 @@ class Client(metaclass=ClientMetaClass):
         Returns:
             A page object with all model versions.
         """
+        if model_name_or_id:
+            logger.warning(
+                "The `model_name_or_id` argument is deprecated. "
+                "Please use the `model` argument instead."
+            )
+            if model is None:
+                model = model_name_or_id
+            else:
+                logger.warning(
+                    "Ignoring `model_name_or_id` argument as `model` argument "
+                    "was also provided."
+                )
+
         model_version_filter_model = ModelVersionFilter(
             page=page,
             size=size,
@@ -6563,7 +6578,7 @@ class Client(metaclass=ClientMetaClass):
             tag=tag,
             tags=tags,
             user=user,
-            model=model_name_or_id,
+            model=model,
             project=project or self.active_project.id,
         )
 

zenml/config/build_configuration.py

@@ -95,6 +95,10 @@ class BuildConfiguration(BaseModel):
         for _, requirements, _ in requirements_files:
             hash_.update(requirements.encode())
 
+        if self.settings.dockerfile:
+            with open(self.settings.dockerfile, "rb") as f:
+                hash_.update(f.read())
+
         return hash_.hexdigest()
 
     def should_include_files(

zenml/config/docker_settings.py

@@ -309,5 +309,5 @@ class DockerSettings(BaseSettings):
         # public attributes are immutable
         frozen=True,
         # prevent extra attributes during model initialization
-        extra="forbid",
+        extra="ignore",
     )

zenml/config/resource_settings.py

@@ -119,5 +119,5 @@ class ResourceSettings(BaseSettings):
         # public attributes are immutable
         frozen=True,
         # prevent extra attributes during model initialization
-        extra="forbid",
+        extra="ignore",
     )

zenml/config/strict_base_model.py

@@ -19,4 +19,4 @@ from pydantic import BaseModel, ConfigDict
 class StrictBaseModel(BaseModel):
     """Immutable pydantic model which prevents extra attributes."""
 
-    model_config = ConfigDict(frozen=True, extra="forbid")
+    model_config = ConfigDict(frozen=True, extra="ignore")

zenml/entrypoints/base_entrypoint_configuration.py

@@ -274,7 +274,7 @@ class BaseEntrypointConfiguration(ABC):
         download_dir = os.path.join(
             code_repo_root, code_reference.subdirectory
         )
-        os.makedirs(download_dir)
+        os.makedirs(download_dir, exist_ok=True)
         repo.download_files(
             commit=code_reference.commit,
             directory=download_dir,

zenml/environment.py

@@ -15,12 +15,11 @@
 
 import os
 import platform
-from pathlib import Path
-from typing import Dict
+import subprocess
+from typing import Any, Dict, List
 
 import distro
 
-from zenml import __version__
 from zenml.constants import INSIDE_ZENML_CONTAINER
 from zenml.enums import EnvironmentType
 from zenml.logger import get_logger
@@ -29,7 +28,7 @@ from zenml.utils.singleton import SingletonMetaClass
 logger = get_logger(__name__)
 
 
-def get_run_environment_dict() -> Dict[str, str]:
+def get_run_environment_dict() -> Dict[str, Any]:
     """Returns a dictionary of the current run environment.
 
     Everything that is returned here will be saved in the DB as
@@ -40,12 +39,24 @@ def get_run_environment_dict() -> Dict[str, str]:
     Returns:
         A dictionary of the current run environment.
     """
-    return {
-        "environment": get_environment(),
+    env_dict: Dict[str, Any] = {
+        "environment": str(get_environment()),
         **Environment.get_system_info(),
         "python_version": Environment.python_version(),
     }
 
+    try:
+        python_packages = Environment.get_python_packages()
+    except RuntimeError:
+        logger.warning("Failed to get list of installed Python packages")
+    else:
+        # TODO: We send the python packages as a string right now to keep
+        # backwards compatibility with old versions. We should update this to
+        # be a list of strings eventually.
+        env_dict["python_packages"] = "\n".join(python_packages)
+
+    return env_dict
+
 
 def get_environment() -> str:
     """Returns a string representing the execution environment of the pipeline.
@@ -88,35 +99,8 @@ def get_environment() -> str:
         return EnvironmentType.NATIVE
 
 
-def get_system_details() -> str:
-    """Returns OS, python and ZenML information.
-
-    Returns:
-        str: OS, python and ZenML information
-    """
-    from zenml.integrations.registry import integration_registry
-
-    info = {
-        "ZenML version": __version__,
-        "Install path": Path(__file__).resolve().parent,
-        "Python version": Environment.python_version(),
-        "Platform information": Environment.get_system_info(),
-        "Environment": get_environment(),
-        "Integrations": integration_registry.get_installed_integrations(),
-    }
-    return "\n".join(
-        "{:>10} {}".format(k + ":", str(v).replace("\n", " "))
-        for k, v in info.items()
-    )
-
-
 class Environment(metaclass=SingletonMetaClass):
-    """Provides environment information.
-
-    Individual environment components can be registered separately to extend
-    the global Environment object with additional information (see
-    `BaseEnvironmentComponent`).
-    """
+    """Provides environment information."""
 
     def __init__(self) -> None:
         """Initializes an Environment instance.
@@ -366,3 +350,22 @@ class Environment(metaclass=SingletonMetaClass):
             "LIGHTNING_CLOUD_URL" in os.environ
             and "LIGHTNING_CLOUDSPACE_HOST" in os.environ
         )
+
+    @staticmethod
+    def get_python_packages() -> List[str]:
+        """Returns a list of installed Python packages.
+
+        Raises:
+            RuntimeError: If the process to get the list of installed packages
+                fails.
+
+        Returns:
+            List of installed packages in pip freeze format.
+        """
+        try:
+            output = subprocess.check_output(["pip", "freeze"]).decode()
+            return output.strip().split("\n")
+        except subprocess.CalledProcessError:
+            raise RuntimeError(
+                "Failed to get list of installed Python packages"
+            )

zenml/exceptions.py

@@ -70,40 +70,6 @@ class DoesNotExistException(ZenMLBaseException):
         super().__init__(message)
 
 
-class PipelineNotSucceededException(ZenMLBaseException):
-    """Raises exception when trying to fetch artifacts from a not succeeded pipeline."""
-
-    def __init__(
-        self,
-        name: str = "",
-        message: str = "{} is not yet completed successfully.",
-    ):
-        """Initializes the exception.
-
-        Args:
-            name: Name of the pipeline.
-            message: Message with details of exception.
-        """
-        super().__init__(message.format(name))
-
-
-class GitException(ZenMLBaseException):
-    """Raises exception when a problem occurs in git resolution."""
-
-    def __init__(
-        self,
-        message: str = "There is a problem with git resolution. "
-        "Please make sure that all relevant files "
-        "are committed.",
-    ):
-        """Initializes the exception.
-
-        Args:
-            message: Message with details of exception.
-        """
-        super().__init__(message)
-
-
 class StepInterfaceError(ZenMLBaseException):
     """Raises exception when interacting with the Step interface in an unsupported way."""
 
@@ -116,50 +82,18 @@ class StepContextError(ZenMLBaseException):
     """Raises exception when interacting with a StepContext in an unsupported way."""
 
 
-class PipelineInterfaceError(ZenMLBaseException):
-    """Raises exception when interacting with the Pipeline interface in an unsupported way."""
-
-
-class ArtifactInterfaceError(ZenMLBaseException):
-    """Raises exception when interacting with the Artifact interface in an unsupported way."""
-
-
 class StackComponentInterfaceError(ZenMLBaseException):
     """Raises exception when interacting with the stack components in an unsupported way."""
 
 
-class StackComponentDeploymentError(ZenMLBaseException):
-    """Raises exception when deploying a stack component fails."""
-
-
 class ArtifactStoreInterfaceError(ZenMLBaseException):
     """Raises exception when interacting with the Artifact Store interface in an unsupported way."""
 
 
-class PipelineConfigurationError(ZenMLBaseException):
-    """Raises exceptions when a pipeline configuration contains invalid values."""
-
-
 class IntegrationError(ZenMLBaseException):
     """Raises exceptions when a requested integration can not be activated."""
 
 
-class DuplicateRunNameError(RuntimeError):
-    """Raises exception when a run with the same name already exists."""
-
-    def __init__(
-        self,
-        message: str = "Unable to run a pipeline with a run name that "
-        "already exists.",
-    ):
-        """Initializes the exception.
-
-        Args:
-            message: Message with details of exception.
-        """
-        super().__init__(message)
-
-
 class ValidationError(ZenMLBaseException):
     """Raised when the Model passed to the ZenStore."""
 
@@ -180,22 +114,10 @@ class StackValidationError(ZenMLBaseException):
     """Raised when a stack configuration is not valid."""
 
 
-class StackComponentValidationError(ZenMLBaseException):
-    """Raised when a stack component configuration is not valid."""
-
-
-class ProvisioningError(ZenMLBaseException):
-    """Raised when an error occurs when provisioning resources for a StackComponent."""
-
-
 class GitNotFoundError(ImportError):
     """Raised when ZenML CLI is used to interact with examples on a machine with no git installation."""
 
 
-class DuplicatedConfigurationError(ZenMLBaseException):
-    """Raised when a configuration parameter is set twice."""
-
-
 class IllegalOperationError(ZenMLBaseException):
     """Raised when an illegal operation is attempted."""
 

zenml/integrations/aws/__init__.py

@@ -41,7 +41,7 @@ class AWSIntegration(Integration):
 
     NAME = AWS
     REQUIREMENTS = [
-        "sagemaker>=2.199.0",
+        "sagemaker>=2.237.3",
         "kubernetes",
         "aws-profile-manager",
     ]

zenml/integrations/evidently/__init__.py

@@ -56,6 +56,7 @@ class EvidentlyIntegration(Integration):
     REQUIREMENTS = [
         "evidently>=0.4.16,<=0.4.22",
         "tenacity!=8.4.0",  # https://github.com/jd/tenacity/issues/471
+        "numpy<2.0.0",  # evidently is not compatible with NumPy 2.0
     ]
     REQUIREMENTS_IGNORED_ON_UNINSTALL = ["tenacity", "pandas"]
 

zenml/integrations/evidently/metrics.py

@@ -361,4 +361,4 @@ class EvidentlyMetricConfig(BaseModel):
                 f"`{self.class_path}`: {str(e)}"
             )
 
-    model_config = ConfigDict(extra="forbid")
+    model_config = ConfigDict(extra="ignore")

zenml/integrations/evidently/tests.py

@@ -345,4 +345,4 @@ class EvidentlyTestConfig(BaseModel):
                 f"`{self.class_path}`: {str(e)}"
             )
 
-    model_config = ConfigDict(extra="forbid")
+    model_config = ConfigDict(extra="ignore")

zenml/integrations/gcp/flavors/vertex_orchestrator_flavor.py

@@ -112,6 +112,11 @@ class VertexOrchestratorConfig(
         network: the full name of the Compute Engine Network to which the job
             should be peered. For example, `projects/12345/global/networks/myVPC`
             If not provided, the job will not be peered with any network.
+        private_service_connect: the full name of a Private Service Connect
+            endpoint to which the job should be peered. For example,
+            `projects/12345/regions/us-central1/networkAttachments/NETWORK_ATTACHMENT_NAME`
+            If not provided, the job will not be peered with any private service
+            connect endpoint.
         cpu_limit: The maximum CPU limit for this operator. This string value
             can be a number (integer value for number of CPUs) as string,
             or a number followed by "m", which means 1/1000. You can specify
@@ -130,6 +135,7 @@ class VertexOrchestratorConfig(
     encryption_spec_key_name: Optional[str] = None
     workload_service_account: Optional[str] = None
     network: Optional[str] = None
+    private_service_connect: Optional[str] = None
 
     # Deprecated
     cpu_limit: Optional[str] = None

zenml/integrations/gcp/orchestrators/vertex_orchestrator.py

@@ -48,7 +48,15 @@ from uuid import UUID
 
 from google.api_core import exceptions as google_exceptions
 from google.cloud import aiplatform
+from google.cloud.aiplatform.compat.services import (
+    pipeline_service_client_v1beta1,
+)
+from google.cloud.aiplatform.compat.types import pipeline_job_v1beta1
 from google.cloud.aiplatform_v1.types import PipelineState
+from google.cloud.aiplatform_v1beta1.types.service_networking import (
+    PscInterfaceConfig,
+)
+from google.protobuf import json_format
 from google_cloud_pipeline_components.v1.custom_job.utils import (
     create_custom_training_job_from_component,
 )
@@ -495,36 +503,16 @@ class VertexOrchestrator(ContainerizedOrchestrator, GoogleCredentialsMixin):
                 )
                 pod_settings = step_settings.pod_settings
                 if pod_settings:
-                    if pod_settings.host_ipc:
-                        logger.warning(
-                            "Host IPC is set to `True` but not supported in "
-                            "this orchestrator. Ignoring..."
-                        )
-                    if pod_settings.affinity:
-                        logger.warning(
-                            "Affinity is set but not supported in Vertex with "
-                            "Kubeflow Pipelines 2.x. Ignoring..."
-                        )
-                    if pod_settings.tolerations:
-                        logger.warning(
-                            "Tolerations are set but not supported in "
-                            "Vertex with Kubeflow Pipelines 2.x. Ignoring..."
-                        )
-                    if pod_settings.volumes:
-                        logger.warning(
-                            "Volumes are set but not supported in Vertex with "
-                            "Kubeflow Pipelines 2.x. Ignoring..."
-                        )
-                    if pod_settings.volume_mounts:
-                        logger.warning(
-                            "Volume mounts are set but not supported in "
-                            "Vertex with Kubeflow Pipelines 2.x. Ignoring..."
-                        )
-                    if pod_settings.env or pod_settings.env_from:
+                    ignored_fields = pod_settings.model_fields_set - {
+                        "node_selectors"
+                    }
+                    if ignored_fields:
                         logger.warning(
-                            "Environment variables are set but not supported "
-                            "in Vertex with Vertex Pipelines 2.x. Ignoring..."
+                            f"The following pod settings are not supported in "
+                            f"Vertex with Vertex Pipelines 2.x and will be "
+                            f"ignored: {list(ignored_fields)}."
                         )
+
                     for key in pod_settings.node_selectors:
                         if (
                             key
@@ -737,6 +725,39 @@ class VertexOrchestrator(ContainerizedOrchestrator, GoogleCredentialsMixin):
                 self.config.network,
             )
 
+        if self.config.private_service_connect:
+            # The PSC setting isn't yet part of the stable v1 API. We need to
+            # temporarily hack the aiplatform.PipelineJob object in two places:
+            # * to use the v1beta1 PipelineJob primitive which supports the
+            #   psc_interface_config field instead of the v1 PipelineJob
+            #   primitive.
+            # * to use the v1beta1 PipelineServiceClient instead of the v1
+            #   PipelineServiceClient.
+            #
+            # We achieve the first by converting the v1 PipelineJob to a
+            # v1beta1 PipelineJob and the second by replacing the v1
+            # PipelineServiceClient with a v1beta1 PipelineServiceClient.
+            #
+            # TODO: Remove this once the v1 stable API is updated to support
+            # the PSC setting.
+            pipeline_job_dict = json_format.MessageToDict(
+                run._gca_resource._pb, preserving_proto_field_name=True
+            )
+            run._gca_resource = pipeline_job_v1beta1.PipelineJob(
+                **pipeline_job_dict,
+                psc_interface_config=PscInterfaceConfig(
+                    network_attachment=self.config.private_service_connect,
+                ),
+            )
+            run.api_client = (
+                pipeline_service_client_v1beta1.PipelineServiceClient(
+                    credentials=run.credentials,
+                    client_options={
+                        "api_endpoint": run.api_client.api_endpoint,
+                    },
+                )
+            )
+
         try:
             if schedule:
                 logger.info(

zenml/integrations/kubeflow/orchestrators/kubeflow_orchestrator.py

@@ -555,36 +555,14 @@ class KubeflowOrchestrator(ContainerizedOrchestrator):
                 )
                 pod_settings = step_settings.pod_settings
                 if pod_settings:
-                    if pod_settings.host_ipc:
+                    ignored_fields = pod_settings.model_fields_set - {
+                        "node_selectors"
+                    }
+                    if ignored_fields:
                         logger.warning(
-                            "Host IPC is set to `True` but not supported in "
-                            "this orchestrator. Ignoring..."
-                        )
-                    if pod_settings.affinity:
-                        logger.warning(
-                            "Affinity is set but not supported in Kubeflow with "
-                            "Kubeflow Pipelines 2.x. Ignoring..."
-                        )
-                    if pod_settings.tolerations:
-                        logger.warning(
-                            "Tolerations are set but not supported in "
-                            "Kubeflow with Kubeflow Pipelines 2.x. Ignoring..."
-                        )
-                    if pod_settings.volumes:
-                        logger.warning(
-                            "Volumes are set but not supported in Kubeflow with "
-                            "Kubeflow Pipelines 2.x. Ignoring..."
-                        )
-                    if pod_settings.volume_mounts:
-                        logger.warning(
-                            "Volume mounts are set but not supported in "
-                            "Kubeflow with Kubeflow Pipelines 2.x. Ignoring..."
-                        )
-                    if pod_settings.env or pod_settings.env_from:
-                        logger.warning(
-                            "Environment variables are set but not supported "
-                            "in Kubeflow with Kubeflow Pipelines 2.x. "
-                            "Ignoring..."
+                            f"The following pod settings are not supported in "
+                            f"Kubeflow with Kubeflow Pipelines 2.x and will be "
+                            f"ignored: {list(ignored_fields)}."
                         )
 
                     # apply pod settings

zenml/integrations/kubernetes/flavors/kubernetes_orchestrator_flavor.py

@@ -15,6 +15,8 @@
 
 from typing import TYPE_CHECKING, Optional, Type
 
+from pydantic import PositiveInt
+
 from zenml.config.base_settings import BaseSettings
 from zenml.constants import KUBERNETES_CLUSTER_RESOURCE_TYPE
 from zenml.integrations.kubernetes import KUBERNETES_ORCHESTRATOR_FLAVOR
@@ -47,6 +49,7 @@ class KubernetesOrchestratorSettings(BaseSettings):
         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
@@ -55,6 +58,7 @@ class KubernetesOrchestratorSettings(BaseSettings):
             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.
     """
 
     synchronous: bool = True
@@ -64,10 +68,12 @@ class KubernetesOrchestratorSettings(BaseSettings):
     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
 
 
 class KubernetesOrchestratorConfig(

zenml/integrations/kubernetes/orchestrators/kubernetes_orchestrator.py

@@ -417,6 +417,9 @@ class KubernetesOrchestrator(ContainerizedOrchestrator):
                 )
 
         pipeline_name = deployment.pipeline_configuration.name
+        settings = cast(
+            KubernetesOrchestratorSettings, self.get_settings(deployment)
+        )
 
         # We already make sure the orchestrator run name has the correct length
         # to make sure we don't cut off the randomized suffix later when
@@ -427,8 +430,16 @@ class KubernetesOrchestrator(ContainerizedOrchestrator):
         orchestrator_run_name = get_orchestrator_run_name(
             pipeline_name, max_length=max_length
         )
+
+        if settings.pod_name_prefix:
+            pod_name = get_orchestrator_run_name(
+                settings.pod_name_prefix, max_length=max_length
+            )
+        else:
+            pod_name = orchestrator_run_name
+
         pod_name = kube_utils.sanitize_pod_name(
-            orchestrator_run_name, namespace=self.config.kubernetes_namespace
+            pod_name, namespace=self.config.kubernetes_namespace
         )
 
         assert stack.container_registry
@@ -455,10 +466,6 @@ class KubernetesOrchestrator(ContainerizedOrchestrator):
             run_id=placeholder_run.id if placeholder_run else None,
         )
 
-        settings = cast(
-            KubernetesOrchestratorSettings, self.get_settings(deployment)
-        )
-
         # Authorize pod to run Kubernetes commands inside the cluster.
         service_account_name = self._get_service_account_name(settings)