diff --git a/comp/core/workloadmeta/collectors/internal/nvml/nvml.go b/comp/core/workloadmeta/collectors/internal/nvml/nvml.go index 06925f50ba0817..4acbd3a4453afd 100644 --- a/comp/core/workloadmeta/collectors/internal/nvml/nvml.go +++ b/comp/core/workloadmeta/collectors/internal/nvml/nvml.go @@ -23,6 +23,7 @@ import ( const ( collectorID = "nvml" componentName = "workloadmeta-nvml" + nvidiaVendor = "nvidia" ) type collector struct { @@ -96,7 +97,7 @@ func (c *collector) Pull(_ context.Context) error { EntityMeta: workloadmeta.EntityMeta{ Name: name, }, - Vendor: "nvidia", + Vendor: nvidiaVendor, Device: name, } diff --git a/comp/core/workloadmeta/collectors/internal/nvml/nvml_test.go b/comp/core/workloadmeta/collectors/internal/nvml/nvml_test.go new file mode 100644 index 00000000000000..0793fc9665687f --- /dev/null +++ b/comp/core/workloadmeta/collectors/internal/nvml/nvml_test.go @@ -0,0 +1,48 @@ +// 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 + +package nvml + +import ( + "context" + "testing" + + "github.com/stretchr/testify/require" + + workloadmeta "github.com/DataDog/datadog-agent/comp/core/workloadmeta/def" + "github.com/DataDog/datadog-agent/pkg/gpu/testutil" +) + +func TestPull(t *testing.T) { + wmetaMock := testutil.GetWorkloadMetaMock(t) + nvmlMock := testutil.GetBasicNvmlMock() + + c := &collector{ + id: collectorID, + catalog: workloadmeta.NodeAgent, + store: wmetaMock, + nvmlLib: nvmlMock, + } + + c.Pull(context.Background()) + + gpus := wmetaMock.ListGPUs() + require.Equal(t, len(testutil.GPUUUIDs), len(gpus)) + + foundIDs := make(map[string]bool) + for _, gpu := range gpus { + foundIDs[gpu.ID] = true + + require.Equal(t, nvidiaVendor, gpu.Vendor) + require.Equal(t, testutil.DefaultGPUName, gpu.Name) + require.Equal(t, testutil.DefaultGPUName, gpu.Device) + } + + for _, uuid := range testutil.GPUUUIDs { + require.True(t, foundIDs[uuid], "GPU with UUID %s not found", uuid) + } +} diff --git a/pkg/gpu/testutil/mocks.go b/pkg/gpu/testutil/mocks.go index 009385125600a3..d90576972a5a81 100644 --- a/pkg/gpu/testutil/mocks.go +++ b/pkg/gpu/testutil/mocks.go @@ -45,6 +45,9 @@ var GPUCores = []int{DefaultGpuCores, 20, 30, 40, 50, 60, 70} // DefaultGpuUUID is the UUID for the default device returned by the mock var DefaultGpuUUID = GPUUUIDs[0] +// DefaultGPUName is the name for the default device returned by the mock +var DefaultGPUName = "Tesla T4" + // GetDeviceMock returns a mock of the nvml.Device with the given UUID. func GetDeviceMock(deviceIdx int) *nvmlmock.Device { return &nvmlmock.Device{ @@ -57,6 +60,9 @@ func GetDeviceMock(deviceIdx int) *nvmlmock.Device { GetUUIDFunc: func() (string, nvml.Return) { return GPUUUIDs[deviceIdx], nvml.SUCCESS }, + GetNameFunc: func() (string, nvml.Return) { + return DefaultGPUName, nvml.SUCCESS + }, } }