torchmonarch-nightly 2025.6.27__cp312-cp312-manylinux2014_x86_64.whl

monarch_supervisor/logging.py

@@ -0,0 +1,103 @@
+# 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
+import logging
+import os
+import socket
+import sys
+from pathlib import Path
+
+logger = logging.getLogger(__name__)
+
+
+def _handle_unhandled_exception(*args):
+    logger.error("Uncaught exception", exc_info=args)
+
+
+_glog_level_to_abbr = {
+    "DEBUG": "V",  # V is for VERBOSE in glog
+    "INFO": "I",
+    "WARNING": "W",
+    "ERROR": "E",
+    "CRITICAL": "C",
+}
+
+
+def fix_exception_lines(tb_lines):
+    formatted_lines = []
+    for line in tb_lines:
+        # Replace the standard file and line format with the custom format
+        if line.startswith("  File"):
+            # Extract the filename and line number
+            parts = line.split(",")
+            file_info = parts[0].strip()[6:-1]  # Remove '  File "' and '"'
+            line_info = parts[1].strip()[5:]  # Remove 'line '
+            new_line = f"  File {file_info}:{line_info}"
+            if len(parts) > 2:
+                new_line += ", " + ",".join(parts[2:]).strip()
+            formatted_lines.append(new_line)
+        else:
+            formatted_lines.append(line.strip())
+    return formatted_lines
+
+
+class _Formatter(logging.Formatter):
+    def __init__(self, suffix):
+        self.suffix = suffix
+
+    def format(self, record):
+        message = record.getMessage()
+        asctime = self.formatTime(record, "%m%d %H:%M:%S")
+
+        lines = message.strip().split("\n")
+        if record.exc_info:
+            exc_info = fix_exception_lines(
+                self.formatException(record.exc_info).split("\n")
+            )
+            lines.extend(exc_info)
+        if record.stack_info:
+            stack_info = self.formatStack(record.stack_info)
+            lines.extend(stack_info.strip().split("\n"))
+
+        shortlevel = _glog_level_to_abbr.get(record.levelname, record.levelname[0])
+
+        prefix = (
+            f"{shortlevel}{asctime}.{int(record.msecs*1000):06d} "
+            f"{record.filename}:"
+            f"{record.lineno}]{self.suffix}"
+        )
+        return "\n".join(f"{prefix} {line}" for line in lines)
+
+
+def initialize_logging(process_name=None):
+    log_folder = os.environ.get("TORCH_MONARCH_LOG_FOLDER")
+    log_level = os.environ.get("TORCH_MONARCH_LOG_LEVEL", "INFO")
+    suffix = "" if process_name is None else f" {process_name}:"
+    handler = None
+    if log_folder is not None:
+        log_folder_path = Path(log_folder)
+        log_folder_path.mkdir(parents=True, exist_ok=True)
+        safe_process_name = (
+            process_name.replace("/", "_") if process_name else "logfile.log"
+        )
+        log_file_name = f"{safe_process_name}.log"
+        log_file_path = log_folder_path / log_file_name
+        handler = logging.FileHandler(log_file_path)
+    else:
+        handler = logging.StreamHandler()
+    handler.setFormatter(_Formatter(suffix))
+    handler.setLevel(log_level)
+    logging.root.setLevel(log_level)
+    logging.root.addHandler(handler)
+    sys.excepthook = _handle_unhandled_exception
+
+
+def gethostname():
+    """Get the hostname of the machine."""
+    hostname = socket.gethostname()
+    hostname = hostname.replace(".facebook.com", "")
+    return hostname

monarch_supervisor/python_executable.py

