Repository: DataDog/datadog-agent
Pull request #34998: [EBPF] gpu: refactor stream collection
Author: login: gjulianm
Labels: name: changelog/no-changelog
description: No changelog entry needed
name: team/ebpf-platform
name: qa/done
description: QA done before merge and regressions are covered by tests
name: long review
description: PR is complex, plan time to review it
Created: 2025-03-11T12:08:27Z | Merged: 2025-03-20T11:51:00Z

<!--
* Contributors are encouraged to read our [CONTRIBUTING](/CONTRIBUTING.md) documentation.
* Both Contributor and Reviewer Checklists are available at https://datadoghq.dev/datadog-agent/guidelines/contributing/#pull-requests.
* The pull request:
  * Should only fix one issue or add one feature at a time.
  * Must update the test suite for the relevant functionality.
  * Should pass all status checks before being reviewed or merged.
* Commit titles should be prefixed with general area of pull request's change.
* Please fill the below sections if possible with relevant information or links.
-->
### What does this PR do?

This PR introduces a `StreamCollection` structure that manages the stream lifetimes, moving that code away from the consumer into its own structure. This brings in several changes:

- Global and non-global streams are now handled transparently by the collection itself. The stream key is now an internal structure, and all the stream metadata is filled in the `streamMetadata` structure, common to all types of streams.
- Container ID is only retrieved when needed using memoized functions. This means that instead of computing it every time we retrieved a stream, it will only be computed once on creation, in order to fill the metadata and/or the visible devices cache. As this is an expensive operation involving a regex, this will bring a significant performance improvement, as the regex part was seen using 10-20% of CPU in some profiles in staging clusters.
- Telemetry for the handlers is moved to the StreamCollection type.
- Handling of finished streams is now entirely owned by `StreamCollection`. Process and stream lifetime still need improvements but those can be done in a separate PR.

### Motivation

