From c5a5d7c08c07704d6a4166048f2fbe831f5b7aca Mon Sep 17 00:00:00 2001 From: Tom White Date: Thu, 5 Sep 2024 16:29:51 +0100 Subject: [PATCH] Align tqdm progress bars (#567) * Align tqdm progress bars * Improve alignment of the colon --- cubed/diagnostics/tqdm.py | 20 ++++++++++++++++---- 1 file changed, 16 insertions(+), 4 deletions(-) diff --git a/cubed/diagnostics/tqdm.py b/cubed/diagnostics/tqdm.py index 03ef4ac9..7cbc9102 100644 --- a/cubed/diagnostics/tqdm.py +++ b/cubed/diagnostics/tqdm.py @@ -17,19 +17,31 @@ def __init__(self, *args, **kwargs): def on_compute_start(self, event): from tqdm.auto import tqdm + # find the maximum display width so we can align bars below + max_op_display_name = ( + max( + len(node["op_display_name"].replace("\n", " ")) + for _, node in visit_nodes(event.dag, event.resume) + ) + + 1 # for the colon + ) + self.pbars = {} - i = 0 - for name, node in visit_nodes(event.dag, event.resume): + for i, (name, node) in enumerate(visit_nodes(event.dag, event.resume)): num_tasks = node["primitive_op"].num_tasks - op_display_name = node["op_display_name"].replace("\n", " ") + op_display_name = node["op_display_name"].replace("\n", " ") + ":" + # note double curlies to get literal { and } for tqdm bar format + bar_format = ( + f"{{desc:{max_op_display_name}}} {{percentage:3.0f}}%|{{bar}}{{r_bar}}" + ) self.pbars[name] = tqdm( *self.args, desc=op_display_name, total=num_tasks, position=i, + bar_format=bar_format, **self.kwargs, ) - i = i + 1 def on_compute_end(self, event): for pbar in self.pbars.values():