[FLINK-18739] Implement MergingSharedSlotProfileRetriever

Input location preferences will be considered for each SharedSlot when allocating a physical slot for it.
Input location preferences of a SharedSlot are the merged input location preferences of all the tasks to run in this SharedSlot.

Inter-ExecutionSlotSharingGroup input location preferences can be respected in this way for ExecutionSlotSharingGroups belonging to different bulks.
If ExecutionSlotSharingGroups belong to the same bulk, the input location preferences are ignored because of possible cyclic dependencies.
Later, we can optimise this case when the declarative resource management for reactive mode is ready.

Intra-ExecutionSlotSharingGroup input location preferences will also be respected when creating ExecutionSlotSharingGroup(s) in LocalInputPreferredSlotSharingStrategy.

This closes apache#13028.

flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/MergingSharedSlotProfileRetrieverFactory.java

@@ -0,0 +1,135 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.runtime.scheduler;
+
+import org.apache.flink.runtime.clusterframework.types.AllocationID;
+import org.apache.flink.runtime.clusterframework.types.ResourceProfile;
+import org.apache.flink.runtime.clusterframework.types.SlotProfile;
+import org.apache.flink.runtime.concurrent.FutureUtils;
+import org.apache.flink.runtime.scheduler.strategy.ExecutionVertexID;
+import org.apache.flink.runtime.taskmanager.TaskManagerLocation;
+import org.apache.flink.util.Preconditions;
+
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.HashSet;
+import java.util.Objects;
+import java.util.Set;
+import java.util.concurrent.CompletableFuture;
+import java.util.function.Function;
+import java.util.stream.Collectors;
+
+/**
+ * Factory for {@link MergingSharedSlotProfileRetriever}.
+ */
+class MergingSharedSlotProfileRetrieverFactory implements SharedSlotProfileRetriever.SharedSlotProfileRetrieverFactory {
+	private final PreferredLocationsRetriever preferredLocationsRetriever;
+
+	private final Function<ExecutionVertexID, ResourceProfile> resourceProfileRetriever;
+
+	private final Function<ExecutionVertexID, AllocationID> priorAllocationIdRetriever;
+
+	MergingSharedSlotProfileRetrieverFactory(
+			PreferredLocationsRetriever preferredLocationsRetriever,
+			Function<ExecutionVertexID, ResourceProfile> resourceProfileRetriever,
+			Function<ExecutionVertexID, AllocationID> priorAllocationIdRetriever) {
+		this.preferredLocationsRetriever = Preconditions.checkNotNull(preferredLocationsRetriever);
+		this.resourceProfileRetriever = Preconditions.checkNotNull(resourceProfileRetriever);
+		this.priorAllocationIdRetriever = Preconditions.checkNotNull(priorAllocationIdRetriever);
+	}
+
+	@Override
+	public SharedSlotProfileRetriever createFromBulk(Set<ExecutionVertexID> bulk) {
+		Set<AllocationID> allPriorAllocationIds = bulk
+			.stream()
+			.map(priorAllocationIdRetriever)
+			.filter(Objects::nonNull)
+			.collect(Collectors.toSet());
+		return new MergingSharedSlotProfileRetriever(allPriorAllocationIds, bulk);
+	}
+
+	/**
+	 * Computes a merged {@link SlotProfile} of an execution slot sharing group within a bulk to schedule.
+	 */
+	private class MergingSharedSlotProfileRetriever implements SharedSlotProfileRetriever {
+		/**
+		 * All previous {@link AllocationID}s of the bulk to schedule.
+		 */
+		private final Set<AllocationID> allBulkPriorAllocationIds;
+
+		/**
+		 * All {@link ExecutionVertexID}s of the bulk.
+		 */
+		private final Set<ExecutionVertexID> producersToIgnore;
+
+		private MergingSharedSlotProfileRetriever(
+				Set<AllocationID> allBulkPriorAllocationIds,
+				Set<ExecutionVertexID> producersToIgnore) {
+			this.allBulkPriorAllocationIds = Preconditions.checkNotNull(allBulkPriorAllocationIds);
+			this.producersToIgnore = Preconditions.checkNotNull(producersToIgnore);
+		}
+
+		/**
+		 * Computes a {@link SlotProfile} of an execution slot sharing group.
+		 *
+		 * <p>The {@link ResourceProfile} of the {@link SlotProfile} is the merged {@link ResourceProfile}s
+		 * of all executions sharing the slot.
+		 *
+		 * <p>The preferred locations of the {@link SlotProfile} is a union of the preferred locations
+		 * of all executions sharing the slot. The input locations within the bulk are ignored to avoid cyclic dependencies
+		 * within the region, e.g. in case of all-to-all pipelined connections, so that the allocations do not block each other.
+		 *
+		 * <p>The preferred {@link AllocationID}s of the {@link SlotProfile} are all previous {@link AllocationID}s
+		 * of all executions sharing the slot.
+		 *
+		 * <p>The {@link SlotProfile} also refers to all previous {@link AllocationID}s
+		 * of all executions within the bulk.
+		 *
+		 * @param executionSlotSharingGroup executions sharing the slot.
+		 * @return a future of the {@link SlotProfile} to allocate for the {@code executionSlotSharingGroup}.
+		 */
+		@Override
+		public CompletableFuture<SlotProfile> getSlotProfileFuture(ExecutionSlotSharingGroup executionSlotSharingGroup) {
+			ResourceProfile totalSlotResourceProfile = ResourceProfile.ZERO;
+			Collection<AllocationID> priorAllocations = new HashSet<>();
+			Collection<CompletableFuture<Collection<TaskManagerLocation>>> preferredLocationsPerExecution = new ArrayList<>();
+			for (ExecutionVertexID execution : executionSlotSharingGroup.getExecutionVertexIds()) {
+				totalSlotResourceProfile = totalSlotResourceProfile.merge(resourceProfileRetriever.apply(execution));
+				priorAllocations.add(priorAllocationIdRetriever.apply(execution));
+				preferredLocationsPerExecution.add(preferredLocationsRetriever
+					.getPreferredLocations(execution, producersToIgnore));
+			}
+
+			CompletableFuture<Collection<TaskManagerLocation>> preferredLocationsFuture = FutureUtils
+				.combineAll(preferredLocationsPerExecution)
+				.thenApply(executionPreferredLocations ->
+					executionPreferredLocations.stream().flatMap(Collection::stream).collect(Collectors.toList()));
+
+			ResourceProfile physicalSlotResourceProfile = totalSlotResourceProfile;
+			return preferredLocationsFuture.thenApply(
+				preferredLocations ->
+					SlotProfile.priorAllocation(
+						physicalSlotResourceProfile,
+						physicalSlotResourceProfile,
+						preferredLocations,
+						priorAllocations,
+						allBulkPriorAllocationIds));
+		}
+	}
+}

flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/SharedSlotProfileRetriever.java

@@ -0,0 +1,44 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.runtime.scheduler;
+
+import org.apache.flink.runtime.clusterframework.types.SlotProfile;
+import org.apache.flink.runtime.scheduler.strategy.ExecutionVertexID;
+
+import java.util.Set;
+import java.util.concurrent.CompletableFuture;
+
+/**
+ * Computes a {@link SlotProfile} to allocate a slot for executions, sharing the slot.
+ */
+@FunctionalInterface
+interface SharedSlotProfileRetriever {
+	/**
+	 * Computes a {@link SlotProfile} of an execution slot sharing group.
+	 *
+	 * @param executionSlotSharingGroup executions sharing the slot.
+	 * @return a future of the {@link SlotProfile} to allocate for the {@code executionSlotSharingGroup}.
+	 */
+	CompletableFuture<SlotProfile> getSlotProfileFuture(ExecutionSlotSharingGroup executionSlotSharingGroup);
+
+	@FunctionalInterface
+	interface SharedSlotProfileRetrieverFactory {
+		SharedSlotProfileRetriever createFromBulk(Set<ExecutionVertexID> bulk);
+	}
+}

flink-runtime/src/test/java/org/apache/flink/runtime/scheduler/MergingSharedSlotProfileRetrieverTest.java

@@ -0,0 +1,183 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.runtime.scheduler;
+
+import org.apache.flink.configuration.MemorySize;
+import org.apache.flink.runtime.clusterframework.types.AllocationID;
+import org.apache.flink.runtime.clusterframework.types.ResourceID;
+import org.apache.flink.runtime.clusterframework.types.ResourceProfile;
+import org.apache.flink.runtime.clusterframework.types.SlotProfile;
+import org.apache.flink.runtime.jobgraph.JobVertexID;
+import org.apache.flink.runtime.scheduler.strategy.ExecutionVertexID;
+import org.apache.flink.runtime.taskmanager.TaskManagerLocation;
+import org.apache.flink.util.FlinkRuntimeException;
+import org.apache.flink.util.Preconditions;
+
+import org.junit.Test;
+
+import java.net.InetAddress;
+import java.net.UnknownHostException;
+import java.util.Arrays;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ExecutionException;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+
+import static org.hamcrest.Matchers.containsInAnyOrder;
+import static org.hamcrest.Matchers.hasSize;
+import static org.hamcrest.Matchers.is;
+import static org.junit.Assert.assertThat;
+
+/**
+ * Tests for {@link org.apache.flink.runtime.scheduler.MergingSharedSlotProfileRetrieverFactory}.
+ */
+public class MergingSharedSlotProfileRetrieverTest {
+
+	private static final PreferredLocationsRetriever EMPTY_PREFERRED_LOCATIONS_RETRIEVER =
+		(executionVertexId, producersToIgnore) -> CompletableFuture.completedFuture(Collections.emptyList());
+
+	@Test
+	public void testGetEmptySlotProfile() throws ExecutionException, InterruptedException {
+		SharedSlotProfileRetriever sharedSlotProfileRetriever = new MergingSharedSlotProfileRetrieverFactory(
+			EMPTY_PREFERRED_LOCATIONS_RETRIEVER,
+			executionVertexID -> ResourceProfile.ZERO,
+			executionVertexID -> new AllocationID()
+		).createFromBulk(Collections.emptySet());
+
+		SlotProfile slotProfile = sharedSlotProfileRetriever.getSlotProfileFuture(new ExecutionSlotSharingGroup()).get();
+
+		assertThat(slotProfile.getTaskResourceProfile(), is(ResourceProfile.ZERO));
+		assertThat(slotProfile.getPhysicalSlotResourceProfile(), is(ResourceProfile.ZERO));
+		assertThat(slotProfile.getPreferredLocations(), hasSize(0));
+		assertThat(slotProfile.getPreferredAllocations(), hasSize(0));
+		assertThat(slotProfile.getPreviousExecutionGraphAllocations(), hasSize(0));
+	}
+
+	@Test
+	public void testResourceProfileOfSlotProfile() throws ExecutionException, InterruptedException {
+		ResourceProfile resourceProfile1 = ResourceProfile
+			.newBuilder()
+			.setCpuCores(1.0)
+			.setTaskHeapMemory(MemorySize.ofMebiBytes(1))
+			.build();
+		ResourceProfile resourceProfile2 = resourceProfile1.multiply(2);
+		ResourceProfile resourceProfile3 = resourceProfile1.multiply(3);
+		List<ResourceProfile> resourceProfiles = Arrays.asList(resourceProfile1, resourceProfile2, resourceProfile3);
+		ResourceProfile totalResourceProfile = resourceProfile1.merge(resourceProfile2);
+
+		SlotProfile slotProfile = getSlotProfile(
+			resourceProfiles,
+			Collections.nCopies(3, new AllocationID()),
+			2);
+
+		assertThat(slotProfile.getTaskResourceProfile(), is(totalResourceProfile));
+		assertThat(slotProfile.getPhysicalSlotResourceProfile(), is(totalResourceProfile));
+	}
+
+	@Test
+	public void testPreferredLocationsOfSlotProfile() throws ExecutionException, InterruptedException {
+		// preferred locations
+		List<ExecutionVertexID> executions = IntStream
+			.range(0, 3)
+			.mapToObj(i -> new ExecutionVertexID(new JobVertexID(), 0))
+			.collect(Collectors.toList());
+
+		List<TaskManagerLocation> allLocations = executions.stream().map(e -> createTaskManagerLocation()).collect(Collectors.toList());
+		Map<ExecutionVertexID, Collection<TaskManagerLocation>> locations = new HashMap<>();
+		locations.put(executions.get(0), Arrays.asList(allLocations.get(0), allLocations.get(1)));
+		locations.put(executions.get(1), Arrays.asList(allLocations.get(1), allLocations.get(2)));
+
+		SlotProfile slotProfile = getSlotProfile(
+			(executionVertexId, producersToIgnore) -> {
+				assertThat(producersToIgnore, containsInAnyOrder(executions.toArray()));
+				return CompletableFuture.completedFuture(locations.get(executionVertexId));
+			},
+			executions,
+			Collections.nCopies(3, ResourceProfile.ZERO),
+			Collections.nCopies(3, new AllocationID()),
+			2);
+
+		assertThat(slotProfile.getPreferredLocations().stream().filter(allLocations.get(0)::equals).count(), is(1L));
+		assertThat(slotProfile.getPreferredLocations().stream().filter(allLocations.get(1)::equals).count(), is(2L));
+		assertThat(slotProfile.getPreferredLocations().stream().filter(allLocations.get(2)::equals).count(), is(1L));
+	}
+
+	@Test
+	public void testAllocationIdsOfSlotProfile() throws ExecutionException, InterruptedException {
+		AllocationID prevAllocationID1 = new AllocationID();
+		AllocationID prevAllocationID2 = new AllocationID();
+		List<AllocationID> prevAllocationIDs = Arrays.asList(prevAllocationID1, prevAllocationID2, new AllocationID());
+
+		SlotProfile slotProfile = getSlotProfile(
+			Collections.nCopies(3, ResourceProfile.ZERO),
+			prevAllocationIDs,
+			2);
+
+		assertThat(slotProfile.getPreferredAllocations(), containsInAnyOrder(prevAllocationID1, prevAllocationID2));
+		assertThat(slotProfile.getPreviousExecutionGraphAllocations(), containsInAnyOrder(prevAllocationIDs.toArray()));
+	}
+
+	private static SlotProfile getSlotProfile(
+			List<ResourceProfile> resourceProfiles,
+			List<AllocationID> prevAllocationIDs,
+			int executionSlotSharingGroupSize) throws ExecutionException, InterruptedException {
+		List<ExecutionVertexID> executions = resourceProfiles.stream()
+			.map(stub -> new ExecutionVertexID(new JobVertexID(), 0))
+			.collect(Collectors.toList());
+		return getSlotProfile(
+			EMPTY_PREFERRED_LOCATIONS_RETRIEVER,
+			executions,
+			resourceProfiles,
+			prevAllocationIDs,
+			executionSlotSharingGroupSize);
+	}
+
+	private static SlotProfile getSlotProfile(
+			PreferredLocationsRetriever preferredLocationsRetriever,
+			List<ExecutionVertexID> executions,
+			List<ResourceProfile> resourceProfiles,
+			List<AllocationID> prevAllocationIDs,
+			int executionSlotSharingGroupSize) throws ExecutionException, InterruptedException {
+		Preconditions.checkArgument(resourceProfiles.size() == prevAllocationIDs.size());
+
+		SharedSlotProfileRetriever sharedSlotProfileRetriever = new MergingSharedSlotProfileRetrieverFactory(
+			preferredLocationsRetriever,
+			executionVertexID -> resourceProfiles.get(executions.indexOf(executionVertexID)),
+			executionVertexID -> prevAllocationIDs.get(executions.indexOf(executionVertexID))
+		).createFromBulk(new HashSet<>(executions));
+
+		ExecutionSlotSharingGroup executionSlotSharingGroup = new ExecutionSlotSharingGroup();
+		executions.stream().limit(executionSlotSharingGroupSize).forEach(executionSlotSharingGroup::addVertex);
+		return sharedSlotProfileRetriever.getSlotProfileFuture(executionSlotSharingGroup).get();
+	}
+
+	private static TaskManagerLocation createTaskManagerLocation() {
+		try {
+			return new TaskManagerLocation(ResourceID.generate(), InetAddress.getByAddress(new byte[] {1, 2, 3, 4}), 8888);
+		} catch (UnknownHostException e) {
+			throw new FlinkRuntimeException("unexpected", e);
+		}
+	}
+}