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

UWFrankGu · web-flow · commit 42193b926199 · 2021-12-17T14:13:57.000-08:00
* Add default jvm opts to tony jobtype and tony core

Signed-off-by: Frank Gu &lt;cgu@linkedin.com&gt;

* Release TonY 0.4.13

Signed-off-by: Frank Gu &lt;cgu@linkedin.com&gt;
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);

Original file line number	Diff line number	Diff line change
`@@ -92,7 +92,7 @@ ext.deps = [`
`92`	`92`
`93`	`93`	`allprojects {`
`94`	`94`	`group = "com.linkedin.tony"`
`95`		`- project.version = "0.4.12"`
	`95`	`+ project.version = "0.4.13"`
`96`	`96`	`}`
`97`	`97`
`98`	`98`	`task sourcesJar(type: Jar) {`
Original file line number	Diff line number	Diff line change
`@@ -440,10 +440,9 @@ public Builder setTonyConf(Configuration tonyConf) {`
`440`	`440`
`441`	`441`	`// Appends default jvm arguments to task executor job`
`442`	`442`	`private String appendDefaultJVMArgs(String jvmArgs) {`
`443`		`- jvmArgs += " -Dlog4j2.formatMsgNoLookups=true";`
`444`		`- String tonyTaskExecutorJavaAgent = tonyConf.get(TonyConfigurationKeys.TASK_EXECUTOR_JAVA_AGENT, "");`
`445`		`- if (!tonyTaskExecutorJavaAgent.isEmpty()) {`
`446`		`- jvmArgs += " -javaagent:" + tonyTaskExecutorJavaAgent;`
	`443`	`+ String tonyTaskDefaultJVMOpts = tonyConf.get(TonyConfigurationKeys.TASK_DEFAULT_JVM_OPTS, "");`
	`444`	`+ if (!tonyTaskDefaultJVMOpts.isEmpty()) {`
	`445`	`+ jvmArgs += " " + tonyTaskDefaultJVMOpts;`
`447`	`446`	`}`
`448`	`447`	`return jvmArgs;`
`449`	`448`	`}`