Add Slurm-job GPU tagging (system probe)

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 {

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/<pid>/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"`
 }

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
 }
 

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)
+}

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

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

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",

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/<pid>/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")),