fix mapping issue + a repro test

Author: DNXie

assets/versions.sh

@@ -14,6 +14,6 @@ PYTORCH_VERSION="2.9.0.dev20250905"
 VLLM_BRANCH="v0.10.0"
 
 # Commit hashes
-MONARCH_COMMIT="2f14096083b1cc1dac6ae15220e4135bc23f9dd3"
+MONARCH_COMMIT="main"
 TORCHTITAN_COMMIT="d0e25450bcac2332359b13fbda430dc701f073d4"
 TORCHSTORE_COMMIT="662299faf4fd50ee30bd9aa3f4ce8c0e2db1d310"

src/forge/actors/generator.py

@@ -154,7 +154,7 @@ async def launch(  # pyright: ignore[reportIncompatibleMethodOverride]
         worker_procs = await get_proc_mesh(process_config=process_config)
 
         # Then, grab a single host from the workers...
-        host_mesh = await host_mesh_from_proc(worker_procs)
+        host_mesh = await host_mesh_from_proc(worker_procs._uid)
         singleton_slice = {k: slice(0, 1) for k in host_mesh.extent.keys()}
         host_mesh = host_mesh.slice(**singleton_slice)
 

src/forge/controller/provisioner.py

@@ -294,6 +294,11 @@ def bootstrap(env: dict[str, str]):
                 per_host={"procs": num_procs},
                 bootstrap=functools.partial(bootstrap, env=env_vars),
             )
+            uid = str(uuid.uuid4())
+            # Generate a unique ID to map procmesh to hostmesh
+            procs._uid = uid
+            print(f"Allocating procmesh with uid={uid}")
+            print(f"Allocating procs._uid: {procs._uid}")
 
             if with_gpus:
                 # Set up environment variables for PyTorch distributed...
@@ -319,28 +324,35 @@ def bootstrap(env: dict[str, str]):
                 self._server_names.append(server_name)
                 self._proc_server_map[procs] = server_name
 
-            self._proc_host_map[procs] = host_mesh
+            self._proc_host_map[uid] = host_mesh
 
         # Spawn LocalFetcherActor for this ProcMesh and register with GlobalLoggingActor.
         # When called, the LocalFetcherActor is broadcast by Monarch to all ranks in the ProcMesh.
         if not FORGE_DISABLE_METRICS.get_value():
             from forge.observability.metric_actors import get_or_create_metric_logger
 
             _ = await get_or_create_metric_logger(procs, process_name=mesh_name)
-        return procs
+
+        print(f"Returning procmesh with uid={uid}")
+        print(f"Returning procs._uid: {procs._uid}")
+        return procs, uid
 
     @endpoint
-    async def host_mesh_from_proc(self, proc_mesh: ProcMesh):
-        if proc_mesh not in self._proc_host_map:
+    async def host_mesh_from_proc(self, uid: str | None):
+        # uid: str | None = getattr(proc_mesh, "_uid", None)
+        print(f"self._proc_host_map: {self._proc_host_map}")
+        print(f"proc_mesh._uid: {uid}")
+        if uid is None or uid not in self._proc_host_map:
             raise ValueError(
                 "The proc mesh was not allocated with an associated hostmesh."
             )
-        return self._proc_host_map[proc_mesh]
+        return self._proc_host_map[uid]
 
     @endpoint
     async def stop_proc_mesh(self, proc_mesh: ProcMesh):
         """Stops a proc mesh."""
-        if proc_mesh not in self._proc_host_map:
+        uid: str | None = getattr(proc_mesh, "_uid", None)
+        if uid is None or uid not in self._proc_host_map:
             logger.warning(
                 f"proc mesh {proc_mesh} was requested to be stopped, but was either already stopped or "
                 "was never registered with the provisioner."
@@ -363,7 +375,7 @@ async def stop_proc_mesh(self, proc_mesh: ProcMesh):
             if proc_mesh in self._proc_server_map:
                 server_name = self._proc_server_map[proc_mesh]
                 commands.kill(server_name)
-            del self._proc_host_map[proc_mesh]
+            del self._proc_host_map[uid]
 
     @endpoint
     def register_service(self, service: "ServiceInterface") -> None:
@@ -464,7 +476,7 @@ async def get_proc_mesh(
 
     """
     provisioner = await get_or_create_provisioner()
-    return await provisioner.get_proc_mesh.call_one(
+    procs, uid = await provisioner.get_proc_mesh.call_one(
         num_procs=process_config.procs,
         with_gpus=process_config.with_gpus,
         num_hosts=process_config.hosts,
@@ -474,17 +486,20 @@ async def get_proc_mesh(
         port=port,
         addr=addr,
     )
+    setattr(procs, "_uid", uid)
+    print(f"Setting procs._uid: {procs._uid}")
+    return procs
 
 
-async def host_mesh_from_proc(proc_mesh: ProcMesh):
+async def host_mesh_from_proc(uid: str | None):
     """Returns the host mesh that allocated the original proc_mesh.
 
     This functionality will be enabled in Monarch, so this is a temporary
     API.
 
     """
     provisioner = await get_or_create_provisioner()
-    return await provisioner.host_mesh_from_proc.call_one(proc_mesh)
+    return await provisioner.host_mesh_from_proc.call_one(uid)
 
 
 async def register_service(service: "ServiceInterface") -> None:

test.py

@@ -0,0 +1,37 @@
+# Copyright (c) Meta Platforms, Inc.
+# All rights reserved.
+#
+# Minimal repro: Provisioner host_mesh_from_proc() UID mapping bug
+#
+# Run this with:
+#   python -m forge.tests.test_provisioner_uid_mapping
+
+import asyncio
+
+# import pytest
+
+from forge.controller.provisioner import (
+    get_or_create_provisioner,
+    get_proc_mesh,
+    stop_proc_mesh,
+)
+from forge.types import ProcessConfig
+
+
+# @pytest.mark.asyncio
+async def test_provisioner_host_mesh_lookup_uid_mapping():
+    prov = await get_or_create_provisioner()
+    pm = await get_proc_mesh(
+        ProcessConfig(procs=1, with_gpus=False, hosts=None, mesh_name="uid_repro")
+    )
+    # UID is attached locally by the helper
+    assert hasattr(pm, "_uid") and pm._uid, "missing _uid on returned ProcMesh"
+    print(f"✅ got ProcMesh with UID {pm._uid}")
+    hm = await prov.host_mesh_from_proc.call_one(pm._uid)  # if pass pm, _uid is None
+    assert hm is not None
+    await stop_proc_mesh(pm)
+    print("✅ repro passed")
+
+
+if __name__ == "__main__":
+    asyncio.run(test_provisioner_host_mesh_lookup_uid_mapping())