Package org.springframework.yarn.am.cluster

Source Code of org.springframework.yarn.am.cluster.AbstractContainerClusterAppmaster$ContainerLaunchContextModifyInterceptor

/*
* Copyright 2014 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.cluster;

import java.util.Arrays;
import java.util.HashMap;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Queue;

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.apache.hadoop.yarn.api.records.Container;
import org.apache.hadoop.yarn.api.records.ContainerLaunchContext;
import org.apache.hadoop.yarn.api.records.ContainerStatus;
import org.apache.hadoop.yarn.api.records.LocalResource;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.core.task.TaskExecutor;
import org.springframework.messaging.support.MessageBuilder;
import org.springframework.scheduling.TaskScheduler;
import org.springframework.util.Assert;
import org.springframework.util.StringUtils;
import org.springframework.yarn.YarnSystemException;
import org.springframework.yarn.am.AbstractEventingAppmaster;
import org.springframework.yarn.am.AppmasterCmOperations;
import org.springframework.yarn.am.AppmasterCmTemplate;
import org.springframework.yarn.am.ContainerLauncherInterceptor;
import org.springframework.yarn.am.allocate.AbstractAllocator;
import org.springframework.yarn.am.allocate.ContainerAllocateData;
import org.springframework.yarn.am.container.AbstractLauncher;
import org.springframework.yarn.am.grid.Grid;
import org.springframework.yarn.am.grid.GridMember;
import org.springframework.yarn.am.grid.GridProjection;
import org.springframework.yarn.am.grid.GridProjectionFactory;
import org.springframework.yarn.am.grid.GridProjectionFactoryLocator;
import org.springframework.yarn.am.grid.ProjectedGrid;
import org.springframework.yarn.am.grid.listener.ProjectedGridListenerAdapter;
import org.springframework.yarn.am.grid.support.AbstractGridProjection;
import org.springframework.yarn.am.grid.support.DefaultGrid;
import org.springframework.yarn.am.grid.support.DefaultGridMember;
import org.springframework.yarn.am.grid.support.DefaultProjectedGrid;
import org.springframework.yarn.am.grid.support.ProjectionData;
import org.springframework.yarn.am.grid.support.ProjectionDataRegistry;
import org.springframework.yarn.am.grid.support.SatisfyStateData;
import org.springframework.yarn.am.monitor.ContainerAware;
import org.springframework.yarn.fs.MultiResourceLocalizer;
import org.springframework.yarn.support.PollingTaskSupport;
import org.springframework.yarn.support.statemachine.StateMachine;
import org.springframework.yarn.support.statemachine.config.EnumStateMachineFactory;
import org.springframework.yarn.support.statemachine.state.State;

/**
* Base implementation of a {@link ContainerClusterAppmaster}.
*
* @author Janne Valkealahti
*
*/
public abstract class AbstractContainerClusterAppmaster extends AbstractEventingAppmaster implements ContainerClusterAppmaster {

  private static final Log log = LogFactory.getLog(AbstractContainerClusterAppmaster.class);

  /** ClusterId to ContainerCluster map */
  final Map<String, ContainerCluster> clusters = new HashMap<String, ContainerCluster>();

  /** ClusterId to Cluster definition id mapping */
  final Map<String, String> clusterIdToRef = new HashMap<String, String>();

  /** Containers scheduled to be killed */
  private final Queue<Container> killQueue = new LinkedList<Container>();

  /** Grid tracking generic grid members */
  private Grid grid;

  /** Projected grid tracking container clusters */
  ProjectedGrid projectedGrid;

  /** Poller for handling periodic tasks */
  private ClusterTaskPoller clusterTaskPoller;

  /** Factory for building on-demand state machines */
  private EnumStateMachineFactory<ClusterState, ClusterEvent> stateMachineFactory;

  /** Locator for factories creating projections */
  private GridProjectionFactoryLocator gridProjectionFactoryLocator;

  /** Projection data registry */
  private ProjectionDataRegistry projectionDataRegistry;

  @Override
  protected void onInit() throws Exception {
    super.onInit();
    grid = doCreateGrid();
    Assert.notNull(grid, "Grid must be set");
    projectedGrid = doCreateProjectedGrid(grid);
    Assert.notNull(projectedGrid, "ProjectedGrid must be set");
    projectedGrid.addProjectedGridListener(new CommandDispatchListener());

    if (getLauncher() instanceof AbstractLauncher) {
      ((AbstractLauncher)getLauncher()).addInterceptor(new ContainerLaunchContextModifyInterceptor());
    }
  }

