mirror of
https://github.com/hanzo-ml/trainer.git
synced 2026-09-24 00:31:34 +00:00
Peter PanandGitHub
15b0a9fc1d
* feat: inject PET envs into init containers via envInjection config - Add EnvInjection and EnvInjectionTarget types to TorchMLPolicySource - Add FindContainerByPodSetName helper in runtime package - Update torch plugin to inject PET envs into additional containers - Update jobset plugin to sync init containers - Add 5 comprehensive unit tests for envInjection scenarios - Add Python API models for envInjection types Signed-off-by: Peter Pan <[email protected]> * fix: add missing +optional markers for kubeapilinter validation Signed-off-by: Peter Pan <[email protected]> * fix: change JobName to optional field for kubeapilinter Change JobName from required to optional with omitempty tag to satisfy kubeapilinter validation. Add default marker to CRD schema for listMapKey requirement. Signed-off-by: Peter Pan <[email protected]> * fix test Signed-off-by: Peter Pan <[email protected]> * fix: address envInjection review feedback Reject missing envInjection targets instead of silently skipping them, align the API schema with required fields, and harden init container propagation so misconfigured runtimes fail predictably. Signed-off-by: Peter Pan <[email protected]> * fix: sync envInjection generated files Add the generated boilerplate and Python model updates expected by make generate so CI does not report stale generated assets. Signed-off-by: Peter Pan <[email protected]> * fix: satisfy kube API lint for envInjection fields Keep the envInjection fields marked as required while preserving omitempty tags required by kubeapilinter. Signed-off-by: Peter Pan <[email protected]> * fix: update envInjection OpenAPI defaults Sync generated OpenAPI and Python model output after restoring omitempty on required envInjection fields. Signed-off-by: Peter Pan <[email protected]> * fix: address latest review comments Signed-off-by: Peter Pan <[email protected]> * chore: remove test plan file per review feedback The KEP README already has a Test Plan section; a separate file with local environment details does not belong in the upstream PR. Signed-off-by: Peter Pan <[email protected]> --------- Signed-off-by: Peter Pan <[email protected]>
679 lines
20 KiB
Go
679 lines
20 KiB
Go
/*
|
|
Copyright 2024 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
|
|
|
|
http://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 runtime
|
|
|
|
import (
|
|
"testing"
|
|
|
|
"github.com/google/go-cmp/cmp"
|
|
"github.com/google/go-cmp/cmp/cmpopts"
|
|
corev1 "k8s.io/api/core/v1"
|
|
"k8s.io/apimachinery/pkg/api/resource"
|
|
batchv1ac "k8s.io/client-go/applyconfigurations/batch/v1"
|
|
corev1ac "k8s.io/client-go/applyconfigurations/core/v1"
|
|
"k8s.io/utils/ptr"
|
|
jobsetv1alpha2ac "sigs.k8s.io/jobset/client-go/applyconfiguration/jobset/v1alpha2"
|
|
|
|
"github.com/kubeflow/trainer/v2/pkg/constants"
|
|
jobsetplgconsts "github.com/kubeflow/trainer/v2/pkg/runtime/framework/plugins/jobset/constants"
|
|
)
|
|
|
|
func TestNewInfo(t *testing.T) {
|
|
cases := map[string]struct {
|
|
infoOpts []InfoOption
|
|
wantInfo *Info
|
|
}{
|
|
"all arguments are specified": {
|
|
infoOpts: []InfoOption{
|
|
WithLabels(map[string]string{
|
|
"labelKey": "labelValue",
|
|
}),
|
|
WithAnnotations(map[string]string{
|
|
"annotationKey": "annotationValue",
|
|
}),
|
|
WithPodSet(constants.DatasetInitializer, ptr.To(constants.DatasetInitializer), 1, corev1.PodSpec{
|
|
Containers: []corev1.Container{{
|
|
Name: constants.DatasetInitializer,
|
|
Resources: corev1.ResourceRequirements{
|
|
Requests: corev1.ResourceList{
|
|
corev1.ResourceCPU: resource.MustParse("10"),
|
|
},
|
|
},
|
|
}},
|
|
InitContainers: []corev1.Container{{
|
|
Name: "setup-initializer",
|
|
RestartPolicy: ptr.To(corev1.ContainerRestartPolicyAlways),
|
|
Resources: corev1.ResourceRequirements{
|
|
Requests: corev1.ResourceList{
|
|
corev1.ResourceCPU: resource.MustParse("5"),
|
|
},
|
|
},
|
|
}},
|
|
}, corev1ac.PodSpec().
|
|
WithContainers(
|
|
corev1ac.Container().
|
|
WithName(constants.DatasetInitializer).
|
|
WithResources(corev1ac.ResourceRequirements().
|
|
WithRequests(corev1.ResourceList{
|
|
corev1.ResourceCPU: resource.MustParse("10"),
|
|
})),
|
|
).
|
|
WithInitContainers(
|
|
corev1ac.Container().
|
|
WithName("setup-initializer").
|
|
WithRestartPolicy(corev1.ContainerRestartPolicyAlways).
|
|
WithResources(corev1ac.ResourceRequirements().
|
|
WithRequests(corev1.ResourceList{
|
|
corev1.ResourceCPU: resource.MustParse("5"),
|
|
})),
|
|
),
|
|
),
|
|
WithPodSet(constants.ModelInitializer, ptr.To(constants.ModelInitializer), 1, corev1.PodSpec{
|
|
Containers: []corev1.Container{{
|
|
Name: constants.ModelInitializer,
|
|
Resources: corev1.ResourceRequirements{
|
|
Requests: corev1.ResourceList{
|
|
corev1.ResourceCPU: resource.MustParse("10"),
|
|
},
|
|
},
|
|
}},
|
|
InitContainers: []corev1.Container{{
|
|
Name: "setup-initializer",
|
|
RestartPolicy: ptr.To(corev1.ContainerRestartPolicyAlways),
|
|
Resources: corev1.ResourceRequirements{
|
|
Requests: corev1.ResourceList{
|
|
corev1.ResourceCPU: resource.MustParse("5"),
|
|
},
|
|
},
|
|
}},
|
|
}, corev1ac.PodSpec().
|
|
WithContainers(
|
|
corev1ac.Container().
|
|
WithName(constants.ModelInitializer).
|
|
WithResources(corev1ac.ResourceRequirements().
|
|
WithRequests(corev1.ResourceList{
|
|
corev1.ResourceCPU: resource.MustParse("10"),
|
|
})),
|
|
).
|
|
WithInitContainers(
|
|
corev1ac.Container().
|
|
WithName("setup-initializer").
|
|
WithRestartPolicy(corev1.ContainerRestartPolicyAlways).
|
|
WithResources(corev1ac.ResourceRequirements().
|
|
WithRequests(corev1.ResourceList{
|
|
corev1.ResourceCPU: resource.MustParse("5"),
|
|
})),
|
|
),
|
|
),
|
|
WithPodSet(constants.Node, ptr.To(constants.AncestorTrainer), 10, corev1.PodSpec{
|
|
Containers: []corev1.Container{{
|
|
Name: constants.Node,
|
|
Resources: corev1.ResourceRequirements{
|
|
Requests: corev1.ResourceList{
|
|
corev1.ResourceCPU: resource.MustParse("15"),
|
|
},
|
|
},
|
|
}},
|
|
InitContainers: []corev1.Container{{
|
|
Name: "preparation",
|
|
RestartPolicy: ptr.To(corev1.ContainerRestartPolicyAlways),
|
|
Resources: corev1.ResourceRequirements{
|
|
Requests: corev1.ResourceList{
|
|
corev1.ResourceCPU: resource.MustParse("25"),
|
|
},
|
|
},
|
|
}},
|
|
}, corev1ac.PodSpec().
|
|
WithContainers(
|
|
corev1ac.Container().
|
|
WithName(constants.Node).
|
|
WithResources(corev1ac.ResourceRequirements().
|
|
WithRequests(corev1.ResourceList{
|
|
corev1.ResourceCPU: resource.MustParse("25"),
|
|
})).
|
|
WithEnv(
|
|
corev1ac.EnvVar().WithName("TEST").WithValue("TEST"),
|
|
).
|
|
WithPorts(corev1ac.ContainerPort().
|
|
WithName("http").
|
|
WithProtocol(corev1.ProtocolTCP).
|
|
WithContainerPort(8080),
|
|
).
|
|
WithVolumeMounts(
|
|
corev1ac.VolumeMount().
|
|
WithName("TEST").
|
|
WithMountPath("/var").
|
|
WithReadOnly(true),
|
|
),
|
|
).
|
|
WithInitContainers(corev1ac.Container().
|
|
WithName("preparation").
|
|
WithResources(corev1ac.ResourceRequirements().
|
|
WithRequests(corev1.ResourceList{
|
|
corev1.ResourceCPU: resource.MustParse("15"),
|
|
})).
|
|
WithRestartPolicy(corev1.ContainerRestartPolicyAlways),
|
|
)),
|
|
WithTemplateSpecObjApply(
|
|
jobsetv1alpha2ac.JobSetSpec().
|
|
WithReplicatedJobs(
|
|
jobsetv1alpha2ac.ReplicatedJob().
|
|
WithName(constants.DatasetInitializer).
|
|
WithTemplate(batchv1ac.JobTemplateSpec().
|
|
WithLabels(nil).
|
|
WithSpec(batchv1ac.JobSpec().
|
|
WithTemplate(corev1ac.PodTemplateSpec().
|
|
WithLabels(nil).
|
|
WithSpec(corev1ac.PodSpec().
|
|
WithContainers(
|
|
corev1ac.Container().
|
|
WithName(constants.DatasetInitializer).
|
|
WithVolumeMounts(
|
|
corev1ac.VolumeMount().
|
|
WithName(jobsetplgconsts.VolumeNameInitializer).
|
|
WithMountPath(constants.DatasetMountPath),
|
|
).
|
|
WithResources(corev1ac.ResourceRequirements()),
|
|
).
|
|
WithVolumes(corev1ac.Volume().
|
|
WithName(jobsetplgconsts.VolumeNameInitializer).
|
|
WithPersistentVolumeClaim(corev1ac.PersistentVolumeClaimVolumeSource().
|
|
WithClaimName(jobsetplgconsts.VolumeNameInitializer),
|
|
),
|
|
),
|
|
),
|
|
),
|
|
),
|
|
),
|
|
jobsetv1alpha2ac.ReplicatedJob().
|
|
WithName(constants.ModelInitializer).
|
|
WithTemplate(batchv1ac.JobTemplateSpec().
|
|
WithLabels(nil).
|
|
WithSpec(batchv1ac.JobSpec().
|
|
WithTemplate(corev1ac.PodTemplateSpec().
|
|
WithLabels(nil).
|
|
WithSpec(corev1ac.PodSpec().
|
|
WithContainers(
|
|
corev1ac.Container().
|
|
WithName(constants.ModelInitializer).
|
|
WithVolumeMounts(
|
|
corev1ac.VolumeMount().
|
|
WithName(jobsetplgconsts.VolumeNameInitializer).
|
|
WithMountPath(constants.ModelMountPath),
|
|
).
|
|
WithResources(corev1ac.ResourceRequirements()),
|
|
).
|
|
WithVolumes(corev1ac.Volume().
|
|
WithName(jobsetplgconsts.VolumeNameInitializer).
|
|
WithPersistentVolumeClaim(corev1ac.PersistentVolumeClaimVolumeSource().
|
|
WithClaimName(jobsetplgconsts.VolumeNameInitializer),
|
|
),
|
|
),
|
|
),
|
|
),
|
|
),
|
|
),
|
|
jobsetv1alpha2ac.ReplicatedJob().
|
|
WithName(constants.Node).
|
|
WithTemplate(batchv1ac.JobTemplateSpec().
|
|
WithLabels(nil).
|
|
WithSpec(batchv1ac.JobSpec().
|
|
WithTemplate(corev1ac.PodTemplateSpec().
|
|
WithLabels(nil).
|
|
WithSpec(corev1ac.PodSpec().
|
|
WithContainers(
|
|
corev1ac.Container().
|
|
WithName(constants.Node).
|
|
WithVolumeMounts(
|
|
corev1ac.VolumeMount().
|
|
WithName(jobsetplgconsts.VolumeNameInitializer).
|
|
WithMountPath(constants.DatasetMountPath),
|
|
corev1ac.VolumeMount().
|
|
WithName(jobsetplgconsts.VolumeNameInitializer).
|
|
WithMountPath(constants.ModelMountPath),
|
|
).
|
|
WithResources(corev1ac.ResourceRequirements()),
|
|
).
|
|
WithVolumes(corev1ac.Volume().
|
|
WithName(jobsetplgconsts.VolumeNameInitializer).
|
|
WithPersistentVolumeClaim(corev1ac.PersistentVolumeClaimVolumeSource().
|
|
WithClaimName(jobsetplgconsts.VolumeNameInitializer),
|
|
),
|
|
),
|
|
),
|
|
),
|
|
),
|
|
),
|
|
),
|
|
),
|
|
},
|
|
wantInfo: &Info{
|
|
Labels: map[string]string{
|
|
"labelKey": "labelValue",
|
|
},
|
|
Annotations: map[string]string{
|
|
"annotationKey": "annotationValue",
|
|
},
|
|
Scheduler: &Scheduler{PodLabels: make(map[string]string)},
|
|
TemplateSpec: TemplateSpec{
|
|
PodSets: []PodSet{
|
|
{
|
|
Name: constants.DatasetInitializer,
|
|
Ancestor: ptr.To(constants.DatasetInitializer),
|
|
Count: ptr.To[int32](1),
|
|
InitContainers: []Container{{
|
|
Name: "setup-initializer",
|
|
}},
|
|
Containers: []Container{{
|
|
Name: constants.DatasetInitializer,
|
|
}},
|
|
SinglePodRequests: corev1.ResourceList{
|
|
corev1.ResourceCPU: resource.MustParse("15"),
|
|
},
|
|
},
|
|
{
|
|
Name: constants.ModelInitializer,
|
|
Ancestor: ptr.To(constants.ModelInitializer),
|
|
Count: ptr.To[int32](1),
|
|
InitContainers: []Container{{
|
|
Name: "setup-initializer",
|
|
}},
|
|
Containers: []Container{{
|
|
Name: constants.ModelInitializer,
|
|
}},
|
|
SinglePodRequests: corev1.ResourceList{
|
|
corev1.ResourceCPU: resource.MustParse("15"),
|
|
},
|
|
},
|
|
{
|
|
Name: constants.Node,
|
|
Ancestor: ptr.To(constants.AncestorTrainer),
|
|
Count: ptr.To[int32](10),
|
|
Containers: []Container{{
|
|
Name: constants.Node,
|
|
Env: []corev1ac.EnvVarApplyConfiguration{{
|
|
Name: ptr.To("TEST"),
|
|
Value: ptr.To("TEST"),
|
|
}},
|
|
Ports: []corev1ac.ContainerPortApplyConfiguration{{
|
|
Name: ptr.To("http"),
|
|
Protocol: ptr.To(corev1.ProtocolTCP),
|
|
ContainerPort: ptr.To[int32](8080),
|
|
}},
|
|
VolumeMounts: []corev1ac.VolumeMountApplyConfiguration{{
|
|
Name: ptr.To("TEST"),
|
|
ReadOnly: ptr.To(true),
|
|
MountPath: ptr.To("/var"),
|
|
}},
|
|
}},
|
|
InitContainers: []Container{{
|
|
Name: "preparation",
|
|
}},
|
|
SinglePodRequests: corev1.ResourceList{
|
|
corev1.ResourceCPU: resource.MustParse("40"),
|
|
},
|
|
},
|
|
},
|
|
ObjApply: jobsetv1alpha2ac.JobSetSpec().
|
|
WithReplicatedJobs(
|
|
jobsetv1alpha2ac.ReplicatedJob().
|
|
WithName(constants.DatasetInitializer).
|
|
WithTemplate(batchv1ac.JobTemplateSpec().
|
|
WithLabels(nil).
|
|
WithSpec(batchv1ac.JobSpec().
|
|
WithTemplate(corev1ac.PodTemplateSpec().
|
|
WithLabels(nil).
|
|
WithSpec(corev1ac.PodSpec().
|
|
WithContainers(
|
|
corev1ac.Container().
|
|
WithName(constants.DatasetInitializer).
|
|
WithVolumeMounts(
|
|
corev1ac.VolumeMount().
|
|
WithName(jobsetplgconsts.VolumeNameInitializer).
|
|
WithMountPath(constants.DatasetMountPath),
|
|
).
|
|
WithResources(corev1ac.ResourceRequirements()),
|
|
).
|
|
WithVolumes(corev1ac.Volume().
|
|
WithName(jobsetplgconsts.VolumeNameInitializer).
|
|
WithPersistentVolumeClaim(corev1ac.PersistentVolumeClaimVolumeSource().
|
|
WithClaimName(jobsetplgconsts.VolumeNameInitializer),
|
|
),
|
|
),
|
|
),
|
|
),
|
|
),
|
|
),
|
|
jobsetv1alpha2ac.ReplicatedJob().
|
|
WithName(constants.ModelInitializer).
|
|
WithTemplate(batchv1ac.JobTemplateSpec().
|
|
WithLabels(nil).
|
|
WithSpec(batchv1ac.JobSpec().
|
|
WithTemplate(corev1ac.PodTemplateSpec().
|
|
WithLabels(nil).
|
|
WithSpec(corev1ac.PodSpec().
|
|
WithContainers(
|
|
corev1ac.Container().
|
|
WithName(constants.ModelInitializer).
|
|
WithVolumeMounts(
|
|
corev1ac.VolumeMount().
|
|
WithName(jobsetplgconsts.VolumeNameInitializer).
|
|
WithMountPath(constants.ModelMountPath),
|
|
).
|
|
WithResources(corev1ac.ResourceRequirements()),
|
|
).
|
|
WithVolumes(corev1ac.Volume().
|
|
WithName(jobsetplgconsts.VolumeNameInitializer).
|
|
WithPersistentVolumeClaim(corev1ac.PersistentVolumeClaimVolumeSource().
|
|
WithClaimName(jobsetplgconsts.VolumeNameInitializer),
|
|
),
|
|
),
|
|
),
|
|
),
|
|
),
|
|
),
|
|
jobsetv1alpha2ac.ReplicatedJob().
|
|
WithName(constants.Node).
|
|
WithTemplate(batchv1ac.JobTemplateSpec().
|
|
WithLabels(nil).
|
|
WithSpec(batchv1ac.JobSpec().
|
|
WithTemplate(corev1ac.PodTemplateSpec().
|
|
WithLabels(nil).
|
|
WithSpec(corev1ac.PodSpec().
|
|
WithContainers(
|
|
corev1ac.Container().
|
|
WithName(constants.Node).
|
|
WithVolumeMounts(
|
|
corev1ac.VolumeMount().
|
|
WithName(jobsetplgconsts.VolumeNameInitializer).
|
|
WithMountPath(constants.DatasetMountPath),
|
|
corev1ac.VolumeMount().
|
|
WithName(jobsetplgconsts.VolumeNameInitializer).
|
|
WithMountPath(constants.ModelMountPath),
|
|
).
|
|
WithResources(corev1ac.ResourceRequirements()),
|
|
).
|
|
WithVolumes(corev1ac.Volume().
|
|
WithName(jobsetplgconsts.VolumeNameInitializer).
|
|
WithPersistentVolumeClaim(corev1ac.PersistentVolumeClaimVolumeSource().
|
|
WithClaimName(jobsetplgconsts.VolumeNameInitializer),
|
|
),
|
|
),
|
|
),
|
|
),
|
|
),
|
|
),
|
|
),
|
|
},
|
|
},
|
|
},
|
|
"all arguments are not specified": {
|
|
wantInfo: &Info{
|
|
Labels: make(map[string]string),
|
|
Annotations: make(map[string]string),
|
|
Scheduler: &Scheduler{PodLabels: make(map[string]string)},
|
|
},
|
|
},
|
|
}
|
|
cmpOpts := []cmp.Option{
|
|
cmpopts.SortMaps(func(a, b string) bool { return a < b }),
|
|
}
|
|
for name, tc := range cases {
|
|
t.Run(name, func(t *testing.T) {
|
|
info := NewInfo(tc.infoOpts...)
|
|
if diff := cmp.Diff(tc.wantInfo, info, cmpOpts...); len(diff) != 0 {
|
|
t.Errorf("Unexpected runtime.Info (-want,+got):\n%s", diff)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestFindContainerByPodSetAncestorContainerName(t *testing.T) {
|
|
cases := map[string]struct {
|
|
info *Info
|
|
psAncestor string
|
|
containerName string
|
|
wantContainer *Container
|
|
}{
|
|
"podSet and container exist": {
|
|
info: &Info{
|
|
TemplateSpec: TemplateSpec{
|
|
PodSets: []PodSet{
|
|
{
|
|
Ancestor: ptr.To("alpha"),
|
|
Containers: []Container{
|
|
{
|
|
Name: "one",
|
|
},
|
|
{
|
|
Name: "two",
|
|
},
|
|
},
|
|
},
|
|
{
|
|
Ancestor: ptr.To("beta"),
|
|
Containers: []Container{{Name: "one"}},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
psAncestor: "alpha",
|
|
containerName: "one",
|
|
wantContainer: &Container{
|
|
Name: "one",
|
|
},
|
|
},
|
|
"podSet exists, but container does not exist": {
|
|
info: &Info{
|
|
TemplateSpec: TemplateSpec{
|
|
PodSets: []PodSet{{
|
|
Ancestor: ptr.To("alpha"),
|
|
Containers: []Container{{Name: "one"}},
|
|
}},
|
|
},
|
|
},
|
|
psAncestor: "alpha",
|
|
containerName: "two",
|
|
wantContainer: nil,
|
|
},
|
|
"podSet does not exist": {
|
|
info: &Info{
|
|
TemplateSpec: TemplateSpec{
|
|
PodSets: []PodSet{{
|
|
Ancestor: ptr.To("alpha"),
|
|
Containers: []Container{{Name: "one"}},
|
|
}},
|
|
},
|
|
},
|
|
psAncestor: "beta",
|
|
containerName: "one",
|
|
wantContainer: nil,
|
|
},
|
|
}
|
|
for name, tc := range cases {
|
|
t.Run(name, func(t *testing.T) {
|
|
got := tc.info.FindContainerByPodSetAncestorContainerName(tc.psAncestor, tc.containerName)
|
|
if diff := cmp.Diff(tc.wantContainer, got); len(diff) != 0 {
|
|
t.Errorf("Unexpected Container (-want,+got):\n%s", diff)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestFindPodSetByName(t *testing.T) {
|
|
cases := map[string]struct {
|
|
info *Info
|
|
psName string
|
|
wantPodSet *PodSet
|
|
}{
|
|
"PodSet exists": {
|
|
info: &Info{
|
|
TemplateSpec: TemplateSpec{
|
|
PodSets: []PodSet{
|
|
{
|
|
Name: "alpha",
|
|
},
|
|
{
|
|
Name: "beta",
|
|
},
|
|
},
|
|
},
|
|
},
|
|
psName: "alpha",
|
|
wantPodSet: &PodSet{
|
|
Name: "alpha",
|
|
},
|
|
},
|
|
"PodSet does not exist": {
|
|
info: &Info{
|
|
TemplateSpec: TemplateSpec{
|
|
PodSets: []PodSet{
|
|
{
|
|
Name: "alpha",
|
|
},
|
|
},
|
|
},
|
|
},
|
|
psName: "beta",
|
|
wantPodSet: nil,
|
|
},
|
|
}
|
|
for name, tc := range cases {
|
|
t.Run(name, func(t *testing.T) {
|
|
got := tc.info.FindPodSetByName(tc.psName)
|
|
if diff := cmp.Diff(tc.wantPodSet, got); len(diff) != 0 {
|
|
t.Errorf("Unexpected PodSet (-want,+got):\n%s", diff)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestFindContainerByPodSetName(t *testing.T) {
|
|
cases := map[string]struct {
|
|
info *Info
|
|
psName string
|
|
containerName string
|
|
wantContainer *Container
|
|
}{
|
|
"regular container exists": {
|
|
info: &Info{TemplateSpec: TemplateSpec{PodSets: []PodSet{
|
|
{
|
|
Name: "node",
|
|
Containers: []Container{{Name: "trainer"}},
|
|
},
|
|
}}},
|
|
psName: "node",
|
|
containerName: "trainer",
|
|
wantContainer: &Container{Name: "trainer"},
|
|
},
|
|
"init container exists": {
|
|
info: &Info{TemplateSpec: TemplateSpec{PodSets: []PodSet{
|
|
{
|
|
Name: "node",
|
|
InitContainers: []Container{{Name: "preflight"}},
|
|
Containers: []Container{{Name: "trainer"}},
|
|
},
|
|
}}},
|
|
psName: "node",
|
|
containerName: "preflight",
|
|
wantContainer: &Container{Name: "preflight"},
|
|
},
|
|
"regular container wins name collision": {
|
|
info: &Info{TemplateSpec: TemplateSpec{PodSets: []PodSet{
|
|
{
|
|
Name: "node",
|
|
InitContainers: []Container{{Name: "shared", Image: "init"}},
|
|
Containers: []Container{{Name: "shared", Image: "main"}},
|
|
},
|
|
}}},
|
|
psName: "node",
|
|
containerName: "shared",
|
|
wantContainer: &Container{Name: "shared", Image: "main"},
|
|
},
|
|
"podSet does not exist": {
|
|
info: &Info{TemplateSpec: TemplateSpec{PodSets: []PodSet{
|
|
{Name: "node", Containers: []Container{{Name: "trainer"}}},
|
|
}}},
|
|
psName: "worker",
|
|
containerName: "trainer",
|
|
wantContainer: nil,
|
|
},
|
|
"container does not exist": {
|
|
info: &Info{TemplateSpec: TemplateSpec{PodSets: []PodSet{
|
|
{Name: "node", Containers: []Container{{Name: "trainer"}}},
|
|
}}},
|
|
psName: "node",
|
|
containerName: "preflight",
|
|
wantContainer: nil,
|
|
},
|
|
}
|
|
for name, tc := range cases {
|
|
t.Run(name, func(t *testing.T) {
|
|
got := tc.info.FindContainerByPodSetName(tc.psName, tc.containerName)
|
|
if diff := cmp.Diff(tc.wantContainer, got); len(diff) != 0 {
|
|
t.Errorf("Unexpected Container (-want,+got):\n%s", diff)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestFindContainerByName(t *testing.T) {
|
|
cases := map[string]struct {
|
|
info *Info
|
|
psAncestor string
|
|
want *PodSet
|
|
}{
|
|
"PodSet exists": {
|
|
info: &Info{
|
|
TemplateSpec: TemplateSpec{
|
|
PodSets: []PodSet{
|
|
{
|
|
Ancestor: ptr.To("alpha"),
|
|
},
|
|
{
|
|
Ancestor: ptr.To("beta"),
|
|
},
|
|
},
|
|
},
|
|
},
|
|
psAncestor: "alpha",
|
|
want: &PodSet{
|
|
Ancestor: ptr.To("alpha"),
|
|
},
|
|
},
|
|
"PodSet does not exist": {
|
|
info: &Info{
|
|
TemplateSpec: TemplateSpec{
|
|
PodSets: []PodSet{{
|
|
Name: "alpha",
|
|
}},
|
|
},
|
|
},
|
|
psAncestor: "beta",
|
|
want: nil,
|
|
},
|
|
}
|
|
for name, tc := range cases {
|
|
t.Run(name, func(t *testing.T) {
|
|
got := tc.info.FindPodSetByAncestor(tc.psAncestor)
|
|
if diff := cmp.Diff(tc.want, got); len(diff) != 0 {
|
|
t.Errorf("Unexpected PodSet (-want,+got):\n%s", diff)
|
|
}
|
|
})
|
|
}
|
|
}
|