[monarch] monarch.actor submodule

This is the proposed `monarch.actor` submodule.

The public API is just a flat namespace, determined by `monarch/actor/__init__.py`.

*Everything else* is a private detail, signified by lots of underscores in filenames. That's a little unforunate, but I think draws some good boundaries, and keeps the API simple to the user.

Some principles that should guide us moving forward:
- **Every symbol has at most two locations**: a private location (which should be exactly the same as its location in the source tree), and optionally a public location in the user-facing API.
- We discourage end-users from depending on private locations via underscores, per python convention.

Differential Revision: [D77709116](https://our.internmc.facebook.com/intern/diff/D77709116/)

ghstack-source-id: 294458318
Pull Request resolved: #437

Author: suo

monarch_extension/src/mesh_controller.rs

@@ -507,7 +507,7 @@ impl History {
                 .getattr("RemoteException")
                 .unwrap();
             let pickle = py
-                .import("monarch.actor_mesh")
+                .import("monarch.actor._actor_mesh")
                 .unwrap()
                 .getattr("_pickle")
                 .unwrap();

monarch_tensor_worker/src/stream.rs

@@ -990,7 +990,7 @@ impl StreamActor {
                 })
                 .and_then(|result| -> Result<PythonMessage, WorkerError> {
                     let pickle = py
-                        .import("monarch.actor_mesh")
+                        .import("monarch.actor._actor_mesh")
                         .unwrap()
                         .getattr("_pickle")
                         .unwrap();

python/monarch/__init__.py

@@ -29,7 +29,8 @@
 
 if TYPE_CHECKING:
     from monarch import timer
-    from monarch.allocator import LocalAllocator, ProcessAllocator
+    from monarch.actor._allocator import LocalAllocator, ProcessAllocator
+    from monarch.actor._shape import NDSlice, Shape
     from monarch.common._coalescing import coalescing
 
     from monarch.common.device_mesh import (
@@ -50,11 +51,9 @@
     from monarch.common.pipe import create_pipe, Pipe, remote_generator
     from monarch.common.remote import remote
     from monarch.common.selection import Selection
-    from monarch.common.shape import NDSlice, Shape
     from monarch.common.stream import get_active_stream, Stream
     from monarch.common.tensor import reduce, reduce_, Tensor
     from monarch.fetch import fetch_shard, inspect, show
-    from monarch.future import ActorFuture
     from monarch.gradient_generator import grad_function, grad_generator
     from monarch.notebook import mast_mesh, reserve_torchx as mast_reserve
     from monarch.python_local_mesh import python_local_mesh
@@ -79,8 +78,8 @@
     "function_resolvers": ("monarch.common.function", "resolvers"),
     "Future": ("monarch.common.future", "Future"),
     "RemoteException": ("monarch.common.invocation", "RemoteException"),
-    "Shape": ("monarch.common.shape", "Shape"),
-    "NDSlice": ("monarch.common.shape", "NDSlice"),
+    "Shape": ("monarch.actor._shape", "Shape"),
+    "NDSlice": ("monarch.actor._shape", "NDSlice"),
     "Selection": ("monarch.common.selection", "Selection"),
     "OpaqueRef": ("monarch.common.opaque_ref", "OpaqueRef"),
     "create_pipe": ("monarch.common.pipe", "create_pipe"),
@@ -112,9 +111,8 @@
     "Simulator": ("monarch.simulator.interface", "Simulator"),
     "world_mesh": ("monarch.world_mesh", "world_mesh"),
     "timer": ("monarch.timer", "timer"),
-    "ProcessAllocator": ("monarch.allocator", "ProcessAllocator"),
-    "LocalAllocator": ("monarch.allocator", "LocalAllocator"),
-    "ActorFuture": ("monarch.future", "ActorFuture"),
+    "ProcessAllocator": ("monarch.actor._allocator", "ProcessAllocator"),
+    "LocalAllocator": ("monarch.actor._allocator", "LocalAllocator"),
     "builtins": ("monarch.builtins", "builtins"),
 }
 
@@ -183,7 +181,6 @@ def __getattr__(name):
     "timer",
     "ProcessAllocator",
     "LocalAllocator",
-    "ActorFuture",
     "builtins",
 ]
 assert sorted(__all__) == sorted(_public_api)

python/monarch/_rust_bindings/__init__.pyi

@@ -0,0 +1,8 @@
+# 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.
+
+# True iff the rust extension was built with the tensor engine feature.
+def has_tensor_engine() -> bool: ...

python/monarch/_rust_bindings/monarch_hyperactor/alloc.pyi

@@ -73,7 +73,7 @@ class RemoteAllocatorBase:
     def __new__(
         cls,
         world_id: str,
-        initializer: "monarch.allocator.RemoteAllocInitializer",  # pyre-ignore[11] defined in monarch/python/monarch/allocator.py
+        initializer: "monarch.actor._allocator.RemoteAllocInitializer",  # pyre-ignore[11]
         heartbeat_interval: timedelta = timedelta(seconds=5),
     ) -> Self:
         """

python/monarch/_testing.py

@@ -13,10 +13,10 @@
 from typing import Any, Callable, Dict, Generator, Literal, Optional
 
 import monarch_supervisor
+from monarch.actor._shape import NDSlice
 from monarch.common.client import Client
 from monarch.common.device_mesh import DeviceMesh
 from monarch.common.invocation import DeviceException, RemoteException
-from monarch.common.shape import NDSlice
 from monarch.controller.backend import ProcessBackend
 from monarch.mesh_controller import spawn_tensor_engine
 from monarch.proc_mesh import proc_mesh, ProcMesh

python/monarch/actor/__init__.py

@@ -0,0 +1,38 @@
+# 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.
+
+"""
+Monarch Actor API
+"""
+
+from monarch.actor._actor_mesh import (
+    Accumulator,
+    Actor,
+    ActorError,
+    current_actor_name,
+    current_rank,
+    current_size,
+    endpoint,
+    MonarchContext,
+    ValueMesh,
+)
+from monarch.actor._future import Future
+from monarch.actor._proc_mesh import proc_mesh, ProcMesh
+
+__all__ = [
+    "Accumulator",
+    "Actor",
+    "ActorError",
+    "current_actor_name",
+    "current_rank",
+    "current_size",
+    "endpoint",
+    "MonarchContext",
+    "ValueMesh",
+    "proc_mesh",
+    "ProcMesh",
+    "Future",
+]