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

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.",
             )

zenml/zen_server/routers/workspaces_endpoints.py

@@ -1009,6 +1009,8 @@ def create_run_metadata(
             verify_models.append(zen_store().get_artifact_version(resource.id))
         elif resource.type == MetadataResourceTypes.MODEL_VERSION:
             verify_models.append(zen_store().get_model_version(resource.id))
+        elif resource.type == MetadataResourceTypes.SCHEDULE:
+            verify_models.append(zen_store().get_schedule(resource.id))
         else:
             raise RuntimeError(f"Unknown resource type: {resource.type}")
 

zenml/zen_server/zen_server_api.py

@@ -22,7 +22,7 @@ To run this file locally, execute:
 
 import os
 from asyncio.log import logger
-from datetime import datetime, timedelta, timezone
+from datetime import datetime, timedelta
 from genericpath import isfile
 from typing import Any, List, Set
 
@@ -54,6 +54,7 @@ from zenml.constants import (
 )
 from zenml.enums import AuthScheme, SourceContextTypes
 from zenml.models import ServerDeploymentType
+from zenml.utils.time_utils import utc_now
 from zenml.zen_server.cloud_utils import send_pro_tenant_status_update
 from zenml.zen_server.exceptions import error_detail
 from zenml.zen_server.routers import (
@@ -129,8 +130,8 @@ app = FastAPI(
 )
 
 # Initialize last_user_activity
-last_user_activity: datetime = datetime.now(timezone.utc)
-last_user_activity_reported: datetime = datetime.now(timezone.utc) + timedelta(
+last_user_activity: datetime = utc_now()
+last_user_activity_reported: datetime = last_user_activity + timedelta(
     seconds=-DEFAULT_ZENML_SERVER_REPORT_USER_ACTIVITY_TO_DB_SECONDS
 )
 
@@ -306,20 +307,20 @@ async def track_last_user_activity(request: Request, call_next: Any) -> Any:
     global last_user_activity
     global last_user_activity_reported
 
+    now = utc_now()
+
     try:
         if is_user_request(request):
-            last_user_activity = datetime.now(timezone.utc)
+            last_user_activity = now
     except Exception as e:
         logger.debug(
             f"An unexpected error occurred while checking user activity: {e}"
         )
     if (
-        (
-            datetime.now(timezone.utc) - last_user_activity_reported
-        ).total_seconds()
+        (now - last_user_activity_reported).total_seconds()
         > DEFAULT_ZENML_SERVER_REPORT_USER_ACTIVITY_TO_DB_SECONDS
     ):
-        last_user_activity_reported = datetime.now(timezone.utc)
+        last_user_activity_reported = now
         zen_store()._update_last_user_activity_timestamp(
             last_user_activity=last_user_activity
         )

zenml/zen_stores/migrations/versions/25155145c545_separate_actions_and_triggers.py

@@ -6,13 +6,14 @@ Create Date: 2024-05-16 11:29:53.341275
 
 """
 
-from datetime import datetime, timezone
 from uuid import uuid4
 
 import sqlalchemy as sa
 import sqlmodel
 from alembic import op
 
+from zenml.utils.time_utils import utc_now
+
 # revision identifiers, used by Alembic.
 revision = "25155145c545"
 down_revision = "0.58.2"
@@ -42,7 +43,7 @@ def migrate_actions() -> None:
         )
     ).fetchall()
 
-    now = datetime.now(timezone.utc)
+    now = utc_now()
 
     actions_to_insert = []
     trigger_updates = {}

zenml/zen_stores/migrations/versions/3dcc5d20e82f_add_last_user_activity.py

@@ -6,12 +6,12 @@ Create Date: 2024-08-07 14:49:07.623500
 
 """
 
-from datetime import datetime, timezone
-
 import sqlalchemy as sa
 import sqlmodel
 from alembic import op
 
+from zenml.utils.time_utils import utc_now
+
 # revision identifiers, used by Alembic.
 revision = "3dcc5d20e82f"
 down_revision = "026d4577b6a0"
@@ -36,7 +36,7 @@ def upgrade() -> None:
             SET last_user_activity = :last_user_activity
             """
         ),
-        params=(dict(last_user_activity=datetime.now(timezone.utc))),
+        params=(dict(last_user_activity=utc_now())),
     )
 
     with op.batch_alter_table("server_settings", schema=None) as batch_op:

zenml/zen_stores/migrations/versions/46506f72f0ed_add_server_settings.py

@@ -6,13 +6,14 @@ Create Date: 2024-04-17 14:17:08.142652
 
 """
 
-from datetime import datetime, timezone
 from uuid import UUID
 
 import sqlalchemy as sa
 import sqlmodel
 from alembic import op
 
+from zenml.utils.time_utils import utc_now
+
 # revision identifiers, used by Alembic.
 revision = "46506f72f0ed"
 down_revision = "cc9894cb58aa"
@@ -76,7 +77,7 @@ def upgrade() -> None:
                 "display_announcements": True,
                 "display_updates": True,
                 # Set the updated timestamp to the current time
-                "updated": datetime.now(timezone.utc),
+                "updated": utc_now(),
             },
         ],
     )

zenml/zen_stores/migrations/versions/5994f9ad0489_introduce_role_permissions.py

@@ -6,7 +6,6 @@ Create Date: 2022-10-25 23:52:25.935344
 
 """
 
-import datetime
 import uuid
 
 import sqlalchemy as sa
@@ -14,6 +13,8 @@ import sqlmodel
 from alembic import op
 from sqlalchemy import or_, select
 
+from zenml.utils.time_utils import utc_now
+
 # revision identifiers, used by Alembic.
 revision = "5994f9ad0489"
 down_revision = "0.21.1"
@@ -100,6 +101,7 @@ def upgrade() -> None:
     guest_id = str(uuid.uuid4()).replace("-", "")
 
     # Prefill the roles table with the admin and guest role
+    now = utc_now()
     op.bulk_insert(
         sa.Table(
             "roleschema",
@@ -109,14 +111,14 @@ def upgrade() -> None:
             {
                 "id": admin_id,
                 "name": "admin",
-                "created": datetime.datetime.now(datetime.timezone.utc),
-                "updated": datetime.datetime.now(datetime.timezone.utc),
+                "created": now,
+                "updated": now,
             },
             {
                 "id": guest_id,
                 "name": "guest",
-                "created": datetime.datetime.now(datetime.timezone.utc),
-                "updated": datetime.datetime.now(datetime.timezone.utc),
+                "created": now,
+                "updated": now,
             },
         ],
     )
@@ -148,6 +150,7 @@ def upgrade() -> None:
     res = conn.execute(select(userschema.c.id)).fetchall()
     user_ids = [i[0] for i in res]
 
+    now = utc_now()
     for user_id in user_ids:
         op.bulk_insert(
             sa.Table(
@@ -159,8 +162,8 @@ def upgrade() -> None:
                     "id": str(uuid.uuid4()).replace("-", ""),
                     "role_id": admin_id,
                     "user_id": user_id,
-                    "created": datetime.datetime.now(datetime.timezone.utc),
-                    "updated": datetime.datetime.now(datetime.timezone.utc),
+                    "created": now,
+                    "updated": now,
                 }
             ],
         )

zenml/zen_stores/migrations/versions/7500f434b71c_remove_shared_columns.py

@@ -8,13 +8,14 @@ Create Date: 2023-10-16 15:15:34.865337
 
 import base64
 from collections import defaultdict
-from datetime import datetime, timezone
 from typing import Dict, Optional, Set
 from uuid import uuid4
 
 import sqlalchemy as sa
 from alembic import op
 
+from zenml.utils.time_utils import utc_now
+
 # revision identifiers, used by Alembic.
 revision = "7500f434b71c"
 down_revision = "14d687c8fa1c"
@@ -123,7 +124,7 @@ def resolve_duplicate_names() -> None:
     _rename_duplicate_entities(service_connector_table)
 
     workspace_query = sa.select(workspace_table.c.id)
-    utcnow = datetime.now(timezone.utc)
+    utcnow = utc_now()
 
     stack_components = []
     stacks = []

zenml/zen_stores/migrations/versions/a91762e6be36_artifact_version_table.py

@@ -6,13 +6,14 @@ Create Date: 2023-11-25 11:01:09.217299
 
 """
 
-from datetime import datetime, timezone
 from uuid import uuid4
 
 import sqlalchemy as sa
 import sqlmodel
 from alembic import op
 
+from zenml.utils.time_utils import utc_now
+
 # revision identifiers, used by Alembic.
 revision = "a91762e6be36"
 down_revision = "0.50.0"
@@ -58,12 +59,13 @@ def upgrade() -> None:
     unique_artifact_names = {
         name: has_custom_name for name, has_custom_name in artifact_names
     }
