Apply PR comments

Author: Nadine-H

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

@@ -20,7 +20,6 @@
     RunStatus,
     RunTerminationReason,
 )
-from dstack._internal.server.background.metrics import run_metrics
 from dstack._internal.server.db import get_session_ctx
 from dstack._internal.server.models import JobModel, ProjectModel, RunModel
 from dstack._internal.server.services.jobs import (
@@ -29,6 +28,7 @@
     group_jobs_by_replica_latest,
 )
 from dstack._internal.server.services.locking import get_locker
+from dstack._internal.server.services.prometheus.push_metrics import run_metrics
 from dstack._internal.server.services.runs import (
     create_job_model_for_new_submission,
     fmt,
@@ -371,7 +371,7 @@ async def _process_active_run(session: AsyncSession, run_model: RunModel):
             new_status.name,
         )
         if run_model.status == RunStatus.SUBMITTED and new_status == RunStatus.PROVISIONING:
-            current_time = datetime.datetime.now(datetime.timezone.utc)
+            current_time = common.get_current_datetime()
             submit_to_provision_duration = (
                 current_time - run_model.submitted_at.replace(tzinfo=datetime.timezone.utc)
             ).total_seconds()
@@ -381,13 +381,11 @@ async def _process_active_run(session: AsyncSession, run_model: RunModel):
                 submit_to_provision_duration,
             )
             project_name = run_model.project.name
-            run_type = RunSpec.__response__.parse_raw(run_model.run_spec).configuration.type
             run_metrics.log_submit_to_provision_duration(
-                submit_to_provision_duration, project_name, run_type
+                submit_to_provision_duration, project_name, run_spec.configuration.type
             )
 
         if new_status == RunStatus.PENDING:
-            run_spec = RunSpec.__response__.parse_raw(run_model.run_spec)
             run_metrics.increment_pending_runs(run_model.project.name, run_spec.configuration.type)
 
         run_model.status = new_status

src/dstack/_internal/server/routers/prometheus.py

@@ -9,7 +9,7 @@
 from dstack._internal.server import settings
 from dstack._internal.server.db import get_session
 from dstack._internal.server.security.permissions import OptionalServiceAccount
-from dstack._internal.server.services import prometheus
+from dstack._internal.server.services.prometheus import pull_metrics
 from dstack._internal.server.utils.routers import error_not_found
 
 _auth = OptionalServiceAccount(os.getenv("DSTACK_PROMETHEUS_AUTH_TOKEN"))
@@ -27,6 +27,6 @@ async def get_prometheus_metrics(
 ) -> str:
     if not settings.ENABLE_PROMETHEUS_METRICS:
         raise error_not_found()
-    custom_metrics = await prometheus.get_metrics(session=session)
+    custom_metrics = await pull_metrics.get_metrics(session=session)
     prometheus_metrics = generate_latest()
     return custom_metrics + prometheus_metrics.decode()

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

@@ -28,6 +28,7 @@
     get_job_provisioning_data,
     get_run_spec,
 )
+from dstack._internal.utils import common
 
 pytestmark = pytest.mark.usefixtures("image_config_mock")
 
@@ -72,10 +73,14 @@ async def make_run(
 class TestProcessRuns:
     @pytest.mark.asyncio
     @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
+    @patch(
+        "dstack._internal.utils.common.get_current_datetime",
+        return_value=datetime.datetime(2024, 1, 1, 12, 0, 0, tzinfo=datetime.timezone.utc),
+    )
     async def test_submitted_to_provisioning(self, test_db, session: AsyncSession):
         run = await make_run(session, status=RunStatus.SUBMITTED)
         await create_job(session=session, run=run, status=JobStatus.PROVISIONING)
-        current_time = datetime.datetime.now(datetime.timezone.utc)
+        current_time = common.get_current_datetime()
 
         expected_duration = (
             current_time - run.submitted_at.replace(tzinfo=datetime.timezone.utc)
@@ -91,9 +96,7 @@ async def test_submitted_to_provisioning(self, test_db, session: AsyncSession):
             assert args[1] == run.project.name
             assert args[2] == "service"
             # Assert the duration is close to our expected duration (within 0.05 second tolerance)
-            assert abs(args[0] - expected_duration) < 0.05, (
-                f"Expected duration ~{expected_duration:.3f}s, got {args[0]:.3f}s"
-            )
+            assert args[0] == expected_duration
 
         await session.refresh(run)
         assert run.status == RunStatus.PROVISIONING
@@ -233,11 +236,15 @@ async def test_pending_to_submitted(self, test_db, session: AsyncSession):
 class TestProcessRunsReplicas:
     @pytest.mark.asyncio
     @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
+    @patch(
+        "dstack._internal.utils.common.get_current_datetime",
+        return_value=datetime.datetime(2024, 1, 1, 12, 0, 0, tzinfo=datetime.timezone.utc),
+    )
     async def test_submitted_to_provisioning_if_any(self, test_db, session: AsyncSession):
         run = await make_run(session, status=RunStatus.SUBMITTED, replicas=2)
         await create_job(session=session, run=run, status=JobStatus.SUBMITTED, replica_num=0)
         await create_job(session=session, run=run, status=JobStatus.PROVISIONING, replica_num=1)
-        current_time = datetime.datetime.now(datetime.timezone.utc)
+        current_time = common.get_current_datetime()
 
         expected_duration = (
             current_time - run.submitted_at.replace(tzinfo=datetime.timezone.utc)
@@ -253,10 +260,7 @@ async def test_submitted_to_provisioning_if_any(self, test_db, session: AsyncSes
             assert args[1] == run.project.name
             assert args[2] == "service"
             assert isinstance(args[0], float)
-            # Assert the duration is close to our expected duration (within 0.05 second tolerance)
-            assert abs(args[0] - expected_duration) < 0.05, (
-                f"Expected duration ~{expected_duration:.3f}s, got {args[0]:.3f}s"
-            )
+            assert args[0] == expected_duration
 
         await session.refresh(run)
         assert run.status == RunStatus.PROVISIONING

…/_internal/server/services/prometheus.py …rver/services/prometheus/pull_metrics.pysrc/dstack/_internal/server/services/prometheus.py renamed to src/dstack/_internal/server/services/prometheus/pull_metrics.py src/dstack/_internal/server/services/prometheus.py renamed to src/dstack/_internal/server/services/prometheus/pull_metrics.py

@@ -1,6 +1,6 @@
 from unittest.mock import MagicMock
 
-from dstack._internal.server.background.metrics import run_metrics
+from dstack._internal.server.services.prometheus.push_metrics import run_metrics
 
 
 class TestRunMetrics: