update twinkle

Author: Yunnglin

client_tools/client_generator.py

@@ -67,7 +67,7 @@ def _prepare_headers(self, request_headers) -> dict[str, str]:
         headers.pop('host', None)
         headers.pop('content-length', None)
         request_id = request_headers.get('X-Ray-Serve-Request-Id')
-        if request_id is not None:
+        if request_id is not None and not request_headers.get('serve_multiplexed_model_id'):
             headers['serve_multiplexed_model_id'] = request_id
         return headers
 

src/twinkle/server/gateway/proxy.py

@@ -6,7 +6,7 @@
 """
 from __future__ import annotations
 
-from fastapi import FastAPI, HTTPException, Request
+from fastapi import Depends, FastAPI, HTTPException, Request
 from typing import TYPE_CHECKING, Callable
 
 if TYPE_CHECKING:
@@ -17,6 +17,8 @@
 from twinkle.server.utils.validation import get_token_from_request
 from twinkle.utils.logger import get_logger
 from twinkle_client.types.server import DeleteCheckpointResponse, HealthResponse, WeightsInfoRequest
+from twinkle_client.types.session import (CreateSessionRequest, CreateSessionResponse, SessionHeartbeatRequest,
+                                          SessionHeartbeatResponse)
 from twinkle_client.types.training import (CheckpointsListResponse, TrainingRun, TrainingRunsResponse,
                                            WeightsInfoResponse)
 
@@ -30,6 +32,26 @@ def _register_twinkle_routes(app: FastAPI, self_fn: Callable[[], GatewayServer])
     async def healthz(request: Request) -> HealthResponse:
         return HealthResponse(status='ok')
 
+    @app.post('/twinkle/create_session', response_model=CreateSessionResponse)
+    async def create_session(
+            request: Request,
+            body: CreateSessionRequest,
+            self: GatewayServer = Depends(self_fn),
+    ) -> CreateSessionResponse:
+        session_id = self.state.create_session(body.model_dump())
+        return CreateSessionResponse(session_id=session_id)
+
+    @app.post('/twinkle/session_heartbeat', response_model=SessionHeartbeatResponse)
+    async def session_heartbeat(
+            request: Request,
+            body: SessionHeartbeatRequest,
+            self: GatewayServer = Depends(self_fn),
+    ) -> SessionHeartbeatResponse:
+        alive = self.state.touch_session(body.session_id)
+        if not alive:
+            raise HTTPException(status_code=404, detail='Unknown session')
+        return SessionHeartbeatResponse()
+
     @app.get('/twinkle/training_runs', response_model=TrainingRunsResponse)
     async def get_training_runs(request: Request, limit: int = 20, offset: int = 0) -> TrainingRunsResponse:
         token = get_token_from_request(request)

src/twinkle/server/gateway/twinkle_gateway_handlers.py

@@ -1,10 +1,13 @@
 # Copyright (c) ModelScope Contributors. All rights reserved.
 from __future__ import annotations
+from typing import Optional
+
 
 def init_tinker_client(**kwargs) -> None:
     """Initialize Tinker client with Twinkle-specific headers.
 
-    After calling this function, users can directly use:
+    After calling this function, users can directly use::
+
         from tinker import ServiceClient
         client = ServiceClient(base_url='...', api_key='...')
 
@@ -13,39 +16,57 @@ def init_tinker_client(**kwargs) -> None:
     Args:
         **kwargs: Additional keyword arguments (currently unused, reserved for future)
 
-    Example:
-        >>> from twinkle import init_tinker_client
+    Example::
+
+        >>> from twinkle_client import init_tinker_client
         >>> init_tinker_client()
         >>> from tinker import ServiceClient
         >>> client = ServiceClient(base_url='http://localhost:8000', api_key='your_token')
     """
     from twinkle.utils import requires
-    
+
     requires('tinker')
     from twinkle_client.utils.patch_tinker import patch_tinker
 
-    # Apply patches to tinker library (includes header injection)
     patch_tinker()
 
 
-def init_twinkle_client(base_url: str | None = None, api_key: str | None = None, **kwargs) -> TwinkleClient:
+def init_twinkle_client(
+    base_url: Optional[str] = None,
+    api_key: Optional[str] = None,
+    session_heartbeat_interval: int = 30,
+    **kwargs,
+) -> 'TwinkleClient':
     """
-    Initialize a Twinkle client and setup context variables.
+    Initialize a Twinkle client.
+
+    This function:
+
+    * Resolves ``base_url`` and ``api_key`` (env-vars as fallbacks).
+    * Sets both values into the shared context so that all other client objects
+      (``MultiLoraTransformersModel``, ``vLLMSampler``, processor clients) created
+      afterwards automatically inherit the same server configuration.
+    * Creates a server-side session and stores the ``session_id`` in context so
+      every subsequent HTTP request carries it in ``X-Twinkle-Session-Id``.
+    * Starts a background thread that touches the session every
+      ``session_heartbeat_interval`` seconds.
+
+    Args:
+        base_url: Twinkle server base URL.  Falls back to ``TWINKLE_SERVER_URL``.
+        api_key: Authentication token.  Falls back to ``TWINKLE_SERVER_TOKEN``.
+        session_heartbeat_interval: Seconds between session touch calls (default: 30).
+        **kwargs: Additional keyword arguments forwarded to :class:`TwinkleClient`.
+
+    Returns:
+        An initialised :class:`~twinkle_client.manager.TwinkleClient` instance.
     """
-    from .http.utils import get_api_key, get_base_url, set_api_key, set_base_url
-    from .manager import TwinkleClient, TwinkleClientError
-    
-    if base_url is not None:
-        set_base_url(base_url)
-    else:
-        base_url = get_base_url()
-
-    if api_key is not None:
-        set_api_key(api_key)
-    else:
-        api_key = get_api_key()
-
-    return TwinkleClient(base_url=base_url, api_key=api_key, **kwargs)
+    from .manager import TwinkleClient
+    return TwinkleClient(
+        base_url=base_url,
+        api_key=api_key,
+        session_heartbeat_interval=session_heartbeat_interval,
+        **kwargs,
+    )
 
 
 __all__ = ['init_tinker_client', 'init_twinkle_client']

src/twinkle_client/__init__.py

@@ -10,7 +10,7 @@
 # ============================================================================
 
 from typing import Callable, Type, Union
-from twinkle_client.http import http_post, heartbeat_manager
+from twinkle_client.http import http_post
 from twinkle.dataset import Dataset
 from twinkle.processor import InputProcessor
 
@@ -31,13 +31,6 @@ def __init__(self, dataset: Union[Dataset, Callable], **kwargs):
         )
         response.raise_for_status()
         self.processor_id = response.json()['processor_id']
-        heartbeat_manager.register_processor(self.processor_id)
-
-    def __del__(self):
-        try:
-            heartbeat_manager.unregister_processor(self.processor_id)
-        except:
-            pass
 
 
     def __len__(self):

src/twinkle_client/dataloader/dataloader.py

@@ -10,7 +10,7 @@
 # ============================================================================
 
 from typing import Any, Callable, Dict, Type, Union
-from twinkle_client.http import http_post, heartbeat_manager
+from twinkle_client.http import http_post
 from twinkle.dataset import Dataset
 from twinkle.dataset import DatasetMeta
 from twinkle.preprocessor import DataFilter
@@ -34,13 +34,6 @@ def __init__(self, dataset_meta: DatasetMeta, **kwargs):
         )
         response.raise_for_status()
         self.processor_id = response.json()['processor_id']
-        heartbeat_manager.register_processor(self.processor_id)
-
-    def __del__(self):
-        try:
-            heartbeat_manager.unregister_processor(self.processor_id)
-        except:
-            pass
 
 
     def set_template(self, template_func: Union[Template, Type[Template], str], **kwargs):

src/twinkle_client/dataset/base.py

@@ -9,7 +9,7 @@
 #   2. Run: python client_tools/client_generator.py
 # ============================================================================
 
-from twinkle_client.http import http_post, heartbeat_manager
+from twinkle_client.http import http_post
 from twinkle.dataset import Dataset
 from twinkle.dataset import DatasetMeta
 from torch.utils.data import IterableDataset
@@ -31,13 +31,6 @@ def __init__(self, dataset_meta: DatasetMeta, **kwargs):
         )
         response.raise_for_status()
         self.processor_id = response.json()['processor_id']
-        heartbeat_manager.register_processor(self.processor_id)
-
-    def __del__(self):
-        try:
-            heartbeat_manager.unregister_processor(self.processor_id)
-        except:
-            pass
 
 
     def add_dataset(self, dataset_meta: DatasetMeta, **kwargs):

src/twinkle_client/dataset/iterable_dataset.py

@@ -10,7 +10,7 @@
 # ============================================================================
 
 from typing import Type, Union
-from twinkle_client.http import http_post, heartbeat_manager
+from twinkle_client.http import http_post
 from twinkle.dataset import Dataset
 from twinkle.dataset import DatasetMeta
 from twinkle.template import Template
@@ -33,13 +33,6 @@ def __init__(self, dataset_meta: DatasetMeta, packing_interval: int = 128, packi
         )
         response.raise_for_status()
         self.processor_id = response.json()['processor_id']
-        heartbeat_manager.register_processor(self.processor_id)
-
-    def __del__(self):
-        try:
-            heartbeat_manager.unregister_processor(self.processor_id)
-        except:
-            pass
 
 
     def set_template(self, template_cls: Union[Type[Template], str, Template], **kwargs):

src/twinkle_client/dataset/iterable_packing_dataset.py

@@ -9,7 +9,7 @@
 #   2. Run: python client_tools/client_generator.py
 # ============================================================================
 
-from twinkle_client.http import http_post, heartbeat_manager
+from twinkle_client.http import http_post
 from twinkle.dataset import Dataset
 from twinkle.dataset import DatasetMeta
 from .base import Dataset
@@ -31,13 +31,6 @@ def __init__(self, dataset_meta: DatasetMeta, **kwargs):
         )
         response.raise_for_status()
         self.processor_id = response.json()['processor_id']
-        heartbeat_manager.register_processor(self.processor_id)
-
-    def __del__(self):
-        try:
-            heartbeat_manager.unregister_processor(self.processor_id)
-        except:
-            pass
 
 
     def encode(self, **kwargs):

src/twinkle_client/dataset/lazy_dataset.py

@@ -9,7 +9,7 @@
 #   2. Run: python client_tools/client_generator.py
 # ============================================================================
 
-from twinkle_client.http import http_post, heartbeat_manager
+from twinkle_client.http import http_post
 from twinkle.dataset import Dataset
 from twinkle.dataset import DatasetMeta
 from .base import Dataset
@@ -31,13 +31,6 @@ def __init__(self, dataset_meta: DatasetMeta, packing_num_proc: int = 1, **kwarg
         )
         response.raise_for_status()
         self.processor_id = response.json()['processor_id']
-        heartbeat_manager.register_processor(self.processor_id)
-
-    def __del__(self):
-        try:
-            heartbeat_manager.unregister_processor(self.processor_id)
-        except:
-            pass
 
 
     def pack_dataset(self):