Create a HostMesh per replica

Differential Revision: D91890958

Pull Request resolved: #746

Author: daniellepintz

src/forge/controller/launcher.py

@@ -34,11 +34,15 @@ def get_meshes_from_config(cfg: LauncherConfig) -> dict[str, int]:
     meshes: dict[str, int] = {}
 
     # Add services that need remote hosts
+    # Expand services with multiple replicas into per-replica meshes
     for service_name, service_cfg in cfg.services.items():
         hosts = getattr(service_cfg, "hosts", None)
         if hosts and hosts > 0:
-            mesh_name = service_cfg.mesh_name or service_name
-            meshes[mesh_name] = hosts
+            base_mesh_name = service_cfg.mesh_name or service_name
+            num_replicas = service_cfg.num_replicas
+            for replica_idx in range(num_replicas):
+                mesh_name = f"{base_mesh_name}_{replica_idx}"
+                meshes[mesh_name] = hosts
 
     # Add actors that need remote hosts
     for actor_name, actor_cfg in cfg.actors.items():
@@ -78,7 +82,7 @@ async def initialize(self) -> tuple[JobTrait, JobState]:
         # Create a single SlurmJob with all meshes
         logger.info(f"Creating SlurmJob with meshes: {meshes}")
         job = SlurmJob(
-            meshes=meshes,  # e.g., {"generator": 1, "trainer": 2, "ref_model": 1}
+            meshes=meshes,  # e.g., {"generator_0": 1, "generator_1": 1, "trainer": 2}
             slurm_args=slurm_args,
             job_name=self.cfg.job_name + "_workers" or "forge_job",
             time_limit="72:00:00",  # Default to 72 hours

src/forge/controller/provisioner.py

@@ -269,20 +269,12 @@ async def get_host_mesh(self, name: str) -> HostMesh:
             # for local jobs, return the local host
             return this_host()
 
-        # Strip replica suffix (e.g., "generator_0" -> "generator")
-        # Services append _{replica_idx} to mesh names
-        base_name = name
-        if "_" in name:
-            parts = name.rsplit("_", 1)
-            if len(parts) == 2 and parts[1].isdigit():
-                base_name = parts[0]
-
         # _job_state contains all the HostMeshes that were allocated as attributes, accessible by their name
         try:
-            host_mesh = getattr(self._job_state, base_name)
+            host_mesh = getattr(self._job_state, name)
         except AttributeError as e:
             raise RuntimeError(
-                f"Mesh '{name}' (base: '{base_name}') was not pre-allocated. "
+                f"Mesh '{name}' was not pre-allocated. "
                 "Make sure the mesh is defined in the launcher config."
             ) from e