feat(backend): configurable log level for driver / launcher images

backend/src/v2/cmd/driver/main.go

@@ -68,13 +68,21 @@ var (
 	// the value stored in the paths will be either 'true' or 'false'
 	cachedDecisionPath = flag.String("cached_decision_path", "", "Cached Decision output path")
 	conditionPath      = flag.String("condition_path", "", "Condition output path")
+	logLevel           = flag.String("log_level", "1", "The verbosity level to log.")
 )
 
 // func RootDAG(pipelineName string, runID string, component *pipelinespec.ComponentSpec, task *pipelinespec.PipelineTaskSpec, mlmd *metadata.Client) (*Execution, error) {
 
 func main() {
 	flag.Parse()
-	err := drive()
+
+	glog.Infof("Setting log level to: '%s'", *logLevel)
+	err := flag.Set("v", *logLevel)
+	if err != nil {
+		glog.Warningf("Failed to set log level: %s", err.Error())
+	}
+
+	err = drive()
 	if err != nil {
 		glog.Exitf("%v", err)
 	}

backend/src/v2/cmd/launcher-v2/main.go

@@ -41,6 +41,7 @@ var (
 	podUID            = flag.String("pod_uid", "", "Kubernetes Pod UID.")
 	mlmdServerAddress = flag.String("mlmd_server_address", "", "The MLMD gRPC server address.")
 	mlmdServerPort    = flag.String("mlmd_server_port", "8080", "The MLMD gRPC server port.")
+	logLevel          = flag.String("log_level", "1", "The verbosity level to log.")
 )
 
 func main() {
@@ -54,6 +55,12 @@ func run() error {
 	flag.Parse()
 	ctx := context.Background()
 
+	glog.Infof("Setting log level to: '%s'", *logLevel)
+	err := flag.Set("v", *logLevel)
+	if err != nil {
+		glog.Warningf("Failed to set log level: %s", err.Error())
+	}
+
 	if *copy != "" {
 		// copy is used to copy this binary to a shared volume
 		// this is a special command, ignore all other flags by returning

backend/src/v2/compiler/argocompiler/argo_test.go

@@ -18,6 +18,7 @@ import (
 	"flag"
 	"fmt"
 	"io/ioutil"
+	"os"
 	"strings"
 	"testing"
 
@@ -36,6 +37,7 @@ func Test_argo_compiler(t *testing.T) {
 		jobPath          string // path of input PipelineJob to compile
 		platformSpecPath string // path of possible input PlatformSpec to compile
 		argoYAMLPath     string // path of expected output argo workflow YAML
+		envVars          []string
 	}{
 		{
 			jobPath:          "../testdata/hello_world.json",
@@ -57,10 +59,40 @@ func Test_argo_compiler(t *testing.T) {
 			platformSpecPath: "../testdata/create_pod_metadata.json",
 			argoYAMLPath:     "testdata/create_pod_metadata.yaml",
 		},
+		// With envOptions
+		{
+			jobPath:          "../testdata/hello_world.json",
+			platformSpecPath: "",
+			argoYAMLPath:     "testdata/with_logging/hello_world.yaml",
+			envVars:          []string{"DRIVER_LOG_LEVEL=5", "LAUNCHER_LOG_LEVEL=5"},
+		},
+		{
+			jobPath:          "../testdata/importer.json",
+			platformSpecPath: "",
+			argoYAMLPath:     "testdata/with_logging/importer.yaml",
+			envVars:          []string{"DRIVER_LOG_LEVEL=5", "LAUNCHER_LOG_LEVEL=5"},
+		},
 	}
 	for _, tt := range tests {
 		t.Run(fmt.Sprintf("%+v", tt), func(t *testing.T) {
 			job, platformSpec := load(t, tt.jobPath, tt.platformSpecPath)
+			if tt.envVars != nil {
+				for _, envVar := range tt.envVars {
+					parts := strings.Split(strings.ReplaceAll(envVar, " ", ""), "=")
+					err := os.Setenv(parts[0], parts[1])
+					if err != nil {
+						t.Fatalf("Failed to set environment variable '%s' with error: %s", parts[0], err.Error())
+					}
+
+					// Unset after test cases has ended
+					defer func() {
+						err := os.Unsetenv(parts[0])
+						if err != nil {
+							t.Fatalf("Failed to unset env variable '%s' with error: %s", parts[0], err.Error())
+						}
+					}()
+				}
+			}
 			if *update {
 				wf, err := argocompiler.Compile(job, platformSpec, nil)
 				if err != nil {

backend/src/v2/compiler/argocompiler/container.go

@@ -34,6 +34,8 @@ const (
 	LauncherImageEnvVar      = "V2_LAUNCHER_IMAGE"
 	DefaultDriverImage       = "gcr.io/ml-pipeline/kfp-driver@sha256:dc8b56a2eb071f30409828a8884d621092e68385af11a6c06aa9e9fbcfbb19de"
 	DriverImageEnvVar        = "V2_DRIVER_IMAGE"
+	DriverLogLevelEnvVar     = "DRIVER_LOG_LEVEL"
+	LauncherLogLevelEnvVar   = "LAUNCHER_LOG_LEVEL"
 	gcsScratchLocation       = "/gcs"
 	gcsScratchName           = "gcs-scratch"
 	s3ScratchLocation        = "/s3"
@@ -130,6 +132,25 @@ func (c *workflowCompiler) addContainerDriverTemplate() string {
 	if ok {
 		return name
 	}
+
+	args := []string{
+		"--type", "CONTAINER",
+		"--pipeline_name", c.spec.GetPipelineInfo().GetName(),
+		"--run_id", runID(),
+		"--dag_execution_id", inputValue(paramParentDagID),
+		"--component", inputValue(paramComponent),
+		"--task", inputValue(paramTask),
+		"--container", inputValue(paramContainer),
+		"--iteration_index", inputValue(paramIterationIndex),
+		"--cached_decision_path", outputPath(paramCachedDecision),
+		"--pod_spec_patch_path", outputPath(paramPodSpecPatch),
+		"--condition_path", outputPath(paramCondition),
+		"--kubernetes_config", inputValue(paramKubernetesConfig),
+	}
+	if value, ok := os.LookupEnv(DriverLogLevelEnvVar); ok {
+		args = append(args, "--log_level", value)
+	}
+
 	t := &wfapi.Template{
 		Name: name,
 		Inputs: wfapi.Inputs{
@@ -150,22 +171,9 @@ func (c *workflowCompiler) addContainerDriverTemplate() string {
 			},
 		},
 		Container: &k8score.Container{
-			Image:   GetDriverImage(),
-			Command: []string{"driver"},
-			Args: []string{
-				"--type", "CONTAINER",
-				"--pipeline_name", c.spec.GetPipelineInfo().GetName(),
-				"--run_id", runID(),
-				"--dag_execution_id", inputValue(paramParentDagID),
-				"--component", inputValue(paramComponent),
-				"--task", inputValue(paramTask),
-				"--container", inputValue(paramContainer),
-				"--iteration_index", inputValue(paramIterationIndex),
-				"--cached_decision_path", outputPath(paramCachedDecision),
-				"--pod_spec_patch_path", outputPath(paramPodSpecPatch),
-				"--condition_path", outputPath(paramCondition),
-				"--kubernetes_config", inputValue(paramKubernetesConfig),
-			},
+			Image:     c.driverImage,
+			Command:   []string{"driver"},
+			Args:      args,
 			Resources: driverResources,
 		},
 	}
@@ -245,6 +253,13 @@ func (c *workflowCompiler) addContainerExecutorTemplate(refName string) string {
 		},
 	}
 	c.templates[nameContainerExecutor] = container
+
+	args := []string{
+		"--copy", component.KFPLauncherPath,
+	}
+	if value, ok := os.LookupEnv(LauncherLogLevelEnvVar); ok {
+		args = append(args, "--log_level", value)
+	}
 	executor := &wfapi.Template{
 		Name: nameContainerImpl,
 		Inputs: wfapi.Inputs{
@@ -303,8 +318,9 @@ func (c *workflowCompiler) addContainerExecutorTemplate(refName string) string {
 		InitContainers: []wfapi.UserContainer{{
 			Container: k8score.Container{
 				Name:    "kfp-launcher",
-				Image:   GetLauncherImage(),
-				Command: []string{"launcher-v2", "--copy", component.KFPLauncherPath},
+				Image:   c.launcherImage,
+				Command: []string{"launcher-v2"},
+				Args:    args,
 				VolumeMounts: []k8score.VolumeMount{
 					{
 						Name:      volumeNameKFPLauncher,

backend/src/v2/compiler/argocompiler/dag.go

@@ -15,6 +15,7 @@ package argocompiler
 
 import (
 	"fmt"
+	"os"
 	"sort"
 	"strings"
 
@@ -406,6 +407,24 @@ func (c *workflowCompiler) addDAGDriverTemplate() string {
 	if ok {
 		return name
 	}
+
+	args := []string{
+		"--type", inputValue(paramDriverType),
+		"--pipeline_name", c.spec.GetPipelineInfo().GetName(),
+		"--run_id", runID(),
+		"--dag_execution_id", inputValue(paramParentDagID),
+		"--component", inputValue(paramComponent),
+		"--task", inputValue(paramTask),
+		"--runtime_config", inputValue(paramRuntimeConfig),
+		"--iteration_index", inputValue(paramIterationIndex),
+		"--execution_id_path", outputPath(paramExecutionID),
+		"--iteration_count_path", outputPath(paramIterationCount),
+		"--condition_path", outputPath(paramCondition),
+	}
+	if value, ok := os.LookupEnv(DriverLogLevelEnvVar); ok {
+		args = append(args, "--log_level", value)
+	}
+
 	t := &wfapi.Template{
 		Name: name,
 		Inputs: wfapi.Inputs{
@@ -426,21 +445,9 @@ func (c *workflowCompiler) addDAGDriverTemplate() string {
 			},
 		},
 		Container: &k8score.Container{
-			Image:   c.driverImage,
-			Command: []string{"driver"},
-			Args: []string{
-				"--type", inputValue(paramDriverType),
-				"--pipeline_name", c.spec.GetPipelineInfo().GetName(),
-				"--run_id", runID(),
-				"--dag_execution_id", inputValue(paramParentDagID),
-				"--component", inputValue(paramComponent),
-				"--task", inputValue(paramTask),
-				"--runtime_config", inputValue(paramRuntimeConfig),
-				"--iteration_index", inputValue(paramIterationIndex),
-				"--execution_id_path", outputPath(paramExecutionID),
-				"--iteration_count_path", outputPath(paramIterationCount),
-				"--condition_path", outputPath(paramCondition),
-			},
+			Image:     c.driverImage,
+			Command:   []string{"driver"},
+			Args:      args,
 			Resources: driverResources,
 		},
 	}

backend/src/v2/compiler/argocompiler/importer.go

@@ -16,6 +16,7 @@ package argocompiler
 
 import (
 	"fmt"
+	"os"
 
 	wfapi "github.com/argoproj/argo-workflows/v3/pkg/apis/workflow/v1alpha1"
 	"github.com/kubeflow/pipelines/api/v2alpha1/go/pipelinespec"
@@ -64,7 +65,7 @@ func (c *workflowCompiler) addImporterTemplate() string {
 	if _, alreadyExists := c.templates[name]; alreadyExists {
 		return name
 	}
-	launcherArgs := []string{
+	args := []string{
 		"--executor_type", "importer",
 		"--task_spec", inputValue(paramTask),
 		"--component_spec", inputValue(paramComponent),
@@ -81,6 +82,9 @@ func (c *workflowCompiler) addImporterTemplate() string {
 		"--mlmd_server_port",
 		fmt.Sprintf("$(%s)", component.EnvMetadataPort),
 	}
+	if value, ok := os.LookupEnv(LauncherLogLevelEnvVar); ok {
+		args = append(args, "--log_level", value)
+	}
 	importerTemplate := &wfapi.Template{
 		Name: name,
 		Inputs: wfapi.Inputs{
@@ -94,7 +98,7 @@ func (c *workflowCompiler) addImporterTemplate() string {
 		Container: &k8score.Container{
 			Image:     c.launcherImage,
 			Command:   []string{"launcher-v2"},
-			Args:      launcherArgs,
+			Args:      args,
 			EnvFrom:   []k8score.EnvFromSource{metadataEnvFrom},
 			Env:       commonEnvs,
 			Resources: driverResources,

backend/src/v2/compiler/argocompiler/testdata/create_mount_delete_dynamic_pvc.yaml

@@ -155,10 +155,11 @@ spec:
       - mountPath: /.config
         name: dot-config-scratch
     initContainers:
-    - command:
-      - launcher-v2
+    - args:
       - --copy
       - /kfp-launcher/launch
+      command:
+      - launcher-v2
       image: gcr.io/ml-pipeline/kfp-launcher
       name: kfp-launcher
       resources:

backend/src/v2/compiler/argocompiler/testdata/create_pod_metadata.yaml

@@ -143,10 +143,11 @@ spec:
       - mountPath: /.config
         name: dot-config-scratch
     initContainers:
-    - command:
-      - launcher-v2
+    - args:
       - --copy
       - /kfp-launcher/launch
+      command:
+      - launcher-v2
       image: gcr.io/ml-pipeline/kfp-launcher
       name: kfp-launcher
       resources:

backend/src/v2/compiler/argocompiler/testdata/hello_world.yaml

@@ -141,10 +141,11 @@ spec:
       - mountPath: /.config
         name: dot-config-scratch
     initContainers:
-    - command:
-      - launcher-v2
+    - args:
       - --copy
       - /kfp-launcher/launch
+      command:
+      - launcher-v2
       image: gcr.io/ml-pipeline/kfp-launcher
       name: kfp-launcher
       resources: