expose specific-rank selection (#506)

Summary:
Pull Request resolved: #506

only exposed via the low-level `send` API, not a first-class adverb on the endpoints (since I don't know how to make a good API for it).

Reviewed By: dulinriley

Differential Revision: D78172884

fbshipit-source-id: bd4d8d18d99af6039ea19e0aa1a9750ea8c1eae0

Author: suo

python/monarch/_src/actor/actor_mesh.py

@@ -125,7 +125,7 @@ def set(debug_context: "DebugContext") -> None:
 _load_balancing_seed = random.Random(4)
 
 
-Selection = Literal["all", "choose"]  # TODO: replace with real selection objects
+Selection = Literal["all", "choose"] | int  # TODO: replace with real selection objects
 
 
 # standin class for whatever is the serializable python object we use
@@ -202,7 +202,6 @@ def cast(
             idx = _load_balancing_seed.randrange(len(self._shape))
             actor_rank = self._shape.ndslice[idx]
             self._mailbox.post(self._please_replace_me_actor_ids[actor_rank], message)
-            return
         elif selection == "all":
             # replace me with actual remote actor mesh
             call_shape = Shape(
@@ -215,6 +214,16 @@ def cast(
                     call_shape,
                     message,
                 )
+        elif isinstance(selection, int):
+            try:
+                self._mailbox.post(
+                    self._please_replace_me_actor_ids[selection], message
+                )
+            except IndexError:
+                raise IndexError(
+                    f"Tried to send to an out-of-range rank {selection}: "
+                    f"mesh has {len(self._please_replace_me_actor_ids)} elements."
+                )
         else:
             raise ValueError(f"invalid selection: {selection}")
 

python/monarch/actor/__init__.py

@@ -18,6 +18,7 @@
     endpoint,
     MonarchContext,
     Point,
+    port,
     send,
     ValueMesh,
 )
@@ -44,6 +45,7 @@
     "Point",
     "proc_mesh",
     "ProcMesh",
+    "port",
     "send",
     "ValueMesh",
     "debug_client",