update custom route

Author: Yunnglin

cookbook/client/tinker/megatron/server_config_7b.yaml

@@ -22,7 +22,7 @@ applications:
     import_path: server            # Python module to import
     args:
       server_config:
-        per_token_model_limit: 1      # Maximum number of models (adapters) per token (server-globally enforced)
+        per_token_model_limit: 3      # Maximum number of models (adapters) per token (server-globally enforced)
       supported_models:
         - Qwen/Qwen2.5-7B-Instruct
     deployments:
@@ -58,11 +58,12 @@ applications:
       adapter_config:
         adapter_timeout: 30                        # Seconds before idle adapter unload
         adapter_max_lifetime: 36000                # Maximum lifetime of an adapter in seconds (e.g., 10 hours)
+      max_loras: 1                                 # Maximum number of LoRA adapters per model
     deployments:
       - name: ModelManagement
         autoscaling_config:
-          min_replicas: 1
-          max_replicas: 1
+          min_replicas: 2
+          max_replicas: 2
           target_ongoing_requests: 16
         ray_actor_options:
           num_cpus: 0.1

src/twinkle/model/multi_lora.py

@@ -115,7 +115,7 @@ def acquire_lora(self, tenant_adapter_name: str, config: LoraConfig) -> str:
             raise ValueError(f'Lora {tenant_adapter_name} already exists')
         _available_lora = self._get_available_lora()
         if _available_lora is None:
-            raise RuntimeError(f'No lora available for tenant {tenant_adapter_name}')
+            raise RuntimeError(f'No lora available for tenant {tenant_adapter_name}. Max loras: {self.max_loras}')
         if config.r > self.max_r:
             raise RuntimeError(f'Too big rank for lora: {config.r}')
         _available_lora.tenant_config = config

src/twinkle/server/tinker/common/router.py

@@ -0,0 +1,73 @@
+from ray.serve.request_router import (FIFOMixin, MultiplexMixin, PendingRequest, ReplicaID, ReplicaResult,
+                                      RequestRouter, RunningReplica)
+from typing import Dict, List, Optional
+
+from twinkle.server.utils.state import ServerStateProxy, get_server_state
+
+
+class StickyLoraRequestRouter(FIFOMixin, MultiplexMixin, RequestRouter):
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        self.state: ServerStateProxy = get_server_state()
+
+    async def choose_replicas(
+        self,
+        candidate_replicas: List[RunningReplica],
+        pending_request: Optional[PendingRequest] = None,
+    ) -> List[List[RunningReplica]]:
+        """
+        This method chooses the best replica for the request based on
+        multiplexed and avaliable lora count. The algorithm
+        works as follows:
+
+        1. Populate top_ranked_replicas based on available replicas based on
+          multiplex_id (only one replica is chosen)
+        2. Populate and override top_ranked_replicas info based on avalible lora
+          slots of the replica.
+        """
+
+        # Take the best set of replicas for the multiplexed model
+        if (pending_request is not None and pending_request.metadata.multiplexed_model_id):
+            ranked_replicas_multiplex: List[RunningReplica] = (self.rank_replicas_via_multiplex(
+                replicas=candidate_replicas,
+                multiplexed_model_id=pending_request.metadata.multiplexed_model_id,
+            ))[0]
+
+            # If found any replica, return it
+            if ranked_replicas_multiplex:
+                print('[Router] Found replica for multiplexed model !!!')
+                return [ranked_replicas_multiplex]
+
+        # Dictionary to hold the top-ranked replicas
+        top_ranked_replicas: Dict[ReplicaID, RunningReplica] = {}
+
+        # Filter out replicas that are not available (queue length exceed max ongoing request)
+        ranked_replicas_locality = self.select_available_replicas(candidates=candidate_replicas)
+
+        for replica in ranked_replicas_locality:
+            top_ranked_replicas[replica.replica_id] = replica
+
+        # Filter out replicas that exceed max lora count (query from server state)
+        candidate_ids = [r.replica_id.unique_id for r in top_ranked_replicas.values()]
+        available_ids = set(self.state.get_available_replica_ids(candidate_ids))
+        if available_ids:
+            top_ranked_replicas = {
+                rid: r
+                for rid, r in top_ranked_replicas.items() if r.replica_id.unique_id in available_ids
+            }
+
+        if not top_ranked_replicas:
+            # No replica has remaining LoRA capacity – fall back to all candidates
+            print('[Router] No replica has remaining LoRA capacity')
+            return [candidate_replicas]
+
+        print('[Router] StickyLoraRequestRouter choosing replica for request')
+
+        # Take the replica with minimum throughput.
+        min_throughput_replicas = min(
+            [replica for replica in top_ranked_replicas.values()],
+            key=lambda r: r.routing_stats.get('throughput', 0),
+        )
+        return [[min_throughput_replicas]]

