Support running nds_power over spark connect

Signed-off-by: Bobby Wang <[email protected]>

Author: wbo4958

nds/PysparkBenchReport.py

@@ -35,9 +35,8 @@
 import time
 import traceback
 from typing import Callable
-from pyspark.sql import SparkSession
 
-import python_listener
+from pyspark.sql import SparkSession
 
 class PysparkBenchReport:
     """Class to generate json summary report for a benchmark
@@ -57,6 +56,36 @@ def __init__(self, spark_session: SparkSession, query_name) -> None:
             'query': query_name,
         }
 
+    def _is_above_spark_4(self):
+        return self.spark_session.version >= "4.0.0"
+
+    def _register_python_listener(self):
+        # Register PythonListener
+        if self._is_above_spark_4():
+            # is_remote_only is added starting from 4.0.0
+            from pyspark import is_remote_only
+            if is_remote_only():
+                # We can't use Py4J in Spark Connect
+                return None
+
+        listener = None
+        try:
+            import python_listener
+            listener = python_listener.PythonListener()
+            listener.register()
+        except TypeError as e:
+            print("Not found com.nvidia.spark.rapids.listener.Manager", str(e))
+        return listener
+
+    def _get_spark_conf(self):
+        if self._is_above_spark_4():
+            from pyspark import is_remote_only
+            if is_remote_only():
+                return self.spark_session.conf.getAll
+
+        return self.spark_session.sparkContext._conf.getAll()
+
+
     def report_on(self, fn: Callable, warmup_iterations = 0, iterations = 1, *args):
         """Record a function for its running environment, running status etc. and exclude sentive
         information like tokens, secret and password Generate summary in dict format for it.
@@ -67,20 +96,14 @@ def report_on(self, fn: Callable, warmup_iterations = 0, iterations = 1, *args):
         Returns:
             dict: summary of the fn
         """
-        spark_conf = dict(self.spark_session.sparkContext._conf.getAll())
+        spark_conf = dict(self._get_spark_conf())
         env_vars = dict(os.environ)
         redacted = ["TOKEN", "SECRET", "PASSWORD"]
         filtered_env_vars = dict((k, env_vars[k]) for k in env_vars.keys() if not (k in redacted))
         self.summary['env']['envVars'] = filtered_env_vars
         self.summary['env']['sparkConf'] = spark_conf
         self.summary['env']['sparkVersion'] = self.spark_session.version
-        listener = None
-        try:
-            listener = python_listener.PythonListener()
-            listener.register()
-        except TypeError as e:
-            print("Not found com.nvidia.spark.rapids.listener.Manager", str(e))
-            listener = None
+        listener = self._register_python_listener()
         if listener is not None:
             print("TaskFailureListener is registered.")
         try:

nds/jvm_listener/pom.xml

@@ -26,19 +26,35 @@
     <properties>
         <maven.compiler.source>8</maven.compiler.source>
         <maven.compiler.target>8</maven.compiler.target>
+        <spark.version>3.1.2</spark.version>
+        <scala.binary.version>2.12</scala.binary.version>
+        <scala.version>2.12.18</scala.version>
     </properties>
+
+    <profiles>
+        <profile>
+            <id>spark4</id>
+            <properties>
+                <spark.version>4.0.0</spark.version>
+                <scala.binary.version>2.13</scala.binary.version>
+                <maven.compiler.source>17</maven.compiler.source>
+                <maven.compiler.target>17</maven.compiler.target>
+                <scala.version>2.13.16</scala.version>
+            </properties>
+        </profile>
+    </profiles>
     <dependencies>
         <!-- https://mvnrepository.com/artifact/org.apache.spark/spark-core -->
         <dependency>
             <groupId>org.apache.spark</groupId>
-            <artifactId>spark-core_2.12</artifactId>
-            <version>3.1.2</version>
+            <artifactId>spark-core_${scala.binary.version}</artifactId>
+            <version>${spark.version}</version>
         </dependency>
         <!-- https://mvnrepository.com/artifact/org.apache.spark/spark-sql -->
         <dependency>
             <groupId>org.apache.spark</groupId>
-            <artifactId>spark-sql_2.12</artifactId>
-            <version>3.1.2</version>
+            <artifactId>spark-sql_${scala.binary.version}</artifactId>
+            <version>${spark.version}</version>
             <scope>provided</scope>
         </dependency>
     </dependencies>
@@ -50,10 +66,6 @@
                 <groupId>org.apache.maven.plugins</groupId>
                 <artifactId>maven-compiler-plugin</artifactId>
                 <version>3.8.1</version>
-                <configuration>
-                    <source>1.8</source>
-                    <target>1.8</target>
-                </configuration>
             </plugin>
             <plugin>
                 <groupId>org.scala-tools</groupId>

nds/nds_power.py

@@ -146,7 +146,7 @@ def setup_tables(spark_session, input_prefix, input_format, use_decimal, executi
     Returns:
         execution_time_list: a list recording query execution time.
     """
-    spark_app_id = spark_session.sparkContext.applicationId
+    spark_app_id = spark_session.conf.get("spark.app.id")
     # Create TempView for tables
     for table_name in get_schemas(False).keys():
         start = int(time.time() * 1000)
@@ -331,7 +331,7 @@ def run_query_stream(input_prefix,
     if input_format == 'delta' and delta_unmanaged:
         # Register tables for Delta Lake. This is only needed for unmanaged tables.
         execution_time_list = register_delta_tables(spark_session, input_prefix, execution_time_list)
-    spark_app_id = spark_session.sparkContext.applicationId
+    spark_app_id = spark_session.conf.get("spark.app.id")
     if input_format != 'iceberg' and input_format != 'delta' and not hive_external:
         execution_time_list = setup_tables(spark_session, input_prefix, input_format, use_decimal,
                                            execution_time_list)
@@ -347,7 +347,9 @@ def run_query_stream(input_prefix,
     power_start = int(time.time())
     for query_name, q_content in query_dict.items():
         # show query name in Spark web UI
-        spark_session.sparkContext.setJobGroup(query_name, query_name)
+        spark_session.conf.set("spark.job.description", query_name)
+        spark_session.conf.set("spark.jobGroup.id", query_name)
+        spark_session.conf.set("spark.job.interruptOnCancel", "false")
         print("====== Run {} ======".format(query_name))
         q_report = PysparkBenchReport(spark_session, query_name)
         summary = q_report.report_on(run_one_query,warmup_iterations,
@@ -374,10 +376,13 @@ def run_query_stream(input_prefix,
             else:
                 summary_prefix =  os.path.join(json_summary_folder, '')
             q_report.write_summary(prefix=summary_prefix)
+    spark_session.conf.unset("spark.job.description")
+    spark_session.conf.unset("spark.jobGroup.id")
+    spark_session.conf.unset("spark.job.interruptOnCancel")
     power_end = int(time.time())
     power_elapse = int((power_end - power_start)*1000)
     if not keep_sc:
-        spark_session.sparkContext.stop()
+        spark_session.stop()
     total_time_end = time.time()
     total_elapse = int((total_time_end - total_time_start)*1000)
     print("====== Power Test Time: {} milliseconds ======".format(power_elapse))