[monarch] refactor monarch/ callsites to use monarch.actor

as title

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

ghstack-source-id: 294437417
Pull Request resolved: #438

Author: suo

examples/grpo_actor.py

@@ -13,9 +13,8 @@
 import torch
 import torch.nn as nn
 import torch.optim as optim
-from monarch.actor_mesh import Actor, ActorMeshRef, endpoint
 
-from monarch.proc_mesh import proc_mesh
+from monarch.actor import Actor, ActorMeshRef, endpoint, proc_mesh
 from monarch.rdma import RDMABuffer
 from torch.distributions import Categorical, kl_divergence
 

examples/notebooks/ping_pong.ipynb

@@ -16,7 +16,7 @@
             "metadata": {},
             "source": [
                 "## Hello World\n",
-                "Actors are spawned in Process meshes via the `monarch.proc_mesh` API. For those familiar with distributed systems, it can be helpful to think of each Actor as a server with endpoints that can be called."
+                "Actors are spawned in Process meshes via the `monarch.actor` API. For those familiar with distributed systems, it can be helpful to think of each Actor as a server with endpoints that can be called."
             ]
         },
         {
@@ -43,8 +43,8 @@
             "source": [
                 "import asyncio\n",
                 "\n",
-                "from monarch.proc_mesh import proc_mesh, ProcMesh\n",
-                "from monarch.actor_mesh import Actor, endpoint, current_rank\n",
+                "from monarch.actor import proc_mesh, ProcMesh\n",
+                "from monarch.actor import Actor, endpoint, current_rank\n",
                 "\n",
                 "NUM_ACTORS=4\n",
                 "\n",
@@ -168,8 +168,8 @@
             "source": [
                 "import asyncio\n",
                 "\n",
-                "from monarch.proc_mesh import proc_mesh, ProcMesh\n",
-                "from monarch.actor_mesh import Actor, endpoint, current_rank\n",
+                "from monarch.actor import proc_mesh, ProcMesh\n",
+                "from monarch.actor import Actor, endpoint, current_rank\n",
                 "\n",
                 "class ExampleActor(Actor):\n",
                 "    def __init__(self, actor_name):\n",

examples/notebooks/spmd_ddp.ipynb

@@ -24,8 +24,8 @@
                 "import torch.nn as nn\n",
                 "import torch.optim as optim\n",
                 "\n",
-                "from monarch.proc_mesh import proc_mesh\n",
-                "from monarch.actor_mesh import Actor, current_rank, endpoint\n",
+                "from monarch.actor import proc_mesh\n",
+                "from monarch.actor import Actor, current_rank, endpoint\n",
                 "\n",
                 "from torch.nn.parallel import DistributedDataParallel as DDP\n",
                 "\n",

monarch_extension/src/mesh_controller.rs

@@ -30,7 +30,6 @@ use hyperactor::PortRef;
 use hyperactor::cap::CanSend;
 use hyperactor::mailbox::MailboxSenderError;
 use hyperactor_mesh::Mesh;
-use hyperactor_mesh::ProcMesh;
 use hyperactor_mesh::actor_mesh::RootActorMesh;
 use hyperactor_mesh::shared_cell::SharedCell;
 use hyperactor_mesh::shared_cell::SharedCellRef;

python/monarch/_testing.py

@@ -14,12 +14,12 @@
 
 import monarch_supervisor
 from monarch.actor._shape import NDSlice
+from monarch.actor import proc_mesh, ProcMesh
 from monarch.common.client import Client
 from monarch.common.device_mesh import DeviceMesh
 from monarch.common.invocation import DeviceException, RemoteException
 from monarch.controller.backend import ProcessBackend
 from monarch.mesh_controller import spawn_tensor_engine
-from monarch.proc_mesh import proc_mesh, ProcMesh
 from monarch.python_local_mesh import PythonLocalContext
 from monarch.rust_local_mesh import (
     local_mesh,

python/monarch/actor/__init__.py

@@ -12,27 +12,34 @@
     Accumulator,
     Actor,
     ActorError,
+    ActorMeshRef,
     current_actor_name,
     current_rank,
     current_size,
     endpoint,
     MonarchContext,
+    Point,
+    send,
     ValueMesh,
 )
 from monarch.actor._future import Future
-from monarch.actor._proc_mesh import proc_mesh, ProcMesh
+from monarch.actor._proc_mesh import local_proc_mesh, proc_mesh, ProcMesh
 
 __all__ = [
     "Accumulator",
     "Actor",
     "ActorError",
+    "ActorMeshRef",
     "current_actor_name",
     "current_rank",
     "current_size",
     "endpoint",
+    "Future",
+    "local_proc_mesh",
     "MonarchContext",
-    "ValueMesh",
+    "Point",
     "proc_mesh",
     "ProcMesh",
-    "Future",
+    "send",
+    "ValueMesh",
 ]

python/monarch/mesh_controller.py

@@ -34,7 +34,7 @@
     ActorId,
 )
 from monarch.actor._shape import NDSlice
-from monarch.actor_mesh import Port, PortTuple
+from monarch.actor._actor_mesh import Port, PortTuple
 from monarch.common import messages
 from monarch.common.controller_api import TController
 from monarch.common.invocation import Seq

python/tests/error_test_binary.py

@@ -13,8 +13,7 @@
 
 from monarch._rust_bindings.monarch_extension.panic import panicking_function
 
-from monarch.actor_mesh import Actor, endpoint, send
-from monarch.proc_mesh import proc_mesh
+from monarch.actor import Actor, endpoint, proc_mesh, send
 
 
 class ErrorActor(Actor):

python/tests/test_actor_error.py

@@ -10,8 +10,7 @@
 
 import pytest
 from monarch._rust_bindings.monarch_hyperactor.proc_mesh import ProcEvent
-from monarch.actor_mesh import Actor, ActorError, endpoint, send
-from monarch.proc_mesh import local_proc_mesh, proc_mesh
+from monarch.actor import Actor, ActorError, endpoint, local_proc_mesh, proc_mesh
 
 
 class ExceptionActor(Actor):

python/tests/test_allocator.py

@@ -32,14 +32,20 @@
     ChannelAddr,
     ChannelTransport,
 )
+from monarch.actor import (
+    Actor,
+    current_rank,
+    current_size,
+    endpoint,
+    ProcMesh,
+    ValueMesh,
+)
 from monarch.actor._allocator import (
     ALLOC_LABEL_PROC_MESH_NAME,
     RemoteAllocator,
     StaticRemoteAllocInitializer,
     TorchXRemoteAllocInitializer,
 )
-from monarch.actor_mesh import Actor, current_rank, current_size, endpoint, ValueMesh
-from monarch.proc_mesh import ProcMesh
 from monarch.tools.mesh_spec import MeshSpec, ServerSpec
 from monarch.tools.network import get_sockaddr