You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@flink.apache.org by ch...@apache.org on 2023/03/28 11:39:34 UTC
[flink] 08/09: [FLINK-31399] Add mechanics to the AdaptiveScheduler for handling job resource requirements
This is an automated email from the ASF dual-hosted git repository.
chesnay pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/flink.git
commit c6be7add044c00b8365efa822805742fc6fa383c
Author: David Moravek <dm...@apache.org>
AuthorDate: Mon Mar 27 13:23:47 2023 +0200
[FLINK-31399] Add mechanics to the AdaptiveScheduler for handling job resource requirements
---
.../scheduler/DefaultVertexParallelismStore.java | 37 +++
.../flink/runtime/scheduler/SchedulerNG.java | 26 +++
.../scheduler/adaptive/AdaptiveScheduler.java | 49 +++-
.../adaptive/AdaptiveSchedulerFactory.java | 2 +
.../scheduler/adaptive/JobGraphJobInformation.java | 4 +
.../runtime/scheduler/TestingSchedulerNG.java | 43 +++-
.../adaptive/AdaptiveSchedulerBuilder.java | 10 +
.../scheduler/adaptive/AdaptiveSchedulerTest.java | 256 ++++++++++++++++++++-
8 files changed, 413 insertions(+), 14 deletions(-)
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/DefaultVertexParallelismStore.java b/flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/DefaultVertexParallelismStore.java
index a922a9e5cc1..9541e1924a5 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/DefaultVertexParallelismStore.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/DefaultVertexParallelismStore.java
@@ -18,14 +18,51 @@
package org.apache.flink.runtime.scheduler;
+import org.apache.flink.runtime.jobgraph.JobResourceRequirements;
import org.apache.flink.runtime.jobgraph.JobVertexID;
import java.util.HashMap;
import java.util.Map;
import java.util.Optional;
+import java.util.function.Function;
/** Maintains the configured parallelisms for vertices, which should be defined by a scheduler. */
public class DefaultVertexParallelismStore implements MutableVertexParallelismStore {
+
+ private static final Function<Integer, Optional<String>> RESCALE_MAX_REJECT =
+ maxParallelism -> Optional.of("Cannot change the max parallelism.");
+
+ /**
+ * Create a new {@link VertexParallelismStore} that reflects given {@link
+ * JobResourceRequirements}.
+ *
+ * @param oldVertexParallelismStore old vertex parallelism store that serves as a base for the
+ * new one
+ * @param jobResourceRequirements to apply over the old vertex parallelism store
+ * @return new vertex parallelism store iff it was updated
+ */
+ public static Optional<VertexParallelismStore> applyJobResourceRequirements(
+ VertexParallelismStore oldVertexParallelismStore,
+ JobResourceRequirements jobResourceRequirements) {
+ final DefaultVertexParallelismStore newVertexParallelismStore =
+ new DefaultVertexParallelismStore();
+ boolean changed = false;
+ for (final JobVertexID jobVertexId : jobResourceRequirements.getJobVertices()) {
+ final VertexParallelismInformation oldVertexParallelismInfo =
+ oldVertexParallelismStore.getParallelismInfo(jobVertexId);
+ final int parallelism =
+ jobResourceRequirements.getParallelism(jobVertexId).getUpperBound();
+ newVertexParallelismStore.setParallelismInfo(
+ jobVertexId,
+ new DefaultVertexParallelismInfo(
+ parallelism,
+ oldVertexParallelismInfo.getMaxParallelism(),
+ RESCALE_MAX_REJECT));
+ changed |= oldVertexParallelismInfo.getParallelism() != parallelism;
+ }
+ return changed ? Optional.of(newVertexParallelismStore) : Optional.empty();
+ }
+
private final Map<JobVertexID, VertexParallelismInformation> vertexToParallelismInfo =
new HashMap<>();
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/SchedulerNG.java b/flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/SchedulerNG.java
index 70b0c3b8807..6630e3bd48b 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/SchedulerNG.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/SchedulerNG.java
@@ -35,6 +35,7 @@ import org.apache.flink.runtime.executiongraph.TaskExecutionStateTransition;
import org.apache.flink.runtime.io.network.partition.ResultPartitionID;
import org.apache.flink.runtime.jobgraph.IntermediateDataSetID;
import org.apache.flink.runtime.jobgraph.JobGraph;
+import org.apache.flink.runtime.jobgraph.JobResourceRequirements;
import org.apache.flink.runtime.jobgraph.JobVertexID;
import org.apache.flink.runtime.jobgraph.OperatorID;
import org.apache.flink.runtime.jobmanager.PartitionProducerDisposedException;
@@ -48,6 +49,7 @@ import org.apache.flink.runtime.operators.coordination.OperatorCoordinator;
import org.apache.flink.runtime.operators.coordination.OperatorEvent;
import org.apache.flink.runtime.query.KvStateLocation;
import org.apache.flink.runtime.query.UnknownKvStateLocation;
+import org.apache.flink.runtime.scheduler.adaptive.AdaptiveScheduler;
import org.apache.flink.runtime.state.KeyGroupRange;
import org.apache.flink.runtime.taskmanager.TaskExecutionState;
import org.apache.flink.util.AutoCloseableAsync;
@@ -185,4 +187,28 @@ public interface SchedulerNG extends GlobalFailureHandler, AutoCloseableAsync {
*/
CompletableFuture<CoordinationResponse> deliverCoordinationRequestToCoordinator(
OperatorID operator, CoordinationRequest request) throws FlinkException;
+
+ /**
+ * Read current {@link JobResourceRequirements job resource requirements}.
+ *
+ * @return Current resource requirements.
+ */
+ default JobResourceRequirements requestJobResourceRequirements() {
+ throw new UnsupportedOperationException(
+ String.format(
+ "The %s does not support changing the parallelism without a job restart. This feature is currently only expected to work with the %s.",
+ getClass().getSimpleName(), AdaptiveScheduler.class.getSimpleName()));
+ }
+
+ /**
+ * Update {@link JobResourceRequirements job resource requirements}.
+ *
+ * @param jobResourceRequirements new resource requirements
+ */
+ default void updateJobResourceRequirements(JobResourceRequirements jobResourceRequirements) {
+ throw new UnsupportedOperationException(
+ String.format(
+ "The %s does not support changing the parallelism without a job restart. This feature is currently only expected to work with the %s.",
+ getClass().getSimpleName(), AdaptiveScheduler.class.getSimpleName()));
+ }
}
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/adaptive/AdaptiveScheduler.java b/flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/adaptive/AdaptiveScheduler.java
index e3499f231bf..436e53b1851 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/adaptive/AdaptiveScheduler.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/adaptive/AdaptiveScheduler.java
@@ -61,6 +61,7 @@ import org.apache.flink.runtime.io.network.partition.ResultPartitionID;
import org.apache.flink.runtime.jobgraph.IntermediateDataSetID;
import org.apache.flink.runtime.jobgraph.JobEdge;
import org.apache.flink.runtime.jobgraph.JobGraph;
+import org.apache.flink.runtime.jobgraph.JobResourceRequirements;
import org.apache.flink.runtime.jobgraph.JobType;
import org.apache.flink.runtime.jobgraph.JobVertex;
import org.apache.flink.runtime.jobgraph.JobVertexID;
@@ -97,6 +98,7 @@ import org.apache.flink.runtime.scheduler.UpdateSchedulerNgOnInternalFailuresLis
import org.apache.flink.runtime.scheduler.VertexParallelismInformation;
import org.apache.flink.runtime.scheduler.VertexParallelismStore;
import org.apache.flink.runtime.scheduler.adaptive.allocator.JobAllocationsInformation;
+import org.apache.flink.runtime.scheduler.adaptive.allocator.JobInformation;
import org.apache.flink.runtime.scheduler.adaptive.allocator.ReservedSlots;
import org.apache.flink.runtime.scheduler.adaptive.allocator.SlotAllocator;
import org.apache.flink.runtime.scheduler.adaptive.allocator.VertexParallelism;
@@ -168,7 +170,7 @@ public class AdaptiveScheduler
private static final Logger LOG = LoggerFactory.getLogger(AdaptiveScheduler.class);
- private final JobGraphJobInformation jobInformation;
+ private final JobGraph jobGraph;
private final VertexParallelismStore initialParallelismStore;
private final DeclarativeSlotPool declarativeSlotPool;
@@ -217,12 +219,14 @@ public class AdaptiveScheduler
private final DeploymentStateTimeMetrics deploymentTimeMetrics;
private final BoundedFIFOQueue<RootExceptionHistoryEntry> exceptionHistory;
+ private JobGraphJobInformation jobInformation;
private ResourceCounter desiredResources = ResourceCounter.empty();
private final JobManagerJobMetricGroup jobManagerJobMetricGroup;
public AdaptiveScheduler(
JobGraph jobGraph,
+ @Nullable JobResourceRequirements jobResourceRequirements,
Configuration configuration,
DeclarativeSlotPool declarativeSlotPool,
SlotAllocator slotAllocator,
@@ -243,10 +247,18 @@ public class AdaptiveScheduler
assertPreconditions(jobGraph);
+ this.jobGraph = jobGraph;
this.executionMode = configuration.get(JobManagerOptions.SCHEDULER_MODE);
VertexParallelismStore vertexParallelismStore =
computeVertexParallelismStore(jobGraph, executionMode);
+ if (jobResourceRequirements != null) {
+ vertexParallelismStore =
+ DefaultVertexParallelismStore.applyJobResourceRequirements(
+ vertexParallelismStore, jobResourceRequirements)
+ .orElse(vertexParallelismStore);
+ }
+
this.initialParallelismStore = vertexParallelismStore;
this.jobInformation = new JobGraphJobInformation(jobGraph, vertexParallelismStore);
@@ -732,6 +744,35 @@ public class AdaptiveScheduler
+ " does not exist")));
}
+ @Override
+ public JobResourceRequirements requestJobResourceRequirements() {
+ final JobResourceRequirements.Builder builder = JobResourceRequirements.newBuilder();
+ for (JobInformation.VertexInformation vertex : jobInformation.getVertices()) {
+ builder.setParallelismForJobVertex(vertex.getJobVertexID(), 1, vertex.getParallelism());
+ }
+ return builder.build();
+ }
+
+ @Override
+ public void updateJobResourceRequirements(JobResourceRequirements jobResourceRequirements) {
+ if (executionMode == SchedulerExecutionMode.REACTIVE) {
+ throw new UnsupportedOperationException(
+ "Cannot change the parallelism of a job running in reactive mode.");
+ }
+ final Optional<VertexParallelismStore> maybeUpdateVertexParallelismStore =
+ DefaultVertexParallelismStore.applyJobResourceRequirements(
+ jobInformation.getVertexParallelismStore(), jobResourceRequirements);
+ if (maybeUpdateVertexParallelismStore.isPresent()) {
+ this.jobInformation =
+ new JobGraphJobInformation(jobGraph, maybeUpdateVertexParallelismStore.get());
+ declareDesiredResources();
+ state.tryRun(
+ ResourceListener.class,
+ ResourceListener::onNewResourceRequirements,
+ "Current state does not react to desired parallelism changes.");
+ }
+ }
+
// ----------------------------------------------------------------
@Override
@@ -1005,10 +1046,8 @@ public class AdaptiveScheduler
executionGraphWithVertexParallelism.getJobSchedulingPlan();
return slotAllocator
.tryReserveResources(jobSchedulingPlan)
- .map(
- reservedSlots ->
- CreatingExecutionGraph.AssignmentResult.success(
- assignSlotsToExecutionGraph(executionGraph, reservedSlots)))
+ .map(reservedSlots -> assignSlotsToExecutionGraph(executionGraph, reservedSlots))
+ .map(CreatingExecutionGraph.AssignmentResult::success)
.orElseGet(CreatingExecutionGraph.AssignmentResult::notPossible);
}
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/adaptive/AdaptiveSchedulerFactory.java b/flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/adaptive/AdaptiveSchedulerFactory.java
index 2d59d56466d..1d6d91d0476 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/adaptive/AdaptiveSchedulerFactory.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/adaptive/AdaptiveSchedulerFactory.java
@@ -31,6 +31,7 @@ import org.apache.flink.runtime.executiongraph.failover.flip1.RestartBackoffTime
import org.apache.flink.runtime.executiongraph.failover.flip1.RestartBackoffTimeStrategyFactoryLoader;
import org.apache.flink.runtime.io.network.partition.JobMasterPartitionTracker;
import org.apache.flink.runtime.jobgraph.JobGraph;
+import org.apache.flink.runtime.jobgraph.JobResourceRequirements;
import org.apache.flink.runtime.jobmaster.ExecutionDeploymentTracker;
import org.apache.flink.runtime.jobmaster.slotpool.DeclarativeSlotPool;
import org.apache.flink.runtime.jobmaster.slotpool.SlotPoolService;
@@ -123,6 +124,7 @@ public class AdaptiveSchedulerFactory implements SchedulerNGFactory {
return new AdaptiveScheduler(
jobGraph,
+ JobResourceRequirements.readFromJobGraph(jobGraph).orElse(null),
jobMasterConfiguration,
declarativeSlotPool,
slotAllocator,
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/adaptive/JobGraphJobInformation.java b/flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/adaptive/JobGraphJobInformation.java
index 2111b6e9c07..ea1d699c831 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/adaptive/JobGraphJobInformation.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/adaptive/JobGraphJobInformation.java
@@ -84,6 +84,10 @@ public class JobGraphJobInformation implements JobInformation {
return InstantiationUtil.clone(jobGraph);
}
+ public VertexParallelismStore getVertexParallelismStore() {
+ return vertexParallelismStore;
+ }
+
private static final class JobVertexInformation implements JobInformation.VertexInformation {
private final JobVertex jobVertex;
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/scheduler/TestingSchedulerNG.java b/flink-runtime/src/test/java/org/apache/flink/runtime/scheduler/TestingSchedulerNG.java
index c54511b7b3b..bde33ea3f41 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/scheduler/TestingSchedulerNG.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/scheduler/TestingSchedulerNG.java
@@ -32,6 +32,7 @@ import org.apache.flink.runtime.executiongraph.ExecutionAttemptID;
import org.apache.flink.runtime.executiongraph.TaskExecutionStateTransition;
import org.apache.flink.runtime.io.network.partition.ResultPartitionID;
import org.apache.flink.runtime.jobgraph.IntermediateDataSetID;
+import org.apache.flink.runtime.jobgraph.JobResourceRequirements;
import org.apache.flink.runtime.jobgraph.JobVertexID;
import org.apache.flink.runtime.jobgraph.OperatorID;
import org.apache.flink.runtime.jobmanager.PartitionProducerDisposedException;
@@ -65,6 +66,8 @@ public class TestingSchedulerNG implements SchedulerNG {
private final Function<CheckpointType, CompletableFuture<CompletedCheckpoint>>
triggerCheckpointFunction;
private final Consumer<Throwable> handleGlobalFailureConsumer;
+ private final Supplier<JobResourceRequirements> requestJobResourceRequirementsSupplier;
+ private final Consumer<JobResourceRequirements> updateJobResourceRequirementsConsumer;
private TestingSchedulerNG(
CompletableFuture<JobStatus> jobTerminationFuture,
@@ -74,13 +77,17 @@ public class TestingSchedulerNG implements SchedulerNG {
triggerSavepointFunction,
Function<CheckpointType, CompletableFuture<CompletedCheckpoint>>
triggerCheckpointFunction,
- Consumer<Throwable> handleGlobalFailureConsumer) {
+ Consumer<Throwable> handleGlobalFailureConsumer,
+ Supplier<JobResourceRequirements> requestJobResourceRequirementsSupplier,
+ Consumer<JobResourceRequirements> updateJobResourceRequirementsConsumer) {
this.jobTerminationFuture = jobTerminationFuture;
this.startSchedulingRunnable = startSchedulingRunnable;
this.closeAsyncSupplier = closeAsyncSupplier;
this.triggerSavepointFunction = triggerSavepointFunction;
this.triggerCheckpointFunction = triggerCheckpointFunction;
this.handleGlobalFailureConsumer = handleGlobalFailureConsumer;
+ this.requestJobResourceRequirementsSupplier = requestJobResourceRequirementsSupplier;
+ this.updateJobResourceRequirementsConsumer = updateJobResourceRequirementsConsumer;
}
@Override
@@ -232,6 +239,16 @@ public class TestingSchedulerNG implements SchedulerNG {
long checkpointId,
CheckpointMetrics checkpointMetrics) {}
+ @Override
+ public JobResourceRequirements requestJobResourceRequirements() {
+ return requestJobResourceRequirementsSupplier.get();
+ }
+
+ @Override
+ public void updateJobResourceRequirements(JobResourceRequirements jobResourceRequirements) {
+ updateJobResourceRequirementsConsumer.accept(jobResourceRequirements);
+ }
+
public static Builder newBuilder() {
return new Builder();
}
@@ -248,6 +265,14 @@ public class TestingSchedulerNG implements SchedulerNG {
private Function<CheckpointType, CompletableFuture<CompletedCheckpoint>>
triggerCheckpointFunction = (ignored) -> new CompletableFuture<>();
private Consumer<Throwable> handleGlobalFailureConsumer = (ignored) -> {};
+ private Supplier<JobResourceRequirements> requestJobResourceRequirementsSupplier =
+ () -> {
+ throw new UnsupportedOperationException("Not supported.");
+ };
+ private Consumer<JobResourceRequirements> updateJobResourceRequirementsConsumer =
+ ignored -> {
+ throw new UnsupportedOperationException("Not supported.");
+ };
public Builder setJobTerminationFuture(CompletableFuture<JobStatus> jobTerminationFuture) {
this.jobTerminationFuture = jobTerminationFuture;
@@ -284,6 +309,18 @@ public class TestingSchedulerNG implements SchedulerNG {
return this;
}
+ public Builder setRequestJobResourceRequirementsSupplier(
+ Supplier<JobResourceRequirements> requestJobResourceRequirementsSupplier) {
+ this.requestJobResourceRequirementsSupplier = requestJobResourceRequirementsSupplier;
+ return this;
+ }
+
+ public Builder setUpdateJobResourceRequirementsConsumer(
+ Consumer<JobResourceRequirements> updateJobResourceRequirementsConsumer) {
+ this.updateJobResourceRequirementsConsumer = updateJobResourceRequirementsConsumer;
+ return this;
+ }
+
public TestingSchedulerNG build() {
return new TestingSchedulerNG(
jobTerminationFuture,
@@ -291,7 +328,9 @@ public class TestingSchedulerNG implements SchedulerNG {
closeAsyncSupplier,
triggerSavepointFunction,
triggerCheckpointFunction,
- handleGlobalFailureConsumer);
+ handleGlobalFailureConsumer,
+ requestJobResourceRequirementsSupplier,
+ updateJobResourceRequirementsConsumer);
}
}
}
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/scheduler/adaptive/AdaptiveSchedulerBuilder.java b/flink-runtime/src/test/java/org/apache/flink/runtime/scheduler/adaptive/AdaptiveSchedulerBuilder.java
index 8f3c65a63a1..32231edf0b2 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/scheduler/adaptive/AdaptiveSchedulerBuilder.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/scheduler/adaptive/AdaptiveSchedulerBuilder.java
@@ -32,6 +32,7 @@ import org.apache.flink.runtime.executiongraph.failover.flip1.RestartBackoffTime
import org.apache.flink.runtime.io.network.partition.JobMasterPartitionTracker;
import org.apache.flink.runtime.io.network.partition.NoOpJobMasterPartitionTracker;
import org.apache.flink.runtime.jobgraph.JobGraph;
+import org.apache.flink.runtime.jobgraph.JobResourceRequirements;
import org.apache.flink.runtime.jobmaster.DefaultExecutionDeploymentTracker;
import org.apache.flink.runtime.jobmaster.slotpool.DeclarativeSlotPool;
import org.apache.flink.runtime.jobmaster.slotpool.DefaultAllocatedSlotPool;
@@ -58,6 +59,8 @@ public class AdaptiveSchedulerBuilder {
private final ComponentMainThreadExecutor mainThreadExecutor;
+ @Nullable private JobResourceRequirements jobResourceRequirements;
+
private Configuration jobMasterConfiguration = new Configuration();
private ClassLoader userCodeLoader = ClassLoader.getSystemClassLoader();
private CheckpointsCleaner checkpointsCleaner = new CheckpointsCleaner();
@@ -95,6 +98,12 @@ public class AdaptiveSchedulerBuilder {
rpcTimeout);
}
+ public AdaptiveSchedulerBuilder setJobResourceRequirements(
+ JobResourceRequirements jobResourceRequirements) {
+ this.jobResourceRequirements = jobResourceRequirements;
+ return this;
+ }
+
public AdaptiveSchedulerBuilder setJobMasterConfiguration(
final Configuration jobMasterConfiguration) {
this.jobMasterConfiguration = jobMasterConfiguration;
@@ -193,6 +202,7 @@ public class AdaptiveSchedulerBuilder {
return new AdaptiveScheduler(
jobGraph,
+ jobResourceRequirements,
jobMasterConfiguration,
declarativeSlotPool,
slotAllocator == null
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/scheduler/adaptive/AdaptiveSchedulerTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/scheduler/adaptive/AdaptiveSchedulerTest.java
index e9fa605e09b..6dcf71b3032 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/scheduler/adaptive/AdaptiveSchedulerTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/scheduler/adaptive/AdaptiveSchedulerTest.java
@@ -36,6 +36,7 @@ import org.apache.flink.runtime.checkpoint.StandaloneCompletedCheckpointStore;
import org.apache.flink.runtime.checkpoint.TestingCheckpointIDCounter;
import org.apache.flink.runtime.checkpoint.TestingCheckpointRecoveryFactory;
import org.apache.flink.runtime.checkpoint.TestingCompletedCheckpointStore;
+import org.apache.flink.runtime.clusterframework.types.AllocationID;
import org.apache.flink.runtime.clusterframework.types.ResourceProfile;
import org.apache.flink.runtime.concurrent.ComponentMainThreadExecutor;
import org.apache.flink.runtime.concurrent.ComponentMainThreadExecutorServiceAdapter;
@@ -59,13 +60,18 @@ import org.apache.flink.runtime.executiongraph.utils.SimpleAckingTaskManagerGate
import org.apache.flink.runtime.io.network.partition.ResultPartitionID;
import org.apache.flink.runtime.jobgraph.IntermediateDataSetID;
import org.apache.flink.runtime.jobgraph.JobGraph;
+import org.apache.flink.runtime.jobgraph.JobResourceRequirements;
import org.apache.flink.runtime.jobgraph.JobVertex;
+import org.apache.flink.runtime.jobgraph.JobVertexResourceRequirements;
import org.apache.flink.runtime.jobgraph.OperatorID;
import org.apache.flink.runtime.jobgraph.tasks.CheckpointCoordinatorConfiguration;
import org.apache.flink.runtime.jobgraph.tasks.JobCheckpointingSettings;
import org.apache.flink.runtime.jobmanager.PartitionProducerDisposedException;
+import org.apache.flink.runtime.jobmanager.slots.TaskManagerGateway;
+import org.apache.flink.runtime.jobmaster.slotpool.DeclarativeSlotPool;
import org.apache.flink.runtime.jobmaster.slotpool.DefaultAllocatedSlotPool;
import org.apache.flink.runtime.jobmaster.slotpool.DefaultDeclarativeSlotPool;
+import org.apache.flink.runtime.messages.Acknowledge;
import org.apache.flink.runtime.messages.checkpoint.DeclineCheckpoint;
import org.apache.flink.runtime.metrics.MetricNames;
import org.apache.flink.runtime.metrics.MetricRegistry;
@@ -95,6 +101,7 @@ import org.apache.flink.testutils.TestingUtils;
import org.apache.flink.testutils.executor.TestExecutorExtension;
import org.apache.flink.util.FlinkException;
import org.apache.flink.util.IterableUtils;
+import org.apache.flink.util.Preconditions;
import org.apache.flink.util.jackson.JacksonMapperFactory;
import org.junit.jupiter.api.Test;
@@ -102,7 +109,6 @@ import org.junit.jupiter.api.extension.RegisterExtension;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
-import javax.annotation.Nonnull;
import javax.annotation.Nullable;
import java.io.IOException;
@@ -124,6 +130,7 @@ import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.BiConsumer;
+import java.util.function.BiFunction;
import java.util.function.Consumer;
import java.util.stream.Collectors;
@@ -140,6 +147,7 @@ import static org.assertj.core.api.Assertions.assertThatThrownBy;
/** Tests for the {@link AdaptiveScheduler}. */
public class AdaptiveSchedulerTest {
+ private static final Duration DEFAULT_TIMEOUT = Duration.ofHours(1);
private static final int PARALLELISM = 4;
private static final JobVertex JOB_VERTEX = createNoOpVertex("v1", PARALLELISM);
@@ -940,6 +948,137 @@ public class AdaptiveSchedulerTest {
assertThat(resubmittedArchivedVertex.getMaxParallelism()).isEqualTo(expectedMaxParallelism);
}
+ @Test
+ void testRequirementIncreaseTriggersScaleUp() throws Exception {
+ final JobGraph jobGraph = createJobGraph();
+
+ final DefaultDeclarativeSlotPool declarativeSlotPool =
+ createDeclarativeSlotPool(jobGraph.getJobID());
+
+ final AdaptiveScheduler scheduler =
+ createSchedulerWithNoResourceWaitTimeout(jobGraph, declarativeSlotPool);
+
+ final int scaledUpParallelism = PARALLELISM * 2;
+
+ final SubmissionBufferingTaskManagerGateway taskManagerGateway =
+ createSubmissionBufferingTaskManagerGateway(scaledUpParallelism, scheduler);
+
+ startJobWithSlotsMatchingParallelism(
+ scheduler, declarativeSlotPool, taskManagerGateway, PARALLELISM);
+ awaitJobReachingParallelism(taskManagerGateway, scheduler, PARALLELISM);
+
+ JobResourceRequirements newJobResourceRequirements =
+ createRequirementsWithUpperParallelism(scaledUpParallelism);
+ singleThreadMainThreadExecutor.execute(
+ () -> {
+ // first update requirements as otherwise slots are rejected!
+ scheduler.updateJobResourceRequirements(newJobResourceRequirements);
+ offerSlots(
+ declarativeSlotPool,
+ createSlotOffersForResourceRequirements(
+ ResourceCounter.withResource(
+ ResourceProfile.UNKNOWN, PARALLELISM)),
+ taskManagerGateway);
+ });
+
+ awaitJobReachingParallelism(taskManagerGateway, scheduler, scaledUpParallelism);
+ }
+
+ @Test
+ void testRequirementDecreaseTriggersScaleDown() throws Exception {
+ final JobGraph jobGraph = createJobGraph();
+
+ final DefaultDeclarativeSlotPool declarativeSlotPool =
+ createDeclarativeSlotPool(jobGraph.getJobID());
+
+ final AdaptiveScheduler scheduler =
+ createSchedulerWithNoResourceWaitTimeout(jobGraph, declarativeSlotPool);
+
+ final SubmissionBufferingTaskManagerGateway taskManagerGateway =
+ createSubmissionBufferingTaskManagerGateway(PARALLELISM, scheduler);
+
+ startJobWithSlotsMatchingParallelism(
+ scheduler, declarativeSlotPool, taskManagerGateway, PARALLELISM);
+ awaitJobReachingParallelism(taskManagerGateway, scheduler, PARALLELISM);
+
+ int scaledDownParallelism = PARALLELISM - 1;
+ JobResourceRequirements newJobResourceRequirements =
+ createRequirementsWithUpperParallelism(scaledDownParallelism);
+ singleThreadMainThreadExecutor.execute(
+ () -> scheduler.updateJobResourceRequirements(newJobResourceRequirements));
+
+ awaitJobReachingParallelism(taskManagerGateway, scheduler, scaledDownParallelism);
+ }
+
+ private AdaptiveScheduler createSchedulerWithNoResourceWaitTimeout(
+ JobGraph jobGraph, DeclarativeSlotPool declarativeSlotPool) throws Exception {
+ final Configuration configuration = new Configuration();
+ configuration.set(JobManagerOptions.RESOURCE_WAIT_TIMEOUT, Duration.ofMillis(1L));
+
+ return new AdaptiveSchedulerBuilder(jobGraph, singleThreadMainThreadExecutor)
+ .setDeclarativeSlotPool(declarativeSlotPool)
+ .setJobMasterConfiguration(configuration)
+ .build(EXECUTOR_RESOURCE.getExecutor());
+ }
+
+ private SubmissionBufferingTaskManagerGateway createSubmissionBufferingTaskManagerGateway(
+ int parallelism, SchedulerNG scheduler) {
+ SubmissionBufferingTaskManagerGateway taskManagerGateway =
+ new SubmissionBufferingTaskManagerGateway(parallelism);
+ taskManagerGateway.setCancelConsumer(
+ executionAttemptID ->
+ singleThreadMainThreadExecutor.execute(
+ () ->
+ scheduler.updateTaskExecutionState(
+ new TaskExecutionState(
+ executionAttemptID,
+ ExecutionState.CANCELED))));
+ return taskManagerGateway;
+ }
+
+ private void startJobWithSlotsMatchingParallelism(
+ SchedulerNG scheduler,
+ DeclarativeSlotPool declarativeSlotPool,
+ TaskManagerGateway taskManagerGateway,
+ int parallelism) {
+ singleThreadMainThreadExecutor.execute(
+ () -> {
+ scheduler.startScheduling();
+ offerSlots(
+ declarativeSlotPool,
+ createSlotOffersForResourceRequirements(
+ ResourceCounter.withResource(
+ ResourceProfile.UNKNOWN, parallelism)),
+ taskManagerGateway);
+ });
+ }
+
+ private void awaitJobReachingParallelism(
+ SubmissionBufferingTaskManagerGateway taskManagerGateway,
+ SchedulerNG scheduler,
+ int parallelism)
+ throws Exception {
+ // Wait for all tasks to be submitted
+ taskManagerGateway.waitForSubmissions(parallelism);
+
+ final ArchivedExecutionGraph executionGraph =
+ CompletableFuture.supplyAsync(
+ () -> scheduler.requestJob().getArchivedExecutionGraph(),
+ singleThreadMainThreadExecutor)
+ .get();
+
+ assertThat(executionGraph.getJobVertex(JOB_VERTEX.getID()).getParallelism())
+ .isEqualTo(parallelism);
+ }
+
+ private static JobResourceRequirements createRequirementsWithUpperParallelism(int parallelism) {
+ return new JobResourceRequirements(
+ Collections.singletonMap(
+ JOB_VERTEX.getID(),
+ new JobVertexResourceRequirements(
+ new JobVertexResourceRequirements.Parallelism(1, parallelism))));
+ }
+
// ---------------------------------------------------------------------------------------------
// Failure handling tests
// ---------------------------------------------------------------------------------------------
@@ -1479,6 +1618,75 @@ public class AdaptiveSchedulerTest {
}
}
+ @Test
+ public void testUpdateResourceRequirementsInReactiveModeIsNotSupported() throws Exception {
+ final Configuration configuration = new Configuration();
+ configuration.set(JobManagerOptions.SCHEDULER_MODE, SchedulerExecutionMode.REACTIVE);
+ final AdaptiveScheduler scheduler =
+ new AdaptiveSchedulerBuilder(createJobGraph(), mainThreadExecutor)
+ .setJobMasterConfiguration(configuration)
+ .build(EXECUTOR_RESOURCE.getExecutor());
+ assertThatThrownBy(
+ () ->
+ scheduler.updateJobResourceRequirements(
+ JobResourceRequirements.empty()))
+ .isInstanceOf(UnsupportedOperationException.class);
+ }
+
+ @Test
+ public void testRequestDefaultResourceRequirements() throws Exception {
+ final JobGraph jobGraph = createJobGraph();
+ final Configuration configuration = new Configuration();
+ final AdaptiveScheduler scheduler =
+ new AdaptiveSchedulerBuilder(jobGraph, mainThreadExecutor)
+ .setJobMasterConfiguration(configuration)
+ .build(EXECUTOR_RESOURCE.getExecutor());
+ assertThat(scheduler.requestJobResourceRequirements())
+ .isEqualTo(
+ JobResourceRequirements.newBuilder()
+ .setParallelismForJobVertex(
+ JOB_VERTEX.getID(), 1, JOB_VERTEX.getParallelism())
+ .build());
+ }
+
+ @Test
+ public void testRequestDefaultResourceRequirementsInReactiveMode() throws Exception {
+ final JobGraph jobGraph = createJobGraph();
+ final Configuration configuration = new Configuration();
+ configuration.set(JobManagerOptions.SCHEDULER_MODE, SchedulerExecutionMode.REACTIVE);
+ final AdaptiveScheduler scheduler =
+ new AdaptiveSchedulerBuilder(jobGraph, mainThreadExecutor)
+ .setJobMasterConfiguration(configuration)
+ .build(EXECUTOR_RESOURCE.getExecutor());
+ assertThat(scheduler.requestJobResourceRequirements())
+ .isEqualTo(
+ JobResourceRequirements.newBuilder()
+ .setParallelismForJobVertex(
+ JOB_VERTEX.getID(),
+ 1,
+ SchedulerBase.getDefaultMaxParallelism(JOB_VERTEX))
+ .build());
+ }
+
+ @Test
+ public void testRequestUpdatedResourceRequirements() throws Exception {
+ final JobGraph jobGraph = createJobGraph();
+ final Configuration configuration = new Configuration();
+ final AdaptiveScheduler scheduler =
+ new AdaptiveSchedulerBuilder(jobGraph, mainThreadExecutor)
+ .setJobMasterConfiguration(configuration)
+ .build(EXECUTOR_RESOURCE.getExecutor());
+ final JobResourceRequirements newJobResourceRequirements =
+ JobResourceRequirements.newBuilder()
+ .setParallelismForJobVertex(JOB_VERTEX.getID(), 1, 12)
+ .build();
+ assertThat(scheduler.requestJobResourceRequirements())
+ .isNotEqualTo(newJobResourceRequirements);
+ scheduler.updateJobResourceRequirements(newJobResourceRequirements);
+ assertThat(scheduler.requestJobResourceRequirements())
+ .isEqualTo(newJobResourceRequirements);
+ }
+
// ---------------------------------------------------------------------------------------------
// Utils
// ---------------------------------------------------------------------------------------------
@@ -1505,14 +1713,18 @@ public class AdaptiveSchedulerTest {
executionAttemptId, ExecutionState.CANCELED)));
}
- @Nonnull
private static DefaultDeclarativeSlotPool createDeclarativeSlotPool(JobID jobId) {
+ return createDeclarativeSlotPool(jobId, DEFAULT_TIMEOUT);
+ }
+
+ private static DefaultDeclarativeSlotPool createDeclarativeSlotPool(
+ JobID jobId, Duration idleSlotTimeout) {
return new DefaultDeclarativeSlotPool(
jobId,
new DefaultAllocatedSlotPool(),
ignored -> {},
- Time.minutes(10),
- Time.minutes(10));
+ Time.fromDuration(idleSlotTimeout),
+ Time.fromDuration(DEFAULT_TIMEOUT));
}
private static JobGraph createJobGraph() {
@@ -1542,17 +1754,32 @@ public class AdaptiveSchedulerTest {
public static class SubmissionBufferingTaskManagerGateway
extends SimpleAckingTaskManagerGateway {
final BlockingQueue<TaskDeploymentDescriptor> submittedTasks;
+ final BlockingQueue<AllocationID> freedSlots;
public SubmissionBufferingTaskManagerGateway(int capacity) {
submittedTasks = new ArrayBlockingQueue<>(capacity);
- setSubmitConsumer(ignored -> {});
+ freedSlots = new ArrayBlockingQueue<>(capacity);
+ initializeFunctions();
}
@Override
public void setSubmitConsumer(Consumer<TaskDeploymentDescriptor> submitConsumer) {
super.setSubmitConsumer(
- ((Consumer<TaskDeploymentDescriptor>) submittedTasks::offer)
- .andThen(submitConsumer));
+ taskDeploymentDescriptor -> {
+ Preconditions.checkState(submittedTasks.offer(taskDeploymentDescriptor));
+ submitConsumer.accept(taskDeploymentDescriptor);
+ });
+ }
+
+ @Override
+ public void setFreeSlotFunction(
+ BiFunction<AllocationID, Throwable, CompletableFuture<Acknowledge>>
+ freeSlotFunction) {
+ super.setFreeSlotFunction(
+ (allocationID, throwable) -> {
+ Preconditions.checkState(freedSlots.offer(allocationID));
+ return freeSlotFunction.apply(allocationID, throwable);
+ });
}
/**
@@ -1570,6 +1797,21 @@ public class AdaptiveSchedulerTest {
}
return descriptors;
}
+
+ public List<AllocationID> waitForFreedSlots(int numFreedSlots) throws InterruptedException {
+ final List<AllocationID> allocationIds = new ArrayList<>();
+ for (int i = 0; i < numFreedSlots; i++) {
+ allocationIds.add(freedSlots.take());
+ }
+ return allocationIds;
+ }
+
+ private void initializeFunctions() {
+ setSubmitConsumer(ignored -> {});
+ setFreeSlotFunction(
+ (allocationId, throwable) ->
+ CompletableFuture.completedFuture(Acknowledge.get()));
+ }
}
private static class StateInstanceFactory