[monarch] monarch.actor submodule

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/actor/__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
+"""