torchmonarch-nightly 2025.6.27__cp312-cp312-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(), indent=2, 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,251 @@
+# 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 logging
+import os
+import time
+from datetime import timedelta
+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, AppState, CfgVal
+from torchx.specs.builders import parse_args
+from torchx.util.types import decode, decode_optional
+
+logger: logging.Logger = logging.getLogger(__name__)
+
+
+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
+
+    # host status grouped by mesh (role) names
+    replica_status = {r.role: r.replicas for r in status.roles}
+
+    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"
+
+        # null-guard since some schedulers do not fill replica_status
+        if host_status := replica_status.get(role.name):
+            spec.hostnames = [h.hostname for h in host_status]
+
+        mesh_specs.append(spec)
+
+    return ServerSpec(name=appdef.name, state=status.state, meshes=mesh_specs)
+
+
+_5_SECONDS = timedelta(seconds=5)
+
+
+async def server_ready(
+    server_handle: str, check_interval: timedelta = _5_SECONDS
+) -> Optional[ServerSpec]:
+    """Waits until the server's job is in RUNNING state to returns the server spec.
+    Returns `None` if the server does not exist.
+
+    NOTE: Certain fields such as `hostnames` is only filled (and valid) when the server is RUNNING.
+
+    Usage:
+
+    .. code-block:: python
+
+        server_info = await server_ready("slurm:///123")
+        if not server_info:
+            print(f"Job does not exist")
+        else:
+            if server_info.is_running:
+                for mesh in server_info.meshes:
+                    connect_to(mesh.hostnames)
+            else:
+                print(f"Job in {server_info.state} state. Hostnames are not available")
+
+    """
+
+    while True:
+        server_spec = info(server_handle)
+
+        if not server_spec:  # server not found
+            return None
+
+        if server_spec.state <= AppState.PENDING:  # UNSUBMITTED or SUBMITTED or PENDING
+            # NOTE: TorchX currently does not have async APIs so need to loop-on-interval
+            # TODO maybe inverse exponential backoff instead of constant interval?
+            check_interval_seconds = check_interval.total_seconds()
+            logger.info(
+                "waiting for %s to be %s (current: %s), will check again in %g seconds...",
+                server_handle,
+                AppState.RUNNING,
+                server_spec.state,
+                check_interval_seconds,
+            )
+            time.sleep(check_interval_seconds)
+            continue
+        else:
+            return server_spec
+
+
+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,58 @@
+# 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,
+    program: str = "monarch_bootstrap",  # installed with monarch wheel (as console script)
+) -> 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)
+        program: path to the binary that the remote process allocator spawns on an allocation request
+    """
+
+    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",  # run "cargo install monarch_hyperactor" to get this binary
+            args=[
+                f"--port={port}",
+                f"--program={program}",
+            ],
+            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