zenml-nightly 0.73.0.dev20250123__py3-none-any.whl → 0.73.0.dev20250125__py3-none-any.whl

zenml/integrations/kubernetes/orchestrators/kube_utils.py

@@ -31,7 +31,6 @@ Internal interface: no backwards compatibility guarantees.
 Adjusted from https://github.com/tensorflow/tfx/blob/master/tfx/utils/kube_utils.py.
 """
 
-import datetime
 import enum
 import re
 import time
@@ -47,6 +46,7 @@ from zenml.integrations.kubernetes.orchestrators.manifest_utils import (
     build_service_account_manifest,
 )
 from zenml.logger import get_logger
+from zenml.utils.time_utils import utc_now
 
 logger = get_logger(__name__)
 
@@ -248,7 +248,7 @@ def wait_pod(
     Returns:
         The pod object which meets the exit condition.
     """
-    start_time = datetime.datetime.now(datetime.timezone.utc)
+    start_time = utc_now()
 
     # Link to exponential back-off algorithm used here:
     # https://cloud.google.com/storage/docs/exponential-backoff
@@ -288,7 +288,7 @@ def wait_pod(
             return resp
 
         # Check if wait timed out.
-        elapse_time = datetime.datetime.now(datetime.timezone.utc) - start_time
+        elapse_time = utc_now() - start_time
         if elapse_time.seconds >= timeout_sec and timeout_sec != 0:
             raise RuntimeError(
                 f"Waiting for pod `{namespace}:{pod_name}` timed out after "

zenml/integrations/kubernetes/service_connectors/kubernetes_service_connector.py

@@ -504,13 +504,17 @@ class KubernetesServiceConnector(ServiceConnector):
                 ).decode("utf-8")
                 if kube_config.ssl_ca_cert
                 else None,
-                cluster_name=kube_config.host.strip("https://").split(":")[0],
+                cluster_name=kube_config.host.removeprefix("https://").split(
+                    ":"
+                )[0],
                 insecure=kube_config.verify_ssl is False,
             )
         else:
             token: Optional[str] = None
             if kube_config.api_key:
