src/main/java/org/apache/sysml/yarn/DMLYarnClient.java - systemds - Git at Google

 /*
  * 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
  *
  *   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.sysml.yarn;

 import java.io.BufferedReader;
 import java.io.File;
 import java.io.IOException;
 import java.io.InputStreamReader;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.Map;

 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
 import org.apache.hadoop.fs.FSDataInputStream;
 import org.apache.hadoop.fs.FSDataOutputStream;
 import org.apache.hadoop.fs.FileStatus;
 import org.apache.hadoop.fs.FileSystem;
 import org.apache.hadoop.fs.FileUtil;
 import org.apache.hadoop.fs.Path;
 import org.apache.hadoop.mapreduce.v2.util.MRApps;
 import org.apache.hadoop.yarn.api.ApplicationConstants;
 import org.apache.hadoop.yarn.api.ApplicationConstants.Environment;
 import org.apache.hadoop.yarn.api.records.ApplicationId;
 import org.apache.hadoop.yarn.api.records.ApplicationReport;
 import org.apache.hadoop.yarn.api.records.ApplicationSubmissionContext;
 import org.apache.hadoop.yarn.api.records.ContainerLaunchContext;
 import org.apache.hadoop.yarn.api.records.FinalApplicationStatus;
 import org.apache.hadoop.yarn.api.records.LocalResource;
 import org.apache.hadoop.yarn.api.records.LocalResourceType;
 import org.apache.hadoop.yarn.api.records.LocalResourceVisibility;
 import org.apache.hadoop.yarn.api.records.Resource;
 import org.apache.hadoop.yarn.api.records.YarnApplicationState;
 import org.apache.hadoop.yarn.client.api.YarnClient;
 import org.apache.hadoop.yarn.client.api.YarnClientApplication;
 import org.apache.hadoop.yarn.conf.YarnConfiguration;
 import org.apache.hadoop.yarn.util.ConverterUtils;
 import org.apache.hadoop.yarn.util.Records;

 import org.apache.sysml.conf.DMLConfig;
 import org.apache.sysml.parser.ParseException;
 import org.apache.sysml.runtime.DMLRuntimeException;
 import org.apache.sysml.runtime.DMLScriptException;
 import org.apache.sysml.runtime.controlprogram.parfor.stat.Timing;
 import org.apache.sysml.runtime.util.MapReduceTool;

 /**
  * NOTES:
  *   * Security: By default, submitted applications are ran as user 'yarn'.
  *     In order to allow for security and relative filenames on hdfs (/user/<username>/.),
  *     we can configure the LinuxContainerExecutor in yarn-site.xml, which runs the
  *     application as the user who submits the application.
  *   * SystemML.jar file dependency: We need to submit the SystemML.jar along with the
  *     application. Unfortunately, hadoop jar unpacks the jar such that we dont get a handle
  *     to the original jar filename. We currently parse the constant IBM_JAVA_COMMAND_LINE
  *     to get the jar filename. For robustness, we fall back to repackaging the unzipped files
  *     to a jar if this constant does not exist.
  *
  */
 public class DMLYarnClient
 {

 	private static final Log LOG = LogFactory.getLog(DMLYarnClient.class);

 	//Internal configuration parameters
 	// environment variable to obtain the original jar filename
 	public static final String JARFILE_ENV_CONST = "IBM_JAVA_COMMAND_LINE";
 	// environment variable to obtain default jvm arguments
 	public static final String JVMOPTS_ENV_CONST = "HADOOP_OPTS";
 	// environment variable to obtain mapred home (for robustness only)
 	public static final String MAPRED_HOME_ENV_CONST = "HADOOP_MAPRED_HOME";
 	public static final String HADOOP_HOME_ENV_CONST = "HADOOP_HOME";
 	// default of 1 core since YARN scheduler does not take the number of cores into account yet
 	public static final int NUM_CORES = 1;
 	// factor for compute virtual memory to request based on given max heap size
 	// (if the absolute overhead is larger than a threshold, we use this threshold as a constant overhead)
 	public static final double MEM_FACTOR = 1.5;
 	public static final int MAX_MEM_OVERHEAD = 2*1024; //2GB
 	// default application state report (in milliseconds)
 	public static final int APP_STATE_INTERVAL = 200;
 	// default application master name
 	public static final String APPMASTER_NAME = "SystemML-AM";
 	// default dml script file name for hdfs script serialization
 	public static final String DML_SCRIPT_NAME = "script.dml";
 	// default dml config file name for hdfs config serialization
 	public static final String DML_CONFIG_NAME = "config.xml";
 	// default SystemML jar file name for hdfs jar copy
 	public static final String DML_JAR_NAME = "SystemML.jar";
 	// default dml stop message file name for hdfs message serialization
 	public static final String DML_STOPMSG_NAME = "stop_msg.txt";


 	private String _dmlScript = null;
 	private DMLConfig _dmlConfig = null;
 	private String[] _args = null;

 	//hdfs file names local resources
 	private String _hdfsJarFile   = null;
 	private String _hdfsDMLScript = null;
 	private String _hdfsDMLConfig = null;

 	/**
 	 * Protected since only supposed to be accessed via proxy in same package.
 	 * This is to ensure robustness in case of missing yarn libraries.
 	 *
 	 * @param dmlScriptStr
 	 * @param conf
 	 * @param args
 	 */
 	protected DMLYarnClient(String dmlScriptStr, DMLConfig conf, String[] args)
 	{
 		_dmlScript = dmlScriptStr;
 		_dmlConfig = conf;
 		_args = args;
 	}


 	/**
 	 * Method to launch the dml yarn app master and execute the given dml script
 	 * with the given configuration and jar file.
 	 *
 	 * NOTE: on launching the yarn app master, we do not explicitly probe if we
 	 *	  are running on a yarn or MR1 cluster. In case of MR1, already the class
 	 *	  YarnConfiguration will not be found and raise a classnotfound. In case of any
 	 *	  exception we fall back to run CP directly in the client process.
 	 *
 	 * @return true if dml program successfully executed as yarn app master
 	 * @throws IOException
 	 */
 	protected boolean launchDMLYarnAppmaster()
 		throws IOException, DMLScriptException
 	{
 		boolean ret = false;
 		String hdfsWD = null;

 		try
 		{
 			Timing time = new Timing(true);

 			// load yarn configuration
 			YarnConfiguration yconf = new YarnConfiguration();

 			// create yarn client
 			YarnClient yarnClient = YarnClient.createYarnClient();
 			yarnClient.init(yconf);
 			yarnClient.start();

 			// create application and get the ApplicationID
 			YarnClientApplication app = yarnClient.createApplication();
 			ApplicationSubmissionContext appContext = app.getApplicationSubmissionContext();
 			ApplicationId appId = appContext.getApplicationId();
 			LOG.debug("Created application (applicationID: "+appId+")");

 			// prepare hdfs working directory via ApplicationID
 			// copy script, config, jar file to hdfs
 			hdfsWD = DMLAppMasterUtils.constructHDFSWorkingDir(_dmlConfig, appId);
 			copyResourcesToHdfsWorkingDir(yconf, hdfsWD);

 			//construct command line argument
 			String command = constructAMCommand(_args, _dmlConfig);
 			LOG.debug("Constructed application master command: \n"+command);

 			// set up the container launch context for the application master
 			ContainerLaunchContext amContainer = Records.newRecord(ContainerLaunchContext.class);
 			amContainer.setCommands( Collections.singletonList(command) );
 			amContainer.setLocalResources( constructLocalResourceMap(yconf) );
 			amContainer.setEnvironment( constructEnvionmentMap(yconf) );

 			// Set up resource type requirements for ApplicationMaster
 			int memHeap = _dmlConfig.getIntValue(DMLConfig.YARN_APPMASTERMEM);
 			int memAlloc = (int) computeMemoryAllocation(memHeap);
 			Resource capability = Records.newRecord(Resource.class);
 			capability.setMemory( memAlloc );
 			capability.setVirtualCores( NUM_CORES );
 			LOG.debug("Requested application resources: memory="+memAlloc+", vcores="+NUM_CORES);

 			// Finally, set-up ApplicationSubmissionContext for the application
 			String qname = _dmlConfig.getTextValue(DMLConfig.YARN_APPQUEUE);
 			appContext.setApplicationName(APPMASTER_NAME); // application name
 			appContext.setAMContainerSpec(amContainer);
 			appContext.setResource(capability);
 			appContext.setQueue(qname); // queue
 			LOG.debug("Configured application meta data: name="+APPMASTER_NAME+", queue="+qname);

 			// submit application (non-blocking)
 			yarnClient.submitApplication(appContext);

 			// Check application status periodically (and output web ui address)
 			ApplicationReport appReport = yarnClient.getApplicationReport(appId);
 			LOG.info("Application tracking-URL: "+appReport.getTrackingUrl());
 			YarnApplicationState appState = appReport.getYarnApplicationState();
 			YarnApplicationState oldState = appState;
 			LOG.info("Application state: " + appState);
 			while( appState != YarnApplicationState.FINISHED
 					&& appState != YarnApplicationState.KILLED
 					&& appState != YarnApplicationState.FAILED )
 			{
 				Thread.sleep(APP_STATE_INTERVAL); //wait for 200ms
 				appReport = yarnClient.getApplicationReport(appId);
 				appState = appReport.getYarnApplicationState();
 				if( appState != oldState ) {
 					oldState = appState;
 					LOG.info("Application state: " + appState);
 				}
 			}
 			//check final status (failed or succeeded)
 			FinalApplicationStatus finalState = appReport.getFinalApplicationStatus();
 			LOG.info("Application final status: " + finalState);

 			//show application and total runtime
 			double appRuntime = (double)(appReport.getFinishTime() - appReport.getStartTime()) / 1000;
 			LOG.info( "Application runtime: " + appRuntime + " sec." );
 			LOG.info( "Total runtime: " + String.format("%.3f", time.stop()/1000) + " sec.");

 			//raised script-level error in case of failed final status
 			if( finalState != FinalApplicationStatus.SUCCEEDED )
 			{
 				//propagate script-level stop call message
 				String stop_msg = readMessageToHDFSWorkingDir(_dmlConfig, yconf, appId);
 				if( stop_msg != null )
 					throw new DMLScriptException(stop_msg);

 				//generic failure message
 				throw new DMLRuntimeException("DML yarn app master finished with final status: "+finalState+".");
 			}

 			ret = true;
 		}
 		catch(DMLScriptException ex) {
 			//rethrow DMLScriptException to propagate stop call
 			throw ex;
 		}
 		catch(Exception ex) {
 			LOG.error("Failed to run DML yarn app master.", ex);
 			ret = false;
 		}
 		finally
 		{
 			//cleanup working directory
 			if( hdfsWD != null )
 				MapReduceTool.deleteFileIfExistOnHDFS(hdfsWD);
 		}

 		return ret;
 	}

 	/**
 	 *
 	 * @param appId
 	 * @throws ParseException
 	 * @throws IOException
 	 * @throws DMLRuntimeException
 	 * @throws InterruptedException
 	 */
 	@SuppressWarnings("deprecation")
 	private void copyResourcesToHdfsWorkingDir( YarnConfiguration yconf, String hdfsWD )
 		throws ParseException, IOException, DMLRuntimeException, InterruptedException
 	{
 		FileSystem fs = FileSystem.get(yconf);

 		//create working directory
 		MapReduceTool.createDirIfNotExistOnHDFS(hdfsWD, DMLConfig.DEFAULT_SHARED_DIR_PERMISSION);

 		//serialize the dml config to HDFS file
 		//NOTE: we do not modify and ship the absolute scratch space path of the current user
 		//because this might result in permission issues if the app master is run with a different user
 		//(runtime plan migration during resource reoptimizations now needs to use qualified names
 		//for shipping/reading intermediates) TODO modify resource reoptimizer on prototype integration.
 		Path confPath = new Path(hdfsWD, DML_CONFIG_NAME);
 		FSDataOutputStream fout = fs.create(confPath, true);
 		//_dmlConfig.makeQualifiedScratchSpacePath();
 		fout.writeBytes(_dmlConfig.serializeDMLConfig() + "\n");
 		fout.close();
 		_hdfsDMLConfig = confPath.makeQualified(fs).toString();
 		LOG.debug("DML config written to HDFS file: "+_hdfsDMLConfig+"");

 		//serialize the dml script to HDFS file
 		Path scriptPath = new Path(hdfsWD, DML_SCRIPT_NAME);
 		FSDataOutputStream fout2 = fs.create(scriptPath, true);
 		fout2.writeBytes(_dmlScript);
 		fout2.close();
 		_hdfsDMLScript = scriptPath.makeQualified(fs).toString();
 		LOG.debug("DML script written to HDFS file: "+_hdfsDMLScript+"");

 		// copy local jar file to HDFS (try to get the original jar filename)
 		String fname = getLocalJarFileNameFromEnvConst();
 		if( fname == null ){
 			//get location of unpacked jar classes and repackage (if required)
 			String lclassFile = DMLYarnClient.class.getProtectionDomain().getCodeSource().getLocation().getPath().toString();
 			File flclassFile = new File( lclassFile );
 			if( !flclassFile.isDirectory() ) //called w/ jar
 				fname = lclassFile;
 			else //called w/ unpacked jar (need to be repackaged)
 				fname = createJar(lclassFile);
 		}
 		Path srcPath = new Path(fname);
 		Path dstPath = new Path(hdfsWD, srcPath.getName());
 		FileUtil.copy(FileSystem.getLocal(yconf), srcPath, fs, dstPath, false, true, yconf);
 		_hdfsJarFile = dstPath.makeQualified(fs).toString();
 		LOG.debug("Jar file copied from local file: "+srcPath.toString()+" to HDFS file: "+dstPath.toString());
 	}

 	/**
 	 *
 	 * @return null if the constant does not exists
 	 */
 	private String getLocalJarFileNameFromEnvConst()
 	{
 		String fname = null;

 		try
 		{
 			//parse environment constants
 			Map<String, String> env = System.getenv();
 			if( env.containsKey(JARFILE_ENV_CONST) ){
 				String tmp = env.get(JARFILE_ENV_CONST);
 				String[] tmpargs = tmp.split(" ");
 				for( int i=0; i<tmpargs.length && fname==null; i++ )
 					if( tmpargs[i]!=null && tmpargs[i].endsWith("RunJar") )
 						fname = tmpargs[i+1];
 			}
 		}
 		catch(Exception ex)
 		{
 			LOG.warn("Failed to parse environment variables ("+ex.getMessage()+")");
 			fname = null; //indicate to use fallback strategy
 		}

 		//give warning that we fallback to alternative jar shipping method
 		if( LOG.isDebugEnabled() && fname == null ) {
 			LOG.warn("Failed to find jar file via environment variable '"+JARFILE_ENV_CONST+"', fallback to jar packaging.");
 		}

 		return fname;
 	}

 	/**
 	 * This is our fallback strategy for obtaining our SystemML.jar that we need
 	 * to submit as resource for the yarn application. We repackage the unzipped
 	 * jar to a temporary jar and later copy it to hdfs.
 	 *
 	 * @param dir
 	 * @return
 	 * @throws IOException
 	 * @throws InterruptedException
 	 */
 	private String createJar( String dir )
 		throws IOException, InterruptedException
 	{
 		//construct jar command
 		String jarname = dir+"/"+DML_JAR_NAME;
 		File fdir = new File(dir);
 		File[] tmp = fdir.listFiles();
 		StringBuilder flist = new StringBuilder();
 		for( File ftmp : tmp ) {
 			flist.append(ftmp.getName());
 			flist.append(" ");
 		}

 		//get jdk home (property 'java.home' gives jre-home of parent jdk or standalone)
 		String javahome = System.getProperty("java.home");
 		File fjdkhome = new File(new File(javahome).getParent() + File.separator  + "bin");
 		String jarPrefix = "";
 		if( fjdkhome.exists() ) { //exists if jdk
 			jarPrefix = fjdkhome.getAbsolutePath();
 			jarPrefix += File.separator;
 		}
 		if( jarPrefix.isEmpty() )
 			LOG.warn("Failed to find jdk home of running jre (java.home="+javahome+").");

 		//execute jar command
 		String command = jarPrefix+"jar cf "+jarname+" "+flist.subSequence(0, flist.length()-1);
 		LOG.debug("Packaging jar of unzipped files: "+command);
 		Process child = Runtime.getRuntime().exec(command, null, fdir);
 		int c = 0;
 		while ((c = child.getInputStream().read()) != -1)
 			System.out.print((char) c);
 		while ((c = child.getErrorStream().read()) != -1)
 			System.err.print((char) c);
 		child.waitFor();

 		return jarname;
 	}

 	/**
 	 *
 	 * @param args
 	 * @param conf
 	 * @return
 	 */
 	private String constructAMCommand( String[] args, DMLConfig conf )
 	{
 		//start command
 		StringBuilder command = new StringBuilder();
 		command.append(Environment.JAVA_HOME.$$() + "/bin/java");

 		//add client jvm arguments (concatenation of HADOOP_CLIENT_OPTS and HADOOP_OPTS)
 		if( System.getenv().containsKey(JVMOPTS_ENV_CONST) ) {
 			String externalArgs = System.getenv(JVMOPTS_ENV_CONST);
 			//safe parsing and replacement of redundant Xmx, Xms, Xmn
 			if( externalArgs != null ) {
 				String[] parts = externalArgs.split(" ");
 				for( int i=0; i<parts.length; i++ )
 					if( !(   parts[i].startsWith("-Xmx")
 						  || parts[i].startsWith("-Xms")
 						  || parts[i].startsWith("-Xmn") ) )
 					{
 						command.append(" ");
 						command.append(parts[i]);
 					}
 			}
 		}

 		//add jvm heap configuration (specify xmn for default gcpolicy:gencon)
 		int memHeap = conf.getIntValue(DMLConfig.YARN_APPMASTERMEM);
 		command.append(" -Xmx"+memHeap+"m");
 		command.append(" -Xms"+memHeap+"m");
 		command.append(" -Xmn"+(int)(memHeap/10)+"m");
 		command.append(' ');
 		command.append(DMLAppMaster.class.getName());

 		//add command line args (modify script and config file path)
 		for( int i=0; i<_args.length; i++ )
 		{
 			String arg = _args[i];
 			command.append(' ');
 			if( i>0 && _args[i-1].equals("-f") ){
 				command.append(_hdfsDMLScript);
 				command.append(" -config=" + _hdfsDMLConfig);
 			}
 			else if( _args[i].startsWith("-config") ){
 				//ignore because config added with -f
 			}
 			else
 				command.append(arg);
 		}

 		//setup stdout and stderr logs
 		command.append(" 1>" + ApplicationConstants.LOG_DIR_EXPANSION_VAR + "/stdout");
 		command.append(" 2>" + ApplicationConstants.LOG_DIR_EXPANSION_VAR + "/stderr");

 		return command.toString();
 	}

 	/**
 	 *
 	 * @param yconf
 	 * @param path
 	 * @param lpath
 	 * @return
 	 * @throws IOException
 	 */
 	private Map<String, LocalResource> constructLocalResourceMap(YarnConfiguration yconf)
 		throws IOException
 	{
 		Map<String, LocalResource> rMap = new HashMap<String, LocalResource>();
 		Path path = new Path(_hdfsJarFile);

 		LocalResource resource = Records.newRecord(LocalResource.class);
 		FileStatus jarStat = FileSystem.get(yconf).getFileStatus(path);
 		resource.setResource(ConverterUtils.getYarnUrlFromPath(path));
 		resource.setSize(jarStat.getLen());
 		resource.setTimestamp(jarStat.getModificationTime());
 		resource.setType(LocalResourceType.FILE);
 		resource.setVisibility(LocalResourceVisibility.PUBLIC);

 		rMap.put(DML_JAR_NAME, resource);
 		return rMap;
 	}

 	/**
 	 *
 	 * @param yconf
 	 * @return
 	 * @throws IOException
 	 */
 	private Map<String, String> constructEnvionmentMap(YarnConfiguration yconf)
 		throws IOException
 	{
 		Map<String, String> eMap = new HashMap<String, String>();

 		//setup default app master environment
 		StringBuilder classpath = new StringBuilder();
 		for (String value : yconf.getStrings(
 				YarnConfiguration.YARN_APPLICATION_CLASSPATH,
 				YarnConfiguration.DEFAULT_YARN_APPLICATION_CLASSPATH))
 		{
 			if( classpath.length() > 0 )
 				classpath.append(File.pathSeparator);
 			classpath.append( value.trim() );
 		}

 		//setup mapreduce appmaster environment (for robustness if not included in default environment)
 		//for example, by default HDP 2.2 did not include mapred client libraries in this configuration
 		//note: we cannot use MRConfigurationNames.MR_APPLICATION_CLASSPATH because it refers to HDFS and $PWD that needs to be setup
 		Map<String, String> env = System.getenv();
 		String mapred_home = null;
 		//get mapred home via alternative environment variables
 		if( env.containsKey(MAPRED_HOME_ENV_CONST) ) {
 			mapred_home = env.get(MAPRED_HOME_ENV_CONST);
 		}
 		else if ( env.containsKey(HADOOP_HOME_ENV_CONST) ){
 			String tmp = env.get(HADOOP_HOME_ENV_CONST);
 			mapred_home = tmp + File.separator + ".." + File.separator + "hadoop-mapreduce";
 		}
 		//concatenate mapred home libs to classpath
 		if( mapred_home != null ) {
 			if( classpath.length() > 0 )
 				classpath.append( File.pathSeparator );
 			classpath.append( mapred_home + File.separator + "*" );
 			classpath.append( File.pathSeparator );
 			classpath.append( mapred_home + File.separator + "lib" + File.separator + "*" );
 		}

 		eMap.put(Environment.CLASSPATH.name(), classpath.toString());
 		MRApps.setClasspath(eMap, yconf);

 		LOG.debug("Constructed environment classpath: "+classpath.toString());

 		return eMap;
 	}

 	/**
 	 *
 	 * @param conf
 	 * @param yconf
 	 * @param appId
 	 * @return
 	 */
 	private String readMessageToHDFSWorkingDir(DMLConfig conf, YarnConfiguration yconf, ApplicationId appId)
 	{
 		String ret = null;

 		//construct working directory (consistent with client)
 		String hdfsWD = DMLAppMasterUtils.constructHDFSWorkingDir(conf, appId);
 		Path msgPath = new Path(hdfsWD, DMLYarnClient.DML_STOPMSG_NAME);

 		//write given message to hdfs
 		try {
 			FileSystem fs = FileSystem.get(yconf);
 			if( fs.exists(msgPath) )
 			{
 				FSDataInputStream fin = fs.open(msgPath);
 				BufferedReader br = new BufferedReader(new InputStreamReader(fin));
 				ret = br.readLine();
 				fin.close();
 				LOG.debug("Stop message read from HDFS file "+msgPath+": "+ret );
 			}
 		}
 		catch(Exception ex) {
 			LOG.error("Failed to read stop message from HDFS file: "+msgPath, ex);
 		}

 		return ret;
 	}

 	/**
 	 *
 	 * @param heapsize
 	 * @return
 	 */
 	public static long computeMemoryAllocation( long heapsize )
 	{
 		long ret = heapsize;

 		if( heapsize * MEM_FACTOR - heapsize < MAX_MEM_OVERHEAD )
 			ret = (long) (heapsize * MEM_FACTOR);
 		else
 			ret = heapsize + MAX_MEM_OVERHEAD;

 		return ret;
 	}
 }
	/*
	* 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
	*
	* 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.sysml.yarn;

	import java.io.BufferedReader;
	import java.io.File;
	import java.io.IOException;
	import java.io.InputStreamReader;
	import java.util.Collections;
	import java.util.HashMap;
	import java.util.Map;

	import org.apache.commons.logging.Log;
	import org.apache.commons.logging.LogFactory;
	import org.apache.hadoop.fs.FSDataInputStream;
	import org.apache.hadoop.fs.FSDataOutputStream;
	import org.apache.hadoop.fs.FileStatus;
	import org.apache.hadoop.fs.FileSystem;
	import org.apache.hadoop.fs.FileUtil;
	import org.apache.hadoop.fs.Path;
	import org.apache.hadoop.mapreduce.v2.util.MRApps;
	import org.apache.hadoop.yarn.api.ApplicationConstants;
	import org.apache.hadoop.yarn.api.ApplicationConstants.Environment;
	import org.apache.hadoop.yarn.api.records.ApplicationId;
	import org.apache.hadoop.yarn.api.records.ApplicationReport;
	import org.apache.hadoop.yarn.api.records.ApplicationSubmissionContext;
	import org.apache.hadoop.yarn.api.records.ContainerLaunchContext;
	import org.apache.hadoop.yarn.api.records.FinalApplicationStatus;
	import org.apache.hadoop.yarn.api.records.LocalResource;
	import org.apache.hadoop.yarn.api.records.LocalResourceType;
	import org.apache.hadoop.yarn.api.records.LocalResourceVisibility;
	import org.apache.hadoop.yarn.api.records.Resource;
	import org.apache.hadoop.yarn.api.records.YarnApplicationState;
	import org.apache.hadoop.yarn.client.api.YarnClient;
	import org.apache.hadoop.yarn.client.api.YarnClientApplication;
	import org.apache.hadoop.yarn.conf.YarnConfiguration;
	import org.apache.hadoop.yarn.util.ConverterUtils;
	import org.apache.hadoop.yarn.util.Records;

	import org.apache.sysml.conf.DMLConfig;
	import org.apache.sysml.parser.ParseException;
	import org.apache.sysml.runtime.DMLRuntimeException;
	import org.apache.sysml.runtime.DMLScriptException;
	import org.apache.sysml.runtime.controlprogram.parfor.stat.Timing;
	import org.apache.sysml.runtime.util.MapReduceTool;

	/**
	* NOTES:
	* * Security: By default, submitted applications are ran as user 'yarn'.
	* In order to allow for security and relative filenames on hdfs (/user/<username>/.),
	* we can configure the LinuxContainerExecutor in yarn-site.xml, which runs the
	* application as the user who submits the application.
	* * SystemML.jar file dependency: We need to submit the SystemML.jar along with the
	* application. Unfortunately, hadoop jar unpacks the jar such that we dont get a handle
	* to the original jar filename. We currently parse the constant IBM_JAVA_COMMAND_LINE
	* to get the jar filename. For robustness, we fall back to repackaging the unzipped files
	* to a jar if this constant does not exist.
	*
	*/
	public class DMLYarnClient
	{

	private static final Log LOG = LogFactory.getLog(DMLYarnClient.class);

	//Internal configuration parameters
	// environment variable to obtain the original jar filename
	public static final String JARFILE_ENV_CONST = "IBM_JAVA_COMMAND_LINE";
	// environment variable to obtain default jvm arguments
	public static final String JVMOPTS_ENV_CONST = "HADOOP_OPTS";
	// environment variable to obtain mapred home (for robustness only)
	public static final String MAPRED_HOME_ENV_CONST = "HADOOP_MAPRED_HOME";
	public static final String HADOOP_HOME_ENV_CONST = "HADOOP_HOME";
	// default of 1 core since YARN scheduler does not take the number of cores into account yet
	public static final int NUM_CORES = 1;
	// factor for compute virtual memory to request based on given max heap size
	// (if the absolute overhead is larger than a threshold, we use this threshold as a constant overhead)
	public static final double MEM_FACTOR = 1.5;
	public static final int MAX_MEM_OVERHEAD = 2*1024; //2GB
	// default application state report (in milliseconds)
	public static final int APP_STATE_INTERVAL = 200;
	// default application master name
	public static final String APPMASTER_NAME = "SystemML-AM";
	// default dml script file name for hdfs script serialization
	public static final String DML_SCRIPT_NAME = "script.dml";
	// default dml config file name for hdfs config serialization
	public static final String DML_CONFIG_NAME = "config.xml";
	// default SystemML jar file name for hdfs jar copy
	public static final String DML_JAR_NAME = "SystemML.jar";
	// default dml stop message file name for hdfs message serialization
	public static final String DML_STOPMSG_NAME = "stop_msg.txt";


	private String _dmlScript = null;
	private DMLConfig _dmlConfig = null;
	private String[] _args = null;

	//hdfs file names local resources
	private String _hdfsJarFile = null;
	private String _hdfsDMLScript = null;
	private String _hdfsDMLConfig = null;

	/**
	* Protected since only supposed to be accessed via proxy in same package.
	* This is to ensure robustness in case of missing yarn libraries.
	*
	* @param dmlScriptStr
	* @param conf
	* @param args
	*/
	protected DMLYarnClient(String dmlScriptStr, DMLConfig conf, String[] args)
	{
	_dmlScript = dmlScriptStr;
	_dmlConfig = conf;
	_args = args;
	}


	/**
	* Method to launch the dml yarn app master and execute the given dml script
	* with the given configuration and jar file.
	*
	* NOTE: on launching the yarn app master, we do not explicitly probe if we
	* are running on a yarn or MR1 cluster. In case of MR1, already the class
	* YarnConfiguration will not be found and raise a classnotfound. In case of any
	* exception we fall back to run CP directly in the client process.
	*
	* @return true if dml program successfully executed as yarn app master
	* @throws IOException
	*/
	protected boolean launchDMLYarnAppmaster()
	throws IOException, DMLScriptException
	{
	boolean ret = false;
	String hdfsWD = null;

	try
	{
	Timing time = new Timing(true);

	// load yarn configuration
	YarnConfiguration yconf = new YarnConfiguration();

	// create yarn client
	YarnClient yarnClient = YarnClient.createYarnClient();
	yarnClient.init(yconf);
	yarnClient.start();

	// create application and get the ApplicationID
	YarnClientApplication app = yarnClient.createApplication();
	ApplicationSubmissionContext appContext = app.getApplicationSubmissionContext();
	ApplicationId appId = appContext.getApplicationId();
	LOG.debug("Created application (applicationID: "+appId+")");

	// prepare hdfs working directory via ApplicationID
	// copy script, config, jar file to hdfs
	hdfsWD = DMLAppMasterUtils.constructHDFSWorkingDir(_dmlConfig, appId);
	copyResourcesToHdfsWorkingDir(yconf, hdfsWD);

	//construct command line argument
	String command = constructAMCommand(_args, _dmlConfig);
	LOG.debug("Constructed application master command: \n"+command);

	// set up the container launch context for the application master
	ContainerLaunchContext amContainer = Records.newRecord(ContainerLaunchContext.class);
	amContainer.setCommands( Collections.singletonList(command) );
	amContainer.setLocalResources( constructLocalResourceMap(yconf) );
	amContainer.setEnvironment( constructEnvionmentMap(yconf) );

	// Set up resource type requirements for ApplicationMaster
	int memHeap = _dmlConfig.getIntValue(DMLConfig.YARN_APPMASTERMEM);
	int memAlloc = (int) computeMemoryAllocation(memHeap);
	Resource capability = Records.newRecord(Resource.class);
	capability.setMemory( memAlloc );
	capability.setVirtualCores( NUM_CORES );
	LOG.debug("Requested application resources: memory="+memAlloc+", vcores="+NUM_CORES);

	// Finally, set-up ApplicationSubmissionContext for the application
	String qname = _dmlConfig.getTextValue(DMLConfig.YARN_APPQUEUE);
	appContext.setApplicationName(APPMASTER_NAME); // application name
	appContext.setAMContainerSpec(amContainer);
	appContext.setResource(capability);
	appContext.setQueue(qname); // queue
	LOG.debug("Configured application meta data: name="+APPMASTER_NAME+", queue="+qname);

	// submit application (non-blocking)
	yarnClient.submitApplication(appContext);

	// Check application status periodically (and output web ui address)
	ApplicationReport appReport = yarnClient.getApplicationReport(appId);
	LOG.info("Application tracking-URL: "+appReport.getTrackingUrl());
	YarnApplicationState appState = appReport.getYarnApplicationState();
	YarnApplicationState oldState = appState;
	LOG.info("Application state: " + appState);
	while( appState != YarnApplicationState.FINISHED
	&& appState != YarnApplicationState.KILLED
	&& appState != YarnApplicationState.FAILED )
	{
	Thread.sleep(APP_STATE_INTERVAL); //wait for 200ms
	appReport = yarnClient.getApplicationReport(appId);
	appState = appReport.getYarnApplicationState();
	if( appState != oldState ) {
	oldState = appState;
	LOG.info("Application state: " + appState);
	}
	}
	//check final status (failed or succeeded)
	FinalApplicationStatus finalState = appReport.getFinalApplicationStatus();
	LOG.info("Application final status: " + finalState);

	//show application and total runtime
	double appRuntime = (double)(appReport.getFinishTime() - appReport.getStartTime()) / 1000;
	LOG.info( "Application runtime: " + appRuntime + " sec." );
	LOG.info( "Total runtime: " + String.format("%.3f", time.stop()/1000) + " sec.");

	//raised script-level error in case of failed final status
	if( finalState != FinalApplicationStatus.SUCCEEDED )
	{
	//propagate script-level stop call message
	String stop_msg = readMessageToHDFSWorkingDir(_dmlConfig, yconf, appId);
	if( stop_msg != null )
	throw new DMLScriptException(stop_msg);

	//generic failure message
	throw new DMLRuntimeException("DML yarn app master finished with final status: "+finalState+".");
	}

	ret = true;
	}
	catch(DMLScriptException ex) {
	//rethrow DMLScriptException to propagate stop call
	throw ex;
	}
	catch(Exception ex) {
	LOG.error("Failed to run DML yarn app master.", ex);
	ret = false;
	}
	finally
	{
	//cleanup working directory
	if( hdfsWD != null )
	MapReduceTool.deleteFileIfExistOnHDFS(hdfsWD);
	}

	return ret;
	}

	/**
	*
	* @param appId
	* @throws ParseException
	* @throws IOException
	* @throws DMLRuntimeException
	* @throws InterruptedException
	*/
	@SuppressWarnings("deprecation")
	private void copyResourcesToHdfsWorkingDir( YarnConfiguration yconf, String hdfsWD )
	throws ParseException, IOException, DMLRuntimeException, InterruptedException
	{
	FileSystem fs = FileSystem.get(yconf);

	//create working directory
	MapReduceTool.createDirIfNotExistOnHDFS(hdfsWD, DMLConfig.DEFAULT_SHARED_DIR_PERMISSION);

	//serialize the dml config to HDFS file
	//NOTE: we do not modify and ship the absolute scratch space path of the current user
	//because this might result in permission issues if the app master is run with a different user
	//(runtime plan migration during resource reoptimizations now needs to use qualified names
	//for shipping/reading intermediates) TODO modify resource reoptimizer on prototype integration.
	Path confPath = new Path(hdfsWD, DML_CONFIG_NAME);
	FSDataOutputStream fout = fs.create(confPath, true);
	//_dmlConfig.makeQualifiedScratchSpacePath();
	fout.writeBytes(_dmlConfig.serializeDMLConfig() + "\n");
	fout.close();
	_hdfsDMLConfig = confPath.makeQualified(fs).toString();
	LOG.debug("DML config written to HDFS file: "+_hdfsDMLConfig+"");

	//serialize the dml script to HDFS file
	Path scriptPath = new Path(hdfsWD, DML_SCRIPT_NAME);
	FSDataOutputStream fout2 = fs.create(scriptPath, true);
	fout2.writeBytes(_dmlScript);
	fout2.close();
	_hdfsDMLScript = scriptPath.makeQualified(fs).toString();
	LOG.debug("DML script written to HDFS file: "+_hdfsDMLScript+"");

	// copy local jar file to HDFS (try to get the original jar filename)
	String fname = getLocalJarFileNameFromEnvConst();
	if( fname == null ){
	//get location of unpacked jar classes and repackage (if required)
	String lclassFile = DMLYarnClient.class.getProtectionDomain().getCodeSource().getLocation().getPath().toString();
	File flclassFile = new File( lclassFile );
	if( !flclassFile.isDirectory() ) //called w/ jar
	fname = lclassFile;
	else //called w/ unpacked jar (need to be repackaged)
	fname = createJar(lclassFile);
	}
	Path srcPath = new Path(fname);
	Path dstPath = new Path(hdfsWD, srcPath.getName());
	FileUtil.copy(FileSystem.getLocal(yconf), srcPath, fs, dstPath, false, true, yconf);
	_hdfsJarFile = dstPath.makeQualified(fs).toString();
	LOG.debug("Jar file copied from local file: "+srcPath.toString()+" to HDFS file: "+dstPath.toString());
	}

	/**
	*
	* @return null if the constant does not exists
	*/
	private String getLocalJarFileNameFromEnvConst()
	{
	String fname = null;

	try
	{
	//parse environment constants
	Map<String, String> env = System.getenv();
	if( env.containsKey(JARFILE_ENV_CONST) ){
	String tmp = env.get(JARFILE_ENV_CONST);
	String[] tmpargs = tmp.split(" ");
	for( int i=0; i<tmpargs.length && fname==null; i++ )
	if( tmpargs[i]!=null && tmpargs[i].endsWith("RunJar") )
	fname = tmpargs[i+1];
	}
	}
	catch(Exception ex)
	{
	LOG.warn("Failed to parse environment variables ("+ex.getMessage()+")");
	fname = null; //indicate to use fallback strategy
	}

	//give warning that we fallback to alternative jar shipping method
	if( LOG.isDebugEnabled() && fname == null ) {
	LOG.warn("Failed to find jar file via environment variable '"+JARFILE_ENV_CONST+"', fallback to jar packaging.");
	}

	return fname;
	}

	/**
	* This is our fallback strategy for obtaining our SystemML.jar that we need
	* to submit as resource for the yarn application. We repackage the unzipped
	* jar to a temporary jar and later copy it to hdfs.
	*
	* @param dir
	* @return
	* @throws IOException
	* @throws InterruptedException
	*/
	private String createJar( String dir )
	throws IOException, InterruptedException
	{
	//construct jar command
	String jarname = dir+"/"+DML_JAR_NAME;
	File fdir = new File(dir);
	File[] tmp = fdir.listFiles();
	StringBuilder flist = new StringBuilder();
	for( File ftmp : tmp ) {
	flist.append(ftmp.getName());
	flist.append(" ");
	}

	//get jdk home (property 'java.home' gives jre-home of parent jdk or standalone)
	String javahome = System.getProperty("java.home");
	File fjdkhome = new File(new File(javahome).getParent() + File.separator + "bin");
	String jarPrefix = "";
	if( fjdkhome.exists() ) { //exists if jdk
	jarPrefix = fjdkhome.getAbsolutePath();
	jarPrefix += File.separator;
	}
	if( jarPrefix.isEmpty() )
	LOG.warn("Failed to find jdk home of running jre (java.home="+javahome+").");

	//execute jar command
	String command = jarPrefix+"jar cf "+jarname+" "+flist.subSequence(0, flist.length()-1);
	LOG.debug("Packaging jar of unzipped files: "+command);
	Process child = Runtime.getRuntime().exec(command, null, fdir);
	int c = 0;
	while ((c = child.getInputStream().read()) != -1)
	System.out.print((char) c);
	while ((c = child.getErrorStream().read()) != -1)
	System.err.print((char) c);
	child.waitFor();

	return jarname;
	}

	/**
	*
	* @param args
	* @param conf
	* @return
	*/
	private String constructAMCommand( String[] args, DMLConfig conf )
	{
	//start command
	StringBuilder command = new StringBuilder();
	command.append(Environment.JAVA_HOME.$$() + "/bin/java");

	//add client jvm arguments (concatenation of HADOOP_CLIENT_OPTS and HADOOP_OPTS)
	if( System.getenv().containsKey(JVMOPTS_ENV_CONST) ) {
	String externalArgs = System.getenv(JVMOPTS_ENV_CONST);
	//safe parsing and replacement of redundant Xmx, Xms, Xmn
	if( externalArgs != null ) {
	String[] parts = externalArgs.split(" ");
	for( int i=0; i<parts.length; i++ )
	if( !( parts[i].startsWith("-Xmx")
	\|\| parts[i].startsWith("-Xms")
	\|\| parts[i].startsWith("-Xmn") ) )
	{
	command.append(" ");
	command.append(parts[i]);
	}
	}
	}

	//add jvm heap configuration (specify xmn for default gcpolicy:gencon)
	int memHeap = conf.getIntValue(DMLConfig.YARN_APPMASTERMEM);
	command.append(" -Xmx"+memHeap+"m");
	command.append(" -Xms"+memHeap+"m");
	command.append(" -Xmn"+(int)(memHeap/10)+"m");
	command.append(' ');
	command.append(DMLAppMaster.class.getName());

	//add command line args (modify script and config file path)
	for( int i=0; i<_args.length; i++ )
	{
	String arg = _args[i];
	command.append(' ');
	if( i>0 && _args[i-1].equals("-f") ){
	command.append(_hdfsDMLScript);
	command.append(" -config=" + _hdfsDMLConfig);
	}
	else if( _args[i].startsWith("-config") ){
	//ignore because config added with -f
	}
	else
	command.append(arg);
	}

	//setup stdout and stderr logs
	command.append(" 1>" + ApplicationConstants.LOG_DIR_EXPANSION_VAR + "/stdout");
	command.append(" 2>" + ApplicationConstants.LOG_DIR_EXPANSION_VAR + "/stderr");

	return command.toString();
	}

	/**
	*
	* @param yconf
	* @param path
	* @param lpath
	* @return
	* @throws IOException
	*/
	private Map<String, LocalResource> constructLocalResourceMap(YarnConfiguration yconf)
	throws IOException
	{
	Map<String, LocalResource> rMap = new HashMap<String, LocalResource>();
	Path path = new Path(_hdfsJarFile);

	LocalResource resource = Records.newRecord(LocalResource.class);
	FileStatus jarStat = FileSystem.get(yconf).getFileStatus(path);
	resource.setResource(ConverterUtils.getYarnUrlFromPath(path));
	resource.setSize(jarStat.getLen());
	resource.setTimestamp(jarStat.getModificationTime());
	resource.setType(LocalResourceType.FILE);
	resource.setVisibility(LocalResourceVisibility.PUBLIC);

	rMap.put(DML_JAR_NAME, resource);
	return rMap;
	}

	/**
	*
	* @param yconf
	* @return
	* @throws IOException
	*/
	private Map<String, String> constructEnvionmentMap(YarnConfiguration yconf)
	throws IOException
	{
	Map<String, String> eMap = new HashMap<String, String>();

	//setup default app master environment
	StringBuilder classpath = new StringBuilder();
	for (String value : yconf.getStrings(
	YarnConfiguration.YARN_APPLICATION_CLASSPATH,
	YarnConfiguration.DEFAULT_YARN_APPLICATION_CLASSPATH))
	{
	if( classpath.length() > 0 )
	classpath.append(File.pathSeparator);
	classpath.append( value.trim() );
	}

	//setup mapreduce appmaster environment (for robustness if not included in default environment)
	//for example, by default HDP 2.2 did not include mapred client libraries in this configuration
	//note: we cannot use MRConfigurationNames.MR_APPLICATION_CLASSPATH because it refers to HDFS and $PWD that needs to be setup
	Map<String, String> env = System.getenv();
	String mapred_home = null;
	//get mapred home via alternative environment variables
	if( env.containsKey(MAPRED_HOME_ENV_CONST) ) {
	mapred_home = env.get(MAPRED_HOME_ENV_CONST);
	}
	else if ( env.containsKey(HADOOP_HOME_ENV_CONST) ){
	String tmp = env.get(HADOOP_HOME_ENV_CONST);
	mapred_home = tmp + File.separator + ".." + File.separator + "hadoop-mapreduce";
	}
	//concatenate mapred home libs to classpath
	if( mapred_home != null ) {
	if( classpath.length() > 0 )
	classpath.append( File.pathSeparator );
	classpath.append( mapred_home + File.separator + "*" );
	classpath.append( File.pathSeparator );
	classpath.append( mapred_home + File.separator + "lib" + File.separator + "*" );
	}

	eMap.put(Environment.CLASSPATH.name(), classpath.toString());
	MRApps.setClasspath(eMap, yconf);

	LOG.debug("Constructed environment classpath: "+classpath.toString());

	return eMap;
	}

	/**
	*
	* @param conf
	* @param yconf
	* @param appId
	* @return
	*/
	private String readMessageToHDFSWorkingDir(DMLConfig conf, YarnConfiguration yconf, ApplicationId appId)
	{
	String ret = null;

	//construct working directory (consistent with client)
	String hdfsWD = DMLAppMasterUtils.constructHDFSWorkingDir(conf, appId);
	Path msgPath = new Path(hdfsWD, DMLYarnClient.DML_STOPMSG_NAME);

	//write given message to hdfs
	try {
	FileSystem fs = FileSystem.get(yconf);
	if( fs.exists(msgPath) )
	{
	FSDataInputStream fin = fs.open(msgPath);
	BufferedReader br = new BufferedReader(new InputStreamReader(fin));
	ret = br.readLine();
	fin.close();
	LOG.debug("Stop message read from HDFS file "+msgPath+": "+ret );
	}
	}
	catch(Exception ex) {
	LOG.error("Failed to read stop message from HDFS file: "+msgPath, ex);
	}

	return ret;
	}

	/**
	*
	* @param heapsize
	* @return
	*/
	public static long computeMemoryAllocation( long heapsize )
	{
	long ret = heapsize;

	if( heapsize * MEM_FACTOR - heapsize < MAX_MEM_OVERHEAD )
	ret = (long) (heapsize * MEM_FACTOR);
	else
	ret = heapsize + MAX_MEM_OVERHEAD;

	return ret;
	}
	}