From a69d75d4663747dd75ee3b0928ad4842f4d5399c Mon Sep 17 00:00:00 2001 From: ByronHsu Date: Mon, 6 Mar 2023 04:28:38 -0800 Subject: [PATCH] If primaryContainerName=="primary", container content will be duplicated (#326) * If primaryContainerName=="primary", container content will be duplicated Signed-off-by: byhsu * address comments Signed-off-by: byhsu --------- Signed-off-by: byhsu Co-authored-by: byhsu --- .../pluginmachinery/flytek8s/pod_helper.go | 21 ++++++++++--------- .../flytek8s/pod_helper_test.go | 11 ++++++++-- 2 files changed, 20 insertions(+), 12 deletions(-) diff --git a/flyteplugins/go/tasks/pluginmachinery/flytek8s/pod_helper.go b/flyteplugins/go/tasks/pluginmachinery/flytek8s/pod_helper.go index e5cf7cefad..bd49e3c1b8 100755 --- 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 index 59ea6c4dd1..58a5e012a9 100755 --- 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]