Java Code Examples for org.apache.flink.runtime.jobmanager.scheduler.Locality#LOCAL

@Nonnull
static SingleLogicalSlot createSingleLogicalSlot(SlotOwner slotOwner, TaskManagerGateway taskManagerGateway, SlotRequestId slotRequestId) {
	TaskManagerLocation location = new TaskManagerLocation(
		ResourceID.generate(), InetAddress.getLoopbackAddress(), 12345);

	SimpleSlotContext slotContext = new SimpleSlotContext(
		new AllocationID(),
		location,
		0,
		taskManagerGateway);

	return new SingleLogicalSlot(
		slotRequestId,
		slotContext,
		null,
		Locality.LOCAL,
		slotOwner);
}
 

private void initializeLocation(ExecutionVertex vertex, TaskManagerLocation location) throws Exception {
	// we need a bit of reflection magic to initialize the location without going through
	// scheduling paths. we choose to do that, rather than the alternatives:
	//  - mocking the scheduler created fragile tests that break whenever the scheduler is adjusted
	//  - exposing test methods in the ExecutionVertex leads to undesirable setters 

	SlotContext slotContext = new SimpleSlotContext(
		new AllocationID(),
		location,
		0,
		mock(TaskManagerGateway.class));



	LogicalSlot slot = new SingleLogicalSlot(
		new SlotRequestId(),
		slotContext,
		null,
		Locality.LOCAL,
		mock(SlotOwner.class));

	if (!vertex.getCurrentExecutionAttempt().tryAssignResource(slot)) {
		throw new FlinkException("Could not assign resource.");
	}
}
 

@Nonnull
static SingleLogicalSlot createSingleLogicalSlot(SlotOwner slotOwner, TaskManagerGateway taskManagerGateway, SlotRequestId slotRequestId) {
	TaskManagerLocation location = new TaskManagerLocation(
		ResourceID.generate(), InetAddress.getLoopbackAddress(), 12345);

	SimpleSlotContext slotContext = new SimpleSlotContext(
		new AllocationID(),
		location,
		0,
		taskManagerGateway);

	return new SingleLogicalSlot(
		slotRequestId,
		slotContext,
		null,
		Locality.LOCAL,
		slotOwner);
}
 

private void initializeLocation(ExecutionVertex vertex, TaskManagerLocation location) throws Exception {
	// we need a bit of reflection magic to initialize the location without going through
	// scheduling paths. we choose to do that, rather than the alternatives:
	//  - mocking the scheduler created fragile tests that break whenever the scheduler is adjusted
	//  - exposing test methods in the ExecutionVertex leads to undesirable setters 

	SlotContext slotContext = new SimpleSlotContext(
		new AllocationID(),
		location,
		0,
		mock(TaskManagerGateway.class));



	LogicalSlot slot = new SingleLogicalSlot(
		new SlotRequestId(),
		slotContext,
		null,
		Locality.LOCAL,
		mock(SlotOwner.class));

	if (!vertex.getCurrentExecutionAttempt().tryAssignResource(slot)) {
		throw new FlinkException("Could not assign resource.");
	}
}
 

@Nonnull
static SingleLogicalSlot createSingleLogicalSlot(SlotOwner slotOwner, TaskManagerGateway taskManagerGateway, SlotRequestId slotRequestId) {
	TaskManagerLocation location = new TaskManagerLocation(
		ResourceID.generate(), InetAddress.getLoopbackAddress(), 12345);

	SimpleSlotContext slotContext = new SimpleSlotContext(
		new AllocationID(),
		location,
		0,
		taskManagerGateway);

	return new SingleLogicalSlot(
		slotRequestId,
		slotContext,
		null,
		Locality.LOCAL,
		slotOwner);
}
 

private SingleLogicalSlot createSingleLogicalSlot(SlotOwner slotOwner) {
	return new SingleLogicalSlot(
		new SlotRequestId(),
		new DummySlotContext(),
		null,
		Locality.LOCAL,
		slotOwner);
}
 

private SingleLogicalSlot createSingleLogicalSlot(SlotOwner slotOwner) {
	return new SingleLogicalSlot(
		new SlotRequestId(),
		createSlotContext(),
		null,
		Locality.LOCAL,
		slotOwner);
}
 

private SingleLogicalSlot createSingleLogicalSlot(SlotOwner slotOwner) {
	return new SingleLogicalSlot(
		new SlotRequestId(),
		createSlotContext(),
		null,
		Locality.LOCAL,
		slotOwner);
}
 

/**
 * Tests that the logical task slot futures are completed once the slot context
 * future is completed.
 */
@Test
public void testSlotContextFutureCompletion() throws Exception {
	final TestingAllocatedSlotActions allocatedSlotActions = new TestingAllocatedSlotActions();

	final SlotSharingManager slotSharingManager = new SlotSharingManager(
		SLOT_SHARING_GROUP_ID,
		allocatedSlotActions,
		SLOT_OWNER);

	final SlotContext slotContext = new SimpleSlotContext(
		new AllocationID(),
		new LocalTaskManagerLocation(),
		0,
		new SimpleAckingTaskManagerGateway());

	CompletableFuture<SlotContext> slotContextFuture = new CompletableFuture<>();
	SlotSharingManager.MultiTaskSlot rootSlot = slotSharingManager.createRootSlot(
		new SlotRequestId(),
		slotContextFuture,
		new SlotRequestId());

	Locality locality1 = Locality.LOCAL;
	SlotSharingManager.SingleTaskSlot singleTaskSlot1 = rootSlot.allocateSingleTaskSlot(
		new SlotRequestId(),
		new AbstractID(),
		locality1);

	Locality locality2 = Locality.HOST_LOCAL;
	SlotSharingManager.SingleTaskSlot singleTaskSlot2 = rootSlot.allocateSingleTaskSlot(
		new SlotRequestId(),
		new AbstractID(),
		locality2);

	CompletableFuture<LogicalSlot> logicalSlotFuture1 = singleTaskSlot1.getLogicalSlotFuture();
	CompletableFuture<LogicalSlot> logicalSlotFuture2 = singleTaskSlot2.getLogicalSlotFuture();
	assertFalse(logicalSlotFuture1.isDone());
	assertFalse(logicalSlotFuture2.isDone());

	slotContextFuture.complete(slotContext);

	assertTrue(logicalSlotFuture1.isDone());
	assertTrue(logicalSlotFuture2.isDone());

	final LogicalSlot logicalSlot1 = logicalSlotFuture1.get();
	final LogicalSlot logicalSlot2 = logicalSlotFuture2.get();

	assertEquals(logicalSlot1.getAllocationId(), slotContext.getAllocationId());
	assertEquals(logicalSlot2.getAllocationId(), slotContext.getAllocationId());
	assertEquals(locality1, logicalSlot1.getLocality());
	assertEquals(locality2, logicalSlot2.getLocality());

	Locality locality3 = Locality.NON_LOCAL;
	SlotSharingManager.SingleTaskSlot singleTaskSlot3 = rootSlot.allocateSingleTaskSlot(
		new SlotRequestId(),
		new AbstractID(),
		locality3);

	CompletableFuture<LogicalSlot> logicalSlotFuture3 = singleTaskSlot3.getLogicalSlotFuture();

	assertTrue(logicalSlotFuture3.isDone());
	LogicalSlot logicalSlot3 = logicalSlotFuture3.get();

	assertEquals(locality3, logicalSlot3.getLocality());
	assertEquals(slotContext.getAllocationId(), logicalSlot3.getAllocationId());
}
 

/**
 * Tests that the logical task slot futures are completed once the slot context
 * future is completed.
 */
@Test
public void testSlotContextFutureCompletion() throws Exception {
	final TestingAllocatedSlotActions allocatedSlotActions = new TestingAllocatedSlotActions();

	final SlotSharingManager slotSharingManager = new SlotSharingManager(
		SLOT_SHARING_GROUP_ID,
		allocatedSlotActions,
		SLOT_OWNER);

	final SlotContext slotContext = new SimpleSlotContext(
		new AllocationID(),
		new LocalTaskManagerLocation(),
		0,
		new SimpleAckingTaskManagerGateway());

	CompletableFuture<SlotContext> slotContextFuture = new CompletableFuture<>();
	SlotSharingManager.MultiTaskSlot rootSlot = slotSharingManager.createRootSlot(
		new SlotRequestId(),
		slotContextFuture,
		new SlotRequestId());

	Locality locality1 = Locality.LOCAL;
	SlotSharingManager.SingleTaskSlot singleTaskSlot1 = rootSlot.allocateSingleTaskSlot(
		new SlotRequestId(),
		ResourceProfile.UNKNOWN,
		new AbstractID(),
		locality1);

	Locality locality2 = Locality.HOST_LOCAL;
	SlotSharingManager.SingleTaskSlot singleTaskSlot2 = rootSlot.allocateSingleTaskSlot(
		new SlotRequestId(),
		ResourceProfile.UNKNOWN,
		new AbstractID(),
		locality2);

	CompletableFuture<LogicalSlot> logicalSlotFuture1 = singleTaskSlot1.getLogicalSlotFuture();
	CompletableFuture<LogicalSlot> logicalSlotFuture2 = singleTaskSlot2.getLogicalSlotFuture();
	assertFalse(logicalSlotFuture1.isDone());
	assertFalse(logicalSlotFuture2.isDone());

	slotContextFuture.complete(slotContext);

	assertTrue(logicalSlotFuture1.isDone());
	assertTrue(logicalSlotFuture2.isDone());

	final LogicalSlot logicalSlot1 = logicalSlotFuture1.get();
	final LogicalSlot logicalSlot2 = logicalSlotFuture2.get();

	assertEquals(logicalSlot1.getAllocationId(), slotContext.getAllocationId());
	assertEquals(logicalSlot2.getAllocationId(), slotContext.getAllocationId());
	assertEquals(locality1, logicalSlot1.getLocality());
	assertEquals(locality2, logicalSlot2.getLocality());

	Locality locality3 = Locality.NON_LOCAL;
	SlotSharingManager.SingleTaskSlot singleTaskSlot3 = rootSlot.allocateSingleTaskSlot(
		new SlotRequestId(),
		ResourceProfile.UNKNOWN,
		new AbstractID(),
		locality3);

	CompletableFuture<LogicalSlot> logicalSlotFuture3 = singleTaskSlot3.getLogicalSlotFuture();

	assertTrue(logicalSlotFuture3.isDone());
	LogicalSlot logicalSlot3 = logicalSlotFuture3.get();

	assertEquals(locality3, logicalSlot3.getLocality());
	assertEquals(slotContext.getAllocationId(), logicalSlot3.getAllocationId());
}
 

/**
 * Allocates a co-located {@link SlotSharingManager.MultiTaskSlot} for the given {@link CoLocationConstraint}.
 *
 * <p>The returned {@link SlotSharingManager.MultiTaskSlot} can be uncompleted.
 *
 * @param coLocationConstraint for which to allocate a {@link SlotSharingManager.MultiTaskSlot}
 * @param multiTaskSlotManager responsible for the slot sharing group for which to allocate the slot
 * @param slotProfile specifying the requirements for the requested slot
 * @param allocationTimeout timeout before the slot allocation times out
 * @return A {@link SlotAndLocality} which contains the allocated{@link SlotSharingManager.MultiTaskSlot}
 * 		and its locality wrt the given location preferences
 */
