fix processor

Author: Yunnglin

cookbook/client/server/transformer/server_config.yaml

@@ -24,7 +24,7 @@ applications:
         - Qwen/Qwen3.5-4B
     deployments:
       - name: TinkerCompatServer
-        max_ongoing_requests: 10
+        max_ongoing_requests: 50
         autoscaling_config:
           min_replicas: 1                # Minimum number of replicas
           max_replicas: 1                # Maximum number of replicas
@@ -107,14 +107,14 @@ applications:
     route_prefix: /api/v1/processor
     import_path: processor
     args:
-      ncpu_proc_per_node: 1           # 每节点 CPU 进程数
+      ncpu_proc_per_node: 2           # 每节点 CPU 进程数
       device_group:
         name: model
-        ranks: 1
+        ranks: 2
         device_type: CPU
       device_mesh:
         device_type: CPU
-        dp_size: 1                    # 数据并行大小
+        dp_size: 2                    # 数据并行大小
     deployments:
       - name: ProcessorManagement
         autoscaling_config:

setup.cfg

@@ -22,7 +22,7 @@ ignore-words-list = patten,nd,ty,mot,hist,formating,winn,gool,datas,wan,confids
 [flake8]
 max-line-length = 120
 select = B,E,F,P,T4,W,B9
-ignore = F401,F403,F405,F821,W503,E251,W504,E126
+ignore = F401,F403,F405,F821,W503,E251,W504,E126,E125
 exclude = docs/src,*.pyi,.git,peft.py
 
 [darglint]

src/twinkle/hub/hub.py

@@ -374,7 +374,7 @@ def push_to_hub(cls,
             ignore_patterns = []
         if revision is None or revision == 'main':
             revision = 'master'
-        return push_to_hub(
+        result = push_to_hub(
             repo_id,
             folder_path,
             token or cls.ms_token,
@@ -383,6 +383,8 @@ def push_to_hub(cls,
             ignore_file_pattern=ignore_patterns,
             revision=revision,
             tag=path_in_repo)
+        if not result:
+            raise Exception('Failed to push to hub')
 
     @classmethod
     def load_dataset(cls,

src/twinkle/server/common/datum.py

@@ -1,5 +1,4 @@
 # Copyright (c) ModelScope Contributors. All rights reserved.
-# Moved from tinker/common/datum.py — logic unchanged.
 from __future__ import annotations
 
 import numpy as np

src/twinkle/server/common/router.py

@@ -56,7 +56,7 @@ async def choose_replicas(
 
         # 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))
+        available_ids = set(await self.state.get_available_replica_ids(candidate_ids))
         if available_ids:
             top_ranked_replicas = {
                 rid: r

src/twinkle/server/gateway/tinker_gateway_handlers.py

@@ -61,7 +61,7 @@ async def create_session(
     async def session_heartbeat(
         request: Request, body: types.SessionHeartbeatRequest, self: GatewayServer = Depends(self_fn)
     ) -> types.SessionHeartbeatResponse:  # noqa: E125
-        alive = self.state.touch_session(body.session_id)
+        alive = await self.state.touch_session(body.session_id)
         if not alive:
             raise HTTPException(status_code=404, detail='Unknown session')
         return types.SessionHeartbeatResponse()
@@ -84,7 +84,7 @@ async def retrieve_future(request: Request,
         start = asyncio.get_event_loop().time()
 
         while True:
-            record = self.state.get_future(request_id)
+            record = await self.state.get_future(request_id)
 
             if record is None:
                 return {'type': 'try_again'}
@@ -103,7 +103,7 @@ async def retrieve_future(request: Request,
 
             await asyncio.sleep(poll_interval)
 
-        record = self.state.get_future(request_id)
+        record = await self.state.get_future(request_id)
         if not record:
             return {'type': 'try_again'}
 
@@ -207,7 +207,7 @@ async def publish_checkpoint(request: Request,
 
         checkpoint_name = checkpoint_id.split('/')[-1]
         hub_model_id = f'{username}/{run_id}_{checkpoint_name}'
-        HubOperation.async_push_to_hub(repo_id=hub_model_id, folder_path=checkpoint_dir, token=token, private=True)
+        HubOperation.push_to_hub(repo_id=hub_model_id, folder_path=checkpoint_dir, token=token, private=True)
 
         return Response(status_code=204)
 

src/twinkle/server/gateway/twinkle_gateway_handlers.py

@@ -43,7 +43,7 @@ async def session_heartbeat(
             body: types.SessionHeartbeatRequest,
             self: GatewayServer = Depends(self_fn),
     ) -> types.SessionHeartbeatResponse:
-        alive = self.state.touch_session(body.session_id)
+        alive = await self.state.touch_session(body.session_id)
         if not alive:
             raise HTTPException(status_code=404, detail='Unknown session')
         return types.SessionHeartbeatResponse()

src/twinkle/server/processor/app.py

@@ -1,155 +1,133 @@
 # Copyright (c) ModelScope Contributors. All rights reserved.
 """
-Processor management application (moved from twinkle/processor.py).
+Processor management application.
 
 Provides a Ray Serve deployment for managing distributed processors
 (datasets, dataloaders, preprocessors, rewards, templates, weight loaders, etc.).
+
+Follows the same structural pattern as model/app.py:
+- ProcessorManagement is a top-level class inheriting ProcessorManagerMixin
+- Routes are registered in build_processor_app() via _register_processor_routes()
+- serve.ingress(app)(ProcessorManagement) applied before deployment
+- Sticky session routing via @serve.multiplexed keyed on session ID
 """
-import importlib
+from __future__ import annotations
+
 import os
-import uuid
-from fastapi import FastAPI, HTTPException, Request
+from fastapi import FastAPI, Request
 from ray import serve
-from typing import Any, Dict
+from typing import Any, Dict, Optional
 
 import twinkle
-import twinkle_client.types as types
 from twinkle import DeviceGroup, DeviceMesh, get_logger
-from twinkle.server.common.serialize import deserialize_object
+from twinkle.server.utils.processor_manager import ProcessorManagerMixin
 from twinkle.server.utils.state import ServerStateProxy, get_server_state
 from twinkle.server.utils.validation import verify_request_token
+from .twinkle_handlers import _register_processor_routes
 
 logger = get_logger()
 
 
+class ProcessorManagement(ProcessorManagerMixin):
+    """Processor management service.
+
+    Manages lifecycle and invocation of distributed processor objects
+    (datasets, dataloaders, rewards, templates, etc.).
+
+    Lifecycle is handled by ProcessorManagerMixin:
+    - Processors are registered with a session ID on creation.
+    - A background thread expires processors whose session has timed out.
+    - Per-user processor limit is enforced at registration.
+    - Sticky session routing ensures session requests hit the same replica.
+    """
+
+    def __init__(self,
+                 ncpu_proc_per_node: int,
+                 device_group: dict[str, Any],
+                 device_mesh: dict[str, Any],
+                 nproc_per_node: int = 1,
+                 processor_config: dict[str, Any] | None = None):
+        self.device_group = DeviceGroup(**device_group)
+        twinkle.initialize(
+            mode='ray',
+            nproc_per_node=nproc_per_node,
+            groups=[self.device_group],
+            lazy_collect=False,
+            ncpu_proc_per_node=ncpu_proc_per_node)
+        if 'mesh_dim_names' in device_mesh:
+            self.device_mesh = DeviceMesh(**device_mesh)
+        else:
+            self.device_mesh = DeviceMesh.from_sizes(**device_mesh)
+
+        # processor objects keyed by processor_id
+        self.resource_dict: dict[str, Any] = {}
+        self.state: ServerStateProxy = get_server_state()
+
+        _cfg = processor_config or {}
+        _env_limit = int(os.environ.get('TWINKLE_PER_USER_PROCESSOR_LIMIT', 20))
+        self._init_processor_manager(
+            processor_timeout=float(_cfg.get('processor_timeout', 1800.0)),
+            per_token_processor_limit=int(_cfg.get('per_token_processor_limit', _env_limit)),
+        )
+        self.start_processor_countdown()
+
+    @serve.multiplexed(max_num_models_per_replica=100)
+    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)
+
+    def _on_processor_expired(self, processor_id: str) -> None:
+        """Called by the countdown thread when a processor's session expires."""
+        self.resource_dict.pop(processor_id, None)
+        self.unregister_processor(processor_id)
+
+
 def build_processor_app(ncpu_proc_per_node: int,
-                        device_group: Dict[str, Any],
-                        device_mesh: Dict[str, Any],
-                        deploy_options: Dict[str, Any],
+                        device_group: dict[str, Any],
+                        device_mesh: dict[str, Any],
+                        deploy_options: dict[str, Any],
                         nproc_per_node: int = 1,
+                        processor_config: dict[str, Any] | None = None,
                         **kwargs):
     """Build the processor management application.
 
+    Follows the same pattern as build_model_app(): FastAPI app and routes are
+    built here BEFORE serve.ingress so that the frozen app contains the full
+    route table visible to ProxyActor.
+
     Args:
-        ncpu_proc_per_node: Number of CPU processes per node
-        device_group: Device group configuration dict
-        device_mesh: Device mesh configuration dict
-        deploy_options: Ray Serve deployment options
-        nproc_per_node: Number of GPU processes per node (default 1, not used for CPU-only tasks)
-        **kwargs: Additional arguments
+        ncpu_proc_per_node: Number of CPU processes per node.
+        device_group: Device group configuration dict.
+        device_mesh: Device mesh configuration dict.
+        deploy_options: Ray Serve deployment options.
+        nproc_per_node: Number of GPU processes per node (default 1).
+        processor_config: Optional lifecycle configuration dict.
+            Supported keys:
+            - ``processor_timeout`` (float): Session inactivity timeout seconds. Default 1800.0.
+            - ``per_token_processor_limit`` (int): Max processors per user.
+              Overrides ``TWINKLE_PER_USER_PROCESSOR_LIMIT`` env var when provided.
+        **kwargs: Additional arguments.
 
     Returns:
-        Ray Serve deployment bound with configuration
+        Ray Serve deployment bound with configuration.
     """
+    # Build the FastAPI app and register all routes BEFORE serve.ingress so that
+    # the frozen app contains the complete route table (visible to ProxyActor).
     app = FastAPI()
 
     @app.middleware('http')
     async def verify_token(request: Request, call_next):
         return await verify_request_token(request=request, call_next=call_next)
 
-    processors = ['dataset', 'dataloader', 'preprocessor', 'processor', 'reward', 'template', 'weight_loader']
-
-    @serve.deployment(name='ProcessorManagement')
-    @serve.ingress(app)
-    class ProcessorManagement:
-        """Processor management service.
-
-        Manages lifecycle and invocation of distributed processor objects
-        (datasets, dataloaders, rewards, templates, etc.).
-        """
-
-        def __init__(self,
-                     ncpu_proc_per_node: int,
-                     device_group: Dict[str, Any],
-                     device_mesh: Dict[str, Any],
-                     nproc_per_node: int = 1):
-            self.device_group = DeviceGroup(**device_group)
-            twinkle.initialize(
-                mode='ray',
-                nproc_per_node=nproc_per_node,
-                groups=[self.device_group],
-                lazy_collect=False,
-                ncpu_proc_per_node=ncpu_proc_per_node)
-            if 'mesh_dim_names' in device_mesh:
-                self.device_mesh = DeviceMesh(**device_mesh)
-            else:
-                self.device_mesh = DeviceMesh.from_sizes(**device_mesh)
-            self.resource_dict = {}
-            self.state: ServerStateProxy = get_server_state()
-            self.per_token_processor_limit = int(os.environ.get('TWINKLE_PER_USER_PROCESSOR_LIMIT', 20))
-            self.key_token_dict = {}
-
-        def assert_processor_exists(self, processor_id: str):
-            assert processor_id and processor_id in self.resource_dict, f'Processor {processor_id} not found'
-
-        @app.post('/twinkle/create', response_model=types.ProcessorCreateResponse)
-        def create(self, request: Request, body: types.ProcessorCreateRequest) -> types.ProcessorCreateResponse:
-            processor_type_name = body.processor_type
-            class_type = body.class_type
-            _kwargs = body.model_extra or {}
-
-            assert processor_type_name in processors, f'Invalid processor type: {processor_type_name}'
-            processor_module = importlib.import_module(f'twinkle.{processor_type_name}')
-            assert hasattr(processor_module, class_type), f'Class {class_type} not found in {processor_type_name}'
-            processor_id = str(uuid.uuid4().hex)
-            self.key_token_dict[processor_id] = request.state.token
-
-            _kwargs.pop('remote_group', None)
-            _kwargs.pop('device_mesh', None)
-
-            resolved_kwargs = {}
-            for key, value in _kwargs.items():
-                if isinstance(value, str) and value.startswith('pid:'):
-                    ref_id = value[4:]
-                    resolved_kwargs[key] = self.resource_dict[ref_id]
-                else:
-                    value = deserialize_object(value)
-                    resolved_kwargs[key] = value
-
-            processor = getattr(processor_module, class_type)(
-                remote_group=self.device_group.name,
-                device_mesh=self.device_mesh,
-                instance_id=processor_id,
-                **resolved_kwargs)
-            self.resource_dict[processor_id] = processor
-            return types.ProcessorCreateResponse(processor_id='pid:' + processor_id)
-
-        @app.post('/twinkle/call', response_model=types.ProcessorCallResponse)
-        def call(self, body: types.ProcessorCallRequest) -> types.ProcessorCallResponse:
-            processor_id = body.processor_id
-            function_name = body.function
-            _kwargs = body.model_extra or {}
-            processor_id = processor_id[4:]
-            self.assert_processor_exists(processor_id=processor_id)
-            processor = self.resource_dict.get(processor_id)
-            function = getattr(processor, function_name, None)
-
-            assert function is not None, f'`{function_name}` not found in {processor.__class__}'
-            assert hasattr(function, '_execute'), f'Cannot call inner method of {processor.__class__}'
-
-            resolved_kwargs = {}
-            for key, value in _kwargs.items():
-                if isinstance(value, str) and value.startswith('pid:'):
-                    ref_id = value[4:]
-                    resolved_kwargs[key] = self.resource_dict[ref_id]
-                else:
-                    value = deserialize_object(value)
-                    resolved_kwargs[key] = value
-
-            # Special handling for __next__ to catch StopIteration
-            if function_name == '__next__':
-                try:
-                    result = function(**resolved_kwargs)
-                    return types.ProcessorCallResponse(result=result)
-                except StopIteration:
-                    # HTTP 410 Gone signals iterator exhausted
-                    raise HTTPException(status_code=410, detail='Iterator exhausted')
-
-            result = function(**resolved_kwargs)
-            if function_name == '__iter__':
-                return types.ProcessorCallResponse(result='ok')
-            else:
-                return types.ProcessorCallResponse(result=result)
-
-    return ProcessorManagement.options(**deploy_options).bind(
-        ncpu_proc_per_node, device_group, device_mesh, nproc_per_node=nproc_per_node)
+    def get_self() -> ProcessorManagement:
+        return serve.get_replica_context().servable_object
+
+    _register_processor_routes(app, get_self)
+
+    ProcessorManagementWithIngress = serve.ingress(app)(ProcessorManagement)
+    DeploymentClass = serve.deployment(name='ProcessorManagement')(ProcessorManagementWithIngress)
+    return DeploymentClass.options(**deploy_options).bind(ncpu_proc_per_node, device_group, device_mesh, nproc_per_node,
+                                                          processor_config)