You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@flink.apache.org by tr...@apache.org on 2020/10/09 14:50:03 UTC
[flink] 08/09: [FLINK-18721][yarn] Introduce
YarnResourceManagerDriver
This is an automated email from the ASF dual-hosted git repository.
trohrmann pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/flink.git
commit 8ee269c96445c317c4bcc767fcaf0608136a16b6
Author: Yangze Guo <ka...@gmail.com>
AuthorDate: Mon Aug 31 10:05:12 2020 +0800
[FLINK-18721][yarn] Introduce YarnResourceManagerDriver
---
.../KubernetesResourceManagerFactory.java | 2 +-
.../active/ActiveResourceManagerFactory.java | 4 +-
.../active/ResourceManagerDriverTestBase.java | 2 +-
.../src/main/java/org/apache/flink/yarn/Utils.java | 24 +
.../flink/yarn/YarnResourceManagerDriver.java | 611 +++++++++++++++++++++
.../flink/yarn/TestingYarnAMRMClientAsync.java | 46 +-
.../flink/yarn/TestingYarnNMClientAsync.java | 46 +-
.../flink/yarn/YarnResourceManagerDriverTest.java | 559 +++++++++++++++++++
8 files changed, 1274 insertions(+), 20 deletions(-)
diff --git a/flink-kubernetes/src/main/java/org/apache/flink/kubernetes/entrypoint/KubernetesResourceManagerFactory.java b/flink-kubernetes/src/main/java/org/apache/flink/kubernetes/entrypoint/KubernetesResourceManagerFactory.java
index 2266c0d..110141d 100644
--- a/flink-kubernetes/src/main/java/org/apache/flink/kubernetes/entrypoint/KubernetesResourceManagerFactory.java
+++ b/flink-kubernetes/src/main/java/org/apache/flink/kubernetes/entrypoint/KubernetesResourceManagerFactory.java
@@ -47,7 +47,7 @@ public class KubernetesResourceManagerFactory extends ActiveResourceManagerFacto
}
@Override
- protected ResourceManagerDriver<KubernetesWorkerNode> createResourceManagerDriver(Configuration configuration) {
+ protected ResourceManagerDriver<KubernetesWorkerNode> createResourceManagerDriver(Configuration configuration, String webInterfaceUrl, String rpcAddress) {
final KubernetesResourceManagerDriverConfiguration kubernetesResourceManagerDriverConfiguration =
new KubernetesResourceManagerDriverConfiguration(
configuration.getString(KubernetesConfigOptions.CLUSTER_ID),
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/resourcemanager/active/ActiveResourceManagerFactory.java b/flink-runtime/src/main/java/org/apache/flink/runtime/resourcemanager/active/ActiveResourceManagerFactory.java
index f0c68a0..32b05a6 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/resourcemanager/active/ActiveResourceManagerFactory.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/resourcemanager/active/ActiveResourceManagerFactory.java
@@ -95,7 +95,7 @@ public abstract class ActiveResourceManagerFactory<WorkerType extends ResourceID
Executor ioExecutor) {
return new ActiveResourceManager<>(
- createResourceManagerDriver(configuration),
+ createResourceManagerDriver(configuration, webInterfaceUrl, rpcService.getAddress()),
configuration,
rpcService,
resourceId,
@@ -110,5 +110,5 @@ public abstract class ActiveResourceManagerFactory<WorkerType extends ResourceID
ioExecutor);
}
- protected abstract ResourceManagerDriver<WorkerType> createResourceManagerDriver(Configuration configuration);
+ protected abstract ResourceManagerDriver<WorkerType> createResourceManagerDriver(Configuration configuration, String webInterfaceUrl, String rpcAddress);
}
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/resourcemanager/active/ResourceManagerDriverTestBase.java b/flink-runtime/src/test/java/org/apache/flink/runtime/resourcemanager/active/ResourceManagerDriverTestBase.java
index 476a692..d1598ae 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/resourcemanager/active/ResourceManagerDriverTestBase.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/resourcemanager/active/ResourceManagerDriverTestBase.java
@@ -180,7 +180,7 @@ public abstract class ResourceManagerDriverTestBase<WorkerType extends ResourceI
assertThat(Thread.currentThread().getName(), is(MAIN_THREAD_NAME));
}
- protected abstract void prepareRunTest();
+ protected abstract void prepareRunTest() throws Exception;
protected abstract ResourceManagerDriver<WorkerType> createResourceManagerDriver();
diff --git a/flink-yarn/src/main/java/org/apache/flink/yarn/Utils.java b/flink-yarn/src/main/java/org/apache/flink/yarn/Utils.java
index 5ae0216..c5a4056 100644
--- a/flink-yarn/src/main/java/org/apache/flink/yarn/Utils.java
+++ b/flink-yarn/src/main/java/org/apache/flink/yarn/Utils.java
@@ -560,6 +560,30 @@ public final class Utils {
ExternalResourceUtils.getExternalResources(flinkConfig, YarnConfigOptions.EXTERNAL_RESOURCE_YARN_CONFIG_KEY_SUFFIX));
}
+ static TaskExecutorProcessSpecContainerResourceAdapter createTaskExecutorProcessSpecContainerResourceAdapter(
+ org.apache.flink.configuration.Configuration flinkConfig,
+ YarnConfiguration yarnConfig) {
+
+ Resource unitResource = getUnitResource(yarnConfig);
+
+ return new TaskExecutorProcessSpecContainerResourceAdapter(
+ yarnConfig.getInt(
+ YarnConfiguration.RM_SCHEDULER_MINIMUM_ALLOCATION_MB,
+ YarnConfiguration.DEFAULT_RM_SCHEDULER_MINIMUM_ALLOCATION_MB),
+ yarnConfig.getInt(
+ YarnConfiguration.RM_SCHEDULER_MINIMUM_ALLOCATION_VCORES,
+ YarnConfiguration.DEFAULT_RM_SCHEDULER_MINIMUM_ALLOCATION_VCORES),
+ yarnConfig.getInt(
+ YarnConfiguration.RM_SCHEDULER_MAXIMUM_ALLOCATION_MB,
+ YarnConfiguration.DEFAULT_RM_SCHEDULER_MAXIMUM_ALLOCATION_MB),
+ yarnConfig.getInt(
+ YarnConfiguration.RM_SCHEDULER_MAXIMUM_ALLOCATION_VCORES,
+ YarnConfiguration.DEFAULT_RM_SCHEDULER_MAXIMUM_ALLOCATION_VCORES),
+ unitResource.getMemory(),
+ unitResource.getVirtualCores(),
+ ExternalResourceUtils.getExternalResources(flinkConfig, YarnConfigOptions.EXTERNAL_RESOURCE_YARN_CONFIG_KEY_SUFFIX));
+ }
+
@VisibleForTesting
static Resource getUnitResource(YarnConfiguration yarnConfig) {
final int unitMemMB, unitVcore;
diff --git a/flink-yarn/src/main/java/org/apache/flink/yarn/YarnResourceManagerDriver.java b/flink-yarn/src/main/java/org/apache/flink/yarn/YarnResourceManagerDriver.java
new file mode 100644
index 0000000..bb4ec26
--- /dev/null
+++ b/flink-yarn/src/main/java/org/apache/flink/yarn/YarnResourceManagerDriver.java
@@ -0,0 +1,611 @@
+/*
+ * 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.yarn;
+
+import org.apache.flink.annotation.VisibleForTesting;
+import org.apache.flink.configuration.Configuration;
+import org.apache.flink.configuration.GlobalConfiguration;
+import org.apache.flink.configuration.TaskManagerOptions;
+import org.apache.flink.configuration.TaskManagerOptionsInternal;
+import org.apache.flink.runtime.clusterframework.ApplicationStatus;
+import org.apache.flink.runtime.clusterframework.BootstrapTools;
+import org.apache.flink.runtime.clusterframework.ContaineredTaskManagerParameters;
+import org.apache.flink.runtime.clusterframework.TaskExecutorProcessSpec;
+import org.apache.flink.runtime.clusterframework.types.ResourceID;
+import org.apache.flink.runtime.concurrent.FutureUtils;
+import org.apache.flink.runtime.resourcemanager.active.AbstractResourceManagerDriver;
+import org.apache.flink.runtime.resourcemanager.active.ResourceManagerDriver;
+import org.apache.flink.runtime.resourcemanager.exceptions.ResourceManagerException;
+import org.apache.flink.runtime.webmonitor.history.HistoryServerUtils;
+import org.apache.flink.util.ExceptionUtils;
+import org.apache.flink.util.Preconditions;
+import org.apache.flink.yarn.configuration.YarnConfigOptions;
+import org.apache.flink.yarn.configuration.YarnConfigOptionsInternal;
+import org.apache.flink.yarn.configuration.YarnResourceManagerConfiguration;
+
+import org.apache.hadoop.yarn.api.protocolrecords.RegisterApplicationMasterResponse;
+import org.apache.hadoop.yarn.api.records.Container;
+import org.apache.hadoop.yarn.api.records.ContainerId;
+import org.apache.hadoop.yarn.api.records.ContainerLaunchContext;
+import org.apache.hadoop.yarn.api.records.ContainerStatus;
+import org.apache.hadoop.yarn.api.records.FinalApplicationStatus;
+import org.apache.hadoop.yarn.api.records.NodeReport;
+import org.apache.hadoop.yarn.api.records.Priority;
+import org.apache.hadoop.yarn.api.records.Resource;
+import org.apache.hadoop.yarn.api.records.ResourceRequest;
+import org.apache.hadoop.yarn.client.api.AMRMClient;
+import org.apache.hadoop.yarn.client.api.async.AMRMClientAsync;
+import org.apache.hadoop.yarn.client.api.async.NMClientAsync;
+import org.apache.hadoop.yarn.conf.YarnConfiguration;
+import org.apache.hadoop.yarn.exceptions.YarnException;
+
+import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
+
+import java.io.IOException;
+import java.net.URL;
+import java.nio.ByteBuffer;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Iterator;
+import java.util.LinkedList;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.Queue;
+import java.util.Set;
+import java.util.concurrent.CompletableFuture;
+import java.util.stream.Collectors;
+
+/**
+ * Implementation of {@link ResourceManagerDriver} for Yarn deployment.
+ */
+public class YarnResourceManagerDriver extends AbstractResourceManagerDriver<YarnWorkerNode> {
+
+ private static final Priority RM_REQUEST_PRIORITY = Priority.newInstance(1);
+
+ /** Environment variable name of the hostname given by YARN.
+ * In task executor we use the hostnames given by YARN consistently throughout akka */
+ private static final String ENV_FLINK_NODE_ID = "_FLINK_NODE_ID";
+
+ static final String ERROR_MESSAGE_ON_SHUTDOWN_REQUEST = "Received shutdown request from YARN ResourceManager.";
+
+ private final YarnConfiguration yarnConfig;
+
+ /** The process environment variables. */
+ private final YarnResourceManagerConfiguration configuration;
+
+ /** Default heartbeat interval between this resource manager and the YARN ResourceManager. */
+ private final int yarnHeartbeatIntervalMillis;
+
+ /** The heartbeat interval while the resource master is waiting for containers. */
+ private final int containerRequestHeartbeatIntervalMillis;
+
+ /** Request resource futures, keyed by container's TaskExecutorProcessSpec. */
+ private final Map<TaskExecutorProcessSpec, Queue<CompletableFuture<YarnWorkerNode>>> requestResourceFutures;
+
+ private final TaskExecutorProcessSpecContainerResourceAdapter taskExecutorProcessSpecContainerResourceAdapter;
+
+ private final RegisterApplicationMasterResponseReflector registerApplicationMasterResponseReflector;
+
+ private final YarnResourceManagerClientFactory yarnResourceManagerClientFactory;
+
+ private final YarnNodeManagerClientFactory yarnNodeManagerClientFactory;
+
+ /** Client to communicate with the Resource Manager (YARN's master). */
+ private AMRMClientAsync<AMRMClient.ContainerRequest> resourceManagerClient;
+
+ /** Client to communicate with the Node manager and launch TaskExecutor processes. */
+ private NMClientAsync nodeManagerClient;
+
+ private TaskExecutorProcessSpecContainerResourceAdapter.MatchingStrategy matchingStrategy;
+
+ public YarnResourceManagerDriver(
+ Configuration flinkConfig,
+ YarnResourceManagerConfiguration configuration,
+ YarnResourceManagerClientFactory yarnResourceManagerClientFactory,
+ YarnNodeManagerClientFactory yarnNodeManagerClientFactory) {
+ super(flinkConfig, GlobalConfiguration.loadConfiguration(configuration.getCurrentDir()));
+
+ this.yarnConfig = new YarnConfiguration();
+ this.requestResourceFutures = new HashMap<>();
+ this.configuration = configuration;
+
+ final int yarnHeartbeatIntervalMS = flinkConfig.getInteger(
+ YarnConfigOptions.HEARTBEAT_DELAY_SECONDS) * 1000;
+
+ final long yarnExpiryIntervalMS = yarnConfig.getLong(
+ YarnConfiguration.RM_AM_EXPIRY_INTERVAL_MS,
+ YarnConfiguration.DEFAULT_RM_AM_EXPIRY_INTERVAL_MS);
+
+ if (yarnHeartbeatIntervalMS >= yarnExpiryIntervalMS) {
+ log.warn("The heartbeat interval of the Flink Application master ({}) is greater " +
+ "than YARN's expiry interval ({}). The application is likely to be killed by YARN.",
+ yarnHeartbeatIntervalMS, yarnExpiryIntervalMS);
+ }
+ yarnHeartbeatIntervalMillis = yarnHeartbeatIntervalMS;
+ containerRequestHeartbeatIntervalMillis = flinkConfig.getInteger(YarnConfigOptions.CONTAINER_REQUEST_HEARTBEAT_INTERVAL_MILLISECONDS);
+
+ this.taskExecutorProcessSpecContainerResourceAdapter = Utils.createTaskExecutorProcessSpecContainerResourceAdapter(flinkConfig, yarnConfig);
+ this.registerApplicationMasterResponseReflector = new RegisterApplicationMasterResponseReflector(log);
+
+ this.matchingStrategy = flinkConfig.getBoolean(YarnConfigOptionsInternal.MATCH_CONTAINER_VCORES) ?
+ TaskExecutorProcessSpecContainerResourceAdapter.MatchingStrategy.MATCH_VCORE :
+ TaskExecutorProcessSpecContainerResourceAdapter.MatchingStrategy.IGNORE_VCORE;
+
+ this.yarnResourceManagerClientFactory = yarnResourceManagerClientFactory;
+ this.yarnNodeManagerClientFactory = yarnNodeManagerClientFactory;
+ }
+
+ // ------------------------------------------------------------------------
+ // ResourceManagerDriver
+ // ------------------------------------------------------------------------
+
+ @Override
+ protected void initializeInternal() throws Exception {
+ final YarnContainerEventHandler yarnContainerEventHandler = new YarnContainerEventHandler();
+ try {
+ resourceManagerClient = yarnResourceManagerClientFactory.createResourceManagerClient(
+ yarnHeartbeatIntervalMillis,
+ yarnContainerEventHandler);
+ resourceManagerClient.init(yarnConfig);
+ resourceManagerClient.start();
+
+ final RegisterApplicationMasterResponse registerApplicationMasterResponse = registerApplicationMaster();
+ getContainersFromPreviousAttempts(registerApplicationMasterResponse);
+ updateMatchingStrategy(registerApplicationMasterResponse);
+ } catch (Exception e) {
+ throw new ResourceManagerException("Could not start resource manager client.", e);
+ }
+
+ nodeManagerClient = yarnNodeManagerClientFactory.createNodeManagerClient(yarnContainerEventHandler);
+ nodeManagerClient.init(yarnConfig);
+ nodeManagerClient.start();
+ }
+
+ @Override
+ public CompletableFuture<Void> terminate() {
+ // shut down all components
+ Exception exception = null;
+
+ if (resourceManagerClient != null) {
+ try {
+ resourceManagerClient.stop();
+ } catch (Exception e) {
+ exception = e;
+ }
+ }
+
+ if (nodeManagerClient != null) {
+ try {
+ nodeManagerClient.stop();
+ } catch (Exception e) {
+ exception = ExceptionUtils.firstOrSuppressed(e, exception);
+ }
+ }
+
+ return exception == null ?
+ FutureUtils.completedVoidFuture() :
+ FutureUtils.completedExceptionally(exception);
+ }
+
+ @Override
+ public void deregisterApplication(ApplicationStatus finalStatus, @Nullable String optionalDiagnostics) {
+ // first, de-register from YARN
+ final FinalApplicationStatus yarnStatus = getYarnStatus(finalStatus);
+ log.info("Unregister application from the YARN Resource Manager with final status {}.", yarnStatus);
+
+ final Optional<URL> historyServerURL = HistoryServerUtils.getHistoryServerURL(flinkConfig);
+
+ final String appTrackingUrl = historyServerURL.map(URL::toString).orElse("");
+
+ try {
+ resourceManagerClient.unregisterApplicationMaster(yarnStatus, optionalDiagnostics, appTrackingUrl);
+ } catch (YarnException | IOException e) {
+ log.error("Could not unregister the application master.", e);
+ }
+
+ Utils.deleteApplicationFiles(configuration.getYarnFiles());
+ }
+
+ @Override
+ public CompletableFuture<YarnWorkerNode> requestResource(TaskExecutorProcessSpec taskExecutorProcessSpec) {
+ final Optional<Resource> containerResourceOptional = getContainerResource(taskExecutorProcessSpec);
+ final CompletableFuture<YarnWorkerNode> requestResourceFuture = new CompletableFuture<>();
+
+ if (containerResourceOptional.isPresent()) {
+ resourceManagerClient.addContainerRequest(getContainerRequest(containerResourceOptional.get()));
+
+ // make sure we transmit the request fast and receive fast news of granted allocations
+ resourceManagerClient.setHeartbeatInterval(containerRequestHeartbeatIntervalMillis);
+
+ requestResourceFutures.computeIfAbsent(taskExecutorProcessSpec, ignore -> new LinkedList<>()).add(requestResourceFuture);
+
+ log.info("Requesting new TaskExecutor container with resource {}.", taskExecutorProcessSpec);
+ } else {
+ requestResourceFuture.completeExceptionally(
+ new ResourceManagerException(String.format("Could not compute the container Resource from the given TaskExecutorProcessSpec %s.", taskExecutorProcessSpec)));
+ }
+
+ return requestResourceFuture;
+ }
+
+ @Override
+ public void releaseResource(YarnWorkerNode workerNode) {
+ final Container container = workerNode.getContainer();
+ log.info("Stopping container {}.", workerNode.getResourceID().getStringWithMetadata());
+ nodeManagerClient.stopContainerAsync(container.getId(), container.getNodeId());
+ resourceManagerClient.releaseAssignedContainer(container.getId());
+ }
+
+ // ------------------------------------------------------------------------
+ // Internal
+ // ------------------------------------------------------------------------
+
+ private void onContainersOfResourceAllocated(Resource resource, List<Container> containers) {
+ final List<TaskExecutorProcessSpec> pendingTaskExecutorProcessSpecs =
+ taskExecutorProcessSpecContainerResourceAdapter.getTaskExecutorProcessSpec(resource, matchingStrategy).stream()
+ .flatMap(spec -> Collections.nCopies(getNumRequestedNotAllocatedWorkersFor(spec), spec).stream())
+ .collect(Collectors.toList());
+
+ int numPending = pendingTaskExecutorProcessSpecs.size();
+ log.info("Received {} containers with resource {}, {} pending container requests.",
+ containers.size(),
+ resource,
+ numPending);
+
+ final Iterator<Container> containerIterator = containers.iterator();
+ final Iterator<TaskExecutorProcessSpec> pendingTaskExecutorProcessSpecIterator = pendingTaskExecutorProcessSpecs.iterator();
+ final Iterator<AMRMClient.ContainerRequest> pendingRequestsIterator =
+ getPendingRequestsAndCheckConsistency(resource, pendingTaskExecutorProcessSpecs.size()).iterator();
+
+ int numAccepted = 0;
+ while (containerIterator.hasNext() && pendingTaskExecutorProcessSpecIterator.hasNext()) {
+ final TaskExecutorProcessSpec taskExecutorProcessSpec = pendingTaskExecutorProcessSpecIterator.next();
+ final Container container = containerIterator.next();
+ final AMRMClient.ContainerRequest pendingRequest = pendingRequestsIterator.next();
+ final ResourceID resourceId = getContainerResourceId(container);
+ final CompletableFuture<YarnWorkerNode> requestResourceFuture =
+ Preconditions.checkNotNull(
+ Preconditions.checkNotNull(
+ requestResourceFutures.get(taskExecutorProcessSpec),
+ "The requestResourceFuture for TasExecutorProcessSpec %s should not be null.", taskExecutorProcessSpec).poll(),
+ "The requestResourceFuture queue for TasExecutorProcessSpec %s should not be empty.", taskExecutorProcessSpec);
+ if (requestResourceFutures.get(taskExecutorProcessSpec).isEmpty()) {
+ requestResourceFutures.remove(taskExecutorProcessSpec);
+ }
+
+ startTaskExecutorInContainer(container, taskExecutorProcessSpec, resourceId, requestResourceFuture);
+ removeContainerRequest(pendingRequest);
+
+ numAccepted++;
+ }
+ numPending -= numAccepted;
+
+ int numExcess = 0;
+ while (containerIterator.hasNext()) {
+ returnExcessContainer(containerIterator.next());
+ numExcess++;
+ }
+
+ log.info("Accepted {} requested containers, returned {} excess containers, {} pending container requests of resource {}.",
+ numAccepted, numExcess, numPending, resource);
+ }
+
+ private int getNumRequestedNotAllocatedWorkers() {
+ return requestResourceFutures.values().stream().mapToInt(Queue::size).sum();
+ }
+
+ private int getNumRequestedNotAllocatedWorkersFor(TaskExecutorProcessSpec taskExecutorProcessSpec) {
+ return requestResourceFutures.getOrDefault(taskExecutorProcessSpec, new LinkedList<>()).size();
+ }
+
+ private void removeContainerRequest(AMRMClient.ContainerRequest pendingContainerRequest) {
+ log.info("Removing container request {}.", pendingContainerRequest);
+ resourceManagerClient.removeContainerRequest(pendingContainerRequest);
+ }
+
+ private void returnExcessContainer(Container excessContainer) {
+ log.info("Returning excess container {}.", excessContainer.getId());
+ resourceManagerClient.releaseAssignedContainer(excessContainer.getId());
+ }
+
+ private void startTaskExecutorInContainer(Container container, TaskExecutorProcessSpec taskExecutorProcessSpec, ResourceID resourceId, CompletableFuture<YarnWorkerNode> requestResourceFuture) {
+ final YarnWorkerNode yarnWorkerNode = new YarnWorkerNode(container, resourceId);
+
+ try {
+ // Context information used to start a TaskExecutor Java process
+ ContainerLaunchContext taskExecutorLaunchContext = createTaskExecutorLaunchContext(
+ resourceId,
+ container.getNodeId().getHost(),
+ taskExecutorProcessSpec);
+
+ nodeManagerClient.startContainerAsync(container, taskExecutorLaunchContext);
+ requestResourceFuture.complete(yarnWorkerNode);
+ } catch (Throwable t) {
+ requestResourceFuture.completeExceptionally(t);
+ }
+ }
+
+ private Collection<AMRMClient.ContainerRequest> getPendingRequestsAndCheckConsistency(Resource resource, int expectedNum) {
+ final Collection<Resource> equivalentResources = taskExecutorProcessSpecContainerResourceAdapter.getEquivalentContainerResource(resource, matchingStrategy);
+ final List<? extends Collection<AMRMClient.ContainerRequest>> matchingRequests =
+ equivalentResources.stream()
+ .flatMap(equivalentResource -> resourceManagerClient.getMatchingRequests(
+ RM_REQUEST_PRIORITY,
+ ResourceRequest.ANY,
+ equivalentResource).stream())
+ .collect(Collectors.toList());
+
+ final Collection<AMRMClient.ContainerRequest> matchingContainerRequests;
+
+ if (matchingRequests.isEmpty()) {
+ matchingContainerRequests = Collections.emptyList();
+ } else {
+ final Collection<AMRMClient.ContainerRequest> collection = matchingRequests.get(0);
+ matchingContainerRequests = new ArrayList<>(collection);
+ }
+
+ Preconditions.checkState(
+ matchingContainerRequests.size() == expectedNum,
+ "The RMClient's and YarnResourceManagers internal state about the number of pending container requests for resource %s has diverged. " +
+ "Number client's pending container requests %s != Number RM's pending container requests %s.",
+ resource, matchingContainerRequests.size(), expectedNum);
+
+ return matchingContainerRequests;
+ }
+
+ private ContainerLaunchContext createTaskExecutorLaunchContext(
+ ResourceID containerId,
+ String host,
+ TaskExecutorProcessSpec taskExecutorProcessSpec) throws Exception {
+
+ // init the ContainerLaunchContext
+ final String currDir = configuration.getCurrentDir();
+
+ final ContaineredTaskManagerParameters taskManagerParameters =
+ ContaineredTaskManagerParameters.create(flinkConfig, taskExecutorProcessSpec);
+
+ log.info("TaskExecutor {} will be started on {} with {}.",
+ containerId.getStringWithMetadata(),
+ host,
+ taskExecutorProcessSpec);
+
+ final Configuration taskManagerConfig = BootstrapTools.cloneConfiguration(flinkConfig);
+ taskManagerConfig.set(TaskManagerOptions.TASK_MANAGER_RESOURCE_ID, containerId.getResourceIdString());
+ taskManagerConfig.set(TaskManagerOptionsInternal.TASK_MANAGER_RESOURCE_ID_METADATA, containerId.getMetadata());
+
+ final String taskManagerDynamicProperties =
+ BootstrapTools.getDynamicPropertiesAsString(flinkClientConfig, taskManagerConfig);
+
+ log.debug("TaskManager configuration: {}", taskManagerConfig);
+
+ final ContainerLaunchContext taskExecutorLaunchContext = Utils.createTaskExecutorContext(
+ flinkConfig,
+ yarnConfig,
+ configuration,
+ taskManagerParameters,
+ taskManagerDynamicProperties,
+ currDir,
+ YarnTaskExecutorRunner.class,
+ log);
+
+ taskExecutorLaunchContext.getEnvironment()
+ .put(ENV_FLINK_NODE_ID, host);
+ return taskExecutorLaunchContext;
+ }
+
+ @VisibleForTesting
+ Optional<Resource> getContainerResource(TaskExecutorProcessSpec taskExecutorProcessSpec) {
+ return taskExecutorProcessSpecContainerResourceAdapter.tryComputeContainerResource(taskExecutorProcessSpec);
+ }
+
+ private RegisterApplicationMasterResponse registerApplicationMaster() throws Exception {
+ final int restPort;
+ final String webInterfaceUrl = configuration.getWebInterfaceUrl();
+ final String rpcAddress = configuration.getRpcAddress();
+
+ if (webInterfaceUrl != null) {
+ final int lastColon = webInterfaceUrl.lastIndexOf(':');
+
+ if (lastColon == -1) {
+ restPort = -1;
+ } else {
+ restPort = Integer.parseInt(webInterfaceUrl.substring(lastColon + 1));
+ }
+ } else {
+ restPort = -1;
+ }
+
+ return resourceManagerClient.registerApplicationMaster(rpcAddress, restPort, webInterfaceUrl);
+ }
+
+ private void getContainersFromPreviousAttempts(final RegisterApplicationMasterResponse registerApplicationMasterResponse) {
+ final List<Container> containersFromPreviousAttempts =
+ registerApplicationMasterResponseReflector.getContainersFromPreviousAttempts(registerApplicationMasterResponse);
+ final List<YarnWorkerNode> recoveredWorkers = new ArrayList<>();
+
+ log.info("Recovered {} containers from previous attempts ({}).", containersFromPreviousAttempts.size(), containersFromPreviousAttempts);
+
+ for (Container container : containersFromPreviousAttempts) {
+ final YarnWorkerNode worker = new YarnWorkerNode(container, getContainerResourceId(container));
+ recoveredWorkers.add(worker);
+ }
+
+ // Should not invoke resource event handler on the main thread executor.
+ // We are in the initializing thread. The main thread executor is not yet ready.
+ getResourceEventHandler().onPreviousAttemptWorkersRecovered(recoveredWorkers);
+ }
+
+ private void updateMatchingStrategy(final RegisterApplicationMasterResponse registerApplicationMasterResponse) {
+ final Optional<Set<String>> schedulerResourceTypesOptional =
+ registerApplicationMasterResponseReflector.getSchedulerResourceTypeNames(registerApplicationMasterResponse);
+
+ if (schedulerResourceTypesOptional.isPresent()) {
+ Set<String> types = schedulerResourceTypesOptional.get();
+ log.info("Register application master response contains scheduler resource types: {}.", types);
+ matchingStrategy = types.contains("CPU") ?
+ TaskExecutorProcessSpecContainerResourceAdapter.MatchingStrategy.MATCH_VCORE :
+ TaskExecutorProcessSpecContainerResourceAdapter.MatchingStrategy.IGNORE_VCORE;
+ } else {
+ log.info("Register application master response does not contain scheduler resource types, use '{}'.",
+ YarnConfigOptionsInternal.MATCH_CONTAINER_VCORES.key());
+ }
+ log.info("Container matching strategy: {}.", matchingStrategy);
+ }
+
+ // ------------------------------------------------------------------------
+ // Utility methods
+ // ------------------------------------------------------------------------
+
+ /**
+ * Converts a Flink application status enum to a YARN application status enum.
+ * @param status The Flink application status.
+ * @return The corresponding YARN application status.
+ */
+ private FinalApplicationStatus getYarnStatus(ApplicationStatus status) {
+ if (status == null) {
+ return FinalApplicationStatus.UNDEFINED;
+ }
+ else {
+ switch (status) {
+ case SUCCEEDED:
+ return FinalApplicationStatus.SUCCEEDED;
+ case FAILED:
+ return FinalApplicationStatus.FAILED;
+ case CANCELED:
+ return FinalApplicationStatus.KILLED;
+ default:
+ return FinalApplicationStatus.UNDEFINED;
+ }
+ }
+ }
+
+ @Nonnull
+ @VisibleForTesting
+ static AMRMClient.ContainerRequest getContainerRequest(Resource containerResource) {
+ return new AMRMClient.ContainerRequest(
+ containerResource,
+ null,
+ null,
+ RM_REQUEST_PRIORITY);
+ }
+
+ @VisibleForTesting
+ private static ResourceID getContainerResourceId(Container container) {
+ return new ResourceID(container.getId().toString(), container.getNodeId().toString());
+ }
+
+ private Map<Resource, List<Container>> groupContainerByResource(List<Container> containers) {
+ return containers.stream().collect(Collectors.groupingBy(Container::getResource));
+ }
+
+ // ------------------------------------------------------------------------
+ // Event handlers
+ // ------------------------------------------------------------------------
+
+ class YarnContainerEventHandler implements AMRMClientAsync.CallbackHandler, NMClientAsync.CallbackHandler {
+
+ @Override
+ public void onContainersCompleted(List<ContainerStatus> statuses) {
+ getMainThreadExecutor().execute(() -> {
+ log.debug("YARN ResourceManager reported the following containers completed: {}.", statuses);
+ for (final ContainerStatus containerStatus : statuses) {
+
+ final String containerId = containerStatus.getContainerId().toString();
+ getResourceEventHandler().onWorkerTerminated(new ResourceID(containerId), containerStatus.getDiagnostics());
+ }
+ }
+ );
+ }
+
+ @Override
+ public void onContainersAllocated(List<Container> containers) {
+ getMainThreadExecutor().execute(() -> {
+ log.info("Received {} containers.", containers.size());
+
+ for (Map.Entry<Resource, List<Container>> entry : groupContainerByResource(containers).entrySet()) {
+ onContainersOfResourceAllocated(entry.getKey(), entry.getValue());
+ }
+
+ // if we are waiting for no further containers, we can go to the
+ // regular heartbeat interval
+ if (getNumRequestedNotAllocatedWorkers() <= 0) {
+ resourceManagerClient.setHeartbeatInterval(yarnHeartbeatIntervalMillis);
+ }
+ });
+ }
+
+ @Override
+ public void onShutdownRequest() {
+ getResourceEventHandler().onError(new ResourceManagerException(ERROR_MESSAGE_ON_SHUTDOWN_REQUEST));
+ }
+
+ @Override
+ public void onNodesUpdated(List<NodeReport> list) {
+ // We are not interested in node updates
+ }
+
+ @Override
+ public float getProgress() {
+ // Temporarily need not record the total size of asked and allocated containers
+ return 1;
+ }
+
+ @Override
+ public void onError(Throwable throwable) {
+ getResourceEventHandler().onError(throwable);
+ }
+
+ @Override
+ public void onContainerStarted(ContainerId containerId, Map<String, ByteBuffer> map) {
+ log.debug("Succeeded to call YARN Node Manager to start container {}.", containerId);
+ }
+
+ @Override
+ public void onContainerStatusReceived(ContainerId containerId, ContainerStatus containerStatus) {
+ // We are not interested in getting container status
+ }
+
+ @Override
+ public void onContainerStopped(ContainerId containerId) {
+ log.debug("Succeeded to call YARN Node Manager to stop container {}.", containerId);
+ }
+
+ @Override
+ public void onStartContainerError(ContainerId containerId, Throwable throwable) {
+ getMainThreadExecutor().execute(() -> {
+ resourceManagerClient.releaseAssignedContainer(containerId);
+ getResourceEventHandler().onWorkerTerminated(new ResourceID(containerId.toString()), throwable.getMessage());
+ });
+ }
+
+ @Override
+ public void onGetContainerStatusError(ContainerId containerId, Throwable throwable) {
+ // We are not interested in getting container status
+ }
+
+ @Override
+ public void onStopContainerError(ContainerId containerId, Throwable throwable) {
+ log.warn("Error while calling YARN Node Manager to stop container {}.", containerId, throwable);
+ }
+ }
+}
diff --git a/flink-yarn/src/test/java/org/apache/flink/yarn/TestingYarnAMRMClientAsync.java b/flink-yarn/src/test/java/org/apache/flink/yarn/TestingYarnAMRMClientAsync.java
index d16dca0..1256844 100644
--- a/flink-yarn/src/test/java/org/apache/flink/yarn/TestingYarnAMRMClientAsync.java
+++ b/flink-yarn/src/test/java/org/apache/flink/yarn/TestingYarnAMRMClientAsync.java
@@ -52,6 +52,9 @@ public class TestingYarnAMRMClientAsync extends AMRMClientAsyncImpl<AMRMClient.C
private final Consumer<Integer> setHeartbeatIntervalConsumer;
private final TriFunction<String, Integer, String, RegisterApplicationMasterResponse> registerApplicationMasterFunction;
private final TriConsumer<FinalApplicationStatus, String, String> unregisterApplicationMasterConsumer;
+ private final Runnable clientInitRunnable;
+ private final Runnable clientStartRunnable;
+ private final Runnable clientStopRunnable;
private TestingYarnAMRMClientAsync(
CallbackHandler callbackHandler,
@@ -61,7 +64,10 @@ public class TestingYarnAMRMClientAsync extends AMRMClientAsyncImpl<AMRMClient.C
BiConsumer<ContainerId, CallbackHandler> releaseAssignedContainerConsumer,
Consumer<Integer> setHeartbeatIntervalConsumer,
TriFunction<String, Integer, String, RegisterApplicationMasterResponse> registerApplicationMasterFunction,
- TriConsumer<FinalApplicationStatus, String, String> unregisterApplicationMasterConsumer) {
+ TriConsumer<FinalApplicationStatus, String, String> unregisterApplicationMasterConsumer,
+ Runnable clientInitRunnable,
+ Runnable clientStartRunnable,
+ Runnable clientStopRunnable) {
super(0, callbackHandler);
this.setHeartbeatIntervalConsumer = Preconditions.checkNotNull(setHeartbeatIntervalConsumer);
this.addContainerRequestConsumer = Preconditions.checkNotNull(addContainerRequestConsumer);
@@ -70,6 +76,9 @@ public class TestingYarnAMRMClientAsync extends AMRMClientAsyncImpl<AMRMClient.C
this.registerApplicationMasterFunction = Preconditions.checkNotNull(registerApplicationMasterFunction);
this.getMatchingRequestsFunction = Preconditions.checkNotNull(getMatchingRequestsFunction);
this.unregisterApplicationMasterConsumer = Preconditions.checkNotNull(unregisterApplicationMasterConsumer);
+ this.clientInitRunnable = Preconditions.checkNotNull(clientInitRunnable);
+ this.clientStartRunnable = Preconditions.checkNotNull(clientStartRunnable);
+ this.clientStopRunnable = Preconditions.checkNotNull(clientStopRunnable);
}
@Override
@@ -116,18 +125,18 @@ public class TestingYarnAMRMClientAsync extends AMRMClientAsyncImpl<AMRMClient.C
// ------------------------------------------------------------------------
@Override
- protected void serviceInit(Configuration conf) throws Exception {
- // noop
+ public void init(Configuration conf) {
+ clientInitRunnable.run();
}
@Override
- protected void serviceStart() throws Exception {
- // noop
+ public void start() {
+ clientStartRunnable.run();
}
@Override
- protected void serviceStop() throws Exception {
- // noop
+ public void stop() {
+ clientStopRunnable.run();
}
/**
@@ -143,6 +152,9 @@ public class TestingYarnAMRMClientAsync extends AMRMClientAsyncImpl<AMRMClient.C
private TriFunction<String, Integer, String, RegisterApplicationMasterResponse> registerApplicationMasterFunction =
(ignored1, ignored2, ignored3) -> new TestingRegisterApplicationMasterResponse(Collections::emptyList);
private TriConsumer<FinalApplicationStatus, String, String> unregisterApplicationMasterConsumer = (ignored1, ignored2, ignored3) -> {};
+ private Runnable clientInitRunnable = () -> {};
+ private Runnable clientStartRunnable = () -> {};
+ private Runnable clientStopRunnable = () -> {};
private Builder() {}
@@ -181,6 +193,21 @@ public class TestingYarnAMRMClientAsync extends AMRMClientAsyncImpl<AMRMClient.C
return this;
}
+ Builder setClientInitRunnable(Runnable clientInitRunnable) {
+ this.clientInitRunnable = clientInitRunnable;
+ return this;
+ }
+
+ Builder setClientStartRunnable(Runnable clientStartRunnable) {
+ this.clientStartRunnable = clientStartRunnable;
+ return this;
+ }
+
+ Builder setClientStopRunnable(Runnable clientStopRunnable) {
+ this.clientStopRunnable = clientStopRunnable;
+ return this;
+ }
+
public TestingYarnAMRMClientAsync build(CallbackHandler callbackHandler) {
return new TestingYarnAMRMClientAsync(
callbackHandler,
@@ -190,7 +217,10 @@ public class TestingYarnAMRMClientAsync extends AMRMClientAsyncImpl<AMRMClient.C
releaseAssignedContainerConsumer,
setHeartbeatIntervalConsumer,
registerApplicationMasterFunction,
- unregisterApplicationMasterConsumer);
+ unregisterApplicationMasterConsumer,
+ clientInitRunnable,
+ clientStartRunnable,
+ clientStopRunnable);
}
}
}
diff --git a/flink-yarn/src/test/java/org/apache/flink/yarn/TestingYarnNMClientAsync.java b/flink-yarn/src/test/java/org/apache/flink/yarn/TestingYarnNMClientAsync.java
index 5765515..de2c8a4 100644
--- a/flink-yarn/src/test/java/org/apache/flink/yarn/TestingYarnNMClientAsync.java
+++ b/flink-yarn/src/test/java/org/apache/flink/yarn/TestingYarnNMClientAsync.java
@@ -36,14 +36,23 @@ class TestingYarnNMClientAsync extends NMClientAsyncImpl {
private final TriConsumer<Container, ContainerLaunchContext, CallbackHandler> startContainerAsyncConsumer;
private final TriConsumer<ContainerId, NodeId, CallbackHandler> stopContainerAsyncConsumer;
+ private final Runnable clientInitRunnable;
+ private final Runnable clientStartRunnable;
+ private final Runnable clientStopRunnable;
private TestingYarnNMClientAsync(
final CallbackHandler callbackHandler,
TriConsumer<Container, ContainerLaunchContext, CallbackHandler> startContainerAsyncConsumer,
- TriConsumer<ContainerId, NodeId, CallbackHandler> stopContainerAsyncConsumer) {
+ TriConsumer<ContainerId, NodeId, CallbackHandler> stopContainerAsyncConsumer,
+ Runnable clientInitRunnable,
+ Runnable clientStartRunnable,
+ Runnable clientStopRunnable) {
super(callbackHandler);
this.startContainerAsyncConsumer = Preconditions.checkNotNull(startContainerAsyncConsumer);
this.stopContainerAsyncConsumer = Preconditions.checkNotNull(stopContainerAsyncConsumer);
+ this.clientInitRunnable = Preconditions.checkNotNull(clientInitRunnable);
+ this.clientStartRunnable = Preconditions.checkNotNull(clientStartRunnable);
+ this.clientStopRunnable = Preconditions.checkNotNull(clientStopRunnable);
}
@Override
@@ -65,18 +74,18 @@ class TestingYarnNMClientAsync extends NMClientAsyncImpl {
// ------------------------------------------------------------------------
@Override
- protected void serviceInit(Configuration conf) throws Exception {
- // noop
+ public void init(Configuration conf) {
+ clientInitRunnable.run();
}
@Override
- protected void serviceStart() throws Exception {
- // noop
+ public void start() {
+ clientStartRunnable.run();
}
@Override
- protected void serviceStop() throws Exception {
- // noop
+ public void stop() {
+ clientStopRunnable.run();
}
/**
@@ -85,6 +94,9 @@ class TestingYarnNMClientAsync extends NMClientAsyncImpl {
public static class Builder {
private TriConsumer<Container, ContainerLaunchContext, CallbackHandler> startContainerAsyncConsumer = (ignored1, ignored2, ignored3) -> {};
private TriConsumer<ContainerId, NodeId, CallbackHandler> stopContainerAsyncConsumer = (ignored1, ignored2, ignored3) -> {};
+ private Runnable clientInitRunnable = () -> {};
+ private Runnable clientStartRunnable = () -> {};
+ private Runnable clientStopRunnable = () -> {};
private Builder() {}
@@ -98,11 +110,29 @@ class TestingYarnNMClientAsync extends NMClientAsyncImpl {
return this;
}
+ Builder setClientInitRunnable(Runnable clientInitRunnable) {
+ this.clientInitRunnable = clientInitRunnable;
+ return this;
+ }
+
+ Builder setClientStartRunnable(Runnable clientStartRunnable) {
+ this.clientStartRunnable = clientStartRunnable;
+ return this;
+ }
+
+ Builder setClientStopRunnable(Runnable clientStopRunnable) {
+ this.clientStopRunnable = clientStopRunnable;
+ return this;
+ }
+
public TestingYarnNMClientAsync build(CallbackHandler callbackHandler) {
return new TestingYarnNMClientAsync(
callbackHandler,
startContainerAsyncConsumer,
- stopContainerAsyncConsumer);
+ stopContainerAsyncConsumer,
+ clientInitRunnable,
+ clientStartRunnable,
+ clientStopRunnable);
}
}
}
diff --git a/flink-yarn/src/test/java/org/apache/flink/yarn/YarnResourceManagerDriverTest.java b/flink-yarn/src/test/java/org/apache/flink/yarn/YarnResourceManagerDriverTest.java
new file mode 100644
index 0000000..1bfc3a5
--- /dev/null
+++ b/flink-yarn/src/test/java/org/apache/flink/yarn/YarnResourceManagerDriverTest.java
@@ -0,0 +1,559 @@
+/*
+ * 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.yarn;
+
+import org.apache.flink.api.common.resources.CPUResource;
+import org.apache.flink.configuration.MemorySize;
+import org.apache.flink.configuration.TaskManagerOptions;
+import org.apache.flink.runtime.clusterframework.ApplicationStatus;
+import org.apache.flink.runtime.clusterframework.BootstrapTools;
+import org.apache.flink.runtime.clusterframework.TaskExecutorProcessSpec;
+import org.apache.flink.runtime.clusterframework.types.ResourceID;
+import org.apache.flink.runtime.resourcemanager.active.ResourceManagerDriver;
+import org.apache.flink.runtime.resourcemanager.active.ResourceManagerDriverTestBase;
+import org.apache.flink.runtime.resourcemanager.exceptions.ResourceManagerException;
+import org.apache.flink.runtime.util.HadoopUtils;
+import org.apache.flink.util.ExceptionUtils;
+import org.apache.flink.yarn.configuration.YarnResourceManagerConfiguration;
+
+import org.apache.flink.shaded.guava18.com.google.common.collect.ImmutableList;
+
+import org.apache.hadoop.fs.Path;
+import org.apache.hadoop.yarn.api.ApplicationConstants;
+import org.apache.hadoop.yarn.api.records.ApplicationAttemptId;
+import org.apache.hadoop.yarn.api.records.ApplicationId;
+import org.apache.hadoop.yarn.api.records.Container;
+import org.apache.hadoop.yarn.api.records.ContainerId;
+import org.apache.hadoop.yarn.api.records.ContainerLaunchContext;
+import org.apache.hadoop.yarn.api.records.ContainerState;
+import org.apache.hadoop.yarn.api.records.ContainerStatus;
+import org.apache.hadoop.yarn.api.records.LocalResourceType;
+import org.apache.hadoop.yarn.api.records.LocalResourceVisibility;
+import org.apache.hadoop.yarn.api.records.NodeId;
+import org.apache.hadoop.yarn.api.records.Priority;
+import org.apache.hadoop.yarn.api.records.Resource;
+import org.apache.hadoop.yarn.client.api.async.AMRMClientAsync;
+import org.apache.hadoop.yarn.client.api.async.NMClientAsync;
+import org.apache.hadoop.yarn.conf.YarnConfiguration;
+import org.junit.Assume;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.TemporaryFolder;
+
+import java.io.File;
+import java.nio.file.Files;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+
+import static org.apache.flink.configuration.GlobalConfiguration.FLINK_CONF_FILENAME;
+import static org.apache.flink.yarn.YarnConfigKeys.ENV_APP_ID;
+import static org.apache.flink.yarn.YarnConfigKeys.ENV_CLIENT_HOME_DIR;
+import static org.apache.flink.yarn.YarnConfigKeys.ENV_CLIENT_SHIP_FILES;
+import static org.apache.flink.yarn.YarnConfigKeys.ENV_FLINK_CLASSPATH;
+import static org.apache.flink.yarn.YarnConfigKeys.ENV_HADOOP_USER_NAME;
+import static org.apache.flink.yarn.YarnConfigKeys.FLINK_DIST_JAR;
+import static org.apache.flink.yarn.YarnConfigKeys.FLINK_YARN_FILES;
+import static org.apache.flink.yarn.YarnResourceManagerDriver.ERROR_MESSAGE_ON_SHUTDOWN_REQUEST;
+import static org.hamcrest.Matchers.is;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNotEquals;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertThat;
+import static org.junit.Assert.assertTrue;
+
+/**
+ * Tests for {@link YarnResourceManagerDriver}.
+ */
+public class YarnResourceManagerDriverTest extends ResourceManagerDriverTestBase<YarnWorkerNode> {
+ private static final Resource testingResource = Resource.newInstance(YarnConfiguration.DEFAULT_RM_SCHEDULER_MINIMUM_ALLOCATION_MB, YarnConfiguration.DEFAULT_RM_SCHEDULER_MINIMUM_ALLOCATION_VCORES);
+ private static final Container testingContainer = createTestingContainerWithResource(testingResource, 0);
+ private static final TaskExecutorProcessSpec testingTaskExecutorProcessSpec =
+ new TaskExecutorProcessSpec(
+ new CPUResource(1),
+ MemorySize.ZERO,
+ MemorySize.ZERO,
+ MemorySize.ofMebiBytes(256),
+ MemorySize.ofMebiBytes(256),
+ MemorySize.ofMebiBytes(256),
+ MemorySize.ofMebiBytes(256),
+ MemorySize.ZERO,
+ MemorySize.ZERO);
+
+ @Rule
+ public TemporaryFolder folder = new TemporaryFolder();
+
+ @Override
+ protected Context createContext() {
+ return new Context();
+ }
+
+ @Test
+ public void testShutdownRequestCausesFatalError() throws Exception {
+ new Context() {{
+ final CompletableFuture<Throwable> throwableCompletableFuture = new CompletableFuture<>();
+ resourceEventHandlerBuilder.setOnErrorConsumer(throwableCompletableFuture::complete);
+ runTest(() -> {
+ resourceManagerClientCallbackHandler.onShutdownRequest();
+
+ Throwable throwable = throwableCompletableFuture.get(TIMEOUT_SEC, TimeUnit.SECONDS);
+ assertThat(ExceptionUtils.findThrowable(throwable, ResourceManagerException.class).isPresent(), is(true));
+ assertThat(ExceptionUtils.findThrowableWithMessage(throwable, ERROR_MESSAGE_ON_SHUTDOWN_REQUEST).isPresent(), is(true));
+ });
+ }};
+ }
+
+ /**
+ * Tests that application files are deleted when the YARN application master is de-registered.
+ */
+ @Test
+ public void testDeleteApplicationFiles() throws Exception {
+ new Context() {{
+ final File applicationDir = folder.newFolder(".flink");
+ env.put(FLINK_YARN_FILES, applicationDir.getCanonicalPath());
+
+ runTest(() -> {
+ getDriver().deregisterApplication(ApplicationStatus.SUCCEEDED, null);
+ assertFalse("YARN application directory was not removed", Files.exists(applicationDir.toPath()));
+ });
+ }};
+ }
+
+ @Test
+ public void testOnContainerAllocated() throws Exception {
+ new Context() {{
+ addContainerRequestFutures.add(new CompletableFuture<>());
+
+ testingYarnAMRMClientAsyncBuilder.setAddContainerRequestConsumer((ignored1, ignored2) ->
+ addContainerRequestFutures.get(addContainerRequestFuturesNumCompleted.getAndIncrement()).complete(null));
+
+ runTest(() -> {
+ runInMainThread(() -> getDriver().requestResource(testingTaskExecutorProcessSpec));
+ resourceManagerClientCallbackHandler.onContainersAllocated(ImmutableList.of(testingContainer));
+
+ verifyFutureCompleted(addContainerRequestFutures.get(0));
+ verifyFutureCompleted(removeContainerRequestFuture);
+ verifyFutureCompleted(startContainerAsyncFuture);
+ });
+ }};
+ }
+
+ @Test
+ public void testOnContainerCompleted() throws Exception {
+ new Context() {{
+ addContainerRequestFutures.add(new CompletableFuture<>());
+ addContainerRequestFutures.add(new CompletableFuture<>());
+
+ testingYarnAMRMClientAsyncBuilder.setAddContainerRequestConsumer((ignored1, ignored2) ->
+ addContainerRequestFutures.get(addContainerRequestFuturesNumCompleted.getAndIncrement()).complete(null));
+ resourceEventHandlerBuilder.setOnWorkerTerminatedConsumer((ignore1, ignore2) -> getDriver().requestResource(testingTaskExecutorProcessSpec));
+
+ runTest(() -> {
+ runInMainThread(() -> getDriver().requestResource(testingTaskExecutorProcessSpec));
+ resourceManagerClientCallbackHandler.onContainersAllocated(ImmutableList.of(testingContainer));
+ ContainerStatus testingContainerStatus = createTestingContainerCompletedStatus(testingContainer.getId());
+ resourceManagerClientCallbackHandler.onContainersCompleted(ImmutableList.of(testingContainerStatus));
+
+ verifyFutureCompleted(addContainerRequestFutures.get(1));
+ });
+ }};
+ }
+
+ @Test
+ public void testOnStartContainerError() throws Exception {
+ new Context() {{
+ addContainerRequestFutures.add(new CompletableFuture<>());
+ addContainerRequestFutures.add(new CompletableFuture<>());
+
+ testingYarnAMRMClientAsyncBuilder.setAddContainerRequestConsumer((ignored1, ignored2) ->
+ addContainerRequestFutures.get(addContainerRequestFuturesNumCompleted.getAndIncrement()).complete(null));
+ resourceEventHandlerBuilder.setOnWorkerTerminatedConsumer((ignore1, ignore2) -> getDriver().requestResource(testingTaskExecutorProcessSpec));
+
+ runTest(() -> {
+ runInMainThread(() -> getDriver().requestResource(testingTaskExecutorProcessSpec));
+ resourceManagerClientCallbackHandler.onContainersAllocated(ImmutableList.of(testingContainer));
+ nodeManagerClientCallbackHandler.onStartContainerError(testingContainer.getId(), new Exception("start error"));
+
+ verifyFutureCompleted(releaseAssignedContainerFuture);
+ verifyFutureCompleted(addContainerRequestFutures.get(1));
+ });
+ }};
+ }
+
+ @Test
+ public void testStartTaskExecutorProcessVariousSpec_SameContainerResource() throws Exception{
+ final TaskExecutorProcessSpec taskExecutorProcessSpec1 =
+ new TaskExecutorProcessSpec(
+ new CPUResource(1),
+ MemorySize.ZERO,
+ MemorySize.ZERO,
+ MemorySize.ofMebiBytes(100),
+ MemorySize.ZERO,
+ MemorySize.ofMebiBytes(100),
+ MemorySize.ofMebiBytes(100),
+ MemorySize.ZERO,
+ MemorySize.ZERO
+ );
+ final TaskExecutorProcessSpec taskExecutorProcessSpec2 =
+ new TaskExecutorProcessSpec(
+ new CPUResource(1),
+ MemorySize.ZERO,
+ MemorySize.ZERO,
+ MemorySize.ofMebiBytes(99),
+ MemorySize.ZERO,
+ MemorySize.ofMebiBytes(100),
+ MemorySize.ofMebiBytes(100),
+ MemorySize.ZERO,
+ MemorySize.ZERO
+ );
+
+ new Context() {{
+ final String startCommand1 = TaskManagerOptions.TASK_HEAP_MEMORY.key() + "=" + (100L << 20);
+ final String startCommand2 = TaskManagerOptions.TASK_HEAP_MEMORY.key() + "=" + (99L << 20);
+ final CompletableFuture<Void> startContainerAsyncCommandFuture1 = new CompletableFuture<>();
+ final CompletableFuture<Void> startContainerAsyncCommandFuture2 = new CompletableFuture<>();
+ prepareForTestStartTaskExecutorProcessVariousSpec(
+ startCommand1,
+ startCommand2,
+ startContainerAsyncCommandFuture1,
+ startContainerAsyncCommandFuture2,
+ taskExecutorProcessSpec1);
+
+ testingYarnAMRMClientAsyncBuilder.setGetMatchingRequestsFunction(ignored ->
+ Collections.singletonList(ImmutableList.of(
+ YarnResourceManagerDriver.getContainerRequest(((YarnResourceManagerDriver) getDriver()).getContainerResource(taskExecutorProcessSpec1).get()),
+ YarnResourceManagerDriver.getContainerRequest(((YarnResourceManagerDriver) getDriver()).getContainerResource(taskExecutorProcessSpec2).get()))));
+
+ runTest(() -> {
+ final Resource containerResource = ((YarnResourceManagerDriver) getDriver()).getContainerResource(taskExecutorProcessSpec1).get();
+ // Make sure two worker resource spec will be normalized to the same container resource
+ assertEquals(containerResource, ((YarnResourceManagerDriver) getDriver()).getContainerResource(taskExecutorProcessSpec2).get());
+
+ runInMainThread(() -> getDriver().requestResource(taskExecutorProcessSpec1));
+ runInMainThread(() -> getDriver().requestResource(taskExecutorProcessSpec2));
+
+ // Verify both containers requested
+ verifyFutureCompleted(addContainerRequestFutures.get(0));
+ verifyFutureCompleted(addContainerRequestFutures.get(1));
+
+ // Mock that both containers are allocated
+ Container container1 = createTestingContainerWithResource(containerResource);
+ Container container2 = createTestingContainerWithResource(containerResource);
+ resourceManagerClientCallbackHandler.onContainersAllocated(ImmutableList.of(container1, container2));
+
+ // Verify workers with both spec are started.
+ verifyFutureCompleted(startContainerAsyncCommandFuture1);
+ verifyFutureCompleted(startContainerAsyncCommandFuture2);
+
+ // Mock that one container is completed, while the worker is still pending
+ ContainerStatus testingContainerStatus = createTestingContainerCompletedStatus(container1.getId());
+ resourceManagerClientCallbackHandler.onContainersCompleted(Collections.singletonList(testingContainerStatus));
+
+ // Verify that only one more container is requested.
+ verifyFutureCompleted(addContainerRequestFutures.get(2));
+ assertFalse(addContainerRequestFutures.get(3).isDone());
+ });
+ }};
+ }
+
+ @Test
+ public void testStartWorkerVariousSpec_DifferentContainerResource() throws Exception{
+ final TaskExecutorProcessSpec taskExecutorProcessSpec1 =
+ new TaskExecutorProcessSpec(
+ new CPUResource(1),
+ MemorySize.ZERO,
+ MemorySize.ZERO,
+ MemorySize.ofMebiBytes(50),
+ MemorySize.ofMebiBytes(50),
+ MemorySize.ofMebiBytes(50),
+ MemorySize.ofMebiBytes(50),
+ MemorySize.ZERO,
+ MemorySize.ZERO
+ );
+ final TaskExecutorProcessSpec taskExecutorProcessSpec2 =
+ new TaskExecutorProcessSpec(
+ new CPUResource(2),
+ MemorySize.ZERO,
+ MemorySize.ZERO,
+ MemorySize.ofMebiBytes(500),
+ MemorySize.ofMebiBytes(500),
+ MemorySize.ofMebiBytes(500),
+ MemorySize.ofMebiBytes(500),
+ MemorySize.ZERO,
+ MemorySize.ZERO
+ );
+
+ new Context() {{
+ final String startCommand1 = TaskManagerOptions.TASK_HEAP_MEMORY.key() + "=" + (50L << 20);
+ final String startCommand2 = TaskManagerOptions.TASK_HEAP_MEMORY.key() + "=" + (100L << 20);
+ final CompletableFuture<Void> startContainerAsyncCommandFuture1 = new CompletableFuture<>();
+ final CompletableFuture<Void> startContainerAsyncCommandFuture2 = new CompletableFuture<>();
+ prepareForTestStartTaskExecutorProcessVariousSpec(
+ startCommand1,
+ startCommand2,
+ startContainerAsyncCommandFuture1,
+ startContainerAsyncCommandFuture2,
+ taskExecutorProcessSpec1);
+
+ testingYarnAMRMClientAsyncBuilder.setGetMatchingRequestsFunction(tuple -> {
+ if (tuple.f2.getVirtualCores() == 1) {
+ return Collections.singletonList(
+ Collections.singletonList(YarnResourceManagerDriver.getContainerRequest(((YarnResourceManagerDriver) getDriver()).getContainerResource(taskExecutorProcessSpec1).get())));
+ } else if (tuple.f2.getVirtualCores() == 2) {
+ return Collections.singletonList(
+ Collections.singletonList(YarnResourceManagerDriver.getContainerRequest(((YarnResourceManagerDriver) getDriver()).getContainerResource(taskExecutorProcessSpec2).get())));
+ }
+ return null;
+ });
+
+ runTest(() -> {
+ final Resource containerResource1 = ((YarnResourceManagerDriver) getDriver()).getContainerResource(taskExecutorProcessSpec1).get();
+ final Resource containerResource2 = ((YarnResourceManagerDriver) getDriver()).getContainerResource(taskExecutorProcessSpec2).get();
+ // Make sure two worker resource spec will be normalized to different container resources
+ assertNotEquals(containerResource1, containerResource2);
+
+ runInMainThread(() -> getDriver().requestResource(taskExecutorProcessSpec1));
+ runInMainThread(() -> getDriver().requestResource(taskExecutorProcessSpec2));
+
+ // Verify both containers requested
+ verifyFutureCompleted(addContainerRequestFutures.get(0));
+ verifyFutureCompleted(addContainerRequestFutures.get(1));
+
+ // Mock that container 1 is allocated
+ Container container1 = createTestingContainerWithResource(containerResource1);
+ resourceManagerClientCallbackHandler.onContainersAllocated(Collections.singletonList(container1));
+
+ // Verify that only worker with spec1 is started.
+ verifyFutureCompleted(startContainerAsyncCommandFuture1);
+ assertFalse(startContainerAsyncCommandFuture2.isDone());
+
+ // Mock that container 1 is completed, while the worker is still pending
+ ContainerStatus testingContainerStatus = createTestingContainerCompletedStatus(container1.getId());
+ resourceManagerClientCallbackHandler.onContainersCompleted(Collections.singletonList(testingContainerStatus));
+
+ // Verify that only container 1 is requested again
+ verifyFutureCompleted(addContainerRequestFutures.get(2));
+ assertThat(addContainerRequestFutures.get(2).get(), is(containerResource1));
+ assertFalse(addContainerRequestFutures.get(3).isDone());
+ });
+ }};
+ }
+
+ private boolean containsStartCommand(ContainerLaunchContext containerLaunchContext, String command) {
+ return containerLaunchContext.getCommands().stream().anyMatch(str -> str.contains(command));
+ }
+
+ private static Container createTestingContainerWithResource(Resource resource, int containerIdx) {
+ final ContainerId containerId = ContainerId.newInstance(
+ ApplicationAttemptId.newInstance(
+ ApplicationId.newInstance(System.currentTimeMillis(), 1),
+ 1),
+ containerIdx);
+ final NodeId nodeId = NodeId.newInstance("container", 1234);
+ return new TestingContainer(containerId, nodeId, resource, Priority.UNDEFINED);
+ }
+
+ private class Context extends ResourceManagerDriverTestBase<YarnWorkerNode>.Context {
+ private final CompletableFuture<Void> stopAndCleanupClusterFuture = new CompletableFuture<>();
+ private final CompletableFuture<Resource> createTaskManagerContainerFuture = new CompletableFuture<>();
+ private final CompletableFuture<Void> stopContainerAsyncFuture = new CompletableFuture<>();
+ final List<CompletableFuture<Resource>> addContainerRequestFutures = new ArrayList<>();
+ final AtomicInteger addContainerRequestFuturesNumCompleted = new AtomicInteger(0);
+ final CompletableFuture<Void> removeContainerRequestFuture = new CompletableFuture<>();
+ final CompletableFuture<Void> releaseAssignedContainerFuture = new CompletableFuture<>();
+ final CompletableFuture<Void> startContainerAsyncFuture = new CompletableFuture<>();
+ final CompletableFuture<Void> resourceManagerClientInitFuture = new CompletableFuture<>();
+ final CompletableFuture<Void> resourceManagerClientStartFuture = new CompletableFuture<>();
+ final CompletableFuture<Void> resourceManagerClientStopFuture = new CompletableFuture<>();
+ final CompletableFuture<Void> nodeManagerClientInitFuture = new CompletableFuture<>();
+ final CompletableFuture<Void> nodeManagerClientStartFuture = new CompletableFuture<>();
+ final CompletableFuture<Void> nodeManagerClientStopFuture = new CompletableFuture<>();
+
+ AMRMClientAsync.CallbackHandler resourceManagerClientCallbackHandler;
+ NMClientAsync.CallbackHandler nodeManagerClientCallbackHandler;
+ TestingYarnNMClientAsync testingYarnNMClientAsync;
+ TestingYarnAMRMClientAsync testingYarnAMRMClientAsync;
+ final TestingYarnNMClientAsync.Builder testingYarnNMClientAsyncBuilder = TestingYarnNMClientAsync.builder()
+ .setStartContainerAsyncConsumer((ignored1, ignored2, ignored3) -> startContainerAsyncFuture.complete(null))
+ .setStopContainerAsyncConsumer((ignored1, ignored2, ignored3) -> stopContainerAsyncFuture.complete(null))
+ .setClientInitRunnable(() -> nodeManagerClientInitFuture.complete(null))
+ .setClientStartRunnable(() -> nodeManagerClientStartFuture.complete(null))
+ .setClientStopRunnable(() -> nodeManagerClientStopFuture.complete(null));
+ final TestingYarnAMRMClientAsync.Builder testingYarnAMRMClientAsyncBuilder = TestingYarnAMRMClientAsync.builder()
+ .setAddContainerRequestConsumer((request, handler) -> {
+ createTaskManagerContainerFuture.complete(request.getCapability());
+ resourceManagerClientCallbackHandler
+ .onContainersAllocated(Collections.singletonList(testingContainer));
+ })
+ .setGetMatchingRequestsFunction(ignored ->
+ Collections.singletonList(Collections.singletonList(YarnResourceManagerDriver.getContainerRequest(testingResource))))
+ .setRemoveContainerRequestConsumer((request, handler) -> removeContainerRequestFuture.complete(null))
+ .setReleaseAssignedContainerConsumer((ignored1, ignored2) -> releaseAssignedContainerFuture.complete(null))
+ .setUnregisterApplicationMasterConsumer((ignore1, ignore2, ignore3) -> stopAndCleanupClusterFuture.complete(null))
+ .setClientInitRunnable(() -> resourceManagerClientInitFuture.complete(null))
+ .setClientStartRunnable(() -> resourceManagerClientStartFuture.complete(null))
+ .setClientStopRunnable(() -> resourceManagerClientStopFuture.complete(null));
+ final TestingYarnResourceManagerClientFactory testingYarnResourceManagerClientFactory =
+ new TestingYarnResourceManagerClientFactory(((integer, handler) -> {
+ resourceManagerClientCallbackHandler = handler;
+ testingYarnAMRMClientAsync = testingYarnAMRMClientAsyncBuilder.build(handler);
+ return testingYarnAMRMClientAsync;
+ }));
+ final TestingYarnNodeManagerClientFactory testingYarnNodeManagerClientFactory =
+ new TestingYarnNodeManagerClientFactory((handler -> {
+ nodeManagerClientCallbackHandler = handler;
+ testingYarnNMClientAsync = testingYarnNMClientAsyncBuilder.build(handler);
+ return testingYarnNMClientAsync;
+ }));
+
+ final Map<String, String> env = new HashMap<>();
+
+ private int containerIdx = 0;
+
+ @Override
+ protected void prepareRunTest() throws Exception {
+ File root = folder.getRoot();
+ File home = new File(root, "home");
+ boolean created = home.mkdir();
+ assertTrue(created);
+
+ env.put(ENV_APP_ID, "foo");
+ env.put(ENV_CLIENT_HOME_DIR, home.getAbsolutePath());
+ env.put(ENV_CLIENT_SHIP_FILES, "");
+ env.put(ENV_FLINK_CLASSPATH, "");
+ env.put(ENV_HADOOP_USER_NAME, "foo");
+ env.putIfAbsent(FLINK_YARN_FILES, "");
+ env.put(FLINK_DIST_JAR, new YarnLocalResourceDescriptor(
+ "flink.jar",
+ new Path("/tmp/flink.jar"),
+ 0,
+ System.currentTimeMillis(),
+ LocalResourceVisibility.APPLICATION,
+ LocalResourceType.FILE).toString());
+ env.put(ApplicationConstants.Environment.PWD.key(), home.getAbsolutePath());
+
+ BootstrapTools.writeConfiguration(flinkConfig, new File(home.getAbsolutePath(), FLINK_CONF_FILENAME));
+ }
+
+ @Override
+ protected void preparePreviousAttemptWorkers() {
+ testingYarnAMRMClientAsyncBuilder.setRegisterApplicationMasterFunction(
+ (ignored1, ignored2, ignored3) -> new TestingRegisterApplicationMasterResponse(() -> Collections.singletonList(testingContainer)));
+ }
+
+ @Override
+ protected ResourceManagerDriver<YarnWorkerNode> createResourceManagerDriver() {
+ return new YarnResourceManagerDriver(
+ flinkConfig,
+ new YarnResourceManagerConfiguration(env, "localhost:9000", null),
+ testingYarnResourceManagerClientFactory,
+ testingYarnNodeManagerClientFactory);
+ }
+
+ @Override
+ protected void validateInitialization() throws Exception {
+ assertNotNull(testingYarnAMRMClientAsync);
+ assertNotNull(testingYarnNMClientAsync);
+ verifyFutureCompleted(nodeManagerClientInitFuture);
+ verifyFutureCompleted(nodeManagerClientStartFuture);
+ verifyFutureCompleted(resourceManagerClientInitFuture);
+ verifyFutureCompleted(resourceManagerClientStartFuture);
+ }
+
+ @Override
+ protected void validateWorkersRecoveredFromPreviousAttempt(Collection<YarnWorkerNode> workers) {
+ Assume.assumeTrue(HadoopUtils.isMinHadoopVersion(2, 2));
+ assertThat(workers.size(), is(1));
+
+ final ResourceID resourceId = workers.iterator().next().getResourceID();
+ assertThat(resourceId.toString(), is(testingContainer.getId().toString()));
+ }
+
+ @Override
+ protected void validateTermination() throws Exception {
+ verifyFutureCompleted(nodeManagerClientStopFuture);
+ verifyFutureCompleted(resourceManagerClientStopFuture);
+ }
+
+ @Override
+ protected void validateDeregisterApplication() throws Exception {
+ verifyFutureCompleted(stopAndCleanupClusterFuture);
+ }
+
+ @Override
+ protected void validateRequestedResources(Collection<TaskExecutorProcessSpec> taskExecutorProcessSpecs) throws Exception {
+ assertThat(taskExecutorProcessSpecs.size(), is(1));
+
+ final Resource resource = createTaskManagerContainerFuture.get(TIMEOUT_SEC, TimeUnit.SECONDS);
+
+ assertThat(resource.getMemory(), is(YarnConfiguration.DEFAULT_RM_SCHEDULER_MINIMUM_ALLOCATION_MB));
+ assertThat(resource.getVirtualCores(), is(YarnConfiguration.DEFAULT_RM_SCHEDULER_MINIMUM_ALLOCATION_VCORES));
+ verifyFutureCompleted(removeContainerRequestFuture);
+ }
+
+ @Override
+ protected void validateReleaseResources(Collection<YarnWorkerNode> workerNodes) throws Exception {
+ assertThat(workerNodes.size(), is(1));
+ verifyFutureCompleted(stopContainerAsyncFuture);
+ verifyFutureCompleted(releaseAssignedContainerFuture);
+ }
+
+ ContainerStatus createTestingContainerCompletedStatus(final ContainerId containerId) {
+ return new TestingContainerStatus(containerId, ContainerState.COMPLETE, "Test exit", -1);
+ }
+
+ Container createTestingContainerWithResource(Resource resource) {
+ return YarnResourceManagerDriverTest.createTestingContainerWithResource(resource, containerIdx++);
+ }
+
+ <T> void verifyFutureCompleted(CompletableFuture<T> future) throws Exception {
+ future.get(TIMEOUT_SEC, TimeUnit.SECONDS);
+ }
+
+ void prepareForTestStartTaskExecutorProcessVariousSpec(
+ String startCommand1,
+ String startCommand2,
+ CompletableFuture<Void> startContainerAsyncCommandFuture1,
+ CompletableFuture<Void> startContainerAsyncCommandFuture2,
+ TaskExecutorProcessSpec taskExecutorProcessSpec) {
+ addContainerRequestFutures.add(new CompletableFuture<>());
+ addContainerRequestFutures.add(new CompletableFuture<>());
+ addContainerRequestFutures.add(new CompletableFuture<>());
+ addContainerRequestFutures.add(new CompletableFuture<>());
+
+ testingYarnAMRMClientAsyncBuilder.setAddContainerRequestConsumer((request, ignored) ->
+ addContainerRequestFutures.get(addContainerRequestFuturesNumCompleted.getAndIncrement()).complete(request.getCapability()));
+ testingYarnNMClientAsyncBuilder.setStartContainerAsyncConsumer((ignored1, context, ignored3) -> {
+ if (containsStartCommand(context, startCommand1)) {
+ startContainerAsyncCommandFuture1.complete(null);
+ } else if (containsStartCommand(context, startCommand2)) {
+ startContainerAsyncCommandFuture2.complete(null);
+ }
+ });
+ resourceEventHandlerBuilder.setOnWorkerTerminatedConsumer((ignore1, ignore2) -> getDriver().requestResource(taskExecutorProcessSpec));
+ }
+ }
+}