zenml-nightly 0.55.0.dev20240124__py3-none-any.whl → 0.72.0.dev20250116__py3-none-any.whl

zenml/config/pipeline_configurations.py

@@ -12,12 +12,15 @@
 #  or implied. See the License for the specific language governing
 #  permissions and limitations under the License.
 """Pipeline configuration classes."""
-from typing import TYPE_CHECKING, Any, Dict, Optional
 
-from pydantic import validator
+from datetime import datetime
+from typing import TYPE_CHECKING, Any, Dict, List, Optional
+
+from pydantic import SerializeAsAny, field_validator
 
 from zenml.config.constants import DOCKER_SETTINGS_KEY
-from zenml.config.source import Source, convert_source_validator
+from zenml.config.retry_config import StepRetryConfig
+from zenml.config.source import SourceWithValidator
 from zenml.config.strict_base_model import StrictBaseModel
 from zenml.model.model import Model
 
@@ -36,16 +39,33 @@ class PipelineConfigurationUpdate(StrictBaseModel):
     enable_artifact_metadata: Optional[bool] = None
     enable_artifact_visualization: Optional[bool] = None
     enable_step_logs: Optional[bool] = None
-    settings: Dict[str, BaseSettings] = {}
+    settings: Dict[str, SerializeAsAny[BaseSettings]] = {}
+    tags: Optional[List[str]] = None
     extra: Dict[str, Any] = {}
-    failure_hook_source: Optional[Source] = None
-    success_hook_source: Optional[Source] = None
+    failure_hook_source: Optional[SourceWithValidator] = None
+    success_hook_source: Optional[SourceWithValidator] = None
     model: Optional[Model] = None
     parameters: Optional[Dict[str, Any]] = None
+    retry: Optional[StepRetryConfig] = None
+    substitutions: Dict[str, str] = {}
+
+    def _get_full_substitutions(
+        self, start_time: Optional[datetime]
+    ) -> Dict[str, str]:
+        """Returns the full substitutions dict.
+
+        Args:
+            start_time: Start time of the pipeline run.
 
-    _convert_source = convert_source_validator(
-        "failure_hook_source", "success_hook_source"
-    )
+        Returns:
+            The full substitutions dict including date and time.
+        """
+        if start_time is None:
+            start_time = datetime.utcnow()
+        ret = self.substitutions.copy()
+        ret.setdefault("date", start_time.strftime("%Y_%m_%d"))
+        ret.setdefault("time", start_time.strftime("%H_%M_%S_%f"))
+        return ret
 
 
 class PipelineConfiguration(PipelineConfigurationUpdate):
@@ -53,7 +73,8 @@ class PipelineConfiguration(PipelineConfigurationUpdate):
 
     name: str
 
-    @validator("name")
+    @field_validator("name")
+    @classmethod
     def ensure_pipeline_name_allowed(cls, name: str) -> str:
         """Ensures the pipeline name is allowed.
 
@@ -85,4 +106,4 @@ class PipelineConfiguration(PipelineConfigurationUpdate):
         model_or_dict: SettingsOrDict = self.settings.get(
             DOCKER_SETTINGS_KEY, {}
         )
-        return DockerSettings.parse_obj(model_or_dict)
+        return DockerSettings.model_validate(model_or_dict)

zenml/config/pipeline_run_configuration.py

@@ -12,11 +12,16 @@
 #  or implied. See the License for the specific language governing
 #  permissions and limitations under the License.
 """Pipeline run configuration class."""
-from typing import Any, Dict, Optional, Union
+
+from typing import Any, Dict, List, Optional, Union
 from uuid import UUID
 
+from pydantic import Field, SerializeAsAny
+
 from zenml.config.base_settings import BaseSettings
+from zenml.config.retry_config import StepRetryConfig
 from zenml.config.schedule import Schedule
+from zenml.config.source import SourceWithValidator
 from zenml.config.step_configurations import StepConfigurationUpdate
 from zenml.config.strict_base_model import StrictBaseModel
 from zenml.model.model import Model
@@ -35,9 +40,16 @@ class PipelineRunConfiguration(
     enable_artifact_visualization: Optional[bool] = None
     enable_step_logs: Optional[bool] = None
     schedule: Optional[Schedule] = None
-    build: Union[PipelineBuildBase, UUID, None] = None
+    build: Union[PipelineBuildBase, UUID, None] = Field(
+        default=None, union_mode="left_to_right"
+    )
     steps: Dict[str, StepConfigurationUpdate] = {}
-    settings: Dict[str, BaseSettings] = {}
+    settings: Dict[str, SerializeAsAny[BaseSettings]] = {}
+    tags: Optional[List[str]] = None
     extra: Dict[str, Any] = {}
     model: Optional[Model] = None
     parameters: Optional[Dict[str, Any]] = None
+    retry: Optional[StepRetryConfig] = None
+    failure_hook_source: Optional[SourceWithValidator] = None
+    success_hook_source: Optional[SourceWithValidator] = None
+    substitutions: Dict[str, str] = {}

zenml/config/pipeline_spec.py

@@ -12,14 +12,14 @@
 #  or implied. See the License for the specific language governing
 #  permissions and limitations under the License.
 """Pipeline configuration classes."""
+
 import json
 from typing import Any, Dict, List, Optional
 
-from pydantic.json import pydantic_encoder
-
-from zenml.config.source import Source
+from zenml.config.source import Source, SourceWithValidator
 from zenml.config.step_configurations import StepSpec
 from zenml.config.strict_base_model import StrictBaseModel
+from zenml.utils.json_utils import pydantic_encoder
 
 
 class PipelineSpec(StrictBaseModel):
@@ -34,7 +34,7 @@ class PipelineSpec(StrictBaseModel):
     # - 0.4: New Pipeline class, the upstream steps and
     #   inputs in the step specs refer to the pipeline parameter names
     version: str = "0.4"
-    source: Optional[Source] = None
+    source: Optional[SourceWithValidator] = None
     parameters: Dict[str, Any] = {}
     steps: List[StepSpec]
 
@@ -60,13 +60,13 @@ class PipelineSpec(StrictBaseModel):
         """
         from packaging import version
 
-        dict_ = self.dict()
+        dict_ = self.model_dump()
 
         if self.source:
             dict_["source"] = self.source.import_path
 
         for step_dict in dict_["steps"]:
