From 23a2bc3ba370e28a5ef5f0e0df1e40ec8f7888b9 Mon Sep 17 00:00:00 2001
From: Dharmit Dalvi <dharmitdalvi2@gmail.com>
Date: Tue, 5 Dec 2023 16:02:06 -0500
Subject: [PATCH] feat(backend): Parameterizing v2 Launcher and Driver Images
 (#10269)

* Parameterize v2 launcher image

* Parameterize v2 driver image
---
 backend/src/v2/compiler/argocompiler/argo.go  |  4 +--
 .../src/v2/compiler/argocompiler/container.go | 25 +++++++++++++++++--
 2 files changed, 25 insertions(+), 4 deletions(-)

diff --git a/backend/src/v2/compiler/argocompiler/argo.go b/backend/src/v2/compiler/argocompiler/argo.go
index d7c488972ae..a5cfed5faef 100644
--- a/backend/src/v2/compiler/argocompiler/argo.go
+++ b/backend/src/v2/compiler/argocompiler/argo.go
@@ -116,8 +116,8 @@ func Compile(jobArg *pipelinespec.PipelineJob, kubernetesSpecArg *pipelinespec.S
 		wf:        wf,
 		templates: make(map[string]*wfapi.Template),
 		// TODO(chensun): release process and update the images.
-		driverImage:   "gcr.io/ml-pipeline/kfp-driver@sha256:8e60086b04d92b657898a310ca9757631d58547e76bbbb8bfc376d654bef1707",
-		launcherImage: "gcr.io/ml-pipeline/kfp-launcher@sha256:50151a8615c8d6907aa627902dce50a2619fd231f25d1e5c2a72737a2ea4001e",
+		launcherImage: GetLauncherImage(),
+		driverImage:   GetDriverImage(),
 		job:           job,
 		spec:          spec,
 		executors:     deploy.GetExecutors(),
diff --git a/backend/src/v2/compiler/argocompiler/container.go b/backend/src/v2/compiler/argocompiler/container.go
index 4b115da9ff3..03d85f18019 100644
--- a/backend/src/v2/compiler/argocompiler/container.go
+++ b/backend/src/v2/compiler/argocompiler/container.go
@@ -16,6 +16,7 @@ package argocompiler
 
 import (
 	wfapi "github.com/argoproj/argo-workflows/v3/pkg/apis/workflow/v1alpha1"
+	"os"
 	"github.com/kubeflow/pipelines/api/v2alpha1/go/pipelinespec"
 	"github.com/kubeflow/pipelines/backend/src/v2/component"
 	k8score "k8s.io/api/core/v1"
@@ -23,6 +24,10 @@ import (
 
 const (
 	volumeNameKFPLauncher = "kfp-launcher"
+	DefaultLauncherImage = "gcr.io/ml-pipeline/kfp-launcher@sha256:80cf120abd125db84fa547640fd6386c4b2a26936e0c2b04a7d3634991a850a4"
+	LauncherImageEnvVar   = "V2_LAUNCHER_IMAGE"
+	DefaultDriverImage = "gcr.io/ml-pipeline/kfp-driver@sha256:8e60086b04d92b657898a310ca9757631d58547e76bbbb8bfc376d654bef1707"
+	DriverImageEnvVar   = "V2_DRIVER_IMAGE"
 )
 
 func (c *workflowCompiler) Container(name string, component *pipelinespec.ComponentSpec, container *pipelinespec.PipelineDeploymentConfig_PipelineContainerSpec) error {
@@ -52,6 +57,22 @@ type containerDriverInputs struct {
 	kubernetesConfig string // optional, used when Kubernetes config is not empty
 }
 
+func GetLauncherImage() string {
+    launcherImage := os.Getenv(LauncherImageEnvVar)
+    if launcherImage == "" {
+        launcherImage = DefaultLauncherImage
+    }
+    return launcherImage
+}
+
+func GetDriverImage() string {
+    driverImage := os.Getenv(DriverImageEnvVar)
+    if driverImage == "" {
+        driverImage = DefaultDriverImage
+    }
+    return driverImage
+}
+
 func (c *workflowCompiler) containerDriverTask(name string, inputs containerDriverInputs) (*wfapi.DAGTask, *containerDriverOutputs) {
 	dagTask := &wfapi.DAGTask{
 		Name:     name,
@@ -111,7 +132,7 @@ func (c *workflowCompiler) addContainerDriverTemplate() string {
 			},
 		},
 		Container: &k8score.Container{
-			Image:   c.driverImage,
+			Image:   GetDriverImage(),
 			Command: []string{"driver"},
 			Args: []string{
 				"--type", "CONTAINER",
@@ -226,7 +247,7 @@ func (c *workflowCompiler) addContainerExecutorTemplate() string {
 		InitContainers: []wfapi.UserContainer{{
 			Container: k8score.Container{
 				Name:    "kfp-launcher",
-				Image:   c.launcherImage,
+				Image:   GetLauncherImage(),
 				Command: []string{"launcher-v2", "--copy", component.KFPLauncherPath},
 				VolumeMounts: []k8score.VolumeMount{{
 					Name:      volumeNameKFPLauncher,