torchmonarch-nightly 2025.6.27__cp312-cp312-manylinux2014_x86_64.whl

monarch_supervisor/_testing.py

@@ -0,0 +1,44 @@
+# 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 typing import NamedTuple
+
+from monarch_supervisor import get_message_queue
+
+
+class Reply(NamedTuple):
+    a: int
+    b: int
+    x: int
+
+
+def reply_hello(a, b, x):
+    q = get_message_queue()
+    q.send(Reply(a, b, x))
+
+
+def echo():
+    q = get_message_queue()
+    i = 0
+    while True:
+        sender, m = q.recv()
+        if m == "exit":
+            break
+        assert m == i
+        q.send(m)
+        i += 1
+
+
+class Mapper:
+    def map(self, items):
+        return sum(x * 2 for x in items)
+
+    def reduce(self, items):
+        return sum(items)
+
+    def finish(self, result):
+        return result

monarch_supervisor/function_call.py

@@ -0,0 +1,30 @@
+# 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 importlib.util
+import sys
+
+from monarch_supervisor import _FunctionCall, get_message_queue
+
+if __name__ == "__main__":
+    q = get_message_queue()
+    _, call = q.recv()
+    assert isinstance(call, _FunctionCall)
+    filename, *rest = call.target.split(":", 1)
+    if not rest:
+        modulename, funcname = filename.rsplit(".", 1)
+        module = importlib.import_module(modulename)
+    else:
+        spec = importlib.util.spec_from_file_location("__entry__", filename)
+        assert spec is not None and spec.loader is not None
+        module = importlib.util.module_from_spec(spec)
+        # pyre-ignore[16]
+        spec.loader.exec_module(module)
+        sys.modules["__entry__"] = module
+        funcname = rest[0]
+    func = getattr(module, funcname)
+    func(*call.args, **call.kwargs)

monarch_supervisor/host.py