+    now = utc_now()
     for name, has_custom_name in unique_artifact_names.items():
         conn.execute(
             artifacts.insert().values(
                 id=uuid4().hex,
-                created=datetime.now(timezone.utc),
-                updated=datetime.now(timezone.utc),
+                created=now,
+                updated=now,
                 name=name,
                 has_custom_name=has_custom_name,
             )

zenml/zen_stores/schemas/action_schemas.py

@@ -15,7 +15,6 @@
 
 import base64
 import json
-from datetime import datetime, timezone
 from typing import TYPE_CHECKING, Any, List, Optional
 from uuid import UUID
 
@@ -31,6 +30,7 @@ from zenml.models import (
     ActionResponseResources,
     ActionUpdate,
 )
+from zenml.utils.time_utils import utc_now
 from zenml.zen_stores.schemas.base_schemas import NamedSchema
 from zenml.zen_stores.schemas.schema_utils import build_foreign_key_field
 from zenml.zen_stores.schemas.user_schemas import UserSchema
@@ -140,7 +140,7 @@ class ActionSchema(NamedSchema, table=True):
             else:
                 setattr(self, field, value)
 
-        self.updated = datetime.now(timezone.utc)
+        self.updated = utc_now()
         return self
 
     def to_model(

zenml/zen_stores/schemas/api_key_schemas.py

@@ -13,7 +13,7 @@
 #  permissions and limitations under the License.
 """SQLModel implementation of user tables."""
 
-from datetime import datetime, timezone
+from datetime import datetime
 from secrets import token_hex
 from typing import Any, Optional, Tuple
 from uuid import UUID
@@ -32,6 +32,7 @@ from zenml.models import (
     APIKeyRotateRequest,
     APIKeyUpdate,
 )
+from zenml.utils.time_utils import utc_now
 from zenml.zen_stores.schemas.base_schemas import NamedSchema
 from zenml.zen_stores.schemas.schema_utils import build_foreign_key_field
 from zenml.zen_stores.schemas.user_schemas import UserSchema
@@ -100,7 +101,7 @@ class APIKeySchema(NamedSchema, table=True):
         """
         key = cls._generate_jwt_secret_key()
         hashed_key = cls._get_hashed_key(key)
-        now = datetime.now(timezone.utc)
+        now = utc_now()
         return (
             cls(
                 name=request.name,
@@ -197,7 +198,7 @@ class APIKeySchema(NamedSchema, table=True):
             if hasattr(self, field):
                 setattr(self, field, value)
 
-        self.updated = datetime.now(timezone.utc)
+        self.updated = utc_now()
         return self
 
     def internal_update(self, update: APIKeyInternalUpdate) -> "APIKeySchema":
@@ -230,7 +231,7 @@ class APIKeySchema(NamedSchema, table=True):
         Returns:
             The updated `APIKeySchema` and the new un-hashed key.
         """
-        self.updated = datetime.now(timezone.utc)
+        self.updated = utc_now()
         self.previous_key = self.key
         self.retain_period = rotate_request.retain_period_minutes
         new_key = self._generate_jwt_secret_key()

zenml/zen_stores/schemas/artifact_schemas.py

@@ -13,7 +13,6 @@
 #  permissions and limitations under the License.
 """SQLModel implementation of artifact table."""
 
-from datetime import datetime, timezone
 from typing import TYPE_CHECKING, Any, List, Optional
 from uuid import UUID
 
@@ -41,6 +40,7 @@ from zenml.models import (
     ArtifactVersionUpdate,
 )
 from zenml.models.v2.core.artifact import ArtifactRequest
+from zenml.utils.time_utils import utc_now
 from zenml.zen_stores.schemas.base_schemas import BaseSchema, NamedSchema
 from zenml.zen_stores.schemas.component_schemas import StackComponentSchema
 from zenml.zen_stores.schemas.schema_utils import build_foreign_key_field
@@ -163,7 +163,7 @@ class ArtifactSchema(NamedSchema, table=True):
         Returns:
             The updated `ArtifactSchema`.
         """
-        self.updated = datetime.now(timezone.utc)
+        self.updated = utc_now()
         if artifact_update.name:
             self.name = artifact_update.name
             self.has_custom_name = True
@@ -401,5 +401,5 @@ class ArtifactVersionSchema(BaseSchema, RunMetadataInterface, table=True):
         Returns:
             The updated `ArtifactVersionSchema`.
         """
-        self.updated = datetime.now(timezone.utc)
+        self.updated = utc_now()
         return self

zenml/zen_stores/schemas/base_schemas.py

@@ -13,12 +13,14 @@
 #  permissions and limitations under the License.
 """Base classes for SQLModel schemas."""
 
-from datetime import datetime, timezone
+from datetime import datetime
 from typing import TYPE_CHECKING, Any, TypeVar
 from uuid import UUID, uuid4
 
 from sqlmodel import Field, SQLModel
 
+from zenml.utils.time_utils import utc_now
+
 if TYPE_CHECKING:
     from zenml.models.v2.base.base import BaseResponse
 
@@ -29,12 +31,8 @@ class BaseSchema(SQLModel):
     """Base SQL Model for ZenML entities."""
 
     id: UUID = Field(default_factory=uuid4, primary_key=True)
-    created: datetime = Field(
-        default_factory=lambda: datetime.now(timezone.utc)
-    )
-    updated: datetime = Field(
-        default_factory=lambda: datetime.now(timezone.utc)
-    )
+    created: datetime = Field(default_factory=utc_now)
+    updated: datetime = Field(default_factory=utc_now)
 
     def to_model(
         self,

zenml/zen_stores/schemas/code_repository_schemas.py

@@ -14,7 +14,6 @@
 """SQL Model Implementations for code repositories."""
 
 import json
-from datetime import datetime, timezone
 from typing import Any, Optional
 from uuid import UUID
 
@@ -32,6 +31,7 @@ from zenml.models import (
     CodeRepositoryResponseMetadata,
     CodeRepositoryUpdate,
 )
+from zenml.utils.time_utils import utc_now
 from zenml.zen_stores.schemas.base_schemas import BaseSchema, NamedSchema
 from zenml.zen_stores.schemas.schema_utils import build_foreign_key_field
 from zenml.zen_stores.schemas.user_schemas import UserSchema
@@ -151,7 +151,7 @@ class CodeRepositorySchema(NamedSchema, table=True):
         if update.logo_url:
             self.logo_url = update.logo_url
 
-        self.updated = datetime.now(timezone.utc)
+        self.updated = utc_now()
         return self