[Refactoring docs](https://docs.google.com/document/d/1rpmxZQcpa2U5y1WqA2EipbflecZDlealDsb8-q_ondc/edit?tab=t.0#heading=h.tv4upu6c82xe). Improve correctness of the stream handling and ensure it's owned by a single struct.

### Describe how you validated your changes
<!--
Validate your changes before merge, ensuring that:
* Your PR is tested by static / unit / integrations / e2e tests
* Your PR description details which e2e tests cover your changes, if any
* The PR description contains details of how you validated your changes. If you validated changes manually and not through automated tests, add context on why automated tests did not fit your changes validation.

If you want additional validation by a second person, you can ask reviewers to do it. Describe how to set up an environment for manual tests in the PR description. Manual validation is expected to happen on every commit before merge.

Any manual validation step should then map to an automated test. Manual validation should not substitute automation, minus exceptions not supported by test tooling yet.
-->

E2E and unit tests passing.

### Possible Drawbacks / Trade-offs

### Additional Notes
<!--
* Anything else we should know when reviewing?
* Include benchmarking information here whenever possible.
* Include info about alternatives that were considered and why the proposed
  version was chosen.
-->


Diff:
diff --git a/pkg/gpu/consumer.go b/pkg/gpu/consumer.go
index b17dee804c8e..fbc42823e19b 100644
--- a/pkg/gpu/consumer.go
+++ b/pkg/gpu/consumer.go
@@ -14,16 +14,12 @@ import (
 	"time"
 	"unsafe"
 
-	"github.com/NVIDIA/go-nvml/pkg/nvml"
-	"golang.org/x/sys/unix"
-
 	"github.com/DataDog/datadog-agent/comp/core/telemetry"
 	ddebpf "github.com/DataDog/datadog-agent/pkg/ebpf"
 	"github.com/DataDog/datadog-agent/pkg/gpu/config"
 	gpuebpf "github.com/DataDog/datadog-agent/pkg/gpu/ebpf"
 	"github.com/DataDog/datadog-agent/pkg/process/monitor"
 	"github.com/DataDog/datadog-agent/pkg/status/health"
-	"github.com/DataDog/datadog-agent/pkg/util/cgroups"
 	"github.com/DataDog/datadog-agent/pkg/util/kernel"
 	"github.com/DataDog/datadog-agent/pkg/util/log"
 )
@@ -33,48 +29,36 @@ const telemetryEventErrorUnknownType = "unknown_type"
 const telemetryEventTypeUnknown = "unknown"
 const telemetryEventHeader = "header"
 
-type nonGlobalStreamKey struct {
-	pid    uint32
-	stream uint64
-}
-
 // cudaEventConsumer is responsible for consuming CUDA events from the eBPF probe, and delivering them
 // to the appropriate stream handler.
 type cudaEventConsumer struct {
-	eventHandler        ddebpf.EventHandler
-	once                sync.Once
-	closed              chan struct{}
-	streamHandlers      map[streamKey]*StreamHandler
-	nonGlobalStreamKeys map[nonGlobalStreamKey]streamKey // TODO: Move this to a separate class that deals with stream assignment
-	wg                  sync.WaitGroup
-	running             atomic.Bool
-	sysCtx              *systemContext
-	cfg                 *config.Config
-	telemetry           *cudaEventConsumerTelemetry
-	debugCollector      *eventCollector
+	eventHandler   ddebpf.EventHandler
+	once           sync.Once
+	closed         chan struct{}
+	streamHandlers *streamCollection
+	wg             sync.WaitGroup
+	running        atomic.Bool
+	sysCtx         *systemContext
+	cfg            *config.Config
+	telemetry      *cudaEventConsumerTelemetry
+	debugCollector *eventCollector
 }
 
 type cudaEventConsumerTelemetry struct {
-	activeHandlers     telemetry.Gauge
-	removedHandlers    telemetry.Counter
-	events             telemetry.Counter
-	eventErrors        telemetry.Counter
-	finalizedProcesses telemetry.Counter
-	missingContainers  telemetry.Counter
-	missingDevices     telemetry.Counter
+	events      telemetry.Counter
+	eventErrors telemetry.Counter
 }
 
 // newCudaEventConsumer creates a new CUDA event consumer.
-func newCudaEventConsumer(sysCtx *systemContext, eventHandler ddebpf.EventHandler, cfg *config.Config, telemetry telemetry.Component) *cudaEventConsumer {
+func newCudaEventConsumer(sysCtx *systemContext, streamHandlers *streamCollection, eventHandler ddebpf.EventHandler, cfg *config.Config, telemetry telemetry.Component) *cudaEventConsumer {
 	return &cudaEventConsumer{
-		eventHandler:        eventHandler,
-		closed:              make(chan struct{}),
-		streamHandlers:      make(map[streamKey]*StreamHandler),
-		nonGlobalStreamKeys: make(map[nonGlobalStreamKey]streamKey),
-		cfg:                 cfg,
-		sysCtx:              sysCtx,
-		telemetry:           newCudaEventConsumerTelemetry(telemetry),
-		debugCollector:      newEventCollector(),
+		eventHandler:   eventHandler,
+		closed:         make(chan struct{}),
+		cfg:            cfg,
+		sysCtx:         sysCtx,
+		streamHandlers: streamHandlers,
+		telemetry:      newCudaEventConsumerTelemetry(telemetry),
+		debugCollector: newEventCollector(),
 	}
 }
 
@@ -82,13 +66,8 @@ func newCudaEventConsumerTelemetry(tm telemetry.Component) *cudaEventConsumerTel
 	subsystem := gpuTelemetryModule + "__consumer"
 
 	return &cudaEventConsumerTelemetry{
-		activeHandlers:     tm.NewGauge(subsystem, "active_handlers", nil, "Number of active stream handlers"),
-		removedHandlers:    tm.NewCounter(subsystem, "removed_handlers", nil, "Number of removed stream handlers"),
-		events:             tm.NewCounter(subsystem, "events", []string{"event_type"}, "Number of processed CUDA events received by the consumer"),
-		eventErrors:        tm.NewCounter(subsystem, "events__errors", []string{"event_type", "error"}, "Number of CUDA events that couldn't be processed due to an error"),
-		finalizedProcesses: tm.NewCounter(subsystem, "finalized_processes", nil, "Number of finalized processes"),
-		missingContainers:  tm.NewCounter(subsystem, "missing_containers", []string{"reason"}, "Number of missing containers"),
-		missingDevices:     tm.NewCounter(subsystem, "missing_devices", nil, "Number of failures to get GPU devices for a stream"),
+		events:      tm.NewCounter(subsystem, "events", []string{"event_type"}, "Number of processed CUDA events received by the consumer"),
+		eventErrors: tm.NewCounter(subsystem, "events__errors", []string{"event_type", "error"}, "Number of CUDA events that couldn't be processed due to an error"),
 	}
 }
 
@@ -199,7 +178,7 @@ func handleTypedEvent[K any](c *cudaEventConsumer, handler func(*K), eventType g
 
 func (c *cudaEventConsumer) handleStreamEvent(header *gpuebpf.CudaEventHeader, data unsafe.Pointer, dataLen int) error {
 	eventType := gpuebpf.CudaEventType(header.Type)
-	streamHandler, err := c.getStreamHandler(header)
+	streamHandler, err := c.streamHandlers.getStream(header)
 
 	if err != nil {
 		return fmt.Errorf("error getting stream handler: %w", err)
@@ -241,88 +220,7 @@ func (c *cudaEventConsumer) handleGlobalEvent(header *gpuebpf.CudaEventHeader, d
 }
 
 func (c *cudaEventConsumer) handleProcessExit(pid uint32) {
-	for key, handler := range c.streamHandlers {
-		if key.pid == pid {
-			log.Debugf("Process %d ended, marking stream %d as ended", pid, key.stream)
-			// the probe is responsible for deleting the stream handler
-			_ = handler.markEnd()
-			c.telemetry.finalizedProcesses.Inc()
-		}
-	}
-}
-
-func (c *cudaEventConsumer) getStreamKey(header *gpuebpf.CudaEventHeader) (streamKey, error) {
-	pid, tid := getPidTidFromHeader(header)
-
-	var nonGlobalKey nonGlobalStreamKey
-	if header.Stream_id != 0 {
-		// Non-global stream, check if we have created it before
-		nonGlobalKey.pid = pid
-		nonGlobalKey.stream = header.Stream_id
-
-		if key, ok := c.nonGlobalStreamKeys[nonGlobalKey]; ok {
-			return key, nil
-		}
-	}
-
-	cgroup := unix.ByteSliceToString(header.Cgroup[:])
-	containerID, err := cgroups.ContainerFilter("", cgroup)
-	if err != nil {
-		// We don't want to return an error here, as we can still process the event without the container ID
-		log.Warnf("error getting container ID for cgroup %s: %s", cgroup, err)
-		c.telemetry.missingContainers.Inc("error")
-	} else if containerID == "" {
-		c.telemetry.missingContainers.Inc("missing")
-	}
-
-	key := streamKey{
-		pid:         pid,
-		stream:      header.Stream_id,
-		gpuUUID:     "",
-		containerID: containerID,
-	}
-
-	gpuDevice, err := c.sysCtx.getCurrentActiveGpuDevice(int(pid), int(tid), containerID)
-	if err != nil {
-		c.telemetry.missingDevices.Inc()
-		return streamKey{}, fmt.Errorf("Error getting GPU device for process %d: %w", pid, err)
-	}
-
-	var ret nvml.Return
-	key.gpuUUID, ret = gpuDevice.GetUUID()
-	if ret != nvml.SUCCESS {
-		return streamKey{}, fmt.Errorf("Error getting GPU UUID for process %d: %v", pid, nvml.ErrorString(ret))
-	}
-
-	if header.Stream_id != 0 {
-		c.nonGlobalStreamKeys[nonGlobalKey] = key
-	}
-
-	return key, nil
-}
-
-func (c *cudaEventConsumer) getStreamHandler(header *gpuebpf.CudaEventHeader) (*StreamHandler, error) {
-	key, err := c.getStreamKey(header)
-	if err != nil {
-		return nil, err
-	}
-
-	if _, ok := c.streamHandlers[key]; !ok {
-		smVersion, ok := c.sysCtx.deviceSmVersions[key.gpuUUID]
-		if !ok {
-			if key.gpuUUID != "" {
-				// Only warn when we have a device, otherwise it's expected to not find the SM version
-				// if the device UUID is empty
-				log.Warnf("SM version not found for device %s, using default", key.gpuUUID)
-			}
-			smVersion = noSmVersion
-		}
-
-		c.streamHandlers[key] = newStreamHandler(key.pid, key.containerID, smVersion, c.sysCtx)
-		c.telemetry.activeHandlers.Set(float64(len(c.streamHandlers)))
-	}
-
-	return c.streamHandlers[key], nil
+	c.streamHandlers.markProcessStreamsAsEnded(pid)
 }
 
 func (c *cudaEventConsumer) checkClosedProcesses() {
@@ -332,24 +230,9 @@ func (c *cudaEventConsumer) checkClosedProcesses() {
 		return nil
 	})
 
-	for key, handler := range c.streamHandlers {
-		if _, ok := seenPIDs[key.pid]; !ok {
-			log.Debugf("Process %d ended, marking stream %d as ended", key.pid, key.stream)
-			_ = handler.markEnd()
-		}
-	}
-}
-
-func (c *cudaEventConsumer) cleanFinishedHandlers() {
-	for key, handler := range c.streamHandlers {
-		if handler.processEnded {
-			delete(c.streamHandlers, key)
-
-			if key.stream != 0 {
-				delete(c.nonGlobalStreamKeys, nonGlobalStreamKey{pid: key.pid, stream: key.stream})
-			}
+	for handler := range c.streamHandlers.allStreams() {
+		if _, ok := seenPIDs[handler.metadata.pid]; !ok {
+			c.streamHandlers.markProcessStreamsAsEnded(handler.metadata.pid)
 		}
 	}
-
-	c.telemetry.activeHandlers.Set(float64(len(c.streamHandlers)))
 }
diff --git a/pkg/gpu/consumer_test.go b/pkg/gpu/consumer_test.go
index 227b399df947..6093c7f67185 100644
--- a/pkg/gpu/consumer_test.go
+++ b/pkg/gpu/consumer_test.go
@@ -26,7 +26,8 @@ func TestConsumerCanStartAndStop(t *testing.T) {
 	cfg := config.New()
 	ctx, err := getSystemContext(testutil.GetBasicNvmlMock(), kernel.ProcFSRoot(), testutil.GetWorkloadMetaMock(t), testutil.GetTelemetryMock(t))
 	require.NoError(t, err)
-	consumer := newCudaEventConsumer(ctx, handler, cfg, testutil.GetTelemetryMock(t))
+	streamHandlers := newStreamCollection(ctx, testutil.GetTelemetryMock(t))
+	consumer := newCudaEventConsumer(ctx, streamHandlers, handler, cfg, testutil.GetTelemetryMock(t))
 
 	consumer.Start()
 	require.Eventually(t, func() bool { return consumer.running.Load() }, 100*time.Millisecond, 10*time.Millisecond)
@@ -39,7 +40,9 @@ func TestGetStreamKeyUpdatesCorrectlyWhenChangingDevice(t *testing.T) {
 	ctx, err := getSystemContext(testutil.GetBasicNvmlMock(), kernel.ProcFSRoot(), testutil.GetWorkloadMetaMock(t), testutil.GetTelemetryMock(t))
 	require.NoError(t, err)
 
-	consumer := newCudaEventConsumer(ctx, nil, nil, testutil.GetTelemetryMock(t))
+	handlers := newStreamCollection(ctx, testutil.GetTelemetryMock(t))
+
+	consumer := newCudaEventConsumer(ctx, handlers, nil, nil, testutil.GetTelemetryMock(t))
 
 	pid := uint32(1)
 	pidTgid := uint64(pid)<<32 + uint64(pid)
@@ -50,34 +53,34 @@ func TestGetStreamKeyUpdatesCorrectlyWhenChangingDevice(t *testing.T) {
 		Stream_id: streamID,
 	}
 
-	globalStream := uint64(0)
+	globalStreamID := uint64(0)
 	headerGlobalStream := gpuebpf.CudaEventHeader{
 		Pid_tgid:  pidTgid,
-		Stream_id: globalStream,
+		Stream_id: globalStreamID,
 	}
 
 	// Configure the visible devices for our process
 	ctx.visibleDevicesCache[int(pid)] = []nvml.Device{testutil.GetDeviceMock(0), testutil.GetDeviceMock(1)}
 
-	streamKey, err := consumer.getStreamKey(&headerStreamSpecific)
+	stream, err := handlers.getStream(&headerStreamSpecific)
 	require.NoError(t, err)
-	require.NotNil(t, streamKey)
-	require.Equal(t, pid, streamKey.pid)
-	require.Equal(t, streamID, streamKey.stream)
+	require.NotNil(t, stream)
+	require.Equal(t, pid, stream.metadata.pid)
+	require.Equal(t, streamID, stream.metadata.streamID)
 
 	// The stream should have the default selected device, which is 0
-	require.Equal(t, testutil.GPUUUIDs[0], streamKey.gpuUUID)
+	require.Equal(t, testutil.GPUUUIDs[0], stream.metadata.gpuUUID)
 
 	// Check the same thing happens with the global stream
-	globalStreamKey, err := consumer.getStreamKey(&headerGlobalStream)
+	globalStream, err := handlers.getStream(&headerGlobalStream)
 
 	require.NoError(t, err)
-	require.NotNil(t, globalStreamKey)
-	require.Equal(t, pid, globalStreamKey.pid)
-	require.Equal(t, globalStream, globalStreamKey.stream)
+	require.NotNil(t, globalStream)
+	require.Equal(t, pid, globalStream.metadata.pid)
+	require.Equal(t, globalStreamID, globalStream.metadata.streamID)
 
 	// Again, this should be on the default device
-	require.Equal(t, testutil.GPUUUIDs[0], globalStreamKey.gpuUUID)
+	require.Equal(t, testutil.GPUUUIDs[0], globalStream.metadata.gpuUUID)
 
 	// Now we trigger a setDevice event on this process
 	setDev := gpuebpf.CudaSetDeviceEvent{
@@ -90,18 +93,18 @@ func TestGetStreamKeyUpdatesCorrectlyWhenChangingDevice(t *testing.T) {
 
 	// The stream key for the specific stream should not change, as streams are per-device
 	// and cannot change devices during its lifetime
-	streamKey, err = consumer.getStreamKey(&headerStreamSpecific)
+	stream, err = handlers.getStream(&headerStreamSpecific)
 	require.NoError(t, err)
-	require.NotNil(t, streamKey)
-	require.Equal(t, pid, streamKey.pid)
-	require.Equal(t, streamID, streamKey.stream)
-	require.Equal(t, testutil.GPUUUIDs[0], streamKey.gpuUUID)
+	require.NotNil(t, stream)
+	require.Equal(t, pid, stream.metadata.pid)
+	require.Equal(t, streamID, stream.metadata.streamID)
+	require.Equal(t, testutil.GPUUUIDs[0], stream.metadata.gpuUUID)
 
 	// The global stream should change, as the global stream can change devices
-	globalStreamKey, err = consumer.getStreamKey(&headerGlobalStream)
+	globalStream, err = handlers.getStream(&headerGlobalStream)
 	require.NoError(t, err)
-	require.NotNil(t, globalStreamKey)
-	require.Equal(t, pid, globalStreamKey.pid)
-	require.Equal(t, globalStream, globalStreamKey.stream)
-	require.Equal(t, testutil.GPUUUIDs[1], globalStreamKey.gpuUUID)
+	require.NotNil(t, globalStream)
+	require.Equal(t, pid, globalStream.metadata.pid)
+	require.Equal(t, globalStreamID, globalStream.metadata.streamID)
+	require.Equal(t, testutil.GPUUUIDs[1], globalStream.metadata.gpuUUID)
 }
diff --git a/pkg/gpu/context.go b/pkg/gpu/context.go
index acb6ba0b8b14..44ef63eca5d0 100644
--- a/pkg/gpu/context.go
+++ b/pkg/gpu/context.go
@@ -382,9 +382,15 @@ func (ctx *systemContext) filterDevicesForContainer(devices []nvml.Device, conta
 
 // getCurrentActiveGpuDevice returns the active GPU device for a given process and thread, based on the
 // last selection (via cudaSetDevice) this thread made and the visible devices for the process.
-func (ctx *systemContext) getCurrentActiveGpuDevice(pid int, tid int, containerID string) (nvml.Device, error) {
+// This function caches the visible devices for the process in the visibleDevicesCache map, so it only
+// does the expensive operations of looking into the process state and filtering devices one time for each process
+// containerIDFunc is a function that returns the container ID for the given process. As retrieving the container ID
+// might be expensive, we pass a function that can be called to retrieve it only when needed
+func (ctx *systemContext) getCurrentActiveGpuDevice(pid int, tid int, containerIDFunc func() string) (nvml.Device, error) {
 	visibleDevices, ok := ctx.visibleDevicesCache[pid]
 	if !ok {
+		containerID := containerIDFunc()
+
 		// Order is important! We need to filter the devices for the container
 		// first. In a container setting, the environment variable acts as a
 		// filter on the devices that are available to the process, not on the
diff --git a/pkg/gpu/context_test.go b/pkg/gpu/context_test.go
index a9fbc300bd08..0a3ce5752f1a 100644
--- a/pkg/gpu/context_test.go
+++ b/pkg/gpu/context_test.go
@@ -198,7 +198,7 @@ func TestGetCurrentActiveGpuDevice(t *testing.T) {
 			}
 
 			for i, idx := range c.expectedDeviceIdx {
-				activeDevice, err := sysCtx.getCurrentActiveGpuDevice(c.pid, c.pid+i, c.containerID)
+				activeDevice, err := sysCtx.getCurrentActiveGpuDevice(c.pid, c.pid+i, func() string { return c.containerID })
 				require.NoError(t, err)
 				testutil.RequireDevicesEqual(t, sysCtx.gpuDevices[idx], activeDevice, "invalid device at index %d (real index is %d, selected index is %d)", i, idx, c.configuredDeviceIdx[i])
 			}
diff --git a/pkg/gpu/probe.go b/pkg/gpu/probe.go
index 07d951c8fe95..3a3df3105243 100644
--- a/pkg/gpu/probe.go
+++ b/pkg/gpu/probe.go
@@ -115,6 +115,7 @@ type Probe struct {
 	eventHandler     ddebpf.EventHandler
 	telemetry        *probeTelemetry
 	mapCleanerEvents *ddebpf.MapCleaner[gpuebpf.CudaEventKey, gpuebpf.CudaEventValue]
+	streamHandlers   *streamCollection
 }
 
 type probeTelemetry struct {
@@ -186,9 +187,9 @@ func NewProbe(cfg *config.Config, deps ProbeDependencies) (*Probe, error) {
 		return nil, fmt.Errorf("error creating uprobes attacher: %w", err)
 	}
 
-	p.consumer = newCudaEventConsumer(sysCtx, p.eventHandler, p.cfg, deps.Telemetry)
-	//TODO: decouple this to avoid sharing streamHandlers between consumer and statsGenerator
-	p.statsGenerator = newStatsGenerator(sysCtx, p.consumer.streamHandlers, deps.Telemetry)
+	p.streamHandlers = newStreamCollection(sysCtx, deps.Telemetry)
+	p.consumer = newCudaEventConsumer(sysCtx, p.streamHandlers, p.eventHandler, p.cfg, deps.Telemetry)
+	p.statsGenerator = newStatsGenerator(sysCtx, p.streamHandlers, deps.Telemetry)
 
 	if err = p.start(); err != nil {
 		return nil, err
@@ -237,7 +238,7 @@ func (p *Probe) GetAndFlush() (*model.GPUStats, error) {
 
 func (p *Probe) cleanupFinished() {
 	p.statsGenerator.cleanupFinishedAggregators()
-	p.consumer.cleanFinishedHandlers()
+	p.streamHandlers.clean()
 }
 
 func (p *Probe) initRCGPU(cfg *config.Config) error {
diff --git a/pkg/gpu/probe_test.go b/pkg/gpu/probe_test.go
index aa4accc2dfd5..0fdff81e1c22 100644
--- a/pkg/gpu/probe_test.go
+++ b/pkg/gpu/probe_test.go
@@ -82,9 +82,9 @@ func (s *probeTestSuite) TestCanReceiveEvents() {
 	var handlerStream, handlerGlobal *StreamHandler
 	require.Eventually(t, func() bool {
 		handlerStream, handlerGlobal = nil, nil // Ensure we see both handlers in the same iteration
-		for key, h := range probe.consumer.streamHandlers {
-			if key.pid == uint32(cmd.Process.Pid) {
-				if key.stream == 0 {
+		for h := range probe.streamHandlers.allStreams() {
+			if h.metadata.pid == uint32(cmd.Process.Pid) {
+				if h.metadata.streamID == 0 {
 					handlerGlobal = h
 				} else {
 					handlerStream = h
@@ -92,7 +92,7 @@ func (s *probeTestSuite) TestCanReceiveEvents() {
 			}
 		}
 
-		return len(probe.consumer.streamHandlers) == 2 && handlerStream != nil && handlerGlobal != nil && len(handlerStream.kernelSpans) > 0 && len(handlerGlobal.allocations) > 0
+		return len(probe.streamHandlers.globalStreams) == 1 && len(probe.streamHandlers.streams) == 1 && handlerStream != nil && handlerGlobal != nil && len(handlerStream.kernelSpans) > 0 && len(handlerGlobal.allocations) > 0
 	}, 3*time.Second, 100*time.Millisecond, "stream and global handlers not found: existing is %v", probe.consumer.streamHandlers)
 
 	// Check that we're receiving the events we expect
@@ -153,8 +153,8 @@ func (s *probeTestSuite) TestCanGenerateStats() {
 	//TODO: change this check to  count telemetry counter of the consumer (once added).
 	// we are expecting 2 different streamhandlers because cudasample generates 3 events in total for 2 different streams (stream 0 and stream 30)
 	require.Eventually(t, func() bool {
-		return len(probe.consumer.streamHandlers) == 2
-	}, 3*time.Second, 100*time.Millisecond, "stream handlers count mismatch: expected: 2, got: %d", len(probe.consumer.streamHandlers))
+		return probe.streamHandlers.streamCount() == 2
+	}, 3*time.Second, 100*time.Millisecond, "stream handlers count mismatch: expected: 2, got: %d", probe.streamHandlers.streamCount())
 
 	stats, err := probe.GetAndFlush()
 	require.NoError(t, err)
@@ -188,8 +188,8 @@ func (s *probeTestSuite) TestMultiGPUSupport() {
 	//TODO: change this check to  count telemetry counter of the consumer (once added).
 	// we are expecting 2 different streamhandlers because cudasample generates 3 events in total for 2 different streams (stream 0 and stream 30)
 	require.Eventually(t, func() bool {
-		return len(probe.consumer.streamHandlers) == 2
-	}, 3*time.Second, 100*time.Millisecond, "stream handlers count mismatch: expected: 2, got: %d", len(probe.consumer.streamHandlers))
+		return probe.streamHandlers.streamCount() == 2
+	}, 3*time.Second, 100*time.Millisecond, "stream handlers count mismatch: expected: 2, got: %d", probe.streamHandlers.streamCount())
 
 	stats, err := probe.GetAndFlush()
 	require.NoError(t, err)
@@ -210,9 +210,9 @@ func (s *probeTestSuite) TestDetectsContainer() {
 	pid, cid := testutil.RunSampleInDocker(t, testutil.CudaSample, testutil.MinimalDockerImage)
 
 	// Check that the stream handlers have the correct container ID assigned
-	for key, handler := range probe.consumer.streamHandlers {
-		if key.pid == uint32(pid) {
-			require.Equal(t, cid, handler.containerID)
+	for handler := range probe.streamHandlers.allStreams() {
+		if handler.metadata.pid == uint32(pid) {
+			require.Equal(t, cid, handler.metadata.containerID)
 		}
 	}
 
diff --git a/pkg/gpu/stats.go b/pkg/gpu/stats.go
index de34eee718cf..182387f9a262 100644
--- a/pkg/gpu/stats.go
+++ b/pkg/gpu/stats.go
@@ -21,7 +21,7 @@ import (
 // statsGenerator connects to the active stream handlers and generates stats for the GPU monitoring, by distributing
 // the data to the aggregators which are responsible for computing the metrics.
 type statsGenerator struct {
-	streamHandlers      map[streamKey]*StreamHandler   // streamHandlers contains the map of active stream handlers.
+	streamHandlers      *streamCollection              // streamHandlers contains the map of active stream handlers.
 	lastGenerationKTime int64                          // lastGenerationTime is the kernel time of the last stats generation.
 	currGenerationKTime int64                          // currGenerationTime is the kernel time of the current stats generation.
 	aggregators         map[model.StatsKey]*aggregator // aggregators contains the map of aggregators
@@ -36,7 +36,7 @@ type statsGeneratorTelemetry struct {
 	aggregators telemetry.Gauge
 }
 
-func newStatsGenerator(sysCtx *systemContext, streamHandlers map[streamKey]*StreamHandler, tm telemetry.Component) *statsGenerator {
+func newStatsGenerator(sysCtx *systemContext, streamHandlers *streamCollection, tm telemetry.Component) *statsGenerator {
 	currKTime, _ := ddebpf.NowNanoseconds()
 	return &statsGenerator{
 		streamHandlers:      streamHandlers,
@@ -62,10 +62,10 @@ func newStatsGeneratorTelemetry(tm telemetry.Component) *statsGeneratorTelemetry
 func (g *statsGenerator) getStats(nowKtime int64) *model.GPUStats {
 	g.currGenerationKTime = nowKtime
 
-	for key, handler := range g.streamHandlers {
-		aggr, err := g.getOrCreateAggregator(key)
+	for handler := range g.streamHandlers.allStreams() {
+		aggr, err := g.getOrCreateAggregator(handler.metadata)
 		if err != nil {
-			log.Errorf("Error getting or creating aggregator for key %v: %s", key, err)
+			log.Errorf("Error getting or creating aggregator for handler metadata %v: %s", handler.metadata, err)
 			continue
 		}
 
@@ -127,7 +127,7 @@ func (g *statsGenerator) getDeviceCores(deviceUUID string) (uint64, error) {
 	return uint64(maxThreads), nil
 }
 
-func (g *statsGenerator) getOrCreateAggregator(sKey streamKey) (*aggregator, error) {
+func (g *statsGenerator) getOrCreateAggregator(sKey streamMetadata) (*aggregator, error) {
 	aggKey := model.StatsKey{
 		PID:         sKey.pid,
 		DeviceUUID:  sKey.gpuUUID,
diff --git a/pkg/gpu/stats_test.go b/pkg/gpu/stats_test.go
index 015c72995ed3..dcbceeebe3b8 100644
--- a/pkg/gpu/stats_test.go
+++ b/pkg/gpu/stats_test.go
@@ -30,7 +30,7 @@ func getMetricsEntry(key model.StatsKey, stats *model.GPUStats) *model.Utilizati
 	return nil
 }
 
-func getStatsGeneratorForTest(t *testing.T) (*statsGenerator, map[streamKey]*StreamHandler, int64) {
+func getStatsGeneratorForTest(t *testing.T) (*statsGenerator, *streamCollection, int64) {
 	sysCtx, err := getSystemContext(testutil.GetBasicNvmlMock(), kernel.ProcFSRoot(), testutil.GetWorkloadMetaMock(t), testutil.GetTelemetryMock(t))
 	require.NoError(t, err)
 	require.NotNil(t, sysCtx)
@@ -38,7 +38,7 @@ func getStatsGeneratorForTest(t *testing.T) (*statsGenerator, map[streamKey]*Str
 	ktime, err := ddebpf.NowNanoseconds()
 	require.NoError(t, err)
 
-	streamHandlers := make(map[streamKey]*StreamHandler)
+	streamHandlers := newStreamCollection(sysCtx, testutil.GetTelemetryMock(t))
 	statsGen := newStatsGenerator(sysCtx, streamHandlers, testutil.GetTelemetryMock(t))
 	statsGen.lastGenerationKTime = ktime
 	statsGen.currGenerationKTime = ktime
@@ -47,6 +47,37 @@ func getStatsGeneratorForTest(t *testing.T) (*statsGenerator, map[streamKey]*Str
 	return statsGen, streamHandlers, ktime
 }
 
+func addStream(streamHandlers *streamCollection, pid uint32, streamID uint64, gpuUUID string, containerID string) *StreamHandler {
+	key := streamKey{pid: pid, stream: streamID}
+	metadata := streamMetadata{
+		pid:         pid,
+		streamID:    streamID,
+		containerID: containerID,
+		gpuUUID:     gpuUUID,
+	}
+
+	stream := newStreamHandler(metadata, streamHandlers.sysCtx)
+	streamHandlers.streams[key] = stream
+
+	return stream
+}
+
+func addGlobalStream(streamHandlers *streamCollection, pid uint32, gpuUUID string, containerID string) *StreamHandler {
+	streamID := uint64(0)
+	key := globalStreamKey{pid: pid, gpuUUID: gpuUUID}
+	metadata := streamMetadata{
+		pid:         pid,
+		streamID:    streamID,
+		containerID: containerID,
+		gpuUUID:     gpuUUID,
+	}
+
+	stream := newStreamHandler(metadata, streamHandlers.sysCtx)
+	streamHandlers.globalStreams[key] = stream
+
+	return stream
+}
+
 func TestGetStatsWithOnlyCurrentStreamData(t *testing.T) {
 	statsGen, streamHandlers, ktime := getStatsGeneratorForTest(t)
 
@@ -54,34 +85,30 @@ func TestGetStatsWithOnlyCurrentStreamData(t *testing.T) {
 	pid := uint32(1)
 	streamID := uint64(120)
 	pidTgid := uint64(pid)<<32 + uint64(pid)
-	skeyKern := streamKey{pid: pid, stream: streamID, gpuUUID: testutil.DefaultGpuUUID}
 	shmemSize := uint64(10)
-	streamHandlers[skeyKern] = &StreamHandler{
-		processEnded: false,
-		kernelLaunches: []enrichedKernelLaunch{
-			{
-				CudaKernelLaunch: gpuebpf.CudaKernelLaunch{
-					Header:          gpuebpf.CudaEventHeader{Ktime_ns: uint64(startKtime), Pid_tgid: pidTgid, Stream_id: streamID},
-					Kernel_addr:     0,
-					Shared_mem_size: shmemSize,
-					Grid_size:       gpuebpf.Dim3{X: 1, Y: 1, Z: 1},
-					Block_size:      gpuebpf.Dim3{X: 1, Y: 1, Z: 1},
-				},
+	stream := addStream(streamHandlers, pid, streamID, testutil.DefaultGpuUUID, "")
+	stream.processEnded = false
+	stream.kernelLaunches = []enrichedKernelLaunch{
+		{
+			CudaKernelLaunch: gpuebpf.CudaKernelLaunch{
+				Header:          gpuebpf.CudaEventHeader{Ktime_ns: uint64(startKtime), Pid_tgid: pidTgid, Stream_id: streamID},
+				Kernel_addr:     0,
+				Shared_mem_size: shmemSize,
+				Grid_size:       gpuebpf.Dim3{X: 1, Y: 1, Z: 1},
+				Block_size:      gpuebpf.Dim3{X: 1, Y: 1, Z: 1},
 			},
 		},
 	}
 
 	allocSize := uint64(10)
-	skeyAlloc := streamKey{pid: pid, stream: 0, gpuUUID: testutil.DefaultGpuUUID}
-	streamHandlers[skeyAlloc] = &StreamHandler{
-		processEnded: false,
-		memAllocEvents: map[uint64]gpuebpf.CudaMemEvent{
-			0: {
-				Header: gpuebpf.CudaEventHeader{Ktime_ns: uint64(startKtime), Pid_tgid: pidTgid, Stream_id: streamID},
-				Addr:   0,
-				Size:   allocSize,
-				Type:   gpuebpf.CudaMemAlloc,
-			},
+	stream = addGlobalStream(streamHandlers, pid, testutil.DefaultGpuUUID, "")
+	stream.processEnded = false
+	stream.memAllocEvents = map[uint64]gpuebpf.CudaMemEvent{
+		0: {
+			Header: gpuebpf.CudaEventHeader{Ktime_ns: uint64(startKtime), Pid_tgid: pidTgid, Stream_id: streamID},
+			Addr:   0,
+			Size:   allocSize,
+			Type:   gpuebpf.CudaMemAlloc,
 		},
 	}
 
@@ -92,7 +119,7 @@ func TestGetStatsWithOnlyCurrentStreamData(t *testing.T) {
 
 	metricsKey := model.StatsKey{PID: pid, DeviceUUID: testutil.DefaultGpuUUID}
 	metrics := getMetricsEntry(metricsKey, stats)
-	require.NotNil(t, metrics)
+	require.NotNil(t, metrics, "did not find metrics for key %+v", metricsKey)
 	require.Equal(t, allocSize*2, metrics.Memory.CurrentBytes)
 	require.Equal(t, allocSize*2, metrics.Memory.MaxBytes)
 
@@ -109,32 +136,28 @@ func TestGetStatsWithOnlyPastStreamData(t *testing.T) {
 
 	pid := uint32(1)
 	streamID := uint64(120)
-	skeyKern := streamKey{pid: pid, stream: streamID, gpuUUID: testutil.DefaultGpuUUID}
 	numThreads := uint64(5)
-	streamHandlers[skeyKern] = &StreamHandler{
-		processEnded: false,
-		kernelSpans: []*kernelSpan{
-			{
-				startKtime:     uint64(startKtime),
-				endKtime:       uint64(endKtime),
-				avgThreadCount: numThreads,
-				numKernels:     10,
-			},
+	stream := addStream(streamHandlers, pid, streamID, testutil.DefaultGpuUUID, "")
+	stream.processEnded = false
+	stream.kernelSpans = []*kernelSpan{
+		{
+			startKtime:     uint64(startKtime),
+			endKtime:       uint64(endKtime),
+			avgThreadCount: numThreads,
+			numKernels:     10,
 		},
 	}
 
 	allocSize := uint64(10)
-	skeyAlloc := streamKey{pid: pid, stream: 0, gpuUUID: testutil.DefaultGpuUUID}
-	streamHandlers[skeyAlloc] = &StreamHandler{
-		processEnded: false,
-		allocations: []*memoryAllocation{
-			{
-				startKtime: uint64(startKtime),
-				endKtime:   uint64(endKtime),
-				size:       allocSize,
-				isLeaked:   false,
-				allocType:  globalMemAlloc,
-			},
+	stream = addGlobalStream(streamHandlers, pid, testutil.DefaultGpuUUID, "")
+	stream.processEnded = false
+	stream.allocations = []*memoryAllocation{
+		{
+			startKtime: uint64(startKtime),
+			endKtime:   uint64(endKtime),
+			size:       allocSize,
+			isLeaked:   false,
+			allocType:  globalMemAlloc,
 		},
 	}
 
@@ -162,53 +185,51 @@ func TestGetStatsWithPastAndCurrentData(t *testing.T) {
 
 	pid := uint32(1)
 	streamID := uint64(120)
-	skeyKern := streamKey{pid: pid, stream: streamID, gpuUUID: testutil.DefaultGpuUUID}
 	pidTgid := uint64(pid)<<32 + uint64(pid)
 	numThreads := uint64(5)
 	shmemSize := uint64(10)
-	streamHandlers[skeyKern] = &StreamHandler{
-		processEnded: false,
-		kernelLaunches: []enrichedKernelLaunch{
-			{
-				CudaKernelLaunch: gpuebpf.CudaKernelLaunch{
-					Header:          gpuebpf.CudaEventHeader{Ktime_ns: uint64(startKtime), Pid_tgid: pidTgid, Stream_id: streamID},
-					Kernel_addr:     0,
-					Shared_mem_size: shmemSize,
-					Grid_size:       gpuebpf.Dim3{X: 1, Y: 1, Z: 1},
-					Block_size:      gpuebpf.Dim3{X: 1, Y: 1, Z: 1},
-				},
+	stream := addStream(streamHandlers, pid, streamID, testutil.DefaultGpuUUID, "")
+	stream.processEnded = false
+	stream.kernelLaunches = []enrichedKernelLaunch{
+		{
+			CudaKernelLaunch: gpuebpf.CudaKernelLaunch{
+				Header:          gpuebpf.CudaEventHeader{Ktime_ns: uint64(startKtime), Pid_tgid: pidTgid, Stream_id: streamID},
+				Kernel_addr:     0,
+				Shared_mem_size: shmemSize,
+				Grid_size:       gpuebpf.Dim3{X: 1, Y: 1, Z: 1},
+				Block_size:      gpuebpf.Dim3{X: 1, Y: 1, Z: 1},
 			},
 		},
-		kernelSpans: []*kernelSpan{
-			{
-				startKtime:     uint64(startKtime),
-				endKtime:       uint64(endKtime),
-				avgThreadCount: numThreads,
-				numKernels:     10,
-			},
+	}
+
+	stream.kernelSpans = []*kernelSpan{
+		{
+			startKtime:     uint64(startKtime),
+			endKtime:       uint64(endKtime),
+			avgThreadCount: numThreads,
+			numKernels:     10,
 		},
 	}
 
 	allocSize := uint64(10)
-	skeyAlloc := streamKey{pid: pid, stream: 0, gpuUUID: testutil.DefaultGpuUUID}
-	streamHandlers[skeyAlloc] = &StreamHandler{
-		processEnded: false,
-		allocations: []*memoryAllocation{
-			{
-				startKtime: uint64(startKtime),
-				endKtime:   uint64(endKtime),
-				size:       allocSize,
-				isLeaked:   false,
-				allocType:  globalMemAlloc,
-			},
+	stream = addGlobalStream(streamHandlers, pid, testutil.DefaultGpuUUID, "")
+	stream.processEnded = false
+	stream.allocations = []*memoryAllocation{
+		{
+			startKtime: uint64(startKtime),
+			endKtime:   uint64(endKtime),
+			size:       allocSize,
+			isLeaked:   false,
+			allocType:  globalMemAlloc,
 		},
-		memAllocEvents: map[uint64]gpuebpf.CudaMemEvent{
-			0: {
-				Header: gpuebpf.CudaEventHeader{Ktime_ns: uint64(startKtime), Pid_tgid: pidTgid, Stream_id: streamID},
-				Addr:   0,
-				Size:   allocSize,
-				Type:   gpuebpf.CudaMemAlloc,
-			},
+	}
+
+	stream.memAllocEvents = map[uint64]gpuebpf.CudaMemEvent{
+		0: {
+			Header: gpuebpf.CudaEventHeader{Ktime_ns: uint64(startKtime), Pid_tgid: pidTgid, Stream_id: streamID},
+			Addr:   0,
+			Size:   allocSize,
+			Type:   gpuebpf.CudaMemAlloc,
 		},
 	}
 
@@ -242,16 +263,14 @@ func TestGetStatsMultiGPU(t *testing.T) {
 	// Add kernels for all devices
 	for i, uuid := range testutil.GPUUUIDs {
 		streamID := uint64(i)
-		streamKey := streamKey{pid: pid, stream: streamID, gpuUUID: uuid}
-		streamHandlers[streamKey] = &StreamHandler{
-			processEnded: false,
-			kernelSpans: []*kernelSpan{
-				{
-					startKtime:     uint64(startKtime),
-					endKtime:       uint64(endKtime),
-					avgThreadCount: numThreads,
-					numKernels:     10,
-				},
+		stream := addStream(streamHandlers, pid, streamID, uuid, "")
+		stream.processEnded = false
+		stream.kernelSpans = []*kernelSpan{
+			{
+				startKtime:     uint64(startKtime),
+				endKtime:       uint64(endKtime),
+				avgThreadCount: numThreads,
+				numKernels:     10,
 			},
 		}
 	}
diff --git a/pkg/gpu/stream.go b/pkg/gpu/stream.go
index 540ff73f8b33..213320dace76 100644
--- a/pkg/gpu/stream.go
+++ b/pkg/gpu/stream.go
@@ -28,15 +28,13 @@ const noSmVersion uint32 = 0
 // StreamHandler is responsible for receiving events from a single CUDA stream and generating
 // kernel spans and memory allocations from them.
 type StreamHandler struct {
-	pid            uint32
+	metadata       streamMetadata
 	kernelLaunches []enrichedKernelLaunch
 	memAllocEvents map[uint64]gpuebpf.CudaMemEvent
 	kernelSpans    []*kernelSpan
 	allocations    []*memoryAllocation
 	processEnded   bool // A marker to indicate that the process has ended, and this handler should be flushed
 	sysCtx         *systemContext
-	containerID    string
-	smVersion      uint32 // The SM version of the GPU this stream is running on, for kernel data attaching
 }
 
 // enrichedKernelLaunch is a kernel launch event with the kernel data attached.
@@ -47,12 +45,23 @@ type enrichedKernelLaunch struct {
 	kernel *cuda.CubinKernel
 }
 
-// streamKey is a unique identifier for a CUDA stream
-type streamKey struct {
-	pid         uint32
-	stream      uint64
-	gpuUUID     string
+// streamMetadata contains metadata about a CUDA stream
+type streamMetadata struct {
+	// pid is the PID of the process that is running this stream
+	pid uint32
+
+	// streamID is the ID of the CUDA stream
+	streamID uint64
+
+	// gpuUUID is the UUID of the GPU this stream is running on
+	gpuUUID string
+
+	// containerID is the container ID of the process that is running this stream. Might be empty if the container ID is not available
+	// or if the process is not running inside a container
 	containerID string
+
+	// smVersion is the SM version of the GPU this stream is running on, for kernel data attaching
+	smVersion uint32
 }
 
 // streamData contains kernel spans and allocations for a stream
@@ -117,13 +126,11 @@ type kernelSpan struct {
 	avgMemoryUsage map[memAllocType]uint64
 }
 
-func newStreamHandler(pid uint32, containerID string, smVersion uint32, sysCtx *systemContext) *StreamHandler {
+func newStreamHandler(metadata streamMetadata, sysCtx *systemContext) *StreamHandler {
 	return &StreamHandler{
 		memAllocEvents: make(map[uint64]gpuebpf.CudaMemEvent),
-		pid:            pid,
 		sysCtx:         sysCtx,
-		containerID:    containerID,
-		smVersion:      smVersion,
+		metadata:       metadata,
 	}
 }
 
@@ -137,7 +144,7 @@ func (sh *StreamHandler) handleKernelLaunch(event *gpuebpf.CudaKernelLaunch) {
 	err := sh.tryAttachKernelData(enrichedLaunch)
 	if err != nil {
 		if logLimitErrorAttach.ShouldLog() {
-			log.Warnf("Error attaching kernel data for PID %d: %v", sh.pid, err)
+			log.Warnf("Error attaching kernel data for PID %d: %v", sh.metadata.pid, err)
 		}
 	}
 
@@ -155,12 +162,13 @@ func findEntryInMaps(procMaps []*procfs.ProcMap, addr uintptr) *procfs.ProcMap {
 }
 
 func (sh *StreamHandler) tryAttachKernelData(event *enrichedKernelLaunch) error {
-	if sh.sysCtx == nil || !sh.sysCtx.fatbinParsingEnabled || sh.smVersion == noSmVersion {
-		// Fatbin parsing is disabled or we don't have SM version, so we can't attach kernel data
+	if sh.sysCtx == nil || !sh.sysCtx.fatbinParsingEnabled || sh.metadata.smVersion == noSmVersion {
+		// No system context or we don't have a SM version to use,
+		// kernel data attaching is disabled
 		return nil
 	}
 
-	maps, err := sh.sysCtx.getProcessMemoryMaps(int(sh.pid))
+	maps, err := sh.sysCtx.getProcessMemoryMaps(int(sh.metadata.pid))
 	if err != nil {
 		return fmt.Errorf("error reading process memory maps: %w", err)
 	}
@@ -172,7 +180,7 @@ func (sh *StreamHandler) tryAttachKernelData(event *enrichedKernelLaunch) error
 
 	offsetInFile := uint64(int64(event.Kernel_addr) - int64(entry.StartAddr) + entry.Offset)
 
-	binaryPath := path.Join(sh.sysCtx.procRoot, strconv.Itoa(int(sh.pid)), "root", entry.Pathname)
+	binaryPath := path.Join(sh.sysCtx.procRoot, strconv.Itoa(int(sh.metadata.pid)), "root", entry.Pathname)
 	fileData, err := sh.sysCtx.getCudaSymbols(binaryPath)
 	if err != nil {
 		return fmt.Errorf("error getting file %s data: %w", binaryPath, err)
@@ -183,7 +191,7 @@ func (sh *StreamHandler) tryAttachKernelData(event *enrichedKernelLaunch) error
 		return fmt.Errorf("could not find symbol for address 0x%x in file %s", event.Kernel_addr, binaryPath)
 	}
 
-	kern := fileData.Fatbin.GetKernel(symbol, sh.smVersion)
+	kern := fileData.Fatbin.GetKernel(symbol, sh.metadata.smVersion)
 	if kern == nil {
 		return fmt.Errorf("could not find kernel for symbol %s in file %s", symbol, binaryPath)
 	}
@@ -373,7 +381,7 @@ func (sh *StreamHandler) markEnd() error {
 		sh.allocations = append(sh.allocations, &data)
 	}
 
-	sh.sysCtx.removeProcess(int(sh.pid))
+	sh.sysCtx.removeProcess(int(sh.metadata.pid))
 
 	return nil
 }
diff --git a/pkg/gpu/stream_collection.go b/pkg/gpu/stream_collection.go
new file mode 100644
index 000000000000..ad4607f832a5
--- /dev/null
+++ b/pkg/gpu/stream_collection.go
@@ -0,0 +1,240 @@
+// 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 2024-present Datadog, Inc.
+
+//go:build linux_bpf
+
+package gpu
+
+import (
+	"fmt"
+	"iter"
+
+	"github.com/NVIDIA/go-nvml/pkg/nvml"
+	"golang.org/x/sys/unix"
+
+	"github.com/DataDog/datadog-agent/comp/core/telemetry"
+	gpuebpf "github.com/DataDog/datadog-agent/pkg/gpu/ebpf"
+	"github.com/DataDog/datadog-agent/pkg/util/cgroups"
+	"github.com/DataDog/datadog-agent/pkg/util/funcs"
+	"github.com/DataDog/datadog-agent/pkg/util/log"
+)
+
+// streamKey is a unique identifier for a CUDA stream that is not global.
+// The streams are created with a specific GPU UUID, which means that pid+stream uniquely
+// identify the GPU.
+type streamKey struct {
+	pid    uint32
+	stream uint64
+}
+
+// globalStreamKey is a unique identifier for a CUDA stream that is global.
+// Global streams depend on which GPU is active when they are used, so the GPU UUID
+// needs to be part of the key
+type globalStreamKey struct {
+	pid     uint32
+	gpuUUID string
+}
+
+type streamCollection struct {
+	streams       map[streamKey]*StreamHandler
+	globalStreams map[globalStreamKey]*StreamHandler
+	sysCtx        *systemContext
+	telemetry     *streamCollectionTelemetry
+}
+
+type streamCollectionTelemetry struct {
+	missingContainers  telemetry.Counter
+	missingDevices     telemetry.Counter
+	finalizedProcesses telemetry.Counter
+	activeHandlers     telemetry.Gauge
+	removedHandlers    telemetry.Counter
+}
+
+func newStreamCollection(sysCtx *systemContext, telemetry telemetry.Component) *streamCollection {
+	return &streamCollection{
+		streams:       make(map[streamKey]*StreamHandler),
+		globalStreams: make(map[globalStreamKey]*StreamHandler),
+		sysCtx:        sysCtx,
+		telemetry:     newStreamCollectionTelemetry(telemetry),
+	}
+}
+
+func newStreamCollectionTelemetry(tm telemetry.Component) *streamCollectionTelemetry {
+	subsystem := gpuTelemetryModule + "__streams"
+
+	return &streamCollectionTelemetry{
+		missingContainers:  tm.NewCounter(subsystem, "missing_containers", []string{"reason"}, "Number of missing containers"),
+		missingDevices:     tm.NewCounter(subsystem, "missing_devices", nil, "Number of failures to get GPU devices for a stream"),
+		finalizedProcesses: tm.NewCounter(subsystem, "finalized_processes", nil, "Number of processes that have ended"),
+		activeHandlers:     tm.NewGauge(subsystem, "active_handlers", nil, "Number of active stream handlers"),
+		removedHandlers:    tm.NewCounter(subsystem, "removed_handlers", []string{"device"}, "Number of removed stream handlers"),
+	}
+}
+
+// getStream returns a StreamHandler for a given CUDA stream.
+func (sc *streamCollection) getStream(header *gpuebpf.CudaEventHeader) (*StreamHandler, error) {
+	if header.Stream_id == 0 {
+		return sc.getGlobalStream(header)
+	}
+
+	return sc.getNonGlobalStream(header)
+}
+
+func (sc *streamCollection) getGlobalStream(header *gpuebpf.CudaEventHeader) (*StreamHandler, error) {
+	pid, tid := getPidTidFromHeader(header)
+	memoizedContainerID := sc.memoizedContainerID(header)
+
+	// Global streams depend on which GPU is active when they are used, so we need to get the current active GPU device
+	// The expensive step here is the container ID parsing, but we don't always need to do it so we pass a function
+	// that can be called to retrieve it only when needed
+	gpuUUID, err := sc.getActiveDevice(int(pid), int(tid), memoizedContainerID)
+	if err != nil {
+		sc.telemetry.missingDevices.Inc()
+		return nil, err
+	}
+
+	key := globalStreamKey{
+		pid:     pid,
+		gpuUUID: gpuUUID,
+	}
+
+	stream, ok := sc.globalStreams[key]
+	if !ok {
+		stream = sc.createStreamHandler(header, &gpuUUID, memoizedContainerID)
+		sc.globalStreams[key] = stream
+		sc.telemetry.activeHandlers.Set(float64(sc.streamCount()))
+	}
+
+	return stream, nil
+}
+
+func (sc *streamCollection) getNonGlobalStream(header *gpuebpf.CudaEventHeader) (*StreamHandler, error) {
+	pid, _ := getPidTidFromHeader(header)
+
+	key := streamKey{
+		pid:    pid,
+		stream: header.Stream_id,
+	}
+
+	stream, ok := sc.streams[key]
+	if !ok {
+		stream = sc.createStreamHandler(header, nil, sc.memoizedContainerID(header))
+		sc.streams[key] = stream
+		sc.telemetry.activeHandlers.Set(float64(sc.streamCount()))
+	}
+
+	return stream, nil
+}
+
+// createStreamHandler creates a new StreamHandler for a given CUDA stream.
+// If the GPU UUID is not provided (it's nil), it will be retrieved from the system context.
+func (sc *streamCollection) createStreamHandler(header *gpuebpf.CudaEventHeader, gpuUUID *string, containerIDFunc func() string) *StreamHandler {
+	pid, tid := getPidTidFromHeader(header)
+	metadata := streamMetadata{
+		pid:         pid,
+		streamID:    header.Stream_id,
+		containerID: containerIDFunc(),
+	}
+
+	if gpuUUID != nil {
+		metadata.gpuUUID = *gpuUUID
+	} else {
+		var err error
+		metadata.gpuUUID, err = sc.getActiveDevice(int(pid), int(tid), containerIDFunc)
+		if err != nil {
+			log.Warnf("error getting GPU UUID for process %d: %s", pid, err)
+			sc.telemetry.missingDevices.Inc()
+			return nil
+		}
+	}
+
+	metadata.smVersion = sc.sysCtx.deviceSmVersions[metadata.gpuUUID]
+
+	return newStreamHandler(metadata, sc.sysCtx)
+}
+
+// memoizedContainerID returns a function that memoizes the container ID for a given CUDA stream.
+// It's memoized because it's an expensive operation that we don't always need to do.
+func (sc *streamCollection) memoizedContainerID(header *gpuebpf.CudaEventHeader) func() string {
+	return funcs.MemoizeNoErrorUnsafe(func() string {
+		cgroup := unix.ByteSliceToString(header.Cgroup[:])
+		containerID, err := cgroups.ContainerFilter("", cgroup)
+		if err != nil {
+			log.Warnf("error getting container ID for cgroup %s: %s", cgroup, err)
+
+			sc.telemetry.missingContainers.Inc("error")
+			return ""
+		}
+
+		if containerID == "" {
+			sc.telemetry.missingContainers.Inc("missing")
+		}
+
+		return containerID
+	})
+}
+
+// getActiveDevice returns the GPU UUID for the current active GPU device for a given process and thread.
+func (sc *streamCollection) getActiveDevice(pid int, tid int, containerIDFunc func() string) (string, error) {
+	device, err := sc.sysCtx.getCurrentActiveGpuDevice(pid, tid, containerIDFunc)
+	if err != nil {
+		return "", err
+	}
+
+	uuid, ret := device.GetUUID()
+	if ret != nvml.SUCCESS {
+		return "", fmt.Errorf("error getting GPU UUID for process %d: %v", pid, nvml.ErrorString(ret))
+	}
+
+	return uuid, nil
+}
+
+func (sc *streamCollection) allStreams() iter.Seq[*StreamHandler] {
+	return func(yield func(*StreamHandler) bool) {
+		for _, stream := range sc.streams {
+			if !yield(stream) {
+				return
+			}
+		}
+
+		for _, stream := range sc.globalStreams {
+			if !yield(stream) {
+				return
+			}
+		}
+	}
+}
+
+func (sc *streamCollection) streamCount() int {
+	return len(sc.streams) + len(sc.globalStreams)
+}
+
+func (sc *streamCollection) markProcessStreamsAsEnded(pid uint32) {
+	for handler := range sc.allStreams() {
+		if handler.metadata.pid == pid {
+			log.Debugf("Process %d ended, marking stream %d as ended", pid, handler.metadata.streamID)
+			_ = handler.markEnd()
+			sc.telemetry.finalizedProcesses.Inc()
+		}
+	}
+}
+
+func (sc *streamCollection) clean() {
+	for key, handler := range sc.streams {
+		if handler.processEnded {
+			delete(sc.streams, key)
+			sc.telemetry.removedHandlers.Inc(handler.metadata.gpuUUID)
+		}
+	}
+
+	for key, handler := range sc.globalStreams {
+		if handler.processEnded {
+			delete(sc.globalStreams, key)
+			sc.telemetry.removedHandlers.Inc(handler.metadata.gpuUUID)
+		}
+	}
+
+	sc.telemetry.activeHandlers.Set(float64(sc.streamCount()))
+}
diff --git a/pkg/gpu/stream_collection_test.go b/pkg/gpu/stream_collection_test.go
new file mode 100644
index 000000000000..0e225ade7133
--- /dev/null
+++ b/pkg/gpu/stream_collection_test.go
@@ -0,0 +1,85 @@
+// 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 2025-present Datadog, Inc.
+
+//go:build linux_bpf
+
+package gpu
+
+import (
+	"testing"
+
+	"github.com/NVIDIA/go-nvml/pkg/nvml"
+	"github.com/stretchr/testify/require"
+
+	gpuebpf "github.com/DataDog/datadog-agent/pkg/gpu/ebpf"
+	"github.com/DataDog/datadog-agent/pkg/gpu/testutil"
+	"github.com/DataDog/datadog-agent/pkg/util/kernel"
+)
+
+func TestStreamKeyUpdatesCorrectlyWhenChangingDevice(t *testing.T) {
+	ctx, err := getSystemContext(testutil.GetBasicNvmlMock(), kernel.ProcFSRoot(), testutil.GetWorkloadMetaMock(t), testutil.GetTelemetryMock(t))
+	require.NoError(t, err)
+
+	handlers := newStreamCollection(ctx, testutil.GetTelemetryMock(t))
+
+	pid := uint32(1)
+	pidTgid := uint64(pid)<<32 + uint64(pid)
+
+	streamID := uint64(120)
+	headerStreamSpecific := gpuebpf.CudaEventHeader{
+		Pid_tgid:  pidTgid,
+		Stream_id: streamID,
+	}
+
+	globalStreamID := uint64(0)
+	headerGlobalStream := gpuebpf.CudaEventHeader{
+		Pid_tgid:  pidTgid,
+		Stream_id: globalStreamID,
+	}
+
+	// Configure the visible devices for our process
+	ctx.visibleDevicesCache[int(pid)] = []nvml.Device{testutil.GetDeviceMock(0), testutil.GetDeviceMock(1)}
+
+	stream, err := handlers.getStream(&headerStreamSpecific)
+	require.NoError(t, err)
+	require.NotNil(t, stream)
+	require.Equal(t, pid, stream.metadata.pid)
+	require.Equal(t, streamID, stream.metadata.streamID)
+
+	// The stream should have the default selected device, which is 0
+	defaultDevice := 0
+	require.Equal(t, testutil.GPUUUIDs[defaultDevice], stream.metadata.gpuUUID)
+
+	// Check the same thing happens with the global stream
+	globalStream, err := handlers.getStream(&headerGlobalStream)
+	require.NoError(t, err)
+	require.NotNil(t, globalStream)
+	require.Equal(t, pid, globalStream.metadata.pid)
+	require.Equal(t, globalStreamID, globalStream.metadata.streamID)
+
+	// Again, this should be on the default device
+	require.Equal(t, testutil.GPUUUIDs[defaultDevice], globalStream.metadata.gpuUUID)
+
+	// Now we change the device for the specific stream
+	selectedDevice := 1
+	ctx.selectedDeviceByPIDAndTID[int(pid)] = map[int]int32{int(pid): int32(selectedDevice)}
+
+	// The stream key for the specific stream should not change, as streams are per-device
+	// and cannot change devices during its lifetime
+	stream, err = handlers.getStream(&headerStreamSpecific)
+	require.NoError(t, err)
+	require.NotNil(t, stream)
+	require.Equal(t, pid, stream.metadata.pid)
+	require.Equal(t, streamID, stream.metadata.streamID)
+	require.Equal(t, testutil.GPUUUIDs[defaultDevice], stream.metadata.gpuUUID)
+
+	// The global stream should change, as the global stream can change devices
+	globalStream, err = handlers.getStream(&headerGlobalStream)
+	require.NoError(t, err)
+	require.NotNil(t, globalStream)
+	require.Equal(t, pid, globalStream.metadata.pid)
+	require.Equal(t, globalStreamID, globalStream.metadata.streamID)
+	require.Equal(t, testutil.GPUUUIDs[selectedDevice], globalStream.metadata.gpuUUID)
+}
diff --git a/pkg/gpu/stream_test.go b/pkg/gpu/stream_test.go
index 5bb301adbfef..12a7a479b695 100644
--- a/pkg/gpu/stream_test.go
+++ b/pkg/gpu/stream_test.go
@@ -30,7 +30,7 @@ func getSystemContextForTest(t *testing.T) *systemContext {
 }
 
 func TestKernelLaunchesHandled(t *testing.T) {
-	stream := newStreamHandler(0, "", 90, getSystemContextForTest(t))
+	stream := newStreamHandler(streamMetadata{}, getSystemContextForTest(t))
 
 	kernStartTime := uint64(1)
 	launch := &gpuebpf.CudaKernelLaunch{
@@ -87,7 +87,7 @@ func TestKernelLaunchesHandled(t *testing.T) {
 }
 
 func TestMemoryAllocationsHandled(t *testing.T) {
-	stream := newStreamHandler(0, "", 90, getSystemContextForTest(t))
+	stream := newStreamHandler(streamMetadata{}, getSystemContextForTest(t))
 
 	memAllocTime := uint64(1)
 	memFreeTime := uint64(2)
@@ -156,7 +156,7 @@ func TestMemoryAllocationsHandled(t *testing.T) {
 }
 
 func TestMemoryAllocationsDetectLeaks(t *testing.T) {
-	stream := newStreamHandler(0, "", 90, getSystemContextForTest(t))
+	stream := newStreamHandler(streamMetadata{}, getSystemContextForTest(t))
 
 	memAllocTime := uint64(1)
 	memAddr := uint64(42)
@@ -189,7 +189,7 @@ func TestMemoryAllocationsDetectLeaks(t *testing.T) {
 }
 
 func TestMemoryAllocationsNoCrashOnInvalidFree(t *testing.T) {
-	stream := newStreamHandler(0, "", 90, getSystemContextForTest(t))
+	stream := newStreamHandler(streamMetadata{}, getSystemContextForTest(t))
 
 	memAllocTime := uint64(1)
 	memFreeTime := uint64(2)
@@ -231,7 +231,7 @@ func TestMemoryAllocationsNoCrashOnInvalidFree(t *testing.T) {
 }
 
 func TestMemoryAllocationsMultipleAllocsHandled(t *testing.T) {
-	stream := newStreamHandler(0, "", 90, getSystemContextForTest(t))
+	stream := newStreamHandler(streamMetadata{}, getSystemContextForTest(t))
 
 	memAllocTime1, memAllocTime2 := uint64(1), uint64(10)
 	memFreeTime1, memFreeTime2 := uint64(15), uint64(20)
@@ -368,7 +368,7 @@ func TestKernelLaunchEnrichment(t *testing.T) {
 				},
 			}
 
-			stream := newStreamHandler(uint32(pid), "", smVersion, sysCtx)
+			stream := newStreamHandler(streamMetadata{pid: uint32(pid), smVersion: smVersion}, sysCtx)
 
 			kernStartTime := uint64(1)
 			launch := &gpuebpf.CudaKernelLaunch{
diff --git a/pkg/util/funcs/memoize_noerror_unsafe.go b/pkg/util/funcs/memoize_noerror_unsafe.go
new file mode 100644
index 000000000000..f5e216777a15
--- /dev/null
+++ b/pkg/util/funcs/memoize_noerror_unsafe.go
@@ -0,0 +1,31 @@
+// 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 2016-present Datadog, Inc.
+
+package funcs
+
+type memoizedNoErrorUnsafeFunc[T any] struct {
+	fn     func() T
+	result T
+	done   bool
+}
+
+func (mf *memoizedNoErrorUnsafeFunc[T]) do() T {
+	if !mf.done {
+		mf.result = mf.fn()
+		mf.done = true
+	}
+	return mf.result
+}
+
+// MemoizeNoErrorUnsafe the result of a function call.
+// This is a thread unsafe version of MemoizeNoError.
+//
+// fn is only ever called once, but this implementation is not thread-safe.
+// Use MemoizeNoError if you need thread safety. Use this if you need to memoize
+// a function that only gets called in a single thread, to avoid the overhead of
+// the mutex.
+func MemoizeNoErrorUnsafe[T any](fn func() T) func() T {
+	return (&memoizedNoErrorUnsafeFunc[T]{fn: fn}).do
+}


Reviews:
user: login: gjulianm
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-18T10:38:38Z
user: login: val06
body: left some small comments
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-18T12:19:07Z
user: login: gjulianm
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-19T11:20:42Z
user: login: gjulianm
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-19T11:24:31Z
user: login: gjulianm
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-19T11:24:35Z
user: login: gjulianm
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-19T11:28:47Z
user: login: val06
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-03-19T11:45:46Z

Review comments:
diff_hunk: @@ -241,88 +220,7 @@ func (c *cudaEventConsumer) handleGlobalEvent(header *gpuebpf.CudaEventHeader, d
 }
 
 func (c *cudaEventConsumer) handleProcessExit(pid uint32) {
-	for key, handler := range c.streamHandlers {
-		if key.pid == pid {
-			log.Debugf("Process %d ended, marking stream %d as ended", pid, key.stream)
-			// the probe is responsible for deleting the stream handler
-			_ = handler.markEnd()
-			c.telemetry.finalizedProcesses.Inc()
-		}
-	}
-}
-
-func (c *cudaEventConsumer) getStreamKey(header *gpuebpf.CudaEventHeader) (streamKey, error) {
-	pid, tid := getPidTidFromHeader(header)
-
-	var nonGlobalKey nonGlobalStreamKey
-	if header.Stream_id != 0 {
-		// Non-global stream, check if we have created it before
-		nonGlobalKey.pid = pid
-		nonGlobalKey.stream = header.Stream_id
-
-		if key, ok := c.nonGlobalStreamKeys[nonGlobalKey]; ok {
-			return key, nil
-		}
-	}
-
-	cgroup := unix.ByteSliceToString(header.Cgroup[:])
-	containerID, err := cgroups.ContainerFilter("", cgroup)
-	if err != nil {
-		// We don't want to return an error here, as we can still process the event without the container ID
-		log.Warnf("error getting container ID for cgroup %s: %s", cgroup, err)
-		c.telemetry.missingContainers.Inc("error")
-	} else if containerID == "" {
-		c.telemetry.missingContainers.Inc("missing")
-	}
-
-	key := streamKey{
-		pid:         pid,
-		stream:      header.Stream_id,
-		gpuUUID:     "",
-		containerID: containerID,
-	}
-
-	gpuDevice, err := c.sysCtx.getCurrentActiveGpuDevice(int(pid), int(tid), containerID)
-	if err != nil {
-		c.telemetry.missingDevices.Inc()
-		return streamKey{}, fmt.Errorf("Error getting GPU device for process %d: %w", pid, err)
-	}
-
-	var ret nvml.Return
-	key.gpuUUID, ret = gpuDevice.GetUUID()
-	if ret != nvml.SUCCESS {
-		return streamKey{}, fmt.Errorf("Error getting GPU UUID for process %d: %v", pid, nvml.ErrorString(ret))
-	}
-
-	if header.Stream_id != 0 {
-		c.nonGlobalStreamKeys[nonGlobalKey] = key
-	}
-
-	return key, nil
-}
-
-func (c *cudaEventConsumer) getStreamHandler(header *gpuebpf.CudaEventHeader) (*StreamHandler, error) {
-	key, err := c.getStreamKey(header)
-	if err != nil {
-		return nil, err
-	}
-
-	if _, ok := c.streamHandlers[key]; !ok {
-		smVersion, ok := c.sysCtx.deviceSmVersions[key.gpuUUID]
-		if !ok {
-			if key.gpuUUID != "" {
-				// Only warn when we have a device, otherwise it's expected to not find the SM version
-				// if the device UUID is empty
-				log.Warnf("SM version not found for device %s, using default", key.gpuUUID)
-			}
-			smVersion = noSmVersion
-		}
-
-		c.streamHandlers[key] = newStreamHandler(key.pid, key.containerID, smVersion, c.sysCtx)
-		c.telemetry.activeHandlers.Set(float64(len(c.streamHandlers)))
-	}
-
-	return c.streamHandlers[key], nil
+	c.streamHandlers.markProcessStreamsAsEnded(pid)
 }
 
 func (c *cudaEventConsumer) checkClosedProcesses() {
path: pkg/gpu/consumer.go
user: login: gjulianm
body: This is the part that we should probably move away from the consumer in the future.
created_at: 2025-03-18T10:38:38Z
author_association: CONTRIBUTOR
diff_hunk: @@ -241,88 +220,7 @@ func (c *cudaEventConsumer) handleGlobalEvent(header *gpuebpf.CudaEventHeader, d
 }
 
 func (c *cudaEventConsumer) handleProcessExit(pid uint32) {
-	for key, handler := range c.streamHandlers {
-		if key.pid == pid {
-			log.Debugf("Process %d ended, marking stream %d as ended", pid, key.stream)
-			// the probe is responsible for deleting the stream handler
-			_ = handler.markEnd()
-			c.telemetry.finalizedProcesses.Inc()
-		}
-	}
-}
-
-func (c *cudaEventConsumer) getStreamKey(header *gpuebpf.CudaEventHeader) (streamKey, error) {
-	pid, tid := getPidTidFromHeader(header)
-
-	var nonGlobalKey nonGlobalStreamKey
-	if header.Stream_id != 0 {
-		// Non-global stream, check if we have created it before
-		nonGlobalKey.pid = pid
-		nonGlobalKey.stream = header.Stream_id
-
-		if key, ok := c.nonGlobalStreamKeys[nonGlobalKey]; ok {
-			return key, nil
-		}
-	}
-
-	cgroup := unix.ByteSliceToString(header.Cgroup[:])
-	containerID, err := cgroups.ContainerFilter("", cgroup)
-	if err != nil {
-		// We don't want to return an error here, as we can still process the event without the container ID
-		log.Warnf("error getting container ID for cgroup %s: %s", cgroup, err)
-		c.telemetry.missingContainers.Inc("error")
-	} else if containerID == "" {
-		c.telemetry.missingContainers.Inc("missing")
-	}
-
-	key := streamKey{
-		pid:         pid,
-		stream:      header.Stream_id,
-		gpuUUID:     "",
-		containerID: containerID,
-	}
-
-	gpuDevice, err := c.sysCtx.getCurrentActiveGpuDevice(int(pid), int(tid), containerID)
-	if err != nil {
-		c.telemetry.missingDevices.Inc()
-		return streamKey{}, fmt.Errorf("Error getting GPU device for process %d: %w", pid, err)
-	}
-
-	var ret nvml.Return
-	key.gpuUUID, ret = gpuDevice.GetUUID()
-	if ret != nvml.SUCCESS {
-		return streamKey{}, fmt.Errorf("Error getting GPU UUID for process %d: %v", pid, nvml.ErrorString(ret))
-	}
-
-	if header.Stream_id != 0 {
-		c.nonGlobalStreamKeys[nonGlobalKey] = key
-	}
-
-	return key, nil
-}
-
-func (c *cudaEventConsumer) getStreamHandler(header *gpuebpf.CudaEventHeader) (*StreamHandler, error) {
-	key, err := c.getStreamKey(header)
-	if err != nil {
-		return nil, err
-	}
-
-	if _, ok := c.streamHandlers[key]; !ok {
-		smVersion, ok := c.sysCtx.deviceSmVersions[key.gpuUUID]
-		if !ok {
-			if key.gpuUUID != "" {
-				// Only warn when we have a device, otherwise it's expected to not find the SM version
-				// if the device UUID is empty
-				log.Warnf("SM version not found for device %s, using default", key.gpuUUID)
-			}
-			smVersion = noSmVersion
-		}
-
-		c.streamHandlers[key] = newStreamHandler(key.pid, key.containerID, smVersion, c.sysCtx)
-		c.telemetry.activeHandlers.Set(float64(len(c.streamHandlers)))
-	}
-
-	return c.streamHandlers[key], nil
+	c.streamHandlers.markProcessStreamsAsEnded(pid)
 }
 
 func (c *cudaEventConsumer) checkClosedProcesses() {
path: pkg/gpu/consumer.go
user: login: val06
body: yeah that makes sense, it can be encapsulated in the streamCollection. do you plan to address it in a separate PR?
created_at: 2025-03-18T10:50:06Z
author_association: CONTRIBUTOR
diff_hunk: @@ -47,41 +47,68 @@ func getStatsGeneratorForTest(t *testing.T) (*statsGenerator, map[streamKey]*Str
 	return statsGen, streamHandlers, ktime
 }
 
+func addNonglobalStream(streamHandlers *streamCollection, pid uint32, streamID uint64, gpuUUID string, containerID string) *StreamHandler {
path: pkg/gpu/stats_test.go
user: login: val06
body: addNonGlobalStream
created_at: 2025-03-18T12:10:03Z
author_association: CONTRIBUTOR
diff_hunk: @@ -47,41 +47,68 @@ func getStatsGeneratorForTest(t *testing.T) (*statsGenerator, map[streamKey]*Str
 	return statsGen, streamHandlers, ktime
 }
 
+func addNonglobalStream(streamHandlers *streamCollection, pid uint32, streamID uint64, gpuUUID string, containerID string) *StreamHandler {
path: pkg/gpu/stats_test.go
user: login: val06
body: actually i would just call it `addStream`
created_at: 2025-03-18T12:10:29Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,240 @@
+// 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 2024-present Datadog, Inc.
+
+//go:build linux_bpf
+
+package gpu
+
+import (
+	"fmt"
+	"iter"
+
+	"github.com/NVIDIA/go-nvml/pkg/nvml"
+	"golang.org/x/sys/unix"
+
+	"github.com/DataDog/datadog-agent/comp/core/telemetry"
+	gpuebpf "github.com/DataDog/datadog-agent/pkg/gpu/ebpf"
+	"github.com/DataDog/datadog-agent/pkg/util/cgroups"
+	"github.com/DataDog/datadog-agent/pkg/util/funcs"
+	"github.com/DataDog/datadog-agent/pkg/util/log"
+)
+
+// nonglobalStreamKey is a unique identifier for a CUDA stream that is not global.
+// The streams are created with a specific GPU UUID, which means that pid+stream uniquely
+// identify the GPU.
+type nonglobalStreamKey struct {
path: pkg/gpu/stream_collection.go
user: login: val06
body: let's call it `streamKey`
created_at: 2025-03-18T12:12:00Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,240 @@
+// 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 2024-present Datadog, Inc.
+
+//go:build linux_bpf
+
+package gpu
+
+import (
+	"fmt"
+	"iter"
+
+	"github.com/NVIDIA/go-nvml/pkg/nvml"
+	"golang.org/x/sys/unix"
+
+	"github.com/DataDog/datadog-agent/comp/core/telemetry"
+	gpuebpf "github.com/DataDog/datadog-agent/pkg/gpu/ebpf"
+	"github.com/DataDog/datadog-agent/pkg/util/cgroups"
+	"github.com/DataDog/datadog-agent/pkg/util/funcs"
+	"github.com/DataDog/datadog-agent/pkg/util/log"
+)
+
+// nonglobalStreamKey is a unique identifier for a CUDA stream that is not global.
+// The streams are created with a specific GPU UUID, which means that pid+stream uniquely
+// identify the GPU.
+type nonglobalStreamKey struct {
path: pkg/gpu/stream_collection.go
user: login: val06
body: you can keep the comment saying that this is a key for non global streams
created_at: 2025-03-18T12:12:16Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,240 @@
+// 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 2024-present Datadog, Inc.
+
+//go:build linux_bpf
+
+package gpu
+
+import (
+	"fmt"
+	"iter"
+
+	"github.com/NVIDIA/go-nvml/pkg/nvml"
+	"golang.org/x/sys/unix"
+
+	"github.com/DataDog/datadog-agent/comp/core/telemetry"
+	gpuebpf "github.com/DataDog/datadog-agent/pkg/gpu/ebpf"
+	"github.com/DataDog/datadog-agent/pkg/util/cgroups"
+	"github.com/DataDog/datadog-agent/pkg/util/funcs"
+	"github.com/DataDog/datadog-agent/pkg/util/log"
+)
+
+// nonglobalStreamKey is a unique identifier for a CUDA stream that is not global.
+// The streams are created with a specific GPU UUID, which means that pid+stream uniquely
+// identify the GPU.
+type nonglobalStreamKey struct {
+	pid    uint32
+	stream uint64
+}
+
+// globalStreamKey is a unique identifier for a CUDA stream that is global.
+// Global streams depend on which GPU is active when they are used, so the GPU UUID
+// needs to be part of the key
+type globalStreamKey struct {
+	pid     uint32
+	gpuUUID string
+}
+
+type streamCollection struct {
+	nonglobalStreams map[nonglobalStreamKey]*StreamHandler
path: pkg/gpu/stream_collection.go
user: login: val06
body: just `streams`
created_at: 2025-03-18T12:12:35Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,240 @@
+// 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 2024-present Datadog, Inc.
+
+//go:build linux_bpf
+
+package gpu
+
+import (
+	"fmt"
+	"iter"
+
+	"github.com/NVIDIA/go-nvml/pkg/nvml"
+	"golang.org/x/sys/unix"
+
+	"github.com/DataDog/datadog-agent/comp/core/telemetry"
+	gpuebpf "github.com/DataDog/datadog-agent/pkg/gpu/ebpf"
+	"github.com/DataDog/datadog-agent/pkg/util/cgroups"
+	"github.com/DataDog/datadog-agent/pkg/util/funcs"
+	"github.com/DataDog/datadog-agent/pkg/util/log"
+)
+
+// nonglobalStreamKey is a unique identifier for a CUDA stream that is not global.
+// The streams are created with a specific GPU UUID, which means that pid+stream uniquely
+// identify the GPU.
+type nonglobalStreamKey struct {
+	pid    uint32
+	stream uint64
+}
+
+// globalStreamKey is a unique identifier for a CUDA stream that is global.
+// Global streams depend on which GPU is active when they are used, so the GPU UUID
+// needs to be part of the key
+type globalStreamKey struct {
+	pid     uint32
+	gpuUUID string
+}
+
+type streamCollection struct {
+	nonglobalStreams map[nonglobalStreamKey]*StreamHandler
+	globalStreams    map[globalStreamKey]*StreamHandler
+	sysCtx           *systemContext
+	telemetry        *streamCollectionTelemetry
+}
+
+type streamCollectionTelemetry struct {
+	missingContainers  telemetry.Counter
+	missingDevices     telemetry.Counter
+	finalizedProcesses telemetry.Counter
+	activeHandlers     telemetry.Gauge
+	removedHandlers    telemetry.Counter
+}
+
+func newStreamCollection(sysCtx *systemContext, telemetry telemetry.Component) *streamCollection {
+	return &streamCollection{
+		nonglobalStreams: make(map[nonglobalStreamKey]*StreamHandler),
+		globalStreams:    make(map[globalStreamKey]*StreamHandler),
+		sysCtx:           sysCtx,
+		telemetry:        newStreamCollectionTelemetry(telemetry),
+	}
+}
+
+func newStreamCollectionTelemetry(tm telemetry.Component) *streamCollectionTelemetry {
+	subsystem := gpuTelemetryModule + "__streams"
+
+	return &streamCollectionTelemetry{
+		missingContainers:  tm.NewCounter(subsystem, "missing_containers", []string{"reason"}, "Number of missing containers"),
+		missingDevices:     tm.NewCounter(subsystem, "missing_devices", nil, "Number of failures to get GPU devices for a stream"),
+		finalizedProcesses: tm.NewCounter(subsystem, "finalized_processes", nil, "Number of processes that have ended"),
+		activeHandlers:     tm.NewGauge(subsystem, "active_handlers", nil, "Number of active stream handlers"),
+		removedHandlers:    tm.NewCounter(subsystem, "removed_handlers", nil, "Number of removed stream handlers"),
path: pkg/gpu/stream_collection.go
user: login: val06
body: add `device` tag to those telemetry counters
created_at: 2025-03-18T12:15:21Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,240 @@
+// 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 2024-present Datadog, Inc.
+
+//go:build linux_bpf
+
+package gpu
+
+import (
+	"fmt"
+	"iter"
+
+	"github.com/NVIDIA/go-nvml/pkg/nvml"
+	"golang.org/x/sys/unix"
+
+	"github.com/DataDog/datadog-agent/comp/core/telemetry"
+	gpuebpf "github.com/DataDog/datadog-agent/pkg/gpu/ebpf"
+	"github.com/DataDog/datadog-agent/pkg/util/cgroups"
+	"github.com/DataDog/datadog-agent/pkg/util/funcs"
+	"github.com/DataDog/datadog-agent/pkg/util/log"
+)
+
+// nonglobalStreamKey is a unique identifier for a CUDA stream that is not global.
+// The streams are created with a specific GPU UUID, which means that pid+stream uniquely
+// identify the GPU.
+type nonglobalStreamKey struct {
+	pid    uint32
+	stream uint64
+}
+
+// globalStreamKey is a unique identifier for a CUDA stream that is global.
+// Global streams depend on which GPU is active when they are used, so the GPU UUID
+// needs to be part of the key
+type globalStreamKey struct {
+	pid     uint32
+	gpuUUID string
+}
+
+type streamCollection struct {
+	nonglobalStreams map[nonglobalStreamKey]*StreamHandler
+	globalStreams    map[globalStreamKey]*StreamHandler
+	sysCtx           *systemContext
+	telemetry        *streamCollectionTelemetry
+}
+
+type streamCollectionTelemetry struct {
+	missingContainers  telemetry.Counter
+	missingDevices     telemetry.Counter
+	finalizedProcesses telemetry.Counter
+	activeHandlers     telemetry.Gauge
+	removedHandlers    telemetry.Counter
+}
+
+func newStreamCollection(sysCtx *systemContext, telemetry telemetry.Component) *streamCollection {
+	return &streamCollection{
+		nonglobalStreams: make(map[nonglobalStreamKey]*StreamHandler),
+		globalStreams:    make(map[globalStreamKey]*StreamHandler),
+		sysCtx:           sysCtx,
+		telemetry:        newStreamCollectionTelemetry(telemetry),
+	}
+}
+
+func newStreamCollectionTelemetry(tm telemetry.Component) *streamCollectionTelemetry {
+	subsystem := gpuTelemetryModule + "__streams"
+
+	return &streamCollectionTelemetry{
+		missingContainers:  tm.NewCounter(subsystem, "missing_containers", []string{"reason"}, "Number of missing containers"),
+		missingDevices:     tm.NewCounter(subsystem, "missing_devices", nil, "Number of failures to get GPU devices for a stream"),
+		finalizedProcesses: tm.NewCounter(subsystem, "finalized_processes", nil, "Number of processes that have ended"),
+		activeHandlers:     tm.NewGauge(subsystem, "active_handlers", nil, "Number of active stream handlers"),
+		removedHandlers:    tm.NewCounter(subsystem, "removed_handlers", nil, "Number of removed stream handlers"),
+	}
+}
+
+// getStream returns a StreamHandler for a given CUDA stream.
+func (sc *streamCollection) getStream(header *gpuebpf.CudaEventHeader) (*StreamHandler, error) {
+	if header.Stream_id == 0 {
+		return sc.getGlobalStream(header)
+	}
+
+	return sc.getNonglobalStream(header)
+}
+
+func (sc *streamCollection) getGlobalStream(header *gpuebpf.CudaEventHeader) (*StreamHandler, error) {
+	pid, tid := getPidTidFromHeader(header)
+	memoizedContainerID := sc.memoizedContainerID(header)
+
+	// Global streams depend on which GPU is active when they are used, so we need to get the current active GPU device
+	// The expensive step here is the container ID parsing, but we don't always need to do it so we pass a function
+	// that can be called to retrieve it only when needed
+	gpuUUID, err := sc.getActiveDevice(int(pid), int(tid), memoizedContainerID)
+	if err != nil {
+		sc.telemetry.missingDevices.Inc()
+		return nil, err
+	}
+
+	key := globalStreamKey{
+		pid:     pid,
+		gpuUUID: gpuUUID,
+	}
+
+	stream, ok := sc.globalStreams[key]
+	if !ok {
+		stream = sc.createStream(header, &gpuUUID, memoizedContainerID)
+		sc.globalStreams[key] = stream
+		sc.telemetry.activeHandlers.Set(float64(sc.streamCount()))
+	}
+
+	return stream, nil
+}
+
+func (sc *streamCollection) getNonglobalStream(header *gpuebpf.CudaEventHeader) (*StreamHandler, error) {
path: pkg/gpu/stream_collection.go
user: login: val06
body: `getNonGlobalStream`
created_at: 2025-03-18T12:16:44Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,240 @@
+// 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 2024-present Datadog, Inc.
+
+//go:build linux_bpf
+
+package gpu
+
+import (
+	"fmt"
+	"iter"
+
+	"github.com/NVIDIA/go-nvml/pkg/nvml"
+	"golang.org/x/sys/unix"
+
+	"github.com/DataDog/datadog-agent/comp/core/telemetry"
+	gpuebpf "github.com/DataDog/datadog-agent/pkg/gpu/ebpf"
+	"github.com/DataDog/datadog-agent/pkg/util/cgroups"
+	"github.com/DataDog/datadog-agent/pkg/util/funcs"
+	"github.com/DataDog/datadog-agent/pkg/util/log"
+)
+
+// nonglobalStreamKey is a unique identifier for a CUDA stream that is not global.
+// The streams are created with a specific GPU UUID, which means that pid+stream uniquely
+// identify the GPU.
+type nonglobalStreamKey struct {
+	pid    uint32
+	stream uint64
+}
+
+// globalStreamKey is a unique identifier for a CUDA stream that is global.
+// Global streams depend on which GPU is active when they are used, so the GPU UUID
+// needs to be part of the key
+type globalStreamKey struct {
+	pid     uint32
+	gpuUUID string
+}
+
+type streamCollection struct {
+	nonglobalStreams map[nonglobalStreamKey]*StreamHandler
+	globalStreams    map[globalStreamKey]*StreamHandler
+	sysCtx           *systemContext
+	telemetry        *streamCollectionTelemetry
+}
+
+type streamCollectionTelemetry struct {
+	missingContainers  telemetry.Counter
+	missingDevices     telemetry.Counter
+	finalizedProcesses telemetry.Counter
+	activeHandlers     telemetry.Gauge
+	removedHandlers    telemetry.Counter
+}
+
+func newStreamCollection(sysCtx *systemContext, telemetry telemetry.Component) *streamCollection {
+	return &streamCollection{
+		nonglobalStreams: make(map[nonglobalStreamKey]*StreamHandler),
+		globalStreams:    make(map[globalStreamKey]*StreamHandler),
+		sysCtx:           sysCtx,
+		telemetry:        newStreamCollectionTelemetry(telemetry),
+	}
+}
+
+func newStreamCollectionTelemetry(tm telemetry.Component) *streamCollectionTelemetry {
+	subsystem := gpuTelemetryModule + "__streams"
+
+	return &streamCollectionTelemetry{
+		missingContainers:  tm.NewCounter(subsystem, "missing_containers", []string{"reason"}, "Number of missing containers"),
+		missingDevices:     tm.NewCounter(subsystem, "missing_devices", nil, "Number of failures to get GPU devices for a stream"),
+		finalizedProcesses: tm.NewCounter(subsystem, "finalized_processes", nil, "Number of processes that have ended"),
+		activeHandlers:     tm.NewGauge(subsystem, "active_handlers", nil, "Number of active stream handlers"),
+		removedHandlers:    tm.NewCounter(subsystem, "removed_handlers", nil, "Number of removed stream handlers"),
+	}
+}
+
+// getStream returns a StreamHandler for a given CUDA stream.
+func (sc *streamCollection) getStream(header *gpuebpf.CudaEventHeader) (*StreamHandler, error) {
+	if header.Stream_id == 0 {
+		return sc.getGlobalStream(header)
+	}
+
+	return sc.getNonglobalStream(header)
+}
+
+func (sc *streamCollection) getGlobalStream(header *gpuebpf.CudaEventHeader) (*StreamHandler, error) {
+	pid, tid := getPidTidFromHeader(header)
+	memoizedContainerID := sc.memoizedContainerID(header)
+
+	// Global streams depend on which GPU is active when they are used, so we need to get the current active GPU device
+	// The expensive step here is the container ID parsing, but we don't always need to do it so we pass a function
+	// that can be called to retrieve it only when needed
+	gpuUUID, err := sc.getActiveDevice(int(pid), int(tid), memoizedContainerID)
+	if err != nil {
+		sc.telemetry.missingDevices.Inc()
+		return nil, err
+	}
+
+	key := globalStreamKey{
+		pid:     pid,
+		gpuUUID: gpuUUID,
+	}
+
+	stream, ok := sc.globalStreams[key]
+	if !ok {
+		stream = sc.createStream(header, &gpuUUID, memoizedContainerID)
+		sc.globalStreams[key] = stream
+		sc.telemetry.activeHandlers.Set(float64(sc.streamCount()))
+	}
+
+	return stream, nil
+}
+
+func (sc *streamCollection) getNonglobalStream(header *gpuebpf.CudaEventHeader) (*StreamHandler, error) {
+	pid, _ := getPidTidFromHeader(header)
+
+	key := nonglobalStreamKey{
+		pid:    pid,
+		stream: header.Stream_id,
+	}
+
+	stream, ok := sc.nonglobalStreams[key]
+	if !ok {
+		stream = sc.createStream(header, nil, sc.memoizedContainerID(header))
+		sc.nonglobalStreams[key] = stream
+		sc.telemetry.activeHandlers.Set(float64(sc.streamCount()))
+	}
+
+	return stream, nil
+}
+
+// createStream creates a new StreamHandler for a given CUDA stream.
+// If the GPU UUID is not provided (it's nil), it will be retrieved from the system context.
+func (sc *streamCollection) createStream(header *gpuebpf.CudaEventHeader, gpuUUID *string, containerIDFunc func() string) *StreamHandler {
path: pkg/gpu/stream_collection.go
user: login: val06
body: nit: rename to `createStreamHandler`
created_at: 2025-03-18T12:17:08Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,240 @@
+// 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 2024-present Datadog, Inc.
+
+//go:build linux_bpf
+
+package gpu
+
+import (
+	"fmt"
+	"iter"
+
+	"github.com/NVIDIA/go-nvml/pkg/nvml"
+	"golang.org/x/sys/unix"
+
+	"github.com/DataDog/datadog-agent/comp/core/telemetry"
+	gpuebpf "github.com/DataDog/datadog-agent/pkg/gpu/ebpf"
+	"github.com/DataDog/datadog-agent/pkg/util/cgroups"
+	"github.com/DataDog/datadog-agent/pkg/util/funcs"
+	"github.com/DataDog/datadog-agent/pkg/util/log"
+)
+
+// nonglobalStreamKey is a unique identifier for a CUDA stream that is not global.
+// The streams are created with a specific GPU UUID, which means that pid+stream uniquely
+// identify the GPU.
+type nonglobalStreamKey struct {
+	pid    uint32
+	stream uint64
+}
+
+// globalStreamKey is a unique identifier for a CUDA stream that is global.
+// Global streams depend on which GPU is active when they are used, so the GPU UUID
+// needs to be part of the key
+type globalStreamKey struct {
+	pid     uint32
+	gpuUUID string
+}
+
+type streamCollection struct {
+	nonglobalStreams map[nonglobalStreamKey]*StreamHandler
+	globalStreams    map[globalStreamKey]*StreamHandler
+	sysCtx           *systemContext
+	telemetry        *streamCollectionTelemetry
+}
+
+type streamCollectionTelemetry struct {
+	missingContainers  telemetry.Counter
+	missingDevices     telemetry.Counter
+	finalizedProcesses telemetry.Counter
+	activeHandlers     telemetry.Gauge
+	removedHandlers    telemetry.Counter
+}
+
+func newStreamCollection(sysCtx *systemContext, telemetry telemetry.Component) *streamCollection {
+	return &streamCollection{
+		nonglobalStreams: make(map[nonglobalStreamKey]*StreamHandler),
+		globalStreams:    make(map[globalStreamKey]*StreamHandler),
+		sysCtx:           sysCtx,
+		telemetry:        newStreamCollectionTelemetry(telemetry),
+	}
+}
+
+func newStreamCollectionTelemetry(tm telemetry.Component) *streamCollectionTelemetry {
+	subsystem := gpuTelemetryModule + "__streams"
+
+	return &streamCollectionTelemetry{
+		missingContainers:  tm.NewCounter(subsystem, "missing_containers", []string{"reason"}, "Number of missing containers"),
+		missingDevices:     tm.NewCounter(subsystem, "missing_devices", nil, "Number of failures to get GPU devices for a stream"),
+		finalizedProcesses: tm.NewCounter(subsystem, "finalized_processes", nil, "Number of processes that have ended"),
+		activeHandlers:     tm.NewGauge(subsystem, "active_handlers", nil, "Number of active stream handlers"),
+		removedHandlers:    tm.NewCounter(subsystem, "removed_handlers", nil, "Number of removed stream handlers"),
+	}
+}
+
+// getStream returns a StreamHandler for a given CUDA stream.
+func (sc *streamCollection) getStream(header *gpuebpf.CudaEventHeader) (*StreamHandler, error) {
+	if header.Stream_id == 0 {
+		return sc.getGlobalStream(header)
+	}
+
+	return sc.getNonglobalStream(header)
+}
+
+func (sc *streamCollection) getGlobalStream(header *gpuebpf.CudaEventHeader) (*StreamHandler, error) {
+	pid, tid := getPidTidFromHeader(header)
+	memoizedContainerID := sc.memoizedContainerID(header)
+
+	// Global streams depend on which GPU is active when they are used, so we need to get the current active GPU device
+	// The expensive step here is the container ID parsing, but we don't always need to do it so we pass a function
+	// that can be called to retrieve it only when needed
+	gpuUUID, err := sc.getActiveDevice(int(pid), int(tid), memoizedContainerID)
+	if err != nil {
+		sc.telemetry.missingDevices.Inc()
+		return nil, err
+	}
+
+	key := globalStreamKey{
+		pid:     pid,
+		gpuUUID: gpuUUID,
+	}
+
+	stream, ok := sc.globalStreams[key]
+	if !ok {
+		stream = sc.createStream(header, &gpuUUID, memoizedContainerID)
+		sc.globalStreams[key] = stream
+		sc.telemetry.activeHandlers.Set(float64(sc.streamCount()))
+	}
+
+	return stream, nil
+}
+
+func (sc *streamCollection) getNonglobalStream(header *gpuebpf.CudaEventHeader) (*StreamHandler, error) {
+	pid, _ := getPidTidFromHeader(header)
+
+	key := nonglobalStreamKey{
+		pid:    pid,
+		stream: header.Stream_id,
+	}
+
+	stream, ok := sc.nonglobalStreams[key]
+	if !ok {
+		stream = sc.createStream(header, nil, sc.memoizedContainerID(header))
+		sc.nonglobalStreams[key] = stream
+		sc.telemetry.activeHandlers.Set(float64(sc.streamCount()))
+	}
+
+	return stream, nil
+}
+
+// createStream creates a new StreamHandler for a given CUDA stream.
+// If the GPU UUID is not provided (it's nil), it will be retrieved from the system context.
+func (sc *streamCollection) createStream(header *gpuebpf.CudaEventHeader, gpuUUID *string, containerIDFunc func() string) *StreamHandler {
+	pid, tid := getPidTidFromHeader(header)
+	metadata := streamMetadata{
+		pid:         pid,
+		streamID:    header.Stream_id,
+		containerID: containerIDFunc(),
+	}
+
+	if gpuUUID != nil {
+		metadata.gpuUUID = *gpuUUID
+	} else {
+		var err error
+		metadata.gpuUUID, err = sc.getActiveDevice(int(pid), int(tid), containerIDFunc)
+		if err != nil {
+			log.Warnf("error getting GPU UUID for process %d: %s", pid, err)
+			sc.telemetry.missingDevices.Inc()
+			return nil
+		}
+	}
+
+	metadata.smVersion = sc.sysCtx.deviceSmVersions[metadata.gpuUUID]
+
+	return newStreamHandler(metadata, sc.sysCtx)
+}
+
+// memoizedContainerID returns a function that memoizes the container ID for a given CUDA stream.
+// It's memoized because it's an expensive operation that we don't always need to do.
+func (sc *streamCollection) memoizedContainerID(header *gpuebpf.CudaEventHeader) func() string {
+	return funcs.MemoizeNoErrorUnsafe(func() string {
+		cgroup := unix.ByteSliceToString(header.Cgroup[:])
+		containerID, err := cgroups.ContainerFilter("", cgroup)
+		if err != nil {
+			log.Warnf("error getting container ID for cgroup %s: %s", cgroup, err)
+
+			sc.telemetry.missingContainers.Inc("error")
+			return ""
+		}
+
+		if containerID == "" {
+			sc.telemetry.missingContainers.Inc("missing")
+		}
+
+		return containerID
+	})
+}
+
+// getActiveDevice returns the GPU UUID for the current active GPU device for a given process and thread.
+func (sc *streamCollection) getActiveDevice(pid int, tid int, containerIDFunc func() string) (string, error) {
+	device, err := sc.sysCtx.getCurrentActiveGpuDevice(pid, tid, containerIDFunc)
+	if err != nil {
+		return "", err
+	}
+
+	uuid, ret := device.GetUUID()
+	if ret != nvml.SUCCESS {
+		return "", fmt.Errorf("error getting GPU UUID for process %d: %v", pid, nvml.ErrorString(ret))
+	}
+
+	return uuid, nil
+}
+
+func (sc *streamCollection) allStreams() iter.Seq[*StreamHandler] {
+	return func(yield func(*StreamHandler) bool) {
+		for _, stream := range sc.nonglobalStreams {
+			if !yield(stream) {
+				return
+			}
+		}
+
+		for _, stream := range sc.globalStreams {
+			if !yield(stream) {
+				return
+			}
+		}
+	}
+}
+
+func (sc *streamCollection) streamCount() int {
+	return len(sc.nonglobalStreams) + len(sc.globalStreams)
+}
+
+func (sc *streamCollection) markProcessStreamsAsEnded(pid uint32) {
path: pkg/gpu/stream_collection.go
user: login: val06
body: rename to `finalizeProcessStreams`
created_at: 2025-03-18T12:18:26Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,240 @@
+// 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 2024-present Datadog, Inc.
+
+//go:build linux_bpf
+
+package gpu
+
+import (
+	"fmt"
+	"iter"
+
+	"github.com/NVIDIA/go-nvml/pkg/nvml"
+	"golang.org/x/sys/unix"
+
+	"github.com/DataDog/datadog-agent/comp/core/telemetry"
+	gpuebpf "github.com/DataDog/datadog-agent/pkg/gpu/ebpf"
+	"github.com/DataDog/datadog-agent/pkg/util/cgroups"
+	"github.com/DataDog/datadog-agent/pkg/util/funcs"
+	"github.com/DataDog/datadog-agent/pkg/util/log"
+)
+
+// nonglobalStreamKey is a unique identifier for a CUDA stream that is not global.
+// The streams are created with a specific GPU UUID, which means that pid+stream uniquely
+// identify the GPU.
+type nonglobalStreamKey struct {
+	pid    uint32
+	stream uint64
+}
+
+// globalStreamKey is a unique identifier for a CUDA stream that is global.
+// Global streams depend on which GPU is active when they are used, so the GPU UUID
+// needs to be part of the key
+type globalStreamKey struct {
+	pid     uint32
+	gpuUUID string
+}
+
+type streamCollection struct {
+	nonglobalStreams map[nonglobalStreamKey]*StreamHandler
+	globalStreams    map[globalStreamKey]*StreamHandler
+	sysCtx           *systemContext
+	telemetry        *streamCollectionTelemetry
+}
+
+type streamCollectionTelemetry struct {
+	missingContainers  telemetry.Counter
+	missingDevices     telemetry.Counter
+	finalizedProcesses telemetry.Counter
+	activeHandlers     telemetry.Gauge
+	removedHandlers    telemetry.Counter
+}
+
+func newStreamCollection(sysCtx *systemContext, telemetry telemetry.Component) *streamCollection {
+	return &streamCollection{
+		nonglobalStreams: make(map[nonglobalStreamKey]*StreamHandler),
+		globalStreams:    make(map[globalStreamKey]*StreamHandler),
+		sysCtx:           sysCtx,
+		telemetry:        newStreamCollectionTelemetry(telemetry),
+	}
+}
+
+func newStreamCollectionTelemetry(tm telemetry.Component) *streamCollectionTelemetry {
+	subsystem := gpuTelemetryModule + "__streams"
+
+	return &streamCollectionTelemetry{
+		missingContainers:  tm.NewCounter(subsystem, "missing_containers", []string{"reason"}, "Number of missing containers"),
+		missingDevices:     tm.NewCounter(subsystem, "missing_devices", nil, "Number of failures to get GPU devices for a stream"),
+		finalizedProcesses: tm.NewCounter(subsystem, "finalized_processes", nil, "Number of processes that have ended"),
+		activeHandlers:     tm.NewGauge(subsystem, "active_handlers", nil, "Number of active stream handlers"),
+		removedHandlers:    tm.NewCounter(subsystem, "removed_handlers", nil, "Number of removed stream handlers"),
+	}
+}
+
+// getStream returns a StreamHandler for a given CUDA stream.
+func (sc *streamCollection) getStream(header *gpuebpf.CudaEventHeader) (*StreamHandler, error) {
+	if header.Stream_id == 0 {
+		return sc.getGlobalStream(header)
+	}
+
+	return sc.getNonglobalStream(header)
+}
+
+func (sc *streamCollection) getGlobalStream(header *gpuebpf.CudaEventHeader) (*StreamHandler, error) {
+	pid, tid := getPidTidFromHeader(header)
+	memoizedContainerID := sc.memoizedContainerID(header)
+
+	// Global streams depend on which GPU is active when they are used, so we need to get the current active GPU device
+	// The expensive step here is the container ID parsing, but we don't always need to do it so we pass a function
+	// that can be called to retrieve it only when needed
+	gpuUUID, err := sc.getActiveDevice(int(pid), int(tid), memoizedContainerID)
+	if err != nil {
+		sc.telemetry.missingDevices.Inc()
+		return nil, err
+	}
+
+	key := globalStreamKey{
+		pid:     pid,
+		gpuUUID: gpuUUID,
+	}
+
+	stream, ok := sc.globalStreams[key]
+	if !ok {
+		stream = sc.createStream(header, &gpuUUID, memoizedContainerID)
+		sc.globalStreams[key] = stream
+		sc.telemetry.activeHandlers.Set(float64(sc.streamCount()))
+	}
+
+	return stream, nil
+}
+
+func (sc *streamCollection) getNonglobalStream(header *gpuebpf.CudaEventHeader) (*StreamHandler, error) {
+	pid, _ := getPidTidFromHeader(header)
+
+	key := nonglobalStreamKey{
+		pid:    pid,
+		stream: header.Stream_id,
+	}
+
+	stream, ok := sc.nonglobalStreams[key]
+	if !ok {
+		stream = sc.createStream(header, nil, sc.memoizedContainerID(header))
+		sc.nonglobalStreams[key] = stream
+		sc.telemetry.activeHandlers.Set(float64(sc.streamCount()))
+	}
+
+	return stream, nil
+}
+
+// createStream creates a new StreamHandler for a given CUDA stream.
+// If the GPU UUID is not provided (it's nil), it will be retrieved from the system context.
+func (sc *streamCollection) createStream(header *gpuebpf.CudaEventHeader, gpuUUID *string, containerIDFunc func() string) *StreamHandler {
+	pid, tid := getPidTidFromHeader(header)
+	metadata := streamMetadata{
+		pid:         pid,
+		streamID:    header.Stream_id,
+		containerID: containerIDFunc(),
+	}
+
+	if gpuUUID != nil {
+		metadata.gpuUUID = *gpuUUID
+	} else {
+		var err error
+		metadata.gpuUUID, err = sc.getActiveDevice(int(pid), int(tid), containerIDFunc)
+		if err != nil {
+			log.Warnf("error getting GPU UUID for process %d: %s", pid, err)
+			sc.telemetry.missingDevices.Inc()
+			return nil
+		}
+	}
+
+	metadata.smVersion = sc.sysCtx.deviceSmVersions[metadata.gpuUUID]
+
+	return newStreamHandler(metadata, sc.sysCtx)
+}
+
+// memoizedContainerID returns a function that memoizes the container ID for a given CUDA stream.
+// It's memoized because it's an expensive operation that we don't always need to do.
+func (sc *streamCollection) memoizedContainerID(header *gpuebpf.CudaEventHeader) func() string {
+	return funcs.MemoizeNoErrorUnsafe(func() string {
+		cgroup := unix.ByteSliceToString(header.Cgroup[:])
+		containerID, err := cgroups.ContainerFilter("", cgroup)
+		if err != nil {
+			log.Warnf("error getting container ID for cgroup %s: %s", cgroup, err)
+
+			sc.telemetry.missingContainers.Inc("error")
+			return ""
+		}
+
+		if containerID == "" {
+			sc.telemetry.missingContainers.Inc("missing")
+		}
+
+		return containerID
+	})
+}
+
+// getActiveDevice returns the GPU UUID for the current active GPU device for a given process and thread.
+func (sc *streamCollection) getActiveDevice(pid int, tid int, containerIDFunc func() string) (string, error) {
+	device, err := sc.sysCtx.getCurrentActiveGpuDevice(pid, tid, containerIDFunc)
+	if err != nil {
+		return "", err
+	}
+
+	uuid, ret := device.GetUUID()
+	if ret != nvml.SUCCESS {
+		return "", fmt.Errorf("error getting GPU UUID for process %d: %v", pid, nvml.ErrorString(ret))
+	}
+
+	return uuid, nil
+}
+
+func (sc *streamCollection) allStreams() iter.Seq[*StreamHandler] {
+	return func(yield func(*StreamHandler) bool) {
+		for _, stream := range sc.nonglobalStreams {
+			if !yield(stream) {
+				return
+			}
+		}
+
+		for _, stream := range sc.globalStreams {
+			if !yield(stream) {
+				return
+			}
+		}
+	}
+}
+
+func (sc *streamCollection) streamCount() int {
+	return len(sc.nonglobalStreams) + len(sc.globalStreams)
+}
+
+func (sc *streamCollection) markProcessStreamsAsEnded(pid uint32) {
+	for handler := range sc.allStreams() {
+		if handler.metadata.pid == pid {
+			log.Debugf("Process %d ended, marking stream %d as ended", pid, handler.metadata.streamID)
+			_ = handler.markEnd()
+			sc.telemetry.finalizedProcesses.Inc()
+		}
+	}
+}
+
+func (sc *streamCollection) cleanupFinishedHandlers() {
path: pkg/gpu/stream_collection.go
user: login: val06
body: why not just `clean` ?
created_at: 2025-03-18T12:18:46Z
author_association: CONTRIBUTOR
diff_hunk: @@ -241,88 +220,7 @@ func (c *cudaEventConsumer) handleGlobalEvent(header *gpuebpf.CudaEventHeader, d
 }
 
 func (c *cudaEventConsumer) handleProcessExit(pid uint32) {
-	for key, handler := range c.streamHandlers {
-		if key.pid == pid {
-			log.Debugf("Process %d ended, marking stream %d as ended", pid, key.stream)
-			// the probe is responsible for deleting the stream handler
-			_ = handler.markEnd()
-			c.telemetry.finalizedProcesses.Inc()
-		}
-	}
-}
-
-func (c *cudaEventConsumer) getStreamKey(header *gpuebpf.CudaEventHeader) (streamKey, error) {
-	pid, tid := getPidTidFromHeader(header)
-
-	var nonGlobalKey nonGlobalStreamKey
-	if header.Stream_id != 0 {
-		// Non-global stream, check if we have created it before
-		nonGlobalKey.pid = pid
-		nonGlobalKey.stream = header.Stream_id
-
-		if key, ok := c.nonGlobalStreamKeys[nonGlobalKey]; ok {
-			return key, nil
-		}
-	}
-
-	cgroup := unix.ByteSliceToString(header.Cgroup[:])
-	containerID, err := cgroups.ContainerFilter("", cgroup)
-	if err != nil {
-		// We don't want to return an error here, as we can still process the event without the container ID
-		log.Warnf("error getting container ID for cgroup %s: %s", cgroup, err)
-		c.telemetry.missingContainers.Inc("error")
-	} else if containerID == "" {
-		c.telemetry.missingContainers.Inc("missing")
-	}
-
-	key := streamKey{
-		pid:         pid,
-		stream:      header.Stream_id,
-		gpuUUID:     "",
-		containerID: containerID,
-	}
-
-	gpuDevice, err := c.sysCtx.getCurrentActiveGpuDevice(int(pid), int(tid), containerID)
-	if err != nil {
-		c.telemetry.missingDevices.Inc()
-		return streamKey{}, fmt.Errorf("Error getting GPU device for process %d: %w", pid, err)
-	}
-
-	var ret nvml.Return
-	key.gpuUUID, ret = gpuDevice.GetUUID()
-	if ret != nvml.SUCCESS {
-		return streamKey{}, fmt.Errorf("Error getting GPU UUID for process %d: %v", pid, nvml.ErrorString(ret))
-	}
-
-	if header.Stream_id != 0 {
-		c.nonGlobalStreamKeys[nonGlobalKey] = key
-	}
-
-	return key, nil
-}
-
-func (c *cudaEventConsumer) getStreamHandler(header *gpuebpf.CudaEventHeader) (*StreamHandler, error) {
-	key, err := c.getStreamKey(header)
-	if err != nil {
-		return nil, err
-	}
-
-	if _, ok := c.streamHandlers[key]; !ok {
-		smVersion, ok := c.sysCtx.deviceSmVersions[key.gpuUUID]
-		if !ok {
-			if key.gpuUUID != "" {
-				// Only warn when we have a device, otherwise it's expected to not find the SM version
-				// if the device UUID is empty
-				log.Warnf("SM version not found for device %s, using default", key.gpuUUID)
-			}
-			smVersion = noSmVersion
-		}
-
-		c.streamHandlers[key] = newStreamHandler(key.pid, key.containerID, smVersion, c.sysCtx)
-		c.telemetry.activeHandlers.Set(float64(len(c.streamHandlers)))
-	}
-
-	return c.streamHandlers[key], nil
+	c.streamHandlers.markProcessStreamsAsEnded(pid)
 }
 
 func (c *cudaEventConsumer) checkClosedProcesses() {
path: pkg/gpu/consumer.go
user: login: gjulianm
body: Yes, once we get a better handler for process lifetime 
created_at: 2025-03-19T11:20:41Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,240 @@
+// 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 2024-present Datadog, Inc.
+
+//go:build linux_bpf
+
+package gpu
+
+import (
+	"fmt"
+	"iter"
+
+	"github.com/NVIDIA/go-nvml/pkg/nvml"
+	"golang.org/x/sys/unix"
+
+	"github.com/DataDog/datadog-agent/comp/core/telemetry"
+	gpuebpf "github.com/DataDog/datadog-agent/pkg/gpu/ebpf"
+	"github.com/DataDog/datadog-agent/pkg/util/cgroups"
+	"github.com/DataDog/datadog-agent/pkg/util/funcs"
+	"github.com/DataDog/datadog-agent/pkg/util/log"
+)
+
+// nonglobalStreamKey is a unique identifier for a CUDA stream that is not global.
+// The streams are created with a specific GPU UUID, which means that pid+stream uniquely
+// identify the GPU.
+type nonglobalStreamKey struct {
+	pid    uint32
+	stream uint64
+}
+
+// globalStreamKey is a unique identifier for a CUDA stream that is global.
+// Global streams depend on which GPU is active when they are used, so the GPU UUID
+// needs to be part of the key
+type globalStreamKey struct {
+	pid     uint32
+	gpuUUID string
+}
+
+type streamCollection struct {
+	nonglobalStreams map[nonglobalStreamKey]*StreamHandler
+	globalStreams    map[globalStreamKey]*StreamHandler
+	sysCtx           *systemContext
+	telemetry        *streamCollectionTelemetry
+}
+
+type streamCollectionTelemetry struct {
+	missingContainers  telemetry.Counter
+	missingDevices     telemetry.Counter
+	finalizedProcesses telemetry.Counter
+	activeHandlers     telemetry.Gauge
+	removedHandlers    telemetry.Counter
+}
+
+func newStreamCollection(sysCtx *systemContext, telemetry telemetry.Component) *streamCollection {
+	return &streamCollection{
+		nonglobalStreams: make(map[nonglobalStreamKey]*StreamHandler),
+		globalStreams:    make(map[globalStreamKey]*StreamHandler),
+		sysCtx:           sysCtx,
+		telemetry:        newStreamCollectionTelemetry(telemetry),
+	}
+}
+
+func newStreamCollectionTelemetry(tm telemetry.Component) *streamCollectionTelemetry {
+	subsystem := gpuTelemetryModule + "__streams"
+
+	return &streamCollectionTelemetry{
+		missingContainers:  tm.NewCounter(subsystem, "missing_containers", []string{"reason"}, "Number of missing containers"),
+		missingDevices:     tm.NewCounter(subsystem, "missing_devices", nil, "Number of failures to get GPU devices for a stream"),
+		finalizedProcesses: tm.NewCounter(subsystem, "finalized_processes", nil, "Number of processes that have ended"),
+		activeHandlers:     tm.NewGauge(subsystem, "active_handlers", nil, "Number of active stream handlers"),
+		removedHandlers:    tm.NewCounter(subsystem, "removed_handlers", nil, "Number of removed stream handlers"),
+	}
+}
+
+// getStream returns a StreamHandler for a given CUDA stream.
+func (sc *streamCollection) getStream(header *gpuebpf.CudaEventHeader) (*StreamHandler, error) {
+	if header.Stream_id == 0 {
+		return sc.getGlobalStream(header)
+	}
+
+	return sc.getNonglobalStream(header)
+}
+
+func (sc *streamCollection) getGlobalStream(header *gpuebpf.CudaEventHeader) (*StreamHandler, error) {
+	pid, tid := getPidTidFromHeader(header)
+	memoizedContainerID := sc.memoizedContainerID(header)
+
+	// Global streams depend on which GPU is active when they are used, so we need to get the current active GPU device
+	// The expensive step here is the container ID parsing, but we don't always need to do it so we pass a function
+	// that can be called to retrieve it only when needed
+	gpuUUID, err := sc.getActiveDevice(int(pid), int(tid), memoizedContainerID)
+	if err != nil {
+		sc.telemetry.missingDevices.Inc()
+		return nil, err
+	}
+
+	key := globalStreamKey{
+		pid:     pid,
+		gpuUUID: gpuUUID,
+	}
+
+	stream, ok := sc.globalStreams[key]
+	if !ok {
+		stream = sc.createStream(header, &gpuUUID, memoizedContainerID)
+		sc.globalStreams[key] = stream
+		sc.telemetry.activeHandlers.Set(float64(sc.streamCount()))
+	}
+
+	return stream, nil
+}
+
+func (sc *streamCollection) getNonglobalStream(header *gpuebpf.CudaEventHeader) (*StreamHandler, error) {
+	pid, _ := getPidTidFromHeader(header)
+
+	key := nonglobalStreamKey{
+		pid:    pid,
+		stream: header.Stream_id,
+	}
+
+	stream, ok := sc.nonglobalStreams[key]
+	if !ok {
+		stream = sc.createStream(header, nil, sc.memoizedContainerID(header))
+		sc.nonglobalStreams[key] = stream
+		sc.telemetry.activeHandlers.Set(float64(sc.streamCount()))
+	}
+
+	return stream, nil
+}
+
+// createStream creates a new StreamHandler for a given CUDA stream.
+// If the GPU UUID is not provided (it's nil), it will be retrieved from the system context.
+func (sc *streamCollection) createStream(header *gpuebpf.CudaEventHeader, gpuUUID *string, containerIDFunc func() string) *StreamHandler {
+	pid, tid := getPidTidFromHeader(header)
+	metadata := streamMetadata{
+		pid:         pid,
+		streamID:    header.Stream_id,
+		containerID: containerIDFunc(),
+	}
+
+	if gpuUUID != nil {
+		metadata.gpuUUID = *gpuUUID
+	} else {
+		var err error
+		metadata.gpuUUID, err = sc.getActiveDevice(int(pid), int(tid), containerIDFunc)
+		if err != nil {
+			log.Warnf("error getting GPU UUID for process %d: %s", pid, err)
+			sc.telemetry.missingDevices.Inc()
+			return nil
+		}
+	}
+
+	metadata.smVersion = sc.sysCtx.deviceSmVersions[metadata.gpuUUID]
+
+	return newStreamHandler(metadata, sc.sysCtx)
+}
+
+// memoizedContainerID returns a function that memoizes the container ID for a given CUDA stream.
+// It's memoized because it's an expensive operation that we don't always need to do.
+func (sc *streamCollection) memoizedContainerID(header *gpuebpf.CudaEventHeader) func() string {
+	return funcs.MemoizeNoErrorUnsafe(func() string {
+		cgroup := unix.ByteSliceToString(header.Cgroup[:])
+		containerID, err := cgroups.ContainerFilter("", cgroup)
+		if err != nil {
+			log.Warnf("error getting container ID for cgroup %s: %s", cgroup, err)
+
+			sc.telemetry.missingContainers.Inc("error")
+			return ""
+		}
+
+		if containerID == "" {
+			sc.telemetry.missingContainers.Inc("missing")
+		}
+
+		return containerID
+	})
+}
+
+// getActiveDevice returns the GPU UUID for the current active GPU device for a given process and thread.
+func (sc *streamCollection) getActiveDevice(pid int, tid int, containerIDFunc func() string) (string, error) {
+	device, err := sc.sysCtx.getCurrentActiveGpuDevice(pid, tid, containerIDFunc)
+	if err != nil {
+		return "", err
+	}
+
+	uuid, ret := device.GetUUID()
+	if ret != nvml.SUCCESS {
+		return "", fmt.Errorf("error getting GPU UUID for process %d: %v", pid, nvml.ErrorString(ret))
+	}
+
+	return uuid, nil
+}
+
+func (sc *streamCollection) allStreams() iter.Seq[*StreamHandler] {
+	return func(yield func(*StreamHandler) bool) {
+		for _, stream := range sc.nonglobalStreams {
+			if !yield(stream) {
+				return
+			}
+		}
+
+		for _, stream := range sc.globalStreams {
+			if !yield(stream) {
+				return
+			}
+		}
+	}
+}
+
+func (sc *streamCollection) streamCount() int {
+	return len(sc.nonglobalStreams) + len(sc.globalStreams)
+}
+
+func (sc *streamCollection) markProcessStreamsAsEnded(pid uint32) {
path: pkg/gpu/stream_collection.go
user: login: gjulianm
body: In this case I prefer this name as it's consistent with the `markEnd` function and it signals that there's no actual action being taken here, other than setting some flags
created_at: 2025-03-19T11:24:29Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,240 @@
+// 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 2024-present Datadog, Inc.
+
+//go:build linux_bpf
+
+package gpu
+
+import (
+	"fmt"
+	"iter"
+
+	"github.com/NVIDIA/go-nvml/pkg/nvml"
+	"golang.org/x/sys/unix"
+
+	"github.com/DataDog/datadog-agent/comp/core/telemetry"
+	gpuebpf "github.com/DataDog/datadog-agent/pkg/gpu/ebpf"
+	"github.com/DataDog/datadog-agent/pkg/util/cgroups"
+	"github.com/DataDog/datadog-agent/pkg/util/funcs"
+	"github.com/DataDog/datadog-agent/pkg/util/log"
+)
+
+// nonglobalStreamKey is a unique identifier for a CUDA stream that is not global.
+// The streams are created with a specific GPU UUID, which means that pid+stream uniquely
+// identify the GPU.
+type nonglobalStreamKey struct {
+	pid    uint32
+	stream uint64
+}
+
+// globalStreamKey is a unique identifier for a CUDA stream that is global.
+// Global streams depend on which GPU is active when they are used, so the GPU UUID
+// needs to be part of the key
+type globalStreamKey struct {
+	pid     uint32
+	gpuUUID string
+}
+
+type streamCollection struct {
+	nonglobalStreams map[nonglobalStreamKey]*StreamHandler
+	globalStreams    map[globalStreamKey]*StreamHandler
+	sysCtx           *systemContext
+	telemetry        *streamCollectionTelemetry
+}
+
+type streamCollectionTelemetry struct {
+	missingContainers  telemetry.Counter
+	missingDevices     telemetry.Counter
+	finalizedProcesses telemetry.Counter
+	activeHandlers     telemetry.Gauge
+	removedHandlers    telemetry.Counter
+}
+
+func newStreamCollection(sysCtx *systemContext, telemetry telemetry.Component) *streamCollection {
+	return &streamCollection{
+		nonglobalStreams: make(map[nonglobalStreamKey]*StreamHandler),
+		globalStreams:    make(map[globalStreamKey]*StreamHandler),
+		sysCtx:           sysCtx,
+		telemetry:        newStreamCollectionTelemetry(telemetry),
+	}
+}
+
+func newStreamCollectionTelemetry(tm telemetry.Component) *streamCollectionTelemetry {
+	subsystem := gpuTelemetryModule + "__streams"
+
+	return &streamCollectionTelemetry{
+		missingContainers:  tm.NewCounter(subsystem, "missing_containers", []string{"reason"}, "Number of missing containers"),
+		missingDevices:     tm.NewCounter(subsystem, "missing_devices", nil, "Number of failures to get GPU devices for a stream"),
+		finalizedProcesses: tm.NewCounter(subsystem, "finalized_processes", nil, "Number of processes that have ended"),
+		activeHandlers:     tm.NewGauge(subsystem, "active_handlers", nil, "Number of active stream handlers"),
+		removedHandlers:    tm.NewCounter(subsystem, "removed_handlers", nil, "Number of removed stream handlers"),
+	}
+}
+
+// getStream returns a StreamHandler for a given CUDA stream.
+func (sc *streamCollection) getStream(header *gpuebpf.CudaEventHeader) (*StreamHandler, error) {
+	if header.Stream_id == 0 {
+		return sc.getGlobalStream(header)
+	}
+
+	return sc.getNonglobalStream(header)
+}
+
+func (sc *streamCollection) getGlobalStream(header *gpuebpf.CudaEventHeader) (*StreamHandler, error) {
+	pid, tid := getPidTidFromHeader(header)
+	memoizedContainerID := sc.memoizedContainerID(header)
+
+	// Global streams depend on which GPU is active when they are used, so we need to get the current active GPU device
+	// The expensive step here is the container ID parsing, but we don't always need to do it so we pass a function
+	// that can be called to retrieve it only when needed
+	gpuUUID, err := sc.getActiveDevice(int(pid), int(tid), memoizedContainerID)
+	if err != nil {
+		sc.telemetry.missingDevices.Inc()
+		return nil, err
+	}
+
+	key := globalStreamKey{
+		pid:     pid,
+		gpuUUID: gpuUUID,
+	}
+
+	stream, ok := sc.globalStreams[key]
+	if !ok {
+		stream = sc.createStream(header, &gpuUUID, memoizedContainerID)
+		sc.globalStreams[key] = stream
+		sc.telemetry.activeHandlers.Set(float64(sc.streamCount()))
+	}
+
+	return stream, nil
+}
+
+func (sc *streamCollection) getNonglobalStream(header *gpuebpf.CudaEventHeader) (*StreamHandler, error) {
+	pid, _ := getPidTidFromHeader(header)
+
+	key := nonglobalStreamKey{
+		pid:    pid,
+		stream: header.Stream_id,
+	}
+
+	stream, ok := sc.nonglobalStreams[key]
+	if !ok {
+		stream = sc.createStream(header, nil, sc.memoizedContainerID(header))
+		sc.nonglobalStreams[key] = stream
+		sc.telemetry.activeHandlers.Set(float64(sc.streamCount()))
+	}
+
+	return stream, nil
+}
+
+// createStream creates a new StreamHandler for a given CUDA stream.
+// If the GPU UUID is not provided (it's nil), it will be retrieved from the system context.
+func (sc *streamCollection) createStream(header *gpuebpf.CudaEventHeader, gpuUUID *string, containerIDFunc func() string) *StreamHandler {
+	pid, tid := getPidTidFromHeader(header)
+	metadata := streamMetadata{
+		pid:         pid,
+		streamID:    header.Stream_id,
+		containerID: containerIDFunc(),
+	}
+
+	if gpuUUID != nil {
+		metadata.gpuUUID = *gpuUUID
+	} else {
+		var err error
+		metadata.gpuUUID, err = sc.getActiveDevice(int(pid), int(tid), containerIDFunc)
+		if err != nil {
+			log.Warnf("error getting GPU UUID for process %d: %s", pid, err)
+			sc.telemetry.missingDevices.Inc()
+			return nil
+		}
+	}
+
+	metadata.smVersion = sc.sysCtx.deviceSmVersions[metadata.gpuUUID]
+
+	return newStreamHandler(metadata, sc.sysCtx)
+}
+
+// memoizedContainerID returns a function that memoizes the container ID for a given CUDA stream.
+// It's memoized because it's an expensive operation that we don't always need to do.
+func (sc *streamCollection) memoizedContainerID(header *gpuebpf.CudaEventHeader) func() string {
+	return funcs.MemoizeNoErrorUnsafe(func() string {
+		cgroup := unix.ByteSliceToString(header.Cgroup[:])
+		containerID, err := cgroups.ContainerFilter("", cgroup)
+		if err != nil {
+			log.Warnf("error getting container ID for cgroup %s: %s", cgroup, err)
+
+			sc.telemetry.missingContainers.Inc("error")
+			return ""
+		}
+
+		if containerID == "" {
+			sc.telemetry.missingContainers.Inc("missing")
+		}
+
+		return containerID
+	})
+}
+
+// getActiveDevice returns the GPU UUID for the current active GPU device for a given process and thread.
+func (sc *streamCollection) getActiveDevice(pid int, tid int, containerIDFunc func() string) (string, error) {
+	device, err := sc.sysCtx.getCurrentActiveGpuDevice(pid, tid, containerIDFunc)
+	if err != nil {
+		return "", err
+	}
+
+	uuid, ret := device.GetUUID()
+	if ret != nvml.SUCCESS {
+		return "", fmt.Errorf("error getting GPU UUID for process %d: %v", pid, nvml.ErrorString(ret))
+	}
+
+	return uuid, nil
+}
+
+func (sc *streamCollection) allStreams() iter.Seq[*StreamHandler] {
+	return func(yield func(*StreamHandler) bool) {
+		for _, stream := range sc.nonglobalStreams {
+			if !yield(stream) {
+				return
+			}
+		}
+
+		for _, stream := range sc.globalStreams {
+			if !yield(stream) {
+				return
+			}
+		}
+	}
+}
+
+func (sc *streamCollection) streamCount() int {
+	return len(sc.nonglobalStreams) + len(sc.globalStreams)
+}
+
+func (sc *streamCollection) markProcessStreamsAsEnded(pid uint32) {
+	for handler := range sc.allStreams() {
+		if handler.metadata.pid == pid {
+			log.Debugf("Process %d ended, marking stream %d as ended", pid, handler.metadata.streamID)
+			_ = handler.markEnd()
+			sc.telemetry.finalizedProcesses.Inc()
+		}
+	}
+}
+
+func (sc *streamCollection) cleanupFinishedHandlers() {
path: pkg/gpu/stream_collection.go
user: login: gjulianm
body: Renamed
created_at: 2025-03-19T11:24:34Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,240 @@
+// 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 2024-present Datadog, Inc.
+
+//go:build linux_bpf
+
+package gpu
+
+import (
+	"fmt"
+	"iter"
+
+	"github.com/NVIDIA/go-nvml/pkg/nvml"
+	"golang.org/x/sys/unix"
+
+	"github.com/DataDog/datadog-agent/comp/core/telemetry"
+	gpuebpf "github.com/DataDog/datadog-agent/pkg/gpu/ebpf"
+	"github.com/DataDog/datadog-agent/pkg/util/cgroups"
+	"github.com/DataDog/datadog-agent/pkg/util/funcs"
+	"github.com/DataDog/datadog-agent/pkg/util/log"
+)
+
+// nonglobalStreamKey is a unique identifier for a CUDA stream that is not global.
+// The streams are created with a specific GPU UUID, which means that pid+stream uniquely
+// identify the GPU.
+type nonglobalStreamKey struct {
+	pid    uint32
+	stream uint64
+}
+
+// globalStreamKey is a unique identifier for a CUDA stream that is global.
+// Global streams depend on which GPU is active when they are used, so the GPU UUID
+// needs to be part of the key
+type globalStreamKey struct {
+	pid     uint32
+	gpuUUID string
+}
+
+type streamCollection struct {
+	nonglobalStreams map[nonglobalStreamKey]*StreamHandler
+	globalStreams    map[globalStreamKey]*StreamHandler
+	sysCtx           *systemContext
+	telemetry        *streamCollectionTelemetry
+}
+
+type streamCollectionTelemetry struct {
+	missingContainers  telemetry.Counter
+	missingDevices     telemetry.Counter
+	finalizedProcesses telemetry.Counter
+	activeHandlers     telemetry.Gauge
+	removedHandlers    telemetry.Counter
+}
+
+func newStreamCollection(sysCtx *systemContext, telemetry telemetry.Component) *streamCollection {
+	return &streamCollection{
+		nonglobalStreams: make(map[nonglobalStreamKey]*StreamHandler),
+		globalStreams:    make(map[globalStreamKey]*StreamHandler),
+		sysCtx:           sysCtx,
+		telemetry:        newStreamCollectionTelemetry(telemetry),
+	}
+}
+
+func newStreamCollectionTelemetry(tm telemetry.Component) *streamCollectionTelemetry {
+	subsystem := gpuTelemetryModule + "__streams"
+
+	return &streamCollectionTelemetry{
+		missingContainers:  tm.NewCounter(subsystem, "missing_containers", []string{"reason"}, "Number of missing containers"),
+		missingDevices:     tm.NewCounter(subsystem, "missing_devices", nil, "Number of failures to get GPU devices for a stream"),
+		finalizedProcesses: tm.NewCounter(subsystem, "finalized_processes", nil, "Number of processes that have ended"),
+		activeHandlers:     tm.NewGauge(subsystem, "active_handlers", nil, "Number of active stream handlers"),
+		removedHandlers:    tm.NewCounter(subsystem, "removed_handlers", nil, "Number of removed stream handlers"),
path: pkg/gpu/stream_collection.go
user: login: gjulianm
body: We're not splitting the maps per-device, so we cannot get active_handlers per device without re-iterating the maps and building the counts manually. I added it to removedHandlers only as in that case it is easy.
created_at: 2025-03-19T11:28:44Z
author_association: CONTRIBUTOR

PR comments:
user: login: agent-platform-auto-pr[bot]
created_at: 2025-03-11T12:37:13Z
body: ## Uncompressed package size comparison

Comparison with [ancestor](https://github.com/DataDog/datadog-agent/commit/c4e0a62ea3426ac25a83e64bd5c88b5f98ea13e4) `c4e0a62ea3426ac25a83e64bd5c88b5f98ea13e4`


<details open>
  <summary> Diff per package </summary>

|package|diff|status|size|ancestor|threshold|
|--|--|--|--|--|--|
|datadog-agent-arm64-deb|0.01MB|⚠️|798.57MB|798.56MB|0.50MB|
|datadog-agent-aarch64-rpm|0.01MB|⚠️|808.35MB|808.33MB|0.50MB|
|datadog-agent-amd64-deb|0.01MB|⚠️|807.58MB|807.57MB|0.50MB|
|datadog-agent-x86_64-rpm|0.01MB|⚠️|817.37MB|817.37MB|0.50MB|
|datadog-agent-x86_64-suse|0.01MB|⚠️|817.37MB|817.37MB|0.50MB|
|datadog-dogstatsd-amd64-deb|0.00MB|✅|39.34MB|39.34MB|0.50MB|
|datadog-dogstatsd-x86_64-rpm|0.00MB|✅|39.42MB|39.42MB|0.50MB|
|datadog-dogstatsd-x86_64-suse|0.00MB|✅|39.42MB|39.42MB|0.50MB|
|datadog-dogstatsd-arm64-deb|0.00MB|✅|37.87MB|37.87MB|0.50MB|
|datadog-heroku-agent-amd64-deb|0.00MB|✅|444.37MB|444.37MB|0.50MB|
|datadog-iot-agent-amd64-deb|0.00MB|✅|62.23MB|62.23MB|0.50MB|
|datadog-iot-agent-x86_64-rpm|0.00MB|✅|62.30MB|62.30MB|0.50MB|
|datadog-iot-agent-x86_64-suse|0.00MB|✅|62.30MB|62.30MB|0.50MB|
|datadog-iot-agent-arm64-deb|0.00MB|✅|59.46MB|59.46MB|0.50MB|
|datadog-iot-agent-aarch64-rpm|0.00MB|✅|59.53MB|59.53MB|0.50MB|

</details>

## Decision
⚠️ Warning



author_association: CONTRIBUTOR
user: login: cit-pr-commenter[bot]
created_at: 2025-03-11T12:56:41Z
body: ## Regression Detector

# Regression Detector Results

[Metrics dashboard](https://app.datadoghq.com/dashboard/ykh-ua8-vcu/SMP-Regression-Detector-Metrics?fromUser=true&refresh_mode=paused&tpl_var_run-id%5B0%5D=8629a13e-006f-4735-a458-d208faee55ca&view=spans&from_ts=1742465026000&to_ts=1742465626000&live=false)  
[Target profiles](https://app.datadoghq.com/profiling/explorer?query=env%3Asingle-machine-performance%20service%3Adatadog-agent%20job_id%3A8629a13e-006f-4735-a458-d208faee55ca&agg_m=count&agg_m_source=base&agg_t=count&fromUser=false&viz=stream&start=1742457826000&end=1742469226000&paused=true)  
Run ID: 8629a13e-006f-4735-a458-d208faee55ca

Baseline: c4e0a62ea3426ac25a83e64bd5c88b5f98ea13e4
Comparison: e6ab8a57a5b92c5cc687d86c90514414da694e04
[Diff](c4e0a62ea3426ac25a83e64bd5c88b5f98ea13e4...e6ab8a57a5b92c5cc687d86c90514414da694e04)


## Optimization Goals: ✅ No significant changes detected


<details>
<summary><h2>
Fine details of change detection per experiment
</h2></summary>

| perf | experiment                                   | goal               | Δ mean % | Δ mean % CI    | trials | links                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                          |
|------|----------------------------------------------|--------------------|----------|----------------|--------|--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
| ➖   | uds_dogstatsd_to_api_cpu                     | % cpu utilization  | +1.10    | [+0.28, +1.91] | 1      | [Logs](https://app.datadoghq.com/logs?query=experiment%3Auds_dogstatsd_to_api_cpu%20run_id%3A8629a13e-006f-4735-a458-d208faee55ca&agg_m=count&agg_m_source=base&agg_q=%40span.url&agg_q_source=base&agg_t=count&fromUser=true&index=single-machine-performance-target-logs&messageDisplay=inline&refresh_mode=paused&storage=hot&stream_sort=time%2Cdesc&top_n=100&top_o=top&viz=stream&x_missing=true&from_ts=1742457826000&to_ts=1742469226000&live=false)                                                                                                                                                                                                                                                                                                                                                                   |
| ➖   | quality_gate_idle_all_features               | memory utilization | +0.99    | [+0.89, +1.09] | 1      | [Logs](https://app.datadoghq.com/logs?query=experiment%3Aquality_gate_idle_all_features%20run_id%3A8629a13e-006f-4735-a458-d208faee55ca&agg_m=count&agg_m_source=base&agg_q=%40span.url&agg_q_source=base&agg_t=count&fromUser=true&index=single-machine-performance-target-logs&messageDisplay=inline&refresh_mode=paused&storage=hot&stream_sort=time%2Cdesc&top_n=100&top_o=top&viz=stream&x_missing=true&from_ts=1742457826000&to_ts=1742469226000&live=false) [bounds checks dashboard](https://app.datadoghq.com/dashboard/vz3-jd5-bdi?fromUser=true&refresh_mode=paused&tpl_var_experiment%5B0%5D=quality_gate_idle_all_features&tpl_var_job_id%5B0%5D=8629a13e-006f-4735-a458-d208faee55ca&tpl_var_run-id%5B0%5D=8629a13e-006f-4735-a458-d208faee55ca&view=spans&from_ts=1742465026000&to_ts=1742465626000&live=false) |
| ➖   | quality_gate_idle                            | memory utilization | +0.83    | [+0.77, +0.88] | 1      | [Logs](https://app.datadoghq.com/logs?query=experiment%3Aquality_gate_idle%20run_id%3A8629a13e-006f-4735-a458-d208faee55ca&agg_m=count&agg_m_source=base&agg_q=%40span.url&agg_q_source=base&agg_t=count&fromUser=true&index=single-machine-performance-target-logs&messageDisplay=inline&refresh_mode=paused&storage=hot&stream_sort=time%2Cdesc&top_n=100&top_o=top&viz=stream&x_missing=true&from_ts=1742457826000&to_ts=1742469226000&live=false) [bounds checks dashboard](https://app.datadoghq.com/dashboard/vz3-jd5-bdi?fromUser=true&refresh_mode=paused&tpl_var_experiment%5B0%5D=quality_gate_idle&tpl_var_job_id%5B0%5D=8629a13e-006f-4735-a458-d208faee55ca&tpl_var_run-id%5B0%5D=8629a13e-006f-4735-a458-d208faee55ca&view=spans&from_ts=1742465026000&to_ts=1742465626000&live=false)                           |
| ➖   | uds_dogstatsd_20mb_12k_contexts_20_senders   | memory utilization | +0.16    | [+0.09, +0.24] | 1      | [Logs](https://app.datadoghq.com/logs?query=experiment%3Auds_dogstatsd_20mb_12k_contexts_20_senders%20run_id%3A8629a13e-006f-4735-a458-d208faee55ca&agg_m=count&agg_m_source=base&agg_q=%40span.url&agg_q_source=base&agg_t=count&fromUser=true&index=single-machine-performance-target-logs&messageDisplay=inline&refresh_mode=paused&storage=hot&stream_sort=time%2Cdesc&top_n=100&top_o=top&viz=stream&x_missing=true&from_ts=1742457826000&to_ts=1742469226000&live=false)                                                                                                                                                                                                                                                                                                                                                 |
| ➖   | file_to_blackhole_1000ms_latency_linear_load | egress throughput  | +0.12    | [-0.35, +0.58] | 1      | [Logs](https://app.datadoghq.com/logs?query=experiment%3Afile_to_blackhole_1000ms_latency_linear_load%20run_id%3A8629a13e-006f-4735-a458-d208faee55ca&agg_m=count&agg_m_source=base&agg_q=%40span.url&agg_q_source=base&agg_t=count&fromUser=true&index=single-machine-performance-target-logs&messageDisplay=inline&refresh_mode=paused&storage=hot&stream_sort=time%2Cdesc&top_n=100&top_o=top&viz=stream&x_missing=true&from_ts=1742457826000&to_ts=1742469226000&live=false)                                                                                                                                                                                                                                                                                                                                               |
| ➖   | file_to_blackhole_100ms_latency              | egress throughput  | +0.02    | [-0.66, +0.69] | 1      | [Logs](https://app.datadoghq.com/logs?query=experiment%3Afile_to_blackhole_100ms_latency%20run_id%3A8629a13e-006f-4735-a458-d208faee55ca&agg_m=count&agg_m_source=base&agg_q=%40span.url&agg_q_source=base&agg_t=count&fromUser=true&index=single-machine-performance-target-logs&messageDisplay=inline&refresh_mode=paused&storage=hot&stream_sort=time%2Cdesc&top_n=100&top_o=top&viz=stream&x_missing=true&from_ts=1742457826000&to_ts=1742469226000&live=false)                                                                                                                                                                                                                                                                                                                                                            |
| ➖   | file_to_blackhole_0ms_latency                | egress throughput  | +0.01    | [-0.78, +0.81] | 1      | [Logs](https://app.datadoghq.com/logs?query=experiment%3Afile_to_blackhole_0ms_latency%20run_id%3A8629a13e-006f-4735-a458-d208faee55ca&agg_m=count&agg_m_source=base&agg_q=%40span.url&agg_q_source=base&agg_t=count&fromUser=true&index=single-machine-performance-target-logs&messageDisplay=inline&refresh_mode=paused&storage=hot&stream_sort=time%2Cdesc&top_n=100&top_o=top&viz=stream&x_missing=true&from_ts=1742457826000&to_ts=1742469226000&live=false)                                                                                                                                                                                                                                                                                                                                                              |
| ➖   | file_to_blackhole_500ms_latency              | egress throughput  | +0.01    | [-0.77, +0.78] | 1      | [Logs](https://app.datadoghq.com/logs?query=experiment%3Afile_to_blackhole_500ms_latency%20run_id%3A8629a13e-006f-4735-a458-d208faee55ca&agg_m=count&agg_m_source=base&agg_q=%40span.url&agg_q_source=base&agg_t=count&fromUser=true&index=single-machine-performance-target-logs&messageDisplay=inline&refresh_mode=paused&storage=hot&stream_sort=time%2Cdesc&top_n=100&top_o=top&viz=stream&x_missing=true&from_ts=1742457826000&to_ts=1742469226000&live=false)                                                                                                                                                                                                                                                                                                                                                            |
| ➖   | file_to_blackhole_300ms_latency              | egress throughput  | +0.00    | [-0.63, +0.64] | 1      | [Logs](https://app.datadoghq.com/logs?query=experiment%3Afile_to_blackhole_300ms_latency%20run_id%3A8629a13e-006f-4735-a458-d208faee55ca&agg_m=count&agg_m_source=base&agg_q=%40span.url&agg_q_source=base&agg_t=count&fromUser=true&index=single-machine-performance-target-logs&messageDisplay=inline&refresh_mode=paused&storage=hot&stream_sort=time%2Cdesc&top_n=100&top_o=top&viz=stream&x_missing=true&from_ts=1742457826000&to_ts=1742469226000&live=false)                                                                                                                                                                                                                                                                                                                                                            |
| ➖   | uds_dogstatsd_to_api                         | ingress throughput | +0.00    | [-0.29, +0.30] | 1      | [Logs](https://app.datadoghq.com/logs?query=experiment%3Auds_dogstatsd_to_api%20run_id%3A8629a13e-006f-4735-a458-d208faee55ca&agg_m=count&agg_m_source=base&agg_q=%40span.url&agg_q_source=base&agg_t=count&fromUser=true&index=single-machine-performance-target-logs&messageDisplay=inline&refresh_mode=paused&storage=hot&stream_sort=time%2Cdesc&top_n=100&top_o=top&viz=stream&x_missing=true&from_ts=1742457826000&to_ts=1742469226000&live=false)                                                                                                                                                                                                                                                                                                                                                                       |
| ➖   | tcp_dd_logs_filter_exclude                   | ingress throughput | +0.00    | [-0.02, +0.02] | 1      | [Logs](https://app.datadoghq.com/logs?query=experiment%3Atcp_dd_logs_filter_exclude%20run_id%3A8629a13e-006f-4735-a458-d208faee55ca&agg_m=count&agg_m_source=base&agg_q=%40span.url&agg_q_source=base&agg_t=count&fromUser=true&index=single-machine-performance-target-logs&messageDisplay=inline&refresh_mode=paused&storage=hot&stream_sort=time%2Cdesc&top_n=100&top_o=top&viz=stream&x_missing=true&from_ts=1742457826000&to_ts=1742469226000&live=false)                                                                                                                                                                                                                                                                                                                                                                 |
| ➖   | file_to_blackhole_0ms_latency_http1          | egress throughput  | -0.00    | [-0.84, +0.84] | 1      | [Logs](https://app.datadoghq.com/logs?query=experiment%3Afile_to_blackhole_0ms_latency_http1%20run_id%3A8629a13e-006f-4735-a458-d208faee55ca&agg_m=count&agg_m_source=base&agg_q=%40span.url&agg_q_source=base&agg_t=count&fromUser=true&index=single-machine-performance-target-logs&messageDisplay=inline&refresh_mode=paused&storage=hot&stream_sort=time%2Cdesc&top_n=100&top_o=top&viz=stream&x_missing=true&from_ts=1742457826000&to_ts=1742469226000&live=false)                                                                                                                                                                                                                                                                                                                                                        |
| ➖   | file_to_blackhole_0ms_latency_http2          | egress throughput  | -0.01    | [-0.84, +0.81] | 1      | [Logs](https://app.datadoghq.com/logs?query=experiment%3Afile_to_blackhole_0ms_latency_http2%20run_id%3A8629a13e-006f-4735-a458-d208faee55ca&agg_m=count&agg_m_source=base&agg_q=%40span.url&agg_q_source=base&agg_t=count&fromUser=true&index=single-machine-performance-target-logs&messageDisplay=inline&refresh_mode=paused&storage=hot&stream_sort=time%2Cdesc&top_n=100&top_o=top&viz=stream&x_missing=true&from_ts=1742457826000&to_ts=1742469226000&live=false)                                                                                                                                                                                                                                                                                                                                                        |
| ➖   | file_to_blackhole_1000ms_latency             | egress throughput  | -0.11    | [-0.88, +0.66] | 1      | [Logs](https://app.datadoghq.com/logs?query=experiment%3Afile_to_blackhole_1000ms_latency%20run_id%3A8629a13e-006f-4735-a458-d208faee55ca&agg_m=count&agg_m_source=base&agg_q=%40span.url&agg_q_source=base&agg_t=count&fromUser=true&index=single-machine-performance-target-logs&messageDisplay=inline&refresh_mode=paused&storage=hot&stream_sort=time%2Cdesc&top_n=100&top_o=top&viz=stream&x_missing=true&from_ts=1742457826000&to_ts=1742469226000&live=false)                                                                                                                                                                                                                                                                                                                                                           |
| ➖   | file_tree                                    | memory utilization | -0.51    | [-0.64, -0.37] | 1      | [Logs](https://app.datadoghq.com/logs?query=experiment%3Afile_tree%20run_id%3A8629a13e-006f-4735-a458-d208faee55ca&agg_m=count&agg_m_source=base&agg_q=%40span.url&agg_q_source=base&agg_t=count&fromUser=true&index=single-machine-performance-target-logs&messageDisplay=inline&refresh_mode=paused&storage=hot&stream_sort=time%2Cdesc&top_n=100&top_o=top&viz=stream&x_missing=true&from_ts=1742457826000&to_ts=1742469226000&live=false)                                                                                                                                                                                                                                                                                                                                                                                  |
| ➖   | tcp_syslog_to_blackhole                      | ingress throughput | -0.77    | [-0.83, -0.71] | 1      | [Logs](https://app.datadoghq.com/logs?query=experiment%3Atcp_syslog_to_blackhole%20run_id%3A8629a13e-006f-4735-a458-d208faee55ca&agg_m=count&agg_m_source=base&agg_q=%40span.url&agg_q_source=base&agg_t=count&fromUser=true&index=single-machine-performance-target-logs&messageDisplay=inline&refresh_mode=paused&storage=hot&stream_sort=time%2Cdesc&top_n=100&top_o=top&viz=stream&x_missing=true&from_ts=1742457826000&to_ts=1742469226000&live=false)                                                                                                                                                                                                                                                                                                                                                                    |
| ➖   | quality_gate_logs                            | % cpu utilization  | -1.47    | [-4.24, +1.31] | 1      | [Logs](https://app.datadoghq.com/logs?query=experiment%3Aquality_gate_logs%20run_id%3A8629a13e-006f-4735-a458-d208faee55ca&agg_m=count&agg_m_source=base&agg_q=%40span.url&agg_q_source=base&agg_t=count&fromUser=true&index=single-machine-performance-target-logs&messageDisplay=inline&refresh_mode=paused&storage=hot&stream_sort=time%2Cdesc&top_n=100&top_o=top&viz=stream&x_missing=true&from_ts=1742457826000&to_ts=1742469226000&live=false)                                                                                                                                                                                                                                                                                                                                                                          |
</details>


<details>
<summary><h2>
Bounds Checks: ✅ Passed
</h2></summary>

| perf | experiment                                   | bounds_check_name  | replicates_passed | links                                                                                                                                                                                                                                                                                                                                                       |
|------|----------------------------------------------|--------------------|-------------------|-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
| ✅   | file_to_blackhole_0ms_latency                | lost_bytes         | 10/10             |                                                                                                                                                                                                                                                                                                                                                             |
| ✅   | file_to_blackhole_0ms_latency                | memory_usage       | 10/10             |                                                                                                                                                                                                                                                                                                                                                             |
| ✅   | file_to_blackhole_0ms_latency_http1          | lost_bytes         | 10/10             |                                                                                                                                                                                                                                                                                                                                                             |
| ✅   | file_to_blackhole_0ms_latency_http1          | memory_usage       | 10/10             |                                                                                                                                                                                                                                                                                                                                                             |
| ✅   | file_to_blackhole_0ms_latency_http2          | lost_bytes         | 10/10             |                                                                                                                                                                                                                                                                                                                                                             |
| ✅   | file_to_blackhole_0ms_latency_http2          | memory_usage       | 10/10             |                                                                                                                                                                                                                                                                                                                                                             |
| ✅   | file_to_blackhole_1000ms_latency             | memory_usage       | 10/10             |                                                                                                                                                                                                                                                                                                                                                             |
| ✅   | file_to_blackhole_1000ms_latency_linear_load | memory_usage       | 10/10             |                                                                                                                                                                                                                                                                                                                                                             |
| ✅   | file_to_blackhole_100ms_latency              | lost_bytes         | 10/10             |                                                                                                                                                                                                                                                                                                                                                             |
| ✅   | file_to_blackhole_100ms_latency              | memory_usage       | 10/10             |                                                                                                                                                                                                                                                                                                                                                             |
| ✅   | file_to_blackhole_300ms_latency              | lost_bytes         | 10/10             |                                                                                                                                                                                                                                                                                                                                                             |
| ✅   | file_to_blackhole_300ms_latency              | memory_usage       | 10/10             |                                                                                                                                                                                                                                                                                                                                                             |
| ✅   | file_to_blackhole_500ms_latency              | lost_bytes         | 10/10             |                                                                                                                                                                                                                                                                                                                                                             |
| ✅   | file_to_blackhole_500ms_latency              | memory_usage       | 10/10             |                                                                                                                                                                                                                                                                                                                                                             |
| ✅   | quality_gate_idle                            | intake_connections | 10/10             | [bounds checks dashboard](https://app.datadoghq.com/dashboard/vz3-jd5-bdi?fromUser=true&refresh_mode=paused&tpl_var_experiment%5B0%5D=quality_gate_idle&tpl_var_job_id%5B0%5D=8629a13e-006f-4735-a458-d208faee55ca&tpl_var_run-id%5B0%5D=8629a13e-006f-4735-a458-d208faee55ca&view=spans&from_ts=1742465026000&to_ts=1742465626000&live=false)              |
| ✅   | quality_gate_idle                            | memory_usage       | 10/10             | [bounds checks dashboard](https://app.datadoghq.com/dashboard/vz3-jd5-bdi?fromUser=true&refresh_mode=paused&tpl_var_experiment%5B0%5D=quality_gate_idle&tpl_var_job_id%5B0%5D=8629a13e-006f-4735-a458-d208faee55ca&tpl_var_run-id%5B0%5D=8629a13e-006f-4735-a458-d208faee55ca&view=spans&from_ts=1742465026000&to_ts=1742465626000&live=false)              |
| ✅   | quality_gate_idle_all_features               | intake_connections | 10/10             | [bounds checks dashboard](https://app.datadoghq.com/dashboard/vz3-jd5-bdi?fromUser=true&refresh_mode=paused&tpl_var_experiment%5B0%5D=quality_gate_idle_all_features&tpl_var_job_id%5B0%5D=8629a13e-006f-4735-a458-d208faee55ca&tpl_var_run-id%5B0%5D=8629a13e-006f-4735-a458-d208faee55ca&view=spans&from_ts=1742465026000&to_ts=1742465626000&live=false) |
| ✅   | quality_gate_idle_all_features               | memory_usage       | 10/10             | [bounds checks dashboard](https://app.datadoghq.com/dashboard/vz3-jd5-bdi?fromUser=true&refresh_mode=paused&tpl_var_experiment%5B0%5D=quality_gate_idle_all_features&tpl_var_job_id%5B0%5D=8629a13e-006f-4735-a458-d208faee55ca&tpl_var_run-id%5B0%5D=8629a13e-006f-4735-a458-d208faee55ca&view=spans&from_ts=1742465026000&to_ts=1742465626000&live=false) |
| ✅   | quality_gate_logs                            | intake_connections | 10/10             |                                                                                                                                                                                                                                                                                                                                                             |
| ✅   | quality_gate_logs                            | lost_bytes         | 10/10             |                                                                                                                                                                                                                                                                                                                                                             |
| ✅   | quality_gate_logs                            | memory_usage       | 10/10             |                                                                                                                                                                                                                                                                                                                                                             |
</details>


<details>
<summary><h2>
Explanation
</h2></summary>

**Confidence level:** 90.00%
**Effect size tolerance:** |Δ mean %| ≥ 5.00%

Performance changes are noted in the **perf** column of each table:

* ✅ = significantly better comparison variant performance
* ❌ = significantly worse comparison variant performance
* ➖ = no significant change in performance

A regression test is an A/B test of target performance in a repeatable rig, where "performance" is measured as "comparison variant minus baseline variant" for an optimization goal (e.g., ingress throughput). Due to intrinsic variability in measuring that goal, we can only estimate its mean value for each experiment; we report uncertainty in that value as a 90.00% confidence interval denoted "Δ mean % CI".

For each experiment, we decide whether a change in performance is a "regression" -- a change worth investigating further -- if all of the following criteria are true:

1. Its estimated |Δ mean %| ≥ 5.00%, indicating the change is big enough to merit a closer look.

2. Its 90.00% confidence interval "Δ mean % CI" does not contain zero, indicating that *if our statistical model is accurate*, there is at least a 90.00% chance there is a difference in performance between baseline and comparison variants.

3. Its configuration does not mark it "erratic".
</details>

## CI Pass/Fail Decision

✅ **Passed.** All Quality Gates passed.

- **quality_gate_idle**, bounds check **memory_usage**: 10/10 replicas passed. Gate passed.
- **quality_gate_idle**, bounds check **intake_connections**: 10/10 replicas passed. Gate passed.
- **quality_gate_idle_all_features**, bounds check **intake_connections**: 10/10 replicas passed. Gate passed.
- **quality_gate_idle_all_features**, bounds check **memory_usage**: 10/10 replicas passed. Gate passed.
- **quality_gate_logs**, bounds check **lost_bytes**: 10/10 replicas passed. Gate passed.
- **quality_gate_logs**, bounds check **memory_usage**: 10/10 replicas passed. Gate passed.
- **quality_gate_logs**, bounds check **intake_connections**: 10/10 replicas passed. Gate passed.
author_association: NONE
user: login: agent-platform-auto-pr[bot]
created_at: 2025-03-11T16:54:52Z
body: ## Test changes on VM

Use this command from [test-infra-definitions](https://github.com/DataDog/test-infra-definitions) to manually test this PR changes on a VM:
```sh
dda inv aws.create-vm --pipeline-id=59424257 --os-family=ubuntu
```
Note: This applies to commit **e6ab8a57**
author_association: CONTRIBUTOR
user: login: agent-platform-auto-pr[bot]
created_at: 2025-03-11T16:59:57Z
body: ## Static quality checks ✅

Please find below the results from static quality gates


<details>
<summary>Successful checks</summary>

### Info

|Result|Quality gate|On disk size|On disk size limit|On wire size|On wire size limit|
|----|----|----|----|----|----|
|✅|static_quality_gate_agent_deb_amd64|781.72 MiB|801.8 MiB|190.72 MiB|202.62 MiB|
|✅|static_quality_gate_agent_deb_arm64|773.05 MiB|793.14 MiB|173.08 MiB|184.51 MiB|
|✅|static_quality_gate_agent_rpm_amd64|781.6 MiB|801.79 MiB|193.32 MiB|205.03 MiB|
|✅|static_quality_gate_agent_rpm_arm64|773.08 MiB|793.09 MiB|174.55 MiB|186.44 MiB|
|✅|static_quality_gate_agent_suse_amd64|781.63 MiB|801.81 MiB|193.32 MiB|205.03 MiB|
|✅|static_quality_gate_agent_suse_arm64|773.09 MiB|793.14 MiB|174.55 MiB|186.44 MiB|
|✅|static_quality_gate_dogstatsd_deb_amd64|37.59 MiB|47.67 MiB|9.75 MiB|19.78 MiB|
|✅|static_quality_gate_dogstatsd_deb_arm64|36.19 MiB|46.27 MiB|8.45 MiB|18.49 MiB|
|✅|static_quality_gate_dogstatsd_rpm_amd64|37.59 MiB|47.67 MiB|9.76 MiB|19.79 MiB|
|✅|static_quality_gate_dogstatsd_suse_amd64|37.59 MiB|47.67 MiB|9.76 MiB|19.79 MiB|
|✅|static_quality_gate_iot_agent_deb_amd64|59.43 MiB|69.0 MiB|14.92 MiB|24.8 MiB|
|✅|static_quality_gate_iot_agent_deb_arm64|56.78 MiB|66.4 MiB|12.87 MiB|22.8 MiB|
|✅|static_quality_gate_iot_agent_rpm_amd64|59.43 MiB|69.0 MiB|14.95 MiB|24.8 MiB|
|✅|static_quality_gate_iot_agent_rpm_arm64|56.78 MiB|66.4 MiB|12.88 MiB|22.8 MiB|
|✅|static_quality_gate_iot_agent_suse_amd64|59.42 MiB|69.0 MiB|14.95 MiB|24.8 MiB|
|✅|static_quality_gate_docker_agent_amd64|866.32 MiB|886.12 MiB|291.44 MiB|304.21 MiB|
|✅|static_quality_gate_docker_agent_arm64|880.99 MiB|900.79 MiB|277.86 MiB|290.47 MiB|
|✅|static_quality_gate_docker_agent_jmx_amd64|1.04 GiB|1.06 GiB|366.53 MiB|379.33 MiB|
|✅|static_quality_gate_docker_agent_jmx_arm64|1.04 GiB|1.06 GiB|348.94 MiB|361.55 MiB|
|✅|static_quality_gate_docker_dogstatsd_amd64|45.74 MiB|55.78 MiB|17.25 MiB|27.28 MiB|
|✅|static_quality_gate_docker_dogstatsd_arm64|44.36 MiB|54.45 MiB|16.12 MiB|26.16 MiB|
|✅|static_quality_gate_docker_cluster_agent_amd64|264.38 MiB|274.78 MiB|106.08 MiB|116.28 MiB|
|✅|static_quality_gate_docker_cluster_agent_arm64|280.33 MiB|290.82 MiB|100.89 MiB|111.12 MiB|

</details>

author_association: CONTRIBUTOR
user: login: gjulianm
created_at: 2025-03-19T12:02:39Z
body: /merge
author_association: CONTRIBUTOR
user: login: dd-devflow[bot]
created_at: 2025-03-19T12:02:43Z
body: <!-- "FeedbackWorkflowId":"0cff9586-2f07-4405-be89-24461f49d3a5_36" -->
[View all feedbacks in Devflow UI.](https://devflow.static-app.us1.prod.dog/logs/0cff9586-2f07-4405-be89-24461f49d3a5_36)
`2025-03-19 12:02:42 UTC` :information_source: Start processing command `/merge`

---
`2025-03-19 12:02:48 UTC` :information_source: **MergeQueue:** waiting for PR to be ready

This merge request is not mergeable yet, because of pending checks/missing approvals. It will be added to the queue as soon as checks pass and/or get approvals.
**Note:** if you pushed new commits since the last approval, you may need additional approval.
You can remove it from the waiting list with `/remove` command.

---
`2025-03-19 16:03:04 UTC` :warning: **MergeQueue:** This merge request was unqueued

devflow unqueued this merge request: It did not become mergeable within the expected time

author_association: NONE
user: login: gjulianm
created_at: 2025-03-20T11:24:40Z
body: /merge
author_association: CONTRIBUTOR
user: login: dd-devflow[bot]
created_at: 2025-03-20T11:24:44Z
body: <!-- "FeedbackWorkflowId":"f414ce38-3f1b-48d5-ab97-5dbae884db88_36" -->
[View all feedbacks in Devflow UI.](https://devflow.static-app.us1.prod.dog/logs/f414ce38-3f1b-48d5-ab97-5dbae884db88_36)
`2025-03-20 11:24:43 UTC` :information_source: Start processing command `/merge`

---
`2025-03-20 11:24:49 UTC` :information_source: **MergeQueue:** pull request added to the queue

The expected merge time in `main` is approximately 52m (p90).

---
`2025-03-20 11:51:02 UTC` :information_source: **MergeQueue:** This merge request was merged

author_association: NONE