[RaySGD] Convert the head worker to a local model (#7746)

Why are these changes needed?

Running a worker on head (locally, not as a Ray actor) allows for easier handling of stateful stuff like logging and for easier debugging.

python/ray/util/sgd/torch/training_operator.py

@@ -1,10 +1,11 @@
 import collections
 
+from tqdm import tqdm
 import torch
 
 from ray.util.sgd.utils import (TimerCollection, AverageMeterCollection,
                                 NUM_SAMPLES)
-from ray.util.sgd.torch.constants import (SCHEDULER_STEP_EPOCH,
+from ray.util.sgd.torch.constants import (SCHEDULER_STEP_EPOCH, NUM_STEPS,
                                           SCHEDULER_STEP_BATCH, SCHEDULER_STEP)
 
 amp = None
@@ -55,7 +56,8 @@ class TrainingOperator:
                  world_rank,
                  criterion=None,
                  schedulers=None,
-                 use_fp16=False):
+                 use_fp16=False,
+                 use_tqdm=False):
         # You are not expected to override this method.
         self._models = models  # List of models
         assert isinstance(models, collections.Iterable), (
@@ -74,6 +76,7 @@ class TrainingOperator:
                     type(schedulers)))
         self._config = config
         self._use_fp16 = use_fp16
+        self._use_tqdm = use_tqdm
         self.global_step = 0
 
         if type(self) is TrainingOperator:
@@ -84,12 +87,8 @@ class TrainingOperator:
                         "TrainingOperator if using multi-scheduler, "
                         "multi-model or multi-optimizer training/validation.")
         self.timers = TimerCollection()
-        self.reporters = []
         self.setup(config)
 
-    def set_reporters(self, reporters):
-        self.reporters = reporters
-
     def _set_timers(self, timers):
         """Passes in the timers from the Runner."""
         self.timers = timers
@@ -142,8 +141,19 @@ class TrainingOperator:
         Returns:
             A dict of metrics from training.
         """
-        for r in self.reporters:
-            r.on_epoch_begin(info, self)
+        if self.use_tqdm and self.world_rank == 0:
+            desc = ""
+            if info is not None and "epoch_idx" in info:
+                if "num_epochs" in info:
+                    desc = "{}/{}e".format(info["epoch_idx"] + 1,
+                                           info["num_epochs"])
+                else:
+                    desc = "{}e".format(info["epoch_idx"] + 1)
+            _progress_bar = tqdm(
+                total=info[NUM_STEPS] or len(self.train_loader),
+                desc=desc,
+                unit="batch",
+                leave=False)
 
         metric_meters = AverageMeterCollection()
 
@@ -156,8 +166,10 @@ class TrainingOperator:
             batch_info.update(info)
             metrics = self.train_batch(batch, batch_info=batch_info)
 
-            for r in self.reporters:
-                r.on_batch_end(batch_info, metrics, self)
+            if self.use_tqdm and self.world_rank == 0:
+                _progress_bar.n = batch_idx + 1
+                if "train_loss" in metrics:
+                    _progress_bar.set_postfix({"loss": metrics["train_loss"]})
 
             if self.scheduler and batch_info.get(
                     SCHEDULER_STEP) == SCHEDULER_STEP_BATCH:
@@ -376,6 +388,11 @@ class TrainingOperator:
         """Whether the model and optimizer have been FP16 enabled."""
         return self._use_fp16
 
+    @property
+    def use_tqdm(self):
+        """Whether tqdm progress bars are enabled."""
+        return self._use_tqdm
+
 
 class _TestingOperator(TrainingOperator):
     def train_epoch(self, iterator, info):