Fix custom gpu resource name specification

flyteorg · Dec 29, 2024 · cb764ad · cb764ad
1 parent 61838b4
commit cb764ad
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 1 deletion.
diff --git a/Makefile b/Makefile
@@ -18,6 +18,10 @@ cmd/single/dist: export FLYTECONSOLE_VERSION ?= latest
 cmd/single/dist:
 	script/get_flyteconsole_dist.sh
 
+.PHONY: run
+run: cmd/single/dist
+	POD_NAMESPACE=flyte go run -tags console cmd/main.go start --config ~/.flyte/flyte-single-binary-local.yaml
+
 .PHONY: compile
 compile: cmd/single/dist
 	go build -tags console -v -o flyte -ldflags=$(LD_FLAGS) ./cmd/

diff --git a/flyteplugins/go/tasks/pluginmachinery/flytek8s/utils.go b/flyteplugins/go/tasks/pluginmachinery/flytek8s/utils.go
@@ -7,6 +7,7 @@ import (
 
 	"github.com/flyteorg/flyte/flyteidl/gen/pb-go/flyteidl/core"
 	pluginmachinery_core "github.com/flyteorg/flyte/flyteplugins/go/tasks/pluginmachinery/core"
+	"github.com/flyteorg/flyte/flyteplugins/go/tasks/pluginmachinery/flytek8s/config"
 )
 
 func ToK8sEnvVar(env []*core.KeyValuePair) []v1.EnvVar {
@@ -20,6 +21,7 @@ func ToK8sEnvVar(env []*core.KeyValuePair) []v1.EnvVar {
 // TODO we should modify the container resources to contain a map of enum values?
 // Also we should probably create tolerations / taints, but we could do that as a post process
 func ToK8sResourceList(resources []*core.Resources_ResourceEntry) (v1.ResourceList, error) {
+	gpuResourceName := config.GetK8sPluginConfig().GpuResourceName
 	k8sResources := make(v1.ResourceList, len(resources))
 	for _, r := range resources {
 		rVal := r.GetValue()
@@ -38,7 +40,7 @@ func ToK8sResourceList(resources []*core.Resources_ResourceEntry) (v1.ResourceLi
 			}
 		case core.Resources_GPU:
 			if !v.IsZero() {
-				k8sResources[ResourceNvidiaGPU] = v
+				k8sResources[gpuResourceName] = v
 			}
 		case core.Resources_EPHEMERAL_STORAGE:
 			if !v.IsZero() {

diff --git a/flyteplugins/go/tasks/pluginmachinery/flytek8s/utils_test.go b/flyteplugins/go/tasks/pluginmachinery/flytek8s/utils_test.go
@@ -9,6 +9,7 @@ import (
 
 	"github.com/flyteorg/flyte/flyteidl/gen/pb-go/flyteidl/core"
 	"github.com/flyteorg/flyte/flyteplugins/go/tasks/pluginmachinery/core/mocks"
+	"github.com/flyteorg/flyte/flyteplugins/go/tasks/pluginmachinery/flytek8s/config"
 )
 
 func TestToK8sEnvVar(t *testing.T) {
@@ -44,6 +45,18 @@ func TestToK8sResourceList(t *testing.T) {
 		assert.Equal(t, resource.MustParse("1024Mi"), r[v1.ResourceMemory])
 		assert.Equal(t, resource.MustParse("1024Mi"), r[v1.ResourceEphemeralStorage])
 	}
+	{
+		gpuResourceName := v1.ResourceName("amd.com/gpu")
+		cfg := config.GetK8sPluginConfig()
+		cfg.GpuResourceName = gpuResourceName
+		assert.NoError(t, config.SetK8sPluginConfig(cfg))
+		r, err := ToK8sResourceList([]*core.Resources_ResourceEntry{
+			{Name: core.Resources_GPU, Value: "1"},
+		})
+		assert.NoError(t, err)
+		assert.NotEmpty(t, r)
+		assert.Equal(t, resource.MustParse("1"), r[gpuResourceName])
+	}
 	{
 		r, err := ToK8sResourceList([]*core.Resources_ResourceEntry{})
 		assert.NoError(t, err)