-            step_dict["source"] = Source.parse_obj(
+            step_dict["source"] = Source.model_validate(
                 step_dict["source"]
             ).import_path
 

zenml/config/resource_settings.py

@@ -16,7 +16,8 @@
 from enum import Enum
 from typing import Optional, Union
 
-from pydantic import Extra, Field, NonNegativeInt, PositiveFloat
+from pydantic import Field, NonNegativeInt, PositiveFloat
+from pydantic_settings import SettingsConfigDict
 
 from zenml.config.base_settings import BaseSettings
 
@@ -70,7 +71,7 @@ class ResourceSettings(BaseSettings):
 
     cpu_count: Optional[PositiveFloat] = None
     gpu_count: Optional[NonNegativeInt] = None
-    memory: Optional[str] = Field(regex=MEMORY_REGEX)
+    memory: Optional[str] = Field(pattern=MEMORY_REGEX, default=None)
 
     @property
     def empty(self) -> bool:
@@ -82,7 +83,7 @@ class ResourceSettings(BaseSettings):
         # To detect whether this config is empty (= no values specified), we
         # check if there are any attributes which are explicitly set to any
         # value other than `None`.
-        return len(self.dict(exclude_unset=True, exclude_none=True)) == 0
+        return len(self.model_dump(exclude_unset=True, exclude_none=True)) == 0
 
     def get_memory(
         self, unit: Union[str, ByteUnit] = ByteUnit.GB
@@ -114,11 +115,9 @@ class ResourceSettings(BaseSettings):
             # Should never happen due to the regex validation
             raise ValueError(f"Unable to parse memory unit from '{memory}'.")
 
-    class Config:
-        """Pydantic configuration class."""
-
+    model_config = SettingsConfigDict(
         # public attributes are immutable
-        allow_mutation = False
-
+        frozen=True,
         # prevent extra attributes during model initialization
-        extra = Extra.forbid
+        extra="forbid",
+    )

zenml/config/retry_config.py

@@ -0,0 +1,27 @@
+#  Copyright (c) ZenML GmbH 2024. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at:
+#
+#       https://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
+#  or implied. See the License for the specific language governing
+#  permissions and limitations under the License.
+"""Retry configuration for a step."""
+
+from zenml.config.strict_base_model import StrictBaseModel
+
+
+class StepRetryConfig(StrictBaseModel):
+    """Retry configuration for a step.
+
+    Delay is an integer (specified in seconds).
+    """
+
+    max_retries: int = 1
+    delay: int = 0  # in seconds
+    backoff: int = 0

zenml/config/schedule.py

@@ -14,9 +14,9 @@
 """Class for defining a pipeline schedule."""
 
 import datetime
-from typing import Any, Dict, Optional
+from typing import Optional
 
-from pydantic import BaseModel, root_validator
+from pydantic import BaseModel, model_validator
 
 from zenml.logger import get_logger
 
@@ -43,6 +43,8 @@ class Schedule(BaseModel):
             schedules the latest interval if more than one interval is ready to
             be scheduled. Usually, if your pipeline handles backfill
             internally, you should turn catchup off to avoid duplicate backfill.
+        run_once_start_time: datetime object to indicate when to run the
+            pipeline once. This is useful for one-off runs.
     """
 
     name: Optional[str] = None
@@ -51,16 +53,12 @@ class Schedule(BaseModel):
     end_time: Optional[datetime.datetime] = None
     interval_second: Optional[datetime.timedelta] = None
     catchup: bool = False
+    run_once_start_time: Optional[datetime.datetime] = None
 
-    @root_validator
-    def _ensure_cron_or_periodic_schedule_configured(
-        cls, values: Dict[str, Any]
-    ) -> Dict[str, Any]:
+    @model_validator(mode="after")
+    def _ensure_cron_or_periodic_schedule_configured(self) -> "Schedule":
         """Ensures that the cron expression or start time + interval are set.
 
-        Args:
-            values: All attributes of the schedule.
-
         Returns:
             All schedule attributes.
 
@@ -68,12 +66,9 @@ class Schedule(BaseModel):
             ValueError: If no cron expression or start time + interval were
                 provided.
         """
-        cron_expression = values.get("cron_expression")
-        periodic_schedule = values.get("start_time") and values.get(
-            "interval_second"
-        )
+        periodic_schedule = self.start_time and self.interval_second
 
-        if cron_expression and periodic_schedule:
+        if self.cron_expression and periodic_schedule:
             logger.warning(
                 "This schedule was created with a cron expression as well as "
                 "values for `start_time` and `interval_seconds`. The resulting "
@@ -81,12 +76,26 @@ class Schedule(BaseModel):
                 "but will usually ignore the interval and use the cron "
                 "expression."
             )
-            return values
-        elif cron_expression or periodic_schedule:
-            return values
+
+            return self
+        elif self.cron_expression and self.run_once_start_time:
+            logger.warning(
+                "This schedule was created with a cron expression as well as "
+                "a value for `run_once_start_time`. The resulting behavior "
+                "depends on the concrete orchestrator implementation but will "
+                "usually ignore the `run_once_start_time`."
+            )
+            return self
+        elif (
+            self.cron_expression
+            or periodic_schedule
+            or self.run_once_start_time
+        ):
+            return self
         else:
             raise ValueError(
-                "Either a cron expression or start time and interval seconds "
+                "Either a cron expression, a start time and interval seconds "
+                "or a run once start time "
                 "need to be set for a valid schedule."
             )
 

zenml/config/secret_reference_mixin.py

@@ -12,12 +12,14 @@
 #  or implied. See the License for the specific language governing
 #  permissions and limitations under the License.
 """Secret reference mixin implementation."""
+
 from typing import TYPE_CHECKING, Any, Set
 
 from pydantic import BaseModel
 
 from zenml.logger import get_logger
 from zenml.utils import secret_utils
+from zenml.utils.pydantic_utils import has_validators
 
 logger = get_logger(__name__)
 
@@ -44,7 +46,7 @@ class SecretReferenceMixin(BaseModel):
         """
         for key, value in kwargs.items():
             try:
-                field = self.__class__.__fields__[key]
+                field = self.__class__.model_fields[key]
             except KeyError:
                 # Value for a private attribute or non-existing field, this
                 # will fail during the upcoming pydantic validation
@@ -64,7 +66,7 @@ class SecretReferenceMixin(BaseModel):
                         "but future versions of ZenML will require you to pass "
                         "in sensitive information as secrets. Check out the "
                         "documentation on how to configure values with secrets "
-                        "here: https://docs.zenml.io/user-guide/advanced-guide/secret-management"
+                        "here: https://docs.zenml.io/getting-started/deploying-zenml/secret-management"
                     )
                 continue
 
@@ -74,7 +76,9 @@ class SecretReferenceMixin(BaseModel):
                     "not allowed."
                 )
 
