/*
* 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.twill.internal.yarn;
import com.google.common.base.Function;
import com.google.common.base.Preconditions;
import com.google.common.collect.ArrayListMultimap;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.Iterables;
import com.google.common.collect.Lists;
import com.google.common.collect.Multimap;
import com.google.common.util.concurrent.AbstractIdleService;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.yarn.api.ApplicationConstants;
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.ContainerStatus;
import org.apache.hadoop.yarn.api.records.FinalApplicationStatus;
import org.apache.hadoop.yarn.api.records.Resource;
import org.apache.hadoop.yarn.ipc.YarnRPC;
import org.apache.hadoop.yarn.util.ConverterUtils;
import org.apache.twill.internal.ProcessLauncher;
import org.apache.twill.internal.appmaster.RunnableProcessLauncher;
import org.apache.twill.internal.yarn.ports.AMRMClient;
import org.apache.twill.internal.yarn.ports.AMRMClientImpl;
import org.apache.twill.internal.yarn.ports.AllocationResponse;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import java.net.InetSocketAddress;
import java.net.URL;
import java.util.List;
import java.util.UUID;
/**
*
*/
public final class Hadoop20YarnAMClient extends AbstractIdleService implements YarnAMClient {
private static final Logger LOG = LoggerFactory.getLogger(Hadoop20YarnAMClient.class);
private static final Function<ContainerStatus, YarnContainerStatus> STATUS_TRANSFORM;
static {
STATUS_TRANSFORM = new Function<ContainerStatus, YarnContainerStatus>() {
@Override
public YarnContainerStatus apply(ContainerStatus status) {
return new Hadoop20YarnContainerStatus(status);
}
};
}
private final ContainerId containerId;
private final Multimap<String, AMRMClient.ContainerRequest> containerRequests;
private final AMRMClient amrmClient;
private final YarnNMClient nmClient;
private InetSocketAddress trackerAddr;
private URL trackerUrl;
private Resource maxCapability;
private Resource minCapability;
public Hadoop20YarnAMClient(Configuration conf) {
String masterContainerId = System.getenv().get(ApplicationConstants.AM_CONTAINER_ID_ENV);
Preconditions.checkArgument(masterContainerId != null,
"Missing %s from environment", ApplicationConstants.AM_CONTAINER_ID_ENV);
this.containerId = ConverterUtils.toContainerId(masterContainerId);
this.containerRequests = ArrayListMultimap.create();
this.amrmClient = new AMRMClientImpl(containerId.getApplicationAttemptId());
this.amrmClient.init(conf);
this.nmClient = new Hadoop20YarnNMClient(YarnRPC.create(conf), conf);
}
@Override
protected void startUp() throws Exception {
Preconditions.checkNotNull(trackerAddr, "Tracker address not set.");
Preconditions.checkNotNull(trackerUrl, "Tracker URL not set.");
amrmClient.start();
String url = String.format("%s:%d",
trackerUrl.getHost(),
trackerUrl.getPort() == -1 ? trackerUrl.getDefaultPort() : trackerUrl.getPort());
RegisterApplicationMasterResponse response = amrmClient.registerApplicationMaster(trackerAddr.getHostName(),
trackerAddr.getPort(),
url);
maxCapability = response.getMaximumResourceCapability();
minCapability = response.getMinimumResourceCapability();
}
@Override
protected void shutDown() throws Exception {
amrmClient.unregisterApplicationMaster(FinalApplicationStatus.SUCCEEDED, null, trackerUrl.toString());
amrmClient.stop();
}
@Override
public ContainerId getContainerId() {
return containerId;
}
@Override
public String getHost() {
return System.getenv().get(ApplicationConstants.NM_HOST_ENV);
}
@Override
public void setTracker(InetSocketAddress trackerAddr, URL trackerUrl) {
this.trackerAddr = trackerAddr;
this.trackerUrl = trackerUrl;
}
@Override
public synchronized void allocate(float progress, AllocateHandler handler) throws Exception {
AllocationResponse response = amrmClient.allocate(progress);
List<ProcessLauncher<YarnContainerInfo>> launchers
= Lists.newArrayListWithCapacity(response.getAllocatedContainers().size());
for (Container container : response.getAllocatedContainers()) {
launchers.add(new RunnableProcessLauncher(new Hadoop20YarnContainerInfo(container), nmClient));
}
if (!launchers.isEmpty()) {
handler.acquired(launchers);
// If no process has been launched through the given launcher, return the container.
for (ProcessLauncher<YarnContainerInfo> l : launchers) {
// This cast always works.
RunnableProcessLauncher launcher = (RunnableProcessLauncher) l;
if (!launcher.isLaunched()) {
Container container = launcher.getContainerInfo().getContainer();
LOG.info("Nothing to run in container, releasing it: {}", container);
amrmClient.releaseAssignedContainer(container.getId());
}
}
}
List<YarnContainerStatus> completed = ImmutableList.copyOf(
Iterables.transform(response.getCompletedContainersStatuses(), STATUS_TRANSFORM));
if (!completed.isEmpty()) {
handler.completed(completed);
}
}
@Override
public ContainerRequestBuilder addContainerRequest(Resource capability) {
return addContainerRequest(capability, 1);
}
@Override
public ContainerRequestBuilder addContainerRequest(Resource capability, int count) {
return new ContainerRequestBuilder(adjustCapability(capability), count) {
@Override
public String apply() {
synchronized (Hadoop20YarnAMClient.this) {
String id = UUID.randomUUID().toString();
String[] hosts = this.hosts.isEmpty() ? null : this.hosts.toArray(new String[this.hosts.size()]);
String[] racks = this.racks.isEmpty() ? null : this.racks.toArray(new String[this.racks.size()]);
for (int i = 0; i < count; i++) {
AMRMClient.ContainerRequest request = new AMRMClient.ContainerRequest(capability, hosts, racks,
priority, 1);
containerRequests.put(id, request);
amrmClient.addContainerRequest(request);
}
return id;
}
}
};
}
@Override
public synchronized void completeContainerRequest(String id) {
for (AMRMClient.ContainerRequest request : containerRequests.removeAll(id)) {
amrmClient.removeContainerRequest(request);
}
}
private Resource adjustCapability(Resource resource) {
int cores = YarnUtils.getVirtualCores(resource);
int updatedCores = Math.max(Math.min(cores, YarnUtils.getVirtualCores(maxCapability)),
YarnUtils.getVirtualCores(minCapability));
// Try and set the virtual cores, which older versions of YARN don't support this.
if (cores != updatedCores && YarnUtils.setVirtualCores(resource, updatedCores)) {
LOG.info("Adjust virtual cores requirement from {} to {}.", cores, updatedCores);
}
int updatedMemory = Math.min(resource.getMemory(), maxCapability.getMemory());
int minMemory = minCapability.getMemory();
updatedMemory = (int) Math.ceil(((double) updatedMemory / minMemory)) * minMemory;
if (resource.getMemory() != updatedMemory) {
resource.setMemory(updatedMemory);
LOG.info("Adjust memory requirement from {} to {} MB.", resource.getMemory(), updatedMemory);
}
return resource;
}
}