diff --git a/INSTALL_CLIENT.ps1 b/INSTALL_CLIENT.ps1 index 83846259..787cf32a 100644 --- a/INSTALL_CLIENT.ps1 +++ b/INSTALL_CLIENT.ps1 @@ -311,5 +311,5 @@ Write-Host " conda activate $EnvName" -ForegroundColor Yellow Write-Host "" Write-Host "Example usage (see cookbook/client/tinker/):" Write-Host ' $env:MODELSCOPE_TOKEN = "your-token"' -ForegroundColor Yellow -Write-Host " python cookbook/client/tinker/modelscope_service/self_cognition.py" -ForegroundColor Yellow +Write-Host " python cookbook/client/tinker/self_cognition.py" -ForegroundColor Yellow Write-Host "" diff --git a/INSTALL_CLIENT.sh b/INSTALL_CLIENT.sh index ecc7e90c..ea193dda 100644 --- a/INSTALL_CLIENT.sh +++ b/INSTALL_CLIENT.sh @@ -233,5 +233,5 @@ echo " conda activate $ENV_NAME" echo "" echo "Example usage (see cookbook/client/tinker/):" echo " export MODELSCOPE_TOKEN='your-token'" -echo " python cookbook/client/tinker/modelscope_service/self_cognition.py" +echo " python cookbook/client/tinker/self_cognition.py" echo "" diff --git a/README.md b/README.md index 49aa7deb..c82f1c15 100644 --- a/README.md +++ b/README.md @@ -100,10 +100,8 @@ sh INSTALL_MEGATRON.sh | DPO multi-LoRA training | transformers | [Script](cookbook/rl/dpo/dpo_multi_lora.py) | | GKD on-policy distillation | megatron | [Script](cookbook/rl/gkd/gkd_on_policy.py) | | GKD off-policy distillation | megatron | [Script](cookbook/rl/gkd/gkd_off_policy.py) | -| Tinker client finetuning (self-host) | transformers | [Script](cookbook/client/tinker/self_host) | -| Tinker client finetuning (ModelScope) | transformers | [Script](cookbook/client/tinker/modelscope) | -| Twinkle client finetuning (self-host) | transformers | [Script](cookbook/client/twinkle/self_host) | -| Twinkle client finetuning (ModelScope) | transformers | [Script](cookbook/client/twinkle/modelscope) | +| Tinker client finetuning | transformers | [Script](cookbook/client/tinker) | +| Twinkle client finetuning | transformers | [Script](cookbook/client/twinkle) | | Server startup scripts | transformers/megatron | [Script](cookbook/client/server) | ## Changelog diff --git a/README_ZH.md b/README_ZH.md index dcd95e39..31cab0f9 100644 --- a/README_ZH.md +++ b/README_ZH.md @@ -94,10 +94,8 @@ sh INSTALL_MEGATRON.sh | DPO 多 LoRA 训练 | transformers | [脚本](cookbook/rl/dpo/dpo_multi_lora.py) | | GKD 在线蒸馏 | megatron | [脚本](cookbook/rl/gkd/gkd_on_policy.py) | | GKD 离线蒸馏 | megatron | [脚本](cookbook/rl/gkd/gkd_off_policy.py) | -| Tinker 客户端微调（自部署） | transformers | [脚本](cookbook/client/tinker/self_host) | -| Tinker 客户端微调（ModelScope） | transformers | [脚本](cookbook/client/tinker/modelscope) | -| Twinkle 客户端微调（自部署） | transformers | [脚本](cookbook/client/twinkle/self_host) | -| Twinkle 客户端微调（ModelScope） | transformers | [脚本](cookbook/client/twinkle/modelscope) | +| Tinker 客户端微调 | transformers | [脚本](cookbook/client/tinker) | +| Twinkle 客户端微调 | transformers | [脚本](cookbook/client/twinkle) | | 服务端启动脚本 | transformers/megatron | [脚本](cookbook/client/server) | Twinkle✨支持相同的算法接口运行在单GPU、torchrun多机、Ray、Client等各场景下。其算法过程是外露的，非常便于修改和调试。完整的框架介绍请查看[快速开始](https://modelscope.github.io/twinkle-web/zh/docs/usage-guide/quick-start/) diff --git a/client_tools/client_generator.py b/client_tools/client_generator.py deleted file mode 100644 index bc2a5288..00000000 --- a/client_tools/client_generator.py +++ /dev/null @@ -1,918 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -import ast -from pathlib import Path -from typing import Dict, List, Set, Tuple - -AUTO_GEN_WARNING = """# ============================================================================ -# WARNING: AUTO-GENERATED FILE - DO NOT MODIFY MANUALLY! -# ============================================================================ -# This file is automatically generated by client_tools/client_generator.py -# Any manual changes will be overwritten when the generator runs again. -# -# To update this file: -# 1. Modify the source files in src/twinkle/ -# 2. Run: python client_tools/client_generator.py -# ============================================================================ -""" - - -def generate_processors(): - """Generate client wrappers for all classes with @remote_function methods.""" - - # Module mapping: module_name -> directory in src/twinkle - module_mapping = { - 'dataloader': 'dataloader', - 'dataset': 'dataset', - 'processor': 'processor', - 'reward': 'reward', - 'template': 'template', - 'weight_loader': 'weight_loader', - } - - # Map module names to processor types in the server - processor_type_mapping = { - 'dataloader': 'dataloader', - 'dataset': 'dataset', - 'processor': 'processor', - 'reward': 'reward', - 'template': 'template', - 'weight_loader': 'weight_loader', - } - - # Get the project root directory - project_root = Path(__file__).parent.parent - src_twinkle_path = project_root / 'src' / 'twinkle' - src_client_path = project_root / 'src' / 'twinkle_client' - - def get_method_signature(func_node: ast.FunctionDef) -> str: - """Extract method signature from AST node.""" - args = [] - - # Regular arguments - for i, arg in enumerate(func_node.args.args): - if arg.arg == 'self': - continue - - # Get argument name - arg_str = arg.arg - - # Get type annotation if available - if arg.annotation: - try: - arg_str += f': {ast.unparse(arg.annotation)}' - except: - pass - - # Get default value if available - defaults_offset = len(func_node.args.args) - len(func_node.args.defaults) - if i >= defaults_offset: - default_idx = i - defaults_offset - try: - default_val = ast.unparse(func_node.args.defaults[default_idx]) - arg_str += f' = {default_val}' - except: - pass - - args.append(arg_str) - - # *args - if func_node.args.vararg: - vararg_str = f'*{func_node.args.vararg.arg}' - if func_node.args.vararg.annotation: - try: - vararg_str += f': {ast.unparse(func_node.args.vararg.annotation)}' - except: - pass - args.append(vararg_str) - - # **kwargs - if func_node.args.kwarg: - kwarg_str = f'**{func_node.args.kwarg.arg}' - if func_node.args.kwarg.annotation: - try: - kwarg_str += f': {ast.unparse(func_node.args.kwarg.annotation)}' - except: - pass - args.append(kwarg_str) - - return ', '.join(args) - - def extract_typing_imports(signatures: List[str]) -> Set[str]: - """Extract required typing imports from signatures.""" - typing_patterns = { - 'Union[': 'Union', - 'Optional[': 'Optional', - 'List[': 'List', - 'Dict[': 'Dict', - 'Tuple[': 'Tuple', - 'Type[': 'Type', - 'Any': 'Any', - 'Callable': 'Callable', - 'Literal[': 'Literal', - 'Required[': 'Required', - 'Set[': 'Set', - 'TypedDict': 'TypedDict', - } - - all_text = ' '.join(signatures) - return {name for pattern, name in typing_patterns.items() if pattern in all_text} - - def extract_twinkle_imports(signatures: List[str]) -> Set[str]: - """Extract required twinkle imports from signatures.""" - twinkle_patterns = { - 'InputFeature': ['from twinkle.data_format import InputFeature'], - 'Trajectory': ['from twinkle.data_format import Trajectory'], - 'DataFilter': ['from twinkle.preprocessor import DataFilter'], - 'Preprocessor': ['from twinkle.preprocessor import Preprocessor'], - 'DatasetMeta': ['from twinkle.dataset import DatasetMeta'], - 'Dataset': ['from twinkle.dataset import Dataset'], - 'DeviceMesh': ['from twinkle import DeviceMesh'], - 'Template': ['from twinkle.template import Template'], - 'template.Template': ['from twinkle.template import Template', 'from twinkle import template'], - 'processor.InputProcessor': - ['from twinkle.processor import InputProcessor', 'from twinkle import processor'], - 'InputProcessor': ['from twinkle.processor import InputProcessor'], - } - - all_text = ' '.join(signatures) - imports = set() - for pattern, stmts in twinkle_patterns.items(): - if pattern in all_text: - imports.update(stmts) - - return imports - - def parse_params_from_signature(signature: str) -> List[str]: - """Parse parameter names from signature, handling nested brackets.""" - params = [] - current = '' - depth = 0 - - for char in signature + ',': - if char in '[(': - depth += 1 - elif char in '])': - depth -= 1 - - if char == ',' and depth == 0: - name = current.split(':')[0].split('=')[0].strip() - if name and name != 'self' and not name.startswith('*'): - params.append(name) - current = '' - else: - current += char - - return params - - def find_classes_with_remote_methods(file_path: Path) -> List[Tuple[str, str, List[Tuple[str, str]]]]: - """Find all classes that have @remote_function decorated methods.""" - try: - with open(file_path, 'r', encoding='utf-8') as f: - tree = ast.parse(f.read(), filename=str(file_path)) - except Exception as e: - print(f'Error parsing {file_path}: {e}') - return [] - - def has_remote_decorator(func: ast.FunctionDef) -> bool: - for dec in func.decorator_list: - if isinstance(dec, ast.Name) and dec.id == 'remote_function': - return True - if isinstance(dec, ast.Call): - func_node = dec.func - if isinstance(func_node, ast.Name) and func_node.id == 'remote_function': - return True - if isinstance(func_node, ast.Attribute) and func_node.attr == 'remote_function': - return True - return False - - def is_public_or_dunder(name: str) -> bool: - return (name.startswith('__') and name.endswith('__')) or not name.startswith('_') - - def get_base_name(node: ast.ClassDef) -> str: - if not node.bases: - return 'object' - base = node.bases[0] - if isinstance(base, ast.Name): - return base.id - if isinstance(base, ast.Attribute): - return base.attr - return 'object' - - classes_found = [] - for node in ast.walk(tree): - if not isinstance(node, ast.ClassDef): - continue - - methods = [ - (item.name, get_method_signature(item)) for item in node.body - if isinstance(item, ast.FunctionDef) and has_remote_decorator(item) and is_public_or_dunder(item.name) - ] - - # Extract __init__ signature separately (it may not have @remote_function) - init_signature = '' - for item in node.body: - if isinstance(item, ast.FunctionDef) and item.name == '__init__': - init_signature = get_method_signature(item) - break - - if methods: - classes_found.append((node.name, get_base_name(node), methods, init_signature)) - - return classes_found - - def generate_client_class(class_name: str, - base_class_name: str, - methods: List[Tuple[str, str]], - module_name: str, - processor_type: str, - source_filename: str, - has_base_file: bool, - init_signature: str = '') -> str: - """Generate client wrapper class code.""" - - def build_imports() -> Tuple[List[str], str]: - # Include both method signatures and __init__ signature for import detection. - # Use the timeout-injected signature for ``encode`` so the ``Optional`` - # import introduced by ``_inject_timeout`` is detected. - signatures = [] - for name, sig in methods: - if name == 'encode' and class_name in ('Dataset', 'LazyDataset'): - signatures.append(_inject_timeout(sig)) - else: - signatures.append(sig) - if init_signature: - signatures.append(init_signature) - - typing_imports = extract_typing_imports(signatures) - twinkle_imports = extract_twinkle_imports(signatures) - - lines = [] - if typing_imports: - lines.append(f"from typing import {', '.join(sorted(typing_imports))}") - lines.extend([ - 'from twinkle_client.http import http_post', - ]) - lines.extend(sorted(twinkle_imports)) - - if source_filename == 'base': - inheritance = 'object' - elif base_class_name == 'IterableDataset': - lines.append('from torch.utils.data import IterableDataset') - inheritance = 'IterableDataset' - elif has_base_file and base_class_name != 'object': - lines.append(f'from .base import {base_class_name}') - inheritance = base_class_name - else: - inheritance = 'object' - - lines.append('') - return lines, inheritance - - def _inject_timeout(signature: str) -> str: - """Insert `timeout: Optional[int] = 600` before any **kwargs in the signature.""" - if 'timeout' in signature: - return signature - if ', **' in signature: - pre, post = signature.rsplit(', **', 1) - return f'{pre}, timeout: Optional[int] = 600, **{post}' - if signature.startswith('**'): - return f'timeout: Optional[int] = 600, {signature}' - if signature: - return f'{signature}, timeout: Optional[int] = 600' - return 'timeout: Optional[int] = 600' - - def build_method(name: str, signature: str) -> str: - param_names = parse_params_from_signature(signature) - kwargs_dict = '{' + ', '.join(f"'{p}': {p}" for p in param_names) + '}' if param_names else '{}' - wants_timeout = name == 'encode' and class_name in ('Dataset', 'LazyDataset') - effective_sig = _inject_timeout(signature) if wants_timeout else signature - sig_part = f', {effective_sig}' if effective_sig else '' - if 'kwargs' in sig_part: - extra_args = '\n **kwargs' - else: - extra_args = '' - ret = 'self' if name == '__iter__' else 'response.json()["result"]' - timeout_kwarg = ',\n timeout=timeout' if wants_timeout else '' - - code = f''' - def {name}(self{sig_part}): - response = http_post( - url=f'{{self.server_url}}/call', - json_data={{ - 'processor_id': self.processor_id, - 'function': '{name}', - **{kwargs_dict},{extra_args} - }}{timeout_kwarg} - ) - response.raise_for_status() - return {ret} - ''' - if name == '__iter__': - code += ''' - def __next__(self): - response = http_post( - url=f'{self.server_url}/call', - json_data={ - 'processor_id': self.processor_id, - 'function': '__next__', - } - ) - response.raise_for_status() - return response.json()["result"] - ''' - return code - - import_lines, inheritance = build_imports() - - # Build __init__ method with actual signature - if init_signature: - # Extract parameter names from signature (excluding **kwargs) - param_names = parse_params_from_signature(init_signature) - init_params = f'self, {init_signature}' if init_signature else 'self' - - # Check if signature has **kwargs - has_kwargs = '**' in init_signature - - # Extract the **kwargs name if present - kwargs_name = None - if has_kwargs: - # Find the **kwargs parameter name - for part in init_signature.split(','): - part = part.strip() - if part.startswith('**'): - # Extract name after **, before : or end - kwargs_name = part[2:].split(':')[0].strip() - break - - # Build kwargs dict for HTTP request - if param_names: - kwargs_items = ', '.join([f"'{p}': {p}" for p in param_names]) - if has_kwargs and kwargs_name: - # Include both named params and **kwargs - kwargs_dict = f'{{{kwargs_items}}}, **{kwargs_name}' - else: - kwargs_dict = f'{{{kwargs_items}}}' - else: - if has_kwargs and kwargs_name: - kwargs_dict = kwargs_name - else: - kwargs_dict = '{}' - else: - # Fallback to **kwargs if no __init__ found - init_params = 'self, **kwargs' - kwargs_dict = 'kwargs' - - class_template = f'''{AUTO_GEN_WARNING} -{chr(10).join(import_lines)} -class {class_name}({inheritance}): - """Client wrapper for {class_name} that calls server HTTP endpoints.""" - - def __init__({init_params}): - from twinkle_client.http import get_base_url - - self.server_url = f'{{get_base_url()}}/processor/twinkle' - response = http_post( - url=f'{{self.server_url}}/create', - json_data={{ - 'processor_type': '{processor_type}', - 'class_type': '{class_name}', - **{kwargs_dict} - }} - ) - response.raise_for_status() - self.processor_id = response.json()['processor_id'] - - ''' - - method_codes = [build_method(name, sig) for name, sig in methods] - - return class_template + '\n'.join(method_codes) - - def scan_modules(src_twinkle_path: Path, module_mapping: Dict[str, str]) -> Dict: - """Scan all modules for classes with @remote_function methods.""" - print('Scanning src/twinkle modules for classes with @remote_function methods...') - - module_files = {} - for module_name, module_dir in module_mapping.items(): - module_path = src_twinkle_path / module_dir - if not module_path.exists(): - continue - - print(f' Scanning {module_name}...') - for py_file in module_path.glob('*.py'): - if py_file.name.startswith('_'): - continue - - if classes := find_classes_with_remote_methods(py_file): - module_files.setdefault(module_name, {}).setdefault(py_file.stem, []).extend(classes) - - return module_files - - def write_client_files(module_files: Dict, src_client_path: Path, processor_type_mapping: Dict[str, str]) -> None: - """Generate and write client files.""" - print('\nGenerating client classes...') - - for module_name, source_files in module_files.items(): - client_module_path = src_client_path / module_name - client_module_path.mkdir(parents=True, exist_ok=True) - - processor_type = processor_type_mapping.get(module_name, module_name) - has_base_file = 'base' in source_files - - for source_filename, classes in source_files.items(): - client_file = client_module_path / f'{source_filename}.py' - print(f' Writing {client_file}...') - - code = '\n\n'.join( - generate_client_class(class_name, base_class_name, methods, module_name, processor_type, - source_filename, has_base_file, init_signature) - for class_name, base_class_name, methods, init_signature in classes) - client_file.write_text(code, encoding='utf-8') - - def write_init_files(module_files: Dict, src_client_path: Path) -> None: - """Generate __init__.py files for each module.""" - print('\nGenerating __init__.py files...') - - for module_name, source_files in module_files.items(): - init_file = src_client_path / module_name / '__init__.py' - print(f' Writing {init_file}...') - - init_lines = [ - f'from .{source_filename} import {class_name}' - for source_filename, classes in sorted(source_files.items()) for class_name, _, _, _ in classes - ] - init_content = AUTO_GEN_WARNING + '\n'.join(sorted(init_lines)) + '\n' - init_file.write_text(init_content, encoding='utf-8') - - module_files = scan_modules(src_twinkle_path, module_mapping) - write_client_files(module_files, src_client_path, processor_type_mapping) - write_init_files(module_files, src_client_path) - print('\nProcessor client generation complete!') - return module_files - - -def generate_models(): - """Generate client wrapper for Model management.""" - from pathlib import Path - - project_root = Path(__file__).parent.parent - src_client_path = project_root / 'src' / 'twinkle_client' - client_module_path = src_client_path / 'model' - client_module_path.mkdir(parents=True, exist_ok=True) - - model_code = AUTO_GEN_WARNING + '''from typing import Any, Dict, Optional -from pathlib import Path -import time -from twinkle_client.http import http_get, http_post -from twinkle_client.types.model import ( - CalculateLossResponse, - CalculateMetricResponse, - ClipGradNormResponse, - ForwardBackwardResponse, - ForwardResponse, - GetStateDictResponse, - GetTrainConfigsResponse, - SaveResponse, - TrainingProgressResponse, -) - - -class MultiLoraTransformersModel: - """Client wrapper for TwinkleModel that calls server HTTP endpoints. - - This client manages adapters and sends training/inference requests to the model server. - The server-side session (managed by TwinkleClient) keeps the model alive. - """ - - def __init__(self, model_id: str, **kwargs): - """Initialize model client.""" - from twinkle_client.http import get_base_url - self.server_url = get_base_url() - - if '://' in model_id: - model_id = model_id.split('://')[1] - self.model_id = model_id - self.server_url = f'{self.server_url}/model/{model_id}/twinkle' - self.adapter_name = None - response = http_post( - url=f'{self.server_url}/create', - ) - response.raise_for_status() - - def add_adapter_to_model(self, adapter_name: str, config: Dict[str, Any], **kwargs) -> None: - """Add a new adapter to the model.""" - save_dir = kwargs.get('save_dir') - if save_dir: - kwargs['save_dir'] = Path(save_dir).expanduser().resolve().as_posix() - response = http_post( - url=f'{self.server_url}/add_adapter_to_model', - json_data={'adapter_name': adapter_name, 'config': config, **kwargs} - ) - response.raise_for_status() - self.adapter_name = adapter_name - - def forward(self, inputs: Any, **kwargs) -> ForwardResponse: - """Execute forward pass on the model.""" - response = http_post( - url=f'{self.server_url}/forward', - 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.""" - response = http_post( - url=f'{self.server_url}/forward_only', - json_data={'inputs': inputs, 'adapter_name': self.adapter_name, **kwargs} - ) - response.raise_for_status() - return ForwardResponse(**response.json()) - - def calculate_loss(self, **kwargs) -> CalculateLossResponse: - """Calculate loss from model outputs.""" - response = http_post( - url=f'{self.server_url}/calculate_loss', - json_data={'adapter_name': self.adapter_name, **kwargs} - ) - response.raise_for_status() - return CalculateLossResponse(**response.json()) - - def get_train_configs(self, **kwargs) -> GetTrainConfigsResponse: - """Get training configs.""" - response = http_post( - url=f'{self.server_url}/get_train_configs', - json_data={'adapter_name': self.adapter_name, **kwargs} - ) - response.raise_for_status() - return GetTrainConfigsResponse(**response.json()) - - def backward(self, **kwargs) -> None: - """Execute backward pass.""" - response = http_post( - url=f'{self.server_url}/backward', - json_data={'adapter_name': self.adapter_name, **kwargs} - ) - response.raise_for_status() - - def forward_backward(self, inputs: Any, **kwargs) -> ForwardBackwardResponse: - """Execute combined forward and backward pass.""" - response = http_post( - url=f'{self.server_url}/forward_backward', - json_data={'inputs': inputs, 'adapter_name': self.adapter_name, **kwargs} - ) - response.raise_for_status() - return ForwardBackwardResponse(**response.json()) - - def step(self, **kwargs) -> None: - """Execute optimizer step.""" - response = http_post( - url=f'{self.server_url}/step', - json_data={'adapter_name': self.adapter_name, **kwargs} - ) - response.raise_for_status() - - def zero_grad(self, **kwargs) -> None: - """Zero out gradients.""" - response = http_post( - url=f'{self.server_url}/zero_grad', - json_data={'adapter_name': self.adapter_name, **kwargs} - ) - response.raise_for_status() - - def lr_step(self, **kwargs) -> None: - """Execute learning rate scheduler step.""" - response = http_post( - url=f'{self.server_url}/lr_step', - json_data={'adapter_name': self.adapter_name, **kwargs} - ) - response.raise_for_status() - - def clip_grad_norm(self, max_grad_norm: float = 1.0, norm_type: int = 2, **kwargs) -> ClipGradNormResponse: - """Clip gradient norm.""" - response = http_post( - url=f'{self.server_url}/clip_grad_norm', - json_data={'max_grad_norm': max_grad_norm, 'norm_type': norm_type, 'adapter_name': self.adapter_name, **kwargs} - ) - response.raise_for_status() - return ClipGradNormResponse(**response.json()) - - def clip_grad_and_step(self, max_grad_norm: float = 1.0, norm_type: int = 2, **kwargs) -> None: - """Clip gradient norm and execute optimizer step in one call.""" - response = http_post( - url=f'{self.server_url}/clip_grad_and_step', - json_data={'max_grad_norm': max_grad_norm, 'norm_type': norm_type, 'adapter_name': self.adapter_name, **kwargs} - ) - response.raise_for_status() - - def set_loss(self, loss_cls: str, **kwargs) -> None: - """Set the loss function.""" - response = http_post( - url=f'{self.server_url}/set_loss', - json_data={'loss_cls': loss_cls, 'adapter_name': self.adapter_name, **kwargs} - ) - response.raise_for_status() - - def set_optimizer(self, optimizer_cls: str, **kwargs) -> None: - """Set the optimizer.""" - response = http_post( - url=f'{self.server_url}/set_optimizer', - json_data={'optimizer_cls': optimizer_cls, 'adapter_name': self.adapter_name, **kwargs} - ) - response.raise_for_status() - - def set_lr_scheduler(self, scheduler_cls: str, **kwargs) -> None: - """Set the learning rate scheduler.""" - response = http_post( - url=f'{self.server_url}/set_lr_scheduler', - json_data={'scheduler_cls': scheduler_cls, 'adapter_name': self.adapter_name, **kwargs} - ) - response.raise_for_status() - - def save(self, name: str, **kwargs) -> SaveResponse: - """Save model checkpoint.""" - response = http_post( - url=f'{self.server_url}/save', - json_data={'name': name, 'adapter_name': self.adapter_name, **kwargs} - ) - response.raise_for_status() - return SaveResponse(**response.json()) - - def load(self, name: str, **kwargs) -> None: - """Load model checkpoint.""" - response = http_post( - url=f'{self.server_url}/load', - json_data={'name': name, 'adapter_name': self.adapter_name, **kwargs} - ) - response.raise_for_status() - - def resume_from_checkpoint(self, name: str, *, resume_only_model: bool = False, **kwargs) -> Dict[str, Any]: - response = http_post( - url=f'{self.server_url}/resume_from_checkpoint', - json_data={'name': name, 'adapter_name': self.adapter_name, - 'resume_only_model': resume_only_model, **kwargs} - ) - response.raise_for_status() - return TrainingProgressResponse(**response.json()).result - - def apply_patch(self, patch_cls: str, **kwargs) -> None: - """Apply a patch to the model.""" - response = http_post( - url=f'{self.server_url}/apply_patch', - json_data={'patch_cls': patch_cls, 'adapter_name': self.adapter_name, **kwargs} - ) - response.raise_for_status() - - def add_metric(self, metric_cls: str, is_training: Optional[bool] = None, **kwargs) -> None: - """Add a metric to the model.""" - response = http_post( - url=f'{self.server_url}/add_metric', - json_data={'metric_cls': metric_cls, 'is_training': is_training, 'adapter_name': self.adapter_name, **kwargs} - ) - response.raise_for_status() - - def set_template(self, template_cls: str, **kwargs) -> None: - """Set the template for data processing.""" - response = http_post( - url=f'{self.server_url}/set_template', - json_data={'template_cls': template_cls, 'adapter_name': self.adapter_name, 'model_id': self.model_id, **kwargs} - ) - response.raise_for_status() - - def set_processor(self, processor_cls: str, **kwargs) -> None: - """Set the input processor.""" - response = http_post( - url=f'{self.server_url}/set_processor', - json_data={'processor_cls': processor_cls, 'adapter_name': self.adapter_name, **kwargs} - ) - response.raise_for_status() - - def calculate_metric(self, is_training: bool = True, **kwargs) -> CalculateMetricResponse: - """Calculate metrics from model outputs.""" - response = http_post( - url=f'{self.server_url}/calculate_metric', - json_data={'is_training': is_training, 'adapter_name': self.adapter_name, **kwargs} - ) - response.raise_for_status() - return CalculateMetricResponse(**response.json()) - - def get_state_dict(self, **kwargs) -> GetStateDictResponse: - """Get model state dictionary.""" - response = http_post( - url=f'{self.server_url}/get_state_dict', - json_data={'adapter_name': self.adapter_name, **kwargs} - ) - response.raise_for_status() - return GetStateDictResponse(**response.json()) - - def upload_to_hub( - self, - checkpoint_dir: str, - hub_model_id: str, - hub_token: Optional[str] = None, - async_upload: bool = True, - poll_interval: float = 5.0, - ) -> None: - """Upload model checkpoint to hub. - - Submits the upload task to the server and polls for completion. - Blocks until the upload finishes or raises on failure. - - Args: - checkpoint_dir: The directory path of the checkpoint to upload. - hub_model_id: The hub model id. - hub_token: The hub token (optional). - async_upload: Deprecated, has no effect. The server always runs the - upload in the background and the client polls for completion. - poll_interval: Seconds between status poll requests (default: 5). - """ - response = http_post( - url=f'{self.server_url}/upload_to_hub', - json_data={ - 'checkpoint_dir': checkpoint_dir, - 'hub_model_id': hub_model_id, - 'hub_token': hub_token, - } - ) - response.raise_for_status() - request_id = response.json().get('request_id') - if not request_id: - return - - print(f'[upload_to_hub] Upload started (task {request_id}), waiting for completion...') - while True: - status_resp = http_get(url=f'{self.server_url}/upload_status/{request_id}') - status_resp.raise_for_status() - data = status_resp.json() - status = data.get('status', 'unknown') - if status == 'completed': - print(f'[upload_to_hub] Upload completed successfully.') - return - elif status == 'failed': - error = data.get('error', 'Unknown error') - raise RuntimeError(f'[upload_to_hub] Upload failed: {error}') - else: - print(f'[upload_to_hub] Status: {status}...') - time.sleep(poll_interval) -''' - - # Write the model client file - client_file = client_module_path / 'multi_lora_transformers.py' - print(f'Generating {client_file}...') - with open(client_file, 'w', encoding='utf-8') as f: - f.write(model_code) - - # Create/overwrite __init__.py - init_file = client_module_path / '__init__.py' - init_content = AUTO_GEN_WARNING + 'from .multi_lora_transformers import MultiLoraTransformersModel\n' - print(f'Writing {init_file}...') - with open(init_file, 'w', encoding='utf-8') as f: - f.write(init_content) - - print('Model client generation complete!') - - -def generate_samplers(): - """Generate client wrapper for Sampler management.""" - from pathlib import Path - - project_root = Path(__file__).parent.parent - src_client_path = project_root / 'src' / 'twinkle_client' - client_module_path = src_client_path / 'sampler' - client_module_path.mkdir(parents=True, exist_ok=True) - - sampler_code = AUTO_GEN_WARNING + '''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 - - -# Intentionally does NOT subclass ``twinkle.sampler.base.Sampler``: importing -# that base pulls ``twinkle.sampler.__init__`` → ``VLLMEngine`` → torch + zmq, -# which the mock / CPU-only client environments don't have. -class vLLMSampler: - """Client wrapper for Sampler that calls server HTTP endpoints. - - This client manages sampling operations and adapter synchronization with the sampler server. - The server-side session (managed by TwinkleClient) keeps the sampler alive. - """ - - 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() - - self.adapter_name = None - if '://' in model_id: - model_id = model_id.split('://')[1] - self.model_id = model_id - self.server_url = f'{self.server_url}/sampler/{model_id}/twinkle' - response = http_post( - url=f'{self.server_url}/create', - json_data=kwargs - ) - response.raise_for_status() - - def add_adapter_to_sampler(self, adapter_name: str, config: PeftConfig, **kwargs) -> AddAdapterResponse: - """Add a new adapter to the sampler.""" - if isinstance(config, PeftConfig): - config = config.__dict__ - response = http_post( - url=f'{self.server_url}/add_adapter_to_sampler', - json_data={'adapter_name': adapter_name, 'config': config, **kwargs} - ) - response.raise_for_status() - self.adapter_name = adapter_name - return AddAdapterResponse(**response.json()) - - def sample( - 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]: - """Sample from the model. - - Args: - inputs: List of Trajectory or InputFeature to sample from. - sampling_params: Sampling parameters dict. - adapter_name: Adapter name for LoRA inference. - adapter_uri: Adapter URI (twinkle:// path or local path) for LoRA inference. - num_samples: Number of completions to generate per prompt. - - Returns: - SampleResponseModel with 'sequences' list, each containing tokens, logprobs, stop_reason. - """ - json_data = { - 'inputs': inputs, - 'sampling_params': sampling_params, - 'adapter_name': adapter_name, - 'num_samples': num_samples, - } - if adapter_uri is not None: - json_data['adapter_uri'] = adapter_uri - - response = http_post( - url=f'{self.server_url}/sample', - json_data=json_data - ) - response.raise_for_status() - return [SampleResponseModel(**r) for r in response.json()['samples']] - - def set_template(self, template_cls: str, adapter_name: str = '', **kwargs) -> SetTemplateResponse: - """Set the template for encoding trajectories.""" - response = http_post( - url=f'{self.server_url}/set_template', - json_data={'template_cls': template_cls, 'adapter_name': adapter_name, **kwargs} - ) - response.raise_for_status() - return SetTemplateResponse(**response.json()) - - def apply_patch(self, patch_cls: str, **kwargs) -> None: - """Apply a patch to the model.""" - response = http_post( - url=f'{self.server_url}/apply_patch', - json_data={'patch_cls': patch_cls, 'adapter_name': self.adapter_name, **kwargs} - ) - response.raise_for_status() -''' - - # Write the sampler client file - client_file = client_module_path / 'vllm_sampler.py' - print(f'Generating {client_file}...') - with open(client_file, 'w', encoding='utf-8') as f: - f.write(sampler_code) - - # Create/overwrite __init__.py - init_file = client_module_path / '__init__.py' - init_content = AUTO_GEN_WARNING + 'from .vllm_sampler import vLLMSampler\n' - print(f'Writing {init_file}...') - with open(init_file, 'w', encoding='utf-8') as f: - f.write(init_content) - - print('Sampler client generation complete!') - - -if __name__ == '__main__': - print('Starting client code generation...\n') - print('=' * 60) - - # Generate processor-based clients - print('\n[1/3] Generating processor-based clients...') - generate_processors() - - # Generate model client - print('\n' + '=' * 60) - print('\n[2/3] Generating model client...') - generate_models() - - # Generate sampler client - print('\n' + '=' * 60) - print('\n[3/3] Generating sampler client...') - generate_samplers() - - print('\n' + '=' * 60) - print('\n✓ All client code generation complete!\n') diff --git a/cookbook/client/server/megatron/entrypoint.sh b/cookbook/client/server/megatron/entrypoint.sh index 879e3890..1fcff4a8 100755 --- a/cookbook/client/server/megatron/entrypoint.sh +++ b/cookbook/client/server/megatron/entrypoint.sh @@ -9,6 +9,8 @@ SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" RUN_SCRIPT="$SCRIPT_DIR/run.sh" TWINKLE_WORK_DIR="${TWINKLE_WORK_DIR:-/dashscope/caches/application/twinkle}" TEMP_DIR="${TWINKLE_TEMP_DIR:-/dashscope/caches/application/ray_logs}" +export MODELSCOPE_CACHE="${MODELSCOPE_CACHE:-/dashscope/caches/application/.cache}" +export FLA_TILELANG="${FLA_TILELANG:-0}" LOG_FILE="$TWINKLE_WORK_DIR/run.log" TWINKLE_HEALTH_URL="${TWINKLE_HEALTH_URL:-http://127.0.0.1:9000/api/v1/healthz}" TWINKLE_DEEP_HEALTH_URL="${TWINKLE_DEEP_HEALTH_URL:-http://127.0.0.1:9000/api/v1/twinkle/healthz/deep}" @@ -21,6 +23,8 @@ RESTART_BACKOFF_SECONDS="${TWINKLE_ENTRYPOINT_RESTART_BACKOFF_SECONDS:-10}" CHILD_PID="" +mkdir -p "$TWINKLE_WORK_DIR" "$MODELSCOPE_CACHE" + print_warning() { echo -e "\033[33m[WARNING]\033[0m $1" } diff --git a/cookbook/client/tinker/modelscope/dpo.py b/cookbook/client/tinker/dpo.py similarity index 97% rename from cookbook/client/tinker/modelscope/dpo.py rename to cookbook/client/tinker/dpo.py index cfbe916e..6091e808 100644 --- a/cookbook/client/tinker/modelscope/dpo.py +++ b/cookbook/client/tinker/dpo.py @@ -39,9 +39,9 @@ # --------------------------------------------------------------------------- # Configuration # --------------------------------------------------------------------------- -base_model = 'Qwen/Qwen3.6-27B' -base_url = 'http://www.modelscope.cn/twinkle' -api_key = os.environ.get('MODELSCOPE_TOKEN') +base_model = os.environ.get('TWINKLE_MODEL_ID', 'Qwen/Qwen3.5-4B') +base_url = os.environ.get('TWINKLE_SERVER_URL', 'http://localhost:8000') +api_key = os.environ.get('TWINKLE_SERVER_TOKEN', 'EMPTY_TOKEN') dataset_id = 'ms://hjh0119/shareAI-Llama3-DPO-zh-en-emoji' batch_size = 4 diff --git a/cookbook/client/tinker/self_host/lora.py b/cookbook/client/tinker/lora.py similarity index 95% rename from cookbook/client/tinker/self_host/lora.py rename to cookbook/client/tinker/lora.py index 292e9115..2b2bbe86 100644 --- a/cookbook/client/tinker/self_host/lora.py +++ b/cookbook/client/tinker/lora.py @@ -10,6 +10,8 @@ import dotenv dotenv.load_dotenv('.env') +import os + # Step 2: Initialize Tinker client before importing ServiceClient from twinkle import init_tinker_client @@ -19,12 +21,8 @@ from tinker import ServiceClient service_client = ServiceClient( - # BASE_URL can be a local server endpoint such as http://localhost:8000, or - # points to a previously deployed remote server, or - # modelscope server such as 'http://www.modelscope.cn/twinkle' - base_url='http://localhost:8000', - # API_KEY can be empty or a meaninful one according to sever configuration - api_key='EMPTY-TOKEN' + base_url=os.environ.get('TWINKLE_SERVER_URL', 'http://localhost:8000'), + api_key=os.environ.get('TWINKLE_SERVER_TOKEN', 'EMPTY_TOKEN'), ) # Step 4: List models available on the server to verify the connection @@ -61,7 +59,7 @@ # Step 6: Create or resume a training client. # If resume_path is set, it restores both model weights and optimizer state. -base_model = 'Qwen/Qwen3.5-4B' +base_model = os.environ.get('TWINKLE_MODEL_ID', 'Qwen/Qwen3.5-4B') if not resume_path: training_client = service_client.create_lora_training_client(base_model=base_model) else: diff --git a/cookbook/client/tinker/modelscope/sample.py b/cookbook/client/tinker/modelscope/sample.py deleted file mode 100644 index 7e03e0bd..00000000 --- a/cookbook/client/tinker/modelscope/sample.py +++ /dev/null @@ -1,69 +0,0 @@ -# Tinker-Compatible Client - Sampling / Inference Example -# -# This script demonstrates how to use a previously trained LoRA checkpoint -# for text generation (sampling) via the Tinker-compatible client API. -# The server must be running first (see server.py and server_config.yaml). - -import os -from tinker import types - -from twinkle.data_format import Message, Trajectory -from twinkle.template import Template -from twinkle import init_tinker_client - -# Step 1: Initialize Tinker client -init_tinker_client() - -from tinker import ServiceClient - -base_model = 'Qwen/Qwen3.6-27B' -base_url = 'http://www.modelscope.cn/twinkle' - -# Step 2: Define the base model and connect to the server -service_client = ServiceClient( - base_url=base_url, - api_key=os.environ.get('MODELSCOPE_TOKEN') -) - -# Step 3: Create a sampling client by loading weights from a saved checkpoint. -# The model_path is a twinkle:// URI pointing to a previously saved LoRA checkpoint. -# The server will load the base model and apply the LoRA adapter weights. -sampling_client = service_client.create_sampling_client( - model_path='twinkle://xxx-Qwen_Qwen3.6-35B-A3B-xxx/weights/twinkle-lora-1', - base_model=base_model -) - -# Step 4: Load the tokenizer locally to encode the prompt and decode the results -print(f'Using model {base_model}') - -template = Template(model_id=f'ms://{base_model}') - -trajectory = Trajectory( - messages=[ - Message(role='system', content='You are a helpful assistant'), - Message(role='user', content='Who are you?'), - ] -) - -input_feature = template.encode(trajectory, add_generation_prompt=True) - -input_ids = input_feature['input_ids'].tolist() - -# Step 5: Prepare the prompt and sampling parameters -prompt = types.ModelInput.from_ints(input_ids) -params = types.SamplingParams( - max_tokens=128, # Maximum number of tokens to generate - temperature=0.7, - stop=['\n'] # Stop generation when a newline character is produced -) - -# Step 6: Send the sampling request to the server. -# num_samples=1 generates 1 independent completions for the same prompt. -print('Sampling...') -future = sampling_client.sample(prompt=prompt, sampling_params=params, num_samples=1) -result = future.result() - -# Step 7: Decode and print the generated responses -print('Responses:') -for i, seq in enumerate(result.sequences): - print(f'{i}: {repr(template.decode(seq.tokens))}') diff --git a/cookbook/client/tinker/modelscope/self_cognition.py b/cookbook/client/tinker/modelscope/self_cognition.py deleted file mode 100644 index f74ec073..00000000 --- a/cookbook/client/tinker/modelscope/self_cognition.py +++ /dev/null @@ -1,136 +0,0 @@ -# Tinker-Compatible Client - Self-Cognition Training & Evaluation Example -# -# This script demonstrates two workflows using the Tinker-compatible client: -# 1. train(): Fine-tune a model on a self-cognition dataset so it learns -# a custom identity (name, author). -# 2. eval(): Load a trained checkpoint and sample from it to verify -# that the model has learned the custom identity. -# The server must be running first (see server.py and server_config.yaml). -import os -from tqdm import tqdm -from tinker import types -from twinkle import init_tinker_client -from twinkle.data_format import Message, Trajectory -from twinkle.template import Template -from twinkle.dataloader import DataLoader -from twinkle.dataset import Dataset, DatasetMeta -from twinkle.preprocessor import SelfCognitionProcessor -from twinkle.server.common import input_feature_to_datum - -# Initialize the Tinker client before importing ServiceClient -init_tinker_client() - -from tinker import ServiceClient - -# The base model to fine-tune / evaluate -base_model = 'Qwen/Qwen3.6-27B' -base_url = 'http://www.modelscope.cn/twinkle' - - -def train(): - # Step 1: Prepare the dataset - - # Load the self-cognition dataset from ModelScope (first 500 examples) - dataset = Dataset(dataset_meta=DatasetMeta('ms://swift/self-cognition', data_slice=range(500))) - - # Apply the chat template matching the base model (max 256 tokens per sample) - dataset.set_template('Qwen3_5Template', model_id=f'ms://{base_model}', max_length=256) - - # Replace placeholder names with custom model/author identity - dataset.map(SelfCognitionProcessor('twinkle模型', 'twinkle团队'), load_from_cache_file=False) - - # Tokenize and encode the dataset into model-ready input features - dataset.encode(batched=True, load_from_cache_file=False) - - # Wrap the dataset into a DataLoader that yields batches of size 8 - dataloader = DataLoader(dataset=dataset, batch_size=8) - - # Step 2: Initialize the training client - - - service_client = ServiceClient( - base_url=base_url, - api_key=os.environ.get('MODELSCOPE_TOKEN') - ) - - # Create a LoRA training client for the base model (rank=16 for the LoRA adapter) - training_client = service_client.create_lora_training_client(base_model=base_model, rank=16) - - # Step 3: Run the training loop - - for epoch in range(3): - print(f'Epoch {epoch}') - for step, batch in tqdm(enumerate(dataloader)): - # Convert each InputFeature into a Datum for the Tinker API - input_datum = [input_feature_to_datum(input_feature) for input_feature in batch] - - # Send data to server: forward + backward pass (computes gradients) - fwdbwd_future = training_client.forward_backward(input_datum, 'cross_entropy') - - # Optimizer step: update model weights with Adam - optim_future = training_client.optim_step(types.AdamParams(learning_rate=1e-4)) - - # Wait for both operations to complete - fwdbwd_result = fwdbwd_future.result() - optim_result = optim_future.result() - - # Compute weighted average log-loss per token for monitoring - # logprobs = np.concatenate([output['logprobs'].tolist() for output in fwdbwd_result.loss_fn_outputs]) - # weights = np.concatenate([example.loss_fn_inputs['weights'].tolist() for example in input_datum]) - # print(f'Loss per token: {-np.dot(logprobs, weights) / weights.sum():.4f}') - print(f'Training Metrics: {optim_result}') - - # Save a checkpoint after each epoch - save_future = training_client.save_state(f'twinkle-lora-{epoch}') - save_result = save_future.result() - print(f'Saved checkpoint to {save_result.path}') - - -def eval(): - # Step 1: Load the trained LoRA checkpoint for inference - - # Path to a previously saved LoRA checkpoint (twinkle:// URI) - weight_path = 'twinkle://20260212_174205-Qwen_Qwen2_5-7B-Instruct-51edc9ed/weights/twinkle-lora-2' - - service_client = ServiceClient(base_url=base_url, api_key=os.environ.get('MODELSCOPE_TOKEN')) - sampling_client = service_client.create_sampling_client(model_path=weight_path, base_model=base_model) - - # Step 2: Prepare the chat prompt - - # Build a multi-turn conversation to test the model's self-cognition - template = Template(model_id=f'ms://{base_model}') - - trajectory = Trajectory( - messages=[ - Message(role='system', content='You are a helpful assistant'), - Message(role='user', content='你是谁？'), - ] - ) - - input_feature = template.batch_encode([trajectory], add_generation_prompt=True)[0] - - input_ids = input_feature['input_ids'].tolist() - - # Step 3: Generate responses - - prompt = types.ModelInput.from_ints(input_ids) - params = types.SamplingParams( - max_tokens=50, # Maximum tokens to generate - temperature=0.2, # Low temperature for more focused responses - stop=['\n'] # Stop at newline - ) - - # Sample 8 independent completions - print('Sampling...') - future = sampling_client.sample(prompt=prompt, sampling_params=params, num_samples=8) - result = future.result() - - # Decode and print each response - print('Responses:') - for i, seq in enumerate(result.sequences): - print(f'{i}: {repr(template.decode(seq.tokens))}') - - -if __name__ == '__main__': - train() # Uncomment to run training - # eval() # Run evaluation / inference diff --git a/cookbook/client/tinker/modelscope/short_math_grpo.py b/cookbook/client/tinker/modelscope/short_math_grpo.py deleted file mode 100644 index 780df444..00000000 --- a/cookbook/client/tinker/modelscope/short_math_grpo.py +++ /dev/null @@ -1,325 +0,0 @@ -# Tinker-Compatible Client - GSM8K GRPO Training Example -# -# This script demonstrates GSM8K math problem training using the -# Tinker-compatible client API with save_weights_for_sampler for weight sync. -# Instead of calling sync_weights directly, it periodically saves weights and -# creates a sampling client for generation. -# -# Flow: -# 1. Prepare GSM8K dataset (client-side) -# 2. Initialize Tinker-compatible training & sampling clients -# 3. Training loop: -# a. Every SYNC_INTERVAL steps: save_weights_for_sampler → sampling_client -# b. Sample completions from the sampling client -# c. Compute rewards and advantages (client-side) -# d. Train on sampled data weighted by advantages -# e. Optimizer step -# -# The server must be running first (see server.py and server_config.yaml). -# Requires both model and sampler services to be configured. -import gc -import numpy as np -import os -import re -from tinker import types -from typing import List, Tuple, Dict, Any - -from twinkle import init_tinker_client -from twinkle import get_logger -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.reward.base import Reward -from twinkle.metric import CompletionRewardMetric -from twinkle.template import Qwen3_5Template - -logger = get_logger() - -# ========== Configuration ========== -BASE_MODEL = 'Qwen/Qwen3.6-27B' -NUM_GENERATIONS = 4 -MAX_NEW_TOKENS = 4096 -LEARNING_RATE = 2e-5 -MAX_STEPS = 1000 -BATCH_SIZE = 2 -TEMPERATURE = 1.0 -SYNC_INTERVAL = 1 # Save weights for sampler every N steps -LORA_RANK = 16 -DATA_NUM = 2000 # Number of Math samples to use - -SYSTEM_PROMPT = ('You are a helpful math assistant. Solve the problem with minimal but correct reasoning ' - 'and put your final answer within \\boxed{}.') - - -# ========== Reward Functions ========== -class GSM8KBrevityReward(Reward): - """Brevity reward: rewards shorter completions that contain a valid answer. - - Returns 0.0 if no valid answer format (\\boxed{} or ####). - Otherwise returns higher score for shorter completions (1.0 at <=200 chars). - """ - - def __call__(self, trajectories: List[Dict[str, Any]], **kwargs) -> List[float]: - rewards = [] - for traj in trajectories: - messages = traj.get('messages', []) - completion = '' - for msg in reversed(messages): - if msg.get('role') == 'assistant': - completion = msg.get('content', '') - break - - has_answer = bool( - re.search(r'\\boxed\{[^}]+\}', completion) - or re.search(r'####\s*[\-\d,\.]+', completion) - ) - - if not has_answer: - rewards.append(0.0) - else: - length = len(completion) - if length <= 200: - rewards.append(1.0) - else: - rewards.append(max(0.0, 1.0 - (length - 200) / 3000)) - return rewards - - -# ========== Dataset ========== -def create_gsm8k_dataset(): - """Create GSM8K dataset.""" - dataset = Dataset(DatasetMeta('ms://modelscope/gsm8k', subset_name='main', split='train', data_slice=range(DATA_NUM))) - dataset.set_template('Qwen3_5Template', model_id=f'ms://{BASE_MODEL}', max_length=4096, - truncation_strategy='delete', enable_thinking=True) - dataset.map(GSM8KProcessor(system=SYSTEM_PROMPT)) - dataset.encode(add_generation_prompt=True) - return dataset - - -def compute_rewards( - trajectories: List[Dict[str, Any]], -) -> Tuple[List[float], List[float], List[float]]: - """Compute accuracy and brevity rewards for GSM8K.""" - accuracy_reward_fn = GSM8KAccuracyReward() - brevity_reward_fn = GSM8KBrevityReward() - - accuracy_rewards = accuracy_reward_fn(trajectories) - brevity_rewards = brevity_reward_fn(trajectories) - total_rewards = [a + b for a, b in zip(accuracy_rewards, brevity_rewards)] - return total_rewards, brevity_rewards, accuracy_rewards - - -def main(): - logger.info('Starting GSM8K GRPO training...') - - # Step 1: Prepare dataset and dataloader (client-side) - dataset = create_gsm8k_dataset() - dataloader = DataLoader(dataset=dataset, batch_size=BATCH_SIZE, num_workers=0) - template = Qwen3_5Template(model_id=f'ms://{BASE_MODEL}') - - logger.info('Dataset and template initialized') - - # Step 2: Initialize the Tinker-compatible client - logger.info('Connecting to Tinker server...') - init_tinker_client() - - from tinker import ServiceClient - service_client = ServiceClient( - base_url='http://www.modelscope.cn/twinkle', - api_key=os.environ.get('MODELSCOPE_TOKEN') - ) - - logger.info('Creating LoRA training client...') - # Create a LoRA training client for GRPO - training_client = service_client.create_lora_training_client( - base_model=BASE_MODEL, - rank=LORA_RANK, - ) - - logger.info('Training client created successfully') - - # Step 3: Setup metrics and advantage function - advantage_fn = GRPOAdvantage() - metrics = CompletionRewardMetric() - - sampling_params = types.SamplingParams( - max_tokens=MAX_NEW_TOKENS, - temperature=TEMPERATURE, - top_p=0.95, - ) - - # The sampling client is created on-demand via save_weights_for_sampler - sampling_client = None - - step = 0 - for batch in dataloader: - if step >= MAX_STEPS: - break - - metrics.reset() - prompts = batch if isinstance(batch, list) else [batch] - - # ========== 1. Save weights for sampler (instead of sync_weights) ========== - if step % SYNC_INTERVAL == 0: - logger.info(f'Step {step}: Saving weights for sampler...') - - sampling_client = training_client.save_weights_and_get_sampling_client() - logger.info(f'Step {step}: Sampling client ready') - - if sampling_client is None: - logger.warning('No sampling client available, skipping step') - step += 1 - continue - - # ========== 2. Sample completions ========== - # Convert input features to token prompts for the sampling client - all_sequences = [] - all_user_data = [] - for prompt_feature in prompts: - input_ids = prompt_feature['input_ids'] - if hasattr(input_ids, 'tolist'): - input_ids = input_ids.tolist() - prompt = types.ModelInput.from_ints(input_ids) - future = sampling_client.sample( - prompt=prompt, - sampling_params=sampling_params, - num_samples=NUM_GENERATIONS, - ) - result = future.result() - # Store both sequences and user data - for _ in range(NUM_GENERATIONS): - all_user_data.append(prompt_feature.get('user_data', [])) - all_sequences.extend(result.sequences) - - if not all_sequences: - logger.warning(f'Step {step}: No valid samples, skipping') - step += 1 - continue - - # ========== 3. Build trajectories and collect logprobs ========== - trajectories = [] - old_logps_list = [] - completion_lengths = [] - - for idx, seq in enumerate(all_sequences): - decoded_text = template.decode(seq.tokens, skip_special_tokens=True) - # Use the corresponding user data for this sequence - trajectories.append({ - 'messages': [ - { - 'role': 'system', - 'content': SYSTEM_PROMPT - }, - { - 'role': 'user', - 'content': 'Math problem' - }, # Placeholder - { - 'role': 'assistant', - 'content': decoded_text - } - ], - 'user_data': - all_user_data[idx] - }) - old_logps_list.append([lp for lp in seq.logprobs] if seq.logprobs else []) - completion_lengths.append(len(seq.tokens)) - - # ========== 4. Compute rewards ========== - total_rewards, brevity_rewards, accuracy_rewards = compute_rewards(trajectories) - metrics.accumulate( - completion_lengths=completion_lengths, - rewards={ - 'total': total_rewards, - 'brevity': brevity_rewards, - 'accuracy': accuracy_rewards, - }) - - # ========== 5. Compute advantages ========== - advantages = advantage_fn( - total_rewards, - num_generations=NUM_GENERATIONS, - scale='group', - ).tolist() - - frac_zero_std = (1.0 if all(abs(a) < 1e-8 for a in advantages) else 0.0) - if frac_zero_std == 1.0: - logger.info(f'Step {step}: All advantages are zero, skipping training') - step += 1 - continue - - # ========== 6. Train the policies with GRPO loss ========== - # Train the policies with the Advantage-Regularized policy - # gradient (GRPO) loss function. - # - # The GRPO loss function requires: - # 1. logprobs: The log probabilities of the tokens under the current policy - # 2. advantages: The advantage values for each completion - # - # The training data is constructed with: - # - model_input: The full prompt + completion tokens - # - target_tokens: The shifted tokens for next-token prediction - # - logprobs: The log probabilities from the sampling step - # - advantages: The computed advantage values - training_data = [] - for i, seq in enumerate(all_sequences): - # Build a Datum from the completion tokens with logprobs and advantages - prompt_feature = prompts[i // NUM_GENERATIONS] - prompt_ids = prompt_feature['input_ids'] - if hasattr(prompt_ids, 'tolist'): - prompt_ids = prompt_ids.tolist() - - sampled_tokens = list(seq.tokens) - logprobs = seq.logprobs if seq.logprobs else [0.0] * len(sampled_tokens) - advantage = float(advantages[i]) - - ob_len = len(prompt_ids) - 1 - input_tokens = prompt_ids + sampled_tokens[:-1] - target_tokens = [0] * ob_len + sampled_tokens - weights = [0] * ob_len + [1] * len(sampled_tokens) - padded_advantages = [0.0] * ob_len + [advantage] * len(sampled_tokens) - padded_logprobs = [0.0] * ob_len + logprobs - - datum = types.Datum( - model_input=types.ModelInput.from_ints(input_tokens), - loss_fn_inputs={ - 'target_tokens': target_tokens, - 'weights': weights, - 'logprobs': types.TensorData.from_numpy(np.array(padded_logprobs, dtype=np.float32)), - 'advantages': types.TensorData.from_numpy(np.array(padded_advantages, dtype=np.float32)), - }, - ) - training_data.append(datum) - - if not training_data: - logger.info(f'Step {step}: No training data constructed, skipping') - step += 1 - continue - - # Forward-backward pass with importance_sampling (GRPO) loss - # The training data already contains logprobs and advantages for the GRPO loss - fwdbwd_result = training_client.forward_backward(training_data, 'importance_sampling').result() - - optim_result = training_client.optim_step(types.AdamParams(learning_rate=LEARNING_RATE)).result() - - gc.collect() - - # ========== 7. Log ========== - log_dict = metrics.calculate() - if optim_result.metrics: - log_dict.update(optim_result.metrics) - log_dict['train/frac_reward_zero_std'] = frac_zero_std - log_dict['train/num_training_samples'] = len(training_data) - logger.info(f'Step {step}: {log_dict}') - step += 1 - - # Save final checkpoint - save_future = training_client.save_state('gsm8k-grpo-final') - save_result = save_future.result() - logger.info(f'Saved final checkpoint to {save_result.path}') - - -if __name__ == '__main__': - main() diff --git a/cookbook/client/tinker/self_host/multi_modal.py b/cookbook/client/tinker/multi_modal.py similarity index 97% rename from cookbook/client/tinker/self_host/multi_modal.py rename to cookbook/client/tinker/multi_modal.py index 6454c1ec..dae6a4dd 100644 --- a/cookbook/client/tinker/self_host/multi_modal.py +++ b/cookbook/client/tinker/multi_modal.py @@ -38,8 +38,9 @@ # ============================================================================= # Step 3: Configuration # ============================================================================= -base_model = 'Qwen/Qwen3.5-4B' # Multimodal vision-language model -base_url = 'http://localhost:8000' +base_model = os.environ.get('TWINKLE_MODEL_ID', 'Qwen/Qwen3.5-4B') +base_url = os.environ.get('TWINKLE_SERVER_URL', 'http://localhost:8000') +api_key = os.environ.get('TWINKLE_SERVER_TOKEN', 'EMPTY_TOKEN') # ============================================================================= # Step 4: Define Preprocessor (identical to Twinkle version) @@ -100,7 +101,7 @@ def train(): logger.info(f'Connecting to Tinker server at {base_url}') service_client = ServiceClient( base_url=base_url, - api_key=os.environ.get('MODELSCOPE_TOKEN', 'EMPTY-TOKEN') + api_key=api_key, ) training_client = service_client.create_lora_training_client( diff --git a/cookbook/client/tinker/self_host/sample.py b/cookbook/client/tinker/sample.py similarity index 91% rename from cookbook/client/tinker/self_host/sample.py rename to cookbook/client/tinker/sample.py index 1f4dfb27..2d94c95a 100644 --- a/cookbook/client/tinker/self_host/sample.py +++ b/cookbook/client/tinker/sample.py @@ -17,10 +17,10 @@ from tinker import ServiceClient # Step 2: Define the base model and connect to the server -base_model = 'Qwen/Qwen3.5-4B' +base_model = os.environ.get('TWINKLE_MODEL_ID', 'Qwen/Qwen3.5-4B') service_client = ServiceClient( - base_url='http://localhost:8000', - api_key='EMPTY_TOKEN' + base_url=os.environ.get('TWINKLE_SERVER_URL', 'http://localhost:8000'), + api_key=os.environ.get('TWINKLE_SERVER_TOKEN', 'EMPTY_TOKEN'), ) # Step 3: Create a sampling client by loading weights from a saved checkpoint. diff --git a/cookbook/client/tinker/self_host/self_cognition.py b/cookbook/client/tinker/self_cognition.py similarity index 96% rename from cookbook/client/tinker/self_host/self_cognition.py rename to cookbook/client/tinker/self_cognition.py index 6d33b6c8..200f7f13 100644 --- a/cookbook/client/tinker/self_host/self_cognition.py +++ b/cookbook/client/tinker/self_cognition.py @@ -24,9 +24,9 @@ from tinker import ServiceClient # The base model to fine-tune / evaluate -base_model = 'Qwen/Qwen3.5-4B' -base_url = 'http://localhost:8000' -api_key = 'EMPTY_API_KEY' +base_model = os.environ.get('TWINKLE_MODEL_ID', 'Qwen/Qwen3.5-4B') +base_url = os.environ.get('TWINKLE_SERVER_URL', 'http://localhost:8000') +api_key = os.environ.get('TWINKLE_SERVER_TOKEN', 'EMPTY_TOKEN') def train(): diff --git a/cookbook/client/tinker/self_host/dpo.py b/cookbook/client/tinker/self_host/dpo.py deleted file mode 100644 index 51474ca0..00000000 --- a/cookbook/client/tinker/self_host/dpo.py +++ /dev/null @@ -1,207 +0,0 @@ -# Tinker-Compatible Client - DPO (Direct Preference Optimization) Training with LoRA -# -# This script demonstrates how to fine-tune a language model using DPO -# through the Tinker-compatible client API. -# -# Training flow per step: -# 1. forward_backward with 'cross_entropy' + disable_lora=True -# → base-model forward pass; LoRA weights are NOT in the computation graph -# so backward accumulates zero LoRA gradients (safe to discard). -# 2. Attach returned per-token ref logps to each datum's loss_fn_inputs. -# 3. forward_backward with 'importance_sampling' -# → server detects ref_logps and switches to DPOLoss + DPOMetric. -# 4. optim_step → update LoRA, DPO metrics returned automatically. -# -# The server must be running first (see server.py and server_config.yaml). - -import os -import numpy as np -import torch -from tqdm import tqdm -from typing import Any, Dict, List - -import swanlab - -from tinker import types -from twinkle import init_tinker_client, get_logger -from twinkle.dataset import Dataset, DatasetMeta, LazyDataset -from twinkle.dataloader import DataLoader -from twinkle.preprocessor import EmojiDPOProcessor -from twinkle.server.common import input_feature_to_datum - -logger = get_logger() - -# Initialize the Tinker client before importing ServiceClient -init_tinker_client() - -from tinker import ServiceClient # noqa: E402 (must follow init_tinker_client) - -# --------------------------------------------------------------------------- -# Configuration -# --------------------------------------------------------------------------- -base_model = 'Qwen/Qwen3.5-4B' -base_url = 'http://localhost:8000' -api_key = 'EMPTY_API_KEY' -dataset_id = 'ms://hjh0119/shareAI-Llama3-DPO-zh-en-emoji' - -batch_size = 4 -learning_rate = 1e-4 -dpo_beta = 0.1 -sft_weight = 1.0 -max_length = 2048 -lora_rank = 8 -system_prompt = 'You are a helpful assistant.' -use_swanlab = False - - -# --------------------------------------------------------------------------- -# Dataset helpers (reused from twinkle/self_host/dpo.py) -# --------------------------------------------------------------------------- - -def create_dpo_dataset(): - """Create DPO dataset with positive/negative format.""" - dataset = LazyDataset(DatasetMeta(dataset_id, data_slice=range(5000))) - dataset.set_template('Qwen3_5Template', model_id=f'ms://{base_model}', max_length=max_length) - dataset.map( - EmojiDPOProcessor, - init_args={'system': system_prompt}, - ) - # EmojiDPOProcessor returns {'positive': InputFeature, 'negative': InputFeature, ...} - # encode handles this format automatically - dataset.encode() - return dataset - - -def prepare_dpo_batch(batch: List[Dict[str, Any]]) -> List[Dict[str, Any]]: - """Reorganise batch into DP-safe interleaved format [pos_1, neg_1, pos_2, neg_2, ...]. - - Args: - batch: List of rows, each with 'positive' and 'negative' InputFeatures. - - Returns: - Interleaved list so each DP worker slice contains complete pairs. - """ - result = [] - for row in batch: - base_fields = {k: v for k, v in row.items() if k not in ('positive', 'negative')} - pos_sample = {**base_fields, **row['positive']} - neg_sample = {**base_fields, **row['negative']} - result.append(pos_sample) - result.append(neg_sample) - return result - - -# --------------------------------------------------------------------------- -# Training -# --------------------------------------------------------------------------- - -def train(): - # Step 0: Initialize SwanLab if enabled - if use_swanlab: - swanlab.login(api_key=os.environ['SWANLAB_API_KEY']) - swanlab.init( - project='twinkle-dpo', - experiment_name='dpo-lora-training', - config={ - 'base_model': base_model, - 'batch_size': batch_size, - 'learning_rate': learning_rate, - 'dpo_beta': dpo_beta, - 'sft_weight': sft_weight, - 'max_length': max_length, - 'lora_rank': lora_rank, - }, - ) - logger.info('SwanLab initialized') - - # Step 1: Prepare dataset & dataloader - logger.info('Loading DPO dataset...') - dataset = create_dpo_dataset() - dataloader = DataLoader(dataset=dataset, batch_size=batch_size) - logger.info(f'Dataset ready: {len(dataloader)} steps per epoch') - - # Step 2: Connect to server and create LoRA training client - service_client = ServiceClient(base_url=base_url, api_key=api_key) - training_client = service_client.create_lora_training_client( - base_model=base_model, - rank=lora_rank, - ) - logger.info(f'LoRA training client created (rank={lora_rank})') - logger.info(f'Starting DPO training: beta={dpo_beta}, lr={learning_rate}') - - # Step 3: Training loop - for step, batch in tqdm(enumerate(dataloader), total=len(dataloader)): - # Normalise numpy / torch tensors to plain Python lists for serialisation - for row in batch: - for key in list(row.keys()): - if isinstance(row[key], np.ndarray): - row[key] = row[key].tolist() - elif isinstance(row[key], torch.Tensor): - row[key] = row[key].cpu().numpy().tolist() - - # Build interleaved [pos, neg, pos, neg, ...] batch - dpo_batch = prepare_dpo_batch(batch) - - # Convert each InputFeature dict to a Tinker Datum - input_datums = [input_feature_to_datum(row) for row in dpo_batch] - - # ----------------------------------------------------------------- - # A. Reference forward pass (base model, disable_lora=True) - # LoRA weights are outside the computation graph → backward - # produces zero LoRA gradients, so this call is safe. - # ----------------------------------------------------------------- - ref_result = training_client.forward( - input_datums, - 'cross_entropy', - loss_fn_config={'disable_lora': True}, - ).result() - - # ----------------------------------------------------------------- - # B. Attach per-token ref logps to each datum's loss_fn_inputs - # ----------------------------------------------------------------- - for datum, ref_out in zip(input_datums, ref_result.loss_fn_outputs): - ref_logprobs_np = np.array(ref_out['logprobs'].tolist(), dtype=np.float32) - datum.loss_fn_inputs['ref_logps'] = types.TensorData.from_numpy(ref_logprobs_np) - - # ----------------------------------------------------------------- - # C. DPO forward_backward - # Server detects ref_logps → sets DPOLoss + DPOMetric automatically. - # Optional DPO hyper-params can be forwarded via loss_fn_config. - # (e.g. beta, sft_weight, not support dpo_loss_type for tinker) - # ----------------------------------------------------------------- - fwdbwd_result = training_client.forward_backward( - input_datums, - 'importance_sampling', - loss_fn_config={ - 'dpo_beta': dpo_beta, - 'dpo_sft_weight': sft_weight, - }, - ).result() - - # ----------------------------------------------------------------- - # D. Optimizer step — DPOMetric is calculated automatically on the - # server and returned inside optim_result.metrics. - # ----------------------------------------------------------------- - optim_result = training_client.optim_step( - types.AdamParams(learning_rate=learning_rate) - ).result() - - logger.info(f'[Step {step}] metrics={optim_result.metrics}') - - # Log metrics to SwanLab - if use_swanlab and optim_result.metrics: - swanlab.log(optim_result.metrics, step=step) - - # Step 4: Save checkpoint - save_result = training_client.save_state('dpo-lora-final').result() - logger.info(f'Saved checkpoint: {save_result.path}') - - # Step 5: (Optional) Upload to ModelScope Hub - # YOUR_USER_NAME = 'your_username' - # hub_model_id = f'{YOUR_USER_NAME}/twinkle-tinker-dpo-lora' - # training_client.publish_checkpoint_from_tinker_path(save_result.path).result() - # logger.info(f'Uploaded checkpoint to hub: {hub_model_id}') - - -if __name__ == '__main__': - train() diff --git a/cookbook/client/tinker/self_host/short_math_grpo.py b/cookbook/client/tinker/short_math_grpo.py similarity index 98% rename from cookbook/client/tinker/self_host/short_math_grpo.py rename to cookbook/client/tinker/short_math_grpo.py index f87fe812..10eabf41 100644 --- a/cookbook/client/tinker/self_host/short_math_grpo.py +++ b/cookbook/client/tinker/short_math_grpo.py @@ -38,7 +38,7 @@ logger = get_logger() # ========== Configuration ========== -BASE_MODEL = 'Qwen/Qwen3.5-4B' +BASE_MODEL = os.environ.get('TWINKLE_MODEL_ID', 'Qwen/Qwen3.5-4B') NUM_GENERATIONS = 4 MAX_NEW_TOKENS = 4096 LEARNING_RATE = 2e-5 @@ -127,8 +127,8 @@ def main(): from tinker import ServiceClient service_client = ServiceClient( - base_url='http://localhost:8000', - api_key='EMPTY_TOKEN' + base_url=os.environ.get('TWINKLE_SERVER_URL', 'http://localhost:8000'), + api_key=os.environ.get('TWINKLE_SERVER_TOKEN', 'EMPTY_TOKEN'), ) logger.info('Creating LoRA training client...') diff --git a/cookbook/client/tinker/self_host/upload_to_hub.py b/cookbook/client/tinker/upload_to_hub.py similarity index 92% rename from cookbook/client/tinker/self_host/upload_to_hub.py rename to cookbook/client/tinker/upload_to_hub.py index 18088e56..da39cc02 100644 --- a/cookbook/client/tinker/self_host/upload_to_hub.py +++ b/cookbook/client/tinker/upload_to_hub.py @@ -18,15 +18,17 @@ dotenv.load_dotenv('.env') +import os + from twinkle import get_logger, init_twinkle_client from twinkle_client.model import MultiLoraTransformersModel logger = get_logger() # ── Configuration ───────────────────────────────────────────────────────────── -base_model = 'Qwen/Qwen3.5-4B' -base_url = 'http://localhost:8000' -api_key = 'EMPTY_TOKEN' # token used for model training / server access +base_model = os.environ.get('TWINKLE_MODEL_ID', 'Qwen/Qwen3.5-4B') +base_url = os.environ.get('TWINKLE_SERVER_URL', 'http://localhost:8000') +api_key = os.environ.get('TWINKLE_SERVER_TOKEN', 'EMPTY_TOKEN') # Checkpoint to upload: the twinkle:// path returned by training_client.save_state(), # e.g. 'twinkle://20260301_142318-Qwen_Qwen3-4B-199d2cdb/weights/my-lora-epoch-0' diff --git a/cookbook/client/twinkle/self_host/dpo.py b/cookbook/client/twinkle/dpo.py similarity index 96% rename from cookbook/client/twinkle/self_host/dpo.py rename to cookbook/client/twinkle/dpo.py index acec9a09..b7fedcd8 100644 --- a/cookbook/client/twinkle/self_host/dpo.py +++ b/cookbook/client/twinkle/dpo.py @@ -27,8 +27,9 @@ logger = get_logger() # Configuration (direct values, not from env) -base_model = 'Qwen/Qwen3.5-4B' -base_url = 'http://localhost:8000' +base_model = os.environ.get('TWINKLE_MODEL_ID', 'Qwen/Qwen3.5-4B') +base_url = os.environ.get('TWINKLE_SERVER_URL', 'http://localhost:8000') +api_key = os.environ.get('TWINKLE_SERVER_TOKEN', 'EMPTY_TOKEN') dataset_id = 'ms://hjh0119/shareAI-Llama3-DPO-zh-en-emoji' batch_size = 4 @@ -44,7 +45,7 @@ # Step 2: Initialize the Twinkle client to communicate with the remote server. # - base_url: the address of the running Twinkle server # - api_key: authentication token (loaded from environment variable) -client = init_twinkle_client(base_url=base_url, api_key=os.environ.get('MODELSCOPE_TOKEN')) +client = init_twinkle_client(base_url=base_url, api_key=api_key) # Step 3: Query the server for existing training runs and their checkpoints. # This is useful for resuming a previous training session. diff --git a/cookbook/client/twinkle/embedding.py b/cookbook/client/twinkle/embedding.py new file mode 100644 index 00000000..304a321f --- /dev/null +++ b/cookbook/client/twinkle/embedding.py @@ -0,0 +1,126 @@ +# Twinkle Client - Embedding (InfoNCE contrastive) Training Example +# +# Client counterpart of the ray-local core-lib example +# ``cookbook/exp/embedding/train_embedding_full_ddp.py`` and the client section +# of ``docs/source_zh/使用指引/Embedding训练.md``. Instead of building a local +# ``TransformersModel``/``MegatronModel`` + Ray, it drives the SAME training flow +# over HTTP through the Twinkle client: +# +# set_processor('InputProcessor') +# -> set_loss('InfonceLoss', ...) +# -> set_optimizer('Adam', ...) +# -> add_metric('EmbeddingMetric', is_training=True) +# -> loop: forward_backward(inputs=mb, task='embedding') + clip_grad_and_step +# -> calculate_metric(is_training=True) +# +# Data format (anchor/positive contrastive pairs): +# inputs: [anchor_0, positive_0, anchor_1, positive_1, ...] +# labels: [ 1, 0, 1, 0, ...] +# (labels==1 marks a new group's anchor; labels==0 marks its positive) +# +# The server must be running first (see server.py / server_config.yaml). Only the +# model service is required (no sampler). Works against both the transformers and +# megatron backends: each single-sequence feature carries ``position_ids`` so the +# megatron mrope model gets valid positions (transformers derives them internally). + +import dotenv + +dotenv.load_dotenv('.env') + +import os +from typing import Any, Dict, List + +from peft import LoraConfig + +from twinkle import get_logger, init_twinkle_client +from twinkle.template import Qwen3_5Template +from twinkle_client.model import MultiLoraTransformersModel + +logger = get_logger() + +# ========== Configuration ========== +BASE_MODEL = os.environ.get('TWINKLE_MODEL_ID', 'Qwen/Qwen3.5-4B') +MODEL_ID = f'ms://{BASE_MODEL}' +ADAPTER_NAME = 'emb_adapter' +TEMPERATURE = 0.07 +LEARNING_RATE = 1e-4 +MAX_STEPS = 10 +MAX_LENGTH = 512 + +# Small self-contained anchor/positive corpus. Each anchor's positive shares its +# meaning, so InfoNCE has a clear contrastive signal to learn from. +PAIRS: List[tuple] = [ + ('What is the capital of France?', 'Paris is the capital city of France.'), + ('How do plants make their food?', 'Plants use photosynthesis to produce food from sunlight.'), + ('What is the boiling point of water?', 'Water boils at 100 degrees Celsius at sea level.'), + ('Who wrote the play Hamlet?', 'William Shakespeare wrote the tragedy Hamlet.'), + ('What is the speed of light?', 'Light travels at about 300,000 kilometers per second.'), + ('What language is spoken in Brazil?', 'The main language spoken in Brazil is Portuguese.'), +] + + +def build_minibatch(tokenizer) -> List[Dict[str, Any]]: + """Build one interleaved [anchor, positive, ...] embedding minibatch. + + Each sample is a plain-dict feature with: + * ``input_ids`` -- tokenized text (plain Python list, JSON-serializable) + * ``attention_mask`` -- all ones + * ``labels`` -- a single scalar: 1 for an anchor, 0 for its positive + * ``position_ids`` -- ``range(len)`` so the megatron mrope model gets valid + positions (transformers derives them internally) + """ + minibatch: List[Dict[str, Any]] = [] + for anchor_text, positive_text in PAIRS: + for text, is_anchor in ((anchor_text, True), (positive_text, False)): + input_ids = tokenizer.encode(text, add_special_tokens=True) + minibatch.append({ + 'input_ids': list(input_ids), + 'attention_mask': [1] * len(input_ids), + 'labels': [1 if is_anchor else 0], + 'position_ids': list(range(len(input_ids))), + }) + return minibatch + + +def train(): + # Step 1: connect to the running Twinkle server. + init_twinkle_client( + base_url=os.environ.get('TWINKLE_SERVER_URL', 'http://localhost:8000'), + api_key=os.environ.get('TWINKLE_SERVER_TOKEN', 'EMPTY_TOKEN'), + ) + + # Step 2: build the client model with a fresh LoRA adapter. + model = MultiLoraTransformersModel(model_id=MODEL_ID) + model.add_adapter_to_model(ADAPTER_NAME, LoraConfig(target_modules='all-linear')) + model.set_template('Qwen3_5Template', model_id=MODEL_ID) + + # Step 3: configure the embedding-training pipeline (ORDER MATTERS). + # Pass class names as strings; the server resolves them to core-lib classes. + model.set_processor('InputProcessor') + model.set_loss('InfonceLoss', temperature=TEMPERATURE, use_batch=True, hard_negatives=None) + # 'Adam' works on both transformers and megatron backends ('AdamW' is + # transformers-only; megatron expects 'Adam'/'default'/'MegatronOptimizer'). + model.set_optimizer('Adam', lr=LEARNING_RATE) + model.add_metric('EmbeddingMetric', is_training=True) + + # Step 4: build a small contrastive minibatch with a local tokenizer. + template = Qwen3_5Template(model_id=MODEL_ID, max_length=MAX_LENGTH, enable_thinking=False) + minibatch = build_minibatch(template.tokenizer) + logger.info(f'Built minibatch: {len(minibatch)} samples ({len(PAIRS)} anchor/positive pairs)') + + # Step 5: training loop. task='embedding' selects the embedding pooling + + # InfoNCE path on the server (TransformersEmbeddingPatch/MegatronEmbeddingPatch + # is applied and rolled back automatically inside the forward). + for step in range(MAX_STEPS): + model.forward_backward(inputs=minibatch, task='embedding') + model.clip_grad_and_step(max_grad_norm=1.0) + metric = model.calculate_metric(is_training=True) + metric = getattr(metric, 'result', metric) + logger.info(f'Step {step}: {metric}') + + logger.info('Embedding training finished. Healthy training shows pos_sim rising ' + 'and neg_sim stable/decreasing.') + + +if __name__ == '__main__': + train() diff --git a/cookbook/client/twinkle/modelscope/dpo.py b/cookbook/client/twinkle/modelscope/dpo.py deleted file mode 100644 index e9451e31..00000000 --- a/cookbook/client/twinkle/modelscope/dpo.py +++ /dev/null @@ -1,204 +0,0 @@ -# Twinkle Client - DPO (Direct Preference Optimization) Training with LoRA -# -# This script demonstrates how to fine-tune a language model using DPO -# through the Twinkle client-server architecture. -# The server must be running first (see server.py and server_config.yaml). - -# Step 1: Load environment variables from a .env file (e.g., API tokens) -import dotenv -import os -from typing import Any, Dict, List - -dotenv.load_dotenv('.env') -import numpy as np -import torch -from peft import LoraConfig - -from twinkle import get_logger -from twinkle.dataset import Dataset, DatasetMeta -from twinkle_client import init_twinkle_client -from twinkle.dataloader import DataLoader -from twinkle_client.model import MultiLoraTransformersModel -from twinkle.preprocessor import EmojiDPOProcessor - -logger = get_logger() - -# Configuration (direct values, not from env) -base_model = 'Qwen/Qwen3.6-27B' -base_url = 'http://www.modelscope.cn/twinkle' -dataset_id = 'ms://hjh0119/shareAI-Llama3-DPO-zh-en-emoji' - -batch_size = 4 -gradient_accumulation_steps = 2 -learning_rate = 1e-4 -dpo_beta = 0.1 -sft_weight = 1.0 -loss_type = 'sigmoid' -max_length = 2048 -adapter_name = 'default' -system_prompt = 'You are a helpful assistant.' - -# Step 2: Initialize the Twinkle client to communicate with the remote server. -# - base_url: the address of the running Twinkle server -# - api_key: authentication token (loaded from environment variable) -client = init_twinkle_client(base_url=base_url, api_key=os.environ.get('MODELSCOPE_TOKEN')) - -# Step 3: Query the server for existing training runs and their checkpoints. -# This is useful for resuming a previous training session. -runs = client.list_training_runs() - -resume_path = None -for run in runs: - logger.info(run.model_dump_json(indent=2)) - # List all saved checkpoints for this training run - checkpoints = client.list_checkpoints(run.training_run_id) - - for checkpoint in checkpoints: - logger.info(checkpoint.model_dump_json(indent=2)) - # Uncomment the line below to resume from a specific checkpoint: - # resume_path = checkpoint.twinkle_path - - -def create_dpo_dataset(): - """Create DPO dataset with positive/negative format.""" - dataset = Dataset(DatasetMeta(dataset_id, data_slice=range(100))) - dataset.set_template('Qwen3_5Template', model_id=f'ms://{base_model}', max_length=max_length) - dataset.map( - EmojiDPOProcessor, - init_args={ - 'system': system_prompt, - } - ) - # DPO preprocessor returns {'positive': [...], 'negative': [...]} - # batch_encode handles this format automatically - dataset.encode() - return dataset - - -def prepare_dpo_batch(batch: List[Dict[str, Any]]) -> List[Dict[str, Any]]: - """Prepare DPO batch: reorganize batch for training with DP-safe interleaving. - - Args: - batch: List of rows, each with 'positive' and 'negative' InputFeatures - and other fields (question, etc.) - - Returns: - List interleaved as [pos_1, neg_1, pos_2, neg_2, ...] to ensure each DP - worker gets complete positive/negative pairs after slicing. - Each item contains all original fields plus the InputFeature fields. - """ - result = [] - - for row in batch: - # Get base fields (excluding positive/negative) - base_fields = {k: v for k, v in row.items() if k not in ('positive', 'negative')} - - # Positive sample: merge base fields with positive InputFeature - pos_sample = {**base_fields, **row['positive']} - # Negative sample: merge base fields with negative InputFeature - neg_sample = {**base_fields, **row['negative']} - - # Interleave: [pos, neg] per pair for DP-safe slicing - result.append(pos_sample) - result.append(neg_sample) - - return result - - -def train(): - # Step 4: Prepare the dataset - - # Load the DPO dataset from ModelScope - dataset = create_dpo_dataset() - - # Wrap the dataset into a DataLoader that yields batches - dataloader = DataLoader(dataset=dataset, batch_size=batch_size) - - # Step 5: Configure the model - - # Create a multi-LoRA Transformers model pointing to the base model on ModelScope - model = MultiLoraTransformersModel(model_id=f'ms://{base_model}') - - # Define LoRA configuration: apply low-rank adapters to all linear layers - lora_config = LoraConfig( - target_modules='all-linear', - r=8, - lora_alpha=32, - lora_dropout=0.05, - ) - - # Attach the LoRA adapter named 'default' to the model. - # gradient_accumulation_steps means gradients are accumulated over micro-batches - # before an optimizer step, effectively increasing the batch size. - model.add_adapter_to_model(adapter_name, lora_config, gradient_accumulation_steps=gradient_accumulation_steps) - - # Set the same chat template used during data preprocessing - model.set_template('Qwen3_5Template') - - # Set the input processor (pads sequences on the right side) - model.set_processor('InputProcessor', padding_side='right') - - # Use DPO loss for preference optimization - model.set_loss('DPOLoss', beta=dpo_beta, loss_type=loss_type, reference_free=False, sft_weight=sft_weight) - - # Add DPO metric for logging - model.add_metric('DPOMetric', beta=dpo_beta) - - # Use Adam optimizer with a learning rate of 1e-4 - model.set_optimizer('Adam', lr=learning_rate) - - # Step 6: Optionally resume from a previous checkpoint - if resume_path: - logger.info(f'Resuming training from {resume_path}') - model.load(resume_path, load_optimizer=True) - - # Step 7: Run the training loop - logger.info(model.get_train_configs().model_dump()) - - optim_step = 0 - max_steps = len(dataloader) - logger.info(f'Starting LoRA DPO training: loss_type={loss_type}, beta={dpo_beta}, lr={learning_rate}') - logger.info(f'Using base model (disable_lora=True) as reference model') - - for batch in dataloader: - # batch is List[Dict] with 'positive' and 'negative' keys - # Convert numpy/torch tensors to lists for serialization - for row in batch: - for key in row: - if isinstance(row[key], np.ndarray): - row[key] = row[key].tolist() - elif isinstance(row[key], torch.Tensor): - row[key] = row[key].cpu().numpy().tolist() - - dpo_batch = prepare_dpo_batch(batch) - - # Get reference outputs using base model (without LoRA adapter) - # disable_lora=True tells the model to skip LoRA and use base weights - ref_outputs = model.forward_only(inputs=dpo_batch, disable_lora=True) - model.forward_backward(inputs=dpo_batch, ref_outputs=ref_outputs.result) - model.clip_grad_and_step() - - optim_step += 1 - - # Logging - if optim_step % gradient_accumulation_steps == 0: - metrics = model.calculate_metric(is_training=True) - logger.info(f'[Step {optim_step // gradient_accumulation_steps}/{max_steps}] {metrics}') - - # Step 8: Save the trained checkpoint - twinkle_path = model.save(name='dpo-lora-final', save_optimizer=True) - logger.info(f'Saved checkpoint: {twinkle_path}') - - # Step 9: Upload the checkpoint to ModelScope Hub - # YOUR_USER_NAME = "your_username" - # hub_model_id = f'{YOUR_USER_NAME}/twinkle-dpo-lora' - # model.upload_to_hub( - # checkpoint_dir=twinkle_path, - # hub_model_id=hub_model_id, - # async_upload=False - # ) - # logger.info(f"Uploaded checkpoint to hub: {hub_model_id}") - - -if __name__ == '__main__': - train() diff --git a/cookbook/client/twinkle/modelscope/multi_modal.py b/cookbook/client/twinkle/modelscope/multi_modal.py deleted file mode 100644 index 106d85d4..00000000 --- a/cookbook/client/twinkle/modelscope/multi_modal.py +++ /dev/null @@ -1,168 +0,0 @@ -# Twinkle Client - Transformers LoRA Training Example -# -# This script demonstrates how to fine-tune a language model using LoRA -# (Low-Rank Adaptation) through the Twinkle client-server architecture. -# The server must be running first (see server.py and server_config.yaml). - -# Step 1: Load environment variables from a .env file (e.g., API tokens) -import dotenv -import os -from twinkle.data_format import Trajectory, Message -from twinkle.preprocessor import Preprocessor - -dotenv.load_dotenv('.env') -import numpy as np -import torch -from peft import LoraConfig - -from twinkle import get_logger -from twinkle.dataset import DatasetMeta -from twinkle_client import init_twinkle_client -from twinkle.dataloader import DataLoader -from twinkle.dataset import LazyDataset -from twinkle_client.model import MultiLoraTransformersModel - -logger = get_logger() - -base_model = 'Qwen/Qwen3.6-27B' -base_url = 'http://www.modelscope.cn/twinkle' - -# Step 2: Initialize the Twinkle client to communicate with the remote server. -# - base_url: the address of the running Twinkle server -# - api_key: authentication token (loaded from environment variable) -client = init_twinkle_client(base_url=base_url, api_key=os.environ.get('MODELSCOPE_TOKEN')) - -# Step 3: Query the server for existing training runs and their checkpoints. -# This is useful for resuming a previous training session. -runs = client.list_training_runs() - -resume_path = None -for run in runs: - logger.info(run.model_dump_json(indent=2)) - # List all saved checkpoints for this training run - checkpoints = client.list_checkpoints(run.training_run_id) - - for checkpoint in checkpoints: - logger.info(checkpoint.model_dump_json(indent=2)) - # Uncomment the line below to resume from a specific checkpoint: - # resume_path = checkpoint.twinkle_path - - -class LatexOCRProcessor(Preprocessor): - - def __call__(self, rows): - rows = self.map_col_to_row(rows) - rows = [self.preprocess(row) for row in rows] - rows = self.map_row_to_col(rows) - return rows - - def preprocess(self, row) -> Trajectory: - return Trajectory( - messages=[ - Message(role='user', content='Using LaTeX to perform OCR on the image.', images=[row['image']]), - Message(role='assistant', content=row['text']), - ] - ) - - -def train(): - # Step 4: Prepare the dataset - - # Load the latex dataset from ModelScope - dataset = LazyDataset(dataset_meta=DatasetMeta('ms://AI-ModelScope/LaTeX_OCR', data_slice=range(500))) - - # Apply a chat template so the data matches the model's expected input format - dataset.set_template('Qwen3_5Template', model_id=f'ms://{base_model}', max_length=512) - - # Replace placeholder names in the dataset with custom model/author names - dataset.map(LatexOCRProcessor) - - # Tokenize and encode the dataset into model-ready input features - dataset.encode(batched=True) - - # Wrap the dataset into a DataLoader that yields batches of size 4 - dataloader = DataLoader(dataset=dataset, batch_size=4) - - # Step 5: Configure the model - - # Create a multi-LoRA Transformers model pointing to the base model on ModelScope - model = MultiLoraTransformersModel(model_id=f'ms://{base_model}') - - # Define LoRA configuration: apply low-rank adapters to all linear layers - lora_config = LoraConfig(target_modules='all-linear') - - # Attach the LoRA adapter named 'default' to the model. - # gradient_accumulation_steps=2 means gradients are accumulated over 2 micro-batches - # before an optimizer step, effectively doubling the batch size. - model.add_adapter_to_model('default', lora_config, gradient_accumulation_steps=2) - - # Set the same chat template used during data preprocessing - model.set_template('Qwen3_5Template') - - # Set the input processor (pads sequences on the right side) - model.set_processor('InputProcessor', padding_side='right') - - # Use cross-entropy loss for language modeling - model.set_loss('CrossEntropyLoss') - - # Use Adam optimizer with a learning rate of 1e-4 (Only support Adam optimizer if server use megatron) - model.set_optimizer('Adam', lr=1e-4) - - # Use a linear learning rate scheduler (Do not support LR scheduler if server use megatron) - # model.set_lr_scheduler('LinearLR') - - # Step 6: Optionally resume from a previous checkpoint - if resume_path: - logger.info(f'Resuming training from {resume_path}') - model.load(resume_path, load_optimizer=True) - - # Step 7: Run the training loop - logger.info(model.get_train_configs().model_dump()) - - for epoch in range(3): - logger.info(f'Starting epoch {epoch}') - for step, batch in enumerate(dataloader): - for sample in batch: - for key in sample: - if isinstance(sample[key], np.ndarray): - sample[key] = sample[key].tolist() - elif isinstance(sample[key], torch.Tensor): - sample[key] = sample[key].cpu().numpy().tolist() - # Forward pass + backward pass (computes gradients) - model.forward_backward(inputs=batch) - - # Step - model.clip_grad_and_step() - # Equal to the following steps: - # # Clip gradients to prevent exploding gradients (max norm = 1.0) - # model.clip_grad_norm(1.0) - # # Perform one optimizer step (update model weights) - # model.step() - # # Reset gradients to zero for the next iteration - # model.zero_grad() - # # Advance the learning rate scheduler by one step - # model.lr_step() - - # Log the loss every 2 steps (aligned with gradient accumulation) - if step % 2 == 0: - # Print metric - metric = model.calculate_metric(is_training=True) - logger.info(f'Current is step {step} of {len(dataloader)}, metric: {metric.result}') - - # Step 8: Save the trained checkpoint - twinkle_path = model.save(name=f'twinkle-epoch-{epoch}', save_optimizer=True) - logger.info(f'Saved checkpoint: {twinkle_path}') - - # Step 9: Upload the checkpoint to ModelScope Hub - # YOUR_USER_NAME = "your_username" - # hub_model_id = f'{YOUR_USER_NAME}/twinkle-multi-modal' - # model.upload_to_hub( - # checkpoint_dir=twinkle_path, - # hub_model_id=hub_model_id, - # async_upload=False - # ) - # logger.info(f"Uploaded checkpoint to hub: {hub_model_id}") - - -if __name__ == '__main__': - train() diff --git a/cookbook/client/twinkle/modelscope/self_cognition.py b/cookbook/client/twinkle/modelscope/self_cognition.py deleted file mode 100644 index 5acd8a9a..00000000 --- a/cookbook/client/twinkle/modelscope/self_cognition.py +++ /dev/null @@ -1,142 +0,0 @@ -# Twinkle Client - Transformers LoRA Training Example -# -# This script demonstrates how to fine-tune a language model using LoRA -# (Low-Rank Adaptation) through the Twinkle client-server architecture. -# The server must be running first (see server.py and server_config.yaml). - -# Step 1: Load environment variables from a .env file (e.g., API tokens) -import dotenv - -dotenv.load_dotenv('.env') - -import os -from peft import LoraConfig - -from twinkle import get_logger -from twinkle.dataset import DatasetMeta -from twinkle import init_twinkle_client -from twinkle.dataloader import DataLoader -from twinkle.dataset import Dataset -from twinkle_client.model import MultiLoraTransformersModel - -logger = get_logger() - -base_model = 'Qwen/Qwen3.6-27B' -base_url = 'http://www.modelscope.cn/twinkle' - -# Step 2: Initialize the Twinkle client to communicate with the remote server. -# - base_url: the address of the running Twinkle server -# - api_key: authentication token (loaded from environment variable) -client = init_twinkle_client(base_url=base_url, api_key=os.environ.get('MODELSCOPE_TOKEN')) - -# Step 3: Query the server for existing training runs and their checkpoints. -# This is useful for resuming a previous training session. -runs = client.list_training_runs() - -resume_path = None -for run in runs: - logger.info(run.model_dump_json(indent=2)) - # List all saved checkpoints for this training run - checkpoints = client.list_checkpoints(run.training_run_id) - - for checkpoint in checkpoints: - logger.info(checkpoint.model_dump_json(indent=2)) - # Uncomment the line below to resume from a specific checkpoint: - # resume_path = checkpoint.twinkle_path - - -def train(): - # Step 4: Prepare the dataset - - # Load the self-cognition dataset from ModelScope - dataset = Dataset(dataset_meta=DatasetMeta('ms://swift/self-cognition', data_slice=range(500))) - - # Apply a chat template so the data matches the model's expected input format - dataset.set_template('Qwen3_5Template', model_id=f'ms://{base_model}', max_length=512) - - # Replace placeholder names in the dataset with custom model/author names - dataset.map('SelfCognitionProcessor', init_args={'model_name': 'twinkle模型', 'model_author': 'ModelScope社区'}) - - # Tokenize and encode the dataset into model-ready input features - dataset.encode(batched=True) - - # Wrap the dataset into a DataLoader that yields batches of size 4 - dataloader = DataLoader(dataset=dataset, batch_size=4) - - # Step 5: Configure the model - - # Create a multi-LoRA Transformers model pointing to the base model on ModelScope - model = MultiLoraTransformersModel(model_id=f'ms://{base_model}') - - # Define LoRA configuration: apply low-rank adapters to all linear layers - lora_config = LoraConfig(target_modules='all-linear') - - # Attach the LoRA adapter named 'default' to the model. - # gradient_accumulation_steps=2 means gradients are accumulated over 2 micro-batches - # before an optimizer step, effectively doubling the batch size. - model.add_adapter_to_model('default', lora_config, gradient_accumulation_steps=2) - - # Set the same chat template used during data preprocessing - model.set_template('Qwen3_5Template') - - # Set the input processor (pads sequences on the right side) - model.set_processor('InputProcessor', padding_side='right') - - # Use cross-entropy loss for language modeling - model.set_loss('CrossEntropyLoss') - - # Use Adam optimizer with a learning rate of 1e-4 (Only support Adam optimizer if server use megatron) - model.set_optimizer('Adam', lr=1e-4) - - # Use a linear learning rate scheduler (Do not support LR scheduler if server use megatron) - # model.set_lr_scheduler('LinearLR') - - # Step 6: Optionally resume from a previous checkpoint - if resume_path: - logger.info(f'Resuming training from {resume_path}') - model.load(resume_path, load_optimizer=True) - - # Step 7: Run the training loop - logger.info(model.get_train_configs().model_dump()) - - for epoch in range(3): - logger.info(f'Starting epoch {epoch}') - for step, batch in enumerate(dataloader): - # Forward pass + backward pass (computes gradients) - model.forward_backward(inputs=batch) - - # Step - model.clip_grad_and_step() - # Equal to the following steps: - # # Clip gradients to prevent exploding gradients (max norm = 1.0) - # model.clip_grad_norm(1.0) - # # Perform one optimizer step (update model weights) - # model.step() - # # Reset gradients to zero for the next iteration - # model.zero_grad() - # # Advance the learning rate scheduler by one step - # model.lr_step() - - # Log the loss every 2 steps (aligned with gradient accumulation) - if step % 2 == 0: - # Print metric - metric = model.calculate_metric(is_training=True) - logger.info(f'Current is step {step} of {len(dataloader)}, metric: {metric.result}') - - # Step 8: Save the trained checkpoint - twinkle_path = model.save(name=f'twinkle-epoch-{epoch}', save_optimizer=True) - logger.info(f'Saved checkpoint: {twinkle_path}') - - # Step 9: Upload the checkpoint to ModelScope Hub - # YOUR_USER_NAME = "your_username" - # hub_model_id = f'{YOUR_USER_NAME}/twinkle-self-cognition' - # model.upload_to_hub( - # checkpoint_dir=twinkle_path, - # hub_model_id=hub_model_id, - # async_upload=False - # ) - # logger.info(f"Uploaded checkpoint to hub: {hub_model_id}") - - -if __name__ == '__main__': - train() diff --git a/cookbook/client/twinkle/self_host/multi_modal.py b/cookbook/client/twinkle/multi_modal.py similarity index 95% rename from cookbook/client/twinkle/self_host/multi_modal.py rename to cookbook/client/twinkle/multi_modal.py index f1908220..3a5a45b6 100644 --- a/cookbook/client/twinkle/self_host/multi_modal.py +++ b/cookbook/client/twinkle/multi_modal.py @@ -24,13 +24,14 @@ logger = get_logger() -base_model = 'Qwen/Qwen3.5-4B' -base_url = 'http://localhost:8000' +base_model = os.environ.get('TWINKLE_MODEL_ID', 'Qwen/Qwen3.5-4B') +base_url = os.environ.get('TWINKLE_SERVER_URL', 'http://localhost:8000') +api_key = os.environ.get('TWINKLE_SERVER_TOKEN', 'EMPTY_TOKEN') # Step 2: Initialize the Twinkle client to communicate with the remote server. # - base_url: the address of the running Twinkle server # - api_key: authentication token (loaded from environment variable) -client = init_twinkle_client(base_url=base_url, api_key=os.environ.get('MODELSCOPE_TOKEN')) +client = init_twinkle_client(base_url=base_url, api_key=api_key) # Step 3: Query the server for existing training runs and their checkpoints. # This is useful for resuming a previous training session. diff --git a/cookbook/client/twinkle/multi_turn_rollout.py b/cookbook/client/twinkle/multi_turn_rollout.py new file mode 100644 index 00000000..199dca38 --- /dev/null +++ b/cookbook/client/twinkle/multi_turn_rollout.py @@ -0,0 +1,221 @@ +# Twinkle Client - Multi-turn agentic rollout Example +# +# Client counterpart of the ray-local core-lib example +# ``cookbook/rl/multi_turn/multi_turn_grpo.py``. It drives the same multi-turn +# "sample -> environment action -> observation -> sample" loop over HTTP with +# ``twinkle_client.rollout.ClientMultiTurnRollout`` and an embedded OpenEnv +# environment for each trajectory. +# +# Differences from the ray-local example (by design): +# * The client does not need a distributed Ray EnvPool. It creates one local +# ``OpenEnv`` instance per trajectory and exposes it through ``EnvTool``. +# * ``ClientMultiTurnRollout`` samples with ``num_samples=1``; the GRPO batch uses +# ``BATCH_SIZE * NUM_GENERATIONS`` independent environment trajectories, matching +# the ray-local example. +# +# The server must be running first (see server.py / server_config.yaml). Both the +# model and sampler services must be configured. +# +# NOTE on weight sync: ``ClientMultiTurnRollout`` samples with the sampler's +# currently loaded server-side weights. A full RL loop would periodically +# ``model.save(is_sampler=True)`` and point the sampler at the saved adapter; that +# sync is intentionally omitted here to keep the rollout example focused. + +import os +from typing import Any, Dict, List, Tuple + +import dotenv +from peft import LoraConfig + +from twinkle import get_logger, init_twinkle_client +from twinkle.advantage import GRPOAdvantage +from twinkle.data_format import SamplingParams +from twinkle.template import Qwen3_5Template +from twinkle_agentic.envs import EnvTool, OpenEnv +from twinkle_agentic.tools.tool_manager import ToolManager +from twinkle_client.model import MultiLoraTransformersModel +from twinkle_client.rollout import ClientMultiTurnRollout +from twinkle_client.sampler import vLLMSampler + +dotenv.load_dotenv('.env') + +logger = get_logger() + +# ========== Configuration ========== +BASE_MODEL = os.environ.get('TWINKLE_MODEL_ID', 'Qwen/Qwen3.5-4B') +MODEL_ID = f'ms://{BASE_MODEL}' +ADAPTER_NAME = 'default' +NUM_GENERATIONS = 2 # GRPO group size (rollout runs num_samples=1 per trajectory) +BATCH_SIZE = 2 +MAX_NEW_TOKENS = 512 +MAX_TURNS = 4 +LEARNING_RATE = 1e-5 +MAX_STEPS = 3 +OPENENV_NAME = os.environ.get('OPENENV_NAME', 'openspiel_env') +OPENENV_GAME = os.environ.get('OPENENV_GAME', 'blackjack') + +BLACKJACK_TOOL_SCHEMA = [ + { + 'type': 'function', + 'function': { + 'name': 'play', + 'description': 'Take an action in the blackjack game.', + 'parameters': { + 'type': 'object', + 'properties': { + 'action': { + 'type': 'string', + 'enum': ['hit', 'stand'], + 'description': 'Choose "hit" to draw a card or "stand" to keep the current hand.', + }, + }, + 'required': ['action'], + }, + }, + }, +] + +BLACKJACK_ACTION_MAP = {'hit': 0, 'stand': 1} + +SYSTEM_PROMPT = """You are a skilled blackjack player. You will be told your current hand and the dealer's visible card. + +Your goal is to win the game by getting as close to 21 as possible without going over. + +Use the `play` tool to choose either `hit` or `stand`. Reason briefly before each action. Once the environment reports that the game is over, give a short final answer without calling another tool.""" + + +def blackjack_action_mapper(tool_name: str, arguments: Dict[str, Any]) -> Dict[str, Any]: + """Convert play(action=...) into the action format expected by OpenSpiel.""" + action = str(arguments.get('action', 'stand')).lower().strip() + return { + 'action_id': BLACKJACK_ACTION_MAP.get(action, BLACKJACK_ACTION_MAP['stand']), + 'game_name': OPENENV_GAME, + } + + +def create_env_tool(env: OpenEnv) -> EnvTool: + """Expose the blackjack OpenEnv action through ToolManager.""" + function = BLACKJACK_TOOL_SCHEMA[0]['function'] + return EnvTool( + env=env, + tool_name=function['name'], + description=function['description'], + parameters=function['parameters'], + ) + + +def prepare_trajectories( + n_trajectories: int, +) -> Tuple[List[Dict[str, Any]], List[ToolManager], List[List[EnvTool]], List[OpenEnv]]: + """Create and reset one independent OpenEnv instance per trajectory.""" + trajectories = [] + tool_managers = [] + env_tools_list = [] + environments = [] + + try: + for _ in range(n_trajectories): + env = OpenEnv( + env_name=OPENENV_NAME, + env_kwargs={'game_name': OPENENV_GAME}, + action_mapper=blackjack_action_mapper, + ) + environments.append(env) + initial_observation = env.reset().observation + + env_tools = [create_env_tool(env)] + tool_manager = ToolManager(env_tools) + trajectories.append({ + 'messages': [ + {'role': 'system', 'content': SYSTEM_PROMPT}, + {'role': 'user', 'content': initial_observation}, + ], + 'tools': tool_manager.tool_infos(), + }) + tool_managers.append(tool_manager) + env_tools_list.append(env_tools) + except Exception: + for env in environments: + env.close() + raise + + return trajectories, tool_managers, env_tools_list, environments + + +def extract_rewards(env_tools_list: List[List[EnvTool]]) -> List[float]: + """Read the terminal reward produced by each OpenEnv episode.""" + return [env_tools[0].episode_reward if env_tools else 0.0 for env_tools in env_tools_list] + + +def train(): + # Step 1: connect to the running Twinkle server. + init_twinkle_client( + base_url=os.environ.get('TWINKLE_SERVER_URL', 'http://localhost:8000'), + api_key=os.environ.get('TWINKLE_SERVER_TOKEN', 'EMPTY_TOKEN'), + ) + + # Step 2: training model (GRPO), mirroring the ray-local example's config. + model = MultiLoraTransformersModel(model_id=MODEL_ID) + model.add_adapter_to_model(ADAPTER_NAME, LoraConfig(target_modules='all-linear', r=16, lora_alpha=32)) + model.set_loss('GRPOLoss', epsilon=0.2) + model.set_optimizer('Adam', lr=LEARNING_RATE) + model.set_processor('InputProcessor') + model.set_template('Qwen3_5Template', model_id=MODEL_ID, enable_thinking=False) + + # Step 3: client sampler (HTTP). + sampler = vLLMSampler(model_id=MODEL_ID) + sampler.set_template('Qwen3_5Template', model_id=MODEL_ID, enable_thinking=False) + + # Step 4: multi-turn rollout. Each call receives trajectory-bound ToolManagers + # backed by independent OpenEnv instances. + rollout_template = Qwen3_5Template(model_id=MODEL_ID, max_length=8192, enable_thinking=False) + rollout_template.truncation_strategy = 'delete' + + rollout = ClientMultiTurnRollout( + sampler=sampler, + template=rollout_template, + sampling_params=SamplingParams( + max_tokens=MAX_NEW_TOKENS, num_samples=1, logprobs=1, temperature=1.0, top_p=0.95), + max_turns=MAX_TURNS, + ) + + advantage_fn = GRPOAdvantage() + + for step in range(MAX_STEPS): + # 1. Reset independent OpenEnv episodes and run the batched rollout. + n_trajectories = BATCH_SIZE * NUM_GENERATIONS + trajectories, tool_managers, env_tools_list, environments = prepare_trajectories(n_trajectories) + try: + rolled = rollout(trajectories, tool_manager=tool_managers) + rewards = extract_rewards(env_tools_list) + finally: + for env in environments: + env.close() + + # 2. Read back per-trajectory logprobs (top-1 chosen-token logprob). + all_inputs: List[Dict[str, Any]] = [] + all_old_logps: List[List[float]] = [] + for traj in rolled: + logprobs = traj.get('logprobs') or [] + all_old_logps.append([lp[0][1] for lp in logprobs]) + all_inputs.append(traj) + + avg_turns = sum(t.get('turns') or 0 for t in rolled) / len(rolled) + logger.info(f'[Step {step}] avg_reward={sum(rewards) / len(rewards):.3f} avg_turns={avg_turns:.1f}') + + # 3. GRPO advantages (group-relative within NUM_GENERATIONS). + advantages = advantage_fn(rewards, num_generations=NUM_GENERATIONS, scale='group').tolist() + if all(abs(a) < 1e-8 for a in advantages): + logger.info(f'[Step {step}] all advantages zero, skipping update') + continue + + # 4. Policy update over the rolled-out trajectories. + model.forward_backward(inputs=all_inputs, advantages=advantages, old_logps=all_old_logps) + model.clip_grad_and_step() + logger.info(f'[Step {step}] {model.calculate_metric(is_training=True).result}') + + logger.info('Multi-turn rollout example finished.') + + +if __name__ == '__main__': + train() diff --git a/cookbook/client/twinkle/self_host/sample.py b/cookbook/client/twinkle/sample.py similarity index 93% rename from cookbook/client/twinkle/self_host/sample.py rename to cookbook/client/twinkle/sample.py index f7925d4f..5bbfd424 100644 --- a/cookbook/client/twinkle/self_host/sample.py +++ b/cookbook/client/twinkle/sample.py @@ -22,7 +22,7 @@ logger = get_logger() -MODEL_ID = 'Qwen/Qwen3.5-4B' +MODEL_ID = os.environ.get('TWINKLE_MODEL_ID', 'Qwen/Qwen3.5-4B') # Optional: adapter URI for LoRA inference # This can be a twinkle:// path from a training run checkpoint @@ -34,8 +34,8 @@ def sample(): # Step 2: Initialize the Twinkle client to communicate with the remote server. client = init_twinkle_client( - base_url='http://127.0.0.1:8000', - api_key='EMPTY_API_KEY', + base_url=os.environ.get('TWINKLE_SERVER_URL', 'http://localhost:8000'), + api_key=os.environ.get('TWINKLE_SERVER_TOKEN', 'EMPTY_TOKEN'), ) # Step 3: Create the sampler client pointing to the model on the server diff --git a/cookbook/client/twinkle/self_host/self_cognition.py b/cookbook/client/twinkle/self_cognition.py similarity index 96% rename from cookbook/client/twinkle/self_host/self_cognition.py rename to cookbook/client/twinkle/self_cognition.py index 997a7732..f5250b29 100644 --- a/cookbook/client/twinkle/self_host/self_cognition.py +++ b/cookbook/client/twinkle/self_cognition.py @@ -21,9 +21,9 @@ logger = get_logger() -base_model = 'Qwen/Qwen3.5-4B' -base_url = 'http://localhost:8000' -api_key = 'EMPTY_API_KEY' +base_model = os.environ.get('TWINKLE_MODEL_ID', 'Qwen/Qwen3.5-4B') +base_url = os.environ.get('TWINKLE_SERVER_URL', 'http://localhost:8000') +api_key = os.environ.get('TWINKLE_SERVER_TOKEN', 'EMPTY_TOKEN') save_dir = '/tmp/twinkle_sft_output' diff --git a/cookbook/client/twinkle/self_host/short_math_grpo.py b/cookbook/client/twinkle/short_math_grpo.py similarity index 97% rename from cookbook/client/twinkle/self_host/short_math_grpo.py rename to cookbook/client/twinkle/short_math_grpo.py index 871d4599..1e90d38c 100644 --- a/cookbook/client/twinkle/self_host/short_math_grpo.py +++ b/cookbook/client/twinkle/short_math_grpo.py @@ -80,7 +80,8 @@ def __call__(self, trajectories: List[Dict[str, Any]], **kwargs) -> List[float]: return rewards # ========== Configuration ========== -MODEL_ID = 'ms://Qwen/Qwen3.5-4B' +BASE_MODEL = os.environ.get('TWINKLE_MODEL_ID', 'Qwen/Qwen3.5-4B') +MODEL_ID = f'ms://{BASE_MODEL}' NUM_GENERATIONS = 4 MAX_NEW_TOKENS = 1024 LEARNING_RATE = 2e-5 @@ -141,8 +142,8 @@ def train(): # Step 1: Initialize the Twinkle client client = init_twinkle_client( - base_url='http://127.0.0.1:8000', - api_key='EMPTY_TOKEN', + base_url=os.environ.get('TWINKLE_SERVER_URL', 'http://localhost:8000'), + api_key=os.environ.get('TWINKLE_SERVER_TOKEN', 'EMPTY_TOKEN'), ) # Step 2: Prepare dataset and dataloader diff --git a/cookbook/client/twinkle/self_host/upload_to_hub.py b/cookbook/client/twinkle/upload_to_hub.py similarity index 92% rename from cookbook/client/twinkle/self_host/upload_to_hub.py rename to cookbook/client/twinkle/upload_to_hub.py index 0505d27d..2c780e25 100644 --- a/cookbook/client/twinkle/self_host/upload_to_hub.py +++ b/cookbook/client/twinkle/upload_to_hub.py @@ -18,15 +18,17 @@ dotenv.load_dotenv('.env') +import os + from twinkle import get_logger, init_twinkle_client from twinkle_client.model import MultiLoraTransformersModel logger = get_logger() # ── Configuration ───────────────────────────────────────────────────────────── -base_model = 'Qwen/Qwen3.5-4B' -base_url = 'http://localhost:8000' -api_key = 'EMPTY_TOKEN' # token used for model training / server access +base_model = os.environ.get('TWINKLE_MODEL_ID', 'Qwen/Qwen3.5-4B') +base_url = os.environ.get('TWINKLE_SERVER_URL', 'http://localhost:8000') +api_key = os.environ.get('TWINKLE_SERVER_TOKEN', 'EMPTY_TOKEN') # Checkpoint to upload: either a twinkle:// path or a local directory path. # Example twinkle:// path (from model.save()): diff --git a/docs/source_en/Usage Guide/Embedding-Training.md b/docs/source_en/Usage Guide/Embedding-Training.md index d86769c3..d0517e18 100644 --- a/docs/source_en/Usage Guide/Embedding-Training.md +++ b/docs/source_en/Usage Guide/Embedding-Training.md @@ -106,6 +106,62 @@ for step, batch in enumerate(dataloader): --- +## Train via Twinkle Client + +Besides using the bare-library `TransformersModel` directly, you can also drive embedding training on the server over HTTP through `twinkle_client`'s `MultiLoraTransformersModel`. Usage is identical to the bare library — you just call the existing methods in the correct order. + +### Call Order + +Client-side embedding training follows the same call order as the bare library: + +``` +set_processor('InputProcessor') + -> set_loss('InfonceLoss', ...) + -> add_metric('EmbeddingMetric', is_training=True) + -> loop: forward_backward(inputs=mb, task='embedding') + clip_grad_and_step(...) + -> calculate_metric(is_training=True) +``` + +The difference from the bare library is that the client passes **class-name strings** (e.g. `'InfonceLoss'`, `'EmbeddingMetric'`, `'InputProcessor'`), which the server resolves to the corresponding core-lib classes. + +### Example + +```python +from peft import LoraConfig +from twinkle_client import init_twinkle_client +from twinkle_client.model import MultiLoraTransformersModel + +# --- Connect to the running Twinkle server --- +init_twinkle_client(base_url='http://127.0.0.1:8000', api_key='EMPTY_TOKEN') + +# --- Build the client model with a LoRA adapter for embedding --- +model = MultiLoraTransformersModel(model_id='ms://Qwen/Qwen3.5-4B') +model.add_adapter_to_model('emb_adapter', LoraConfig(target_modules='all-linear')) +model.set_template('Qwen3_5Template') + +# --- Configure the embedding-training pipeline (order matters) --- +# NOTE: pass class names as strings; the server resolves them to core-lib classes. +model.set_processor('InputProcessor') +model.set_loss('InfonceLoss', temperature=0.07, use_batch=True, hard_negatives=None) +model.add_metric('EmbeddingMetric', is_training=True) + +# --- Training loop --- +for step, mb in enumerate(minibatches): + # `task='embedding'` is forwarded through the /twinkle/* protocol as an extra + # kwarg and selects the embedding pooling + InfoNCE loss path on the server. + model.forward_backward(inputs=mb, task='embedding') + model.clip_grad_and_step(max_grad_norm=1.0) + +# --- Read back embedding metrics (pos_sim / neg_sim / loss) --- +metric = model.calculate_metric(is_training=True) +``` + +**No manual apply_patch needed.** As long as you pass `task='embedding'` to `forward_backward` (or `forward_only`), the server automatically switches to embedding mode for that single `forward` call and rolls back afterwards — you do **not** need to call `apply_patch(...)` explicitly. So after a `forward_backward(task='embedding')`, calling `forward_only` without `task` on the same adapter still returns normal vocab-dimension `logits`, unaffected. + +The only three explicit setup steps on the client side are: `set_processor('InputProcessor')`, `set_loss('InfonceLoss', ...)`, `add_metric('EmbeddingMetric', is_training=True)`. Parameter meanings and recommended values are in "Key Parameters" above; returned metrics are in "Monitoring" below. + +--- + ## Monitoring The `EmbeddingMetric` reports key training signals: diff --git a/docs/source_en/Usage Guide/Quick-Start.md b/docs/source_en/Usage Guide/Quick-Start.md index 44c4dedf..69953a9a 100644 --- a/docs/source_en/Usage Guide/Quick-Start.md +++ b/docs/source_en/Usage Guide/Quick-Start.md @@ -473,7 +473,7 @@ python train.py A major feature of Twinkle is support for multi-tenant mixed training. Specifically, multiple users can use a single base model for LoRA training, which can greatly reduce server-side deployment costs. -Checkpoint resumption is also supported in client-server training. The recommended flow is to call `model.resume_from_checkpoint(resume_path)` to restore weights and optimizer state, then call `dataloader.resume_from_checkpoint(progress['consumed_train_samples'])` to skip consumed data. See [Twinkle-Client](./Server%20and%20Client/Twinkle-Client.md) and [self_cognition.py](https://github.com/modelscope/twinkle/blob/main/cookbook/client/twinkle/self_host/self_cognition.py). +Checkpoint resumption is also supported in client-server training. The recommended flow is to call `model.resume_from_checkpoint(resume_path)` to restore weights and optimizer state, then call `dataloader.resume_from_checkpoint(progress['consumed_train_samples'])` to skip consumed data. See [Twinkle-Client](./Server%20and%20Client/Twinkle-Client.md) and [self_cognition.py](https://github.com/modelscope/twinkle/blob/main/cookbook/client/twinkle/self_cognition.py). Suppose we start a service using eight GPUs. First, we need to start the Ray cluster: diff --git a/docs/source_en/Usage Guide/Server and Client/Overview.md b/docs/source_en/Usage Guide/Server and Client/Overview.md index a2944977..3c45673f 100644 --- a/docs/source_en/Usage Guide/Server and Client/Overview.md +++ b/docs/source_en/Usage Guide/Server and Client/Overview.md @@ -65,29 +65,22 @@ cookbook/ │ │ └── mock/ # Mock backend (CPU-only quick start) │ │ └── server_config.yaml ├── twinkle/ # Twinkle Client examples -│ ├── self_host/ # Self-hosted Server -│ │ ├── dpo.py # DPO training client -│ │ ├── multi_modal.py # Multi-modal training client -│ │ ├── sample.py # Inference sampling client -│ │ ├── self_congnition.py # Self-cognition training client -│ │ └── short_math_grpo.py # GRPO math training client -│ └── modelscope/ # ModelScope managed service -│ ├── dpo.py -│ ├── multi_modal.py -│ └── self_congnition.py +│ ├── dpo.py # DPO training client +│ ├── embedding.py # Embedding training client +│ ├── multi_modal.py # Multi-modal training client +│ ├── multi_turn_rollout.py # Multi-turn rollout client +│ ├── sample.py # Inference sampling client +│ ├── self_cognition.py # Self-cognition training client +│ ├── short_math_grpo.py # GRPO math training client +│ └── upload_to_hub.py # Checkpoint upload client └── tinker/ # Tinker Client examples - ├── self_host/ # Self-hosted Server - │ ├── dpo.py # DPO training client - │ ├── lora.py # LoRA training client - │ ├── multi_modal.py # Multi-modal training client - │ ├── sample.py # Inference sampling client - │ ├── self_cognition.py # Self-cognition training client - │ └── short_math_grpo.py # GRPO math training client - └── modelscope/ # ModelScope managed service - ├── dpo.py - ├── sample.py - ├── self_cognition.py - └── short_math_grpo.py + ├── dpo.py # DPO training client + ├── lora.py # LoRA training client + ├── multi_modal.py # Multi-modal training client + ├── sample.py # Inference sampling client + ├── self_cognition.py # Self-cognition training client + ├── short_math_grpo.py # GRPO math training client + └── upload_to_hub.py # Checkpoint upload client ``` Running steps: @@ -96,9 +89,22 @@ Running steps: # 1. Start Server first twinkle-server launch -c cookbook/client/server/transformer/server_config.yaml -# 2. Run Client in another terminal (Tinker Client example) -python cookbook/client/tinker/self_host/self_cognition.py +# 2. Configure the client in another terminal +export TWINKLE_SERVER_URL=http://localhost:8000 +export TWINKLE_SERVER_TOKEN=EMPTY_TOKEN +export TWINKLE_MODEL_ID=Qwen/Qwen3.5-4B + +# 3. Run a Tinker Client example +python cookbook/client/tinker/self_cognition.py # Or use Twinkle Client -python cookbook/client/twinkle/self_host/self_cognition.py +python cookbook/client/twinkle/self_cognition.py +``` + +The same examples work with the ModelScope managed service by changing only the environment: + +```bash +export TWINKLE_SERVER_URL=https://www.modelscope.cn/twinkle +export TWINKLE_SERVER_TOKEN="$MODELSCOPE_TOKEN" +export TWINKLE_MODEL_ID=Qwen/Qwen3.6-27B ``` diff --git a/docs/source_en/Usage Guide/Server and Client/Twinkle-Client.md b/docs/source_en/Usage Guide/Server and Client/Twinkle-Client.md index 5a118a61..598f215e 100644 --- a/docs/source_en/Usage Guide/Server and Client/Twinkle-Client.md +++ b/docs/source_en/Usage Guide/Server and Client/Twinkle-Client.md @@ -186,7 +186,7 @@ For checkpoint resumption, the recommended client-side flow is: 2. Call `model.resume_from_checkpoint(resume_path)` to restore weights, optimizer, scheduler, RNG, and progress metadata. 3. Call `dataloader.resume_from_checkpoint(progress['consumed_train_samples'])` to skip already-consumed samples. -This matches the end-to-end example in `cookbook/client/twinkle/self_host/self_cognition.py`. +This matches the end-to-end example in `cookbook/client/twinkle/self_cognition.py`. ## Differences with Megatron Backend @@ -202,3 +202,97 @@ model.set_lr_scheduler('default', lr_decay_steps=1000, max_lr=1e-4) ``` The rest of the data processing, training loop, checkpoint saving, and other code remains exactly the same. + +## Trainable Multi-turn Rollout (ClientMultiTurnRollout) + +The examples above are all single-turn training. If you want to do **multi-turn agentic RL with tool use** (e.g. GRPO) and need training-ready token-level alignment info, use `twinkle_client.rollout.ClientMultiTurnRollout`. It drives the "sample → call tool → stitch context → sample again" multi-turn loop on the client side, samples over HTTP each round (`/twinkle/sample`), and produces a trainable result with `logprobs` per trajectory that can be fed directly into GRPO and other RL training. + +### Dependencies and Constraints + +- **Local Template**: bridge-token stitching (rendering tool turns + the next generation prompt) requires a local `Template` instance on the client. +- **vLLMSampler**: the client sampler pointing at the server's Sampler service. +- **ToolManager** (optional): register your tools; if a trajectory produces tool_calls but no tool_manager is provided, a `ValueError` is raised at dispatch. +- **`num_samples=1`**: each trajectory is sampled once. For a GRPO group, replicate the same prompt into `NUM_GENERATIONS` independent trajectories. + +### Minimal Example + +```python +from peft import LoraConfig +from twinkle import init_twinkle_client +from twinkle.advantage import GRPOAdvantage +from twinkle.data_format import SamplingParams +from twinkle.template import Qwen3_5Template +from twinkle_agentic.tools.tool_manager import ToolManager +from twinkle_client.model import MultiLoraTransformersModel +from twinkle_client.rollout import ClientMultiTurnRollout +from twinkle_client.sampler import vLLMSampler + +MODEL_ID = 'ms://Qwen/Qwen3.5-4B' +NUM_GENERATIONS = 2 # GRPO group size (rollout samples num_samples=1 per trajectory) + +init_twinkle_client(base_url='http://127.0.0.1:8000', api_key='EMPTY_TOKEN') + +# Training model (GRPO) +model = MultiLoraTransformersModel(model_id=MODEL_ID) +model.add_adapter_to_model('default', LoraConfig(target_modules='all-linear', r=16, lora_alpha=32)) +model.set_loss('GRPOLoss', epsilon=0.2) +model.set_optimizer('Adam', lr=1e-5) +model.set_processor('InputProcessor') +model.set_template('Qwen3_5Template', model_id=MODEL_ID, enable_thinking=False) + +# Client sampler (HTTP) +sampler = vLLMSampler(model_id=MODEL_ID) +sampler.set_template('Qwen3_5Template', model_id=MODEL_ID, enable_thinking=False) + +# Multi-turn rollout: needs a local Template (bridge stitching) and a ToolManager +rollout_template = Qwen3_5Template(model_id=MODEL_ID, max_length=8192, enable_thinking=False) +rollout_template.truncation_strategy = 'delete' +tool_manager = ToolManager([MyCalculatorTool()]) # your tools + +rollout = ClientMultiTurnRollout( + sampler=sampler, + template=rollout_template, + tool_manager=tool_manager, + sampling_params=SamplingParams(max_tokens=512, num_samples=1, logprobs=1, temperature=1.0, top_p=0.95), + max_turns=4, +) +advantage_fn = GRPOAdvantage() + +for step in range(3): + # 1. Batched multi-turn rollout: replicate each prompt into NUM_GENERATIONS trajectories + trajectories = build_trajectories(tool_manager.tool_infos()) # see cookbook + rolled = rollout(trajectories, tool_manager=tool_manager) + + # 2. Read back token-level logprobs (top-1) and rewards + all_inputs, all_old_logps = [], [] + for traj in rolled: + all_old_logps.append([lp[0][1] for lp in (traj.get('logprobs') or [])]) + all_inputs.append(traj) + rewards = compute_rewards(rolled) # see cookbook + + # 3. GRPO advantages (group-relative) + advantages = advantage_fn(rewards, num_generations=NUM_GENERATIONS, scale='group').tolist() + + # 4. Policy update + model.forward_backward(inputs=all_inputs, advantages=advantages, old_logps=all_old_logps) + model.clip_grad_and_step() +``` + +### Output Fields + +Each returned trajectory has the following top-level fields appended to the original dict: + +| Field | Meaning | +|:------|:--------| +| `messages` | Full multi-turn conversation (including assistant tool_calls and tool-response turns) | +| `logprobs` | Top-1 logprob per trainable token; `None` if the round sampled no logprobs | +| `turns` | Number of turns actually taken (`<= max_turns`) | +| `stop_reason` | One of `'stop'` / `'length'` / `'max_turns'` | +| `truncated` | Whether truncated due to `max_turns` or the length cap | + +### Common Errors + +- A trajectory triggered a tool call but no `tool_manager` was provided → raises `ValueError`. Pass `tool_manager` at construction or per call. +- The sampler's network / timeout errors are **raised as-is** (not swallowed); handle retry/backoff outside your loop. + +See `cookbook/client/twinkle/multi_turn_rollout.py` for a full runnable example. diff --git "a/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/Embedding\350\256\255\347\273\203.md" "b/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/Embedding\350\256\255\347\273\203.md" index 94ea86eb..7e3829ef 100644 --- "a/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/Embedding\350\256\255\347\273\203.md" +++ "b/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/Embedding\350\256\255\347\273\203.md" @@ -106,6 +106,62 @@ for step, batch in enumerate(dataloader): --- +## 通过 Twinkle Client 训练 + +除了直接使用裸库 `TransformersModel`，你也可以通过 `twinkle_client` 的 `MultiLoraTransformersModel` 以 HTTP 方式驱动服务端完成 Embedding 训练。用法与裸库一致，只需按正确顺序调用现有方法。 + +### 完整调用顺序 + +client 化的 Embedding 训练遵循与裸库一致的调用顺序： + +``` +set_processor('InputProcessor') + -> set_loss('InfonceLoss', ...) + -> add_metric('EmbeddingMetric', is_training=True) + -> loop: forward_backward(inputs=mb, task='embedding') + clip_grad_and_step(...) + -> calculate_metric(is_training=True) +``` + +与裸库不同的是，client 侧传入的是**类名字符串**（如 `'InfonceLoss'`、`'EmbeddingMetric'`、`'InputProcessor'`），由服务端解析为对应的核心库类。 + +### 示例 + +```python +from peft import LoraConfig +from twinkle_client import init_twinkle_client +from twinkle_client.model import MultiLoraTransformersModel + +# --- Connect to the running Twinkle server --- +init_twinkle_client(base_url='http://127.0.0.1:8000', api_key='EMPTY_TOKEN') + +# --- Build the client model with a LoRA adapter for embedding --- +model = MultiLoraTransformersModel(model_id='ms://Qwen/Qwen3.5-4B') +model.add_adapter_to_model('emb_adapter', LoraConfig(target_modules='all-linear')) +model.set_template('Qwen3_5Template') + +# --- Configure the embedding-training pipeline (order matters) --- +# NOTE: pass class names as strings; the server resolves them to core-lib classes. +model.set_processor('InputProcessor') +model.set_loss('InfonceLoss', temperature=0.07, use_batch=True, hard_negatives=None) +model.add_metric('EmbeddingMetric', is_training=True) + +# --- Training loop --- +for step, mb in enumerate(minibatches): + # `task='embedding'` is forwarded through the /twinkle/* protocol as an extra + # kwarg and selects the embedding pooling + InfoNCE loss path on the server. + model.forward_backward(inputs=mb, task='embedding') + model.clip_grad_and_step(max_grad_norm=1.0) + +# --- Read back embedding metrics (pos_sim / neg_sim / loss) --- +metric = model.calculate_metric(is_training=True) +``` + +**无需手动 apply_patch。** 只要在 `forward_backward`（或 `forward_only`）中传入 `task='embedding'`，服务端就会在这一次 `forward` 期间自动切换到 embedding 模式、并在结束后自动回滚——你**不需要**显式调用 `apply_patch(...)`。因此同一个 adapter 在 `forward_backward(task='embedding')` 之后，再调用不带 `task` 的 `forward_only` 仍会返回正常的词表维度 `logits`，互不影响。 + +client 侧真正需要显式配置的只有三步：`set_processor('InputProcessor')`、`set_loss('InfonceLoss', ...)`、`add_metric('EmbeddingMetric', is_training=True)`；各参数含义与取值建议见上文「关键参数」，返回指标见下文「监控指标」。 + +--- + ## 监控指标 `EmbeddingMetric` 报告关键训练信号： diff --git "a/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/\345\277\253\351\200\237\345\274\200\345\247\213.md" "b/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/\345\277\253\351\200\237\345\274\200\345\247\213.md" index b4bbcd49..0cfaa4bb 100644 --- "a/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/\345\277\253\351\200\237\345\274\200\345\247\213.md" +++ "b/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/\345\277\253\351\200\237\345\274\200\345\247\213.md" @@ -472,7 +472,7 @@ python train.py ``` ### 远程训练 -client-server 训练场景同样支持断点续训。推荐流程是调用 `model.resume_from_checkpoint(resume_path)` 恢复权重和优化器状态，再调用 `dataloader.resume_from_checkpoint(progress['consumed_train_samples'])` 跳过已消费数据。详细示例可参考 [Twinkle客户端](./服务端和客户端/Twinkle客户端.md) 和 [self_cognition.py](https://github.com/modelscope/twinkle/blob/main/cookbook/client/twinkle/self_host/self_cognition.py)。 +client-server 训练场景同样支持断点续训。推荐流程是调用 `model.resume_from_checkpoint(resume_path)` 恢复权重和优化器状态，再调用 `dataloader.resume_from_checkpoint(progress['consumed_train_samples'])` 跳过已消费数据。详细示例可参考 [Twinkle客户端](./服务端和客户端/Twinkle客户端.md) 和 [self_cognition.py](https://github.com/modelscope/twinkle/blob/main/cookbook/client/twinkle/self_cognition.py)。 Twinkle 的一大特色是支持多租户用户混合训练。具体来说，多个用户可以使用一个基模进行 LoRA 训练，这样可以极大减小服务端部署成本。 diff --git "a/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/\346\234\215\345\212\241\347\253\257\345\222\214\345\256\242\346\210\267\347\253\257/Twinkle\345\256\242\346\210\267\347\253\257.md" "b/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/\346\234\215\345\212\241\347\253\257\345\222\214\345\256\242\346\210\267\347\253\257/Twinkle\345\256\242\346\210\267\347\253\257.md" index 4c3fa09e..668156d7 100644 --- "a/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/\346\234\215\345\212\241\347\253\257\345\222\214\345\256\242\346\210\267\347\253\257/Twinkle\345\256\242\346\210\267\347\253\257.md" +++ "b/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/\346\234\215\345\212\241\347\253\257\345\222\214\345\256\242\346\210\267\347\253\257/Twinkle\345\256\242\346\210\267\347\253\257.md" @@ -186,7 +186,7 @@ Twinkle Client 场景下，推荐的断点续训流程是： 2. 调用 `model.resume_from_checkpoint(resume_path)` 恢复权重、优化器、调度器、随机数状态和训练进度元数据。 3. 使用返回结果中的 `consumed_train_samples` 调用 `dataloader.resume_from_checkpoint(...)`，跳过已经训练过的数据。 -完整示例可直接参考 `cookbook/client/twinkle/self_host/self_cognition.py`。 +完整示例可直接参考 `cookbook/client/twinkle/self_cognition.py`。 ## Megatron 后端的差异 @@ -202,3 +202,97 @@ model.set_lr_scheduler('default', lr_decay_steps=1000, max_lr=1e-4) ``` 其余数据处理、训练循环、检查点保存等代码完全相同。 + +## 可训练多轮 Rollout（ClientMultiTurnRollout） + +前面的示例都是单轮训练。如果你要做**带工具调用的多轮 agentic RL**（如 GRPO），并且需要产出可直接用于训练的 token 级对齐信息，可使用 `twinkle_client.rollout.ClientMultiTurnRollout`。它在客户端侧驱动 “采样 → 调用工具 → 拼接上下文 → 再采样” 的多轮循环，每轮采样走 HTTP（`/twinkle/sample`），每条 trajectory 产出带 `logprobs` 的可训练结果，可直接用于 GRPO 等 RL 训练。 + +### 依赖与约束 + +- **本地 Template**：bridge token 拼接（渲染工具轮 + 下一轮生成提示）需要在客户端本地持有一个 `Template` 实例。 +- **vLLMSampler**：指向服务端 Sampler 服务的客户端采样器。 +- **ToolManager**（可选）：注册你的工具；若某条 trajectory 产生了 tool_calls 但未提供 tool_manager，会在派发时抛 `ValueError`。 +- **`num_samples=1`**：当前每条 trajectory 只采样一次。做 GRPO group 时，把同一个 prompt 复制成 `NUM_GENERATIONS` 条独立 trajectory 即可。 + +### 最小示例 + +```python +from peft import LoraConfig +from twinkle import init_twinkle_client +from twinkle.advantage import GRPOAdvantage +from twinkle.data_format import SamplingParams +from twinkle.template import Qwen3_5Template +from twinkle_agentic.tools.tool_manager import ToolManager +from twinkle_client.model import MultiLoraTransformersModel +from twinkle_client.rollout import ClientMultiTurnRollout +from twinkle_client.sampler import vLLMSampler + +MODEL_ID = 'ms://Qwen/Qwen3.5-4B' +NUM_GENERATIONS = 2 # GRPO group size（rollout 每条采样 num_samples=1） + +init_twinkle_client(base_url='http://127.0.0.1:8000', api_key='EMPTY_TOKEN') + +# 训练模型（GRPO） +model = MultiLoraTransformersModel(model_id=MODEL_ID) +model.add_adapter_to_model('default', LoraConfig(target_modules='all-linear', r=16, lora_alpha=32)) +model.set_loss('GRPOLoss', epsilon=0.2) +model.set_optimizer('Adam', lr=1e-5) +model.set_processor('InputProcessor') +model.set_template('Qwen3_5Template', model_id=MODEL_ID, enable_thinking=False) + +# 客户端采样器（HTTP） +sampler = vLLMSampler(model_id=MODEL_ID) +sampler.set_template('Qwen3_5Template', model_id=MODEL_ID, enable_thinking=False) + +# 多轮 rollout：需要本地 Template（bridge 拼接）与 ToolManager +rollout_template = Qwen3_5Template(model_id=MODEL_ID, max_length=8192, enable_thinking=False) +rollout_template.truncation_strategy = 'delete' +tool_manager = ToolManager([MyCalculatorTool()]) # 你的工具 + +rollout = ClientMultiTurnRollout( + sampler=sampler, + template=rollout_template, + tool_manager=tool_manager, + sampling_params=SamplingParams(max_tokens=512, num_samples=1, logprobs=1, temperature=1.0, top_p=0.95), + max_turns=4, +) +advantage_fn = GRPOAdvantage() + +for step in range(3): + # 1. 批量多轮 rollout：每个 prompt 复制成 NUM_GENERATIONS 条 trajectory + trajectories = build_trajectories(tool_manager.tool_infos()) # 见 cookbook + rolled = rollout(trajectories, tool_manager=tool_manager) + + # 2. 读回 token 级 logprobs（top-1）与 reward + all_inputs, all_old_logps = [], [] + for traj in rolled: + all_old_logps.append([lp[0][1] for lp in (traj.get('logprobs') or [])]) + all_inputs.append(traj) + rewards = compute_rewards(rolled) # 见 cookbook + + # 3. GRPO 优势（组内相对） + advantages = advantage_fn(rewards, num_generations=NUM_GENERATIONS, scale='group').tolist() + + # 4. 策略更新 + model.forward_backward(inputs=all_inputs, advantages=advantages, old_logps=all_old_logps) + model.clip_grad_and_step() +``` + +### 输出字段 + +每条返回的 trajectory 在原 dict 基础上追加以下顶层字段： + +| 字段 | 含义 | +|:----|:-----| +| `messages` | 完整多轮对话（含 assistant 的 tool_calls 与 tool 响应轮） | +| `logprobs` | 各可训练 token 的 top-1 logprob；本轮未采样 logprobs 时为 `None` | +| `turns` | 实际经历的轮数（`<= max_turns`） | +| `stop_reason` | `'stop'` / `'length'` / `'max_turns'` 之一 | +| `truncated` | 是否因 `max_turns` 或长度上限被截断 | + +### 常见错误 + +- 某条 trajectory 触发了工具调用，但没有提供 `tool_manager` → 抛 `ValueError`。构造时或按调用传入 `tool_manager` 即可。 +- 采样器的网络 / 超时错误会**原样抛出**（不被吞掉），重试、退避请在你的循环外层处理。 + +完整可运行示例见 `cookbook/client/twinkle/multi_turn_rollout.py`。 diff --git "a/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/\346\234\215\345\212\241\347\253\257\345\222\214\345\256\242\346\210\267\347\253\257/\346\246\202\350\277\260.md" "b/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/\346\234\215\345\212\241\347\253\257\345\222\214\345\256\242\346\210\267\347\253\257/\346\246\202\350\277\260.md" index 48b89947..5d3044e0 100644 --- "a/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/\346\234\215\345\212\241\347\253\257\345\222\214\345\256\242\346\210\267\347\253\257/\346\246\202\350\277\260.md" +++ "b/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/\346\234\215\345\212\241\347\253\257\345\222\214\345\256\242\346\210\267\347\253\257/\346\246\202\350\277\260.md" @@ -65,29 +65,22 @@ cookbook/ │ │ └── mock/ # Mock 后端（CPU-only 快速启动） │ │ └── server_config.yaml ├── twinkle/ # Twinkle Client 示例 -│ ├── self_host/ # 自托管 Server -│ │ ├── dpo.py # DPO 训练客户端 -│ │ ├── multi_modal.py # 多模态训练客户端 -│ │ ├── sample.py # 推理采样客户端 -│ │ ├── self_congnition.py # 自我认知训练客户端 -│ │ └── short_math_grpo.py # GRPO 数学训练客户端 -│ └── modelscope/ # ModelScope 托管服务 -│ ├── dpo.py -│ ├── multi_modal.py -│ └── self_congnition.py +│ ├── dpo.py # DPO 训练客户端 +│ ├── embedding.py # Embedding 训练客户端 +│ ├── multi_modal.py # 多模态训练客户端 +│ ├── multi_turn_rollout.py # 多轮 rollout 客户端 +│ ├── sample.py # 推理采样客户端 +│ ├── self_cognition.py # 自我认知训练客户端 +│ ├── short_math_grpo.py # GRPO 数学训练客户端 +│ └── upload_to_hub.py # Checkpoint 上传客户端 └── tinker/ # Tinker Client 示例 - ├── self_host/ # 自托管 Server - │ ├── dpo.py # DPO 训练客户端 - │ ├── lora.py # LoRA 训练客户端 - │ ├── multi_modal.py # 多模态训练客户端 - │ ├── sample.py # 推理采样客户端 - │ ├── self_cognition.py # 自我认知训练客户端 - │ └── short_math_grpo.py # GRPO 数学训练客户端 - └── modelscope/ # ModelScope 托管服务 - ├── dpo.py - ├── sample.py - ├── self_cognition.py - └── short_math_grpo.py + ├── dpo.py # DPO 训练客户端 + ├── lora.py # LoRA 训练客户端 + ├── multi_modal.py # 多模态训练客户端 + ├── sample.py # 推理采样客户端 + ├── self_cognition.py # 自我认知训练客户端 + ├── short_math_grpo.py # GRPO 数学训练客户端 + └── upload_to_hub.py # Checkpoint 上传客户端 ``` 运行步骤： @@ -96,9 +89,22 @@ cookbook/ # 1. 先启动 Server twinkle-server launch -c cookbook/client/server/transformer/server_config.yaml -# 2. 在另一个终端运行 Client（以 Tinker Client 为例） -python cookbook/client/tinker/self_host/self_cognition.py +# 2. 在另一个终端配置客户端 +export TWINKLE_SERVER_URL=http://localhost:8000 +export TWINKLE_SERVER_TOKEN=EMPTY_TOKEN +export TWINKLE_MODEL_ID=Qwen/Qwen3.5-4B + +# 3. 运行 Tinker Client 示例 +python cookbook/client/tinker/self_cognition.py # 或使用 Twinkle Client -python cookbook/client/twinkle/self_host/self_cognition.py +python cookbook/client/twinkle/self_cognition.py +``` + +同一组示例也可用于 ModelScope 托管服务，只需切换环境变量： + +```bash +export TWINKLE_SERVER_URL=https://www.modelscope.cn/twinkle +export TWINKLE_SERVER_TOKEN="$MODELSCOPE_TOKEN" +export TWINKLE_MODEL_ID=Qwen/Qwen3.6-27B ``` diff --git a/src/twinkle/model/megatron/megatron.py b/src/twinkle/model/megatron/megatron.py index a15df9af..89d0a171 100644 --- a/src/twinkle/model/megatron/megatron.py +++ b/src/twinkle/model/megatron/megatron.py @@ -698,7 +698,13 @@ def set_loss(self, loss_cls: Union[Loss, Type[Loss], str, Callable[[InputFeature """ adapter_name = kwargs.pop('adapter_name', self._get_default_group()) optimizer_config = self.optimizer_group[adapter_name] - optimizer_config.loss_instance = construct_class(loss_cls, Loss, twinkle.loss, **kwargs) + loss = construct_class(loss_cls, Loss, twinkle.loss, **kwargs) + # Gather over the DP group, not WORLD: under PP the loss runs only on the last + # pipeline stage, so a WORLD-group all-gather (e.g. InfonceLoss in-batch negatives) + # would deadlock on earlier-stage ranks that never enter the loss. Mirrors add_metric. + if getattr(loss, 'process_group', None) is None and getattr(optimizer_config, '_dp_group', None) is not None: + loss.process_group = optimizer_config._dp_group + optimizer_config.loss_instance = loss @remote_function() def add_metric(self, metric_cls: Union[Metric, str], is_training: Optional[bool] = None, **kwargs): diff --git a/src/twinkle/processor/base.py b/src/twinkle/processor/base.py index 7b1f9df5..ae37dbdd 100644 --- a/src/twinkle/processor/base.py +++ b/src/twinkle/processor/base.py @@ -194,7 +194,24 @@ def _postprocess_embedding(self, self.framework == 'transformers' and sp_strategy is not None and getattr(sp_strategy, 'enabled', False) and getattr(sp_strategy, 'world_size', 1) > 1) - ref_pos = sp_strategy.real_position_ids if sp_enabled else inputs['position_ids'] + if sp_enabled: + ref_pos = sp_strategy.real_position_ids + elif inputs.get('position_ids') is not None: + ref_pos = inputs['position_ids'] + else: + # Standard padded batch: HF omits position_ids (it derives them internally). + # Rebuild a position_ids-like reference (valid token -> its index, padding -> -1) + # so the last-valid-token pooling below, which keys off ``ref_pos >= 0`` for + # validity and a single ``== 0`` per row for packing detection, still works. + T = features.shape[1] + arange = torch.arange(T, device=features.device) + am = inputs.get('attention_mask') + if am is None: + ref_pos = arange.expand(features.shape[0], T) + else: + if am.dim() >= 3: + am = am.reshape(am.shape[0], -1)[:, :T] + ref_pos = torch.where(am.bool(), arange, arange.new_tensor(-1)) if ref_pos.dim() == 3: ref_pos = ref_pos[0] # Accept both flash-attn (cu_seq_lens_q) and Megatron (cu_seqlens_q) conventions. @@ -721,7 +738,8 @@ def is_mm_position_ids(position_ids): result[key] = result[key].reshape(len(values), num_axes, -1).permute(1, 0, 2).contiguous() elif isinstance(values[0], torch.Tensor): concat = False if (key == 'routed_experts') else None - result[key] = InputProcessor._pad_sequence(values, self.padding_map[key], self.padding_side, concat) + result[key] = InputProcessor._pad_sequence(values, self.padding_map.get(key, 0), self.padding_side, + concat) if result[key].dim() == 1: result[key] = result[key].unsqueeze(0) else: diff --git a/src/twinkle/server/sampler/backends/mock_sampler.py b/src/twinkle/server/sampler/backends/mock_sampler.py index ec5c86f9..300e5b0c 100644 --- a/src/twinkle/server/sampler/backends/mock_sampler.py +++ b/src/twinkle/server/sampler/backends/mock_sampler.py @@ -14,7 +14,7 @@ import numpy as np import time -from typing import Any +from typing import Any, Iterable from twinkle import remote_class, remote_function from twinkle.data_format import SampledSequence, SampleResponse, SamplingParams @@ -26,13 +26,53 @@ @remote_class() class MockSampler: - """Deterministic mock sampler for CPU-only testing.""" + """Deterministic mock sampler for CPU-only testing. - def __init__(self, model_id: str, *, seed: int = 0, vocab_size: int = 32, **kwargs: Any) -> None: + Beyond the base sampler surface, this backend exposes a few opt-in knobs so + a local (CPU-only) multi-turn rollout e2e can be driven without a GPU: + + - ``stop_reason`` (default ``'length'``): the ``SampledSequence.stop_reason`` + emitted for every sampled sequence. ``'length'`` terminates a multi-turn + loop immediately; ``'stop'`` lets the loop proceed to tool-call parsing. + - ``tool_call_text`` (default ``None``): when set, this text is emitted as + ``SampledSequence.decoded`` on the configured turns so the rollout's + ``template.parse_tool_call`` can produce a tool call and exercise the + "sample -> tool -> bridge -> sample" control flow. + - ``tool_call_turns`` (default ``(1,)`` when ``tool_call_text`` is set): + the 1-based turn indices on which ``tool_call_text`` is injected. A turn + corresponds to one ``sample()`` invocation (one multi-turn round). + + All three knobs may be supplied at construction time or overridden per call + via ``sample()`` keyword arguments. When none are configured the backend + keeps its previous behaviour exactly (``stop_reason='length'``, + ``decoded=None``), so existing callers are unaffected. Outputs stay + deterministic and CPU-only. + """ + + def __init__( + self, + model_id: str, + *, + seed: int = 0, + vocab_size: int = 32, + stop_reason: str = 'length', + tool_call_text: str | None = None, + tool_call_turns: Iterable[int] | None = None, + **kwargs: Any, + ) -> None: self.model_id = model_id self._seed = int(seed) self._vocab_size = int(vocab_size) self._adapters: dict[str, Any] = {} + # Multi-turn control-flow knobs (see class docstring). + self._stop_reason = str(stop_reason) + self._tool_call_text = tool_call_text + self._tool_call_turns = self._normalize_turns(tool_call_turns, tool_call_text) + # Per-round counter: incremented once per ``sample()`` call so that + # ``tool_call_turns`` can address individual multi-turn rounds. Only + # consulted when tool-call injection is active, so the default path + # stays fully stateless and deterministic. + self._round = 0 # Surface (rather than silently swallow) extra ctor kwargs: a real # backend signature drift then shows up as a visible DEBUG warning in # the mock e2e instead of being discarded without trace. @@ -64,9 +104,19 @@ def sample( raise ValueError(f'max_tokens must be >= 1, got {max_tokens!r} ' '(set sampling_params.max_tokens to a positive integer)') + # Resolve per-call multi-turn knobs (explicit sample() kwargs override + # ctor defaults) and advance the round counter that ``tool_call_turns`` + # addresses. + stop_reason = self._resolve_stop_reason(kwargs) + tool_call_text = self._resolve_tool_call_text(kwargs) + tool_call_turns = self._resolve_tool_call_turns(kwargs, tool_call_text) + self._round += 1 + inject_tool_call = tool_call_text is not None and self._round in tool_call_turns + decoded = tool_call_text if inject_tool_call else None + normalized = self._normalize_inputs(inputs) responses: list[SampleResponse] = [] - for prompt_idx, _ in enumerate(normalized): + for prompt_idx, pif in enumerate(normalized): sequences: list[SampledSequence] = [] for sample_idx in range(num_samples): seed = stable_seed(self.model_id, adapter_name, self._seed, prompt_idx, sample_idx) @@ -78,11 +128,14 @@ def sample( # mock returns top-1 with the chosen token. The tinker handler # flattens this to a single chosen-token logprob. logprobs = [[(tok, float(lp))] for tok, lp in zip(tokens, logprobs_per_token)] - sequences.append(SampledSequence( - stop_reason='length', - tokens=tokens, - logprobs=logprobs, - )) + sequences.append( + SampledSequence( + stop_reason=stop_reason, + tokens=tokens, + logprobs=logprobs, + decoded=decoded, + new_input_feature=self._build_new_input_feature(pif, tokens), + )) responses.append(SampleResponse(sequences=sequences)) return responses @@ -150,3 +203,65 @@ def _resolve_max_tokens(params: SamplingParams | None) -> int | None: if params is None: return None return getattr(params, 'max_tokens', None) + + @staticmethod + def _normalize_turns(turns: Iterable[int] | None, tool_call_text: str | None) -> frozenset[int]: + """Normalise the ``tool_call_turns`` config into a ``frozenset[int]``. + + When ``tool_call_text`` is set but no turns are given, default to + ``{1}`` — inject the tool call exactly once, on the first round. + """ + if turns is None: + return frozenset({1}) if tool_call_text is not None else frozenset() + return frozenset(int(t) for t in turns) + + def _resolve_stop_reason(self, kwargs: dict[str, Any]) -> str: + if 'stop_reason' in kwargs and kwargs['stop_reason'] is not None: + return str(kwargs['stop_reason']) + return self._stop_reason + + def _resolve_tool_call_text(self, kwargs: dict[str, Any]) -> str | None: + if 'tool_call_text' in kwargs: + return kwargs['tool_call_text'] + return self._tool_call_text + + def _resolve_tool_call_turns( + self, + kwargs: dict[str, Any], + tool_call_text: str | None, + ) -> frozenset[int]: + if 'tool_call_turns' in kwargs and kwargs['tool_call_turns'] is not None: + return self._normalize_turns(kwargs['tool_call_turns'], tool_call_text) + if tool_call_text is not None and not self._tool_call_turns: + return frozenset({1}) + return self._tool_call_turns + + @staticmethod + def _build_new_input_feature(pif: Any, tokens: list[int]) -> dict[str, Any]: + """Deterministically append the freshly sampled ``tokens`` to ``pif``. + + Produces a plain-dict ``InputFeature`` that carries the running context + for the next multi-turn round: ``input_ids`` is the prior prompt plus + this round's sampled tokens, and ``labels`` marks the sampled tokens as + trainable (their own ids) while prior/context positions stay ``-100``. + This mirrors the shape a real sampler's ``concat_input_feature`` yields, + which the multi-turn rollout relies on (it reads + ``new_input_feature.input_ids`` and counts trainable ``labels``). + + The mock does not depend on a template, so the append is a simple, + deterministic list concatenation performed entirely on the CPU. + """ + feat: dict[str, Any] = dict(pif) if isinstance(pif, dict) else {} + prev_ids = list(feat.get('input_ids') or []) + prev_labels = feat.get('labels') + if prev_labels is not None and len(prev_labels) == len(prev_ids): + labels = list(prev_labels) + else: + # No (or misaligned) prior labels: treat the entire prior context as + # non-trainable so only this round's sampled tokens count. + labels = [-100] * len(prev_ids) + + feat['input_ids'] = prev_ids + list(tokens) + feat['labels'] = labels + list(tokens) + feat['length'] = len(feat['input_ids']) + return feat diff --git a/src/twinkle/server/telemetry/worker_init.py b/src/twinkle/server/telemetry/worker_init.py index 40edc662..dcf53b90 100644 --- a/src/twinkle/server/telemetry/worker_init.py +++ b/src/twinkle/server/telemetry/worker_init.py @@ -79,7 +79,7 @@ def flush_telemetry_safely() -> None: worker deployment's FastAPI lifespan shutdown so buffered OTLP batches (traces / metrics / logs) flush on graceful termination. A telemetry-shutdown failure MUST NOT mask the user-facing shutdown path, - so every error here is swallowed (Requirement 21.3). + so every error here is swallowed. """ try: from twinkle.server.telemetry import shutdown_telemetry diff --git a/src/twinkle_agentic/rollout/__init__.py b/src/twinkle_agentic/rollout/__init__.py index 52c0fb12..67c589e0 100644 --- a/src/twinkle_agentic/rollout/__init__.py +++ b/src/twinkle_agentic/rollout/__init__.py @@ -1,6 +1,7 @@ # Copyright (c) ModelScope Contributors. All rights reserved. from .api_multi_turn import APIMultiTurnRollout from .base import Rollout +from .bridge import extend_with_bridge from .multi_turn import MultiTurnRollout from .multi_turn_condense import MultiTurnCondenseRollout @@ -9,4 +10,5 @@ 'MultiTurnCondenseRollout', 'MultiTurnRollout', 'Rollout', + 'extend_with_bridge', ] diff --git a/src/twinkle_agentic/rollout/bridge.py b/src/twinkle_agentic/rollout/bridge.py new file mode 100644 index 00000000..2663f9ed --- /dev/null +++ b/src/twinkle_agentic/rollout/bridge.py @@ -0,0 +1,156 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Shared, pure bridge-token stitching logic for multi-turn rollouts. + +This module hosts :func:`extend_with_bridge`, a ``self``-free function that +appends tool messages and the next generation prompt to a running +``InputFeature`` (``pif``) as ``-100`` "bridge" tokens. It is shared between +the core-library ``MultiTurnRollout`` and the client-side rollout so the two +paths cannot drift. + +The logic was lifted verbatim from ``MultiTurnRollout._extend_with_bridge`` and +``MultiTurnRollout._append_bridge_tokens``; every ``self.template`` access was +rewritten to use the ``template`` parameter. No Ray decorators +(``@remote_function`` / ``@remote_class``) are applied here. +""" +import numpy as np +from typing import Any, Dict, List, Optional + +from twinkle.template.base import Template + + +def _to_plain(obj: Any) -> Any: + """Recursively convert numpy arrays/scalars to plain Python lists/numbers. + + Mirrors ``vllm_sampler._convert_ndarray_to_list`` but lives locally so we + do not depend on a private symbol. + """ + if isinstance(obj, np.ndarray): + return obj.tolist() + if isinstance(obj, np.integer): + return int(obj) + if isinstance(obj, np.floating): + return float(obj) + if isinstance(obj, np.bool_): + return bool(obj) + if isinstance(obj, dict): + return {k: _to_plain(v) for k, v in obj.items()} + if isinstance(obj, (list, tuple)): + conv = [_to_plain(x) for x in obj] + return type(obj)(conv) if isinstance(obj, tuple) else conv + return obj + + +def extend_with_bridge( + pif: Dict[str, Any], + tool_messages: List[Dict[str, Any]], + template: Template, +) -> Optional[Dict[str, Any]]: + """Append tool messages and the next generation prompt as -100 bridge. + + Strategy: compute the bridge ENTIRELY in template space. Render + ``messages_before`` and ``messages_before + tool_messages`` with the + same chat template and take ``s_after[len(s_before):]`` as the delta. + + We deliberately do NOT diff against ``tokenizer.decode(pif.input_ids)`` + because raw vLLM output and canonical template rendering differ in + whitespace (e.g. Qwen inserts ``\\n\\n`` between assistant content and + a ```` block, while the model generates only ``\\n``). Such + cosmetic divergences would break a ``startswith`` alignment but do not + affect training correctness: history tokens stay in ``pif.input_ids`` + verbatim; only the newly appended bridge is tokenized from the + canonical template output. + + Returns ``None`` when the trajectory exceeds ``max_length`` and the + template's truncation strategy is ``'delete'``. + """ + tokenizer = template.tokenizer + + messages_before = list(pif.get('messages') or []) + messages_after = messages_before + list(tool_messages) + + enable_thinking = getattr(template, 'enable_thinking', False) + s_before = tokenizer.apply_chat_template( + messages_before, tokenize=False, add_generation_prompt=False, enable_thinking=enable_thinking) + s_after = tokenizer.apply_chat_template( + messages_after, tokenize=False, add_generation_prompt=True, enable_thinking=enable_thinking) + + if not s_after.startswith(s_before): + raise RuntimeError('Canonical chat_template output for messages_after is not a ' + 'prefix-extension of messages_before; cannot compute bridge ' + 'delta. This indicates the template is non-monotonic in the ' + 'message list (e.g. reorders / rewrites earlier turns).\n' + f's_before tail: {s_before[-80:]!r}\n' + f's_after at same offset: ' + f'{s_after[max(0, len(s_before) - 80):len(s_before) + 80]!r}') + bridge_text = s_after[len(s_before):] + if not bridge_text: + raise RuntimeError('Bridge text computation returned empty string; ' + 'tool turn would add no tokens (template misconfiguration?).') + + bridge_ids = tokenizer.encode(bridge_text, add_special_tokens=False) + if not bridge_ids: + raise RuntimeError(f'Bridge text tokenised to empty id list: {bridge_text!r}') + + new_pif = _append_bridge_tokens(pif, bridge_ids, template) + if new_pif is None: + # Trajectory exceeds max_length and strategy is 'delete' + return None + new_pif['messages'] = messages_after + return new_pif + + +def _append_bridge_tokens( + pif: Dict[str, Any], + bridge_ids: List[int], + template: Template, +) -> Optional[Dict[str, Any]]: + """Append bridge tokens with labels = -100. + + Mirrors the unroll-append-reroll pattern of + :meth:`Template.concat_input_feature` so that ``labels`` semantics + stay consistent with the sampler-produced pif. + + Shallow copy is deliberately used: every mutation below is a + top-level key reassignment, never an in-place change to nested + tensors. Multimodal payloads (``images``, ``pixel_values``, + ``image_grid_thw`` ...) are shared by reference so we avoid + re-copying image buffers every turn. + """ + result = dict(pif) + + input_ids = list(result['input_ids']) + labels = list(result.get('labels') or []) + # labels arrive in output/shifted order (post _roll_labels). Unroll by + # one position (shift right by 1) to get back to input order. + if labels: + if len(labels) != len(input_ids): + raise RuntimeError(f'labels length ({len(labels)}) != input_ids length ' + f'({len(input_ids)}); cannot safely append bridge tokens.') + labels = labels[-1:] + labels[:-1] + else: + labels = [-100] * len(input_ids) + + input_ids = input_ids + list(bridge_ids) + labels = labels + [-100] * len(bridge_ids) + + result['input_ids'] = input_ids + result['labels'] = labels + + if 'mm_token_type_ids' in result: + import torch + mm = result['mm_token_type_ids'] + if not isinstance(mm, torch.Tensor): + mm = torch.as_tensor(mm) + # Pad along the last (sequence) dim — handles 1D [T] and 2D [1, T] uniformly. + leading_shape = mm.shape[:-1] + pad = torch.zeros((*leading_shape, len(bridge_ids)), dtype=mm.dtype, device=mm.device) + result['mm_token_type_ids'] = torch.cat([mm, pad], dim=-1) + + # Replay the post pipeline: refresh attention_mask / position_ids / + # length and re-roll labels back into output/shifted order. + refreshed_list = template._invoke_post_pipeline([result]) + if not refreshed_list: + # truncation_strategy='delete': trajectory exceeds max_length + return None + result.update(refreshed_list[0]) + return _to_plain(result) diff --git a/src/twinkle_agentic/rollout/multi_turn.py b/src/twinkle_agentic/rollout/multi_turn.py index 3b4035ed..786e4261 100644 --- a/src/twinkle_agentic/rollout/multi_turn.py +++ b/src/twinkle_agentic/rollout/multi_turn.py @@ -1,6 +1,5 @@ # Copyright (c) ModelScope Contributors. All rights reserved. import json -import numpy as np import os import re import time @@ -12,28 +11,7 @@ from twinkle.template.base import Template from twinkle_agentic.tools.tool_manager import ToolManager from .base import Rollout - - -def _to_plain(obj: Any) -> Any: - """Recursively convert numpy arrays/scalars to plain Python lists/numbers. - - Mirrors ``vllm_sampler._convert_ndarray_to_list`` but lives locally so we - do not depend on a private symbol. - """ - if isinstance(obj, np.ndarray): - return obj.tolist() - if isinstance(obj, np.integer): - return int(obj) - if isinstance(obj, np.floating): - return float(obj) - if isinstance(obj, np.bool_): - return bool(obj) - if isinstance(obj, dict): - return {k: _to_plain(v) for k, v in obj.items()} - if isinstance(obj, (list, tuple)): - conv = [_to_plain(x) for x in obj] - return type(obj)(conv) if isinstance(obj, tuple) else conv - return obj +from .bridge import _to_plain, extend_with_bridge @remote_class() @@ -214,7 +192,7 @@ def __call__(self, trajectories: List[Trajectory], **kwargs) -> List[Trajectory] # outstanding tool turns. Done serially: bridge computation is # a cheap decode-diff-encode on python strings / token lists. for global_idx, tool_messages in pending_bridges: - extended = self._extend_with_bridge(pifs[global_idx], tool_messages) + extended = extend_with_bridge(pifs[global_idx], tool_messages, self.template) if extended is None: # Trajectory exceeded max_length, mark as done (deleted) truncated[global_idx] = True @@ -403,114 +381,3 @@ def _unwrap_response_list(resps, expected: int) -> List[SampleResponse]: if not r.sequences: raise RuntimeError(f'SampleResponse at batch index {i} has no sequences') return resps - - def _extend_with_bridge( - self, - pif: Dict[str, Any], - tool_messages: List[Dict[str, Any]], - ) -> Dict[str, Any]: - """Append tool messages and the next generation prompt as -100 bridge. - - Strategy: compute the bridge ENTIRELY in template space. Render - ``messages_before`` and ``messages_before + tool_messages`` with the - same chat template and take ``s_after[len(s_before):]`` as the delta. - - We deliberately do NOT diff against ``tokenizer.decode(pif.input_ids)`` - because raw vLLM output and canonical template rendering differ in - whitespace (e.g. Qwen inserts ``\\n\\n`` between assistant content and - a ```` block, while the model generates only ``\\n``). Such - cosmetic divergences would break a ``startswith`` alignment but do not - affect training correctness: history tokens stay in ``pif.input_ids`` - verbatim; only the newly appended bridge is tokenized from the - canonical template output. - """ - tokenizer = self.template.tokenizer - - messages_before = list(pif.get('messages') or []) - messages_after = messages_before + list(tool_messages) - - enable_thinking = getattr(self.template, 'enable_thinking', False) - s_before = tokenizer.apply_chat_template( - messages_before, tokenize=False, add_generation_prompt=False, enable_thinking=enable_thinking) - s_after = tokenizer.apply_chat_template( - messages_after, tokenize=False, add_generation_prompt=True, enable_thinking=enable_thinking) - - if not s_after.startswith(s_before): - raise RuntimeError('Canonical chat_template output for messages_after is not a ' - 'prefix-extension of messages_before; cannot compute bridge ' - 'delta. This indicates the template is non-monotonic in the ' - 'message list (e.g. reorders / rewrites earlier turns).\n' - f's_before tail: {s_before[-80:]!r}\n' - f's_after at same offset: ' - f'{s_after[max(0, len(s_before) - 80):len(s_before) + 80]!r}') - bridge_text = s_after[len(s_before):] - if not bridge_text: - raise RuntimeError('Bridge text computation returned empty string; ' - 'tool turn would add no tokens (template misconfiguration?).') - - bridge_ids = tokenizer.encode(bridge_text, add_special_tokens=False) - if not bridge_ids: - raise RuntimeError(f'Bridge text tokenised to empty id list: {bridge_text!r}') - - new_pif = self._append_bridge_tokens(pif, bridge_ids) - if new_pif is None: - # Trajectory exceeds max_length and strategy is 'delete' - return None - new_pif['messages'] = messages_after - return new_pif - - def _append_bridge_tokens( - self, - pif: Dict[str, Any], - bridge_ids: List[int], - ) -> Dict[str, Any]: - """Append bridge tokens with labels = -100. - - Mirrors the unroll-append-reroll pattern of - :meth:`Template.concat_input_feature` so that ``labels`` semantics - stay consistent with the sampler-produced pif. - - Shallow copy is deliberately used: every mutation below is a - top-level key reassignment, never an in-place change to nested - tensors. Multimodal payloads (``images``, ``pixel_values``, - ``image_grid_thw`` ...) are shared by reference so we avoid - re-copying image buffers every turn. - """ - result = dict(pif) - - input_ids = list(result['input_ids']) - labels = list(result.get('labels') or []) - # labels arrive in output/shifted order (post _roll_labels). Unroll by - # one position (shift right by 1) to get back to input order. - if labels: - if len(labels) != len(input_ids): - raise RuntimeError(f'labels length ({len(labels)}) != input_ids length ' - f'({len(input_ids)}); cannot safely append bridge tokens.') - labels = labels[-1:] + labels[:-1] - else: - labels = [-100] * len(input_ids) - - input_ids = input_ids + list(bridge_ids) - labels = labels + [-100] * len(bridge_ids) - - result['input_ids'] = input_ids - result['labels'] = labels - - if 'mm_token_type_ids' in result: - import torch - mm = result['mm_token_type_ids'] - if not isinstance(mm, torch.Tensor): - mm = torch.as_tensor(mm) - # Pad along the last (sequence) dim — handles 1D [T] and 2D [1, T] uniformly. - leading_shape = mm.shape[:-1] - pad = torch.zeros((*leading_shape, len(bridge_ids)), dtype=mm.dtype, device=mm.device) - result['mm_token_type_ids'] = torch.cat([mm, pad], dim=-1) - - # Replay the post pipeline: refresh attention_mask / position_ids / - # length and re-roll labels back into output/shifted order. - refreshed_list = self.template._invoke_post_pipeline([result]) - if not refreshed_list: - # truncation_strategy='delete': trajectory exceeds max_length - return None - result.update(refreshed_list[0]) - return _to_plain(result) diff --git a/src/twinkle_client/dataloader/__init__.py b/src/twinkle_client/dataloader/__init__.py index 341d0b77..b94ed5b8 100644 --- a/src/twinkle_client/dataloader/__init__.py +++ b/src/twinkle_client/dataloader/__init__.py @@ -1,11 +1 @@ -# ============================================================================ -# WARNING: AUTO-GENERATED FILE - DO NOT MODIFY MANUALLY! -# ============================================================================ -# This file is automatically generated by client_tools/client_generator.py -# Any manual changes will be overwritten when the generator runs again. -# -# To update this file: -# 1. Modify the source files in src/twinkle/ -# 2. Run: python client_tools/client_generator.py -# ============================================================================ from .dataloader import DataLoader diff --git a/src/twinkle_client/dataloader/dataloader.py b/src/twinkle_client/dataloader/dataloader.py index 4164940d..86e2bbf4 100644 --- a/src/twinkle_client/dataloader/dataloader.py +++ b/src/twinkle_client/dataloader/dataloader.py @@ -1,13 +1,3 @@ -# ============================================================================ -# WARNING: AUTO-GENERATED FILE - DO NOT MODIFY MANUALLY! -# ============================================================================ -# This file is automatically generated by client_tools/client_generator.py -# Any manual changes will be overwritten when the generator runs again. -# -# To update this file: -# 1. Modify the source files in src/twinkle/ -# 2. Run: python client_tools/client_generator.py -# ============================================================================ from typing import Callable, Type, Union from twinkle_client.http import http_post diff --git a/src/twinkle_client/dataset/__init__.py b/src/twinkle_client/dataset/__init__.py index ad90b90a..ba37b1fe 100644 --- a/src/twinkle_client/dataset/__init__.py +++ b/src/twinkle_client/dataset/__init__.py @@ -1,13 +1,3 @@ -# ============================================================================ -# WARNING: AUTO-GENERATED FILE - DO NOT MODIFY MANUALLY! -# ============================================================================ -# This file is automatically generated by client_tools/client_generator.py -# Any manual changes will be overwritten when the generator runs again. -# -# To update this file: -# 1. Modify the source files in src/twinkle/ -# 2. Run: python client_tools/client_generator.py -# ============================================================================ from .base import Dataset from .iterable_dataset import IterableDataset from .iterable_packing_dataset import IterablePackingDataset diff --git a/src/twinkle_client/dataset/base.py b/src/twinkle_client/dataset/base.py index 845b11cb..bec2d430 100644 --- a/src/twinkle_client/dataset/base.py +++ b/src/twinkle_client/dataset/base.py @@ -1,13 +1,3 @@ -# ============================================================================ -# WARNING: AUTO-GENERATED FILE - DO NOT MODIFY MANUALLY! -# ============================================================================ -# This file is automatically generated by client_tools/client_generator.py -# Any manual changes will be overwritten when the generator runs again. -# -# To update this file: -# 1. Modify the source files in src/twinkle/ -# 2. Run: python client_tools/client_generator.py -# ============================================================================ from typing import Any, Callable, Dict, Optional, Type, Union from twinkle_client.http import http_post diff --git a/src/twinkle_client/dataset/iterable_dataset.py b/src/twinkle_client/dataset/iterable_dataset.py index 95ba9e44..0646ea33 100644 --- a/src/twinkle_client/dataset/iterable_dataset.py +++ b/src/twinkle_client/dataset/iterable_dataset.py @@ -1,13 +1,3 @@ -# ============================================================================ -# WARNING: AUTO-GENERATED FILE - DO NOT MODIFY MANUALLY! -# ============================================================================ -# This file is automatically generated by client_tools/client_generator.py -# Any manual changes will be overwritten when the generator runs again. -# -# To update this file: -# 1. Modify the source files in src/twinkle/ -# 2. Run: python client_tools/client_generator.py -# ============================================================================ from twinkle_client.http import http_post from twinkle.dataset import Dataset diff --git a/src/twinkle_client/dataset/iterable_packing_dataset.py b/src/twinkle_client/dataset/iterable_packing_dataset.py index f774090f..f4b6b5d1 100644 --- a/src/twinkle_client/dataset/iterable_packing_dataset.py +++ b/src/twinkle_client/dataset/iterable_packing_dataset.py @@ -1,13 +1,3 @@ -# ============================================================================ -# WARNING: AUTO-GENERATED FILE - DO NOT MODIFY MANUALLY! -# ============================================================================ -# This file is automatically generated by client_tools/client_generator.py -# Any manual changes will be overwritten when the generator runs again. -# -# To update this file: -# 1. Modify the source files in src/twinkle/ -# 2. Run: python client_tools/client_generator.py -# ============================================================================ from typing import Type, Union from twinkle_client.http import http_post diff --git a/src/twinkle_client/dataset/lazy_dataset.py b/src/twinkle_client/dataset/lazy_dataset.py index c3b70c85..7bf49c70 100644 --- a/src/twinkle_client/dataset/lazy_dataset.py +++ b/src/twinkle_client/dataset/lazy_dataset.py @@ -1,13 +1,3 @@ -# ============================================================================ -# WARNING: AUTO-GENERATED FILE - DO NOT MODIFY MANUALLY! -# ============================================================================ -# This file is automatically generated by client_tools/client_generator.py -# Any manual changes will be overwritten when the generator runs again. -# -# To update this file: -# 1. Modify the source files in src/twinkle/ -# 2. Run: python client_tools/client_generator.py -# ============================================================================ from typing import Any, Callable, Dict, Optional, Type, Union from twinkle_client.http import http_post diff --git a/src/twinkle_client/dataset/packing_dataset.py b/src/twinkle_client/dataset/packing_dataset.py index 1f7ff066..85577767 100644 --- a/src/twinkle_client/dataset/packing_dataset.py +++ b/src/twinkle_client/dataset/packing_dataset.py @@ -1,13 +1,3 @@ -# ============================================================================ -# WARNING: AUTO-GENERATED FILE - DO NOT MODIFY MANUALLY! -# ============================================================================ -# This file is automatically generated by client_tools/client_generator.py -# Any manual changes will be overwritten when the generator runs again. -# -# To update this file: -# 1. Modify the source files in src/twinkle/ -# 2. Run: python client_tools/client_generator.py -# ============================================================================ from twinkle_client.http import http_post from twinkle.dataset import Dataset diff --git a/src/twinkle_client/model/__init__.py b/src/twinkle_client/model/__init__.py index 507cc4cb..94e3538a 100644 --- a/src/twinkle_client/model/__init__.py +++ b/src/twinkle_client/model/__init__.py @@ -1,11 +1 @@ -# ============================================================================ -# WARNING: AUTO-GENERATED FILE - DO NOT MODIFY MANUALLY! -# ============================================================================ -# This file is automatically generated by client_tools/client_generator.py -# Any manual changes will be overwritten when the generator runs again. -# -# To update this file: -# 1. Modify the source files in src/twinkle/ -# 2. Run: python client_tools/client_generator.py -# ============================================================================ from .multi_lora_transformers import MultiLoraTransformersModel diff --git a/src/twinkle_client/model/multi_lora_transformers.py b/src/twinkle_client/model/multi_lora_transformers.py index 2508bbde..c628c353 100644 --- a/src/twinkle_client/model/multi_lora_transformers.py +++ b/src/twinkle_client/model/multi_lora_transformers.py @@ -1,13 +1,3 @@ -# ============================================================================ -# WARNING: AUTO-GENERATED FILE - DO NOT MODIFY MANUALLY! -# ============================================================================ -# This file is automatically generated by client_tools/client_generator.py -# Any manual changes will be overwritten when the generator runs again. -# -# To update this file: -# 1. Modify the source files in src/twinkle/ -# 2. Run: python client_tools/client_generator.py -# ============================================================================ from typing import Any, Dict, Optional from pathlib import Path import time diff --git a/src/twinkle_client/processor/__init__.py b/src/twinkle_client/processor/__init__.py index 1f8acd8f..677da196 100644 --- a/src/twinkle_client/processor/__init__.py +++ b/src/twinkle_client/processor/__init__.py @@ -1,11 +1 @@ -# ============================================================================ -# WARNING: AUTO-GENERATED FILE - DO NOT MODIFY MANUALLY! -# ============================================================================ -# This file is automatically generated by client_tools/client_generator.py -# Any manual changes will be overwritten when the generator runs again. -# -# To update this file: -# 1. Modify the source files in src/twinkle/ -# 2. Run: python client_tools/client_generator.py -# ============================================================================ from .base import InputProcessor diff --git a/src/twinkle_client/processor/base.py b/src/twinkle_client/processor/base.py index 048ace5e..37786527 100644 --- a/src/twinkle_client/processor/base.py +++ b/src/twinkle_client/processor/base.py @@ -1,13 +1,3 @@ -# ============================================================================ -# WARNING: AUTO-GENERATED FILE - DO NOT MODIFY MANUALLY! -# ============================================================================ -# This file is automatically generated by client_tools/client_generator.py -# Any manual changes will be overwritten when the generator runs again. -# -# To update this file: -# 1. Modify the source files in src/twinkle/ -# 2. Run: python client_tools/client_generator.py -# ============================================================================ from typing import List, Literal, Optional, Union from twinkle_client.http import http_post diff --git a/src/twinkle_client/rollout/__init__.py b/src/twinkle_client/rollout/__init__.py new file mode 100644 index 00000000..66e39339 --- /dev/null +++ b/src/twinkle_client/rollout/__init__.py @@ -0,0 +1,4 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from .multi_turn import ClientMultiTurnRollout + +__all__ = ['ClientMultiTurnRollout'] diff --git a/src/twinkle_client/rollout/multi_turn.py b/src/twinkle_client/rollout/multi_turn.py new file mode 100644 index 00000000..55c5800b --- /dev/null +++ b/src/twinkle_client/rollout/multi_turn.py @@ -0,0 +1,291 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Client-side multi-turn agentic rollout orchestration. + +This module hosts :class:`ClientMultiTurnRollout`, a hand-maintained multi-turn +rollout orchestrator whose algorithmic structure mirrors +``twinkle_agentic.rollout.multi_turn.MultiTurnRollout`` but issues sampling over +HTTP via ``twinkle_client.sampler.vLLMSampler.sample()`` instead of holding a +Ray actor handle. + +Design notes: + * It deliberately does NOT subclass ``MultiTurnRollout``. That class is + decorated with ``@remote_class`` / ``@remote_function`` for Ray remote + dispatch and assumes ``self.sampler`` is a Ray actor handle, which does + not match the HTTP-client semantics here. + * The ``tool_manager`` type is reused directly from + ``twinkle_agentic.tools.tool_manager.ToolManager`` (imported, not copied). + * Bridge-token stitching is reused from + ``twinkle_agentic.rollout.bridge.extend_with_bridge``. +""" +import dataclasses +from typing import Any, Dict, List, Optional + +from twinkle.data_format import Trajectory +from twinkle.data_format.sampling import SamplingParams +from twinkle.template.base import Template +from twinkle_agentic.rollout.bridge import extend_with_bridge +from twinkle_agentic.tools.tool_manager import ToolManager +from twinkle_client.sampler import vLLMSampler + + +class ClientMultiTurnRollout: + """Agentic multi-turn rollout with tool use, driven over HTTP. + + Mirrors the per-trajectory state machine of + ``twinkle_agentic.rollout.multi_turn.MultiTurnRollout`` but issues sampling + via ``vLLMSampler.sample()`` (an HTTP call to ``/twinkle/sample``) rather + than a Ray actor call. + """ + + def __init__( + self, + sampler: vLLMSampler, + template: Template, + tool_manager: Optional[ToolManager] = None, + sampling_params: Optional[SamplingParams] = None, + max_turns: int = 6, + max_trajectory_tokens: Optional[int] = None, + ): + # Validation aligned with MultiTurnRollout.__init__. + if template is None: + raise ValueError('ClientMultiTurnRollout requires a local Template instance') + if max_turns < 1: + raise ValueError(f'max_turns must be >= 1, got {max_turns}') + if max_trajectory_tokens is not None and max_trajectory_tokens < 1: + raise ValueError(f'max_trajectory_tokens must be >= 1 or None, got ' + f'{max_trajectory_tokens}') + + self.sampler = sampler + self.template = template + self.tool_manager = tool_manager + self.sampling_params = sampling_params or SamplingParams() + self.max_turns = max_turns + self.max_trajectory_tokens = max_trajectory_tokens + + if self.sampling_params.num_samples != 1: + raise ValueError(f'ClientMultiTurnRollout currently supports num_samples=1 only, ' + f'got {self.sampling_params.num_samples}') + + def __call__(self, trajectories: List[Trajectory], **kwargs) -> List[Trajectory]: + """Run the batched multi-turn rollout state machine over HTTP. + + Structurally mirrors ``MultiTurnRollout.__call__`` but issues each + round's sampling through ``vLLMSampler.sample()`` (an HTTP POST to + ``/twinkle/sample``) rather than a Ray actor call. Every round makes a + SINGLE batched HTTP call for all currently-live trajectories so the + sampler can run them in parallel; finished trajectories are parked and + excluded from later batches. + + Returns a ``List[Trajectory]`` of the same length and order as the + input, each augmented with ``messages`` / ``logprobs`` / ``turns`` / + ``stop_reason`` / ``truncated`` fields. + + Exception handling and boundary truncation contract: + * ``new_input_feature`` missing / lacking ``input_ids`` -> RuntimeError + carrying both the batch index and the trajectory index. + * per-round ``len(seq.logprobs) != len(seq.tokens)`` -> RuntimeError + carrying the specific counts. + * final per-trajectory ``len(all_logprobs[i]) != count(labels != -100)`` + -> RuntimeError (protects downstream GRPO old_logps alignment). + * ``vLLMSampler.sample()`` network/timeout errors propagate unchanged + (never wrapped or swallowed). + * tool_calls produced with no ``tool_manager`` -> ValueError. + * ``max_turns == 1`` with a first-round tool call -> the trajectory is + marked ``truncated=True, stop_reason='max_turns'`` and sampling stops. + """ + if isinstance(trajectories, dict): + raise TypeError('ClientMultiTurnRollout.__call__ expects a List[Trajectory]; ' + 'wrap a single trajectory as [trajectory].') + trajectories = list(trajectories) + n = len(trajectories) + if n == 0: + return [] + + sampling_params = self._as_sampling_params_dict( + kwargs.get('sampling_params', self.sampling_params)) + tool_managers = self._resolve_tool_managers( + kwargs.get('tool_manager', self.tool_manager), n) + + # 1. Encode each trajectory once; ``pifs[i]`` is the live per-turn + # state for trajectory ``i``. ``vLLMSampler.sample`` is responsible for + # JSON-serialising the feature (ndarray / tensor -> list) before the + # HTTP POST, so no conversion is needed here. + pifs: List[Dict[str, Any]] = [] + for traj in trajectories: + pif = self.template.encode(traj, add_generation_prompt=True) + pif.setdefault('messages', list(traj.get('messages', []))) + pifs.append(pif) + + all_logprobs: List[List[Any]] = [[] for _ in range(n)] + stop_reasons: List[Optional[str]] = [None] * n + turns: List[int] = [0] * n + truncated: List[bool] = [False] * n + done: List[bool] = [False] * n + + for _ in range(self.max_turns): + active = [i for i in range(n) if not done[i]] + if not active: + break + + # 2. One batched HTTP sample call for all currently-live + # trajectories. No device_mesh / min_batch_size padding: an HTTP + # client has no Ray DP ranks to align against. + # + # Passthrough contract: ``vLLMSampler.sample()`` may raise + # network / timeout / HTTP errors (e.g. requests exceptions). We + # deliberately do NOT wrap this call in try/except -- such errors + # propagate unchanged to the caller so ret/backoff policy stays an + # upstream concern (retry/backoff) and failures are never + # silently swallowed. + batch_pifs = [pifs[i] for i in active] + resps = self.sampler.sample(batch_pifs, sampling_params=sampling_params) + + pending_bridges: List[tuple] = [] # (global_idx, tool_messages) + for local_idx, global_idx in enumerate(active): + turns[global_idx] += 1 + seq = resps[local_idx].sequences[0] + + # ``new_input_feature`` is the running pif for the next round; + # the /twinkle/sample response contract guarantees it is set and + # carries ``input_ids``. A missing feature makes the next round + # impossible, so raise a batch/trajectory-indexed RuntimeError. + if seq.new_input_feature is None or 'input_ids' not in seq.new_input_feature: + raise RuntimeError( + f'Sampler returned a sequence without new_input_feature.input_ids at ' + f'batch index {local_idx} (trajectory {global_idx}); ' + f'cannot continue multi-turn.') + + pifs[global_idx] = dict(seq.new_input_feature) + # Per-round logprobs/token alignment guard: each sampled token + # must carry exactly one logprob entry. Mirrors the core-lib + # ``len(seq.logprobs) != len(seq.tokens)`` semantic so client and + # Ray paths cannot drift on this invariant. + if seq.logprobs is not None: + if len(seq.logprobs) != len(seq.tokens): + raise RuntimeError( + f'logprobs length ({len(seq.logprobs)}) does not match sampled ' + f'token count ({len(seq.tokens)}) at turn {turns[global_idx]} ' + f'(trajectory {global_idx})') + all_logprobs[global_idx].extend(seq.logprobs) + stop_reasons[global_idx] = seq.stop_reason + + # 3. Termination conditions. + if seq.stop_reason == 'length': + done[global_idx] = True + continue + + # 3a. Sequence-length cap. + if (self.max_trajectory_tokens is not None and len( + pifs[global_idx].get('input_ids') or []) >= self.max_trajectory_tokens): + truncated[global_idx] = True + done[global_idx] = True + continue + + # 3b. Parse tool calls from the freshly sampled assistant turn. + _msgs = pifs[global_idx].get('messages') or [] + _last_msg = _msgs[-1] if _msgs else None + tool_calls = (_last_msg.get('tool_calls') if isinstance(_last_msg, dict) else None) + if not tool_calls: + tool_calls = self.template.parse_tool_call(seq.decoded or '') + if not tool_calls: + done[global_idx] = True + continue + + # 3c. Hit the turn cap while still wanting to call a tool: force + # truncation. Also covers the ``max_turns == 1`` edge, where + # the very first sampled turn trips this branch. + if turns[global_idx] >= self.max_turns: + truncated[global_idx] = True + stop_reasons[global_idx] = 'max_turns' + done[global_idx] = True + continue + + # 4. Dispatch tools for this trajectory via its ToolManager. + tool_manager = tool_managers[global_idx] + if tool_manager is None: + raise ValueError( + f'trajectory {global_idx} produced tool_calls but no tool_manager ' + f'was provided (at construction time or as a per-call kwarg).') + tool_messages = [{ + 'role': 'tool', + 'content': tool_manager(tc), + } for tc in tool_calls] + pending_bridges.append((global_idx, tool_messages)) + + # Stitch bridge tokens (tool turns + next generation prompt) for + # every trajectory with outstanding tool turns. Reuses the shared + # pure function so client and core-lib paths cannot drift. + for global_idx, tool_messages in pending_bridges: + extended = extend_with_bridge(pifs[global_idx], tool_messages, self.template) + if extended is None: + # Trajectory exceeded max_length (truncation strategy 'delete'). + truncated[global_idx] = True + done[global_idx] = True + else: + pifs[global_idx] = extended + + # 4b. Final logprobs/labels alignment guard. For every trajectory that + # collected logprobs, the total logprob count must equal the number + # of trainable positions (labels != -100) in the final pif. This is + # the same invariant grpo._pad_and_align_to_batch relies on; a + # mismatch would silently corrupt GRPO old_logps alignment, so we + # fail loudly with the specific numbers. + for i in range(n): + if not all_logprobs[i]: + continue + labels_i = pifs[i].get('labels') or [] + trainable_i = sum(1 for label in labels_i if label != -100) + if len(all_logprobs[i]) != trainable_i: + raise RuntimeError(f'logprobs/labels misaligned for trajectory {i}: ' + f'{len(all_logprobs[i])} logprobs vs {trainable_i} ' + f'trainable labels (labels != -100). This invariant is ' + f'required by grpo._pad_and_align_to_batch; a mismatch ' + f'would silently corrupt GRPO old_logps alignment.') + + # 5. Merge pif fields into each trajectory dict at TOP LEVEL, preserving + # input length and order. + outs: List[Trajectory] = [] + for i, traj in enumerate(trajectories): + out = dict(traj) + out.update(pifs[i]) + out['messages'] = list(pifs[i].get('messages') or out.get('messages', [])) + out['logprobs'] = all_logprobs[i] if all_logprobs[i] else None + out['turns'] = turns[i] + out['stop_reason'] = stop_reasons[i] + out['truncated'] = truncated[i] + outs.append(out) + return outs + + # ------------------------------------------------------------------ private + + @staticmethod + def _as_sampling_params_dict(sampling_params) -> Optional[Dict[str, Any]]: + """Coerce ``sampling_params`` into the ``Optional[Dict]`` that + ``vLLMSampler.sample()`` expects. + + ``self.sampling_params`` is a core-lib :class:`SamplingParams` dataclass, + while the HTTP sampler wants a plain dict. A per-call kwarg override may + be either a dataclass or an already-built dict. + """ + if sampling_params is None: + return None + if isinstance(sampling_params, dict): + return sampling_params + if dataclasses.is_dataclass(sampling_params): + return dataclasses.asdict(sampling_params) + return sampling_params + + @staticmethod + def _resolve_tool_managers(arg, n: int) -> List[Optional[ToolManager]]: + """Broadcast a single ``ToolManager`` or validate a per-trajectory list. + + Unlike the core-lib rollout, ``None`` is tolerated here and broadcast as + ``[None] * n``; the ValueError is raised lazily at the tool-dispatch site + only when a trajectory actually produces tool_calls. + """ + if isinstance(arg, list): + if len(arg) != n: + raise ValueError(f'per-call tool_manager list length ({len(arg)}) does ' + f'not match number of trajectories ({n})') + return list(arg) + return [arg] * n diff --git a/src/twinkle_client/sampler/__init__.py b/src/twinkle_client/sampler/__init__.py index 724d41ef..06b961b3 100644 --- a/src/twinkle_client/sampler/__init__.py +++ b/src/twinkle_client/sampler/__init__.py @@ -1,11 +1 @@ -# ============================================================================ -# WARNING: AUTO-GENERATED FILE - DO NOT MODIFY MANUALLY! -# ============================================================================ -# This file is automatically generated by client_tools/client_generator.py -# Any manual changes will be overwritten when the generator runs again. -# -# To update this file: -# 1. Modify the source files in src/twinkle/ -# 2. Run: python client_tools/client_generator.py -# ============================================================================ from .vllm_sampler import vLLMSampler diff --git a/src/twinkle_client/sampler/vllm_sampler.py b/src/twinkle_client/sampler/vllm_sampler.py index 610163b5..0d553bb3 100644 --- a/src/twinkle_client/sampler/vllm_sampler.py +++ b/src/twinkle_client/sampler/vllm_sampler.py @@ -1,13 +1,3 @@ -# ============================================================================ -# WARNING: AUTO-GENERATED FILE - DO NOT MODIFY MANUALLY! -# ============================================================================ -# This file is automatically generated by client_tools/client_generator.py -# Any manual changes will be overwritten when the generator runs again. -# -# To update this file: -# 1. Modify the source files in src/twinkle/ -# 2. Run: python client_tools/client_generator.py -# ============================================================================ from typing import Any, Dict, List, Optional, Union from twinkle_client.http import http_post from twinkle_client.types.sampler import AddAdapterResponse, SampleResponseModel, SetTemplateResponse @@ -18,6 +8,28 @@ # Intentionally does NOT subclass ``twinkle.sampler.base.Sampler``: importing # that base pulls ``twinkle.sampler.__init__`` → ``VLLMEngine`` → torch + zmq, # which the mock / CPU-only client environments don't have. +def _json_safe(obj: Any) -> Any: + """Recursively coerce numpy arrays / torch tensors to JSON-serialisable lists. + + ``sample()`` accepts pre-encoded ``InputFeature`` dicts (e.g. from a multi-turn + rollout's ``template.encode``) whose values are numpy arrays or torch tensors; + these are not JSON-serialisable and would break the HTTP POST. Detection is by + 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 + + class vLLMSampler: """Client wrapper for Sampler that calls server HTTP endpoints. @@ -74,7 +86,7 @@ def sample( SampleResponseModel with 'sequences' list, each containing tokens, logprobs, stop_reason. """ json_data = { - 'inputs': inputs, + 'inputs': _json_safe(inputs), 'sampling_params': sampling_params, 'adapter_name': adapter_name, 'num_samples': num_samples, diff --git a/tests/server/config/test_server_config.py b/tests/server/config/test_server_config.py index d7a7f296..93644cbe 100644 --- a/tests/server/config/test_server_config.py +++ b/tests/server/config/test_server_config.py @@ -2,14 +2,10 @@ """Property + unit tests for the typed ``ServerConfig``. Properties covered: -- # Feature: server-config-observability-refactor, - Property 12: Valid configuration yields a fully validated instance -- # Feature: server-config-observability-refactor, - Property 13: Any constraint violation is rejected with the offending field named -- # Feature: server-config-observability-refactor, - Property 14: Configuration round-trip fidelity -- # Feature: server-config-observability-refactor, - Property 15: Legacy / unknown field names are rejected +- Valid configuration yields a fully validated instance +- Any constraint violation is rejected with the offending field named +- Configuration round-trip fidelity +- Legacy / unknown field names are rejected """ from __future__ import annotations @@ -66,12 +62,12 @@ 'applications': st.lists(_MODEL_APP, min_size=0, max_size=3), }) -# ---------- Property 12: valid → fully validated ----------------------------- # +# ---------- valid → fully validated ---------------------------------------- # @settings(max_examples=100) @given(payload=_VALID_CONFIG) -def test_property_12_valid_payload_yields_full_instance(payload: dict) -> None: +def test_valid_payload_yields_full_instance(payload: dict) -> None: cfg = ServerConfig.model_validate(payload) assert isinstance(cfg, ServerConfig) assert all(isinstance(a, ApplicationSpec) for a in cfg.applications) @@ -80,17 +76,17 @@ def test_property_12_valid_payload_yields_full_instance(payload: dict) -> None: assert cfg.task_queue.rps_limit >= 0 -# ---------- Property 13: violation → field-named error ---------------------- # +# ---------- violation → field-named error ---------------------------------- # -def test_property_13_redis_mode_missing_url() -> None: +def test_redis_mode_missing_url() -> None: with pytest.raises(ValidationError) as exc: ServerConfig.model_validate({'persistence': {'mode': 'redis'}}) msg = str(exc.value) assert 'persistence.redis_url' in msg or 'redis_url' in msg -def test_property_13_file_mode_missing_path() -> None: +def test_file_mode_missing_path() -> None: with pytest.raises(ValidationError) as exc: ServerConfig.model_validate({'persistence': {'mode': 'file'}}) msg = str(exc.value) @@ -99,7 +95,7 @@ def test_property_13_file_mode_missing_path() -> None: @settings(max_examples=100) @given(bad_backend=st.text(min_size=1, max_size=8).filter(lambda s: s not in ('mock', 'transformers', 'megatron'))) -def test_property_13_bad_backend_names_field(bad_backend: str) -> None: +def test_bad_backend_names_field(bad_backend: str) -> None: payload = { 'applications': [{ 'name': 'm', @@ -120,7 +116,7 @@ def test_property_13_bad_backend_names_field(bad_backend: str) -> None: @settings(max_examples=100) @given(bad_max_input_tokens=st.integers(max_value=0, min_value=-1000)) -def test_property_13_nested_field_constraint_violation_named(bad_max_input_tokens: int) -> None: +def test_nested_field_constraint_violation_named(bad_max_input_tokens: int) -> None: """Nested-section constraints (here ``task_queue.max_input_tokens``) are enforced together with cross-field ones and the offending path is visible in the error.""" @@ -135,7 +131,7 @@ def test_property_13_nested_field_constraint_violation_named(bad_max_input_token @settings(max_examples=100) @given(payload=_VALID_CONFIG) -def test_property_14_round_trip_fidelity(payload: dict) -> None: +def test_round_trip_fidelity(payload: dict) -> None: cfg = ServerConfig.model_validate(payload) dumped = cfg.to_yaml_dict() re_loaded = ServerConfig.model_validate(dumped) @@ -143,14 +139,14 @@ def test_property_14_round_trip_fidelity(payload: dict) -> None: assert re_loaded.model_dump() == cfg.model_dump() -# ---------- Property 15: legacy/unknown rejected ----------------------------- # +# ---------- legacy/unknown rejected ---------------------------------------- # @pytest.mark.parametrize( 'legacy_field', ['telemetry_config', 'persistence_config'], ) -def test_property_15_legacy_field_rejected(legacy_field: str) -> None: +def test_legacy_field_rejected(legacy_field: str) -> None: payload = {legacy_field: {}} with pytest.raises(ValidationError) as exc: ServerConfig.model_validate(payload) @@ -161,7 +157,7 @@ def test_property_15_legacy_field_rejected(legacy_field: str) -> None: @settings(max_examples=100) @given(unknown=st.text(min_size=1, max_size=20).filter(lambda s: not s.startswith('_'))) -def test_property_15_unknown_field_rejected(unknown: str) -> None: +def test_unknown_field_rejected(unknown: str) -> None: known = { 'ray_namespace', 'proxy_location', @@ -178,7 +174,7 @@ def test_property_15_unknown_field_rejected(unknown: str) -> None: @pytest.mark.parametrize('section', ['telemetry', 'persistence']) -def test_property_15_unknown_nested_field_rejected(section: str) -> None: +def test_unknown_nested_field_rejected(section: str) -> None: """Nested config sections also reject unknown keys (defends against typos inside ``telemetry: {...}`` / ``persistence: {...}``).""" payload = {section: {'unknown_typo': 1}} @@ -187,7 +183,7 @@ def test_property_15_unknown_nested_field_rejected(section: str) -> None: assert any('unknown_typo' in err['loc'] for err in exc.value.errors()) -# ---------- 3.11: from_yaml error paths + launcher dict rejection ---------- # +# ---------- from_yaml error paths + launcher dict rejection ---------------- # def test_from_yaml_missing_path(tmp_path: Path) -> None: diff --git a/tests/server/integration/test_mock_mode_startup.py b/tests/server/integration/test_mock_mode_startup.py index 0e793643..be8610ce 100644 --- a/tests/server/integration/test_mock_mode_startup.py +++ b/tests/server/integration/test_mock_mode_startup.py @@ -152,7 +152,7 @@ def test_mock_mode_reaches_ready_under_30s_and_is_deterministic(ray_cluster) -> # dict — this mirrors what the production launcher does at # ``launcher/server_launcher.py:161`` and is required since the # builders + deployment ``__init__`` accept ``TaskQueueConfig`` directly - # (Task 27 removed the ``from_dict`` revival path). + # (the ``from_dict`` revival path was removed). args = {k: v for k, v in dict(app_spec.args).items() if v is not None} if app_spec.import_path == 'server': # Gateway's ServiceProxy reads http_options.port to build internal diff --git a/tests/server/model/test_mock_model.py b/tests/server/model/test_mock_model.py index 26ef51a9..5c24450f 100644 --- a/tests/server/model/test_mock_model.py +++ b/tests/server/model/test_mock_model.py @@ -58,12 +58,12 @@ @pytest.mark.parametrize('method_name', _REQUIRED_METHODS) -def test_property_1_required_method_present(method_name: str) -> None: +def test_required_method_present(method_name: str) -> None: m = TwinkleCompatMockModel('mid') assert callable(getattr(m, method_name)), method_name -def test_property_1_constructor_does_not_raise() -> None: +def test_constructor_does_not_raise() -> None: TwinkleCompatMockModel('mid') @@ -75,7 +75,7 @@ def test_property_1_constructor_does_not_raise() -> None: seq_lens=st.lists(st.integers(min_value=1, max_value=12), min_size=1, max_size=5), seed=st.integers(min_value=0, max_value=99), ) -def test_property_2_forward_only_deterministic_and_shaped(seq_lens: list, seed: int) -> None: +def test_forward_only_deterministic_and_shaped(seq_lens: list, seed: int) -> None: inputs = [{'tokens': list(range(n))} for n in seq_lens] a = TwinkleCompatMockModel('mid', seed=seed) b = TwinkleCompatMockModel('mid', seed=seed) @@ -90,7 +90,7 @@ def test_property_2_forward_only_deterministic_and_shaped(seq_lens: list, seed: @settings(max_examples=100) @given(seq_lens=st.lists(st.integers(min_value=1, max_value=8), min_size=1, max_size=4)) -def test_property_2_tinker_forward_backward_loss_is_finite(seq_lens: list) -> None: +def test_tinker_forward_backward_loss_is_finite(seq_lens: list) -> None: m = TwinkleCompatMockModel('mid') inputs = [{'tokens': list(range(n))} for n in seq_lens] result, loss = m.tinker_forward_backward(inputs=inputs, adapter_name='a', loss_fn='cross_entropy') @@ -106,7 +106,7 @@ def test_property_2_tinker_forward_backward_loss_is_finite(seq_lens: list) -> No @given( name=st.text( min_size=1, max_size=12, alphabet=st.characters(whitelist_categories=('L', 'N'), whitelist_characters='_-'))) -def test_property_3_adapter_add_remove_round_trip(name: str) -> None: +def test_adapter_add_remove_round_trip(name: str) -> None: m = TwinkleCompatMockModel('mid') assert not m.has_adapter(name) m.add_adapter(name, rank=4) @@ -120,7 +120,7 @@ def test_property_3_adapter_add_remove_round_trip(name: str) -> None: @settings(max_examples=100) @given(name=st.text(min_size=1, max_size=12)) -def test_property_4_remove_absent_raises(name: str) -> None: +def test_remove_absent_raises(name: str) -> None: m = TwinkleCompatMockModel('mid') pre = dict(m._adapters) with pytest.raises(KeyError): @@ -131,14 +131,14 @@ def test_property_4_remove_absent_raises(name: str) -> None: # ---------- Model backend dispatch ---------------------------------------- # -def test_property_10_mock_dispatch_returns_mock_model() -> None: +def test_mock_dispatch_returns_mock_model() -> None: m = MODEL_SELECTOR.construct(MODEL_SELECTOR.validate('mock'), {'model_id': 'mid'}) assert isinstance(m, TwinkleCompatMockModel) @settings(max_examples=100) @given(bad=st.text(min_size=1, max_size=10).filter(lambda s: s not in _MODEL_BACKENDS)) -def test_property_10_invalid_backend_raises_config_error(bad: str) -> None: +def test_invalid_backend_raises_config_error(bad: str) -> None: """Validation runs BEFORE any backend import / instantiation.""" with pytest.raises(ConfigError) as exc: MODEL_SELECTOR.validate(bad) @@ -148,7 +148,7 @@ def test_property_10_invalid_backend_raises_config_error(bad: str) -> None: @pytest.mark.parametrize('value', [None, '']) -def test_property_10_absent_or_empty_backend_raises(value) -> None: +def test_absent_or_empty_backend_raises(value) -> None: with pytest.raises(ConfigError) as exc: MODEL_SELECTOR.validate(value) assert exc.value.field == 'backend' diff --git a/tests/server/sampler/test_mock_sampler.py b/tests/server/sampler/test_mock_sampler.py index ea1c0766..afa72dda 100644 --- a/tests/server/sampler/test_mock_sampler.py +++ b/tests/server/sampler/test_mock_sampler.py @@ -29,7 +29,7 @@ @pytest.mark.parametrize('method', _REQUIRED_METHODS) -def test_property_5_required_method_present(method: str) -> None: +def test_required_method_present(method: str) -> None: s = MockSampler('mid') assert callable(getattr(s, method)) @@ -42,7 +42,7 @@ def test_property_5_required_method_present(method: str) -> None: max_tokens=st.integers(min_value=1, max_value=20), num_samples=st.integers(min_value=1, max_value=4), ) -def test_property_6_output_length_and_logprob_count(max_tokens: int, num_samples: int) -> None: +def test_output_length_and_logprob_count(max_tokens: int, num_samples: int) -> None: s = MockSampler('mid') inp = InputFeature(input_ids=[1, 2, 3]) responses = s.sample(inp, SamplingParams(max_tokens=max_tokens), adapter_name='a', num_samples=num_samples) @@ -63,7 +63,7 @@ def test_property_6_output_length_and_logprob_count(max_tokens: int, num_samples num_samples=st.integers(min_value=1, max_value=3), adapter=st.sampled_from(['', 'a', 'lora-1']), ) -def test_property_7_determinism(max_tokens: int, num_samples: int, adapter: str) -> None: +def test_determinism(max_tokens: int, num_samples: int, adapter: str) -> None: s = MockSampler('mid', seed=42) inp = InputFeature(input_ids=[1, 2, 3]) r1 = s.sample(inp, SamplingParams(max_tokens=max_tokens), adapter_name=adapter, num_samples=num_samples) @@ -76,7 +76,7 @@ def test_property_7_determinism(max_tokens: int, num_samples: int, adapter: str) @settings(max_examples=50) @given(bad=st.integers(max_value=0, min_value=-1000)) -def test_property_8_max_tokens_lt_1_raises(bad: int) -> None: +def test_max_tokens_lt_1_raises(bad: int) -> None: s = MockSampler('mid') inp = InputFeature(input_ids=[1]) with pytest.raises(ValueError) as exc: @@ -84,7 +84,7 @@ def test_property_8_max_tokens_lt_1_raises(bad: int) -> None: assert 'max_tokens' in str(exc.value) -def test_property_8_no_sampling_params_raises() -> None: +def test_no_sampling_params_raises() -> None: s = MockSampler('mid') inp = InputFeature(input_ids=[1]) with pytest.raises(ValueError): @@ -98,7 +98,7 @@ def test_property_8_no_sampling_params_raises() -> None: @given( name=st.text( min_size=1, max_size=12, alphabet=st.characters(whitelist_categories=('L', 'N'), whitelist_characters='_-'))) -def test_property_9_add_adapter_to_sampler(name: str) -> None: +def test_add_adapter_to_sampler(name: str) -> None: s = MockSampler('mid') assert not s.has_adapter(name) s.add_adapter_to_sampler(name, {'rank': 4}) @@ -109,14 +109,14 @@ def test_property_9_add_adapter_to_sampler(name: str) -> None: # ---------- Sampler backend dispatch -------------------------------------- # -def test_property_11_mock_dispatch_returns_mock_sampler() -> None: +def test_mock_dispatch_returns_mock_sampler() -> None: s = SAMPLER_SELECTOR.construct(SAMPLER_SELECTOR.validate('mock'), {'model_id': 'mid'}) assert isinstance(s, MockSampler) @settings(max_examples=100) @given(bad=st.text(min_size=1, max_size=10).filter(lambda s: s not in _SAMPLER_TYPES)) -def test_property_11_invalid_sampler_type_raises_config_error(bad: str) -> None: +def test_invalid_sampler_type_raises_config_error(bad: str) -> None: """Validation runs BEFORE any sampler import / instantiation.""" with pytest.raises(ConfigError) as exc: SAMPLER_SELECTOR.validate(bad) @@ -126,7 +126,7 @@ def test_property_11_invalid_sampler_type_raises_config_error(bad: str) -> None: @pytest.mark.parametrize('value', [None, '']) -def test_property_11_absent_or_empty_sampler_type_raises(value) -> None: +def test_absent_or_empty_sampler_type_raises(value) -> None: with pytest.raises(ConfigError) as exc: SAMPLER_SELECTOR.validate(value) assert exc.value.field == 'sampler_type' @@ -160,3 +160,63 @@ def test_sample_stream_rejects_bad_max_tokens() -> None: inp = InputFeature(input_ids=[1]) with pytest.raises(ValueError): list(s.sample_stream(inp, SamplingParams(max_tokens=0))) + + +# ---------- Multi-turn contract knobs -------------------------- # + + +def test_default_new_input_feature_appends_sampled_tokens() -> None: + """Default behaviour now carries a non-empty new_input_feature that appends + the round's sampled tokens onto the input pif's input_ids.""" + s = MockSampler('mid') + inp = InputFeature(input_ids=[1, 2, 3]) + seq = s.sample(inp, SamplingParams(max_tokens=4))[0].sequences[0] + assert seq.new_input_feature is not None + assert 'input_ids' in seq.new_input_feature + assert seq.new_input_feature['input_ids'] == [1, 2, 3] + list(seq.tokens) + # Prior context stays non-trainable; sampled tokens are trainable (their ids). + assert seq.new_input_feature['labels'] == [-100, -100, -100] + list(seq.tokens) + assert seq.new_input_feature['length'] == 3 + len(seq.tokens) + + +def test_default_backward_compatible_stop_reason_and_decoded() -> None: + """With no knobs configured, stop_reason stays 'length' and decoded None.""" + s = MockSampler('mid') + inp = InputFeature(input_ids=[1]) + seq = s.sample(inp, SamplingParams(max_tokens=2))[0].sequences[0] + assert seq.stop_reason == 'length' + assert seq.decoded is None + + +def test_configurable_stop_reason_via_ctor() -> None: + s = MockSampler('mid', stop_reason='stop') + inp = InputFeature(input_ids=[1]) + seq = s.sample(inp, SamplingParams(max_tokens=2))[0].sequences[0] + assert seq.stop_reason == 'stop' + + +def test_stop_reason_per_call_kwarg_overrides_ctor() -> None: + s = MockSampler('mid', stop_reason='stop') + inp = InputFeature(input_ids=[1]) + seq = s.sample(inp, SamplingParams(max_tokens=2), stop_reason='length')[0].sequences[0] + assert seq.stop_reason == 'length' + + +def test_tool_call_text_injected_only_on_configured_turns() -> None: + """tool_call_text is emitted as decoded on turn 1 only (default), driving a + 'sample -> tool -> next round -> terminate' loop.""" + s = MockSampler('mid', stop_reason='stop', tool_call_text='x') + inp = InputFeature(input_ids=[1, 2]) + # Round 1: tool call injected. + seq1 = s.sample(inp, SamplingParams(max_tokens=2))[0].sequences[0] + assert seq1.decoded == 'x' + # Round 2: no injection -> decoded None (loop can terminate on empty parse). + seq2 = s.sample(inp, SamplingParams(max_tokens=2))[0].sequences[0] + assert seq2.decoded is None + + +def test_tool_call_turns_multiple_rounds() -> None: + s = MockSampler('mid', stop_reason='stop', tool_call_text='TC', tool_call_turns=(1, 3)) + inp = InputFeature(input_ids=[1]) + decoded = [s.sample(inp, SamplingParams(max_tokens=1))[0].sequences[0].decoded for _ in range(3)] + assert decoded == ['TC', None, 'TC'] diff --git a/tests/server/state/test_leader_election.py b/tests/server/state/test_leader_election.py index 57891e1d..f1f8955c 100644 --- a/tests/server/state/test_leader_election.py +++ b/tests/server/state/test_leader_election.py @@ -1,6 +1,6 @@ """Tests for the cleanup-leader election loop and metrics gauge in ``ServerState``. -Validates: +Covers: - exactly one ``ServerState`` instance over the same backend becomes leader; - the cleanup task is gated on leadership (non-leaders never call cleanup); - when the leader stops, another instance can take over; @@ -100,7 +100,7 @@ async def test_renew_keeps_leader() -> None: @pytest.mark.asyncio async def test_leader_recovers_after_renewal_failure() -> None: - """Regression (Requirement 19): when a leader's renewal raises, it releases + """Regression: when a leader's renewal raises, it releases the lease best-effort so the very next election tick re-acquires leadership without waiting LEASE_TTL. @@ -140,7 +140,7 @@ async def flaky_update_atomic(*args, **kwargs): @pytest.mark.asyncio async def test_renewal_failure_does_not_steal_other_leader_lease() -> None: """A non-leader whose election attempt raises must NOT delete a lease that - another replica legitimately holds (Requirement 19.3).""" + another replica legitimately holds.""" backend = RayActorBackend() leader = ServerState(backend=backend, cleanup_interval=600.0) follower = ServerState(backend=backend, cleanup_interval=600.0) diff --git a/tests/server/state/test_redis_integration.py b/tests/server/state/test_redis_integration.py index d5a49fa8..e730703a 100644 --- a/tests/server/state/test_redis_integration.py +++ b/tests/server/state/test_redis_integration.py @@ -104,7 +104,7 @@ async def _cleanup() -> None: @pytest.mark.asyncio -async def test_property_26_replica_write_via_a_visible_via_b(make_state) -> None: +async def test_replica_write_via_a_visible_via_b(make_state) -> None: """One worker registers a replica; a second worker on the same shared backend sees the same capacity / availability view.""" a = make_state() @@ -118,7 +118,7 @@ async def test_property_26_replica_write_via_a_visible_via_b(make_state) -> None @pytest.mark.asyncio -async def test_property_26_model_write_visible(make_state) -> None: +async def test_model_write_visible(make_state) -> None: a = make_state() b = make_state() rid = f'r-{uuid.uuid4().hex[:6]}' @@ -132,7 +132,7 @@ async def test_property_26_model_write_visible(make_state) -> None: @pytest.mark.asyncio -async def test_property_26_session_and_config(make_state) -> None: +async def test_session_and_config(make_state) -> None: a = make_state() b = make_state() sid = await a.create_session({'session_id': f'sess-{uuid.uuid4().hex[:6]}'}) @@ -146,7 +146,7 @@ async def test_property_26_session_and_config(make_state) -> None: @pytest.mark.asyncio -async def test_property_27_concurrent_config_writes_no_torn_records(make_state) -> None: +async def test_concurrent_config_writes_no_torn_records(make_state) -> None: """Many concurrent writes of distinct keys complete and every record equals one of the writes (no torn / partial value).""" a = make_state() @@ -168,7 +168,7 @@ async def writer(state: ServerState, items: dict) -> None: @pytest.mark.asyncio -async def test_property_27_concurrent_same_key_lands_one_of_committed(make_state) -> None: +async def test_concurrent_same_key_lands_one_of_committed(make_state) -> None: """Two writers race on the same key — final value equals one of the writes; no torn record.""" a = make_state() @@ -182,7 +182,7 @@ async def test_property_27_concurrent_same_key_lands_one_of_committed(make_state @pytest.mark.asyncio -async def test_property_27_concurrent_replica_registration(make_state) -> None: +async def test_concurrent_replica_registration(make_state) -> None: a = make_state() b = make_state() rid = f'r-{uuid.uuid4().hex[:6]}' diff --git a/tests/server/telemetry/test_tracing_and_correlation.py b/tests/server/telemetry/test_tracing_and_correlation.py index 01b43b0a..e4120421 100644 --- a/tests/server/telemetry/test_tracing_and_correlation.py +++ b/tests/server/telemetry/test_tracing_and_correlation.py @@ -3,11 +3,11 @@ keys, and ``ResourceMetricsCollector``. Properties covered: -- # Feature: server-config-observability-refactor, Property 19: Business-layer span lifecycle -- # Feature: server-config-observability-refactor, Property 20: Span exception handling -- # Feature: server-config-observability-refactor, Property 21: Tracing graceful-degradation equivalence -- # Feature: server-config-observability-refactor, Property 22: Correlation attribute attachment -- # Feature: server-config-observability-refactor, Property 23: Correlation prefix invariant +- Business-layer span lifecycle +- Span exception handling +- Tracing graceful-degradation equivalence +- Correlation attribute attachment +- Correlation prefix invariant """ from __future__ import annotations @@ -24,11 +24,11 @@ @pytest.mark.parametrize('key', CORRELATION_KEYS) -def test_property_23_prefix_invariant(key: str) -> None: +def test_prefix_invariant(key: str) -> None: assert key.startswith(PREFIX), key -def test_property_23_helper_constants_complete() -> None: +def test_helper_constants_complete() -> None: expected = { 'twinkle.session_id', 'twinkle.model_id', @@ -40,7 +40,7 @@ def test_property_23_helper_constants_complete() -> None: assert set(CORRELATION_KEYS) == expected -# ---------- Property 22: attachment of present-only values ------------------ # +# ---------- attachment of present-only values ------------------------------ # class _RecordingSpan: @@ -63,14 +63,14 @@ def set_attribute(self, key: str, value: object) -> None: correlation.TOKEN_ID: st.one_of(st.none(), st.text(min_size=1, max_size=8)), }, )) -def test_property_22_set_correlation_attrs_skips_none(payload: dict) -> None: +def test_set_correlation_attrs_skips_none(payload: dict) -> None: span = _RecordingSpan() set_correlation_attrs(span, payload) expected = {k: v for k, v in payload.items() if v is not None} assert span.attrs == expected -def test_property_22_noop_span_safe() -> None: +def test_noop_span_safe() -> None: """``set_correlation_attrs`` is a no-op on a NoOp span (no SDK installed).""" span = _NoopSpan() set_correlation_attrs(span, {correlation.SESSION_ID: 's1'}) @@ -79,10 +79,10 @@ def test_property_22_noop_span_safe() -> None: set_correlation_attrs(span, None) -# ---------- Property 21: NoOp degradation equivalence ---------------------- # +# ---------- NoOp degradation equivalence ----------------------------------- # -def test_property_21_noop_yields_same_result_as_active() -> None: +def test_noop_yields_same_result_as_active() -> None: """When OTEL is absent, ``traced_operation`` runs the body and returns the body's result identically to when OTEL is active.""" with mock.patch('twinkle.server.telemetry.tracing._OTEL_AVAILABLE', False): @@ -92,7 +92,7 @@ def test_property_21_noop_yields_same_result_as_active() -> None: assert result == 10 -def test_property_21_noop_propagates_exceptions() -> None: +def test_noop_propagates_exceptions() -> None: """NoOp path still re-raises the original exception unchanged.""" with mock.patch('twinkle.server.telemetry.tracing._OTEL_AVAILABLE', False): with pytest.raises(RuntimeError, match='boom'): @@ -100,7 +100,7 @@ def test_property_21_noop_propagates_exceptions() -> None: raise RuntimeError('boom') -# ---------- Property 19/20: span lifecycle + exception handling ----------- # +# ---------- span lifecycle + exception handling ---------------------------- # def _otel_available() -> bool: @@ -111,7 +111,7 @@ def _otel_available() -> bool: return True -def test_property_19_span_lifecycle(in_memory_span_exporter) -> None: +def test_span_lifecycle(in_memory_span_exporter) -> None: """When OTEL is present, a span is started before and ended after the block.""" in_memory_span_exporter.clear() with mock.patch('twinkle.server.telemetry.tracing._OTEL_AVAILABLE', True): @@ -124,7 +124,7 @@ def test_property_19_span_lifecycle(in_memory_span_exporter) -> None: assert matches[-1].attributes.get(correlation.SESSION_ID) == 's1' -def test_property_20_exception_recorded_and_reraised(in_memory_span_exporter) -> None: +def test_exception_recorded_and_reraised(in_memory_span_exporter) -> None: """Exception inside the block is recorded on the span and re-raised.""" in_memory_span_exporter.clear() with mock.patch('twinkle.server.telemetry.tracing._OTEL_AVAILABLE', True): @@ -268,8 +268,6 @@ def test_inbound_traceparent_continues_trace() -> None: runs nested inside the middleware's SERVER span), which makes the assertion independent of whichever global tracer provider another test may have installed. - - Validates Requirement 15: inbound HTTP trace-context continuity. """ if not _otel_available(): pytest.skip('OTEL SDK not installed in test env') @@ -304,7 +302,7 @@ async def ping() -> dict: def test_inbound_without_traceparent_starts_new_trace() -> None: """A request with no trace headers still completes without raising and runs - the handler under a span on a fresh trace (Requirement 15.3).""" + the handler under a span on a fresh trace.""" if not _otel_available(): pytest.skip('OTEL SDK not installed in test env') diff --git a/tests/server/utils/task_queue/test_config.py b/tests/server/utils/task_queue/test_config.py index 7c818473..3126c507 100644 --- a/tests/server/utils/task_queue/test_config.py +++ b/tests/server/utils/task_queue/test_config.py @@ -3,7 +3,7 @@ Pins constraint enforcement, default-value defaulting, and ``extra='forbid'`` behavior. The class is constructed directly from validated YAML by -``ApplicationSpec`` (typed end-to-end after Task 27), so there is no +``ApplicationSpec`` (typed end-to-end), so there is no ``from_dict`` revival path to exercise. """ from __future__ import annotations @@ -26,7 +26,7 @@ 'max_input_tokens': 16000, } -# ---------- Property 16: constraint enforcement ----------------------------- # +# ---------- constraint enforcement ----------------------------------------- # _CONSTRAINED_GE0_FLOATS = ['rps_limit', 'tps_limit', 'queue_timeout', 'token_cleanup_interval'] @@ -36,7 +36,7 @@ field=st.sampled_from(_CONSTRAINED_GE0_FLOATS), bad_value=st.floats(max_value=-1e-6, min_value=-1e6, allow_nan=False, allow_infinity=False), ) -def test_property_16_ge0_floats_reject_negative(field: str, bad_value: float) -> None: +def test_ge0_floats_reject_negative(field: str, bad_value: float) -> None: """Non-negative float fields reject any negative input.""" with pytest.raises(ValidationError) as exc: TaskQueueConfig(**{field: bad_value}) @@ -45,7 +45,7 @@ def test_property_16_ge0_floats_reject_negative(field: str, bad_value: float) -> @settings(max_examples=100) @given(bad_value=st.floats(max_value=0.0, min_value=-1e6, allow_nan=False, allow_infinity=False)) -def test_property_16_window_seconds_rejects_zero_and_negative(bad_value: float) -> None: +def test_window_seconds_rejects_zero_and_negative(bad_value: float) -> None: """``window_seconds`` must be strictly > 0.""" with pytest.raises(ValidationError) as exc: TaskQueueConfig(window_seconds=bad_value) @@ -54,7 +54,7 @@ def test_property_16_window_seconds_rejects_zero_and_negative(bad_value: float) @settings(max_examples=100) @given(bad_value=st.integers(max_value=0, min_value=-1_000_000)) -def test_property_16_max_input_tokens_rejects_lt_1(bad_value: int) -> None: +def test_max_input_tokens_rejects_lt_1(bad_value: int) -> None: """``max_input_tokens`` must be an integer ≥ 1.""" with pytest.raises(ValidationError) as exc: TaskQueueConfig(max_input_tokens=bad_value) @@ -70,8 +70,8 @@ def test_property_16_max_input_tokens_rejects_lt_1(bad_value: int) -> None: cleanup=st.floats(min_value=0.0, max_value=1e6, allow_nan=False, allow_infinity=False), mit=st.integers(min_value=1, max_value=10_000_000), ) -def test_property_16_valid_values_accepted(rps: float, tps: float, win: float, qt: float, cleanup: float, - mit: int) -> None: +def test_valid_values_accepted(rps: float, tps: float, win: float, qt: float, cleanup: float, + mit: int) -> None: """Any value satisfying the constraints constructs successfully.""" cfg = TaskQueueConfig( rps_limit=rps, diff --git a/tests/twinkle_agentic/test_bridge.py b/tests/twinkle_agentic/test_bridge.py new file mode 100644 index 00000000..c3c9155e --- /dev/null +++ b/tests/twinkle_agentic/test_bridge.py @@ -0,0 +1,253 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Unit tests for :func:`twinkle_agentic.rollout.bridge.extend_with_bridge`. + +These tests target the pure, ``self``-free bridge-stitching function directly +(rather than through ``MultiTurnRollout``). They exercise: + + - normal bridge append: tool messages + next generation prompt are + appended as ``-100`` bridge tokens, history is preserved verbatim, and + ``messages`` is updated to ``messages_before + tool_messages``. + - RuntimeError when the template's ``s_after`` rendering is NOT a + prefix-extension of ``s_before`` (non-monotonic template). + - RuntimeError when the computed bridge text is empty (template adds no + tokens for the tool turn). + - RuntimeError when ``labels`` length != ``input_ids`` length in the pif. + +The fakes mirror the char-level FakeTokenizer / FakeTemplate infrastructure in +``test_multi_turn_rollout.py``: ``decode(encode(s)) == s`` for any mix of raw +chars and registered specials, and ``_invoke_post_pipeline`` replays the +label-roll / attention-mask semantics the real Template applies. +""" +from __future__ import annotations + +from typing import Any + +import pytest + +from twinkle_agentic.rollout.bridge import extend_with_bridge + + +# ============================================================================= +# Fakes (mirrors test_multi_turn_rollout.py) +# ============================================================================= +class FakeTokenizer: + """Char-level tokenizer with atomic special tokens. + + Guarantees ``decode(encode(s)) == s`` for any mix of raw chars and + registered specials, which keeps the template-space bridge diff exact. + """ + SPECIALS = ('<|im_start|>', '<|im_end|>') + + def __init__(self) -> None: + self._s2i: dict[str, int] = {} + self._i2s: dict[int, str] = {} + for s in self.SPECIALS: + self._add(s) + + def _add(self, tok: str) -> int: + if tok not in self._s2i: + i = len(self._s2i) + self._s2i[tok] = i + self._i2s[i] = tok + return self._s2i[tok] + + def encode(self, text: str, add_special_tokens: bool = False) -> list[int]: + ids: list[int] = [] + i = 0 + while i < len(text): + matched = False + for sp in self.SPECIALS: + if text.startswith(sp, i): + ids.append(self._add(sp)) + i += len(sp) + matched = True + break + if not matched: + ids.append(self._add(text[i])) + i += 1 + return ids + + def decode(self, ids: list[int], skip_special_tokens: bool = False) -> str: + specials = set(self.SPECIALS) + toks = [self._i2s[int(i)] for i in ids] + if skip_special_tokens: + toks = [t for t in toks if t not in specials] + return ''.join(toks) + + def apply_chat_template( + self, + messages: list[dict[str, Any]], + tokenize: bool = False, + add_generation_prompt: bool = False, + **_, + ): + s = '' + for m in messages: + role = m['role'] + content = m['content'] + s += f'<|im_start|>{role}\n{content}<|im_end|>\n' + if add_generation_prompt: + s += '<|im_start|>assistant\n' + if tokenize: + return self.encode(s) + return s + + +class NonMonotonicTokenizer(FakeTokenizer): + """Renders a length-prefix that changes with message count. + + Because the leading ``[]`` marker differs between ``messages_before`` + and ``messages_after``, ``s_after`` is NOT a prefix-extension of + ``s_before`` — this is the non-monotonic template case that + ``extend_with_bridge`` must reject. + """ + + def apply_chat_template(self, messages, tokenize=False, add_generation_prompt=False, **_): + s = f'[{len(messages)}]' + s += super().apply_chat_template(messages, tokenize=False, add_generation_prompt=add_generation_prompt) + if tokenize: + return self.encode(s) + return s + + +class EmptyBridgeTokenizer(FakeTokenizer): + """Always renders the same constant string regardless of messages. + + ``s_after == s_before`` ⇒ the computed bridge text is empty, which + ``extend_with_bridge`` must reject. + """ + + def apply_chat_template(self, messages, tokenize=False, add_generation_prompt=False, **_): + s = 'CONSTANT' + if tokenize: + return self.encode(s) + return s + + +class FakeTemplate: + """Minimal Template mirroring the parts ``extend_with_bridge`` touches.""" + model_id = 'qwen-fake' + truncation_strategy = 'right' + + def __init__(self, tokenizer: FakeTokenizer) -> None: + self.tokenizer = tokenizer + + def encode(self, trajectory: dict[str, Any], add_generation_prompt: bool = False) -> dict[str, Any]: + messages = trajectory.get('messages', []) + s = self.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=add_generation_prompt) + input_ids = self.tokenizer.encode(s, add_special_tokens=False) + pif: dict[str, Any] = dict(trajectory) + pif['input_ids'] = input_ids + pif['labels'] = [-100] * len(input_ids) + return self._invoke_post_pipeline([pif])[0] + + def _invoke_post_pipeline(self, inputs: list[dict[str, Any]]) -> list[dict[str, Any]]: + out = [] + for pif in inputs: + pif = dict(pif) + input_ids = list(pif['input_ids']) + labels = list(pif.get('labels') or []) + if labels: + if len(labels) != len(input_ids): + raise RuntimeError(f'FakeTemplate post_pipeline: labels({len(labels)}) ' + f'!= input_ids({len(input_ids)})') + # np.roll(labels, -1): shift LEFT by 1 (output/shifted order) + labels = labels[1:] + labels[:1] + pif['input_ids'] = input_ids + pif['labels'] = labels + pif['attention_mask'] = [1] * len(input_ids) + pif['position_ids'] = list(range(len(input_ids))) + pif['length'] = len(input_ids) + out.append(pif) + return out + + +# ============================================================================= +# Helpers +# ============================================================================= +def _make_pif(template: FakeTemplate, messages: list[dict[str, Any]]) -> dict[str, Any]: + """Build a post-pipeline pif for ``messages`` in inference mode.""" + return template.encode({'messages': list(messages)}, add_generation_prompt=True) + + +def _count_trainable(labels: list[int]) -> int: + return sum(1 for label in labels if label != -100) + + +# ============================================================================= +# Tests +# ============================================================================= +def test_normal_bridge_append(): + """Tool messages + generation prompt are appended as -100 bridge tokens. + + The pre-existing history is preserved verbatim (prefix of the new + ``input_ids``), the appended positions are all masked (-100), and + ``messages`` is updated to ``messages_before + tool_messages``. + """ + tokenizer = FakeTokenizer() + template = FakeTemplate(tokenizer) + + messages = [{'role': 'user', 'content': 'What is the weather?'}] + pif = _make_pif(template, messages) + before_ids = list(pif['input_ids']) + before_trainable = _count_trainable(pif['labels']) + + tool_messages = [{'role': 'tool', 'content': 'sunny'}] + new_pif = extend_with_bridge(pif, tool_messages, template) + + assert new_pif is not None + # History preserved verbatim as a prefix. + assert new_pif['input_ids'][:len(before_ids)] == before_ids + # Bridge actually added tokens. + assert len(new_pif['input_ids']) > len(before_ids) + # All newly appended positions are masked (-100 → not trainable). + assert _count_trainable(new_pif['labels']) == before_trainable + # input_ids / labels stay aligned. + assert len(new_pif['input_ids']) == len(new_pif['labels']) + # messages updated to before + tool. + assert new_pif['messages'] == messages + tool_messages + + # The bridge delta equals the template-space difference exactly. + s_before = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=False) + s_after = tokenizer.apply_chat_template( + messages + tool_messages, tokenize=False, add_generation_prompt=True) + expected_bridge_ids = tokenizer.encode(s_after[len(s_before):], add_special_tokens=False) + assert new_pif['input_ids'][len(before_ids):] == expected_bridge_ids + + +def test_non_prefix_extension_raises(): + """``s_after`` not a prefix-extension of ``s_before`` → RuntimeError.""" + tokenizer = NonMonotonicTokenizer() + template = FakeTemplate(tokenizer) + + messages = [{'role': 'user', 'content': 'hello'}] + pif = _make_pif(template, messages) + + with pytest.raises(RuntimeError, match='prefix-extension'): + extend_with_bridge(pif, [{'role': 'tool', 'content': 'x'}], template) + + +def test_empty_bridge_text_raises(): + """Bridge text computing to empty string → RuntimeError.""" + tokenizer = EmptyBridgeTokenizer() + template = FakeTemplate(tokenizer) + + messages = [{'role': 'user', 'content': 'hello'}] + pif = _make_pif(template, messages) + + with pytest.raises(RuntimeError, match='empty string'): + extend_with_bridge(pif, [{'role': 'tool', 'content': 'x'}], template) + + +def test_labels_length_mismatch_raises(): + """``labels`` length != ``input_ids`` length in the pif → RuntimeError.""" + tokenizer = FakeTokenizer() + template = FakeTemplate(tokenizer) + + messages = [{'role': 'user', 'content': 'hello'}] + pif = _make_pif(template, messages) + # Corrupt the pif: drop one label so lengths disagree. + pif['labels'] = list(pif['labels'])[:-1] + + with pytest.raises(RuntimeError, match='labels length'): + extend_with_bridge(pif, [{'role': 'tool', 'content': 'x'}], template) diff --git a/tests/twinkle_client/__init__.py b/tests/twinkle_client/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/twinkle_client/test_client_multi_turn_rollout.py b/tests/twinkle_client/test_client_multi_turn_rollout.py new file mode 100644 index 00000000..1ff69f5f --- /dev/null +++ b/tests/twinkle_client/test_client_multi_turn_rollout.py @@ -0,0 +1,563 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Property-based tests for +:class:`twinkle_client.rollout.multi_turn.ClientMultiTurnRollout`. + +These tests are 100% CPU-only and do NOT require a GPU or a running server. +They reuse the char-level Fake Tokenizer / Template infrastructure style from +``tests/twinkle_agentic/test_multi_turn_rollout.py`` but adapt the fake sampler +to the ``twinkle_client`` HTTP contract: ``FakeClientSampler.sample()`` mirrors +``vLLMSampler.sample()`` and returns ``List[SampleResponseModel]`` (pydantic, +from ``twinkle_client.types.sampler``) whose ``sequences[0]`` carries a populated +``new_input_feature`` so the multi-turn loop can proceed round after round. + +Properties covered: + * Output length & order preservation + * stop_reason value range + * logprobs / trainable-label alignment + * Actual turns never exceed max_turns + * Forced truncation at the max_turns edge +""" +from __future__ import annotations + +import copy +import json +import re +from collections import defaultdict +from typing import Any, Dict, List, Optional + +import pytest +from hypothesis import given, settings +from hypothesis import strategies as st + +from twinkle.data_format.sampling import SamplingParams +from twinkle_agentic.tools.base import Tool +from twinkle_agentic.tools.tool_manager import ToolManager +from twinkle_client.rollout.multi_turn import ClientMultiTurnRollout +from twinkle_client.types.sampler import SampledSequenceModel, SampleResponseModel + + +# ============================================================================= +# Fakes (tokenizer / template mirror the twinkle_agentic test infra) +# ============================================================================= +class FakeTokenizer: + """Char-level tokenizer with atomic special tokens. + + Guarantees ``decode(encode(s)) == s`` for any mix of raw chars and + registered specials, which is what makes ``extend_with_bridge``'s + template-space delta computation deterministic in the test. + """ + SPECIALS = ('<|im_start|>', '<|im_end|>') + + def __init__(self) -> None: + self._s2i: Dict[str, int] = {} + self._i2s: Dict[int, str] = {} + for s in self.SPECIALS: + self._add(s) + + def _add(self, tok: str) -> int: + if tok not in self._s2i: + i = len(self._s2i) + self._s2i[tok] = i + self._i2s[i] = tok + return self._s2i[tok] + + def encode(self, text: str, add_special_tokens: bool = False) -> List[int]: + ids: List[int] = [] + i = 0 + while i < len(text): + matched = False + for sp in self.SPECIALS: + if text.startswith(sp, i): + ids.append(self._add(sp)) + i += len(sp) + matched = True + break + if not matched: + ids.append(self._add(text[i])) + i += 1 + return ids + + def decode(self, ids: List[int], skip_special_tokens: bool = False) -> str: + specials = set(self.SPECIALS) + toks = [self._i2s[int(i)] for i in ids] + if skip_special_tokens: + toks = [t for t in toks if t not in specials] + return ''.join(toks) + + def apply_chat_template( + self, + messages: List[Dict[str, Any]], + tokenize: bool = False, + add_generation_prompt: bool = False, + **_, + ): + s = '' + for m in messages: + role = m['role'] + content = m['content'] + s += f'<|im_start|>{role}\n{content}<|im_end|>\n' + if add_generation_prompt: + s += '<|im_start|>assistant\n' + if tokenize: + return self.encode(s) + return s + + +class FakeTemplate: + """Minimal Template mirroring the parts ClientMultiTurnRollout touches.""" + model_id = 'qwen-fake' + truncation_strategy = 'right' + enable_thinking = False + + def __init__(self, tokenizer: FakeTokenizer) -> None: + self.tokenizer = tokenizer + + def encode(self, trajectory: Dict[str, Any], add_generation_prompt: bool = False) -> Dict[str, Any]: + messages = trajectory.get('messages', []) + s = self.tokenizer.apply_chat_template( + messages, tokenize=False, add_generation_prompt=add_generation_prompt) + input_ids = self.tokenizer.encode(s, add_special_tokens=False) + pif: Dict[str, Any] = dict(trajectory) # preserve top-level fields (incl. _tid) + pif['input_ids'] = input_ids + pif['labels'] = [-100] * len(input_ids) # inference mode + return self._invoke_post_pipeline([pif])[0] + + def _invoke_post_pipeline(self, inputs: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + out = [] + for pif in inputs: + pif = dict(pif) + input_ids = list(pif['input_ids']) + labels = list(pif.get('labels') or []) + if labels: + if len(labels) != len(input_ids): + raise RuntimeError(f'FakeTemplate post_pipeline: labels({len(labels)}) ' + f'!= input_ids({len(input_ids)})') + # np.roll(labels, -1): shift LEFT by 1 (output/shifted order) + labels = labels[1:] + labels[:1] + pif['input_ids'] = input_ids + pif['labels'] = labels + pif['attention_mask'] = [1] * len(input_ids) + pif['position_ids'] = list(range(len(input_ids))) + pif['length'] = len(input_ids) + out.append(pif) + return out + + def parse_tool_call(self, decoded: str) -> List[Dict[str, Any]]: + matches = re.findall(r'\s*([\s\S]*?)\s*', decoded or '') + results: List[Dict[str, Any]] = [] + for m in matches: + try: + d = json.loads(m) + except json.JSONDecodeError: + continue + name = d.get('name') or d.get('tool_name') + if not name: + continue + results.append({ + 'type': 'function', + 'function': { + 'name': name, + 'arguments': d.get('arguments', {}), + }, + }) + return results + + def concat_input_feature(self, pif: Dict[str, Any], new_tokens: List[int]) -> Dict[str, Any]: + result = copy.deepcopy(pif) + prompt_ids = list(result['input_ids']) + labels = list(result.get('labels') or []) + if labels: + # Unroll (shift RIGHT by 1): reverse the post_pipeline roll + labels = labels[-1:] + labels[:-1] + else: + labels = [-100] * len(prompt_ids) + input_ids = prompt_ids + list(new_tokens) + labels = labels + list(new_tokens) # assistant tokens trainable + result['input_ids'] = input_ids + result['labels'] = labels + result = self._invoke_post_pipeline([result])[0] + response_text = self.tokenizer.decode(new_tokens, skip_special_tokens=True) + messages = list(result.get('messages') or []) + messages.append({'role': 'assistant', 'content': response_text}) + result['messages'] = messages + return result + + +class FakeClientSampler: + """Script-driven sampler mirroring ``vLLMSampler.sample()``. + + Adapts the ``twinkle_agentic`` fake sampler to the client HTTP contract: + ``sample()`` accepts ``(inputs, sampling_params=, ...)`` and returns + ``List[SampleResponseModel]`` (pydantic), each sequence carrying a populated + ``new_input_feature`` so the multi-turn loop can proceed. + + Each trajectory is identified by a hidden ``_tid`` field that survives + ``encode`` / ``concat_input_feature`` / ``extend_with_bridge`` (all preserve + top-level keys), so we can look up its scripted turns regardless of how the + active set shrinks across rounds. + """ + + def __init__(self, template: FakeTemplate, scripts: Dict[int, List[Dict[str, Any]]]) -> None: + self.template = template + self.scripts = scripts + self.turn_counters: Dict[int, int] = defaultdict(int) + self.sample_calls = 0 + + def sample(self, inputs, sampling_params: Optional[Dict[str, Any]] = None, **kwargs): + if isinstance(inputs, dict): + inputs = [inputs] + assert isinstance(inputs, list), f'expects a list, got {type(inputs).__name__}' + # Contract check: the rollout coerces core-lib SamplingParams into a + # plain dict before calling the HTTP sampler. + assert sampling_params is None or isinstance(sampling_params, dict) + + responses: List[SampleResponseModel] = [] + for pif in inputs: + tid = pif['_tid'] + script = self.scripts[tid] + idx = self.turn_counters[tid] + self.turn_counters[tid] += 1 + self.sample_calls += 1 + + if idx < len(script): + turn = script[idx] + else: + # Defensive fallback: terminate cleanly if over-sampled. + turn = {'kind': 'terminal', 'stop_reason': 'stop', 'logprobs': False} + + if turn['kind'] == 'tool': + decoded = _tool_call_text('search', {'q': f't{idx}'}) + stop_reason = 'stop' + else: + decoded = f'final-{idx}' + stop_reason = turn['stop_reason'] + + raw = decoded + '<|im_end|>' + tokens = self.template.tokenizer.encode(raw, add_special_tokens=False) + logprobs = ([[(int(t), -0.1)] for t in tokens] if turn['logprobs'] else None) + new_pif = self.template.concat_input_feature(pif, tokens) + + seq = SampledSequenceModel( + stop_reason=stop_reason, + tokens=tokens, + logprobs=logprobs, + decoded=decoded, + new_input_feature=new_pif, + ) + responses.append(SampleResponseModel(sequences=[seq])) + return responses + + +class EchoTool(Tool): + """Echoes its arguments as a JSON string.""" + + def __init__(self, name: str = 'search'): + self._name = name + + def __call__(self, tool_name: str, arguments: Dict[str, Any]) -> str: + return f'echo[{tool_name}]:{json.dumps(arguments, sort_keys=True)}' + + def tool_info(self): + return { + 'type': 'function', + 'function': { + 'name': self._name, + 'description': 'echo test tool', + 'parameters': {}, + }, + } + + +# ============================================================================= +# Helpers +# ============================================================================= +def _tool_call_text(name: str, arguments: Dict[str, Any]) -> str: + return '' + json.dumps({'name': name, 'arguments': arguments}) + '' + + +def _count_trainable(labels: List[int]) -> int: + return sum(1 for label in labels if label != -100) + + +def _make_tool_manager() -> ToolManager: + mgr = ToolManager({}) + mgr.register(EchoTool('search')) + return mgr + + +def _build_from_scripts(scripts_spec: List[Dict[str, Any]]): + """Turn a list of per-trajectory specs into (rollout inputs, sampler). + + ``scripts_spec[k]`` = {'num_tools': int, 'terminal': 'stop'|'length', + 'logprobs': bool}. Each trajectory's script is + ``num_tools`` tool-call turns followed by one terminal turn. + """ + tokenizer = FakeTokenizer() + template = FakeTemplate(tokenizer) + + scripts: Dict[int, List[Dict[str, Any]]] = {} + trajectories: List[Dict[str, Any]] = [] + for tid, spec in enumerate(scripts_spec): + turns: List[Dict[str, Any]] = [] + for _ in range(spec['num_tools']): + turns.append({'kind': 'tool', 'stop_reason': 'stop', 'logprobs': spec['logprobs']}) + turns.append({'kind': 'terminal', 'stop_reason': spec['terminal'], 'logprobs': spec['logprobs']}) + scripts[tid] = turns + trajectories.append({'messages': [{'role': 'user', 'content': f'q{tid}'}], '_tid': tid}) + + sampler = FakeClientSampler(template, scripts) + return trajectories, sampler, template + + +# ============================================================================= +# Hypothesis strategies +# ============================================================================= +_MAX_TOOLS = 5 + + +def _traj_spec(): + return st.fixed_dictionaries({ + 'num_tools': st.integers(min_value=0, max_value=_MAX_TOOLS), + 'terminal': st.sampled_from(['stop', 'length']), + 'logprobs': st.booleans(), + }) + + +def _batch_specs(min_size: int = 1, max_size: int = 4): + return st.lists(_traj_spec(), min_size=min_size, max_size=max_size) + + +# ============================================================================= +# Multi-turn output length & order preservation +# ============================================================================= +@settings(deadline=None, max_examples=60) +@given(scripts_spec=_batch_specs(min_size=0, max_size=5), max_turns=st.integers(min_value=1, max_value=6)) +def test_output_length_and_order_preserved(scripts_spec, max_turns): + """``__call__`` returns a list of the same length as the input, in the exact + same order (verified via the hidden per-trajectory ``_tid`` tag).""" + trajectories, sampler, template = _build_from_scripts(scripts_spec) + rollout = ClientMultiTurnRollout( + sampler=sampler, template=template, tool_manager=_make_tool_manager(), max_turns=max_turns) + + outs = rollout(copy.deepcopy(trajectories)) + + assert len(outs) == len(trajectories) + for i, out in enumerate(outs): + assert out['_tid'] == i, 'output order must match input order' + + +# ============================================================================= +# stop_reason value range +# ============================================================================= +@settings(deadline=None, max_examples=60) +@given(scripts_spec=_batch_specs(), max_turns=st.integers(min_value=1, max_value=6)) +def test_stop_reason_value_range(scripts_spec, max_turns): + """Every returned trajectory's ``stop_reason`` is one of + ``{'length', 'stop', 'max_turns'}``.""" + trajectories, sampler, template = _build_from_scripts(scripts_spec) + rollout = ClientMultiTurnRollout( + sampler=sampler, template=template, tool_manager=_make_tool_manager(), max_turns=max_turns) + + outs = rollout(copy.deepcopy(trajectories)) + + for out in outs: + assert out['stop_reason'] in {'length', 'stop', 'max_turns'}, out['stop_reason'] + + +# ============================================================================= +# logprobs / trainable-label alignment +# ============================================================================= +@settings(deadline=None, max_examples=60) +@given(scripts_spec=_batch_specs(), max_turns=st.integers(min_value=1, max_value=6)) +def test_logprobs_align_with_trainable_labels(scripts_spec, max_turns): + """For every returned trajectory with a non-empty ``logprobs`` list, its length + equals the number of trainable labels (``label != -100``).""" + trajectories, sampler, template = _build_from_scripts(scripts_spec) + rollout = ClientMultiTurnRollout( + sampler=sampler, template=template, tool_manager=_make_tool_manager(), max_turns=max_turns) + + outs = rollout(copy.deepcopy(trajectories)) + + for out in outs: + logprobs = out.get('logprobs') + if logprobs: + trainable = _count_trainable(out.get('labels') or []) + assert len(logprobs) == trainable, ( + f'logprobs({len(logprobs)}) != trainable labels({trainable})') + + +# ============================================================================= +# Actual turns never exceed max_turns +# ============================================================================= +@settings(deadline=None, max_examples=60) +@given(scripts_spec=_batch_specs(), max_turns=st.integers(min_value=1, max_value=6)) +def test_turns_do_not_exceed_max_turns(scripts_spec, max_turns): + """Every returned trajectory's ``turns`` count is <= the configured + ``max_turns``.""" + trajectories, sampler, template = _build_from_scripts(scripts_spec) + rollout = ClientMultiTurnRollout( + sampler=sampler, template=template, tool_manager=_make_tool_manager(), max_turns=max_turns) + + outs = rollout(copy.deepcopy(trajectories)) + + for out in outs: + assert out['turns'] <= max_turns, f"turns({out['turns']}) > max_turns({max_turns})" + + +# ============================================================================= +# Forced truncation at the max_turns edge +# ============================================================================= +@settings(deadline=None, max_examples=60) +@given(logprobs_flags=st.lists(st.booleans(), min_size=1, max_size=5)) +def test_max_turns_one_forces_truncation(logprobs_flags): + """With ``max_turns == 1`` and a first-round tool_call, every trajectory is + marked ``truncated=True`` and ``stop_reason='max_turns'`` and stops after + exactly one turn.""" + # Every trajectory emits a tool_call on its first (and only allowed) turn. + scripts_spec = [{'num_tools': 3, 'terminal': 'stop', 'logprobs': lp} for lp in logprobs_flags] + trajectories, sampler, template = _build_from_scripts(scripts_spec) + rollout = ClientMultiTurnRollout( + sampler=sampler, template=template, tool_manager=_make_tool_manager(), max_turns=1) + + outs = rollout(copy.deepcopy(trajectories)) + + assert len(outs) == len(trajectories) + for out in outs: + assert out['truncated'] is True + assert out['stop_reason'] == 'max_turns' + assert out['turns'] == 1 + + +# ============================================================================= +# Deterministic unit tests: exception paths & dependency reuse (non-hypothesis) +# +# These cover the failure/edge contract described in the module docstring of +# ``twinkle_client.rollout.multi_turn``. +# All are CPU-only and reuse the Fake infra above. +# ============================================================================= +class NetworkError(Exception): + """Stand-in for a requests-style transport error (connection reset/timeout).""" + + +class _NullFeatureSampler: + """Sampler that violates the contract by returning ``new_input_feature=None``. + + Mirrors ``vLLMSampler.sample()`` shape (returns ``List[SampleResponseModel]``) + but every sequence lacks ``new_input_feature``, which must make the multi-turn + loop raise a batch/trajectory-indexed ``RuntimeError``. + """ + + def __init__(self, template: FakeTemplate) -> None: + self.template = template + self.sample_calls = 0 + + def sample(self, inputs, sampling_params=None, **kwargs): + if isinstance(inputs, dict): + inputs = [inputs] + self.sample_calls += 1 + responses: List[SampleResponseModel] = [] + for _ in inputs: + seq = SampledSequenceModel( + stop_reason='stop', + tokens=[0, 1], + logprobs=None, + decoded='final', + new_input_feature=None, # contract violation under test + ) + responses.append(SampleResponseModel(sequences=[seq])) + return responses + + +class _NetworkFailingSampler: + """Sampler whose ``sample()`` always raises a network-like exception. + + Used to assert the rollout NEVER wraps or swallows transport errors coming + from ``vLLMSampler.sample()``; they must propagate unchanged. + """ + + def __init__(self, template: FakeTemplate, exc: Exception) -> None: + self.template = template + self._exc = exc + self.sample_calls = 0 + + def sample(self, inputs, sampling_params=None, **kwargs): + self.sample_calls += 1 + raise self._exc + + +def test_missing_new_input_feature_raises_indexed_runtime_error(): + """new_input_feature=None -> RuntimeError naming batch AND trajectory index.""" + trajectories, _script_sampler, template = _build_from_scripts( + [{'num_tools': 0, 'terminal': 'stop', 'logprobs': False}]) + sampler = _NullFeatureSampler(template) + rollout = ClientMultiTurnRollout( + sampler=sampler, template=template, tool_manager=_make_tool_manager(), max_turns=3) + + with pytest.raises(RuntimeError) as excinfo: + rollout(copy.deepcopy(trajectories)) + + msg = str(excinfo.value) + # Message must carry both the batch index and the trajectory index so the + # failure is localizable in a batched HTTP round. + assert 'batch index 0' in msg, msg + assert 'trajectory 0' in msg, msg + assert 'new_input_feature' in msg, msg + + +def test_tool_calls_without_tool_manager_raises_value_error(): + """tool_calls produced but tool_manager missing -> ValueError.""" + # One tool-call turn then a terminal turn; max_turns=2 so the tool-dispatch + # site (not the max_turns truncation edge) is what fails. + trajectories, sampler, template = _build_from_scripts( + [{'num_tools': 1, 'terminal': 'stop', 'logprobs': False}]) + rollout = ClientMultiTurnRollout( + sampler=sampler, template=template, tool_manager=None, max_turns=2) + + with pytest.raises(ValueError) as excinfo: + rollout(copy.deepcopy(trajectories)) + + msg = str(excinfo.value) + assert 'tool_manager' in msg, msg + assert 'trajectory 0' in msg, msg + + +def test_tool_calls_without_tool_manager_via_per_call_kwarg_raises_value_error(): + """Passing tool_manager=None as a per-call kwarg also raises at dispatch.""" + trajectories, sampler, template = _build_from_scripts( + [{'num_tools': 1, 'terminal': 'stop', 'logprobs': False}]) + # Constructed WITH a manager, but the per-call override nulls it out. + rollout = ClientMultiTurnRollout( + sampler=sampler, template=template, tool_manager=_make_tool_manager(), max_turns=2) + + with pytest.raises(ValueError): + rollout(copy.deepcopy(trajectories), tool_manager=None) + + +def test_sampler_network_error_propagates_unchanged(): + """vLLMSampler.sample() network error propagates unchanged (not swallowed/wrapped).""" + trajectories, _script_sampler, template = _build_from_scripts( + [{'num_tools': 0, 'terminal': 'stop', 'logprobs': False}]) + sentinel = NetworkError('simulated connection reset by peer') + sampler = _NetworkFailingSampler(template, sentinel) + rollout = ClientMultiTurnRollout( + sampler=sampler, template=template, tool_manager=_make_tool_manager(), max_turns=3) + + with pytest.raises(NetworkError) as excinfo: + rollout(copy.deepcopy(trajectories)) + + # Same exact exception object, neither re-wrapped nor replaced. + assert excinfo.value is sentinel + assert str(excinfo.value) == 'simulated connection reset by peer' + assert sampler.sample_calls == 1 + + +def test_dependencies_are_reused_not_reimplemented(): + """ClientMultiTurnRollout imports (does not copy) ToolManager & extend_with_bridge.""" + import twinkle_agentic.rollout.bridge as bridge_mod + import twinkle_agentic.tools.tool_manager as tool_manager_mod + import twinkle_client.rollout.multi_turn as m + + # Same object identity => the symbols are imported from the shared core-lib + # modules rather than re-defined locally. + assert m.ToolManager is tool_manager_mod.ToolManager + assert m.extend_with_bridge is bridge_mod.extend_with_bridge