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

zenml/integrations/aws/orchestrators/sagemaker_orchestrator.py

@@ -807,14 +807,20 @@ class SagemakerOrchestrator(ContainerizedOrchestrator):
             settings=settings,
         )
 
-    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 SageMaker, so step_statuses_dict will always be None.
 
         Raises:
             AssertionError: If the run was not executed by to this orchestrator.
@@ -855,18 +861,21 @@ class SagemakerOrchestrator(ContainerizedOrchestrator):
         # Map the potential outputs to ZenML ExecutionStatus. Potential values:
         # https://docs.aws.amazon.com/sagemaker/latest/APIReference/API_DescribePipelineExecution.html
         if status == "Executing":
-            return ExecutionStatus.RUNNING
+            pipeline_status = ExecutionStatus.RUNNING
         elif status == "Stopping":
-            return ExecutionStatus.STOPPING
+            pipeline_status = ExecutionStatus.STOPPING
         elif status == "Stopped":
-            return ExecutionStatus.STOPPED
+            pipeline_status = ExecutionStatus.STOPPED
         elif status == "Failed":
-            return ExecutionStatus.FAILED
+            pipeline_status = ExecutionStatus.FAILED
         elif status == "Succeeded":
-            return ExecutionStatus.COMPLETED
+            pipeline_status = ExecutionStatus.COMPLETED
         else:
             raise ValueError("Unknown status for the pipeline execution.")
 
+        # SageMaker doesn't support step-level status fetching yet
+        return pipeline_status, None
+
     def compute_metadata(
         self,
         execution_arn: str,

zenml/integrations/azure/orchestrators/azureml_orchestrator.py

@@ -451,14 +451,20 @@ class AzureMLOrchestrator(ContainerizedOrchestrator):
             )
             return {}
 
-    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 (not supported for AzureML).
 
         Returns:
-            the actual status of the pipeline execution.
+            A tuple of (pipeline_status, step_statuses_dict).
+            Step statuses are not supported for AzureML, so step_statuses_dict will always be None.
 
         Raises:
             AssertionError: If the run was not executed by to this orchestrator.
@@ -512,23 +518,23 @@ class AzureMLOrchestrator(ContainerizedOrchestrator):
             "Preparing",
             "Queued",
         ]:
-            return ExecutionStatus.INITIALIZING
+            pipeline_status = ExecutionStatus.INITIALIZING
         elif status in ["Running", "Finalizing"]:
-            return ExecutionStatus.RUNNING
+            pipeline_status = ExecutionStatus.RUNNING
         elif status == "CancelRequested":
-            return ExecutionStatus.STOPPING
+            pipeline_status = ExecutionStatus.STOPPING
         elif status == "Canceled":
-            return ExecutionStatus.STOPPED
-        elif status in [
-            "Failed",
-            "NotResponding",
-        ]:
-            return ExecutionStatus.FAILED
+            pipeline_status = ExecutionStatus.STOPPED
+        elif status in ["Failed", "NotResponding"]:
+            pipeline_status = ExecutionStatus.FAILED
         elif status == "Completed":
-            return ExecutionStatus.COMPLETED
+            pipeline_status = ExecutionStatus.COMPLETED
         else:
             raise ValueError("Unknown status for the pipeline job.")
 
+        # AzureML doesn't support step-level status fetching yet
+        return pipeline_status, None
+
     def compute_metadata(self, job: Any) -> Dict[str, MetadataType]:
         """Generate run metadata based on the generated AzureML PipelineJob.
 

zenml/integrations/bentoml/flavors/bentoml_model_deployer_flavor.py

@@ -15,6 +15,8 @@
 
 from typing import TYPE_CHECKING, Optional, Type
 
+from pydantic import Field
+
 from zenml.integrations.bentoml import BENTOML_MODEL_DEPLOYER_FLAVOR
 from zenml.model_deployers.base_model_deployer import (
     BaseModelDeployerConfig,
@@ -28,7 +30,11 @@ if TYPE_CHECKING:
 class BentoMLModelDeployerConfig(BaseModelDeployerConfig):
     """Configuration for the BentoMLModelDeployer."""
 
-    service_path: str = ""
+    service_path: str = Field(
+        "",
+        description="Path to the BentoML service directory. "
+        "If not provided, a default service path will be used.",
+    )
 
 
 class BentoMLModelDeployerFlavor(BaseModelDeployerFlavor):

zenml/integrations/databricks/flavors/databricks_orchestrator_flavor.py

@@ -15,6 +15,8 @@
 
 from typing import TYPE_CHECKING, Dict, Optional, Tuple, Type
 
+from pydantic import Field
+
 from zenml.config.base_settings import BaseSettings
 from zenml.integrations.databricks import DATABRICKS_ORCHESTRATOR_FLAVOR
 from zenml.logger import get_logger
@@ -42,31 +44,64 @@ class DatabricksAvailabilityType(StrEnum):
 class DatabricksOrchestratorSettings(BaseSettings):
     """Databricks orchestrator base settings.
 
