[SGD] v2 prototype: BackendExecutor and TorchBackend implementation

[amogkam]

Why are these changes needed?

Related issue number

Checks

• I've run scripts/format.sh to lint the changes in this PR.
• I've included any doc changes needed for https://docs.ray.io/en/master/.
• I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/
• Testing Strategy
  • Unit tests
  • Release tests
  • This PR is not tested :(

[matthewdeng]

Trainer -> BackendExecutor

[amogkam]

The user isn't aware of BackendExecutor right?

[richardliaw]

There's a lot of state that gets tossed around between Backend, BackendExecutor, TorchConfig.

Can we instead reduce the places where we're keeping track of state?

here's one attempt. the tldr is that TorchBackend becomes a purely "functional" callback and doesn't need to hold state itself.

diff --git a/python/ray/util/sgd/v2/backends/backend.py b/python/ray/util/sgd/v2/backends/backend.py
index 99ea5a20b..734425c0c 100644
--- a/python/ray/util/sgd/v2/backends/backend.py
+++ b/python/ray/util/sgd/v2/backends/backend.py
@@ -59,12 +59,14 @@ class BackendExecutor:
         self._num_gpus_per_worker = num_gpus_per_worker
 
         self.worker_group = DeactivatedWorkerGroup()
+        self._backend = get_backend(self._backend_config)
 
     def start(self):
         """Starts the worker group."""
         self.worker_group = WorkerGroup(self._num_workers,
                                         self._num_cpus_per_worker,
                                         self._num_gpus_per_worker)
+        self._backend.on_start(self.worker_group, self.backend_config)
 
     def execute(self, train_func: Callable[[], T]) -> Iterator[Any]:
         """Executes training function on all workers and yield results.
@@ -156,6 +158,7 @@ class BackendExecutor:
 
     def shutdown(self):
         """Shuts down the workers in the worker group."""
+        self._backend.on_shutdown(self.worker_group)
         self.worker_group.shutdown()
         self.worker_group = DeactivatedWorkerGroup()
 
diff --git a/python/ray/util/sgd/v2/backends/torch.py b/python/ray/util/sgd/v2/backends/torch.py
index 53d72e098..6170f2b94 100644
--- a/python/ray/util/sgd/v2/backends/torch.py
+++ b/python/ray/util/sgd/v2/backends/torch.py
@@ -91,57 +91,57 @@ def shutdown_torch():
         torch.cuda.empty_cache()
 
 
-class TorchExecutor(BackendExecutor):
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        self._backend_config.validate(name="torch")
-
-        if self._backend_config.backend is None:
-            if self._num_gpus_per_worker > 0:
-                self.backend = "nccl"
-            else:
-                self.backend = "gloo"
-
-    def start(self):
-        super().start()
-        if self._num_workers > 1:
-
+class TorchBackend:
+    # def __init__(self, backend_config, use_gpu=False):
+    #     self._backend_config = backend_config
+
+    #     # Can we actually just resolve this in config dataclass?
+    #     self._backend_config.validate(name="torch")
+
+    #     # Can we actually just resolve this too in config dataclass?
+    #     if self._backend_config.backend is None:
+    #         if use_gpu:
+    #             self.backend = "nccl"
+    #         else:
+    #             self.backend = "gloo"
+
+    def on_start(worker_group, backend_config):
+        if len(worker_group) > 1:
             def get_address():
                 addr = ray.util.get_node_ip_address()
                 port = find_free_port()
                 return addr, port
 
-            master_addr, master_port = self.worker_group.execute_single(
+            master_addr, master_port = worker_group.execute_single(
                 0, get_address)
 
-            if self._backend_config.init_method == "env":
+            if backend_config.init_method == "env":
 
                 def set_env_vars(addr, port):
                     os.environ["MASTER_ADDR"] = addr
                     os.environ["MASTER_PORT"] = str(port)
 
-                self.worker_group.execute(
+                worker_group.execute(
                     set_env_vars, addr=master_addr, port=master_port)
                 url = "env://"
-            elif self._backend_config == "tcp":
+            elif backend_config == "tcp":
                 url = f"tcp://{master_addr}:{master_port}"
             else:
                 raise ValueError(
                     f"The provided init_method ("
-                    f"{self._backend_config.init_method} is not supported.")
+                    f"{backend_config.init_method} is not supported.")
 
-            for i in range(len(self.worker_group)):
-                self.worker_group.execute_single(
+            for i in range(len(worker_group)):
+                worker_group.execute_single(
                     i,
                     setup_torch_process_group,
-                    backend=self.backend,
+                    backend=self._backend_config.backend,
                     world_rank=i,
-                    world_size=len(self.worker_group),
+                    world_size=len(worker_group),
                     init_method=url,
-                    timeout_s=self._backend_config.timeout_s)
+                    timeout_s=backend_config.timeout_s)
 
-    def shutdown(self):
-        self.worker_group.execute_single(
+    def on_shutdown(worker_group):
+        worker_group.execute_single(
             0, torch.distributed.destroy_process_group)
-        self.worker_group.execute(shutdown_torch)
-        super().shutdown()
+        worker_group.execute(shutdown_torch)

[amogkam]

Ok I addressed all the comments and removed all the reporting functionality. It would be great if you guys could take another look.

[matthewdeng]

Can we move tests for individual backends into their own test files?

[richardliaw]

looks good! merge when tests pass.