Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Allow tools to display percent progress in the monitor #258

Merged
merged 3 commits into from
Jan 6, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
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
Loading