From ab7492dbc27e05b8f2178644ce3f4b6f6e636767 Mon Sep 17 00:00:00 2001
From: tqchen <tianqi.tchen@gmail.com>
Date: Mon, 22 Dec 2014 03:24:00 -0800
Subject: [PATCH] add support for yarn

---
 src/allreduce_base.cc | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/src/allreduce_base.cc b/src/allreduce_base.cc
index 6fb4a11cb..30fcfa49f 100644
--- a/src/allreduce_base.cc
+++ b/src/allreduce_base.cc
@@ -43,6 +43,9 @@ void AllreduceBase::Init(void) {
     }
     // handling for hadoop
     const char *num_task = getenv("mapred_map_tasks");
+    if (num_task == NULL) {
+      num_task = getenv("mapreduce_job_maps");
+    }
     if (hadoop_mode != 0) {
       utils::Check(num_task != NULL, "hadoop_mode is set but cannot find mapred_map_tasks");      
     }