Add dask plugin #patch

go.mod

@@ -10,8 +10,9 @@ require (
 	github.com/aws/aws-sdk-go-v2 v1.2.0
 	github.com/aws/aws-sdk-go-v2/config v1.0.0
 	github.com/aws/aws-sdk-go-v2/service/athena v1.0.0
+	github.com/bstadlbauer/dask-k8s-operator-go-client v0.1.0
 	github.com/coocood/freecache v1.1.1
-	github.com/flyteorg/flyteidl v1.3.1
+	github.com/flyteorg/flyteidl v1.3.2
 	github.com/flyteorg/flytestdlib v1.0.11
 	github.com/go-test/deep v1.0.7
 	github.com/golang/protobuf v1.5.2
@@ -35,7 +36,7 @@ require (
 	k8s.io/api v0.24.1
 	k8s.io/apimachinery v0.24.1
 	k8s.io/client-go v0.24.1
-	k8s.io/utils v0.0.0-20220210201930-3a6ce19ff2f9
+	k8s.io/utils v0.0.0-20220728103510-ee6ede2d64ed
 	sigs.k8s.io/controller-runtime v0.12.1
 )
 
@@ -66,7 +67,7 @@ require (
 	github.com/cespare/xxhash v1.1.0 // indirect
 	github.com/cespare/xxhash/v2 v2.1.2 // indirect
 	github.com/davecgh/go-spew v1.1.1 // indirect
-	github.com/emicklei/go-restful v2.9.6+incompatible // indirect
+	github.com/emicklei/go-restful/v3 v3.8.0 // indirect
 	github.com/evanphx/json-patch v4.12.0+incompatible // indirect
 	github.com/fatih/color v1.13.0 // indirect
 	github.com/flyteorg/stow v0.3.6 // indirect
@@ -125,10 +126,10 @@ require (
 	gopkg.in/ini.v1 v1.66.4 // indirect
 	gopkg.in/yaml.v2 v2.4.0 // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect
-	k8s.io/klog/v2 v2.60.1 // indirect
-	k8s.io/kube-openapi v0.0.0-20220328201542-3ee0da9b0b42 // indirect
-	sigs.k8s.io/json v0.0.0-20220525155127-227cbc7cc124 // indirect
-	sigs.k8s.io/structured-merge-diff/v4 v4.2.1 // indirect
+	k8s.io/klog/v2 v2.70.1 // indirect
+	k8s.io/kube-openapi v0.0.0-20220803162953-67bda5d908f1 // indirect
+	sigs.k8s.io/json v0.0.0-20220713155537-f223a00ba0e2 // indirect
+	sigs.k8s.io/structured-merge-diff/v4 v4.2.3 // indirect
 	sigs.k8s.io/yaml v1.3.0 // indirect
 )
 

go.sum

@@ -188,6 +188,8 @@ github.com/beorn7/perks v1.0.1/go.mod h1:G2ZrVWU2WbWT9wwq4/hrbKbnv/1ERSJQ0ibhJ6r
 github.com/bgentry/speakeasy v0.1.0/go.mod h1:+zsyZBPWlz7T6j88CTgSN5bM796AkVf0kBD4zp0CCIs=
 github.com/bketelsen/crypt v0.0.3-0.20200106085610-5cbc8cc4026c/go.mod h1:MKsuJmJgSg28kpZDP6UIiPt0e0Oz0kqKNGyRaWEPv84=
 github.com/blang/semver v3.5.0+incompatible/go.mod h1:kRBLl5iJ+tD4TcOOxsy/0fnwebNt5EWlYSAyrTnjyyk=
+github.com/bstadlbauer/dask-k8s-operator-go-client v0.1.0 h1:PMUenya6FhDLW6WjWFdJ0l3uRj7eSxAcFfZ8EoEWgs0=
+github.com/bstadlbauer/dask-k8s-operator-go-client v0.1.0/go.mod h1:QPyKMRVI9NicWoMJqokH7eDGRqo7QR7Lu4931uxcS1Q=
 github.com/casbin/casbin/v2 v2.1.2/go.mod h1:YcPU1XXisHhLzuxH9coDNf2FbKpjGlbCg3n9yuLkIJQ=
 github.com/cenkalti/backoff v2.2.1+incompatible/go.mod h1:90ReRw6GdpyfrHakVjL/QHaoyV4aDUVVkXQJJJ3NXXM=
 github.com/census-instrumentation/opencensus-proto v0.2.1/go.mod h1:f6KPmirojxKA12rnyqOA5BBL4O983OfeGPqjHWSTneU=
@@ -256,8 +258,8 @@ github.com/edsrzf/mmap-go v1.0.0/go.mod h1:YO35OhQPt3KJa3ryjFM5Bs14WD66h8eGKpfaB
 github.com/elazarl/goproxy v0.0.0-20180725130230-947c36da3153/go.mod h1:/Zj4wYkgs4iZTTu3o/KG3Itv/qCCa8VVMlb3i9OVuzc=
 github.com/emicklei/go-restful v0.0.0-20170410110728-ff4f55a20633/go.mod h1:otzb+WCGbkyDHkqmQmT5YD2WR4BBwUdeQoFo8l/7tVs=
 github.com/emicklei/go-restful v2.9.5+incompatible/go.mod h1:otzb+WCGbkyDHkqmQmT5YD2WR4BBwUdeQoFo8l/7tVs=
-github.com/emicklei/go-restful v2.9.6+incompatible h1:tfrHha8zJ01ywiOEC1miGY8st1/igzWB8OmvPgoYX7w=
-github.com/emicklei/go-restful v2.9.6+incompatible/go.mod h1:otzb+WCGbkyDHkqmQmT5YD2WR4BBwUdeQoFo8l/7tVs=
+github.com/emicklei/go-restful/v3 v3.8.0 h1:eCZ8ulSerjdAiaNpF7GxXIE7ZCMo1moN1qX+S609eVw=
+github.com/emicklei/go-restful/v3 v3.8.0/go.mod h1:6n3XBCmQQb25CM2LCACGz8ukIrRry+4bhvbpWn3mrbc=
 github.com/envoyproxy/go-control-plane v0.6.9/go.mod h1:SBwIajubJHhxtWwsL9s8ss4safvEdbitLhGGK48rN6g=
 github.com/envoyproxy/go-control-plane v0.9.0/go.mod h1:YTl/9mNaCwkRvm6d1a2C3ymFceY/DCBVvsKhRF0iEA4=
 github.com/envoyproxy/go-control-plane v0.9.1-0.20191026205805-5f8ba28d4473/go.mod h1:YTl/9mNaCwkRvm6d1a2C3ymFceY/DCBVvsKhRF0iEA4=
@@ -281,8 +283,8 @@ github.com/fatih/color v1.10.0/go.mod h1:ELkj/draVOlAH/xkhN6mQ50Qd0MPOk5AAr3maGE
 github.com/fatih/color v1.13.0 h1:8LOYc1KYPPmyKMuN8QV2DNRWNbLo6LZ0iLs8+mlH53w=
 github.com/fatih/color v1.13.0/go.mod h1:kLAiJbzzSOZDVNGyDpeOxJ47H46qBXwg5ILebYFFOfk=
 github.com/fatih/structtag v1.2.0/go.mod h1:mBJUNpUnHmRKrKlQQlmCrh5PuhftFbNv8Ys4/aAZl94=
-github.com/flyteorg/flyteidl v1.3.1 h1:KIE7hXRmhBHttpyvfWo6X2LYLtx66Oj/nbKMG2V5bs0=
-github.com/flyteorg/flyteidl v1.3.1/go.mod h1:OJAq333OpInPnMhvVz93AlEjmlQ+t0FAD4aakIYE4OU=
+github.com/flyteorg/flyteidl v1.3.2 h1:s4DC8go2ou5LtZ+CFcS31r0mhv3baelNV81C1KZS26U=
+github.com/flyteorg/flyteidl v1.3.2/go.mod h1:OJAq333OpInPnMhvVz93AlEjmlQ+t0FAD4aakIYE4OU=
 github.com/flyteorg/flytestdlib v1.0.0/go.mod h1:QSVN5wIM1lM9d60eAEbX7NwweQXW96t5x4jbyftn89c=
 github.com/flyteorg/flytestdlib v1.0.11 h1:f7B8x2/zMuimEVi4Jx0zqzvNhdi7aq7+ZWoqHsbp4F4=
 github.com/flyteorg/flytestdlib v1.0.11/go.mod h1:nIBmBHtjTJvhZEn3e/EwVC/iMkR2tUX8hEiXjRBpH/s=
@@ -1519,19 +1521,22 @@ k8s.io/klog/v2 v2.0.0/go.mod h1:PBfzABfn139FHAV07az/IF9Wp1bkk3vpT2XSJ76fSDE=
 k8s.io/klog/v2 v2.2.0/go.mod h1:Od+F08eJP+W3HUb4pSrPpgp9DGU4GzlpG/TmITuYh/Y=
 k8s.io/klog/v2 v2.4.0/go.mod h1:Od+F08eJP+W3HUb4pSrPpgp9DGU4GzlpG/TmITuYh/Y=
 k8s.io/klog/v2 v2.5.0/go.mod h1:hy9LJ/NvuK+iVyP4Ehqva4HxZG/oXyIS3n3Jmire4Ec=
-k8s.io/klog/v2 v2.60.1 h1:VW25q3bZx9uE3vvdL6M8ezOX79vA2Aq1nEWLqNQclHc=
 k8s.io/klog/v2 v2.60.1/go.mod h1:y1WjHnz7Dj687irZUWR/WLkLc5N1YHtjLdmgWjndZn0=
+k8s.io/klog/v2 v2.70.1 h1:7aaoSdahviPmR+XkS7FyxlkkXs6tHISSG03RxleQAVQ=
+k8s.io/klog/v2 v2.70.1/go.mod h1:y1WjHnz7Dj687irZUWR/WLkLc5N1YHtjLdmgWjndZn0=
 k8s.io/kube-openapi v0.0.0-20200121204235-bf4fb3bd569c/go.mod h1:GRQhZsXIAJ1xR0C9bd8UpWHZ5plfAS9fzPjJuQ6JL3E=
 k8s.io/kube-openapi v0.0.0-20200410145947-61e04a5be9a6/go.mod h1:GRQhZsXIAJ1xR0C9bd8UpWHZ5plfAS9fzPjJuQ6JL3E=
 k8s.io/kube-openapi v0.0.0-20201113171705-d219536bb9fd/go.mod h1:WOJ3KddDSol4tAGcJo0Tvi+dK12EcqSLqcWsryKMpfM=
-k8s.io/kube-openapi v0.0.0-20220328201542-3ee0da9b0b42 h1:Gii5eqf+GmIEwGNKQYQClCayuJCe2/4fZUvF7VG99sU=
 k8s.io/kube-openapi v0.0.0-20220328201542-3ee0da9b0b42/go.mod h1:Z/45zLw8lUo4wdiUkI+v/ImEGAvu3WatcZl3lPMR4Rk=
+k8s.io/kube-openapi v0.0.0-20220803162953-67bda5d908f1 h1:MQ8BAZPZlWk3S9K4a9NCkIFQtZShWqoha7snGixVgEA=
+k8s.io/kube-openapi v0.0.0-20220803162953-67bda5d908f1/go.mod h1:C/N6wCaBHeBHkHUesQOQy2/MZqGgMAFPqGsGQLdbZBU=
 k8s.io/utils v0.0.0-20200324210504-a9aa75ae1b89/go.mod h1:sZAwmy6armz5eXlNoLmJcl4F1QuKu7sr+mFQ0byX7Ew=
 k8s.io/utils v0.0.0-20200603063816-c1c6865ac451/go.mod h1:jPW/WVKK9YHAvNhRxK0md/EJ228hCsBRufyofKtW8HA=
 k8s.io/utils v0.0.0-20201110183641-67b214c5f920/go.mod h1:jPW/WVKK9YHAvNhRxK0md/EJ228hCsBRufyofKtW8HA=
 k8s.io/utils v0.0.0-20210802155522-efc7438f0176/go.mod h1:jPW/WVKK9YHAvNhRxK0md/EJ228hCsBRufyofKtW8HA=
-k8s.io/utils v0.0.0-20220210201930-3a6ce19ff2f9 h1:HNSDgDCrr/6Ly3WEGKZftiE7IY19Vz2GdbOCyI4qqhc=
 k8s.io/utils v0.0.0-20220210201930-3a6ce19ff2f9/go.mod h1:jPW/WVKK9YHAvNhRxK0md/EJ228hCsBRufyofKtW8HA=
+k8s.io/utils v0.0.0-20220728103510-ee6ede2d64ed h1:jAne/RjBTyawwAy0utX5eqigAwz/lQhTmy+Hr/Cpue4=
+k8s.io/utils v0.0.0-20220728103510-ee6ede2d64ed/go.mod h1:jPW/WVKK9YHAvNhRxK0md/EJ228hCsBRufyofKtW8HA=
 rsc.io/binaryregexp v0.2.0/go.mod h1:qTv7/COck+e2FymRvadv62gMdZztPaShugOCi3I+8D8=
 rsc.io/quote/v3 v3.1.0/go.mod h1:yEA65RcK8LyAZtP9Kv3t0HmxON59tX3rD+tICJqUlj0=
 rsc.io/sampler v1.3.0/go.mod h1:T1hPZKmBbMNahiBKFy5HrXp6adAjACjK9JXDnKaTXpA=
@@ -1541,14 +1546,15 @@ sigs.k8s.io/controller-runtime v0.12.1 h1:4BJY01xe9zKQti8oRjj/NeHKRXthf1YkYJAgLO
 sigs.k8s.io/controller-runtime v0.12.1/go.mod h1:BKhxlA4l7FPK4AQcsuL4X6vZeWnKDXez/vp1Y8dxTU0=
 sigs.k8s.io/controller-tools v0.3.0/go.mod h1:enhtKGfxZD1GFEoMgP8Fdbu+uKQ/cq1/WGJhdVChfvI=
 sigs.k8s.io/json v0.0.0-20211208200746-9f7c6b3444d2/go.mod h1:B+TnT182UBxE84DiCz4CVE26eOSDAeYCpfDnC2kdKMY=
-sigs.k8s.io/json v0.0.0-20220525155127-227cbc7cc124 h1:2sgAQQcY0dEW2SsQwTXhQV4vO6+rSslYx8K3XmM5hqQ=
-sigs.k8s.io/json v0.0.0-20220525155127-227cbc7cc124/go.mod h1:B+TnT182UBxE84DiCz4CVE26eOSDAeYCpfDnC2kdKMY=
+sigs.k8s.io/json v0.0.0-20220713155537-f223a00ba0e2 h1:iXTIw73aPyC+oRdyqqvVJuloN1p0AC/kzH07hu3NE+k=
+sigs.k8s.io/json v0.0.0-20220713155537-f223a00ba0e2/go.mod h1:B8JuhiUyNFVKdsE8h686QcCxMaH6HrOAZj4vswFpcB0=
 sigs.k8s.io/structured-merge-diff/v3 v3.0.0-20200116222232-67a7b8c61874/go.mod h1:PlARxl6Hbt/+BC80dRLi1qAmnMqwqDg62YvvVkZjemw=
 sigs.k8s.io/structured-merge-diff/v3 v3.0.0/go.mod h1:PlARxl6Hbt/+BC80dRLi1qAmnMqwqDg62YvvVkZjemw=
 sigs.k8s.io/structured-merge-diff/v4 v4.0.2/go.mod h1:bJZC9H9iH24zzfZ/41RGcq60oK1F7G282QMXDPYydCw=
 sigs.k8s.io/structured-merge-diff/v4 v4.0.3/go.mod h1:bJZC9H9iH24zzfZ/41RGcq60oK1F7G282QMXDPYydCw=
-sigs.k8s.io/structured-merge-diff/v4 v4.2.1 h1:bKCqE9GvQ5tiVHn5rfn1r+yao3aLQEaLzkkmAkf+A6Y=
 sigs.k8s.io/structured-merge-diff/v4 v4.2.1/go.mod h1:j/nl6xW8vLS49O8YvXW1ocPhZawJtm+Yrr7PPRQ0Vg4=
+sigs.k8s.io/structured-merge-diff/v4 v4.2.3 h1:PRbqxJClWWYMNV1dhaG4NsibJbArud9kFxnAMREiWFE=
+sigs.k8s.io/structured-merge-diff/v4 v4.2.3/go.mod h1:qjx8mGObPmV2aSZepjQjbmb2ihdVs8cGKBraizNC69E=
 sigs.k8s.io/yaml v1.1.0/go.mod h1:UJmg0vDUVViEyp3mgSv9WPwZCDxu4rQW1olrI1uml+o=
 sigs.k8s.io/yaml v1.2.0/go.mod h1:yfXDCHCao9+ENCvLSE62v9VSji2MKu5jeNfTrofGhJc=
 sigs.k8s.io/yaml v1.3.0 h1:a2VclLzOGrwOHDiV8EfBGhvjHvP46CtW5j6POvhYGGo=

go/tasks/pluginmachinery/flytek8s/utils.go

@@ -3,7 +3,9 @@ package flytek8s
 import (
 	"github.com/flyteorg/flyteidl/gen/pb-go/flyteidl/core"
 	pluginmachinery_core "github.com/flyteorg/flyteplugins/go/tasks/pluginmachinery/core"
+	"github.com/pkg/errors"
 	v1 "k8s.io/api/core/v1"
+	"k8s.io/apimachinery/pkg/api/resource"
 )
 
 func ToK8sEnvVar(env []*core.KeyValuePair) []v1.EnvVar {
@@ -14,6 +16,60 @@ func ToK8sEnvVar(env []*core.KeyValuePair) []v1.EnvVar {
 	return envVars
 }
 
+// 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) {
+	k8sResources := make(v1.ResourceList, len(resources))
+	for _, r := range resources {
+		rVal := r.Value
+		v, err := resource.ParseQuantity(rVal)
+		if err != nil {
+			return nil, errors.Wrap(err, "Failed to parse resource as a valid quantity.")
+		}
+		switch r.Name {
+		case core.Resources_CPU:
+			if !v.IsZero() {
+				k8sResources[v1.ResourceCPU] = v
+			}
+		case core.Resources_MEMORY:
+			if !v.IsZero() {
+				k8sResources[v1.ResourceMemory] = v
+			}
+		case core.Resources_STORAGE:
+			if !v.IsZero() {
+				k8sResources[v1.ResourceStorage] = v
+			}
+		case core.Resources_GPU:
+			if !v.IsZero() {
+				k8sResources[ResourceNvidiaGPU] = v
+			}
+		case core.Resources_EPHEMERAL_STORAGE:
+			if !v.IsZero() {
+				k8sResources[v1.ResourceEphemeralStorage] = v
+			}
+		}
+	}
+	return k8sResources, nil
+}
+
+func ToK8sResourceRequirements(resources *core.Resources) (*v1.ResourceRequirements, error) {
+	res := &v1.ResourceRequirements{}
+	if resources == nil {
+		return res, nil
+	}
+	req, err := ToK8sResourceList(resources.Requests)
+	if err != nil {
+		return res, err
+	}
+	lim, err := ToK8sResourceList(resources.Limits)
+	if err != nil {
+		return res, err
+	}
+	res.Limits = lim
+	res.Requests = req
+	return res, nil
+}
+
 func GetServiceAccountNameFromTaskExecutionMetadata(taskExecutionMetadata pluginmachinery_core.TaskExecutionMetadata) string {
 	var serviceAccount string
 	securityContext := taskExecutionMetadata.GetSecurityContext()

go/tasks/pluginmachinery/flytek8s/utils_test.go

@@ -5,10 +5,118 @@ import (
 
 	"github.com/flyteorg/flyteidl/gen/pb-go/flyteidl/core"
 	"github.com/flyteorg/flyteplugins/go/tasks/pluginmachinery/core/mocks"
+	v1 "k8s.io/api/core/v1"
+	"k8s.io/apimachinery/pkg/api/resource"
 
 	"github.com/stretchr/testify/assert"
 )
 
+func TestToK8sEnvVar(t *testing.T) {
+	e := ToK8sEnvVar([]*core.KeyValuePair{
+		{Key: "k1", Value: "v1"},
+		{Key: "k2", Value: "v2"},
+	})
+
+	assert.NotEmpty(t, e)
+	assert.Equal(t, []v1.EnvVar{
+		{Name: "k1", Value: "v1"},
+		{Name: "k2", Value: "v2"},
+	}, e)
+
+	e = ToK8sEnvVar(nil)
+	assert.Empty(t, e)
+}
+
+func TestToK8sResourceList(t *testing.T) {
+	{
+		r, err := ToK8sResourceList([]*core.Resources_ResourceEntry{
+			{Name: core.Resources_CPU, Value: "250m"},
+			{Name: core.Resources_GPU, Value: "1"},
+			{Name: core.Resources_MEMORY, Value: "1024Mi"},
+			{Name: core.Resources_STORAGE, Value: "1024Mi"},
+			{Name: core.Resources_EPHEMERAL_STORAGE, Value: "1024Mi"},
+		})
+
+		assert.NoError(t, err)
+		assert.NotEmpty(t, r)
+		assert.NotNil(t, r[v1.ResourceCPU])
+		assert.Equal(t, resource.MustParse("250m"), r[v1.ResourceCPU])
+		assert.Equal(t, resource.MustParse("1"), r[ResourceNvidiaGPU])
+		assert.Equal(t, resource.MustParse("1024Mi"), r[v1.ResourceMemory])
+		assert.Equal(t, resource.MustParse("1024Mi"), r[v1.ResourceStorage])
+		assert.Equal(t, resource.MustParse("1024Mi"), r[v1.ResourceEphemeralStorage])
+	}
+	{
+		r, err := ToK8sResourceList([]*core.Resources_ResourceEntry{})
+		assert.NoError(t, err)
+		assert.Empty(t, r)
+	}
+	{
+		_, err := ToK8sResourceList([]*core.Resources_ResourceEntry{
+			{Name: core.Resources_CPU, Value: "250x"},
+		})
+		assert.Error(t, err)
+	}
+
+}
+
+func TestToK8sResourceRequirements(t *testing.T) {
+
+	{
+		r, err := ToK8sResourceRequirements(nil)
+		assert.NoError(t, err)
+		assert.NotNil(t, r)
+		assert.Empty(t, r.Limits)
+		assert.Empty(t, r.Requests)
+	}
+	{
+		r, err := ToK8sResourceRequirements(&core.Resources{
+			Requests: nil,
+			Limits:   nil,
+		})
+		assert.NoError(t, err)
+		assert.NotNil(t, r)
+		assert.Empty(t, r.Limits)
+		assert.Empty(t, r.Requests)
+	}
+	{
+		r, err := ToK8sResourceRequirements(&core.Resources{
+			Requests: []*core.Resources_ResourceEntry{
+				{Name: core.Resources_CPU, Value: "250m"},
+			},
+			Limits: []*core.Resources_ResourceEntry{
+				{Name: core.Resources_CPU, Value: "1024m"},
+			},
+		})
+		assert.NoError(t, err)
+		assert.NotNil(t, r)
+		assert.Equal(t, resource.MustParse("250m"), r.Requests[v1.ResourceCPU])
+		assert.Equal(t, resource.MustParse("1024m"), r.Limits[v1.ResourceCPU])
+	}
+	{
+		_, err := ToK8sResourceRequirements(&core.Resources{
+			Requests: []*core.Resources_ResourceEntry{
+				{Name: core.Resources_CPU, Value: "blah"},
+			},
+			Limits: []*core.Resources_ResourceEntry{
+				{Name: core.Resources_CPU, Value: "1024m"},
+			},
+		})
+		assert.Error(t, err)
+	}
+	{
+		_, err := ToK8sResourceRequirements(&core.Resources{
+			Requests: []*core.Resources_ResourceEntry{
+				{Name: core.Resources_CPU, Value: "250m"},
+			},
+			Limits: []*core.Resources_ResourceEntry{
+				{Name: core.Resources_CPU, Value: "blah"},
+			},
+		})
+		assert.Error(t, err)
+	}
+}
+
 func TestGetServiceAccountNameFromTaskExecutionMetadata(t *testing.T) {
 	mockTaskExecMetadata := mocks.TaskExecutionMetadata{}
 	mockTaskExecMetadata.OnGetSecurityContext().Return(core.SecurityContext{