From f3e867ed975121417e74956017d44a4fcde3cbd0 Mon Sep 17 00:00:00 2001
From: tqchen <tianqi.tchen@gmail.com>
Date: Sun, 15 Mar 2015 22:38:51 -0700
Subject: [PATCH] add option queue

---
 tracker/rabit_yarn.py                             | 3 +++
 yarn/build.sh                                     | 4 ++++
 yarn/src/org/apache/hadoop/yarn/rabit/Client.java | 6 +++++-
 3 files changed, 12 insertions(+), 1 deletion(-)

diff --git a/tracker/rabit_yarn.py b/tracker/rabit_yarn.py
index 9af56d4e3..9b2141559 100755
--- a/tracker/rabit_yarn.py
+++ b/tracker/rabit_yarn.py
@@ -39,6 +39,8 @@ parser.add_argument('-hip', '--host_ip', default='auto', type=str,
                     help = 'host IP address if cannot be automatically guessed, specify the IP of submission machine')
 parser.add_argument('-v', '--verbose', default=0, choices=[0, 1], type=int,
                     help = 'print more messages into the console')
+parser.add_argument('-q', '--queue', default='default', type=str,
+                    help = 'the queue we want to submit the job to')
 parser.add_argument('-ac', '--auto_file_cache', default=1, choices=[0, 1], type=int,
                     help = 'whether automatically cache the files in the command to hadoop localfile, this is on by default')
 parser.add_argument('-f', '--files', default = [], action='append',
@@ -122,6 +124,7 @@ def submit_yarn(nworker, worker_args, worker_env):
         cmd += ' -file %s' % f
     cmd += ' -jobname %s ' % args.jobname
     cmd += ' -tempdir %s ' % args.tempdir
+    cmd += ' -queue %s ' % args.queue
     cmd += (' '.join(['./run_hdfs_prog.py'] + args.command + worker_args))
     if args.verbose != 0:
         print cmd
diff --git a/yarn/build.sh b/yarn/build.sh
index 12764357c..8908cafdd 100755
--- a/yarn/build.sh
+++ b/yarn/build.sh
@@ -1,4 +1,8 @@
 #!/bin/bash
+if [ ! -d bin ]; then
+    mkdir bin
+fi
+
 CPATH=`${HADOOP_HOME}/bin/hadoop classpath`
 javac -cp $CPATH -d bin src/org/apache/hadoop/yarn/rabit/*
 jar cf rabit-yarn.jar -C bin . 
diff --git a/yarn/src/org/apache/hadoop/yarn/rabit/Client.java b/yarn/src/org/apache/hadoop/yarn/rabit/Client.java
index 40470fe19..4d252d76a 100644
--- a/yarn/src/org/apache/hadoop/yarn/rabit/Client.java
+++ b/yarn/src/org/apache/hadoop/yarn/rabit/Client.java
@@ -48,6 +48,8 @@ public class Client {
     private String userName = "";
     // job name
     private String jobName = "";
+    // queue
+    private String queue = "default";
     /**
      * constructor
      * @throws IOException
@@ -158,6 +160,8 @@ public class Client {
                 this.jobName = args[++i];
             } else if(args[i].equals("-tempdir")) {
                 this.tempdir = args[++i];
+            } else if(args[i].equals("-queue")) {
+                this.queue = args[++i];
             } else {
                 sargs.append(" ");
                 sargs.append(args[i]);
@@ -210,7 +214,7 @@ public class Client {
         appContext.setApplicationName(jobName + ":RABIT-YARN");
         appContext.setAMContainerSpec(amContainer);
         appContext.setResource(capability);
-        appContext.setQueue("default");
+        appContext.setQueue(queue);
         //appContext.setUser(userName);
         LOG.info("Submitting application " + appId);      
         yarnClient.submitApplication(appContext);