@@ -0,0 +1,42 @@
+# 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 importlib.resources
+import os
+import sys
+
+try:
+    from __manifest__ import fbmake  # noqa
+
+    IN_PAR = True
+except ImportError:
+    IN_PAR = False
+
+PYTHON_EXECUTABLE: str
+if IN_PAR:
+    # The worker bootstrap binary will import this supervisor lib. When that
+    # happens don't try to search for the bootstrap binary again, just use the
+    # current executable.
+    import __main__ as main_module  # @manual
+
+    if hasattr(main_module, "__MONARCH_TENSOR_WORKER_ENV__"):
+        PYTHON_EXECUTABLE = os.environ["FB_XAR_INVOKED_NAME"]
+    else:
+        try:
+            with importlib.resources.path(
+                "monarch_tensor_worker_env", "worker_env"
+            ) as path:
+                if not path.exists():
+                    raise ImportError()
+                PYTHON_EXECUTABLE = str(path)
+        except ImportError:
+            raise ImportError(
+                "Monarch worker env not found, please define a custom 'monarch_tensor_worker_env' or "
+                "add '//monarch/python/monarch_supervisor/worker:default_worker_env' "
+                "to your binary dependencies in TARGETS"
+            )
+else:
+    PYTHON_EXECUTABLE = sys.executable

tests/dispatch_bench.py

@@ -0,0 +1,124 @@
+# 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
+import logging
+import sys
+
+import torch
+import torch.utils.benchmark as benchmark
+
+# this function helps get a local device mesh for testing
+from monarch._testing import mock_mesh
+from monarch.builtins.log import set_logging_level_remote
+
+from monarch.common._coalescing import coalescing
+from monarch.common.remote import remote
+from monarch.fetch import fetch_shard
+from monarch.python_local_mesh import python_local_mesh
+from monarch_supervisor.logging import initialize_logging
+from tests.dispatch_bench_helper import run_loop, run_loop_local
+
+NITER = 10000
+DEFAULT_TENSOR_SIZE = (100, 100)
+
+initialize_logging("dispatch_bench")
+
+
+# user-defined remote functions
+log = remote("monarch.worker._testing_function.log", propagate="inspect")
+
+
+def local_run():
+    run_loop_local(NITER, DEFAULT_TENSOR_SIZE)
+
+
+def dispatch_to_worker(device_mesh, n_iter, tensor_size):
+    with device_mesh.activate():
+        result = run_loop_local(n_iter, tensor_size)
+        local_result = fetch_shard(result, {"host": 0, "gpu": 0})
+    local_result = local_result.result()
+
+
+def dispatch_to_worker_remote_function(device_mesh, n_iter, tensor_size):
+    with device_mesh.activate():
+        result = run_loop(n_iter, tensor_size)
+        local_result = fetch_shard(result, {"host": 0, "gpu": 0})
+    local_result = local_result.result()
+
+
+def dispatch_to_worker_coalescing(device_mesh, n_iter, tensor_size):
+    with device_mesh.activate():
+        with coalescing():
+            result = run_loop_local(n_iter, tensor_size)
+        local_result = fetch_shard(result, {"host": 0, "gpu": 0})
+    local_result = local_result.result()
+
+
+def main():
+    mocked = False
+    torch.set_default_device("cuda")
+    if mocked:
+        device_mesh = mock_mesh(hosts=1, gpus=1)
+    else:
+        device_mesh = python_local_mesh(hosts=1, gpus=1)
+
+    with device_mesh.activate():
+        torch.set_default_device("cuda")
+        set_logging_level_remote(logging.WARNING)
+
+    # bench 1: local compute only
+    t0 = benchmark.Timer(
+        stmt="run_loop_local(niter, tensor_size)",
+        setup="from __main__ import run_loop_local",
+        globals={"niter": NITER, "tensor_size": DEFAULT_TENSOR_SIZE},
+    )
+    local_only_results = t0.blocked_autorange(min_run_time=10)
+    print(local_only_results)
+
+    t1 = benchmark.Timer(
+        stmt="dispatch_to_worker(device_mesh, niter, tensor_size)",
+        setup="from __main__ import dispatch_to_worker",
+        globals={
+            "device_mesh": device_mesh,
+            "niter": NITER,
+            "tensor_size": DEFAULT_TENSOR_SIZE,
+        },
+    )
+    dispatch_to_worker_results = t1.blocked_autorange(min_run_time=10)
+    print(dispatch_to_worker_results)
+
+    t2 = benchmark.Timer(
+        stmt="dispatch_to_worker_remote_function(device_mesh, niter, tensor_size)",
+        setup="from __main__ import dispatch_to_worker_remote_function",
+        globals={
+            "device_mesh": device_mesh,
+            "niter": NITER,
+            "tensor_size": DEFAULT_TENSOR_SIZE,
+        },
+    )
+    dispatch_to_worker_remote_function_results = t2.blocked_autorange(min_run_time=10)
+    print(dispatch_to_worker_remote_function_results)
+
+    t3 = benchmark.Timer(
+        stmt="dispatch_to_worker_coalescing(device_mesh, niter, tensor_size)",
+        setup="from __main__ import dispatch_to_worker_coalescing",
+        globals={
+            "device_mesh": device_mesh,
+            "niter": NITER,
+            "tensor_size": DEFAULT_TENSOR_SIZE,
+        },
+    )
+    dispatch_to_worker_coalescing_results = t3.blocked_autorange(min_run_time=10)
+    print(dispatch_to_worker_coalescing_results)
+
+    device_mesh.exit()
+
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

tests/dispatch_bench_helper.py

@@ -0,0 +1,25 @@
+# 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
+import torch
+
+from monarch.common.remote import remote
+
+
+def run_loop_local(n_iters, tensor_shape=(2, 2)):
+    local = torch.zeros(*tensor_shape)
+    ones = torch.ones(*tensor_shape)
+    for _ in range(n_iters):
+        local = ones + local
+    return local
+
+
+def _run_loop(*args, **kwargs):
+    return torch.ones(args[1])
+
+
+run_loop = remote("tests.dispatch_bench_helper.run_loop_local", propagate=_run_loop)

tests/error_test_binary.py

@@ -0,0 +1,180 @@
+# 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 asyncio
+import ctypes
+import sys
+
+import click
+
+from monarch._rust_bindings.monarch_extension.panic import panicking_function
+
+from monarch.actor_mesh import Actor, endpoint, send
+from monarch.proc_mesh import proc_mesh
+
+
+class ErrorActor(Actor):
+    """An actor that has endpoints cause segfaults."""
+
+    @endpoint
+    async def cause_segfault(self) -> None:
+        """Endpoint that causes a segmentation fault."""
+        # Create a C function pointer to an invalid memory address
+        # This will reliably cause a segmentation fault when called
+        function_type = ctypes.CFUNCTYPE(None)
+        # Use a non-zero but invalid address to avoid ctypes null pointer checks
+        invalid_address = 0xDEADBEEF
+        invalid_function = function_type(invalid_address)
+        # Calling this function will cause a segfault
+        invalid_function()
+
+    @endpoint
+    async def cause_panic(self) -> None:
+        """Endpoint that calls a Rust function that panics."""
+        panicking_function()
+
+    @endpoint
+    async def await_then_error(self) -> None:
+        await asyncio.sleep(0.1)
+        await asyncio.sleep(0.1)
+        raise RuntimeError("oh noez")
+
+
+class ErrorActorSync(Actor):
+    """An actor that has endpoints cause segfaults."""
+
+    @endpoint  # pyre-ignore
+    def cause_segfault(self) -> None:
+        """Endpoint that causes a segmentation fault."""
+        # Create a C function pointer to an invalid memory address
+        # This will reliably cause a segmentation fault when called
+        function_type = ctypes.CFUNCTYPE(None)
+        # Use a non-zero but invalid address to avoid ctypes null pointer checks
+        invalid_address = 0xDEADBEEF
+        invalid_function = function_type(invalid_address)
+        # Calling this function will cause a segfault
+        invalid_function()
+
+    @endpoint  # pyre-ignore
+    def cause_panic(self) -> None:
+        """Endpoint that calls a Rust function that panics."""
+        panicking_function()
+
+
+def _run_error_test_sync(num_procs, sync_endpoint, endpoint_name):
+    proc = proc_mesh(gpus=num_procs).get()
+    if sync_endpoint:
+        actor_class = ErrorActorSync
+    else:
+        actor_class = ErrorActor
+    error_actor = proc.spawn("error_actor", actor_class).get()
+
+    # This output is checked in the test to make sure that the process actually got here
+    print("I actually ran")
+    sys.stdout.flush()
+
+    if endpoint_name == "cause_segfault":
+        endpoint = error_actor.cause_segfault
+    elif endpoint_name == "cause_panic":
+        endpoint = error_actor.cause_panic
+    else:
+        raise ValueError(f"Unknown endpoint name: {endpoint_name}")
+
+    # Exercise both call() and call_one() in our tests, to check that error
+    # aggregation behavior is consistent.
+    if num_procs == 1:
+        endpoint.call_one().get()
+    else:
+        endpoint.call().get()
+
+
+def _run_error_test(num_procs, sync_endpoint, endpoint_name):
+    import asyncio
+
+    if sync_endpoint:
+        actor_class = ErrorActorSync
+    else:
+        actor_class = ErrorActor
+
+    async def run_test():
+        proc = await proc_mesh(gpus=num_procs)
+        error_actor = await proc.spawn("error_actor", actor_class)
+
+        # This output is checked in the test to make sure that the process actually got here
+        print("I actually ran")
+        sys.stdout.flush()
+
+        if endpoint_name == "cause_segfault":
+            endpoint = error_actor.cause_segfault
+        elif endpoint_name == "cause_panic":
+            endpoint = error_actor.cause_panic
+        else:
+            raise ValueError(f"Unknown endpoint name: {endpoint_name}")
+
+        # Exercise both call() and call_one() in our tests, to check that error
+        # aggregation behavior is consistent.
+        if num_procs == 1:
+            await endpoint.call_one()
+        else:
+            await endpoint.call()
+
+    asyncio.run(run_test())
+
+
+@click.group()
+def main():
+    pass
+
+
+@main.command("error-endpoint")
+@click.option("--num-procs", type=int, required=True)
+@click.option("--sync-test-impl", type=bool, required=True)
+@click.option("--sync-endpoint", type=bool, required=True)
+@click.option("--endpoint-name", type=str, required=True)
+def error_endpoint(num_procs, sync_test_impl, sync_endpoint, endpoint_name):
+    print(
+        f"Running segfault test: {num_procs=} {sync_test_impl=} {sync_endpoint=}, {endpoint_name=}"
+    )
+
+    if sync_test_impl:
+        _run_error_test_sync(num_procs, sync_endpoint, endpoint_name)
+    else:
+        _run_error_test(num_procs, sync_endpoint, endpoint_name)
+
+
+@main.command("error-bootstrap")
+def error_bootstrap():
+    print("I actually ran")
+    sys.stdout.flush()
+
+    proc_mesh(gpus=4, env={"MONARCH_ERROR_DURING_BOOTSTRAP_FOR_TESTING": "1"}).get()
+
+
+async def _error_unmonitored():
+    print("I actually ran")
+    sys.stdout.flush()
+
+    proc = await proc_mesh(gpus=1)
+    actor = await proc.spawn("error_actor", ErrorActor)
+
+    # fire and forget
+    send(actor.await_then_error, (), {}, None, "all")
+
+    # Wait. Eventually a supervision event will get propagated and the process
+    # will exit.
+    #
+    # If an event is not delivered, the test will time out before this sleep
+    # finishes.
+    await asyncio.sleep(300)
+
+
+@main.command("error-unmonitored")
+def error_unmonitored():
+    asyncio.run(_error_unmonitored())
+
+
+if __name__ == "__main__":
+    main()

tests/simulator/test_profiling.py

@@ -0,0 +1,136 @@
+# 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
+import unittest
+
+import pytest
+
+import torch
+
+from monarch.common import messages
+from monarch.simulator.profiling import RuntimeEstimator, RuntimeProfiler, TimingType
+
+
+# pyre-ignore-all-errors[6]
+# pyre-ignore-all-errors[16]
+class TestRuntimeEstimator(unittest.TestCase):
+    def test_user_manual_setting(self):
+        runtime = RuntimeEstimator()
+
+        input_tensor = torch.rand(10, 10)
+        input_tensor.ref = 1
+        input_tensor._fake = None
+        output_tensor = torch.rand(10, 10)
+        output_tensor.ref = 2
+        output_tensor._fake = None
+
+        send_tensor = messages.SendTensor(
+            result=output_tensor,
+            from_ranks=[1],
+            to_ranks=[2],
+            tensor=input_tensor,
+            factory=None,
+            from_stream=None,
+            to_stream=None,
+        )
+        reduce = messages.Reduce(
+            result=output_tensor,
+            local_tensor=input_tensor,
+            factory=None,
+            source_mesh=None,
+            stream=None,
+            dims=None,
+            reduction=None,
+            scatter=False,
+            inplace=False,
+            out=None,
+        )
+        call_function = messages.CallFunction(
+            ident=1,
+            result=None,
+            mutates=None,
+            function=None,
+            args=None,
+            kwargs=None,
+            stream=None,
+            device_mesh=None,
+            remote_process_groups=None,
+        )
+
+        self.assertEqual(runtime.get_runtime(send_tensor), 100_000)
+        self.assertEqual(runtime.get_runtime(reduce), 100_000)
+        self.assertEqual(runtime.get_runtime(call_function), 10_000)
+        self.assertEqual(runtime.get_runtime("kernel_launch"), 500)
+        self.assertEqual(runtime.get_runtime("wait_event"), 500)
+
+        runtime.set_custom_timing(
+            {
+                TimingType.SEND_TENSOR: 1_000,
+                TimingType.REDUCE: 2_000,
+                TimingType.CALL_FUNCTION: 3_000,
+                TimingType.KERNEL_LAUNCH: 4_000,
+                TimingType.WAIT_EVENT: 5_000,
+            }
+        )
+        self.assertEqual(runtime.get_runtime(send_tensor), 1_000)
+        self.assertEqual(runtime.get_runtime(reduce), 2_000)
+        self.assertEqual(runtime.get_runtime(call_function), 3_000)
+        self.assertEqual(runtime.get_runtime("kernel_launch"), 4_000)
+        self.assertEqual(runtime.get_runtime("wait_event"), 5_000)
+
+        runtime.set_custom_timing(
+            {
+                TimingType.SEND_TENSOR: lambda msg: 4_000,
+                TimingType.REDUCE: lambda msg: 5_000,
+                TimingType.CALL_FUNCTION: lambda msg: 6_000,
+                TimingType.KERNEL_LAUNCH: lambda: 8_000,
+                TimingType.WAIT_EVENT: lambda: 9_000,
+            }
+        )
+        self.assertEqual(runtime.get_runtime(send_tensor), 4_000)
+        self.assertEqual(runtime.get_runtime(reduce), 5_000)
+        self.assertEqual(runtime.get_runtime(call_function), 6_000)
+        self.assertEqual(runtime.get_runtime("kernel_launch"), 8_000)
+        self.assertEqual(runtime.get_runtime("wait_event"), 9_000)
+
+    @pytest.mark.oss_skip
+    def test_runtime_profiler(self) -> None:
+        m1 = torch.rand(1000, 2000).cuda()
+        m2 = torch.rand(2000, 4000).cuda()
+        m1.ref = 1
+        m2.ref = 2
+        msg = messages.CallFunction(
+            ident=1,
+            result=None,
+            mutates=None,
+            function=torch.ops.aten.mm.default,
+            args=(m1, m2),
+            kwargs=None,
+            stream=None,
+            device_mesh=None,
+            remote_process_groups=None,
+        )
+        profiler = RuntimeProfiler()
+
+        ret = profiler.profile_cmd(msg, ranks=[0])[0]
+        self.assertEqual(ret[0].factory.size, (1000, 4000))
+        # Should be at least 0.1 ms
+        self.assertTrue(ret[1] > 100)
+        # Should be at most 100 ms
+        self.assertTrue(ret[1] < 100_000)
+
+        # Change the cached profiling result to verify if cached mechanism works
+        key = next(iter(profiler.cached.keys()))
+        profiler.cached[key][0] = (profiler.cached[key][0][0], 987_654_321)
+
+        ret = profiler.profile_cmd(msg, ranks=[0])[0]
+        self.assertEqual(ret[0].factory.size, (1000, 4000))
+        self.assertEqual(ret[1], 987_654_321)
+
+
+if __name__ == "__main__":
+    unittest.main()