torchmonarch-nightly 2025.8.2__cp313-cp313-manylinux2014_x86_64.whl → 2025.9.4__cp313-cp313-manylinux2014_x86_64.whl

tests/test_allocator.py

@@ -15,13 +15,11 @@ import os
 import subprocess
 import sys
 import unittest
-from datetime import timedelta
 from time import sleep
 from typing import Generator, Optional
 from unittest import mock
 
 import cloudpickle
-import pytest
 
 import torch
 import torch.distributed as dist
@@ -55,8 +53,6 @@ from monarch.tools.network import get_sockaddr
 from torch.distributed.elastic.utils.distributed import get_free_port
 from torchx.specs import AppState
 
-_100_MILLISECONDS = timedelta(milliseconds=100)
-
 SERVER_READY = "monarch.tools.commands.server_ready"
 UNUSED = "__UNUSED__"
 
@@ -170,10 +166,9 @@ class TestSetupActorInAllocator(unittest.IsolatedAsyncioTestCase):
 
         spec = AllocSpec(AllocConstraints(), gpus=1, hosts=1)
         allocator = LocalAllocator()
-        alloc = await allocator.allocate(spec)
-
-        proc_mesh = await ProcMesh.from_alloc(alloc, setup=setup_multiple_env_vars)
+        alloc = allocator.allocate(spec)
 
+        proc_mesh = ProcMesh.from_alloc(alloc, setup=setup_multiple_env_vars)
         try:
             actor = await proc_mesh.spawn("env_check", EnvCheckActor)
 
@@ -197,9 +192,9 @@ class TestSetupActorInAllocator(unittest.IsolatedAsyncioTestCase):
 
         spec = AllocSpec(AllocConstraints(), gpus=1, hosts=1)
         allocator = LocalAllocator()
-        alloc = await allocator.allocate(spec)
+        alloc = allocator.allocate(spec)
 
-        proc_mesh = await ProcMesh.from_alloc(alloc, setup=setup_with_rank)
+        proc_mesh = ProcMesh.from_alloc(alloc, setup=setup_with_rank)
 
         try:
             actor = await proc_mesh.spawn("env_check", EnvCheckActor)
@@ -253,10 +248,11 @@ class TestRemoteAllocator(unittest.IsolatedAsyncioTestCase):
                 allocator = RemoteAllocator(
                     world_id="test_remote_allocator",
                     initializer=StaticRemoteAllocInitializer(host1, host2),
-                    heartbeat_interval=_100_MILLISECONDS,
                 )
-                alloc = await allocator.allocate(spec)
-                await ProcMesh.from_alloc(alloc)
+                alloc = allocator.allocate(spec)
+                await alloc.initialized
+                pm = ProcMesh.from_alloc(alloc)
+                await pm.initialized
 
     async def test_call_allocate_twice(self) -> None:
         class DeletingAllocInitializer(StaticRemoteAllocInitializer):
@@ -275,15 +271,17 @@ class TestRemoteAllocator(unittest.IsolatedAsyncioTestCase):
             allocator = RemoteAllocator(
                 world_id="test_remote_allocator",
                 initializer=initializer,
-                heartbeat_interval=_100_MILLISECONDS,
             )
 
             spec = AllocSpec(AllocConstraints(), host=1, gpu=1)
 
-            await allocator.allocate(spec)
+            alloc = allocator.allocate(spec)
+            await alloc.initialized
+
             self.assertEqual([host1], initializer.addrs)
 
-            await allocator.allocate(spec)
+            alloc = allocator.allocate(spec)
+            await alloc.initialized
             self.assertEqual([], initializer.addrs)
 
     async def test_throws_when_initializer_returns_empty_addrs(self) -> None:
@@ -301,9 +299,10 @@ class TestRemoteAllocator(unittest.IsolatedAsyncioTestCase):
             allocator = RemoteAllocator(
                 world_id="test_remote_allocator",
                 initializer=empty_initializer,
-                heartbeat_interval=_100_MILLISECONDS,
             )
-            await allocator.allocate(AllocSpec(AllocConstraints(), host=1, gpu=1))
+            await allocator.allocate(
+                AllocSpec(AllocConstraints(), host=1, gpu=1)
+            ).initialized
 
     async def test_allocate_2d_mesh(self) -> None:
         hosts = 2
@@ -316,10 +315,9 @@ class TestRemoteAllocator(unittest.IsolatedAsyncioTestCase):
             allocator = RemoteAllocator(
                 world_id="test_remote_allocator",
                 initializer=StaticRemoteAllocInitializer(host1, host2),
-                heartbeat_interval=_100_MILLISECONDS,
             )
-            alloc = await allocator.allocate(spec)
-            proc_mesh = await ProcMesh.from_alloc(alloc)
+            alloc = allocator.allocate(spec)
+            proc_mesh = ProcMesh.from_alloc(alloc)
             actor = await proc_mesh.spawn("test_actor", TestActor)
 
             values = await actor.compute_world_size.call(
@@ -335,11 +333,10 @@ class TestRemoteAllocator(unittest.IsolatedAsyncioTestCase):
             allocator = RemoteAllocator(
                 world_id="test_remote_allocator",
                 initializer=StaticRemoteAllocInitializer(host1, host2),
-                heartbeat_interval=_100_MILLISECONDS,
             )
 
-            alloc = await allocator.allocate(spec)
-            proc_mesh = await ProcMesh.from_alloc(alloc)
+            alloc = allocator.allocate(spec)
+            proc_mesh = ProcMesh.from_alloc(alloc)
             # XXX - it is not clear why this trying to use
             # async code in a sync context.
             with fake_sync_state():
@@ -348,7 +345,7 @@ class TestRemoteAllocator(unittest.IsolatedAsyncioTestCase):
             with self.assertRaises(
                 RuntimeError, msg="`ProcMesh` has already been stopped"
             ):
-                proc_mesh.spawn("test_actor", TestActor).get()
+                proc_mesh.spawn("test_actor", TestActor).initialized.get()
             del actor
 
     async def test_wrong_address(self) -> None:
@@ -362,14 +359,14 @@ class TestRemoteAllocator(unittest.IsolatedAsyncioTestCase):
             allocator = RemoteAllocator(
                 world_id="test_remote_allocator",
                 initializer=StaticRemoteAllocInitializer(wrong_host),
-                heartbeat_interval=_100_MILLISECONDS,
             )
-            alloc = await allocator.allocate(spec)
+            alloc = allocator.allocate(spec)
+            await alloc.initialized
 
             with self.assertRaisesRegex(
                 Exception, r"no process has ever been allocated.*"
             ):
-                await ProcMesh.from_alloc(alloc)
+                await ProcMesh.from_alloc(alloc).initialized
 
     async def test_init_failure(self) -> None:
         class FailInitActor(Actor):
@@ -385,10 +382,9 @@ class TestRemoteAllocator(unittest.IsolatedAsyncioTestCase):
             allocator = RemoteAllocator(
                 world_id="helloworld",
                 initializer=StaticRemoteAllocInitializer(host1, host2),
-                heartbeat_interval=_100_MILLISECONDS,
             )
             spec = AllocSpec(AllocConstraints(), host=2, gpu=2)
-            proc_mesh = await ProcMesh.from_alloc(await allocator.allocate(spec))
+            proc_mesh = ProcMesh.from_alloc(allocator.allocate(spec))
             actor_mesh = await proc_mesh.spawn("actor", FailInitActor)
 
             with self.assertRaisesRegex(
@@ -405,10 +401,9 @@ class TestRemoteAllocator(unittest.IsolatedAsyncioTestCase):
             allocator = RemoteAllocator(
                 world_id="test_remote_allocator",
                 initializer=StaticRemoteAllocInitializer(host1, host2),
-                heartbeat_interval=_100_MILLISECONDS,
             )
-            alloc = await allocator.allocate(spec)
-            proc_mesh = await ProcMesh.from_alloc(alloc)
+            alloc = allocator.allocate(spec)
+            proc_mesh = ProcMesh.from_alloc(alloc)
             actor = await proc_mesh.spawn("test_actor", TestActor)
 
             await proc_mesh.stop()
@@ -416,7 +411,7 @@ class TestRemoteAllocator(unittest.IsolatedAsyncioTestCase):
             with self.assertRaises(
                 RuntimeError, msg="`ProcMesh` has already been stopped"
             ):
-                await proc_mesh.spawn("test_actor", TestActor)
+                await proc_mesh.spawn("test_actor", TestActor).initialized
 
             # TODO(agallagher): It'd be nice to test that this just fails
             # immediately, trying to access the wrapped actor mesh, but right
@@ -431,10 +426,9 @@ class TestRemoteAllocator(unittest.IsolatedAsyncioTestCase):
             allocator = RemoteAllocator(
                 world_id="test_remote_allocator",
                 initializer=StaticRemoteAllocInitializer(host1, host2),
-                heartbeat_interval=_100_MILLISECONDS,
             )
-            alloc = await allocator.allocate(spec)
-            proc_mesh = await ProcMesh.from_alloc(alloc)
+            alloc = allocator.allocate(spec)
+            proc_mesh = ProcMesh.from_alloc(alloc)
             with self.assertRaises(ValueError, msg="foo"):
                 async with proc_mesh:
                     actor = await proc_mesh.spawn("test_actor", TestActor)
@@ -444,7 +438,7 @@ class TestRemoteAllocator(unittest.IsolatedAsyncioTestCase):
             with self.assertRaises(
                 RuntimeError, msg="`ProcMesh` has already been stopped"
             ):
-                await proc_mesh.spawn("test_actor", TestActor)
+                await proc_mesh.spawn("test_actor", TestActor).initialized
 
             # TODO(agallagher): It'd be nice to test that this just fails
             # immediately, trying to access the wrapped actor mesh, but right
@@ -467,11 +461,10 @@ class TestRemoteAllocator(unittest.IsolatedAsyncioTestCase):
             allocator = RemoteAllocator(
                 world_id="test_remote_allocator",
                 initializer=StaticRemoteAllocInitializer(host1, host2),
-                heartbeat_interval=_100_MILLISECONDS,
             )
-            alloc = await allocator.allocate(spec)
-            proc_mesh = await ProcMesh.from_alloc(alloc, setup=setup_env_vars)
-
+            alloc = allocator.allocate(spec)
+            proc_mesh = ProcMesh.from_alloc(alloc, setup=setup_env_vars)
+            await proc_mesh.initialized
             try:
                 actor = await proc_mesh.spawn("env_check", EnvCheckActor)
 
@@ -494,10 +487,9 @@ class TestRemoteAllocator(unittest.IsolatedAsyncioTestCase):
             allocator = RemoteAllocator(
                 world_id="test_remote_allocator",
                 initializer=StaticRemoteAllocInitializer(host1, host2),
-                heartbeat_interval=_100_MILLISECONDS,
             )
-            alloc = await allocator.allocate(spec)
-            proc_mesh = await ProcMesh.from_alloc(alloc)
+            alloc = allocator.allocate(spec)
+            proc_mesh = ProcMesh.from_alloc(alloc)
             # We can nest multiple context managers on the same mesh, the innermost
             # one closes the mesh and it cannot be used after that.
             async with proc_mesh:
@@ -507,7 +499,7 @@ class TestRemoteAllocator(unittest.IsolatedAsyncioTestCase):
                 with self.assertRaises(
                     RuntimeError, msg="`ProcMesh` has already been stopped"
                 ):
-                    await proc_mesh.spawn("test_actor", TestActor)
+                    await proc_mesh.spawn("test_actor", TestActor).initialized
                 # Exiting a second time should not raise an error.
 
             # TODO(agallagher): It'd be nice to test that this just fails
@@ -524,13 +516,12 @@ class TestRemoteAllocator(unittest.IsolatedAsyncioTestCase):
             allocator = RemoteAllocator(
                 world_id="test_remote_allocator",
                 initializer=StaticRemoteAllocInitializer(host1),
-                heartbeat_interval=_100_MILLISECONDS,
             )
             with self.assertRaisesRegex(
                 Exception, "no process has ever been allocated on"
             ):
-                alloc = await allocator.allocate(spec)
-                await ProcMesh.from_alloc(alloc)
+                alloc = allocator.allocate(spec)
+                await ProcMesh.from_alloc(alloc).initialized
 
     async def test_stacked_1d_meshes(self) -> None:
         # create two stacked actor meshes on the same host
@@ -540,19 +531,17 @@ class TestRemoteAllocator(unittest.IsolatedAsyncioTestCase):
             allocator_a = RemoteAllocator(
                 world_id="a",
                 initializer=StaticRemoteAllocInitializer(host1_a),
-                heartbeat_interval=_100_MILLISECONDS,
             )
             allocator_b = RemoteAllocator(
                 world_id="b",
                 initializer=StaticRemoteAllocInitializer(host1_b),
-                heartbeat_interval=_100_MILLISECONDS,
             )
 
             spec_a = AllocSpec(AllocConstraints(), host=1, gpu=2)
             spec_b = AllocSpec(AllocConstraints(), host=1, gpu=6)
 
-            proc_mesh_a = await ProcMesh.from_alloc(await allocator_a.allocate(spec_a))
-            proc_mesh_b = await ProcMesh.from_alloc(await allocator_b.allocate(spec_b))
+            proc_mesh_a = ProcMesh.from_alloc(allocator_a.allocate(spec_a))
+            proc_mesh_b = ProcMesh.from_alloc(allocator_b.allocate(spec_b))
 
             actor_a = await proc_mesh_a.spawn("actor_a", TestActor)
             actor_b = await proc_mesh_b.spawn("actor_b", TestActor)
@@ -576,7 +565,9 @@ class TestRemoteAllocator(unittest.IsolatedAsyncioTestCase):
                 RuntimeError,
                 r"slurm:///123 does not exist or is in a terminal state",
             ):
-                await allocator.allocate(AllocSpec(AllocConstraints(), host=1, gpu=1))
+                await allocator.allocate(
+                    AllocSpec(AllocConstraints(), host=1, gpu=1)
+                ).initialized
 
     async def test_torchx_remote_alloc_initializer_no_match_label_gt_1_meshes(
         self,
@@ -599,9 +590,10 @@ class TestRemoteAllocator(unittest.IsolatedAsyncioTestCase):
                 RuntimeError,
                 r"2 proc meshes in slurm:///123, please specify the mesh name as a match label `procmesh.monarch.meta.com/name`",
             ):
-                await allocator.allocate(AllocSpec(AllocConstraints(), host=1, gpu=1))
+                await allocator.allocate(
+                    AllocSpec(AllocConstraints(), host=1, gpu=1)
+                ).initialized
 
-    @pytest.mark.oss_skip  # pyre-ignore[56] TODO T228752279
     async def test_torchx_remote_alloc_initializer_no_match_label_1_mesh(self) -> None:
         server = ServerSpec(
             name=UNUSED,
@@ -612,30 +604,26 @@ class TestRemoteAllocator(unittest.IsolatedAsyncioTestCase):
                     name="x",
                     num_hosts=1,
                     transport="tcp",
-                    hostnames=["localhost"],
+                    hostnames=["0.0.0.0"],
                 )
             ],
         )
         port = get_free_port()
-        with remote_process_allocator(addr=f"tcp!{get_sockaddr('localhost', port)}"):
+        with remote_process_allocator(addr=f"tcp!{get_sockaddr('0.0.0.0', port)}"):
             with mock.patch(SERVER_READY, return_value=server):
                 initializer = TorchXRemoteAllocInitializer("local:///test", port=port)
                 allocator = RemoteAllocator(
                     world_id="test",
                     initializer=initializer,
-                    heartbeat_interval=_100_MILLISECONDS,
-                )
-                alloc = await allocator.allocate(
-                    AllocSpec(AllocConstraints(), host=1, gpu=4)
                 )
-                proc_mesh = await ProcMesh.from_alloc(alloc)
+                alloc = allocator.allocate(AllocSpec(AllocConstraints(), host=1, gpu=4))
+                proc_mesh = ProcMesh.from_alloc(alloc)
                 actor = await proc_mesh.spawn("test_actor", TestActor)
                 results = await actor.compute_world_size.call(
                     master_addr="0.0.0.0", master_port=get_free_port()
                 )
                 self.assert_computed_world_size(results, 4)  # 1x4 mesh
 
-    @pytest.mark.oss_skip  # pyre-ignore[56] TODO T228752279
     async def test_torchx_remote_alloc_initializer_with_match_label(self) -> None:
         server = ServerSpec(
             name=UNUSED,
@@ -646,20 +634,19 @@ class TestRemoteAllocator(unittest.IsolatedAsyncioTestCase):
                     name="x",
                     num_hosts=1,
                     transport="tcp",
-                    hostnames=["localhost"],
+                    hostnames=["0.0.0.0"],
                 )
             ],
         )
         port = get_free_port()
-        with remote_process_allocator(addr=f"tcp!{get_sockaddr('localhost', port)}"):
+        with remote_process_allocator(addr=f"tcp!{get_sockaddr('0.0.0.0', port)}"):
             with mock.patch(SERVER_READY, return_value=server):
                 initializer = TorchXRemoteAllocInitializer("local:///test", port=port)
                 allocator = RemoteAllocator(
                     world_id="test",
                     initializer=initializer,
-                    heartbeat_interval=_100_MILLISECONDS,
                 )
-                alloc = await allocator.allocate(
+                alloc = allocator.allocate(
                     AllocSpec(
                         AllocConstraints(
                             match_labels={ALLOC_LABEL_PROC_MESH_NAME: "x"}
@@ -668,7 +655,7 @@ class TestRemoteAllocator(unittest.IsolatedAsyncioTestCase):
                         gpu=3,
                     )
                 )
-                proc_mesh = await ProcMesh.from_alloc(alloc)
+                proc_mesh = ProcMesh.from_alloc(alloc)
                 actor = await proc_mesh.spawn("test_actor", TestActor)
                 results = await actor.compute_world_size.call(
                     master_addr="0.0.0.0", master_port=get_free_port()
@@ -698,7 +685,7 @@ class TestRemoteAllocator(unittest.IsolatedAsyncioTestCase):
             with self.assertRaisesRegex(RuntimeError, r"'y' not found in job: test"):
                 initializer = TorchXRemoteAllocInitializer("local:///test")
                 allocator = RemoteAllocator(world_id="test", initializer=initializer)
-                alloc = await allocator.allocate(
+                alloc = allocator.allocate(
                     AllocSpec(
                         AllocConstraints(
                             match_labels={ALLOC_LABEL_PROC_MESH_NAME: "y"}
@@ -707,7 +694,8 @@ class TestRemoteAllocator(unittest.IsolatedAsyncioTestCase):
                         gpu=1,
                     )
                 )
-                await ProcMesh.from_alloc(alloc)
+                await alloc.initialized
+                await ProcMesh.from_alloc(alloc).initialized
 
     async def test_log(self) -> None:
         # create a mesh to log to both stdout and stderr
@@ -716,12 +704,11 @@ class TestRemoteAllocator(unittest.IsolatedAsyncioTestCase):
             allocator = RemoteAllocator(
                 world_id="test_actor_logger",
                 initializer=StaticRemoteAllocInitializer(host),
-                heartbeat_interval=_100_MILLISECONDS,
             )
 
             spec = AllocSpec(AllocConstraints(), host=1, gpu=2)
 
-            proc_mesh = await ProcMesh.from_alloc(await allocator.allocate(spec))
+            proc_mesh = ProcMesh.from_alloc(allocator.allocate(spec))
 
             # Generate aggregated log every 1 second.
             await proc_mesh.logging_option(True, 1)