Merge pull request #156 from mila-iqia/SARC-384-sanity-check-prometheus

[SARC-384] Sanity check: les gpu_types des fichiers de conf des clusters sont-ils les mêmes que ceux renvoyés par Prometheus ?

Author: abergeron

sarc/alerts/usage_alerts/prometheus_gpu_types.py

@@ -0,0 +1,69 @@
+import logging
+
+from sarc.config import config
+from sarc.jobs.node_gpu_mapping import get_node_to_gpu
+
+logger = logging.getLogger(__name__)
+
+
+def check_prometheus_vs_slurmconfig(cluster_name=None):
+    """
+    Check if GPU types from Prometheus are the same as the ones in slurm config files.
+
+    To get Prometheus GPU types, we make a Prometheus query:
+    `slurm_job_utilization_gpu_memory{cluster=<cluster name>}`,
+    then we collect GPU types from results.
+
+    To get slurm config GPU types, we collect all GPU types
+    from latest node => GPU mapping in database.
+    Node => GPU mappings should have been collected using command line
+    `sarc acquire slurmconfig -c <cluster name>`.
+
+    Parameters
+    ----------
+    cluster_name: str
+        Name of cluster to check. If None, all clusters are checked.
+    """
+    if cluster_name is None:
+        clusters = config().clusters.values()
+    else:
+        clusters = [config().clusters[cluster_name]]
+
+    for cluster in clusters:
+        # We only check clusters which have a prometheus_url
+        if not cluster.prometheus_url:
+            continue
+
+        # Get slurm config GPU types from latest
+        # node => GPU mappings stored in database
+        slurmconfig_gpu_types = set()
+        mapping = get_node_to_gpu(cluster.name)
+        if mapping:
+            for gpu_types in mapping.node_to_gpu.values():
+                # If gpu_types is a string, split it on commas
+                if isinstance(gpu_types, str):
+                    gpu_types = gpu_types.split(",")
+                # If a GPU type is in format gpu:<name>:<count>, get <name>
+                for gpu_type in gpu_types:
+                    if gpu_type.startswith("gpu:") and gpu_type.count(":") == 2:
+                        _, gpu_type, _ = gpu_type.split(":")
+                    slurmconfig_gpu_types.add(gpu_type)
+        if not slurmconfig_gpu_types:
+            # Warn if there is no slurm config GPUs available.
+            logger.warning(
+                f"[prometheus][{cluster.name}] cannot find GPU types from slurm config file. "
+                f"You may need to call `sarc acquire slurmconfig -c {cluster.name}`"
+            )
+        else:
+            # Get Prometheus GPU types
+            query = f'slurm_job_utilization_gpu_memory{{cluster="{cluster.name}"}}'
+            results = cluster.prometheus.custom_query(query)
+            prometheus_gpu_types = {result["metric"]["gpu_type"] for result in results}
+
+            # Warn for each prometheus GPU not found in slurm config GPUs.
+            only_in_prometheus = prometheus_gpu_types - slurmconfig_gpu_types
+            for gpu_type in only_in_prometheus:
+                logger.warning(
+                    f"[prometheus][{cluster.name}] gpu_type not found in slurm config file: {gpu_type}. "
+                    f"Expected: {', '.join(sorted(slurmconfig_gpu_types))}"
+                )

tests/functional/usage_alerts/test_prometheus_gpu_types.py

