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

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
 
 

zenml/zen_stores/schemas/component_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
 
@@ -30,6 +29,7 @@ from zenml.models import (
     ComponentResponseResources,
     ComponentUpdate,
 )
+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.service_connector_schemas import (
@@ -171,7 +171,7 @@ class StackComponentSchema(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/device_schemas.py

@@ -13,7 +13,7 @@
 #  permissions and limitations under the License.
 """SQLModel implementation for authorized OAuth2 devices."""
 
-from datetime import datetime, timedelta, timezone
+from datetime import datetime, timedelta
 from secrets import token_hex
 from typing import Any, Optional, Tuple
 from uuid import UUID
@@ -31,6 +31,7 @@ from zenml.models import (
     OAuthDeviceResponseMetadata,
     OAuthDeviceUpdate,
 )
+from zenml.utils.time_utils import utc_now
 from zenml.zen_stores.schemas.base_schemas import BaseSchema
 from zenml.zen_stores.schemas.schema_utils import build_foreign_key_field
 from zenml.zen_stores.schemas.user_schemas import UserSchema
@@ -115,7 +116,7 @@ class OAuthDeviceSchema(BaseSchema, table=True):
         device_code = cls._generate_device_code()
         hashed_user_code = cls._get_hashed_code(user_code)
         hashed_device_code = cls._get_hashed_code(device_code)
-        now = datetime.now(timezone.utc)
+        now = utc_now()
         return (
             cls(
                 client_id=request.client_id,
@@ -159,7 +160,7 @@ class OAuthDeviceSchema(BaseSchema, table=True):
         elif device_update.locked is False:
             self.status = OAuthDeviceStatus.ACTIVE.value
 
-        self.updated = datetime.now(timezone.utc)
+        self.updated = utc_now()
         return self
 
     def internal_update(
@@ -174,7 +175,7 @@ class OAuthDeviceSchema(BaseSchema, table=True):
             The updated `OAuthDeviceSchema` and the new user code and device
             code, if they were generated.
         """
-        now = datetime.now(timezone.utc)
+        now = utc_now()
         user_code: Optional[str] = None
         device_code: Optional[str] = None
 

zenml/zen_stores/schemas/event_source_schemas.py

@@ -15,7 +15,6 @@
 
 import base64
 import json
-from datetime import datetime, timezone
 from typing import TYPE_CHECKING, Any, List, Optional, cast
 from uuid import UUID
 
@@ -32,6 +31,7 @@ from zenml.models import (
     Page,
 )
 from zenml.utils.json_utils import pydantic_encoder
+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
@@ -184,5 +184,5 @@ class EventSourceSchema(NamedSchema, table=True):
                 )
             else:
                 setattr(self, field, value)
-        self.updated = datetime.now(timezone.utc)
+        self.updated = utc_now()
         return self

zenml/zen_stores/schemas/flavor_schemas.py

@@ -14,7 +14,6 @@
 """SQL Model Implementations for Flavors."""
 
 import json
-from datetime import datetime, timezone
 from typing import Any, Optional
 from uuid import UUID
 
@@ -28,6 +27,7 @@ from zenml.models import (
     FlavorResponseMetadata,
     FlavorUpdate,
 )
+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
@@ -103,7 +103,7 @@ class FlavorSchema(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/model_schemas.py

@@ -13,7 +13,6 @@
 #  permissions and limitations under the License.
 """SQLModel implementation of model tables."""
 
-from datetime import datetime, timezone
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, cast
 from uuid import UUID, uuid4
 
@@ -52,6 +51,7 @@ from zenml.models import (
     ModelVersionResponseResources,
     Page,
 )
+from zenml.utils.time_utils import utc_now
 from zenml.zen_stores.schemas.artifact_schemas import ArtifactVersionSchema
 from zenml.zen_stores.schemas.base_schemas import BaseSchema, NamedSchema
 from zenml.zen_stores.schemas.constants import MODEL_VERSION_TABLENAME
@@ -224,7 +224,7 @@ class ModelSchema(NamedSchema, table=True):
             exclude_unset=True, exclude_none=True
         ).items():
             setattr(self, field, value)
-        self.updated = datetime.now(timezone.utc)
+        self.updated = utc_now()
         return self
 
 
@@ -504,7 +504,7 @@ class ModelVersionSchema(NamedSchema, RunMetadataInterface, table=True):
             self.name = target_name
         if target_description is not None:
             self.description = target_description
-        self.updated = datetime.now(timezone.utc)
+        self.updated = utc_now()
         return self
 
 

zenml/zen_stores/schemas/pipeline_run_schemas.py

@@ -14,7 +14,7 @@
 """SQLModel implementation of pipeline run tables."""
 
 import json
-from datetime import datetime, timezone
+from datetime import datetime
 from typing import TYPE_CHECKING, Any, Dict, List, Optional
 from uuid import UUID
 
@@ -37,6 +37,7 @@ from zenml.models import (
     RunMetadataEntry,
 )
 from zenml.models.v2.core.pipeline_run import PipelineRunResponseResources
+from zenml.utils.time_utils import utc_now
 from zenml.zen_stores.schemas.base_schemas import NamedSchema
 from zenml.zen_stores.schemas.constants import MODEL_VERSION_TABLENAME
 from zenml.zen_stores.schemas.pipeline_build_schemas import PipelineBuildSchema
@@ -437,7 +438,7 @@ class PipelineRunSchema(NamedSchema, RunMetadataInterface, table=True):
         if run_update.model_version_id and self.model_version_id is None:
             self.model_version_id = run_update.model_version_id
 
-        self.updated = datetime.now(timezone.utc)
+        self.updated = utc_now()
         return self
 
     def update_placeholder(
@@ -478,7 +479,7 @@ class PipelineRunSchema(NamedSchema, RunMetadataInterface, table=True):
         self.orchestrator_environment = orchestrator_environment
         self.status = request.status.value
 
-        self.updated = datetime.now(timezone.utc)
+        self.updated = utc_now()
 
         return self
 

zenml/zen_stores/schemas/pipeline_schemas.py

@@ -13,7 +13,6 @@
 #  permissions and limitations under the License.
 """SQL Model Implementations for Pipelines and Pipeline Runs."""
 
-from datetime import datetime, timezone
 from typing import TYPE_CHECKING, Any, List, Optional
 from uuid import UUID
 
@@ -29,6 +28,7 @@ from zenml.models import (
     PipelineResponseResources,
     PipelineUpdate,
 )
+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
@@ -185,5 +185,5 @@ class PipelineSchema(NamedSchema, table=True):
             The updated `PipelineSchema`.
         """
         self.description = pipeline_update.description
-        self.updated = datetime.now(timezone.utc)
+        self.updated = utc_now()
         return self

zenml/zen_stores/schemas/run_template_schemas.py

@@ -13,7 +13,6 @@
 #  permissions and limitations under the License.
 """SQLModel implementation of run template tables."""
 
-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 (
     RunTemplateResponseResources,
     RunTemplateUpdate,
 )
+from zenml.utils.time_utils import utc_now
 from zenml.zen_stores.schemas.base_schemas import BaseSchema
 from zenml.zen_stores.schemas.schema_utils import build_foreign_key_field
 from zenml.zen_stores.schemas.user_schemas import UserSchema
@@ -156,7 +156,7 @@ class RunTemplateSchema(BaseSchema, table=True):
         ).items():
             setattr(self, field, value)
 
-        self.updated = datetime.now(timezone.utc)
+        self.updated = utc_now()
         return self
 
     def to_model(

zenml/zen_stores/schemas/schedule_schema.py

@@ -13,7 +13,7 @@
 #  permissions and limitations under the License.
 """SQL Model Implementations for Pipeline Schedules."""
 
-from datetime import datetime, timedelta, timezone
+from datetime import datetime, timedelta
 from typing import TYPE_CHECKING, Any, List, Optional
 from uuid import UUID
 
@@ -27,6 +27,7 @@ from zenml.models import (
     ScheduleResponseMetadata,
     ScheduleUpdate,
 )
+from zenml.utils.time_utils import utc_now
 from zenml.zen_stores.schemas.base_schemas import NamedSchema
 from zenml.zen_stores.schemas.component_schemas import StackComponentSchema
 from zenml.zen_stores.schemas.pipeline_schemas import PipelineSchema
@@ -167,7 +168,7 @@ class ScheduleSchema(NamedSchema, RunMetadataInterface, table=True):
             )
         if schedule_update.catchup is not None:
             self.catchup = schedule_update.catchup
-        self.updated = datetime.now(timezone.utc)
+        self.updated = utc_now()
         return self
 
     def to_model(