  @Override
  protected void doStart() {
    super.doStart();
    if (projectionDataRegistry != null) {
      Map<String, ProjectionData> defaults = projectionDataRegistry.getProjectionDatas();
      for (Entry<String, ProjectionData> entry : defaults.entrySet()) {
        // no type means it's missing projection settings
        // so assume it's a blueprint and don't try to create
        // cluster automatically
        if (StringUtils.hasText(entry.getValue().getType())) {
          createContainerCluster(entry.getKey(), entry.getValue());
          startContainerCluster(entry.getKey());
        }
      }
    }
  }

  @Override
  protected void doStop() {
    if (clusterTaskPoller != null) {
      clusterTaskPoller.stop();
      clusterTaskPoller = null;
    }
    super.doStop();
  }

  @Override
  public void submitApplication() {
    log.info("Submitting application");
    registerAppmaster();
    start();
    if(getAllocator() instanceof AbstractAllocator) {
      ((AbstractAllocator)getAllocator()).setApplicationAttemptId(getApplicationAttemptId());
    }
    clusterTaskPoller = new ClusterTaskPoller(getTaskScheduler(), getTaskExecutor());
    clusterTaskPoller.init();
    clusterTaskPoller.start();
  }

  @Override
  protected void onContainerAllocated(Container container) {
    if (getMonitor() instanceof ContainerAware) {
      ((ContainerAware)getMonitor()).onContainer(Arrays.asList(container));
    }
    DefaultGridMember member = new DefaultGridMember(container);
    if (grid.addMember(member)) {
      ContainerCluster cluster = findContainerClusterByContainer(container);
      if (cluster != null) {
        getLauncher().launchContainer(container, onContainerLaunchCommands(container, cluster, getCommands(clusterIdToRef.get(cluster.getId()))));
      } else {
        getLauncher().launchContainer(container, getCommands());
      }
    } else {
      getAllocator().releaseContainers(Arrays.asList(container));
    }
  }

  @Override
  protected void onContainerLaunched(Container container) {
    if (getMonitor() instanceof ContainerAware) {
      ((ContainerAware)getMonitor()).onContainer(Arrays.asList(container));
    }
  }

  @Override
  protected void onContainerCompleted(ContainerStatus status) {
    super.onContainerCompleted(status);

    boolean removed = grid.removeMember(status.getContainerId());
    if (!removed) {
      // force allocation for all clusters if we got completed
      // container unknown to a grid. might be an indication
      // that satisfy states are not met.
      requestAllocationForAll();
    }

    if (getMonitor() instanceof ContainerAware) {
      ((ContainerAware)getMonitor()).onContainerStatus(Arrays.asList(status));
    }
  }

  @Override
  public Map<String, ContainerCluster> getContainerClusters() {
    return clusters;
  }

  @Override
  public ContainerCluster createContainerCluster(String clusterId, ProjectionData projectionData) {
    return createContainerCluster(clusterId, clusterId, projectionData, null);
  }

  @Override
  public ContainerCluster createContainerCluster(String clusterId, String clusterDef, ProjectionData projectionData, Map<String, Object> extraProperties) {

    if (clusterDef == null) {
      clusterDef = clusterId;
    }

    GridProjectionFactory gridProjectionFactory = gridProjectionFactoryLocator.getGridProjectionFactory(projectionData.getType());
    if (gridProjectionFactory == null) {
      throw new IllegalArgumentException("Projection type " + projectionData.getType()
          + " not know to gridProjectionFactoryLocator=[" + gridProjectionFactoryLocator + "]");
    }

    GridProjection projection = gridProjectionFactory.getGridProjection(projectionData);
    if (projection == null) {
      throw new IllegalArgumentException("Unable to build projection using type " + projectionData.getType());
    }

    ProjectionData p = projectionDataRegistry.getProjectionDatas().get(clusterDef);

    if (projection instanceof AbstractGridProjection) {
      ((AbstractGridProjection) projection).setConfiguration(getConfiguration());
      if (p != null && ((AbstractGridProjection) projection).getPriority() == null) {
        ((AbstractGridProjection) projection).setPriority(p.getPriority());
      }
    }

    StateMachine<State<ClusterState, ClusterEvent>, ClusterEvent> stateMachine = stateMachineFactory
        .getStateMachine();
    DefaultContainerCluster cluster = new DefaultContainerCluster(clusterId, projection, stateMachine, extraProperties);
    clusters.put(cluster.getId(), cluster);
    projectedGrid.addProjection(cluster.getGridProjection());
    clusterIdToRef.put(clusterId, clusterDef);
    return cluster;
  }

  @Override
  public void startContainerCluster(String id) {
    ContainerCluster cluster = clusters.get(id);
    if (cluster != null) {
      StateMachine<State<ClusterState, ClusterEvent>, ClusterEvent> stateMachine = cluster.getStateMachine();
      stateMachine.sendEvent(ClusterEvent.START);
      stateMachine.sendEvent(MessageBuilder.withPayload(ClusterEvent.CONFIGURE)
          .setHeader("containercluster", cluster).setHeader("appmaster", AbstractContainerClusterAppmaster.this).build());
    }
  }

  @Override
  public void stopContainerCluster(String id) {
    ContainerCluster cluster = clusters.get(id);
    if (cluster != null) {
      StateMachine<State<ClusterState, ClusterEvent>, ClusterEvent> stateMachine = cluster.getStateMachine();
      stateMachine.sendEvent(MessageBuilder.withPayload(ClusterEvent.STOP).setHeader("containercluster", cluster)
          .setHeader("appmaster", AbstractContainerClusterAppmaster.this).build());
    }
  }

  @Override
  public void destroyContainerCluster(String id) {
    ContainerCluster cluster = clusters.get(id);
    if (cluster != null) {
      StateMachine<State<ClusterState, ClusterEvent>, ClusterEvent> stateMachine = cluster.getStateMachine();
      stateMachine.sendEvent(MessageBuilder.withPayload(ClusterEvent.DESTROY)
          .setHeader("containercluster", cluster)
          .setHeader("appmaster", AbstractContainerClusterAppmaster.this).build());
    }
  }

  @Override
  public void modifyContainerCluster(String id, ProjectionData data) {
    ContainerCluster cluster = clusters.get(id);
    if (cluster != null) {
      StateMachine<State<ClusterState, ClusterEvent>, ClusterEvent> stateMachine = cluster.getStateMachine();
      stateMachine.sendEvent(MessageBuilder.withPayload(ClusterEvent.CONFIGURE).setHeader("projectiondata", data)
          .setHeader("containercluster", cluster).setHeader("appmaster", AbstractContainerClusterAppmaster.this).build());
    }
  }

  @Autowired
  public void setStateMachineFactory(EnumStateMachineFactory<ClusterState, ClusterEvent> stateMachineFactory) {
    this.stateMachineFactory = stateMachineFactory;
  }

  /**
   * Sets the {@link GridProjectionFactoryLocator} used to find factories
   * which are creating an instances of {@link GridProjection}s.
   *
   * @param gridProjectionFactoryLocator the grid projection factory locator
   */
  @Autowired
  public void setGridProjectionFactoryLocator(GridProjectionFactoryLocator gridProjectionFactoryLocator) {
    log.info("Setting gridProjectionFactoryLocator=" + gridProjectionFactoryLocator);
    this.gridProjectionFactoryLocator = gridProjectionFactoryLocator;
  }

  @Autowired(required = false)
  public void setProjectionDataRegistry(ProjectionDataRegistry projectionDataRegistry) {
    log.info("Setting projectionDataRegistry=" + projectionDataRegistry);
    this.projectionDataRegistry = projectionDataRegistry;
  }

  protected Grid doCreateGrid() {
    return new DefaultGrid();
  }

  protected ProjectedGrid doCreateProjectedGrid(Grid grid) {
    return new DefaultProjectedGrid(grid);
  }

  protected void handleSatisfyStateData(ContainerCluster cluster, SatisfyStateData satisfyData) {
    if (satisfyData.getAllocateData() != null) {
      // who set the allocator settings matching id from allocate data???
      // if not configures, fall back to defaults???
      ContainerAllocateData allocateData = satisfyData.getAllocateData();
      // should not set id here!!!
      allocateData.setId(clusterIdToRef.get(cluster.getId()));
      getAllocator().allocateContainers(allocateData);
    }
    for (GridMember member : satisfyData.getRemoveData()) {
      log.info("Queued container to be killed: " + member.getContainer().getId());
      killContainer(member.getContainer());
    }
  }

  protected void killContainer(Container container) {
    killQueue.add(container);
  }

