feat(backend): Parameterizing v2 Launcher Image. Resolves #1358 (#1359)

kubeflow · Oct 4, 2023 · ec56e80 · ec56e80
1 parent 4acb2c2
commit ec56e80
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 1 deletion.
diff --git a/backend/src/v2/compiler/tektoncompiler/container.go b/backend/src/v2/compiler/tektoncompiler/container.go
@@ -37,6 +37,7 @@ const (
 	MetadataGPRCServicePort = "8080"
 	MLPipelineServiceHost   = "ml-pipeline.kubeflow.svc.cluster.local"
 	MLPipelineServicePort   = "8887"
+	LauncherImage           = "gcr.io/ml-pipeline/kfp-launcher@sha256:80cf120abd125db84fa547640fd6386c4b2a26936e0c2b04a7d3634991a850a4"
 )
 
 var (
@@ -45,6 +46,7 @@ var (
 	metadataGRPCServicePort = MetadataGPRCServicePort
 	mlPipelineServiceHost   = MLPipelineServiceHost
 	mlPipelineServicePort   = MLPipelineServicePort
+	launcherImage           = LauncherImage
 )
 
 func initEnvVars() {
@@ -68,6 +70,11 @@ func initEnvVars() {
 	if metadataGRPCServicePort == "" {
 		metadataGRPCServicePort = MetadataGPRCServicePort
 	}
+
+	launcherImage = os.Getenv("V2_LAUNCHER_IMAGE")
+	if launcherImage == "" {
+		launcherImage = LauncherImage
+	}
 	envVarInit = true
 }
 
@@ -99,6 +106,13 @@ func GetMLPipelinePort() string {
 	return mlPipelineServicePort
 }
 
+func GetLauncherImage() string {
+	if !envVarInit {
+		initEnvVars()
+	}
+	return launcherImage
+}
+
 // add KubernetesSpec for the container of the component
 func (c *pipelinerunCompiler) AddKubernetesSpec(name string, kubernetesSpec *structpb.Struct) error {
 	err := c.saveKubernetesSpec(name, kubernetesSpec)

diff --git a/backend/src/v2/compiler/tektoncompiler/tekton.go b/backend/src/v2/compiler/tektoncompiler/tekton.go
@@ -113,7 +113,7 @@ func Compile(jobArg *pipelinespec.PipelineJob, kubernetesSpecArg *pipelinespec.S
 	c := &pipelinerunCompiler{
 		pr: pr,
 		// TODO(chensun): release process and update the images.
-		launcherImage: "gcr.io/ml-pipeline/kfp-launcher@sha256:80cf120abd125db84fa547640fd6386c4b2a26936e0c2b04a7d3634991a850a4",
+		launcherImage: GetLauncherImage(),
 		job:           job,
 		spec:          spec,
 		dagStack:      make([]string, 0, 10),