torchmonarch-nightly 2025.6.27__cp312-cp312-manylinux2014_x86_64.whl

tests/test_rust_backend.py

@@ -0,0 +1,217 @@
+# 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-unsafe
+
+from contextlib import contextmanager
+from typing import Generator
+from unittest import TestCase
+
+import monarch
+
+import pytest
+import torch
+import torch.utils._python_dispatch
+from monarch import fetch_shard, no_mesh, remote, Stream
+from monarch.common.device_mesh import DeviceMesh
+from monarch.rust_local_mesh import local_meshes, LoggingLocation, SocketType
+from torch.nn.attention import sdpa_kernel, SDPBackend
+from torch.nn.functional import scaled_dot_product_attention
+
+
+def simple_all_reduce(*args, **kwargs):
+    return torch.ones(args[0].shape)
+
+
+simple_all_reduce = remote(
+    "monarch.worker._testing_function.simple_all_reduce_local",
+    propagate=simple_all_reduce,
+)
+
+
+@contextmanager
+def local_mesh(
+    hosts: int = 1, gpu_per_host: int = 2, activate: bool = True
+) -> Generator[DeviceMesh, None, None]:
+    with monarch.rust_local_mesh.local_mesh(
+        hosts=hosts,
+        gpus_per_host=gpu_per_host,
+        socket_type=SocketType.UNIX,
+        logging_location=LoggingLocation.DEFAULT,
+    ) as dm:
+        try:
+            if activate:
+                with dm.activate():
+                    yield dm
+            else:
+                yield dm
+            dm.exit()
+        except Exception:
+            dm.client._shutdown = True
+            raise
+
+
+# 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.
+@pytest.mark.timeout(120)
+@pytest.mark.skipif(
+    torch.cuda.device_count() < 2,
+    reason="Not enough GPUs, this test requires at least 2 GPUs",
+)
+class TestRustBackend(TestCase):
+    def test_local_mesh_setup(self):
+        with local_mesh():
+            t = torch.zeros(3, 4)
+            t.add_(1)
+            fut = fetch_shard(t)
+
+            with no_mesh.activate():
+                local_t = fut.result()
+        assert torch.equal(local_t, torch.ones(3, 4))
+
+    def test_result_in_mesh(self):
+        with local_mesh():
+            t = torch.ones(3, 4)
+            t.add_(-1)
+            # Assert calling result() is fine within an active mesh.
+            local_t = fetch_shard(t).result()
+        assert torch.equal(local_t, torch.zeros(3, 4))
+
+    def test_errors(self):
+        t = torch.rand(3, 4)
+        with local_mesh(2, 2) as dm:
+            y = torch.rand(3, 4)
+            with pytest.raises(TypeError, match="LOCAL_TENSOR"):
+                t.add(y)
+            with pytest.raises(TypeError, match="WRONG_MESH"):
+                sub_mesh = dm(host=0)
+                with sub_mesh.activate():
+                    x = torch.rand(3, 4)
+                    x.add(y)
+            other = Stream("other")
+            t = torch.rand(10).cuda()
+            with pytest.raises(TypeError, match="WRONG_STREAM"):
+                with other.activate():
+                    t = t.reduce("host", "sum")
+
+    def test_multi_hosts(self):
+        with local_mesh(hosts=2, gpu_per_host=2):
+            t = torch.rand(3, 4).cuda()
+            local_t1 = fetch_shard(t, {"host": 1, "gpu": 0}).result()
+            local_t2 = fetch_shard(t, {"host": 1, "gpu": 0}).result()
+            local_t3 = fetch_shard(t, {"host": 0, "gpu": 1}).result()
+        assert torch.equal(local_t1, local_t2)
+        assert not torch.equal(local_t1, local_t3)
+
+    def test_fetch_preprocess(self):
+        with local_mesh():
+            assert (
+                "an argument processed"
+                == remote("monarch.worker._testing_function.do_some_processing")
+                .call_on_shard_and_fetch("an argument")
+                .result()
+            )
+
+    def test_brutal_shutdown(self):
+        with monarch.rust_local_mesh.local_mesh(
+            hosts=1, gpus_per_host=1, socket_type=SocketType.UNIX
+        ) as dm:
+            dm.exit()
+            dm.deactivate()
+
+    def test_results_filtering(self):
+        with local_mesh(gpu_per_host=1):
+            query = torch.rand(1, 1, 1, 1, dtype=torch.float16, device="cuda")
+            key = torch.rand(1, 1, 1, 1, dtype=torch.float16, device="cuda")
+            value = torch.rand(1, 1, 1, 1, dtype=torch.float16, device="cuda")
+            with sdpa_kernel(backends=[SDPBackend.FLASH_ATTENTION]):
+                # This function will send 9 results. Only 5 of them will be set.
+                t = scaled_dot_product_attention(query, key, value)
+                fut = fetch_shard(t)
+                local_tensor = fut.result()
+            assert len(local_tensor) == 1
+
+    def test_live_function(self):
+        with local_mesh():
+
+            @remote
+            def has_nan(t):
+                return torch.isnan(t).any().item()
+
+            t = torch.rand(3, 4)
+            res = has_nan.call_on_shard_and_fetch(
+                t, shard={"host": 0, "gpu": 0}
+            ).result()
+
+        self.assertFalse(res)
+
+    def test_multiple_global_meshes(self):
+        """
+        This test is to validate we can have a single client process
+        connecting to multiple global meshes. The global meshes are distinct
+        from each other to provide native failure domain isolation.
+        """
+        replicas = 4
+        with local_meshes(
+            meshes=replicas,
+            hosts_per_mesh=1,
+            gpus_per_host=1,
+            socket_type=SocketType.UNIX,
+            logging_location=LoggingLocation.DEFAULT,
+        ) as groups:
+            results = []
+            for i, group in enumerate(groups):
+                with group.activate():
+                    t = torch.ones(i + 1)
+                    results.append(fetch_shard(t).result())
+            for i in range(replicas):
+                assert torch.equal(results[i], torch.ones(i + 1))
+
+            for group in groups:
+                group.exit()
+                group.deactivate()
+
+    def test_get_world_status(self) -> None:
+        with local_mesh(gpu_per_host=2) as mesh:
+            mesh_info = mesh.get_info()
+
+            self.assertIsNotNone(mesh_info.mesh_labels)
+            self.assertEqual(len(mesh_info.devices_labels), 2)
+
+    def test_ivalue_problems(self) -> None:
+        with local_mesh(hosts=1, gpu_per_host=1):
+            from typing import cast
+
+            from monarch.common.messages import CallFunction, CommandGroup
+
+            a = cast(monarch.Tensor, torch.rand(3, 4))
+            result = monarch.Tensor(a._fake, a.mesh, a.stream)
+            msg = CallFunction(
+                0,
+                result,
+                (),
+                monarch.common.function.ResolvableFunctionFromPath(
+                    "torch.ops.aten.mul.Tensor"
+                ),
+                (2, a),
+                {},
+                a.stream._to_ref(a.mesh.client),
+                a.mesh,
+                [],
+            )
+            # Internally, this will call CallFunction(...).to_rust_message().
+            # The 2 arg will be converted to an IValue tensor via rust + C++.
+            # Then when the CommandGroup message gets converted to rust, it
+            # will attempt to clone the rust CallFunction message, which will
+            # attempt to clone the IValue tensor, which will cause a crash.
+            # Upon attempting to clone the IValue tensor, our custom __torch_dispatch__
+            # intercepts the following two calls:
+            #   aten._to_copy.default () (2,) {'dtype': torch.float64, 'device': device(type='cpu')}
+            #   aten.clone.default () (2,) {}
+
+            with torch.utils._python_dispatch._disable_current_modes():
+                CommandGroup([msg]).to_rust_message()

tests/test_signal_safe_block_on.py

@@ -0,0 +1,103 @@
+#!/usr/bin/env python3
+# 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
+
+"""
+Integration test for signal_safe_block_on.
+
+This test spawns a Python binary that calls a Rust function which sleeps indefinitely.
+The test then sends SIGINT to the process and confirms that it exits properly,
+verifying that signal_safe_block_on correctly handles signals.
+"""
+
+import importlib.resources
+import os
+import signal
+import subprocess
+import time
+import unittest
+
+import pytest
+
+
+# oss_skip: importlib not pulling resource correctly in git CI, needs to be revisited
+class TestSignalSafeBlockOn(unittest.TestCase):
+    # pyre-ignore[56]
+    @pytest.mark.oss_skip
+    def test_sigint_handling(self) -> None:
+        """
+        Test that a process using signal_safe_block_on can be interrupted with SIGINT.
+
+        This test:
+        1. Spawns a subprocess running sleep_binary.py
+        2. Waits for it to start
+        3. Sends SIGINT to the process
+        4. Verifies that the process exits within a reasonable timeout
+
+        To validate that it will behave in the same way as a ctl-c in the shell,
+        we launch the process in it's own process group and send the signal to the process
+        group instead of the process itself.
+        """
+        test_bin = importlib.resources.files("monarch.python.tests").joinpath(
+            "test_bin"
+        )
+        # Start the subprocess
+        process = subprocess.Popen(
+            [str(test_bin)],
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            text=True,
+            start_new_session=True,
+        )
+
+        gpig = os.getpgid(process.pid)
+
+        try:
+            # Wait for the process to start and print its startup message
+            start_time = time.time()
+            startup_timeout = 10  # seconds
+
+            while time.time() - start_time < startup_timeout:
+                if process.stdout and "Starting sleep_binary" in (
+                    process.stdout.readline() or ""
+                ):
+                    break
+                time.sleep(0.1)
+            else:
+                self.fail("Subprocess did not start properly within timeout")
+
+            # Give the process a moment to enter the sleep_indefinitely_for_unit_tests function
+            time.sleep(1)
+
+            # Send SIGINT to the process
+            os.killpg(gpig, signal.SIGINT)
+
+            # Wait for the process to exit with a timeout
+            exit_timeout = 5  # seconds
+            exit_time = time.time()
+
+            while time.time() - exit_time < exit_timeout:
+                if process.poll() is not None:
+                    # Process has exited
+                    break
+                time.sleep(0.1)
+            else:
+                self.fail("Process did not exit after receiving SIGINT")
+
+            # Check that the process exited with code 0 (clean exit)
+            self.assertEqual(process.returncode, 0, "Process did not exit cleanly")
+
+        finally:
+            # Clean up in case the test fails
+            if process.poll() is None:
+                process.kill()
+                process.wait()
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_sim_backend.py

@@ -0,0 +1,54 @@
+# 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-unsafe
+
+from contextlib import contextmanager
+from typing import Generator, Optional
+from unittest import TestCase
+
+import pytest
+
+import torch
+from monarch import fetch_shard
+from monarch.common.device_mesh import DeviceMesh
+from monarch.sim_mesh import sim_mesh
+
+
+@contextmanager
+def local_sim_mesh(
+    hosts: int = 1,
+    # TODO: support multiple gpus in a mesh.
+    gpu_per_host: int = 1,
+    activate: bool = True,
+    proxy_addr: Optional[str] = None,
+) -> Generator[DeviceMesh, None, None]:
+    dms = sim_mesh(
+        n_meshes=1, hosts=hosts, gpus_per_host=gpu_per_host, proxy_addr=proxy_addr
+    )
+    dm = dms[0]
+    try:
+        if activate:
+            with dm.activate():
+                yield dm
+        else:
+            yield dm
+        dm.exit()
+    except Exception:
+        dm.client._shutdown = True
+        raise
+
+
+# oss_skip: importlib not pulling resource correctly in git CI, needs to be revisited
+@pytest.mark.oss_skip
+class TestSimBackend(TestCase):
+    def test_local_mesh_setup(self):
+        with local_sim_mesh():
+            t = torch.zeros(3, 4)
+            t.add_(1)
+            local_t = fetch_shard(t).result()
+        # consider support specifying the return value in the mock worker.
+        assert local_t is not None

tests/test_tensor_engine.py

@@ -0,0 +1,52 @@
+# 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 monarch
+import pytest
+import torch
+from monarch.mesh_controller import spawn_tensor_engine
+from monarch.proc_mesh import proc_mesh
+
+
+two_gpu = pytest.mark.skipif(
+    torch.cuda.device_count() < 2,
+    reason="Not enough GPUs, this test requires at least 2 GPUs",
+)
+
+
+@two_gpu
+def test_tensor_engine() -> None:
+    pm = proc_mesh(gpus=2).get()
+
+    dm = spawn_tensor_engine(pm)
+    with dm.activate():
+        r = monarch.inspect(2 * torch.zeros(3, 4))
+
+    fm = dm.flatten("all")
+    with fm.activate():
+        f = monarch.inspect(2 * torch.zeros(3, 4), all=1)
+
+    assert torch.allclose(torch.zeros(3, 4), r)
+    assert torch.allclose(torch.zeros(3, 4), f)
+
+    dm.exit()
+
+
+@two_gpu
+def test_proc_mesh_tensor_engine() -> None:
+    pm = proc_mesh(gpus=2).get()
+    with pm.activate():
+        f = 10 * pm.rank_tensor("gpus").cuda()
+        a = monarch.inspect(f, hosts=0, gpus=0)
+        b = monarch.inspect(f, hosts=0, gpus=1)
+
+    one = pm.slice(gpus=1)
+    with one.activate():
+        sliced_b = monarch.slice_mesh(f, gpus=1).to_mesh(one)
+        c = monarch.inspect(sliced_b * 10)
+    assert a == 0
+    assert b == 10
+    assert c == 100

torchmonarch_nightly-2025.6.27.dist-info/METADATA

@@ -0,0 +1,94 @@
+Metadata-Version: 2.4
+Name: torchmonarch-nightly
+Version: 2025.6.27
+Summary: Monarch: Single controller library
+Author: Meta
+Author-email: oncall+monarch@xmail.facebook.com
+License: BSD-3-Clause
+Requires-Python: >= 3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch
+Requires-Dist: pyzmq
+Requires-Dist: requests
+Requires-Dist: numpy
+Requires-Dist: pyre-extensions
+Requires-Dist: cloudpickle
+Requires-Dist: torchx-nightly
+Dynamic: author
+Dynamic: author-email
+Dynamic: description
+Dynamic: description-content-type
+Dynamic: license
+Dynamic: license-file
+Dynamic: requires-dist
+Dynamic: requires-python
+Dynamic: summary
+
+# Monarch 🦋
+
+**Monarch** is a distributed execution engine for PyTorch. Our overall goal is
+to deliver the high-quality user experience that people get from single-GPU
+PyTorch, but at cluster scale.
+
+> ⚠️ **Early Development Warning** Monarch is currently in an experimental
+> stage. You should expect bugs, incomplete features, and APIs that may change
+> in future versions. The project welcomes bugfixes, but to make sure things are
+> well coordinated you should discuss any significant change before starting the
+> work. It's recommended that you signal your intention to contribute in the
+> issue tracker, either by filing a new issue or by claiming an existing one.
+
+Note: Monarch is currently only supported on Linux systems
+
+## Installation
+
+`pip install torchmonarch-nightly`
+
+or manually
+
+```sh
+
+# Create and activate the conda environment
+conda create -n monarchenv python=3.10 -y
+conda activate monarchenv
+
+# Install nightly rust toolchain
+curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh
+rustup toolchain install nightly
+rustup default nightly
+
+# Install non-python dependencies
+conda install libunwind -y
+
+# Install the correct cuda and cuda-toolkit versions for your machine
+sudo dnf install cuda-toolkit-12-0 cuda-12-0
+
+# Install clang-dev and nccl-dev
+sudo dnf install clang-devel libnccl-devel
+# Or, in some envrionments, the following may be necessary instead
+conda install -c conda-forge clangdev nccl
+conda update -n monarchenv --all -c conda-forge -y
+
+# Install build dependencies
+pip install -r build-requirements.txt
+# Install test dependencies
+pip install -r python/tests/requirements.txt
+
+# Build and install Monarch
+pip install --no-build-isolation .
+# or setup for development
+pip install --no-build-isolation -e .
+
+# Run unit tests. consider -s for more verbose output
+pytest python/tests/ -v -m "not oss_skip"
+```
+
+## Running examples
+
+Check out the `examples/` directory for demonstrations of how to use Monarch's APIs.
+
+We'll be adding more examples as we stabilize and polish functionality!
+
+## License
+
+Monarch is BSD-3 licensed, as found in the [LICENSE](LICENSE) file.