Fix: Correct resource reporting for single gpu on shared gpu worker (#…

…824)
beam-cloud · Jan 3, 2025 · ac7c555 · ac7c555
1 parent 60f3eee
commit ac7c555
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 8 deletions.
diff --git a/pkg/worker/events.go b/pkg/worker/events.go
@@ -16,7 +16,7 @@ func (w *Worker) collectAndSendContainerMetrics(ctx context.Context, request *ty
 	ticker := time.NewTicker(w.config.Monitoring.ContainerMetricsInterval)
 	defer ticker.Stop()
 
-	monitor := NewProcessMonitor(containerPid, spec.Linux.Resources.Devices)
+	monitor := NewProcessMonitor(containerPid, spec.Linux.Resources.Devices, w.containerCudaManager.GetContainerGPUDevices(request.ContainerId))
 
 	for {
 		select {
@@ -74,10 +74,11 @@ type ProcessMonitor struct {
 	lastIO        process.IOCountersStat
 	lastNetIO     net.IOCountersStat
 	gpuInfoClient GPUInfoClient
+	gpuDeviceIds  []int
 }
 
-func NewProcessMonitor(pid int, devices []specs.LinuxDeviceCgroup) *ProcessMonitor {
-	return &ProcessMonitor{pid: int32(pid), devices: devices, gpuInfoClient: &NvidiaInfoClient{}}
+func NewProcessMonitor(pid int, devices []specs.LinuxDeviceCgroup, gpuDeviceIds []int) *ProcessMonitor {
+	return &ProcessMonitor{pid: int32(pid), devices: devices, gpuInfoClient: &NvidiaInfoClient{}, gpuDeviceIds: gpuDeviceIds}
 }
 
 func (m *ProcessMonitor) GetStatistics() (*ProcessStats, error) {
@@ -112,12 +113,8 @@ func (m *ProcessMonitor) GetStatistics() (*ProcessStats, error) {
 
 func (m *ProcessMonitor) fetchGPUMemory() *GPUInfoStat {
 	stat := &GPUInfoStat{}
-	availableDevices, err := m.gpuInfoClient.AvailableGPUDevices()
-	if err != nil {
-		return stat
-	}
 
-	for _, device := range availableDevices {
+	for _, device := range m.gpuDeviceIds {
 		stats, err := m.gpuInfoClient.GetGPUMemoryUsage(device)
 		if err == nil {
 			stat.MemoryUsed += uint64(stats.UsedCapacity)

diff --git a/pkg/worker/nvidia.go b/pkg/worker/nvidia.go
@@ -21,6 +21,7 @@ var (
 
 type GPUManager interface {
 	AssignGPUDevices(containerId string, gpuCount uint32) (*AssignedGpuDevices, error)
+	GetContainerGPUDevices(containerId string) []int
 	UnassignGPUDevices(containerId string)
 	InjectEnvVars(env []string, options *ContainerOptions) ([]string, bool)
 	InjectMounts(mounts []specs.Mount) []specs.Mount
@@ -120,6 +121,15 @@ func (c *ContainerNvidiaManager) AssignGPUDevices(containerId string, gpuCount u
 	}, nil
 }
 
+func (c *ContainerNvidiaManager) GetContainerGPUDevices(containerId string) []int {
+	gpuDevices, ok := c.gpuAllocationMap.Get(containerId)
+	if !ok {
+		return []int{}
+	}
+
+	return gpuDevices
+}
+
 func (c *ContainerNvidiaManager) chooseDevices(containerId string, requestedGpuCount uint32) ([]int, error) {
 	c.mu.Lock()
 	defer c.mu.Unlock()