Repository: ludwig-ai/ludwig
Pull request: Fix progress bar ray
# Code Pull Requests

Please provide the following:

- Fix #1528 
- Moves all the tqdm logic into the ray driver and uses a Callback from ray train to trigger tqdm stuff. Need to remove some of the coordinator logic into the LudwigProgressbar Class since all child processes need to run ray.train.report if one of them does (otherwise it hangs)




Problem statement:
[ray] Fix training progress bars
Progress bars are not rendering properly in the terminal for Ray, instead printing one line per step.

This solution here is to use something like Callbacks to send updates back to the head / driver to print.

Hints from issue discussion:
Proposed general solution from the Ray side, but looks like it still has some way to go before being landed and released: https://github.com/ray-project/ray/pull/21977

Commit message: [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

Patch to review:
diff --git a/ludwig/backend/ray.py b/ludwig/backend/ray.py
index b5cd3c3666e..ebb6c67d93b 100644
--- a/ludwig/backend/ray.py
+++ b/ludwig/backend/ray.py
@@ -26,6 +26,7 @@
 import pandas as pd
 import ray
 import torch
+import tqdm
 from ray import ObjectRef
 from ray.data.dataset_pipeline import DatasetPipeline
 from ray.data.extensions import TensorDtype
@@ -189,7 +190,7 @@ def train_fn(
         device = get_torch_device()
         model = model.to(device)
 
-        trainer = RemoteTrainer(model=model, horovod=hvd, **executable_kwargs)
+        trainer = RemoteTrainer(model=model, horovod=hvd, report_tqdm_to_ray=True, **executable_kwargs)
         results = trainer.train(train_shard, val_shard, test_shard, **kwargs)
 
         if results is not None:
@@ -272,6 +273,30 @@ def tune_learning_rate_fn(
         hvd.shutdown()
 
 
+class TqdmCallback(rt.TrainingCallback):
+    def __init__(self):
+        super().__init__()
+        self.progess_bars = {}
+
+    def process_results(self, results, **info):
+        progress_bar_opts = results[0].get("progress_bar")
+        if not progress_bar_opts:
+            return
+        # Skip commands received by non-coordinators
+        if not progress_bar_opts["is_coordinator"]:
+            return
+        _id = progress_bar_opts["id"]
+        action = progress_bar_opts.pop("action")
+        if action == "create":
+            progress_bar_config = progress_bar_opts.get("config")
+            self.progess_bars[_id] = tqdm.tqdm(**progress_bar_config)
+        elif action == "close":
+            self.progess_bars[_id].close()
+        elif action == "update":
+            update_by = progress_bar_opts.pop("update_by")
+            self.progess_bars[_id].update(update_by)
+
+
 class RayTrainerV2(BaseTrainer):
     def __init__(self, model, trainer_kwargs, data_loader_kwargs, executable_kwargs):
         self.model = model.cpu()
@@ -315,6 +340,7 @@ def train(
             results, self._validation_field, self._validation_metric = runner.run(
                 lambda config: train_fn(**config),
                 config={"executable_kwargs": executable_kwargs, "model_ref": ray.put(self.model), **kwargs},
+                callbacks=[TqdmCallback()],
                 dataset=dataset,
             )[0]
 
@@ -326,6 +352,8 @@ def train(
         return results
 
     def train_online(self, *args, **kwargs):
+        # TODO: When this is implemented we also need to update the
+        # Tqdm flow to report back the callback
         raise NotImplementedError()
 
     def tune_batch_size(
@@ -534,7 +562,7 @@ def eval_fn(
         device = get_torch_device()
         model = model.to(device)
 
-        predictor = RemotePredictor(model=model, horovod=hvd, **predictor_kwargs)
+        predictor = RemotePredictor(model=model, horovod=hvd, report_tqdm_to_ray=True, **predictor_kwargs)
         return predictor.batch_evaluation(eval_shard, **kwargs)
     finally:
         torch.cuda.empty_cache()
diff --git a/ludwig/models/predictor.py b/ludwig/models/predictor.py
index bca0d29bbee..ca596238437 100644
--- a/ludwig/models/predictor.py
+++ b/ludwig/models/predictor.py
@@ -10,7 +10,6 @@
 import pandas as pd
 import psutil
 import torch
-from tqdm import tqdm
 
 from ludwig.constants import COMBINED, LAST_HIDDEN, LOGITS
 from ludwig.data.dataset.base import Dataset
@@ -22,6 +21,7 @@
     TEST_STATISTICS_FILE_NAME,
 )
 from ludwig.models.ecd import ECD
+from ludwig.progress_bar import LudwigProgressBar
 from ludwig.utils.data_utils import flatten_df, from_numpy_dataset, save_csv, save_json
 from ludwig.utils.horovod_utils import return_first
 from ludwig.utils.print_utils import repr_ordered_dict
@@ -67,9 +67,10 @@ def __exit__(self, exc_type, exc_val, exc_tb):
 class Predictor(BasePredictor):
     """Predictor is a class that uses a model to predict and evaluate."""
 
-    def __init__(self, model: ECD, batch_size=128, horovod=None, **kwargs):
+    def __init__(self, model: ECD, batch_size=128, horovod=None, report_tqdm_to_ray=False, **kwargs):
         self._batch_size = batch_size
         self._horovod = horovod
+        self.report_tqdm_to_ray = report_tqdm_to_ray
 
         self.device = get_torch_device()
         self.model = model.to(self.device)
@@ -85,26 +86,21 @@ def batch_predict(
         with torch.no_grad():
             with dataset.initialize_batcher(self._batch_size, should_shuffle=False, horovod=self._horovod) as batcher:
 
-                progress_bar = None
-                if self.is_coordinator():
-                    progress_bar = tqdm(
-                        desc="Prediction" if dataset_name is None else f"Prediction {dataset_name: <5.5}",
-                        total=batcher.steps_per_epoch,
-                        file=sys.stdout,
-                        disable=is_progressbar_disabled(),
-                    )
-
+                progress_bar_config = {
+                    "desc": "Prediction" if dataset_name is None else f"Prediction {dataset_name: <5.5}",
+                    "total": batcher.steps_per_epoch,
+                    "file": sys.stdout,
+                    "disable": is_progressbar_disabled(),
+                }
+                progress_bar = LudwigProgressBar(self.report_tqdm_to_ray, progress_bar_config, self.is_coordinator())
                 predictions = defaultdict(list)
                 while not batcher.last_batch():
                     batch = batcher.next_batch()
                     preds = self._predict(self.model, batch)
                     self._accumulate_preds(preds, predictions)
+                    progress_bar.update(1)
 
-                    if self.is_coordinator():
-                        progress_bar.update(1)
-
-                if self.is_coordinator():
-                    progress_bar.close()
+                progress_bar.close()
 
         # consolidate predictions from each batch to a single tensor
         self._concat_preds(predictions)
@@ -164,16 +160,14 @@ def batch_evaluation(self, dataset, collect_predictions=False, dataset_name=None
 
         with torch.no_grad():
             with dataset.initialize_batcher(self._batch_size, should_shuffle=False, horovod=self._horovod) as batcher:
-
-                progress_bar = None
-                if self.is_coordinator():
-                    progress_bar = tqdm(
-                        desc="Evaluation" if dataset_name is None else f"Evaluation {dataset_name: <5.5}",
-                        total=batcher.steps_per_epoch,
-                        file=sys.stdout,
-                        disable=is_progressbar_disabled(),
-                        position=0,  # Necessary to disable extra new line artifacts in training logs.
-                    )
+                progress_bar_config = {
+                    "desc": "Evaluation" if dataset_name is None else f"Evaluation {dataset_name: <5.5}",
+                    "total": batcher.steps_per_epoch,
+                    "file": sys.stdout,
+                    "disable": is_progressbar_disabled(),
+                    "position": 0,  # Necessary to disable extra new line artifacts in training logs.
+                }
+                progress_bar = LudwigProgressBar(self.report_tqdm_to_ray, progress_bar_config, self.is_coordinator())
 
                 predictions = defaultdict(list)
                 while not batcher.last_batch():
@@ -201,14 +195,13 @@ def batch_evaluation(self, dataset, collect_predictions=False, dataset_name=None
                                     key = f"{of_name}_{pred_name}"
                                     predictions[key].append(pred_values)
 
+                    progress_bar.update(1)
                     if self.is_coordinator():
-                        progress_bar.update(1)
                         logger.debug(
-                            f"evaluation for {dataset_name}: completed batch {progress_bar.n} "
+                            f"evaluation for {dataset_name}: completed batch {progress_bar.total_steps} "
                             f"memory used: {psutil.Process(os.getpid()).memory_info()[0] / 1e6:0.2f}MB"
                         )
 
-            if self.is_coordinator():
                 progress_bar.close()
 
             # consolidate predictions from each batch to a single tensor
@@ -232,12 +225,13 @@ def batch_collect_activations(self, layer_names, dataset, bucketing_field=None):
 
         with torch.no_grad():
             with dataset.initialize_batcher(self._batch_size, should_shuffle=False) as batcher:
-                progress_bar = tqdm(
-                    desc="Collecting Tensors",
-                    total=batcher.steps_per_epoch,
-                    file=sys.stdout,
-                    disable=is_progressbar_disabled(),
-                )
+                progress_bar_config = {
+                    "desc": "Collecting Tensors",
+                    "total": batcher.steps_per_epoch,
+                    "file": sys.stdout,
+                    "disable": is_progressbar_disabled(),
+                }
+                progress_bar = LudwigProgressBar(self.report_tqdm_to_ray, progress_bar_config, True)
 
                 collected_tensors = []
                 while not batcher.last_batch():
@@ -249,7 +243,6 @@ def batch_collect_activations(self, layer_names, dataset, bucketing_field=None):
                     }
                     outputs = self.model(inputs)
                     collected_tensors = [(concat_name, tensor) for concat_name, tensor in outputs.items()]
-
                     progress_bar.update(1)
 
                 progress_bar.close()
diff --git a/ludwig/models/trainer.py b/ludwig/models/trainer.py
index 8b2e4af634a..737f7e01c88 100644
--- a/ludwig/models/trainer.py
+++ b/ludwig/models/trainer.py
@@ -32,7 +32,6 @@
 import torch
 from tabulate import tabulate
 from torch.utils.tensorboard import SummaryWriter
-from tqdm import tqdm
 
 from ludwig.constants import COMBINED, LOSS, TEST, TRAINING, VALIDATION
 from ludwig.data.dataset.base import Dataset
@@ -47,6 +46,7 @@
 from ludwig.models.predictor import Predictor
 from ludwig.modules.metric_modules import get_improved_fun, get_initial_validation_value
 from ludwig.modules.optimization_modules import create_clipper, create_optimizer
+from ludwig.progress_bar import LudwigProgressBar
 from ludwig.schema.trainer import TrainerConfig
 from ludwig.utils import time_utils
 from ludwig.utils.checkpoint_utils import Checkpoint, CheckpointManager
@@ -127,6 +127,7 @@ def __init__(
         skip_save_progress: bool = False,
         skip_save_log: bool = False,
         callbacks: List = None,
+        report_tqdm_to_ray=False,
         random_seed: float = default_random_seed,
         horovod: Optional[Dict] = None,
         device: Optional[str] = None,
@@ -155,6 +156,7 @@ def __init__(
         :param callbacks: List of `ludwig.callbacks.Callback` objects that provide hooks into the Ludwig pipeline.
                 (default: None).
         :type callbacks: list
+        :param report_tqdm_to_ray: Enables using the ray based tqdm Callback for progress bar reporting
         :param random_seed: Default initialization for the random seeds (default: 42).
         :type random_seed: Float
         :param horovod: Horovod parameters (default: None).
@@ -210,6 +212,7 @@ def __init__(
         self.random_seed = random_seed
         self.horovod = horovod
         self.received_sigint = False
+        self.report_tqdm_to_ray = report_tqdm_to_ray
         self.callbacks = callbacks or []
         self.device = device
         if self.device is None:
@@ -817,14 +820,13 @@ def train(self, training_set, validation_set=None, test_set=None, save_path="mod
 
                     logger.info(f"Starting with step {progress_tracker.steps}, epoch: {progress_tracker.epoch}")
 
-                progress_bar = None
-                if self.is_coordinator():
-                    progress_bar = tqdm(
-                        desc="Training",
-                        total=self.total_steps,
-                        file=sys.stdout,
-                        disable=is_progressbar_disabled(),
-                    )
+                progress_bar_config = {
+                    "desc": "Training",
+                    "total": self.total_steps,
+                    "disable": is_progressbar_disabled(),
+                    "file": sys.stdout,
+                }
+                progress_bar = LudwigProgressBar(self.report_tqdm_to_ray, progress_bar_config, self.is_coordinator())
 
                 while progress_tracker.steps < self.total_steps:
                     # note that batch size may change over epochs
@@ -994,10 +996,10 @@ def _train_loop(
                 )
 
             progress_tracker.steps += 1
+            progress_bar.update(1)
             if self.is_coordinator():
-                progress_bar.update(1)
                 logger.debug(
-                    f"training: completed batch {progress_bar.n} "
+                    f"training: completed batch {progress_bar.total_steps} "
                     f"memory used: "
                     f"{psutil.Process(os.getpid()).memory_info()[0] / 1e6:0.2f}MB"
                 )
@@ -1039,12 +1041,13 @@ def train_online(self, dataset):
         ) as batcher:
 
             # training step loop
-            progress_bar = tqdm(
-                desc="Training online",
-                total=batcher.steps_per_epoch,
-                file=sys.stdout,
-                disable=is_progressbar_disabled(),
-            )
+            progress_bar_config = {
+                "desc": "Training online",
+                "total": batcher.steps_per_epoch,
+                "file": sys.stdout,
+                "disable": is_progressbar_disabled(),
+            }
+            progress_bar = LudwigProgressBar(self.report_tqdm_to_ray, progress_bar_config, True)
 
             while not batcher.last_batch():
                 batch = batcher.next_batch()
@@ -1100,7 +1103,9 @@ def append_metrics(self, dataset_name, results, metrics_log, tables, progress_tr
         return metrics_log, tables
 
     def evaluation(self, dataset, dataset_name, metrics_log, tables, batch_size, progress_tracker):
-        predictor = Predictor(self.model, batch_size=batch_size, horovod=self.horovod)
+        predictor = Predictor(
+            self.model, batch_size=batch_size, horovod=self.horovod, report_tqdm_to_ray=self.report_tqdm_to_ray
+        )
         metrics, predictions = predictor.batch_evaluation(dataset, collect_predictions=False, dataset_name=dataset_name)
 
         self.append_metrics(dataset_name, metrics, metrics_log, tables, progress_tracker)
diff --git a/ludwig/progress_bar.py b/ludwig/progress_bar.py
new file mode 100644
index 00000000000..7725c84c1a6
--- /dev/null
+++ b/ludwig/progress_bar.py
@@ -0,0 +1,73 @@
+import uuid
+
+import tqdm
+
+try:
+    import ray.train as rt  # noqa: E402
+except ImportError:
+    rt = None
+
+
+class LudwigProgressBarActions:
+    CREATE = "create"
+    UPDATE = "update"
+    CLOSE = "close"
+
+
+class LudwigProgressBar:
+    def __init__(self, report_to_ray, config, is_coordinator):
+        if report_to_ray and rt is None:
+            raise ValueError("Set report_to_ray=True but ray is not installed. Run `pip install ray`")
+
+        self.id = str(uuid.uuid4())[-8:]
+
+        self.report_to_ray = report_to_ray
+        self.is_coordinator = is_coordinator
+        self.config = config
+
+        self.total_steps = 0
+        self.progress_bar = None
+
+        if not self.report_to_ray:
+            if self.is_coordinator:
+                self.progress_bar = tqdm.tqdm(**config)
+        else:
+            if "file" in self.config:
+                self.config.pop("file")
+            # All processes need to call ray.train.report since ray has a lock that blocks
+            # a process when calling report if there are processes that haven't called it. Similar
+            # to a distributed checkpoint. Therefore we pass the flag to the driver
+            rt.report(
+                progress_bar={
+                    "id": self.id,
+                    "config": self.config,
+                    "action": LudwigProgressBarActions.CREATE,
+                    "is_coordinator": self.is_coordinator,
+                }
+            )
+
+    def update(self, steps):
+        self.total_steps += steps
+        if self.progress_bar:
+            self.progress_bar.update(steps)
+        elif self.report_to_ray:
+            rt.report(
+                progress_bar={
+                    "id": self.id,
+                    "update_by": steps,
+                    "is_coordinator": self.is_coordinator,
+                    "action": LudwigProgressBarActions.UPDATE,
+                }
+            )
+
+    def close(self):
+        if self.progress_bar:
+            self.progress_bar.close()
+        elif self.report_to_ray:
+            rt.report(
+                progress_bar={
+                    "id": self.id,
+                    "is_coordinator": self.is_coordinator,
+                    "action": LudwigProgressBarActions.CLOSE,
+                }
+            )
diff --git a/tests/integration_tests/test_ray.py b/tests/integration_tests/test_ray.py
index 08da4b03bb7..f72687c4259 100644
--- a/tests/integration_tests/test_ray.py
+++ b/tests/integration_tests/test_ray.py
@@ -43,7 +43,6 @@
     number_feature,
     sequence_feature,
     set_feature,
-    spawn,
     text_feature,
     timeseries_feature,
     train_with_backend,
@@ -147,7 +146,6 @@ def split(data_parquet):
     return train_fname, val_fname, test_fname
 
 
-@spawn
 def run_test_with_features(
     input_features,
     output_features,
@@ -462,3 +460,19 @@ def test_tune_batch_size_lr():
             model = run_api_experiment(config, dataset=dataset_parquet, backend_config=backend_config)
             assert model.config[TRAINER]["batch_size"] != "auto"
             assert model.config[TRAINER]["learning_rate"] != "auto"
+
+
+@pytest.mark.distributed
+def test_ray_progress_bar():
+    # This is a simple test that is just meant to make sure that the progress bar isn't breaking
+    input_features = [
+        sequence_feature(reduce_output="sum"),
+    ]
+    output_features = [
+        binary_feature(bool2str=["No", "Yes"]),
+    ]
+    run_test_with_features(
+        input_features,
+        output_features,
+        df_engine="dask",
+    )


Code review comments:
diff_hunk: @@ -0,0 +1,73 @@
+import uuid
+
+import tqdm
+
+try:
+    import ray.train as rt  # noqa: E402
path: ludwig/progress_bar.py
text: @user1:
Is the `noqa: E402` necessary?

@author:
I though E402 also applies if it is wrapped in a try except 

@user1:
It seems like it may not be, i.e. https://github.com/ludwig-ai/ludwig/blob/master/ludwig/utils/data_utils.py#L45-L52

One way to test whether the annotation is needed (I think) is to try removing the annotation and see if flake8 (should be enabled automatically in VSCode for the Ludwig repo) complains.

@author:
You are right my friend

diff_hunk: @@ -272,6 +273,30 @@ def tune_learning_rate_fn(
         hvd.shutdown()
 
 
+class TqdmCallback(rt.TrainingCallback):
+    def __init__(self):
+        super().__init__()
+        self.progess_bars = {}
+
+    def process_results(self, results, **info):
path: ludwig/backend/ray.py
text: @user1:
Add type hints and a docstring. What is results?

diff_hunk: @@ -1039,12 +1041,13 @@ def train_online(self, dataset):
         ) as batcher:
 
             # training step loop
-            progress_bar = tqdm(
-                desc="Training online",
-                total=batcher.steps_per_epoch,
-                file=sys.stdout,
-                disable=is_progressbar_disabled(),
-            )
+            progress_bar_config = {
+                "desc": "Training online",
+                "total": batcher.steps_per_epoch,
+                "file": sys.stdout,
+                "disable": is_progressbar_disabled(),
+            }
+            progress_bar = LudwigProgressBar(self.report_tqdm_to_ray, progress_bar_config, True)
path: ludwig/models/trainer.py
text: @user1:
Why is this constructed with `True` as the argument for `is_coordinator` directly instead of `self.is_coordinator()`?

@author:
good find