merged master

Signed-off-by: Kevin Su <[email protected]>

.github/workflows/master.yml

@@ -34,7 +34,7 @@ jobs:
       - name: Set up Go
         uses: actions/setup-go@v2
         with:
-          go-version: '1.18'
+          go-version: '1.19'
       - name: Run GoReleaser
         uses: goreleaser/goreleaser-action@v2
         with:
@@ -52,7 +52,7 @@ jobs:
       - name: Set up Go
         uses: actions/setup-go@v2
         with:
-          go-version: '1.18'
+          go-version: '1.19'
       - name: Unit Tests
         run: make install && make test_unit_codecov
       - name: Push CodeCov

.github/workflows/pull_request.yml

@@ -12,7 +12,7 @@ jobs:
       - name: Set up Go
         uses: actions/setup-go@v2
         with:
-          go-version: '1.18'
+          go-version: '1.19'
       - name: Unit Tests
         run: make install && make test_unit_codecov
       - name: Push CodeCov
@@ -30,6 +30,6 @@ jobs:
       - uses: actions/checkout@v1
       - uses: actions/setup-go@v2
         with:
-          go-version: '1.18'
+          go-version: '1.19'
       - name: Go generate and diff
         run: DELTA_CHECK=true make generate

go.mod

@@ -1,6 +1,6 @@
 module github.com/flyteorg/flyteplugins
 
