Autoset UTC timezone for datetimes loaded from the db (#2922)

* Load datetimes with utc timezone from db

* Use tz aware datetimes for sqlalchemy filters

Author: r4victor

src/dstack/_internal/server/background/tasks/process_fleets.py

@@ -40,7 +40,7 @@ async def _process_next_fleet():
                     FleetModel.deleted == False,
                     FleetModel.id.not_in(lockset),
                     FleetModel.last_processed_at
-                    < get_current_datetime().replace(tzinfo=None) - MIN_PROCESSING_INTERVAL,
+                    < get_current_datetime() - MIN_PROCESSING_INTERVAL,
                 )
                 .order_by(FleetModel.last_processed_at.asc())
                 .limit(1)

src/dstack/_internal/server/background/tasks/process_idle_volumes.py

@@ -82,8 +82,7 @@ def _should_delete_volume(volume: VolumeModel) -> bool:
 
 def _get_idle_time(volume: VolumeModel) -> datetime.timedelta:
     last_used = volume.last_job_processed_at or volume.created_at
-    last_used_utc = last_used.replace(tzinfo=datetime.timezone.utc)
-    idle_time = get_current_datetime() - last_used_utc
+    idle_time = get_current_datetime() - last_used
     return max(idle_time, datetime.timedelta(0))
 
 

src/dstack/_internal/server/background/tasks/process_instances.py

@@ -104,7 +104,10 @@
 from dstack._internal.server.services.runner import client as runner_client
 from dstack._internal.server.services.runner.client import HealthStatus
 from dstack._internal.server.services.runner.ssh import runner_ssh_tunnel
-from dstack._internal.utils.common import get_current_datetime, run_async
+from dstack._internal.utils.common import (
+    get_current_datetime,
+    run_async,
+)
 from dstack._internal.utils.logging import get_logger
 from dstack._internal.utils.network import get_ip_from_network, is_ip_among_addresses
 from dstack._internal.utils.ssh import (
@@ -149,7 +152,7 @@ async def _process_next_instance():
                     ),
                     InstanceModel.id.not_in(lockset),
                     InstanceModel.last_processed_at
-                    < get_current_datetime().replace(tzinfo=None) - MIN_PROCESSING_INTERVAL,
+                    < get_current_datetime() - MIN_PROCESSING_INTERVAL,
                 )
                 .options(lazyload(InstanceModel.jobs))
                 .order_by(InstanceModel.last_processed_at.asc())
@@ -461,7 +464,7 @@ def _deploy_instance(
 
 async def _create_instance(session: AsyncSession, instance: InstanceModel) -> None:
     if instance.last_retry_at is not None:
-        last_retry = instance.last_retry_at.replace(tzinfo=datetime.timezone.utc)
+        last_retry = instance.last_retry_at
         if get_current_datetime() < last_retry + timedelta(minutes=1):
             return
 
@@ -801,7 +804,7 @@ async def _check_instance(instance: InstanceModel) -> None:
             instance.name,
             extra={"instance_name": instance.name},
         )
-        deadline = instance.termination_deadline.replace(tzinfo=datetime.timezone.utc)
+        deadline = instance.termination_deadline
         if get_current_datetime() > deadline:
             instance.status = InstanceStatus.TERMINATING
             instance.termination_reason = "Termination deadline"
@@ -956,18 +959,12 @@ async def _terminate(instance: InstanceModel) -> None:
 
 def _next_termination_retry_at(instance: InstanceModel) -> datetime.datetime:
     assert instance.last_termination_retry_at is not None
-    return (
-        instance.last_termination_retry_at.replace(tzinfo=datetime.timezone.utc)
-        + TERMINATION_RETRY_TIMEOUT
-    )
+    return instance.last_termination_retry_at + TERMINATION_RETRY_TIMEOUT
 
 
 def _get_termination_deadline(instance: InstanceModel) -> datetime.datetime:
     assert instance.first_termination_retry_at is not None
-    return (
-        instance.first_termination_retry_at.replace(tzinfo=datetime.timezone.utc)
-        + TERMINATION_RETRY_MAX_DURATION
-    )
+    return instance.first_termination_retry_at + TERMINATION_RETRY_MAX_DURATION
 
 
 def _need_to_wait_fleet_provisioning(instance: InstanceModel) -> bool:
@@ -1102,27 +1099,26 @@ async def _create_placement_group(
 
 
 def _get_instance_idle_duration(instance: InstanceModel) -> datetime.timedelta:
-    last_time = instance.created_at.replace(tzinfo=datetime.timezone.utc)
+    last_time = instance.created_at
     if instance.last_job_processed_at is not None:
-        last_time = instance.last_job_processed_at.replace(tzinfo=datetime.timezone.utc)
+        last_time = instance.last_job_processed_at
     return get_current_datetime() - last_time
 
 
 def _get_retry_duration_deadline(instance: InstanceModel, retry: Retry) -> datetime.datetime:
-    return instance.created_at.replace(tzinfo=datetime.timezone.utc) + timedelta(
-        seconds=retry.duration
-    )
+    return instance.created_at + timedelta(seconds=retry.duration)
 
 
 def _get_provisioning_deadline(
     instance: InstanceModel,
     job_provisioning_data: JobProvisioningData,
 ) -> datetime.datetime:
+    assert instance.started_at is not None
     timeout_interval = get_provisioning_timeout(
         backend_type=job_provisioning_data.get_base_backend(),
         instance_type_name=job_provisioning_data.instance_type.name,
     )
-    return instance.started_at.replace(tzinfo=datetime.timezone.utc) + timeout_interval
+    return instance.started_at + timeout_interval
 
 
 def _ssh_keys_to_pkeys(ssh_keys: list[SSHKey]) -> list[PKey]:

src/dstack/_internal/server/background/tasks/process_running_jobs.py

@@ -2,7 +2,7 @@
 import re
 import uuid
 from collections.abc import Iterable
-from datetime import timedelta, timezone
+from datetime import timedelta
 from typing import Dict, List, Optional
 
 from sqlalchemy import select
@@ -108,8 +108,7 @@ async def _process_next_running_job():
                     RunModel.status.not_in([RunStatus.TERMINATING]),
                     JobModel.id.not_in(lockset),
                     JobModel.last_processed_at
-                    < common_utils.get_current_datetime().replace(tzinfo=None)
-                    - MIN_PROCESSING_INTERVAL,
+                    < common_utils.get_current_datetime() - MIN_PROCESSING_INTERVAL,
                 )
                 .order_by(JobModel.last_processed_at.asc())
                 .limit(1)
@@ -801,7 +800,7 @@ def _should_terminate_job_due_to_disconnect(job_model: JobModel) -> bool:
         return False
     return (
         common_utils.get_current_datetime()
-        > job_model.disconnected_at.replace(tzinfo=timezone.utc) + JOB_DISCONNECTED_RETRY_TIMEOUT
+        > job_model.disconnected_at + JOB_DISCONNECTED_RETRY_TIMEOUT
     )
 
 

src/dstack/_internal/server/background/tasks/process_runs.py

@@ -56,7 +56,7 @@ async def process_runs(batch_size: int = 1):
 async def _process_next_run():
     run_lock, run_lockset = get_locker(get_db().dialect_name).get_lockset(RunModel.__tablename__)
     job_lock, job_lockset = get_locker(get_db().dialect_name).get_lockset(JobModel.__tablename__)
-    now = common.get_current_datetime().replace(tzinfo=None)
+    now = common.get_current_datetime()
     async with get_session_ctx() as session:
         async with run_lock, job_lock:
             res = await session.execute(
@@ -370,9 +370,7 @@ async def _process_active_run(session: AsyncSession, run_model: RunModel):
         )
         if run_model.status == RunStatus.SUBMITTED and new_status == RunStatus.PROVISIONING:
             current_time = common.get_current_datetime()
-            submit_to_provision_duration = (
-                current_time - run_model.submitted_at.replace(tzinfo=datetime.timezone.utc)
-            ).total_seconds()
+            submit_to_provision_duration = (current_time - run_model.submitted_at).total_seconds()
             logger.info(
                 "%s: run took %.2f seconds from submission to provisioning.",
                 fmt(run_model),

src/dstack/_internal/server/background/tasks/process_terminating_jobs.py

@@ -18,7 +18,10 @@
 )
 from dstack._internal.server.services.locking import get_locker
 from dstack._internal.server.services.logging import fmt