-                token = kube_config.api_key["authorization"].strip("Bearer ")
+                token = kube_config.api_key["authorization"].removeprefix(
+                    "Bearer "
+                )
 
             auth_method = KubernetesAuthenticationMethods.TOKEN
             auth_config = KubernetesTokenConfig(

zenml/integrations/whylogs/data_validators/whylogs_data_validator.py

@@ -34,6 +34,7 @@ from zenml.integrations.whylogs.secret_schemas.whylabs_secret_schema import (
 )
 from zenml.logger import get_logger
 from zenml.stack.authentication_mixin import AuthenticationMixin
+from zenml.utils.time_utils import utc_now
 
 logger = get_logger(__name__)
 
@@ -97,9 +98,7 @@ class WhylogsDataValidator(BaseDataValidator, AuthenticationMixin):
         """
         results = why.log(pandas=dataset)
         profile = results.profile()
-        dataset_timestamp = dataset_timestamp or datetime.datetime.now(
-            datetime.timezone.utc
-        )
+        dataset_timestamp = dataset_timestamp or utc_now()
         profile.set_dataset_timestamp(dataset_timestamp=dataset_timestamp)
         return profile.view()
 

zenml/logging/step_logging.py

@@ -13,7 +13,6 @@
 #  permissions and limitations under the License.
 """ZenML logging handler."""
 
-import datetime
 import os
 import re
 import sys
@@ -35,6 +34,7 @@ from zenml.logging import (
     STEP_LOGS_STORAGE_MAX_MESSAGES,
     STEP_LOGS_STORAGE_MERGE_INTERVAL_SECONDS,
 )
+from zenml.utils.time_utils import utc_now
 from zenml.zen_stores.base_zen_store import BaseZenStore
 
 # Get the logger
@@ -303,9 +303,9 @@ class StepLogsStorage:
                             "w",
                         ) as file:
                             for message in self.buffer:
-                                timestamp = datetime.datetime.now(
-                                    datetime.timezone.utc
-                                ).strftime("%Y-%m-%d %H:%M:%S")
+                                timestamp = utc_now().strftime(
+                                    "%Y-%m-%d %H:%M:%S"
+                                )
                                 file.write(
                                     f"[{timestamp} UTC] {remove_ansi_escape_codes(message)}\n"
                                 )
@@ -314,9 +314,9 @@ class StepLogsStorage:
                             self.logs_uri, "a"
                         ) as file:
                             for message in self.buffer:
-                                timestamp = datetime.datetime.now(
-                                    datetime.timezone.utc
-                                ).strftime("%Y-%m-%d %H:%M:%S")
+                                timestamp = utc_now().strftime(
+                                    "%Y-%m-%d %H:%M:%S"
+                                )
                                 file.write(
                                     f"[{timestamp} UTC] {remove_ansi_escape_codes(message)}\n"
                                 )

zenml/login/credentials.py

@@ -13,7 +13,7 @@
 #  permissions and limitations under the License.
 """ZenML login credentials models."""
 
-from datetime import datetime, timedelta, timezone
+from datetime import datetime, timedelta
 from typing import Any, Dict, Optional, Union
 from urllib.parse import urlparse
 from uuid import UUID
@@ -27,6 +27,7 @@ from zenml.models.v2.misc.server_models import ServerDeploymentType
 from zenml.services.service_status import ServiceState
 from zenml.utils.enum_utils import StrEnum
 from zenml.utils.string_utils import get_human_readable_time
+from zenml.utils.time_utils import to_local_tz, utc_now
 
 
 class ServerType(StrEnum):
@@ -71,7 +72,7 @@ class APIToken(BaseModel):
         expires_at = self.expires_at_with_leeway
         if not expires_at:
             return False
-        return expires_at < datetime.now(timezone.utc)
+        return expires_at < utc_now(tz_aware=expires_at)
 
     model_config = ConfigDict(
         # Allow extra attributes to allow backwards compatibility
@@ -223,7 +224,7 @@ class ServerCredentials(BaseModel):
         expires_at = self.api_token.expires_at_with_leeway
         if not expires_at:
             return "never expires"
-        if expires_at < datetime.now(timezone.utc):
+        if expires_at < utc_now(tz_aware=expires_at):
             return "expired at " + self.expires_at
 
         return f"valid until {self.expires_at} (in {self.expires_in})"
@@ -242,7 +243,7 @@ class ServerCredentials(BaseModel):
             return "never"
 
         # Convert the date in the local timezone
-        local_expires_at = expires_at.astimezone()
+        local_expires_at = to_local_tz(expires_at)
         return local_expires_at.strftime("%Y-%m-%d %H:%M:%S %Z")
 
     @property
@@ -259,7 +260,7 @@ class ServerCredentials(BaseModel):
             return "never"
 
         # Get the time remaining until the token expires
-        expires_in = expires_at - datetime.now(timezone.utc)
+        expires_in = expires_at - utc_now(tz_aware=expires_at)
         return get_human_readable_time(expires_in.total_seconds())
 
     @property

zenml/login/credentials_store.py

@@ -14,7 +14,7 @@
 """ZenML login credentials store support."""
 
 import os
-from datetime import datetime, timedelta, timezone
+from datetime import timedelta
 from typing import Dict, List, Optional, Tuple, Union, cast
 
 from zenml.config.global_config import GlobalConfiguration
@@ -29,6 +29,7 @@ from zenml.login.pro.tenant.models import TenantRead
 from zenml.models import OAuthTokenResponse, ServerModel
 from zenml.utils import yaml_utils
 from zenml.utils.singleton import SingletonMetaClass
+from zenml.utils.time_utils import utc_now_tz_aware
 
 logger = get_logger(__name__)
 
@@ -190,7 +191,7 @@ class CredentialsStore(metaclass=SingletonMetaClass):
                 not credential.api_token.expires_at
                 or credential.api_token.expires_at
                 + timedelta(seconds=TOKEN_STORE_EVICTION_TIME)
-                > datetime.now(timezone.utc)
+                > utc_now_tz_aware()
             )
         }
         yaml_utils.write_yaml(credentials_file, credentials_store)
@@ -477,7 +478,7 @@ class CredentialsStore(metaclass=SingletonMetaClass):
         """
         self.check_and_reload_from_file()
         if token_response.expires_in:
-            expires_at = datetime.now(timezone.utc) + timedelta(
+            expires_at = utc_now_tz_aware() + timedelta(
                 seconds=token_response.expires_in
             )
             # Best practice to calculate the leeway depending on the token

zenml/models/v2/core/api_key.py

@@ -13,7 +13,7 @@
 #  permissions and limitations under the License.
 """Models representing API keys."""
 
-from datetime import datetime, timedelta, timezone
+from datetime import datetime, timedelta
 from typing import TYPE_CHECKING, ClassVar, List, Optional, Type, Union
 from uuid import UUID
 
@@ -35,6 +35,7 @@ from zenml.models.v2.base.base import (
 )
 from zenml.models.v2.base.filter import AnyQuery, BaseFilter
 from zenml.utils.string_utils import b64_decode, b64_encode
+from zenml.utils.time_utils import utc_now
 
 if TYPE_CHECKING:
     from zenml.models.v2.base.filter import AnySchema
@@ -319,7 +320,9 @@ class APIKeyInternalResponse(APIKeyResponse):
             and self.retain_period_minutes > 0
         ):
             # check if the previous key is still valid
-            if datetime.now(timezone.utc) - self.last_rotated < timedelta(
+            if utc_now(
+                tz_aware=self.last_rotated
+            ) - self.last_rotated < timedelta(
                 minutes=self.retain_period_minutes
             ):
                 key_hash = self.previous_key

zenml/models/v2/core/schedule.py

@@ -14,13 +14,14 @@
 """Models representing schedules."""
 
 import datetime
-from typing import Optional, Union
+from typing import Dict, Optional, Union
 from uuid import UUID
 
 from pydantic import Field, model_validator
 
 from zenml.constants import STR_FIELD_MAX_LENGTH
 from zenml.logger import get_logger
+from zenml.metadata.metadata_types import MetadataType
 from zenml.models.v2.base.base import BaseUpdate
 from zenml.models.v2.base.scoped import (
     WorkspaceScopedFilter,
@@ -30,6 +31,7 @@ from zenml.models.v2.base.scoped import (
     WorkspaceScopedResponseMetadata,
     WorkspaceScopedResponseResources,
 )
+from zenml.utils.time_utils import to_utc_timezone
 
 logger = get_logger(__name__)
 
@@ -136,6 +138,11 @@ class ScheduleResponseMetadata(WorkspaceScopedResponseMetadata):
     orchestrator_id: Optional[UUID]
     pipeline_id: Optional[UUID]
 
+    run_metadata: Dict[str, MetadataType] = Field(
+        title="Metadata associated with this schedule.",
+        default={},
+    )
+
 
 class ScheduleResponseResources(WorkspaceScopedResponseResources):
     """Class for all resource models associated with the schedule entity."""
@@ -176,7 +183,7 @@ class ScheduleResponse(
         if not self.start_time:
             return None
 
-        return self.start_time.astimezone(datetime.timezone.utc).isoformat()
+        return to_utc_timezone(self.start_time).isoformat()
 
     @property
     def utc_end_time(self) -> Optional[str]:
@@ -188,7 +195,7 @@ class ScheduleResponse(
         if not self.end_time:
             return None
 
-        return self.end_time.astimezone(datetime.timezone.utc).isoformat()
+        return to_utc_timezone(self.end_time).isoformat()
 
     # Body and metadata properties
     @property
@@ -272,6 +279,15 @@ class ScheduleResponse(
         """
         return self.get_metadata().pipeline_id
 
+    @property
+    def run_metadata(self) -> Dict[str, MetadataType]:
+        """The `run_metadata` property.
+
+        Returns:
+            the value of the property.
+        """
+        return self.get_metadata().run_metadata
+
 
 # ------------------ Filter Model ------------------
 

zenml/orchestrators/publish_utils.py

@@ -13,7 +13,6 @@
 #  permissions and limitations under the License.
 """Utilities to publish pipeline and step runs."""
 
-from datetime import datetime, timezone
 from typing import TYPE_CHECKING, Dict, List
 
 from zenml.client import Client
@@ -25,6 +24,7 @@ from zenml.models import (
     StepRunResponse,
     StepRunUpdate,
 )
+from zenml.utils.time_utils import utc_now
 
 if TYPE_CHECKING:
     from uuid import UUID
@@ -48,7 +48,7 @@ def publish_successful_step_run(
         step_run_id=step_run_id,
         step_run_update=StepRunUpdate(
             status=ExecutionStatus.COMPLETED,
-            end_time=datetime.now(timezone.utc),
+            end_time=utc_now(),
             outputs=output_artifact_ids,
         ),
     )
@@ -67,7 +67,7 @@ def publish_failed_step_run(step_run_id: "UUID") -> "StepRunResponse":
         step_run_id=step_run_id,
         step_run_update=StepRunUpdate(
             status=ExecutionStatus.FAILED,
-            end_time=datetime.now(timezone.utc),
+            end_time=utc_now(),
         ),
     )
 
@@ -87,7 +87,7 @@ def publish_failed_pipeline_run(
         run_id=pipeline_run_id,
         run_update=PipelineRunUpdate(
             status=ExecutionStatus.FAILED,
-            end_time=datetime.now(timezone.utc),
+            end_time=utc_now(),
         ),
     )
 

zenml/orchestrators/step_launcher.py

@@ -16,7 +16,6 @@
 import os
 import time
 from contextlib import nullcontext
-from datetime import datetime, timezone
 from functools import partial
 from typing import TYPE_CHECKING, Any, Callable, Dict, Tuple
 
@@ -45,6 +44,7 @@ from zenml.orchestrators import utils as orchestrator_utils
 from zenml.orchestrators.step_runner import StepRunner
 from zenml.stack import Stack
 from zenml.utils import string_utils
+from zenml.utils.time_utils import utc_now
 
 if TYPE_CHECKING:
     from zenml.step_operators import BaseStepOperator
@@ -201,7 +201,7 @@ class StepLauncher:
                         f"Failed preparing step `{self._step_name}`."
                     )
                     step_run_request.status = ExecutionStatus.FAILED
-                    step_run_request.end_time = datetime.now(timezone.utc)
+                    step_run_request.end_time = utc_now()
                     raise
                 finally:
                     step_run = Client().zen_store.create_run_step(
@@ -305,7 +305,7 @@ class StepLauncher:
             The created or existing pipeline run,
             and a boolean indicating whether the run was created or reused.
         """
-        start_time = datetime.now(timezone.utc)
+        start_time = utc_now()
         run_name = string_utils.format_name_template(
             name_template=self._deployment.run_name_template,
             substitutions=self._deployment.pipeline_configuration._get_full_substitutions(

zenml/orchestrators/step_run_utils.py

@@ -13,7 +13,6 @@
 #  permissions and limitations under the License.
 """Utilities for creating step runs."""
 
-from datetime import datetime, timezone
 from typing import Dict, List, Optional, Set, Tuple
 
 from zenml.client import Client
@@ -31,6 +30,7 @@ from zenml.models import (
 )
 from zenml.orchestrators import cache_utils, input_utils, utils
 from zenml.stack import Stack
+from zenml.utils.time_utils import utc_now
 
 logger = get_logger(__name__)
 
@@ -75,7 +75,7 @@ class StepRunRequestFactory:
             pipeline_run_id=self.pipeline_run.id,
             deployment=self.deployment.id,
             status=ExecutionStatus.RUNNING,
-            start_time=datetime.now(timezone.utc),
+            start_time=utc_now(),
             user=Client().active_user.id,
             workspace=Client().active_workspace.id,
         )

zenml/pipelines/run_utils.py

@@ -1,7 +1,6 @@
 """Utility functions for running pipelines."""
 
 import time
-from datetime import datetime, timezone
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Set, Union
 from uuid import UUID
 
@@ -25,6 +24,7 @@ from zenml.models import (
 from zenml.orchestrators.publish_utils import publish_failed_pipeline_run
 from zenml.stack import Flavor, Stack
 from zenml.utils import code_utils, notebook_utils, source_utils, string_utils
+from zenml.utils.time_utils import utc_now
 from zenml.zen_stores.base_zen_store import BaseZenStore
 
 if TYPE_CHECKING:
@@ -65,7 +65,7 @@ def create_placeholder_run(
 
     if deployment.schedule:
         return None
-    start_time = datetime.now(timezone.utc)
+    start_time = utc_now()
     run_request = PipelineRunRequest(
         name=string_utils.format_name_template(
             name_template=deployment.run_name_template,

zenml/service_connectors/service_connector.py

@@ -55,6 +55,7 @@ from zenml.models import (
     UserResponse,
     WorkspaceResponse,
 )
+from zenml.utils.time_utils import utc_now
 
 logger = get_logger(__name__)
 
@@ -794,13 +795,14 @@ class ServiceConnector(BaseModel, metaclass=ServiceConnectorMeta):
                 "connector configuration is not valid: name and ID must be set"
             )
 
+        now = utc_now()
         model = ServiceConnectorResponse(
             id=id,
             name=name,
             body=ServiceConnectorResponseBody(
                 user=user,
-                created=datetime.now(timezone.utc),
-                updated=datetime.now(timezone.utc),
+                created=now,
+                updated=now,
                 description=description,
                 connector_type=self.get_type(),
                 auth_method=self.auth_method,
@@ -845,14 +847,15 @@ class ServiceConnector(BaseModel, metaclass=ServiceConnectorMeta):
             if self.expires_skew_tolerance is not None
             else SERVICE_CONNECTOR_SKEW_TOLERANCE_SECONDS
         )
-        delta = expires_at - datetime.now(timezone.utc)
+        now = utc_now(tz_aware=expires_at)
+        delta = expires_at - now
         result = delta < timedelta(seconds=0)
 
         logger.debug(
             f"Checking if connector {self.name} has expired.\n"
             f"Expires at: {self.expires_at}\n"
             f"Expires at (+skew): {expires_at}\n"
-            f"Current UTC time: {datetime.now(timezone.utc)}\n"
+            f"Current UTC time: {now}\n"
             f"Delta: {delta}\n"
             f"Result: {result}\n"
         )

zenml/stack/stack.py

@@ -16,7 +16,7 @@
 import itertools
 import json
 import os
-from datetime import datetime, timezone
+from datetime import datetime
 from typing import (
     TYPE_CHECKING,
     AbstractSet,
@@ -45,6 +45,7 @@ from zenml.logger import get_logger
 from zenml.metadata.metadata_types import MetadataType
 from zenml.models import StackResponse
 from zenml.utils import pagination_utils, settings_utils
+from zenml.utils.time_utils import utc_now
 
 if TYPE_CHECKING:
     from zenml.alerter import BaseAlerter
@@ -732,7 +733,6 @@ class Stack:
             and not skip_default_image_builder
             and not self.image_builder
         ):
-            from datetime import datetime
             from uuid import uuid4
 
             from zenml.image_builders import (
@@ -743,6 +743,7 @@ class Stack:
 
             flavor = LocalImageBuilderFlavor()
 
+            now = utc_now()
             image_builder = LocalImageBuilder(
                 id=uuid4(),
                 name="temporary_default",
@@ -751,8 +752,8 @@ class Stack:
                 config=LocalImageBuilderConfig(),
                 user=Client().active_user.id,
                 workspace=Client().active_workspace.id,
-                created=datetime.now(timezone.utc),
-                updated=datetime.now(timezone.utc),
+                created=now,
+                updated=now,
             )
 
             self._image_builder = image_builder

zenml/stack/stack_component.py

@@ -29,7 +29,11 @@ from zenml.config.step_run_info import StepRunInfo
 from zenml.enums import StackComponentType
 from zenml.exceptions import AuthorizationException
 from zenml.logger import get_logger
-from zenml.models import ServiceConnectorRequirements, StepRunResponse
+from zenml.models import (
+    PipelineRunResponse,
+    ServiceConnectorRequirements,
+    StepRunResponse,
+)
 from zenml.utils import (
     pydantic_utils,
     secret_utils,
@@ -496,6 +500,7 @@ class StackComponent:
             "StepRunInfo",
             "PipelineDeploymentBase",
             "PipelineDeploymentResponse",
+            "PipelineRunResponse",
         ],
     ) -> "BaseSettings":
         """Gets settings for this stack component.
@@ -527,7 +532,10 @@ class StackComponent:
 
         all_settings = (
             container.config.settings
-            if isinstance(container, (Step, StepRunResponse, StepRunInfo))
+            if isinstance(
+                container,
+                (Step, StepRunResponse, StepRunInfo, PipelineRunResponse),
+            )
             else container.pipeline_configuration.settings
         )
 

zenml/stack_deployments/stack_deployment.py

@@ -26,6 +26,7 @@ from zenml.models import (
     StackDeploymentConfig,
     StackDeploymentInfo,
 )
+from zenml.utils.time_utils import to_utc_timezone
 
 STACK_DEPLOYMENT_TERRAFORM = "terraform"
 
@@ -199,9 +200,7 @@ class ZenMLCloudStackDeployment(BaseModel):
         # Get all stacks created after the start date
 
         if date_start and date_start.tzinfo:
-            date_start = date_start.astimezone(datetime.timezone.utc).replace(
-                tzinfo=None
-            )
+            date_start = to_utc_timezone(date_start).replace(tzinfo=None)
         stacks = client.list_stacks(
             created=f"gt:{str(date_start.replace(microsecond=0))}"
             if date_start

zenml/utils/string_utils.py

@@ -17,12 +17,12 @@ import base64
 import functools
 import random
 import string
-from datetime import datetime, timezone
 from typing import Any, Callable, Dict, Optional, TypeVar, cast
 
 from pydantic import BaseModel
 
 from zenml.constants import BANNED_NAME_CHARACTERS
+from zenml.utils.time_utils import utc_now
 
 V = TypeVar("V", bound=Any)
 
@@ -180,7 +180,7 @@ def format_name_template(
             start_time = None
 
         if start_time is None:
-            start_time = datetime.now(timezone.utc)
+            start_time = utc_now()
         substitutions.setdefault("date", start_time.strftime("%Y_%m_%d"))
         substitutions.setdefault("time", start_time.strftime("%H_%M_%S_%f"))