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

zenml/integrations/gcp/orchestrators/vertex_orchestrator.py

@@ -893,14 +893,20 @@ class VertexOrchestrator(ContainerizedOrchestrator, GoogleCredentialsMixin):
 
         return dynamic_component
 
-    def fetch_status(self, run: "PipelineRunResponse") -> ExecutionStatus:
+    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: Whether to fetch steps.
 
         Returns:
-            the actual status of the pipeline job.
+            A tuple of (pipeline_status, step_statuses_dict).
+            Step statuses are not supported for Vertex, so step_statuses_dict will always be None.
 
         Raises:
             AssertionError: If the run was not executed by to this orchestrator.
@@ -943,28 +949,31 @@ 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 == PipelineState.PIPELINE_STATE_UNSPECIFIED:
-            return run.status
+            pipeline_status = run.status
         elif status in [
             PipelineState.PIPELINE_STATE_QUEUED,
             PipelineState.PIPELINE_STATE_PENDING,
         ]:
-            return ExecutionStatus.INITIALIZING
+            pipeline_status = ExecutionStatus.INITIALIZING
         elif status in [
             PipelineState.PIPELINE_STATE_RUNNING,
             PipelineState.PIPELINE_STATE_PAUSED,
         ]:
-            return ExecutionStatus.RUNNING
+            pipeline_status = ExecutionStatus.RUNNING
         elif status == PipelineState.PIPELINE_STATE_SUCCEEDED:
-            return ExecutionStatus.COMPLETED
+            pipeline_status = ExecutionStatus.COMPLETED
         elif status == PipelineState.PIPELINE_STATE_CANCELLING:
-            return ExecutionStatus.STOPPING
+            pipeline_status = ExecutionStatus.STOPPING
         elif status == PipelineState.PIPELINE_STATE_CANCELLED:
-            return ExecutionStatus.STOPPED
+            pipeline_status = ExecutionStatus.STOPPED
         elif status == PipelineState.PIPELINE_STATE_FAILED:
-            return ExecutionStatus.FAILED
+            pipeline_status = ExecutionStatus.FAILED
         else:
             raise ValueError("Unknown status for the pipeline job.")
 
+        # Vertex doesn't support step-level status fetching yet
+        return pipeline_status, None
+
     def compute_metadata(
         self, job: aiplatform.PipelineJob
     ) -> Dict[str, MetadataType]:

zenml/integrations/huggingface/__init__.py

