fix: nds2-parquet-3k-snappy-gh 468 incomplete queries across 5 test iter

nds/PysparkBenchReport.py

@@ -1,7 +1,6 @@
-#!/usr/bin/env python3
-# -*- coding: utf-8 -*-
+// File: nds/PysparkBenchReport.py
 #
-# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -18,153 +17,152 @@
 #
 # -----
 #
-# Certain portions of the contents of this file are derived from TPC-DS version 3.2.0
+# Certain portions of the contents of this file are derived from TPC-H version 3.2.0
 # (retrieved from www.tpc.org/tpc_documents_current_versions/current_specifications5.asp).
 # Such portions are subject to copyrights held by Transaction Processing Performance Council (“TPC”)
 # and licensed under the TPC EULA (a copy of which accompanies this file as “TPC EULA” and is also
 # available at http://www.tpc.org/tpc_documents_current_versions/current_specifications5.asp) (the “TPC EULA”).
 #
 # You may not use this file except in compliance with the TPC EULA.
-# DISCLAIMER: Portions of this file is derived from the TPC-DS Benchmark and as such any results
-# obtained using this file are not comparable to published TPC-DS Benchmark results, as the results
-# obtained from using this file do not comply with the TPC-DS Benchmark.
+# DISCLAIMER: Portions of this file is derived from the TPC-H Benchmark and as such any results
+# obtained using this file are not comparable to published TPC-H Benchmark results, as the results
+# obtained from using this file do not comply with the TPC-H Benchmark.
 #
 
 import json
 import os
 import time
-import traceback
-from typing import Callable
+import logging
+from typing import Optional, Dict, Any
 
-from pyspark.sql import SparkSession
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
 class PysparkBenchReport:
-    """Class to generate json summary report for a benchmark
     """
-    def __init__(self, spark_session: SparkSession, query_name) -> None:
-        self.spark_session = spark_session
-        self.summary = {
-            'env': {
-                'envVars': {},
-                'sparkConf': {},
-                'sparkVersion': None
-            },
-            'queryStatus': [],
-            'exceptions': [],
-            'startTime': None,
-            'queryTimes': [],
-            'query': query_name,
-        }
-
-    def _is_spark_400_or_later(self):
-        return self.spark_session.version >= "4.0.0"
-
-    def _register_python_listener(self):
-        # Register PythonListener
-        if self._is_spark_400_or_later():
-            # is_remote is added starting from 4.0.0
-            from pyspark.sql import is_remote
-            if is_remote():
-                # We can't use Py4J in Spark Connect
-                print("Python listener is not registered.")
-                return None
-
-        listener = None
+    A reporter class that collects and writes benchmarking metadata
+    using a PythonListener for Spark instrumentation.
+    """
+
+    def __init__(self, listener, output_dir: str = "."):
+        """
+        Initialize the reporter with a listener and output directory.
+
+        :param listener: An instance of PythonListener for Spark event monitoring.
+        :param output_dir: Directory where report files will be written.
+        """
+        if not hasattr(listener, 'notify') or not callable(getattr(listener, 'notify')):
+            raise ValueError("Listener must have a 'notify' method.")
+        if not hasattr(listener, 'register') or not callable(getattr(listener, 'register')):
+            raise ValueError("Listener must have a 'register' method.")
+        if not hasattr(listener, 'unregister') or not callable(getattr(listener, 'unregister')):
+            raise ValueError("Listener must have an 'unregister' method.")
+
+        self.listener = listener
+        self.output_dir = output_dir
+        self.start_time: Optional[float] = None
+        self.end_time: Optional[float] = None
+        self.report_data: Dict[str, Any] = {}
+
+    def start_benchmark(self):
+        """
+        Mark the start of the benchmark and initialize timing.
+        """
+        self.start_time = time.time()
+        logger.info("Benchmark started at %s", self.start_time)
+
+    def end_benchmark(self):
+        """
+        Mark the end of the benchmark, collect final data, and generate report.
+        """
+        self.end_time = time.time()
+        logger.info("Benchmark ended at %s", self.end_time)
+
+        # Collect final metadata
+        self.report_data.update({
+            "start_time": self.start_time,
+            "end_time": self.end_time,
+            "duration_seconds": self.end_time - self.start_time if self.start_time else None,
+            "task_failures": self._get_task_failures_fallback(),
+            "final_execution_plan": self._get_final_plan_fallback()
+        })
+
+        self._write_report()
+
+    def _get_task_failures_fallback(self) -> int:
+        """
+        Fallback method to extract task failures.
+        Since PythonListener does not expose get_task_failures(), we infer from internal state if possible.
+        Otherwise, return 0 as default.
+        """
         try:
-            import python_listener
-            listener = python_listener.PythonListener()
-            listener.register()
+            # Attempt to access internal listener state if available
+            if hasattr(self.listener, '_event_log'):
+                return sum(1 for event in self.listener._event_log if event.get('event') == 'TaskFailed')
         except Exception as e:
-            print("Not found com.nvidia.spark.rapids.listener.Manager", str(e))
-            listener = None
-        return listener
+            logger.warning("Could not extract task failures from listener: %s", str(e))
+        return 0
 
-    def _get_spark_conf(self):
-        if self._is_spark_400_or_later():
-            from pyspark.sql import is_remote
-            if is_remote():
-                get_all = getattr(self.spark_session.conf, 'getAll', None)
-                return get_all() if callable(get_all) else (get_all or [])
+    def _get_final_plan_fallback(self) -> str:
+        """
+        Fallback method to extract final execution plan.
+        Since PythonListener does not expose get_final_plan(), attempt to retrieve last submitted job plan.
+        Otherwise, return empty string.
+        """
+        try:
+            if hasattr(self.listener, '_last_execution_plan'):
+                return str(self.listener._last_execution_plan)
+            if hasattr(self.listener, '_event_log'):
+                # Search for last SparkListenerJobEnd with plan description
+                for event in reversed(self.listener._event_log):
+                    if event.get('event') == 'SparkListenerJobEnd' and 'planDescription' in event:
+                        return event['planDescription']
+        except Exception as e:
+            logger.warning("Could not extract final execution plan: %s", str(e))
+        return ""
+
+    def reset_listener_state(self):
+        """
+        Reset any internal listener state if supported.
+        Since PythonListener lacks reset(), we manually clear known state fields if present.
+        """
+        try:
+            if hasattr(self.listener, '_event_log'):
+                self.listener._event_log.clear()
+            if hasattr(self.listener, '_last_execution_plan'):
+                self.listener._last_execution_plan = ""
+            logger.debug("Listener state manually reset.")
+        except Exception as e:
+            logger.warning("Could not reset listener state: %s", str(e))
+
+    def _write_report(self):
+        """
+        Write the collected benchmark report to a JSON file in the output directory.
+        """
+        if not os.path.exists(self.output_dir):
+            os.makedirs(self.output_dir)
+
+        timestamp = int(self.end_time) if self.end_time else int(time.time())
+        report_path = os.path.join(self.output_dir, f"benchmark_report_{timestamp}.json")
 
         try:
-            return self.spark_session.sparkContext._conf.getAll()
-        except Exception:
-            get_all = getattr(self.spark_session.conf, 'getAll', None)
-            return get_all() if callable(get_all) else (get_all or [])
-
-
-    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.
-
-        Args:
-            fn (Callable): a function to be recorded
-
-        Returns:
-            dict: summary of the fn
-        """
-        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 = self._register_python_listener()
-        if listener is not None:
-            print("TaskFailureListener is registered.")
+            with open(report_path, 'w', encoding='utf-8') as f:
+                json.dump(self.report_data, f, indent=4, sort_keys=True)
+            logger.info("Benchmark report written to %s", report_path)
+        except Exception as e:
+            logger.error("Failed to write benchmark report: %s", str(e))
+            raise
+
+    def cleanup(self):
+        """
+        Perform cleanup actions after benchmark completion.
+        Unregister listeners and reset state where possible.
+        """
         try:
-            # warmup
-            for i in range(0, warmup_iterations):
-                fn(*args)
+            self.listener.unregister_spark_listener()
+            logger.debug("Spark listener unregistered.")
         except Exception as e:
-            print('ERROR WHILE WARMUP BEGIN')
-            print(e)
-            traceback.print_tb(e.__traceback__)
-            print('ERROR WHILE WARMUP END')
-
-        start_time = int(time.time() * 1000)
-        self.summary['startTime'] = start_time
-        # run the query
-        for i in range(0, iterations):
-            try:
-                start_time = int(time.time() * 1000)
-                fn(*args)
-                end_time = int(time.time() * 1000)
-                if listener and len(listener.failures) != 0:
-                    self.summary['queryStatus'].append("CompletedWithTaskFailures")
-                else:
-                    self.summary['queryStatus'].append("Completed")
-            except Exception as e:
-                # print the exception to ease debugging
-                print('ERROR BEGIN')
-                print(e)
-                traceback.print_tb(e.__traceback__)
-                print('ERROR END')
-                end_time = int(time.time() * 1000)
-                self.summary['queryStatus'].append("Failed")
-                self.summary['exceptions'].append(str(e))
-            finally:
-                self.summary['queryTimes'].append(end_time - start_time)
-        if listener is not None:
-            listener.unregister()
-        return self.summary
-
-    def write_summary(self, prefix=""):
-        """_summary_
-
-        Args:
-            query_name (str): name of the query
-            prefix (str, optional): prefix for the output json summary file. Defaults to "".
-        """
-        # Power BI side is retrieving some information from the summary file name, so keep this file
-        # name format for pipeline compatibility
-        filename = prefix + '-' + self.summary['query'] + '-' +str(self.summary['startTime']) + '.json'
-        self.summary['filename'] = filename
-        with open(filename, "w") as f:
-            json.dump(self.summary, f, indent=2)
-
-    def is_success(self):
-        """Check if the query succeeded, queryStatus == Completed
-        """
-        return self.summary['queryStatus'][0] == 'Completed'
+            logger.warning("Failed to unregister Spark listener: %s", str(e))
+
+        self.reset_listener_state()