Fix spinner animation blocking diagnostic prompt (#631)

Signed-off-by: Partho Sarthi <[email protected]>

user_tools/src/spark_rapids_pytools/common/utilities.py

@@ -379,23 +379,25 @@ class ToolsSpinner:
     A class to manage the spinner animation.
     Reference: https://stackoverflow.com/a/66558182
 
-    :param in_debug_mode: Flag indicating if running in debug (verbose) mode. Defaults to False.
+    :param enabled: Flag indicating if the spinner is enabled. Defaults to True.
     """
-    in_debug_mode: bool = field(default=False, init=True)
-    pixel_spinner: PixelSpinner = field(default=PixelSpinner('Processing...'), init=False)
+    enabled: bool = field(default=True, init=True)
+    pixel_spinner: PixelSpinner = field(default=PixelSpinner('Processing...', hide_cursor=False), init=False)
     end: str = field(default='Processing Completed!', init=False)
     timeout: float = field(default=0.1, init=False)
     completed: bool = field(default=False, init=False)
     spinner_thread: threading.Thread = field(default=None, init=False)
+    pause_event: threading.Event = field(default=threading.Event(), init=False)
 
     def _spinner_animation(self):
         while not self.completed:
             self.pixel_spinner.next()
             time.sleep(self.timeout)
+            while self.pause_event.is_set():
+                self.pause_event.wait(self.timeout)
 
     def start(self):
-        # Don't start if in debug mode
-        if not self.in_debug_mode:
+        if self.enabled:
             self.spinner_thread = threading.Thread(target=self._spinner_animation, daemon=True)
             self.spinner_thread.start()
         return self
@@ -404,6 +406,16 @@ def stop(self):
         self.completed = True
         print(f'\r\n{self.end}', flush=True)
 
+    def pause(self, insert_newline=False):
+        if self.enabled:
+            if insert_newline:
+                # Print a newline for visual separation
+                print()
+            self.pause_event.set()
+
+    def resume(self):
+        self.pause_event.clear()
+
     def __enter__(self):
         return self.start()
 

user_tools/src/spark_rapids_pytools/rapids/diagnostic.py

@@ -40,16 +40,20 @@ def _process_custom_args(self):
 
         self.thread_num = thread_num
         self.logger.debug('Set thread number as: %d', self.thread_num)
-
-        self.logger.warning('This operation will collect sensitive information from your cluster, '
-                            'such as OS & HW info, Yarn/Spark configurations and log files etc.')
+        log_message = ('This operation will collect sensitive information from your cluster, '
+                       'such as OS & HW info, Yarn/Spark configurations and log files etc.')
         yes = self.wrapper_options.get('yes', False)
         if yes:
+            self.logger.warning(log_message)
             self.logger.info('Confirmed by command line option.')
         else:
+            # Pause the spinner for user prompt
+            self.spinner.pause(insert_newline=True)
+            print(log_message)
             user_input = input('Do you want to continue (yes/no): ')
             if user_input.lower() not in ['yes', 'y']:
                 raise RuntimeError('User canceled the operation.')
+            self.spinner.resume()
 
     def requires_cluster_connection(self) -> bool:
         return True

user_tools/src/spark_rapids_pytools/rapids/rapids_tool.py

@@ -59,6 +59,7 @@ class RapidsTool(object):
     name: str = field(default=None, init=False)
     ctxt: ToolContext = field(default=None, init=False)
     logger: Logger = field(default=None, init=False)
+    spinner: ToolsSpinner = field(default=None, init=False)
 
     def pretty_name(self):
         return self.name.capitalize()
@@ -272,7 +273,9 @@ def _verify_exec_cluster(self):
         self._handle_non_running_exec_cluster(msg)
 
     def launch(self):
-        with ToolsSpinner(in_debug_mode=ToolLogging.is_debug_mode_enabled()):
+        # Spinner should not be enabled in debug mode
+        enable_spinner = not ToolLogging.is_debug_mode_enabled()
+        with ToolsSpinner(enabled=enable_spinner) as self.spinner:
             self._init_tool()
             self._connect_to_execution_cluster()
             self._process_arguments()