Skip to content

Commit

Permalink
Allow tools to display percent progress in the monitor (#258)
Browse files Browse the repository at this point in the history
  • Loading branch information
jeremyfowers authored Jan 6, 2025
1 parent 3ff3e16 commit 2b9a783
Show file tree
Hide file tree
Showing 2 changed files with 49 additions and 7 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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
46 changes: 40 additions & 6 deletions src/turnkeyml/tools/tool.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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:
Expand Down Expand Up @@ -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,
Expand All @@ -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
Expand All @@ -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:
"""
Expand Down Expand Up @@ -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:
Expand Down

0 comments on commit 2b9a783

Please sign in to comment.