@@ -49,7 +49,7 @@ class HuggingfaceIntegration(Integration):
             A list of requirements.
         """
         requirements = [
-            "datasets>=2.16.0",
+            "datasets>=2.16.0,<4.0.0",
             "huggingface_hub>0.19.0",
             "accelerate",
             "bitsandbytes>=0.41.3",

zenml/integrations/hyperai/flavors/hyperai_orchestrator_flavor.py

@@ -15,6 +15,8 @@
 
 from typing import TYPE_CHECKING, Dict, Optional, Type
 
+from pydantic import Field
+
 from zenml.config.base_settings import BaseSettings
 from zenml.integrations.hyperai import HYPERAI_RESOURCE_TYPE
 from zenml.logger import get_logger
@@ -31,42 +33,38 @@ logger = get_logger(__name__)
 
 
 class HyperAIOrchestratorSettings(BaseSettings):
-    """HyperAI orchestrator settings.
-
-    Attributes:
-        mounts_from_to: A dictionary mapping from paths on the HyperAI instance
-            to paths within the Docker container. This allows users to mount
-            directories from the HyperAI instance into the Docker container that runs
-            on it.
-    """
+    """HyperAI orchestrator settings."""
 
-    mounts_from_to: Dict[str, str] = {}
+    mounts_from_to: Dict[str, str] = Field(
+        default_factory=dict,
+        description="A dictionary mapping from paths on the HyperAI instance "
+        "to paths within the Docker container. This allows users to mount "
+        "directories from the HyperAI instance into the Docker container.",
+    )
 
 
 class HyperAIOrchestratorConfig(
     BaseOrchestratorConfig, HyperAIOrchestratorSettings
 ):
-    """Configuration for the HyperAI orchestrator.
-
-    Attributes:
-        container_registry_autologin: If True, the orchestrator will attempt to
-            automatically log in to the container registry specified in the stack
-            configuration on the HyperAI instance. This is useful if the container
-            registry requires authentication and the HyperAI instance has not been
-            manually logged in to the container registry. Defaults to `False`.
-        automatic_cleanup_pipeline_files: If True, the orchestrator will
-            automatically clean up old pipeline files that are on the HyperAI
-            instance. Pipeline files will be cleaned up if they are 7 days old or
-            older. Defaults to `True`.
-        gpu_enabled_in_container: If True, the orchestrator will enable GPU
-            support in the Docker container that runs on the HyperAI instance.
-            Defaults to `True`.
-
-    """
-
-    container_registry_autologin: bool = False
-    automatic_cleanup_pipeline_files: bool = True
-    gpu_enabled_in_container: bool = True
+    """Configuration for the HyperAI orchestrator."""
+
+    container_registry_autologin: bool = Field(
+        False,
+        description="If True, the orchestrator will attempt to automatically "
+        "log in to the container registry specified in the stack configuration "
+        "on the HyperAI instance.",
+    )
+    automatic_cleanup_pipeline_files: bool = Field(
+        True,
+        description="If True, the orchestrator will automatically clean up old "
+        "pipeline files that are on the HyperAI instance. Pipeline files will "
+        "be cleaned up if they are 7 days old or older.",
+    )
+    gpu_enabled_in_container: bool = Field(
+        True,
+        description="If True, the orchestrator will enable GPU support in the "
+        "Docker container that runs on the HyperAI instance.",
+    )
 
     @property
     def is_remote(self) -> bool:

zenml/integrations/kaniko/flavors/kaniko_image_builder_flavor.py

@@ -15,7 +15,7 @@
 
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Type
 
-from pydantic import PositiveInt
+from pydantic import Field, PositiveInt
 
 from zenml.image_builders import BaseImageBuilderConfig, BaseImageBuilderFlavor
 from zenml.integrations.kaniko import KANIKO_IMAGE_BUILDER_FLAVOR
@@ -39,47 +39,63 @@ class KanikoImageBuilderConfig(BaseImageBuilderConfig):
     configure secrets and environment variables so that the Kaniko build
     container is able to push to the container registry (and optionally access
     the artifact store to upload the build context).
-
-    Attributes:
-        kubernetes_context: The Kubernetes context in which to run the Kaniko
-            pod.
-        kubernetes_namespace: The Kubernetes namespace in which to run the
-            Kaniko pod. This namespace will not be created and must already
-            exist.
-        executor_image: The image of the Kaniko executor to use.
-        pod_running_timeout: The timeout to wait until the pod is running
-            in seconds. Defaults to `300`.
-        env: `env` section of the Kubernetes container spec.
-        env_from: `envFrom` section of the Kubernetes container spec.
-        volume_mounts: `volumeMounts` section of the Kubernetes container spec.
-        volumes: `volumes` section of the Kubernetes pod spec.
-        service_account_name: Name of the Kubernetes service account to use.
-        store_context_in_artifact_store: If `True`, the build context will be
-            stored in the artifact store. If `False`, the build context will be
-            streamed over stdin of the `kubectl` process that runs the build.
-            In case the artifact store is used, the container running the build
-            needs read access to the artifact store.
-        executor_args: Additional arguments to forward to the Kaniko executor.
-            See https://github.com/GoogleContainerTools/kaniko#additional-flags
-            for a full list of available arguments.
-            Example: `["--compressed-caching=false"]`
-
     """
 
