Infra audit: add single-GPU routing and best-fit worker packing

pkg/common/config.default.yaml

@@ -285,6 +285,7 @@ worker:
   # non-standard k8s job spec
   imagePVCName: beta9-images
   jobResourcesEnforced: false
+  preferBestFit: false
   containerResourceLimits:
     cpuEnforced: true
     cpuAffinityEnforced: false

pkg/scheduler/pool.go

@@ -76,6 +76,28 @@ type WorkerPoolCapacity struct {
 	PendingGpu    uint
 }
 
+func workerNodeSelector(jobSpec types.WorkerPoolJobSpecConfig, gpuCount uint32) map[string]string {
+	var selector map[string]string
+	if len(jobSpec.NodeSelector) > 0 {
+		selector = make(map[string]string, len(jobSpec.NodeSelector))
+		for key, value := range jobSpec.NodeSelector {
+			selector[key] = value
+		}
+	}
+
+	if gpuCount != 1 || len(jobSpec.SingleGPUNodeSelector) == 0 {
+		return selector
+	}
+
+	if selector == nil {
+		selector = make(map[string]string, len(jobSpec.SingleGPUNodeSelector))
+	}
+	for key, value := range jobSpec.SingleGPUNodeSelector {
+		selector[key] = value
+	}
+	return selector
+}
+
 type WorkerPoolControllerOptions struct {
 	Name            string
 	Context         context.Context

pkg/scheduler/pool_job_spec_test.go

@@ -0,0 +1,185 @@
+package scheduler
+
+import (
+	"testing"
+
+	"github.com/beam-cloud/beta9/pkg/types"
+	"github.com/tj/assert"
+)
+
+func TestWorkerNodeSelector(t *testing.T) {
+	jobSpec := types.WorkerPoolJobSpecConfig{
+		NodeSelector: map[string]string{
+			"karpenter.sh/nodepool": "gpu",
+			"kubernetes.io/arch":    "amd64",
+		},
+		SingleGPUNodeSelector: map[string]string{
+			"karpenter.sh/nodepool":                "gpu-single",
+			"karpenter.k8s.aws/instance-gpu-count": "1",
+		},
+	}
+
+	tests := []struct {
+		name     string
+		gpuCount uint32
+		want     map[string]string
+	}{
+		{
+			name: "CPU uses base selector",
+			want: map[string]string{
+				"karpenter.sh/nodepool": "gpu",
+				"kubernetes.io/arch":    "amd64",
+			},
+		},
+		{
+			name:     "single GPU overlays selector",
+			gpuCount: 1,
+			want: map[string]string{
+				"karpenter.sh/nodepool":                "gpu-single",
+				"kubernetes.io/arch":                   "amd64",
+				"karpenter.k8s.aws/instance-gpu-count": "1",
+			},
+		},
+		{
+			name:     "multiple GPUs use base selector",
+			gpuCount: 2,
+			want: map[string]string{
+				"karpenter.sh/nodepool": "gpu",
+				"kubernetes.io/arch":    "amd64",
+			},
+		},
+		{
+			name:     "four GPUs use base selector",
+			gpuCount: 4,
+			want: map[string]string{
+				"karpenter.sh/nodepool": "gpu",
+				"kubernetes.io/arch":    "amd64",
+			},
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			got := workerNodeSelector(jobSpec, test.gpuCount)
+			assert.Equal(t, test.want, got)
+
+			got["mutated"] = "true"
+			assert.Equal(t, "gpu", jobSpec.NodeSelector["karpenter.sh/nodepool"])
+			assert.Equal(t, "gpu-single", jobSpec.SingleGPUNodeSelector["karpenter.sh/nodepool"])
+			_, baseMutated := jobSpec.NodeSelector["mutated"]
+			_, singleMutated := jobSpec.SingleGPUNodeSelector["mutated"]
+			assert.False(t, baseMutated)
+			assert.False(t, singleMutated)
+		})
+	}
+}
+
+func TestWorkerNodeSelectorSingleGPUWithoutBase(t *testing.T) {
+	jobSpec := types.WorkerPoolJobSpecConfig{
+		SingleGPUNodeSelector: map[string]string{
+			"karpenter.sh/nodepool": "gpu-single",
+		},
+	}
+
+	assert.Equal(t, map[string]string{
+		"karpenter.sh/nodepool": "gpu-single",
+	}, workerNodeSelector(jobSpec, 1))
+	assert.Nil(t, workerNodeSelector(jobSpec, 0))
+}
+
+func TestWorkerJobBuildersUseGPUCountSelectorsWithoutMutatingConfig(t *testing.T) {
+	jobSpec := types.WorkerPoolJobSpecConfig{
+		NodeSelector: map[string]string{
+			"karpenter.sh/nodepool": "gpu",
+			"kubernetes.io/arch":    "amd64",
+		},
+		SingleGPUNodeSelector: map[string]string{
+			"karpenter.sh/nodepool":                "gpu-single",
+			"karpenter.k8s.aws/instance-gpu-count": "1",
+		},
+	}
+	workerConfig := types.WorkerConfig{
+		DefaultWorkerCPURequest:    1000,
+		DefaultWorkerMemoryRequest: 1024,
+	}
+	localController := &LocalKubernetesWorkerPoolController{
+		name:             "local-gpu",
+		config:           types.AppConfig{Worker: workerConfig},
+		workerPoolConfig: types.WorkerPoolConfig{JobSpec: jobSpec},
+	}
+	providerController := &ProviderWorkerPoolController{
+		name:             "provider-gpu",
+		config:           types.AppConfig{Worker: workerConfig},
+		workerPoolConfig: types.WorkerPoolConfig{JobSpec: jobSpec},
+	}
+
+	builders := []struct {
+		name  string
+		build func(uint32) map[string]string
+	}{
+		{
+			name: "local Job",
+			build: func(gpuCount uint32) map[string]string {
+				job, _ := localController.createWorkerJob("worker-local", 1000, 1024, "A10G", gpuCount, "token")
+				return job.Spec.Template.Spec.NodeSelector
+			},
+		},
+		{
+			name: "provider Job",
+			build: func(gpuCount uint32) map[string]string {
+				job, _ := providerController.buildWorkerJob("worker-provider", "machine-provider", 1000, 1024, "A10G", gpuCount, nil)
+				return job.Spec.Template.Spec.NodeSelector
+			},
+		},
+	}
+	tests := []struct {
+		name     string
+		gpuCount uint32
+		want     map[string]string
+	}{
+		{
+			name:     "one GPU overlays selector",
+			gpuCount: 1,
+			want: map[string]string{
+				"karpenter.sh/nodepool":                "gpu-single",
+				"kubernetes.io/arch":                   "amd64",
+				"karpenter.k8s.aws/instance-gpu-count": "1",
+			},
+		},
+		{
+			name:     "two GPUs retain base selector",
+			gpuCount: 2,
+			want: map[string]string{
+				"karpenter.sh/nodepool": "gpu",
+				"kubernetes.io/arch":    "amd64",
+			},
+		},
+		{
+			name:     "four GPUs retain base selector",
+			gpuCount: 4,
+			want: map[string]string{
+				"karpenter.sh/nodepool": "gpu",
+				"kubernetes.io/arch":    "amd64",
+			},
+		},
+	}
+
+	for _, builder := range builders {
+		t.Run(builder.name, func(t *testing.T) {
+			for _, test := range tests {
+				t.Run(test.name, func(t *testing.T) {
+					selector := builder.build(test.gpuCount)
+					assert.Equal(t, test.want, selector)
+
+					selector["mutated"] = "true"
+					assert.Equal(t, "gpu", jobSpec.NodeSelector["karpenter.sh/nodepool"])
+					assert.Equal(t, "gpu-single", jobSpec.SingleGPUNodeSelector["karpenter.sh/nodepool"])
+					_, baseMutated := jobSpec.NodeSelector["mutated"]
+					_, singleMutated := jobSpec.SingleGPUNodeSelector["mutated"]
+					assert.False(t, baseMutated)
+					assert.False(t, singleMutated)
+				})
+			}
+		})
+	}
+}

