diff --git a/master/internal/rm/kubernetesrm/jobs_test.go b/master/internal/rm/kubernetesrm/jobs_test.go index 590524e221aa..cc357fa02318 100644 --- a/master/internal/rm/kubernetesrm/jobs_test.go +++ b/master/internal/rm/kubernetesrm/jobs_test.go @@ -14,6 +14,7 @@ import ( _ "k8s.io/client-go/plugin/pkg/client/auth" "github.com/determined-ai/determined/master/internal/mocks" + "github.com/determined-ai/determined/master/pkg/model" ) func TestGetNonDetPods(t *testing.T) { @@ -168,6 +169,54 @@ func TestAllTaintsTolerated(t *testing.T) { } } +func TestNodeSelectorsAffinityMatch(t *testing.T) { + cases := []struct { + name string + tcd *model.TaskContainerDefaultsConfig + node *k8sV1.Node + match bool + }{ + {"no task containers default", nil, nil, false}, + {"no node labels", &model.TaskContainerDefaultsConfig{CPUPodSpec: &k8sV1.Pod{}}, nil, false}, + { + "gpu pod spec defined + match", + &model.TaskContainerDefaultsConfig{ + GPUPodSpec: &k8sV1.Pod{}, + }, &k8sV1.Node{}, true, + }, + {"cpu pod spec defined + match", &model.TaskContainerDefaultsConfig{ + CPUPodSpec: &k8sV1.Pod{}, + }, &k8sV1.Node{}, true}, + {"both gpu & cpu pod spec defined + match", &model.TaskContainerDefaultsConfig{ + CPUPodSpec: &k8sV1.Pod{}, + GPUPodSpec: &k8sV1.Pod{}, + }, &k8sV1.Node{}, true}, + {"incorrectly formatted task spec, error NewNodeSelector", &model.TaskContainerDefaultsConfig{ + CPUPodSpec: &k8sV1.Pod{}, + }, &k8sV1.Node{}, true}, + {"no match", nil, nil, false}, + } + + for _, tt := range cases { + t.Run(tt.name, func(t *testing.T) { + j := newTestJobsService(t) + match := j.podsCanBeScheduledOnNode(tt.tcd, tt.node) + require.Equal(t, tt.match, match) + }) + } +} + +func TestGetAgentsWithNodeSelectors(t *testing.T) { + cases := []struct { + name string + match bool + }{} + + for _, tt := range cases { + t.Run(tt.name, func(t *testing.T) {}) + } +} + var taintFooBar = k8sV1.Taint{ Key: "foo", Value: "bar",