zenml-nightly 0.73.0.dev20250127__py3-none-any.whl → 0.73.0.dev20250129__py3-none-any.whl

zenml/models/v2/core/schedule.py

@@ -13,11 +13,11 @@
 #  permissions and limitations under the License.
 """Models representing schedules."""
 
-import datetime
+from datetime import datetime, timedelta, timezone
 from typing import Dict, Optional, Union
 from uuid import UUID
 
-from pydantic import Field, model_validator
+from pydantic import Field, field_validator, model_validator
 
 from zenml.constants import STR_FIELD_MAX_LENGTH
 from zenml.logger import get_logger
@@ -46,15 +46,36 @@ class ScheduleRequest(WorkspaceScopedRequest):
     active: bool
 
     cron_expression: Optional[str] = None
-    start_time: Optional[datetime.datetime] = None
-    end_time: Optional[datetime.datetime] = None
-    interval_second: Optional[datetime.timedelta] = None
+    start_time: Optional[datetime] = None
+    end_time: Optional[datetime] = None
+    interval_second: Optional[timedelta] = None
     catchup: bool = False
-    run_once_start_time: Optional[datetime.datetime] = None
+    run_once_start_time: Optional[datetime] = None
 
     orchestrator_id: Optional[UUID]
     pipeline_id: Optional[UUID]
 
