Add tolerations for extended resources (#6033)

--------- Signed-off-by: troychiu <[email protected]>
flyteorg · Nov 27, 2024 · 1dcc8bf · 1dcc8bf
1 parent ab04192
commit 1dcc8bf
Show file tree

Hide file tree

Showing 10 changed files with 242 additions and 14 deletions.
diff --git a/flyteplugins/go/tasks/pluginmachinery/flytek8s/config/config.go b/flyteplugins/go/tasks/pluginmachinery/flytek8s/config/config.go
@@ -64,8 +64,9 @@ var (
 		DefaultPodTemplateResync: config2.Duration{
 			Duration: 30 * time.Second,
 		},
-		UpdateBaseBackoffDuration: 10,
-		UpdateBackoffRetries:      5,
+		UpdateBaseBackoffDuration:          10,
+		UpdateBackoffRetries:               5,
+		AddTolerationsForExtendedResources: []string{},
 	}
 
 	// K8sPluginConfigSection provides a singular top level config section for all plugins.
@@ -214,6 +215,9 @@ type K8sPluginConfig struct {
 
 	// Number of retries for exponential backoff when updating a resource.
 	UpdateBackoffRetries int `json:"update-backoff-retries" pflag:",Number of retries for exponential backoff when updating a resource."`
+
+	// Extended resources that should be added to the tolerations automatically.
+	AddTolerationsForExtendedResources []string `json:"add-tolerations-for-extended-resources" pflag:",Name of the extended resources for which tolerations should be added."`
 }
 
 // FlyteCoPilotConfig specifies configuration for the Flyte CoPilot system. FlyteCoPilot, allows running flytekit-less containers

diff --git a/flyteplugins/go/tasks/pluginmachinery/flytek8s/config/k8spluginconfig_flags.go b/flyteplugins/go/tasks/pluginmachinery/flytek8s/config/k8spluginconfig_flags.go
diff --git a/flyteplugins/go/tasks/pluginmachinery/flytek8s/config/k8spluginconfig_flags_test.go b/flyteplugins/go/tasks/pluginmachinery/flytek8s/config/k8spluginconfig_flags_test.go
diff --git a/flyteplugins/go/tasks/pluginmachinery/flytek8s/pod_helper.go b/flyteplugins/go/tasks/pluginmachinery/flytek8s/pod_helper.go
@@ -11,6 +11,7 @@ import (
 	"github.com/imdario/mergo"
 	v1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/util/sets"
 
 	"github.com/flyteorg/flyte/flyteidl/gen/pb-go/flyteidl/core"
 	pluginserrors "github.com/flyteorg/flyte/flyteplugins/go/tasks/errors"
@@ -445,6 +446,54 @@ func ApplyContainerImageOverride(podSpec *v1.PodSpec, containerImage string, pri
 	}
 }
 
+func addTolerationInPodSpec(podSpec *v1.PodSpec, toleration *v1.Toleration) *v1.PodSpec {
+	podTolerations := podSpec.Tolerations
+
+	var newTolerations []v1.Toleration
+	for i := range podTolerations {
+		if toleration.MatchToleration(&podTolerations[i]) {
+			return podSpec
+		}
+		newTolerations = append(newTolerations, podTolerations[i])
+	}
+	newTolerations = append(newTolerations, *toleration)
+	podSpec.Tolerations = newTolerations
+	return podSpec
+}
+
+func AddTolerationsForExtendedResources(podSpec *v1.PodSpec) *v1.PodSpec {
+	if podSpec == nil {
+		podSpec = &v1.PodSpec{}
+	}
+
+	resources := sets.NewString()
+	for _, container := range podSpec.Containers {
+		for _, extendedResource := range config.GetK8sPluginConfig().AddTolerationsForExtendedResources {
+			if _, ok := container.Resources.Requests[v1.ResourceName(extendedResource)]; ok {
+				resources.Insert(extendedResource)
+			}
+		}
+	}
+
+	for _, container := range podSpec.InitContainers {
+		for _, extendedResource := range config.GetK8sPluginConfig().AddTolerationsForExtendedResources {
+			if _, ok := container.Resources.Requests[v1.ResourceName(extendedResource)]; ok {
+				resources.Insert(extendedResource)
+			}
+		}
+	}
+
+	for _, resource := range resources.List() {
+		addTolerationInPodSpec(podSpec, &v1.Toleration{
+			Key:      resource,
+			Operator: v1.TolerationOpExists,
+			Effect:   v1.TaintEffectNoSchedule,
+		})
+	}
+
+	return podSpec
+}
+
 // ToK8sPodSpec builds a PodSpec and ObjectMeta based on the definition passed by the TaskExecutionContext. This
 // involves parsing the raw PodSpec definition and applying all Flyte configuration options.
 func ToK8sPodSpec(ctx context.Context, tCtx pluginsCore.TaskExecutionContext) (*v1.PodSpec, *metav1.ObjectMeta, string, error) {
@@ -460,6 +509,8 @@ func ToK8sPodSpec(ctx context.Context, tCtx pluginsCore.TaskExecutionContext) (*
 		return nil, nil, "", err
 	}
 
+	podSpec = AddTolerationsForExtendedResources(podSpec)
+
 	return podSpec, objectMeta, primaryContainerName, nil
 }
 

diff --git a/flyteplugins/go/tasks/pluginmachinery/flytek8s/pod_helper_test.go b/flyteplugins/go/tasks/pluginmachinery/flytek8s/pod_helper_test.go
@@ -3,6 +3,7 @@ package flytek8s
 import (
 	"context"
 	"encoding/json"
+	"fmt"
 	"io/ioutil"
 	"path/filepath"
 	"reflect"
@@ -2244,3 +2245,112 @@ func TestAddFlyteCustomizationsToContainer_SetConsoleUrl(t *testing.T) {
 		})
 	}
 }
+
+func TestAddTolerationsForExtendedResources(t *testing.T) {
+	gpuResourceName := v1.ResourceName("nvidia.com/gpu")
+	addTolerationResourceName := v1.ResourceName("foo/bar")
+	noTolerationResourceName := v1.ResourceName("foo/baz")
+	assert.NoError(t, config.SetK8sPluginConfig(&config.K8sPluginConfig{
+		GpuResourceName: gpuResourceName,
+		AddTolerationsForExtendedResources: []string{
+			gpuResourceName.String(),
+			addTolerationResourceName.String(),
+		},
+	}))
+
+	podSpec := &v1.PodSpec{
+		Containers: []v1.Container{
+			v1.Container{
+				Resources: v1.ResourceRequirements{
+					Requests: v1.ResourceList{
+						gpuResourceName:           resource.MustParse("1"),
+						addTolerationResourceName: resource.MustParse("1"),
+						noTolerationResourceName:  resource.MustParse("1"),
+					},
+				},
+			},
+		},
+		Tolerations: []v1.Toleration{
+			{
+				Key:      "foo",
+				Operator: v1.TolerationOpExists,
+				Effect:   v1.TaintEffectNoSchedule,
+			},
+		},
+	}
+
+	podSpec = AddTolerationsForExtendedResources(podSpec)
+	fmt.Printf("%v\n", podSpec.Tolerations)
+	assert.Equal(t, 3, len(podSpec.Tolerations))
+	assert.Equal(t, addTolerationResourceName.String(), podSpec.Tolerations[1].Key)
+	assert.Equal(t, v1.TolerationOpExists, podSpec.Tolerations[1].Operator)
+	assert.Equal(t, v1.TaintEffectNoSchedule, podSpec.Tolerations[1].Effect)
+	assert.Equal(t, gpuResourceName.String(), podSpec.Tolerations[2].Key)
+	assert.Equal(t, v1.TolerationOpExists, podSpec.Tolerations[2].Operator)
+	assert.Equal(t, v1.TaintEffectNoSchedule, podSpec.Tolerations[2].Effect)
+
+	podSpec = &v1.PodSpec{
+		InitContainers: []v1.Container{
+			v1.Container{
+				Resources: v1.ResourceRequirements{
+					Requests: v1.ResourceList{
+						gpuResourceName:           resource.MustParse("1"),
+						addTolerationResourceName: resource.MustParse("1"),
+						noTolerationResourceName:  resource.MustParse("1"),
+					},
+				},
+			},
+		},
+		Tolerations: []v1.Toleration{
+			{
+				Key:      "foo",
+				Operator: v1.TolerationOpExists,
+				Effect:   v1.TaintEffectNoSchedule,
+			},
+		},
+	}
+
+	podSpec = AddTolerationsForExtendedResources(podSpec)
+	assert.Equal(t, 3, len(podSpec.Tolerations))
+	assert.Equal(t, addTolerationResourceName.String(), podSpec.Tolerations[1].Key)
+	assert.Equal(t, v1.TolerationOpExists, podSpec.Tolerations[1].Operator)
+	assert.Equal(t, v1.TaintEffectNoSchedule, podSpec.Tolerations[1].Effect)
+	assert.Equal(t, gpuResourceName.String(), podSpec.Tolerations[2].Key)
+	assert.Equal(t, v1.TolerationOpExists, podSpec.Tolerations[2].Operator)
+	assert.Equal(t, v1.TaintEffectNoSchedule, podSpec.Tolerations[2].Effect)
+
+	podSpec = &v1.PodSpec{
+		Containers: []v1.Container{
+			v1.Container{
+				Resources: v1.ResourceRequirements{
+					Requests: v1.ResourceList{
+						gpuResourceName:           resource.MustParse("1"),
+						addTolerationResourceName: resource.MustParse("1"),
+						noTolerationResourceName:  resource.MustParse("1"),
+					},
+				},
+			},
+		},
+		Tolerations: []v1.Toleration{
+			{
+				Key:      "foo",
+				Operator: v1.TolerationOpExists,
+				Effect:   v1.TaintEffectNoSchedule,
+			},
+			{
+				Key:      gpuResourceName.String(),
+				Operator: v1.TolerationOpExists,
+				Effect:   v1.TaintEffectNoSchedule,
+			},
+		},
+	}
+
+	podSpec = AddTolerationsForExtendedResources(podSpec)
+	assert.Equal(t, 3, len(podSpec.Tolerations))
+	assert.Equal(t, gpuResourceName.String(), podSpec.Tolerations[1].Key)
+	assert.Equal(t, v1.TolerationOpExists, podSpec.Tolerations[1].Operator)
+	assert.Equal(t, v1.TaintEffectNoSchedule, podSpec.Tolerations[1].Effect)
+	assert.Equal(t, addTolerationResourceName.String(), podSpec.Tolerations[2].Key)
+	assert.Equal(t, v1.TolerationOpExists, podSpec.Tolerations[2].Operator)
+	assert.Equal(t, v1.TaintEffectNoSchedule, podSpec.Tolerations[2].Effect)
+}
diff --git a/flyteplugins/go/tasks/plugins/k8s/dask/dask_test.go b/flyteplugins/go/tasks/plugins/k8s/dask/dask_test.go
@@ -525,9 +525,10 @@ func TestBuildResouceDaskUsePodTemplate(t *testing.T) {
 
 func TestBuildResourceDaskExtendedResources(t *testing.T) {
 	assert.NoError(t, config.SetK8sPluginConfig(&config.K8sPluginConfig{
-		GpuDeviceNodeLabel:        "gpu-node-label",
-		GpuPartitionSizeNodeLabel: "gpu-partition-size",
-		GpuResourceName:           flytek8s.ResourceNvidiaGPU,
+		GpuDeviceNodeLabel:                 "gpu-node-label",
+		GpuPartitionSizeNodeLabel:          "gpu-partition-size",
+		GpuResourceName:                    flytek8s.ResourceNvidiaGPU,
+		AddTolerationsForExtendedResources: []string{"nvidia.com/gpu"},
 	}))
 
 	fixtures := []struct {
@@ -569,6 +570,11 @@ func TestBuildResourceDaskExtendedResources(t *testing.T) {
 					Operator: v1.TolerationOpEqual,
 					Effect:   v1.TaintEffectNoSchedule,
 				},
+				{
+					Key:      "nvidia.com/gpu",
+					Operator: v1.TolerationOpExists,
+					Effect:   v1.TaintEffectNoSchedule,
+				},
 			},
 		},
 		{
@@ -620,6 +626,11 @@ func TestBuildResourceDaskExtendedResources(t *testing.T) {
 					Operator: v1.TolerationOpEqual,
 					Effect:   v1.TaintEffectNoSchedule,
 				},
+				{
+					Key:      "nvidia.com/gpu",
+					Operator: v1.TolerationOpExists,
+					Effect:   v1.TaintEffectNoSchedule,
+				},
 			},
 		},
 	}

diff --git a/flyteplugins/go/tasks/plugins/k8s/kfoperators/mpi/mpi_test.go b/flyteplugins/go/tasks/plugins/k8s/kfoperators/mpi/mpi_test.go
@@ -368,9 +368,10 @@ func TestBuildResourceMPIForWrongInput(t *testing.T) {
 
 func TestBuildResourceMPIExtendedResources(t *testing.T) {
 	assert.NoError(t, flytek8sConfig.SetK8sPluginConfig(&flytek8sConfig.K8sPluginConfig{
-		GpuDeviceNodeLabel:        "gpu-node-label",
-		GpuPartitionSizeNodeLabel: "gpu-partition-size",
-		GpuResourceName:           flytek8s.ResourceNvidiaGPU,
+		GpuDeviceNodeLabel:                 "gpu-node-label",
+		GpuPartitionSizeNodeLabel:          "gpu-partition-size",
+		GpuResourceName:                    flytek8s.ResourceNvidiaGPU,
+		AddTolerationsForExtendedResources: []string{"nvidia.com/gpu"},
 	}))
 
 	fixtures := []struct {
@@ -412,6 +413,11 @@ func TestBuildResourceMPIExtendedResources(t *testing.T) {
 					Operator: corev1.TolerationOpEqual,
 					Effect:   corev1.TaintEffectNoSchedule,
 				},
+				{
+					Key:      "nvidia.com/gpu",
+					Operator: corev1.TolerationOpExists,
+					Effect:   corev1.TaintEffectNoSchedule,
+				},
 			},
 		},
 		{
@@ -463,6 +469,11 @@ func TestBuildResourceMPIExtendedResources(t *testing.T) {
 					Operator: corev1.TolerationOpEqual,
 					Effect:   corev1.TaintEffectNoSchedule,
 				},
+				{
+					Key:      "nvidia.com/gpu",
+					Operator: corev1.TolerationOpExists,
+					Effect:   corev1.TaintEffectNoSchedule,
+				},
 			},
 		},
 	}