src/twinkle/server/tinker/model.py

@@ -13,6 +13,7 @@
 from fastapi import FastAPI, Request
 from peft import LoraConfig
 from ray import serve
+from ray.serve.config import RequestRouterConfig
 from tinker import types
 from typing import Any, Dict, Optional
 
@@ -21,9 +22,10 @@
 from twinkle.server.utils.adapter_manager import AdapterManagerMixin
 from twinkle.server.utils.state import ServerStateProxy, get_server_state
 from twinkle.server.utils.task_queue import TaskQueueConfig, TaskQueueMixin
-from twinkle.server.utils.validation import verify_request_token
+from twinkle.server.utils.validation import get_token_from_request, verify_request_token
 from twinkle.utils.logger import get_logger
 from .common.io_utils import create_checkpoint_manager, create_training_run_manager
+from .common.router import StickyLoraRequestRouter
 
 logger = get_logger()
 
@@ -62,7 +64,10 @@ async def verify_token(request: Request, call_next):
         """Middleware to verify authentication token for all requests."""
         return await verify_request_token(request=request, call_next=call_next)
 
-    @serve.deployment(name='ModelManagement')
+    @serve.deployment(
+        name='ModelManagement',
+        request_router_config=RequestRouterConfig(request_router_class=StickyLoraRequestRouter, ),
+    )
     @serve.ingress(app)
     class ModelManagement(TaskQueueMixin, AdapterManagerMixin):
         """Model management service handling training operations.
@@ -99,28 +104,31 @@ def __init__(self,
             else:
                 self.device_mesh = DeviceMesh.from_sizes(**device_mesh)
             self.use_megatron = use_megatron
-            replica_context = serve.get_replica_context()
-            replica_id = replica_context.replica_id.unique_id
+            self.replica_id = serve.get_replica_context().replica_id.unique_id
+            self.max_loras = kwargs.get('max_loras', 5)
             # Initialize model immediately - choose backend based on use_megatron
             if use_megatron:
                 from .common.megatron_model import TwinkleCompatMegatronModel
                 self.model = TwinkleCompatMegatronModel(
                     model_id=model_id,
                     device_mesh=self.device_mesh,
                     remote_group=self.device_group.name,
-                    instance_id=replica_id,
+                    instance_id=self.replica_id,
                     **kwargs)
             else:
                 from .common.transformers_model import TwinkleCompatTransformersModel
                 self.model = TwinkleCompatTransformersModel(
                     model_id=model_id,
                     device_mesh=self.device_mesh,
                     remote_group=self.device_group.name,
-                    instance_id=replica_id,
+                    instance_id=self.replica_id,
                     **kwargs)
             self.base_model = model_id
             self.state: ServerStateProxy = get_server_state()
 
+            # Register this replica so the router can track capacity
+            self.state.register_replica(self.replica_id, self.max_loras)
+
             # Initialize task queue
             self._init_task_queue(TaskQueueConfig.from_dict(queue_config))
 
@@ -136,9 +144,18 @@ def __init__(self,
             4. Direct call actor instead of http or handler in server.py
         """
 
