(monarch/tools) add ServerSpec.host0() method that can be used to set MASTER_ADDR for PTD (#419)

Summary:
Pull Request resolved: #419

Simplifies getting host0's hostname to set as MASTER_ADDR for torch.distributed.

Usage:

```
from monarch.tools.commands import get_or_create

server_info = await get_or_create(...)

# allocate proc mesh -> create actor (code omitted for brevity)...
trainer_actor.call(
  MASTER_ADDR=server_info.host0("trainer") # trainer mesh's 1st host
  MASTER_PORT=29500,
  ...
)
```

Reviewed By: suo

Differential Revision: D77690685

fbshipit-source-id: 1f37c370ac33282d3b9e34ee983a74acc91d68f6

Author: kiukchung

python/monarch/tools/mesh_spec.py

@@ -11,6 +11,7 @@
 
 from monarch.tools.network import get_sockaddr
 from torchx import specs
+from torchx.specs.api import is_terminal
 
 DEFAULT_REMOTE_ALLOCATOR_PORT = 26600
 
@@ -133,6 +134,53 @@ def server_handle(self) -> str:
     def is_running(self) -> bool:
         return self.state == specs.AppState.RUNNING
 
+    def host0(self, mesh_name: str) -> str:
+        """The hostname of the first node in the given mesh.
+        The return value of this method can be used to set `MASTER_ADDR` env var for torch.distributed.
+
+        NOTE: the state of this server must be RUNNING for this method to return a valid value.
+
+        Usage:
+
+        .. code-block::python
+            from monarch.tools.commands import get_or_create
+
+            server_info = await get_or_create(...)
+            assert server_info.is_running
+
+            # allocate proc mesh -> create actor (code omitted for brevity)...
+
+            trainer_actor.call(
+                MASTER_ADDR=server_info.host0("trainer") # trainer mesh's 1st host
+                MASTER_PORT=29500,
+                ...
+            )
+
+        NOTE: The ordering of the hostnames is exactly the same as what comes back from the underlying
+        scheduler's `describe_job` or `list_*` API. Please find the exact semantics in the
+        respective scheduler's implementation in https://github.com/pytorch/torchx/tree/main/torchx/schedulers.
+        """
+        mesh_spec = self.get_mesh_spec(mesh_name)
+        if self.is_running:
+            # hostnames are only valid when the server is RUNNING
+            if not mesh_spec.hostnames:
+                raise RuntimeError(f"{self.server_handle} does not have any hosts")
+            return mesh_spec.hostnames[0]
+        elif self.state in [specs.AppState.SUBMITTED, specs.AppState.PENDING]:
+            raise RuntimeError(
+                f"{self.server_handle} is {self.state}."
+                f" Use `monarch.tools.commands.server_ready()` to wait for the server to be {specs.AppState.RUNNING}"
+            )
+        elif is_terminal(self.state):
+            raise RuntimeError(
+                f"{self.server_handle} is {self.state}."
+                " Use `monarch.tools.commands.get_or_create()` to create a new server"
+            )
+        else:
+            raise RuntimeError(
+                f"{self.server_handle} is in an invalid state: {self.state}. Please report this as a bug"
+            )
+
     def get_mesh_spec(self, mesh_name: str) -> MeshSpec:
         for mesh_spec in self.meshes:
             if mesh_spec.name == mesh_name:

python/tests/tools/test_mesh_spec.py

@@ -180,6 +180,17 @@ def get_test_server_spec(self) -> ServerSpec:
             ],
         )
 
+    def test_server_handle(self) -> None:
+        unused = specs.AppState.RUNNING
+
+        server = ServerSpec(name="foo", scheduler="slurm", meshes=[], state=unused)
+        self.assertEqual("slurm:///foo", server.server_handle)
+
+        server = ServerSpec(
+            name="foo", scheduler="slurm", namespace="prod", meshes=[], state=unused
+        )
+        self.assertEqual("slurm://prod/foo", server.server_handle)
+
     def test_get_mesh_spec(self) -> None:
         server_spec = self.get_test_server_spec()
         mesh_spec = server_spec.get_mesh_spec("trainer")
@@ -196,3 +207,55 @@ def test_get_mesh_spec_not_found(self) -> None:
             r"Mesh: 'worker' not found in job: monarch-foo-1a2b3c. Try one of: \['trainer', 'generator'\]",
         ):
             server_spec.get_mesh_spec("worker")
+
+    def _1_mesh_2_host_server_spec(self, state: specs.AppState) -> ServerSpec:
+        return ServerSpec(
+            name="foo",
+            scheduler="slurm",
+            meshes=[
+                MeshSpec(
+                    name="trainer",
+                    num_hosts=2,
+                    hostnames=["compute-node-0", "compute-node-1"],
+                )
+            ],
+            state=state,
+        )
+
+    def test_node0(self) -> None:
+        server = self._1_mesh_2_host_server_spec(specs.AppState.RUNNING)
+        self.assertEqual("compute-node-0", server.host0("trainer"))
+
+    def test_node0_server_in_terminal_state(self) -> None:
+        for terminal_state in [
+            specs.AppState.FAILED,
+            specs.AppState.SUCCEEDED,
+            specs.AppState.CANCELLED,
+        ]:
+            with self.subTest(terminal_state=terminal_state):
+                server = self._1_mesh_2_host_server_spec(terminal_state)
+                with self.assertRaisesRegex(
+                    RuntimeError,
+                    r"Use `monarch.tools.commands.get_or_create\(\)` to create a new server",
+                ):
+                    server.host0("trainer")
+
+    def test_node0_server_in_pending_state(self) -> None:
+        for pending_state in [specs.AppState.SUBMITTED, specs.AppState.PENDING]:
+            with self.subTest(pending_state=pending_state):
+                server = self._1_mesh_2_host_server_spec(pending_state)
+                with self.assertRaisesRegex(
+                    RuntimeError,
+                    r"Use `monarch.tools.commands.server_ready\(\)` to wait for the server to be RUNNING",
+                ):
+                    server.host0("trainer")
+
+    def test_node0_server_in_illegal_tate(self) -> None:
+        for illegal_state in [specs.AppState.UNSUBMITTED, specs.AppState.UNKNOWN]:
+            with self.subTest(illegal_state=illegal_state):
+                server = self._1_mesh_2_host_server_spec(illegal_state)
+                with self.assertRaisesRegex(
+                    RuntimeError,
+                    r"Please report this as a bug",
+                ):
+                    server.host0("trainer")