cleanup

beam-cloud · Jan 3, 2025 · 574bf28 · 574bf28
1 parent c9fb429
commit 574bf28
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 6 deletions.
diff --git a/pkg/worker/events.go b/pkg/worker/events.go
@@ -16,8 +16,7 @@ func (w *Worker) collectAndSendContainerMetrics(ctx context.Context, request *ty
 	ticker := time.NewTicker(w.config.Monitoring.ContainerMetricsInterval)
 	defer ticker.Stop()
 
-	gpuDeviceIds, _ := w.containerCudaManager.GetContainerGPUDevices(request.ContainerId)
-	monitor := NewProcessMonitor(containerPid, spec.Linux.Resources.Devices, gpuDeviceIds)
+	monitor := NewProcessMonitor(containerPid, spec.Linux.Resources.Devices, w.containerCudaManager.GetContainerGPUDevices(request.ContainerId))
 
 	for {
 		select {

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