Replace retention window with TTL

Author: un-def

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

@@ -1,15 +1,14 @@
 import asyncio
 import json
-from datetime import timedelta
 from typing import Dict, List, Optional
 
-from sqlalchemy import BigInteger, delete, func, select
+from sqlalchemy import delete, select
 from sqlalchemy.orm import joinedload
 
 from dstack._internal.core.consts import DSTACK_RUNNER_HTTP_PORT
 from dstack._internal.core.models.runs import JobStatus
 from dstack._internal.server import settings
-from dstack._internal.server.db import get_db, get_session_ctx
+from dstack._internal.server.db import get_session_ctx
 from dstack._internal.server.models import InstanceModel, JobMetricsPoint, JobModel
 from dstack._internal.server.schemas.runner import MetricsResponse
 from dstack._internal.server.services.instances import get_instance_ssh_private_keys
@@ -43,22 +42,13 @@ async def collect_metrics():
 
 
 async def delete_metrics():
-    now = get_current_datetime()
-    running_timestamp_micro_cutoff = int(now.timestamp() * 1_000_000) - (
-        settings.SERVER_METRICS_TTL_SECONDS * 1_000_000
+    now_timestamp_micro = int(get_current_datetime().timestamp() * 1_000_000)
+    running_timestamp_micro_cutoff = (
+        now_timestamp_micro - settings.SERVER_METRICS_RUNNING_TTL_SECONDS * 1_000_000
+    )
+    finished_timestamp_micro_cutoff = (
+        now_timestamp_micro - settings.SERVER_METRICS_FINISHED_TTL_SECONDS * 1_000_000
     )
-    retention_window = timedelta(seconds=settings.SERVER_METRICS_WINDOW_SECONDS)
-    retention_window_micro = int(retention_window.total_seconds() * 1_000_000)
-    dialect_name = get_db().dialect_name
-    if dialect_name == "sqlite":
-        last_processed_at_epoch = func.unixepoch(JobModel.last_processed_at)
-    elif dialect_name == "postgresql":
-        last_processed_at_epoch = func.date_part("epoch", JobModel.last_processed_at)
-    else:
-        raise ValueError(f"unsupported dialect: {dialect_name}")
-    # Optimization - only check recently finished jobs,
-    # where last_processed_at > now() - retention_window * 2
-    finished_last_processed_cutoff = now - retention_window * 2
     async with get_session_ctx() as session:
         await asyncio.gather(
             session.execute(
@@ -73,19 +63,10 @@ async def delete_metrics():
                 delete(JobMetricsPoint).where(
                     JobMetricsPoint.job_id.in_(
                         select(JobModel.id).where(
-                            JobModel.status.in_(JobStatus.finished_statuses()),
-                            JobModel.last_processed_at > finished_last_processed_cutoff,
-                        )
-                    ),
-                    JobMetricsPoint.timestamp_micro
-                    < (
-                        select(
-                            last_processed_at_epoch.cast(BigInteger) * 1_000_000
-                            - retention_window_micro
+                            JobModel.status.in_(JobStatus.finished_statuses())
                         )
-                        .where(JobModel.id == JobMetricsPoint.job_id)
-                        .scalar_subquery()
                     ),
+                    JobMetricsPoint.timestamp_micro < finished_timestamp_micro_cutoff,
                 )
             ),
         )

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

@@ -870,10 +870,10 @@ def _validate_run_spec_and_set_defaults(run_spec: RunSpec):
     if (
         run_spec.merged_profile.utilization_policy is not None
         and run_spec.merged_profile.utilization_policy.time_window
-        > settings.SERVER_METRICS_TTL_SECONDS
+        > settings.SERVER_METRICS_RUNNING_TTL_SECONDS
     ):
         raise ServerClientError(
-            f"Maximum utilization_policy.time_window is {settings.SERVER_METRICS_TTL_SECONDS}s"
+            f"Maximum utilization_policy.time_window is {settings.SERVER_METRICS_RUNNING_TTL_SECONDS}s"
         )
     set_resources_defaults(run_spec.configuration.resources)
 

src/dstack/_internal/server/settings.py

@@ -1,4 +1,5 @@
 import os
+import warnings
 from pathlib import Path
 
 DSTACK_DIR_PATH = Path("~/.dstack/").expanduser()
@@ -45,8 +46,25 @@
 
 SERVER_GCP_LOGGING_PROJECT = os.getenv("DSTACK_SERVER_GCP_LOGGING_PROJECT")
 
-SERVER_METRICS_TTL_SECONDS = int(os.getenv("DSTACK_SERVER_METRICS_TTL_SECONDS", 3600))
-SERVER_METRICS_WINDOW_SECONDS = int(os.getenv("DSTACK_SERVER_METRICS_WINDOW_SECONDS", 1800))
+SERVER_METRICS_RUNNING_TTL_SECONDS: int
+_SERVER_METRICS_RUNNING_TTL_SECONDS = os.getenv("DSTACK_SERVER_METRICS_RUNNING_TTL_SECONDS")
+if _SERVER_METRICS_RUNNING_TTL_SECONDS is None:
+    _SERVER_METRICS_RUNNING_TTL_SECONDS = os.getenv("DSTACK_SERVER_METRICS_TTL_SECONDS")
+    if _SERVER_METRICS_RUNNING_TTL_SECONDS is not None:
+        warnings.warn(
+            (
+                "DSTACK_SERVER_METRICS_TTL_SECONDS is deprecated,"
+                " use DSTACK_SERVER_METRICS_RUNNING_TTL_SECONDS instead"
+            ),
+            DeprecationWarning,
+        )
+    else:
+        _SERVER_METRICS_RUNNING_TTL_SECONDS = 3600
+SERVER_METRICS_RUNNING_TTL_SECONDS = int(_SERVER_METRICS_RUNNING_TTL_SECONDS)
+del _SERVER_METRICS_RUNNING_TTL_SECONDS
+SERVER_METRICS_FINISHED_TTL_SECONDS = int(
+    os.getenv("DSTACK_SERVER_METRICS_FINISHED_TTL_SECONDS", 7 * 24 * 3600)
+)
 
 DEFAULT_PROJECT_NAME = "main"
 

src/tests/_internal/server/background/tasks/test_process_metrics.py

@@ -114,7 +114,6 @@ async def test_deletes_old_metrics_running_job(self, test_db, session: AsyncSess
             session=session,
             run=run,
             status=JobStatus.RUNNING,
-            last_processed_at=datetime(2023, 1, 2, 3, 5, 0, tzinfo=timezone.utc),
         )
         await create_job_metrics_point(
             session=session,
@@ -131,7 +130,9 @@ async def test_deletes_old_metrics_running_job(self, test_db, session: AsyncSess
             job_model=job,
             timestamp=datetime(2023, 1, 2, 3, 5, 10, tzinfo=timezone.utc),
         )
-        with patch.object(settings, "SERVER_METRICS_TTL_SECONDS", 15):
+        with patch.multiple(
+            settings, SERVER_METRICS_RUNNING_TTL_SECONDS=15, SERVER_METRICS_FINISHED_TTL_SECONDS=0
+        ):
             await delete_metrics()
         res = await session.execute(select(JobMetricsPoint))
         points = res.scalars().all()
@@ -161,24 +162,25 @@ async def test_deletes_old_metrics_finished_job(self, test_db, session: AsyncSes
             session=session,
             run=run,
             status=JobStatus.FAILED,
-            last_processed_at=datetime(2023, 1, 2, 3, 5, 0, tzinfo=timezone.utc),
         )
         await create_job_metrics_point(
             session=session,
             job_model=job,
-            timestamp=datetime(2023, 1, 2, 3, 4, 30, tzinfo=timezone.utc),
+            timestamp=datetime(2023, 1, 2, 3, 4, 10, tzinfo=timezone.utc),
         )
         await create_job_metrics_point(
             session=session,
             job_model=job,
-            timestamp=datetime(2023, 1, 2, 3, 4, 40, tzinfo=timezone.utc),
+            timestamp=datetime(2023, 1, 2, 3, 4, 20, tzinfo=timezone.utc),
         )
         last_metric = await create_job_metrics_point(
             session=session,
             job_model=job,
-            timestamp=datetime(2023, 1, 2, 3, 4, 50, tzinfo=timezone.utc),
+            timestamp=datetime(2023, 1, 2, 3, 5, 10, tzinfo=timezone.utc),
         )
-        with patch.object(settings, "SERVER_METRICS_WINDOW_SECONDS", 15):
+        with patch.multiple(
+            settings, SERVER_METRICS_RUNNING_TTL_SECONDS=0, SERVER_METRICS_FINISHED_TTL_SECONDS=15
+        ):
             await delete_metrics()
         res = await session.execute(select(JobMetricsPoint))
         points = res.scalars().all()