TaskMigrationDescheduler { Optional<Pair<Node, List<Task>>> nextBestMatch() { if (evictionQuotaTracker.getSystemEvictionQuota() <= 0) { return Optional.empty(); } return evacuatedAgentsAllocationTracker.getRemovableAgentsById().values().stream() .map(i -> Pair.of(i, computeFitness(i))) .filter(p -> p.getRight().getLeft() > 0) .max(Comparator.comparingDouble(p -> p.getRight().getLeft())) .map(p -> { Node agent = p.getLeft(); List<Task> tasks = p.getRight().getRight(); tasks.forEach(task -> { evacuatedAgentsAllocationTracker.descheduled(task); evictionQuotaTracker.consumeQuota(task.getJobId()); }); return Pair.of(agent, tasks); }); } TaskMigrationDescheduler(Map<String, TaskRelocationPlan> plannedAheadTaskRelocationPlans,
                             EvacuatedAgentsAllocationTracker evacuatedAgentsAllocationTracker,
                             EvictionQuotaTracker evictionQuotaTracker,
                             Map<String, Job<?>> jobsById,
                             Map<String, Task> tasksById,
                             TitusRuntime titusRuntime);   }
@Test public void testSelfMigration() { Task job1Task0 = jobOperations.getTasks("job1").get(0); dataGenerator.place("removable1", job1Task0); dataGenerator.setQuota("job1", 1); TaskRelocationPlan job1Task0Plan = TaskRelocationPlan.newBuilder() .withTaskId(job1Task0.getId()) .withRelocationTime(Long.MAX_VALUE / 2) .build(); Optional<Pair<Node, List<Task>>> results = newDescheduler(Collections.singletonMap(job1Task0.getId(), job1Task0Plan)).nextBestMatch(); assertThat(results).isEmpty(); }
@Test public void testSelfMigrationAfterDeadline() { Task job1Task0 = jobOperations.getTasks("job1").get(0); dataGenerator.place("removable1", job1Task0); dataGenerator.setQuota("job1", 1); TaskRelocationPlan job1Task0Plan = TaskRelocationPlan.newBuilder() .withTaskId(job1Task0.getId()) .withRelocationTime(clock.wallTime() - 1) .build(); Optional<Pair<Node, List<Task>>> results = newDescheduler(Collections.singletonMap(job1Task0.getId(), job1Task0Plan)).nextBestMatch(); assertThat(results).isNotEmpty(); }
@Test public void testFitness() { List<Node> removableAgents = nodeDataResolver.resolve().values().stream() .filter(n -> n.getServerGroupId().equals("removable1")) .collect(Collectors.toList());; String agent1 = removableAgents.get(0).getId(); String agent2 = removableAgents.get(1).getId(); List<Task> tasksOfJob1 = jobOperations.getTasks("job1"); dataGenerator.placeOnAgent(agent1, tasksOfJob1.get(0), tasksOfJob1.get(1)); dataGenerator.placeOnAgent(agent2, tasksOfJob1.get(2)); dataGenerator.setQuota("job1", 1); Optional<Pair<Node, List<Task>>> results = newDescheduler(Collections.emptyMap()).nextBestMatch(); assertThat(results).isPresent(); assertThat(results.get().getLeft().getId()).isEqualTo(agent2); }
@Test public void testAgentInstanceRequiredMigration() { Task job1Task0 = jobOperations.getTasks("job1").get(0); dataGenerator.place("active1", job1Task0); dataGenerator.setQuota("job1", 1); job1Task0 = jobOperations.findTaskById(job1Task0.getId()).get().getRight(); dataGenerator.addInstanceAttribute(job1Task0.getTaskContext().get(TaskAttributes.TASK_ATTRIBUTES_AGENT_INSTANCE_ID), RelocationAttributes.RELOCATION_REQUIRED, "true"); Optional<Pair<Node, List<Task>>> results = newDescheduler(Collections.emptyMap()).nextBestMatch(); assertThat(results).isNotEmpty(); }