diff --git a/src/zenml/analytics/context.py b/src/zenml/analytics/context.py index a2e75a75906..0dd4c8fb84a 100644 --- a/src/zenml/analytics/context.py +++ b/src/zenml/analytics/context.py @@ -17,7 +17,6 @@ The base functionalities are adapted to work with the ZenML analytics server. """ -import datetime import locale from types import TracebackType from typing import TYPE_CHECKING, Any, Dict, List, Optional, Type, Union @@ -32,6 +31,7 @@ ) from zenml.environment import Environment, get_environment from zenml.logger import get_logger +from zenml.utils.time_utils import utc_now_tz_aware if TYPE_CHECKING: from zenml.analytics.enums import AnalyticsEvent @@ -284,11 +284,7 @@ def track( try: # Timezone as tzdata - tz = ( - datetime.datetime.now(datetime.timezone.utc) - .astimezone() - .tzname() - ) + tz = utc_now_tz_aware().astimezone().tzname() if tz is not None: properties.update({"timezone": tz}) diff --git a/src/zenml/cli/service_connectors.py b/src/zenml/cli/service_connectors.py index 3958777b5bb..68d201e2dff 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, timezone +from datetime import datetime from typing import Any, Dict, List, Optional, Union, cast from uuid import UUID @@ -25,7 +25,6 @@ is_sorted_or_filtered, list_options, print_page_info, - seconds_to_human_readable, ) from zenml.client import Client from zenml.console import console @@ -37,6 +36,7 @@ ServiceConnectorResourcesModel, ServiceConnectorResponse, ) +from zenml.utils.time_utils import seconds_to_human_readable, utc_now # Service connectors @@ -292,7 +292,7 @@ def prompt_expires_at( default_str = "" if default is not None: seconds = int( - (default - datetime.now(timezone.utc)).total_seconds() + (default - utc_now(tz_aware=default)).total_seconds() ) default_str = ( f" [{str(default)} i.e. in " @@ -309,7 +309,7 @@ def prompt_expires_at( assert expires_at is not None assert isinstance(expires_at, datetime) - if expires_at < datetime.now(timezone.utc): + if expires_at < utc_now(tz_aware=expires_at): cli_utils.warning( "The expiration time must be in the future. Please enter a " "later date and time." @@ -317,7 +317,7 @@ def prompt_expires_at( continue seconds = int( - (expires_at - datetime.now(timezone.utc)).total_seconds() + (expires_at - utc_now(tz_aware=expires_at)).total_seconds() ) confirm = click.confirm( diff --git a/src/zenml/cli/stack.py b/src/zenml/cli/stack.py index b022df5be77..b1cab129a2c 100644 --- a/src/zenml/cli/stack.py +++ b/src/zenml/cli/stack.py @@ -17,7 +17,6 @@ import re import time import webbrowser -from datetime import datetime, timezone from typing import ( TYPE_CHECKING, Any, @@ -77,6 +76,7 @@ ) from zenml.utils import requirements_utils from zenml.utils.dashboard_utils import get_component_url, get_stack_url +from zenml.utils.time_utils import utc_now_tz_aware from zenml.utils.yaml_utils import read_yaml, write_yaml if TYPE_CHECKING: @@ -1575,7 +1575,7 @@ def deploy( ): raise click.Abort() - date_start = datetime.now(timezone.utc) + date_start = utc_now_tz_aware() webbrowser.open(deployment_config.deployment_url) console.print( diff --git a/src/zenml/cli/utils.py b/src/zenml/cli/utils.py index bda4d966c9d..e382b0cfce8 100644 --- a/src/zenml/cli/utils.py +++ b/src/zenml/cli/utils.py @@ -14,7 +14,6 @@ """Utility functions for the CLI.""" import contextlib -import datetime import json import os import platform @@ -1581,6 +1580,7 @@ def print_components_table( configurations.append(component_config) print_table(configurations) + def print_service_connectors_table( client: "Client", connectors: Sequence["ServiceConnectorResponse"], diff --git a/src/zenml/config/pipeline_configurations.py b/src/zenml/config/pipeline_configurations.py index f33c6711819..0285d0d7b1c 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, timezone +from datetime import datetime from typing import TYPE_CHECKING, Any, Dict, List, Optional from pydantic import SerializeAsAny, field_validator @@ -23,6 +23,7 @@ from zenml.config.source import SourceWithValidator from zenml.config.strict_base_model import StrictBaseModel from zenml.model.model import Model +from zenml.utils.time_utils import utc_now if TYPE_CHECKING: from zenml.config import DockerSettings @@ -61,7 +62,7 @@ def _get_full_substitutions( The full substitutions dict including date and time. """ if start_time is None: - start_time = datetime.now(timezone.utc) + start_time = utc_now() 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 4b93a6da17c..ddbc095ada0 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, timezone +from datetime import datetime, timedelta from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Tuple from zenml import EventSourceResponse @@ -28,6 +28,7 @@ TriggerExecutionResponse, TriggerResponse, ) +from zenml.utils.time_utils import utc_now from zenml.zen_server.auth import AuthContext from zenml.zen_server.jwt import JWTToken @@ -134,9 +135,7 @@ def trigger_action( ) expires: Optional[datetime] = None if trigger.action.auth_window: - expires = datetime.now(timezone.utc) + timedelta( - minutes=trigger.action.auth_window - ) + expires = utc_now() + timedelta(minutes=trigger.action.auth_window) encoded_token = token.encode(expires=expires) auth_context = AuthContext( user=trigger.action.service_account, diff --git a/src/zenml/integrations/airflow/orchestrators/airflow_orchestrator.py b/src/zenml/integrations/airflow/orchestrators/airflow_orchestrator.py index fcb46d143e3..4195ba26837 100644 --- a/src/zenml/integrations/airflow/orchestrators/airflow_orchestrator.py +++ b/src/zenml/integrations/airflow/orchestrators/airflow_orchestrator.py @@ -42,6 +42,7 @@ from zenml.orchestrators.utils import get_orchestrator_run_name from zenml.stack import StackValidator from zenml.utils import io_utils +from zenml.utils.time_utils import utc_now if TYPE_CHECKING: from zenml.config import ResourceSettings @@ -408,8 +409,7 @@ def _translate_schedule( if schedule: if schedule.cron_expression: start_time = schedule.start_time or ( - datetime.datetime.now(datetime.timezone.utc) - - datetime.timedelta(7) + utc_now() - datetime.timedelta(7) ) return { "schedule": schedule.cron_expression, @@ -429,7 +429,6 @@ 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.now(datetime.timezone.utc) - - datetime.timedelta(7), + "start_date": utc_now() - datetime.timedelta(7), "catchup": False, } diff --git a/src/zenml/integrations/aws/orchestrators/sagemaker_orchestrator.py b/src/zenml/integrations/aws/orchestrators/sagemaker_orchestrator.py index 52ba430b082..164192f9797 100644 --- a/src/zenml/integrations/aws/orchestrators/sagemaker_orchestrator.py +++ b/src/zenml/integrations/aws/orchestrators/sagemaker_orchestrator.py @@ -15,7 +15,7 @@ import os import re -from datetime import datetime, timezone +from datetime import timezone from typing import ( TYPE_CHECKING, Any, @@ -64,6 +64,7 @@ from zenml.orchestrators.utils import get_orchestrator_run_name from zenml.stack import StackValidator from zenml.utils.env_utils import split_environment_variables +from zenml.utils.time_utils import utc_now if TYPE_CHECKING: from zenml.models import PipelineDeploymentResponse, PipelineRunResponse @@ -553,8 +554,7 @@ def prepare_or_run_pipeline( enabled=True, ) next_execution = ( - deployment.schedule.start_time - or datetime.now(timezone.utc) + deployment.schedule.start_time or utc_now() ) + deployment.schedule.interval_second else: # One-time schedule diff --git a/src/zenml/integrations/aws/service_connectors/aws_service_connector.py b/src/zenml/integrations/aws/service_connectors/aws_service_connector.py index aead0098e12..793e782f185 100644 --- a/src/zenml/integrations/aws/service_connectors/aws_service_connector.py +++ b/src/zenml/integrations/aws/service_connectors/aws_service_connector.py @@ -66,6 +66,7 @@ ) from zenml.utils.enum_utils import StrEnum from zenml.utils.secret_utils import PlainSerializedSecretStr +from zenml.utils.time_utils import utc_now_tz_aware logger = get_logger(__name__) @@ -711,7 +712,7 @@ def get_boto3_session( return session, None # Refresh expired sessions - now = datetime.datetime.now(datetime.timezone.utc) + now = utc_now_tz_aware() expires_at = expires_at.replace(tzinfo=datetime.timezone.utc) # check if the token expires in the near future if expires_at > now + datetime.timedelta( @@ -959,9 +960,7 @@ def _authenticate( # determine the expiration time of the temporary credentials # from the boto3 session, so we assume the default IAM role # expiration date is used - expiration_time = datetime.datetime.now( - tz=datetime.timezone.utc - ) + datetime.timedelta( + expiration_time = utc_now_tz_aware() + datetime.timedelta( seconds=DEFAULT_IAM_ROLE_TOKEN_EXPIRATION ) return session, expiration_time @@ -1673,9 +1672,7 @@ def _auto_configure( # expiration time of the temporary credentials from the # boto3 session, so we assume the default IAM role # expiration period is used - expires_at = datetime.datetime.now( - tz=datetime.timezone.utc - ) + datetime.timedelta( + expires_at = utc_now_tz_aware() + datetime.timedelta( seconds=DEFAULT_IAM_ROLE_TOKEN_EXPIRATION ) @@ -1720,9 +1717,7 @@ def _auto_configure( aws_secret_access_key=credentials["SecretAccessKey"], aws_session_token=credentials["SessionToken"], ) - expires_at = datetime.datetime.now( - tz=datetime.timezone.utc - ) + datetime.timedelta( + expires_at = utc_now_tz_aware() + datetime.timedelta( seconds=DEFAULT_STS_TOKEN_EXPIRATION ) @@ -2130,9 +2125,9 @@ def _get_connector_client( # Kubernetes authentication tokens issued by AWS EKS have a fixed # expiration time of 15 minutes # source: https://aws.github.io/aws-eks-best-practices/security/docs/iam/#controlling-access-to-eks-clusters - expires_at = datetime.datetime.now( - tz=datetime.timezone.utc - ) + datetime.timedelta(minutes=EKS_KUBE_API_TOKEN_EXPIRATION) + expires_at = utc_now_tz_aware() + datetime.timedelta( + minutes=EKS_KUBE_API_TOKEN_EXPIRATION + ) # get cluster details cluster_arn = cluster["cluster"]["arn"] diff --git a/src/zenml/integrations/azure/service_connectors/azure_service_connector.py b/src/zenml/integrations/azure/service_connectors/azure_service_connector.py index 515efb4671e..37c11c627d4 100644 --- a/src/zenml/integrations/azure/service_connectors/azure_service_connector.py +++ b/src/zenml/integrations/azure/service_connectors/azure_service_connector.py @@ -58,6 +58,7 @@ ) from zenml.utils.enum_utils import StrEnum from zenml.utils.secret_utils import PlainSerializedSecretStr +from zenml.utils.time_utils import to_local_tz, utc_now # Configure the logging level for azure.identity logging.getLogger("azure.identity").setLevel(logging.WARNING) @@ -171,12 +172,7 @@ def __init__(self, token: str, expires_at: datetime.datetime): self.token = token # Convert the expiration time from UTC to local time - expires_at.replace(tzinfo=datetime.timezone.utc) - expires_at = expires_at.astimezone( - datetime.datetime.now().astimezone().tzinfo - ) - - self.expires_on = int(expires_at.timestamp()) + self.expires_on = int(to_local_tz(expires_at).timestamp()) def get_token(self, *scopes: str, **kwargs: Any) -> Any: """Get token. @@ -604,11 +600,9 @@ def get_azure_credential( return session, None # Refresh expired sessions - now = datetime.datetime.now(datetime.timezone.utc) - expires_at = expires_at.replace(tzinfo=datetime.timezone.utc) # check if the token expires in the near future - if expires_at > now + datetime.timedelta( + if expires_at > utc_now(tz_aware=expires_at) + datetime.timedelta( minutes=AZURE_SESSION_EXPIRATION_BUFFER ): return session, expires_at diff --git a/src/zenml/integrations/gcp/service_connectors/gcp_service_connector.py b/src/zenml/integrations/gcp/service_connectors/gcp_service_connector.py index 959cdd7e56d..6a91b760f55 100644 --- a/src/zenml/integrations/gcp/service_connectors/gcp_service_connector.py +++ b/src/zenml/integrations/gcp/service_connectors/gcp_service_connector.py @@ -78,6 +78,7 @@ from zenml.utils.enum_utils import StrEnum from zenml.utils.pydantic_utils import before_validator_handler from zenml.utils.secret_utils import PlainSerializedSecretStr +from zenml.utils.time_utils import utc_now logger = get_logger(__name__) @@ -1124,10 +1125,9 @@ def get_session( return session, None # Refresh expired sessions - now = datetime.datetime.now(datetime.timezone.utc) - expires_at = expires_at.replace(tzinfo=datetime.timezone.utc) + # check if the token expires in the near future - if expires_at > now + datetime.timedelta( + if expires_at > utc_now(tz_aware=expires_at) + datetime.timedelta( minutes=GCP_SESSION_EXPIRATION_BUFFER ): return session, expires_at diff --git a/src/zenml/integrations/kubernetes/orchestrators/kube_utils.py b/src/zenml/integrations/kubernetes/orchestrators/kube_utils.py index e97142df086..0b74341c4cf 100644 --- a/src/zenml/integrations/kubernetes/orchestrators/kube_utils.py +++ b/src/zenml/integrations/kubernetes/orchestrators/kube_utils.py @@ -31,7 +31,6 @@ Adjusted from https://github.com/tensorflow/tfx/blob/master/tfx/utils/kube_utils.py. """ -import datetime import enum import re import time @@ -47,6 +46,7 @@ build_service_account_manifest, ) from zenml.logger import get_logger +from zenml.utils.time_utils import utc_now logger = get_logger(__name__) @@ -248,7 +248,7 @@ def wait_pod( Returns: The pod object which meets the exit condition. """ - start_time = datetime.datetime.now(datetime.timezone.utc) + start_time = utc_now() # Link to exponential back-off algorithm used here: # https://cloud.google.com/storage/docs/exponential-backoff @@ -288,7 +288,7 @@ def wait_pod( return resp # Check if wait timed out. - elapse_time = datetime.datetime.now(datetime.timezone.utc) - start_time + elapse_time = utc_now() - start_time if elapse_time.seconds >= timeout_sec and timeout_sec != 0: raise RuntimeError( f"Waiting for pod `{namespace}:{pod_name}` timed out after " 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 99ed49d0df9..d35828ab1be 100644 --- a/src/zenml/integrations/whylogs/data_validators/whylogs_data_validator.py +++ b/src/zenml/integrations/whylogs/data_validators/whylogs_data_validator.py @@ -34,6 +34,7 @@ ) from zenml.logger import get_logger from zenml.stack.authentication_mixin import AuthenticationMixin +from zenml.utils.time_utils import utc_now logger = get_logger(__name__) @@ -97,9 +98,7 @@ def data_profiling( """ results = why.log(pandas=dataset) profile = results.profile() - dataset_timestamp = dataset_timestamp or datetime.datetime.now( - datetime.timezone.utc - ) + dataset_timestamp = dataset_timestamp or utc_now() profile.set_dataset_timestamp(dataset_timestamp=dataset_timestamp) return profile.view() diff --git a/src/zenml/logging/step_logging.py b/src/zenml/logging/step_logging.py index b31ea363041..938a95c50dc 100644 --- a/src/zenml/logging/step_logging.py +++ b/src/zenml/logging/step_logging.py @@ -13,7 +13,6 @@ # permissions and limitations under the License. """ZenML logging handler.""" -import datetime import os import re import sys @@ -35,6 +34,7 @@ STEP_LOGS_STORAGE_MAX_MESSAGES, STEP_LOGS_STORAGE_MERGE_INTERVAL_SECONDS, ) +from zenml.utils.time_utils import utc_now from zenml.zen_stores.base_zen_store import BaseZenStore # Get the logger @@ -303,9 +303,9 @@ def save_to_file(self, force: bool = False) -> None: "w", ) as file: for message in self.buffer: - timestamp = datetime.datetime.now( - datetime.timezone.utc - ).strftime("%Y-%m-%d %H:%M:%S") + timestamp = utc_now().strftime( + "%Y-%m-%d %H:%M:%S" + ) file.write( f"[{timestamp} UTC] {remove_ansi_escape_codes(message)}\n" ) @@ -314,9 +314,9 @@ def save_to_file(self, force: bool = False) -> None: self.logs_uri, "a" ) as file: for message in self.buffer: - timestamp = datetime.datetime.now( - datetime.timezone.utc - ).strftime("%Y-%m-%d %H:%M:%S") + timestamp = utc_now().strftime( + "%Y-%m-%d %H:%M:%S" + ) file.write( f"[{timestamp} UTC] {remove_ansi_escape_codes(message)}\n" ) diff --git a/src/zenml/login/credentials.py b/src/zenml/login/credentials.py index ca76eddafb2..ed401c6faf9 100644 --- a/src/zenml/login/credentials.py +++ b/src/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.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 utc_now class ServerType(StrEnum): @@ -71,7 +72,7 @@ def expired(self) -> bool: 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 @@ def auth_status(self) -> str: 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})" @@ -259,7 +260,7 @@ def expires_in(self) -> str: 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 diff --git a/src/zenml/login/credentials_store.py b/src/zenml/login/credentials_store.py index efb9c0f895d..d7658dc65a1 100644 --- a/src/zenml/login/credentials_store.py +++ b/src/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.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 @@ def _save_credentials(self) -> None: 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 @@ def set_token( """ 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 diff --git a/src/zenml/models/v2/core/api_key.py b/src/zenml/models/v2/core/api_key.py index d6855add252..99a10c7e5eb 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, 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.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 @@ def verify_key( 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 diff --git a/src/zenml/orchestrators/publish_utils.py b/src/zenml/orchestrators/publish_utils.py index 778736b37ed..f47156c15de 100644 --- a/src/zenml/orchestrators/publish_utils.py +++ b/src/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 @@ 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(), ), ) diff --git a/src/zenml/orchestrators/step_launcher.py b/src/zenml/orchestrators/step_launcher.py index cddf12e8baf..e07c72445be 100644 --- a/src/zenml/orchestrators/step_launcher.py +++ b/src/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.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 @@ def launch(self) -> None: 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 @@ 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.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( diff --git a/src/zenml/orchestrators/step_run_utils.py b/src/zenml/orchestrators/step_run_utils.py index 5b7355c46bf..2177f3ce5d5 100644 --- a/src/zenml/orchestrators/step_run_utils.py +++ b/src/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.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 @@ 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.now(timezone.utc), + start_time=utc_now(), 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 0d198f4e40f..f978b2c8265 100644 --- a/src/zenml/pipelines/run_utils.py +++ b/src/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.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, diff --git a/src/zenml/service_connectors/service_connector.py b/src/zenml/service_connectors/service_connector.py index 1e190643dfb..64d22015db8 100644 --- a/src/zenml/service_connectors/service_connector.py +++ b/src/zenml/service_connectors/service_connector.py @@ -55,6 +55,7 @@ UserResponse, WorkspaceResponse, ) +from zenml.utils.time_utils import utc_now logger = get_logger(__name__) @@ -794,13 +795,14 @@ def to_response_model( "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 @@ def has_expired(self) -> bool: 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" ) diff --git a/src/zenml/stack/stack.py b/src/zenml/stack/stack.py index 90b164a09f3..a29b5b08977 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, timezone +from datetime import datetime from typing import ( TYPE_CHECKING, AbstractSet, @@ -45,6 +45,7 @@ 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 @@ def validate_image_builder(self) -> None: 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 @@ def validate_image_builder(self) -> None: flavor = LocalImageBuilderFlavor() + now = utc_now() image_builder = LocalImageBuilder( id=uuid4(), name="temporary_default", @@ -751,8 +752,8 @@ def validate_image_builder(self) -> None: 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 diff --git a/src/zenml/utils/string_utils.py b/src/zenml/utils/string_utils.py index 8c0838c7a39..49c391234d3 100644 --- a/src/zenml/utils/string_utils.py +++ b/src/zenml/utils/string_utils.py @@ -17,12 +17,12 @@ 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")) diff --git a/src/zenml/utils/time_utils.py b/src/zenml/utils/time_utils.py index b39ef19c559..6a88840a8e2 100644 --- a/src/zenml/utils/time_utils.py +++ b/src/zenml/utils/time_utils.py @@ -14,31 +14,58 @@ """Time utils.""" from datetime import datetime, timedelta, timezone -from typing import Optional +from typing import Optional, Union -def utc_now(tz: bool = False) -> datetime: +def utc_now(tz_aware: Union[bool, datetime] = False) -> datetime: """Get the current UTC time. Args: - tz: Whether to return a timezone-aware datetime. + 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 match the timezone of the input datetime. Returns: - The current UTC time. + 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. """ - if tz: - return datetime.now(timezone.utc) - else: - return datetime.now(timezone.utc).replace(tzinfo=None) + 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() -> datetime: +def utc_now_tz_aware() -> datetime: """Get the current timezone-aware UTC time. Returns: The current UTC time. """ - return utc_now(tz=True) + 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 UTC. + + 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 seconds_to_human_readable(time_seconds: int) -> str: @@ -84,8 +111,7 @@ def expires_in( Returns: Human readable string. """ - now = utc_now_tz() - expires_at = expires_at.replace(tzinfo=timezone.utc) + now = utc_now(tz_aware=expires_at) if skew_tolerance: expires_at -= timedelta(seconds=skew_tolerance) if expires_at < now: diff --git a/src/zenml/zen_server/auth.py b/src/zenml/zen_server/auth.py index 94de2940080..466d8a302ae 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, 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 @@ 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 diff --git a/src/zenml/zen_server/cloud_utils.py b/src/zenml/zen_server/cloud_utils.py index 0aff9135636..5cdbf5702f1 100644 --- a/src/zenml/zen_server/cloud_utils.py +++ b/src/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 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 @@ def _fetch_auth_token(self) -> str: 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 @@ def _fetch_auth_token(self) -> str: ) 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 diff --git a/src/zenml/zen_server/routers/devices_endpoints.py b/src/zenml/zen_server/routers/devices_endpoints.py index 95bc091ce7a..6cfe28fc83d 100644 --- a/src/zenml/zen_server/routers/devices_endpoints.py +++ b/src/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 @@ 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.", ) diff --git a/src/zenml/zen_server/zen_server_api.py b/src/zenml/zen_server/zen_server_api.py index 060e9db59cb..fe3a72a6713 100644 --- a/src/zenml/zen_server/zen_server_api.py +++ b/src/zenml/zen_server/zen_server_api.py @@ -22,7 +22,7 @@ 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.enums import AuthScheme, SourceContextTypes from zenml.models import ServerDeploymentType +from zenml.utils.time_utils import utc_now_tz_aware 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 @@ def relative_path(rel: str) -> str: ) # 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_tz_aware() +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_tz_aware() + 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 ) 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 2ecf44dd788..7d945a4e217 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,13 +6,14 @@ """ -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 = {} diff --git a/src/zenml/zen_stores/migrations/versions/3dcc5d20e82f_add_last_user_activity.py b/src/zenml/zen_stores/migrations/versions/3dcc5d20e82f_add_last_user_activity.py index b7079c1f2e3..6ef031a8608 100644 --- a/src/zenml/zen_stores/migrations/versions/3dcc5d20e82f_add_last_user_activity.py +++ b/src/zenml/zen_stores/migrations/versions/3dcc5d20e82f_add_last_user_activity.py @@ -6,12 +6,12 @@ """ -from datetime import datetime, timezone - import sqlalchemy as sa import sqlmodel from alembic import op +from zenml.utils.time_utils import utc_now_tz_aware + # 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_tz_aware())), ) with op.batch_alter_table("server_settings", schema=None) as batch_op: 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 32f8578b60b..014995f91ce 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,13 +6,14 @@ """ -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(), }, ], ) 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 fe7d77d263b..d38dafe73bf 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 @@ -6,7 +6,6 @@ """ -import datetime import uuid import sqlalchemy as sa @@ -14,6 +13,8 @@ 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, } ], ) 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 75fa38551a0..5f391ff957d 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,13 +8,14 @@ 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 = [] 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 370e2253f5b..d3473789c3f 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,13 +6,14 @@ """ -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, ) diff --git a/src/zenml/zen_stores/schemas/action_schemas.py b/src/zenml/zen_stores/schemas/action_schemas.py index 4a57cbb7ec5..a25c78bdbb0 100644 --- a/src/zenml/zen_stores/schemas/action_schemas.py +++ b/src/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 @@ 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 @@ def update(self, action_update: "ActionUpdate") -> "ActionSchema": else: setattr(self, field, value) - self.updated = datetime.now(timezone.utc) + self.updated = utc_now() 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 a1535bf8869..89ac3804a71 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, timezone +from datetime import datetime from secrets import token_hex from typing import Any, Optional, Tuple from uuid import UUID @@ -32,6 +32,7 @@ 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 @@ def from_request( """ 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 @@ def update(self, update: APIKeyUpdate) -> "APIKeySchema": 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 @@ def rotate( 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() diff --git a/src/zenml/zen_stores/schemas/artifact_schemas.py b/src/zenml/zen_stores/schemas/artifact_schemas.py index ddfa6e51366..90870ba329a 100644 --- a/src/zenml/zen_stores/schemas/artifact_schemas.py +++ b/src/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 @@ 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 @@ def update(self, artifact_update: ArtifactUpdate) -> "ArtifactSchema": 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 @@ def update( Returns: The updated `ArtifactVersionSchema`. """ - self.updated = datetime.now(timezone.utc) + self.updated = utc_now() return self diff --git a/src/zenml/zen_stores/schemas/base_schemas.py b/src/zenml/zen_stores/schemas/base_schemas.py index e0c1e1f7500..953032a2687 100644 --- a/src/zenml/zen_stores/schemas/base_schemas.py +++ b/src/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=lambda: utc_now()) + updated: datetime = Field(default_factory=lambda: utc_now()) 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 86a0c5d24e7..677947a46ae 100644 --- a/src/zenml/zen_stores/schemas/code_repository_schemas.py +++ b/src/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 @@ 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 @@ def update(self, update: "CodeRepositoryUpdate") -> "CodeRepositorySchema": if update.logo_url: self.logo_url = update.logo_url - self.updated = datetime.now(timezone.utc) + self.updated = utc_now() return self diff --git a/src/zenml/zen_stores/schemas/component_schemas.py b/src/zenml/zen_stores/schemas/component_schemas.py index 05b9cd4b9be..d4346cd5782 100644 --- a/src/zenml/zen_stores/schemas/component_schemas.py +++ b/src/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 @@ 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 @@ def update( else: setattr(self, field, value) - self.updated = datetime.now(timezone.utc) + self.updated = utc_now() 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 68ee10f7488..e77ea45e02d 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, 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 @@ 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 @@ 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.now(timezone.utc) + now = utc_now() return ( cls( client_id=request.client_id, @@ -159,7 +160,7 @@ def update(self, device_update: OAuthDeviceUpdate) -> "OAuthDeviceSchema": 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 @@ def internal_update( 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 diff --git a/src/zenml/zen_stores/schemas/event_source_schemas.py b/src/zenml/zen_stores/schemas/event_source_schemas.py index 5ee8d0742f2..6406e90b266 100644 --- a/src/zenml/zen_stores/schemas/event_source_schemas.py +++ b/src/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 @@ 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 @@ def update(self, update: EventSourceUpdate) -> "EventSourceSchema": ) else: setattr(self, field, value) - self.updated = datetime.now(timezone.utc) + self.updated = utc_now() return self diff --git a/src/zenml/zen_stores/schemas/flavor_schemas.py b/src/zenml/zen_stores/schemas/flavor_schemas.py index a17d74501f7..9f4c4441038 100644 --- a/src/zenml/zen_stores/schemas/flavor_schemas.py +++ b/src/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 @@ 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 @@ def update(self, flavor_update: "FlavorUpdate") -> "FlavorSchema": else: setattr(self, field, value) - self.updated = datetime.now(timezone.utc) + self.updated = utc_now() 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 3572a662fd0..a327e682b1a 100644 --- a/src/zenml/zen_stores/schemas/model_schemas.py +++ b/src/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 @@ 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 @@ def update( 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 @@ def update( 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 diff --git a/src/zenml/zen_stores/schemas/pipeline_run_schemas.py b/src/zenml/zen_stores/schemas/pipeline_run_schemas.py index 4283b4f2c8d..ac14279f504 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, timezone +from datetime import datetime from typing import TYPE_CHECKING, Any, Dict, List, Optional from uuid import UUID @@ -37,6 +37,7 @@ 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 @@ 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.now(timezone.utc) + self.updated = utc_now() return self def update_placeholder( @@ -478,7 +479,7 @@ def update_placeholder( self.orchestrator_environment = orchestrator_environment self.status = request.status.value - self.updated = datetime.now(timezone.utc) + self.updated = utc_now() return self diff --git a/src/zenml/zen_stores/schemas/pipeline_schemas.py b/src/zenml/zen_stores/schemas/pipeline_schemas.py index 76fc021e753..56e62ce9a4d 100644 --- a/src/zenml/zen_stores/schemas/pipeline_schemas.py +++ b/src/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 @@ 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 @@ def update(self, pipeline_update: "PipelineUpdate") -> "PipelineSchema": The updated `PipelineSchema`. """ self.description = pipeline_update.description - self.updated = datetime.now(timezone.utc) + self.updated = utc_now() 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 5ffaed01460..8646d55363e 100644 --- a/src/zenml/zen_stores/schemas/run_template_schemas.py +++ b/src/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 @@ 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 @@ def update(self, update: RunTemplateUpdate) -> "RunTemplateSchema": ).items(): setattr(self, field, value) - self.updated = datetime.now(timezone.utc) + self.updated = utc_now() 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 577e62f34be..c08fba6da94 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, timezone +from datetime import datetime, timedelta from typing import TYPE_CHECKING, Any, List, Optional from uuid import UUID @@ -27,6 +27,7 @@ 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 @@ def update(self, schedule_update: ScheduleUpdate) -> "ScheduleSchema": ) 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( diff --git a/src/zenml/zen_stores/schemas/secret_schemas.py b/src/zenml/zen_stores/schemas/secret_schemas.py index 1d17664648b..a439d0a2f80 100644 --- a/src/zenml/zen_stores/schemas/secret_schemas.py +++ b/src/zenml/zen_stores/schemas/secret_schemas.py @@ -15,7 +15,6 @@ import base64 import json -from datetime import datetime, timezone from typing import Any, Dict, Optional, cast from uuid import UUID @@ -35,6 +34,7 @@ SecretResponseMetadata, SecretUpdate, ) +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 @@ -209,7 +209,7 @@ def update( else: setattr(self, field, value) - self.updated = datetime.now(timezone.utc) + self.updated = utc_now() 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 7ee48ca3688..48550f42ff6 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, timezone +from datetime import datetime from typing import Any, Optional, Set from uuid import UUID @@ -27,6 +27,7 @@ ServerSettingsResponseResources, ServerSettingsUpdate, ) +from zenml.utils.time_utils import utc_now class ServerSettingsSchema(SQLModel, table=True): @@ -42,12 +43,8 @@ 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=lambda: datetime.now(timezone.utc) - ) - updated: datetime = Field( - default_factory=lambda: datetime.now(timezone.utc) - ) + last_user_activity: datetime = Field(default_factory=utc_now) + updated: datetime = Field(default_factory=utc_now) def update( self, settings_update: ServerSettingsUpdate @@ -67,7 +64,7 @@ def update( if hasattr(self, field): setattr(self, field, value) - self.updated = datetime.now(timezone.utc) + self.updated = utc_now() return self @@ -87,7 +84,7 @@ def update_onboarding_state( ) new_state = old_state.union(completed_steps) self.onboarding_state = json.dumps(list(new_state)) - self.updated = datetime.now(timezone.utc) + self.updated = utc_now() 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 9573e573a03..d9d5fd4293f 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, timezone +from datetime import datetime from typing import TYPE_CHECKING, Any, Dict, List, Optional, cast from uuid import UUID @@ -29,6 +29,7 @@ ServiceConnectorResponseMetadata, ServiceConnectorUpdate, ) +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 @@ -227,7 +228,7 @@ def update( else: setattr(self, field, value) self.secret_id = secret_id - self.updated = datetime.now(timezone.utc) + self.updated = utc_now() 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 29f1a97544e..75e8c6e92bc 100644 --- a/src/zenml/zen_stores/schemas/service_schemas.py +++ b/src/zenml/zen_stores/schemas/service_schemas.py @@ -15,7 +15,6 @@ import base64 import json -from datetime import datetime, timezone from typing import Any, Optional from uuid import UUID @@ -32,6 +31,7 @@ ServiceUpdate, ) from zenml.utils.dict_utils import dict_to_bytes +from zenml.utils.time_utils import utc_now from zenml.zen_stores.schemas.base_schemas import NamedSchema from zenml.zen_stores.schemas.model_schemas import ModelVersionSchema from zenml.zen_stores.schemas.pipeline_run_schemas import PipelineRunSchema @@ -210,7 +210,7 @@ def update( ) else: setattr(self, field, value) - self.updated = datetime.now(timezone.utc) + self.updated = utc_now() return self @classmethod diff --git a/src/zenml/zen_stores/schemas/stack_schemas.py b/src/zenml/zen_stores/schemas/stack_schemas.py index e4c68d833f9..58e19979b26 100644 --- a/src/zenml/zen_stores/schemas/stack_schemas.py +++ b/src/zenml/zen_stores/schemas/stack_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 @@ -27,6 +26,7 @@ StackResponseMetadata, StackUpdate, ) +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 @@ -134,7 +134,7 @@ def update( else: setattr(self, field, value) - self.updated = datetime.now(timezone.utc) + self.updated = utc_now() 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 6064b5e04c9..53998488e0a 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, timezone +from datetime import datetime from typing import TYPE_CHECKING, Any, Dict, List, Optional from uuid import UUID @@ -43,6 +43,7 @@ StepRunInputResponse, StepRunResponseResources, ) +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_deployment_schemas import ( @@ -358,7 +359,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.now(timezone.utc) + self.updated = utc_now() return self diff --git a/src/zenml/zen_stores/schemas/tag_schemas.py b/src/zenml/zen_stores/schemas/tag_schemas.py index 6228c59f2f4..e24a4504319 100644 --- a/src/zenml/zen_stores/schemas/tag_schemas.py +++ b/src/zenml/zen_stores/schemas/tag_schemas.py @@ -13,7 +13,6 @@ # permissions and limitations under the License. """SQLModel implementation of tag tables.""" -from datetime import datetime, timezone from typing import Any, List from uuid import UUID @@ -30,6 +29,7 @@ TagResponseBody, TagUpdate, ) +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 @@ -103,7 +103,7 @@ def update(self, update: TagUpdate) -> "TagSchema": else: setattr(self, field, value) - self.updated = datetime.now(timezone.utc) + self.updated = utc_now() return self diff --git a/src/zenml/zen_stores/schemas/trigger_schemas.py b/src/zenml/zen_stores/schemas/trigger_schemas.py index 7388fefe996..693b63df6b2 100644 --- a/src/zenml/zen_stores/schemas/trigger_schemas.py +++ b/src/zenml/zen_stores/schemas/trigger_schemas.py @@ -15,7 +15,6 @@ import base64 import json -from datetime import datetime, timezone from typing import Any, List, Optional, cast from uuid import UUID @@ -38,6 +37,7 @@ TriggerUpdate, ) from zenml.utils.json_utils import pydantic_encoder +from zenml.utils.time_utils import utc_now from zenml.zen_stores.schemas.action_schemas import ActionSchema from zenml.zen_stores.schemas.base_schemas import BaseSchema, NamedSchema from zenml.zen_stores.schemas.event_source_schemas import EventSourceSchema @@ -133,7 +133,7 @@ def update(self, trigger_update: "TriggerUpdate") -> "TriggerSchema": else: setattr(self, field, value) - self.updated = datetime.now(timezone.utc) + self.updated = utc_now() return self @classmethod diff --git a/src/zenml/zen_stores/schemas/user_schemas.py b/src/zenml/zen_stores/schemas/user_schemas.py index 40b6eee0a91..cfe0c19db1e 100644 --- a/src/zenml/zen_stores/schemas/user_schemas.py +++ b/src/zenml/zen_stores/schemas/user_schemas.py @@ -14,7 +14,6 @@ """SQLModel implementation of user tables.""" import json -from datetime import datetime, timezone from typing import TYPE_CHECKING, Any, List, Optional from uuid import UUID @@ -33,6 +32,7 @@ UserResponseMetadata, UserUpdate, ) +from zenml.utils.time_utils import utc_now from zenml.zen_stores.schemas.base_schemas import NamedSchema if TYPE_CHECKING: @@ -225,7 +225,7 @@ def update_user(self, user_update: UserUpdate) -> "UserSchema": else: setattr(self, field, value) - self.updated = datetime.now(timezone.utc) + self.updated = utc_now() return self def update_service_account( @@ -245,7 +245,7 @@ def update_service_account( ).items(): setattr(self, field, value) - self.updated = datetime.now(timezone.utc) + self.updated = utc_now() 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 37cda2e92fb..f5e62f8bf77 100644 --- a/src/zenml/zen_stores/schemas/workspace_schemas.py +++ b/src/zenml/zen_stores/schemas/workspace_schemas.py @@ -13,7 +13,6 @@ # permissions and limitations under the License. """SQL Model Implementations for Workspaces.""" -from datetime import datetime, timezone from typing import TYPE_CHECKING, Any, List from sqlmodel import Relationship @@ -25,6 +24,7 @@ WorkspaceResponseMetadata, WorkspaceUpdate, ) +from zenml.utils.time_utils import utc_now from zenml.zen_stores.schemas.base_schemas import NamedSchema if TYPE_CHECKING: @@ -168,7 +168,7 @@ def update(self, workspace_update: WorkspaceUpdate) -> "WorkspaceSchema": ).items(): setattr(self, field, value) - self.updated = datetime.now(timezone.utc) + self.updated = utc_now() 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 4f9301ae5d0..cbffb630cc9 100644 --- a/src/zenml/zen_stores/sql_zen_store.py +++ b/src/zenml/zen_stores/sql_zen_store.py @@ -309,6 +309,7 @@ random_str, validate_name, ) +from zenml.utils.time_utils import utc_now from zenml.zen_stores import template_utils from zenml.zen_stores.base_zen_store import ( BaseZenStore, @@ -4048,7 +4049,7 @@ def delete_expired_authorized_devices(self) -> None: # Delete devices that have expired if ( device.expires is not None - and device.expires < datetime.now() + and device.expires < utc_now() and device.user_id is None ): session.delete(device) @@ -8634,7 +8635,7 @@ def _update_pipeline_run_status( ExecutionStatus.COMPLETED, ExecutionStatus.FAILED, }: - run_update.end_time = datetime.now(timezone.utc) + run_update.end_time = utc_now() if pipeline_run.start_time and isinstance( pipeline_run.start_time, datetime ):