monarch.actor submodule (#437)

Summary:
Pull Request resolved: #437

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.
ghstack-source-id: 294760596
exported-using-ghexport

Reviewed By: mariusae

Differential Revision: D77709116

fbshipit-source-id: ca9a0190ac28cdead452859396bc7d1d9596c880

Author: suo

monarch_extension/src/mesh_controller.rs

@@ -511,7 +511,7 @@ impl History {
                 .getattr("RemoteException")
                 .unwrap();
             let pickle = py
-                .import("monarch.actor_mesh")
+                .import("monarch._src.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._src.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._src.actor.allocator import LocalAllocator, ProcessAllocator
+    from monarch._src.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._src.actor.shape", "Shape"),
+    "NDSlice": ("monarch._src.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._src.actor.allocator", "ProcessAllocator"),
+    "LocalAllocator": ("monarch._src.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._src.actor.allocator.RemoteAllocInitializer",  # pyre-ignore[11]
         heartbeat_interval: timedelta = timedelta(seconds=5),
     ) -> Self:
         """

python/monarch/_src/__init__.py

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