Allow tools to display percent progress in the monitor (onnx#258)

Author: jeremyfowers

examples/cli/plugins/example_tool/turnkeyml_plugin_example_tool/tool.py

@@ -4,14 +4,16 @@
 and build_model() to have any custom behavior.
 
 In this the tool simply passes the build state to the next
-tool in the sequence (i.e., this example is a no-op). 
+tool in the sequence (i.e., this example is a no-op). It also
+spends a few seconds updating the monitor's percent progress indicator.
 
 After you install the plugin, you can tell `turnkey` to use this sequence with:
 
     turnkey -i INPUT_SCRIPT export-pytorch example-plugin-tool 
 """
 
 import argparse
+from time import sleep
 from turnkeyml.tools import Tool
 from turnkeyml.state import State
 
@@ -39,4 +41,10 @@ def parser(add_help: bool = True) -> argparse.ArgumentParser:
         return parser
 
     def run(self, state: State):
+        self.set_percent_progress(0.0)
+        total = 15  # seconds
+        for i in range(total):
+            sleep(1)
+            percent_progress = (i + 1) / float(total) * 100
+            self.set_percent_progress(percent_progress)
         return state

src/turnkeyml/tools/tool.py

@@ -6,7 +6,7 @@
 import textwrap as _textwrap
 import re
 from typing import Tuple, Dict
-from multiprocessing import Process
+from multiprocessing import Process, Queue
 import psutil
 import turnkeyml.common.printing as printing
 import turnkeyml.common.exceptions as exp
@@ -15,13 +15,25 @@
 from turnkeyml.state import State
 
 
-def _spinner(message):
+def _spinner(message, q: Queue):
+    """
+    Displays a moving "..." indicator so that the user knows that the
+    Tool is still working. Tools can optionally use a multiprocessing
+    Queue to display the percent progress of the Tool.
+    """
+    percent_complete = None
+
     try:
         parent_process = psutil.Process(pid=os.getppid())
         while parent_process.status() == psutil.STATUS_RUNNING:
             for cursor in ["   ", ".  ", ".. ", "..."]:
                 time.sleep(0.5)
-                status = f"      {message}{cursor}\r"
+                if not q.empty():
+                    percent_complete = q.get()
+                if percent_complete is not None:
+                    status = f"      {message} ({percent_complete:.1f}%){cursor}\r"
+                else:
+                    status = f"      {message}{cursor}\r"
                 sys.stdout.write(status)
                 sys.stdout.flush()
     except psutil.NoSuchProcess:
@@ -146,17 +158,22 @@ def status_line(self, successful, verbosity):
                 success_tick = "+"
                 fail_tick = "x"
 
+            if self.percent_progress is None:
+                progress_indicator = ""
+            else:
+                progress_indicator = f" ({self.percent_progress:.1f}%)"
+
             if successful is None:
                 # Initialize the message
                 printing.logn(f"      {self.monitor_message}   ")
             elif successful:
                 # Print success message
                 printing.log(f"    {success_tick} ", c=printing.Colors.OKGREEN)
-                printing.logn(self.monitor_message + "   ")
+                printing.logn(self.monitor_message + progress_indicator + "   ")
             else:
                 # successful == False, print failure message
                 printing.log(f"    {fail_tick} ", c=printing.Colors.FAIL)
-                printing.logn(self.monitor_message + "   ")
+                printing.logn(self.monitor_message + progress_indicator + "   ")
 
     def __init__(
         self,
@@ -169,6 +186,8 @@ def __init__(
         self.duration_key = f"{fs.Keys.TOOL_DURATION}:{self.__class__.unique_name}"
         self.monitor_message = monitor_message
         self.progress = None
+        self.progress_queue = None
+        self.percent_progress = None
         self.logfile_path = None
         # Tools can disable build.Logger, which captures all stdout and stderr from
         # the Tool, by setting enable_logger=False
@@ -192,6 +211,18 @@ def parser() -> argparse.ArgumentParser:
         line interface for this Tool.
         """
 
+    def set_percent_progress(self, percent_progress: float):
+        """
+        Update the progress monitor with a percent progress to let the user
+        know how much progress the Tool has made.
+        """
+
+        if not isinstance(percent_progress, float):
+            raise ValueError(f"Input argument must be a float, got {percent_progress}")
+
+        self.progress_queue.put(percent_progress)
+        self.percent_progress = percent_progress
+
     # pylint: disable=unused-argument
     def parse(self, state: State, args, known_only=True) -> argparse.Namespace:
         """
@@ -250,7 +281,10 @@ def run_helper(
         )
 
         if monitor:
-            self.progress = Process(target=_spinner, args=[self.monitor_message])
+            self.progress_queue = Queue()
+            self.progress = Process(
+                target=_spinner, args=(self.monitor_message, self.progress_queue)
+            )
             self.progress.start()
 
         try: