LocalitySchedulingConstraint implements SchedulingConstraint { @Override public boolean testSchedulability(final ExecutorRepresenter executor, final Task task) { if (task.getTaskIncomingEdges().isEmpty()) { final Set<String> sourceLocations; try { sourceLocations = getSourceDataLocations(task.getIrVertexIdToReadable().values()); } catch (final UnsupportedOperationException e) { return true; } catch (final Exception e) { throw new RuntimeException(e); } if (sourceLocations.size() == 0) { return true; } return sourceLocations.contains(executor.getNodeName()); } else { final List<String> intermediateLocations = getIntermediateDataLocations(task); if (intermediateLocations.isEmpty()) { return true; } else { return intermediateLocations.contains(executor.getExecutorId()); } } } @Inject private  LocalitySchedulingConstraint(final BlockManagerMaster blockManagerMaster); @Override boolean testSchedulability(final ExecutorRepresenter executor, final Task task);  }
@Test public void testSourceLocationAwareSchedulingNotAvailable() throws InjectionException { final SchedulingConstraint schedulingConstraint = injector .getInstance(LocalitySchedulingConstraint.class); final Task task = CreateTask.withReadablesWithSourceLocations( Collections.singletonList(Collections.singletonList(SITE_0))); final ExecutorRepresenter e0 = mockExecutorRepresenter(SITE_1); final ExecutorRepresenter e1 = mockExecutorRepresenter(SITE_1); assertEquals(Collections.emptySet(), Arrays.asList(e0, e1).stream() .filter(e -> schedulingConstraint.testSchedulability(e, task)) .collect(Collectors.toSet())); }
@Test public void testSourceLocationAwareSchedulingWithMultiSource() throws InjectionException { final SchedulingConstraint schedulingConstraint = injector .getInstance(LocalitySchedulingConstraint.class); final Task task0 = CreateTask.withReadablesWithSourceLocations( Collections.singletonList(Collections.singletonList(SITE_1))); final Task task1 = CreateTask.withReadablesWithSourceLocations( Collections.singletonList(Arrays.asList(SITE_0, SITE_1, SITE_2))); final Task task2 = CreateTask.withReadablesWithSourceLocations( Arrays.asList(Collections.singletonList(SITE_0), Collections.singletonList(SITE_1), Arrays.asList(SITE_1, SITE_2))); final Task task3 = CreateTask.withReadablesWithSourceLocations( Arrays.asList(Collections.singletonList(SITE_1), Collections.singletonList(SITE_0), Arrays.asList(SITE_0, SITE_2))); final ExecutorRepresenter e = mockExecutorRepresenter(SITE_1); for (final Task task : new HashSet<>(Arrays.asList(task0, task1, task2, task3))) { assertTrue(schedulingConstraint.testSchedulability(e, task)); } }