Refactor service spawning: add ForgeActor.options().as_service() API

apps/vllm/main.py

@@ -13,7 +13,6 @@
 
 from forge.actors.policy import Policy
 from forge.cli.config import parse
-from forge.controller.service import ServiceConfig, shutdown_service, spawn_service
 
 from omegaconf import DictConfig
 from src.forge.data.utils import exclude_service
@@ -28,10 +27,8 @@ async def run(cfg: DictConfig):
 
     print("Spawning service...")
 
-    policy = await spawn_service(
-        ServiceConfig(**cfg.policy.service),
-        Policy,
-        **exclude_service(cfg.policy),
+    policy = await Policy.options(**cfg.policy.service).as_service(
+        **exclude_service(cfg.policy)
     )
 
     async with policy.session():
@@ -48,7 +45,7 @@ async def run(cfg: DictConfig):
 
         print("\nShutting down...")
 
-    await shutdown_service(policy)
+    await policy.shutdown()
 
 
 @parse

src/forge/controller/actor.py

@@ -8,11 +8,13 @@
 
 import math
 import sys
+from typing import Type
 
 from monarch.actor import Actor, current_rank, current_size, endpoint
 
 from forge.controller.proc_mesh import get_proc_mesh, stop_proc_mesh
-from forge.types import ProcessConfig
+
+from forge.types import ProcessConfig, ServiceConfig
 
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.DEBUG)
@@ -41,6 +43,79 @@ def __init__(self, *args, **kwargs):
         self.logger.root.addHandler(stdout_handler)
         super().__init__(*args, **kwargs)
 
+    @classmethod
+    def options(
+        cls, *, num_replicas: int, procs_per_replica: int, **service_kwargs
+    ) -> Type["ConfiguredService"]:
+        """
+        Returns a ConfiguredService class that wraps this ForgeActor in a Service.
+
+        Usage:
+            service = await MyForgeActor.options(num_replicas=1, procs_per_replica=2).as_service(...)
+            await service.shutdown()
+        """
+        from forge.controller.service import Service, ServiceInterface
+
+        cfg = ServiceConfig(
+            num_replicas=num_replicas,
+            procs_per_replica=procs_per_replica,
+            **service_kwargs,
+        )
+
+        class ConfiguredService:
+            """
+            A wrapper around Service that binds a ForgeActor class.
+            Provides:
+              - as_service(): spawns the actor inside the service
+              - shutdown(): stops the service
+            """
+
+            _actor_def = cls
+            _service_interface: ServiceInterface | None
+
+            def __init__(self) -> None:
+                self._service_interface = None
+
+            @classmethod
+            async def as_service(cls, **actor_kwargs) -> "ConfiguredService":
+                """
+                Spawn the actor inside a Service with the given configuration.
+
+                Args:
+                    **actor_kwargs: arguments to pass to the ForgeActor constructor
+
+                Returns:
+                    self: so that methods like .shutdown() can be called
+                """
+                self = cls()
+                logger.info("Spawning Service Actor for %s", self._actor_def.__name__)
+                service = Service(cfg, self._actor_def, actor_kwargs)
+                await service.__initialize__()
+                self._service_interface = ServiceInterface(service, self._actor_def)
+                return self
+
+            async def shutdown(self):
+                """
+                Gracefully stops the service if it has been started.
+                """
+                if self._service_interface is None:
+                    raise RuntimeError("Service not started yet")
+                await self._service_interface._service.stop()
+                self._service_interface = None
+
+            def __getattr__(self, item):
+                """
+                Delegate attribute access to the ServiceInterface instance.
+                This makes ConfiguredService behave like a ServiceInterface.
+                """
+                if self._service_interface is None:
+                    raise AttributeError(
+                        f"Service not started yet; cannot access '{item}'"
+                    )
+                return getattr(self._service_interface, item)
+
+        return ConfiguredService
+
     @endpoint
     async def setup(self):
         """Sets up the actor.