Add default jvm opts to tony jobtype and tony core (#629)

* Add default jvm opts to tony jobtype and tony core Signed-off-by: Frank Gu <cgu@linkedin.com> * Release TonY 0.4.13 Signed-off-by: Frank Gu <cgu@linkedin.com>
tony-framework · Dec 17, 2021 · 42193b9 · 42193b9
1 parent d198bed
commit 42193b9
Show file tree

Hide file tree

Showing 7 changed files with 45 additions and 7 deletions.
diff --git a/build.gradle b/build.gradle
@@ -92,7 +92,7 @@ ext.deps = [
 
 allprojects {
   group = "com.linkedin.tony"
-  project.version = "0.4.12"
+  project.version = "0.4.13"
 }
 
 task sourcesJar(type: Jar) {

diff --git a/tony-azkaban/src/main/java/com/linkedin/tony/azkaban/TonyJob.java b/tony-azkaban/src/main/java/com/linkedin/tony/azkaban/TonyJob.java
@@ -38,6 +38,9 @@ public class TonyJob extends HadoopJavaJob {
   private static final String TONY_CONF_PREFIX = "tony.";
   public static final String TONY_APPLICATION_TAGS =
       TONY_CONF_PREFIX + "application.tags";
+  public static final String TONY_DEFAULT_JVM_OPTS = TONY_CONF_PREFIX + "default.jvm.opts";
+  public static final String TONY_TASK_AM_JVM_OPTS = TONY_CONF_PREFIX + "task.am.jvm.opts";
+  public static final String TONY_TASK_EXECUTOR_JVM_OPTS = TONY_CONF_PREFIX + "task.executor.jvm.opts";
   private String tonyXml;
   private File tonyConfFile;
   private final Configuration tonyConf;
@@ -64,9 +67,37 @@ private Configuration getJobConfiguration() {
     String applicationTags =
         HadoopJobUtils.constructHadoopTags(getJobProps(), tagKeys);
     tonyConf.set(TONY_APPLICATION_TAGS, applicationTags);
+
+    // inject default jvm options
+    updateDefaultTonyJvmOpts(tonyConf);
     return tonyConf;
   }
 
+  private void updateDefaultTonyJvmOpts(Configuration conf) {
+    String defaultTonyJvmOpts = getJobProps().get(TONY_DEFAULT_JVM_OPTS);
+    if (defaultTonyJvmOpts == null) {
+      return;
+    }
+
+    // tony am jvm opts
+    String tonyTaskAmJvmOpts = conf.get(TONY_TASK_AM_JVM_OPTS, "");
+    if (!tonyTaskAmJvmOpts.isEmpty()) {
+      tonyTaskAmJvmOpts += " " + defaultTonyJvmOpts;
+    } else {
+      tonyTaskAmJvmOpts += defaultTonyJvmOpts;
+    }
+    conf.set(TONY_TASK_AM_JVM_OPTS, tonyTaskAmJvmOpts);
+
+    // tony executor jvm opts
+    String tonyTaskExecutorJvmOpts = conf.get(TONY_TASK_EXECUTOR_JVM_OPTS, "");
+    if (!tonyTaskExecutorJvmOpts.isEmpty()) {
+      tonyTaskExecutorJvmOpts += " " + defaultTonyJvmOpts;
+    } else {
+      tonyTaskExecutorJvmOpts += defaultTonyJvmOpts;
+    }
+    conf.set(TONY_TASK_EXECUTOR_JVM_OPTS, tonyTaskExecutorJvmOpts);
+  }
+
   public Configuration getTonyJobConf() {
     return tonyConf;
   }

diff --git a/tony-azkaban/src/test/java/com/linkedin/tony/azkaban/TestTonyJob.java b/tony-azkaban/src/test/java/com/linkedin/tony/azkaban/TestTonyJob.java
@@ -95,6 +95,8 @@ public void testFlowInfoPropagation() {
     jobProps.put(CommonJobProperties.FLOW_ID, "1");
     jobProps.put(CommonJobProperties.EXEC_ID, "0");
     jobProps.put(TonyJob.AZKABAN_WEB_HOST, "localhost");
+    jobProps.put(TonyJob.TONY_TASK_EXECUTOR_JVM_OPTS, "-Duser.jvm.opts=opts");
+    jobProps.put(TonyJob.TONY_DEFAULT_JVM_OPTS, "-Dlog4j2.formatMsgNoLookups=true");
 
     final TonyJob tonyJob = new TonyJob("test_tony_job", new Props(), jobProps, log) {
       @Override
@@ -119,6 +121,8 @@ public String getWorkingDirectory() {
     Assert.assertEquals(parsedTags.get(CommonJobProperties.FLOW_ID), "1");
     Assert.assertEquals(parsedTags.get(CommonJobProperties.PROJECT_NAME), "unit_test");
     Assert.assertEquals(parsedTags.get(TonyJob.AZKABAN_WEB_HOST), "localhost");
+    Assert.assertEquals(conf.get(TonyJob.TONY_TASK_AM_JVM_OPTS), "-Dlog4j2.formatMsgNoLookups=true");
+    Assert.assertEquals(conf.get(TonyJob.TONY_TASK_EXECUTOR_JVM_OPTS), "-Duser.jvm.opts=opts -Dlog4j2.formatMsgNoLookups=true");
   }
 
   @Test

diff --git a/tony-core/src/main/java/com/linkedin/tony/TonyClient.java b/tony-core/src/main/java/com/linkedin/tony/TonyClient.java
@@ -933,6 +933,10 @@ String buildCommand(long amMemory) {
     if (org.apache.commons.lang3.StringUtils.isNotBlank(amJvm)) {
       arguments.add(amJvm);
     }
+    String defaultJvm = tonyConf.get(TonyConfigurationKeys.TASK_DEFAULT_JVM_OPTS, "");
+    if (!defaultJvm.isEmpty()) {
+      arguments.add(defaultJvm);
+    }
     // Set class name
     arguments.add("com.linkedin.tony.ApplicationMaster");
 

diff --git a/tony-core/src/main/java/com/linkedin/tony/TonyConfigurationKeys.java b/tony-core/src/main/java/com/linkedin/tony/TonyConfigurationKeys.java
@@ -155,7 +155,7 @@ private TonyConfigurationKeys() {
   public static final String TASK_EXECUTOR_JVM_OPTS = TONY_TASK_PREFIX + "executor.jvm.opts";
   public static final String DEFAULT_TASK_EXECUTOR_JVM_OPTS = "-Xmx1536m";
 
-  public static final String TASK_EXECUTOR_JAVA_AGENT = TONY_TASK_PREFIX + "executor.java.agent";
+  public static final String TASK_DEFAULT_JVM_OPTS = TONY_TASK_PREFIX + "default.jvm.opts";
 
   public static final String TASK_HEARTBEAT_INTERVAL_MS = TONY_TASK_PREFIX + "heartbeat-interval-ms";
   public static final int DEFAULT_TASK_HEARTBEAT_INTERVAL_MS = 1000;

diff --git a/tony-core/src/main/java/com/linkedin/tony/TonySession.java b/tony-core/src/main/java/com/linkedin/tony/TonySession.java
@@ -440,10 +440,9 @@ public Builder setTonyConf(Configuration tonyConf) {
 
     // Appends default jvm arguments to task executor job
     private String appendDefaultJVMArgs(String jvmArgs) {
-      jvmArgs += " -Dlog4j2.formatMsgNoLookups=true";
-      String tonyTaskExecutorJavaAgent = tonyConf.get(TonyConfigurationKeys.TASK_EXECUTOR_JAVA_AGENT, "");
-      if (!tonyTaskExecutorJavaAgent.isEmpty()) {
-        jvmArgs += " -javaagent:" + tonyTaskExecutorJavaAgent;
+      String tonyTaskDefaultJVMOpts = tonyConf.get(TonyConfigurationKeys.TASK_DEFAULT_JVM_OPTS, "");
+      if (!tonyTaskDefaultJVMOpts.isEmpty()) {
+        jvmArgs += " " + tonyTaskDefaultJVMOpts;
       }
       return jvmArgs;
     }

diff --git a/tony-core/src/test/java/com/linkedin/tony/TestTonyConfigurationFields.java b/tony-core/src/test/java/com/linkedin/tony/TestTonyConfigurationFields.java
@@ -49,7 +49,7 @@ public void initializeMemberVariables() {
     configurationPropsToSkipCompare.add(TonyConfigurationKeys.APPLICATION_TRAINING_STAGE);
     configurationPropsToSkipCompare.add(TonyConfigurationKeys.APPLICATION_HADOOP_LOCATION);
     configurationPropsToSkipCompare.add(TonyConfigurationKeys.APPLICATION_HADOOP_CLASSPATH);
-    configurationPropsToSkipCompare.add(TonyConfigurationKeys.TASK_EXECUTOR_JAVA_AGENT);
+    configurationPropsToSkipCompare.add(TonyConfigurationKeys.TASK_DEFAULT_JVM_OPTS);
     configurationPropsToSkipCompare.add(TonyConfigurationKeys.TENSORBOARD_LOG_DIR);
     configurationPropsToSkipCompare.add(TonyConfigurationKeys.PYTHON_EXEC_PATH);
     configurationPropsToSkipCompare.add(TonyConfigurationKeys.TB_VCORE);