-            requires_validation = field.pre_validators or field.post_validators
+            requires_validation = has_validators(
+                pydantic_class=self.__class__, field_name=key
+            )
             if requires_validation:
                 raise ValueError(
                     f"Passing the attribute `{key}` as a secret reference is "
@@ -145,6 +149,6 @@ class SecretReferenceMixin(BaseModel):
         """
         return {
             secret_utils.parse_secret_reference(v)
-            for v in self.dict().values()
+            for v in self.model_dump().values()
             if secret_utils.is_secret_reference(v)
         }

zenml/config/secrets_store_config.py

@@ -13,10 +13,9 @@
 #  permissions and limitations under the License.
 """Functionality to support ZenML secrets store configurations."""
 
+from typing import Optional
 
-from typing import Any, Dict, Optional
-
-from pydantic import BaseModel, root_validator
+from pydantic import BaseModel, ConfigDict, model_validator
 
 from zenml.enums import SecretsStoreType
 from zenml.logger import get_logger
@@ -42,45 +41,38 @@ class SecretsStoreConfiguration(BaseModel):
     type: SecretsStoreType
     class_path: Optional[str] = None
 
-    @root_validator
-    def validate_custom(cls, values: Dict[str, Any]) -> Dict[str, Any]:
+    @model_validator(mode="after")
+    def validate_custom(self) -> "SecretsStoreConfiguration":
         """Validate that class_path is set for custom secrets stores.
 
-        Args:
-            values: Dict representing user-specified runtime settings.
-
         Returns:
             Validated settings.
 
         Raises:
-            ValueError: If class_path is not set when using an custom secrets
+            ValueError: If class_path is not set when using a custom secrets
                 store.
         """
-        if not values.get("type"):
-            return values
-        if values["type"] == SecretsStoreType.CUSTOM:
-            if values["class_path"] is None:
+        if not self.type:
+            return self
+        if self.type == SecretsStoreType.CUSTOM:
+            if self.class_path is None:
                 raise ValueError(
                     "A class_path must be set when using a custom secrets "
                     "store implementation."
                 )
-        elif values["class_path"] is not None:
+        elif self.class_path is not None:
             raise ValueError(
                 f"The class_path attribute is not supported for the "
-                f"{values['type']} secrets store type."
+                f"{self.type} secrets store type."
             )
 
-        return values
-
-    class Config:
-        """Pydantic configuration class."""
+        return self
 
+    model_config = ConfigDict(
         # Validate attributes when assigning them. We need to set this in order
         # to have a mix of mutable and immutable attributes
-        validate_assignment = True
+        validate_assignment=True,
         # Allow extra attributes to be set in the base class. The concrete
         # classes are responsible for validating the attributes.
-        extra = "allow"
-        # all attributes with leading underscore are private and therefore
-        # are mutable and not included in serialization
-        underscore_attrs_are_private = True
+        extra="allow",
+    )