org.apache.flink.runtime.taskmanager.TaskManagerLocation Java Examples

public SimpleSlotProvider(JobID jobId, int numSlots, TaskManagerGateway taskManagerGateway) {
	checkNotNull(jobId, "jobId");
	checkArgument(numSlots >= 0, "numSlots must be >= 0");

	this.slots = new ArrayDeque<>(numSlots);

	for (int i = 0; i < numSlots; i++) {
		SimpleSlotContext as = new SimpleSlotContext(
			new AllocationID(),
			new TaskManagerLocation(ResourceID.generate(), InetAddress.getLoopbackAddress(), 10000 + i),
			0,
			taskManagerGateway);
		slots.add(as);
	}

	allocatedSlots = new HashMap<>(slots.size());
}
 

private CompletableFuture<Collection<TaskManagerLocation>> combineLocations(
		final CompletableFuture<Collection<TaskManagerLocation>> locationsCombinedAlready,
		final Collection<CompletableFuture<TaskManagerLocation>> locationsToCombine) {

	final CompletableFuture<Set<TaskManagerLocation>> uniqueLocationsFuture =
		FutureUtils.combineAll(locationsToCombine).thenApply(HashSet::new);

	return locationsCombinedAlready.thenCombine(
		uniqueLocationsFuture,
		(locationsOnOneEdge, locationsOnAnotherEdge) -> {
			if ((!locationsOnOneEdge.isEmpty() && locationsOnAnotherEdge.size() > locationsOnOneEdge.size())
				|| locationsOnAnotherEdge.isEmpty()) {
				return locationsOnOneEdge;
			} else {
				return locationsOnAnotherEdge;
			}
		});
}
 

public ArchivedExecution(
		StringifiedAccumulatorResult[] userAccumulators, IOMetrics ioMetrics,
		ExecutionAttemptID attemptId, int attemptNumber, ExecutionState state, String failureCause,
		TaskManagerLocation assignedResourceLocation, AllocationID assignedAllocationID,  int parallelSubtaskIndex,
		long[] stateTimestamps) {
	this.userAccumulators = userAccumulators;
	this.ioMetrics = ioMetrics;
	this.failureCause = failureCause;
	this.assignedResourceLocation = assignedResourceLocation;
	this.attemptNumber = attemptNumber;
	this.attemptId = attemptId;
	this.state = state;
	this.stateTimestamps = stateTimestamps;
	this.parallelSubtaskIndex = parallelSubtaskIndex;
	this.assignedAllocationID = assignedAllocationID;
}
 

public SimpleSlotProvider(int numSlots, TaskManagerGateway taskManagerGateway) {
	checkArgument(numSlots >= 0, "numSlots must be >= 0");

	this.slots = new ArrayDeque<>(numSlots);

	for (int i = 0; i < numSlots; i++) {
		SimpleSlotContext as = new SimpleSlotContext(
			new AllocationID(),
			new TaskManagerLocation(ResourceID.generate(), InetAddress.getLoopbackAddress(), 10000 + i),
			0,
			taskManagerGateway,
			ResourceProfile.ANY);
		slots.add(as);
	}

	allocatedSlots = new HashMap<>(slots.size());
}
 

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.");
	}
}
 

JobManagerRetryingRegistration(
		Logger log,
		RpcService rpcService,
		String targetName,
		Class<JobMasterGateway> targetType,
		String targetAddress,
		JobMasterId jobMasterId,
		RetryingRegistrationConfiguration retryingRegistrationConfiguration,
		String taskManagerRpcAddress,
		TaskManagerLocation taskManagerLocation) {
	super(
		log,
		rpcService,
		targetName,
		targetType,
		targetAddress,
		jobMasterId,
		retryingRegistrationConfiguration);

	this.taskManagerRpcAddress = taskManagerRpcAddress;
	this.taskManagerLocation = Preconditions.checkNotNull(taskManagerLocation);
}
 

/**
 * Tests that it can get the task manager location in an Execution.
 */
