feat: manage multiple GPU bindings of an application (#5)

* feat: add an API to switch GPUBinding for app

* fix: add periodic GC for GPU bindings

Author: dkeven

cmd/scheduler/main.go

@@ -120,6 +120,7 @@ func start() error {
 
 	// start monitor metrics
 	go sher.RegisterFromNodeAnnotations()
+	go sher.CleanupGPUBindingsLoop()
 	go initMetrics(config.MetricsBindAddress)
 
 	// start http server
@@ -130,6 +131,7 @@ func start() error {
 	router.GET("/healthz", routes.HealthzRoute())
 
 	router.GET("/gpus", routes.ListGPUDetails(sher))
+	router.PUT("/gpus/assignments/bulk", routes.BulkManageAssignments(sher))
 	router.POST("/gpus/:id/mode", routes.SwitchGPUMode(sher))
 	router.POST("/gpus/:id/assign", routes.AssignGPUToApp(sher))
 	router.POST("/gpus/:id/unassign", routes.UnassignGPUFromApp(sher))

pkg/scheduler/routes/gpu_manage.go

@@ -52,6 +52,17 @@ type UnassignGPURequest struct {
 	AppName string `json:"appName"`
 }
 
+type SwitchAssignItem struct {
+	ID     string             `json:"id"`
+	Memory *resource.Quantity `json:"memory,omitempty"`
+}
+
+type SwitchAssignRequest struct {
+	AppName  string             `json:"appName"`
+	Unassign []SwitchAssignItem `json:"unassign,omitempty"`
+	Assign   []SwitchAssignItem `json:"assign,omitempty"`
+}
+
 func ListGPUInfos(s *scheduler.Scheduler) httprouter.Handle {
 	return func(w http.ResponseWriter, r *http.Request, _ httprouter.Params) {
 		klog.Infoln("Listing all GPUs")
@@ -545,3 +556,306 @@ func UnassignGPUFromApp(s *scheduler.Scheduler) httprouter.Handle {
 		w.WriteHeader(http.StatusOK)
 	}
 }
+
+// SwitchAssign performs an atomic switch of GPU assignments for an app:
+// - Unassigns specified GPU IDs if currently bound to the app (ignores non-existent bindings)
+// - Assigns specified GPU IDs (with optional memory for mem-slicing mode) to the app
+// - Enforces single-node binding policy across the app's final bindings
+// - For exclusive GPUs, evicts existing app bindings and restarts their pods
+// - Restarts the target app's pods only if its binding relationship changes
+func BulkManageAssignments(s *scheduler.Scheduler) httprouter.Handle {
+	return func(w http.ResponseWriter, r *http.Request, _ httprouter.Params) {
+		var req SwitchAssignRequest
+		if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
+			http.Error(w, fmt.Sprintf("failed to decode request: %v", err), http.StatusBadRequest)
+			return
+		}
+		if req.AppName == "" {
+			http.Error(w, "AppName is required", http.StatusBadRequest)
+			return
+		}
+
+		klog.Infof("SwitchAssign request for app %s: unassign=%v assign=%v", req.AppName, req.Unassign, req.Assign)
+		util.GPUManageLock.Lock()
+		defer util.GPUManageLock.Unlock()
+
+		nodes, err := s.ListNodes()
+		if err != nil {
+			klog.Errorln(err)
+			http.Error(w, fmt.Sprintf("failed to list nodes: %v", err), http.StatusInternalServerError)
+			return
+		}
+
+		// Build device maps
+		uuidToDevice := make(map[string]util.DeviceInfo)
+		uuidToNodeName := make(map[string]string)
+		for _, node := range nodes {
+			for _, device := range node.Devices {
+				uuidToDevice[device.ID] = device
+				uuidToNodeName[device.ID] = node.Node.Name
+			}
+		}
+
+		bindings, err := s.ListGPUBindings()
+		if err != nil {
+			klog.Errorln(err)
+			http.Error(w, err.Error(), http.StatusInternalServerError)
+			return
+		}
+
+		// Group existing bindings
+		currentBindingsByUUID := make(map[string]*v1alpha1.GPUBinding)
+		bindingsByUUID := make(map[string][]*v1alpha1.GPUBinding)
+		for _, b := range bindings {
+			bindingsByUUID[b.Spec.UUID] = append(bindingsByUUID[b.Spec.UUID], b)
+			if b.Spec.AppName == req.AppName {
+				currentBindingsByUUID[b.Spec.UUID] = b
+			}
+		}
+
+		// Build a quick lookup for GPUs that will be assigned, so unassign won't remove them
+		assignIDs := make(map[string]struct{})
+		for _, it := range req.Assign {
+			if it.ID != "" {
+				assignIDs[it.ID] = struct{}{}
+			}
+		}
+
+		// Plan unassignments (skip any UUID that is also in the assign list)
+		toUnassignNames := make([]string, 0)
+		unassignSet := make(map[string]struct{})
+		for _, it := range req.Unassign {
+			if it.ID == "" {
+				continue
+			}
+			if _, willAssign := assignIDs[it.ID]; willAssign {
+				continue
+			}
+			if b := currentBindingsByUUID[it.ID]; b != nil {
+				toUnassignNames = append(toUnassignNames, b.Name)
+				unassignSet[it.ID] = struct{}{}
+			}
+		}
+
+		// Plan assignments (patches for mem changes, creates for new bindings, evictions for exclusive)
+		type patchItem struct {
+			old *v1alpha1.GPUBinding
+			mem *resource.Quantity
+		}
+		patches := make([]patchItem, 0)
+		type createItem struct {
+			binding *v1alpha1.GPUBinding
+		}
+		creates := make([]createItem, 0)
+		evictUUIDs := make(map[string]struct{})
+		deleteOtherBindingNames := make([]string, 0)
+
+		seenAssign := make(map[string]struct{})
+		for _, it := range req.Assign {
+			if it.ID == "" {
+				continue
+			}
+			if _, duplicated := seenAssign[it.ID]; duplicated {
+				continue
+			}
+			seenAssign[it.ID] = struct{}{}
+
+			dev, ok := uuidToDevice[it.ID]
+			if !ok {
+				http.Error(w, fmt.Sprintf("GPU %s not found", it.ID), http.StatusNotFound)
+				return
+			}
+
+			if existing := currentBindingsByUUID[it.ID]; existing != nil {
+				// Already bound to this app
+				if dev.ShareMode != util.ShareModeMemSlicing {
+					// In exclusive/time-slicing, reassign to same GPU is a no-op
+					continue
+				}
+				// mem-slicing: treat missing/zero or unchanged memory as no-op
+				if it.Memory == nil || it.Memory.Value() == 0 ||
+					(existing.Spec.Memory != nil && it.Memory.Value() == existing.Spec.Memory.Value()) {
+					continue
+				}
+				// validate memory availability excluding this app's current allocation
+				totalUsed := int64(0)
+				for _, b := range bindingsByUUID[it.ID] {
+					if b.Spec.AppName == req.AppName {
+						continue
+					}
+					if b.Spec.Memory != nil {
+						totalUsed += b.Spec.Memory.Value()
+					}
+				}
+				if totalUsed+it.Memory.Value() > int64(dev.Devmem) {
+					err = fmt.Errorf("not enough memory on GPU %s, available: %d, request: %d for app %s",
+						it.ID, int64(dev.Devmem)-totalUsed, it.Memory.Value(), req.AppName)
+					klog.Warningln(err)
+					http.Error(w, err.Error(), http.StatusConflict)
+					return
+				}
+				patches = append(patches, patchItem{old: existing, mem: it.Memory})
+				continue
+			}
+
+			// Not currently bound to this app
+			if dev.ShareMode == util.ShareModeMemSlicing {
+				if it.Memory == nil || it.Memory.Value() == 0 {
+					err = fmt.Errorf("memory allocation is required for GPU %s in memory slicing mode, refuse assigning to app %s", it.ID, req.AppName)
+					klog.Warningln(err)
+					http.Error(w, err.Error(), http.StatusBadRequest)
+					return
+				}
+				totalUsed := int64(0)
+				for _, b := range bindingsByUUID[it.ID] {
+					if b.Spec.Memory != nil {
+						totalUsed += b.Spec.Memory.Value()
+					}
+				}
+				if totalUsed+it.Memory.Value() > int64(dev.Devmem) {
+					err = fmt.Errorf("not enough memory available on GPU %s, available: %d, request: %d, refuse assigning to app %s",
+						it.ID, int64(dev.Devmem)-totalUsed, it.Memory.Value(), req.AppName)
+					klog.Warningln(err)
+					http.Error(w, err.Error(), http.StatusConflict)
+					return
+				}
+			} else if dev.ShareMode == util.ShareModeExclusive {
+				// Plan eviction of other app(s) holding this GPU
+				for _, b := range bindingsByUUID[it.ID] {
+					if b.Spec.AppName != req.AppName {
+						deleteOtherBindingNames = append(deleteOtherBindingNames, b.Name)
+						evictUUIDs[it.ID] = struct{}{}
+					}
+				}
+			}
+
+			// Prepare new binding
+			mem := it.Memory
+			if dev.ShareMode != util.ShareModeMemSlicing {
+				mem = nil
+			}
+			newBinding := &v1alpha1.GPUBinding{
+				ObjectMeta: metav1.ObjectMeta{
+					Name: strings.ToLower(fmt.Sprintf("%s-%s-%d", req.AppName, it.ID, time.Now().Unix())),
+				},
+				Spec: v1alpha1.GPUBindingSpec{
+					UUID:    it.ID,
+					AppName: req.AppName,
+					PodSelector: &metav1.LabelSelector{
+						MatchLabels: map[string]string{
+							util.AppNameLabelKey: req.AppName,
+						},
+					},
+					Memory: mem,
+				},
+			}
+			creates = append(creates, createItem{binding: newBinding})
+		}
+
+		// Determine final set of UUIDs for the app after changes (for node policy check)
+		finalUUIDSet := make(map[string]struct{})
+		for uuid := range currentBindingsByUUID {
+			if _, toUn := unassignSet[uuid]; !toUn {
+				finalUUIDSet[uuid] = struct{}{}
+			}
+		}
+		for _, c := range creates {
+			finalUUIDSet[c.binding.Spec.UUID] = struct{}{}
+		}
+
+		// If no effective change to app's own bindings, return OK without restarting its pods
+		if len(toUnassignNames) == 0 && len(patches) == 0 && len(creates) == 0 {
+			w.WriteHeader(http.StatusOK)
+			return
+		}
+
+		// Enforce single-node binding policy
+		nodeSet := make(map[string]struct{})
+		for uuid := range finalUUIDSet {
+			nodeName := uuidToNodeName[uuid]
+			if nodeName == "" {
+				http.Error(w, fmt.Sprintf("GPU %s not found", uuid), http.StatusNotFound)
+				return
+			}
+			nodeSet[nodeName] = struct{}{}
+		}
+		if len(nodeSet) > 1 {
+			err = fmt.Errorf("app %s binding spans multiple nodes which is not allowed", req.AppName)
+			klog.Warningln(err)
+			http.Error(w, err.Error(), http.StatusConflict)
+			return
+		}
+
+		// Execute plan
+		// 1) Evict other apps for exclusive GPUs
+		if len(evictUUIDs) > 0 {
+			pods := s.ListPodsInfo()
+			for _, pod := range pods {
+				for _, pdev := range pod.Devices {
+					for _, cdevs := range pdev {
+						for _, cdev := range cdevs {
+							if _, needEvict := evictUUIDs[cdev.UUID]; needEvict {
+								klog.Infof("Evicting pod %s/%s occupying exclusive GPU %s", pod.Namespace, pod.Name, cdev.UUID)
+								if err := ctrlclient.IgnoreNotFound(client.GetClient().CoreV1().Pods(pod.Namespace).Delete(r.Context(), pod.Name, metav1.DeleteOptions{})); err != nil {
+									err = fmt.Errorf("failed to delete existing pod occupying GPU %s/%s: %v", pod.Namespace, pod.Name, err)
+									klog.Errorln(err)
+									http.Error(w, err.Error(), http.StatusInternalServerError)
+									return
+								}
+							}
+						}
+					}
+				}
+			}
+			for _, name := range deleteOtherBindingNames {
+				if err := ctrlclient.IgnoreNotFound(util.DeleteGPUBinding(r.Context(), name)); err != nil {
+					err = fmt.Errorf("failed to delete existing GPUBinding %s: %v", name, err)
+					klog.Errorln(err)
+					http.Error(w, err.Error(), http.StatusInternalServerError)
+					return
+				}
+			}
+		}
+
+		// 2) Restart this app's pods due to binding changes
+		if err := ctrlclient.IgnoreNotFound(util.DeletePodsBelongToApp(r.Context(), req.AppName)); err != nil {
+			err = fmt.Errorf("failed to delete existing pods of app %s: %v", req.AppName, err)
+			klog.Errorln(err)
+			http.Error(w, err.Error(), http.StatusInternalServerError)
+			return
+		}
+
+		// 3) Apply unassignments
+		for _, name := range toUnassignNames {
+			if err := ctrlclient.IgnoreNotFound(util.DeleteGPUBinding(r.Context(), name)); err != nil {
+				err = fmt.Errorf("failed to delete GPUBinding %s: %v", name, err)
+				klog.Errorln(err)
+				http.Error(w, err.Error(), http.StatusInternalServerError)
+				return
+			}
+		}
+
+		// 4) Apply memory patches
+		for _, p := range patches {
+			newBinding := p.old.DeepCopy()
+			newBinding.Spec.Memory = p.mem
+			if err := client.GPUClient.Patch(r.Context(), newBinding, ctrlclient.MergeFrom(p.old)); err != nil {
+				err = fmt.Errorf("failed to patch GPUBinding %s: %v", p.old.Name, err)
+				klog.Errorln(err)
+				http.Error(w, err.Error(), http.StatusInternalServerError)
+				return
+			}
+		}
+
+		// 5) Create new bindings
+		for _, c := range creates {
+			if err := s.CreateGPUBinding(r.Context(), c.binding); err != nil {
+				klog.Errorln(err)
+				http.Error(w, err.Error(), http.StatusInternalServerError)
+				return
+			}
+		}
+
+		w.WriteHeader(http.StatusOK)
+	}
+}