-        # @serve.multiplexed(max_num_models_per_replica=kwargs.get('max_loras', 5))
-        # async def get_multiplexed_adapter(self, request_id: str):
-        # return request_id
+        @serve.multiplexed(max_num_models_per_replica=kwargs.get('max_loras', 5))
+        async def _sticky_entry(self, sticky_key: str):
+            return sticky_key
+
+        async def _ensure_sticky(self):
+            sticky_key = serve.get_multiplexed_model_id()
+            await self._sticky_entry(sticky_key)
+
+        async def _on_request_start(self, request: Request) -> str:
+            await self._ensure_sticky()
+            token = get_token_from_request(request)
+            return token
 
         def _cleanup_adapter(self, adapter_name: str) -> None:
             """Common adapter cleanup logic used by both manual unload and automatic expiration.
@@ -188,12 +205,13 @@ async def create_model(self, request: Request, body: types.CreateModelRequest) -
             Returns:
                 UntypedAPIFuture wrapping CreateModelResponse with model_id
             """
+            token = await self._on_request_start(request)
 
             async def _create_adapter():
                 model_id = None
                 try:
                     # Register a new model_id for each create_model call
-                    model_id = self.state.register_model(body.model_dump(), token=request.state.token)
+                    model_id = self.state.register_model(body.model_dump(), token=token, replica_id=self.replica_id)
 
                     # Create a new LoRA adapter for the model
                     if body.lora_config:
@@ -203,7 +221,7 @@ async def _create_adapter():
                         adapter_name = self.get_adapter_name(adapter_name=model_id)
 
                         # Register adapter FIRST
-                        self.register_adapter(adapter_name, request.state.token, session_id=body.session_id)
+                        self.register_adapter(adapter_name, token, session_id=body.session_id)
 
                         # Create adapter AFTER successful registration
                         self.model.add_adapter_to_model(adapter_name=adapter_name, config_or_dir=lora_cfg)
@@ -215,7 +233,7 @@ async def _create_adapter():
                         # Fresh adapter has no accumulated gradients.
                         self.set_adapter_state(adapter_name, 'grad_ready', False)
 
-                    training_run_manager = create_training_run_manager(request.state.token)
+                    training_run_manager = create_training_run_manager(token)
                     training_run_manager.save(model_id, body)
 
                     return types.CreateModelResponse(model_id=model_id)
@@ -233,7 +251,7 @@ async def _create_adapter():
 
             return await self.schedule_task(
                 _create_adapter,
-                token=request.state.token,
+                token=token,
                 task_type='create_model',
             )
 
@@ -248,9 +266,10 @@ async def get_info(self, request: Request, body: types.GetInfoRequest) -> types.
             Returns:
                 GetInfoResponse with model metadata (name, lora_rank, etc.)
             """
+            token = await self._on_request_start(request)
             # Note: get_info doesn't require token for reading metadata in tinker
             # Using a default token or None since this is read-only
-            training_run_manager = create_training_run_manager(request.state.token)
+            training_run_manager = create_training_run_manager(token)
             metadata = training_run_manager.get(str(body.model_id))
             model_name = metadata.base_model if metadata else model_id
             lora_rank = None
@@ -279,6 +298,7 @@ async def unload_model(self, request: Request, body: types.UnloadModelRequest) -
             Returns:
                 UntypedAPIFuture wrapping UnloadModelResponse
             """
+            token = await self._on_request_start(request)
 
             async def _do_unload():
                 # Only remove adapter, not the base model
@@ -290,7 +310,7 @@ async def _do_unload():
             return await self.schedule_task(
                 _do_unload,
                 model_id=body.model_id,
-                token=request.state.token,
+                token=token,
                 task_type='unload_model',
             )
 