@@ -0,0 +1,386 @@
+# 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 ctypes
+import io
+import logging
+import os
+import signal
+import socket
+import subprocess
+import sys
+import time
+import traceback
+import uuid
+from contextlib import nullcontext
+from pathlib import Path
+from random import random
+from string import Template
+from typing import Any, Callable, Dict, List, Mapping, Optional
+
+import zmq
+from monarch_supervisor import (
+    _FunctionCall,
+    HEARTBEAT_INTERVAL,
+    HEARTBEAT_LIVENESS,
+    pickle_dumps,
+    pickle_loads,
+    ProcessFailedToStart,
+)
+from monarch_supervisor.logging import gethostname, initialize_logging
+from monarch_supervisor.python_executable import PYTHON_EXECUTABLE
+
+logger: logging.Logger = logging.getLogger(__name__)
+ABORT_INTERVAL = 5
+LOG_PSTREE_INTERVAL: int = 60 * 10
+__NR_pidfd_open = 434
+libc = ctypes.CDLL(None)
+
+
+# older libc do not have this syscall
+def pidfd_open(pid: int) -> int:
+    return libc.syscall(__NR_pidfd_open, pid, 0)
+
+
+# objects in this file represent Host/Process
+# on the host machine itself.
+
+# main package has Host/Process used by
+# the supervisor.
+
+
+class Process:
+    def __init__(
+        self,
+        name: str,
+        logfilename: Optional[str],
+        proc_comm: zmq.Socket,
+        proc_id: int,
+        rank: int,
+        processes_per_host: int,
+        world_size: int,
+        popen: Mapping[str, Any],
+        proc_addr: str,
+        start_new_session: bool,
+    ) -> None:
+        self.proc_id = proc_id
+        self.proc_comm = proc_comm
+        self.deferred_sends: Optional[List[bytes]] = []
+        local_config = {
+            "RANK": str(rank),
+            "WORLD_SIZE": str(world_size),
+            "LOCAL_RANK": str(rank % processes_per_host),
+            "LOCAL_WORLD_SIZE": str(processes_per_host),
+            "SUPERVISOR_PIPE": proc_addr,
+            "SUPERVISOR_IDENT": str(proc_id),
+        }
+
+        environ = dict(os.environ)
+        if popen["env"] is not None:
+            environ.update(
+                {
+                    k: Template(v).safe_substitute(local_config)
+                    for k, v in popen["env"].items()
+                }
+            )
+        args = popen["args"]
+        if isinstance(args, _FunctionCall):
+            self._send(pickle_dumps(args))
+            args = [PYTHON_EXECUTABLE, "-m", "monarch_supervisor.function_call"]
+
+        environ.update(local_config)
+        popen = {**popen, "env": environ, "args": args}
+        try:
+            if logfilename is None:
+                logcontext = nullcontext()
+            else:
+                try:
+                    logcontext = open(logfilename, "a")
+                except FileNotFoundError:
+                    Path(logfilename).parent.mkdir(exist_ok=True, parents=True)
+                    logcontext = open(logfilename, "a")
+            with logcontext as logfile:
+                self.subprocess: subprocess.Popen[str] = subprocess.Popen(
+                    **popen,
+                    start_new_session=start_new_session,
+                    stdout=logfile,
+                    stderr=logfile,
+                )
+        except Exception:
+            s = io.StringIO()
+            traceback.print_exc(file=s)
+            logger.warning(f"Process failed to start: {s.getvalue()}\n")
+            raise ProcessFailedToStart(s.getvalue())
+        self.proc_id_bytes: bytes = proc_id.to_bytes(8, byteorder="little")
+
+    def _send(self, msg: bytes) -> None:
+        if self.deferred_sends is not None:
+            self.deferred_sends.append(msg)
+        else:
+            self.proc_comm.send_multipart([self.proc_id_bytes, msg])
+
+    def _notify_connected(self) -> None:
+        deferred_sends = self.deferred_sends
+        if deferred_sends is not None:
+            for msg in deferred_sends:
+                self.proc_comm.send_multipart([self.proc_id_bytes, msg])
+            self.deferred_sends = None
+
+
+class Host:
+    """
+    Represents a host (Host Manager) that can be supervised.
+    Starts an event loop listening for commands from the supervisor, including launching/killing processes.
+    """
+
+    def __init__(self, supervisor_port: str, start_new_session: bool = True) -> None:
+        self.context: zmq.Context = zmq.Context(1)
+        self.supervisor_comm: zmq.Socket = self._socket(zmq.DEALER)
+        self.supervisor_comm.setsockopt(zmq.IPV6, True)
+        logger.info("Connecting to %s", supervisor_port)
+        self.supervisor_comm.connect(supervisor_port)
+
+        # tell the supervisor we exist, and provide
+        # hostname for debugging.
+        self.supervisor_comm.send(
+            pickle_dumps(("_hostname", None, socket.gethostname()))
+        )
+
+        self.poller = zmq.Poller()
+        self.poller.register(self.supervisor_comm, zmq.POLLIN)
+
+        # optional way to send messages to processes.
+        # all processes on this host will use the same
+        # socket.
+        self.proc_comm: zmq.Socket = self._socket(zmq.ROUTER)
+
+        self.proc_addr = f"ipc:///tmp/proc-{uuid.uuid4()}"
+        self.proc_comm.bind(self.proc_addr)
+        self.poller.register(self.proc_comm, zmq.POLLIN)
+
+        self.process_table: Dict[bytes, Process] = {}
+        self.fd_to_on_exit: Dict[int, Callable[[], None]] = {}
+        self._launches = 0
+        self.has_shutdown = False
+        self.exits: List[bytes] = []
+        self.start_new_session = start_new_session
+
+    def _socket(self, kind: int) -> zmq.Socket:
+        socket = self.context.socket(kind)
+        socket.setsockopt(zmq.SNDHWM, 0)
+        socket.setsockopt(zmq.RCVHWM, 0)
+        return socket
+
+    def heartbeat(self) -> None:
+        self.supervisor_comm.send(b"")
+
+    # TODO: validate these are valid messages to send
+
+    def launch(
+        self,
+        proc_id: int,
+        rank: int,
+        processes_per_rank: int,
+        world_size: int,
+        popen: Mapping[str, object],
+        name: str,
+        simulate: bool,
+        log_file: Optional[str],
+    ) -> None:
+        self._launches += 1
+        if simulate:
+            self.supervisor_comm.send(pickle_dumps(("_started", proc_id, 2)))
+            self.supervisor_comm.send(pickle_dumps(("_exited", proc_id, 0)))
+            return
+        try:
+            logger.info(f"starting new process proc_id: {proc_id}")
+            process = Process(
+                name,
+                log_file,
+                self.proc_comm,
+                proc_id,
+                rank,
+                processes_per_rank,
+                world_size,
+                popen,
+                self.proc_addr,
+                self.start_new_session,
+            )
+            self.process_table[process.proc_id_bytes] = process
+            self.on_subprocess_exit(
+                process.subprocess, lambda: self.process_exit(process)
+            )
+            reply = process.subprocess.pid
+        except ProcessFailedToStart as e:
+            reply = str(e)
+        self.supervisor_comm.send(pickle_dumps(("_started", proc_id, reply)))
+
+    def process_exit(self, process: Process) -> None:
+        self.process_table.pop(process.proc_id_bytes)
+        # we do not allow descendents to outlive the parent
+        # if any remain this kill will clean them up
+        self.kill(process.subprocess.pid, signal.SIGKILL)
+        returncode = process.subprocess.wait()
+        if not self.has_shutdown:
+            self.exits.append(pickle_dumps(("_exited", process.proc_id, returncode)))
+
+    def kill(self, pid: int, sig: int) -> None:
+        if self.start_new_session:
+            os.killpg(pid, sig)
+        else:
+            os.kill(pid, sig)
+
+    def on_subprocess_exit(
+        self, subprocess: subprocess.Popen, on_exit: Callable[[], Any]
+    ) -> None:
+        fd: int = pidfd_open(subprocess.pid)
+        self.fd_to_on_exit[fd] = on_exit
+        self.poller.register(fd, zmq.POLLIN)
+
+    def send(self, _proc_id: int, msg: bytes) -> None:
+        proc_id = _proc_id.to_bytes(8, byteorder="little")
+        if proc_id in self.process_table:
+            process = self.process_table[proc_id]
+            process._send(msg)
+
+    def signal(self, _proc_id: int, sig: int, group: bool) -> None:
+        proc_id = _proc_id.to_bytes(8, byteorder="little")
+        if proc_id in self.process_table:
+            process = self.process_table[proc_id]
+            if group and self.start_new_session:
+                os.killpg(process.subprocess.pid, sig)
+            else:
+                process.subprocess.send_signal(sig)
+
+    def _fd_exit(self, fd: int) -> None:
+        on_exit = self.fd_to_on_exit.pop(fd)
+        self.poller.unregister(fd)
+        os.close(fd)
+        on_exit()
+
+    def shutdown(self) -> None:
+        if self.has_shutdown:
+            return
+        self.has_shutdown = True
+        for proc in self.process_table.values():
+            self.kill(proc.subprocess.pid, signal.SIGTERM)
+        expiry = time.time() + ABORT_INTERVAL
+        ttl = ABORT_INTERVAL
+        while ttl > 0 and self.process_table:
+            for s, _ in self.poller.poll(timeout=int(1000 * ttl)):
+                if isinstance(s, int):
+                    self._fd_exit(s)
+            ttl = time.time() - expiry
+        if self.process_table:
+            for proc in self.process_table.values():
+                self.kill(proc.subprocess.pid, signal.SIGKILL)
+
+        self.proc_comm.close(linger=0)
+        self.supervisor_comm.close(linger=0)
+        self.context.term()
+
+    def abort(self, with_error: Optional[str] = None) -> None:
+        self.shutdown()
+        if with_error:
+            logger.error("exiting with error: %s", with_error)
+            raise ConnectionAbortedError(with_error)
+        else:
+            logger.warning("exiting cleanly.")
+            sys.exit(0)
+
+    def run_event_loop_forever(self) -> None:
+        log_pstree_info_at = time.time() + LOG_PSTREE_INTERVAL
+        supervisor_expiry = None
+        heartbeat_at = 0
+        while True:
+            timeout = (
+                -1
+                if supervisor_expiry is None
+                else int(max(1000 * (heartbeat_at - time.time()) + 1, 0))
+            )
+            proc_comm_processed = False
+            for s, _ in self.poller.poll(timeout=timeout):
+                if isinstance(s, int):
+                    # we register a file descriptor to the poller, which is a raw
+                    # int file description that becomes ready when the a subprocess exits
+                    # see pidfd_open.
+                    self._fd_exit(s)
+                elif s is self.supervisor_comm:
+                    if supervisor_expiry is None:
+                        logging.info("Connected to supervisor")
+                        # first heartbeat is set to between 0 to HEARTBEAT_INTERVAL
+                        # to spread out the heartbeats from hosts that all start
+                        # at the same time.
+                        heartbeat_at = time.time() + HEARTBEAT_INTERVAL * random()
+
+                    supervisor_expiry = (
+                        time.time() + HEARTBEAT_INTERVAL * HEARTBEAT_LIVENESS
+                    )
+                    # pyre-ignore[29]
+                    msg = self.supervisor_comm.recv()
+                    if msg:
+                        cmd, *args = pickle_loads(msg)
+                        getattr(self, cmd)(*args)
+                elif s is self.proc_comm:
+                    proc_comm_processed = True
+                    proc_id_bytes, msg = self.proc_comm.recv_multipart()
+                    process = self.process_table.get(proc_id_bytes)
+                    # it is possible for the process to have already exited before
+                    # we get its messages, so process_table will be empty
+                    if process is not None:
+                        process._notify_connected()
+                    if len(msg):
+                        proc_id = int.from_bytes(proc_id_bytes, byteorder="little")
+                        self.supervisor_comm.send(
+                            pickle_dumps(("_response", proc_id, msg))
+                        )
+            if not proc_comm_processed and self.exits:
+                for exit in self.exits:
+                    self.supervisor_comm.send(exit)
+                self.exits.clear()
+
+            if supervisor_expiry is not None:
+                t = time.time()
+                if t > heartbeat_at:
+                    heartbeat_at = t + HEARTBEAT_INTERVAL
+                    self.heartbeat()
+                if t > supervisor_expiry:
+                    self.abort(
+                        f"No messages from supervisor for {HEARTBEAT_INTERVAL*HEARTBEAT_LIVENESS} seconds, aborting."
+                    )
+                if t > log_pstree_info_at:
+                    log_pstree = subprocess.Popen(
+                        [
+                            os.getenv("FB_XAR_INVOKED_NAME", default=sys.executable),
+                            "-m",
+                            "monarch_supervisor.log_pstree",
+                            str(os.getpid()),
+                        ]
+                    )
+                    self.on_subprocess_exit(log_pstree, log_pstree.wait)
+                    log_pstree_info_at = t + LOG_PSTREE_INTERVAL
+
+
+def main(addr: str) -> None:
+    manager: Host = Host(addr)
+
+    def handler(signal: int, frame: object) -> None:
+        manager.shutdown()
+        sys.exit(1)
+
+    signal.signal(signal.SIGINT, handler)
+    signal.signal(signal.SIGTERM, handler)
+    try:
+        manager.run_event_loop_forever()
+    finally:
+        manager.shutdown()
+
+
+if __name__ == "__main__":
+    (addr,) = sys.argv[1:]
+    initialize_logging(f"{gethostname()} pid {os.getpid()} host-manager")
+    main(addr)

monarch_supervisor/launchers.py

@@ -0,0 +1,145 @@
+# 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 io
+import json
+import logging
+import os
+import signal
+import socket
+import subprocess
+import sys
+import time
+import traceback
+from typing import Callable, Optional, Sequence, Tuple
+
+from . import Context, Host
+
+from .host import main
+from .logging import gethostname, initialize_logging
+from .python_executable import PYTHON_EXECUTABLE
+
+# Default port leveraging one from the reserved range for torchelastic
+PORT: str = os.environ.get("SUPERVISOR_PORT", "29401")
+
+logger: logging.Logger = logging.getLogger(__name__)
+
+
+NON_RETRYABLE_FAILURE: int = 100
+JOB_RESTART_SCOPE_ESCALATION: int = 101
+TW_USER_METADATA_HOSTNAMES_LIST_KEY: str = "TW_USER_METADATA_HOSTNAMES_LIST_KEY"
+TW_USER_METADATA_FILE_PATH: str = "TW_USER_METADATA_FILE_PATH"
+
+
+def _write_reply_file(msg: str, reply_file: Optional[str] = None) -> None:
+    if reply_file is None:
+        reply_file = os.environ["MAST_HPC_TASK_FAILURE_REPLY_FILE"]
+    job_attempt = int(os.environ["MAST_HPC_JOB_ATTEMPT_INDEX"])
+    logger.info(
+        f"Supervisor writing a reply file with JOB_RESTART_SCOPE_ESCALATION to {reply_file} (attempt {job_attempt})."
+    )
+    with open(reply_file, "w") as f:
+        timestamp_ns = time.time_ns()
+        error_data = {
+            "message": msg,
+            "errorCode": JOB_RESTART_SCOPE_ESCALATION,
+            "timestamp": int(timestamp_ns // 1e9),
+            "timestamp_us": int(timestamp_ns // 1e3),
+        }
+        json.dump(error_data, f)
+
+
+def mast(supervise: Callable[[Context, Sequence[Host]], None]) -> None:
+    """
+    This function is the entrypoint for starting the supervisor when
+    running on MAST. Each host should call `mast(supervise)` where
+    `supervise` is the supervisor policy function for the job.
+    Supervisor will be called only on the supervisor machine with
+    `supervisor(n_hosts_in_task, port)` where `n_hosts_in_task` is
+    the number of hosts reserved in the task group, and `port` is the
+    port that supervisor should listen on.
+
+    The supervise function can then create a supervisor Context object,
+    request up to n_hosts_in_tasks hosts, and then
+    """
+
+    hostnames = get_hostnames()
+    N = len(hostnames)
+    my_host_name = (os.environ.get("HOSTNAME") or socket.getfqdn()).removesuffix(
+        ".facebook.com"
+    )
+    # Get first host in the task group
+    is_supervisor = my_host_name == hostnames[0]
+    initialize_logging(
+        "supervisor" if is_supervisor else f"{gethostname()} host-manager"
+    )
+
+    supervisor_addr = f"tcp://{socket.getfqdn(hostnames[0])}:{PORT}"
+    logger.info(
+        "hostname %s, supervisor host is %s, supervisor=%s",
+        my_host_name,
+        hostnames[0],
+        is_supervisor,
+    )
+
+    if is_supervisor:
+        _write_reply_file(
+            "Supervisor deadman's switch. "
+            "This reply file is written when the supervisor starts and deleted right before a successful exit. "
+            "It is used to cause the whole job to restart if for some reason the "
+            "supervisor host is unscheduled without it throwing an exception itself."
+        )
+        # local host manager on supervisor machine
+        host_process = subprocess.Popen(
+            [PYTHON_EXECUTABLE, "-m", "monarch_supervisor.host", supervisor_addr]
+        )
+        try:
+            ctx = Context(port=int(PORT))
+            hosts: Tuple[Host, ...] = ctx.request_hosts(n=N)
+            supervise(ctx, hosts)
+            ctx.shutdown()
+            logger.info("Supervisor shutdown complete.")
+        except BaseException:
+            ty, e, st = sys.exc_info()
+            s = io.StringIO()
+            traceback.print_tb(st, file=s)
+            _write_reply_file(
+                f"{ty.__name__ if ty is not None else 'None'}: {str(e)}\n{s.getvalue()}"
+            )
+            host_process.send_signal(signal.SIGINT)
+            raise
+        return_code = host_process.wait(timeout=10)
+        if return_code != 0:
+            # Host manager may have been instructed to write a reply file, so
+            # we do not write a reply file here which would clobber it.
+            logger.warning(
+                f"Host manager on supervisor returned non-zero code: {return_code}."
+            )
+            sys.exit(return_code)
+        else:
+            # successful exit, so we remove the deadman's switch reply file we wrote earlier.
+            reply_file = os.environ["MAST_HPC_TASK_FAILURE_REPLY_FILE"]
+            os.unlink(reply_file)
+    else:
+        # host manager on non-supervisor machine
+        main(supervisor_addr)
+
+
+def get_hostnames() -> Sequence[str]:
+    """
+    Get the list of hostnames for the current task group.
+    """
+    tw_metatdata_file = os.environ.get(TW_USER_METADATA_FILE_PATH)
+    hostnames_key = os.environ.get(TW_USER_METADATA_HOSTNAMES_LIST_KEY)
+    if tw_metatdata_file and hostnames_key:
+        with open(tw_metatdata_file, "r") as f:
+            data = json.load(f)
+            hostnames_str = data["userAttributes"][hostnames_key]
+            return hostnames_str.split(",")
+
+    return os.environ["MAST_HPC_TASK_GROUP_HOSTNAMES"].split(",")

monarch_supervisor/log_pstree.py

@@ -0,0 +1,48 @@
+# 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 logging
+import subprocess
+import sys
+from typing import Optional
+
+from monarch_supervisor.logging import gethostname, initialize_logging
+
+pid: str
+logger: logging.Logger = logging.getLogger(__name__)
+
+
+def extract_pss(pid: str) -> Optional[str]:
+    try:
+        with open(f"/proc/{pid}/smaps_rollup", "r") as f:
+            for line in f.readlines():
+                if line.startswith("Pss:"):  # Check if the line starts with 'Pss:'
+                    return " ".join(line.split()[1:3])
+    except Exception:
+        pass
+    return None
+
+
+def log_pstree_output(pid: int) -> None:
+    pstree_output = subprocess.check_output(["pstree", "-Tap", str(pid)]).decode(
+        "utf-8"
+    )
+    lines = pstree_output.split("\n")
+    logger.info("Process Info")
+    for line in lines:
+        if not line.strip():
+            continue
+        parts = line.split(",")
+        pids = parts[1].split()[0]
+        mem = extract_pss(pids)
+        logger.info(f"{line} {mem}")
+
+
+if __name__ == "__main__":
+    (pid,) = sys.argv[1:]
+    initialize_logging(f"{gethostname()} host-manager")
+    log_pstree_output(int(pid))