Fix ReplicaType key for filtering pods and services (#152)

kubeflow · Aug 12, 2021 · 7d3b7d9 · 7d3b7d9
1 parent 6c9fe11
commit 7d3b7d9
Show file tree

Hide file tree

Showing 4 changed files with 91 additions and 2 deletions.
diff --git a/pkg/controller.v1/common/pod.go b/pkg/controller.v1/common/pod.go
@@ -258,7 +258,7 @@ func (jc *JobController) FilterPodsForReplicaType(pods []*v1.Pod, replicaType ap
  var result []*v1.Pod
 
  selector := labels.SelectorFromValidatedSet(labels.Set{
- apiv1.ReplicaIndexLabel: string(replicaType),
+ apiv1.ReplicaTypeLabel: string(replicaType),
  })
 
  // TODO(#149): Remove deprecated selector.

diff --git a/pkg/controller.v1/common/pod_test.go b/pkg/controller.v1/common/pod_test.go
@@ -146,3 +146,48 @@ func TestCalculatePodSliceSize(t *testing.T) {
  assert.Equal(t, tc.expectedSize, result)
  }
 }
+
+func TestFilterPodsForReplicaType(t *testing.T) {
+ pods := []*v1.Pod{
+ {
+ ObjectMeta: metav1.ObjectMeta{
+ Name: "a",
+ Labels: map[string]string{apiv1.ReplicaTypeLabel: "foo"},
+ },
+ },
+ {
+ ObjectMeta: metav1.ObjectMeta{
+ Name: "b",
+ Labels: map[string]string{apiv1.ReplicaTypeLabel: "bar"},
+ },
+ },
+ {
+ ObjectMeta: metav1.ObjectMeta{
+ Name: "c",
+ Labels: map[string]string{apiv1.ReplicaTypeLabelDeprecated: "foo"},
+ },
+ },
+ {
+ ObjectMeta: metav1.ObjectMeta{
+ Name: "d",
+ Labels: map[string]string{apiv1.ReplicaTypeLabelDeprecated: "bar"},
+ },
+ },
+ {
+ ObjectMeta: metav1.ObjectMeta{
+ Name: "e",
+ Labels: map[string]string{
+ apiv1.ReplicaTypeLabel: "foo",
+ apiv1.ReplicaTypeLabelDeprecated: "bar",
+ },
+ },
+ },
+ }
+ c := &JobController{}
+ got, err := c.FilterPodsForReplicaType(pods, "foo")
+ if err != nil {
+ t.Fatalf("FilterPodsForReplicaType returned error: %v", err)
+ }
+ want := []*v1.Pod{pods[0], pods[2], pods[4]}
+ assert.Equal(t, want, got)
+}
diff --git a/pkg/controller.v1/common/service.go b/pkg/controller.v1/common/service.go
@@ -142,7 +142,7 @@ func (jc *JobController) FilterServicesForReplicaType(services []*v1.Service, re
  var result []*v1.Service
 
  selector := labels.SelectorFromValidatedSet(labels.Set{
- apiv1.ReplicaIndexLabel: string(replicaType),
+ apiv1.ReplicaTypeLabel: string(replicaType),
  })
 
  // TODO(#149): Remove deprecated selector.

diff --git a/pkg/controller.v1/common/service_test.go b/pkg/controller.v1/common/service_test.go
@@ -208,3 +208,47 @@ func TestJobController_CreateNewService(t *testing.T) {
  })
  }
 }
+func TestFilterServicesForReplicaType(t *testing.T) {
+ services := []*v1.Service{
+ {
+ ObjectMeta: metav1.ObjectMeta{
+ Name: "a",
+ Labels: map[string]string{apiv1.ReplicaTypeLabel: "foo"},
+ },
+ },
+ {
+ ObjectMeta: metav1.ObjectMeta{
+ Name: "b",
+ Labels: map[string]string{apiv1.ReplicaTypeLabel: "bar"},
+ },
+ },
+ {
+ ObjectMeta: metav1.ObjectMeta{
+ Name: "c",
+ Labels: map[string]string{apiv1.ReplicaTypeLabelDeprecated: "foo"},
+ },
+ },
+ {
+ ObjectMeta: metav1.ObjectMeta{
+ Name: "d",
+ Labels: map[string]string{apiv1.ReplicaTypeLabelDeprecated: "bar"},
+ },
+ },
+ {
+ ObjectMeta: metav1.ObjectMeta{
+ Name: "e",
+ Labels: map[string]string{
+ apiv1.ReplicaTypeLabel: "foo",
+ apiv1.ReplicaTypeLabelDeprecated: "bar",
+ },
+ },
+ },
+ }
+ c := &JobController{}
+ got, err := c.FilterServicesForReplicaType(services, "foo")
+ if err != nil {
+ t.Fatalf("FilterPodsForReplicaType returned error: %v", err)
+ }
+ want := []*v1.Service{services[0], services[2], services[4]}
+ assert.Equal(t, want, got)
+}