Resolve SLURM scontrol namespace PIDs (DataDog#24384)

* Resolve SLURM scontrol namespace PIDs

* Add SLURM changelog entry

* Track multiple SLURM host PID matches

* Tag SLURM host and namespace PIDs

* Refine SLURM PID match handling

* Sync SLURM generated config example

* Document SLURM PID match limitation

* Fix SLURM namespace PID collision handling

Author: gjulianm

slurm/assets/configuration/spec.yaml

@@ -49,12 +49,20 @@ files:
         example: True
     - name: collect_scontrol_stats
       description: |
-        Whether or not to collect statistics from the scontrol command. This is mainly used in the worker 
+        Whether or not to collect statistics from the scontrol command. This is mainly used in the worker
         node to collect the list of running jobs along with their PIDs.
       fleet_configurable: true
       value:
         type: boolean
         example: False
+    - name: resolve_scontrol_host_pids
+      description: |
+        Whether or not to resolve PIDs reported by scontrol through the host proc filesystem before collecting
+        process tags. Enable this when scontrol reports namespace PIDs instead of host PIDs.
+      fleet_configurable: true
+      value:
+        type: boolean
+        example: False
     - name: collect_gpu_stats
       description: Whether or not to collect GPU statistics when Slurm is configured to use GPUs.
       fleet_configurable: true
@@ -70,8 +78,8 @@ files:
     - name: sinfo_collection_level
       description: |
         The level of detail to collect from the sinfo command. The default is 'basic'. Available options are 1, 2 and
-        3. Level 1 collects data only for partitions. Level 2 collects data from individual nodes. Level 3 
-        collects data from from individual nodes as well but is more verbose and includes data such as CPU and 
+        3. Level 1 collects data only for partitions. Level 2 collects data from individual nodes. Level 3
+        collects data from from individual nodes as well but is more verbose and includes data such as CPU and
         memory usage as reported from the OS, as well as additional tags.
       fleet_configurable: true
       value:
@@ -187,7 +195,7 @@ files:
           https://docs.datadoghq.com/developers/write_agent_check/#collection-interval
 
           Most Slurm metrics are collected from calling the different binaries. Depending on the size of the Slurm cluster,
-          this can be a very expensive operation. It is recommended to set this to a higher value than the default 15 
+          this can be a very expensive operation. It is recommended to set this to a higher value than the default 15
           seconds, but this can be adjusted based on the size of the cluster and the desired granularity of the metrics.
         min_collection_interval.value.display_default: 60
         min_collection_interval.value.default: 60
@@ -196,7 +204,7 @@ files:
   - template: logs
     example:
     - type: file
-      path: /var/log/slurm/slurmd.log 
+      path: /var/log/slurm/slurmd.log
       source: slurm
       service: slurm
     - type: file

slurm/changelog.d/24384.fixed

@@ -0,0 +1 @@
+Fix process tag enrichment when SLURM scontrol reports namespace PIDs.

slurm/datadog_checks/slurm/check.py

@@ -5,6 +5,7 @@
 import re
 import subprocess
 import time
+from dataclasses import dataclass
 from datetime import timedelta
 
 from datadog_checks.base import AgentCheck, is_affirmative
@@ -51,6 +52,12 @@ def parse_duration(time_str):
         return None
 
 
+@dataclass
+class ProcessPidMatch:
+    host_pid: str
+    namespace_pids: list[str]
+
+
 class SlurmCheck(AgentCheck, ConfigMixin):
     # This will be the prefix of every metric and service check the integration sends
     __NAMESPACE__ = 'slurm'
@@ -70,6 +77,7 @@ def __init__(self, name, init_config, instances):
         self.collect_sacct_stats = is_affirmative(self.instance.get('collect_sacct_stats', True))
         self.collect_scontrol_stats = is_affirmative(self.instance.get('collect_scontrol_stats', False))
         self.collect_seff_stats = is_affirmative(self.instance.get('collect_seff_stats', False))
+        self.resolve_scontrol_host_pids = is_affirmative(self.instance.get('resolve_scontrol_host_pids', False))
 
         # Additional configurations
         self.gpu_stats = is_affirmative(self.instance.get('collect_gpu_stats', False))
@@ -589,6 +597,9 @@ def process_scontrol(self, output):
 
         # Cache for job details to avoid duplicate calls
         job_details_cache = {}
+        host_pid_matches_by_namespace_pid: dict[str, list[ProcessPidMatch]] = (
+            self._get_host_pid_matches_by_namespace_pid() if self.resolve_scontrol_host_pids else {}
+        )
 
         for line in lines[1:]:
             tags = [f"slurm_node_name:{slurm_node.strip()}"]
@@ -597,14 +608,17 @@ def process_scontrol(self, output):
 
             for header, value in zip(headers, fields):
                 new_header = SCONTROL_TAG_MAPPING.get(header, f"slurm_{header.lower()}")
-                tags.append(f"{new_header}:{value}")
 
                 if new_header == "pid":
-                    # Example gpu tags being returned:
-                    # ['gpu_vendor:nvidia', 'gpu_device:tesla_v100', 'gpu_uuid:gpu_xxxx...']
-                    pidtags = tagger.tag(f"process://{value}", tagger.ORCHESTRATOR)
-                    if pidtags:  # Guard against tagger.tag returning None
-                        tags.extend(pidtags)
+                    host_pid_match = self._resolve_scontrol_host_pid(value, host_pid_matches_by_namespace_pid)
+                    tags.extend(self._get_process_tags(host_pid_match.host_pid))
+                    for namespace_pid in host_pid_match.namespace_pids:
+                        if namespace_pid == host_pid_match.host_pid:
+                            continue
+                        tags.append(f"nspid:{namespace_pid}")
+                    value = host_pid_match.host_pid
+
+                tags.append(f"{new_header}:{value}")
 
                 if header == "JOBID" and value.isdigit():
                     job_id = value
@@ -617,6 +631,78 @@ def process_scontrol(self, output):
 
             self.gauge("scontrol.jobs.info", 1, tags=tags + self.tags)
 
+    def _resolve_scontrol_host_pid(
+        self, namespace_pid: str, host_pid_matches_by_namespace_pid: dict[str, list[ProcessPidMatch]]
+    ) -> ProcessPidMatch:
+        host_pid_matches = host_pid_matches_by_namespace_pid.get(namespace_pid)
+        if not host_pid_matches:
+            return ProcessPidMatch(host_pid=namespace_pid, namespace_pids=[])
+
+        if len(host_pid_matches) > 1:
+            namespace_match = host_pid_matches[0]
+            for host_pid_match in host_pid_matches:
+                # Prefer a translated namespace PID, but do not disambiguate by container yet.
+                if namespace_pid != host_pid_match.host_pid and namespace_pid in host_pid_match.namespace_pids:
+                    namespace_match = host_pid_match
+                    break
+
+            matches = [
+                {"host_pid": match.host_pid, "namespace_pids": match.namespace_pids} for match in host_pid_matches
+            ]
+            self.log.debug(
+                "Found multiple host PID matches for scontrol namespace PID %s: %s. Using host PID %s.",
+                namespace_pid,
+                matches,
+                namespace_match.host_pid,
+            )
+            return namespace_match
+
+        return host_pid_matches[0]
+
+    def _get_process_tags(self, pid):
+        # Example gpu tags being returned:
+        # ['gpu_vendor:nvidia', 'gpu_device:tesla_v100', 'gpu_uuid:gpu_xxxx...']
+        pidtags = tagger.tag(f"process://{pid}", tagger.ORCHESTRATOR)
+        if pidtags:  # Guard against tagger.tag returning None
+            return pidtags
+
+        return []
+
+    def _get_host_pid_matches_by_namespace_pid(self):
+        host_proc = os.environ.get("HOST_PROC", "/host/proc")
+        host_pid_matches_by_namespace_pid = {}
+
+        try:
+            proc_entries = os.scandir(host_proc)
+        except OSError as e:
+            self.log.debug("Unable to scan host proc path '%s': %s", host_proc, e)
+            return host_pid_matches_by_namespace_pid
+
+        with proc_entries:
+            for entry in proc_entries:
+                if not entry.name.isdigit():
+                    continue
+
+                pid_match = ProcessPidMatch(
+                    host_pid=entry.name,
+                    namespace_pids=self._read_namespace_pids(entry.path, entry.name),
+                )
+                for namespace_pid in pid_match.namespace_pids:
+                    host_pid_matches_by_namespace_pid.setdefault(namespace_pid, []).append(pid_match)
+
+        return host_pid_matches_by_namespace_pid
+
+    def _read_namespace_pids(self, proc_path, host_pid):
+        try:
+            with open(os.path.join(proc_path, "status")) as status_file:
+                for line in status_file:
+                    if line.startswith("NSpid:"):
+                        return line.split()[1:]
+        except OSError as e:
+            self.log.debug("Unable to read process status for PID %s: %s", host_pid, e)
+
+        return [host_pid]
+
     def _enrich_scontrol_tags(self, job_id):
         # Tries to enrich the scontrol job with additional details from squeue.
         try:

slurm/datadog_checks/slurm/config_models/defaults.py

@@ -84,6 +84,10 @@ def instance_min_collection_interval():
     return 60
 
 
+def instance_resolve_scontrol_host_pids():
+    return False
+
+
 def instance_sacct_path():
     return '/usr/bin/sacct'
 

slurm/datadog_checks/slurm/config_models/instance.py

@@ -58,6 +58,7 @@ class InstanceConfig(BaseModel):
     enable_legacy_tags_normalization: Optional[bool] = None
     metric_patterns: Optional[MetricPatterns] = None
     min_collection_interval: Optional[float] = None
+    resolve_scontrol_host_pids: Optional[bool] = None
     sacct_path: Optional[str] = None
     scontrol_path: Optional[str] = None
     sdiag_path: Optional[str] = None

slurm/datadog_checks/slurm/data/conf.yaml.example

@@ -46,11 +46,17 @@ instances:
     # collect_sshare_stats: true
 
     ## @param collect_scontrol_stats - boolean - optional - default: false
-    ## Whether or not to collect statistics from the scontrol command. This is mainly used in the worker 
+    ## Whether or not to collect statistics from the scontrol command. This is mainly used in the worker
     ## node to collect the list of running jobs along with their PIDs.
     #
     # collect_scontrol_stats: false
 
+    ## @param resolve_scontrol_host_pids - boolean - optional - default: false
+    ## Whether or not to resolve PIDs reported by scontrol through the host proc filesystem before collecting
+    ## process tags. Enable this when scontrol reports namespace PIDs instead of host PIDs.
+    #
+    # resolve_scontrol_host_pids: false
+
     ## @param collect_gpu_stats - boolean - optional - default: false
     ## Whether or not to collect GPU statistics when Slurm is configured to use GPUs.
     #
@@ -63,8 +69,8 @@ instances:
 
     ## @param sinfo_collection_level - integer - optional - default: 1
     ## The level of detail to collect from the sinfo command. The default is 'basic'. Available options are 1, 2 and
-    ## 3. Level 1 collects data only for partitions. Level 2 collects data from individual nodes. Level 3 
-    ## collects data from from individual nodes as well but is more verbose and includes data such as CPU and 
+    ## 3. Level 1 collects data only for partitions. Level 2 collects data from individual nodes. Level 3
+    ## collects data from from individual nodes as well but is more verbose and includes data such as CPU and
     ## memory usage as reported from the OS, as well as additional tags.
     #
     # sinfo_collection_level: 1
@@ -125,7 +131,7 @@ instances:
     ## https://docs.datadoghq.com/developers/write_agent_check/#collection-interval
     ##
     ## Most Slurm metrics are collected from calling the different binaries. Depending on the size of the Slurm cluster,
-    ## this can be a very expensive operation. It is recommended to set this to a higher value than the default 15 
+    ## this can be a very expensive operation. It is recommended to set this to a higher value than the default 15
     ## seconds, but this can be adjusted based on the size of the cluster and the desired granularity of the metrics.
     #
     min_collection_interval: 60