From c84e6337f3de42f57d537656a71864705a374df5 Mon Sep 17 00:00:00 2001 From: Yunnglin Date: Tue, 13 Oct 2026 20:21:40 +0800 Subject: [PATCH 1/5] fix(server): enable concurrent ordinary sampler requests --- .../Usage Guide/Server and Client/Server.md | 7 + .../\346\234\215\345\212\241\347\253\257.md" | 7 + .../sampler/vllm_sampler/vllm_sampler.py | 15 +- src/twinkle/server/sampler/app.py | 37 +- src/twinkle/server/sampler/concurrency.py | 116 ++++++ src/twinkle/server/sampler/tinker_handlers.py | 48 ++- .../server/sampler/twinkle_handlers.py | 71 ++-- src/twinkle/server/sampler/weights.py | 54 ++- src/twinkle/server/task_queue/config.py | 2 + src/twinkle/server/task_queue/mixin.py | 55 ++- src/twinkle/server/task_queue/types.py | 8 + src/twinkle/server/task_queue/worker.py | 126 ++++-- .../async_rl/vllm_sampler_tq.py | 110 +++-- tests/sampler/test_vllm_sampling_lifecycle.py | 36 ++ .../server/integration/sampler_e2e_helpers.py | 191 +++++++++ .../test_sampler_concurrency_e2e.py | 279 +++++++++++++ tests/server/sampler/sampling_test_helpers.py | 147 +++++++ .../sampler/test_resolve_sampler_weights.py | 2 +- .../sampler/test_sampling_concurrency.py | 392 ++++++++++++++++++ .../server/sampler/test_stream_guarantees.py | 5 +- tests/server/sampler/test_tinker_handlers.py | 39 +- .../server/sampler/test_twinkle_async_rows.py | 2 + tests/server/utils/test_task_queue_mixin.py | 211 ++++++++++ .../test_vllm_sampler_tq_generation.py | 58 ++- 24 files changed, 1833 insertions(+), 185 deletions(-) create mode 100644 src/twinkle/server/sampler/concurrency.py create mode 100644 tests/sampler/test_vllm_sampling_lifecycle.py create mode 100644 tests/server/integration/sampler_e2e_helpers.py create mode 100644 tests/server/integration/test_sampler_concurrency_e2e.py create mode 100644 tests/server/sampler/sampling_test_helpers.py create mode 100644 tests/server/sampler/test_sampling_concurrency.py diff --git a/docs/source_en/Usage Guide/Server and Client/Server.md b/docs/source_en/Usage Guide/Server and Client/Server.md index dd0064aba..4b24dc592 100644 --- a/docs/source_en/Usage Guide/Server and Client/Server.md +++ b/docs/source_en/Usage Guide/Server and Client/Server.md @@ -479,3 +479,10 @@ raises. During that window newly arriving tasks fail fast with a `server`/503 er Each persisted future stores its immutable `absolute_deadline` when it is created. Cleanup therefore reaches the same decision regardless of which deployment process holds the cleanup lease. + +### Sampler concurrency and weight changes + +- Ordinary Tinker and Twinkle `sample()` requests support concurrency without changing response schemas or future polling. `queue_config.concurrent_task_limit` limits unfinished concurrent tasks per sampler replica and defaults to `128`; capacity exhaustion returns `server`/503. Training compute remains serial. +- During ordinary sampling, full-checkpoint switches wait for existing ordinary and DataPlane generation to finish. Loading a full checkpoint and generating with it cannot overlap these requests. LoRA uses adapter checkpoints; full checkpoints are for full-parameter or merged models. +- After an ordinary sampling request times out, backend cleanup may still be running. New requests are temporarily refused during cleanup. If resource release or the integrity of weights after a failed update cannot be confirmed, the replica keeps returning `server`/503 and must be restarted. +- Unexpected backend cancellation leaves a `failed` future and requires a replica restart. Ordinary generation errors allow subsequent requests once cleanup finishes. diff --git "a/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/\346\234\215\345\212\241\347\253\257\345\222\214\345\256\242\346\210\267\347\253\257/\346\234\215\345\212\241\347\253\257.md" "b/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/\346\234\215\345\212\241\347\253\257\345\222\214\345\256\242\346\210\267\347\253\257/\346\234\215\345\212\241\347\253\257.md" index 81fcacfe5..1b72302bf 100644 --- "a/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/\346\234\215\345\212\241\347\253\257\345\222\214\345\256\242\346\210\267\347\253\257/\346\234\215\345\212\241\347\253\257.md" +++ "b/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/\346\234\215\345\212\241\347\253\257\345\222\214\345\256\242\346\210\267\347\253\257/\346\234\215\345\212\241\347\253\257.md" @@ -466,3 +466,10 @@ twinkle-server check-config -c server_config.yaml 任务记录进入终态后，per-replica 准入闸门额外保持关闭的最长时长为 `max(0, Collect_Width × R − T)`：此时记录已是终态，但泄漏的 executor 线程可能仍持有闸门，直到其 `ray.get` 返回或抛出。在该窗口内新到达的任务会以 `server`/503 错误快速失败。 每条持久化 future 在创建时写入不可变的 `absolute_deadline`，因此无论哪个 deployment 进程持有 cleanup lease，清理结果都由任务自身契约决定。 + +### Sampler 并发与权重切换 + +- Tinker 和 Twinkle 的普通 `sample()` 支持并发，响应格式与 future 轮询接口保持不变。`queue_config.concurrent_task_limit` 限制单个 sampler 副本尚未结束的并发任务数，默认为 `128`；满载时返回 `server`/503。训练计算任务仍然串行。 +- 普通采样中，切换全量 checkpoint 会等待已有普通生成和 DataPlane 生成结束；全量 checkpoint 加载与使用它生成期间不允许这些请求交错执行。LoRA 使用 adapter checkpoint；全量 checkpoint 用于全参数或合并后的模型。 +- 普通采样请求超时后，后端可能仍在清理。清理期间会暂时拒绝新请求；无法确认资源释放，或权重更新失败后无法确认权重完整性时，副本持续返回 `server`/503，需要重启。 +- 后端意外取消时，future 会进入 `failed` 终态，副本需要重启；普通生成错误在清理完成后不影响后续请求。 diff --git a/src/twinkle/sampler/vllm_sampler/vllm_sampler.py b/src/twinkle/sampler/vllm_sampler/vllm_sampler.py index 6438fa660..148e8df4b 100644 --- a/src/twinkle/sampler/vllm_sampler/vllm_sampler.py +++ b/src/twinkle/sampler/vllm_sampler/vllm_sampler.py @@ -36,10 +36,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). @@ -289,6 +289,15 @@ 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.""" + results = await asyncio.gather(*samples, return_exceptions=True) + for result in results: + if isinstance(result, BaseException): + raise result + return results + @remote_function(dispatch='slice_dp', collect='flatten', lazy_collect=False, enable_continous_work=True) def sample( self, @@ -375,7 +384,7 @@ async def _sample_all(): disable_lora=use_base_model, ) 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 diff --git a/src/twinkle/server/sampler/app.py b/src/twinkle/server/sampler/app.py index b10c5894b..a3f4fc977 100644 --- a/src/twinkle/server/sampler/app.py +++ b/src/twinkle/server/sampler/app.py @@ -7,6 +7,7 @@ """ from __future__ import annotations +import asyncio from fastapi import FastAPI, Request from ray import serve from typing import Any @@ -19,8 +20,10 @@ 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.task_queue.types import BACKEND_TIMEOUT_EXCEPTIONS from twinkle.server.utils import 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 @@ -73,7 +76,7 @@ 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 """ @@ -128,12 +131,17 @@ def __init__(self, collect_width=len(actors) if actors else 1, ) self.sampler._ray_get_timeout = self.task_queue_config.effective_execution_timeout + self.sampling_guard = SamplerConcurrency(self.task_queue_config.concurrent_task_limit) + # Serialize only the one-time Tinker template initialization. + self.tinker_setup_lock = asyncio.Lock() + self.tinker_template_set = False 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) @@ -141,6 +149,31 @@ async def shutdown(self) -> None: await self.shutdown_task_queue() await self.data_plane.close() + async def call_backend(self, fn, /, *args, admit: bool = True, mutation: bool = False, **kwargs): + + if not admit: + return await super().call_backend(fn, *args, admit=False, **kwargs) + + async def call(): + try: + return await super(SamplerManagement, self).call_backend(fn, *args, admit=admit, **kwargs) + except asyncio.CancelledError: + self.sampling_guard.poison() + raise + except BACKEND_TIMEOUT_EXCEPTIONS: + self.sampling_guard.poison() + raise + except Exception: + if mutation: + self.sampling_guard.poison() + raise + + if self.sampling_guard.owns_job(): + return await call() + # Control-plane operations also own their backend lifetime. Sampling + # calls inside a job reuse its access rather than acquiring another lock. + return await self.sampling_guard.run(call, exclusive=True) + @serve.multiplexed(max_num_models_per_replica=5) async def _sticky_entry(self, sticky_key: str): return sticky_key diff --git a/src/twinkle/server/sampler/concurrency.py b/src/twinkle/server/sampler/concurrency.py new file mode 100644 index 000000000..87608729a --- /dev/null +++ b/src/twinkle/server/sampler/concurrency.py @@ -0,0 +1,116 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Let each background sampling job own access until its backend has exited.""" +from __future__ import annotations + +import asyncio +from collections.abc import Callable, Coroutine +from contextvars import ContextVar +from typing import Any + +from twinkle.server.exceptions import ResourceExhaustedError +from twinkle.server.task_queue.types import BackendBusyError + + +class SamplerConcurrency: + """Stable weights admit readers; changing base weights requires one writer. + + A job owns its entire load/generate/cleanup sequence. Cancelling its waiter + leaves that job running, so a request timeout cannot release live weights. + All bookkeeping runs on the replica's event loop. + """ + + def __init__(self, limit: int = 128): + self._limit = limit + self._readers = 0 + self._writer = False + self._waiting_writers = 0 + self._changed = asyncio.Event() + self._jobs: set[asyncio.Task] = set() + self._draining: set[asyncio.Task] = set() + # Nested polling tasks inherit the owner, without passing lock handles + # through every backend call. Membership prevents use after job exit. + self._owner: ContextVar[asyncio.Task | None] = ContextVar('sampler_job', default=None) + self._uncertain = False + self._closed = False + + def owns_job(self) -> bool: + owner = self._owner.get() + return owner is not None and not owner.done() and owner in self._jobs + + def _notify(self) -> None: + self._changed.set() + self._changed = asyncio.Event() + + def check_available(self, *, admitted: bool = False) -> None: + if self._closed: + raise BackendBusyError('The sampler is shutting down.') + if self._uncertain: + raise BackendBusyError('Sampler weights or backend completion are uncertain; restart this replica.') + if self._draining and not admitted: + raise BackendBusyError('The sampler is waiting for backend resource release.') + + def poison(self) -> None: + self._uncertain = True + self._notify() + + def close(self) -> None: + self._closed = True + self._notify() + + async def _acquire(self, exclusive: bool) -> None: + if exclusive: + self._waiting_writers += 1 + try: + while True: + self.check_available(admitted=True) + if asyncio.current_task() in self._draining: + raise BackendBusyError('Sampling request exited before its job could start.') + if not self._writer and (self._readers == 0 if exclusive else self._waiting_writers == 0): + if exclusive: + self._writer = True + else: + self._readers += 1 + return + await self._changed.wait() + finally: + if exclusive: + self._waiting_writers -= 1 + self._notify() + + async def _execute(self, factory, exclusive: bool): + await self._acquire(exclusive) + token = self._owner.set(asyncio.current_task()) + try: + return await factory() + finally: + self._owner.reset(token) + if exclusive: + self._writer = False + else: + self._readers -= 1 + self._notify() + + def _finished(self, job: asyncio.Task) -> None: + self._jobs.discard(job) + self._draining.discard(job) + if not job.cancelled(): + job.exception() # Observe failures even when its waiter already left. + self._notify() + + async def run(self, factory: Callable[[], Coroutine[Any, Any, Any]], *, exclusive: bool = False): + """Wait for one owned job without transferring cancellation to it.""" + self.check_available() + # Unlike the future queue's admission count, this capacity is retained + # until actual backend cleanup, including jobs outliving request timeout. + if len(self._jobs) >= self._limit: + raise ResourceExhaustedError('Concurrent sampling capacity is unavailable on this replica.') + job = asyncio.create_task(self._execute(factory, exclusive)) + self._jobs.add(job) + job.add_done_callback(self._finished) + try: + return await asyncio.shield(job) + except asyncio.CancelledError: + if not job.done(): + self._draining.add(job) + self._notify() + raise diff --git a/src/twinkle/server/sampler/tinker_handlers.py b/src/twinkle/server/sampler/tinker_handlers.py index 923458ce2..635a9ea36 100644 --- a/src/twinkle/server/sampler/tinker_handlers.py +++ b/src/twinkle/server/sampler/tinker_handlers.py @@ -4,7 +4,7 @@ Registered by ``_register_tinker_sampler_routes(app, self_fn)`` -- module-level route registration closing over ``self_fn`` via ``Depends``, not a mixin: there is no inheritance relationship with the deployment class. Provides POST /tinker/asample using -schedule_task() returning UntypedAPIFuture. +schedule_concurrent_task() returning UntypedAPIFuture. """ from __future__ import annotations @@ -20,7 +20,7 @@ from twinkle.data_format import SamplingParams from twinkle.server.checkpoint import create_checkpoint_manager -from twinkle.server.sampler.weights import resolve_sampler_weights +from twinkle.server.sampler.weights import sample_with_weights from twinkle.server.task_queue.types import UserTaskError from twinkle.server.utils import get_template_for_model from twinkle.utils.logger import get_logger @@ -70,15 +70,6 @@ async def _do_sample(): # Extract prompt token IDs from ModelInput prompt_inputs = {'input_ids': body.prompt.to_ints()} - # Set template for sampler based on model type - template = get_template_for_model(self.model_id) - await self.call_backend(self.sampler.set_template, template, model_id=self.model_id) - # Reset prefix cache unconditionally on every tinker request (by - # design): the tinker dialect does not signal whether weights - # changed, so it always invalidates. This differs from the twinkle - # endpoints, which reset only when an adapter_uri is supplied. - await self.call_backend(self.sampler.reset_prefix_cache) - # Get model_path from body or sampling session model_path = body.model_path if not model_path and body.sampling_session_id: @@ -113,18 +104,28 @@ async def _do_sample(): logprobs=1, ) - # LoRA adapter dir vs full-parameter checkpoint (shared helper); - # a full checkpoint is loaded into the base model and yields no - # LoRA path. - lora_path = await resolve_sampler_weights(self, adapter_uri) - - responses = await self.call_backend( - self.sampler.sample, - inputs=[prompt_inputs] * body.num_samples, - sampling_params=sampling_params, - adapter_path=lora_path, - ) + # Requests run concurrently, so sampler setup is serialized here. + # The template depends only on the model and is set once: every + # set_template or reset_prefix_cache call first waits for all + # in-flight generations (infra ``_cw_barrier``), so doing it per + # request would serialize sampling again. No prefix-cache reset is + # needed: vLLM caches LoRA requests under their own adapter name, + # and load_full_weights_from_path resets the cache itself. + async with self.tinker_setup_lock: + if not self.tinker_template_set: + template = get_template_for_model(self.model_id) + await self.call_backend(self.sampler.set_template, template, model_id=self.model_id) + self.tinker_template_set = True + + async def generate(lora_path): + return await self.call_backend( + self.sampler.sample, + inputs=[prompt_inputs] * body.num_samples, + sampling_params=sampling_params, + adapter_path=lora_path, + ) + responses = await sample_with_weights(self, adapter_uri, generate) tinker_sequences = [] for response in responses: # twinkle logprobs are ``List[List[Tuple[int, float]]]`` @@ -159,7 +160,8 @@ async def _do_sample(): raise input_tokens = len(body.prompt.to_ints()) - return await self.schedule_task( + # Not the serial compute queue: vLLM batches concurrent requests itself. + return await self.schedule_concurrent_task( _do_sample, token=token, input_tokens=input_tokens, diff --git a/src/twinkle/server/sampler/twinkle_handlers.py b/src/twinkle/server/sampler/twinkle_handlers.py index 2aca4b22b..2bef2ae42 100644 --- a/src/twinkle/server/sampler/twinkle_handlers.py +++ b/src/twinkle/server/sampler/twinkle_handlers.py @@ -29,7 +29,7 @@ from twinkle.protocol.types import sampler as sampler_types from twinkle.server.exceptions import EndpointUnavailableError, RequestRejectedError from twinkle.server.lifecycle.submit import backend_kwargs, resolve_twinkle_adapter_name, to_backend_inputs -from twinkle.server.sampler.weights import resolve_sampler_weights +from twinkle.server.sampler.weights import sample_with_weights from twinkle.server.task_errors import task_error_payload from twinkle.server.telemetry.correlation import MODEL_ID from twinkle.server.telemetry.tracing import traced_operation @@ -166,11 +166,15 @@ async def _stream_queue(q, sentinel, request_id: str, total_timeout: float, sing async def _await_generation(service: SamplerManagement, submission_id: str, timeout: float): """Poll one admitted generation through the backend boundary.""" collected = False + loop = asyncio.get_running_loop() + deadline = loop.time() + timeout async def poll(): nonlocal collected poll_interval = 0.01 while True: + if loop.time() >= deadline: + raise asyncio.TimeoutError try: states = _submission_states(await service.call_backend(service.sampler.get_generation_status, submission_id)) @@ -196,14 +200,24 @@ async def poll(): poll_interval = min(poll_interval * 1.5, 0.25) try: - return await asyncio.wait_for(poll(), timeout=timeout) + return await poll() finally: if not collected: + + async def cancel_and_wait(): + while True: + states = _submission_states(await service.call_backend(service.sampler.cancel_generation, + submission_id)) + if states and all(state.get('status') in ('completed', 'missing') for state in states): + return + await asyncio.sleep(0.01) + try: - await asyncio.wait_for( - service.call_backend(service.sampler.cancel_generation, submission_id), timeout=4.0) - except Exception: - logger.warning('Failed to cancel generation %s', submission_id, exc_info=True) + await asyncio.wait_for(cancel_and_wait(), timeout=4.0) + except BaseException: + service.sampling_guard.poison() + logger.warning('Generation cleanup did not confirm completion: %s', submission_id, exc_info=True) + # Preserve the original generation error, including cancellation. def _register_twinkle_sampler_routes(app: FastAPI, self_fn: Callable[[], SamplerManagement]) -> None: @@ -230,7 +244,7 @@ async def sample(request: Request, body: types.SampleRequest, async def _task(): # Resolve adapter - adapter_path = None + resolved_uri = None adapter_name = body.adapter_name or '' full_adapter_name = resolve_twinkle_adapter_name(request, adapter_name) or '' @@ -238,9 +252,6 @@ async def _task(): from twinkle.server.checkpoint import create_checkpoint_manager checkpoint_manager = create_checkpoint_manager(token, client_type='twinkle') _, resolved_uri = checkpoint_manager.parse_adapter_uri(body.adapter_uri) - # Reset prefix cache only when new weights are loaded. - await self.call_backend(self.sampler.reset_prefix_cache) - adapter_path = await resolve_sampler_weights(self, resolved_uri) # Parse inputs (shared seam; batch form) inputs = to_backend_inputs(body.inputs) @@ -251,19 +262,22 @@ async def _task(): params = SamplingParams.from_dict(body.sampling_params) # Sample - responses = await self.call_backend( - self.sampler.sample, - inputs, - params, - adapter_name=full_adapter_name, - adapter_path=adapter_path, - ) + async def generate(adapter_path): + return await self.call_backend( + self.sampler.sample, + inputs, + params, + adapter_name=full_adapter_name, + adapter_path=adapter_path, + ) + + responses = await sample_with_weights(self, resolved_uri, generate) return types.SampleResponseModelList(samples=_to_sample_response_models(responses)).model_dump() # Calculate metrics for queue scheduling. The body is wire-validated, so the # entries are models and ``input_ids`` is absent or a list of ints. input_tokens = sum(len(getattr(entry, 'input_ids', None) or ()) for entry in body.inputs) - return await self.submit_and_peek(_task, token=token, input_tokens=input_tokens, task_type='sample') + return await self.submit_concurrent_and_peek(_task, token=token, input_tokens=input_tokens, task_type='sample') @app.post('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/twinkle/sample_to_data_plane', response_model=types.TaskEnvelope) async def sample_to_data_plane( @@ -298,15 +312,18 @@ async def _generate_and_store(): # vLLM async engine owns generation concurrency, so the whole # admit -> await -> store sequence runs as one background future # (outside the serial compute queue) and its result is the DataRef. - await self.call_backend( - self.sampler.submit_generation, - submission_id, - inputs, - params, - adapter_name=full_adapter_name, - adapter_path=adapter_path, - ) - responses = await _await_generation(self, submission_id, self.task_queue_config.effective_execution_timeout) + async def generate(): + await self.call_backend( + self.sampler.submit_generation, + submission_id, + inputs, + params, + adapter_name=full_adapter_name, + adapter_path=adapter_path, + ) + return await _await_generation(self, submission_id, self.task_queue_config.effective_execution_timeout) + + responses = await self.sampling_guard.run(generate) rows, tags = _build_rollout_rows_and_tags( _to_sample_response_models(responses), group_ids=body.group_ids, diff --git a/src/twinkle/server/sampler/weights.py b/src/twinkle/server/sampler/weights.py index 2e5005df8..5c8363e00 100644 --- a/src/twinkle/server/sampler/weights.py +++ b/src/twinkle/server/sampler/weights.py @@ -1,35 +1,29 @@ # Copyright (c) ModelScope Contributors. All rights reserved. -"""Shared sampler weight-resolution rule (F004 / P004). - -The "a resolved checkpoint is either a LoRA adapter dir or a full-parameter HF -checkpoint" rule was re-derived from a filesystem probe in three sampler handlers -(``sample`` / ``sample_stream`` on the twinkle dialect and ``asample`` on tinker), -and the copies had begun to diverge. It lives here so the storage-layout decision -has a single owner. - -Prefix-cache invalidation is deliberately NOT handled here: each caller keeps its -own ``reset_prefix_cache`` policy (the tinker endpoint resets unconditionally on -every request; the twinkle endpoints reset only when an ``adapter_uri`` is -present), which is an observable behaviour difference this helper must not erase. -""" +"""Run checkpoint loading and generation as one owned sampling job.""" from __future__ import annotations import os -from typing import Any - - -async def resolve_sampler_weights(service: Any, resolved_uri: str | None) -> str | None: - """Resolve a checkpoint path into a LoRA adapter path (or load full weights). - - Returns the LoRA ``adapter_path`` when ``resolved_uri`` is a directory holding - an ``adapter_config.json``. Otherwise the path is a full-parameter checkpoint: - it is loaded into the sampler base model via ``load_full_weights_from_path`` and - ``None`` is returned (no LoRA adapter to pass). ``None``/empty input returns - ``None`` unchanged. - """ - if not resolved_uri: - return None - if os.path.exists(os.path.join(resolved_uri, 'adapter_config.json')): - return resolved_uri - await service.call_backend(service.sampler.load_full_weights_from_path, resolved_uri) + + +def is_full_checkpoint(resolved_uri: str | None) -> bool: + return bool(resolved_uri and not os.path.exists(os.path.join(resolved_uri, 'adapter_config.json'))) + + +async def resolve_sampler_weights(service, resolved_uri: str | None, *, full_checkpoint=None) -> str | None: + """Load base weights or return the adapter path.""" + full_checkpoint = is_full_checkpoint(resolved_uri) if full_checkpoint is None else full_checkpoint + if not full_checkpoint: + return resolved_uri or None + await service.call_backend(service.sampler.load_full_weights_from_path, resolved_uri, mutation=True) return None + + +async def sample_with_weights(service, resolved_uri: str | None, generate): + """Keep a full checkpoint's load and generation in the same exclusive job.""" + full_checkpoint = is_full_checkpoint(resolved_uri) + + async def job(): + adapter_path = await resolve_sampler_weights(service, resolved_uri, full_checkpoint=full_checkpoint) + return await generate(adapter_path) + + return await service.sampling_guard.run(job, exclusive=full_checkpoint) diff --git a/src/twinkle/server/task_queue/config.py b/src/twinkle/server/task_queue/config.py index 449c694cc..76f1813e8 100644 --- a/src/twinkle/server/task_queue/config.py +++ b/src/twinkle/server/task_queue/config.py @@ -32,6 +32,7 @@ class TaskQueueConfig(BaseModel): token_cleanup_multiplier: Multiplier for token cleanup threshold. token_cleanup_interval: How often to run cleanup task (seconds). max_input_tokens: Maximum allowed input tokens per request. + concurrent_task_limit: Maximum admitted concurrent compute tasks per replica. inline_fast_path_timeout: Upper bound (seconds) on how long submit briefly polls the record so a millisecond-scale control-plane op (step / zero_grad) completes in a single HTTP round trip instead of forcing a retrieve. @@ -49,6 +50,7 @@ class TaskQueueConfig(BaseModel): token_cleanup_multiplier: float = Field(default=10.0, ge=0) token_cleanup_interval: float = Field(default=60.0, ge=0) max_input_tokens: int = Field(default=16000, ge=1) + concurrent_task_limit: int = Field(default=128, ge=1) inline_fast_path_timeout: float = Field(default=0.05, gt=0) @property diff --git a/src/twinkle/server/task_queue/mixin.py b/src/twinkle/server/task_queue/mixin.py index 31bdff261..f26ec9c5c 100644 --- a/src/twinkle/server/task_queue/mixin.py +++ b/src/twinkle/server/task_queue/mixin.py @@ -34,7 +34,7 @@ class TaskQueueMixin: - """Mixin providing two task execution paths. + """Mixin providing three task execution paths. Execution paths --------------- @@ -42,7 +42,12 @@ class TaskQueueMixin: Single background worker, serial execution, round-robin across queues. Use for GPU operations: forward, backward, step, save, load, etc. - 2. Background task (schedule_background_task): + 2. Concurrent compute task (schedule_concurrent_task): + Same admission, status tracking and failure handling as the compute + queue, but starts immediately. Use for a backend that batches + concurrent requests itself, such as vLLM sampling. + + 3. Background task (schedule_background_task): asyncio.create_task, runs concurrently with compute queue. Use for I/O or an engine that provides its own safe concurrency and batching. Status is still tracked; clients can poll the same status endpoints. @@ -305,11 +310,57 @@ async def schedule_task( request_id=request_id, ) + async def schedule_concurrent_task( + self, + coro_factory: Callable[[], Coroutine], + model_id: str | None = None, + token: str | None = None, + input_tokens: int = 0, + task_type: str | None = None, + request_id: str | None = None, + ) -> dict[str, Any]: + """Schedule a compute task that runs now instead of in serial order. + + Use for a backend that batches concurrent requests itself, such as the + vLLM sampler: through the serial queue it would only ever see one + request at a time. Preflight checks, the persisted future, the execution + timeout and failure mapping are the same as :meth:`schedule_task`. + + Returns: + {'request_id': str, 'model_id': str | None} + """ + request_id = request_id or f'req_{uuid.uuid4().hex}' + + # Decision_Boundary: raises RequestRejectedError before any state write. + await self._perform_preflight_checks(model_id=model_id, token=token, input_tokens=input_tokens) + + await self._compute_worker.submit_detached( + QueuedTask( + request_id=request_id, + coro_factory=coro_factory, + model_id=model_id, + token=token, + input_tokens=input_tokens, + task_type=task_type, + created_at=time.monotonic()), + self._queue_key(model_id=model_id, token=token), + replica_id=getattr(self, 'replica_id', None), + absolute_deadline=self._future_deadline()) + logger.info(f'[TaskQueue] Task {request_id} started concurrently, type={task_type or "unknown"}, ' + f'model_id={model_id}, input_tokens={input_tokens}') + return {'request_id': request_id, 'model_id': model_id} + # Poll cadence *inside* the Inline_Fast_Path window. Much smaller than the window # itself so a task that finishes early is noticed promptly; the window bound # (config.inline_fast_path_timeout) is what actually caps submit latency. _INLINE_FAST_PATH_POLL = 0.005 + async def submit_concurrent_and_peek(self, coro_factory: Callable[[], Coroutine], + **schedule_kwargs) -> TaskEnvelope: + """Concurrent compute submission with the existing Twinkle future envelope.""" + ref = await self.schedule_concurrent_task(coro_factory, **schedule_kwargs) + return await self._peek_terminal(ref['request_id'], fallback_status='pending') + async def _peek_terminal(self, request_id: str, *, fallback_status: str) -> TaskEnvelope: """Poll a record up to the Inline_Fast_Path window; return a terminal envelope if it settled, else a non-terminal envelope with ``fallback_status``.""" diff --git a/src/twinkle/server/task_queue/types.py b/src/twinkle/server/task_queue/types.py index 9085ff8ec..936a433a3 100644 --- a/src/twinkle/server/task_queue/types.py +++ b/src/twinkle/server/task_queue/types.py @@ -9,10 +9,18 @@ """ from __future__ import annotations +import asyncio from collections.abc import Callable, Coroutine from dataclasses import dataclass from enum import Enum +# Ray_Get_Timeout is classified the same as asyncio.TimeoutError: 504/Server. +try: + from ray.exceptions import GetTimeoutError as _RayGetTimeout + BACKEND_TIMEOUT_EXCEPTIONS: tuple[type[BaseException], ...] = (asyncio.TimeoutError, _RayGetTimeout) +except Exception: # pragma: no cover - ray always present in server runtime + BACKEND_TIMEOUT_EXCEPTIONS = (asyncio.TimeoutError, ) + class TaskStatus(Enum): """Task lifecycle status.""" diff --git a/src/twinkle/server/task_queue/worker.py b/src/twinkle/server/task_queue/worker.py index 8d07fd4a3..439f0c3e6 100644 --- a/src/twinkle/server/task_queue/worker.py +++ b/src/twinkle/server/task_queue/worker.py @@ -17,15 +17,15 @@ from twinkle.protocol.types.errors import ErrorCategory from twinkle.server.exceptions import (BatchSizeError, EndpointUnavailableError, FullModeBusyError, InputTokensExceededError, RateLimitExceededError, RequestRejectedError, - ResourceNotFoundError, ResourceQuotaExceededError, StateBackendError, - TrainModeMismatchError, TwinkleServerError) + ResourceExhaustedError, ResourceNotFoundError, ResourceQuotaExceededError, + StateBackendError, TrainModeMismatchError, TwinkleServerError) from twinkle.server.state.backend.base import ConcurrencyError from twinkle.server.state.models import FutureFailureRecord from twinkle.server.telemetry.correlation import MODEL_ID, TOKEN_ID from twinkle.server.telemetry.tracing import traced_operation from twinkle.utils.logger import get_logger from .config import TaskQueueConfig -from .types import BackendBusyError, QueuedTask, QueueState, TaskStatus, UserTaskError +from .types import BACKEND_TIMEOUT_EXCEPTIONS, BackendBusyError, QueuedTask, QueueState, TaskStatus, UserTaskError if TYPE_CHECKING: from twinkle.server.state import ServerState @@ -55,14 +55,6 @@ def _reason_code_for_server_error(exc: TwinkleServerError) -> str: return 'internal_error' -# Ray_Get_Timeout is classified the same as asyncio.TimeoutError: 504/Server. -try: - from ray.exceptions import GetTimeoutError as _RayGetTimeout - _TIMEOUT_EXCEPTIONS: tuple[type[BaseException], ...] = (asyncio.TimeoutError, _RayGetTimeout) -except Exception: # pragma: no cover - ray always present in server runtime - _TIMEOUT_EXCEPTIONS = (asyncio.TimeoutError, ) - - class ComputeWorker: """Serial background worker that processes GPU compute tasks. @@ -70,9 +62,10 @@ class ComputeWorker: queues so that no single user's long-running task (e.g. save/load) can starve other users waiting in different queues. - Only one task is executed at a time to preserve serial GPU semantics. - Upload and other pure-I/O tasks should NOT be submitted here; use the - background-task path in TaskQueueMixin instead. + Only one queued task is executed at a time to preserve serial GPU semantics; + :meth:`run_detached` runs a task outside that order. Upload and other + pure-I/O tasks should NOT be submitted here; use the background-task path + in TaskQueueMixin instead. """ def __init__( @@ -98,6 +91,10 @@ def __init__( self._started = False self._stopping = False self._start_lock = asyncio.Lock() + # Strong references to tasks started by run_detached(); the event loop + # only keeps weak ones. + self._detached_tasks: dict[asyncio.Task, QueuedTask] = {} + self._detached_reserved = 0 async def ensure_started(self) -> None: """Ensure the background worker coroutine is running.""" @@ -111,7 +108,7 @@ async def ensure_started(self) -> None: self._started = True async def stop(self) -> None: - """Cancel the worker and wait for it to exit cleanly.""" + """Cancel the worker and any detached tasks, and wait for them to exit.""" self._stopping = True if self._worker_task and not self._worker_task.done(): self._worker_task.cancel() @@ -119,6 +116,17 @@ async def stop(self) -> None: await self._worker_task except asyncio.CancelledError: pass + detached = list(self._detached_tasks.items()) + for running, _ in detached: + running.cancel() + await asyncio.gather(*(running for running, _ in detached), return_exceptions=True) + for running, task in detached: + if running.cancelled() and not await self._is_record_terminal(task.request_id): + await self._store_task_failed( + task, + 'Sampler task stopped during replica shutdown.', + QueueState.ACTIVE.value, + reason_code='orphaned_replica') self._worker_task = None self._started = False self.task_queues.clear() @@ -265,18 +273,69 @@ async def _fail_timed_out_task(self, task: QueuedTask, waited: float, q: asyncio q.task_done() async def _execute_task(self, task: QueuedTask, queue_key: str, q: asyncio.Queue) -> None: - """Execute a single task: update status, run coroutine, record metrics. + """Execute a dequeued task, then release its queue slot.""" + try: + await self._process_task(task, queue_key) + finally: + q.task_done() - Handles execution timeout, general exceptions, and always calls - q.task_done() in the finally block. + async def submit_detached(self, task: QueuedTask, queue_key: str, *, replica_id: str | None, + absolute_deadline: float) -> None: + """Own admission, future persistence and startup as one lifecycle.""" + self.reserve_detached() + wrote_status = False + try: + await self._state.store_future_status( + task.request_id, + TaskStatus.PENDING.value, + task.model_id, + queue_state=QueueState.ACTIVE.value, + replica_id=replica_id, + absolute_deadline=absolute_deadline) + wrote_status = True + self.run_detached(task, queue_key) + except BaseException: + self.release_detached() + if wrote_status: + await self._store_task_failed( + task, + 'Replica stopped before the concurrent task could start.', + QueueState.ACTIVE.value, + reason_code='orphaned_replica') + raise + + def reserve_detached(self) -> None: + """Reserve before the future write, so capacity rejection leaves no record.""" + if self._stopping or self._detached_reserved >= self._config.concurrent_task_limit: + raise ResourceExhaustedError('Concurrent compute capacity is unavailable on this replica.') + self._detached_reserved += 1 + + def release_detached(self) -> None: + self._detached_reserved -= 1 + + def run_detached(self, task: QueuedTask, queue_key: str) -> None: + """Start ``task`` now, beside the serial loop instead of in its turn. + + For backends that batch concurrent requests themselves (vLLM). The task + gets the same status records, execution timeout and failure mapping as a + queued one. """ - await self._state.store_future_status( - task.request_id, - TaskStatus.RUNNING.value, - task.model_id, - queue_state=QueueState.ACTIVE.value, - ) + if self._stopping: + raise ResourceExhaustedError('The replica is shutting down.') + running = asyncio.create_task(self._process_task(task, queue_key)) + self._detached_tasks[running] = task + + def done(completed: asyncio.Task) -> None: + self._detached_tasks.pop(completed, None) + self.release_detached() + + running.add_done_callback(done) + async def _process_task(self, task: QueuedTask, queue_key: str) -> None: + """Run a single task: update status, run coroutine, record metrics. + + Handles execution timeout and general exceptions. + """ task_type = task.task_type or 'unknown' exec_start = time.monotonic() task_status = 'completed' @@ -293,6 +352,12 @@ async def _execute_task(self, task: QueuedTask, queue_key: str, q: asyncio.Queue handler_attrs = {TOKEN_ID: task.token, MODEL_ID: task.model_id} handler_span_name = f'{self._deployment_name or "deployment"}.{task_type}' try: + await self._state.store_future_status( + task.request_id, + TaskStatus.RUNNING.value, + task.model_id, + queue_state=QueueState.ACTIVE.value, + ) with traced_operation('task_queue.execute', attrs=queue_attrs): logger.debug(f'[ComputeWorker] Task {task.request_id} executing, ' f'type={task_type}, queue_key={queue_key}') @@ -310,7 +375,17 @@ async def _execute_task(self, task: QueuedTask, queue_key: str, q: asyncio.Queue result=result, queue_state=QueueState.ACTIVE.value, ) - except _TIMEOUT_EXCEPTIONS: + except asyncio.CancelledError: + task_status = 'cancelled' + exec_time = time.monotonic() - exec_start + if not await self._is_record_terminal(task.request_id): + await self._store_task_failed( + task, + 'Task stopped during replica shutdown.' if self._stopping else 'Backend generation was cancelled.', + QueueState.ACTIVE.value, + reason_code='orphaned_replica' if self._stopping else 'backend_cancelled') + raise + except BACKEND_TIMEOUT_EXCEPTIONS: task_status = 'timeout' exec_time = time.monotonic() - exec_start error = (f'Backend call timed out (bound {self._config.effective_execution_timeout}s), ' @@ -380,7 +455,6 @@ async def _execute_task(self, task: QueuedTask, queue_key: str, q: asyncio.Queue traceback_text=traceback.format_exc(), ) finally: - q.task_done() self._record_execution_time(task_type, exec_time) self._record_tasks_total(task_type, task_status) diff --git a/src/twinkle_agentic/async_rl/vllm_sampler_tq.py b/src/twinkle_agentic/async_rl/vllm_sampler_tq.py index 4f2d0342d..657fa0300 100644 --- a/src/twinkle_agentic/async_rl/vllm_sampler_tq.py +++ b/src/twinkle_agentic/async_rl/vllm_sampler_tq.py @@ -17,7 +17,7 @@ from twinkle.data_format import SampledSequence, SampleResponse, SamplingParams, user_data_get from twinkle.hub import HubOperation from twinkle.metric import MetricBuffer, MetricRecord -from twinkle.sampler.vllm_sampler import vLLMSampler +from twinkle.sampler.vllm_sampler.vllm_sampler import SAMPLER_MAX_CONCURRENCY, vLLMSampler from .data_plane import TQDataPlane from .metrics import rollout_metrics from .types import LoraContext, PromptGroup, RolloutOutput, RolloutPolicy @@ -85,6 +85,14 @@ def _compute_reward_metrics( return dict(metric_payload(rollout_rows, rewards=rewards, context=context)) +@dataclass +class _GenerationSubmission: + """Separate a cancellation request from actual coroutine completion.""" + future: Future | None = None + task: asyncio.Task | None = None + cancel_requested: bool = False + + @dataclass(frozen=True) class _GeneratedSample: response: SampleResponse @@ -113,13 +121,14 @@ class _PromptGroupRolloutStats: policy_versions: tuple[int, ...] -@remote_class() +@remote_class(max_concurrency=SAMPLER_MAX_CONCURRENCY) class VLLMSamplerTQ(vLLMSampler): - """vLLM sampler that writes async RL rollout results directly to TransferQueue. + """vLLM sampler for native async RL and client-server generation. - ``sample()`` is intentionally fire-and-forget: it schedules generation work - on the sampler actor's vLLM event loop and returns submission metadata - without waiting for any prompt group to finish. + ``submit_prompt_groups()`` schedules native rollouts and returns metadata; + their results go to TransferQueue. Inherited ``sample()`` returns completed + responses. The client-server path uses ``submit_generation()``, status + polling and ``collect_generation()`` for non-blocking admission. """ def __init__( @@ -157,7 +166,7 @@ def __init__( # path. Unlike ``_background_submissions`` above, their results must # remain available until SamplerManagement collects them and writes # them to the opaque client DataPlane. - self._generation_submissions: dict[str, Future[list[SampleResponse]]] = {} + self._generation_submissions: dict[str, _GenerationSubmission] = {} self.metric_buffer = MetricBuffer() self._failure: str | None = None @@ -249,27 +258,34 @@ def submit_generation( """ if submission_id in self._generation_submissions: raise KeyError(f'generation submission already exists: {submission_id}') - future = self._submit_in_loop( - self._generate_inputs( + submission = _GenerationSubmission() + + async def generate(): + submission.task = asyncio.current_task() + if submission.cancel_requested: + raise asyncio.CancelledError + return await self._generate_inputs( inputs, sampling_params, adapter_name=adapter_name, adapter_path=adapter_path, - use_base_model=use_base_model, - )) - self._generation_submissions[submission_id] = future + use_base_model=use_base_model) + + submission.future = self._submit_in_loop(generate()) + self._generation_submissions[submission_id] = submission return {'submission_id': submission_id, 'status': 'running'} @remote_function(dispatch='all', collect='none', lazy_collect=False) def get_generation_status(self, submission_id: str) -> dict[str, Any]: """Return this DP worker's submission state without waiting.""" - future = self._generation_submissions.get(submission_id) - if future is None: + submission = self._generation_submissions.get(submission_id) + if submission is None: return { 'submission_id': submission_id, 'status': 'missing', 'error': f'unknown generation submission: {submission_id}', } + future = submission.future if future.cancelled(): return {'submission_id': submission_id, 'status': 'cancelled'} if not future.done(): @@ -286,9 +302,10 @@ def get_generation_status(self, submission_id: str) -> dict[str, Any]: @remote_function(dispatch='all', collect='flatten', lazy_collect=False) def collect_generation(self, submission_id: str) -> list[SampleResponse]: """Consume completed responses from every DP worker.""" - future = self._generation_submissions.get(submission_id) - if future is None: + submission = self._generation_submissions.get(submission_id) + if submission is None: raise KeyError(f'unknown generation submission: {submission_id}') + future = submission.future if not future.done(): raise RuntimeError(f'generation submission is still running: {submission_id}') try: @@ -298,30 +315,34 @@ def collect_generation(self, submission_id: str) -> list[SampleResponse]: @remote_function(dispatch='all', collect='none', lazy_collect=False) def cancel_generation(self, submission_id: str) -> dict[str, Any]: - """Cancel and forget one generation submission on every DP worker.""" - future = self._generation_submissions.pop(submission_id, None) - if future is None: + """Request cancellation; forget only after vLLM's coroutine has exited. + + Cancelling run_coroutine_threadsafe's Future reports success before the + coroutine's abort/finally runs. Cancel the loop-owned Task instead so the + Future becomes terminal only after actual cleanup. + """ + submission = self._generation_submissions.get(submission_id) + if submission is None: return {'submission_id': submission_id, 'status': 'missing'} - was_done = future.done() - cancelled = future.cancel() - if cancelled: - status = 'cancelled' - elif was_done: - status = 'completed' - else: - status = 'cancellation_requested' - return { - 'submission_id': submission_id, - 'status': status, - } + if submission.future.done(): + self._generation_submissions.pop(submission_id, None) + return {'submission_id': submission_id, 'status': 'completed'} + if not submission.cancel_requested: + submission.cancel_requested = True + + def cancel(): + if submission.task is not None and not submission.task.done(): + submission.task.cancel() + + self._async_loop.call_soon_threadsafe(cancel) + return {'submission_id': submission_id, 'status': 'cancellation_requested'} @remote_function(dispatch='all', collect='none', lazy_collect=False) def cancel_all_generations(self) -> dict[str, int]: - """Cancel all retained CS submissions during replica shutdown.""" - submissions = list(self._generation_submissions.values()) - self._generation_submissions.clear() - cancelled = sum(future.cancel() for future in submissions if not future.done()) - return {'submissions': len(submissions), 'cancelled': cancelled} + """Request cancellation of retained submissions during shutdown.""" + ids = list(self._generation_submissions) + requested = sum(self.cancel_generation(key)['status'] == 'cancellation_requested' for key in ids) + return {'submissions': len(ids), 'cancelled': requested} @remote_function(dispatch='slice_dp', collect='flatten', lazy_collect=False) def evaluate( @@ -388,14 +409,15 @@ async def _generate_inputs( logger.warning(f'Failed to pre-load LoRA from {local_adapter_path}, ' 'sampling will proceed without LoRA') - return await asyncio.gather(*(self._sample_single( - feat, - sampling_params, - lora_request=lora_request, - multi_modal_data=multi_modal_data, - logprobs_only=logprobs_only, - disable_lora=use_base_model, - ) for feat, multi_modal_data in zip(encoded_inputs, multi_modal_data_list))) + return await self._gather_samples( + self._sample_single( + feat, + sampling_params, + lora_request=lora_request, + multi_modal_data=multi_modal_data, + logprobs_only=logprobs_only, + disable_lora=use_base_model, + ) for feat, multi_modal_data in zip(encoded_inputs, multi_modal_data_list)) def _on_submission_done(self, submission_id: str): diff --git a/tests/sampler/test_vllm_sampling_lifecycle.py b/tests/sampler/test_vllm_sampling_lifecycle.py new file mode 100644 index 000000000..abe5c1bb8 --- /dev/null +++ b/tests/sampler/test_vllm_sampling_lifecycle.py @@ -0,0 +1,36 @@ +"""A batch owns generation until all of its prompts have exited.""" +import pytest + + +@pytest.mark.asyncio +async def test_failed_prompt_waits_for_sibling_generation_to_exit(): + import asyncio + from twinkle.data_format import SamplingParams + from twinkle.sampler.vllm_sampler import vLLMSampler + sampler = object.__new__(vLLMSampler) + sampler.template = None + sampler._run_in_loop = lambda coroutine: coroutine + entered = asyncio.Event() + release = asyncio.Event() + finished = asyncio.Event() + async def single(feat, *args, **kwargs): + if feat['input_ids'][0] == 1: + raise ValueError('bad prompt') + entered.set() + await release.wait() + finished.set() + return object() + sampler._sample_single = single + generation = asyncio.create_task(sampler.sample( + [{'input_ids': [1]}, {'input_ids': [2]}], SamplingParams(max_tokens=1))) + try: + await entered.wait() + await asyncio.sleep(.01) + assert not generation.done(), 'a failed prompt returned while its sibling was still generating' + release.set() + with pytest.raises(ValueError, match='bad prompt'): + await generation + assert finished.is_set() + finally: + release.set() + await asyncio.gather(generation, return_exceptions=True) diff --git a/tests/server/integration/sampler_e2e_helpers.py b/tests/server/integration/sampler_e2e_helpers.py new file mode 100644 index 000000000..f09f80f6a --- /dev/null +++ b/tests/server/integration/sampler_e2e_helpers.py @@ -0,0 +1,191 @@ +"""Shared clients and model fixtures for real sampler E2E tests.""" +from __future__ import annotations + +import concurrent.futures +import json +import math +import os +import threading +import time +from pathlib import Path + +import pytest + +@pytest.fixture(scope='session', autouse=True) +def _ray_runtime(): + yield + + + +@pytest.fixture(autouse=True) +def _reset_canonical_state_actor(): + yield + + + +def _report(**values): + print('SAMPLER_CONCURRENCY ' + json.dumps(values, sort_keys=True), flush=True) + + + +@pytest.fixture(scope='module') +def clients(): + from transformers import AutoTokenizer + from twinkle import init_tinker_client, init_twinkle_client + init_tinker_client() + from tinker import ServiceClient + model_id = os.environ['TWINKLE_TEST_MODEL_ID'] + model_path = os.environ['TWINKLE_TEST_MODEL_PATH'] + url = os.environ['TWINKLE_SERVER_URL'] + tokenizer = AutoTokenizer.from_pretrained(model_path, local_files_only=True) + service = ServiceClient(base_url=url, api_key='EMPTY_TOKEN') + native = init_twinkle_client(base_url=url, api_key='EMPTY_TOKEN').sampler(model_id) + return { + 'tokenizer': tokenizer, 'service': service, 'model_id': model_id, 'model_path': model_path, + 'url': url, 'tinker': service.create_sampling_client(base_model=model_id), 'twinkle': native, + 'single_throughput': {}, + } + + + +def _prompt(clients, text): + return clients['tokenizer'].apply_chat_template( + [{'role': 'user', 'content': text}], tokenize=True, add_generation_prompt=True, + enable_thinking=False, return_dict=False) + + + +def _sample(clients, dialect, ids, max_tokens=128, *, adapter=None, num_samples=1, temperature=0): + if dialect == 'tinker': + from tinker import types + result = (adapter or clients['tinker']).sample( + prompt=types.ModelInput.from_ints(ids), num_samples=num_samples, + sampling_params=types.SamplingParams(max_tokens=max_tokens, temperature=temperature, top_p=1), + ).result(timeout=300) + sequences = [(list(seq.tokens), list(seq.logprobs)) for seq in result.sequences] + else: + result = clients['twinkle'].sample( + inputs=[{'input_ids': ids}], num_samples=num_samples, adapter_uri=adapter, + sampling_params={'max_tokens': max_tokens, 'temperature': temperature, 'top_p': 1, 'logprobs': 1}) + assert result[0].prompt_token_ids == ids + sequences = [(list(seq.tokens), [position[0][1] for position in seq.logprobs]) + for seq in result[0].sequences] + assert len(sequences) == num_samples + for tokens, logprobs in sequences: + assert tokens and len(tokens) == len(logprobs) + assert all(math.isfinite(value) for value in logprobs) + return sequences + + + +def _parallel(jobs): + barrier = threading.Barrier(len(jobs)) + def run(job): + barrier.wait(timeout=30) + return job() + started = time.perf_counter() + with concurrent.futures.ThreadPoolExecutor(max_workers=len(jobs)) as pool: + results = list(pool.map(run, jobs)) + return results, time.perf_counter() - started + + + +@pytest.fixture(scope='module') +def adapters(clients): + import torch + from peft import LoraConfig + from safetensors import safe_open + from safetensors.torch import save_file + projections = {} + for shard in Path(clients['model_path']).glob('*.safetensors'): + with safe_open(str(shard), framework='pt', device='cpu') as weights: + for name in weights.keys(): + if name.endswith('.self_attn.q_proj.weight'): + projections[name[:-len('.weight')]] = weights.get_slice(name).get_shape() + if not projections: + pytest.skip('The cached model has no q_proj layers for this LoRA fixture') + root = Path(os.environ['TWINKLE_TEST_ARTIFACT_DIR']) / 'lora-isolation' + paths = [] + rank = 8 + for index in range(2): + path = root / str(index) + path.mkdir(parents=True, exist_ok=True) + LoraConfig(r=rank, lora_alpha=rank, target_modules=['q_proj'], bias='none', + task_type='CAUSAL_LM', inference_mode=True, + base_model_name_or_path=clients['model_id']).save_pretrained(path) + generator = torch.Generator().manual_seed(100 + index) + weights = {} + for name, (output_size, input_size) in projections.items(): + prefix = 'base_model.model.' + name + weights[prefix + '.lora_A.weight'] = torch.randn(rank, input_size, generator=generator) * 0.05 + weights[prefix + '.lora_B.weight'] = torch.randn(output_size, rank, generator=generator) * 0.05 + save_file(weights, str(path / 'adapter_model.safetensors')) + paths.append(str(path)) + return paths + + + +@pytest.fixture(scope='module') +def full_checkpoints(clients): + """Two complete HF checkpoints with different, observable model weights.""" + import torch + from safetensors.torch import load_file, save_file + source = Path(clients['model_path']) + target = Path(os.environ['TWINKLE_TEST_ARTIFACT_DIR']) / 'full-checkpoint-B' + target.mkdir(parents=True, exist_ok=True) + for file in source.iterdir(): + if file.is_file() and not (target / file.name).exists(): + (target / file.name).symlink_to(file) + index = json.loads((source / 'model.safetensors.index.json').read_text()) + shard = index['weight_map']['model.embed_tokens.weight'] + output = target / shard + if output.is_symlink(): + weights = load_file(str(source / shard), device='cpu') + weights['model.embed_tokens.weight'] = torch.roll(weights['model.embed_tokens.weight'], 257, 0).contiguous() + output.unlink() + save_file(weights, str(output), metadata={'format': 'pt'}) + return [str(source), str(target)] + + + +def _checkpoint_handles(clients, dialect, paths): + if dialect == 'tinker': + return [clients['service'].create_sampling_client(base_model=clients['model_id'], model_path=path) + for path in paths] + return paths + + + +def _assert_adapter_isolation(clients, adapters, dialect): + handles = ([None] + [clients['service'].create_sampling_client(base_model=clients['model_id'], model_path=path) + for path in adapters] if dialect == 'tinker' else [None, *adapters]) + ids = _prompt(clients, 'Choose one word: apple, banana, orange, purple.') + solo = [_sample(clients, dialect, ids, 8, adapter=handle)[0] for handle in handles] + # Default vLLM kernels do not guarantee batch-invariant continuations or + # logprobs. Homogeneous PPU batches also diverge after this stable prefix. + # Require distinct adapter fingerprints, then check routing against them; + # _sample still validates every generated token's finite, aligned logprob. + fingerprints = [tuple(tokens[:2]) for tokens, _logprobs in solo] + assert all(len(fingerprint) == 2 for fingerprint in fingerprints) + assert len(set(fingerprints)) == 3, 'the adapters must change the result, rather than silently falling back to base' + results, elapsed = _parallel([ + lambda i=i: _sample(clients, dialect, ids, 8, adapter=handles[i % 3]) for i in range(24)]) + for i, result in enumerate(results): + tokens, _logprobs = result[0] + assert tuple(tokens[:2]) == fingerprints[i % 3] + _report(check='lora_isolation', dialect=dialect, requests=24, variants=3, seconds=elapsed, + fingerprints=fingerprints) + + +def _real_sampler_actors(): + import ray + from ray.util.state import list_actors + ray.init(address=os.environ['RAY_ADDRESS'], namespace='default', ignore_reinit_error=True, + log_to_driver=False) + actors = list_actors(address=os.environ['TWINKLE_TEST_RAY_DASHBOARD_URL'], + filters=[('class_name', '=', os.environ.get('TWINKLE_TEST_SAMPLER_CLASS', 'VLLMSamplerTQ')), + ('state', '=', 'ALIVE')], + detail=True) + assert len(actors) == int(os.environ.get('TWINKLE_TEST_DP', '1')) + assert all(actor.state == 'ALIVE' for actor in actors) + return [ray.get_actor(actor.name, namespace='default') for actor in actors] diff --git a/tests/server/integration/test_sampler_concurrency_e2e.py b/tests/server/integration/test_sampler_concurrency_e2e.py new file mode 100644 index 000000000..2193b9b0b --- /dev/null +++ b/tests/server/integration/test_sampler_concurrency_e2e.py @@ -0,0 +1,279 @@ +"""Real Tinker/Twinkle concurrency, weight isolation and lifecycle checks.""" +from __future__ import annotations + +import concurrent.futures +import json +import math +import os +import threading +import time +from contextlib import contextmanager +from pathlib import Path + +import pytest + +from tests.server.integration.sampler_e2e_helpers import ( + _assert_adapter_isolation, _checkpoint_handles, _parallel, _prompt, _ray_runtime, + _report, _reset_canonical_state_actor, _sample, _real_sampler_actors, adapters, clients, full_checkpoints, +) + +pytestmark = pytest.mark.skipif(os.environ.get('TWINKLE_TEST_GPU_E2E', '0') != '1', + reason='Requires a running real vLLM sampler server') + + +@pytest.mark.parametrize('count', [1, 8, 32]) +@pytest.mark.parametrize('dialect', ['tinker', 'twinkle']) +def test_both_dialects_batch_independent_requests(clients, count, dialect): + ids = _prompt(clients, 'Write a detailed story of at least 300 words about an astronaut. Keep writing.') + _sample(clients, dialect, ids, 8) + if count != 1 and dialect not in clients['single_throughput']: + started = time.perf_counter() + single = _sample(clients, dialect, ids) + clients['single_throughput'][dialect] = len(single[0][0]) / (time.perf_counter() - started) + results, elapsed = _parallel([lambda: _sample(clients, dialect, ids)] * count) + tokens = sum(len(result[0][0]) for result in results) + throughput = tokens / elapsed + _report(check='benchmark', dialect=dialect, count=count, tokens=tokens, seconds=elapsed, + tokens_per_second=throughput) + if count == 1: + clients['single_throughput'][dialect] = throughput + else: + assert throughput > 1.5 * clients['single_throughput'][dialect], 'independent requests still appear serial' + + + +@pytest.mark.parametrize('dialect', ['tinker', 'twinkle']) +def test_concurrent_request_mapping(clients, dialect): + words = ['apple', 'banana', 'orange', 'purple', 'yellow', 'green', 'blue', 'red'] + prompts = [_prompt(clients, 'Reply with exactly this word and nothing else: ' + word) for word in words] + solo = [_sample(clients, dialect, ids, 16)[0][0] for ids in prompts] + jobs = [lambda i=i: _sample(clients, dialect, prompts[i % len(prompts)], 16) for i in range(32)] + results, elapsed = _parallel(jobs) + assert all(result[0][0] == solo[i % len(solo)] for i, result in enumerate(results)) + _report(check='mapping', dialect=dialect, count=32, unique_prompts=8, seconds=elapsed, solo_equal=True) + + + +def test_mixed_dialects_and_multiple_samples(clients): + ids = _prompt(clients, 'Reply with exactly this word and nothing else: apple') + solo = _sample(clients, 'tinker', ids, 16)[0][0] + jobs = [lambda i=i: _sample(clients, 'tinker' if i % 2 else 'twinkle', ids, 16) + for i in range(16)] + results, elapsed = _parallel(jobs) + assert all(tokens == solo for result in results for tokens, _logprobs in result) + _report(check='mixed', requests=16, sequences=16, seconds=elapsed, solo_equal=True) + # vLLM requires n=1 for greedy sampling. Multiple completions exercise the + # stochastic contract: count, token/logprob alignment and finite logprobs. + results, elapsed = _parallel([ + lambda i=i: _sample(clients, 'tinker' if i % 2 else 'twinkle', ids, 16, + num_samples=4, temperature=0.7) for i in range(16)]) + assert sum(len(result) for result in results) == 64 + _report(check='mixed_multiple_samples', requests=16, sequences=64, seconds=elapsed) + + + +@pytest.mark.parametrize('dialect', ['tinker', 'twinkle']) +def test_base_and_two_lora_adapters_are_isolated(clients, adapters, dialect): + _assert_adapter_isolation(clients, adapters, dialect) + + +def test_full_checkpoints_are_isolated_across_dialects(clients, full_checkpoints): + ids = _prompt(clients, 'Reply with exactly this word and nothing else: apple') + tinker = _checkpoint_handles(clients, 'tinker', full_checkpoints) + solo = [_sample(clients, 'tinker', ids, 8, adapter=handle)[0][0] for handle in tinker] + assert solo[0] != solo[1] + try: + jobs = [lambda i=i: _sample(clients, 'tinker' if i % 2 else 'twinkle', ids, 8, + adapter=tinker[i % 2] if i % 2 else full_checkpoints[i % 2]) for i in range(8)] + results, elapsed = _parallel(jobs) + assert all(result[0][0] == solo[i % 2] for i, result in enumerate(results)) + _report(check='full_checkpoint_mixed_isolation', requests=8, seconds=elapsed, a=solo[0], b=solo[1]) + finally: + _sample(clients, 'twinkle', ids, 2, adapter=full_checkpoints[0]) + + + +@pytest.mark.parametrize('dialect', ['tinker', 'twinkle']) +def test_invalid_generation_does_not_disable_sampler(clients, dialect): + ids = _prompt(clients, 'Reply with exactly this word and nothing else: apple') + with pytest.raises(Exception): + _sample(clients, dialect, [10**9], 4) + assert _sample(clients, dialect, ids, 4)[0][0] + _report(check='failure_recovery', dialect=dialect) + + + +def test_async_generation_cleanup_and_recovery_on_real_engine(clients): + import ray + from ray.util.state import list_actors + ray.init(address=os.environ['RAY_ADDRESS'], namespace='default', ignore_reinit_error=True, + log_to_driver=False) + actors = list_actors(address=os.environ.get('TWINKLE_TEST_RAY_DASHBOARD_URL', 'http://127.0.0.1:8395'), + filters=[('class_name', '=', 'VLLMSamplerTQ'), ('state', '=', 'ALIVE')], detail=True) + assert len(actors) == int(os.environ.get('TWINKLE_TEST_DP', '1')) + assert all(actor.state == 'ALIVE' for actor in actors) + actor = ray.get_actor(actors[0].name, namespace='default') + ids = _prompt(clients, 'Write a very long story of at least 2000 words about an astronaut.') + def call(method, *args): + return ray.get(getattr(actor, method).remote(*args), timeout=30) + def wait_status(key, terminal): + deadline = time.monotonic() + 30 + while time.monotonic() < deadline: + state = call('get_generation_status', key) + if state['status'] in terminal: + return state + time.sleep(.01) + raise AssertionError(f'{key} did not terminate: {state}') + try: + call('submit_generation', 'e2e-bad', [{'input_ids': [10**9]}], + {'max_tokens': 8, 'temperature': 0}) + failed = wait_status('e2e-bad', {'failed'}) + assert call('cancel_generation', 'e2e-bad')['status'] == 'completed' + # A failed prompt in ordinary sample() must wait for its live sibling. + batch = actor.sample.remote( + [{'input_ids': ids}, {'input_ids': [10**9]}], + {'max_tokens': 1024, 'temperature': 0}) + ready, _ = ray.wait([batch], timeout=.2) + assert not ready, 'ordinary batch returned before its sibling generation exited' + with pytest.raises(ray.exceptions.RayTaskError): + ray.get(batch, timeout=30) + + call('submit_generation', 'e2e-cancel', [{'input_ids': ids}], + {'max_tokens': 1024, 'temperature': 0}) + time.sleep(.1) + assert call('get_generation_status', 'e2e-cancel')['status'] == 'running' + assert call('cancel_generation', 'e2e-cancel')['status'] == 'cancellation_requested' + wait_status('e2e-cancel', {'cancelled'}) + assert call('cancel_generation', 'e2e-cancel')['status'] == 'completed' + call('submit_generation', 'e2e-after', [{'input_ids': ids}], + {'max_tokens': 8, 'temperature': 0}) + wait_status('e2e-after', {'completed'}) + responses = call('collect_generation', 'e2e-after') + assert responses[0].sequences[0].tokens + _report(check='async_cleanup_recovery', failure=failed, recovered=True) + finally: + for key in ['e2e-bad', 'e2e-cancel', 'e2e-after']: + call('cancel_generation', key) + ray.shutdown() + + + + + + + + + + + + + + + +def test_data_plane_http_generation_and_failure_recovery(clients): + from twinkle_client.data_plane import DataPlaneClient + store = DataPlaneClient(transport=clients['twinkle']._transport) + ids = _prompt(clients, 'Reply with exactly this word and nothing else: apple') + expected = _sample(clients, 'twinkle', ids, 2)[0][0] + with pytest.raises(Exception): + clients['twinkle'].sample_to_data_plane( + inputs=[{'input_ids': [10**9]}], sampling_params={'max_tokens': 2, 'logprobs': 1}) + + def generate_and_fetch(): + ref = clients['twinkle'].sample_to_data_plane( + inputs=[{'input_ids': ids}], num_samples=2, group_ids=['e2e-group'], policy_version=7, + sampling_params={'max_tokens': 2, 'temperature': .7, 'logprobs': 1}) + try: + batch = store.get_batch(ref) + assert ref.size == 2 and len(batch.rows) == len(batch.tags) == 2 + assert all(0 < len(row['tokens']) <= 2 for row in batch.rows) + assert all(len(row['sampled_logprobs']) == len(row['tokens']) + and all(math.isfinite(value) for value in row['sampled_logprobs']) for row in batch.rows) + assert [tag['generation_idx'] for tag in batch.tags] == [0, 1] + assert all(tag['group_id'] == 'e2e-group' and tag['rollout_policy_version'] == 7 for tag in batch.tags) + return len(batch.rows) + finally: + store.release(ref) + + results, elapsed = _parallel([ + generate_and_fetch, generate_and_fetch, + lambda: _sample(clients, 'tinker', ids, 2)[0][0], + lambda: _sample(clients, 'twinkle', ids, 2)[0][0], + ]) + assert results == [2, 2, expected, expected] + _report(check='data_plane_http', requests=4, sequences=6, recovered=True, seconds=elapsed) + + + + + + +def test_backend_self_cancellation_terminates_both_sdk_futures(clients): + """Cancel at the driver boundary after both requests performed real PPU generation. + + Submit both requests before cancellation quarantines the replica. This test + runs last because uncertain backend cancellation requires a replica restart. + """ + import ray + from ray.util.state import list_actors + actors = _real_sampler_actors() + replicas = [actor for actor in list_actors(address=os.environ['TWINKLE_TEST_RAY_DASHBOARD_URL'], detail=True) + if actor.state == 'ALIVE' and 'SamplerManagement' in actor.class_name + and actor.class_name.startswith('ServeReplica')] + assert len(replicas) == 1, [(actor.class_name, actor.name) for actor in replicas] + replica = ray.get_actor(replicas[0].name, namespace='serve') + + def install(actor): + from concurrent.futures import CancelledError + service = actor._replica_impl._user_callable_wrapper._callable + original = service.sampler.sample + barrier = threading.Barrier(2) + def cancelled(*args, **kwargs): + result = original(*args, **kwargs) + assert result[0].sequences[0].tokens + barrier.wait(timeout=15) + raise CancelledError('test: cancelled after real engine cleanup') + service.sampler._e2e_original_sample = original + service.sampler.sample = cancelled + + def restore(actor): + service = actor._replica_impl._user_callable_wrapper._callable + service.sampler.sample = service.sampler._e2e_original_sample + del service.sampler._e2e_original_sample + + def status(actor): + service = actor._replica_impl._user_callable_wrapper._callable + return {'jobs': len(service.sampling_guard._jobs), + 'detached': len(service._compute_worker._detached_tasks), + 'reserved': service._compute_worker._detached_reserved, + 'quarantined': service.sampling_guard._uncertain} + + ids = _prompt(clients, 'Reply with exactly this word and nothing else: apple') + ray.get(replica.__ray_call__.remote(install), timeout=10) + try: + def cancelled_sample(dialect): + started = time.monotonic() + with pytest.raises(Exception, match='Backend generation was cancelled') as error: + _sample(clients, dialect, ids, 2) + assert time.monotonic() - started < 15, 'SDK future did not receive a terminal failure promptly' + return {'dialect': dialect, 'error': str(error.value)} + results, elapsed = _parallel([lambda: cancelled_sample('tinker'), lambda: cancelled_sample('twinkle')]) + deadline = time.monotonic() + 5 + while True: + state = ray.get(replica.__ray_call__.remote(status), timeout=5) + if state['jobs'] == state['detached'] == state['reserved'] == 0: + break + assert time.monotonic() < deadline, state + time.sleep(.01) + assert state['quarantined'] + engine_requests = ray.get([actor.__ray_call__.remote( + lambda sampler: len(sampler.engine.engine.output_processor.request_states)) for actor in actors], timeout=10) + assert engine_requests == [0] * len(actors) + for dialect in ('tinker', 'twinkle'): + with pytest.raises(Exception): + _sample(clients, dialect, ids, 2) + _report(check='backend_self_cancellation', results=results, state=state, + engine_requests=engine_requests, seconds=elapsed) + finally: + ray.get(replica.__ray_call__.remote(restore), timeout=10) + ray.shutdown() diff --git a/tests/server/sampler/sampling_test_helpers.py b/tests/server/sampler/sampling_test_helpers.py new file mode 100644 index 000000000..953c235d8 --- /dev/null +++ b/tests/server/sampler/sampling_test_helpers.py @@ -0,0 +1,147 @@ +"""Shared route fixtures for sampler concurrency and stream lifecycle tests.""" +from __future__ import annotations + +import asyncio +import queue +import threading +from types import SimpleNamespace + +import pytest +from fastapi import FastAPI +from ray.exceptions import GetTimeoutError +from starlette.requests import Request +from tinker import types as tinker_types + +from twinkle.data_format import SampledSequence, SampleResponse +from twinkle.protocol.types import SampleRequest +from twinkle.server.sampler.app import SamplerManagement +from twinkle.server.sampler.concurrency import SamplerConcurrency +from twinkle.server.sampler.tinker_handlers import _register_tinker_sampler_routes +from twinkle.server.sampler.twinkle_handlers import _register_twinkle_sampler_routes +from twinkle.server.task_queue.backend_gate import BackendGate +from twinkle.server.task_queue.config import TaskQueueConfig +from twinkle.server.task_queue.types import BackendBusyError +from twinkle.server.task_queue.worker import ComputeWorker + + +async def _until(predicate): + async def wait(): + while not predicate(): + await asyncio.sleep(0.001) + await asyncio.wait_for(wait(), 3) + + + +class _State: + def __init__(self): + self.records = {} + + async def store_future_status(self, request_id, status, model_id=None, **kwargs): + self.records[request_id] = {'status': status, **kwargs} + + async def get_future(self, request_id): + return self.records.get(request_id) + + + +class _Limiter: + async def check_and_record(self, *args): + return True, None + + + +class _Sampler: + def __init__(self): + self.weight = 0 + self.loads = [] + self.active = 0 + self.peak = 0 + self.entered = threading.Event() + self.release = threading.Event() + self.release.set() + self._lock = threading.Lock() + self.template_calls = 0 + + def set_template(self, *args, **kwargs): + self.template_calls += 1 + + def load_full_weights_from_path(self, path): + self.weight = int(path.rsplit('/', 1)[-1]) + self.loads.append(self.weight) + + def sample(self, inputs, sampling_params=None, **kwargs): + with self._lock: + self.active += 1 + self.peak = max(self.peak, self.active) + self.entered.set() + try: + if not self.release.wait(3): + raise TimeoutError('test did not release generation') + token = self.weight or inputs[0]['input_ids'][0] + return [SampleResponse( + prompt_token_ids=list(item['input_ids']), + sequences=[SampledSequence(stop_reason='length', tokens=[token], logprobs=[[(token, -0.1)]])], + ) for item in inputs] + finally: + with self._lock: + self.active -= 1 + + + +class _Service(SamplerManagement): + def __init__(self, *, timeout=2, limit=128): + self.model_id = 'mock-model' + self.replica_id = 'sampler-test' + self.state = _State() + self.sampler = _Sampler() + self.sampling_guard = SamplerConcurrency(limit) + self.tinker_setup_lock = asyncio.Lock() + self.tinker_template_set = True + self._task_queue_config = TaskQueueConfig(execution_timeout=timeout, concurrent_task_limit=limit) + self._rate_limiter = _Limiter() + self._task_metrics = None + self._deployment_name = None + self._event_loop = None + self._backend_gate = BackendGate() + self._compute_worker = ComputeWorker( + state=self.state, config=self._task_queue_config, task_metrics=None, deployment_name=None) + app = FastAPI() + _register_tinker_sampler_routes(app, lambda: self) + _register_twinkle_sampler_routes(app, lambda: self) + self.endpoints = {route.path: route.endpoint for route in app.routes} + + async def _on_request_start(self, request): + return 'token' + + async def submit(self, dialect, token=7, path=None): + request = Request({'type': 'http', 'headers': []}) + if dialect == 'tinker': + body = tinker_types.SampleRequest( + prompt=tinker_types.ModelInput.from_ints([token]), model_path=path, + sampling_params=tinker_types.SamplingParams(max_tokens=1, temperature=0)) + ref = await self.endpoints['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/tinker/asample'](request, body, self) + return ref['request_id'] + body = SampleRequest(inputs=[{'input_ids': [token]}], adapter_uri=path, + sampling_params={'max_tokens': 1, 'temperature': 0}) + envelope = await self.endpoints['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/twinkle/sample'](request, body, self) + return envelope.request_id + + async def result(self, request_id): + await _until(lambda: self.state.records[request_id]['status'] in ('completed', 'failed')) + return self.state.records[request_id] + + async def close(self): + self.sampler.release.set() + self.sampling_guard.close() + await self._compute_worker.stop() + await _until(lambda: not self.sampling_guard._jobs) + self._backend_gate.shutdown() + + + +def _tokens(record, dialect): + assert record['status'] == 'completed', record + result = record['result'] + if dialect == 'tinker': + return result.sequences[0].tokens + return result['samples'][0]['sequences'][0]['tokens'] diff --git a/tests/server/sampler/test_resolve_sampler_weights.py b/tests/server/sampler/test_resolve_sampler_weights.py index 1d98c76c9..ecaaa18f2 100644 --- a/tests/server/sampler/test_resolve_sampler_weights.py +++ b/tests/server/sampler/test_resolve_sampler_weights.py @@ -24,7 +24,7 @@ async def call_backend(self, fn, *args, **kwargs): # The helper only routes load_full_weights_from_path through call_backend. if getattr(fn, '__name__', None) == 'load_full_weights_from_path': self.full_weight_loads.append(args[0]) - return fn(*args, **kwargs) + return fn(*args) @pytest.mark.asyncio diff --git a/tests/server/sampler/test_sampling_concurrency.py b/tests/server/sampler/test_sampling_concurrency.py new file mode 100644 index 000000000..fd7afb87d --- /dev/null +++ b/tests/server/sampler/test_sampling_concurrency.py @@ -0,0 +1,392 @@ +from __future__ import annotations + +import asyncio +import queue +import threading +from types import SimpleNamespace + +import pytest +from fastapi import FastAPI +from ray.exceptions import GetTimeoutError +from starlette.requests import Request +from tinker import types as tinker_types + +from twinkle.data_format import SampledSequence, SampleResponse +from twinkle.protocol.types import SampleRequest +from twinkle.server.sampler.app import SamplerManagement +from twinkle.server.sampler.concurrency import SamplerConcurrency +from twinkle.server.sampler.tinker_handlers import _register_tinker_sampler_routes +from twinkle.server.sampler.twinkle_handlers import _register_twinkle_sampler_routes +from twinkle.server.task_queue.backend_gate import BackendGate +from twinkle.server.task_queue.config import TaskQueueConfig +from twinkle.server.task_queue.types import BackendBusyError +from twinkle.server.task_queue.worker import ComputeWorker + + +from tests.server.sampler.sampling_test_helpers import _until, _State, _Limiter, _Sampler, _Service, _tokens + +@pytest.mark.asyncio +@pytest.mark.parametrize('dialects', [('tinker', 'tinker'), ('twinkle', 'twinkle'), ('tinker', 'twinkle')]) +@pytest.mark.parametrize('count', [1, 8, 32]) +async def test_real_routes_overlap_and_keep_request_results(dialects, count): + service = _Service() + service.sampler.release.clear() + try: + jobs = [dialects[i % 2] for i in range(count)] + refs = await asyncio.gather(*(service.submit(dialect, token=i + 1) for i, dialect in enumerate(jobs))) + await _until(lambda: service.sampler.peak >= min(count, 2)) + service.sampler.release.set() + records = await asyncio.gather(*(service.result(ref) for ref in refs)) + assert [_tokens(record, dialect) for record, dialect in zip(records, jobs)] == [[i + 1] for i in range(count)] + assert service._compute_worker._worker_task is None + finally: + await service.close() + + + +@pytest.mark.asyncio +@pytest.mark.parametrize('first,second', [('tinker', 'tinker'), ('twinkle', 'twinkle'), + ('tinker', 'twinkle'), ('twinkle', 'tinker')]) +async def test_full_checkpoint_generation_is_one_exclusive_transaction(tmp_path, first, second): + for number in ['1', '2']: + (tmp_path / number).mkdir() + service = _Service() + service.sampler.release.clear() + try: + a = await service.submit(first, path=str(tmp_path / '1')) + await _until(service.sampler.entered.is_set) + b = await service.submit(second, path=str(tmp_path / '2')) + await _until(lambda: service.sampling_guard._waiting_writers == 1) + assert service.sampler.loads == [1] + service.sampler.release.set() + assert _tokens(await service.result(a), first) == [1] + assert _tokens(await service.result(b), second) == [2] + assert service.sampler.loads == [1, 2] + finally: + await service.close() + + + +@pytest.mark.asyncio +async def test_mutation_waits_for_both_dialects_generation(): + service = _Service() + service.sampler.release.clear() + mutation = None + try: + a, b = await asyncio.gather(service.submit('tinker'), service.submit('twinkle')) + await _until(lambda: service.sampler.peak == 2) + mutation = asyncio.create_task(service.call_backend(service.sampler.set_template, 'template')) + await _until(lambda: service.sampling_guard._waiting_writers == 1) + assert service.sampler.template_calls == 0 + service.sampler.release.set() + await asyncio.gather(service.result(a), service.result(b), mutation) + assert service.sampler.template_calls == 1 + finally: + service.sampler.release.set() + if mutation is not None: + await asyncio.gather(mutation, return_exceptions=True) + await service.close() + + + +@pytest.mark.asyncio +async def test_timeout_keeps_actual_load_exclusive_and_recovers_after_exit(tmp_path): + for number in ['1', '2']: + (tmp_path / number).mkdir() + service = _Service(timeout=0.05) + load_entered = threading.Event() + load_release = threading.Event() + original = service.sampler.load_full_weights_from_path + + def slow_load(path): + load_entered.set() + if not load_release.wait(3): + raise AssertionError('test did not release loading') + original(path) + + service.sampler.load_full_weights_from_path = slow_load + try: + ref = await service.submit('tinker', path=str(tmp_path / '1')) + await _until(load_entered.is_set) + failed = await service.result(ref) + assert failed['failure'].reason_code == 'execution_timeout' + assert service.sampling_guard._writer + next_ref = await service.submit('twinkle', path=str(tmp_path / '2')) + rejected = await service.result(next_ref) + assert rejected['failure'].reason_code == 'backend_gate_unavailable' + assert service.sampler.loads == [] + load_release.set() + await _until(lambda: not service.sampling_guard._writer) + ref = await service.submit('twinkle', path=str(tmp_path / '2')) + assert _tokens(await service.result(ref), 'twinkle') == [2] + finally: + load_release.set() + await service.close() + + + +@pytest.mark.asyncio +async def test_ray_timeout_refuses_unconfirmed_backend_release(): + service = _Service() + service.sampler.sample = lambda *a, **k: (_ for _ in ()).throw(GetTimeoutError('actor still running')) + try: + with pytest.raises(GetTimeoutError): + await service.call_backend(service.sampler.sample) + with pytest.raises(BackendBusyError): + await service.sampling_guard.run(lambda: asyncio.sleep(0)) + finally: + await service.close() + + + +@pytest.mark.asyncio +async def test_waiting_writer_blocks_new_readers_until_the_writer_finishes(): + guard = SamplerConcurrency() + readers_release = asyncio.Event() + writer_release = asyncio.Event() + order = [] + + async def reader(): + order.append('first-reader') + await readers_release.wait() + + async def writer(): + order.append('writer') + await writer_release.wait() + + async def next_reader(): + order.append('next-reader') + + first = asyncio.create_task(guard.run(reader)) + await _until(lambda: guard._readers == 1) + write = asyncio.create_task(guard.run(writer, exclusive=True)) + await _until(lambda: guard._waiting_writers == 1) + second = asyncio.create_task(guard.run(next_reader)) + await asyncio.sleep(0) + assert order == ['first-reader'] + readers_release.set() + await _until(lambda: guard._writer) + assert order == ['first-reader', 'writer'] + writer_release.set() + await asyncio.gather(first, write, second) + assert order == ['first-reader', 'writer', 'next-reader'] + assert guard._readers == 0 + + + +@pytest.mark.asyncio +async def test_shutdown_cancels_futures_without_freeing_running_backend_early(): + service = _Service() + service.sampler.release.clear() + try: + a, b = await asyncio.gather(service.submit('tinker'), service.submit('twinkle')) + await _until(lambda: service.sampler.peak == 2) + service.sampling_guard.close() + await service._compute_worker.stop() + assert service.state.records[a]['failure'].reason_code == 'orphaned_replica' + assert service.state.records[b]['failure'].reason_code == 'orphaned_replica' + assert service.sampling_guard._readers == 2 + with pytest.raises(BackendBusyError): + await service.call_backend(service.sampler.set_template, 'template') + finally: + await service.close() + + + +@pytest.mark.asyncio +async def test_data_plane_generation_blocks_mutations_through_collection(): + service = _Service() + admitted = threading.Event() + completed = threading.Event() + stored = [] + service.data_plane = SimpleNamespace(enabled=True) + + def submit_generation(*args, **kwargs): + admitted.set() + + def status(submission_id): + return {'status': 'completed' if completed.is_set() else 'running'} + + service.sampler.submit_generation = submit_generation + service.sampler.get_generation_status = status + service.sampler.collect_generation = lambda submission_id: service.sampler.sample([{'input_ids': [1]}]) + service.sampler.cancel_generation = lambda submission_id: None + + async def put(rows, **kwargs): + from twinkle.protocol.types import DataRef + stored.extend(rows) + return DataRef(ref_id='ref', size=len(rows), fields=['tokens'], kind='rollout') + + service.data_plane.put = put + service.submit_background_and_peek = lambda factory, **kwargs: factory() + from twinkle.protocol.types import DataPlaneSampleRequest + body = DataPlaneSampleRequest(inputs=[{'input_ids': [1]}], sampling_params={'max_tokens': 1}) + request = Request({'type': 'http', 'headers': []}) + generation = asyncio.create_task(service.endpoints['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/twinkle/sample_to_data_plane'](request, body, service)) + mutation = None + try: + await _until(admitted.is_set) + mutation = asyncio.create_task(service.call_backend(service.sampler.set_template, 'template')) + await _until(lambda: service.sampling_guard._waiting_writers == 1) + assert service.sampler.template_calls == 0 + completed.set() + await asyncio.gather(generation, mutation) + assert stored[0]['tokens'] == [1] + assert service.sampler.template_calls == 1 + finally: + completed.set() + await asyncio.gather(generation, *([mutation] if mutation is not None else []), return_exceptions=True) + await service.close() + + + +@pytest.mark.asyncio +@pytest.mark.parametrize('dialect', ['tinker', 'twinkle']) +async def test_partial_checkpoint_failure_quarantines_changed_weights(tmp_path, dialect): + service = _Service() + def partial_load(path): + service.sampler.weight = 99 + raise RuntimeError('load failed after a weight was changed') + service.sampler.load_full_weights_from_path = partial_load + try: + ref = await service.submit(dialect, path=str(tmp_path)) + assert (await service.result(ref))['status'] == 'failed' + ref = await service.submit('twinkle') + record = await service.result(ref) + assert record['failure'].reason_code == 'backend_gate_unavailable' + assert service.sampler.active == 0 + finally: + await service.close() + + + +@pytest.mark.asyncio +async def test_shutdown_before_task_first_step_marks_future_failed(): + from twinkle.server.task_queue.types import QueuedTask + service = _Service() + worker = service._compute_worker + worker.reserve_detached() + await service.state.store_future_status('not-started', 'pending') + worker.run_detached(QueuedTask('not-started', lambda: asyncio.sleep(0), None, None, 0, 'sample', 0), 'q') + await worker.stop() + assert service.state.records['not-started']['failure'].reason_code == 'orphaned_replica' + assert worker._detached_reserved == 0 + await service.close() + + + +@pytest.mark.asyncio +async def test_completed_generation_failure_does_not_quarantine_replica(): + from twinkle.server.sampler.twinkle_handlers import _await_generation + service = _Service() + service.sampler.get_generation_status = lambda _: {'status': 'failed', 'error': 'invalid input'} + service.sampler.cancel_generation = lambda _: {'status': 'completed'} + try: + with pytest.raises(RuntimeError, match='invalid input'): + await service.sampling_guard.run(lambda: _await_generation(service, 'bad', 1)) + ref = await service.submit('twinkle') + assert _tokens(await service.result(ref), 'twinkle') == [7] + finally: + await service.close() + + + +@pytest.mark.asyncio +async def test_timeout_keeps_job_capacity_and_does_not_overwrite_failed_future(): + service = _Service(timeout=0.05, limit=1) + service.sampler.release.clear() + try: + ref = await service.submit('tinker') + await _until(service.sampler.entered.is_set) + failed = await service.result(ref) + assert failed['failure'].reason_code == 'execution_timeout' + assert len(service.sampling_guard._jobs) == 1 + with pytest.raises(BackendBusyError): + await service.sampling_guard.run(lambda: asyncio.sleep(0)) + service.sampler.release.set() + await _until(lambda: not service.sampling_guard._jobs) + assert service.state.records[ref] is failed + assert (await service.result(await service.submit('twinkle')))['status'] == 'completed' + finally: + await service.close() + + + +@pytest.mark.asyncio +async def test_nested_status_poll_reuses_job_access(): + service = _Service() + try: + async def job(): + # A separate polling Task inherits the job's access context. + await asyncio.wait_for(service.call_backend(service.sampler.set_template, 'test'), timeout=1) + await service.sampling_guard.run(job, exclusive=True) + assert service.sampler.template_calls == 1 + assert not service.sampling_guard._writer + finally: + await service.close() + + + +@pytest.mark.asyncio +async def test_admitted_writer_waits_for_abandoned_reader_to_finish(): + guard = SamplerConcurrency() + release = asyncio.Event() + wrote = asyncio.Event() + async def read(): + await release.wait() + async def write(): + wrote.set() + reader = asyncio.create_task(guard.run(read)) + await _until(lambda: guard._readers == 1) + writer = asyncio.create_task(guard.run(write, exclusive=True)) + await _until(lambda: guard._waiting_writers == 1) + reader.cancel() + await asyncio.gather(reader, return_exceptions=True) + assert not wrote.is_set() + release.set() + await asyncio.wait_for(writer, 1) + assert wrote.is_set() + + + +@pytest.mark.asyncio +async def test_abandoned_waiting_writer_never_starts_and_releases_capacity(): + guard = SamplerConcurrency(limit=2) + release = asyncio.Event() + wrote = asyncio.Event() + async def write(): + wrote.set() + reader = asyncio.create_task(guard.run(release.wait)) + await _until(lambda: guard._readers == 1) + writer = asyncio.create_task(guard.run(write, exclusive=True)) + await _until(lambda: guard._waiting_writers == 1) + writer.cancel() + await asyncio.gather(writer, return_exceptions=True) + await _until(lambda: len(guard._jobs) == 1) + assert not wrote.is_set() and guard._waiting_writers == 0 + release.set() + await reader + + + +@pytest.mark.asyncio +@pytest.mark.parametrize('dialect', ['tinker', 'twinkle']) +async def test_backend_cancellation_leaves_terminal_failure(dialect): + from concurrent.futures import CancelledError + service = _Service() + + def cancelled(*args, **kwargs): + raise CancelledError('engine coroutine cancelled') + + service.sampler.sample = cancelled + try: + ref = await service.submit(dialect) + record = await service.result(ref) + assert record['status'] == 'failed' + assert record['failure'].reason_code == 'backend_cancelled' + await _until(lambda: not service._compute_worker._detached_tasks) + assert service._compute_worker._detached_reserved == 0 + await service._compute_worker.stop() + assert service.state.records[ref] == record + finally: + await service.close() diff --git a/tests/server/sampler/test_stream_guarantees.py b/tests/server/sampler/test_stream_guarantees.py index d14895751..ff4907e4b 100644 --- a/tests/server/sampler/test_stream_guarantees.py +++ b/tests/server/sampler/test_stream_guarantees.py @@ -89,6 +89,8 @@ async def test_stream_timeout_returns_error_payload_and_closes_queue(): class _GenerationService: def __init__(self) -> None: + from twinkle.server.sampler.concurrency import SamplerConcurrency + self.sampling_guard = SamplerConcurrency() self.cancelled = False self.sampler = SimpleNamespace( get_generation_status=lambda _submission_id: {'status': 'running'}, @@ -99,8 +101,9 @@ def __init__(self) -> None: async def call_backend(self, fn, /, *args, **kwargs): return fn(*args, **kwargs) - def _cancel(self, _submission_id: str) -> None: + def _cancel(self, _submission_id: str): self.cancelled = True + return {'status': 'completed'} @pytest.mark.asyncio diff --git a/tests/server/sampler/test_tinker_handlers.py b/tests/server/sampler/test_tinker_handlers.py index 788183490..2e9cb3133 100644 --- a/tests/server/sampler/test_tinker_handlers.py +++ b/tests/server/sampler/test_tinker_handlers.py @@ -1,3 +1,4 @@ +import asyncio import pytest from fastapi import FastAPI from starlette.requests import Request @@ -5,6 +6,7 @@ from twinkle.data_format import SampledSequence, SampleResponse from twinkle.server.sampler.tinker_handlers import _register_tinker_sampler_routes +from twinkle.server.sampler.concurrency import SamplerConcurrency class _DummyState: @@ -21,12 +23,13 @@ class _DummySampler: def __init__(self): self.adapter_paths = [] self.sampling_params = [] + self.calls = [] def set_template(self, *args, **kwargs): - return None + self.calls.append('set_template') def reset_prefix_cache(self): - return None + self.calls.append('reset_prefix_cache') def sample(self, inputs, sampling_params=None, adapter_name='', *, adapter_path=None, **kwargs): self.adapter_paths.append(adapter_path) @@ -47,15 +50,22 @@ def __init__(self): self.model_id = 'mock-model' self.state = _DummyState() self.sampler = _DummySampler() + self.tinker_setup_lock = asyncio.Lock() + self.tinker_template_set = False + self.sampling_guard = SamplerConcurrency() async def _on_request_start(self, request): return 'EMPTY_TOKEN' - async def schedule_task(self, task, **kwargs): + async def schedule_concurrent_task(self, task, **kwargs): return await task() async def call_backend(self, fn, /, *args, **kwargs): - return fn(*args, **kwargs) + async def call(): + return fn(*args, **kwargs) + if self.sampling_guard.owns_job(): + return await call() + return await self.sampling_guard.run(call, exclusive=True) @pytest.mark.asyncio @@ -83,3 +93,24 @@ async def test_tinker_asample_allows_base_model_session_without_model_path(tempe assert management.sampler.sampling_params[0].logprobs == 1 assert management.sampler.sampling_params[0].temperature == temperature assert management.sampler.sampling_params[0].max_tokens == (max_tokens if max_tokens is not None else 256) + + +@pytest.mark.asyncio +async def test_tinker_asample_sets_template_once_without_prefix_cache_reset(): + """Both calls would wait for every in-flight generation, serializing requests.""" + management = _DummyManagement() + app = FastAPI() + _register_tinker_sampler_routes(app, lambda: management) + route = next(route for route in app.routes if getattr(route, 'path', None) == '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/tinker/asample') + body = types.SampleRequest( + prompt=types.ModelInput.from_ints([1, 2, 3]), + sampling_params=types.SamplingParams(max_tokens=2), + sampling_session_id='base-session', + base_model='mock-model', + ) + request = Request({'type': 'http', 'headers': []}) + + await asyncio.gather(*(route.endpoint(request, body, management) for _ in range(3))) + + assert management.sampler.calls == ['set_template'] + assert len(management.sampler.adapter_paths) == 3 diff --git a/tests/server/sampler/test_twinkle_async_rows.py b/tests/server/sampler/test_twinkle_async_rows.py index 4df488130..b98ee2fd9 100644 --- a/tests/server/sampler/test_twinkle_async_rows.py +++ b/tests/server/sampler/test_twinkle_async_rows.py @@ -64,6 +64,8 @@ def test_async_sampler_rejects_group_id_count_mismatch() -> None: class _SamplerManagement: def __init__(self): + from twinkle.server.sampler.concurrency import SamplerConcurrency + self.sampling_guard = SamplerConcurrency() self.sampler = self self.data_plane = self self.enabled = True diff --git a/tests/server/utils/test_task_queue_mixin.py b/tests/server/utils/test_task_queue_mixin.py index bdb830650..61d682f3e 100644 --- a/tests/server/utils/test_task_queue_mixin.py +++ b/tests/server/utils/test_task_queue_mixin.py @@ -360,3 +360,214 @@ async def _factory(): env = await queue.submit_and_peek(_factory, task_type='step', request_id='req_fixed123') assert env.request_id == 'req_fixed123' assert env.status == 'completed' + + +@pytest.mark.asyncio +async def test_concurrent_tasks_run_side_by_side(): + queue = _DummyQueue() + queue.enable_compute_worker() + started = [] + both_started = asyncio.Event() + + async def work(index): + started.append(index) + if len(started) == 2: + both_started.set() + # Under serial execution the first task would wait here forever. + await asyncio.wait_for(both_started.wait(), timeout=1) + return {'index': index} + + try: + refs = [ + await queue.schedule_concurrent_task(lambda i=i: work(i), token='token1', task_type='sample') + for i in range(2) + ] + for _ in range(100): + records = [await queue.state.get_future(ref['request_id']) for ref in refs] + if all(record['status'] == 'completed' for record in records): + break + await asyncio.sleep(0.01) + finally: + await queue._compute_worker.stop() + + assert [record['result'] for record in records] == [{'index': 0}, {'index': 1}] + # The serial worker loop is never involved. + assert queue._compute_worker._worker_task is None + + +@pytest.mark.asyncio +async def test_concurrent_task_keeps_queue_failure_mapping(): + queue = _DummyQueue() + queue.enable_compute_worker() + + async def work(): + raise UserTaskError('invalid request') + + try: + await queue.schedule_concurrent_task(work, token='token1') + for _ in range(100): + failed = [kwargs for args, kwargs in queue.state.records if args[1] == 'failed'] + if failed: + break + await asyncio.sleep(0) + finally: + await queue._compute_worker.stop() + + assert [args[1] for args, _ in queue.state.records] == ['pending', 'running', 'failed'] + failure = failed[-1]['failure'] + assert failure.reason_code == 'request_rejected' + assert failure.attribution == 'user' + + +@pytest.mark.asyncio +async def test_concurrent_task_preflight_rejection_raises_without_writing(): + from twinkle.server.exceptions import InputTokensExceededError + queue = _DummyQueue() + queue.enable_compute_worker() + + async def work(): + raise AssertionError('preflight rejection must not execute the task') + + with pytest.raises(InputTokensExceededError): + await queue.schedule_concurrent_task( + work, token='token1', input_tokens=queue.task_queue_config.max_input_tokens + 1) + + assert queue.state.records == [] + assert not queue._compute_worker._detached_tasks + + +@pytest.mark.asyncio +async def test_worker_stop_cancels_concurrent_task(): + queue = _DummyQueue() + queue.enable_compute_worker() + entered = asyncio.Event() + cancelled = asyncio.Event() + + async def work(): + entered.set() + try: + await asyncio.Event().wait() + finally: + cancelled.set() + + await queue.schedule_concurrent_task(work, token='token1') + await asyncio.wait_for(entered.wait(), timeout=1) + await asyncio.wait_for(queue._compute_worker.stop(), timeout=1) + assert cancelled.is_set() + assert not queue._compute_worker._detached_tasks + + +@pytest.mark.asyncio +async def test_concurrent_capacity_rejects_before_future_write_and_is_reusable(): + from twinkle.server.exceptions import ResourceExhaustedError + queue = _DummyQueue() + queue._task_queue_config = TaskQueueConfig(concurrent_task_limit=1) + queue.enable_compute_worker() + release = asyncio.Event() + + async def work(): + await release.wait() + return {'ok': True} + + try: + first = await queue.schedule_concurrent_task(work, token='token1') + with pytest.raises(ResourceExhaustedError): + await queue.schedule_concurrent_task(work, token='token2') + assert set(queue.state._latest) == {first['request_id']} + release.set() + for _ in range(100): + if queue._compute_worker._detached_reserved == 0: + break + await asyncio.sleep(0.001) + envelope = await queue.submit_concurrent_and_peek(work, token='token2') + assert envelope.status == 'completed' + finally: + release.set() + await queue._compute_worker.stop() + + +@pytest.mark.asyncio +async def test_concurrent_reservation_is_released_when_state_write_fails(): + queue = _DummyQueue() + queue.enable_compute_worker() + + async def fail(*args, **kwargs): + raise RuntimeError('state unavailable') + + queue.state.store_future_status = fail + with pytest.raises(RuntimeError, match='state unavailable'): + await queue.schedule_concurrent_task(lambda: asyncio.sleep(0), token='token1') + assert queue._compute_worker._detached_reserved == 0 + assert not queue._compute_worker._detached_tasks + + +@pytest.mark.asyncio +async def test_stop_between_future_write_and_start_leaves_terminal_record(): + from twinkle.server.exceptions import ResourceExhaustedError + queue = _DummyQueue() + queue.enable_compute_worker() + original = queue.state.store_future_status + + async def write_then_stop(*args, **kwargs): + await original(*args, **kwargs) + if args[1] == 'pending': + await queue._compute_worker.stop() + + queue.state.store_future_status = write_then_stop + with pytest.raises(ResourceExhaustedError): + await queue.schedule_concurrent_task(lambda: asyncio.sleep(0), token='token1', request_id='req-stop') + assert queue.state._latest['req-stop']['status'] == 'failed' + assert queue._compute_worker._detached_reserved == 0 + + +@pytest.mark.asyncio +async def test_shutdown_cancel_before_task_start_leaves_failed_future(): + queue = _DummyQueue() + queue.enable_compute_worker() + executed = False + async def work(): + nonlocal executed + executed = True + ref = await queue.schedule_concurrent_task(work, token='token1') + await queue._compute_worker.stop() # No yield between admission and shutdown. + record = queue.state._latest[ref['request_id']] + assert record['status'] == 'failed' + assert record['failure']['reason_code'] == 'orphaned_replica' + assert not executed and queue._compute_worker._detached_reserved == 0 + + +@pytest.mark.asyncio +async def test_cancel_preserves_existing_terminal_record(): + queue = _DummyQueue() + queue.enable_compute_worker() + entered = asyncio.Event() + async def work(): + entered.set() + await asyncio.Event().wait() + ref = await queue.schedule_concurrent_task(work, token='token1') + await entered.wait() + await queue.state.store_future_status(ref['request_id'], 'failed', None, reason='already terminal') + record = queue.state._latest[ref['request_id']].copy() + await queue._compute_worker.stop() + assert queue.state._latest[ref['request_id']] == record + + +@pytest.mark.asyncio +async def test_cancel_during_running_status_write_leaves_failed_future(): + queue = _DummyQueue() + queue.enable_compute_worker() + entered = asyncio.Event() + original = queue.state.store_future_status + async def write(*args, **kwargs): + if args[1] == 'running': + entered.set() + await asyncio.Event().wait() + await original(*args, **kwargs) + queue.state.store_future_status = write + ref = await queue.schedule_concurrent_task(lambda: asyncio.sleep(0), token='token1') + await entered.wait() + await queue._compute_worker.stop() + record = queue.state._latest[ref['request_id']] + assert record['status'] == 'failed' + assert record['failure']['reason_code'] == 'orphaned_replica' + assert [args[1] for args, _ in queue.state.records] == ['pending', 'failed'] diff --git a/tests/twinkle_agentic/test_vllm_sampler_tq_generation.py b/tests/twinkle_agentic/test_vllm_sampler_tq_generation.py index 3b81f387d..4b4343282 100644 --- a/tests/twinkle_agentic/test_vllm_sampler_tq_generation.py +++ b/tests/twinkle_agentic/test_vllm_sampler_tq_generation.py @@ -13,7 +13,7 @@ from twinkle.server.sampler.twinkle_handlers import _await_generation from twinkle_agentic.async_rl import LoraContext from twinkle_agentic.async_rl.types import PartitionAdmission, PromptGroup, RolloutPolicy -from twinkle_agentic.async_rl.vllm_sampler_tq import (VLLMSamplerTQ, _dispatch_generation, _GeneratedSample, +from twinkle_agentic.async_rl.vllm_sampler_tq import (_GenerationSubmission, VLLMSamplerTQ, _dispatch_generation, _GeneratedSample, _PromptGroupRolloutStats) @@ -174,7 +174,7 @@ def test_generation_failure_is_isolated_and_consumed() -> None: sampler = _bare_sampler() failed = Future() failed.set_exception(ValueError('bad prompt')) - sampler._generation_submissions['failed'] = failed + sampler._generation_submissions['failed'] = _GenerationSubmission(future=failed) state = sampler.get_generation_status('failed') assert state['status'] == 'failed' @@ -185,29 +185,43 @@ def test_generation_failure_is_isolated_and_consumed() -> None: assert 'failed' not in sampler._generation_submissions -def test_generation_can_be_cancelled_without_waiting() -> None: - sampler = _bare_sampler() - pending = Future() - sampler._generation_submissions['pending'] = pending - - state = sampler.cancel_generation('pending') - - assert state == {'submission_id': 'pending', 'status': 'cancelled'} - assert pending.cancelled() - assert 'pending' not in sampler._generation_submissions +def test_cancellation_waits_for_actual_coroutine_cleanup() -> None: + async def run(): + sampler = _bare_sampler() + sampler._async_loop = asyncio.get_running_loop() + entered = asyncio.Event() + release = asyncio.Event() + async def generate(*args, **kwargs): + try: + entered.set() + await asyncio.Future() + finally: + await release.wait() + sampler._generate_inputs = generate + sampler.submit_generation('pending', [], SamplingParams()) + await entered.wait() + assert sampler.cancel_generation('pending')['status'] == 'cancellation_requested' + await asyncio.sleep(.01) + assert sampler.get_generation_status('pending')['status'] == 'running' + assert sampler.cancel_generation('pending')['status'] == 'cancellation_requested' + release.set() + for _ in range(100): + if sampler.get_generation_status('pending')['status'] == 'cancelled': + break + await asyncio.sleep(.001) + assert sampler.cancel_generation('pending')['status'] == 'completed' + assert not sampler._generation_submissions + asyncio.run(run()) -def test_all_generations_are_cancelled_on_shutdown() -> None: +def test_all_completed_generations_are_forgotten_on_shutdown() -> None: sampler = _bare_sampler() - first = Future() - second = Future() - sampler._generation_submissions.update(first=first, second=second) - - state = sampler.cancel_all_generations() - - assert state == {'submissions': 2, 'cancelled': 2} - assert first.cancelled() - assert second.cancelled() + first, second = Future(), Future() + first.set_result([]) + second.set_exception(ValueError('bad input')) + sampler._generation_submissions.update( + first=_GenerationSubmission(future=first), second=_GenerationSubmission(future=second)) + assert sampler.cancel_all_generations() == {'submissions': 2, 'cancelled': 0} assert sampler._generation_submissions == {} From 1b2cb366043a62a90bb99b51e371c34b0afba0bb Mon Sep 17 00:00:00 2001 From: Yunnglin Date: Tue, 13 Oct 2026 21:33:08 +0800 Subject: [PATCH 2/5] test(sampler): simplify concurrency fixtures and remove unused imports --- .../server/integration/sampler_e2e_helpers.py | 12 +- .../test_sampler_concurrency_e2e.py | 26 -- tests/server/sampler/sampling_test_helpers.py | 26 +- .../sampler/test_sampling_concurrency.py | 239 +++++++----------- tests/server/utils/test_task_queue_mixin.py | 107 ++++---- 5 files changed, 163 insertions(+), 247 deletions(-) diff --git a/tests/server/integration/sampler_e2e_helpers.py b/tests/server/integration/sampler_e2e_helpers.py index f09f80f6a..4dbc31ce3 100644 --- a/tests/server/integration/sampler_e2e_helpers.py +++ b/tests/server/integration/sampler_e2e_helpers.py @@ -13,21 +13,20 @@ @pytest.fixture(scope='session', autouse=True) def _ray_runtime(): + """Use the external server; do not start the parent conftest's local Ray.""" yield - @pytest.fixture(autouse=True) def _reset_canonical_state_actor(): + """Keep external server state intact across E2E requests.""" yield - def _report(**values): print('SAMPLER_CONCURRENCY ' + json.dumps(values, sort_keys=True), flush=True) - @pytest.fixture(scope='module') def clients(): from transformers import AutoTokenizer @@ -47,14 +46,12 @@ def clients(): } - def _prompt(clients, text): return clients['tokenizer'].apply_chat_template( [{'role': 'user', 'content': text}], tokenize=True, add_generation_prompt=True, enable_thinking=False, return_dict=False) - def _sample(clients, dialect, ids, max_tokens=128, *, adapter=None, num_samples=1, temperature=0): if dialect == 'tinker': from tinker import types @@ -77,7 +74,6 @@ def _sample(clients, dialect, ids, max_tokens=128, *, adapter=None, num_samples= return sequences - def _parallel(jobs): barrier = threading.Barrier(len(jobs)) def run(job): @@ -89,7 +85,6 @@ def run(job): return results, time.perf_counter() - started - @pytest.fixture(scope='module') def adapters(clients): import torch @@ -124,7 +119,6 @@ def adapters(clients): return paths - @pytest.fixture(scope='module') def full_checkpoints(clients): """Two complete HF checkpoints with different, observable model weights.""" @@ -147,7 +141,6 @@ def full_checkpoints(clients): return [str(source), str(target)] - def _checkpoint_handles(clients, dialect, paths): if dialect == 'tinker': return [clients['service'].create_sampling_client(base_model=clients['model_id'], model_path=path) @@ -155,7 +148,6 @@ def _checkpoint_handles(clients, dialect, paths): return paths - def _assert_adapter_isolation(clients, adapters, dialect): handles = ([None] + [clients['service'].create_sampling_client(base_model=clients['model_id'], model_path=path) for path in adapters] if dialect == 'tinker' else [None, *adapters]) diff --git a/tests/server/integration/test_sampler_concurrency_e2e.py b/tests/server/integration/test_sampler_concurrency_e2e.py index 2193b9b0b..9ca2336e4 100644 --- a/tests/server/integration/test_sampler_concurrency_e2e.py +++ b/tests/server/integration/test_sampler_concurrency_e2e.py @@ -1,14 +1,10 @@ """Real Tinker/Twinkle concurrency, weight isolation and lifecycle checks.""" from __future__ import annotations -import concurrent.futures -import json import math import os import threading import time -from contextlib import contextmanager -from pathlib import Path import pytest @@ -41,7 +37,6 @@ def test_both_dialects_batch_independent_requests(clients, count, dialect): assert throughput > 1.5 * clients['single_throughput'][dialect], 'independent requests still appear serial' - @pytest.mark.parametrize('dialect', ['tinker', 'twinkle']) def test_concurrent_request_mapping(clients, dialect): words = ['apple', 'banana', 'orange', 'purple', 'yellow', 'green', 'blue', 'red'] @@ -53,7 +48,6 @@ def test_concurrent_request_mapping(clients, dialect): _report(check='mapping', dialect=dialect, count=32, unique_prompts=8, seconds=elapsed, solo_equal=True) - def test_mixed_dialects_and_multiple_samples(clients): ids = _prompt(clients, 'Reply with exactly this word and nothing else: apple') solo = _sample(clients, 'tinker', ids, 16)[0][0] @@ -71,7 +65,6 @@ def test_mixed_dialects_and_multiple_samples(clients): _report(check='mixed_multiple_samples', requests=16, sequences=64, seconds=elapsed) - @pytest.mark.parametrize('dialect', ['tinker', 'twinkle']) def test_base_and_two_lora_adapters_are_isolated(clients, adapters, dialect): _assert_adapter_isolation(clients, adapters, dialect) @@ -92,7 +85,6 @@ def test_full_checkpoints_are_isolated_across_dialects(clients, full_checkpoints _sample(clients, 'twinkle', ids, 2, adapter=full_checkpoints[0]) - @pytest.mark.parametrize('dialect', ['tinker', 'twinkle']) def test_invalid_generation_does_not_disable_sampler(clients, dialect): ids = _prompt(clients, 'Reply with exactly this word and nothing else: apple') @@ -102,7 +94,6 @@ def test_invalid_generation_does_not_disable_sampler(clients, dialect): _report(check='failure_recovery', dialect=dialect) - def test_async_generation_cleanup_and_recovery_on_real_engine(clients): import ray from ray.util.state import list_actors @@ -157,19 +148,6 @@ def wait_status(key, terminal): ray.shutdown() - - - - - - - - - - - - - def test_data_plane_http_generation_and_failure_recovery(clients): from twinkle_client.data_plane import DataPlaneClient store = DataPlaneClient(transport=clients['twinkle']._transport) @@ -204,10 +182,6 @@ def generate_and_fetch(): _report(check='data_plane_http', requests=4, sequences=6, recovered=True, seconds=elapsed) - - - - def test_backend_self_cancellation_terminates_both_sdk_futures(clients): """Cancel at the driver boundary after both requests performed real PPU generation. diff --git a/tests/server/sampler/sampling_test_helpers.py b/tests/server/sampler/sampling_test_helpers.py index 953c235d8..63564d363 100644 --- a/tests/server/sampler/sampling_test_helpers.py +++ b/tests/server/sampler/sampling_test_helpers.py @@ -2,13 +2,10 @@ from __future__ import annotations import asyncio -import queue import threading -from types import SimpleNamespace -import pytest +import pytest_asyncio from fastapi import FastAPI -from ray.exceptions import GetTimeoutError from starlette.requests import Request from tinker import types as tinker_types @@ -20,7 +17,6 @@ from twinkle.server.sampler.twinkle_handlers import _register_twinkle_sampler_routes from twinkle.server.task_queue.backend_gate import BackendGate from twinkle.server.task_queue.config import TaskQueueConfig -from twinkle.server.task_queue.types import BackendBusyError from twinkle.server.task_queue.worker import ComputeWorker @@ -31,7 +27,6 @@ async def wait(): await asyncio.wait_for(wait(), 3) - class _State: def __init__(self): self.records = {} @@ -43,13 +38,11 @@ async def get_future(self, request_id): return self.records.get(request_id) - class _Limiter: async def check_and_record(self, *args): return True, None - class _Sampler: def __init__(self): self.weight = 0 @@ -87,7 +80,6 @@ def sample(self, inputs, sampling_params=None, **kwargs): self.active -= 1 - class _Service(SamplerManagement): def __init__(self, *, timeout=2, limit=128): self.model_id = 'mock-model' @@ -138,10 +130,24 @@ async def close(self): self._backend_gate.shutdown() - def _tokens(record, dialect): assert record['status'] == 'completed', record result = record['result'] if dialect == 'tinker': return result.sequences[0].tokens return result['samples'][0]['sequences'][0]['tokens'] + + +@pytest_asyncio.fixture +async def make_service(): + """Create services and close them even when a test assertion fails.""" + services = [] + + def create(**kwargs): + service = _Service(**kwargs) + services.append(service) + return service + + yield create + for service in reversed(services): + await service.close() diff --git a/tests/server/sampler/test_sampling_concurrency.py b/tests/server/sampler/test_sampling_concurrency.py index fd7afb87d..23affb0db 100644 --- a/tests/server/sampler/test_sampling_concurrency.py +++ b/tests/server/sampler/test_sampling_concurrency.py @@ -1,75 +1,56 @@ from __future__ import annotations import asyncio -import queue import threading from types import SimpleNamespace import pytest -from fastapi import FastAPI from ray.exceptions import GetTimeoutError from starlette.requests import Request -from tinker import types as tinker_types -from twinkle.data_format import SampledSequence, SampleResponse -from twinkle.protocol.types import SampleRequest -from twinkle.server.sampler.app import SamplerManagement from twinkle.server.sampler.concurrency import SamplerConcurrency -from twinkle.server.sampler.tinker_handlers import _register_tinker_sampler_routes -from twinkle.server.sampler.twinkle_handlers import _register_twinkle_sampler_routes -from twinkle.server.task_queue.backend_gate import BackendGate -from twinkle.server.task_queue.config import TaskQueueConfig from twinkle.server.task_queue.types import BackendBusyError -from twinkle.server.task_queue.worker import ComputeWorker -from tests.server.sampler.sampling_test_helpers import _until, _State, _Limiter, _Sampler, _Service, _tokens +from tests.server.sampler.sampling_test_helpers import _tokens, _until, make_service @pytest.mark.asyncio @pytest.mark.parametrize('dialects', [('tinker', 'tinker'), ('twinkle', 'twinkle'), ('tinker', 'twinkle')]) @pytest.mark.parametrize('count', [1, 8, 32]) -async def test_real_routes_overlap_and_keep_request_results(dialects, count): - service = _Service() +async def test_real_routes_overlap_and_keep_request_results(make_service, dialects, count): + service = make_service() service.sampler.release.clear() - try: - jobs = [dialects[i % 2] for i in range(count)] - refs = await asyncio.gather(*(service.submit(dialect, token=i + 1) for i, dialect in enumerate(jobs))) - await _until(lambda: service.sampler.peak >= min(count, 2)) - service.sampler.release.set() - records = await asyncio.gather(*(service.result(ref) for ref in refs)) - assert [_tokens(record, dialect) for record, dialect in zip(records, jobs)] == [[i + 1] for i in range(count)] - assert service._compute_worker._worker_task is None - finally: - await service.close() - + jobs = [dialects[i % 2] for i in range(count)] + refs = await asyncio.gather(*(service.submit(dialect, token=i + 1) for i, dialect in enumerate(jobs))) + await _until(lambda: service.sampler.peak >= min(count, 2)) + service.sampler.release.set() + records = await asyncio.gather(*(service.result(ref) for ref in refs)) + assert [_tokens(record, dialect) for record, dialect in zip(records, jobs)] == [[i + 1] for i in range(count)] + assert service._compute_worker._worker_task is None @pytest.mark.asyncio @pytest.mark.parametrize('first,second', [('tinker', 'tinker'), ('twinkle', 'twinkle'), ('tinker', 'twinkle'), ('twinkle', 'tinker')]) -async def test_full_checkpoint_generation_is_one_exclusive_transaction(tmp_path, first, second): +async def test_full_checkpoint_generation_is_one_exclusive_transaction(make_service, tmp_path, first, second): for number in ['1', '2']: (tmp_path / number).mkdir() - service = _Service() + service = make_service() service.sampler.release.clear() - try: - a = await service.submit(first, path=str(tmp_path / '1')) - await _until(service.sampler.entered.is_set) - b = await service.submit(second, path=str(tmp_path / '2')) - await _until(lambda: service.sampling_guard._waiting_writers == 1) - assert service.sampler.loads == [1] - service.sampler.release.set() - assert _tokens(await service.result(a), first) == [1] - assert _tokens(await service.result(b), second) == [2] - assert service.sampler.loads == [1, 2] - finally: - await service.close() - + a = await service.submit(first, path=str(tmp_path / '1')) + await _until(service.sampler.entered.is_set) + b = await service.submit(second, path=str(tmp_path / '2')) + await _until(lambda: service.sampling_guard._waiting_writers == 1) + assert service.sampler.loads == [1] + service.sampler.release.set() + assert _tokens(await service.result(a), first) == [1] + assert _tokens(await service.result(b), second) == [2] + assert service.sampler.loads == [1, 2] @pytest.mark.asyncio -async def test_mutation_waits_for_both_dialects_generation(): - service = _Service() +async def test_mutation_waits_for_both_dialects_generation(make_service): + service = make_service() service.sampler.release.clear() mutation = None try: @@ -85,15 +66,13 @@ async def test_mutation_waits_for_both_dialects_generation(): service.sampler.release.set() if mutation is not None: await asyncio.gather(mutation, return_exceptions=True) - await service.close() - @pytest.mark.asyncio -async def test_timeout_keeps_actual_load_exclusive_and_recovers_after_exit(tmp_path): +async def test_timeout_keeps_actual_load_exclusive_and_recovers_after_exit(make_service, tmp_path): for number in ['1', '2']: (tmp_path / number).mkdir() - service = _Service(timeout=0.05) + service = make_service(timeout=0.05) load_entered = threading.Event() load_release = threading.Event() original = service.sampler.load_full_weights_from_path @@ -121,22 +100,16 @@ def slow_load(path): assert _tokens(await service.result(ref), 'twinkle') == [2] finally: load_release.set() - await service.close() - @pytest.mark.asyncio -async def test_ray_timeout_refuses_unconfirmed_backend_release(): - service = _Service() +async def test_ray_timeout_refuses_unconfirmed_backend_release(make_service): + service = make_service() service.sampler.sample = lambda *a, **k: (_ for _ in ()).throw(GetTimeoutError('actor still running')) - try: - with pytest.raises(GetTimeoutError): - await service.call_backend(service.sampler.sample) - with pytest.raises(BackendBusyError): - await service.sampling_guard.run(lambda: asyncio.sleep(0)) - finally: - await service.close() - + with pytest.raises(GetTimeoutError): + await service.call_backend(service.sampler.sample) + with pytest.raises(BackendBusyError): + await service.sampling_guard.run(lambda: asyncio.sleep(0)) @pytest.mark.asyncio @@ -173,29 +146,24 @@ async def next_reader(): assert guard._readers == 0 - @pytest.mark.asyncio -async def test_shutdown_cancels_futures_without_freeing_running_backend_early(): - service = _Service() +async def test_shutdown_cancels_futures_without_freeing_running_backend_early(make_service): + service = make_service() service.sampler.release.clear() - try: - a, b = await asyncio.gather(service.submit('tinker'), service.submit('twinkle')) - await _until(lambda: service.sampler.peak == 2) - service.sampling_guard.close() - await service._compute_worker.stop() - assert service.state.records[a]['failure'].reason_code == 'orphaned_replica' - assert service.state.records[b]['failure'].reason_code == 'orphaned_replica' - assert service.sampling_guard._readers == 2 - with pytest.raises(BackendBusyError): - await service.call_backend(service.sampler.set_template, 'template') - finally: - await service.close() - + a, b = await asyncio.gather(service.submit('tinker'), service.submit('twinkle')) + await _until(lambda: service.sampler.peak == 2) + service.sampling_guard.close() + await service._compute_worker.stop() + assert service.state.records[a]['failure'].reason_code == 'orphaned_replica' + assert service.state.records[b]['failure'].reason_code == 'orphaned_replica' + assert service.sampling_guard._readers == 2 + with pytest.raises(BackendBusyError): + await service.call_backend(service.sampler.set_template, 'template') @pytest.mark.asyncio -async def test_data_plane_generation_blocks_mutations_through_collection(): - service = _Service() +async def test_data_plane_generation_blocks_mutations_through_collection(make_service): + service = make_service() admitted = threading.Event() completed = threading.Event() stored = [] @@ -236,34 +204,28 @@ async def put(rows, **kwargs): finally: completed.set() await asyncio.gather(generation, *([mutation] if mutation is not None else []), return_exceptions=True) - await service.close() - @pytest.mark.asyncio @pytest.mark.parametrize('dialect', ['tinker', 'twinkle']) -async def test_partial_checkpoint_failure_quarantines_changed_weights(tmp_path, dialect): - service = _Service() +async def test_partial_checkpoint_failure_quarantines_changed_weights(make_service, tmp_path, dialect): + service = make_service() def partial_load(path): service.sampler.weight = 99 raise RuntimeError('load failed after a weight was changed') service.sampler.load_full_weights_from_path = partial_load - try: - ref = await service.submit(dialect, path=str(tmp_path)) - assert (await service.result(ref))['status'] == 'failed' - ref = await service.submit('twinkle') - record = await service.result(ref) - assert record['failure'].reason_code == 'backend_gate_unavailable' - assert service.sampler.active == 0 - finally: - await service.close() - + ref = await service.submit(dialect, path=str(tmp_path)) + assert (await service.result(ref))['status'] == 'failed' + ref = await service.submit('twinkle') + record = await service.result(ref) + assert record['failure'].reason_code == 'backend_gate_unavailable' + assert service.sampler.active == 0 @pytest.mark.asyncio -async def test_shutdown_before_task_first_step_marks_future_failed(): +async def test_shutdown_before_task_first_step_marks_future_failed(make_service): from twinkle.server.task_queue.types import QueuedTask - service = _Service() + service = make_service() worker = service._compute_worker worker.reserve_detached() await service.state.store_future_status('not-started', 'pending') @@ -271,60 +233,46 @@ async def test_shutdown_before_task_first_step_marks_future_failed(): await worker.stop() assert service.state.records['not-started']['failure'].reason_code == 'orphaned_replica' assert worker._detached_reserved == 0 - await service.close() - @pytest.mark.asyncio -async def test_completed_generation_failure_does_not_quarantine_replica(): +async def test_completed_generation_failure_does_not_quarantine_replica(make_service): from twinkle.server.sampler.twinkle_handlers import _await_generation - service = _Service() + service = make_service() service.sampler.get_generation_status = lambda _: {'status': 'failed', 'error': 'invalid input'} service.sampler.cancel_generation = lambda _: {'status': 'completed'} - try: - with pytest.raises(RuntimeError, match='invalid input'): - await service.sampling_guard.run(lambda: _await_generation(service, 'bad', 1)) - ref = await service.submit('twinkle') - assert _tokens(await service.result(ref), 'twinkle') == [7] - finally: - await service.close() - + with pytest.raises(RuntimeError, match='invalid input'): + await service.sampling_guard.run(lambda: _await_generation(service, 'bad', 1)) + ref = await service.submit('twinkle') + assert _tokens(await service.result(ref), 'twinkle') == [7] @pytest.mark.asyncio -async def test_timeout_keeps_job_capacity_and_does_not_overwrite_failed_future(): - service = _Service(timeout=0.05, limit=1) +async def test_timeout_keeps_job_capacity_and_does_not_overwrite_failed_future(make_service): + service = make_service(timeout=0.05, limit=1) service.sampler.release.clear() - try: - ref = await service.submit('tinker') - await _until(service.sampler.entered.is_set) - failed = await service.result(ref) - assert failed['failure'].reason_code == 'execution_timeout' - assert len(service.sampling_guard._jobs) == 1 - with pytest.raises(BackendBusyError): - await service.sampling_guard.run(lambda: asyncio.sleep(0)) - service.sampler.release.set() - await _until(lambda: not service.sampling_guard._jobs) - assert service.state.records[ref] is failed - assert (await service.result(await service.submit('twinkle')))['status'] == 'completed' - finally: - await service.close() - + ref = await service.submit('tinker') + await _until(service.sampler.entered.is_set) + failed = await service.result(ref) + assert failed['failure'].reason_code == 'execution_timeout' + assert len(service.sampling_guard._jobs) == 1 + with pytest.raises(BackendBusyError): + await service.sampling_guard.run(lambda: asyncio.sleep(0)) + service.sampler.release.set() + await _until(lambda: not service.sampling_guard._jobs) + assert service.state.records[ref] is failed + assert (await service.result(await service.submit('twinkle')))['status'] == 'completed' @pytest.mark.asyncio -async def test_nested_status_poll_reuses_job_access(): - service = _Service() - try: - async def job(): - # A separate polling Task inherits the job's access context. - await asyncio.wait_for(service.call_backend(service.sampler.set_template, 'test'), timeout=1) - await service.sampling_guard.run(job, exclusive=True) - assert service.sampler.template_calls == 1 - assert not service.sampling_guard._writer - finally: - await service.close() - +async def test_nested_status_poll_reuses_job_access(make_service): + service = make_service() + async def job(): + # A separate polling Task inherits the job's access context. + await asyncio.wait_for(service.call_backend(service.sampler.set_template, 'test'), timeout=1) + await service.sampling_guard.run(job, exclusive=True) + assert service.sampler.template_calls == 1 + assert not service.sampling_guard._writer @pytest.mark.asyncio @@ -348,7 +296,6 @@ async def write(): assert wrote.is_set() - @pytest.mark.asyncio async def test_abandoned_waiting_writer_never_starts_and_releases_capacity(): guard = SamplerConcurrency(limit=2) @@ -368,25 +315,21 @@ async def write(): await reader - @pytest.mark.asyncio @pytest.mark.parametrize('dialect', ['tinker', 'twinkle']) -async def test_backend_cancellation_leaves_terminal_failure(dialect): +async def test_backend_cancellation_leaves_terminal_failure(make_service, dialect): from concurrent.futures import CancelledError - service = _Service() + service = make_service() def cancelled(*args, **kwargs): raise CancelledError('engine coroutine cancelled') service.sampler.sample = cancelled - try: - ref = await service.submit(dialect) - record = await service.result(ref) - assert record['status'] == 'failed' - assert record['failure'].reason_code == 'backend_cancelled' - await _until(lambda: not service._compute_worker._detached_tasks) - assert service._compute_worker._detached_reserved == 0 - await service._compute_worker.stop() - assert service.state.records[ref] == record - finally: - await service.close() + ref = await service.submit(dialect) + record = await service.result(ref) + assert record['status'] == 'failed' + assert record['failure'].reason_code == 'backend_cancelled' + await _until(lambda: not service._compute_worker._detached_tasks) + assert service._compute_worker._detached_reserved == 0 + await service._compute_worker.stop() + assert service.state.records[ref] == record diff --git a/tests/server/utils/test_task_queue_mixin.py b/tests/server/utils/test_task_queue_mixin.py index 61d682f3e..ebc3bc25c 100644 --- a/tests/server/utils/test_task_queue_mixin.py +++ b/tests/server/utils/test_task_queue_mixin.py @@ -1,5 +1,6 @@ import asyncio import pytest +import pytest_asyncio from twinkle.server.task_queue.config import TaskQueueConfig from twinkle.server.task_queue.mixin import TaskQueueMixin @@ -362,10 +363,27 @@ async def _factory(): assert env.status == 'completed' +@pytest_asyncio.fixture +async def make_concurrent_queue(): + """Keep concurrent worker setup and teardown in one place.""" + queues = [] + + def create(config=None): + queue = _DummyQueue() + if config is not None: + queue._task_queue_config = config + queue.enable_compute_worker() + queues.append(queue) + return queue + + yield create + for queue in reversed(queues): + await queue._compute_worker.stop() + + @pytest.mark.asyncio -async def test_concurrent_tasks_run_side_by_side(): - queue = _DummyQueue() - queue.enable_compute_worker() +async def test_concurrent_tasks_run_side_by_side(make_concurrent_queue): + queue = make_concurrent_queue() started = [] both_started = asyncio.Event() @@ -377,18 +395,15 @@ async def work(index): await asyncio.wait_for(both_started.wait(), timeout=1) return {'index': index} - try: - refs = [ - await queue.schedule_concurrent_task(lambda i=i: work(i), token='token1', task_type='sample') - for i in range(2) - ] - for _ in range(100): - records = [await queue.state.get_future(ref['request_id']) for ref in refs] - if all(record['status'] == 'completed' for record in records): - break - await asyncio.sleep(0.01) - finally: - await queue._compute_worker.stop() + refs = [ + await queue.schedule_concurrent_task(lambda i=i: work(i), token='token1', task_type='sample') + for i in range(2) + ] + for _ in range(100): + records = [await queue.state.get_future(ref['request_id']) for ref in refs] + if all(record['status'] == 'completed' for record in records): + break + await asyncio.sleep(0.01) assert [record['result'] for record in records] == [{'index': 0}, {'index': 1}] # The serial worker loop is never involved. @@ -396,22 +411,18 @@ async def work(index): @pytest.mark.asyncio -async def test_concurrent_task_keeps_queue_failure_mapping(): - queue = _DummyQueue() - queue.enable_compute_worker() +async def test_concurrent_task_keeps_queue_failure_mapping(make_concurrent_queue): + queue = make_concurrent_queue() async def work(): raise UserTaskError('invalid request') - try: - await queue.schedule_concurrent_task(work, token='token1') - for _ in range(100): - failed = [kwargs for args, kwargs in queue.state.records if args[1] == 'failed'] - if failed: - break - await asyncio.sleep(0) - finally: - await queue._compute_worker.stop() + await queue.schedule_concurrent_task(work, token='token1') + for _ in range(100): + failed = [kwargs for args, kwargs in queue.state.records if args[1] == 'failed'] + if failed: + break + await asyncio.sleep(0) assert [args[1] for args, _ in queue.state.records] == ['pending', 'running', 'failed'] failure = failed[-1]['failure'] @@ -420,10 +431,9 @@ async def work(): @pytest.mark.asyncio -async def test_concurrent_task_preflight_rejection_raises_without_writing(): +async def test_concurrent_task_preflight_rejection_raises_without_writing(make_concurrent_queue): from twinkle.server.exceptions import InputTokensExceededError - queue = _DummyQueue() - queue.enable_compute_worker() + queue = make_concurrent_queue() async def work(): raise AssertionError('preflight rejection must not execute the task') @@ -437,9 +447,8 @@ async def work(): @pytest.mark.asyncio -async def test_worker_stop_cancels_concurrent_task(): - queue = _DummyQueue() - queue.enable_compute_worker() +async def test_worker_stop_cancels_concurrent_task(make_concurrent_queue): + queue = make_concurrent_queue() entered = asyncio.Event() cancelled = asyncio.Event() @@ -458,11 +467,9 @@ async def work(): @pytest.mark.asyncio -async def test_concurrent_capacity_rejects_before_future_write_and_is_reusable(): +async def test_concurrent_capacity_rejects_before_future_write_and_is_reusable(make_concurrent_queue): from twinkle.server.exceptions import ResourceExhaustedError - queue = _DummyQueue() - queue._task_queue_config = TaskQueueConfig(concurrent_task_limit=1) - queue.enable_compute_worker() + queue = make_concurrent_queue(TaskQueueConfig(concurrent_task_limit=1)) release = asyncio.Event() async def work(): @@ -483,13 +490,11 @@ async def work(): assert envelope.status == 'completed' finally: release.set() - await queue._compute_worker.stop() @pytest.mark.asyncio -async def test_concurrent_reservation_is_released_when_state_write_fails(): - queue = _DummyQueue() - queue.enable_compute_worker() +async def test_concurrent_reservation_is_released_when_state_write_fails(make_concurrent_queue): + queue = make_concurrent_queue() async def fail(*args, **kwargs): raise RuntimeError('state unavailable') @@ -502,10 +507,9 @@ async def fail(*args, **kwargs): @pytest.mark.asyncio -async def test_stop_between_future_write_and_start_leaves_terminal_record(): +async def test_stop_between_future_write_and_start_leaves_terminal_record(make_concurrent_queue): from twinkle.server.exceptions import ResourceExhaustedError - queue = _DummyQueue() - queue.enable_compute_worker() + queue = make_concurrent_queue() original = queue.state.store_future_status async def write_then_stop(*args, **kwargs): @@ -521,9 +525,8 @@ async def write_then_stop(*args, **kwargs): @pytest.mark.asyncio -async def test_shutdown_cancel_before_task_start_leaves_failed_future(): - queue = _DummyQueue() - queue.enable_compute_worker() +async def test_shutdown_cancel_before_task_start_leaves_failed_future(make_concurrent_queue): + queue = make_concurrent_queue() executed = False async def work(): nonlocal executed @@ -537,9 +540,8 @@ async def work(): @pytest.mark.asyncio -async def test_cancel_preserves_existing_terminal_record(): - queue = _DummyQueue() - queue.enable_compute_worker() +async def test_cancel_preserves_existing_terminal_record(make_concurrent_queue): + queue = make_concurrent_queue() entered = asyncio.Event() async def work(): entered.set() @@ -553,9 +555,8 @@ async def work(): @pytest.mark.asyncio -async def test_cancel_during_running_status_write_leaves_failed_future(): - queue = _DummyQueue() - queue.enable_compute_worker() +async def test_cancel_during_running_status_write_leaves_failed_future(make_concurrent_queue): + queue = make_concurrent_queue() entered = asyncio.Event() original = queue.state.store_future_status async def write(*args, **kwargs): From e39acd6266e7ba5c958784471d502e824204710f Mon Sep 17 00:00:00 2001 From: Yunnglin Date: Wed, 14 Oct 2026 12:50:33 +0800 Subject: [PATCH 3/5] fix(sampler): narrow concurrency ownership and recover request failures --- .../Usage Guide/Server and Client/Server.md | 7 +- .../\346\234\215\345\212\241\347\253\257.md" | 7 +- src/twinkle/exceptions.py | 5 + .../sampler/vllm_sampler/vllm_sampler.py | 55 +++++-- src/twinkle/server/sampler/app.py | 66 +++------ .../server/sampler/backends/vllm_async.py | 10 ++ src/twinkle/server/sampler/concurrency.py | 140 ++++++++---------- src/twinkle/server/sampler/tinker_handlers.py | 15 +- .../server/sampler/twinkle_handlers.py | 42 ++---- src/twinkle/server/sampler/weights.py | 30 +++- src/twinkle/server/task_queue/worker.py | 31 +++- .../async_rl/vllm_sampler_tq.py | 93 +++++------- .../sampler/test_full_checkpoint_failures.py | 52 +++++++ tests/sampler/test_token_only_sampling.py | 39 +++++ tests/sampler/test_vllm_sampling_lifecycle.py | 47 +++--- .../test_sampler_concurrency_e2e.py | 130 +++++++++++++--- tests/server/sampler/sampling_test_helpers.py | 4 +- .../sampler/test_backend_concurrency.py | 24 +++ tests/server/sampler/test_mock_sampler.py | 4 +- .../sampler/test_resolve_sampler_weights.py | 5 +- .../sampler/test_sampling_concurrency.py | 123 +++++++++------ .../server/sampler/test_sampling_recovery.py | 61 ++++++++ .../server/sampler/test_stream_guarantees.py | 2 +- tests/server/sampler/test_tinker_handlers.py | 13 +- .../server/sampler/test_twinkle_async_rows.py | 2 +- .../test_vllm_sampler_tq_generation.py | 64 ++++---- 26 files changed, 681 insertions(+), 390 deletions(-) create mode 100644 src/twinkle/exceptions.py create mode 100644 src/twinkle/server/sampler/backends/vllm_async.py create mode 100644 tests/sampler/test_full_checkpoint_failures.py create mode 100644 tests/sampler/test_token_only_sampling.py create mode 100644 tests/server/sampler/test_backend_concurrency.py create mode 100644 tests/server/sampler/test_sampling_recovery.py diff --git a/docs/source_en/Usage Guide/Server and Client/Server.md b/docs/source_en/Usage Guide/Server and Client/Server.md index 4b24dc592..114c7041f 100644 --- a/docs/source_en/Usage Guide/Server and Client/Server.md +++ b/docs/source_en/Usage Guide/Server and Client/Server.md @@ -482,7 +482,6 @@ holds the cleanup lease. ### Sampler concurrency and weight changes -- Ordinary Tinker and Twinkle `sample()` requests support concurrency without changing response schemas or future polling. `queue_config.concurrent_task_limit` limits unfinished concurrent tasks per sampler replica and defaults to `128`; capacity exhaustion returns `server`/503. Training compute remains serial. -- During ordinary sampling, full-checkpoint switches wait for existing ordinary and DataPlane generation to finish. Loading a full checkpoint and generating with it cannot overlap these requests. LoRA uses adapter checkpoints; full checkpoints are for full-parameter or merged models. -- After an ordinary sampling request times out, backend cleanup may still be running. New requests are temporarily refused during cleanup. If resource release or the integrity of weights after a failed update cannot be confirmed, the replica keeps returning `server`/503 and must be restarted. -- Unexpected backend cancellation leaves a `failed` future and requires a replica restart. Ordinary generation errors allow subsequent requests once cleanup finishes. +- Ordinary Tinker and Twinkle `sample()` requests support concurrency with unchanged response schemas and future polling; training compute remains serial. `queue_config.concurrent_task_limit` defaults to `128` and bounds admitted, unfinished tasks per replica; exhaustion returns `server`/503. Actual execution concurrency also depends on the backend thread pool, actor, and vLLM settings. +- Full checkpoint loading and generation exclusively own the weights and wait for existing ordinary and DataPlane generation to finish. LoRA uses adapter checkpoints without replacing base weights on each request. +- A sampling timeout fails only its future; capacity stays reserved until the backend finishes, while other sampling can continue. Invalid checkpoints are rejected before writes. A failed weight write fails the health check so Ray Serve replaces the replica. diff --git "a/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/\346\234\215\345\212\241\347\253\257\345\222\214\345\256\242\346\210\267\347\253\257/\346\234\215\345\212\241\347\253\257.md" "b/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/\346\234\215\345\212\241\347\253\257\345\222\214\345\256\242\346\210\267\347\253\257/\346\234\215\345\212\241\347\253\257.md" index 1b72302bf..de96291a4 100644 --- "a/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/\346\234\215\345\212\241\347\253\257\345\222\214\345\256\242\346\210\267\347\253\257/\346\234\215\345\212\241\347\253\257.md" +++ "b/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/\346\234\215\345\212\241\347\253\257\345\222\214\345\256\242\346\210\267\347\253\257/\346\234\215\345\212\241\347\253\257.md" @@ -469,7 +469,6 @@ twinkle-server check-config -c server_config.yaml ### Sampler 并发与权重切换 -- Tinker 和 Twinkle 的普通 `sample()` 支持并发，响应格式与 future 轮询接口保持不变。`queue_config.concurrent_task_limit` 限制单个 sampler 副本尚未结束的并发任务数，默认为 `128`；满载时返回 `server`/503。训练计算任务仍然串行。 -- 普通采样中，切换全量 checkpoint 会等待已有普通生成和 DataPlane 生成结束；全量 checkpoint 加载与使用它生成期间不允许这些请求交错执行。LoRA 使用 adapter checkpoint；全量 checkpoint 用于全参数或合并后的模型。 -- 普通采样请求超时后，后端可能仍在清理。清理期间会暂时拒绝新请求；无法确认资源释放，或权重更新失败后无法确认权重完整性时，副本持续返回 `server`/503，需要重启。 -- 后端意外取消时，future 会进入 `failed` 终态，副本需要重启；普通生成错误在清理完成后不影响后续请求。 +- Tinker 和 Twinkle 的普通 `sample()` 支持并发，响应格式与 future 轮询接口不变；训练计算仍然串行。`queue_config.concurrent_task_limit` 默认为 `128`，限制单个副本已接收但尚未结束的任务数，满载时返回 `server`/503。这是准入上限，实际执行并发还受后端线程池、actor 和 vLLM 配置限制。 +- 全量 checkpoint 的加载和生成独占权重访问，等待已有普通生成和 DataPlane 生成结束。LoRA 使用 adapter checkpoint，无须每次替换基座权重。 +- 采样超时只让该 future 失败；后台任务实际结束后才释放容量，其他采样可以继续。错误 checkpoint 在写入前拒绝；真正的权重写入失败会使健康检查失败，由 Ray Serve 替换副本。 diff --git a/src/twinkle/exceptions.py b/src/twinkle/exceptions.py new file mode 100644 index 000000000..2a1697aea --- /dev/null +++ b/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.""" diff --git a/src/twinkle/sampler/vllm_sampler/vllm_sampler.py b/src/twinkle/sampler/vllm_sampler/vllm_sampler.py index 148e8df4b..ca9b58840 100644 --- a/src/twinkle/sampler/vllm_sampler/vllm_sampler.py +++ b/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 @@ -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) @@ -292,11 +303,18 @@ async def _sample_single( @staticmethod async def _gather_samples(samples) -> List[SampleResponse]: """Finish every prompt before exposing the batch's result or failure.""" - results = await asyncio.gather(*samples, return_exceptions=True) - for result in results: - if isinstance(result, BaseException): - raise result - return results + 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( @@ -308,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. @@ -322,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. @@ -343,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) @@ -382,6 +406,7 @@ 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 self._gather_samples(tasks) @@ -562,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 @@ -578,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 diff --git a/src/twinkle/server/sampler/app.py b/src/twinkle/server/sampler/app.py index a3f4fc977..f0ee357df 100644 --- a/src/twinkle/server/sampler/app.py +++ b/src/twinkle/server/sampler/app.py @@ -7,7 +7,6 @@ """ from __future__ import annotations -import asyncio from fastapi import FastAPI, Request from ray import serve from typing import Any @@ -20,8 +19,7 @@ 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.task_queue.types import BACKEND_TIMEOUT_EXCEPTIONS -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 @@ -44,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. @@ -80,16 +78,16 @@ class SamplerManagement(LazyCleanupMixin, TaskQueueMixin): - 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'), @@ -130,11 +128,12 @@ 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 - self.sampling_guard = SamplerConcurrency(self.task_queue_config.concurrent_task_limit) - # Serialize only the one-time Tinker template initialization. - self.tinker_setup_lock = asyncio.Lock() - self.tinker_template_set = False + # 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() @@ -147,32 +146,11 @@ async def shutdown(self) -> None: 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 call_backend(self, fn, /, *args, admit: bool = True, mutation: bool = False, **kwargs): - - if not admit: - return await super().call_backend(fn, *args, admit=False, **kwargs) - - async def call(): - try: - return await super(SamplerManagement, self).call_backend(fn, *args, admit=admit, **kwargs) - except asyncio.CancelledError: - self.sampling_guard.poison() - raise - except BACKEND_TIMEOUT_EXCEPTIONS: - self.sampling_guard.poison() - raise - except Exception: - if mutation: - self.sampling_guard.poison() - raise - - if self.sampling_guard.owns_job(): - return await call() - # Control-plane operations also own their backend lifetime. Sampling - # calls inside a job reuse its access rather than acquiring another lock. - return await self.sampling_guard.run(call, exclusive=True) + 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): diff --git a/src/twinkle/server/sampler/backends/vllm_async.py b/src/twinkle/server/sampler/backends/vllm_async.py new file mode 100644 index 000000000..310eb0414 --- /dev/null +++ b/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 diff --git a/src/twinkle/server/sampler/concurrency.py b/src/twinkle/server/sampler/concurrency.py index 87608729a..71d7e7c68 100644 --- a/src/twinkle/server/sampler/concurrency.py +++ b/src/twinkle/server/sampler/concurrency.py @@ -1,116 +1,92 @@ # Copyright (c) ModelScope Contributors. All rights reserved. -"""Let each background sampling job own access until its backend has exited.""" +"""Coordinate sampling with changes to shared sampler state.""" from __future__ import annotations import asyncio -from collections.abc import Callable, Coroutine -from contextvars import ContextVar -from typing import Any -from twinkle.server.exceptions import ResourceExhaustedError from twinkle.server.task_queue.types import BackendBusyError class SamplerConcurrency: - """Stable weights admit readers; changing base weights requires one writer. + """Readers share stable weights; mutations wait for every reader to finish. - A job owns its entire load/generate/cleanup sequence. Cancelling its waiter - leaves that job running, so a request timeout cannot release live weights. - All bookkeeping runs on the replica's event loop. + Jobs retain their lock if an HTTP waiter leaves. Admission and execution + timeouts belong to ComputeWorker, not this lock. """ - def __init__(self, limit: int = 128): - self._limit = limit + def __init__(self): self._readers = 0 self._writer = False self._waiting_writers = 0 - self._changed = asyncio.Event() + self._condition = asyncio.Condition() self._jobs: set[asyncio.Task] = set() - self._draining: set[asyncio.Task] = set() - # Nested polling tasks inherit the owner, without passing lock handles - # through every backend call. Membership prevents use after job exit. - self._owner: ContextVar[asyncio.Task | None] = ContextVar('sampler_job', default=None) - self._uncertain = False + self._weights_error: Exception | None = None self._closed = False - def owns_job(self) -> bool: - owner = self._owner.get() - return owner is not None and not owner.done() and owner in self._jobs - - def _notify(self) -> None: - self._changed.set() - self._changed = asyncio.Event() - - def check_available(self, *, admitted: bool = False) -> None: + def check_available(self): if self._closed: raise BackendBusyError('The sampler is shutting down.') - if self._uncertain: - raise BackendBusyError('Sampler weights or backend completion are uncertain; restart this replica.') - if self._draining and not admitted: - raise BackendBusyError('The sampler is waiting for backend resource release.') + if self._weights_error is not None: + raise BackendBusyError('A checkpoint write failed; the sampler is restarting.') - def poison(self) -> None: - self._uncertain = True - self._notify() + def invalidate_weights(self, error: Exception): + self._weights_error = error - def close(self) -> None: + def close(self): self._closed = True - self._notify() - - async def _acquire(self, exclusive: bool) -> None: - if exclusive: - self._waiting_writers += 1 - try: - while True: - self.check_available(admitted=True) - if asyncio.current_task() in self._draining: - raise BackendBusyError('Sampling request exited before its job could start.') - if not self._writer and (self._readers == 0 if exclusive else self._waiting_writers == 0): - if exclusive: - self._writer = True - else: - self._readers += 1 - return - await self._changed.wait() - finally: - if exclusive: - self._waiting_writers -= 1 - self._notify() async def _execute(self, factory, exclusive: bool): - await self._acquire(exclusive) - token = self._owner.set(asyncio.current_task()) + 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: - self._owner.reset(token) - if exclusive: - self._writer = False - else: - self._readers -= 1 - self._notify() - - def _finished(self, job: asyncio.Task) -> None: - self._jobs.discard(job) - self._draining.discard(job) - if not job.cancelled(): - job.exception() # Observe failures even when its waiter already left. - self._notify() - - async def run(self, factory: Callable[[], Coroutine[Any, Any, Any]], *, exclusive: bool = False): - """Wait for one owned job without transferring cancellation to it.""" + 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() - # Unlike the future queue's admission count, this capacity is retained - # until actual backend cleanup, including jobs outliving request timeout. - if len(self._jobs) >= self._limit: - raise ResourceExhaustedError('Concurrent sampling capacity is unavailable on this replica.') - job = asyncio.create_task(self._execute(factory, exclusive)) + started = False + + async def call(): + nonlocal started + started = True + return await factory() + + job = asyncio.create_task(self._execute(call, exclusive)) self._jobs.add(job) - job.add_done_callback(self._finished) + + 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 job.done(): - self._draining.add(job) - self._notify() + 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) diff --git a/src/twinkle/server/sampler/tinker_handlers.py b/src/twinkle/server/sampler/tinker_handlers.py index 635a9ea36..719ab0621 100644 --- a/src/twinkle/server/sampler/tinker_handlers.py +++ b/src/twinkle/server/sampler/tinker_handlers.py @@ -22,7 +22,6 @@ from twinkle.server.checkpoint import create_checkpoint_manager from twinkle.server.sampler.weights import sample_with_weights from twinkle.server.task_queue.types import UserTaskError -from twinkle.server.utils import get_template_for_model from twinkle.utils.logger import get_logger logger = get_logger() @@ -104,25 +103,13 @@ async def _do_sample(): logprobs=1, ) - # Requests run concurrently, so sampler setup is serialized here. - # The template depends only on the model and is set once: every - # set_template or reset_prefix_cache call first waits for all - # in-flight generations (infra ``_cw_barrier``), so doing it per - # request would serialize sampling again. No prefix-cache reset is - # needed: vLLM caches LoRA requests under their own adapter name, - # and load_full_weights_from_path resets the cache itself. - async with self.tinker_setup_lock: - if not self.tinker_template_set: - template = get_template_for_model(self.model_id) - await self.call_backend(self.sampler.set_template, template, model_id=self.model_id) - self.tinker_template_set = True - async def generate(lora_path): return await self.call_backend( self.sampler.sample, inputs=[prompt_inputs] * body.num_samples, sampling_params=sampling_params, adapter_path=lora_path, + token_only=True, ) responses = await sample_with_weights(self, adapter_uri, generate) diff --git a/src/twinkle/server/sampler/twinkle_handlers.py b/src/twinkle/server/sampler/twinkle_handlers.py index 2bef2ae42..b6cf58d3b 100644 --- a/src/twinkle/server/sampler/twinkle_handlers.py +++ b/src/twinkle/server/sampler/twinkle_handlers.py @@ -10,7 +10,6 @@ import asyncio import json -import traceback import uuid from collections.abc import Callable from fastapi import Depends, FastAPI, Request @@ -29,7 +28,7 @@ from twinkle.protocol.types import sampler as sampler_types from twinkle.server.exceptions import EndpointUnavailableError, RequestRejectedError from twinkle.server.lifecycle.submit import backend_kwargs, resolve_twinkle_adapter_name, to_backend_inputs -from twinkle.server.sampler.weights import sample_with_weights +from twinkle.server.sampler.weights import resolve_sampler_weights, sample_with_weights from twinkle.server.task_errors import task_error_payload from twinkle.server.telemetry.correlation import MODEL_ID from twinkle.server.telemetry.tracing import traced_operation @@ -163,18 +162,14 @@ async def _stream_queue(q, sentinel, request_id: str, total_timeout: float, sing pass -async def _await_generation(service: SamplerManagement, submission_id: str, timeout: float): +async def _await_generation(service: SamplerManagement, submission_id: str): """Poll one admitted generation through the backend boundary.""" collected = False - loop = asyncio.get_running_loop() - deadline = loop.time() + timeout async def poll(): nonlocal collected poll_interval = 0.01 while True: - if loop.time() >= deadline: - raise asyncio.TimeoutError try: states = _submission_states(await service.call_backend(service.sampler.get_generation_status, submission_id)) @@ -204,20 +199,10 @@ async def poll(): finally: if not collected: - async def cancel_and_wait(): - while True: - states = _submission_states(await service.call_backend(service.sampler.cancel_generation, - submission_id)) - if states and all(state.get('status') in ('completed', 'missing') for state in states): - return - await asyncio.sleep(0.01) - try: - await asyncio.wait_for(cancel_and_wait(), timeout=4.0) - except BaseException: - service.sampling_guard.poison() - logger.warning('Generation cleanup did not confirm completion: %s', submission_id, exc_info=True) - # Preserve the original generation error, including cancellation. + await service.call_backend(service.sampler.cancel_generation, submission_id) + except Exception: + logger.warning('Failed to clear generation %s', submission_id, exc_info=True) def _register_twinkle_sampler_routes(app: FastAPI, self_fn: Callable[[], SamplerManagement]) -> None: @@ -321,7 +306,7 @@ async def generate(): adapter_name=full_adapter_name, adapter_path=adapter_path, ) - return await _await_generation(self, submission_id, self.task_queue_config.effective_execution_timeout) + return await _await_generation(self, submission_id) responses = await self.sampling_guard.run(generate) rows, tags = _build_rollout_rows_and_tags( @@ -333,7 +318,7 @@ async def generate(): ref = await self.data_plane.put([json_safe(item) for item in rows], kind='rollout', tags=tags) return ref.model_dump() - return await self.submit_background_and_peek( + return await self.submit_concurrent_and_peek( _generate_and_store, model_id=full_adapter_name or None, task_type='sample_to_data_plane') @app.post('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/twinkle/unload_adapter_paths') @@ -353,7 +338,7 @@ async def unload_adapter_paths( resolved_paths.append(adapter_path) unload = getattr(self.sampler, 'unload_adapter_paths', None) if unload is not None: - await self.call_backend(unload, resolved_paths) + await self.sampling_guard.run(lambda: self.call_backend(unload, resolved_paths), exclusive=True) return {'status': 'ok'} @app.post('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/twinkle/set_template', response_model=sampler_types.SamplerSetTemplateResponse) @@ -364,7 +349,9 @@ async def set_template( ) -> sampler_types.SamplerSetTemplateResponse: """Set the chat template for encoding Trajectory inputs.""" with traced_operation('sampler.set_template'): - await self.call_backend(self.sampler.set_template, body.template_cls, **backend_kwargs(body)) + await self.sampling_guard.run( + lambda: self.call_backend(self.sampler.set_template, body.template_cls, **backend_kwargs(body)), + exclusive=True) return sampler_types.SamplerSetTemplateResponse() @app.post('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/twinkle/add_adapter_to_sampler', response_model=sampler_types.SamplerAddAdapterResponse) @@ -385,7 +372,9 @@ async def add_adapter_to_sampler( config = LoraConfig(**body.config) if isinstance(body.config, dict) else body.config with traced_operation('sampler.add_adapter_to_sampler', attrs={MODEL_ID: self.model_id}): - await self.call_backend(self.sampler.add_adapter_to_sampler, full_adapter_name, config) + await self.sampling_guard.run( + lambda: self.call_backend(self.sampler.add_adapter_to_sampler, full_adapter_name, config), + exclusive=True) return sampler_types.SamplerAddAdapterResponse(adapter_name=full_adapter_name) @@ -397,7 +386,8 @@ async def apply_patch( ) -> None: patch_cls = deserialize_object(body.patch_cls) with traced_operation('sampler.apply_patch'): - await self.call_backend(self.sampler.apply_patch, patch_cls, **backend_kwargs(body)) + await self.sampling_guard.run( + lambda: self.call_backend(self.sampler.apply_patch, patch_cls, **backend_kwargs(body)), exclusive=True) @app.post('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/twinkle/sample_stream') async def sample_stream( diff --git a/src/twinkle/server/sampler/weights.py b/src/twinkle/server/sampler/weights.py index 5c8363e00..f2d1077f7 100644 --- a/src/twinkle/server/sampler/weights.py +++ b/src/twinkle/server/sampler/weights.py @@ -2,19 +2,43 @@ """Run checkpoint loading and generation as one owned sampling job.""" from __future__ import annotations +import glob import os +from safetensors import SafetensorError + +from twinkle.exceptions import FullWeightLoadError +from twinkle.server.task_queue.types import BACKEND_TIMEOUT_EXCEPTIONS, UserTaskError def is_full_checkpoint(resolved_uri: str | None) -> bool: - return bool(resolved_uri and not os.path.exists(os.path.join(resolved_uri, 'adapter_config.json'))) + if not resolved_uri: + return False + if os.path.isfile(os.path.join(resolved_uri, 'adapter_config.json')): + return False + if os.path.isabs(resolved_uri) and not os.path.isdir(resolved_uri): + raise UserTaskError(f'Checkpoint directory does not exist: {resolved_uri}') + if os.path.isdir(resolved_uri) and not (os.path.isfile(os.path.join(resolved_uri, 'model.safetensors.index.json')) + or glob.glob(os.path.join(resolved_uri, '*.safetensors'))): + raise UserTaskError(f'No .safetensors weights found under {resolved_uri}') + return True async def resolve_sampler_weights(service, resolved_uri: str | None, *, full_checkpoint=None) -> str | None: """Load base weights or return the adapter path.""" - full_checkpoint = is_full_checkpoint(resolved_uri) if full_checkpoint is None else full_checkpoint + if full_checkpoint is None: + full_checkpoint = is_full_checkpoint(resolved_uri) + if full_checkpoint: + return await service.sampling_guard.run( + lambda: resolve_sampler_weights(service, resolved_uri, full_checkpoint=True), exclusive=True) if not full_checkpoint: return resolved_uri or None - await service.call_backend(service.sampler.load_full_weights_from_path, resolved_uri, mutation=True) + try: + await service.call_backend(service.sampler.load_full_weights_from_path, resolved_uri) + except (FullWeightLoadError, *BACKEND_TIMEOUT_EXCEPTIONS) as error: + service.sampling_guard.invalidate_weights(error) + raise + except (ValueError, OSError, KeyError, SafetensorError) as error: + raise UserTaskError(str(error)) from error return None diff --git a/src/twinkle/server/task_queue/worker.py b/src/twinkle/server/task_queue/worker.py index 439f0c3e6..307b469c9 100644 --- a/src/twinkle/server/task_queue/worker.py +++ b/src/twinkle/server/task_queue/worker.py @@ -12,6 +12,7 @@ import time import traceback from collections import deque +from dataclasses import replace from typing import TYPE_CHECKING, Any, Callable, Deque from twinkle.protocol.types.errors import ErrorCategory @@ -322,7 +323,35 @@ def run_detached(self, task: QueuedTask, queue_key: str) -> None: """ if self._stopping: raise ResourceExhaustedError('The replica is shutting down.') - running = asyncio.create_task(self._process_task(task, queue_key)) + execution = None + + def start(): + nonlocal execution + execution = asyncio.create_task(task.coro_factory()) + return asyncio.shield(execution) + + async def run(): + try: + await self._process_task(replace(task, coro_factory=start), queue_key) + finally: + # A timeout ends the client future, not the backend operation. + # Retain this task's sole capacity reservation until it exits. + if execution is not None: + if not execution.done(): + execution.cancel() + while not execution.done(): + try: + await asyncio.shield(execution) + except asyncio.CancelledError: + if not execution.done() and not execution.cancelling(): + execution.cancel() + continue + except Exception: + break + if not execution.cancelled(): + execution.exception() + + running = asyncio.create_task(run()) self._detached_tasks[running] = task def done(completed: asyncio.Task) -> None: diff --git a/src/twinkle_agentic/async_rl/vllm_sampler_tq.py b/src/twinkle_agentic/async_rl/vllm_sampler_tq.py index 657fa0300..fcdd21e96 100644 --- a/src/twinkle_agentic/async_rl/vllm_sampler_tq.py +++ b/src/twinkle_agentic/async_rl/vllm_sampler_tq.py @@ -17,7 +17,7 @@ from twinkle.data_format import SampledSequence, SampleResponse, SamplingParams, user_data_get from twinkle.hub import HubOperation from twinkle.metric import MetricBuffer, MetricRecord -from twinkle.sampler.vllm_sampler.vllm_sampler import SAMPLER_MAX_CONCURRENCY, vLLMSampler +from twinkle.sampler.vllm_sampler import vLLMSampler from .data_plane import TQDataPlane from .metrics import rollout_metrics from .types import LoraContext, PromptGroup, RolloutOutput, RolloutPolicy @@ -85,14 +85,6 @@ def _compute_reward_metrics( return dict(metric_payload(rollout_rows, rewards=rewards, context=context)) -@dataclass -class _GenerationSubmission: - """Separate a cancellation request from actual coroutine completion.""" - future: Future | None = None - task: asyncio.Task | None = None - cancel_requested: bool = False - - @dataclass(frozen=True) class _GeneratedSample: response: SampleResponse @@ -121,14 +113,13 @@ class _PromptGroupRolloutStats: policy_versions: tuple[int, ...] -@remote_class(max_concurrency=SAMPLER_MAX_CONCURRENCY) +@remote_class() class VLLMSamplerTQ(vLLMSampler): - """vLLM sampler for native async RL and client-server generation. + """vLLM sampler that writes async RL rollout results directly to TransferQueue. - ``submit_prompt_groups()`` schedules native rollouts and returns metadata; - their results go to TransferQueue. Inherited ``sample()`` returns completed - responses. The client-server path uses ``submit_generation()``, status - polling and ``collect_generation()`` for non-blocking admission. + ``sample()`` is intentionally fire-and-forget: it schedules generation work + on the sampler actor's vLLM event loop and returns submission metadata + without waiting for any prompt group to finish. """ def __init__( @@ -166,7 +157,7 @@ def __init__( # path. Unlike ``_background_submissions`` above, their results must # remain available until SamplerManagement collects them and writes # them to the opaque client DataPlane. - self._generation_submissions: dict[str, _GenerationSubmission] = {} + self._generation_submissions: dict[str, Future[list[SampleResponse]]] = {} self.metric_buffer = MetricBuffer() self._failure: str | None = None @@ -258,34 +249,27 @@ def submit_generation( """ if submission_id in self._generation_submissions: raise KeyError(f'generation submission already exists: {submission_id}') - submission = _GenerationSubmission() - - async def generate(): - submission.task = asyncio.current_task() - if submission.cancel_requested: - raise asyncio.CancelledError - return await self._generate_inputs( + future = self._submit_in_loop( + self._generate_inputs( inputs, sampling_params, adapter_name=adapter_name, adapter_path=adapter_path, - use_base_model=use_base_model) - - submission.future = self._submit_in_loop(generate()) - self._generation_submissions[submission_id] = submission + use_base_model=use_base_model, + )) + self._generation_submissions[submission_id] = future return {'submission_id': submission_id, 'status': 'running'} @remote_function(dispatch='all', collect='none', lazy_collect=False) def get_generation_status(self, submission_id: str) -> dict[str, Any]: """Return this DP worker's submission state without waiting.""" - submission = self._generation_submissions.get(submission_id) - if submission is None: + future = self._generation_submissions.get(submission_id) + if future is None: return { 'submission_id': submission_id, 'status': 'missing', 'error': f'unknown generation submission: {submission_id}', } - future = submission.future if future.cancelled(): return {'submission_id': submission_id, 'status': 'cancelled'} if not future.done(): @@ -302,10 +286,9 @@ def get_generation_status(self, submission_id: str) -> dict[str, Any]: @remote_function(dispatch='all', collect='flatten', lazy_collect=False) def collect_generation(self, submission_id: str) -> list[SampleResponse]: """Consume completed responses from every DP worker.""" - submission = self._generation_submissions.get(submission_id) - if submission is None: + future = self._generation_submissions.get(submission_id) + if future is None: raise KeyError(f'unknown generation submission: {submission_id}') - future = submission.future if not future.done(): raise RuntimeError(f'generation submission is still running: {submission_id}') try: @@ -315,34 +298,30 @@ def collect_generation(self, submission_id: str) -> list[SampleResponse]: @remote_function(dispatch='all', collect='none', lazy_collect=False) def cancel_generation(self, submission_id: str) -> dict[str, Any]: - """Request cancellation; forget only after vLLM's coroutine has exited. - - Cancelling run_coroutine_threadsafe's Future reports success before the - coroutine's abort/finally runs. Cancel the loop-owned Task instead so the - Future becomes terminal only after actual cleanup. - """ - submission = self._generation_submissions.get(submission_id) - if submission is None: + """Cancel and forget one generation submission on every DP worker.""" + future = self._generation_submissions.pop(submission_id, None) + if future is None: return {'submission_id': submission_id, 'status': 'missing'} - if submission.future.done(): - self._generation_submissions.pop(submission_id, None) - return {'submission_id': submission_id, 'status': 'completed'} - if not submission.cancel_requested: - submission.cancel_requested = True - - def cancel(): - if submission.task is not None and not submission.task.done(): - submission.task.cancel() - - self._async_loop.call_soon_threadsafe(cancel) - return {'submission_id': submission_id, 'status': 'cancellation_requested'} + was_done = future.done() + cancelled = future.cancel() + if cancelled: + status = 'cancelled' + elif was_done: + status = 'completed' + else: + status = 'cancellation_requested' + return { + 'submission_id': submission_id, + 'status': status, + } @remote_function(dispatch='all', collect='none', lazy_collect=False) def cancel_all_generations(self) -> dict[str, int]: - """Request cancellation of retained submissions during shutdown.""" - ids = list(self._generation_submissions) - requested = sum(self.cancel_generation(key)['status'] == 'cancellation_requested' for key in ids) - return {'submissions': len(ids), 'cancelled': requested} + """Cancel all retained CS submissions during replica shutdown.""" + submissions = list(self._generation_submissions.values()) + self._generation_submissions.clear() + cancelled = sum(future.cancel() for future in submissions if not future.done()) + return {'submissions': len(submissions), 'cancelled': cancelled} @remote_function(dispatch='slice_dp', collect='flatten', lazy_collect=False) def evaluate( diff --git a/tests/sampler/test_full_checkpoint_failures.py b/tests/sampler/test_full_checkpoint_failures.py new file mode 100644 index 000000000..755e4dd6c --- /dev/null +++ b/tests/sampler/test_full_checkpoint_failures.py @@ -0,0 +1,52 @@ +"""Validate checkpoints before writing; distinguish partial engine writes.""" +import asyncio +import json +from types import SimpleNamespace + +import pytest +import torch +from safetensors.torch import save_file + +from twinkle.exceptions import FullWeightLoadError +from twinkle.sampler.vllm_sampler.vllm_sampler import vLLMSampler + + +@pytest.mark.parametrize('kind', ['missing', 'empty', 'bad_index', 'missing_shard', 'corrupt_shard']) +def test_invalid_checkpoint_does_not_write_weights(tmp_path, kind): + path = tmp_path / 'checkpoint' + if kind != 'missing': + path.mkdir() + if kind == 'bad_index': + (path / 'model.safetensors.index.json').write_text('{') + if kind == 'missing_shard': + (path / 'model.safetensors.index.json').write_text(json.dumps({'weight_map': {'weight': 'absent.safetensors'}})) + if kind == 'corrupt_shard': + (path / 'model.safetensors').write_text('not a checkpoint') + sampler = object.__new__(vLLMSampler) + # This existing local directory avoids a hub request for the missing case. + from unittest.mock import patch + writes = [] + sampler.model_id = str(path) + sampler.engine = SimpleNamespace(update_weights=lambda *a, **k: writes.append(True)) + with patch('twinkle.sampler.vllm_sampler.vllm_sampler.HubOperation.download_model', return_value=str(path)): + with pytest.raises(Exception) as error: + sampler.load_full_weights_from_path(str(path)) + assert not isinstance(error.value, FullWeightLoadError) + assert writes == [] + assert not hasattr(sampler, '_loaded_full_weights_path') + + +def test_partial_write_is_reported_as_weight_failure(tmp_path): + save_file({'weight': torch.ones(1)}, str(tmp_path / 'model.safetensors')) + sampler = object.__new__(vLLMSampler) + writes = [] + async def update(weights, **kwargs): + writes.extend(list(weights)) + raise RuntimeError('engine failed after writing') + sampler.engine = SimpleNamespace(update_weights=update) + sampler._run_in_loop = asyncio.run + with pytest.raises(FullWeightLoadError) as error: + sampler.load_full_weights_from_path(str(tmp_path)) + assert isinstance(error.value.__cause__, RuntimeError) + assert len(writes) == 1 + assert not hasattr(sampler, '_loaded_full_weights_path') diff --git a/tests/sampler/test_token_only_sampling.py b/tests/sampler/test_token_only_sampling.py new file mode 100644 index 000000000..f9aeea903 --- /dev/null +++ b/tests/sampler/test_token_only_sampling.py @@ -0,0 +1,39 @@ +"""Tinker consumes encoded tokens, independently of native template settings.""" +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest + +from twinkle.data_format import SampledSequence, SampleResponse, SamplingParams +from twinkle.sampler.vllm_sampler import vLLMSampler + + +class _OtherClientTemplate: + def get_vllm_input_ids(self, *args): + raise AssertionError('Token-only sampling rewrote encoded input IDs') + + def concat_input_feature(self, *args): + raise AssertionError('Token-only sampling built unused training features') + + def decode(self, *args): + raise AssertionError('Token-only sampling decoded unused text') + + +@pytest.mark.asyncio +@pytest.mark.parametrize('template', [None, _OtherClientTemplate()]) +@pytest.mark.parametrize('logprobs_only', [False, True]) +async def test_token_only_sampling_needs_no_template(template, logprobs_only): + sampler = object.__new__(vLLMSampler) + sampler.template = template + sampler.engine = SimpleNamespace(sample=AsyncMock(return_value=SampleResponse( + prompt_token_ids=[7], + sequences=[SampledSequence(tokens=[8], logprobs=[[(8, -.1)]], stop_reason='length')], + prompt_logprobs=[-.2], + ))) + result = await sampler._sample_single( + {'input_ids': [7]}, SamplingParams(max_tokens=1), token_only=True, logprobs_only=logprobs_only) + assert sampler.engine.sample.call_args.kwargs['prompt'] == [7] + assert result.prompt_token_ids == [7] and result.prompt_logprobs == [-.2] + assert result.sequences[0].tokens == ([] if logprobs_only else [8]) + assert result.sequences[0].logprobs == (None if logprobs_only else [[(8, -.1)]]) + assert result.sequences[0].decoded is None and result.sequences[0].new_input_feature is None diff --git a/tests/sampler/test_vllm_sampling_lifecycle.py b/tests/sampler/test_vllm_sampling_lifecycle.py index abe5c1bb8..08fe609fa 100644 --- a/tests/sampler/test_vllm_sampling_lifecycle.py +++ b/tests/sampler/test_vllm_sampling_lifecycle.py @@ -1,34 +1,39 @@ -"""A batch owns generation until all of its prompts have exited.""" +"""Batch failure cancels siblings and waits for their actual cleanup.""" +import asyncio + import pytest +from twinkle.sampler.vllm_sampler import vLLMSampler @pytest.mark.asyncio -async def test_failed_prompt_waits_for_sibling_generation_to_exit(): - import asyncio - from twinkle.data_format import SamplingParams - from twinkle.sampler.vllm_sampler import vLLMSampler - sampler = object.__new__(vLLMSampler) - sampler.template = None - sampler._run_in_loop = lambda coroutine: coroutine +@pytest.mark.parametrize('cancel_batch', [False, True]) +async def test_batch_waits_for_cancelled_sibling_cleanup(cancel_batch): entered = asyncio.Event() release = asyncio.Event() + cleanup = asyncio.Event() finished = asyncio.Event() - async def single(feat, *args, **kwargs): - if feat['input_ids'][0] == 1: - raise ValueError('bad prompt') - entered.set() - await release.wait() - finished.set() - return object() - sampler._sample_single = single - generation = asyncio.create_task(sampler.sample( - [{'input_ids': [1]}, {'input_ids': [2]}], SamplingParams(max_tokens=1))) + async def failed(): + await entered.wait() + if cancel_batch: + await asyncio.Future() + raise ValueError('bad prompt') + async def sibling(): + try: + entered.set() + await asyncio.Future() + finally: + cleanup.set() + await release.wait() + finished.set() + generation = asyncio.create_task(vLLMSampler._gather_samples([failed(), sibling()])) try: await entered.wait() - await asyncio.sleep(.01) - assert not generation.done(), 'a failed prompt returned while its sibling was still generating' + if cancel_batch: + generation.cancel() + await asyncio.wait_for(cleanup.wait(), 1) + assert not generation.done() release.set() - with pytest.raises(ValueError, match='bad prompt'): + with pytest.raises(asyncio.CancelledError if cancel_batch else ValueError): await generation assert finished.is_set() finally: diff --git a/tests/server/integration/test_sampler_concurrency_e2e.py b/tests/server/integration/test_sampler_concurrency_e2e.py index 9ca2336e4..6b051c161 100644 --- a/tests/server/integration/test_sampler_concurrency_e2e.py +++ b/tests/server/integration/test_sampler_concurrency_e2e.py @@ -94,13 +94,115 @@ def test_invalid_generation_does_not_disable_sampler(clients, dialect): _report(check='failure_recovery', dialect=dialect) +@pytest.mark.parametrize('dialect', ['tinker', 'twinkle']) +def test_invalid_checkpoint_does_not_disable_sampler(clients, tmp_path, dialect): + ids = _prompt(clients, 'Reply with exactly this word and nothing else: apple') + corrupt = tmp_path / 'corrupt' + corrupt.mkdir() + (corrupt / 'model.safetensors').write_text('invalid checkpoint') + for path in (str(tmp_path), str(tmp_path / 'missing'), str(corrupt)): + handle = _checkpoint_handles(clients, dialect, [path])[0] + with pytest.raises(Exception): + _sample(clients, dialect, ids, 2, adapter=handle) + assert _sample(clients, dialect, ids, 2)[0][0] + _report(check='invalid_checkpoint_recovery', dialect=dialect) + + +def test_tinker_tokens_do_not_use_native_template(clients): + import ray + actors = _real_sampler_actors() + ids = _prompt(clients, 'Reply with exactly this word and nothing else: apple') + expected = _sample(clients, 'tinker', ids, 2)[0][0] + + def install(actor): + class OtherClientTemplate: + def get_vllm_input_ids(self, *args): + raise AssertionError('Tinker used the native template') + def concat_input_feature(self, *args): + raise AssertionError('Tinker built unused training features') + def decode(self, *args): + raise AssertionError('Tinker decoded unused text') + actor._e2e_template = actor.template + actor.template = OtherClientTemplate() + + def restore(actor): + actor.template = actor._e2e_template + del actor._e2e_template + + ray.get([actor.__ray_call__.remote(install) for actor in actors], timeout=10) + try: + assert _sample(clients, 'tinker', ids, 2)[0][0] == expected + finally: + ray.get([actor.__ray_call__.remote(restore) for actor in actors], timeout=10) + ray.shutdown() + + +@pytest.mark.parametrize('dialect', ['tinker', 'twinkle']) +def test_sampling_timeout_keeps_backend_owned_without_blocking_readers(clients, dialect): + import ray + from ray.util.state import list_actors + _real_sampler_actors() + replicas = [a for a in list_actors(address=os.environ['TWINKLE_TEST_RAY_DASHBOARD_URL'], detail=True) + if a.state == 'ALIVE' and a.class_name.startswith('ServeReplica') and 'SamplerManagement' in a.class_name] + assert len(replicas) == 1 + replica = ray.get_actor(replicas[0].name, namespace='serve') + + def install(actor): + service = actor._replica_impl._user_callable_wrapper._callable + service._e2e_timeout = service.task_queue_config.execution_timeout + service._e2e_sample = service.sampler.sample + service._e2e_release = threading.Event() + service._e2e_entered = threading.Event() + service.task_queue_config.execution_timeout = .1 + original, entered, release = service._e2e_sample, service._e2e_entered, service._e2e_release + def held(*args, **kwargs): + result = original(*args, **kwargs) + entered.set() + assert release.wait(60) + return result + service.sampler.sample = held + + def allow_next(actor): + service = actor._replica_impl._user_callable_wrapper._callable + service.task_queue_config.execution_timeout = service._e2e_timeout + service.sampler.sample = service._e2e_sample + return service._e2e_entered.is_set(), service._compute_worker._detached_reserved + + def restore(actor): + service = actor._replica_impl._user_callable_wrapper._callable + service._e2e_release.set() + service.task_queue_config.execution_timeout = service._e2e_timeout + service.sampler.sample = service._e2e_sample + for name in ('_e2e_timeout', '_e2e_sample', '_e2e_release', '_e2e_entered'): + delattr(service, name) + + ids = _prompt(clients, 'Reply with exactly this word and nothing else: apple') + ray.get(replica.__ray_call__.remote(install), timeout=10) + try: + with pytest.raises(Exception, match='timed out|timeout'): + _sample(clients, dialect, ids, 2) + deadline = time.monotonic() + 30 + while True: + entered, reserved = ray.get(replica.__ray_call__.remote(allow_next), timeout=10) + if entered: + break + assert time.monotonic() < deadline + time.sleep(.01) + assert reserved == 1 + assert _sample(clients, dialect, ids, 2)[0][0] + _report(check='sampling_timeout_recovery', dialect=dialect, retained_capacity=reserved) + finally: + ray.get(replica.__ray_call__.remote(restore), timeout=10) + ray.shutdown() + + def test_async_generation_cleanup_and_recovery_on_real_engine(clients): import ray from ray.util.state import list_actors ray.init(address=os.environ['RAY_ADDRESS'], namespace='default', ignore_reinit_error=True, log_to_driver=False) actors = list_actors(address=os.environ.get('TWINKLE_TEST_RAY_DASHBOARD_URL', 'http://127.0.0.1:8395'), - filters=[('class_name', '=', 'VLLMSamplerTQ'), ('state', '=', 'ALIVE')], detail=True) + filters=[('class_name', '=', os.environ.get('TWINKLE_TEST_SAMPLER_CLASS', 'ServerVLLMSampler')), ('state', '=', 'ALIVE')], detail=True) assert len(actors) == int(os.environ.get('TWINKLE_TEST_DP', '1')) assert all(actor.state == 'ALIVE' for actor in actors) actor = ray.get_actor(actors[0].name, namespace='default') @@ -124,8 +226,6 @@ def wait_status(key, terminal): batch = actor.sample.remote( [{'input_ids': ids}, {'input_ids': [10**9]}], {'max_tokens': 1024, 'temperature': 0}) - ready, _ = ray.wait([batch], timeout=.2) - assert not ready, 'ordinary batch returned before its sibling generation exited' with pytest.raises(ray.exceptions.RayTaskError): ray.get(batch, timeout=30) @@ -133,9 +233,8 @@ def wait_status(key, terminal): {'max_tokens': 1024, 'temperature': 0}) time.sleep(.1) assert call('get_generation_status', 'e2e-cancel')['status'] == 'running' - assert call('cancel_generation', 'e2e-cancel')['status'] == 'cancellation_requested' - wait_status('e2e-cancel', {'cancelled'}) - assert call('cancel_generation', 'e2e-cancel')['status'] == 'completed' + assert call('cancel_generation', 'e2e-cancel')['status'] == 'cancelled' + assert call('get_generation_status', 'e2e-cancel')['status'] == 'missing' call('submit_generation', 'e2e-after', [{'input_ids': ids}], {'max_tokens': 8, 'temperature': 0}) wait_status('e2e-after', {'completed'}) @@ -183,11 +282,7 @@ def generate_and_fetch(): def test_backend_self_cancellation_terminates_both_sdk_futures(clients): - """Cancel at the driver boundary after both requests performed real PPU generation. - - Submit both requests before cancellation quarantines the replica. This test - runs last because uncertain backend cancellation requires a replica restart. - """ + """A backend cancellation terminates futures and allows subsequent sampling.""" import ray from ray.util.state import list_actors actors = _real_sampler_actors() @@ -212,15 +307,16 @@ def cancelled(*args, **kwargs): def restore(actor): service = actor._replica_impl._user_callable_wrapper._callable - service.sampler.sample = service.sampler._e2e_original_sample - del service.sampler._e2e_original_sample + if hasattr(service.sampler, '_e2e_original_sample'): + service.sampler.sample = service.sampler._e2e_original_sample + del service.sampler._e2e_original_sample def status(actor): service = actor._replica_impl._user_callable_wrapper._callable return {'jobs': len(service.sampling_guard._jobs), 'detached': len(service._compute_worker._detached_tasks), 'reserved': service._compute_worker._detached_reserved, - 'quarantined': service.sampling_guard._uncertain} + 'weights_invalid': service.sampling_guard._weights_error is not None} ids = _prompt(clients, 'Reply with exactly this word and nothing else: apple') ray.get(replica.__ray_call__.remote(install), timeout=10) @@ -239,13 +335,13 @@ def cancelled_sample(dialect): break assert time.monotonic() < deadline, state time.sleep(.01) - assert state['quarantined'] + assert not state['weights_invalid'] engine_requests = ray.get([actor.__ray_call__.remote( lambda sampler: len(sampler.engine.engine.output_processor.request_states)) for actor in actors], timeout=10) assert engine_requests == [0] * len(actors) + ray.get(replica.__ray_call__.remote(restore), timeout=10) for dialect in ('tinker', 'twinkle'): - with pytest.raises(Exception): - _sample(clients, dialect, ids, 2) + assert _sample(clients, dialect, ids, 2)[0][0] _report(check='backend_self_cancellation', results=results, state=state, engine_requests=engine_requests, seconds=elapsed) finally: diff --git a/tests/server/sampler/sampling_test_helpers.py b/tests/server/sampler/sampling_test_helpers.py index 63564d363..17f52d1a9 100644 --- a/tests/server/sampler/sampling_test_helpers.py +++ b/tests/server/sampler/sampling_test_helpers.py @@ -86,9 +86,7 @@ def __init__(self, *, timeout=2, limit=128): self.replica_id = 'sampler-test' self.state = _State() self.sampler = _Sampler() - self.sampling_guard = SamplerConcurrency(limit) - self.tinker_setup_lock = asyncio.Lock() - self.tinker_template_set = True + self.sampling_guard = SamplerConcurrency() self._task_queue_config = TaskQueueConfig(execution_timeout=timeout, concurrent_task_limit=limit) self._rate_limiter = _Limiter() self._task_metrics = None diff --git a/tests/server/sampler/test_backend_concurrency.py b/tests/server/sampler/test_backend_concurrency.py new file mode 100644 index 000000000..3590edb27 --- /dev/null +++ b/tests/server/sampler/test_backend_concurrency.py @@ -0,0 +1,24 @@ +"""Only the server backend opts native TQ workers into ordinary concurrency.""" +import twinkle.infra as infra +from twinkle.infra._ray import RayHelper +from twinkle.utils.device_mesh import DeviceMesh +from twinkle.server.sampler.backends.vllm_async import ServerVLLMSampler +from twinkle.sampler.vllm_sampler.vllm_sampler import SAMPLER_MAX_CONCURRENCY +from twinkle_agentic.async_rl.vllm_sampler_tq import VLLMSamplerTQ + + +def test_server_concurrency_does_not_change_native_async_rl_defaults(monkeypatch): + captured = [] + def create(cls, *args, **kwargs): + captured.append((cls.__name__, kwargs['max_concurrency'])) + return [] + monkeypatch.setattr(infra, '_mode', 'ray') + monkeypatch.setattr(infra, '_device_group', []) + monkeypatch.setattr(RayHelper, 'create_workers', create) + monkeypatch.delenv('WORKER_NAME', raising=False) + monkeypatch.delenv('CLUSTER_NAME', raising=False) + monkeypatch.setenv('TWINKLE_TRUST_REMOTE_CODE', '1') + mesh = object.__new__(DeviceMesh) + for cls in (VLLMSamplerTQ, ServerVLLMSampler): + cls(model_id='test-model', context_manager=None, device_mesh=mesh, remote_group='test-sampler') + assert captured == [('VLLMSamplerTQ', None), ('ServerVLLMSampler', SAMPLER_MAX_CONCURRENCY)] diff --git a/tests/server/sampler/test_mock_sampler.py b/tests/server/sampler/test_mock_sampler.py index f9a2806df..5f432e556 100644 --- a/tests/server/sampler/test_mock_sampler.py +++ b/tests/server/sampler/test_mock_sampler.py @@ -121,7 +121,7 @@ def test_explicit_async_vllm_uses_non_blocking_sampler(monkeypatch) -> None: # in the `async-rl` extra, so skip when it is absent instead of failing a sampler test # for a data-path dependency it does not use. pytest.importorskip('transfer_queue') - from twinkle_agentic.async_rl import vllm_sampler_tq as module + from twinkle.server.sampler.backends import vllm_async as module captured = {} @@ -129,7 +129,7 @@ def construct(**kwargs): captured.update(kwargs) return 'vllm-tq' - monkeypatch.setattr(module, 'VLLMSamplerTQ', construct) + monkeypatch.setattr(module, 'ServerVLLMSampler', construct) sampler = _construct_sampler_backend( 'vllm_async', diff --git a/tests/server/sampler/test_resolve_sampler_weights.py b/tests/server/sampler/test_resolve_sampler_weights.py index ecaaa18f2..bf3b32fb3 100644 --- a/tests/server/sampler/test_resolve_sampler_weights.py +++ b/tests/server/sampler/test_resolve_sampler_weights.py @@ -17,6 +17,8 @@ class _FakeService: """Minimal stand-in exposing the two attributes the helper touches.""" def __init__(self): + from twinkle.server.sampler.concurrency import SamplerConcurrency + self.sampling_guard = SamplerConcurrency() self.sampler = _FakeSampler() self.full_weight_loads: list[str] = [] @@ -38,7 +40,8 @@ async def test_lora_dir_returns_adapter_path_and_loads_no_full_weights(tmp_path) @pytest.mark.asyncio async def test_full_checkpoint_loads_weights_and_returns_none(tmp_path): - # A directory without adapter_config.json is a full-parameter checkpoint. + (tmp_path / 'model.safetensors').touch() + # The fake backend supplies loading; the helper checks the checkpoint layout. svc = _FakeService() result = await resolve_sampler_weights(svc, str(tmp_path)) assert result is None diff --git a/tests/server/sampler/test_sampling_concurrency.py b/tests/server/sampler/test_sampling_concurrency.py index 23affb0db..52a08923f 100644 --- a/tests/server/sampler/test_sampling_concurrency.py +++ b/tests/server/sampler/test_sampling_concurrency.py @@ -5,15 +5,38 @@ from types import SimpleNamespace import pytest -from ray.exceptions import GetTimeoutError from starlette.requests import Request from twinkle.server.sampler.concurrency import SamplerConcurrency +from twinkle.exceptions import FullWeightLoadError from twinkle.server.task_queue.types import BackendBusyError from tests.server.sampler.sampling_test_helpers import _tokens, _until, make_service + +@pytest.mark.asyncio +async def test_standalone_checkpoint_load_waits_for_ordinary_sampling(make_service, tmp_path): + from twinkle.server.sampler.weights import resolve_sampler_weights + checkpoint = tmp_path / '2' + checkpoint.mkdir() + (checkpoint / 'model.safetensors').touch() + service = make_service() + service.sampler.release.clear() + ref = await service.submit('twinkle') + await _until(service.sampler.entered.is_set) + loading = asyncio.create_task(resolve_sampler_weights(service, str(checkpoint))) + try: + await asyncio.sleep(.02) + assert service.sampler.loads == [] + service.sampler.release.set() + assert _tokens(await service.result(ref), 'twinkle') == [7] + await loading + assert service.sampler.loads == [2] + finally: + service.sampler.release.set() + await asyncio.gather(loading, return_exceptions=True) + @pytest.mark.asyncio @pytest.mark.parametrize('dialects', [('tinker', 'tinker'), ('twinkle', 'twinkle'), ('tinker', 'twinkle')]) @pytest.mark.parametrize('count', [1, 8, 32]) @@ -35,6 +58,7 @@ async def test_real_routes_overlap_and_keep_request_results(make_service, dialec async def test_full_checkpoint_generation_is_one_exclusive_transaction(make_service, tmp_path, first, second): for number in ['1', '2']: (tmp_path / number).mkdir() + (tmp_path / number / 'model.safetensors').touch() service = make_service() service.sampler.release.clear() a = await service.submit(first, path=str(tmp_path / '1')) @@ -56,7 +80,7 @@ async def test_mutation_waits_for_both_dialects_generation(make_service): try: a, b = await asyncio.gather(service.submit('tinker'), service.submit('twinkle')) await _until(lambda: service.sampler.peak == 2) - mutation = asyncio.create_task(service.call_backend(service.sampler.set_template, 'template')) + mutation = asyncio.create_task(service.sampling_guard.run(lambda: service.call_backend(service.sampler.set_template, 'template'), exclusive=True)) await _until(lambda: service.sampling_guard._waiting_writers == 1) assert service.sampler.template_calls == 0 service.sampler.release.set() @@ -72,44 +96,35 @@ async def test_mutation_waits_for_both_dialects_generation(make_service): async def test_timeout_keeps_actual_load_exclusive_and_recovers_after_exit(make_service, tmp_path): for number in ['1', '2']: (tmp_path / number).mkdir() + (tmp_path / number / 'model.safetensors').touch() service = make_service(timeout=0.05) - load_entered = threading.Event() - load_release = threading.Event() + entered = threading.Event() + release = threading.Event() original = service.sampler.load_full_weights_from_path - def slow_load(path): - load_entered.set() - if not load_release.wait(3): - raise AssertionError('test did not release loading') + entered.set() + assert release.wait(3) original(path) - service.sampler.load_full_weights_from_path = slow_load try: ref = await service.submit('tinker', path=str(tmp_path / '1')) - await _until(load_entered.is_set) + await _until(entered.is_set) failed = await service.result(ref) assert failed['failure'].reason_code == 'execution_timeout' assert service.sampling_guard._writer - next_ref = await service.submit('twinkle', path=str(tmp_path / '2')) - rejected = await service.result(next_ref) - assert rejected['failure'].reason_code == 'backend_gate_unavailable' + submitting = asyncio.create_task(service.submit('twinkle', path=str(tmp_path / '2'))) + await _until(lambda: service.sampling_guard._waiting_writers == 1) assert service.sampler.loads == [] - load_release.set() - await _until(lambda: not service.sampling_guard._writer) - ref = await service.submit('twinkle', path=str(tmp_path / '2')) - assert _tokens(await service.result(ref), 'twinkle') == [2] + release.set() + next_ref = await submitting + assert _tokens(await service.result(next_ref), 'twinkle') == [2] + await _until(lambda: not service._compute_worker._detached_tasks) + assert service.state.records[ref] is failed + assert service.sampler.loads == [1, 2] finally: - load_release.set() + release.set() -@pytest.mark.asyncio -async def test_ray_timeout_refuses_unconfirmed_backend_release(make_service): - service = make_service() - service.sampler.sample = lambda *a, **k: (_ for _ in ()).throw(GetTimeoutError('actor still running')) - with pytest.raises(GetTimeoutError): - await service.call_backend(service.sampler.sample) - with pytest.raises(BackendBusyError): - await service.sampling_guard.run(lambda: asyncio.sleep(0)) @pytest.mark.asyncio @@ -153,12 +168,20 @@ async def test_shutdown_cancels_futures_without_freeing_running_backend_early(ma a, b = await asyncio.gather(service.submit('tinker'), service.submit('twinkle')) await _until(lambda: service.sampler.peak == 2) service.sampling_guard.close() - await service._compute_worker.stop() - assert service.state.records[a]['failure'].reason_code == 'orphaned_replica' - assert service.state.records[b]['failure'].reason_code == 'orphaned_replica' - assert service.sampling_guard._readers == 2 - with pytest.raises(BackendBusyError): - await service.call_backend(service.sampler.set_template, 'template') + stopping = asyncio.create_task(service._compute_worker.stop()) + try: + await _until(lambda: service.state.records[a]['status'] == 'failed') + assert service.state.records[a]['failure'].reason_code == 'orphaned_replica' + assert service.state.records[b]['failure'].reason_code == 'orphaned_replica' + assert service.sampling_guard._readers == 2 + assert not stopping.done() + with pytest.raises(BackendBusyError): + await service.sampling_guard.run(lambda: asyncio.sleep(0)) + finally: + service.sampler.release.set() + await stopping + assert service.sampling_guard._readers == 0 + assert service._compute_worker._detached_reserved == 0 @pytest.mark.asyncio @@ -186,7 +209,7 @@ async def put(rows, **kwargs): return DataRef(ref_id='ref', size=len(rows), fields=['tokens'], kind='rollout') service.data_plane.put = put - service.submit_background_and_peek = lambda factory, **kwargs: factory() + service.submit_concurrent_and_peek = lambda factory, **kwargs: factory() from twinkle.protocol.types import DataPlaneSampleRequest body = DataPlaneSampleRequest(inputs=[{'input_ids': [1]}], sampling_params={'max_tokens': 1}) request = Request({'type': 'http', 'headers': []}) @@ -194,7 +217,7 @@ async def put(rows, **kwargs): mutation = None try: await _until(admitted.is_set) - mutation = asyncio.create_task(service.call_backend(service.sampler.set_template, 'template')) + mutation = asyncio.create_task(service.sampling_guard.run(lambda: service.call_backend(service.sampler.set_template, 'template'), exclusive=True)) await _until(lambda: service.sampling_guard._waiting_writers == 1) assert service.sampler.template_calls == 0 completed.set() @@ -210,9 +233,10 @@ async def put(rows, **kwargs): @pytest.mark.parametrize('dialect', ['tinker', 'twinkle']) async def test_partial_checkpoint_failure_quarantines_changed_weights(make_service, tmp_path, dialect): service = make_service() + (tmp_path / 'model.safetensors').touch() def partial_load(path): service.sampler.weight = 99 - raise RuntimeError('load failed after a weight was changed') + raise FullWeightLoadError('load failed after a weight was changed') service.sampler.load_full_weights_from_path = partial_load ref = await service.submit(dialect, path=str(tmp_path)) assert (await service.result(ref))['status'] == 'failed' @@ -220,6 +244,8 @@ def partial_load(path): record = await service.result(ref) assert record['failure'].reason_code == 'backend_gate_unavailable' assert service.sampler.active == 0 + with pytest.raises(BackendBusyError, match='checkpoint write failed'): + await service.check_health() @pytest.mark.asyncio @@ -242,24 +268,25 @@ async def test_completed_generation_failure_does_not_quarantine_replica(make_ser service.sampler.get_generation_status = lambda _: {'status': 'failed', 'error': 'invalid input'} service.sampler.cancel_generation = lambda _: {'status': 'completed'} with pytest.raises(RuntimeError, match='invalid input'): - await service.sampling_guard.run(lambda: _await_generation(service, 'bad', 1)) + await service.sampling_guard.run(lambda: _await_generation(service, 'bad')) ref = await service.submit('twinkle') assert _tokens(await service.result(ref), 'twinkle') == [7] @pytest.mark.asyncio async def test_timeout_keeps_job_capacity_and_does_not_overwrite_failed_future(make_service): + from twinkle.server.exceptions import ResourceExhaustedError service = make_service(timeout=0.05, limit=1) service.sampler.release.clear() ref = await service.submit('tinker') await _until(service.sampler.entered.is_set) failed = await service.result(ref) assert failed['failure'].reason_code == 'execution_timeout' - assert len(service.sampling_guard._jobs) == 1 - with pytest.raises(BackendBusyError): - await service.sampling_guard.run(lambda: asyncio.sleep(0)) + assert service._compute_worker._detached_reserved == 1 + with pytest.raises(ResourceExhaustedError): + await service.submit('twinkle') service.sampler.release.set() - await _until(lambda: not service.sampling_guard._jobs) + await _until(lambda: not service._compute_worker._detached_tasks) assert service.state.records[ref] is failed assert (await service.result(await service.submit('twinkle')))['status'] == 'completed' @@ -289,28 +316,28 @@ async def write(): writer = asyncio.create_task(guard.run(write, exclusive=True)) await _until(lambda: guard._waiting_writers == 1) reader.cancel() - await asyncio.gather(reader, return_exceptions=True) - assert not wrote.is_set() + await asyncio.sleep(0) + assert not reader.done() and not wrote.is_set() release.set() + await asyncio.gather(reader, return_exceptions=True) await asyncio.wait_for(writer, 1) assert wrote.is_set() @pytest.mark.asyncio -async def test_abandoned_waiting_writer_never_starts_and_releases_capacity(): - guard = SamplerConcurrency(limit=2) +async def test_cancelled_writer_wait_releases_its_lock_bookkeeping(): + guard = SamplerConcurrency() release = asyncio.Event() wrote = asyncio.Event() - async def write(): - wrote.set() reader = asyncio.create_task(guard.run(release.wait)) await _until(lambda: guard._readers == 1) + async def write(): + wrote.set() writer = asyncio.create_task(guard.run(write, exclusive=True)) await _until(lambda: guard._waiting_writers == 1) writer.cancel() await asyncio.gather(writer, return_exceptions=True) - await _until(lambda: len(guard._jobs) == 1) - assert not wrote.is_set() and guard._waiting_writers == 0 + assert guard._waiting_writers == 0 and not wrote.is_set() release.set() await reader diff --git a/tests/server/sampler/test_sampling_recovery.py b/tests/server/sampler/test_sampling_recovery.py new file mode 100644 index 000000000..28ada50a6 --- /dev/null +++ b/tests/server/sampler/test_sampling_recovery.py @@ -0,0 +1,61 @@ +"""A rejected or timed-out request must not disable unrelated sampling.""" +import json + +import pytest +from ray.exceptions import GetTimeoutError +from starlette.requests import Request + +from twinkle.protocol.types import SampleRequest +from .sampling_test_helpers import _tokens, make_service + + +@pytest.mark.asyncio +@pytest.mark.parametrize('dialect', ['tinker', 'twinkle']) +async def test_invalid_checkpoint_rejects_request_without_disabling_sampling(make_service, tmp_path, dialect): + service = make_service() + ref = await service.submit(dialect, path=str(tmp_path / 'missing')) + failed = await service.result(ref) + assert failed['failure'].reason_code == 'request_rejected' + assert failed['failure'].attribution == 'user' + assert _tokens(await service.result(await service.submit(dialect)), dialect) == [7] + + +@pytest.mark.asyncio +async def test_ordinary_backend_timeout_does_not_disable_replica(make_service): + service = make_service() + sample = service.sampler.sample + service.sampler.sample = lambda *a, **k: (_ for _ in ()).throw(GetTimeoutError('one request timed out')) + with pytest.raises(GetTimeoutError): + await service.call_backend(service.sampler.sample) + service.sampler.sample = sample + assert _tokens(await service.result(await service.submit('twinkle')), 'twinkle') == [7] + + +@pytest.mark.asyncio +async def test_stream_route_resolves_adapter_uri(make_service, tmp_path): + service = make_service() + (tmp_path / 'adapter_config.json').write_text(json.dumps({'r': 1})) + service.sampler.reset_prefix_cache = lambda: None + service.sampler._actors = [] + request = Request({'type': 'http', 'headers': []}) + response = await service.endpoints['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/twinkle/sample_stream']( + request, SampleRequest(inputs=[{'input_ids': [7]}], adapter_uri=str(tmp_path)), service) + rows = [json.loads(row) async for row in response.body_iterator] + assert rows[0]['error_code'] == 503 + + +@pytest.mark.asyncio +async def test_sampling_timeout_keeps_capacity_without_blocking_other_readers(make_service): + import asyncio + from .sampling_test_helpers import _until + service = make_service(timeout=0.05, limit=2) + service.sampler.release.clear() + timed_out = await service.submit('tinker') + await _until(service.sampler.entered.is_set) + assert (await service.result(timed_out))['failure'].reason_code == 'execution_timeout' + submitting = asyncio.create_task(service.submit('twinkle')) + await _until(lambda: service.sampler.peak == 2) + service.sampler.release.set() + assert _tokens(await service.result(await submitting), 'twinkle') == [7] + await _until(lambda: service._compute_worker._detached_reserved == 0) + await service.check_health() diff --git a/tests/server/sampler/test_stream_guarantees.py b/tests/server/sampler/test_stream_guarantees.py index ff4907e4b..ef3aceafd 100644 --- a/tests/server/sampler/test_stream_guarantees.py +++ b/tests/server/sampler/test_stream_guarantees.py @@ -111,6 +111,6 @@ async def test_generation_poll_has_total_timeout_and_cancels(): service = _GenerationService() with pytest.raises(asyncio.TimeoutError): - await _await_generation(service, 'submission', timeout=0.05) + await asyncio.wait_for(_await_generation(service, 'submission'), timeout=0.05) assert service.cancelled is True diff --git a/tests/server/sampler/test_tinker_handlers.py b/tests/server/sampler/test_tinker_handlers.py index 2e9cb3133..2f4a0aad8 100644 --- a/tests/server/sampler/test_tinker_handlers.py +++ b/tests/server/sampler/test_tinker_handlers.py @@ -50,8 +50,6 @@ def __init__(self): self.model_id = 'mock-model' self.state = _DummyState() self.sampler = _DummySampler() - self.tinker_setup_lock = asyncio.Lock() - self.tinker_template_set = False self.sampling_guard = SamplerConcurrency() async def _on_request_start(self, request): @@ -61,12 +59,7 @@ async def schedule_concurrent_task(self, task, **kwargs): return await task() async def call_backend(self, fn, /, *args, **kwargs): - async def call(): - return fn(*args, **kwargs) - if self.sampling_guard.owns_job(): - return await call() - return await self.sampling_guard.run(call, exclusive=True) - + return fn(*args, **kwargs) @pytest.mark.asyncio @pytest.mark.parametrize('temperature,max_tokens', [(1.0, 2), (0.0, 2), (0.0, 0), (1.0, None)]) @@ -96,7 +89,7 @@ async def test_tinker_asample_allows_base_model_session_without_model_path(tempe @pytest.mark.asyncio -async def test_tinker_asample_sets_template_once_without_prefix_cache_reset(): +async def test_tinker_encoded_prompts_do_not_set_template_or_reset_prefix_cache(): """Both calls would wait for every in-flight generation, serializing requests.""" management = _DummyManagement() app = FastAPI() @@ -112,5 +105,5 @@ async def test_tinker_asample_sets_template_once_without_prefix_cache_reset(): await asyncio.gather(*(route.endpoint(request, body, management) for _ in range(3))) - assert management.sampler.calls == ['set_template'] + assert management.sampler.calls == [] assert len(management.sampler.adapter_paths) == 3 diff --git a/tests/server/sampler/test_twinkle_async_rows.py b/tests/server/sampler/test_twinkle_async_rows.py index b98ee2fd9..aa1877474 100644 --- a/tests/server/sampler/test_twinkle_async_rows.py +++ b/tests/server/sampler/test_twinkle_async_rows.py @@ -76,7 +76,7 @@ def __init__(self): async def _on_request_start(self, _request): return 'token' - async def submit_background_and_peek(self, coro_factory, *, model_id=None, task_type=None): + async def submit_concurrent_and_peek(self, coro_factory, *, model_id=None, task_type=None): self.scheduled.append({'model_id': model_id, 'task_type': task_type}) result = await coro_factory() from twinkle.protocol.types.lifecycle import TaskEnvelope diff --git a/tests/twinkle_agentic/test_vllm_sampler_tq_generation.py b/tests/twinkle_agentic/test_vllm_sampler_tq_generation.py index 4b4343282..9567138c4 100644 --- a/tests/twinkle_agentic/test_vllm_sampler_tq_generation.py +++ b/tests/twinkle_agentic/test_vllm_sampler_tq_generation.py @@ -13,7 +13,7 @@ from twinkle.server.sampler.twinkle_handlers import _await_generation from twinkle_agentic.async_rl import LoraContext from twinkle_agentic.async_rl.types import PartitionAdmission, PromptGroup, RolloutPolicy -from twinkle_agentic.async_rl.vllm_sampler_tq import (_GenerationSubmission, VLLMSamplerTQ, _dispatch_generation, _GeneratedSample, +from twinkle_agentic.async_rl.vllm_sampler_tq import (VLLMSamplerTQ, _dispatch_generation, _GeneratedSample, _PromptGroupRolloutStats) @@ -174,7 +174,7 @@ def test_generation_failure_is_isolated_and_consumed() -> None: sampler = _bare_sampler() failed = Future() failed.set_exception(ValueError('bad prompt')) - sampler._generation_submissions['failed'] = _GenerationSubmission(future=failed) + sampler._generation_submissions['failed'] = failed state = sampler.get_generation_status('failed') assert state['status'] == 'failed' @@ -185,43 +185,29 @@ def test_generation_failure_is_isolated_and_consumed() -> None: assert 'failed' not in sampler._generation_submissions -def test_cancellation_waits_for_actual_coroutine_cleanup() -> None: - async def run(): - sampler = _bare_sampler() - sampler._async_loop = asyncio.get_running_loop() - entered = asyncio.Event() - release = asyncio.Event() - async def generate(*args, **kwargs): - try: - entered.set() - await asyncio.Future() - finally: - await release.wait() - sampler._generate_inputs = generate - sampler.submit_generation('pending', [], SamplingParams()) - await entered.wait() - assert sampler.cancel_generation('pending')['status'] == 'cancellation_requested' - await asyncio.sleep(.01) - assert sampler.get_generation_status('pending')['status'] == 'running' - assert sampler.cancel_generation('pending')['status'] == 'cancellation_requested' - release.set() - for _ in range(100): - if sampler.get_generation_status('pending')['status'] == 'cancelled': - break - await asyncio.sleep(.001) - assert sampler.cancel_generation('pending')['status'] == 'completed' - assert not sampler._generation_submissions - asyncio.run(run()) +def test_generation_can_be_cancelled_without_waiting() -> None: + sampler = _bare_sampler() + pending = Future() + sampler._generation_submissions['pending'] = pending + state = sampler.cancel_generation('pending') -def test_all_completed_generations_are_forgotten_on_shutdown() -> None: + assert state == {'submission_id': 'pending', 'status': 'cancelled'} + assert pending.cancelled() + assert 'pending' not in sampler._generation_submissions + + +def test_all_generations_are_cancelled_on_shutdown() -> None: sampler = _bare_sampler() - first, second = Future(), Future() - first.set_result([]) - second.set_exception(ValueError('bad input')) - sampler._generation_submissions.update( - first=_GenerationSubmission(future=first), second=_GenerationSubmission(future=second)) - assert sampler.cancel_all_generations() == {'submissions': 2, 'cancelled': 0} + first = Future() + second = Future() + sampler._generation_submissions.update(first=first, second=second) + + state = sampler.cancel_all_generations() + + assert state == {'submissions': 2, 'cancelled': 2} + assert first.cancelled() + assert second.cancelled() assert sampler._generation_submissions == {} @@ -270,8 +256,8 @@ def cancel_generation(self, submission_id): async def run(): sampler.submit_generation('first') sampler.submit_generation('second') - first = asyncio.create_task(_await_generation(service, 'first', timeout=5)) - second = asyncio.create_task(_await_generation(service, 'second', timeout=5)) + first = asyncio.create_task(_await_generation(service, 'first')) + second = asyncio.create_task(_await_generation(service, 'second')) while len(sampler.submission_order) < 2: await asyncio.sleep(0) assert not first.done() @@ -312,7 +298,7 @@ def cancel_generation(self, _submission_id): sampler = Sampler() sampler.submit_generation('submission') service = _DirectGenerationService(sampler) - result = asyncio.run(_await_generation(service, 'submission', timeout=5)) + result = asyncio.run(_await_generation(service, 'submission')) assert result == ['completed'] assert sampler.status_calls == 2 From addab02051e8cf55faa849e2a86145cfb2892e7e Mon Sep 17 00:00:00 2001 From: Yunnglin Date: Wed, 14 Oct 2026 15:54:47 +0800 Subject: [PATCH 4/5] test(sampler): consolidate concurrency coverage and narrow PR scope --- .../Usage Guide/Server and Client/Server.md | 6 - .../\346\234\215\345\212\241\347\253\257.md" | 6 - .../sampler/test_full_checkpoint_failures.py | 52 ----- tests/sampler/test_token_only_sampling.py | 39 ---- tests/sampler/test_vllm_sampling.py | 123 ++++++++++ tests/sampler/test_vllm_sampling_lifecycle.py | 41 ---- .../server/integration/sampler_e2e_helpers.py | 183 --------------- .../test_sampler_concurrency_e2e.py | 182 ++++++++++++++- tests/server/sampler/sampling_test_helpers.py | 151 ------------ .../sampler/test_backend_concurrency.py | 24 -- tests/server/sampler/test_mock_sampler.py | 25 ++ .../sampler/test_resolve_sampler_weights.py | 2 +- .../sampler/test_sampling_concurrency.py | 221 ++++++++++++++++-- .../server/sampler/test_sampling_recovery.py | 61 ----- .../server/sampler/test_stream_guarantees.py | 5 +- 15 files changed, 524 insertions(+), 597 deletions(-) delete mode 100644 tests/sampler/test_full_checkpoint_failures.py delete mode 100644 tests/sampler/test_token_only_sampling.py create mode 100644 tests/sampler/test_vllm_sampling.py delete mode 100644 tests/sampler/test_vllm_sampling_lifecycle.py delete mode 100644 tests/server/integration/sampler_e2e_helpers.py delete mode 100644 tests/server/sampler/sampling_test_helpers.py delete mode 100644 tests/server/sampler/test_backend_concurrency.py delete mode 100644 tests/server/sampler/test_sampling_recovery.py diff --git a/docs/source_en/Usage Guide/Server and Client/Server.md b/docs/source_en/Usage Guide/Server and Client/Server.md index 114c7041f..dd0064aba 100644 --- a/docs/source_en/Usage Guide/Server and Client/Server.md +++ b/docs/source_en/Usage Guide/Server and Client/Server.md @@ -479,9 +479,3 @@ raises. During that window newly arriving tasks fail fast with a `server`/503 er Each persisted future stores its immutable `absolute_deadline` when it is created. Cleanup therefore reaches the same decision regardless of which deployment process holds the cleanup lease. - -### Sampler concurrency and weight changes - -- Ordinary Tinker and Twinkle `sample()` requests support concurrency with unchanged response schemas and future polling; training compute remains serial. `queue_config.concurrent_task_limit` defaults to `128` and bounds admitted, unfinished tasks per replica; exhaustion returns `server`/503. Actual execution concurrency also depends on the backend thread pool, actor, and vLLM settings. -- Full checkpoint loading and generation exclusively own the weights and wait for existing ordinary and DataPlane generation to finish. LoRA uses adapter checkpoints without replacing base weights on each request. -- A sampling timeout fails only its future; capacity stays reserved until the backend finishes, while other sampling can continue. Invalid checkpoints are rejected before writes. A failed weight write fails the health check so Ray Serve replaces the replica. diff --git "a/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/\346\234\215\345\212\241\347\253\257\345\222\214\345\256\242\346\210\267\347\253\257/\346\234\215\345\212\241\347\253\257.md" "b/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/\346\234\215\345\212\241\347\253\257\345\222\214\345\256\242\346\210\267\347\253\257/\346\234\215\345\212\241\347\253\257.md" index de96291a4..81fcacfe5 100644 --- "a/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/\346\234\215\345\212\241\347\253\257\345\222\214\345\256\242\346\210\267\347\253\257/\346\234\215\345\212\241\347\253\257.md" +++ "b/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/\346\234\215\345\212\241\347\253\257\345\222\214\345\256\242\346\210\267\347\253\257/\346\234\215\345\212\241\347\253\257.md" @@ -466,9 +466,3 @@ twinkle-server check-config -c server_config.yaml 任务记录进入终态后，per-replica 准入闸门额外保持关闭的最长时长为 `max(0, Collect_Width × R − T)`：此时记录已是终态，但泄漏的 executor 线程可能仍持有闸门，直到其 `ray.get` 返回或抛出。在该窗口内新到达的任务会以 `server`/503 错误快速失败。 每条持久化 future 在创建时写入不可变的 `absolute_deadline`，因此无论哪个 deployment 进程持有 cleanup lease，清理结果都由任务自身契约决定。 - -### Sampler 并发与权重切换 - -- Tinker 和 Twinkle 的普通 `sample()` 支持并发，响应格式与 future 轮询接口不变；训练计算仍然串行。`queue_config.concurrent_task_limit` 默认为 `128`，限制单个副本已接收但尚未结束的任务数，满载时返回 `server`/503。这是准入上限，实际执行并发还受后端线程池、actor 和 vLLM 配置限制。 -- 全量 checkpoint 的加载和生成独占权重访问，等待已有普通生成和 DataPlane 生成结束。LoRA 使用 adapter checkpoint，无须每次替换基座权重。 -- 采样超时只让该 future 失败；后台任务实际结束后才释放容量，其他采样可以继续。错误 checkpoint 在写入前拒绝；真正的权重写入失败会使健康检查失败，由 Ray Serve 替换副本。 diff --git a/tests/sampler/test_full_checkpoint_failures.py b/tests/sampler/test_full_checkpoint_failures.py deleted file mode 100644 index 755e4dd6c..000000000 --- a/tests/sampler/test_full_checkpoint_failures.py +++ /dev/null @@ -1,52 +0,0 @@ -"""Validate checkpoints before writing; distinguish partial engine writes.""" -import asyncio -import json -from types import SimpleNamespace - -import pytest -import torch -from safetensors.torch import save_file - -from twinkle.exceptions import FullWeightLoadError -from twinkle.sampler.vllm_sampler.vllm_sampler import vLLMSampler - - -@pytest.mark.parametrize('kind', ['missing', 'empty', 'bad_index', 'missing_shard', 'corrupt_shard']) -def test_invalid_checkpoint_does_not_write_weights(tmp_path, kind): - path = tmp_path / 'checkpoint' - if kind != 'missing': - path.mkdir() - if kind == 'bad_index': - (path / 'model.safetensors.index.json').write_text('{') - if kind == 'missing_shard': - (path / 'model.safetensors.index.json').write_text(json.dumps({'weight_map': {'weight': 'absent.safetensors'}})) - if kind == 'corrupt_shard': - (path / 'model.safetensors').write_text('not a checkpoint') - sampler = object.__new__(vLLMSampler) - # This existing local directory avoids a hub request for the missing case. - from unittest.mock import patch - writes = [] - sampler.model_id = str(path) - sampler.engine = SimpleNamespace(update_weights=lambda *a, **k: writes.append(True)) - with patch('twinkle.sampler.vllm_sampler.vllm_sampler.HubOperation.download_model', return_value=str(path)): - with pytest.raises(Exception) as error: - sampler.load_full_weights_from_path(str(path)) - assert not isinstance(error.value, FullWeightLoadError) - assert writes == [] - assert not hasattr(sampler, '_loaded_full_weights_path') - - -def test_partial_write_is_reported_as_weight_failure(tmp_path): - save_file({'weight': torch.ones(1)}, str(tmp_path / 'model.safetensors')) - sampler = object.__new__(vLLMSampler) - writes = [] - async def update(weights, **kwargs): - writes.extend(list(weights)) - raise RuntimeError('engine failed after writing') - sampler.engine = SimpleNamespace(update_weights=update) - sampler._run_in_loop = asyncio.run - with pytest.raises(FullWeightLoadError) as error: - sampler.load_full_weights_from_path(str(tmp_path)) - assert isinstance(error.value.__cause__, RuntimeError) - assert len(writes) == 1 - assert not hasattr(sampler, '_loaded_full_weights_path') diff --git a/tests/sampler/test_token_only_sampling.py b/tests/sampler/test_token_only_sampling.py deleted file mode 100644 index f9aeea903..000000000 --- a/tests/sampler/test_token_only_sampling.py +++ /dev/null @@ -1,39 +0,0 @@ -"""Tinker consumes encoded tokens, independently of native template settings.""" -from types import SimpleNamespace -from unittest.mock import AsyncMock - -import pytest - -from twinkle.data_format import SampledSequence, SampleResponse, SamplingParams -from twinkle.sampler.vllm_sampler import vLLMSampler - - -class _OtherClientTemplate: - def get_vllm_input_ids(self, *args): - raise AssertionError('Token-only sampling rewrote encoded input IDs') - - def concat_input_feature(self, *args): - raise AssertionError('Token-only sampling built unused training features') - - def decode(self, *args): - raise AssertionError('Token-only sampling decoded unused text') - - -@pytest.mark.asyncio -@pytest.mark.parametrize('template', [None, _OtherClientTemplate()]) -@pytest.mark.parametrize('logprobs_only', [False, True]) -async def test_token_only_sampling_needs_no_template(template, logprobs_only): - sampler = object.__new__(vLLMSampler) - sampler.template = template - sampler.engine = SimpleNamespace(sample=AsyncMock(return_value=SampleResponse( - prompt_token_ids=[7], - sequences=[SampledSequence(tokens=[8], logprobs=[[(8, -.1)]], stop_reason='length')], - prompt_logprobs=[-.2], - ))) - result = await sampler._sample_single( - {'input_ids': [7]}, SamplingParams(max_tokens=1), token_only=True, logprobs_only=logprobs_only) - assert sampler.engine.sample.call_args.kwargs['prompt'] == [7] - assert result.prompt_token_ids == [7] and result.prompt_logprobs == [-.2] - assert result.sequences[0].tokens == ([] if logprobs_only else [8]) - assert result.sequences[0].logprobs == (None if logprobs_only else [[(8, -.1)]]) - assert result.sequences[0].decoded is None and result.sequences[0].new_input_feature is None diff --git a/tests/sampler/test_vllm_sampling.py b/tests/sampler/test_vllm_sampling.py new file mode 100644 index 000000000..928326527 --- /dev/null +++ b/tests/sampler/test_vllm_sampling.py @@ -0,0 +1,123 @@ +"""Token-only sampling, batch cleanup and checkpoint write boundaries.""" +from __future__ import annotations + +import asyncio +import json +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest +import torch +from safetensors.torch import save_file + +from twinkle.data_format import SampledSequence, SampleResponse, SamplingParams +from twinkle.exceptions import FullWeightLoadError +from twinkle.sampler.vllm_sampler.vllm_sampler import vLLMSampler + + +@pytest.mark.parametrize('kind', ['missing', 'empty', 'bad_index', 'missing_shard', 'corrupt_shard']) +def test_invalid_checkpoint_does_not_write_weights(tmp_path, kind): + path = tmp_path / 'checkpoint' + if kind != 'missing': + path.mkdir() + if kind == 'bad_index': + (path / 'model.safetensors.index.json').write_text('{') + if kind == 'missing_shard': + (path / 'model.safetensors.index.json').write_text(json.dumps({'weight_map': {'weight': 'absent.safetensors'}})) + if kind == 'corrupt_shard': + (path / 'model.safetensors').write_text('not a checkpoint') + sampler = object.__new__(vLLMSampler) + # This existing local directory avoids a hub request for the missing case. + from unittest.mock import patch + writes = [] + sampler.model_id = str(path) + sampler.engine = SimpleNamespace(update_weights=lambda *a, **k: writes.append(True)) + with patch('twinkle.sampler.vllm_sampler.vllm_sampler.HubOperation.download_model', return_value=str(path)): + with pytest.raises(Exception) as error: + sampler.load_full_weights_from_path(str(path)) + assert not isinstance(error.value, FullWeightLoadError) + assert writes == [] + assert not hasattr(sampler, '_loaded_full_weights_path') + + +def test_partial_write_is_reported_as_weight_failure(tmp_path): + save_file({'weight': torch.ones(1)}, str(tmp_path / 'model.safetensors')) + sampler = object.__new__(vLLMSampler) + writes = [] + async def update(weights, **kwargs): + writes.extend(list(weights)) + raise RuntimeError('engine failed after writing') + sampler.engine = SimpleNamespace(update_weights=update) + sampler._run_in_loop = asyncio.run + with pytest.raises(FullWeightLoadError) as error: + sampler.load_full_weights_from_path(str(tmp_path)) + assert isinstance(error.value.__cause__, RuntimeError) + assert len(writes) == 1 + assert not hasattr(sampler, '_loaded_full_weights_path') + + +class _OtherClientTemplate: + def get_vllm_input_ids(self, *args): + raise AssertionError('Token-only sampling rewrote encoded input IDs') + + def concat_input_feature(self, *args): + raise AssertionError('Token-only sampling built unused training features') + + def decode(self, *args): + raise AssertionError('Token-only sampling decoded unused text') + + +@pytest.mark.asyncio +@pytest.mark.parametrize('template', [None, _OtherClientTemplate()]) +@pytest.mark.parametrize('logprobs_only', [False, True]) +async def test_token_only_sampling_needs_no_template(template, logprobs_only): + sampler = object.__new__(vLLMSampler) + sampler.template = template + sampler.engine = SimpleNamespace(sample=AsyncMock(return_value=SampleResponse( + prompt_token_ids=[7], + sequences=[SampledSequence(tokens=[8], logprobs=[[(8, -.1)]], stop_reason='length')], + prompt_logprobs=[-.2], + ))) + result = await sampler._sample_single( + {'input_ids': [7]}, SamplingParams(max_tokens=1), token_only=True, logprobs_only=logprobs_only) + assert sampler.engine.sample.call_args.kwargs['prompt'] == [7] + assert result.prompt_token_ids == [7] and result.prompt_logprobs == [-.2] + assert result.sequences[0].tokens == ([] if logprobs_only else [8]) + assert result.sequences[0].logprobs == (None if logprobs_only else [[(8, -.1)]]) + assert result.sequences[0].decoded is None and result.sequences[0].new_input_feature is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize('cancel_batch', [False, True]) +async def test_batch_waits_for_cancelled_sibling_cleanup(cancel_batch): + entered = asyncio.Event() + release = asyncio.Event() + cleanup = asyncio.Event() + finished = asyncio.Event() + async def failed(): + await entered.wait() + if cancel_batch: + await asyncio.Future() + raise ValueError('bad prompt') + async def sibling(): + try: + entered.set() + await asyncio.Future() + finally: + cleanup.set() + await release.wait() + finished.set() + generation = asyncio.create_task(vLLMSampler._gather_samples([failed(), sibling()])) + try: + await entered.wait() + if cancel_batch: + generation.cancel() + await asyncio.wait_for(cleanup.wait(), 1) + assert not generation.done() + release.set() + with pytest.raises(asyncio.CancelledError if cancel_batch else ValueError): + await generation + assert finished.is_set() + finally: + release.set() + await asyncio.gather(generation, return_exceptions=True) diff --git a/tests/sampler/test_vllm_sampling_lifecycle.py b/tests/sampler/test_vllm_sampling_lifecycle.py deleted file mode 100644 index 08fe609fa..000000000 --- a/tests/sampler/test_vllm_sampling_lifecycle.py +++ /dev/null @@ -1,41 +0,0 @@ -"""Batch failure cancels siblings and waits for their actual cleanup.""" -import asyncio - -import pytest -from twinkle.sampler.vllm_sampler import vLLMSampler - - -@pytest.mark.asyncio -@pytest.mark.parametrize('cancel_batch', [False, True]) -async def test_batch_waits_for_cancelled_sibling_cleanup(cancel_batch): - entered = asyncio.Event() - release = asyncio.Event() - cleanup = asyncio.Event() - finished = asyncio.Event() - async def failed(): - await entered.wait() - if cancel_batch: - await asyncio.Future() - raise ValueError('bad prompt') - async def sibling(): - try: - entered.set() - await asyncio.Future() - finally: - cleanup.set() - await release.wait() - finished.set() - generation = asyncio.create_task(vLLMSampler._gather_samples([failed(), sibling()])) - try: - await entered.wait() - if cancel_batch: - generation.cancel() - await asyncio.wait_for(cleanup.wait(), 1) - assert not generation.done() - release.set() - with pytest.raises(asyncio.CancelledError if cancel_batch else ValueError): - await generation - assert finished.is_set() - finally: - release.set() - await asyncio.gather(generation, return_exceptions=True) diff --git a/tests/server/integration/sampler_e2e_helpers.py b/tests/server/integration/sampler_e2e_helpers.py deleted file mode 100644 index 4dbc31ce3..000000000 --- a/tests/server/integration/sampler_e2e_helpers.py +++ /dev/null @@ -1,183 +0,0 @@ -"""Shared clients and model fixtures for real sampler E2E tests.""" -from __future__ import annotations - -import concurrent.futures -import json -import math -import os -import threading -import time -from pathlib import Path - -import pytest - -@pytest.fixture(scope='session', autouse=True) -def _ray_runtime(): - """Use the external server; do not start the parent conftest's local Ray.""" - yield - - -@pytest.fixture(autouse=True) -def _reset_canonical_state_actor(): - """Keep external server state intact across E2E requests.""" - yield - - -def _report(**values): - print('SAMPLER_CONCURRENCY ' + json.dumps(values, sort_keys=True), flush=True) - - -@pytest.fixture(scope='module') -def clients(): - from transformers import AutoTokenizer - from twinkle import init_tinker_client, init_twinkle_client - init_tinker_client() - from tinker import ServiceClient - model_id = os.environ['TWINKLE_TEST_MODEL_ID'] - model_path = os.environ['TWINKLE_TEST_MODEL_PATH'] - url = os.environ['TWINKLE_SERVER_URL'] - tokenizer = AutoTokenizer.from_pretrained(model_path, local_files_only=True) - service = ServiceClient(base_url=url, api_key='EMPTY_TOKEN') - native = init_twinkle_client(base_url=url, api_key='EMPTY_TOKEN').sampler(model_id) - return { - 'tokenizer': tokenizer, 'service': service, 'model_id': model_id, 'model_path': model_path, - 'url': url, 'tinker': service.create_sampling_client(base_model=model_id), 'twinkle': native, - 'single_throughput': {}, - } - - -def _prompt(clients, text): - return clients['tokenizer'].apply_chat_template( - [{'role': 'user', 'content': text}], tokenize=True, add_generation_prompt=True, - enable_thinking=False, return_dict=False) - - -def _sample(clients, dialect, ids, max_tokens=128, *, adapter=None, num_samples=1, temperature=0): - if dialect == 'tinker': - from tinker import types - result = (adapter or clients['tinker']).sample( - prompt=types.ModelInput.from_ints(ids), num_samples=num_samples, - sampling_params=types.SamplingParams(max_tokens=max_tokens, temperature=temperature, top_p=1), - ).result(timeout=300) - sequences = [(list(seq.tokens), list(seq.logprobs)) for seq in result.sequences] - else: - result = clients['twinkle'].sample( - inputs=[{'input_ids': ids}], num_samples=num_samples, adapter_uri=adapter, - sampling_params={'max_tokens': max_tokens, 'temperature': temperature, 'top_p': 1, 'logprobs': 1}) - assert result[0].prompt_token_ids == ids - sequences = [(list(seq.tokens), [position[0][1] for position in seq.logprobs]) - for seq in result[0].sequences] - assert len(sequences) == num_samples - for tokens, logprobs in sequences: - assert tokens and len(tokens) == len(logprobs) - assert all(math.isfinite(value) for value in logprobs) - return sequences - - -def _parallel(jobs): - barrier = threading.Barrier(len(jobs)) - def run(job): - barrier.wait(timeout=30) - return job() - started = time.perf_counter() - with concurrent.futures.ThreadPoolExecutor(max_workers=len(jobs)) as pool: - results = list(pool.map(run, jobs)) - return results, time.perf_counter() - started - - -@pytest.fixture(scope='module') -def adapters(clients): - import torch - from peft import LoraConfig - from safetensors import safe_open - from safetensors.torch import save_file - projections = {} - for shard in Path(clients['model_path']).glob('*.safetensors'): - with safe_open(str(shard), framework='pt', device='cpu') as weights: - for name in weights.keys(): - if name.endswith('.self_attn.q_proj.weight'): - projections[name[:-len('.weight')]] = weights.get_slice(name).get_shape() - if not projections: - pytest.skip('The cached model has no q_proj layers for this LoRA fixture') - root = Path(os.environ['TWINKLE_TEST_ARTIFACT_DIR']) / 'lora-isolation' - paths = [] - rank = 8 - for index in range(2): - path = root / str(index) - path.mkdir(parents=True, exist_ok=True) - LoraConfig(r=rank, lora_alpha=rank, target_modules=['q_proj'], bias='none', - task_type='CAUSAL_LM', inference_mode=True, - base_model_name_or_path=clients['model_id']).save_pretrained(path) - generator = torch.Generator().manual_seed(100 + index) - weights = {} - for name, (output_size, input_size) in projections.items(): - prefix = 'base_model.model.' + name - weights[prefix + '.lora_A.weight'] = torch.randn(rank, input_size, generator=generator) * 0.05 - weights[prefix + '.lora_B.weight'] = torch.randn(output_size, rank, generator=generator) * 0.05 - save_file(weights, str(path / 'adapter_model.safetensors')) - paths.append(str(path)) - return paths - - -@pytest.fixture(scope='module') -def full_checkpoints(clients): - """Two complete HF checkpoints with different, observable model weights.""" - import torch - from safetensors.torch import load_file, save_file - source = Path(clients['model_path']) - target = Path(os.environ['TWINKLE_TEST_ARTIFACT_DIR']) / 'full-checkpoint-B' - target.mkdir(parents=True, exist_ok=True) - for file in source.iterdir(): - if file.is_file() and not (target / file.name).exists(): - (target / file.name).symlink_to(file) - index = json.loads((source / 'model.safetensors.index.json').read_text()) - shard = index['weight_map']['model.embed_tokens.weight'] - output = target / shard - if output.is_symlink(): - weights = load_file(str(source / shard), device='cpu') - weights['model.embed_tokens.weight'] = torch.roll(weights['model.embed_tokens.weight'], 257, 0).contiguous() - output.unlink() - save_file(weights, str(output), metadata={'format': 'pt'}) - return [str(source), str(target)] - - -def _checkpoint_handles(clients, dialect, paths): - if dialect == 'tinker': - return [clients['service'].create_sampling_client(base_model=clients['model_id'], model_path=path) - for path in paths] - return paths - - -def _assert_adapter_isolation(clients, adapters, dialect): - handles = ([None] + [clients['service'].create_sampling_client(base_model=clients['model_id'], model_path=path) - for path in adapters] if dialect == 'tinker' else [None, *adapters]) - ids = _prompt(clients, 'Choose one word: apple, banana, orange, purple.') - solo = [_sample(clients, dialect, ids, 8, adapter=handle)[0] for handle in handles] - # Default vLLM kernels do not guarantee batch-invariant continuations or - # logprobs. Homogeneous PPU batches also diverge after this stable prefix. - # Require distinct adapter fingerprints, then check routing against them; - # _sample still validates every generated token's finite, aligned logprob. - fingerprints = [tuple(tokens[:2]) for tokens, _logprobs in solo] - assert all(len(fingerprint) == 2 for fingerprint in fingerprints) - assert len(set(fingerprints)) == 3, 'the adapters must change the result, rather than silently falling back to base' - results, elapsed = _parallel([ - lambda i=i: _sample(clients, dialect, ids, 8, adapter=handles[i % 3]) for i in range(24)]) - for i, result in enumerate(results): - tokens, _logprobs = result[0] - assert tuple(tokens[:2]) == fingerprints[i % 3] - _report(check='lora_isolation', dialect=dialect, requests=24, variants=3, seconds=elapsed, - fingerprints=fingerprints) - - -def _real_sampler_actors(): - import ray - from ray.util.state import list_actors - ray.init(address=os.environ['RAY_ADDRESS'], namespace='default', ignore_reinit_error=True, - log_to_driver=False) - actors = list_actors(address=os.environ['TWINKLE_TEST_RAY_DASHBOARD_URL'], - filters=[('class_name', '=', os.environ.get('TWINKLE_TEST_SAMPLER_CLASS', 'VLLMSamplerTQ')), - ('state', '=', 'ALIVE')], - detail=True) - assert len(actors) == int(os.environ.get('TWINKLE_TEST_DP', '1')) - assert all(actor.state == 'ALIVE' for actor in actors) - return [ray.get_actor(actor.name, namespace='default') for actor in actors] diff --git a/tests/server/integration/test_sampler_concurrency_e2e.py b/tests/server/integration/test_sampler_concurrency_e2e.py index 6b051c161..697b1eab8 100644 --- a/tests/server/integration/test_sampler_concurrency_e2e.py +++ b/tests/server/integration/test_sampler_concurrency_e2e.py @@ -1,17 +1,189 @@ -"""Real Tinker/Twinkle concurrency, weight isolation and lifecycle checks.""" +"""Real Tinker and Twinkle concurrency, weight isolation and recovery.""" from __future__ import annotations +import concurrent.futures +import json import math import os import threading import time +from pathlib import Path import pytest -from tests.server.integration.sampler_e2e_helpers import ( - _assert_adapter_isolation, _checkpoint_handles, _parallel, _prompt, _ray_runtime, - _report, _reset_canonical_state_actor, _sample, _real_sampler_actors, adapters, clients, full_checkpoints, -) + +@pytest.fixture(scope='session', autouse=True) +def _ray_runtime(): + """Use the external server; do not start the parent conftest's local Ray.""" + yield + + +@pytest.fixture(autouse=True) +def _reset_canonical_state_actor(): + """Keep external server state intact across E2E requests.""" + yield + + +def _report(**values): + print('SAMPLER_CONCURRENCY ' + json.dumps(values, sort_keys=True), flush=True) + + +@pytest.fixture(scope='module') +def clients(): + from transformers import AutoTokenizer + + from twinkle import init_tinker_client, init_twinkle_client + init_tinker_client() + from tinker import ServiceClient + model_id = os.environ['TWINKLE_TEST_MODEL_ID'] + model_path = os.environ['TWINKLE_TEST_MODEL_PATH'] + url = os.environ['TWINKLE_SERVER_URL'] + tokenizer = AutoTokenizer.from_pretrained(model_path, local_files_only=True) + service = ServiceClient(base_url=url, api_key='EMPTY_TOKEN') + native = init_twinkle_client(base_url=url, api_key='EMPTY_TOKEN').sampler(model_id) + return { + 'tokenizer': tokenizer, 'service': service, 'model_id': model_id, 'model_path': model_path, + 'url': url, 'tinker': service.create_sampling_client(base_model=model_id), 'twinkle': native, + 'single_throughput': {}, + } + + +def _prompt(clients, text): + return clients['tokenizer'].apply_chat_template( + [{'role': 'user', 'content': text}], tokenize=True, add_generation_prompt=True, + enable_thinking=False, return_dict=False) + + +def _sample(clients, dialect, ids, max_tokens=128, *, adapter=None, num_samples=1, temperature=0): + if dialect == 'tinker': + from tinker import types + result = (adapter or clients['tinker']).sample( + prompt=types.ModelInput.from_ints(ids), num_samples=num_samples, + sampling_params=types.SamplingParams(max_tokens=max_tokens, temperature=temperature, top_p=1), + ).result(timeout=300) + sequences = [(list(seq.tokens), list(seq.logprobs)) for seq in result.sequences] + else: + result = clients['twinkle'].sample( + inputs=[{'input_ids': ids}], num_samples=num_samples, adapter_uri=adapter, + sampling_params={'max_tokens': max_tokens, 'temperature': temperature, 'top_p': 1, 'logprobs': 1}) + assert result[0].prompt_token_ids == ids + sequences = [(list(seq.tokens), [position[0][1] for position in seq.logprobs]) + for seq in result[0].sequences] + assert len(sequences) == num_samples + for tokens, logprobs in sequences: + assert tokens and len(tokens) == len(logprobs) + assert all(math.isfinite(value) for value in logprobs) + return sequences + + +def _parallel(jobs): + barrier = threading.Barrier(len(jobs)) + def run(job): + barrier.wait(timeout=30) + return job() + started = time.perf_counter() + with concurrent.futures.ThreadPoolExecutor(max_workers=len(jobs)) as pool: + results = list(pool.map(run, jobs)) + return results, time.perf_counter() - started + + +@pytest.fixture(scope='module') +def adapters(clients): + import torch + from peft import LoraConfig + from safetensors import safe_open + from safetensors.torch import save_file + projections = {} + for shard in Path(clients['model_path']).glob('*.safetensors'): + with safe_open(str(shard), framework='pt', device='cpu') as weights: + for name in weights.keys(): + if name.endswith('.self_attn.q_proj.weight'): + projections[name[:-len('.weight')]] = weights.get_slice(name).get_shape() + if not projections: + pytest.skip('The cached model has no q_proj layers for this LoRA fixture') + root = Path(os.environ['TWINKLE_TEST_ARTIFACT_DIR']) / 'lora-isolation' + paths = [] + rank = 8 + for index in range(2): + path = root / str(index) + path.mkdir(parents=True, exist_ok=True) + LoraConfig(r=rank, lora_alpha=rank, target_modules=['q_proj'], bias='none', + task_type='CAUSAL_LM', inference_mode=True, + base_model_name_or_path=clients['model_id']).save_pretrained(path) + generator = torch.Generator().manual_seed(100 + index) + weights = {} + for name, (output_size, input_size) in projections.items(): + prefix = 'base_model.model.' + name + weights[prefix + '.lora_A.weight'] = torch.randn(rank, input_size, generator=generator) * 0.05 + weights[prefix + '.lora_B.weight'] = torch.randn(output_size, rank, generator=generator) * 0.05 + save_file(weights, str(path / 'adapter_model.safetensors')) + paths.append(str(path)) + return paths + + +@pytest.fixture(scope='module') +def full_checkpoints(clients): + """Two complete HF checkpoints with different, observable model weights.""" + import torch + from safetensors.torch import load_file, save_file + source = Path(clients['model_path']) + target = Path(os.environ['TWINKLE_TEST_ARTIFACT_DIR']) / 'full-checkpoint-B' + target.mkdir(parents=True, exist_ok=True) + for file in source.iterdir(): + if file.is_file() and not (target / file.name).exists(): + (target / file.name).symlink_to(file) + index = json.loads((source / 'model.safetensors.index.json').read_text()) + shard = index['weight_map']['model.embed_tokens.weight'] + output = target / shard + if output.is_symlink(): + weights = load_file(str(source / shard), device='cpu') + weights['model.embed_tokens.weight'] = torch.roll(weights['model.embed_tokens.weight'], 257, 0).contiguous() + output.unlink() + save_file(weights, str(output), metadata={'format': 'pt'}) + return [str(source), str(target)] + + +def _checkpoint_handles(clients, dialect, paths): + if dialect == 'tinker': + return [clients['service'].create_sampling_client(base_model=clients['model_id'], model_path=path) + for path in paths] + return paths + + +def _assert_adapter_isolation(clients, adapters, dialect): + handles = ([None] + [clients['service'].create_sampling_client(base_model=clients['model_id'], model_path=path) + for path in adapters] if dialect == 'tinker' else [None, *adapters]) + ids = _prompt(clients, 'Choose one word: apple, banana, orange, purple.') + solo = [_sample(clients, dialect, ids, 8, adapter=handle)[0] for handle in handles] + # Default vLLM kernels do not guarantee batch-invariant continuations or + # logprobs. Homogeneous PPU batches also diverge after this stable prefix. + # Require distinct adapter fingerprints, then check routing against them; + # _sample still validates every generated token's finite, aligned logprob. + fingerprints = [tuple(tokens[:2]) for tokens, _logprobs in solo] + assert all(len(fingerprint) == 2 for fingerprint in fingerprints) + assert len(set(fingerprints)) == 3, 'the adapters must change the result, rather than silently falling back to base' + results, elapsed = _parallel([ + lambda i=i: _sample(clients, dialect, ids, 8, adapter=handles[i % 3]) for i in range(24)]) + for i, result in enumerate(results): + tokens, _logprobs = result[0] + assert tuple(tokens[:2]) == fingerprints[i % 3] + _report(check='lora_isolation', dialect=dialect, requests=24, variants=3, seconds=elapsed, + fingerprints=fingerprints) + + +def _real_sampler_actors(): + import ray + from ray.util.state import list_actors + ray.init(address=os.environ['RAY_ADDRESS'], namespace='default', ignore_reinit_error=True, + log_to_driver=False) + actors = list_actors(address=os.environ['TWINKLE_TEST_RAY_DASHBOARD_URL'], + filters=[('class_name', '=', os.environ.get('TWINKLE_TEST_SAMPLER_CLASS', 'VLLMSamplerTQ')), + ('state', '=', 'ALIVE')], + detail=True) + assert len(actors) == int(os.environ.get('TWINKLE_TEST_DP', '1')) + assert all(actor.state == 'ALIVE' for actor in actors) + return [ray.get_actor(actor.name, namespace='default') for actor in actors] + pytestmark = pytest.mark.skipif(os.environ.get('TWINKLE_TEST_GPU_E2E', '0') != '1', reason='Requires a running real vLLM sampler server') diff --git a/tests/server/sampler/sampling_test_helpers.py b/tests/server/sampler/sampling_test_helpers.py deleted file mode 100644 index 17f52d1a9..000000000 --- a/tests/server/sampler/sampling_test_helpers.py +++ /dev/null @@ -1,151 +0,0 @@ -"""Shared route fixtures for sampler concurrency and stream lifecycle tests.""" -from __future__ import annotations - -import asyncio -import threading - -import pytest_asyncio -from fastapi import FastAPI -from starlette.requests import Request -from tinker import types as tinker_types - -from twinkle.data_format import SampledSequence, SampleResponse -from twinkle.protocol.types import SampleRequest -from twinkle.server.sampler.app import SamplerManagement -from twinkle.server.sampler.concurrency import SamplerConcurrency -from twinkle.server.sampler.tinker_handlers import _register_tinker_sampler_routes -from twinkle.server.sampler.twinkle_handlers import _register_twinkle_sampler_routes -from twinkle.server.task_queue.backend_gate import BackendGate -from twinkle.server.task_queue.config import TaskQueueConfig -from twinkle.server.task_queue.worker import ComputeWorker - - -async def _until(predicate): - async def wait(): - while not predicate(): - await asyncio.sleep(0.001) - await asyncio.wait_for(wait(), 3) - - -class _State: - def __init__(self): - self.records = {} - - async def store_future_status(self, request_id, status, model_id=None, **kwargs): - self.records[request_id] = {'status': status, **kwargs} - - async def get_future(self, request_id): - return self.records.get(request_id) - - -class _Limiter: - async def check_and_record(self, *args): - return True, None - - -class _Sampler: - def __init__(self): - self.weight = 0 - self.loads = [] - self.active = 0 - self.peak = 0 - self.entered = threading.Event() - self.release = threading.Event() - self.release.set() - self._lock = threading.Lock() - self.template_calls = 0 - - def set_template(self, *args, **kwargs): - self.template_calls += 1 - - def load_full_weights_from_path(self, path): - self.weight = int(path.rsplit('/', 1)[-1]) - self.loads.append(self.weight) - - def sample(self, inputs, sampling_params=None, **kwargs): - with self._lock: - self.active += 1 - self.peak = max(self.peak, self.active) - self.entered.set() - try: - if not self.release.wait(3): - raise TimeoutError('test did not release generation') - token = self.weight or inputs[0]['input_ids'][0] - return [SampleResponse( - prompt_token_ids=list(item['input_ids']), - sequences=[SampledSequence(stop_reason='length', tokens=[token], logprobs=[[(token, -0.1)]])], - ) for item in inputs] - finally: - with self._lock: - self.active -= 1 - - -class _Service(SamplerManagement): - def __init__(self, *, timeout=2, limit=128): - self.model_id = 'mock-model' - self.replica_id = 'sampler-test' - self.state = _State() - self.sampler = _Sampler() - self.sampling_guard = SamplerConcurrency() - self._task_queue_config = TaskQueueConfig(execution_timeout=timeout, concurrent_task_limit=limit) - self._rate_limiter = _Limiter() - self._task_metrics = None - self._deployment_name = None - self._event_loop = None - self._backend_gate = BackendGate() - self._compute_worker = ComputeWorker( - state=self.state, config=self._task_queue_config, task_metrics=None, deployment_name=None) - app = FastAPI() - _register_tinker_sampler_routes(app, lambda: self) - _register_twinkle_sampler_routes(app, lambda: self) - self.endpoints = {route.path: route.endpoint for route in app.routes} - - async def _on_request_start(self, request): - return 'token' - - async def submit(self, dialect, token=7, path=None): - request = Request({'type': 'http', 'headers': []}) - if dialect == 'tinker': - body = tinker_types.SampleRequest( - prompt=tinker_types.ModelInput.from_ints([token]), model_path=path, - sampling_params=tinker_types.SamplingParams(max_tokens=1, temperature=0)) - ref = await self.endpoints['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/tinker/asample'](request, body, self) - return ref['request_id'] - body = SampleRequest(inputs=[{'input_ids': [token]}], adapter_uri=path, - sampling_params={'max_tokens': 1, 'temperature': 0}) - envelope = await self.endpoints['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/twinkle/sample'](request, body, self) - return envelope.request_id - - async def result(self, request_id): - await _until(lambda: self.state.records[request_id]['status'] in ('completed', 'failed')) - return self.state.records[request_id] - - async def close(self): - self.sampler.release.set() - self.sampling_guard.close() - await self._compute_worker.stop() - await _until(lambda: not self.sampling_guard._jobs) - self._backend_gate.shutdown() - - -def _tokens(record, dialect): - assert record['status'] == 'completed', record - result = record['result'] - if dialect == 'tinker': - return result.sequences[0].tokens - return result['samples'][0]['sequences'][0]['tokens'] - - -@pytest_asyncio.fixture -async def make_service(): - """Create services and close them even when a test assertion fails.""" - services = [] - - def create(**kwargs): - service = _Service(**kwargs) - services.append(service) - return service - - yield create - for service in reversed(services): - await service.close() diff --git a/tests/server/sampler/test_backend_concurrency.py b/tests/server/sampler/test_backend_concurrency.py deleted file mode 100644 index 3590edb27..000000000 --- a/tests/server/sampler/test_backend_concurrency.py +++ /dev/null @@ -1,24 +0,0 @@ -"""Only the server backend opts native TQ workers into ordinary concurrency.""" -import twinkle.infra as infra -from twinkle.infra._ray import RayHelper -from twinkle.utils.device_mesh import DeviceMesh -from twinkle.server.sampler.backends.vllm_async import ServerVLLMSampler -from twinkle.sampler.vllm_sampler.vllm_sampler import SAMPLER_MAX_CONCURRENCY -from twinkle_agentic.async_rl.vllm_sampler_tq import VLLMSamplerTQ - - -def test_server_concurrency_does_not_change_native_async_rl_defaults(monkeypatch): - captured = [] - def create(cls, *args, **kwargs): - captured.append((cls.__name__, kwargs['max_concurrency'])) - return [] - monkeypatch.setattr(infra, '_mode', 'ray') - monkeypatch.setattr(infra, '_device_group', []) - monkeypatch.setattr(RayHelper, 'create_workers', create) - monkeypatch.delenv('WORKER_NAME', raising=False) - monkeypatch.delenv('CLUSTER_NAME', raising=False) - monkeypatch.setenv('TWINKLE_TRUST_REMOTE_CODE', '1') - mesh = object.__new__(DeviceMesh) - for cls in (VLLMSamplerTQ, ServerVLLMSampler): - cls(model_id='test-model', context_manager=None, device_mesh=mesh, remote_group='test-sampler') - assert captured == [('VLLMSamplerTQ', None), ('ServerVLLMSampler', SAMPLER_MAX_CONCURRENCY)] diff --git a/tests/server/sampler/test_mock_sampler.py b/tests/server/sampler/test_mock_sampler.py index 5f432e556..f39828428 100644 --- a/tests/server/sampler/test_mock_sampler.py +++ b/tests/server/sampler/test_mock_sampler.py @@ -261,3 +261,28 @@ def test_tool_call_turns_multiple_rounds() -> None: inp = InputFeature(input_ids=[1]) decoded = [s.sample(inp, SamplingParams(max_tokens=1))[0].sequences[0].decoded for _ in range(3)] assert decoded == ['TC', None, 'TC'] + + +def test_server_concurrency_does_not_change_native_async_rl_defaults(monkeypatch): + pytest.importorskip('transfer_queue') + import twinkle.infra as infra + from twinkle.infra._ray import RayHelper + from twinkle.sampler.vllm_sampler.vllm_sampler import SAMPLER_MAX_CONCURRENCY + from twinkle.server.sampler.backends.vllm_async import ServerVLLMSampler + from twinkle.utils.device_mesh import DeviceMesh + from twinkle_agentic.async_rl.vllm_sampler_tq import VLLMSamplerTQ + + captured = [] + def create(cls, *args, **kwargs): + captured.append((cls.__name__, kwargs['max_concurrency'])) + return [] + monkeypatch.setattr(infra, '_mode', 'ray') + monkeypatch.setattr(infra, '_device_group', []) + monkeypatch.setattr(RayHelper, 'create_workers', create) + monkeypatch.delenv('WORKER_NAME', raising=False) + monkeypatch.delenv('CLUSTER_NAME', raising=False) + monkeypatch.setenv('TWINKLE_TRUST_REMOTE_CODE', '1') + mesh = object.__new__(DeviceMesh) + for cls in (VLLMSamplerTQ, ServerVLLMSampler): + cls(model_id='test-model', context_manager=None, device_mesh=mesh, remote_group='test-sampler') + assert captured == [('VLLMSamplerTQ', None), ('ServerVLLMSampler', SAMPLER_MAX_CONCURRENCY)] diff --git a/tests/server/sampler/test_resolve_sampler_weights.py b/tests/server/sampler/test_resolve_sampler_weights.py index bf3b32fb3..7bbf9f25d 100644 --- a/tests/server/sampler/test_resolve_sampler_weights.py +++ b/tests/server/sampler/test_resolve_sampler_weights.py @@ -26,7 +26,7 @@ async def call_backend(self, fn, *args, **kwargs): # The helper only routes load_full_weights_from_path through call_backend. if getattr(fn, '__name__', None) == 'load_full_weights_from_path': self.full_weight_loads.append(args[0]) - return fn(*args) + return fn(*args, **kwargs) @pytest.mark.asyncio diff --git a/tests/server/sampler/test_sampling_concurrency.py b/tests/server/sampler/test_sampling_concurrency.py index 52a08923f..f1697ec9c 100644 --- a/tests/server/sampler/test_sampling_concurrency.py +++ b/tests/server/sampler/test_sampling_concurrency.py @@ -1,18 +1,160 @@ +"""Sampler route concurrency, weight ownership and request failure recovery.""" from __future__ import annotations import asyncio +import json import threading from types import SimpleNamespace import pytest +import pytest_asyncio +from fastapi import FastAPI +from ray.exceptions import GetTimeoutError from starlette.requests import Request +from tinker import types as tinker_types -from twinkle.server.sampler.concurrency import SamplerConcurrency +from twinkle.data_format import SampledSequence, SampleResponse from twinkle.exceptions import FullWeightLoadError +from twinkle.protocol.types import SampleRequest +from twinkle.server.sampler.app import SamplerManagement +from twinkle.server.sampler.concurrency import SamplerConcurrency +from twinkle.server.sampler.tinker_handlers import _register_tinker_sampler_routes +from twinkle.server.sampler.twinkle_handlers import _register_twinkle_sampler_routes +from twinkle.server.task_queue.backend_gate import BackendGate +from twinkle.server.task_queue.config import TaskQueueConfig from twinkle.server.task_queue.types import BackendBusyError - - -from tests.server.sampler.sampling_test_helpers import _tokens, _until, make_service +from twinkle.server.task_queue.worker import ComputeWorker + + +async def _until(predicate): + async def wait(): + while not predicate(): + await asyncio.sleep(0.001) + await asyncio.wait_for(wait(), 3) + + +class _State: + def __init__(self): + self.records = {} + + async def store_future_status(self, request_id, status, model_id=None, **kwargs): + self.records[request_id] = {'status': status, **kwargs} + + async def get_future(self, request_id): + return self.records.get(request_id) + + +class _Limiter: + async def check_and_record(self, *args): + return True, None + + +class _Sampler: + def __init__(self): + self.weight = 0 + self.loads = [] + self.active = 0 + self.peak = 0 + self.entered = threading.Event() + self.release = threading.Event() + self.release.set() + self._lock = threading.Lock() + self.template_calls = 0 + + def set_template(self, *args, **kwargs): + self.template_calls += 1 + + def load_full_weights_from_path(self, path): + self.weight = int(path.rsplit('/', 1)[-1]) + self.loads.append(self.weight) + + def sample(self, inputs, sampling_params=None, **kwargs): + with self._lock: + self.active += 1 + self.peak = max(self.peak, self.active) + self.entered.set() + try: + if not self.release.wait(3): + raise TimeoutError('test did not release generation') + token = self.weight or inputs[0]['input_ids'][0] + return [SampleResponse( + prompt_token_ids=list(item['input_ids']), + sequences=[SampledSequence(stop_reason='length', tokens=[token], logprobs=[[(token, -0.1)]])], + ) for item in inputs] + finally: + with self._lock: + self.active -= 1 + + +class _Service(SamplerManagement): + def __init__(self, *, timeout=2, limit=128): + self.model_id = 'mock-model' + self.replica_id = 'sampler-test' + self.state = _State() + self.sampler = _Sampler() + self.sampling_guard = SamplerConcurrency() + self._task_queue_config = TaskQueueConfig(execution_timeout=timeout, concurrent_task_limit=limit) + self._rate_limiter = _Limiter() + self._task_metrics = None + self._deployment_name = None + self._event_loop = None + self._backend_gate = BackendGate() + self._compute_worker = ComputeWorker( + state=self.state, config=self._task_queue_config, task_metrics=None, deployment_name=None) + app = FastAPI() + _register_tinker_sampler_routes(app, lambda: self) + _register_twinkle_sampler_routes(app, lambda: self) + self.endpoints = {route.path: route.endpoint for route in app.routes} + + async def _on_request_start(self, request): + return 'token' + + async def submit(self, dialect, token=7, path=None): + request = Request({'type': 'http', 'headers': []}) + if dialect == 'tinker': + body = tinker_types.SampleRequest( + prompt=tinker_types.ModelInput.from_ints([token]), model_path=path, + sampling_params=tinker_types.SamplingParams(max_tokens=1, temperature=0)) + ref = await self.endpoints['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/tinker/asample'](request, body, self) + return ref['request_id'] + body = SampleRequest(inputs=[{'input_ids': [token]}], adapter_uri=path, + sampling_params={'max_tokens': 1, 'temperature': 0}) + envelope = await self.endpoints['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/twinkle/sample'](request, body, self) + return envelope.request_id + + async def result(self, request_id): + await _until(lambda: self.state.records[request_id]['status'] in ('completed', 'failed')) + return self.state.records[request_id] + + async def close(self): + self.sampler.release.set() + self.sampling_guard.close() + await self._compute_worker.stop() + await _until(lambda: not self.sampling_guard._jobs) + self._backend_gate.shutdown() + + +def _tokens(record, dialect): + assert record['status'] == 'completed', record + result = record['result'] + if dialect == 'tinker': + return result.sequences[0].tokens + return result['samples'][0]['sequences'][0]['tokens'] + + +@pytest_asyncio.fixture +async def make_service(): + """Create services and close them even when a test assertion fails.""" + services = [] + + def create(**kwargs): + service = _Service(**kwargs) + services.append(service) + return service + + yield create + for service in reversed(services): + await service.close() @pytest.mark.asyncio @@ -248,17 +390,6 @@ def partial_load(path): await service.check_health() -@pytest.mark.asyncio -async def test_shutdown_before_task_first_step_marks_future_failed(make_service): - from twinkle.server.task_queue.types import QueuedTask - service = make_service() - worker = service._compute_worker - worker.reserve_detached() - await service.state.store_future_status('not-started', 'pending') - worker.run_detached(QueuedTask('not-started', lambda: asyncio.sleep(0), None, None, 0, 'sample', 0), 'q') - await worker.stop() - assert service.state.records['not-started']['failure'].reason_code == 'orphaned_replica' - assert worker._detached_reserved == 0 @pytest.mark.asyncio @@ -291,15 +422,6 @@ async def test_timeout_keeps_job_capacity_and_does_not_overwrite_failed_future(m assert (await service.result(await service.submit('twinkle')))['status'] == 'completed' -@pytest.mark.asyncio -async def test_nested_status_poll_reuses_job_access(make_service): - service = make_service() - async def job(): - # A separate polling Task inherits the job's access context. - await asyncio.wait_for(service.call_backend(service.sampler.set_template, 'test'), timeout=1) - await service.sampling_guard.run(job, exclusive=True) - assert service.sampler.template_calls == 1 - assert not service.sampling_guard._writer @pytest.mark.asyncio @@ -360,3 +482,54 @@ def cancelled(*args, **kwargs): assert service._compute_worker._detached_reserved == 0 await service._compute_worker.stop() assert service.state.records[ref] == record + + +@pytest.mark.asyncio +@pytest.mark.parametrize('dialect', ['tinker', 'twinkle']) +async def test_invalid_checkpoint_rejects_request_without_disabling_sampling(make_service, tmp_path, dialect): + service = make_service() + ref = await service.submit(dialect, path=str(tmp_path / 'missing')) + failed = await service.result(ref) + assert failed['failure'].reason_code == 'request_rejected' + assert failed['failure'].attribution == 'user' + assert _tokens(await service.result(await service.submit(dialect)), dialect) == [7] + + +@pytest.mark.asyncio +async def test_ordinary_backend_timeout_does_not_disable_replica(make_service): + service = make_service() + sample = service.sampler.sample + service.sampler.sample = lambda *a, **k: (_ for _ in ()).throw(GetTimeoutError('one request timed out')) + with pytest.raises(GetTimeoutError): + await service.call_backend(service.sampler.sample) + service.sampler.sample = sample + assert _tokens(await service.result(await service.submit('twinkle')), 'twinkle') == [7] + + +@pytest.mark.asyncio +async def test_stream_route_resolves_adapter_uri(make_service, tmp_path): + service = make_service() + (tmp_path / 'adapter_config.json').write_text(json.dumps({'r': 1})) + service.sampler.reset_prefix_cache = lambda: None + service.sampler._actors = [] + request = Request({'type': 'http', 'headers': []}) + response = await service.endpoints['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/twinkle/sample_stream']( + request, SampleRequest(inputs=[{'input_ids': [7]}], adapter_uri=str(tmp_path)), service) + rows = [json.loads(row) async for row in response.body_iterator] + assert rows[0]['error_code'] == 503 + + +@pytest.mark.asyncio +async def test_sampling_timeout_keeps_capacity_without_blocking_other_readers(make_service): + import asyncio + service = make_service(timeout=0.05, limit=2) + service.sampler.release.clear() + timed_out = await service.submit('tinker') + await _until(service.sampler.entered.is_set) + assert (await service.result(timed_out))['failure'].reason_code == 'execution_timeout' + submitting = asyncio.create_task(service.submit('twinkle')) + await _until(lambda: service.sampler.peak == 2) + service.sampler.release.set() + assert _tokens(await service.result(await submitting), 'twinkle') == [7] + await _until(lambda: service._compute_worker._detached_reserved == 0) + await service.check_health() diff --git a/tests/server/sampler/test_sampling_recovery.py b/tests/server/sampler/test_sampling_recovery.py deleted file mode 100644 index 28ada50a6..000000000 --- a/tests/server/sampler/test_sampling_recovery.py +++ /dev/null @@ -1,61 +0,0 @@ -"""A rejected or timed-out request must not disable unrelated sampling.""" -import json - -import pytest -from ray.exceptions import GetTimeoutError -from starlette.requests import Request - -from twinkle.protocol.types import SampleRequest -from .sampling_test_helpers import _tokens, make_service - - -@pytest.mark.asyncio -@pytest.mark.parametrize('dialect', ['tinker', 'twinkle']) -async def test_invalid_checkpoint_rejects_request_without_disabling_sampling(make_service, tmp_path, dialect): - service = make_service() - ref = await service.submit(dialect, path=str(tmp_path / 'missing')) - failed = await service.result(ref) - assert failed['failure'].reason_code == 'request_rejected' - assert failed['failure'].attribution == 'user' - assert _tokens(await service.result(await service.submit(dialect)), dialect) == [7] - - -@pytest.mark.asyncio -async def test_ordinary_backend_timeout_does_not_disable_replica(make_service): - service = make_service() - sample = service.sampler.sample - service.sampler.sample = lambda *a, **k: (_ for _ in ()).throw(GetTimeoutError('one request timed out')) - with pytest.raises(GetTimeoutError): - await service.call_backend(service.sampler.sample) - service.sampler.sample = sample - assert _tokens(await service.result(await service.submit('twinkle')), 'twinkle') == [7] - - -@pytest.mark.asyncio -async def test_stream_route_resolves_adapter_uri(make_service, tmp_path): - service = make_service() - (tmp_path / 'adapter_config.json').write_text(json.dumps({'r': 1})) - service.sampler.reset_prefix_cache = lambda: None - service.sampler._actors = [] - request = Request({'type': 'http', 'headers': []}) - response = await service.endpoints['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/twinkle/sample_stream']( - request, SampleRequest(inputs=[{'input_ids': [7]}], adapter_uri=str(tmp_path)), service) - rows = [json.loads(row) async for row in response.body_iterator] - assert rows[0]['error_code'] == 503 - - -@pytest.mark.asyncio -async def test_sampling_timeout_keeps_capacity_without_blocking_other_readers(make_service): - import asyncio - from .sampling_test_helpers import _until - service = make_service(timeout=0.05, limit=2) - service.sampler.release.clear() - timed_out = await service.submit('tinker') - await _until(service.sampler.entered.is_set) - assert (await service.result(timed_out))['failure'].reason_code == 'execution_timeout' - submitting = asyncio.create_task(service.submit('twinkle')) - await _until(lambda: service.sampler.peak == 2) - service.sampler.release.set() - assert _tokens(await service.result(await submitting), 'twinkle') == [7] - await _until(lambda: service._compute_worker._detached_reserved == 0) - await service.check_health() diff --git a/tests/server/sampler/test_stream_guarantees.py b/tests/server/sampler/test_stream_guarantees.py index ef3aceafd..60623c970 100644 --- a/tests/server/sampler/test_stream_guarantees.py +++ b/tests/server/sampler/test_stream_guarantees.py @@ -89,8 +89,6 @@ async def test_stream_timeout_returns_error_payload_and_closes_queue(): class _GenerationService: def __init__(self) -> None: - from twinkle.server.sampler.concurrency import SamplerConcurrency - self.sampling_guard = SamplerConcurrency() self.cancelled = False self.sampler = SimpleNamespace( get_generation_status=lambda _submission_id: {'status': 'running'}, @@ -101,9 +99,8 @@ def __init__(self) -> None: async def call_backend(self, fn, /, *args, **kwargs): return fn(*args, **kwargs) - def _cancel(self, _submission_id: str): + def _cancel(self, _submission_id: str) -> None: self.cancelled = True - return {'status': 'completed'} @pytest.mark.asyncio From 07bfb7eee457b3dceda4fc9f44c9814d65d3d9eb Mon Sep 17 00:00:00 2001 From: Yunnglin Date: Wed, 14 Oct 2026 17:34:50 +0800 Subject: [PATCH 5/5] fix(task-queue): skip tasks cancelled before the running write A client cancel can win the race against the running status write. The terminal record is kept, but the task still started and held detached capacity. Re-read the record after the write and return before starting work. Also default the E2E sampler actor lookup to ServerVLLMSampler. --- src/twinkle/server/task_queue/worker.py | 8 +++++ .../test_sampler_concurrency_e2e.py | 2 +- tests/server/utils/test_task_queue_mixin.py | 31 +++++++++++++++++++ 3 files changed, 40 insertions(+), 1 deletion(-) diff --git a/src/twinkle/server/task_queue/worker.py b/src/twinkle/server/task_queue/worker.py index 307b469c9..93d4e165f 100644 --- a/src/twinkle/server/task_queue/worker.py +++ b/src/twinkle/server/task_queue/worker.py @@ -387,6 +387,14 @@ async def _process_task(self, task: QueuedTask, queue_key: str) -> None: task.model_id, queue_state=QueueState.ACTIVE.value, ) + # A client cancel (or orphan cleanup) can win the race against the + # running write; the terminal record is kept, so do not start work + # nobody will collect. Once running is stored, cancel_if_pending can + # no longer succeed, so this read closes the window. + if await self._is_record_terminal(task.request_id): + task_status = 'cancelled' + logger.info(f'[ComputeWorker] Task {task.request_id} became terminal before start; skipping.') + return with traced_operation('task_queue.execute', attrs=queue_attrs): logger.debug(f'[ComputeWorker] Task {task.request_id} executing, ' f'type={task_type}, queue_key={queue_key}') diff --git a/tests/server/integration/test_sampler_concurrency_e2e.py b/tests/server/integration/test_sampler_concurrency_e2e.py index 697b1eab8..740580eaa 100644 --- a/tests/server/integration/test_sampler_concurrency_e2e.py +++ b/tests/server/integration/test_sampler_concurrency_e2e.py @@ -177,7 +177,7 @@ def _real_sampler_actors(): ray.init(address=os.environ['RAY_ADDRESS'], namespace='default', ignore_reinit_error=True, log_to_driver=False) actors = list_actors(address=os.environ['TWINKLE_TEST_RAY_DASHBOARD_URL'], - filters=[('class_name', '=', os.environ.get('TWINKLE_TEST_SAMPLER_CLASS', 'VLLMSamplerTQ')), + filters=[('class_name', '=', os.environ.get('TWINKLE_TEST_SAMPLER_CLASS', 'ServerVLLMSampler')), ('state', '=', 'ALIVE')], detail=True) assert len(actors) == int(os.environ.get('TWINKLE_TEST_DP', '1')) diff --git a/tests/server/utils/test_task_queue_mixin.py b/tests/server/utils/test_task_queue_mixin.py index ebc3bc25c..033c372b4 100644 --- a/tests/server/utils/test_task_queue_mixin.py +++ b/tests/server/utils/test_task_queue_mixin.py @@ -572,3 +572,34 @@ async def write(*args, **kwargs): assert record['status'] == 'failed' assert record['failure']['reason_code'] == 'orphaned_replica' assert [args[1] for args, _ in queue.state.records] == ['pending', 'failed'] + + +@pytest.mark.asyncio +async def test_cancel_winning_running_write_race_skips_execution(make_concurrent_queue): + queue = make_concurrent_queue() + executed = False + original = queue.state.store_future_status + + async def cancel_before_running(*args, **kwargs): + if args[1] == 'running': + # Mirror FutureManager: cancel_if_pending wins, then the terminal record refuses 'running'. + queue.state._latest[args[0]] = {'status': 'cancelled', 'result': None, 'failure': None} + return + await original(*args, **kwargs) + + async def work(): + nonlocal executed + executed = True + return {'ok': True} + + queue.state.store_future_status = cancel_before_running + ref = await queue.schedule_concurrent_task(work, token='token1') + for _ in range(100): + if queue._compute_worker._detached_reserved == 0: + break + await asyncio.sleep(0.001) + + assert not executed + assert queue.state._latest[ref['request_id']]['status'] == 'cancelled' + assert queue._compute_worker._detached_reserved == 0 + assert not queue._compute_worker._detached_tasks