viz pbar

python/dolma/core/parallel.py

@@ -420,7 +420,7 @@ def _run_all(
 
         with PoolWithDebug(processes=num_processes, debug=self.debug) as pool:
             pbar_queue: QueueType = (manager := get_manager(pool)).Queue()
-            (pbar := self.PROGRESS_BAR_CLS(pbar_queue)).start()
+            (pbar := self.PROGRESS_BAR_CLS(pbar_queue, thread=True)).start()
 
             process_single_fn = partial(self.process_single, queue=pbar_queue)
             results = []

python/dolma/core/progressbar.py

@@ -81,7 +81,7 @@ def __init__(self, queue: QueueType, min_step: int = 1, min_time: float = 1e-1,
         """
         self._logger = get_logger(self.__class__.__name__, "warn")
         self._queue = queue
-        self._last_update_delta_time = 0
+        self._last_update_time = 0
         self._last_update_delta_step = 0
 
         self._update_every_seconds = min_time
@@ -196,6 +196,7 @@ def _update(self):
 
         # reset the steps
         self._last_update_delta_step = 0
+        self._last_update_time = time.time()
 
         # reset the steps
         for k in self.fields():
@@ -208,6 +209,7 @@ def update(self):
         if self._update_every_steps > self._last_update_delta_step:
             return
 
+        time_before_update = self._last_update_time
         self._update()
 
         # check if we wanna update frequency based on steps
@@ -216,8 +218,7 @@ def update(self):
             return
 
         # check if we wanna update frequency based on time
-        self._last_update_delta_time = -(time.time() - self._last_update_delta_time)
-        if self._last_update_delta_time < self._update_every_seconds:
+        if (self._last_update_time - time_before_update) < self._update_every_seconds:
             self._update_every_steps *= 2
             return