Profiling tool: Print UCX and GDS parameters

tools/src/main/scala/com/nvidia/spark/rapids/tool/profiling/CollectInformation.scala

@@ -98,7 +98,7 @@ class CollectInformation(apps: Seq[ApplicationInfo],
     val messageHeader = "\nSpark Rapids parameters set explicitly:\n"
     for (app <- apps) {
       if (app.allDataFrames.contains(s"propertiesDF_${app.index}")) {
-        app.runQuery(query = app.generateRapidsProperties + " order by key",
+        app.runQuery(query = app.generateRapidsUcxGdsProperties + " order by key",
           fileWriter = fileWriter, messageHeader = messageHeader)
       } else {
         fileWriter.foreach(_.write("No Spark Rapids parameters Found!\n"))

tools/src/main/scala/com/nvidia/spark/rapids/tool/profiling/CompareApplications.scala

@@ -92,20 +92,20 @@ class CompareApplications(apps: Seq[ApplicationInfo],
       // For the 1st app
       if (i == 1) {
         withClauseAllKeys += "select distinct key from (" +
-            app.generateRapidsProperties + ") union "
-        query += "(" + app.generateRapidsProperties + s") tmp_$i"
+            app.generateRapidsUcxGdsProperties + ") union "
+        query += "(" + app.generateRapidsUcxGdsProperties + s") tmp_$i"
         query += s" on allKeys.key=tmp_$i.key"
         query += "\n LEFT OUTER JOIN \n"
       } else if (i < apps.size) { // For the 2nd to non-last app(s)
         withClauseAllKeys += "select distinct key from (" +
-            app.generateRapidsProperties + ") union "
-        query += "(" + app.generateRapidsProperties + s") tmp_$i"
+            app.generateRapidsUcxGdsProperties + ") union "
+        query += "(" + app.generateRapidsUcxGdsProperties + s") tmp_$i"
         query += s" on allKeys.key=tmp_$i.key"
         query += "\n LEFT OUTER JOIN \n"
       } else { // For the last app
         withClauseAllKeys += "select distinct key from (" +
-            app.generateRapidsProperties + "))\n"
-        query += "(" + app.generateRapidsProperties + s") tmp_$i"
+            app.generateRapidsUcxGdsProperties + "))\n"
+        query += "(" + app.generateRapidsUcxGdsProperties + s") tmp_$i"
         query += s" on allKeys.key=tmp_$i.key"
       }
       selectValuePart += s",value_app$i"

tools/src/main/scala/org/apache/spark/sql/rapids/tool/profiling/ApplicationInfo.scala

@@ -649,12 +649,14 @@ class ApplicationInfo(
     }
   }
 
-  // Function to generate a query for printing Rapids related Spark properties
-  def generateRapidsProperties: String =
+  // Function to generate a query for printing Rapids/UCX/GDS related Spark properties
+  def generateRapidsUcxGdsProperties: String =
     s"""select key,value as value_app$index
        |from propertiesDF_$index
        |where source ='spark'
        |and key like 'spark.rapids%'
+       |or key like 'spark.executorEnv.UCX%'
+       |or key in ('spark.shuffle.manager','spark.shuffle.service.enabled')
        |""".stripMargin
 
   // Function to generate the SQL string for aggregating task metrics columns.

tools/src/test/scala/com/nvidia/spark/rapids/tool/profiling/ApplicationInfoSuite.scala

@@ -322,4 +322,28 @@ class ApplicationInfoSuite extends FunSuite with Logging {
       tempFile4.delete()
     }
   }
+
+  test("test gds-ucx-parameters") {
+    val apps: ArrayBuffer[ApplicationInfo] = ArrayBuffer[ApplicationInfo]()
+    val appArgs =
+      new ProfileArgs(Array(s"$logDir/gds_ucx_eventlog.zstd"))
+    var index: Int = 1
+    val eventlogPaths = appArgs.eventlog()
+    for (path <- eventlogPaths) {
+      apps += new ApplicationInfo(appArgs.numOutputRows.getOrElse(1000), sparkSession,
+        EventLogPathProcessor.getEventLogInfo(path, sparkSession).head._1, index)
+      index += 1
+    }
+    assert(apps.size == 1)
+    for (app <- apps) {
+      val rows = app.runQuery(query = app.generateRapidsUcxGdsProperties + " order by key",
+        fileWriter = None).collect()
+      assert(rows.length == 22) // 22 properties captured.
+      // verify  ucx parameters are captured.
+      assert(rows(1)(0).equals("spark.executorEnv.UCX_MAX_RNDV_RAILS"))
+
+      //verify gds parameters are captured.
+      assert(rows(11)(0).equals("spark.rapids.memory.gpu.direct.storage.spill.useHostMemory"))
+    }
+  }
 }