diff --git a/.dev_scripts/ci_container_test.sh b/.dev_scripts/ci_container_test.sh index beb5fe530..b0484a71a 100644 --- a/.dev_scripts/ci_container_test.sh +++ b/.dev_scripts/ci_container_test.sh @@ -1,5 +1,5 @@ install_twinkle_with_kernels() { - pip install ".[test,client,server]" -i https://mirrors.aliyun.com/pypi/simple/ || pip install ".[test,client,server]" + pip install ".[test,client,server,async-rl]" -i https://mirrors.aliyun.com/pypi/simple/ || pip install ".[test,client,server,async-rl]" } if [ "$MODELSCOPE_SDK_DEBUG" == "True" ]; then diff --git a/cookbook/client/async_rl/README.md b/cookbook/client/async_rl/README.md new file mode 100644 index 000000000..8441c3ddb --- /dev/null +++ b/cookbook/client/async_rl/README.md @@ -0,0 +1,156 @@ +# Client-Orchestrated Async GRPO + +The client owns the Dataset, Reward, Advantage, rollout partitions, staleness, +training schedule, and policy publication. The server exposes shared +Multi-LoRA Model, vLLM Sampler, and TransferQueue DataPlane components without +owning the GRPO loop. + +This example deploys a dedicated Qwen3.5-4B GRPO server. It does not combine +different RL algorithms in one deployment. For the YAML-managed multi-tenant +runtime, see [`cookbook/rl`](../../rl/README.md). + +## Resources + +The server configuration uses two GPUs: + +| Component | GPUs | Purpose | +|---|---:|---| +| Multi-LoRA Model | 1 | GRPO forward, backward, optimizer step, and checkpoint save | +| Async vLLM Sampler | 1 | Shared continuous-batched rollout generation | + +The DataPlane keeps token tensors and sampled log-probabilities server-side. +The client reads decoded completions for Reward and Advantage, appends those +values to the same `DataRef`, and sends only references to the Model. + +## Quick start + +Install the server, client, and async-RL dependencies: + +```bash +pip install -e '.[async-rl,client,server]' +``` + +Terminal 1 — start the dedicated component server: + +```bash +export TWINKLE_LOCAL_MODEL_PATH=/absolute/path/to/Qwen3.5-4B +export CUDA_VISIBLE_DEVICES=0,1 +bash cookbook/client/async_rl/run_server.sh +``` + +The server script starts a local Ray cluster when needed, validates +`server_config.yaml`, and launches Ray Serve on port 8000. + +Terminal 2 — start one GRPO client: + +```bash +export TWINKLE_SERVER_URL=http://127.0.0.1:8000 +export TWINKLE_SERVER_TOKEN=EMPTY_TOKEN +export TWINKLE_TEMPLATE_MODEL_ID=/absolute/path/to/Qwen3.5-4B +export TWINKLE_DATASET_ID=/absolute/path/to/gsm8k +python cookbook/client/async_rl/client_orchestrated_grpo.py +``` + +`TWINKLE_TEMPLATE_MODEL_ID` is the tokenizer/template model path used in the +client process. The template class is fixed to `Qwen3_5Template`; a model ID +such as `Qwen/Qwen3.5-4B` is not a template class name. + +## Execution model + +Each DataLoader batch becomes a private client-side `_RolloutPartition` bound +to one immutable adapter checkpoint: + +```text +DataLoader + -> Rollout Worker + -> async vLLM Sampler + -> DataPlane DataRef + -> Advantage Worker + -> Trainer Worker + -> save and publish policy +``` + +The workers are independent asyncio tasks: + +- Prompt groups are submitted separately through + `sampler.asample_to_data_plane()`. +- A complete group is rewarded and made trainable without waiting for the rest + of its partition. +- The Trainer consumes `TRAIN_MINI_BATCH_SIZE / NUM_GENERATIONS` ready groups + at a time. +- Partitions train and publish in FIFO order. +- A policy version advances only after every group in its partition has + trained and the new checkpoint has been saved. + +`MAX_STALENESS` controls the number of live partitions. Setting it to zero +still permits rollout/training overlap inside one partition; values above zero +also permit cross-partition overlap. + +## Training metrics + +After every `clip_grad_and_step()`, the Trainer calls: + +```python +model.calculate_metric(is_training=True) +``` + +and prints the returned metrics to client stdout: + +```text +optimizer_step=1 grad_norm=0.42 learning_rate=2e-05 loss=0.031 +``` + +Partition publication is logged separately: + +```text +partition=0 policy=1 optimizer_step=8 staleness=0 +``` + +## Client settings + +| Environment variable | Default | Purpose | +|---|---|---| +| `TWINKLE_SERVER_URL` | `http://localhost:8000` | Server base URL | +| `TWINKLE_SERVER_TOKEN` | `EMPTY_TOKEN` | Request authentication token | +| `TWINKLE_TEMPLATE_MODEL_ID` | `ms://Qwen/Qwen3.5-4B` | Client tokenizer/template source | +| `TWINKLE_DATASET_ID` | `ms://modelscope/gsm8k` | GSM8K dataset source | +| `TWINKLE_ADAPTER_NAME` | `client-grpo` | LoRA adapter name | +| `TWINKLE_MAX_PARTITIONS` | `100` | Maximum DataLoader batches admitted | +| `TWINKLE_MAX_STALENESS` | `2` | Maximum extra live partitions | +| `TWINKLE_ROLLOUT_CONCURRENCY` | `8` | Concurrent prompt-group submissions | +| `TWINKLE_NUM_GENERATIONS` | `4` | Generations per prompt group | +| `TWINKLE_BATCH_SIZE` | `8` | Prompt groups per partition | +| `TWINKLE_TRAIN_MINI_BATCH_SIZE` | `8` | Samples per optimizer step | +| `TWINKLE_MICRO_BATCH_SIZE` | `4` | Model micro-batch size | +| `TWINKLE_MAX_TOKENS_PER_MICRO_BATCH` | `4096` | Dynamic batching token limit | + +The following must hold: + +```text +TWINKLE_TRAIN_MINI_BATCH_SIZE % TWINKLE_NUM_GENERATIONS == 0 +TWINKLE_BATCH_SIZE % ( + TWINKLE_TRAIN_MINI_BATCH_SIZE / TWINKLE_NUM_GENERATIONS +) == 0 +``` + +## Files + +| File | Role | +|---|---| +| `client_orchestrated_grpo.py` | Client-side async GRPO orchestration | +| `run_server.sh` | Start Ray and the dedicated component server | +| `server_config.yaml` | Qwen3.5-4B Model, Sampler, DataPlane, Gateway, and Processor deployment | + +## Troubleshooting + +- **Client receives HTTP 404 for the model** — use the provided server config; + its public route is fixed to `Qwen/Qwen3.5-4B`, matching the client. +- **The process tries to access ModelScope while offline** — set both + `TWINKLE_LOCAL_MODEL_PATH` on the server and + `TWINKLE_TEMPLATE_MODEL_ID`/`TWINKLE_DATASET_ID` on the client to local + paths. +- **No loss is printed** — confirm the Model request completed and look for an + `optimizer_step=...` line. Metrics are calculated after every optimizer + step, not after each rollout group. +- **Only one GPU is active** — verify Ray sees two GPUs and that Model and + Sampler placement groups were assigned to different devices. diff --git a/cookbook/client/async_rl/client_orchestrated_grpo.py b/cookbook/client/async_rl/client_orchestrated_grpo.py new file mode 100644 index 000000000..0f3162569 --- /dev/null +++ b/cookbook/client/async_rl/client_orchestrated_grpo.py @@ -0,0 +1,417 @@ +"""Client-orchestrated async GRPO built from the low-level component APIs.""" +from __future__ import annotations + +import asyncio +import inspect +import os +from collections import deque +from dataclasses import dataclass, field +from typing import Any + +from peft import LoraConfig + +from twinkle.advantage import GRPOAdvantage +from twinkle.dataloader import DataLoader +from twinkle.dataset import Dataset, DatasetMeta +from twinkle.preprocessor.llm import GSM8KProcessor +from twinkle.reward import GSM8KAccuracyReward +from twinkle_client import DataPlaneClient, init_twinkle_client +from twinkle_client.async_rl import Worker, WorkerPipeline +from twinkle_client.common.json_utils import json_safe +from twinkle_client.model import MultiLoraTransformersModel +from twinkle_client.sampler import vLLMSampler + +BASE_MODEL = 'Qwen/Qwen3.5-4B' +MODEL_ID = f'ms://{BASE_MODEL}' +TEMPLATE_MODEL_ID = os.environ.get('TWINKLE_TEMPLATE_MODEL_ID', MODEL_ID) +TEMPLATE_CLS = 'Qwen3_5Template' +DATASET_ID = os.environ.get('TWINKLE_DATASET_ID', 'ms://modelscope/gsm8k') +ADAPTER_NAME = os.environ.get('TWINKLE_ADAPTER_NAME', 'client-grpo') +MAX_PARTITIONS = int(os.environ.get('TWINKLE_MAX_PARTITIONS', '100')) +MAX_STALENESS = int(os.environ.get('TWINKLE_MAX_STALENESS', '2')) +ROLLOUT_CONCURRENCY = int(os.environ.get('TWINKLE_ROLLOUT_CONCURRENCY', '8')) +NUM_GENERATIONS = int(os.environ.get('TWINKLE_NUM_GENERATIONS', '4')) +BATCH_SIZE = int(os.environ.get('TWINKLE_BATCH_SIZE', '8')) +TRAIN_MINI_BATCH_SIZE = int(os.environ.get('TWINKLE_TRAIN_MINI_BATCH_SIZE', '8')) +MICRO_BATCH_SIZE = int(os.environ.get('TWINKLE_MICRO_BATCH_SIZE', '4')) +MAX_TOKENS_PER_MICRO_BATCH = int(os.environ.get('TWINKLE_MAX_TOKENS_PER_MICRO_BATCH', '4096')) + + +@dataclass(frozen=True) +class _Policy: + version: int + adapter_uri: str + + +@dataclass +class _RolloutPartition: + """One DataLoader batch bound to one immutable policy snapshot.""" + + partition_id: int + policy: _Policy + prompts: list[dict[str, Any]] + rollouts: list[asyncio.Task[Any]] + ready: asyncio.Queue['_ReadyGroup'] = field(default_factory=asyncio.Queue) + + +@dataclass +class _ReadyGroup: + group_index: int + ref: Any + + +@dataclass +class _RolloutResult: + partition: _RolloutPartition + group_index: int + prompt: dict[str, Any] + ref: Any + + +class _GRPOState: + + def __init__(self, policy: _Policy): + self.policy = policy + self.live: deque[_RolloutPartition] = deque() + self.input_done = False + self.failure: BaseException | None = None + self.condition = asyncio.Condition() + + async def wait_for_admission(self) -> _Policy: + async with self.condition: + await self.condition.wait_for( + lambda: self.failure is not None or len(self.live) < MAX_STALENESS + 1) + if self.failure is not None: + raise self.failure + return self.policy + + async def add_partition(self, partition: _RolloutPartition) -> None: + async with self.condition: + self.live.append(partition) + self.condition.notify_all() + + async def finish_input(self) -> None: + async with self.condition: + self.input_done = True + self.condition.notify_all() + + async def fail(self, error: BaseException) -> None: + async with self.condition: + if self.failure is None: + self.failure = error + self.condition.notify_all() + + async def oldest_partition(self) -> _RolloutPartition | None: + async with self.condition: + await self.condition.wait_for( + lambda: self.failure is not None or bool(self.live) or self.input_done) + if self.failure is not None: + raise self.failure + return self.live[0] if self.live else None + + async def publish(self, partition: _RolloutPartition, policy: _Policy) -> None: + async with self.condition: + if not self.live or self.live[0] is not partition: + raise RuntimeError(f'partition {partition.partition_id} attempted out-of-order publication') + self.policy = policy + self.live.popleft() + self.condition.notify_all() + + +def create_dataset() -> Dataset: + dataset = Dataset(DatasetMeta(DATASET_ID, subset_name='main', split='train')) + dataset.set_template(TEMPLATE_CLS, model_id=TEMPLATE_MODEL_ID, max_length=2048, enable_thinking=False) + dataset.map(GSM8KProcessor(system='Put the final answer within \\boxed{}.')) + dataset.encode(add_generation_prompt=True) + return dataset + + +async def rollout_group( + sampler: vLLMSampler, + prompt: dict[str, Any], + policy: _Policy, + semaphore: asyncio.Semaphore, + group_id: str, +) -> Any: + """Submit one GRPO group and keep its sample-level TQ DataRef alive.""" + async with semaphore: + return await _submit( + sampler.asample_to_data_plane, + [prompt], + adapter_name=ADAPTER_NAME, + adapter_uri=policy.adapter_uri, + policy_version=policy.version, + group_ids=[group_id], + sampling_params={ + 'max_tokens': 1024, + 'temperature': 1.0, + 'top_p': 0.95, + 'logprobs': 1, + }, + num_samples=NUM_GENERATIONS, + ) + + +def start_partition( + partition_id: int, + batch: list[dict[str, Any]], + policy: _Policy, + sampler: vLLMSampler, + semaphore: asyncio.Semaphore, +) -> _RolloutPartition: + """Capture the snapshot before submitting any rollout in this partition.""" + prompts = json_safe(batch) + return _RolloutPartition( + partition_id=partition_id, + policy=policy, + prompts=prompts, + rollouts=[ + asyncio.create_task( + rollout_group( + sampler, + prompt, + policy, + semaphore, + f'partition-{partition_id}/group-{group_index}', + )) + for group_index, prompt in enumerate(prompts) + ], + ) + + +async def _submit(method, *args, **kwargs): + if inspect.iscoroutinefunction(method): + return await method(*args, **kwargs) + task = await asyncio.to_thread(method, *args, **kwargs) + if inspect.isawaitable(task): + return await task + return task + + +def _checkpoint_path(saved: Any) -> str: + if isinstance(saved, dict): + return str(saved['twinkle_path']) + return str(saved.twinkle_path) + + +class _RolloutWorker(Worker): + + def __init__(self, dataloader, sampler, state: _GRPOState, output: asyncio.Queue, semaphore): + super().__init__('rollout') + self.dataloader = dataloader + self.sampler = sampler + self.state = state + self.output = output + self.semaphore = semaphore + + async def _collect(self, partition, group_index, task): + try: + ref = await task + await self.output.put( + _RolloutResult(partition, group_index, partition.prompts[group_index], ref)) + except BaseException as error: + await self.state.fail(error) + raise + + async def run(self) -> None: + batches = iter(self.dataloader) + collectors: list[asyncio.Task] = [] + try: + for partition_id in range(MAX_PARTITIONS): + policy = await self.state.wait_for_admission() + batch = next(batches, None) + if batch is None: + break + prompts = batch if isinstance(batch, list) else [batch] + if len(prompts) != BATCH_SIZE: + print(f'dropping incomplete final batch with {len(prompts)} prompts') + break + partition = start_partition( + partition_id, prompts, policy, self.sampler, self.semaphore) + await self.state.add_partition(partition) + collectors.extend( + asyncio.create_task(self._collect(partition, index, task)) + for index, task in enumerate(partition.rollouts) + ) + await self.state.finish_input() + await asyncio.gather(*collectors) + await self.output.put(None) + except BaseException as error: + await self.state.fail(error) + for task in collectors: + if not task.done(): + task.cancel() + await asyncio.gather(*collectors, return_exceptions=True) + raise + + +class _AdvantageWorker(Worker): + + def __init__(self, data_plane, state: _GRPOState, source: asyncio.Queue): + super().__init__('advantage') + self.data_plane = data_plane + self.state = state + self.source = source + + async def run(self) -> None: + try: + while True: + result = await self.source.get() + if result is None: + return + group_id = f'partition-{result.partition.partition_id}/group-{result.group_index}' + try: + rows = await self.data_plane.aget(result.ref, fields=['decoded']) + if len(rows) != NUM_GENERATIONS: + raise RuntimeError( + f'group {group_id} expected {NUM_GENERATIONS} generations, ' + f'got {len(rows)}') + trajectories = [] + for row in rows: + messages = list(result.prompt.get('messages') or []) + messages.append({'role': 'assistant', 'content': row.get('decoded') or ''}) + trajectories.append({**result.prompt, 'messages': messages}) + rewards = await asyncio.to_thread(GSM8KAccuracyReward(), trajectories) + advantages = await asyncio.to_thread( + GRPOAdvantage(), rewards, num_generations=NUM_GENERATIONS) + ref = await self.data_plane.aappend( + result.ref, + [{ + 'reward': float(reward), + 'advantage': float(advantage), + } for reward, advantage in zip(rewards, advantages)], + ) + await result.partition.ready.put( + _ReadyGroup(result.group_index, ref)) + except BaseException: + await self.data_plane.arelease(result.ref) + raise + except BaseException as error: + await self.state.fail(error) + raise + + +class _TrainerWorker(Worker): + + def __init__(self, model, data_plane, state: _GRPOState): + super().__init__('trainer') + self.model = model + self.data_plane = data_plane + self.state = state + self.optimizer_step = 0 + + async def _train(self, groups: list[_ReadyGroup]) -> None: + refs = [group.ref for group in groups] + try: + await _submit( + self.model.forward_backward_from_data_plane, + refs, + input_field='train_input', + kwarg_fields={ + 'old_logps': 'sampled_logprobs', + 'advantages': 'advantage', + }, + dynamic_batching=True, + micro_batch_size=MICRO_BATCH_SIZE, + max_tokens_per_micro_batch=MAX_TOKENS_PER_MICRO_BATCH, + ) + await _submit(self.model.clip_grad_and_step, max_grad_norm=1.0) + self.optimizer_step += 1 + metric_response = await _submit(self.model.calculate_metric, is_training=True) + metrics = dict( + metric_response['result'] + if isinstance(metric_response, dict) + else metric_response.result + ) + values = ' '.join(f'{name}={value}' for name, value in sorted(metrics.items())) + print(f'optimizer_step={self.optimizer_step} {values}'.rstrip()) + finally: + await asyncio.gather(*(self.data_plane.arelease(ref) for ref in refs)) + + async def run(self) -> None: + try: + while True: + partition = await self.state.oldest_partition() + if partition is None: + return + staleness = self.state.policy.version - partition.policy.version + if staleness > MAX_STALENESS: + raise RuntimeError( + f'partition {partition.partition_id} staleness {staleness} exceeds {MAX_STALENESS}') + groups_per_step = TRAIN_MINI_BATCH_SIZE // NUM_GENERATIONS + ready = [] + for _ in range(len(partition.rollouts)): + ready.append(await partition.ready.get()) + if len(ready) == groups_per_step: + await self._train(ready) + ready.clear() + if ready: + raise RuntimeError('partition ended with an incomplete train mini-batch') + publish_version = self.state.policy.version + 1 + saved = await _submit(self.model.save, f'policy-{publish_version}') + policy = _Policy(publish_version, _checkpoint_path(saved)) + await self.state.publish(partition, policy) + print( + f'partition={partition.partition_id} policy={policy.version} ' + f'optimizer_step={self.optimizer_step} staleness={staleness}') + except BaseException as error: + await self.state.fail(error) + raise + + +async def run_grpo( + dataloader: DataLoader, + model: MultiLoraTransformersModel, + sampler: vLLMSampler, + data_plane: DataPlaneClient, +) -> None: + """Overlap rollout partitions while training and publishing them in FIFO order.""" + if MAX_STALENESS < 0: + raise ValueError('MAX_STALENESS must be non-negative') + if min(ROLLOUT_CONCURRENCY, NUM_GENERATIONS, BATCH_SIZE, TRAIN_MINI_BATCH_SIZE) <= 0: + raise ValueError('rollout concurrency and all batch sizes must be positive') + if TRAIN_MINI_BATCH_SIZE % NUM_GENERATIONS: + raise ValueError('TRAIN_MINI_BATCH_SIZE must be divisible by NUM_GENERATIONS') + groups_per_step = TRAIN_MINI_BATCH_SIZE // NUM_GENERATIONS + if BATCH_SIZE % groups_per_step: + raise ValueError('BATCH_SIZE * NUM_GENERATIONS must be divisible by TRAIN_MINI_BATCH_SIZE') + + initial = await _submit(model.save, 'policy-0') + state = _GRPOState(_Policy(version=0, adapter_uri=_checkpoint_path(initial))) + semaphore = asyncio.Semaphore(ROLLOUT_CONCURRENCY) + rollout_results: asyncio.Queue = asyncio.Queue() + await WorkerPipeline(( + _RolloutWorker(dataloader, sampler, state, rollout_results, semaphore), + _AdvantageWorker(data_plane, state, rollout_results), + _TrainerWorker(model, data_plane, state), + )).run() + + +async def train() -> None: + client = init_twinkle_client( + base_url=os.environ.get('TWINKLE_SERVER_URL', 'http://localhost:8000'), + api_key=os.environ.get('TWINKLE_SERVER_TOKEN', 'EMPTY_TOKEN'), + ) + try: + model = MultiLoraTransformersModel(MODEL_ID) + sampler = vLLMSampler(MODEL_ID) + data_plane = DataPlaneClient() + + model.add_adapter_to_model( + ADAPTER_NAME, + LoraConfig(target_modules='all-linear', r=8, lora_alpha=32, lora_dropout=0.05), + ) + model.set_loss('GRPOLoss', epsilon=0.2, beta=0.0) + model.set_optimizer('AdamW', lr=2e-5) + model.set_processor('InputProcessor', padding_free=False) + model.set_template(TEMPLATE_CLS, model_id=TEMPLATE_MODEL_ID) + sampler.set_template(TEMPLATE_CLS, model_id=TEMPLATE_MODEL_ID) + + dataloader = DataLoader(dataset=create_dataset(), batch_size=BATCH_SIZE, num_workers=0) + await run_grpo(dataloader, model, sampler, data_plane) + finally: + client.close() + + +if __name__ == '__main__': + asyncio.run(train()) diff --git a/cookbook/client/async_rl/run_server.sh b/cookbook/client/async_rl/run_server.sh new file mode 100755 index 000000000..d3a235d7c --- /dev/null +++ b/cookbook/client/async_rl/run_server.sh @@ -0,0 +1,32 @@ +#!/usr/bin/env bash +set -euo pipefail + +repo_root="$(cd "$(dirname "${BASH_SOURCE[0]}")/../../.." && pwd)" +cd "${repo_root}" + +: "${TWINKLE_LOCAL_MODEL_PATH:?Set TWINKLE_LOCAL_MODEL_PATH to the local Qwen3.5-4B directory}" + +ray_port="${RAY_PORT:-6379}" +ray_address="${RAY_ADDRESS:-127.0.0.1:${ray_port}}" + +if ! command -v ray >/dev/null 2>&1; then + echo "ray is not installed; install the async-RL dependencies first" >&2 + exit 1 +fi +if ! command -v twinkle-server >/dev/null 2>&1; then + echo "twinkle-server is not installed; run: pip install -e '.[async-rl,client]'" >&2 + exit 1 +fi + +if ! ray status --address="${ray_address}" >/dev/null 2>&1; then + ray start \ + --head \ + --port="${ray_port}" \ + --num-gpus="${RAY_NUM_GPUS:-2}" \ + --include-dashboard=false \ + --disable-usage-stats +fi + +config=cookbook/client/async_rl/server_config.yaml +twinkle-server check-config -c "${config}" +exec twinkle-server launch -c "${config}" diff --git a/cookbook/client/async_rl/server_config.yaml b/cookbook/client/async_rl/server_config.yaml new file mode 100644 index 000000000..103ea523f --- /dev/null +++ b/cookbook/client/async_rl/server_config.yaml @@ -0,0 +1,160 @@ +# Twinkle Server Configuration - local Qwen3.5-4B for client-orchestrated async RL + +# Set the absolute Hugging Face-compatible model directory before loading this file: +# export TWINKLE_LOCAL_MODEL_PATH=/absolute/path/to/Qwen3.5-4B +# export TWINKLE_TEMPLATE_MODEL_ID=/absolute/path/to/Qwen3.5-4B # in every client process +# +# The public HTTP model name remains Qwen/Qwen3.5-4B. Both the training Model +# and the vLLM Sampler load the same local directory, so neither component +# downloads the base model independently. + +proxy_location: EveryNode + +http_options: + host: 0.0.0.0 + port: 8000 + +telemetry: + enabled: false + otlp_endpoint: http://localhost:4317 + +persistence: + mode: file + file_path: /tmp/twinkle_state.json + +applications: + + # Gateway and Future query endpoint. + - name: server + route_prefix: /api/v1 + import_path: server + args: + server_config: + per_token_model_limit: 3 + supported_models: + - Qwen/Qwen3.5-4B + deployments: + - name: TinkerCompatServer + max_ongoing_requests: 50 + autoscaling_config: + min_replicas: 1 + max_replicas: 1 + target_ongoing_requests: 128 + ray_actor_options: + num_cpus: 0.1 + runtime_env: + env_vars: + TWINKLE_FAIL_FAST: "0" + + # TransferQueue-backed DataRef service. + - name: data-plane + route_prefix: /api/v1/data-plane + import_path: data_plane + args: + config: + backend: + SimpleStorage: + num_data_storage_units: 2 + deployments: + - name: DataPlaneManagement + num_replicas: 1 + ray_actor_options: + num_cpus: 1 + + # One GPU hosts the training base model and multiple LoRA adapters. + - name: models-Qwen3.5-4B + route_prefix: /api/v1/model/Qwen/Qwen3.5-4B + import_path: model + args: + backend: transformers + model_id: ${oc.env:TWINKLE_LOCAL_MODEL_PATH} + max_loras: 8 + max_length: 10240 + data_plane_url: http://127.0.0.1:8000/api/v1/data-plane + nproc_per_node: 1 + device_group: + name: model + ranks: 1 + device_type: cuda + device_mesh: + device_type: cuda + dp_size: 1 + queue_config: + rps_limit: 100 + tps_limit: 100000 + adapter_config: + adapter_timeout: 30 + deployments: + - name: ModelManagement + autoscaling_config: + min_replicas: 1 + max_replicas: 1 + target_ongoing_requests: 16 + ray_actor_options: + num_cpus: 0.1 + runtime_env: + env_vars: + TWINKLE_TRUST_REMOTE_CODE: "1" + TWINKLE_FAIL_FAST: "0" + + # A second GPU hosts vLLM and loads the same local base model. + - name: sampler-Qwen3.5-4B + route_prefix: /api/v1/sampler/Qwen/Qwen3.5-4B + import_path: sampler + args: + model_id: ${oc.env:TWINKLE_LOCAL_MODEL_PATH} + data_plane_url: http://127.0.0.1:8000/api/v1/data-plane + nproc_per_node: 1 + sampler_type: vllm_async + engine_args: + max_model_len: 4096 + gpu_memory_utilization: 0.5 + enable_lora: true + max_loras: 8 + logprobs_mode: processed_logprobs + device_group: + name: sampler + ranks: 1 + device_type: cuda + device_mesh: + device_type: cuda + dp_size: 1 + queue_config: + rps_limit: 100 + tps_limit: 100000 + deployments: + - name: SamplerManagement + autoscaling_config: + min_replicas: 1 + max_replicas: 1 + target_ongoing_requests: 16 + ray_actor_options: + num_cpus: 0.1 + runtime_env: + env_vars: + TWINKLE_TRUST_REMOTE_CODE: "1" + TWINKLE_FAIL_FAST: "0" + + - name: processor + route_prefix: /api/v1/processor + import_path: processor + args: + ncpu_proc_per_node: 2 + device_group: + name: model + ranks: 2 + device_type: CPU + device_mesh: + device_type: CPU + dp_size: 2 + deployments: + - name: ProcessorManagement + autoscaling_config: + min_replicas: 1 + max_replicas: 1 + target_ongoing_requests: 128 + ray_actor_options: + num_cpus: 0.1 + runtime_env: + env_vars: + TWINKLE_FAIL_FAST: "0" diff --git a/cookbook/rl/async_rl/README.md b/cookbook/rl/async_rl/README.md new file mode 100644 index 000000000..4c8188269 --- /dev/null +++ b/cookbook/rl/async_rl/README.md @@ -0,0 +1,129 @@ +# Async Multi-LoRA GRPO + +One YAML configuration launches two LoRA tenants over a shared training model, +vLLM sampler, and TransferQueue data plane. Rollout, advantage calculation, and +training run as independent workers, while each tenant keeps its own dataset, +reward, optimizer, scheduler, partitions, and policy versions. + +This directory contains the YAML-managed CLI workflow. For client-orchestrated +GRPO over HTTP, see +[`cookbook/client/async_rl`](../../client/async_rl/README.md). + +## Resources + +The default configuration uses three GPUs: + +| Component | GPUs | Purpose | +|---|---:|---| +| Training model | 2 | Native FSDP training for all LoRA tenants | +| vLLM sampler | 1 | Shared rollout generation | + +The example has two GSM8K tenants. Each tenant consumes 128 prompts in eight +partitions: + +```text +16 prompts × 4 generations = 64 samples per partition +64 samples ÷ mini_batch_size 4 = 16 optimizer steps per partition +8 partitions × 16 optimizer steps = 128 optimizer steps per tenant +``` + +## Quick start + +Install the async-RL dependencies: + +```bash +pip install -e '.[async-rl]' +``` + +Set a model and datasets that both training processes and Ray workers can +access: + +```bash +export MODEL_ID=/absolute/path/to/Qwen3.5-4B +export TENANT_A_DATASET_ID=/absolute/path/to/gsm8k +export TENANT_B_DATASET_ID=/absolute/path/to/gsm8k +export CUDA_VISIBLE_DEVICES=0,1,2 +``` + +`MODEL_ID` must be a Hugging Face-compatible directory readable by both +Transformers and vLLM. Dataset values may be local ModelScope-compatible paths +or `ms://...` identifiers. Use local paths when running offline. + +Start a local Ray cluster when needed and launch training: + +```bash +bash cookbook/rl/async_rl/run_async_multi_lora_grpo.sh +``` + +If Ray is already running, the script reuses it. To launch the Python entry +point directly: + +```bash +python cookbook/rl/async_rl/async_multi_lora_grpo.py \ + --config cookbook/rl/async_rl/async_multi_lora_grpo.yaml +``` + +## Scheduling and staleness + +The workers exchange complete prompt groups through TransferQueue: + +```text +RolloutWorker -> TransferQueue -> AdvantageWorker -> TrainerWorker +``` + +`max_staleness` limits live partitions, not mini-batches within one partition: + +- `max_staleness=0` allows one live partition. Training may still overlap with + rollout inside that partition once enough complete prompt groups form a + mini-batch. +- `max_staleness=1` allows two live partitions, so training an older partition + may also overlap with rollout for the next partition. + +The default `mini_batch_size=4` equals one complete four-generation prompt +group, allowing the Trainer to consume a group without waiting for all 16 +groups in the partition. A policy version is published only after every +mini-batch in the partition has trained. + +Batch settings must satisfy: + +```text +partition_samples = rollout.batch_size × rollout.num_generations +partition_samples % train.mini_batch_size == 0 +train.mini_batch_size % rollout.num_generations == 0 +train.mini_batch_size % model_dp == 0 +``` + +## Outputs + +| Output | Default path | +|---|---| +| LoRA checkpoints | `output/async_multi_lora_grpo/` | +| Rollout JSONL | `output/async_multi_lora_grpo/rollouts/` | +| Metrics | `outputs/async_rl/metrics.jsonl` | +| Metrics summary | `outputs/async_rl/summary.json` | + +Set `rollout_output.enabled: false` when benchmarking throughput to avoid +writing one JSONL file per completed prompt group. + +## Files + +| File | Role | +|---|---| +| `async_multi_lora_grpo.py` | Load the YAML and run `AsyncMultiLoraGRPOPipeline` | +| `async_multi_lora_grpo.yaml` | Model, sampler, scheduler, LoRA, dataset, reward, and tenant configuration | +| `run_async_multi_lora_grpo.sh` | Validate required environment variables, start local Ray if needed, and launch training | + +## Troubleshooting + +- **A placement group stays pending** — the default run needs three visible + GPUs. Check `ray status` and make sure Ray was started with at least three + GPUs. +- **The process tries to download a model or dataset** — replace every + `MODEL_ID` and `TENANT_*_DATASET_ID` value with an absolute local path visible + to all Ray workers. +- **Async is slower with `max_staleness=0`** — this setting disables + cross-partition overlap but retains Worker, Ray RPC, and TransferQueue + overhead. Partition-internal overlap also depends on prompt completion + distribution. +- **Rollout files dominate runtime** — disable `rollout_output.enabled` for + performance measurements. diff --git a/cookbook/rl/async_rl/async_multi_lora_grpo.py b/cookbook/rl/async_rl/async_multi_lora_grpo.py new file mode 100644 index 000000000..f6c19b8c4 --- /dev/null +++ b/cookbook/rl/async_rl/async_multi_lora_grpo.py @@ -0,0 +1,21 @@ +"""Launch native-TQ async multi-LoRA GRPO from one YAML configuration.""" + +from __future__ import annotations + +import argparse + +from omegaconf import OmegaConf + +from twinkle_agentic.async_rl import AsyncMultiLoraGRPOPipeline + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument('--config', default='cookbook/rl/async_rl/async_multi_lora_grpo.yaml') + args = parser.parse_args() + config = OmegaConf.to_container(OmegaConf.load(args.config), resolve=True) + print(AsyncMultiLoraGRPOPipeline.from_config(config).run()) + + +if __name__ == '__main__': + main() diff --git a/cookbook/rl/async_rl/async_multi_lora_grpo.yaml b/cookbook/rl/async_rl/async_multi_lora_grpo.yaml new file mode 100644 index 000000000..36a5c303b --- /dev/null +++ b/cookbook/rl/async_rl/async_multi_lora_grpo.yaml @@ -0,0 +1,136 @@ +runtime: + run_id: async_multi_lora_grpo + model_id: ${oc.env:MODEL_ID,ms://Qwen/Qwen3.5-4B} + mode: ray + model_gpus: 2 + sampler_gpus: 1 + sampler_tp: 1 + sampler_max_loras: 4 + seed: 1 + max_staleness: 1 + max_steps: null + allow_partial_rollout: false + rollout_max_retries: 2 + rollout_retry_delay_s: 0.5 + keep_adapter_versions: 2 + output_dir: output/async_multi_lora_grpo + +metrics: + enabled: true + drain_interval_s: 1.0 + queue_capacity: 10000 + close_timeout_s: 10 + jsonl: + enabled: true + path: outputs/async_rl/metrics.jsonl + summary_path: outputs/async_rl/summary.json + batch_size: 64 + flush_interval_s: 2.0 + swanlab: + enabled: false + mode: local + project: twinkle-rl + name: ${runtime.run_id} + log_dir: outputs/swanlab + +template: + cls: Qwen3_5Template + enable_thinking: false + +model: + strategy: native_fsdp + fsdp_config: + reshard_after_forward: true + sequence_parallel_size: 1 + padding_free: false + max_length: 8192 + +sampler: + max_model_len: 8192 + gpu_memory_utilization: 0.8 + max_num_seqs: 64 + max_num_batched_tokens: 8192 + enforce_eager: false + +rollout_output: + enabled: true + output_dir: ${runtime.output_dir}/rollouts + include_token_ids: false + +tq: + polling_mode: true + storage_units: 2 + +scheduler: + rollout: {policy: round_robin, max_consecutive_units: 1} + advantage: {policy: oldest_partition, max_consecutive_units: 1} + train: {policy: sticky, max_consecutive_units: null} + +lora: + target_modules: all-linear + r: 16 + alpha: 32 + dropout: 0.05 + learning_rate: 5.0e-5 + lr_scheduler: + cls: CosineAnnealingLR + # 128 prompts / 16 prompts per partition * 16 groups per partition. + T_max: 128 + eta_min: 0.0 + +loss: + cls: GRPOLoss + epsilon: 0.2 + +lora_contexts: + - tenant_id: tenant_a + training_run_id: gsm8k_async + adapter_name: tenant_a_gsm8k_lora + reward: + class_path: twinkle.reward.GSM8KAccuracyBrevityReward + dataset: + dataset_id: ${oc.env:TENANT_A_DATASET_ID} + subset_name: main + split: train + data_num: 128 + max_length: 8192 + processor: GSM8KProcessor + system_prompt: >- + You are a helpful math assistant. Solve the problem with minimal but + correct reasoning and put your final answer within \boxed{}. + rollout: + batch_size: 16 + num_generations: 4 + max_tokens: 2048 + temperature: 1.0 + top_p: 0.95 + train: + mini_batch_size: 4 + micro_batch_size: 1 + dynamic_batching: false + + - tenant_id: tenant_b + training_run_id: gsm8k_async + adapter_name: tenant_b_gsm8k_lora + reward: + class_path: twinkle.reward.GSM8KAccuracyReward + dataset: + dataset_id: ${oc.env:TENANT_B_DATASET_ID} + subset_name: main + split: train + data_num: 128 + max_length: 8192 + processor: GSM8KProcessor + system_prompt: >- + You are a helpful math assistant. Solve the problem with minimal but + correct reasoning and put your final answer within \boxed{}. + rollout: + batch_size: 16 + num_generations: 4 + max_tokens: 2048 + temperature: 1.0 + top_p: 0.95 + train: + mini_batch_size: 4 + micro_batch_size: 1 + dynamic_batching: false diff --git a/cookbook/rl/async_rl/run_async_multi_lora_grpo.sh b/cookbook/rl/async_rl/run_async_multi_lora_grpo.sh new file mode 100755 index 000000000..7149a2d84 --- /dev/null +++ b/cookbook/rl/async_rl/run_async_multi_lora_grpo.sh @@ -0,0 +1,30 @@ +#!/usr/bin/env bash +set -euo pipefail + +repo_root="$(cd "$(dirname "${BASH_SOURCE[0]}")/../../.." && pwd)" +cd "${repo_root}" + +: "${MODEL_ID:?Set MODEL_ID to a Hugging Face-compatible model directory or model ID}" +: "${TENANT_A_DATASET_ID:?Set TENANT_A_DATASET_ID to a GSM8K dataset path or ID}" +: "${TENANT_B_DATASET_ID:?Set TENANT_B_DATASET_ID to a GSM8K dataset path or ID}" + +ray_port="${RAY_PORT:-6379}" +ray_address="${RAY_ADDRESS:-127.0.0.1:${ray_port}}" +ray_num_gpus="${RAY_NUM_GPUS:-3}" + +if ! command -v ray >/dev/null 2>&1; then + echo "ray is not installed; install the async-RL dependencies first" >&2 + exit 1 +fi + +if ! ray status --address="${ray_address}" >/dev/null 2>&1; then + ray start \ + --head \ + --port="${ray_port}" \ + --num-gpus="${ray_num_gpus}" \ + --include-dashboard=false \ + --disable-usage-stats +fi + +exec python cookbook/rl/async_rl/async_multi_lora_grpo.py \ + --config cookbook/rl/async_rl/async_multi_lora_grpo.yaml diff --git a/pyproject.toml b/pyproject.toml index 27a720456..08d91473f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -31,6 +31,13 @@ rl = [ "vllm>=0.11", "ray[serve]" ] +async-rl = [ + "vllm>=0.11", + "ray[serve]", + "TransferQueue>=0.1.9.dev0", + "math-verify==0.8.0", + "swanlab>=0.8", +] client = [ "textual>=1.0.0", "plotext>=5.2.0", diff --git a/src/twinkle/infra/_ray/ray_helper.py b/src/twinkle/infra/_ray/ray_helper.py index 6eb991383..281e020dd 100644 --- a/src/twinkle/infra/_ray/ray_helper.py +++ b/src/twinkle/infra/_ray/ray_helper.py @@ -351,10 +351,20 @@ def create_workers(worker_cls: Type[T], 'num_cpus': 0.01, } - if device_type == 'GPU': - worker_options['num_gpus'] = 0.01 - else: - # Use custom resource key for non-GPU accelerators (e.g., NPU). + # The placement group already reserves a GPU worker's full + # allocation and ``CUDA_VISIBLE_DEVICES`` above pins that actor + # to the ranks selected by DeviceGroup. Do not additionally + # request a fractional Ray GPU resource for CUDA workers. + # + # With a fractional request, Ray translates the placement-group + # GPU id through the already narrowed visible-device list. For + # the second GPU that becomes index 1 into ``['1']``, killing + # the worker in ``set_visible_accelerator_ids`` before its + # constructor runs. The CPU claim and placement-group strategy + # are sufficient to place the actor in the GPU-owning bundle. + if device_type != 'GPU': + # Use a custom resource key for non-GPU accelerators + # (for example, NPU). worker_options['resources'] = {device_type: 0.01} worker = worker_cls.options(**worker_options).remote(*args, **kwargs) diff --git a/src/twinkle/infra/_ray/resource_manager.py b/src/twinkle/infra/_ray/resource_manager.py index e9cba0f4d..9b8ba0cef 100644 --- a/src/twinkle/infra/_ray/resource_manager.py +++ b/src/twinkle/infra/_ray/resource_manager.py @@ -9,6 +9,15 @@ logger = get_logger() +def _gpu_placement_group_cpus(node_cpu: int, nproc_per_node: int) -> int: + """Return the logical CPUs reserved by one GPU placement group.""" + cpus_per_proc = int(os.environ.get('TWINKLE_GPU_PG_CPUS_PER_PROC', 4)) + cpus_per_proc = max(cpus_per_proc, 1) + required_cpus = max(nproc_per_node * cpus_per_proc, 1) + node_cap = max(node_cpu // 4, 1) + return min(required_cpus, node_cap) + + class ResourceManager: def __init__(self, nproc_per_node: int, ncpu_proc_per_node: int, groups: List[DeviceGroup]): @@ -97,7 +106,10 @@ def __init__(self, nproc_per_node: int, ncpu_proc_per_node: int, groups: List[De except IndexError: node = self.nodes[0] node_cpu = int(node['Resources']['CPU']) - bundles.append({device_type: nproc_per_node, 'CPU': max(node_cpu // 2, 1)}) + bundles.append({ + device_type: nproc_per_node, + 'CPU': _gpu_placement_group_cpus(node_cpu, nproc_per_node), + }) # CPU placement groups: only create when there are actual CPU processes to allocate. if cpu_proc_count > 0: @@ -150,11 +162,21 @@ def get_visible_devices(): return os.environ.get(Platform.get_platform(device_type).visible_device_env()) if self.placement_groups: - self.visible_devices = ray.get([ - get_visible_devices.options(placement_group=pg, runtime_env={ - 'env_vars': self.noset_env() - }).remote() for pg in self.placement_groups - ]) + visible_device_futures = [] + for pg in self.placement_groups: + probe_options = {'placement_group': pg} + if device_type == 'GPU': + # Ask Ray for the GPUs owned by this placement group. A + # no-GPU probe with NOSET_* only sees the node-global + # CUDA_VISIBLE_DEVICES list, so independently initialized + # Model and Sampler groups both select its first entry. + # The probe is short-lived and releases the bundle before + # component workers are created. + probe_options['num_gpus'] = nproc_per_node + else: + probe_options['resources'] = {device_type: nproc_per_node} + visible_device_futures.append(get_visible_devices.options(**probe_options).remote()) + self.visible_devices = ray.get(visible_device_futures) visible_devices = [] for visible_device in self.visible_devices: diff --git a/src/twinkle/loss/base.py b/src/twinkle/loss/base.py index c7f112c42..945ed6353 100644 --- a/src/twinkle/loss/base.py +++ b/src/twinkle/loss/base.py @@ -11,3 +11,9 @@ class Loss: def __call__(self, inputs: InputFeature, outputs: ModelOutput, **kwargs) -> LossOutput: ... + + def micro_batch_scale(self, inputs: list[InputFeature], indices: list[int]) -> float: + if len(indices) == len(inputs): + return 1.0 + raise NotImplementedError( + f'{self.__class__.__name__} does not support micro-batching, including dynamic batching') diff --git a/src/twinkle/loss/chunked_cross_entropy.py b/src/twinkle/loss/chunked_cross_entropy.py index 0df7b6472..1e829a692 100644 --- a/src/twinkle/loss/chunked_cross_entropy.py +++ b/src/twinkle/loss/chunked_cross_entropy.py @@ -129,6 +129,21 @@ def __init__(self, chunk_size: int, ignore_index: int = -100, reduction: str = ' self.reduction = reduction self.dft = dft + def micro_batch_scale(self, inputs, indices): + if self.reduction == 'sum': + return 1.0 + token_counts = [] + for model_input in inputs: + labels = model_input['labels'] + if hasattr(labels, 'ne'): + token_counts.append(int(labels.ne(self.ignore_index).sum().item())) + else: + token_counts.append(sum(int(token != self.ignore_index) for token in labels)) + total_tokens = sum(token_counts) + if total_tokens == 0: + return 0.0 + return sum(token_counts[index] for index in indices) / total_tokens + def __call__(self, inputs, outputs, **kwargs): labels = inputs['labels'] logps = outputs.get('logps') diff --git a/src/twinkle/loss/cross_entropy.py b/src/twinkle/loss/cross_entropy.py index c1b5225d6..8d3627c45 100644 --- a/src/twinkle/loss/cross_entropy.py +++ b/src/twinkle/loss/cross_entropy.py @@ -12,6 +12,21 @@ def __init__(self, ignore_index: int = -100, reduction='mean', dft: bool = False self.reduction = reduction self.dft = dft + def micro_batch_scale(self, inputs, indices): + if self.reduction == 'sum': + return 1.0 + token_counts = [] + for model_input in inputs: + labels = model_input['labels'] + if hasattr(labels, 'ne'): + token_counts.append(int(labels.ne(self.ignore_index).sum().item())) + else: + token_counts.append(sum(int(token != self.ignore_index) for token in labels)) + total_tokens = sum(token_counts) + if total_tokens == 0: + return 0.0 + return sum(token_counts[index] for index in indices) / total_tokens + def __call__(self, inputs, outputs, **kwargs): labels = inputs['labels'] logps = outputs.get('logps') diff --git a/src/twinkle/loss/dpo.py b/src/twinkle/loss/dpo.py index 7a58acdf9..9099225a2 100644 --- a/src/twinkle/loss/dpo.py +++ b/src/twinkle/loss/dpo.py @@ -11,6 +11,7 @@ from twinkle.data_format import LossOutput from twinkle.loss.base import Loss +from twinkle.utils.rl_tensor_utils import align_per_token_values from twinkle.utils.torch_utils import selective_log_softmax if TYPE_CHECKING: @@ -144,49 +145,6 @@ def __init__( self.reference_free = reference_free self.sft_weight = sft_weight - def _align_logps( - self, - logps: 'torch.Tensor', - target_shape: tuple, - device: 'torch.device', - dtype: 'torch.dtype', - ) -> 'torch.Tensor': - """Align log probabilities to target shape. - - Args: - logps: Input log probabilities tensor - target_shape: Target (batch, seq_len) shape - device: Target device - dtype: Target dtype - - Returns: - Aligned tensor of shape target_shape - """ - import torch - - if not torch.is_tensor(logps): - raise TypeError(f'Expected torch.Tensor, got {type(logps)}') - - if logps.dim() == 1: - logps = logps.unsqueeze(0) - - if logps.shape == target_shape: - return logps.to(device=device, dtype=dtype) - - # Handle tensor with different sequence length - if logps.dim() == 2 and logps.shape[0] == target_shape[0]: - batch_size, target_seq_len = target_shape - src_seq_len = logps.shape[1] - logps = logps.to(device=device, dtype=dtype) - if src_seq_len > target_seq_len: - # Truncate right (keep left part) - may happen in Ray result merging - return logps[:, :target_seq_len] - else: - raise ValueError(f'ref_logps seq_len ({src_seq_len}) < target seq_len ({target_seq_len}). ' - f'This should not happen when both models process the same batch.') - - raise ValueError(f'Cannot align ref_logps shape {logps.shape} to target shape {target_shape}') - def _compute_dpo_loss( self, policy_chosen_logps: 'torch.Tensor', @@ -311,13 +269,18 @@ def __call__( # Handle reference log probs if ref_chosen_logps is not None and ref_rejected_logps is not None: # Pre-computed sequence-level reference log probs provided - reference_chosen_logps = ref_chosen_logps.to(device=device, dtype=dtype) - reference_rejected_logps = ref_rejected_logps.to(device=device, dtype=dtype) + reference_chosen_logps = torch.as_tensor(ref_chosen_logps, device=device, dtype=dtype) + reference_rejected_logps = torch.as_tensor(ref_rejected_logps, device=device, dtype=dtype) elif ref_logps is not None: # Per-token reference log probs provided, need to align and sum - if not torch.is_tensor(ref_logps): - ref_logps = torch.as_tensor(ref_logps) - ref_logps_aligned = self._align_logps(ref_logps, labels.shape, device, dtype) + ref_logps_aligned = align_per_token_values( + ref_logps, + tuple(labels.shape), + device=device, + dtype=dtype, + name='ref_logps', + valid_mask=labels != self.ignore_index, + ) ref_chosen, ref_rejected = self._split_chosen_rejected(ref_logps_aligned) reference_chosen_logps = self._compute_sequence_logps(ref_chosen, chosen_labels) reference_rejected_logps = self._compute_sequence_logps(ref_rejected, rejected_labels) diff --git a/src/twinkle/loss/grpo.py b/src/twinkle/loss/grpo.py index 81e0b9208..85d48a82c 100644 --- a/src/twinkle/loss/grpo.py +++ b/src/twinkle/loss/grpo.py @@ -42,6 +42,9 @@ def __init__( self.require_entropy = entropy_coef > 0.0 self.ignore_index = ignore_index + def micro_batch_scale(self, inputs, indices): + return len(indices) / len(inputs) + def _compute_log_importance_weights( self, per_token_logps: 'torch.Tensor', @@ -400,6 +403,19 @@ class CISPOLoss(GRPOLoss): Clamps the IS weight and uses policy gradient. """ + def micro_batch_scale(self, inputs, indices): + token_counts = [] + for model_input in inputs: + labels = model_input['labels'] + if hasattr(labels, 'ne'): + token_counts.append(int(labels.ne(self.ignore_index).sum().item())) + else: + token_counts.append(sum(int(token != self.ignore_index) for token in labels)) + total_tokens = sum(token_counts) + if total_tokens == 0: + return 0.0 + return sum(token_counts[index] for index in indices) / total_tokens + def _compute_per_token_loss( self, ratio: 'torch.Tensor', @@ -431,6 +447,19 @@ class BNPOLoss(GRPOLoss): Normalizes by total completion tokens across batch. """ + def micro_batch_scale(self, inputs, indices): + token_counts = [] + for model_input in inputs: + labels = model_input['labels'] + if hasattr(labels, 'ne'): + token_counts.append(int(labels.ne(self.ignore_index).sum().item())) + else: + token_counts.append(sum(int(token != self.ignore_index) for token in labels)) + total_tokens = sum(token_counts) + if total_tokens == 0: + return 0.0 + return sum(token_counts[index] for index in indices) / total_tokens + def _aggregate_loss( self, per_token_loss: 'torch.Tensor', diff --git a/src/twinkle/metric/__init__.py b/src/twinkle/metric/__init__.py index f6ac5120d..b80bee74a 100644 --- a/src/twinkle/metric/__init__.py +++ b/src/twinkle/metric/__init__.py @@ -1,10 +1,13 @@ # Copyright (c) ModelScope Contributors. All rights reserved. from .accuracy import Accuracy from .base import Metric +from .buffer import MetricBuffer from .completion_and_reward import CompletionRewardMetric from .dpo import DPOMetric from .embedding import EmbeddingMetric from .grpo import CISPOMetric, GRPOMetric, GSPOMetric, PPOMetric from .loss import LossMetric from .ppo import PPOValueMetric +from .reporting import MetricsReporter, create_metrics_reporter from .train_metric import TrainMetric +from .types import MetricRecord diff --git a/src/twinkle/metric/buffer.py b/src/twinkle/metric/buffer.py new file mode 100644 index 000000000..74d532a8c --- /dev/null +++ b/src/twinkle/metric/buffer.py @@ -0,0 +1,26 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Worker-local metric buffering.""" + +from __future__ import annotations + +import threading + +from .types import MetricRecord + + +class MetricBuffer: + """Thread-safe, destructive worker-local metric buffer.""" + + def __init__(self): + self._records: list[MetricRecord] = [] + self._lock = threading.Lock() + + def record(self, record: MetricRecord) -> None: + with self._lock: + self._records.append(record) + + def drain(self) -> list[MetricRecord]: + with self._lock: + records = self._records + self._records = [] + return records diff --git a/src/twinkle/metric/dpo.py b/src/twinkle/metric/dpo.py index 024cb0473..be7ea8c43 100644 --- a/src/twinkle/metric/dpo.py +++ b/src/twinkle/metric/dpo.py @@ -4,6 +4,7 @@ from twinkle.data_format import InputFeature, ModelOutput from twinkle.utils import pad_and_stack_tensors +from twinkle.utils.rl_tensor_utils import align_per_token_values from .base import Metric @@ -33,38 +34,9 @@ def __init__(self, device_mesh, process_group, ignore_index: int = -100, beta: f def _compute_sequence_logps(self, per_token_logps, labels): """Compute sequence-level log probs by summing valid token logps.""" - import torch loss_mask = (labels != self.ignore_index).float() return (per_token_logps * loss_mask).sum(dim=-1) - def _align_logps(self, logps, target_shape, device, dtype): - """Align per-token logps to target shape by padding or truncating. - - Args: - logps: [batch, seq_len] tensor to align - target_shape: Target shape (batch, target_seq_len) - device: Target device - dtype: Target dtype - - Returns: - Aligned tensor with shape matching target_shape - """ - import torch - - if not torch.is_tensor(logps): - logps = torch.as_tensor(logps) - logps = logps.to(device=device, dtype=dtype) - batch_size, src_len = logps.shape - _, target_len = target_shape - - if src_len == target_len: - return logps - elif src_len < target_len: - raise ValueError(f'ref_logps seq_len ({src_len}) < target seq_len ({target_len}). ' - f'This should not happen when both models process the same batch.') - else: - return logps[:, :target_len] - def _split_chosen_rejected(self, tensor): """Split interleaved tensor into chosen and rejected. @@ -121,7 +93,6 @@ def accumulate(self, inputs: Union[InputFeature, List[InputFeature]], outputs: M # Split into chosen and rejected (interleaved format) chosen_logps, rejected_logps = self._split_chosen_rejected(seq_logps) - chosen_labels, rejected_labels = self._split_chosen_rejected(labels) # Accumulate policy logps self.total_chosen_logps += chosen_logps.sum().item() @@ -131,15 +102,18 @@ def accumulate(self, inputs: Union[InputFeature, List[InputFeature]], outputs: M ref_outputs = kwargs.get('ref_outputs') if ref_outputs is not None: ref_logps = ref_outputs.get('logps') - if ref_logps is not None: - if isinstance(ref_logps, list): - if len(ref_logps) == 0: - ref_logps = None - else: - ref_logps = pad_and_stack_tensors(ref_logps) + if isinstance(ref_logps, (list, tuple)) and not ref_logps: + ref_logps = None if ref_logps is not None: # Align ref_logps to match labels shape (handles different seq lengths) - ref_logps = self._align_logps(ref_logps, labels.shape, labels.device, logps.dtype) + ref_logps = align_per_token_values( + ref_logps, + tuple(labels.shape), + device=labels.device, + dtype=logps.dtype, + name='ref_logps', + valid_mask=labels != self.ignore_index, + ) ref_seq_logps = self._compute_sequence_logps(ref_logps, labels) ref_chosen_logps, ref_rejected_logps = self._split_chosen_rejected(ref_seq_logps) diff --git a/src/twinkle/metric/loss.py b/src/twinkle/metric/loss.py index df0ce15c3..90acb647a 100644 --- a/src/twinkle/metric/loss.py +++ b/src/twinkle/metric/loss.py @@ -71,7 +71,7 @@ def calculate(self): self.reset() results = {} if avg_loss is not None: - results['loss'] = f'{avg_loss:.4f}' + results['loss'] = f'{avg_loss:.5f}' if grad_norm > 0: results['grad_norm'] = f'{grad_norm:.6f}' return results diff --git a/src/twinkle/metric/reporting.py b/src/twinkle/metric/reporting.py new file mode 100644 index 000000000..de7f9265e --- /dev/null +++ b/src/twinkle/metric/reporting.py @@ -0,0 +1,556 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Asynchronous JSONL and SwanLab metric reporting.""" + +from __future__ import annotations + +import json +import logging +import math +import os +import re +import threading +import time +from collections import Counter, defaultdict, deque +from collections.abc import Mapping, Sequence +from dataclasses import dataclass, replace +from pathlib import Path +from typing import Any + +from .types import MetricRecord + +logger = logging.getLogger(__name__) + + +def _finite_number(value: Any) -> float | int | None: + if isinstance(value, bool): + return int(value) + if isinstance(value, int): + return value + if isinstance(value, float): + return value if math.isfinite(value) else None + if isinstance(value, str): + match = re.fullmatch(r'\s*tensor\(([-+]?\d+(?:\.\d+)?(?:[eE][-+]?\d+)?)\)\s*', value) + candidate = match.group(1) if match else value + try: + number = float(candidate) + except ValueError: + return None + return number if math.isfinite(number) else None + try: + scalar = value.item() + except (AttributeError, RuntimeError, ValueError): + return None + return _finite_number(scalar) + + +def _safe_name(value: str) -> str: + return re.sub(r'[^A-Za-z0-9_.-]+', '_', value).strip('_') or 'default' + + +@dataclass +class _ScalarSummary: + count: int = 0 + total: float = 0.0 + minimum: float = math.inf + maximum: float = -math.inf + last: float = 0.0 + + def add(self, value: float | int) -> None: + number = float(value) + self.count += 1 + self.total += number + self.minimum = min(self.minimum, number) + self.maximum = max(self.maximum, number) + self.last = number + + def as_dict(self) -> dict[str, float | int]: + return { + 'count': self.count, + 'last': self.last, + 'mean': self.total / self.count, + 'min': self.minimum, + 'max': self.maximum, + } + + +class _SummaryReducer: + + def __init__(self, run_id: str, started_at: float): + self.run_id = run_id + self.started_at = started_at + self.record_counts: Counter[str] = Counter() + self.context_counts: dict[str, Counter[str]] = defaultdict(Counter) + self.context_rollout_groups: Counter[str] = Counter() + self.context_rollout_samples: Counter[str] = Counter() + self.context_trained_samples: Counter[str] = Counter() + self.context_optimizer_steps: dict[str, int] = {} + self.context_policy_versions: dict[str, int] = {} + self.metric_summaries: dict[str, _ScalarSummary] = defaultdict(_ScalarSummary) + self.run_status = 'running' + self.result: dict[str, Any] = {} + + def add(self, record: MetricRecord) -> None: + record_key = f'{record.stage}:{record.status}' + self.record_counts[record_key] += 1 + if record.context_key: + counts = self.context_counts[record.context_key] + counts[record_key] += 1 + if record.optimizer_step is not None: + self.context_optimizer_steps[record.context_key] = record.optimizer_step + if record.policy_version is not None: + previous_version = self.context_policy_versions.get(record.context_key) + self.context_policy_versions[record.context_key] = ( + record.policy_version if previous_version is None else max(previous_version, record.policy_version)) + sample_count = _finite_number(record.values.get('sample_count')) + if sample_count is not None and record.status == 'completed': + if record.stage == 'rollout' and record.attributes.get('scope', 'group') == 'group': + self.context_rollout_groups[record.context_key] += 1 + self.context_rollout_samples[record.context_key] += int(sample_count) + elif record.stage == 'train': + self.context_trained_samples[record.context_key] += int(sample_count) + summarize_values = ( + record.status == 'completed' + and (record.stage != 'rollout' or record.attributes.get('scope', 'group') == 'group')) + if summarize_values: + for name, value in record.values.items(): + number = _finite_number(value) + if number is not None: + self.metric_summaries[f'{record.stage}/{name}'].add(number) + if record.stage == 'run': + self.run_status = record.status + self.result = {**record.values, **record.attributes} + + def as_dict(self, backend_health: Mapping[str, Any]) -> dict[str, Any]: + wall_time = _finite_number(self.result.get('wall_time_s')) + if wall_time is None: + wall_time = time.time() - self.started_at + rollout_groups = sum(self.context_rollout_groups.values()) + train_steps = sum(counts['train:completed'] for counts in self.context_counts.values()) + trained_partitions = sum(counts['partition:completed'] for counts in self.context_counts.values()) + terminal_partitions = _finite_number(self.result.get('trained_partitions')) + if terminal_partitions is not None: + trained_partitions = int(terminal_partitions) + rollout_samples = sum(self.context_rollout_samples.values()) + trained_samples = sum(self.context_trained_samples.values()) + dropped_records = sum(int(item.get('dropped_records', 0)) for item in backend_health.values()) + backend_write_latency_s = sum(float(item.get('write_latency_s', 0.0)) for item in backend_health.values()) + contexts = {} + for context_key, counts in self.context_counts.items(): + contexts[context_key] = { + 'rollout_groups': self.context_rollout_groups[context_key], + 'rollout_samples': self.context_rollout_samples[context_key], + 'train_steps': counts['train:completed'], + 'trained_samples': self.context_trained_samples[context_key], + 'trained_partitions': counts['partition:completed'], + 'optimizer_step': self.context_optimizer_steps.get(context_key), + 'policy_version': self.context_policy_versions.get(context_key), + } + return { + 'run_id': self.run_id, + 'status': self.run_status, + 'wall_time_s': wall_time, + 'record_counts': dict(sorted(self.record_counts.items())), + 'rollout_groups': rollout_groups, + 'rollout_samples': rollout_samples, + 'train_steps': train_steps, + 'trained_samples': trained_samples, + 'trained_partitions': trained_partitions, + 'rollout_groups_per_sec': rollout_groups / wall_time if wall_time > 0 else 0.0, + 'rollout_samples_per_sec': rollout_samples / wall_time if wall_time > 0 else 0.0, + 'train_steps_per_hour': train_steps * 3600 / wall_time if wall_time > 0 else 0.0, + 'trained_samples_per_sec': trained_samples / wall_time if wall_time > 0 else 0.0, + 'train_partitions_per_hour': trained_partitions * 3600 / wall_time if wall_time > 0 else 0.0, + 'dropped_records': dropped_records, + 'backend_write_latency_s': backend_write_latency_s, + 'per_context': contexts, + 'metrics': { + name: summary.as_dict() + for name, summary in sorted(self.metric_summaries.items()) + }, + 'backends': dict(backend_health), + 'result': self.result, + } + + +class _QueuedBackend: + + def __init__( + self, + name: str, + *, + queue_capacity: int, + batch_size: int, + flush_interval_s: float, + ): + if queue_capacity <= 0: + raise ValueError('queue_capacity must be positive') + if batch_size <= 0: + raise ValueError('batch_size must be positive') + if flush_interval_s <= 0: + raise ValueError('flush_interval_s must be positive') + self.name = name + self.queue_capacity = queue_capacity + self.batch_size = batch_size + self.flush_interval_s = flush_interval_s + self._queue: deque[dict[str, Any]] = deque() + self._condition = threading.Condition() + self._closing = False + self._flush_requested = False + self._closed = False + self._disabled = False + self._in_flight = False + self._submitted = 0 + self._written = 0 + self._dropped = 0 + self._write_batches = 0 + self._write_latency_s = 0.0 + self._failures = 0 + self._last_error: str | None = None + self._warning_emitted = False + self._thread = threading.Thread(target=self._run, name=f'twinkle-metrics-{name}', daemon=True) + self._thread.start() + + def submit(self, payload: dict[str, Any]) -> None: + with self._condition: + if self._closing or self._disabled: + self._dropped += 1 + return + if len(self._queue) >= self.queue_capacity: + self._queue.popleft() + self._dropped += 1 + self._queue.append(payload) + self._submitted += 1 + self._condition.notify() + + def flush(self, timeout_s: float | None = None) -> bool: + deadline = None if timeout_s is None else time.monotonic() + timeout_s + with self._condition: + self._flush_requested = True + self._condition.notify_all() + while (self._queue or self._in_flight) and not self._disabled: + remaining = None if deadline is None else deadline - time.monotonic() + if remaining is not None and remaining <= 0: + return False + self._condition.wait(remaining) + self._flush_requested = False + return True + + def close(self, timeout_s: float | None = None) -> bool: + with self._condition: + if self._closed: + return True + self._closing = True + self._condition.notify_all() + self._thread.join(timeout_s) + closed = not self._thread.is_alive() + if closed: + self._closed = True + return closed + + def health(self) -> dict[str, Any]: + with self._condition: + return { + 'enabled': not self._disabled, + 'queue_size': len(self._queue), + 'submitted_records': self._submitted, + 'written_records': self._written, + 'dropped_records': self._dropped, + 'write_batches': self._write_batches, + 'write_latency_s': self._write_latency_s, + 'failure_count': self._failures, + 'last_error': self._last_error, + } + + def _run(self) -> None: + last_write = time.monotonic() + while True: + with self._condition: + while True: + if self._closing and not self._queue: + batch = [] + break + elapsed = time.monotonic() - last_write + should_write = bool(self._queue) and (self._closing or self._flush_requested or len(self._queue) + >= self.batch_size or elapsed >= self.flush_interval_s) + if should_write: + batch = [self._queue.popleft() for _ in range(min(len(self._queue), self.batch_size))] + break + wait_s = (max(0.0, self.flush_interval_s - elapsed) if self._queue else self.flush_interval_s) + self._condition.wait(wait_s) + if not batch: + break + self._in_flight = True + started = time.perf_counter() + try: + self._write_batch(batch) + except Exception as exc: + with self._condition: + self._failures += 1 + self._last_error = f'{type(exc).__name__}: {exc}' + self._dropped += len(batch) + len(self._queue) + self._queue.clear() + self._disabled = True + if not self._warning_emitted: + logger.warning('Metrics backend %s failed and was disabled: %s', self.name, exc) + self._warning_emitted = True + else: + elapsed = time.perf_counter() - started + with self._condition: + self._written += len(batch) + self._write_batches += 1 + self._write_latency_s += elapsed + finally: + last_write = time.monotonic() + with self._condition: + self._in_flight = False + self._condition.notify_all() + if self._disabled: + break + try: + self._close_sink() + except Exception as exc: + with self._condition: + self._failures += 1 + self._last_error = f'{type(exc).__name__}: {exc}' + with self._condition: + self._closed = True + self._condition.notify_all() + + def _write_batch(self, batch: Sequence[dict[str, Any]]) -> None: + raise NotImplementedError + + def _close_sink(self) -> None: + pass + + +class _JSONLBackend(_QueuedBackend): + + def __init__(self, path: str | Path, **kwargs: Any): + self.path = Path(path) + self.path.parent.mkdir(parents=True, exist_ok=True) + self._stream = self.path.open('w', encoding='utf-8') + super().__init__('jsonl', **kwargs) + + def _write_batch(self, batch: Sequence[dict[str, Any]]) -> None: + self._stream.writelines(json.dumps(payload, ensure_ascii=True, default=str) + '\n' for payload in batch) + self._stream.flush() + + def _close_sink(self) -> None: + self._stream.close() + + +class _SwanLabBackend(_QueuedBackend): + + def __init__( + self, + *, + project: str, + experiment_name: str, + log_dir: str | Path, + mode: str, + **kwargs: Any, + ): + import swanlab + + self._swanlab = swanlab + self._swanlab_run = swanlab.init( + project=project, + experiment_name=experiment_name, + logdir=str(log_dir), + mode=mode, + ) + super().__init__('swanlab', **kwargs) + + def _write_batch(self, batch: Sequence[dict[str, Any]]) -> None: + for payload in batch: + prefix = (f'context/{_safe_name(payload["context_key"])}' if payload.get('context_key') else 'global') + stage = _safe_name(payload['stage']) + values = {} + for name, value in payload['values'].items(): + metric_name = str(name) + if metric_name.startswith(f'{stage}/'): + metric_name = metric_name[len(stage) + 1:] + values[f'{prefix}/{stage}/{_safe_name(metric_name)}'] = value + if payload.get('optimizer_step') is not None: + values[f'{prefix}/train/optimizer_step'] = payload['optimizer_step'] + if payload.get('policy_version') is not None: + values[f'{prefix}/policy/version'] = payload['policy_version'] + if payload.get('partition_index') is not None: + values[f'{prefix}/partition/index'] = payload['partition_index'] + if values: + self._swanlab_run.log(values, step=payload['sequence']) + + def _close_sink(self) -> None: + self._swanlab.finish() + + +class MetricsReporter: + """Assign ordering and asynchronously fan metric records out to backends.""" + + def __init__( + self, + *, + run_id: str, + backends: Sequence[_QueuedBackend] = (), + summary_path: str | Path | None = None, + close_timeout_s: float = 10.0, + ): + self.run_id = run_id + self.started_at = time.time() + self.close_timeout_s = close_timeout_s + self.summary_path = Path(summary_path) if summary_path is not None else None + self._backends = tuple(backends) + self._lock = threading.Lock() + self._sequence = 0 + self._closed = False + self._reducer = _SummaryReducer(run_id, self.started_at) + self._initial_errors: dict[str, str] = {} + + def add_backend_error(self, name: str, error: BaseException) -> None: + self._initial_errors[name] = f'{type(error).__name__}: {error}' + + def record(self, record: MetricRecord) -> None: + with self._lock: + if self._closed: + return + self._sequence += 1 + normalized = self._normalize(record, self._sequence) + self._reducer.add(normalized) + payload = self._payload(normalized) + for backend in self._backends: + backend.submit(payload) + + def record_many(self, records: Sequence[MetricRecord]) -> None: + for record in records: + self.record(record) + + def flush(self, timeout_s: float | None = None) -> None: + timeout = self.close_timeout_s if timeout_s is None else timeout_s + deadline = time.monotonic() + timeout + for backend in self._backends: + backend.flush(max(0.0, deadline - time.monotonic())) + + def close(self, timeout_s: float | None = None) -> None: + timeout = self.close_timeout_s if timeout_s is None else timeout_s + with self._lock: + if self._closed: + return + self._closed = True + deadline = time.monotonic() + timeout + for backend in self._backends: + backend.close(max(0.0, deadline - time.monotonic())) + self._write_summary() + + def health(self) -> dict[str, Any]: + backend_health = {backend.name: backend.health() for backend in self._backends} + for name, error in self._initial_errors.items(): + backend_health[name] = { + 'enabled': False, + 'failure_count': 1, + 'last_error': error, + 'dropped_records': self._sequence, + } + return { + 'record_count': self._sequence, + 'dropped_records': sum(int(item.get('dropped_records', 0)) for item in backend_health.values()), + 'backends': backend_health, + } + + def summary(self) -> dict[str, Any]: + return self._reducer.as_dict(self.health()['backends']) + + def _normalize(self, record: MetricRecord, sequence: int) -> MetricRecord: + values: dict[str, float | int] = {} + non_numeric = {} + for name, value in record.values.items(): + number = _finite_number(value) + if number is None: + non_numeric[name] = value + else: + values[name] = number + attributes = dict(record.attributes) + if non_numeric: + attributes['non_numeric_values'] = non_numeric + return replace(record, sequence=sequence, values=values, attributes=attributes) + + def _payload(self, record: MetricRecord) -> dict[str, Any]: + return { + 'timestamp': record.timestamp, + 'elapsed_s': record.timestamp - self.started_at, + 'sequence': record.sequence, + 'run_id': self.run_id, + 'stage': record.stage, + 'context_key': record.context_key, + 'partition_id': record.partition_id, + 'partition_index': record.partition_index, + 'optimizer_step': record.optimizer_step, + 'policy_version': record.policy_version, + 'status': record.status, + 'values': record.values, + 'attributes': record.attributes, + } + + def _write_summary(self) -> None: + if self.summary_path is None: + return + try: + self.summary_path.parent.mkdir(parents=True, exist_ok=True) + temporary_path = self.summary_path.with_suffix(f'{self.summary_path.suffix}.tmp') + with temporary_path.open('w', encoding='utf-8') as stream: + json.dump(self.summary(), stream, ensure_ascii=True, indent=2, default=str) + stream.write('\n') + os.replace(temporary_path, self.summary_path) + except Exception as exc: + logger.warning('Failed to write metrics summary %s: %s', self.summary_path, exc) + + +def create_metrics_reporter(config: Mapping[str, Any] | None, *, run_id: str) -> MetricsReporter | None: + if config is None: + return None + config = dict(config or {}) + if not bool(config.get('enabled', True)): + return None + queue_capacity = int(config.get('queue_capacity', 10000)) + close_timeout_s = float(config.get('close_timeout_s', 10.0)) + jsonl_config = dict(config.get('jsonl') or {}) + swanlab_config = dict(config.get('swanlab') or {}) + backends: list[_QueuedBackend] = [] + backend_errors: dict[str, BaseException] = {} + backend_defaults = { + 'queue_capacity': queue_capacity, + 'batch_size': int(jsonl_config.get('batch_size', 64)), + 'flush_interval_s': float(jsonl_config.get('flush_interval_s', 2.0)), + } + if bool(jsonl_config.get('enabled', True)): + try: + backends.append(_JSONLBackend(jsonl_config['path'], **backend_defaults)) + except Exception as exc: + backend_errors['jsonl'] = exc + logger.warning('JSONL metrics backend could not start: %s', exc) + if bool(swanlab_config.get('enabled', False)) and swanlab_config.get('mode') != 'disabled': + try: + backends.append( + _SwanLabBackend( + project=str(swanlab_config.get('project', 'twinkle-rl')), + experiment_name=str(swanlab_config.get('name', run_id)), + log_dir=swanlab_config.get('log_dir', 'outputs/swanlab'), + mode=str(swanlab_config.get('mode', 'local')), + queue_capacity=queue_capacity, + batch_size=int(swanlab_config.get('batch_size', 16)), + flush_interval_s=float(swanlab_config.get('flush_interval_s', 1.0)), + )) + except Exception as exc: + backend_errors['swanlab'] = exc + logger.warning('SwanLab metrics backend could not start: %s', exc) + summary_path = jsonl_config.get('summary_path') + if summary_path is None and jsonl_config.get('path') is not None: + summary_path = Path(jsonl_config['path']).with_name('summary.json') + reporter = MetricsReporter( + run_id=run_id, + backends=backends, + summary_path=summary_path, + close_timeout_s=close_timeout_s, + ) + for name, error in backend_errors.items(): + reporter.add_backend_error(name, error) + return reporter diff --git a/src/twinkle/metric/types.py b/src/twinkle/metric/types.py new file mode 100644 index 000000000..68741dff7 --- /dev/null +++ b/src/twinkle/metric/types.py @@ -0,0 +1,42 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Transport-neutral metric value types.""" + +from __future__ import annotations + +import time +from dataclasses import dataclass, field +from typing import Any + +METRIC_STAGES = frozenset({ + 'rollout', + 'advantage', + 'train', + 'evaluation', + 'partition', + 'policy', + 'run', +}) +METRIC_STATUSES = frozenset({'submitted', 'completed', 'failed'}) + + +@dataclass(frozen=True) +class MetricRecord: + stage: str + values: dict[str, Any] + timestamp: float = field(default_factory=time.time) + sequence: int | None = None + context_key: str | None = None + partition_id: str | None = None + partition_index: int | None = None + optimizer_step: int | None = None + policy_version: int | None = None + status: str = 'completed' + attributes: dict[str, Any] = field(default_factory=dict) + + def __post_init__(self) -> None: + if self.stage not in METRIC_STAGES: + raise ValueError(f'unsupported metric stage: {self.stage!r}') + if self.status not in METRIC_STATUSES: + raise ValueError(f'unsupported metric status: {self.status!r}') + object.__setattr__(self, 'values', dict(self.values)) + object.__setattr__(self, 'attributes', dict(self.attributes)) diff --git a/src/twinkle/model/__init__.py b/src/twinkle/model/__init__.py index 91401a085..2b367bbd4 100644 --- a/src/twinkle/model/__init__.py +++ b/src/twinkle/model/__init__.py @@ -6,11 +6,13 @@ if TYPE_CHECKING: from .base import TwinkleModel from .megatron import MegatronModel, MultiLoraMegatronModel + from .micro_batch import MicroBatchConfig from .transformers import MultiLoraTransformersModel, TransformersModel, TransformersValueModel else: _import_structure = { 'base': ['TwinkleModel'], + 'micro_batch': ['MicroBatchConfig'], 'transformers': ['TransformersModel', 'MultiLoraTransformersModel', 'TransformersValueModel'], 'megatron': ['MegatronModel', 'MultiLoraMegatronModel'], } diff --git a/src/twinkle/model/micro_batch.py b/src/twinkle/model/micro_batch.py new file mode 100644 index 000000000..f0b47e050 --- /dev/null +++ b/src/twinkle/model/micro_batch.py @@ -0,0 +1,192 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +# Packing algorithms are adapted from AReaL (Apache-2.0). +from __future__ import annotations + +import heapq +import math +from dataclasses import dataclass +from typing import Any, Literal + + +@dataclass(frozen=True) +class MicroBatchConfig: + micro_batch_size: int + dynamic_batching: bool = False + max_tokens_per_micro_batch: int | None = None + packing_algorithm: Literal['ffd', 'kk'] = 'ffd' + + def __post_init__(self): + if self.micro_batch_size <= 0: + raise ValueError(f'micro_batch_size must be positive, got {self.micro_batch_size}') + if self.packing_algorithm not in ('ffd', 'kk'): + raise ValueError(f'packing_algorithm must be ffd or kk, got {self.packing_algorithm!r}') + if self.dynamic_batching and (self.max_tokens_per_micro_batch is None or self.max_tokens_per_micro_batch <= 0): + raise ValueError('max_tokens_per_micro_batch must be positive when dynamic_batching=true') + + @classmethod + def from_kwargs(cls, kwargs: dict[str, Any]) -> MicroBatchConfig | None: + option_names = ( + 'micro_batch_size', + 'dynamic_batching', + 'max_tokens_per_micro_batch', + 'packing_algorithm', + ) + if not any(name in kwargs for name in option_names): + return None + if 'micro_batch_size' not in kwargs: + raise ValueError('micro_batch_size is required when configuring micro-batching') + return cls( + micro_batch_size=int(kwargs.pop('micro_batch_size')), + dynamic_batching=bool(kwargs.pop('dynamic_batching', False)), + max_tokens_per_micro_batch=kwargs.pop('max_tokens_per_micro_batch', None), + packing_algorithm=kwargs.pop('packing_algorithm', 'ffd'), + ) + + +def sequence_length(model_input: dict[str, Any]) -> int: + input_ids = model_input['input_ids'] + return int(input_ids.shape[-1]) if hasattr(input_ids, 'shape') else len(input_ids) + + +def _batch_cost(group: list[int], lengths: list[int], padding_free: bool) -> int: + if not group: + return 0 + values = [lengths[index] for index in group] + return sum(values) if padding_free else max(values) * len(values) + + +def _fits(group: list[int], index: int, lengths: list[int], config: MicroBatchConfig, padding_free: bool) -> bool: + if len(group) >= config.micro_batch_size: + return False + candidate = [*group, index] + return _batch_cost(candidate, lengths, padding_free) <= config.max_tokens_per_micro_batch + + +def _ffd_allocate(lengths: list[int], config: MicroBatchConfig, padding_free: bool, + min_micro_batches: int) -> list[list[int]]: + groups: list[list[int]] = [[] for _ in range(min_micro_batches)] + for index in sorted(range(len(lengths)), key=lengths.__getitem__, reverse=True): + candidates = [ + group_index for group_index, group in enumerate(groups) + if _fits(group, index, lengths, config, padding_free) + ] + if not candidates: + groups.append([index]) + continue + group_index = min( + candidates, + key=lambda candidate: ( + _batch_cost(groups[candidate], lengths, padding_free), + len(groups[candidate]), + ), + ) + groups[group_index].append(index) + return [group for group in groups if group] + + +class _KKSet: + __slots__ = ('total', 'items') + + def __init__(self): + self.total = 0 + self.items: list[int] = [] + + def add(self, index: int, value: int) -> None: + self.items.append(index) + self.total += value + + def merge(self, other: _KKSet) -> None: + self.items.extend(other.items) + self.total += other.total + + def __lt__(self, other: _KKSet) -> bool: + return (self.total, len(self.items), self.items) < (other.total, len(other.items), other.items) + + +class _KKState: + __slots__ = ('sets', ) + + def __init__(self, items: list[tuple[int, int]], group_count: int): + self.sets = [_KKSet() for _ in range(group_count)] + for group, (index, value) in zip(self.sets, items): + group.add(index, value) + self.sets.sort(reverse=True) + + @property + def spread(self) -> int: + return self.sets[0].total - self.sets[-1].total + + def merge(self, other: _KKState) -> None: + for index in range(len(self.sets)): + self.sets[index].merge(other.sets[-1 - index]) + self.sets.sort(reverse=True) + + def __lt__(self, other: _KKState) -> bool: + return self.spread > other.spread + + +def _kk_partition(lengths: list[int], group_count: int) -> list[list[int]]: + queue = [] + for value, index in sorted((value, index) for index, value in enumerate(lengths)): + heapq.heappush(queue, _KKState([(index, value)], group_count)) + while len(queue) > 1: + first = heapq.heappop(queue) + second = heapq.heappop(queue) + first.merge(second) + heapq.heappush(queue, first) + return [group.items for group in queue[0].sets if group.items] + + +def _kk_allocate(lengths: list[int], config: MicroBatchConfig, padding_free: bool, + min_micro_batches: int) -> list[list[int]]: + capacity = config.max_tokens_per_micro_batch + group_count = max(min_micro_batches, math.ceil(sum(lengths) / capacity)) + while group_count <= len(lengths): + groups = _kk_partition(lengths, group_count) + if all( + len(group) <= config.micro_batch_size and _batch_cost(group, lengths, padding_free) <= capacity + for group in groups): + return groups + group_count += 1 + raise ValueError('unable to construct a valid KK micro-batch plan') + + +def plan_micro_batches( + inputs: list[dict[str, Any]], + config: MicroBatchConfig, + *, + padding_free: bool, + min_micro_batches: int = 1, +) -> list[list[int]]: + if not inputs: + raise ValueError('cannot plan micro-batches for empty inputs') + if min_micro_batches <= 0 or min_micro_batches > len(inputs): + raise ValueError(f'invalid min_micro_batches={min_micro_batches} for {len(inputs)} inputs') + if not config.dynamic_batching: + group_count = max(min_micro_batches, math.ceil(len(inputs) / config.micro_batch_size)) + base_size, remainder = divmod(len(inputs), group_count) + groups = [] + start = 0 + for group_index in range(group_count): + size = base_size + int(group_index < remainder) + groups.append(list(range(start, start + size))) + start += size + return groups + lengths = [sequence_length(model_input) for model_input in inputs] + capacity = config.max_tokens_per_micro_batch + oversized = [length for length in lengths if length > capacity] + if oversized: + raise ValueError(f'sequence length {max(oversized)} exceeds max_tokens_per_micro_batch={capacity}') + if config.packing_algorithm == 'ffd': + return _ffd_allocate(lengths, config, padding_free, min_micro_batches) + return _kk_allocate(lengths, config, padding_free, min_micro_batches) + + +def select_batch(value: Any, indices: list[int], batch_size: int) -> Any: + if isinstance(value, list): + return [value[index] for index in indices] if len(value) == batch_size else value + if isinstance(value, tuple): + return tuple(value[index] for index in indices) if len(value) == batch_size else value + if hasattr(value, 'shape') and len(value.shape) > 0 and value.shape[0] == batch_size: + return value[indices] + return value diff --git a/src/twinkle/model/multi_lora.py b/src/twinkle/model/multi_lora.py index f310306c6..ff7766102 100644 --- a/src/twinkle/model/multi_lora.py +++ b/src/twinkle/model/multi_lora.py @@ -201,10 +201,18 @@ def _after(_module): _after(self.module) # self.deactivate_adapter() - def check_length(self, inputs: InputFeature): - total_length = sum(len(_input['input_ids']) for _input in inputs) - if total_length > self.max_length: - raise ValueError(f'Max length exceeds {self.max_length}') + def check_length( + self, + inputs: Union[InputFeature, List[InputFeature]], + ): + if isinstance(inputs, dict): + inputs = [inputs] + for index, item in enumerate(inputs): + if 'input_ids' not in item: + continue + length = len(item['input_ids']) + if length > self.max_length: + raise ValueError(f'Input length {length} exceeds max_length {self.max_length} at sample {index}') def acquire_lora(self, tenant_adapter_name: str, config: LoraConfig) -> str: if self.has_lora(tenant_adapter_name): diff --git a/src/twinkle/model/transformers/multi_lora_transformers.py b/src/twinkle/model/transformers/multi_lora_transformers.py index ea53930de..03f7296b5 100644 --- a/src/twinkle/model/transformers/multi_lora_transformers.py +++ b/src/twinkle/model/transformers/multi_lora_transformers.py @@ -84,6 +84,10 @@ def __init__( self.multi_adapter = MultiLora(max_loras=max_loras, max_r=max_r, max_length=max_length) self.model.gradient_checkpointing_enable() self.model = self.multi_adapter.patch(self.model, target_modules=target_modules, lora_config=self.lora_config) + # PEFT creates LoRA parameters in FP32 even when the base model uses a + # lower-precision dtype. Align the preallocated slots before FSDP2 + # records their parameter and gradient dtypes. + self._ensure_lora_dtype(self.model) self.multi_adapter.save_initial_weights() # Active group for compatibility with single adapter self.active_group = None diff --git a/src/twinkle/model/transformers/strategy/native_fsdp.py b/src/twinkle/model/transformers/strategy/native_fsdp.py index 925c84c82..15c863e7e 100644 --- a/src/twinkle/model/transformers/strategy/native_fsdp.py +++ b/src/twinkle/model/transformers/strategy/native_fsdp.py @@ -116,6 +116,11 @@ def wrap_model(self, model, optimizer=None): if self.device_mesh is None: return model, optimizer fsdp_mesh = _build_fsdp_mesh(self.device_mesh) + if fsdp_mesh is None: + # FSDP has nothing to shard with a single rank, but callers still + # expect the native strategy to place the model on the worker's + # local device before CUDA inputs reach it. + model = model.to(torch.device(Platform.get_local_device())) if fsdp_mesh is not None: ep_enabled = (self.enable_ep and self.ep_fsdp_device_mesh is not None) diff --git a/src/twinkle/model/transformers/transformers.py b/src/twinkle/model/transformers/transformers.py index 19e00022e..d9fe73598 100644 --- a/src/twinkle/model/transformers/transformers.py +++ b/src/twinkle/model/transformers/transformers.py @@ -33,6 +33,7 @@ from twinkle.loss import CrossEntropyLoss, Loss from twinkle.metric import Accuracy, LossMetric, Metric, TrainMetric from twinkle.model.base import TwinkleModel +from twinkle.model.micro_batch import MicroBatchConfig, plan_micro_batches, select_batch from twinkle.model.optimizer_group import BaseOptimizerGroup, TrainStatus from twinkle.model.transformers.moe import apply_expert_parallel from twinkle.model.transformers.strategy import AccelerateStrategy, NativeFSDPStrategy @@ -730,8 +731,13 @@ def backward(self, **kwargs): self.set_grad_scaler(adapter_name=adapter_name) scaler = optimizer_config.scaler - optimizer_config.cur_step += 1 - should_sync = optimizer_config.do_grad_sync(kwargs.get('gradient_accumulation_steps')) + increment_step = kwargs.pop('_increment_step', True) + if increment_step: + optimizer_config.cur_step += 1 + sync_gradients = kwargs.pop('sync_gradients', None) + should_sync = ( + optimizer_config.do_grad_sync(kwargs.get('gradient_accumulation_steps')) + if sync_gradients is None else bool(sync_gradients)) import contextlib no_sync_ctx = contextlib.nullcontext() @@ -755,6 +761,117 @@ def backward(self, **kwargs): optimizer_config.train_status.loss_value = None + def _build_micro_batch_plan(self, inputs, config, optimizer_config): + processor = optimizer_config.processor + assert isinstance(processor, InputProcessor), 'Set a correct `InputProcessor` before forwarding' + optimizer_config._ensure_dp_group() + dp_group = optimizer_config._dp_group + min_micro_batches = 1 + while True: + try: + plan = plan_micro_batches( + inputs, + config, + padding_free=processor.padding_free, + min_micro_batches=min_micro_batches, + ) + planning_error = None + except Exception as exc: + if dp_group is None: + raise + plan = None + planning_error = f'{type(exc).__name__}: {exc}' + + if dp_group is None: + return plan + + local_state = { + 'micro_batch_count': len(plan) if plan is not None else None, + 'input_count': len(inputs), + 'error': planning_error, + } + states = [None] * dist.get_world_size(dp_group) + dist.all_gather_object(states, local_state, group=dp_group) + errors = [ + f'rank {rank}: {state["error"]}' for rank, state in enumerate(states) if state['error'] is not None + ] + if errors: + raise RuntimeError('micro-batch planning failed on one or more model DP ranks: ' + '; '.join(errors)) + + counts = [state['micro_batch_count'] for state in states] + if all(count == len(plan) for count in counts): + return plan + min_micro_batches = max(counts) + if any(min_micro_batches > state['input_count'] for state in states): + raise ValueError('model DP ranks cannot execute the same number of non-empty micro-batches; ' + 'make the input batch divisible by the model data-parallel size') + + def _forward_backward_micro_batch( + self, + *, + inputs, + optimizer_config, + loss_scale, + sync_gradients, + increment_step, + **kwargs, + ): + outputs = self.forward( + inputs=inputs, + router_replay_manual_cleanup=True, + **kwargs, + ) + previous_normalizer = optimizer_config.train_status.num_tokens + loss = self.calculate_loss(**kwargs) + normalizer_delta = optimizer_config.train_status.num_tokens - previous_normalizer + optimizer_config.train_status.loss_value = (optimizer_config.train_status.loss_value * loss_scale) + optimizer_config.train_status.num_tokens = (previous_normalizer + normalizer_delta * loss_scale) + outputs['loss'] = loss * loss_scale + self.backward( + sync_gradients=sync_gradients, + _increment_step=increment_step, + **kwargs, + ) + return outputs + + def _forward_backward_micro_batches( + self, + *, + inputs, + config, + sync_gradients, + loss_scale, + **kwargs, + ): + adapter_name = kwargs.get('adapter_name') + if adapter_name is None: + adapter_name = self._get_default_group() + optimizer_config = self.optimizer_group[adapter_name] + if isinstance(inputs, dict): + inputs = [inputs] + if self._not_encoded(inputs[0]): + assert optimizer_config.template is not None, \ + 'Use set_template to add a template when trying to input `List[Trajectory]`' + inputs = optimizer_config.template.batch_encode(inputs) + + local_batch_size = len(inputs) + plan = self._build_micro_batch_plan(inputs, config, optimizer_config) + outputs = {} + loss_instance = optimizer_config.loss_instance + for micro_batch_index, indices in enumerate(plan): + micro_kwargs = {key: select_batch(value, indices, local_batch_size) for key, value in kwargs.items()} + micro_loss_scale = loss_scale * loss_instance.micro_batch_scale(inputs, indices) + is_last_micro_batch = micro_batch_index == len(plan) - 1 + outputs = self._forward_backward_micro_batch( + inputs=[inputs[index] for index in indices], + optimizer_config=optimizer_config, + loss_scale=micro_loss_scale, + sync_gradients=sync_gradients if is_last_micro_batch else False, + increment_step=is_last_micro_batch, + **micro_kwargs, + ) + return outputs + @remote_function(dispatch='slice_dp', collect=collect_tensor_dict) def forward_backward(self, *, inputs: Union[InputFeature, List[InputFeature], Trajectory, List[Trajectory]], **kwargs): @@ -765,10 +882,26 @@ def forward_backward(self, *, inputs: Union[InputFeature, List[InputFeature], Tr **kwargs: adapter_name: Lora adapter name. gradient_accumulation_steps: Number of gradient accumulation steps. + micro_batch_size: Maximum samples processed per rank in one forward/backward. + dynamic_batching: Pack sequences by token cost instead of fixed sample slices. + max_tokens_per_micro_batch: Per-rank token limit used by dynamic batching. + packing_algorithm: Dynamic packing algorithm. + sync_gradients: Override gradient synchronization on the final micro-batch. + loss_scale: Weight applied to this input batch's loss. Any parameters needed for the specific loss type. Returns: The output of the model forward. """ + micro_batch_config = MicroBatchConfig.from_kwargs(kwargs) + if micro_batch_config is not None: + return self._forward_backward_micro_batches( + inputs=inputs, + config=micro_batch_config, + sync_gradients=kwargs.pop('sync_gradients', None), + loss_scale=float(kwargs.pop('loss_scale', 1.0)), + **kwargs, + ) + outputs = self.forward(inputs=inputs, router_replay_manual_cleanup=True, **kwargs) loss = self.calculate_loss(**kwargs) outputs['loss'] = loss diff --git a/src/twinkle/preprocessor/__init__.py b/src/twinkle/preprocessor/__init__.py index 40d756e3b..58e3a2b05 100644 --- a/src/twinkle/preprocessor/__init__.py +++ b/src/twinkle/preprocessor/__init__.py @@ -2,6 +2,6 @@ from .base import DataFilter, Preprocessor from .dpo import EmojiDPOProcessor from .llm import (AlpacaProcessor, CompetitionMathGRPOProcessor, CompetitionMathProcessor, CountdownProcessor, - GSM8KProcessor, SelfCognitionProcessor) + DAPOMathProcessor, GSM8KProcessor, SelfCognitionProcessor) from .mm import CLEVRProcessor from .olympiad_bench import OlympiadBenchProcessor diff --git a/src/twinkle/preprocessor/llm.py b/src/twinkle/preprocessor/llm.py index 39d3257be..509d6f12a 100644 --- a/src/twinkle/preprocessor/llm.py +++ b/src/twinkle/preprocessor/llm.py @@ -159,3 +159,23 @@ def preprocess(self, row) -> Trajectory: messages=messages, user_data=[('ground_truth', ground_truth)], ) + + +class DAPOMathProcessor(Preprocessor): + """Prepare BytedTsinghua-SIA/DAPO-Math-17k rows for on-policy rollout. + + Required source schema:: + + prompt: list[{"role": str, "content": str}] + reward_model: {"ground_truth": str, ...} + """ + + def __call__(self, rows: Dict[str, List[Any]]) -> Dict[str, List[Any]]: + rows = self.map_col_to_row(rows) + rows = [self.preprocess(row) for row in rows] + return self.map_row_to_col(rows) + + def preprocess(self, row) -> Trajectory: + messages = [Message(role=message['role'], content=message['content']) for message in row['prompt']] + ground_truth = str(row['reward_model']['ground_truth']) + return Trajectory(messages=messages, user_data=[('ground_truth', ground_truth)]) diff --git a/src/twinkle/reward/__init__.py b/src/twinkle/reward/__init__.py index 3ba5babd4..4e8bed086 100644 --- a/src/twinkle/reward/__init__.py +++ b/src/twinkle/reward/__init__.py @@ -1,7 +1,10 @@ # Copyright (c) ModelScope Contributors. All rights reserved. from .base import Reward +from .boxed_math import BoxedMathAccuracyReward +from .dapo_math import DAPOMathAccuracyReward, DAPOMathReward from .format_reward import FormatReward -from .gsm8k import GSM8KAccuracyReward, GSM8KFormatReward +from .gsm8k import (GSM8KAccuracyBrevityReward, GSM8KAccuracyReward, GSM8KBrevityReward, GSM8KFormatReward, + MathVerifyAccuracyReward) from .math_reward import MathReward from .mm_reward import MultiModalAccuracyReward from .olympiad_bench import OlympiadBenchAccuracyReward, OlympiadBenchFormatReward, OlympiadBenchQualityReward diff --git a/src/twinkle/reward/boxed_math.py b/src/twinkle/reward/boxed_math.py new file mode 100644 index 000000000..f9134944f --- /dev/null +++ b/src/twinkle/reward/boxed_math.py @@ -0,0 +1,54 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Accuracy reward for math trajectories with boxed final answers.""" + +from __future__ import annotations + +from decimal import Decimal, InvalidOperation +from typing import Any + +from twinkle.data_format import Trajectory, user_data_get +from twinkle.reward.base import Reward +from twinkle.reward.math_reward import MathReward + + +class BoxedMathAccuracyReward(Reward): + """Compare the final boxed answer with ``user_data.ground_truth``.""" + + def __call__(self, trajectories: list[Trajectory], **kwargs: Any) -> list[float]: + return [self._score(trajectory) for trajectory in trajectories] + + def metric_payload( + self, + trajectories: list[Trajectory], + *, + rewards: list[float], + **kwargs: Any, + ) -> dict[str, float]: + return {'accuracy_reward': sum(rewards) / len(rewards)} + + @classmethod + def _score(cls, trajectory: Trajectory) -> float: + completion = cls._last_assistant_content(trajectory) + if '\\boxed{' not in completion: + return 0.0 + prediction = MathReward.extract_boxed_result(completion).strip() + ground_truth = str(user_data_get(trajectory.get('user_data'), 'ground_truth', '')).strip() + if not prediction or not ground_truth: + return 0.0 + if cls._decimal_equal(prediction, ground_truth): + return 1.0 + return float(MathReward.compare_consecutive(prediction, ground_truth)) + + @staticmethod + def _last_assistant_content(trajectory: Trajectory) -> str: + for message in reversed(trajectory.get('messages', [])): + if message.get('role') == 'assistant': + return str(message.get('content', '')) + return '' + + @staticmethod + def _decimal_equal(first: str, second: str) -> bool: + try: + return Decimal(first.replace(',', '')) == Decimal(second.replace(',', '')) + except InvalidOperation: + return False diff --git a/src/twinkle/reward/dapo_math.py b/src/twinkle/reward/dapo_math.py new file mode 100644 index 000000000..59ecedcfd --- /dev/null +++ b/src/twinkle/reward/dapo_math.py @@ -0,0 +1,150 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Accuracy reward for the Answer-style format used by DAPO-Math.""" + +from __future__ import annotations + +import re +from decimal import Decimal, InvalidOperation +from typing import Any + +from twinkle.data_format import Trajectory, user_data_get +from twinkle.reward.base import Reward +from twinkle.reward.math_reward import MathReward + +_ANSWER_LINE = re.compile(r'^\s*Answer\s*:\s*(.+?)\s*$', re.IGNORECASE | re.MULTILINE) + + +class DAPOMathAccuracyReward(Reward): + """Compare the final ``Answer:`` or boxed result with the ground truth.""" + + def __call__(self, trajectories: list[Trajectory], **kwargs: Any) -> list[float]: + return [self._score(trajectory) for trajectory in trajectories] + + def metric_payload( + self, + trajectories: list[Trajectory], + *, + rewards: list[float], + **kwargs: Any, + ) -> dict[str, float]: + return {'accuracy_reward': sum(rewards) / len(rewards)} + + @classmethod + def _score(cls, trajectory: Trajectory) -> float: + completion = cls._last_assistant_content(trajectory) + prediction = cls.extract_answer(completion) + ground_truth = str(user_data_get(trajectory.get('user_data'), 'ground_truth', '')).strip() + if not prediction or not ground_truth: + return 0.0 + if cls._decimal_equal(prediction, ground_truth): + return 1.0 + return float(MathReward.compare_consecutive(prediction, ground_truth)) + + @staticmethod + def extract_answer(completion: str) -> str: + if '\\boxed{' in completion: + return MathReward.extract_boxed_result(completion).strip() + matches = _ANSWER_LINE.findall(completion) + if not matches: + return '' + answer = matches[-1].strip().rstrip('.') + if len(answer) >= 2 and answer.startswith('$') and answer.endswith('$'): + answer = answer[1:-1].strip() + return answer + + @staticmethod + def _last_assistant_content(trajectory: Trajectory) -> str: + for message in reversed(trajectory.get('messages', [])): + if message.get('role') == 'assistant': + return str(message.get('content', '')) + return '' + + @staticmethod + def _decimal_equal(first: str, second: str) -> bool: + try: + return Decimal(first.replace(',', '')) == Decimal(second.replace(',', '')) + except InvalidOperation: + return False + + +class DAPOMathReward(Reward): + """DAPO math training reward with token-level overlong shaping. + + Accuracy remains a ``0/1`` diagnostic metric, while the optimization score + is ``+1/-1`` with a linear penalty near the response-length limit. + """ + + def __init__( + self, + max_response_length: int, + overlong_buffer_length: int, + overlong_penalty_factor: float = 1.0, + score_tail_chars: int = 300, + ): + if max_response_length <= 0: + raise ValueError('max_response_length must be positive') + if overlong_buffer_length <= 0 or overlong_buffer_length > max_response_length: + raise ValueError('overlong_buffer_length must be in [1, max_response_length]') + if overlong_penalty_factor < 0: + raise ValueError('overlong_penalty_factor must be non-negative') + if score_tail_chars <= 0: + raise ValueError('score_tail_chars must be positive') + self.max_response_length = max_response_length + self.overlong_buffer_length = overlong_buffer_length + self.overlong_penalty_factor = overlong_penalty_factor + self.score_tail_chars = score_tail_chars + + def components(self, trajectories: list[Trajectory]) -> tuple[list[float], list[float]]: + accuracy_rewards = [self._accuracy(trajectory) for trajectory in trajectories] + overlong_rewards = [self._overlong_reward(trajectory) for trajectory in trajectories] + return accuracy_rewards, overlong_rewards + + def __call__(self, trajectories: list[Trajectory], **kwargs: Any) -> list[float]: + accuracy_rewards, overlong_rewards = self.components(trajectories) + return [(1.0 if accuracy else -1.0) + overlong + for accuracy, overlong in zip(accuracy_rewards, overlong_rewards)] + + def metric_payload( + self, + trajectories: list[Trajectory], + *, + rewards: list[float], + **kwargs: Any, + ) -> dict[str, float]: + accuracy_rewards, overlong_rewards = self.components(trajectories) + size = len(trajectories) + if size == 0: + return { + 'total_reward': 0.0, + 'accuracy_reward': 0.0, + 'overlong_reward': 0.0, + 'overlong_ratio': 0.0, + } + return { + 'total_reward': sum(rewards) / size, + 'accuracy_reward': sum(accuracy_rewards) / size, + 'overlong_reward': sum(overlong_rewards) / size, + 'overlong_ratio': sum(value < 0 for value in overlong_rewards) / size, + } + + def _accuracy(self, trajectory: Trajectory) -> float: + completion = DAPOMathAccuracyReward._last_assistant_content(trajectory) + scored_completion = completion[-self.score_tail_chars:] + prediction = DAPOMathAccuracyReward.extract_answer(scored_completion) + ground_truth = str(user_data_get(trajectory.get('user_data'), 'ground_truth', '')).strip() + if not prediction or not ground_truth: + return 0.0 + if DAPOMathAccuracyReward._decimal_equal(prediction, ground_truth): + return 1.0 + return float(MathReward.compare_consecutive(prediction, ground_truth)) + + def _overlong_reward(self, trajectory: Trajectory) -> float: + if 'completion_length' not in trajectory: + raise ValueError('DAPOMathReward requires token-level completion_length on every trajectory') + completion_length = int(trajectory['completion_length']) + expected_length = self.max_response_length - self.overlong_buffer_length + exceed_length = completion_length - expected_length + return min( + -exceed_length / self.overlong_buffer_length * self.overlong_penalty_factor, + 0.0, + ) diff --git a/src/twinkle/reward/gsm8k.py b/src/twinkle/reward/gsm8k.py index 347d49e40..2871e30a7 100644 --- a/src/twinkle/reward/gsm8k.py +++ b/src/twinkle/reward/gsm8k.py @@ -1,6 +1,7 @@ import re from typing import Any, Dict, List +from twinkle.data_format import user_data_get from twinkle.reward.base import Reward @@ -80,6 +81,126 @@ def __call__(self, trajectories: List[Dict[str, Any]], **kwargs) -> List[float]: return rewards +class MathVerifyAccuracyReward(Reward): + """Use the same math-verify parsing and equivalence check as AReaL.""" + + def __init__(self, *, precision: int = 6, try_extract_without_anchor: bool = True): + self.precision = precision + self.try_extract_without_anchor = try_extract_without_anchor + + def __call__(self, trajectories: List[Dict[str, Any]], **kwargs) -> List[float]: + from math_verify.grader import verify + from math_verify.parser import ExprExtractionConfig, LatexExtractionConfig, parse + + extraction_config = ( + ExprExtractionConfig(try_extract_without_anchor=self.try_extract_without_anchor), + LatexExtractionConfig(), + ) + + rewards = [] + for trajectory in trajectories: + completion = '' + for message in reversed(trajectory.get('messages', [])): + if message.get('role') == 'assistant': + completion = str(message.get('content', '')) + break + ground_truth = str(user_data_get(trajectory.get('user_data'), 'ground_truth', '')) + try: + # Disable signal-based timeouts because rollout rewards run on + # the sampler's background event-loop thread. + gold = parse( + ground_truth, + extraction_config=extraction_config, + parsing_timeout=None, + ) + answer = parse( + completion, + extraction_config=extraction_config, + parsing_timeout=None, + ) + if not gold or not answer: + rewards.append(0.0) + continue + correct = verify( + gold, + answer, + float_rounding=self.precision, + timeout_seconds=None, + ) + rewards.append(1.0 if correct else 0.0) + except Exception: + rewards.append(0.0) + return rewards + + def metric_payload( + self, + trajectories: List[Dict[str, Any]], + *, + rewards: List[float], + **kwargs, + ) -> Dict[str, float]: + return {'accuracy_reward': sum(rewards) / len(rewards)} + + +class GSM8KBrevityReward(Reward): + """Reward concise completions that contain a parseable final answer.""" + + def __init__(self, full_reward_length: int = 300, decay_length: int = 3000): + self.full_reward_length = full_reward_length + self.decay_length = decay_length + + def __call__(self, trajectories: List[Dict[str, Any]], **kwargs) -> List[float]: + rewards = [] + for trajectory in trajectories: + completion = '' + for message in reversed(trajectory.get('messages', [])): + if message.get('role') == 'assistant': + completion = message.get('content', '') + break + has_answer = _has_boxed(completion) or bool(re.search(r'####\s*[\-\d,\.]+', completion)) + if not has_answer: + rewards.append(0.0) + continue + excess_length = max(0, len(completion) - self.full_reward_length) + rewards.append(max(0.0, 1.0 - excess_length / self.decay_length)) + return rewards + + +class GSM8KAccuracyBrevityReward(Reward): + """Sum GSM8K answer accuracy and brevity rewards.""" + + def __init__(self, accuracy_weight: float = 1.0, brevity_weight: float = 1.0): + self.accuracy_weight = accuracy_weight + self.brevity_weight = brevity_weight + self.accuracy_reward = GSM8KAccuracyReward() + self.brevity_reward = GSM8KBrevityReward() + + def components(self, trajectories: List[Dict[str, Any]]) -> tuple[list[float], list[float]]: + return self.accuracy_reward(trajectories), self.brevity_reward(trajectories) + + def __call__(self, trajectories: List[Dict[str, Any]], **kwargs) -> List[float]: + accuracy_rewards, brevity_rewards = self.components(trajectories) + return [ + self.accuracy_weight * accuracy + self.brevity_weight * brevity + for accuracy, brevity in zip(accuracy_rewards, brevity_rewards) + ] + + def metric_payload( + self, + trajectories: List[Dict[str, Any]], + *, + rewards: List[float], + **kwargs, + ) -> Dict[str, float]: + accuracy_rewards, brevity_rewards = self.components(trajectories) + size = len(trajectories) + return { + 'total_reward': sum(rewards) / size, + 'accuracy_reward': sum(accuracy_rewards) / size, + 'brevity_reward': sum(brevity_rewards) / size, + } + + class GSM8KFormatReward(Reward): """Format reward: checks if output contains \\boxed{} or #### answer format. diff --git a/src/twinkle/sampler/vllm_sampler/vllm_engine.py b/src/twinkle/sampler/vllm_sampler/vllm_engine.py index b1e1790de..4487a22f1 100644 --- a/src/twinkle/sampler/vllm_sampler/vllm_engine.py +++ b/src/twinkle/sampler/vllm_sampler/vllm_engine.py @@ -1,4 +1,5 @@ # Copyright (c) ModelScope Contributors. All rights reserved. +import asyncio import contextlib import inspect import os @@ -101,6 +102,7 @@ def __init__( self.engine_kwargs = kwargs or {} self._lora_request_cache: Dict[str, Any] = {} + self._lora_load_tasks: Dict[str, asyncio.Task] = {} self._next_lora_id = 1 # Cached LoRARequest for the RL-training synced LoRA. @@ -462,13 +464,26 @@ async def _get_or_load_lora( Returns: ``LoRARequest`` or ``None`` if loading fails. """ - from vllm.lora.request import LoRARequest - - # Fast path: return cached request for this path. if lora_path in self._lora_request_cache: - logger.info(f'Using cached LoRA request for {lora_path}') + logger.debug(f'Using cached LoRA request for {lora_path}') return self._lora_request_cache[lora_path] + load_task = self._lora_load_tasks.get(lora_path) + if load_task is None: + load_task = asyncio.create_task(self._load_lora(lora_path)) + self._lora_load_tasks[lora_path] = load_task + try: + lora_request = await load_task + finally: + if self._lora_load_tasks.get(lora_path) is load_task: + self._lora_load_tasks.pop(lora_path) + if lora_request is not None: + self._lora_request_cache[lora_path] = lora_request + return lora_request + + async def _load_lora(self, lora_path: str): + from vllm.lora.request import LoRARequest + if not os.path.exists(lora_path): logger.error(f'LoRA path does not exist: {lora_path}') return None @@ -487,14 +502,41 @@ async def _get_or_load_lora( lora_path=lora_path, ) + logger.info(f'Loading LoRA from {lora_path}') try: await self.engine.add_lora(lora_request) - self._lora_request_cache[lora_path] = lora_request return lora_request except Exception as e: logger.error(f'Failed to load LoRA from {lora_path}: {e}') return None + async def unload_lora_paths(self, adapter_paths: list[str]) -> None: + """Evict selected LoRA requests without requiring their files to exist.""" + for adapter_path in adapter_paths: + normalized = os.path.abspath(os.path.expanduser(adapter_path)) + request = self._lora_request_cache.pop(normalized, None) + if request is None: + request = self._lora_request_cache.pop(adapter_path, None) + load_task = self._lora_load_tasks.pop(normalized, None) + if load_task is None: + load_task = self._lora_load_tasks.pop(adapter_path, None) + if load_task is not None and not load_task.done(): + load_task.cancel() + await asyncio.gather(load_task, return_exceptions=True) + elif request is None and load_task is not None: + try: + request = load_task.result() + except (asyncio.CancelledError, Exception): + request = None + if request is None: + continue + try: + result = self.engine.remove_lora(request.lora_int_id) + if inspect.isawaitable(result): + await result + except Exception as exc: + logger.warning('Failed to unload LoRA %s: %s', adapter_path, exc) + async def sleep(self, level: int = 2) -> None: """ Offload weights and/or KV cache from GPU memory. diff --git a/src/twinkle/sampler/vllm_sampler/vllm_sampler.py b/src/twinkle/sampler/vllm_sampler/vllm_sampler.py index f12d561ee..82bf94a17 100644 --- a/src/twinkle/sampler/vllm_sampler/vllm_sampler.py +++ b/src/twinkle/sampler/vllm_sampler/vllm_sampler.py @@ -486,6 +486,11 @@ async def _receive_and_load(): self._run_in_loop(_receive_and_load()) + @remote_function(dispatch='all', collect='first', lazy_collect=False) + def unload_adapter_paths(self, adapter_paths: list[str]) -> None: + """Unload policy snapshots from vLLM and clear cached requests.""" + self._run_in_loop(self.engine.unload_lora_paths(adapter_paths)) + @remote_function(dispatch='all', collect='first', lazy_collect=False) def load_full_weights_from_path(self, path: str) -> int: """Load a full (non-LoRA) HF checkpoint into the engine's base model. diff --git a/src/twinkle/server/config/__init__.py b/src/twinkle/server/config/__init__.py index dfdd5176d..d660ec55d 100644 --- a/src/twinkle/server/config/__init__.py +++ b/src/twinkle/server/config/__init__.py @@ -1,13 +1,15 @@ # Copyright (c) ModelScope Contributors. All rights reserved. """Server configuration package — aggregate root and per-deployment specs.""" -from .application_spec import ApplicationSpec, HttpOptions, ModelArgs, ProcessorArgs, SamplerArgs, ServerArgs +from .application_spec import (ApplicationSpec, DataPlaneArgs, HttpOptions, ModelArgs, ProcessorArgs, SamplerArgs, + ServerArgs) from .persistence import PersistenceConfig from .server_config import ServerConfig from .telemetry import TelemetryConfig __all__ = [ 'ApplicationSpec', + 'DataPlaneArgs', 'HttpOptions', 'ModelArgs', 'PersistenceConfig', diff --git a/src/twinkle/server/config/application_spec.py b/src/twinkle/server/config/application_spec.py index 3f8d60cbc..5e6fa89bf 100644 --- a/src/twinkle/server/config/application_spec.py +++ b/src/twinkle/server/config/application_spec.py @@ -1,7 +1,7 @@ # Copyright (c) ModelScope Contributors. All rights reserved. """Per-deployment ``ApplicationSpec`` and typed argument schemas. -Each deployment kind (``server | model | sampler | processor``) carries its +Each deployment kind (``server | model | sampler | processor | data_plane``) carries its own ``args`` block with strict field validation. ``ApplicationSpec`` holds the routing metadata plus the deployment kind and validates ``args`` against the matching ``*Args`` schema in a model validator. @@ -59,6 +59,7 @@ class ModelArgs(_ArgsBase): queue_config: TaskQueueConfig = Field(default_factory=TaskQueueConfig) max_loras: int = 5 max_length: int | None = None + data_plane_url: str | None = None class SamplerArgs(_ArgsBase): @@ -71,9 +72,10 @@ class SamplerArgs(_ArgsBase): nproc_per_node: int = 1 device_group: dict[str, Any] device_mesh: dict[str, Any] - sampler_type: Literal['mock', 'vllm', 'torch'] + sampler_type: Literal['mock', 'vllm', 'vllm_async', 'torch'] engine_args: dict[str, Any] | None = None queue_config: TaskQueueConfig = Field(default_factory=TaskQueueConfig) + data_plane_url: str | None = None class ServerStateArgs(_ArgsBase): @@ -112,11 +114,18 @@ class ProcessorArgs(_ArgsBase): queue_config: TaskQueueConfig = Field(default_factory=TaskQueueConfig) +class DataPlaneArgs(_ArgsBase): + """Args for the TransferQueue-backed client data plane.""" + + config: dict[str, Any] | None = None + + _ARGS_SCHEMA: dict[str, type[_ArgsBase]] = { 'server': ServerArgs, 'model': ModelArgs, 'sampler': SamplerArgs, 'processor': ProcessorArgs, + 'data_plane': DataPlaneArgs, } # ---------- ApplicationSpec ------------------------------------------------ # @@ -138,12 +147,12 @@ class ApplicationSpec(BaseModel): name: str route_prefix: str = '/' - import_path: Literal['server', 'model', 'sampler', 'processor'] + import_path: Literal['server', 'model', 'sampler', 'processor', 'data_plane'] # ``args`` is always populated by the ``mode='before'`` validator below # (which validates the raw block against the schema selected by # ``import_path`` and defaults a missing block to ``{}``), so the field is # required here — the validator runs first and fills it. - args: ServerArgs | ModelArgs | SamplerArgs | ProcessorArgs + args: ServerArgs | ModelArgs | SamplerArgs | ProcessorArgs | DataPlaneArgs deployments: list[dict[str, Any]] = Field(default_factory=list) @model_validator(mode='before') diff --git a/src/twinkle/server/data_plane/__init__.py b/src/twinkle/server/data_plane/__init__.py new file mode 100644 index 000000000..1d11b1ee2 --- /dev/null +++ b/src/twinkle/server/data_plane/__init__.py @@ -0,0 +1,4 @@ +from .app import DataPlaneManagement, build_data_plane_app +from .proxy import DataPlaneProxy + +__all__ = ['DataPlaneManagement', 'DataPlaneProxy', 'build_data_plane_app'] diff --git a/src/twinkle/server/data_plane/app.py b/src/twinkle/server/data_plane/app.py new file mode 100644 index 000000000..271127c4c --- /dev/null +++ b/src/twinkle/server/data_plane/app.py @@ -0,0 +1,44 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from __future__ import annotations + +from fastapi import FastAPI +from typing import Any + +from twinkle.server.deployment import bind_deployment, build_deployment_app +from .store import TQDataRefStore + + +class DataPlaneManagement: + + def __init__(self, config: dict[str, Any] | None = None): + self.store = TQDataRefStore(config) + + +def build_data_plane_app( + deploy_options: dict[str, Any], + config: dict[str, Any] | None = None, +): + from .handlers import register_data_plane_routes + + deploy_options = dict(deploy_options) + autoscaling = deploy_options.get('autoscaling_config') + if autoscaling: + values = autoscaling.model_dump() if hasattr(autoscaling, 'model_dump') else autoscaling + if int(values.get('min_replicas', 1)) != 1 or int(values.get('max_replicas', 1)) != 1: + raise ValueError('data_plane must use exactly one replica') + else: + if int(deploy_options.get('num_replicas', 1)) != 1: + raise ValueError('data_plane must use exactly one replica') + deploy_options.setdefault('num_replicas', 1) + + def register(app: FastAPI, get_self: Any) -> None: + register_data_plane_routes(app, get_self) + + app = build_deployment_app('DataPlane', register) + return bind_deployment( + app, + DataPlaneManagement, + deploy_options, + deployment_name='DataPlaneManagement', + bind_kwargs={'config': config}, + ) diff --git a/src/twinkle/server/data_plane/handlers.py b/src/twinkle/server/data_plane/handlers.py new file mode 100644 index 000000000..aaf214ed7 --- /dev/null +++ b/src/twinkle/server/data_plane/handlers.py @@ -0,0 +1,44 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from __future__ import annotations + +from collections.abc import Callable +from fastapi import Depends, FastAPI +from typing import TYPE_CHECKING + +import twinkle_client.types as types + +if TYPE_CHECKING: + from .app import DataPlaneManagement + + +def register_data_plane_routes(app: FastAPI, self_fn: Callable[[], DataPlaneManagement]) -> None: + + @app.post('/twinkle/put', response_model=types.DataRef) + async def put(body: types.DataPutRequest, self: DataPlaneManagement = Depends(self_fn)) -> types.DataRef: + return await self.store.put( + body.rows, + kind=body.kind, + tags=body.tags, + ) + + @app.post('/twinkle/get', response_model=types.DataRowsResponse) + async def get(body: types.DataGetRequest, self: DataPlaneManagement = Depends(self_fn)) -> types.DataRowsResponse: + rows = await self.store.get( + body.ref, + fields=body.fields, + ) + tags = (await self.store.get_tags(body.ref) if body.include_tags else []) + return types.DataRowsResponse(rows=rows, tags=tags) + + @app.post('/twinkle/append', response_model=types.DataRef) + async def append(body: types.DataAppendRequest, self: DataPlaneManagement = Depends(self_fn)) -> types.DataRef: + return await self.store.append( + body.ref, + body.rows, + tags=body.tags, + ) + + @app.post('/twinkle/release') + async def release(body: types.DataReleaseRequest, self: DataPlaneManagement = Depends(self_fn)) -> dict[str, str]: + await self.store.release(body.ref) + return {'status': 'ok'} diff --git a/src/twinkle/server/data_plane/proxy.py b/src/twinkle/server/data_plane/proxy.py new file mode 100644 index 000000000..0aa2ae7f2 --- /dev/null +++ b/src/twinkle/server/data_plane/proxy.py @@ -0,0 +1,85 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Internal HTTP adapter used by Model and Sampler component deployments.""" +from __future__ import annotations + +import httpx +from typing import Any + +from twinkle_client.http.headers import build_routing_headers +from twinkle_client.types.component import DataRef + + +class DataPlaneProxy: + + def __init__(self, base_url: str | None): + self.base_url = base_url.rstrip('/') if base_url else None + self.client = httpx.AsyncClient(timeout=None) if self.base_url else None + + @property + def enabled(self) -> bool: + return self.client is not None + + async def get( + self, + ref: DataRef, + *, + fields: list[str] | None = None, + ) -> list[dict[str, Any]]: + if self.client is None or self.base_url is None: + raise RuntimeError('data_plane_url is required when a component request uses input_ref') + response = await self.client.post( + f'{self.base_url}/twinkle/get', + json={ + 'ref': ref.model_dump(), + 'fields': fields + }, + headers=build_routing_headers(f'data-ref-{ref.ref_id}'), + ) + response.raise_for_status() + return response.json()['rows'] + + async def put( + self, + rows: list[dict[str, Any]], + *, + kind: str, + tags: list[dict[str, Any]] | None = None, + ) -> DataRef: + if self.client is None or self.base_url is None: + raise RuntimeError('data_plane_url is required to store component output') + response = await self.client.post( + f'{self.base_url}/twinkle/put', + json={ + 'rows': rows, + 'kind': kind, + 'tags': tags + }, + headers=build_routing_headers(f'data-put-{kind}'), + ) + response.raise_for_status() + return DataRef(**response.json()) + + async def append( + self, + ref: DataRef, + rows: list[dict[str, Any]], + *, + tags: list[dict[str, Any]] | None = None, + ) -> DataRef: + if self.client is None or self.base_url is None: + raise RuntimeError('data_plane_url is required to append component output') + response = await self.client.post( + f'{self.base_url}/twinkle/append', + json={ + 'ref': ref.model_dump(), + 'rows': rows, + 'tags': tags, + }, + headers=build_routing_headers(f'data-append-{ref.ref_id}'), + ) + response.raise_for_status() + return DataRef(**response.json()) + + async def close(self) -> None: + if self.client is not None: + await self.client.aclose() diff --git a/src/twinkle/server/data_plane/store.py b/src/twinkle/server/data_plane/store.py new file mode 100644 index 000000000..81fe7fa60 --- /dev/null +++ b/src/twinkle/server/data_plane/store.py @@ -0,0 +1,133 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""TransferQueue KV storage behind opaque client DataRef values.""" +from __future__ import annotations + +import uuid +from typing import Any + +from twinkle.tq_utils import rows_to_tq_fields +from twinkle_client.common.json_utils import json_safe +from twinkle_client.types.component import DataRef + + +def _keys(ref: DataRef) -> list[str]: + return [str(index) for index in range(ref.size)] + + +def _partition(ref: DataRef) -> str: + """Resolve an opaque DataRef to its self-contained physical TQ partition.""" + return f'twinkle-client/{ref.ref_id}' + + +def _input_token_count(rows: list[dict[str, Any]]) -> int: + total = 0 + for row in rows: + input_ids = row.get('input_ids') + if input_ids is None and isinstance(row.get('train_input'), dict): + input_ids = row['train_input'].get('input_ids') + if isinstance(input_ids, (list, tuple)): + total += len(input_ids) + return total + + +def _rows_from_tensordict(data: Any, size: int) -> list[dict[str, Any]]: + rows: list[dict[str, Any]] = [] + fields = list(data.keys()) + for index in range(size): + rows.append({field: json_safe(data[field][index]) for field in fields}) + return rows + + +class TQDataRefStore: + + def __init__(self, config: dict[str, Any] | None = None): + import transfer_queue as tq + if config: + from omegaconf import OmegaConf + tq.init(OmegaConf.create(config)) + else: + tq.init() + + async def put( + self, + rows: list[dict[str, Any]], + *, + kind: str = 'data', + tags: list[dict[str, Any]] | None = None, + ) -> DataRef: + if not rows: + raise ValueError('rows must not be empty') + if tags is not None and len(tags) != len(rows): + raise ValueError(f'tag count {len(tags)} does not match row count {len(rows)}') + import transfer_queue as tq + ref = DataRef( + ref_id=uuid.uuid4().hex, + size=len(rows), + fields=list(rows[0]), + kind=kind, + num_tokens=_input_token_count(rows), + ) + await tq.async_kv_batch_put( + keys=_keys(ref), + partition_id=_partition(ref), + fields=rows_to_tq_fields(rows), + tags=tags, + ) + return ref + + async def get( + self, + ref: DataRef, + *, + fields: list[str] | None = None, + ) -> list[dict[str, Any]]: + import transfer_queue as tq + selected = fields if fields is not None else ref.fields + data = await tq.async_kv_batch_get( + keys=_keys(ref), + partition_id=_partition(ref), + select_fields=selected, + ) + return _rows_from_tensordict(data, ref.size) + + async def get_tags( + self, + ref: DataRef, + ) -> list[dict[str, Any]]: + """Return TQ sample tags in the same order as the rows in ``ref``.""" + import transfer_queue as tq + partition_id = _partition(ref) + partitions = await tq.async_kv_list(partition_id=partition_id) + partition = partitions.get(partition_id, {}) + return [dict(partition.get(key, {})) for key in _keys(ref)] + + async def append( + self, + ref: DataRef, + rows: list[dict[str, Any]], + *, + tags: list[dict[str, Any]] | None = None, + ) -> DataRef: + if len(rows) != ref.size: + raise ValueError(f'append row count {len(rows)} does not match DataRef size {ref.size}') + if not rows: + raise ValueError('rows must not be empty') + if tags is not None and len(tags) != len(rows): + raise ValueError(f'tag count {len(tags)} does not match row count {len(rows)}') + import transfer_queue as tq + await tq.async_kv_batch_put( + keys=_keys(ref), + partition_id=_partition(ref), + fields=rows_to_tq_fields(rows), + tags=tags, + ) + updates: dict[str, Any] = { + 'fields': list(dict.fromkeys([*ref.fields, *rows[0].keys()])), + } + if 'input_ids' in rows[0] or 'train_input' in rows[0]: + updates['num_tokens'] = _input_token_count(rows) + return ref.model_copy(update=updates) + + async def release(self, ref: DataRef) -> None: + import transfer_queue as tq + await tq.async_kv_clear(keys=_keys(ref), partition_id=_partition(ref)) diff --git a/src/twinkle/server/deployment.py b/src/twinkle/server/deployment.py index bb4e3c1cd..ccd700964 100644 --- a/src/twinkle/server/deployment.py +++ b/src/twinkle/server/deployment.py @@ -1,11 +1,11 @@ # Copyright (c) ModelScope Contributors. All rights reserved. """Shared deployment-application construction. -Top-level, central deployment-construction infrastructure shared by all four -deployments (Gateway, Model, Sampler, Processor). It is intentionally NOT under +Top-level deployment-construction infrastructure shared by Gateway, Model, +Sampler, Processor, and DataPlane. It is intentionally NOT under ``utils/`` — it is core to how every deployment is built, not a generic helper. -It consolidates, in one place, the construction logic the four App_Builders +It consolidates, in one place, the construction logic the App_Builders used to repeat: - ``get_servable()`` — the single servable-object accessor; diff --git a/src/twinkle/server/launcher/builder_registry.py b/src/twinkle/server/launcher/builder_registry.py index 5b9eda698..a18be38b1 100644 --- a/src/twinkle/server/launcher/builder_registry.py +++ b/src/twinkle/server/launcher/builder_registry.py @@ -5,8 +5,8 @@ decomposition). No logic change. The operator-facing YAML ``import_path`` literals (``"server"``, ``"model"``, -``"sampler"``, ``"processor"``) are unchanged; only the internal builder -function the ``"server"`` literal resolves to was renamed +``"sampler"``, ``"processor"``, ``"data_plane"``) resolve to internal builder +functions. The function selected by the ``"server"`` literal was renamed (``build_server_app`` → ``build_gateway_app``). """ from __future__ import annotations @@ -19,6 +19,7 @@ 'model': 'build_model_app', 'sampler': 'build_sampler_app', 'processor': 'build_processor_app', + 'data_plane': 'build_data_plane_app', } @@ -28,6 +29,7 @@ def get_builders() -> dict[str, Callable]: Imported lazily so that importing the launcher package does not eagerly pull in every deployment module. """ + from twinkle.server.data_plane import build_data_plane_app from twinkle.server.gateway import build_gateway_app from twinkle.server.model import build_model_app from twinkle.server.processor import build_processor_app @@ -38,6 +40,7 @@ def get_builders() -> dict[str, Callable]: 'build_model_app': build_model_app, 'build_sampler_app': build_sampler_app, 'build_processor_app': build_processor_app, + 'build_data_plane_app': build_data_plane_app, } diff --git a/src/twinkle/server/model/app.py b/src/twinkle/server/model/app.py index 44c1247de..a56eb7ddd 100644 --- a/src/twinkle/server/model/app.py +++ b/src/twinkle/server/model/app.py @@ -97,6 +97,7 @@ def __init__(self, backend: str, adapter_config: dict[str, Any] | None = None, queue_config: TaskQueueConfig | None = None, + data_plane_url: str | None = None, **kwargs): self.backend = backend self.device_group = DeviceGroup(**device_group) @@ -127,6 +128,8 @@ def __init__(self, self.model = MODEL_SELECTOR.construct(backend, ctor_kwargs) self.state: ServerState = get_server_state() + from twinkle.server.data_plane import DataPlaneProxy + self.data_plane = DataPlaneProxy(data_plane_url) self._replica_registered = False # Initialize mixins @@ -173,6 +176,7 @@ async def shutdown(self) -> None: await self.state.unregister_replica(self.replica_id) except Exception: pass + await self.data_plane.close() def check_model_health(self) -> dict: """Probe model actors liveness via a lightweight ping. @@ -245,6 +249,7 @@ def build_model_app(model_id: str, backend: str, adapter_config: dict[str, Any] | None = None, queue_config: TaskQueueConfig | None = None, + data_plane_url: str | None = None, **kwargs): """Build a unified model management application for distributed training. @@ -289,7 +294,8 @@ async def _on_shutdown(servable: Any) -> None: deploy_options, deployment_name='ModelManagement', request_router_config=RequestRouterConfig(request_router_class=StickyLoraRequestRouter), - bind_args=(model_id, nproc_per_node, device_group, device_mesh, backend, adapter_config, queue_config), + bind_args=(model_id, nproc_per_node, device_group, device_mesh, backend, adapter_config, queue_config, + data_plane_url), bind_kwargs=kwargs, ) diff --git a/src/twinkle/server/model/twinkle_handlers.py b/src/twinkle/server/model/twinkle_handlers.py index 5bacdd4b0..582aa5c00 100644 --- a/src/twinkle/server/model/twinkle_handlers.py +++ b/src/twinkle/server/model/twinkle_handlers.py @@ -25,6 +25,8 @@ from twinkle.server.checkpoint import (_resolve_client_save_dir, create_checkpoint_manager, create_training_run_manager, validate_user_path) from twinkle.server.exceptions import FullModeBusyError +from twinkle.server.model.utils import (data_plane_request_shape, merge_forward_kwargs, resolve_data_plane_model_inputs, + select_output_rows) from twinkle.server.utils.validation import get_session_id_from_request from twinkle.utils.logger import get_logger from twinkle_client.common.serialize import deserialize_object @@ -124,6 +126,60 @@ async def _task(): task_type='forward', )) + @app.post('/twinkle/forward_from_data_plane', response_model=types.ForwardResponse) + async def forward_from_data_plane( + request: Request, + body: types.DataPlaneForwardRequest, + self: ModelManagement = Depends(self_fn), + ) -> types.ForwardResponse: + token = await self._on_request_start(request) + adapter_name = _get_twinkle_adapter_name(request, body.adapter_name) + + async def _task(): + self.assert_resource_exists(adapter_name) + raw_inputs, field_kwargs = await resolve_data_plane_model_inputs(body, self.data_plane) + kwargs = merge_forward_kwargs(body.model_extra or {}, field_kwargs) + ret = self.model.forward( + inputs=_parse_inputs(raw_inputs), + adapter_name=adapter_name, + **kwargs, + ) + return {'result': ret} + + input_tokens, batch_size = data_plane_request_shape(body) + return await run_task( + self.schedule_task_and_wait( + _task, + model_id=adapter_name, + token=token, + input_tokens=input_tokens, + batch_size=batch_size, + data_world_size=self.data_world_size, + task_type='forward_from_data_plane', + )) + + @app.post('/twinkle/remove_adapter') + async def remove_adapter( + request: Request, + body: types.AdapterRequest, + self: ModelManagement = Depends(self_fn), + ) -> dict[str, str]: + """Release a drained tenant's in-memory training adapter.""" + token = await self._on_request_start(request) + adapter_name = _get_twinkle_adapter_name(request, body.adapter_name) + + async def _task(): + await self._cleanup_adapter(adapter_name) + return {'status': 'ok'} + + return await run_task( + self.schedule_task_and_wait( + _task, + model_id=adapter_name, + token=token, + task_type='remove_adapter', + )) + @app.post('/twinkle/forward_only', response_model=types.ForwardResponse) async def forward_only( request: Request, @@ -152,6 +208,43 @@ async def _task(): task_type='forward_only', )) + @app.post('/twinkle/forward_only_from_data_plane', response_model=types.ForwardResponse) + async def forward_only_from_data_plane( + request: Request, + body: types.DataPlaneForwardOnlyRequest, + self: ModelManagement = Depends(self_fn), + ) -> types.ForwardResponse: + token = await self._on_request_start(request) + adapter_name = _get_twinkle_adapter_name(request, body.adapter_name) + + async def _task(): + self.assert_resource_exists(adapter_name) + raw_inputs, field_kwargs = await resolve_data_plane_model_inputs(body, self.data_plane) + inputs = _parse_inputs(raw_inputs) + kwargs = merge_forward_kwargs(body.model_extra or {}, field_kwargs) + ret = self.model.forward_only(inputs=inputs, adapter_name=adapter_name, **kwargs) + if body.output_ref is not None: + rows = select_output_rows( + ret, + batch_size=len(inputs), + output_fields=body.output_fields, + ) + output_ref = await self.data_plane.append(body.output_ref, rows) + return {'result': output_ref.model_dump()} + return {'result': ret} + + input_tokens, batch_size = data_plane_request_shape(body) + return await run_task( + self.schedule_task_and_wait( + _task, + model_id=adapter_name, + token=token, + input_tokens=input_tokens, + batch_size=batch_size, + data_world_size=self.data_world_size, + task_type='forward_only_from_data_plane', + )) + @app.post('/twinkle/calculate_loss', response_model=types.CalculateLossResponse) async def calculate_loss( request: Request, @@ -224,6 +317,38 @@ async def _task(): task_type='forward_backward', )) + @app.post('/twinkle/forward_backward_from_data_plane', response_model=types.ForwardBackwardResponse) + async def forward_backward_from_data_plane( + request: Request, + body: types.DataPlaneForwardRequest, + self: ModelManagement = Depends(self_fn), + ) -> types.ForwardBackwardResponse: + token = await self._on_request_start(request) + adapter_name = _get_twinkle_adapter_name(request, body.adapter_name) + + async def _task(): + self.assert_resource_exists(adapter_name) + raw_inputs, field_kwargs = await resolve_data_plane_model_inputs(body, self.data_plane) + kwargs = merge_forward_kwargs(body.model_extra or {}, field_kwargs) + ret = self.model.forward_backward( + inputs=_parse_inputs(raw_inputs), + adapter_name=adapter_name, + **kwargs, + ) + return {'result': ret} + + input_tokens, batch_size = data_plane_request_shape(body) + return await run_task( + self.schedule_task_and_wait( + _task, + model_id=adapter_name, + token=token, + input_tokens=input_tokens, + batch_size=batch_size, + data_world_size=self.data_world_size, + task_type='forward_backward_from_data_plane', + )) + @app.post('/twinkle/clip_grad_norm', response_model=types.ClipGradNormResponse) async def clip_grad_norm( request: Request, diff --git a/src/twinkle/server/model/utils.py b/src/twinkle/server/model/utils.py new file mode 100644 index 000000000..aed4b3c20 --- /dev/null +++ b/src/twinkle/server/model/utils.py @@ -0,0 +1,92 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Utilities shared by DataPlane-backed model routes.""" +from __future__ import annotations + +import asyncio +from typing import Any + +from twinkle_client.common.json_utils import json_safe + + +def model_result_rows(result: Any, batch_size: int) -> list[dict[str, Any]]: + """Keep per-sample model outputs at DataPlane row granularity.""" + if isinstance(result, list) and len(result) == batch_size and all(isinstance(item, dict) for item in result): + return result + if isinstance(result, dict): + batched = {name for name, value in result.items() if isinstance(value, list) and len(value) == batch_size} + if batched: + return [{ + name: value[index] if name in batched else value + for name, value in result.items() + } for index in range(batch_size)] + return [{'result': result}] + + +def value_at_path(value: Any, path: str) -> Any: + for part in path.split('.'): + if not isinstance(value, dict) or part not in value: + raise KeyError(f'data field {path!r} does not exist') + value = value[part] + return value + + +def set_at_path(target: dict[str, Any], path: str, value: Any) -> None: + parts = path.split('.') + current = target + for part in parts[:-1]: + nested = current.setdefault(part, {}) + if not isinstance(nested, dict): + raise ValueError(f'cannot bind nested model argument {path!r}') + current = nested + current[parts[-1]] = value + + +async def resolve_data_plane_model_inputs(body: Any, data_plane: Any) -> tuple[Any, dict[str, Any]]: + """Resolve DataPlane references into model inputs and bound keyword arguments.""" + selected_fields = None + if body.input_field is not None: + selected_fields = list( + dict.fromkeys([ + body.input_field, + *(source.split('.', 1)[0] for source in body.kwarg_fields.values()), + ])) + batches = await asyncio.gather(*(data_plane.get(ref, fields=selected_fields) for ref in body.input_refs)) + rows = [row for batch in batches for row in batch] + if body.input_field is None: + kwarg_roots = {source.split('.', 1)[0] for source in body.kwarg_fields.values()} + inputs = [{key: value for key, value in row.items() if key not in kwarg_roots} for row in rows] + else: + inputs = [value_at_path(row, body.input_field) for row in rows] + + field_kwargs: dict[str, Any] = {} + for target_path, source_path in body.kwarg_fields.items(): + field_value = [value_at_path(row, source_path) for row in rows] + set_at_path(field_kwargs, target_path, field_value) + return inputs, field_kwargs + + +def merge_forward_kwargs(explicit: dict[str, Any], bound: dict[str, Any]) -> dict[str, Any]: + collisions = set(explicit).intersection(bound) + if collisions: + names = ', '.join(sorted(collisions)) + raise ValueError(f'explicit model kwargs conflict with kwarg_fields: {names}') + return {**explicit, **bound} + + +def data_plane_request_shape(body: Any) -> tuple[int, int]: + return ( + sum(ref.num_tokens for ref in body.input_refs), + sum(ref.size for ref in body.input_refs), + ) + + +def select_output_rows( + result: Any, + *, + batch_size: int, + output_fields: dict[str, str], +) -> list[dict[str, Any]]: + rows = model_result_rows(json_safe(result), batch_size) + if len(rows) != batch_size: + raise ValueError(f'model returned {len(rows)} rows for an output_ref of size {batch_size}') + return [{target: value_at_path(row, source) for source, target in output_fields.items()} for row in rows] diff --git a/src/twinkle/server/sampler/app.py b/src/twinkle/server/sampler/app.py index d2cba54e9..8941a40bf 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 @@ -37,6 +38,13 @@ def _make_vllm_sampler(kw: dict[str, Any]) -> Any: return vLLMSampler(**kw) +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 + + return VLLMSamplerTQ(**kw, context_manager=None) + + def _make_torch_sampler(kw: dict[str, Any]) -> Any: from twinkle.sampler import TorchSampler # type: ignore[attr-defined] @@ -49,11 +57,22 @@ def _make_torch_sampler(kw: dict[str, Any]) -> Any: { 'mock': _make_mock_sampler, 'vllm': _make_vllm_sampler, + 'vllm_async': _make_vllm_async_sampler, 'torch': _make_torch_sampler, }, ) +def _construct_sampler_backend( + sampler_type: str, + sampler_kwargs: dict[str, Any], + data_plane_url: str | None, +) -> Any: + # Backend selection is explicit. DataPlane config controls where results + # are stored, not which sampler implementation is instantiated. + return SAMPLER_SELECTOR.construct(sampler_type, sampler_kwargs) + + class SamplerManagement(LazyCleanupMixin, TaskQueueMixin): """Unified sampler management service. @@ -72,6 +91,7 @@ def __init__(self, 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( @@ -101,13 +121,21 @@ def __init__(self, ) else: sampler_kwargs.update(kwargs) - self.sampler = SAMPLER_SELECTOR.construct(sampler_type, sampler_kwargs) + self.sampler = _construct_sampler_backend(sampler_type, sampler_kwargs, data_plane_url) self.state: ServerState = get_server_state() + from twinkle.server.data_plane import DataPlaneProxy + self.data_plane = DataPlaneProxy(data_plane_url) # Initialize task queue mixin self._init_task_queue(queue_config, deployment_name='Sampler') + async def shutdown(self) -> None: + cancel_all = getattr(self.sampler, 'cancel_all_generations', None) + if callable(cancel_all): + await asyncio.to_thread(cancel_all) + await self.data_plane.close() + @serve.multiplexed(max_num_models_per_replica=5) async def _sticky_entry(self, sticky_key: str): return sticky_key @@ -131,6 +159,7 @@ def build_sampler_app(model_id: str, sampler_type: str, engine_args: dict[str, Any] | None = None, queue_config: TaskQueueConfig | None = None, + data_plane_url: str | None = None, **kwargs): """Build a unified sampler application for text generation inference. @@ -143,7 +172,7 @@ def build_sampler_app(model_id: str, device_group: Device group configuration dict device_mesh: Device mesh configuration dict for parallelism deploy_options: Ray Serve deployment options - sampler_type: Sampler selector — ``mock`` | ``vllm`` | ``torch``. + sampler_type: Sampler selector — ``mock`` | ``vllm`` | ``vllm_async`` | ``torch``. Validated up front; bad values raise :class:`ConfigError` before any side effect. engine_args: Additional engine arguments for the sampler @@ -172,6 +201,7 @@ def register_routes(app: FastAPI, get_self: Any) -> None: 'version': '1.0.0', }, attach_replica_id_header=True, + on_shutdown=lambda servable: servable.shutdown(), ) return bind_deployment( @@ -179,7 +209,8 @@ def register_routes(app: FastAPI, get_self: Any) -> None: SamplerManagement, deploy_options, deployment_name='SamplerManagement', - bind_args=(model_id, nproc_per_node, device_group, device_mesh, sampler_type, engine_args, queue_config), + bind_args=(model_id, nproc_per_node, device_group, device_mesh, sampler_type, engine_args, queue_config, + data_plane_url), bind_kwargs=kwargs, ) diff --git a/src/twinkle/server/sampler/backends/mock_sampler.py b/src/twinkle/server/sampler/backends/mock_sampler.py index 735f52639..d8355008b 100644 --- a/src/twinkle/server/sampler/backends/mock_sampler.py +++ b/src/twinkle/server/sampler/backends/mock_sampler.py @@ -84,6 +84,11 @@ def __init__( # ----- Sampler interface --------------------------------------------- # + @remote_function() + def unload_adapter_paths(self, adapter_paths: list[str]) -> None: + """Mirror the production cache-eviction API for control-plane tests.""" + return None + @remote_function() def sample( self, diff --git a/src/twinkle/server/sampler/twinkle_handlers.py b/src/twinkle/server/sampler/twinkle_handlers.py index e49f631dd..708f827f8 100644 --- a/src/twinkle/server/sampler/twinkle_handlers.py +++ b/src/twinkle/server/sampler/twinkle_handlers.py @@ -9,6 +9,7 @@ import asyncio import json import traceback +import uuid from collections.abc import Callable from fastapi import Depends, FastAPI, HTTPException, Request from fastapi.responses import StreamingResponse @@ -27,6 +28,7 @@ from twinkle.server.telemetry.tracing import traced_operation from twinkle.server.utils.validation import get_session_id_from_request from twinkle.utils.logger import get_logger +from twinkle_client.common.json_utils import json_safe logger = get_logger() @@ -57,6 +59,123 @@ def _get_twinkle_sampler_adapter_name(request: Request, adapter_name: str | None return owner_id + '-' + adapter_name +def _build_rollout_rows_and_tags( + sample_models: list[types.SampleResponseModel], + *, + group_ids: list[str] | None, + policy_version: int | None, + adapter_uri: str | None, +) -> tuple[list[dict], list[dict]]: + """Flatten sampler output to one TQ row per generated sequence.""" + resolved_group_ids = group_ids or [uuid.uuid4().hex for _ in sample_models] + if len(resolved_group_ids) != len(sample_models): + raise ValueError(f'group_ids contains {len(resolved_group_ids)} values for ' + f'{len(sample_models)} sampler inputs') + rows = [] + tags = [] + for prompt_index, (response, group_id) in enumerate(zip(sample_models, resolved_group_ids)): + for generation_idx, sequence in enumerate(response.sequences): + sampled_logprobs = [ + 0.0 if not position else float(position[0][1]) for position in (sequence.logprobs or []) + ] + rows.append({ + 'train_input': sequence.new_input_feature, + 'sampled_logprobs': sampled_logprobs, + 'tokens': sequence.tokens, + 'decoded': sequence.decoded, + 'stop_reason': sequence.stop_reason, + 'prompt_logprobs': response.prompt_logprobs, + 'topk_prompt_logprobs': response.topk_prompt_logprobs, + }) + tags.append({ + 'record_type': 'sample', + 'group_id': group_id, + 'prompt_index': prompt_index, + 'generation_idx': generation_idx, + 'rollout_status': 'ROLLOUT_DONE', + 'rollout_policy_version': policy_version, + 'rollout_adapter_uri': adapter_uri, + }) + return rows, tags + + +def _to_sample_response_models(responses) -> list[types.SampleResponseModel]: + """Convert internal sampler responses to the HTTP response schema.""" + sample_models = [] + for response in responses: + sequences = [ + types.SampledSequenceModel( + stop_reason=sequence.stop_reason, + tokens=list(sequence.tokens), + logprobs=list(sequence.logprobs) if sequence.logprobs is not None else None, + decoded=sequence.decoded, + new_input_feature=(_serialize_input_feature(sequence.new_input_feature) + if sequence.new_input_feature is not None else None), + ) for sequence in response.sequences + ] + sample_models.append( + types.SampleResponseModel( + sequences=sequences, + prompt_logprobs=response.prompt_logprobs, + topk_prompt_logprobs=response.topk_prompt_logprobs, + )) + return sample_models + + +def _submission_states(value) -> list[dict]: + """Normalize Twinkle's single-worker unwrapping to a list of states.""" + return value if isinstance(value, list) else [value] + + +async def _await_generation( + sampler, + submission_id: str, +): + """Poll an admitted generation without occupying the sampler admission queue.""" + collected = False + try: + poll_interval = 0.01 + while True: + try: + states = _submission_states(await asyncio.to_thread(sampler.get_generation_status, submission_id)) + except Exception as error: + # A pending read-only actor call can be cancelled by Ray while + # the generation submitted just above remains alive. Treating + # that as a generation failure makes the finally block discard + # otherwise valid rollout work. Retry only Ray's explicit task + # cancellation; actor death and application errors must still + # propagate immediately. + from ray.exceptions import TaskCancelledError + if not isinstance(error, TaskCancelledError): + raise + logger.warning( + 'Generation status poll was cancelled; retrying submission %s', + submission_id, + ) + await asyncio.sleep(poll_interval) + poll_interval = min(poll_interval * 1.5, 0.25) + continue + failed = next( + (state for state in states if state.get('status') not in ('running', 'completed')), + None, + ) + if failed is not None: + error = failed.get('error') or failed.get('status', 'unknown failure') + raise RuntimeError(f'generation {submission_id} failed: {error}') + if states and all(state.get('status') == 'completed' for state in states): + responses = await asyncio.to_thread(sampler.collect_generation, submission_id) + collected = True + return responses + await asyncio.sleep(poll_interval) + poll_interval = min(poll_interval * 1.5, 0.25) + finally: + if not collected: + try: + await asyncio.to_thread(sampler.cancel_generation, submission_id) + except Exception: + logger.warning('Failed to cancel generation %s', submission_id, exc_info=True) + + def _register_twinkle_sampler_routes(app: FastAPI, self_fn: Callable[[], SamplerManagement]) -> None: """Register all /twinkle/* sampler routes on the given FastAPI app. @@ -141,26 +260,7 @@ async def _task(): adapter_name=full_adapter_name, adapter_path=adapter_path, ) - - sample_models = [] - for response in responses: - sequences = [ - types.SampledSequenceModel( - stop_reason=seq.stop_reason, - tokens=list(seq.tokens), - logprobs=list(seq.logprobs) if seq.logprobs is not None else None, - decoded=seq.decoded, - new_input_feature=_serialize_input_feature(seq.new_input_feature) - if seq.new_input_feature is not None else None, - ) for seq in response.sequences - ] - sample_models.append( - types.SampleResponseModel( - sequences=sequences, - prompt_logprobs=response.prompt_logprobs, - topk_prompt_logprobs=response.topk_prompt_logprobs, - )) - return types.SampleResponseModelList(samples=sample_models) + return types.SampleResponseModelList(samples=_to_sample_response_models(responses)) # Calculate metrics for queue scheduling inputs_list = body.inputs if isinstance(body.inputs, list) else [body.inputs] @@ -173,6 +273,98 @@ async def _task(): task_type='sample', )) + @app.post('/twinkle/sample_to_data_plane', response_model=types.DataRef) + async def sample_to_data_plane( + request: Request, + body: types.DataPlaneSampleRequest, + self: SamplerManagement = Depends(self_fn), + ) -> types.DataRef: + """Generate a complete group, store it server-side, and return its DataRef.""" + token = await self._on_request_start(request) + if not self.data_plane.enabled: + raise HTTPException(status_code=503, detail='sample_to_data_plane requires data_plane_url') + if not callable(getattr(self.sampler, 'submit_generation', None)): + raise HTTPException(status_code=503, detail='sampler_type must be vllm_async') + + adapter_path = None + full_adapter_name = _get_twinkle_sampler_adapter_name(request, body.adapter_name) or '' + if body.adapter_uri: + from twinkle.server.checkpoint import create_checkpoint_manager + checkpoint_manager = create_checkpoint_manager(token, client_type='twinkle') + _, adapter_path = checkpoint_manager.parse_adapter_uri(body.adapter_uri) + + inputs = (await self.data_plane.get(body.input_ref) if body.input_ref is not None else body.inputs) + if isinstance(inputs, list) and inputs: + first = inputs[0] + if isinstance(first, dict) and 'input_ids' in first: + inputs = [InputFeature(**item) for item in inputs] + else: + inputs = [Trajectory(**item) for item in inputs] + elif isinstance(inputs, dict): + inputs = [InputFeature(**inputs)] if 'input_ids' in inputs else [Trajectory(**inputs)] + + params_dict = dict(body.sampling_params or {}) + params_dict['num_samples'] = body.num_samples + params = SamplingParams.from_dict(params_dict) + submission_id = uuid.uuid4().hex + + async def _admit(): + await asyncio.to_thread( + self.sampler.submit_generation, + submission_id, + inputs, + params, + adapter_name=full_adapter_name, + adapter_path=adapter_path, + ) + return submission_id + + inline_inputs = body.inputs if isinstance(body.inputs, list) else [body.inputs] + input_tokens = ( + body.input_ref.num_tokens if body.input_ref is not None else sum( + len(item.get('input_ids', [])) for item in inline_inputs if isinstance(item, dict))) + await run_task( + self.schedule_task_and_wait( + _admit, + model_id=full_adapter_name or None, + token=token, + input_tokens=input_tokens, + task_type='sample_admission', + )) + + responses = await _await_generation(self.sampler, submission_id) + rows, tags = _build_rollout_rows_and_tags( + _to_sample_response_models(responses), + group_ids=body.group_ids, + policy_version=body.policy_version, + adapter_uri=body.adapter_uri, + ) + return await self.data_plane.put( + [json_safe(item) for item in rows], + kind='rollout', + tags=tags, + ) + + @app.post('/twinkle/unload_adapter_paths') + async def unload_adapter_paths( + request: Request, + body: types.UnloadAdapterPathsRequest, + self: SamplerManagement = Depends(self_fn), + ) -> dict[str, str]: + """Best-effort eviction of published LoRA snapshots from sampler caches.""" + token = await self._on_request_start(request) + resolved_paths = [] + for adapter_path in body.adapter_paths: + if adapter_path.startswith('twinkle://'): + from twinkle.server.checkpoint import create_checkpoint_manager + checkpoint_manager = create_checkpoint_manager(token, client_type='twinkle') + _, adapter_path = checkpoint_manager.parse_adapter_uri(adapter_path) + resolved_paths.append(adapter_path) + unload = getattr(self.sampler, 'unload_adapter_paths', None) + if unload is not None: + unload(resolved_paths) + return {'status': 'ok'} + @app.post('/twinkle/set_template', response_model=types.SetTemplateResponse) async def set_template( request: Request, diff --git a/src/twinkle/server/utils/task_queue/mixin.py b/src/twinkle/server/utils/task_queue/mixin.py index 7b6f895e0..dcdf805ce 100644 --- a/src/twinkle/server/utils/task_queue/mixin.py +++ b/src/twinkle/server/utils/task_queue/mixin.py @@ -1,10 +1,9 @@ # Copyright (c) ModelScope Contributors. All rights reserved. -""" -TaskQueueMixin: serial compute queue + background-task execution. +"""TaskQueueMixin: serial compute queue plus admitted concurrent tasks. -Two execution paths: - schedule_task() / schedule_task_and_wait() -> serial compute queue (GPU ops) - schedule_background_task() -> fire-and-forget asyncio Task (I/O ops) +``schedule_task`` serializes stateful Model operations. Detached tasks are +used for I/O and for engines such as vLLM that own their compute concurrency +and continuous batching internally. """ from __future__ import annotations @@ -39,8 +38,8 @@ class TaskQueueMixin: Use for GPU operations: forward, backward, step, save, load, etc. 2. Background task (schedule_background_task): - asyncio.create_task, runs concurrently with compute queue. - Use for pure I/O: upload_to_hub, etc. + 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. Requirements @@ -99,6 +98,7 @@ async def _perform_preflight_checks( batch_size: int | None = None, data_world_size: int | None = None, batch_size_multiple: int | None = None, + persist_failure: bool = True, ) -> dict[str, Any] | None: """Run rate-limit and validation checks before queuing a task. @@ -107,69 +107,53 @@ async def _perform_preflight_checks( if not token or not self._task_queue_config.enabled: return None - if input_tokens > self._task_queue_config.max_input_tokens: - error_msg = (f'Input tokens ({input_tokens}) exceed maximum allowed ' - f'({self._task_queue_config.max_input_tokens})') + async def reject(error_msg: str, queue_state: str) -> dict[str, Any]: error_payload = {'error': error_msg, 'category': 'User'} - await self.state.store_future_status( - request_id, - TaskStatus.FAILED.value, - model_id, - result=error_payload, - queue_state=QueueState.UNKNOWN.value, - queue_state_reason=error_msg, - ) - return {'request_id': request_id, 'model_id': model_id} - - if batch_size is not None and data_world_size is not None: - if batch_size < data_world_size: - error_msg = (f'Batch size {batch_size} must be >= data world size {data_world_size}') - error_payload = {'error': error_msg, 'category': 'User'} + if persist_failure: await self.state.store_future_status( request_id, TaskStatus.FAILED.value, model_id, result=error_payload, - queue_state=QueueState.UNKNOWN.value, + queue_state=queue_state, queue_state_reason=error_msg, ) return {'request_id': request_id, 'model_id': model_id} + # Private marker consumed by schedule_task_and_wait(). It is not + # returned by the public polling-style schedule_task() API. + return { + 'request_id': request_id, + 'model_id': model_id, + '_error': error_msg, + } + + if input_tokens > self._task_queue_config.max_input_tokens: + error_msg = (f'Input tokens ({input_tokens}) exceed maximum allowed ' + f'({self._task_queue_config.max_input_tokens})') + return await reject(error_msg, QueueState.UNKNOWN.value) + + if batch_size is not None and data_world_size is not None: + if batch_size < data_world_size: + error_msg = (f'Batch size {batch_size} must be >= data world size {data_world_size}') + return await reject(error_msg, QueueState.UNKNOWN.value) if batch_size_multiple is not None: required_multiple = data_world_size * batch_size_multiple if batch_size % required_multiple != 0: error_msg = (f'Batch size {batch_size} must be divisible by {required_multiple} ' f'so each data-parallel shard gets a multiple of ' f'{batch_size_multiple} examples') - error_payload = {'error': error_msg, 'category': 'User'} - await self.state.store_future_status( - request_id, - TaskStatus.FAILED.value, - model_id, - result=error_payload, - queue_state=QueueState.UNKNOWN.value, - queue_state_reason=error_msg, - ) - return {'request_id': request_id, 'model_id': model_id} + return await reject(error_msg, QueueState.UNKNOWN.value) allowed, reason = await self._rate_limiter.check_and_record(token, input_tokens) if not allowed: if self._task_metrics: self._task_metrics.rate_limit_rejections.inc(tags={'deployment': self._deployment_name}) error_msg = f'Rate limit exceeded: {reason}' - error_payload = {'error': error_msg, 'category': 'User'} - await self.state.store_future_status( - request_id, - TaskStatus.FAILED.value, - model_id, - result=error_payload, - queue_state=QueueState.PAUSED_RATE_LIMIT.value, - queue_state_reason=error_msg, - ) - return {'request_id': request_id, 'model_id': model_id} + return await reject(error_msg, QueueState.PAUSED_RATE_LIMIT.value) return None - async def schedule_task( + async def _schedule_task( self, coro_factory: Callable[[], Coroutine], model_id: str | None = None, @@ -179,37 +163,36 @@ async def schedule_task( data_world_size: int | None = None, batch_size_multiple: int | None = None, task_type: str | None = None, + *, + completion: asyncio.Future[Any] | None = None, + persist_status: bool, ) -> dict[str, Any]: - """Schedule a GPU compute task through the serial compute queue. - - Tasks are processed one at a time in round-robin order across all - per-adapter/per-token queues. Use for any operation that touches GPU - state: forward, backward, step, save, load, add_adapter, etc. - - Args: - coro_factory: Zero-argument callable that creates the coroutine. - model_id: Adapter/model id for queue routing and result association. - token: User token for rate limiting. - input_tokens: Token count for TPS rate limiting. - batch_size: Optional batch size, validated against data_world_size. - data_world_size: Optional data world size for batch validation. - task_type: Label for logging and metrics. - - Returns: - {'request_id': str, 'model_id': str | None} - """ + """Common enqueue path for polling and in-process wait callers.""" request_id = f'req_{uuid.uuid4().hex}' - preflight_result = await self._perform_preflight_checks(request_id, model_id, token, input_tokens, batch_size, - data_world_size, batch_size_multiple) + preflight_result = await self._perform_preflight_checks( + request_id=request_id, + model_id=model_id, + token=token, + input_tokens=input_tokens, + batch_size=batch_size, + data_world_size=data_world_size, + batch_size_multiple=batch_size_multiple, + persist_failure=persist_status, + ) if preflight_result is not None: return preflight_result if self._event_loop is None: self._event_loop = asyncio.get_running_loop() - await self.state.store_future_status( - request_id, TaskStatus.PENDING.value, model_id, queue_state=QueueState.ACTIVE.value) + if persist_status: + await self.state.store_future_status( + request_id, + TaskStatus.PENDING.value, + model_id, + queue_state=QueueState.ACTIVE.value, + ) queue_key = self._queue_key(model_id=model_id, token=token) self._compute_worker.ensure_queue_registered(queue_key) @@ -225,9 +208,16 @@ async def schedule_task( input_tokens=input_tokens, task_type=task_type, created_at=time.monotonic(), + completion=completion, + persist_status=persist_status, )) - await self.state.store_future_status( - request_id, TaskStatus.QUEUED.value, model_id, queue_state=QueueState.ACTIVE.value) + if persist_status: + await self.state.store_future_status( + request_id, + TaskStatus.QUEUED.value, + model_id, + queue_state=QueueState.ACTIVE.value, + ) logger.info(f'[TaskQueue] Task {request_id} queued, type={task_type or "unknown"}, ' f'model_id={model_id}, queue_key={queue_key}, ' f'queue_depth={q.qsize()}, input_tokens={input_tokens}') @@ -240,6 +230,47 @@ async def schedule_task( return {'request_id': request_id, 'model_id': model_id} + async def schedule_task( + self, + coro_factory: Callable[[], Coroutine], + model_id: str | None = None, + token: str | None = None, + input_tokens: int = 0, + batch_size: int | None = None, + data_world_size: int | None = None, + batch_size_multiple: int | None = None, + task_type: str | None = None, + ) -> dict[str, Any]: + """Schedule a GPU compute task through the serial compute queue. + + Tasks are processed one at a time in round-robin order across all + per-adapter/per-token queues. Use for any operation that touches GPU + state: forward, backward, step, save, load, add_adapter, etc. + + Args: + coro_factory: Zero-argument callable that creates the coroutine. + model_id: Adapter/model id for queue routing and result association. + token: User token for rate limiting. + input_tokens: Token count for TPS rate limiting. + batch_size: Optional batch size, validated against data_world_size. + data_world_size: Optional data world size for batch validation. + task_type: Label for logging and metrics. + + Returns: + {'request_id': str, 'model_id': str | None} + """ + return await self._schedule_task( + coro_factory, + model_id=model_id, + token=token, + input_tokens=input_tokens, + batch_size=batch_size, + data_world_size=data_world_size, + batch_size_multiple=batch_size_multiple, + task_type=task_type, + persist_status=True, + ) + async def schedule_task_and_wait( self, coro_factory: Callable[[], Coroutine], @@ -254,12 +285,14 @@ async def schedule_task_and_wait( """Schedule a compute task and block until it completes. Twinkle-side counterpart to schedule_task(). Enqueues the task through - the serial worker, polls until a terminal state, and returns the result. + the same serial worker but delivers the result through an in-process + Future. Large model outputs therefore never enter ServerState. Raises: RuntimeError: If the task fails or scheduling is rejected. """ - future_ref = await self.schedule_task( + completion = asyncio.get_running_loop().create_future() + task_ref = await self._schedule_task( coro_factory, model_id=model_id, token=token, @@ -268,25 +301,13 @@ async def schedule_task_and_wait( data_world_size=data_world_size, batch_size_multiple=batch_size_multiple, task_type=task_type, + completion=completion, + persist_status=False, ) - request_id = future_ref.get('request_id') - if request_id is None: - raise RuntimeError(f'Task scheduling failed: {future_ref}') - - poll_interval = 0.05 - max_poll_interval = 1.0 - while True: - record = await self.state.get_future(request_id) - if record and record.get('status') not in ('pending', 'queued', 'running'): - break - await asyncio.sleep(poll_interval) - poll_interval = min(poll_interval * 2, max_poll_interval) - - if record['status'] == 'failed': - error = record.get('result', {}).get('error', 'Unknown error') + if error := task_ref.get('_error'): + completion.cancel() raise RuntimeError(error) - - return record['result'] + return await completion async def schedule_background_task( self, @@ -294,12 +315,12 @@ async def schedule_background_task( model_id: str | None = None, task_type: str | None = None, ) -> dict[str, Any]: - """Schedule a fire-and-forget background task (bypasses compute queue). + """Schedule a fire-and-forget task outside the serial queue. - Designed for pure I/O operations such as upload_to_hub that do not - require GPU serialization. The task is launched immediately as an - asyncio.create_task so it runs concurrently with the compute queue - without blocking any other user's training operations. + The task is launched immediately as an asyncio task. This is suitable + for pure I/O and for an inference engine such as vLLM that performs its + own safe request concurrency and continuous batching. Stateful Model + operations must continue to use :meth:`schedule_task`. Status is tracked via state.store_future_status so clients can poll progress through the same status endpoints as schedule_task(). @@ -317,7 +338,11 @@ async def schedule_background_task( f'type={task_type or "unknown"}, model_id={model_id}') await self.state.store_future_status( - request_id, TaskStatus.RUNNING.value, model_id, queue_state=QueueState.ACTIVE.value) + request_id, + TaskStatus.RUNNING.value, + model_id, + queue_state=QueueState.ACTIVE.value, + ) async def _run() -> None: try: @@ -327,7 +352,8 @@ async def _run() -> None: TaskStatus.COMPLETED.value, model_id, result=result, - queue_state=QueueState.ACTIVE.value) + queue_state=QueueState.ACTIVE.value, + ) logger.info(f'[TaskQueue] Background task {request_id} completed, type={task_type or "unknown"}') except Exception: error_payload = task_error_payload(traceback.format_exc()) @@ -336,7 +362,8 @@ async def _run() -> None: TaskStatus.FAILED.value, model_id, result=error_payload, - queue_state=QueueState.ACTIVE.value) + queue_state=QueueState.ACTIVE.value, + ) logger.error(f'[TaskQueue] Background task {request_id} FAILED, type={task_type or "unknown"}:\n' f'{traceback.format_exc(limit=3)}') diff --git a/src/twinkle/server/utils/task_queue/types.py b/src/twinkle/server/utils/task_queue/types.py index 4d9a7aa5d..daf8d2bb2 100644 --- a/src/twinkle/server/utils/task_queue/types.py +++ b/src/twinkle/server/utils/task_queue/types.py @@ -9,9 +9,11 @@ """ from __future__ import annotations +import asyncio from collections.abc import Callable, Coroutine from dataclasses import dataclass from enum import Enum +from typing import Any class TaskStatus(Enum): @@ -47,3 +49,9 @@ class QueuedTask: task_type: str | None created_at: float first_rate_limited_at: float | None = None + # ``schedule_task_and_wait`` is an in-process request/response path. Its + # potentially large result is delivered through this Future instead of + # being persisted in ServerState merely for the same process to read it + # back. Polling-style ``schedule_task`` leaves this as ``None``. + completion: asyncio.Future[Any] | None = None + persist_status: bool = True diff --git a/src/twinkle/server/utils/task_queue/worker.py b/src/twinkle/server/utils/task_queue/worker.py index a26888cfe..fdbb36d16 100644 --- a/src/twinkle/server/utils/task_queue/worker.py +++ b/src/twinkle/server/utils/task_queue/worker.py @@ -125,6 +125,16 @@ def _record_queue_metrics(self, task_type: str, queue_wait: float) -> None: # ------------------------------------------------------------------ + @staticmethod + def _complete_result(task: QueuedTask, result: Any) -> None: + if task.completion is not None and not task.completion.done(): + task.completion.set_result(result) + + @staticmethod + def _complete_error(task: QueuedTask, error: str) -> None: + if task.completion is not None and not task.completion.done(): + task.completion.set_exception(RuntimeError(error)) + async def _store_task_failed( self, task: QueuedTask, @@ -133,14 +143,16 @@ async def _store_task_failed( queue_state_reason: str | None = None, ) -> None: """Store FAILED status with a standardised error payload.""" - await self._state.store_future_status( - task.request_id, - TaskStatus.FAILED.value, - task.model_id, - result=task_error_payload(error), - queue_state=queue_state, - queue_state_reason=queue_state_reason, - ) + if task.persist_status: + await self._state.store_future_status( + task.request_id, + TaskStatus.FAILED.value, + task.model_id, + result=task_error_payload(error), + queue_state=queue_state, + queue_state_reason=queue_state_reason, + ) + self._complete_error(task, error) async def fail_queue_tasks(self, queue_key: str, reason: str) -> None: """Drain a queue and mark all pending tasks as FAILED.""" @@ -191,8 +203,13 @@ async def _execute_task(self, task: QueuedTask, queue_key: str, q: asyncio.Queue Handles execution timeout, general exceptions, and always calls q.task_done() in the finally block. """ - await self._state.store_future_status( - task.request_id, TaskStatus.RUNNING.value, task.model_id, queue_state=QueueState.ACTIVE.value) + if task.persist_status: + await self._state.store_future_status( + task.request_id, + TaskStatus.RUNNING.value, + task.model_id, + queue_state=QueueState.ACTIVE.value, + ) task_type = task.task_type or 'unknown' exec_start = time.monotonic() @@ -221,12 +238,15 @@ async def _execute_task(self, task: QueuedTask, queue_key: str, q: asyncio.Queue result = await coro exec_time = time.monotonic() - exec_start logger.info(f'[ComputeWorker] Task {task.request_id} completed in {exec_time:.2f}s, type={task_type}') - await self._state.store_future_status( - task.request_id, - TaskStatus.COMPLETED.value, - task.model_id, - result=result, - queue_state=QueueState.ACTIVE.value) + if task.persist_status: + await self._state.store_future_status( + task.request_id, + TaskStatus.COMPLETED.value, + task.model_id, + result=result, + queue_state=QueueState.ACTIVE.value, + ) + self._complete_result(task, result) except asyncio.TimeoutError: task_status = 'timeout' exec_time = time.monotonic() - exec_start diff --git a/src/twinkle/tq_utils.py b/src/twinkle/tq_utils.py new file mode 100644 index 000000000..f7c34f073 --- /dev/null +++ b/src/twinkle/tq_utils.py @@ -0,0 +1,43 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Small TransferQueue packing helpers shared by both async-RL modes.""" +from __future__ import annotations + +from numbers import Number +from typing import Any + + +def rows_to_tq_fields(rows: list[dict[str, Any]]): + from tensordict import TensorDict + + if not rows: + return TensorDict({}, batch_size=[0]) + field_names = tuple(rows[0].keys()) + expected = set(field_names) + for row_index, row in enumerate(rows): + actual = set(row) + if actual != expected: + missing = sorted(expected - actual) + extra = sorted(actual - expected) + raise ValueError(f'TQ row {row_index} fields mismatch: missing={missing}, extra={extra}') + columns = {field_name: [row[field_name] for row in rows] for field_name in field_names} + return columns_to_tq_fields(columns, len(rows)) + + +def columns_to_tq_fields(columns: dict[str, list[Any]], size: int): + import torch + from tensordict import TensorDict + from tensordict.tensorclass import NonTensorStack + + if size < 0: + raise ValueError(f'TQ field size must be non-negative, got {size}') + packed = {} + for field_name, values in columns.items(): + if not isinstance(values, list): + raise TypeError(f'TQ field {field_name!r} must be a list, got {type(values)!r}') + if len(values) != size: + raise ValueError(f'TQ field {field_name!r} must contain {size} values, got {len(values)}') + if all(isinstance(item, Number) and not isinstance(item, bool) for item in values): + packed[field_name] = torch.tensor(values) + else: + packed[field_name] = NonTensorStack(*values) + return TensorDict(packed, batch_size=[size]) diff --git a/src/twinkle/utils/nccl_safe.py b/src/twinkle/utils/nccl_safe.py index f1e6d4095..3552ce3bb 100644 --- a/src/twinkle/utils/nccl_safe.py +++ b/src/twinkle/utils/nccl_safe.py @@ -94,6 +94,10 @@ def __call__(self, inputs, outputs, **kwargs): type(e).__name__, e, traceback.format_exc()) return _zero_loss(outputs) + def micro_batch_scale(self, inputs, indices): + """Preserve the wrapped loss's micro-batch reduction semantics.""" + return self._loss_instance.micro_batch_scale(inputs, indices) + def _zero_loss(outputs) -> 'LossOutput': """Create a graph-connected zero loss for FSDP compatibility. diff --git a/src/twinkle/utils/rl_tensor_utils.py b/src/twinkle/utils/rl_tensor_utils.py new file mode 100644 index 000000000..5f5eba36f --- /dev/null +++ b/src/twinkle/utils/rl_tensor_utils.py @@ -0,0 +1,96 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Tensor normalization helpers for JSON-compatible RL inputs.""" +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + import torch + + +def align_per_token_values( + values: Any, + target_shape: tuple[int, int], + *, + device: torch.device, + dtype: torch.dtype, + name: str = 'values', + padding_value: float = 0.0, + valid_mask: Any | None = None, +) -> torch.Tensor: + """Convert tensor-like per-token values and align them to a model batch. + + Component HTTP APIs intentionally use JSON-compatible values, so tensors + sent by a client or read through DataPlane arrive as Python lists. The + computation layer calls this helper only for fields it knows are per-token + tensors. Rectangular values are converted directly; ragged rows are padded + before validating and aligning them to ``target_shape``. Missing suffixes + are accepted only when ``valid_mask`` marks those positions as padding. + """ + import torch + + row_lengths: list[int] | None = None + if torch.is_tensor(values): + tensor = values + else: + try: + tensor = torch.as_tensor(values) + except (TypeError, ValueError): + if not isinstance(values, (list, tuple)) or not values: + raise TypeError(f'{name} must be a tensor or a non-empty sequence') from None + + rows = [] + for index, value in enumerate(values): + try: + row = torch.as_tensor(value) + except (TypeError, ValueError) as exc: + raise TypeError(f'{name}[{index}] cannot be converted to a tensor') from exc + if row.ndim == 2 and row.shape[0] == 1: + row = row.squeeze(0) + if row.ndim != 1: + raise ValueError(f'{name}[{index}] must be one-dimensional, got shape {tuple(row.shape)}') + rows.append(row) + row_lengths = [row.numel() for row in rows] + tensor = torch.nn.utils.rnn.pad_sequence( + rows, + batch_first=True, + padding_value=padding_value, + ) + + if tensor.ndim == 1: + tensor = tensor.unsqueeze(0) + if tensor.ndim != 2: + raise ValueError(f'{name} must be two-dimensional, got shape {tuple(tensor.shape)}') + + target_batch_size, target_seq_len = target_shape + batch_size, seq_len = tensor.shape + if batch_size != target_batch_size: + raise ValueError(f'{name} batch size ({batch_size}) does not match target batch size ' + f'({target_batch_size})') + mask = None + if valid_mask is not None: + mask = torch.as_tensor(valid_mask, dtype=torch.bool) + if tuple(mask.shape) != target_shape: + raise ValueError(f'valid_mask shape {tuple(mask.shape)} does not match target shape ' + f'{target_shape}') + + if row_lengths is not None: + for index, row_len in enumerate(row_lengths): + if row_len >= target_seq_len: + continue + if mask is None or bool(mask[index, row_len:].any().item()): + raise ValueError(f'{name}[{index}] has {row_len} tokens but target sequence length ' + f'is {target_seq_len}') + if seq_len < target_seq_len: + if mask is None or bool(mask[:, seq_len:].any().item()): + raise ValueError(f'{name} seq_len ({seq_len}) is smaller than target seq_len ' + f'({target_seq_len})') + tensor = torch.nn.functional.pad( + tensor, + (0, target_seq_len - seq_len), + value=padding_value, + ) + if seq_len > target_seq_len: + tensor = tensor[:, :target_seq_len] + + return tensor.to(device=device, dtype=dtype) diff --git a/src/twinkle_agentic/async_rl/__init__.py b/src/twinkle_agentic/async_rl/__init__.py new file mode 100644 index 000000000..0f79590b0 --- /dev/null +++ b/src/twinkle_agentic/async_rl/__init__.py @@ -0,0 +1,33 @@ +"""Native TransferQueue building blocks for YAML-driven async multi-LoRA RL.""" + +from .context_manager import ContextStatus, LoraContextManager +from .data_plane import TQDataPlane +from .native_tq import ContextGRPOGroupNSampler +from .pipeline import AsyncMultiLoraGRPOConfig, AsyncMultiLoraGRPOPipeline, create_cpu_actor +from .scheduler import ContextSchedulePolicy, ContextScheduler, ScheduleCandidate, SchedulerConfig +from .types import LoraContext, PartitionAdmission, PreparedPartition, PromptGroup, RolloutPolicy +from .vllm_sampler_tq import VLLMSamplerTQ +from .workers import AdvantageWorker, RolloutWorker, TrainerWorker + +__all__ = [ + 'AdvantageWorker', + 'AsyncMultiLoraGRPOConfig', + 'AsyncMultiLoraGRPOPipeline', + 'ContextSchedulePolicy', + 'ContextScheduler', + 'ContextStatus', + 'ContextGRPOGroupNSampler', + 'LoraContext', + 'LoraContextManager', + 'PartitionAdmission', + 'PreparedPartition', + 'PromptGroup', + 'RolloutPolicy', + 'RolloutWorker', + 'ScheduleCandidate', + 'SchedulerConfig', + 'TQDataPlane', + 'TrainerWorker', + 'create_cpu_actor', + 'VLLMSamplerTQ', +] diff --git a/src/twinkle_agentic/async_rl/context_manager.py b/src/twinkle_agentic/async_rl/context_manager.py new file mode 100644 index 000000000..f97e91323 --- /dev/null +++ b/src/twinkle_agentic/async_rl/context_manager.py @@ -0,0 +1,298 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Single-owner control plane for async multi-LoRA RL.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from enum import StrEnum + +from .types import LoraContext, PartitionAdmission, RolloutPolicy + + +class ContextStatus(StrEnum): + ADDING = 'ADDING' + ACTIVE = 'ACTIVE' + DRAINING = 'DRAINING' + EXHAUSTED = 'EXHAUSTED' + FINISHED = 'FINISHED' + REMOVED = 'REMOVED' + FAILED = 'FAILED' + + +@dataclass +class _ContextState: + context: LoraContext + policy: RolloutPolicy + policy_history: list[RolloutPolicy] = field(default_factory=list) + next_step: int = 0 + completed_partitions: int = 0 + status: ContextStatus = ContextStatus.ACTIVE + dataset_exhausted: bool = False + training_partition_id: str | None = None + live_partitions: dict[str, PartitionAdmission] = field(default_factory=dict) + rollout_policy_references: dict[str, int] = field(default_factory=dict) + max_steps: int | None = None + + +class LoraContextManager: + """Ray-safe control plane; it deliberately knows no TQ sample readiness.""" + + def __init__(self, *, max_staleness: int = 0, max_steps: int | None = None): + if max_staleness < 0: + raise ValueError('max_staleness must be non-negative') + if max_steps is not None and max_steps < 0: + raise ValueError('max_steps must be non-negative') + self.max_staleness = int(max_staleness) + self.max_steps = max_steps + self._contexts: dict[str, _ContextState] = {} + self._creation_order = 0 + self._stop_requested = max_steps == 0 + + def register_context(self, + context: LoraContext, + *, + adapter_path: str | None = None, + policy_version: int = 0, + max_steps: int | None = None, + status: ContextStatus = ContextStatus.ACTIVE) -> None: + if context.key in self._contexts: + return + if max_steps is not None and max_steps < 0: + raise ValueError('max_steps must be non-negative') + self._contexts[context.key] = _ContextState( + context=context, + policy=RolloutPolicy(context.key, context.adapter_name, int(policy_version), adapter_path), + policy_history=[RolloutPolicy(context.key, context.adapter_name, int(policy_version), adapter_path)], + status=status, + max_steps=max_steps, + ) + + def reserve_context(self, context: LoraContext, *, max_steps: int | None = None) -> None: + if context.key in self._contexts: + raise KeyError(f'context already exists: {context.key}') + if max_steps is not None and max_steps < 0: + raise ValueError('max_steps must be non-negative') + self.register_context(context, status=ContextStatus.ADDING, max_steps=max_steps) + + def activate_context(self, + context: LoraContext | str, + *, + adapter_path: str | None = None, + policy_version: int = 0) -> None: + state = self._state(context) + if state.status is not ContextStatus.ADDING: + raise RuntimeError(f'{state.context.key} cannot activate from {state.status}') + policy = RolloutPolicy(state.context.key, state.context.adapter_name, int(policy_version), adapter_path) + state.policy = policy + state.policy_history = [policy] + state.status = ContextStatus.ACTIVE + + def request_context_drain(self, context: LoraContext | str) -> None: + state = self._state(context) + if state.status in (ContextStatus.REMOVED, ContextStatus.FAILED): + return + if state.status is ContextStatus.ADDING: + raise RuntimeError(f'{state.context.key} is still being added') + state.status = ContextStatus.DRAINING + + def context_is_drained(self, context: LoraContext | str) -> bool: + return not self._state(context).live_partitions + + def fail_context(self, context: LoraContext | str) -> None: + state = self._state(context) + if state.live_partitions: + raise RuntimeError(f'{state.context.key} still has live partitions') + state.status = ContextStatus.FAILED + + def unregister_context(self, context: LoraContext | str) -> None: + state = self._state(context) + if state.live_partitions: + raise RuntimeError(f'{state.context.key} still has live partitions') + state.status = ContextStatus.REMOVED + self._contexts.pop(state.context.key) + + def context_snapshot(self, context: LoraContext | str) -> dict[str, object]: + state = self._state(context) + return { + 'context': state.context, + 'status': state.status, + 'policy_version': state.policy.version, + 'adapter_path': state.policy.adapter_path, + 'next_step': state.next_step, + 'completed_partitions': state.completed_partitions, + 'live_partitions': len(state.live_partitions), + 'dataset_exhausted': state.dataset_exhausted, + 'max_steps': state.max_steps, + } + + def list_context_snapshots(self) -> list[dict[str, object]]: + return [self.context_snapshot(key) for key in self._contexts] + + def context_adapter_paths(self, context: LoraContext | str) -> list[str]: + return [ + policy.adapter_path for policy in self._state(context).policy_history if policy.adapter_path is not None + ] + + def get_rollout_policy(self, context: LoraContext | str) -> RolloutPolicy: + return self._state(context).policy + + def acquire_rollout_policy(self, context: LoraContext | str) -> RolloutPolicy: + """Pin the current policy while one sampler request is using it.""" + state = self._state(context) + policy = state.policy + if policy.adapter_path is not None: + references = state.rollout_policy_references + references[policy.adapter_path] = references.get(policy.adapter_path, 0) + 1 + return policy + + def release_rollout_policy(self, policy: RolloutPolicy) -> None: + """Release a policy previously returned by :meth:`acquire_rollout_policy`.""" + if policy.adapter_path is None: + return + state = self._state(policy.context_key) + references = state.rollout_policy_references + count = references.get(policy.adapter_path, 0) + if count <= 0: + raise RuntimeError(f'rollout policy was not acquired: {policy.adapter_path}') + if count == 1: + references.pop(policy.adapter_path) + else: + references[policy.adapter_path] = count - 1 + + def request_rollout_partition(self, context: LoraContext | str, *, target_groups: int, + num_generations: int) -> PartitionAdmission | None: + state = self._state(context) + if target_groups <= 0 or num_generations <= 0: + raise ValueError('target_groups and num_generations must be positive') + if state.status is not ContextStatus.ACTIVE or self._stop_requested: + return None + if state.max_steps is not None and state.completed_partitions + len(state.live_partitions) >= state.max_steps: + state.dataset_exhausted = True + state.status = ContextStatus.EXHAUSTED + self._finish_if_drained(state) + return None + if self.max_steps is not None and (self.completed_partitions + len(self.list_live_partitions()) + >= self.max_steps): + return None + oldest_unreleased_step = ( + min(admission.step + for admission in state.live_partitions.values()) if state.live_partitions else state.next_step) + if state.next_step - oldest_unreleased_step > self.max_staleness: + return None + admission = PartitionAdmission( + context=state.context, + partition_id=state.context.partition_id(state.next_step), + step=state.next_step, + target_groups=target_groups, + num_generations=num_generations, + created_order=self._creation_order, + ) + self._creation_order += 1 + state.next_step += 1 + state.live_partitions[admission.partition_id] = admission + return admission + + def list_live_partitions(self) -> list[PartitionAdmission]: + return sorted( + (admission for state in self._contexts.values() for admission in state.live_partitions.values()), + key=lambda admission: admission.created_order, + ) + + def list_trainable_partitions(self) -> list[PartitionAdmission]: + """Return at most one partition per context, in partition step order.""" + partitions = [] + for state in self._contexts.values(): + if not state.live_partitions: + continue + if state.training_partition_id is not None: + partitions.append(state.live_partitions[state.training_partition_id]) + continue + partitions.append(min(state.live_partitions.values(), key=lambda admission: admission.step)) + return sorted(partitions, key=lambda admission: admission.created_order) + + def on_dataset_exhausted(self, context: LoraContext | str) -> None: + state = self._state(context) + state.dataset_exhausted = True + if state.status is ContextStatus.ACTIVE: + state.status = ContextStatus.EXHAUSTED + self._finish_if_drained(state) + + def on_partition_training_started(self, admission: PartitionAdmission) -> None: + state = self._state(admission.context) + if state.training_partition_id not in (None, admission.partition_id): + raise RuntimeError(f'{state.context.key} already trains {state.training_partition_id}') + self._require_live(state, admission) + oldest_partition = min(state.live_partitions.values(), key=lambda candidate: candidate.step) + if oldest_partition.partition_id != admission.partition_id: + raise RuntimeError(f'{admission.partition_id} cannot train before {oldest_partition.partition_id}') + state.training_partition_id = admission.partition_id + + def on_partition_trained(self, admission: PartitionAdmission, *, adapter_path: str) -> RolloutPolicy: + state = self._state(admission.context) + self._require_live(state, admission) + next_policy = RolloutPolicy(state.context.key, state.context.adapter_name, state.policy.version + 1, + adapter_path) + state.policy = next_policy + state.policy_history.append(next_policy) + return next_policy + + def on_partition_cleared(self, admission: PartitionAdmission) -> None: + state = self._state(admission.context) + self._require_live(state, admission) + state.live_partitions.pop(admission.partition_id) + if state.training_partition_id == admission.partition_id: + state.training_partition_id = None + state.completed_partitions += 1 + if state.max_steps is not None and state.completed_partitions >= state.max_steps: + state.dataset_exhausted = True + if state.status is ContextStatus.ACTIVE: + state.status = ContextStatus.EXHAUSTED + if self.max_steps is not None and self.completed_partitions >= self.max_steps: + self._stop_requested = True + self._finish_if_drained(state) + + @property + def completed_partitions(self) -> int: + return sum(state.completed_partitions for state in self._contexts.values()) + + def get_completed_partitions(self) -> int: + return self.completed_partitions + + def is_run_finished(self) -> bool: + if self._stop_requested: + return not self.list_live_partitions() + return bool(self._contexts) and all(state.status is ContextStatus.FINISHED and not state.live_partitions + for state in self._contexts.values()) + + def is_rollout_admission_closed(self) -> bool: + """Whether the producer must stop reading new prompts. + + A global train limit closes admission only. Existing live partitions + remain available to AdvantageWorker and TrainerWorker until drained. + """ + return self._stop_requested or all(state.status is not ContextStatus.ACTIVE + for state in self._contexts.values()) + + def adapter_paths_to_keep(self) -> set[str]: + paths: set[str] = set() + for state in self._contexts.values(): + if state.policy.adapter_path is not None: + paths.add(state.policy.adapter_path) + paths.update(state.rollout_policy_references) + return paths + + def context_status(self, context: LoraContext | str) -> ContextStatus: + return self._state(context).status + + def _finish_if_drained(self, state: _ContextState) -> None: + if state.dataset_exhausted and not state.live_partitions and state.status is ContextStatus.EXHAUSTED: + state.status = ContextStatus.FINISHED + + def _state(self, context: LoraContext | str) -> _ContextState: + key = context if isinstance(context, str) else context.key + return self._contexts[key] + + @staticmethod + def _require_live(state: _ContextState, admission: PartitionAdmission) -> None: + if state.live_partitions.get(admission.partition_id) != admission: + raise KeyError(f'unknown live partition {admission.partition_id}') diff --git a/src/twinkle_agentic/async_rl/data_plane.py b/src/twinkle_agentic/async_rl/data_plane.py new file mode 100644 index 000000000..641947366 --- /dev/null +++ b/src/twinkle_agentic/async_rl/data_plane.py @@ -0,0 +1,274 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""The only async-RL layer that speaks native TransferQueue BatchMeta.""" + +from __future__ import annotations + +from typing import Any, Sequence + +from .native_tq import (AsyncTQClient, append_fields, batch_size_for_groups, clear_partition, fetch_ready_batch, + metadata_size, preallocate_partition, set_sample_tags, split_batch_meta) +from .tq_utils import REQUIRED_MODEL_INPUT_FIELDS, ROLLOUT_TRAIN_FIELDS, columns_to_tq_fields, rows_to_tq_fields +from .types import ClaimedBatch, LoraContext, PartitionAdmission, PreparedPartition, PromptGroup, RolloutOutput + +_REQUIRED_ROLLOUT_FIELDS = frozenset((*REQUIRED_MODEL_INPUT_FIELDS, 'logprobs', 'rewards')) + + +def build_rollout_group_sample_write( + group: PromptGroup, + samples: Sequence[RolloutOutput], + *, + rewards: list[float] | None = None, + expected_num_generations: int, +) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]: + group_samples = [dict(sample) for sample in samples] + if expected_num_generations <= 0: + raise ValueError(f'expected_num_generations must be positive, got {expected_num_generations}') + if len(group_samples) != expected_num_generations: + raise ValueError(f'group {group.group_id} expected {expected_num_generations} rollout samples, ' + f'got {len(group_samples)}') + if rewards is not None and len(rewards) != len(group_samples): + raise ValueError(f'reward count {len(rewards)} does not match sample count {len(group_samples)}') + + sample_fields: list[dict[str, Any]] = [] + sample_tags: list[dict[str, Any]] = [] + generation_indices: list[int] = [] + reward_iter = iter(rewards or []) + for sample_index, trajectory in enumerate(group_samples): + sample = dict(trajectory) + if rewards is not None: + sample['rewards'] = float(next(reward_iter)) + generation_idx = int(sample.get('generation_idx', sample_index)) + sample_key = f'samples/{group.group_id}/{generation_idx}' + generation_indices.append(generation_idx) + logprobs = _require_rollout_logprobs(sample, sample_key=sample_key) + sample['logprobs'] = logprobs + sample_fields.append(_rollout_sample_fields(sample)) + sample_tags.append( + _sample_tag( + context=group.context, + group=group, + sample=sample, + sample_key=sample_key, + generation_idx=generation_idx, + logprobs=logprobs, + )) + + expected_indices = list(range(expected_num_generations)) + if generation_indices != expected_indices: + raise ValueError(f'group {group.group_id} generation_idx must be 0..{expected_num_generations - 1} ' + f'in order, got {generation_indices}') + return sample_fields, sample_tags + + +def _require_rollout_logprobs(sample: dict[str, Any], *, sample_key: str) -> list[float]: + logprobs = sample.get('logprobs') + if not isinstance(logprobs, list): + raise TypeError(f'rollout sample {sample_key!r} logprobs must be list[float], got {type(logprobs)!r}') + values: list[float] = [] + for index, value in enumerate(logprobs): + if isinstance(value, bool) or not isinstance(value, (int, float)): + raise TypeError(f'rollout sample {sample_key!r} logprobs[{index}] must be a float, got {type(value)!r}') + values.append(float(value)) + labels = sample.get('labels') + if labels is not None: + trainable_tokens = sum(1 for label in labels if label != -100) + if len(values) != trainable_tokens: + raise ValueError(f'rollout sample {sample_key!r} logprobs length must match trainable labels: ' + f'{len(values)} != {trainable_tokens}') + return values + + +def _rollout_sample_fields(sample: dict[str, Any]) -> dict[str, Any]: + return {field_name: sample[field_name] for field_name in ROLLOUT_TRAIN_FIELDS if field_name in sample} + + +def _sample_tag( + *, + context: LoraContext, + group: PromptGroup, + sample: dict[str, Any], + sample_key: str, + generation_idx: int, + logprobs: list[float], +) -> dict[str, Any]: + tag = { + 'record_type': 'sample', + 'sample_status': 'success', + 'context_key': context.key, + 'tenant_id': context.tenant_id, + 'training_run_id': context.training_run_id, + 'adapter_name': context.adapter_name, + 'sample_id': sample.get('sample_id', sample_key), + 'group_id': group.group_id, + 'generation_idx': generation_idx, + 'rollout_policy_version': int(sample['rollout_policy_version']), + 'rollout_adapter_path': sample.get('rollout_adapter_path'), + 'logprobs_length': len(logprobs), + } + for field_name in ('rollout_policy_versions', 'initial_policy_version', 'final_policy_version', + 'policy_version_span'): + if field_name in sample: + tag[field_name] = sample[field_name] + trainable_tokens = _trainable_token_count(sample.get('labels')) + if trainable_tokens is not None: + tag['trainable_tokens'] = trainable_tokens + for sample_field, tag_field in ( + ('input_ids', 'input_length'), + ('labels', 'label_length'), + ('attention_mask', 'attention_length'), + ): + length = _safe_len(sample.get(sample_field)) + if length is not None: + tag[tag_field] = length + for field_name in ('stop_reason', 'truncated', 'turns'): + if field_name in sample: + tag[field_name] = sample[field_name] + if 'completion_length' in sample: + tag['completion_length'] = int(sample['completion_length']) + return tag + + +def _trainable_token_count(labels: Any) -> int | None: + if labels is None: + return None + return sum(1 for label in labels if label != -100) + + +def _safe_len(value: Any) -> int | None: + if value is None: + return None + try: + return len(value) + except TypeError: + return None + + +class TQDataPlane: + """Maps partition-level training operations to native TQ operations.""" + + def __init__(self, client: AsyncTQClient | None = None): + self._client = client + + @property + def client(self) -> AsyncTQClient: + if self._client is None: + import transfer_queue as tq + tq.init() + self._client = tq.get_client() + return self._client + + async def prepare_rollout_partition( + self, + admission: PartitionAdmission, + prompts: Sequence[dict[str, Any]], + sampling_params: Any, + ) -> PreparedPartition: + if len(prompts) != admission.target_groups: + raise ValueError(f'{admission.partition_id} expected {admission.target_groups} prompts, got {len(prompts)}') + rows = [dict(prompt) for prompt in prompts for _ in range(admission.num_generations)] + metadata = await preallocate_partition( + self.client, partition_id=admission.partition_id, prompt_fields=rows_to_tq_fields(rows)) + group_batch_metas = split_batch_meta(metadata, admission.num_generations) + groups = [] + for index, (prompt, batch_meta) in enumerate(zip(prompts, group_batch_metas)): + group_id = f'{admission.partition_id}/group_{index}' + await set_sample_tags(self.client, batch_meta, [{ + 'group_id': group_id, + 'generation_idx': generation_idx, + 'rollout_status': 'PENDING', + } for generation_idx in range(admission.num_generations)]) + groups.append( + PromptGroup( + context=admission.context, + partition=admission, + group_id=group_id, + prompt=dict(prompt), + batch_meta=batch_meta, + )) + return PreparedPartition(admission, tuple(groups), sampling_params) + + async def complete_rollout_group( + self, + group: PromptGroup, + *, + rollout_rows: Sequence[RolloutOutput], + rewards: Sequence[float], + submission_id: str, + tag_metrics: dict[str, Any] | None = None, + ) -> None: + expected = group.partition.num_generations + sample_fields, sample_tags = build_rollout_group_sample_write( + group, + rollout_rows, + rewards=list(rewards), + expected_num_generations=expected, + ) + for index, fields in enumerate(sample_fields): + missing = sorted(_REQUIRED_ROLLOUT_FIELDS - set(fields)) + if missing: + raise ValueError(f'rollout sample {group.group_id}/{index} is missing training fields {missing}') + metrics = dict(tag_metrics or {}) + completed_tags = [] + for tag in sample_tags: + completed_tag = dict(tag) + completed_tag.update(metrics) + completed_tag.update({'rollout_status': 'ROLLOUT_DONE', 'submission_id': submission_id}) + completed_tags.append(completed_tag) + await set_sample_tags(self.client, group.batch_meta, completed_tags) + await append_fields(self.client, rows_to_tq_fields(sample_fields), group.batch_meta) + + async def claim_advantage_batch(self, admission: PartitionAdmission, group_count: int) -> ClaimedBatch | None: + metadata = await self._claim(admission, group_count, ['input_ids', 'logprobs', 'rewards'], + self._advantage_task(admission)) + if metadata is None: + return None + return ClaimedBatch( + admission=admission, + data=await self.client.async_get_data(metadata.select_fields(['rewards'])), + batch_meta=metadata, + ) + + async def write_advantages(self, batch: ClaimedBatch, *, advantages: Any, returns: Any) -> None: + size = metadata_size(batch.batch_meta) + fields = columns_to_tq_fields({'advantages': list(advantages), 'returns': list(returns)}, size) + await append_fields(self.client, fields, batch.batch_meta) + + async def claim_training_batch(self, admission: PartitionAdmission, group_count: int) -> ClaimedBatch | None: + metadata = await self._claim( + admission, + group_count, + [*REQUIRED_MODEL_INPUT_FIELDS, 'logprobs', 'rewards', 'advantages', 'returns'], + self._trainer_task(admission), + ) + if metadata is None: + return None + return ClaimedBatch( + admission=admission, + data=await self.client.async_get_data(metadata), + batch_meta=metadata, + sample_tags=tuple(metadata.get_all_custom_meta()), + ) + + async def is_training_consumed(self, admission: PartitionAdmission) -> bool: + return await self.client.async_check_consumption_status(self._trainer_task(admission), admission.partition_id) + + async def clear_partition(self, admission: PartitionAdmission) -> None: + await clear_partition(self.client, admission.partition_id) + + async def _claim(self, admission: PartitionAdmission, groups: int, fields: list[str], task: str) -> Any | None: + return await fetch_ready_batch( + self.client, + data_fields=fields, + batch_size=batch_size_for_groups(groups, admission.num_generations), + partition_id=admission.partition_id, + task_name=task, + num_generations=admission.num_generations, + ) + + @staticmethod + def _advantage_task(admission: PartitionAdmission) -> str: + return f'async_rl/advantage/{admission.context.key}' + + @staticmethod + def _trainer_task(admission: PartitionAdmission) -> str: + return f'async_rl/trainer/{admission.context.key}' diff --git a/src/twinkle_agentic/async_rl/metrics.py b/src/twinkle_agentic/async_rl/metrics.py new file mode 100644 index 000000000..d78ac06c9 --- /dev/null +++ b/src/twinkle_agentic/async_rl/metrics.py @@ -0,0 +1,125 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Stateless metrics specific to async RL policy and advantage semantics.""" + +from __future__ import annotations + +import math +import statistics +from collections.abc import Mapping, Sequence +from typing import Any + + +def _p95(values: list[float]) -> float: + ordered = sorted(values) + return ordered[max(0, (95 * len(ordered) + 99) // 100 - 1)] + + +def rollout_metrics( + *, + rewards: Mapping[str, Sequence[float]] | None = None, + completion_lengths: Sequence[int] = (), + stop_reasons: Sequence[str | None] = (), + rollout_latency_s: float | None = None, +) -> dict[str, float | int]: + """Summarize one RL rollout collection without retaining state.""" + metrics: dict[str, float | int] = {} + reward_counts = [len(values) for values in (rewards or {}).values() if values] + if len(set(reward_counts)) > 1: + raise ValueError(f'reward metric lengths must match, got {reward_counts}') + sample_count = len(completion_lengths) or (reward_counts[0] if reward_counts else 0) + if completion_lengths and reward_counts and any(count != sample_count for count in reward_counts): + raise ValueError(f'reward and completion metric lengths must match: {reward_counts} != {sample_count}') + if stop_reasons and len(stop_reasons) != len(completion_lengths): + raise ValueError('stop reason and completion metric lengths must match: ' + f'{len(stop_reasons)} != {len(completion_lengths)}') + if sample_count: + metrics['sample_count'] = sample_count + if completion_lengths: + lengths = [int(value) for value in completion_lengths] + output_tokens = sum(lengths) + truncated_count = sum(reason == 'length' for reason in stop_reasons) + metrics.update({ + 'completion_length_mean': output_tokens / sample_count, + 'completion_length_p95': _p95(lengths), + 'completion_length_max': max(lengths), + 'completion_truncated_count': truncated_count, + 'completion_truncated_ratio': truncated_count / sample_count, + 'output_tokens': output_tokens, + }) + if rollout_latency_s is not None: + latency = float(rollout_latency_s) + metrics['rollout_latency_s'] = latency + metrics['output_tokens_per_s'] = output_tokens / latency if latency > 0 else 0.0 + elif rollout_latency_s is not None: + metrics['rollout_latency_s'] = float(rollout_latency_s) + + for name, raw_values in (rewards or {}).items(): + values = [float(value) for value in raw_values] + if not values: + continue + prefix = 'reward' if name == 'reward' else f'{name}_reward' + metrics[prefix] = sum(values) / len(values) + metrics[f'{prefix}_std'] = statistics.stdev(values) if len(values) > 1 else 0.0 + return metrics + + +def training_policy_metrics( + sample_tags: tuple[dict[str, Any], ...], + train_policy_version: int, +) -> dict[str, float | int]: + if not sample_tags: + raise ValueError('training batch must contain sample policy tags') + final_versions = [int(tag['final_policy_version']) for tag in sample_tags] + spans = [int(tag['policy_version_span']) for tag in sample_tags] + gaps = [int(train_policy_version) - version for version in final_versions] + if any(gap < 0 for gap in gaps): + raise ValueError( + f'training policy version {train_policy_version} is older than rollout versions {final_versions}') + return { + 'policy_version_gap_mean': sum(gaps) / len(gaps), + 'policy_version_gap_p95': _p95(gaps), + 'policy_version_gap_max': max(gaps), + 'rollout_policy_span_mean': sum(spans) / len(spans), + 'rollout_policy_span_max': max(spans), + } + + +def advantage_signal_metrics( + rewards: Sequence[float], + advantages: Sequence[float], + *, + num_generations: int, + zero_tolerance: float = 1e-8, +) -> dict[str, float | int]: + """Summarize whether GRPO groups provide a useful learning signal.""" + if num_generations <= 0: + raise ValueError(f'num_generations must be positive, got {num_generations}') + if len(rewards) != len(advantages): + raise ValueError(f'rewards and advantages must have equal length: {len(rewards)} != {len(advantages)}') + if len(rewards) == 0 or len(rewards) % num_generations: + raise ValueError(f'advantage metrics require complete groups: sample_count={len(rewards)}, ' + f'num_generations={num_generations}') + + reward_values = [float(value) for value in rewards] + advantage_values = [float(value) for value in advantages] + group_reward_stds: list[float] = [] + zero_advantage_groups = 0 + for start in range(0, len(reward_values), num_generations): + group_rewards = reward_values[start:start + num_generations] + group_advantages = advantage_values[start:start + num_generations] + reward_mean = sum(group_rewards) / num_generations + group_reward_stds.append(math.sqrt(sum((value - reward_mean)**2 for value in group_rewards) / num_generations)) + if max(abs(value) for value in group_advantages) <= zero_tolerance: + zero_advantage_groups += 1 + + advantage_mean = sum(advantage_values) / len(advantage_values) + group_count = len(group_reward_stds) + return { + 'group_count': group_count, + 'group_reward_std_mean': sum(group_reward_stds) / group_count, + 'zero_advantage_group_ratio': zero_advantage_groups / group_count, + 'positive_advantage_ratio': sum(value > zero_tolerance for value in advantage_values) / len(advantage_values), + 'advantage_mean': advantage_mean, + 'advantage_std': + math.sqrt(sum((value - advantage_mean)**2 for value in advantage_values) / len(advantage_values)), + } diff --git a/src/twinkle_agentic/async_rl/native_tq.py b/src/twinkle_agentic/async_rl/native_tq.py new file mode 100644 index 000000000..9bc137e67 --- /dev/null +++ b/src/twinkle_agentic/async_rl/native_tq.py @@ -0,0 +1,186 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Adapters and helpers for the native TransferQueue client API. + +The async RL data path deliberately uses ``BatchMeta`` as its descriptor. A +``kv_list`` result is a diagnostic snapshot, not a queue cursor, and therefore +must not be used to drive the hot path. +""" + +from __future__ import annotations + +from transfer_queue import GRPOGroupNSampler +from typing import Any, Protocol, Sequence + + +class AsyncTQClient(Protocol): + + async def async_get_meta(self, + *, + data_fields: list[str], + batch_size: int, + partition_id: str, + mode: str = 'fetch', + task_name: str | None = None, + sampling_config: dict[str, Any] | None = None) -> Any: + ... + + async def async_get_data(self, metadata: Any) -> Any: + ... + + async def async_put(self, data: Any, metadata: Any | None = None, partition_id: str | None = None) -> Any: + ... + + async def async_clear_partition(self, partition_id: str) -> Any: + ... + + async def async_check_consumption_status(self, task_name: str, partition_id: str) -> bool: + ... + + async def async_set_custom_meta(self, metadata: Any) -> Any: + ... + + +class ContextGRPOGroupNSampler(GRPOGroupNSampler): + """Select complete prompt groups using the request's generation count.""" + + def sample( + self, + ready_indexes: list[int], + batch_size: int, + task_name: str = '', + partition_id: str = '', + *args: Any, + **kwargs: Any, + ) -> tuple[list[int], list[int]]: + group_size = int(kwargs['n_samples_per_prompt']) + if group_size <= 0: + raise ValueError(f'n_samples_per_prompt must be positive, got {group_size}') + if batch_size % group_size: + raise ValueError(f'batch_size ({batch_size}) must be a multiple of n_samples_per_prompt ({group_size})') + + states = self._states.get(partition_id, {}).get(task_name, {}) + dp_rank = kwargs.get('dp_rank') + batch_index = kwargs.get('batch_index') + if dp_rank in states and batch_index in states[dp_rank]: + return states[dp_rank][batch_index] + + ready = sorted(ready_indexes) + selected: list[int] = [] + offset = 0 + while offset <= len(ready) - group_size and len(selected) < batch_size: + group = ready[offset:offset + group_size] + if all(right - left == 1 for left, right in zip(group, group[1:])): + selected.extend(group) + offset += group_size + else: + offset += 1 + + if len(selected) != batch_size: + return [], [] + + result = (selected, selected.copy()) + if dp_rank is not None: + states.setdefault(dp_rank, {})[batch_index] = result + self._states.setdefault(partition_id, {})[task_name] = states + return result + + +def batch_size_for_groups(groups: int, num_generations: int) -> int: + if groups <= 0: + raise ValueError(f'groups must be positive, got {groups}') + if num_generations <= 0: + raise ValueError(f'num_generations must be positive, got {num_generations}') + return groups * num_generations + + +def validate_group_batch_size(batch_size: int, num_generations: int) -> None: + if batch_size <= 0: + raise ValueError(f'batch_size must be positive, got {batch_size}') + if num_generations <= 0: + raise ValueError(f'num_generations must be positive, got {num_generations}') + if batch_size % num_generations: + raise ValueError(f'batch_size={batch_size} must be divisible by num_generations={num_generations}') + + +def metadata_size(metadata: Any) -> int: + """Return the native BatchMeta size.""" + if metadata is None: + return 0 + return int(metadata.size) + + +async def fetch_ready_batch( + client: AsyncTQClient, + *, + data_fields: list[str], + batch_size: int, + partition_id: str, + task_name: str, + num_generations: int, + sampling_config: dict[str, Any] | None = None, +) -> Any | None: + """Fetch one complete group-aligned batch using TQ production status. + + The caller owns the outer service loop. This helper performs one request + only, which keeps shutdown and failure handling explicit and avoids hiding + an unbounded wait in a data-plane utility. + """ + + validate_group_batch_size(batch_size, num_generations) + config = dict(sampling_config or {}) + config['n_samples_per_prompt'] = num_generations + metadata = await client.async_get_meta( + data_fields=list(data_fields), + batch_size=batch_size, + partition_id=partition_id, + mode='fetch', + task_name=task_name, + sampling_config=config, + ) + return metadata if metadata_size(metadata) else None + + +async def append_fields(client: AsyncTQClient, data: Any, metadata: Any) -> Any: + """Append fields to exactly the samples described by ``metadata``.""" + + if metadata_size(metadata) == 0: + raise ValueError('cannot append fields to an empty BatchMeta') + return await client.async_put(data=data, metadata=metadata) + + +async def set_sample_tags(client: AsyncTQClient, metadata: Any, tags: Sequence[dict[str, Any]]) -> None: + """Persist tags through the native metadata API in one controller request.""" + + if metadata_size(metadata) != len(tags): + raise ValueError(f'metadata size {metadata_size(metadata)} does not match tags {len(tags)}') + metadata.update_custom_meta([dict(tag) for tag in tags]) + await client.async_set_custom_meta(metadata) + + +def split_batch_meta(metadata: Any, group_size: int) -> list[Any]: + """Split a preallocated BatchMeta into contiguous prompt-group views.""" + + size = metadata_size(metadata) + if group_size <= 0: + raise ValueError(f'group_size must be positive, got {group_size}') + if size % group_size: + raise ValueError(f'metadata size {size} is not divisible by group_size {group_size}') + return [metadata.select_samples(list(range(start, start + group_size))) for start in range(0, size, group_size)] + + +async def preallocate_partition( + client: AsyncTQClient, + *, + partition_id: str, + prompt_fields: Any, +) -> Any: + """Insert prompt rows once and return the native BatchMeta descriptor.""" + + batch_size = getattr(prompt_fields, 'batch_size', None) + if batch_size is None or len(batch_size) == 0 or int(batch_size[0]) <= 0: + raise ValueError('prompt_fields must be a non-empty batched TensorDict') + return await client.async_put(data=prompt_fields, partition_id=partition_id) + + +async def clear_partition(client: AsyncTQClient, partition_id: str) -> None: + await client.async_clear_partition(partition_id) diff --git a/src/twinkle_agentic/async_rl/pipeline.py b/src/twinkle_agentic/async_rl/pipeline.py new file mode 100644 index 000000000..75a78e023 --- /dev/null +++ b/src/twinkle_agentic/async_rl/pipeline.py @@ -0,0 +1,705 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Driver for the independent async-RL Ray workers.""" + +from __future__ import annotations + +import asyncio +import time +from dataclasses import dataclass +from functools import partial +from pydoc import locate +from typing import Any, Sequence + +from twinkle.metric import MetricRecord, MetricsReporter, create_metrics_reporter +from .context_manager import LoraContextManager +from .data_plane import TQDataPlane +from .scheduler import ContextSchedulePolicy, SchedulerConfig +from .types import LoraContext, PartitionAdmission +from .utils import (TrainBatchConfig, build_native_fsdp_model_kwargs, configure_lora_lr_scheduler, + resolve_context_learning_rate, resolve_context_lora_target_modules, resolve_context_loss_config, + resolve_model_attention_implementation, resolve_sequence_parallel_size, sampler_data_parallel_size, + validate_context_batch_config) +from .workers import AdvantageWorker, RolloutWorker, TrainerWorker + + +@dataclass(frozen=True) +class AsyncMultiLoraGRPOConfig: + metrics_drain_interval_s: float = 1.0 + + +class AsyncMultiLoraGRPOPipeline: + """Owns the production async-RL runtime and drives its worker services. + + ``from_config`` is the real training construction path. The explicit + constructor remains available to fake-TQ tests, where injecting a fake + sampler/model is the point of the test. + """ + + def __init__(self, + *, + context_manager: LoraContextManager, + rollout_worker: RolloutWorker, + advantage_worker: AdvantageWorker, + trainer_worker: TrainerWorker, + metrics: MetricsReporter | None = None, + config: AsyncMultiLoraGRPOConfig = AsyncMultiLoraGRPOConfig(), + sampler: Any | None = None, + model: Any | None = None, + contexts: Sequence[LoraContext] = ()): + self.context_manager = context_manager + self.rollout_worker = rollout_worker + self.advantage_worker = advantage_worker + self.trainer_worker = trainer_worker + self.sampler = sampler + self.model = model + self.contexts = tuple(contexts) + self.metrics = metrics + self.config = config + + @classmethod + def from_config( + cls, + raw_config: dict[str, Any], + *, + persistent: bool = False, + ) -> AsyncMultiLoraGRPOPipeline: + """Build the complete Ray/TQ runtime from the async-RL YAML mapping.""" + from omegaconf import OmegaConf + + raw_config = OmegaConf.to_container(OmegaConf.create(raw_config), resolve=True) + if not isinstance(raw_config, dict): + raise TypeError('async-RL config must resolve to a mapping') + + import ray + import transfer_queue as tq + from peft import LoraConfig + + import twinkle + from twinkle import DeviceGroup, DeviceMesh + from twinkle.data_format import SamplingParams + from twinkle.model import MultiLoraTransformersModel + from twinkle.processor import InputProcessor + from .native_tq import ContextGRPOGroupNSampler + + runtime = raw_config['runtime'] + model_config = raw_config['model'] + lora_config_data = raw_config['lora'] + loss_config_data = raw_config.get('loss') + template_config = raw_config.get('template', {}) + template_cls = template_config.get('cls', 'Qwen3_5Template') + enable_thinking = bool(template_config.get('enable_thinking', False)) + rollout_output_config = dict(raw_config.get('rollout_output') or {}) + sampler_gpus = int(runtime['sampler_gpus']) + sampler_tp = int(runtime['sampler_tp']) + sampler_dp = sampler_data_parallel_size(sampler_gpus, sampler_tp) + model_dp = int(runtime['model_gpus']) + sequence_parallel_size = resolve_sequence_parallel_size( + model_dp, + int(model_config['sequence_parallel_size']), + ) + padding_free = bool(model_config['padding_free']) + attn_implementation = resolve_model_attention_implementation( + model_config, + padding_free=padding_free, + sequence_parallel_size=sequence_parallel_size, + ) + model_max_length = int(model_config['max_length']) + sampler_config = raw_config['sampler'] + total_gpus = model_dp + sampler_gpus + device_groups = [ + DeviceGroup('model', list(range(int(runtime['model_gpus']))), device_type='GPU'), + DeviceGroup( + 'sampler', + list(range(int(runtime['model_gpus']), total_gpus)), + device_type='GPU', + gpus_per_worker=sampler_tp, + ), + ] + twinkle.initialize(mode='ray', nproc_per_node=total_gpus, groups=device_groups, lazy_collect=False) + tq.init( + OmegaConf.create( + { + 'controller': { + 'sampler': ContextGRPOGroupNSampler, + 'polling_mode': bool(raw_config['tq'].get('polling_mode', True)), + }, + 'backend': { + 'SimpleStorage': { + 'num_data_storage_units': raw_config['tq']['storage_units'] + } + }, + }, + flags={'allow_objects': True})) + + model_mesh = DeviceMesh.from_sizes( + world_size=model_dp, + dp_size=model_dp, + ulysses_size=sequence_parallel_size, + ) + model_data_parallel_size = model_mesh.data_world_size + sampler_mesh = DeviceMesh.from_sizes(world_size=sampler_gpus, dp_size=sampler_dp, tp_size=sampler_tp) + model_kwargs = build_native_fsdp_model_kwargs(model_config) + if attn_implementation is not None: + model_kwargs['attn_implementation'] = attn_implementation + model = MultiLoraTransformersModel( + model_id=runtime['model_id'], + device_mesh=model_mesh, + remote_group='model', + max_length=model_max_length, + **model_kwargs, + ) + contexts: list[LoraContext] = [] + prompt_sources: dict[str, Any] = {} + rollout_config: dict[str, dict[str, Any]] = {} + train_batch_configs: dict[str, TrainBatchConfig] = {} + rewards: dict[str, Any] = {} + evaluation_config: dict[str, dict[str, Any]] = {} + evaluation_rewards: dict[str, Any] = {} + initial_paths: dict[str, str] = {} + global_evaluation = dict(raw_config.get('evaluation') or {}) + for item in raw_config['lora_contexts']: + train = item['train'] + context = LoraContext( + item['tenant_id'], + item['training_run_id'], + runtime['model_id'], + item['adapter_name'], + ) + contexts.append(context) + adapter_lora_config = LoraConfig( + target_modules=resolve_context_lora_target_modules(item, lora_config_data), + r=lora_config_data['r'], + lora_alpha=lora_config_data['alpha'], + lora_dropout=lora_config_data['dropout'], + ) + model.add_adapter_to_model( + context.adapter_name, + adapter_lora_config, + gradient_accumulation_steps=1, + ) + model.set_optimizer( + 'AdamW', + lr=resolve_context_learning_rate(train, lora_config_data), + adapter_name=context.adapter_name, + ) + configure_lora_lr_scheduler(model, context.adapter_name, lora_config_data) + loss_cls, loss_kwargs = resolve_context_loss_config(item, loss_config_data) + model.set_loss( + loss_cls, + adapter_name=context.adapter_name, + **loss_kwargs, + ) + model.set_processor( + InputProcessor, + adapter_name=context.adapter_name, + padding_free=padding_free, + ) + model.set_template( + template_cls, + model_id=runtime['model_id'], + adapter_name=context.adapter_name, + enable_thinking=enable_thinking, + max_length=model_max_length, + ) + + rollout = item['rollout'] + rollout_batch_size = int(rollout['batch_size']) + num_generations = int(rollout['num_generations']) + train_batch_config = TrainBatchConfig( + mini_batch_size=int(train['mini_batch_size']), + micro_batch_size=int(train['micro_batch_size']), + dynamic_batching=bool(train.get('dynamic_batching', False)), + max_tokens_per_micro_batch=(int(train['max_tokens_per_micro_batch']) + if train.get('max_tokens_per_micro_batch') is not None else None), + packing_algorithm=str(train.get('packing_algorithm', 'ffd')), + ) + validate_context_batch_config( + context.key, + rollout_groups=rollout_batch_size, + num_generations=num_generations, + train=train_batch_config, + sampler_dp=sampler_dp, + model_dp=model_data_parallel_size, + ) + prompt_sources[context.key] = partial( + _prompt_batches, + item['dataset'], + model_id=runtime['model_id'], + batch_size=rollout_batch_size, + template_cls=template_cls, + enable_thinking=enable_thinking, + ) + rollout_config[context.key] = { + 'context': + context, + 'batch_size': + rollout_batch_size, + 'num_generations': + num_generations, + 'sampling_params': + SamplingParams( + max_tokens=rollout['max_tokens'], + temperature=rollout['temperature'], + top_p=rollout['top_p'], + repetition_penalty=float(rollout.get('repetition_penalty', 1.0)), + logprobs=1, + num_samples=1, + ), + } + train_batch_configs[context.key] = train_batch_config + rewards[context.key] = _reward_for_context( + item.get('reward'), + context_key=context.key, + ) + if bool(global_evaluation.get('enabled', False)): + eval_dataset = item.get('eval_dataset') + if eval_dataset is None: + raise ValueError(f'eval_dataset is required for periodic evaluation of {context.key}') + eval_batch_size = int(global_evaluation.get('batch_size', 16)) + eval_interval = int(global_evaluation.get('interval', 1)) + if eval_batch_size <= 0 or eval_interval <= 0: + raise ValueError('evaluation.batch_size and evaluation.interval must be positive') + eval_sampling = dict(global_evaluation.get('sampling_params') or {}) + evaluation_config[context.key] = { + 'interval': + eval_interval, + 'dataset_name': + eval_dataset.get('name', eval_dataset['dataset_id']), + 'prompt_batches': + partial( + _prompt_batches, + eval_dataset, + model_id=runtime['model_id'], + batch_size=eval_batch_size, + template_cls=template_cls, + enable_thinking=enable_thinking, + full_batches_only=False, + ), + 'sampling_params': + SamplingParams( + max_tokens=int(eval_sampling.get('max_tokens', rollout['max_tokens'])), + temperature=float(eval_sampling.get('temperature', 0.0)), + top_p=float(eval_sampling.get('top_p', 1.0)), + repetition_penalty=float(eval_sampling.get('repetition_penalty', 1.0)), + logprobs=0, + num_samples=1, + ), + } + evaluation_rewards[context.key] = _reward_for_context( + eval_dataset.get('reward'), + context_key=f'{context.key} evaluation', + ) + initial_paths[context.key] = _collect_adapter_path( + model.save( + f'async-{context.adapter_name}-initial', + output_dir=runtime['output_dir'], + adapter_name=context.adapter_name, + ), + operation=f'initial adapter save for {context.key}', + ) + + manager = create_cpu_actor( + LoraContextManager, + max_staleness=runtime['max_staleness'], + max_steps=runtime['max_steps'], + ) + for context in contexts: + ray.get(manager.register_context.remote(context, adapter_path=initial_paths[context.key])) + + from .vllm_sampler_tq import VLLMSamplerTQ + sampler_engine_args = { + 'tensor_parallel_size': sampler_tp, + 'enable_lora': True, + 'max_loras': int(runtime['sampler_max_loras']), + 'max_lora_rank': lora_config_data['r'], + 'max_model_len': int(sampler_config['max_model_len']), + 'gpu_memory_utilization': float(sampler_config['gpu_memory_utilization']), + 'max_num_seqs': int(sampler_config['max_num_seqs']), + 'enforce_eager': bool(sampler_config['enforce_eager']), + 'seed': int(runtime.get('seed', 1)), + } + if sampler_config.get('max_num_batched_tokens') is not None: + sampler_engine_args['max_num_batched_tokens'] = int(sampler_config['max_num_batched_tokens']) + sampler = VLLMSamplerTQ( + model_id=runtime['model_id'], + remote_group='sampler', + device_mesh=sampler_mesh, + engine_args=sampler_engine_args, + reward_registry=rewards, + context_manager=manager, + rollout_max_retries=int(runtime.get('rollout_max_retries', 2)), + rollout_retry_delay_s=float(runtime.get('rollout_retry_delay_s', 0.5)), + rollout_output_dir=(rollout_output_config.get('output_dir') if bool( + rollout_output_config.get('enabled', False)) else None), + rollout_output_include_token_ids=bool(rollout_output_config.get('include_token_ids', False)), + ) + sampler.set_template( + template_cls, + model_id=runtime['model_id'], + enable_thinking=enable_thinking, + max_length=model_max_length, + ) + + rollout_worker = create_cpu_actor( + RolloutWorker, + context_manager=manager, + data_plane=TQDataPlane(), + sampler=sampler, + prompt_batches=prompt_sources, + rollout_config=rollout_config, + scheduler=_scheduler(raw_config['scheduler']['rollout']), + allow_partial_rollout=runtime['allow_partial_rollout'], + persistent=persistent, + ) + advantage_worker = create_cpu_actor( + AdvantageWorker, + context_manager=manager, + data_plane=TQDataPlane(), + advantage_fn=_compute_advantages, + scheduler=_scheduler(raw_config['scheduler']['advantage']), + persistent=persistent, + ) + trainer_worker = create_cpu_actor( + TrainerWorker, + context_manager=manager, + data_plane=TQDataPlane(), + train_fn=partial( + _train_batch, + model, + train_batch_configs, + model_data_parallel_size=model_data_parallel_size, + ), + train_with_config_fn=partial( + _train_batch_with_config, + model, + model_data_parallel_size=model_data_parallel_size, + ), + train_batch_configs=train_batch_configs, + save_adapter=partial(_save_adapter, model, runtime['output_dir']), + mini_batch_sizes={ + key: config.mini_batch_size + for key, config in train_batch_configs.items() + }, + scheduler=_scheduler(raw_config['scheduler']['train']), + keep_adapter_versions=runtime['keep_adapter_versions'], + initial_adapter_paths=initial_paths, + remove_adapter=partial(_remove_adapter_snapshot, sampler), + evaluation_config=evaluation_config, + evaluate_batch=partial(_evaluate_batch, sampler, evaluation_rewards) if evaluation_config else None, + evaluate_with_reward_fn=partial(_evaluate_batch_with_reward, sampler), + evaluation_rewards=evaluation_rewards, + persistent=persistent, + ) + raw_metrics_config = raw_config.get('metrics') + metrics_config = dict(raw_metrics_config or {}) + metrics = create_metrics_reporter( + raw_metrics_config, + run_id=str(runtime.get('run_id', 'async_multi_lora_grpo')), + ) + return cls( + context_manager=manager, + rollout_worker=rollout_worker, + advantage_worker=advantage_worker, + trainer_worker=trainer_worker, + sampler=sampler, + metrics=metrics, + config=AsyncMultiLoraGRPOConfig( + metrics_drain_interval_s=float(metrics_config.get('drain_interval_s', 1.0)), ), + model=model, + contexts=contexts, + ) + + async def run_async(self) -> dict[str, Any]: + started = time.perf_counter() + workers = [self.rollout_worker, self.advantage_worker, self.trainer_worker] + await asyncio.gather(*(worker.start.remote() for worker in workers)) + try: + while True: + await self._drain_metrics() + states = await asyncio.gather(*(worker.get_service_state.remote() for worker in workers)) + failures = [state['failure'] for state in states if state['failure']] + if failures: + raise RuntimeError(f'async RL worker failed: {failures[0]}') + if self.sampler is not None: + await asyncio.to_thread(self.sampler.check_health) + running = any(bool(state['running']) for state in states) + if not running: + if await self.context_manager.is_run_finished.remote(): + break + raise RuntimeError('async RL workers stopped before all contexts were drained') + await asyncio.sleep(self.config.metrics_drain_interval_s) + except Exception as exc: + if self.metrics is not None: + self.metrics.record( + MetricRecord( + stage='run', + status='failed', + values={'wall_time_s': time.perf_counter() - started}, + attributes={'error': f'{type(exc).__name__}: {exc}'}, + )) + self.metrics.flush() + raise + finally: + await asyncio.gather(*(worker.stop.remote() for worker in workers), return_exceptions=True) + await self._drain_metrics() + result = { + 'trained_partitions': await self.context_manager.get_completed_partitions.remote(), + 'wall_time_s': time.perf_counter() - started, + } + if self.metrics is not None: + self.metrics.record(MetricRecord(stage='run', values=result)) + self.metrics.flush() + result['metrics_health'] = self.metrics.health() + return result + + def run(self) -> dict[str, Any]: + try: + return asyncio.run(self.run_async()) + finally: + if self.metrics is not None: + self.metrics.close() + + async def _drain_metrics(self) -> None: + workers = [self.rollout_worker, self.advantage_worker, self.trainer_worker] + for worker in workers: + records = await worker.drain_metric_records.remote() + if self.metrics is not None: + self.metrics.record_many(records) + if self.sampler is not None: + records = await asyncio.to_thread(self.sampler.drain_metric_records) + if self.metrics is not None: + self.metrics.record_many(records) + + +def create_cpu_actor(cls: type, *args: Any, **kwargs: Any) -> Any: + """Deploy a CPU service as one raw Ray actor; local tests use the class directly.""" + + import ray + actor_class = ray.remote( + num_cpus=1, + runtime_env={'env_vars': { + 'TWINKLE_MODE': 'ray' + }}, + )( + cls) + return actor_class.remote(*args, **kwargs) + + +def _scheduler(config: dict[str, Any]) -> SchedulerConfig: + return SchedulerConfig(ContextSchedulePolicy(config['policy']), config.get('max_consecutive_units')) + + +def _prompt_batches( + dataset_config: dict[str, Any], + *, + model_id: str, + batch_size: int, + template_cls: str, + enable_thinking: bool, + full_batches_only: bool = True, +): + """Create a lazy, full-batch-only prompt source for one context.""" + from twinkle.dataloader import DataLoader + from twinkle.dataset import Dataset, DatasetMeta + from twinkle.preprocessor import llm as llm_processors + + def batches(): + data_num = dataset_config.get('data_num') + dataset = Dataset( + DatasetMeta( + dataset_config['dataset_id'], + subset_name=dataset_config.get('subset_name'), + split=dataset_config.get('split', 'train'), + data_slice=range(int(data_num)) if data_num is not None else None, + )) + dataset.set_template( + template_cls, + model_id=model_id, + max_length=dataset_config['max_length'], + enable_thinking=enable_thinking, + ) + processor_name = dataset_config.get('processor', 'GSM8KProcessor') + processor_cls = getattr(llm_processors, processor_name) + if processor_name == 'GSM8KProcessor': + processor = processor_cls(system=dataset_config['system_prompt']) + else: + processor = processor_cls() + dataset.map(processor) + dataset.encode(add_generation_prompt=True) + loader = DataLoader( + dataset=dataset, + batch_size=batch_size, + min_batch_size=batch_size if full_batches_only else 1, + ) + remaining = data_num + remaining = None if remaining is None else int(remaining) + for batch in loader: + if full_batches_only and (len(batch) != batch_size or (remaining is not None and remaining < batch_size)): + return + yield batch + if remaining is not None: + remaining -= batch_size + + return batches() + + +def _evaluate_batch( + sampler: Any, + reward_registry: dict[str, Any], + prompts: Sequence[dict[str, Any]], + admission: PartitionAdmission, + adapter_path: str, + policy_version: int, + sampling_params: Any, +) -> dict[str, Any]: + reward_fn = reward_registry[admission.context.key] + return _evaluate_batch_with_reward( + sampler, + prompts, + admission, + adapter_path, + policy_version, + sampling_params, + reward_fn, + ) + + +def _evaluate_batch_with_reward( + sampler: Any, + prompts: Sequence[dict[str, Any]], + admission: PartitionAdmission, + adapter_path: str, + policy_version: int, + sampling_params: Any, + reward_fn: Any, +) -> dict[str, Any]: + from .utils import sample_responses_to_rollout_rows + + responses = sampler.evaluate( + list(prompts), + sampling_params, + admission.context.adapter_name, + adapter_path, + ) + rows = sample_responses_to_rollout_rows(list(prompts), responses, policy_version=policy_version) + rewards = list(reward_fn(rows, context=admission.context)) + return { + 'rewards': rewards, + 'completion_lengths': [int(row['completion_length']) for row in rows], + } + + +def _reward_for_context( + reward_config: dict[str, Any] | None = None, + *, + context_key: str, +) -> Any: + from twinkle.reward import Reward + + config = dict(reward_config or {}) + class_path = config.get('class_path', '') + reward_cls = locate(class_path) + if not isinstance(reward_cls, type) or not issubclass(reward_cls, Reward): + raise TypeError(f'reward.class_path {class_path!r} for {context_key} must reference a Reward subclass') + return reward_cls(**dict(config.get('kwargs') or {})) + + +def _compute_advantages(data: Any, admission: PartitionAdmission) -> tuple[list[float], list[float]]: + from twinkle.advantage import GRPOAdvantage + rewards = [float(value) for value in data['rewards']] + advantages = GRPOAdvantage()(rewards, num_generations=admission.num_generations, scale='group').tolist() + return advantages, rewards + + +def _train_batch( + model: Any, + train_batch_configs: dict[str, TrainBatchConfig], + data: Any, + admission: PartitionAdmission, + *, + model_data_parallel_size: int = 1, +) -> dict[str, Any]: + config = train_batch_configs[admission.context.key] + return _train_batch_with_config( + model, + data, + admission, + config, + model_data_parallel_size=model_data_parallel_size, + ) + + +def _train_batch_with_config( + model: Any, + data: Any, + admission: PartitionAdmission, + config: TrainBatchConfig, + *, + model_data_parallel_size: int = 1, +) -> dict[str, Any]: + from .tq_utils import REQUIRED_MODEL_INPUT_FIELDS + + size = int(data.batch_size[0]) + inputs = [{name: data[name][index] for name in REQUIRED_MODEL_INPUT_FIELDS} for index in range(size)] + old_logps = list(data['logprobs']) + advantages = list(data['advantages']) + if size != config.mini_batch_size: + raise ValueError(f'train batch for {admission.context.key} has {size} samples; ' + f'expected mini_batch_size={config.mini_batch_size}') + + if size % model_data_parallel_size: + raise ValueError(f'train batch size {size} must be divisible by model DP size ' + f'{model_data_parallel_size}') + model.forward_backward( + inputs=inputs, + old_logps=old_logps, + advantages=advantages, + adapter_name=admission.context.adapter_name, + micro_batch_size=config.micro_batch_size, + dynamic_batching=config.dynamic_batching, + max_tokens_per_micro_batch=config.max_tokens_per_micro_batch, + packing_algorithm=config.packing_algorithm, + sync_gradients=True, + loss_scale=1.0, + ) + + model.clip_grad_and_step(adapter_name=admission.context.adapter_name) + metrics = dict(model.calculate_metric(is_training=True, adapter_name=admission.context.adapter_name)) + metrics['mini_batch_size'] = config.mini_batch_size + metrics['micro_batch_size_per_rank'] = config.micro_batch_size + metrics['dynamic_batching'] = config.dynamic_batching + return metrics + + +def _save_adapter(model: Any, output_dir: str, admission: PartitionAdmission) -> str: + return _collect_adapter_path( + model.save( + f'async-{admission.context.adapter_name}-v{admission.step + 1}', + output_dir=output_dir, + adapter_name=admission.context.adapter_name, + ), + operation=f'adapter save for {admission.partition_id}', + ) + + +def _collect_adapter_path(value: Any, *, operation: str) -> str: + """Collect a lazy model.save result at the async-RL publication boundary.""" + if callable(value) and getattr(value, '_is_lazy_collect', False): + value = value() + return _require_adapter_path(value, operation=operation) + + +def _require_adapter_path(value: Any, *, operation: str) -> str: + """Fail at the save boundary instead of publishing an invalid policy.""" + if not isinstance(value, str) or not value: + raise TypeError(f'{operation} must return a non-empty checkpoint path string, ' + f'got {type(value).__name__}: {value!r}') + return value + + +def _remove_adapter_snapshot(sampler: Any, adapter_path: str) -> None: + """Unload an unreferenced policy from vLLM before deleting its checkpoint.""" + from .workers import _remove_local_adapter + + sampler.unload_adapter_paths([adapter_path]) + _remove_local_adapter(adapter_path) diff --git a/src/twinkle_agentic/async_rl/scheduler.py b/src/twinkle_agentic/async_rl/scheduler.py new file mode 100644 index 000000000..e77c53a93 --- /dev/null +++ b/src/twinkle_agentic/async_rl/scheduler.py @@ -0,0 +1,73 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Context selection policies used independently by rollout/advantage/training.""" + +from __future__ import annotations + +from dataclasses import dataclass +from enum import StrEnum +from typing import Sequence + +from .types import LoraContext, PartitionAdmission + + +class ContextSchedulePolicy(StrEnum): + ROUND_ROBIN = 'round_robin' + STICKY = 'sticky' + OLDEST_PARTITION = 'oldest_partition' + + +@dataclass(frozen=True) +class SchedulerConfig: + policy: ContextSchedulePolicy = ContextSchedulePolicy.ROUND_ROBIN + max_consecutive_units: int | None = 1 + + +@dataclass(frozen=True) +class ScheduleCandidate: + context: LoraContext + partition: PartitionAdmission | None = None + + +class ContextScheduler: + + def __init__(self, config: SchedulerConfig): + self.config = config + self._cursor = 0 + self._sticky_key: str | None = None + self._consecutive = 0 + + def choose(self, candidates: Sequence[ScheduleCandidate]) -> ScheduleCandidate | None: + if not candidates: + return None + if self.config.policy is ContextSchedulePolicy.OLDEST_PARTITION: + return min( + candidates, + key=lambda item: (item.partition.created_order if item.partition else float('inf'), item.context.key)) + if self.config.policy is ContextSchedulePolicy.STICKY and self._sticky_key is not None: + cap = self.config.max_consecutive_units + if cap is None or self._consecutive < cap: + for candidate in candidates: + if candidate.context.key == self._sticky_key: + return candidate + else: + for candidate in candidates: + if candidate.context.key != self._sticky_key: + return candidate + index = self._cursor % len(candidates) + return candidates[index] + + def on_success(self, candidate: ScheduleCandidate) -> None: + if self.config.policy is ContextSchedulePolicy.STICKY and candidate.context.key == self._sticky_key: + self._consecutive += 1 + return + self._sticky_key = candidate.context.key + self._consecutive = 1 + if self.config.policy is ContextSchedulePolicy.ROUND_ROBIN: + self._cursor += 1 + + def on_blocked(self, candidate: ScheduleCandidate) -> None: + if candidate.context.key == self._sticky_key: + self._sticky_key = None + self._consecutive = 0 + if self.config.policy is ContextSchedulePolicy.ROUND_ROBIN: + self._cursor += 1 diff --git a/src/twinkle_agentic/async_rl/tq_utils.py b/src/twinkle_agentic/async_rl/tq_utils.py new file mode 100644 index 000000000..55284567e --- /dev/null +++ b/src/twinkle_agentic/async_rl/tq_utils.py @@ -0,0 +1,29 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from __future__ import annotations + +from twinkle.tq_utils import columns_to_tq_fields, rows_to_tq_fields + +TRANSFORMERS_INPUT_FIELDS = ( + 'input_ids', + 'labels', + 'attention_mask', + 'position_ids', + 'cu_seqlens', + 'completion_mask', + 'pixel_values', + 'image_grid_thw', + 'video_pixel_values', + 'video_grid_thw', + 'input_features', + 'feature_attention_mask', +) +REQUIRED_MODEL_INPUT_FIELDS = ('input_ids', 'labels', 'attention_mask', 'position_ids') +ROLLOUT_TRAIN_FIELDS = (*TRANSFORMERS_INPUT_FIELDS, 'logprobs', 'rewards', 'advantages', 'returns') + +__all__ = [ + 'ROLLOUT_TRAIN_FIELDS', + 'REQUIRED_MODEL_INPUT_FIELDS', + 'TRANSFORMERS_INPUT_FIELDS', + 'columns_to_tq_fields', + 'rows_to_tq_fields', +] diff --git a/src/twinkle_agentic/async_rl/types.py b/src/twinkle_agentic/async_rl/types.py new file mode 100644 index 000000000..69545858b --- /dev/null +++ b/src/twinkle_agentic/async_rl/types.py @@ -0,0 +1,104 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Training-domain descriptors for native-TQ async multi-LoRA RL.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, TypedDict + + +@dataclass(frozen=True) +class LoraContext: + tenant_id: str + training_run_id: str + base_model_id: str + adapter_name: str + tool_profile: str = 'default' + + @property + def key(self) -> str: + return f'{self.tenant_id}/{self.training_run_id}/{self.adapter_name}' + + def partition_id(self, step: int) -> str: + return f'{self.key}/train_{step}' + + +@dataclass(frozen=True) +class RolloutPolicy: + """The immutable policy snapshot used by one rollout group.""" + + context_key: str + adapter_name: str + version: int + adapter_path: str | None + + +@dataclass(frozen=True) +class PartitionAdmission: + """Control-plane admission result for one complete training-data partition. + + A partition is only a training-data batch. Its prompt groups may be + generated by different policy snapshots. + """ + + context: LoraContext + partition_id: str + step: int + target_groups: int + num_generations: int + created_order: int + + @property + def sample_count(self) -> int: + return self.target_groups * self.num_generations + + +@dataclass +class PromptGroup: + """A prompt and the BatchMeta descriptor reserved for its generated samples.""" + + context: LoraContext + partition: PartitionAdmission + group_id: str + prompt: dict[str, Any] + batch_meta: Any + + @property + def partition_id(self) -> str: + return self.partition.partition_id + + @property + def num_samples(self) -> int: + return self.partition.num_generations + + +@dataclass +class PreparedPartition: + """Partition prepared in TQ and ready for sampler submission.""" + + admission: PartitionAdmission + groups: tuple[PromptGroup, ...] + sampling_params: Any + + +@dataclass +class ClaimedBatch: + """A TQ-consumed batch and the BatchMeta descriptor selecting its samples.""" + + admission: PartitionAdmission + data: Any + batch_meta: Any + sample_tags: tuple[dict[str, Any], ...] = () + + +class RolloutOutput(TypedDict, total=False): + logprobs: list[float] + rewards: float + completion_length: int + generation_idx: int + rollout_policy_version: int + rollout_adapter_path: str | None + rollout_policy_versions: list[int] + initial_policy_version: int + final_policy_version: int + policy_version_span: int diff --git a/src/twinkle_agentic/async_rl/utils.py b/src/twinkle_agentic/async_rl/utils.py new file mode 100644 index 000000000..2c8384aaa --- /dev/null +++ b/src/twinkle_agentic/async_rl/utils.py @@ -0,0 +1,208 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Shared configuration helpers for synchronous and asynchronous RL runners.""" + +from __future__ import annotations + +import math +import os +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from typing import Any, Literal + +from twinkle.data_format import SampleResponse +from .types import RolloutOutput + + +@dataclass(frozen=True) +class TrainBatchConfig: + mini_batch_size: int + micro_batch_size: int + dynamic_batching: bool = False + max_tokens_per_micro_batch: int | None = None + packing_algorithm: Literal['ffd', 'kk'] = 'ffd' + + +def _extract_sampled_token_logps(logprobs: Any) -> list[float]: + return [0.0 if not item else float(item[0][1]) for item in logprobs or []] + + +def sample_responses_to_rollout_rows( + sources: list[dict[str, Any]], + responses: list[SampleResponse], + *, + policy_version: int | None, +) -> list[RolloutOutput]: + rows: list[RolloutOutput] = [] + for source, response in zip(sources, responses): + for sequence in response.sequences: + row = dict(source) + row.update(sequence.new_input_feature or {}) + row['logprobs'] = _extract_sampled_token_logps(sequence.logprobs) + row['stop_reason'] = sequence.stop_reason + row['completion_length'] = len(sequence.tokens) + row['rollout_policy_version'] = policy_version + rows.append(row) + return rows + + +def resolve_adapter_path(adapter_path: str) -> str: + path = os.path.abspath(os.path.expanduser(str(adapter_path))) + if not os.path.exists(path): + raise FileNotFoundError(f'local LoRA adapter path does not exist: {path}') + return path + + +def sampler_data_parallel_size(sampler_gpus: int, sampler_tp: int) -> int: + if sampler_gpus <= 0: + raise ValueError(f'sampler_gpus must be positive, got {sampler_gpus}') + if sampler_tp <= 0: + raise ValueError(f'sampler_tp must be positive, got {sampler_tp}') + if sampler_gpus % sampler_tp != 0: + raise ValueError(f'sampler_gpus ({sampler_gpus}) must be divisible by sampler_tp ({sampler_tp})') + return sampler_gpus // sampler_tp + + +def resolve_sequence_parallel_size(model_gpus: int, configured_size: int) -> int: + if configured_size <= 0: + raise ValueError(f'model.sequence_parallel_size must be positive, got {configured_size}') + if model_gpus % configured_size: + raise ValueError(f'runtime.model_gpus ({model_gpus}) must be divisible by ' + f'model.sequence_parallel_size ({configured_size})') + return configured_size + + +def resolve_model_attention_implementation( + model_config: Mapping[str, Any], + *, + padding_free: bool, + sequence_parallel_size: int, +) -> str | None: + implementation = model_config.get('attn_implementation') + if implementation is not None: + implementation = str(implementation) + if padding_free and sequence_parallel_size > 1 and implementation != 'flash_attention_2': + raise ValueError('model.attn_implementation must be flash_attention_2 when ' + 'model.padding_free=true and model.sequence_parallel_size>1') + return implementation + + +def build_native_fsdp_model_kwargs(model_config: Mapping[str, Any]) -> dict[str, Any]: + strategy = str(model_config.get('strategy', 'native_fsdp')) + if strategy != 'native_fsdp': + raise ValueError(f'model.strategy must be native_fsdp for RL training, got {strategy!r}') + return { + 'strategy': strategy, + 'fsdp_config': dict(model_config.get('fsdp_config') or {}), + } + + +def validate_context_batch_config( + context_key: str, + *, + rollout_groups: int, + num_generations: int, + train: TrainBatchConfig, + sampler_dp: int, + model_dp: int, +) -> None: + values = { + 'rollout.batch_size': rollout_groups, + 'rollout.num_generations': num_generations, + 'train.mini_batch_size': train.mini_batch_size, + 'train.micro_batch_size': train.micro_batch_size, + } + for name, value in values.items(): + if value <= 0: + raise ValueError(f'{name} for {context_key} must be positive, got {value}') + if rollout_groups % sampler_dp: + raise ValueError(f'rollout.batch_size for {context_key} must be divisible by sampler DP size ' + f'({sampler_dp}), got {rollout_groups}') + partition_samples = rollout_groups * num_generations + if partition_samples % train.mini_batch_size: + raise ValueError(f'partition for {context_key} has {partition_samples} samples and must be divisible by ' + f'train.mini_batch_size={train.mini_batch_size}') + if train.mini_batch_size % num_generations: + raise ValueError(f'train.mini_batch_size for {context_key} must preserve complete prompt groups: ' + f'{train.mini_batch_size} % {num_generations} != 0') + if train.mini_batch_size % model_dp: + raise ValueError(f'train.mini_batch_size for {context_key} must be divisible by ' + f'model DP size {model_dp}') + samples_per_rank = train.mini_batch_size // model_dp + if train.micro_batch_size > samples_per_rank: + raise ValueError(f'train.micro_batch_size for {context_key} must not exceed the per-rank train batch ' + f'({samples_per_rank}), got {train.micro_batch_size}') + if train.dynamic_batching: + if train.max_tokens_per_micro_batch is None or train.max_tokens_per_micro_batch <= 0: + raise ValueError(f'train.max_tokens_per_micro_batch for {context_key} must be positive when ' + 'train.dynamic_batching=true') + if train.packing_algorithm not in ('ffd', 'kk'): + raise ValueError(f'train.packing_algorithm for {context_key} must be ffd or kk, ' + f'got {train.packing_algorithm!r}') + + +def configure_lora_lr_scheduler( + model: Any, + adapter_name: str, + lora_config: Mapping[str, Any], +) -> None: + scheduler_config = lora_config.get('lr_scheduler') + if scheduler_config is None: + return + scheduler_config = dict(scheduler_config) + scheduler_cls = scheduler_config.pop('cls') + model.set_lr_scheduler( + scheduler_cls, + adapter_name=adapter_name, + **scheduler_config, + ) + + +def resolve_context_learning_rate( + train_config: Mapping[str, Any], + lora_defaults: Mapping[str, Any], +) -> float: + configured = train_config.get('learning_rate', lora_defaults.get('learning_rate')) + if configured is None: + raise ValueError('train.learning_rate or lora.learning_rate must be configured') + learning_rate = float(configured) + if not math.isfinite(learning_rate) or learning_rate <= 0: + raise ValueError(f'train.learning_rate must be a positive finite value, got {configured!r}') + return learning_rate + + +def resolve_context_lora_target_modules( + context_config: Mapping[str, Any], + lora_defaults: Mapping[str, Any], +) -> str | list[str]: + context_lora_config = dict(context_config.get('lora') or {}) + target_modules = context_lora_config.get( + 'target_modules', + lora_defaults.get('target_modules', 'all-linear'), + ) + if isinstance(target_modules, str): + if not target_modules: + raise ValueError('lora.target_modules must not be empty') + return target_modules + if isinstance(target_modules, Sequence) and target_modules: + modules = list(target_modules) + if all(isinstance(module, str) and module for module in modules): + return modules + raise ValueError('lora.target_modules must be a non-empty string or sequence of module names, ' + f'got {target_modules!r}') + + +def resolve_context_loss_config( + context_config: Mapping[str, Any], + loss_defaults: Mapping[str, Any] | None = None, +) -> tuple[str, dict[str, Any]]: + loss_config: dict[str, Any] = { + 'cls': 'GRPOLoss', + 'epsilon': 0.2, + } + loss_config.update(dict(loss_defaults or {})) + loss_config.update(dict(context_config.get('loss') or {})) + + loss_cls = loss_config.pop('cls', None) + if not isinstance(loss_cls, str) or not loss_cls: + raise ValueError(f'loss.cls must be a non-empty string, got {loss_cls!r}') + return loss_cls, loss_config diff --git a/src/twinkle_agentic/async_rl/vllm_sampler_tq.py b/src/twinkle_agentic/async_rl/vllm_sampler_tq.py new file mode 100644 index 000000000..4f2d0342d --- /dev/null +++ b/src/twinkle_agentic/async_rl/vllm_sampler_tq.py @@ -0,0 +1,781 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from __future__ import annotations + +import asyncio +import json +import os +import re +import time +import uuid +from concurrent.futures import Future +from copy import copy +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +from twinkle import DeviceMesh, get_logger, remote_class, remote_function +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 .data_plane import TQDataPlane +from .metrics import rollout_metrics +from .types import LoraContext, PromptGroup, RolloutOutput, RolloutPolicy +from .utils import resolve_adapter_path, sample_responses_to_rollout_rows + +logger = get_logger() + + +def _dispatch_generation( + worker_count: int, + worker_index: int, + args: tuple[Any, ...], + kwargs: dict[str, Any], + **_dispatch_kwargs, +) -> tuple[tuple[Any, ...], dict[str, Any]]: + """Slice CS inputs while allowing a prompt count smaller than DP size.""" + sliced_args = list(args) + sliced_kwargs = dict(kwargs) + if len(sliced_args) > 1: + inputs = sliced_args[1] + target = ('args', 1) + elif 'inputs' in sliced_kwargs: + inputs = sliced_kwargs['inputs'] + target = ('kwargs', 'inputs') + else: + raise ValueError('submit_generation requires inputs') + + input_list = list(inputs) if isinstance(inputs, (list, tuple)) else [inputs] + size, remainder = divmod(len(input_list), worker_count) + start = worker_index * size + min(worker_index, remainder) + stop = (worker_index + 1) * size + min(worker_index + 1, remainder) + shard = input_list[start:stop] + if target[0] == 'args': + sliced_args[target[1]] = shard + else: + sliced_kwargs[target[1]] = shard + return tuple(sliced_args), sliced_kwargs + + +def _path_component(value: str) -> str: + return re.sub(r'[^A-Za-z0-9._-]+', '_', value).strip('._') or 'unknown' + + +def _compute_rewards( + reward_registry: dict[str, Any], + context: LoraContext, + rollout_rows: list[RolloutOutput], +) -> list[float] | None: + reward_fn = reward_registry.get(context.key) + if reward_fn is None: + return None + return list(reward_fn(rollout_rows, context=context)) + + +def _compute_reward_metrics( + reward_registry: dict[str, Any], + context: LoraContext, + rollout_rows: list[RolloutOutput], + rewards: list[float], +) -> dict[str, Any]: + reward_fn = reward_registry.get(context.key) + metric_payload = getattr(reward_fn, 'metric_payload', None) + if metric_payload is None: + return {} + return dict(metric_payload(rollout_rows, rewards=rewards, context=context)) + + +@dataclass(frozen=True) +class _GeneratedSample: + response: SampleResponse + policies: tuple[RolloutPolicy, ...] + attempts: int + was_aborted: bool + resumed_partial_output: bool + + @property + def initial_policy(self) -> RolloutPolicy: + return self.policies[0] + + @property + def final_policy(self) -> RolloutPolicy: + return self.policies[-1] + + @property + def retry_count(self) -> int: + return self.attempts - 1 + + +@dataclass(frozen=True) +class _PromptGroupRolloutStats: + completion_lengths: tuple[int, ...] + stop_reasons: tuple[str | None, ...] + policy_versions: tuple[int, ...] + + +@remote_class() +class VLLMSamplerTQ(vLLMSampler): + """vLLM sampler that writes async RL rollout results directly to TransferQueue. + + ``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__( + self, + model_id: str, + engine_args: dict[str, Any] | None = None, + device_mesh: DeviceMesh | None = None, + *, + context_manager: Any | None = None, + reward_registry: dict[str, Any] | None = None, + rollout_max_retries: int = 2, + rollout_retry_delay_s: float = 0.5, + rollout_output_dir: str | None = None, + rollout_output_include_token_ids: bool = False, + **kwargs, + ): + self.context_manager = context_manager + super().__init__(model_id=model_id, engine_args=engine_args, device_mesh=device_mesh, **kwargs) + # Native YAML async-RL writes rollout groups to TransferQueue. The C/S + # component mode only uses submit_generation/collect_generation and + # stores results through the server DataPlane deployment. + self.data_plane = TQDataPlane() if context_manager is not None else None + self.reward_registry = dict(reward_registry or {}) + self.rollout_max_retries = int(rollout_max_retries) + self.rollout_retry_delay_s = float(rollout_retry_delay_s) + self.rollout_output_dir = ( + Path(rollout_output_dir).expanduser().resolve() if rollout_output_dir is not None else None) + self.rollout_output_include_token_ids = bool(rollout_output_include_token_ids) + if self.rollout_max_retries < 0: + raise ValueError(f'rollout_max_retries must be non-negative, got {self.rollout_max_retries}') + if self.rollout_retry_delay_s < 0: + raise ValueError(f'rollout_retry_delay_s must be non-negative, got {self.rollout_retry_delay_s}') + self._background_submissions: dict[str, Future] = {} + # Generation submissions are used by the client-orchestrated server + # 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.metric_buffer = MetricBuffer() + self._failure: str | None = None + + def _record_metrics( + self, + group: PromptGroup, + values: dict[str, Any], + *, + status: str = 'completed', + attributes: dict[str, Any] | None = None, + policy_version: int | None = None, + ) -> None: + self.metric_buffer.record( + MetricRecord( + stage='rollout', + values=dict(values), + context_key=group.context.key, + partition_id=group.partition_id, + partition_index=group.partition.step, + policy_version=policy_version, + status=status, + attributes=dict(attributes or {}), + )) + + @remote_function(dispatch='all', collect='flatten', lazy_collect=False) + def drain_metric_records(self) -> list[MetricRecord]: + return self.metric_buffer.drain() + + @remote_function(dispatch='all', collect='none', lazy_collect=False) + def check_health(self) -> None: + if self._failure is not None: + raise RuntimeError(self._failure) + + @remote_function(dispatch='all', collect='none', lazy_collect=False) + def register_reward(self, context_key: str, reward: Any) -> None: + if context_key in self.reward_registry: + raise KeyError(f'reward already registered for {context_key}') + self.reward_registry[context_key] = reward + + @remote_function(dispatch='all', collect='none', lazy_collect=False) + def unregister_reward(self, context_key: str) -> None: + self.reward_registry.pop(context_key, None) + + @remote_function(dispatch='slice_dp', collect='none', lazy_collect=False) + def submit_prompt_groups( + self, + groups: list[PromptGroup], + sampling_params: SamplingParams, + allow_partial_rollout: bool = False, + ) -> dict[str, Any]: + """Schedule this DP worker's complete prompt groups and return immediately.""" + if self.context_manager is None: + raise RuntimeError('context_manager is required for native TQ prompt-group sampling') + submission_id = str(uuid.uuid4()) + submitted_at = time.perf_counter() + future = self._submit_in_loop( + self._sample_prompt_groups( + submission_id, + groups, + sampling_params, + bool(allow_partial_rollout), + submitted_at, + )) + self._background_submissions[submission_id] = future + future.add_done_callback(self._on_submission_done(submission_id)) + return { + 'submission_id': submission_id, + 'submitted_prompt_groups': len(groups), + 'submitted_samples': sum(group.num_samples for group in groups), + } + + @remote_function(dispatch=_dispatch_generation, collect='none', lazy_collect=False) + def submit_generation( + self, + submission_id: str, + inputs: Any, + sampling_params: SamplingParams | dict[str, Any] | None = None, + adapter_name: str = '', + adapter_path: str | None = None, + *, + use_base_model: bool = False, + ) -> dict[str, Any]: + """Submit a CS sampling shard without blocking the Ray actor. + + The generated responses stay local to this DP worker until + :meth:`collect_generation` consumes them. This gives the HTTP + service the same fast-admission property as the native TQ rollout path + without exposing PromptGroup or BatchMeta to the client. + """ + if submission_id in self._generation_submissions: + raise KeyError(f'generation submission already exists: {submission_id}') + 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, + )) + 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.""" + 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}', + } + if future.cancelled(): + return {'submission_id': submission_id, 'status': 'cancelled'} + if not future.done(): + return {'submission_id': submission_id, 'status': 'running'} + error = future.exception() + if error is not None: + return { + 'submission_id': submission_id, + 'status': 'failed', + 'error': f'{type(error).__name__}: {error}', + } + return {'submission_id': submission_id, 'status': 'completed'} + + @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: + raise KeyError(f'unknown generation submission: {submission_id}') + if not future.done(): + raise RuntimeError(f'generation submission is still running: {submission_id}') + try: + return future.result() + finally: + self._generation_submissions.pop(submission_id, None) + + @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: + 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, + } + + @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} + + @remote_function(dispatch='slice_dp', collect='flatten', lazy_collect=False) + def evaluate( + self, + inputs: list[dict[str, Any]], + sampling_params: SamplingParams, + adapter_name: str, + adapter_path: str, + ) -> list[SampleResponse]: + """Synchronously evaluate one adapter without writing results to TQ.""" + return super().sample( + inputs, + sampling_params, + adapter_name=adapter_name, + adapter_path=adapter_path, + ) + + def _submit_in_loop(self, coro) -> Future: + return asyncio.run_coroutine_threadsafe(coro, self._async_loop) + + async def _generate_inputs( + self, + inputs: Any, + sampling_params: SamplingParams | dict[str, Any] | None, + *, + adapter_name: str, + adapter_path: str | None, + use_base_model: bool, + ) -> list[SampleResponse]: + """Asynchronous counterpart of ``vLLMSampler.sample`` for CS use.""" + if sampling_params is None: + sampling_params = SamplingParams() + elif isinstance(sampling_params, dict): + sampling_params = SamplingParams.from_dict(sampling_params) + + inputs_list = self._normalize_inputs(inputs) + if not inputs_list: + return [] + + is_trajectory = 'input_ids' not in inputs_list[0] + logprobs_only = False + if sampling_params.max_tokens == 0: + sampling_params = copy(sampling_params) + sampling_params.max_tokens = 1 + logprobs_only = True + + multi_modal_data_list = [self._extract_multi_modal_data(feat) for feat in inputs_list] + if is_trajectory: + if self.template is None: + raise ValueError('Use set_template to add a template when trying to input Trajectory') + encoded_inputs = [ + self.encode_trajectory_for_vllm(trajectory, adapter_name, not logprobs_only) + for trajectory in inputs_list + ] + else: + encoded_inputs = inputs_list + + lora_request = None + if adapter_path is not None: + logger.info(f'Loading LoRA from {adapter_path}') + local_adapter_path = HubOperation.download_model(model_id_or_path=adapter_path) + lora_request = await self.engine._get_or_load_lora(local_adapter_path) + if lora_request is None: + 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))) + + def _on_submission_done(self, submission_id: str): + + def callback(future: Future) -> None: + self._background_submissions.pop(submission_id, None) + error = future.exception() + if error is not None: + self._failure = f'{type(error).__name__}: {error}' + logger.warning('VLLMSamplerTQ background submission failed: submission=%s error=%s', submission_id, + error) + + return callback + + async def _sample_prompt_groups( + self, + submission_id: str, + groups: list[PromptGroup], + sampling_params: SamplingParams, + allow_partial_rollout: bool, + submitted_at: float, + ) -> None: + results = await asyncio.gather( + *(self._run_prompt_group( + submission_id=submission_id, + group=group, + sampling_params=sampling_params, + allow_partial_rollout=allow_partial_rollout, + ) for group in groups), + return_exceptions=True) + failed_group = next( + ((group, result) for group, result in zip(groups, results) if isinstance(result, Exception)), None) + if failed_group is not None: + group, error = failed_group + self._record_metrics( + group, + {}, + status='failed', + attributes={ + 'scope': 'group', + 'group_id': group.group_id, + 'error': str(error) + }, + ) + raise RuntimeError(f'rollout failed for {group.group_id}: {error}') from error + + rollout_stats = [result for result in results if isinstance(result, _PromptGroupRolloutStats)] + metric_rows = [{ + 'completion_length': completion_length, + 'stop_reason': stop_reason, + } for stats in rollout_stats + for completion_length, stop_reason in zip(stats.completion_lengths, stats.stop_reasons)] + policy_versions = [version for stats in rollout_stats for version in stats.policy_versions] + first_group = groups[0] + dp_size = self.device_mesh.dp_world_size or 1 + self._record_metrics( + first_group, + { + 'prompt_group_count': + len(groups), + **rollout_metrics( + completion_lengths=[row['completion_length'] for row in metric_rows], + stop_reasons=[row['stop_reason'] for row in metric_rows], + rollout_latency_s=time.perf_counter() - submitted_at, + ), + 'policy_version_min': + min(policy_versions), + 'policy_version_max': + max(policy_versions), + 'sampler_dp_size': + dp_size, + }, + attributes={'scope': 'partition' if dp_size == 1 else 'shard'}, + policy_version=max(policy_versions), + ) + + async def _run_prompt_group( + self, + *, + submission_id: str, + group: PromptGroup, + sampling_params: SamplingParams, + allow_partial_rollout: bool, + ) -> _PromptGroupRolloutStats: + """Sample all generations for one group, then write that group once.""" + started = asyncio.get_running_loop().time() + num_generations = group.num_samples + sources = [{ + **group.prompt, 'group_id': group.group_id, + 'generation_idx': generation_idx + } for generation_idx in range(num_generations)] + generated_samples = await self._generate_group_samples( + group.context, sources, sampling_params, allow_partial_rollout=allow_partial_rollout) + rows = [] + for source, generated in zip(sources, generated_samples): + sample_rows = sample_responses_to_rollout_rows([source], [generated.response], + policy_version=generated.final_policy.version) + if len(sample_rows) != 1: + raise ValueError(f'generation {source["generation_idx"]} produced {len(sample_rows)} samples') + row = sample_rows[0] + versions = [policy.version for policy in generated.policies] + row.update({ + 'rollout_policy_version': generated.final_policy.version, + 'rollout_adapter_path': generated.final_policy.adapter_path, + 'rollout_policy_versions': versions, + 'initial_policy_version': generated.initial_policy.version, + 'final_policy_version': generated.final_policy.version, + 'policy_version_span': generated.final_policy.version - generated.initial_policy.version, + }) + rows.append(row) + if len(rows) != num_generations: + raise ValueError(f'group {group.group_id} expected {num_generations} rollout samples, got {len(rows)}') + + rewards = _compute_rewards(self.reward_registry, group.context, rows) + if rewards is None: + raise ValueError(f'no reward function registered for context {group.context.key}') + reward_metrics = _compute_reward_metrics(self.reward_registry, group.context, rows, rewards) + if self.data_plane is None: + raise RuntimeError('native TQ data plane is required for prompt-group sampling') + await self.data_plane.complete_rollout_group( + group, + rollout_rows=rows, + rewards=rewards, + submission_id=submission_id, + tag_metrics=reward_metrics, + ) + + rollout_latency_s = asyncio.get_running_loop().time() - started + policy_versions = [policy.version for sample in generated_samples for policy in sample.policies] + if self.rollout_output_dir is not None: + try: + await asyncio.to_thread( + self._write_rollout_group, + submission_id, + group, + generated_samples, + rows, + rewards, + ) + except Exception as error: + logger.warning('Failed to write rollout output for %s: %s', group.group_id, error) + self._record_metrics( + group, + { + **rollout_metrics( + rewards={'reward': rewards}, + completion_lengths=[int(row['completion_length']) for row in rows], + stop_reasons=[row.get('stop_reason') for row in rows], + rollout_latency_s=rollout_latency_s, + ), + 'retry_count': + sum(sample.retry_count for sample in generated_samples), + 'aborted_sample_count': + sum(sample.was_aborted for sample in generated_samples), + 'partial_resumed_sample_count': + sum(sample.resumed_partial_output for sample in generated_samples), + 'policy_version_min': + min(policy_versions), + 'policy_version_max': + max(policy_versions), + **reward_metrics, + }, + attributes={ + 'scope': 'group', + 'group_id': group.group_id + }, + policy_version=max(policy_versions), + ) + return _PromptGroupRolloutStats( + completion_lengths=tuple(int(row['completion_length']) for row in rows), + stop_reasons=tuple(row.get('stop_reason') for row in rows), + policy_versions=tuple(policy_versions), + ) + + def _write_rollout_group( + self, + submission_id: str, + group: PromptGroup, + generated_samples: list[_GeneratedSample], + rows: list[RolloutOutput], + rewards: list[float], + ) -> None: + policy_version = max(int(row['rollout_policy_version']) for row in rows) + partition_name = _path_component(group.partition_id.rsplit('/', 1)[-1]) + group_name = _path_component(group.group_id.rsplit('/', 1)[-1]) + output_dir = self.rollout_output_dir.joinpath( + _path_component(group.context.tenant_id), + _path_component(group.context.training_run_id), + _path_component(group.context.adapter_name), + f'policy_{policy_version}', + ) + output_dir.mkdir(parents=True, exist_ok=True) + output_path = output_dir / f'{partition_name}-{group_name}.jsonl' + temporary_path = output_path.with_suffix(f'.jsonl.{uuid.uuid4().hex}.tmp') + ground_truth = user_data_get(group.prompt.get('user_data'), 'ground_truth') + + with temporary_path.open('w', encoding='utf-8') as stream: + for generated, row, reward in zip(generated_samples, rows, rewards): + response = generated.response + sequence = response.sequences[0] + prompt_token_ids = list(response.prompt_token_ids or []) + completion_token_ids = list(sequence.tokens) + record = { + 'submission_id': submission_id, + 'context_key': group.context.key, + 'tenant_id': group.context.tenant_id, + 'training_run_id': group.context.training_run_id, + 'adapter_name': group.context.adapter_name, + 'partition_id': group.partition_id, + 'group_id': group.group_id, + 'sample_idx': int(row['generation_idx']), + 'seqlen': len(prompt_token_ids) + len(completion_token_ids), + 'prompt_len': len(prompt_token_ids), + 'completion_len': len(completion_token_ids), + 'head_version': int(row['initial_policy_version']), + 'tail_version': int(row['final_policy_version']), + 'policy_versions': list(row['rollout_policy_versions']), + 'adapter_path': row.get('rollout_adapter_path'), + 'reward': float(reward), + 'ground_truth': ground_truth, + 'stop_reason': row.get('stop_reason'), + 'retry_count': generated.retry_count, + 'was_aborted': generated.was_aborted, + 'resumed_partial_output': generated.resumed_partial_output, + 'prompt': self.template.decode(prompt_token_ids, skip_special_tokens=False), + 'completion': self.template.decode(completion_token_ids, skip_special_tokens=False), + } + if self.rollout_output_include_token_ids: + record.update({ + 'prompt_token_ids': prompt_token_ids, + 'completion_token_ids': completion_token_ids, + 'logprobs': list(row['logprobs']), + }) + stream.write(json.dumps(record, ensure_ascii=False, default=str) + '\n') + os.replace(temporary_path, output_path) + + async def _load_lora_for_policy(self, policy: RolloutPolicy) -> Any: + """Load the adapter selected for one group's rollout snapshot.""" + if policy.adapter_path is None: + return None + local_path = await asyncio.to_thread(resolve_adapter_path, policy.adapter_path) + lora_request = await self.engine._get_or_load_lora(local_path) + if lora_request is None: + raise RuntimeError(f'failed to load LoRA adapter from {local_path}') + return lora_request + + async def _generate_group_samples( + self, + context: LoraContext, + sources: list[dict[str, Any]], + sampling_params: SamplingParams, + *, + allow_partial_rollout: bool, + ) -> list[_GeneratedSample]: + logprobs_only = False + if sampling_params.max_tokens == 0: + sampling_params = copy(sampling_params) + sampling_params.max_tokens = 1 + logprobs_only = True + + is_trajectory = 'input_ids' not in sources[0] + multi_modal_data_list = [self._extract_multi_modal_data(source) for source in sources] + if is_trajectory: + template = self.template + assert template is not None, 'Use set_template before sampling trajectories' + encoded_inputs = [ + self.encode_trajectory_for_vllm(source, context.adapter_name, not logprobs_only) for source in sources + ] + else: + encoded_inputs = sources + tasks = [ + self._generate_sample( + context, + feat, + sampling_params, + multi_modal_data=multi_modal_data, + logprobs_only=logprobs_only, + allow_partial_rollout=allow_partial_rollout, + ) for feat, multi_modal_data in zip(encoded_inputs, multi_modal_data_list) + ] + results = await asyncio.gather(*tasks, return_exceptions=True) + failures = [result for result in results if isinstance(result, Exception)] + if failures: + raise failures[0] + return results + + async def _generate_sample( + self, + context: LoraContext, + original_input: dict[str, Any], + sampling_params: SamplingParams, + *, + multi_modal_data: dict[str, Any] | None, + logprobs_only: bool, + allow_partial_rollout: bool, + ) -> _GeneratedSample: + current_input = original_input + partial_responses: list[SampleResponse] = [] + partial_policies: list[RolloutPolicy] = [] + generated_tokens = 0 + last_error: Exception | None = None + was_aborted = False + resumed_partial_output = False + + for attempt in range(self.rollout_max_retries + 1): + policy = await self.context_manager.acquire_rollout_policy.remote(context) + attempt_params = copy(sampling_params) + if allow_partial_rollout and attempt_params.max_tokens is not None: + attempt_params.max_tokens -= generated_tokens + try: + try: + response = await self._sample_single( + current_input, + attempt_params, + lora_request=await self._load_lora_for_policy(policy), + multi_modal_data=multi_modal_data, + logprobs_only=logprobs_only, + ) + sequence = response.sequences[0] + except Exception as exc: + last_error = exc + else: + if sequence.stop_reason not in {'abort', 'error'}: + if not allow_partial_rollout or not partial_responses: + return _GeneratedSample(response, (policy, ), attempt + 1, was_aborted, + resumed_partial_output) + partial_responses.append(response) + partial_policies.append(policy) + return _GeneratedSample( + self._merge_partial_responses(partial_responses), tuple(partial_policies), attempt + 1, + was_aborted, resumed_partial_output) + + last_error = RuntimeError(f'generation stopped with {sequence.stop_reason}') + was_aborted = was_aborted or sequence.stop_reason == 'abort' + if allow_partial_rollout and sequence.tokens: + resumed_partial_output = True + partial_responses.append(response) + partial_policies.append(policy) + generated_tokens += len(sequence.tokens) + current_input = sequence.new_input_feature + if sampling_params.max_tokens is not None and generated_tokens >= sampling_params.max_tokens: + return _GeneratedSample( + self._merge_partial_responses(partial_responses, stop_reason='length'), + tuple(partial_policies), + attempt + 1, + was_aborted, + resumed_partial_output, + ) + elif not allow_partial_rollout: + current_input = original_input + finally: + await self.context_manager.release_rollout_policy.remote(policy) + + if attempt < self.rollout_max_retries: + await asyncio.sleep(self.rollout_retry_delay_s) + + error_detail = f'{type(last_error).__name__}: {last_error}' + error = RuntimeError( + f'generation failed after {self.rollout_max_retries + 1} attempts; last error: {error_detail}') + raise error from last_error + + def _merge_partial_responses( + self, + responses: list[SampleResponse], + *, + stop_reason: str | None = None, + ) -> SampleResponse: + sequences = [response.sequences[0] for response in responses] + tokens = [token for sequence in sequences for token in sequence.tokens] + logprobs = [logprob for sequence in sequences for logprob in (sequence.logprobs or [])] + final_sequence = sequences[-1] + return SampleResponse( + prompt_token_ids=responses[0].prompt_token_ids, + sequences=[ + SampledSequence( + stop_reason=stop_reason or final_sequence.stop_reason, + tokens=tokens, + logprobs=logprobs, + decoded=self.template.decode(tokens), + new_input_feature=final_sequence.new_input_feature, + routed_experts=final_sequence.routed_experts, + ) + ], + ) diff --git a/src/twinkle_agentic/async_rl/workers.py b/src/twinkle_agentic/async_rl/workers.py new file mode 100644 index 000000000..984190c8f --- /dev/null +++ b/src/twinkle_agentic/async_rl/workers.py @@ -0,0 +1,676 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Long-lived, context-scheduling async-RL workers.""" + +from __future__ import annotations + +import asyncio +import os +import shutil +import time +from collections import defaultdict +from collections.abc import Callable, Iterable, Sequence +from typing import Any + +from twinkle.metric import MetricBuffer, MetricRecord +from .context_manager import ContextStatus, LoraContextManager +from .data_plane import TQDataPlane +from .metrics import advantage_signal_metrics, training_policy_metrics +from .scheduler import ContextScheduler, ScheduleCandidate, SchedulerConfig +from .types import LoraContext, PartitionAdmission + + +class _Worker: + """A long-lived Ray service with one privately owned background loop.""" + + def __init__(self): + self._service_task: asyncio.Task[None] | None = None + self._stop_requested = False + self._failure: str | None = None + self.metric_buffer = MetricBuffer() + + async def start(self) -> None: + if self._service_task is not None and not self._service_task.done(): + return + self._stop_requested = False + self._failure = None + self._service_task = asyncio.create_task(self._run_service()) + + async def stop(self) -> None: + self._stop_requested = True + if self._service_task is None or self._service_task.done(): + return + self._service_task.cancel() + try: + await self._service_task + except asyncio.CancelledError: + pass + + async def get_service_state(self) -> dict[str, str | bool | None]: + return { + 'running': self._service_task is not None and not self._service_task.done(), + 'failure': self._failure, + } + + def drain_metric_records(self) -> list[MetricRecord]: + return self.metric_buffer.drain() + + def _record_metric( + self, + stage: str, + *, + context: LoraContext | None = None, + admission: PartitionAdmission | None = None, + partition_id: str | None = None, + values: dict[str, Any] | None = None, + status: str = 'completed', + attributes: dict[str, Any] | None = None, + optimizer_step: int | None = None, + policy_version: int | None = None, + ) -> None: + self.metric_buffer.record( + MetricRecord( + stage=stage, + values=dict(values or {}), + context_key=context.key if context is not None else None, + partition_id=admission.partition_id if admission is not None else partition_id, + partition_index=admission.step if admission is not None else None, + optimizer_step=optimizer_step, + policy_version=policy_version, + status=status, + attributes=dict(attributes or {}), + )) + + async def _run_service(self) -> None: + try: + await self._serve() + except asyncio.CancelledError: + return + except Exception as exc: + self._failure = f'{type(exc).__name__}: {exc}' + + async def _serve(self) -> None: + raise NotImplementedError + + +class RolloutWorker(_Worker): + """Admits full prompt batches and submits them to the sampler without waiting.""" + + def __init__(self, + *, + context_manager: LoraContextManager, + data_plane: TQDataPlane, + sampler: Any, + prompt_batches: dict[str, Iterable[Sequence[dict[str, Any]]] + | Callable[[], Iterable[Sequence[dict[str, Any]]]]], + rollout_config: dict[str, dict[str, Any]], + scheduler: SchedulerConfig, + allow_partial_rollout: bool = False, + persistent: bool = False, + idle_delay_s: float = 0.05): + super().__init__() + self.data_plane = data_plane + self.sampler = sampler + self.context_manager = context_manager + self.idle_delay_s = idle_delay_s + self.rollout_config = rollout_config + self.scheduler = ContextScheduler(scheduler) + self.allow_partial_rollout = allow_partial_rollout + self.persistent = persistent + self._prompt_batch_iterators = { + key: iter(value() if callable(value) else value) + for key, value in prompt_batches.items() + } + self._next_batch_tasks: dict[str, asyncio.Task[Sequence[dict[str, Any]] | None]] = {} + self._exhausted: set[str] = set() + self._contexts_changed = asyncio.Event() + + async def register_context( + self, + context: LoraContext, + prompt_batches: Iterable[Sequence[dict[str, Any]]] | Callable[[], Iterable[Sequence[dict[str, Any]]]], + rollout_config: dict[str, Any], + ) -> None: + key = context.key + if key in self._prompt_batch_iterators: + raise KeyError(f'rollout context already exists: {key}') + self._prompt_batch_iterators[key] = iter(prompt_batches() if callable(prompt_batches) else prompt_batches) + self.rollout_config[key] = dict(rollout_config) + self._exhausted.discard(key) + if self._service_task is not None and not self._service_task.done(): + self._start_next_batch(key) + self._contexts_changed.set() + + async def unregister_context(self, context: LoraContext | str) -> None: + key = context if isinstance(context, str) else context.key + task = self._next_batch_tasks.pop(key, None) + if task is not None: + task.cancel() + await asyncio.gather(task, return_exceptions=True) + iterator = self._prompt_batch_iterators.pop(key, None) + self.rollout_config.pop(key, None) + self._exhausted.discard(key) + close = getattr(iterator, 'close', None) + if callable(close): + await asyncio.to_thread(close) + self._contexts_changed.set() + + def _start_next_batch(self, key: str) -> None: + if key in self._prompt_batch_iterators and key not in self._next_batch_tasks: + self._next_batch_tasks[key] = asyncio.create_task( + asyncio.to_thread(next, self._prompt_batch_iterators[key], None)) + + async def stop(self) -> None: + await super().stop() + pending = list(self._next_batch_tasks.values()) + self._next_batch_tasks.clear() + for task in pending: + task.cancel() + if pending: + await asyncio.gather(*pending, return_exceptions=True) + + async def _serve(self) -> None: + for key in self._prompt_batch_iterators: + self._start_next_batch(key) + while not self._stop_requested: + if not self.persistent and await self.context_manager.is_rollout_admission_closed.remote(): + return + candidates = [] + for key in list(self._prompt_batch_iterators): + if key in self._exhausted: + continue + status = await self.context_manager.context_status.remote(key) + if status is not ContextStatus.ACTIVE: + if status in (ContextStatus.EXHAUSTED, ContextStatus.FINISHED): + self._exhausted.add(key) + continue + config = self.rollout_config[key] + task = self._next_batch_tasks.get(key) + if task is None: + self._start_next_batch(key) + task = self._next_batch_tasks.get(key) + if task is not None and task.done(): + candidates.append(ScheduleCandidate(config['context'])) + candidate = self.scheduler.choose(candidates) + if candidate is None: + if not self.persistent and len(self._exhausted) == len(self._prompt_batch_iterators): + return + self._contexts_changed.clear() + try: + await asyncio.wait_for(self._contexts_changed.wait(), timeout=self.idle_delay_s) + except TimeoutError: + pass + continue + key = candidate.context.key + config = self.rollout_config[key] + batch_task = self._next_batch_tasks[key] + try: + batch = batch_task.result() + except Exception as exc: + self._next_batch_tasks.pop(key) + self._record_metric( + 'rollout', + context=candidate.context, + status='failed', + attributes={'error': f'prompt loading failed: {exc}'}, + ) + raise RuntimeError(f'prompt loading failed for {key}: {exc}') from exc + if batch is None or len(batch) != int(config['batch_size']): + self._next_batch_tasks.pop(key) + self._exhausted.add(key) + await self.context_manager.on_dataset_exhausted.remote(candidate.context) + self.scheduler.on_blocked(candidate) + continue + admission = await self.context_manager.request_rollout_partition.remote( + candidate.context, + target_groups=len(batch), + num_generations=int(config['num_generations']), + ) + if admission is None: + self.scheduler.on_blocked(candidate) + await asyncio.sleep(self.idle_delay_s) + continue + self._next_batch_tasks.pop(key) + submission_started = time.perf_counter() + try: + prepared = await self.data_plane.prepare_rollout_partition( + admission, + list(batch), + config['sampling_params'], + ) + await asyncio.to_thread( + self.sampler.submit_prompt_groups, + list(prepared.groups), + prepared.sampling_params, + self.allow_partial_rollout, + ) + except Exception as exc: + self._record_metric( + 'rollout', + context=admission.context, + admission=admission, + status='failed', + attributes={'error': str(exc)}, + ) + raise RuntimeError(f'rollout submission failed for {admission.partition_id}: {exc}') from exc + self._start_next_batch(key) + self.scheduler.on_success(candidate) + self._record_metric( + 'rollout', + context=admission.context, + admission=admission, + status='submitted', + values={ + 'prompt_count': admission.target_groups, + 'sample_count': admission.sample_count, + 'num_generations': admission.num_generations, + 'rollout_submission_latency_s': time.perf_counter() - submission_started, + }, + attributes={'scope': 'partition'}, + ) + + +class AdvantageWorker(_Worker): + + def __init__(self, + *, + context_manager: LoraContextManager, + data_plane: TQDataPlane, + advantage_fn: Callable[[Any, PartitionAdmission], tuple[Sequence[float], Sequence[float]]], + scheduler: SchedulerConfig, + persistent: bool = False, + idle_delay_s: float = 0.05): + super().__init__() + self.data_plane = data_plane + self.context_manager = context_manager + self.idle_delay_s = idle_delay_s + self.advantage_fn = advantage_fn + self.scheduler = ContextScheduler(scheduler) + self.persistent = persistent + + async def _serve(self) -> None: + while not self._stop_requested: + if not self.persistent and await self.context_manager.is_run_finished.remote(): + return + admissions = await self.context_manager.list_live_partitions.remote() + blocked: set[str] = set() + progressed = False + for _ in range(len(admissions)): + candidates = [ + ScheduleCandidate(admission.context, admission) for admission in admissions + if admission.partition_id not in blocked + ] + candidate = self.scheduler.choose(candidates) + if candidate is None: + break + admission = candidate.partition + batch = await self.data_plane.claim_advantage_batch(admission, 1) + if batch is None: + blocked.add(admission.partition_id) + self.scheduler.on_blocked(candidate) + continue + started = time.perf_counter() + try: + advantages, returns = self.advantage_fn(batch.data, admission) + await self.data_plane.write_advantages(batch, advantages=advantages, returns=returns) + except Exception as exc: + self._record_metric( + 'advantage', + context=admission.context, + admission=admission, + status='failed', + attributes={'error': str(exc)}, + ) + raise RuntimeError(f'advantage failed for {admission.partition_id}: {exc}') from exc + self.scheduler.on_success(candidate) + policy = await self.context_manager.get_rollout_policy.remote(admission.context) + advantage_metrics = advantage_signal_metrics( + batch.data['rewards'], + advantages, + num_generations=admission.num_generations, + ) + advantage_metrics.update({ + 'sample_count': len(advantages), + 'advantage_latency_s': time.perf_counter() - started, + }) + self._record_metric( + 'advantage', + context=admission.context, + admission=admission, + values=advantage_metrics, + policy_version=policy.version, + ) + progressed = True + break + if not progressed: + await asyncio.sleep(self.idle_delay_s) + + +class TrainerWorker(_Worker): + + def __init__(self, + *, + context_manager: LoraContextManager, + data_plane: TQDataPlane, + train_fn: Callable[[Any, PartitionAdmission], dict[str, Any] | None], + save_adapter: Callable[[PartitionAdmission], str], + mini_batch_sizes: dict[str, int], + scheduler: SchedulerConfig, + train_with_config_fn: Callable[[Any, PartitionAdmission, Any], dict[str, Any] | None] | None = None, + train_batch_configs: dict[str, Any] | None = None, + keep_adapter_versions: int = 0, + initial_adapter_paths: dict[str, str] | None = None, + remove_adapter: Callable[[str], None] | None = None, + evaluation_config: dict[str, dict[str, Any]] | None = None, + evaluate_batch: Callable[[Sequence[dict[str, Any]], PartitionAdmission, str, int, Any], dict[str, Any]] + | None = None, + evaluate_with_reward_fn: Callable[[Sequence[dict[str, Any]], PartitionAdmission, str, int, Any, Any], + dict[str, Any]] | None = None, + evaluation_rewards: dict[str, Any] | None = None, + persistent: bool = False, + idle_delay_s: float = 0.05): + super().__init__() + self.data_plane = data_plane + self.context_manager = context_manager + self.idle_delay_s = idle_delay_s + self.train_fn = train_fn + self.train_with_config_fn = train_with_config_fn + self.train_batch_configs = dict(train_batch_configs or {}) + self.save_adapter = save_adapter + self.mini_batch_sizes = mini_batch_sizes + self.scheduler = ContextScheduler(scheduler) + self.keep_adapter_versions = max(0, int(keep_adapter_versions)) + self._adapter_history: dict[str, list[str]] = defaultdict(list) + for context_key, path in (initial_adapter_paths or {}).items(): + if path: + self._adapter_history[context_key].append(path) + self.remove_adapter = remove_adapter or _remove_local_adapter + self._adapter_removal_tasks: set[asyncio.Task[None]] = set() + self.evaluation_config = dict(evaluation_config or {}) + self.evaluate_batch = evaluate_batch + self.evaluate_with_reward_fn = evaluate_with_reward_fn + self.evaluation_rewards = dict(evaluation_rewards or {}) + self._evaluation_batches: dict[str, list[Sequence[dict[str, Any]]]] = {} + self._optimizer_steps: dict[str, int] = defaultdict(int) + self.persistent = persistent + + async def register_context( + self, + context: LoraContext, + *, + mini_batch_size: int, + train_batch_config: Any | None = None, + initial_adapter_path: str | None = None, + evaluation_config: dict[str, Any] | None = None, + evaluation_reward: Any | None = None, + ) -> None: + key = context.key + if key in self.mini_batch_sizes: + raise KeyError(f'trainer context already exists: {key}') + self.mini_batch_sizes[key] = int(mini_batch_size) + if train_batch_config is not None: + self.train_batch_configs[key] = train_batch_config + if initial_adapter_path: + self._adapter_history[key].append(initial_adapter_path) + if evaluation_config is not None: + self.evaluation_config[key] = dict(evaluation_config) + if evaluation_reward is not None: + self.evaluation_rewards[key] = evaluation_reward + + async def unregister_context(self, context: LoraContext | str) -> None: + key = context if isinstance(context, str) else context.key + self.mini_batch_sizes.pop(key, None) + self.train_batch_configs.pop(key, None) + self.evaluation_config.pop(key, None) + self.evaluation_rewards.pop(key, None) + self._evaluation_batches.pop(key, None) + self._adapter_history.pop(key, None) + self._optimizer_steps.pop(key, None) + + async def stop(self) -> None: + await super().stop() + pending = tuple(self._adapter_removal_tasks) + if pending: + await asyncio.gather(*pending) + + async def _serve(self) -> None: + while not self._stop_requested: + if not self.persistent and await self.context_manager.is_run_finished.remote(): + return + admissions = await self.context_manager.list_trainable_partitions.remote() + blocked: set[str] = set() + progressed = False + for _ in range(len(admissions)): + candidates = [ + ScheduleCandidate(admission.context, admission) for admission in admissions + if admission.partition_id not in blocked + ] + candidate = self.scheduler.choose(candidates) + if candidate is None: + break + admission = candidate.partition + mini_batch_size = self.mini_batch_sizes[admission.context.key] + batch = await self.data_plane.claim_training_batch( + admission, + mini_batch_size // admission.num_generations, + ) + if batch is None: + if await self.data_plane.is_training_consumed(admission): + try: + await self.context_manager.on_partition_training_started.remote(admission) + await self._finish_partition(admission) + self.scheduler.on_success(candidate) + progressed = True + break + except Exception as exc: + self._record_metric( + 'train', + context=admission.context, + admission=admission, + status='failed', + attributes={'error': str(exc)}, + ) + raise RuntimeError( + f'training completion failed for {admission.partition_id}: {exc}') from exc + blocked.add(admission.partition_id) + self.scheduler.on_blocked(candidate) + continue + try: + await self.context_manager.on_partition_training_started.remote(admission) + policy = await self.context_manager.get_rollout_policy.remote(admission.context) + sample_count = len(batch.data['input_ids']) + if sample_count != mini_batch_size: + raise RuntimeError( + f'training claim for {admission.partition_id} returned {sample_count} samples; ' + f'expected mini_batch_size={mini_batch_size}') + started = time.perf_counter() + if self.train_with_config_fn is not None: + config = self.train_batch_configs[admission.context.key] + metrics = dict(self.train_with_config_fn(batch.data, admission, config) or {}) + else: + metrics = dict(self.train_fn(batch.data, admission) or {}) + except Exception as exc: + self._record_metric( + 'train', + context=admission.context, + admission=admission, + status='failed', + attributes={'error': str(exc)}, + ) + raise RuntimeError(f'training failed for {admission.partition_id}: {exc}') from exc + self.scheduler.on_success(candidate) + metrics['sample_count'] = sample_count + metrics['reward'] = (sum(float(value) for value in batch.data['rewards']) / sample_count) + metrics['train_latency_s'] = time.perf_counter() - started + metrics.update(training_policy_metrics(batch.sample_tags, policy.version)) + context_key = admission.context.key + self._optimizer_steps[context_key] += 1 + optimizer_step = self._optimizer_steps[context_key] + self._record_metric( + 'train', + context=admission.context, + admission=admission, + values=metrics, + optimizer_step=optimizer_step, + policy_version=policy.version, + ) + progressed = True + break + if not progressed: + await asyncio.sleep(self.idle_delay_s) + + async def _finish_partition(self, admission: PartitionAdmission) -> None: + finalize_started = time.perf_counter() + save_started = time.perf_counter() + adapter_path = self.save_adapter(admission) + adapter_save_latency_s = time.perf_counter() - save_started + publish_started = time.perf_counter() + policy = await self.context_manager.on_partition_trained.remote(admission, adapter_path=adapter_path) + policy_publish_latency_s = time.perf_counter() - publish_started + self._record_metric( + 'policy', + context=admission.context, + admission=admission, + values={ + 'adapter_save_latency_s': adapter_save_latency_s, + 'policy_publish_latency_s': policy_publish_latency_s, + }, + attributes={ + 'operation': 'publish', + 'adapter_path': adapter_path + }, + optimizer_step=self._optimizer_steps[admission.context.key], + policy_version=policy.version, + ) + await self._evaluate_policy(admission, adapter_path, policy.version) + clear_started = time.perf_counter() + await self.data_plane.clear_partition(admission) + tq_clear_latency_s = time.perf_counter() - clear_started + release_started = time.perf_counter() + await self.context_manager.on_partition_cleared.remote(admission) + partition_release_latency_s = time.perf_counter() - release_started + self._adapter_history[admission.context.key].append(adapter_path) + prune_started = time.perf_counter() + await self._prune_adapter_history(admission.context) + adapter_prune_schedule_latency_s = time.perf_counter() - prune_started + self._record_metric( + 'partition', + context=admission.context, + admission=admission, + values={ + 'adapter_save_latency_s': adapter_save_latency_s, + 'policy_publish_latency_s': policy_publish_latency_s, + 'tq_clear_latency_s': tq_clear_latency_s, + 'partition_release_latency_s': partition_release_latency_s, + 'adapter_prune_schedule_latency_s': adapter_prune_schedule_latency_s, + 'partition_finalize_latency_s': time.perf_counter() - finalize_started, + }, + optimizer_step=self._optimizer_steps[admission.context.key], + policy_version=policy.version, + ) + + async def _evaluate_policy(self, admission: PartitionAdmission, adapter_path: str, policy_version: int) -> None: + config = self.evaluation_config.get(admission.context.key) + if config is None or (self.evaluate_batch is None and self.evaluate_with_reward_fn is None): + return + interval = int(config['interval']) + if policy_version % interval: + return + + context_key = admission.context.key + if context_key not in self._evaluation_batches: + source = config['prompt_batches'] + self._evaluation_batches[context_key] = list(source() if callable(source) else source) + batches = self._evaluation_batches[context_key] + started = time.perf_counter() + rewards: list[float] = [] + completion_lengths: list[int] = [] + prompt_count = 0 + for batch in batches: + if self.evaluate_with_reward_fn is not None: + result = await asyncio.to_thread( + self.evaluate_with_reward_fn, + batch, + admission, + adapter_path, + policy_version, + config['sampling_params'], + self.evaluation_rewards[context_key], + ) + else: + result = await asyncio.to_thread( + self.evaluate_batch, + batch, + admission, + adapter_path, + policy_version, + config['sampling_params'], + ) + rewards.extend(float(value) for value in result['rewards']) + completion_lengths.extend(int(value) for value in result['completion_lengths']) + prompt_count += len(batch) + if not rewards: + raise ValueError(f'evaluation dataset is empty for {context_key}') + self._record_metric( + 'evaluation', + context=admission.context, + admission=admission, + values={ + 'accuracy': sum(rewards) / len(rewards), + 'sample_count': len(rewards), + 'prompt_count': prompt_count, + 'completion_length': sum(completion_lengths) / len(completion_lengths), + 'eval_latency_s': time.perf_counter() - started, + }, + attributes={'eval_dataset': config['dataset_name']}, + optimizer_step=self._optimizer_steps[context_key], + policy_version=policy_version, + ) + + async def _prune_adapter_history(self, context: LoraContext) -> None: + protected = set(await self.context_manager.adapter_paths_to_keep.remote()) + context_key = context.key + history = self._adapter_history[context_key] + retained_history = set(history[-self.keep_adapter_versions:]) if self.keep_adapter_versions else set() + retained = protected | retained_history + stale = [path for path in history if path not in retained] + self._adapter_history[context_key] = [path for path in history if path in retained] + for path in stale: + task = asyncio.create_task(self._remove_adapter(context, path)) + self._adapter_removal_tasks.add(task) + task.add_done_callback(self._adapter_removal_tasks.discard) + + async def _remove_adapter(self, context: LoraContext, path: str) -> None: + started = time.perf_counter() + try: + await asyncio.to_thread(self.remove_adapter, path) + except OSError as exc: + self._record_metric( + 'policy', + context=context, + status='failed', + values={ + 'adapter_prune_latency_s': time.perf_counter() - started, + }, + attributes={ + 'operation': 'adapter_prune', + 'adapter_path': path, + 'error': str(exc) + }, + ) + return + self._record_metric( + 'policy', + context=context, + values={ + 'adapter_prune_latency_s': time.perf_counter() - started, + }, + attributes={ + 'operation': 'adapter_prune', + 'adapter_path': path + }, + ) + + +def _remove_local_adapter(path: str) -> None: + if os.path.isdir(path): + shutil.rmtree(path) diff --git a/src/twinkle_client/__init__.py b/src/twinkle_client/__init__.py index a5105d497..bb13f19ad 100644 --- a/src/twinkle_client/__init__.py +++ b/src/twinkle_client/__init__.py @@ -72,4 +72,6 @@ def init_twinkle_client( ) -__all__ = ['init_tinker_client', 'init_twinkle_client'] +from .data_plane import DataPlaneClient + +__all__ = ['DataPlaneClient', 'init_tinker_client', 'init_twinkle_client'] diff --git a/src/twinkle_client/async_rl/__init__.py b/src/twinkle_client/async_rl/__init__.py new file mode 100644 index 000000000..f26758240 --- /dev/null +++ b/src/twinkle_client/async_rl/__init__.py @@ -0,0 +1,4 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from .workers import Worker, WorkerPipeline + +__all__ = ['Worker', 'WorkerPipeline'] diff --git a/src/twinkle_client/async_rl/workers.py b/src/twinkle_client/async_rl/workers.py new file mode 100644 index 000000000..8d105bcf9 --- /dev/null +++ b/src/twinkle_client/async_rl/workers.py @@ -0,0 +1,60 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Small client-side lifecycle primitives for composable async-RL workers. + +Workers remain concrete long-running roles. This module deliberately does not +introduce an algorithm graph or a data-dependency DSL; role implementations +coordinate through ordinary asyncio queues and server-side DataRefs. +""" +from __future__ import annotations + +import asyncio +from abc import ABC, abstractmethod +from collections.abc import Sequence + + +class Worker(ABC): + """One long-running client-side computation role.""" + + def __init__(self, name: str) -> None: + if not name: + raise ValueError('worker name must not be empty') + self.name = name + + @abstractmethod + async def run(self) -> None: + """Run until this role has drained its input or fails.""" + + +class WorkerPipeline: + """Run a concrete set of worker roles and propagate failures as one unit.""" + + def __init__(self, workers: Sequence[Worker]) -> None: + self.workers = tuple(workers) + if not self.workers: + raise ValueError('at least one worker is required') + names = [worker.name for worker in self.workers] + if len(names) != len(set(names)): + raise ValueError(f'worker names must be unique, got {names}') + + async def run(self) -> None: + tasks = { + asyncio.create_task(worker.run(), name=worker.name): worker + for worker in self.workers + } + try: + done, pending = await asyncio.wait(tasks, return_when=asyncio.FIRST_EXCEPTION) + failure = next( + (task.exception() for task in done if not task.cancelled() and task.exception() is not None), + None, + ) + if failure is not None: + for task in pending: + task.cancel() + await asyncio.gather(*pending, return_exceptions=True) + raise failure + await asyncio.gather(*pending) + finally: + for task in tasks: + if not task.done(): + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) diff --git a/src/twinkle_client/common/json_utils.py b/src/twinkle_client/common/json_utils.py new file mode 100644 index 000000000..51c039c15 --- /dev/null +++ b/src/twinkle_client/common/json_utils.py @@ -0,0 +1,33 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Lightweight JSON conversion helpers shared by component clients and servers.""" +from __future__ import annotations + +from collections.abc import Mapping +from numbers import Number +from typing import Any + +from pydantic import BaseModel + + +_PRIMITIVE_TYPES = (str, Number, bool, bytes, type(None)) + + +def json_safe(obj: Any) -> Any: + """Recursively convert models, tensors, and arrays into JSON-compatible values. + + This module intentionally does not import :mod:`twinkle`: component protocol + types are imported while the top-level package is still being initialized. + """ + if isinstance(obj, BaseModel): + return json_safe(obj.model_dump()) + if isinstance(obj, Mapping): + return {key: json_safe(value) for key, value in obj.items()} + if isinstance(obj, (list, tuple, set, frozenset)): + return [json_safe(value) for value in obj] + tolist = getattr(obj, 'tolist', None) + if callable(tolist) and not isinstance(obj, _PRIMITIVE_TYPES): + try: + return json_safe(tolist()) + except Exception: + pass + return obj diff --git a/src/twinkle_client/data_plane.py b/src/twinkle_client/data_plane.py new file mode 100644 index 000000000..281df8315 --- /dev/null +++ b/src/twinkle_client/data_plane.py @@ -0,0 +1,129 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Client for the server-side TransferQueue DataRef service.""" +from __future__ import annotations + +import asyncio +from collections.abc import Callable +from typing import Any, TypeVar + +from twinkle_client.common.json_utils import json_safe +from twinkle_client.http import get_base_url, http_post +from twinkle_client.types.component import DataRef, DataRowsResponse + + +_T = TypeVar('_T') + + +async def _call_in_thread(func: Callable[..., _T], /, *args: Any, **kwargs: Any) -> _T: + """Run one synchronous DataPlane operation without blocking the event loop.""" + return await asyncio.to_thread(func, *args, **kwargs) + + +class DataPlaneClient: + + def __init__(self, server_url: str | None = None): + self.server_url = (server_url or f'{get_base_url()}/data-plane').rstrip('/') + + def put( + self, + rows: list[dict[str, Any]], + *, + kind: str = 'data', + tags: list[dict[str, Any]] | None = None, + ) -> DataRef: + response = http_post( + f'{self.server_url}/twinkle/put', + json_data={'rows': json_safe(rows), 'kind': kind, 'tags': json_safe(tags)}, + ) + response.raise_for_status() + return DataRef(**response.json()) + + async def aput( + self, + rows: list[dict[str, Any]], + *, + kind: str = 'data', + tags: list[dict[str, Any]] | None = None, + ) -> DataRef: + """Asynchronously store rows while preserving :meth:`put` semantics.""" + if tags is None: + return await _call_in_thread(self.put, rows, kind=kind) + return await _call_in_thread(self.put, rows, kind=kind, tags=tags) + + def get(self, ref: DataRef, *, fields: list[str] | None = None) -> list[dict[str, Any]]: + response = http_post( + f'{self.server_url}/twinkle/get', + json_data={'ref': ref.model_dump(), 'fields': fields}, + ) + response.raise_for_status() + return DataRowsResponse(**response.json()).rows + + def get_batch( + self, + ref: DataRef, + *, + fields: list[str] | None = None, + ) -> DataRowsResponse: + response = http_post( + f'{self.server_url}/twinkle/get', + json_data={'ref': ref.model_dump(), 'fields': fields, 'include_tags': True}, + ) + response.raise_for_status() + return DataRowsResponse(**response.json()) + + async def aget(self, ref: DataRef, *, fields: list[str] | None = None) -> list[dict[str, Any]]: + """Asynchronously fetch rows while preserving :meth:`get` semantics.""" + if fields is None: + return await _call_in_thread(self.get, ref) + return await _call_in_thread(self.get, ref, fields=fields) + + async def aget_batch( + self, + ref: DataRef, + *, + fields: list[str] | None = None, + ) -> DataRowsResponse: + if fields is None: + return await _call_in_thread(self.get_batch, ref) + return await _call_in_thread(self.get_batch, ref, fields=fields) + + def append( + self, + ref: DataRef, + rows: list[dict[str, Any]], + *, + tags: list[dict[str, Any]] | None = None, + ) -> DataRef: + response = http_post( + f'{self.server_url}/twinkle/append', + json_data={ + 'ref': ref.model_dump(), + 'rows': json_safe(rows), + 'tags': json_safe(tags), + }, + ) + response.raise_for_status() + return DataRef(**response.json()) + + async def aappend( + self, + ref: DataRef, + rows: list[dict[str, Any]], + *, + tags: list[dict[str, Any]] | None = None, + ) -> DataRef: + """Asynchronously append rows while preserving :meth:`append` semantics.""" + if tags is None: + return await _call_in_thread(self.append, ref, rows) + return await _call_in_thread(self.append, ref, rows, tags=tags) + + def release(self, ref: DataRef) -> None: + response = http_post( + f'{self.server_url}/twinkle/release', + json_data={'ref': ref.model_dump()}, + ) + response.raise_for_status() + + async def arelease(self, ref: DataRef) -> None: + """Asynchronously release a reference while preserving :meth:`release` semantics.""" + await _call_in_thread(self.release, ref) diff --git a/src/twinkle_client/model/multi_lora_transformers.py b/src/twinkle_client/model/multi_lora_transformers.py index bf4ef54af..3471ca6c0 100644 --- a/src/twinkle_client/model/multi_lora_transformers.py +++ b/src/twinkle_client/model/multi_lora_transformers.py @@ -2,6 +2,8 @@ from pathlib import Path import time from twinkle_client.http import http_get, http_post +from twinkle_client.common.json_utils import json_safe +from twinkle_client.types.component import DataRef from twinkle_client.types.model import ( CalculateLossResponse, CalculateMetricResponse, @@ -15,6 +17,16 @@ ) +def _data_ref_payload(inputs: DataRef | list[DataRef]) -> dict[str, Any]: + """Encode one or more opaque references for a DataPlane model endpoint.""" + refs = [inputs] if isinstance(inputs, DataRef) else list(inputs) + if not refs: + raise ValueError('at least one DataRef is required') + if not all(isinstance(item, DataRef) for item in refs): + raise TypeError('data-plane model inputs must contain only DataRef values') + return {'input_refs': [item.model_dump() for item in refs]} + + class MultiLoraTransformersModel: """Client wrapper for TwinkleModel that calls server HTTP endpoints. @@ -26,6 +38,7 @@ def __init__(self, model_id: str, **kwargs): """Initialize model client.""" from twinkle_client.http import get_base_url self.server_url = get_base_url() + kwargs.pop('data_plane_url', None) if '://' in model_id: model_id = model_id.split('://')[1] @@ -54,24 +67,87 @@ def add_adapter_to_model(self, adapter_name: str, config: Optional[Dict[str, Any response.raise_for_status() self.adapter_name = adapter_name + def remove_adapter(self, adapter_name: str | None = None) -> None: + """Release one client-owned adapter from the training component.""" + name = adapter_name or self.adapter_name + response = http_post( + url=f'{self.server_url}/remove_adapter', + json_data={'adapter_name': name}, + ) + response.raise_for_status() + if name == self.adapter_name: + self.adapter_name = None + def forward(self, inputs: Any, **kwargs) -> ForwardResponse: - """Execute forward pass on the model.""" + """Execute forward pass on inline model inputs.""" response = http_post( url=f'{self.server_url}/forward', - json_data={'inputs': inputs, 'adapter_name': self.adapter_name, **kwargs} + json_data={'inputs': inputs, 'adapter_name': self.adapter_name, **kwargs}, ) response.raise_for_status() return ForwardResponse(**response.json()) def forward_only(self, inputs: Any, **kwargs) -> ForwardResponse: - """Execute forward pass without gradient computation.""" + """Execute forward pass without gradient computation on inline inputs.""" response = http_post( url=f'{self.server_url}/forward_only', - json_data={'inputs': inputs, 'adapter_name': self.adapter_name, **kwargs} + json_data={'inputs': inputs, 'adapter_name': self.adapter_name, **kwargs}, + ) + response.raise_for_status() + return ForwardResponse(**response.json()) + + def forward_from_data_plane( + self, + inputs: DataRef | list[DataRef], + *, + input_field: str | None = None, + kwarg_fields: dict[str, str] | None = None, + **kwargs, + ) -> ForwardResponse: + """Execute forward using rows referenced from the server DataPlane.""" + response = http_post( + url=f'{self.server_url}/forward_from_data_plane', + json_data={ + **_data_ref_payload(inputs), + 'adapter_name': self.adapter_name, + 'input_field': input_field, + 'kwarg_fields': kwarg_fields or {}, + **json_safe(kwargs), + }, ) response.raise_for_status() return ForwardResponse(**response.json()) + def forward_only_from_data_plane( + self, + inputs: DataRef | list[DataRef], + *, + input_field: str | None = None, + kwarg_fields: dict[str, str] | None = None, + output_ref: DataRef | None = None, + output_fields: dict[str, str] | None = None, + **kwargs, + ) -> ForwardResponse | DataRef: + """Execute forward-only using DataPlane rows and optionally append outputs.""" + body = { + **_data_ref_payload(inputs), + 'adapter_name': self.adapter_name, + 'input_field': input_field, + 'kwarg_fields': kwarg_fields or {}, + 'output_ref': output_ref.model_dump() if output_ref is not None else None, + 'output_fields': output_fields or {}, + **json_safe(kwargs), + } + response = http_post( + url=f'{self.server_url}/forward_only_from_data_plane', + json_data=body, + ) + response.raise_for_status() + result = ForwardResponse(**response.json()) + if output_ref is not None: + return DataRef(**result.result) + return result + def calculate_loss(self, **kwargs) -> CalculateLossResponse: """Calculate loss from model outputs.""" response = http_post( @@ -99,10 +175,32 @@ def backward(self, **kwargs) -> None: response.raise_for_status() def forward_backward(self, inputs: Any, **kwargs) -> ForwardBackwardResponse: - """Execute combined forward and backward pass.""" + """Execute combined forward and backward pass on inline inputs.""" response = http_post( url=f'{self.server_url}/forward_backward', - json_data={'inputs': inputs, 'adapter_name': self.adapter_name, **kwargs} + json_data={'inputs': inputs, 'adapter_name': self.adapter_name, **kwargs}, + ) + response.raise_for_status() + return ForwardBackwardResponse(**response.json()) + + def forward_backward_from_data_plane( + self, + inputs: DataRef | list[DataRef], + *, + input_field: str | None = None, + kwarg_fields: dict[str, str] | None = None, + **kwargs, + ) -> ForwardBackwardResponse: + """Execute forward/backward using rows referenced from the server DataPlane.""" + response = http_post( + url=f'{self.server_url}/forward_backward_from_data_plane', + json_data={ + **_data_ref_payload(inputs), + 'adapter_name': self.adapter_name, + 'input_field': input_field, + 'kwarg_fields': kwarg_fields or {}, + **json_safe(kwargs), + }, ) response.raise_for_status() return ForwardBackwardResponse(**response.json()) diff --git a/src/twinkle_client/sampler/vllm_sampler.py b/src/twinkle_client/sampler/vllm_sampler.py index 0d553bb32..25ea639c0 100644 --- a/src/twinkle_client/sampler/vllm_sampler.py +++ b/src/twinkle_client/sampler/vllm_sampler.py @@ -1,8 +1,11 @@ +import asyncio from typing import Any, Dict, List, Optional, Union from twinkle_client.http import http_post from twinkle_client.types.sampler import AddAdapterResponse, SampleResponseModel, SetTemplateResponse from peft import PeftConfig from twinkle.data_format import Trajectory, InputFeature +from twinkle_client.common.json_utils import json_safe +from twinkle_client.types.component import DataRef # Intentionally does NOT subclass ``twinkle.sampler.base.Sampler``: importing @@ -17,17 +20,7 @@ def _json_safe(obj: Any) -> Any: duck-typing (``.tolist()``) so this stays free of a hard torch/numpy import, honouring the CPU-only client contract noted above. """ - if isinstance(obj, dict): - return {k: _json_safe(v) for k, v in obj.items()} - if isinstance(obj, (list, tuple)): - return [_json_safe(x) for x in obj] - tolist = getattr(obj, 'tolist', None) - if callable(tolist) and not isinstance(obj, (str, bytes, int, float, bool)): - try: - return _json_safe(tolist()) - except Exception: - return obj - return obj + return json_safe(obj) class vLLMSampler: @@ -41,6 +34,8 @@ def __init__(self, model_id: str, **kwargs): """Create the sampler instance on server.""" from twinkle_client.http import get_base_url self.server_url = get_base_url() + from twinkle_client.data_plane import DataPlaneClient + self.data_plane = DataPlaneClient(kwargs.pop('data_plane_url', None)) self.adapter_name = None if '://' in model_id: @@ -85,6 +80,8 @@ def sample( Returns: SampleResponseModel with 'sequences' list, each containing tokens, logprobs, stop_reason. """ + sampling_params = dict(sampling_params or {}) + sampling_params['num_samples'] = num_samples json_data = { 'inputs': _json_safe(inputs), 'sampling_params': sampling_params, @@ -101,6 +98,84 @@ def sample( response.raise_for_status() return [SampleResponseModel(**r) for r in response.json()['samples']] + def sample_to_data_plane( + self, + inputs: Union[List[Trajectory], List[InputFeature], DataRef], + sampling_params: Optional[Dict[str, Any]] = None, + *, + adapter_name: str = '', + adapter_uri: Optional[str] = None, + policy_version: int | None = None, + group_ids: list[str] | None = None, + num_samples: int = 1, + ) -> DataRef: + """Generate complete prompt groups and keep their rows in the server DataPlane.""" + body = { + 'sampling_params': sampling_params, + 'adapter_name': adapter_name, + 'adapter_uri': adapter_uri, + 'policy_version': policy_version, + 'group_ids': group_ids, + 'num_samples': num_samples, + } + body['input_ref' if isinstance(inputs, DataRef) else 'inputs'] = ( + inputs.model_dump() if isinstance(inputs, DataRef) else _json_safe(inputs)) + response = http_post( + url=f'{self.server_url}/sample_to_data_plane', + json_data=json_safe(body), + ) + response.raise_for_status() + return DataRef(**response.json()) + + async def asample( + self, + inputs: Union[List[Trajectory], List[InputFeature]], + sampling_params: Optional[Dict[str, Any]] = None, + adapter_name: str = '', + adapter_uri: Optional[str] = None, + num_samples: int = 1, + ) -> List[SampleResponseModel]: + """Asynchronous convenience wrapper for the materialized sample API.""" + return await asyncio.to_thread( + self.sample, + inputs, + sampling_params, + adapter_name=adapter_name, + adapter_uri=adapter_uri, + num_samples=num_samples, + ) + + async def asample_to_data_plane( + self, + inputs: Union[List[Trajectory], List[InputFeature], DataRef], + sampling_params: Optional[Dict[str, Any]] = None, + *, + adapter_name: str = '', + adapter_uri: Optional[str] = None, + policy_version: int | None = None, + group_ids: list[str] | None = None, + num_samples: int = 1, + ) -> DataRef: + """Asynchronously sample and return the opaque server-side result reference.""" + return await asyncio.to_thread( + self.sample_to_data_plane, + inputs, + sampling_params, + adapter_name=adapter_name, + adapter_uri=adapter_uri, + policy_version=policy_version, + group_ids=group_ids, + num_samples=num_samples, + ) + + def unload_adapter_paths(self, adapter_paths: list[str]) -> None: + """Evict policy snapshots that are no longer referenced by this client.""" + response = http_post( + url=f'{self.server_url}/unload_adapter_paths', + json_data={'adapter_paths': adapter_paths}, + ) + response.raise_for_status() + def set_template(self, template_cls: str, adapter_name: str = '', **kwargs) -> SetTemplateResponse: """Set the template for encoding trajectories.""" response = http_post( diff --git a/src/twinkle_client/types/__init__.py b/src/twinkle_client/types/__init__.py index 49673b0e9..1c25324a1 100644 --- a/src/twinkle_client/types/__init__.py +++ b/src/twinkle_client/types/__init__.py @@ -16,6 +16,8 @@ ClipGradNormResponse, CreateRequest, CreateResponse, + DataPlaneForwardOnlyRequest, + DataPlaneForwardRequest, ForwardBackwardResponse, ForwardOnlyRequest, ForwardRequest, @@ -92,3 +94,13 @@ ) from .checkpoint import ResolvedLoadPath +from .component import ( + DataAppendRequest, + DataGetRequest, + DataPlaneSampleRequest, + DataPutRequest, + DataRef, + DataReleaseRequest, + DataRowsResponse, + UnloadAdapterPathsRequest, +) diff --git a/src/twinkle_client/types/component.py b/src/twinkle_client/types/component.py new file mode 100644 index 000000000..d7e9ec2a0 --- /dev/null +++ b/src/twinkle_client/types/component.py @@ -0,0 +1,69 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Protocol types for directly orchestrating asynchronous server components.""" +from __future__ import annotations + +from typing import Any + +from pydantic import BaseModel, Field, model_validator + + +class DataRef(BaseModel): + """Opaque reference to rows stored in the server-side TransferQueue.""" + + ref_id: str + size: int + fields: list[str] = Field(default_factory=list) + kind: str = 'data' + num_tokens: int = 0 + + +class DataPutRequest(BaseModel): + rows: list[dict[str, Any]] + kind: str = 'data' + tags: list[dict[str, Any]] | None = None + + +class DataGetRequest(BaseModel): + ref: DataRef + fields: list[str] | None = None + include_tags: bool = False + + +class DataAppendRequest(BaseModel): + ref: DataRef + rows: list[dict[str, Any]] + tags: list[dict[str, Any]] | None = None + + +class DataReleaseRequest(BaseModel): + ref: DataRef + + +class DataRowsResponse(BaseModel): + rows: list[dict[str, Any]] + tags: list[dict[str, Any]] = Field(default_factory=list) + + +class DataPlaneSampleRequest(BaseModel): + inputs: Any = None + input_ref: DataRef | None = None + sampling_params: dict[str, Any] | None = None + adapter_name: str = '' + adapter_uri: str | None = None + policy_version: int | None = None + group_ids: list[str] | None = None + num_samples: int = 1 + + @model_validator(mode='after') + def validate_input(self) -> 'DataPlaneSampleRequest': + if (self.inputs is None) == (self.input_ref is None): + raise ValueError('exactly one of inputs and input_ref must be provided') + if self.group_ids is not None and self.inputs is not None: + size = len(self.inputs) if isinstance(self.inputs, list) else 1 + if len(self.group_ids) != size: + raise ValueError('group_ids must contain one value per sampler input') + return self + + +class UnloadAdapterPathsRequest(BaseModel): + adapter_paths: list[str] diff --git a/src/twinkle_client/types/model.py b/src/twinkle_client/types/model.py index 83f489cde..3d4f4cf45 100644 --- a/src/twinkle_client/types/model.py +++ b/src/twinkle_client/types/model.py @@ -4,9 +4,11 @@ These models are used by both the server-side handler and the twinkle client. """ -from pydantic import BaseModel, field_validator +from pydantic import BaseModel, Field, field_validator, model_validator from typing import Any, Dict, List, Optional, Union +from .component import DataRef + class CreateRequest(BaseModel): @@ -30,6 +32,27 @@ class Config: extra = 'allow' +class DataPlaneForwardRequest(BaseModel): + input_refs: List[DataRef] = Field(min_length=1) + input_field: str | None = None + kwarg_fields: Dict[str, str] = Field(default_factory=dict) + adapter_name: str + + class Config: + extra = 'allow' + + +class DataPlaneForwardOnlyRequest(DataPlaneForwardRequest): + output_ref: DataRef | None = None + output_fields: Dict[str, str] = Field(default_factory=dict) + + @model_validator(mode='after') + def validate_output(self) -> 'DataPlaneForwardOnlyRequest': + if (self.output_ref is None) != (len(self.output_fields) == 0): + raise ValueError('output_ref and output_fields must be configured together') + return self + + class AdapterRequest(BaseModel): adapter_name: str diff --git a/tests/infra/test_resource_manager.py b/tests/infra/test_resource_manager.py new file mode 100644 index 000000000..ee635ce4e --- /dev/null +++ b/tests/infra/test_resource_manager.py @@ -0,0 +1,20 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. +from twinkle.infra._ray.resource_manager import _gpu_placement_group_cpus + + +def test_gpu_placement_group_cpus_uses_per_process_default(monkeypatch): + monkeypatch.delenv('TWINKLE_GPU_PG_CPUS_PER_PROC', raising=False) + + assert _gpu_placement_group_cpus(node_cpu=192, nproc_per_node=2) == 8 + + +def test_gpu_placement_group_cpus_respects_node_cap(monkeypatch): + monkeypatch.setenv('TWINKLE_GPU_PG_CPUS_PER_PROC', '16') + + assert _gpu_placement_group_cpus(node_cpu=192, nproc_per_node=8) == 48 + + +def test_gpu_placement_group_cpus_can_be_configured(monkeypatch): + monkeypatch.setenv('TWINKLE_GPU_PG_CPUS_PER_PROC', '8') + + assert _gpu_placement_group_cpus(node_cpu=192, nproc_per_node=2) == 16 diff --git a/tests/loss/test_dpo.py b/tests/loss/test_dpo.py index f12d5ea3a..89b9ddac4 100644 --- a/tests/loss/test_dpo.py +++ b/tests/loss/test_dpo.py @@ -42,6 +42,15 @@ def test_basic_dpo_sigmoid(self): assert isinstance(result, dict) and 'loss' in result assert result['loss'].dim() == 0 + def test_json_ref_outputs_match_tensor_ref_outputs(self): + loss_fn = DPOLoss(beta=0.1, loss_type='sigmoid') + inputs, outputs, ref_logps = _make_preference_batch() + + tensor_result = loss_fn(inputs, outputs, ref_outputs={'logps': ref_logps}) + json_result = loss_fn(inputs, outputs, ref_outputs={'logps': ref_logps.tolist()}) + + torch.testing.assert_close(json_result['loss'], tensor_result['loss']) + def test_dpo_hinge(self): loss_fn = DPOLoss(beta=0.1, loss_type='hinge') inputs, outputs, ref_logps = _make_preference_batch() diff --git a/tests/loss/test_grpo_gkd.py b/tests/loss/test_grpo_gkd.py index 541b867a0..3d0f55120 100644 --- a/tests/loss/test_grpo_gkd.py +++ b/tests/loss/test_grpo_gkd.py @@ -82,6 +82,25 @@ def test_grpo_list_advantages(self): result = loss_fn(inputs, outputs, old_logps=old_logps, advantages=adv_list) assert torch.isfinite(result['loss']) + def test_grpo_weights_sequences_equally(self): + labels = torch.tensor([ + [1, -100, -100], + [1, 1, 1], + ]) + logps = torch.zeros_like(labels, dtype=torch.float32) + inputs = {'labels': labels} + outputs = {'logps': logps} + advantages = torch.tensor([[1.0], [3.0]]) + + result = GRPOLoss()( + inputs, + outputs, + old_logps=logps, + advantages=advantages, + ) + + assert result['loss'].item() == pytest.approx(-2.0) + def test_grpo_entropy_coef(self): loss_fn = GRPOLoss(epsilon=0.2, entropy_coef=0.01) inputs, outputs, old_logps, _, advantages = _make_rl_batch() diff --git a/tests/metric/test_metrics.py b/tests/metric/test_metrics.py index e4651ee86..691da2c06 100644 --- a/tests/metric/test_metrics.py +++ b/tests/metric/test_metrics.py @@ -271,6 +271,26 @@ def test_dpo_metric_with_ref(self): assert 'rewards/chosen' in result assert 'rewards/accuracies' in result + def test_dpo_metric_accepts_json_ref_logps(self): + labels = torch.tensor([[1, 2, -100], [3, 4, -100]]) + logps = torch.randn(2, 3) + ref_logps = torch.randn(2, 3) + tensor_metric = _no_dist_metric(DPOMetric, beta=0.1) + json_metric = _no_dist_metric(DPOMetric, beta=0.1) + + tensor_metric.accumulate( + {'labels': labels}, + {'logps': logps}, + ref_outputs={'logps': ref_logps}, + ) + json_metric.accumulate( + {'labels': labels}, + {'logps': logps}, + ref_outputs={'logps': ref_logps.tolist()}, + ) + + assert json_metric.calculate() == pytest.approx(tensor_metric.calculate()) + def test_dpo_metric_no_logps_skips(self): m = _no_dist_metric(DPOMetric, beta=0.1) m.accumulate({'labels': torch.tensor([[1, 2]])}, {}) diff --git a/tests/model/test_micro_batch.py b/tests/model/test_micro_batch.py new file mode 100644 index 000000000..8d21c1dbb --- /dev/null +++ b/tests/model/test_micro_batch.py @@ -0,0 +1,261 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from types import SimpleNamespace + +import pytest + +from twinkle.loss import CrossEntropyLoss, GRPOLoss +from twinkle.loss.base import Loss +from twinkle.model.micro_batch import MicroBatchConfig, plan_micro_batches +from twinkle.model.transformers.transformers import TransformersModel +from twinkle.processor import InputProcessor +from twinkle.utils.nccl_safe import safe_loss + + +@pytest.mark.parametrize('packing_algorithm', ['ffd', 'kk']) +def test_dynamic_micro_batch_plan_preserves_samples_and_limits_cost(packing_algorithm): + lengths = [10, 9, 8, 7, 4, 3, 2, 1] + inputs = [{'input_ids': list(range(length))} for length in lengths] + config = MicroBatchConfig( + micro_batch_size=3, + dynamic_batching=True, + max_tokens_per_micro_batch=18, + packing_algorithm=packing_algorithm, + ) + + plan = plan_micro_batches(inputs, config, padding_free=False) + + assert sorted(index for batch in plan for index in batch) == list(range(len(inputs))) + for batch in plan: + assert len(batch) <= config.micro_batch_size + assert max(lengths[index] for index in batch) * len(batch) <= 18 + + +def test_padding_free_dynamic_batching_uses_unpadded_token_cost(): + lengths = [10, 8, 6, 4] + inputs = [{'input_ids': list(range(length))} for length in lengths] + config = MicroBatchConfig( + micro_batch_size=3, + dynamic_batching=True, + max_tokens_per_micro_batch=18, + ) + + plan = plan_micro_batches(inputs, config, padding_free=True) + + assert sorted(index for batch in plan for index in batch) == list(range(len(inputs))) + assert all(sum(lengths[index] for index in batch) <= 18 for batch in plan) + + +def test_token_mean_loss_weight_uses_valid_label_count(): + inputs = [ + {'labels': [1, 2, -100, -100]}, + {'labels': [3, 4, 5, -100]}, + {'labels': [6, -100, -100, -100]}, + ] + + loss = CrossEntropyLoss(reduction='mean') + first_weight = loss.micro_batch_scale(inputs, [0, 2]) + second_weight = loss.micro_batch_scale(inputs, [1]) + + assert first_weight == .5 + assert second_weight == .5 + + +def test_sample_mean_and_token_sum_micro_batch_scales(): + inputs = [ + {'labels': [1, -100]}, + {'labels': [2, 3]}, + {'labels': [4, -100]}, + {'labels': [5, 6]}, + ] + + assert GRPOLoss().micro_batch_scale(inputs, [0]) == .25 + assert CrossEntropyLoss(reduction='sum').micro_batch_scale(inputs, [0]) == 1.0 + + +def test_safe_loss_preserves_wrapped_micro_batch_scale(): + inputs = [ + {'labels': [1, -100]}, + {'labels': [2, 3]}, + {'labels': [4, -100]}, + {'labels': [5, 6]}, + ] + + assert safe_loss(GRPOLoss()).micro_batch_scale(inputs, [0, 2]) == .5 + + +def test_loss_without_micro_batch_semantics_fails_when_split(): + with pytest.raises(NotImplementedError, match='does not support micro-batching'): + Loss().micro_batch_scale([{}, {}], [0]) + + +def test_transformers_forward_backward_keeps_original_default_path(): + class ModelHarness: + def __init__(self): + self.calls = [] + + def forward(self, *, inputs, **_kwargs): + self.calls.append(('forward', inputs)) + return {} + + def calculate_loss(self, **_kwargs): + self.calls.append(('loss', None)) + return 2.0 + + def backward(self, **_kwargs): + self.calls.append(('backward', None)) + + model = ModelHarness() + outputs = TransformersModel.forward_backward.__wrapped__( + model, + inputs=[{'input_ids': [1, 2]}], + ) + + assert [name for name, _ in model.calls] == ['forward', 'loss', 'backward'] + assert outputs['loss'] == 2.0 + + +def test_transformers_forward_backward_executes_real_micro_batches(): + class OptimizerConfig: + def __init__(self): + self.processor = InputProcessor(padding_free=False) + self.template = None + self.train_status = SimpleNamespace(loss_value=None, num_tokens=0.0) + self.loss_instance = SimpleNamespace( + micro_batch_scale=lambda inputs, indices: len(indices) / len(inputs), + ) + self._dp_group = None + + def _ensure_dp_group(self): + pass + + class ModelHarness: + _build_micro_batch_plan = TransformersModel._build_micro_batch_plan + _forward_backward_micro_batch = TransformersModel._forward_backward_micro_batch + _forward_backward_micro_batches = TransformersModel._forward_backward_micro_batches + + def __init__(self): + self.optimizer_group = {'adapter': OptimizerConfig()} + self.forward_batches = [] + self.backward_calls = [] + + def _get_default_group(self): + return 'adapter' + + @staticmethod + def _not_encoded(_inputs): + return False + + def forward(self, *, inputs, **_kwargs): + self.forward_batches.append([item['sample_id'] for item in inputs]) + return {} + + def calculate_loss(self, **_kwargs): + self.optimizer_group['adapter'].train_status.loss_value = 2.0 + self.optimizer_group['adapter'].train_status.num_tokens += 1.0 + return 2.0 + + def backward(self, *, sync_gradients, **_kwargs): + loss = self.optimizer_group['adapter'].train_status.loss_value + self.backward_calls.append((sync_gradients, loss)) + self.optimizer_group['adapter'].train_status.loss_value = None + + model = ModelHarness() + inputs = [ + { + 'sample_id': index, + 'input_ids': list(range(index + 1)), + } + for index in range(4) + ] + + outputs = TransformersModel.forward_backward.__wrapped__( + model, + inputs=inputs, + adapter_name='adapter', + micro_batch_size=2, + sync_gradients=True, + ) + + assert model.forward_batches == [[0, 1], [2, 3]] + assert model.backward_calls == [(False, 1.0), (True, 1.0)] + assert model.optimizer_group['adapter'].train_status.num_tokens == 1.0 + + +def test_fixed_micro_batch_plan_can_match_a_larger_dp_micro_batch_count(): + inputs = [{'input_ids': [index]} for index in range(4)] + + plan = plan_micro_batches( + inputs, + MicroBatchConfig(micro_batch_size=2), + padding_free=False, + min_micro_batches=3, + ) + + assert plan == [[0, 1], [2], [3]] + + +def test_dp_micro_batch_planning_propagates_remote_rank_error(monkeypatch): + from twinkle.model.transformers import transformers as module + + class OptimizerConfig: + processor = InputProcessor(padding_free=False) + _dp_group = object() + + @staticmethod + def _ensure_dp_group(): + pass + + def all_gather(states, local_state, *, group): + assert group is OptimizerConfig._dp_group + states[:] = [ + local_state, + { + 'micro_batch_count': None, + 'input_count': 2, + 'error': 'ValueError: sequence length 20 exceeds the limit', + }, + ] + + monkeypatch.setattr(module.dist, 'get_world_size', lambda _group: 2) + monkeypatch.setattr(module.dist, 'all_gather_object', all_gather) + + with pytest.raises(RuntimeError, match='rank 1.*sequence length 20'): + TransformersModel._build_micro_batch_plan( + object(), + [{'input_ids': [1]}, {'input_ids': [2]}], + MicroBatchConfig(micro_batch_size=1), + OptimizerConfig(), + ) + +def test_dp_micro_batch_planning_rejects_common_count_on_all_ranks(monkeypatch): + from twinkle.model.transformers import transformers as module + + class OptimizerConfig: + processor = InputProcessor(padding_free=False) + _dp_group = object() + + @staticmethod + def _ensure_dp_group(): + pass + + def all_gather(states, local_state, *, group): + assert group is OptimizerConfig._dp_group + states[:] = [ + local_state, + { + 'micro_batch_count': 3, + 'input_count': 3, + 'error': None, + }, + ] + + monkeypatch.setattr(module.dist, 'get_world_size', lambda _group: 2) + monkeypatch.setattr(module.dist, 'all_gather_object', all_gather) + + with pytest.raises(ValueError, match='same number of non-empty micro-batches'): + TransformersModel._build_micro_batch_plan( + object(), + [{'input_ids': [1]}, {'input_ids': [2]}], + MicroBatchConfig(micro_batch_size=1), + OptimizerConfig(), + ) diff --git a/tests/model/test_multi_lora.py b/tests/model/test_multi_lora.py new file mode 100644 index 000000000..f2795a5ad --- /dev/null +++ b/tests/model/test_multi_lora.py @@ -0,0 +1,28 @@ +import pytest + +from twinkle.model.multi_lora import MultiLora + + +def test_check_length_checks_each_sample_independently(): + multi_lora = MultiLora(max_length=4) + + multi_lora.check_length([ + {'input_ids': [1, 2, 3]}, + {'input_ids': [4, 5, 6]}, + ]) + + +def test_check_length_accepts_a_single_input_feature(): + multi_lora = MultiLora(max_length=4) + + multi_lora.check_length({'input_ids': [1, 2, 3, 4]}) + + +def test_check_length_reports_the_oversized_sample(): + multi_lora = MultiLora(max_length=4) + + with pytest.raises(ValueError, match=r'Input length 5 exceeds max_length 4 at sample 1'): + multi_lora.check_length([ + {'input_ids': [1, 2]}, + {'input_ids': [1, 2, 3, 4, 5]}, + ]) diff --git a/tests/model/test_multi_lora_dtype.py b/tests/model/test_multi_lora_dtype.py new file mode 100644 index 000000000..611008a40 --- /dev/null +++ b/tests/model/test_multi_lora_dtype.py @@ -0,0 +1,22 @@ +import torch +from peft import LoraConfig +from torch import nn + +from twinkle.model.multi_lora import MultiLora +from twinkle.model.transformers.transformers import TransformersModel + + +def test_multi_lora_dtype_matches_bf16_base_before_fsdp_wrap(): + model = nn.Sequential(nn.Linear(4, 4, dtype=torch.bfloat16)) + multi_lora = MultiLora(max_loras=2, max_r=4) + model = multi_lora.patch( + model, + target_modules=['0'], + lora_config=LoraConfig(r=4, lora_alpha=8, target_modules=['0']), + ) + + assert {param.dtype for name, param in model.named_parameters() if 'lora_' in name} == {torch.float32} + + TransformersModel._ensure_lora_dtype(None, model) + + assert {param.dtype for name, param in model.named_parameters() if 'lora_' in name} == {torch.bfloat16} diff --git a/tests/preprocessor/test_math_rl_processors.py b/tests/preprocessor/test_math_rl_processors.py new file mode 100644 index 000000000..5e43405c3 --- /dev/null +++ b/tests/preprocessor/test_math_rl_processors.py @@ -0,0 +1,19 @@ +from twinkle.preprocessor import DAPOMathProcessor + + +def test_dapo_math_processor_preserves_prompt_and_ground_truth(): + row = { + 'prompt': [{ + 'role': 'user', + 'content': 'Solve the problem and put the answer in a box.', + }], + 'reward_model': { + 'ground_truth': '34', + 'style': 'rule-lighteval/MATH_v2', + }, + } + + trajectory = DAPOMathProcessor().preprocess(row) + + assert trajectory['messages'] == row['prompt'] + assert trajectory['user_data'] == [('ground_truth', '34')] diff --git a/tests/sampler/test_vllm_lora_loading.py b/tests/sampler/test_vllm_lora_loading.py new file mode 100644 index 000000000..f63c9c8ad --- /dev/null +++ b/tests/sampler/test_vllm_lora_loading.py @@ -0,0 +1,65 @@ +import asyncio +from unittest.mock import MagicMock + +from twinkle.sampler.vllm_sampler.vllm_engine import VLLMEngine + + +def test_concurrent_lora_requests_share_one_load_task(): + async def run(): + engine = VLLMEngine.__new__(VLLMEngine) + engine._lora_request_cache = {} + engine._lora_load_tasks = {} + request = object() + load_count = 0 + + async def load_lora(_path): + nonlocal load_count + load_count += 1 + await asyncio.sleep(.01) + return request + + engine._load_lora = load_lora + results = await asyncio.gather(*(engine._get_or_load_lora('/adapter') for _ in range(8))) + + assert load_count == 1 + assert results == [request] * 8 + assert engine._lora_request_cache == {'/adapter': request} + assert engine._lora_load_tasks == {} + + asyncio.run(run()) + + +def test_unload_lora_accepts_synchronous_engine_api(): + async def run(): + engine = VLLMEngine.__new__(VLLMEngine) + request = MagicMock(lora_int_id=7) + engine._lora_request_cache = {'/adapter': request} + engine._lora_load_tasks = {} + engine.engine = MagicMock() + engine.engine.remove_lora.return_value = True + + await engine.unload_lora_paths(['/adapter']) + + engine.engine.remove_lora.assert_called_once_with(7) + assert engine._lora_request_cache == {} + + asyncio.run(run()) + + +def test_unload_lora_removes_a_just_completed_load(): + async def run(): + engine = VLLMEngine.__new__(VLLMEngine) + request = MagicMock(lora_int_id=9) + load_task = asyncio.create_task(asyncio.sleep(0, result=request)) + await load_task + engine._lora_request_cache = {} + engine._lora_load_tasks = {'/adapter': load_task} + engine.engine = MagicMock() + engine.engine.remove_lora.return_value = None + + await engine.unload_lora_paths(['/adapter']) + + engine.engine.remove_lora.assert_called_once_with(9) + assert engine._lora_load_tasks == {} + + asyncio.run(run()) diff --git a/tests/server/config/test_server_config.py b/tests/server/config/test_server_config.py index 93644cbef..bebd63c9f 100644 --- a/tests/server/config/test_server_config.py +++ b/tests/server/config/test_server_config.py @@ -235,6 +235,32 @@ def test_launcher_accepts_typed_config() -> None: assert launcher.config is cfg +def test_data_plane_application_uses_its_own_strict_args_schema() -> None: + app = ApplicationSpec.model_validate({ + 'name': 'data-plane', + 'route_prefix': '/api/v1/data-plane', + 'import_path': 'data_plane', + 'args': { + 'config': { + 'backend': { + 'SimpleStorage': { + 'num_data_storage_units': 2, + }, + }, + }, + }, + }) + assert app.import_path == 'data_plane' + assert app.args.config['backend']['SimpleStorage']['num_data_storage_units'] == 2 + + with pytest.raises(ValidationError): + ApplicationSpec.model_validate({ + 'name': 'data-plane', + 'import_path': 'data_plane', + 'args': {'unknown': True}, + }) + + def test_cookbook_examples_load() -> None: """Migrated cookbook configs all parse with the new field names.""" here = Path(__file__).resolve().parents[3] diff --git a/tests/server/contract/client_api_baseline.json b/tests/server/contract/client_api_baseline.json index dcafe9470..65f9db147 100644 --- a/tests/server/contract/client_api_baseline.json +++ b/tests/server/contract/client_api_baseline.json @@ -1,4 +1,48 @@ { + "data_plane": { + "paths": { + "/twinkle/append": { + "POST": { + "operationId": "append_twinkle_append_post", + "parameters": [], + "responses": [ + "200", + "422" + ] + } + }, + "/twinkle/get": { + "POST": { + "operationId": "get_twinkle_get_post", + "parameters": [], + "responses": [ + "200", + "422" + ] + } + }, + "/twinkle/put": { + "POST": { + "operationId": "put_twinkle_put_post", + "parameters": [], + "responses": [ + "200", + "422" + ] + } + }, + "/twinkle/release": { + "POST": { + "operationId": "release_twinkle_release_post", + "parameters": [], + "responses": [ + "200", + "422" + ] + } + } + } + }, "gateway": { "paths": { "/asample": { @@ -733,6 +777,16 @@ ] } }, + "/twinkle/forward_from_data_plane": { + "POST": { + "operationId": "forward_from_data_plane_twinkle_forward_from_data_plane_post", + "parameters": [], + "responses": [ + "200", + "422" + ] + } + }, "/twinkle/forward_backward": { "POST": { "operationId": "forward_backward_twinkle_forward_backward_post", @@ -743,6 +797,16 @@ ] } }, + "/twinkle/forward_backward_from_data_plane": { + "POST": { + "operationId": "forward_backward_from_data_plane_twinkle_forward_backward_from_data_plane_post", + "parameters": [], + "responses": [ + "200", + "422" + ] + } + }, "/twinkle/forward_only": { "POST": { "operationId": "forward_only_twinkle_forward_only_post", @@ -753,6 +817,16 @@ ] } }, + "/twinkle/forward_only_from_data_plane": { + "POST": { + "operationId": "forward_only_from_data_plane_twinkle_forward_only_from_data_plane_post", + "parameters": [], + "responses": [ + "200", + "422" + ] + } + }, "/twinkle/get_state_dict": { "POST": { "operationId": "get_state_dict_twinkle_get_state_dict_post", @@ -793,6 +867,16 @@ ] } }, + "/twinkle/remove_adapter": { + "POST": { + "operationId": "remove_adapter_twinkle_remove_adapter_post", + "parameters": [], + "responses": [ + "200", + "422" + ] + } + }, "/twinkle/resume_from_checkpoint": { "POST": { "operationId": "resume_from_checkpoint_twinkle_resume_from_checkpoint_post", @@ -990,6 +1074,16 @@ ] } }, + "/twinkle/sample_to_data_plane": { + "POST": { + "operationId": "sample_to_data_plane_twinkle_sample_to_data_plane_post", + "parameters": [], + "responses": [ + "200", + "422" + ] + } + }, "/twinkle/sample_stream": { "POST": { "operationId": "sample_stream_twinkle_sample_stream_post", @@ -1009,6 +1103,16 @@ "422" ] } + }, + "/twinkle/unload_adapter_paths": { + "POST": { + "operationId": "unload_adapter_paths_twinkle_unload_adapter_paths_post", + "parameters": [], + "responses": [ + "200", + "422" + ] + } } } } diff --git a/tests/server/data_plane/test_proxy.py b/tests/server/data_plane/test_proxy.py new file mode 100644 index 000000000..550998db2 --- /dev/null +++ b/tests/server/data_plane/test_proxy.py @@ -0,0 +1,61 @@ +from __future__ import annotations + +import pytest + +from twinkle.server.data_plane.proxy import DataPlaneProxy +from twinkle_client.http.headers import H_AUTH, H_AUTH_TWINKLE, H_REQUEST_ID +from twinkle_client.types import DataRef + + +class _Response: + + def __init__(self, payload): + self.payload = payload + + def raise_for_status(self) -> None: + return None + + def json(self): + return self.payload + + +class _Client: + + def __init__(self): + self.calls = [] + + async def post(self, url, **kwargs): + self.calls.append((url, kwargs)) + if url.endswith('/get'): + return _Response({'rows': [{'value': 1}]}) + return _Response({ + 'ref_id': 'output', + 'size': 1, + 'fields': ['value'], + 'kind': 'model-output', + }) + + +@pytest.mark.asyncio +async def test_proxy_routes_by_data_ref_without_tenant_identity() -> None: + proxy = DataPlaneProxy.__new__(DataPlaneProxy) + proxy.base_url = 'http://data-plane' + proxy.client = _Client() + ref = DataRef(ref_id='input', size=1, fields=['value']) + + assert await proxy.get(ref, fields=['value']) == [{'value': 1}] + output = await proxy.put([{'value': 2}], kind='model-output') + appended = await proxy.append(ref, [{'value': 3}]) + + assert output.ref_id == 'output' + assert appended.ref_id == 'output' + get_headers = proxy.client.calls[0][1]['headers'] + put_headers = proxy.client.calls[1][1]['headers'] + append_headers = proxy.client.calls[2][1]['headers'] + assert proxy.client.calls[0][1]['json']['fields'] == ['value'] + assert get_headers[H_REQUEST_ID] == 'data-ref-input' + assert put_headers[H_REQUEST_ID] == 'data-put-model-output' + assert append_headers[H_REQUEST_ID] == 'data-append-input' + assert get_headers[H_AUTH] == get_headers[H_AUTH_TWINKLE] == '' + assert put_headers[H_AUTH] == put_headers[H_AUTH_TWINKLE] == '' + assert append_headers[H_AUTH] == append_headers[H_AUTH_TWINKLE] == '' diff --git a/tests/server/data_plane/test_store.py b/tests/server/data_plane/test_store.py new file mode 100644 index 000000000..3c7ec1633 --- /dev/null +++ b/tests/server/data_plane/test_store.py @@ -0,0 +1,117 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from __future__ import annotations + +import pytest + +from twinkle.server.data_plane.store import TQDataRefStore, _partition + + +@pytest.mark.asyncio +async def test_data_ref_round_trip_append_release_and_ref_isolation(monkeypatch) -> None: + import transfer_queue as tq + + records = {} + + async def batch_put(*, keys, partition_id, fields, tags=None): + storage_key = (partition_id, tuple(keys)) + if storage_key in records: + records[storage_key].update(fields) + else: + records[storage_key] = fields.clone() + if tags is not None: + current = tag_records.setdefault(storage_key, [{} for _ in tags]) + for existing, update in zip(current, tags): + existing.update(update) + + async def batch_get(*, keys, partition_id, select_fields): + data = records[(partition_id, tuple(keys))] + return data.select(*select_fields) + + async def clear(*, keys, partition_id): + records.pop((partition_id, tuple(keys))) + + async def kv_list(partition_id): + result = {} + for (stored_partition, keys), tags in tag_records.items(): + if stored_partition == partition_id: + result[stored_partition] = dict(zip(keys, tags)) + return result + + monkeypatch.setattr(tq, 'async_kv_batch_put', batch_put) + monkeypatch.setattr(tq, 'async_kv_batch_get', batch_get) + monkeypatch.setattr(tq, 'async_kv_clear', clear) + monkeypatch.setattr(tq, 'async_kv_list', kv_list) + + # Bypass tq.init(): this test exercises only the DataRef mapping layer. + store = TQDataRefStore.__new__(TQDataRefStore) + tag_records = {} + rows = [ + {'input_ids': [1, 2], 'answer': 'a'}, + {'input_ids': [3], 'answer': 'b'}, + ] + tags = [{'group_id': 'g0', 'generation_idx': 0}, {'group_id': 'g0', 'generation_idx': 1}] + ref = await store.put(rows, kind='train', tags=tags) + + assert ref.size == 2 + assert ref.fields == ['input_ids', 'answer'] + assert ref.num_tokens == 3 + assert await store.get(ref) == rows + assert await store.get_tags(ref) == tags + + ref = await store.append( + ref, + [{'reward': 1.0}, {'reward': -1.0}], + tags=[{'status': 'ready'}, {'status': 'ready'}], + ) + assert ref.fields == ['input_ids', 'answer', 'reward'] + assert ref.num_tokens == 3 + assert await store.get( + ref, + fields=['answer', 'reward'], + ) == [ + {'answer': 'a', 'reward': 1.0}, + {'answer': 'b', 'reward': -1.0}, + ] + assert await store.get_tags(ref) == [ + {'group_id': 'g0', 'generation_idx': 0, 'status': 'ready'}, + {'group_id': 'g0', 'generation_idx': 1, 'status': 'ready'}, + ] + + ref = await store.append( + ref, + [{'input_ids': [4, 5, 6]}, {'input_ids': [7, 8]}], + ) + assert ref.num_tokens == 5 + + nested_ref = await store.put([ + {'train_input': {'input_ids': [1, 2, 3]}}, + {'train_input': {'input_ids': [4]}}, + ], kind='rollout') + assert nested_ref.num_tokens == 4 + + with pytest.raises(KeyError): + await store.get(ref.model_copy(update={'ref_id': 'another-ref'})) + + await store.release(ref) + await store.release(nested_ref) + assert records == {} + + +@pytest.mark.asyncio +async def test_append_rejects_row_count_mismatch() -> None: + from twinkle_client.types import DataRef + + store = TQDataRefStore.__new__(TQDataRefStore) + ref = DataRef(ref_id='r', size=2, fields=['x']) + with pytest.raises(ValueError, match='row count'): + await store.append(ref, [{'reward': 1.0}]) + + +def test_partition_is_stable_and_scoped_by_data_ref() -> None: + from twinkle_client.types import DataRef + + first = DataRef(ref_id='a', size=1, fields=['x']) + same = DataRef(ref_id='a', size=99, fields=['other']) + other = DataRef(ref_id='b', size=1, fields=['x']) + assert _partition(first) == _partition(same) + assert _partition(first) != _partition(other) diff --git a/tests/server/model/test_twinkle_async_inputs.py b/tests/server/model/test_twinkle_async_inputs.py new file mode 100644 index 000000000..7a13b73c5 --- /dev/null +++ b/tests/server/model/test_twinkle_async_inputs.py @@ -0,0 +1,139 @@ +from __future__ import annotations + +import pytest +from fastapi import FastAPI +from starlette.requests import Request + +import twinkle_client.types as types +from twinkle.server.model.twinkle_handlers import _register_twinkle_routes +from twinkle.server.model.utils import model_result_rows + + +def test_model_result_rows_keeps_one_output_row_per_sample() -> None: + assert model_result_rows( + {'logps': [[-1.0], [-2.0]], 'loss': 0.25}, + batch_size=2, + ) == [ + {'logps': [-1.0], 'loss': 0.25}, + {'logps': [-2.0], 'loss': 0.25}, + ] + + +class _SchedulingManagement: + + def __init__(self): + self.data_world_size = 2 + self.scheduled = [] + self.model_calls = [] + self.model = self + self.data_plane = self + self.rows = { + 'data-a': [{ + 'train_input': {'input_ids': [index]}, + 'sampled_logprobs': [-0.1], + 'advantage': 1.0, + } for index in range(4)], + 'data-b': [{ + 'train_input': {'input_ids': [index]}, + 'sampled_logprobs': [-0.2], + 'advantage': -1.0, + } for index in range(4, 8)], + } + + async def _on_request_start(self, _request): + return 'token' + + def assert_resource_exists(self, _adapter_name): + return None + + def resolve_model_adapter_name(self, adapter_name): + return adapter_name + + def forward_backward(self, *, inputs, adapter_name, **kwargs): + self.model_calls.append((inputs, adapter_name, kwargs)) + return {'loss': 1.0} + + async def get(self, ref, *, fields=None): + rows = self.rows[ref.ref_id] + if fields is None: + return rows + return [{field: row[field] for field in fields} for row in rows] + + async def schedule_task_and_wait(self, task, **kwargs): + self.scheduled.append(kwargs) + return await task() + + +@pytest.mark.asyncio +async def test_forward_backward_resolves_multiple_data_refs_and_field_kwargs() -> None: + management = _SchedulingManagement() + app = FastAPI() + _register_twinkle_routes(app, lambda: management) + route = next( + route for route in app.routes + if getattr(route, 'path', None) == '/twinkle/forward_backward_from_data_plane' + ) + request = Request({'type': 'http', 'headers': []}) + request.state.session_id = 'session' + body = types.DataPlaneForwardRequest( + adapter_name='adapter', + input_refs=[ + types.DataRef(ref_id='data-a', size=4, num_tokens=4), + types.DataRef(ref_id='data-b', size=4, num_tokens=4), + ], + input_field='train_input', + kwarg_fields={ + 'old_logps': 'sampled_logprobs', + 'advantages': 'advantage', + }, + ) + + await route.endpoint(request, body, management) + + assert management.scheduled[-1]['batch_size'] == 8 + assert management.scheduled[-1]['data_world_size'] == 2 + inputs, adapter_name, forwarded_kwargs = management.model_calls[-1] + assert adapter_name == 'session-adapter' + assert [row['input_ids'] for row in inputs] == [[index] for index in range(8)] + assert forwarded_kwargs['old_logps'] == [[-0.1]] * 4 + [[-0.2]] * 4 + assert forwarded_kwargs['advantages'] == [1.0] * 4 + [-1.0] * 4 + + +@pytest.mark.asyncio +async def test_forward_backward_binds_nested_dpo_ref_logps_without_coercion() -> None: + management = _SchedulingManagement() + management.rows['dpo'] = [ + { + 'input_ids': [1, 2, 3], + 'labels': [-100, 2, 3], + 'ref_logps': [-0.1, -0.2, -0.3], + }, + { + 'input_ids': [1, 4, 5], + 'labels': [-100, 4, 5], + 'ref_logps': [-0.4, -0.5, -0.6], + }, + ] + app = FastAPI() + _register_twinkle_routes(app, lambda: management) + route = next( + route for route in app.routes + if getattr(route, 'path', None) == '/twinkle/forward_backward_from_data_plane' + ) + request = Request({'type': 'http', 'headers': []}) + request.state.session_id = 'session' + body = types.DataPlaneForwardRequest( + adapter_name='adapter', + input_refs=[types.DataRef(ref_id='dpo', size=2, num_tokens=6)], + kwarg_fields={'ref_outputs.logps': 'ref_logps'}, + ) + + await route.endpoint(request, body, management) + + inputs, adapter_name, forwarded_kwargs = management.model_calls[-1] + assert adapter_name == 'session-adapter' + assert [row['input_ids'] for row in inputs] == [[1, 2, 3], [1, 4, 5]] + assert forwarded_kwargs['ref_outputs']['logps'] == [ + [-0.1, -0.2, -0.3], + [-0.4, -0.5, -0.6], + ] diff --git a/tests/server/sampler/test_mock_sampler.py b/tests/server/sampler/test_mock_sampler.py index afa72ddac..8efae14b5 100644 --- a/tests/server/sampler/test_mock_sampler.py +++ b/tests/server/sampler/test_mock_sampler.py @@ -17,7 +17,7 @@ from twinkle.data_format import InputFeature, SamplingParams from twinkle.server.exceptions import ConfigError -from twinkle.server.sampler.app import SAMPLER_SELECTOR +from twinkle.server.sampler.app import SAMPLER_SELECTOR, _construct_sampler_backend from twinkle.server.sampler.backends.mock_sampler import MockSampler _SAMPLER_TYPES = tuple(SAMPLER_SELECTOR.builders) @@ -114,6 +114,41 @@ def test_mock_dispatch_returns_mock_sampler() -> None: assert isinstance(s, MockSampler) +def test_explicit_async_vllm_uses_non_blocking_sampler(monkeypatch) -> None: + from twinkle_agentic.async_rl import vllm_sampler_tq as module + + captured = {} + + def construct(**kwargs): + captured.update(kwargs) + return 'vllm-tq' + + monkeypatch.setattr(module, 'VLLMSamplerTQ', construct) + + sampler = _construct_sampler_backend( + 'vllm_async', + {'model_id': 'local-model'}, + None, + ) + + assert sampler == 'vllm-tq' + assert captured == {'model_id': 'local-model', 'context_manager': None} + + +def test_standard_vllm_is_independent_of_data_plane(monkeypatch) -> None: + calls = [] + monkeypatch.setattr( + SAMPLER_SELECTOR, + 'construct', + lambda sampler_type, kwargs: calls.append((sampler_type, kwargs)) or 'standard-vllm', + ) + + sampler = _construct_sampler_backend('vllm', {'model_id': 'local-model'}, 'http://data-plane') + + assert sampler == 'standard-vllm' + assert calls == [('vllm', {'model_id': 'local-model'})] + + @settings(max_examples=100) @given(bad=st.text(min_size=1, max_size=10).filter(lambda s: s not in _SAMPLER_TYPES)) def test_invalid_sampler_type_raises_config_error(bad: str) -> None: diff --git a/tests/server/sampler/test_twinkle_async_rows.py b/tests/server/sampler/test_twinkle_async_rows.py new file mode 100644 index 000000000..30b7016b2 --- /dev/null +++ b/tests/server/sampler/test_twinkle_async_rows.py @@ -0,0 +1,145 @@ +from __future__ import annotations + +import pytest +from fastapi import FastAPI +from starlette.requests import Request + +import twinkle_client.types as types +from twinkle.data_format import SampledSequence, SampleResponse +from twinkle.server.sampler.twinkle_handlers import ( + _register_twinkle_sampler_routes, + _build_rollout_rows_and_tags, +) + + +def _response(tokens: list[int]) -> types.SampleResponseModel: + return types.SampleResponseModel( + sequences=[ + types.SampledSequenceModel( + stop_reason='stop', + tokens=[token], + logprobs=[[(token, -0.1)]], + new_input_feature={'input_ids': [token], 'labels': [token]}, + ) + for token in tokens + ], + prompt_logprobs=[-0.2], + ) + + +def test_async_sampler_flattens_generations_to_tagged_tq_rows() -> None: + rows, tags = _build_rollout_rows_and_tags( + [_response([10, 11]), _response([20, 21])], + group_ids=['group-a', 'group-b'], + policy_version=7, + adapter_uri='twinkle://policy-7', + ) + + assert [row['tokens'] for row in rows] == [[10], [11], [20], [21]] + assert [row['sampled_logprobs'] for row in rows] == [[-0.1]] * 4 + assert [row['train_input']['input_ids'] for row in rows] == [[10], [11], [20], [21]] + assert all('new_input_feature' not in row for row in rows) + assert [(tag['group_id'], tag['generation_idx']) for tag in tags] == [ + ('group-a', 0), + ('group-a', 1), + ('group-b', 0), + ('group-b', 1), + ] + assert {tag['rollout_policy_version'] for tag in tags} == {7} + assert {tag['rollout_adapter_uri'] for tag in tags} == {'twinkle://policy-7'} + + +def test_async_sampler_rejects_group_id_count_mismatch() -> None: + with pytest.raises(ValueError, match='group_ids contains 1 values for 2'): + _build_rollout_rows_and_tags( + [_response([10]), _response([20])], + group_ids=['only-one'], + policy_version=0, + adapter_uri=None, + ) + + +class _SamplerManagement: + + def __init__(self): + self.sampler = self + self.data_plane = self + self.enabled = True + self.scheduled = [] + self.put_rows = None + + async def _on_request_start(self, _request): + return 'token' + + async def schedule_task_and_wait(self, task, **kwargs): + self.scheduled.append(kwargs) + return await task() + + def submit_generation(self, submission_id, inputs, params, **kwargs): + self.submission_id = submission_id + self.inputs = inputs + self.params = params + self.generation_kwargs = kwargs + + def get_generation_status(self, submission_id): + assert submission_id == self.submission_id + return {'status': 'completed'} + + def collect_generation(self, submission_id): + assert submission_id == self.submission_id + return [SampleResponse(sequences=[SampledSequence( + stop_reason='stop', + tokens=[7], + logprobs=[[(7, -0.25)]], + decoded='answer', + new_input_feature={'input_ids': [1, 7], 'labels': [-100, 7]}, + )])] + + def cancel_generation(self, _submission_id): + raise AssertionError('completed generation must not be cancelled') + + async def put(self, rows, *, kind, tags): + self.put_rows = rows + self.put_tags = tags + return types.DataRef(ref_id='rollout-ref', size=len(rows), fields=list(rows[0]), kind=kind) + + +@pytest.mark.asyncio +async def test_sample_to_data_plane_returns_ref_after_short_admission() -> None: + management = _SamplerManagement() + app = FastAPI() + _register_twinkle_sampler_routes(app, lambda: management) + route = next( + route for route in app.routes + if getattr(route, 'path', None) == '/twinkle/sample_to_data_plane' + ) + request = Request({'type': 'http', 'headers': []}) + request.state.session_id = 'session' + body = types.DataPlaneSampleRequest( + inputs=[{'input_ids': [1]}], + adapter_name='adapter', + group_ids=['group-1'], + policy_version=3, + num_samples=1, + sampling_params={'max_tokens': 4}, + ) + + ref = await route.endpoint(request, body, management) + + assert ref.ref_id == 'rollout-ref' + assert management.scheduled == [{ + 'model_id': 'session-adapter', + 'token': 'token', + 'input_tokens': 1, + 'task_type': 'sample_admission', + }] + assert management.put_rows == [{ + 'train_input': {'input_ids': [1, 7], 'labels': [-100, 7]}, + 'sampled_logprobs': [-0.25], + 'tokens': [7], + 'decoded': 'answer', + 'stop_reason': 'stop', + 'prompt_logprobs': None, + 'topk_prompt_logprobs': None, + }] + assert management.put_tags[0]['group_id'] == 'group-1' diff --git a/tests/server/test_app_builders_characterization.py b/tests/server/test_app_builders_characterization.py index 31f9e7267..d2a7f5523 100644 --- a/tests/server/test_app_builders_characterization.py +++ b/tests/server/test_app_builders_characterization.py @@ -1,8 +1,9 @@ # Copyright (c) ModelScope Contributors. All rights reserved. -"""Characterization tests for the four App_Builders. +"""Characterization tests for the component App_Builders. These freeze the externally observable behavior of ``build_gateway_app``, -``build_model_app``, ``build_sampler_app`` and ``build_processor_app`` BEFORE +``build_model_app``, ``build_sampler_app``, ``build_processor_app``, and the +DataPlane builder. The original four were captured before the Shared_App_Scaffold is extracted, so the extraction can be shown to be behavior-preserving. For each builder they assert, as fixed expectations: @@ -15,7 +16,7 @@ stack is wired in front of the routes; 3. the **bound deployment** identified by the name passed to ``serve.deployment`` (``GatewayServer``, ``ModelManagement``, ``SamplerManagement``, - ``ProcessorManagement``). + ``ProcessorManagement``, ``DataPlaneManagement``). They MUST NOT assert internal object identity or internal middleware-stack structure. They are built by capturing the FastAPI app the builder @@ -187,6 +188,29 @@ def test_processor_builder_characterization(monkeypatch) -> None: _assert_middleware_lifo_order(res.app, expect_cleanup=False) +def test_data_plane_builder_characterization(monkeypatch) -> None: + from twinkle.server.data_plane import app as data_plane_mod + + res = _capture_builder( + monkeypatch, + data_plane_mod.build_data_plane_app, + deploy_options={}, + ) + assert res.deployment_name == 'DataPlaneManagement' + assert _route_set(res.app) == _baseline_route_set('data_plane') + _assert_auth_middleware_effect(res.app) + _assert_middleware_lifo_order(res.app, expect_cleanup=False) + + +def test_data_plane_builder_rejects_multiple_replicas() -> None: + from twinkle.server.data_plane import app as data_plane_mod + + with pytest.raises(ValueError, match='exactly one replica'): + data_plane_mod.build_data_plane_app( + deploy_options={'num_replicas': 2}, + ) + + # ----- black-box middleware-effect oracle ---------------------------------- # diff --git a/tests/server/utils/test_task_queue_mixin.py b/tests/server/utils/test_task_queue_mixin.py index c579a3b9a..f0bdbf963 100644 --- a/tests/server/utils/test_task_queue_mixin.py +++ b/tests/server/utils/test_task_queue_mixin.py @@ -1,7 +1,10 @@ +import asyncio + import pytest from twinkle.server.utils.task_queue.config import TaskQueueConfig from twinkle.server.utils.task_queue.mixin import TaskQueueMixin +from twinkle.server.utils.task_queue.worker import ComputeWorker class _DummyState: @@ -28,6 +31,15 @@ def __init__(self): self._task_metrics = None self._deployment_name = 'test' + def enable_compute_worker(self): + self._compute_worker = ComputeWorker( + state=self.state, + config=self._task_queue_config, + task_metrics=None, + deployment_name=self._deployment_name, + ) + self._event_loop = None + @pytest.mark.asyncio async def test_preflight_rejects_batch_without_per_dp_multiple(): @@ -65,3 +77,113 @@ async def test_preflight_accepts_batch_with_per_dp_multiple(): assert result is None assert queue.state.records == [] + + +@pytest.mark.asyncio +async def test_background_task_tracks_status(): + queue = _DummyQueue() + + async def work(): + return {'ok': True} + + await queue.schedule_background_task( + work, + model_id='model1', + ) + await asyncio.sleep(0) + + assert [args[1] for args, _ in queue.state.records] == ['running', 'completed'] + assert queue.state.records[-1][1]['result'] == {'ok': True} + + +@pytest.mark.asyncio +async def test_schedule_task_and_wait_returns_large_result_without_persisting_it(): + queue = _DummyQueue() + queue.enable_compute_worker() + result = {'logps': [[float(index) for index in range(128)]]} + + async def work(): + return result + + try: + actual = await queue.schedule_task_and_wait( + work, + model_id='model1', + token='token1', + task_type='forward_backward', + ) + finally: + await queue._compute_worker.stop() + + assert actual is result + assert queue.state.records == [] + + +@pytest.mark.asyncio +async def test_polling_schedule_task_still_persists_its_result(): + queue = _DummyQueue() + queue.enable_compute_worker() + result = {'value': 42} + + async def work(): + return result + + try: + await queue.schedule_task(work, model_id='model1', token='token1') + for _ in range(100): + completed = [ + kwargs + for args, kwargs in queue.state.records + if args[1] == 'completed' + ] + if completed: + break + await asyncio.sleep(0) + finally: + await queue._compute_worker.stop() + + assert completed[-1]['result'] is result + + +@pytest.mark.asyncio +async def test_schedule_task_and_wait_propagates_failure_without_persisting_it(): + queue = _DummyQueue() + queue.enable_compute_worker() + + async def work(): + raise ValueError('model failed') + + try: + with pytest.raises(RuntimeError, match='ValueError: model failed'): + await queue.schedule_task_and_wait( + work, + model_id='model1', + token='token1', + task_type='forward_backward', + ) + finally: + await queue._compute_worker.stop() + + assert queue.state.records == [] + + +@pytest.mark.asyncio +async def test_schedule_task_and_wait_reports_preflight_failure_without_persisting_it(): + queue = _DummyQueue() + queue.enable_compute_worker() + + async def work(): + raise AssertionError('preflight rejection must not execute the task') + + with pytest.raises(RuntimeError, match='Batch size 2 must be divisible by 4'): + await queue.schedule_task_and_wait( + work, + model_id='model1', + token='token1', + batch_size=2, + data_world_size=2, + batch_size_multiple=2, + ) + + assert queue.state.records == [] + assert queue._compute_worker._worker_task is None diff --git a/tests/transformers/test_native_fsdp_strategy.py b/tests/transformers/test_native_fsdp_strategy.py new file mode 100644 index 000000000..79402c785 --- /dev/null +++ b/tests/transformers/test_native_fsdp_strategy.py @@ -0,0 +1,25 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from unittest.mock import MagicMock, patch + +import torch + +from twinkle.model.transformers.strategy.native_fsdp import NativeFSDPStrategy + + +def test_single_rank_native_fsdp_places_model_on_local_device(): + device_mesh = MagicMock() + device_mesh.mesh_dim_names = None + model = MagicMock() + moved_model = MagicMock() + model.to.return_value = moved_model + strategy = NativeFSDPStrategy(device_mesh=device_mesh, enable_ep=False) + + with patch( + 'twinkle.model.transformers.strategy.native_fsdp.Platform.get_local_device', + return_value='cuda:0', + ): + wrapped_model, wrapped_optimizer = strategy.wrap_model(model) + + model.to.assert_called_once_with(torch.device('cuda:0')) + assert wrapped_model is moved_model + assert wrapped_optimizer is None diff --git a/tests/twinkle_agentic/test_async_rl_config.py b/tests/twinkle_agentic/test_async_rl_config.py new file mode 100644 index 000000000..64d6f5760 --- /dev/null +++ b/tests/twinkle_agentic/test_async_rl_config.py @@ -0,0 +1,226 @@ +from __future__ import annotations + +import pytest + +from twinkle_agentic.async_rl.pipeline import _reward_for_context +from twinkle_agentic.async_rl.utils import ( + TrainBatchConfig, + build_native_fsdp_model_kwargs, + configure_lora_lr_scheduler, + resolve_context_learning_rate, + resolve_context_lora_target_modules, + resolve_context_loss_config, + resolve_model_attention_implementation, + resolve_sequence_parallel_size, + sampler_data_parallel_size, + validate_context_batch_config, +) + + +def test_sequence_parallel_size_must_divide_model_gpus(): + assert resolve_sequence_parallel_size(2, 1) == 1 + assert resolve_sequence_parallel_size(2, 2) == 2 + + with pytest.raises(ValueError, match='must be divisible'): + resolve_sequence_parallel_size(2, 3) + + +def test_padding_free_sequence_parallel_requires_flash_attention(): + assert resolve_model_attention_implementation( + {'attn_implementation': 'flash_attention_2'}, + padding_free=True, + sequence_parallel_size=2, + ) == 'flash_attention_2' + + with pytest.raises(ValueError, match='model.attn_implementation'): + resolve_model_attention_implementation({}, padding_free=True, sequence_parallel_size=2) + + +@pytest.mark.parametrize( + ('sampler_gpus', 'sampler_tp', 'expected_dp'), + [(8, 2, 4), (1, 1, 1)], +) +def test_sampler_data_parallel_size(sampler_gpus, sampler_tp, expected_dp): + assert sampler_data_parallel_size(sampler_gpus, sampler_tp) == expected_dp + + +def test_sampler_parallelism_rejects_incomplete_tp_group(): + with pytest.raises(ValueError, match='must be divisible'): + sampler_data_parallel_size(3, 2) + + +def test_lora_lr_scheduler_uses_shared_adapter_config(): + calls = [] + + class Model: + def set_lr_scheduler(self, scheduler_cls, **kwargs): + calls.append((scheduler_cls, kwargs)) + + configure_lora_lr_scheduler( + Model(), + 'tenant_lora', + { + 'lr_scheduler': { + 'cls': 'CosineAnnealingLR', + 'T_max': 2000, + 'eta_min': 0.0, + }, + }, + ) + + assert calls == [('CosineAnnealingLR', { + 'adapter_name': 'tenant_lora', + 'T_max': 2000, + 'eta_min': 0.0, + })] + + +def test_context_learning_rate_overrides_global_default(): + assert resolve_context_learning_rate({'learning_rate': 5e-6}, {'learning_rate': 1e-6}) == pytest.approx(5e-6) + assert resolve_context_learning_rate({}, {'learning_rate': 1e-6}) == pytest.approx(1e-6) + + +@pytest.mark.parametrize('value', [0, -1e-6, float('inf')]) +def test_context_learning_rate_rejects_invalid_values(value): + with pytest.raises(ValueError, match='positive finite'): + resolve_context_learning_rate({'learning_rate': value}, {'learning_rate': 1e-6}) + + +def test_context_lora_target_modules_override_global_default(): + defaults = {'target_modules': 'all-linear'} + + assert resolve_context_lora_target_modules({}, defaults) == 'all-linear' + assert resolve_context_lora_target_modules( + {'lora': {'target_modules': ['q_proj', 'v_proj']}}, + defaults, + ) == ['q_proj', 'v_proj'] + + +@pytest.mark.parametrize('value', ['', [], [None], {'q_proj': True}]) +def test_context_lora_target_modules_reject_invalid_values(value): + with pytest.raises(ValueError, match='target_modules'): + resolve_context_lora_target_modules( + {'lora': {'target_modules': value}}, + {'target_modules': 'all-linear'}, + ) + + +def test_context_loss_config_overrides_global_defaults(): + loss_cls, loss_kwargs = resolve_context_loss_config( + {'loss': {'cls': 'GSPOLoss', 'epsilon_high': 0.3}}, + {'cls': 'GRPOLoss', 'epsilon': 0.2}, + ) + + assert loss_cls == 'GSPOLoss' + assert loss_kwargs == {'epsilon': 0.2, 'epsilon_high': 0.3} + + +def test_context_loss_config_uses_grpo_defaults(): + assert resolve_context_loss_config({}) == ('GRPOLoss', {'epsilon': 0.2}) + + +def test_context_loss_config_rejects_empty_class_name(): + with pytest.raises(ValueError, match='loss.cls'): + resolve_context_loss_config({'loss': {'cls': ''}}) + + +def test_rl_model_kwargs_enforce_native_fsdp(): + assert build_native_fsdp_model_kwargs({}) == { + 'strategy': 'native_fsdp', + 'fsdp_config': {}, + } + assert build_native_fsdp_model_kwargs({ + 'strategy': 'native_fsdp', + 'fsdp_config': {'reshard_after_forward': False}, + }) == { + 'strategy': 'native_fsdp', + 'fsdp_config': {'reshard_after_forward': False}, + } + with pytest.raises(ValueError, match='must be native_fsdp'): + build_native_fsdp_model_kwargs({'strategy': 'accelerate'}) + + +def test_reward_factory_loads_class_and_resolved_kwargs(): + reward = _reward_for_context( + { + 'class_path': 'twinkle.reward.DAPOMathReward', + 'kwargs': { + 'max_response_length': 8192, + 'overlong_buffer_length': 4096, + 'overlong_penalty_factor': 1.0, + 'score_tail_chars': 300, + }, + }, + context_key='tenant/run/adapter', + ) + + assert reward.max_response_length == 8192 + assert reward.overlong_buffer_length == 4096 + + +def test_reward_factory_rejects_non_reward_class(): + with pytest.raises(TypeError, match='Reward subclass'): + _reward_for_context( + {'class_path': 'collections.Counter'}, + context_key='tenant/run/adapter', + ) + + +def test_context_batch_config_accepts_group_aligned_dp_batches(): + validate_context_batch_config( + 'tenant/run/adapter', + rollout_groups=8, + num_generations=4, + train=TrainBatchConfig(mini_batch_size=8, micro_batch_size=2), + sampler_dp=2, + model_dp=2, + ) + + +def test_context_batch_config_allows_training_group_to_span_model_dp_ranks(): + validate_context_batch_config( + 'tenant/run/adapter', + rollout_groups=8, + num_generations=4, + train=TrainBatchConfig(mini_batch_size=16, micro_batch_size=2), + sampler_dp=1, + model_dp=8, + ) + + +def test_context_batch_config_rejects_partition_tail_and_undersized_rank_batch(): + with pytest.raises(ValueError, match='complete prompt groups'): + validate_context_batch_config( + 'tenant/run/adapter', + rollout_groups=6, + num_generations=4, + train=TrainBatchConfig(mini_batch_size=6, micro_batch_size=1), + sampler_dp=2, + model_dp=2, + ) + + with pytest.raises(ValueError, match='per-rank train batch'): + validate_context_batch_config( + 'tenant/run/adapter', + rollout_groups=8, + num_generations=2, + train=TrainBatchConfig(mini_batch_size=2, micro_batch_size=2), + sampler_dp=2, + model_dp=2, + ) + + +def test_context_batch_config_requires_token_limit_for_dynamic_batching(): + with pytest.raises(ValueError, match='max_tokens_per_micro_batch'): + validate_context_batch_config( + 'tenant/run/adapter', + rollout_groups=8, + num_generations=4, + train=TrainBatchConfig( + mini_batch_size=8, + micro_batch_size=2, + dynamic_batching=True, + ), + sampler_dp=1, + model_dp=1, + ) diff --git a/tests/twinkle_agentic/test_async_rl_data_plane.py b/tests/twinkle_agentic/test_async_rl_data_plane.py new file mode 100644 index 000000000..05208849b --- /dev/null +++ b/tests/twinkle_agentic/test_async_rl_data_plane.py @@ -0,0 +1,130 @@ +from __future__ import annotations + +import asyncio + +import pytest + +from twinkle_agentic.async_rl import LoraContext, TQDataPlane +from twinkle_agentic.async_rl.data_plane import build_rollout_group_sample_write +from twinkle_agentic.async_rl.types import PartitionAdmission, PromptGroup + + +def _context() -> LoraContext: + return LoraContext('tenant', 'run_adapter', 'model', 'adapter') + + +def test_rollout_sample_tags_use_new_context_descriptor_only(): + context = _context() + admission = PartitionAdmission(context, context.partition_id(0), 0, 1, 2, 0) + group = PromptGroup(context, admission, f'{admission.partition_id}/group_0', {}, batch_meta=None) + fields, tags = build_rollout_group_sample_write( + group, + [ + { + 'generation_idx': 0, + 'labels': [-100, 1], + 'logprobs': [-.1], + 'rollout_policy_version': 3, + 'rollout_adapter_path': 'adapter-v3', + }, + { + 'generation_idx': 1, + 'labels': [-100, 2], + 'logprobs': [-.2], + 'rollout_policy_version': 4, + 'rollout_adapter_path': 'adapter-v4', + }, + ], + rewards=[1., 0.], + expected_num_generations=2, + ) + assert [row['rewards'] for row in fields] == [1., 0.] + assert [tag['generation_idx'] for tag in tags] == [0, 1] + assert all(tag['context_key'] == context.key for tag in tags) + assert [tag['rollout_policy_version'] for tag in tags] == [3, 4] + + +def test_data_plane_completes_rollout_with_full_training_trajectory(): + class Metadata: + def __init__(self): + self.size = 2 + self.custom_meta = [{}, {}] + + def update_custom_meta(self, updates): + for tag, update in zip(self.custom_meta, updates): + tag.update(update) + + class Client: + def __init__(self): + self.written = None + self.calls = [] + + async def async_put(self, data, metadata=None, partition_id=None): + self.calls.append('fields') + self.written = data + return metadata + + async def async_set_custom_meta(self, _metadata): + self.calls.append('tags') + + context = _context() + admission = PartitionAdmission(context, context.partition_id(0), 0, 1, 2, 0) + metadata = Metadata() + group = PromptGroup(context, admission, f'{admission.partition_id}/group_0', {}, metadata) + client = Client() + rows = [{ + 'input_ids': [1, 2, token], + 'labels': [-100, -100, token], + 'attention_mask': [1, 1, 1], + 'position_ids': [0, 1, 2], + 'logprobs': [-.1], + 'generation_idx': generation_idx, + 'rollout_policy_version': 3, + 'rollout_policy_versions': [3], + 'initial_policy_version': 3, + 'final_policy_version': 3, + 'policy_version_span': 0, + 'rollout_adapter_path': 'adapter-v3', + 'completion_length': 1, + } for generation_idx, token in enumerate((7, 8))] + + asyncio.run( + TQDataPlane(client).complete_rollout_group( + group, + rollout_rows=rows, + rewards=[1., 0.], + submission_id='submission', + )) + + assert set(client.written.keys()) == { + 'input_ids', 'labels', 'attention_mask', 'position_ids', 'logprobs', 'rewards' + } + assert client.calls == ['tags', 'fields'] + assert [tag['rollout_status'] for tag in metadata.custom_meta] == ['ROLLOUT_DONE', 'ROLLOUT_DONE'] + assert [tag['submission_id'] for tag in metadata.custom_meta] == ['submission', 'submission'] + + +def test_data_plane_rejects_rollout_without_complete_model_inputs(): + context = _context() + admission = PartitionAdmission(context, context.partition_id(0), 0, 1, 1, 0) + metadata = type('Metadata', (), {'size': 1})() + group = PromptGroup(context, admission, f'{admission.partition_id}/group_0', {}, metadata) + row = { + 'input_ids': [1, 2], + 'labels': [-100, 2], + 'logprobs': [-.1], + 'generation_idx': 0, + 'rollout_policy_version': 0, + } + + with pytest.raises(ValueError) as error: + asyncio.run( + TQDataPlane(object()).complete_rollout_group( + group, + rollout_rows=[row], + rewards=[1.], + submission_id='submission', + )) + + assert 'attention_mask' in str(error.value) + assert 'position_ids' in str(error.value) diff --git a/tests/twinkle_agentic/test_async_rl_metrics.py b/tests/twinkle_agentic/test_async_rl_metrics.py new file mode 100644 index 000000000..34f1b60e7 --- /dev/null +++ b/tests/twinkle_agentic/test_async_rl_metrics.py @@ -0,0 +1,343 @@ +from __future__ import annotations + +import json +import sys +import threading +import time +import types + +import pytest + +from twinkle.metric import ( + CompletionRewardMetric, + MetricBuffer, + MetricRecord, + create_metrics_reporter, +) +from twinkle_agentic.async_rl.metrics import advantage_signal_metrics, rollout_metrics +from twinkle.metric.reporting import MetricsReporter, _QueuedBackend + + +def test_advantage_signal_metrics_report_zero_and_nonzero_groups(): + metrics = advantage_signal_metrics( + rewards=[1.0, 1.0, 0.0, 1.0], + advantages=[0.0, 0.0, -1.0, 1.0], + num_generations=2, + ) + assert metrics['group_count'] == 2 + assert metrics['group_reward_std_mean'] == pytest.approx(0.25) + assert metrics['zero_advantage_group_ratio'] == pytest.approx(0.5) + assert metrics['positive_advantage_ratio'] == pytest.approx(0.25) + + +def test_advantage_signal_metrics_reject_incomplete_groups(): + with pytest.raises(ValueError, match='complete groups'): + advantage_signal_metrics([1.0, 0.0, 1.0], [1.0, -1.0, 0.0], num_generations=2) + + +def test_rollout_metrics_include_rewards_tokens_and_truncation(): + metrics = rollout_metrics( + rewards={'accuracy': [1.0, 0.0]}, + completion_lengths=[3, 7], + stop_reasons=['stop', 'length'], + rollout_latency_s=2.0, + ) + assert metrics == { + 'sample_count': 2, + 'completion_length_mean': 5.0, + 'completion_length_p95': 7, + 'completion_length_max': 7, + 'completion_truncated_count': 1, + 'completion_truncated_ratio': 0.5, + 'output_tokens': 10, + 'rollout_latency_s': 2.0, + 'output_tokens_per_s': 5.0, + 'accuracy_reward': 0.5, + 'accuracy_reward_std': pytest.approx(2**-0.5), + } + + +def test_completion_reward_metric_preserves_model_metric_contract(): + metric = CompletionRewardMetric() + metric.accumulate( + rewards={'accuracy': [1.0, 0.0]}, + completion_lengths=[3, 7], + generate_time=2.0, + weight_sync_time=0.25, + ) + result = metric.calculate() + assert result == { + 'profiling/Time taken: move_model_to_sampler': 0.25, + 'profiling/Time taken: generate': 2.0, + 'train/accuracy_reward': 0.5, + 'train/accuracy_reward_std': pytest.approx(2**-0.5), + 'train/completion_length': 5.0, + } + + +def test_metric_buffer_drain_is_atomic_and_destructive(): + buffer = MetricBuffer() + + def produce(start): + for value in range(start, start + 50): + buffer.record(MetricRecord(stage='train', values={'loss': value})) + + threads = [threading.Thread(target=produce, args=(index * 50,)) for index in range(4)] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + records = buffer.drain() + assert len(records) == 200 + assert buffer.drain() == [] + + +def test_reporter_writes_new_jsonl_schema_and_summary(tmp_path): + path = tmp_path / 'metrics.jsonl' + summary_path = tmp_path / 'summary.json' + reporter = create_metrics_reporter({ + 'queue_capacity': 100, + 'jsonl': { + 'path': path, + 'summary_path': summary_path, + 'batch_size': 2, + 'flush_interval_s': 60, + }, + }, run_id='test') + reporter.record(MetricRecord( + stage='rollout', + status='completed', + context_key='tenant/run/adapter', + partition_id='tenant/run/adapter/train_5', + partition_index=5, + policy_version=0, + values={'sample_count': 4, 'reward': 0.5}, + attributes={'scope': 'group', 'group_id': 'group_0'}, + )) + reporter.record(MetricRecord( + stage='rollout', + status='completed', + context_key='tenant/run/adapter', + partition_id='tenant/run/adapter/train_5', + partition_index=5, + policy_version=0, + values={'sample_count': 4, 'reward': 0.25}, + attributes={'scope': 'partition'}, + )) + reporter.record(MetricRecord( + stage='train', + context_key='tenant/run/adapter', + partition_id='tenant/run/adapter/train_5', + partition_index=5, + optimizer_step=1, + policy_version=0, + values={'sample_count': 4, 'loss': '0.25'}, + )) + reporter.record(MetricRecord( + stage='partition', + context_key='tenant/run/adapter', + partition_index=5, + policy_version=1, + values={}, + )) + reporter.record(MetricRecord(stage='run', values={'trained_partitions': 1, 'wall_time_s': 10.0})) + reporter.close() + + records = [json.loads(line) for line in path.read_text().splitlines()] + assert [record['sequence'] for record in records] == [1, 2, 3, 4, 5] + assert records[2]['stage'] == 'train' + assert records[2]['optimizer_step'] == 1 + assert records[2]['values']['loss'] == 0.25 + assert 'event' not in records[2] + summary = json.loads(summary_path.read_text()) + assert summary['status'] == 'completed' + assert summary['trained_partitions'] == 1 + assert summary['per_context']['tenant/run/adapter']['optimizer_step'] == 1 + assert summary['metrics']['train/loss']['mean'] == 0.25 + assert summary['metrics']['rollout/reward']['count'] == 1 + assert summary['metrics']['rollout/reward']['mean'] == 0.5 + + +def test_summary_policy_version_does_not_regress_for_out_of_order_records(): + reporter = MetricsReporter(run_id='test') + reporter.record(MetricRecord( + stage='policy', + context_key='tenant/run/adapter', + policy_version=5, + values={}, + )) + reporter.record(MetricRecord( + stage='rollout', + context_key='tenant/run/adapter', + policy_version=3, + values={'sample_count': 4}, + attributes={'scope': 'group'}, + )) + + assert reporter.summary()['per_context']['tenant/run/adapter']['policy_version'] == 5 + reporter.close() + + +def test_jsonl_backend_waits_for_batch_threshold(tmp_path): + path = tmp_path / 'metrics.jsonl' + reporter = create_metrics_reporter({ + 'jsonl': { + 'path': path, + 'summary_path': tmp_path / 'summary.json', + 'batch_size': 2, + 'flush_interval_s': 60, + }, + }, run_id='test') + reporter.record(MetricRecord(stage='train', values={'loss': 1.0})) + time.sleep(0.05) + assert path.read_text() == '' + reporter.record(MetricRecord(stage='train', values={'loss': 2.0})) + deadline = time.monotonic() + 1 + while not path.read_text() and time.monotonic() < deadline: + time.sleep(0.01) + assert len(path.read_text().splitlines()) == 2 + reporter.close() + + +def test_reporter_assigns_monotonic_sequence_across_threads(tmp_path): + reporter = create_metrics_reporter({ + 'jsonl': { + 'path': tmp_path / 'metrics.jsonl', + 'summary_path': tmp_path / 'summary.json', + }, + }, run_id='test') + + def produce(): + for _ in range(25): + reporter.record(MetricRecord(stage='train', values={'loss': 1.0})) + + threads = [threading.Thread(target=produce) for _ in range(4)] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + reporter.close() + records = [json.loads(line) for line in (tmp_path / 'metrics.jsonl').read_text().splitlines()] + assert [record['sequence'] for record in records] == list(range(1, 101)) + + +def test_swanlab_backend_uses_sequence_and_context_metric_names(monkeypatch, tmp_path): + logged = [] + + class Run: + def log(self, values, step): + logged.append((values, step)) + + fake_swanlab = types.SimpleNamespace(init=lambda **_kwargs: Run(), finish=lambda: None) + monkeypatch.setitem(sys.modules, 'swanlab', fake_swanlab) + reporter = create_metrics_reporter({ + 'jsonl': {'enabled': False}, + 'swanlab': { + 'enabled': True, + 'mode': 'local', + 'project': 'test', + 'name': 'test-run', + 'log_dir': tmp_path, + 'batch_size': 1, + }, + }, run_id='test') + reporter.record(MetricRecord( + stage='train', + context_key='tenant/run/adapter', + optimizer_step=3, + policy_version=2, + partition_index=1, + values={'loss': 0.5}, + )) + reporter.close() + values, sequence = logged[0] + assert sequence == 1 + assert values['context/tenant_run_adapter/train/loss'] == 0.5 + assert values['context/tenant_run_adapter/train/optimizer_step'] == 3 + assert values['context/tenant_run_adapter/policy/version'] == 2 + assert values['context/tenant_run_adapter/partition/index'] == 1 + + +def test_backend_queue_drops_oldest_record_without_blocking_reporter(): + release = threading.Event() + + class BlockingBackend(_QueuedBackend): + def _write_batch(self, batch): + release.wait(2) + + backend = BlockingBackend( + 'blocking', + queue_capacity=2, + batch_size=2, + flush_interval_s=60, + ) + reporter = MetricsReporter(run_id='test', backends=[backend]) + for index in range(5): + reporter.record(MetricRecord(stage='train', values={'loss': index})) + assert reporter.health()['backends']['blocking']['dropped_records'] >= 1 + release.set() + reporter.close() + + +def test_backend_failure_is_nonfatal_and_reported(): + class FailedBackend(_QueuedBackend): + def _write_batch(self, batch): + raise OSError('disk unavailable') + + backend = FailedBackend( + 'failed', + queue_capacity=4, + batch_size=1, + flush_interval_s=1, + ) + reporter = MetricsReporter(run_id='test', backends=[backend]) + reporter.record(MetricRecord(stage='train', values={'loss': 1.0})) + reporter.flush() + reporter.record(MetricRecord(stage='train', values={'loss': 2.0})) + health = reporter.health()['backends']['failed'] + assert health['enabled'] is False + assert health['failure_count'] == 1 + assert 'disk unavailable' in health['last_error'] + reporter.close() + + +def test_swanlab_failure_does_not_prevent_jsonl(monkeypatch, tmp_path): + class FailedRun: + def log(self, values, step): + raise RuntimeError('swanlab unavailable') + + fake_swanlab = types.SimpleNamespace(init=lambda **_kwargs: FailedRun(), finish=lambda: None) + monkeypatch.setitem(sys.modules, 'swanlab', fake_swanlab) + path = tmp_path / 'metrics.jsonl' + reporter = create_metrics_reporter({ + 'jsonl': { + 'path': path, + 'summary_path': tmp_path / 'summary.json', + 'batch_size': 1, + }, + 'swanlab': { + 'enabled': True, + 'mode': 'local', + 'log_dir': tmp_path, + 'batch_size': 1, + }, + }, run_id='test') + reporter.record(MetricRecord(stage='train', values={'loss': 1.0})) + reporter.close() + assert len(path.read_text().splitlines()) == 1 + assert reporter.health()['backends']['swanlab']['failure_count'] == 1 + + +def test_jsonl_startup_failure_is_nonfatal(tmp_path): + reporter = create_metrics_reporter({ + 'jsonl': { + 'path': tmp_path, + 'summary_path': tmp_path / 'summary.json', + }, + }, run_id='test') + reporter.record(MetricRecord(stage='train', values={'loss': 1.0})) + health = reporter.health() + assert health['record_count'] == 1 + assert health['backends']['jsonl']['enabled'] is False + assert health['backends']['jsonl']['failure_count'] == 1 + reporter.close() diff --git a/tests/twinkle_agentic/test_async_rl_native_tq.py b/tests/twinkle_agentic/test_async_rl_native_tq.py new file mode 100644 index 000000000..ef8ec0174 --- /dev/null +++ b/tests/twinkle_agentic/test_async_rl_native_tq.py @@ -0,0 +1,653 @@ +from __future__ import annotations + +import asyncio +import inspect + +import pytest + +from twinkle.data_format import SampledSequence, SampleResponse +from twinkle.metric import MetricRecord +from twinkle_agentic.async_rl import (AsyncMultiLoraGRPOPipeline, ContextSchedulePolicy, ContextScheduler, + ContextStatus, LoraContext, LoraContextManager, ScheduleCandidate, + SchedulerConfig, TQDataPlane, TrainerWorker) +from twinkle_agentic.async_rl.metrics import training_policy_metrics +from twinkle_agentic.async_rl.native_tq import ContextGRPOGroupNSampler +from twinkle_agentic.async_rl.pipeline import ( + _collect_adapter_path, + _require_adapter_path, + _train_batch, + create_cpu_actor, +) +from twinkle_agentic.async_rl.types import PartitionAdmission, PreparedPartition +from twinkle_agentic.async_rl.utils import ( + TrainBatchConfig, + sample_responses_to_rollout_rows, +) +from twinkle_agentic.async_rl.workers import RolloutWorker + + +class LocalActorHandle: + def __init__(self, target): + self.target = target + + def __getattr__(self, name): + method = getattr(self.target, name) + + class RemoteMethod: + async def remote(_, *args, **kwargs): + result = method(*args, **kwargs) + return await result if inspect.isawaitable(result) else result + + return RemoteMethod() + + +def test_cpu_service_actor_uses_twinkle_ray_mode(monkeypatch): + import ray + + captured = {} + + class ActorClass: + + @staticmethod + def remote(*args, **kwargs): + captured['actor_args'] = args + captured['actor_kwargs'] = kwargs + return 'actor' + + def fake_remote(**options): + captured['options'] = options + return lambda cls: ActorClass + + monkeypatch.setattr(ray, 'remote', fake_remote) + + assert create_cpu_actor(object, 'value', enabled=True) == 'actor' + assert captured['options'] == { + 'num_cpus': 1, + 'runtime_env': { + 'env_vars': { + 'TWINKLE_MODE': 'ray' + } + }, + } + assert captured['actor_args'] == ('value', ) + assert captured['actor_kwargs'] == {'enabled': True} + + +def test_train_batch_preserves_position_ids_from_tq(): + class Batch(dict): + batch_size = (1, ) + + class Model: + inputs = None + + def forward_backward(self, *, inputs, **_kwargs): + self.inputs = inputs + + def clip_grad_and_step(self, **_kwargs): + return None + + def calculate_metric(self, **_kwargs): + return {} + + context = _context() + admission = PartitionAdmission(context, context.partition_id(0), 0, 1, 1, 0) + data = Batch({ + 'input_ids': [[1, 2]], + 'labels': [[-100, 2]], + 'attention_mask': [[1, 1]], + 'position_ids': [[0, 1]], + 'logprobs': [[-.1]], + 'advantages': [1.], + 'rewards': [1.], + }) + model = Model() + + _train_batch(model, {context.key: TrainBatchConfig(1, 1)}, data, admission) + + assert model.inputs == [{ + 'input_ids': [1, 2], + 'labels': [-100, 2], + 'attention_mask': [1, 1], + 'position_ids': [0, 1], + }] + + +def test_train_batch_accumulates_real_micro_batches_before_one_optimizer_step(): + class Batch(dict): + batch_size = (4, ) + + class Model: + def __init__(self): + self.calls = [] + self.optimizer_steps = 0 + + def forward_backward(self, **kwargs): + self.calls.append(kwargs) + return lambda: {} + + def clip_grad_and_step(self, **_kwargs): + self.optimizer_steps += 1 + + def calculate_metric(self, **_kwargs): + return {} + + context = _context() + admission = PartitionAdmission(context, context.partition_id(0), 0, 1, 4, 0) + data = Batch({ + 'input_ids': [[index] for index in range(4)], + 'labels': [[index] for index in range(4)], + 'attention_mask': [[1] for _ in range(4)], + 'position_ids': [[0] for _ in range(4)], + 'logprobs': [[-.1] for _ in range(4)], + 'advantages': [1., 2., 3., 4.], + 'rewards': [1., 1., 1., 1.], + }) + model = Model() + + metrics = _train_batch( + model, + {context.key: TrainBatchConfig(4, 1)}, + data, + admission, + model_data_parallel_size=1, + ) + + assert [len(call['inputs']) for call in model.calls] == [4] + assert model.calls[0]['advantages'] == [1., 2., 3., 4.] + assert model.calls[0]['micro_batch_size'] == 1 + assert model.calls[0]['loss_scale'] == 1.0 + assert model.optimizer_steps == 1 + assert metrics['micro_batch_size_per_rank'] == 1 + + +def _context(name: str = 'adapter') -> LoraContext: + return LoraContext('tenant', f'run_{name}', 'model', name) + + +def _sample_response(tokens, stop_reason, input_ids): + return SampleResponse( + prompt_token_ids=[1, 2], + sequences=[ + SampledSequence( + stop_reason=stop_reason, + tokens=tokens, + logprobs=[[(token, -.1)] for token in tokens], + new_input_feature={ + 'input_ids': input_ids, + 'labels': [-100, -100, *tokens], + }, + ) + ], + ) + + +def test_evaluation_rows_do_not_require_rollout_group_metadata(): + prompt = {'input_ids': [1, 2], 'labels': [-100, -100]} + + rows = sample_responses_to_rollout_rows( + [prompt], + [_sample_response([3], 'stop', [1, 2, 3])], + policy_version=10, + ) + + assert len(rows) == 1 + assert 'group_id' not in rows[0] + assert 'generation_idx' not in rows[0] + assert rows[0]['rollout_policy_version'] == 10 + + +def test_training_rows_preserve_rollout_group_metadata(): + source = { + 'input_ids': [1, 2], + 'labels': [-100, -100], + 'group_id': 'partition/group_0', + 'generation_idx': 2, + } + + rows = sample_responses_to_rollout_rows( + [source], + [_sample_response([3], 'stop', [1, 2, 3])], + policy_version=4, + ) + + assert rows[0]['group_id'] == 'partition/group_0' + assert rows[0]['generation_idx'] == 2 + + +def test_adapter_path_rejects_uncollected_remote_result(): + def lazy_result(): + return '/tmp/policy' + + with pytest.raises(TypeError, match='must return a non-empty checkpoint path string'): + _require_adapter_path(lazy_result, operation='test save') + + +def test_adapter_path_collects_lazy_remote_result(): + def lazy_result(): + return '/tmp/policy' + + lazy_result._is_lazy_collect = True + assert _collect_adapter_path(lazy_result, operation='test save') == '/tmp/policy' + + +def test_training_policy_metrics_use_final_version_and_partial_span(): + metrics = training_policy_metrics(( + { + 'final_policy_version': 3, + 'policy_version_span': 1 + }, + { + 'final_policy_version': 4, + 'policy_version_span': 0 + }, + ), train_policy_version=5) + + assert metrics == { + 'policy_version_gap_mean': 1.5, + 'policy_version_gap_p95': 2, + 'policy_version_gap_max': 2, + 'rollout_policy_span_mean': 0.5, + 'rollout_policy_span_max': 1, + } + + +def test_training_policy_metrics_reject_future_rollout_version(): + try: + training_policy_metrics(({ + 'final_policy_version': 6, + 'policy_version_span': 0 + }, ), train_policy_version=5) + except ValueError as exc: + assert 'older than rollout versions' in str(exc) + else: + raise AssertionError('expected a future rollout policy version to fail') + + +def test_unified_staleness_admission(): + context = _context() + zero = LoraContextManager(max_staleness=0) + zero.register_context(context) + first = zero.request_rollout_partition(context, target_groups=1, num_generations=2) + assert first is not None + assert zero.request_rollout_partition(context, target_groups=1, num_generations=2) is None + + one = LoraContextManager(max_staleness=1) + one.register_context(context) + assert one.request_rollout_partition(context, target_groups=1, num_generations=2) is not None + assert one.request_rollout_partition(context, target_groups=1, num_generations=2) is not None + assert one.request_rollout_partition(context, target_groups=1, num_generations=2) is None + + +def test_rollout_worker_retains_prefetched_batch_until_admission_succeeds(): + context = _context() + + class AdmissionGate: + def __init__(self): + self.blocked = True + self.attempts = 0 + self.accepted = False + + def is_rollout_admission_closed(self): + return False + + def context_status(self, _context): + return ContextStatus.ACTIVE + + def request_rollout_partition(self, _context, *, target_groups, num_generations): + self.attempts += 1 + if self.blocked or self.accepted: + return None + self.accepted = True + return PartitionAdmission(context, context.partition_id(0), 0, target_groups, num_generations, 0) + + class DataPlane: + async def prepare_rollout_partition(self, admission, _prompts, sampling_params): + return PreparedPartition(admission, (), sampling_params) + + class Sampler: + def __init__(self, loop): + self.submitted = asyncio.Event() + self.loop = loop + + def submit_prompt_groups(self, _groups, _sampling_params, _allow_partial_rollout): + self.loop.call_soon_threadsafe(self.submitted.set) + + loaded_batches = [] + + def batches(): + for value in (1, 2): + loaded_batches.append(value) + yield [{'input_ids': [value]}] + + async def run(): + manager = AdmissionGate() + sampler = Sampler(asyncio.get_running_loop()) + worker = RolloutWorker( + context_manager=LocalActorHandle(manager), + data_plane=DataPlane(), + sampler=sampler, + prompt_batches={context.key: batches()}, + rollout_config={ + context.key: { + 'context': context, + 'batch_size': 1, + 'num_generations': 2, + 'sampling_params': {}, + } + }, + scheduler=SchedulerConfig(ContextSchedulePolicy.ROUND_ROBIN, 1), + idle_delay_s=.001, + ) + await worker.start() + while manager.attempts == 0: + await asyncio.sleep(.001) + prefetched_task = worker._next_batch_tasks[context.key] + await asyncio.sleep(.01) + assert worker._next_batch_tasks[context.key] is prefetched_task + assert loaded_batches == [1] + + manager.blocked = False + await asyncio.wait_for(sampler.submitted.wait(), timeout=1) + while loaded_batches == [1]: + await asyncio.sleep(.001) + assert loaded_batches == [1, 2] + await worker.stop() + + asyncio.run(run()) + + +def test_partition_clear_releases_capacity_only_after_publish(): + context = _context() + manager = LoraContextManager(max_staleness=0) + manager.register_context(context, adapter_path='initial') + admission = manager.request_rollout_partition(context, target_groups=1, num_generations=2) + manager.on_partition_training_started(admission) + policy = manager.on_partition_trained(admission, adapter_path='v1') + assert policy.version == 1 + assert manager.request_rollout_partition(context, target_groups=1, num_generations=2) is None + manager.on_partition_cleared(admission) + assert manager.request_rollout_partition(context, target_groups=1, num_generations=2) is not None + + +def test_context_trains_partitions_in_step_order(): + context = _context() + manager = LoraContextManager(max_staleness=1) + manager.register_context(context) + first = manager.request_rollout_partition(context, target_groups=1, num_generations=2) + second = manager.request_rollout_partition(context, target_groups=1, num_generations=2) + + assert manager.list_trainable_partitions() == [first] + manager.on_partition_training_started(first) + assert manager.list_trainable_partitions() == [first] + + try: + manager.on_partition_training_started(second) + except RuntimeError as exc: + assert f'already trains {first.partition_id}' in str(exc) + else: + raise AssertionError('expected the next partition to remain blocked') + + manager.on_partition_trained(first, adapter_path='v1') + manager.on_partition_cleared(first) + assert manager.list_trainable_partitions() == [second] + manager.on_partition_training_started(second) + + +def test_scheduler_supports_round_robin_sticky_and_oldest(): + a, b = _context('a'), _context('b') + candidates = [ScheduleCandidate(a), ScheduleCandidate(b)] + round_robin = ContextScheduler(SchedulerConfig(ContextSchedulePolicy.ROUND_ROBIN, 1)) + assert round_robin.choose(candidates).context == a + round_robin.on_success(candidates[0]) + assert round_robin.choose(candidates).context == b + + sticky = ContextScheduler(SchedulerConfig(ContextSchedulePolicy.STICKY, None)) + sticky.on_success(candidates[1]) + assert sticky.choose(candidates).context == b + sticky.on_blocked(candidates[1]) + assert sticky.choose(candidates).context == a + + capped = ContextScheduler(SchedulerConfig(ContextSchedulePolicy.STICKY, 1)) + first = capped.choose(candidates) + capped.on_success(first) + assert capped.choose(candidates).context == b + + manager = LoraContextManager(max_staleness=2) + manager.register_context(a) + manager.register_context(b) + old = manager.request_rollout_partition(a, target_groups=1, num_generations=2) + new = manager.request_rollout_partition(b, target_groups=1, num_generations=2) + oldest = ContextScheduler(SchedulerConfig(ContextSchedulePolicy.OLDEST_PARTITION, 1)) + assert oldest.choose([ScheduleCandidate(b, new), ScheduleCandidate(a, old)]).partition == old + + +def test_context_group_sampler_uses_request_generation_count(): + sampler = ContextGRPOGroupNSampler() + + selected, consumed = sampler.sample( + [0, 1, 4, 5, 6, 7], + batch_size=4, + partition_id='train_0', + task_name='advantage/context', + n_samples_per_prompt=4, + ) + + assert selected == [4, 5, 6, 7] + assert consumed == selected + + selected, consumed = sampler.sample( + [8, 9, 12], + batch_size=2, + partition_id='train_1', + task_name='advantage/context', + n_samples_per_prompt=2, + ) + + assert selected == [8, 9] + assert consumed == selected + + +def test_context_finishes_after_exhaustion_and_clear(): + context = _context() + manager = LoraContextManager() + manager.register_context(context) + admission = manager.request_rollout_partition(context, target_groups=1, num_generations=2) + manager.on_dataset_exhausted(context) + assert not manager.is_run_finished() + manager.on_partition_training_started(admission) + manager.on_partition_trained(admission, adapter_path='v1') + manager.on_partition_cleared(admission) + assert manager.is_run_finished() + + +def test_pipeline_fails_fast_when_a_worker_service_fails(): + context = _context() + manager = LoraContextManager() + manager.register_context(context) + + class FailedWorker: + async def start(self): + return None + + async def stop(self): + return None + + async def get_service_state(self): + return {'running': False, 'failure': 'CUDA out of memory'} + + def drain_metric_records(self): + return [] + + worker = LocalActorHandle(FailedWorker()) + pipeline = AsyncMultiLoraGRPOPipeline( + context_manager=LocalActorHandle(manager), + rollout_worker=worker, + advantage_worker=worker, + trainer_worker=worker, + ) + + try: + asyncio.run(pipeline.run_async()) + except RuntimeError as exc: + assert 'CUDA out of memory' in str(exc) + else: + raise AssertionError('expected worker failure to fail the pipeline') + + +def test_pipeline_drains_actor_metric_buffers_when_reporting_is_disabled(): + class BufferedWorker: + def __init__(self): + self.drain_count = 0 + + def drain_metric_records(self): + self.drain_count += 1 + return [MetricRecord(stage='train', values={'loss': 1.0})] + + class BufferedSampler: + def __init__(self): + self.drain_count = 0 + + def drain_metric_records(self): + self.drain_count += 1 + return [MetricRecord(stage='rollout', values={'sample_count': 1})] + + workers = [BufferedWorker() for _ in range(3)] + sampler = BufferedSampler() + pipeline = AsyncMultiLoraGRPOPipeline( + context_manager=object(), + rollout_worker=LocalActorHandle(workers[0]), + advantage_worker=LocalActorHandle(workers[1]), + trainer_worker=LocalActorHandle(workers[2]), + sampler=sampler, + metrics=None, + ) + + asyncio.run(pipeline._drain_metrics()) + + assert [worker.drain_count for worker in workers] == [1, 1, 1] + assert sampler.drain_count == 1 + + +def test_global_max_steps_limits_admission_and_closes_after_completion(): + first, second = _context('a'), _context('b') + manager = LoraContextManager(max_staleness=1, max_steps=1) + manager.register_context(first) + manager.register_context(second) + first_admission = manager.request_rollout_partition(first, target_groups=1, num_generations=2) + assert manager.request_rollout_partition(second, target_groups=1, num_generations=2) is None + manager.on_partition_training_started(first_admission) + manager.on_partition_trained(first_admission, adapter_path='v1') + manager.on_partition_cleared(first_admission) + assert manager.is_rollout_admission_closed() + assert manager.is_run_finished() + + +def test_zero_max_steps_finishes_without_admission(): + context = _context() + manager = LoraContextManager(max_steps=0) + manager.register_context(context) + assert manager.request_rollout_partition(context, target_groups=1, num_generations=2) is None + assert manager.is_rollout_admission_closed() + assert manager.is_run_finished() + + +def test_checkpoint_retention_preserves_current_policy_and_history_window(): + context = _context() + manager = LoraContextManager() + manager.register_context(context, adapter_path='initial') + removed = [] + worker = TrainerWorker( + context_manager=LocalActorHandle(manager), + data_plane=TQDataPlane(), + train_fn=lambda _data, _admission: {}, + save_adapter=lambda _admission: 'unused', + mini_batch_sizes={context.key: 2}, + scheduler=SchedulerConfig(ContextSchedulePolicy.STICKY, None), + keep_adapter_versions=1, + initial_adapter_paths={context.key: 'initial'}, + remove_adapter=removed.append, + ) + admission = manager.request_rollout_partition(context, target_groups=1, num_generations=2) + manager.on_partition_training_started(admission) + manager.on_partition_trained(admission, adapter_path='current') + manager.on_partition_cleared(admission) + worker._adapter_history[context.key].append('current') + async def prune(): + await worker._prune_adapter_history(context) + await worker.stop() + + asyncio.run(prune()) + assert removed == ['initial'] + assert worker._adapter_history[context.key] == ['current'] + prune_events = [ + record for record in worker.drain_metric_records() + if record.stage == 'policy' and record.attributes.get('operation') == 'adapter_prune' + ] + assert len(prune_events) == 1 + assert prune_events[0].context_key == context.key + assert prune_events[0].attributes['adapter_path'] == 'initial' + assert prune_events[0].values['adapter_prune_latency_s'] >= 0 + + +def test_policy_retention_keeps_only_current_and_actively_referenced_paths(): + context = _context() + manager = LoraContextManager(max_staleness=1) + manager.register_context(context, adapter_path='initial') + admission = manager.request_rollout_partition(context, target_groups=1, num_generations=2) + + acquired = manager.acquire_rollout_policy(context) + manager.on_partition_training_started(admission) + manager.on_partition_trained(admission, adapter_path='current') + + assert manager.adapter_paths_to_keep() == {'initial', 'current'} + manager.release_rollout_policy(acquired) + assert manager.adapter_paths_to_keep() == {'current'} + + +def test_trainer_periodically_evaluates_published_policy(): + context = _context() + manager = LoraContextManager() + manager.register_context(context, adapter_path='initial') + admission = manager.request_rollout_partition(context, target_groups=1, num_generations=2) + calls = [] + + def evaluate_batch(batch, evaluated_admission, adapter_path, policy_version, sampling_params): + calls.append((list(batch), evaluated_admission, adapter_path, policy_version, sampling_params)) + return { + 'rewards': [1.0] * len(batch), + 'completion_lengths': [10] * len(batch), + } + + worker = TrainerWorker( + context_manager=LocalActorHandle(manager), + data_plane=TQDataPlane(), + train_fn=lambda _data, _admission: {}, + save_adapter=lambda _admission: 'unused', + mini_batch_sizes={context.key: 2}, + scheduler=SchedulerConfig(ContextSchedulePolicy.STICKY, None), + evaluation_config={ + context.key: { + 'interval': 5, + 'dataset_name': 'validation', + 'prompt_batches': lambda: [[{'input_ids': [1]}], [{'input_ids': [2]}]], + 'sampling_params': 'params', + } + }, + evaluate_batch=evaluate_batch, + ) + worker._optimizer_steps[context.key] = 50 + + async def evaluate(): + await worker._evaluate_policy(admission, 'adapter-v4', 4) + await worker._evaluate_policy(admission, 'adapter-v5', 5) + + asyncio.run(evaluate()) + assert len(calls) == 2 + records = [record for record in worker.drain_metric_records() if record.stage == 'evaluation'] + assert len(records) == 1 + assert records[0].policy_version == 5 + assert records[0].optimizer_step == 50 + assert records[0].values['accuracy'] == 1.0 + assert records[0].values['prompt_count'] == 2 + assert records[0].values['sample_count'] == 2 + assert records[0].values['completion_length'] == 10 diff --git a/tests/twinkle_agentic/test_vllm_sampler_tq_generation.py b/tests/twinkle_agentic/test_vllm_sampler_tq_generation.py new file mode 100644 index 000000000..e64669b59 --- /dev/null +++ b/tests/twinkle_agentic/test_vllm_sampler_tq_generation.py @@ -0,0 +1,495 @@ +from __future__ import annotations + +import asyncio +import inspect +import json +import time +from concurrent.futures import Future + +import pytest + +from twinkle import DeviceMesh +from twinkle.data_format import SampledSequence, SampleResponse, SamplingParams +from twinkle.infra import _dispatch_args +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, + _GeneratedSample, + _PromptGroupRolloutStats, + _dispatch_generation, +) + + +class LocalActorHandle: + def __init__(self, target): + self.target = target + + def __getattr__(self, name): + method = getattr(self.target, name) + + class RemoteMethod: + async def remote(_, *args, **kwargs): + result = method(*args, **kwargs) + return await result if inspect.isawaitable(result) else result + + return RemoteMethod() + + +class PolicyProvider: + def __init__(self, policies): + self.policies = iter(policies) + self.released = [] + + def get_rollout_policy(self, _context): + return next(self.policies) + + def acquire_rollout_policy(self, context): + return self.get_rollout_policy(context) + + def release_rollout_policy(self, policy): + self.released.append(policy) + + +class GenerationHarness: + _merge_partial_responses = VLLMSamplerTQ._merge_partial_responses + + def __init__(self, policies, responses): + self.context_manager = LocalActorHandle(PolicyProvider(policies)) + self.responses = iter(responses) + self.rollout_max_retries = 1 + self.rollout_retry_delay_s = 0 + self.calls = [] + self.template = type('Template', (), {'decode': staticmethod(lambda tokens: str(tokens))})() + + async def _load_lora_for_policy(self, policy): + return policy.version + + async def _sample_single(self, feat, sampling_params, *, lora_request, multi_modal_data, logprobs_only): + self.calls.append((list(feat['input_ids']), sampling_params.max_tokens, lora_request)) + return next(self.responses) + + +def _context(name: str = 'adapter') -> LoraContext: + return LoraContext('tenant', f'run_{name}', 'model', name) + + +def _sample_response(tokens, stop_reason, input_ids): + return SampleResponse( + prompt_token_ids=[1, 2], + sequences=[ + SampledSequence( + stop_reason=stop_reason, + tokens=tokens, + logprobs=[[(token, -.1)] for token in tokens], + new_input_feature={ + 'input_ids': input_ids, + 'labels': [-100, -100, *tokens], + }, + ) + ], + ) + + +def _bare_sampler() -> VLLMSamplerTQ: + sampler = object.__new__(VLLMSamplerTQ) + sampler._generation_submissions = {} + return sampler + + +def test_generation_dispatch_allows_one_prompt_with_multiple_dp_workers() -> None: + assert VLLMSamplerTQ.submit_generation._dispatch is _dispatch_generation + assert VLLMSamplerTQ.submit_prompt_groups._dispatch == 'slice_dp' + shards = [ + _dispatch_generation( + 3, + worker_index, + ('submission', [{'input_ids': [1]}], 'params'), + {}, + )[0][1] + for worker_index in range(3) + ] + + assert shards == [[{'input_ids': [1]}], [], []] + + +def test_generation_submission_returns_before_generation_finishes() -> None: + sampler = _bare_sampler() + pending = Future() + submitted_coroutines = [] + + def submit(coro): + submitted_coroutines.append(coro) + coro.close() + return pending + + sampler._submit_in_loop = submit + + result = sampler.submit_generation( + 'submission-1', + [{'input_ids': [1]}], + SamplingParams(max_tokens=4), + ) + + assert result == {'submission_id': 'submission-1', 'status': 'running'} + assert not pending.done() + assert len(submitted_coroutines) == 1 + assert sampler.get_generation_status('submission-1')['status'] == 'running' + + responses = [object()] + pending.set_result(responses) + assert sampler.get_generation_status('submission-1')['status'] == 'completed' + assert sampler.collect_generation('submission-1') == responses + assert 'submission-1' not in sampler._generation_submissions + + +def test_generation_keeps_one_response_per_prompt() -> None: + sampler = _bare_sampler() + sampler.template = None + + async def sample_single(feat, _params, **_kwargs): + await asyncio.sleep(0) + return feat['input_ids'][0] + + sampler._sample_single = sample_single + responses = asyncio.run( + sampler._generate_inputs( + [{'input_ids': [10]}, {'input_ids': [20]}], + SamplingParams(max_tokens=4), + adapter_name='', + adapter_path=None, + use_base_model=False, + )) + + assert responses == [10, 20] + + +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 + + state = sampler.get_generation_status('failed') + assert state['status'] == 'failed' + assert state['error'] == 'ValueError: bad prompt' + + with pytest.raises(ValueError, match='bad prompt'): + sampler.collect_generation('failed') + 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_all_generations_are_cancelled_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() + assert sampler._generation_submissions == {} + + +def test_native_prompt_group_sampling_requires_context_manager() -> None: + sampler = _bare_sampler() + sampler.context_manager = None + + with pytest.raises(RuntimeError, match='context_manager is required'): + sampler.submit_prompt_groups([], SamplingParams(max_tokens=4)) + + +def test_server_waiter_admits_later_submission_before_first_finishes() -> None: + + class Sampler: + + def __init__(self): + self.futures: dict[str, Future] = {} + self.submission_order = [] + + def submit_generation(self, submission_id, *_args, **_kwargs): + self.submission_order.append(submission_id) + self.futures[submission_id] = Future() + + def get_generation_status(self, submission_id): + future = self.futures[submission_id] + return {'status': 'completed' if future.done() else 'running'} + + def collect_generation(self, submission_id): + return self.futures[submission_id].result() + + def cancel_generation(self, submission_id): + self.futures.pop(submission_id, None) + + sampler = Sampler() + + async def run(): + sampler.submit_generation('first') + sampler.submit_generation('second') + first = asyncio.create_task( + _await_generation(sampler, 'first')) + second = asyncio.create_task( + _await_generation(sampler, 'second')) + while len(sampler.submission_order) < 2: + await asyncio.sleep(0) + assert not first.done() + sampler.futures['second'].set_result(['short']) + assert await second == ['short'] + assert not first.done() + sampler.futures['first'].set_result(['long']) + assert await first == ['long'] + + asyncio.run(run()) + assert set(sampler.submission_order) == {'first', 'second'} + + +def test_server_waiter_retries_cancelled_status_poll() -> None: + from ray.exceptions import TaskCancelledError + + class Sampler: + + def __init__(self): + self.status_calls = 0 + self.cancelled = False + + def submit_generation(self, *_args, **_kwargs): + return None + + def get_generation_status(self, _submission_id): + self.status_calls += 1 + if self.status_calls == 1: + raise TaskCancelledError() + return {'status': 'completed'} + + def collect_generation(self, _submission_id): + return ['completed'] + + def cancel_generation(self, _submission_id): + self.cancelled = True + + sampler = Sampler() + sampler.submit_generation('submission') + result = asyncio.run( + _await_generation(sampler, 'submission')) + + assert result == ['completed'] + assert sampler.status_calls == 2 + assert not sampler.cancelled + + +def test_sampler_dp_dispatch_slices_complete_groups_without_duplication(): + mesh = DeviceMesh.from_sizes(world_size=4, dp_size=2, tp_size=2) + groups = ['group_0', 'group_1', 'group_2', 'group_3'] + dispatched = _dispatch_args( + workers=['dp_0', 'dp_1'], + dispatch='slice_dp', + execute='all', + device_mesh=mesh, + args=(groups, 'sampling_params', False), + kwargs={}, + ) + + assert [worker for worker, _, _ in dispatched] == ['dp_0', 'dp_1'] + assert [args[0] for _, args, _ in dispatched] == [groups[:2], groups[2:]] + assert [group for _, args, _ in dispatched for group in args[0]] == groups + + +@pytest.mark.parametrize(('dp_size', 'expected_scope'), [(1, 'partition'), (2, 'shard')]) +def test_sampler_reports_submission_throughput_at_partition_or_shard_scope(dp_size, expected_scope): + context = _context() + admission = PartitionAdmission(context, context.partition_id(0), 0, 2, 2, 0) + groups = [ + PromptGroup(context, admission, f'{admission.partition_id}/group_{index}', {}, object()) + for index in range(2) + ] + + class RolloutMetricsHarness: + def __init__(self): + self.device_mesh = DeviceMesh.from_sizes(world_size=dp_size, dp_size=dp_size) + self.events = [] + + async def _run_prompt_group(self, *, group, **_kwargs): + index = int(group.group_id.rsplit('_', 1)[1]) + lengths = ((10, 20), (30, 40))[index] + reasons = (('stop', 'length'), ('stop', 'stop'))[index] + return _PromptGroupRolloutStats(lengths, reasons, (index + 1, index + 1)) + + def _record_metrics(self, group, values, **kwargs): + self.events.append((group, values, kwargs)) + + sampler = RolloutMetricsHarness() + asyncio.run( + VLLMSamplerTQ._sample_prompt_groups( + sampler, + 'submission', + groups, + SamplingParams(max_tokens=64), + False, + time.perf_counter() - 1, + )) + + recorded_group, metrics, record_options = sampler.events[-1] + assert recorded_group.context == context + assert recorded_group.partition_id == admission.partition_id + assert record_options['attributes']['scope'] == expected_scope + assert metrics['prompt_group_count'] == 2 + assert metrics['sample_count'] == 4 + assert metrics['output_tokens'] == 100 + assert metrics['completion_length_mean'] == 25 + assert metrics['completion_truncated_count'] == 1 + assert metrics['policy_version_min'] == 1 + assert metrics['policy_version_max'] == 2 + assert metrics['sampler_dp_size'] == dp_size + assert metrics['output_tokens_per_s'] == pytest.approx(100 / metrics['rollout_latency_s']) + + +def test_sampler_writes_one_atomic_rollout_file_per_prompt_group(tmp_path): + context = _context() + admission = PartitionAdmission(context, context.partition_id(3), 3, 1, 2, 0) + group = PromptGroup( + context, + admission, + f'{admission.partition_id}/group_0', + {'user_data': [('ground_truth', '"42"')]}, + object(), + ) + policy = RolloutPolicy(context.key, context.adapter_name, 7, '/tmp/adapter-v7') + generated = [ + _GeneratedSample( + SampleResponse( + sequences=[SampledSequence('stop', [20 + index], decoded=f'completion-{index}')], + prompt_token_ids=[10, 11], + ), + (policy,), + attempts=1, + was_aborted=False, + resumed_partial_output=False, + ) + for index in range(2) + ] + rows = [ + { + 'generation_idx': index, + 'rollout_policy_version': 7, + 'initial_policy_version': 7, + 'final_policy_version': 7, + 'rollout_policy_versions': [7], + 'rollout_adapter_path': '/tmp/adapter-v7', + 'stop_reason': 'stop', + 'logprobs': [-0.1], + } + for index in range(2) + ] + + class Template: + @staticmethod + def decode(token_ids, **_kwargs): + return ' '.join(map(str, token_ids)) + + sampler = object.__new__(VLLMSamplerTQ) + sampler.rollout_output_dir = tmp_path + sampler.rollout_output_include_token_ids = False + sampler.template = Template() + + sampler._write_rollout_group('submission-1', group, generated, rows, [1.0, 0.0]) + sampler._write_rollout_group('submission-2', group, generated, rows, [1.0, 0.0]) + + output_path = ( + tmp_path + / context.tenant_id + / context.training_run_id + / context.adapter_name + / 'policy_7' + / 'train_3-group_0.jsonl' + ) + records = [json.loads(line) for line in output_path.read_text().splitlines()] + assert len(records) == 2 + assert records[0]['submission_id'] == 'submission-2' + assert records[0]['prompt'] == '10 11' + assert records[0]['completion'] == '20' + assert records[0]['ground_truth'] == '42' + assert records[0]['reward'] == 1.0 + assert records[0]['head_version'] == 7 + assert records[0]['tail_version'] == 7 + assert 'prompt_token_ids' not in records[0] + + +def test_aborted_generation_restarts_from_original_prompt_when_partial_is_disabled(): + context = _context() + policies = [ + RolloutPolicy(context.key, context.adapter_name, 3, 'adapter-v3'), + RolloutPolicy(context.key, context.adapter_name, 4, 'adapter-v4'), + ] + sampler = GenerationHarness( + policies, + [ + _sample_response([7], 'abort', [1, 2, 7]), + _sample_response([8], 'stop', [1, 2, 8]), + ], + ) + generated = asyncio.run( + VLLMSamplerTQ._generate_sample( + sampler, + context, + {'input_ids': [1, 2], 'labels': [-100, -100]}, + SamplingParams(max_tokens=4, logprobs=1), + multi_modal_data=None, + logprobs_only=False, + allow_partial_rollout=False, + )) + + assert sampler.calls == [([1, 2], 4, 3), ([1, 2], 4, 4)] + assert generated.response.sequences[0].tokens == [8] + assert [policy.version for policy in generated.policies] == [4] + assert generated.retry_count == 1 + assert generated.was_aborted + assert not generated.resumed_partial_output + + +def test_aborted_generation_continues_from_partial_tokens_when_enabled(): + context = _context() + policies = [ + RolloutPolicy(context.key, context.adapter_name, 3, 'adapter-v3'), + RolloutPolicy(context.key, context.adapter_name, 4, 'adapter-v4'), + ] + sampler = GenerationHarness( + policies, + [ + _sample_response([7], 'abort', [1, 2, 7]), + _sample_response([8], 'stop', [1, 2, 7, 8]), + ], + ) + generated = asyncio.run( + VLLMSamplerTQ._generate_sample( + sampler, + context, + {'input_ids': [1, 2], 'labels': [-100, -100]}, + SamplingParams(max_tokens=4, logprobs=1), + multi_modal_data=None, + logprobs_only=False, + allow_partial_rollout=True, + )) + + assert sampler.calls == [([1, 2], 4, 3), ([1, 2, 7], 3, 4)] + assert generated.response.sequences[0].tokens == [7, 8] + assert [policy.version for policy in generated.policies] == [3, 4] + assert generated.initial_policy.version == 3 + assert generated.final_policy.version == 4 + assert generated.retry_count == 1 + assert generated.was_aborted + assert generated.resumed_partial_output diff --git a/tests/twinkle_client/test_async_components.py b/tests/twinkle_client/test_async_components.py new file mode 100644 index 000000000..99b91fcc9 --- /dev/null +++ b/tests/twinkle_client/test_async_components.py @@ -0,0 +1,172 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from __future__ import annotations + +import asyncio + +from twinkle_client.types import DataRef + + +class _Response: + + def __init__(self, payload, status_code: int = 200): + self._payload = payload + self.status_code = status_code + + def raise_for_status(self) -> None: + if self.status_code >= 400: + raise RuntimeError(self.status_code) + + def json(self): + return self._payload + + +def test_model_forward_backward_sends_multiple_data_refs(monkeypatch) -> None: + import twinkle_client.http as http_module + from twinkle_client.model import multi_lora_transformers as module + + calls = [] + + def post(*, url, json_data=None, **_kwargs): + calls.append((url, json_data)) + if url.endswith('/create'): + return _Response({}) + return _Response({'result': {'loss': 1.0}}) + + monkeypatch.setattr(http_module, 'get_base_url', lambda: 'http://server/api/v1') + monkeypatch.setattr(module, 'http_post', post) + + model = module.MultiLoraTransformersModel('ms://base') + model.adapter_name = 'adapter' + refs = [ + DataRef(ref_id='data-1', size=2, fields=['train_input']), + DataRef(ref_id='data-2', size=2, fields=['train_input']), + ] + model.forward_backward_from_data_plane( + refs, + input_field='train_input', + kwarg_fields={'advantages': 'advantage'}, + ) + + url, body = calls[-1] + assert url.endswith('/model/base/twinkle/forward_backward_from_data_plane') + assert body['input_refs'] == [ref.model_dump() for ref in refs] + assert body['input_field'] == 'train_input' + assert body['kwarg_fields'] == {'advantages': 'advantage'} + assert body['adapter_name'] == 'adapter' + + +def test_model_inline_forward_methods_keep_the_original_endpoints(monkeypatch) -> None: + import twinkle_client.http as http_module + from twinkle_client.model import multi_lora_transformers as module + + calls = [] + + def post(*, url, json_data=None, **_kwargs): + calls.append((url, json_data)) + return _Response({} if url.endswith('/create') else {'result': {}}) + + monkeypatch.setattr(http_module, 'get_base_url', lambda: 'http://server/api/v1') + monkeypatch.setattr(module, 'http_post', post) + + model = module.MultiLoraTransformersModel('ms://base') + model.adapter_name = 'adapter' + inputs = [{'input_ids': [1, 2]}] + model.forward(inputs, return_logits=True) + model.forward_only(inputs, disable_lora=True) + model.forward_backward(inputs, micro_batch_size=1) + + assert [url.rsplit('/', 1)[-1] for url, _ in calls[-3:]] == [ + 'forward', + 'forward_only', + 'forward_backward', + ] + assert all(body['inputs'] == inputs for _, body in calls[-3:]) + assert all('input_refs' not in body for _, body in calls[-3:]) + + +def test_model_data_plane_forward_uses_a_separate_api(monkeypatch) -> None: + import twinkle_client.http as http_module + from twinkle_client.model import multi_lora_transformers as module + + calls = [] + + def post(url, json_data=None, **_kwargs): + calls.append((url, json_data)) + return _Response({} if url.endswith('/create') else {'result': {'value': 1}}) + + monkeypatch.setattr(http_module, 'get_base_url', lambda: 'http://server/api/v1') + monkeypatch.setattr(module, 'http_post', post) + + model = module.MultiLoraTransformersModel('ms://base') + ref = DataRef(ref_id='data-1', size=2, fields=['train_input']) + model.forward_from_data_plane(ref, input_field='train_input') + + url, body = calls[-1] + assert url.endswith('/model/base/twinkle/forward_from_data_plane') + assert body['input_refs'] == [ref.model_dump()] + assert body['input_field'] == 'train_input' + + +def test_model_data_plane_forward_only_can_append_selected_outputs(monkeypatch) -> None: + import twinkle_client.http as http_module + from twinkle_client.model import multi_lora_transformers as module + + ref = DataRef(ref_id='data-1', size=2, fields=['input_ids']) + updated_ref = ref.model_copy(update={'fields': ['input_ids', 'ref_logps']}) + calls = [] + + def post(*, url, json_data=None, **_kwargs): + calls.append((url, json_data)) + return _Response({} if url.endswith('/create') else {'result': updated_ref.model_dump()}) + + monkeypatch.setattr(http_module, 'get_base_url', lambda: 'http://server/api/v1') + monkeypatch.setattr(module, 'http_post', post) + + model = module.MultiLoraTransformersModel('ms://base') + result = model.forward_only_from_data_plane( + ref, + output_ref=ref, + output_fields={'logps': 'ref_logps'}, + disable_lora=True, + ) + + url, body = calls[-1] + assert url.endswith('/model/base/twinkle/forward_only_from_data_plane') + assert body['input_refs'] == [ref.model_dump()] + assert body['output_ref'] == ref.model_dump() + assert result == updated_ref + + +def test_sampler_async_data_plane_path_returns_reference_without_materializing(monkeypatch) -> None: + import twinkle_client.http as http_module + from twinkle_client.sampler import vllm_sampler as module + + output_ref = DataRef( + ref_id='rollout-1', + size=4, + fields=['train_input', 'sampled_logprobs', 'decoded'], + kind='rollout', + ) + calls = [] + + def post(*, url, json_data=None, **_kwargs): + calls.append((url, json_data)) + if url.endswith('/create'): + return _Response({}) + return _Response(output_ref.model_dump()) + + monkeypatch.setattr(http_module, 'get_base_url', lambda: 'http://server/api/v1') + monkeypatch.setattr(module, 'http_post', post) + sampler = module.vLLMSampler('ms://base') + + result = asyncio.run(sampler.asample_to_data_plane( + [{'input_ids': [1]}], + num_samples=4, + group_ids=['group-1'], + )) + + assert result == output_ref + url, body = calls[-1] + assert url.endswith('/sampler/base/twinkle/sample_to_data_plane') + assert body['num_samples'] == 4 + assert body['group_ids'] == ['group-1'] diff --git a/tests/twinkle_client/test_async_rl_workers.py b/tests/twinkle_client/test_async_rl_workers.py new file mode 100644 index 000000000..7a79cb8a5 --- /dev/null +++ b/tests/twinkle_client/test_async_rl_workers.py @@ -0,0 +1,69 @@ +from __future__ import annotations + +import asyncio + +import pytest + +from twinkle_client.async_rl import Worker, WorkerPipeline + + +class _FunctionWorker(Worker): + + def __init__(self, name, function): + super().__init__(name) + self.function = function + + async def run(self) -> None: + await self.function() + + +def test_worker_pipeline_runs_roles_concurrently() -> None: + producer_started = asyncio.Event() + consumer_started = asyncio.Event() + + async def producer(): + producer_started.set() + await consumer_started.wait() + + async def consumer(): + consumer_started.set() + await producer_started.wait() + + asyncio.run(WorkerPipeline(( + _FunctionWorker('producer', producer), + _FunctionWorker('consumer', consumer), + )).run()) + + +def test_worker_pipeline_cancels_peer_when_one_role_fails() -> None: + waiting = asyncio.Event() + cancelled = asyncio.Event() + + async def peer(): + waiting.set() + try: + await asyncio.Event().wait() + finally: + cancelled.set() + + async def failure(): + await waiting.wait() + raise RuntimeError('role failed') + + with pytest.raises(RuntimeError, match='role failed'): + asyncio.run(WorkerPipeline(( + _FunctionWorker('peer', peer), + _FunctionWorker('failure', failure), + )).run()) + assert cancelled.is_set() + + +def test_worker_pipeline_rejects_duplicate_role_names() -> None: + async def noop(): + return None + + with pytest.raises(ValueError, match='unique'): + WorkerPipeline(( + _FunctionWorker('same', noop), + _FunctionWorker('same', noop), + )) diff --git a/tests/twinkle_client/test_client_orchestrated_grpo.py b/tests/twinkle_client/test_client_orchestrated_grpo.py new file mode 100644 index 000000000..2302e3bfc --- /dev/null +++ b/tests/twinkle_client/test_client_orchestrated_grpo.py @@ -0,0 +1,199 @@ +from __future__ import annotations + +import asyncio +import importlib.util +import sys +from pathlib import Path + +from twinkle_client.types import DataRef + + +MODULE_PATH = ( + Path(__file__).parents[2] / 'cookbook' / 'client' / 'async_rl' / 'client_orchestrated_grpo.py' +) + + +def _load_module(): + spec = importlib.util.spec_from_file_location('client_orchestrated_grpo', MODULE_PATH) + module = importlib.util.module_from_spec(spec) + assert spec.loader is not None + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module + + +def test_rollout_and_train_overlap_with_fifo_policy_publication(monkeypatch, capsys) -> None: + module = _load_module() + monkeypatch.setattr(module, 'BATCH_SIZE', 2) + monkeypatch.setattr(module, 'NUM_GENERATIONS', 2) + monkeypatch.setattr(module, 'TRAIN_MINI_BATCH_SIZE', 2) + monkeypatch.setattr(module, 'MAX_STALENESS', 1) + monkeypatch.setattr(module, 'MAX_PARTITIONS', 3) + + first_train_started = asyncio.Event() + rollout_snapshots = [] + events = [] + + async def fake_rollout(_sampler, prompt, policy, _semaphore, _group_id): + name = prompt['name'] + rollout_snapshots.append((name, policy.version, policy.adapter_uri)) + events.append(f'rollout-start:{name}') + if name == 'p0-g1': + await first_train_started.wait() + events.append(f'rollout-done:{name}') + return DataRef( + ref_id=name, + size=module.NUM_GENERATIONS, + fields=['train_input', 'sampled_logprobs', 'decoded'], + kind='rollout', + ) + + monkeypatch.setattr(module, 'rollout_group', fake_rollout) + monkeypatch.setattr(module, 'GSM8KAccuracyReward', lambda: lambda rows: [1.0] * len(rows)) + monkeypatch.setattr( + module, + 'GRPOAdvantage', + lambda: lambda rewards, **_kwargs: [1.0, -1.0], + ) + + class FakeModel: + def __init__(self): + self.saved = [] + self.steps = 0 + self.forward_backward_kwargs = [] + + async def save(self, name): + self.saved.append(name) + return {'twinkle_path': f'/checkpoints/{name}'} + + async def forward_backward_from_data_plane(self, _refs, **kwargs): + self.forward_backward_kwargs.append(kwargs) + events.append('train') + first_train_started.set() + + async def clip_grad_and_step(self, **_kwargs): + self.steps += 1 + + async def calculate_metric(self, **_kwargs): + return {'result': {'loss': 1.0 / self.steps, 'grad_norm': 0.5}} + + class FakeDataPlane: + def __init__(self): + self.released = [] + + async def aget(self, ref, *, fields=None): + assert fields == ['decoded'] + return [{'decoded': f'{ref.ref_id}-{index}'} for index in range(ref.size)] + + async def aappend(self, ref, rows, **_kwargs): + return ref.model_copy(update={'fields': [*ref.fields, *rows[0]]}) + + async def arelease(self, ref): + self.released.append(ref) + + async def run(): + model = FakeModel() + data_plane = FakeDataPlane() + batches = [ + [{'name': 'p0-g0'}, {'name': 'p0-g1'}], + [{'name': 'p1-g0'}, {'name': 'p1-g1'}], + [{'name': 'p2-g0'}, {'name': 'p2-g1'}], + ] + await module.run_grpo(batches, model, object(), data_plane) + return model, data_plane + + model, data_plane = asyncio.run(run()) + + assert events.index('train') < events.index('rollout-done:p0-g1') + assert model.saved == ['policy-0', 'policy-1', 'policy-2', 'policy-3'] + assert model.steps == 6 + assert all(kwargs['input_field'] == 'train_input' for kwargs in model.forward_backward_kwargs) + assert all(kwargs['kwarg_fields'] == { + 'old_logps': 'sampled_logprobs', + 'advantages': 'advantage', + } for kwargs in model.forward_backward_kwargs) + assert len(data_plane.released) == 6 + + snapshots = {name: (version, uri) for name, version, uri in rollout_snapshots} + assert snapshots['p0-g0'] == (0, '/checkpoints/policy-0') + assert snapshots['p1-g0'] == (0, '/checkpoints/policy-0') + assert snapshots['p2-g0'][0] in (1, 2) + assert snapshots['p2-g0'][1] == f'/checkpoints/policy-{snapshots["p2-g0"][0]}' + output = capsys.readouterr().out + assert 'optimizer_step=1' in output + assert 'loss=1.0' in output + assert 'grad_norm=0.5' in output + + +def test_younger_rollout_failure_stops_admission(monkeypatch) -> None: + module = _load_module() + monkeypatch.setattr(module, 'BATCH_SIZE', 1) + monkeypatch.setattr(module, 'NUM_GENERATIONS', 1) + monkeypatch.setattr(module, 'TRAIN_MINI_BATCH_SIZE', 1) + monkeypatch.setattr(module, 'MAX_STALENESS', 1) + monkeypatch.setattr(module, 'MAX_PARTITIONS', 3) + + started = [] + + async def fake_rollout(_sampler, prompt, _policy, _semaphore, _group_id): + name = prompt['name'] + started.append(name) + if name == 'p1': + raise RuntimeError('rollout failed') + return DataRef( + ref_id=name, + size=1, + fields=['train_input', 'sampled_logprobs', 'decoded'], + kind='rollout', + ) + + monkeypatch.setattr(module, 'rollout_group', fake_rollout) + monkeypatch.setattr(module, 'GSM8KAccuracyReward', lambda: lambda rows: [1.0]) + monkeypatch.setattr(module, 'GRPOAdvantage', lambda: lambda rewards, **_kwargs: [1.0]) + + class FakeModel: + def __init__(self): + self.saved = [] + + async def save(self, name): + self.saved.append(name) + return {'twinkle_path': name} + + async def forward_backward(self, _ref, **_kwargs): + return None + + async def clip_grad_and_step(self, **_kwargs): + return None + + async def calculate_metric(self, **_kwargs): + return {'result': {'loss': 1.0}} + + class FakeDataPlane: + async def aget(self, ref, *, fields=None): + assert fields == ['decoded'] + return [{'decoded': ref.ref_id}] + + async def aappend(self, ref, rows, **_kwargs): + return ref.model_copy(update={'fields': [*ref.fields, *rows[0]]}) + + async def arelease(self, _ref): + return None + + async def run(): + model = FakeModel() + try: + await module.run_grpo( + [[{'name': 'p0'}], [{'name': 'p1'}], [{'name': 'p2'}]], + model, + object(), + FakeDataPlane(), + ) + except RuntimeError as error: + assert str(error) == 'rollout failed' + else: + raise AssertionError('expected the younger rollout failure') + return model + + model = asyncio.run(run()) + assert started == ['p0', 'p1'] + assert model.saved == ['policy-0'] diff --git a/tests/twinkle_client/test_data_plane_async.py b/tests/twinkle_client/test_data_plane_async.py new file mode 100644 index 000000000..fd616f347 --- /dev/null +++ b/tests/twinkle_client/test_data_plane_async.py @@ -0,0 +1,102 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from __future__ import annotations + +import asyncio +import threading + +import pytest + +from twinkle_client.data_plane import DataPlaneClient +from twinkle_client.types import DataRef, DataRowsResponse + + +def test_async_convenience_methods_delegate_to_sync_operations(monkeypatch) -> None: + client = DataPlaneClient('http://server/data-plane') + original_ref = DataRef(ref_id='data-1', size=1, fields=['value']) + appended_ref = DataRef(ref_id='data-1', size=2, fields=['value']) + calls = [] + caller_thread = threading.get_ident() + + def put(rows, *, kind='data'): + calls.append(('put', rows, kind, threading.get_ident())) + return original_ref + + def get(ref, *, fields=None): + calls.append(('get', ref, fields, threading.get_ident())) + return [{'value': 1}] + + def append(ref, rows): + calls.append(('append', ref, rows, threading.get_ident())) + return appended_ref + + def release(ref): + calls.append(('release', ref, threading.get_ident())) + + monkeypatch.setattr(client, 'put', put) + monkeypatch.setattr(client, 'get', get) + monkeypatch.setattr(client, 'append', append) + monkeypatch.setattr(client, 'release', release) + + async def run(): + assert await client.aput([{'value': 1}], kind='rollout') == original_ref + assert await client.aget(original_ref, fields=['value']) == [{'value': 1}] + assert await client.aappend(original_ref, [{'value': 2}]) == appended_ref + assert await client.arelease(appended_ref) is None + + asyncio.run(run()) + + assert [call[:-1] for call in calls] == [ + ('put', [{'value': 1}], 'rollout'), + ('get', original_ref, ['value']), + ('append', original_ref, [{'value': 2}]), + ('release', appended_ref), + ] + assert all(call[-1] != caller_thread for call in calls) + + +def test_async_convenience_method_propagates_sync_error(monkeypatch) -> None: + client = DataPlaneClient('http://server/data-plane') + + def fail(_rows, *, kind='data'): + raise RuntimeError(f'put failed for {kind}') + + monkeypatch.setattr(client, 'put', fail) + + with pytest.raises(RuntimeError, match='put failed for rollout'): + asyncio.run(client.aput([], kind='rollout')) + + +def test_async_tagged_methods_and_batch_read_delegate_to_sync_operations(monkeypatch) -> None: + client = DataPlaneClient('http://server/data-plane') + ref = DataRef(ref_id='data-1', size=1, fields=['value']) + tags = [{'group_id': 'group-1'}] + calls = [] + + def put(rows, *, kind='data', tags=None): + calls.append(('put', rows, kind, tags)) + return ref + + def get_batch(value, *, fields=None): + calls.append(('get_batch', value, fields)) + return DataRowsResponse(rows=[{'value': 1}], tags=tags) + + def append(value, rows, *, tags=None): + calls.append(('append', value, rows, tags)) + return value + + monkeypatch.setattr(client, 'put', put) + monkeypatch.setattr(client, 'get_batch', get_batch) + monkeypatch.setattr(client, 'append', append) + + async def run(): + assert await client.aput([{'value': 1}], tags=tags) == ref + assert await client.aget_batch(ref) == DataRowsResponse(rows=[{'value': 1}], tags=tags) + assert await client.aappend(ref, [{'reward': 1.0}], tags=tags) == ref + + asyncio.run(run()) + + assert calls == [ + ('put', [{'value': 1}], 'data', tags), + ('get_batch', ref, None), + ('append', ref, [{'reward': 1.0}], tags), + ] diff --git a/tests/utils/test_rl_tensor_utils.py b/tests/utils/test_rl_tensor_utils.py new file mode 100644 index 000000000..d5d4486d1 --- /dev/null +++ b/tests/utils/test_rl_tensor_utils.py @@ -0,0 +1,58 @@ +import pytest +import torch + +from twinkle.utils.rl_tensor_utils import align_per_token_values + + +def test_align_per_token_values_accepts_json_rows() -> None: + actual = align_per_token_values( + [[-0.1, -0.2], [-0.3, -0.4]], + (2, 2), + device=torch.device('cpu'), + dtype=torch.float32, + name='ref_logps', + ) + + torch.testing.assert_close( + actual, + torch.tensor([[-0.1, -0.2], [-0.3, -0.4]]), + ) + + +def test_align_per_token_values_pads_ragged_json_rows() -> None: + actual = align_per_token_values( + [[-0.1, -0.2, -0.3], [-0.4, -0.5]], + (2, 3), + device=torch.device('cpu'), + dtype=torch.float32, + name='ref_logps', + valid_mask=torch.tensor([[True, True, True], [True, True, False]]), + ) + + torch.testing.assert_close( + actual, + torch.tensor([[-0.1, -0.2, -0.3], [-0.4, -0.5, 0.0]]), + ) + + +def test_align_per_token_values_rejects_missing_valid_tokens() -> None: + with pytest.raises(ValueError, match=r'ref_logps\[1\] has 2 tokens'): + align_per_token_values( + [[-0.1, -0.2, -0.3], [-0.4, -0.5]], + (2, 3), + device=torch.device('cpu'), + dtype=torch.float32, + name='ref_logps', + valid_mask=torch.ones(2, 3, dtype=torch.bool), + ) + + +def test_align_per_token_values_rejects_short_batches() -> None: + with pytest.raises(ValueError, match='batch size'): + align_per_token_values( + [[-0.1, -0.2]], + (2, 2), + device=torch.device('cpu'), + dtype=torch.float32, + name='ref_logps', + )