Enable get_or_spawn_controller from inside an actor endpoint (#1347)

Summary:
Pull Request resolved: #1347

This diff enables get_or_spawn_controller API being called from inside an actor endpoint, and adds unit tests for this functionality.
ghstack-source-id: 313701938
exported-using-ghexport

Reviewed By: samlurye

Differential Revision: D83109145

fbshipit-source-id: 19a8084982060877f827bd73ee8b39adacd5af77

Author: shuhuayu

python/monarch/_src/actor/proc_mesh.py

@@ -343,7 +343,12 @@ async def task() -> HyProcMesh:
                 )._class
                 is _ControllerController
             ), "Expected v0 _ControllerController, got v1 _ControllerController"
-            pm._controller_controller = instance._controller_controller  # type: ignore
+            if instance._controller_controller is None:
+                pm._controller_controller = _get_controller_controller()[1]
+            else:
+                pm._controller_controller = cast(
+                    _ControllerController, instance._controller_controller
+                )
             instance._add_child(pm)
 
         async def task(

python/tests/test_python_actors.py

@@ -48,7 +48,11 @@
     fake_in_process_host,
     HostMesh,
 )
-from monarch._src.actor.proc_mesh import _get_bootstrap_args, ProcMesh
+from monarch._src.actor.proc_mesh import (
+    _get_bootstrap_args,
+    get_or_spawn_controller,
+    ProcMesh,
+)
 from monarch._src.actor.v1.host_mesh import (
     _bootstrap_cmd,
     fake_in_process_host as fake_in_process_host_v1,
@@ -1640,3 +1644,25 @@ def test_cuda_is_not_initialized_in_a_new_proc():
         pytest.skip("cannot find cuda")
     proc = this_host().spawn_procs().spawn("is_init", IsInit)
     assert not proc.is_cuda_initialized.call_one().get()
+
+
+class SpawningActorFromEndpointActor(Actor):
+    def __init__(self, root="None"):
+        self._root = root
+
+    @endpoint
+    def return_root(self):
+        return self._root
+
+    @endpoint
+    async def spawning_from_endpoint(self, name, root) -> None:
+        await get_or_spawn_controller(name, SpawningActorFromEndpointActor, root=root)
+
+
+@pytest.mark.timeout(60)
+def test_get_or_spawn_controller_inside_actor_endpoint():
+    actor_1 = get_or_spawn_controller("actor_1", SpawningActorFromEndpointActor).get()
+    actor_1.spawning_from_endpoint.call_one("actor_2", root="actor_1").get()
+    actor_2 = get_or_spawn_controller("actor_2", SpawningActorFromEndpointActor).get()
+    # verify that actor_2 was spawned from actor_1 with the correct root
+    assert actor_2.return_root.call_one().get() == "actor_1"