+    @field_validator(
+        "start_time", "end_time", "run_once_start_time", mode="after"
+    )
+    @classmethod
+    def _ensure_tzunaware_utc(
+        cls, value: Optional[datetime]
+    ) -> Optional[datetime]:
+        """Ensures that all datetimes are timezone unaware and in UTC time.
+
+        Args:
+            value: The datetime.
+
+        Returns:
+            The datetime in UTC time without timezone.
+        """
+        if value and value.tzinfo:
+            value = value.astimezone(timezone.utc)
+            value = value.replace(tzinfo=None)
+
+        return value
+
     @model_validator(mode="after")
     def _ensure_cron_or_periodic_schedule_configured(
         self,
@@ -108,11 +129,11 @@ class ScheduleUpdate(BaseUpdate):
     name: Optional[str] = None
     active: Optional[bool] = None
     cron_expression: Optional[str] = None
-    start_time: Optional[datetime.datetime] = None
-    end_time: Optional[datetime.datetime] = None
-    interval_second: Optional[datetime.timedelta] = None
+    start_time: Optional[datetime] = None
+    end_time: Optional[datetime] = None
+    interval_second: Optional[timedelta] = None
     catchup: Optional[bool] = None
-    run_once_start_time: Optional[datetime.datetime] = None
+    run_once_start_time: Optional[datetime] = None
     orchestrator_id: Optional[UUID] = None
     pipeline_id: Optional[UUID] = None
 
@@ -125,11 +146,11 @@ class ScheduleResponseBody(WorkspaceScopedResponseBody):
 
     active: bool
     cron_expression: Optional[str] = None
-    start_time: Optional[datetime.datetime] = None
-    end_time: Optional[datetime.datetime] = None
-    interval_second: Optional[datetime.timedelta] = None
+    start_time: Optional[datetime] = None
+    end_time: Optional[datetime] = None
+    interval_second: Optional[timedelta] = None
     catchup: bool = False
-    run_once_start_time: Optional[datetime.datetime] = None
+    run_once_start_time: Optional[datetime] = None
 
 
 class ScheduleResponseMetadata(WorkspaceScopedResponseMetadata):
@@ -217,7 +238,7 @@ class ScheduleResponse(
         return self.get_body().cron_expression
 
     @property
-    def start_time(self) -> Optional[datetime.datetime]:
+    def start_time(self) -> Optional[datetime]:
         """The `start_time` property.
 
         Returns:
@@ -226,7 +247,7 @@ class ScheduleResponse(
         return self.get_body().start_time
 
     @property
-    def end_time(self) -> Optional[datetime.datetime]:
+    def end_time(self) -> Optional[datetime]:
         """The `end_time` property.
 
         Returns:
@@ -235,7 +256,7 @@ class ScheduleResponse(
         return self.get_body().end_time
 
     @property
-    def run_once_start_time(self) -> Optional[datetime.datetime]:
+    def run_once_start_time(self) -> Optional[datetime]:
         """The `run_once_start_time` property.
 
         Returns:
@@ -244,7 +265,7 @@ class ScheduleResponse(
         return self.get_body().run_once_start_time
 
     @property
-    def interval_second(self) -> Optional[datetime.timedelta]:
+    def interval_second(self) -> Optional[timedelta]:
         """The `interval_second` property.
 
         Returns:
@@ -313,10 +334,10 @@ class ScheduleFilter(WorkspaceScopedFilter):
         default=None,
         description="The cron expression, describing the schedule",
     )
-    start_time: Optional[Union[datetime.datetime, str]] = Field(
+    start_time: Optional[Union[datetime, str]] = Field(
         default=None, description="Start time", union_mode="left_to_right"
     )
-    end_time: Optional[Union[datetime.datetime, str]] = Field(
+    end_time: Optional[Union[datetime, str]] = Field(
         default=None, description="End time", union_mode="left_to_right"
     )
     interval_second: Optional[Optional[float]] = Field(
@@ -332,7 +353,7 @@ class ScheduleFilter(WorkspaceScopedFilter):
         default=None,
         description="Name of the schedule",
     )
-    run_once_start_time: Optional[Union[datetime.datetime, str]] = Field(
+    run_once_start_time: Optional[Union[datetime, str]] = Field(
         default=None,
         description="The time at which the schedule should run once",
         union_mode="left_to_right",

zenml/pipelines/build_utils.py

@@ -30,7 +30,6 @@ from zenml.code_repositories import BaseCodeRepository
 from zenml.logger import get_logger
 from zenml.models import (
     BuildItem,
-    CodeReferenceRequest,
     PipelineBuildBase,
     PipelineBuildRequest,
     PipelineBuildResponse,
@@ -362,6 +361,7 @@ def create_pipeline_build(
             item_key = checksums[checksum]
             image_name_or_digest = images[item_key].image
             contains_code = images[item_key].contains_code
+            requires_code_download = images[item_key].requires_code_download
             dockerfile = images[item_key].dockerfile
             requirements = images[item_key].requirements
         else:
@@ -373,7 +373,7 @@ def create_pipeline_build(
             include_files = build_config.should_include_files(
                 code_repository=code_repository,
             )
-            download_files = build_config.should_download_files(
+            requires_code_download = build_config.should_download_files(
                 code_repository=code_repository,
             )
             pass_code_repo = (
@@ -391,7 +391,6 @@ def create_pipeline_build(
                 tag=tag,
                 stack=stack,
                 include_files=include_files,
-                download_files=download_files,
                 entrypoint=build_config.entrypoint,
                 extra_files=build_config.extra_files,
                 code_repository=code_repository if pass_code_repo else None,
@@ -404,7 +403,7 @@ def create_pipeline_build(
             requirements=requirements,
             settings_checksum=checksum,
             contains_code=contains_code,
-            requires_code_download=download_files,
+            requires_code_download=requires_code_download,
         )
         checksums[checksum] = combined_key
 
@@ -537,6 +536,14 @@ def verify_local_repository_context(
         )
         code_repository = BaseCodeRepository.from_model(model)
 
+        if will_download_from_code_repository(
+            deployment=deployment, local_repo_context=local_repo_context
+        ):
+            logger.info(
+                "Using code repository `%s` to download code for this run.",
+                model.name,
+            )
+
     return code_repository
 
 
@@ -695,14 +702,15 @@ def compute_stack_checksum(stack: StackResponse) -> str:
 def should_upload_code(
     deployment: PipelineDeploymentBase,
     build: Optional[PipelineBuildResponse],
-    code_reference: Optional[CodeReferenceRequest],
+    can_download_from_code_repository: bool,
 ) -> bool:
     """Checks whether the current code should be uploaded for the deployment.
 
     Args:
         deployment: The deployment.
         build: The build for the deployment.
-        code_reference: The code reference for the deployment.
+        can_download_from_code_repository: Whether the code can be downloaded
+            from a code repository.
 
     Returns:
         Whether the current code should be uploaded for the deployment.
@@ -718,7 +726,7 @@ def should_upload_code(
         docker_settings = step.config.docker_settings
 
         if (
-            code_reference
+            can_download_from_code_repository
             and docker_settings.allow_download_from_code_repository
         ):
             # No upload needed for this step
@@ -728,3 +736,31 @@ def should_upload_code(
             return True
 
     return False
+
+
+def will_download_from_code_repository(
+    deployment: PipelineDeploymentBase,
+    local_repo_context: "LocalRepositoryContext",
+) -> bool:
+    """Checks whether a code repository will be used to download code.
+
+    Args:
+        deployment: The deployment.
+        local_repo_context: The local repository context.
+
+    Returns:
+        Whether a code repository will be used to download code.
+    """
+    if not build_required(deployment=deployment):
+        return False
+
+    if local_repo_context.has_local_changes:
+        return False
+
+    for step in deployment.step_configurations.values():
+        docker_settings = step.config.docker_settings
+
+        if docker_settings.allow_download_from_code_repository:
+            return True
+
+    return False

zenml/pipelines/pipeline_definition.py

@@ -701,6 +701,7 @@ To avoid this consider setting pipeline parameters only in one place (config or
         code_repository = build_utils.verify_local_repository_context(
             deployment=deployment, local_repo_context=local_repo_context
         )
+        can_download_from_code_repository = code_repository is not None
 
         if prevent_build_reuse:
             logger.warning(
@@ -737,7 +738,7 @@ To avoid this consider setting pipeline parameters only in one place (config or
         if build_utils.should_upload_code(
             deployment=deployment,
             build=build_model,
-            code_reference=code_reference,
+            can_download_from_code_repository=can_download_from_code_repository,
         ):
             code_archive = code_utils.CodeArchive(
                 root=source_utils.get_source_root()

zenml/stack/flavor.py

@@ -13,11 +13,13 @@
 #  permissions and limitations under the License.
 """Base ZenML Flavor implementation."""
 
+import os
 from abc import abstractmethod
 from typing import Any, Dict, Optional, Type, cast
 
 from zenml.client import Client
 from zenml.enums import StackComponentType
+from zenml.exceptions import CustomFlavorImportError
 from zenml.models import (
     FlavorRequest,
     FlavorResponse,
@@ -126,10 +128,34 @@ class Flavor:
         Args:
             flavor_model: The model to load from.
 
+        Raises:
+            CustomFlavorImportError: If the custom flavor can't be imported.
+            ImportError: If the flavor can't be imported.
+
         Returns:
             The loaded flavor.
         """
-        flavor = source_utils.load(flavor_model.source)()
+        try:
+            flavor = source_utils.load(flavor_model.source)()
+        except (ModuleNotFoundError, ImportError, NotImplementedError) as err:
+            if flavor_model.is_custom:
+                flavor_module, _ = flavor_model.source.rsplit(".")
+                expected_file_path = os.path.join(
+                    source_utils.get_source_root(),
+                    flavor_module.replace(".", os.path.sep),
+                )
+                raise CustomFlavorImportError(
+                    f"Couldn't import custom flavor {flavor_model.name}: "
+                    f"{err}. Make sure the custom flavor class "
+                    f"`{flavor_model.source}` is importable. If it is part of "
+                    "a library, make sure it is installed. If "
+                    "it is a local code file, make sure it exists at "
+                    f"`{expected_file_path}.py`."
+                )
+            else:
+                raise ImportError(
+                    f"Couldn't import flavor {flavor_model.name}: {err}"
+                )
         return cast(Flavor, flavor)
 
     def to_model(

zenml/stack/stack_component.py

@@ -402,16 +402,10 @@ class StackComponent:
         Raises:
             ImportError: If the flavor can't be imported.
         """
-        flavor_model = component_model.flavor
-
-        try:
-            from zenml.stack import Flavor
+        from zenml.stack import Flavor
 
-            flavor = Flavor.from_model(flavor_model)
-        except (ModuleNotFoundError, ImportError, NotImplementedError) as err:
-            raise ImportError(
-                f"Couldn't import flavor {flavor_model.name}: {err}"
-            )
+        flavor_model = component_model.flavor
+        flavor = Flavor.from_model(flavor_model)
 
         configuration = flavor.config_class(**component_model.configuration)
 

zenml/utils/code_repository_utils.py

@@ -109,9 +109,16 @@ def find_active_code_repository(
     for model in depaginate(list_method=Client().list_code_repositories):
         try:
             repo = BaseCodeRepository.from_model(model)
-        except Exception:
+        except ImportError:
             logger.debug(
-                "Failed to instantiate code repository class.", exc_info=True
+                "Failed to import code repository class.", exc_info=True
+            )
+            continue
+        except Exception as e:
+            logger.warning(
+                "Failed to instantiate or login to code repository `%s`: %s",
+                model.name,
+                e,
             )
             continue
 

zenml/utils/code_utils.py

@@ -249,16 +249,7 @@ def download_and_extract_code(code_path: str, extract_dir: str) -> None:
     Args:
         code_path: Path where the code is uploaded.
         extract_dir: Directory where to code should be extracted to.
-
-    Raises:
-        RuntimeError: If the code is stored in an artifact store which is
-            not active.
     """
-    artifact_store = Client().active_stack.artifact_store
-
-    if not code_path.startswith(artifact_store.path):
-        raise RuntimeError("Code stored in different artifact store.")
-
     download_path = os.path.basename(code_path)
     fileio.copy(code_path, download_path)
 
@@ -266,17 +257,30 @@ def download_and_extract_code(code_path: str, extract_dir: str) -> None:
     os.remove(download_path)
 
 
-def download_code_from_artifact_store(code_path: str) -> None:
+def download_code_from_artifact_store(
+    code_path: str, artifact_store: "BaseArtifactStore"
+) -> None:
     """Download code from the artifact store.
 
     Args:
         code_path: Path where the code is stored.
+        artifact_store: The artifact store to use for the download.
+
+    Raises:
+        RuntimeError: If the code is stored in an artifact store which is
+            not active.
     """
     logger.info("Downloading code from artifact store path `%s`.", code_path)
 
-    # Do not remove this line, we need to instantiate the artifact store to
-    # register the filesystem needed for the file download
-    _ = Client().active_stack.artifact_store
+    if not code_path.startswith(artifact_store.path):
+        raise RuntimeError(
+            "The code is not stored in the artifact store "
+            f"{artifact_store.name} that was passed to download it."
+        )
+
+    # Make sure we register the artifact store filesystem here so the
+    # fileio.copy call will pick up the right credentials
+    artifact_store._register()
 
     extract_dir = os.path.abspath("code")
     os.makedirs(extract_dir)

zenml/utils/dashboard_utils.py

@@ -13,6 +13,7 @@
 #  permissions and limitations under the License.
 """Utility class to help with interacting with the dashboard."""
 
+import os
 from typing import Optional
 from uuid import UUID
 
@@ -156,9 +157,27 @@ def show_dashboard(url: str) -> None:
         display(IFrame(src=url, width="100%", height=720))
 
     elif environment in (EnvironmentType.NATIVE, EnvironmentType.WSL):
-        if constants.handle_bool_env_var(
+        open_dashboard = True
+
+        if constants.ENV_AUTO_OPEN_DASHBOARD in os.environ:
+            logger.warning(
+                "The `%s` environment variable is deprecated, use the `%s` "
+                "environment variable instead.",
+                constants.ENV_AUTO_OPEN_DASHBOARD,
+                constants.ENV_ZENML_AUTO_OPEN_DASHBOARD,
+            )
+
+        if not constants.handle_bool_env_var(
             constants.ENV_AUTO_OPEN_DASHBOARD, default=True
         ):
+            open_dashboard = False
+
+        if not constants.handle_bool_env_var(
+            constants.ENV_ZENML_AUTO_OPEN_DASHBOARD, default=True
+        ):
+            open_dashboard = False
+
+        if open_dashboard:
             try:
                 import webbrowser
 
@@ -169,14 +188,16 @@ def show_dashboard(url: str) -> None:
                 logger.info(
                     "Automatically opening the dashboard in your "
                     "browser. To disable this, set the env variable "
-                    "AUTO_OPEN_DASHBOARD=false."
+                    "`%s=false`.",
+                    constants.ENV_ZENML_AUTO_OPEN_DASHBOARD,
                 )
             except Exception as e:
                 logger.error(e)
         else:
             logger.info(
                 "To open the dashboard in a browser automatically, "
-                "set the env variable AUTO_OPEN_DASHBOARD=true."
+                "set the env variable `%s=true`.",
+                constants.ENV_ZENML_AUTO_OPEN_DASHBOARD,
             )
 
     else:

zenml/utils/deprecation_utils.py

@@ -132,14 +132,16 @@ def deprecate_pydantic_attributes(
             if replacement_attribute is None:
                 _warn(
                     message=f"The attribute `{deprecated_attribute}` of class "
-                    f"`{cls.__name__}` will be deprecated soon.",
+                    f"`{cls.__name__}` is deprecated and will be removed in "
+                    "the future.",
                     attribute=deprecated_attribute,
                 )
                 continue
 
             _warn(
                 message=f"The attribute `{deprecated_attribute}` of class "
-                f"`{cls.__name__}` will be deprecated soon. Use the "
+                f"`{cls.__name__}` is deprecated and will be removed in the "
+                "future. Use the "
                 f"attribute `{replacement_attribute}` instead.",
                 attribute=deprecated_attribute,
             )

zenml/utils/pipeline_docker_image_builder.py

@@ -78,7 +78,6 @@ class PipelineDockerImageBuilder:
         tag: str,
         stack: "Stack",
         include_files: bool,
-        download_files: bool,
         entrypoint: Optional[str] = None,
         extra_files: Optional[Dict[str, str]] = None,
         code_repository: Optional["BaseCodeRepository"] = None,
@@ -93,7 +92,6 @@ class PipelineDockerImageBuilder:
             tag: The tag to use for the image.
             stack: The stack on which the pipeline will be deployed.
             include_files: Whether to include files in the build context.
-            download_files: Whether to download files in the build context.
             entrypoint: Entrypoint to use for the final image. If left empty,
                 no entrypoint will be included in the image.
             extra_files: Extra files to add to the build context. Keys are the
@@ -165,7 +163,6 @@ class PipelineDockerImageBuilder:
                 docker_settings.apt_packages,
                 docker_settings.environment,
                 include_files,
-                download_files,
                 entrypoint,
                 extra_files,
             ]
@@ -590,7 +587,7 @@ class PipelineDockerImageBuilder:
             f"ENV {ENV_ZENML_LOGGING_COLORS_DISABLED}={str(handle_bool_env_var(ENV_ZENML_LOGGING_COLORS_DISABLED, False))}"
         )
         for key, value in docker_settings.environment.items():
-            lines.append(f"ENV {key.upper()}={value}")
+            lines.append(f"ENV {key.upper()}='{value}'")
 
         if apt_packages:
             apt_packages = " ".join(f"'{p}'" for p in apt_packages)

zenml/zen_stores/schemas/code_repository_schemas.py

@@ -151,6 +151,9 @@ class CodeRepositorySchema(NamedSchema, table=True):
         if update.logo_url:
             self.logo_url = update.logo_url
 
+        if update.config:
+            self.config = json.dumps(update.config)
+
         self.updated = utc_now()
         return self
 

{zenml_nightly-0.73.0.dev20250127.dist-info → zenml_nightly-0.73.0.dev20250129.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: zenml-nightly
-Version: 0.73.0.dev20250127
+Version: 0.73.0.dev20250129
 Summary: ZenML: Write production-ready ML code.
 License: Apache-2.0
 Keywords: machine learning,production,pipeline,mlops,devops