Fix busy waiting

Author: Logiquo

pyhealth/datasets/base_dataset.py

@@ -709,10 +709,11 @@ def _task_transform(self, task: BaseTask, output_dir: Path, num_workers: int) ->
                 with ctx.Pool(processes=num_workers, initializer=_task_transform_init, initargs=(queue,)) as pool:
                     result = pool.map_async(_task_transform_fn, args_list) # type: ignore
                     with tqdm(total=len(patient_ids)) as progress:
-                        # TODO: this is busy-waiting, can we do better?
                         while not result.ready():
-                            while not queue.empty():
-                                progress.update(queue.get())
+                            try:
+                                progress.update(queue.get(timeout=1))
+                            except:
+                                pass
 
                         # remaining items
                         while not queue.empty():
@@ -759,8 +760,10 @@ def _proc_transform(self, task_df: Path, output_dir: Path, num_workers: int) ->
                     result = pool.map_async(_proc_transform_fn, args_list) # type: ignore
                     with tqdm(total=num_samples) as progress:
                         while not result.ready():
-                            while not queue.empty():
-                                progress.update(queue.get())
+                            try:
+                                progress.update(queue.get(timeout=1))
+                            except:
+                                pass
 
                         # remaining items
                         while not queue.empty():