diff --git a/slurm/assets/configuration/spec.yaml b/slurm/assets/configuration/spec.yaml index 43d1c5966f9eb..4bcf02b4f9b1b 100644 --- a/slurm/assets/configuration/spec.yaml +++ b/slurm/assets/configuration/spec.yaml @@ -85,6 +85,60 @@ files: value: type: integer example: 1 + - name: slurm_rest_api_url + description: | + Base URL of the Slurm REST API (slurmrestd), for example http://slurm-restapi:6820. + When set, the check collects metrics over the REST API instead of executing the Slurm + CLI binaries. This lets the check run without co-located Slurm binaries or munge/auth + access (for example, as a cluster check against a containerized Slurm deployment such + as SUNK/Slinky). Leave unset to use the default CLI collection mode. slurmrestd is + optional in Slurm, so this mode requires it to be enabled on the cluster. + fleet_configurable: true + value: + type: string + display_default: null + example: http://slurm-restapi:6820 + - name: slurm_openapi_version + description: | + The slurmrestd OpenAPI data_parser version to query (for example v0.0.42). Only used + when slurm_rest_api_url is set. + fleet_configurable: true + value: + type: string + example: v0.0.42 + - name: slurm_rest_api_token + description: | + The Slurm JWT (auth/jwt) used to authenticate to slurmrestd, sent in the + X-SLURM-USER-TOKEN header. Only used when slurm_rest_api_url is set. Prefer + slurm_rest_api_token_file for tokens that are rotated. + fleet_configurable: false + value: + type: string + secret: true + example: + - name: slurm_rest_api_token_file + description: | + Path to a file containing the Slurm JWT used to authenticate to slurmrestd. The file is + re-read on every check run so rotated tokens are picked up. Takes precedence over + slurm_rest_api_token. Only used when slurm_rest_api_url is set. + fleet_configurable: true + value: + type: string + display_default: null + example: /etc/datadog-agent/slurm.d/token + - name: slurm_rest_api_user + description: | + Slurm username to send in the X-SLURM-USER-NAME header, in addition to the JWT token. + Only needed when slurmrestd sits behind an authenticating proxy that uses one shared + privileged token to act on behalf of multiple users (see + https://slurm.schedmd.com/jwt.html). Not needed for a JWT generated directly for a + single user (e.g. via `scontrol token`), since the token's own `sun` claim already + identifies the user in that case. Only used when slurm_rest_api_url is set. + fleet_configurable: true + value: + type: string + display_default: null + example: datadog - name: sinfo_path description: Full path to the sinfo binary. fleet_configurable: true diff --git a/slurm/assets/service_checks.json b/slurm/assets/service_checks.json index fe51488c7066f..6de8b78ea3a5e 100644 --- a/slurm/assets/service_checks.json +++ b/slurm/assets/service_checks.json @@ -1 +1,14 @@ -[] +[ + { + "agent_version": "7.70.0", + "integration": "Slurm", + "check": "slurm.rest.can_connect", + "statuses": [ + "ok", + "critical" + ], + "groups": [], + "name": "Slurm REST API can connect", + "description": "Returns `CRITICAL` if the Agent cannot reach the Slurm REST API (slurmrestd) when `slurm_rest_api_url` is configured, otherwise returns `OK`." + } +] diff --git a/slurm/changelog.d/24566.added b/slurm/changelog.d/24566.added new file mode 100644 index 0000000000000..4e72b66558395 --- /dev/null +++ b/slurm/changelog.d/24566.added @@ -0,0 +1 @@ +Add an optional slurmrestd (REST API) collection mode. When `slurm_rest_api_url` is set, the check collects scheduler (`sdiag`), partition, and node (`sinfo`) metrics, including per-node and per-partition GPU counts when `collect_gpu_stats` is enabled, from the Slurm REST API over HTTP with a JWT token instead of executing the Slurm CLI binaries, enabling a no-sidecar deployment (for example as a cluster check) for containerized Slurm such as SUNK/Slinky. diff --git a/slurm/datadog_checks/slurm/check.py b/slurm/datadog_checks/slurm/check.py index 7cbee9bd28062..7d4aed7598ad8 100644 --- a/slurm/datadog_checks/slurm/check.py +++ b/slurm/datadog_checks/slurm/check.py @@ -33,6 +33,13 @@ SSHARE_MAP, SSHARE_PARAMS, ) +from .rest import ( + SlurmRestAPIClient, + iter_node_metrics, + iter_partition_gpu_metrics, + iter_partition_metrics, + iter_sdiag_metrics, +) def get_subprocess_output(cmd): @@ -83,6 +90,22 @@ def __init__(self, name, init_config, instances): self.gpu_stats = is_affirmative(self.instance.get('collect_gpu_stats', False)) self.sinfo_collection_level = self.instance.get('sinfo_collection_level', 1) + # REST (slurmrestd) mode. When slurm_rest_api_url is set, collect over the Slurm REST + # API instead of the CLI binaries (see rest.py). Works without co-located binaries, + # munge/auth, or a login node -- suitable for a cluster check against containerized + # Slurm (SUNK/Slinky). slurmrestd is optional and handled gracefully if unreachable. + self.rest_api_url = self.instance.get('slurm_rest_api_url') + self.use_rest_api = bool(self.rest_api_url) + if self.use_rest_api: + self.openapi_version = self.instance.get('slurm_openapi_version', 'v0.0.42') + self._rest_token = self.instance.get('slurm_rest_api_token') + self._rest_token_file = self.instance.get('slurm_rest_api_token_file') + self.rest_client = SlurmRestAPIClient(self.http, self.rest_api_url, self.openapi_version, self.log) + # Only needed behind an authenticating proxy that reuses one privileged token for + # multiple users (https://slurm.schedmd.com/jwt.html); a direct per-user JWT (e.g. + # via `scontrol token`) already identifies the user via its own `sun` claim. + self.rest_client.user = self.instance.get('slurm_rest_api_user') + # Binary paths self.slurm_binaries_dir = self.init_config.get('slurm_binaries_dir', '/usr/bin/') @@ -133,7 +156,84 @@ def __init__(self, name, init_config, instances): self.debug_sacct_stats = is_affirmative(self.instance.get('debug_sacct_stats', False)) self.debug_scontrol_stats = is_affirmative(self.instance.get('debug_scontrol_stats', False)) + def _get_rest_token(self): + if self._rest_token_file: + try: + with open(self._rest_token_file) as f: + return f.read().strip() + except OSError as e: + self.log.error("Could not read slurm_rest_api_token_file %s: %s", self._rest_token_file, e) + return None + return self._rest_token + + def _check_rest(self): + base_tags = list(self.tags) + # Read the token once per run (rotation-safe without re-reading the file per request). + self.rest_client.token = self._get_rest_token() + + connected = self.rest_client.get('ping') is not None + self.service_check('rest.can_connect', self.OK if connected else self.CRITICAL, tags=base_tags) + if not connected: + self.log.warning( + "Could not reach slurmrestd at %s. slurmrestd is optional in Slurm; ensure it is " + "enabled and the JWT token is valid. Skipping this run.", + self.rest_api_url, + ) + return + + if self.collect_sdiag_stats: + self.gauge('sdiag.enabled', 1, tags=base_tags) + for metric, value, tags in iter_sdiag_metrics(self.rest_client.get('diag')): + self.gauge(metric, value, tags=base_tags + tags) + + if self.collect_sinfo_stats: + # Mirror CLI semantics: level 1 (default) collects partition metrics; node metrics + # require level >= 2. + self.gauge('sinfo.partition.enabled', 1, tags=base_tags) + for metric, value, tags in iter_partition_metrics(self.rest_client.get('partitions')): + self.gauge(metric, value, tags=base_tags + tags) + + # /nodes is needed for node metrics (level >= 2) and, when GPU stats are enabled, for + # partition GPU totals (aggregated from nodes, since /partitions has no GPU field). + # Fetch it at most once per run and reuse. + nodes_payload = None + if self.gpu_stats: + nodes_payload = self.rest_client.get('nodes') + for metric, value, tags in iter_partition_gpu_metrics(nodes_payload): + self.gauge(metric, value, tags=base_tags + tags) + + if self.sinfo_collection_level > 1: + self.gauge('sinfo.node.enabled', 1, tags=base_tags) + if nodes_payload is None: + nodes_payload = self.rest_client.get('nodes') + for metric, value, tags in iter_node_metrics(nodes_payload, collect_gpu=self.gpu_stats): + self.gauge(metric, value, tags=base_tags + tags) + + # These collectors are not (yet) supported over REST. scontrol is node-local OS data not + # exposed by slurmrestd; squeue/sacct/sshare REST support is a follow-up. Warn instead of + # silently dropping so a CLI user migrating to REST is not surprised by missing data. + unsupported = [ + name + for name, enabled in ( + ('collect_squeue_stats', self.collect_squeue_stats), + ('collect_sacct_stats', self.collect_sacct_stats), + ('collect_sshare_stats', self.collect_sshare_stats), + ('collect_scontrol_stats', self.collect_scontrol_stats), + ) + if enabled + ] + if unsupported: + self.log.warning( + "The following collectors are enabled but not supported in REST mode and are " + "skipped: %s. scontrol is node-local; squeue/sacct/sshare REST support is a follow-up.", + ', '.join(unsupported), + ) + def check(self, _): + if self.use_rest_api: + self._check_rest() + return + self.collect_metadata() commands = [] diff --git a/slurm/datadog_checks/slurm/config_models/defaults.py b/slurm/datadog_checks/slurm/config_models/defaults.py index 6b3c29a63c19b..668ea6d5c0933 100644 --- a/slurm/datadog_checks/slurm/config_models/defaults.py +++ b/slurm/datadog_checks/slurm/config_models/defaults.py @@ -112,6 +112,10 @@ def instance_sinfo_path(): return '/usr/bin/sinfo' +def instance_slurm_openapi_version(): + return 'v0.0.42' + + def instance_squeue_path(): return '/usr/bin/squeue' diff --git a/slurm/datadog_checks/slurm/config_models/instance.py b/slurm/datadog_checks/slurm/config_models/instance.py index 96f9edc790ce5..f2ff1da85c6d6 100644 --- a/slurm/datadog_checks/slurm/config_models/instance.py +++ b/slurm/datadog_checks/slurm/config_models/instance.py @@ -66,6 +66,11 @@ class InstanceConfig(BaseModel): service: Optional[str] = None sinfo_collection_level: Optional[int] = None sinfo_path: Optional[str] = None + slurm_openapi_version: Optional[str] = None + slurm_rest_api_token: Optional[str] = None + slurm_rest_api_token_file: Optional[str] = None + slurm_rest_api_url: Optional[str] = None + slurm_rest_api_user: Optional[str] = None squeue_path: Optional[str] = None sshare_path: Optional[str] = None tags: Optional[tuple[str, ...]] = None diff --git a/slurm/datadog_checks/slurm/data/conf.yaml.example b/slurm/datadog_checks/slurm/data/conf.yaml.example index dfe0019be650b..1102ae267571b 100644 --- a/slurm/datadog_checks/slurm/data/conf.yaml.example +++ b/slurm/datadog_checks/slurm/data/conf.yaml.example @@ -75,6 +75,46 @@ instances: # # sinfo_collection_level: 1 + ## @param slurm_rest_api_url - string - optional + ## Base URL of the Slurm REST API (slurmrestd), for example http://slurm-restapi:6820. + ## When set, the check collects metrics over the REST API instead of executing the Slurm + ## CLI binaries. This lets the check run without co-located Slurm binaries or munge/auth + ## access (for example, as a cluster check against a containerized Slurm deployment such + ## as SUNK/Slinky). Leave unset to use the default CLI collection mode. slurmrestd is + ## optional in Slurm, so this mode requires it to be enabled on the cluster. + # + # slurm_rest_api_url: http://slurm-restapi:6820 + + ## @param slurm_openapi_version - string - optional - default: v0.0.42 + ## The slurmrestd OpenAPI data_parser version to query (for example v0.0.42). Only used + ## when slurm_rest_api_url is set. + # + # slurm_openapi_version: v0.0.42 + + ## @param slurm_rest_api_token - string - optional + ## The Slurm JWT (auth/jwt) used to authenticate to slurmrestd, sent in the + ## X-SLURM-USER-TOKEN header. Only used when slurm_rest_api_url is set. Prefer + ## slurm_rest_api_token_file for tokens that are rotated. + # + # slurm_rest_api_token: + + ## @param slurm_rest_api_token_file - string - optional + ## Path to a file containing the Slurm JWT used to authenticate to slurmrestd. The file is + ## re-read on every check run so rotated tokens are picked up. Takes precedence over + ## slurm_rest_api_token. Only used when slurm_rest_api_url is set. + # + # slurm_rest_api_token_file: /etc/datadog-agent/slurm.d/token + + ## @param slurm_rest_api_user - string - optional + ## Slurm username to send in the X-SLURM-USER-NAME header, in addition to the JWT token. + ## Only needed when slurmrestd sits behind an authenticating proxy that uses one shared + ## privileged token to act on behalf of multiple users (see + ## https://slurm.schedmd.com/jwt.html). Not needed for a JWT generated directly for a + ## single user (e.g. via `scontrol token`), since the token's own `sun` claim already + ## identifies the user in that case. Only used when slurm_rest_api_url is set. + # + # slurm_rest_api_user: datadog + ## @param sinfo_path - string - optional - default: /usr/bin/sinfo ## Full path to the sinfo binary. # diff --git a/slurm/datadog_checks/slurm/rest.py b/slurm/datadog_checks/slurm/rest.py new file mode 100644 index 0000000000000..857bdca1b67fa --- /dev/null +++ b/slurm/datadog_checks/slurm/rest.py @@ -0,0 +1,272 @@ +# (C) Datadog, Inc. 2024-present +# All rights reserved +# Licensed under a 3-clause BSD style license (see LICENSE) +"""slurmrestd (REST) collection mode for the Slurm check. + +Optional alternative to the CLI collectors: query the Slurm REST API (slurmrestd) over HTTP +with a JWT token (auth/jwt) instead of executing the Slurm CLI binaries. This lets the check +run without co-located Slurm binaries or munge/auth access -- for example as a cluster check +against a containerized Slurm deployment (SUNK/Slinky), including deployments that have NO +login node: the check talks to the slurmrestd service over the network from wherever it runs, +and the JWT is provided out-of-band (token/token_file), so no login pod is required. + +slurmrestd is optional in Slurm and may be disabled. SlurmRestAPIClient.get never raises on +transport or API errors; it logs and returns None so the caller can emit a can_connect service +check and skip metrics for that run. + +Metric names and tags mirror the CLI collectors so REST and CLI are schema-compatible. +""" + +# slurmrestd /diag `statistics` key -> slurm.* metric name (matches the CLI sdiag metric names). +# Note: `schedule_cycle_total` is the scheduling-cycle COUNT (verified: total=340, sum=3255, +# mean=9), which matches the CLI "Total cycles:" value. +SDIAG_REST_METRIC_MAP = { + 'server_thread_count': 'sdiag.server_thread_count', + 'agent_queue_size': 'sdiag.agent_queue_size', + 'agent_count': 'sdiag.agent_count', + 'agent_thread_count': 'sdiag.agent_thread_count', + 'dbd_agent_queue_size': 'sdiag.dbd_agent_queue_size', + 'schedule_queue_length': 'sdiag.last_queue_length', + 'jobs_submitted': 'sdiag.jobs_submitted', + 'jobs_started': 'sdiag.jobs_started', + 'jobs_completed': 'sdiag.jobs_completed', + 'jobs_failed': 'sdiag.jobs_failed', + 'jobs_canceled': 'sdiag.jobs_canceled', + 'jobs_pending': 'sdiag.jobs_pending', + 'jobs_running': 'sdiag.jobs_running', + 'schedule_cycle_last': 'sdiag.last_cycle', + 'schedule_cycle_max': 'sdiag.max_cycle', + 'schedule_cycle_total': 'sdiag.total_cycles', + 'schedule_cycle_mean': 'sdiag.mean_cycle', + 'schedule_cycle_mean_depth': 'sdiag.mean_depth_cycle', + 'schedule_cycle_per_minute': 'sdiag.cycles_per_minute', + # Backfill scheduler statistics. + 'bf_backfilled_jobs': 'sdiag.backfill.total_jobs_since_start', + 'bf_last_backfilled_jobs': 'sdiag.backfill.total_jobs_since_cycle_start', + 'bf_backfilled_het_jobs': 'sdiag.backfill.total_heterogeneous_components', + 'bf_cycle_counter': 'sdiag.backfill.total_cycles', + 'bf_cycle_last': 'sdiag.backfill.last_cycle', + 'bf_cycle_max': 'sdiag.backfill.max_cycle', + 'bf_cycle_mean': 'sdiag.backfill.mean_cycle', + 'bf_last_depth': 'sdiag.backfill.last_depth_cycle', + 'bf_last_depth_try': 'sdiag.backfill.last_depth_try_schedule', + 'bf_depth_mean': 'sdiag.backfill.depth_mean', + 'bf_depth_mean_try': 'sdiag.backfill.depth_mean_try_depth', + 'bf_queue_len': 'sdiag.backfill.last_queue_length', + 'bf_queue_len_mean': 'sdiag.backfill.queue_length_mean', + 'bf_table_size': 'sdiag.backfill.last_table_size', + 'bf_table_size_mean': 'sdiag.backfill.mean_table_size', +} + + +def unwrap_number(value): + """slurmrestd wraps some numeric fields as {"set": bool, "infinite": bool, "number": N}.""" + if isinstance(value, dict): + if not value.get('set', True) or value.get('infinite'): + return None + return value.get('number') + return value + + +def is_number(value): + return isinstance(value, (int, float)) and not isinstance(value, bool) + + +def gpu_count_from_gres(gres: str | None) -> tuple[str | None, int | None]: + """Parse a GRES string into (gpu_type, gpu_count). + + Handles the shapes slurmrestd emits for node ``gres``/``gres_used``: ``"gpu:1"`` (no type), + ``"gpu:tesla_t4:4"`` (with type), and a ``(IDX:...)`` allocation suffix on the used string + (``"gpu:tesla_t4:2(IDX:0-1)"``). Comma-separated multi-GRES uses the first gpu entry. Returns + (None, None) when there is no gpu GRES or the count is unparseable. + """ + if not gres: + return None, None + for segment in gres.split(','): + head = segment.split('(')[0].strip() + parts = head.split(':') + if len(parts) < 2 or parts[0] != 'gpu': + continue + gpu_type = parts[1] if len(parts) >= 3 else None + try: + return gpu_type, int(parts[-1]) + except ValueError: + return None, None + return None, None + + +class SlurmRestAPIClient: + """Thin slurmrestd HTTP client. Never raises on transport/API errors -- returns None. + + `token` is set by the caller once per check run (so a rotated token file is read at most + once per run rather than once per request). `user` is optional and only needed behind an + authenticating proxy that reuses one privileged token for multiple users -- a direct + per-user JWT already identifies the user via its own claims, per + https://slurm.schedmd.com/jwt.html. + """ + + def __init__(self, http, base_url, openapi_version, log): + self.http = http + self.base_url = base_url.rstrip('/') + self.openapi_version = openapi_version + self.log = log + self.token = None + self.user = None + + def _headers(self): + headers = {} + if self.token: + headers['X-SLURM-USER-TOKEN'] = self.token + if self.user: + headers['X-SLURM-USER-NAME'] = self.user + return headers + + def get(self, resource): + """GET /slurm//. Returns the parsed dict, or None on any failure.""" + url = f'{self.base_url}/slurm/{self.openapi_version}/{resource}' + try: + response = self.http.get(url, extra_headers=self._headers()) + response.raise_for_status() + payload = response.json() + except Exception as e: + self.log.error("slurmrestd request to %s failed: %s", url, e) + return None + if not isinstance(payload, dict): + self.log.error("slurmrestd returned an unexpected (non-object) payload for %s", resource) + return None + errors = payload.get('errors') or [] + if errors: + self.log.error( + "slurmrestd returned errors for %s (is slurmrestd enabled and the token valid?): %s", + resource, + errors, + ) + return None + return payload + + +def iter_sdiag_metrics(payload): + """Yield (metric_name, value, tags) for scheduler diagnostics from a /diag payload.""" + statistics = (payload or {}).get('statistics', {}) + for key, metric in SDIAG_REST_METRIC_MAP.items(): + value = unwrap_number(statistics.get(key)) + if is_number(value): + yield metric, value, [] + + +def iter_partition_metrics(payload): + """Yield (metric_name, value, tags) per partition from a /partitions payload. + + Mirrors the CLI level-1 (default) collection, which emits partition-scoped metrics. Known + gap: the CLI path also emits a per-partition node-state breakdown + (partition.node.allocated/idle/other/total, from `sinfo`'s AIOT columns); REST does not, + since that requires a per-partition node-state count not present in the /partitions payload. + """ + for partition in (payload or {}).get('partitions', []): + name = (partition.get('name') or '').strip() + if not name: + continue + tags = [ + f'slurm_partition_name:{name}', + f'slurm_cluster_name:{partition.get("cluster") or "null"}', + ] + nodes = partition.get('nodes') if isinstance(partition.get('nodes'), dict) else {} + node_total = unwrap_number(nodes.get('total')) + if is_number(node_total): + yield 'partition.nodes.count', node_total, tags + info_tags = list(tags) + node_list = nodes.get('configured') + if node_list: + info_tags.append(f'slurm_partition_node_list:{node_list}') + yield 'partition.info', 1, info_tags + + +def iter_node_metrics(payload, collect_gpu=False): + """Yield (metric_name, value, tags) per node from a /nodes payload. + + Emits once per (node, partition) to mirror the CLI `sinfo -N` (one row per node per + partition) so `slurm_partition_name` is single-valued. Descriptive tags (state, features) + are attached to `node.info` only -- matching the CLI, where the scalar node metrics do not + carry the state tag. When collect_gpu is set, node.gpu_total/gpu_used are emitted from the + node's gres/gres_used strings, carrying the CLI's slurm_node_gpu_type tag. + """ + for node in (payload or {}).get('nodes', []): + name = node.get('name') + if not name: + continue + gpu_metrics = {} + gpu_tag = [] + if collect_gpu: + gpu_type, gpu_total = gpu_count_from_gres(node.get('gres')) + _, gpu_used = gpu_count_from_gres(node.get('gres_used')) + gpu_metrics = {'node.gpu_total': gpu_total, 'node.gpu_used': gpu_used} + gpu_tag = [f'slurm_node_gpu_type:{gpu_type or "null"}'] + cpu_load = unwrap_number(node.get('cpu_load')) + # Every numeric field is routed through unwrap_number: slurmrestd's data_parser wraps + # some integer fields as {"set","infinite","number"} depending on OpenAPI/build version, + # and unwrap_number is a no-op on plain numbers, so this is safe regardless of whether a + # given field happens to be wrapped on a particular cluster. + scalar_metrics = { + 'node.cpu.total': unwrap_number(node.get('cpus')), + 'node.cpu.allocated': unwrap_number(node.get('alloc_cpus')), + 'node.cpu.idle': unwrap_number(node.get('alloc_idle_cpus')), + # slurmrestd reports cpu_load scaled x100 (e.g. 305 == load 3.05); match the CLI value. + 'node.cpu_load': cpu_load / 100.0 if is_number(cpu_load) else None, + 'node.memory': unwrap_number(node.get('real_memory')), + 'node.alloc_mem': unwrap_number(node.get('alloc_memory')), + 'node.free_mem': unwrap_number(node.get('free_mem')), + 'node.tmp_disk': unwrap_number(node.get('temporary_disk')), + } + identity_tags = [ + f'slurm_node_name:{name}', + # Note: the /nodes payload keys this `cluster_name`, while /partitions keys the same + # concept `cluster` -- different slurmrestd schema keys for the same field, not a bug. + f'slurm_cluster_name:{node.get("cluster_name") or "null"}', + ] + states = [s.lower() for s in (node.get('state') or [])] + features = ','.join(node.get('active_features') or []) + partition_tags = [f'slurm_partition_name:{p}' for p in (node.get('partitions') or [])] or [None] + + for partition_tag in partition_tags: + tags = identity_tags + ([partition_tag] if partition_tag else []) + for metric, value in scalar_metrics.items(): + if is_number(value): + yield metric, value, tags + for metric, value in gpu_metrics.items(): + if is_number(value): + yield metric, value, tags + gpu_tag + info_tags = tags + [ + 'slurm_node_state:' + (','.join(states) if states else 'unknown'), + f'slurm_node_active_features:{features}', + ] + yield 'node.info', 1, info_tags + + +def iter_partition_gpu_metrics(payload): + """Yield (metric_name, value, tags) for per-partition GPU totals, aggregated from a /nodes + payload. slurmrestd's /partitions payload carries no GPU field, so partition.gpu_total and + partition.gpu_used are summed from each node's gres/gres_used over the node's partition + membership. gpu_type is intentionally omitted (unlike the per-node metric): a partition can + span nodes with differing GPU types, so a single type tag on the aggregate would be misleading. + """ + totals = {} + for node in (payload or {}).get('nodes', []): + _, node_total = gpu_count_from_gres(node.get('gres')) + _, node_used = gpu_count_from_gres(node.get('gres_used')) + if node_total is None and node_used is None: + continue + cluster = node.get('cluster_name') or 'null' + for partition in node.get('partitions') or []: + agg = totals.setdefault(partition, {'total': 0, 'used': 0, 'seen': False, 'cluster': cluster}) + if node_total is not None: + agg['total'] += node_total + agg['seen'] = True + if node_used is not None: + agg['used'] += node_used + agg['seen'] = True + for partition, agg in totals.items(): + if not agg['seen']: + continue + tags = [f'slurm_partition_name:{partition}', f'slurm_cluster_name:{agg["cluster"]}'] + yield 'partition.gpu_total', agg['total'], tags + yield 'partition.gpu_used', agg['used'], tags diff --git a/slurm/tests/test_rest.py b/slurm/tests/test_rest.py new file mode 100644 index 0000000000000..ccd2d6d2a9414 --- /dev/null +++ b/slurm/tests/test_rest.py @@ -0,0 +1,390 @@ +# (C) Datadog, Inc. 2024-present +# All rights reserved +# Licensed under a 3-clause BSD style license (see LICENSE) +from unittest.mock import MagicMock + +import pytest + +from datadog_checks.slurm import SlurmCheck +from datadog_checks.slurm.rest import ( + SlurmRestAPIClient, + gpu_count_from_gres, + is_number, + iter_node_metrics, + iter_partition_gpu_metrics, + iter_partition_metrics, + iter_sdiag_metrics, + unwrap_number, +) + +DIAG_PAYLOAD = { + "statistics": { + "server_thread_count": 3, + "jobs_submitted": 10, + "jobs_running": 2, + "jobs_pending": 1, + "schedule_cycle_last": 55, + "schedule_cycle_max": 900, + "schedule_cycle_total": 340, + "schedule_cycle_mean": 9, + "bf_backfilled_jobs": 7, + "bf_cycle_counter": 4, + "bf_queue_len": 5, + "bf_active": False, # boolean must be ignored + "gettimeofday_latency": 12, # unmapped must be ignored + }, + "errors": [], +} + +NODES_PAYLOAD = { + "nodes": [ + { + "name": "slinky-0", + "cluster_name": "slurm_slurm", + "partitions": ["all"], + "state": ["IDLE", "DYNAMIC_NORM"], + "cpus": 16, + "alloc_cpus": 4, + "alloc_idle_cpus": 12, + "cpu_load": 305, # slurmrestd scales x100 -> load 3.05 + "real_memory": 62919, + "alloc_memory": 2048, + "free_mem": {"set": True, "infinite": False, "number": 2296}, + "temporary_disk": 0, + "active_features": ["cpu"], + } + ], + "errors": [], +} + +PARTITIONS_PAYLOAD = { + "partitions": [ + { + "name": "all", + "cluster": "slurm_slurm", + "nodes": {"total": 1, "configured": "slinky-0"}, + "cpus": {"total": 16}, + } + ], + "errors": [], +} + + +def test_unwrap_number(): + assert unwrap_number(42) == 42 + assert unwrap_number({"set": True, "infinite": False, "number": 7}) == 7 + assert unwrap_number({"set": False, "infinite": False, "number": 0}) is None + assert unwrap_number({"set": True, "infinite": True, "number": 0}) is None + + +def test_is_number_rejects_bool(): + assert is_number(3) and is_number(1.5) + assert not is_number(True) and not is_number(None) and not is_number("x") + + +def test_iter_sdiag_metrics_maps_known_keys(): + metrics = {name: value for name, value, _tags in iter_sdiag_metrics(DIAG_PAYLOAD)} + assert metrics["sdiag.server_thread_count"] == 3 + assert metrics["sdiag.jobs_submitted"] == 10 + assert metrics["sdiag.last_cycle"] == 55 + assert metrics["sdiag.total_cycles"] == 340 + assert metrics["sdiag.mean_cycle"] == 9 + assert metrics["sdiag.backfill.total_jobs_since_start"] == 7 + assert metrics["sdiag.backfill.total_cycles"] == 4 + assert metrics["sdiag.backfill.last_queue_length"] == 5 + # Booleans and unmapped keys are not emitted. + assert not any(name.endswith("bf_active") for name in metrics) + + +def test_iter_sdiag_metrics_handles_empty(): + assert list(iter_sdiag_metrics(None)) == [] + assert list(iter_sdiag_metrics({})) == [] + + +def test_iter_node_metrics(): + emitted = list(iter_node_metrics(NODES_PAYLOAD)) + metrics = {name: value for name, value, _tags in emitted} + assert metrics["node.cpu.total"] == 16 + assert metrics["node.cpu.allocated"] == 4 + assert metrics["node.cpu.idle"] == 12 + assert metrics["node.cpu_load"] == pytest.approx(3.05) # scaled down from 305 + assert metrics["node.memory"] == 62919 + assert metrics["node.alloc_mem"] == 2048 + assert metrics["node.free_mem"] == 2296 # unwrapped + assert metrics["node.info"] == 1 + # node.cpu.other is intentionally not emitted in REST mode (structurally ~0, misleading). + assert "node.cpu.other" not in metrics + + tags_by_metric = {name: t for name, _v, t in emitted} + # Scalar metrics carry identity + partition, but NOT state (state lives on node.info, per CLI). + scalar_tags = tags_by_metric["node.cpu.total"] + assert "slurm_node_name:slinky-0" in scalar_tags + assert "slurm_cluster_name:slurm_slurm" in scalar_tags + assert "slurm_partition_name:all" in scalar_tags + assert not any(t.startswith("slurm_node_state:") for t in scalar_tags) + # node.info carries the full (joined) state. + info_tags = tags_by_metric["node.info"] + assert "slurm_node_state:idle,dynamic_norm" in info_tags + + +def test_iter_node_metrics_multiple_partitions_emits_per_partition(): + payload = {"nodes": [{"name": "n0", "cpus": 8, "partitions": ["a", "b"], "state": ["IDLE"]}]} + emitted = [(m, tuple(t)) for m, _v, t in iter_node_metrics(payload) if m == "node.cpu.total"] + # One series per (node, partition), each with a single-valued partition tag. + assert len(emitted) == 2 + part_tags = {next(t for t in tags if t.startswith("slurm_partition_name:")) for _m, tags in emitted} + assert part_tags == {"slurm_partition_name:a", "slurm_partition_name:b"} + + +def test_iter_node_metrics_missing_fields(): + # A node missing most fields must not crash and must still emit node.info. + emitted = list(iter_node_metrics({"nodes": [{"name": "n0"}]})) + metrics = {name for name, _v, _t in emitted} + assert "node.info" in metrics + assert "node.cpu.total" not in metrics # cpus missing -> skipped, no crash + + +def test_iter_node_metrics_handles_wrapped_numeric_fields(): + # Some slurmrestd builds/data_parser versions wrap fields that are plain ints on others + # (see NODES_PAYLOAD above, e.g. free_mem). Every numeric node field must go through + # unwrap_number, not just the ones observed wrapped in our reference cluster -- otherwise a + # differently-wrapped field is silently dropped (is_number(dict) is False) instead of parsed. + payload = { + "nodes": [ + { + "name": "n0", + "cpus": {"set": True, "infinite": False, "number": 8}, + "alloc_cpus": {"set": True, "infinite": False, "number": 2}, + "alloc_idle_cpus": {"set": True, "infinite": False, "number": 6}, + "real_memory": {"set": True, "infinite": False, "number": 32000}, + "alloc_memory": {"set": True, "infinite": False, "number": 1024}, + "state": ["IDLE"], + } + ] + } + metrics = {name: value for name, value, _tags in iter_node_metrics(payload)} + assert metrics["node.cpu.total"] == 8 + assert metrics["node.cpu.allocated"] == 2 + assert metrics["node.cpu.idle"] == 6 + assert metrics["node.memory"] == 32000 + assert metrics["node.alloc_mem"] == 1024 + + +def test_iter_partition_metrics(): + emitted = list(iter_partition_metrics(PARTITIONS_PAYLOAD)) + metrics = {name: (value, tags) for name, value, tags in emitted} + assert metrics["partition.nodes.count"][0] == 1 + assert "slurm_partition_name:all" in metrics["partition.nodes.count"][1] + assert metrics["partition.info"][0] == 1 + assert "slurm_partition_node_list:slinky-0" in metrics["partition.info"][1] + + +def test_iter_partition_metrics_handles_wrapped_node_total(): + payload = { + "partitions": [ + { + "name": "all", + "nodes": {"total": {"set": True, "infinite": False, "number": 3}, "configured": "n0,n1,n2"}, + } + ] + } + metrics = {name: value for name, value, _tags in iter_partition_metrics(payload)} + assert metrics["partition.nodes.count"] == 3 + + +def test_gpu_count_from_gres(): + # No type (slurmrestd's observed shape), with type, and a used-string IDX suffix. + assert gpu_count_from_gres("gpu:1") == (None, 1) + assert gpu_count_from_gres("gpu:0") == (None, 0) + assert gpu_count_from_gres("gpu:tesla_t4:4") == ("tesla_t4", 4) + assert gpu_count_from_gres("gpu:tesla_t4:2(IDX:0-1)") == ("tesla_t4", 2) + # Multi-GRES: the gpu entry is picked out regardless of position. + assert gpu_count_from_gres("mps:100,gpu:3") == (None, 3) + # No gpu / empty / unparseable. + assert gpu_count_from_gres("") == (None, None) + assert gpu_count_from_gres(None) == (None, None) + assert gpu_count_from_gres("mps:100") == (None, None) + + +def test_iter_node_metrics_gpu_disabled_by_default(): + payload = {"nodes": [{"name": "n0", "partitions": ["all"], "gres": "gpu:4", "gres_used": "gpu:1"}]} + metrics = {name for name, _value, _tags in iter_node_metrics(payload)} + assert "node.gpu_total" not in metrics + assert "node.gpu_used" not in metrics + + +def test_iter_node_metrics_gpu(): + payload = { + "nodes": [ + { + "name": "n0", + "cluster_name": "c", + "partitions": ["all"], + "gres": "gpu:tesla_t4:4", + "gres_used": "gpu:tesla_t4:1(IDX:0)", + } + ] + } + emitted = {name: (value, tags) for name, value, tags in iter_node_metrics(payload, collect_gpu=True)} + assert emitted["node.gpu_total"][0] == 4 + assert emitted["node.gpu_used"][0] == 1 + assert "slurm_node_gpu_type:tesla_t4" in emitted["node.gpu_total"][1] + assert "slurm_node_name:n0" in emitted["node.gpu_total"][1] + + +def test_iter_partition_gpu_metrics_aggregates_across_nodes(): + # Two nodes in partition "all", one also in "gpu"; partition totals are summed from node gres. + payload = { + "nodes": [ + {"name": "n0", "cluster_name": "c", "partitions": ["all", "gpu"], "gres": "gpu:4", "gres_used": "gpu:1"}, + {"name": "n1", "cluster_name": "c", "partitions": ["all"], "gres": "gpu:2", "gres_used": "gpu:2"}, + {"name": "n2", "cluster_name": "c", "partitions": ["all"]}, # no gres -> contributes nothing + ] + } + emitted = {} + for name, value, tags in iter_partition_gpu_metrics(payload): + emitted[(name, tuple(t for t in tags if t.startswith("slurm_partition_name")))] = value + + assert emitted[("partition.gpu_total", ("slurm_partition_name:all",))] == 6 + assert emitted[("partition.gpu_used", ("slurm_partition_name:all",))] == 3 + assert emitted[("partition.gpu_total", ("slurm_partition_name:gpu",))] == 4 + assert emitted[("partition.gpu_used", ("slurm_partition_name:gpu",))] == 1 + + +def test_iter_partition_gpu_metrics_no_gpu_nodes_emits_nothing(): + payload = {"nodes": [{"name": "n0", "partitions": ["all"]}, {"name": "n1", "partitions": ["all"]}]} + assert list(iter_partition_gpu_metrics(payload)) == [] + + +def test_rest_client_returns_none_on_exception(): + http = MagicMock() + http.get.side_effect = ConnectionError("refused") + client = SlurmRestAPIClient(http, "http://x:6820", "v0.0.42", MagicMock()) + assert client.get("nodes") is None + + +def test_rest_client_returns_none_when_payload_has_errors(): + http = MagicMock() + response = MagicMock() + response.json.return_value = {"errors": [{"error": "Protocol authentication error"}]} + http.get.return_value = response + client = SlurmRestAPIClient(http, "http://x:6820", "v0.0.42", MagicMock()) + assert client.get("nodes") is None + + +def test_rest_client_returns_none_on_non_dict_payload(): + # slurmrestd/proxy returning a JSON list or scalar must not raise (never-raise guarantee). + http = MagicMock() + response = MagicMock() + response.json.return_value = [1, 2, 3] + http.get.return_value = response + client = SlurmRestAPIClient(http, "http://x:6820", "v0.0.42", MagicMock()) + assert client.get("nodes") is None + + +def test_rest_client_sends_token_header_and_returns_payload(): + http = MagicMock() + response = MagicMock() + response.json.return_value = {"nodes": [], "errors": []} + http.get.return_value = response + client = SlurmRestAPIClient(http, "http://x:6820/", "v0.0.42", MagicMock()) + client.token = "jwt-abc" + assert client.get("nodes") == {"nodes": [], "errors": []} + _args, kwargs = http.get.call_args + assert kwargs["extra_headers"]["X-SLURM-USER-TOKEN"] == "jwt-abc" + + +def test_rest_client_omits_user_header_by_default(): + # Direct per-user JWTs (e.g. via `scontrol token`) already identify the user via the + # token's own claims -- X-SLURM-USER-NAME must not be sent unless explicitly configured. + http = MagicMock() + response = MagicMock() + response.json.return_value = {"nodes": [], "errors": []} + http.get.return_value = response + client = SlurmRestAPIClient(http, "http://x:6820/", "v0.0.42", MagicMock()) + client.token = "jwt-abc" + client.get("nodes") + _args, kwargs = http.get.call_args + assert "X-SLURM-USER-NAME" not in kwargs["extra_headers"] + + +def test_rest_client_sends_user_header_when_configured(): + # Needed behind an authenticating proxy that reuses one privileged token for multiple users. + http = MagicMock() + response = MagicMock() + response.json.return_value = {"nodes": [], "errors": []} + http.get.return_value = response + client = SlurmRestAPIClient(http, "http://x:6820/", "v0.0.42", MagicMock()) + client.token = "jwt-abc" + client.user = "datadog" + client.get("nodes") + _args, kwargs = http.get.call_args + assert kwargs["extra_headers"]["X-SLURM-USER-NAME"] == "datadog" + assert kwargs["extra_headers"]["X-SLURM-USER-TOKEN"] == "jwt-abc" + + +def test_get_rest_token_from_file(tmp_path): + token_file = tmp_path / "token" + token_file.write_text("jwt-from-file\n") + check = SlurmCheck( + "slurm", + {}, + [{"slurm_rest_api_url": "http://x:6820", "slurm_rest_api_token_file": str(token_file)}], + ) + assert check._get_rest_token() == "jwt-from-file" + + +def test_get_rest_token_missing_file_returns_none(): + check = SlurmCheck( + "slurm", + {}, + [{"slurm_rest_api_url": "http://x:6820", "slurm_rest_api_token_file": "/no/such/token"}], + ) + assert check._get_rest_token() is None + + +def test_cli_mode_when_rest_url_absent(): + check = SlurmCheck("slurm", {}, [{}]) + assert check.use_rest_api is False + + +@pytest.mark.parametrize("ping_ok", [True, False]) +def test_check_rest_mode(aggregator, dd_run_check, ping_ok): + instance = { + "slurm_rest_api_url": "http://slurm-restapi:6820", + "slurm_rest_api_token": "jwt-abc", + "collect_sinfo_stats": True, + "collect_sdiag_stats": True, + "sinfo_collection_level": 2, # level >= 2 also collects node metrics + } + check = SlurmCheck("slurm", {}, [instance]) + assert check.use_rest_api is True + + def fake_get(resource): + if not ping_ok: + return None + return { + "ping": {"errors": []}, + "diag": DIAG_PAYLOAD, + "nodes": NODES_PAYLOAD, + "partitions": PARTITIONS_PAYLOAD, + }.get(resource) + + check.rest_client = MagicMock() + check.rest_client.get.side_effect = fake_get + + dd_run_check(check) + + if ping_ok: + aggregator.assert_service_check("slurm.rest.can_connect", SlurmCheck.OK) + aggregator.assert_metric("slurm.sdiag.jobs_submitted", value=10) + aggregator.assert_metric("slurm.partition.nodes.count", value=1) + aggregator.assert_metric("slurm.sinfo.partition.enabled", value=1) + aggregator.assert_metric("slurm.sinfo.node.enabled", value=1) + aggregator.assert_metric("slurm.node.cpu.total", value=16) + aggregator.assert_metric("slurm.node.cpu_load", value=pytest.approx(3.05)) + aggregator.assert_metric("slurm.node.free_mem", value=2296) + aggregator.assert_metric("slurm.node.info", value=1) + else: + aggregator.assert_service_check("slurm.rest.can_connect", SlurmCheck.CRITICAL) + assert not aggregator.metric_names