fix(server): enable concurrent Tinker and Twinkle sampling

src/twinkle/exceptions.py

@@ -0,0 +1,5 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+
+
+class FullWeightLoadError(RuntimeError):
+    """A checkpoint failed after the engine began replacing base weights."""

src/twinkle/sampler/vllm_sampler/vllm_sampler.py

@@ -10,6 +10,7 @@
 from twinkle import DeviceMesh, get_logger, remote_class, remote_function, requires
 from twinkle.checkpoint_engine import CheckpointEngineMixin
 from twinkle.data_format import InputFeature, SampledSequence, SampleResponse, SamplingParams, Trajectory
+from twinkle.exceptions import FullWeightLoadError
 from twinkle.hub import HubOperation
 from twinkle.patch import Patch, apply_patch
 from twinkle.patch.vllm_lora_weights import VLLMLoraWeights
@@ -36,10 +37,10 @@ def _convert_ndarray_to_list(obj: Any) -> Any:
     return obj
 
 
-_MAX_CONCURRENCY = max(1, int(os.environ.get('TWINKLE_SAMPLER_MAX_CONCURRENCY') or 24))
+SAMPLER_MAX_CONCURRENCY = max(1, int(os.environ.get('TWINKLE_SAMPLER_MAX_CONCURRENCY') or 24))
 
 
-@remote_class(max_concurrency=_MAX_CONCURRENCY)
+@remote_class(max_concurrency=SAMPLER_MAX_CONCURRENCY)
 class vLLMSampler(Sampler, CheckpointEngineMixin):
     """A vLLM-based sampler using VLLMEngine (AsyncLLM).
 
@@ -232,10 +233,12 @@ async def _sample_single(
         multi_modal_data: Optional[Dict[str, Any]] = None,
         logprobs_only: bool = False,
         disable_lora: bool = False,
+        token_only: bool = False,
     ) -> SampleResponse:
         """Sample a single input asynchronously."""
         response = await self.engine.sample(
-            prompt=self.template.get_vllm_input_ids(feat['input_ids']) if self.template else feat['input_ids'],
+            prompt=self.template.get_vllm_input_ids(feat['input_ids'])
+            if self.template and not token_only else feat['input_ids'],
             sampling_params=sampling_params,
             lora_request=lora_request,
             multi_modal_data=multi_modal_data,
@@ -257,11 +260,19 @@ async def _sample_single(
         # labels when they are missing or length-mismatched. Present, aligned labels are
         # left untouched (preserving provenance), and the logprobs-only path -- which
         # never concatenates -- is not touched.
-        if not logprobs_only and 'input_ids' in feat and len(feat.get('labels') or []) != len(feat['input_ids']):
+        if not token_only and not logprobs_only and 'input_ids' in feat and len(feat.get('labels') or []) != len(
+                feat['input_ids']):
             feat['labels'] = [-100] * len(feat['input_ids'])
         sequences = []
         for seq in response.sequences:
-            if logprobs_only:
+            if token_only:
+                sampled_seq = SampledSequence(
+                    stop_reason=seq.stop_reason,
+                    tokens=[] if logprobs_only else seq.tokens,
+                    logprobs=None if logprobs_only else seq.logprobs,
+                    sampling_mask=seq.sampling_mask,
+                )
+            elif logprobs_only:
                 new_input_feature = _convert_ndarray_to_list(feat)
                 if seq.routed_experts is not None:
                     new_input_feature['routed_experts'] = _convert_ndarray_to_list(seq.routed_experts)
@@ -289,6 +300,22 @@ async def _sample_single(
             prompt_logprobs=response.prompt_logprobs,
             topk_prompt_logprobs=response.topk_prompt_logprobs)
 
+    @staticmethod
+    async def _gather_samples(samples) -> List[SampleResponse]:
+        """Finish every prompt before exposing the batch's result or failure."""
+        tasks = [asyncio.create_task(sample) for sample in samples]
+        batch = asyncio.gather(*tasks)
+        try:
+            return await asyncio.shield(batch)
+        except BaseException:
+            for task in tasks:
+                if not task.done():
+                    task.cancel()
+            await asyncio.gather(*tasks, return_exceptions=True)
+            if not batch.cancelled():
+                batch.exception()
+            raise
+
     @remote_function(dispatch='slice_dp', collect='flatten', lazy_collect=False, enable_continous_work=True)
     def sample(
         self,
@@ -299,6 +326,7 @@ def sample(
         *,
         return_encoded: bool = False,
         use_base_model: bool = False,
+        token_only: bool = False,
     ) -> List[SampleResponse]:
         """Sample responses for given inputs.
 
@@ -313,6 +341,9 @@ def sample(
 
             adapter_path: Optional LoRA adapter path.
 
+            token_only: Return token IDs and logprobs without template processing
+                or training features. Inputs must already contain input_ids.
+
             num_samples: Number of completions to generate per input prompt.
                 When > 1, returns num_samples sequences for each input.
 
@@ -334,6 +365,8 @@ def sample(
 
         # Check if inputs are Trajectory (not encoded) - aligned with Model.forward logic
         is_trajectory = 'input_ids' not in inputs_list[0]
+        if token_only and is_trajectory:
+            raise ValueError('Token-only sampling requires encoded input_ids.')
         logprobs_only = False
         if sampling_params.max_tokens == 0:
             sampling_params = copy(sampling_params)
@@ -373,9 +406,10 @@ async def _sample_all():
                     multi_modal_data=multi_modal_data,
                     logprobs_only=logprobs_only,
                     disable_lora=use_base_model,
+                    token_only=token_only,
                 ) for feat, multi_modal_data in zip(encoded_inputs, multi_modal_data_list)
             ]
-            return await asyncio.gather(*tasks)
+            return await self._gather_samples(tasks)
 
         sample_results = self._run_in_loop(_sample_all())
         return sample_results
@@ -553,6 +587,11 @@ def load_full_weights_from_path(self, path: Optional[str] = None) -> int:
             shards = sorted(glob.glob(os.path.join(resolved, '*.safetensors')))
         if not shards:
             raise FileNotFoundError(f'No .safetensors weights found under {resolved}')
+        # Check every shard header before changing any live weights.
+        for shard in shards:
+            with safe_open(shard, framework='pt', device='cpu') as f:
+                if not list(f.keys()):
+                    raise ValueError(f'Empty checkpoint shard: {shard}')
 
         def _iter_weights():
             # safe_open + get_tensor reads one tensor at a time (mmap-backed), so peak
@@ -569,11 +608,12 @@ async def _load():
             self.engine.invalidate_synced_lora()
 
         logger.info(f'Loading full-parameter weights into sampler base model from {resolved}')
-        self._run_in_loop(_load())
+        try:
+            self._run_in_loop(_load())
+            self.reset_prefix_cache()
+        except Exception as error:
+            raise FullWeightLoadError(f'Failed to install checkpoint {resolved}') from error
         self._loaded_full_weights_path = resolved
-        # Prefixes cached under the previous weights would decode against a model
-        # that no longer exists; drop them before the next sample().
-        self.reset_prefix_cache()
         logger.info(f'Reloaded base weights from {resolved} ({len(shards)} shard(s))')
         return 1
 

src/twinkle/server/sampler/app.py

@@ -19,8 +19,9 @@
 from twinkle.server.state import ServerState, get_server_state
 from twinkle.server.task_queue.config import TaskQueueConfig
 from twinkle.server.task_queue.mixin import TaskQueueMixin
-from twinkle.server.utils import wrap_builder_with_device_group_env
+from twinkle.server.utils import get_template_for_model, wrap_builder_with_device_group_env
 from twinkle.utils.logger import get_logger
+from .concurrency import SamplerConcurrency
 from .tinker_handlers import _register_tinker_sampler_routes
 from .twinkle_handlers import _register_twinkle_sampler_routes
 
@@ -41,9 +42,9 @@ def _make_vllm_sampler(kw: dict[str, Any]) -> Any:
 
 def _make_vllm_async_sampler(kw: dict[str, Any]) -> Any:
     """Construct the vLLM backend with non-blocking generation admission."""
-    from twinkle_agentic.async_rl.vllm_sampler_tq import VLLMSamplerTQ
+    from .backends.vllm_async import ServerVLLMSampler
 
-    return VLLMSamplerTQ(**kw, context_manager=None)
+    return ServerVLLMSampler(**kw, context_manager=None)
 
 
 # Single validate-then-dispatch selector for the sampler backend.
@@ -73,20 +74,20 @@ class SamplerManagement(LazyCleanupMixin, TaskQueueMixin):
     Manages:
     - mock or vLLM sampler initialization and lifecycle
     - Tinker inference requests (/tinker/asample) with rate limiting via TaskQueueMixin
-    - Twinkle inference requests (/twinkle/*) calling sampler directly
+    - Concurrent Twinkle inference requests (/twinkle/*) with future polling
     - Template configuration for trajectory encoding
     """
 
-    def __init__(self,
-                 model_id: str,
-                 nproc_per_node: int,
-                 device_group: dict[str, Any],
-                 device_mesh: dict[str, Any],
-                 sampler_type: str,
-                 engine_args: dict[str, Any] | None = None,
-                 queue_config: TaskQueueConfig | None = None,
-                 data_plane_url: str | None = None,
-                 **kwargs):
+    async def __init__(self,
+                       model_id: str,
+                       nproc_per_node: int,
+                       device_group: dict[str, Any],
+                       device_mesh: dict[str, Any],
+                       sampler_type: str,
+                       engine_args: dict[str, Any] | None = None,
+                       queue_config: TaskQueueConfig | None = None,
+                       data_plane_url: str | None = None,
+                       **kwargs):
         self.device_group = DeviceGroup(**device_group)
         self.device_mesh = init_twinkle_runtime(
             is_mock=(sampler_type == 'mock'),
@@ -127,20 +128,30 @@ def __init__(self,
             deployment_name='Sampler',
             collect_width=len(actors) if actors else 1,
         )
-        self.sampler._ray_get_timeout = self.task_queue_config.effective_execution_timeout
+        # ComputeWorker bounds the client future. Keep collecting the actual
+        # backend result so a timed-out request cannot release live weights.
+        self.sampler._ray_get_timeout = None
+        self.sampling_guard = SamplerConcurrency()
+        if sampler_type != 'mock':
+            await self.call_backend(self.sampler.set_template, get_template_for_model(model_id), model_id=model_id)
 
     async def shutdown(self) -> None:
+        self.sampling_guard.close()
         try:
             cancel_all = getattr(self.sampler, 'cancel_all_generations', None)
             if callable(cancel_all):
-                await self.call_backend(cancel_all)
+                await self.call_backend(cancel_all, admit=False)
         finally:
             try:
                 await self.state.unregister_replica(self.replica_id)
             finally:
                 await self.shutdown_task_queue()
+                await self.sampling_guard.wait_closed()
                 await self.data_plane.close()
 
+    async def check_health(self) -> None:
+        self.sampling_guard.check_available()
+
     @serve.multiplexed(max_num_models_per_replica=5)
     async def _sticky_entry(self, sticky_key: str):
         return sticky_key

src/twinkle/server/sampler/backends/vllm_async.py

@@ -0,0 +1,10 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+"""Server sampling concurrency without changing native async RL defaults."""
+from twinkle import remote_class
+from twinkle.sampler.vllm_sampler.vllm_sampler import SAMPLER_MAX_CONCURRENCY
+from twinkle_agentic.async_rl.vllm_sampler_tq import VLLMSamplerTQ
+
+
+@remote_class(max_concurrency=SAMPLER_MAX_CONCURRENCY)
+class ServerVLLMSampler(VLLMSamplerTQ):
+    pass

src/twinkle/server/sampler/concurrency.py

@@ -0,0 +1,92 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+"""Coordinate sampling with changes to shared sampler state."""
+from __future__ import annotations
+
+import asyncio
+
+from twinkle.server.task_queue.types import BackendBusyError
+
+
+class SamplerConcurrency:
+    """Readers share stable weights; mutations wait for every reader to finish.
+
+    Jobs retain their lock if an HTTP waiter leaves. Admission and execution
+    timeouts belong to ComputeWorker, not this lock.
+    """
+
+    def __init__(self):
+        self._readers = 0
+        self._writer = False
+        self._waiting_writers = 0
+        self._condition = asyncio.Condition()
+        self._jobs: set[asyncio.Task] = set()
+        self._weights_error: Exception | None = None
+        self._closed = False
+
+    def check_available(self):
+        if self._closed:
+            raise BackendBusyError('The sampler is shutting down.')
+        if self._weights_error is not None:
+            raise BackendBusyError('A checkpoint write failed; the sampler is restarting.')
+
+    def invalidate_weights(self, error: Exception):
+        self._weights_error = error
+
+    def close(self):
+        self._closed = True
+
+    async def _execute(self, factory, exclusive: bool):
+        async with self._condition:
+            if exclusive:
+                self._waiting_writers += 1
+            try:
+                await self._condition.wait_for(lambda: not self._writer and
+                                               (self._readers == 0 if exclusive else self._waiting_writers == 0))
+                self.check_available()
+                if exclusive:
+                    self._writer = True
+                else:
+                    self._readers += 1
+            finally:
+                if exclusive:
+                    self._waiting_writers -= 1
+                    self._condition.notify_all()
+        try:
+            return await factory()
+        finally:
+            async with self._condition:
+                if exclusive:
+                    self._writer = False
+                else:
+                    self._readers -= 1
+                self._condition.notify_all()
+
+    async def run(self, factory, *, exclusive: bool = False):
+        self.check_available()
+        started = False
+
+        async def call():
+            nonlocal started
+            started = True
+            return await factory()
+
+        job = asyncio.create_task(self._execute(call, exclusive))
+        self._jobs.add(job)
+
+        def finished(task):
+            self._jobs.discard(task)
+            if not task.cancelled():
+                task.exception()
+
+        job.add_done_callback(finished)
+        try:
+            return await asyncio.shield(job)
+        except asyncio.CancelledError:
+            if not started:
+                job.cancel()
+            # The blocking backend may still be using the locked weights.
+            await asyncio.gather(job, return_exceptions=True)
+            raise
+
+    async def wait_closed(self):
+        await asyncio.gather(*self._jobs, return_exceptions=True)