diff --git a/docs/book/how-to/pipeline-development/develop-locally/keep-your-dashboard-server-clean.md b/docs/book/how-to/pipeline-development/develop-locally/keep-your-dashboard-server-clean.md index 956b5ab6b47..1f1f09dff78 100644 --- a/docs/book/how-to/pipeline-development/develop-locally/keep-your-dashboard-server-clean.md +++ b/docs/book/how-to/pipeline-development/develop-locally/keep-your-dashboard-server-clean.md @@ -59,7 +59,7 @@ def delete_recent_pipeline_runs(): zc = Client() # Calculate the timestamp for 24 hours ago - twenty_four_hours_ago = datetime.datetime.utcnow() - datetime.timedelta(hours=24) + twenty_four_hours_ago = datetime.datetime.now(timezone.utc) - datetime.timedelta(hours=24) # Format the timestamp as required by ZenML time_filter = twenty_four_hours_ago.strftime("%Y-%m-%d %H:%M:%S") diff --git a/src/zenml/cli/service_connectors.py b/src/zenml/cli/service_connectors.py index bbfa86a45ec..82d276b9b5b 100644 --- a/src/zenml/cli/service_connectors.py +++ b/src/zenml/cli/service_connectors.py @@ -13,7 +13,7 @@ # permissions and limitations under the License. """Service connector CLI commands.""" -from datetime import datetime +from datetime import datetime, timezone from typing import Any, Dict, List, Optional, Union, cast from uuid import UUID @@ -291,7 +291,9 @@ def prompt_expires_at( while True: default_str = "" if default is not None: - seconds = int((default - datetime.utcnow()).total_seconds()) + seconds = int( + (default - datetime.now(timezone.utc)).total_seconds() + ) default_str = ( f" [{str(default)} i.e. in " f"{seconds_to_human_readable(seconds)}]" @@ -307,14 +309,16 @@ def prompt_expires_at( assert expires_at is not None assert isinstance(expires_at, datetime) - if expires_at < datetime.utcnow(): + if expires_at < datetime.now(timezone.utc): cli_utils.warning( "The expiration time must be in the future. Please enter a " "later date and time." ) continue - seconds = int((expires_at - datetime.utcnow()).total_seconds()) + seconds = int( + (expires_at - datetime.now(timezone.utc)).total_seconds() + ) confirm = click.confirm( f"Credentials will be valid until {str(expires_at)} UTC (i.e. " diff --git a/src/zenml/cli/stack.py b/src/zenml/cli/stack.py index c99a5f346ec..5bd33984426 100644 --- a/src/zenml/cli/stack.py +++ b/src/zenml/cli/stack.py @@ -17,7 +17,7 @@ import re import time import webbrowser -from datetime import datetime +from datetime import datetime, timezone from typing import ( TYPE_CHECKING, Any, @@ -1576,7 +1576,7 @@ def deploy( ): raise click.Abort() - date_start = datetime.utcnow() + date_start = datetime.now(timezone.utc) webbrowser.open(deployment_config.deployment_url) console.print( diff --git a/src/zenml/config/pipeline_configurations.py b/src/zenml/config/pipeline_configurations.py index e3004f821b4..f33c6711819 100644 --- a/src/zenml/config/pipeline_configurations.py +++ b/src/zenml/config/pipeline_configurations.py @@ -13,7 +13,7 @@ # permissions and limitations under the License. """Pipeline configuration classes.""" -from datetime import datetime +from datetime import datetime, timezone from typing import TYPE_CHECKING, Any, Dict, List, Optional from pydantic import SerializeAsAny, field_validator @@ -61,7 +61,7 @@ def _get_full_substitutions( The full substitutions dict including date and time. """ if start_time is None: - start_time = datetime.utcnow() + start_time = datetime.now(timezone.utc) ret = self.substitutions.copy() ret.setdefault("date", start_time.strftime("%Y_%m_%d")) ret.setdefault("time", start_time.strftime("%H_%M_%S_%f")) diff --git a/src/zenml/event_hub/base_event_hub.py b/src/zenml/event_hub/base_event_hub.py index 23af0bd8465..4b93a6da17c 100644 --- a/src/zenml/event_hub/base_event_hub.py +++ b/src/zenml/event_hub/base_event_hub.py @@ -14,7 +14,7 @@ """Base class for event hub implementations.""" from abc import ABC, abstractmethod -from datetime import datetime, timedelta +from datetime import datetime, timedelta, timezone from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Tuple from zenml import EventSourceResponse @@ -134,7 +134,7 @@ def trigger_action( ) expires: Optional[datetime] = None if trigger.action.auth_window: - expires = datetime.utcnow() + timedelta( + expires = datetime.now(timezone.utc) + timedelta( minutes=trigger.action.auth_window ) encoded_token = token.encode(expires=expires) diff --git a/src/zenml/integrations/airflow/orchestrators/airflow_orchestrator.py b/src/zenml/integrations/airflow/orchestrators/airflow_orchestrator.py index d6f9bde5ccb..fcb46d143e3 100644 --- a/src/zenml/integrations/airflow/orchestrators/airflow_orchestrator.py +++ b/src/zenml/integrations/airflow/orchestrators/airflow_orchestrator.py @@ -408,7 +408,8 @@ def _translate_schedule( if schedule: if schedule.cron_expression: start_time = schedule.start_time or ( - datetime.datetime.utcnow() - datetime.timedelta(7) + datetime.datetime.now(datetime.timezone.utc) + - datetime.timedelta(7) ) return { "schedule": schedule.cron_expression, @@ -428,6 +429,7 @@ def _translate_schedule( "schedule": "@once", # set a start time in the past and disable catchup so airflow # runs the dag immediately - "start_date": datetime.datetime.utcnow() - datetime.timedelta(7), + "start_date": datetime.datetime.now(datetime.timezone.utc) + - datetime.timedelta(7), "catchup": False, } diff --git a/src/zenml/integrations/kubernetes/orchestrators/kube_utils.py b/src/zenml/integrations/kubernetes/orchestrators/kube_utils.py index 28d363c22a5..e97142df086 100644 --- a/src/zenml/integrations/kubernetes/orchestrators/kube_utils.py +++ b/src/zenml/integrations/kubernetes/orchestrators/kube_utils.py @@ -248,7 +248,7 @@ def wait_pod( Returns: The pod object which meets the exit condition. """ - start_time = datetime.datetime.utcnow() + start_time = datetime.datetime.now(datetime.timezone.utc) # 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.utcnow() - start_time + elapse_time = datetime.datetime.now(datetime.timezone.utc) - start_time if elapse_time.seconds >= timeout_sec and timeout_sec != 0: raise RuntimeError( f"Waiting for pod `{namespace}:{pod_name}` timed out after " diff --git a/src/zenml/integrations/whylogs/data_validators/whylogs_data_validator.py b/src/zenml/integrations/whylogs/data_validators/whylogs_data_validator.py index 5f039fe6a66..99ed49d0df9 100644 --- a/src/zenml/integrations/whylogs/data_validators/whylogs_data_validator.py +++ b/src/zenml/integrations/whylogs/data_validators/whylogs_data_validator.py @@ -97,7 +97,9 @@ def data_profiling( """ results = why.log(pandas=dataset) profile = results.profile() - dataset_timestamp = dataset_timestamp or datetime.datetime.utcnow() + dataset_timestamp = dataset_timestamp or datetime.datetime.now( + datetime.timezone.utc + ) profile.set_dataset_timestamp(dataset_timestamp=dataset_timestamp) return profile.view() diff --git a/src/zenml/models/v2/core/api_key.py b/src/zenml/models/v2/core/api_key.py index 321707bd4ad..d6855add252 100644 --- a/src/zenml/models/v2/core/api_key.py +++ b/src/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 +from datetime import datetime, timedelta, timezone from typing import TYPE_CHECKING, ClassVar, List, Optional, Type, Union from uuid import UUID @@ -319,7 +319,7 @@ def verify_key( and self.retain_period_minutes > 0 ): # check if the previous key is still valid - if datetime.utcnow() - self.last_rotated < timedelta( + if datetime.now(timezone.utc) - self.last_rotated < timedelta( minutes=self.retain_period_minutes ): key_hash = self.previous_key diff --git a/src/zenml/orchestrators/publish_utils.py b/src/zenml/orchestrators/publish_utils.py index 0d5cea792ae..778736b37ed 100644 --- a/src/zenml/orchestrators/publish_utils.py +++ b/src/zenml/orchestrators/publish_utils.py @@ -13,7 +13,7 @@ # permissions and limitations under the License. """Utilities to publish pipeline and step runs.""" -from datetime import datetime +from datetime import datetime, timezone from typing import TYPE_CHECKING, Dict, List from zenml.client import Client @@ -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.utcnow(), + end_time=datetime.now(timezone.utc), 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.utcnow(), + end_time=datetime.now(timezone.utc), ), ) @@ -87,7 +87,7 @@ def publish_failed_pipeline_run( run_id=pipeline_run_id, run_update=PipelineRunUpdate( status=ExecutionStatus.FAILED, - end_time=datetime.utcnow(), + end_time=datetime.now(timezone.utc), ), ) diff --git a/src/zenml/orchestrators/step_launcher.py b/src/zenml/orchestrators/step_launcher.py index 1141172bf31..cddf12e8baf 100644 --- a/src/zenml/orchestrators/step_launcher.py +++ b/src/zenml/orchestrators/step_launcher.py @@ -16,7 +16,7 @@ import os import time from contextlib import nullcontext -from datetime import datetime +from datetime import datetime, timezone from functools import partial from typing import TYPE_CHECKING, Any, Callable, Dict, Tuple @@ -201,7 +201,7 @@ def launch(self) -> None: f"Failed preparing step `{self._step_name}`." ) step_run_request.status = ExecutionStatus.FAILED - step_run_request.end_time = datetime.utcnow() + step_run_request.end_time = datetime.now(timezone.utc) raise finally: step_run = Client().zen_store.create_run_step( @@ -305,7 +305,7 @@ def _create_or_reuse_run(self) -> Tuple[PipelineRunResponse, bool]: The created or existing pipeline run, and a boolean indicating whether the run was created or reused. """ - start_time = datetime.utcnow() + start_time = datetime.now(timezone.utc) run_name = string_utils.format_name_template( name_template=self._deployment.run_name_template, substitutions=self._deployment.pipeline_configuration._get_full_substitutions( diff --git a/src/zenml/orchestrators/step_run_utils.py b/src/zenml/orchestrators/step_run_utils.py index 6451a4cc0a4..5b7355c46bf 100644 --- a/src/zenml/orchestrators/step_run_utils.py +++ b/src/zenml/orchestrators/step_run_utils.py @@ -13,7 +13,7 @@ # permissions and limitations under the License. """Utilities for creating step runs.""" -from datetime import datetime +from datetime import datetime, timezone from typing import Dict, List, Optional, Set, Tuple from zenml.client import Client @@ -75,7 +75,7 @@ def create_request(self, invocation_id: str) -> StepRunRequest: pipeline_run_id=self.pipeline_run.id, deployment=self.deployment.id, status=ExecutionStatus.RUNNING, - start_time=datetime.utcnow(), + start_time=datetime.now(timezone.utc), user=Client().active_user.id, workspace=Client().active_workspace.id, ) diff --git a/src/zenml/pipelines/run_utils.py b/src/zenml/pipelines/run_utils.py index 0f0a32564c6..0d198f4e40f 100644 --- a/src/zenml/pipelines/run_utils.py +++ b/src/zenml/pipelines/run_utils.py @@ -1,7 +1,7 @@ """Utility functions for running pipelines.""" import time -from datetime import datetime +from datetime import datetime, timezone from typing import TYPE_CHECKING, Any, Dict, List, Optional, Set, Union from uuid import UUID @@ -65,7 +65,7 @@ def create_placeholder_run( if deployment.schedule: return None - start_time = datetime.utcnow() + start_time = datetime.now(timezone.utc) run_request = PipelineRunRequest( name=string_utils.format_name_template( name_template=deployment.run_name_template, diff --git a/src/zenml/service_connectors/service_connector.py b/src/zenml/service_connectors/service_connector.py index b2de590a460..1e190643dfb 100644 --- a/src/zenml/service_connectors/service_connector.py +++ b/src/zenml/service_connectors/service_connector.py @@ -799,8 +799,8 @@ def to_response_model( name=name, body=ServiceConnectorResponseBody( user=user, - created=datetime.utcnow(), - updated=datetime.utcnow(), + created=datetime.now(timezone.utc), + updated=datetime.now(timezone.utc), description=description, connector_type=self.get_type(), auth_method=self.auth_method, diff --git a/src/zenml/stack/stack.py b/src/zenml/stack/stack.py index 06cf6445377..90b164a09f3 100644 --- a/src/zenml/stack/stack.py +++ b/src/zenml/stack/stack.py @@ -16,7 +16,7 @@ import itertools import json import os -from datetime import datetime +from datetime import datetime, timezone from typing import ( TYPE_CHECKING, AbstractSet, @@ -751,8 +751,8 @@ def validate_image_builder(self) -> None: config=LocalImageBuilderConfig(), user=Client().active_user.id, workspace=Client().active_workspace.id, - created=datetime.utcnow(), - updated=datetime.utcnow(), + created=datetime.now(timezone.utc), + updated=datetime.now(timezone.utc), ) self._image_builder = image_builder diff --git a/src/zenml/utils/git_utils.py b/src/zenml/utils/git_utils.py index 2bd35ac160d..7cf4b99703a 100644 --- a/src/zenml/utils/git_utils.py +++ b/src/zenml/utils/git_utils.py @@ -16,7 +16,7 @@ import os from typing import Optional -from git import GitCommandError +from git.exc import GitCommandError from git.repo import Repo diff --git a/src/zenml/utils/string_utils.py b/src/zenml/utils/string_utils.py index 38615662851..8c0838c7a39 100644 --- a/src/zenml/utils/string_utils.py +++ b/src/zenml/utils/string_utils.py @@ -17,7 +17,7 @@ import functools import random import string -from datetime import datetime +from datetime import datetime, timezone from typing import Any, Callable, Dict, Optional, TypeVar, cast from pydantic import BaseModel @@ -180,7 +180,7 @@ def format_name_template( start_time = None if start_time is None: - start_time = datetime.utcnow() + start_time = datetime.now(timezone.utc) substitutions.setdefault("date", start_time.strftime("%Y_%m_%d")) substitutions.setdefault("time", start_time.strftime("%H_%M_%S_%f")) diff --git a/src/zenml/zen_server/auth.py b/src/zenml/zen_server/auth.py index 80290f091cc..913680aa4ce 100644 --- a/src/zenml/zen_server/auth.py +++ b/src/zenml/zen_server/auth.py @@ -14,7 +14,7 @@ """Authentication module for ZenML server.""" from contextvars import ContextVar -from datetime import datetime, timedelta +from datetime import datetime, timedelta, timezone from typing import Callable, Optional, Union from urllib.parse import urlencode from uuid import UUID @@ -315,7 +315,7 @@ def authenticate_credentials( if ( device_model.expires - and datetime.utcnow() >= device_model.expires + and datetime.now(timezone.utc) >= device_model.expires ): error = ( f"Authentication error: device {decoded_token.device_id} " @@ -552,7 +552,10 @@ def authenticate_device(client_id: UUID, device_code: str) -> AuthContext: error_description=error, ) - if device_model.expires and datetime.utcnow() >= device_model.expires: + if ( + device_model.expires + and datetime.now(timezone.utc) >= device_model.expires + ): error = ( f"Authentication error: device for client ID {client_id} has " "expired" @@ -835,17 +838,21 @@ def generate_access_token( if expires_in == 0: expires_in = None elif expires_in is not None: - expires = datetime.utcnow() + timedelta(seconds=expires_in) + expires = datetime.now(timezone.utc) + 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.utcnow().timestamp()), 0 + int( + expires.timestamp() + - datetime.now(timezone.utc).timestamp() + ), + 0, ) elif config.jwt_token_expire_minutes: - expires = datetime.utcnow() + timedelta( + expires = datetime.now(timezone.utc) + timedelta( minutes=config.jwt_token_expire_minutes ) expires_in = config.jwt_token_expire_minutes * 60 diff --git a/src/zenml/zen_server/routers/devices_endpoints.py b/src/zenml/zen_server/routers/devices_endpoints.py index ebd6f44404e..95bc091ce7a 100644 --- a/src/zenml/zen_server/routers/devices_endpoints.py +++ b/src/zenml/zen_server/routers/devices_endpoints.py @@ -13,7 +13,7 @@ # permissions and limitations under the License. """Endpoint definitions for code repositories.""" -from datetime import datetime +from datetime import datetime, timezone from typing import Optional from uuid import UUID @@ -219,7 +219,9 @@ def verify_authorized_device( ) # Check if the device verification has expired. - if device_model.expires and device_model.expires < datetime.utcnow(): + if device_model.expires and device_model.expires < datetime.now( + timezone.utc + ): raise ValueError( "Invalid request: device verification expired.", ) diff --git a/src/zenml/zen_server/routers/stack_deployment_endpoints.py b/src/zenml/zen_server/routers/stack_deployment_endpoints.py index a98a9de54f5..28b2ccfc01b 100644 --- a/src/zenml/zen_server/routers/stack_deployment_endpoints.py +++ b/src/zenml/zen_server/routers/stack_deployment_endpoints.py @@ -114,9 +114,9 @@ def get_stack_deployment_config( assert token is not None # A new API token is generated for the stack deployment - expires = datetime.datetime.utcnow() + datetime.timedelta( - minutes=STACK_DEPLOYMENT_API_TOKEN_EXPIRATION - ) + expires = datetime.datetime.now( + datetime.timezone.utc + ) + datetime.timedelta(minutes=STACK_DEPLOYMENT_API_TOKEN_EXPIRATION) api_token = token.encode(expires=expires) return stack_deployment_class( diff --git a/src/zenml/zen_stores/migrations/versions/25155145c545_separate_actions_and_triggers.py b/src/zenml/zen_stores/migrations/versions/25155145c545_separate_actions_and_triggers.py index a6317cd0fdb..2ecf44dd788 100644 --- a/src/zenml/zen_stores/migrations/versions/25155145c545_separate_actions_and_triggers.py +++ b/src/zenml/zen_stores/migrations/versions/25155145c545_separate_actions_and_triggers.py @@ -6,7 +6,7 @@ """ -from datetime import datetime +from datetime import datetime, timezone from uuid import uuid4 import sqlalchemy as sa @@ -42,7 +42,7 @@ def migrate_actions() -> None: ) ).fetchall() - now = datetime.utcnow() + now = datetime.now(timezone.utc) actions_to_insert = [] trigger_updates = {} diff --git a/src/zenml/zen_stores/migrations/versions/46506f72f0ed_add_server_settings.py b/src/zenml/zen_stores/migrations/versions/46506f72f0ed_add_server_settings.py index b7255ff1de2..32f8578b60b 100644 --- a/src/zenml/zen_stores/migrations/versions/46506f72f0ed_add_server_settings.py +++ b/src/zenml/zen_stores/migrations/versions/46506f72f0ed_add_server_settings.py @@ -6,7 +6,7 @@ """ -from datetime import datetime +from datetime import datetime, timezone from uuid import UUID import sqlalchemy as sa @@ -76,7 +76,7 @@ def upgrade() -> None: "display_announcements": True, "display_updates": True, # Set the updated timestamp to the current time - "updated": datetime.utcnow(), + "updated": datetime.now(timezone.utc), }, ], ) diff --git a/src/zenml/zen_stores/migrations/versions/5994f9ad0489_introduce_role_permissions.py b/src/zenml/zen_stores/migrations/versions/5994f9ad0489_introduce_role_permissions.py index 05d4dbe68d0..fe7d77d263b 100644 --- a/src/zenml/zen_stores/migrations/versions/5994f9ad0489_introduce_role_permissions.py +++ b/src/zenml/zen_stores/migrations/versions/5994f9ad0489_introduce_role_permissions.py @@ -109,14 +109,14 @@ def upgrade() -> None: { "id": admin_id, "name": "admin", - "created": datetime.datetime.utcnow(), - "updated": datetime.datetime.utcnow(), + "created": datetime.datetime.now(datetime.timezone.utc), + "updated": datetime.datetime.now(datetime.timezone.utc), }, { "id": guest_id, "name": "guest", - "created": datetime.datetime.utcnow(), - "updated": datetime.datetime.utcnow(), + "created": datetime.datetime.now(datetime.timezone.utc), + "updated": datetime.datetime.now(datetime.timezone.utc), }, ], ) @@ -159,8 +159,8 @@ def upgrade() -> None: "id": str(uuid.uuid4()).replace("-", ""), "role_id": admin_id, "user_id": user_id, - "created": datetime.datetime.utcnow(), - "updated": datetime.datetime.utcnow(), + "created": datetime.datetime.now(datetime.timezone.utc), + "updated": datetime.datetime.now(datetime.timezone.utc), } ], ) diff --git a/src/zenml/zen_stores/migrations/versions/7500f434b71c_remove_shared_columns.py b/src/zenml/zen_stores/migrations/versions/7500f434b71c_remove_shared_columns.py index d5d58437d77..75fa38551a0 100644 --- a/src/zenml/zen_stores/migrations/versions/7500f434b71c_remove_shared_columns.py +++ b/src/zenml/zen_stores/migrations/versions/7500f434b71c_remove_shared_columns.py @@ -8,7 +8,7 @@ import base64 from collections import defaultdict -from datetime import datetime +from datetime import datetime, timezone from typing import Dict, Optional, Set from uuid import uuid4 @@ -123,7 +123,7 @@ def resolve_duplicate_names() -> None: _rename_duplicate_entities(service_connector_table) workspace_query = sa.select(workspace_table.c.id) - utcnow = datetime.utcnow() + utcnow = datetime.now(timezone.utc) stack_components = [] stacks = [] diff --git a/src/zenml/zen_stores/migrations/versions/a91762e6be36_artifact_version_table.py b/src/zenml/zen_stores/migrations/versions/a91762e6be36_artifact_version_table.py index 06ae7620c93..370e2253f5b 100644 --- a/src/zenml/zen_stores/migrations/versions/a91762e6be36_artifact_version_table.py +++ b/src/zenml/zen_stores/migrations/versions/a91762e6be36_artifact_version_table.py @@ -6,7 +6,7 @@ """ -from datetime import datetime +from datetime import datetime, timezone from uuid import uuid4 import sqlalchemy as sa @@ -62,8 +62,8 @@ def upgrade() -> None: conn.execute( artifacts.insert().values( id=uuid4().hex, - created=datetime.utcnow(), - updated=datetime.utcnow(), + created=datetime.now(timezone.utc), + updated=datetime.now(timezone.utc), name=name, has_custom_name=has_custom_name, ) diff --git a/src/zenml/zen_stores/schemas/action_schemas.py b/src/zenml/zen_stores/schemas/action_schemas.py index 092ae57d1c9..4a57cbb7ec5 100644 --- a/src/zenml/zen_stores/schemas/action_schemas.py +++ b/src/zenml/zen_stores/schemas/action_schemas.py @@ -15,7 +15,7 @@ import base64 import json -from datetime import datetime +from datetime import datetime, timezone from typing import TYPE_CHECKING, Any, List, Optional from uuid import UUID @@ -140,7 +140,7 @@ def update(self, action_update: "ActionUpdate") -> "ActionSchema": else: setattr(self, field, value) - self.updated = datetime.utcnow() + self.updated = datetime.now(timezone.utc) return self def to_model( diff --git a/src/zenml/zen_stores/schemas/api_key_schemas.py b/src/zenml/zen_stores/schemas/api_key_schemas.py index 47f63f6ca2c..a1535bf8869 100644 --- a/src/zenml/zen_stores/schemas/api_key_schemas.py +++ b/src/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 +from datetime import datetime, timezone from secrets import token_hex from typing import Any, Optional, Tuple from uuid import UUID @@ -100,7 +100,7 @@ def from_request( """ key = cls._generate_jwt_secret_key() hashed_key = cls._get_hashed_key(key) - now = datetime.utcnow() + now = datetime.now(timezone.utc) return ( cls( name=request.name, @@ -197,7 +197,7 @@ def update(self, update: APIKeyUpdate) -> "APIKeySchema": if hasattr(self, field): setattr(self, field, value) - self.updated = datetime.utcnow() + self.updated = datetime.now(timezone.utc) return self def internal_update(self, update: APIKeyInternalUpdate) -> "APIKeySchema": @@ -230,7 +230,7 @@ def rotate( Returns: The updated `APIKeySchema` and the new un-hashed key. """ - self.updated = datetime.utcnow() + self.updated = datetime.now(timezone.utc) self.previous_key = self.key self.retain_period = rotate_request.retain_period_minutes new_key = self._generate_jwt_secret_key() diff --git a/src/zenml/zen_stores/schemas/artifact_schemas.py b/src/zenml/zen_stores/schemas/artifact_schemas.py index 02e842a5fb5..df84b35d1bd 100644 --- a/src/zenml/zen_stores/schemas/artifact_schemas.py +++ b/src/zenml/zen_stores/schemas/artifact_schemas.py @@ -13,7 +13,7 @@ # permissions and limitations under the License. """SQLModel implementation of artifact table.""" -from datetime import datetime +from datetime import datetime, timezone from typing import TYPE_CHECKING, Any, List, Optional from uuid import UUID @@ -164,7 +164,7 @@ def update(self, artifact_update: ArtifactUpdate) -> "ArtifactSchema": Returns: The updated `ArtifactSchema`. """ - self.updated = datetime.utcnow() + self.updated = datetime.now(timezone.utc) if artifact_update.name: self.name = artifact_update.name self.has_custom_name = True @@ -401,5 +401,5 @@ def update( Returns: The updated `ArtifactVersionSchema`. """ - self.updated = datetime.utcnow() + self.updated = datetime.now(timezone.utc) return self diff --git a/src/zenml/zen_stores/schemas/base_schemas.py b/src/zenml/zen_stores/schemas/base_schemas.py index cd8d3f3a7bd..e0c1e1f7500 100644 --- a/src/zenml/zen_stores/schemas/base_schemas.py +++ b/src/zenml/zen_stores/schemas/base_schemas.py @@ -13,7 +13,7 @@ # permissions and limitations under the License. """Base classes for SQLModel schemas.""" -from datetime import datetime +from datetime import datetime, timezone from typing import TYPE_CHECKING, Any, TypeVar from uuid import UUID, uuid4 @@ -29,8 +29,12 @@ class BaseSchema(SQLModel): """Base SQL Model for ZenML entities.""" id: UUID = Field(default_factory=uuid4, primary_key=True) - created: datetime = Field(default_factory=datetime.utcnow) - updated: datetime = Field(default_factory=datetime.utcnow) + created: datetime = Field( + default_factory=lambda: datetime.now(timezone.utc) + ) + updated: datetime = Field( + default_factory=lambda: datetime.now(timezone.utc) + ) def to_model( self, diff --git a/src/zenml/zen_stores/schemas/code_repository_schemas.py b/src/zenml/zen_stores/schemas/code_repository_schemas.py index 6bd4511edf2..86a0c5d24e7 100644 --- a/src/zenml/zen_stores/schemas/code_repository_schemas.py +++ b/src/zenml/zen_stores/schemas/code_repository_schemas.py @@ -14,7 +14,7 @@ """SQL Model Implementations for code repositories.""" import json -from datetime import datetime +from datetime import datetime, timezone from typing import Any, Optional from uuid import UUID @@ -151,7 +151,7 @@ def update(self, update: "CodeRepositoryUpdate") -> "CodeRepositorySchema": if update.logo_url: self.logo_url = update.logo_url - self.updated = datetime.utcnow() + self.updated = datetime.now(timezone.utc) return self diff --git a/src/zenml/zen_stores/schemas/component_schemas.py b/src/zenml/zen_stores/schemas/component_schemas.py index eb38cee7e0d..05b9cd4b9be 100644 --- a/src/zenml/zen_stores/schemas/component_schemas.py +++ b/src/zenml/zen_stores/schemas/component_schemas.py @@ -15,7 +15,7 @@ import base64 import json -from datetime import datetime +from datetime import datetime, timezone from typing import TYPE_CHECKING, Any, List, Optional from uuid import UUID @@ -171,7 +171,7 @@ def update( else: setattr(self, field, value) - self.updated = datetime.utcnow() + self.updated = datetime.now(timezone.utc) return self def to_model( diff --git a/src/zenml/zen_stores/schemas/device_schemas.py b/src/zenml/zen_stores/schemas/device_schemas.py index 52f4ebc9c60..68ee10f7488 100644 --- a/src/zenml/zen_stores/schemas/device_schemas.py +++ b/src/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 +from datetime import datetime, timedelta, timezone from secrets import token_hex from typing import Any, Optional, Tuple from uuid import UUID @@ -115,7 +115,7 @@ def from_request( 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.utcnow() + now = datetime.now(timezone.utc) return ( cls( client_id=request.client_id, @@ -159,7 +159,7 @@ def update(self, device_update: OAuthDeviceUpdate) -> "OAuthDeviceSchema": elif device_update.locked is False: self.status = OAuthDeviceStatus.ACTIVE.value - self.updated = datetime.utcnow() + self.updated = datetime.now(timezone.utc) return self def internal_update( @@ -174,7 +174,7 @@ def internal_update( The updated `OAuthDeviceSchema` and the new user code and device code, if they were generated. """ - now = datetime.utcnow() + now = datetime.now(timezone.utc) user_code: Optional[str] = None device_code: Optional[str] = None diff --git a/src/zenml/zen_stores/schemas/event_source_schemas.py b/src/zenml/zen_stores/schemas/event_source_schemas.py index 46a2d08158f..5ee8d0742f2 100644 --- a/src/zenml/zen_stores/schemas/event_source_schemas.py +++ b/src/zenml/zen_stores/schemas/event_source_schemas.py @@ -15,7 +15,7 @@ import base64 import json -from datetime import datetime +from datetime import datetime, timezone from typing import TYPE_CHECKING, Any, List, Optional, cast from uuid import UUID @@ -184,5 +184,5 @@ def update(self, update: EventSourceUpdate) -> "EventSourceSchema": ) else: setattr(self, field, value) - self.updated = datetime.utcnow() + self.updated = datetime.now(timezone.utc) return self diff --git a/src/zenml/zen_stores/schemas/flavor_schemas.py b/src/zenml/zen_stores/schemas/flavor_schemas.py index 44a996bb23c..a17d74501f7 100644 --- a/src/zenml/zen_stores/schemas/flavor_schemas.py +++ b/src/zenml/zen_stores/schemas/flavor_schemas.py @@ -14,7 +14,7 @@ """SQL Model Implementations for Flavors.""" import json -from datetime import datetime +from datetime import datetime, timezone from typing import Any, Optional from uuid import UUID @@ -103,7 +103,7 @@ def update(self, flavor_update: "FlavorUpdate") -> "FlavorSchema": else: setattr(self, field, value) - self.updated = datetime.utcnow() + self.updated = datetime.now(timezone.utc) return self def to_model( diff --git a/src/zenml/zen_stores/schemas/model_schemas.py b/src/zenml/zen_stores/schemas/model_schemas.py index 41c186c75ca..eb589d75fad 100644 --- a/src/zenml/zen_stores/schemas/model_schemas.py +++ b/src/zenml/zen_stores/schemas/model_schemas.py @@ -13,7 +13,7 @@ # permissions and limitations under the License. """SQLModel implementation of model tables.""" -from datetime import datetime +from datetime import datetime, timezone from typing import TYPE_CHECKING, Any, Dict, List, Optional, cast from uuid import UUID, uuid4 @@ -225,7 +225,7 @@ def update( exclude_unset=True, exclude_none=True ).items(): setattr(self, field, value) - self.updated = datetime.utcnow() + self.updated = datetime.now(timezone.utc) return self @@ -504,7 +504,7 @@ def update( self.name = target_name if target_description is not None: self.description = target_description - self.updated = datetime.utcnow() + self.updated = datetime.now(timezone.utc) return self diff --git a/src/zenml/zen_stores/schemas/pipeline_run_schemas.py b/src/zenml/zen_stores/schemas/pipeline_run_schemas.py index 67236f0ab7d..42c7cd91900 100644 --- a/src/zenml/zen_stores/schemas/pipeline_run_schemas.py +++ b/src/zenml/zen_stores/schemas/pipeline_run_schemas.py @@ -14,7 +14,7 @@ """SQLModel implementation of pipeline run tables.""" import json -from datetime import datetime +from datetime import datetime, timezone from typing import TYPE_CHECKING, Any, Dict, List, Optional from uuid import UUID @@ -431,7 +431,7 @@ def update(self, run_update: "PipelineRunUpdate") -> "PipelineRunSchema": 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.utcnow() + self.updated = datetime.now(timezone.utc) return self def update_placeholder( @@ -472,7 +472,7 @@ def update_placeholder( self.orchestrator_environment = orchestrator_environment self.status = request.status.value - self.updated = datetime.utcnow() + self.updated = datetime.now(timezone.utc) return self diff --git a/src/zenml/zen_stores/schemas/pipeline_schemas.py b/src/zenml/zen_stores/schemas/pipeline_schemas.py index 3719a64b207..a8aefd45b94 100644 --- a/src/zenml/zen_stores/schemas/pipeline_schemas.py +++ b/src/zenml/zen_stores/schemas/pipeline_schemas.py @@ -13,7 +13,7 @@ # permissions and limitations under the License. """SQL Model Implementations for Pipelines and Pipeline Runs.""" -from datetime import datetime +from datetime import datetime, timezone from typing import TYPE_CHECKING, Any, List, Optional from uuid import UUID @@ -183,5 +183,5 @@ def update(self, pipeline_update: "PipelineUpdate") -> "PipelineSchema": The updated `PipelineSchema`. """ self.description = pipeline_update.description - self.updated = datetime.utcnow() + self.updated = datetime.now(timezone.utc) return self diff --git a/src/zenml/zen_stores/schemas/run_template_schemas.py b/src/zenml/zen_stores/schemas/run_template_schemas.py index c2869e099f4..4914c2759c7 100644 --- a/src/zenml/zen_stores/schemas/run_template_schemas.py +++ b/src/zenml/zen_stores/schemas/run_template_schemas.py @@ -13,7 +13,7 @@ # permissions and limitations under the License. """SQLModel implementation of run template tables.""" -from datetime import datetime +from datetime import datetime, timezone from typing import TYPE_CHECKING, Any, List, Optional from uuid import UUID @@ -154,7 +154,7 @@ def update(self, update: RunTemplateUpdate) -> "RunTemplateSchema": ).items(): setattr(self, field, value) - self.updated = datetime.utcnow() + self.updated = datetime.now(timezone.utc) return self def to_model( diff --git a/src/zenml/zen_stores/schemas/schedule_schema.py b/src/zenml/zen_stores/schemas/schedule_schema.py index 632744ce3eb..5a56765bec4 100644 --- a/src/zenml/zen_stores/schemas/schedule_schema.py +++ b/src/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 +from datetime import datetime, timedelta, timezone from typing import TYPE_CHECKING, Any, Optional from uuid import UUID @@ -153,7 +153,7 @@ def update(self, schedule_update: ScheduleUpdate) -> "ScheduleSchema": ) if schedule_update.catchup is not None: self.catchup = schedule_update.catchup - self.updated = datetime.utcnow() + self.updated = datetime.now(timezone.utc) return self def to_model( diff --git a/src/zenml/zen_stores/schemas/secret_schemas.py b/src/zenml/zen_stores/schemas/secret_schemas.py index e5322d69a30..1d17664648b 100644 --- a/src/zenml/zen_stores/schemas/secret_schemas.py +++ b/src/zenml/zen_stores/schemas/secret_schemas.py @@ -15,7 +15,7 @@ import base64 import json -from datetime import datetime +from datetime import datetime, timezone from typing import Any, Dict, Optional, cast from uuid import UUID @@ -209,7 +209,7 @@ def update( else: setattr(self, field, value) - self.updated = datetime.utcnow() + self.updated = datetime.now(timezone.utc) return self def to_model( diff --git a/src/zenml/zen_stores/schemas/server_settings_schemas.py b/src/zenml/zen_stores/schemas/server_settings_schemas.py index 1ed5b03157f..7ee48ca3688 100644 --- a/src/zenml/zen_stores/schemas/server_settings_schemas.py +++ b/src/zenml/zen_stores/schemas/server_settings_schemas.py @@ -14,7 +14,7 @@ """SQLModel implementation for the server settings table.""" import json -from datetime import datetime +from datetime import datetime, timezone from typing import Any, Optional, Set from uuid import UUID @@ -42,8 +42,12 @@ class ServerSettingsSchema(SQLModel, table=True): display_announcements: Optional[bool] = Field(nullable=True) display_updates: Optional[bool] = Field(nullable=True) onboarding_state: Optional[str] = Field(nullable=True) - last_user_activity: datetime = Field(default_factory=datetime.utcnow) - updated: datetime = Field(default_factory=datetime.utcnow) + last_user_activity: datetime = Field( + default_factory=lambda: datetime.now(timezone.utc) + ) + updated: datetime = Field( + default_factory=lambda: datetime.now(timezone.utc) + ) def update( self, settings_update: ServerSettingsUpdate @@ -63,7 +67,7 @@ def update( if hasattr(self, field): setattr(self, field, value) - self.updated = datetime.utcnow() + self.updated = datetime.now(timezone.utc) return self @@ -83,7 +87,7 @@ def update_onboarding_state( ) new_state = old_state.union(completed_steps) self.onboarding_state = json.dumps(list(new_state)) - self.updated = datetime.utcnow() + self.updated = datetime.now(timezone.utc) return self diff --git a/src/zenml/zen_stores/schemas/service_connector_schemas.py b/src/zenml/zen_stores/schemas/service_connector_schemas.py index e59c50f836d..9573e573a03 100644 --- a/src/zenml/zen_stores/schemas/service_connector_schemas.py +++ b/src/zenml/zen_stores/schemas/service_connector_schemas.py @@ -15,7 +15,7 @@ import base64 import json -from datetime import datetime +from datetime import datetime, timezone from typing import TYPE_CHECKING, Any, Dict, List, Optional, cast from uuid import UUID @@ -227,7 +227,7 @@ def update( else: setattr(self, field, value) self.secret_id = secret_id - self.updated = datetime.utcnow() + self.updated = datetime.now(timezone.utc) return self def to_model( diff --git a/src/zenml/zen_stores/schemas/service_schemas.py b/src/zenml/zen_stores/schemas/service_schemas.py index 5abb5fcae5e..29f1a97544e 100644 --- a/src/zenml/zen_stores/schemas/service_schemas.py +++ b/src/zenml/zen_stores/schemas/service_schemas.py @@ -15,7 +15,7 @@ import base64 import json -from datetime import datetime +from datetime import datetime, timezone from typing import Any, Optional from uuid import UUID @@ -210,7 +210,7 @@ def update( ) else: setattr(self, field, value) - self.updated = datetime.utcnow() + self.updated = datetime.now(timezone.utc) return self @classmethod diff --git a/src/zenml/zen_stores/schemas/stack_schemas.py b/src/zenml/zen_stores/schemas/stack_schemas.py index 4307d2cb6ff..e4c68d833f9 100644 --- a/src/zenml/zen_stores/schemas/stack_schemas.py +++ b/src/zenml/zen_stores/schemas/stack_schemas.py @@ -15,7 +15,7 @@ import base64 import json -from datetime import datetime +from datetime import datetime, timezone from typing import TYPE_CHECKING, Any, List, Optional from uuid import UUID @@ -134,7 +134,7 @@ def update( else: setattr(self, field, value) - self.updated = datetime.utcnow() + self.updated = datetime.now(timezone.utc) return self def to_model( diff --git a/src/zenml/zen_stores/schemas/step_run_schemas.py b/src/zenml/zen_stores/schemas/step_run_schemas.py index ea01de1ab24..acb192a82cd 100644 --- a/src/zenml/zen_stores/schemas/step_run_schemas.py +++ b/src/zenml/zen_stores/schemas/step_run_schemas.py @@ -14,7 +14,7 @@ """SQLModel implementation of step run tables.""" import json -from datetime import datetime +from datetime import datetime, timezone from typing import TYPE_CHECKING, Any, Dict, List, Optional from uuid import UUID @@ -360,7 +360,7 @@ def update(self, step_update: "StepRunUpdate") -> "StepRunSchema": if value and self.model_version_id is None: self.model_version_id = value - self.updated = datetime.utcnow() + self.updated = datetime.now(timezone.utc) return self diff --git a/src/zenml/zen_stores/schemas/tag_schemas.py b/src/zenml/zen_stores/schemas/tag_schemas.py index d1e6a0483b4..c1bf397f81f 100644 --- a/src/zenml/zen_stores/schemas/tag_schemas.py +++ b/src/zenml/zen_stores/schemas/tag_schemas.py @@ -13,7 +13,7 @@ # permissions and limitations under the License. """SQLModel implementation of tag tables.""" -from datetime import datetime +from datetime import datetime, timezone from typing import TYPE_CHECKING, Any, List from uuid import UUID @@ -113,7 +113,7 @@ def update(self, update: TagUpdate) -> "TagSchema": else: setattr(self, field, value) - self.updated = datetime.utcnow() + self.updated = datetime.now(timezone.utc) return self diff --git a/src/zenml/zen_stores/schemas/trigger_schemas.py b/src/zenml/zen_stores/schemas/trigger_schemas.py index 63265b385f0..7388fefe996 100644 --- a/src/zenml/zen_stores/schemas/trigger_schemas.py +++ b/src/zenml/zen_stores/schemas/trigger_schemas.py @@ -15,7 +15,7 @@ import base64 import json -from datetime import datetime +from datetime import datetime, timezone from typing import Any, List, Optional, cast from uuid import UUID @@ -133,7 +133,7 @@ def update(self, trigger_update: "TriggerUpdate") -> "TriggerSchema": else: setattr(self, field, value) - self.updated = datetime.utcnow() + self.updated = datetime.now(timezone.utc) return self @classmethod diff --git a/src/zenml/zen_stores/schemas/user_schemas.py b/src/zenml/zen_stores/schemas/user_schemas.py index 2541c40deee..40b6eee0a91 100644 --- a/src/zenml/zen_stores/schemas/user_schemas.py +++ b/src/zenml/zen_stores/schemas/user_schemas.py @@ -14,7 +14,7 @@ """SQLModel implementation of user tables.""" import json -from datetime import datetime +from datetime import datetime, timezone from typing import TYPE_CHECKING, Any, List, Optional from uuid import UUID @@ -225,7 +225,7 @@ def update_user(self, user_update: UserUpdate) -> "UserSchema": else: setattr(self, field, value) - self.updated = datetime.utcnow() + self.updated = datetime.now(timezone.utc) return self def update_service_account( @@ -245,7 +245,7 @@ def update_service_account( ).items(): setattr(self, field, value) - self.updated = datetime.utcnow() + self.updated = datetime.now(timezone.utc) return self def to_model( diff --git a/src/zenml/zen_stores/schemas/workspace_schemas.py b/src/zenml/zen_stores/schemas/workspace_schemas.py index 37d911954b8..37cda2e92fb 100644 --- a/src/zenml/zen_stores/schemas/workspace_schemas.py +++ b/src/zenml/zen_stores/schemas/workspace_schemas.py @@ -13,7 +13,7 @@ # permissions and limitations under the License. """SQL Model Implementations for Workspaces.""" -from datetime import datetime +from datetime import datetime, timezone from typing import TYPE_CHECKING, Any, List from sqlmodel import Relationship @@ -168,7 +168,7 @@ def update(self, workspace_update: WorkspaceUpdate) -> "WorkspaceSchema": ).items(): setattr(self, field, value) - self.updated = datetime.utcnow() + self.updated = datetime.now(timezone.utc) return self def to_model( diff --git a/src/zenml/zen_stores/sql_zen_store.py b/src/zenml/zen_stores/sql_zen_store.py index 19bdda8b28f..9c535e573ca 100644 --- a/src/zenml/zen_stores/sql_zen_store.py +++ b/src/zenml/zen_stores/sql_zen_store.py @@ -8628,10 +8628,19 @@ def _update_pipeline_run_status( ExecutionStatus.COMPLETED, ExecutionStatus.FAILED, }: - run_update.end_time = datetime.utcnow() + run_update.end_time = datetime.now(timezone.utc) if pipeline_run.start_time and isinstance( pipeline_run.start_time, datetime ): + # We need to ensure both datetimes are timezone-aware to avoid TypeError when subtracting + # Now calculate the duration time + if pipeline_run.start_time.tzinfo is None: + pipeline_run.start_time = ( + pipeline_run.start_time.replace( + tzinfo=timezone.utc + ) + ) + duration_time = ( run_update.end_time - pipeline_run.start_time )