AgentManagementConstraint implements SystemConstraint { @Override public Result evaluate(TaskRequest taskRequest, VirtualMachineCurrentState targetVM, TaskTrackerState taskTrackerState) { V3QueueableTask v3QueueableTask = (V3QueueableTask) taskRequest; Optional<AgentInstance> instanceOpt = SchedulerUtils.findInstance(agentManagementService, schedulerConfiguration.getInstanceAttributeName(), targetVM); if (!instanceOpt.isPresent()) { return Failure.INSTANCE_NOT_FOUND.toResult(); } AgentInstance instance = instanceOpt.get(); String instanceGroupId = instance.getInstanceGroupId(); Optional<AgentInstanceGroup> instanceGroupOpt = agentManagementService.findInstanceGroup(instanceGroupId); if (!instanceGroupOpt.isPresent()) { return Failure.INSTANCE_GROUP_NOT_FOUND.toResult(); } AgentInstanceGroup instanceGroup = instanceGroupOpt.get(); Result instanceGroupEvaluationResult = evaluateInstanceGroup(instanceGroup); if (instanceGroupEvaluationResult != VALID) { return instanceGroupEvaluationResult; } Result instanceEvaluationResult = evaluateInstance(instance); if (instanceEvaluationResult != VALID) { return instanceEvaluationResult; } return evaluateTask(v3QueueableTask, instanceGroup, instance); } @Inject  AgentManagementConstraint(SchedulerConfiguration schedulerConfiguration,
                                     AgentManagementService agentManagementService,
                                     AgentStatusMonitor agentStatusMonitor); static boolean isAgentManagementConstraintReason(String reason); @Override String getName(); @Override Result evaluate(TaskRequest taskRequest, VirtualMachineCurrentState targetVM, TaskTrackerState taskTrackerState); static final String NAME; }
