feat(scheduler): implement naive scheduler release with resource tracking

0x5457 · 0x5457 · commit 4d3bc2df6920 · 2024-12-06T08:15:51.000Z
- Update Release method in naive scheduler to properly track and restore resources
- Add requeue duration for pending connections
- Fix status updates in TensorFusionConnection controller
- Update tests to verify partial and full resource releases
- Make QosClass optional in TensorFusionConnection status
diff --git a/api/v1/tensorfusionconnection_types.go b/api/v1/tensorfusionconnection_types.go
@@ -48,7 +48,7 @@ const (
 type TensorFusionConnectionStatus struct {
 	Phase         TensorFusionConnectionPhase `json:"phase"`
 	ConnectionURL string                      `json:"connectionURL"`
-	QosClass      string                      `json:"qosClass"`
+	QosClass      string                      `json:"qosClass,omitempty"`
 	Node          string                      `json:"node,omitempty"`
 }
 
diff --git a/cmd/main.go b/cmd/main.go
@@ -152,15 +152,16 @@ func main() {
 
 	ctx := context.Background()
 	config := config.NewDefaultConfig()
+	scheduler := scheduler.NewNaiveScheduler()
 	if err = (&controller.TensorFusionConnectionReconciler{
-		Client: mgr.GetClient(),
-		Scheme: mgr.GetScheme(),
+		Client:    mgr.GetClient(),
+		Scheme:    mgr.GetScheme(),
+		Scheduler: scheduler,
 	}).SetupWithManager(mgr); err != nil {
 		setupLog.Error(err, "unable to create controller", "controller", "TensorFusionConnection")
 		os.Exit(1)
 	}
 
-	scheduler := scheduler.NewNaiveScheduler()
 	if err = (&controller.GPUNodeReconciler{
 		Client:    mgr.GetClient(),
 		Scheme:    mgr.GetScheme(),
diff --git a/config/crd/bases/tensor-fusion.ai_tensorfusionconnections.yaml b/config/crd/bases/tensor-fusion.ai_tensorfusionconnections.yaml
@@ -100,7 +100,6 @@ spec:
             required:
             - connectionURL
             - phase
-            - qosClass
             type: object
         type: object
     served: true
diff --git a/config/manager/kustomization.yaml b/config/manager/kustomization.yaml
@@ -1,2 +1,4 @@
+namespace: tensor-fusion
+
 resources:
 - manager.yaml
diff --git a/config/samples/v1_gpunode.yaml b/config/samples/v1_gpunode.yaml
@@ -5,5 +5,13 @@ metadata:
     app.kubernetes.io/name: tensor-fusion-operator
     app.kubernetes.io/managed-by: kustomize
   name: gpunode-sample
-spec:
-  # TODO(user): Add fields here
+  namespace: tensor-fusion
+status:
+  capacity:
+    tflops: '200'
+    vram: 100Gi
+  available:
+    tflops: '200'
+    vram: 100Gi
+  devices: []
+  
diff --git a/config/samples/v1_tensorfusionconnection.yaml b/config/samples/v1_tensorfusionconnection.yaml
@@ -5,5 +5,12 @@ metadata:
     app.kubernetes.io/name: tensor-fusion-operator
     app.kubernetes.io/managed-by: kustomize
   name: tensorfusionconnection-sample
+  namespace: tensor-fusion
 spec:
-  # TODO(user): Add fields here
+  resources:
+    limit:
+      tflops: '100'
+      vram: 8Gi
+    request:
+      tflops: '20'
+      vram: 9Gi
diff --git a/config/webhook/kustomization.yaml b/config/webhook/kustomization.yaml
@@ -1,6 +1,9 @@
+namespace: tensor-fusion
+
 resources:
 - manifests.yaml
 - service.yaml
 
 configurations:
 - kustomizeconfig.yaml
+
diff --git a/go.mod b/go.mod
@@ -7,6 +7,7 @@ require (
 	github.com/gin-gonic/gin v1.10.0
 	github.com/onsi/ginkgo/v2 v2.19.0
 	github.com/onsi/gomega v1.33.1
+	gomodules.xyz/jsonpatch/v2 v2.4.0
 	k8s.io/api v0.31.0
 	k8s.io/apimachinery v0.31.0
 	k8s.io/client-go v0.31.0
@@ -97,11 +98,11 @@ require (
 	golang.org/x/text v0.16.0 // indirect
 	golang.org/x/time v0.3.0 // indirect
 	golang.org/x/tools v0.21.1-0.20240508182429-e35e4ccd0d2d // indirect
-	gomodules.xyz/jsonpatch/v2 v2.4.0 // indirect
 	google.golang.org/genproto/googleapis/api v0.0.0-20240528184218-531527333157 // indirect
 	google.golang.org/genproto/googleapis/rpc v0.0.0-20240701130421-f6361c86f094 // indirect
 	google.golang.org/grpc v1.65.0 // indirect
 	google.golang.org/protobuf v1.34.2 // indirect
+	gopkg.in/evanphx/json-patch.v4 v4.12.0 // indirect
 	gopkg.in/inf.v0 v0.9.1 // indirect
 	gopkg.in/yaml.v2 v2.4.0 // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect
diff --git a/internal/constants/constants.go b/internal/constants/constants.go
@@ -1,5 +1,7 @@
 package constants
 
+import "time"
+
 const (
 	// TensorFusionDomain is the domain prefix used for all tensor-fusion.ai related annotations and finalizers
 	TensorFusionDomain = "tensor-fusion.ai"
@@ -12,4 +14,6 @@ const (
 	EnableContainerAnnotationFormat = TensorFusionDomain + "/enable-%s"
 	TFLOPSContainerAnnotationFormat = TensorFusionDomain + "/tflops-%s"
 	VRAMContainerAnnotationFormat   = TensorFusionDomain + "/vram-%s"
+
+	PendingRequeueDuration = time.Second * 3
 )
diff --git a/internal/controller/tensorfusionconnection_controller.go b/internal/controller/tensorfusionconnection_controller.go
@@ -96,14 +96,16 @@ func (r *TensorFusionConnectionReconciler) Reconcile(ctx context.Context, req ct
 	// If status is not set or pending, try to schedule
 	if connection.Status.Phase == "" || connection.Status.Phase == tfv1.TensorFusionConnectionPending {
 		// Try to get an available node from scheduler
-		node, err := r.Scheduler.Schedule(connection.Spec.Resources.Request)
+		var err error
+		node, err = r.Scheduler.Schedule(connection.Spec.Resources.Request)
 		if err != nil {
-			log.Error(err, "Failed to schedule connection")
+			log.Info(err.Error())
 			connection.Status.Phase = tfv1.TensorFusionConnectionPending
 		} else if node != nil {
 			connection.Status.Phase = tfv1.TensorFusionConnectionRunning
 			connection.Status.ConnectionURL = worker.GenerateConnectionURL(node, connection)
-			connection.Status.Node = node.Name // Store the node name for cleanup
+			// Store the node name for cleanup
+			connection.Status.Node = node.Name
 		} else {
 			connection.Status.Phase = tfv1.TensorFusionConnectionPending
 		}
@@ -113,6 +115,9 @@ func (r *TensorFusionConnectionReconciler) Reconcile(ctx context.Context, req ct
 		return ctrl.Result{}, err
 	}
 
+	if connection.Status.Phase == tfv1.TensorFusionConnectionPending {
+		return ctrl.Result{RequeueAfter: constants.PendingRequeueDuration}, nil
+	}
 	return ctrl.Result{}, nil
 }
 
@@ -124,7 +129,7 @@ func (r *TensorFusionConnectionReconciler) handleDeletion(ctx context.Context, c
 
 	// Get the node
 	node := &tfv1.GPUNode{}
-	if err := r.Get(ctx, client.ObjectKey{Name: connection.Status.Node}, node); err != nil {
+	if err := r.Get(ctx, client.ObjectKey{Name: connection.Status.Node, Namespace: connection.Namespace}, node); err != nil {
 		if errors.IsNotFound(err) {
 			// Node is already gone, nothing to do
 			return nil
@@ -133,11 +138,11 @@ func (r *TensorFusionConnectionReconciler) handleDeletion(ctx context.Context, c
 	}
 
 	// Release the resources
-	if err := r.Scheduler.Release(node); err != nil {
+	if err := r.Scheduler.Release(connection.Spec.Resources.Request, node); err != nil {
 		return err
 	}
 
-	return nil
+	return r.MustUpdateStatus(ctx, connection, node)
 }
 
 // Helper functions to handle finalizers
@@ -172,8 +177,7 @@ func (r *TensorFusionConnectionReconciler) MustUpdateStatus(ctx context.Context,
 		}
 
 		// Update the status fields we care about
-		latestConnection.Status.Phase = connection.Status.Phase
-		latestConnection.Status.ConnectionURL = connection.Status.ConnectionURL
+		latestConnection.Status = connection.Status
 
 		// Update the connection status
 		if err := r.Status().Update(ctx, latestConnection); err != nil {
diff --git a/internal/scheduler/naive.go b/internal/scheduler/naive.go
@@ -60,7 +60,7 @@ func (s *NaiveScheduler) OnDelete(node *tfv1.GPUNode) {
 }
 
 // Release implements Scheduler interface
-func (s *NaiveScheduler) Release(node *tfv1.GPUNode) error {
+func (s *NaiveScheduler) Release(request tfv1.Resource, node *tfv1.GPUNode) error {
 	s.Lock()
 	defer s.Unlock()
 
@@ -69,7 +69,10 @@ func (s *NaiveScheduler) Release(node *tfv1.GPUNode) error {
 		return fmt.Errorf("node %s not found", node.Name)
 	}
 
-	// Reset the node's available resources to its capacity
-	existingNode.Status.Available = existingNode.Status.Capacity
+	// Add back the released resources
+	existingNode.Status.Available.Tflops.Add(request.Tflops)
+	existingNode.Status.Available.Vram.Add(request.Vram)
+	// output the updated node
+	node.Status.Available = existingNode.Status.Available
 	return nil
 }
diff --git a/internal/scheduler/naive_test.go b/internal/scheduler/naive_test.go
@@ -168,14 +168,18 @@ func TestNaiveScheduler_NodeOperations(t *testing.T) {
 
 func TestNaiveScheduler_Release(t *testing.T) {
 	tests := []struct {
-		name      string
-		node      *tfv1.GPUNode
-		schedule  *tfv1.Resource
-		wantError bool
+		name                string
+		node               *tfv1.GPUNode
+		schedule           *tfv1.Resource
+		release            *tfv1.Resource
+		wantError          bool
+		wantRemainingTflops string
+		wantRemainingVram   string
 	}{
 		{
 			name:      "release non-existent node",
 			node:      createGPUNode("node1", "100", "16Gi"),
+			release:   &tfv1.Resource{},
 			wantError: true,
 		},
 		{
@@ -199,7 +203,42 @@ func TestNaiveScheduler_Release(t *testing.T) {
 				Tflops: resource.MustParse("50"),
 				Vram:   resource.MustParse("8Gi"),
 			},
-			wantError: false,
+			release: &tfv1.Resource{
+				Tflops: resource.MustParse("50"),
+				Vram:   resource.MustParse("8Gi"),
+			},
+			wantError:          false,
+			wantRemainingTflops: "100",
+			wantRemainingVram:   "16Gi",
+		},
+		{
+			name: "partial release",
+			node: &tfv1.GPUNode{
+				ObjectMeta: metav1.ObjectMeta{
+					Name: "node1",
+				},
+				Status: tfv1.GPUNodeStatus{
+					Capacity: tfv1.Resource{
+						Tflops: resource.MustParse("100"),
+						Vram:   resource.MustParse("16Gi"),
+					},
+					Available: tfv1.Resource{
+						Tflops: resource.MustParse("100"),
+						Vram:   resource.MustParse("16Gi"),
+					},
+				},
+			},
+			schedule: &tfv1.Resource{
+				Tflops: resource.MustParse("60"),
+				Vram:   resource.MustParse("10Gi"),
+			},
+			release: &tfv1.Resource{
+				Tflops: resource.MustParse("30"),
+				Vram:   resource.MustParse("5Gi"),
+			},
+			wantError:          false,
+			wantRemainingTflops: "70",
+			wantRemainingVram:   "11Gi",
 		},
 	}
 
@@ -220,26 +259,34 @@ func TestNaiveScheduler_Release(t *testing.T) {
 					}
 
 					// Verify resources were allocated
-					if node.Status.Available.Tflops.Cmp(resource.MustParse("50")) != 0 ||
-						node.Status.Available.Vram.Cmp(resource.MustParse("8Gi")) != 0 {
+					expectedTflops := tt.node.Status.Capacity.Tflops.DeepCopy()
+					expectedVram := tt.node.Status.Capacity.Vram.DeepCopy()
+					expectedTflops.Sub(tt.schedule.Tflops)
+					expectedVram.Sub(tt.schedule.Vram)
+					if node.Status.Available.Tflops.Cmp(expectedTflops) != 0 ||
+						node.Status.Available.Vram.Cmp(expectedVram) != 0 {
 						t.Errorf("Schedule() did not allocate resources correctly")
 						return
 					}
 				}
 			}
 
-			err := s.Release(tt.node)
+			err := s.Release(*tt.release, tt.node)
 			if (err != nil) != tt.wantError {
 				t.Errorf("Release() error = %v, wantError %v", err, tt.wantError)
 				return
 			}
 
 			if !tt.wantError {
-				// Verify resources were restored
+				// Verify resources were restored correctly
 				node := s.nodes[tt.node.Name]
-				if node.Status.Available.Tflops.Cmp(node.Status.Capacity.Tflops) != 0 ||
-					node.Status.Available.Vram.Cmp(node.Status.Capacity.Vram) != 0 {
-					t.Errorf("Release() did not restore resources correctly")
+				if node.Status.Available.Tflops.String() != tt.wantRemainingTflops ||
+					node.Status.Available.Vram.String() != tt.wantRemainingVram {
+					t.Errorf("Release() resources incorrect, got tflops=%v vram=%v, want tflops=%v vram=%v",
+						node.Status.Available.Tflops.String(),
+						node.Status.Available.Vram.String(),
+						tt.wantRemainingTflops,
+						tt.wantRemainingVram)
 				}
 			}
 		})
diff --git a/internal/scheduler/scheduler.go b/internal/scheduler/scheduler.go
@@ -12,7 +12,7 @@ type Scheduler interface {
 	Schedule(request tfv1.Resource) (*tfv1.GPUNode, error)
 
 	// Release frees the allocated resources of a node
-	Release(node *tfv1.GPUNode) error
+	Release(request tfv1.Resource, node *tfv1.GPUNode) error
 
 	// OnAdd is called when a new node is added
 	OnAdd(node *tfv1.GPUNode)
diff --git a/internal/webhook/v1/pod_webhook_test.go b/internal/webhook/v1/pod_webhook_test.go

Original file line number	Diff line number	Diff line change
`@@ -48,7 +48,7 @@ const (`
`48`	`48`	`type TensorFusionConnectionStatus struct {`
`49`	`49`	Phase TensorFusionConnectionPhase `json:"phase"`
`50`	`50`	ConnectionURL string `json:"connectionURL"`
`51`		- QosClass string `json:"qosClass"`
	`51`	+ QosClass string `json:"qosClass,omitempty"`
`52`	`52`	Node string `json:"node,omitempty"`
`53`	`53`	`}`
`54`	`54`