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

zenml/integrations/lightning/flavors/lightning_orchestrator_flavor.py

@@ -15,6 +15,8 @@
 
 from typing import TYPE_CHECKING, List, Optional, Type
 
+from pydantic import Field
+
 from zenml.config.base_settings import BaseSettings
 from zenml.integrations.lightning import LIGHTNING_ORCHESTRATOR_FLAVOR
 from zenml.logger import get_logger
@@ -33,33 +35,47 @@ logger = get_logger(__name__)
 class LightningOrchestratorSettings(BaseSettings):
     """Lightning orchestrator base settings.
 
-    Attributes:
-        main_studio_name: Main studio name.
-        machine_type: Machine type.
-        user_id: User id.
-        api_key: api_key.
-        username: Username.
-        teamspace: Teamspace.
-        organization: Organization.
-        custom_commands: Custom commands to run.
-        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.
+    Configuration for executing pipelines on Lightning AI platform.
+    Field descriptions are defined inline using Field() descriptors.
     """
 
-    # Resources
-    main_studio_name: Optional[str] = None
-    machine_type: Optional[str] = None
-    user_id: Optional[str] = SecretField(default=None)
-    api_key: Optional[str] = SecretField(default=None)
-    username: Optional[str] = None
-    teamspace: Optional[str] = None
-    organization: Optional[str] = None
-    custom_commands: Optional[List[str]] = None
-    synchronous: bool = True
+    # Lightning AI Platform Configuration
+    main_studio_name: Optional[str] = Field(
+        default=None,
+        description="Lightning AI studio instance name where the pipeline will execute.",
+    )
+    machine_type: Optional[str] = Field(
+        default=None,
+        description="Compute instance type for pipeline execution. "
+        "Refer to Lightning AI documentation for available options.",
+    )
+    user_id: Optional[str] = SecretField(
+        default=None, description="Lightning AI user ID for authentication."
+    )
+    api_key: Optional[str] = SecretField(
+        default=None,
+        description="Lightning AI API key for platform authentication.",
+    )
+    username: Optional[str] = Field(
+        default=None, description="Lightning AI platform username."
+    )
+    teamspace: Optional[str] = Field(
+        default=None,
+        description="Lightning AI teamspace for collaborative pipeline execution.",
+    )
+    organization: Optional[str] = Field(
+        default=None,
+        description="Lightning AI organization name for enterprise accounts.",
+    )
+    custom_commands: Optional[List[str]] = Field(
+        default=None,
+        description="Additional shell commands to execute in the Lightning AI environment.",
+    )
+    synchronous: bool = Field(
+        default=True,
+        description="Whether to wait for pipeline completion. "
+        "When `False`, execution continues asynchronously after submission.",
+    )
 
 
 class LightningOrchestratorConfig(

zenml/integrations/mlflow/flavors/mlflow_experiment_tracker_flavor.py

@@ -15,7 +15,7 @@
 
 from typing import TYPE_CHECKING, Any, Dict, Optional, Type
 
-from pydantic import model_validator
+from pydantic import Field, model_validator
 
 from zenml.config.base_settings import BaseSettings
 from zenml.experiment_trackers.base_experiment_tracker import (
@@ -59,56 +59,63 @@ def is_databricks_tracking_uri(tracking_uri: str) -> bool:
 
 
 class MLFlowExperimentTrackerSettings(BaseSettings):
-    """Settings for the MLflow experiment tracker.
+    """Settings for the MLflow experiment tracker."""
 
-    Attributes:
-        experiment_name: The MLflow experiment name.
-        nested: If `True`, will create a nested sub-run for the step.
-        tags: Tags for the Mlflow run.
-    """
-
-    experiment_name: Optional[str] = None
-    nested: bool = False
-    tags: Dict[str, Any] = {}
+    experiment_name: Optional[str] = Field(
+        None,
+        description="The MLflow experiment name to use for tracking runs.",
+    )
+    nested: bool = Field(
+        False,
+        description="If `True`, will create a nested sub-run for the step.",
+    )
+    tags: Dict[str, Any] = Field(
+        default_factory=dict,
+        description="Tags to attach to the MLflow run for categorization and filtering.",
+    )
 
 
 class MLFlowExperimentTrackerConfig(
     BaseExperimentTrackerConfig, MLFlowExperimentTrackerSettings
 ):
-    """Config for the MLflow experiment tracker.
-
-    Attributes:
-        tracking_uri: The uri of the mlflow tracking server. If no uri is set,
-            your stack must contain a `LocalArtifactStore` and ZenML will
-            point MLflow to a subdirectory of your artifact store instead.
-        tracking_username: Username for authenticating with the MLflow
-            tracking server. When a remote tracking uri is specified,
-            either `tracking_token` or `tracking_username` and
-            `tracking_password` must be specified.
-        tracking_password: Password for authenticating with the MLflow
-            tracking server. When a remote tracking uri is specified,
-            either `tracking_token` or `tracking_username` and
-            `tracking_password` must be specified.
-        tracking_token: Token for authenticating with the MLflow
-            tracking server. When a remote tracking uri is specified,
-            either `tracking_token` or `tracking_username` and
-            `tracking_password` must be specified.
-        tracking_insecure_tls: Skips verification of TLS connection to the
-            MLflow tracking server if set to `True`.
-        databricks_host: The host of the Databricks workspace with the MLflow
-            managed server to connect to. This is only required if
-            `tracking_uri` value is set to `"databricks"`.
-        enable_unity_catalog: If `True`, will enable the Databricks Unity Catalog for
-            logging and registering models.
-    """
+    """Config for the MLflow experiment tracker."""
 
-    tracking_uri: Optional[str] = None
-    tracking_username: Optional[str] = SecretField(default=None)
-    tracking_password: Optional[str] = SecretField(default=None)
-    tracking_token: Optional[str] = SecretField(default=None)
-    tracking_insecure_tls: bool = False
-    databricks_host: Optional[str] = None
-    enable_unity_catalog: bool = False
+    tracking_uri: Optional[str] = Field(
+        None,
+        description="The URI of the MLflow tracking server. If no URI is set, "
+        "your stack must contain a LocalArtifactStore and ZenML will point "
+        "MLflow to a subdirectory of your artifact store instead.",
+    )
+    tracking_username: Optional[str] = SecretField(
+        default=None,
+        description="Username for authenticating with the MLflow tracking server. "
+        "Required when using a remote tracking URI along with tracking_password.",
+    )
+    tracking_password: Optional[str] = SecretField(
+        default=None,
+        description="Password for authenticating with the MLflow tracking server. "
+        "Required when using a remote tracking URI along with tracking_username.",
+    )
+    tracking_token: Optional[str] = SecretField(
+        default=None,
+        description="Token for authenticating with the MLflow tracking server. "
+        "Alternative to username/password authentication for remote tracking URIs.",
+    )
+    tracking_insecure_tls: bool = Field(
+        False,
+        description="Skips verification of TLS connection to the MLflow tracking "
+        "server if set to `True`. Use with caution in production environments.",
+    )
+    databricks_host: Optional[str] = Field(
+        None,
+        description="The host of the Databricks workspace with the MLflow managed "
+        "server to connect to. Required when tracking_uri is set to 'databricks'.",
+    )
+    enable_unity_catalog: bool = Field(
+        False,
+        description="If `True`, will enable the Databricks Unity Catalog for "
+        "logging and registering models.",
+    )
 
     @model_validator(mode="after")
     def _ensure_authentication_if_necessary(

zenml/integrations/mlflow/flavors/mlflow_model_deployer_flavor.py

@@ -15,6 +15,8 @@
 
 from typing import TYPE_CHECKING, Optional, Type
 
+from pydantic import Field
+
 from zenml.integrations.mlflow import MLFLOW_MODEL_DEPLOYER_FLAVOR
 from zenml.model_deployers.base_model_deployer import (
     BaseModelDeployerConfig,
@@ -28,12 +30,15 @@ if TYPE_CHECKING:
 class MLFlowModelDeployerConfig(BaseModelDeployerConfig):
     """Configuration for the MLflow model deployer.
 
-    Attributes:
-        service_path: the path where the local MLflow deployment service
-            configuration, PID and log files are stored.
+    Configuration for local MLflow model serving.
+    Field descriptions are defined inline using Field() descriptors.
     """
 
-    service_path: str = ""
+    service_path: str = Field(
+        default="",
+        description="Local directory for MLflow deployment service files "
+        "(configuration, PID, and logs). Uses temp directory if empty.",
+    )
 
     @property
     def is_local(self) -> bool:

zenml/integrations/neptune/flavors/neptune_experiment_tracker_flavor.py

@@ -21,6 +21,8 @@ __all__ = [
 
 from typing import TYPE_CHECKING, Optional, Set, Type
 
+from pydantic import Field
+
 from zenml.config.base_settings import BaseSettings
 from zenml.experiment_trackers.base_experiment_tracker import (
     BaseExperimentTrackerConfig,
@@ -40,24 +42,23 @@ class NeptuneExperimentTrackerConfig(BaseExperimentTrackerConfig):
 
     If attributes are left as None, the neptune.init_run() method
     will try to find the relevant values in the environment
-
-    Attributes:
-        project: name of the Neptune project you want to log the metadata to
-        api_token: your Neptune API token
     """
 
-    project: Optional[str] = None
-    api_token: Optional[str] = SecretField(default=None)
+    project: Optional[str] = Field(
+        None,
+        description="Name of the Neptune project you want to log the metadata to.",
+    )
+    api_token: Optional[str] = SecretField(
+        default=None, description="Your Neptune API token for authentication."
+    )
 
 
 class NeptuneExperimentTrackerSettings(BaseSettings):
-    """Settings for the Neptune experiment tracker.
+    """Settings for the Neptune experiment tracker."""
 
-    Attributes:
-        tags: Tags for the Neptune run.
-    """
-
-    tags: Set[str] = set()
+    tags: Set[str] = Field(
+        default_factory=set, description="Tags for the Neptune run."
+    )
 
 
 class NeptuneExperimentTrackerFlavor(BaseExperimentTrackerFlavor):

zenml/integrations/s3/flavors/s3_artifact_store_flavor.py

@@ -24,7 +24,7 @@ from typing import (
     Type,
 )
 
-from pydantic import field_validator
+from pydantic import Field, field_validator
 
 from zenml.artifact_stores import (
     BaseArtifactStoreConfig,
@@ -64,12 +64,37 @@ class S3ArtifactStoreConfig(
 
     SUPPORTED_SCHEMES: ClassVar[Set[str]] = {"s3://"}
 
-    key: Optional[str] = SecretField(default=None)
-    secret: Optional[str] = SecretField(default=None)
-    token: Optional[str] = SecretField(default=None)
-    client_kwargs: Optional[Dict[str, Any]] = None
-    config_kwargs: Optional[Dict[str, Any]] = None
-    s3_additional_kwargs: Optional[Dict[str, Any]] = None
+    key: Optional[str] = SecretField(
+        default=None,
+        description="AWS access key ID for authentication. "
+        "If not provided, credentials will be inferred from the environment.",
+    )
+    secret: Optional[str] = SecretField(
+        default=None,
+        description="AWS secret access key for authentication. "
+        "If not provided, credentials will be inferred from the environment.",
+    )
+    token: Optional[str] = SecretField(
+        default=None,
+        description="AWS session token for temporary credentials. "
+        "If not provided, credentials will be inferred from the environment.",
+    )
+    client_kwargs: Optional[Dict[str, Any]] = Field(
+        None,
+        description="Additional keyword arguments to pass to the S3 client. "
+        "For example, to connect to a custom S3-compatible endpoint: "
+        "{'endpoint_url': 'http://minio:9000'}",
+    )
+    config_kwargs: Optional[Dict[str, Any]] = Field(
+        None,
+        description="Additional keyword arguments to pass to the S3 client configuration. "
+        "For example: {'region_name': 'us-west-2', 'signature_version': 's3v4'}",
+    )
+    s3_additional_kwargs: Optional[Dict[str, Any]] = Field(
+        None,
+        description="Additional keyword arguments for S3 operations. "
+        "For example: {'ACL': 'bucket-owner-full-control'}",
+    )
 
     _bucket: Optional[str] = None
 

zenml/integrations/vllm/flavors/vllm_model_deployer_flavor.py

@@ -15,6 +15,8 @@
 
 from typing import TYPE_CHECKING, Optional, Type
 
+from pydantic import Field
+
 from zenml.integrations.vllm import VLLM_MODEL_DEPLOYER
 from zenml.model_deployers.base_model_deployer import (
     BaseModelDeployerConfig,
@@ -28,7 +30,11 @@ if TYPE_CHECKING:
 class VLLMModelDeployerConfig(BaseModelDeployerConfig):
     """Configuration for vLLM Inference model deployer."""
 
-    service_path: str = ""
+    service_path: str = Field(
+        "",
+        description="The path where the local vLLM deployment service "
+        "configuration, PID and log files are stored.",
+    )
 
 
 class VLLMModelDeployerFlavor(BaseModelDeployerFlavor):

zenml/integrations/wandb/flavors/wandb_experiment_tracker_flavor.py

@@ -23,7 +23,7 @@ from typing import (
     cast,
 )
 
-from pydantic import BaseModel, field_validator
+from pydantic import BaseModel, Field, field_validator
 
 from zenml.config.base_settings import BaseSettings
 from zenml.experiment_trackers.base_experiment_tracker import (
@@ -40,19 +40,24 @@ if TYPE_CHECKING:
 
 
 class WandbExperimentTrackerSettings(BaseSettings):
-    """Settings for the Wandb experiment tracker.
+    """Settings for the Wandb experiment tracker."""
 
-    Attributes:
-        run_name: The Wandb run name.
-        tags: Tags for the Wandb run.
-        settings: Settings for the Wandb run.
-        enable_weave: Whether to enable Weave integration.
-    """
-
-    run_name: Optional[str] = None
-    tags: List[str] = []
-    settings: Dict[str, Any] = {}
-    enable_weave: bool = False
+    run_name: Optional[str] = Field(
+        None,
+        description="The Wandb run name to use for tracking experiments."
+    )
+    tags: List[str] = Field(
+        default_factory=list,
+        description="Tags to attach to the Wandb run for categorization and filtering."
+    )
+    settings: Dict[str, Any] = Field(
+        default_factory=dict,
+        description="Additional settings for the Wandb run configuration."
+    )
+    enable_weave: bool = Field(
+        False,
+        description="Whether to enable Weave integration for enhanced experiment tracking."
+    )
 
     @field_validator("settings", mode="before")
     @classmethod
@@ -89,18 +94,22 @@ class WandbExperimentTrackerSettings(BaseSettings):
 class WandbExperimentTrackerConfig(
     BaseExperimentTrackerConfig, WandbExperimentTrackerSettings
 ):
-    """Config for the Wandb experiment tracker.
-
-    Attributes:
-        entity: Name of an existing wandb entity.
-        project_name: Name of an existing wandb project to log to.
-        api_key: API key to should be authorized to log to the configured wandb
-            entity and project.
-    """
-
-    api_key: str = SecretField()
-    entity: Optional[str] = None
-    project_name: Optional[str] = None
+    """Config for the Wandb experiment tracker."""
+
+    api_key: str = SecretField(
+        description="API key that should be authorized to log to the configured "
+        "Wandb entity and project. Required for authentication."
+    )
+    entity: Optional[str] = Field(
+        None,
+        description="Name of an existing Wandb entity (team or user account) "
+        "to log experiments to."
+    )
+    project_name: Optional[str] = Field(
+        None,
+        description="Name of an existing Wandb project to log experiments to. "
+        "If not specified, a default project will be used."
+    )
 
 
 class WandbExperimentTrackerFlavor(BaseExperimentTrackerFlavor):

zenml/integrations/whylogs/flavors/whylogs_data_validator_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.data_validators.base_data_validator import (
     BaseDataValidatorConfig,
@@ -28,17 +30,18 @@ if TYPE_CHECKING:
 
 
 class WhylogsDataValidatorSettings(BaseSettings):
-    """Settings for the Whylogs data validator.
-
-    Attributes:
-        enable_whylabs: If set to `True` for a step, all the whylogs data
-            profile views returned by the step will automatically be uploaded
-            to the Whylabs platform if Whylabs credentials are configured.
-        dataset_id: Dataset ID to use when uploading profiles to Whylabs.
-    """
-
-    enable_whylabs: bool = False
-    dataset_id: Optional[str] = None
+    """Settings for the Whylogs data validator."""
+
+    enable_whylabs: bool = Field(
+        False,
+        description="If set to `True` for a step, all the whylogs data "
+        "profile views returned by the step will automatically be uploaded "
+        "to the Whylabs platform if Whylabs credentials are configured.",
+    )
+    dataset_id: Optional[str] = Field(
+        None,
+        description="Dataset ID to use when uploading profiles to Whylabs.",
+    )
 
 
 class WhylogsDataValidatorConfig(

zenml/logging/step_logging.py

@@ -49,6 +49,7 @@ from zenml.models import (
     PipelineDeploymentResponse,
     PipelineRunUpdate,
 )
+from zenml.utils.io_utils import sanitize_remote_path
 from zenml.utils.time_utils import utc_now
 from zenml.zen_stores.base_zen_store import BaseZenStore
 
@@ -101,15 +102,14 @@ def prepare_logs_uri(
 
     # Delete the file if it already exists
     if artifact_store.config.IS_IMMUTABLE_FILESYSTEM:
-        logs_uri_folder = os.path.join(logs_base_uri, log_key)
-        if artifact_store.exists(logs_uri_folder):
+        logs_uri = os.path.join(logs_base_uri, log_key)
+        if artifact_store.exists(logs_uri):
             logger.warning(
-                f"Logs directory {logs_uri_folder} already exists! Removing old log directory..."
+                f"Logs directory {logs_uri} already exists! Removing old log directory..."
             )
-            artifact_store.rmtree(logs_uri_folder)
+            artifact_store.rmtree(logs_uri)
 
-        artifact_store.makedirs(logs_uri_folder)
-        return logs_uri_folder
+        artifact_store.makedirs(logs_uri)
     else:
         logs_uri = os.path.join(logs_base_uri, f"{log_key}{LOGS_EXTENSION}")
         if artifact_store.exists(logs_uri):
@@ -117,7 +117,8 @@ def prepare_logs_uri(
                 f"Logs file {logs_uri} already exists! Removing old log file..."
             )
             artifact_store.remove(logs_uri)
-        return logs_uri
+
+    return sanitize_remote_path(logs_uri)
 
 
 def fetch_logs(

zenml/models/v2/core/logs.py

@@ -34,7 +34,8 @@ class LogsRequest(BaseRequest):
     """Request model for logs."""
 
     uri: str = Field(title="The uri of the logs file")
-    source: str = Field(title="The source of the logs file")
+    # TODO: Remove default value when not supporting clients <0.84.0 anymore
+    source: str = Field(default="", title="The source of the logs file")
     artifact_store_id: UUID = Field(
         title="The artifact store ID to associate the logs with.",
     )

zenml/models/v2/core/pipeline_run.py

@@ -24,7 +24,6 @@ from typing import (
     Type,
     TypeVar,
     Union,
-    cast,
 )
 from uuid import UUID
 
@@ -333,64 +332,6 @@ class PipelineRunResponse(
 
         return get_artifacts_versions_of_pipeline_run(self, only_produced=True)
 
-    def refresh_run_status(self) -> "PipelineRunResponse":
-        """Method to refresh the status of a run if it is initializing/running.
-
-        Returns:
-            The updated pipeline.
-
-        Raises:
-            ValueError: If the stack of the run response is None.
-        """
-        if self.status in [
-            ExecutionStatus.INITIALIZING,
-            ExecutionStatus.RUNNING,
-        ]:
-            # Check if the stack still accessible
-            if self.stack is None:
-                raise ValueError(
-                    "The stack that this pipeline run response was executed on"
-                    "is either not accessible or has been deleted."
-                )
-
-            # Create the orchestrator instance
-            from zenml.enums import StackComponentType
-            from zenml.orchestrators.base_orchestrator import BaseOrchestrator
-            from zenml.stack.stack_component import StackComponent
-
-            # Check if the stack still accessible
-            orchestrator_list = self.stack.components.get(
-                StackComponentType.ORCHESTRATOR, []
-            )
-            if len(orchestrator_list) == 0:
-                raise ValueError(
-                    "The orchestrator that this pipeline run response was "
-                    "executed with is either not accessible or has been deleted."
-                )
-
-            orchestrator = cast(
-                BaseOrchestrator,
-                StackComponent.from_model(
-                    component_model=orchestrator_list[0]
-                ),
-            )
-
-            # Fetch the status
-            status = orchestrator.fetch_status(run=self)
-
-            # If it is different from the current status, update it
-            if status != self.status:
-                from zenml.client import Client
-                from zenml.models import PipelineRunUpdate
-
-                client = Client()
-                return client.zen_store.update_run(
-                    run_id=self.id,
-                    run_update=PipelineRunUpdate(status=status),
-                )
-
-        return self
-
     # Body and metadata properties
     @property
     def status(self) -> ExecutionStatus:

zenml/orchestrators/base_orchestrator.py

@@ -22,6 +22,7 @@ from typing import (
     Dict,
     Iterator,
     Optional,
+    Tuple,
     Type,
     cast,
 )
@@ -443,11 +444,16 @@ class BaseOrchestrator(StackComponent, ABC):
         """Cleans up the active run."""
         self._active_deployment = None
 
-    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: A pipeline run response to fetch its status.
+            include_steps: If True, also fetch the status of individual steps.
 
         Raises:
             NotImplementedError: If any orchestrator inheriting from the base

zenml/pipelines/build_utils.py

@@ -38,7 +38,7 @@ from zenml.models import (
     StackResponse,
 )
 from zenml.stack import Stack
-from zenml.utils import source_utils
+from zenml.utils import docker_utils, source_utils
 from zenml.utils.pipeline_docker_image_builder import (
     PipelineDockerImageBuilder,
 )
@@ -378,6 +378,7 @@ def create_pipeline_build(
             if build_config.step_name:
                 tag += f"-{build_config.step_name}"
             tag += f"-{build_config.key}"
+            tag = docker_utils.sanitize_tag(tag)
 
             include_files = build_config.should_include_files(
                 code_repository=code_repository,

zenml/stack/authentication_mixin.py

@@ -15,7 +15,7 @@
 
 from typing import Optional, Type, TypeVar, cast
 
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
 
 from zenml.client import Client
 from zenml.models import SecretResponse
@@ -30,12 +30,13 @@ class AuthenticationConfigMixin(StackComponentConfig):
     Any stack component that implements `AuthenticationMixin` should have a
     config that inherits from this class.
 
-    Attributes:
-        authentication_secret: Name of the secret that stores the
-            authentication credentials.
+    Field descriptions are defined inline using Field() descriptors.
     """
 
-    authentication_secret: Optional[str] = None
+    authentication_secret: Optional[str] = Field(
+        default=None,
+        description="Name of the ZenML secret containing authentication credentials.",
+    )
 
 
 class AuthenticationMixin(StackComponent):