-go 1.18
+go 1.19
 
 require (
 	github.com/GoogleCloudPlatform/spark-on-k8s-operator v0.0.0-20200723154620-6f35a1152625
@@ -12,8 +12,8 @@ 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.2
-	github.com/flyteorg/flytestdlib v1.0.11
+	github.com/flyteorg/flyteidl v1.3.6
+	github.com/flyteorg/flytestdlib v1.0.15
 	github.com/go-test/deep v1.0.7
 	github.com/golang/protobuf v1.5.2
 	github.com/hashicorp/golang-lru v0.5.4
@@ -136,6 +136,4 @@ require (
 
 replace github.com/aws/amazon-sagemaker-operator-for-k8s => github.com/aws/amazon-sagemaker-operator-for-k8s v1.0.1-0.20210303003444-0fb33b1fd49d
 
-replace github.com/flyteorg/flyteplugins => github.com/flyteorg/flyteplugins v1.0.29-0.20230224022650-a2881f35dd25
-
 replace github.com/flyteorg/flyteidl => github.com/flyteorg/flyteidl v1.3.9-0.20230224194627-a1df35060476

go.sum

@@ -234,8 +234,8 @@ 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.9-0.20230224194627-a1df35060476 h1:mA3Ry5YjNu5BqjnCTbA+lFRTRFjGKEMDALRhLTtBuuU=
 github.com/flyteorg/flyteidl v1.3.9-0.20230224194627-a1df35060476/go.mod h1:Pkt2skI1LiHs/2ZoekBnyPhuGOFMiuul6HHcKGZBsbM=
-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=
+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=
 github.com/flyteorg/stow v0.3.6/go.mod h1:5dfBitPM004dwaZdoVylVjxFT4GWAgI0ghAndhNUzCo=
 github.com/form3tech-oss/jwt-go v3.2.2+incompatible/go.mod h1:pbq4aXjuKjdthFRnoDwaVPLA+WlJuPGy+QneDUgJi2k=

go/tasks/pluginmachinery/core/resource_manager.go

@@ -37,18 +37,20 @@ type ResourceRegistrar interface {
 
 // ResourceManager Interface
 // 1. Terms and definitions
+//
 //   - Resource: resource is an abstraction of anything that has a limited quota of units and can be claimed in a
 //     single unit or multiple units at once. At Flyte's current state, a resource means a logical
 //     separation (e.g., a cluster) of an external service that allows a limited number of outstanding
 //     requests to be sent to.
+//
 //   - Token: Flyte uses a token to serve as the placeholder to represent a unit of resource. Flyte resource manager
 //     manages resources by managing the tokens of the resources.
 //
-//  2. Description
+//     2. Description
 //     ResourceManager provides a task-type-specific pooling system for Flyte Tasks. Plugin writers can optionally
 //     request for resources in their tasks, in single quantity.
 //
-//  3. Usage
+//     3. Usage
 //     A Flyte plugin registers the resources and the desired quota of each resource with ResourceRegistrar at the
 //     setup time of Flyte Propeller. At the end of the setup time, Flyte Propeller builds a ResourceManager based on
 //     these registration requests.
@@ -63,7 +65,7 @@ type ResourceRegistrar interface {
 //     the resource manager to release the token by calling ResourceManager's ReleaseResource(), and the token will be
 //     erased from the corresponding pool.
 //
-//  4. Example
+//     4. Example
 //     Flyte has a built-on Qubole plugin that allows Flyte tasks to send out Hive commands to Qubole.
 //     In the plugin, a single Qubole cluster is a resource, and sending out a single Hive command to a Qubole cluster consumes
 //     a token of the corresponding resource. The resource allocation is achieved by the Qubole plugin calling

go/tasks/pluginmachinery/flytek8s/container_helper.go

@@ -3,18 +3,19 @@ package flytek8s
 import (
 	"context"
 
-	"github.com/flyteorg/flytestdlib/logger"
+	"github.com/flyteorg/flyteidl/gen/pb-go/flyteidl/core"
 
+	"github.com/flyteorg/flyteplugins/go/tasks/errors"
+	pluginscore "github.com/flyteorg/flyteplugins/go/tasks/pluginmachinery/core"
 	"github.com/flyteorg/flyteplugins/go/tasks/pluginmachinery/core/template"
-	"k8s.io/apimachinery/pkg/util/validation"
+	"github.com/flyteorg/flyteplugins/go/tasks/pluginmachinery/flytek8s/config"
+
+	"github.com/flyteorg/flytestdlib/logger"
 
-	"github.com/flyteorg/flyteidl/gen/pb-go/flyteidl/core"
 	v1 "k8s.io/api/core/v1"
 	"k8s.io/apimachinery/pkg/api/resource"
 	"k8s.io/apimachinery/pkg/util/rand"
-
-	"github.com/flyteorg/flyteplugins/go/tasks/errors"
-	"github.com/flyteorg/flyteplugins/go/tasks/pluginmachinery/flytek8s/config"
+	"k8s.io/apimachinery/pkg/util/validation"
 )
 
 const resourceGPU = "gpu"
@@ -193,22 +194,16 @@ func ApplyResourceOverrides(resources, platformResources v1.ResourceRequirements
 	return resources
 }
 
-// ToK8sContainer transforms a task template target of type core.Container into a bare-bones kubernetes container, which
-// can be further modified with flyte-specific customizations specified by various static and run-time attributes.
-func ToK8sContainer(ctx context.Context, taskContainer *core.Container, iFace *core.TypedInterface, parameters template.Parameters) (*v1.Container, error) {
-	// Perform preliminary validations
-	if parameters.TaskExecMetadata.GetOverrides() == nil {
-		return nil, errors.Errorf(errors.BadTaskSpecification, "platform/compiler error, overrides not set for task")
-	}
-	if parameters.TaskExecMetadata.GetOverrides() == nil || parameters.TaskExecMetadata.GetOverrides().GetResources() == nil {
-		return nil, errors.Errorf(errors.BadTaskSpecification, "resource requirements not found for container task, required!")
-	}
+// BuildRawContainer constructs a Container based on the definition passed by the taskContainer and
+// TaskExecutionMetadata.
+func BuildRawContainer(ctx context.Context, taskContainer *core.Container, taskExecMetadata pluginscore.TaskExecutionMetadata) (*v1.Container, error) {
 	// Make the container name the same as the pod name, unless it violates K8s naming conventions
 	// Container names are subject to the DNS-1123 standard
-	containerName := parameters.TaskExecMetadata.GetTaskExecutionID().GetGeneratedName()
+	containerName := taskExecMetadata.GetTaskExecutionID().GetGeneratedName()
 	if errs := validation.IsDNS1123Label(containerName); len(errs) > 0 {
 		containerName = rand.String(4)
 	}
+
 	container := &v1.Container{
 		Name:                     containerName,
 		Image:                    taskContainer.GetImage(),
@@ -217,12 +212,49 @@ func ToK8sContainer(ctx context.Context, taskContainer *core.Container, iFace *c
 		Env:                      ToK8sEnvVar(taskContainer.GetEnv()),
 		TerminationMessagePolicy: v1.TerminationMessageFallbackToLogsOnError,
 	}
-	if err := AddCoPilotToContainer(ctx, config.GetK8sPluginConfig().CoPilot, container, iFace, taskContainer.DataConfig); err != nil {
+
+	return container, nil
+}
+
+// ToK8sContainer builds a Container based on the definition passed by the TaskExecutionContext. This involves applying
+// all Flyte configuration including k8s plugins and resource requests.
+func ToK8sContainer(ctx context.Context, tCtx pluginscore.TaskExecutionContext) (*v1.Container, error) {
+	taskTemplate, err := tCtx.TaskReader().Read(ctx)
+	if err != nil {
+		logger.Warnf(ctx, "failed to read task information when trying to construct container, err: %s", err.Error())
+		return nil, err
+	}
+
+	// validate arguments
+	if taskTemplate.GetContainer() == nil {
+		return nil, errors.Errorf(errors.BadTaskSpecification, "unable to create container with no definition in TaskTemplate")
+	}
+	if tCtx.TaskExecutionMetadata().GetOverrides() == nil || tCtx.TaskExecutionMetadata().GetOverrides().GetResources() == nil {
+		return nil, errors.Errorf(errors.BadTaskSpecification, "resource requirements not found for container task, required!")
+	}
+
+	// build raw container
+	container, err := BuildRawContainer(ctx, taskTemplate.GetContainer(), tCtx.TaskExecutionMetadata())
+	if err != nil {
 		return nil, err
 	}
+
 	if container.SecurityContext == nil && config.GetK8sPluginConfig().DefaultSecurityContext != nil {
 		container.SecurityContext = config.GetK8sPluginConfig().DefaultSecurityContext.DeepCopy()
 	}
+
+	// add flyte resource customizations to the container
+	templateParameters := template.Parameters{
+		TaskExecMetadata: tCtx.TaskExecutionMetadata(),
+		Inputs:           tCtx.InputReader(),
+		OutputPath:       tCtx.OutputWriter(),
+		Task:             tCtx.TaskReader(),
+	}
+
+	if err := AddFlyteCustomizationsToContainer(ctx, templateParameters, ResourceCustomizationModeAssignResources, container); err != nil {
+		return nil, err
+	}
+
 	return container, nil
 }
 

go/tasks/pluginmachinery/flytek8s/container_helper_test.go

@@ -335,26 +335,45 @@ func TestMergeResources_PartialResourceKeys(t *testing.T) {
 }
 
 func TestToK8sContainer(t *testing.T) {
-	taskContainer := &core.Container{
-		Image: "myimage",
-		Args: []string{
-			"arg1",
-			"arg2",
-			"arg3",
-		},
-		Command: []string{
-			"com1",
-			"com2",
-			"com3",
-		},
-		Env: []*core.KeyValuePair{
-			{
-				Key:   "k",
-				Value: "v",
+	taskTemplate := &core.TaskTemplate{
+		Type: "test",
+		Target: &core.TaskTemplate_Container{
+			Container: &core.Container{
+				Image: "myimage",
+				Args: []string{
+					"arg1",
+					"arg2",
+					"arg3",
+				},
+				Command: []string{
+					"com1",
+					"com2",
+					"com3",
+				},
+				Env: []*core.KeyValuePair{
+					{
+						Key:   "k",
+						Value: "v",
+					},
+				},
 			},
 		},
 	}
 
+	taskReader := &mocks.TaskReader{}
+	taskReader.On("Read", mock.Anything).Return(taskTemplate, nil)
+
+	inputReader := &mocks2.InputReader{}
+	inputReader.OnGetInputPath().Return(storage.DataReference("test-data-reference"))
+	inputReader.OnGetInputPrefixPath().Return(storage.DataReference("test-data-reference-prefix"))
+	inputReader.OnGetMatch(mock.Anything).Return(&core.LiteralMap{}, nil)
+
+	outputWriter := &mocks2.OutputWriter{}
+	outputWriter.OnGetOutputPrefixPath().Return("")
+	outputWriter.OnGetRawOutputPrefix().Return("")
+	outputWriter.OnGetCheckpointPrefix().Return("/checkpoint")
+	outputWriter.OnGetPreviousCheckpointsPrefix().Return("/prev")
+
 	mockTaskExecMetadata := mocks.TaskExecutionMetadata{}
 	mockTaskOverrides := mocks.TaskOverrides{}
 	mockTaskOverrides.OnGetResources().Return(&v1.ResourceRequirements{
@@ -364,12 +383,16 @@ func TestToK8sContainer(t *testing.T) {
 	})
 	mockTaskExecMetadata.OnGetOverrides().Return(&mockTaskOverrides)
 	mockTaskExecutionID := mocks.TaskExecutionID{}
+	mockTaskExecutionID.OnGetID().Return(core.TaskExecutionIdentifier{})
 	mockTaskExecutionID.OnGetGeneratedName().Return("gen_name")
 	mockTaskExecMetadata.OnGetTaskExecutionID().Return(&mockTaskExecutionID)
+	mockTaskExecMetadata.OnGetPlatformResources().Return(&v1.ResourceRequirements{})
 
-	templateParameters := template.Parameters{
-		TaskExecMetadata: &mockTaskExecMetadata,
-	}
+	tCtx := &mocks.TaskExecutionContext{}
+	tCtx.OnTaskExecutionMetadata().Return(&mockTaskExecMetadata)
+	tCtx.OnInputReader().Return(inputReader)
+	tCtx.OnTaskReader().Return(taskReader)
+	tCtx.OnOutputWriter().Return(outputWriter)
 
 	cfg := config.GetK8sPluginConfig()
 	allow := false
@@ -378,7 +401,7 @@ func TestToK8sContainer(t *testing.T) {
 	}
 	assert.NoError(t, config.SetK8sPluginConfig(cfg))
 
-	container, err := ToK8sContainer(context.TODO(), taskContainer, nil, templateParameters)
+	container, err := ToK8sContainer(context.TODO(), tCtx)
 	assert.NoError(t, err)
 	assert.Equal(t, container.Image, "myimage")
 	assert.EqualValues(t, []string{

go/tasks/pluginmachinery/flytek8s/copilot.go

@@ -175,7 +175,7 @@ func AddCoPilotToContainer(ctx context.Context, cfg config.FlyteCoPilotConfig, c
 	c.SecurityContext.Capabilities.Add = append(c.SecurityContext.Capabilities.Add, pTraceCapability)
 
 	if iFace != nil {
-		if iFace.Inputs != nil {
+		if iFace.Inputs != nil && len(iFace.Inputs.Variables) > 0 {
 			inPath := cfg.DefaultInputDataPath
 			if pilot.GetInputPath() != "" {
 				inPath = pilot.GetInputPath()
@@ -187,7 +187,7 @@ func AddCoPilotToContainer(ctx context.Context, cfg config.FlyteCoPilotConfig, c
 			})
 		}
 
-		if iFace.Outputs != nil {
+		if iFace.Outputs != nil && len(iFace.Outputs.Variables) > 0 {
 			outPath := cfg.DefaultOutputPath
 			if pilot.GetOutputPath() != "" {
 				outPath = pilot.GetOutputPath()
@@ -210,7 +210,7 @@ func AddCoPilotToPod(ctx context.Context, cfg config.FlyteCoPilotConfig, coPilot
 	shareProcessNamespaceEnabled := true
 	coPilotPod.ShareProcessNamespace = &shareProcessNamespaceEnabled
 	if iFace != nil {
-		if iFace.Inputs != nil {
+		if iFace.Inputs != nil && len(iFace.Inputs.Variables) > 0 {
 			inPath := cfg.DefaultInputDataPath
 			if pilot.GetInputPath() != "" {
 				inPath = pilot.GetInputPath()
@@ -240,7 +240,7 @@ func AddCoPilotToPod(ctx context.Context, cfg config.FlyteCoPilotConfig, coPilot
 			coPilotPod.InitContainers = append(coPilotPod.InitContainers, downloader)
 		}
 
-		if iFace.Outputs != nil {
+		if iFace.Outputs != nil && len(iFace.Outputs.Variables) > 0 {
 			outPath := cfg.DefaultOutputPath
 			if pilot.GetOutputPath() != "" {
 				outPath = pilot.GetOutputPath()
@@ -268,7 +268,6 @@ func AddCoPilotToPod(ctx context.Context, cfg config.FlyteCoPilotConfig, coPilot
 			}
 			coPilotPod.Containers = append(coPilotPod.Containers, sidecar)
 		}
-
 	}
 
 	return nil