torchmonarch-nightly 2025.6.27__cp311-cp311-manylinux2014_x86_64.whl

monarch/tools/mesh_spec.py

@@ -0,0 +1,165 @@
+# 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 string
+from dataclasses import dataclass, field
+from typing import Any, Optional
+
+from monarch.tools.network import get_sockaddr
+from torchx import specs
+
+DEFAULT_REMOTE_ALLOCATOR_PORT = 26600
+
+_TAG_MESHES_PREFIX = "monarch/meshes/${mesh_name}/"
+_TAG_HOST_TYPE: str = _TAG_MESHES_PREFIX + "host_type"
+_TAG_GPUS: str = _TAG_MESHES_PREFIX + "gpus"
+_TAG_TRANSPORT: str = _TAG_MESHES_PREFIX + "transport"
+
+_UNSET_INT = -1
+_UNSET_STR = "__UNSET__"
+
+
+@dataclass
+class MeshSpec:
+    """Doubles as the 'input' specifications of how to setup the mesh role
+    when submitting the job and as the 'info' (describe) API's return value.
+    """
+
+    name: str
+    num_hosts: int
+    host_type: str = _UNSET_STR
+    gpus: int = _UNSET_INT
+    # NOTE: using str over monarch._rust_bindings.monarch_hyperactor.channel.ChannelTransport enum
+    #  b/c the rust binding doesn't have Python enum semantics, hence doesn't serialize well
+    transport: str = "tcp"
+    port: int = DEFAULT_REMOTE_ALLOCATOR_PORT
+    hostnames: list[str] = field(default_factory=list)
+
+    def server_addrs(
+        self, transport: Optional[str] = None, port: Optional[int] = None
+    ) -> list[str]:
+        """
+        Returns the hostnames (servers) in channel address format.
+        `transport` and `port` is typically taken from this mesh spec's fields, but
+        the caller can override them when calling this function.
+        """
+
+        transport = transport or self.transport
+        port = port or self.port
+
+        if transport == "tcp":
+            # need to resolve hostnames to ip address for TCP
+            return [
+                f"tcp!{get_sockaddr(hostname, port)}" for hostname in self.hostnames
+            ]
+        elif transport == "metatls":
+            return [f"metatls!{hostname}:{port}" for hostname in self.hostnames]
+        else:
+            raise ValueError(
+                f"Unsupported transport: {transport}. Must be one of: 'tcp' or 'metatls'"
+            )
+
+
+def _tag(mesh_name: str, tag_template: str) -> str:
+    return string.Template(tag_template).substitute(mesh_name=mesh_name)
+
+
+def tag_as_metadata(mesh_spec: MeshSpec, appdef: specs.AppDef) -> None:
+    appdef.metadata[_tag(mesh_spec.name, _TAG_HOST_TYPE)] = mesh_spec.host_type
+    appdef.metadata[_tag(mesh_spec.name, _TAG_GPUS)] = str(mesh_spec.gpus)
+    appdef.metadata[_tag(mesh_spec.name, _TAG_TRANSPORT)] = mesh_spec.transport
+
+
+def mesh_spec_from_metadata(appdef: specs.AppDef, mesh_name: str) -> Optional[MeshSpec]:
+    for role in appdef.roles:
+        if role.name == mesh_name:
+            return MeshSpec(
+                name=mesh_name,
+                num_hosts=role.num_replicas,
+                host_type=appdef.metadata.get(
+                    _tag(mesh_name, _TAG_HOST_TYPE), _UNSET_STR
+                ),
+                gpus=int(
+                    appdef.metadata.get(_tag(mesh_name, _TAG_GPUS), str(_UNSET_INT))
+                ),
+                transport=appdef.metadata.get(_tag(mesh_name, _TAG_TRANSPORT), "tcp"),
+                port=role.port_map.get("mesh", DEFAULT_REMOTE_ALLOCATOR_PORT),
+            )
+
+    return None
+
+
+def mesh_spec_from_str(mesh_spec_str: str) -> MeshSpec:
+    """Parses the given string into a MeshSpec.
+
+    Args:
+        mesh_spec_str: A string representation of the mesh specification
+            in the format 'NAME:NUM_HOSTS:HOST_TYPE' (e.g. 'trainer:8:gpu.medium').
+    """
+    parts = mesh_spec_str.split(":")
+    assert (
+        len(parts) == 3
+    ), f"`{mesh_spec_str}` is not of the form 'NAME:NUM_HOSTS:HOST_TYPE'"
+
+    name, num_hosts, host_type = parts
+    gpus = specs.resource(h=host_type).gpu
+
+    assert num_hosts.isdigit(), f"`{num_hosts}` is not a number in: {mesh_spec_str}"
+
+    return MeshSpec(name, int(num_hosts), host_type, gpus)
+
+
+@dataclass
+class ServerSpec:
+    """Holds information (as returned by the 'describe' API of the scheduler)
+    about the monarch server. This is the return value of ``monarch.tools.commands.info` API.
+    """
+
+    name: str
+    state: specs.AppState
+    meshes: list[MeshSpec]
+
+    @property
+    def is_running(self) -> bool:
+        return self.state == specs.AppState.RUNNING
+
+    def get_mesh_spec(self, mesh_name: str) -> MeshSpec:
+        for mesh_spec in self.meshes:
+            if mesh_spec.name == mesh_name:
+                return mesh_spec
+
+        raise ValueError(
+            f"Mesh: '{mesh_name}' not found in job: {self.name}. Try one of: {self.get_mesh_names()}"
+        )
+
+    def get_mesh_names(self) -> list[str]:
+        return [m.name for m in self.meshes]
+
+    def to_json(self) -> dict[str, Any]:
+        """Returns the JSON form of this struct that can be printed to console by:
+
+        .. code-block:: python
+
+            import json
+
+            server_spec = ServerSpec(...)
+            print(json.dumps(server_spec, indent=2))
+        """
+
+        return {
+            "name": self.name,
+            "state": self.state.name,
+            "meshes": {
+                mesh.name: {
+                    "host_type": mesh.host_type,
+                    "hosts": mesh.num_hosts,
+                    "gpus": mesh.gpus,
+                    "hostnames": mesh.hostnames,
+                }
+                for mesh in self.meshes
+            },
+        }