@Test public void instanceGroupNotFound() { AgentInstance instance = createAgentInstance(""); when(agentManagementService.findAgentInstance(INSTANCE_ID)).thenReturn(Optional.of(instance)); Result result = agentManagementConstraint.evaluate(createTaskRequest(), createVirtualMachineCurrentStateMock(INSTANCE_ID), mock(TaskTrackerState.class)); assertThat(result.isSuccessful()).isFalse(); assertThat(result.getFailureReason()).isEqualToIgnoringCase("Instance group not found"); }
@Test public void instanceGroupNotActive() { AgentInstance instance = createAgentInstance(INSTANCE_GROUP_ID); when(agentManagementService.findAgentInstance(INSTANCE_ID)).thenReturn(Optional.of(instance)); AgentInstanceGroup agentInstanceGroup = createAgentInstanceGroup(InstanceGroupLifecycleState.Inactive, Tier.Flex); when(agentManagementService.findInstanceGroup(INSTANCE_GROUP_ID)).thenReturn(Optional.of(agentInstanceGroup)); Result result = agentManagementConstraint.evaluate(createTaskRequest(), createVirtualMachineCurrentStateMock(INSTANCE_ID), mock(TaskTrackerState.class)); assertThat(result.isSuccessful()).isFalse(); assertThat(result.getFailureReason()).isEqualToIgnoringCase("Instance group is not active or phased out"); }
@Test public void instanceGroupSystemNoPlacement() { AgentInstance instance = createAgentInstance(INSTANCE_GROUP_ID); when(agentManagementService.findAgentInstance(INSTANCE_ID)).thenReturn(Optional.of(instance)); AgentInstanceGroup agentInstanceGroup = createAgentInstanceGroup(InstanceGroupLifecycleState.Active, Tier.Flex, Collections.singletonMap(SchedulerAttributes.SYSTEM_NO_PLACEMENT, "true")); when(agentManagementService.findInstanceGroup(INSTANCE_GROUP_ID)).thenReturn(Optional.of(agentInstanceGroup)); Result result = agentManagementConstraint.evaluate(createTaskRequest(), createVirtualMachineCurrentStateMock(INSTANCE_ID), mock(TaskTrackerState.class)); assertThat(result.isSuccessful()).isFalse(); assertThat(result.getFailureReason()).isEqualToIgnoringCase("Cannot place on instance group or agent instance due to systemNoPlacement attribute"); }
@Test public void instanceGroupNoPlacement() { AgentInstance instance = createAgentInstance(INSTANCE_GROUP_ID); when(agentManagementService.findAgentInstance(INSTANCE_ID)).thenReturn(Optional.of(instance)); AgentInstanceGroup agentInstanceGroup = createAgentInstanceGroup(InstanceGroupLifecycleState.Active, Tier.Flex, Collections.singletonMap(SchedulerAttributes.NO_PLACEMENT, "true")); when(agentManagementService.findInstanceGroup(INSTANCE_GROUP_ID)).thenReturn(Optional.of(agentInstanceGroup)); Result result = agentManagementConstraint.evaluate(createTaskRequest(), createVirtualMachineCurrentStateMock(INSTANCE_ID), mock(TaskTrackerState.class)); assertThat(result.isSuccessful()).isFalse(); assertThat(result.getFailureReason()).isEqualToIgnoringCase("Cannot place on instance group or agent instance due to noPlacement attribute"); }
@Test public void instanceGroupTierMismatch() { AgentInstance instance = createAgentInstance(INSTANCE_GROUP_ID); when(agentManagementService.findAgentInstance(INSTANCE_ID)).thenReturn(Optional.of(instance)); AgentInstanceGroup agentInstanceGroup = createAgentInstanceGroup(InstanceGroupLifecycleState.Active, Tier.Critical); when(agentManagementService.findInstanceGroup(INSTANCE_GROUP_ID)).thenReturn(Optional.of(agentInstanceGroup)); Result result = agentManagementConstraint.evaluate(createTaskRequest(), createVirtualMachineCurrentStateMock(INSTANCE_ID), mock(TaskTrackerState.class)); assertThat(result.isSuccessful()).isFalse(); assertThat(result.getFailureReason()).isEqualToIgnoringCase("Task cannot run on instance group tier"); }
@Test public void instanceGroupDoesNotHaveGpus() { AgentInstance instance = createAgentInstance(INSTANCE_GROUP_ID); when(agentManagementService.findAgentInstance(INSTANCE_ID)).thenReturn(Optional.of(instance)); AgentInstanceGroup agentInstanceGroup = createAgentInstanceGroup(InstanceGroupLifecycleState.Active, Tier.Flex); when(agentManagementService.findInstanceGroup(INSTANCE_GROUP_ID)).thenReturn(Optional.of(agentInstanceGroup)); JobDescriptor<BatchJobExt> descriptor = oneTaskBatchJobDescriptor() .but(jd -> jd.getContainer().but(c -> c.toBuilder().withContainerResources( c.getContainerResources().toBuilder().withGpu(1).build() ))); Job<BatchJobExt> gpuBatchJob = JobGenerator.batchJobs(descriptor).getValue(); TaskRequest taskRequest = createTaskRequest(gpuBatchJob); Result result = agentManagementConstraint.evaluate(taskRequest, createVirtualMachineCurrentStateMock(INSTANCE_ID), mock(TaskTrackerState.class)); assertThat(result.isSuccessful()).isFalse(); assertThat(result.getFailureReason()).isEqualToIgnoringCase("Instance group does not have gpus"); }
@Test public void instanceOnlyRunsGpuTasks() { AgentInstance instance = createAgentInstance(INSTANCE_GROUP_ID); when(agentManagementService.findAgentInstance(INSTANCE_ID)).thenReturn(Optional.of(instance)); AgentInstanceGroup agentInstanceGroup = createAgentInstanceGroup(InstanceGroupLifecycleState.Active, Tier.Flex, 1); when(agentManagementService.findInstanceGroup(INSTANCE_GROUP_ID)).thenReturn(Optional.of(agentInstanceGroup)); JobDescriptor<BatchJobExt> descriptor = oneTaskBatchJobDescriptor() .but(jd -> jd.getContainer().but(c -> c.toBuilder().withContainerResources( c.getContainerResources().toBuilder().withGpu(0).build() ))); Job<BatchJobExt> gpuBatchJob = JobGenerator.batchJobs(descriptor).getValue(); TaskRequest taskRequest = createTaskRequest(gpuBatchJob); Result result = agentManagementConstraint.evaluate(taskRequest, createVirtualMachineCurrentStateMock(INSTANCE_ID), mock(TaskTrackerState.class)); assertThat(result.isSuccessful()).isFalse(); assertThat(result.getFailureReason()).isEqualToIgnoringCase("Instance group does not run non gpu tasks"); }
@Test public void taintsAndInstanceGroupTolerationsMatch() { AgentInstance instance = createAgentInstance(INSTANCE_GROUP_ID); when(agentManagementService.findAgentInstance(INSTANCE_ID)).thenReturn(Optional.of(instance)); AgentInstanceGroup agentInstanceGroup = createAgentInstanceGroup(InstanceGroupLifecycleState.Active, Tier.Flex, Collections.singletonMap(SchedulerAttributes.TAINTS, "b,c")); when(agentManagementService.findInstanceGroup(INSTANCE_GROUP_ID)).thenReturn(Optional.of(agentInstanceGroup)); JobDescriptor<BatchJobExt> descriptor = oneTaskBatchJobDescriptor() .but(jd -> jd.toBuilder().withAttributes(Collections.singletonMap(JOB_PARAMETER_ATTRIBUTES_TOLERATIONS, "a,b,c"))); Job<BatchJobExt> job = JobGenerator.batchJobs(descriptor).getValue(); TaskRequest taskRequest = createTaskRequest(job); Result result = agentManagementConstraint.evaluate(taskRequest, createVirtualMachineCurrentStateMock(INSTANCE_ID), mock(TaskTrackerState.class)); assertThat(result.isSuccessful()).isTrue(); }
@Test public void taintsAndInstanceTolerationsMatch() { AgentInstance instance = createAgentInstance(INSTANCE_GROUP_ID, Collections.singletonMap(SchedulerAttributes.TAINTS, "b,c")); when(agentManagementService.findAgentInstance(INSTANCE_ID)).thenReturn(Optional.of(instance)); AgentInstanceGroup agentInstanceGroup = createAgentInstanceGroup(InstanceGroupLifecycleState.Active, Tier.Flex); when(agentManagementService.findInstanceGroup(INSTANCE_GROUP_ID)).thenReturn(Optional.of(agentInstanceGroup)); JobDescriptor<BatchJobExt> descriptor = oneTaskBatchJobDescriptor() .but(jd -> jd.toBuilder().withAttributes(Collections.singletonMap(JOB_PARAMETER_ATTRIBUTES_TOLERATIONS, "a,b,c"))); Job<BatchJobExt> job = JobGenerator.batchJobs(descriptor).getValue(); TaskRequest taskRequest = createTaskRequest(job); Result result = agentManagementConstraint.evaluate(taskRequest, createVirtualMachineCurrentStateMock(INSTANCE_ID), mock(TaskTrackerState.class)); assertThat(result.isSuccessful()).isTrue(); }
@Test public void taintsAndTolerationsDoNotMatch() { AgentInstance instance = createAgentInstance(INSTANCE_GROUP_ID, Collections.singletonMap(SchedulerAttributes.TAINTS, "b,d")); when(agentManagementService.findAgentInstance(INSTANCE_ID)).thenReturn(Optional.of(instance)); AgentInstanceGroup agentInstanceGroup = createAgentInstanceGroup(InstanceGroupLifecycleState.Active, Tier.Flex, Collections.singletonMap(SchedulerAttributes.TAINTS, "b,c")); when(agentManagementService.findInstanceGroup(INSTANCE_GROUP_ID)).thenReturn(Optional.of(agentInstanceGroup)); JobDescriptor<BatchJobExt> descriptor = oneTaskBatchJobDescriptor() .but(jd -> jd.toBuilder().withAttributes(Collections.singletonMap(JOB_PARAMETER_ATTRIBUTES_TOLERATIONS, "a,b"))); Job<BatchJobExt> job = JobGenerator.batchJobs(descriptor).getValue(); TaskRequest taskRequest = createTaskRequest(job); Result result = agentManagementConstraint.evaluate(taskRequest, createVirtualMachineCurrentStateMock(INSTANCE_ID), mock(TaskTrackerState.class)); assertThat(result.isSuccessful()).isFalse(); assertThat(result.getFailureReason()).isEqualToIgnoringCase("Cannot place on instance group or agent instance due to toleration attribute not matching taint attribute"); }