-    Attributes:
-        spark_version: Spark version.
-        num_workers: Number of workers.
-        node_type_id: Node type id.
-        policy_id: Policy id.
-        autotermination_minutes: Autotermination minutes.
-        autoscale: Autoscale.
-        single_user_name: Single user name.
-        spark_conf: Spark configuration.
-        spark_env_vars: Spark environment variables.
-        schedule_timezone: Schedule timezone.
+    Configuration for Databricks cluster and Spark execution settings.
+    Field descriptions are defined inline using Field() descriptors.
     """
 
-    # Resources
-    spark_version: Optional[str] = None
-    num_workers: Optional[int] = None
-    node_type_id: Optional[str] = None
-    policy_id: Optional[str] = None
-    autotermination_minutes: Optional[int] = None
-    autoscale: Tuple[int, int] = (0, 1)
-    single_user_name: Optional[str] = None
-    spark_conf: Optional[Dict[str, str]] = None
-    spark_env_vars: Optional[Dict[str, str]] = None
-    schedule_timezone: Optional[str] = None
-    availability_type: Optional[DatabricksAvailabilityType] = None
+    # Cluster Configuration
+    spark_version: Optional[str] = Field(
+        default=None,
+        description="Apache Spark version for the Databricks cluster. "
+        "Uses workspace default if not specified. Example: '3.2.x-scala2.12'",
+    )
+    num_workers: Optional[int] = Field(
+        default=None,
+        description="Fixed number of worker nodes. Cannot be used with autoscaling.",
+    )
+    node_type_id: Optional[str] = Field(
+        default=None,
+        description="Databricks node type identifier. "
+        "Refer to Databricks documentation for available instance types. "
+        "Example: 'i3.xlarge'",
+    )
+    policy_id: Optional[str] = Field(
+        default=None,
+        description="Databricks cluster policy ID for governance and cost control.",
+    )
+    autotermination_minutes: Optional[int] = Field(
+        default=None,
+        description="Minutes of inactivity before automatic cluster termination. "
+        "Helps control costs by shutting down idle clusters.",
+    )
+    autoscale: Tuple[int, int] = Field(
+        default=(0, 1),
+        description="Cluster autoscaling bounds as (min_workers, max_workers). "
+        "Automatically adjusts cluster size based on workload.",
+    )
+    single_user_name: Optional[str] = Field(
+        default=None,
+        description="Databricks username for single-user cluster access mode.",
+    )
+    spark_conf: Optional[Dict[str, str]] = Field(
+        default=None,
+        description="Custom Spark configuration properties as key-value pairs. "
+        "Example: {'spark.sql.adaptive.enabled': 'true', 'spark.sql.adaptive.coalescePartitions.enabled': 'true'}",
+    )
+    spark_env_vars: Optional[Dict[str, str]] = Field(
+        default=None,
+        description="Environment variables for the Spark driver and executors. "
+        "Example: {'SPARK_WORKER_MEMORY': '4g', 'SPARK_DRIVER_MEMORY': '2g'}",
+    )
+    schedule_timezone: Optional[str] = Field(
+        default=None,
+        description="Timezone for scheduled pipeline execution. "
+        "Uses IANA timezone format (e.g., 'America/New_York').",
+    )
+    availability_type: Optional[DatabricksAvailabilityType] = Field(
+        default=None,
+        description="Instance availability type: ON_DEMAND (guaranteed), SPOT (cost-optimized), "
+        "or SPOT_WITH_FALLBACK (spot with on-demand backup).",
+    )
 
 
 class DatabricksOrchestratorConfig(

zenml/integrations/feast/flavors/feast_feature_store_flavor.py

@@ -15,6 +15,8 @@
 
 from typing import TYPE_CHECKING, Optional, Type
 
+from pydantic import Field
+
 from zenml.feature_stores.base_feature_store import (
     BaseFeatureStoreConfig,
     BaseFeatureStoreFlavor,
@@ -26,11 +28,22 @@ if TYPE_CHECKING:
 
 
 class FeastFeatureStoreConfig(BaseFeatureStoreConfig):
-    """Config for Feast feature store."""
-
-    online_host: str = "localhost"
-    online_port: int = 6379
-    feast_repo: str
+    """Config for Feast feature store.
+
+    Configuration for connecting to Feast feature stores.
+    Field descriptions are defined inline using Field() descriptors.
+    """
+
+    online_host: str = Field(
+        default="localhost",
+        description="Online feature store host address (typically Redis server).",
+    )
+    online_port: int = Field(
+        default=6379, description="Online feature store port number."
+    )
+    feast_repo: str = Field(
+        description="Local filesystem path to the Feast repository with feature definitions."
+    )
 
     @property
     def is_local(self) -> bool:

zenml/integrations/gcp/flavors/vertex_experiment_tracker_flavor.py

@@ -16,7 +16,7 @@
 import re
 from typing import TYPE_CHECKING, Any, Dict, Optional, Type, Union
 
-from pydantic import field_validator
+from pydantic import Field, field_validator
 
 from zenml.config.base_settings import BaseSettings
 from zenml.experiment_trackers.base_experiment_tracker import (
@@ -40,15 +40,15 @@ if TYPE_CHECKING:
 
 
 class VertexExperimentTrackerSettings(BaseSettings):
-    """Settings for the VertexAI experiment tracker.
+    """Settings for the VertexAI experiment tracker."""
 
