Fix orca ray pytorch example (intel-analytics#3007)

* fix horovod pytorch exampe * fix bug * fix process group * fix style * fix tests * fix test * fix tests * revert ray context change
yangw1234 · Sep 26, 2021 · 3b8bcbf · 3b8bcbf
1 parent 3af2b86
commit 3b8bcbf
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 4 deletions.
diff --git a/python/orca/src/bigdl/orca/learn/pytorch/pytorch_ray_estimator.py b/python/orca/src/bigdl/orca/learn/pytorch/pytorch_ray_estimator.py
@@ -101,7 +101,7 @@ def __init__(
         if backend == "pytorch":
             cores_per_node = ray_ctx.ray_node_cpu_cores // workers_per_node
             num_nodes = ray_ctx.num_ray_nodes * workers_per_node
-            RemoteRunner = ray.remote(num_cpus=1)(TorchRunner)
+            RemoteRunner = ray.remote(num_cpus=cores_per_node)(TorchRunner)
             self.remote_workers = [
                 RemoteRunner.remote(**params) for i in range(num_nodes)
             ]
@@ -110,9 +110,10 @@ def __init__(
                 for i, worker in enumerate(self.remote_workers)
             ])
 
-            ip = ray.services.get_node_ip_address()
-            port = utils.find_free_port()
-            address = "tcp://{ip}:{port}".format(ip=ip, port=port)
+            head_worker = self.remote_workers[0]
+            address = ray.get(head_worker.setup_address.remote())
+
+            logger.info(f"initializing pytorch process group on {address}")
 
             ray.get([
                 worker.setup_torch_distribute.remote(address, i, num_nodes)

diff --git a/python/orca/src/bigdl/orca/learn/pytorch/torch_runner.py b/python/orca/src/bigdl/orca/learn/pytorch/torch_runner.py
@@ -47,6 +47,7 @@
 from zoo.orca.learn.pytorch.constants import SCHEDULER_STEP, NUM_STEPS
 from zoo.orca.learn.pytorch.training_operator import TrainingOperator
 from zoo.orca.learn.pytorch import utils
+from zoo.orca.learn.pytorch.utils import find_free_port
 
 logger = logging.getLogger(__name__)
 
@@ -124,6 +125,11 @@ def setup_horovod(self):
         self.setup_components_horovod()
         self.setup_operator()
 
+    def setup_address(self):
+        ip = ray.services.get_node_ip_address()
+        port = find_free_port()
+        return f"tcp://{ip}:{port}"
+
     def setup_torch_distribute(self, url, world_rank, world_size):
         import torch.distributed as dist
         dist.init_process_group(