Auto-track and globally shut down all Forge actors and services

apps/grpo/main.py

@@ -516,23 +516,10 @@ async def continuous_training():
         training_task.cancel()
     finally:
         print("Shutting down...")
-
         # give mlogger time to shutdown backends, otherwise they can stay running.
         # TODO (felipemello) find more elegant solution
         await mlogger.shutdown.call_one()
         await asyncio.sleep(2)
-
-        await asyncio.gather(
-            DatasetActor.shutdown(dataloader),
-            policy.shutdown(),
-            RLTrainer.shutdown(trainer),
-            ReplayBuffer.shutdown(replay_buffer),
-            ComputeAdvantages.shutdown(compute_advantages),
-            ref_model.shutdown(),
-            reward_actor.shutdown(),
-        )
-        # TODO - add a global shutdown that implicitly shuts down all services
-        # and remote allocations
         await shutdown()
 
 

src/forge/actors/policy.py

@@ -157,7 +157,7 @@ async def launch(  # pyright: ignore[reportIncompatibleMethodOverride]
 
     @classmethod
     async def shutdown(  # pyright: ignore[reportIncompatibleMethodOverride]
-        cls: type["Policy"], actor: "Policy"
+        cls: type["Policy"], actor: "Policy", quiet: bool = False
     ):
         assert (
             actor._policy_proc is not None

src/forge/controller/actor.py

@@ -12,7 +12,12 @@
 
 from monarch.actor import Actor, current_rank, current_size, endpoint
 
-from forge.controller.provisioner import get_proc_mesh, stop_proc_mesh
+from forge.controller.provisioner import (
+    get_proc_mesh,
+    register_actor,
+    register_service,
+    stop_proc_mesh,
+)
 
 from forge.types import ProcessConfig, ServiceConfig
 
@@ -81,11 +86,11 @@ def options(
 
         # Pre-configure a single actor
         actor = await MyForgeActor.options(procs=1, hosts=1).as_actor(...)
-        await actor.shutdown()
+        await MyForgeActor.shutdown(actor)
 
         # Default usage without calling options
         actor = await MyForgeActor.as_actor(...)
-        await actor.shutdown()
+        await MyForgeActor.shutdown(actor)
         """
 
         attrs = {
@@ -127,7 +132,10 @@ async def as_service(
         logger.info("Spawning Service for %s", cls.__name__)
         service = Service(cfg, cls, actor_args, actor_kwargs)
         await service.__initialize__()
-        return ServiceInterface(service, cls)
+        service_interface = ServiceInterface(service, cls)
+        # Register this service with the provisioner so it can cleanly shut this down
+        await register_service(service_interface)
+        return service_interface
 
     @endpoint
     async def setup(self):
@@ -145,7 +153,7 @@ async def setup(self):
         pass
 
     @classmethod
-    async def launch(cls, *args, **kwargs) -> "ForgeActor":
+    async def launch(cls, *args, **kwargs) -> "ActorMesh":
         """Provisions and deploys a new actor.
 
         This method is used by `Service` to provision a new replica.
@@ -185,6 +193,8 @@ async def as_actor(cls: Type[T], *args, **actor_kwargs) -> T:
         """
         logger.info("Spawning single actor %s", cls.__name__)
         actor = await cls.launch(*args, **actor_kwargs)
+        # Register this actor with the provisioner so it can cleanly shut this down
+        await register_actor(actor)
         return actor
 
     @classmethod

src/forge/controller/provisioner.py

@@ -13,8 +13,10 @@
 import socket
 import uuid
 
+from monarch._src.actor.actor_mesh import ActorMesh
+
 from monarch._src.actor.shape import Extent, NDSlice, Shape
-from monarch.actor import Actor, endpoint, ProcMesh
+from monarch.actor import Actor, endpoint
 
 from monarch.tools import commands
 
@@ -141,6 +143,9 @@ def __init__(self, cfg: ProvisionerConfig | None = None):
         if not self.launcher:
             logger.warning("Launcher not provided, remote allocations will not work.")
 
+        self._registered_actors: list["ForgeActor"] = []
+        self._registered_services: list["ServiceInterface"] = []
+
     async def initialize(self):
         """Call this after creating the instance"""
         if self.launcher is not None:
@@ -359,8 +364,55 @@ async def stop_proc_mesh(self, proc_mesh: ProcMesh):
                 commands.kill(server_name)
             del self._proc_host_map[proc_mesh]
 
+    def register_service(self, service: "ServiceInterface") -> None:
+        """Registers a service allocation for cleanup."""
+        # Import ServiceInterface here instead of at top-level to avoid circular import
+        from forge.controller.service import ServiceInterface
+
+        if not isinstance(service, ServiceInterface):
+            raise TypeError(
+                f"register_service expected ServiceInterface, got {type(service)}"
+            )
+
+        self._registered_services.append(service)
+
+    def register_actor(self, actor: "ForgeActor") -> None:
+        """Registers a single actor allocation for cleanup."""
+
+        if not isinstance(actor, ActorMesh):
+            raise TypeError(f"register_actor expected ActorMesh, got {type(actor)}")
+
+        self._registered_actors.append(actor)
+
+    async def shutdown_all_allocations(self):
+        """Gracefully shut down all tracked actors and services."""
+        logger.info(
+            f"Shutting down {len(self._registered_services)} service(s) and {len(self._registered_actors)} actor(s)..."
+        )
+        # --- ServiceInterface ---
+        for service in reversed(self._registered_services):
+            try:
+                await service.shutdown()
+
+            except Exception as e:
+                logger.warning(f"Failed to shut down {service}: {e}")
+
+        # --- Actor instance (ForgeActor or underlying ActorMesh) ---
+        for actor in reversed(self._registered_actors):
+            try:
+                # Get the class to call shutdown on (ForgeActor or its bound class)
+                actor_cls = getattr(actor, "_class", None) or actor.__class__
+                await actor_cls.shutdown(actor)
+
+            except Exception as e:
+                logger.warning(f"Failed to shut down {actor}: {e}")
+
+        self._registered_actors.clear()
+        self._registered_services.clear()
+
     async def shutdown(self):
         """Tears down all remaining remote allocations."""
+        await self.shutdown_all_allocations()
         async with self._lock:
             for server_name in self._server_names:
                 commands.kill(server_name)
@@ -429,6 +481,18 @@ async def host_mesh_from_proc(proc_mesh: ProcMesh):
     return await provisioner.host_mesh_from_proc(proc_mesh)
 
 
+async def register_service(service: "ServiceInterface") -> None:
+    """Registers a service allocation with the global provisioner."""
+    provisioner = await _get_provisioner()
+    provisioner.register_service(service)
+
+
+async def register_actor(actor: "ForgeActor") -> None:
+    """Registers an actor allocation with the global provisioner."""
+    provisioner = await _get_provisioner()
+    provisioner.register_actor(actor)
+
+
 async def stop_proc_mesh(proc_mesh: ProcMesh):
     provisioner = await _get_provisioner()
     return await provisioner.stop_proc_mesh(proc_mesh=proc_mesh)

tests/sandbox/rl_trainer/main.py

@@ -233,7 +233,6 @@ async def continuous_training():
         print("Training interrupted by user")
     finally:
         print("Shutting down trainer...")
-        await RLTrainer.shutdown(trainer)
         await mlogger.shutdown.call_one()
         await shutdown()
         print("Trainer shutdown complete.")

tests/sandbox/toy_rl/sumdigits.py

@@ -533,15 +533,6 @@ async def continuous_training():
         training_task.cancel()
     finally:
         print("Shutting down...")
-        await asyncio.gather(
-            DatasetActor.shutdown(dataloader),
-            policy.shutdown(),
-            Trainer.shutdown(trainer),
-            ReplayBuffer.shutdown(replay_buffer),
-            reward_actor.shutdown(),
-        )
-        # TODO - add a global shutdown that implicitly shuts down all services
-        # and remote allocations
         await shutdown()
 
 

tests/sandbox/toy_rl/toy_metrics/main.py

@@ -112,12 +112,6 @@ async def main():
     # shutdown
     await mlogger.shutdown.call_one()
     await asyncio.sleep(2)
-
-    await asyncio.gather(
-        trainer.shutdown(),
-        generator.shutdown(),
-    )
-
     await shutdown()
 
 

tests/sandbox/vllm/main.py

@@ -66,7 +66,6 @@ async def run(cfg: DictConfig):
         print("-" * 80)
 
     print("\nShutting down...")
-    await policy.shutdown()
     await shutdown()