From 3f45131887602551036e83a86a71c9f7badf7b52 Mon Sep 17 00:00:00 2001 From: "Thomas J. Fan" Date: Fri, 15 Mar 2024 16:40:10 -0400 Subject: [PATCH] Migrate over to using datetime.now(timezone.utc) (#2270) Signed-off-by: Thomas J. Fan --- flytekit/bin/entrypoint.py | 2 +- flytekit/core/context_manager.py | 2 +- flytekit/core/mock_stats.py | 6 ++++-- flytekit/core/utils.py | 4 ++-- flytekit/experimental/eager_function.py | 10 +++++----- flytekit/remote/remote.py | 4 ++-- plugins/flytekit-deck-standard/tests/test_renderer.py | 4 ++-- .../flytekitplugins/great_expectations/schema.py | 2 +- .../flytekitplugins/great_expectations/task.py | 2 +- .../tests/test_flytekitplugins_iap.py | 4 ++-- tests/flytekit/integration/remote/test_remote.py | 4 ++-- tests/flytekit/unit/core/test_context_manager.py | 4 ++-- tests/flytekit/unit/deck/test_deck.py | 4 ++-- tests/flytekit/unit/models/test_literals.py | 2 +- 14 files changed, 28 insertions(+), 26 deletions(-) diff --git a/flytekit/bin/entrypoint.py b/flytekit/bin/entrypoint.py index bfb4fd860b..92f56409ec 100644 --- a/flytekit/bin/entrypoint.py +++ b/flytekit/bin/entrypoint.py @@ -249,7 +249,7 @@ def setup_execution( domain=exe_domain, name=exe_name, ), - execution_date=_datetime.datetime.utcnow(), + execution_date=_datetime.datetime.now(_datetime.timezone.utc), stats=_get_stats( cfg=StatsConfig.auto(), # Stats metric path will be: diff --git a/flytekit/core/context_manager.py b/flytekit/core/context_manager.py index 5c6341712a..f70f10bc94 100644 --- a/flytekit/core/context_manager.py +++ b/flytekit/core/context_manager.py @@ -937,7 +937,7 @@ def initialize(): default_user_space_params = ExecutionParameters( execution_id=WorkflowExecutionIdentifier.promote_from_model(default_execution_id), task_id=_identifier.Identifier(_identifier.ResourceType.TASK, "local", "local", "local", "local"), - execution_date=_datetime.datetime.utcnow(), + execution_date=_datetime.datetime.now(_datetime.timezone.utc), stats=mock_stats.MockStats(), logging=user_space_logger, tmp_dir=user_space_path, diff --git a/flytekit/core/mock_stats.py b/flytekit/core/mock_stats.py index dedce6ae7c..18763fa74a 100644 --- a/flytekit/core/mock_stats.py +++ b/flytekit/core/mock_stats.py @@ -57,8 +57,10 @@ def __init__(self, mock_stats, metric, tags): self._tags = tags def __enter__(self): - self._timer = _datetime.datetime.utcnow() + self._timer = _datetime.datetime.now(_datetime.timezone.utc) def __exit__(self, exc_type, exc_val, exc_tb): - self._mock_stats.gauge(self._metric, _datetime.datetime.utcnow() - self._timer, tags=self._tags) + self._mock_stats.gauge( + self._metric, _datetime.datetime.now(_datetime.timezone.utc) - self._timer, tags=self._tags + ) self._timer = None diff --git a/flytekit/core/utils.py b/flytekit/core/utils.py index 17cdfb3de9..b5a415d13d 100644 --- a/flytekit/core/utils.py +++ b/flytekit/core/utils.py @@ -311,7 +311,7 @@ def wrapper(*args, **kwargs): return wrapper def __enter__(self): - self.start_time = datetime.datetime.utcnow() + self.start_time = datetime.datetime.now(datetime.timezone.utc) self._start_wall_time = _time.perf_counter() self._start_process_time = _time.process_time() return self @@ -323,7 +323,7 @@ def __exit__(self, exc_type, exc_val, exc_tb): """ from flytekit.core.context_manager import FlyteContextManager - end_time = datetime.datetime.utcnow() + end_time = datetime.datetime.now(datetime.timezone.utc) end_wall_time = _time.perf_counter() end_process_time = _time.process_time() diff --git a/flytekit/experimental/eager_function.py b/flytekit/experimental/eager_function.py index 7a3cee897d..d47a2baef2 100644 --- a/flytekit/experimental/eager_function.py +++ b/flytekit/experimental/eager_function.py @@ -2,7 +2,7 @@ import inspect import signal from contextlib import asynccontextmanager -from datetime import datetime, timedelta +from datetime import datetime, timedelta, timezone from functools import partial, wraps from typing import List, Optional @@ -179,9 +179,9 @@ async def __call__(self, **kwargs): self.async_stack.set_node(node) poll_interval = self._poll_interval or timedelta(seconds=30) - time_to_give_up = datetime.max if self._timeout is None else datetime.utcnow() + self._timeout + time_to_give_up = datetime.max if self._timeout is None else datetime.now(timezone.utc) + self._timeout - while datetime.utcnow() < time_to_give_up: + while datetime.now(timezone.utc) < time_to_give_up: execution = self.remote.sync(execution) if execution.closure.phase in {WorkflowExecutionPhase.FAILED}: raise EagerException(f"Error executing {self.entity.name} with error: {execution.closure.error}") @@ -208,9 +208,9 @@ async def terminate(self): ) poll_interval = self._poll_interval or timedelta(seconds=6) - time_to_give_up = datetime.max if self._timeout is None else datetime.utcnow() + self._timeout + time_to_give_up = datetime.max if self._timeout is None else datetime.now(timezone.utc) + self._timeout - while datetime.utcnow() < time_to_give_up: + while datetime.now(timezone.utc) < time_to_give_up: execution = self.remote.sync(execution) if execution.is_done: break diff --git a/flytekit/remote/remote.py b/flytekit/remote/remote.py index 992685abb8..4fd17fe40b 100644 --- a/flytekit/remote/remote.py +++ b/flytekit/remote/remote.py @@ -1661,9 +1661,9 @@ def wait( :param sync_nodes: passed along to the sync call for the workflow execution """ poll_interval = poll_interval or timedelta(seconds=30) - time_to_give_up = datetime.max if timeout is None else datetime.utcnow() + timeout + time_to_give_up = datetime.max if timeout is None else datetime.now() + timeout - while datetime.utcnow() < time_to_give_up: + while datetime.now() < time_to_give_up: execution = self.sync_execution(execution, sync_nodes=sync_nodes) if execution.is_done: return execution diff --git a/plugins/flytekit-deck-standard/tests/test_renderer.py b/plugins/flytekit-deck-standard/tests/test_renderer.py index c114333d58..dbe157cefc 100644 --- a/plugins/flytekit-deck-standard/tests/test_renderer.py +++ b/plugins/flytekit-deck-standard/tests/test_renderer.py @@ -22,8 +22,8 @@ [ dict( Name="foo", - Start=datetime.datetime.utcnow(), - Finish=datetime.datetime.utcnow() + datetime.timedelta(microseconds=1000), + Start=datetime.datetime.now(datetime.timezone.utc), + Finish=datetime.datetime.now(datetime.timezone.utc) + datetime.timedelta(microseconds=1000), WallTime=1.0, ProcessTime=1.0, ) diff --git a/plugins/flytekit-greatexpectations/flytekitplugins/great_expectations/schema.py b/plugins/flytekit-greatexpectations/flytekitplugins/great_expectations/schema.py index 3413cdcdd3..4cd0d5a3e9 100644 --- a/plugins/flytekit-greatexpectations/flytekitplugins/great_expectations/schema.py +++ b/plugins/flytekit-greatexpectations/flytekitplugins/great_expectations/schema.py @@ -290,7 +290,7 @@ def to_python_value( run_id = ge.core.run_identifier.RunIdentifier( **{ "run_name": ge_conf.datasource_name + "_run", - "run_time": datetime.datetime.utcnow(), + "run_time": datetime.datetime.now(datetime.timezone.utc), } ) diff --git a/plugins/flytekit-greatexpectations/flytekitplugins/great_expectations/task.py b/plugins/flytekit-greatexpectations/flytekitplugins/great_expectations/task.py index bd04f18782..a39baacf34 100644 --- a/plugins/flytekit-greatexpectations/flytekitplugins/great_expectations/task.py +++ b/plugins/flytekit-greatexpectations/flytekitplugins/great_expectations/task.py @@ -216,7 +216,7 @@ def execute(self, **kwargs) -> Any: run_id = ge.core.run_identifier.RunIdentifier( **{ "run_name": self._datasource_name + "_run", - "run_time": datetime.datetime.utcnow(), + "run_time": datetime.datetime.now(datetime.timezone.utc), } ) diff --git a/plugins/flytekit-identity-aware-proxy/tests/test_flytekitplugins_iap.py b/plugins/flytekit-identity-aware-proxy/tests/test_flytekitplugins_iap.py index 766ff646ab..9b0136331d 100644 --- a/plugins/flytekit-identity-aware-proxy/tests/test_flytekitplugins_iap.py +++ b/plugins/flytekit-identity-aware-proxy/tests/test_flytekitplugins_iap.py @@ -1,5 +1,5 @@ import uuid -from datetime import datetime, timedelta +from datetime import datetime, timedelta, timezone from unittest.mock import MagicMock, patch import click @@ -58,7 +58,7 @@ def test_get_gcp_secret_manager_secret_not_found(): def create_mock_token(aud: str, expires_in: timedelta = None): """Create a mock JWT token with a certain audience, expiration time, and random JTI.""" - exp = datetime.utcnow() + expires_in + exp = datetime.now(timezone.utc) + expires_in jti = "test_token" + str(uuid.uuid4()) payload = {"exp": exp, "aud": aud, "jti": jti} diff --git a/tests/flytekit/integration/remote/test_remote.py b/tests/flytekit/integration/remote/test_remote.py index a1c137dc48..3398c771f3 100644 --- a/tests/flytekit/integration/remote/test_remote.py +++ b/tests/flytekit/integration/remote/test_remote.py @@ -81,10 +81,10 @@ def test_monitor_workflow_execution(register): ) poll_interval = datetime.timedelta(seconds=1) - time_to_give_up = datetime.datetime.utcnow() + datetime.timedelta(seconds=60) + time_to_give_up = datetime.datetime.now(datetime.timezone.utc) + datetime.timedelta(seconds=60) execution = remote.sync_execution(execution, sync_nodes=True) - while datetime.datetime.utcnow() < time_to_give_up: + while datetime.datetime.now(datetime.timezone.utc) < time_to_give_up: if execution.is_done: break diff --git a/tests/flytekit/unit/core/test_context_manager.py b/tests/flytekit/unit/core/test_context_manager.py index ad27ad1852..6d519e2da6 100644 --- a/tests/flytekit/unit/core/test_context_manager.py +++ b/tests/flytekit/unit/core/test_context_manager.py @@ -1,6 +1,6 @@ import base64 import os -from datetime import datetime +from datetime import datetime, timezone from pathlib import Path from unittest.mock import Mock @@ -251,7 +251,7 @@ def test_exec_params(): ep = ExecutionParameters( execution_id=id_models.WorkflowExecutionIdentifier("p", "d", "n"), task_id=id_models.Identifier(id_models.ResourceType.TASK, "local", "local", "local", "local"), - execution_date=datetime.utcnow(), + execution_date=datetime.now(timezone.utc), stats=mock_stats.MockStats(), logging=None, tmp_dir="/tmp", diff --git a/tests/flytekit/unit/deck/test_deck.py b/tests/flytekit/unit/deck/test_deck.py index da23ed188b..45056ae283 100644 --- a/tests/flytekit/unit/deck/test_deck.py +++ b/tests/flytekit/unit/deck/test_deck.py @@ -32,8 +32,8 @@ def test_deck(): def test_timeline_deck(): time_info = dict( Name="foo", - Start=datetime.datetime.utcnow(), - Finish=datetime.datetime.utcnow() + datetime.timedelta(microseconds=1000), + Start=datetime.datetime.now(datetime.timezone.utc), + Finish=datetime.datetime.now(datetime.timezone.utc) + datetime.timedelta(microseconds=1000), WallTime=1.0, ProcessTime=1.0, ) diff --git a/tests/flytekit/unit/models/test_literals.py b/tests/flytekit/unit/models/test_literals.py index 0035f8ec1d..d32628d036 100644 --- a/tests/flytekit/unit/models/test_literals.py +++ b/tests/flytekit/unit/models/test_literals.py @@ -103,7 +103,7 @@ def test_boolean_primitive(): def test_datetime_primitive(): - dt = datetime.utcnow().replace(tzinfo=timezone.utc) + dt = datetime.now(timezone.utc) obj = literals.Primitive(datetime=dt) assert obj.integer is None assert obj.boolean is None