upgrade yarn scheduler driver memory to ByteAmount

heron/common/src/java/com/twitter/heron/common/basics/ByteAmountUnit.java

@@ -0,0 +1,21 @@
+// Copyright 2017 Twitter. All rights reserved.
+//
+// Licensed 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 com.twitter.heron.common.basics;
+
+public enum ByteAmountUnit {
+  BYTE,
+  MB,
+  GB
+}

heron/common/src/java/com/twitter/heron/common/basics/TypeUtils.java

@@ -81,7 +81,28 @@ public static ByteAmount getByteAmount(Object o) {
       return (ByteAmount) o;
     }
     try {
-      return ByteAmount.fromBytes(getLong(o));
+      return getByteAmount(o, ByteAmountUnit.BYTE);
+    } catch (IllegalArgumentException e) {
+      throw new IllegalArgumentException("Don't know how to convert " + o + " to ByteAmount", e);
+    }
+  }
+
+  public static ByteAmount getByteAmount(Object o, ByteAmountUnit unit) {
+    if (o != null && o instanceof ByteAmount) {
+      return (ByteAmount) o;
+    }
+    try {
+      long amount = getLong(o);
+      switch (unit) {
+        case BYTE:
+          return ByteAmount.fromBytes(amount);
+        case MB:
+          return ByteAmount.fromMegabytes(amount);
+        case GB:
+          return ByteAmount.fromGigabytes(amount);
+        default:
+          return ByteAmount.fromBytes(amount);
+      }
     } catch (IllegalArgumentException e) {
       throw new IllegalArgumentException("Don't know how to convert " + o + " to ByteAmount", e);
     }

heron/schedulers/src/java/com/twitter/heron/scheduler/yarn/HeronMasterDriver.java

@@ -99,9 +99,7 @@
  */
 @Unit
 public class HeronMasterDriver {
-  static final int TM_MEM_SIZE_MB = 1024;
   static final int TMASTER_CONTAINER_ID = 0;
-  static final int MB = 1024 * 1024;
   private static final Logger LOG = Logger.getLogger(HeronMasterDriver.class.getName());
   private final String topologyPackageName;
   private final String heronCorePackageName;

heron/schedulers/src/java/com/twitter/heron/scheduler/yarn/YarnContext.java

@@ -14,6 +14,8 @@
 
 package com.twitter.heron.scheduler.yarn;
 
+import com.twitter.heron.common.basics.ByteAmount;
+import com.twitter.heron.common.basics.ByteAmountUnit;
 import com.twitter.heron.spi.common.Config;
 import com.twitter.heron.spi.common.Context;
 
@@ -27,8 +29,8 @@ public static String heronYarnQueue(Config cfg) {
         YarnKey.HERON_SCHEDULER_YARN_QUEUE.getDefaultString());
   }
 
-  public static int heronDriverMemoryMb(Config cfg) {
-    return cfg.getIntegerValue(YarnKey.YARN_SCHEDULER_DRIVER_MEMORY_MB.value(),
-        YarnKey.YARN_SCHEDULER_DRIVER_MEMORY_MB.getDefaultInt());
+  public static ByteAmount heronDriverMemory(Config cfg) {
+    return cfg.getByteAmountValue(YarnKey.YARN_SCHEDULER_DRIVER_MEMORY_MB.value(),
+        (ByteAmount) YarnKey.YARN_SCHEDULER_DRIVER_MEMORY_MB.getDefault(), ByteAmountUnit.MB);
   }
 }

heron/schedulers/src/java/com/twitter/heron/scheduler/yarn/YarnKey.java

@@ -14,6 +14,7 @@
 
 package com.twitter.heron.scheduler.yarn;
 
+import com.twitter.heron.common.basics.ByteAmount;
 import com.twitter.heron.spi.common.Key;
 
 /**
@@ -23,7 +24,8 @@ public enum YarnKey {
   // yarn queue for submitting and launching the topology
   HERON_SCHEDULER_YARN_QUEUE("heron.scheduler.yarn.queue", "default"),
   // the amount of memory topology's driver (yarn application master) needs
-  YARN_SCHEDULER_DRIVER_MEMORY_MB("heron.scheduler.yarn.driver.memory.mb", 2048);
+  YARN_SCHEDULER_DRIVER_MEMORY_MB("heron.scheduler.yarn.driver.memory.mb",
+      ByteAmount.fromMegabytes(2048));
 
   private final String value;
   private final Key.Type type;
@@ -35,9 +37,9 @@ public enum YarnKey {
     this.defaultValue = defaultValue;
   }
 
-  YarnKey(String value, Integer defaultValue) {
+  YarnKey(String value, ByteAmount defaultValue) {
     this.value = value;
-    this.type = Key.Type.INTEGER;
+    this.type = Key.Type.BYTE_AMOUNT;
     this.defaultValue = defaultValue;
   }
 
@@ -56,12 +58,4 @@ public String getDefaultString() {
     }
     return (String) this.defaultValue;
   }
-
-  public int getDefaultInt() {
-    if (type != Key.Type.INTEGER) {
-      throw new IllegalAccessError(String.format(
-          "Config Key %s is type %s, getDefaultInt() not supported", this.name(), this.type));
-    }
-    return (Integer) this.defaultValue;
-  }
 }

heron/schedulers/src/java/com/twitter/heron/scheduler/yarn/YarnLauncher.java

@@ -34,6 +34,7 @@
 import org.apache.reef.tang.annotations.Unit;
 import org.apache.reef.tang.exceptions.InjectionException;
 
+import com.twitter.heron.common.basics.ByteAmount;
 import com.twitter.heron.scheduler.yarn.HeronMasterDriver.ContainerAllocationHandler;
 import com.twitter.heron.scheduler.yarn.HeronMasterDriver.FailedContainerHandler;
 import com.twitter.heron.scheduler.yarn.HeronMasterDriver.HeronSchedulerLauncher;
@@ -63,7 +64,7 @@ public class YarnLauncher implements ILauncher {
   private String role;
   private String env;
   private String queue;
-  private int driverMemory;
+  private ByteAmount driverMemory;
   private ArrayList<String> libJars = new ArrayList<>();
 
   @Override
@@ -75,7 +76,7 @@ public void initialize(Config config, Config runtime) {
     role = Context.role(config);
     env = Context.environ(config);
     queue = YarnContext.heronYarnQueue(config);
-    driverMemory = YarnContext.heronDriverMemoryMb(config);
+    driverMemory = YarnContext.heronDriverMemory(config);
 
     try {
       // In addition to jar for REEF's driver implementation, jar for packing and state manager
@@ -161,7 +162,7 @@ Configuration getHMDriverConf() {
         .set(HeronDriverConfiguration.HTTP_PORT, 0)
         .set(HeronDriverConfiguration.VERBOSE, false)
         .set(YarnDriverConfiguration.QUEUE, queue)
-        .set(DriverConfiguration.DRIVER_MEMORY, driverMemory)
+        .set(DriverConfiguration.DRIVER_MEMORY, (int) driverMemory.asMegabytes())
         .build();
   }
 

heron/spi/src/java/com/twitter/heron/spi/common/Config.java

@@ -21,19 +21,20 @@
 import java.util.logging.Logger;
 
 import com.twitter.heron.common.basics.ByteAmount;
+import com.twitter.heron.common.basics.ByteAmountUnit;
 import com.twitter.heron.common.basics.DryRunFormatType;
 import com.twitter.heron.common.basics.PackageType;
 import com.twitter.heron.common.basics.TypeUtils;
 
 /**
  * Config is an Immutable Map of <String, Object> The get/set API that uses Key objects
  * should be favored over Strings. Usage of the String API should be refactored out.
- *
+ * <p>
  * A newly created Config object holds configs that might include wildcard tokens, like
  * ${HERON_HOME}/bin, ${HERON_LIB}/packing/*. Token substitution can be done by converting that
  * config to a local or cluster config by using the {@code Config.toLocalMode} or
  * {@code Config.toClusterMode} methods.
- *
+ * <p>
  * Local mode is for a config to be used to run Heron locally, where HERON_HOME might be an install
  * dir on the local host (e.g. HERON_HOME=/usr/bin/heron). Cluster mode is to be used when building
  * configs for a remote process run on a service, where all directories might be relative to the
@@ -103,11 +104,11 @@ private static Config expand(Config config) {
    * Recursively expand each config value until token substitution is exhausted. We must recurse
    * to handle the case where field expansion requires multiple iterations, due to new tokens being
    * introduced as we replace. For example:
-   *
-   *   ${HERON_BIN}/heron-executor        gets expanded to
-   *   ${HERON_HOME}/bin/heron-executor   gets expanded to
-   *   /usr/local/heron/bin/heron-executor
-   *
+   * <p>
+   * ${HERON_BIN}/heron-executor        gets expanded to
+   * ${HERON_HOME}/bin/heron-executor   gets expanded to
+   * /usr/local/heron/bin/heron-executor
+   * <p>
    * If break logic is when another round does not reduce the number of tokens, since it means we
    * couldn't find a valid replacement.
    */
@@ -240,6 +241,14 @@ public ByteAmount getByteAmountValue(Key key) {
     return TypeUtils.getByteAmount(value);
   }
 
+  public ByteAmount getByteAmountValue(String key, ByteAmount defaultValue, ByteAmountUnit unit) {
+    Object value = get(key);
+    if (value != null) {
+      return TypeUtils.getByteAmount(value, unit);
+    }
+    return defaultValue;
+  }
+
   DryRunFormatType getDryRunFormatType(Key key) {
     return (DryRunFormatType) get(key);
   }