diff --git a/api/config/v1/config.go b/api/config/v1/config.go index c102c42390..2f5cbc5a8a 100644 --- a/api/config/v1/config.go +++ b/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 diff --git a/api/config/v1/replicas.go b/api/config/v1/replicas.go index 0e00dd57aa..d28783a992 100644 --- a/api/config/v1/replicas.go +++ b/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"] diff --git a/cmd/mps-control-daemon/mps/daemon.go b/cmd/mps-control-daemon/mps/daemon.go index 0351289cc5..1d591a9ae5 100644 --- a/cmd/mps-control-daemon/mps/daemon.go +++ b/cmd/mps-control-daemon/mps/daemon.go @@ -249,14 +249,15 @@ 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) @@ -264,17 +265,18 @@ func (m *Daemon) perDevicePinnedDeviceMemoryLimits() map[string]string { 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" } diff --git a/deployments/container/Dockerfile b/deployments/container/Dockerfile index 54d4c08f07..79504ff7a2 100644 --- a/deployments/container/Dockerfile +++ b/deployments/container/Dockerfile @@ -30,7 +30,7 @@ RUN set -eux; \ 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 diff --git a/internal/plugin/mps.go b/internal/plugin/mps.go index 73ed94edfd..da5829173b 100644 --- a/internal/plugin/mps.go +++ b/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, ",") + } } diff --git a/internal/plugin/server.go b/internal/plugin/server.go index 6b55589b3e..606f593974 100644 --- a/internal/plugin/server.go +++ b/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. diff --git a/internal/rm/rm.go b/internal/rm/rm.go index e94f158472..76e23a7851 100644 --- a/internal/rm/rm.go +++ b/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) } } diff --git a/internal/rm/rm_test.go b/internal/rm/rm_test.go index 063ed40345..47f4a16db0 100644 --- a/internal/rm/rm_test.go +++ b/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",