zenml-nightly 0.73.0.dev20250124__py3-none-any.whl → 0.73.0.dev20250126__py3-none-any.whl

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

@@ -31,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__)
 
@@ -182,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]:
@@ -194,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

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_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"))
 

zenml/utils/time_utils.py

@@ -0,0 +1,138 @@
+#  Copyright (c) ZenML GmbH 2025. 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.
+"""Time utils."""
+
+from datetime import datetime, timedelta, timezone
+from typing import Optional, Union
+
+
+def utc_now(tz_aware: Union[bool, datetime] = False) -> datetime:
+    """Get the current time in the UTC timezone.
+
+    Args:
+        tz_aware: Use this flag to control whether the returned datetime is
+            timezone-aware or timezone-naive. If a datetime is provided, the
+            returned datetime will be timezone-aware if and only if the input
+            datetime is also timezone-aware.
+
+    Returns:
+        The current UTC time. If tz_aware is a datetime, the returned datetime
+        will be timezone-aware only if the input datetime is also timezone-aware.
+        If tz_aware is a boolean, the returned datetime will be timezone-aware
+        if True, and timezone-naive if False.
+    """
+    now = datetime.now(timezone.utc)
+    if (
+        isinstance(tz_aware, bool)
+        and tz_aware is False
+        or isinstance(tz_aware, datetime)
+        and tz_aware.tzinfo is None
+    ):
+        return now.replace(tzinfo=None)
+
+    return now
+
+
+def utc_now_tz_aware() -> datetime:
+    """Get the current timezone-aware UTC time.
+
+    Returns:
+        The current UTC time.
+    """
+    return utc_now(tz_aware=True)
+
+
+def to_local_tz(dt: datetime) -> datetime:
+    """Convert a datetime to the local timezone.
+
+    If the input datetime is timezone-naive, it will be assumed to be in the UTC
+    timezone.
+
+    Args:
+        dt: datetime to convert.
+
+    Returns:
+        Datetime in the local timezone.
+    """
+    if dt.tzinfo is None:
+        dt = dt.replace(tzinfo=timezone.utc)
+    return dt.astimezone()
+
+
+def to_utc_timezone(dt: datetime) -> datetime:
+    """Convert a datetime to the UTC timezone.
+
+    If the input datetime is timezone-naive, it will be assumed to be in the UTC
+    timezone.
+
+    Args:
+        dt: datetime to convert.
+
+    Returns:
+        Datetime in the UTC timezone.
+    """
+    if dt.tzinfo is None:
+        dt = dt.replace(tzinfo=timezone.utc)
+    return dt.astimezone(timezone.utc)
+
+
+def seconds_to_human_readable(time_seconds: int) -> str:
+    """Converts seconds to human-readable format.
+
+    Args:
+        time_seconds: Seconds to convert.
+
+    Returns:
+        Human readable string.
+    """
+    seconds = time_seconds % 60
+    minutes = (time_seconds // 60) % 60
+    hours = (time_seconds // 3600) % 24
+    days = time_seconds // 86400
+    tokens = []
+    if days:
+        tokens.append(f"{days}d")
+    if hours:
+        tokens.append(f"{hours}h")
+    if minutes:
+        tokens.append(f"{minutes}m")
+    if seconds:
+        tokens.append(f"{seconds}s")
+
+    return "".join(tokens)
+
+
+def expires_in(
+    expires_at: datetime,
+    expired_str: str,
+    skew_tolerance: Optional[int] = None,
+) -> str:
+    """Returns a human-readable string of the time until an expiration.
+
+    Args:
+        expires_at: Expiration time.
+        expired_str: String to return if the expiration is in the past.
+        skew_tolerance: Seconds of skew tolerance to subtract from the
+            expiration time. If the expiration is within the skew tolerance,
+            the function will return the expired string.
+
+    Returns:
+        Human readable string.
+    """
+    now = utc_now(tz_aware=expires_at)
+    if skew_tolerance:
+        expires_at -= timedelta(seconds=skew_tolerance)
+    if expires_at < now:
+        return expired_str
+    return seconds_to_human_readable(int((expires_at - now).total_seconds()))

zenml/zen_server/auth.py

@@ -14,7 +14,7 @@
 """Authentication module for ZenML server."""
 
 from contextvars import ContextVar
-from datetime import datetime, timedelta, timezone
+from datetime import datetime, timedelta
 from typing import Callable, Optional, Union
 from urllib.parse import urlencode, urlparse
 from uuid import UUID, uuid4
@@ -62,6 +62,7 @@ from zenml.models import (
     UserResponse,
     UserUpdate,
 )
+from zenml.utils.time_utils import utc_now
 from zenml.zen_server.cache import cache_result
 from zenml.zen_server.csrf import CSRFToken
 from zenml.zen_server.exceptions import http_exception_from_error
@@ -350,7 +351,8 @@ def authenticate_credentials(
 
             if (
                 device_model.expires
-                and datetime.now(timezone.utc) >= device_model.expires
+                and utc_now(tz_aware=device_model.expires)
+                >= device_model.expires
             ):
                 error = (
                     f"Authentication error: device {decoded_token.device_id} "
@@ -589,7 +591,7 @@ def authenticate_device(client_id: UUID, device_code: str) -> AuthContext:
 
     if (
         device_model.expires
-        and datetime.now(timezone.utc) >= device_model.expires
+        and utc_now(tz_aware=device_model.expires) >= device_model.expires
     ):
         error = (
             f"Authentication error: device for client ID {client_id} has "
@@ -892,21 +894,18 @@ def generate_access_token(
     if expires_in == 0:
         expires_in = None
     elif expires_in is not None:
-        expires = datetime.now(timezone.utc) + timedelta(seconds=expires_in)
+        expires = utc_now() + timedelta(seconds=expires_in)
     elif device:
         # If a device was used for authentication, the token will expire
         # at the same time as the device.
         expires = device.expires
         if expires:
             expires_in = max(
-                int(
-                    expires.timestamp()
-                    - datetime.now(timezone.utc).timestamp()
-                ),
+                int(expires.timestamp() - utc_now().timestamp()),
                 0,
             )
     elif config.jwt_token_expire_minutes:
-        expires = datetime.now(timezone.utc) + timedelta(
+        expires = utc_now() + timedelta(
             minutes=config.jwt_token_expire_minutes
         )
         expires_in = config.jwt_token_expire_minutes * 60

zenml/zen_server/cloud_utils.py

@@ -1,6 +1,6 @@
 """Utils concerning anything concerning the cloud control plane backend."""
 
-from datetime import datetime, timedelta, timezone
+from datetime import datetime, timedelta
 from typing import Any, Dict, Optional
 
 import requests
@@ -8,6 +8,7 @@ from requests.adapters import HTTPAdapter, Retry
 
 from zenml.config.server_config import ServerProConfiguration
 from zenml.exceptions import SubscriptionUpgradeRequiredError
+from zenml.utils.time_utils import utc_now
 from zenml.zen_server.utils import get_zenml_headers, server_config
 
 _cloud_connection: Optional["ZenMLCloudConnection"] = None
@@ -185,8 +186,7 @@ class ZenMLCloudConnection:
         if (
             self._token is not None
             and self._token_expires_at is not None
-            and datetime.now(timezone.utc) + timedelta(minutes=5)
-            < self._token_expires_at
+            and utc_now() + timedelta(minutes=5) < self._token_expires_at
         ):
             return self._token
 
@@ -227,9 +227,7 @@ class ZenMLCloudConnection:
             )
 
         self._token = access_token
-        self._token_expires_at = datetime.now(timezone.utc) + timedelta(
-            seconds=expires_in
-        )
+        self._token_expires_at = utc_now() + timedelta(seconds=expires_in)
 
         assert self._token is not None
         return self._token

zenml/zen_server/routers/devices_endpoints.py

@@ -13,7 +13,6 @@
 #  permissions and limitations under the License.
 """Endpoint definitions for code repositories."""
 
-from datetime import datetime, timezone
 from typing import Optional
 from uuid import UUID
 
@@ -36,6 +35,7 @@ from zenml.models import (
     OAuthDeviceVerificationRequest,
     Page,
 )
+from zenml.utils.time_utils import utc_now
 from zenml.zen_server.auth import AuthContext, authorize
 from zenml.zen_server.exceptions import error_response
 from zenml.zen_server.utils import (
@@ -219,9 +219,7 @@ def verify_authorized_device(
             )
 
         # Check if the device verification has expired.
-        if device_model.expires and device_model.expires < datetime.now(
-            timezone.utc
-        ):
+        if device_model.expires and device_model.expires < utc_now():
             raise ValueError(
                 "Invalid request: device verification expired.",
             )