- merge tests

- fix formatting

Signed-off-by: arpechenin <[email protected]>

Author: arpechenin

backend/src/driver/api/request.go

@@ -1,3 +1,18 @@
+// Copyright 2025 The Kubeflow Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// Package api provides HTTP DTOs used by the driver server.
 package api
 
 type DriverPluginArgs struct {
@@ -6,8 +21,8 @@ type DriverPluginArgs struct {
 	Container            string `json:"container,omitempty"`
 	DagExecutionID       string `json:"dag_execution_id"`
 	IterationIndex       string `json:"iteration_index"`
-	HttpProxy            string `json:"http_proxy"`
-	HttpsProxy           string `json:"https_proxy"`
+	HTTPProxy            string `json:"http_proxy"`
+	HTTPSProxy           string `json:"https_proxy"`
 	NoProxy              string `json:"no_proxy"`
 	KubernetesConfig     string `json:"kubernetes_config,omitempty"`
 	RuntimeConfig        string `json:"runtime_config,omitempty"`
@@ -20,7 +35,7 @@ type DriverPluginArgs struct {
 	Task                 string `json:"task"`
 	Type                 string `json:"type"`
 	CacheDisabledFlag    bool   `json:"cache_disabled"`
-	ExecutionIdPath      string `json:"execution_id_path"`
+	ExecutionIDPath      string `json:"execution_id_path"`
 	IterationCountPath   string `json:"iteration_count_path"`
 	ConditionPath        string `json:"condition_path"`
 	PodSpecPathPath      string `json:"pod_spec_patch_path"`

backend/src/driver/api/response.go

@@ -1,3 +1,17 @@
+// Copyright 2025 The Kubeflow Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
 package api
 
 type DriverResponse struct {

backend/src/driver/execution_paths.go

@@ -1,3 +1,17 @@
+// Copyright 2025 The Kubeflow Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
 package main
 
 type ExecutionPaths struct {

backend/src/driver/main.go

@@ -11,6 +11,7 @@
 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 // See the License for the specific language governing permissions and
 // limitations under the License.
+
 package main
 
 import (
@@ -37,7 +38,7 @@ import (
 
 const (
 	unsetProxyArgValue = "unset"
-	ROOT_DAG           = "ROOT_DAG"
+	RootDag            = "ROOT_DAG"
 	DAG                = "DAG"
 	CONTAINER          = "CONTAINER"
 )
@@ -76,13 +77,13 @@ func init() {
 	flag.Set("stderrthreshold", "WARNING")
 }
 
-func parseExecConfigJson(k8sExecConfigJson *string) (*kubernetesplatform.KubernetesExecutorConfig, error) {
+func parseExecConfigJSON(k8sExecConfigJSON *string) (*kubernetesplatform.KubernetesExecutorConfig, error) {
 	var k8sExecCfg *kubernetesplatform.KubernetesExecutorConfig
-	if *k8sExecConfigJson != "" {
-		glog.Infof("input kubernetesConfig:%s\n", prettyPrint(*k8sExecConfigJson))
+	if *k8sExecConfigJSON != "" {
+		glog.Infof("input kubernetesConfig:%s\n", prettyPrint(*k8sExecConfigJSON))
 		k8sExecCfg = &kubernetesplatform.KubernetesExecutorConfig{}
-		if err := util.UnmarshalString(*k8sExecConfigJson, k8sExecCfg); err != nil {
-			return nil, fmt.Errorf("failed to unmarshal Kubernetes config, error: %w\nKubernetesConfig: %v", err, k8sExecConfigJson)
+		if err := util.UnmarshalString(*k8sExecConfigJSON, k8sExecCfg); err != nil {
+			return nil, fmt.Errorf("failed to unmarshal Kubernetes config, error: %w\nKubernetesConfig: %v", err, k8sExecConfigJSON)
 		}
 	}
 	return k8sExecCfg, nil
@@ -102,7 +103,7 @@ func handleExecution(execution *driver.Execution, driverType string, executionPa
 			return fmt.Errorf("failed to write iteration count to file: %w", err)
 		}
 	} else {
-		if driverType == ROOT_DAG {
+		if driverType == RootDag {
 			if err := writeFile(executionPaths.IterationCount, []byte("0")); err != nil {
 				return fmt.Errorf("failed to write iteration count to file: %w", err)
 			}
@@ -119,7 +120,7 @@ func handleExecution(execution *driver.Execution, driverType string, executionPa
 		}
 	} else {
 		// nil is a valid value for Condition
-		if driverType == ROOT_DAG || driverType == CONTAINER {
+		if driverType == RootDag || driverType == CONTAINER {
 			if err := writeFile(executionPaths.Condition, []byte("nil")); err != nil {
 				return fmt.Errorf("failed to write condition to file: %w", err)
 			}

backend/src/driver/main_test.go

@@ -1,3 +1,17 @@
+// Copyright 2025 The Kubeflow Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
 package main
 
 import (
@@ -45,7 +59,7 @@ func TestSpecParsing(t *testing.T) {
 
 	for _, tc := range tt {
 		t.Logf("Running test case: %s", tc.name)
-		cfg, err := parseExecConfigJson(tc.input)
+		cfg, err := parseExecConfigJSON(tc.input)
 		assert.Equal(t, tc.wantErr, err != nil)
 		assert.True(t, proto.Equal(tc.expected, cfg))
 	}
@@ -77,7 +91,7 @@ func Test_handleExecutionRootDAG(t *testing.T) {
 		Condition:      "condition.txt",
 	}
 
-	err := handleExecution(execution, ROOT_DAG, executionPaths)
+	err := handleExecution(execution, RootDag, executionPaths)
 
 	if err != nil {
 		t.Errorf("Unexpected error: %v", err)

backend/src/driver/rpc_handler.go

@@ -1,3 +1,17 @@
+// Copyright 2025 The Kubeflow Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
 package main
 
 import (
@@ -94,13 +108,14 @@ func parseDriverRequestArgs(r *http.Request) (*api.DriverPluginArgs, error) {
 	if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
 		return nil, fmt.Errorf("failed to parse driver request body: %v", err)
 	}
-	if body.Template == nil {
+	switch {
+	case body.Template == nil:
 		return nil, fmt.Errorf("driver request body.Template is empty")
-	} else if body.Template.Plugin == nil {
+	case body.Template.Plugin == nil:
 		return nil, fmt.Errorf("driver request body.Template.Plugin is empty")
-	} else if body.Template.Plugin.DriverPlugin == nil {
+	case body.Template.Plugin.DriverPlugin == nil:
 		return nil, fmt.Errorf("driver request body.Template.Plugin.DriverPlugin is empty")
-	} else if body.Template.Plugin.DriverPlugin.Args == nil {
+	case body.Template.Plugin.DriverPlugin.Args == nil:
 		return nil, fmt.Errorf("driver request body.Template.Plugin.Args is empty")
 	}
 	args := body.Template.Plugin.DriverPlugin.Args
@@ -118,7 +133,7 @@ func drive(args api.DriverPluginArgs) (execution *driver.Execution, err error) {
 	}()
 	ctx := context.Background()
 
-	proxy.InitializeConfig(args.HttpProxy, args.HttpsProxy, args.NoProxy)
+	proxy.InitializeConfig(args.HTTPProxy, args.HTTPSProxy, args.NoProxy)
 
 	glog.Infof("input ComponentSpec:%s\n", prettyPrint(args.Component))
 	componentSpec := &pipelinespec.ComponentSpec{}
@@ -149,7 +164,7 @@ func drive(args api.DriverPluginArgs) (execution *driver.Execution, err error) {
 			return nil, fmt.Errorf("failed to unmarshal runtime config, error: %w\nruntimeConfig: %v", err, args.RuntimeConfig)
 		}
 	}
-	k8sExecCfg, err := parseExecConfigJson(&args.KubernetesConfig)
+	k8sExecCfg, err := parseExecConfigJSON(&args.KubernetesConfig)
 	if err != nil {
 		return nil, err
 	}
@@ -160,7 +175,9 @@ func drive(args api.DriverPluginArgs) (execution *driver.Execution, err error) {
 	var tlsCfg *tls.Config
 	if args.MetadataTLSEnabled {
 		tlsCfg, err = util.GetTLSConfig(args.CACertPath)
-		return nil, fmt.Errorf("unable to drive driver: failed to load TLS configuration: %v", err)
+		if err != nil {
+			return nil, fmt.Errorf("unable to drive driver: failed to load TLS configuration: %v", err)
+		}
 	}
 	client, err := newMlmdClient(tlsCfg)
 	if err != nil {
@@ -180,23 +197,28 @@ func drive(args api.DriverPluginArgs) (execution *driver.Execution, err error) {
 		return nil, fmt.Errorf("failed to parse iteration index, error: %w", err)
 	}
 	options := driver.Options{
-		PipelineName:   args.PipelineName,
-		RunID:          args.RunID,
-		RunName:        args.RunName,
-		RunDisplayName: args.RunDisplayName,
-		Namespace:      namespace,
-		Component:      componentSpec,
-		Task:           taskSpec,
-		DAGExecutionID: dagExecutionID,
-		IterationIndex: iterationIndex,
-		PublishLogs:    args.PublishLogs,
-		CacheDisabled:  args.CacheDisabledFlag,
-		DriverType:     args.Type,
-		TaskName:       args.TaskName,
+		PipelineName:         args.PipelineName,
+		RunID:                args.RunID,
+		RunName:              args.RunName,
+		RunDisplayName:       args.RunDisplayName,
+		Namespace:            namespace,
+		Component:            componentSpec,
+		Task:                 taskSpec,
+		DAGExecutionID:       dagExecutionID,
+		IterationIndex:       iterationIndex,
+		PublishLogs:          args.PublishLogs,
+		CacheDisabled:        args.CacheDisabledFlag,
+		DriverType:           args.Type,
+		TaskName:             args.TaskName,
+		MLPipelineTLSEnabled: args.MlPipelineTLSEnabled,
+		MLMDServerAddress:    *mlmdServerAddress,
+		MLMDServerPort:       *mlmdServerPort,
+		CaCertPath:           args.CACertPath,
 	}
+
 	var driverErr error
 	switch args.Type {
-	case ROOT_DAG:
+	case RootDag:
 		options.RuntimeConfig = runtimeConfig
 		execution, driverErr = driver.RootDAG(ctx, options, client)
 	case DAG:
@@ -224,19 +246,16 @@ func drive(args api.DriverPluginArgs) (execution *driver.Execution, err error) {
 }
 
 func validate(args api.DriverPluginArgs) error {
-	if args.Type == "" {
+	switch {
+	case args.Type == "":
 		return fmt.Errorf("argument type must be specified")
-	}
-	if args.HttpProxy == unsetProxyArgValue {
+	case args.HTTPProxy == unsetProxyArgValue:
 		return fmt.Errorf("argument http_proxy is required but can be an empty value")
-	}
-	if args.HttpsProxy == unsetProxyArgValue {
+	case args.HTTPSProxy == unsetProxyArgValue:
 		return fmt.Errorf("argument https_proxy is required but can be an empty value")
-	}
-	if args.NoProxy == unsetProxyArgValue {
+	case args.NoProxy == unsetProxyArgValue:
 		return fmt.Errorf("argument no_proxy is required but can be an empty value")
 	}
-	// validation responsibility lives in driver itself, so we do not validate all other args
 	return nil
 }
 
@@ -251,18 +270,17 @@ func extractOutputParameters(execution *driver.Execution, driverType string) []a
 			Value: fmt.Sprint(execution.ID),
 		})
 	}
-	if execution.IterationCount != nil {
+	switch {
+	case execution.IterationCount != nil:
 		outputs = append(outputs, api.Parameter{
 			Name:  "iteration-count",
-			Value: fmt.Sprint(execution.IterationCount),
+			Value: fmt.Sprint(*execution.IterationCount),
+		})
+	case driverType == RootDag:
+		outputs = append(outputs, api.Parameter{
+			Name:  "iteration-count",
+			Value: "0",
 		})
-	} else {
-		if driverType == ROOT_DAG {
-			outputs = append(outputs, api.Parameter{
-				Name:  "iteration-count",
-				Value: fmt.Sprint(0),
-			})
-		}
 	}
 	if execution.Cached != nil {
 		outputs = append(outputs, api.Parameter{