Replace freezegun with time-machine

pyproject.toml

@@ -193,7 +193,7 @@ test = [
     "pytest-unordered>=0.7.0",
     "pytest-xdist>=3.6.1",
     "requests-mock>=1.12.1",
-    "freezegun>=1.5.1",
+    "time-machine>=3.5.1",
     "httpbin>=0.10.2", # indirect to make compatible with Werkzeug 3
     "openai>=1.68.2",
     "testcontainers>=4.15.0",

src/dstack/_internal/proxy/gateway/services/registry.py

@@ -1,6 +1,5 @@
 import asyncio
 from asyncio import Lock
-from datetime import datetime
 from pathlib import Path
 from typing import Iterable, Optional
 
@@ -26,6 +25,7 @@
     ServiceConnection,
     ServiceConnectionPool,
 )
+from dstack._internal.utils.common import get_current_datetime
 from dstack._internal.utils.logging import get_logger
 
 ACCESS_LOG_PATH = Path("/var/log/nginx/dstack.access.log")
@@ -97,7 +97,7 @@ async def register_service(
                 models.ChatModel(
                     project_name=project_name,
                     name=model.name,
-                    created_at=datetime.now(),
+                    created_at=get_current_datetime(),
                     run_name=run_name,
                     format_spec=model_schema_to_format_spec(model),
                 ),

src/tests/_internal/cli/utils/test_run.py

@@ -4,7 +4,7 @@
 from unittest.mock import Mock
 
 import pytest
-from freezegun import freeze_time
+import time_machine
 from rich.table import Table
 from rich.text import Text
 from sqlalchemy import select
@@ -206,7 +206,7 @@ async def create_run_with_job(
 @pytest.mark.usefixtures("test_db", "image_config_mock")
 @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
 class TestGetRunsTable:
-    @freeze_time(datetime(2026, 1, 1))
+    @time_machine.travel(datetime(2026, 1, 1), tick=False)
     async def test_simple_run(self, session: AsyncSession):
         api_run = await create_run_with_job(
             session=session, submitted_at=datetime(2023, 1, 1, tzinfo=timezone.utc)

src/tests/_internal/proxy/gateway/routers/test_registry.py

@@ -1,12 +1,12 @@
 import re
 import uuid
-from datetime import datetime
+from datetime import datetime, timezone
 from pathlib import Path
 from typing import Optional
 
 import httpx
 import pytest
-from freezegun import freeze_time
+import time_machine
 
 from dstack._internal.core.errors import SSHError
 from dstack._internal.proxy.gateway.app import make_app
@@ -209,7 +209,7 @@ async def test_register_same_domain_error(self, tmp_path: Path, system_mocks: Mo
         assert (tmp_path / "sites-enabled" / "443-test-run.gtw.test.conf").exists()
         assert system_mocks.reload_nginx.call_count == 1
 
-    @freeze_time(datetime(2024, 12, 12, 0, 30))
+    @time_machine.travel(datetime(2024, 12, 12, 0, 30, tzinfo=timezone.utc), tick=False)
     async def test_register_with_model(self, tmp_path: Path, system_mocks: Mocks) -> None:
         repo = GatewayProxyRepo()
         client = make_client(tmp_path, repo=repo)
@@ -225,7 +225,7 @@ async def test_register_with_model(self, tmp_path: Path, system_mocks: Mocks) ->
             ChatModel(
                 project_name="test-proj",
                 name="test-model",
-                created_at=datetime(2024, 12, 12, 0, 30),
+                created_at=datetime(2024, 12, 12, 0, 30, tzinfo=timezone.utc),
                 run_name="test-run",
                 format_spec=OpenAIChatModelFormat(prefix="/v1"),
             )

src/tests/_internal/proxy/gateway/services/test_stats.py

@@ -3,14 +3,14 @@
 from textwrap import dedent
 
 import pytest
-from freezegun import freeze_time
+import time_machine
 
 from dstack._internal.proxy.gateway.schemas.stats import Stat
 from dstack._internal.proxy.gateway.services.stats import StatsCollector
 
 
 @pytest.mark.asyncio
-@freeze_time(datetime(2024, 12, 6, 12, 10, tzinfo=timezone.utc))
+@time_machine.travel(datetime(2024, 12, 6, 12, 10, tzinfo=timezone.utc), tick=False)
 @pytest.mark.parametrize(
     ("access_log", "expected_result"),
     [
@@ -128,7 +128,7 @@ async def test_collect_stats(access_log: str, expected_result: dict, tmp_path: P
 
 
 @pytest.mark.asyncio
-@freeze_time(datetime(2024, 12, 6, 12, 10, tzinfo=timezone.utc))
+@time_machine.travel(datetime(2024, 12, 6, 12, 10, tzinfo=timezone.utc), tick=False)
 async def test_collect_stats_after_log_update(tmp_path: Path) -> None:
     access_log_path = tmp_path / "dstack.access.log"
     collector = StatsCollector(access_log_path)

src/tests/_internal/server/background/pipeline_tasks/test_instances/test_termination.py

@@ -4,7 +4,7 @@
 from unittest.mock import AsyncMock, Mock, patch
 
 import pytest
-from freezegun import freeze_time
+import time_machine
 from sqlalchemy.ext.asyncio import AsyncSession
 
 from dstack._internal.core.errors import BackendError, NotYetTerminated
@@ -126,7 +126,7 @@ async def test_terminate_retry(
         await session.commit()
 
         with (
-            freeze_time(initial_time + dt.timedelta(minutes=1)),
+            time_machine.travel(initial_time + dt.timedelta(minutes=1), tick=False),
             self.mock_terminate_in_backend(error=error) as mock,
         ):
             await process_instance(session, worker, instance)
@@ -135,7 +135,7 @@ async def test_terminate_retry(
         assert instance.status == InstanceStatus.TERMINATING
 
         with (
-            freeze_time(initial_time + dt.timedelta(minutes=2)),
+            time_machine.travel(initial_time + dt.timedelta(minutes=2), tick=False),
             self.mock_terminate_in_backend(error=None) as mock,
         ):
             await process_instance(session, worker, instance)
@@ -162,7 +162,7 @@ async def test_terminate_not_retries_if_too_early(
         await session.commit()
 
         with (
-            freeze_time(initial_time + dt.timedelta(minutes=1)),
+            time_machine.travel(initial_time + dt.timedelta(minutes=1), tick=False),
             self.mock_terminate_in_backend(error=BackendError("err")) as mock,
         ):
             await process_instance(session, worker, instance)
@@ -174,7 +174,7 @@ async def test_terminate_not_retries_if_too_early(
         await session.commit()
 
         with (
-            freeze_time(initial_time + dt.timedelta(minutes=1, seconds=11)),
+            time_machine.travel(initial_time + dt.timedelta(minutes=1, seconds=11), tick=False),
             self.mock_terminate_in_backend(error=None) as mock,
         ):
             await process_instance(session, worker, instance)
@@ -201,7 +201,7 @@ async def test_terminate_on_termination_deadline(
         await session.commit()
 
         with (
-            freeze_time(initial_time + dt.timedelta(minutes=1)),
+            time_machine.travel(initial_time + dt.timedelta(minutes=1), tick=False),
             self.mock_terminate_in_backend(error=BackendError("err")) as mock,
         ):
             await process_instance(session, worker, instance)
@@ -210,7 +210,7 @@ async def test_terminate_on_termination_deadline(
         assert instance.status == InstanceStatus.TERMINATING
 
         with (
-            freeze_time(initial_time + dt.timedelta(minutes=15, seconds=55)),
+            time_machine.travel(initial_time + dt.timedelta(minutes=15, seconds=55), tick=False),
             self.mock_terminate_in_backend(error=None) as mock,
         ):
             await process_instance(session, worker, instance)

src/tests/_internal/server/background/pipeline_tasks/test_running_jobs.py

@@ -8,7 +8,7 @@
 from unittest.mock import AsyncMock, MagicMock, Mock, patch
 
 import pytest
-from freezegun import freeze_time
+import time_machine
 from sqlalchemy import select
 from sqlalchemy.ext.asyncio import AsyncSession
 from sqlalchemy.orm import selectinload
@@ -1281,7 +1281,7 @@ async def test_pulling_shim_failed(
 
         with (
             patch("dstack._internal.server.services.runner.pool.SSHTunnel") as ssh_tunnel_cls,
-            freeze_time(job.disconnected_at + timedelta(minutes=5)),
+            time_machine.travel(job.disconnected_at + timedelta(minutes=5), tick=False),
         ):
             from dstack._internal.core.errors import SSHError
 
@@ -1622,7 +1622,7 @@ async def test_running_job_disconnect_retries_then_terminates(
 
         with (
             patch("dstack._internal.server.services.runner.pool.SSHTunnel") as ssh_tunnel_cls,
-            freeze_time(job.disconnected_at + timedelta(minutes=5)),
+            time_machine.travel(job.disconnected_at + timedelta(minutes=5), tick=False),
         ):
             ssh_tunnel_cls.side_effect = SSHError
             await _process_job(session, worker, job)
@@ -1858,7 +1858,7 @@ async def test_max_duration_enforced_by_server(
             patch(
                 "dstack._internal.server.services.runner.client.RunnerClient.from_address"
             ) as runner_client_cls,
-            freeze_time(now),
+            time_machine.travel(now, tick=False),
         ):
             runner_client_mock = runner_client_cls.return_value
             runner_client_mock.pull.return_value = PullResponse(
@@ -1909,7 +1909,7 @@ async def test_max_duration_enforced_by_server(
             ),
         ],
     )
-    @freeze_time(datetime(2023, 1, 1, 12, 30, tzinfo=timezone.utc))
+    @time_machine.travel(datetime(2023, 1, 1, 12, 30, tzinfo=timezone.utc), tick=False)
     async def test_gpu_utilization(
         self,
         test_db,
@@ -3148,7 +3148,7 @@ async def test_router_failed_terminates_worker(self):
             result.job_update_map.get("termination_reason_message") or ""
         )
 
-    @freeze_time("2023-01-01 12:00:00Z")
+    @time_machine.travel(datetime(2023, 1, 1, 12, tzinfo=timezone.utc), tick=False)
     async def test_router_not_provisioned_within_timeout_defers(self):
         context = self._make_context(
             submitted_at=datetime(2023, 1, 1, 11, 45, 0, tzinfo=timezone.utc),
@@ -3162,7 +3162,7 @@ async def test_router_not_provisioned_within_timeout_defers(self):
         assert out is None
         assert result.job_update_map == {}
 
-    @freeze_time("2023-01-01 12:00:00Z")
+    @time_machine.travel(datetime(2023, 1, 1, 12, tzinfo=timezone.utc), tick=False)
     async def test_router_not_provisioned_past_timeout_terminates(self):
         context = self._make_context(
             submitted_at=datetime(2023, 1, 1, 10, 0, 0, tzinfo=timezone.utc),
@@ -3262,7 +3262,7 @@ def _make_context(self) -> _ProcessContext:
             instance_access_revoked=False,
         )
 
-    @freeze_time("2023-01-01 12:00:00Z")
+    @time_machine.travel(datetime(2023, 1, 1, 12, tzinfo=timezone.utc), tick=False)
     async def test_submitted_sibling_defers(self):
         context = self._make_context()
         result = _ProcessResult()
@@ -3345,7 +3345,7 @@ async def _fake_session_ctx():
             ),
         )
 
-    @freeze_time("2023-01-01 12:00:00Z")
+    @time_machine.travel(datetime(2023, 1, 1, 12, tzinfo=timezone.utc), tick=False)
     async def test_groups_ip_not_ready_defers(self):
         context = self._make_context()
         result = _ProcessResult()
@@ -3556,7 +3556,7 @@ async def _fake_session_ctx():
             ),
         )
 
-    @freeze_time("2023-01-01 12:00:00Z")
+    @time_machine.travel(datetime(2023, 1, 1, 12, tzinfo=timezone.utc), tick=False)
     async def test_replica_ip_not_ready_defers(self):
         context = self._make_context(router_ip="")
         result = _ProcessResult()
@@ -3568,7 +3568,7 @@ async def test_replica_ip_not_ready_defers(self):
         assert result.job_update_map == {}
         assert context.job.job_spec.commands == ["echo ${{ groups[0].replicas[0].IP_ADDRESS }}"]
 
-    @freeze_time("2023-01-01 12:00:00Z")
+    @time_machine.travel(datetime(2023, 1, 1, 12, tzinfo=timezone.utc), tick=False)
     async def test_replica_ip_ready_substitutes(self):
         context = self._make_context(router_ip="10.0.0.5")
         result = _ProcessResult()
@@ -3579,7 +3579,7 @@ async def test_replica_ip_ready_substitutes(self):
         assert out is not None
         assert context.job.job_spec.commands == ["echo 10.0.0.5"]
 
-    @freeze_time("2023-01-01 12:00:00Z")
+    @time_machine.travel(datetime(2023, 1, 1, 12, tzinfo=timezone.utc), tick=False)
     async def test_nodes_ref_in_service_terminates(self):
         context = self._make_context(
             router_ip="10.0.0.5",

src/tests/_internal/server/background/pipeline_tasks/test_runs/test_termination.py

@@ -4,7 +4,7 @@
 from unittest.mock import AsyncMock, patch
 
 import pytest
-from freezegun import freeze_time
+import time_machine
 from sqlalchemy.ext.asyncio import AsyncSession
 
 from dstack._internal.core.models.configurations import TaskConfiguration
@@ -187,7 +187,7 @@ async def test_finishes_non_scheduled_run_when_all_jobs_are_finished(
         assert run.lock_expires_at is None
         assert run.lock_owner is None
 
-    @freeze_time(datetime(2023, 1, 2, 3, 10, tzinfo=timezone.utc))
+    @time_machine.travel(datetime(2023, 1, 2, 3, 10, tzinfo=timezone.utc), tick=False)
     async def test_reschedules_scheduled_run_and_clears_fleet(
         self, test_db, session: AsyncSession, worker: RunWorker
     ) -> None:

src/tests/_internal/server/background/scheduled_tasks/test_events.py

@@ -2,7 +2,7 @@
 from unittest.mock import patch
 
 import pytest
-from freezegun import freeze_time
+import time_machine
 from sqlalchemy.ext.asyncio import AsyncSession
 
 from dstack._internal.server import settings
@@ -16,7 +16,7 @@
 async def test_deletes_old_events(test_db, session: AsyncSession) -> None:
     user = await create_user(session=session)
     for i in range(10):
-        with freeze_time(datetime(2026, 1, 1, i)):
+        with time_machine.travel(datetime(2026, 1, 1, i), tick=False):
             events.emit(
                 session,
                 message=f"Event {i}",
@@ -30,7 +30,7 @@ async def test_deletes_old_events(test_db, session: AsyncSession) -> None:
 
     with (
         patch.multiple(settings, SERVER_EVENTS_TTL_SECONDS=5 * 3600),
-        freeze_time(datetime(2026, 1, 1, 10)),
+        time_machine.travel(datetime(2026, 1, 1, 10), tick=False),
     ):
         await delete_events()
 

src/tests/_internal/server/background/scheduled_tasks/test_metrics.py

@@ -2,7 +2,7 @@
 from unittest.mock import patch
 
 import pytest
-from freezegun import freeze_time
+import time_machine
 from sqlalchemy import select
 from sqlalchemy.ext.asyncio import AsyncSession
 
@@ -93,7 +93,7 @@ async def test_collects_metrics(self, test_db, session: AsyncSession):
 class TestDeleteMetrics:
     @pytest.mark.asyncio
     @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
-    @freeze_time(datetime(2023, 1, 2, 3, 5, 20, tzinfo=timezone.utc))
+    @time_machine.travel(datetime(2023, 1, 2, 3, 5, 20, tzinfo=timezone.utc), tick=False)
     async def test_deletes_old_metrics_running_job(self, test_db, session: AsyncSession):
         user = await create_user(session=session, global_role=GlobalRole.USER)
         project = await create_project(session=session, owner=user)
@@ -141,7 +141,7 @@ async def test_deletes_old_metrics_running_job(self, test_db, session: AsyncSess
 
     @pytest.mark.asyncio
     @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
-    @freeze_time(datetime(2023, 1, 2, 3, 5, 20, tzinfo=timezone.utc))
+    @time_machine.travel(datetime(2023, 1, 2, 3, 5, 20, tzinfo=timezone.utc), tick=False)
     async def test_deletes_old_metrics_finished_job(self, test_db, session: AsyncSession):
         user = await create_user(session=session, global_role=GlobalRole.USER)
         project = await create_project(session=session, owner=user)

src/tests/_internal/server/background/scheduled_tasks/test_probes.py

@@ -2,7 +2,7 @@
 from unittest.mock import patch
 
 import pytest
-from freezegun import freeze_time
+import time_machine
 from sqlalchemy.ext.asyncio import AsyncSession
 
 from dstack._internal.core.models.configurations import ProbeConfig, ServiceConfiguration
@@ -138,7 +138,7 @@ async def test_schedules_probe_execution(self, test_db, session: AsyncSession) -
             session, job, probe_num=1, due=datetime(2025, 1, 1, 0, 0, 0, tzinfo=timezone.utc)
         )
         processing_time = datetime(2025, 1, 1, 0, 0, 1, tzinfo=timezone.utc)
-        with freeze_time(processing_time):
+        with time_machine.travel(processing_time, tick=False):
             with patch(
                 "dstack._internal.server.background.scheduled_tasks.probes.PROBES_SCHEDULER"
             ) as scheduler_mock: