org.apache.myriad.scheduler.fgs.NMHeartBeatHandler.java Source code

Java tutorial

Introduction

Here is the source code for org.apache.myriad.scheduler.fgs.NMHeartBeatHandler.java

Source

/**
 * 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
 * <p/>
 * http://www.apache.org/licenses/LICENSE-2.0
 * <p/>
 * 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.myriad.scheduler.fgs;

import com.google.common.annotations.VisibleForTesting;
import java.util.ArrayList;
import java.util.List;
import javax.inject.Inject;
import org.apache.hadoop.yarn.api.records.ContainerState;
import org.apache.hadoop.yarn.api.records.ContainerStatus;
import org.apache.hadoop.yarn.api.records.NodeId;
import org.apache.hadoop.yarn.api.records.Resource;
import org.apache.hadoop.yarn.server.resourcemanager.RMContext;
import org.apache.hadoop.yarn.server.resourcemanager.rmcontainer.RMContainer;
import org.apache.hadoop.yarn.server.resourcemanager.rmnode.RMNode;
import org.apache.hadoop.yarn.server.resourcemanager.rmnode.RMNodeEvent;
import org.apache.hadoop.yarn.server.resourcemanager.rmnode.RMNodeStatusEvent;
import org.apache.hadoop.yarn.server.resourcemanager.scheduler.AbstractYarnScheduler;
import org.apache.hadoop.yarn.util.resource.Resources;
import org.apache.mesos.Protos;
import org.apache.mesos.Protos.Offer;
import org.apache.myriad.configuration.NodeManagerConfiguration;
import org.apache.myriad.scheduler.MyriadDriver;
import org.apache.myriad.scheduler.SchedulerUtils;
import org.apache.myriad.scheduler.yarn.interceptor.BaseInterceptor;
import org.apache.myriad.scheduler.yarn.interceptor.InterceptorRegistry;
import org.apache.myriad.state.SchedulerState;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * Handles node manager heartbeat.
 */
public class NMHeartBeatHandler extends BaseInterceptor {
    @VisibleForTesting
    Logger logger = LoggerFactory.getLogger(NMHeartBeatHandler.class);

    private final AbstractYarnScheduler yarnScheduler;
    private final MyriadDriver myriadDriver;
    private final YarnNodeCapacityManager yarnNodeCapacityMgr;
    private final OfferLifecycleManager offerLifecycleMgr;
    private final NodeStore nodeStore;
    private final SchedulerState state;
    private final NodeManagerConfiguration conf;

    @Inject
    public NMHeartBeatHandler(InterceptorRegistry registry, AbstractYarnScheduler yarnScheduler,
            MyriadDriver myriadDriver, YarnNodeCapacityManager yarnNodeCapacityMgr,
            OfferLifecycleManager offerLifecycleMgr, NodeStore nodeStore, SchedulerState state,
            NodeManagerConfiguration conf) {

        if (registry != null) {
            registry.register(this);
        }

        this.yarnScheduler = yarnScheduler;
        this.myriadDriver = myriadDriver;
        this.yarnNodeCapacityMgr = yarnNodeCapacityMgr;
        this.offerLifecycleMgr = offerLifecycleMgr;
        this.nodeStore = nodeStore;
        this.state = state;
        this.conf = conf;
    }

    @Override
    public CallBackFilter getCallBackFilter() {
        return new CallBackFilter() {
            @Override
            public boolean allowCallBacksForNode(NodeId nodeManager) {
                return SchedulerUtils.isEligibleForFineGrainedScaling(nodeManager.getHost(), state);
            }
        };
    }

    @Override
    public void beforeRMNodeEventHandled(RMNodeEvent event, RMContext context) {
        switch (event.getType()) {
        case STARTED:
            // Since the RMNode was just started, it should not have a non-zero capacity
            RMNode rmNode = context.getRMNodes().get(event.getNodeId());

            if (isNonZeroCapacityNode(rmNode)) {
                Resource totalCapability = rmNode.getTotalCapability();
                logger.warn(
                        "FineGrainedScaling feature got invoked for a NM with non-zero capacity. Host: {}, Mem: {}, CPU: {}. Setting the "
                                + "NM's capacity to (0G,0CPU)",
                        rmNode.getHostName(), totalCapability.getMemory(), totalCapability.getVirtualCores());
                totalCapability.setMemory(0);
                totalCapability.setVirtualCores(0);
            }
            break;

        case STATUS_UPDATE:
            handleStatusUpdate(event, context);
            break;

        default:
            break;
        }
    }

