yarn/src/main/java/org/apache/hama/bsp/JobImpl.java - hama - Git at Google

 /**
  * Licensed to the Apache Software Foundation (ASF) under one
  * or more contributor license agreements.  See the NOTICE file
  * regarding copyright ownership.  The ASF licenses this file
  * distributed with this work for additional information
  * 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.hama.bsp;

 import java.io.IOException;
 import java.net.InetSocketAddress;
 import java.security.PrivilegedAction;
 import java.util.*;

 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
 import org.apache.hadoop.conf.Configuration;
 import org.apache.hadoop.fs.Path;
 import org.apache.hadoop.net.NetUtils;
 import org.apache.hadoop.security.UserGroupInformation;
 import org.apache.hadoop.yarn.api.ApplicationConstants;
 import org.apache.hadoop.yarn.api.ApplicationMasterProtocol;
 import org.apache.hadoop.yarn.api.ContainerManagementProtocol;
 import org.apache.hadoop.yarn.api.protocolrecords.AllocateRequest;
 import org.apache.hadoop.yarn.api.protocolrecords.AllocateResponse;
 import org.apache.hadoop.yarn.api.records.*;
 import org.apache.hadoop.yarn.api.records.Token;
 import org.apache.hadoop.yarn.client.api.NMTokenCache;
 import org.apache.hadoop.yarn.exceptions.YarnException;
 import org.apache.hadoop.yarn.ipc.YarnRPC;
 import org.apache.hadoop.yarn.util.ConverterUtils;
 import org.apache.hadoop.yarn.util.Records;
 import org.apache.hama.bsp.BSPTaskLauncher.BSPTaskStatus;

 public class JobImpl implements Job {

   private static final Log LOG = LogFactory.getLog(JobImpl.class);
   private static final int DEFAULT_MEMORY_MB = 256;

   private Configuration conf;
   private BSPJobID jobId;
   private int numBSPTasks;
   private int priority = 0;
   private String childOpts;
   private int taskMemoryInMb;
   private Path jobFile;

   private JobState state;
   private BSPPhase phase;

   private ApplicationAttemptId appAttemptId;
   private YarnRPC yarnRPC;
   private ApplicationMasterProtocol resourceManager;

   private List<Container> allocatedContainers;
   private List<ContainerId> releasedContainers = Collections.emptyList();

   private Map<Integer, BSPTaskLauncher> launchers = new HashMap<Integer, BSPTaskLauncher>();
   private Deque<BSPTaskLauncher> completionQueue = new LinkedList<BSPTaskLauncher>();

   private int lastResponseID = 0;

   private int getMemoryRequirements() {
     String newMemoryProperty = conf.get("bsp.child.mem.in.mb");
     if (newMemoryProperty == null) {
       LOG.warn("\"bsp.child.mem.in.mb\" was not set! Try parsing the child opts...");
       return getMemoryFromOptString(childOpts);
     } else {
       return Integer.valueOf(newMemoryProperty);
     }
   }

   public JobImpl(ApplicationAttemptId appAttemptId,
                  Configuration jobConfiguration, YarnRPC yarnRPC, ApplicationMasterProtocol amrmRPC,
                  String jobFile, BSPJobID jobId) {
     super();
     this.appAttemptId = appAttemptId;
     this.yarnRPC = yarnRPC;
     this.resourceManager = amrmRPC;
     this.jobFile = new Path(jobFile);
     this.state = JobState.NEW;
     this.jobId = jobId;
     this.conf = jobConfiguration;
     this.numBSPTasks = conf.getInt("bsp.peers.num", 1);
     this.childOpts = conf.get("bsp.child.java.opts");

     this.taskMemoryInMb = getMemoryRequirements();
   }

   // This really needs a testcase
   private static int getMemoryFromOptString(String opts) {
     if (opts == null) {
       return DEFAULT_MEMORY_MB;
     }

     if (!opts.contains("-Xmx")) {
       LOG.info("No \"-Xmx\" option found in child opts, using default amount of memory!");
       return DEFAULT_MEMORY_MB;
     } else {
       // e.G: -Xmx512m

       int startIndex = opts.indexOf("-Xmx") + 4;
       String xmxString = opts.substring(startIndex);
       char qualifier = xmxString.charAt(xmxString.length() - 1);
       int memory = Integer.valueOf(xmxString.substring(0,
           xmxString.length() - 1));
       if (qualifier == 'm') {
         return memory;
       } else if (qualifier == 'g') {
         return memory * 1024;
       } else {
         throw new IllegalArgumentException(
             "Memory Limit in child opts was not set! \"bsp.child.java.opts\" String was: "
                 + opts);
       }
     }
   }

   @Override
   public JobState startJob() throws Exception {

     this.allocatedContainers = new ArrayList<Container>(numBSPTasks);
     NMTokenCache nmTokenCache = new NMTokenCache();
     while (allocatedContainers.size() < numBSPTasks) {
       AllocateRequest req = AllocateRequest.newInstance(lastResponseID, 0.0f,
           createBSPTaskRequest(numBSPTasks - allocatedContainers.size(), taskMemoryInMb,
               priority), releasedContainers, null);

       AllocateResponse allocateResponse = resourceManager.allocate(req);
       for (NMToken token : allocateResponse.getNMTokens()) {
         nmTokenCache.setToken(token.getNodeId().toString(), token.getToken());
       }

       LOG.info("Got response ID: " + allocateResponse.getResponseId()
           + " with num of containers: "
           + allocateResponse.getAllocatedContainers().size()
           + " and following resources: "
           + allocateResponse.getAvailableResources().getMemory() + "mb");
       this.lastResponseID = allocateResponse.getResponseId();

       this.allocatedContainers.addAll(allocateResponse.getAllocatedContainers());

       LOG.info("Waiting to allocate " + (numBSPTasks - allocatedContainers.size()) + " more containers...");

       Thread.sleep(1000l);
     }

     LOG.info("Got " + allocatedContainers.size() + " containers!");

     int id = 0;
     for (Container allocatedContainer : allocatedContainers) {
       LOG.info("Launching task on a new container." + ", containerId="
           + allocatedContainer.getId() + ", containerNode="
           + allocatedContainer.getNodeId().getHost() + ":"
           + allocatedContainer.getNodeId().getPort() + ", containerNodeURI="
           + allocatedContainer.getNodeHttpAddress() + ", containerResourceMemory"
           + allocatedContainer.getResource().getMemory());

       // Connect to ContainerManager on the allocated container
       String user = conf.get("bsp.user.name");
       if (user == null) {
         user = System.getenv(ApplicationConstants.Environment.USER.name());
       }

       ContainerManagementProtocol cm = null;
       try {
         cm = getContainerManagementProtocolProxy(yarnRPC,
             nmTokenCache.getToken(allocatedContainer.getNodeId().toString()), allocatedContainer.getNodeId(), user);
       } catch (Exception e) {
         LOG.error("Failed to create ContainerManager...");
         if (cm != null)
           yarnRPC.stopProxy(cm, conf);
         e.printStackTrace();
       }

       BSPTaskLauncher runnableLaunchContainer = new BSPTaskLauncher(id,
           allocatedContainer, cm, conf, jobFile, jobId);

       launchers.put(id, runnableLaunchContainer);
       runnableLaunchContainer.start();
       completionQueue.add(runnableLaunchContainer);
       id++;
     }

     LOG.info("Waiting for tasks to finish...");
     state = JobState.RUNNING;
     int completed = 0;

     while (completed != numBSPTasks) {
       for (BSPTaskLauncher task : completionQueue) {
         BSPTaskStatus returnedTask = task.poll();
         if(returnedTask != null && returnedTask.getExitStatus() == 0) {
           LOG.info("Task \"" + returnedTask.getId()
               + "\" sucessfully finished!");
           completed++;
           LOG.info("Waiting for " + (numBSPTasks - completed)
               + " tasks to finish!");
         }

         if(returnedTask != null && returnedTask.getExitStatus() != 0) {
           LOG.error("Task with id \"" + returnedTask.getId() + "\" failed!");
           completionQueue.add(task);
           state = JobState.FAILED;
           return state;
         }
       }
       Thread.sleep(1000L);
     }

     state = JobState.SUCCESS;
     return state;
   }

   /**
    *
    * @param rpc
    * @param nmToken
    * @param nodeId
    * @param user
    * @return
    */
   protected ContainerManagementProtocol getContainerManagementProtocolProxy(
       final YarnRPC rpc, Token nmToken, NodeId nodeId, String user) {
     ContainerManagementProtocol proxy;
     UserGroupInformation ugi = UserGroupInformation.createRemoteUser(user);
     final InetSocketAddress addr =
         NetUtils.createSocketAddr(nodeId.getHost(), nodeId.getPort());
     if (nmToken != null) {
       ugi.addToken(ConverterUtils.convertFromYarn(nmToken, addr));
     }

     proxy = ugi
         .doAs(new PrivilegedAction<ContainerManagementProtocol>() {
           @Override
           public ContainerManagementProtocol run() {
             return (ContainerManagementProtocol) rpc.getProxy(
                 ContainerManagementProtocol.class,
                 addr, conf);
           }
         });
     return proxy;
   }

   /**
    * Makes a lookup for the taskid and stops its container and task. It also
    * removes the task from the launcher so that we won't have to stop it twice.
    *
    * @param id
    * @throws YarnException
    */
   private void cleanupTask(int id) throws YarnException, IOException {
     BSPTaskLauncher bspTaskLauncher = launchers.get(id);
     bspTaskLauncher.stopAndCleanup();
     launchers.remove(id);
     completionQueue.remove(bspTaskLauncher);
   }

   private List<ResourceRequest> createBSPTaskRequest(int numTasks,
       int memoryInMb, int priority) {

     List<ResourceRequest> reqList = new ArrayList<ResourceRequest>(numTasks);
     for (int i = 0; i < numTasks; i++) {
       // Resource Request
       ResourceRequest rsrcRequest = Records.newRecord(ResourceRequest.class);

       // setup requirements for hosts
       // whether a particular rack/host is needed
       // useful for applications that are sensitive
       // to data locality
       rsrcRequest.setResourceName("*");

       // set the priority for the request
       Priority pri = Records.newRecord(Priority.class);
       pri.setPriority(priority);
       rsrcRequest.setPriority(pri);

       // Set up resource type requirements
       // For now, only memory is supported so we set memory requirements
       Resource capability = Records.newRecord(Resource.class);
       capability.setMemory(memoryInMb);
       rsrcRequest.setCapability(capability);

       // set no. of containers needed
       // matching the specifications
       rsrcRequest.setNumContainers(numBSPTasks);
       reqList.add(rsrcRequest);
     }
     return reqList;
   }

   @Override
   public void cleanup() throws YarnException, IOException {
     for (BSPTaskLauncher launcher : completionQueue) {
       LOG.info("cleanup tasks !!!");
       launcher.stopAndCleanup();
     }
   }

   @Override
   public JobState getState() {
     return state;
   }

   @Override
   public int getTotalBSPTasks() {
     return numBSPTasks;
   }

   @Override
   public BSPPhase getBSPPhase() {
     return phase;
   }

 }
	/**
	* Licensed to the Apache Software Foundation (ASF) under one
	* or more contributor license agreements. See the NOTICE file
	* regarding copyright ownership. The ASF licenses this file
	* distributed with this work for additional information
	* 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.hama.bsp;

	import java.io.IOException;
	import java.net.InetSocketAddress;
	import java.security.PrivilegedAction;
	import java.util.*;

	import org.apache.commons.logging.Log;
	import org.apache.commons.logging.LogFactory;
	import org.apache.hadoop.conf.Configuration;
	import org.apache.hadoop.fs.Path;
	import org.apache.hadoop.net.NetUtils;
	import org.apache.hadoop.security.UserGroupInformation;
	import org.apache.hadoop.yarn.api.ApplicationConstants;
	import org.apache.hadoop.yarn.api.ApplicationMasterProtocol;
	import org.apache.hadoop.yarn.api.ContainerManagementProtocol;
	import org.apache.hadoop.yarn.api.protocolrecords.AllocateRequest;
	import org.apache.hadoop.yarn.api.protocolrecords.AllocateResponse;
	import org.apache.hadoop.yarn.api.records.*;
	import org.apache.hadoop.yarn.api.records.Token;
	import org.apache.hadoop.yarn.client.api.NMTokenCache;
	import org.apache.hadoop.yarn.exceptions.YarnException;
	import org.apache.hadoop.yarn.ipc.YarnRPC;
	import org.apache.hadoop.yarn.util.ConverterUtils;
	import org.apache.hadoop.yarn.util.Records;
	import org.apache.hama.bsp.BSPTaskLauncher.BSPTaskStatus;

	public class JobImpl implements Job {

	private static final Log LOG = LogFactory.getLog(JobImpl.class);
	private static final int DEFAULT_MEMORY_MB = 256;

	private Configuration conf;
	private BSPJobID jobId;
	private int numBSPTasks;
	private int priority = 0;
	private String childOpts;
	private int taskMemoryInMb;
	private Path jobFile;

	private JobState state;
	private BSPPhase phase;

	private ApplicationAttemptId appAttemptId;
	private YarnRPC yarnRPC;
	private ApplicationMasterProtocol resourceManager;

	private List<Container> allocatedContainers;
	private List<ContainerId> releasedContainers = Collections.emptyList();

	private Map<Integer, BSPTaskLauncher> launchers = new HashMap<Integer, BSPTaskLauncher>();
	private Deque<BSPTaskLauncher> completionQueue = new LinkedList<BSPTaskLauncher>();

	private int lastResponseID = 0;

	private int getMemoryRequirements() {
	String newMemoryProperty = conf.get("bsp.child.mem.in.mb");
	if (newMemoryProperty == null) {
	LOG.warn("\"bsp.child.mem.in.mb\" was not set! Try parsing the child opts...");
	return getMemoryFromOptString(childOpts);
	} else {
	return Integer.valueOf(newMemoryProperty);
	}
	}

	public JobImpl(ApplicationAttemptId appAttemptId,
	Configuration jobConfiguration, YarnRPC yarnRPC, ApplicationMasterProtocol amrmRPC,
	String jobFile, BSPJobID jobId) {
	super();
	this.appAttemptId = appAttemptId;
	this.yarnRPC = yarnRPC;
	this.resourceManager = amrmRPC;
	this.jobFile = new Path(jobFile);
	this.state = JobState.NEW;
	this.jobId = jobId;
	this.conf = jobConfiguration;
	this.numBSPTasks = conf.getInt("bsp.peers.num", 1);
	this.childOpts = conf.get("bsp.child.java.opts");

	this.taskMemoryInMb = getMemoryRequirements();
	}

	// This really needs a testcase
	private static int getMemoryFromOptString(String opts) {
	if (opts == null) {
	return DEFAULT_MEMORY_MB;
	}

	if (!opts.contains("-Xmx")) {
	LOG.info("No \"-Xmx\" option found in child opts, using default amount of memory!");
	return DEFAULT_MEMORY_MB;
	} else {
	// e.G: -Xmx512m

	int startIndex = opts.indexOf("-Xmx") + 4;
	String xmxString = opts.substring(startIndex);
	char qualifier = xmxString.charAt(xmxString.length() - 1);
	int memory = Integer.valueOf(xmxString.substring(0,
	xmxString.length() - 1));
	if (qualifier == 'm') {
	return memory;
	} else if (qualifier == 'g') {
	return memory * 1024;
	} else {
	throw new IllegalArgumentException(
	"Memory Limit in child opts was not set! \"bsp.child.java.opts\" String was: "
	+ opts);
	}
	}
	}

	@Override
	public JobState startJob() throws Exception {

	this.allocatedContainers = new ArrayList<Container>(numBSPTasks);
	NMTokenCache nmTokenCache = new NMTokenCache();
	while (allocatedContainers.size() < numBSPTasks) {
	AllocateRequest req = AllocateRequest.newInstance(lastResponseID, 0.0f,
	createBSPTaskRequest(numBSPTasks - allocatedContainers.size(), taskMemoryInMb,
	priority), releasedContainers, null);

	AllocateResponse allocateResponse = resourceManager.allocate(req);
	for (NMToken token : allocateResponse.getNMTokens()) {
	nmTokenCache.setToken(token.getNodeId().toString(), token.getToken());
	}

	LOG.info("Got response ID: " + allocateResponse.getResponseId()
	+ " with num of containers: "
	+ allocateResponse.getAllocatedContainers().size()
	+ " and following resources: "
	+ allocateResponse.getAvailableResources().getMemory() + "mb");
	this.lastResponseID = allocateResponse.getResponseId();

	this.allocatedContainers.addAll(allocateResponse.getAllocatedContainers());

	LOG.info("Waiting to allocate " + (numBSPTasks - allocatedContainers.size()) + " more containers...");

	Thread.sleep(1000l);
	}

	LOG.info("Got " + allocatedContainers.size() + " containers!");

	int id = 0;
	for (Container allocatedContainer : allocatedContainers) {
	LOG.info("Launching task on a new container." + ", containerId="
	+ allocatedContainer.getId() + ", containerNode="
	+ allocatedContainer.getNodeId().getHost() + ":"
	+ allocatedContainer.getNodeId().getPort() + ", containerNodeURI="
	+ allocatedContainer.getNodeHttpAddress() + ", containerResourceMemory"
	+ allocatedContainer.getResource().getMemory());

	// Connect to ContainerManager on the allocated container
	String user = conf.get("bsp.user.name");
	if (user == null) {
	user = System.getenv(ApplicationConstants.Environment.USER.name());
	}

	ContainerManagementProtocol cm = null;
	try {
	cm = getContainerManagementProtocolProxy(yarnRPC,
	nmTokenCache.getToken(allocatedContainer.getNodeId().toString()), allocatedContainer.getNodeId(), user);
	} catch (Exception e) {
	LOG.error("Failed to create ContainerManager...");
	if (cm != null)
	yarnRPC.stopProxy(cm, conf);
	e.printStackTrace();
	}

	BSPTaskLauncher runnableLaunchContainer = new BSPTaskLauncher(id,
	allocatedContainer, cm, conf, jobFile, jobId);

	launchers.put(id, runnableLaunchContainer);
	runnableLaunchContainer.start();
	completionQueue.add(runnableLaunchContainer);
	id++;
	}

	LOG.info("Waiting for tasks to finish...");
	state = JobState.RUNNING;
	int completed = 0;

	while (completed != numBSPTasks) {
	for (BSPTaskLauncher task : completionQueue) {
	BSPTaskStatus returnedTask = task.poll();
	if(returnedTask != null && returnedTask.getExitStatus() == 0) {
	LOG.info("Task \"" + returnedTask.getId()
	+ "\" sucessfully finished!");
	completed++;
	LOG.info("Waiting for " + (numBSPTasks - completed)
	+ " tasks to finish!");
	}

	if(returnedTask != null && returnedTask.getExitStatus() != 0) {
	LOG.error("Task with id \"" + returnedTask.getId() + "\" failed!");
	completionQueue.add(task);
	state = JobState.FAILED;
	return state;
	}
	}
	Thread.sleep(1000L);
	}

	state = JobState.SUCCESS;
	return state;
	}

	/**
	*
	* @param rpc
	* @param nmToken
	* @param nodeId
	* @param user
	* @return
	*/
	protected ContainerManagementProtocol getContainerManagementProtocolProxy(
	final YarnRPC rpc, Token nmToken, NodeId nodeId, String user) {
	ContainerManagementProtocol proxy;
	UserGroupInformation ugi = UserGroupInformation.createRemoteUser(user);
	final InetSocketAddress addr =
	NetUtils.createSocketAddr(nodeId.getHost(), nodeId.getPort());
	if (nmToken != null) {
	ugi.addToken(ConverterUtils.convertFromYarn(nmToken, addr));
	}

	proxy = ugi
	.doAs(new PrivilegedAction<ContainerManagementProtocol>() {
	@Override
	public ContainerManagementProtocol run() {
	return (ContainerManagementProtocol) rpc.getProxy(
	ContainerManagementProtocol.class,
	addr, conf);
	}
	});
	return proxy;
	}

	/**
	* Makes a lookup for the taskid and stops its container and task. It also
	* removes the task from the launcher so that we won't have to stop it twice.
	*
	* @param id
	* @throws YarnException
	*/
	private void cleanupTask(int id) throws YarnException, IOException {
	BSPTaskLauncher bspTaskLauncher = launchers.get(id);
	bspTaskLauncher.stopAndCleanup();
	launchers.remove(id);
	completionQueue.remove(bspTaskLauncher);
	}

	private List<ResourceRequest> createBSPTaskRequest(int numTasks,
	int memoryInMb, int priority) {

	List<ResourceRequest> reqList = new ArrayList<ResourceRequest>(numTasks);
	for (int i = 0; i < numTasks; i++) {
	// Resource Request
	ResourceRequest rsrcRequest = Records.newRecord(ResourceRequest.class);

	// setup requirements for hosts
	// whether a particular rack/host is needed
	// useful for applications that are sensitive
	// to data locality
	rsrcRequest.setResourceName("*");

	// set the priority for the request
	Priority pri = Records.newRecord(Priority.class);
	pri.setPriority(priority);
	rsrcRequest.setPriority(pri);

	// Set up resource type requirements
	// For now, only memory is supported so we set memory requirements
	Resource capability = Records.newRecord(Resource.class);
	capability.setMemory(memoryInMb);
	rsrcRequest.setCapability(capability);

	// set no. of containers needed
	// matching the specifications
	rsrcRequest.setNumContainers(numBSPTasks);
	reqList.add(rsrcRequest);
	}
	return reqList;
	}

	@Override
	public void cleanup() throws YarnException, IOException {
	for (BSPTaskLauncher launcher : completionQueue) {
	LOG.info("cleanup tasks !!!");
	launcher.stopAndCleanup();
	}
	}

	@Override
	public JobState getState() {
	return state;
	}

	@Override
	public int getTotalBSPTasks() {
	return numBSPTasks;
	}

	@Override
	public BSPPhase getBSPPhase() {
	return phase;
	}

	}