feat(mps): grant full GPU resources when a container requests all replicas

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: radu-malliu

cmd/mps-control-daemon/mps/daemon.go

@@ -249,32 +249,34 @@ func (d *Daemon) setComputeMode(mode computeMode) error {
 	return nil
 }
 
-// perDevicePinnedMemoryLimits returns the pinned memory limits for each device.
+// perDevicePinnedDeviceMemoryLimits returns the pinned memory limits for each
+// device as the full physical memory. Per-client throttling for shared
+// (non-full-node) requests is applied by the device plugin via the
+// CUDA_MPS_PINNED_DEVICE_MEM_LIMIT env var in the Allocate response; the MPS
+// server default must be at least that high or the env var is clamped down.
 func (m *Daemon) perDevicePinnedDeviceMemoryLimits() map[string]string {
 	totalMemoryInBytesPerDevice := make(map[string]uint64)
-	replicasPerDevice := make(map[string]uint64)
 	for _, device := range m.Devices() {
-		index := device.Index
-		totalMemoryInBytesPerDevice[index] = device.TotalMemory
-		replicasPerDevice[index] += 1
+		totalMemoryInBytesPerDevice[device.Index] = device.TotalMemory
 	}
 
 	limits := make(map[string]string)
 	for index, totalMemory := range totalMemoryInBytesPerDevice {
 		if totalMemory == 0 {
 			continue
 		}
-		replicas := replicasPerDevice[index]
-		limits[index] = fmt.Sprintf("%vM", totalMemory/replicas/1024/1024)
+		limits[index] = fmt.Sprintf("%vM", totalMemory/1024/1024)
 	}
 	return limits
 }
 
+// activeThreadPercentage returns the server-side default active thread
+// percentage. Set to 100 so it does not cap per-client env-var overrides;
+// per-client throttling is applied via CUDA_MPS_ACTIVE_THREAD_PERCENTAGE in
+// the Allocate response.
 func (m *Daemon) activeThreadPercentage() string {
 	if len(m.Devices()) == 0 {
 		return ""
 	}
-	replicasPerDevice := len(m.Devices()) / len(m.Devices().GetUUIDs())
-
-	return fmt.Sprintf("%d", 100/replicasPerDevice)
+	return "100"
 }

internal/plugin/mps.go

@@ -19,6 +19,8 @@ package plugin
 import (
 	"errors"
 	"fmt"
+	"sort"
+	"strings"
 
 	"k8s.io/klog/v2"
 	pluginapi "k8s.io/kubelet/pkg/apis/deviceplugin/v1beta1"
@@ -71,7 +73,7 @@ func (m *mpsOptions) waitForDaemon() error {
 	return nil
 }
 
-func (m *mpsOptions) updateReponse(response *pluginapi.ContainerAllocateResponse) {
+func (m *mpsOptions) updateReponse(response *pluginapi.ContainerAllocateResponse, grantedCount int) {
 	if m == nil || !m.enabled {
 		return
 	}
@@ -88,4 +90,54 @@ func (m *mpsOptions) updateReponse(response *pluginapi.ContainerAllocateResponse
 			HostPath:      m.hostRoot.ShmDir(m.resourceName),
 		},
 	)
+
+	// The MPS control daemon is configured with per-device defaults at full
+	// hardware. When a container has been granted every replica the plugin
+	// advertises on this node, kubelet's accounting guarantees no other pod
+	// holds any replica concurrently, so it can use the full daemon defaults —
+	// no client-side env vars are needed.
+	//
+	// For any other (non-full-node) grant, inject the 1/replicas per-device
+	// caps via the client env vars. The MPS server clamps the per-client value
+	// to no more than the daemon default; since the default is now full, these
+	// env vars are the effective cap. This preserves the historical behavior
+	// where every non-full grant got 1/replicas memory and thread percentage,
+	// regardless of how many replicas were granted.
+	devices := m.daemon.Devices()
+	total := len(devices)
+	if total == 0 || grantedCount >= total {
+		return
+	}
+
+	replicasByIndex := make(map[string]uint64)
+	totalMemoryByIndex := make(map[string]uint64)
+	for _, device := range devices {
+		replicasByIndex[device.Index]++
+		totalMemoryByIndex[device.Index] = device.TotalMemory
+	}
+
+	// All physical GPUs managed by this daemon share the same replicas value
+	// (it comes from a single config); pick any.
+	var replicasPerGPU uint64
+	for _, n := range replicasByIndex {
+		replicasPerGPU = n
+		break
+	}
+	if replicasPerGPU == 0 {
+		return
+	}
+
+	response.Envs["CUDA_MPS_ACTIVE_THREAD_PERCENTAGE"] = fmt.Sprintf("%d", 100/replicasPerGPU)
+
+	limits := make([]string, 0, len(totalMemoryByIndex))
+	for index, totalMemory := range totalMemoryByIndex {
+		if totalMemory == 0 {
+			continue
+		}
+		limits = append(limits, fmt.Sprintf("%s=%dM", index, totalMemory/replicasPerGPU/1024/1024))
+	}
+	if len(limits) > 0 {
+		sort.Strings(limits)
+		response.Envs["CUDA_MPS_PINNED_DEVICE_MEM_LIMIT"] = strings.Join(limits, ",")
+	}
 }

internal/plugin/server.go

@@ -330,7 +330,7 @@ func (plugin *nvidiaDevicePlugin) getAllocateResponse(requestIds []string) (*plu
 		}
 	}
 	if plugin.mps.enabled {
-		plugin.updateResponseForMPS(response)
+		plugin.updateResponseForMPS(response, len(requestIds))
 	}
 
 	// The following modifications are only made if at least one non-CDI device
@@ -361,8 +361,10 @@ func (plugin *nvidiaDevicePlugin) getAllocateResponse(requestIds []string) (*plu
 // updateResponseForMPS ensures that the ContainerAllocate response contains the information required to use MPS.
 // This includes per-resource pipe and log directories as well as a global daemon-specific shm
 // and assumes that an MPS control daemon has already been started.
-func (plugin nvidiaDevicePlugin) updateResponseForMPS(response *pluginapi.ContainerAllocateResponse) {
-	plugin.mps.updateReponse(response)
+// grantedCount is the number of replica IDs kubelet allocated to this container and
+// is used to decide whether to inject per-client limit env vars.
+func (plugin nvidiaDevicePlugin) updateResponseForMPS(response *pluginapi.ContainerAllocateResponse, grantedCount int) {
+	plugin.mps.updateReponse(response, grantedCount)
 }
 
 // updateResponseForCDI updates the specified response for the given device IDs.