    @VisibleForTesting
    protected boolean isNonZeroCapacityNode(RMNode node) {
        Resource resource = node.getTotalCapability();
        return (resource.getMemory() != 0 || resource.getVirtualCores() != 0);
    }

    @VisibleForTesting
    protected void handleStatusUpdate(RMNodeEvent event, RMContext context) {
        if (!(event instanceof RMNodeStatusEvent)) {
            logger.error("{} not an instance of {}", event.getClass().getName(), RMNodeStatusEvent.class.getName());
            return;
        }

        RMNodeStatusEvent statusEvent = (RMNodeStatusEvent) event;
        RMNode rmNode = context.getRMNodes().get(event.getNodeId());
        String hostName = rmNode.getNodeID().getHost();

        Node host = nodeStore.getNode(hostName);
        if (host != null) {
            host.snapshotRunningContainers();
        }

        /*
         * Set the new node capacity which is the sum of the current node resources plus those offered by Mesos. 
         * If the sum is greater than the max capacity of the node, reject the offer.
         */
        Resource offeredResources = getNewResourcesOfferedByMesos(hostName);
        Resource currentResources = getResourcesUnderUse(statusEvent);

        if (offerWithinResourceLimits(currentResources, offeredResources)) {
            yarnNodeCapacityMgr.setNodeCapacity(rmNode, Resources.add(currentResources, offeredResources));
            logger.info("Updated resources for {} with {} cores and {} memory", rmNode.getNode().getName(),
                    offeredResources.getVirtualCores(), offeredResources.getMemory());
        } else {
            logger.info("Did not update {} with {} cores and {} memory, over max cpu cores and/or max memory",
                    rmNode.getNode().getName(), offeredResources.getVirtualCores(), offeredResources.getMemory());
        }
    }

    @VisibleForTesting
    protected boolean offerWithinResourceLimits(Resource currentResources, Resource offeredResources) {
        int newMemory = currentResources.getMemory() + offeredResources.getMemory();
        int newCores = currentResources.getVirtualCores() + offeredResources.getVirtualCores();

        return (newMemory <= conf.getJvmMaxMemoryMB() && newCores <= conf.getMaxCpus());
    }

    @VisibleForTesting
    protected Resource getNewResourcesOfferedByMesos(String hostname) {
        OfferFeed feed = offerLifecycleMgr.getOfferFeed(hostname);
        List<Offer> offers = new ArrayList<>();
        Protos.Offer offer;

        while ((offer = feed.poll()) != null) {
            offers.add(offer);
            offerLifecycleMgr.markAsConsumed(offer);
        }

        Resource fromMesosOffers = OfferUtils.getYarnResourcesFromMesosOffers(offers);

        if (logger.isDebugEnabled()) {
            logger.debug("NM on host {} got {} CPUs and {} memory from mesos", hostname,
                    fromMesosOffers.getVirtualCores(), fromMesosOffers.getMemory());
        }

        return fromMesosOffers;
    }

    @VisibleForTesting
    protected Resource getResourcesUnderUse(RMNodeStatusEvent statusEvent) {
        Resource usedResources = Resource.newInstance(0, 0);
        for (ContainerStatus status : statusEvent.getContainers()) {
            if (containerInUse(status)) {
                RMContainer rmContainer = yarnScheduler.getRMContainer(status.getContainerId());
                // (sdaingade) This check is needed as RMContainer information may not be populated
                // immediately after a RM restart.
                if (rmContainer != null) {
                    Resources.addTo(usedResources, rmContainer.getAllocatedResource());
                }
            }
        }
        return usedResources;
    }

    private boolean containerInUse(ContainerStatus status) {
        return (status.getState() == ContainerState.NEW || status.getState() == ContainerState.RUNNING);
    }
}