Fix actor mesh typechecking (#504)

Summary:
Pull Request resolved: #504

spawn needs to return the actor class, otherwise no endpoint types are inferred.

ActorMeshRef is an internal API detail. Client code should use the name of the class of the actor for references to that actor.
ghstack-source-id: 295681253
exported-using-ghexport

Reviewed By: colin2328, suo

Differential Revision: D78170765

fbshipit-source-id: 8139489108308c4e7a7cf9e87efaac74d31d10e7

Author: zdevito

examples/grpo_actor.py

@@ -8,13 +8,13 @@
 import copy
 import random
 from dataclasses import dataclass
-from typing import Dict, List, Optional, Tuple
+from typing import Any, Dict, List, Optional, Tuple
 
 import torch
 import torch.nn as nn
 import torch.optim as optim
 
-from monarch.actor import Actor, ActorMeshRef, endpoint, proc_mesh
+from monarch.actor import Actor, endpoint, proc_mesh
 from monarch.rdma import RDMABuffer
 from torch.distributions import Categorical, kl_divergence
 
@@ -152,7 +152,7 @@ async def sample_from(self, k: int) -> List[TrajectorySlice]:
 class Scorer(Actor):
     """Evaluates actions and assigns rewards to trajectory slices."""
 
-    def __init__(self, trajectory_queue: ActorMeshRef, replay_buffer: ActorMeshRef):
+    def __init__(self, trajectory_queue: Any, replay_buffer: Any):
         """Initialize the scorer.
 
         Args:
@@ -216,7 +216,7 @@ async def stop(self) -> None:
 class Learner(Actor):
     """Updates policy based on collected experiences using PPO algorithm."""
 
-    def __init__(self, replay_buffer: ActorMeshRef):
+    def __init__(self, replay_buffer: Any):
         """Initialize the learner.
 
         Args:
@@ -238,10 +238,10 @@ def __init__(self, replay_buffer: ActorMeshRef):
         self.policy_version = 0
         self.replay_buffer = replay_buffer
         self.batch_size = 2
-        self.generators: Optional[ActorMeshRef] = None
+        self.generators: Optional[Any] = None
 
     @endpoint
-    async def init_generators(self, generators: ActorMeshRef) -> None:
+    async def init_generators(self, generators: Any) -> None:
         """Set the generators service for weight updates.
 
         Args:

python/monarch/_src/actor/actor_mesh.py

@@ -751,7 +751,7 @@ def _new_with_shape(self, shape: Shape) -> "ActorMeshRef":
         )
 
 
-class ActorMeshRef(MeshTrait, Generic[T]):
+class ActorMeshRef(MeshTrait):
     def __init__(
         self, Class: Type[T], actor_mesh_ref: _ActorMeshRefImpl, mailbox: Mailbox
     ) -> None:

python/monarch/_src/actor/proc_mesh.py

@@ -128,9 +128,7 @@ def _new_with_shape(self, shape: Shape) -> "ProcMesh":
         )
         return ProcMesh(self._proc_mesh, _mock_shape=shape, _device_mesh=device_mesh)
 
-    def spawn(
-        self, name: str, Class: Type[T], *args: Any, **kwargs: Any
-    ) -> Future[ActorMeshRef[T]]:
+    def spawn(self, name: str, Class: Type[T], *args: Any, **kwargs: Any) -> Future[T]:
         if self._mock_shape is not None:
             raise NotImplementedError("NYI: spawn on slice of a proc mesh.")
         return Future(
@@ -408,12 +406,12 @@ def _get_debug_proc_mesh() -> "ProcMesh":
     return _debug_proc_mesh
 
 
-_debug_client_mesh: Optional[ActorMeshRef[DebugClient]] = None
+_debug_client_mesh: Optional[DebugClient] = None
 
 
 # Lazy init for the same reason as above. This is defined in proc_mesh.py
 # instead of debugger.py for circular import reasons.
-def debug_client() -> ActorMeshRef[DebugClient]:
+def debug_client() -> DebugClient:
     global _debug_client_mesh
     if _debug_client_mesh is None:
         _debug_client_mesh = (

python/monarch/actor/__init__.py

@@ -12,7 +12,6 @@
     Accumulator,
     Actor,
     ActorError,
-    ActorMeshRef,
     current_actor_name,
     current_rank,
     current_size,
@@ -35,7 +34,6 @@
     "Accumulator",
     "Actor",
     "ActorError",
-    "ActorMeshRef",
     "current_actor_name",
     "current_rank",
     "current_size",