-    kubernetes_context: str
-    kubernetes_namespace: str = "zenml-kaniko"
-    executor_image: str = DEFAULT_KANIKO_EXECUTOR_IMAGE
-    pod_running_timeout: PositiveInt = DEFAULT_KANIKO_POD_RUNNING_TIMEOUT
-
-    env: List[Dict[str, Any]] = []
-    env_from: List[Dict[str, Any]] = []
-    volume_mounts: List[Dict[str, Any]] = []
-    volumes: List[Dict[str, Any]] = []
-    service_account_name: Optional[str] = None
-
-    store_context_in_artifact_store: bool = False
-
-    executor_args: List[str] = []
+    kubernetes_context: str = Field(
+        ...,
+        description="The Kubernetes context in which to run the Kaniko pod.",
+    )
+    kubernetes_namespace: str = Field(
+        "zenml-kaniko",
+        description="The Kubernetes namespace in which to run the Kaniko pod. "
+        "This namespace will not be created and must already exist.",
+    )
+    executor_image: str = Field(
+        DEFAULT_KANIKO_EXECUTOR_IMAGE,
+        description="The image of the Kaniko executor to use for building container images.",
+    )
+    pod_running_timeout: PositiveInt = Field(
+        DEFAULT_KANIKO_POD_RUNNING_TIMEOUT,
+        description="The timeout to wait until the pod is running in seconds.",
+    )
+
+    env: List[Dict[str, Any]] = Field(
+        default_factory=list,
+        description="Environment variables section of the Kubernetes container spec. "
+        "Used to configure secrets and environment variables for registry access.",
+    )
+    env_from: List[Dict[str, Any]] = Field(
+        default_factory=list,
+        description="EnvFrom section of the Kubernetes container spec. "
+        "Used to load environment variables from ConfigMaps or Secrets.",
+    )
+    volume_mounts: List[Dict[str, Any]] = Field(
+        default_factory=list,
+        description="VolumeMounts section of the Kubernetes container spec. "
+        "Used to mount volumes containing credentials or other data.",
+    )
+    volumes: List[Dict[str, Any]] = Field(
+        default_factory=list,
+        description="Volumes section of the Kubernetes pod spec. "
+        "Used to define volumes for credentials or other data.",
+    )
+    service_account_name: Optional[str] = Field(
+        None,
+        description="Name of the Kubernetes service account to use for the Kaniko pod. "
+        "This service account should have the necessary permissions for building and pushing images.",
+    )
+
+    store_context_in_artifact_store: bool = Field(
+        False,
+        description="If `True`, the build context will be stored in the artifact store. "
+        "If `False`, the build context will be streamed over stdin of the kubectl process.",
+    )
+
+    executor_args: List[str] = Field(
+        default_factory=list,
+        description="Additional arguments to forward to the Kaniko executor. "
+        "See Kaniko documentation for available flags, e.g. ['--compressed-caching=false'].",
+    )
 
 
 class KanikoImageBuilderFlavor(BaseImageBuilderFlavor):

zenml/integrations/kubeflow/flavors/kubeflow_orchestrator_flavor.py

@@ -15,7 +15,7 @@
 
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Type, cast
 
-from pydantic import model_validator
+from pydantic import Field, model_validator
 
 from zenml.config.base_settings import BaseSettings
 from zenml.constants import KUBERNETES_CLUSTER_RESOURCE_TYPE
@@ -36,34 +36,43 @@ DEFAULT_KFP_UI_PORT = 8080
 
 
 class KubeflowOrchestratorSettings(BaseSettings):
