Repository: kaito-project/kaito

Issue:
LabelSelector is not being used to schedule finetuning job/pods on selected nodes
**Describe the bug**
I am trying to schedule fine tuning on BYO nodes and observed LabelSelector is not used to schedule finetuning job/pods on selected nodes. When fine tuning workspace is created, it creates a new node pool with selected SKU and pod is scheduled on other node pools that has matching SKU machines available with different labels.

On AKS, FineTuning workspace has label selectors but I don’t see any nodeaffinity or node selectors settings in the fine-tuning job or pod. Seems tuning preset doesn’t take into account label selectors [https://github.com/kaito-project/kaito/blob/9b79ccc4334f74ef5c312b646703b423ec802f6e/pkg/workspace/tuning/preset_tuning.go](https://nam06.safelinks.protection.outlook.com/?url=https%3A%2F%2Fgithub.com%2Fkaito-project%2Fkaito%2Fblob%2F9b79ccc4334f74ef5c312b646703b423ec802f6e%2Fpkg%2Fworkspace%2Ftuning%2Fpreset_tuning.go&data=05%7C02%7Cv-vakataria%40microsoft.com%7C29acbe9f4d7b460fdcf608dde1c17b8a%7C72f988bf86f141af91ab2d7cd011db47%7C1%7C0%7C638914943300039017%7CUnknown%7CTWFpbGZsb3d8eyJFbXB0eU1hcGkiOnRydWUsIlYiOiIwLjAuMDAwMCIsIlAiOiJXaW4zMiIsIkFOIjoiTWFpbCIsIldUIjoyfQ%3D%3D%7C0%7C%7C%7C&sdata=x%2FghMS63ViABQAczirfB3BzNqFKmBCMzmBBM6ScdDUM%3D&reserved=0)

Inference workspaces are hosted correctly as per label selectors, only finetuning workspace has this issue.

**Steps To Reproduce**
Pre-provision node pool for fine tuning, Use LabelSelectors to schedule on pre-provisioned node pool. If there are other node pools with requested SKU then pods will be scheduled randomly on any of these pools ignoring label selectors.

**Expected behavior**
Fine tuning should be scheduled on matching label selectors node pool only


Discussion:

Solution patch:
diff --git a/pkg/workspace/tuning/preset_tuning.go b/pkg/workspace/tuning/preset_tuning.go
index 2627eda..931cebc 100644
--- a/pkg/workspace/tuning/preset_tuning.go
+++ b/pkg/workspace/tuning/preset_tuning.go
@@ -29,7 +29,6 @@ import (
 
 	kaitov1beta1 "github.com/kaito-project/kaito/api/v1beta1"
 	pkgmodel "github.com/kaito-project/kaito/pkg/model"
-	"github.com/kaito-project/kaito/pkg/sku"
 	"github.com/kaito-project/kaito/pkg/utils"
 	"github.com/kaito-project/kaito/pkg/utils/consts"
 	"github.com/kaito-project/kaito/pkg/utils/generator"
@@ -167,7 +166,7 @@ func CreatePresetTuning(ctx context.Context, workspaceObj *kaitov1beta1.Workspac
 	}
 
 	podSpec, err := generator.GenerateManifest(gctx,
-		GenerateBasicTuningPodSpec(gpuConfig, skuNumGPUs),
+		GenerateBasicTuningPodSpec(skuNumGPUs),
 		SetTrainingResultVolume,
 		SetTrainingInput,
 		SetTrainingOutputImagePush,
@@ -191,7 +190,7 @@ func CreatePresetTuning(ctx context.Context, workspaceObj *kaitov1beta1.Workspac
 	return jobObj, nil
 }
 
-func GenerateBasicTuningPodSpec(gpuConfig *sku.GPUConfig, skuNumGPUs int) func(*generator.WorkspaceGeneratorContext, *corev1.PodSpec) error {
+func GenerateBasicTuningPodSpec(skuNumGPUs int) func(*generator.WorkspaceGeneratorContext, *corev1.PodSpec) error {
 	return func(ctx *generator.WorkspaceGeneratorContext, spec *corev1.PodSpec) error {
 		// additional volume
 		var volumes []corev1.Volume
@@ -256,6 +255,29 @@ func GenerateBasicTuningPodSpec(gpuConfig *sku.GPUConfig, skuNumGPUs int) func(*
 		}
 		spec.Volumes = volumes
 		spec.RestartPolicy = corev1.RestartPolicyNever
+
+		// Add node affinity based on label selector from workspace resource
+		nodeRequirements := make([]corev1.NodeSelectorRequirement, 0, len(ctx.Workspace.Resource.LabelSelector.MatchLabels))
+		for key, value := range ctx.Workspace.Resource.LabelSelector.MatchLabels {
+			nodeRequirements = append(nodeRequirements, corev1.NodeSelectorRequirement{
+				Key:      key,
+				Operator: corev1.NodeSelectorOpIn,
+				Values:   []string{value},
+			})
+		}
+
+		spec.Affinity = &corev1.Affinity{
+			NodeAffinity: &corev1.NodeAffinity{
+				RequiredDuringSchedulingIgnoredDuringExecution: &corev1.NodeSelector{
+					NodeSelectorTerms: []corev1.NodeSelectorTerm{
+						{
+							MatchExpressions: nodeRequirements,
+						},
+					},
+				},
+			},
+		}
+
 		return nil
 	}
 }



Test patch:
diff --git a/pkg/workspace/tuning/preset_tuning_test.go b/pkg/workspace/tuning/preset_tuning_test.go
index 4e249e0..e0b7f8c 100644
--- a/pkg/workspace/tuning/preset_tuning_test.go
+++ b/pkg/workspace/tuning/preset_tuning_test.go
@@ -26,6 +26,7 @@ import (
 	"sigs.k8s.io/controller-runtime/pkg/client/fake"
 
 	kaitov1beta1 "github.com/kaito-project/kaito/api/v1beta1"
+	pkgmodel "github.com/kaito-project/kaito/pkg/model"
 	"github.com/kaito-project/kaito/pkg/utils"
 	"github.com/kaito-project/kaito/pkg/utils/generator"
 	"github.com/kaito-project/kaito/pkg/workspace/image"
@@ -687,3 +688,93 @@ func TestPrepareTrainingOutput(t *testing.T) {
 		})
 	}
 }
+
+func TestGenerateBasicTuningPodSpec_NodeAffinity(t *testing.T) {
+	workspace := &kaitov1beta1.Workspace{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:      "test-workspace",
+			Namespace: "default",
+		},
+		Resource: kaitov1beta1.ResourceSpec{
+			LabelSelector: &metav1.LabelSelector{
+				MatchLabels: map[string]string{
+					"accelerator":   "nvidia-gpu",
+					"instance-type": "gpu-node",
+				},
+			},
+		},
+		Tuning: &kaitov1beta1.TuningSpec{
+			Preset: &kaitov1beta1.PresetSpec{
+				PresetMeta: kaitov1beta1.PresetMeta{
+					Name: "test-model",
+				},
+			},
+		},
+	}
+
+	gctx := &generator.WorkspaceGeneratorContext{
+		Ctx:       context.Background(),
+		Workspace: workspace,
+		Model:     &mockModel{},
+	}
+
+	// Create a basic pod spec
+	podSpec := &corev1.PodSpec{
+		Containers: []corev1.Container{
+			{
+				Name: workspace.Name,
+			},
+		},
+	}
+
+	// Get the modifier function
+	modifier := GenerateBasicTuningPodSpec(1)
+	err := modifier(gctx, podSpec)
+
+	// Verify no error
+	assert.NoError(t, err)
+
+	// Verify node affinity is set correctly
+	assert.NotNil(t, podSpec.Affinity)
+	assert.NotNil(t, podSpec.Affinity.NodeAffinity)
+	assert.NotNil(t, podSpec.Affinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution)
+
+	nodeSelector := podSpec.Affinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution
+	assert.Len(t, nodeSelector.NodeSelectorTerms, 1)
+	assert.Len(t, nodeSelector.NodeSelectorTerms[0].MatchExpressions, 2)
+
+	// Verify the match expressions contain the expected labels
+	expectedLabels := map[string]string{
+		"accelerator":   "nvidia-gpu",
+		"instance-type": "gpu-node",
+	}
+
+	for _, req := range nodeSelector.NodeSelectorTerms[0].MatchExpressions {
+		assert.Equal(t, corev1.NodeSelectorOpIn, req.Operator)
+		assert.Len(t, req.Values, 1)
+		assert.Equal(t, expectedLabels[req.Key], req.Values[0])
+	}
+}
+
+// mockModel implements pkgmodel.Model for testing
+type mockModel struct{}
+
+func (m *mockModel) GetInferenceParameters() *pkgmodel.PresetParam {
+	return &pkgmodel.PresetParam{
+		GPUCountRequirement: "1",
+	}
+}
+
+func (m *mockModel) GetTuningParameters() *pkgmodel.PresetParam {
+	return &pkgmodel.PresetParam{
+		GPUCountRequirement: "1",
+	}
+}
+
+func (m *mockModel) SupportDistributedInference() bool {
+	return false
+}
+
+func (m *mockModel) SupportTuning() bool {
+	return true
+}