hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/main/java/org/apache/hadoop/yarn/client/cli/SchedConfCLI.java - hadoop - 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.hadoop.yarn.client.cli;

 import com.google.common.annotations.VisibleForTesting;
 import com.sun.jersey.api.client.Client;
 import com.sun.jersey.api.client.ClientResponse;
 import com.sun.jersey.api.client.WebResource;
 import org.apache.commons.cli.CommandLine;
 import org.apache.commons.cli.GnuParser;
 import org.apache.commons.cli.MissingArgumentException;
 import org.apache.commons.cli.Options;
 import org.apache.hadoop.classification.InterfaceAudience.Public;
 import org.apache.hadoop.classification.InterfaceStability.Unstable;
 import org.apache.hadoop.conf.Configured;
 import org.apache.hadoop.util.Tool;
 import org.apache.hadoop.yarn.conf.YarnConfiguration;
 import org.apache.hadoop.yarn.webapp.dao.QueueConfigInfo;
 import org.apache.hadoop.yarn.webapp.dao.SchedConfUpdateInfo;
 import org.apache.hadoop.yarn.webapp.util.WebAppUtils;
 import org.apache.hadoop.yarn.webapp.util.YarnWebServiceUtils;

 import javax.ws.rs.core.MediaType;
 import javax.ws.rs.core.Response.Status;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;

 /**
  * CLI for modifying scheduler configuration.
  */
 @Public
 @Unstable
 public class SchedConfCLI extends Configured implements Tool {

   private static final String ADD_QUEUES_OPTION = "addQueues";
   private static final String REMOVE_QUEUES_OPTION = "removeQueues";
   private static final String UPDATE_QUEUES_OPTION = "updateQueues";
   private static final String GLOBAL_OPTIONS = "globalUpdates";
   private static final String HELP_CMD = "help";

   private static final String CONF_ERR_MSG = "Specify configuration key " +
       "value as confKey=confVal.";

   public SchedConfCLI() {
     super(new YarnConfiguration());
   }

   public static void main(String[] args) throws Exception {
     SchedConfCLI cli = new SchedConfCLI();
     int exitCode = cli.run(args);
     System.exit(exitCode);
   }

   @Override
   public int run(String[] args) throws Exception {
     Options opts = new Options();

     opts.addOption("add", ADD_QUEUES_OPTION, true,
         "Add queues with configurations");
     opts.addOption("remove", REMOVE_QUEUES_OPTION, true,
         "Remove queues");
     opts.addOption("update", UPDATE_QUEUES_OPTION, true,
         "Update queue configurations");
     opts.addOption("global", GLOBAL_OPTIONS, true,
         "Update global scheduler configurations");
     opts.addOption("h", HELP_CMD, false, "Displays help for all commands.");

     int exitCode = -1;
     CommandLine parsedCli = null;
     try {
       parsedCli = new GnuParser().parse(opts, args);
     } catch (MissingArgumentException ex) {
       System.err.println("Missing argument for options");
       printUsage();
       return exitCode;
     }

     if (parsedCli.hasOption(HELP_CMD)) {
       printUsage();
       return 0;
     }

     boolean hasOption = false;
     SchedConfUpdateInfo updateInfo = new SchedConfUpdateInfo();
     try {
       if (parsedCli.hasOption(ADD_QUEUES_OPTION)) {
         hasOption = true;
         addQueues(parsedCli.getOptionValue(ADD_QUEUES_OPTION), updateInfo);
       }
       if (parsedCli.hasOption(REMOVE_QUEUES_OPTION)) {
         hasOption = true;
         removeQueues(parsedCli.getOptionValue(REMOVE_QUEUES_OPTION),
             updateInfo);
       }
       if (parsedCli.hasOption(UPDATE_QUEUES_OPTION)) {
         hasOption = true;
         updateQueues(parsedCli.getOptionValue(UPDATE_QUEUES_OPTION),
             updateInfo);
       }
       if (parsedCli.hasOption(GLOBAL_OPTIONS)) {
         hasOption = true;
         globalUpdates(parsedCli.getOptionValue(GLOBAL_OPTIONS), updateInfo);
       }
     } catch (IllegalArgumentException e) {
       System.err.println(e.getMessage());
       return -1;
     }

     if (!hasOption) {
       System.err.println("Invalid Command Usage: ");
       printUsage();
       return -1;
     }

     Client webServiceClient = Client.create();
     WebResource webResource = webServiceClient.resource(WebAppUtils.
         getRMWebAppURLWithScheme(getConf()));
     ClientResponse response = webResource.path("ws").path("v1").path("cluster")
         .path("scheduler-conf").accept(MediaType.APPLICATION_JSON)
         .entity(YarnWebServiceUtils.toJson(updateInfo,
             SchedConfUpdateInfo.class), MediaType.APPLICATION_JSON)
         .put(ClientResponse.class);
     if (response != null) {
       if (response.getStatus() == Status.OK.getStatusCode()) {
         System.out.println("Configuration changed successfully.");
         return 0;
       } else {
         System.err.println("Configuration change unsuccessful: "
             + response.getEntity(String.class));
       }
     } else {
       System.err.println("Configuration change unsuccessful: null response");
     }
     return -1;
   }

   @VisibleForTesting
   void addQueues(String args, SchedConfUpdateInfo updateInfo) {
     if (args == null) {
       return;
     }
     ArrayList<QueueConfigInfo> queueConfigInfos = new ArrayList<>();
     for (String arg : args.split(";")) {
       queueConfigInfos.add(getQueueConfigInfo(arg));
     }
     updateInfo.setAddQueueInfo(queueConfigInfos);
   }

   @VisibleForTesting
   void removeQueues(String args, SchedConfUpdateInfo updateInfo) {
     if (args == null) {
       return;
     }
     List<String> queuesToRemove = Arrays.asList(args.split(";"));
     updateInfo.setRemoveQueueInfo(new ArrayList<>(queuesToRemove));
   }

   @VisibleForTesting
   void updateQueues(String args, SchedConfUpdateInfo updateInfo) {
     if (args == null) {
       return;
     }
     ArrayList<QueueConfigInfo> queueConfigInfos = new ArrayList<>();
     for (String arg : args.split(";")) {
       queueConfigInfos.add(getQueueConfigInfo(arg));
     }
     updateInfo.setUpdateQueueInfo(queueConfigInfos);
   }

   @VisibleForTesting
   void globalUpdates(String args, SchedConfUpdateInfo updateInfo) {
     if (args == null) {
       return;
     }
     HashMap<String, String> globalUpdates = new HashMap<>();
     for (String globalUpdate : args.split(",")) {
       putKeyValuePair(globalUpdates, globalUpdate);
     }
     updateInfo.setGlobalParams(globalUpdates);
   }

   private QueueConfigInfo getQueueConfigInfo(String arg) {
     String[] args = arg.split(":");
     String queuePath = args[0];
     Map<String, String> queueConfigs = new HashMap<>();
     if (args.length > 1) {
       String[] queueArgs = args[1].split(",");
       for (int i = 0; i < queueArgs.length; ++i) {
         putKeyValuePair(queueConfigs, queueArgs[i]);
       }
     }
     return new QueueConfigInfo(queuePath, queueConfigs);
   }

   private void putKeyValuePair(Map<String, String> kv, String args) {
     String[] argParts = args.split("=");
     if (argParts.length == 1) {
       if (argParts[0].isEmpty() || !args.contains("=")) {
         throw new IllegalArgumentException(CONF_ERR_MSG);
       } else {
         // key specified, but no value e.g. "confKey="
         kv.put(argParts[0], null);
       }
     } else if (argParts.length > 2) {
       throw new IllegalArgumentException(CONF_ERR_MSG);
     } else {
       if (argParts[0].isEmpty()) {
         throw new IllegalArgumentException(CONF_ERR_MSG);
       }
       kv.put(argParts[0], argParts[1]);
     }
   }

   private void printUsage() {
     System.out.println("yarn schedulerconf [-add "
         + "\"queueAddPath1:confKey1=confVal1,confKey2=confVal2;"
         + "queueAddPath2:confKey3=confVal3\"] "
         + "[-remove \"queueRemovePath1;queueRemovePath2\"] "
         + "[-update \"queueUpdatePath1:confKey1=confVal1\"] "
         + "[-global globalConfKey1=globalConfVal1,"
         + "globalConfKey2=globalConfVal2]\n"
         + "Example (adding queues): yarn schedulerconf -add "
         + "\"root.a.a1:capacity=100,maximum-capacity=100;root.a.a2:capacity=0,"
         + "maximum-capacity=0\"\n"
         + "Example (removing queues): yarn schedulerconf -remove \"root.a.a1;"
         + "root.a.a2\"\n"
         + "Example (updating queues): yarn schedulerconf -update \"root.a.a1"
         + ":capacity=25,maximum-capacity=25;root.a.a2:capacity=75,"
         + "maximum-capacity=75\"\n"
         + "Example (global scheduler update): yarn schedulerconf "
         + "-global yarn.scheduler.capacity.maximum-applications=10000\n"
         + "Note: This is an alpha feature, the syntax/options are subject to "
         + "change, please run at your own risk.");
   }
 }
	/**
	* 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.hadoop.yarn.client.cli;

	import com.google.common.annotations.VisibleForTesting;
	import com.sun.jersey.api.client.Client;
	import com.sun.jersey.api.client.ClientResponse;
	import com.sun.jersey.api.client.WebResource;
	import org.apache.commons.cli.CommandLine;
	import org.apache.commons.cli.GnuParser;
	import org.apache.commons.cli.MissingArgumentException;
	import org.apache.commons.cli.Options;
	import org.apache.hadoop.classification.InterfaceAudience.Public;
	import org.apache.hadoop.classification.InterfaceStability.Unstable;
	import org.apache.hadoop.conf.Configured;
	import org.apache.hadoop.util.Tool;
	import org.apache.hadoop.yarn.conf.YarnConfiguration;
	import org.apache.hadoop.yarn.webapp.dao.QueueConfigInfo;
	import org.apache.hadoop.yarn.webapp.dao.SchedConfUpdateInfo;
	import org.apache.hadoop.yarn.webapp.util.WebAppUtils;
	import org.apache.hadoop.yarn.webapp.util.YarnWebServiceUtils;

	import javax.ws.rs.core.MediaType;
	import javax.ws.rs.core.Response.Status;
	import java.util.ArrayList;
	import java.util.Arrays;
	import java.util.HashMap;
	import java.util.List;
	import java.util.Map;

	/**
	* CLI for modifying scheduler configuration.
	*/
	@Public
	@Unstable
	public class SchedConfCLI extends Configured implements Tool {

	private static final String ADD_QUEUES_OPTION = "addQueues";
	private static final String REMOVE_QUEUES_OPTION = "removeQueues";
	private static final String UPDATE_QUEUES_OPTION = "updateQueues";
	private static final String GLOBAL_OPTIONS = "globalUpdates";
	private static final String HELP_CMD = "help";

	private static final String CONF_ERR_MSG = "Specify configuration key " +
	"value as confKey=confVal.";

	public SchedConfCLI() {
	super(new YarnConfiguration());
	}

	public static void main(String[] args) throws Exception {
	SchedConfCLI cli = new SchedConfCLI();
	int exitCode = cli.run(args);
	System.exit(exitCode);
	}

	@Override
	public int run(String[] args) throws Exception {
	Options opts = new Options();

	opts.addOption("add", ADD_QUEUES_OPTION, true,
	"Add queues with configurations");
	opts.addOption("remove", REMOVE_QUEUES_OPTION, true,
	"Remove queues");
	opts.addOption("update", UPDATE_QUEUES_OPTION, true,
	"Update queue configurations");
	opts.addOption("global", GLOBAL_OPTIONS, true,
	"Update global scheduler configurations");
	opts.addOption("h", HELP_CMD, false, "Displays help for all commands.");

	int exitCode = -1;
	CommandLine parsedCli = null;
	try {
	parsedCli = new GnuParser().parse(opts, args);
	} catch (MissingArgumentException ex) {
	System.err.println("Missing argument for options");
	printUsage();
	return exitCode;
	}

	if (parsedCli.hasOption(HELP_CMD)) {
	printUsage();
	return 0;
	}

	boolean hasOption = false;
	SchedConfUpdateInfo updateInfo = new SchedConfUpdateInfo();
	try {
	if (parsedCli.hasOption(ADD_QUEUES_OPTION)) {
	hasOption = true;
	addQueues(parsedCli.getOptionValue(ADD_QUEUES_OPTION), updateInfo);
	}
	if (parsedCli.hasOption(REMOVE_QUEUES_OPTION)) {
	hasOption = true;
	removeQueues(parsedCli.getOptionValue(REMOVE_QUEUES_OPTION),
	updateInfo);
	}
	if (parsedCli.hasOption(UPDATE_QUEUES_OPTION)) {
	hasOption = true;
	updateQueues(parsedCli.getOptionValue(UPDATE_QUEUES_OPTION),
	updateInfo);
	}
	if (parsedCli.hasOption(GLOBAL_OPTIONS)) {
	hasOption = true;
	globalUpdates(parsedCli.getOptionValue(GLOBAL_OPTIONS), updateInfo);
	}
	} catch (IllegalArgumentException e) {
	System.err.println(e.getMessage());
	return -1;
	}

	if (!hasOption) {
	System.err.println("Invalid Command Usage: ");
	printUsage();
	return -1;
	}

	Client webServiceClient = Client.create();
	WebResource webResource = webServiceClient.resource(WebAppUtils.
	getRMWebAppURLWithScheme(getConf()));
	ClientResponse response = webResource.path("ws").path("v1").path("cluster")
	.path("scheduler-conf").accept(MediaType.APPLICATION_JSON)
	.entity(YarnWebServiceUtils.toJson(updateInfo,
	SchedConfUpdateInfo.class), MediaType.APPLICATION_JSON)
	.put(ClientResponse.class);
	if (response != null) {
	if (response.getStatus() == Status.OK.getStatusCode()) {
	System.out.println("Configuration changed successfully.");
	return 0;
	} else {
	System.err.println("Configuration change unsuccessful: "
	+ response.getEntity(String.class));
	}
	} else {
	System.err.println("Configuration change unsuccessful: null response");
	}
	return -1;
	}

	@VisibleForTesting
	void addQueues(String args, SchedConfUpdateInfo updateInfo) {
	if (args == null) {
	return;
	}
	ArrayList<QueueConfigInfo> queueConfigInfos = new ArrayList<>();
	for (String arg : args.split(";")) {
	queueConfigInfos.add(getQueueConfigInfo(arg));
	}
	updateInfo.setAddQueueInfo(queueConfigInfos);
	}

	@VisibleForTesting
	void removeQueues(String args, SchedConfUpdateInfo updateInfo) {
	if (args == null) {
	return;
	}
	List<String> queuesToRemove = Arrays.asList(args.split(";"));
	updateInfo.setRemoveQueueInfo(new ArrayList<>(queuesToRemove));
	}

	@VisibleForTesting
	void updateQueues(String args, SchedConfUpdateInfo updateInfo) {
	if (args == null) {
	return;
	}
	ArrayList<QueueConfigInfo> queueConfigInfos = new ArrayList<>();
	for (String arg : args.split(";")) {
	queueConfigInfos.add(getQueueConfigInfo(arg));
	}
	updateInfo.setUpdateQueueInfo(queueConfigInfos);
	}

	@VisibleForTesting
	void globalUpdates(String args, SchedConfUpdateInfo updateInfo) {
	if (args == null) {
	return;
	}
	HashMap<String, String> globalUpdates = new HashMap<>();
	for (String globalUpdate : args.split(",")) {
	putKeyValuePair(globalUpdates, globalUpdate);
	}
	updateInfo.setGlobalParams(globalUpdates);
	}

	private QueueConfigInfo getQueueConfigInfo(String arg) {
	String[] args = arg.split(":");
	String queuePath = args[0];
	Map<String, String> queueConfigs = new HashMap<>();
	if (args.length > 1) {
	String[] queueArgs = args[1].split(",");
	for (int i = 0; i < queueArgs.length; ++i) {
	putKeyValuePair(queueConfigs, queueArgs[i]);
	}
	}
	return new QueueConfigInfo(queuePath, queueConfigs);
	}

	private void putKeyValuePair(Map<String, String> kv, String args) {
	String[] argParts = args.split("=");
	if (argParts.length == 1) {
	if (argParts[0].isEmpty() \|\| !args.contains("=")) {
	throw new IllegalArgumentException(CONF_ERR_MSG);
	} else {
	// key specified, but no value e.g. "confKey="
	kv.put(argParts[0], null);
	}
	} else if (argParts.length > 2) {
	throw new IllegalArgumentException(CONF_ERR_MSG);
	} else {
	if (argParts[0].isEmpty()) {
	throw new IllegalArgumentException(CONF_ERR_MSG);
	}
	kv.put(argParts[0], argParts[1]);
	}
	}

	private void printUsage() {
	System.out.println("yarn schedulerconf [-add "
	+ "\"queueAddPath1:confKey1=confVal1,confKey2=confVal2;"
	+ "queueAddPath2:confKey3=confVal3\"] "
	+ "[-remove \"queueRemovePath1;queueRemovePath2\"] "
	+ "[-update \"queueUpdatePath1:confKey1=confVal1\"] "
	+ "[-global globalConfKey1=globalConfVal1,"
	+ "globalConfKey2=globalConfVal2]\n"
	+ "Example (adding queues): yarn schedulerconf -add "
	+ "\"root.a.a1:capacity=100,maximum-capacity=100;root.a.a2:capacity=0,"
	+ "maximum-capacity=0\"\n"
	+ "Example (removing queues): yarn schedulerconf -remove \"root.a.a1;"
	+ "root.a.a2\"\n"
	+ "Example (updating queues): yarn schedulerconf -update \"root.a.a1"
	+ ":capacity=25,maximum-capacity=25;root.a.a2:capacity=75,"
	+ "maximum-capacity=75\"\n"
	+ "Example (global scheduler update): yarn schedulerconf "
	+ "-global yarn.scheduler.capacity.maximum-applications=10000\n"
	+ "Note: This is an alpha feature, the syntax/options are subject to "
	+ "change, please run at your own risk.");
	}
	}