From 130f295e085d95e8205d882174a5667d29b3b1f2 Mon Sep 17 00:00:00 2001
From: Andrew Or <andrewor14@gmail.com>
Date: Tue, 5 Aug 2014 22:12:28 -0700
Subject: [PATCH] Handle spark.driver.memory too

---
 bin/spark-submit | 24 +++++++++++++++---------
 1 file changed, 15 insertions(+), 9 deletions(-)

diff --git a/bin/spark-submit b/bin/spark-submit
index 3f4a3840b4d9b..e86032e006dd6 100755
--- a/bin/spark-submit
+++ b/bin/spark-submit
@@ -46,30 +46,36 @@ PROPERTIES_FILE=${PROPERTIES_FILE:-"$SPARK_HOME/conf/spark-defaults.conf"}
 # time the JVM has started.
 
 if [ $DEPLOY_MODE == "client" ]; then
-  if [ -n "$DRIVER_MEMORY" ]; then
-    export SPARK_DRIVER_MEMORY=$DRIVER_MEMORY
-  fi
   # We parse the default properties file here, assuming each line is
   # a key value pair delimited either by white space or "=" sign. All
   # spark.driver.* configs must be processed now before it's too late.
   if [ -f "$PROPERTIES_FILE" ]; then
     echo "Using properties file $PROPERTIES_FILE"
+    DRIVER_MEMORY_CONF="spark.driver.memory"
     DRIVER_EXTRA_JAVA_OPTS="spark.driver.extraJavaOptions"
     DRIVER_EXTRA_CLASSPATH="spark.driver.extraClassPath"
     DRIVER_EXTRA_LIBRARY_PATH="spark.driver.extraLibraryPath"
     # Remove "=" sign and double quotes around the value, if any
-    DRIVER_EXTRA_JAVA_OPTS=$(sed "/^#/ d" "$PROPERTIES_FILE" | grep "$DRIVER_EXTRA_JAVA_OPTS" | \
-      sed "s/$DRIVER_EXTRA_JAVA_OPTS//g" | sed "s/^=//g" | sed "s/^\"\(.*\)\"$/\1/g")
-    DRIVER_EXTRA_CLASSPATH=$(sed "/^#/ d" "$PROPERTIES_FILE" | grep "$DRIVER_EXTRA_CLASSPATH" | \
-      sed "s/$DRIVER_EXTRA_CLASSPATH//g" | sed "s/^=//g" | sed "s/^\"\(.*\)\"$/\1/g")
-    DRIVER_EXTRA_LIBRARY_PATH=$(sed "/^#/ d" "$PROPERTIES_FILE" | grep "$DRIVER_EXTRA_LIBRARY_PATH" | \
-      sed "s/$DRIVER_EXTRA_LIBRARY_PATH//g" | sed "s/^=//g" | sed "s/^\"\(.*\)\"$/\1/g")
+    DRIVER_MEMORY_CONF=$(sed "/^#/ d" "$PROPERTIES_FILE" | grep "$DRIVER_MEMORY_CONF" | sed "s/$DRIVER_MEMORY_CONF//g" | \
+      sed "s/^=//g" | sed "s/^\"\(.*\)\"$/\1/g" | sed "s/^[[:space:]]*//g" | sed "s/[[:space:]]*$//g")
+    DRIVER_EXTRA_JAVA_OPTS=$(sed "/^#/ d" "$PROPERTIES_FILE" | grep "$DRIVER_EXTRA_JAVA_OPTS" | sed "s/$DRIVER_EXTRA_JAVA_OPTS//g" | \
+      sed "s/^=//g" | sed "s/^\"\(.*\)\"$/\1/g" | sed "s/^[[:space:]]*//g" | sed "s/[[:space:]]*$//g")
+    DRIVER_EXTRA_CLASSPATH=$(sed "/^#/ d" "$PROPERTIES_FILE" | grep "$DRIVER_EXTRA_CLASSPATH" | sed "s/$DRIVER_EXTRA_CLASSPATH//g" | \
+      sed "s/^=//g" | sed "s/^\"\(.*\)\"$/\1/g" | sed "s/^[[:space:]]*//g" | sed "s/[[:space:]]*$//g")
+    DRIVER_EXTRA_LIBRARY_PATH=$(sed "/^#/ d" "$PROPERTIES_FILE" | grep "$DRIVER_EXTRA_LIBRARY_PATH" | sed "s/$DRIVER_EXTRA_LIBRARY_PATH//g" | \
+      sed "s/^=//g" | sed "s/^\"\(.*\)\"$/\1/g" | sed "s/^[[:space:]]*//g" | sed "s/[[:space:]]*$//g")
     export SPARK_SUBMIT_OPTS="$SPARK_SUBMIT_OPTS $DRIVER_EXTRA_JAVA_OPTS"
     export SPARK_SUBMIT_CLASSPATH="$SPARK_SUBMIT_CLASSPATH:$DRIVER_EXTRA_CLASSPATH"
     export SPARK_SUBMIT_LIBRARY_PATH="$SPARK_SUBMIT_LIBRARY_PATH:$DRIVER_EXTRA_LIBRARY_PATH"
   else
     echo "Warning: properties file $PROPERTIES_FILE does not exist."
   fi
+
+  # Favor command line memory over config memory
+  DRIVER_MEMORY=${DRIVER_MEMORY:-"$DRIVER_MEMORY_CONF"}
+  if [ -n "$DRIVER_MEMORY" ]; then
+    export SPARK_DRIVER_MEMORY=$DRIVER_MEMORY
+  fi
 fi
 
 exec $SPARK_HOME/bin/spark-class org.apache.spark.deploy.SparkSubmit "${ORIG_ARGS[@]}"