# 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 os
import subprocess
import sys
import tempfile
from typing import Any, Awaitable, Callable, Dict, Literal, Optional, Tuple
from monarch._rust_bindings.monarch_hyperactor.host_mesh import (
_spawn_admin as _hy_spawn_admin,
BootstrapCommand,
HostMesh as HyHostMesh,
PyMeshAdminRef,
)
from monarch._rust_bindings.monarch_hyperactor.proc_mesh import ProcMesh as HyProcMesh
from monarch._rust_bindings.monarch_hyperactor.pytokio import PythonTask, Shared
from monarch._rust_bindings.monarch_hyperactor.shape import Extent, Point, Region
from monarch._src.actor.actor_mesh import _Lazy, context
from monarch._src.actor.future import Future
from monarch._src.actor.proc_mesh import _get_bootstrap_args, ProcMesh
from monarch._src.actor.shape import MeshTrait, NDSlice, Shape
from monarch.tools.config.workspace import Workspace
[docs]def default_bootstrap_cmd() -> BootstrapCommand:
"""Get the default bootstrap command for the current environment.
Returns a BootstrapCommand configured with the current Python executable
and environment. This can be used as a base for customization with
``with_env()`` or by modifying its attributes directly.
Returns:
BootstrapCommand: The default bootstrap command.
"""
cmd, args, bootstrap_env = _get_bootstrap_args()
return BootstrapCommand(
cmd,
None,
args if args else [],
bootstrap_env,
)
[docs]def this_host() -> "HostMesh":
"""
The current machine.
This is just shorthand for looking it up via the context
"""
return this_proc().host_mesh
[docs]def this_proc() -> "ProcMesh":
"""
The current singleton process that this specific actor is
running on
"""
return context().actor_instance.proc
[docs]class HostMesh(MeshTrait):
"""
HostMesh represents a collection of compute hosts that can be used to spawn
processes and actors.
Can be used as an async context manager, which will shut down the hosts on
exit:
```
host_mesh = job.state().hosts
async with host_mesh:
# spawn proc meshes, actor meshes, etc.
# shutdown() is called automatically on exit.
```
If you don't want to shutdown the hosts, you don't need to use it as a
context manager.
"""
def __init__(
self,
hy_host_mesh: Shared[HyHostMesh],
region: Region,
stream_logs: bool,
is_fake_in_process: bool,
code_sync_proc_mesh: Optional["_Lazy[ProcMesh]"],
) -> None:
self._inner_host_mesh: Optional[Shared[HyHostMesh]] = hy_host_mesh
self._region = region
self._stream_logs = stream_logs
self._is_fake_in_process = is_fake_in_process
self._code_sync_proc_mesh: Optional["_Lazy[ProcMesh]"] = code_sync_proc_mesh
self._pending_spawns: list[Shared[HyProcMesh]] = []
self._proc_meshes: list["ProcMesh"] = []
[docs] def spawn_procs(
self,
per_host: Dict[str, int] | None = None,
bootstrap: Callable[[], None] | Callable[[], Awaitable[None]] | None = None,
name: str | None = None,
proc_bind: list[dict[str, str]] | None = None,
bootstrap_command: (
BootstrapCommand | Callable[[Point], BootstrapCommand] | None
) = None,
) -> "ProcMesh":
"""Spawn a ProcMesh onto this host mesh.
Args:
per_host: shape of procs per host, e.g. ``{"gpus": 4}``.
bootstrap: optional setup callable run on each proc.
name: optional name for the proc mesh.
proc_bind: optional per-process CPU/NUMA binding config.
Length must equal ``math.prod(per_host.values())``.
Each dict maps binding keys (``cpunodebind``,
``membind``, ``physcpubind``, ``cpus``) to values.
bootstrap_command: optional BootstrapCommand or callable that
returns a BootstrapCommand for each coordinate. The callable
receives a ``Point`` (combined coordinate across host and
per_host dimensions). This allows full customization of the
bootstrap command per coordinate.
"""
if not per_host:
per_host = {}
if not name:
name = "anon"
import math
procs_per_host = math.prod(per_host.values()) if per_host else 1
if proc_bind is not None and len(proc_bind) != procs_per_host:
raise ValueError(
f"proc_bind length ({len(proc_bind)}) must equal "
f"procs_per_host ({procs_per_host})"
)
return self._spawn_nonblocking(
name,
Extent(list(per_host.keys()), list(per_host.values())),
bootstrap,
True,
proc_bind,
bootstrap_command,
)
def _spawn_nonblocking(
self,
name: str,
per_host: Extent,
setup: Callable[[], None] | Callable[[], Awaitable[None]] | None,
_attach_controller_controller: bool,
proc_bind: list[dict[str, str]] | None = None,
bootstrap_command: (
BootstrapCommand | Callable[[Point], BootstrapCommand] | None
) = None,
) -> "ProcMesh":
if set(per_host.labels) & set(self._labels):
# The rust side will catch this too, but this lets us fail fast
raise ValueError(
f"per_host labels {per_host.labels} overlap with host labels {self._labels}"
)
# This is checked inside the task as well, but we can pre-emptively raise
# earlier.
if self._inner_host_mesh is None:
raise RuntimeError("HostMesh has already been shut down")
per_rank_bootstrap: Callable[[Point], BootstrapCommand] | None = None
if bootstrap_command is not None:
if isinstance(bootstrap_command, BootstrapCommand):
# Uniform BootstrapCommand - convert to callable
def make_uniform_bootstrap(
cmd: BootstrapCommand,
) -> Callable[[Point], BootstrapCommand]:
def bootstrap_callable(point: Point) -> BootstrapCommand:
return cmd
return bootstrap_callable
per_rank_bootstrap = make_uniform_bootstrap(bootstrap_command)
else:
# Callable that returns BootstrapCommand
per_rank_bootstrap = bootstrap_command
async def task() -> HyProcMesh:
hy_host_mesh = await self._hy_host_mesh
return await hy_host_mesh.spawn_nonblocking(
context().actor_instance._as_rust(),
name,
per_host,
proc_bind,
per_rank_bootstrap,
)
spawn_shared = PythonTask.from_coroutine(task()).spawn()
self._pending_spawns.append(spawn_shared)
pm = ProcMesh.from_host_mesh(
self,
spawn_shared,
Extent(
self._labels + tuple(per_host.labels),
self.region.slice().sizes + list(per_host.sizes),
).region,
setup,
_attach_controller_controller,
)
self._proc_meshes.append(pm)
return pm
@property
def _ndslice(self) -> NDSlice:
return self.region.slice()
@property
def _labels(self) -> Tuple[str, ...]:
return tuple(self.region.labels)
def _new_with_shape(self, shape: Shape) -> "HostMesh":
if shape.region == self._region:
return self
sliced_hy_hm: Shared[HyHostMesh]
if (hm := self._hy_host_mesh.poll()) is not None:
sliced_hy_hm = Shared.from_value(hm.sliced(shape.region))
else:
async def task() -> HyHostMesh:
return (await self._hy_host_mesh).sliced(shape.region)
sliced_hy_hm = PythonTask.from_coroutine(task()).spawn()
return HostMesh(
sliced_hy_hm,
shape.region,
self.stream_logs,
self.is_fake_in_process,
None,
)
@property
def region(self) -> Region:
return self._region
@property
def stream_logs(self) -> bool:
return self._stream_logs
@classmethod
def _from_initialized_hy_host_mesh(
cls,
hy_host_mesh: HyHostMesh,
region: Region,
stream_logs: bool,
is_fake_in_process: bool,
) -> "HostMesh":
return HostMesh(
Shared.from_value(hy_host_mesh),
region,
stream_logs,
is_fake_in_process,
None,
)
@classmethod
def _from_rust(cls, hy_host_mesh: HyHostMesh) -> "HostMesh":
"""
Create a HostMesh from a Rust HyHostMesh.
This is used when the host was bootstrapped via bootstrap_host()
instead of being allocated through an allocator.
"""
return cls._from_initialized_hy_host_mesh(
hy_host_mesh,
hy_host_mesh.region,
stream_logs=False,
is_fake_in_process=False,
)
[docs] def with_python_executable(self, python_executable: str) -> "HostMesh":
"""
Return a new HostMesh that will use the given Python executable when
spawning procs. Procs spawned from this mesh will also inherit this
Python executable when they call ``this_host().spawn_procs(...)``.
Args:
python_executable: Path to the Python executable to use.
Returns:
A new HostMesh configured to use the specified Python executable.
"""
_, _, bootstrap_env = _get_bootstrap_args()
bootstrap_cmd: BootstrapCommand = BootstrapCommand(
python_executable,
None,
["-m", "monarch._src.actor.bootstrap_main"],
bootstrap_env,
)
async def task() -> HyHostMesh:
return (await self._hy_host_mesh).with_bootstrap(bootstrap_cmd)
return HostMesh(
PythonTask.from_coroutine(task()).spawn(),
self._region,
self._stream_logs,
self._is_fake_in_process,
None,
)
# pyrefly: ignore [invalid-annotation]
def __reduce_ex__(self, protocol: ...) -> Tuple[Any, Tuple[Any, ...]]:
# A pending host mesh has no HostMeshRef yet. When mesh-reference
# collection is active, reserve an out-of-band slot (filled sender-side
# once the mesh resolves) and reconstruct from the popped mesh;
# otherwise fall through to the ordinary reduce.
if self._hy_host_mesh.poll() is None:
from monarch._rust_bindings.monarch_hyperactor.pickle import (
reserve_mesh_reference,
)
from monarch._src.actor.pickle import _MeshSlot
if reserve_mesh_reference(self._hy_host_mesh):
return HostMesh._from_initialized_hy_host_mesh, (
_MeshSlot(),
self._region,
self.stream_logs,
self.is_fake_in_process,
)
return HostMesh, (
self._hy_host_mesh,
self._region,
self.stream_logs,
self.is_fake_in_process,
None,
)
@property
def is_fake_in_process(self) -> bool:
return self._is_fake_in_process
# pyrefly: ignore [bad-override]
def __eq__(self, other: "HostMesh") -> bool:
# Should we include code sync proc mesh?
return (
self._initialized_mesh() == other._initialized_mesh()
and self._region == other._region
and self.stream_logs == other.stream_logs
and self.is_fake_in_process == other.is_fake_in_process
)
def _initialized_mesh(self) -> HyHostMesh:
return self._hy_host_mesh.poll() or self._hy_host_mesh.block_on()
async def _flush_pending_spawns(self) -> None:
for shared in self._pending_spawns:
try:
await shared
except Exception:
pass
self._pending_spawns.clear()
for pm in self._proc_meshes:
await pm._flush_pending_actor_spawns()
try:
await pm._logging_manager.flush_async()
except Exception:
pass
[docs] def shutdown(self) -> Future[None]:
"""
Shutdown the host mesh and all of its processes. It will throw an exception
if this host mesh is a *reference* rather than *owned*, which can happen
if this `HostMesh` object was received from a remote actor or if it was
produced by slicing.
After shutting down, the hosts in this mesh will be unusable, and no new
HostMeshes will be able to connect to them.
If you want to stop everything on the host but keep them available for
new clients, use `stop()` instead.
This is run automatically on __aexit__ when used as an async context manager.
Returns:
Future[None]: A future that completes when the host mesh has been shut down.
"""
async def task() -> None:
await self._flush_pending_spawns()
hy_mesh = await self._hy_host_mesh
await hy_mesh.shutdown(context().actor_instance._as_rust())
# Remove the inner host mesh to clean up associated memory.
self._inner_host_mesh = None
return Future(coro=task())
[docs] def stop(self) -> Future[None]:
"""
Stop the host mesh, releasing all resources but keeping worker
processes alive for reconnection. A new HostMesh can be created that
points to the same hosts.
Like `shutdown`, this throws if the host mesh is a reference
rather than owned.
Returns:
Future[None]: A future that completes when the host mesh has been stopped.
"""
async def task() -> None:
await self._flush_pending_spawns()
hy_mesh = await self._hy_host_mesh
await hy_mesh.stop(context().actor_instance._as_rust())
return Future(coro=task())
async def __aenter__(self) -> "HostMesh":
if self._inner_host_mesh is None:
raise RuntimeError("HostMesh has already been shut down")
return self
async def __aexit__(
self, exc_type: object, exc_val: object, exc_tb: object
) -> None:
# In case there are multiple nested "async with" statements, we only
# want it to close once.
if self._inner_host_mesh is not None:
await self.shutdown()
[docs] async def sync_workspace(
self,
workspace: Workspace,
conda: bool = False,
auto_reload: bool = False,
) -> None:
"""
Sync local code changes to the remote hosts.
Args:
workspace: The workspace to sync.
conda: If True, also sync the currently activated conda env.
auto_reload: If True, automatically reload the workspace on changes.
"""
if self._code_sync_proc_mesh:
await self._code_sync_proc_mesh.get()._sync_workspace(
workspace, conda, auto_reload
)
else:
raise RuntimeError(
"cannot call sync_workspace on a sliced host mesh or one that was sent over an actor endpoint"
)
@property
def initialized(self) -> Future[Literal[True]]:
"""
Future completes with 'True' when the `HostMesh` has initialized.
Because `HostMesh` are remote objects, there is no guarantee that the `HostMesh` is
still usable after this completes, only that at some point in the past it was usable.
"""
hm: Shared[HyHostMesh] = self._hy_host_mesh
async def task() -> Literal[True]:
await hm
return True
return Future(coro=task())
@property
def _hy_host_mesh(self) -> Shared[HyHostMesh]:
if self._inner_host_mesh is None:
raise RuntimeError("HostMesh has already been shut down")
return self._inner_host_mesh
def _spawn_admin(
host_meshes: list["HostMesh"],
admin_addr: Optional[str] = None,
telemetry_url: Optional[str] = None,
) -> "Future[tuple[str, PyMeshAdminRef]]":
"""
Spawn a MeshAdminAgent aggregating topology across one or more HostMeshes.
Returns ``(admin_url, admin_ref)`` where ``admin_ref`` is an opaque
capability token for immediate use only.
Args:
host_meshes: One or more HostMeshes whose hosts the admin
will aggregate for introspection. Must not be empty.
admin_addr: Optional socket address for the admin HTTP server.
When ``None``, reads ``MESH_ADMIN_ADDR`` from config.
telemetry_url: Optional base URL of the Monarch telemetry dashboard.
Returns:
Future[tuple[str, PyMeshAdminRef]]: (admin_url, admin_ref).
Raises:
ValueError: If host_meshes is empty.
"""
if not host_meshes:
raise ValueError("_spawn_admin requires at least one HostMesh")
async def task() -> tuple[str, PyMeshAdminRef]:
hy_meshes = [await m._hy_host_mesh for m in host_meshes]
admin_url, admin_ref = await _hy_spawn_admin(
hy_meshes, context().actor_instance._as_rust(), admin_addr, telemetry_url
)
os.environ["MONARCH_ADMIN_URL"] = admin_url
return admin_url, admin_ref
return Future(coro=task())
[docs]def hosts_from_config(name: str) -> HostMesh:
"""
Get the host mesh 'name' from the monarch configuration for the project.
This config can be modified so that the same code can create meshes from scheduler sources,
and different sizes etc.
WARNING: This function is a standin so that our getting_started example code works. The real implementation
needs an RFC design.
"""
num_hosts = 2
tmpdir = tempfile.mkdtemp(prefix="monarch_hosts_from_config_")
workers = []
for i in range(num_hosts):
addr = f"ipc://{tmpdir}/{name}_{i}"
env = {**os.environ}
cmd = [
sys.executable,
"-c",
"from monarch.actor import run_worker_loop_forever; "
f'run_worker_loop_forever(address="{addr}", '
'ca="trust_all_connections")',
]
subprocess.Popen(cmd, env=env, start_new_session=True)
workers.append(addr)
from monarch._src.actor.bootstrap import attach_to_workers
return attach_to_workers(
name=name,
ca="trust_all_connections",
# pyrefly: ignore [bad-argument-type]
workers=workers,
)
