feat: maintain all GPUs in memory instead of accessing the apiserver … (#186)

* feat: maintain all GPUs in memory instead of accessing the apiserver every time

* go mod tidy

* fix test

* fix lint

* feat: add sync interval parameter to GPU allocator

* fix test

* feat: update GPU allocator setup to return readiness channel

* fix: disable metricsserver for mgr in suite test

* feat: add allocator cleanup and initialization in main and suite tests

* fix: change readyCh channel to buffered in GPU allocator setup

* feat: implement dirty queue for tracking modified GPUs in GpuAllocator

* fix: if the gpu update fails, put it back in the queue

Author: 0x5457

.vscode/settings.json

@@ -20,8 +20,8 @@
         "CUDA",
         "cycjimmy",
         "dylib",
-        "essd",
         "envtest",
+        "essd",
         "Eventf",
         "finalizer",
         "Finalizers",
@@ -30,6 +30,7 @@
         "Gomega",
         "gopsutil",
         "gosec",
+        "gpuallocator",
         "gpunode",
         "gpunodeclasses",
         "gpunodes",

cmd/main.go

@@ -36,14 +36,15 @@ import (
 	"sigs.k8s.io/controller-runtime/pkg/client"
 	"sigs.k8s.io/controller-runtime/pkg/healthz"
 	"sigs.k8s.io/controller-runtime/pkg/log/zap"
+	"sigs.k8s.io/controller-runtime/pkg/manager"
 	"sigs.k8s.io/controller-runtime/pkg/metrics/filters"
 	metricsserver "sigs.k8s.io/controller-runtime/pkg/metrics/server"
 	"sigs.k8s.io/controller-runtime/pkg/webhook"
 
 	tfv1 "github.com/NexusGPU/tensor-fusion/api/v1"
 	"github.com/NexusGPU/tensor-fusion/internal/config"
 	"github.com/NexusGPU/tensor-fusion/internal/controller"
-	"github.com/NexusGPU/tensor-fusion/internal/scheduler"
+	"github.com/NexusGPU/tensor-fusion/internal/gpuallocator"
 	"github.com/NexusGPU/tensor-fusion/internal/server"
 	"github.com/NexusGPU/tensor-fusion/internal/server/router"
 	webhookcorev1 "github.com/NexusGPU/tensor-fusion/internal/webhook/v1"
@@ -169,7 +170,12 @@ func main() {
 
 	ctx := context.Background()
 
-	scheduler := scheduler.NewScheduler(mgr.GetClient())
+	// Initialize GPU allocator and set up watches
+	allocator := gpuallocator.NewGpuAllocator(ctx, mgr.GetClient(), 10*time.Second)
+	if _, err = allocator.SetupWithManager(ctx, mgr); err != nil {
+		setupLog.Error(err, "unable to set up GPU allocator watches")
+		os.Exit(1)
+	}
 	if err = (&controller.TensorFusionConnectionReconciler{
 		Client:   mgr.GetClient(),
 		Scheme:   mgr.GetScheme(),
@@ -270,7 +276,7 @@ func main() {
 	if err = (&controller.TensorFusionWorkloadReconciler{
 		Client:    mgr.GetClient(),
 		Scheme:    mgr.GetScheme(),
-		Scheduler: scheduler,
+		Allocator: allocator,
 		Recorder:  mgr.GetEventRecorderFor("tensorfusionworkload"),
 		GpuInfos:  &gpuInfos,
 	}).SetupWithManager(mgr); err != nil {
@@ -307,6 +313,21 @@ func main() {
 		}
 	}()
 
+	// cleanup function to stop the allocator
+	err = mgr.Add(manager.RunnableFunc(func(ctx context.Context) error {
+		// wait for the context to be done
+		<-ctx.Done()
+		setupLog.Info("stopping allocator")
+		if allocator != nil {
+			allocator.Stop()
+		}
+		return nil
+	}))
+	if err != nil {
+		setupLog.Error(err, "unable to add allocator cleanup to manager")
+		os.Exit(1)
+	}
+
 	setupLog.Info("starting manager")
 	if err := mgr.Start(ctrl.SetupSignalHandler()); err != nil {
 		setupLog.Error(err, "problem running manager")

internal/controller/suite_test.go

@@ -46,7 +46,7 @@ import (
 	tfv1 "github.com/NexusGPU/tensor-fusion/api/v1"
 	"github.com/NexusGPU/tensor-fusion/internal/config"
 	"github.com/NexusGPU/tensor-fusion/internal/constants"
-	scheduler "github.com/NexusGPU/tensor-fusion/internal/scheduler"
+	"github.com/NexusGPU/tensor-fusion/internal/gpuallocator"
 	"github.com/NexusGPU/tensor-fusion/internal/utils"
 	// +kubebuilder:scaffold:imports
 )
@@ -59,6 +59,7 @@ var k8sClient client.Client
 var testEnv *envtest.Environment
 var ctx context.Context
 var cancel context.CancelFunc
+var allocator *gpuallocator.GpuAllocator
 
 const (
 	timeout  = time.Second * 10
@@ -178,7 +179,10 @@ var _ = BeforeSuite(func() {
 	}).SetupWithManager(mgr)
 	Expect(err).ToNot(HaveOccurred())
 
-	scheduler := scheduler.NewScheduler(mgr.GetClient())
+	allocator = gpuallocator.NewGpuAllocator(ctx, mgr.GetClient(), 3*time.Second)
+	_, err = allocator.SetupWithManager(ctx, mgr)
+	Expect(err).ToNot(HaveOccurred())
+
 	err = (&TensorFusionConnectionReconciler{
 		Client:   mgr.GetClient(),
 		Scheme:   mgr.GetScheme(),
@@ -195,7 +199,7 @@ var _ = BeforeSuite(func() {
 	err = (&TensorFusionWorkloadReconciler{
 		Client:    mgr.GetClient(),
 		Scheme:    mgr.GetScheme(),
-		Scheduler: scheduler,
+		Allocator: allocator,
 		Recorder:  mgr.GetEventRecorderFor("TensorFusionWorkload"),
 		GpuInfos:  config.MockGpuInfo(),
 	}).SetupWithManager(mgr)
@@ -211,6 +215,7 @@ var _ = BeforeSuite(func() {
 
 var _ = AfterSuite(func() {
 	By("tearing down the test environment")
+	allocator.Stop()
 	cancel()
 	err := testEnv.Stop()
 	Expect(err).NotTo(HaveOccurred())

internal/controller/tensorfusionworkload_controller.go

@@ -33,8 +33,8 @@ import (
 	tfv1 "github.com/NexusGPU/tensor-fusion/api/v1"
 	"github.com/NexusGPU/tensor-fusion/internal/config"
 	"github.com/NexusGPU/tensor-fusion/internal/constants"
+	"github.com/NexusGPU/tensor-fusion/internal/gpuallocator"
 	"github.com/NexusGPU/tensor-fusion/internal/metrics"
-	scheduler "github.com/NexusGPU/tensor-fusion/internal/scheduler"
 	"github.com/NexusGPU/tensor-fusion/internal/utils"
 	"github.com/NexusGPU/tensor-fusion/internal/worker"
 	"github.com/lithammer/shortuuid/v4"
@@ -46,7 +46,7 @@ import (
 type TensorFusionWorkloadReconciler struct {
 	client.Client
 	Scheme    *runtime.Scheme
-	Scheduler scheduler.Scheduler
+	Allocator *gpuallocator.GpuAllocator
 	Recorder  record.EventRecorder
 	GpuInfos  *[]config.GpuInfo
 }
@@ -315,7 +315,7 @@ func (r *TensorFusionWorkloadReconciler) handlePodGPUCleanup(ctx context.Context
 	}
 
 	// Release GPU resources
-	if err := r.Scheduler.Release(ctx, workload.Spec.Resources.Requests, gpu); err != nil {
+	if err := r.Allocator.Dealloc(ctx, workload.Spec.Resources.Requests, gpu); err != nil {
 		log.Error(err, "Failed to release GPU resources, will retry", "gpu", gpuName, "pod", pod.Name)
 		return false, err
 	}
@@ -344,16 +344,19 @@ func (r *TensorFusionWorkloadReconciler) scaleUpWorkers(ctx context.Context, wor
 	// Create worker pods
 	for range count {
 		// Schedule GPU for the worker
-		gpu, err := r.Scheduler.Schedule(ctx, workload.Spec.PoolName, workload.Spec.Resources.Requests)
+		gpus, err := r.Allocator.Alloc(ctx, workload.Spec.PoolName, workload.Spec.Resources.Requests, 1)
 		if err != nil {
 			r.Recorder.Eventf(workload, corev1.EventTypeWarning, "ScheduleGPUFailed", "Failed to schedule GPU: %v", err)
 			return ctrl.Result{RequeueAfter: constants.PendingRequeueDuration}, nil
 		}
 
+		// Use the first GPU from the allocated array
+		gpu := gpus[0]
+
 		pod, err := r.tryStartWorker(ctx, workerGenerator, gpu, workload, hash)
 		if err != nil {
 			// Try to release the GPU resource if pod creation fails
-			releaseErr := r.Scheduler.Release(ctx, workload.Spec.Resources.Requests, gpu)
+			releaseErr := r.Allocator.Dealloc(ctx, workload.Spec.Resources.Requests, gpu)
 			if releaseErr != nil {
 				log.Error(releaseErr, "Failed to release GPU after pod creation failure")
 			}

internal/controller/tensorfusionworkload_controller_test.go

@@ -158,11 +158,15 @@ var _ = Describe("TensorFusionWorkload Controller", func() {
 			checkWorkerPodCount(workload)
 			checkWorkloadStatus(workload)
 
-			gpuList := tfEnv.GetPoolGpuList(0)
-			updatedGPU, ok := lo.Find(gpuList.Items, func(gpu tfv1.GPU) bool {
-				return gpu.Status.Available.Tflops.Equal(resource.MustParse("1990")) && gpu.Status.Available.Vram.Equal(resource.MustParse("1992Gi"))
-			})
-			Expect(ok).Should(BeTrue())
+			var updatedGPU tfv1.GPU
+			Eventually(func(g Gomega) bool {
+				gpuList := tfEnv.GetPoolGpuList(0)
+				ok := false
+				updatedGPU, ok = lo.Find(gpuList.Items, func(gpu tfv1.GPU) bool {
+					return gpu.Status.Available.Tflops.Equal(resource.MustParse("1990")) && gpu.Status.Available.Vram.Equal(resource.MustParse("1992Gi"))
+				})
+				return ok
+			}, timeout, interval).Should(BeTrue())
 
 			Expect(k8sClient.Get(ctx, key, workload)).Should(Succeed())
 			workloadCopy := workload.DeepCopy()

internal/scheduler/filter.go renamed to internal/gpuallocator/filter/filter.go

@@ -1,4 +1,4 @@
-package scheduler
+package filter
 
 import (
 	"context"