add sessionloader and leastloadedloader and tests

Author: DNXie

src/forge/controller/service/__init__.py

@@ -6,17 +6,22 @@
 
 from .interface import ServiceInterface, Session, SessionContext
 from .metrics import ServiceMetrics
-from .replica import Replica, ReplicaMetrics
+from .replica import Replica, ReplicaMetrics, ReplicaState
+from .router import LeastLoadedRouter, RoundRobinRouter, SessionRouter
 from .service import Service, ServiceActor, ServiceConfig
 
 __all__ = [
     "Replica",
     "ReplicaMetrics",
+    "ReplicaState",
     "Service",
     "ServiceConfig",
     "ServiceInterface",
     "ServiceMetrics",
     "Session",
     "SessionContext",
     "ServiceActor",
+    "LeastLoadedRouter",
+    "RoundRobinRouter",
+    "SessionRouter",
 ]

src/forge/controller/service/interface.py

@@ -14,7 +14,7 @@
 import logging
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
-from typing import Generic, List, ParamSpec, TypeVar
+from typing import Dict, Generic, List, ParamSpec, TypeVar
 
 from monarch._src.actor.endpoint import EndpointProperty
 
@@ -287,7 +287,10 @@ class Router(ABC):
 
     @abstractmethod
     def get_replica(
-        self, replicas: List[Replica], sess_id: str | None = None
+        self,
+        replicas: List[Replica],
+        sess_id: str | None = None,
+        session_map: Dict[str, int] | None = None,
     ) -> Replica:
         """Select a replica from the list based on routing logic."""
         pass

src/forge/controller/service/router.py

@@ -4,11 +4,15 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-from typing import List
+import logging
+from typing import Dict, List
 
 from .interface import Router
 from .replica import Replica
 
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.DEBUG)
+
 
 class RoundRobinRouter(Router):
     """Round-robin router for stateless requests."""
@@ -17,7 +21,10 @@ def __init__(self):
         self._next_idx = 0
 
     def get_replica(
-        self, replicas: List[Replica], sess_id: str | None = None
+        self,
+        replicas: List[Replica],
+        sess_id: str | None = None,
+        session_map: Dict[str, int] | None = None,
     ) -> Replica:
         healthy_replicas = [r for r in replicas if r.healthy]
         if not healthy_replicas:
