TEZ-4564: TezClient to expose Tez AM host:port

apache · May 16, 2024 · cefab38 · cefab38
1 parent a1fcddb
commit cefab38
Show file tree

Hide file tree

Showing 4 changed files with 30 additions and 0 deletions.
diff --git a/tez-api/src/main/java/org/apache/tez/client/FrameworkClient.java b/tez-api/src/main/java/org/apache/tez/client/FrameworkClient.java
@@ -55,6 +55,9 @@
 public abstract class FrameworkClient {
   protected static final Logger LOG = LoggerFactory.getLogger(FrameworkClient.class);
 
+  protected String amHost = null;
+  protected int amPort = -1;
+
   public static FrameworkClient createFrameworkClient(TezConfiguration tezConf) {
 
     boolean isLocal = tezConf.getBoolean(TezConfiguration.TEZ_LOCAL_MODE, TezConfiguration.TEZ_LOCAL_MODE_DEFAULT);
@@ -94,6 +97,20 @@ public abstract ApplicationId submitApplication(ApplicationSubmissionContext app
 
   public abstract ApplicationReport getApplicationReport(ApplicationId appId) throws YarnException, IOException;
 
+  public String getAmHost() {
+    if (amHost == null) {
+      throw new IllegalStateException("AM Host is not set when called, please check FrameworkClient implementation");
+    }
+    return amHost;
+  }
+
+  public int getAmPort() {
+    if (amPort == -1) {
+      throw new IllegalStateException("AM Port is not set when called, please check FrameworkClient implementation");
+    }
+    return amPort;
+  }
+
   public abstract boolean isRunning() throws IOException;
 
   public TezAppMasterStatus getAMStatus(Configuration conf, ApplicationId appId,

diff --git a/tez-api/src/main/java/org/apache/tez/client/TezClient.java b/tez-api/src/main/java/org/apache/tez/client/TezClient.java
@@ -1286,4 +1286,12 @@ public static ApplicationId appIdfromString(String appIdStr) {
               + appIdStr, n);
     }
   }
+
+  public String getAmHost() {
+    return frameworkClient.getAmHost();
+  }
+
+  public int getAmPort() {
+    return frameworkClient.getAmPort();
+  }
 }
diff --git a/tez-api/src/main/java/org/apache/tez/client/TezYarnClient.java b/tez-api/src/main/java/org/apache/tez/client/TezYarnClient.java
@@ -100,6 +100,8 @@ public ApplicationReport getApplicationReport(ApplicationId appId) throws YarnEx
       throw new ApplicationNotFoundException("YARN reports no state for application "
           + appId);
     }
+    this.amHost = report.getHost();
+    this.amPort = report.getRpcPort();
     return report;
   }
 

diff --git a/tez-dag/src/main/java/org/apache/tez/client/LocalClient.java b/tez-dag/src/main/java/org/apache/tez/client/LocalClient.java
@@ -204,6 +204,9 @@ public ApplicationReport getApplicationReport(ApplicationId appId) {
     report.setProgress(dagAppMaster.getProgress());
     report.setAMRMToken(null);
 
+    this.amHost = dagAppMaster.getAppNMHost();
+    this.amPort = dagAppMaster.getRpcPort();
+
     return report;
   }