/*
* Copyright 2014-2016 the original author or authors.
*
* Licensed 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.springframework.yarn.am.allocate;
import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Map.Entry;
import java.util.Queue;
import java.util.Set;
import java.util.concurrent.ConcurrentLinkedQueue;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.apache.hadoop.yarn.api.protocolrecords.AllocateRequest;
import org.apache.hadoop.yarn.api.protocolrecords.AllocateResponse;
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.Priority;
import org.apache.hadoop.yarn.api.records.Resource;
import org.apache.hadoop.yarn.api.records.ResourceRequest;
import org.apache.hadoop.yarn.util.Records;
import org.springframework.util.StringUtils;
import org.springframework.yarn.am.allocate.AllocationGroup.ContainerAllocationValues;
import org.springframework.yarn.am.allocate.DefaultAllocateCountTracker.AllocateCountInfo;
import org.springframework.yarn.listener.CompositeContainerAllocatorListener;
import org.springframework.yarn.listener.ContainerAllocatorListener;
import org.springframework.yarn.support.compat.ResourceCompat;
/**
* Default allocator which polls resource manager, requests new containers
* and acts as a heart beat sender at the same time.
*
* @author Janne Valkealahti
*
*/
public class DefaultContainerAllocator extends AbstractPollingAllocator implements ContainerAllocator {
private static final Log log = LogFactory.getLog(DefaultContainerAllocator.class);
/** Listener dispatcher for events */
private CompositeContainerAllocatorListener allocatorListener = new CompositeContainerAllocatorListener();
/** Container request priority */
private int priority = 0;
/** Container label expression */
private String labelExpression;
/** Resource capability as of cores */
private int virtualcores = 1;
/** Resource capability as of memory */
private int memory = 64;
/** Locality relaxing */
private boolean locality = false;
/** Increasing counter for rpc request id*/
private AtomicInteger requestId = new AtomicInteger();
/** Current progress reported during allocate requests */
private float applicationProgress = 0;
/** Queued container id's to be released */
private Queue<ContainerId> releaseContainers = new ConcurrentLinkedQueue<ContainerId>();
/** Internal set of containers marked as garbage by allocate tracker */
private Set<ContainerId> garbageContainers = new HashSet<ContainerId>();
/** Flag helping to avoid allocation garbage */
private AtomicBoolean allocationDirty = new AtomicBoolean();
/** Empty list for requests without container asks */
private final List<ResourceRequest> EMPTY = new ArrayList<ResourceRequest>();
/** Groups for allocation tracking */
private final AllocationGroups allocationGroups = new AllocationGroups();
@Override
protected void onInit() throws Exception {
super.onInit();
internalInit();
}
@Override
public void allocateContainers(int count) {
if (log.isDebugEnabled()) {
log.debug("Incoming count: " + count);
}
ContainerAllocateData cad = new ContainerAllocateData();
cad.addAny(count);
allocateContainers(cad);
}
@Override
public void addListener(ContainerAllocatorListener listener) {
allocatorListener.register(listener);
}
@Override
public void allocateContainers(ContainerAllocateData containerAllocateData) {
log.info("Incoming containerAllocateData: " + containerAllocateData);
String id = StringUtils.hasText(containerAllocateData.getId()) ? containerAllocateData.getId() : "";
AllocationGroup group = allocationGroups.get(id);
if (group == null) {
group = allocationGroups.get("");
}
boolean dirty = false;
ContainerAllocateData cad = containerAllocateData.byAny();
if (cad.hasData()) {
dirty = true;
DefaultAllocateCountTracker tracker = group.getAllocateCountTracker(AllocationGroup.GROUP_ANY);
if (tracker == null) {
allocationGroups.reserve(id, AllocationGroup.GROUP_ANY);
tracker = new DefaultAllocateCountTracker(AllocationGroup.GROUP_ANY, getConfiguration());
group.setAllocateCountTracker(AllocationGroup.GROUP_ANY, tracker);
}
log.info("State allocateCountTracker before adding allocation data: " + tracker);
tracker.addContainers(cad);
log.info("State allocateCountTracker after adding allocation data: " + tracker);
}
cad = containerAllocateData.byHosts();
if (cad.hasData()) {
dirty = true;
DefaultAllocateCountTracker tracker = group.getAllocateCountTracker(AllocationGroup.GROUP_HOST);
if (tracker == null) {
allocationGroups.reserve(id, AllocationGroup.GROUP_HOST);
tracker = new DefaultAllocateCountTracker(AllocationGroup.GROUP_HOST, getConfiguration());
group.setAllocateCountTracker(AllocationGroup.GROUP_HOST, tracker);
}
log.info("State allocateCountTracker before adding allocation data: " + tracker);
tracker.addContainers(cad);
log.info("State allocateCountTracker after adding allocation data: " + tracker);
}
cad = containerAllocateData.byRacks();
if (cad.hasData()) {
dirty = true;
DefaultAllocateCountTracker tracker = group.getAllocateCountTracker(AllocationGroup.GROUP_RACK);
if (tracker == null) {
allocationGroups.reserve(id, AllocationGroup.GROUP_RACK);
tracker = new DefaultAllocateCountTracker(AllocationGroup.GROUP_RACK, getConfiguration());
group.setAllocateCountTracker(AllocationGroup.GROUP_RACK, tracker);
}
log.info("State allocateCountTracker before adding allocation data: " + tracker);
tracker.addContainers(cad);
log.info("State allocateCountTracker after adding allocation data: " + tracker);
}
if (dirty) {
allocationDirty.set(true);
}
}
@Override
public void releaseContainers(List<Container> containers) {
for (Container container : containers) {
releaseContainer(container.getId());
}
}
@Override
public void releaseContainer(ContainerId containerId) {
log.info("Adding new container to be released containerId=" + containerId);
releaseContainers.add(containerId);
}
@Override
protected AllocateResponse doContainerRequest() {
List<ResourceRequest> requestedContainers = null;
if (allocationDirty.getAndSet(false)) {
requestedContainers = createRequests();
} else {
requestedContainers = EMPTY;
}
// add pending containers to be released
List<ContainerId> release = new ArrayList<ContainerId>();
ContainerId element = null;
while ((element = releaseContainers.poll()) != null) {
release.add(element);
}
if (log.isDebugEnabled()) {
log.debug("Requesting containers using " + requestedContainers.size() + " requests.");
for (ResourceRequest resourceRequest : requestedContainers) {
log.debug("ResourceRequest: " + resourceRequest + " with count=" +
resourceRequest.getNumContainers() + " with hostName=" + resourceRequest.getResourceName());
}
log.debug("Releasing containers " + release.size());
for (ContainerId cid : release) {
log.debug("Release container=" + cid);
}
log.debug("Request id will be: " + requestId.get());
}
// build the allocation request
AllocateRequest request = Records.newRecord(AllocateRequest.class);
request.setResponseId(requestId.get());
request.setAskList(requestedContainers);
request.setReleaseList(release);
request.setProgress(applicationProgress);
// do request and return response
AllocateResponse allocate = getRmTemplate().allocate(request);
requestId.set(allocate.getResponseId());
return allocate;
}
@Override
protected List<Container> preProcessAllocatedContainers(List<Container> containers) {
// checking if we have standing allocation counts,
// if not assume as garbage and send to release queue.
// what's left will be out of our hand and expected to be
// processed by the listener and eventually send back
// to us as a released container.
List<Container> preProcessed = new ArrayList<Container>();
for (Container container : containers) {
DefaultAllocateCountTracker tracker = allocationGroups.get(container.getPriority().getPriority())
.getAllocateCountTracker(container.getPriority().getPriority());
if (log.isDebugEnabled()) {
log.debug("State allocateCountTracker before handling allocated container: " + tracker);
}
Container processed = tracker.processAllocatedContainer(container);
if (processed != null) {
preProcessed.add(processed);
} else {
garbageContainers.add(container.getId());
releaseContainers.add(container.getId());
}
if (log.isDebugEnabled()) {
log.debug("State allocateCountTracker after handling allocated container: " + tracker);
}
}
return preProcessed;
}
@Override
protected void handleAllocatedContainers(List<Container> containers) {
allocatorListener.allocated(containers);
}
@Override
protected void handleCompletedContainers(List<ContainerStatus> containerStatuses) {
// strip away containers which were already marked
// garbage by allocate tracker. system
// never knew those even exist and might create mess
// with monitor component. monitor only sees
// complete status which is also the case for garbage
// when it's released.
List<ContainerStatus> garbageFree = new ArrayList<ContainerStatus>();
for (ContainerStatus status : containerStatuses) {
if (!garbageContainers.contains(status.getContainerId())) {
garbageFree.add(status);
}
}
allocatorListener.completed(garbageFree);
}
@Override
public void setProgress(float progress) {
applicationProgress = progress;
}
/**
* Sets the allocation values for given identifier.
*
* @param id the allocation identifier
* @param priority the base priority
* @param labelExpression the label expression
* @param virtualcores the cpu count
* @param memory the memory
* @param locality the locality flag
*/
public void setAllocationValues(String id, Integer priority, String labelExpression, Integer virtualcores, Integer memory, Boolean locality) {
if (log.isTraceEnabled()) {
log.trace("setAllocationValues 1: id=" + id + " priority=" + priority + " labelExpression=" + labelExpression + " cores="
+ virtualcores + " memory=" + memory + " locality=" + locality);
}
id = StringUtils.hasText(id) ? id : "";
AllocationGroup group = allocationGroups.add(id, priority);
group.setContainerAllocationValues(new ContainerAllocationValues(priority, labelExpression, virtualcores, memory, locality));
}
/**
* Gets the priority for container request.
*
* @return the priority
*/
public int getPriority() {
return priority;
}
/**
* Sets the priority for container request.
*
* @param priority the new priority
*/
public void setPriority(int priority) {
this.priority = priority;
}
/**
* Sets the label expression.
*
* @param labelExpression the new label expression
*/
public void setLabelExpression(String labelExpression) {
this.labelExpression = labelExpression;
}
/**
* Gets the virtualcores for container request.
*
* @return the virtualcores
*/
public int getVirtualcores() {
return virtualcores;
}
/**
* Sets the virtualcores for container request defining
* <em>number of virtual cpu cores</em> of the resource.
*
* @param virtualcores the new virtualcores
*/
public void setVirtualcores(int virtualcores) {
this.virtualcores = virtualcores;
}
/**
* Gets the memory for container request.
*
* @return the memory
*/
public int getMemory() {
return memory;
}
/**
* Sets the memory for container request defining
* <em>memory</em> of the resource.
*
* @param memory the new memory
*/
public void setMemory(int memory) {
this.memory = memory;
}
/**
* Checks if is locality relax flag is enabled.
*
* @return true, if is locality is enabled.
*/
public boolean isLocality() {
return locality;
}
/**
* Sets the flag telling if resource allocation
* should not be relaxed. Setting this to true
* means that relaxing is not used. Default value
* is false.
*
* @param locality the new locality relax flag
*/
public void setLocality(boolean locality) {
this.locality = locality;
}
private List<ResourceRequest> createRequests() {
List<ResourceRequest> requestedContainers = new ArrayList<ResourceRequest>();
for (AllocationGroup group : allocationGroups.getGroups()) {
for (DefaultAllocateCountTracker tracker : group.getAllocateCountTrackers()) {
AllocateCountInfo allocateCounts = tracker.getAllocateCounts();
Integer pri = group.getPriority(tracker.getId());
ContainerAllocationValues allocationValues = group.getContainerAllocationValues();
if (pri == null) {
pri = allocationValues.priority;
}
if (log.isTraceEnabled()) {
log.trace("trace 1 " + allocationValues.locality);
log.trace("trace 2 tracker id:" + tracker.getId());
}
boolean hostsAdded = false;
for (Entry<String, Integer> entry : allocateCounts.hostsInfo.entrySet()) {
if (log.isTraceEnabled()) {
log.trace("trace 3 entry key=" + entry.getKey() + " value=" + entry.getValue());
}
requestedContainers.add(getContainerResourceRequest(allocationValues, pri, entry.getValue(), entry.getKey(),
true));
hostsAdded = true;
}
for (Entry<String, Integer> entry : allocateCounts.racksInfo.entrySet()) {
if (log.isTraceEnabled()) {
log.trace("trace 4 entry key=" + entry.getKey() + " value=" + entry.getValue());
}
requestedContainers.add(getContainerResourceRequest(allocationValues, pri, entry.getValue(), entry.getKey(),
(hostsAdded && allocationValues.locality) ? false : true));
}
for (Entry<String, Integer> entry : allocateCounts.anysInfo.entrySet()) {
if (log.isTraceEnabled()) {
log.trace("trace 5 entry key=" + entry.getKey() + " value=" + entry.getValue());
}
// need to force locality flag for any in case of
// any group
boolean localityForAny = tracker.getId().equals(AllocationGroup.GROUP_ANY) ? true : !allocationValues.locality;
requestedContainers.add(getContainerResourceRequest(allocationValues, pri, entry.getValue(), entry.getKey(),
localityForAny));
}
}
}
return requestedContainers;
}
/**
* Internal init which should only configure this class and
* should not touch parent classes.
*/
private void internalInit() {
setAllocationValues(null, priority, labelExpression, virtualcores, memory, locality);
for (DefaultAllocateCountTracker tracker : allocationGroups.getAllocateCountTrackers()) {
tracker.setConfiguration(getConfiguration());
}
}
/**
* Utility method creating a {@link ResourceRequest}.
*
* @param allocationValues
* @param priority
* @param numContainers
* @param hostName
* @param relaxLocality
* @return request to be sent to resource manager
*/
private ResourceRequest getContainerResourceRequest(ContainerAllocationValues allocationValues, int priority,
int numContainers, String hostName, boolean relaxLocality) {
ResourceRequest request = Records.newRecord(ResourceRequest.class);
request.setRelaxLocality(relaxLocality);
request.setResourceName(hostName);
request.setNumContainers(numContainers);
request.setNodeLabelExpression(allocationValues.labelExpression);
Priority pri = Records.newRecord(Priority.class);
pri.setPriority(priority);
request.setPriority(pri);
Resource capability = Records.newRecord(Resource.class);
capability.setMemory(allocationValues.memory);
ResourceCompat.setVirtualCores(capability, allocationValues.virtualcores);
request.setCapability(capability);
return request;
}
}