torchmonarch-nightly 2025.6.27__cp312-cp312-manylinux2014_x86_64.whl

tests/test_device_mesh.py

@@ -0,0 +1,132 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# 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 pytest
+
+from monarch import DeviceMesh, NDSlice
+from monarch.common.client import Client
+from monarch.simulator.mock_controller import MockController
+
+
+class TestDeviceMesh:
+    def test_mesh_index(self) -> None:
+        fake_processes = NDSlice(offset=0, sizes=[2, 3, 4], strides=[12, 4, 1])
+        ctrl = MockController(1, False)
+        client = Client(ctrl, ctrl.world_size, ctrl.gpu_per_host)
+        dm = DeviceMesh(client, fake_processes, ("a", "b", "c"))
+        assert 0 == dm(a=0, b=0, c=0).processes[0]
+        x = dm(a=0, b=0)
+        assert x.processes[:] == fake_processes[0:4]
+        assert x.names == ("c",)
+        assert x.processes.sizes[0] == 4
+        x = dm(c=slice(None, None, 2))
+        assert x.processes[:] == fake_processes[::2]
+        x = dm(b=2, c=3)
+        assert x.processes[:] == (11, 23)
+        client.shutdown()
+
+    def test_mesh_reshape(self) -> None:
+        fake_processes = NDSlice(offset=0, sizes=[60, 24], strides=[24, 1])
+        ctrl = MockController(1, False)
+        client = Client(ctrl, ctrl.world_size, ctrl.gpu_per_host)
+        dm = DeviceMesh(client, fake_processes, ("host", "gpu"))
+        dm2 = dm.split(host=("dp", "pp"), gpu=("tp",), pp=4)
+        assert dm2.names == ("dp", "pp", "tp")
+        assert dm2.processes.sizes == [15, 4, 24]
+        assert dm2.processes.strides == [4 * 24, 24, 1]
+
+        dm3 = dm.rename(host="dp", gpu="tp")
+        assert dm3.names == ("dp", "tp")
+        assert dm.processes.strides == dm3.processes.strides
+        dm4 = dm.split(host=("dp", "pp"), gpu=("tp",), dp=4)
+        assert dm4.processes.sizes == [4, 15, 24]
+        dm5 = dm.split(host=("dp", "pp"), dp=60)
+        assert dm5.processes.sizes == [60, 1, 24]
+        dm6 = dm.split(host=("dp", "pp"), pp=60)
+        assert dm6.processes.sizes == [1, 60, 24]
+
+        with pytest.raises(ValueError, match="Cannot infer size"):
+            dm2 = dm.split(host=("dp", "pp"))
+
+        with pytest.raises(ValueError, match="unused size constraints"):
+            dm2 = dm.split(host=("dp", "pp"), pp=4, ddp=3)
+
+        dm2 = dm.rename(host="dp")
+        assert dm2.names == ("dp", "gpu")
+
+        with pytest.raises(ValueError, match="Duplicate dimension name"):
+            dm2 = dm.split(host=("dp", "pp"), gpu=("pp",), dp=3)
+
+        with pytest.raises(ValueError, match="evenly divided"):
+            dm2 = dm.split(host=("dp", "pp"), dp=7)
+
+        client.shutdown()
+
+    def test_flatten(self) -> None:
+        fake_processes = NDSlice(offset=0, sizes=[60, 24], strides=[24, 1])
+        ctrl = MockController(1, False)
+        client = Client(ctrl, ctrl.world_size, ctrl.gpu_per_host)
+        dm = DeviceMesh(client, fake_processes, ("host", "gpu"))
+        dm2 = dm.flatten("gpu")
+        assert dm2.names == ("gpu",)
+        assert dm2.processes.sizes == [60 * 24]
+        assert dm2.processes.strides == [1]
+        client.shutdown()
+
+        good_cases = [
+            NDSlice(offset=0, sizes=[100], strides=[1]),
+            NDSlice(offset=100, sizes=[8, 4, 2], strides=[8, 2, 1]),
+            NDSlice(offset=1, sizes=[4, 2], strides=[2, 1]),
+        ]
+        for slice in good_cases:
+            ctrl = MockController(1, False)
+            client = Client(ctrl, ctrl.world_size, ctrl.gpu_per_host)
+            dm = DeviceMesh(
+                client, slice, tuple(f"dim{d}" for d in range(len(slice.sizes)))
+            )
+            dm2 = dm.flatten("outer")
+            assert dm2.names == ("outer",)
+            assert dm2.processes.strides == [1]
+            assert list(slice) == list(dm2.processes)
+            client.shutdown()
+
+        # Test some bad ones (sparse slices).
+        bad_cases = [
+            NDSlice(offset=0, sizes=[100], strides=[2]),
+            NDSlice(offset=0, sizes=[64, 32], strides=[64, 1]),
+        ]
+        for slice in bad_cases:
+            with pytest.raises(ValueError, match="cannot flatten sparse mesh"):
+                ctrl = MockController(1, False)
+                client = Client(ctrl, ctrl.world_size, ctrl.gpu_per_host)
+                dm = DeviceMesh(
+                    client, slice, tuple(f"dim{d}" for d in range(len(slice.sizes)))
+                )
+                dm.flatten("bad_dim")
+                client.shutdown()
+
+    def test_worker_mesh_init(self) -> None:
+        from monarch.worker.worker import DeviceMesh as WorkerDeviceMesh
+
+        processes = NDSlice(offset=0, sizes=[3, 4], strides=[4, 1])
+        wdm = WorkerDeviceMesh(0, ("a", "b"), processes, rank=1)
+        a, b = wdm.dims["a"], wdm.dims["b"]
+        assert b.members == [0, 1, 2, 3]
+        assert b.rank == 1
+
+        assert a.members == [1, 5, 9]
+        assert a.rank == 0
+
+        wdm = WorkerDeviceMesh(0, ("a", "b"), processes, rank=6)
+        a, b = wdm.dims["a"], wdm.dims["b"]
+        assert b.members == [4, 5, 6, 7]
+        assert b.rank == 2
+        assert a.members == [2, 6, 10]
+        assert a.rank == 1
+
+        processes = NDSlice(offset=0, sizes=[3, 4, 2], strides=[8, 2, 1])
+        wdm = WorkerDeviceMesh(0, ("a", "b", "c"), processes, rank=10)

tests/test_fault_tolerance.py

@@ -0,0 +1,398 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import random
+import time
+from typing import List, Optional
+
+import pytest
+import torch
+
+try:
+    from later.unittest import TestCase
+except ModuleNotFoundError:
+    from unittest import TestCase
+
+from monarch import fetch_shard, no_mesh, remote
+from monarch.common.device_mesh import DeviceMesh, DeviceMeshStatus
+from monarch.common.invocation import DeviceException, RemoteException
+from monarch.rust_backend_mesh import MeshWorld, PoolDeviceMeshProvider
+from monarch.rust_local_mesh import (
+    Bootstrap,
+    local_mesh_provider,
+    local_meshes_and_bootstraps,
+    LoggingLocation,
+    SocketType,
+    SupervisionParams,
+)
+
+
+def _do_bogus_tensor_work(
+    x: torch.Tensor, y: torch.Tensor, fail_rank: Optional[int] = None
+) -> torch.Tensor:
+    return x + y  # real function actually does x @ y
+
+
+do_bogus_tensor_work = remote(
+    "monarch.worker._testing_function.do_bogus_tensor_work",
+    propagate=_do_bogus_tensor_work,
+)
+
+
+def mesh_provider(
+    meshes: int = 2,
+    hosts_per_mesh: int = 1,
+    gpus_per_host: int = 1,
+    # pyre-fixme[11]: Annotation `DeviceMeshProvider` is not defined as a type.
+) -> tuple[PoolDeviceMeshProvider, Bootstrap]:
+    return local_mesh_provider(
+        meshes=meshes,
+        hosts_per_mesh=hosts_per_mesh,
+        gpus_per_host=gpus_per_host,
+        socket_type=SocketType.UNIX,
+        logging_location=LoggingLocation.DEFAULT,
+        supervision_params=SupervisionParams(
+            update_timeout_in_sec=10,  # Fail fast
+            query_interval_in_sec=1,
+            update_interval_in_sec=1,
+        ),
+        auto_epoch=True,
+    )
+
+
+def local_meshes(
+    meshes: int = 2,
+    hosts_per_mesh: int = 1,
+    gpus_per_host: int = 1,
+) -> tuple[list[DeviceMesh], Bootstrap]:
+    return local_meshes_and_bootstraps(
+        meshes=meshes,
+        hosts_per_mesh=hosts_per_mesh,
+        gpus_per_host=gpus_per_host,
+        socket_type=SocketType.UNIX,
+        logging_location=LoggingLocation.DEFAULT,
+        supervision_params=SupervisionParams(
+            update_timeout_in_sec=10,  # Fail fast
+            query_interval_in_sec=1,
+            update_interval_in_sec=1,
+        ),
+        auto_epoch=True,
+    )
+
+
+# Set global timeout--sandcastle's timeout is 600s. A test that sandcastle times
+# out is not counted as a failure, so we set a more restrictive timeout to
+# ensure we see a hard failure in CI.
+# The timeout is set to 250s as the failover takes longer than other tests.
+@pytest.mark.timeout(250)
+class TestFaultTolerance(TestCase):
+    def test_mesh_provider(self) -> None:
+        # Create multiple meshes using mesh provider
+        replicas = 4
+        provider, bootstrap = mesh_provider(meshes=replicas)
+        meshes: list[DeviceMesh] = []
+        while len(meshes) < replicas:
+            dm = provider.new_mesh()
+            meshes.append(dm)
+
+        statuses = provider._root_client.world_status()
+        for _, status in statuses.items():
+            assert (
+                DeviceMeshStatus(status) != DeviceMeshStatus.UNHEALTHY
+            ), f"unexpected unhealthy mesh; world status: {statuses}"
+
+        # Check that all meshes are initially live
+        for mesh in meshes:
+            with mesh.activate():
+                t = torch.ones(1)
+                local_t = fetch_shard(t).result()
+            assert torch.equal(local_t, torch.ones(1))
+
+        # Simulate a failure by killing one of the processes
+        bootstrap.processes[-1].kill()
+
+        # Find unhealthy mesh
+        # Mix user and device errors
+        unhealthy_meshes = []
+        for mesh in meshes:
+            with mesh.activate():
+                # Send a call to trigger a failure
+                x = torch.rand(3, 4)
+                y = torch.rand(3, 4)
+                z = do_bogus_tensor_work(x, y)
+                try:
+                    _ = fetch_shard(z).result()
+                except RemoteException:
+                    pass
+                except DeviceException as e:
+                    # Device error
+                    unhealthy_meshes.append(mesh)
+                    mesh.exit(e)
+
+        self.assertEqual(len(unhealthy_meshes), 1)
+
+        # World status will transition to unhealthy
+        has_unhealth = False
+        unhealthy_statuses = []
+        while not has_unhealth:
+            statuses = provider._root_client.world_status()
+            for _, status in statuses.items():
+                if DeviceMeshStatus(status) == DeviceMeshStatus.UNHEALTHY:
+                    has_unhealth = True
+                    unhealthy_statuses = statuses
+                    break
+            time.sleep(1)
+
+        # Unhealthy worlds will be evicted
+        has_unhealth = True
+        healthy_statuses = []
+        while has_unhealth:
+            has_unhealth = False
+            statuses = provider._root_client.world_status()
+            healthy_statuses = statuses
+            for _, status in statuses.items():
+                if DeviceMeshStatus(status) == DeviceMeshStatus.UNHEALTHY:
+                    has_unhealth = True
+                    break
+            time.sleep(1)
+
+        # A worker world will be evicted
+        self.assertEqual(len(healthy_statuses), len(unhealthy_statuses) - 1)
+
+    def test_worker_supervision_failure(self) -> None:
+        meshes, bootstrap = local_meshes(meshes=1)
+        assert len(meshes) == 1
+        mesh = meshes[0]
+
+        # Check the mesh initially functional
+        with mesh.activate():
+            t = torch.ones(1)
+            local_t = fetch_shard(t).result()
+        assert torch.equal(local_t, torch.ones(1))
+
+        # Simulate a failure by killing one of the processes
+        bootstrap.processes[-1].kill()
+
+        # A device error will be raised
+        with mesh.activate():
+            t = torch.ones(1)
+            with self.assertRaisesRegex(DeviceException, r"crashed"):
+                local_t = fetch_shard(t).result()
+
+    def test_multi_mesh_failure_isolation(self) -> None:
+        replicas = 4
+        provider, bootstrap = mesh_provider(meshes=replicas)
+        meshes: list[DeviceMesh] = []
+        while len(meshes) < replicas:
+            dm = provider.new_mesh()
+            meshes.append(dm)
+
+        # Check the meshes initially functional
+        for mesh in meshes:
+            with mesh.activate():
+                t = torch.ones(1)
+                local_t = fetch_shard(t).result()
+            assert torch.equal(local_t, torch.ones(1))
+
+        initial_size = len(provider._root_client.world_status())
+
+        # Simulate a failure by killing one of the processes
+        bootstrap.processes[-1].kill()
+
+        # Mix user and device errors
+        healthy_meshes = []
+        unhealthy_meshes = []
+        for mesh in meshes:
+            with mesh.activate():
+                # Send a call to trigger a failure
+                x = torch.rand(3, 4)
+                y = torch.rand(3, 4)
+                z = do_bogus_tensor_work(x, y)
+                try:
+                    _ = fetch_shard(z).result()
+                except RemoteException:
+                    # User error
+                    fetch_shard(x).result()
+                    healthy_meshes.append(mesh)
+                except DeviceException as e:
+                    # Device error
+                    unhealthy_meshes.append(mesh)
+                    mesh.exit(e)
+
+        self.assertEqual(len(healthy_meshes), replicas - 1)
+        self.assertEqual(len(unhealthy_meshes), 1)
+
+        while True:
+            size = len(provider._root_client.world_status())
+            if size == initial_size - 2:
+                break
+
+        # The healthy meshes should still be functional
+        for mesh in healthy_meshes:
+            with mesh.activate():
+                t = torch.ones(1)
+                local_t = fetch_shard(t).result()
+            assert torch.equal(local_t, torch.ones(1))
+
+    def test_out_of_order_receive(self) -> None:
+        meshes, _ = local_meshes(meshes=8)
+
+        # Check the meshes initially functional
+        ts = []
+        for i, mesh in enumerate(meshes):
+            with mesh.activate():
+                t = torch.ones(i + 1)
+                ts.append(t)
+
+        # Shuffle the meshes to makes sure the client is able to dispatch results in order
+        indices = list(range(len(meshes)))
+        shuffled_meshes = list(zip(indices, meshes, ts))
+        random.shuffle(shuffled_meshes)
+        for i, mesh, t in shuffled_meshes:
+            with mesh.activate():
+                local_t = fetch_shard(t).result()
+            assert torch.equal(local_t, torch.ones(i + 1))
+
+    def test_mesh_shrink_and_grow(self) -> None:
+        # Create multiple meshes using mesh provider
+        replicas = 4
+        provider, bootstrap = mesh_provider(meshes=replicas)
+        meshes: list[DeviceMesh] = []
+        while len(meshes) < replicas:
+            dm = provider.new_mesh()
+            meshes.append(dm)
+
+        worlds = len(provider._root_client.world_status())
+        assigned_meshes = provider._mesh_map
+        assert len(assigned_meshes) == replicas
+
+        # Happy path
+        for i, mesh in enumerate(meshes):
+            with mesh.activate():
+                t = torch.ones(i + 1)
+                local_t = fetch_shard(t).result()
+            assert torch.equal(local_t, torch.ones(i + 1))
+
+        # Kill a worker
+        mesh_to_kill: MeshWorld = list(bootstrap.mesh_worlds.keys())[1]
+        procs = bootstrap.mesh_worlds[mesh_to_kill]
+        assert len(procs) == 2
+        procs[-1].kill()
+
+        # The killed mesh will become unhealthy
+        healthy_meshes = []
+        unhealthy_meshes = []
+        for i, mesh in enumerate(meshes):
+            with mesh.activate():
+                try:
+                    t = torch.ones(i + 1)
+                    local_t = fetch_shard(t).result()
+                    with no_mesh.activate():
+                        assert torch.equal(local_t, torch.ones(i + 1))
+                    healthy_meshes.append(mesh)
+                except DeviceException as e:
+                    unhealthy_meshes.append(mesh)
+                    mesh.exit(e)
+        assert len(healthy_meshes) == replicas - 1
+        assert len(unhealthy_meshes) == 1
+
+        # Restart the mesh
+        for proc in procs:
+            proc.kill()
+
+        # We should be able to acquire a new mesh without waiting for the old mesh to be evicted
+        (worker_world, controller_id) = mesh_to_kill
+        bootstrap.launch_mesh(controller_id=controller_id, worker_world=worker_world)
+
+        dm = provider.new_mesh()
+        healthy_meshes.append(dm)
+
+        # We could have 4 or 5 meshes depending on if the unhealthy mesh is evicted
+        assigned_meshes = provider._mesh_map
+        assert len(assigned_meshes) >= replicas
+
+        # We are happy again
+        assert len(healthy_meshes) == replicas
+        for i, mesh in enumerate(healthy_meshes):
+            with mesh.activate():
+                t = torch.ones(i + 1)
+                local_t = fetch_shard(t).result()
+            assert torch.equal(local_t, torch.ones(i + 1))
+
+        # Old world should be evicted and new world should be spawned. So we ended up with the same number of worlds.
+        while len((provider._root_client.world_status())) != worlds:
+            # We expect to evict both controller and worker worlds from the same mesh.
+            time.sleep(1)
+
+        # Eventually, we only have 4 healthy meshes
+        assigned_meshes = provider._mesh_map
+        while len(assigned_meshes) != replicas:
+            with self.assertRaisesRegex(
+                TimeoutError, r"Could not find a healthy world"
+            ):
+                _ = provider.new_mesh(timeout_in_sec=1)
+            assigned_meshes = provider._mesh_map
+            time.sleep(1)
+
+    def test_kill_controller(self) -> None:
+        # Create multiple meshes using mesh provider
+        replicas = 2
+        provider, bootstrap = mesh_provider(meshes=replicas)
+        meshes: list[DeviceMesh] = []
+        while len(meshes) < replicas:
+            dm = provider.new_mesh()
+            meshes.append(dm)
+
+        # Happy path
+        for i, mesh in enumerate(meshes):
+            with mesh.activate():
+                t = torch.ones(i + 1)
+                local_t = fetch_shard(t).result()
+            assert torch.equal(local_t, torch.ones(i + 1))
+
+        # Kill a controller
+        mesh_to_kill: MeshWorld = list(bootstrap.mesh_worlds.keys())[1]
+        procs = bootstrap.mesh_worlds[mesh_to_kill]
+        assert len(procs) == 2
+        procs[0].kill()
+
+        # We should be able to detect the failure
+        healthy_meshes = []
+        detected_failure = False
+        for i, mesh in enumerate(meshes):
+            with mesh.activate():
+                try:
+                    t = torch.ones(i + 1)
+                    local_t = fetch_shard(t).result()
+                    with no_mesh.activate():
+                        assert torch.equal(local_t, torch.ones(i + 1))
+                    healthy_meshes.append(mesh)
+                except DeviceException:
+                    detected_failure = True
+        assert len(healthy_meshes) == replicas - 1
+        assert detected_failure
+
+    def test_late_client_attaching(self) -> None:
+        provider, _ = mesh_provider(meshes=1)
+
+        # Wait for the meshes to be healthy
+        healthy_meshes = 0
+        while healthy_meshes < 2:
+            healthy_meshes = 0
+            statuses = provider._root_client.world_status()
+            for _, status in statuses.items():
+                if DeviceMeshStatus(status) == DeviceMeshStatus.LIVE:
+                    healthy_meshes += 1
+            time.sleep(1)
+
+        # Sleep long enough to allow those "hidden messages" to be sent
+        time.sleep(15)
+
+        # Those "hidden messages" should not cause a trouble before a client is ready
+        mesh = provider.new_mesh()
+        with mesh.activate():
+            t = torch.ones(1)
+            fetch_shard(t).result()

tests/test_future.py

@@ -0,0 +1,94 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# 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 time
+from typing import Callable
+
+import pytest
+from monarch import Future, RemoteException
+from monarch._rust_bindings.monarch_hyperactor.proc import (  # @manual=//monarch/monarch_extension:monarch_extension
+    ActorId,
+)
+from monarch.common import future
+from monarch.common.client import Client
+
+
+class TestFuture:
+    def test_future(self, monkeypatch: pytest.MonkeyPatch) -> None:
+        the_time: int = 0
+        the_messages: list[tuple[int | float, Callable[[], None]]] = []
+
+        class MockClient(Client):
+            def __init__(self):
+                pass
+
+            def handle_next_message(self, timeout) -> bool:
+                nonlocal the_time
+                if not the_messages:
+                    return False
+                time, action = the_messages[0]
+                if timeout is None or time <= the_time + timeout:
+                    the_time = time
+                    action()
+                    the_messages.pop(0)
+                    return True
+                else:
+                    the_time += timeout
+                    return False
+
+            def _request_status(self):
+                pass
+
+        client: Client = MockClient()
+
+        def mock_time() -> int:
+            return the_time
+
+        monkeypatch.setattr(time, "time", mock_time)
+        f = Future(client)
+        the_messages = [(1, lambda: f._set_result(4))]
+        assert not f.done()
+        with pytest.raises(TimeoutError):
+            f.result(timeout=0.5)
+        assert 4 == f.result(timeout=1)
+        assert f.exception() is None
+        assert f.done()
+        f = Future(client)
+        the_messages = [(1, lambda: None), (2, lambda: f._set_result(3))]
+        the_time = 0
+        assert 3 == f.result()
+        f = Future(client)
+        re = RemoteException(
+            0, Exception(), None, [], [], ActorId.from_string("unknown[0].unknown[0]")
+        )
+
+        the_messages = [(1, lambda: None), (2, lambda: f._set_result(re))]
+        the_time = 0
+        assert f.exception() is not None
+
+        f = Future(client)
+        the_messages = [(0, lambda: None), (0.2, lambda: f._set_result(7))]
+        the_time = 0
+        assert 7 == f.result(timeout=0.3)
+
+        fs = []
+
+        def setup() -> None:
+            nonlocal fs, the_messages
+            fs = [Future(client) for _ in range(4)]
+
+            # To avoid closure binding gotcha.
+            def set_at_time(f: Future, time: int) -> tuple[int, Callable[[], None]]:
+                return (time, lambda: f._set_result(time))
+
+            the_messages = [set_at_time(f, time) for time, f in enumerate(fs)]
+
+        setup()
+        assert {f.result() for f in future.stream(fs, timeout=2)} == {0, 1, 2}
+
+        setup()
+        assert {f.result() for f in future.stream(fs, timeout=3)} == {0, 1, 2, 3}