@Test
public void testGetTaskManagerLocationWhenScheduled() throws Exception {
	final JobVertex jobVertex = ExecutionGraphTestUtils.createNoOpVertex(1);

	final TestingLogicalSlot testingLogicalSlot = new TestingLogicalSlotBuilder().createTestingLogicalSlot();
	final ExecutionGraph eg = ExecutionGraphTestUtils.createSimpleTestGraph(new JobID(), jobVertex);
	final ExecutionGraphToInputsLocationsRetrieverAdapter inputsLocationsRetriever =
			new ExecutionGraphToInputsLocationsRetrieverAdapter(eg);

	final ExecutionVertex onlyExecutionVertex = eg.getAllExecutionVertices().iterator().next();
	onlyExecutionVertex.deployToSlot(testingLogicalSlot);

	ExecutionVertexID executionVertexId = new ExecutionVertexID(jobVertex.getID(), 0);
	Optional<CompletableFuture<TaskManagerLocation>> taskManagerLocationOptional =
			inputsLocationsRetriever.getTaskManagerLocation(executionVertexId);

	assertTrue(taskManagerLocationOptional.isPresent());

	final CompletableFuture<TaskManagerLocation> taskManagerLocationFuture = taskManagerLocationOptional.get();
	assertThat(taskManagerLocationFuture.get(), is(testingLogicalSlot.getTaskManagerLocation()));
}
 

@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 CompletableFuture<Collection<TaskManagerLocation>> getPreferredLocationsBasedOnInputs(
		final ExecutionVertexID executionVertexId,
		final Set<ExecutionVertexID> producersToIgnore) {

	CompletableFuture<Collection<TaskManagerLocation>> preferredLocations =
		CompletableFuture.completedFuture(Collections.emptyList());

	final Collection<Collection<ExecutionVertexID>> allProducers =
		inputsLocationsRetriever.getConsumedResultPartitionsProducers(executionVertexId);
	for (Collection<ExecutionVertexID> producers : allProducers) {
		final Collection<CompletableFuture<TaskManagerLocation>> locationsFutures =
			getInputLocationFutures(producersToIgnore, producers);

		preferredLocations = combineLocations(preferredLocations, locationsFutures);
	}
	return preferredLocations;
}
 

@Override
public Optional<SlotInfoAndLocality> selectBestSlotForProfile(
	@Nonnull Collection<SlotInfoAndResources> availableSlots,
	@Nonnull SlotProfile slotProfile) {

	Collection<TaskManagerLocation> locationPreferences = slotProfile.getPreferredLocations();

	if (availableSlots.isEmpty()) {
		return Optional.empty();
	}

	final ResourceProfile resourceProfile = slotProfile.getPhysicalSlotResourceProfile();

	// if we have no location preferences, we can only filter by the additional requirements.
	return locationPreferences.isEmpty() ?
		selectWithoutLocationPreference(availableSlots, resourceProfile) :
		selectWitLocationPreference(availableSlots, locationPreferences, resourceProfile);
}
 

/**
 * Creates a channel deployment descriptor by completing the partition location.
 *
 * @see InputChannelDeploymentDescriptor
 */
public InputChannelDeploymentDescriptor createInputChannelDeploymentDescriptor(Execution consumerExecution) {
	checkNotNull(consumerExecution, "consumerExecution");

	TaskManagerLocation consumerLocation = consumerExecution.getAssignedResourceLocation();
	checkNotNull(consumerLocation, "Consumer connection info null");

	final ResultPartitionLocation partitionLocation;

	if (consumerLocation.equals(partitionTaskManagerLocation)) {
		partitionLocation = ResultPartitionLocation.createLocal();
	}
	else {
		partitionLocation = ResultPartitionLocation.createRemote(
				new ConnectionID(partitionTaskManagerLocation, partitionConnectionIndex));
	}

	return new InputChannelDeploymentDescriptor(partitionID, partitionLocation);
}
 

private InputGateDeploymentDescriptor createInputGateDeploymentDescriptor(
		TaskManagerLocation senderLocation,
		int gateIndex,
		ResourceID localLocation) {

	final ShuffleDescriptor[] channelDescriptors = new ShuffleDescriptor[channels];
	for (int channelIndex = 0; channelIndex < channels; ++channelIndex) {
		channelDescriptors[channelIndex] = createShuffleDescriptor(
			localMode, partitionIds[gateIndex], localLocation, senderLocation, channelIndex);
	}

	return new InputGateDeploymentDescriptor(
		dataSetID,
		ResultPartitionType.PIPELINED_BOUNDED,
		// 0 is used because TestRemoteInputChannel and TestLocalInputChannel will
		// ignore this and use channelIndex instead when requesting a subpartition
		0,
		channelDescriptors);
}
 

/**
 * This test validates that vertices with too many input streams do not have a location
 * preference any more.
 */
@Test
public void testNoLocalityInputLargeAllToAll() throws Exception {
	final int parallelism = 100;

	final ExecutionGraph graph = createTestGraph(parallelism, true);

	// set the location for all sources to a distinct location
	for (int i = 0; i < parallelism; i++) {
		ExecutionVertex source = graph.getAllVertices().get(sourceVertexId).getTaskVertices()[i];
		TaskManagerLocation location = new TaskManagerLocation(
				ResourceID.generate(), InetAddress.getLoopbackAddress(), 10000 + i);
		initializeLocation(source, location);
	}

	// validate that the target vertices have no location preference
	for (int i = 0; i < parallelism; i++) {
		ExecutionVertex target = graph.getAllVertices().get(targetVertexId).getTaskVertices()[i];

		Iterator<CompletableFuture<TaskManagerLocation>> preference = target.getPreferredLocations().iterator();
		assertFalse(preference.hasNext());
	}
}
 

/**
 * Tests that any preferred locations are calculated.
 */
@Test
public void testAnyPreferredLocationCalculation() throws Exception {
	final TaskManagerLocation taskManagerLocation1 = new LocalTaskManagerLocation();
	final TaskManagerLocation taskManagerLocation3 = new LocalTaskManagerLocation();

	final CompletableFuture<TaskManagerLocation> locationFuture1 = CompletableFuture.completedFuture(taskManagerLocation1);
	final CompletableFuture<TaskManagerLocation> locationFuture2 = new CompletableFuture<>();
	final CompletableFuture<TaskManagerLocation> locationFuture3 = CompletableFuture.completedFuture(taskManagerLocation3);

	final Execution execution = getExecution(Arrays.asList(locationFuture1, locationFuture2, locationFuture3));

	CompletableFuture<Collection<TaskManagerLocation>> preferredLocationsFuture = execution.calculatePreferredLocations(LocationPreferenceConstraint.ANY);

	assertTrue(preferredLocationsFuture.isDone());

	final Collection<TaskManagerLocation> preferredLocations = preferredLocationsFuture.get();

	assertThat(preferredLocations, containsInAnyOrder(taskManagerLocation1, taskManagerLocation3));
}
 

TestingLogicalSlot(
		TaskManagerLocation taskManagerLocation,
		TaskManagerGateway taskManagerGateway,
		int slotNumber,
		AllocationID allocationId,
		SlotRequestId slotRequestId,
		@Nullable SlotSharingGroupId slotSharingGroupId,
		boolean automaticallyCompleteReleaseFuture,
		SlotOwner slotOwner) {

	this.taskManagerLocation = Preconditions.checkNotNull(taskManagerLocation);
	this.taskManagerGateway = Preconditions.checkNotNull(taskManagerGateway);
	this.payloadReference = new AtomicReference<>();
	this.slotNumber = slotNumber;
	this.allocationId = Preconditions.checkNotNull(allocationId);
	this.slotRequestId = Preconditions.checkNotNull(slotRequestId);
	this.slotSharingGroupId = slotSharingGroupId;
	this.releaseFuture = new CompletableFuture<>();
	this.automaticallyCompleteReleaseFuture = automaticallyCompleteReleaseFuture;
	this.slotOwner = Preconditions.checkNotNull(slotOwner);
}
 

/**
 * Creates a new simple slot that belongs to the given shared slot and
 * is identified by the given ID.
 *
 * @param owner The component from which this slot is allocated.
 * @param location The location info of the TaskManager where the slot was allocated from
 * @param slotNumber The number of the simple slot in its parent shared slot.
 * @param taskManagerGateway to communicate with the associated task manager.
 * @param parent The parent shared slot.
 * @param groupID The ID that identifies the group that the slot belongs to.
 */