private SlotSharingManager.MultiTaskSlotLocality allocateCoLocatedMultiTaskSlot(
	CoLocationConstraint coLocationConstraint,
	SlotSharingManager multiTaskSlotManager,
	SlotProfile slotProfile,
	@Nullable Time allocationTimeout) throws NoResourceAvailableException {
	final SlotRequestId coLocationSlotRequestId = coLocationConstraint.getSlotRequestId();

	if (coLocationSlotRequestId != null) {
		// we have a slot assigned --> try to retrieve it
		final SlotSharingManager.TaskSlot taskSlot = multiTaskSlotManager.getTaskSlot(coLocationSlotRequestId);

		if (taskSlot != null) {
			Preconditions.checkState(taskSlot instanceof SlotSharingManager.MultiTaskSlot);

			SlotSharingManager.MultiTaskSlot multiTaskSlot = (SlotSharingManager.MultiTaskSlot) taskSlot;

			if (multiTaskSlot.mayHaveEnoughResourcesToFulfill(slotProfile.getTaskResourceProfile())) {
				return SlotSharingManager.MultiTaskSlotLocality.of(multiTaskSlot, Locality.LOCAL);
			}

			throw new NoResourceAvailableException("Not enough resources in the slot for all co-located tasks.");
		} else {
			// the slot may have been cancelled in the mean time
			coLocationConstraint.setSlotRequestId(null);
		}
	}

	if (coLocationConstraint.isAssigned()) {
		// refine the preferred locations of the slot profile
		slotProfile = SlotProfile.priorAllocation(
			slotProfile.getTaskResourceProfile(),
			slotProfile.getPhysicalSlotResourceProfile(),
			Collections.singleton(coLocationConstraint.getLocation()),
			slotProfile.getPreferredAllocations(),
			slotProfile.getPreviousExecutionGraphAllocations());
	}

	// get a new multi task slot
	SlotSharingManager.MultiTaskSlotLocality multiTaskSlotLocality = allocateMultiTaskSlot(
		coLocationConstraint.getGroupId(),
		multiTaskSlotManager,
		slotProfile,
		allocationTimeout);

	// check whether we fulfill the co-location constraint
	if (coLocationConstraint.isAssigned() && multiTaskSlotLocality.getLocality() != Locality.LOCAL) {
		multiTaskSlotLocality.getMultiTaskSlot().release(
			new FlinkException("Multi task slot is not local and, thus, does not fulfill the co-location constraint."));

		throw new NoResourceAvailableException("Could not allocate a local multi task slot for the " +
			"co location constraint " + coLocationConstraint + '.');
	}

	final SlotRequestId slotRequestId = new SlotRequestId();
	final SlotSharingManager.MultiTaskSlot coLocationSlot =
		multiTaskSlotLocality.getMultiTaskSlot().allocateMultiTaskSlot(
			slotRequestId,
			coLocationConstraint.getGroupId());

	// mark the requested slot as co-located slot for other co-located tasks
	coLocationConstraint.setSlotRequestId(slotRequestId);

	// lock the co-location constraint once we have obtained the allocated slot
	coLocationSlot.getSlotContextFuture().whenComplete(
		(SlotContext slotContext, Throwable throwable) -> {
			if (throwable == null) {
				// check whether we are still assigned to the co-location constraint
				if (Objects.equals(coLocationConstraint.getSlotRequestId(), slotRequestId)) {
					coLocationConstraint.lockLocation(slotContext.getTaskManagerLocation());
				} else {
					log.debug("Failed to lock colocation constraint {} because assigned slot " +
							"request {} differs from fulfilled slot request {}.",
						coLocationConstraint.getGroupId(),
						coLocationConstraint.getSlotRequestId(),
						slotRequestId);
				}
			} else {
				log.debug("Failed to lock colocation constraint {} because the slot " +
						"allocation for slot request {} failed.",
					coLocationConstraint.getGroupId(),
					coLocationConstraint.getSlotRequestId(),
					throwable);
			}
		});

	return SlotSharingManager.MultiTaskSlotLocality.of(coLocationSlot, multiTaskSlotLocality.getLocality());
}
 

