From 403276168f9ce62a2ccd461849fb556514148009 Mon Sep 17 00:00:00 2001 From: Minyi Zhu Date: Thu, 16 Jul 2026 23:37:41 -0400 Subject: [PATCH] Add Slurm-job tagging for GPU process metrics (system-probe) Tags GPU process metrics (gpu.process.*, gpu.memory.*) with the owning Slurm job's identity (slurm_job_id, slurm_job_name, slurm_job_partition) on Slurm-on-Kubernetes deployments, so GPU usage can be attributed to Slurm jobs without a compute sidecar. Identity is resolved in the system-probe GPU module, which already holds SYS_PTRACE, by reading SLURM_JOB_* from /proc//environ for the union of eBPF-tracked and NVML-visible GPU PIDs (device.GetComputeRunningProcesses) -- exactly the PID set the core agent emits process metrics for. It is shipped in model.GPUStats.SlurmInfoByPID over the existing socket; the core agent looks it up by PID in the workload tag cache and tags each GPU process metric. The core agent performs no privileged reads and needs no SYS_PTRACE. Opt-in via the system-probe setting gpu_monitoring.enable_slurm_job_tagging (default false). Resolution is stateless; the PID -> identity mapping is dropped whenever the system-probe stats refresh fails, so a reused PID is never misattributed to a previous job's tags. --- pkg/collector/corechecks/gpu/gpu.go | 10 +- pkg/collector/corechecks/gpu/model/model.go | 19 ++++ pkg/collector/corechecks/gpu/tags.go | 44 ++++++++ pkg/collector/corechecks/gpu/tags_test.go | 93 ++++++++++++++++ .../schema/yaml/system-probe_schema.yaml | 4 + pkg/config/setup/system_probe_settings.go | 1 + pkg/gpu/BUILD.bazel | 2 + pkg/gpu/config/config.go | 4 + pkg/gpu/context.go | 78 +++++++++++++ pkg/gpu/context_test.go | 86 +++++++++++++++ pkg/gpu/stats.go | 1 + pkg/process/util/slurm/BUILD.bazel | 37 +++++++ pkg/process/util/slurm/slurm.go | 104 ++++++++++++++++++ pkg/process/util/slurm/slurm_test.go | 79 +++++++++++++ ...lurm-job-gpu-tagging-8f3c1a7e9d2b4051.yaml | 16 +++ 15 files changed, 577 insertions(+), 1 deletion(-) create mode 100644 pkg/process/util/slurm/BUILD.bazel create mode 100644 pkg/process/util/slurm/slurm.go create mode 100644 pkg/process/util/slurm/slurm_test.go create mode 100644 releasenotes/notes/slurm-job-gpu-tagging-8f3c1a7e9d2b4051.yaml diff --git a/pkg/collector/corechecks/gpu/gpu.go b/pkg/collector/corechecks/gpu/gpu.go index 69cf7afe5525..a6f127027699 100644 --- a/pkg/collector/corechecks/gpu/gpu.go +++ b/pkg/collector/corechecks/gpu/gpu.go @@ -38,7 +38,6 @@ const ( gpuMetricsNs = "gpu." ) -// logLimitCheck is used to limit the number of times we log messages about streams and cuda events, as that can be very verbose var logLimitCheck = log.NewLogLimit(20, 10*time.Minute) var _ check.IssueAwareCheck = (*Check)(nil) @@ -170,6 +169,7 @@ func (c *Check) Configure(senderManager sender.SenderManager, _ uint64, config, return fmt.Errorf("error creating workload tag cache: %w", err) } c.workloadTagCache = workloadTagCache + c.deviceEvtGatherer = nvidia.NewDeviceEventsGatherer() // Compute whether we should prefer system-probe process metrics @@ -305,6 +305,14 @@ func (c *Check) Run() error { } // Continue with NVML-only metrics, SP collectors will return empty metrics } + + // Refresh the PID -> Slurm job identity mapping from the system-probe payload. Slurm + // identity is resolved system-probe-side (where SYS_PTRACE is held) and applies to every + // GPU metric for that PID, including NVML-sourced ones, via the shared workload tag cache. + // GetStats() is nil when the refresh above failed; passing that through clears the mapping, + // matching SystemProbeCache's fail-safe of dropping stale data. Keeping a stale mapping + // could misattribute a reused PID to the previous job's slurm_* tags. + c.workloadTagCache.SetSlurmInfo(slurmInfoFromStats(c.spCache.GetStats())) } if c.prmCache != nil { diff --git a/pkg/collector/corechecks/gpu/model/model.go b/pkg/collector/corechecks/gpu/model/model.go index b1e6487a740b..6faa3e122cc1 100644 --- a/pkg/collector/corechecks/gpu/model/model.go +++ b/pkg/collector/corechecks/gpu/model/model.go @@ -42,6 +42,18 @@ type ProcessStatsKey struct { ContainerID string `json:"container_id"` } +// SlurmInfo is the Slurm job identity resolved for a process by the system-probe GPU module +// from /proc//environ. An empty JobID means the process is not a Slurm job. It is only +// populated when Slurm job tagging is enabled in the GPU module (gpu_monitoring.enable_slurm_job_tagging). +type SlurmInfo struct { + // JobID is the numeric Slurm job ID, from SLURM_JOB_ID. Empty if the process is not a Slurm job. + JobID string `json:"job_id,omitempty"` + // JobName is the Slurm job name, from SLURM_JOB_NAME. Only set when JobID is set. + JobName string `json:"job_name,omitempty"` + // Partition is the Slurm partition, from SLURM_JOB_PARTITION. Only set when JobID is set. + Partition string `json:"partition,omitempty"` +} + // ProcessStatsTuple is a single entry in the GPUStats array, as we cannot use a complex key in the map type ProcessStatsTuple struct { Key ProcessStatsKey @@ -68,4 +80,11 @@ type DeviceStatsTuple struct { type GPUStats struct { ProcessMetrics []ProcessStatsTuple `json:"process_metrics"` // Per-process metrics DeviceMetrics []DeviceStatsTuple `json:"device_metrics"` // Device-level metrics + + // SlurmInfoByPID maps each GPU-using process's host PID to its owning Slurm job identity. + // It is populated (system-probe-side, where SYS_PTRACE is held) only when Slurm job tagging is + // enabled, and covers the union of eBPF-tracked processes and NVML-visible compute processes, + // so the core agent can tag every GPU process metric it emits without reading /proc itself. + // Nil when tagging is disabled or no GPU process resolved to a Slurm job. + SlurmInfoByPID map[uint32]SlurmInfo `json:"slurm_info_by_pid,omitempty"` } diff --git a/pkg/collector/corechecks/gpu/tags.go b/pkg/collector/corechecks/gpu/tags.go index 0b05e824d726..f6ebd37d3942 100644 --- a/pkg/collector/corechecks/gpu/tags.go +++ b/pkg/collector/corechecks/gpu/tags.go @@ -17,6 +17,7 @@ import ( taggertypes "github.com/DataDog/datadog-agent/comp/core/tagger/types" telemetry "github.com/DataDog/datadog-agent/comp/core/telemetry/def" workloadmeta "github.com/DataDog/datadog-agent/comp/core/workloadmeta/def" + "github.com/DataDog/datadog-agent/pkg/collector/corechecks/gpu/model" agenterrors "github.com/DataDog/datadog-agent/pkg/errors" "github.com/DataDog/datadog-agent/pkg/gpu/config/consts" proccontainers "github.com/DataDog/datadog-agent/pkg/process/util/containers" @@ -43,6 +44,7 @@ type WorkloadTagCache struct { tagger tagger.Component wmeta workloadmeta.Component containerProvider proccontainers.ContainerProvider // containerProvider is used as a fallback to get a PID -> CID mapping when workloadmeta does not have the process data + pidToSlurm map[int]model.SlurmInfo // pidToSlurm maps PIDs to the owning Slurm job identity, resolved system-probe-side; nil/empty means no Slurm tagging. pidToCid map[int]string // pidToCid is the mapping of PIDs to container IDs, retrieved from the container provider until it is invalidated. telemetry *workloadTagCacheTelemetry // telemetry is the telemetry component for the workload tag cache } @@ -149,6 +151,36 @@ func (c *WorkloadTagCache) SetContainerProvider(p proccontainers.ContainerProvid c.pidToCid = nil } +// slurmInfoFromStats converts the system-probe GPU stats payload's PID -> Slurm identity map +// (keyed by host PID as uint32) into the int-keyed map used by the tag cache, keeping only +// processes that resolved to an actual Slurm job. Returns nil for a nil/empty payload so the +// common (no-Slurm) case allocates nothing. +func slurmInfoFromStats(stats *model.GPUStats) map[int]model.SlurmInfo { + if stats == nil || len(stats.SlurmInfoByPID) == 0 { + return nil + } + + pidToSlurm := make(map[int]model.SlurmInfo, len(stats.SlurmInfoByPID)) + for pid, info := range stats.SlurmInfoByPID { + if info.JobID == "" { + continue + } + pidToSlurm[int(pid)] = info + } + if len(pidToSlurm) == 0 { + return nil + } + return pidToSlurm +} + +// SetSlurmInfo replaces the PID -> Slurm job identity mapping for the current run. The mapping is +// built from the system-probe GPU stats payload (which resolves Slurm identity where it already +// holds SYS_PTRACE). An empty or nil map means no process has a resolved Slurm identity, so no +// Slurm tags are added. +func (c *WorkloadTagCache) SetSlurmInfo(pidToSlurm map[int]model.SlurmInfo) { + c.pidToSlurm = pidToSlurm +} + // MarkStale marks all entries in the cache as stale. That way, on the next calls to GetWorkloadTags, we will // try to rebuild them, anf if we can't we will return stale data. func (c *WorkloadTagCache) MarkStale() { @@ -263,6 +295,18 @@ func (c *WorkloadTagCache) buildProcessTags(processID string) ([]string, error) tags = append(tags, containerTags...) } + // Slurm job identity is resolved system-probe-side and delivered via the GPU stats payload + // (see SetSlurmInfo); the core agent only looks it up here, so no SYS_PTRACE is needed here. + if info, ok := c.pidToSlurm[int(pid)]; ok && info.JobID != "" { + tags = append(tags, "slurm_job_id:"+info.JobID) + if info.JobName != "" { + tags = append(tags, "slurm_job_name:"+info.JobName) + } + if info.Partition != "" { + tags = append(tags, "slurm_job_partition:"+info.Partition) + } + } + return tags, multiErr } diff --git a/pkg/collector/corechecks/gpu/tags_test.go b/pkg/collector/corechecks/gpu/tags_test.go index 8e2901b64506..538af0b08a71 100644 --- a/pkg/collector/corechecks/gpu/tags_test.go +++ b/pkg/collector/corechecks/gpu/tags_test.go @@ -12,6 +12,7 @@ import ( "math/rand" "slices" "strconv" + "strings" "testing" "time" @@ -25,6 +26,7 @@ import ( "github.com/DataDog/datadog-agent/comp/core/telemetry/def" workloadmeta "github.com/DataDog/datadog-agent/comp/core/workloadmeta/def" workloadmetamock "github.com/DataDog/datadog-agent/comp/core/workloadmeta/mock" + "github.com/DataDog/datadog-agent/pkg/collector/corechecks/gpu/model" agenterrors "github.com/DataDog/datadog-agent/pkg/errors" "github.com/DataDog/datadog-agent/pkg/gpu/testutil" mock_containers "github.com/DataDog/datadog-agent/pkg/process/util/containers/mocks" @@ -1292,3 +1294,94 @@ func TestNewWorkloadTagCache_DuplicateSubsystemPanics(t *testing.T) { _, _ = NewWorkloadTagCacheWithSubsystem("dup", tg, wmeta, cp, tm, defaultCacheSize) }, "constructing a second cache with the same subsystem must panic on Prometheus duplicate registration") } + +func TestBuildProcessTags_SlurmJobIDAdditive(t *testing.T) { + cache, mocks := setupWorkloadTagCache(t) + + pid := int32(4242) + mocks.workloadMeta.Set(&workloadmeta.Process{ + EntityID: workloadmeta.EntityID{Kind: workloadmeta.KindProcess, ID: strconv.FormatInt(int64(pid), 10)}, + NsPid: pid, + Owner: nil, + }) + // containerID="" triggers the existing container-provider fallback path; slurm tagging must + // still happen additively regardless of what that fallback returns. + mocks.containerProvider.EXPECT().GetPidToCid(time.Duration(0)).Return(map[int]string{}) + + cache.SetSlurmInfo(map[int]model.SlurmInfo{ + int(pid): {JobID: "3", JobName: "gpuhold", Partition: "gpu"}, + }) + + tags, err := cache.buildProcessTags(strconv.FormatInt(int64(pid), 10)) + require.NoError(t, err) + assert.Contains(t, tags, "slurm_job_id:3") + assert.Contains(t, tags, "slurm_job_name:gpuhold") + assert.Contains(t, tags, "slurm_job_partition:gpu") +} + +func TestBuildProcessTags_SlurmPidNotInMap(t *testing.T) { + cache, mocks := setupWorkloadTagCache(t) + + // The Slurm map holds a different PID, so the process under test resolves no Slurm identity. + // This is the normal "not a Slurm job" case: no slurm_* tags, no error. + pid := int32(8686) + mocks.workloadMeta.Set(&workloadmeta.Process{ + EntityID: workloadmeta.EntityID{Kind: workloadmeta.KindProcess, ID: strconv.FormatInt(int64(pid), 10)}, + NsPid: pid, + Owner: nil, + }) + mocks.containerProvider.EXPECT().GetPidToCid(time.Duration(0)).Return(map[int]string{}) + + cache.SetSlurmInfo(map[int]model.SlurmInfo{ + 9999: {JobID: "1", JobName: "other", Partition: "gpu"}, + }) + + tags, err := cache.buildProcessTags(strconv.FormatInt(int64(pid), 10)) + require.NoError(t, err) + for _, tag := range tags { + assert.False(t, strings.HasPrefix(tag, "slurm_")) + } +} + +func TestBuildProcessTags_SlurmInfoNilByDefault(t *testing.T) { + cache, mocks := setupWorkloadTagCache(t) + // SetSlurmInfo is never called: behavior must be identical to today (no slurm_* tags, + // no panics) — a nil map lookup is safe. + + pid := int32(6464) + mocks.workloadMeta.Set(&workloadmeta.Process{ + EntityID: workloadmeta.EntityID{Kind: workloadmeta.KindProcess, ID: strconv.FormatInt(int64(pid), 10)}, + NsPid: pid, + Owner: nil, + }) + mocks.containerProvider.EXPECT().GetPidToCid(time.Duration(0)).Return(map[int]string{}) + + tags, err := cache.buildProcessTags(strconv.FormatInt(int64(pid), 10)) + require.NoError(t, err) + for _, tag := range tags { + assert.False(t, strings.HasPrefix(tag, "slurm_")) + } +} + +func TestSlurmInfoFromStats(t *testing.T) { + // nil payload -> nil map, no allocation, no panic. + assert.Nil(t, slurmInfoFromStats(nil)) + + // Empty/absent SlurmInfoByPID -> nil map. + assert.Nil(t, slurmInfoFromStats(&model.GPUStats{})) + + // A map with only empty-JobID entries -> nil map. + assert.Nil(t, slurmInfoFromStats(&model.GPUStats{SlurmInfoByPID: map[uint32]model.SlurmInfo{ + 2: {JobName: "no-id"}, + }})) + + // Only entries with a non-empty JobID are kept; keys are converted uint32 -> int. + stats := &model.GPUStats{SlurmInfoByPID: map[uint32]model.SlurmInfo{ + 10: {JobID: "7", JobName: "a", Partition: "gpu"}, + 20: {}, + }} + got := slurmInfoFromStats(stats) + assert.Equal(t, map[int]model.SlurmInfo{ + 10: {JobID: "7", JobName: "a", Partition: "gpu"}, + }, got) +} diff --git a/pkg/config/schema/yaml/system-probe_schema.yaml b/pkg/config/schema/yaml/system-probe_schema.yaml index a595c3626583..35944eef0f4d 100644 --- a/pkg/config/schema/yaml/system-probe_schema.yaml +++ b/pkg/config/schema/yaml/system-probe_schema.yaml @@ -2444,6 +2444,10 @@ properties: node_type: setting type: boolean default: false + enable_slurm_job_tagging: + node_type: setting + type: boolean + default: false fatbin_request_queue_size: node_type: setting type: integer diff --git a/pkg/config/setup/system_probe_settings.go b/pkg/config/setup/system_probe_settings.go index c5b778a19e75..2746cc3e6fcb 100644 --- a/pkg/config/setup/system_probe_settings.go +++ b/pkg/config/setup/system_probe_settings.go @@ -364,6 +364,7 @@ func initMainSystemProbeConfig(cfg pkgconfigmodel.Setup) { cfg.BindEnvAndSetDefault("gpu_monitoring.configure_cgroup_perms", false) cfg.BindEnvAndSetDefault("gpu_monitoring.prm_endpoint_enabled", true) cfg.BindEnvAndSetDefault("gpu_monitoring.enable_fatbin_parsing", false) + cfg.BindEnvAndSetDefault("gpu_monitoring.enable_slurm_job_tagging", false) cfg.BindEnvAndSetDefault("gpu_monitoring.fatbin_request_queue_size", 100) cfg.BindEnvAndSetDefault("gpu_monitoring.ring_buffer_pages_per_device", 32) // 32 pages = 128KB by default per device cfg.BindEnvAndSetDefault("gpu_monitoring.ringbuffer_wakeup_size", 3000) // 3000 bytes is about ~10-20 events depending on the specific type diff --git a/pkg/gpu/BUILD.bazel b/pkg/gpu/BUILD.bazel index e39582d00402..33df75e86f0b 100644 --- a/pkg/gpu/BUILD.bazel +++ b/pkg/gpu/BUILD.bazel @@ -38,6 +38,7 @@ go_library( "//pkg/gpu/safenvml", "//pkg/network/usm/sharedlibraries", "//pkg/network/usm/utils", + "//pkg/process/util/slurm", "//pkg/security/utils/lru/simplelru", "//pkg/status/health", "//pkg/system-probe/config", @@ -97,6 +98,7 @@ go_test( "//pkg/gpu/safenvml", "//pkg/gpu/safenvml/testutil", "//pkg/gpu/testutil", + "//pkg/process/util/slurm", "//pkg/util/gpu", "//pkg/util/kernel", "@com_github_stretchr_testify//assert", diff --git a/pkg/gpu/config/config.go b/pkg/gpu/config/config.go index 3f2b58d5f869..a1360518e932 100644 --- a/pkg/gpu/config/config.go +++ b/pkg/gpu/config/config.go @@ -36,6 +36,9 @@ type Config struct { ConfigureCgroupPerms bool // EnableFatbinParsing indicates whether the probe should enable fatbin parsing. EnableFatbinParsing bool + // EnableSlurmJobTagging indicates whether the probe should resolve the owning Slurm job + // identity for each GPU process (from /proc//environ) and attach it to process metrics. + EnableSlurmJobTagging bool // KernelCacheQueueSize is the size of the kernel cache queue for parsing requests KernelCacheQueueSize int // RingBufferSizePagesPerDevice is the number of pages to use for the ring buffer per device. @@ -86,6 +89,7 @@ func New() *Config { PRMEndpointEnabled: spCfg.GetBool(sysconfig.FullKeyPath(consts.GPUNS, "prm_endpoint_enabled")), ConfigureCgroupPerms: spCfg.GetBool(sysconfig.FullKeyPath(consts.GPUNS, "configure_cgroup_perms")), EnableFatbinParsing: spCfg.GetBool(sysconfig.FullKeyPath(consts.GPUNS, "enable_fatbin_parsing")), + EnableSlurmJobTagging: spCfg.GetBool(sysconfig.FullKeyPath(consts.GPUNS, "enable_slurm_job_tagging")), KernelCacheQueueSize: spCfg.GetInt(sysconfig.FullKeyPath(consts.GPUNS, "fatbin_request_queue_size")), RingBufferSizePagesPerDevice: spCfg.GetInt(sysconfig.FullKeyPath(consts.GPUNS, "ring_buffer_pages_per_device")), RingBufferWakeupSize: spCfg.GetInt(sysconfig.FullKeyPath(consts.GPUNS, "ringbuffer_wakeup_size")), diff --git a/pkg/gpu/context.go b/pkg/gpu/context.go index 790451b79eda..48f6da8ba2b7 100644 --- a/pkg/gpu/context.go +++ b/pkg/gpu/context.go @@ -18,11 +18,13 @@ import ( "github.com/DataDog/datadog-agent/comp/core/telemetry/def" workloadmeta "github.com/DataDog/datadog-agent/comp/core/workloadmeta/def" + "github.com/DataDog/datadog-agent/pkg/collector/corechecks/gpu/model" dderrors "github.com/DataDog/datadog-agent/pkg/errors" "github.com/DataDog/datadog-agent/pkg/gpu/config" "github.com/DataDog/datadog-agent/pkg/gpu/containers" "github.com/DataDog/datadog-agent/pkg/gpu/cuda" ddnvml "github.com/DataDog/datadog-agent/pkg/gpu/safenvml" + slurmutil "github.com/DataDog/datadog-agent/pkg/process/util/slurm" "github.com/DataDog/datadog-agent/pkg/util/kernel" "github.com/DataDog/datadog-agent/pkg/util/ktime" "github.com/DataDog/datadog-agent/pkg/util/log" @@ -67,6 +69,10 @@ type systemContext struct { // deviceCacheRefreshInterval is the minimum time passing between device cache refreshes deviceCacheRefreshInterval time.Duration + + // slurmProvider resolves the owning Slurm job identity for a process. It is nil unless + // gpu_monitoring.enable_slurm_job_tagging is set; otherwise Slurm tagging is disabled. + slurmProvider slurmutil.Provider } type systemContextOptions struct { @@ -134,6 +140,10 @@ func getSystemContext(optList ...systemContextOption) (*systemContext, error) { lastDeviceCacheRefreshTime: time.Now(), } + if opts.config.EnableSlurmJobTagging { + ctx.slurmProvider = slurmutil.InitSharedProvider() + } + var err error ctx.timeResolver, err = ktime.NewResolver() if err != nil { @@ -161,6 +171,74 @@ func (ctx *systemContext) removeProcess(pid int) { } } +// getSlurmInfo returns the owning Slurm job identity for pid, resolved from /proc//environ. +// Returns the zero value when Slurm job tagging is disabled or the process is not a Slurm job. +// A permission error (the agent missing SYS_PTRACE) is logged, rate-limited, as it is the only +// signal distinguishing a misconfiguration from a process that simply is not a Slurm job. +// +// The result is intentionally not cached: there are only a handful of GPU processes per node and +// this runs once per stats-flush interval, so a small environ read per process is cheaper than a +// shared cross-goroutine cache (the stats-flush and consumer goroutines would both touch it). +func (ctx *systemContext) getSlurmInfo(pid uint32) model.SlurmInfo { + if ctx.slurmProvider == nil { + return model.SlurmInfo{} + } + + resolved, err := ctx.slurmProvider.GetSlurmInfo(int32(pid)) + if err != nil { + if logLimitProbe.ShouldLog() { + log.Warnf("could not resolve slurm info for pid %d: %v", pid, err) + } + return model.SlurmInfo{} + } + + return model.SlurmInfo{JobID: resolved.JobID, JobName: resolved.JobName, Partition: resolved.Partition} +} + +// resolveSlurmInfoForGPUProcesses returns the Slurm job identity for every GPU process, keyed by +// host PID. It resolves the union of the eBPF-tracked PIDs (from processMetrics) and the +// NVML-visible compute processes (device.GetComputeRunningProcesses). The NVML set is exactly the +// PID set the core agent's NVML collector emits process.memory.usage for -- including jobs that are +// resident but idle, or were already running when the eBPF probe attached -- so this gives precise, +// not accidental, coverage of every submitted GPU process metric. Resolving here keeps SYS_PTRACE +// confined to system-probe. Only a handful of processes hold GPU memory, so the per-PID environ +// read stays small. Returns nil when Slurm job tagging is disabled or no process is a Slurm job. +func (ctx *systemContext) resolveSlurmInfoForGPUProcesses(processMetrics []model.ProcessStatsTuple) map[uint32]model.SlurmInfo { + if ctx.slurmProvider == nil { + return nil + } + + pids := make(map[uint32]struct{}, len(processMetrics)) + for _, entry := range processMetrics { + pids[entry.Key.PID] = struct{}{} + } + + // NVML compute processes: the core agent's NVML collector emits process.memory.usage for these + // PIDs, so they must be resolvable even when eBPF has no active stream for them. + if devices, err := ctx.deviceCache.AllPhysicalDevices(); err == nil { + for _, dev := range devices { + procs, procErr := dev.GetComputeRunningProcesses() + if procErr != nil { + continue + } + for _, p := range procs { + pids[p.Pid] = struct{}{} + } + } + } + + result := make(map[uint32]model.SlurmInfo, len(pids)) + for pid := range pids { + if info := ctx.getSlurmInfo(pid); info.JobID != "" { + result[pid] = info + } + } + if len(result) == 0 { + return nil + } + return result +} + // filterDevicesForContainer filters the available GPU devices for the given // container. If the ID is not empty, we check the assignment of GPU resources // to the container and return only the devices that are available to the diff --git a/pkg/gpu/context_test.go b/pkg/gpu/context_test.go index bd24b786a6b6..9b13394b3134 100644 --- a/pkg/gpu/context_test.go +++ b/pkg/gpu/context_test.go @@ -8,6 +8,7 @@ package gpu import ( + "errors" "strconv" "strings" "testing" @@ -15,9 +16,11 @@ import ( "github.com/stretchr/testify/require" workloadmeta "github.com/DataDog/datadog-agent/comp/core/workloadmeta/def" + "github.com/DataDog/datadog-agent/pkg/collector/corechecks/gpu/model" ddnvml "github.com/DataDog/datadog-agent/pkg/gpu/safenvml" nvmltestutil "github.com/DataDog/datadog-agent/pkg/gpu/safenvml/testutil" "github.com/DataDog/datadog-agent/pkg/gpu/testutil" + slurmutil "github.com/DataDog/datadog-agent/pkg/process/util/slurm" gpuutil "github.com/DataDog/datadog-agent/pkg/util/gpu" "github.com/DataDog/datadog-agent/pkg/util/kernel" ) @@ -287,3 +290,86 @@ func TestGetCurrentActiveGpuDevice(t *testing.T) { }) } } + +// fakeSlurmProvider is a test slurm.Provider that returns canned data without touching /proc. +// If infoByPID is set it is consulted per PID (unknown PIDs resolve to the zero value); otherwise +// the flat info/err are returned for every PID. +type fakeSlurmProvider struct { + info slurmutil.SlurmInfo + err error + infoByPID map[int32]slurmutil.SlurmInfo + calls int +} + +func (f *fakeSlurmProvider) GetSlurmInfo(pid int32) (slurmutil.SlurmInfo, error) { + f.calls++ + if f.infoByPID != nil { + return f.infoByPID[pid], nil + } + return f.info, f.err +} + +func TestGetSlurmInfo(t *testing.T) { + t.Run("nil provider returns zero value", func(t *testing.T) { + sysCtx := getTestSystemContext(t) + require.Nil(t, sysCtx.slurmProvider) + require.Equal(t, model.SlurmInfo{}, sysCtx.getSlurmInfo(1234)) + }) + + t.Run("resolves per pid on every call", func(t *testing.T) { + sysCtx := getTestSystemContext(t) + provider := &fakeSlurmProvider{info: slurmutil.SlurmInfo{JobID: "7", JobName: "gpuhold", Partition: "gpu"}} + sysCtx.slurmProvider = provider + + want := model.SlurmInfo{JobID: "7", JobName: "gpuhold", Partition: "gpu"} + require.Equal(t, want, sysCtx.getSlurmInfo(4242)) + require.Equal(t, want, sysCtx.getSlurmInfo(4242)) + // No caching: resolution is intentionally stateless, so the provider is hit every call. + require.Equal(t, 2, provider.calls) + }) + + t.Run("permission error returns zero value", func(t *testing.T) { + sysCtx := getTestSystemContext(t) + provider := &fakeSlurmProvider{err: errors.New("permission denied")} + sysCtx.slurmProvider = provider + + require.Equal(t, model.SlurmInfo{}, sysCtx.getSlurmInfo(8686)) + require.Equal(t, model.SlurmInfo{}, sysCtx.getSlurmInfo(8686)) + require.Equal(t, 2, provider.calls) + }) +} + +func TestResolveSlurmInfoForGPUProcesses(t *testing.T) { + // PID 10 is an eBPF-tracked process (in processMetrics); PID 20 is tracked but not a Slurm job. + // NVML-visible PIDs come from the mock device cache and are keyed by the same provider map, so a + // PID that is not a Slurm job (or not present) is excluded regardless of which source found it. + procMetrics := []model.ProcessStatsTuple{ + {Key: model.ProcessStatsKey{PID: 10}}, + {Key: model.ProcessStatsKey{PID: 20}}, + } + + t.Run("nil provider returns nil (tagging disabled)", func(t *testing.T) { + sysCtx := getTestSystemContext(t) + require.Nil(t, sysCtx.resolveSlurmInfoForGPUProcesses(procMetrics)) + }) + + t.Run("resolves GPU-process PIDs, drops non-Slurm ones", func(t *testing.T) { + sysCtx := getTestSystemContext(t) + // Only PID 10 is a Slurm job; PID 20 (and any NVML-enumerated PID from the mock device + // cache) resolves to the zero value and must be excluded. + sysCtx.slurmProvider = &fakeSlurmProvider{infoByPID: map[int32]slurmutil.SlurmInfo{ + 10: {JobID: "42", JobName: "train", Partition: "gpu"}, + }} + + got := sysCtx.resolveSlurmInfoForGPUProcesses(procMetrics) + require.Equal(t, model.SlurmInfo{JobID: "42", JobName: "train", Partition: "gpu"}, got[10]) + _, has20 := got[20] + require.False(t, has20, "a process that is not a Slurm job must not appear in the map") + }) + + t.Run("no Slurm jobs returns nil", func(t *testing.T) { + sysCtx := getTestSystemContext(t) + sysCtx.slurmProvider = &fakeSlurmProvider{infoByPID: map[int32]slurmutil.SlurmInfo{}} + require.Nil(t, sysCtx.resolveSlurmInfoForGPUProcesses(procMetrics)) + }) +} diff --git a/pkg/gpu/stats.go b/pkg/gpu/stats.go index f96549e443d3..b3141b312e3b 100644 --- a/pkg/gpu/stats.go +++ b/pkg/gpu/stats.go @@ -120,6 +120,7 @@ func (g *statsGenerator) getStats(nowKtime int64) (*model.GPUStats, error) { stats := &model.GPUStats{ ProcessMetrics: processMetrics, DeviceMetrics: deviceMetrics, + SlurmInfoByPID: g.sysCtx.resolveSlurmInfoForGPUProcesses(processMetrics), } g.telemetry.aggregators.Set(float64(len(g.aggregators))) diff --git a/pkg/process/util/slurm/BUILD.bazel b/pkg/process/util/slurm/BUILD.bazel new file mode 100644 index 000000000000..e1d2c3d7b73e --- /dev/null +++ b/pkg/process/util/slurm/BUILD.bazel @@ -0,0 +1,37 @@ +load("@rules_go//go:def.bzl", "go_library", "go_test") + +go_library( + name = "slurm", + srcs = ["slurm.go"], + importpath = "github.com/DataDog/datadog-agent/pkg/process/util/slurm", + visibility = ["//visibility:public"], + deps = select({ + "@rules_go//go/platform:android": [ + "//pkg/security/utils", + ], + "@rules_go//go/platform:linux": [ + "//pkg/security/utils", + ], + "//conditions:default": [], + }), +) + +go_test( + name = "slurm_test", + srcs = ["slurm_test.go"], + embed = [":slurm"], + gotags = ["test"], + deps = select({ + "@rules_go//go/platform:android": [ + "//pkg/util/kernel", + "@com_github_stretchr_testify//assert", + "@com_github_stretchr_testify//require", + ], + "@rules_go//go/platform:linux": [ + "//pkg/util/kernel", + "@com_github_stretchr_testify//assert", + "@com_github_stretchr_testify//require", + ], + "//conditions:default": [], + }), +) diff --git a/pkg/process/util/slurm/slurm.go b/pkg/process/util/slurm/slurm.go new file mode 100644 index 000000000000..2fd572218b9f --- /dev/null +++ b/pkg/process/util/slurm/slurm.go @@ -0,0 +1,104 @@ +// Unless explicitly stated otherwise all files in this repository are licensed +// under the Apache License Version 2.0. +// This product includes software developed at Datadog (https://www.datadoghq.com/). +// Copyright 2026-present Datadog, Inc. + +//go:build linux + +// Package slurm resolves Slurm job identity for a process from /proc//environ. Slurm sets +// SLURM_JOB_ID (and related SLURM_JOB_* vars) unconditionally at task launch, so reading them +// back needs only the SYS_PTRACE capability -- no slurm binaries, no munge, and no hostPID. +package slurm + +import ( + "fmt" + "os" + "strings" + "sync" + + secutils "github.com/DataDog/datadog-agent/pkg/security/utils" +) + +// SlurmInfo is the Slurm job identity resolved for a PID. +type SlurmInfo struct { + // JobID is the numeric Slurm job ID, from SLURM_JOB_ID. Empty if unresolved. + JobID string + // JobName is the Slurm job name, from SLURM_JOB_NAME. Only set when JobID is set. + JobName string + // Partition is the Slurm partition, from SLURM_JOB_PARTITION. Only set when JobID is set. + Partition string +} + +// Provider resolves Slurm job identity for a PID. +type Provider interface { + // GetSlurmInfo returns the Slurm job identity for pid. A zero-value SlurmInfo with a nil + // error means the process is not a Slurm job (expected, not an error). A non-nil error + // means the read was denied (e.g. the agent is missing the SYS_PTRACE capability needed to + // read another process's /proc//environ) and callers should treat it as a + // misconfiguration signal rather than a normal "not a Slurm job" miss. + GetSlurmInfo(pid int32) (SlurmInfo, error) +} + +var ( + initOnce sync.Once + shared Provider +) + +// InitSharedProvider initializes and returns the shared Provider singleton. Safe to call +// multiple times; only the first call constructs the provider. +func InitSharedProvider() Provider { + initOnce.Do(func() { + shared = &procProvider{} + }) + return shared +} + +// GetSharedProvider returns the shared Provider singleton, or nil if InitSharedProvider was +// never called. +func GetSharedProvider() Provider { + return shared +} + +type procProvider struct{} + +const ( + envJobIDPrefix = "SLURM_JOB_ID=" + envJobNamePrefix = "SLURM_JOB_NAME=" + envPartitionPrefix = "SLURM_JOB_PARTITION=" +) + +// GetSlurmInfo implements Provider. +func (*procProvider) GetSlurmInfo(pid int32) (SlurmInfo, error) { + var info SlurmInfo + + // maxEnvVars=0: EnvVars always collects every priority-prefixed var regardless of this cap + // (the cap only bounds the second pass over non-matching vars), so passing 0 returns just + // the SLURM_-prefixed entries instead of materializing the whole process environment. + envs, _, envErr := secutils.EnvVars([]string{"SLURM_"}, uint32(pid), 0) + if envErr == nil { + for _, e := range envs { + switch { + case strings.HasPrefix(e, envJobIDPrefix): + info.JobID = strings.TrimPrefix(e, envJobIDPrefix) + case strings.HasPrefix(e, envJobNamePrefix): + info.JobName = strings.TrimPrefix(e, envJobNamePrefix) + case strings.HasPrefix(e, envPartitionPrefix): + info.Partition = strings.TrimPrefix(e, envPartitionPrefix) + } + } + } + + if info.JobID != "" { + return info, nil + } + + // No job ID resolved. Distinguish "not a Slurm job" (expected, silent) from + // "misconfigured, can't read /proc" (a real problem callers should count) — the latter + // requires SYS_PTRACE in the agent's securityContext to read another process's + // /proc//environ. + if os.IsPermission(envErr) { + return SlurmInfo{}, fmt.Errorf("permission denied resolving slurm info for pid %d: %w", pid, envErr) + } + + return SlurmInfo{}, nil +} diff --git a/pkg/process/util/slurm/slurm_test.go b/pkg/process/util/slurm/slurm_test.go new file mode 100644 index 000000000000..f29c0ffe767e --- /dev/null +++ b/pkg/process/util/slurm/slurm_test.go @@ -0,0 +1,79 @@ +// Unless explicitly stated otherwise all files in this repository are licensed +// under the Apache License Version 2.0. +// This product includes software developed at Datadog (https://www.datadoghq.com/). +// Copyright 2026-present Datadog, Inc. + +//go:build linux && test + +package slurm + +import ( + "os" + "path/filepath" + "strconv" + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/DataDog/datadog-agent/pkg/util/kernel" +) + +// writeFakeEnviron writes a /proc//environ file (null-separated) under procRoot. +func writeFakeEnviron(tb testing.TB, procRoot string, pid int, vars []string) { + pidDir := filepath.Join(procRoot, strconv.Itoa(pid)) + require.NoError(tb, os.MkdirAll(pidDir, 0755)) + content := strings.Join(vars, "\x00") + "\x00" + require.NoError(tb, os.WriteFile(filepath.Join(pidDir, "environ"), []byte(content), 0644)) +} + +func TestGetSlurmInfo_EnvironPrimarySignal(t *testing.T) { + procRoot := t.TempDir() + kernel.WithFakeProcFS(t, procRoot) + + writeFakeEnviron(t, procRoot, 4242, []string{ + "PATH=/usr/bin", + "SLURM_JOB_ID=3", + "SLURM_JOB_NAME=gpuhold", + "SLURM_JOB_PARTITION=gpu", + "SLURM_VERSION=26.05.1", + }) + + p := &procProvider{} + info, err := p.GetSlurmInfo(4242) + require.NoError(t, err) + assert.Equal(t, "3", info.JobID) + assert.Equal(t, "gpuhold", info.JobName) + assert.Equal(t, "gpu", info.Partition) +} + +func TestGetSlurmInfo_NotASlurmJob(t *testing.T) { + procRoot := t.TempDir() + kernel.WithFakeProcFS(t, procRoot) + + writeFakeEnviron(t, procRoot, 6464, []string{"PATH=/usr/bin"}) + + p := &procProvider{} + info, err := p.GetSlurmInfo(6464) + require.NoError(t, err) // absence is not an error + assert.Empty(t, info.JobID) +} + +func TestGetSlurmInfo_PermissionDeniedIsAnError(t *testing.T) { + if os.Geteuid() == 0 { + t.Skip("chmod-based permission denial doesn't apply when running as root") + } + + procRoot := t.TempDir() + kernel.WithFakeProcFS(t, procRoot) + + pidDir := filepath.Join(procRoot, "7575") + require.NoError(t, os.MkdirAll(pidDir, 0755)) + // environ exists but is unreadable, so the only signal (SYS_PTRACE-gated) fails permission. + require.NoError(t, os.WriteFile(filepath.Join(pidDir, "environ"), []byte("x"), 0000)) + + p := &procProvider{} + _, err := p.GetSlurmInfo(7575) + require.Error(t, err) +} diff --git a/releasenotes/notes/slurm-job-gpu-tagging-8f3c1a7e9d2b4051.yaml b/releasenotes/notes/slurm-job-gpu-tagging-8f3c1a7e9d2b4051.yaml new file mode 100644 index 000000000000..84d0d7fac916 --- /dev/null +++ b/releasenotes/notes/slurm-job-gpu-tagging-8f3c1a7e9d2b4051.yaml @@ -0,0 +1,16 @@ +# Each section from every release note are combined when the +# CHANGELOG.rst is rendered. So the text needs to be worded so that +# it does not depend on any information only available in another +# section. This may mean repeating some details, but each section +# must be readable independently of the other. +# +# Each section note must be formatted as reStructuredText. +--- +features: + - | + GPU: Added an opt-in ``gpu_monitoring.enable_slurm_job_tagging`` system-probe config option + that tags GPU-process metrics with the owning Slurm job's identity (``slurm_job_id``, + ``slurm_job_name``, ``slurm_job_partition``) on Slurm-on-Kubernetes deployments. The identity + is resolved by system-probe from ``/proc//environ`` (reusing the ``SYS_PTRACE`` capability + it already holds) and delivered to the GPU check over the existing system-probe socket, so no + additional capability is granted to the core agent.