[Data] - Rich Progress bar doesn't render during execution for Train

[goutamvenkat-anyscale]

What happened + What you expected to happen

Output:

(SplitCoordinator pid=190940) ✔️  Dataset train_115_0 execution finished in 5.44 seconds 1… 1.02k/1.… 199.58 ro
(SplitCoordinator pid=190940)   │ Active/total resources: Active & requested resources: 0/111 CPU, 480.0B/153…
(SplitCoordinator pid=190940)   │                                                                             
(SplitCoordinator pid=190940)   ├─   ReadRange 100% ━━━━━━━━━━ 1.02k/1.02k [ 0:00:05 , 200.47 row/s ]         
(SplitCoordinator pid=190940)   │    0.0 CPU, 0.0B object store; Tasks: 0; Queued blocks: 0;                  
(SplitCoordinator pid=190940)   ├─   split(1, equal=True) 100% ━━━━━━━━ 1.02k/1.02k [ 0:00:05 , 200.42 row/s ]
(SplitCoordinator pid=190940)   │    0.0 CPU, 480.0B object store; Tasks: 0: [0/255 objects local]; Queued bl…
(SplitCoordinator pid=190940)                                                                                 ✔️  Dataset train_115_0 execution finished in 5.44 seconds
(RayTrainWorker pid=486040, ip=10.0.236.116) Exiting prefetcher's background thread
(TrainController pid=190675) [State Transition] RUNNING -> SHUTTING_DOWN.
(TrainController pid=190675) [State Transition] SHUTTING_DOWN -> FINISHED.

Note: the finished log does show up at the very end of execution – but nothing shows up while execution is ongoing.

Versions / Dependencies

Ray 2.50.1

Reproduction script

import ray.train
import ray.data
from ray.train.v2.api.data_parallel_trainer import DataParallelTrainer


def train_fn(config):
    list(ray.train.get_dataset_shard("train").iter_torch_batches(batch_size=32))
    
ray.data.DataContext.get_current().enable_rich_progress_bars = True

ds = ray.data.range(1024)

trainer = DataParallelTrainer(train_fn, datasets={"train": ds})
trainer.fit()

Issue Severity

Low: It annoys or frustrates me.

[kyuds]

will look into this

[kyuds]

offline discussion with @goutamvenkat-anyscale -- reason for this issue is ray train will work on ray workers --> aka different process. Before, we handled this with tqdm_ray, and so we'll default to tqdm_ray before we figure out something that will work for distributed rich. I'll be submitting a PR for this within the day.