@@ -27,3 +34,57 @@ def get_replica(
         replica = healthy_replicas[self._next_idx]
 
         return replica
+
+
+class LeastLoadedRouter(Router):
+    """Always routes to the replica with the lowest current load."""
+
+    def get_replica(
+        self,
+        replicas: List["Replica"],
+        sess_id: str | None = None,
+        session_map: Dict[str, int] | None = None,
+    ) -> "Replica":
+        healthy_replicas = [r for r in replicas if r.healthy]
+        if not healthy_replicas:
+            raise RuntimeError("No healthy replicas available for session assignment")
+        return min(healthy_replicas, key=lambda r: r.current_load)
+
+
+class SessionRouter(Router):
+    """Session-based routing: sticky sessions with a fallback router."""
+
+    def __init__(self, fallback_router: Router):
+        self.fallback_router = fallback_router
+
+    def get_replica(
+        self,
+        replicas: List["Replica"],
+        sess_id: str | None = None,
+        session_map: Dict[str, int] | None = None,
+    ) -> "Replica":
+        if sess_id is None:
+            raise ValueError("SessionRouter requires a session ID")
+
+        if session_map is None:
+            raise ValueError("Session map must be provided for SessionRouter")
+
+        # Check if session already has a replica
+        if sess_id in session_map:
+            replica_idx = session_map[sess_id]
+            # Find the replica with this index
+            for r in replicas:
+                if r.idx == replica_idx and r.healthy:
+                    return r
+            # If the replica is no longer healthy, remove from session map and reassign
+            del session_map[sess_id]
+
+        # Use fallback router to assign a new replica
+        replica = self.fallback_router.get_replica(replicas, sess_id, session_map)
+        session_map[sess_id] = replica.idx
+        logger.debug(
+            "Assigning session %s to replica %d",
+            sess_id,
+            replica.idx,
+        )
+        return replica

src/forge/controller/service/service.py

@@ -36,16 +36,20 @@
 import logging
 import pprint
 import uuid
-from typing import Dict, List
+from typing import Dict, List, Type
 
 from monarch.actor import Actor, endpoint
 
-from forge.controller.service.interface import _session_context, Session
+from forge.controller.service.interface import _session_context, Router, Session
 
 from forge.controller.service.metrics import ServiceMetrics
 from forge.controller.service.replica import Replica, ServiceRequest
 
-from forge.controller.service.router import RoundRobinRouter
+from forge.controller.service.router import (
+    LeastLoadedRouter,
+    RoundRobinRouter,
+    SessionRouter,
+)
 from forge.types import ServiceConfig
 
 logger = logging.getLogger(__name__)
@@ -64,6 +68,13 @@ class Service:
         actor_def: Actor class definition to instantiate on each replica
         *actor_args: Positional arguments passed to actor constructor
         **actor_kwargs: Keyword arguments passed to actor constructor
+        router_cls (Type[Router], optional): Router class used for non-session
+            calls. Defaults to RoundRobinRouter. Examples include RoundRobinRouter
+            or LeastLoadedRouter. The router is instantiated internally.
+        fallback_router_cls: Router class used as a fallback when a session
+                             cannot be mapped to an existing replica. Defaults
+                             to LeastLoadedRouter.
+
 
     Attributes:
         _cfg: Service configuration
@@ -73,16 +84,24 @@ class Service:
         _endpoints: Dynamically registered actor endpoints
     """
 
-    def __init__(self, cfg: ServiceConfig, actor_def, actor_kwargs: dict):
+    def __init__(
+        self,
+        cfg: ServiceConfig,
+        actor_def,
+        actor_kwargs: dict,
+        router_cls: Type["Router"] = RoundRobinRouter,
+        fallback_router_cls: Type["Router"] = LeastLoadedRouter,
+    ):
         self._cfg = cfg
         self._replicas = []
         self._actor_def = actor_def
         self._actor_kwargs = actor_kwargs
+        self.router_cls = router_cls
+        self.fallback_router_cls = fallback_router_cls
 
         self._active_sessions = []
         self._id_session_map = {}
         self._session_replica_map: Dict[str, int] = {}
-        self._router = RoundRobinRouter()
 
         # Initialize metrics collection
         self._metrics = ServiceMetrics()
@@ -95,6 +114,12 @@ def __init__(self, cfg: ServiceConfig, actor_def, actor_kwargs: dict):
     async def __initialize__(self):
         """Initializes the service and starts the health loop."""
         logger.debug(f"Starting service up with {self._cfg.num_replicas} replicas.")
+
+        # Initialize the routers
+        self._default_router = self.router_cls()
+        self._session_router = SessionRouter(fallback_router=self.fallback_router_cls())
+
+        # Initialize all replicas
         replicas = []
         num_replicas = self._cfg.num_replicas
         for i in range(num_replicas):
@@ -457,36 +482,15 @@ async def _health_loop(self, poll_rate_s: float):
 
             await asyncio.sleep(poll_rate_s)
 
-    def _get_least_loaded_replica(self) -> "Replica":
-        """Get the replica with the lowest load."""
-        healthy_replicas = [r for r in self._replicas if r.healthy]
-        if not healthy_replicas:
-            raise RuntimeError("No healthy replicas available for session assignment")
-
-        # Use the replica's current_load property
-        return min(healthy_replicas, key=lambda replica: replica.current_load)
-
     async def _get_replica(self, sess_id: str | None) -> "Replica":
         """Get a replica for the given session ID."""
         if sess_id is None:
             # No session, use the default router
-            return self._router.get_replica(self._replicas)
-
-        # Session-based routing
-        if sess_id in self._session_replica_map:
-            replica_idx = self._session_replica_map[sess_id]
-            # Find the replica with this index
-            for replica in self._replicas:
-                if replica.idx == replica_idx and replica.healthy:
-                    return replica
-            # If the replica is no longer healthy, remove from session map and reassign
-            del self._session_replica_map[sess_id]
+            return self._default_router.get_replica(self._replicas)
 
-        # New session, assign to least loaded replica
-        replica = self._get_least_loaded_replica()
-        self._session_replica_map[sess_id] = replica.idx
-        logger.debug("Assigning session %s to replica %d", sess_id, replica.idx)
-        return replica
+        return self._session_router.get_replica(
+            self._replicas, sess_id, self._session_replica_map
+        )
 
     async def stop(self):
         logger.debug("Stopping service...")