public SimpleSlot(
		SlotOwner owner,
		TaskManagerLocation location,
		int slotNumber,
		TaskManagerGateway taskManagerGateway,
		@Nullable SharedSlot parent,
		@Nullable AbstractID groupID) {

	super(
		parent != null ?
			parent.getSlotContext() :
			new SimpleSlotContext(
				NO_ALLOCATION_ID,
				location,
				slotNumber,
				taskManagerGateway),
		owner,
		slotNumber,
		parent,
		groupID);
}
 

@Override
public CompletableFuture<Collection<SlotOffer>> offerSlots(
		final ResourceID taskManagerId,
		final Collection<SlotOffer> slots,
		final Time timeout) {

	Tuple2<TaskManagerLocation, TaskExecutorGateway> taskManager = registeredTaskManagers.get(taskManagerId);

	if (taskManager == null) {
		return FutureUtils.completedExceptionally(new Exception("Unknown TaskManager " + taskManagerId));
	}

	final TaskManagerLocation taskManagerLocation = taskManager.f0;
	final TaskExecutorGateway taskExecutorGateway = taskManager.f1;

	final RpcTaskManagerGateway rpcTaskManagerGateway = new RpcTaskManagerGateway(taskExecutorGateway, getFencingToken());

	return CompletableFuture.completedFuture(
		slotPool.offerSlots(
			taskManagerLocation,
			rpcTaskManagerGateway,
			slots));
}
 

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 slot = new SimpleSlotContext(
		new AllocationID(),
		location,
		0,
		mock(TaskManagerGateway.class));

	SimpleSlot simpleSlot = new SimpleSlot(slot, mock(SlotOwner.class), 0);

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

/**
 * Gets the location preferences of the vertex's current task execution, as determined by the locations
 * of the predecessors from which it receives input data.
 * If there are more than MAX_DISTINCT_LOCATIONS_TO_CONSIDER different locations of source data, this
 * method returns {@code null} to indicate no location preference.
 *
 * @return The preferred locations based in input streams, or an empty iterable,
 *         if there is no input-based preference.
 */
public Collection<CompletableFuture<TaskManagerLocation>> getPreferredLocationsBasedOnInputs() {
	// otherwise, base the preferred locations on the input connections
	if (inputEdges == null) {
		return Collections.emptySet();
	}
	else {
		Set<CompletableFuture<TaskManagerLocation>> locations = new HashSet<>(getTotalNumberOfParallelSubtasks());
		Set<CompletableFuture<TaskManagerLocation>> inputLocations = new HashSet<>(getTotalNumberOfParallelSubtasks());

		// go over all inputs
		for (int i = 0; i < inputEdges.length; i++) {
			inputLocations.clear();
			ExecutionEdge[] sources = inputEdges[i];
			if (sources != null) {
				// go over all input sources
				for (int k = 0; k < sources.length; k++) {
					// look-up assigned slot of input source
					CompletableFuture<TaskManagerLocation> locationFuture = sources[k].getSource().getProducer().getCurrentTaskManagerLocationFuture();
					// add input location
					inputLocations.add(locationFuture);
					// inputs which have too many distinct sources are not considered
					if (inputLocations.size() > MAX_DISTINCT_LOCATIONS_TO_CONSIDER) {
						inputLocations.clear();
						break;
					}
				}
			}
			// keep the locations of the input with the least preferred locations
			if (locations.isEmpty() || // nothing assigned yet
					(!inputLocations.isEmpty() && inputLocations.size() < locations.size())) {
				// current input has fewer preferred locations
				locations.clear();
				locations.addAll(inputLocations);
			}
		}

		return locations.isEmpty() ? Collections.emptyList() : locations;
	}
}
 

@Test
public void nonColocationFollowsCoLocation() throws Exception {
	JobVertexID jid1 = new JobVertexID();
	JobVertexID jid2 = new JobVertexID();

	TaskManagerLocation loc1 = testingSlotProvider.addTaskManager(1);
	TaskManagerLocation loc2 = testingSlotProvider.addTaskManager(1);

	assertEquals(2, testingSlotProvider.getNumberOfAvailableSlots());

	SlotSharingGroup sharingGroup = new SlotSharingGroup();

	CoLocationGroup ccg = new CoLocationGroup();
	CoLocationConstraint cc1 = new CoLocationConstraint(ccg);
	CoLocationConstraint cc2 = new CoLocationConstraint(ccg);

	LogicalSlot s1 = testingSlotProvider.allocateSlot(
			new ScheduledUnit(getTestVertexWithLocation(jid1, 0, 2, sharingGroup, loc1), sharingGroup.getSlotSharingGroupId(), cc1), false, slotProfileForLocation(loc1), TestingUtils.infiniteTime()).get();
	LogicalSlot s2 = testingSlotProvider.allocateSlot(
			new ScheduledUnit(getTestVertexWithLocation(jid1, 1, 2, sharingGroup, loc2), sharingGroup.getSlotSharingGroupId(), cc2), false, slotProfileForLocation(loc2), TestingUtils.infiniteTime()).get();

	LogicalSlot s3 = testingSlotProvider.allocateSlot(
			new ScheduledUnit(getTestVertexWithLocation(jid2, 0, 2, sharingGroup, loc1), sharingGroup.getSlotSharingGroupId()), false, slotProfileForLocation(loc1), TestingUtils.infiniteTime()).get();
	LogicalSlot s4 = testingSlotProvider.allocateSlot(
			new ScheduledUnit(getTestVertexWithLocation(jid2, 1, 2, sharingGroup, loc1), sharingGroup.getSlotSharingGroupId()), false, slotProfileForLocation(loc1), TestingUtils.infiniteTime()).get();

	// check that each slot got two
	assertEquals(s1.getTaskManagerLocation(), s3.getTaskManagerLocation());
	assertEquals(s2.getTaskManagerLocation(), s4.getTaskManagerLocation());

	s1.releaseSlot();
	s2.releaseSlot();
	s3.releaseSlot();
	s4.releaseSlot();

	assertEquals(2, testingSlotProvider.getNumberOfAvailableSlots());
}
 

public static Execution getTestVertex(Instance... preferredInstances) {
	List<TaskManagerLocation> locations = new ArrayList<>(preferredInstances.length);
	for (Instance i : preferredInstances) {
		locations.add(i.getTaskManagerLocation());
	}
	return getTestVertex(locations);
}
 

public AllocatedSlot(
		AllocationID allocationId,
		TaskManagerLocation location,
		int physicalSlotNumber,
		ResourceProfile resourceProfile,
		TaskManagerGateway taskManagerGateway) {
	this.allocationId = checkNotNull(allocationId);
	this.taskManagerLocation = checkNotNull(location);
	this.physicalSlotNumber = physicalSlotNumber;
	this.resourceProfile = checkNotNull(resourceProfile);
	this.taskManagerGateway = checkNotNull(taskManagerGateway);

	payloadReference = new AtomicReference<>(null);
}
 

private ExecutionVertexSchedulingRequirements(
		ExecutionVertexID executionVertexId,
		@Nullable AllocationID previousAllocationId,
		ResourceProfile resourceProfile,
		@Nullable SlotSharingGroupId slotSharingGroupId,
		@Nullable CoLocationConstraint coLocationConstraint,
		Collection<TaskManagerLocation> preferredLocations) {
	this.executionVertexId = checkNotNull(executionVertexId);
	this.previousAllocationId = previousAllocationId;
	this.resourceProfile = checkNotNull(resourceProfile);
	this.slotSharingGroupId = slotSharingGroupId;
	this.coLocationConstraint = coLocationConstraint;
	this.preferredLocations = checkNotNull(preferredLocations);
}
 

private static void generateArchivedJob() throws Exception {
	// Attempt
	StringifiedAccumulatorResult acc1 = new StringifiedAccumulatorResult("name1", "type1", "value1");
	StringifiedAccumulatorResult acc2 = new StringifiedAccumulatorResult("name2", "type2", "value2");
	TaskManagerLocation location = new TaskManagerLocation(new ResourceID("hello"), InetAddress.getLocalHost(), 1234);
	AllocationID allocationID = new AllocationID(42L, 43L);
	originalAttempt = new ArchivedExecutionBuilder()
		.setStateTimestamps(new long[]{1, 2, 3, 4, 5, 6, 7, 8, 9})
		.setParallelSubtaskIndex(1)
		.setAttemptNumber(0)
		.setAssignedResourceLocation(location)
		.setAssignedAllocationID(allocationID)
		.setUserAccumulators(new StringifiedAccumulatorResult[]{acc1, acc2})
		.setState(ExecutionState.FINISHED)
		.setFailureCause("attemptException")
		.build();
	// Subtask
	originalSubtask = new ArchivedExecutionVertexBuilder()
		.setSubtaskIndex(originalAttempt.getParallelSubtaskIndex())
		.setTaskNameWithSubtask("hello(1/1)")
		.setCurrentExecution(originalAttempt)
		.build();
	// Task
	originalTask = new ArchivedExecutionJobVertexBuilder()
		.setTaskVertices(new ArchivedExecutionVertex[]{originalSubtask})
		.build();
	// Job
	Map<JobVertexID, ArchivedExecutionJobVertex> tasks = new HashMap<>();
	tasks.put(originalTask.getJobVertexId(), originalTask);
	originalJob = new ArchivedExecutionGraphBuilder()
		.setJobID(new JobID())
		.setTasks(tasks)
		.setFailureCause(new ErrorInfo(new Exception("jobException"), originalAttempt.getStateTimestamp(ExecutionState.FAILED)))
		.setState(JobStatus.FINISHED)
		.setStateTimestamps(new long[]{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11})
		.setArchivedUserAccumulators(new StringifiedAccumulatorResult[]{acc1, acc2})
		.build();
}
 

@Test
public void testCancelAllSlots() {
	try {
		ResourceID resourceID = ResourceID.generate();
		HardwareDescription hardwareDescription = new HardwareDescription(4, 2L*1024*1024*1024, 1024*1024*1024, 512*1024*1024);
		InetAddress address = InetAddress.getByName("127.0.0.1");
		TaskManagerLocation connection = new TaskManagerLocation(resourceID, address, 10001);

		Instance instance = new Instance(
			new ActorTaskManagerGateway(DummyActorGateway.INSTANCE),
			connection,
			new InstanceID(),
			hardwareDescription,
			3);

		assertEquals(3, instance.getNumberOfAvailableSlots());

		SimpleSlot slot1 = instance.allocateSimpleSlot();
		SimpleSlot slot2 = instance.allocateSimpleSlot();
		SimpleSlot slot3 = instance.allocateSimpleSlot();

		instance.cancelAndReleaseAllSlots();

		assertEquals(3, instance.getNumberOfAvailableSlots());

		assertTrue(slot1.isCanceled());
		assertTrue(slot2.isCanceled());
		assertTrue(slot3.isCanceled());
	}
	catch (Exception e) {
		e.printStackTrace();
		fail(e.getMessage());
	}
}
 

/**
 * Tests that the lazy scheduling flag is correctly forwarded to the produced partition descriptors.
 */
@Test
public void testTddProducedPartitionsLazyScheduling() throws Exception {
	for (ScheduleMode scheduleMode: ScheduleMode.values()) {
		ExecutionJobVertex jobVertex = ExecutionGraphTestUtils.getExecutionJobVertex(
			new JobVertexID(),
			new DirectScheduledExecutorService(),
			scheduleMode);

		IntermediateResult result =
			new IntermediateResult(new IntermediateDataSetID(), jobVertex, 1, ResultPartitionType.PIPELINED);

		ExecutionAttemptID attemptID = new ExecutionAttemptID();
		ExecutionVertex vertex =
			new ExecutionVertex(jobVertex, 0, new IntermediateResult[]{result}, Time.minutes(1));
		TaskDeploymentDescriptorFactory tddFactory =
			TaskDeploymentDescriptorFactory.fromExecutionVertex(vertex, 1);

		ExecutionEdge mockEdge = createMockExecutionEdge(1);

		result.getPartitions()[0].addConsumerGroup();
		result.getPartitions()[0].addConsumer(mockEdge, 0);

		TaskManagerLocation location =
			new TaskManagerLocation(ResourceID.generate(), InetAddress.getLoopbackAddress(), 1);

		TaskDeploymentDescriptor tdd = tddFactory.createDeploymentDescriptor(
			new AllocationID(),
			0,
			null,
			Execution.registerProducedPartitions(vertex, location, attemptID, scheduleMode.allowLazyDeployment()).get().values());

		Collection<ResultPartitionDeploymentDescriptor> producedPartitions = tdd.getProducedPartitions();

		assertEquals(1, producedPartitions.size());
		ResultPartitionDeploymentDescriptor desc = producedPartitions.iterator().next();
		assertEquals(scheduleMode.allowLazyDeployment(), desc.sendScheduleOrUpdateConsumersMessage());
	}
}
 

public SlotProfile(
	@Nonnull ResourceProfile resourceProfile,
	@Nonnull Collection<TaskManagerLocation> preferredLocations,
	@Nonnull Collection<AllocationID> preferredAllocations) {

	this(resourceProfile, preferredLocations, preferredAllocations, Collections.emptySet());
}
 

private Collection<TaskManagerLocation> getPreferredLocationsBasedOnState(
		final ExecutionVertexID executionVertexId) {

	return stateLocationRetriever.getStateLocation(executionVertexId)
		.map(Collections::singleton)
		.orElse(Collections.emptySet());
}
 

public ArchivedExecution build() throws UnknownHostException {
	return new ArchivedExecution(
		userAccumulators != null ? userAccumulators : new StringifiedAccumulatorResult[0],
		ioMetrics != null ? ioMetrics : new TestIOMetrics(),
		attemptId != null ? attemptId : new ExecutionAttemptID(),
		attemptNumber,
		state != null ? state : ExecutionState.FINISHED,
		failureCause != null ? failureCause : "(null)",
		assignedResourceLocation != null ? assignedResourceLocation : new TaskManagerLocation(new ResourceID("tm"), InetAddress.getLocalHost(), 1234),
		assignedAllocationID != null ? assignedAllocationID : new AllocationID(0L, 0L),
		parallelSubtaskIndex,
		stateTimestamps != null ? stateTimestamps : new long[]{1, 2, 3, 4, 5, 5, 5, 5}
	);
}
 

public static SubtaskExecutionAttemptDetailsInfo create(AccessExecution execution, MutableIOMetrics ioMetrics) {
	final ExecutionState status = execution.getState();
	final long now = System.currentTimeMillis();

	final TaskManagerLocation location = execution.getAssignedResourceLocation();
	final String locationString = location == null ? "(unassigned)" : location.getHostname();

	long startTime = execution.getStateTimestamp(ExecutionState.DEPLOYING);
	if (startTime == 0) {
		startTime = -1;
	}
	final long endTime = status.isTerminal() ? execution.getStateTimestamp(status) : -1;
	final long duration = startTime > 0 ? ((endTime > 0 ? endTime : now) - startTime) : -1;

	final IOMetricsInfo ioMetricsInfo = new IOMetricsInfo(
		ioMetrics.getNumBytesInLocal() + ioMetrics.getNumBytesInRemote(),
		ioMetrics.isNumBytesInLocalComplete() && ioMetrics.isNumBytesInRemoteComplete(),
		ioMetrics.getNumBytesOut(),
		ioMetrics.isNumBytesOutComplete(),
		ioMetrics.getNumRecordsIn(),
		ioMetrics.isNumRecordsInComplete(),
		ioMetrics.getNumRecordsOut(),
		ioMetrics.isNumRecordsOutComplete());

	return new SubtaskExecutionAttemptDetailsInfo(
		execution.getParallelSubtaskIndex(),
		status,
		execution.getAttemptNumber(),
		locationString,
		startTime,
		endTime,
		duration,
		ioMetricsInfo
	);
}