initial commit for hostmesh v1

Author: allenwang28

src/forge/controller/launcher.py

@@ -15,15 +15,19 @@
 import monarch
 
 import torchx.specs as specs
-
 from monarch._rust_bindings.monarch_hyperactor.alloc import AllocConstraints
+from monarch._rust_bindings.monarch_hyperactor.channel import ChannelTransport
+
+from monarch._rust_bindings.monarch_hyperactor.config import configure
 from monarch._src.actor.allocator import RemoteAllocator, TorchXRemoteAllocInitializer
 from monarch.actor import Actor, endpoint, ProcMesh
 from monarch.tools import commands
 from monarch.tools.commands import info
 from monarch.tools.components import hyperactor
 from monarch.tools.config import Config, Workspace
 
+from forge.env import MONARCH_HOSTMESH_V1
+
 from forge.types import Launcher, LauncherConfig
 
 _MAST_AVAILABLE = False
@@ -116,7 +120,8 @@ async def remote_setup(self, procs: ProcMesh) -> tuple[str, int]:
 
 class Slurmlauncher(BaseLauncher):
     async def initialize(self) -> None:
-        pass
+        if MONARCH_HOSTMESH_V1.get_value():
+            configure(default_transport=ChannelTransport.Tcp)
 
     async def get_allocator(self, name: str, num_hosts: int) -> tuple[Any, Any, str]:
         appdef = hyperactor.host_mesh(
@@ -172,6 +177,9 @@ def __init__(self, cfg: LauncherConfig | None = None):
         self.job_name = self.cfg.job_name or self.create_job_name()
 
     async def initialize(self) -> None:
+        if MONARCH_HOSTMESH_V1.get_value():
+            configure(default_transport=ChannelTransport.MetaTlsWithHostname)
+
         await self.launch_mast_job()
 
     async def get_allocator(self, name: str, num_hosts: int) -> tuple[Any, Any, str]:

src/forge/controller/provisioner.py

@@ -4,6 +4,8 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+from forge.env import MONARCH_HOSTMESH_V1
+
 """Remote and local resource manager for allocation and provisioning."""
 import asyncio
 import functools
@@ -14,7 +16,8 @@
 import uuid
 
 from monarch._src.actor.shape import NDSlice, Shape
-from monarch.actor import Actor, endpoint, HostMesh, ProcMesh, this_host
+from monarch.actor import Actor, endpoint, ProcMesh
+
 from monarch.tools import commands
 
 from forge.controller.launcher import BaseLauncher, get_launcher
@@ -27,6 +30,14 @@
 logger.setLevel(logging.DEBUG)
 
 
+if MONARCH_HOSTMESH_V1.get_value():
+    from monarch._src.actor.v1.host_mesh import HostMesh, this_host
+
+    logger.info("Using Monarch HostMesh v1...")
+else:
+    from monarch.actor import HostMesh, this_host
+
+
 def _get_port() -> str:
     with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
         s.bind(("localhost", 0))
@@ -148,14 +159,29 @@ async def create_host_mesh(self, name: str, num_hosts: int) -> HostMesh:
         alloc, alloc_constraints, server_name = await self.launcher.get_allocator(
             name, num_hosts
         )
-        return (
-            HostMesh(
+
+        if MONARCH_HOSTMESH_V1.get_value():
+            # We are asking Monarch to allocate a single process on
+            # every host, reflected in the Extent we provide below.
+
+            # Technically, this is ["hosts", "procs"] but to reduce
+            # confusion on its relationship with procs elsewhere,
+            # we call it "no_dim".
+
+            # TODO - remove this once Monarch supports HostMesh without it.
+            host_mesh = HostMesh.allocate_nonblocking(
+                name=name,
+                extent=Extent(["hosts", "no_dim"], [num_hosts, 1]),
+                allocator=alloc,
+                alloc_constraints=alloc_constraints,
+            )
+        else:
+            host_mesh = HostMesh(
                 Shape(["hosts"], NDSlice.new_row_major([num_hosts])),
                 allocator=alloc,
                 alloc_constraints=alloc_constraints,
-            ),
-            server_name,
-        )
+            )
+        return host_mesh, server_name
 
     async def get_proc_mesh(
         self,
@@ -256,7 +282,7 @@ def bootstrap(env: dict[str, str]):
                     env_vars[env_var.name] = str(env_var.get_value())
 
             procs = host_mesh.spawn_procs(
-                per_host={"gpus": num_procs},
+                per_host={"procs": num_procs},
                 bootstrap=functools.partial(bootstrap, env=env_vars),
             )
 

src/forge/env.py

@@ -82,12 +82,6 @@ def get_value(self) -> Any:
     description="Sets Monarch's stderr log level, i.e. set to 'info' or 'debug'",
 )
 
-TORCHSTORE_USE_RDMA = EnvVar(
-    name="TORCHSTORE_RDMA_ENABLED",
-    default=False,
-    description="Whether or not to use RDMA in TorchStore.",
-)
-
 RUST_BACKTRACE = EnvVar(
     name="RUST_BACKTRACE",
     default="full",
@@ -106,6 +100,18 @@ def get_value(self) -> Any:
     description="Sets the maximum frame length for Monarch's actor message delivery in bytes.",
 )
 
+MONARCH_HOSTMESH_V1 = EnvVar(
+    name="MONARCH_HOSTMESH_V1",
+    default=True,
+    description="Whether or not to use Monarch's experimental hostmesh v1 APIs",
+)
+
+TORCHSTORE_USE_RDMA = EnvVar(
+    name="TORCHSTORE_RDMA_ENABLED",
+    default=True,
+    description="Whether or not to use RDMA in TorchStore.",
+)
+
 
 @functools.cache
 def all_env_vars() -> list[EnvVar]: