torchmonarch-nightly 2025.6.4__cp310-cp310-manylinux2014_x86_64.whl

monarch/tools/cli.py

@@ -0,0 +1,167 @@
+# 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 argparse
+import json
+import sys
+
+from monarch.tools.commands import (
+    bounce,
+    component_args_from_cli,
+    create,
+    info,
+    kill,
+    stop,
+    torchx_runner,
+)
+from monarch.tools.config import (  # @manual=//monarch/python/monarch/tools/config/meta:defaults
+    Config,
+    defaults,
+)
+from torchx.specs.finder import get_component
+
+
+def config_from_cli_args(args: argparse.Namespace) -> Config:
+    config = defaults.config(args.scheduler, args.workspace)
+
+    if args.scheduler_args:
+        with torchx_runner() as runner:
+            opts = runner.scheduler_run_opts(config.scheduler)
+            for cfg_str in args.scheduler_args:
+                parsed_cfg = opts.cfg_from_str(cfg_str)
+                config.scheduler_args.update(parsed_cfg)
+
+    config.dryrun = args.dryrun
+    return config
+
+
+class CreateCmd:
+    def add_arguments(self, subparser: argparse.ArgumentParser) -> None:
+        subparser.add_argument(
+            "-s",
+            "--scheduler",
+            type=str,
+            help="Scheduler to submit to",
+        )
+        subparser.add_argument(
+            "-cfg",
+            "--scheduler_args",
+            default=[],
+            action="append",
+            help="Scheduler args (e.g. `-cfg cluster=foo -cfg user=bar`)",
+        )
+        subparser.add_argument(
+            "--dryrun",
+            action="store_true",
+            default=False,
+            help="Just prints the scheduler request",
+        )
+        subparser.add_argument(
+            "--workspace",
+            help="The local directory to build into the job's image and make available on the job."
+            " Pass --workspace='' to disable any default workspaces configured for the scheduler",
+        )
+        subparser.add_argument(
+            "--component",
+            help="A custom TorchX component to use",
+        )
+        subparser.add_argument(
+            "-arg",
+            "--component_args",
+            default=[],
+            action="append",
+            help="Arguments to the component fn (e.g. `-arg a=b -arg c=d` to pass as `component_fn(a=b, c=d)`)",
+        )
+
+    def run(self, args: argparse.Namespace) -> None:
+        config = config_from_cli_args(args)
+
+        component_fn = (
+            get_component(args.component).fn
+            if args.component
+            else defaults.component_fn(config.scheduler)
+        )
+        component_args = component_args_from_cli(component_fn, args.component_args)
+        handle = create(config, component_fn)(**component_args)
+        print(handle)
+
+
+class CommonArguments:
+    @staticmethod
+    def add_server_handle(subparser: argparse.ArgumentParser) -> None:
+        subparser.add_argument(
+            "server_handle",
+            type=str,
+            help="monarch server handle (e.g. slurm:///job_id)",
+        )
+
+
+class InfoCmd:
+    def add_arguments(self, subparser: argparse.ArgumentParser) -> None:
+        CommonArguments.add_server_handle(subparser)
+
+    def run(self, args: argparse.Namespace) -> None:
+        server_spec = info(args.server_handle)
+        if server_spec is None:
+            print(
+                f"Server: {args.server_handle} does not exist",
+                file=sys.stderr,
+            )
+        else:
+            json.dump(server_spec.to_json(), fp=sys.stdout)
+
+
+class KillCmd:
+    def add_arguments(self, subparser: argparse.ArgumentParser) -> None:
+        CommonArguments.add_server_handle(subparser)
+
+    def run(self, args: argparse.Namespace) -> None:
+        kill(args.server_handle)
+
+
+class BounceCmd:
+    def add_arguments(self, subparser: argparse.ArgumentParser) -> None:
+        CommonArguments.add_server_handle(subparser)
+
+    def run(self, args: argparse.Namespace) -> None:
+        bounce(args.server_handle)
+
+
+class StopCmd:
+    def add_arguments(self, subparser: argparse.ArgumentParser) -> None:
+        CommonArguments.add_server_handle(subparser)
+
+    def run(self, args: argparse.Namespace) -> None:
+        stop(args.server_handle)
+
+
+def get_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(description="Monarch CLI")
+    subparser = parser.add_subparsers(title="COMMANDS")
+
+    for cmd_name, cmd in {
+        "create": CreateCmd(),
+        "info": InfoCmd(),
+        "kill": KillCmd(),
+        # --- placeholder subcommands (not yet implemented) ---
+        "bounce": BounceCmd(),
+        "stop": StopCmd(),
+    }.items():
+        cmd_parser = subparser.add_parser(cmd_name)
+        cmd.add_arguments(cmd_parser)
+        cmd_parser.set_defaults(func=cmd.run)
+    return parser
+
+
+def main(argv: list[str] = sys.argv[1:]) -> None:
+    parser = get_parser()
+    args = parser.parse_args(argv)
+    args.func(args)
+
+
+if __name__ == "__main__":
+    main()

monarch/tools/commands.py

@@ -0,0 +1,189 @@
+# 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 argparse
+import functools
+import inspect
+import os
+from typing import Any, Callable, Mapping, Optional, Union
+
+from monarch.tools.config import (  # @manual=//monarch/python/monarch/tools/config/meta:defaults
+    Config,
+    defaults,
+)
+
+from monarch.tools.mesh_spec import mesh_spec_from_metadata, ServerSpec
+
+from torchx.runner import Runner
+from torchx.specs import AppDef, AppDryRunInfo, CfgVal
+from torchx.specs.builders import parse_args
+from torchx.util.types import decode, decode_optional
+
+
+def torchx_runner() -> Runner:
+    # namespace is currently unused so make it empty str
+    # so that server handle is short (e.g. slurm:///job-id)
+    _EMPTY_NS = ""
+    return Runner(_EMPTY_NS, defaults.scheduler_factories())
+
+
+def component_args_from_cli(
+    component_fn: Callable[..., AppDef], component_args: list[str]
+) -> dict[str, Any]:
+    """Parses component function's arguments from 'argname=argvalue' strings.
+
+    Returns: component arguments kwarg-ified.
+    """
+
+    cli_fied_component_args = []
+    for arg in component_args:
+        argname = arg.split("=")[0]
+        # torchx auto-generates an argparse parser for component function based
+        # type-hints and docstring as if the component was a CLI itself so we have to
+        # CLI arg-ify the component arguments by adding a "-" for
+        # single-char argnames (short arg) and "--" for multi-char (long arg)
+        cli_fied_component_args.append(f"-{arg}" if len(argname) == 1 else f"--{arg}")
+
+    parsed_args: argparse.Namespace = parse_args(component_fn, cli_fied_component_args)
+
+    # TODO kiuk@ logic below needs to move into torchx.specs.builders.parse_args()
+    #  which is copied from torchx.specs.builders.materialize_appdef()
+    #  parse_args() returns all the component parameters parsed from cli inputs
+    #  as a string. Additional parameter type matching needs to be done (as below)
+    #  to turn the CLI inputs to component function arguments.
+    component_kwargs = {}
+
+    parameters = inspect.signature(component_fn).parameters
+    for param_name, parameter in parameters.items():
+        arg_value = getattr(parsed_args, param_name)
+        parameter_type = parameter.annotation
+        parameter_type = decode_optional(parameter_type)
+        arg_value = decode(arg_value, parameter_type)
+        if parameter.kind == inspect.Parameter.VAR_POSITIONAL:
+            raise TypeError(
+                f"component fn param `{param_name}` is a '*arg' which is not supported; consider changing the type to a list"
+            )
+        elif parameter.kind == inspect.Parameter.VAR_KEYWORD:
+            raise TypeError(
+                f"component fn param `{param_name}` is a '**kwargs' which is not supported; consider changing the type to a dict or explicitly declare the params"
+            )
+        else:
+            component_kwargs[param_name] = arg_value
+
+    return component_kwargs
+
+
+def create(
+    config: Config,
+    component_fn: Optional[Callable[..., AppDef]] = None,
+) -> Callable[..., Union[str, AppDryRunInfo]]:
+    """Creates a monarch server by submitting it as a job to the target scheduler.
+
+    Note that this function returns a `Callable` that has to be called with the
+    same arguments that one would call the `component_fn` to actually submit
+    the job that runs the monarch server.
+
+    Usage:
+
+    .. doc-test::
+
+        from monarch.tools.config import defaults
+
+        config = defaults.config(scheduler="slurm")
+        config.scheduler_args.update(
+            {
+                "partition": "prod",
+                "mail-user": "foo@bar.com",
+                "mail-type": "FAIL",
+            }
+        )
+        config.dryrun = True
+
+        create(default_config)(host_type="gpu.medium", num_hosts=4)
+
+
+    Args:
+        scheduler: where to submit a job that runs the server
+        scheduler_args: scheduler configs
+        component_fn: a function that returns the AppDef (job def).
+            If not provided, defaults to the configured default for the scheduler
+            (in most cases ``monarch.tools.components.hyperactor.proc_mesh``)
+    """
+    scheduler: str = config.scheduler
+    cfg: Mapping[str, CfgVal] = config.scheduler_args
+    component: Callable[..., AppDef] = component_fn or defaults.component_fn(scheduler)
+
+    @functools.wraps(component)
+    def _run(*args: Any, **kwargs: Any) -> Union[str, AppDryRunInfo]:
+        # for logging call-site context in application metadata
+        os.environ["TORCHX_CONTEXT_NAME"] = os.getenv("TORCHX_CONTEXT_NAME", "monarch")
+
+        appdef = component(*args, **kwargs)
+
+        with torchx_runner() as runner:
+            info = runner.dryrun(appdef, scheduler, cfg, config.workspace)
+
+            info_json_fmt = AppDryRunInfo(
+                info.request,
+                fmt=defaults.dryrun_info_formatter(info),
+            )
+            info_json_fmt._app = info._app
+            info_json_fmt._cfg = info._cfg
+            info_json_fmt._scheduler = info._scheduler
+
+            if config.dryrun:
+                return info_json_fmt
+            else:
+                server_handle = runner.schedule(info)
+                return server_handle
+
+    return _run
+
+
+def info(server_handle: str) -> Optional[ServerSpec]:
+    """Calls the ``describe`` API on the scheduler hosting the server to get
+    information about it.
+
+    Returns ``None`` if the server's job is not found in the scheduler's
+    control-plane. This can happen if the job does not exist
+    (e.g. typo in the server_handle) or the job already exited a long time ago.
+
+    NOTE: This function can return non-empty info for jobs that have
+    exited recently.
+    """
+    with torchx_runner() as runner:
+        status = runner.status(server_handle)
+        if status is None:
+            return None
+
+        appdef = runner.describe(server_handle)
+        if appdef is None:
+            return None
+
+    mesh_specs = []
+    for role in appdef.roles:
+        spec = mesh_spec_from_metadata(appdef, role.name)
+        assert spec is not None, "cannot be 'None' since we iterate over appdef's roles"
+        mesh_specs.append(spec)
+
+    return ServerSpec(name=appdef.name, state=status.state, meshes=mesh_specs)
+
+
+def kill(server_handle: str) -> None:
+    with torchx_runner() as runner:
+        runner.cancel(server_handle)
+
+
+def bounce(server_handle: str) -> None:
+    """(re)starts the server's processes without tearing down the server's job."""
+    raise NotImplementedError("`bounce` is not yet implemented")
+
+
+def stop(server_handle: str) -> None:
+    """Stops the server's unix processes without tearing down the server's job."""
+    raise NotImplementedError("`stop` is not yet implemented")

monarch/tools/components/__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

monarch/tools/components/hyperactor.py

@@ -0,0 +1,57 @@
+# 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 getpass
+from typing import Optional
+
+from monarch.tools import mesh_spec
+from monarch.tools.mesh_spec import mesh_spec_from_str
+from torchx import specs
+
+_DEFAULT_MESHES = ["mesh_0:1:gpu.small"]
+
+_USER: str = getpass.getuser()
+
+__version__ = "latest"  # TODO get version from monarch.__version_
+
+
+def proc_mesh(
+    name: str = f"monarch-{_USER}",
+    image: str = f"ghcr.io/pytorch-labs/monarch:{__version__}",  # TODO docker needs to be built and pushed to ghcr
+    meshes: list[str] = _DEFAULT_MESHES,
+    env: Optional[dict[str, str]] = None,
+    port: int = mesh_spec.DEFAULT_REMOTE_ALLOCATOR_PORT,
+) -> specs.AppDef:
+    """
+    Args:
+        name: the name of the monarch server job
+        image: docker image to run the job on, for slurm, image is the dir the job is run from
+        meshes: list of mesh specs of the form "{name}:{num_hosts}:{host_type}"
+        env: environment variables to be passed to the main command (e.g. ENV1=v1,ENV2=v2,ENV3=v3)
+        port: the port that the remote process allocator runs on (must be reachable from the client)
+    """
+
+    appdef = specs.AppDef(name)
+
+    for mesh in [mesh_spec_from_str(mesh) for mesh in meshes]:
+        mesh_role = specs.Role(
+            name=mesh.name,
+            image=image,
+            entrypoint="process_allocator",  # 'cargo install monarch_hyperactor' to get this binary
+            args=[
+                "mesh-worker",
+                f"--port={port}",
+                "--program=monarch_bootstrap",  # installed with monarch wheel (as console script)
+            ],
+            num_replicas=mesh.num_hosts,
+            resource=specs.resource(h=mesh.host_type),
+            env=env or {},
+            port_map={"mesh": port},
+        )
+        appdef.roles.append(mesh_role)
+
+    return appdef

monarch/tools/config/__init__.py

@@ -0,0 +1,20 @@
+# 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
+from dataclasses import dataclass, field
+from typing import Any, Optional
+
+
+NOT_SET: str = "__NOT_SET__"
+
+
+@dataclass
+class Config:
+    scheduler: str = NOT_SET
+    scheduler_args: dict[str, Any] = field(default_factory=dict)
+    workspace: Optional[str] = None
+    dryrun: bool = False

monarch/tools/config/defaults.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-strict
+
+"""Defines defaults for ``monarch.tools``"""
+
+from typing import Callable, Optional
+
+from monarch.tools.components import hyperactor
+from monarch.tools.config import Config
+
+from torchx import specs
+from torchx.schedulers import (
+    docker_scheduler,
+    kubernetes_scheduler,
+    local_scheduler,
+    SchedulerFactory,
+    slurm_scheduler,
+)
+
+
+def component_fn(scheduler: str) -> Callable[..., specs.AppDef]:
+    """The default TorchX component function for the scheduler"""
+    return hyperactor.proc_mesh
+
+
+def scheduler_factories() -> dict[str, SchedulerFactory]:
+    """Supported schedulers (name -> scheduler static factory method)"""
+    return {  # pyre-ignore[7]
+        # --- local schedulers (no multi-host support) ---
+        "local_cwd": local_scheduler.create_scheduler,
+        "local_docker": docker_scheduler.create_scheduler,
+        # --- remote schedulers (yes multi-host support) ---
+        "slurm": slurm_scheduler.create_scheduler,
+        "k8s": kubernetes_scheduler.create_scheduler,
+    }
+
+
+def config(scheduler: str, workspace: Optional[str] = None) -> Config:
+    """The default :py:class:`~monarch.tools.config.Config` to use when submitting to the provided ``scheduler``."""
+    return Config(scheduler=scheduler, workspace=workspace)
+
+
+def dryrun_info_formatter(dryrun_info: specs.AppDryRunInfo) -> Callable[..., str]:
+    """Used to attach a formatter to the dryrun info when running
+    :py:function:`~monarch.tools.commands.create` in ``dryrun`` mode so that
+    the returned ``AppDryrunInfo`` can be printed to console.
+    """
+    # no-op, use the default formatter already attached to the dryrun info
+    return dryrun_info._fmt

monarch/tools/mesh_spec.py

@@ -0,0 +1,121 @@
+# 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
+from typing import Any, Optional
+
+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"
+
+
+@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
+    gpus: int
+    port: int = DEFAULT_REMOTE_ALLOCATOR_PORT
+
+
+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)
+
+
+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), ""),
+                gpus=int(appdef.metadata.get(_tag(mesh_name, _TAG_GPUS), "-1")),
+                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]
+
+    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,
+                }
+                for mesh in self.meshes
+            },
+        }

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