From d100f5d3e74d0e9bff3f9db7f0bc515b4996f36a Mon Sep 17 00:00:00 2001 From: Yunnglin Date: Tue, 14 Jul 2026 10:31:58 +0800 Subject: [PATCH 01/10] feat: client-side embedding training & multi-turn rollout support Expose the core-library embedding training (InfoNCE) and trainable multi-turn rollout capabilities through twinkle_client, and clean up the client_generator technical debt. - Remove client_tools/client_generator.py; convert affected client modules (dataloader/dataset/processor/model/sampler) to hand-maintained sources by stripping their AUTO-GENERATED headers. - Extract bridge-token stitching into the shared pure function twinkle_agentic/rollout/bridge.py::extend_with_bridge; MultiTurnRollout now calls it instead of duplicating the logic. - Add ClientMultiTurnRollout (src/twinkle_client/rollout/) driving HTTP sampling via vLLMSampler, reusing ToolManager and extend_with_bridge. - Enhance MockSampler with opt-in multi-turn knobs (new_input_feature, configurable stop_reason / tool_call_text / tool_call_turns) for local CPU-only multi-turn e2e; backward compatible by default. - Add embedding e2e scaffolding + local mock protocol-link validation and GPU-gated numeric/SP-CP tests; add multi-turn property/unit tests, local mock e2e, and GPU-gated alignment / Gateway-logprobs tests. - Document client embedding call sequence and the Tinker Gateway indirect multi-turn approach with its capability boundaries. Local CPU suite: 98 passed, GPU-gated cases skip cleanly. --- client_tools/client_generator.py | 918 ------------- .../Embedding\350\256\255\347\273\203.md" | 66 + ...71\345\256\242\346\210\267\347\253\257.md" | 65 + .../server/sampler/backends/mock_sampler.py | 150 ++- src/twinkle_agentic/rollout/__init__.py | 2 + src/twinkle_agentic/rollout/bridge.py | 156 +++ src/twinkle_agentic/rollout/multi_turn.py | 114 +- src/twinkle_client/dataloader/__init__.py | 10 - src/twinkle_client/dataloader/dataloader.py | 10 - src/twinkle_client/dataset/__init__.py | 10 - src/twinkle_client/dataset/base.py | 10 - .../dataset/iterable_dataset.py | 10 - .../dataset/iterable_packing_dataset.py | 10 - src/twinkle_client/dataset/lazy_dataset.py | 10 - src/twinkle_client/dataset/packing_dataset.py | 10 - src/twinkle_client/model/__init__.py | 10 - .../model/multi_lora_transformers.py | 10 - src/twinkle_client/processor/__init__.py | 10 - src/twinkle_client/processor/base.py | 10 - src/twinkle_client/rollout/__init__.py | 4 + src/twinkle_client/rollout/multi_turn.py | 289 +++++ src/twinkle_client/sampler/__init__.py | 10 - src/twinkle_client/sampler/vllm_sampler.py | 10 - ...test_api_multi_turn_rollout_no_logprobs.py | 261 ++++ .../test_client_multi_turn_rollout_e2e.py | 597 +++++++++ ...test_client_multi_turn_rollout_mock_e2e.py | 577 +++++++++ tests/server/sampler/test_mock_sampler.py | 60 + tests/server/test_embedding_e2e.py | 1139 +++++++++++++++++ tests/server/test_embedding_e2e_sp_cp.py | 446 +++++++ tests/twinkle_agentic/test_bridge.py | 253 ++++ tests/twinkle_client/__init__.py | 0 .../test_client_multi_turn_rollout.py | 588 +++++++++ 32 files changed, 4646 insertions(+), 1179 deletions(-) delete mode 100644 client_tools/client_generator.py create mode 100644 src/twinkle_agentic/rollout/bridge.py create mode 100644 src/twinkle_client/rollout/__init__.py create mode 100644 src/twinkle_client/rollout/multi_turn.py create mode 100644 tests/server/gateway/test_api_multi_turn_rollout_no_logprobs.py create mode 100644 tests/server/integration/test_client_multi_turn_rollout_e2e.py create mode 100644 tests/server/integration/test_client_multi_turn_rollout_mock_e2e.py create mode 100644 tests/server/test_embedding_e2e.py create mode 100644 tests/server/test_embedding_e2e_sp_cp.py create mode 100644 tests/twinkle_agentic/test_bridge.py create mode 100644 tests/twinkle_client/__init__.py create mode 100644 tests/twinkle_client/test_client_multi_turn_rollout.py 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/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..f5df4e9f 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,72 @@ for step, batch in enumerate(dataloader): --- +## 通过 Twinkle Client 训练 + +除了直接使用裸库 `TransformersModel`，你也可以通过 `twinkle_client` 的 `MultiLoraTransformersModel` 以 HTTP 方式驱动服务端完成 Embedding 训练。协议层（`/twinkle/*`）已经具备完成 Embedding 训练所需的全部能力，**无需引入任何新的高层封装函数**（例如 `setup_embedding_training`），只需按正确顺序调用现有方法即可。 + +### 完整调用顺序 + +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) +``` + +### 关于 `TransformersEmbeddingPatch` 的自动应用 + +Embedding 训练依赖 `TransformersEmbeddingPatch` 把 causal LM 的 `lm_head` 替换为 identity 输出，从而得到 per-token hidden states 用于池化。**你不需要显式调用 `apply_patch(TransformersEmbeddingPatch())`**：当你在 `forward_backward`（或 `forward_only`）中传入 `task='embedding'` 时，服务端的 `_resolve_task_context` 会在该次 `forward` 调用期间自动应用该 patch，并在调用结束后自动回滚。 + +这意味着： + +- 同一个 adapter 在一次 `forward_backward(task='embedding')` 之后，紧接着调用不带 `task` 参数的 `forward_only` 仍会返回正常的词表维度 `logits`，不会残留上一次 embedding 任务的 identity hidden states。 +- 你在 client 侧真正需要显式调用的前置步骤只有三个：`set_processor('InputProcessor')`、`set_loss('InfonceLoss', ...)` 与 `add_metric('EmbeddingMetric', is_training=True)`。 + +### 说明 + +- 本节不引入任何新的高层封装函数，仅说明现有 `MultiLoraTransformersModel` 方法的正确调用顺序。 +- 各参数（`temperature`、`use_batch`、`hard_negatives` 等）的含义与取值建议参见上文「关键参数」；`calculate_metric` 返回的指标含义参见下文「监控指标」。 + +--- + ## 监控指标 `EmbeddingMetric` 报告关键训练信号： 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/Tinker\345\205\274\345\256\271\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/Tinker\345\205\274\345\256\271\345\256\242\346\210\267\347\253\257.md" index a1f7e064..2bc16c17 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/Tinker\345\205\274\345\256\271\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/Tinker\345\205\274\345\256\271\345\256\242\346\210\267\347\253\257.md" @@ -195,3 +195,68 @@ rest_client = service_client.create_rest_client() rest_client.publish_checkpoint_from_tinker_path(save_result.path).result() print("Published checkpoint to ModelScope Hub") ``` + +## 多轮工具调用（间接方案） + +Tinker 兼容客户端本身不提供多轮 agentic rollout 编排能力。如果你在用 Tinker Client 训练/管理 checkpoint，同时又需要多轮工具调用能力（**仅用于生成或评测**），可以在完全不改动 Tinker SDK 代码的前提下，额外使用一个标准的 OpenAI 兼容客户端指向 Gateway 的 `/chat/completions` 端点，并复用 `twinkle_agentic.rollout.APIMultiTurnRollout` 与 `ToolManager`。 + +这是一条纯增量的使用路径：它复用的是 Twinkle Server 上一个独立于 Tinker 协议命名空间的、已经存在并跑通的 OpenAI 兼容端点，因此不涉及对 Tinker SDK 源码或 `patch_tinker.py` 中现有 monkeypatch 范围的任何修改。 + +```python +import os +from twinkle.data_format.sampling import SamplingParams +from twinkle_agentic.protocol.openai import OpenAI +from twinkle_agentic.rollout import APIMultiTurnRollout +from twinkle_agentic.tools.tool_manager import ToolManager +# from your_project.tools import CalculatorTool, SearchTool + +# An OpenAI-compatible client pointing at the Gateway route, +# NOT at tinker's base_url. This is a plain additive usage of an +# already-existing endpoint; the tinker SDK is untouched. +api = OpenAI( + model='Qwen/Qwen3.5-4B', + api_key=os.environ.get('MODELSCOPE_TOKEN', 'EMPTY_API_KEY'), + base_url='http://localhost:8000/api/v1', # Gateway /chat/completions route +) + +# Register your tools; ToolManager is reused directly from twinkle_agentic. +tool_manager = ToolManager([ + # CalculatorTool(), + # SearchTool(), +]) + +rollout = APIMultiTurnRollout( + api=api, + tool_manager=tool_manager, + sampling_params=SamplingParams(temperature=0.7, max_tokens=512), + max_turns=6, + concurrency=8, +) + +trajectories_in = [ + {'messages': [{'role': 'user', 'content': 'What is 12 * 34?'}]}, +] +trajectories_out = rollout(trajectories_in) + +# trajectories_out[i]['messages'] contains the full conversation, including +# assistant tool_calls and tool-response turns. +# trajectories_out[i]['stop_reason'] is one of 'stop' | 'length' | 'max_turns' | 'api_error'. +for traj in trajectories_out: + for message in traj['messages']: + print(message['role'], ':', message.get('content')) +``` + +### 能力边界：不可用于 RL 训练 + +该间接方案返回的 Trajectory **不包含** `logprobs` 字段（端到端验证已确认走 Gateway `/chat/completions` 的 `APIMultiTurnRollout` 产出的 trajectory 中不存在 `logprobs`，也没有 `input_ids` 等训练所需字段）。因此该路径**仅适合生成数据、离线评测等场景，不能直接喂给依赖 token 级对齐信息的 RL 训练（如 GRPO）**。 + +如果多轮 rollout 的产出需要直接用于 RL 训练（需要 token 级 `logprobs` 与 `labels` 对齐），请使用 Twinkle 客户端侧的 `ClientMultiTurnRollout`（通过 `/twinkle/sample` 返回 `new_input_feature`/`logprobs`），详见《Twinkle 客户端》文档。 + +### 为什么 Tinker 客户端不支持 embedding 训练与可训练多轮 rollout + +这两项能力无法通过 Tinker 兼容客户端提供，根本原因在于 Tinker 协议的数据类型层面缺少所需结构： + +- **Embedding 训练**：`tinker.types.Datum` 不具备对比学习所需的样本分组结构（anchor/positive 成对分组），无法承载 embedding 训练的输入语义。 +- **可训练多轮 rollout**：`tinker.types.SampledSequence` 不具备 `new_input_feature` 字段，无法在多轮之间传递并累积 token 级对齐信息，因此产出无法用于训练。 + +由于 Tinker 是第三方 SDK、不可修改其源码，即使想补齐这些字段也需要改动 SDK 本身（不允许）。因此上述能力仅在 Twinkle 客户端侧支持，Tinker 兼容客户端只提供上文所述的"仅生成/评测"间接多轮方案。 diff --git a/src/twinkle/server/sampler/backends/mock_sampler.py b/src/twinkle/server/sampler/backends/mock_sampler.py index ec5c86f9..f04accb4 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 @@ -23,16 +23,61 @@ logger = get_logger() +# Recognised ctor / per-call knobs that tune multi-turn control flow. Kept in a +# set so the ctor can log only *genuinely* unknown kwargs (a real backend +# signature drift) while still accepting these opt-in multi-turn knobs. +_MULTI_TURN_KNOBS = ('stop_reason', 'tool_call_text', 'tool_call_turns') + @remote_class() class MockSampler: - """Deterministic mock sampler for CPU-only testing.""" + """Deterministic mock sampler for CPU-only testing. + + 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). - def __init__(self, model_id: str, *, seed: int = 0, vocab_size: int = 32, **kwargs: Any) -> None: + All three knobs may be supplied at construction time or overridden per call + via ``sample()`` keyword arguments or matching attributes on + ``sampling_params``. 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 +109,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 (precedence: explicit sample() + # kwargs > sampling_params attributes > ctor defaults) and advance the + # round counter that ``tool_call_turns`` addresses. + stop_reason = self._resolve_stop_reason(sampling_params, kwargs) + tool_call_text = self._resolve_tool_call_text(sampling_params, kwargs) + tool_call_turns = self._resolve_tool_call_turns(sampling_params, 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 +133,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 +208,77 @@ 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, params: SamplingParams | None, kwargs: dict[str, Any]) -> str: + if 'stop_reason' in kwargs and kwargs['stop_reason'] is not None: + return str(kwargs['stop_reason']) + override = getattr(params, 'stop_reason', None) + if override is not None: + return str(override) + return self._stop_reason + + def _resolve_tool_call_text(self, params: SamplingParams | None, kwargs: dict[str, Any]) -> str | None: + if 'tool_call_text' in kwargs: + return kwargs['tool_call_text'] + override = getattr(params, 'tool_call_text', None) + if override is not None: + return override + return self._tool_call_text + + def _resolve_tool_call_turns( + self, + params: SamplingParams | None, + 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) + override = getattr(params, 'tool_call_turns', None) + if override is not None: + return self._normalize_turns(override, tool_call_text) + # Fall back to the ctor-configured turns, but keep the "default to {1} + # when text is injected via a per-call override" behaviour consistent. + 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_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..c3c82dd5 100644 --- a/src/twinkle_agentic/rollout/multi_turn.py +++ b/src/twinkle_agentic/rollout/multi_turn.py @@ -12,6 +12,7 @@ from twinkle.template.base import Template from twinkle_agentic.tools.tool_manager import ToolManager from .base import Rollout +from .bridge import extend_with_bridge def _to_plain(obj: Any) -> Any: @@ -214,7 +215,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 +404,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..cd42bafd --- /dev/null +++ b/src/twinkle_client/rollout/multi_turn.py @@ -0,0 +1,289 @@ +# 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`` (see task 8.2). +""" +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``. + 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 (task + # 8.3), 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 (requirement 3.10). + """ + 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..270da067 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 diff --git a/tests/server/gateway/test_api_multi_turn_rollout_no_logprobs.py b/tests/server/gateway/test_api_multi_turn_rollout_no_logprobs.py new file mode 100644 index 00000000..76b68cdf --- /dev/null +++ b/tests/server/gateway/test_api_multi_turn_rollout_no_logprobs.py @@ -0,0 +1,261 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""GPU-gated E2E: ``APIMultiTurnRollout`` over the Gateway ``/chat/completions`` +endpoint proves the returned Trajectory carries NO ``logprobs`` field. + +Purpose +------- +This is the "generation-only, not trainable" counterpart to the trainable +``ClientMultiTurnRollout`` path validated in task 9.2. It drives the SAME +tool-calling scenario, but through the OpenAI-compatible route: + + OpenAI client -> Gateway POST /chat/completions -> /twinkle/sample + +and asserts that the resulting Trajectory does **not** contain a ``logprobs`` +field (or that it is ``None``/absent). The OpenAI chat-completions protocol only +surfaces assistant ``content`` / ``tool_calls`` / ``finish_reason`` — it never +carries the token-level ``logprobs`` + ``new_input_feature`` alignment info that +GRPO training requires. This test locks in the accuracy of that limitation +statement (Requirement 5.2 / 5.5): the Gateway indirection is fine for +generation/evaluation but cannot feed token-aligned RL training. + +What is "real" here +------------------- + * REAL (GPU): the running Twinkle server (gateway + real vLLM sampler for + ``Qwen/Qwen3.5-4B``), the ``/chat/completions`` -> ``/twinkle/sample`` hop, + and the ``openai`` pip client issuing actual HTTP requests. + * ``APIMultiTurnRollout`` + a small ``ToolManager`` (calculator tool) drive + the multi-turn control flow client-side. + +============================================================================== +GPU + CONDA ENVIRONMENT REQUIREMENT (MANDATORY) +============================================================================== +This case requires a GPU and an already-running GPU server (see +tests/server/start_e2e_server.py). It is gated behind ``TWINKLE_TEST_GPU_E2E=1`` +and SKIPS cleanly on machines without a GPU (e.g. local dev). Run on GPU CI in +the ``twinkle`` conda env: + + TWINKLE_TEST_GPU_E2E=1 conda run -n twinkle pytest \ + tests/server/gateway/test_api_multi_turn_rollout_no_logprobs.py -v + +Locally (no GPU / variable unset) it collects and skips without error: + + conda run -n twinkle pytest \ + tests/server/gateway/test_api_multi_turn_rollout_no_logprobs.py -v +============================================================================== +""" +from __future__ import annotations + +import json +import os +import sys +from typing import Any, Dict + +# Ensure project root is importable for both pytest and direct execution. +_PROJECT_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', '..')) +if _PROJECT_ROOT not in sys.path: + sys.path.insert(0, _PROJECT_ROOT) + +import pytest + +# Reuse the shared GPU-e2e gate and server coordinates. +from tests.server.integration.e2e_helpers import ( + API_KEY, + BASE_MODEL, + BASE_URL, + log, + wait_for_server, +) +from tests.server.test_embedding_e2e import gpu_e2e_enabled + +# The gateway (OpenAI-compatible) routes are mounted under the ``server`` app's +# route prefix ``/api/v1`` (see tests/server/config/server_config_4b_e2e.yaml). +# The ``openai`` client appends ``/chat/completions`` to this base URL. +GATEWAY_BASE_URL = f'{BASE_URL}/api/v1' + +# Model/adapter name routed by the gateway. The e2e config serves the base model +# ``Qwen/Qwen3.5-4B`` directly. +GATEWAY_MODEL = BASE_MODEL + + +# ═══════════════════════════════════════════════════════════════════════════ +# Tool scenario (mirrors the 9.2 tool-calling scenario): a small calculator. +# ═══════════════════════════════════════════════════════════════════════════ +TOOL_NAME = 'calculator' + + +def _make_tool_manager(): + """Build a ToolManager with a single deterministic calculator tool.""" + from twinkle_agentic.tools.base import Tool + from twinkle_agentic.tools.tool_manager import ToolManager + + class CalculatorTool(Tool): + """Adds/subtracts/multiplies two integers; returns the result as text.""" + + def __call__(self, tool_name: str, arguments: Dict[str, Any]) -> str: + a = int(arguments.get('a', 0)) + b = int(arguments.get('b', 0)) + op = arguments.get('operation', 'add') + if op == 'add': + return str(a + b) + if op == 'subtract': + return str(a - b) + if op == 'multiply': + return str(a * b) + return f'Error: unknown operation {op!r}' + + def tool_info(self): + return { + 'type': 'function', + 'function': { + 'name': TOOL_NAME, + 'description': 'Compute a + b, a - b, or a * b for two integers.', + 'parameters': { + 'type': 'object', + 'properties': { + 'a': {'type': 'integer', 'description': 'first operand'}, + 'b': {'type': 'integer', 'description': 'second operand'}, + 'operation': { + 'type': 'string', + 'enum': ['add', 'subtract', 'multiply'], + 'description': 'the arithmetic operation', + }, + }, + 'required': ['a', 'b', 'operation'], + }, + }, + } + + mgr = ToolManager({}) + mgr.register(CalculatorTool()) + return mgr + + +def _make_trajectory() -> Dict[str, Any]: + """A single math trajectory that nudges the model toward the calculator tool.""" + return { + 'messages': [ + { + 'role': 'system', + 'content': ('You are a precise assistant. When a calculation is needed, ' + 'call the calculator tool instead of computing yourself.'), + }, + {'role': 'user', 'content': 'What is 123 multiplied by 7? Use the calculator tool.'}, + ], + } + + +def _build_rollout(max_turns: int = 3): + """Construct an APIMultiTurnRollout wired to the Gateway /chat/completions.""" + from twinkle.data_format.sampling import SamplingParams + from twinkle_agentic.protocol.openai import OpenAI + from twinkle_agentic.rollout import APIMultiTurnRollout + + api = OpenAI(model=GATEWAY_MODEL, api_key=API_KEY, base_url=GATEWAY_BASE_URL) + # temperature=0 keeps the control flow as reproducible as the server allows. + sampling_params = SamplingParams(num_samples=1, temperature=0.0, max_tokens=256) + return APIMultiTurnRollout( + api=api, + tool_manager=_make_tool_manager(), + sampling_params=sampling_params, + max_turns=max_turns, + ) + + +# ═══════════════════════════════════════════════════════════════════════════ +# Fixtures +# ═══════════════════════════════════════════════════════════════════════════ +@pytest.fixture(scope='module') +def gpu_gateway_ready(): + """Ensure a GPU server is up before running; skip cleanly when GPU e2e is off.""" + if not gpu_e2e_enabled(): + pytest.skip('Set TWINKLE_TEST_GPU_E2E=1 to run the Gateway APIMultiTurnRollout ' + 'no-logprobs E2E (requires a running GPU server)') + wait_for_server() + yield + + +# ═══════════════════════════════════════════════════════════════════════════ +# Task 9.3 — APIMultiTurnRollout via Gateway: Trajectory carries NO logprobs. +# ═══════════════════════════════════════════════════════════════════════════ +def test_api_multi_turn_rollout_trajectory_has_no_logprobs(gpu_gateway_ready): + """The Gateway /chat/completions path yields Trajectories without logprobs. + + Runs the same tool-calling scenario as task 9.2 but through + ``APIMultiTurnRollout`` + an OpenAI-compatible client pointed at the Gateway. + Asserts the returned Trajectory does NOT expose a ``logprobs`` field (absent + or ``None``), confirming the "generation-only, not trainable" limitation is + accurate: token-level alignment info never crosses the OpenAI protocol. + + GPU-gated (TWINKLE_TEST_GPU_E2E=1 + running GPU server); skipped locally. + + Validates: Requirements 5.2, 5.5, 7.2 + """ + rollout = _build_rollout(max_turns=3) + trajectories = [_make_trajectory()] + + outs = rollout(trajectories) + + # Same length/order contract the multi-turn rollouts guarantee. + assert len(outs) == 1 + out = outs[0] + log(f'[api-rollout] stop_reason={out.get("stop_reason")} turns={out.get("turns")}') + + # Control flow terminated within the APIMultiTurnRollout stop-reason vocabulary. + assert out.get('stop_reason') in {'stop', 'length', 'max_turns', 'api_error'}, out.get('stop_reason') + assert isinstance(out.get('turns'), int) and out['turns'] >= 1 + + # A functional run should not have surfaced an API error. + assert out.get('stop_reason') != 'api_error', f"unexpected api_error: {out.get('error')!r}" + + # Core assertion (Requirement 5.2 / 5.5): NO trainable token-level logprobs. + assert out.get('logprobs') is None, ( + f'Gateway APIMultiTurnRollout Trajectory unexpectedly carries logprobs: ' + f'{out.get("logprobs")!r}. The OpenAI /chat/completions path is ' + f'generation-only and must not expose token-level logprobs.') + + # The per-message conversation also must not smuggle token-level logprobs + # into any assistant turn (only content / tool_calls / finish_reason are set). + for msg in out.get('messages', []): + assert 'logprobs' not in msg, ( + f'assistant/tool message unexpectedly carries logprobs: {msg!r}') + + +def test_api_multi_turn_rollout_tool_call_scenario_shape(gpu_gateway_ready): + """The tool-calling scenario produces a well-formed multi-turn conversation. + + Complements the no-logprobs assertion by checking the conversation shape: + messages accumulate across turns and, when the model invokes the calculator, + a matching ``role='tool'`` message is stitched back in. This mirrors the 9.2 + scenario so the two paths are compared on the same workload. Tool invocation + itself is model-dependent, so it is asserted conditionally. + + Validates: Requirements 5.2, 5.5, 7.2 + """ + rollout = _build_rollout(max_turns=3) + + outs = rollout([_make_trajectory()]) + out = outs[0] + + messages = out.get('messages', []) + # The full conversation includes at least the seed system+user turns plus one + # assistant reply. + assert len(messages) >= 3, messages + assert messages[0]['role'] == 'system' + assert messages[1]['role'] == 'user' + assert any(m['role'] == 'assistant' for m in messages) + + # If the model emitted a tool call, a paired tool result must follow it, and + # the calculator's deterministic answer (123 * 7 = 861) should appear. + assistant_with_tool = next( + (m for m in messages if m['role'] == 'assistant' and m.get('tool_calls')), None) + if assistant_with_tool is not None: + tool_msgs = [m for m in messages if m['role'] == 'tool'] + assert tool_msgs, 'assistant emitted tool_calls but no tool result was stitched back' + # Whichever calculator call was made, its numeric result is plain text + # (no logprobs), reinforcing the generation-only contract. + for tm in tool_msgs: + assert isinstance(tm.get('content'), str) + assert 'logprobs' not in tm + + # Regardless of tool invocation, the no-logprobs invariant still holds. + assert out.get('logprobs') is None diff --git a/tests/server/integration/test_client_multi_turn_rollout_e2e.py b/tests/server/integration/test_client_multi_turn_rollout_e2e.py new file mode 100644 index 00000000..b063f3ac --- /dev/null +++ b/tests/server/integration/test_client_multi_turn_rollout_e2e.py @@ -0,0 +1,597 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Behavioural-alignment E2E: ``ClientMultiTurnRollout`` vs bare ``MultiTurnRollout``. + +This module validates that the client-side, HTTP-driven +:class:`twinkle_client.rollout.multi_turn.ClientMultiTurnRollout` produces the +SAME multi-turn control flow as the bare-library, Ray-actor-driven +:class:`twinkle_agentic.rollout.multi_turn.MultiTurnRollout` when both are pointed +at the *same real sampler weights* (Qwen3.5-4B) with the same prompt, the same +tools, and greedy decoding (``temperature=0``). + +Why this must be GPU-gated +--------------------------- +Real numeric semantics (actual token sampling + per-token ``logprobs`` on real +model weights) cannot be reproduced by the CPU-only mock sampler used in the +local mock E2E (task 9.1). Two things in particular are only observable with a +real sampler and are therefore asserted here under GPU gating: + + * The two rollout paths agree on the ``messages`` structure and on *when* + tool calls are triggered (greedy decoding makes the control flow + deterministic across both transports, even though we allow token-level + non-determinism in principle). + * ``logprobs`` are actually populated, so the strict invariant + ``len(logprobs) == count(labels != -100)`` can be checked per trajectory on + both paths (the mock path cannot exercise real logprobs numerics). + +Topology on GPU CI +------------------ + * CLIENT path: connects to an ALREADY-RUNNING Twinkle e2e server (start it + first with ``tests/server/start_e2e_server.py``, which serves + ``tests/server/config/server_config_4b_e2e.yaml`` including the + ``sampler-Qwen3.5-4B`` application). Sampling goes over HTTP through + :class:`twinkle_client.sampler.vLLMSampler`. This mirrors the convention of + the GPU cases in ``tests/server/test_embedding_e2e.py``. + * BARE path: builds its OWN bare-library Ray-actor + :class:`twinkle.sampler.vLLMSampler` (``remote_group='sampler'`` + + ``DeviceMesh``) in-process, mirroring the standalone GPU sampler tests + (``tests/sampler/test_weight_sync.py``, ``tests/sampler/align_swift.py``) + and the multi-turn cookbook (``cookbook/rl/multi_turn/multi_turn_grpo.py``). + This is what the task means by "directly holding the corresponding Ray actor + sampler". The runner must therefore provision enough GPUs for BOTH the + server's sampler and this in-test bare sampler. + +============================================================================== +CONDA ENVIRONMENT REQUIREMENT (MANDATORY) — GPU REQUIRED +============================================================================== +This whole file is gated behind ``TWINKLE_TEST_GPU_E2E=1`` and is skipped +automatically on machines without a GPU (so it collects/skips cleanly during +local, CPU-only development). Run it inside the ``twinkle`` conda env on a GPU +host, with the e2e server already running: + + # 1) start the server (separate shell) + conda run -n twinkle python tests/server/start_e2e_server.py + + # 2) run the alignment test + TWINKLE_TEST_GPU_E2E=1 conda run -n twinkle pytest \ + tests/server/integration/test_client_multi_turn_rollout_e2e.py -v +============================================================================== +""" +from __future__ import annotations + +import copy +import json +import os +import sys +from typing import Any, Dict, List, Optional + +# Ensure project root is importable for both pytest and direct execution. +_PROJECT_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', '..')) +if _PROJECT_ROOT not in sys.path: + sys.path.insert(0, _PROJECT_ROOT) + +import pytest + +from twinkle.data_format.sampling import SamplingParams +from twinkle_agentic.tools.base import Tool +from twinkle_agentic.tools.tool_manager import ToolManager + +# Reuse the shared GPU e2e helpers (running-server URL, session init, model id). +from tests.server.integration.e2e_helpers import ( + MODEL_ID, + init_twinkle_client_session, + log, + wait_for_server, +) + + +# ═══════════════════════════════════════════════════════════════════════════ +# GPU gating +# +# Mirrors the ``gpu_e2e_enabled`` gate used by tests/server/test_embedding_e2e.py. +# Defined locally (not imported) so this file collects/skips cleanly even if the +# embedding e2e module is being edited concurrently. +# ═══════════════════════════════════════════════════════════════════════════ +def gpu_e2e_enabled() -> bool: + """Return True only when GPU e2e tests are explicitly enabled.""" + return os.environ.get('TWINKLE_TEST_GPU_E2E', '0') == '1' + + +# ═══════════════════════════════════════════════════════════════════════════ +# Configuration +# ═══════════════════════════════════════════════════════════════════════════ + +# Greedy decoding (temperature=0) makes both transports produce the SAME token +# stream for the SAME weights, so the multi-turn control flow is deterministic +# and directly comparable. ``logprobs=1`` forces per-token logprobs so the +# strict logprobs/trainable-label alignment can be asserted on both paths. +GREEDY_MAX_TOKENS = 128 +GREEDY_SEED = 0 + +# 3-6 rounds as required by the task. +ALIGN_MAX_TURNS = 6 + +# Template used locally by BOTH rollouts for encode / parse_tool_call / bridge. +TEMPLATE_CLS = 'Qwen3_5Template' + + +# ═══════════════════════════════════════════════════════════════════════════ +# A tiny ToolManager with 1-2 simple, deterministic tools (echo + calculator). +# ═══════════════════════════════════════════════════════════════════════════ +class EchoTool(Tool): + """Echoes its arguments back as a JSON string (deterministic).""" + + NAME = 'echo' + + def __call__(self, tool_name: str, arguments: Dict[str, Any]) -> str: + return f'echo[{tool_name}]:{json.dumps(arguments, sort_keys=True, ensure_ascii=False)}' + + def tool_info(self): + return { + 'type': 'function', + 'function': { + 'name': self.NAME, + 'description': 'Echo the given text back verbatim.', + 'parameters': { + 'type': 'object', + 'properties': {'text': {'type': 'string', 'description': 'Text to echo.'}}, + 'required': ['text'], + }, + }, + } + + +class CalculatorTool(Tool): + """Evaluates ``a b`` for the four basic operators (deterministic).""" + + NAME = 'calculator' + _OPS = {'add': lambda a, b: a + b, 'sub': lambda a, b: a - b, + 'mul': lambda a, b: a * b, 'div': lambda a, b: (a / b if b else float('inf'))} + + def __call__(self, tool_name: str, arguments: Dict[str, Any]) -> str: + op = str(arguments.get('op', 'add')) + a = float(arguments.get('a', 0)) + b = float(arguments.get('b', 0)) + fn = self._OPS.get(op, self._OPS['add']) + return json.dumps({'result': fn(a, b)}, ensure_ascii=False) + + def tool_info(self): + return { + 'type': 'function', + 'function': { + 'name': self.NAME, + 'description': 'Compute a basic arithmetic operation on two numbers.', + 'parameters': { + 'type': 'object', + 'properties': { + 'op': {'type': 'string', 'enum': ['add', 'sub', 'mul', 'div']}, + 'a': {'type': 'number'}, + 'b': {'type': 'number'}, + }, + 'required': ['op', 'a', 'b'], + }, + }, + } + + +def make_tool_manager() -> ToolManager: + """Build a small ToolManager holding the echo + calculator tools.""" + mgr = ToolManager({}) + mgr.register(EchoTool()) + mgr.register(CalculatorTool()) + return mgr + + +def tool_schema() -> List[Dict[str, Any]]: + """Tool schema advertised to the model in the trajectory ``tools`` field.""" + return [EchoTool().tool_info(), CalculatorTool().tool_info()] + + +SYSTEM_PROMPT = ( + 'You are a helpful assistant with access to tools. When a computation is ' + 'needed, call the `calculator` tool with the appropriate operator and ' + 'operands. When asked to repeat text, call the `echo` tool. Prefer using a ' + 'tool over answering directly, then give a short final answer.') + + +def build_alignment_trajectories() -> List[Dict[str, Any]]: + """Build a small, fixed batch of tool-inviting trajectories. + + Each trajectory carries a hidden ``_tid`` so output order can be checked, a + system prompt describing the tools, and the ``tools`` schema so the template + renders tool definitions into the prompt. + """ + users = [ + 'Please compute 21 plus 34 using your tools, then tell me the result.', + 'Use a tool to echo the exact text "ping", then confirm what you echoed.', + ] + trajectories: List[Dict[str, Any]] = [] + for tid, content in enumerate(users): + trajectories.append({ + 'messages': [ + {'role': 'system', 'content': SYSTEM_PROMPT}, + {'role': 'user', 'content': content}, + ], + 'tools': tool_schema(), + '_tid': tid, + }) + return trajectories + + +# ═══════════════════════════════════════════════════════════════════════════ +# Pure comparison helpers (collection-safe: unit-tested locally without a GPU). +# ═══════════════════════════════════════════════════════════════════════════ +def count_trainable_labels(labels: Optional[List[int]]) -> int: + """Number of trainable positions (``label != -100``) in a labels list.""" + return sum(1 for label in (labels or []) if label != -100) + + +def message_role_signature(messages: List[Dict[str, Any]]) -> List[str]: + """Ordered list of message roles — captures turn structure & tool timing. + + The presence and position of ``tool`` role messages encodes exactly when a + tool call was triggered and answered, which is the control-flow signal we + compare across the two transports. + """ + return [str(m.get('role')) for m in (messages or [])] + + +def tool_turn_indices(messages: List[Dict[str, Any]]) -> List[int]: + """Assistant-turn indices (0-based over assistant messages) that were + immediately followed by a tool response. + + This makes "when tool_calls fired" explicit and independent of textual + content: assistant turn ``k`` is a tool-calling turn iff the next message is + a ``tool`` message. + """ + roles = message_role_signature(messages) + indices: List[int] = [] + assistant_idx = -1 + for i, role in enumerate(roles): + if role == 'assistant': + assistant_idx += 1 + if i + 1 < len(roles) and roles[i + 1] == 'tool': + indices.append(assistant_idx) + return indices + + +def control_flow_fingerprint(traj: Dict[str, Any]) -> Dict[str, Any]: + """Extract the transport-independent control-flow fingerprint of an output. + + Deliberately EXCLUDES token ids / raw logprob values (which may differ) and + keeps only the control-flow-relevant fields the task requires to match: + role signature, tool-call timing, stop_reason, turns, truncated. + """ + return { + 'roles': message_role_signature(traj.get('messages') or []), + 'tool_turns': tool_turn_indices(traj.get('messages') or []), + 'stop_reason': traj.get('stop_reason'), + 'turns': traj.get('turns'), + 'truncated': bool(traj.get('truncated')), + } + + +def assert_logprobs_align_with_labels(traj: Dict[str, Any], label: str) -> None: + """Assert ``len(logprobs) == count(labels != -100)`` for a rollout output. + + Only meaningful when logprobs were requested/collected; a ``None``/empty + logprobs list means "not trainable" and is skipped (the rollout itself + already guards the non-empty case, this re-asserts it at the e2e boundary). + """ + logprobs = traj.get('logprobs') + if not logprobs: + return + trainable = count_trainable_labels(traj.get('labels')) + assert len(logprobs) == trainable, ( + f'[{label}] logprobs({len(logprobs)}) != trainable labels({trainable}) ' + f'(labels != -100)') + + +# Allowed stop reasons (Property 4 in the design doc / task 8.4). +_ALLOWED_STOP_REASONS = {'length', 'stop', 'max_turns'} + + +def assert_properties_3_4_6_7(outs: List[Dict[str, Any]], n_inputs: int, max_turns: int, label: str) -> None: + """Re-assert the multi-turn Properties 3/4/6/7 on a rollout output list. + + * Property 3 — output length & order preserved (via hidden ``_tid``). + * Property 4 — stop_reason within ``{'length','stop','max_turns'}``. + * Property 6 — actual turns never exceed max_turns. + * Property 7 — a ``max_turns`` truncation implies ``truncated=True``. + """ + # Property 3: same length and order. + assert len(outs) == n_inputs, f'[{label}] expected {n_inputs} outputs, got {len(outs)}' + for i, out in enumerate(outs): + assert out.get('_tid') == i, f'[{label}] output order mismatch at index {i}' + # Property 4: stop_reason value range. + assert out.get('stop_reason') in _ALLOWED_STOP_REASONS, \ + f"[{label}] stop_reason={out.get('stop_reason')!r} not in {_ALLOWED_STOP_REASONS}" + # Property 6: turns bounded by max_turns. + assert out.get('turns') is not None and out['turns'] <= max_turns, \ + f"[{label}] turns({out.get('turns')}) > max_turns({max_turns})" + # Property 7: max_turns truncation implies truncated=True. + if out.get('stop_reason') == 'max_turns': + assert out.get('truncated') is True, \ + f'[{label}] stop_reason==max_turns but truncated is not True' + + +# ═══════════════════════════════════════════════════════════════════════════ +# GPU-only sampler / rollout builders (never called at collection time). +# ═══════════════════════════════════════════════════════════════════════════ +def _build_local_template(): + """Build the real Qwen3.5 template used locally by BOTH rollouts.""" + from twinkle.template import Qwen3_5Template + template = Qwen3_5Template(MODEL_ID, max_length=8192, enable_thinking=False) + # Multi-turn bridge stitching does not support 'split'. + template.truncation_strategy = 'delete' + return template + + +def _greedy_sampling_params() -> SamplingParams: + """Greedy, deterministic sampling params shared by both paths.""" + return SamplingParams( + max_tokens=GREEDY_MAX_TOKENS, + temperature=0.0, + num_samples=1, + logprobs=1, + seed=GREEDY_SEED, + ) + + +def _build_client_rollout(): + """Build a ClientMultiTurnRollout wired to the running e2e server (HTTP).""" + from twinkle_client.sampler import vLLMSampler as ClientVLLMSampler + from twinkle_client.rollout.multi_turn import ClientMultiTurnRollout + + sampler = ClientVLLMSampler(model_id=MODEL_ID) + sampler.set_template(TEMPLATE_CLS, model_id=MODEL_ID, enable_thinking=False) + + return ClientMultiTurnRollout( + sampler=sampler, + template=_build_local_template(), + tool_manager=make_tool_manager(), + sampling_params=_greedy_sampling_params(), + max_turns=ALIGN_MAX_TURNS, + ) + + +def _build_bare_rollout(): + """Build a bare-library MultiTurnRollout holding its own Ray-actor sampler. + + Mirrors the standalone GPU sampler tests / multi-turn cookbook: initialise + Twinkle in Ray mode with a dedicated ``sampler`` device group and construct + a Ray-actor :class:`twinkle.sampler.vLLMSampler` (``remote_group='sampler'``). + """ + import twinkle + from twinkle import DeviceGroup, DeviceMesh + from twinkle.sampler import vLLMSampler as RayVLLMSampler + from twinkle_agentic.rollout.multi_turn import MultiTurnRollout + + sampler_gpus = int(os.environ.get('TWINKLE_ALIGN_SAMPLER_GPUS', '1')) + twinkle.initialize( + mode='ray', + nproc_per_node=sampler_gpus, + groups=[DeviceGroup(name='sampler', ranks=list(range(sampler_gpus)), device_type='GPU')], + lazy_collect=False, + ) + sampler_mesh = DeviceMesh.from_sizes(world_size=sampler_gpus, dp_size=sampler_gpus) + sampler = RayVLLMSampler( + model_id=MODEL_ID, + engine_args={ + 'gpu_memory_utilization': 0.5, + 'max_model_len': 4096, + 'enable_lora': True, + }, + device_mesh=sampler_mesh, + remote_group='sampler', + ) + sampler.set_template(TEMPLATE_CLS, model_id=MODEL_ID, enable_thinking=False) + + return MultiTurnRollout( + sampler=sampler, + template=_build_local_template(), + tool_manager=make_tool_manager(), + sampling_params=_greedy_sampling_params(), + max_turns=ALIGN_MAX_TURNS, + ) + + +# ═══════════════════════════════════════════════════════════════════════════ +# GPU fixture: build both rollouts once for the module (skips without GPU). +# ═══════════════════════════════════════════════════════════════════════════ +@pytest.fixture(scope='module') +def aligned_rollouts(): + """Yield ``(client_rollout, bare_rollout)`` built against the same weights. + + Skipped automatically unless ``TWINKLE_TEST_GPU_E2E=1``. Requires the e2e + server to be already running (client path) and enough GPUs for the in-test + bare Ray-actor sampler. + """ + if not gpu_e2e_enabled(): + pytest.skip('Set TWINKLE_TEST_GPU_E2E=1 to run the real-sampler multi-turn ' + 'alignment E2E (requires a running server + GPU).') + + wait_for_server() + init_twinkle_client_session() + + log('Building client (HTTP) rollout...') + client_rollout = _build_client_rollout() + log('Building bare-library (Ray actor) rollout...') + bare_rollout = _build_bare_rollout() + + yield client_rollout, bare_rollout + + +# ═══════════════════════════════════════════════════════════════════════════ +# GPU test: the two paths agree on control flow, and logprobs/labels align. +# ═══════════════════════════════════════════════════════════════════════════ +def test_client_and_bare_multi_turn_control_flow_aligned(aligned_rollouts): + """ClientMultiTurnRollout and bare MultiTurnRollout agree on control flow. + + Both paths run the SAME prompt / tools / greedy sampling against the SAME + real Qwen3.5-4B weights for 3-6 rounds. We assert: + + * Both satisfy Properties 3/4/6/7 (length/order, stop_reason range, + turns <= max_turns, max_turns => truncated). + * The transport-independent control-flow fingerprint (message role + signature, tool-call timing, stop_reason, turns, truncated) is IDENTICAL + between the two paths for every trajectory. Token-level sampling + non-determinism is allowed, but the control flow must match. + * On each path, ``len(logprobs) == count(labels != -100)`` for every + trajectory that collected logprobs (real-sampler numeric alignment). + + Validates: Requirements 3.12, 7.2 + """ + client_rollout, bare_rollout = aligned_rollouts + trajectories = build_alignment_trajectories() + n = len(trajectories) + + client_outs = client_rollout(copy.deepcopy(trajectories)) + bare_outs = bare_rollout(copy.deepcopy(trajectories)) + + # Per-path structural properties (3/4/6/7). + assert_properties_3_4_6_7(client_outs, n, ALIGN_MAX_TURNS, label='client') + assert_properties_3_4_6_7(bare_outs, n, ALIGN_MAX_TURNS, label='bare') + + # Per-path real-sampler logprobs/labels alignment (strict equality). + for out in client_outs: + assert_logprobs_align_with_labels(out, label='client') + for out in bare_outs: + assert_logprobs_align_with_labels(out, label='bare') + + # Cross-path control-flow equality (allowing token-level non-determinism). + for i in range(n): + client_fp = control_flow_fingerprint(client_outs[i]) + bare_fp = control_flow_fingerprint(bare_outs[i]) + assert client_fp == bare_fp, ( + f'control-flow mismatch for trajectory {i}:\n' + f' client={client_fp}\n bare ={bare_fp}') + + # The two paths must also agree on the exact trainable-label count, which + # (with greedy decoding on identical weights) is the strongest available + # cross-path numeric-alignment signal short of bit-identical logprobs. + client_trainable = count_trainable_labels(client_outs[i].get('labels')) + bare_trainable = count_trainable_labels(bare_outs[i].get('labels')) + assert client_trainable == bare_trainable, ( + f'trainable-label count mismatch for trajectory {i}: ' + f'client={client_trainable} vs bare={bare_trainable}') + + +# ═══════════════════════════════════════════════════════════════════════════ +# Collection-safe unit tests for the pure comparison helpers (run locally, no +# GPU). These prove the module imports and the fingerprint logic behaves, so the +# file is meaningful even when the GPU test above is skipped. +# ═══════════════════════════════════════════════════════════════════════════ +def test_gpu_gate_reflects_env_flag(): + """The GPU gate reflects TWINKLE_TEST_GPU_E2E without side effects.""" + assert gpu_e2e_enabled() == (os.environ.get('TWINKLE_TEST_GPU_E2E', '0') == '1') + + +def test_count_trainable_labels(): + assert count_trainable_labels(None) == 0 + assert count_trainable_labels([]) == 0 + assert count_trainable_labels([-100, -100]) == 0 + assert count_trainable_labels([-100, 5, 7, -100, 9]) == 3 + + +def test_message_role_signature_and_tool_turns(): + messages = [ + {'role': 'system', 'content': 's'}, + {'role': 'user', 'content': 'u'}, + {'role': 'assistant', 'content': 'call'}, # assistant turn 0 -> tool + {'role': 'tool', 'content': 't'}, + {'role': 'assistant', 'content': 'final'}, # assistant turn 1 -> no tool + ] + assert message_role_signature(messages) == ['system', 'user', 'assistant', 'tool', 'assistant'] + # Only the first assistant turn is immediately followed by a tool response. + assert tool_turn_indices(messages) == [0] + + +def test_control_flow_fingerprint_ignores_tokens_and_logprob_values(): + """Two outputs with identical control flow but different tokens/logprobs + produce an identical fingerprint (token-level non-determinism allowed).""" + base_messages = [ + {'role': 'user', 'content': 'q'}, + {'role': 'assistant', 'content': 'call'}, + {'role': 'tool', 'content': 't'}, + {'role': 'assistant', 'content': 'final'}, + ] + a = { + 'messages': base_messages, + 'stop_reason': 'stop', 'turns': 2, 'truncated': False, + 'labels': [-100, 1, 2], 'logprobs': [[(1, -0.1)], [(2, -0.2)]], + } + b = { + # Same roles/timing, DIFFERENT token content and logprob values. + 'messages': [dict(m) for m in base_messages], + 'stop_reason': 'stop', 'turns': 2, 'truncated': False, + 'labels': [-100, 9, 8], 'logprobs': [[(9, -1.0)], [(8, -2.0)]], + } + assert control_flow_fingerprint(a) == control_flow_fingerprint(b) + + +def test_assert_logprobs_align_with_labels_pass_and_fail(): + # Passing case: 2 logprob entries, 2 trainable labels. + ok = {'labels': [-100, 3, 4], 'logprobs': [[(3, -0.1)], [(4, -0.2)]]} + assert_logprobs_align_with_labels(ok, label='unit') # no raise + + # Empty/None logprobs is a no-op (not trainable). + assert_logprobs_align_with_labels({'labels': [-100, 3], 'logprobs': None}, label='unit') + + # Mismatch raises. + bad = {'labels': [-100, 3, 4], 'logprobs': [[(3, -0.1)]]} + with pytest.raises(AssertionError): + assert_logprobs_align_with_labels(bad, label='unit') + + +def test_assert_properties_3_4_6_7_detects_violations(): + good = [ + {'_tid': 0, 'stop_reason': 'stop', 'turns': 1, 'truncated': False}, + {'_tid': 1, 'stop_reason': 'max_turns', 'turns': 3, 'truncated': True}, + ] + assert_properties_3_4_6_7(good, n_inputs=2, max_turns=3, label='unit') # no raise + + # Property 4 violation: unknown stop_reason. + with pytest.raises(AssertionError): + assert_properties_3_4_6_7( + [{'_tid': 0, 'stop_reason': 'weird', 'turns': 1, 'truncated': False}], + n_inputs=1, max_turns=3, label='unit') + + # Property 6 violation: turns exceed max_turns. + with pytest.raises(AssertionError): + assert_properties_3_4_6_7( + [{'_tid': 0, 'stop_reason': 'stop', 'turns': 5, 'truncated': False}], + n_inputs=1, max_turns=3, label='unit') + + # Property 7 violation: max_turns but not truncated. + with pytest.raises(AssertionError): + assert_properties_3_4_6_7( + [{'_tid': 0, 'stop_reason': 'max_turns', 'turns': 3, 'truncated': False}], + n_inputs=1, max_turns=3, label='unit') + + # Property 3 violation: order/length mismatch. + with pytest.raises(AssertionError): + assert_properties_3_4_6_7( + [{'_tid': 1, 'stop_reason': 'stop', 'turns': 1, 'truncated': False}], + n_inputs=1, max_turns=3, label='unit') + + +def test_build_alignment_trajectories_shape(): + trajs = build_alignment_trajectories() + assert len(trajs) == 2 + for i, traj in enumerate(trajs): + assert traj['_tid'] == i + assert traj['messages'][0]['role'] == 'system' + assert traj['messages'][1]['role'] == 'user' + assert isinstance(traj['tools'], list) and traj['tools'] + names = {t['function']['name'] for t in traj['tools']} + assert names == {'echo', 'calculator'} + + +def test_tool_manager_dispatch_is_deterministic(): + """The echo/calculator tools dispatch deterministically via ToolManager.""" + mgr = make_tool_manager() + calc_call = {'type': 'function', 'function': {'name': 'calculator', + 'arguments': {'op': 'add', 'a': 21, 'b': 34}}} + echo_call = {'type': 'function', 'function': {'name': 'echo', 'arguments': {'text': 'ping'}}} + assert json.loads(mgr(calc_call))['result'] == 55 + assert mgr(echo_call) == 'echo[echo]:{"text": "ping"}' diff --git a/tests/server/integration/test_client_multi_turn_rollout_mock_e2e.py b/tests/server/integration/test_client_multi_turn_rollout_mock_e2e.py new file mode 100644 index 00000000..39ddb7b5 --- /dev/null +++ b/tests/server/integration/test_client_multi_turn_rollout_mock_e2e.py @@ -0,0 +1,577 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Local CPU-only multi-turn control-flow E2E for ``ClientMultiTurnRollout``. + +This test boots a REAL CPU-only Twinkle server (Ray Serve, in-process) whose +sampler backend is the enhanced :class:`MockSampler` (task 8.6), then drives +:class:`twinkle_client.rollout.multi_turn.ClientMultiTurnRollout` over ACTUAL +HTTP through :class:`twinkle_client.sampler.vLLMSampler`. No GPU is required. + +What is "real" vs "test double" here +------------------------------------ + * REAL: the HTTP server (gateway + sampler apps on Ray Serve), the + ``/twinkle/sample`` protocol hop, ``vLLMSampler.sample()`` and the + enhanced ``MockSampler`` producing ``new_input_feature`` / configurable + ``stop_reason`` / configurable tool-call text over the wire. + * TEST DOUBLE (client-local only): a lightweight char-level Template used by + the client to ``encode`` trajectories and ``parse_tool_call`` the sampled + text, plus a small ``ToolManager`` with an echo tool. The Template never + crosses the network; encoding, tool parsing and bridge stitching are + client-side concerns. A real HF Template would require a downloaded + tokenizer (network/model weights) which is unavailable in a local offline + CPU environment, so a deterministic char-level double is used instead — + exactly the established pattern from + ``tests/twinkle_client/test_client_multi_turn_rollout.py`` (task 8.4). + +Why the sampler knobs are set at CONSTRUCTION time +-------------------------------------------------- +The multi-turn knobs (``stop_reason`` / ``tool_call_text`` / ``tool_call_turns``) +CANNOT be injected per-call through the HTTP layer: the ``/twinkle/sample`` +handler rebuilds ``sampling_params`` via ``SamplingParams.from_dict`` which +filters the payload down to the known dataclass fields, dropping any extra +knobs before they reach ``MockSampler.sample``. Setting them per call would +require server-side protocol changes beyond this task's scope. Therefore each +behaviour is realised as a SEPARATE sampler app whose ``MockSampler`` is +constructed with the desired knobs: + + * ``mock-tool`` — ``stop_reason='stop'`` + a tool call injected on EVERY + round (``tool_call_turns`` covers a wide range). This deterministically + exercises the "sample -> tool -> bridge -> sample -> ... -> max_turns + truncation" control flow regardless of the sampler's monotonic per-call + round counter (so the module-scoped server can be reused across cases). + * ``mock-stop`` — ``stop_reason='stop'`` with NO tool call, i.e. natural + single-turn termination with ``stop_reason == 'stop'``. + * ``mock-length`` — ``stop_reason='length'``, i.e. immediate termination with + ``stop_reason == 'length'``. + +============================================================================== +CONDA ENVIRONMENT REQUIREMENT (MANDATORY) — NO GPU REQUIRED +============================================================================== +ALL cases in this file are CPU-only and MUST run inside the ``twinkle`` conda env: + + conda run -n twinkle pytest \ + tests/server/integration/test_client_multi_turn_rollout_mock_e2e.py -v +============================================================================== +""" +from __future__ import annotations + +import copy +import json +import os +import re +import subprocess +import sys +import time +from typing import Any, Dict, List, Optional + +# Ensure project root is importable for both pytest and direct execution. +_PROJECT_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', '..')) +if _PROJECT_ROOT not in sys.path: + sys.path.insert(0, _PROJECT_ROOT) + +import pytest + +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 + +# ═══════════════════════════════════════════════════════════════════════════ +# Sampler model ids (one app per multi-turn behaviour — see module docstring). +# ═══════════════════════════════════════════════════════════════════════════ +MODEL_TOOL = 'mock-tool' # stop_reason='stop' + tool call on every round +MODEL_STOP = 'mock-stop' # stop_reason='stop', no tool call +MODEL_LENGTH = 'mock-length' # stop_reason='length' + +# Tool call text the mock emits as SampledSequence.decoded on injected rounds. +# Must match the client Template's parse_tool_call (Qwen {...}). +TOOL_NAME = 'echo' +TOOL_CALL_TEXT = '' + json.dumps({'name': TOOL_NAME, 'arguments': {'text': 'hi'}}) + '' + +# Wide range so "inject a tool call on every round" holds even as the mock's +# monotonic per-call round counter advances across reused test cases. +_ALWAYS_INJECT_TURNS = list(range(1, 201)) + +# Sampling params carried on every round. ``max_tokens`` MUST be set: the mock +# rejects max_tokens < 1, and the rollout's default SamplingParams leaves it None. +SAMPLE_MAX_TOKENS = 4 + + +# ═══════════════════════════════════════════════════════════════════════════ +# Client-local test doubles: char-level tokenizer + Template + echo tool. +# (Mirrors tests/twinkle_client/test_client_multi_turn_rollout.py — task 8.4.) +# ═══════════════════════════════════════════════════════════════════════════ +class _FakeTokenizer: + """Char-level tokenizer with atomic special tokens. + + Guarantees ``decode(encode(s)) == s`` so ``extend_with_bridge``'s + template-space delta computation is deterministic. + """ + 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: + s += f"<|im_start|>{m['role']}\n{m['content']}<|im_end|>\n" + if add_generation_prompt: + s += '<|im_start|>assistant\n' + return self.encode(s) if tokenize else s + + +class _FakeTemplate: + """Minimal Template mirroring the parts ClientMultiTurnRollout touches. + + A client-local test double (never crosses the network). Implements exactly + ``encode`` / ``parse_tool_call`` / ``_invoke_post_pipeline`` / + ``enable_thinking`` / ``tokenizer`` — everything ``extend_with_bridge`` and + the rollout need. + """ + 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)})') + labels = labels[1:] + labels[:1] # np.roll(labels, -1): shift LEFT by 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 + + +class EchoTool(Tool): + """Echoes its arguments as a JSON string.""" + + def __init__(self, name: str = TOOL_NAME): + 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': {}}, + } + + +def _make_tool_manager() -> ToolManager: + mgr = ToolManager({}) + mgr.register(EchoTool(TOOL_NAME)) + return mgr + + +def _make_template() -> _FakeTemplate: + return _FakeTemplate(_FakeTokenizer()) + + +def _make_trajectories(n: int) -> List[Dict[str, Any]]: + """Build ``n`` single-user-message trajectories tagged with a hidden _tid.""" + return [{'messages': [{'role': 'user', 'content': f'q{tid}'}], '_tid': tid} for tid in range(n)] + + +# ═══════════════════════════════════════════════════════════════════════════ +# CPU-only mock server config. +# +# Built as plain dicts and fed DIRECTLY to the app builders (build_gateway_app / +# build_sampler_app) instead of through ``ServerConfig.model_validate``. The +# typed ``SamplerArgs`` schema uses ``extra='forbid'`` and does NOT declare the +# mock multi-turn knobs (stop_reason / tool_call_text / tool_call_turns), so +# routing them through the validated config would be rejected. The builders +# themselves accept ``**kwargs`` and forward them to the MockSampler ctor, so we +# skip validation and call them directly (this is test-only wiring). +# ═══════════════════════════════════════════════════════════════════════════ +def _sampler_app(name: str, model_id: str, *, extra_args: Dict[str, Any]) -> Dict[str, Any]: + # NOTE: ``queue_config`` is intentionally omitted. When calling the builder + # directly (bypassing ServerConfig validation) a raw dict would NOT be + # coerced into a TaskQueueConfig; leaving it unset lets the sampler app + # construct a default TaskQueueConfig, which is fine for this CPU test. + args: Dict[str, Any] = { + 'sampler_type': 'mock', + 'model_id': model_id, + 'nproc_per_node': 1, + 'device_group': {'name': f'sampler_{model_id}', 'ranks': 1, 'device_type': 'CPU'}, + 'device_mesh': {'device_type': 'CPU', 'dp_size': 1}, + } + args.update(extra_args) + return { + 'name': name, + 'route_prefix': f'/api/v1/sampler/{model_id}', + 'import_path': 'sampler', + 'args': args, + } + + +def _build_applications() -> List[Dict[str, Any]]: + """Return the plain-dict application specs (gateway + 3 sampler apps).""" + return [ + { + 'name': 'server', + 'route_prefix': '/api/v1', + 'import_path': 'server', + 'args': { + 'server_config': {'per_token_model_limit': 3}, + 'supported_models': [MODEL_TOOL, MODEL_STOP, MODEL_LENGTH], + }, + }, + _sampler_app('sampler-mock-tool', MODEL_TOOL, extra_args={ + 'stop_reason': 'stop', + 'tool_call_text': TOOL_CALL_TEXT, + 'tool_call_turns': _ALWAYS_INJECT_TURNS, + }), + _sampler_app('sampler-mock-stop', MODEL_STOP, extra_args={'stop_reason': 'stop'}), + _sampler_app('sampler-mock-length', MODEL_LENGTH, extra_args={'stop_reason': 'length'}), + ] + + +# File-backed persistence path (gateway + samplers run as separate replicas, so +# ``memory`` mode is insufficient — cross-process visibility requires a file). +_PERSISTENCE_FILE = '/tmp/twinkle_state_multi_turn_mock.json' + + +# ═══════════════════════════════════════════════════════════════════════════ +# In-process CPU-only Ray Serve harness (mirrors MockEmbeddingServerHarness). +# ═══════════════════════════════════════════════════════════════════════════ +class MultiTurnMockServerHarness: + """Boots the CPU-only mock multi-turn server in-process via Ray Serve. + + Manages its own local Ray head node (independent of the session-scoped Ray + fixture), starts Ray Serve on a randomized port, and runs the gateway plus + all three sampler apps declared by ``_build_server_config_dict``. No GPU. + """ + + READY_BUDGET_SECONDS = 90.0 + RAY_NODE_CPUS = 8 + + def __init__(self) -> None: + self.host = '127.0.0.1' + self.port = 18900 + (os.getpid() % 700) + self.base_url = f'http://{self.host}:{self.port}' + self._started = False + + @staticmethod + def _run_ray_command(*args: str) -> None: + ray_bin = os.path.join(os.path.dirname(sys.executable), 'ray') + result = subprocess.run( + [ray_bin, *args], check=False, + stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True) + if result.returncode != 0: + raise RuntimeError(f'ray {" ".join(args)} failed ({result.returncode}):\n{result.stdout}') + + def start(self) -> str: + import ray + from ray import serve + + from twinkle.server.config.persistence import PersistenceConfig + from twinkle.server.gateway import build_gateway_app + from twinkle.server.sampler import build_sampler_app + + # File-backed persistence so gateway + sampler replicas share state. + persistence_env = PersistenceConfig(mode='file', file_path=_PERSISTENCE_FILE).to_env_vars() + for k, v in persistence_env.items(): + os.environ[k] = v + + if ray.is_initialized(): + ray.shutdown() + self._run_ray_command('stop', '--force') + self._run_ray_command( + 'start', '--head', '--port=0', f'--num-cpus={self.RAY_NODE_CPUS}', + '--num-gpus=0', '--include-dashboard=false', '--disable-usage-stats') + ray.init( + address='auto', + runtime_env={'env_vars': persistence_env} if persistence_env else None, + ) + self._started = True + + serve.start(http_options={'host': self.host, 'port': self.port}) + # Call the builders DIRECTLY with plain-dict args (see _build_applications): + # this bypasses the typed SamplerArgs schema (extra='forbid') so the mock + # multi-turn knobs reach the MockSampler ctor. + builders = {'server': build_gateway_app, 'sampler': build_sampler_app} + deploy_options: Dict[str, Any] = {'ray_actor_options': {'num_cpus': 0.1}} + for app_spec in _build_applications(): + builder = builders[app_spec['import_path']] + args = dict(app_spec['args']) + if app_spec['import_path'] == 'server': + args.setdefault('http_options', {'host': self.host, 'port': self.port}) + bound = builder(deploy_options=deploy_options, **args) + serve.run(bound, name=app_spec['name'], route_prefix=app_spec['route_prefix']) + + self._wait_until_healthy(serve, self.READY_BUDGET_SECONDS) + return self.base_url + + def _wait_until_healthy(self, serve_module: Any, timeout: float) -> None: + deadline = time.monotonic() + timeout + last: Dict[str, Any] = {} + while time.monotonic() < deadline: + status = serve_module.status() + last = {name: app.status for name, app in status.applications.items()} + if last and all(s == 'RUNNING' for s in last.values()): + return + time.sleep(0.5) + raise TimeoutError(f'Mock multi-turn server not RUNNING within {timeout}s: {last}') + + def stop(self) -> None: + if not self._started: + return + try: + import ray + from ray import serve + try: + serve.shutdown() + except Exception: + pass + try: + ray.shutdown() + except Exception: + pass + finally: + try: + self._run_ray_command('stop', '--force') + except Exception: + pass + self._started = False + + +# ═══════════════════════════════════════════════════════════════════════════ +# Fixtures +# ═══════════════════════════════════════════════════════════════════════════ +@pytest.fixture(scope='module') +def mock_multi_turn_server(): + """Boot the CPU-only mock multi-turn server ONCE for the module.""" + harness = MultiTurnMockServerHarness() + base_url = harness.start() + try: + from twinkle_client import init_twinkle_client + init_twinkle_client(base_url=base_url, api_key='EMPTY_TOKEN') + yield base_url + finally: + harness.stop() + + +def _make_sampler(model_id: str): + """Create a real vLLMSampler bound to the given mock sampler app.""" + from twinkle_client.sampler import vLLMSampler + return vLLMSampler(model_id=model_id) + + +def _make_rollout(model_id: str, *, tool_manager: Optional[ToolManager], max_turns: int) -> ClientMultiTurnRollout: + return ClientMultiTurnRollout( + sampler=_make_sampler(model_id), + template=_make_template(), + tool_manager=tool_manager, + sampling_params=SamplingParams(max_tokens=SAMPLE_MAX_TOKENS, num_samples=1), + max_turns=max_turns, + ) + + +# ═══════════════════════════════════════════════════════════════════════════ +# Multi-turn control flow over real HTTP (tool app: tool call on every round). +# +# With ``stop_reason='stop'`` and a tool call injected every round, the loop +# runs "sample -> tool -> bridge -> sample -> ..." until it hits ``max_turns`` +# and force-truncates. Exercises Property 3 (len/order), Property 4 +# ('max_turns' in the allowed set), and Property 6 (turns <= max_turns). +# ═══════════════════════════════════════════════════════════════════════════ +@pytest.mark.parametrize('n_traj,max_turns', [(1, 3), (3, 2), (2, 4)]) +def test_multi_turn_tool_loop_over_http(mock_multi_turn_server, n_traj, max_turns): + """Full multi-turn control flow over real HTTP terminates at max_turns. + + Boots a real CPU server, drives ClientMultiTurnRollout via vLLMSampler HTTP + calls through several "sample -> tool -> bridge -> sample" rounds, and + checks the batch invariants. + + Validates: Requirements 3.2 (Property 3), 3.3 (Property 4), 3.5 (Property 6), 7.1, 7.2 + """ + rollout = _make_rollout(MODEL_TOOL, tool_manager=_make_tool_manager(), max_turns=max_turns) + trajectories = _make_trajectories(n_traj) + + outs = rollout(copy.deepcopy(trajectories)) + + # Property 3: output list is same length and order as the input. + assert len(outs) == n_traj + for i, out in enumerate(outs): + assert out['_tid'] == i, 'output order must match input order' + + for out in outs: + # Property 4: stop_reason is within the allowed set. + assert out['stop_reason'] in {'length', 'stop', 'max_turns'}, out['stop_reason'] + # Property 6: actual turns never exceed the configured max_turns. + assert out['turns'] <= max_turns, f"turns({out['turns']}) > max_turns({max_turns})" + # A tool call on every round means the loop must hit the turn cap. + assert out['stop_reason'] == 'max_turns' + assert out['truncated'] is True + assert out['turns'] == max_turns + # Multi-turn actually stitched tool turns via the shared bridge helper: + # the running context grew past the initial prompt encoding. + assert out['messages'][0]['role'] == 'user' + if max_turns >= 2: + # At least one tool message was appended by extend_with_bridge. + assert any(m['role'] == 'tool' for m in out['messages']) + + +# ═══════════════════════════════════════════════════════════════════════════ +# Property 7: max_turns == 1 with a first-round tool call forces truncation. +# ═══════════════════════════════════════════════════════════════════════════ +def test_property7_max_turns_one_forces_truncation(mock_multi_turn_server): + """max_turns==1 + first-round tool call -> truncated=True, stop_reason='max_turns'. + + Validates: Requirements 3.6 (Property 7), 7.1, 7.2 + """ + rollout = _make_rollout(MODEL_TOOL, tool_manager=_make_tool_manager(), max_turns=1) + trajectories = _make_trajectories(3) + + outs = rollout(copy.deepcopy(trajectories)) + + assert len(outs) == 3 + for i, out in enumerate(outs): + assert out['_tid'] == i + assert out['truncated'] is True + assert out['stop_reason'] == 'max_turns' + assert out['turns'] == 1 + + +# ═══════════════════════════════════════════════════════════════════════════ +# Property 4: natural termination reasons ('stop' and 'length') over real HTTP. +# ═══════════════════════════════════════════════════════════════════════════ +def test_natural_stop_termination_over_http(mock_multi_turn_server): + """No tool call + stop_reason='stop' -> single-turn natural termination. + + Validates: Requirements 3.2 (Property 3), 3.3 (Property 4), 3.5 (Property 6), 7.1, 7.2 + """ + rollout = _make_rollout(MODEL_STOP, tool_manager=_make_tool_manager(), max_turns=4) + trajectories = _make_trajectories(2) + + outs = rollout(copy.deepcopy(trajectories)) + + assert len(outs) == 2 + for i, out in enumerate(outs): + assert out['_tid'] == i + assert out['stop_reason'] == 'stop' + assert out['truncated'] is False + assert out['turns'] == 1 + assert out['turns'] <= 4 + + +def test_length_termination_over_http(mock_multi_turn_server): + """stop_reason='length' -> immediate termination on the first round. + + Validates: Requirements 3.3 (Property 4), 3.5 (Property 6), 7.1, 7.2 + """ + rollout = _make_rollout(MODEL_LENGTH, tool_manager=_make_tool_manager(), max_turns=4) + trajectories = _make_trajectories(2) + + outs = rollout(copy.deepcopy(trajectories)) + + assert len(outs) == 2 + for out in outs: + assert out['stop_reason'] == 'length' + assert out['truncated'] is False + assert out['turns'] == 1 + + +# ═══════════════════════════════════════════════════════════════════════════ +# Exception path (task 8.3): tool_calls produced but no tool_manager -> ValueError. +# +# Note on ``new_input_feature=None``: the enhanced MockSampler ALWAYS populates +# new_input_feature, so that specific 8.3 error path is not reproducible against +# a real mock server and is covered by the unit tests in +# tests/twinkle_client/test_client_multi_turn_rollout.py instead. The +# tool_manager-missing path IS reachable over real HTTP and is asserted here. +# ═══════════════════════════════════════════════════════════════════════════ +def test_tool_calls_without_tool_manager_raises_value_error_over_http(mock_multi_turn_server): + """A tool call with no tool_manager raises ValueError over the real HTTP path. + + Validates: Requirements 3.10, 7.1, 7.2 + """ + rollout = _make_rollout(MODEL_TOOL, tool_manager=None, max_turns=3) + trajectories = _make_trajectories(1) + + 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 diff --git a/tests/server/sampler/test_mock_sampler.py b/tests/server/sampler/test_mock_sampler.py index ea1c0766..85c467ab 100644 --- a/tests/server/sampler/test_mock_sampler.py +++ b/tests/server/sampler/test_mock_sampler.py @@ -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 (task 8.6) -------------------------- # + + +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/test_embedding_e2e.py b/tests/server/test_embedding_e2e.py new file mode 100644 index 00000000..e910ffee --- /dev/null +++ b/tests/server/test_embedding_e2e.py @@ -0,0 +1,1139 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Embedding training E2E test scaffolding (Twinkle client path). + +This module provides the reusable scaffolding for validating the embedding +training pipeline end to end through the Twinkle client: + + set_processor('InputProcessor') + -> set_loss('InfonceLoss', temperature=..., use_batch=True) + -> add_metric('EmbeddingMetric', is_training=True) + -> loop: forward_backward(inputs=mb, task='embedding') + clip_grad_and_step(...) + -> calculate_metric(is_training=True) + +It exposes TWO switchable backend entrypoints (fixtures): + + * ``mock_embedding_model`` — local CPU-only entrypoint backed by + ``TwinkleCompatMockModel`` (src/twinkle/server/model/backends/mock_model.py). + Requires NO GPU. Boots an in-process Ray Serve cluster from the CPU-only + mock server config fixture. Consumed by task 4.2 (protocol/link validation). + + * ``gpu_embedding_model`` — real transformers model entrypoint. Gated behind + the ``TWINKLE_TEST_GPU_E2E=1`` environment variable (skipped otherwise) and + connects to an already-running GPU server (see + tests/server/start_e2e_server.py). Consumed by tasks 4.3-4.5 (real numeric + validation). + +Plus two reusable helpers: + + * ``build_synthetic_contrastive_dataset(...)`` — builds a small synthetic + anchor/positive contrastive-learning dataset (even number of samples, with + ``labels = [1, 0, 1, 0, ...]`` anchor/positive semantics) in the embedding + pooling input format consumed by ``InputProcessor``. + * ``run_embedding_training(...)`` — issues the design-document "verification + call sequence" against a client model and returns the observed losses and + the final metric result. + +============================================================================== +CONDA ENVIRONMENT REQUIREMENT (MANDATORY) +============================================================================== +ALL test cases in this file MUST be run inside the ``twinkle`` conda env, e.g.: + + conda run -n twinkle pytest tests/server/test_embedding_e2e.py -v + +The local, CPU-only mock cases (task 4.2) run without a GPU and must pass in +that environment: + + conda run -n twinkle pytest tests/server/test_embedding_e2e.py -v -k mock + +The GPU-gated cases (tasks 4.3-4.5) additionally require ``TWINKLE_TEST_GPU_E2E=1`` +and a running GPU server; they are skipped automatically when the variable is +unset: + + TWINKLE_TEST_GPU_E2E=1 conda run -n twinkle pytest tests/server/test_embedding_e2e.py -v -k gpu +============================================================================== +""" +from __future__ import annotations + +import math +import os +import subprocess +import sys +import time +from typing import Any, Dict, List, Optional + +# Ensure project root is importable for both pytest and direct execution. +_PROJECT_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..')) +if _PROJECT_ROOT not in sys.path: + sys.path.insert(0, _PROJECT_ROOT) + +import pytest + +# Reuse the existing GPU e2e helpers (server URL, session init, real model id). +from tests.server.integration.e2e_helpers import ( + BASE_URL, + MODEL_ID, + init_twinkle_client_session, + log, + wait_for_server, +) + +# ═══════════════════════════════════════════════════════════════════════════ +# Configuration +# ═══════════════════════════════════════════════════════════════════════════ + +# Local CPU-only mock backend model id (see tests/server/fixtures/server_config_mock.yaml). +MOCK_MODEL_ID = 'mock-model' + +# InfoNCE temperature used by the verification call sequence. +DEFAULT_TEMPERATURE = 0.07 + +# Default synthetic-dataset shape (kept tiny for CPU speed). +DEFAULT_NUM_PAIRS = 4 +DEFAULT_SEQ_LEN = 8 +DEFAULT_VOCAB_SIZE = 32 + +# Number of training steps used by the default verification loop. +DEFAULT_TRAIN_STEPS = 4 + + +def gpu_e2e_enabled() -> bool: + """Return True when GPU e2e tests are explicitly enabled.""" + return os.environ.get('TWINKLE_TEST_GPU_E2E', '0') == '1' + + +# ═══════════════════════════════════════════════════════════════════════════ +# Synthetic contrastive-learning dataset +# ═══════════════════════════════════════════════════════════════════════════ + +def build_synthetic_contrastive_dataset( + num_pairs: int = DEFAULT_NUM_PAIRS, + *, + seq_len: int = DEFAULT_SEQ_LEN, + vocab_size: int = DEFAULT_VOCAB_SIZE, + seed: int = 0, +) -> List[Dict[str, Any]]: + """Build a small synthetic anchor/positive contrastive dataset. + + Produces ``2 * num_pairs`` samples interleaved as + ``[anchor_0, positive_0, anchor_1, positive_1, ...]`` so that the per-sample + ``labels`` scalar carries the ``[1, 0, 1, 0, ...]`` anchor/positive semantics + expected by embedding pooling (anchor -> 1, positive -> 0). Each sample is a + feature dict in the embedding pooling input format consumed by + ``InputProcessor`` (``input_ids`` + ``attention_mask`` + ``labels``). + + Anchor and its positive within a pair share correlated token content so a + real InfoNCE loss has a meaningful contrastive signal to learn from, while a + fixed ``seed`` keeps the dataset deterministic across runs. + + Args: + num_pairs: Number of anchor/positive pairs (result has ``2*num_pairs`` samples). + seq_len: Token sequence length per sample. + vocab_size: Upper bound (exclusive) for synthetic token ids. + seed: RNG seed for deterministic generation. + + Returns: + A list of feature dicts of length ``2 * num_pairs``. + """ + if num_pairs < 1: + raise ValueError(f'num_pairs must be >= 1, got {num_pairs}') + if seq_len < 1: + raise ValueError(f'seq_len must be >= 1, got {seq_len}') + + import numpy as np + + rng = np.random.default_rng(seed) + samples: List[Dict[str, Any]] = [] + for pair_idx in range(num_pairs): + # Shared "concept" tokens make anchor/positive semantically correlated. + base = rng.integers(low=1, high=vocab_size, size=seq_len) + # Anchor: the base sequence. + anchor_ids = base.copy() + # Positive: perturb a couple of positions so it is similar but not identical. + positive_ids = base.copy() + n_perturb = max(1, seq_len // 4) + perturb_pos = rng.choice(seq_len, size=n_perturb, replace=False) + positive_ids[perturb_pos] = rng.integers(low=1, high=vocab_size, size=n_perturb) + + samples.append(_make_embedding_feature(anchor_ids.tolist(), label=1)) + samples.append(_make_embedding_feature(positive_ids.tolist(), label=0)) + + return samples + + +def _make_embedding_feature(input_ids: List[int], *, label: int) -> Dict[str, Any]: + """Build a single embedding pooling feature dict. + + ``labels`` is a single scalar per sample (``[label]``) — this mirrors the + anchor/positive grouping used by the bare-library embedding example + (cookbook/exp/embedding/train_embedding_full_ddp.py) where anchors carry + ``labels=[1]`` and positives carry ``labels=[0]``. + """ + seq_len = len(input_ids) + return { + 'input_ids': list(input_ids), + 'attention_mask': [1] * seq_len, + 'labels': [int(label)], + } + + +def iter_minibatches(dataset: List[Dict[str, Any]], batch_size: int) -> List[List[Dict[str, Any]]]: + """Split a dataset into contiguous minibatches. + + ``batch_size`` should be even so anchor/positive pairs stay together inside + each minibatch (InfoNCE assumes paired samples within a batch). + """ + if batch_size < 2 or batch_size % 2 != 0: + raise ValueError(f'batch_size must be a positive even number, got {batch_size}') + return [dataset[i:i + batch_size] for i in range(0, len(dataset), batch_size)] + + +# ═══════════════════════════════════════════════════════════════════════════ +# Verification call sequence +# ═══════════════════════════════════════════════════════════════════════════ + +def configure_embedding_adapter( + model: Any, + *, + temperature: float = DEFAULT_TEMPERATURE, + use_batch: bool = True, + hard_negatives: Optional[int] = None, +) -> None: + """Apply the embedding-training configuration steps to a client model. + + Order matches the design document verification call sequence: + set_processor('InputProcessor') + -> set_loss('InfonceLoss', temperature=..., use_batch=True) + -> add_metric('EmbeddingMetric', is_training=True) + + Note: ``TransformersEmbeddingPatch`` is applied/rolled back automatically + inside ``forward`` via ``_resolve_task_context`` when ``task='embedding'`` is + passed, so there is intentionally no ``apply_patch(...)`` call here. + """ + model.set_processor('InputProcessor') + model.set_loss('InfonceLoss', temperature=temperature, use_batch=use_batch, hard_negatives=hard_negatives) + model.add_metric('EmbeddingMetric', is_training=True) + + +def extract_loss(fwd_bwd_response: Any) -> float: + """Best-effort extraction of the scalar loss from a forward_backward response. + + Handles both backend shapes: + * real transformers backend: ``result`` is a ModelOutput/dict with a + ``'loss'`` key (see TransformersModel.forward_backward). + * mock backend: ``result`` is ``[records, loss]`` (see + TwinkleCompatMockModel.forward_backward). + """ + result = getattr(fwd_bwd_response, 'result', fwd_bwd_response) + if isinstance(result, dict) and 'loss' in result: + return float(result['loss']) + if isinstance(result, (list, tuple)) and result and isinstance(result[-1], (int, float)): + return float(result[-1]) + raise AssertionError(f'Could not extract loss from forward_backward response: {result!r}') + + +def run_embedding_training( + model: Any, + minibatches: List[List[Dict[str, Any]]], + *, + temperature: float = DEFAULT_TEMPERATURE, + max_grad_norm: float = 1.0, + configure: bool = True, +) -> Dict[str, Any]: + """Run the design-document embedding-training verification call sequence. + + Steps: + (optional) set_processor -> set_loss -> add_metric + loop over minibatches: forward_backward(task='embedding') + clip_grad_and_step + calculate_metric(is_training=True) + + Args: + model: A configured client model (mock or GPU entrypoint). + minibatches: List of minibatches, each a list of embedding feature dicts. + temperature: InfoNCE temperature (used only when ``configure=True``). + max_grad_norm: Grad clipping threshold passed to ``clip_grad_and_step``. + configure: When True, apply ``configure_embedding_adapter`` first. + + Returns: + Dict with keys ``losses`` (list[float]) and ``metric`` (final metric result). + """ + if configure: + configure_embedding_adapter(model, temperature=temperature) + + losses: List[float] = [] + for step, mb in enumerate(minibatches): + fwd_bwd = model.forward_backward(inputs=mb, task='embedding') + loss = extract_loss(fwd_bwd) + losses.append(loss) + model.clip_grad_and_step(max_grad_norm=max_grad_norm) + log(f'[embedding step {step}] loss={loss:.6f}') + + metric = model.calculate_metric(is_training=True) + metric_result = getattr(metric, 'result', metric) + return {'losses': losses, 'metric': metric_result} + + +# ═══════════════════════════════════════════════════════════════════════════ +# Client model builder (shared by both entrypoints) +# ═══════════════════════════════════════════════════════════════════════════ + +def build_embedding_client_model(model_id: str, *, adapter_name: str = 'emb_adapter') -> Any: + """Create a ``MultiLoraTransformersModel`` with a LoRA adapter for embedding. + + The InfoNCE loss / EmbeddingMetric / processor configuration is applied + separately by ``configure_embedding_adapter`` so that callers can inspect or + reorder the verification call sequence. + """ + from peft import LoraConfig + from twinkle_client.model import MultiLoraTransformersModel + + model = MultiLoraTransformersModel(model_id=model_id) + model.add_adapter_to_model(adapter_name, LoraConfig(target_modules='all-linear')) + model.set_template('Qwen3_5Template') + return model + + +# ═══════════════════════════════════════════════════════════════════════════ +# Local CPU-only mock server harness +# ═══════════════════════════════════════════════════════════════════════════ + +class MockEmbeddingServerHarness: + """Boots the CPU-only mock server in-process via Ray Serve. + + Mirrors the proven boot sequence in + tests/server/integration/test_mock_mode_startup.py: it manages its own local + Ray head node (so it does not depend on the session-scoped Ray fixture), + starts Ray Serve on a randomized port, and runs every application declared in + the CPU-only mock server config fixture (server + mock model + mock sampler). + + No GPU is required. + """ + + READY_BUDGET_SECONDS = 60.0 + RAY_NODE_CPUS = 8 + + def __init__(self) -> None: + self.host = '127.0.0.1' + self.port = 18100 + (os.getpid() % 800) + self.base_url = f'http://{self.host}:{self.port}' + self._started = False + + # ----- Ray lifecycle ------------------------------------------------- # + + @staticmethod + def _run_ray_command(*args: str) -> None: + ray_bin = os.path.join(os.path.dirname(sys.executable), 'ray') + result = subprocess.run( + [ray_bin, *args], + check=False, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + ) + if result.returncode != 0: + raise RuntimeError(f'ray {" ".join(args)} failed ({result.returncode}):\n{result.stdout}') + + def start(self) -> str: + """Boot Ray + Serve + all mock apps; return the base URL when ready.""" + import ray + from ray import serve + + from tests.server.fixtures import MOCK_SERVER_CONFIG + from twinkle.server.config import ServerConfig + from twinkle.server.gateway import build_gateway_app + from twinkle.server.model import build_model_app + from twinkle.server.sampler import build_sampler_app + + cfg = ServerConfig.from_yaml(MOCK_SERVER_CONFIG) + + persistence_env: Dict[str, str] = {} + if cfg.persistence is not None: + persistence_env = cfg.persistence.to_env_vars() + for k, v in persistence_env.items(): + os.environ[k] = v + + if ray.is_initialized(): + ray.shutdown() + self._run_ray_command('stop', '--force') + self._run_ray_command( + 'start', '--head', '--port=0', f'--num-cpus={self.RAY_NODE_CPUS}', + '--num-gpus=0', '--include-dashboard=false', '--disable-usage-stats') + ray.init( + address='auto', + runtime_env={'env_vars': persistence_env} if persistence_env else None, + ) + self._started = True + + serve.start(http_options={'host': self.host, 'port': self.port}) + builders = { + 'server': build_gateway_app, + 'model': build_model_app, + 'sampler': build_sampler_app, + } + for app_spec in cfg.applications: + builder = builders[app_spec.import_path] + args = {k: v for k, v in dict(app_spec.args).items() if v is not None} + if app_spec.import_path == 'server': + http_opts = cfg.http_options.model_dump() + http_opts['host'] = self.host + http_opts['port'] = self.port + args.setdefault('http_options', http_opts) + deploy_options: Dict[str, Any] = {'ray_actor_options': {'num_cpus': 0.1}} + for raw in app_spec.deployments: + if isinstance(raw, dict): + deploy_options = { + k: v + for k, v in raw.items() if k not in ('name', 'ray_actor_options', 'autoscaling_config') + } + deploy_options['ray_actor_options'] = {'num_cpus': 0.1} + break + bound = builder(deploy_options=deploy_options, **args) + serve.run(bound, name=app_spec.name, route_prefix=app_spec.route_prefix) + + self._wait_until_healthy(serve, self.READY_BUDGET_SECONDS) + return self.base_url + + def _wait_until_healthy(self, serve_module: Any, timeout: float) -> None: + deadline = time.monotonic() + timeout + last: Dict[str, Any] = {} + while time.monotonic() < deadline: + status = serve_module.status() + last = {name: app.status for name, app in status.applications.items()} + if last and all(s == 'RUNNING' for s in last.values()): + return + time.sleep(0.5) + raise TimeoutError(f'Mock server not RUNNING within {timeout}s: {last}') + + def stop(self) -> None: + if not self._started: + return + try: + import ray + from ray import serve + try: + serve.shutdown() + except Exception: + pass + try: + ray.shutdown() + except Exception: + pass + finally: + try: + self._run_ray_command('stop', '--force') + except Exception: + pass + self._started = False + + +# ═══════════════════════════════════════════════════════════════════════════ +# Fixtures: the two switchable backend entrypoints +# ═══════════════════════════════════════════════════════════════════════════ + +@pytest.fixture(scope='module') +def mock_embedding_model(): + """Local CPU-only mock entrypoint (no GPU required). + + Boots the mock server in-process and yields a client model configured with a + LoRA adapter, ready for the embedding verification call sequence. Consumed by + task 4.2. + """ + harness = MockEmbeddingServerHarness() + base_url = harness.start() + try: + from twinkle_client import init_twinkle_client + init_twinkle_client(base_url=base_url, api_key='EMPTY_TOKEN') + model = build_embedding_client_model(MOCK_MODEL_ID) + yield model + finally: + harness.stop() + + +@pytest.fixture(scope='module') +def gpu_embedding_model(): + """GPU real-model entrypoint (gated by TWINKLE_TEST_GPU_E2E=1). + + Connects to an already-running GPU server (see start_e2e_server.py) and + yields a client model configured with a LoRA adapter. Consumed by tasks + 4.3-4.5. Skipped automatically when GPU e2e is not enabled. + """ + if not gpu_e2e_enabled(): + pytest.skip('Set TWINKLE_TEST_GPU_E2E=1 to run real GPU embedding E2E tests (requires running server)') + + wait_for_server() + init_twinkle_client_session() + model = build_embedding_client_model(MODEL_ID) + yield model + + +# ═══════════════════════════════════════════════════════════════════════════ +# Smoke tests — verify the scaffolding collects/imports and helpers behave. +# (Real protocol/numeric assertions live in tasks 4.2-4.5.) +# ═══════════════════════════════════════════════════════════════════════════ + +def test_synthetic_dataset_shape_and_labels(): + """Synthetic dataset is even-length with [1,0,1,0,...] anchor/positive labels.""" + num_pairs = 3 + dataset = build_synthetic_contrastive_dataset(num_pairs, seq_len=6, vocab_size=16, seed=123) + + assert len(dataset) == 2 * num_pairs + labels = [sample['labels'][0] for sample in dataset] + assert labels == [1, 0] * num_pairs + + for sample in dataset: + assert set(sample.keys()) == {'input_ids', 'attention_mask', 'labels'} + assert len(sample['input_ids']) == 6 + assert len(sample['attention_mask']) == 6 + assert all(m == 1 for m in sample['attention_mask']) + + +def test_synthetic_dataset_is_deterministic(): + """Same seed -> identical dataset (keeps mock/GPU comparisons reproducible).""" + a = build_synthetic_contrastive_dataset(2, seq_len=5, seed=7) + b = build_synthetic_contrastive_dataset(2, seq_len=5, seed=7) + assert a == b + + +def test_iter_minibatches_keeps_pairs_together(): + """Minibatching requires an even batch size and preserves order.""" + dataset = build_synthetic_contrastive_dataset(4, seq_len=4, seed=1) + batches = iter_minibatches(dataset, batch_size=4) + assert [len(b) for b in batches] == [4, 4] + assert batches[0] + batches[1] == dataset + + with pytest.raises(ValueError): + iter_minibatches(dataset, batch_size=3) + + +def test_extract_loss_handles_both_backend_shapes(): + """extract_loss understands mock ([records, loss]) and real ({'loss': ...}).""" + + class _Resp: + def __init__(self, result): + self.result = result + + # Mock backend shape: [records, loss] + assert extract_loss(_Resp([[{'logprobs': [0.1]}], 0.42])) == pytest.approx(0.42) + # Real backend shape: dict with 'loss' + assert extract_loss(_Resp({'loss': 1.5, 'logits': None})) == pytest.approx(1.5) + + with pytest.raises(AssertionError): + extract_loss(_Resp({'no_loss_here': 1})) + + +def test_gpu_fixture_gating_env_flag(): + """The GPU gate reflects TWINKLE_TEST_GPU_E2E without side effects.""" + assert gpu_e2e_enabled() == (os.environ.get('TWINKLE_TEST_GPU_E2E', '0') == '1') + + +# ═══════════════════════════════════════════════════════════════════════════ +# Task 4.2 — Local CPU protocol/link validation against the mock backend. +# +# These cases boot the CPU-only mock server (module-scoped ``mock_embedding_model`` +# fixture — booted ONCE, never rebooted per example) and drive the full embedding +# verification call sequence over HTTP: +# +# set_processor('InputProcessor') +# -> set_loss('InfonceLoss', temperature=..., use_batch=True) +# -> add_metric('EmbeddingMetric', is_training=True) +# -> forward_backward(inputs=mb, task='embedding') +# -> calculate_metric(is_training=True) +# +# They assert each HTTP request/response hop is well-formed, that ``task='embedding'`` +# is transmitted through the protocol layer to the mock backend without protocol +# errors, and — Property 1 (Validates: Requirements 1.1) — that the loss returned by +# ``forward_backward(task='embedding')`` is a finite number (math.isfinite: not NaN, +# not Inf). Here the mock loss layer validates the numeric-finiteness contract at the +# protocol-link level; real-model numeric semantics live in the GPU-gated tasks 4.3-4.5. +# +# Run locally (no GPU) via: +# conda run -n twinkle pytest tests/server/test_embedding_e2e.py -v -k mock +# ═══════════════════════════════════════════════════════════════════════════ + +from hypothesis import HealthCheck, given, settings +from hypothesis import strategies as st + + +def test_mock_embedding_protocol_link_full_sequence(mock_embedding_model): + """Full ordered call sequence runs over HTTP against the mock backend. + + Validates that every hop of the design-document verification call sequence + (set_processor -> set_loss -> add_metric -> forward_backward(task='embedding') + -> calculate_metric) completes without any protocol-layer exception, that + ``task='embedding'`` is accepted through the /twinkle/* protocol and reaches + the mock backend, and — Property 1 — that every returned loss is finite. + + Validates: Requirements 1.1, 1.7, 7.1, 7.2 + """ + model = mock_embedding_model + + dataset = build_synthetic_contrastive_dataset(DEFAULT_NUM_PAIRS, seq_len=DEFAULT_SEQ_LEN, seed=0) + minibatches = iter_minibatches(dataset, batch_size=2 * DEFAULT_NUM_PAIRS) + + result = run_embedding_training(model, minibatches) + + losses = result['losses'] + # One loss per minibatch, and the sequence actually issued forward_backward calls. + assert len(losses) == len(minibatches) + assert losses, 'expected at least one forward_backward loss' + + # Property 1: every embedding forward_backward loss is a finite number. + for step, loss in enumerate(losses): + assert math.isfinite(loss), f'loss at step {step} is not finite: {loss!r}' + + # calculate_metric(is_training=True) returned a well-formed metric mapping. + metric = result['metric'] + assert isinstance(metric, dict), f'expected metric dict, got {type(metric)!r}' + + +def test_mock_embedding_task_embedding_is_passed_through(mock_embedding_model): + """``task='embedding'`` traverses the HTTP protocol layer to the mock backend. + + The mock backend intentionally ignores ``task`` semantically (it only exercises + dispatch), so we assert the protocol contract instead: a ``forward_backward`` + carrying ``task='embedding'`` is accepted end-to-end and returns a well-formed + response whose extracted loss is finite. This confirms the extra kwarg is + transported (via ``ForwardRequest.model_extra``) rather than rejected by the + /twinkle/forward_backward endpoint. + + Validates: Requirements 1.1, 1.7, 7.1, 7.2 + """ + model = mock_embedding_model + configure_embedding_adapter(model) + + dataset = build_synthetic_contrastive_dataset(2, seq_len=DEFAULT_SEQ_LEN, seed=42) + minibatch = dataset # single minibatch of 4 samples (2 pairs) + + response = model.forward_backward(inputs=minibatch, task='embedding') + + # Response is the ForwardBackwardResponse pydantic model with a ``result`` field. + assert hasattr(response, 'result'), f'malformed forward_backward response: {response!r}' + loss = extract_loss(response) + assert math.isfinite(loss), f'forward_backward(task="embedding") loss not finite: {loss!r}' + + +@settings( + max_examples=12, + deadline=None, + # ``mock_embedding_model`` is module-scoped (booted once and reused across all + # generated examples); suppress the health check so hypothesis does not object + # to the shared fixture and — critically — does NOT reboot the server per example. + suppress_health_check=[HealthCheck.function_scoped_fixture], +) +@given( + num_pairs=st.integers(min_value=1, max_value=4), + seq_len=st.integers(min_value=2, max_value=12), + seed=st.integers(min_value=0, max_value=10_000), +) +def test_mock_embedding_loss_is_finite_property(mock_embedding_model, num_pairs, seq_len, seed): + """Property 1: forward_backward(task='embedding') loss is always finite. + + Varied contrastive dataset shapes (num_pairs, seq_len) and seeds are generated + within a SINGLE module-scoped mock server instance — the server is never + rebooted per example. For every generated minibatch, the loss returned by the + mock backend over the HTTP protocol link must be a finite number + (math.isfinite: not NaN, not Inf). + + Validates: Requirements 1.1 + """ + model = mock_embedding_model + # Idempotent no-op configuration on the mock backend; keeps the case self-contained. + configure_embedding_adapter(model) + + dataset = build_synthetic_contrastive_dataset(num_pairs, seq_len=seq_len, seed=seed) + # 2*num_pairs is a positive even batch size, so every anchor/positive pair stays + # together inside a single minibatch (InfoNCE assumes paired samples per batch). + minibatches = iter_minibatches(dataset, batch_size=2 * num_pairs) + + for step, mb in enumerate(minibatches): + response = model.forward_backward(inputs=mb, task='embedding') + loss = extract_loss(response) + assert math.isfinite(loss), ( + f'non-finite loss {loss!r} at step {step} ' + f'(num_pairs={num_pairs}, seq_len={seq_len}, seed={seed})') + + +# ═══════════════════════════════════════════════════════════════════════════ +# Task 4.3 — Single-GPU real-model loss finiteness (Property 1, GPU-gated). +# +# This case exercises the REAL transformers model backend through the Twinkle +# client HTTP path (module-scoped ``gpu_embedding_model`` fixture, gated behind +# TWINKLE_TEST_GPU_E2E=1 and an already-running GPU server; skipped otherwise). +# For a non-empty contrastive minibatch, the real-model +# ``forward_backward(task='embedding')`` loss must be a finite number +# (math.isfinite: not NaN, not Inf). Unlike the mock cases in task 4.2 (which +# validate the numeric-finiteness contract only at the protocol-link level), this +# asserts Property 1 against real InfoNCE numeric semantics on real model weights. +# +# Run on GPU via: +# TWINKLE_TEST_GPU_E2E=1 conda run -n twinkle pytest \ +# tests/server/test_embedding_e2e.py -v -k gpu +# ═══════════════════════════════════════════════════════════════════════════ + + +def test_gpu_embedding_real_model_loss_is_finite(gpu_embedding_model): + """Property 1: real-model forward_backward(task='embedding') loss is finite. + + Drives the design-document embedding verification call sequence + (set_processor -> set_loss('InfonceLoss', ...) -> add_metric('EmbeddingMetric') + -> forward_backward(task='embedding') + clip_grad_and_step) against the REAL + transformers backend and asserts that every loss returned by + ``forward_backward(task='embedding')`` over a non-empty contrastive minibatch + is a finite number (math.isfinite: not NaN, not Inf). + + GPU-gated (TWINKLE_TEST_GPU_E2E=1 + running GPU server): skipped automatically + on machines without a GPU, so this asserts real numeric semantics only on GPU + CI while collecting/skipping cleanly locally. + + Validates: Requirements 1.1 + """ + model = gpu_embedding_model + + dataset = build_synthetic_contrastive_dataset(DEFAULT_NUM_PAIRS, seq_len=DEFAULT_SEQ_LEN, seed=0) + minibatches = iter_minibatches(dataset, batch_size=2 * DEFAULT_NUM_PAIRS) + + # Property 1 requires a non-empty minibatch to actually exercise the loss path. + assert minibatches, 'expected at least one minibatch' + assert all(mb for mb in minibatches), 'expected every minibatch to be non-empty' + + result = run_embedding_training(model, minibatches) + + losses = result['losses'] + # One loss per minibatch, and the sequence actually issued forward_backward calls. + assert len(losses) == len(minibatches) + assert losses, 'expected at least one forward_backward loss' + + # Property 1: every real-model embedding forward_backward loss is finite. + for step, loss in enumerate(losses): + assert math.isfinite(loss), f'real-model loss at step {step} is not finite: {loss!r}' + + +# ═══════════════════════════════════════════════════════════════════════════ +# Task 4.4 — Single-GPU: TransformersEmbeddingPatch auto-rollback (Property 2). +# +# GPU-gated (TWINKLE_TEST_GPU_E2E=1). Skipped automatically on this GPU-less dev +# machine, and on any environment where the variable is unset. +# +# Property 2 (Validates: Requirements 1.2): TransformersEmbeddingPatch is applied +# *and rolled back automatically* inside a single forward call via +# ``_resolve_task_context(model, task='embedding')`` (src/twinkle/model/transformers/ +# transformers.py). While the patch is active, ``lm_head`` is swapped for identity +# and a forward hook replaces the model output with per-token hidden states +# (src/twinkle/patch/transformers_emb.py::_output_features_hook), i.e. the "logits" +# would carry the HIDDEN-STATE dimension. After the embedding forward_backward +# returns, the patch MUST be reverted, so the very next ``forward_only`` WITHOUT a +# ``task`` argument (defaults to 'causal_lm') must produce real language-model +# logits whose trailing dimension is the VOCABULARY size — never the leftover +# identity hidden states. +# +# This property depends on the real model's patch/unpatch semantics, which the +# CPU-only mock backend cannot exhibit, so it lives behind the GPU gate. +# +# Run on GPU (with a running GPU server) in the ``twinkle`` conda env: +# TWINKLE_TEST_GPU_E2E=1 conda run -n twinkle pytest tests/server/test_embedding_e2e.py -v -k gpu +# ═══════════════════════════════════════════════════════════════════════════ + + +def _innermost_dim(nested: Any) -> Optional[int]: + """Return the length of the innermost (last-axis) list of a nested list. + + ``forward_only`` returns logits that ``to_cpu_safe_output`` has converted from + a torch tensor of shape ``[B, T, V]`` (or a per-sample list of ``[T, V]``) into + plain nested Python lists. Descending to the innermost list of scalars yields + the trailing dimension ``V`` (the vocabulary size for real logits, or the hidden + size if the embedding patch had leaked). Returns ``None`` when the structure is + not a nested list (e.g. logits were suppressed to ``None``). + """ + cur = nested + while isinstance(cur, list) and cur and isinstance(cur[0], list): + cur = cur[0] + return len(cur) if isinstance(cur, list) else None + + +def _build_causal_sample(seq_len: int = DEFAULT_SEQ_LEN) -> Dict[str, Any]: + """Build a well-formed causal-LM feature (input_ids/attention_mask/labels aligned). + + Distinct from the embedding features (which carry a single scalar ``labels``): + here ``labels`` has the same length as ``input_ids`` so a causal ``forward_only`` + can compute logps without shape mismatch, and we can read back real vocab-dim + logits. Small, deterministic token ids keep it valid for any real tokenizer. + """ + input_ids = [(i % 7) + 1 for i in range(seq_len)] + return { + 'input_ids': list(input_ids), + 'attention_mask': [1] * seq_len, + 'labels': list(input_ids), + } + + +def test_gpu_embedding_patch_auto_rollback_property(gpu_embedding_model): + """Property 2: TransformersEmbeddingPatch auto-rolls back after an embedding step. + + Flow (single real GPU adapter, via the ``gpu_embedding_model`` fixture): + + 1. Configure the adapter for embedding training (set_processor -> set_loss + 'InfonceLoss' -> add_metric 'EmbeddingMetric'). + 2. Establish a baseline: call ``forward_only(return_logits=True)`` WITHOUT a + ``task`` argument BEFORE any embedding step. Since the patch is never + applied outside a ``task='embedding'`` forward, these baseline logits are + genuine vocabulary-dimension logits — the ground-truth "vocab dim". + 3. Run one real ``forward_backward(inputs=mb, task='embedding')``. This applies + TransformersEmbeddingPatch for the duration of that forward and must revert + it on the way out. + 4. Immediately call ``forward_only(return_logits=True)`` again WITHOUT ``task``. + + Assertions: + * The post-embedding logits exist and are 3D nested lists (a leaked patch + would instead make ``forward_only`` fail — the feature hook returns only + ``{'features': ...}`` with no ``logits`` key — or yield the hidden-state + dimension). + * The post-embedding logits' trailing dimension equals the baseline vocabulary + dimension, proving ``lm_head``/the forward hook were restored (not left as + identity emitting hidden states). + + Validates: Requirements 1.2, 7.2 + """ + model = gpu_embedding_model + configure_embedding_adapter(model) + + causal_sample = _build_causal_sample(seq_len=DEFAULT_SEQ_LEN) + + # (2) Baseline vocab-dim logits with the patch NEVER applied. + baseline_resp = model.forward_only(inputs=[causal_sample], return_logits=True) + baseline_result = getattr(baseline_resp, 'result', baseline_resp) + assert isinstance(baseline_result, dict), f'malformed forward_only response: {baseline_result!r}' + baseline_logits = baseline_result.get('logits') + assert baseline_logits is not None, 'baseline forward_only returned no logits (return_logits was set)' + vocab_dim = _innermost_dim(baseline_logits) + assert isinstance(vocab_dim, int) and vocab_dim > 1, ( + f'baseline logits trailing dim is not a valid vocab size: {vocab_dim!r}') + + # (3) One real embedding step: applies + must auto-rollback the patch. + emb_dataset = build_synthetic_contrastive_dataset(2, seq_len=DEFAULT_SEQ_LEN, seed=0) + fwd_bwd = model.forward_backward(inputs=emb_dataset, task='embedding') + emb_loss = extract_loss(fwd_bwd) + assert math.isfinite(emb_loss), f'embedding forward_backward loss not finite: {emb_loss!r}' + + # (4) Post-embedding causal forward_only WITHOUT task: patch must be gone. + post_resp = model.forward_only(inputs=[causal_sample], return_logits=True) + post_result = getattr(post_resp, 'result', post_resp) + assert isinstance(post_result, dict), ( + f'post-embedding forward_only returned malformed result (patch may have leaked): {post_result!r}') + post_logits = post_result.get('logits') + assert post_logits is not None, ( + 'post-embedding forward_only returned no logits — the embedding patch likely ' + 'did NOT roll back (feature hook suppresses the logits key)') + post_dim = _innermost_dim(post_logits) + + # Property 2: trailing dim is the vocab dim (identical to baseline), NOT the + # leftover identity hidden-state dim. + assert post_dim == vocab_dim, ( + f'TransformersEmbeddingPatch did NOT auto-rollback: post-embedding logits ' + f'trailing dim {post_dim!r} != baseline vocab dim {vocab_dim!r} — logits appear ' + f'to reuse identity hidden states from the embedding task') + + +# ═══════════════════════════════════════════════════════════════════════════ +# Task 4.5 — Single-GPU: bare-library parity + pos_sim upward trend (GPU-gated). +# +# These cases exercise the REAL transformers model backend and assert real +# numeric semantics that the mock backend cannot express, so they live behind +# the GPU gate (module-scoped ``gpu_embedding_model`` fixture — gated by +# TWINKLE_TEST_GPU_E2E=1 and a running GPU server; skipped automatically +# otherwise). Two properties are validated: +# +# * Requirement 1.5 — Twinkle client HTTP path vs. bare-library path parity: +# the same small synthetic contrastive dataset is trained for the same number +# of steps through (a) the Twinkle client HTTP path (task 4.1 helpers) and +# (b) the bare-library training path used by +# ``cookbook/exp/embedding/train_embedding_full_ddp.py``. The two losses must +# stay within the SAME ORDER OF MAGNITUDE (digit-for-digit equality is NOT +# required). +# +# * Requirement 1.6 — pos_sim upward trend: after several real training steps +# on a dataset with a clear contrastive signal, ``calculate_metric(is_training= +# True)`` must report an increasing anchor-positive cosine similarity +# (``pos_sim``) relative to the untrained baseline. +# +# NOTE on the "bare-library path": the published script +# ``cookbook/exp/embedding/train_embedding_full_ddp.py`` orchestrates an 8-GPU +# online-compression pipeline (Ray, vLLM condenser, real datasets) that cannot be +# executed verbatim inside a single-GPU test. This module instead reproduces the +# script's ESSENTIAL bare-library embedding-training primitives IN-PROCESS with +# the very same core-library classes it uses — +# TransformersModel/MultiLoraTransformersModel +# + set_processor(InputProcessor) +# + set_loss(InfonceLoss, temperature=..., use_batch=True, hard_negatives=None) +# + set_optimizer('AdamW', ...) +# + add_metric(EmbeddingMetric, is_training=True) +# + loop: forward_backward(inputs=mb, task='embedding') + clip_grad_and_step(...) +# + calculate_metric(is_training=True) +# — running on the SAME synthetic dataset and the SAME number of steps as the HTTP +# path. A LoRA adapter is used (instead of full fine-tuning) purely to keep the +# in-process GPU memory footprint tractable next to the running server; the loss +# numerics being compared come from the identical InfoNCE + embedding-pooling code +# path exercised by the bare script. +# +# Run on GPU via: +# TWINKLE_TEST_GPU_E2E=1 conda run -n twinkle pytest \ +# tests/server/test_embedding_e2e.py -v -k gpu +# ═══════════════════════════════════════════════════════════════════════════ + +# Dataset / training shape for the parity + trend cases (kept small for GPU speed +# while still giving InfoNCE a real contrastive signal to learn from). +CMP_NUM_PAIRS = 6 # -> 12 samples +CMP_BATCH_SIZE = 4 # -> 3 minibatches == 3 training steps for BOTH paths +TREND_NUM_PAIRS = 6 # -> 12-sample fixed eval minibatch +TREND_STEPS = 20 # real training steps to expose the pos_sim trend +TREND_LR = 1e-3 # aggressive LR so the trend is visible within few steps + + +def _metric_dict(metric_response: Any) -> Dict[str, Any]: + """Normalize a calculate_metric result to a plain dict. + + Handles the client HTTP shape (``CalculateMetricResponse`` with a ``result`` + attribute) and the bare-library local shape (a plain dict returned directly). + """ + result = getattr(metric_response, 'result', metric_response) + assert isinstance(result, dict), f'expected metric dict, got {type(result)!r}: {result!r}' + return result + + +def _parse_metric_float(metric_response: Any, key: str) -> float: + """Parse a single float metric value (EmbeddingMetric formats values as strings).""" + result = _metric_dict(metric_response) + assert key in result, f'metric {key!r} missing from result: {result!r}' + return float(result[key]) + + +def _order_of_magnitude(value: float) -> int: + """Return floor(log10(|value|)); 0 is treated as magnitude 0.""" + if value == 0: + return 0 + return int(math.floor(math.log10(abs(value)))) + + +def _assert_same_order_of_magnitude(a: float, b: float, *, label: str) -> None: + """Assert two positive losses share the same order of magnitude. + + "Same order of magnitude" is interpreted as a bounded ratio in [0.1, 10] + (equivalently, their base-10 exponents differ by at most 1). Digit-for-digit + equality is intentionally NOT required (Requirement 1.5). + """ + assert math.isfinite(a) and math.isfinite(b), f'[{label}] non-finite losses: {a!r}, {b!r}' + assert a > 0 and b > 0, f'[{label}] expected positive InfoNCE losses, got {a!r}, {b!r}' + ratio = a / b + assert 0.1 <= ratio <= 10.0, ( + f'[{label}] losses differ by more than one order of magnitude: ' + f'client={a:.6f}, bare={b:.6f}, ratio={ratio:.4f} ' + f'(oom client={_order_of_magnitude(a)}, bare={_order_of_magnitude(b)})') + log(f'[{label}] same order of magnitude: client={a:.6f}, bare={b:.6f}, ratio={ratio:.4f}') + + +def run_bare_library_embedding_training( + minibatches: List[List[Dict[str, Any]]], + *, + model_id: str = MODEL_ID, + temperature: float = DEFAULT_TEMPERATURE, + lr: float = 1e-4, + max_grad_norm: float = 1.0, + adapter_name: str = 'bare_emb_adapter', +) -> Dict[str, Any]: + """Run the bare-library embedding-training path IN-PROCESS on a single GPU. + + Reproduces the essential core-library primitives used by + ``cookbook/exp/embedding/train_embedding_full_ddp.py`` (see the module section + comment for why the full multi-GPU script cannot be executed verbatim): + + model = MultiLoraTransformersModel(model_id=...) + model.add_adapter_to_model(adapter, LoraConfig(target_modules='all-linear')) + model.set_processor(InputProcessor) + model.set_loss(InfonceLoss, temperature=..., use_batch=True, hard_negatives=None) + model.set_optimizer('AdamW', lr=...) + model.add_metric(EmbeddingMetric, is_training=True) + for mb in minibatches: + model.forward_backward(inputs=mb, task='embedding') + model.clip_grad_and_step(...) + model.calculate_metric(is_training=True) + + The model runs in ``local`` mode (the default when ``TWINKLE_MODE`` is unset), + so the ``@remote_function`` decorators call straight through and every method + requires an explicit ``adapter_name`` kwarg (unlike the client wrapper which + tracks it internally). + + Args: + minibatches: The SAME minibatches trained by the HTTP path (same order/steps). + model_id: Base model id (defaults to the shared real-model id). + temperature: InfoNCE temperature (matches the client path). + lr: AdamW learning rate. + max_grad_norm: Grad clipping threshold. + adapter_name: LoRA adapter name for the in-process model. + + Returns: + Dict with ``losses`` (list[float]) and ``metric`` (final metric dict). + """ + import gc + + from peft import LoraConfig + + from twinkle.loss import InfonceLoss + from twinkle.metric import EmbeddingMetric + from twinkle.model import MultiLoraTransformersModel + from twinkle.processor import InputProcessor + + model = None + try: + model = MultiLoraTransformersModel(model_id=model_id) + model.add_adapter_to_model(adapter_name, LoraConfig(target_modules='all-linear')) + model.set_processor(InputProcessor, adapter_name=adapter_name) + model.set_loss( + InfonceLoss, temperature=temperature, use_batch=True, hard_negatives=None, + adapter_name=adapter_name) + model.set_optimizer(optimizer_cls='AdamW', lr=lr, adapter_name=adapter_name) + model.add_metric(EmbeddingMetric, is_training=True, adapter_name=adapter_name) + + losses: List[float] = [] + for step, mb in enumerate(minibatches): + outputs = model.forward_backward(inputs=mb, task='embedding', adapter_name=adapter_name) + loss = extract_loss(outputs) + losses.append(loss) + model.clip_grad_and_step(max_grad_norm=max_grad_norm, adapter_name=adapter_name) + log(f'[bare embedding step {step}] loss={loss:.6f}') + + metric = model.calculate_metric(is_training=True, adapter_name=adapter_name) + return {'losses': losses, 'metric': _metric_dict(metric)} + finally: + # Free the in-process model's GPU memory so it does not linger next to the + # running server (the test process holds a full second copy of the weights). + try: + del model + except Exception: + pass + gc.collect() + try: + import torch + if torch.cuda.is_available(): + torch.cuda.empty_cache() + except Exception: + pass + + +def test_gpu_embedding_loss_matches_bare_library_order_of_magnitude(gpu_embedding_model): + """Requirement 1.5: client HTTP path and bare-library path losses match in magnitude. + + Trains the SAME small synthetic contrastive dataset for the SAME number of + steps through two independent paths: + + * the Twinkle client HTTP path (task 4.1 helpers), against the real GPU + server, and + * the bare-library training path reproduced in-process from + ``cookbook/exp/embedding/train_embedding_full_ddp.py``. + + Both paths use a fresh, untrained LoRA adapter, the same InfoNCE temperature, + the same AdamW learning rate, and the same minibatch sequence, so the observed + losses are directly comparable. The assertion requires only that the mean + losses stay within the SAME ORDER OF MAGNITUDE — digit-for-digit equality is + explicitly NOT required (LoRA weights are randomly initialized independently on + each path, so the two runs are not bit-identical). + + GPU-gated (TWINKLE_TEST_GPU_E2E=1 + running GPU server): skipped automatically + on machines without a GPU, so it collects/skips cleanly locally and asserts + real numeric semantics only on GPU CI. + + Validates: Requirements 1.5, 7.2 + """ + # ``gpu_embedding_model`` guarantees the GPU gate + a live server/session; use a + # dedicated fresh client adapter so the comparison is order-independent of other + # GPU cases that may have already trained the shared fixture adapter. + client_model = build_embedding_client_model(MODEL_ID, adapter_name='emb_cmp_adapter') + + dataset = build_synthetic_contrastive_dataset(CMP_NUM_PAIRS, seq_len=DEFAULT_SEQ_LEN, seed=0) + minibatches = iter_minibatches(dataset, batch_size=CMP_BATCH_SIZE) + assert minibatches, 'expected at least one minibatch' + assert all(mb for mb in minibatches), 'expected every minibatch to be non-empty' + + # ---- Client HTTP path -------------------------------------------------- # + # Configure explicitly (processor + loss + metric) and add an optimizer with a + # matching LR, then train WITHOUT re-configuring inside run_embedding_training. + configure_embedding_adapter(client_model, temperature=DEFAULT_TEMPERATURE) + client_model.set_optimizer('AdamW', lr=1e-4) + client_result = run_embedding_training(client_model, minibatches, configure=False) + client_losses = client_result['losses'] + + # ---- Bare-library path (in-process, same dataset + steps) -------------- # + bare_result = run_bare_library_embedding_training( + minibatches, temperature=DEFAULT_TEMPERATURE, lr=1e-4) + bare_losses = bare_result['losses'] + + # Both paths issued exactly one loss per (identical) minibatch. + assert len(client_losses) == len(minibatches), ( + f'client produced {len(client_losses)} losses for {len(minibatches)} minibatches') + assert len(bare_losses) == len(minibatches), ( + f'bare produced {len(bare_losses)} losses for {len(minibatches)} minibatches') + assert client_losses and bare_losses, 'both paths must produce at least one loss' + + for step, loss in enumerate(client_losses): + assert math.isfinite(loss), f'client loss at step {step} not finite: {loss!r}' + for step, loss in enumerate(bare_losses): + assert math.isfinite(loss), f'bare loss at step {step} not finite: {loss!r}' + + # Same order of magnitude on the mean loss across the identical step sequence. + client_mean = sum(client_losses) / len(client_losses) + bare_mean = sum(bare_losses) / len(bare_losses) + _assert_same_order_of_magnitude(client_mean, bare_mean, label='embedding-loss-parity') + + +def test_gpu_embedding_pos_sim_increases_after_training(gpu_embedding_model): + """Requirement 1.6: pos_sim rises after several real training steps. + + Uses a fresh, untrained LoRA adapter on the real GPU model and repeatedly + trains on a fixed synthetic contrastive minibatch (anchors and their positives + share correlated tokens, so InfoNCE has a clear signal). ``pos_sim`` (the + anchor-positive cosine similarity reported by ``EmbeddingMetric``) is measured + on the SAME fixed minibatch before training and after each step, and must show + an upward trend: the average of the last few measurements exceeds the average + of the first few. + + This is a real numeric-semantics property that the mock backend cannot express + (mock embeddings carry no contrastive signal), hence the GPU gate. + + GPU-gated (TWINKLE_TEST_GPU_E2E=1 + running GPU server): skipped automatically + on machines without a GPU. + + Validates: Requirements 1.6, 7.2 + """ + # Fresh dedicated adapter -> clean untrained baseline, independent of other cases. + model = build_embedding_client_model(MODEL_ID, adapter_name='emb_trend_adapter') + configure_embedding_adapter(model, temperature=DEFAULT_TEMPERATURE) + model.set_optimizer('AdamW', lr=TREND_LR) + + dataset = build_synthetic_contrastive_dataset(TREND_NUM_PAIRS, seq_len=DEFAULT_SEQ_LEN, seed=0) + eval_mb = dataset # single fixed minibatch reused for every measurement + assert eval_mb, 'expected a non-empty eval minibatch' + # Sanity: the dataset must contain anchors (label==1) for pos_sim to be defined. + assert any(sample['labels'][0] == 1 for sample in eval_mb), 'eval minibatch has no anchors' + + pos_sims: List[float] = [] + for step in range(TREND_STEPS + 1): + # Measure pos_sim on the CURRENT (pre-step) weights, then take a train step. + model.forward_backward(inputs=eval_mb, task='embedding') + metric = model.calculate_metric(is_training=True) + pos_sims.append(_parse_metric_float(metric, 'pos_sim')) + model.clip_grad_and_step(max_grad_norm=1.0) + + # Every measurement must be finite and a valid cosine similarity in [-1, 1]. + for step, ps in enumerate(pos_sims): + assert math.isfinite(ps), f'pos_sim at measurement {step} not finite: {ps!r}' + assert -1.0001 <= ps <= 1.0001, f'pos_sim at measurement {step} out of range: {ps!r}' + + # Upward trend: average of the last 3 measurements exceeds the first 3 + # (mirrors the robust first-avg/last-avg comparison used elsewhere in the + # e2e helpers, tolerating per-step noise). + assert len(pos_sims) >= 6, f'need >=6 pos_sim measurements, got {len(pos_sims)}' + first_avg = sum(pos_sims[:3]) / 3 + last_avg = sum(pos_sims[-3:]) / 3 + assert last_avg > first_avg, ( + f'pos_sim did NOT increase after training: ' + f'first_3_avg={first_avg:.4f} >= last_3_avg={last_avg:.4f} ' + f'(full trace: {[round(p, 4) for p in pos_sims]})') + log(f'[embedding pos_sim trend] {first_avg:.4f} -> {last_avg:.4f} ' + f'(baseline={pos_sims[0]:.4f}, final={pos_sims[-1]:.4f})') diff --git a/tests/server/test_embedding_e2e_sp_cp.py b/tests/server/test_embedding_e2e_sp_cp.py new file mode 100644 index 00000000..e849fa5b --- /dev/null +++ b/tests/server/test_embedding_e2e_sp_cp.py @@ -0,0 +1,446 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Embedding training E2E validation under SP/CP distributed configs (task 5.1). + +This module reuses the synthetic contrastive dataset and the embedding +verification call sequence from ``tests/server/test_embedding_e2e.py`` and reruns +the embedding training flow under a MULTI-GPU device mesh that enables sequence +parallel (SP, ``ulysses_size > 1``) and/or context parallel (CP, ``cp_size > 1``), +then asserts that the resulting ``pos_sim`` (anchor-positive cosine similarity) +stays within a reasonable tolerance of the single-card baseline. + +The distributed device mesh is constructed with +``DeviceMesh.from_sizes(fsdp_size=..., ulysses_size=..., cp_size=...)`` (see +``src/twinkle/utils/device_mesh.py``) and driven across ranks with +``torch.multiprocessing.spawn``, mirroring the multi-rank harness in +``tests/transformers/test_sequence_parallel_and_cp.py``. Each rank builds a real +``MultiLoraTransformersModel`` bound to that mesh; SP is enabled automatically +inside the model whenever ``device_mesh.ulysses_size > 1``. + +============================================================================== +ENVIRONMENT CONSTRAINTS (MANDATORY — READ BEFORE RUNNING) +============================================================================== +1. MULTI-GPU required. These cases spawn ``world_size`` distributed workers, one + per GPU (SP-only needs >= 2 GPUs; combined CP+SP needs >= 4 GPUs). They are + gated behind ``TWINKLE_TEST_GPU_E2E=1`` AND an actual multi-GPU CUDA runtime, + and are SKIPPED automatically on a machine without enough GPUs (including this + GPU-less dev machine). + +2. ``twinkle`` conda env required. Run every case inside the ``twinkle`` conda + env, e.g.: + + TWINKLE_TEST_GPU_E2E=1 conda run -n twinkle \ + pytest tests/server/test_embedding_e2e_sp_cp.py -v + + On a GPU-less machine the file still collects/imports cleanly and skips: + + conda run -n twinkle pytest tests/server/test_embedding_e2e_sp_cp.py -v +============================================================================== +""" +from __future__ import annotations + +import json +import math +import os +import socket +import sys +import tempfile +from typing import Any, Dict, List, Optional + +# Ensure project root is importable for both pytest and direct execution. +_PROJECT_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..')) +if _PROJECT_ROOT not in sys.path: + sys.path.insert(0, _PROJECT_ROOT) + +import pytest + +# Reuse the task 4.1-4.5 scaffolding (dataset, minibatching, model id, metric +# parsing, loss extraction, GPU gate) so the SP/CP flow exercises the SAME +# synthetic dataset and the SAME verification call sequence as the single-card +# path — the only difference is the distributed device mesh. +from tests.server.test_embedding_e2e import ( + DEFAULT_SEQ_LEN, + DEFAULT_TEMPERATURE, + MODEL_ID, + _parse_metric_float, + build_synthetic_contrastive_dataset, + extract_loss, + gpu_e2e_enabled, + iter_minibatches, +) +from tests.server.integration.e2e_helpers import log + +# ═══════════════════════════════════════════════════════════════════════════ +# Configuration +# ═══════════════════════════════════════════════════════════════════════════ + +# Dataset / training shape (kept small for GPU speed while still giving InfoNCE a +# real contrastive signal). Mirrors the pos_sim-trend shape used in task 4.5. +SP_CP_NUM_PAIRS = 6 # -> 12 samples +SP_CP_BATCH_SIZE = 4 # -> 3 minibatches == 3 training steps +SP_CP_TRAIN_STEPS = 8 # real training steps before measuring pos_sim +SP_CP_LR = 1e-3 # aggressive LR so weights actually move within few steps +SP_CP_SEED = 1234 # identical seed on baseline + distributed runs + +# ``pos_sim`` is a cosine similarity in [-1, 1]. SP/CP pooling is mathematically +# equivalent to the single-card path, so trained values should agree closely; +# a loose absolute tolerance accommodates bf16 + distributed reduction noise. +POS_SIM_ATOL = 5e-2 + +# Distributed configs exercised by the parametrized case. +# +# In the transformers backend, sequence/context parallel is driven by +# ``device_mesh.ulysses_size`` (see TransformersModel._decide_strategy: SP is +# enabled whenever ``ulysses_size > 1``). Ulysses degree is *derived* from the +# data ranks and therefore does NOT multiply the mesh world size — so +# ``from_sizes(fsdp_size=W, ulysses_size=U)`` yields exactly ``W`` ranks with a +# single ulysses group of size ``U``. Whether that group behaves as pure SP or a +# CP (ring-attention) mode is derived internally from ``ulysses_size`` vs. the +# model's attention-head count (see tests/transformers/test_sequence_parallel_and_cp.py), +# which is why both SP and CP are exercised through ``ulysses_size`` here. +# +# Each config is skipped unless its required GPU count is available. +DIST_CONFIGS = [ + # SP: 2 GPUs, ulysses_size=2 (sequence parallel enabled over 2 ranks). + {'name': 'sp2', 'world_size': 2, 'ulysses_size': 2, 'cp_size': None}, + # SP over 4 GPUs: ulysses_size=4 (larger sequence/context-parallel group). + {'name': 'sp4', 'world_size': 4, 'ulysses_size': 4, 'cp_size': None}, +] + + +def _cuda_device_count() -> int: + """Return the CUDA device count, or 0 when torch/CUDA is unavailable.""" + try: + import torch + if torch.cuda.is_available(): + return int(torch.cuda.device_count()) + except Exception: + pass + return 0 + + +def multi_gpu_e2e_enabled(min_devices: int = 2) -> bool: + """True only when GPU e2e is enabled AND enough CUDA devices are present.""" + return gpu_e2e_enabled() and _cuda_device_count() >= min_devices + + +def _find_free_port() -> int: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.bind(('127.0.0.1', 0)) + return int(sock.getsockname()[1]) + + +# ═══════════════════════════════════════════════════════════════════════════ +# Distributed worker: build an SP/CP device mesh, train, report pos_sim. +# +# Runs inside every spawned rank. Must be a module-level function so it is +# picklable by ``torch.multiprocessing.spawn`` (start method: spawn). +# ═══════════════════════════════════════════════════════════════════════════ + + +def _init_distributed(rank: int, world_size: int, port: int): + """Initialize the process group + pin this rank to its CUDA device.""" + import torch + import torch.distributed as dist + + os.environ['RANK'] = str(rank) + os.environ['WORLD_SIZE'] = str(world_size) + os.environ['LOCAL_RANK'] = str(rank) + os.environ['LOCAL_WORLD_SIZE'] = str(world_size) + os.environ['MASTER_ADDR'] = '127.0.0.1' + os.environ['MASTER_PORT'] = str(port) + + torch.cuda.set_device(rank) + device = torch.device(f'cuda:{rank}') + if not dist.is_initialized(): + dist.init_process_group( + backend='nccl', + rank=rank, + world_size=world_size, + init_method=f'tcp://127.0.0.1:{port}', + ) + return device + + +def _seed_everything(seed: int) -> None: + import torch + torch.manual_seed(seed) + if torch.cuda.is_available(): + torch.cuda.manual_seed_all(seed) + + +def _build_embedding_model_with_mesh( + world_size: int, + ulysses_size: Optional[int], + cp_size: Optional[int], + *, + adapter_name: str, + temperature: float, + lr: float, +): + """Build a real ``MultiLoraTransformersModel`` bound to an SP/CP device mesh. + + Uses the SAME bare-library primitives as the single-card path in task 4.5 + (InputProcessor + InfonceLoss + EmbeddingMetric + AdamW), the only difference + being the distributed ``DeviceMesh`` passed to the constructor. SP is enabled + automatically by the transformers backend whenever ``ulysses_size > 1``. + """ + from peft import LoraConfig + + from twinkle.loss import InfonceLoss + from twinkle.metric import EmbeddingMetric + from twinkle.model import MultiLoraTransformersModel + from twinkle.processor import InputProcessor + from twinkle.utils import DeviceMesh + + # fsdp_size = world_size shards the model across all ranks; ulysses_size + # enables sequence/context parallel derived from those data ranks (it does + # not change the world size). cp_size, when provided, adds an explicit CP + # mesh dimension (and therefore multiplies the world size). + mesh = DeviceMesh.from_sizes( + fsdp_size=world_size, + dp_size=1, + ulysses_size=ulysses_size, + cp_size=cp_size, + device_type='cuda', + ) + + model = MultiLoraTransformersModel(model_id=MODEL_ID, device_mesh=mesh) + model.add_adapter_to_model(adapter_name, LoraConfig(target_modules='all-linear')) + model.set_processor(InputProcessor, adapter_name=adapter_name) + model.set_loss( + InfonceLoss, temperature=temperature, use_batch=True, hard_negatives=None, + adapter_name=adapter_name) + model.set_optimizer(optimizer_cls='AdamW', lr=lr, adapter_name=adapter_name) + model.add_metric(EmbeddingMetric, is_training=True, adapter_name=adapter_name) + return model, adapter_name + + +def _distributed_embedding_pos_sim_worker( + rank: int, + world_size: int, + port: int, + ulysses_size: Optional[int], + cp_size: Optional[int], + seed: int, + result_path: str, +) -> None: + """Spawned worker: run embedding training under the given mesh, emit pos_sim. + + Every rank trains identically (FSDP-sharded, SP/CP-parallel). Rank 0 writes + the final ``pos_sim`` (plus the per-step losses) to ``result_path`` as JSON so + the parent process can compare it against the single-card baseline. + """ + import gc + + import torch + import torch.distributed as dist + + device = _init_distributed(rank, world_size, port) + model = None + try: + _seed_everything(seed) + + # SAME synthetic dataset + minibatch sequence as the single-card path. + dataset = build_synthetic_contrastive_dataset( + SP_CP_NUM_PAIRS, seq_len=DEFAULT_SEQ_LEN, seed=0) + minibatches = iter_minibatches(dataset, batch_size=SP_CP_BATCH_SIZE) + + adapter_name = f'emb_sp_cp_adapter_u{ulysses_size}_c{cp_size}' + model, adapter_name = _build_embedding_model_with_mesh( + world_size, ulysses_size, cp_size, + adapter_name=adapter_name, temperature=DEFAULT_TEMPERATURE, lr=SP_CP_LR) + + losses: List[float] = [] + # Repeat the (short) minibatch sequence until SP_CP_TRAIN_STEPS steps run. + step = 0 + while step < SP_CP_TRAIN_STEPS: + for mb in minibatches: + if step >= SP_CP_TRAIN_STEPS: + break + outputs = model.forward_backward(inputs=mb, task='embedding', adapter_name=adapter_name) + losses.append(extract_loss(outputs)) + model.clip_grad_and_step(max_grad_norm=1.0, adapter_name=adapter_name) + step += 1 + + metric = model.calculate_metric(is_training=True, adapter_name=adapter_name) + pos_sim = _parse_metric_float(metric, 'pos_sim') + + if rank == 0: + with open(result_path, 'w') as f: + json.dump({'pos_sim': pos_sim, 'losses': losses}, f) + log(f'[sp_cp worker u={ulysses_size} c={cp_size}] pos_sim={pos_sim:.6f} ' + f'losses={[round(x, 4) for x in losses]}') + dist.barrier() + finally: + try: + del model + except Exception: + pass + gc.collect() + try: + if torch.cuda.is_available(): + torch.cuda.empty_cache() + except Exception: + pass + if dist.is_initialized(): + dist.destroy_process_group() + + +def _run_embedding_pos_sim( + world_size: int, + ulysses_size: Optional[int], + cp_size: Optional[int], + *, + seed: int = SP_CP_SEED, +) -> Dict[str, Any]: + """Spawn ``world_size`` workers, run embedding training, return rank-0 result. + + ``world_size == 1`` with ``ulysses_size in (None, 1)`` yields the single-card + baseline (SP/CP disabled). Larger world sizes exercise the SP/CP mesh. The + spawn path is identical for both so the ONLY variable is the device mesh. + """ + import torch.multiprocessing as mp + + port = _find_free_port() + fd, result_path = tempfile.mkstemp(prefix='emb_sp_cp_', suffix='.json') + os.close(fd) + try: + mp.spawn( + _distributed_embedding_pos_sim_worker, + args=(world_size, port, ulysses_size, cp_size, seed, result_path), + nprocs=world_size, + join=True, + ) + with open(result_path) as f: + return json.load(f) + finally: + try: + os.remove(result_path) + except OSError: + pass + + +# ═══════════════════════════════════════════════════════════════════════════ +# Single-card baseline (module-scoped: computed once, reused by every config). +# ═══════════════════════════════════════════════════════════════════════════ + + +@pytest.fixture(scope='module') +def single_card_pos_sim(): + """Baseline ``pos_sim`` from a single-card (SP/CP-disabled) embedding run. + + Gated behind the multi-GPU e2e requirement so it skips cleanly on machines + without enough GPUs. Runs through the same spawn harness with + ``world_size=1`` / ``ulysses_size=1`` so the baseline and the distributed + runs differ ONLY in the device mesh. + """ + if not multi_gpu_e2e_enabled(min_devices=2): + pytest.skip( + 'Requires TWINKLE_TEST_GPU_E2E=1 and >= 2 CUDA devices ' + '(multi-GPU SP/CP embedding e2e).') + + result = _run_embedding_pos_sim(world_size=1, ulysses_size=1, cp_size=None) + pos_sim = result['pos_sim'] + log(f'[sp_cp baseline] single-card pos_sim={pos_sim:.6f}') + return pos_sim + + +# ═══════════════════════════════════════════════════════════════════════════ +# Task 5.1 — SP/CP pos_sim consistency (multi-GPU, GPU-gated). +# ═══════════════════════════════════════════════════════════════════════════ + + +@pytest.mark.parametrize('config', DIST_CONFIGS, ids=[c['name'] for c in DIST_CONFIGS]) +def test_gpu_sp_cp_pos_sim_matches_single_card(config, single_card_pos_sim): + """SP/CP pos_sim agrees with the single-card baseline within tolerance. + + Reruns the task 4.1-4.5 embedding verification flow on the SAME synthetic + contrastive dataset under a distributed device mesh built with + ``DeviceMesh.from_sizes(fsdp_size=world_size, ulysses_size=..., cp_size=...)`` + (SP when ``ulysses_size > 1``, CP when ``cp_size > 1``), then asserts the + resulting ``pos_sim`` stays within ``POS_SIM_ATOL`` of the single-card value. + + Multi-GPU + GPU-gated (TWINKLE_TEST_GPU_E2E=1): each config is skipped unless + its required GPU count is available, so the file collects/skips cleanly on a + GPU-less machine and asserts real distributed numeric semantics only on + multi-GPU CI. + + Validates: Requirements 1.3, 7.2 + """ + world_size = int(config['world_size']) + if not multi_gpu_e2e_enabled(min_devices=world_size): + pytest.skip( + f"Config {config['name']!r} needs TWINKLE_TEST_GPU_E2E=1 and " + f'>= {world_size} CUDA devices (have {_cuda_device_count()}).') + + result = _run_embedding_pos_sim( + world_size=world_size, + ulysses_size=config['ulysses_size'], + cp_size=config['cp_size'], + ) + dist_pos_sim = result['pos_sim'] + + # pos_sim is a valid cosine similarity. + assert math.isfinite(dist_pos_sim), f"non-finite pos_sim: {dist_pos_sim!r}" + assert -1.0001 <= dist_pos_sim <= 1.0001, f'pos_sim out of range: {dist_pos_sim!r}' + + delta = abs(dist_pos_sim - single_card_pos_sim) + log(f"[sp_cp {config['name']}] dist_pos_sim={dist_pos_sim:.6f} " + f'baseline={single_card_pos_sim:.6f} delta={delta:.6f}') + assert delta <= POS_SIM_ATOL, ( + f"[{config['name']}] SP/CP pos_sim diverged from single-card baseline: " + f'dist={dist_pos_sim:.6f}, baseline={single_card_pos_sim:.6f}, ' + f'delta={delta:.6f} > atol={POS_SIM_ATOL}') + + +# ═══════════════════════════════════════════════════════════════════════════ +# Local (GPU-less) collection guards — verify the file imports/collects/skips +# cleanly and the mesh-config helpers behave, WITHOUT requiring any GPU. +# ═══════════════════════════════════════════════════════════════════════════ + + +def test_multi_gpu_gate_reflects_env_and_devices(): + """The multi-GPU gate is False unless BOTH the env flag and >=N GPUs exist.""" + expected = gpu_e2e_enabled() and _cuda_device_count() >= 2 + assert multi_gpu_e2e_enabled(min_devices=2) == expected + + +def test_dist_configs_are_well_formed(): + """Every declared SP/CP config enables SP and/or CP with a matching world size.""" + for cfg in DIST_CONFIGS: + assert cfg['world_size'] >= 2, cfg + sp = (cfg['ulysses_size'] or 1) > 1 + cp = (cfg['cp_size'] or 1) > 1 + assert sp or cp, f'config enables neither SP nor CP: {cfg}' + # ulysses/cp parallelism must fit within the world size. + assert (cfg['ulysses_size'] or 1) <= cfg['world_size'], cfg + assert (cfg['cp_size'] or 1) <= cfg['world_size'], cfg + + +def test_device_mesh_from_sizes_builds_sp_cp_mesh(): + """DeviceMesh.from_sizes wires ulysses_size/cp_size into a multi-rank mesh. + + Pure CPU construction (no distributed init, no GPU): confirms the exact + ``from_sizes(fsdp_size=..., ulysses_size=..., cp_size=...)`` usage the workers + rely on. ``ulysses_size`` is derived from the data ranks and does NOT change + the mesh world size (SP degree lives inside the existing ranks), whereas + ``cp_size`` adds an explicit mesh dimension that multiplies the world size. + """ + from twinkle.utils import DeviceMesh + + # SP mesh (2 ranks, ulysses_size=2): ulysses does not inflate world_size. + sp_mesh = DeviceMesh.from_sizes(fsdp_size=2, dp_size=1, ulysses_size=2, device_type='cuda') + assert sp_mesh.world_size == 2 + assert sp_mesh.ulysses_size == 2 + + # SP mesh over 4 ranks (ulysses_size=4): still exactly fsdp_size ranks. + sp4_mesh = DeviceMesh.from_sizes(fsdp_size=4, dp_size=1, ulysses_size=4, device_type='cuda') + assert sp4_mesh.world_size == 4 + assert sp4_mesh.ulysses_size == 4 + + # CP mesh dimension (cp_size=2) multiplies the world size and is recorded as + # an explicit 'cp' dim: fsdp(2) x dp(1) x cp(2) == 4 ranks. + cp_mesh = DeviceMesh.from_sizes(fsdp_size=2, dp_size=1, cp_size=2, device_type='cuda') + assert cp_mesh.world_size == 4 + assert cp_mesh.has_dim('cp') + assert cp_mesh.get_dim_size('cp') == 2 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..91dd072b --- /dev/null +++ b/tests/twinkle_client/test_client_multi_turn_rollout.py @@ -0,0 +1,588 @@ +# 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 (see design doc "Correctness Properties"): + * Property 3 - output length & order preservation (Validates: Requirements 3.2) + * Property 4 - stop_reason value range (Validates: Requirements 3.3) + * Property 5 - logprobs / trainable-label alignment (Validates: Requirements 3.4) + * Property 6 - actual turns never exceed max_turns (Validates: Requirements 3.5) + * Property 7 - forced truncation at the max_turns edge (Validates: Requirements 3.6) +""" +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) + + +# ============================================================================= +# Property 3: multi-turn output length & order preservation (Req 3.2) +# ============================================================================= +@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_property3_output_length_and_order_preserved(scripts_spec, max_turns): + """**Validates: Requirements 3.2** + + ``__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' + + +# ============================================================================= +# Property 4: stop_reason value range (Req 3.3) +# ============================================================================= +@settings(deadline=None, max_examples=60) +@given(scripts_spec=_batch_specs(), max_turns=st.integers(min_value=1, max_value=6)) +def test_property4_stop_reason_value_range(scripts_spec, max_turns): + """**Validates: Requirements 3.3** + + 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'] + + +# ============================================================================= +# Property 5: logprobs / trainable-label alignment (Req 3.4) +# ============================================================================= +@settings(deadline=None, max_examples=60) +@given(scripts_spec=_batch_specs(), max_turns=st.integers(min_value=1, max_value=6)) +def test_property5_logprobs_align_with_trainable_labels(scripts_spec, max_turns): + """**Validates: Requirements 3.4** + + 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})') + + +# ============================================================================= +# Property 6: actual turns never exceed max_turns (Req 3.5) +# ============================================================================= +@settings(deadline=None, max_examples=60) +@given(scripts_spec=_batch_specs(), max_turns=st.integers(min_value=1, max_value=6)) +def test_property6_turns_do_not_exceed_max_turns(scripts_spec, max_turns): + """**Validates: Requirements 3.5** + + 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})" + + +# ============================================================================= +# Property 7: forced truncation at the max_turns edge (Req 3.6) +# ============================================================================= +@settings(deadline=None, max_examples=60) +@given(logprobs_flags=st.lists(st.booleans(), min_size=1, max_size=5)) +def test_property7_max_turns_one_forces_truncation(logprobs_flags): + """**Validates: Requirements 3.6** + + 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`` and requirements 3.7-3.10, 4.3-4.5, 7.1. +# 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. + + _Requirements: 3.7_ + """ + 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. + + _Requirements: 3.10, 4.5_ + """ + # 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. + + _Requirements: 3.10, 4.5_ + """ + 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). + + _Requirements: 3.9_ + """ + 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. + + _Requirements: 4.3, 4.4, 7.1_ + """ + 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 From 47fda646ac8e207ca2d4e292f343d5fc6fb417d5 Mon Sep 17 00:00:00 2001 From: Yunnglin Date: Thu, 16 Jul 2026 14:50:49 +0800 Subject: [PATCH 02/10] fix(megatron): PP-aware DP-group in set_loss + dedup _to_plain + client cookbook examples - megatron set_loss: inject _dp_group into loss.process_group so InfonceLoss all-gather stays within DP ranks under PP (prevents NCCL deadlock on earlier pipeline stages that never enter the loss). Verified: pp=2/dp=2 local + server. - dedup _to_plain: remove duplicated helper from twinkle_agentic/rollout/multi_turn.py, import from bridge.py (single source). Also remove unused numpy import. - cookbook: add client embedding training example (cookbook/client/twinkle/self_host/embedding.py) and client multi-turn rollout GRPO example (multi_turn_rollout.py). Both verified against live megatron server (pp=2/dp=2). - tests/config: add no-sampler megatron config for faster embedding-only server boots. - processor/base.py, twinkle_client/rollout, sampler, tests: branch feature updates for embedding + multi-turn client support. --- .../client/twinkle/self_host/embedding.py | 121 +++++++++++ .../twinkle/self_host/multi_turn_rollout.py | 199 ++++++++++++++++++ src/twinkle/model/megatron/megatron.py | 8 +- src/twinkle/processor/base.py | 19 +- src/twinkle_agentic/rollout/multi_turn.py | 25 +-- src/twinkle_client/rollout/multi_turn.py | 4 +- src/twinkle_client/sampler/vllm_sampler.py | 24 ++- ...ver_config_4b_e2e_megatron_no_sampler.yaml | 97 +++++++++ tests/server/test_embedding_e2e.py | 14 +- tests/server/test_embedding_e2e_sp_cp.py | 7 +- 10 files changed, 486 insertions(+), 32 deletions(-) create mode 100644 cookbook/client/twinkle/self_host/embedding.py create mode 100644 cookbook/client/twinkle/self_host/multi_turn_rollout.py create mode 100644 tests/server/config/server_config_4b_e2e_megatron_no_sampler.yaml diff --git a/cookbook/client/twinkle/self_host/embedding.py b/cookbook/client/twinkle/self_host/embedding.py new file mode 100644 index 00000000..e71c66b9 --- /dev/null +++ b/cookbook/client/twinkle/self_host/embedding.py @@ -0,0 +1,121 @@ +# 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') + +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 ========== +MODEL_ID = 'ms://Qwen/Qwen3.5-4B' +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='http://127.0.0.1:8000', api_key='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/self_host/multi_turn_rollout.py b/cookbook/client/twinkle/self_host/multi_turn_rollout.py new file mode 100644 index 00000000..b4141323 --- /dev/null +++ b/cookbook/client/twinkle/self_host/multi_turn_rollout.py @@ -0,0 +1,199 @@ +# Twinkle Client - Multi-turn agentic rollout Example +# +# Client counterpart of the ray-local core-lib example +# ``cookbook/rl/multi_turn/multi_turn_grpo.py``. Instead of a Ray ``MultiTurnRollout`` +# + ``EnvPool``, it drives the SAME multi-turn "sample -> tool -> bridge -> sample" +# loop over HTTP with ``twinkle_client.rollout.ClientMultiTurnRollout`` and the +# client ``vLLMSampler``. +# +# Differences from the ray-local example (by design): +# * EnvPool is a Ray-only component and is NOT usable from the client, so this +# example uses a self-contained Python tool (a small calculator) wrapped in a +# ``ToolManager`` instead of an interactive environment. +# * ``ClientMultiTurnRollout`` samples with ``num_samples=1``; for a GRPO group we +# replicate each prompt ``NUM_GENERATIONS`` times as separate trajectories +# (exactly what the ray-local example does with ``BATCH_SIZE * NUM_GENERATIONS``). +# +# 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 dotenv + +dotenv.load_dotenv('.env') + +from typing import Any, Dict, List + +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.tools.base import Tool +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 + +logger = get_logger() + +# ========== Configuration ========== +MODEL_ID = 'ms://Qwen/Qwen3.5-4B' +ADAPTER_NAME = 'default' +NUM_GENERATIONS = 2 # GRPO group size (rollout runs num_samples=1 per trajectory) +MAX_NEW_TOKENS = 512 +MAX_TURNS = 4 +LEARNING_RATE = 1e-5 +MAX_STEPS = 3 + +SYSTEM_PROMPT = ('You are a careful assistant. When a calculation is needed, call the ' + '`calculator` tool with a single arithmetic ``expression`` (e.g. "12*7+3") ' + 'and then give the final numeric answer.') + +# A few arithmetic questions; each is expanded into NUM_GENERATIONS trajectories. +QUESTIONS: List[Dict[str, Any]] = [ + {'q': 'What is 12 * 7 + 3?', 'answer': 87}, + {'q': 'Compute (45 - 6) * 2.', 'answer': 78}, +] + + +# ========== A self-contained tool (replaces EnvPool) ========== +class CalculatorTool(Tool): + """Evaluate a single arithmetic ``expression`` string and return the result.""" + + _ALLOWED = set('0123456789+-*/(). ') + + def tool_info(self) -> Dict[str, Any]: + return { + 'type': 'function', + 'function': { + 'name': 'calculator', + 'description': 'Evaluate one arithmetic expression and return its numeric value.', + 'parameters': { + 'type': 'object', + 'properties': { + 'expression': { + 'type': 'string', + 'description': 'The arithmetic expression to evaluate, e.g. "12*7+3".', + } + }, + 'required': ['expression'], + }, + }, + } + + def __call__(self, tool_name: str, arguments: Dict[str, Any]) -> str: + expr = str(arguments.get('expression', '')).strip() + if not expr: + return 'Error: missing "expression".' + if not set(expr) <= self._ALLOWED: + return 'Error: expression may only contain digits, spaces and + - * / ( ) .' + try: + # Safe: the whitelist above forbids names, calls and attribute access. + return str(eval(expr, {'__builtins__': {}}, {})) # noqa: S307 + except Exception as e: # noqa + return f'Error: could not evaluate {expr!r}: {type(e).__name__}' + + +def build_trajectories(tool_schema: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """Expand each question into NUM_GENERATIONS identical trajectories (GRPO group).""" + trajectories: List[Dict[str, Any]] = [] + for item in QUESTIONS: + for _ in range(NUM_GENERATIONS): + trajectories.append({ + 'messages': [ + {'role': 'system', 'content': SYSTEM_PROMPT}, + {'role': 'user', 'content': item['q']}, + ], + 'tools': tool_schema, + }) + return trajectories + + +def compute_rewards(trajectories: List[Dict[str, Any]]) -> List[float]: + """+1 if the correct numeric answer appears in the final assistant message.""" + rewards: List[float] = [] + n_per_q = NUM_GENERATIONS + for i, traj in enumerate(trajectories): + expected = str(QUESTIONS[i // n_per_q]['answer']) + final = '' + for msg in reversed(traj.get('messages', [])): + if msg.get('role') == 'assistant': + final = str(msg.get('content') or '') + break + rewards.append(1.0 if expected in final else 0.0) + return rewards + + +def train(): + # Step 1: connect to the running Twinkle server. + init_twinkle_client(base_url='http://127.0.0.1:8000', api_key='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. Needs a LOCAL template for bridge-token stitching + # (rendering tool turns + the next generation prompt) 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([CalculatorTool()]) + tool_schema = tool_manager.tool_infos() + + rollout = ClientMultiTurnRollout( + sampler=sampler, + template=rollout_template, + tool_manager=tool_manager, + 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. Run the batched multi-turn rollout (sample -> tool -> bridge -> sample). + trajectories = build_trajectories(tool_schema) + rolled = rollout(trajectories, tool_manager=tool_manager) + + # 2. Read back per-trajectory logprobs (top-1 chosen-token logprob) and rewards. + 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) + rewards = compute_rewards(rolled) + + 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/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..27d91dc8 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. diff --git a/src/twinkle_agentic/rollout/multi_turn.py b/src/twinkle_agentic/rollout/multi_turn.py index c3c82dd5..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,29 +11,7 @@ from twinkle.template.base import Template from twinkle_agentic.tools.tool_manager import ToolManager from .base import Rollout -from .bridge import extend_with_bridge - - -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() diff --git a/src/twinkle_client/rollout/multi_turn.py b/src/twinkle_client/rollout/multi_turn.py index cd42bafd..3a8c7353 100644 --- a/src/twinkle_client/rollout/multi_turn.py +++ b/src/twinkle_client/rollout/multi_turn.py @@ -107,7 +107,9 @@ def __call__(self, trajectories: List[Trajectory], **kwargs) -> List[Trajectory] kwargs.get('tool_manager', self.tool_manager), n) # 1. Encode each trajectory once; ``pifs[i]`` is the live per-turn - # state for trajectory ``i``. + # 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) diff --git a/src/twinkle_client/sampler/vllm_sampler.py b/src/twinkle_client/sampler/vllm_sampler.py index 270da067..0d553bb3 100644 --- a/src/twinkle_client/sampler/vllm_sampler.py +++ b/src/twinkle_client/sampler/vllm_sampler.py @@ -8,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. @@ -64,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/server_config_4b_e2e_megatron_no_sampler.yaml b/tests/server/config/server_config_4b_e2e_megatron_no_sampler.yaml new file mode 100644 index 00000000..a6eb93e7 --- /dev/null +++ b/tests/server/config/server_config_4b_e2e_megatron_no_sampler.yaml @@ -0,0 +1,97 @@ +# Twinkle Server Configuration - E2E Test (4B, Megatron DP=2 PP=2) - NO SAMPLER +# +# Same as server_config_4b_e2e_megatron.yaml but WITHOUT the vLLM sampler +# application. Embedding training (and any forward/backward-only workflow) does +# not need a sampler, so dropping it avoids loading a ~40GB vLLM engine and cuts +# server startup time significantly. +# +# NOTE: with no sampler the extra 1-GPU Ray worker started by +# tests/server/start_e2e_server.py is simply unused (the model's PP=2 x DP=2 +# still fits on the head node's 4 GPUs). + +proxy_location: EveryNode + +http_options: + host: 0.0.0.0 + port: 9000 + +applications: + + - name: server + route_prefix: /api/v1 + import_path: server + args: + server_config: + per_token_model_limit: 30 + supported_models: + - Qwen/Qwen3.5-4B + deployments: + - name: TinkerCompatServer + max_ongoing_requests: 50 + autoscaling_config: + min_replicas: 1 + max_replicas: 1 + target_ongoing_requests: 128 + ray_actor_options: + num_cpus: 0.1 + runtime_env: + env_vars: + TWINKLE_FAIL_FAST: "0" + + - name: models-Qwen3.5-4B + route_prefix: /api/v1/model/Qwen/Qwen3.5-4B + import_path: model + args: + backend: megatron + model_id: "ms://Qwen/Qwen3.5-4B" + max_length: 10240 + nproc_per_node: 4 + device_group: + name: model + ranks: 4 + device_type: cuda + device_mesh: + device_type: cuda + dp_size: 2 + pp_size: 2 + queue_config: + rps_limit: 100 + tps_limit: 100000 + adapter_config: + adapter_timeout: 30 + deployments: + - name: ModelManagement + autoscaling_config: + min_replicas: 1 + max_replicas: 1 + target_ongoing_requests: 16 + ray_actor_options: + num_cpus: 0.1 + runtime_env: + env_vars: + TWINKLE_TRUST_REMOTE_CODE: "1" + TWINKLE_FAIL_FAST: "0" + + - name: processor + route_prefix: /api/v1/processor + import_path: processor + args: + ncpu_proc_per_node: 2 + device_group: + name: model + ranks: 2 + device_type: CPU + device_mesh: + device_type: CPU + dp_size: 2 + deployments: + - name: ProcessorManagement + autoscaling_config: + min_replicas: 1 + max_replicas: 1 + target_ongoing_requests: 128 + ray_actor_options: + num_cpus: 0.1 + runtime_env: + env_vars: + TWINKLE_FAIL_FAST: "0" diff --git a/tests/server/test_embedding_e2e.py b/tests/server/test_embedding_e2e.py index e910ffee..68856427 100644 --- a/tests/server/test_embedding_e2e.py +++ b/tests/server/test_embedding_e2e.py @@ -695,7 +695,11 @@ def test_gpu_embedding_real_model_loss_is_finite(gpu_embedding_model): assert minibatches, 'expected at least one minibatch' assert all(mb for mb in minibatches), 'expected every minibatch to be non-empty' - result = run_embedding_training(model, minibatches) + # Configure the embedding pipeline and attach an optimizer before training: + # run_embedding_training issues clip_grad_and_step, which requires a set optimizer. + configure_embedding_adapter(model, temperature=DEFAULT_TEMPERATURE) + model.set_optimizer('AdamW', lr=1e-4) + result = run_embedding_training(model, minibatches, configure=False) losses = result['losses'] # One loss per minibatch, and the sequence actually issued forward_backward calls. @@ -796,9 +800,13 @@ def test_gpu_embedding_patch_auto_rollback_property(gpu_embedding_model): configure_embedding_adapter(model) causal_sample = _build_causal_sample(seq_len=DEFAULT_SEQ_LEN) + # Use two identical samples so the baseline/post forward_only can be split across + # a multi-DP model server (dp>=2); the assertions only inspect the logits' trailing + # vocab dimension, which is invariant to batch size. + causal_batch = [causal_sample, causal_sample] # (2) Baseline vocab-dim logits with the patch NEVER applied. - baseline_resp = model.forward_only(inputs=[causal_sample], return_logits=True) + baseline_resp = model.forward_only(inputs=causal_batch, return_logits=True) baseline_result = getattr(baseline_resp, 'result', baseline_resp) assert isinstance(baseline_result, dict), f'malformed forward_only response: {baseline_result!r}' baseline_logits = baseline_result.get('logits') @@ -814,7 +822,7 @@ def test_gpu_embedding_patch_auto_rollback_property(gpu_embedding_model): assert math.isfinite(emb_loss), f'embedding forward_backward loss not finite: {emb_loss!r}' # (4) Post-embedding causal forward_only WITHOUT task: patch must be gone. - post_resp = model.forward_only(inputs=[causal_sample], return_logits=True) + post_resp = model.forward_only(inputs=causal_batch, return_logits=True) post_result = getattr(post_resp, 'result', post_resp) assert isinstance(post_result, dict), ( f'post-embedding forward_only returned malformed result (patch may have leaked): {post_result!r}') diff --git a/tests/server/test_embedding_e2e_sp_cp.py b/tests/server/test_embedding_e2e_sp_cp.py index e849fa5b..9b676d12 100644 --- a/tests/server/test_embedding_e2e_sp_cp.py +++ b/tests/server/test_embedding_e2e_sp_cp.py @@ -204,7 +204,12 @@ def _build_embedding_model_with_mesh( device_type='cuda', ) - model = MultiLoraTransformersModel(model_id=MODEL_ID, device_mesh=mesh) + # find_unused_parameters=True is REQUIRED for embedding training: only the last + # hidden states contribute to the InfoNCE loss, so some parameters receive no + # gradient and DDP would otherwise abort the reduction (see Embedding训练.md). + model = MultiLoraTransformersModel( + model_id=MODEL_ID, device_mesh=mesh, + ddp_config={'find_unused_parameters': True}) model.add_adapter_to_model(adapter_name, LoraConfig(target_modules='all-linear')) model.set_processor(InputProcessor, adapter_name=adapter_name) model.set_loss( From 50af737b59fa6fb991220aae39d8cd3dd6da4834 Mon Sep 17 00:00:00 2001 From: Yunnglin Date: Fri, 17 Jul 2026 10:52:40 +0800 Subject: [PATCH 03/10] refactor(mock_sampler): remove unused _MULTI_TURN_KNOBS The tuple was defined but never referenced anywhere; the constructor already declares stop_reason / tool_call_text / tool_call_turns as keyword-only args, so Python's native argument binding handles them. Addresses PR #247 review. --- src/twinkle/server/sampler/backends/mock_sampler.py | 5 ----- 1 file changed, 5 deletions(-) diff --git a/src/twinkle/server/sampler/backends/mock_sampler.py b/src/twinkle/server/sampler/backends/mock_sampler.py index f04accb4..127a6dc6 100644 --- a/src/twinkle/server/sampler/backends/mock_sampler.py +++ b/src/twinkle/server/sampler/backends/mock_sampler.py @@ -23,11 +23,6 @@ logger = get_logger() -# Recognised ctor / per-call knobs that tune multi-turn control flow. Kept in a -# set so the ctor can log only *genuinely* unknown kwargs (a real backend -# signature drift) while still accepting these opt-in multi-turn knobs. -_MULTI_TURN_KNOBS = ('stop_reason', 'tool_call_text', 'tool_call_turns') - @remote_class() class MockSampler: From c71607489c004062c9b338a313e13416d1e38cf8 Mon Sep 17 00:00:00 2001 From: Yunnglin Date: Mon, 20 Jul 2026 15:07:14 +0800 Subject: [PATCH 04/10] fix --- .../server/sampler/backends/mock_sampler.py | 36 +++----- src/twinkle/server/telemetry/worker_init.py | 2 +- src/twinkle_client/rollout/multi_turn.py | 4 +- tests/server/config/test_server_config.py | 38 ++++----- ...test_api_multi_turn_rollout_no_logprobs.py | 20 ++--- .../test_client_multi_turn_rollout_e2e.py | 52 ++++++------ ...test_client_multi_turn_rollout_mock_e2e.py | 48 ++++------- .../integration/test_mock_mode_startup.py | 2 +- tests/server/model/test_mock_model.py | 18 ++-- tests/server/sampler/test_mock_sampler.py | 20 ++--- tests/server/state/test_leader_election.py | 4 +- tests/server/state/test_redis_integration.py | 12 +-- .../telemetry/test_tracing_and_correlation.py | 36 ++++---- tests/server/test_embedding_e2e.py | 82 ++++++++----------- tests/server/test_embedding_e2e_sp_cp.py | 14 ++-- tests/server/utils/task_queue/test_config.py | 14 ++-- .../test_client_multi_turn_rollout.py | 54 +++++------- 17 files changed, 196 insertions(+), 260 deletions(-) diff --git a/src/twinkle/server/sampler/backends/mock_sampler.py b/src/twinkle/server/sampler/backends/mock_sampler.py index 127a6dc6..300e5b0c 100644 --- a/src/twinkle/server/sampler/backends/mock_sampler.py +++ b/src/twinkle/server/sampler/backends/mock_sampler.py @@ -43,10 +43,10 @@ class MockSampler: 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 or matching attributes on - ``sampling_params``. 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. + 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__( @@ -104,12 +104,12 @@ 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 (precedence: explicit sample() - # kwargs > sampling_params attributes > ctor defaults) and advance the - # round counter that ``tool_call_turns`` addresses. - stop_reason = self._resolve_stop_reason(sampling_params, kwargs) - tool_call_text = self._resolve_tool_call_text(sampling_params, kwargs) - tool_call_turns = self._resolve_tool_call_turns(sampling_params, kwargs, tool_call_text) + # 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 @@ -215,35 +215,23 @@ def _normalize_turns(turns: Iterable[int] | None, tool_call_text: str | 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, params: SamplingParams | None, kwargs: dict[str, Any]) -> str: + 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']) - override = getattr(params, 'stop_reason', None) - if override is not None: - return str(override) return self._stop_reason - def _resolve_tool_call_text(self, params: SamplingParams | None, kwargs: dict[str, Any]) -> str | None: + def _resolve_tool_call_text(self, kwargs: dict[str, Any]) -> str | None: if 'tool_call_text' in kwargs: return kwargs['tool_call_text'] - override = getattr(params, 'tool_call_text', None) - if override is not None: - return override return self._tool_call_text def _resolve_tool_call_turns( self, - params: SamplingParams | None, 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) - override = getattr(params, 'tool_call_turns', None) - if override is not None: - return self._normalize_turns(override, tool_call_text) - # Fall back to the ctor-configured turns, but keep the "default to {1} - # when text is injected via a per-call override" behaviour consistent. if tool_call_text is not None and not self._tool_call_turns: return frozenset({1}) return self._tool_call_turns 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_client/rollout/multi_turn.py b/src/twinkle_client/rollout/multi_turn.py index 3a8c7353..d488736d 100644 --- a/src/twinkle_client/rollout/multi_turn.py +++ b/src/twinkle_client/rollout/multi_turn.py @@ -15,7 +15,7 @@ * 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`` (see task 8.2). + ``twinkle_agentic.rollout.bridge.extend_with_bridge``. """ import dataclasses from typing import Any, Dict, List, Optional @@ -281,7 +281,7 @@ def _resolve_tool_managers(arg, n: int) -> List[Optional[ToolManager]]: 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 (requirement 3.10). + only when a trajectory actually produces tool_calls. """ if isinstance(arg, list): if len(arg) != n: 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/gateway/test_api_multi_turn_rollout_no_logprobs.py b/tests/server/gateway/test_api_multi_turn_rollout_no_logprobs.py index 76b68cdf..789c4045 100644 --- a/tests/server/gateway/test_api_multi_turn_rollout_no_logprobs.py +++ b/tests/server/gateway/test_api_multi_turn_rollout_no_logprobs.py @@ -5,7 +5,7 @@ Purpose ------- This is the "generation-only, not trainable" counterpart to the trainable -``ClientMultiTurnRollout`` path validated in task 9.2. It drives the SAME +``ClientMultiTurnRollout`` path. It drives the SAME tool-calling scenario, but through the OpenAI-compatible route: OpenAI client -> Gateway POST /chat/completions -> /twinkle/sample @@ -15,7 +15,7 @@ surfaces assistant ``content`` / ``tool_calls`` / ``finish_reason`` — it never carries the token-level ``logprobs`` + ``new_input_feature`` alignment info that GRPO training requires. This test locks in the accuracy of that limitation -statement (Requirement 5.2 / 5.5): the Gateway indirection is fine for +statement: the Gateway indirection is fine for generation/evaluation but cannot feed token-aligned RL training. What is "real" here @@ -175,20 +175,18 @@ def gpu_gateway_ready(): # ═══════════════════════════════════════════════════════════════════════════ -# Task 9.3 — APIMultiTurnRollout via Gateway: Trajectory carries NO logprobs. +# APIMultiTurnRollout via Gateway: Trajectory carries NO logprobs. # ═══════════════════════════════════════════════════════════════════════════ def test_api_multi_turn_rollout_trajectory_has_no_logprobs(gpu_gateway_ready): """The Gateway /chat/completions path yields Trajectories without logprobs. - Runs the same tool-calling scenario as task 9.2 but through + Runs the same tool-calling scenario as the trainable client rollout but through ``APIMultiTurnRollout`` + an OpenAI-compatible client pointed at the Gateway. Asserts the returned Trajectory does NOT expose a ``logprobs`` field (absent or ``None``), confirming the "generation-only, not trainable" limitation is accurate: token-level alignment info never crosses the OpenAI protocol. GPU-gated (TWINKLE_TEST_GPU_E2E=1 + running GPU server); skipped locally. - - Validates: Requirements 5.2, 5.5, 7.2 """ rollout = _build_rollout(max_turns=3) trajectories = [_make_trajectory()] @@ -207,7 +205,7 @@ def test_api_multi_turn_rollout_trajectory_has_no_logprobs(gpu_gateway_ready): # A functional run should not have surfaced an API error. assert out.get('stop_reason') != 'api_error', f"unexpected api_error: {out.get('error')!r}" - # Core assertion (Requirement 5.2 / 5.5): NO trainable token-level logprobs. + # Core assertion: NO trainable token-level logprobs. assert out.get('logprobs') is None, ( f'Gateway APIMultiTurnRollout Trajectory unexpectedly carries logprobs: ' f'{out.get("logprobs")!r}. The OpenAI /chat/completions path is ' @@ -225,11 +223,9 @@ def test_api_multi_turn_rollout_tool_call_scenario_shape(gpu_gateway_ready): Complements the no-logprobs assertion by checking the conversation shape: messages accumulate across turns and, when the model invokes the calculator, - a matching ``role='tool'`` message is stitched back in. This mirrors the 9.2 - scenario so the two paths are compared on the same workload. Tool invocation - itself is model-dependent, so it is asserted conditionally. - - Validates: Requirements 5.2, 5.5, 7.2 + a matching ``role='tool'`` message is stitched back in. This mirrors the + trainable client rollout scenario so the two paths are compared on the same + workload. Tool invocation itself is model-dependent, so it is asserted conditionally. """ rollout = _build_rollout(max_turns=3) diff --git a/tests/server/integration/test_client_multi_turn_rollout_e2e.py b/tests/server/integration/test_client_multi_turn_rollout_e2e.py index b063f3ac..9ce78589 100644 --- a/tests/server/integration/test_client_multi_turn_rollout_e2e.py +++ b/tests/server/integration/test_client_multi_turn_rollout_e2e.py @@ -12,7 +12,7 @@ --------------------------- Real numeric semantics (actual token sampling + per-token ``logprobs`` on real model weights) cannot be reproduced by the CPU-only mock sampler used in the -local mock E2E (task 9.1). Two things in particular are only observable with a +local mock E2E. Two things in particular are only observable with a real sampler and are therefore asserted here under GPU gating: * The two rollout paths agree on the ``messages`` structure and on *when* @@ -286,29 +286,29 @@ def assert_logprobs_align_with_labels(traj: Dict[str, Any], label: str) -> None: f'(labels != -100)') -# Allowed stop reasons (Property 4 in the design doc / task 8.4). +# Allowed stop reasons. _ALLOWED_STOP_REASONS = {'length', 'stop', 'max_turns'} -def assert_properties_3_4_6_7(outs: List[Dict[str, Any]], n_inputs: int, max_turns: int, label: str) -> None: - """Re-assert the multi-turn Properties 3/4/6/7 on a rollout output list. +def assert_multi_turn_invariants(outs: List[Dict[str, Any]], n_inputs: int, max_turns: int, label: str) -> None: + """Re-assert the multi-turn output invariants on a rollout output list. - * Property 3 — output length & order preserved (via hidden ``_tid``). - * Property 4 — stop_reason within ``{'length','stop','max_turns'}``. - * Property 6 — actual turns never exceed max_turns. - * Property 7 — a ``max_turns`` truncation implies ``truncated=True``. + * Output length & order preserved (via hidden ``_tid``). + * stop_reason within ``{'length','stop','max_turns'}``. + * Actual turns never exceed max_turns. + * A ``max_turns`` truncation implies ``truncated=True``. """ - # Property 3: same length and order. + # Same length and order. assert len(outs) == n_inputs, f'[{label}] expected {n_inputs} outputs, got {len(outs)}' for i, out in enumerate(outs): assert out.get('_tid') == i, f'[{label}] output order mismatch at index {i}' - # Property 4: stop_reason value range. + # stop_reason value range. assert out.get('stop_reason') in _ALLOWED_STOP_REASONS, \ f"[{label}] stop_reason={out.get('stop_reason')!r} not in {_ALLOWED_STOP_REASONS}" - # Property 6: turns bounded by max_turns. + # Turns bounded by max_turns. assert out.get('turns') is not None and out['turns'] <= max_turns, \ f"[{label}] turns({out.get('turns')}) > max_turns({max_turns})" - # Property 7: max_turns truncation implies truncated=True. + # max_turns truncation implies truncated=True. if out.get('stop_reason') == 'max_turns': assert out.get('truncated') is True, \ f'[{label}] stop_reason==max_turns but truncated is not True' @@ -438,8 +438,6 @@ def test_client_and_bare_multi_turn_control_flow_aligned(aligned_rollouts): non-determinism is allowed, but the control flow must match. * On each path, ``len(logprobs) == count(labels != -100)`` for every trajectory that collected logprobs (real-sampler numeric alignment). - - Validates: Requirements 3.12, 7.2 """ client_rollout, bare_rollout = aligned_rollouts trajectories = build_alignment_trajectories() @@ -448,9 +446,9 @@ def test_client_and_bare_multi_turn_control_flow_aligned(aligned_rollouts): client_outs = client_rollout(copy.deepcopy(trajectories)) bare_outs = bare_rollout(copy.deepcopy(trajectories)) - # Per-path structural properties (3/4/6/7). - assert_properties_3_4_6_7(client_outs, n, ALIGN_MAX_TURNS, label='client') - assert_properties_3_4_6_7(bare_outs, n, ALIGN_MAX_TURNS, label='bare') + # Per-path structural invariants. + assert_multi_turn_invariants(client_outs, n, ALIGN_MAX_TURNS, label='client') + assert_multi_turn_invariants(bare_outs, n, ALIGN_MAX_TURNS, label='bare') # Per-path real-sampler logprobs/labels alignment (strict equality). for out in client_outs: @@ -543,34 +541,34 @@ def test_assert_logprobs_align_with_labels_pass_and_fail(): assert_logprobs_align_with_labels(bad, label='unit') -def test_assert_properties_3_4_6_7_detects_violations(): +def test_assert_multi_turn_invariants_detects_violations(): good = [ {'_tid': 0, 'stop_reason': 'stop', 'turns': 1, 'truncated': False}, {'_tid': 1, 'stop_reason': 'max_turns', 'turns': 3, 'truncated': True}, ] - assert_properties_3_4_6_7(good, n_inputs=2, max_turns=3, label='unit') # no raise + assert_multi_turn_invariants(good, n_inputs=2, max_turns=3, label='unit') # no raise - # Property 4 violation: unknown stop_reason. + # stop_reason value-range violation: unknown stop_reason. with pytest.raises(AssertionError): - assert_properties_3_4_6_7( + assert_multi_turn_invariants( [{'_tid': 0, 'stop_reason': 'weird', 'turns': 1, 'truncated': False}], n_inputs=1, max_turns=3, label='unit') - # Property 6 violation: turns exceed max_turns. + # Turn-bound violation: turns exceed max_turns. with pytest.raises(AssertionError): - assert_properties_3_4_6_7( + assert_multi_turn_invariants( [{'_tid': 0, 'stop_reason': 'stop', 'turns': 5, 'truncated': False}], n_inputs=1, max_turns=3, label='unit') - # Property 7 violation: max_turns but not truncated. + # Truncation-flag violation: max_turns but not truncated. with pytest.raises(AssertionError): - assert_properties_3_4_6_7( + assert_multi_turn_invariants( [{'_tid': 0, 'stop_reason': 'max_turns', 'turns': 3, 'truncated': False}], n_inputs=1, max_turns=3, label='unit') - # Property 3 violation: order/length mismatch. + # Order/length violation. with pytest.raises(AssertionError): - assert_properties_3_4_6_7( + assert_multi_turn_invariants( [{'_tid': 1, 'stop_reason': 'stop', 'turns': 1, 'truncated': False}], n_inputs=1, max_turns=3, label='unit') diff --git a/tests/server/integration/test_client_multi_turn_rollout_mock_e2e.py b/tests/server/integration/test_client_multi_turn_rollout_mock_e2e.py index 39ddb7b5..89e57e22 100644 --- a/tests/server/integration/test_client_multi_turn_rollout_mock_e2e.py +++ b/tests/server/integration/test_client_multi_turn_rollout_mock_e2e.py @@ -2,7 +2,7 @@ """Local CPU-only multi-turn control-flow E2E for ``ClientMultiTurnRollout``. This test boots a REAL CPU-only Twinkle server (Ray Serve, in-process) whose -sampler backend is the enhanced :class:`MockSampler` (task 8.6), then drives +sampler backend is the enhanced :class:`MockSampler`, then drives :class:`twinkle_client.rollout.multi_turn.ClientMultiTurnRollout` over ACTUAL HTTP through :class:`twinkle_client.sampler.vLLMSampler`. No GPU is required. @@ -20,7 +20,7 @@ tokenizer (network/model weights) which is unavailable in a local offline CPU environment, so a deterministic char-level double is used instead — exactly the established pattern from - ``tests/twinkle_client/test_client_multi_turn_rollout.py`` (task 8.4). + ``tests/twinkle_client/test_client_multi_turn_rollout.py``. Why the sampler knobs are set at CONSTRUCTION time -------------------------------------------------- @@ -98,7 +98,7 @@ # ═══════════════════════════════════════════════════════════════════════════ # Client-local test doubles: char-level tokenizer + Template + echo tool. -# (Mirrors tests/twinkle_client/test_client_multi_turn_rollout.py — task 8.4.) +# (Mirrors tests/twinkle_client/test_client_multi_turn_rollout.py.) # ═══════════════════════════════════════════════════════════════════════════ class _FakeTokenizer: """Char-level tokenizer with atomic special tokens. @@ -452,8 +452,8 @@ def _make_rollout(model_id: str, *, tool_manager: Optional[ToolManager], max_tur # # With ``stop_reason='stop'`` and a tool call injected every round, the loop # runs "sample -> tool -> bridge -> sample -> ..." until it hits ``max_turns`` -# and force-truncates. Exercises Property 3 (len/order), Property 4 -# ('max_turns' in the allowed set), and Property 6 (turns <= max_turns). +# and force-truncates. Exercises the output length/order, stop_reason value +# range ('max_turns' in the allowed set), and turns <= max_turns invariants. # ═══════════════════════════════════════════════════════════════════════════ @pytest.mark.parametrize('n_traj,max_turns', [(1, 3), (3, 2), (2, 4)]) def test_multi_turn_tool_loop_over_http(mock_multi_turn_server, n_traj, max_turns): @@ -462,23 +462,21 @@ def test_multi_turn_tool_loop_over_http(mock_multi_turn_server, n_traj, max_turn Boots a real CPU server, drives ClientMultiTurnRollout via vLLMSampler HTTP calls through several "sample -> tool -> bridge -> sample" rounds, and checks the batch invariants. - - Validates: Requirements 3.2 (Property 3), 3.3 (Property 4), 3.5 (Property 6), 7.1, 7.2 """ rollout = _make_rollout(MODEL_TOOL, tool_manager=_make_tool_manager(), max_turns=max_turns) trajectories = _make_trajectories(n_traj) outs = rollout(copy.deepcopy(trajectories)) - # Property 3: output list is same length and order as the input. + # Output list is same length and order as the input. assert len(outs) == n_traj for i, out in enumerate(outs): assert out['_tid'] == i, 'output order must match input order' for out in outs: - # Property 4: stop_reason is within the allowed set. + # stop_reason is within the allowed set. assert out['stop_reason'] in {'length', 'stop', 'max_turns'}, out['stop_reason'] - # Property 6: actual turns never exceed the configured max_turns. + # Actual turns never exceed the configured max_turns. assert out['turns'] <= max_turns, f"turns({out['turns']}) > max_turns({max_turns})" # A tool call on every round means the loop must hit the turn cap. assert out['stop_reason'] == 'max_turns' @@ -493,13 +491,10 @@ def test_multi_turn_tool_loop_over_http(mock_multi_turn_server, n_traj, max_turn # ═══════════════════════════════════════════════════════════════════════════ -# Property 7: max_turns == 1 with a first-round tool call forces truncation. +# max_turns == 1 with a first-round tool call forces truncation. # ═══════════════════════════════════════════════════════════════════════════ -def test_property7_max_turns_one_forces_truncation(mock_multi_turn_server): - """max_turns==1 + first-round tool call -> truncated=True, stop_reason='max_turns'. - - Validates: Requirements 3.6 (Property 7), 7.1, 7.2 - """ +def test_max_turns_one_forces_truncation(mock_multi_turn_server): + """max_turns==1 + first-round tool call -> truncated=True, stop_reason='max_turns'.""" rollout = _make_rollout(MODEL_TOOL, tool_manager=_make_tool_manager(), max_turns=1) trajectories = _make_trajectories(3) @@ -514,13 +509,10 @@ def test_property7_max_turns_one_forces_truncation(mock_multi_turn_server): # ═══════════════════════════════════════════════════════════════════════════ -# Property 4: natural termination reasons ('stop' and 'length') over real HTTP. +# Natural termination reasons ('stop' and 'length') over real HTTP. # ═══════════════════════════════════════════════════════════════════════════ def test_natural_stop_termination_over_http(mock_multi_turn_server): - """No tool call + stop_reason='stop' -> single-turn natural termination. - - Validates: Requirements 3.2 (Property 3), 3.3 (Property 4), 3.5 (Property 6), 7.1, 7.2 - """ + """No tool call + stop_reason='stop' -> single-turn natural termination.""" rollout = _make_rollout(MODEL_STOP, tool_manager=_make_tool_manager(), max_turns=4) trajectories = _make_trajectories(2) @@ -536,10 +528,7 @@ def test_natural_stop_termination_over_http(mock_multi_turn_server): def test_length_termination_over_http(mock_multi_turn_server): - """stop_reason='length' -> immediate termination on the first round. - - Validates: Requirements 3.3 (Property 4), 3.5 (Property 6), 7.1, 7.2 - """ + """stop_reason='length' -> immediate termination on the first round.""" rollout = _make_rollout(MODEL_LENGTH, tool_manager=_make_tool_manager(), max_turns=4) trajectories = _make_trajectories(2) @@ -553,19 +542,16 @@ def test_length_termination_over_http(mock_multi_turn_server): # ═══════════════════════════════════════════════════════════════════════════ -# Exception path (task 8.3): tool_calls produced but no tool_manager -> ValueError. +# Exception path: tool_calls produced but no tool_manager -> ValueError. # # Note on ``new_input_feature=None``: the enhanced MockSampler ALWAYS populates -# new_input_feature, so that specific 8.3 error path is not reproducible against +# new_input_feature, so that specific error path is not reproducible against # a real mock server and is covered by the unit tests in # tests/twinkle_client/test_client_multi_turn_rollout.py instead. The # tool_manager-missing path IS reachable over real HTTP and is asserted here. # ═══════════════════════════════════════════════════════════════════════════ def test_tool_calls_without_tool_manager_raises_value_error_over_http(mock_multi_turn_server): - """A tool call with no tool_manager raises ValueError over the real HTTP path. - - Validates: Requirements 3.10, 7.1, 7.2 - """ + """A tool call with no tool_manager raises ValueError over the real HTTP path.""" rollout = _make_rollout(MODEL_TOOL, tool_manager=None, max_turns=3) trajectories = _make_trajectories(1) 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 85c467ab..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' @@ -162,7 +162,7 @@ def test_sample_stream_rejects_bad_max_tokens() -> None: list(s.sample_stream(inp, SamplingParams(max_tokens=0))) -# ---------- Multi-turn contract knobs (task 8.6) -------------------------- # +# ---------- Multi-turn contract knobs -------------------------- # def test_default_new_input_feature_appends_sampled_tokens() -> None: diff --git a/tests/server/state/test_leader_election.py b/tests/server/state/test_leader_election.py index 57891e1d..29d6122d 100644 --- a/tests/server/state/test_leader_election.py +++ b/tests/server/state/test_leader_election.py @@ -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/test_embedding_e2e.py b/tests/server/test_embedding_e2e.py index 68856427..240e36fd 100644 --- a/tests/server/test_embedding_e2e.py +++ b/tests/server/test_embedding_e2e.py @@ -15,13 +15,13 @@ * ``mock_embedding_model`` — local CPU-only entrypoint backed by ``TwinkleCompatMockModel`` (src/twinkle/server/model/backends/mock_model.py). Requires NO GPU. Boots an in-process Ray Serve cluster from the CPU-only - mock server config fixture. Consumed by task 4.2 (protocol/link validation). + mock server config fixture. Consumed by the local CPU protocol/link validation. * ``gpu_embedding_model`` — real transformers model entrypoint. Gated behind the ``TWINKLE_TEST_GPU_E2E=1`` environment variable (skipped otherwise) and connects to an already-running GPU server (see - tests/server/start_e2e_server.py). Consumed by tasks 4.3-4.5 (real numeric - validation). + tests/server/start_e2e_server.py). Consumed by the real numeric + validation cases. Plus two reusable helpers: @@ -29,7 +29,7 @@ anchor/positive contrastive-learning dataset (even number of samples, with ``labels = [1, 0, 1, 0, ...]`` anchor/positive semantics) in the embedding pooling input format consumed by ``InputProcessor``. - * ``run_embedding_training(...)`` — issues the design-document "verification + * ``run_embedding_training(...)`` — issues the "verification call sequence" against a client model and returns the observed losses and the final metric result. @@ -40,12 +40,12 @@ conda run -n twinkle pytest tests/server/test_embedding_e2e.py -v -The local, CPU-only mock cases (task 4.2) run without a GPU and must pass in +The local, CPU-only mock cases run without a GPU and must pass in that environment: conda run -n twinkle pytest tests/server/test_embedding_e2e.py -v -k mock -The GPU-gated cases (tasks 4.3-4.5) additionally require ``TWINKLE_TEST_GPU_E2E=1`` +The GPU-gated cases additionally require ``TWINKLE_TEST_GPU_E2E=1`` and a running GPU server; they are skipped automatically when the variable is unset: @@ -200,7 +200,7 @@ def configure_embedding_adapter( ) -> None: """Apply the embedding-training configuration steps to a client model. - Order matches the design document verification call sequence: + Order matches the verification call sequence: set_processor('InputProcessor') -> set_loss('InfonceLoss', temperature=..., use_batch=True) -> add_metric('EmbeddingMetric', is_training=True) @@ -435,7 +435,7 @@ def mock_embedding_model(): Boots the mock server in-process and yields a client model configured with a LoRA adapter, ready for the embedding verification call sequence. Consumed by - task 4.2. + the local CPU protocol/link validation cases. """ harness = MockEmbeddingServerHarness() base_url = harness.start() @@ -526,7 +526,7 @@ def test_gpu_fixture_gating_env_flag(): # ═══════════════════════════════════════════════════════════════════════════ -# Task 4.2 — Local CPU protocol/link validation against the mock backend. +# Local CPU protocol/link validation against the mock backend. # # These cases boot the CPU-only mock server (module-scoped ``mock_embedding_model`` # fixture — booted ONCE, never rebooted per example) and drive the full embedding @@ -540,10 +540,10 @@ def test_gpu_fixture_gating_env_flag(): # # They assert each HTTP request/response hop is well-formed, that ``task='embedding'`` # is transmitted through the protocol layer to the mock backend without protocol -# errors, and — Property 1 (Validates: Requirements 1.1) — that the loss returned by +# errors, and that the loss returned by # ``forward_backward(task='embedding')`` is a finite number (math.isfinite: not NaN, # not Inf). Here the mock loss layer validates the numeric-finiteness contract at the -# protocol-link level; real-model numeric semantics live in the GPU-gated tasks 4.3-4.5. +# protocol-link level; real-model numeric semantics live in the GPU-gated cases. # # Run locally (no GPU) via: # conda run -n twinkle pytest tests/server/test_embedding_e2e.py -v -k mock @@ -556,13 +556,11 @@ def test_gpu_fixture_gating_env_flag(): def test_mock_embedding_protocol_link_full_sequence(mock_embedding_model): """Full ordered call sequence runs over HTTP against the mock backend. - Validates that every hop of the design-document verification call sequence + Validates that every hop of the verification call sequence (set_processor -> set_loss -> add_metric -> forward_backward(task='embedding') -> calculate_metric) completes without any protocol-layer exception, that ``task='embedding'`` is accepted through the /twinkle/* protocol and reaches - the mock backend, and — Property 1 — that every returned loss is finite. - - Validates: Requirements 1.1, 1.7, 7.1, 7.2 + the mock backend, and that every returned loss is finite. """ model = mock_embedding_model @@ -576,7 +574,7 @@ def test_mock_embedding_protocol_link_full_sequence(mock_embedding_model): assert len(losses) == len(minibatches) assert losses, 'expected at least one forward_backward loss' - # Property 1: every embedding forward_backward loss is a finite number. + # Every embedding forward_backward loss is a finite number. for step, loss in enumerate(losses): assert math.isfinite(loss), f'loss at step {step} is not finite: {loss!r}' @@ -594,8 +592,6 @@ def test_mock_embedding_task_embedding_is_passed_through(mock_embedding_model): response whose extracted loss is finite. This confirms the extra kwarg is transported (via ``ForwardRequest.model_extra``) rather than rejected by the /twinkle/forward_backward endpoint. - - Validates: Requirements 1.1, 1.7, 7.1, 7.2 """ model = mock_embedding_model configure_embedding_adapter(model) @@ -625,15 +621,13 @@ def test_mock_embedding_task_embedding_is_passed_through(mock_embedding_model): seed=st.integers(min_value=0, max_value=10_000), ) def test_mock_embedding_loss_is_finite_property(mock_embedding_model, num_pairs, seq_len, seed): - """Property 1: forward_backward(task='embedding') loss is always finite. + """forward_backward(task='embedding') loss is always finite. Varied contrastive dataset shapes (num_pairs, seq_len) and seeds are generated within a SINGLE module-scoped mock server instance — the server is never rebooted per example. For every generated minibatch, the loss returned by the mock backend over the HTTP protocol link must be a finite number (math.isfinite: not NaN, not Inf). - - Validates: Requirements 1.1 """ model = mock_embedding_model # Idempotent no-op configuration on the mock backend; keeps the case self-contained. @@ -653,16 +647,16 @@ def test_mock_embedding_loss_is_finite_property(mock_embedding_model, num_pairs, # ═══════════════════════════════════════════════════════════════════════════ -# Task 4.3 — Single-GPU real-model loss finiteness (Property 1, GPU-gated). +# Single-GPU real-model loss finiteness (GPU-gated). # # This case exercises the REAL transformers model backend through the Twinkle # client HTTP path (module-scoped ``gpu_embedding_model`` fixture, gated behind # TWINKLE_TEST_GPU_E2E=1 and an already-running GPU server; skipped otherwise). # For a non-empty contrastive minibatch, the real-model # ``forward_backward(task='embedding')`` loss must be a finite number -# (math.isfinite: not NaN, not Inf). Unlike the mock cases in task 4.2 (which +# (math.isfinite: not NaN, not Inf). Unlike the mock cases (which # validate the numeric-finiteness contract only at the protocol-link level), this -# asserts Property 1 against real InfoNCE numeric semantics on real model weights. +# asserts finiteness against real InfoNCE numeric semantics on real model weights. # # Run on GPU via: # TWINKLE_TEST_GPU_E2E=1 conda run -n twinkle pytest \ @@ -671,9 +665,9 @@ def test_mock_embedding_loss_is_finite_property(mock_embedding_model, num_pairs, def test_gpu_embedding_real_model_loss_is_finite(gpu_embedding_model): - """Property 1: real-model forward_backward(task='embedding') loss is finite. + """Real-model forward_backward(task='embedding') loss is finite. - Drives the design-document embedding verification call sequence + Drives the embedding verification call sequence (set_processor -> set_loss('InfonceLoss', ...) -> add_metric('EmbeddingMetric') -> forward_backward(task='embedding') + clip_grad_and_step) against the REAL transformers backend and asserts that every loss returned by @@ -683,15 +677,13 @@ def test_gpu_embedding_real_model_loss_is_finite(gpu_embedding_model): GPU-gated (TWINKLE_TEST_GPU_E2E=1 + running GPU server): skipped automatically on machines without a GPU, so this asserts real numeric semantics only on GPU CI while collecting/skipping cleanly locally. - - Validates: Requirements 1.1 """ model = gpu_embedding_model dataset = build_synthetic_contrastive_dataset(DEFAULT_NUM_PAIRS, seq_len=DEFAULT_SEQ_LEN, seed=0) minibatches = iter_minibatches(dataset, batch_size=2 * DEFAULT_NUM_PAIRS) - # Property 1 requires a non-empty minibatch to actually exercise the loss path. + # A non-empty minibatch is required to actually exercise the loss path. assert minibatches, 'expected at least one minibatch' assert all(mb for mb in minibatches), 'expected every minibatch to be non-empty' @@ -706,18 +698,18 @@ def test_gpu_embedding_real_model_loss_is_finite(gpu_embedding_model): assert len(losses) == len(minibatches) assert losses, 'expected at least one forward_backward loss' - # Property 1: every real-model embedding forward_backward loss is finite. + # Every real-model embedding forward_backward loss is finite. for step, loss in enumerate(losses): assert math.isfinite(loss), f'real-model loss at step {step} is not finite: {loss!r}' # ═══════════════════════════════════════════════════════════════════════════ -# Task 4.4 — Single-GPU: TransformersEmbeddingPatch auto-rollback (Property 2). +# Single-GPU: TransformersEmbeddingPatch auto-rollback. # # GPU-gated (TWINKLE_TEST_GPU_E2E=1). Skipped automatically on this GPU-less dev # machine, and on any environment where the variable is unset. # -# Property 2 (Validates: Requirements 1.2): TransformersEmbeddingPatch is applied +# TransformersEmbeddingPatch is applied # *and rolled back automatically* inside a single forward call via # ``_resolve_task_context(model, task='embedding')`` (src/twinkle/model/transformers/ # transformers.py). While the patch is active, ``lm_head`` is swapped for identity @@ -770,7 +762,7 @@ def _build_causal_sample(seq_len: int = DEFAULT_SEQ_LEN) -> Dict[str, Any]: def test_gpu_embedding_patch_auto_rollback_property(gpu_embedding_model): - """Property 2: TransformersEmbeddingPatch auto-rolls back after an embedding step. + """TransformersEmbeddingPatch auto-rolls back after an embedding step. Flow (single real GPU adapter, via the ``gpu_embedding_model`` fixture): @@ -793,8 +785,6 @@ def test_gpu_embedding_patch_auto_rollback_property(gpu_embedding_model): * The post-embedding logits' trailing dimension equals the baseline vocabulary dimension, proving ``lm_head``/the forward hook were restored (not left as identity emitting hidden states). - - Validates: Requirements 1.2, 7.2 """ model = gpu_embedding_model configure_embedding_adapter(model) @@ -832,7 +822,7 @@ def test_gpu_embedding_patch_auto_rollback_property(gpu_embedding_model): 'did NOT roll back (feature hook suppresses the logits key)') post_dim = _innermost_dim(post_logits) - # Property 2: trailing dim is the vocab dim (identical to baseline), NOT the + # Trailing dim is the vocab dim (identical to baseline), NOT the # leftover identity hidden-state dim. assert post_dim == vocab_dim, ( f'TransformersEmbeddingPatch did NOT auto-rollback: post-embedding logits ' @@ -841,7 +831,7 @@ def test_gpu_embedding_patch_auto_rollback_property(gpu_embedding_model): # ═══════════════════════════════════════════════════════════════════════════ -# Task 4.5 — Single-GPU: bare-library parity + pos_sim upward trend (GPU-gated). +# Single-GPU: bare-library parity + pos_sim upward trend (GPU-gated). # # These cases exercise the REAL transformers model backend and assert real # numeric semantics that the mock backend cannot express, so they live behind @@ -849,15 +839,15 @@ def test_gpu_embedding_patch_auto_rollback_property(gpu_embedding_model): # TWINKLE_TEST_GPU_E2E=1 and a running GPU server; skipped automatically # otherwise). Two properties are validated: # -# * Requirement 1.5 — Twinkle client HTTP path vs. bare-library path parity: +# * Twinkle client HTTP path vs. bare-library path parity: # the same small synthetic contrastive dataset is trained for the same number -# of steps through (a) the Twinkle client HTTP path (task 4.1 helpers) and +# of steps through (a) the Twinkle client HTTP path and # (b) the bare-library training path used by # ``cookbook/exp/embedding/train_embedding_full_ddp.py``. The two losses must # stay within the SAME ORDER OF MAGNITUDE (digit-for-digit equality is NOT # required). # -# * Requirement 1.6 — pos_sim upward trend: after several real training steps +# * pos_sim upward trend: after several real training steps # on a dataset with a clear contrastive signal, ``calculate_metric(is_training= # True)`` must report an increasing anchor-positive cosine similarity # (``pos_sim``) relative to the untrained baseline. @@ -925,7 +915,7 @@ def _assert_same_order_of_magnitude(a: float, b: float, *, label: str) -> None: "Same order of magnitude" is interpreted as a bounded ratio in [0.1, 10] (equivalently, their base-10 exponents differ by at most 1). Digit-for-digit - equality is intentionally NOT required (Requirement 1.5). + equality is intentionally NOT required. """ assert math.isfinite(a) and math.isfinite(b), f'[{label}] non-finite losses: {a!r}, {b!r}' assert a > 0 and b > 0, f'[{label}] expected positive InfoNCE losses, got {a!r}, {b!r}' @@ -1026,12 +1016,12 @@ def run_bare_library_embedding_training( def test_gpu_embedding_loss_matches_bare_library_order_of_magnitude(gpu_embedding_model): - """Requirement 1.5: client HTTP path and bare-library path losses match in magnitude. + """Client HTTP path and bare-library path losses match in magnitude. Trains the SAME small synthetic contrastive dataset for the SAME number of steps through two independent paths: - * the Twinkle client HTTP path (task 4.1 helpers), against the real GPU + * the Twinkle client HTTP path, against the real GPU server, and * the bare-library training path reproduced in-process from ``cookbook/exp/embedding/train_embedding_full_ddp.py``. @@ -1046,8 +1036,6 @@ def test_gpu_embedding_loss_matches_bare_library_order_of_magnitude(gpu_embeddin GPU-gated (TWINKLE_TEST_GPU_E2E=1 + running GPU server): skipped automatically on machines without a GPU, so it collects/skips cleanly locally and asserts real numeric semantics only on GPU CI. - - Validates: Requirements 1.5, 7.2 """ # ``gpu_embedding_model`` guarantees the GPU gate + a live server/session; use a # dedicated fresh client adapter so the comparison is order-independent of other @@ -1091,7 +1079,7 @@ def test_gpu_embedding_loss_matches_bare_library_order_of_magnitude(gpu_embeddin def test_gpu_embedding_pos_sim_increases_after_training(gpu_embedding_model): - """Requirement 1.6: pos_sim rises after several real training steps. + """pos_sim rises after several real training steps. Uses a fresh, untrained LoRA adapter on the real GPU model and repeatedly trains on a fixed synthetic contrastive minibatch (anchors and their positives @@ -1106,8 +1094,6 @@ def test_gpu_embedding_pos_sim_increases_after_training(gpu_embedding_model): GPU-gated (TWINKLE_TEST_GPU_E2E=1 + running GPU server): skipped automatically on machines without a GPU. - - Validates: Requirements 1.6, 7.2 """ # Fresh dedicated adapter -> clean untrained baseline, independent of other cases. model = build_embedding_client_model(MODEL_ID, adapter_name='emb_trend_adapter') diff --git a/tests/server/test_embedding_e2e_sp_cp.py b/tests/server/test_embedding_e2e_sp_cp.py index 9b676d12..06c6a316 100644 --- a/tests/server/test_embedding_e2e_sp_cp.py +++ b/tests/server/test_embedding_e2e_sp_cp.py @@ -1,5 +1,5 @@ # Copyright (c) ModelScope Contributors. All rights reserved. -"""Embedding training E2E validation under SP/CP distributed configs (task 5.1). +"""Embedding training E2E validation under SP/CP distributed configs. This module reuses the synthetic contrastive dataset and the embedding verification call sequence from ``tests/server/test_embedding_e2e.py`` and reruns @@ -53,7 +53,7 @@ import pytest -# Reuse the task 4.1-4.5 scaffolding (dataset, minibatching, model id, metric +# Reuse the single-card scaffolding (dataset, minibatching, model id, metric # parsing, loss extraction, GPU gate) so the SP/CP flow exercises the SAME # synthetic dataset and the SAME verification call sequence as the single-card # path — the only difference is the distributed device mesh. @@ -74,7 +74,7 @@ # ═══════════════════════════════════════════════════════════════════════════ # Dataset / training shape (kept small for GPU speed while still giving InfoNCE a -# real contrastive signal). Mirrors the pos_sim-trend shape used in task 4.5. +# real contrastive signal). Mirrors the pos_sim-trend shape used on the single card. SP_CP_NUM_PAIRS = 6 # -> 12 samples SP_CP_BATCH_SIZE = 4 # -> 3 minibatches == 3 training steps SP_CP_TRAIN_STEPS = 8 # real training steps before measuring pos_sim @@ -179,7 +179,7 @@ def _build_embedding_model_with_mesh( ): """Build a real ``MultiLoraTransformersModel`` bound to an SP/CP device mesh. - Uses the SAME bare-library primitives as the single-card path in task 4.5 + Uses the SAME bare-library primitives as the single-card path (InputProcessor + InfonceLoss + EmbeddingMetric + AdamW), the only difference being the distributed ``DeviceMesh`` passed to the constructor. SP is enabled automatically by the transformers backend whenever ``ulysses_size > 1``. @@ -351,7 +351,7 @@ def single_card_pos_sim(): # ═══════════════════════════════════════════════════════════════════════════ -# Task 5.1 — SP/CP pos_sim consistency (multi-GPU, GPU-gated). +# SP/CP pos_sim consistency (multi-GPU, GPU-gated). # ═══════════════════════════════════════════════════════════════════════════ @@ -359,7 +359,7 @@ def single_card_pos_sim(): def test_gpu_sp_cp_pos_sim_matches_single_card(config, single_card_pos_sim): """SP/CP pos_sim agrees with the single-card baseline within tolerance. - Reruns the task 4.1-4.5 embedding verification flow on the SAME synthetic + Reruns the embedding verification flow on the SAME synthetic contrastive dataset under a distributed device mesh built with ``DeviceMesh.from_sizes(fsdp_size=world_size, ulysses_size=..., cp_size=...)`` (SP when ``ulysses_size > 1``, CP when ``cp_size > 1``), then asserts the @@ -369,8 +369,6 @@ def test_gpu_sp_cp_pos_sim_matches_single_card(config, single_card_pos_sim): its required GPU count is available, so the file collects/skips cleanly on a GPU-less machine and asserts real distributed numeric semantics only on multi-GPU CI. - - Validates: Requirements 1.3, 7.2 """ world_size = int(config['world_size']) if not multi_gpu_e2e_enabled(min_devices=world_size): 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_client/test_client_multi_turn_rollout.py b/tests/twinkle_client/test_client_multi_turn_rollout.py index 91dd072b..e3598c3a 100644 --- a/tests/twinkle_client/test_client_multi_turn_rollout.py +++ b/tests/twinkle_client/test_client_multi_turn_rollout.py @@ -10,12 +10,12 @@ 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 (see design doc "Correctness Properties"): - * Property 3 - output length & order preservation (Validates: Requirements 3.2) - * Property 4 - stop_reason value range (Validates: Requirements 3.3) - * Property 5 - logprobs / trainable-label alignment (Validates: Requirements 3.4) - * Property 6 - actual turns never exceed max_turns (Validates: Requirements 3.5) - * Property 7 - forced truncation at the max_turns edge (Validates: Requirements 3.6) +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 @@ -328,14 +328,12 @@ def _batch_specs(min_size: int = 1, max_size: int = 4): # ============================================================================= -# Property 3: multi-turn output length & order preservation (Req 3.2) +# 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_property3_output_length_and_order_preserved(scripts_spec, max_turns): - """**Validates: Requirements 3.2** - - ``__call__`` returns a list of the same length as the input, in the exact +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( @@ -349,14 +347,12 @@ def test_property3_output_length_and_order_preserved(scripts_spec, max_turns): # ============================================================================= -# Property 4: stop_reason value range (Req 3.3) +# 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_property4_stop_reason_value_range(scripts_spec, max_turns): - """**Validates: Requirements 3.3** - - Every returned trajectory's ``stop_reason`` is one of +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( @@ -369,14 +365,12 @@ def test_property4_stop_reason_value_range(scripts_spec, max_turns): # ============================================================================= -# Property 5: logprobs / trainable-label alignment (Req 3.4) +# 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_property5_logprobs_align_with_trainable_labels(scripts_spec, max_turns): - """**Validates: Requirements 3.4** - - For every returned trajectory with a non-empty ``logprobs`` list, its length +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( @@ -393,14 +387,12 @@ def test_property5_logprobs_align_with_trainable_labels(scripts_spec, max_turns) # ============================================================================= -# Property 6: actual turns never exceed max_turns (Req 3.5) +# 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_property6_turns_do_not_exceed_max_turns(scripts_spec, max_turns): - """**Validates: Requirements 3.5** - - Every returned trajectory's ``turns`` count is <= the configured +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( @@ -413,14 +405,12 @@ def test_property6_turns_do_not_exceed_max_turns(scripts_spec, max_turns): # ============================================================================= -# Property 7: forced truncation at the max_turns edge (Req 3.6) +# 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_property7_max_turns_one_forces_truncation(logprobs_flags): - """**Validates: Requirements 3.6** - - With ``max_turns == 1`` and a first-round tool_call, every trajectory is +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. @@ -442,7 +432,7 @@ def test_property7_max_turns_one_forces_truncation(logprobs_flags): # 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`` and requirements 3.7-3.10, 4.3-4.5, 7.1. +# ``twinkle_client.rollout.multi_turn``. # All are CPU-only and reuse the Fake infra above. # ============================================================================= class NetworkError(Exception): From 2d8e4732d3a141f5d65e040204f413d7c3fd9f79 Mon Sep 17 00:00:00 2001 From: Yunnglin Date: Mon, 20 Jul 2026 17:45:02 +0800 Subject: [PATCH 05/10] =?UTF-8?q?docs:=20=E8=A1=A5=E5=85=85=20ClientMultiT?= =?UTF-8?q?urnRollout=20=E4=B8=8E=20client=20=E4=BE=A7=20embedding=20?= =?UTF-8?q?=E6=8C=87=E5=8D=97=EF=BC=8C=E5=B9=B6=E7=B2=BE=E7=AE=80=E6=B5=8B?= =?UTF-8?q?=E8=AF=95?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - docs(zh/en): Twinkle 客户端文档新增「可训练多轮 Rollout（ClientMultiTurnRollout）」章节；Embedding 文档新增「通过 Twinkle Client 训练」章节 - docs(zh): Tinker 兼容客户端文档移除多轮 rollout 与 embedding 内容（Tinker 协议不支持），精简实现细节叙述，交叉引用改为 markdown 链接 - test: 移除冗余的 GPU/e2e 及重复的多轮测试，保留 CPU 单元测试 (test_bridge / test_client_multi_turn_rollout)；删除随之孤立的 e2e 服务配置 - chore: 清理注释/docstring 中残留的 spec 工作流编号 --- .../Usage Guide/Embedding-Training.md | 56 + .../Server and Client/Twinkle-Client.md | 94 ++ .../Embedding\350\256\255\347\273\203.md" | 16 +- ...71\345\256\242\346\210\267\347\253\257.md" | 65 - ...le\345\256\242\346\210\267\347\253\257.md" | 94 ++ src/twinkle_client/rollout/multi_turn.py | 4 +- ...ver_config_4b_e2e_megatron_no_sampler.yaml | 97 -- ...test_api_multi_turn_rollout_no_logprobs.py | 257 ---- .../test_client_multi_turn_rollout_e2e.py | 595 --------- ...test_client_multi_turn_rollout_mock_e2e.py | 563 -------- tests/server/state/test_leader_election.py | 2 +- tests/server/test_embedding_e2e.py | 1133 ----------------- tests/server/test_embedding_e2e_sp_cp.py | 449 ------- .../test_client_multi_turn_rollout.py | 25 +- 14 files changed, 255 insertions(+), 3195 deletions(-) delete mode 100644 tests/server/config/server_config_4b_e2e_megatron_no_sampler.yaml delete mode 100644 tests/server/gateway/test_api_multi_turn_rollout_no_logprobs.py delete mode 100644 tests/server/integration/test_client_multi_turn_rollout_e2e.py delete mode 100644 tests/server/integration/test_client_multi_turn_rollout_mock_e2e.py delete mode 100644 tests/server/test_embedding_e2e.py delete mode 100644 tests/server/test_embedding_e2e_sp_cp.py 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/Server and Client/Twinkle-Client.md b/docs/source_en/Usage Guide/Server and Client/Twinkle-Client.md index 5a118a61..2e63a297 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 @@ -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/self_host/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 f5df4e9f..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" @@ -108,7 +108,7 @@ for step, batch in enumerate(dataloader): ## 通过 Twinkle Client 训练 -除了直接使用裸库 `TransformersModel`，你也可以通过 `twinkle_client` 的 `MultiLoraTransformersModel` 以 HTTP 方式驱动服务端完成 Embedding 训练。协议层（`/twinkle/*`）已经具备完成 Embedding 训练所需的全部能力，**无需引入任何新的高层封装函数**（例如 `setup_embedding_training`），只需按正确顺序调用现有方法即可。 +除了直接使用裸库 `TransformersModel`，你也可以通过 `twinkle_client` 的 `MultiLoraTransformersModel` 以 HTTP 方式驱动服务端完成 Embedding 训练。用法与裸库一致，只需按正确顺序调用现有方法。 ### 完整调用顺序 @@ -156,19 +156,9 @@ for step, mb in enumerate(minibatches): metric = model.calculate_metric(is_training=True) ``` -### 关于 `TransformersEmbeddingPatch` 的自动应用 +**无需手动 apply_patch。** 只要在 `forward_backward`（或 `forward_only`）中传入 `task='embedding'`，服务端就会在这一次 `forward` 期间自动切换到 embedding 模式、并在结束后自动回滚——你**不需要**显式调用 `apply_patch(...)`。因此同一个 adapter 在 `forward_backward(task='embedding')` 之后，再调用不带 `task` 的 `forward_only` 仍会返回正常的词表维度 `logits`，互不影响。 -Embedding 训练依赖 `TransformersEmbeddingPatch` 把 causal LM 的 `lm_head` 替换为 identity 输出，从而得到 per-token hidden states 用于池化。**你不需要显式调用 `apply_patch(TransformersEmbeddingPatch())`**：当你在 `forward_backward`（或 `forward_only`）中传入 `task='embedding'` 时，服务端的 `_resolve_task_context` 会在该次 `forward` 调用期间自动应用该 patch，并在调用结束后自动回滚。 - -这意味着： - -- 同一个 adapter 在一次 `forward_backward(task='embedding')` 之后，紧接着调用不带 `task` 参数的 `forward_only` 仍会返回正常的词表维度 `logits`，不会残留上一次 embedding 任务的 identity hidden states。 -- 你在 client 侧真正需要显式调用的前置步骤只有三个：`set_processor('InputProcessor')`、`set_loss('InfonceLoss', ...)` 与 `add_metric('EmbeddingMetric', is_training=True)`。 - -### 说明 - -- 本节不引入任何新的高层封装函数，仅说明现有 `MultiLoraTransformersModel` 方法的正确调用顺序。 -- 各参数（`temperature`、`use_batch`、`hard_negatives` 等）的含义与取值建议参见上文「关键参数」；`calculate_metric` 返回的指标含义参见下文「监控指标」。 +client 侧真正需要显式配置的只有三步：`set_processor('InputProcessor')`、`set_loss('InfonceLoss', ...)`、`add_metric('EmbeddingMetric', is_training=True)`；各参数含义与取值建议见上文「关键参数」，返回指标见下文「监控指标」。 --- 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/Tinker\345\205\274\345\256\271\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/Tinker\345\205\274\345\256\271\345\256\242\346\210\267\347\253\257.md" index 2bc16c17..a1f7e064 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/Tinker\345\205\274\345\256\271\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/Tinker\345\205\274\345\256\271\345\256\242\346\210\267\347\253\257.md" @@ -195,68 +195,3 @@ rest_client = service_client.create_rest_client() rest_client.publish_checkpoint_from_tinker_path(save_result.path).result() print("Published checkpoint to ModelScope Hub") ``` - -## 多轮工具调用（间接方案） - -Tinker 兼容客户端本身不提供多轮 agentic rollout 编排能力。如果你在用 Tinker Client 训练/管理 checkpoint，同时又需要多轮工具调用能力（**仅用于生成或评测**），可以在完全不改动 Tinker SDK 代码的前提下，额外使用一个标准的 OpenAI 兼容客户端指向 Gateway 的 `/chat/completions` 端点，并复用 `twinkle_agentic.rollout.APIMultiTurnRollout` 与 `ToolManager`。 - -这是一条纯增量的使用路径：它复用的是 Twinkle Server 上一个独立于 Tinker 协议命名空间的、已经存在并跑通的 OpenAI 兼容端点，因此不涉及对 Tinker SDK 源码或 `patch_tinker.py` 中现有 monkeypatch 范围的任何修改。 - -```python -import os -from twinkle.data_format.sampling import SamplingParams -from twinkle_agentic.protocol.openai import OpenAI -from twinkle_agentic.rollout import APIMultiTurnRollout -from twinkle_agentic.tools.tool_manager import ToolManager -# from your_project.tools import CalculatorTool, SearchTool - -# An OpenAI-compatible client pointing at the Gateway route, -# NOT at tinker's base_url. This is a plain additive usage of an -# already-existing endpoint; the tinker SDK is untouched. -api = OpenAI( - model='Qwen/Qwen3.5-4B', - api_key=os.environ.get('MODELSCOPE_TOKEN', 'EMPTY_API_KEY'), - base_url='http://localhost:8000/api/v1', # Gateway /chat/completions route -) - -# Register your tools; ToolManager is reused directly from twinkle_agentic. -tool_manager = ToolManager([ - # CalculatorTool(), - # SearchTool(), -]) - -rollout = APIMultiTurnRollout( - api=api, - tool_manager=tool_manager, - sampling_params=SamplingParams(temperature=0.7, max_tokens=512), - max_turns=6, - concurrency=8, -) - -trajectories_in = [ - {'messages': [{'role': 'user', 'content': 'What is 12 * 34?'}]}, -] -trajectories_out = rollout(trajectories_in) - -# trajectories_out[i]['messages'] contains the full conversation, including -# assistant tool_calls and tool-response turns. -# trajectories_out[i]['stop_reason'] is one of 'stop' | 'length' | 'max_turns' | 'api_error'. -for traj in trajectories_out: - for message in traj['messages']: - print(message['role'], ':', message.get('content')) -``` - -### 能力边界：不可用于 RL 训练 - -该间接方案返回的 Trajectory **不包含** `logprobs` 字段（端到端验证已确认走 Gateway `/chat/completions` 的 `APIMultiTurnRollout` 产出的 trajectory 中不存在 `logprobs`，也没有 `input_ids` 等训练所需字段）。因此该路径**仅适合生成数据、离线评测等场景，不能直接喂给依赖 token 级对齐信息的 RL 训练（如 GRPO）**。 - -如果多轮 rollout 的产出需要直接用于 RL 训练（需要 token 级 `logprobs` 与 `labels` 对齐），请使用 Twinkle 客户端侧的 `ClientMultiTurnRollout`（通过 `/twinkle/sample` 返回 `new_input_feature`/`logprobs`），详见《Twinkle 客户端》文档。 - -### 为什么 Tinker 客户端不支持 embedding 训练与可训练多轮 rollout - -这两项能力无法通过 Tinker 兼容客户端提供，根本原因在于 Tinker 协议的数据类型层面缺少所需结构： - -- **Embedding 训练**：`tinker.types.Datum` 不具备对比学习所需的样本分组结构（anchor/positive 成对分组），无法承载 embedding 训练的输入语义。 -- **可训练多轮 rollout**：`tinker.types.SampledSequence` 不具备 `new_input_feature` 字段，无法在多轮之间传递并累积 token 级对齐信息，因此产出无法用于训练。 - -由于 Tinker 是第三方 SDK、不可修改其源码，即使想补齐这些字段也需要改动 SDK 本身（不允许）。因此上述能力仅在 Twinkle 客户端侧支持，Tinker 兼容客户端只提供上文所述的"仅生成/评测"间接多轮方案。 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..cad5cbb9 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" @@ -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/self_host/multi_turn_rollout.py`。 diff --git a/src/twinkle_client/rollout/multi_turn.py b/src/twinkle_client/rollout/multi_turn.py index d488736d..55c5800b 100644 --- a/src/twinkle_client/rollout/multi_turn.py +++ b/src/twinkle_client/rollout/multi_turn.py @@ -192,8 +192,8 @@ def __call__(self, trajectories: List[Trajectory], **kwargs) -> List[Trajectory] continue # 3c. Hit the turn cap while still wanting to call a tool: force - # truncation. Also covers the ``max_turns == 1`` edge (task - # 8.3), where the very first sampled turn trips this branch. + # 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' diff --git a/tests/server/config/server_config_4b_e2e_megatron_no_sampler.yaml b/tests/server/config/server_config_4b_e2e_megatron_no_sampler.yaml deleted file mode 100644 index a6eb93e7..00000000 --- a/tests/server/config/server_config_4b_e2e_megatron_no_sampler.yaml +++ /dev/null @@ -1,97 +0,0 @@ -# Twinkle Server Configuration - E2E Test (4B, Megatron DP=2 PP=2) - NO SAMPLER -# -# Same as server_config_4b_e2e_megatron.yaml but WITHOUT the vLLM sampler -# application. Embedding training (and any forward/backward-only workflow) does -# not need a sampler, so dropping it avoids loading a ~40GB vLLM engine and cuts -# server startup time significantly. -# -# NOTE: with no sampler the extra 1-GPU Ray worker started by -# tests/server/start_e2e_server.py is simply unused (the model's PP=2 x DP=2 -# still fits on the head node's 4 GPUs). - -proxy_location: EveryNode - -http_options: - host: 0.0.0.0 - port: 9000 - -applications: - - - name: server - route_prefix: /api/v1 - import_path: server - args: - server_config: - per_token_model_limit: 30 - supported_models: - - Qwen/Qwen3.5-4B - deployments: - - name: TinkerCompatServer - max_ongoing_requests: 50 - autoscaling_config: - min_replicas: 1 - max_replicas: 1 - target_ongoing_requests: 128 - ray_actor_options: - num_cpus: 0.1 - runtime_env: - env_vars: - TWINKLE_FAIL_FAST: "0" - - - name: models-Qwen3.5-4B - route_prefix: /api/v1/model/Qwen/Qwen3.5-4B - import_path: model - args: - backend: megatron - model_id: "ms://Qwen/Qwen3.5-4B" - max_length: 10240 - nproc_per_node: 4 - device_group: - name: model - ranks: 4 - device_type: cuda - device_mesh: - device_type: cuda - dp_size: 2 - pp_size: 2 - queue_config: - rps_limit: 100 - tps_limit: 100000 - adapter_config: - adapter_timeout: 30 - deployments: - - name: ModelManagement - autoscaling_config: - min_replicas: 1 - max_replicas: 1 - target_ongoing_requests: 16 - ray_actor_options: - num_cpus: 0.1 - runtime_env: - env_vars: - TWINKLE_TRUST_REMOTE_CODE: "1" - TWINKLE_FAIL_FAST: "0" - - - name: processor - route_prefix: /api/v1/processor - import_path: processor - args: - ncpu_proc_per_node: 2 - device_group: - name: model - ranks: 2 - device_type: CPU - device_mesh: - device_type: CPU - dp_size: 2 - deployments: - - name: ProcessorManagement - autoscaling_config: - min_replicas: 1 - max_replicas: 1 - target_ongoing_requests: 128 - ray_actor_options: - num_cpus: 0.1 - runtime_env: - env_vars: - TWINKLE_FAIL_FAST: "0" diff --git a/tests/server/gateway/test_api_multi_turn_rollout_no_logprobs.py b/tests/server/gateway/test_api_multi_turn_rollout_no_logprobs.py deleted file mode 100644 index 789c4045..00000000 --- a/tests/server/gateway/test_api_multi_turn_rollout_no_logprobs.py +++ /dev/null @@ -1,257 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""GPU-gated E2E: ``APIMultiTurnRollout`` over the Gateway ``/chat/completions`` -endpoint proves the returned Trajectory carries NO ``logprobs`` field. - -Purpose -------- -This is the "generation-only, not trainable" counterpart to the trainable -``ClientMultiTurnRollout`` path. It drives the SAME -tool-calling scenario, but through the OpenAI-compatible route: - - OpenAI client -> Gateway POST /chat/completions -> /twinkle/sample - -and asserts that the resulting Trajectory does **not** contain a ``logprobs`` -field (or that it is ``None``/absent). The OpenAI chat-completions protocol only -surfaces assistant ``content`` / ``tool_calls`` / ``finish_reason`` — it never -carries the token-level ``logprobs`` + ``new_input_feature`` alignment info that -GRPO training requires. This test locks in the accuracy of that limitation -statement: the Gateway indirection is fine for -generation/evaluation but cannot feed token-aligned RL training. - -What is "real" here -------------------- - * REAL (GPU): the running Twinkle server (gateway + real vLLM sampler for - ``Qwen/Qwen3.5-4B``), the ``/chat/completions`` -> ``/twinkle/sample`` hop, - and the ``openai`` pip client issuing actual HTTP requests. - * ``APIMultiTurnRollout`` + a small ``ToolManager`` (calculator tool) drive - the multi-turn control flow client-side. - -============================================================================== -GPU + CONDA ENVIRONMENT REQUIREMENT (MANDATORY) -============================================================================== -This case requires a GPU and an already-running GPU server (see -tests/server/start_e2e_server.py). It is gated behind ``TWINKLE_TEST_GPU_E2E=1`` -and SKIPS cleanly on machines without a GPU (e.g. local dev). Run on GPU CI in -the ``twinkle`` conda env: - - TWINKLE_TEST_GPU_E2E=1 conda run -n twinkle pytest \ - tests/server/gateway/test_api_multi_turn_rollout_no_logprobs.py -v - -Locally (no GPU / variable unset) it collects and skips without error: - - conda run -n twinkle pytest \ - tests/server/gateway/test_api_multi_turn_rollout_no_logprobs.py -v -============================================================================== -""" -from __future__ import annotations - -import json -import os -import sys -from typing import Any, Dict - -# Ensure project root is importable for both pytest and direct execution. -_PROJECT_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', '..')) -if _PROJECT_ROOT not in sys.path: - sys.path.insert(0, _PROJECT_ROOT) - -import pytest - -# Reuse the shared GPU-e2e gate and server coordinates. -from tests.server.integration.e2e_helpers import ( - API_KEY, - BASE_MODEL, - BASE_URL, - log, - wait_for_server, -) -from tests.server.test_embedding_e2e import gpu_e2e_enabled - -# The gateway (OpenAI-compatible) routes are mounted under the ``server`` app's -# route prefix ``/api/v1`` (see tests/server/config/server_config_4b_e2e.yaml). -# The ``openai`` client appends ``/chat/completions`` to this base URL. -GATEWAY_BASE_URL = f'{BASE_URL}/api/v1' - -# Model/adapter name routed by the gateway. The e2e config serves the base model -# ``Qwen/Qwen3.5-4B`` directly. -GATEWAY_MODEL = BASE_MODEL - - -# ═══════════════════════════════════════════════════════════════════════════ -# Tool scenario (mirrors the 9.2 tool-calling scenario): a small calculator. -# ═══════════════════════════════════════════════════════════════════════════ -TOOL_NAME = 'calculator' - - -def _make_tool_manager(): - """Build a ToolManager with a single deterministic calculator tool.""" - from twinkle_agentic.tools.base import Tool - from twinkle_agentic.tools.tool_manager import ToolManager - - class CalculatorTool(Tool): - """Adds/subtracts/multiplies two integers; returns the result as text.""" - - def __call__(self, tool_name: str, arguments: Dict[str, Any]) -> str: - a = int(arguments.get('a', 0)) - b = int(arguments.get('b', 0)) - op = arguments.get('operation', 'add') - if op == 'add': - return str(a + b) - if op == 'subtract': - return str(a - b) - if op == 'multiply': - return str(a * b) - return f'Error: unknown operation {op!r}' - - def tool_info(self): - return { - 'type': 'function', - 'function': { - 'name': TOOL_NAME, - 'description': 'Compute a + b, a - b, or a * b for two integers.', - 'parameters': { - 'type': 'object', - 'properties': { - 'a': {'type': 'integer', 'description': 'first operand'}, - 'b': {'type': 'integer', 'description': 'second operand'}, - 'operation': { - 'type': 'string', - 'enum': ['add', 'subtract', 'multiply'], - 'description': 'the arithmetic operation', - }, - }, - 'required': ['a', 'b', 'operation'], - }, - }, - } - - mgr = ToolManager({}) - mgr.register(CalculatorTool()) - return mgr - - -def _make_trajectory() -> Dict[str, Any]: - """A single math trajectory that nudges the model toward the calculator tool.""" - return { - 'messages': [ - { - 'role': 'system', - 'content': ('You are a precise assistant. When a calculation is needed, ' - 'call the calculator tool instead of computing yourself.'), - }, - {'role': 'user', 'content': 'What is 123 multiplied by 7? Use the calculator tool.'}, - ], - } - - -def _build_rollout(max_turns: int = 3): - """Construct an APIMultiTurnRollout wired to the Gateway /chat/completions.""" - from twinkle.data_format.sampling import SamplingParams - from twinkle_agentic.protocol.openai import OpenAI - from twinkle_agentic.rollout import APIMultiTurnRollout - - api = OpenAI(model=GATEWAY_MODEL, api_key=API_KEY, base_url=GATEWAY_BASE_URL) - # temperature=0 keeps the control flow as reproducible as the server allows. - sampling_params = SamplingParams(num_samples=1, temperature=0.0, max_tokens=256) - return APIMultiTurnRollout( - api=api, - tool_manager=_make_tool_manager(), - sampling_params=sampling_params, - max_turns=max_turns, - ) - - -# ═══════════════════════════════════════════════════════════════════════════ -# Fixtures -# ═══════════════════════════════════════════════════════════════════════════ -@pytest.fixture(scope='module') -def gpu_gateway_ready(): - """Ensure a GPU server is up before running; skip cleanly when GPU e2e is off.""" - if not gpu_e2e_enabled(): - pytest.skip('Set TWINKLE_TEST_GPU_E2E=1 to run the Gateway APIMultiTurnRollout ' - 'no-logprobs E2E (requires a running GPU server)') - wait_for_server() - yield - - -# ═══════════════════════════════════════════════════════════════════════════ -# APIMultiTurnRollout via Gateway: Trajectory carries NO logprobs. -# ═══════════════════════════════════════════════════════════════════════════ -def test_api_multi_turn_rollout_trajectory_has_no_logprobs(gpu_gateway_ready): - """The Gateway /chat/completions path yields Trajectories without logprobs. - - Runs the same tool-calling scenario as the trainable client rollout but through - ``APIMultiTurnRollout`` + an OpenAI-compatible client pointed at the Gateway. - Asserts the returned Trajectory does NOT expose a ``logprobs`` field (absent - or ``None``), confirming the "generation-only, not trainable" limitation is - accurate: token-level alignment info never crosses the OpenAI protocol. - - GPU-gated (TWINKLE_TEST_GPU_E2E=1 + running GPU server); skipped locally. - """ - rollout = _build_rollout(max_turns=3) - trajectories = [_make_trajectory()] - - outs = rollout(trajectories) - - # Same length/order contract the multi-turn rollouts guarantee. - assert len(outs) == 1 - out = outs[0] - log(f'[api-rollout] stop_reason={out.get("stop_reason")} turns={out.get("turns")}') - - # Control flow terminated within the APIMultiTurnRollout stop-reason vocabulary. - assert out.get('stop_reason') in {'stop', 'length', 'max_turns', 'api_error'}, out.get('stop_reason') - assert isinstance(out.get('turns'), int) and out['turns'] >= 1 - - # A functional run should not have surfaced an API error. - assert out.get('stop_reason') != 'api_error', f"unexpected api_error: {out.get('error')!r}" - - # Core assertion: NO trainable token-level logprobs. - assert out.get('logprobs') is None, ( - f'Gateway APIMultiTurnRollout Trajectory unexpectedly carries logprobs: ' - f'{out.get("logprobs")!r}. The OpenAI /chat/completions path is ' - f'generation-only and must not expose token-level logprobs.') - - # The per-message conversation also must not smuggle token-level logprobs - # into any assistant turn (only content / tool_calls / finish_reason are set). - for msg in out.get('messages', []): - assert 'logprobs' not in msg, ( - f'assistant/tool message unexpectedly carries logprobs: {msg!r}') - - -def test_api_multi_turn_rollout_tool_call_scenario_shape(gpu_gateway_ready): - """The tool-calling scenario produces a well-formed multi-turn conversation. - - Complements the no-logprobs assertion by checking the conversation shape: - messages accumulate across turns and, when the model invokes the calculator, - a matching ``role='tool'`` message is stitched back in. This mirrors the - trainable client rollout scenario so the two paths are compared on the same - workload. Tool invocation itself is model-dependent, so it is asserted conditionally. - """ - rollout = _build_rollout(max_turns=3) - - outs = rollout([_make_trajectory()]) - out = outs[0] - - messages = out.get('messages', []) - # The full conversation includes at least the seed system+user turns plus one - # assistant reply. - assert len(messages) >= 3, messages - assert messages[0]['role'] == 'system' - assert messages[1]['role'] == 'user' - assert any(m['role'] == 'assistant' for m in messages) - - # If the model emitted a tool call, a paired tool result must follow it, and - # the calculator's deterministic answer (123 * 7 = 861) should appear. - assistant_with_tool = next( - (m for m in messages if m['role'] == 'assistant' and m.get('tool_calls')), None) - if assistant_with_tool is not None: - tool_msgs = [m for m in messages if m['role'] == 'tool'] - assert tool_msgs, 'assistant emitted tool_calls but no tool result was stitched back' - # Whichever calculator call was made, its numeric result is plain text - # (no logprobs), reinforcing the generation-only contract. - for tm in tool_msgs: - assert isinstance(tm.get('content'), str) - assert 'logprobs' not in tm - - # Regardless of tool invocation, the no-logprobs invariant still holds. - assert out.get('logprobs') is None diff --git a/tests/server/integration/test_client_multi_turn_rollout_e2e.py b/tests/server/integration/test_client_multi_turn_rollout_e2e.py deleted file mode 100644 index 9ce78589..00000000 --- a/tests/server/integration/test_client_multi_turn_rollout_e2e.py +++ /dev/null @@ -1,595 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""Behavioural-alignment E2E: ``ClientMultiTurnRollout`` vs bare ``MultiTurnRollout``. - -This module validates that the client-side, HTTP-driven -:class:`twinkle_client.rollout.multi_turn.ClientMultiTurnRollout` produces the -SAME multi-turn control flow as the bare-library, Ray-actor-driven -:class:`twinkle_agentic.rollout.multi_turn.MultiTurnRollout` when both are pointed -at the *same real sampler weights* (Qwen3.5-4B) with the same prompt, the same -tools, and greedy decoding (``temperature=0``). - -Why this must be GPU-gated ---------------------------- -Real numeric semantics (actual token sampling + per-token ``logprobs`` on real -model weights) cannot be reproduced by the CPU-only mock sampler used in the -local mock E2E. Two things in particular are only observable with a -real sampler and are therefore asserted here under GPU gating: - - * The two rollout paths agree on the ``messages`` structure and on *when* - tool calls are triggered (greedy decoding makes the control flow - deterministic across both transports, even though we allow token-level - non-determinism in principle). - * ``logprobs`` are actually populated, so the strict invariant - ``len(logprobs) == count(labels != -100)`` can be checked per trajectory on - both paths (the mock path cannot exercise real logprobs numerics). - -Topology on GPU CI ------------------- - * CLIENT path: connects to an ALREADY-RUNNING Twinkle e2e server (start it - first with ``tests/server/start_e2e_server.py``, which serves - ``tests/server/config/server_config_4b_e2e.yaml`` including the - ``sampler-Qwen3.5-4B`` application). Sampling goes over HTTP through - :class:`twinkle_client.sampler.vLLMSampler`. This mirrors the convention of - the GPU cases in ``tests/server/test_embedding_e2e.py``. - * BARE path: builds its OWN bare-library Ray-actor - :class:`twinkle.sampler.vLLMSampler` (``remote_group='sampler'`` + - ``DeviceMesh``) in-process, mirroring the standalone GPU sampler tests - (``tests/sampler/test_weight_sync.py``, ``tests/sampler/align_swift.py``) - and the multi-turn cookbook (``cookbook/rl/multi_turn/multi_turn_grpo.py``). - This is what the task means by "directly holding the corresponding Ray actor - sampler". The runner must therefore provision enough GPUs for BOTH the - server's sampler and this in-test bare sampler. - -============================================================================== -CONDA ENVIRONMENT REQUIREMENT (MANDATORY) — GPU REQUIRED -============================================================================== -This whole file is gated behind ``TWINKLE_TEST_GPU_E2E=1`` and is skipped -automatically on machines without a GPU (so it collects/skips cleanly during -local, CPU-only development). Run it inside the ``twinkle`` conda env on a GPU -host, with the e2e server already running: - - # 1) start the server (separate shell) - conda run -n twinkle python tests/server/start_e2e_server.py - - # 2) run the alignment test - TWINKLE_TEST_GPU_E2E=1 conda run -n twinkle pytest \ - tests/server/integration/test_client_multi_turn_rollout_e2e.py -v -============================================================================== -""" -from __future__ import annotations - -import copy -import json -import os -import sys -from typing import Any, Dict, List, Optional - -# Ensure project root is importable for both pytest and direct execution. -_PROJECT_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', '..')) -if _PROJECT_ROOT not in sys.path: - sys.path.insert(0, _PROJECT_ROOT) - -import pytest - -from twinkle.data_format.sampling import SamplingParams -from twinkle_agentic.tools.base import Tool -from twinkle_agentic.tools.tool_manager import ToolManager - -# Reuse the shared GPU e2e helpers (running-server URL, session init, model id). -from tests.server.integration.e2e_helpers import ( - MODEL_ID, - init_twinkle_client_session, - log, - wait_for_server, -) - - -# ═══════════════════════════════════════════════════════════════════════════ -# GPU gating -# -# Mirrors the ``gpu_e2e_enabled`` gate used by tests/server/test_embedding_e2e.py. -# Defined locally (not imported) so this file collects/skips cleanly even if the -# embedding e2e module is being edited concurrently. -# ═══════════════════════════════════════════════════════════════════════════ -def gpu_e2e_enabled() -> bool: - """Return True only when GPU e2e tests are explicitly enabled.""" - return os.environ.get('TWINKLE_TEST_GPU_E2E', '0') == '1' - - -# ═══════════════════════════════════════════════════════════════════════════ -# Configuration -# ═══════════════════════════════════════════════════════════════════════════ - -# Greedy decoding (temperature=0) makes both transports produce the SAME token -# stream for the SAME weights, so the multi-turn control flow is deterministic -# and directly comparable. ``logprobs=1`` forces per-token logprobs so the -# strict logprobs/trainable-label alignment can be asserted on both paths. -GREEDY_MAX_TOKENS = 128 -GREEDY_SEED = 0 - -# 3-6 rounds as required by the task. -ALIGN_MAX_TURNS = 6 - -# Template used locally by BOTH rollouts for encode / parse_tool_call / bridge. -TEMPLATE_CLS = 'Qwen3_5Template' - - -# ═══════════════════════════════════════════════════════════════════════════ -# A tiny ToolManager with 1-2 simple, deterministic tools (echo + calculator). -# ═══════════════════════════════════════════════════════════════════════════ -class EchoTool(Tool): - """Echoes its arguments back as a JSON string (deterministic).""" - - NAME = 'echo' - - def __call__(self, tool_name: str, arguments: Dict[str, Any]) -> str: - return f'echo[{tool_name}]:{json.dumps(arguments, sort_keys=True, ensure_ascii=False)}' - - def tool_info(self): - return { - 'type': 'function', - 'function': { - 'name': self.NAME, - 'description': 'Echo the given text back verbatim.', - 'parameters': { - 'type': 'object', - 'properties': {'text': {'type': 'string', 'description': 'Text to echo.'}}, - 'required': ['text'], - }, - }, - } - - -class CalculatorTool(Tool): - """Evaluates ``a b`` for the four basic operators (deterministic).""" - - NAME = 'calculator' - _OPS = {'add': lambda a, b: a + b, 'sub': lambda a, b: a - b, - 'mul': lambda a, b: a * b, 'div': lambda a, b: (a / b if b else float('inf'))} - - def __call__(self, tool_name: str, arguments: Dict[str, Any]) -> str: - op = str(arguments.get('op', 'add')) - a = float(arguments.get('a', 0)) - b = float(arguments.get('b', 0)) - fn = self._OPS.get(op, self._OPS['add']) - return json.dumps({'result': fn(a, b)}, ensure_ascii=False) - - def tool_info(self): - return { - 'type': 'function', - 'function': { - 'name': self.NAME, - 'description': 'Compute a basic arithmetic operation on two numbers.', - 'parameters': { - 'type': 'object', - 'properties': { - 'op': {'type': 'string', 'enum': ['add', 'sub', 'mul', 'div']}, - 'a': {'type': 'number'}, - 'b': {'type': 'number'}, - }, - 'required': ['op', 'a', 'b'], - }, - }, - } - - -def make_tool_manager() -> ToolManager: - """Build a small ToolManager holding the echo + calculator tools.""" - mgr = ToolManager({}) - mgr.register(EchoTool()) - mgr.register(CalculatorTool()) - return mgr - - -def tool_schema() -> List[Dict[str, Any]]: - """Tool schema advertised to the model in the trajectory ``tools`` field.""" - return [EchoTool().tool_info(), CalculatorTool().tool_info()] - - -SYSTEM_PROMPT = ( - 'You are a helpful assistant with access to tools. When a computation is ' - 'needed, call the `calculator` tool with the appropriate operator and ' - 'operands. When asked to repeat text, call the `echo` tool. Prefer using a ' - 'tool over answering directly, then give a short final answer.') - - -def build_alignment_trajectories() -> List[Dict[str, Any]]: - """Build a small, fixed batch of tool-inviting trajectories. - - Each trajectory carries a hidden ``_tid`` so output order can be checked, a - system prompt describing the tools, and the ``tools`` schema so the template - renders tool definitions into the prompt. - """ - users = [ - 'Please compute 21 plus 34 using your tools, then tell me the result.', - 'Use a tool to echo the exact text "ping", then confirm what you echoed.', - ] - trajectories: List[Dict[str, Any]] = [] - for tid, content in enumerate(users): - trajectories.append({ - 'messages': [ - {'role': 'system', 'content': SYSTEM_PROMPT}, - {'role': 'user', 'content': content}, - ], - 'tools': tool_schema(), - '_tid': tid, - }) - return trajectories - - -# ═══════════════════════════════════════════════════════════════════════════ -# Pure comparison helpers (collection-safe: unit-tested locally without a GPU). -# ═══════════════════════════════════════════════════════════════════════════ -def count_trainable_labels(labels: Optional[List[int]]) -> int: - """Number of trainable positions (``label != -100``) in a labels list.""" - return sum(1 for label in (labels or []) if label != -100) - - -def message_role_signature(messages: List[Dict[str, Any]]) -> List[str]: - """Ordered list of message roles — captures turn structure & tool timing. - - The presence and position of ``tool`` role messages encodes exactly when a - tool call was triggered and answered, which is the control-flow signal we - compare across the two transports. - """ - return [str(m.get('role')) for m in (messages or [])] - - -def tool_turn_indices(messages: List[Dict[str, Any]]) -> List[int]: - """Assistant-turn indices (0-based over assistant messages) that were - immediately followed by a tool response. - - This makes "when tool_calls fired" explicit and independent of textual - content: assistant turn ``k`` is a tool-calling turn iff the next message is - a ``tool`` message. - """ - roles = message_role_signature(messages) - indices: List[int] = [] - assistant_idx = -1 - for i, role in enumerate(roles): - if role == 'assistant': - assistant_idx += 1 - if i + 1 < len(roles) and roles[i + 1] == 'tool': - indices.append(assistant_idx) - return indices - - -def control_flow_fingerprint(traj: Dict[str, Any]) -> Dict[str, Any]: - """Extract the transport-independent control-flow fingerprint of an output. - - Deliberately EXCLUDES token ids / raw logprob values (which may differ) and - keeps only the control-flow-relevant fields the task requires to match: - role signature, tool-call timing, stop_reason, turns, truncated. - """ - return { - 'roles': message_role_signature(traj.get('messages') or []), - 'tool_turns': tool_turn_indices(traj.get('messages') or []), - 'stop_reason': traj.get('stop_reason'), - 'turns': traj.get('turns'), - 'truncated': bool(traj.get('truncated')), - } - - -def assert_logprobs_align_with_labels(traj: Dict[str, Any], label: str) -> None: - """Assert ``len(logprobs) == count(labels != -100)`` for a rollout output. - - Only meaningful when logprobs were requested/collected; a ``None``/empty - logprobs list means "not trainable" and is skipped (the rollout itself - already guards the non-empty case, this re-asserts it at the e2e boundary). - """ - logprobs = traj.get('logprobs') - if not logprobs: - return - trainable = count_trainable_labels(traj.get('labels')) - assert len(logprobs) == trainable, ( - f'[{label}] logprobs({len(logprobs)}) != trainable labels({trainable}) ' - f'(labels != -100)') - - -# Allowed stop reasons. -_ALLOWED_STOP_REASONS = {'length', 'stop', 'max_turns'} - - -def assert_multi_turn_invariants(outs: List[Dict[str, Any]], n_inputs: int, max_turns: int, label: str) -> None: - """Re-assert the multi-turn output invariants on a rollout output list. - - * Output length & order preserved (via hidden ``_tid``). - * stop_reason within ``{'length','stop','max_turns'}``. - * Actual turns never exceed max_turns. - * A ``max_turns`` truncation implies ``truncated=True``. - """ - # Same length and order. - assert len(outs) == n_inputs, f'[{label}] expected {n_inputs} outputs, got {len(outs)}' - for i, out in enumerate(outs): - assert out.get('_tid') == i, f'[{label}] output order mismatch at index {i}' - # stop_reason value range. - assert out.get('stop_reason') in _ALLOWED_STOP_REASONS, \ - f"[{label}] stop_reason={out.get('stop_reason')!r} not in {_ALLOWED_STOP_REASONS}" - # Turns bounded by max_turns. - assert out.get('turns') is not None and out['turns'] <= max_turns, \ - f"[{label}] turns({out.get('turns')}) > max_turns({max_turns})" - # max_turns truncation implies truncated=True. - if out.get('stop_reason') == 'max_turns': - assert out.get('truncated') is True, \ - f'[{label}] stop_reason==max_turns but truncated is not True' - - -# ═══════════════════════════════════════════════════════════════════════════ -# GPU-only sampler / rollout builders (never called at collection time). -# ═══════════════════════════════════════════════════════════════════════════ -def _build_local_template(): - """Build the real Qwen3.5 template used locally by BOTH rollouts.""" - from twinkle.template import Qwen3_5Template - template = Qwen3_5Template(MODEL_ID, max_length=8192, enable_thinking=False) - # Multi-turn bridge stitching does not support 'split'. - template.truncation_strategy = 'delete' - return template - - -def _greedy_sampling_params() -> SamplingParams: - """Greedy, deterministic sampling params shared by both paths.""" - return SamplingParams( - max_tokens=GREEDY_MAX_TOKENS, - temperature=0.0, - num_samples=1, - logprobs=1, - seed=GREEDY_SEED, - ) - - -def _build_client_rollout(): - """Build a ClientMultiTurnRollout wired to the running e2e server (HTTP).""" - from twinkle_client.sampler import vLLMSampler as ClientVLLMSampler - from twinkle_client.rollout.multi_turn import ClientMultiTurnRollout - - sampler = ClientVLLMSampler(model_id=MODEL_ID) - sampler.set_template(TEMPLATE_CLS, model_id=MODEL_ID, enable_thinking=False) - - return ClientMultiTurnRollout( - sampler=sampler, - template=_build_local_template(), - tool_manager=make_tool_manager(), - sampling_params=_greedy_sampling_params(), - max_turns=ALIGN_MAX_TURNS, - ) - - -def _build_bare_rollout(): - """Build a bare-library MultiTurnRollout holding its own Ray-actor sampler. - - Mirrors the standalone GPU sampler tests / multi-turn cookbook: initialise - Twinkle in Ray mode with a dedicated ``sampler`` device group and construct - a Ray-actor :class:`twinkle.sampler.vLLMSampler` (``remote_group='sampler'``). - """ - import twinkle - from twinkle import DeviceGroup, DeviceMesh - from twinkle.sampler import vLLMSampler as RayVLLMSampler - from twinkle_agentic.rollout.multi_turn import MultiTurnRollout - - sampler_gpus = int(os.environ.get('TWINKLE_ALIGN_SAMPLER_GPUS', '1')) - twinkle.initialize( - mode='ray', - nproc_per_node=sampler_gpus, - groups=[DeviceGroup(name='sampler', ranks=list(range(sampler_gpus)), device_type='GPU')], - lazy_collect=False, - ) - sampler_mesh = DeviceMesh.from_sizes(world_size=sampler_gpus, dp_size=sampler_gpus) - sampler = RayVLLMSampler( - model_id=MODEL_ID, - engine_args={ - 'gpu_memory_utilization': 0.5, - 'max_model_len': 4096, - 'enable_lora': True, - }, - device_mesh=sampler_mesh, - remote_group='sampler', - ) - sampler.set_template(TEMPLATE_CLS, model_id=MODEL_ID, enable_thinking=False) - - return MultiTurnRollout( - sampler=sampler, - template=_build_local_template(), - tool_manager=make_tool_manager(), - sampling_params=_greedy_sampling_params(), - max_turns=ALIGN_MAX_TURNS, - ) - - -# ═══════════════════════════════════════════════════════════════════════════ -# GPU fixture: build both rollouts once for the module (skips without GPU). -# ═══════════════════════════════════════════════════════════════════════════ -@pytest.fixture(scope='module') -def aligned_rollouts(): - """Yield ``(client_rollout, bare_rollout)`` built against the same weights. - - Skipped automatically unless ``TWINKLE_TEST_GPU_E2E=1``. Requires the e2e - server to be already running (client path) and enough GPUs for the in-test - bare Ray-actor sampler. - """ - if not gpu_e2e_enabled(): - pytest.skip('Set TWINKLE_TEST_GPU_E2E=1 to run the real-sampler multi-turn ' - 'alignment E2E (requires a running server + GPU).') - - wait_for_server() - init_twinkle_client_session() - - log('Building client (HTTP) rollout...') - client_rollout = _build_client_rollout() - log('Building bare-library (Ray actor) rollout...') - bare_rollout = _build_bare_rollout() - - yield client_rollout, bare_rollout - - -# ═══════════════════════════════════════════════════════════════════════════ -# GPU test: the two paths agree on control flow, and logprobs/labels align. -# ═══════════════════════════════════════════════════════════════════════════ -def test_client_and_bare_multi_turn_control_flow_aligned(aligned_rollouts): - """ClientMultiTurnRollout and bare MultiTurnRollout agree on control flow. - - Both paths run the SAME prompt / tools / greedy sampling against the SAME - real Qwen3.5-4B weights for 3-6 rounds. We assert: - - * Both satisfy Properties 3/4/6/7 (length/order, stop_reason range, - turns <= max_turns, max_turns => truncated). - * The transport-independent control-flow fingerprint (message role - signature, tool-call timing, stop_reason, turns, truncated) is IDENTICAL - between the two paths for every trajectory. Token-level sampling - non-determinism is allowed, but the control flow must match. - * On each path, ``len(logprobs) == count(labels != -100)`` for every - trajectory that collected logprobs (real-sampler numeric alignment). - """ - client_rollout, bare_rollout = aligned_rollouts - trajectories = build_alignment_trajectories() - n = len(trajectories) - - client_outs = client_rollout(copy.deepcopy(trajectories)) - bare_outs = bare_rollout(copy.deepcopy(trajectories)) - - # Per-path structural invariants. - assert_multi_turn_invariants(client_outs, n, ALIGN_MAX_TURNS, label='client') - assert_multi_turn_invariants(bare_outs, n, ALIGN_MAX_TURNS, label='bare') - - # Per-path real-sampler logprobs/labels alignment (strict equality). - for out in client_outs: - assert_logprobs_align_with_labels(out, label='client') - for out in bare_outs: - assert_logprobs_align_with_labels(out, label='bare') - - # Cross-path control-flow equality (allowing token-level non-determinism). - for i in range(n): - client_fp = control_flow_fingerprint(client_outs[i]) - bare_fp = control_flow_fingerprint(bare_outs[i]) - assert client_fp == bare_fp, ( - f'control-flow mismatch for trajectory {i}:\n' - f' client={client_fp}\n bare ={bare_fp}') - - # The two paths must also agree on the exact trainable-label count, which - # (with greedy decoding on identical weights) is the strongest available - # cross-path numeric-alignment signal short of bit-identical logprobs. - client_trainable = count_trainable_labels(client_outs[i].get('labels')) - bare_trainable = count_trainable_labels(bare_outs[i].get('labels')) - assert client_trainable == bare_trainable, ( - f'trainable-label count mismatch for trajectory {i}: ' - f'client={client_trainable} vs bare={bare_trainable}') - - -# ═══════════════════════════════════════════════════════════════════════════ -# Collection-safe unit tests for the pure comparison helpers (run locally, no -# GPU). These prove the module imports and the fingerprint logic behaves, so the -# file is meaningful even when the GPU test above is skipped. -# ═══════════════════════════════════════════════════════════════════════════ -def test_gpu_gate_reflects_env_flag(): - """The GPU gate reflects TWINKLE_TEST_GPU_E2E without side effects.""" - assert gpu_e2e_enabled() == (os.environ.get('TWINKLE_TEST_GPU_E2E', '0') == '1') - - -def test_count_trainable_labels(): - assert count_trainable_labels(None) == 0 - assert count_trainable_labels([]) == 0 - assert count_trainable_labels([-100, -100]) == 0 - assert count_trainable_labels([-100, 5, 7, -100, 9]) == 3 - - -def test_message_role_signature_and_tool_turns(): - messages = [ - {'role': 'system', 'content': 's'}, - {'role': 'user', 'content': 'u'}, - {'role': 'assistant', 'content': 'call'}, # assistant turn 0 -> tool - {'role': 'tool', 'content': 't'}, - {'role': 'assistant', 'content': 'final'}, # assistant turn 1 -> no tool - ] - assert message_role_signature(messages) == ['system', 'user', 'assistant', 'tool', 'assistant'] - # Only the first assistant turn is immediately followed by a tool response. - assert tool_turn_indices(messages) == [0] - - -def test_control_flow_fingerprint_ignores_tokens_and_logprob_values(): - """Two outputs with identical control flow but different tokens/logprobs - produce an identical fingerprint (token-level non-determinism allowed).""" - base_messages = [ - {'role': 'user', 'content': 'q'}, - {'role': 'assistant', 'content': 'call'}, - {'role': 'tool', 'content': 't'}, - {'role': 'assistant', 'content': 'final'}, - ] - a = { - 'messages': base_messages, - 'stop_reason': 'stop', 'turns': 2, 'truncated': False, - 'labels': [-100, 1, 2], 'logprobs': [[(1, -0.1)], [(2, -0.2)]], - } - b = { - # Same roles/timing, DIFFERENT token content and logprob values. - 'messages': [dict(m) for m in base_messages], - 'stop_reason': 'stop', 'turns': 2, 'truncated': False, - 'labels': [-100, 9, 8], 'logprobs': [[(9, -1.0)], [(8, -2.0)]], - } - assert control_flow_fingerprint(a) == control_flow_fingerprint(b) - - -def test_assert_logprobs_align_with_labels_pass_and_fail(): - # Passing case: 2 logprob entries, 2 trainable labels. - ok = {'labels': [-100, 3, 4], 'logprobs': [[(3, -0.1)], [(4, -0.2)]]} - assert_logprobs_align_with_labels(ok, label='unit') # no raise - - # Empty/None logprobs is a no-op (not trainable). - assert_logprobs_align_with_labels({'labels': [-100, 3], 'logprobs': None}, label='unit') - - # Mismatch raises. - bad = {'labels': [-100, 3, 4], 'logprobs': [[(3, -0.1)]]} - with pytest.raises(AssertionError): - assert_logprobs_align_with_labels(bad, label='unit') - - -def test_assert_multi_turn_invariants_detects_violations(): - good = [ - {'_tid': 0, 'stop_reason': 'stop', 'turns': 1, 'truncated': False}, - {'_tid': 1, 'stop_reason': 'max_turns', 'turns': 3, 'truncated': True}, - ] - assert_multi_turn_invariants(good, n_inputs=2, max_turns=3, label='unit') # no raise - - # stop_reason value-range violation: unknown stop_reason. - with pytest.raises(AssertionError): - assert_multi_turn_invariants( - [{'_tid': 0, 'stop_reason': 'weird', 'turns': 1, 'truncated': False}], - n_inputs=1, max_turns=3, label='unit') - - # Turn-bound violation: turns exceed max_turns. - with pytest.raises(AssertionError): - assert_multi_turn_invariants( - [{'_tid': 0, 'stop_reason': 'stop', 'turns': 5, 'truncated': False}], - n_inputs=1, max_turns=3, label='unit') - - # Truncation-flag violation: max_turns but not truncated. - with pytest.raises(AssertionError): - assert_multi_turn_invariants( - [{'_tid': 0, 'stop_reason': 'max_turns', 'turns': 3, 'truncated': False}], - n_inputs=1, max_turns=3, label='unit') - - # Order/length violation. - with pytest.raises(AssertionError): - assert_multi_turn_invariants( - [{'_tid': 1, 'stop_reason': 'stop', 'turns': 1, 'truncated': False}], - n_inputs=1, max_turns=3, label='unit') - - -def test_build_alignment_trajectories_shape(): - trajs = build_alignment_trajectories() - assert len(trajs) == 2 - for i, traj in enumerate(trajs): - assert traj['_tid'] == i - assert traj['messages'][0]['role'] == 'system' - assert traj['messages'][1]['role'] == 'user' - assert isinstance(traj['tools'], list) and traj['tools'] - names = {t['function']['name'] for t in traj['tools']} - assert names == {'echo', 'calculator'} - - -def test_tool_manager_dispatch_is_deterministic(): - """The echo/calculator tools dispatch deterministically via ToolManager.""" - mgr = make_tool_manager() - calc_call = {'type': 'function', 'function': {'name': 'calculator', - 'arguments': {'op': 'add', 'a': 21, 'b': 34}}} - echo_call = {'type': 'function', 'function': {'name': 'echo', 'arguments': {'text': 'ping'}}} - assert json.loads(mgr(calc_call))['result'] == 55 - assert mgr(echo_call) == 'echo[echo]:{"text": "ping"}' diff --git a/tests/server/integration/test_client_multi_turn_rollout_mock_e2e.py b/tests/server/integration/test_client_multi_turn_rollout_mock_e2e.py deleted file mode 100644 index 89e57e22..00000000 --- a/tests/server/integration/test_client_multi_turn_rollout_mock_e2e.py +++ /dev/null @@ -1,563 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""Local CPU-only multi-turn control-flow E2E for ``ClientMultiTurnRollout``. - -This test boots a REAL CPU-only Twinkle server (Ray Serve, in-process) whose -sampler backend is the enhanced :class:`MockSampler`, then drives -:class:`twinkle_client.rollout.multi_turn.ClientMultiTurnRollout` over ACTUAL -HTTP through :class:`twinkle_client.sampler.vLLMSampler`. No GPU is required. - -What is "real" vs "test double" here ------------------------------------- - * REAL: the HTTP server (gateway + sampler apps on Ray Serve), the - ``/twinkle/sample`` protocol hop, ``vLLMSampler.sample()`` and the - enhanced ``MockSampler`` producing ``new_input_feature`` / configurable - ``stop_reason`` / configurable tool-call text over the wire. - * TEST DOUBLE (client-local only): a lightweight char-level Template used by - the client to ``encode`` trajectories and ``parse_tool_call`` the sampled - text, plus a small ``ToolManager`` with an echo tool. The Template never - crosses the network; encoding, tool parsing and bridge stitching are - client-side concerns. A real HF Template would require a downloaded - tokenizer (network/model weights) which is unavailable in a local offline - CPU environment, so a deterministic char-level double is used instead — - exactly the established pattern from - ``tests/twinkle_client/test_client_multi_turn_rollout.py``. - -Why the sampler knobs are set at CONSTRUCTION time --------------------------------------------------- -The multi-turn knobs (``stop_reason`` / ``tool_call_text`` / ``tool_call_turns``) -CANNOT be injected per-call through the HTTP layer: the ``/twinkle/sample`` -handler rebuilds ``sampling_params`` via ``SamplingParams.from_dict`` which -filters the payload down to the known dataclass fields, dropping any extra -knobs before they reach ``MockSampler.sample``. Setting them per call would -require server-side protocol changes beyond this task's scope. Therefore each -behaviour is realised as a SEPARATE sampler app whose ``MockSampler`` is -constructed with the desired knobs: - - * ``mock-tool`` — ``stop_reason='stop'`` + a tool call injected on EVERY - round (``tool_call_turns`` covers a wide range). This deterministically - exercises the "sample -> tool -> bridge -> sample -> ... -> max_turns - truncation" control flow regardless of the sampler's monotonic per-call - round counter (so the module-scoped server can be reused across cases). - * ``mock-stop`` — ``stop_reason='stop'`` with NO tool call, i.e. natural - single-turn termination with ``stop_reason == 'stop'``. - * ``mock-length`` — ``stop_reason='length'``, i.e. immediate termination with - ``stop_reason == 'length'``. - -============================================================================== -CONDA ENVIRONMENT REQUIREMENT (MANDATORY) — NO GPU REQUIRED -============================================================================== -ALL cases in this file are CPU-only and MUST run inside the ``twinkle`` conda env: - - conda run -n twinkle pytest \ - tests/server/integration/test_client_multi_turn_rollout_mock_e2e.py -v -============================================================================== -""" -from __future__ import annotations - -import copy -import json -import os -import re -import subprocess -import sys -import time -from typing import Any, Dict, List, Optional - -# Ensure project root is importable for both pytest and direct execution. -_PROJECT_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', '..')) -if _PROJECT_ROOT not in sys.path: - sys.path.insert(0, _PROJECT_ROOT) - -import pytest - -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 - -# ═══════════════════════════════════════════════════════════════════════════ -# Sampler model ids (one app per multi-turn behaviour — see module docstring). -# ═══════════════════════════════════════════════════════════════════════════ -MODEL_TOOL = 'mock-tool' # stop_reason='stop' + tool call on every round -MODEL_STOP = 'mock-stop' # stop_reason='stop', no tool call -MODEL_LENGTH = 'mock-length' # stop_reason='length' - -# Tool call text the mock emits as SampledSequence.decoded on injected rounds. -# Must match the client Template's parse_tool_call (Qwen {...}). -TOOL_NAME = 'echo' -TOOL_CALL_TEXT = '' + json.dumps({'name': TOOL_NAME, 'arguments': {'text': 'hi'}}) + '' - -# Wide range so "inject a tool call on every round" holds even as the mock's -# monotonic per-call round counter advances across reused test cases. -_ALWAYS_INJECT_TURNS = list(range(1, 201)) - -# Sampling params carried on every round. ``max_tokens`` MUST be set: the mock -# rejects max_tokens < 1, and the rollout's default SamplingParams leaves it None. -SAMPLE_MAX_TOKENS = 4 - - -# ═══════════════════════════════════════════════════════════════════════════ -# Client-local test doubles: char-level tokenizer + Template + echo tool. -# (Mirrors tests/twinkle_client/test_client_multi_turn_rollout.py.) -# ═══════════════════════════════════════════════════════════════════════════ -class _FakeTokenizer: - """Char-level tokenizer with atomic special tokens. - - Guarantees ``decode(encode(s)) == s`` so ``extend_with_bridge``'s - template-space delta computation is deterministic. - """ - 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: - s += f"<|im_start|>{m['role']}\n{m['content']}<|im_end|>\n" - if add_generation_prompt: - s += '<|im_start|>assistant\n' - return self.encode(s) if tokenize else s - - -class _FakeTemplate: - """Minimal Template mirroring the parts ClientMultiTurnRollout touches. - - A client-local test double (never crosses the network). Implements exactly - ``encode`` / ``parse_tool_call`` / ``_invoke_post_pipeline`` / - ``enable_thinking`` / ``tokenizer`` — everything ``extend_with_bridge`` and - the rollout need. - """ - 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)})') - labels = labels[1:] + labels[:1] # np.roll(labels, -1): shift LEFT by 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 - - -class EchoTool(Tool): - """Echoes its arguments as a JSON string.""" - - def __init__(self, name: str = TOOL_NAME): - 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': {}}, - } - - -def _make_tool_manager() -> ToolManager: - mgr = ToolManager({}) - mgr.register(EchoTool(TOOL_NAME)) - return mgr - - -def _make_template() -> _FakeTemplate: - return _FakeTemplate(_FakeTokenizer()) - - -def _make_trajectories(n: int) -> List[Dict[str, Any]]: - """Build ``n`` single-user-message trajectories tagged with a hidden _tid.""" - return [{'messages': [{'role': 'user', 'content': f'q{tid}'}], '_tid': tid} for tid in range(n)] - - -# ═══════════════════════════════════════════════════════════════════════════ -# CPU-only mock server config. -# -# Built as plain dicts and fed DIRECTLY to the app builders (build_gateway_app / -# build_sampler_app) instead of through ``ServerConfig.model_validate``. The -# typed ``SamplerArgs`` schema uses ``extra='forbid'`` and does NOT declare the -# mock multi-turn knobs (stop_reason / tool_call_text / tool_call_turns), so -# routing them through the validated config would be rejected. The builders -# themselves accept ``**kwargs`` and forward them to the MockSampler ctor, so we -# skip validation and call them directly (this is test-only wiring). -# ═══════════════════════════════════════════════════════════════════════════ -def _sampler_app(name: str, model_id: str, *, extra_args: Dict[str, Any]) -> Dict[str, Any]: - # NOTE: ``queue_config`` is intentionally omitted. When calling the builder - # directly (bypassing ServerConfig validation) a raw dict would NOT be - # coerced into a TaskQueueConfig; leaving it unset lets the sampler app - # construct a default TaskQueueConfig, which is fine for this CPU test. - args: Dict[str, Any] = { - 'sampler_type': 'mock', - 'model_id': model_id, - 'nproc_per_node': 1, - 'device_group': {'name': f'sampler_{model_id}', 'ranks': 1, 'device_type': 'CPU'}, - 'device_mesh': {'device_type': 'CPU', 'dp_size': 1}, - } - args.update(extra_args) - return { - 'name': name, - 'route_prefix': f'/api/v1/sampler/{model_id}', - 'import_path': 'sampler', - 'args': args, - } - - -def _build_applications() -> List[Dict[str, Any]]: - """Return the plain-dict application specs (gateway + 3 sampler apps).""" - return [ - { - 'name': 'server', - 'route_prefix': '/api/v1', - 'import_path': 'server', - 'args': { - 'server_config': {'per_token_model_limit': 3}, - 'supported_models': [MODEL_TOOL, MODEL_STOP, MODEL_LENGTH], - }, - }, - _sampler_app('sampler-mock-tool', MODEL_TOOL, extra_args={ - 'stop_reason': 'stop', - 'tool_call_text': TOOL_CALL_TEXT, - 'tool_call_turns': _ALWAYS_INJECT_TURNS, - }), - _sampler_app('sampler-mock-stop', MODEL_STOP, extra_args={'stop_reason': 'stop'}), - _sampler_app('sampler-mock-length', MODEL_LENGTH, extra_args={'stop_reason': 'length'}), - ] - - -# File-backed persistence path (gateway + samplers run as separate replicas, so -# ``memory`` mode is insufficient — cross-process visibility requires a file). -_PERSISTENCE_FILE = '/tmp/twinkle_state_multi_turn_mock.json' - - -# ═══════════════════════════════════════════════════════════════════════════ -# In-process CPU-only Ray Serve harness (mirrors MockEmbeddingServerHarness). -# ═══════════════════════════════════════════════════════════════════════════ -class MultiTurnMockServerHarness: - """Boots the CPU-only mock multi-turn server in-process via Ray Serve. - - Manages its own local Ray head node (independent of the session-scoped Ray - fixture), starts Ray Serve on a randomized port, and runs the gateway plus - all three sampler apps declared by ``_build_server_config_dict``. No GPU. - """ - - READY_BUDGET_SECONDS = 90.0 - RAY_NODE_CPUS = 8 - - def __init__(self) -> None: - self.host = '127.0.0.1' - self.port = 18900 + (os.getpid() % 700) - self.base_url = f'http://{self.host}:{self.port}' - self._started = False - - @staticmethod - def _run_ray_command(*args: str) -> None: - ray_bin = os.path.join(os.path.dirname(sys.executable), 'ray') - result = subprocess.run( - [ray_bin, *args], check=False, - stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True) - if result.returncode != 0: - raise RuntimeError(f'ray {" ".join(args)} failed ({result.returncode}):\n{result.stdout}') - - def start(self) -> str: - import ray - from ray import serve - - from twinkle.server.config.persistence import PersistenceConfig - from twinkle.server.gateway import build_gateway_app - from twinkle.server.sampler import build_sampler_app - - # File-backed persistence so gateway + sampler replicas share state. - persistence_env = PersistenceConfig(mode='file', file_path=_PERSISTENCE_FILE).to_env_vars() - for k, v in persistence_env.items(): - os.environ[k] = v - - if ray.is_initialized(): - ray.shutdown() - self._run_ray_command('stop', '--force') - self._run_ray_command( - 'start', '--head', '--port=0', f'--num-cpus={self.RAY_NODE_CPUS}', - '--num-gpus=0', '--include-dashboard=false', '--disable-usage-stats') - ray.init( - address='auto', - runtime_env={'env_vars': persistence_env} if persistence_env else None, - ) - self._started = True - - serve.start(http_options={'host': self.host, 'port': self.port}) - # Call the builders DIRECTLY with plain-dict args (see _build_applications): - # this bypasses the typed SamplerArgs schema (extra='forbid') so the mock - # multi-turn knobs reach the MockSampler ctor. - builders = {'server': build_gateway_app, 'sampler': build_sampler_app} - deploy_options: Dict[str, Any] = {'ray_actor_options': {'num_cpus': 0.1}} - for app_spec in _build_applications(): - builder = builders[app_spec['import_path']] - args = dict(app_spec['args']) - if app_spec['import_path'] == 'server': - args.setdefault('http_options', {'host': self.host, 'port': self.port}) - bound = builder(deploy_options=deploy_options, **args) - serve.run(bound, name=app_spec['name'], route_prefix=app_spec['route_prefix']) - - self._wait_until_healthy(serve, self.READY_BUDGET_SECONDS) - return self.base_url - - def _wait_until_healthy(self, serve_module: Any, timeout: float) -> None: - deadline = time.monotonic() + timeout - last: Dict[str, Any] = {} - while time.monotonic() < deadline: - status = serve_module.status() - last = {name: app.status for name, app in status.applications.items()} - if last and all(s == 'RUNNING' for s in last.values()): - return - time.sleep(0.5) - raise TimeoutError(f'Mock multi-turn server not RUNNING within {timeout}s: {last}') - - def stop(self) -> None: - if not self._started: - return - try: - import ray - from ray import serve - try: - serve.shutdown() - except Exception: - pass - try: - ray.shutdown() - except Exception: - pass - finally: - try: - self._run_ray_command('stop', '--force') - except Exception: - pass - self._started = False - - -# ═══════════════════════════════════════════════════════════════════════════ -# Fixtures -# ═══════════════════════════════════════════════════════════════════════════ -@pytest.fixture(scope='module') -def mock_multi_turn_server(): - """Boot the CPU-only mock multi-turn server ONCE for the module.""" - harness = MultiTurnMockServerHarness() - base_url = harness.start() - try: - from twinkle_client import init_twinkle_client - init_twinkle_client(base_url=base_url, api_key='EMPTY_TOKEN') - yield base_url - finally: - harness.stop() - - -def _make_sampler(model_id: str): - """Create a real vLLMSampler bound to the given mock sampler app.""" - from twinkle_client.sampler import vLLMSampler - return vLLMSampler(model_id=model_id) - - -def _make_rollout(model_id: str, *, tool_manager: Optional[ToolManager], max_turns: int) -> ClientMultiTurnRollout: - return ClientMultiTurnRollout( - sampler=_make_sampler(model_id), - template=_make_template(), - tool_manager=tool_manager, - sampling_params=SamplingParams(max_tokens=SAMPLE_MAX_TOKENS, num_samples=1), - max_turns=max_turns, - ) - - -# ═══════════════════════════════════════════════════════════════════════════ -# Multi-turn control flow over real HTTP (tool app: tool call on every round). -# -# With ``stop_reason='stop'`` and a tool call injected every round, the loop -# runs "sample -> tool -> bridge -> sample -> ..." until it hits ``max_turns`` -# and force-truncates. Exercises the output length/order, stop_reason value -# range ('max_turns' in the allowed set), and turns <= max_turns invariants. -# ═══════════════════════════════════════════════════════════════════════════ -@pytest.mark.parametrize('n_traj,max_turns', [(1, 3), (3, 2), (2, 4)]) -def test_multi_turn_tool_loop_over_http(mock_multi_turn_server, n_traj, max_turns): - """Full multi-turn control flow over real HTTP terminates at max_turns. - - Boots a real CPU server, drives ClientMultiTurnRollout via vLLMSampler HTTP - calls through several "sample -> tool -> bridge -> sample" rounds, and - checks the batch invariants. - """ - rollout = _make_rollout(MODEL_TOOL, tool_manager=_make_tool_manager(), max_turns=max_turns) - trajectories = _make_trajectories(n_traj) - - outs = rollout(copy.deepcopy(trajectories)) - - # Output list is same length and order as the input. - assert len(outs) == n_traj - for i, out in enumerate(outs): - assert out['_tid'] == i, 'output order must match input order' - - for out in outs: - # stop_reason is within the allowed set. - assert out['stop_reason'] in {'length', 'stop', 'max_turns'}, out['stop_reason'] - # Actual turns never exceed the configured max_turns. - assert out['turns'] <= max_turns, f"turns({out['turns']}) > max_turns({max_turns})" - # A tool call on every round means the loop must hit the turn cap. - assert out['stop_reason'] == 'max_turns' - assert out['truncated'] is True - assert out['turns'] == max_turns - # Multi-turn actually stitched tool turns via the shared bridge helper: - # the running context grew past the initial prompt encoding. - assert out['messages'][0]['role'] == 'user' - if max_turns >= 2: - # At least one tool message was appended by extend_with_bridge. - assert any(m['role'] == 'tool' for m in out['messages']) - - -# ═══════════════════════════════════════════════════════════════════════════ -# max_turns == 1 with a first-round tool call forces truncation. -# ═══════════════════════════════════════════════════════════════════════════ -def test_max_turns_one_forces_truncation(mock_multi_turn_server): - """max_turns==1 + first-round tool call -> truncated=True, stop_reason='max_turns'.""" - rollout = _make_rollout(MODEL_TOOL, tool_manager=_make_tool_manager(), max_turns=1) - trajectories = _make_trajectories(3) - - outs = rollout(copy.deepcopy(trajectories)) - - assert len(outs) == 3 - for i, out in enumerate(outs): - assert out['_tid'] == i - assert out['truncated'] is True - assert out['stop_reason'] == 'max_turns' - assert out['turns'] == 1 - - -# ═══════════════════════════════════════════════════════════════════════════ -# Natural termination reasons ('stop' and 'length') over real HTTP. -# ═══════════════════════════════════════════════════════════════════════════ -def test_natural_stop_termination_over_http(mock_multi_turn_server): - """No tool call + stop_reason='stop' -> single-turn natural termination.""" - rollout = _make_rollout(MODEL_STOP, tool_manager=_make_tool_manager(), max_turns=4) - trajectories = _make_trajectories(2) - - outs = rollout(copy.deepcopy(trajectories)) - - assert len(outs) == 2 - for i, out in enumerate(outs): - assert out['_tid'] == i - assert out['stop_reason'] == 'stop' - assert out['truncated'] is False - assert out['turns'] == 1 - assert out['turns'] <= 4 - - -def test_length_termination_over_http(mock_multi_turn_server): - """stop_reason='length' -> immediate termination on the first round.""" - rollout = _make_rollout(MODEL_LENGTH, tool_manager=_make_tool_manager(), max_turns=4) - trajectories = _make_trajectories(2) - - outs = rollout(copy.deepcopy(trajectories)) - - assert len(outs) == 2 - for out in outs: - assert out['stop_reason'] == 'length' - assert out['truncated'] is False - assert out['turns'] == 1 - - -# ═══════════════════════════════════════════════════════════════════════════ -# Exception path: tool_calls produced but no tool_manager -> ValueError. -# -# Note on ``new_input_feature=None``: the enhanced MockSampler ALWAYS populates -# new_input_feature, so that specific error path is not reproducible against -# a real mock server and is covered by the unit tests in -# tests/twinkle_client/test_client_multi_turn_rollout.py instead. The -# tool_manager-missing path IS reachable over real HTTP and is asserted here. -# ═══════════════════════════════════════════════════════════════════════════ -def test_tool_calls_without_tool_manager_raises_value_error_over_http(mock_multi_turn_server): - """A tool call with no tool_manager raises ValueError over the real HTTP path.""" - rollout = _make_rollout(MODEL_TOOL, tool_manager=None, max_turns=3) - trajectories = _make_trajectories(1) - - 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 diff --git a/tests/server/state/test_leader_election.py b/tests/server/state/test_leader_election.py index 29d6122d..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; diff --git a/tests/server/test_embedding_e2e.py b/tests/server/test_embedding_e2e.py deleted file mode 100644 index 240e36fd..00000000 --- a/tests/server/test_embedding_e2e.py +++ /dev/null @@ -1,1133 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""Embedding training E2E test scaffolding (Twinkle client path). - -This module provides the reusable scaffolding for validating the embedding -training pipeline end to end through the Twinkle client: - - set_processor('InputProcessor') - -> set_loss('InfonceLoss', temperature=..., use_batch=True) - -> add_metric('EmbeddingMetric', is_training=True) - -> loop: forward_backward(inputs=mb, task='embedding') + clip_grad_and_step(...) - -> calculate_metric(is_training=True) - -It exposes TWO switchable backend entrypoints (fixtures): - - * ``mock_embedding_model`` — local CPU-only entrypoint backed by - ``TwinkleCompatMockModel`` (src/twinkle/server/model/backends/mock_model.py). - Requires NO GPU. Boots an in-process Ray Serve cluster from the CPU-only - mock server config fixture. Consumed by the local CPU protocol/link validation. - - * ``gpu_embedding_model`` — real transformers model entrypoint. Gated behind - the ``TWINKLE_TEST_GPU_E2E=1`` environment variable (skipped otherwise) and - connects to an already-running GPU server (see - tests/server/start_e2e_server.py). Consumed by the real numeric - validation cases. - -Plus two reusable helpers: - - * ``build_synthetic_contrastive_dataset(...)`` — builds a small synthetic - anchor/positive contrastive-learning dataset (even number of samples, with - ``labels = [1, 0, 1, 0, ...]`` anchor/positive semantics) in the embedding - pooling input format consumed by ``InputProcessor``. - * ``run_embedding_training(...)`` — issues the "verification - call sequence" against a client model and returns the observed losses and - the final metric result. - -============================================================================== -CONDA ENVIRONMENT REQUIREMENT (MANDATORY) -============================================================================== -ALL test cases in this file MUST be run inside the ``twinkle`` conda env, e.g.: - - conda run -n twinkle pytest tests/server/test_embedding_e2e.py -v - -The local, CPU-only mock cases run without a GPU and must pass in -that environment: - - conda run -n twinkle pytest tests/server/test_embedding_e2e.py -v -k mock - -The GPU-gated cases additionally require ``TWINKLE_TEST_GPU_E2E=1`` -and a running GPU server; they are skipped automatically when the variable is -unset: - - TWINKLE_TEST_GPU_E2E=1 conda run -n twinkle pytest tests/server/test_embedding_e2e.py -v -k gpu -============================================================================== -""" -from __future__ import annotations - -import math -import os -import subprocess -import sys -import time -from typing import Any, Dict, List, Optional - -# Ensure project root is importable for both pytest and direct execution. -_PROJECT_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..')) -if _PROJECT_ROOT not in sys.path: - sys.path.insert(0, _PROJECT_ROOT) - -import pytest - -# Reuse the existing GPU e2e helpers (server URL, session init, real model id). -from tests.server.integration.e2e_helpers import ( - BASE_URL, - MODEL_ID, - init_twinkle_client_session, - log, - wait_for_server, -) - -# ═══════════════════════════════════════════════════════════════════════════ -# Configuration -# ═══════════════════════════════════════════════════════════════════════════ - -# Local CPU-only mock backend model id (see tests/server/fixtures/server_config_mock.yaml). -MOCK_MODEL_ID = 'mock-model' - -# InfoNCE temperature used by the verification call sequence. -DEFAULT_TEMPERATURE = 0.07 - -# Default synthetic-dataset shape (kept tiny for CPU speed). -DEFAULT_NUM_PAIRS = 4 -DEFAULT_SEQ_LEN = 8 -DEFAULT_VOCAB_SIZE = 32 - -# Number of training steps used by the default verification loop. -DEFAULT_TRAIN_STEPS = 4 - - -def gpu_e2e_enabled() -> bool: - """Return True when GPU e2e tests are explicitly enabled.""" - return os.environ.get('TWINKLE_TEST_GPU_E2E', '0') == '1' - - -# ═══════════════════════════════════════════════════════════════════════════ -# Synthetic contrastive-learning dataset -# ═══════════════════════════════════════════════════════════════════════════ - -def build_synthetic_contrastive_dataset( - num_pairs: int = DEFAULT_NUM_PAIRS, - *, - seq_len: int = DEFAULT_SEQ_LEN, - vocab_size: int = DEFAULT_VOCAB_SIZE, - seed: int = 0, -) -> List[Dict[str, Any]]: - """Build a small synthetic anchor/positive contrastive dataset. - - Produces ``2 * num_pairs`` samples interleaved as - ``[anchor_0, positive_0, anchor_1, positive_1, ...]`` so that the per-sample - ``labels`` scalar carries the ``[1, 0, 1, 0, ...]`` anchor/positive semantics - expected by embedding pooling (anchor -> 1, positive -> 0). Each sample is a - feature dict in the embedding pooling input format consumed by - ``InputProcessor`` (``input_ids`` + ``attention_mask`` + ``labels``). - - Anchor and its positive within a pair share correlated token content so a - real InfoNCE loss has a meaningful contrastive signal to learn from, while a - fixed ``seed`` keeps the dataset deterministic across runs. - - Args: - num_pairs: Number of anchor/positive pairs (result has ``2*num_pairs`` samples). - seq_len: Token sequence length per sample. - vocab_size: Upper bound (exclusive) for synthetic token ids. - seed: RNG seed for deterministic generation. - - Returns: - A list of feature dicts of length ``2 * num_pairs``. - """ - if num_pairs < 1: - raise ValueError(f'num_pairs must be >= 1, got {num_pairs}') - if seq_len < 1: - raise ValueError(f'seq_len must be >= 1, got {seq_len}') - - import numpy as np - - rng = np.random.default_rng(seed) - samples: List[Dict[str, Any]] = [] - for pair_idx in range(num_pairs): - # Shared "concept" tokens make anchor/positive semantically correlated. - base = rng.integers(low=1, high=vocab_size, size=seq_len) - # Anchor: the base sequence. - anchor_ids = base.copy() - # Positive: perturb a couple of positions so it is similar but not identical. - positive_ids = base.copy() - n_perturb = max(1, seq_len // 4) - perturb_pos = rng.choice(seq_len, size=n_perturb, replace=False) - positive_ids[perturb_pos] = rng.integers(low=1, high=vocab_size, size=n_perturb) - - samples.append(_make_embedding_feature(anchor_ids.tolist(), label=1)) - samples.append(_make_embedding_feature(positive_ids.tolist(), label=0)) - - return samples - - -def _make_embedding_feature(input_ids: List[int], *, label: int) -> Dict[str, Any]: - """Build a single embedding pooling feature dict. - - ``labels`` is a single scalar per sample (``[label]``) — this mirrors the - anchor/positive grouping used by the bare-library embedding example - (cookbook/exp/embedding/train_embedding_full_ddp.py) where anchors carry - ``labels=[1]`` and positives carry ``labels=[0]``. - """ - seq_len = len(input_ids) - return { - 'input_ids': list(input_ids), - 'attention_mask': [1] * seq_len, - 'labels': [int(label)], - } - - -def iter_minibatches(dataset: List[Dict[str, Any]], batch_size: int) -> List[List[Dict[str, Any]]]: - """Split a dataset into contiguous minibatches. - - ``batch_size`` should be even so anchor/positive pairs stay together inside - each minibatch (InfoNCE assumes paired samples within a batch). - """ - if batch_size < 2 or batch_size % 2 != 0: - raise ValueError(f'batch_size must be a positive even number, got {batch_size}') - return [dataset[i:i + batch_size] for i in range(0, len(dataset), batch_size)] - - -# ═══════════════════════════════════════════════════════════════════════════ -# Verification call sequence -# ═══════════════════════════════════════════════════════════════════════════ - -def configure_embedding_adapter( - model: Any, - *, - temperature: float = DEFAULT_TEMPERATURE, - use_batch: bool = True, - hard_negatives: Optional[int] = None, -) -> None: - """Apply the embedding-training configuration steps to a client model. - - Order matches the verification call sequence: - set_processor('InputProcessor') - -> set_loss('InfonceLoss', temperature=..., use_batch=True) - -> add_metric('EmbeddingMetric', is_training=True) - - Note: ``TransformersEmbeddingPatch`` is applied/rolled back automatically - inside ``forward`` via ``_resolve_task_context`` when ``task='embedding'`` is - passed, so there is intentionally no ``apply_patch(...)`` call here. - """ - model.set_processor('InputProcessor') - model.set_loss('InfonceLoss', temperature=temperature, use_batch=use_batch, hard_negatives=hard_negatives) - model.add_metric('EmbeddingMetric', is_training=True) - - -def extract_loss(fwd_bwd_response: Any) -> float: - """Best-effort extraction of the scalar loss from a forward_backward response. - - Handles both backend shapes: - * real transformers backend: ``result`` is a ModelOutput/dict with a - ``'loss'`` key (see TransformersModel.forward_backward). - * mock backend: ``result`` is ``[records, loss]`` (see - TwinkleCompatMockModel.forward_backward). - """ - result = getattr(fwd_bwd_response, 'result', fwd_bwd_response) - if isinstance(result, dict) and 'loss' in result: - return float(result['loss']) - if isinstance(result, (list, tuple)) and result and isinstance(result[-1], (int, float)): - return float(result[-1]) - raise AssertionError(f'Could not extract loss from forward_backward response: {result!r}') - - -def run_embedding_training( - model: Any, - minibatches: List[List[Dict[str, Any]]], - *, - temperature: float = DEFAULT_TEMPERATURE, - max_grad_norm: float = 1.0, - configure: bool = True, -) -> Dict[str, Any]: - """Run the design-document embedding-training verification call sequence. - - Steps: - (optional) set_processor -> set_loss -> add_metric - loop over minibatches: forward_backward(task='embedding') + clip_grad_and_step - calculate_metric(is_training=True) - - Args: - model: A configured client model (mock or GPU entrypoint). - minibatches: List of minibatches, each a list of embedding feature dicts. - temperature: InfoNCE temperature (used only when ``configure=True``). - max_grad_norm: Grad clipping threshold passed to ``clip_grad_and_step``. - configure: When True, apply ``configure_embedding_adapter`` first. - - Returns: - Dict with keys ``losses`` (list[float]) and ``metric`` (final metric result). - """ - if configure: - configure_embedding_adapter(model, temperature=temperature) - - losses: List[float] = [] - for step, mb in enumerate(minibatches): - fwd_bwd = model.forward_backward(inputs=mb, task='embedding') - loss = extract_loss(fwd_bwd) - losses.append(loss) - model.clip_grad_and_step(max_grad_norm=max_grad_norm) - log(f'[embedding step {step}] loss={loss:.6f}') - - metric = model.calculate_metric(is_training=True) - metric_result = getattr(metric, 'result', metric) - return {'losses': losses, 'metric': metric_result} - - -# ═══════════════════════════════════════════════════════════════════════════ -# Client model builder (shared by both entrypoints) -# ═══════════════════════════════════════════════════════════════════════════ - -def build_embedding_client_model(model_id: str, *, adapter_name: str = 'emb_adapter') -> Any: - """Create a ``MultiLoraTransformersModel`` with a LoRA adapter for embedding. - - The InfoNCE loss / EmbeddingMetric / processor configuration is applied - separately by ``configure_embedding_adapter`` so that callers can inspect or - reorder the verification call sequence. - """ - from peft import LoraConfig - from twinkle_client.model import MultiLoraTransformersModel - - model = MultiLoraTransformersModel(model_id=model_id) - model.add_adapter_to_model(adapter_name, LoraConfig(target_modules='all-linear')) - model.set_template('Qwen3_5Template') - return model - - -# ═══════════════════════════════════════════════════════════════════════════ -# Local CPU-only mock server harness -# ═══════════════════════════════════════════════════════════════════════════ - -class MockEmbeddingServerHarness: - """Boots the CPU-only mock server in-process via Ray Serve. - - Mirrors the proven boot sequence in - tests/server/integration/test_mock_mode_startup.py: it manages its own local - Ray head node (so it does not depend on the session-scoped Ray fixture), - starts Ray Serve on a randomized port, and runs every application declared in - the CPU-only mock server config fixture (server + mock model + mock sampler). - - No GPU is required. - """ - - READY_BUDGET_SECONDS = 60.0 - RAY_NODE_CPUS = 8 - - def __init__(self) -> None: - self.host = '127.0.0.1' - self.port = 18100 + (os.getpid() % 800) - self.base_url = f'http://{self.host}:{self.port}' - self._started = False - - # ----- Ray lifecycle ------------------------------------------------- # - - @staticmethod - def _run_ray_command(*args: str) -> None: - ray_bin = os.path.join(os.path.dirname(sys.executable), 'ray') - result = subprocess.run( - [ray_bin, *args], - check=False, - stdout=subprocess.PIPE, - stderr=subprocess.STDOUT, - text=True, - ) - if result.returncode != 0: - raise RuntimeError(f'ray {" ".join(args)} failed ({result.returncode}):\n{result.stdout}') - - def start(self) -> str: - """Boot Ray + Serve + all mock apps; return the base URL when ready.""" - import ray - from ray import serve - - from tests.server.fixtures import MOCK_SERVER_CONFIG - from twinkle.server.config import ServerConfig - from twinkle.server.gateway import build_gateway_app - from twinkle.server.model import build_model_app - from twinkle.server.sampler import build_sampler_app - - cfg = ServerConfig.from_yaml(MOCK_SERVER_CONFIG) - - persistence_env: Dict[str, str] = {} - if cfg.persistence is not None: - persistence_env = cfg.persistence.to_env_vars() - for k, v in persistence_env.items(): - os.environ[k] = v - - if ray.is_initialized(): - ray.shutdown() - self._run_ray_command('stop', '--force') - self._run_ray_command( - 'start', '--head', '--port=0', f'--num-cpus={self.RAY_NODE_CPUS}', - '--num-gpus=0', '--include-dashboard=false', '--disable-usage-stats') - ray.init( - address='auto', - runtime_env={'env_vars': persistence_env} if persistence_env else None, - ) - self._started = True - - serve.start(http_options={'host': self.host, 'port': self.port}) - builders = { - 'server': build_gateway_app, - 'model': build_model_app, - 'sampler': build_sampler_app, - } - for app_spec in cfg.applications: - builder = builders[app_spec.import_path] - args = {k: v for k, v in dict(app_spec.args).items() if v is not None} - if app_spec.import_path == 'server': - http_opts = cfg.http_options.model_dump() - http_opts['host'] = self.host - http_opts['port'] = self.port - args.setdefault('http_options', http_opts) - deploy_options: Dict[str, Any] = {'ray_actor_options': {'num_cpus': 0.1}} - for raw in app_spec.deployments: - if isinstance(raw, dict): - deploy_options = { - k: v - for k, v in raw.items() if k not in ('name', 'ray_actor_options', 'autoscaling_config') - } - deploy_options['ray_actor_options'] = {'num_cpus': 0.1} - break - bound = builder(deploy_options=deploy_options, **args) - serve.run(bound, name=app_spec.name, route_prefix=app_spec.route_prefix) - - self._wait_until_healthy(serve, self.READY_BUDGET_SECONDS) - return self.base_url - - def _wait_until_healthy(self, serve_module: Any, timeout: float) -> None: - deadline = time.monotonic() + timeout - last: Dict[str, Any] = {} - while time.monotonic() < deadline: - status = serve_module.status() - last = {name: app.status for name, app in status.applications.items()} - if last and all(s == 'RUNNING' for s in last.values()): - return - time.sleep(0.5) - raise TimeoutError(f'Mock server not RUNNING within {timeout}s: {last}') - - def stop(self) -> None: - if not self._started: - return - try: - import ray - from ray import serve - try: - serve.shutdown() - except Exception: - pass - try: - ray.shutdown() - except Exception: - pass - finally: - try: - self._run_ray_command('stop', '--force') - except Exception: - pass - self._started = False - - -# ═══════════════════════════════════════════════════════════════════════════ -# Fixtures: the two switchable backend entrypoints -# ═══════════════════════════════════════════════════════════════════════════ - -@pytest.fixture(scope='module') -def mock_embedding_model(): - """Local CPU-only mock entrypoint (no GPU required). - - Boots the mock server in-process and yields a client model configured with a - LoRA adapter, ready for the embedding verification call sequence. Consumed by - the local CPU protocol/link validation cases. - """ - harness = MockEmbeddingServerHarness() - base_url = harness.start() - try: - from twinkle_client import init_twinkle_client - init_twinkle_client(base_url=base_url, api_key='EMPTY_TOKEN') - model = build_embedding_client_model(MOCK_MODEL_ID) - yield model - finally: - harness.stop() - - -@pytest.fixture(scope='module') -def gpu_embedding_model(): - """GPU real-model entrypoint (gated by TWINKLE_TEST_GPU_E2E=1). - - Connects to an already-running GPU server (see start_e2e_server.py) and - yields a client model configured with a LoRA adapter. Consumed by tasks - 4.3-4.5. Skipped automatically when GPU e2e is not enabled. - """ - if not gpu_e2e_enabled(): - pytest.skip('Set TWINKLE_TEST_GPU_E2E=1 to run real GPU embedding E2E tests (requires running server)') - - wait_for_server() - init_twinkle_client_session() - model = build_embedding_client_model(MODEL_ID) - yield model - - -# ═══════════════════════════════════════════════════════════════════════════ -# Smoke tests — verify the scaffolding collects/imports and helpers behave. -# (Real protocol/numeric assertions live in tasks 4.2-4.5.) -# ═══════════════════════════════════════════════════════════════════════════ - -def test_synthetic_dataset_shape_and_labels(): - """Synthetic dataset is even-length with [1,0,1,0,...] anchor/positive labels.""" - num_pairs = 3 - dataset = build_synthetic_contrastive_dataset(num_pairs, seq_len=6, vocab_size=16, seed=123) - - assert len(dataset) == 2 * num_pairs - labels = [sample['labels'][0] for sample in dataset] - assert labels == [1, 0] * num_pairs - - for sample in dataset: - assert set(sample.keys()) == {'input_ids', 'attention_mask', 'labels'} - assert len(sample['input_ids']) == 6 - assert len(sample['attention_mask']) == 6 - assert all(m == 1 for m in sample['attention_mask']) - - -def test_synthetic_dataset_is_deterministic(): - """Same seed -> identical dataset (keeps mock/GPU comparisons reproducible).""" - a = build_synthetic_contrastive_dataset(2, seq_len=5, seed=7) - b = build_synthetic_contrastive_dataset(2, seq_len=5, seed=7) - assert a == b - - -def test_iter_minibatches_keeps_pairs_together(): - """Minibatching requires an even batch size and preserves order.""" - dataset = build_synthetic_contrastive_dataset(4, seq_len=4, seed=1) - batches = iter_minibatches(dataset, batch_size=4) - assert [len(b) for b in batches] == [4, 4] - assert batches[0] + batches[1] == dataset - - with pytest.raises(ValueError): - iter_minibatches(dataset, batch_size=3) - - -def test_extract_loss_handles_both_backend_shapes(): - """extract_loss understands mock ([records, loss]) and real ({'loss': ...}).""" - - class _Resp: - def __init__(self, result): - self.result = result - - # Mock backend shape: [records, loss] - assert extract_loss(_Resp([[{'logprobs': [0.1]}], 0.42])) == pytest.approx(0.42) - # Real backend shape: dict with 'loss' - assert extract_loss(_Resp({'loss': 1.5, 'logits': None})) == pytest.approx(1.5) - - with pytest.raises(AssertionError): - extract_loss(_Resp({'no_loss_here': 1})) - - -def test_gpu_fixture_gating_env_flag(): - """The GPU gate reflects TWINKLE_TEST_GPU_E2E without side effects.""" - assert gpu_e2e_enabled() == (os.environ.get('TWINKLE_TEST_GPU_E2E', '0') == '1') - - -# ═══════════════════════════════════════════════════════════════════════════ -# Local CPU protocol/link validation against the mock backend. -# -# These cases boot the CPU-only mock server (module-scoped ``mock_embedding_model`` -# fixture — booted ONCE, never rebooted per example) and drive the full embedding -# verification call sequence over HTTP: -# -# set_processor('InputProcessor') -# -> set_loss('InfonceLoss', temperature=..., use_batch=True) -# -> add_metric('EmbeddingMetric', is_training=True) -# -> forward_backward(inputs=mb, task='embedding') -# -> calculate_metric(is_training=True) -# -# They assert each HTTP request/response hop is well-formed, that ``task='embedding'`` -# is transmitted through the protocol layer to the mock backend without protocol -# errors, and that the loss returned by -# ``forward_backward(task='embedding')`` is a finite number (math.isfinite: not NaN, -# not Inf). Here the mock loss layer validates the numeric-finiteness contract at the -# protocol-link level; real-model numeric semantics live in the GPU-gated cases. -# -# Run locally (no GPU) via: -# conda run -n twinkle pytest tests/server/test_embedding_e2e.py -v -k mock -# ═══════════════════════════════════════════════════════════════════════════ - -from hypothesis import HealthCheck, given, settings -from hypothesis import strategies as st - - -def test_mock_embedding_protocol_link_full_sequence(mock_embedding_model): - """Full ordered call sequence runs over HTTP against the mock backend. - - Validates that every hop of the verification call sequence - (set_processor -> set_loss -> add_metric -> forward_backward(task='embedding') - -> calculate_metric) completes without any protocol-layer exception, that - ``task='embedding'`` is accepted through the /twinkle/* protocol and reaches - the mock backend, and that every returned loss is finite. - """ - model = mock_embedding_model - - dataset = build_synthetic_contrastive_dataset(DEFAULT_NUM_PAIRS, seq_len=DEFAULT_SEQ_LEN, seed=0) - minibatches = iter_minibatches(dataset, batch_size=2 * DEFAULT_NUM_PAIRS) - - result = run_embedding_training(model, minibatches) - - losses = result['losses'] - # One loss per minibatch, and the sequence actually issued forward_backward calls. - assert len(losses) == len(minibatches) - assert losses, 'expected at least one forward_backward loss' - - # Every embedding forward_backward loss is a finite number. - for step, loss in enumerate(losses): - assert math.isfinite(loss), f'loss at step {step} is not finite: {loss!r}' - - # calculate_metric(is_training=True) returned a well-formed metric mapping. - metric = result['metric'] - assert isinstance(metric, dict), f'expected metric dict, got {type(metric)!r}' - - -def test_mock_embedding_task_embedding_is_passed_through(mock_embedding_model): - """``task='embedding'`` traverses the HTTP protocol layer to the mock backend. - - The mock backend intentionally ignores ``task`` semantically (it only exercises - dispatch), so we assert the protocol contract instead: a ``forward_backward`` - carrying ``task='embedding'`` is accepted end-to-end and returns a well-formed - response whose extracted loss is finite. This confirms the extra kwarg is - transported (via ``ForwardRequest.model_extra``) rather than rejected by the - /twinkle/forward_backward endpoint. - """ - model = mock_embedding_model - configure_embedding_adapter(model) - - dataset = build_synthetic_contrastive_dataset(2, seq_len=DEFAULT_SEQ_LEN, seed=42) - minibatch = dataset # single minibatch of 4 samples (2 pairs) - - response = model.forward_backward(inputs=minibatch, task='embedding') - - # Response is the ForwardBackwardResponse pydantic model with a ``result`` field. - assert hasattr(response, 'result'), f'malformed forward_backward response: {response!r}' - loss = extract_loss(response) - assert math.isfinite(loss), f'forward_backward(task="embedding") loss not finite: {loss!r}' - - -@settings( - max_examples=12, - deadline=None, - # ``mock_embedding_model`` is module-scoped (booted once and reused across all - # generated examples); suppress the health check so hypothesis does not object - # to the shared fixture and — critically — does NOT reboot the server per example. - suppress_health_check=[HealthCheck.function_scoped_fixture], -) -@given( - num_pairs=st.integers(min_value=1, max_value=4), - seq_len=st.integers(min_value=2, max_value=12), - seed=st.integers(min_value=0, max_value=10_000), -) -def test_mock_embedding_loss_is_finite_property(mock_embedding_model, num_pairs, seq_len, seed): - """forward_backward(task='embedding') loss is always finite. - - Varied contrastive dataset shapes (num_pairs, seq_len) and seeds are generated - within a SINGLE module-scoped mock server instance — the server is never - rebooted per example. For every generated minibatch, the loss returned by the - mock backend over the HTTP protocol link must be a finite number - (math.isfinite: not NaN, not Inf). - """ - model = mock_embedding_model - # Idempotent no-op configuration on the mock backend; keeps the case self-contained. - configure_embedding_adapter(model) - - dataset = build_synthetic_contrastive_dataset(num_pairs, seq_len=seq_len, seed=seed) - # 2*num_pairs is a positive even batch size, so every anchor/positive pair stays - # together inside a single minibatch (InfoNCE assumes paired samples per batch). - minibatches = iter_minibatches(dataset, batch_size=2 * num_pairs) - - for step, mb in enumerate(minibatches): - response = model.forward_backward(inputs=mb, task='embedding') - loss = extract_loss(response) - assert math.isfinite(loss), ( - f'non-finite loss {loss!r} at step {step} ' - f'(num_pairs={num_pairs}, seq_len={seq_len}, seed={seed})') - - -# ═══════════════════════════════════════════════════════════════════════════ -# Single-GPU real-model loss finiteness (GPU-gated). -# -# This case exercises the REAL transformers model backend through the Twinkle -# client HTTP path (module-scoped ``gpu_embedding_model`` fixture, gated behind -# TWINKLE_TEST_GPU_E2E=1 and an already-running GPU server; skipped otherwise). -# For a non-empty contrastive minibatch, the real-model -# ``forward_backward(task='embedding')`` loss must be a finite number -# (math.isfinite: not NaN, not Inf). Unlike the mock cases (which -# validate the numeric-finiteness contract only at the protocol-link level), this -# asserts finiteness against real InfoNCE numeric semantics on real model weights. -# -# Run on GPU via: -# TWINKLE_TEST_GPU_E2E=1 conda run -n twinkle pytest \ -# tests/server/test_embedding_e2e.py -v -k gpu -# ═══════════════════════════════════════════════════════════════════════════ - - -def test_gpu_embedding_real_model_loss_is_finite(gpu_embedding_model): - """Real-model forward_backward(task='embedding') loss is finite. - - Drives the embedding verification call sequence - (set_processor -> set_loss('InfonceLoss', ...) -> add_metric('EmbeddingMetric') - -> forward_backward(task='embedding') + clip_grad_and_step) against the REAL - transformers backend and asserts that every loss returned by - ``forward_backward(task='embedding')`` over a non-empty contrastive minibatch - is a finite number (math.isfinite: not NaN, not Inf). - - GPU-gated (TWINKLE_TEST_GPU_E2E=1 + running GPU server): skipped automatically - on machines without a GPU, so this asserts real numeric semantics only on GPU - CI while collecting/skipping cleanly locally. - """ - model = gpu_embedding_model - - dataset = build_synthetic_contrastive_dataset(DEFAULT_NUM_PAIRS, seq_len=DEFAULT_SEQ_LEN, seed=0) - minibatches = iter_minibatches(dataset, batch_size=2 * DEFAULT_NUM_PAIRS) - - # A non-empty minibatch is required to actually exercise the loss path. - assert minibatches, 'expected at least one minibatch' - assert all(mb for mb in minibatches), 'expected every minibatch to be non-empty' - - # Configure the embedding pipeline and attach an optimizer before training: - # run_embedding_training issues clip_grad_and_step, which requires a set optimizer. - configure_embedding_adapter(model, temperature=DEFAULT_TEMPERATURE) - model.set_optimizer('AdamW', lr=1e-4) - result = run_embedding_training(model, minibatches, configure=False) - - losses = result['losses'] - # One loss per minibatch, and the sequence actually issued forward_backward calls. - assert len(losses) == len(minibatches) - assert losses, 'expected at least one forward_backward loss' - - # Every real-model embedding forward_backward loss is finite. - for step, loss in enumerate(losses): - assert math.isfinite(loss), f'real-model loss at step {step} is not finite: {loss!r}' - - -# ═══════════════════════════════════════════════════════════════════════════ -# Single-GPU: TransformersEmbeddingPatch auto-rollback. -# -# GPU-gated (TWINKLE_TEST_GPU_E2E=1). Skipped automatically on this GPU-less dev -# machine, and on any environment where the variable is unset. -# -# TransformersEmbeddingPatch is applied -# *and rolled back automatically* inside a single forward call via -# ``_resolve_task_context(model, task='embedding')`` (src/twinkle/model/transformers/ -# transformers.py). While the patch is active, ``lm_head`` is swapped for identity -# and a forward hook replaces the model output with per-token hidden states -# (src/twinkle/patch/transformers_emb.py::_output_features_hook), i.e. the "logits" -# would carry the HIDDEN-STATE dimension. After the embedding forward_backward -# returns, the patch MUST be reverted, so the very next ``forward_only`` WITHOUT a -# ``task`` argument (defaults to 'causal_lm') must produce real language-model -# logits whose trailing dimension is the VOCABULARY size — never the leftover -# identity hidden states. -# -# This property depends on the real model's patch/unpatch semantics, which the -# CPU-only mock backend cannot exhibit, so it lives behind the GPU gate. -# -# Run on GPU (with a running GPU server) in the ``twinkle`` conda env: -# TWINKLE_TEST_GPU_E2E=1 conda run -n twinkle pytest tests/server/test_embedding_e2e.py -v -k gpu -# ═══════════════════════════════════════════════════════════════════════════ - - -def _innermost_dim(nested: Any) -> Optional[int]: - """Return the length of the innermost (last-axis) list of a nested list. - - ``forward_only`` returns logits that ``to_cpu_safe_output`` has converted from - a torch tensor of shape ``[B, T, V]`` (or a per-sample list of ``[T, V]``) into - plain nested Python lists. Descending to the innermost list of scalars yields - the trailing dimension ``V`` (the vocabulary size for real logits, or the hidden - size if the embedding patch had leaked). Returns ``None`` when the structure is - not a nested list (e.g. logits were suppressed to ``None``). - """ - cur = nested - while isinstance(cur, list) and cur and isinstance(cur[0], list): - cur = cur[0] - return len(cur) if isinstance(cur, list) else None - - -def _build_causal_sample(seq_len: int = DEFAULT_SEQ_LEN) -> Dict[str, Any]: - """Build a well-formed causal-LM feature (input_ids/attention_mask/labels aligned). - - Distinct from the embedding features (which carry a single scalar ``labels``): - here ``labels`` has the same length as ``input_ids`` so a causal ``forward_only`` - can compute logps without shape mismatch, and we can read back real vocab-dim - logits. Small, deterministic token ids keep it valid for any real tokenizer. - """ - input_ids = [(i % 7) + 1 for i in range(seq_len)] - return { - 'input_ids': list(input_ids), - 'attention_mask': [1] * seq_len, - 'labels': list(input_ids), - } - - -def test_gpu_embedding_patch_auto_rollback_property(gpu_embedding_model): - """TransformersEmbeddingPatch auto-rolls back after an embedding step. - - Flow (single real GPU adapter, via the ``gpu_embedding_model`` fixture): - - 1. Configure the adapter for embedding training (set_processor -> set_loss - 'InfonceLoss' -> add_metric 'EmbeddingMetric'). - 2. Establish a baseline: call ``forward_only(return_logits=True)`` WITHOUT a - ``task`` argument BEFORE any embedding step. Since the patch is never - applied outside a ``task='embedding'`` forward, these baseline logits are - genuine vocabulary-dimension logits — the ground-truth "vocab dim". - 3. Run one real ``forward_backward(inputs=mb, task='embedding')``. This applies - TransformersEmbeddingPatch for the duration of that forward and must revert - it on the way out. - 4. Immediately call ``forward_only(return_logits=True)`` again WITHOUT ``task``. - - Assertions: - * The post-embedding logits exist and are 3D nested lists (a leaked patch - would instead make ``forward_only`` fail — the feature hook returns only - ``{'features': ...}`` with no ``logits`` key — or yield the hidden-state - dimension). - * The post-embedding logits' trailing dimension equals the baseline vocabulary - dimension, proving ``lm_head``/the forward hook were restored (not left as - identity emitting hidden states). - """ - model = gpu_embedding_model - configure_embedding_adapter(model) - - causal_sample = _build_causal_sample(seq_len=DEFAULT_SEQ_LEN) - # Use two identical samples so the baseline/post forward_only can be split across - # a multi-DP model server (dp>=2); the assertions only inspect the logits' trailing - # vocab dimension, which is invariant to batch size. - causal_batch = [causal_sample, causal_sample] - - # (2) Baseline vocab-dim logits with the patch NEVER applied. - baseline_resp = model.forward_only(inputs=causal_batch, return_logits=True) - baseline_result = getattr(baseline_resp, 'result', baseline_resp) - assert isinstance(baseline_result, dict), f'malformed forward_only response: {baseline_result!r}' - baseline_logits = baseline_result.get('logits') - assert baseline_logits is not None, 'baseline forward_only returned no logits (return_logits was set)' - vocab_dim = _innermost_dim(baseline_logits) - assert isinstance(vocab_dim, int) and vocab_dim > 1, ( - f'baseline logits trailing dim is not a valid vocab size: {vocab_dim!r}') - - # (3) One real embedding step: applies + must auto-rollback the patch. - emb_dataset = build_synthetic_contrastive_dataset(2, seq_len=DEFAULT_SEQ_LEN, seed=0) - fwd_bwd = model.forward_backward(inputs=emb_dataset, task='embedding') - emb_loss = extract_loss(fwd_bwd) - assert math.isfinite(emb_loss), f'embedding forward_backward loss not finite: {emb_loss!r}' - - # (4) Post-embedding causal forward_only WITHOUT task: patch must be gone. - post_resp = model.forward_only(inputs=causal_batch, return_logits=True) - post_result = getattr(post_resp, 'result', post_resp) - assert isinstance(post_result, dict), ( - f'post-embedding forward_only returned malformed result (patch may have leaked): {post_result!r}') - post_logits = post_result.get('logits') - assert post_logits is not None, ( - 'post-embedding forward_only returned no logits — the embedding patch likely ' - 'did NOT roll back (feature hook suppresses the logits key)') - post_dim = _innermost_dim(post_logits) - - # Trailing dim is the vocab dim (identical to baseline), NOT the - # leftover identity hidden-state dim. - assert post_dim == vocab_dim, ( - f'TransformersEmbeddingPatch did NOT auto-rollback: post-embedding logits ' - f'trailing dim {post_dim!r} != baseline vocab dim {vocab_dim!r} — logits appear ' - f'to reuse identity hidden states from the embedding task') - - -# ═══════════════════════════════════════════════════════════════════════════ -# Single-GPU: bare-library parity + pos_sim upward trend (GPU-gated). -# -# These cases exercise the REAL transformers model backend and assert real -# numeric semantics that the mock backend cannot express, so they live behind -# the GPU gate (module-scoped ``gpu_embedding_model`` fixture — gated by -# TWINKLE_TEST_GPU_E2E=1 and a running GPU server; skipped automatically -# otherwise). Two properties are validated: -# -# * Twinkle client HTTP path vs. bare-library path parity: -# the same small synthetic contrastive dataset is trained for the same number -# of steps through (a) the Twinkle client HTTP path and -# (b) the bare-library training path used by -# ``cookbook/exp/embedding/train_embedding_full_ddp.py``. The two losses must -# stay within the SAME ORDER OF MAGNITUDE (digit-for-digit equality is NOT -# required). -# -# * pos_sim upward trend: after several real training steps -# on a dataset with a clear contrastive signal, ``calculate_metric(is_training= -# True)`` must report an increasing anchor-positive cosine similarity -# (``pos_sim``) relative to the untrained baseline. -# -# NOTE on the "bare-library path": the published script -# ``cookbook/exp/embedding/train_embedding_full_ddp.py`` orchestrates an 8-GPU -# online-compression pipeline (Ray, vLLM condenser, real datasets) that cannot be -# executed verbatim inside a single-GPU test. This module instead reproduces the -# script's ESSENTIAL bare-library embedding-training primitives IN-PROCESS with -# the very same core-library classes it uses — -# TransformersModel/MultiLoraTransformersModel -# + set_processor(InputProcessor) -# + set_loss(InfonceLoss, temperature=..., use_batch=True, hard_negatives=None) -# + set_optimizer('AdamW', ...) -# + add_metric(EmbeddingMetric, is_training=True) -# + loop: forward_backward(inputs=mb, task='embedding') + clip_grad_and_step(...) -# + calculate_metric(is_training=True) -# — running on the SAME synthetic dataset and the SAME number of steps as the HTTP -# path. A LoRA adapter is used (instead of full fine-tuning) purely to keep the -# in-process GPU memory footprint tractable next to the running server; the loss -# numerics being compared come from the identical InfoNCE + embedding-pooling code -# path exercised by the bare script. -# -# Run on GPU via: -# TWINKLE_TEST_GPU_E2E=1 conda run -n twinkle pytest \ -# tests/server/test_embedding_e2e.py -v -k gpu -# ═══════════════════════════════════════════════════════════════════════════ - -# Dataset / training shape for the parity + trend cases (kept small for GPU speed -# while still giving InfoNCE a real contrastive signal to learn from). -CMP_NUM_PAIRS = 6 # -> 12 samples -CMP_BATCH_SIZE = 4 # -> 3 minibatches == 3 training steps for BOTH paths -TREND_NUM_PAIRS = 6 # -> 12-sample fixed eval minibatch -TREND_STEPS = 20 # real training steps to expose the pos_sim trend -TREND_LR = 1e-3 # aggressive LR so the trend is visible within few steps - - -def _metric_dict(metric_response: Any) -> Dict[str, Any]: - """Normalize a calculate_metric result to a plain dict. - - Handles the client HTTP shape (``CalculateMetricResponse`` with a ``result`` - attribute) and the bare-library local shape (a plain dict returned directly). - """ - result = getattr(metric_response, 'result', metric_response) - assert isinstance(result, dict), f'expected metric dict, got {type(result)!r}: {result!r}' - return result - - -def _parse_metric_float(metric_response: Any, key: str) -> float: - """Parse a single float metric value (EmbeddingMetric formats values as strings).""" - result = _metric_dict(metric_response) - assert key in result, f'metric {key!r} missing from result: {result!r}' - return float(result[key]) - - -def _order_of_magnitude(value: float) -> int: - """Return floor(log10(|value|)); 0 is treated as magnitude 0.""" - if value == 0: - return 0 - return int(math.floor(math.log10(abs(value)))) - - -def _assert_same_order_of_magnitude(a: float, b: float, *, label: str) -> None: - """Assert two positive losses share the same order of magnitude. - - "Same order of magnitude" is interpreted as a bounded ratio in [0.1, 10] - (equivalently, their base-10 exponents differ by at most 1). Digit-for-digit - equality is intentionally NOT required. - """ - assert math.isfinite(a) and math.isfinite(b), f'[{label}] non-finite losses: {a!r}, {b!r}' - assert a > 0 and b > 0, f'[{label}] expected positive InfoNCE losses, got {a!r}, {b!r}' - ratio = a / b - assert 0.1 <= ratio <= 10.0, ( - f'[{label}] losses differ by more than one order of magnitude: ' - f'client={a:.6f}, bare={b:.6f}, ratio={ratio:.4f} ' - f'(oom client={_order_of_magnitude(a)}, bare={_order_of_magnitude(b)})') - log(f'[{label}] same order of magnitude: client={a:.6f}, bare={b:.6f}, ratio={ratio:.4f}') - - -def run_bare_library_embedding_training( - minibatches: List[List[Dict[str, Any]]], - *, - model_id: str = MODEL_ID, - temperature: float = DEFAULT_TEMPERATURE, - lr: float = 1e-4, - max_grad_norm: float = 1.0, - adapter_name: str = 'bare_emb_adapter', -) -> Dict[str, Any]: - """Run the bare-library embedding-training path IN-PROCESS on a single GPU. - - Reproduces the essential core-library primitives used by - ``cookbook/exp/embedding/train_embedding_full_ddp.py`` (see the module section - comment for why the full multi-GPU script cannot be executed verbatim): - - model = MultiLoraTransformersModel(model_id=...) - model.add_adapter_to_model(adapter, LoraConfig(target_modules='all-linear')) - model.set_processor(InputProcessor) - model.set_loss(InfonceLoss, temperature=..., use_batch=True, hard_negatives=None) - model.set_optimizer('AdamW', lr=...) - model.add_metric(EmbeddingMetric, is_training=True) - for mb in minibatches: - model.forward_backward(inputs=mb, task='embedding') - model.clip_grad_and_step(...) - model.calculate_metric(is_training=True) - - The model runs in ``local`` mode (the default when ``TWINKLE_MODE`` is unset), - so the ``@remote_function`` decorators call straight through and every method - requires an explicit ``adapter_name`` kwarg (unlike the client wrapper which - tracks it internally). - - Args: - minibatches: The SAME minibatches trained by the HTTP path (same order/steps). - model_id: Base model id (defaults to the shared real-model id). - temperature: InfoNCE temperature (matches the client path). - lr: AdamW learning rate. - max_grad_norm: Grad clipping threshold. - adapter_name: LoRA adapter name for the in-process model. - - Returns: - Dict with ``losses`` (list[float]) and ``metric`` (final metric dict). - """ - import gc - - from peft import LoraConfig - - from twinkle.loss import InfonceLoss - from twinkle.metric import EmbeddingMetric - from twinkle.model import MultiLoraTransformersModel - from twinkle.processor import InputProcessor - - model = None - try: - model = MultiLoraTransformersModel(model_id=model_id) - model.add_adapter_to_model(adapter_name, LoraConfig(target_modules='all-linear')) - model.set_processor(InputProcessor, adapter_name=adapter_name) - model.set_loss( - InfonceLoss, temperature=temperature, use_batch=True, hard_negatives=None, - adapter_name=adapter_name) - model.set_optimizer(optimizer_cls='AdamW', lr=lr, adapter_name=adapter_name) - model.add_metric(EmbeddingMetric, is_training=True, adapter_name=adapter_name) - - losses: List[float] = [] - for step, mb in enumerate(minibatches): - outputs = model.forward_backward(inputs=mb, task='embedding', adapter_name=adapter_name) - loss = extract_loss(outputs) - losses.append(loss) - model.clip_grad_and_step(max_grad_norm=max_grad_norm, adapter_name=adapter_name) - log(f'[bare embedding step {step}] loss={loss:.6f}') - - metric = model.calculate_metric(is_training=True, adapter_name=adapter_name) - return {'losses': losses, 'metric': _metric_dict(metric)} - finally: - # Free the in-process model's GPU memory so it does not linger next to the - # running server (the test process holds a full second copy of the weights). - try: - del model - except Exception: - pass - gc.collect() - try: - import torch - if torch.cuda.is_available(): - torch.cuda.empty_cache() - except Exception: - pass - - -def test_gpu_embedding_loss_matches_bare_library_order_of_magnitude(gpu_embedding_model): - """Client HTTP path and bare-library path losses match in magnitude. - - Trains the SAME small synthetic contrastive dataset for the SAME number of - steps through two independent paths: - - * the Twinkle client HTTP path, against the real GPU - server, and - * the bare-library training path reproduced in-process from - ``cookbook/exp/embedding/train_embedding_full_ddp.py``. - - Both paths use a fresh, untrained LoRA adapter, the same InfoNCE temperature, - the same AdamW learning rate, and the same minibatch sequence, so the observed - losses are directly comparable. The assertion requires only that the mean - losses stay within the SAME ORDER OF MAGNITUDE — digit-for-digit equality is - explicitly NOT required (LoRA weights are randomly initialized independently on - each path, so the two runs are not bit-identical). - - GPU-gated (TWINKLE_TEST_GPU_E2E=1 + running GPU server): skipped automatically - on machines without a GPU, so it collects/skips cleanly locally and asserts - real numeric semantics only on GPU CI. - """ - # ``gpu_embedding_model`` guarantees the GPU gate + a live server/session; use a - # dedicated fresh client adapter so the comparison is order-independent of other - # GPU cases that may have already trained the shared fixture adapter. - client_model = build_embedding_client_model(MODEL_ID, adapter_name='emb_cmp_adapter') - - dataset = build_synthetic_contrastive_dataset(CMP_NUM_PAIRS, seq_len=DEFAULT_SEQ_LEN, seed=0) - minibatches = iter_minibatches(dataset, batch_size=CMP_BATCH_SIZE) - assert minibatches, 'expected at least one minibatch' - assert all(mb for mb in minibatches), 'expected every minibatch to be non-empty' - - # ---- Client HTTP path -------------------------------------------------- # - # Configure explicitly (processor + loss + metric) and add an optimizer with a - # matching LR, then train WITHOUT re-configuring inside run_embedding_training. - configure_embedding_adapter(client_model, temperature=DEFAULT_TEMPERATURE) - client_model.set_optimizer('AdamW', lr=1e-4) - client_result = run_embedding_training(client_model, minibatches, configure=False) - client_losses = client_result['losses'] - - # ---- Bare-library path (in-process, same dataset + steps) -------------- # - bare_result = run_bare_library_embedding_training( - minibatches, temperature=DEFAULT_TEMPERATURE, lr=1e-4) - bare_losses = bare_result['losses'] - - # Both paths issued exactly one loss per (identical) minibatch. - assert len(client_losses) == len(minibatches), ( - f'client produced {len(client_losses)} losses for {len(minibatches)} minibatches') - assert len(bare_losses) == len(minibatches), ( - f'bare produced {len(bare_losses)} losses for {len(minibatches)} minibatches') - assert client_losses and bare_losses, 'both paths must produce at least one loss' - - for step, loss in enumerate(client_losses): - assert math.isfinite(loss), f'client loss at step {step} not finite: {loss!r}' - for step, loss in enumerate(bare_losses): - assert math.isfinite(loss), f'bare loss at step {step} not finite: {loss!r}' - - # Same order of magnitude on the mean loss across the identical step sequence. - client_mean = sum(client_losses) / len(client_losses) - bare_mean = sum(bare_losses) / len(bare_losses) - _assert_same_order_of_magnitude(client_mean, bare_mean, label='embedding-loss-parity') - - -def test_gpu_embedding_pos_sim_increases_after_training(gpu_embedding_model): - """pos_sim rises after several real training steps. - - Uses a fresh, untrained LoRA adapter on the real GPU model and repeatedly - trains on a fixed synthetic contrastive minibatch (anchors and their positives - share correlated tokens, so InfoNCE has a clear signal). ``pos_sim`` (the - anchor-positive cosine similarity reported by ``EmbeddingMetric``) is measured - on the SAME fixed minibatch before training and after each step, and must show - an upward trend: the average of the last few measurements exceeds the average - of the first few. - - This is a real numeric-semantics property that the mock backend cannot express - (mock embeddings carry no contrastive signal), hence the GPU gate. - - GPU-gated (TWINKLE_TEST_GPU_E2E=1 + running GPU server): skipped automatically - on machines without a GPU. - """ - # Fresh dedicated adapter -> clean untrained baseline, independent of other cases. - model = build_embedding_client_model(MODEL_ID, adapter_name='emb_trend_adapter') - configure_embedding_adapter(model, temperature=DEFAULT_TEMPERATURE) - model.set_optimizer('AdamW', lr=TREND_LR) - - dataset = build_synthetic_contrastive_dataset(TREND_NUM_PAIRS, seq_len=DEFAULT_SEQ_LEN, seed=0) - eval_mb = dataset # single fixed minibatch reused for every measurement - assert eval_mb, 'expected a non-empty eval minibatch' - # Sanity: the dataset must contain anchors (label==1) for pos_sim to be defined. - assert any(sample['labels'][0] == 1 for sample in eval_mb), 'eval minibatch has no anchors' - - pos_sims: List[float] = [] - for step in range(TREND_STEPS + 1): - # Measure pos_sim on the CURRENT (pre-step) weights, then take a train step. - model.forward_backward(inputs=eval_mb, task='embedding') - metric = model.calculate_metric(is_training=True) - pos_sims.append(_parse_metric_float(metric, 'pos_sim')) - model.clip_grad_and_step(max_grad_norm=1.0) - - # Every measurement must be finite and a valid cosine similarity in [-1, 1]. - for step, ps in enumerate(pos_sims): - assert math.isfinite(ps), f'pos_sim at measurement {step} not finite: {ps!r}' - assert -1.0001 <= ps <= 1.0001, f'pos_sim at measurement {step} out of range: {ps!r}' - - # Upward trend: average of the last 3 measurements exceeds the first 3 - # (mirrors the robust first-avg/last-avg comparison used elsewhere in the - # e2e helpers, tolerating per-step noise). - assert len(pos_sims) >= 6, f'need >=6 pos_sim measurements, got {len(pos_sims)}' - first_avg = sum(pos_sims[:3]) / 3 - last_avg = sum(pos_sims[-3:]) / 3 - assert last_avg > first_avg, ( - f'pos_sim did NOT increase after training: ' - f'first_3_avg={first_avg:.4f} >= last_3_avg={last_avg:.4f} ' - f'(full trace: {[round(p, 4) for p in pos_sims]})') - log(f'[embedding pos_sim trend] {first_avg:.4f} -> {last_avg:.4f} ' - f'(baseline={pos_sims[0]:.4f}, final={pos_sims[-1]:.4f})') diff --git a/tests/server/test_embedding_e2e_sp_cp.py b/tests/server/test_embedding_e2e_sp_cp.py deleted file mode 100644 index 06c6a316..00000000 --- a/tests/server/test_embedding_e2e_sp_cp.py +++ /dev/null @@ -1,449 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -"""Embedding training E2E validation under SP/CP distributed configs. - -This module reuses the synthetic contrastive dataset and the embedding -verification call sequence from ``tests/server/test_embedding_e2e.py`` and reruns -the embedding training flow under a MULTI-GPU device mesh that enables sequence -parallel (SP, ``ulysses_size > 1``) and/or context parallel (CP, ``cp_size > 1``), -then asserts that the resulting ``pos_sim`` (anchor-positive cosine similarity) -stays within a reasonable tolerance of the single-card baseline. - -The distributed device mesh is constructed with -``DeviceMesh.from_sizes(fsdp_size=..., ulysses_size=..., cp_size=...)`` (see -``src/twinkle/utils/device_mesh.py``) and driven across ranks with -``torch.multiprocessing.spawn``, mirroring the multi-rank harness in -``tests/transformers/test_sequence_parallel_and_cp.py``. Each rank builds a real -``MultiLoraTransformersModel`` bound to that mesh; SP is enabled automatically -inside the model whenever ``device_mesh.ulysses_size > 1``. - -============================================================================== -ENVIRONMENT CONSTRAINTS (MANDATORY — READ BEFORE RUNNING) -============================================================================== -1. MULTI-GPU required. These cases spawn ``world_size`` distributed workers, one - per GPU (SP-only needs >= 2 GPUs; combined CP+SP needs >= 4 GPUs). They are - gated behind ``TWINKLE_TEST_GPU_E2E=1`` AND an actual multi-GPU CUDA runtime, - and are SKIPPED automatically on a machine without enough GPUs (including this - GPU-less dev machine). - -2. ``twinkle`` conda env required. Run every case inside the ``twinkle`` conda - env, e.g.: - - TWINKLE_TEST_GPU_E2E=1 conda run -n twinkle \ - pytest tests/server/test_embedding_e2e_sp_cp.py -v - - On a GPU-less machine the file still collects/imports cleanly and skips: - - conda run -n twinkle pytest tests/server/test_embedding_e2e_sp_cp.py -v -============================================================================== -""" -from __future__ import annotations - -import json -import math -import os -import socket -import sys -import tempfile -from typing import Any, Dict, List, Optional - -# Ensure project root is importable for both pytest and direct execution. -_PROJECT_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..')) -if _PROJECT_ROOT not in sys.path: - sys.path.insert(0, _PROJECT_ROOT) - -import pytest - -# Reuse the single-card scaffolding (dataset, minibatching, model id, metric -# parsing, loss extraction, GPU gate) so the SP/CP flow exercises the SAME -# synthetic dataset and the SAME verification call sequence as the single-card -# path — the only difference is the distributed device mesh. -from tests.server.test_embedding_e2e import ( - DEFAULT_SEQ_LEN, - DEFAULT_TEMPERATURE, - MODEL_ID, - _parse_metric_float, - build_synthetic_contrastive_dataset, - extract_loss, - gpu_e2e_enabled, - iter_minibatches, -) -from tests.server.integration.e2e_helpers import log - -# ═══════════════════════════════════════════════════════════════════════════ -# Configuration -# ═══════════════════════════════════════════════════════════════════════════ - -# Dataset / training shape (kept small for GPU speed while still giving InfoNCE a -# real contrastive signal). Mirrors the pos_sim-trend shape used on the single card. -SP_CP_NUM_PAIRS = 6 # -> 12 samples -SP_CP_BATCH_SIZE = 4 # -> 3 minibatches == 3 training steps -SP_CP_TRAIN_STEPS = 8 # real training steps before measuring pos_sim -SP_CP_LR = 1e-3 # aggressive LR so weights actually move within few steps -SP_CP_SEED = 1234 # identical seed on baseline + distributed runs - -# ``pos_sim`` is a cosine similarity in [-1, 1]. SP/CP pooling is mathematically -# equivalent to the single-card path, so trained values should agree closely; -# a loose absolute tolerance accommodates bf16 + distributed reduction noise. -POS_SIM_ATOL = 5e-2 - -# Distributed configs exercised by the parametrized case. -# -# In the transformers backend, sequence/context parallel is driven by -# ``device_mesh.ulysses_size`` (see TransformersModel._decide_strategy: SP is -# enabled whenever ``ulysses_size > 1``). Ulysses degree is *derived* from the -# data ranks and therefore does NOT multiply the mesh world size — so -# ``from_sizes(fsdp_size=W, ulysses_size=U)`` yields exactly ``W`` ranks with a -# single ulysses group of size ``U``. Whether that group behaves as pure SP or a -# CP (ring-attention) mode is derived internally from ``ulysses_size`` vs. the -# model's attention-head count (see tests/transformers/test_sequence_parallel_and_cp.py), -# which is why both SP and CP are exercised through ``ulysses_size`` here. -# -# Each config is skipped unless its required GPU count is available. -DIST_CONFIGS = [ - # SP: 2 GPUs, ulysses_size=2 (sequence parallel enabled over 2 ranks). - {'name': 'sp2', 'world_size': 2, 'ulysses_size': 2, 'cp_size': None}, - # SP over 4 GPUs: ulysses_size=4 (larger sequence/context-parallel group). - {'name': 'sp4', 'world_size': 4, 'ulysses_size': 4, 'cp_size': None}, -] - - -def _cuda_device_count() -> int: - """Return the CUDA device count, or 0 when torch/CUDA is unavailable.""" - try: - import torch - if torch.cuda.is_available(): - return int(torch.cuda.device_count()) - except Exception: - pass - return 0 - - -def multi_gpu_e2e_enabled(min_devices: int = 2) -> bool: - """True only when GPU e2e is enabled AND enough CUDA devices are present.""" - return gpu_e2e_enabled() and _cuda_device_count() >= min_devices - - -def _find_free_port() -> int: - with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: - sock.bind(('127.0.0.1', 0)) - return int(sock.getsockname()[1]) - - -# ═══════════════════════════════════════════════════════════════════════════ -# Distributed worker: build an SP/CP device mesh, train, report pos_sim. -# -# Runs inside every spawned rank. Must be a module-level function so it is -# picklable by ``torch.multiprocessing.spawn`` (start method: spawn). -# ═══════════════════════════════════════════════════════════════════════════ - - -def _init_distributed(rank: int, world_size: int, port: int): - """Initialize the process group + pin this rank to its CUDA device.""" - import torch - import torch.distributed as dist - - os.environ['RANK'] = str(rank) - os.environ['WORLD_SIZE'] = str(world_size) - os.environ['LOCAL_RANK'] = str(rank) - os.environ['LOCAL_WORLD_SIZE'] = str(world_size) - os.environ['MASTER_ADDR'] = '127.0.0.1' - os.environ['MASTER_PORT'] = str(port) - - torch.cuda.set_device(rank) - device = torch.device(f'cuda:{rank}') - if not dist.is_initialized(): - dist.init_process_group( - backend='nccl', - rank=rank, - world_size=world_size, - init_method=f'tcp://127.0.0.1:{port}', - ) - return device - - -def _seed_everything(seed: int) -> None: - import torch - torch.manual_seed(seed) - if torch.cuda.is_available(): - torch.cuda.manual_seed_all(seed) - - -def _build_embedding_model_with_mesh( - world_size: int, - ulysses_size: Optional[int], - cp_size: Optional[int], - *, - adapter_name: str, - temperature: float, - lr: float, -): - """Build a real ``MultiLoraTransformersModel`` bound to an SP/CP device mesh. - - Uses the SAME bare-library primitives as the single-card path - (InputProcessor + InfonceLoss + EmbeddingMetric + AdamW), the only difference - being the distributed ``DeviceMesh`` passed to the constructor. SP is enabled - automatically by the transformers backend whenever ``ulysses_size > 1``. - """ - from peft import LoraConfig - - from twinkle.loss import InfonceLoss - from twinkle.metric import EmbeddingMetric - from twinkle.model import MultiLoraTransformersModel - from twinkle.processor import InputProcessor - from twinkle.utils import DeviceMesh - - # fsdp_size = world_size shards the model across all ranks; ulysses_size - # enables sequence/context parallel derived from those data ranks (it does - # not change the world size). cp_size, when provided, adds an explicit CP - # mesh dimension (and therefore multiplies the world size). - mesh = DeviceMesh.from_sizes( - fsdp_size=world_size, - dp_size=1, - ulysses_size=ulysses_size, - cp_size=cp_size, - device_type='cuda', - ) - - # find_unused_parameters=True is REQUIRED for embedding training: only the last - # hidden states contribute to the InfoNCE loss, so some parameters receive no - # gradient and DDP would otherwise abort the reduction (see Embedding训练.md). - model = MultiLoraTransformersModel( - model_id=MODEL_ID, device_mesh=mesh, - ddp_config={'find_unused_parameters': True}) - model.add_adapter_to_model(adapter_name, LoraConfig(target_modules='all-linear')) - model.set_processor(InputProcessor, adapter_name=adapter_name) - model.set_loss( - InfonceLoss, temperature=temperature, use_batch=True, hard_negatives=None, - adapter_name=adapter_name) - model.set_optimizer(optimizer_cls='AdamW', lr=lr, adapter_name=adapter_name) - model.add_metric(EmbeddingMetric, is_training=True, adapter_name=adapter_name) - return model, adapter_name - - -def _distributed_embedding_pos_sim_worker( - rank: int, - world_size: int, - port: int, - ulysses_size: Optional[int], - cp_size: Optional[int], - seed: int, - result_path: str, -) -> None: - """Spawned worker: run embedding training under the given mesh, emit pos_sim. - - Every rank trains identically (FSDP-sharded, SP/CP-parallel). Rank 0 writes - the final ``pos_sim`` (plus the per-step losses) to ``result_path`` as JSON so - the parent process can compare it against the single-card baseline. - """ - import gc - - import torch - import torch.distributed as dist - - device = _init_distributed(rank, world_size, port) - model = None - try: - _seed_everything(seed) - - # SAME synthetic dataset + minibatch sequence as the single-card path. - dataset = build_synthetic_contrastive_dataset( - SP_CP_NUM_PAIRS, seq_len=DEFAULT_SEQ_LEN, seed=0) - minibatches = iter_minibatches(dataset, batch_size=SP_CP_BATCH_SIZE) - - adapter_name = f'emb_sp_cp_adapter_u{ulysses_size}_c{cp_size}' - model, adapter_name = _build_embedding_model_with_mesh( - world_size, ulysses_size, cp_size, - adapter_name=adapter_name, temperature=DEFAULT_TEMPERATURE, lr=SP_CP_LR) - - losses: List[float] = [] - # Repeat the (short) minibatch sequence until SP_CP_TRAIN_STEPS steps run. - step = 0 - while step < SP_CP_TRAIN_STEPS: - for mb in minibatches: - if step >= SP_CP_TRAIN_STEPS: - break - outputs = model.forward_backward(inputs=mb, task='embedding', adapter_name=adapter_name) - losses.append(extract_loss(outputs)) - model.clip_grad_and_step(max_grad_norm=1.0, adapter_name=adapter_name) - step += 1 - - metric = model.calculate_metric(is_training=True, adapter_name=adapter_name) - pos_sim = _parse_metric_float(metric, 'pos_sim') - - if rank == 0: - with open(result_path, 'w') as f: - json.dump({'pos_sim': pos_sim, 'losses': losses}, f) - log(f'[sp_cp worker u={ulysses_size} c={cp_size}] pos_sim={pos_sim:.6f} ' - f'losses={[round(x, 4) for x in losses]}') - dist.barrier() - finally: - try: - del model - except Exception: - pass - gc.collect() - try: - if torch.cuda.is_available(): - torch.cuda.empty_cache() - except Exception: - pass - if dist.is_initialized(): - dist.destroy_process_group() - - -def _run_embedding_pos_sim( - world_size: int, - ulysses_size: Optional[int], - cp_size: Optional[int], - *, - seed: int = SP_CP_SEED, -) -> Dict[str, Any]: - """Spawn ``world_size`` workers, run embedding training, return rank-0 result. - - ``world_size == 1`` with ``ulysses_size in (None, 1)`` yields the single-card - baseline (SP/CP disabled). Larger world sizes exercise the SP/CP mesh. The - spawn path is identical for both so the ONLY variable is the device mesh. - """ - import torch.multiprocessing as mp - - port = _find_free_port() - fd, result_path = tempfile.mkstemp(prefix='emb_sp_cp_', suffix='.json') - os.close(fd) - try: - mp.spawn( - _distributed_embedding_pos_sim_worker, - args=(world_size, port, ulysses_size, cp_size, seed, result_path), - nprocs=world_size, - join=True, - ) - with open(result_path) as f: - return json.load(f) - finally: - try: - os.remove(result_path) - except OSError: - pass - - -# ═══════════════════════════════════════════════════════════════════════════ -# Single-card baseline (module-scoped: computed once, reused by every config). -# ═══════════════════════════════════════════════════════════════════════════ - - -@pytest.fixture(scope='module') -def single_card_pos_sim(): - """Baseline ``pos_sim`` from a single-card (SP/CP-disabled) embedding run. - - Gated behind the multi-GPU e2e requirement so it skips cleanly on machines - without enough GPUs. Runs through the same spawn harness with - ``world_size=1`` / ``ulysses_size=1`` so the baseline and the distributed - runs differ ONLY in the device mesh. - """ - if not multi_gpu_e2e_enabled(min_devices=2): - pytest.skip( - 'Requires TWINKLE_TEST_GPU_E2E=1 and >= 2 CUDA devices ' - '(multi-GPU SP/CP embedding e2e).') - - result = _run_embedding_pos_sim(world_size=1, ulysses_size=1, cp_size=None) - pos_sim = result['pos_sim'] - log(f'[sp_cp baseline] single-card pos_sim={pos_sim:.6f}') - return pos_sim - - -# ═══════════════════════════════════════════════════════════════════════════ -# SP/CP pos_sim consistency (multi-GPU, GPU-gated). -# ═══════════════════════════════════════════════════════════════════════════ - - -@pytest.mark.parametrize('config', DIST_CONFIGS, ids=[c['name'] for c in DIST_CONFIGS]) -def test_gpu_sp_cp_pos_sim_matches_single_card(config, single_card_pos_sim): - """SP/CP pos_sim agrees with the single-card baseline within tolerance. - - Reruns the embedding verification flow on the SAME synthetic - contrastive dataset under a distributed device mesh built with - ``DeviceMesh.from_sizes(fsdp_size=world_size, ulysses_size=..., cp_size=...)`` - (SP when ``ulysses_size > 1``, CP when ``cp_size > 1``), then asserts the - resulting ``pos_sim`` stays within ``POS_SIM_ATOL`` of the single-card value. - - Multi-GPU + GPU-gated (TWINKLE_TEST_GPU_E2E=1): each config is skipped unless - its required GPU count is available, so the file collects/skips cleanly on a - GPU-less machine and asserts real distributed numeric semantics only on - multi-GPU CI. - """ - world_size = int(config['world_size']) - if not multi_gpu_e2e_enabled(min_devices=world_size): - pytest.skip( - f"Config {config['name']!r} needs TWINKLE_TEST_GPU_E2E=1 and " - f'>= {world_size} CUDA devices (have {_cuda_device_count()}).') - - result = _run_embedding_pos_sim( - world_size=world_size, - ulysses_size=config['ulysses_size'], - cp_size=config['cp_size'], - ) - dist_pos_sim = result['pos_sim'] - - # pos_sim is a valid cosine similarity. - assert math.isfinite(dist_pos_sim), f"non-finite pos_sim: {dist_pos_sim!r}" - assert -1.0001 <= dist_pos_sim <= 1.0001, f'pos_sim out of range: {dist_pos_sim!r}' - - delta = abs(dist_pos_sim - single_card_pos_sim) - log(f"[sp_cp {config['name']}] dist_pos_sim={dist_pos_sim:.6f} " - f'baseline={single_card_pos_sim:.6f} delta={delta:.6f}') - assert delta <= POS_SIM_ATOL, ( - f"[{config['name']}] SP/CP pos_sim diverged from single-card baseline: " - f'dist={dist_pos_sim:.6f}, baseline={single_card_pos_sim:.6f}, ' - f'delta={delta:.6f} > atol={POS_SIM_ATOL}') - - -# ═══════════════════════════════════════════════════════════════════════════ -# Local (GPU-less) collection guards — verify the file imports/collects/skips -# cleanly and the mesh-config helpers behave, WITHOUT requiring any GPU. -# ═══════════════════════════════════════════════════════════════════════════ - - -def test_multi_gpu_gate_reflects_env_and_devices(): - """The multi-GPU gate is False unless BOTH the env flag and >=N GPUs exist.""" - expected = gpu_e2e_enabled() and _cuda_device_count() >= 2 - assert multi_gpu_e2e_enabled(min_devices=2) == expected - - -def test_dist_configs_are_well_formed(): - """Every declared SP/CP config enables SP and/or CP with a matching world size.""" - for cfg in DIST_CONFIGS: - assert cfg['world_size'] >= 2, cfg - sp = (cfg['ulysses_size'] or 1) > 1 - cp = (cfg['cp_size'] or 1) > 1 - assert sp or cp, f'config enables neither SP nor CP: {cfg}' - # ulysses/cp parallelism must fit within the world size. - assert (cfg['ulysses_size'] or 1) <= cfg['world_size'], cfg - assert (cfg['cp_size'] or 1) <= cfg['world_size'], cfg - - -def test_device_mesh_from_sizes_builds_sp_cp_mesh(): - """DeviceMesh.from_sizes wires ulysses_size/cp_size into a multi-rank mesh. - - Pure CPU construction (no distributed init, no GPU): confirms the exact - ``from_sizes(fsdp_size=..., ulysses_size=..., cp_size=...)`` usage the workers - rely on. ``ulysses_size`` is derived from the data ranks and does NOT change - the mesh world size (SP degree lives inside the existing ranks), whereas - ``cp_size`` adds an explicit mesh dimension that multiplies the world size. - """ - from twinkle.utils import DeviceMesh - - # SP mesh (2 ranks, ulysses_size=2): ulysses does not inflate world_size. - sp_mesh = DeviceMesh.from_sizes(fsdp_size=2, dp_size=1, ulysses_size=2, device_type='cuda') - assert sp_mesh.world_size == 2 - assert sp_mesh.ulysses_size == 2 - - # SP mesh over 4 ranks (ulysses_size=4): still exactly fsdp_size ranks. - sp4_mesh = DeviceMesh.from_sizes(fsdp_size=4, dp_size=1, ulysses_size=4, device_type='cuda') - assert sp4_mesh.world_size == 4 - assert sp4_mesh.ulysses_size == 4 - - # CP mesh dimension (cp_size=2) multiplies the world size and is recorded as - # an explicit 'cp' dim: fsdp(2) x dp(1) x cp(2) == 4 ranks. - cp_mesh = DeviceMesh.from_sizes(fsdp_size=2, dp_size=1, cp_size=2, device_type='cuda') - assert cp_mesh.world_size == 4 - assert cp_mesh.has_dim('cp') - assert cp_mesh.get_dim_size('cp') == 2 diff --git a/tests/twinkle_client/test_client_multi_turn_rollout.py b/tests/twinkle_client/test_client_multi_turn_rollout.py index e3598c3a..1ff69f5f 100644 --- a/tests/twinkle_client/test_client_multi_turn_rollout.py +++ b/tests/twinkle_client/test_client_multi_turn_rollout.py @@ -486,10 +486,7 @@ def sample(self, inputs, sampling_params=None, **kwargs): def test_missing_new_input_feature_raises_indexed_runtime_error(): - """new_input_feature=None -> RuntimeError naming batch AND trajectory index. - - _Requirements: 3.7_ - """ + """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) @@ -508,10 +505,7 @@ def test_missing_new_input_feature_raises_indexed_runtime_error(): def test_tool_calls_without_tool_manager_raises_value_error(): - """tool_calls produced but tool_manager missing -> ValueError. - - _Requirements: 3.10, 4.5_ - """ + """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( @@ -528,10 +522,7 @@ def test_tool_calls_without_tool_manager_raises_value_error(): 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. - - _Requirements: 3.10, 4.5_ - """ + """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. @@ -543,10 +534,7 @@ def test_tool_calls_without_tool_manager_via_per_call_kwarg_raises_value_error() def test_sampler_network_error_propagates_unchanged(): - """vLLMSampler.sample() network error propagates unchanged (not swallowed/wrapped). - - _Requirements: 3.9_ - """ + """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') @@ -564,10 +552,7 @@ def test_sampler_network_error_propagates_unchanged(): def test_dependencies_are_reused_not_reimplemented(): - """ClientMultiTurnRollout imports (does not copy) ToolManager & extend_with_bridge. - - _Requirements: 4.3, 4.4, 7.1_ - """ + """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 From 3848c6da0fb1f07353d4dbfd04d92315bf1f4190 Mon Sep 17 00:00:00 2001 From: Yunnglin Date: Tue, 21 Jul 2026 10:21:03 +0800 Subject: [PATCH 06/10] Consolidate client examples and use OpenEnv rollout --- INSTALL_CLIENT.ps1 | 2 +- INSTALL_CLIENT.sh | 2 +- README.md | 6 +- README_ZH.md | 6 +- .../client/tinker/{modelscope => }/dpo.py | 6 +- .../client/tinker/{self_host => }/lora.py | 12 +- cookbook/client/tinker/modelscope/sample.py | 69 ---- .../tinker/modelscope/self_cognition.py | 136 -------- .../tinker/modelscope/short_math_grpo.py | 325 ------------------ .../tinker/{self_host => }/multi_modal.py | 7 +- .../client/tinker/{self_host => }/sample.py | 6 +- .../tinker/{self_host => }/self_cognition.py | 6 +- cookbook/client/tinker/self_host/dpo.py | 207 ----------- .../tinker/{self_host => }/short_math_grpo.py | 6 +- .../tinker/{self_host => }/upload_to_hub.py | 8 +- .../client/twinkle/{self_host => }/dpo.py | 7 +- .../twinkle/{self_host => }/embedding.py | 9 +- cookbook/client/twinkle/modelscope/dpo.py | 204 ----------- .../twinkle/modelscope/self_cognition.py | 142 -------- .../twinkle/{modelscope => }/multi_modal.py | 7 +- cookbook/client/twinkle/multi_turn_rollout.py | 221 ++++++++++++ .../client/twinkle/{self_host => }/sample.py | 6 +- .../twinkle/{self_host => }/self_cognition.py | 6 +- .../client/twinkle/self_host/multi_modal.py | 168 --------- .../twinkle/self_host/multi_turn_rollout.py | 199 ----------- .../{self_host => }/short_math_grpo.py | 7 +- .../twinkle/{self_host => }/upload_to_hub.py | 8 +- docs/source_en/Usage Guide/Quick-Start.md | 2 +- .../Usage Guide/Server and Client/Overview.md | 56 +-- .../Server and Client/Twinkle-Client.md | 4 +- ...53\351\200\237\345\274\200\345\247\213.md" | 2 +- ...le\345\256\242\346\210\267\347\253\257.md" | 4 +- .../\346\246\202\350\277\260.md" | 56 +-- 33 files changed, 351 insertions(+), 1561 deletions(-) rename cookbook/client/tinker/{modelscope => }/dpo.py (97%) rename cookbook/client/tinker/{self_host => }/lora.py (95%) delete mode 100644 cookbook/client/tinker/modelscope/sample.py delete mode 100644 cookbook/client/tinker/modelscope/self_cognition.py delete mode 100644 cookbook/client/tinker/modelscope/short_math_grpo.py rename cookbook/client/tinker/{self_host => }/multi_modal.py (97%) rename cookbook/client/tinker/{self_host => }/sample.py (91%) rename cookbook/client/tinker/{self_host => }/self_cognition.py (96%) delete mode 100644 cookbook/client/tinker/self_host/dpo.py rename cookbook/client/tinker/{self_host => }/short_math_grpo.py (98%) rename cookbook/client/tinker/{self_host => }/upload_to_hub.py (92%) rename cookbook/client/twinkle/{self_host => }/dpo.py (96%) rename cookbook/client/twinkle/{self_host => }/embedding.py (94%) delete mode 100644 cookbook/client/twinkle/modelscope/dpo.py delete mode 100644 cookbook/client/twinkle/modelscope/self_cognition.py rename cookbook/client/twinkle/{modelscope => }/multi_modal.py (95%) create mode 100644 cookbook/client/twinkle/multi_turn_rollout.py rename cookbook/client/twinkle/{self_host => }/sample.py (93%) rename cookbook/client/twinkle/{self_host => }/self_cognition.py (96%) delete mode 100644 cookbook/client/twinkle/self_host/multi_modal.py delete mode 100644 cookbook/client/twinkle/self_host/multi_turn_rollout.py rename cookbook/client/twinkle/{self_host => }/short_math_grpo.py (97%) rename cookbook/client/twinkle/{self_host => }/upload_to_hub.py (92%) 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/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/self_host/embedding.py b/cookbook/client/twinkle/embedding.py similarity index 94% rename from cookbook/client/twinkle/self_host/embedding.py rename to cookbook/client/twinkle/embedding.py index e71c66b9..304a321f 100644 --- a/cookbook/client/twinkle/self_host/embedding.py +++ b/cookbook/client/twinkle/embedding.py @@ -27,6 +27,7 @@ dotenv.load_dotenv('.env') +import os from typing import Any, Dict, List from peft import LoraConfig @@ -38,7 +39,8 @@ logger = get_logger() # ========== 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}' ADAPTER_NAME = 'emb_adapter' TEMPERATURE = 0.07 LEARNING_RATE = 1e-4 @@ -82,7 +84,10 @@ def build_minibatch(tokenizer) -> List[Dict[str, Any]]: def train(): # Step 1: connect to the running Twinkle server. - init_twinkle_client(base_url='http://127.0.0.1:8000', api_key='EMPTY_TOKEN') + 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) 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/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/modelscope/multi_modal.py b/cookbook/client/twinkle/multi_modal.py similarity index 95% rename from cookbook/client/twinkle/modelscope/multi_modal.py rename to cookbook/client/twinkle/multi_modal.py index 106d85d4..3a5a45b6 100644 --- a/cookbook/client/twinkle/modelscope/multi_modal.py +++ b/cookbook/client/twinkle/multi_modal.py @@ -24,13 +24,14 @@ logger = get_logger() -base_model = 'Qwen/Qwen3.6-27B' -base_url = 'http://www.modelscope.cn/twinkle' +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/multi_modal.py b/cookbook/client/twinkle/self_host/multi_modal.py deleted file mode 100644 index f1908220..00000000 --- a/cookbook/client/twinkle/self_host/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.5-4B' -base_url = 'http://localhost:8000' - -# 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/self_host/multi_turn_rollout.py b/cookbook/client/twinkle/self_host/multi_turn_rollout.py deleted file mode 100644 index b4141323..00000000 --- a/cookbook/client/twinkle/self_host/multi_turn_rollout.py +++ /dev/null @@ -1,199 +0,0 @@ -# Twinkle Client - Multi-turn agentic rollout Example -# -# Client counterpart of the ray-local core-lib example -# ``cookbook/rl/multi_turn/multi_turn_grpo.py``. Instead of a Ray ``MultiTurnRollout`` -# + ``EnvPool``, it drives the SAME multi-turn "sample -> tool -> bridge -> sample" -# loop over HTTP with ``twinkle_client.rollout.ClientMultiTurnRollout`` and the -# client ``vLLMSampler``. -# -# Differences from the ray-local example (by design): -# * EnvPool is a Ray-only component and is NOT usable from the client, so this -# example uses a self-contained Python tool (a small calculator) wrapped in a -# ``ToolManager`` instead of an interactive environment. -# * ``ClientMultiTurnRollout`` samples with ``num_samples=1``; for a GRPO group we -# replicate each prompt ``NUM_GENERATIONS`` times as separate trajectories -# (exactly what the ray-local example does with ``BATCH_SIZE * NUM_GENERATIONS``). -# -# 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 dotenv - -dotenv.load_dotenv('.env') - -from typing import Any, Dict, List - -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.tools.base import Tool -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 - -logger = get_logger() - -# ========== Configuration ========== -MODEL_ID = 'ms://Qwen/Qwen3.5-4B' -ADAPTER_NAME = 'default' -NUM_GENERATIONS = 2 # GRPO group size (rollout runs num_samples=1 per trajectory) -MAX_NEW_TOKENS = 512 -MAX_TURNS = 4 -LEARNING_RATE = 1e-5 -MAX_STEPS = 3 - -SYSTEM_PROMPT = ('You are a careful assistant. When a calculation is needed, call the ' - '`calculator` tool with a single arithmetic ``expression`` (e.g. "12*7+3") ' - 'and then give the final numeric answer.') - -# A few arithmetic questions; each is expanded into NUM_GENERATIONS trajectories. -QUESTIONS: List[Dict[str, Any]] = [ - {'q': 'What is 12 * 7 + 3?', 'answer': 87}, - {'q': 'Compute (45 - 6) * 2.', 'answer': 78}, -] - - -# ========== A self-contained tool (replaces EnvPool) ========== -class CalculatorTool(Tool): - """Evaluate a single arithmetic ``expression`` string and return the result.""" - - _ALLOWED = set('0123456789+-*/(). ') - - def tool_info(self) -> Dict[str, Any]: - return { - 'type': 'function', - 'function': { - 'name': 'calculator', - 'description': 'Evaluate one arithmetic expression and return its numeric value.', - 'parameters': { - 'type': 'object', - 'properties': { - 'expression': { - 'type': 'string', - 'description': 'The arithmetic expression to evaluate, e.g. "12*7+3".', - } - }, - 'required': ['expression'], - }, - }, - } - - def __call__(self, tool_name: str, arguments: Dict[str, Any]) -> str: - expr = str(arguments.get('expression', '')).strip() - if not expr: - return 'Error: missing "expression".' - if not set(expr) <= self._ALLOWED: - return 'Error: expression may only contain digits, spaces and + - * / ( ) .' - try: - # Safe: the whitelist above forbids names, calls and attribute access. - return str(eval(expr, {'__builtins__': {}}, {})) # noqa: S307 - except Exception as e: # noqa - return f'Error: could not evaluate {expr!r}: {type(e).__name__}' - - -def build_trajectories(tool_schema: List[Dict[str, Any]]) -> List[Dict[str, Any]]: - """Expand each question into NUM_GENERATIONS identical trajectories (GRPO group).""" - trajectories: List[Dict[str, Any]] = [] - for item in QUESTIONS: - for _ in range(NUM_GENERATIONS): - trajectories.append({ - 'messages': [ - {'role': 'system', 'content': SYSTEM_PROMPT}, - {'role': 'user', 'content': item['q']}, - ], - 'tools': tool_schema, - }) - return trajectories - - -def compute_rewards(trajectories: List[Dict[str, Any]]) -> List[float]: - """+1 if the correct numeric answer appears in the final assistant message.""" - rewards: List[float] = [] - n_per_q = NUM_GENERATIONS - for i, traj in enumerate(trajectories): - expected = str(QUESTIONS[i // n_per_q]['answer']) - final = '' - for msg in reversed(traj.get('messages', [])): - if msg.get('role') == 'assistant': - final = str(msg.get('content') or '') - break - rewards.append(1.0 if expected in final else 0.0) - return rewards - - -def train(): - # Step 1: connect to the running Twinkle server. - init_twinkle_client(base_url='http://127.0.0.1:8000', api_key='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. Needs a LOCAL template for bridge-token stitching - # (rendering tool turns + the next generation prompt) 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([CalculatorTool()]) - tool_schema = tool_manager.tool_infos() - - rollout = ClientMultiTurnRollout( - sampler=sampler, - template=rollout_template, - tool_manager=tool_manager, - 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. Run the batched multi-turn rollout (sample -> tool -> bridge -> sample). - trajectories = build_trajectories(tool_schema) - rolled = rollout(trajectories, tool_manager=tool_manager) - - # 2. Read back per-trajectory logprobs (top-1 chosen-token logprob) and rewards. - 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) - rewards = compute_rewards(rolled) - - 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/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/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 2e63a297..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 @@ -295,4 +295,4 @@ Each returned trajectory has the following top-level fields appended to the orig - 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/self_host/multi_turn_rollout.py` for a full runnable example. +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/\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 cad5cbb9..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 后端的差异 @@ -295,4 +295,4 @@ for step in range(3): - 某条 trajectory 触发了工具调用，但没有提供 `tool_manager` → 抛 `ValueError`。构造时或按调用传入 `tool_manager` 即可。 - 采样器的网络 / 超时错误会**原样抛出**（不被吞掉），重试、退避请在你的循环外层处理。 -完整可运行示例见 `cookbook/client/twinkle/self_host/multi_turn_rollout.py`。 +完整可运行示例见 `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 ``` From 70b2f193b84fd6ebd13077ac008dd8fe3345ef67 Mon Sep 17 00:00:00 2001 From: Yunnglin Date: Tue, 21 Jul 2026 10:33:19 +0800 Subject: [PATCH 07/10] Configure Megatron entrypoint runtime defaults --- cookbook/client/server/megatron/entrypoint.sh | 4 ++++ 1 file changed, 4 insertions(+) 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" } From d7d6175848c32c10446d03f4b134e5369840708a Mon Sep 17 00:00:00 2001 From: Yunnglin Date: Tue, 21 Jul 2026 14:26:32 +0800 Subject: [PATCH 08/10] fix(processor): tolerate non-model keys in collate padding Use padding_map.get(key, 0) instead of hard index in _collate_macro_batch, matching the padding-free branch and the existing .get idiom elsewhere in this file. Fixes KeyError('logprobs') when rollout trajectories carrying metadata (e.g. logprobs) reach the non-padding-free collate path; such keys are dropped downstream by the to_transformers_dict whitelist anyway. --- src/twinkle/processor/base.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/twinkle/processor/base.py b/src/twinkle/processor/base.py index 27d91dc8..552f803b 100644 --- a/src/twinkle/processor/base.py +++ b/src/twinkle/processor/base.py @@ -738,7 +738,7 @@ 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: From 76eb55ecfcdbb4515b6e5862470d5196ef4fe05c Mon Sep 17 00:00:00 2001 From: Yunnglin Date: Tue, 21 Jul 2026 14:37:07 +0800 Subject: [PATCH 09/10] style(processor): wrap long _pad_sequence call to satisfy flake8 (E501) --- src/twinkle/processor/base.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/twinkle/processor/base.py b/src/twinkle/processor/base.py index 552f803b..956d38c9 100644 --- a/src/twinkle/processor/base.py +++ b/src/twinkle/processor/base.py @@ -738,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.get(key, 0), 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: From 85b2b1b3e8c95ad816d2bec442c94c9e7c02be4f Mon Sep 17 00:00:00 2001 From: Yunnglin Date: Tue, 21 Jul 2026 14:43:09 +0800 Subject: [PATCH 10/10] fix lint --- src/twinkle/processor/base.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/twinkle/processor/base.py b/src/twinkle/processor/base.py index 956d38c9..ae37dbdd 100644 --- a/src/twinkle/processor/base.py +++ b/src/twinkle/processor/base.py @@ -738,8 +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.get(key, 0), 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: