DefaultPodAffinityFactory implements PodAffinityFactory { @Override public Pair<V1Affinity, Map<String, String>> buildV1Affinity(Job<?> job, Task task) { return new Processor(job).build(); } @Inject  DefaultPodAffinityFactory(DirectKubeConfiguration configuration,
                                     PodResourcePoolResolver podResourcePoolResolver); @Override Pair<V1Affinity, Map<String, String>> buildV1Affinity(Job<?> job, Task task);  }
@Test public void testInstanceTypeAffinity() { Pair<V1Affinity, Map<String, String>> affinityWithAnnotations = factory.buildV1Affinity( newJobWithHardConstraint(JobConstraints.MACHINE_TYPE.toUpperCase(), "r5.metal"), JobGenerator.oneBatchTask() ); V1NodeSelector nodeSelector = affinityWithAnnotations.getLeft().getNodeAffinity().getRequiredDuringSchedulingIgnoredDuringExecution(); assertThat(nodeSelector.getNodeSelectorTerms()).hasSize(1); }
@Test public void testKubeBackendAffinity() { Pair<V1Affinity, Map<String, String>> affinityWithAnnotations = factory.buildV1Affinity( newJobWithHardConstraint(JobConstraints.KUBE_BACKEND, "kublet"), JobGenerator.oneBatchTask() ); V1NodeSelector nodeSelector = affinityWithAnnotations.getLeft().getNodeAffinity().getRequiredDuringSchedulingIgnoredDuringExecution(); assertThat(nodeSelector.getNodeSelectorTerms()).hasSize(1); V1NodeSelectorRequirement requirement = nodeSelector.getNodeSelectorTerms().get(0).getMatchExpressions().get(0); assertThat(requirement.getKey()).isEqualTo(KubeConstants.TAINT_KUBE_BACKEND); assertThat(requirement.getValues()).contains("kublet"); }
@Test public void testEmptyInstanceTypeIsIgnored() { Pair<V1Affinity, Map<String, String>> affinityWithAnnotations = factory.buildV1Affinity( newJobWithHardConstraint(JobConstraints.MACHINE_TYPE, ""), JobGenerator.oneBatchTask() ); assertThat(affinityWithAnnotations.getLeft().getNodeAffinity()).isNull(); }
@Test public void testDefaultGpuInstanceAssignment() { Pair<V1Affinity, Map<String, String>> affinityWithAnnotations = factory.buildV1Affinity( newGpuJob(Collections.emptyMap()), JobGenerator.oneBatchTask() ); V1NodeSelector nodeSelector = affinityWithAnnotations.getLeft().getNodeAffinity().getRequiredDuringSchedulingIgnoredDuringExecution(); assertThat(nodeSelector.getNodeSelectorTerms()).hasSize(1); assertThat(nodeSelector.getNodeSelectorTerms().get(0).getMatchExpressions().get(0).getValues().get(0)).isEqualTo(DEFAULT_GPU_INSTANCE_TYPE); }
@Test public void testSpecificGpuInstanceAssignment() { Pair<V1Affinity, Map<String, String>> affinityWithAnnotations = factory.buildV1Affinity(newGpuJob(Collections.singletonMap( JobConstraints.MACHINE_TYPE, SPECIFIC_GPU_INSTANCE_TYPE )), JobGenerator.oneBatchTask()); V1NodeSelector nodeSelector = affinityWithAnnotations.getLeft().getNodeAffinity().getRequiredDuringSchedulingIgnoredDuringExecution(); assertThat(nodeSelector.getNodeSelectorTerms()).hasSize(1); assertThat(nodeSelector.getNodeSelectorTerms().get(0).getMatchExpressions().get(0).getValues().get(0)).isEqualTo(SPECIFIC_GPU_INSTANCE_TYPE); }
@Test public void testResourcePoolAffinity() { Job<BatchJobExt> job = JobGenerator.oneBatchJob(); job = job.toBuilder().withJobDescriptor(JobFunctions.appendJobDescriptorAttributes(job.getJobDescriptor(), Collections.singletonMap(JobAttributes.JOB_PARAMETER_RESOURCE_POOLS, "elastic")) ).build(); Pair<V1Affinity, Map<String, String>> affinityWithAnnotations = factory.buildV1Affinity(job, JobGenerator.oneBatchTask()); V1NodeSelector nodeSelector = affinityWithAnnotations.getLeft().getNodeAffinity().getRequiredDuringSchedulingIgnoredDuringExecution(); assertThat(nodeSelector.getNodeSelectorTerms()).hasSize(1); assertThat(nodeSelector.getNodeSelectorTerms().get(0).getMatchExpressions().get(0).getKey()).isEqualTo(KubeConstants.NODE_LABEL_RESOURCE_POOL); assertThat(nodeSelector.getNodeSelectorTerms().get(0).getMatchExpressions().get(0).getValues().get(0)).isEqualTo("elastic"); }