@@ -0,0 +1,121 @@
+import re
+from datetime import datetime
+from unittest.mock import MagicMock
+
+import pytest
+
+from sarc.alerts.usage_alerts.prometheus_gpu_types import (
+    check_prometheus_vs_slurmconfig,
+)
+from sarc.config import config
+
+TESTING_DATA = {
+    "00_hyrule": {
+        "cluster": "hyrule",
+        "message": "No node_to_gpu, so no slurm config data: should complain.",
+        "node_to_gpu": {},
+        "prometheus": [],
+    },
+    "01_gerudo": {
+        "cluster": "gerudo",
+        "message": "No node_to_gpu, so no slurm config data: should complain.",
+        "node_to_gpu": {},
+        "prometheus": [{"metric": {"gpu_type": "phantom_gpu"}}],
+    },
+    "10_patate": {
+        "cluster": "patate",
+        "message": "No prometheus data to check: no warning.",
+        "node_to_gpu": {"node0": "gpu0"},
+        "prometheus": [],
+    },
+    "11_0_fromage": {
+        "cluster": "fromage",
+        "message": "Both slurm config data and prometheus data available, but prometheus GPU not in slurm config: warning.",
+        "node_to_gpu": {"node0": "gpu0"},
+        "prometheus": [{"metric": {"gpu_type": "phantom_gpu"}}],
+    },
+    "11_1_raisin": {
+        "cluster": "raisin",
+        "message": "Both slurm config data and prometheus data available, and prometheus GPU in slurm config: no warning.",
+        "node_to_gpu": {"node0": "phantom_gpu"},
+        "prometheus": [{"metric": {"gpu_type": "phantom_gpu"}}],
+    },
+}
+
+
+@pytest.mark.usefixtures("empty_read_write_db", "tzlocal_is_mtl")
+@pytest.mark.parametrize("params", TESTING_DATA.values(), ids=TESTING_DATA.keys())
+def test_check_prometheus_vs_slurmconfig(params, monkeypatch, caplog, file_regression):
+    """Test each case from TEST_DATA (one test per cluster)."""
+
+    from prometheus_api_client import PrometheusConnect
+
+    # Mock PrometheusConnect.custom_query() to prevent a real call to Prometheus
+    monkeypatch.setattr(
+        PrometheusConnect,
+        "custom_query",
+        MagicMock(return_value=params["prometheus"]),
+    )
+    # Add node_to_gpu entry in db if necessary
+    if params["node_to_gpu"]:
+        db = config().mongo.database_instance
+        db.node_gpu_mapping.insert_one(
+            {
+                "cluster_name": params["cluster"],
+                "since": datetime.now(),
+                "node_to_gpu": params["node_to_gpu"],
+            }
+        )
+
+    check_prometheus_vs_slurmconfig(cluster_name=params["cluster"])
+    file_regression.check(
+        params["message"]
+        + "\n\n"
+        + re.sub(
+            r"WARNING +sarc\.alerts\.usage_alerts\.prometheus_gpu_types:prometheus_gpu_types.py:[0-9]+ +",
+            "",
+            caplog.text,
+        )
+    )
+
+
+@pytest.mark.usefixtures("empty_read_write_db", "tzlocal_is_mtl")
+def test_check_prometheus_vs_slurmconfig_all(monkeypatch, caplog, file_regression):
+    """Test all data at once (all clusters)."""
+
+    from prometheus_api_client import PrometheusConnect
+
+    def _gen_fake_custom_query(self_, query_: str):
+        return [
+            prom_data
+            for params in TESTING_DATA.values()
+            if params["cluster"] in query_
+            for prom_data in params["prometheus"]
+        ]
+
+    monkeypatch.setattr(
+        PrometheusConnect,
+        "custom_query",
+        _gen_fake_custom_query,
+    )
+
+    db = config().mongo.database_instance
+    collection = db.node_gpu_mapping
+    for params in TESTING_DATA.values():
+        if params["node_to_gpu"]:
+            collection.insert_one(
+                {
+                    "cluster_name": params["cluster"],
+                    "since": datetime.now(),
+                    "node_to_gpu": params["node_to_gpu"],
+                }
+            )
+
+    check_prometheus_vs_slurmconfig()
+    file_regression.check(
+        re.sub(
+            r"WARNING +sarc\.alerts\.usage_alerts\.prometheus_gpu_types:prometheus_gpu_types.py:[0-9]+ +",
+            "",
+            caplog.text,
+        )
+    )

tests/functional/usage_alerts/test_prometheus_gpu_types/test_check_prometheus_vs_slurmconfig_00_hyrule_.txt

@@ -0,0 +1,3 @@
+No node_to_gpu, so no slurm config data: should complain.
+
+[prometheus][hyrule] cannot find GPU types from slurm config file. You may need to call `sarc acquire slurmconfig -c hyrule`

tests/functional/usage_alerts/test_prometheus_gpu_types/test_check_prometheus_vs_slurmconfig_01_gerudo_.txt

@@ -0,0 +1,3 @@
+No node_to_gpu, so no slurm config data: should complain.
+
+[prometheus][gerudo] cannot find GPU types from slurm config file. You may need to call `sarc acquire slurmconfig -c gerudo`

tests/functional/usage_alerts/test_prometheus_gpu_types/test_check_prometheus_vs_slurmconfig_10_patate_.txt

@@ -0,0 +1,2 @@
+No prometheus data to check: no warning.
+

tests/functional/usage_alerts/test_prometheus_gpu_types/test_check_prometheus_vs_slurmconfig_11_0_fromage_.txt

@@ -0,0 +1,3 @@
+Both slurm config data and prometheus data available, but prometheus GPU not in slurm config: warning.
+
+[prometheus][fromage] gpu_type not found in slurm config file: phantom_gpu. Expected: gpu0

tests/functional/usage_alerts/test_prometheus_gpu_types/test_check_prometheus_vs_slurmconfig_11_1_raisin_.txt

@@ -0,0 +1,2 @@
+Both slurm config data and prometheus data available, and prometheus GPU in slurm config: no warning.
+

tests/functional/usage_alerts/test_prometheus_gpu_types/test_check_prometheus_vs_slurmconfig_all.txt

@@ -0,0 +1,4 @@
+[prometheus][fromage] gpu_type not found in slurm config file: phantom_gpu. Expected: gpu0
+[prometheus][gerudo] cannot find GPU types from slurm config file. You may need to call `sarc acquire slurmconfig -c gerudo`
+[prometheus][hyrule] cannot find GPU types from slurm config file. You may need to call `sarc acquire slurmconfig -c hyrule`
+[prometheus][mila] cannot find GPU types from slurm config file. You may need to call `sarc acquire slurmconfig -c mila`