update async

Author: Yunnglin

src/twinkle/server/gateway/server.py

@@ -61,8 +61,8 @@ def _validate_base_model(self, base_model: str) -> None:
                 detail=f"Base model '{base_model}' is not supported. "
                 f"Supported models: {', '.join(supported_model_names)}")
 
-    def _get_base_model(self, model_id: str) -> str:
-        metadata = self.state.get_model_metadata(model_id)
+    async def _get_base_model(self, model_id: str) -> str:
+        metadata = await self.state.get_model_metadata(model_id)
         if metadata and metadata.get('base_model'):
             return metadata['base_model']
         raise HTTPException(status_code=404, detail=f'Model {model_id} not found')

src/twinkle/server/gateway/tinker_gateway_handlers.py

@@ -56,7 +56,7 @@ async def create_session(
             body: types.CreateSessionRequest,
             self: GatewayServer = Depends(self_fn),
     ) -> types.CreateSessionResponse:
-        session_id = self.state.create_session(body.model_dump())
+        session_id = await self.state.create_session(body.model_dump())
         return types.CreateSessionResponse(session_id=session_id)
 
     @app.post('/session_heartbeat')
@@ -72,7 +72,7 @@ async def session_heartbeat(
     async def create_sampling_session(
         request: Request, body: types.CreateSamplingSessionRequest, self: GatewayServer = Depends(self_fn)
     ) -> types.CreateSamplingSessionResponse:  # noqa: E125
-        sampling_session_id = self.state.create_sampling_session(body.model_dump())
+        sampling_session_id = await self.state.create_sampling_session(body.model_dump())
         return types.CreateSamplingSessionResponse(sampling_session_id=sampling_session_id)
 
     @app.post('/retrieve_future')
@@ -223,44 +223,44 @@ async def create_model(request: Request, body: types.CreateModelRequest,
 
     @app.post('/get_info')
     async def get_info(request: Request, body: types.GetInfoRequest, self: GatewayServer = Depends(self_fn)) -> Any:
-        return await self.proxy.proxy_to_model(request, 'get_info', self._get_base_model(body.model_id))
+        return await self.proxy.proxy_to_model(request, 'get_info', await self._get_base_model(body.model_id))
 
     @app.post('/unload_model')
     async def unload_model(request: Request, body: types.UnloadModelRequest,
                            self: GatewayServer = Depends(self_fn)) -> Any:
-        return await self.proxy.proxy_to_model(request, 'unload_model', self._get_base_model(body.model_id))
+        return await self.proxy.proxy_to_model(request, 'unload_model', await self._get_base_model(body.model_id))
 
     @app.post('/forward')
     async def forward(request: Request, body: types.ForwardRequest, self: GatewayServer = Depends(self_fn)) -> Any:
-        return await self.proxy.proxy_to_model(request, 'forward', self._get_base_model(body.model_id))
+        return await self.proxy.proxy_to_model(request, 'forward', await self._get_base_model(body.model_id))
 
     @app.post('/forward_backward')
     async def forward_backward(request: Request,
                                body: types.ForwardBackwardRequest,
                                self: GatewayServer = Depends(self_fn)) -> Any:
-        return await self.proxy.proxy_to_model(request, 'forward_backward', self._get_base_model(body.model_id))
+        return await self.proxy.proxy_to_model(request, 'forward_backward', await self._get_base_model(body.model_id))
 
     @app.post('/optim_step')
     async def optim_step(request: Request, body: types.OptimStepRequest, self: GatewayServer = Depends(self_fn)) -> Any:
-        return await self.proxy.proxy_to_model(request, 'optim_step', self._get_base_model(body.model_id))
+        return await self.proxy.proxy_to_model(request, 'optim_step', await self._get_base_model(body.model_id))
 
     @app.post('/save_weights')
     async def save_weights(request: Request, body: types.SaveWeightsRequest,
                            self: GatewayServer = Depends(self_fn)) -> Any:
-        return await self.proxy.proxy_to_model(request, 'save_weights', self._get_base_model(body.model_id))
+        return await self.proxy.proxy_to_model(request, 'save_weights', await self._get_base_model(body.model_id))
 
     @app.post('/load_weights')
     async def load_weights(request: Request, body: types.LoadWeightsRequest,
                            self: GatewayServer = Depends(self_fn)) -> Any:
-        return await self.proxy.proxy_to_model(request, 'load_weights', self._get_base_model(body.model_id))
+        return await self.proxy.proxy_to_model(request, 'load_weights', await self._get_base_model(body.model_id))
 
     # --- Sampler Proxy Endpoints ---
 
     @app.post('/asample')
     async def asample(request: Request, body: types.SampleRequest, self: GatewayServer = Depends(self_fn)) -> Any:
         base_model = body.base_model
         if not base_model and body.sampling_session_id:
-            session = self.state.get_sampling_session(body.sampling_session_id)
+            session = await self.state.get_sampling_session(body.sampling_session_id)
             if session:
                 base_model = session.get('base_model')
         return await self.proxy.proxy_to_sampler(request, 'asample', base_model)
@@ -271,4 +271,5 @@ async def save_weights_for_sampler(
             body: types.SaveWeightsForSamplerRequest,
             self: GatewayServer = Depends(self_fn),
     ) -> Any:
-        return await self.proxy.proxy_to_model(request, 'save_weights_for_sampler', self._get_base_model(body.model_id))
+        return await self.proxy.proxy_to_model(request, 'save_weights_for_sampler', await
+                                               self._get_base_model(body.model_id))

src/twinkle/server/gateway/twinkle_gateway_handlers.py

@@ -41,7 +41,7 @@ async def create_session(
             body: types.CreateSessionRequest,
             self: GatewayServer = Depends(self_fn),
     ) -> types.CreateSessionResponse:
-        session_id = self.state.create_session(body.model_dump())
+        session_id = await self.state.create_session(body.model_dump())
         return types.CreateSessionResponse(session_id=session_id)
 
     @app.post('/twinkle/session_heartbeat', response_model=types.SessionHeartbeatResponse)

src/twinkle/server/model/app.py

@@ -78,13 +78,19 @@ def __init__(self,
                 **kwargs)
 
         self.state: ServerStateProxy = get_server_state()
-        self.state.register_replica(self.replica_id, self.max_loras)
+        self._replica_registered = False
 
         # Initialize mixins
         self._init_task_queue(TaskQueueConfig.from_dict(queue_config))
         self._init_adapter_manager(**adapter_config)
         self.start_adapter_countdown()
 
+    async def _ensure_replica_registered(self):
+        """Lazily register replica on first async request."""
+        if not self._replica_registered:
+            await self.state.register_replica(self.replica_id, self.max_loras)
+            self._replica_registered = True
+
     @serve.multiplexed(max_num_models_per_replica=5)
     async def _sticky_entry(self, sticky_key: str):
         return sticky_key
@@ -95,22 +101,30 @@ async def _ensure_sticky(self):
 
     async def _on_request_start(self, request: Request) -> str:
         await self._ensure_sticky()
+        await self._ensure_replica_registered()
         token = get_token_from_request(request)
         return token
 
     def __del__(self):
-        self.state.unregister_replica(self.replica_id)
-
-    def _cleanup_adapter(self, adapter_name: str) -> None:
+        try:
+            # Best-effort cleanup; event loop may already be closed
+            import asyncio
+            loop = asyncio.get_event_loop()
+            if loop.is_running():
+                asyncio.create_task(self.state.unregister_replica(self.replica_id))
+        except Exception:
+            pass
+
+    async def _cleanup_adapter(self, adapter_name: str) -> None:
         if self.get_adapter_info(adapter_name):
             self.clear_adapter_state(adapter_name)
             self.model.remove_adapter(adapter_name)
             self.unregister_adapter(adapter_name)
-            self.state.unload_model(adapter_name)
+            await self.state.unload_model(adapter_name)
 
-    def _on_adapter_expired(self, adapter_name: str) -> None:
+    async def _on_adapter_expired(self, adapter_name: str) -> None:
         self.fail_pending_tasks_for_model(adapter_name, reason='Adapter expired')
-        self._cleanup_adapter(adapter_name)
+        await self._cleanup_adapter(adapter_name)
 
 
 def build_model_app(model_id: str,

src/twinkle/server/model/tinker_handlers.py

@@ -40,7 +40,7 @@ async def create_model(
         async def _create_adapter():
             _model_id = None
             try:
-                _model_id = self.state.register_model(body.model_dump(), token=token, replica_id=self.replica_id)
+                _model_id = await self.state.register_model(body.model_dump(), token=token, replica_id=self.replica_id)
                 if body.lora_config:
                     lora_cfg = LoraConfig(r=body.lora_config.rank, target_modules='all-linear')
                     adapter_name = self.get_adapter_name(adapter_name=_model_id)
@@ -56,7 +56,7 @@ async def _create_adapter():
             except Exception:
                 if _model_id:
                     adapter_name = self.get_adapter_name(adapter_name=_model_id)
-                    self._cleanup_adapter(adapter_name)
+                    await self._cleanup_adapter(adapter_name)
                 logger.error(traceback.format_exc())
                 return types.RequestFailedResponse(
                     error=traceback.format_exc(),
@@ -95,7 +95,7 @@ async def unload_model(
 
         async def _do_unload():
             adapter_name = self.get_adapter_name(adapter_name=body.model_id)
-            self._cleanup_adapter(adapter_name)
+            await self._cleanup_adapter(adapter_name)
             return types.UnloadModelResponse(model_id=body.model_id)
 
         return await self.schedule_task(_do_unload, model_id=body.model_id, token=token, task_type='unload_model')
@@ -260,10 +260,10 @@ async def _do_save_for_sampler():
                     name=checkpoint_name, output_dir=save_dir, adapter_name=adapter_name, save_optimizer=False)
                 payload = body.model_dump()
                 payload['model_path'] = tinker_path
-                metadata = self.state.get_model_metadata(body.model_id) or {}
+                metadata = await self.state.get_model_metadata(body.model_id) or {}
                 if metadata.get('base_model'):
                     payload['base_model'] = metadata['base_model']
-                sampling_session_id = self.state.create_sampling_session(payload)
+                sampling_session_id = await self.state.create_sampling_session(payload)
                 return types.SaveWeightsForSamplerResponseInternal(path=None, sampling_session_id=sampling_session_id)
             except Exception:
                 logger.error(traceback.format_exc())

src/twinkle/server/sampler/app.py

@@ -101,7 +101,7 @@ async def _on_request_start(self, request: Request) -> str:
         token = get_token_from_request(request)
         return token
 
-    def _on_adapter_expired(self, adapter_name: str, token: str = None) -> None:
+    async def _on_adapter_expired(self, adapter_name: str, token: str = None) -> None:
         """Handle expired adapters by removing them from the sampler."""
         try:
             self.sampler.remove_adapter(adapter_name)

src/twinkle/server/sampler/tinker_handlers.py

@@ -51,7 +51,7 @@ async def _do_sample():
                 # Get model_path from body or sampling session
                 model_path = body.model_path
                 if not model_path and body.sampling_session_id:
-                    session = self.state.get_sampling_session(body.sampling_session_id)
+                    session = await self.state.get_sampling_session(body.sampling_session_id)
                     if session:
                         model_path = session.get('model_path')
 

src/twinkle/server/utils/adapter_manager.py

@@ -107,8 +107,21 @@ def _is_session_alive(self, session_id: str) -> bool:
         if not session_id:
             return True  # No session association means always alive
 
-        # Get session last heartbeat through proxy
-        last_heartbeat = self.state.get_session_last_heartbeat(session_id)
+        # Get session last heartbeat through proxy (async method called from sync thread)
+        import asyncio
+        try:
+            loop = asyncio.get_event_loop()
+            if loop.is_running():
+                # If loop is running, we can't use asyncio.run()
+                # Use run_coroutine_threadsafe instead
+                future = asyncio.run_coroutine_threadsafe(self.state.get_session_last_heartbeat(session_id), loop)
+                last_heartbeat = future.result(timeout=5.0)
+            else:
+                last_heartbeat = asyncio.run(self.state.get_session_last_heartbeat(session_id))
+        except Exception as e:
+            logger.warning(f'[AdapterManager] Failed to check session liveness: {e}')
+            return True  # Assume alive on error
+
         if last_heartbeat is None:
             return False  # Session doesn't exist
 
@@ -182,7 +195,7 @@ def get_adapter_info(self, adapter_name: str) -> dict[str, Any] | None:
         """
         return self._adapter_records.get(adapter_name)
 
-    def _on_adapter_expired(self, adapter_name: str) -> None:
+    async def _on_adapter_expired(self, adapter_name: str) -> None:
         """Hook method called when an adapter expires.
 
         This method must be overridden by inheriting classes to handle
@@ -277,7 +290,18 @@ def _adapter_countdown_loop(self) -> None:
                 for adapter_name, _token, session_id in expired_adapters:
                     success = False
                     try:
-                        self._on_adapter_expired(adapter_name)
+                        # Call async _on_adapter_expired from sync thread
+                        import asyncio
+                        try:
+                            loop = asyncio.get_event_loop()
+                            if loop.is_running():
+                                future = asyncio.run_coroutine_threadsafe(self._on_adapter_expired(adapter_name), loop)
+                                future.result(timeout=10.0)
+                            else:
+                                asyncio.run(self._on_adapter_expired(adapter_name))
+                        except Exception as async_e:
+                            logger.warning(f'[AdapterManager] Async call failed for {adapter_name}: {async_e}')
+                            raise
                         logger.info(f'[AdapterManager] Adapter {adapter_name} expired '
                                     f'(reason=session_expired, session={session_id})')
                         success = True

src/twinkle/server/utils/processor_manager.py

@@ -125,7 +125,18 @@ def _is_session_alive(self, session_id: str) -> bool:
         """Check if a session is still alive via state proxy."""
         if not session_id:
             return True
-        last_heartbeat = self.state.get_session_last_heartbeat(session_id)
+        # Get session last heartbeat through proxy (async method called from sync thread)
+        import asyncio
+        try:
+            loop = asyncio.get_event_loop()
+            if loop.is_running():
+                future = asyncio.run_coroutine_threadsafe(self.state.get_session_last_heartbeat(session_id), loop)
+                last_heartbeat = future.result(timeout=5.0)
+            else:
+                last_heartbeat = asyncio.run(self.state.get_session_last_heartbeat(session_id))
+        except Exception as e:
+            logger.warning(f'[ProcessorManager] Failed to check session liveness: {e}')
+            return True  # Assume alive on error
         if last_heartbeat is None:
             return False
         return (time.time() - last_heartbeat) < self._processor_timeout

src/twinkle/server/utils/ray_serve_patch.py

@@ -63,7 +63,7 @@ def _patched_setup_request_context_and_handle(
             multiplexed_model_id = value.decode()
             handle = handle.options(multiplexed_model_id=multiplexed_model_id)
             request_context_info['multiplexed_model_id'] = multiplexed_model_id
-            logger.info(f'[Ray Serve Patch] Matched multiplexed_model_id: {multiplexed_model_id}')
+            logger.debug(f'[Ray Serve Patch] Matched multiplexed_model_id: {multiplexed_model_id}')
 
         # Original logic for other headers (unchanged)
         if decoded_key == 'x-request-id':
@@ -91,8 +91,8 @@ def _apply_patch_in_worker_process():
         HTTPProxy.setup_request_context_and_handle = _patched_setup_request_context_and_handle
         _patch_applied = True
 
-        logger.info('[Ray Serve Patch] Applied in worker process: '
-                    'HTTPProxy.setup_request_context_and_handle patched')
+        logger.debug('[Ray Serve Patch] Applied in worker process: '
+                     'HTTPProxy.setup_request_context_and_handle patched')
     except ImportError:
         # Ray Serve not available in this worker
         pass