/*
* 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.apache.tez.dag.app.rm;
import java.io.IOException;
import java.util.Arrays;
import java.util.LinkedList;
import java.util.List;
import java.util.Random;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.atomic.AtomicInteger;
import com.google.common.base.Preconditions;
import com.google.common.primitives.Ints;
import org.apache.hadoop.conf.Configuration;
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.NodeId;
import org.apache.hadoop.yarn.api.records.Priority;
import org.apache.hadoop.yarn.api.records.Resource;
import org.apache.tez.common.TezUtils;
import org.apache.tez.dag.api.TezUncheckedException;
import org.apache.tez.serviceplugins.api.TaskAttemptEndReason;
import org.apache.tez.service.TezTestServiceConfConstants;
import org.apache.tez.serviceplugins.api.TaskScheduler;
import org.apache.tez.serviceplugins.api.TaskSchedulerContext;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
public class TezTestServiceTaskSchedulerService extends TaskScheduler {
private static final Logger
LOG = LoggerFactory.getLogger(TezTestServiceTaskSchedulerService.class);
private final List<String> serviceHosts;
private final ContainerFactory containerFactory;
private final Random random = new Random();
// Currently all services must be running on the same port.
private final int containerPort;
private final ConcurrentMap<Object, ContainerId> runningTasks =
new ConcurrentHashMap<Object, ContainerId>();
// AppIdIdentifier to avoid conflicts with other containers in the system.
// Per instance
private final int memoryPerInstance;
private final int coresPerInstance;
private final int executorsPerInstance;
// Per Executor Thread
private final Resource resourcePerContainer;
// Not registering with the RM. Assuming the main TezScheduler will always run (except local mode),
// and take care of YARN registration.
public TezTestServiceTaskSchedulerService(TaskSchedulerContext taskSchedulerContext) {
// Accepting configuration here to allow setting up fields as final
super(taskSchedulerContext);
this.serviceHosts = new LinkedList<String>();
this.containerFactory = new ContainerFactory(taskSchedulerContext.getApplicationAttemptId(),
taskSchedulerContext.getCustomClusterIdentifier());
Configuration conf = null;
try {
conf = TezUtils.createConfFromUserPayload(taskSchedulerContext.getInitialUserPayload());
} catch (IOException e) {
throw new TezUncheckedException(e);
}
this.memoryPerInstance = conf
.getInt(TezTestServiceConfConstants.TEZ_TEST_SERVICE_MEMORY_PER_INSTANCE_MB, -1);
Preconditions.checkArgument(memoryPerInstance > 0,
TezTestServiceConfConstants.TEZ_TEST_SERVICE_MEMORY_PER_INSTANCE_MB +
" must be configured");
this.executorsPerInstance = conf.getInt(
TezTestServiceConfConstants.TEZ_TEST_SERVICE_NUM_EXECUTORS_PER_INSTANCE,
-1);
Preconditions.checkArgument(executorsPerInstance > 0,
TezTestServiceConfConstants.TEZ_TEST_SERVICE_NUM_EXECUTORS_PER_INSTANCE +
" must be configured");
this.coresPerInstance = conf
.getInt(TezTestServiceConfConstants.TEZ_TEST_SERVICE_VCPUS_PER_INSTANCE,
executorsPerInstance);
this.containerPort = conf.getInt(TezTestServiceConfConstants.TEZ_TEST_SERVICE_RPC_PORT, -1);
Preconditions.checkArgument(executorsPerInstance > 0,
TezTestServiceConfConstants.TEZ_TEST_SERVICE_RPC_PORT + " must be configured");
int memoryPerContainer = (int) (memoryPerInstance / (float) executorsPerInstance);
int coresPerContainer = (int) (coresPerInstance / (float) executorsPerInstance);
this.resourcePerContainer = Resource.newInstance(memoryPerContainer, coresPerContainer);
String[] hosts = conf.getTrimmedStrings(TezTestServiceConfConstants.TEZ_TEST_SERVICE_HOSTS);
if (hosts == null || hosts.length == 0) {
hosts = new String[]{"localhost"};
}
for (String host : hosts) {
serviceHosts.add(host);
}
LOG.info("Running with configuration: " +
"memoryPerInstance=" + memoryPerInstance +
", vcoresPerInstance=" + coresPerInstance +
", executorsPerInstance=" + executorsPerInstance +
", resourcePerContainerInferred=" + resourcePerContainer +
", hosts=" + serviceHosts.toString());
}
@Override
public Resource getAvailableResources() {
// TODO This needs information about all running executors, and the amount of memory etc available across the cluster.
return Resource
.newInstance(Ints.checkedCast(serviceHosts.size() * memoryPerInstance),
serviceHosts.size() * coresPerInstance);
}
@Override
public int getClusterNodeCount() {
return serviceHosts.size();
}
@Override
public void dagComplete() {
}
@Override
public Resource getTotalResources() {
return Resource
.newInstance(Ints.checkedCast(serviceHosts.size() * memoryPerInstance),
serviceHosts.size() * coresPerInstance);
}
@Override
public void blacklistNode(NodeId nodeId) {
LOG.info("BlacklistNode not supported");
}
@Override
public void unblacklistNode(NodeId nodeId) {
LOG.info("unBlacklistNode not supported");
}
@Override
public void allocateTask(Object task, Resource capability, String[] hosts, String[] racks,
Priority priority, Object containerSignature, Object clientCookie) {
String host = selectHost(hosts);
Container container =
containerFactory.createContainer(resourcePerContainer, priority, host, containerPort);
runningTasks.put(task, container.getId());
getContext().taskAllocated(task, clientCookie, container);
}
@Override
public void allocateTask(Object task, Resource capability, ContainerId containerId,
Priority priority, Object containerSignature, Object clientCookie) {
String host = selectHost(null);
Container container =
containerFactory.createContainer(resourcePerContainer, priority, host, containerPort);
runningTasks.put(task, container.getId());
getContext().taskAllocated(task, clientCookie, container);
}
@Override
public boolean deallocateTask(Object task, boolean taskSucceeded, TaskAttemptEndReason endReason, String diagnostics) {
ContainerId containerId = runningTasks.remove(task);
if (containerId == null) {
LOG.error("Could not determine ContainerId for task: " + task +
" . Could have hit a race condition. Ignoring." +
" The query may hang since this \"unknown\" container is now taking up a slot permanently");
return false;
}
getContext().containerBeingReleased(containerId);
return true;
}
@Override
public Object deallocateContainer(ContainerId containerId) {
LOG.info("Ignoring deallocateContainer for containerId: " + containerId);
return null;
}
@Override
public void setShouldUnregister() {
}
@Override
public boolean hasUnregistered() {
// Nothing to do. No registration involved.
return true;
}
private String selectHost(String[] requestedHosts) {
String host;
if (requestedHosts != null && requestedHosts.length > 0) {
Arrays.sort(requestedHosts);
host = requestedHosts[0];
LOG.info("Selected host: " + host + " from requested hosts: " + Arrays.toString(requestedHosts));
} else {
host = serviceHosts.get(random.nextInt(serviceHosts.size()));
LOG.info("Selected random host: " + host + " since the request contained no host information");
}
return host;
}
static class ContainerFactory {
AtomicInteger nextId;
final ApplicationAttemptId customAppAttemptId;
public ContainerFactory(ApplicationAttemptId appAttemptId, long appIdLong) {
this.nextId = new AtomicInteger(1);
ApplicationId appId = ApplicationId
.newInstance(appIdLong, appAttemptId.getApplicationId().getId());
this.customAppAttemptId = ApplicationAttemptId
.newInstance(appId, appAttemptId.getAttemptId());
}
@SuppressWarnings("deprecation")
public Container createContainer(Resource capability, Priority priority, String hostname, int port) {
ContainerId containerId = ContainerId.newInstance(customAppAttemptId, nextId.getAndIncrement());
NodeId nodeId = NodeId.newInstance(hostname, port);
String nodeHttpAddress = "hostname:0";
Container container = Container.newInstance(containerId,
nodeId,
nodeHttpAddress,
capability,
priority,
null);
return container;
}
}
}