If primaryContainerName=="primary", container content will be duplica…

…ted (flyteorg#326) * If primaryContainerName=="primary", container content will be duplicated Signed-off-by: byhsu <[email protected]> * address comments Signed-off-by: byhsu <[email protected]> --------- Signed-off-by: byhsu <[email protected]> Co-authored-by: byhsu <[email protected]>
eapolinario · Mar 6, 2023 · a69d75d · a69d75d
1 parent c2c633e
commit a69d75d
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 12 deletions.
diff --git a/flyteplugins/go/tasks/pluginmachinery/flytek8s/pod_helper.go b/flyteplugins/go/tasks/pluginmachinery/flytek8s/pod_helper.go
@@ -326,23 +326,24 @@ func mergePodSpecs(basePodSpec *v1.PodSpec, podSpec *v1.PodSpec, primaryContaine
 		return nil, errors.New("neither the basePodSpec or the podSpec can be nil")
 	}
 
+	// extract defaultContainerTemplate and primaryContainerTemplate
+	var defaultContainerTemplate, primaryContainerTemplate *v1.Container
+	for i := 0; i < len(basePodSpec.Containers); i++ {
+		if basePodSpec.Containers[i].Name == defaultContainerTemplateName {
+			defaultContainerTemplate = &basePodSpec.Containers[i]
+		} else if basePodSpec.Containers[i].Name == primaryContainerTemplateName {
+			primaryContainerTemplate = &basePodSpec.Containers[i]
+		}
+	}
+
 	// merge PodTemplate PodSpec with podSpec
 	var mergedPodSpec *v1.PodSpec = basePodSpec.DeepCopy()
 	if err := mergo.Merge(mergedPodSpec, podSpec, mergo.WithOverride, mergo.WithAppendSlice); err != nil {
 		return nil, err
 	}
 
-	// merge template Containers
+	// merge PodTemplate containers
 	var mergedContainers []v1.Container
-	var defaultContainerTemplate, primaryContainerTemplate *v1.Container
-	for i := 0; i < len(mergedPodSpec.Containers); i++ {
-		if mergedPodSpec.Containers[i].Name == defaultContainerTemplateName {
-			defaultContainerTemplate = &mergedPodSpec.Containers[i]
-		} else if mergedPodSpec.Containers[i].Name == primaryContainerTemplateName {
-			primaryContainerTemplate = &mergedPodSpec.Containers[i]
-		}
-	}
-
 	for _, container := range podSpec.Containers {
 		// if applicable start with defaultContainerTemplate
 		var mergedContainer *v1.Container

diff --git a/flyteplugins/go/tasks/pluginmachinery/flytek8s/pod_helper_test.go b/flyteplugins/go/tasks/pluginmachinery/flytek8s/pod_helper_test.go
@@ -1244,7 +1244,13 @@ func TestMergePodSpecs(t *testing.T) {
 	podSpec := v1.PodSpec{
 		Containers: []v1.Container{
 			v1.Container{
-				Name: "foo",
+				Name: "primary",
+				VolumeMounts: []v1.VolumeMount{
+					{
+						Name:      "nccl",
+						MountPath: "abc",
+					},
+				},
 			},
 			v1.Container{
 				Name: "bar",
@@ -1292,7 +1298,7 @@ func TestMergePodSpecs(t *testing.T) {
 		},
 	}
 
-	mergedPodSpec, err := mergePodSpecs(&podTemplateSpec, &podSpec, "foo")
+	mergedPodSpec, err := mergePodSpecs(&podTemplateSpec, &podSpec, "primary")
 	assert.Nil(t, err)
 
 	// validate a PodTemplate-only field
@@ -1310,6 +1316,7 @@ func TestMergePodSpecs(t *testing.T) {
 	primaryContainer := mergedPodSpec.Containers[0]
 	assert.Equal(t, podSpec.Containers[0].Name, primaryContainer.Name)
 	assert.Equal(t, primaryContainerTemplate.TerminationMessagePath, primaryContainer.TerminationMessagePath)
+	assert.Equal(t, 1, len(primaryContainer.VolumeMounts))
 
 	// validate default container
 	defaultContainer := mergedPodSpec.Containers[1]