Changes to kd_monarch to reflect default port changes to remote-process-allocator and monarch CLI component

Summary:
Updates `kd_monarch` with the capability to dynamically query for information needed to create and use the `MastAllocator` in `controller.py`.

Specifically it makes it possible to get these fields from the provided `ALLOCATOR_JOB_NAME`

1. `(hosts, gpus)` ProcMesh dimensions
2. RemoteProcessAllocator `port` (no need to hard code or have the user provide)
3. Task group name

Also changes the remote-process-allocator (aka `hyperactor mesh-worker`) entrypoint's binary name in `kd_monarch_pkg` to be `monarch_bootstrap` to be consistent with its equivalent in OSS (hence makes the TorchX AppDef portable)

Next:
~~[6/n] Have kd_monarch use the default component (the custom mast.py is no longer needed). Update the README with updated instructions.~~
[7/n] Remove rust CLI in favor of all-python (we delegate to torchx for most things anyways)
[8/n] Add E2E unittest using the local_cwd scheduler (actually run a mini-trainer actor)
[9/n] Write an oss hyperactor mesh-worker entrypoint binary
[10/n] Author a Dockerfile that sets up the environment (much like fbpkgs do it for internal runs)
[11/n] Author a TorchXAllocator

Reviewed By: vidhyav

Differential Revision: D75162358

fbshipit-source-id: d2756ffb529b3069baf94cd9dfdde00b8ef2ede4

Author: kiukchung

tools/mesh_spec.py

@@ -28,6 +28,7 @@ class MeshSpec:
     num_hosts: int
     host_type: str
     gpus: int
+    port: int = DEFAULT_REMOTE_ALLOCATOR_PORT
 
 
 def _tag(mesh_name: str, tag_template: str) -> str:
@@ -47,6 +48,7 @@ def mesh_spec_from_metadata(appdef: specs.AppDef, mesh_name: str) -> Optional[Me
                 num_hosts=role.num_replicas,
                 host_type=appdef.metadata.get(_tag(mesh_name, _TAG_HOST_TYPE), ""),
                 gpus=int(appdef.metadata.get(_tag(mesh_name, _TAG_GPUS), "-1")),
+                port=role.port_map.get("mesh", DEFAULT_REMOTE_ALLOCATOR_PORT),
             )
 
     return None
@@ -82,6 +84,18 @@ class ServerSpec:
     state: specs.AppState
     meshes: list[MeshSpec]
 
+    def get_mesh_spec(self, mesh_name: str) -> MeshSpec:
+        for mesh_spec in self.meshes:
+            if mesh_spec.name == mesh_name:
+                return mesh_spec
+
+        raise ValueError(
+            f"Mesh: '{mesh_name}' not found in job: {self.name}. Try one of: {self.get_mesh_names()}"
+        )
+
+    def get_mesh_names(self) -> list[str]:
+        return [m.name for m in self.meshes]
+
     def to_json(self) -> dict[str, Any]:
         """Returns the JSON form of this struct that can be printed to console by:
 

tools/tests/commands_test.py

@@ -12,7 +12,9 @@
 from monarch.tools import commands
 from monarch.tools.commands import component_args_from_cli
 
-from monarch.tools.config import defaults
+from monarch.tools.config import (  # @manual=//monarch/tools/config/meta:defaults
+    defaults,
+)
 from monarch.tools.mesh_spec import MeshSpec, ServerSpec
 from torchx.specs import AppDef, AppDryRunInfo, AppState, AppStatus, Role
 
@@ -68,7 +70,14 @@ def test_info(
 
         appdef = AppDef(
             name="monarch_test_123",
-            roles=[Role(name="trainer", image="__unused__", num_replicas=4)],
+            roles=[
+                Role(
+                    name="trainer",
+                    image="__unused__",
+                    num_replicas=4,
+                    port_map={"mesh": 26501},
+                )
+            ],
             metadata={
                 "monarch/meshes/trainer/host_type": "gpu.medium",
                 "monarch/meshes/trainer/gpus": "2",
@@ -86,6 +95,7 @@ def test_info(
                         num_hosts=4,
                         host_type="gpu.medium",
                         gpus=2,
+                        port=26501,
                     )
                 ],
             ),

tools/tests/mesh_spec_test.py

@@ -4,6 +4,7 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+# pyre-strict
 import json
 import unittest
 from dataclasses import asdict
@@ -12,6 +13,7 @@
     mesh_spec_from_metadata,
     mesh_spec_from_str,
     MeshSpec,
+    ServerSpec,
     tag_as_metadata,
 )
 
@@ -87,7 +89,37 @@ def test_mesh_spec_can_dump_as_json(self) -> None:
   "name": "trainer",
   "num_hosts": 4,
   "host_type": "gpu.medium",
-  "gpus": 2
+  "gpus": 2,
+  "port": 26600
 }
 """
         self.assertEqual(expected.strip("\n"), json.dumps(asdict(mesh_spec), indent=2))
+
+
+class ServerSpecTest(unittest.TestCase):
+    def get_test_server_spec(self) -> ServerSpec:
+        return ServerSpec(
+            name="monarch-foo-1a2b3c",
+            state=specs.AppState.RUNNING,
+            meshes=[
+                MeshSpec(name="trainer", num_hosts=4, host_type="gpu.medium", gpus=2),
+                MeshSpec(name="generator", num_hosts=8, host_type="gpu.small", gpus=1),
+            ],
+        )
+
+    def test_get_mesh_spec(self) -> None:
+        server_spec = self.get_test_server_spec()
+        mesh_spec = server_spec.get_mesh_spec("trainer")
+
+        self.assertEqual("trainer", mesh_spec.name)
+        self.assertEqual(4, mesh_spec.num_hosts)
+        self.assertEqual(2, mesh_spec.gpus)
+        self.assertEqual("gpu.medium", mesh_spec.host_type)
+
+    def test_get_mesh_spec_not_found(self) -> None:
+        server_spec = self.get_test_server_spec()
+        with self.assertRaisesRegex(
+            ValueError,
+            r"Mesh: 'worker' not found in job: monarch-foo-1a2b3c. Try one of: \['trainer', 'generator'\]",
+        ):
+            server_spec.get_mesh_spec("worker")