Allow users to specify a different cluster address when initializing Ray (#915)

* Allow users to specify a different cluster address when initializing Ray.

Author: Jun Gong

alpa/api.py

@@ -23,6 +23,7 @@
 
 
 def init(cluster: str = "ray",
+         cluster_address: Optional[str] = None,
          num_nodes: Optional[int] = None,
          num_devices_per_node: Optional[int] = None,
          namespace: Optional[str] = "alpa_default_space"):
@@ -40,6 +41,12 @@ def init(cluster: str = "ray",
         Possible choices: {"local", "ray"}.
         "local" means using all local devices on a single node.
         "ray" means using all devices in a ray cluster.
+      cluster_address: Address of the distributed cluster.
+        If cluster is "ray", this parameter can be used to specify a different
+          address that will be used to initialize the ray cluster.
+          E.g., "ray://123.45.67.89:10001". If not specified, "auto" will be
+          used instead.
+        Ignored if cluster is "local".
       num_nodes: The number of nodes.
       num_devices_per_node: The number of devices per node.
     """
@@ -49,7 +56,8 @@ def init(cluster: str = "ray",
         return
     is_initialized = True
 
-    init_global_cluster(cluster, num_nodes, num_devices_per_node, namespace)
+    init_global_cluster(cluster, cluster_address, num_nodes,
+                        num_devices_per_node, namespace)
 
 
 def shutdown():

alpa/device_mesh.py

@@ -2304,6 +2304,7 @@ def profile_all(self, *args, **kwargs):
 
 
 def init_global_cluster(cluster: str,
+                        cluster_address: Optional[str] = None,
                         num_nodes: Optional[int] = None,
                         num_devices_per_node: Optional[int] = None,
                         namespace: Optional[str] = None):
@@ -2313,7 +2314,8 @@ def init_global_cluster(cluster: str,
         global_physical_mesh = LocalPhysicalDeviceMesh()
     elif cluster == "ray":
         if not ray.is_initialized():
-            ray.init(address="auto",
+            ray_addr = cluster_address if cluster_address else "auto"
+            ray.init(address=ray_addr,
                      ignore_reinit_error=True,
                      namespace=namespace)
         update_jax_platform("cpu")