-from dstack._internal.utils.common import get_current_datetime, get_or_error
+from dstack._internal.utils.common import (
+    get_current_datetime,
+    get_or_error,
+)
 from dstack._internal.utils.logging import get_logger
 
 logger = get_logger(__name__)
@@ -43,7 +46,10 @@ async def _process_next_terminating_job():
                 .where(
                     JobModel.id.not_in(job_lockset),
                     JobModel.status == JobStatus.TERMINATING,
-                    or_(JobModel.remove_at.is_(None), JobModel.remove_at < get_current_datetime()),
+                    or_(
+                        JobModel.remove_at.is_(None),
+                        JobModel.remove_at < get_current_datetime(),
+                    ),
                 )
                 .order_by(JobModel.last_processed_at.asc())
                 .limit(1)

src/dstack/_internal/server/models.py

@@ -1,6 +1,6 @@
 import enum
 import uuid
-from datetime import datetime
+from datetime import datetime, timezone
 from typing import Callable, List, Optional, Union
 
 from sqlalchemy import (
@@ -51,9 +51,10 @@
 
 class NaiveDateTime(TypeDecorator):
     """
-    A custom type decorator that ensures datetime objects are offset-naive when stored in the database.
-    This is needed because we use datetimes in UTC only and store them as offset-naive.
-    Some databases (e.g. Postgres) throw an error if the timezone is set.
+    A custom type decorator that ensures datetime objects are offset-naive when stored in the database
+    and offset-aware with UTC timezone when loaded from the database.
+    This is because we use datetimes in UTC everywhere, and
+    some databases (e.g. Postgres) throw an error if the timezone is set.
     """
 
     impl = DateTime
@@ -65,7 +66,9 @@ def process_bind_param(self, value, dialect):
         return value
 
     def process_result_value(self, value, dialect):
-        return value
+        if value is None:
+            return None
+        return value.replace(tzinfo=timezone.utc)
 
 
 class DecryptedString(CoreModel):

src/dstack/_internal/server/services/fleets.py

@@ -1,6 +1,6 @@
 import uuid
 from collections.abc import Callable
-from datetime import datetime, timezone
+from datetime import datetime
 from functools import wraps
 from typing import List, Literal, Optional, Tuple, TypeVar, Union, cast
 
@@ -600,7 +600,7 @@ def fleet_model_to_fleet(
         name=fleet_model.name,
         project_name=fleet_model.project.name,
         spec=spec,
-        created_at=fleet_model.created_at.replace(tzinfo=timezone.utc),
+        created_at=fleet_model.created_at,
         status=fleet_model.status,
         status_message=fleet_model.status_message,
         instances=instances,

src/dstack/_internal/server/services/gateways/__init__.py

@@ -1,7 +1,7 @@
 import asyncio
 import datetime
 import uuid
-from datetime import timedelta, timezone
+from datetime import timedelta
 from functools import partial
 from typing import List, Optional, Sequence
 
@@ -557,7 +557,7 @@ def gateway_model_to_gateway(gateway_model: GatewayModel) -> Gateway:
         region=gateway_model.region,
         wildcard_domain=gateway_model.wildcard_domain,
         default=gateway_model.project.default_gateway_id == gateway_model.id,
-        created_at=gateway_model.created_at.replace(tzinfo=timezone.utc),
+        created_at=gateway_model.created_at,
         status=gateway_model.status,
         status_message=gateway_model.status_message,
         configuration=configuration,

src/dstack/_internal/server/services/instances.py

@@ -1,6 +1,6 @@
 import uuid
 from collections.abc import Container, Iterable
-from datetime import datetime, timezone
+from datetime import datetime
 from typing import Dict, List, Literal, Optional, Union
 
 import gpuhunt
@@ -62,7 +62,7 @@ def instance_model_to_instance(instance_model: InstanceModel) -> Instance:
         status=instance_model.status,
         unreachable=instance_model.unreachable,
         termination_reason=instance_model.termination_reason,
-        created=instance_model.created_at.replace(tzinfo=timezone.utc),
+        created=instance_model.created_at,
         total_blocks=instance_model.total_blocks,
         busy_blocks=instance_model.busy_blocks,
     )