#minor Simplify and revisit task resource assignment (#215)

go/tasks/config_load_test.go

@@ -4,6 +4,8 @@ import (
 	"context"
 	"testing"
 
+	"k8s.io/apimachinery/pkg/api/resource"
+
 	sagemakerConfig "github.com/flyteorg/flyteplugins/go/tasks/plugins/k8s/sagemaker/config"
 
 	"github.com/flyteorg/flytestdlib/config"
@@ -68,8 +70,10 @@ func TestLoadConfig(t *testing.T) {
 
 		assert.Equal(t, []v1.Toleration{tolGPU}, k8sConfig.ResourceTolerations[v1.ResourceName("nvidia.com/gpu")])
 		assert.Equal(t, []v1.Toleration{tolStorage}, k8sConfig.ResourceTolerations[v1.ResourceStorage])
-		assert.Equal(t, "1000m", k8sConfig.DefaultCPURequest)
-		assert.Equal(t, "1024Mi", k8sConfig.DefaultMemoryRequest)
+		expectedCPU := resource.MustParse("1000m")
+		assert.True(t, expectedCPU.Equal(k8sConfig.DefaultCPURequest))
+		expectedMemory := resource.MustParse("1024Mi")
+		assert.True(t, expectedMemory.Equal(k8sConfig.DefaultMemoryRequest))
 		assert.Equal(t, map[string]string{"x/interruptible": "true"}, k8sConfig.InterruptibleNodeSelector)
 		assert.Equal(t, "x/flyte", k8sConfig.InterruptibleTolerations[0].Key)
 		assert.Equal(t, "interruptible", k8sConfig.InterruptibleTolerations[0].Value)

go/tasks/pluginmachinery/core/exec_metadata.go

@@ -35,4 +35,5 @@ type TaskExecutionMetadata interface {
 	GetK8sServiceAccount() string
 	GetSecurityContext() core.SecurityContext
 	IsInterruptible() bool
+	GetPlatformResources() *v1.ResourceRequirements
 }

go/tasks/pluginmachinery/flytek8s/config/config.go

@@ -8,6 +8,8 @@ package config
 import (
 	"time"
 
+	"k8s.io/apimachinery/pkg/api/resource"
+
 	config2 "github.com/flyteorg/flytestdlib/config"
 	v1 "k8s.io/api/core/v1"
 
@@ -17,8 +19,13 @@ import (
 //go:generate pflags K8sPluginConfig --default-var=defaultK8sConfig
 
 const k8sPluginConfigSectionKey = "k8s"
-const defaultCPURequest = "1000m"
-const defaultMemoryRequest = "1024Mi"
+
+// ResourceNvidiaGPU is the name of the Nvidia GPU resource.
+// Copied from: k8s.io/autoscaler/cluster-autoscaler/utils/gpu/gpu.go
+const ResourceNvidiaGPU v1.ResourceName = "nvidia.com/gpu"
+
+var defaultCPURequest = resource.MustParse("1000m")
+var defaultMemoryRequest = resource.MustParse("1024Mi")
 
 var (
 	defaultK8sConfig = K8sPluginConfig{
@@ -43,6 +50,7 @@ var (
 		CreateContainerErrorGracePeriod: config2.Duration{
 			Duration: time.Minute * 3,
 		},
+		GpuResourceName: ResourceNvidiaGPU,
 	}
 
 	// K8sPluginConfigSection provides a singular top level config section for all plugins.
@@ -69,9 +77,9 @@ type K8sPluginConfig struct {
 	DefaultEnvVarsFromEnv map[string]string `json:"default-env-vars-from-env" pflag:"-,Additional environment variable that should be injected into every resource"`
 
 	// default cpu requests for a container
-	DefaultCPURequest string `json:"default-cpus" pflag:",Defines a default value for cpu for containers if not specified."`
+	DefaultCPURequest resource.Quantity `json:"default-cpus" pflag:",Defines a default value for cpu for containers if not specified."`
 	// default memory requests for a container
-	DefaultMemoryRequest string `json:"default-memory" pflag:",Defines a default value for memory for containers if not specified."`
+	DefaultMemoryRequest resource.Quantity `json:"default-memory" pflag:",Defines a default value for memory for containers if not specified."`
 
 	// Default Tolerations that will be added to every Pod that is created by Flyte. These can be used in heterogenous clusters, where one wishes to keep all pods created by Flyte on a separate
 	// set of nodes.
@@ -118,6 +126,9 @@ type K8sPluginConfig struct {
 	// error persists past this grace period, it will be inferred to be a permanent
 	// one, and the corresponding task marked as failed
 	CreateContainerErrorGracePeriod config2.Duration `json:"create-container-error-grace-period" pflag:"-,Time to wait for transient CreateContainerError errors to be resolved."`
+
+	// The name of the GPU resource to use when the task resource requests GPUs.
+	GpuResourceName v1.ResourceName `json:"gpu-resource-name" pflag:",The name of the GPU resource to use when the task resource requests GPUs."`
 }
 
 type FlyteCoPilotConfig struct {