From f0db51179f215546beac2e70fe6e8e94b92fcd4e Mon Sep 17 00:00:00 2001 From: Dan Rammer Date: Mon, 13 Mar 2023 13:00:49 -0500 Subject: [PATCH] Supporting DataLoadingConfig on K8sPod type (#325) * supporting data loading config for k8s pod type Signed-off-by: Daniel Rammer * should be checking dataLoadingConfig for nil Signed-off-by: Daniel Rammer * updated flyteidl Signed-off-by: Daniel Rammer --------- Signed-off-by: Daniel Rammer --- go.mod | 2 +- go.sum | 4 ++-- go/tasks/pluginmachinery/flytek8s/pod_helper.go | 15 +++++++++++---- 3 files changed, 14 insertions(+), 7 deletions(-) diff --git a/go.mod b/go.mod index 48e906f323..e06d9fa50b 100644 --- a/go.mod +++ b/go.mod @@ -12,7 +12,7 @@ require ( 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.6 + github.com/flyteorg/flyteidl v1.3.12 github.com/flyteorg/flytestdlib v1.0.15 github.com/go-test/deep v1.0.7 github.com/golang/protobuf v1.5.2 diff --git a/go.sum b/go.sum index 559ec0ef58..209c816b55 100644 --- a/go.sum +++ b/go.sum @@ -232,8 +232,8 @@ github.com/evanphx/json-patch v4.12.0+incompatible/go.mod h1:50XU6AFN0ol/bzJsmQL github.com/fatih/color v1.7.0/go.mod h1:Zm6kSWBoL9eyXnKyktHP6abPY2pDugNf5KwzbycvMj4= github.com/fatih/color v1.13.0 h1:8LOYc1KYPPmyKMuN8QV2DNRWNbLo6LZ0iLs8+mlH53w= github.com/fatih/color v1.13.0/go.mod h1:kLAiJbzzSOZDVNGyDpeOxJ47H46qBXwg5ILebYFFOfk= -github.com/flyteorg/flyteidl v1.3.6 h1:PI846AdnrQZ84pxRVAzA3WGihv+xXmjQHO91nj/kV9g= -github.com/flyteorg/flyteidl v1.3.6/go.mod h1:Pkt2skI1LiHs/2ZoekBnyPhuGOFMiuul6HHcKGZBsbM= +github.com/flyteorg/flyteidl v1.3.12 h1:RTcxCrqKU235cWuy+j3gkmqPJOaaYEcJaT6fsRjoS8Q= +github.com/flyteorg/flyteidl v1.3.12/go.mod h1:Pkt2skI1LiHs/2ZoekBnyPhuGOFMiuul6HHcKGZBsbM= github.com/flyteorg/flytestdlib v1.0.15 h1:kv9jDQmytbE84caY+pkZN8trJU2ouSAmESzpTEhfTt0= github.com/flyteorg/flytestdlib v1.0.15/go.mod h1:ghw/cjY0sEWIIbyCtcJnL/Gt7ZS7gf9SUi0CCPhbz3s= github.com/flyteorg/stow v0.3.6 h1:jt50ciM14qhKBaIrB+ppXXY+SXB59FNREFgTJqCyqIk= diff --git a/go/tasks/pluginmachinery/flytek8s/pod_helper.go b/go/tasks/pluginmachinery/flytek8s/pod_helper.go index bd49e3c1b8..6297b07c37 100755 --- a/go/tasks/pluginmachinery/flytek8s/pod_helper.go +++ b/go/tasks/pluginmachinery/flytek8s/pod_helper.go @@ -219,14 +219,21 @@ func ApplyFlytePodConfiguration(ctx context.Context, tCtx pluginsCore.TaskExecut } // add copilot configuration to primaryContainer and PodSpec (if necessary) - if taskTemplate.GetContainer() != nil { - if err := AddCoPilotToContainer(ctx, config.GetK8sPluginConfig().CoPilot, primaryContainer, - taskTemplate.Interface, taskTemplate.GetContainer().DataConfig); err != nil { + var dataLoadingConfig *core.DataLoadingConfig + if container := taskTemplate.GetContainer(); container != nil { + dataLoadingConfig = container.GetDataConfig() + } else if pod := taskTemplate.GetK8SPod(); pod != nil { + dataLoadingConfig = pod.GetDataConfig() + } + + if dataLoadingConfig != nil { + if err := AddCoPilotToContainer(ctx, config.GetK8sPluginConfig().CoPilot, + primaryContainer, taskTemplate.Interface, dataLoadingConfig); err != nil { return nil, nil, err } if err := AddCoPilotToPod(ctx, config.GetK8sPluginConfig().CoPilot, podSpec, taskTemplate.GetInterface(), - tCtx.TaskExecutionMetadata(), tCtx.InputReader(), tCtx.OutputWriter(), taskTemplate.GetContainer().GetDataConfig()); err != nil { + tCtx.TaskExecutionMetadata(), tCtx.InputReader(), tCtx.OutputWriter(), dataLoadingConfig); err != nil { return nil, nil, err } }