@Nonnull
private Optional<SlotInfoAndLocality> selectWitLocationPreference(
	@Nonnull Collection<SlotInfoAndResources> availableSlots,
	@Nonnull Collection<TaskManagerLocation> locationPreferences,
	@Nonnull ResourceProfile resourceProfile) {

	// we build up two indexes, one for resource id and one for host names of the preferred locations.
	final Map<ResourceID, Integer> preferredResourceIDs = new HashMap<>(locationPreferences.size());
	final Map<String, Integer> preferredFQHostNames = new HashMap<>(locationPreferences.size());

	for (TaskManagerLocation locationPreference : locationPreferences) {
		preferredResourceIDs.merge(locationPreference.getResourceID(), 1, Integer::sum);
		preferredFQHostNames.merge(locationPreference.getFQDNHostname(), 1, Integer::sum);
	}

	SlotInfoAndResources bestCandidate = null;
	Locality bestCandidateLocality = Locality.UNKNOWN;
	double bestCandidateScore = Double.NEGATIVE_INFINITY;

	for (SlotInfoAndResources candidate : availableSlots) {

		if (candidate.getRemainingResources().isMatching(resourceProfile)) {

			// this gets candidate is local-weigh
			int localWeigh = preferredResourceIDs.getOrDefault(
				candidate.getSlotInfo().getTaskManagerLocation().getResourceID(), 0);

			// this gets candidate is host-local-weigh
			int hostLocalWeigh = preferredFQHostNames.getOrDefault(
				candidate.getSlotInfo().getTaskManagerLocation().getFQDNHostname(), 0);

			double candidateScore = calculateCandidateScore(localWeigh, hostLocalWeigh, candidate.getTaskExecutorUtilization());
			if (candidateScore > bestCandidateScore) {
				bestCandidateScore = candidateScore;
				bestCandidate = candidate;
				bestCandidateLocality = localWeigh > 0 ?
					Locality.LOCAL : hostLocalWeigh > 0 ?
					Locality.HOST_LOCAL : Locality.NON_LOCAL;
			}
		}
	}

	// at the end of the iteration, we return the candidate with best possible locality or null.
	return bestCandidate != null ?
		Optional.of(SlotInfoAndLocality.of(bestCandidate.getSlotInfo(), bestCandidateLocality)) :
		Optional.empty();
}
 

/**
 * Tests that the logical task slot futures are completed once the slot context
 * future is completed.
 */
@Test
public void testSlotContextFutureCompletion() throws Exception {
	final SlotSharingManager slotSharingManager = createTestingSlotSharingManager();

	final SlotContext slotContext = createSimpleSlotContext();

	CompletableFuture<SlotContext> slotContextFuture = new CompletableFuture<>();
	SlotSharingManager.MultiTaskSlot rootSlot = slotSharingManager.createRootSlot(
		new SlotRequestId(),
		slotContextFuture,
		new SlotRequestId());

	Locality locality1 = Locality.LOCAL;
	SlotSharingManager.SingleTaskSlot singleTaskSlot1 = rootSlot.allocateSingleTaskSlot(
		new SlotRequestId(),
		ResourceProfile.UNKNOWN,
		new AbstractID(),
		locality1);

	Locality locality2 = Locality.HOST_LOCAL;
	SlotSharingManager.SingleTaskSlot singleTaskSlot2 = rootSlot.allocateSingleTaskSlot(
		new SlotRequestId(),
		ResourceProfile.UNKNOWN,
		new AbstractID(),
		locality2);

	CompletableFuture<LogicalSlot> logicalSlotFuture1 = singleTaskSlot1.getLogicalSlotFuture();
	CompletableFuture<LogicalSlot> logicalSlotFuture2 = singleTaskSlot2.getLogicalSlotFuture();
	assertFalse(logicalSlotFuture1.isDone());
	assertFalse(logicalSlotFuture2.isDone());

	slotContextFuture.complete(slotContext);

	assertTrue(logicalSlotFuture1.isDone());
	assertTrue(logicalSlotFuture2.isDone());

	final LogicalSlot logicalSlot1 = logicalSlotFuture1.get();
	final LogicalSlot logicalSlot2 = logicalSlotFuture2.get();

	assertEquals(logicalSlot1.getAllocationId(), slotContext.getAllocationId());
	assertEquals(logicalSlot2.getAllocationId(), slotContext.getAllocationId());
	assertEquals(locality1, logicalSlot1.getLocality());
	assertEquals(locality2, logicalSlot2.getLocality());

	Locality locality3 = Locality.NON_LOCAL;
	SlotSharingManager.SingleTaskSlot singleTaskSlot3 = rootSlot.allocateSingleTaskSlot(
		new SlotRequestId(),
		ResourceProfile.UNKNOWN,
		new AbstractID(),
		locality3);

	CompletableFuture<LogicalSlot> logicalSlotFuture3 = singleTaskSlot3.getLogicalSlotFuture();

	assertTrue(logicalSlotFuture3.isDone());
	LogicalSlot logicalSlot3 = logicalSlotFuture3.get();

	assertEquals(locality3, logicalSlot3.getLocality());
	assertEquals(slotContext.getAllocationId(), logicalSlot3.getAllocationId());
}