diff --git a/flyteplugins/go/tasks/plugins/k8s/kfoperators/pytorch/pytorch_test.go b/flyteplugins/go/tasks/plugins/k8s/kfoperators/pytorch/pytorch_test.go
@@ -473,9 +473,10 @@ func TestBuildResourcePytorchContainerImage(t *testing.T) {
 
 func TestBuildResourcePytorchExtendedResources(t *testing.T) {
 	assert.NoError(t, flytek8sConfig.SetK8sPluginConfig(&flytek8sConfig.K8sPluginConfig{
-		GpuDeviceNodeLabel:        "gpu-node-label",
-		GpuPartitionSizeNodeLabel: "gpu-partition-size",
-		GpuResourceName:           flytek8s.ResourceNvidiaGPU,
+		GpuDeviceNodeLabel:                 "gpu-node-label",
+		GpuPartitionSizeNodeLabel:          "gpu-partition-size",
+		GpuResourceName:                    flytek8s.ResourceNvidiaGPU,
+		AddTolerationsForExtendedResources: []string{"nvidia.com/gpu"},
 	}))
 
 	fixtures := []struct {
@@ -517,6 +518,11 @@ func TestBuildResourcePytorchExtendedResources(t *testing.T) {
 					Operator: corev1.TolerationOpEqual,
 					Effect:   corev1.TaintEffectNoSchedule,
 				},
+				{
+					Key:      "nvidia.com/gpu",
+					Operator: corev1.TolerationOpExists,
+					Effect:   corev1.TaintEffectNoSchedule,
+				},
 			},
 		},
 		{
@@ -568,6 +574,11 @@ func TestBuildResourcePytorchExtendedResources(t *testing.T) {
 					Operator: corev1.TolerationOpEqual,
 					Effect:   corev1.TaintEffectNoSchedule,
 				},
+				{
+					Key:      "nvidia.com/gpu",
+					Operator: corev1.TolerationOpExists,
+					Effect:   corev1.TaintEffectNoSchedule,
+				},
 			},
 		},
 	}

diff --git a/flyteplugins/go/tasks/plugins/k8s/ray/ray.go b/flyteplugins/go/tasks/plugins/k8s/ray/ray.go
@@ -374,6 +374,8 @@ func buildHeadPodTemplate(primaryContainer *v1.Container, basePodSpec *v1.PodSpe
 		return v1.PodTemplateSpec{}, err
 	}
 
+	basePodSpec = flytek8s.AddTolerationsForExtendedResources(basePodSpec)
+
 	podTemplateSpec := v1.PodTemplateSpec{
 		Spec:       *basePodSpec,
 		ObjectMeta: *objectMeta,
@@ -502,6 +504,8 @@ func buildWorkerPodTemplate(primaryContainer *v1.Container, basePodSpec *v1.PodS
 		return v1.PodTemplateSpec{}, err
 	}
 
+	basePodSpec = flytek8s.AddTolerationsForExtendedResources(basePodSpec)
+
 	podTemplateSpec := v1.PodTemplateSpec{
 		Spec:       *basePodSpec,
 		ObjectMeta: *objectMetadata,