Skip to content

Commit

Permalink
Merge pull request #313 from Edwinhr716/wrapper-refactor
Browse files Browse the repository at this point in the history 
Refactored wrappers into its own package
  • Loading branch information
@k8s-ci-robot
k8s-ci-robot authored Jan 9, 2025
2 parents 08c92c6 + 6098516 commit 47a918d
Show file tree
Hide file tree
Showing 11 changed files with 372 additions and 484 deletions.
48 changes: 24 additions & 24 deletions pkg/controllers/leaderworkerset_controller_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -34,22 +34,22 @@ import (

"sigs.k8s.io/controller-runtime/pkg/client/fake"
revisionutils "sigs.k8s.io/lws/pkg/utils/revision"
testutils "sigs.k8s.io/lws/test/testutils"
"sigs.k8s.io/lws/test/wrappers"
)

func TestLeaderStatefulSetApplyConfig(t *testing.T) {
client := fake.NewClientBuilder().Build()
lws1 := testutils.BuildBasicLeaderWorkerSet("test-sample", "default").
LeaderTemplateSpec(testutils.MakeLeaderPodSpec()).
WorkerTemplateSpec(testutils.MakeWorkerPodSpec()).Obj()
lws1 := wrappers.BuildBasicLeaderWorkerSet("test-sample", "default").
LeaderTemplateSpec(wrappers.MakeLeaderPodSpec()).
WorkerTemplateSpec(wrappers.MakeWorkerPodSpec()).Obj()
cr1, err := revisionutils.NewRevision(context.TODO(), client, lws1, "")
if err != nil {
t.Fatal(err)
}
revisionKey1 := revisionutils.GetRevisionKey(cr1)

lws2 := testutils.BuildBasicLeaderWorkerSet("test-sample", "default").
WorkerTemplateSpec(testutils.MakeWorkerPodSpec()).Obj()
lws2 := wrappers.BuildBasicLeaderWorkerSet("test-sample", "default").
WorkerTemplateSpec(wrappers.MakeWorkerPodSpec()).Obj()
cr2, err := revisionutils.NewRevision(context.TODO(), client, lws2, "")
if err != nil {
t.Fatal(err)
Expand All @@ -65,15 +65,15 @@ func TestLeaderStatefulSetApplyConfig(t *testing.T) {
{
name: "1 replica, size 1, with empty leader template, exclusive placement disabled",
revisionKey: revisionKey2,
lws: testutils.BuildBasicLeaderWorkerSet("test-sample", "default").
lws: wrappers.BuildBasicLeaderWorkerSet("test-sample", "default").
Replica(1).
RolloutStrategy(leaderworkerset.RolloutStrategy{
Type: leaderworkerset.RollingUpdateStrategyType,
RollingUpdateConfiguration: &leaderworkerset.RollingUpdateConfiguration{
MaxUnavailable: intstr.FromInt32(1),
},
}).
WorkerTemplateSpec(testutils.MakeWorkerPodSpec()).
WorkerTemplateSpec(wrappers.MakeWorkerPodSpec()).
Size(1).
RestartPolicy(leaderworkerset.RecreateGroupOnPodRestart).Obj(),
wantApplyConfig: &appsapplyv1.StatefulSetApplyConfiguration{
Expand Down Expand Up @@ -131,7 +131,7 @@ func TestLeaderStatefulSetApplyConfig(t *testing.T) {
{
name: "1 replica, size 2 , with empty leader template, exclusive placement enabled",
revisionKey: revisionKey2,
lws: testutils.BuildBasicLeaderWorkerSet("test-sample", "default").
lws: wrappers.BuildBasicLeaderWorkerSet("test-sample", "default").
Annotation(map[string]string{
"leaderworkerset.sigs.k8s.io/exclusive-topology": "topologyKey",
}).Replica(1).
Expand All @@ -141,7 +141,7 @@ func TestLeaderStatefulSetApplyConfig(t *testing.T) {
MaxUnavailable: intstr.FromInt32(1),
},
}).
WorkerTemplateSpec(testutils.MakeWorkerPodSpec()).
WorkerTemplateSpec(wrappers.MakeWorkerPodSpec()).
Size(2).
RestartPolicy(leaderworkerset.RecreateGroupOnPodRestart).Obj(),
wantApplyConfig: &appsapplyv1.StatefulSetApplyConfiguration{
Expand Down Expand Up @@ -200,7 +200,7 @@ func TestLeaderStatefulSetApplyConfig(t *testing.T) {
{
name: "2 replica, size 2, with leader template, exclusive placement enabled",
revisionKey: revisionKey1,
lws: testutils.BuildBasicLeaderWorkerSet("test-sample", "default").Annotation(map[string]string{
lws: wrappers.BuildBasicLeaderWorkerSet("test-sample", "default").Annotation(map[string]string{
"leaderworkerset.sigs.k8s.io/exclusive-topology": "topologyKey",
}).Replica(2).
RolloutStrategy(leaderworkerset.RolloutStrategy{
Expand All @@ -209,8 +209,8 @@ func TestLeaderStatefulSetApplyConfig(t *testing.T) {
MaxUnavailable: intstr.FromInt32(1),
},
}).
WorkerTemplateSpec(testutils.MakeWorkerPodSpec()).
LeaderTemplateSpec(testutils.MakeLeaderPodSpec()).
WorkerTemplateSpec(wrappers.MakeWorkerPodSpec()).
LeaderTemplateSpec(wrappers.MakeLeaderPodSpec()).
Size(2).
RestartPolicy(leaderworkerset.RecreateGroupOnPodRestart).Obj(),
wantApplyConfig: &appsapplyv1.StatefulSetApplyConfiguration{
Expand Down Expand Up @@ -268,7 +268,7 @@ func TestLeaderStatefulSetApplyConfig(t *testing.T) {
{
name: "2 maxUnavailable, 1 maxSurge, with empty leader template, exclusive placement disabled",
revisionKey: revisionKey2,
lws: testutils.BuildBasicLeaderWorkerSet("test-sample", "default").
lws: wrappers.BuildBasicLeaderWorkerSet("test-sample", "default").
Replica(1).
RolloutStrategy(leaderworkerset.RolloutStrategy{
Type: leaderworkerset.RollingUpdateStrategyType,
Expand All @@ -277,7 +277,7 @@ func TestLeaderStatefulSetApplyConfig(t *testing.T) {
MaxSurge: intstr.FromInt32(1),
},
}).
WorkerTemplateSpec(testutils.MakeWorkerPodSpec()).
WorkerTemplateSpec(wrappers.MakeWorkerPodSpec()).
Size(1).
RestartPolicy(leaderworkerset.RecreateGroupOnPodRestart).Obj(),
wantApplyConfig: &appsapplyv1.StatefulSetApplyConfiguration{
Expand Down Expand Up @@ -335,7 +335,7 @@ func TestLeaderStatefulSetApplyConfig(t *testing.T) {
{
name: "1 replica, size 2, with leader template, exclusive placement enabled, subgroupsize enabled",
revisionKey: revisionKey1,
lws: testutils.BuildBasicLeaderWorkerSet("test-sample", "default").Annotation(map[string]string{
lws: wrappers.BuildBasicLeaderWorkerSet("test-sample", "default").Annotation(map[string]string{
leaderworkerset.SubGroupExclusiveKeyAnnotationKey: "topologyKey",
}).SubGroupSize(2).Replica(1).
RolloutStrategy(leaderworkerset.RolloutStrategy{
Expand All @@ -344,8 +344,8 @@ func TestLeaderStatefulSetApplyConfig(t *testing.T) {
MaxUnavailable: intstr.FromInt32(1),
},
}).
WorkerTemplateSpec(testutils.MakeWorkerPodSpec()).
LeaderTemplateSpec(testutils.MakeLeaderPodSpec()).
WorkerTemplateSpec(wrappers.MakeWorkerPodSpec()).
LeaderTemplateSpec(wrappers.MakeLeaderPodSpec()).
Size(2).
RestartPolicy(leaderworkerset.RecreateGroupOnPodRestart).Obj(),
wantApplyConfig: &appsapplyv1.StatefulSetApplyConfiguration{
Expand Down Expand Up @@ -483,44 +483,44 @@ func TestSetCondition(t *testing.T) {
{
name: "Different condition type, same condition status",
condition: metav1.Condition{Type: "Progressing", Status: "True"},
lws: testutils.BuildBasicLeaderWorkerSet("test-sample", "default").
lws: wrappers.BuildBasicLeaderWorkerSet("test-sample", "default").
Conditions([]metav1.Condition{{Type: "Available", Status: "True"}}).
Obj(),
expectedShouldUpdate: true,
},
{
name: "Same condition type, different condition status",
condition: metav1.Condition{Type: "Progressing", Status: "True"},
lws: testutils.BuildBasicLeaderWorkerSet("test-sample", "default").
lws: wrappers.BuildBasicLeaderWorkerSet("test-sample", "default").
Conditions([]metav1.Condition{{Type: "Progressing", Status: "False"}}).
Obj(),
expectedShouldUpdate: true,
},
{
name: "Different conditio type, new condition status is true",
condition: metav1.Condition{Type: "Progressing", Status: "True"},
lws: testutils.BuildBasicLeaderWorkerSet("test-sample", "default").
lws: wrappers.BuildBasicLeaderWorkerSet("test-sample", "default").
Conditions([]metav1.Condition{{Type: "Available", Status: "False"}}).
Obj(),
expectedShouldUpdate: true,
},
{
name: "No initial condition",
condition: metav1.Condition{Type: "Progressing", Status: "True"},
lws: testutils.BuildBasicLeaderWorkerSet("test-sample", "default").Obj(),
lws: wrappers.BuildBasicLeaderWorkerSet("test-sample", "default").Obj(),
expectedShouldUpdate: true,
},
{
name: "Different condition type, new condition status is false",
condition: metav1.Condition{Type: "Progressing", Status: "False"},
lws: testutils.BuildBasicLeaderWorkerSet("test-sample", "default").
lws: wrappers.BuildBasicLeaderWorkerSet("test-sample", "default").
Conditions([]metav1.Condition{{Type: "Available", Status: "True"}}).
Obj(),
},
{
name: "Same condition type, Same condition status",
condition: metav1.Condition{Type: "Progressing", Status: "False"},
lws: testutils.BuildBasicLeaderWorkerSet("test-sample", "default").
lws: wrappers.BuildBasicLeaderWorkerSet("test-sample", "default").
Conditions([]metav1.Condition{{Type: "Progressing", Status: "False"}}).
Obj(),
},
Expand Down
16 changes: 8 additions & 8 deletions pkg/controllers/pod_controller_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -31,13 +31,13 @@ import (
"sigs.k8s.io/controller-runtime/pkg/client/fake"
leaderworkerset "sigs.k8s.io/lws/api/leaderworkerset/v1"
revisionutils "sigs.k8s.io/lws/pkg/utils/revision"
testutils "sigs.k8s.io/lws/test/testutils"
"sigs.k8s.io/lws/test/wrappers"
)

func TestConstructWorkerStatefulSetApplyConfiguration(t *testing.T) {
client := fake.NewClientBuilder().Build()

lws := testutils.BuildBasicLeaderWorkerSet("test-sample", "default").Replica(1).WorkerTemplateSpec(testutils.MakeWorkerPodSpec()).Size(1).Obj()
lws := wrappers.BuildBasicLeaderWorkerSet("test-sample", "default").Replica(1).WorkerTemplateSpec(wrappers.MakeWorkerPodSpec()).Size(1).Obj()
updateRevision, err := revisionutils.NewRevision(context.TODO(), client, lws, "")
if err != nil {
t.Fatal(err)
Expand Down Expand Up @@ -67,9 +67,9 @@ func TestConstructWorkerStatefulSetApplyConfiguration(t *testing.T) {
},
},
},
lws: testutils.BuildBasicLeaderWorkerSet("test-sample", "default").
lws: wrappers.BuildBasicLeaderWorkerSet("test-sample", "default").
Replica(1).
WorkerTemplateSpec(testutils.MakeWorkerPodSpec()).
WorkerTemplateSpec(wrappers.MakeWorkerPodSpec()).
Size(1).Obj(),
wantStatefulSetConfig: &appsapplyv1.StatefulSetApplyConfiguration{
TypeMetaApplyConfiguration: metaapplyv1.TypeMetaApplyConfiguration{
Expand Down Expand Up @@ -141,9 +141,9 @@ func TestConstructWorkerStatefulSetApplyConfiguration(t *testing.T) {
},
},
},
lws: testutils.BuildBasicLeaderWorkerSet("test-sample", "default").
lws: wrappers.BuildBasicLeaderWorkerSet("test-sample", "default").
Replica(1).
WorkerTemplateSpec(testutils.MakeWorkerPodSpec()).
WorkerTemplateSpec(wrappers.MakeWorkerPodSpec()).
Annotation(map[string]string{
"leaderworkerset.sigs.k8s.io/exclusive-topology": "topologyKey",
}).Size(2).Obj(),
Expand Down Expand Up @@ -218,9 +218,9 @@ func TestConstructWorkerStatefulSetApplyConfiguration(t *testing.T) {
},
},
},
lws: testutils.BuildBasicLeaderWorkerSet("test-sample", "default").
lws: wrappers.BuildBasicLeaderWorkerSet("test-sample", "default").
Replica(1).
WorkerTemplateSpec(testutils.MakeWorkerPodSpec()).
WorkerTemplateSpec(wrappers.MakeWorkerPodSpec()).
Annotation(map[string]string{
leaderworkerset.SubGroupExclusiveKeyAnnotationKey: "topologyKey",
}).Size(2).SubGroupSize(2).Obj(),
Expand Down
72 changes: 9 additions & 63 deletions pkg/utils/accelerators/tpu_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@ import (
"github.com/google/go-cmp/cmp"
"k8s.io/apimachinery/pkg/api/resource"
leaderworkerset "sigs.k8s.io/lws/api/leaderworkerset/v1"
"sigs.k8s.io/lws/test/wrappers"

corev1 "k8s.io/api/core/v1"
v1 "k8s.io/apimachinery/pkg/apis/meta/v1"
Expand All @@ -40,7 +41,7 @@ func TestAddTPUVariables(t *testing.T) {
{
name: "Worker Index is 0",
pod: &corev1.Pod{
Spec: MakeLeaderPodSpecWithTPUResource(),
Spec: wrappers.MakeLeaderPodSpecWithTPUResource(),
ObjectMeta: v1.ObjectMeta{
Name: "test-sample-1",
Namespace: "default",
Expand All @@ -58,7 +59,7 @@ func TestAddTPUVariables(t *testing.T) {
{
name: "Worker Index is non-zero, size is above 2",
pod: &corev1.Pod{
Spec: MakeLeaderPodSpecWithTPUResource(),
Spec: wrappers.MakeLeaderPodSpecWithTPUResource(),
ObjectMeta: v1.ObjectMeta{
Name: "test-sample-1-3",
Namespace: "default",
Expand Down Expand Up @@ -117,7 +118,7 @@ func TestAddTPUVariablesSubGroup(t *testing.T) {
{
name: "Leader requests TPU resources",
pod: &corev1.Pod{
Spec: MakeLeaderPodSpecWithTPUResource(),
Spec: wrappers.MakeLeaderPodSpecWithTPUResource(),
ObjectMeta: v1.ObjectMeta{
Name: "test-sample-1-3",
Namespace: "default",
Expand All @@ -138,7 +139,7 @@ func TestAddTPUVariablesSubGroup(t *testing.T) {
{
name: "Leader requests TPU resources, worker with subgroup index > 0",
pod: &corev1.Pod{
Spec: MakeLeaderPodSpecWithTPUResource(),
Spec: wrappers.MakeLeaderPodSpecWithTPUResource(),
ObjectMeta: v1.ObjectMeta{
Name: "test-sample-1-7",
Namespace: "default",
Expand All @@ -159,7 +160,7 @@ func TestAddTPUVariablesSubGroup(t *testing.T) {
{
name: "Leader does not request TPU resources, worker with subgroup index > 0",
pod: &corev1.Pod{
Spec: MakeLeaderPodSpecWithTPUResource(),
Spec: wrappers.MakeLeaderPodSpecWithTPUResource(),
ObjectMeta: v1.ObjectMeta{
Name: "test-sample-1-5",
Namespace: "default",
Expand Down Expand Up @@ -204,7 +205,7 @@ func TestGetContainerRequestingTPUs(t *testing.T) {
}{
{
name: "Single Container with TPU Resource",
podSpec: MakeLeaderPodSpecWithTPUResource(),
podSpec: wrappers.MakeLeaderPodSpecWithTPUResource(),
expectedContainer: &corev1.Container{
Name: "worker",
Image: "busybox",
Expand All @@ -217,7 +218,7 @@ func TestGetContainerRequestingTPUs(t *testing.T) {
},
{
name: "Multiple Containers, one with TPU Resource",
podSpec: MakeLeaderPodSpecWithTPUResourceMultipleContainers(),
podSpec: wrappers.MakeLeaderPodSpecWithTPUResourceMultipleContainers(),
expectedContainer: &corev1.Container{
Name: "worker",
Image: "busybox",
Expand All @@ -230,7 +231,7 @@ func TestGetContainerRequestingTPUs(t *testing.T) {
},
{
name: "Container without TPU Resource",
podSpec: MakeLeaderPodSpec(),
podSpec: wrappers.MakeLeaderPodSpec(),
expectedContainer: nil,
},
}
Expand All @@ -243,58 +244,3 @@ func TestGetContainerRequestingTPUs(t *testing.T) {
})
}
}

func MakeLeaderPodSpec() corev1.PodSpec {
return corev1.PodSpec{
Containers: []corev1.Container{
{
Name: "worker",
Image: "busybox",
},
},
}
}

func MakeLeaderPodSpecWithTPUResource() corev1.PodSpec {
return corev1.PodSpec{
Containers: []corev1.Container{
{
Name: "worker",
Image: "busybox",
Resources: corev1.ResourceRequirements{
Limits: map[corev1.ResourceName]resource.Quantity{
corev1.ResourceName("google.com/tpu"): resource.MustParse("4"),
},
},
},
},
Subdomain: "default",
}
}

func MakeLeaderPodSpecWithTPUResourceMultipleContainers() corev1.PodSpec {
return corev1.PodSpec{
Containers: []corev1.Container{
{
Name: "worker",
Image: "busybox",
Resources: corev1.ResourceRequirements{
Limits: map[corev1.ResourceName]resource.Quantity{
corev1.ResourceName("google.com/tpu"): resource.MustParse("4"),
},
},
},
{
Name: "leader",
Image: "nginx:1.14.2",
Ports: []corev1.ContainerPort{
{
ContainerPort: 8080,
Protocol: "TCP",
},
},
},
},
Subdomain: "default",
}
}
Loading

0 comments on commit 47a918d

Please sign in to comment.