  protected AppmasterCmOperations getCmTemplate(Container container) {
    try {
      AppmasterCmTemplate template = new AppmasterCmTemplate(getConfiguration(), container);
      template.afterPropertiesSet();
      return template;
    } catch (Exception e) {
      throw new YarnSystemException("Unable to create AppmasterCmTemplate", e);
    }
  }

  /**
   * Called when a container is launched for sub classes to do
   * a final modifications in these commands. Default implementation
   * returns commands as is.
   *
   * @param cluster the cluster id
   * @param commands the original commands
   * @return modified commands
   */
  protected List<String> onContainerLaunchCommands(Container container, ContainerCluster cluster, List<String> commands) {
    return commands;
  }

  /**
   * Periodic task callback called by a {@link ClusterTaskPoller}.
   */
  private void doTask() {
    // kill containers not needed
    handleKillQueue();
  }

  /**
   * Kill all containers from queue.
   */
  private void handleKillQueue() {
    Container toKill = null;
    while ((toKill = killQueue.poll()) != null) {
      log.info("Killing container: " + toKill);
      getCmTemplate(toKill).stopContainers();
    }
  }

  private ContainerCluster findContainerClusterByContainer(Container container) {
    // TODO: make finding cluster more clever
    for (Entry<String, ContainerCluster> entry : clusters.entrySet()) {
      for (GridMember member : entry.getValue().getGridProjection().getMembers()) {
        if (member.getContainer().equals(container)) {
          return entry.getValue();
        }
      }
    }
    return null;
  }

  private void requestAllocationForAll() {
    for (ContainerCluster cluster : clusters.values()) {
      StateMachine<State<ClusterState, ClusterEvent>, ClusterEvent> stateMachine = cluster.getStateMachine();
      stateMachine.sendEvent(MessageBuilder.withPayload(ClusterEvent.CONFIGURE)
          .setHeader("containercluster", cluster)
          .setHeader("appmaster", AbstractContainerClusterAppmaster.this).build());
    }

  }

  /**
   * Internal poller handling cluster allocation request scheduling.
   */
  private class ClusterTaskPoller extends PollingTaskSupport<Void> {

    public ClusterTaskPoller(TaskScheduler taskScheduler, TaskExecutor taskExecutor) {
      super(taskScheduler, taskExecutor);
    }

    @Override
    protected Void doPoll() {
      doTask();
      return null;
    }

  }

  private class CommandDispatchListener extends ProjectedGridListenerAdapter {

    @Override
    public void memberRemoved(GridProjection projection, GridMember member) {
      log.info("memberRemoved projection=" + projection + " member=" + member);

      for (ContainerCluster cluster : clusters.values()) {
        if (cluster.getGridProjection().equals(projection)) {
          StateMachine<State<ClusterState, ClusterEvent>, ClusterEvent> stateMachine = cluster.getStateMachine();
          stateMachine.sendEvent(MessageBuilder.withPayload(ClusterEvent.CONFIGURE)
              .setHeader("containercluster", cluster).setHeader("appmaster", AbstractContainerClusterAppmaster.this).build());
        }
      }

    }
  }

  private class ContainerLaunchContextModifyInterceptor implements ContainerLauncherInterceptor {

    @Override
    public ContainerLaunchContext preLaunch(Container container, ContainerLaunchContext context) {
      ContainerCluster cluster = findContainerClusterByContainer(container);
      if (getResourceLocalizer() instanceof MultiResourceLocalizer) {
        if (cluster != null) {
          MultiResourceLocalizer loc = (MultiResourceLocalizer) getResourceLocalizer();
          Map<String, LocalResource> resources = loc.getResources(clusterIdToRef.get(cluster.getId()));
          context.setLocalResources(resources);
        }
      } else {
        log.warn("Can't use container specific local resources because MultiResourceLocalizer expected instead of "
            + getResourceLocalizer());
      }

      if (cluster != null) {
        Map<String, String> environment = new HashMap<String, String>(context.getEnvironment());
        environment.putAll(getEnvironment(clusterIdToRef.get(cluster.getId())));
        context.setEnvironment(environment);
      }

      return context;
    }
  }

}
TOP

Related Classes of org.springframework.yarn.am.cluster.AbstractContainerClusterAppmaster$ContainerLaunchContextModifyInterceptor

TOP
Copyright © 2018 www.massapi.com. All rights reserved.
All source code are property of their respective owners. Java is a trademark of Sun Microsystems, Inc and owned by ORACLE Inc. Contact coftware#gmail.com.