FairSchedulerShell.java

package org.apache.hadoop.mapred;

import java.io.IOException;
import java.net.InetSocketAddress;

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.conf.Configured;
import org.apache.hadoop.ipc.RPC;
import org.apache.hadoop.mapred.protocal.FairSchedulerProtocol;
import org.apache.hadoop.mapreduce.TaskType;
import org.apache.hadoop.net.NetUtils;
import org.apache.hadoop.security.UserGroupInformation;
import org.apache.hadoop.util.Tool;
import org.apache.hadoop.util.ToolRunner;

/**
 * Contains a main function that uses RPC client to control the slots of
 * {@link FairScheduler}
 */
public class FairSchedulerShell extends Configured implements Tool {
  private static final Log LOG = LogFactory.getLog(FairSchedulerShell.class);
  private FairSchedulerProtocol client = null;

  public FairSchedulerShell() throws IOException {
    super(null);
  }

  @Override
  public void setConf(Configuration conf) {
    super.setConf(conf);
    if (conf != null) {
      try {
        initialize();
      } catch (IOException e) {
        // Failed to connect to server
        e.printStackTrace();
      }
    }
  }

  private void initialize() throws IOException {
    if (client != null) {
      return;
    }
    InetSocketAddress address = FairScheduler.getAddress(getConf());
    UserGroupInformation ugi = UserGroupInformation.getCurrentUGI();
    client = createClient(address, getConf(), ugi);
  }
  
  private static FairSchedulerProtocol createClient(InetSocketAddress addr,
      Configuration conf, UserGroupInformation ugi)
      throws IOException {
    return (FairSchedulerProtocol) RPC.getProxy(FairSchedulerProtocol.class,
        FairSchedulerProtocol.versionID, addr, ugi, conf,
        NetUtils.getSocketFactory(conf, FairSchedulerProtocol.class));
  }

  public int getFSMaxSlots(String trackerName, TaskType type)
      throws IOException {
    return client.getFSMaxSlots(trackerName, type);
  }


  public void setFSMaxSlots(String trackerName, TaskType type, int slots)
      throws IOException {
    client.setFSMaxSlots(trackerName, type, slots);
  }

  public void resetFSMaxSlots() throws IOException {
    client.resetFSMaxSlots();
  }

  public int[] getPoolMaxTasks(String pool) throws IOException {
    return client.getPoolMaxTasks(pool);
  }

  public int[] getPoolRunningTasks(String pool) throws IOException {
    return client.getPoolRunningTasks(pool);
  }

  /**
   * Displays format of commands.
   */
  private static void printUsage(String cmd) {
    String prefix = "Usage: java " + FairSchedulerShell.class.getSimpleName();
    if ("-getfsmaxslots".equalsIgnoreCase(cmd)) {
      System.err.println(prefix + " [-getfsmaxslots trackerName1, trackerName2 ... ]");
    } else if ("-setfsmaxslots".equalsIgnoreCase(cmd)) {
      System.err.println(prefix + " [-setfsmaxslots trackerName #maps #reduces trackerName2 #maps #reduces ... ]");
    } else {
      System.err.println(prefix);
      System.err.println("          [-setfsmaxslots trackerName #maps #reduces trackerName2 #maps #reduces ... ]");
      System.err.println("          [-getfsmaxslots trackerName1, trackerName2 ... ]");
      System.err.println("          [-resetfsmaxslots]");
      System.err.println("          [-help [cmd]]");
      System.err.println();
      ToolRunner.printGenericCommandUsage(System.err);
    }
  }

  @Override
  public int run(String[] arg) throws Exception {
    int curr = 0;
    if (arg.length == 0) {
      printUsage("");
      return -1;
    }
    if (arg[curr].equalsIgnoreCase("-help")) {
      curr += 1;
      if (arg.length == curr) {
        printUsage("");
        return 0;
      }
      printUsage(arg[curr]);
      return 0;
    }
    if (arg[curr].equalsIgnoreCase("-resetfsmaxslots")) {
      resetFSMaxSlots();
      return 0;
    }
    if (arg[curr].equalsIgnoreCase("-getfsmaxslots")) {
      curr += 1;
      if (arg.length == curr) {
        printUsage(arg[0]);
        return -1;
      }
      for (int i = curr; i < arg.length; ++i) {
        String trackerName = arg[i];
        int maxMap = getFSMaxSlots(trackerName, TaskType.MAP);
        int maxReduce = getFSMaxSlots(trackerName, TaskType.REDUCE);
        System.out.println(trackerName + " " + maxMap + " " +  maxReduce);
      }
      return 0;
    }
    if (arg[curr].equalsIgnoreCase("-setfsmaxslots")) {
      curr += 1;
      if (arg.length <= curr + 2) {
        printUsage(arg[0]);
        return -1;
      }
      for (int i = curr; i < arg.length; i += 3) {
        String trackerName = arg[i];
        int maxMap = Integer.parseInt(arg[i + 1]);
        int maxReduce = Integer.parseInt(arg[i + 2]);
        setFSMaxSlots(trackerName, TaskType.MAP, maxMap);
        setFSMaxSlots(trackerName, TaskType.REDUCE, maxReduce);
        System.out.println("set " + trackerName + " " +
            maxMap + " " +  maxReduce);
      }
      return 0;
    }
    if (arg[curr].equalsIgnoreCase("-pooltasks")) {
      curr += 1;
      if (arg.length == curr) {
        printUsage(arg[0]);
        return -1;
      }
      String pool = arg[curr];
      int[] result = getPoolRunningTasks(pool);
      System.out.println(pool + " " + result[0] + " " + result[1]);
      return 0;
    }
    if (arg[curr].equalsIgnoreCase("-poolmaxtasks")) {
      curr += 1;
      if (arg.length == curr) {
        printUsage(arg[0]);
        return -1;
      }
      String pool = arg[curr];
      int[] result = getPoolMaxTasks(pool);
      System.out.println(pool + " " + result[0] + " " + result[1]);
      return 0;
    }
    printUsage("");
    return -1;
  }

  public static void main(String argv[]) throws Exception {
    FairSchedulerShell shell = null;
    try {
      shell = new FairSchedulerShell();
      int res = ToolRunner.run(shell, argv);
      System.exit(res);
    } catch (RPC.VersionMismatch v) {
      System.err.println("Version Mismatch between client and server" +
                         "... command aborted.");
      System.exit(-1);
    } catch (IOException e) {
      System.err.
        println("Bad connection to FairScheduler. command aborted.");
      System.err.println(e.getMessage());
      System.exit(-1);
    }
  }

}