@@ -307,6 +327,7 @@ async def forward(self, request: Request, body: types.ForwardRequest) -> types.U
             Returns:
                 UntypedAPIFuture wrapping ForwardBackwardOutput with loss
             """
+            token = await self._on_request_start(request)
 
             async def _do_forward():
                 try:
@@ -340,7 +361,7 @@ async def _do_forward():
             return await self.schedule_task(
                 _do_forward,
                 model_id=body.model_id,
-                token=request.state.token,
+                token=token,
                 input_tokens=input_tokens,
                 batch_size=batch_size,
                 data_world_size=self.device_mesh.data_world_size,
@@ -364,6 +385,7 @@ async def forward_backward(self, request: Request,
             Returns:
                 UntypedAPIFuture wrapping ForwardBackwardOutput with loss and metrics
             """
+            token = await self._on_request_start(request)
 
             async def _do_forward_backward():
                 try:
@@ -405,7 +427,7 @@ async def _do_forward_backward():
             return await self.schedule_task(
                 _do_forward_backward,
                 model_id=body.model_id,
-                token=request.state.token,
+                token=token,
                 input_tokens=input_tokens,
                 batch_size=batch_size,
                 data_world_size=self.device_mesh.data_world_size,
@@ -425,6 +447,7 @@ async def optim_step(self, request: Request, body: types.OptimStepRequest) -> ty
             Returns:
                 UntypedAPIFuture wrapping OptimStepResponse
             """
+            token = await self._on_request_start(request)
 
             async def _do_optim():
                 try:
@@ -455,7 +478,7 @@ async def _do_optim():
             return await self.schedule_task(
                 _do_optim,
                 model_id=body.model_id,
-                token=request.state.token,
+                token=token,
                 task_type='optim_step',
             )
 
@@ -473,6 +496,7 @@ async def save_weights(self, request: Request, body: types.SaveWeightsRequest) -
             Returns:
                 UntypedAPIFuture wrapping SaveWeightsResponse with saved path
             """
+            token = await self._on_request_start(request)
 
             async def _do_save():
                 try:
@@ -482,8 +506,6 @@ async def _do_save():
                     # Touch adapter to reset inactivity counter
                     self.touch_adapter(adapter_name)
 
-                    # Extract token from request for user isolation
-                    token = request.state.token
                     checkpoint_manager = create_checkpoint_manager(token)
 
                     # get save dir with token-based isolation
@@ -506,7 +528,7 @@ async def _do_save():
             return await self.schedule_task(
                 _do_save,
                 model_id=body.model_id,
-                token=request.state.token,
+                token=token,
                 task_type='save_weights',
             )
 
@@ -525,6 +547,7 @@ async def save_weights_for_sampler(self, request: Request,
             Returns:
                 UntypedAPIFuture wrapping SaveWeightsForSamplerResponseInternal
             """
+            token = await self._on_request_start(request)
 
             async def _do_save_for_sampler():
                 try:
@@ -535,8 +558,6 @@ async def _do_save_for_sampler():
                     # Touch adapter to reset inactivity counter
                     self.touch_adapter(adapter_name)
 
-                    # Extract token from request for user isolation
-                    token = request.state.token
                     checkpoint_manager = create_checkpoint_manager(token)
 
                     # get save dir with token-based isolation
@@ -571,7 +592,7 @@ async def _do_save_for_sampler():
             return await self.schedule_task(
                 _do_save_for_sampler,
                 model_id=body.model_id,
-                token=request.state.token,
+                token=token,
                 task_type='save_weights_for_sampler',
             )
 
@@ -589,6 +610,7 @@ async def load_weights(self, request: Request, body: types.LoadWeightsRequest) -
             Returns:
                 UntypedAPIFuture wrapping LoadWeightsResponse
             """
+            token = await self._on_request_start(request)
 
             async def _do_load():
                 try:
@@ -600,9 +622,6 @@ async def _do_load():
                     # Touch adapter to reset inactivity counter
                     self.touch_adapter(adapter_name)
 
-                    # Extract token from request for user isolation
-                    token = request.state.token
-
                     weight_path = body.path
                     load_optimizer = body.optimizer
 
@@ -625,7 +644,7 @@ async def _do_load():
             return await self.schedule_task(
                 _do_load,
                 model_id=body.model_id,
-                token=request.state.token,
+                token=token,
                 task_type='load_weights',
             )