Rmalliu/feat/multi share mps

api/config/v1/config.go

@@ -68,12 +68,11 @@ func NewConfig(c *cli.Context, flags []cli.Flag) (*Config, error) {
 		config.Flags.NvidiaDevRoot = config.Flags.NvidiaDriverRoot
 	}
 
-	// We explicitly set sharing.mps.failRequestsGreaterThanOne = true
-	// This can be relaxed in certain cases -- such as a single GPU -- but
-	// requires additional logic around when it's OK to combine requests and
-	// makes the semantics of a request unclear.
-	if config.Sharing.MPS != nil {
-		config.Sharing.MPS.FailRequestsGreaterThanOne = true
+	// Default sharing.mps.failRequestsGreaterThanOne to true if not explicitly set in config.
+	// Set it to false in the device plugin config to allow pods to request more than one MPS-shared GPU unit.
+	if config.Sharing.MPS != nil && config.Sharing.MPS.FailRequestsGreaterThanOne == nil {
+		t := true
+		config.Sharing.MPS.FailRequestsGreaterThanOne = &t
 	}
 
 	return config, nil

api/config/v1/replicas.go

@@ -28,7 +28,7 @@ import (
 // ReplicatedResources defines generic options for replicating devices.
 type ReplicatedResources struct {
 	RenameByDefault            bool                 `json:"renameByDefault,omitempty"            yaml:"renameByDefault,omitempty"`
-	FailRequestsGreaterThanOne bool                 `json:"failRequestsGreaterThanOne,omitempty" yaml:"failRequestsGreaterThanOne,omitempty"`
+	FailRequestsGreaterThanOne *bool                `json:"failRequestsGreaterThanOne,omitempty" yaml:"failRequestsGreaterThanOne,omitempty"`
 	Resources                  []ReplicatedResource `json:"resources,omitempty"                  yaml:"resources,omitempty"`
 }
 
@@ -179,14 +179,11 @@ func (s *ReplicatedResources) UnmarshalJSON(b []byte) error {
 		return err
 	}
 
-	failRequestsGreaterThanOne, exists := ts["failRequestsGreaterThanOne"]
-	if !exists {
-		failRequestsGreaterThanOne = []byte(`false`)
-	}
-
-	err = json.Unmarshal(failRequestsGreaterThanOne, &s.FailRequestsGreaterThanOne)
-	if err != nil {
-		return err
+	if failRequestsGreaterThanOne, exists := ts["failRequestsGreaterThanOne"]; exists {
+		err = json.Unmarshal(failRequestsGreaterThanOne, &s.FailRequestsGreaterThanOne)
+		if err != nil {
+			return err
+		}
 	}
 
 	resources, exists := ts["resources"]

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"
 }

deployments/container/Dockerfile

@@ -30,11 +30,11 @@
         aarch64) ARCH='arm64' ;; \
         *) echo "unsupported architecture" ; exit 1 ;; \
     esac; \
-    wget -nv -O - https://storage.googleapis.com/golang/go${GOLANG_VERSION}.linux-${ARCH}.tar.gz \
+    wget -nv -O - https://go.dev/dl/go${GOLANG_VERSION}.linux-${ARCH}.tar.gz \
     | tar -C /usr/local -xz
 
 ENV GOPATH /go
 ENV PATH $GOPATH/bin:/usr/local/go/bin:$PATH
 
 WORKDIR /build
 COPY . .

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.

internal/rm/rm.go

@@ -74,20 +74,14 @@ func (r *resourceManager) ValidateRequest(ids AnnotatedIDs) error {
 	// error out if more than one resource is being allocated.
 	includesReplicas := ids.AnyHasAnnotations()
 	numRequestedDevices := len(ids)
+	failRequestsGTOne := r.config.Sharing.ReplicatedResources().FailRequestsGreaterThanOne
 	switch r.config.Sharing.SharingStrategy() {
 	case spec.SharingStrategyTimeSlicing:
-		if includesReplicas && numRequestedDevices > 1 && r.config.Sharing.ReplicatedResources().FailRequestsGreaterThanOne {
+		if includesReplicas && numRequestedDevices > 1 && failRequestsGTOne != nil && *failRequestsGTOne {
 			return fmt.Errorf("%w: maximum request size for shared resources is 1; found %d", errInvalidRequest, numRequestedDevices)
 		}
 	case spec.SharingStrategyMPS:
-		// For MPS sharing, we explicitly ignore the FailRequestsGreaterThanOne
-		// value in the sharing settings.
-		// This setting was added to timeslicing after the initial release and
-		// is set to `false` to maintain backward compatibility with existing
-		// deployments. If we do extend MPS to allow multiple devices to be
-		// requested, the MPS API will be extended separately from the
-		// time-slicing API.
-		if includesReplicas && numRequestedDevices > 1 {
+		if includesReplicas && numRequestedDevices > 1 && failRequestsGTOne != nil && *failRequestsGTOne {
 			return fmt.Errorf("%w: maximum request size for shared resources is 1; found %d", errInvalidRequest, numRequestedDevices)
 		}
 	}

internal/rm/rm_test.go

@@ -20,6 +20,7 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/require"
+	"k8s.io/utils/ptr"
 
 	spec "github.com/NVIDIA/k8s-device-plugin/api/config/v1"
 )
@@ -94,7 +95,7 @@ func TestValidateRequest(t *testing.T) {
 			description: "timeslicing with two devices -- failRequestsGreaterThanOne",
 			sharing: spec.Sharing{
 				TimeSlicing: spec.ReplicatedResources{
-					FailRequestsGreaterThanOne: true,
+					FailRequestsGreaterThanOne: ptr.To(true),
 					Resources: []spec.ReplicatedResource{
 						{
 							Name:     "nvidia.com/gpu",
@@ -151,13 +152,12 @@ func TestValidateRequest(t *testing.T) {
 				"device1::1": nil,
 			},
 			requestDevicesIDs: []string{"device0::1", "device1::0"},
-			expectedError:     errInvalidRequest,
 		},
 		{
 			description: "MPS with two devices -- failRequestsGreaterThanOne",
 			sharing: spec.Sharing{
 				MPS: &spec.ReplicatedResources{
-					FailRequestsGreaterThanOne: true,
+					FailRequestsGreaterThanOne: ptr.To(true),
 					Resources: []spec.ReplicatedResource{
 						{
 							Name:     "nvidia.com/gpu",