fix: add tolerations for custom pod template (#161)

Co-authored-by: guofei <[email protected]>
Azure · Nov 10, 2023 · 6c4fa45 · 6c4fa45
1 parent 30f2442
commit 6c4fa45
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 4 deletions.
diff --git a/pkg/inference/template_inference.go b/pkg/inference/template_inference.go
@@ -11,7 +11,7 @@ import (
 )
 
 func CreateTemplateInference(ctx context.Context, workspaceObj *kaitov1alpha1.Workspace, kubeClient client.Client) (client.Object, error) {
-	depObj := resources.GenerateDeploymentManifestWithPodTemplate(ctx, workspaceObj)
+	depObj := resources.GenerateDeploymentManifestWithPodTemplate(ctx, workspaceObj, tolerations)
 	err := resources.CreateResource(ctx, client.Object(depObj), kubeClient)
 	if client.IgnoreAlreadyExists(err) != nil {
 		return nil, err

diff --git a/pkg/resources/manifests.go b/pkg/resources/manifests.go
@@ -259,7 +259,7 @@ func GenerateDeploymentManifest(ctx context.Context, workspaceObj *kaitov1alpha1
 	}
 }
 
-func GenerateDeploymentManifestWithPodTemplate(ctx context.Context, workspaceObj *kaitov1alpha1.Workspace) *appsv1.Deployment {
+func GenerateDeploymentManifestWithPodTemplate(ctx context.Context, workspaceObj *kaitov1alpha1.Workspace, tolerations []corev1.Toleration) *appsv1.Deployment {
 	nodeRequirements := make([]corev1.NodeSelectorRequirement, 0, len(workspaceObj.Resource.LabelSelector.MatchLabels))
 	for key, value := range workspaceObj.Resource.LabelSelector.MatchLabels {
 		nodeRequirements = append(nodeRequirements, corev1.NodeSelectorRequirement{
@@ -292,7 +292,13 @@ func GenerateDeploymentManifestWithPodTemplate(ctx context.Context, workspaceObj
 		},
 	}
 
-	controller := true
+	// append tolerations
+	if templateCopy.Spec.Tolerations == nil {
+		templateCopy.Spec.Tolerations = tolerations
+	} else {
+		templateCopy.Spec.Tolerations = append(templateCopy.Spec.Tolerations, tolerations...)
+	}
+
 	return &appsv1.Deployment{
 		ObjectMeta: v1.ObjectMeta{
 			Name:      workspaceObj.Name,

diff --git a/pkg/resources/manifests_test.go b/pkg/resources/manifests_test.go
@@ -108,7 +108,7 @@ func TestGenerateDeploymentManifestWithPodTemplate(t *testing.T) {
 
 		workspace := utils.MockWorkspaceWithInferenceTemplate
 
-		obj := GenerateDeploymentManifestWithPodTemplate(context.TODO(), workspace)
+		obj := GenerateDeploymentManifestWithPodTemplate(context.TODO(), workspace, nil)
 
 		appSelector := map[string]string{
 			kaitov1alpha1.LabelWorkspaceName: workspace.Name,