monarch/tools/network.py

@@ -0,0 +1,69 @@
+# 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 socket
+from typing import Optional
+
+logger: logging.Logger = logging.getLogger(__name__)
+
+
+def get_sockaddr(hostname: str, port: int) -> str:
+    """Returns either an IPv6 or IPv4 socket address (that supports TCP) of the given hostname and port.
+    The socket address is of the form:
+      1. `{ipv4.address}:{port}` (e.g. `127.0.0.1:8080`)
+      2. `[{ipv6:address}]:{port}` (e.g. `[::1]:8080`)
+
+    The hostname is resolved to an IPv6 (or IPv4 if IPv6 is not available on the host) address that
+    supports `SOCK_STREAM` (TCP).
+
+    Raises a `RuntimeError` if neither ipv6 or ipv4 ip can be resolved from hostname.
+    """
+
+    def resolve_sockaddr(family: socket.AddressFamily) -> Optional[str]:
+        try:
+            # patternlint-disable-next-line python-dns-deps (only used for oss)
+            addrs = socket.getaddrinfo(hostname, port, family, type=socket.SOCK_STREAM)
+            if addrs:
+                family, _, _, _, sockaddr = addrs[0]  # use the first address
+
+                # sockaddr is a tuple (ipv4) or a 4-tuple (ipv6)
+                # in both cases the first element is the ip addr
+                ipaddr = str(sockaddr[0])
+
+                if family == socket.AF_INET6:
+                    socket_address = f"[{ipaddr}]:{port}"
+                else:  # socket.AF_INET
+                    socket_address = f"{ipaddr}:{port}"
+
+                logger.info(
+                    "resolved %s address `%s` for `%s:%d`",
+                    family.name,
+                    socket_address,
+                    hostname,
+                    port,
+                )
+
+                return socket_address
+        except socket.gaierror as e:
+            logger.info(
+                "no %s address that can bind TCP sockets for `%s:%d` (error: %s)",
+                family.name,
+                hostname,
+                port,
+                e,
+            )
+        return None
+
+    for family in [socket.AF_INET6, socket.AF_INET]:
+        if ipaddr := resolve_sockaddr(family):
+            return ipaddr
+
+    raise RuntimeError(
+        f"Unable to resolve `{hostname}` to ipv6 or ipv4 address that can bind TCP socket."
+        " Check the network configuration on the host."
+    )

monarch/worker/__init__.py

@@ -0,0 +1,7 @@
+# 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