feat(backend): ARTIFACT_COPY_STEP_TEMPLATE - custom template for copy…

… step (kubeflow#336) Issue: kubeflow/kfp-tekton#333
red-hat-data-services · Oct 29, 2020 · 905079b · 905079b
1 parent 64a9dbb
commit 905079b
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 20 deletions.
diff --git a/backend/src/apiserver/common/config.go b/backend/src/apiserver/common/config.go
@@ -20,6 +20,7 @@ import (
 
 	"github.com/golang/glog"
 	"github.com/spf13/viper"
+	corev1 "k8s.io/api/core/v1"
 )
 
 const (
@@ -37,6 +38,7 @@ const (
 	ArtifactEndpointScheme              string = "ARTIFACT_ENDPOINT_SCHEME"
 	ArtifactScript                      string = "ARTIFACT_SCRIPT"
 	ArtifactImage                       string = "ARTIFACT_IMAGE"
+	ArtifactCopyStepTemplate            string = "ARTIFACT_COPY_STEP_TEMPLATE"
 	InjectDefaultScript                 string = "INJECT_DEFAULT_SCRIPT"
 )
 
@@ -108,6 +110,14 @@ func GetArtifactImage() string {
 	return GetStringConfigWithDefault(ArtifactImage, DefaultArtifactImage)
 }
 
+func GetCopyStepTemplate() *corev1.Container {
+	var tpl corev1.Container
+	if err := viper.UnmarshalKey(ArtifactCopyStepTemplate, &tpl); err != nil {
+		glog.Fatalf("Invalid '%s', %v", ArtifactCopyStepTemplate, err)
+	}
+	return &tpl
+}
+
 func GetBoolFromStringWithDefault(value string, defaultValue bool) bool {
 	boolVal, err := strconv.ParseBool(value)
 	if err != nil {

diff --git a/backend/src/apiserver/resource/resource_manager.go b/backend/src/apiserver/resource/resource_manager.go
@@ -1281,6 +1281,7 @@ func (r *ResourceManager) injectArchivalStep(workflow util.Workflow, artifactIte
 		trackArtifacts := common.IsTrackArtifacts()
 		stripEOF := common.IsStripEOF()
 		injectDefaultScript := common.IsInjectDefaultScript()
+		copyStepTemplate := common.GetCopyStepTemplate()
 
 		if (hasArtifacts && len(artifacts) > 0 && trackArtifacts) || archiveLogs || (hasArtifacts && len(artifacts) > 0 && stripEOF) {
 			artifactScript := common.GetArtifactScript()
@@ -1319,27 +1320,30 @@ func (r *ResourceManager) injectArchivalStep(workflow util.Workflow, artifactIte
 			}
 
 			// Define post-processing step
-			step := workflowapi.Step{Container: corev1.Container{
-				Name:  "copy-artifacts",
-				Image: common.GetArtifactImage(),
-				Env: []corev1.EnvVar{
-					r.getObjectFieldSelector("ARTIFACT_BUCKET", "metadata.annotations['tekton.dev/artifact_bucket']"),
-					r.getObjectFieldSelector("ARTIFACT_ENDPOINT", "metadata.annotations['tekton.dev/artifact_endpoint']"),
-					r.getObjectFieldSelector("ARTIFACT_ENDPOINT_SCHEME", "metadata.annotations['tekton.dev/artifact_endpoint_scheme']"),
-					r.getObjectFieldSelector("ARTIFACT_ITEMS", "metadata.annotations['tekton.dev/artifact_items']"),
-					r.getObjectFieldSelector("PIPELINETASK", "metadata.labels['tekton.dev/pipelineTask']"),
-					r.getObjectFieldSelector("PIPELINERUN", "metadata.labels['tekton.dev/pipelineRun']"),
-					r.getObjectFieldSelector("PODNAME", "metadata.name"),
-					r.getObjectFieldSelector("NAMESPACE", "metadata.namespace"),
-					r.getSecretKeySelector("AWS_ACCESS_KEY_ID", "mlpipeline-minio-artifact", "accesskey"),
-					r.getSecretKeySelector("AWS_SECRET_ACCESS_KEY", "mlpipeline-minio-artifact", "secretkey"),
-					r.getEnvVar("ARCHIVE_LOGS", strconv.FormatBool(archiveLogs)),
-					r.getEnvVar("TRACK_ARTIFACTS", strconv.FormatBool(trackArtifacts)),
-					r.getEnvVar("STRIP_EOF", strconv.FormatBool(stripEOF)),
-				},
-			},
-				Script: artifactScript,
+			container := *copyStepTemplate
+			if container.Name == "" {
+				container.Name = "copy-artifacts"
+			}
+			if container.Image == "" {
+				container.Image = common.GetArtifactImage()
 			}
+			container.Env = append(container.Env,
+				r.getObjectFieldSelector("ARTIFACT_BUCKET", "metadata.annotations['tekton.dev/artifact_bucket']"),
+				r.getObjectFieldSelector("ARTIFACT_ENDPOINT", "metadata.annotations['tekton.dev/artifact_endpoint']"),
+				r.getObjectFieldSelector("ARTIFACT_ENDPOINT_SCHEME", "metadata.annotations['tekton.dev/artifact_endpoint_scheme']"),
+				r.getObjectFieldSelector("ARTIFACT_ITEMS", "metadata.annotations['tekton.dev/artifact_items']"),
+				r.getObjectFieldSelector("PIPELINETASK", "metadata.labels['tekton.dev/pipelineTask']"),
+				r.getObjectFieldSelector("PIPELINERUN", "metadata.labels['tekton.dev/pipelineRun']"),
+				r.getObjectFieldSelector("PODNAME", "metadata.name"),
+				r.getObjectFieldSelector("NAMESPACE", "metadata.namespace"),
+				r.getSecretKeySelector("AWS_ACCESS_KEY_ID", "mlpipeline-minio-artifact", "accesskey"),
+				r.getSecretKeySelector("AWS_SECRET_ACCESS_KEY", "mlpipeline-minio-artifact", "secretkey"),
+				r.getEnvVar("ARCHIVE_LOGS", strconv.FormatBool(archiveLogs)),
+				r.getEnvVar("TRACK_ARTIFACTS", strconv.FormatBool(trackArtifacts)),
+				r.getEnvVar("STRIP_EOF", strconv.FormatBool(stripEOF)),
+			)
+
+			step := workflowapi.Step{Container: container, Script: artifactScript}
 			task.TaskSpec.Steps = append(task.TaskSpec.Steps, step)
 		}
 	}