-    """Settings for the Kubeflow 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`. This setting only
-            has an effect when specified on the pipeline and will be ignored if
-            specified on steps.
-        timeout: How many seconds to wait for synchronous runs.
-        client_args: Arguments to pass when initializing the KFP client.
-        client_username: Username to generate a session cookie for the kubeflow client. Both `client_username`
-        and `client_password` need to be set together.
-        client_password: Password to generate a session cookie for the kubeflow client. Both `client_username`
-        and `client_password` need to be set together.
-        user_namespace: The user namespace to use when creating experiments
-            and runs.
-        pod_settings: Pod settings to apply.
-    """
-
-    synchronous: bool = True
-    timeout: int = 1200
-
-    client_args: Dict[str, Any] = {}
-    client_username: Optional[str] = SecretField(default=None)
-    client_password: Optional[str] = SecretField(default=None)
-    user_namespace: Optional[str] = None
-    pod_settings: Optional[KubernetesPodSettings] = None
+    """Settings for the Kubeflow orchestrator."""
+
+    synchronous: bool = Field(
+        True,
+        description="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.",
+    )
+    timeout: int = Field(
+        1200, description="How many seconds to wait for synchronous runs."
+    )
+
+    client_args: Dict[str, Any] = Field(
+        default_factory=dict,
+        description="Arguments to pass when initializing the KFP client. "
+        "Example: {'host': 'https://kubeflow.example.com', 'client_id': 'kubeflow-oidc-authservice', 'existing_token': 'your-auth-token'}",
+    )
+    client_username: Optional[str] = SecretField(
+        default=None,
+        description="Username to generate a session cookie for the kubeflow client. "
+        "Both `client_username` and `client_password` need to be set together.",
+    )
+    client_password: Optional[str] = SecretField(
+        default=None,
+        description="Password to generate a session cookie for the kubeflow client. "
+        "Both `client_username` and `client_password` need to be set together.",
+    )
+    user_namespace: Optional[str] = Field(
+        None,
+        description="The user namespace to use when creating experiments and runs. "
+        "Example: 'my-experiments' or 'team-alpha'",
+    )
+    pod_settings: Optional[KubernetesPodSettings] = Field(
+        None,
+        description="Pod settings to apply to the orchestrator and step pods.",
+    )
 
     @model_validator(mode="before")
     @classmethod
@@ -131,25 +140,27 @@ class KubeflowOrchestratorSettings(BaseSettings):
 class KubeflowOrchestratorConfig(
     BaseOrchestratorConfig, KubeflowOrchestratorSettings
 ):
-    """Configuration for the Kubeflow orchestrator.
-
-    Attributes:
-        kubeflow_hostname: The hostname to use to talk to the Kubeflow Pipelines
-            API. If not set, the hostname will be derived from the Kubernetes
-            API proxy. Mandatory when connecting to a multi-tenant Kubeflow
-            Pipelines deployment.
-        kubeflow_namespace: The Kubernetes namespace in which Kubeflow
-            Pipelines is deployed. Defaults to `kubeflow`.
-        kubernetes_context: Name of a kubernetes context to run
-            pipelines in. Not applicable when connecting to a multi-tenant
-            Kubeflow Pipelines deployment (i.e. when `kubeflow_hostname` is
-            set) or if the stack component is linked to a Kubernetes service
-            connector.
-    """
-
-    kubeflow_hostname: Optional[str] = None
-    kubeflow_namespace: str = "kubeflow"
-    kubernetes_context: Optional[str] = None  # TODO: Potential setting
+    """Configuration for the Kubeflow orchestrator."""
+
+    kubeflow_hostname: Optional[str] = Field(
+        None,
+        description="The hostname to use to talk to the Kubeflow Pipelines API. "
+        "If not set, the hostname will be derived from the Kubernetes API proxy. "
+        "Mandatory when connecting to a multi-tenant Kubeflow Pipelines deployment. "
+        "Example: 'https://kubeflow.example.com' or 'kubeflow.company.com'",
+    )
+    kubeflow_namespace: str = Field(
+        "kubeflow",
+        description="The Kubernetes namespace in which Kubeflow Pipelines is deployed.",
+    )
+    kubernetes_context: Optional[str] = Field(
+        None,
+        description="Name of a kubernetes context to run pipelines in. "
+        "Not applicable when connecting to a multi-tenant Kubeflow Pipelines "
+        "deployment (i.e. when `kubeflow_hostname` is set) or if the stack "
+        "component is linked to a Kubernetes service connector. "
+        "Example: 'my-cluster' or 'production-cluster'",
+    )
 
     @model_validator(mode="before")
     @classmethod