Enable custom training job in SageMaker plugin

go/tasks/pluginmachinery/k8s/plugin.go

@@ -3,6 +3,8 @@ package k8s
 import (
 	"context"
 
+	"github.com/lyft/flytestdlib/storage"
+
 	"github.com/lyft/flyteplugins/go/tasks/pluginmachinery/io"
 
 	"k8s.io/apimachinery/pkg/runtime"
@@ -48,6 +50,12 @@ type PluginContext interface {
 
 	// Provides an output sync of type io.OutputWriter
 	OutputWriter() io.OutputWriter
+
+	// Returns a handle to the currently configured storage backend that can be used to communicate with the tasks or write metadata
+	DataStore() *storage.DataStore
+
+	// Returns the max allowed dataset size that the outputwriter will accept
+	MaxDatasetSizeBytes() int64
 }
 
 // Defines a simplified interface to author plugins for k8s resources.

go/tasks/pluginmachinery/utils/template.go

@@ -65,12 +65,7 @@ func transformVarNameToStringVal(ctx context.Context, varName string, inputs *co
 	return v, nil
 }
 
-func replaceTemplateCommandArgs(ctx context.Context, commandTemplate string, in io.InputReader, out io.OutputFilePaths) (string, error) {
-	val := inputFileRegex.ReplaceAllString(commandTemplate, in.GetInputPath().String())
-	val = outputRegex.ReplaceAllString(val, out.GetOutputPrefixPath().String())
-	val = inputPrefixRegex.ReplaceAllString(val, in.GetInputPrefixPath().String())
-	val = rawOutputDataPrefixRegex.ReplaceAllString(val, out.GetRawOutputPrefix().String())
-
+func replaceInputVarsTemplateCommandArgs(ctx context.Context, in io.InputReader, val string) (string, error) {
 	inputs, err := in.Get(ctx)
 	if err != nil {
 		return val, errors.Wrapf(err, "unable to read inputs")
@@ -98,6 +93,14 @@ func replaceTemplateCommandArgs(ctx context.Context, commandTemplate string, in
 	return val, nil
 }
 
+func replaceTemplateCommandArgs(ctx context.Context, commandTemplate string, in io.InputReader, out io.OutputFilePaths) (string, error) {
+	val := inputFileRegex.ReplaceAllString(commandTemplate, in.GetInputPath().String())
+	val = outputRegex.ReplaceAllString(val, out.GetOutputPrefixPath().String())
+	val = inputPrefixRegex.ReplaceAllString(val, in.GetInputPrefixPath().String())
+
+	return replaceInputVarsTemplateCommandArgs(ctx, in, val)
+}
+
 func serializePrimitive(p *core.Primitive) (string, error) {
 	switch o := p.Value.(type) {
 	case *core.Primitive_Integer:

go/tasks/plugins/awsutils/awsutils.go

@@ -4,7 +4,9 @@ import "context"
 
 func GetRole(_ context.Context, roleAnnotationKey string, annotations map[string]string) string {
 	if len(roleAnnotationKey) > 0 {
-		return annotations[roleAnnotationKey]
+		if role, found := annotations[roleAnnotationKey]; found {
+			return role
+		}
 	}
 
 	return ""

go/tasks/plugins/k8s/sagemaker/constants.go

@@ -5,6 +5,11 @@ const (
 	trainingJobTaskType     = "sagemaker_training_job_task"
 )
 
+const (
+	customTrainingJobTaskPluginID = "sagemaker_custom_training"
+	customTrainingJobTaskType     = "sagemaker_custom_training_job_task"
+)
+
 const (
 	hyperparameterTuningJobTaskPluginID = "sagemaker_hyperparameter_tuning"
 	hyperparameterTuningJobTaskType     = "sagemaker_hyperparameter_tuning_job_task"
@@ -13,3 +18,16 @@ const (
 const (
 	TEXTCSVInputContentType string = "text/csv"
 )
+
+const (
+	FlyteSageMakerEnvVarKeyPrefix         string = "__FLYTE_ENV_VAR_"
+	FlyteSageMakerKeySuffix               string = "__"
+	FlyteSageMakerCmdKeyPrefix            string = "__FLYTE_CMD_"
+	FlyteSageMakerCmdDummyValue           string = "__FLYTE_CMD_DUMMY_VALUE__"
+	FlytesageMakerEnvVarKeyStatsdDisabled string = "FLYTE_STATSD_DISABLED"
+)
+
+const (
+	TrainingJobOutputPathSubDir    = "training_outputs"
+	HyperparameterOutputPathSubDir = "hyperparameter_tuning_outputs"
+)

go/tasks/plugins/k8s/sagemaker/outputs.go

@@ -0,0 +1,46 @@
+package sagemaker
+
+import (
+	"context"
+
+	"github.com/lyft/flyteplugins/go/tasks/pluginmachinery/io"
+	"github.com/lyft/flyteplugins/go/tasks/pluginmachinery/ioutils"
+	"github.com/lyft/flytestdlib/logger"
+	"github.com/lyft/flytestdlib/storage"
+)
+
+type JobOutputPaths struct {
+	io.OutputFilePaths
+	jobName      string
+	store        storage.ReferenceConstructor
+	outputPrefix storage.DataReference
+}
+
+func constructPath(store storage.ReferenceConstructor, base storage.DataReference, suffix string) storage.DataReference {
+	res, err := store.ConstructReference(context.Background(), base, suffix)
+	if err != nil {
+		logger.Error(context.Background(), "Failed to construct path. Base[%v] Error: %v", base, err)
+	}
+
+	return res
+}
+
+func (s JobOutputPaths) GetOutputPrefixPath() storage.DataReference {
+	return constructPath(s.store, s.outputPrefix, s.jobName)
+}
+
+func (s JobOutputPaths) GetOutputPath() storage.DataReference {
+	return constructPath(s.store, s.GetOutputPrefixPath(), ioutils.OutputsSuffix)
+}
+
+func (s JobOutputPaths) GetErrorPath() storage.DataReference {
+	return constructPath(s.store, s.GetOutputPrefixPath(), ioutils.ErrorsSuffix)
+}
+
+func NewJobOutputPaths(_ context.Context, store storage.ReferenceConstructor, outputPrefix storage.DataReference, jobName string) JobOutputPaths {
+	return JobOutputPaths{
+		jobName:      jobName,
+		store:        store,
+		outputPrefix: outputPrefix,
+	}
+}