pkg/scheduler/pool_local.go

@@ -253,7 +253,7 @@ func (wpc *LocalKubernetesWorkerPoolController) createWorkerJob(workerId string,
 			HostNetwork:                   wpc.config.Worker.HostNetwork,
 			ImagePullSecrets:              imagePullSecrets,
 			RestartPolicy:                 corev1.RestartPolicyOnFailure,
-			NodeSelector:                  wpc.workerPoolConfig.JobSpec.NodeSelector,
+			NodeSelector:                  workerNodeSelector(wpc.workerPoolConfig.JobSpec, workerGpuCount),
 			Containers:                    containers,
 			Volumes:                       wpc.getWorkerVolumes(workerMemory),
 			EnableServiceLinks:            ptr.To(false),

pkg/scheduler/pool_provider.go

@@ -329,6 +329,25 @@ func (wpc *ProviderWorkerPoolController) createWorkerOnMachine(workerId, machine
 }
 
 func (wpc *ProviderWorkerPoolController) createWorkerJob(workerId, machineId string, cpu int64, memory int64, gpuType string, gpuCount uint32, token string) (*batchv1.Job, *types.Worker, error) {
+	workerGpuType := wpc.workerPoolConfig.GPUType
+	env, err := wpc.getWorkerEnvironment(workerId, machineId, cpu, memory, workerGpuType, gpuCount, token)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	job, worker := wpc.buildWorkerJob(workerId, machineId, cpu, memory, workerGpuType, gpuCount, env)
+	return job, worker, nil
+}
+
+func (wpc *ProviderWorkerPoolController) buildWorkerJob(
+	workerId string,
+	machineId string,
+	workerCpu int64,
+	workerMemory int64,
+	workerGpuType string,
+	workerGpuCount uint32,
+	env []corev1.EnvVar,
+) (*batchv1.Job, *types.Worker) {
 	jobName := fmt.Sprintf("%s-%s-%s", Beta9WorkerJobPrefix, wpc.name, workerId)
 	labels := map[string]string{
 		"app":                       Beta9WorkerLabelValue,
@@ -339,11 +358,6 @@ func (wpc *ProviderWorkerPoolController) createWorkerJob(workerId, machineId str
 		PrometheusScrapeKey:         strconv.FormatBool(wpc.config.Monitoring.Prometheus.ScrapeWorkers),
 	}
 
-	workerCpu := cpu
-	workerMemory := memory
-	workerGpuType := wpc.workerPoolConfig.GPUType
-	workerGpuCount := gpuCount
-
 	workerImage := fmt.Sprintf("%s/%s:%s",
 		wpc.config.Worker.ImageRegistry,
 		wpc.config.Worker.ImageName,
@@ -353,18 +367,13 @@ func (wpc *ProviderWorkerPoolController) createWorkerJob(workerId, machineId str
 	resources := corev1.ResourceRequirements{}
 	if workerGpuType != "" {
 		resources.Requests = corev1.ResourceList{
-			"nvidia.com/gpu": *resource.NewQuantity(int64(gpuCount), resource.DecimalSI),
+			"nvidia.com/gpu": *resource.NewQuantity(int64(workerGpuCount), resource.DecimalSI),
 		}
 		resources.Limits = corev1.ResourceList{
-			"nvidia.com/gpu": *resource.NewQuantity(int64(gpuCount), resource.DecimalSI),
+			"nvidia.com/gpu": *resource.NewQuantity(int64(workerGpuCount), resource.DecimalSI),
 		}
 	}
 
-	env, err := wpc.getWorkerEnvironment(workerId, machineId, workerCpu, workerMemory, workerGpuType, workerGpuCount, token)
-	if err != nil {
-		return nil, nil, err
-	}
-
 	containers := []corev1.Container{
 		{
 			Name:    defaultContainerName,
@@ -393,7 +402,7 @@ func (wpc *ProviderWorkerPoolController) createWorkerJob(workerId, machineId str
 			HostNetwork:                   true,
 			ImagePullSecrets:              imagePullSecrets,
 			RestartPolicy:                 corev1.RestartPolicyOnFailure,
-			NodeSelector:                  wpc.workerPoolConfig.JobSpec.NodeSelector,
+			NodeSelector:                  workerNodeSelector(wpc.workerPoolConfig.JobSpec, workerGpuCount),
 			Containers:                    containers,
 			Volumes:                       wpc.getWorkerVolumes(workerMemory),
 			EnableServiceLinks:            ptr.To(false),
@@ -433,7 +442,7 @@ func (wpc *ProviderWorkerPoolController) createWorkerJob(workerId, machineId str
 		Runtime:       wpc.workerPoolConfig.ContainerRuntime,
 		BuildVersion:  wpc.config.Worker.ImageTag,
 		Preemptable:   wpc.workerPoolConfig.Preemptable,
-	}, nil
+	}
 }
 
 func (wpc *ProviderWorkerPoolController) getWorkerEnvironment(workerId, machineId string, cpu int64, memory int64, gpuType string, gpuCount uint32, token string) ([]corev1.EnvVar, error) {

pkg/scheduler/scheduler.go

@@ -1164,7 +1164,15 @@ func (s *Scheduler) selectWorkerFromWorkersByStatus(workers []*types.Worker, req
 		if scoredWorkers[i].score != scoredWorkers[j].score {
 			return scoredWorkers[i].score > scoredWorkers[j].score
 		}
-		return workerFreeCapacityScore(scoredWorkers[i].worker, request) > workerFreeCapacityScore(scoredWorkers[j].worker, request)
+		iCapacity := workerFreeCapacityScore(scoredWorkers[i].worker, request)
+		jCapacity := workerFreeCapacityScore(scoredWorkers[j].worker, request)
+		if !s.config.Worker.PreferBestFit {
+			return iCapacity > jCapacity
+		}
+		if iCapacity != jCapacity {
+			return iCapacity < jCapacity
+		}
+		return scoredWorkers[i].worker.Id < scoredWorkers[j].worker.Id
 	})
 
 	return scoredWorkers[0].worker, nil