Add new metrics (#2434)

un-def · web-flow · commit e31b609ab9e5 · 2025-03-19T06:17:21.000Z
* `memory_total_bytes`
* `gpu_memory_total_bytes`
* `cpus_detected_num`
diff --git a/src/dstack/_internal/server/routers/metrics.py b/src/dstack/_internal/server/routers/metrics.py
@@ -40,14 +40,16 @@ async def get_job_metrics(
     By default, returns one latest sample. To control time window/number of samples, use
     `limit`, `after`, `before`.
 
-    Supported metrics: [
-        "cpu_usage_percent",
-        "memory_usage_bytes",
-        "memory_working_set_bytes",
-        "gpus_detected_num",
-        "gpu_memory_usage_bytes_gpu{i}",
-        "gpu_util_percent_gpu{i}"
-    ]
+    Supported metrics (all optional):
+    * `cpus_detected_num`
+    * `cpu_usage_percent`
+    * `memory_total_bytes`
+    * `memory_usage_bytes`
+    * `memory_working_set_bytes`
+    * `gpus_detected_num`
+    * `gpu_memory_total_bytes`
+    * `gpu_memory_usage_bytes_gpu{i}`
+    * `gpu_util_percent_gpu{i}`
     """
     _, project = user_project
 
diff --git a/src/dstack/_internal/server/services/metrics.py b/src/dstack/_internal/server/services/metrics.py
@@ -7,8 +7,11 @@
 from sqlalchemy import select
 from sqlalchemy.ext.asyncio import AsyncSession
 
+from dstack._internal.core.models.instances import Resources
 from dstack._internal.core.models.metrics import JobMetrics, Metric
 from dstack._internal.server.models import JobMetricsPoint, JobModel
+from dstack._internal.server.services.jobs import get_job_provisioning_data, get_job_runtime_data
+from dstack._internal.utils.common import get_or_error
 from dstack._internal.utils.logging import get_logger
 
 logger = get_logger(__name__)
@@ -47,17 +50,34 @@ async def get_job_metrics(
     # we need at least 2 points to calculate cpu_usage_percent
     if len(points) < 2:
         return JobMetrics(metrics=[])
-    return _calculate_job_metrics(points)
+    return _calculate_job_metrics(job_model, points)
 
 
-def _calculate_job_metrics(points: Sequence[JobMetricsPoint]) -> JobMetrics:
+def _calculate_job_metrics(job_model: JobModel, points: Sequence[JobMetricsPoint]) -> JobMetrics:
     timestamps: list[datetime] = []
     cpu_usage_points: list[int] = []
     memory_usage_points: list[int] = []
     memory_working_set_points: list[int] = []
     gpus_memory_usage_points: defaultdict[int, list[int]] = defaultdict(list)
     gpus_util_points: defaultdict[int, list[int]] = defaultdict(list)
 
+    cpus_detected_num: Optional[int] = None
+    memory_total: Optional[int] = None
+    gpu_memory_total: Optional[int] = None
+    resources: Optional[Resources] = None
+    jrd = get_job_runtime_data(job_model)
+    if jrd is not None and jrd.offer is not None:
+        resources = jrd.offer.instance.resources
+    else:
+        jpd = get_job_provisioning_data(job_model)
+        if jpd is not None:
+            resources = jpd.instance_type.resources
+    if resources is not None:
+        cpus_detected_num = resources.cpus
+        memory_total = resources.memory_mib * 1024 * 1024
+        if len(resources.gpus) > 0:
+            gpu_memory_total = resources.gpus[0].memory_mib * 1024 * 1024
+
     gpus_detected_num: Optional[int] = None
     gpus_detected_num_mismatch: bool = False
     for point, prev_point in zip(points, points[1:]):
@@ -93,25 +113,23 @@ def _calculate_job_metrics(points: Sequence[JobMetricsPoint]) -> JobMetrics:
             values=memory_working_set_points,
         ),
     ]
+    if cpus_detected_num is not None:
+        metrics.append(_make_constant_metric("cpus_detected_num", timestamps, cpus_detected_num))
+    if memory_total is not None:
+        metrics.append(_make_constant_metric("memory_total_bytes", timestamps, memory_total))
     if gpus_detected_num_mismatch:
         # If number of GPUs changed in the time window, skip GPU metrics altogether, otherwise
         # results can be unpredictable (e.g, one GPU takes place of another, as they are
         # identified by an array index only).
         logger.warning("gpus_detected_num mismatch, skipping GPU metrics")
     else:
         metrics.append(
-            # As gpus_detected_num expected to be constant, we add only two points — the latest
-            # and the earliest in the batch
-            Metric(
-                name="gpus_detected_num",
-                timestamps=[timestamps[0], timestamps[-1]]
-                if len(timestamps) > 1
-                else [timestamps[0]],
-                values=[gpus_detected_num, gpus_detected_num]
-                if len(timestamps) > 1
-                else [gpus_detected_num],
-            )
+            _make_constant_metric("gpus_detected_num", timestamps, get_or_error(gpus_detected_num))
         )
+        if gpu_memory_total is not None:
+            metrics.append(
+                _make_constant_metric("gpu_memory_total_bytes", timestamps, gpu_memory_total)
+            )
         for index, gpu_memory_usage_points in gpus_memory_usage_points.items():
             metrics.append(
                 Metric(
@@ -131,6 +149,14 @@ def _calculate_job_metrics(points: Sequence[JobMetricsPoint]) -> JobMetrics:
     return JobMetrics(metrics=metrics)
 
 
+def _make_constant_metric(name: str, timestamps: list[datetime], value: float) -> Metric:
+    return Metric(
+        name=name,
+        timestamps=timestamps,
+        values=[value] * len(timestamps),
+    )
+
+
 def _get_cpu_usage(last_point: JobMetricsPoint, prev_point: JobMetricsPoint) -> int:
     window = last_point.timestamp_micro - prev_point.timestamp_micro
     if window == 0:
diff --git a/src/dstack/_internal/server/testing/common.py b/src/dstack/_internal/server/testing/common.py
@@ -329,13 +329,20 @@ def get_job_provisioning_data(
     backend: BackendType = BackendType.AWS,
     region: str = "us-east-1",
     gpu_count: int = 0,
+    gpu_memory_gib: float = 16,
     cpu_count: int = 1,
     memory_gib: float = 0.5,
     spot: bool = False,
     hostname: str = "127.0.0.4",
     internal_ip: Optional[str] = "127.0.0.4",
 ) -> JobProvisioningData:
-    gpus = [Gpu(name="T4", memory_mib=16384, vendor=gpuhunt.AcceleratorVendor.NVIDIA)] * gpu_count
+    gpus = [
+        Gpu(
+            name="T4",
+            memory_mib=int(gpu_memory_gib * 1024),
+            vendor=gpuhunt.AcceleratorVendor.NVIDIA,
+        )
+    ] * gpu_count
     return JobProvisioningData(
         backend=backend,
         instance_type=InstanceType(
@@ -597,6 +604,7 @@ def get_instance_offer_with_availability(
     backend: BackendType = BackendType.AWS,
     region: str = "eu-west",
     gpu_count: int = 0,
+    gpu_memory_gib: float = 16,
     cpu_count: int = 2,
     memory_gib: float = 12,
     disk_gib: float = 100.0,
@@ -605,7 +613,13 @@ def get_instance_offer_with_availability(
     total_blocks: int = 1,
     availability_zones: Optional[List[str]] = None,
 ):
-    gpus = [Gpu(name="T4", memory_mib=16384, vendor=gpuhunt.AcceleratorVendor.NVIDIA)] * gpu_count
+    gpus = [
+        Gpu(
+            name="T4",
+            memory_mib=int(gpu_memory_gib * 1024),
+            vendor=gpuhunt.AcceleratorVendor.NVIDIA,
+        )
+    ] * gpu_count
     return InstanceOfferWithAvailability(
         backend=backend,
         instance=InstanceType(
diff --git a/src/tests/_internal/server/routers/test_metrics.py b/src/tests/_internal/server/routers/test_metrics.py
@@ -14,6 +14,9 @@
     create_run,
     create_user,
     get_auth_headers,
+    get_instance_offer_with_availability,
+    get_job_provisioning_data,
+    get_job_runtime_data,
 )
 
 pytestmark = pytest.mark.usefixtures("image_config_mock")
@@ -51,9 +54,18 @@ async def test_returns_metrics(self, test_db, session: AsyncSession, client: Asy
             repo=repo,
             user=user,
         )
+        jpd = get_job_provisioning_data(
+            cpu_count=128, memory_gib=256, gpu_count=2, gpu_memory_gib=32
+        )
+        offer = get_instance_offer_with_availability(
+            cpu_count=64, memory_gib=128, gpu_count=1, gpu_memory_gib=32
+        )
+        jrd = get_job_runtime_data(offer=offer)
         job = await create_job(
             session=session,
             run=run,
+            job_provisioning_data=jpd,
+            job_runtime_data=jrd,
         )
         await create_job_metrics_point(
             session=session,
@@ -108,11 +120,26 @@ async def test_returns_metrics(self, test_db, session: AsyncSession, client: Asy
                     "timestamps": ["2023-01-02T03:04:25+00:00"],
                     "values": [512],
                 },
+                {
+                    "name": "cpus_detected_num",
+                    "timestamps": ["2023-01-02T03:04:25+00:00"],
+                    "values": [64],
+                },
+                {
+                    "name": "memory_total_bytes",
+                    "timestamps": ["2023-01-02T03:04:25+00:00"],
+                    "values": [137438953472],
+                },
                 {
                     "name": "gpus_detected_num",
                     "timestamps": ["2023-01-02T03:04:25+00:00"],
                     "values": [1],
                 },
+                {
+                    "name": "gpu_memory_total_bytes",
+                    "timestamps": ["2023-01-02T03:04:25+00:00"],
+                    "values": [34359738368],
+                },
                 {
                     "name": "gpu_memory_usage_bytes_gpu0",
                     "timestamps": ["2023-01-02T03:04:25+00:00"],
diff --git a/src/tests/_internal/server/services/test_metrics.py b/src/tests/_internal/server/services/test_metrics.py
@@ -12,6 +12,7 @@
     create_repo,
     create_run,
     create_user,
+    get_job_provisioning_data,
 )
 
 
@@ -129,9 +130,13 @@ async def test_get_metrics(
             repo=repo,
             user=user,
         )
+        jpd = get_job_provisioning_data(
+            cpu_count=64, memory_gib=128, gpu_count=2, gpu_memory_gib=32
+        )
         job = await create_job(
             session=session,
             run=run,
+            job_provisioning_data=jpd,
         )
         for dt, _cpu, _mem, _mem_ws, _gpu0_mem, _gpu0_util, _gpu1_mem, _gpu1_util in self.points:
             await create_job_metrics_point(
@@ -151,11 +156,10 @@ async def test_get_metrics(
             Metric(name="cpu_usage_percent", timestamps=ts, values=cpu),
             Metric(name="memory_usage_bytes", timestamps=ts, values=mem),
             Metric(name="memory_working_set_bytes", timestamps=ts, values=mem_ws),
-            Metric(
-                name="gpus_detected_num",
-                timestamps=[ts[0], ts[-1]] if len(ts) > 1 else ts,
-                values=[2, 2] if len(ts) > 1 else [2],
-            ),
+            Metric(name="cpus_detected_num", timestamps=ts, values=[64] * len(ts)),
+            Metric(name="memory_total_bytes", timestamps=ts, values=[137438953472] * len(ts)),
+            Metric(name="gpus_detected_num", timestamps=ts, values=[2] * len(ts)),
+            Metric(name="gpu_memory_total_bytes", timestamps=ts, values=[34359738368] * len(ts)),
             Metric(name="gpu_memory_usage_bytes_gpu0", timestamps=ts, values=gpu0_mem),
             Metric(name="gpu_memory_usage_bytes_gpu1", timestamps=ts, values=gpu1_mem),
             Metric(name="gpu_util_percent_gpu0", timestamps=ts, values=gpu0_util),
