From 45ff94c7b154d443d93746be12c481268f8e74f1 Mon Sep 17 00:00:00 2001 From: Bas Nijholt Date: Fri, 22 Dec 2023 11:11:18 -0800 Subject: [PATCH 1/3] Allow manually driving progress bars --- adaptive_scheduler/utils.py | 45 +++++++++++++++++++++++-------------- 1 file changed, 28 insertions(+), 17 deletions(-) diff --git a/adaptive_scheduler/utils.py b/adaptive_scheduler/utils.py index 9a1532ae..e1ce20e0 100644 --- a/adaptive_scheduler/utils.py +++ b/adaptive_scheduler/utils.py @@ -1241,24 +1241,10 @@ def _remove_completed_paths( return n_completed -async def _track_file_creation_progress( +def _initialize_progress_for_paths( paths_dict: dict[str, set[Path | tuple[Path, ...]]], progress: Progress, - interval: float = 1, -) -> None: - """Asynchronously track and update the progress of file creation. - - Parameters - ---------- - paths_dict - A dictionary with keys representing categories and values being sets of file paths to monitor. - progress - The Progress object from the rich library for displaying progress. - interval - The time interval (in seconds) at which to update the progress. The interval is dynamically - adjusted to be at least 50 times the time it takes to update the progress. This ensures that - updating the progress does not take up a significant amount of time. - """ +) -> tuple[dict[str, TaskID], TaskID | None, int]: # create total_files and add_total_progress before updating paths_dict total_files = sum(len(paths) for paths in paths_dict.values()) add_total_progress = len(paths_dict) > 1 @@ -1283,6 +1269,31 @@ async def _track_file_creation_progress( total=n_remaining + n_done, completed=n_done, ) + return task_ids, total_task, total_files + + +async def _track_file_creation_progress( + paths_dict: dict[str, set[Path | tuple[Path, ...]]], + progress: Progress, + interval: float = 1, +) -> None: + """Asynchronously track and update the progress of file creation. + + Parameters + ---------- + paths_dict + A dictionary with keys representing categories and values being sets of file paths to monitor. + progress + The Progress object from the rich library for displaying progress. + interval + The time interval (in seconds) at which to update the progress. The interval is dynamically + adjusted to be at least 50 times the time it takes to update the progress. This ensures that + updating the progress does not take up a significant amount of time. + """ + task_ids, total_task, total_files = _initialize_progress_for_paths( + paths_dict, + progress, + ) try: progress.start() # Start the progress display total_processed = 0 @@ -1308,7 +1319,7 @@ async def _track_file_creation_progress( def track_file_creation_progress( paths_dict: dict[str, set[Path | tuple[Path, ...]]], - interval: int = 1, + interval: float = 1, ) -> asyncio.Task: """Initialize and asynchronously track the progress of file creation. From b245e2e0135f4b92804132e02ee98923172d7f21 Mon Sep 17 00:00:00 2001 From: Bas Nijholt Date: Fri, 22 Dec 2023 12:01:29 -0800 Subject: [PATCH 2/3] Add FileCreatingProgress class --- adaptive_scheduler/utils.py | 58 +++++++++++++++++++++++++++---------- 1 file changed, 42 insertions(+), 16 deletions(-) diff --git a/adaptive_scheduler/utils.py b/adaptive_scheduler/utils.py index e1ce20e0..38633c7a 100644 --- a/adaptive_scheduler/utils.py +++ b/adaptive_scheduler/utils.py @@ -1199,7 +1199,7 @@ def _update_progress_for_paths( progress: Progress, total_task: TaskID | None, task_ids: dict[str, TaskID], -) -> int: +) -> bool: """Update progress bars for each set of paths.""" n_completed = _remove_completed_paths(paths_dict) total_completed = sum(n_completed.values()) @@ -1207,7 +1207,11 @@ def _update_progress_for_paths( progress.update(task_ids[key], advance=n_done) if total_task is not None: progress.update(total_task, advance=total_completed) - return total_completed + progress.refresh() + if not any(paths_dict.values()): + progress.stop() + return True + return False def _remove_completed_paths( @@ -1243,8 +1247,10 @@ def _remove_completed_paths( def _initialize_progress_for_paths( paths_dict: dict[str, set[Path | tuple[Path, ...]]], - progress: Progress, -) -> tuple[dict[str, TaskID], TaskID | None, int]: +) -> tuple[Progress, dict[str, TaskID], TaskID | None, int]: + columns = (*Progress.get_default_columns(), TimeElapsedColumn()) + progress = Progress(*columns, auto_refresh=False) + # create total_files and add_total_progress before updating paths_dict total_files = sum(len(paths) for paths in paths_dict.values()) add_total_progress = len(paths_dict) > 1 @@ -1269,12 +1275,11 @@ def _initialize_progress_for_paths( total=n_remaining + n_done, completed=n_done, ) - return task_ids, total_task, total_files + return progress, task_ids, total_task, total_files async def _track_file_creation_progress( paths_dict: dict[str, set[Path | tuple[Path, ...]]], - progress: Progress, interval: float = 1, ) -> None: """Asynchronously track and update the progress of file creation. @@ -1290,25 +1295,21 @@ async def _track_file_creation_progress( adjusted to be at least 50 times the time it takes to update the progress. This ensures that updating the progress does not take up a significant amount of time. """ - task_ids, total_task, total_files = _initialize_progress_for_paths( + progress, task_ids, total_task, total_files = _initialize_progress_for_paths( paths_dict, - progress, ) try: progress.start() # Start the progress display - total_processed = 0 while True: t_start = time.time() - total_processed += _update_progress_for_paths( + is_done = _update_progress_for_paths( paths_dict, progress, total_task, task_ids, ) - if total_processed >= total_files: - progress.refresh() # Final refresh to ensure 100% + if is_done: break # Exit loop if all files are processed - progress.refresh() # Sleep for at least 50 times the update time t_update = time.time() - t_start await asyncio.sleep(max(interval, 50 * t_update)) @@ -1317,6 +1318,32 @@ async def _track_file_creation_progress( progress.stop() # Stop the progress display, regardless of what happens +class FileCreatingProgress: + """A class for tracking the progress of file creation. + + Alternative to the async function `track_file_creation_progress`. + + This class needs to be manually updated by calling the `update` method. + """ + + def __init__(self, paths_dict: dict[str, set[Path | tuple[Path, ...]]]) -> None: + """Initialize the progress bar.""" + get_console().clear_live() + _init = _initialize_progress_for_paths(paths_dict) + self.paths_dict = paths_dict + self.progress, self.task_ids, self.total_task, self.total_files = _init + self.progress.start() + + def update(self) -> None: + """Update the progress bar.""" + _update_progress_for_paths( + self.paths_dict, + self.progress, + self.total_task, + self.task_ids, + ) + + def track_file_creation_progress( paths_dict: dict[str, set[Path | tuple[Path, ...]]], interval: float = 1, @@ -1361,8 +1388,7 @@ def track_file_creation_progress( >>> task = track_file_creation_progress(paths_dict) """ get_console().clear_live() # avoid LiveError, only 1 live render allowed at a time - columns = (*Progress.get_default_columns(), TimeElapsedColumn()) - progress = Progress(*columns, auto_refresh=False) - coro = _track_file_creation_progress(paths_dict, progress, interval) + + coro = _track_file_creation_progress(paths_dict, interval) ioloop = asyncio.get_event_loop() return ioloop.create_task(coro) From a7171f3e1a9340e9b8637204fe0cd0484b1d84f8 Mon Sep 17 00:00:00 2001 From: Bas Nijholt Date: Fri, 22 Dec 2023 12:23:46 -0800 Subject: [PATCH 3/3] Add spinner col --- adaptive_scheduler/utils.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/adaptive_scheduler/utils.py b/adaptive_scheduler/utils.py index 38633c7a..36d9d72f 100644 --- a/adaptive_scheduler/utils.py +++ b/adaptive_scheduler/utils.py @@ -42,7 +42,9 @@ from adaptive.notebook_integration import in_ipynb from rich.console import Console from rich.progress import ( + MofNCompleteColumn, Progress, + SpinnerColumn, TimeElapsedColumn, get_console, ) @@ -1248,7 +1250,12 @@ def _remove_completed_paths( def _initialize_progress_for_paths( paths_dict: dict[str, set[Path | tuple[Path, ...]]], ) -> tuple[Progress, dict[str, TaskID], TaskID | None, int]: - columns = (*Progress.get_default_columns(), TimeElapsedColumn()) + columns = ( + SpinnerColumn(), + *Progress.get_default_columns(), + TimeElapsedColumn(), + MofNCompleteColumn(), + ) progress = Progress(*columns, auto_refresh=False) # create total_files and add_total_progress before updating paths_dict