-    Attributes:
-        experiment: The VertexAI experiment name.
-        experiment_tensorboard: The VertexAI experiment tensorboard.
-    """
-
-    experiment: Optional[str] = None
-    experiment_tensorboard: Optional[Union[str, bool]] = None
+    experiment: Optional[str] = Field(
+        None, description="The VertexAI experiment name."
+    )
+    experiment_tensorboard: Optional[Union[str, bool]] = Field(
+        None,
+        description="The VertexAI experiment tensorboard instance to use.",
+    )
 
     @field_validator("experiment", mode="before")
     def _validate_experiment(cls, value: str) -> str:
@@ -76,39 +76,7 @@ class VertexExperimentTrackerConfig(
     GoogleCredentialsConfigMixin,
     VertexExperimentTrackerSettings,
 ):
-    """Config for the VertexAI experiment tracker.
-
-    Attributes:
-        location: Optional. The default location to use when making API calls. If not
-            set defaults to us-central1.
-        staging_bucket: Optional. The default staging bucket to use to stage artifacts
-            when making API calls. In the form gs://...
-        network:
-            Optional. The full name of the Compute Engine network to which jobs
-            and resources should be peered. E.g. "projects/12345/global/networks/myVPC".
-            Private services access must already be configured for the network.
-            If specified, all eligible jobs and resources created will be peered
-            with this VPC.
-        encryption_spec_key_name:
-            Optional. The Cloud KMS resource identifier of the customer
-            managed encryption key used to protect a resource. Has the
-            form:
-            ``projects/my-project/locations/my-region/keyRings/my-kr/cryptoKeys/my-key``.
-            The key needs to be in the same region as where the compute
-            resource is created.
-        api_endpoint (str):
-                Optional. The desired API endpoint,
-                e.g., us-central1-aiplatform.googleapis.com
-        api_key (str):
-            Optional. The API key to use for service calls.
-            NOTE: Not all services support API keys.
-        api_transport (str):
-            Optional. The transport method which is either 'grpc' or 'rest'.
-            NOTE: "rest" transport functionality is currently in a
-            beta state (preview).
-        request_metadata:
-            Optional. Additional gRPC metadata to send with every client request.
-    """
+    """Config for the VertexAI experiment tracker."""
 
     location: Optional[str] = None
     staging_bucket: Optional[str] = None

zenml/integrations/gcp/flavors/vertex_orchestrator_flavor.py

@@ -15,6 +15,8 @@
 
 from typing import TYPE_CHECKING, Dict, Optional, Tuple, Type
 
+from pydantic import Field
+
 from zenml.config.base_settings import BaseSettings
 from zenml.integrations.gcp import (
     GCP_RESOURCE_TYPE,
@@ -36,35 +38,37 @@ if TYPE_CHECKING:
 
 
 class VertexOrchestratorSettings(BaseSettings):
-    """Settings for the Vertex 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`.
-        labels: Labels to assign to the pipeline job.
-        node_selector_constraint: Each constraint is a key-value pair label.
-            For the container to be eligible to run on a node, the node must have
-            each of the constraints appeared as labels.
-            For example a GPU type can be providing by one of the following tuples:
-                - ("cloud.google.com/gke-accelerator", "NVIDIA_TESLA_A100")
-                - ("cloud.google.com/gke-accelerator", "NVIDIA_TESLA_K80")
-                - ("cloud.google.com/gke-accelerator", "NVIDIA_TESLA_P4")
-                - ("cloud.google.com/gke-accelerator", "NVIDIA_TESLA_P100")
-                - ("cloud.google.com/gke-accelerator", "NVIDIA_TESLA_T4")
-                - ("cloud.google.com/gke-accelerator", "NVIDIA_TESLA_V100")
-            Hint: the selected region (location) must provide the requested accelerator
-            (see https://cloud.google.com/compute/docs/gpus/gpu-regions-zones).
-        pod_settings: Pod settings to apply.
-    """
-
-    labels: Dict[str, str] = {}
-    synchronous: bool = True
-    node_selector_constraint: Optional[Tuple[str, str]] = None
-    pod_settings: Optional[KubernetesPodSettings] = None
-
-    custom_job_parameters: Optional[VertexCustomJobParameters] = None
+    """Settings for the Vertex orchestrator."""
+
+    labels: Dict[str, str] = Field(
+        default_factory=dict,
+        description="Labels to assign to the pipeline job. "
+        "Example: {'environment': 'production', 'team': 'ml-ops'}",
+    )
+    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.",
+    )
+    node_selector_constraint: Optional[Tuple[str, str]] = Field(
+        None,
+        description="Each constraint is a key-value pair label. For the container "
+        "to be eligible to run on a node, the node must have each of the "
+        "constraints appeared as labels. For example, a GPU type can be provided "
+        "by ('cloud.google.com/gke-accelerator', 'NVIDIA_TESLA_T4')."
+        "Hint: the selected region (location) must provide the requested accelerator"
+        "(see https://cloud.google.com/compute/docs/gpus/gpu-regions-zones).",
+    )
+    pod_settings: Optional[KubernetesPodSettings] = Field(
+        None,
+        description="Pod settings to apply to the orchestrator and step pods.",
+    )
+
+    custom_job_parameters: Optional[VertexCustomJobParameters] = Field(
+        None, description="Custom parameters for the Vertex AI custom job."
+    )
 
     _node_selector_deprecation = (
         deprecation_utils.deprecate_pydantic_attributes(
@@ -78,71 +82,74 @@ class VertexOrchestratorConfig(
     GoogleCredentialsConfigMixin,
     VertexOrchestratorSettings,
 ):
-    """Configuration for the Vertex orchestrator.
-
-    Attributes:
-        location: Name of GCP region where the pipeline job will be executed.
-            Vertex AI Pipelines is available in the following regions:
-            https://cloud.google.com/vertex-ai/docs/general/locations#feature-availability
-        pipeline_root: a Cloud Storage URI that will be used by the Vertex AI
-            Pipelines. If not provided but the artifact store in the stack used
-            to execute the pipeline is a
-            `zenml.integrations.gcp.artifact_stores.GCPArtifactStore`,
-            then a subdirectory of the artifact store will be used.
-        encryption_spec_key_name: The Cloud KMS resource identifier of the
-            customer managed encryption key used to protect the job. Has the form:
-            `projects/<PROJECT>/locations/<REGION>/keyRings/<KR>/cryptoKeys/<KEY>`
-            . The key needs to be in the same region as where the compute
-            resource is created.
-        workload_service_account: the service account for workload run-as
-            account. Users submitting jobs must have act-as permission on this
-            run-as account. If not provided, the Compute Engine default service
-            account for the GCP project in which the pipeline is running is
-            used.
-        function_service_account: the service account for cloud function run-as
-            account, for scheduled pipelines. This service account must have
-            the act-as permission on the workload_service_account.
-            If not provided, the Compute Engine default service account for the
-            GCP project in which the pipeline is running is used.
-        scheduler_service_account: the service account used by the Google Cloud
-            Scheduler to trigger and authenticate to the pipeline Cloud Function
-            on a schedule. If not provided, the Compute Engine default service
-            account for the GCP project in which the pipeline is running is
-            used.
-        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
-            at most 96 CPUs.
-            (see. https://cloud.google.com/vertex-ai/docs/pipelines/machine-types)
-        memory_limit: The maximum memory limit for this operator. This string
-            value can be a number, or a number followed by "K" (kilobyte),
-            "M" (megabyte), or "G" (gigabyte). At most 624GB is supported.
-        gpu_limit: The GPU limit (positive number) for the operator.
-            For more information about GPU resources, see:
-            https://cloud.google.com/vertex-ai/docs/training/configure-compute#specifying_gpus
-    """
-
-    location: str
-    pipeline_root: Optional[str] = None
-    encryption_spec_key_name: Optional[str] = None
-    workload_service_account: Optional[str] = None
-    network: Optional[str] = None
-    private_service_connect: Optional[str] = None
+    """Configuration for the Vertex orchestrator."""
+
+    location: str = Field(
+        ...,
+        description="Name of GCP region where the pipeline job will be executed. "
+        "Vertex AI Pipelines is available in specific regions: "
+        "https://cloud.google.com/vertex-ai/docs/general/locations#feature-availability",
+    )
+    pipeline_root: Optional[str] = Field(
+        None,
+        description="A Cloud Storage URI that will be used by the Vertex AI Pipelines. "
+        "If not provided but the artifact store in the stack is a GCPArtifactStore, "
+        "then a subdirectory of the artifact store will be used.",
+    )
+    encryption_spec_key_name: Optional[str] = Field(
+        None,
+        description="The Cloud KMS resource identifier of the customer managed "
+        "encryption key used to protect the job. Has the form: "
+        "projects/<PROJECT>/locations/<REGION>/keyRings/<KR>/cryptoKeys/<KEY>. "
+        "The key needs to be in the same region as where the compute resource is created.",
+    )
+    workload_service_account: Optional[str] = Field(
+        None,
+        description="The service account for workload run-as account. Users submitting "
+        "jobs must have act-as permission on this run-as account. If not provided, "
+        "the Compute Engine default service account for the GCP project is used.",
+    )
+    network: Optional[str] = Field(
+        None,
+        description="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: Optional[str] = Field(
+        None,
+        description="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.",
+    )
 
     # Deprecated
-    cpu_limit: Optional[str] = None
-    memory_limit: Optional[str] = None
-    gpu_limit: Optional[int] = None
-    function_service_account: Optional[str] = None
-    scheduler_service_account: Optional[str] = None
+    cpu_limit: Optional[str] = Field(
+        None,
+        description="DEPRECATED: The maximum CPU limit for this operator. "
+        "Use custom_job_parameters or pod_settings instead.",
+    )
+    memory_limit: Optional[str] = Field(
+        None,
+        description="DEPRECATED: The maximum memory limit for this operator. "
+        "Use custom_job_parameters or pod_settings instead.",
+    )
+    gpu_limit: Optional[int] = Field(
+        None,
+        description="DEPRECATED: The GPU limit for the operator. "
+        "Use custom_job_parameters or pod_settings instead.",
+    )
+    function_service_account: Optional[str] = Field(
+        None,
+        description="DEPRECATED: The service account for cloud function run-as account, "
+        "for scheduled pipelines. This functionality is no longer supported.",
+    )
+    scheduler_service_account: Optional[str] = Field(
+        None,
+        description="DEPRECATED: The service account used by the Google Cloud Scheduler "
+        "to trigger and authenticate to the pipeline Cloud Function on a schedule. "
+        "This functionality is no longer supported.",
+    )
 
     _resource_deprecation = deprecation_utils.deprecate_pydantic_attributes(
         "cpu_limit",

zenml/integrations/gcp/google_credentials_mixin.py

@@ -15,6 +15,8 @@
 
 from typing import TYPE_CHECKING, Optional, Tuple, cast
 
+from pydantic import Field
+
 from zenml.logger import get_logger
 from zenml.stack.stack_component import StackComponent, StackComponentConfig
 
@@ -28,16 +30,19 @@ logger = get_logger(__name__)
 class GoogleCredentialsConfigMixin(StackComponentConfig):
     """Config mixin for Google Cloud Platform credentials.
 
-    Attributes:
-        project: GCP project name. If `None`, the project will be inferred from
-            the environment.
-        service_account_path: path to the service account credentials file to be
-            used for authentication. If not provided, the default credentials
-            will be used.
+    Provides common GCP authentication configuration for stack components.
+    Field descriptions are defined inline using Field() descriptors.
     """
 
-    project: Optional[str] = None
-    service_account_path: Optional[str] = None
+    project: Optional[str] = Field(
+        default=None,
+        description="Google Cloud Project ID. Auto-detected from environment if not specified.",
+    )
+    service_account_path: Optional[str] = Field(
+        default=None,
+        description="Path to service account JSON key file for authentication. "
+        "Uses Application Default Credentials if not provided.",
+    )
 
 
 class GoogleCredentialsMixin(StackComponent):