diff --git a/cookbook/client/async_rl/README.md b/cookbook/client/async_rl/README.md deleted file mode 100644 index 8441c3ddb..000000000 --- a/cookbook/client/async_rl/README.md +++ /dev/null @@ -1,156 +0,0 @@ -# Client-Orchestrated Async GRPO - -The client owns the Dataset, Reward, Advantage, rollout partitions, staleness, -training schedule, and policy publication. The server exposes shared -Multi-LoRA Model, vLLM Sampler, and TransferQueue DataPlane components without -owning the GRPO loop. - -This example deploys a dedicated Qwen3.5-4B GRPO server. It does not combine -different RL algorithms in one deployment. For the YAML-managed multi-tenant -runtime, see [`cookbook/rl`](../../rl/README.md). - -## Resources - -The server configuration uses two GPUs: - -| Component | GPUs | Purpose | -|---|---:|---| -| Multi-LoRA Model | 1 | GRPO forward, backward, optimizer step, and checkpoint save | -| Async vLLM Sampler | 1 | Shared continuous-batched rollout generation | - -The DataPlane keeps token tensors and sampled log-probabilities server-side. -The client reads decoded completions for Reward and Advantage, appends those -values to the same `DataRef`, and sends only references to the Model. - -## Quick start - -Install the server, client, and async-RL dependencies: - -```bash -pip install -e '.[async-rl,client,server]' -``` - -Terminal 1 — start the dedicated component server: - -```bash -export TWINKLE_LOCAL_MODEL_PATH=/absolute/path/to/Qwen3.5-4B -export CUDA_VISIBLE_DEVICES=0,1 -bash cookbook/client/async_rl/run_server.sh -``` - -The server script starts a local Ray cluster when needed, validates -`server_config.yaml`, and launches Ray Serve on port 8000. - -Terminal 2 — start one GRPO client: - -```bash -export TWINKLE_SERVER_URL=http://127.0.0.1:8000 -export TWINKLE_SERVER_TOKEN=EMPTY_TOKEN -export TWINKLE_TEMPLATE_MODEL_ID=/absolute/path/to/Qwen3.5-4B -export TWINKLE_DATASET_ID=/absolute/path/to/gsm8k -python cookbook/client/async_rl/client_orchestrated_grpo.py -``` - -`TWINKLE_TEMPLATE_MODEL_ID` is the tokenizer/template model path used in the -client process. The template class is fixed to `Qwen3_5Template`; a model ID -such as `Qwen/Qwen3.5-4B` is not a template class name. - -## Execution model - -Each DataLoader batch becomes a private client-side `_RolloutPartition` bound -to one immutable adapter checkpoint: - -```text -DataLoader - -> Rollout Worker - -> async vLLM Sampler - -> DataPlane DataRef - -> Advantage Worker - -> Trainer Worker - -> save and publish policy -``` - -The workers are independent asyncio tasks: - -- Prompt groups are submitted separately through - `sampler.asample_to_data_plane()`. -- A complete group is rewarded and made trainable without waiting for the rest - of its partition. -- The Trainer consumes `TRAIN_MINI_BATCH_SIZE / NUM_GENERATIONS` ready groups - at a time. -- Partitions train and publish in FIFO order. -- A policy version advances only after every group in its partition has - trained and the new checkpoint has been saved. - -`MAX_STALENESS` controls the number of live partitions. Setting it to zero -still permits rollout/training overlap inside one partition; values above zero -also permit cross-partition overlap. - -## Training metrics - -After every `clip_grad_and_step()`, the Trainer calls: - -```python -model.calculate_metric(is_training=True) -``` - -and prints the returned metrics to client stdout: - -```text -optimizer_step=1 grad_norm=0.42 learning_rate=2e-05 loss=0.031 -``` - -Partition publication is logged separately: - -```text -partition=0 policy=1 optimizer_step=8 staleness=0 -``` - -## Client settings - -| Environment variable | Default | Purpose | -|---|---|---| -| `TWINKLE_SERVER_URL` | `http://localhost:8000` | Server base URL | -| `TWINKLE_SERVER_TOKEN` | `EMPTY_TOKEN` | Request authentication token | -| `TWINKLE_TEMPLATE_MODEL_ID` | `ms://Qwen/Qwen3.5-4B` | Client tokenizer/template source | -| `TWINKLE_DATASET_ID` | `ms://modelscope/gsm8k` | GSM8K dataset source | -| `TWINKLE_ADAPTER_NAME` | `client-grpo` | LoRA adapter name | -| `TWINKLE_MAX_PARTITIONS` | `100` | Maximum DataLoader batches admitted | -| `TWINKLE_MAX_STALENESS` | `2` | Maximum extra live partitions | -| `TWINKLE_ROLLOUT_CONCURRENCY` | `8` | Concurrent prompt-group submissions | -| `TWINKLE_NUM_GENERATIONS` | `4` | Generations per prompt group | -| `TWINKLE_BATCH_SIZE` | `8` | Prompt groups per partition | -| `TWINKLE_TRAIN_MINI_BATCH_SIZE` | `8` | Samples per optimizer step | -| `TWINKLE_MICRO_BATCH_SIZE` | `4` | Model micro-batch size | -| `TWINKLE_MAX_TOKENS_PER_MICRO_BATCH` | `4096` | Dynamic batching token limit | - -The following must hold: - -```text -TWINKLE_TRAIN_MINI_BATCH_SIZE % TWINKLE_NUM_GENERATIONS == 0 -TWINKLE_BATCH_SIZE % ( - TWINKLE_TRAIN_MINI_BATCH_SIZE / TWINKLE_NUM_GENERATIONS -) == 0 -``` - -## Files - -| File | Role | -|---|---| -| `client_orchestrated_grpo.py` | Client-side async GRPO orchestration | -| `run_server.sh` | Start Ray and the dedicated component server | -| `server_config.yaml` | Qwen3.5-4B Model, Sampler, DataPlane, Gateway, and Processor deployment | - -## Troubleshooting - -- **Client receives HTTP 404 for the model** — use the provided server config; - its public route is fixed to `Qwen/Qwen3.5-4B`, matching the client. -- **The process tries to access ModelScope while offline** — set both - `TWINKLE_LOCAL_MODEL_PATH` on the server and - `TWINKLE_TEMPLATE_MODEL_ID`/`TWINKLE_DATASET_ID` on the client to local - paths. -- **No loss is printed** — confirm the Model request completed and look for an - `optimizer_step=...` line. Metrics are calculated after every optimizer - step, not after each rollout group. -- **Only one GPU is active** — verify Ray sees two GPUs and that Model and - Sampler placement groups were assigned to different devices. diff --git a/cookbook/client/async_rl/multi_turn_grpo.py b/cookbook/client/async_rl/multi_turn_grpo.py new file mode 100644 index 000000000..45b206ea8 --- /dev/null +++ b/cookbook/client/async_rl/multi_turn_grpo.py @@ -0,0 +1,468 @@ +"""Client-orchestrated multi-turn GRPO with independent AgentENV workspaces.""" +from __future__ import annotations + +import asyncio +import json +import math +import os +import posixpath +import uuid +from dataclasses import dataclass +from typing import Any + +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 AgentEnv, EnvTool +from twinkle_agentic.envs.base import DEFAULT_TOOLS, StepResult +from twinkle_agentic.tools.tool_manager import ToolManager +from twinkle_client.rollout import ClientMultiTurnRollout + +dotenv.load_dotenv('.env') +logger = get_logger() + +BASE_MODEL = os.environ.get('TWINKLE_MODEL_ID', 'Qwen/Qwen3.5-4B') +MODEL_ID = f'ms://{BASE_MODEL}' +# Tokenizer/template source; point it at a local model directory when the client runs offline. +TEMPLATE_MODEL_ID = os.environ.get('TWINKLE_TEMPLATE_MODEL_ID', MODEL_ID) +ADAPTER_NAME = 'client-multiturn-grpo' +AGENTENV_TEMPLATE = os.environ.get('AGENTENV_TEMPLATE', 'python:3.11') +NUM_GENERATIONS = 4 +ACTIVE_GROUPS = 2 +MAX_GROUPS = 4 +MAX_TURNS = 6 + + +@dataclass(frozen=True) +class CodeTask: + """One small code-repair task: the buggy file, a visible smoke test and hidden scoring cases.""" + filename: str + function: str + solution: str + smoke_test: str + prompt: str + hidden_cases: tuple + + +TASK = CodeTask( + filename='solution.py', + function='smallest_positive', + solution='''def smallest_positive(values): + positives = [value for value in values if value > 0] + return max(positives) if positives else None +''', + smoke_test='''from solution import smallest_positive +assert smallest_positive([3, 1, 2]) == 1 +assert smallest_positive([-1, 0]) is None +''', + prompt='''A small Python workspace contains solution.py and smoke_test.py. +Read the files, run the smoke test, fix the boundary-condition bug, and rerun the test. +Use the file and command tools, then give a short final answer. Do not edit the test to make it pass. +''', + hidden_cases=(([3, 1, 2], 1), ([-2, 4, 1], 1), ([0, -1], None), ([7], 7), ([5, 3, 9], 3)), +) + +_VERIFIER = '''import json + +call_id = __CALL_ID__ +cases = __CASES__ +payload = {"call_id": call_id, "completed": True, "passed": 0, "total": len(cases)} +try: + # Compile the current source directly: a stale __pycache__ from the smoke test must not be scored. + with open(__FILENAME__, encoding="utf-8") as handle: + namespace = {"__name__": "candidate_solution"} + exec(compile(handle.read(), __FILENAME__, "exec"), namespace) + function = namespace[__FUNCTION__] + for values, expected in cases: + try: + actual = function(values) + except Exception: + continue + if actual == expected: + payload["passed"] += 1 +except Exception as error: + payload["candidate_error"] = type(error).__name__ +print("TWINKLE_VERIFY=" + json.dumps(payload, sort_keys=True)) +''' + + +class _TaskInfrastructureError(RuntimeError): + + def __init__(self, stage: str, sandbox_id: str | None, detail: str = ''): + self.stage = stage + self.sandbox_id = sandbox_id + message = f'AgentENV infrastructure failure at {stage} (sandbox={sandbox_id})' + if detail: + message = f'{message}: {detail}' + super().__init__(message) + + +class _TaskAgentEnv(AgentEnv): + """Keep provider/verification failures distinct from candidate test failures.""" + + def __init__(self, task: CodeTask, *args, **kwargs): + super().__init__(*args, **kwargs) + self.task = task + self._task_infrastructure_error: _TaskInfrastructureError | None = None + + def _local_refusal(self, tool_name: str, arguments: dict[str, Any]) -> str | None: + if tool_name not in {'run_command', 'read_file', 'write_file'}: + return f'Error: unsupported task tool {tool_name!r}.' + if tool_name == 'run_command' and not arguments.get('command'): + return "Error: 'command' argument is required." + if tool_name in {'read_file', 'write_file'}: + path = arguments.get('path') + if not path: + return "Error: 'path' argument is required." + resolved = path if str(path).startswith('/') else posixpath.join(self.workspace, str(path)) + normalized = posixpath.normpath(resolved) + if normalized != self.workspace and not normalized.startswith(self.workspace.rstrip('/') + '/'): + return 'Error: task files must stay inside the workspace.' + return None + + def _record_infrastructure_error(self, stage: str, detail: str = '') -> None: + if self._task_infrastructure_error is None: + self._task_infrastructure_error = _TaskInfrastructureError(stage, self.sandbox_id, detail) + + def step(self, tool_name: str, arguments: dict[str, Any] | None = None) -> StepResult: + arguments = arguments or {} + refusal = self._local_refusal(tool_name, arguments) + if refusal is not None: + return StepResult(observation=refusal, done=False) + result = super().step(tool_name, arguments) + if result.info.get('error'): + self._record_infrastructure_error(f'tool:{tool_name}', str(result.info['error'])) + return result + + def ensure_idle(self) -> None: + super().ensure_idle() + if self._task_infrastructure_error is not None: + raise self._task_infrastructure_error + + def verify(self) -> float: + """Run the trusted checker in the same workspace and require its own result marker.""" + call_id = uuid.uuid4().hex + source = ( + _VERIFIER.replace('__CALL_ID__', repr(call_id)).replace('__CASES__', repr(list(self.task.hidden_cases))) + .replace('__FILENAME__', repr(self.task.filename)).replace('__FUNCTION__', repr(self.task.function))) + path = f'/tmp/twinkle_scripts/verify-{call_id}.py' + try: + self.sandbox.files.write(path, source) + except Exception as error: + self._record_infrastructure_error('verify_upload', str(error)) + raise self._task_infrastructure_error from error + exit_code, output = self._execute(f'python3 {path}', self.workspace, self._command_timeout) + try: + super().ensure_idle() + except Exception as error: + self._record_infrastructure_error('verify_command', str(error)) + raise self._task_infrastructure_error from error + if exit_code != 0: + self._record_infrastructure_error('verify_command', output) + raise self._task_infrastructure_error + marker = 'TWINKLE_VERIFY=' + rows = [row[len(marker):] for row in output.splitlines() if row.startswith(marker)] + try: + if len(rows) != 1: + raise ValueError(f'expected one result marker, got {len(rows)}') + result = json.loads(rows[0]) + if result.get('call_id') != call_id or result.get('completed') is not True: + raise ValueError('verification result does not match this call') + passed, total = result['passed'], result['total'] + if total != len(self.task.hidden_cases) or not isinstance(passed, int) or not 0 <= passed <= total: + raise ValueError('verification case statistics are invalid') + return passed / total + except (KeyError, TypeError, ValueError, json.JSONDecodeError) as error: + self._record_infrastructure_error('verify_result', str(error)) + raise self._task_infrastructure_error from error + + +@dataclass +class _Attempt: + env: _TaskAgentEnv + manager: ToolManager + trajectory: dict[str, Any] + + +def _checkpoint_path(saved: Any) -> str: + return str(saved['twinkle_path'] if isinstance(saved, dict) else saved.twinkle_path) + + +async def _await_operation(operation, *args, stop_event: asyncio.Event, on_result=None, **kwargs): + """Wait for one blocking operation without losing it when the caller is cancelled.""" + work = asyncio.create_task(asyncio.to_thread(operation, *args, **kwargs)) + task = asyncio.current_task() + cancel_error: asyncio.CancelledError | None = None + while True: + try: + result = await asyncio.shield(work) + except asyncio.CancelledError as error: + externally_cancelled = task is not None and task.cancelling() > 0 + if cancel_error is None and externally_cancelled: + cancel_error = error + stop_event.set() + if work.done(): + if cancel_error is None: + raise + try: + result = work.result() + if on_result is not None: + on_result(result) + except BaseException as drain_error: + raise cancel_error from drain_error + raise cancel_error + if cancel_error is None: + cancel_error = error + stop_event.set() + except BaseException as error: + if cancel_error is not None: + raise cancel_error from error + raise + else: + if on_result is not None: + on_result(result) + if cancel_error is not None: + raise cancel_error + return result + + +def _ensure_running(stop_event: asyncio.Event) -> None: + if stop_event.is_set(): + raise asyncio.CancelledError() + + +def _create_attempt(task: CodeTask, on_created) -> _Attempt: + env = _TaskAgentEnv(task, template=AGENTENV_TEMPLATE, api_url=os.environ.get('AGENTENV_API_URL')) + on_created(env) + env.reset() + for path, content in [(task.filename, task.solution), ('smoke_test.py', task.smoke_test)]: + result = env.step('write_file', {'path': path, 'content': content}) + env.ensure_idle() + if result.info.get('error'): + raise _TaskInfrastructureError('setup', env.sandbox_id, str(result.info['error'])) + manager = ToolManager(EnvTool.from_schemas(env, DEFAULT_TOOLS)) + trajectory = { + 'messages': [{'role': 'user', 'content': task.prompt}], + 'tools': manager.tool_infos(), + } + return _Attempt(env, manager, trajectory) + + +def _validate_training_rows(rolled: list[dict[str, Any]]) -> None: + for index, trajectory in enumerate(rolled): + if trajectory.get('stop_reason') in {'error', 'abort'} or trajectory.get('truncated'): + raise RuntimeError(f'attempt {index} is not a complete valid rollout') + tokens = list(trajectory.get('response_token_ids') or []) + logprobs = list(trajectory.get('rollout_logprobs') or []) + if not tokens or len(tokens) != len(logprobs) or not all(math.isfinite(value) for value in logprobs): + raise RuntimeError(f'attempt {index} has missing or misaligned rollout logprobs') + + +async def _rollout_group(rollout, policy_uri: str, environments: list, stop_event: asyncio.Event, + task: CodeTask = TASK): + attempts = [] + for _ in range(NUM_GENERATIONS): + _ensure_running(stop_event) + await _await_operation( + _create_attempt, + task, + environments.append, + stop_event=stop_event, + on_result=attempts.append, + ) + rolled = await rollout.acall( + [attempt.trajectory for attempt in attempts], + tool_manager=[attempt.manager for attempt in attempts], + adapter_name=ADAPTER_NAME, + adapter_uri=policy_uri, + ) + _ensure_running(stop_event) + _validate_training_rows(rolled) + rewards = [] + for attempt in attempts: + _ensure_running(stop_event) + rewards.append(await _await_operation(attempt.env.verify, stop_event=stop_event)) + return rolled, rewards + + +async def _train_group(model, rolled, rewards, advantage_fn, stop_event: asyncio.Event, on_metrics=None): + advantages = await _await_operation( + lambda: advantage_fn(rewards, num_generations=NUM_GENERATIONS, scale='group').tolist(), + stop_event=stop_event, + ) + if all(abs(value) < 1e-8 for value in advantages): + logger.info(f'all advantages are zero; skipping update rewards={rewards}') + return False + _ensure_running(stop_event) + await _await_operation( + model.forward_backward, + inputs=rolled, + advantages=advantages, + old_logps=[trajectory['rollout_logprobs'] for trajectory in rolled], + stop_event=stop_event, + ) + _ensure_running(stop_event) + await _await_operation(model.clip_grad_and_step, max_grad_norm=1.0, stop_event=stop_event) + metric_response = await _await_operation(model.calculate_metric, is_training=True, stop_event=stop_event) + metrics = metric_response['result'] if isinstance(metric_response, dict) else metric_response.result + if on_metrics is not None: + on_metrics(metrics) + logger.info(f'rewards={rewards} advantages={advantages} metrics={metrics}') + return True + + +async def _cleanup_resources(clients, environments, active_tasks, stop_event) -> list[str]: + issues = [] + stop_event.set() + for task in active_tasks: + task.cancel() + if active_tasks: + await asyncio.gather(*active_tasks, return_exceptions=True) + for env in environments: + try: + await _await_operation(env.ensure_idle, stop_event=stop_event) + except BaseException as error: + issues.append(f'quarantined environment {env.sandbox_id}: {error}') + logger.warning(issues[-1]) + continue + try: + await _await_operation(env.close, stop_event=stop_event) + except BaseException as error: + issues.append(f'could not confirm environment close for {env.sandbox_id}: {error}') + logger.warning(issues[-1]) + for client in clients: + try: + await _await_operation(client.close, stop_event=stop_event) + except BaseException as error: + issues.append(f'could not close client: {error}') + logger.warning(issues[-1]) + return issues + + +async def _drain_cleanup(clients, environments, active_tasks, stop_event, *, + propagate_cancel: bool = True) -> list[str]: + work = asyncio.create_task(_cleanup_resources(clients, environments, active_tasks, stop_event)) + cancel_error = None + while True: + try: + issues = await asyncio.shield(work) + break + except asyncio.CancelledError as error: + if cancel_error is None: + cancel_error = error + except BaseException as error: + logger.warning(f'cleanup failed: {error}') + issues = [f'cleanup failed: {error}'] + break + if propagate_cancel and cancel_error is not None: + raise cancel_error + return issues + + +async def train() -> None: + stop_event = asyncio.Event() + # Resources are registered as soon as they exist so cleanup can close them after any cancellation. + clients = [] + environments = [] + snapshots = [] + active_tasks = set() + primary_error = None + try: + client = await _await_operation( + init_twinkle_client, + base_url=os.environ.get('TWINKLE_SERVER_URL', 'http://localhost:8000'), + api_key=os.environ.get('TWINKLE_SERVER_TOKEN', 'EMPTY_TOKEN'), + stop_event=stop_event, + on_result=clients.append, + ) + model = await _await_operation(client.model, MODEL_ID, stop_event=stop_event) + sampler = await _await_operation(client.sampler, MODEL_ID, stop_event=stop_event) + await _await_operation( + model.add_adapter_to_model, + ADAPTER_NAME, + LoraConfig(target_modules='all-linear', r=8, lora_alpha=32), + stop_event=stop_event, + ) + await _await_operation(model.set_loss, 'GRPOLoss', epsilon=0.2, beta=0.0, stop_event=stop_event) + # Adam works with both the Transformers and Megatron training backends. + await _await_operation(model.set_optimizer, 'Adam', lr=2e-5, stop_event=stop_event) + await _await_operation(model.set_processor, 'InputProcessor', stop_event=stop_event) + await _await_operation( + model.set_template, 'Qwen3_5Template', model_id=TEMPLATE_MODEL_ID, stop_event=stop_event) + await _await_operation( + sampler.set_template, 'Qwen3_5Template', model_id=TEMPLATE_MODEL_ID, stop_event=stop_event) + template = await _await_operation( + Qwen3_5Template, + model_id=TEMPLATE_MODEL_ID, + max_length=8192, + enable_thinking=False, + stop_event=stop_event, + ) + template.truncation_strategy = 'delete' + rollout = ClientMultiTurnRollout( + sampler, + template, + sampling_params=SamplingParams(max_tokens=512, temperature=1.0, top_p=0.95, logprobs=1), + max_turns=MAX_TURNS, + ) + initial = await _await_operation( + model.save, + 'policy-0', + save_optimizer=False, + stop_event=stop_event, + on_result=snapshots.append, + ) + policy_uri = _checkpoint_path(initial) + advantage_fn = GRPOAdvantage() + next_group = 0 + version = 0 + semaphore = asyncio.Semaphore(ACTIVE_GROUPS) + + async def run_group(uri): + async with semaphore: + return await _rollout_group(rollout, uri, environments, stop_event) + + while next_group < MAX_GROUPS or active_tasks: + while next_group < MAX_GROUPS and len(active_tasks) < ACTIVE_GROUPS and not stop_event.is_set(): + task = asyncio.create_task(run_group(policy_uri)) + active_tasks.add(task) + next_group += 1 + if not active_tasks: + break + done, _ = await asyncio.wait(active_tasks, return_when=asyncio.FIRST_COMPLETED) + for task in done: + active_tasks.remove(task) + rolled, rewards = task.result() + _ensure_running(stop_event) + updated = await _train_group(model, rolled, rewards, advantage_fn, stop_event) + if updated: + _ensure_running(stop_event) + version += 1 + saved = await _await_operation( + model.save, + f'policy-{version}', + save_optimizer=False, + stop_event=stop_event, + on_result=snapshots.append, + ) + policy_uri = _checkpoint_path(saved) + except BaseException as error: + primary_error = error + stop_event.set() + raise + finally: + await _drain_cleanup( + clients, + environments, + active_tasks, + stop_event, + propagate_cancel=primary_error is None, + ) + + +if __name__ == '__main__': + asyncio.run(train()) diff --git a/cookbook/client/twinkle/multi_turn_rollout.py b/cookbook/client/twinkle/multi_turn_rollout.py index 942f572b6..c23cce3b6 100644 --- a/cookbook/client/twinkle/multi_turn_rollout.py +++ b/cookbook/client/twinkle/multi_turn_rollout.py @@ -16,18 +16,16 @@ # 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. +# Each rollout call names one immutable saved adapter URI. Training saves a new +# policy after each update, so an in-flight episode never switches weights. -import dotenv import os -from peft import LoraConfig from typing import Any, Dict, List, Tuple -from twinkle import get_logger -from twinkle import init_twinkle_client +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 @@ -185,24 +183,30 @@ def train(): ) advantage_fn = GRPOAdvantage() + policy_uri = model.save(name='policy-0', save_optimizer=False).twinkle_path 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) + rolled = rollout( + trajectories, + tool_manager=tool_managers, + adapter_name=ADAPTER_NAME, + adapter_uri=policy_uri, + ) rewards = extract_rewards(env_tools_list) finally: for env in environments: + env.ensure_idle() env.close() - # 2. Read back per-trajectory logprobs (top-1 chosen-token logprob). + # 2. Reuse the ledger's sampled-token logprobs without decoding top-k entries. 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_old_logps.append(traj['rollout_logprobs']) all_inputs.append(traj) avg_turns = sum(t.get('turns') or 0 for t in rolled) / len(rolled) @@ -217,6 +221,7 @@ def train(): # 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() + policy_uri = model.save(name=f'policy-{step + 1}', save_optimizer=False).twinkle_path logger.info(f'[Step {step}] {model.calculate_metric(is_training=True).result}') logger.info('Multi-turn rollout example finished.') diff --git a/cookbook/client/twinkle/prepare_seq_cls_model.py b/cookbook/client/twinkle/prepare_seq_cls_model.py index 7844c21bf..1f80dc966 100644 --- a/cookbook/client/twinkle/prepare_seq_cls_model.py +++ b/cookbook/client/twinkle/prepare_seq_cls_model.py @@ -35,10 +35,11 @@ # The classifier scores the last non-padding token; batched inputs need the pad id to find it. text_config.pad_token_id = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else tokenizer.eos_token_id -torch.manual_seed(42) -# The load report lists the vision weights as UNEXPECTED and score.weight as MISSING: -# both are expected, the vision tower is dropped and the new head is initialized here. +# The load report lists the vision weights as UNEXPECTED and score.weight as MISSING. model = AutoModelForSequenceClassification.from_pretrained(model_dir, config=text_config, torch_dtype='auto') +# Initialize the new head explicitly instead of relying on implicit missing-weight initialization. +torch.manual_seed(42) +torch.nn.init.normal_(model.score.weight, mean=0.0, std=text_config.initializer_range) model.save_pretrained(args.output_dir) tokenizer.save_pretrained(args.output_dir) print(f'Saved {type(model).__name__} with {args.num_labels} label(s) to {args.output_dir}') # noqa: T201 diff --git a/cookbook/client/twinkle/reranker.py b/cookbook/client/twinkle/reranker.py index d7fa7098d..3af500135 100644 --- a/cookbook/client/twinkle/reranker.py +++ b/cookbook/client/twinkle/reranker.py @@ -84,9 +84,10 @@ def train(): # Step 5: training loop. task='reranker' returns one score per pair. for step in range(MAX_STEPS): - output = model.forward_backward(inputs=batch, task='reranker') + model.forward_backward(inputs=batch, task='reranker') model.clip_grad_and_step(max_grad_norm=1.0) - logger.info(f'Step {step}: loss={output.result["loss"]:.4f}') + metric = model.calculate_metric(is_training=True) + logger.info(f'Step {step}: {metric.result}') # Step 6: relevant documents should now score higher, then save the adapter. scores = model.forward_only(inputs=batch, task='reranker', return_logits=True).result['logits'] diff --git a/cookbook/client/twinkle/seq_cls.py b/cookbook/client/twinkle/seq_cls.py index c01f17444..84b516ba2 100644 --- a/cookbook/client/twinkle/seq_cls.py +++ b/cookbook/client/twinkle/seq_cls.py @@ -90,9 +90,10 @@ def train(): # Step 5: training loop. task='seq_cls' returns one [num_labels] score row per text. for step in range(MAX_STEPS): - output = model.forward_backward(inputs=batch, task='seq_cls') + model.forward_backward(inputs=batch, task='seq_cls') model.clip_grad_and_step(max_grad_norm=1.0) - logger.info(f'Step {step}: loss={output.result["loss"]:.4f}') + metric = model.calculate_metric(is_training=True) + logger.info(f'Step {step}: {metric.result}') # Step 6: check the predictions, then save the adapter and optimizer state. logits = model.forward_only(inputs=batch, task='seq_cls', return_logits=True).result['logits'] diff --git a/cookbook/rl/async_rl/README.md b/cookbook/rl/async_rl/README.md index 4c8188269..1d778663e 100644 --- a/cookbook/rl/async_rl/README.md +++ b/cookbook/rl/async_rl/README.md @@ -6,8 +6,8 @@ training run as independent workers, while each tenant keeps its own dataset, reward, optimizer, scheduler, partitions, and policy versions. This directory contains the YAML-managed CLI workflow. For client-orchestrated -GRPO over HTTP, see -[`cookbook/client/async_rl`](../../client/async_rl/README.md). +GRPO over HTTP, see the +[Twinkle Client guide](../../../docs/source_en/Usage%20Guide/Server%20and%20Client/Twinkle-Client.md). ## Resources diff --git a/docs/source_en/Components/Sampler/TorchSampler.md b/docs/source_en/Components/Sampler/TorchSampler.md index 8b4282657..5dc3ec0b1 100644 --- a/docs/source_en/Components/Sampler/TorchSampler.md +++ b/docs/source_en/Components/Sampler/TorchSampler.md @@ -2,4 +2,4 @@ The current release does not provide `TorchSampler`. The former `from twinkle.sampler import TorchSampler` example is unavailable, and `sampler_type: torch` is rejected by server configuration validation. -Use [vLLMSampler](vLLMSampler.md) for sampling. For server-side async RL, use `sampler_type: vllm_async`; see the [client-orchestrated async RL example](https://github.com/modelscope/twinkle/blob/main/cookbook/client/async_rl/README.md). +Use [vLLMSampler](vLLMSampler.md) for sampling. For server-side async RL, use `sampler_type: vllm_async`; see the [Client-side Async GRPO](../../Usage%20Guide/Server%20and%20Client/Twinkle-Client.md#client-async-grpo). 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 e7539e9a9..575467026 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 @@ -207,96 +207,73 @@ 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) +(client-async-grpo)= +## Client-side Async GRPO -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. +The client can train while sampling continues. After [installing Twinkle](../Installation.md), choose an example and follow the steps below: -### Dependencies and Constraints +| Task | Example | +|---|---| +| Single-turn math questions (GSM8K) | [client_orchestrated_grpo.py](https://github.com/modelscope/twinkle/blob/main/cookbook/client/async_rl/client_orchestrated_grpo.py) | +| Multi-turn code repair with tools (text only) | [multi_turn_grpo.py](https://github.com/modelscope/twinkle/blob/main/cookbook/client/async_rl/multi_turn_grpo.py) | -- **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. -- **Safe text-only trajectories**: after every turn, the client verifies that state has no media and `feature.input_ids == engine_input_ids`. The template may support multimodality; only trajectories carrying media or divergent dual encodings raise `NotImplementedError` until the server exposes a `state_id`. +### 1. Start the Server -### Minimal Example +The example uses Qwen3.5-4B with Transformers for training and vLLM for sampling, requiring two GPUs. Prepare a local model directory, then run these commands from the repository root: -```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.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) - -client = init_twinkle_client(base_url='http://127.0.0.1:8000', api_key='EMPTY_TOKEN') - -# Training model (GRPO): bind the current transport, session, and authentication context -model = client.model(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() +```bash +pip install -e '.[async-rl,client,server]' +export TWINKLE_LOCAL_MODEL_PATH=/absolute/path/to/Qwen3.5-4B +export CUDA_VISIBLE_DEVICES=0,1 +bash cookbook/client/async_rl/run_server.sh ``` -### Output Fields +The server listens on `http://127.0.0.1:8000` by default. Skip this step if you already have a suitable server; see [server_config.yaml](https://github.com/modelscope/twinkle/blob/main/cookbook/client/async_rl/server_config.yaml) for its configuration. -Each returned trajectory has the following top-level fields appended to the original dict: +### 2. Run the Math Example -| 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 | +Keep the server running and use another terminal, also at the repository root: -### Common Errors +```bash +export TWINKLE_SERVER_URL=http://127.0.0.1:8000 +export TWINKLE_SERVER_TOKEN=EMPTY_TOKEN +export TWINKLE_TEMPLATE_MODEL_ID=/absolute/path/to/Qwen3.5-4B +python cookbook/client/async_rl/client_orchestrated_grpo.py +``` + +The example downloads GSM8K by default. For offline use, also set `TWINKLE_DATASET_ID=/absolute/path/to/gsm8k`. When connecting to an existing server, use its address and authentication token. The template directory must match the server's model. + +### 3. Run the Code Repair Example (Optional) + +The multi-turn example requires an accessible AgentENV service and a working Python environment template. Keep the client settings above and run: + +```bash +pip install 'e2b>=2.7' +export AGENTENV_API_URL=http://your-agentenv-gateway:8080 +export AGENTENV_TEMPLATE=your-python-template +python cookbook/client/async_rl/multi_turn_grpo.py +``` -- 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. +The model reads files, runs tests and repairs code, then trains on the test results. Each attempt uses an independent workspace; the example handles initialization, scoring and cleanup. For environment setup, see [Agentic RL Deployment and Training](../Agentic-RL-Deployment-and-Training.md). + +### Use Your Own Task + +- Math tasks: change the dataset and reward function in the script. +- Code tasks: edit `TASK` at the top of `multi_turn_grpo.py`, including the candidate file, tests and prompt. +- Start with the default training settings; edit the constants at the top of the script when needed. Training logs show rewards and metrics, and LoRA checkpoints are saved. + +## Trainable Multi-turn Rollout (ClientMultiTurnRollout) + +For a custom multi-turn training loop, use `twinkle_client.rollout.ClientMultiTurnRollout`. It supports text-only tool use and returns complete trajectories ready for training. + +Inside an async training loop where the rollout, trajectories, tools and policy checkpoint are already initialized: + +```python +rolled = await rollout.acall( + trajectories, + tool_manager=tool_managers, + adapter_uri=policy_uri, +) +``` -See `cookbook/client/twinkle/multi_turn_rollout.py` for a full runnable example. +For synchronous use, call `rollout(...)` with the same arguments and result format. Give each trajectory the `ToolManager` for its independent environment; rollout does not close environments you pass in. See [multi_turn_rollout.py](https://github.com/modelscope/twinkle/blob/main/cookbook/client/twinkle/multi_turn_rollout.py) for the full synchronous example. 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 daed53ffe..cab5c164d 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" @@ -206,96 +206,73 @@ model.set_lr_scheduler('default', lr_decay_steps=1000, max_lr=1e-4) 其余数据处理、训练循环、检查点保存等代码完全相同。 -## 可训练多轮 Rollout(ClientMultiTurnRollout) +(client-async-grpo)= +## 客户端异步 GRPO -前面的示例都是单轮训练。如果你要做**带工具调用的多轮 agentic RL**(如 GRPO),并且需要产出可直接用于训练的 token 级对齐信息,可使用 `twinkle_client.rollout.ClientMultiTurnRollout`。它在客户端侧驱动 “采样 → 调用工具 → 拼接上下文 → 再采样” 的多轮循环,每轮采样走 HTTP(`/twinkle/sample`),每条 trajectory 产出带 `logprobs` 的可训练结果,可直接用于 GRPO 等 RL 训练。 +客户端可以在采样的同时进行训练。完成 [Twinkle 安装](../安装.md) 后,选择一个示例,按下面的步骤启动即可: -### 依赖与约束 +| 场景 | 示例 | +|---|---| +| 单轮数学问答(GSM8K) | [client_orchestrated_grpo.py](https://github.com/modelscope/twinkle/blob/main/cookbook/client/async_rl/client_orchestrated_grpo.py) | +| 带工具调用的多轮代码修复(纯文本) | [multi_turn_grpo.py](https://github.com/modelscope/twinkle/blob/main/cookbook/client/async_rl/multi_turn_grpo.py) | -- **本地 Template**:bridge token 拼接(渲染工具轮 + 下一轮生成提示)需要在客户端本地持有一个 `Template` 实例。 -- **vLLMSampler**:指向服务端 Sampler 服务的客户端采样器。 -- **ToolManager**(可选):注册你的工具;若某条 trajectory 产生了 tool_calls 但未提供 tool_manager,会在派发时抛 `ValueError`。 -- **`num_samples=1`**:当前每条 trajectory 只采样一次。做 GRPO group 时,把同一个 prompt 复制成 `NUM_GENERATIONS` 条独立 trajectory 即可。 -- **仅限可证明安全的纯文本轨迹**:客户端会逐轮验证 state 不含媒体且 `feature.input_ids == engine_input_ids`。模板可以具备多模态能力;只有当前轨迹实际携带媒体或两种编码不一致时才会抛 `NotImplementedError`,等待服务端 `state_id`。 +### 1. 启动服务端 -### 最小示例 +示例使用 Qwen3.5-4B,训练后端为 Transformers,采样后端为 vLLM,需要两张 GPU。准备好本地模型目录后,在仓库根目录执行: -```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.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) - -client = init_twinkle_client(base_url='http://127.0.0.1:8000', api_key='EMPTY_TOKEN') - -# 训练模型(GRPO):通过当前 client 显式绑定 transport、会话与认证上下文 -model = client.model(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() +```bash +pip install -e '.[async-rl,client,server]' +export TWINKLE_LOCAL_MODEL_PATH=/absolute/path/to/Qwen3.5-4B +export CUDA_VISIBLE_DEVICES=0,1 +bash cookbook/client/async_rl/run_server.sh ``` -### 输出字段 +服务默认监听 `http://127.0.0.1:8000`。如果已有可用的服务端,可以跳过启动步骤;服务配置见 [server_config.yaml](https://github.com/modelscope/twinkle/blob/main/cookbook/client/async_rl/server_config.yaml)。 -每条返回的 trajectory 在原 dict 基础上追加以下顶层字段: +### 2. 运行单轮数学任务 -| 字段 | 含义 | -|:----|:-----| -| `messages` | 完整多轮对话(含 assistant 的 tool_calls 与 tool 响应轮) | -| `logprobs` | 各可训练 token 的 top-1 logprob;本轮未采样 logprobs 时为 `None` | -| `turns` | 实际经历的轮数(`<= max_turns`) | -| `stop_reason` | `'stop'` / `'length'` / `'max_turns'` 之一 | -| `truncated` | 是否因 `max_turns` 或长度上限被截断 | +保持服务端运行,在另一个终端中从仓库根目录执行: -### 常见错误 +```bash +export TWINKLE_SERVER_URL=http://127.0.0.1:8000 +export TWINKLE_SERVER_TOKEN=EMPTY_TOKEN +export TWINKLE_TEMPLATE_MODEL_ID=/absolute/path/to/Qwen3.5-4B +python cookbook/client/async_rl/client_orchestrated_grpo.py +``` + +示例默认下载 GSM8K 数据集。离线运行时,额外设置 `TWINKLE_DATASET_ID=/absolute/path/to/gsm8k`;连接已有服务端时,将地址和令牌换成对应配置。模板目录应与服务端模型匹配。 + +### 3. 运行多轮代码任务(可选) + +多轮示例需要已有的 AgentENV 服务和可用的 Python 环境模板。沿用上面的客户端配置,执行: + +```bash +pip install 'e2b>=2.7' +export AGENTENV_API_URL=http://your-agentenv-gateway:8080 +export AGENTENV_TEMPLATE=your-python-template +python cookbook/client/async_rl/multi_turn_grpo.py +``` -- 某条 trajectory 触发了工具调用,但没有提供 `tool_manager` → 抛 `ValueError`。构造时或按调用传入 `tool_manager` 即可。 -- 采样器的网络 / 超时错误会**原样抛出**(不被吞掉),重试、退避请在你的循环外层处理。 +示例会让模型读取文件、运行测试并修复代码,再根据测试结果训练。每次尝试使用独立的工作区,初始化、评分和关闭环境都已在示例中处理;环境准备见 [Agentic RL 部署与训练](../Agentic%20RL部署与训练.md)。 + +### 修改为自己的任务 + +- 数学任务:修改脚本中的数据集和奖励函数。 +- 代码任务:修改 `multi_turn_grpo.py` 顶部的 `TASK`,包括待修复文件、测试和提示词。 +- 先使用示例默认训练参数;需要调整时,再修改脚本顶部的配置。训练日志会输出奖励和指标,并保存 LoRA 检查点。 + +## 可训练多轮 Rollout(ClientMultiTurnRollout) + +需要自行编写多轮训练循环时,可使用 `twinkle_client.rollout.ClientMultiTurnRollout`。它支持纯文本工具调用,返回可直接用于训练的完整轨迹。 + +在已创建 rollout、轨迹、工具和策略检查点的异步训练循环中: + +```python +rolled = await rollout.acall( + trajectories, + tool_manager=tool_managers, + adapter_uri=policy_uri, +) +``` -完整可运行示例见 `cookbook/client/twinkle/multi_turn_rollout.py`。 +同步调用使用 `rollout(...)`,参数和结果相同。每条轨迹使用独立环境对应的 `ToolManager`;rollout 不会关闭你传入的环境。完整同步示例见 [multi_turn_rollout.py](https://github.com/modelscope/twinkle/blob/main/cookbook/client/twinkle/multi_turn_rollout.py)。 diff --git "a/docs/source_zh/\347\273\204\344\273\266/\351\207\207\346\240\267\345\231\250/TorchSampler.md" "b/docs/source_zh/\347\273\204\344\273\266/\351\207\207\346\240\267\345\231\250/TorchSampler.md" index e77efd8b8..e511feebb 100644 --- "a/docs/source_zh/\347\273\204\344\273\266/\351\207\207\346\240\267\345\231\250/TorchSampler.md" +++ "b/docs/source_zh/\347\273\204\344\273\266/\351\207\207\346\240\267\345\231\250/TorchSampler.md" @@ -2,4 +2,4 @@ 当前版本不提供 `TorchSampler`。旧的 `from twinkle.sampler import TorchSampler` 示例已不可用,服务端配置校验也会拒绝 `sampler_type: torch`。 -采样请使用 [vLLMSampler](vLLMSampler.md)。服务端异步 RL 使用 `sampler_type: vllm_async`,参见 [客户端编排异步 RL 示例](https://github.com/modelscope/twinkle/blob/main/cookbook/client/async_rl/README.md)。 +采样请使用 [vLLMSampler](vLLMSampler.md)。服务端异步 RL 使用 `sampler_type: vllm_async`,参见 [客户端异步 GRPO](../../使用指引/服务端和客户端/Twinkle客户端.md#client-async-grpo)。 diff --git a/src/twinkle/model/megatron/megatron.py b/src/twinkle/model/megatron/megatron.py index b69dfdb63..889e853cb 100644 --- a/src/twinkle/model/megatron/megatron.py +++ b/src/twinkle/model/megatron/megatron.py @@ -1189,7 +1189,8 @@ def _get_trainable_parameters(self, adapter_name: str = _default_adapter_name) - is_default = adapter_name == _default_adapter_name # Same pattern family as the transformers backend: LoRA deltas plus the fully trained # modules_to_save copies -- missing it here would freeze those copies out of the optimizer. - pattern = re.compile(rf'\.(?:lora_\w+|modules_to_save)\.{re.escape(adapter_name)}') + # The trailing boundary keeps lora_2 from also selecting lora_20..lora_29. + pattern = re.compile(rf'\.(?:lora_\w+|modules_to_save)\.{re.escape(adapter_name)}(?:\.|$)') params = {} model = self.strategy.unwrap_model(self.model) diff --git a/src/twinkle/model/transformers/transformers.py b/src/twinkle/model/transformers/transformers.py index 3b527678d..aa0b82743 100644 --- a/src/twinkle/model/transformers/transformers.py +++ b/src/twinkle/model/transformers/transformers.py @@ -1934,8 +1934,9 @@ def set_optimizer(self, optimizer_cls: Union[Type[Optimizer], str, Optimizer], * def _get_trainable_parameters(self, adapter_name=_default_adapter_name): is_default = adapter_name == _default_adapter_name # Previously the pattern did not match nn.Parameter() weights, causing EP LoRA - # parameters to be missed by the optimizer. - pattern = re.compile(rf'\.(?:lora_\w+|modules_to_save)\.{re.escape(adapter_name)}') + # parameters to be missed by the optimizer. The trailing boundary keeps lora_2 from + # also selecting lora_20..lora_29 when many adapter slots are resident. + pattern = re.compile(rf'\.(?:lora_\w+|modules_to_save)\.{re.escape(adapter_name)}(?:\.|$)') params = {} model = self.strategy.unwrap_model(self.model) for name, param in model.named_parameters(): diff --git a/src/twinkle/server/model/app.py b/src/twinkle/server/model/app.py index 908dddbd8..80c4d50ee 100644 --- a/src/twinkle/server/model/app.py +++ b/src/twinkle/server/model/app.py @@ -101,6 +101,10 @@ def _make_model_class(base: type) -> type: def _make_transformers_model(kw: dict[str, Any]) -> Any: train_mode = kw.pop('train_mode', 'lora') + # The head comes from the persisted config checked by ``_validate_model_task``; HF model + # constructors do not accept these deployment keys, unlike Megatron's bridge config. + kw.pop('task_type', None) + kw.pop('num_labels', None) # ``model_cls`` arrives as plain data and is resolved here, so the class object never reaches the # constructor -- it would either be rejected by the remote-class layer in safe mode # (``TWINKLE_TRUST_REMOTE_CODE=0``) or leak into every worker's constructor kwargs. The schema only diff --git a/src/twinkle_agentic/rollout/base.py b/src/twinkle_agentic/rollout/base.py index a75db775d..e6231590c 100644 --- a/src/twinkle_agentic/rollout/base.py +++ b/src/twinkle_agentic/rollout/base.py @@ -1,13 +1,12 @@ # Copyright (c) ModelScope Contributors. All rights reserved. from __future__ import annotations - import threading import uuid from abc import ABC, abstractmethod from collections.abc import Callable -from concurrent.futures import Future, ThreadPoolExecutor, as_completed +from concurrent.futures import CancelledError, Future, ThreadPoolExecutor, as_completed, wait from dataclasses import dataclass -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple +from typing import TYPE_CHECKING, Any, Optional from twinkle.data_format import Trajectory from twinkle.data_format.sampling import SamplingParams @@ -22,10 +21,10 @@ # Termination reasons surfaced via ``trajectory['stop_reason']``. The sampler # path takes the first three from the sampler itself; the API path has to name # them, and one vocabulary for both is what lets a consumer read either. -STOP_NO_TOOL = 'stop' -STOP_LENGTH = 'length' -STOP_MAX_TURNS = 'max_turns' -STOP_GENERATION_ERROR = 'generation_error' +STOP_NO_TOOL = "stop" +STOP_LENGTH = "length" +STOP_MAX_TURNS = "max_turns" +STOP_GENERATION_ERROR = "generation_error" # Runaway guard: a ``followup_fn`` is expected to return None eventually. This # only bounds a callback that never does, so one bad hook cannot spin forever. @@ -62,6 +61,40 @@ def __eq__(self, other: Any) -> bool: return isinstance(other, SampleHandle) and self.submission_id == other.submission_id +class _RolloutExecution: + """Process-local cooperative stop state for one blocking rollout call.""" + + def __init__(self) -> None: + self.cancel_event = threading.Event() + self._lock = threading.Lock() + self._first_error: BaseException | None = None + + @property + def cancelled(self) -> bool: + return self.cancel_event.is_set() + + @property + def first_error(self) -> BaseException | None: + with self._lock: + return self._first_error + + def cancel(self) -> None: + self.cancel_event.set() + + def checkpoint(self) -> None: + """Linearize an operation start with cancellation under one lock.""" + with self._lock: + if self.cancel_event.is_set(): + raise CancelledError() + + def fail(self, error: BaseException) -> None: + with self._lock: + cooperative = isinstance(error, (CancelledError, InterruptedError)) and self.cancel_event.is_set() + if self._first_error is None and not cooperative: + self._first_error = error + self.cancel_event.set() + + class Rollout(ABC): """A batch of trajectories in, the same batch with the model's turns appended. @@ -90,9 +123,9 @@ class Rollout(ABC): # own setup still answers these attributes instead of raising from a base # method it inherited. max_turns: int = 1 - sampling_params: Optional[SamplingParams] = None - tracer: Optional[TraceWriter] = None - concurrency: Optional[int] = None + sampling_params: SamplingParams | None = None + tracer: TraceWriter | None = None + concurrency: int | None = None # ------------------------------------------------------------------ setup @@ -100,9 +133,9 @@ def _init_common( self, *, max_turns: int = 1, - sampling_params: Optional[SamplingParams] = None, - concurrency: Optional[int] = None, - tracer: Optional[TraceWriter] = None, + sampling_params: SamplingParams | None = None, + concurrency: int | None = None, + tracer: TraceWriter | None = None, ) -> None: """Validate and store the options every rollout takes. @@ -111,16 +144,15 @@ def _init_common( decides it is done -- leaves it alone. """ if max_turns < 1: - raise ValueError(f'max_turns must be >= 1, got {max_turns}') + raise ValueError(f"max_turns must be >= 1, got {max_turns}") sp = sampling_params or SamplingParams() if sp.num_samples != 1: # n>1 would fork the conversation at turn 1 and there is no defined # way to continue a forked trajectory: ask for several rollouts by # passing the trajectory several times instead. - raise ValueError(f'{type(self).__name__} supports num_samples=1 only, ' - f'got {sp.num_samples}') + raise ValueError(f"{type(self).__name__} supports num_samples=1 only, got {sp.num_samples}") if concurrency is not None and concurrency < 1: - raise ValueError(f'concurrency must be >= 1 or None, got {concurrency}') + raise ValueError(f"concurrency must be >= 1 or None, got {concurrency}") self.max_turns = max_turns self.sampling_params = sp # None means one thread per trajectory. A cap below the batch size costs @@ -130,7 +162,7 @@ def _init_common( # ------------------------------------------------------------------- drive - def __call__(self, trajectories: List[Trajectory], **kwargs) -> List[Trajectory]: + def __call__(self, trajectories: list[Trajectory], **kwargs) -> list[Trajectory]: """Run one episode per trajectory and return them in the input order. Order is restored from the future map rather than from completion order, @@ -138,33 +170,62 @@ def __call__(self, trajectories: List[Trajectory], **kwargs) -> List[Trajectory] GRPO group is a slice of this list. """ if isinstance(trajectories, dict): - raise TypeError(f'{type(self).__name__}.__call__ expects a List[Trajectory]; ' - 'wrap a single trajectory as [trajectory].') + raise TypeError( + f"{type(self).__name__}.__call__ expects a List[Trajectory]; wrap a single trajectory as [trajectory]." + ) trajectories = list(trajectories) n = len(trajectories) if n == 0: return [] + execution = kwargs.pop("_execution", None) or _RolloutExecution() ctx = self._resolve_call(kwargs, n) - outs: List[Optional[Trajectory]] = [None] * n + ctx["_execution"] = execution + outs: list[Trajectory | None] = [None] * n workers = min(n, self.concurrency or n) + + def run_one(index: int) -> Trajectory: + execution.checkpoint() + try: + return self._run_one(trajectories[index], index, ctx) + except Exception as error: + execution.fail(error) + raise + if workers == 1: - # No pool for a single episode: a thread would only make the - # traceback of a failing one harder to read. - outs = [self._run_one(trajectories[i], i, ctx) for i in range(n)] + for index in range(n): + execution.checkpoint() + outs[index] = run_one(index) else: - with ThreadPoolExecutor(max_workers=workers) as pool: - futures = {pool.submit(self._run_one, trajectories[i], i, ctx): i for i in range(n)} - for fut in as_completed(futures): - outs[futures[fut]] = fut.result() - - result: List[Trajectory] = [o if o is not None else dict(trajectories[i]) for i, o in enumerate(outs)] + pool = ThreadPoolExecutor(max_workers=workers) + futures: dict[Future, int] = {} + try: + try: + for index in range(n): + if execution.cancelled: + break + futures[pool.submit(run_one, index)] = index + for future in as_completed(futures): + try: + outs[futures[future]] = future.result() + except Exception as error: + execution.fail(error) + finally: + wait(futures) + finally: + pool.shutdown(wait=True) + + error = execution.first_error + if error is not None: + raise error + execution.checkpoint() + result: list[Trajectory] = [o if o is not None else dict(trajectories[i]) for i, o in enumerate(outs)] if self.tracer is not None: - self.tracer.write(result, global_step=kwargs.get('global_step')) + self.tracer.write(result, global_step=kwargs.get("global_step")) return result @abstractmethod - def _run_one(self, trajectory: Trajectory, index: int, ctx: Dict[str, Any]) -> Trajectory: + def _run_one(self, trajectory: Trajectory, index: int, ctx: dict[str, Any]) -> Trajectory: """One trajectory, start to finish, in its own thread. ``ctx`` is whatever ``_resolve_call`` produced; ``index`` is the @@ -176,11 +237,17 @@ def _run_one(self, trajectory: Trajectory, index: int, ctx: Dict[str, Any]) -> T # -------------------------------------------------------------- streaming @abstractmethod - def admit(self, prompt: List[Dict[str, Any]], trajectory_idx: int = 0, *, prompt_idx: Any = 0, - sampling_params: Optional[Dict[str, Any]] = None, - prompt_extras: Optional[Dict[str, Any]] = None, - adapter_path: Optional[str] = None, - policy_version: int = 0) -> SampleHandle: + def admit( + self, + prompt: list[dict[str, Any]], + trajectory_idx: int = 0, + *, + prompt_idx: Any = 0, + sampling_params: dict[str, Any] | None = None, + prompt_extras: dict[str, Any] | None = None, + adapter_path: str | None = None, + policy_version: int = 0, + ) -> SampleHandle: """Admit ONE trajectory without blocking; poll/collect/stop complete the protocol. The streaming contract the async driver default-binds: the returned handle @@ -191,7 +258,7 @@ def admit(self, prompt: List[Dict[str, Any]], trajectory_idx: int = 0, *, prompt raise NotImplementedError() @abstractmethod - def poll(self, handles: List[SampleHandle]) -> List[SampleHandle]: + def poll(self, handles: list[SampleHandle]) -> list[SampleHandle]: """Non-blocking as-completed query over admitted handles. Returns the subset whose episodes finished; a failed episode raises here, @@ -218,12 +285,12 @@ def stop(self, submission_id: str) -> None: def _new_handle( self, - prompt: List[Dict[str, Any]], + prompt: list[dict[str, Any]], trajectory_idx: int, *, prompt_idx: Any, - prompt_extras: Optional[Dict[str, Any]], - adapter_path: Optional[str], + prompt_extras: dict[str, Any] | None, + adapter_path: str | None, policy_version: int, ) -> SampleHandle: """The submit-time facts every admission records, on the protocol's handle type.""" @@ -245,7 +312,7 @@ def _admit_local(self, handle: SampleHandle, run_episode: Callable[[], Any]) -> also runs on a Ray proxy whose ``__init__`` never executed. The future is not popped here: poll probes it, collect pops it, cancel/stop drop it. """ - futures: Dict[str, Future] = self.__dict__.setdefault('_episode_futures', {}) + futures: dict[str, Future] = self.__dict__.setdefault("_episode_futures", {}) future: Future = Future() futures[handle.submission_id] = future @@ -255,7 +322,7 @@ def _work() -> None: except Exception as exc: # surfaced at poll/collect; the thread must not die silently future.set_exception(exc) - threading.Thread(target=_work, name=f'episode-{handle.submission_id[:8]}', daemon=True).start() + threading.Thread(target=_work, name=f"episode-{handle.submission_id[:8]}", daemon=True).start() return handle def _episode_done(self, handle: SampleHandle) -> bool: @@ -263,21 +330,19 @@ def _episode_done(self, handle: SampleHandle) -> bool: A failed episode raises here, fail-loudly, rather than waiting for its collect. """ - future = getattr(self, '_episode_futures', {}).get(handle.submission_id) + future = getattr(self, "_episode_futures", {}).get(handle.submission_id) if future is None or not future.done(): return False error = future.exception() if error is not None: - raise RuntimeError( - f'{type(self).__name__} episode {handle.submission_id} failed: {error}') from error + raise RuntimeError(f"{type(self).__name__} episode {handle.submission_id} failed: {error}") from error return True def _pop_episode_future(self, handle: SampleHandle) -> Future: """Take the episode's future at collect; an unknown submission is a caller bug.""" - future = getattr(self, '_episode_futures', {}).pop(handle.submission_id, None) + future = getattr(self, "_episode_futures", {}).pop(handle.submission_id, None) if future is None: - raise RuntimeError( - f'collect called for an unknown or already-collected submission {handle.submission_id}.') + raise RuntimeError(f"collect called for an unknown or already-collected submission {handle.submission_id}.") return future @staticmethod @@ -290,9 +355,9 @@ def _stamp_admission(sample: RolloutSample, handle: SampleHandle) -> RolloutSamp def _drop_episode_future(self, submission_id: str) -> None: """Forget an admitted episode's future; its thread runs on, discarded.""" - self.__dict__.get('_episode_futures', {}).pop(submission_id, None) + self.__dict__.get("_episode_futures", {}).pop(submission_id, None) - def _resolve_call(self, kwargs: Dict[str, Any], n: int) -> Dict[str, Any]: + def _resolve_call(self, kwargs: dict[str, Any], n: int) -> dict[str, Any]: """Fold per-call ``**kwargs`` over the constructor defaults, once. Done before the pool starts so a bad argument raises from the caller's @@ -302,7 +367,7 @@ def _resolve_call(self, kwargs: Dict[str, Any], n: int) -> Dict[str, Any]: return {} @staticmethod - def _unpack_followup(followup: Any) -> Tuple[str, Optional[SamplingParams]]: + def _unpack_followup(followup: Any) -> tuple[str, SamplingParams | None]: """``followup_fn`` may answer with text, or text plus its own budget.""" if isinstance(followup, tuple): text, params = followup @@ -310,7 +375,7 @@ def _unpack_followup(followup: Any) -> Tuple[str, Optional[SamplingParams]]: return followup, None @staticmethod - def _broadcast(arg, n: int, *, name: str, required: bool = False, per_trajectory: bool = False) -> List[Any]: + def _broadcast(arg, n: int, *, name: str, required: bool = False, per_trajectory: bool = False) -> list[Any]: """One value shared by the batch, or a list already aligned 1:1 with it. A list of the wrong length is refused rather than zipped short: the @@ -324,16 +389,20 @@ def _broadcast(arg, n: int, *, name: str, required: bool = False, per_trajectory """ if arg is None: if required: - raise ValueError(f'{name} is required but was not provided. ' - 'Pass it at construction time or as a per-call kwarg.') + raise ValueError( + f"{name} is required but was not provided. Pass it at construction time or as a per-call kwarg." + ) return [None] * n if isinstance(arg, list): if len(arg) != n: - raise ValueError(f'per-call {name} list length ({len(arg)}) does ' - f'not match number of trajectories ({n})') + raise ValueError( + f"per-call {name} list length ({len(arg)}) does not match number of trajectories ({n})" + ) return list(arg) if per_trajectory and n > 1: - raise ValueError(f'{name} holds per-episode state and cannot be shared by ' - f'{n} trajectories running in parallel threads: pass a list ' - f'of {n}, one per trajectory.') + raise ValueError( + f"{name} holds per-episode state and cannot be shared by " + f"{n} trajectories running in parallel threads: pass a list " + f"of {n}, one per trajectory." + ) return [arg] * n diff --git a/src/twinkle_agentic/rollout/multi_turn.py b/src/twinkle_agentic/rollout/multi_turn.py index 1ea97a44f..2dcf6d282 100644 --- a/src/twinkle_agentic/rollout/multi_turn.py +++ b/src/twinkle_agentic/rollout/multi_turn.py @@ -2,6 +2,7 @@ import contextlib import copy import re +import uuid from collections.abc import Callable from dataclasses import dataclass, field, replace from typing import Any, Literal, Optional @@ -15,6 +16,7 @@ from twinkle_agentic.protocol.api_sampler import APIGenerationError, APISampler from twinkle_agentic.protocol.base import API from twinkle_agentic.tools.tool_manager import ToolManager +from twinkle_agentic.utils.leases import Leases from .base import MAX_FOLLOWUPS, STOP_GENERATION_ERROR, STOP_NO_TOOL, Rollout, SampleHandle from .ledger import MessageLedger, TurnLedger, last_assistant_text from .trace import TraceWriter @@ -759,8 +761,13 @@ def admit( runs the episode on a daemon-thread Future. """ handle = self._new_handle( - prompt, trajectory_idx, prompt_idx=prompt_idx, prompt_extras=prompt_extras, - adapter_path=adapter_path, policy_version=policy_version) + prompt, + trajectory_idx, + prompt_idx=prompt_idx, + prompt_extras=prompt_extras, + adapter_path=adapter_path, + policy_version=policy_version, + ) generate_kwargs = dict( num_samples=1, sampling_params=sampling_params, @@ -867,6 +874,12 @@ def _resolve_call(self, kwargs: dict[str, Any], n: int) -> dict[str, Any]: if self.sampler is not None and kwargs.get("allow_partial_rollout", False): adapter_kwargs["allow_partial_rollout"] = True sampling_params = kwargs.get("sampling_params", self.sampling_params) + if isinstance(sampling_params, dict): + sampling_params = SamplingParams.from_dict(dict(sampling_params)) + elif isinstance(sampling_params, SamplingParams): + sampling_params = replace(sampling_params) + else: + raise TypeError("sampling_params must be SamplingParams or a dict") if sampling_params.num_samples != 1: raise ValueError(f"MultiTurnRollout supports num_samples=1 only, got {sampling_params.num_samples}") response_callback = kwargs.get("response_callback", self.response_callback) @@ -926,12 +939,27 @@ def _new_ledger(self, index: int): backend=self._template_backend, ) + @staticmethod + def _lease(pool, execution): + if isinstance(pool, Leases): + return pool.lease(cancel_event=execution.cancel_event) + return pool.lease() + + @staticmethod + def _validate_ledger(ledger, ctx: dict[str, Any], index: int) -> None: + validate = getattr(ctx["response_callback"], "_validate_state", None) + if callable(validate): + validate(ledger.state, ledger.input_feature, index) + def _run_one(self, trajectory: Trajectory, index: int, ctx: dict[str, Any]) -> Trajectory: + execution = ctx["_execution"] + execution.checkpoint() if self._per_episode_tools: # A sandbox env holds one workspace: the trajectory drives alone with its # tools bound to the env it leased; the lease blocks when the pool empties # (the sandbox's own concurrency bound) and is held for the whole episode. - with self.env_pool.lease() as env: + with self._lease(self.env_pool, execution) as env: + execution.checkpoint() scoped = dict(ctx) scoped["tool_managers"] = [self._tool_manager_for(env, self.tool_plugins)] * len(ctx["tool_managers"]) return self._run_with_harness(trajectory, index, scoped) @@ -943,11 +971,14 @@ def _run_with_harness(self, trajectory: Trajectory, index: int, ctx: dict[str, A # next one. Anything else is used as passed. harness = ctx["harnesses"][index] if harness is not None and hasattr(harness, "lease"): - with harness.lease() as leased: + with self._lease(harness, ctx["_execution"]) as leased: + ctx["_execution"].checkpoint() return self._run_episode(trajectory, index, ctx, leased) return self._run_episode(trajectory, index, ctx, harness) def _run_episode(self, trajectory: Trajectory, index: int, ctx: dict[str, Any], harness=None) -> Trajectory: + execution = ctx["_execution"] + execution.checkpoint() tool_manager: ToolManager = ctx["tool_managers"][index] followup_fn = ctx["followup_fn"] # Each episode owns its metadata; mutating the shared dict would let @@ -971,6 +1002,7 @@ def _run_episode(self, trajectory: Trajectory, index: int, ctx: dict[str, Any], if "tools" not in trajectory and tool_manager is not None: opening_tools = list(tool_manager.tool_infos() or []) ledger.open(trajectory, tools=opening_tools) + self._validate_ledger(ledger, ctx, index) stop_reason: str | None = None generation_error: str | None = None @@ -1010,15 +1042,18 @@ def grant_followup() -> bool: nonlocal followups, budget, params, truncated if followup_fn is None or followups >= MAX_FOLLOWUPS: return False + execution.checkpoint() followup = followup_fn( ledger.merge(trajectory, turns=ledger.turns, stop_reason=stop_reason, truncated=truncated), followups ) + execution.checkpoint() if followup is None: return False text, next_params = self._unpack_followup(followup) if not ledger.observe([{"role": "user", "content": text}]): truncated = True return False + self._validate_ledger(ledger, ctx, index) # Follow-up stages are answers, so an API must not see tool schemas. ledger.clear_tools() followups += 1 @@ -1028,11 +1063,13 @@ def grant_followup() -> bool: return True while spent < budget: + execution.checkpoint() spent += 1 # 2. One request. The callback chooses the local sampler or the API # adapter, but both paths return exactly one SampledSequence. try: + execution.checkpoint() seq = response_callback( self.sampler, self.api, @@ -1050,10 +1087,12 @@ def grant_followup() -> bool: generation_error = str(exc) truncated = True break + execution.checkpoint() if not isinstance(seq, SampledSequence): raise TypeError(f"response_callback must return SampledSequence, got {type(seq).__name__}") ledger.record(seq) + self._validate_ledger(ledger, ctx, index) stop_reason = seq.stop_reason msgs = ledger.messages @@ -1102,6 +1141,7 @@ def grant_followup() -> bool: if not ledger.observe([_malformed_tool_message(parse_errors)]): truncated = True break + self._validate_ledger(ledger, ctx, index) continue # The episode is over as far as the model is concerned. Give the # caller one chance to say otherwise -- see ``followup_fn`` for @@ -1125,7 +1165,9 @@ def grant_followup() -> bool: # messages the model did not write. if tool_manager is None: raise ValueError("the model emitted tool_calls but this trajectory has no ToolManager") + execution.checkpoint() observations = self._run_tools(tool_manager, tool_calls) + execution.checkpoint() if self.stop_after_stuck_turns: keys = [_call_key(tc) for tc in tool_calls] all_repeats = bool(keys) and all(k in seen_calls for k in keys) @@ -1141,6 +1183,8 @@ def grant_followup() -> bool: else: tool_messages = _default_tool_messages(tool_calls, observations) overflowed = not ledger.observe(tool_messages) + if not overflowed: + self._validate_ledger(ledger, ctx, index) if overflowed: truncated = True # Checked after the messages are appended, so the turns that ended @@ -1161,6 +1205,7 @@ def grant_followup() -> bool: # consumers (VLLMSampler with ``'input_ids' in inputs``) see an encoded # InputFeature and skip re-encoding. The ledger audits its own account # on the way out -- one logprob per trainable token, or it raises. + execution.checkpoint() out = ledger.merge( trajectory, turns=ledger.turns, diff --git a/src/twinkle_agentic/tools/tool_manager.py b/src/twinkle_agentic/tools/tool_manager.py index c7f944ab1..a920ddae3 100644 --- a/src/twinkle_agentic/tools/tool_manager.py +++ b/src/twinkle_agentic/tools/tool_manager.py @@ -1,7 +1,9 @@ # Copyright (c) ModelScope Contributors. All rights reserved. -import json +from collections.abc import Iterable from concurrent.futures import ThreadPoolExecutor, as_completed -from typing import Any, Dict, Iterable, List, Optional, Tuple, Union +from typing import Any + +import json from twinkle.data_format import ToolCall from twinkle.data_format.message import Tool as ToolInfo @@ -11,50 +13,49 @@ logger = get_logger() -def _extract_name(info: Any) -> Optional[str]: +def _extract_name(info: Any) -> str | None: """Read ``function.name`` from an OpenAI-shaped tool / tool-call dict.""" if not isinstance(info, dict): return None - fn = info.get('function') + fn = info.get("function") if isinstance(fn, dict): - name = fn.get('name') + name = fn.get("name") if isinstance(name, str) and name: return name return None -def _unpack_tool_call(tool_call: Any) -> Tuple[Optional[str], Dict[str, Any], Optional[str]]: +def _unpack_tool_call(tool_call: Any) -> tuple[str | None, dict[str, Any], str | None]: """Split an OpenAI-shaped tool_call into ``(name, args, error)``. These dicts come from :meth:`twinkle.template.base.Template.parse_tool_call`. ``error`` is set when the payload cannot be executed. """ if not isinstance(tool_call, dict): - return None, {}, f'Error: tool_call must be an object, got {type(tool_call).__name__}.' - fn = tool_call.get('function') + return None, {}, f"Error: tool_call must be an object, got {type(tool_call).__name__}." + fn = tool_call.get("function") if not isinstance(fn, dict): return None, {}, 'Error: tool_call missing "function" object.' - name = fn.get('name') + name = fn.get("name") if not name: return None, {}, 'Error: tool_call missing "function.name".' - raw_args = fn.get('arguments') + raw_args = fn.get("arguments") if raw_args is None: return str(name), {}, None if isinstance(raw_args, str): try: args = json.loads(raw_args) if raw_args.strip() else {} except json.JSONDecodeError as e: - return str(name), {}, f'Error: invalid JSON in arguments: {e}' + return str(name), {}, f"Error: invalid JSON in arguments: {e}" if not isinstance(args, dict): return str(name), {}, 'Error: "arguments" JSON must be an object.' return str(name), args, None if isinstance(raw_args, dict): return str(name), raw_args, None - return None, {}, (f'Error: "arguments" must be a JSON string or object, ' - f'got {type(raw_args).__name__}.') + return None, {}, (f'Error: "arguments" must be a JSON string or object, got {type(raw_args).__name__}.') -def _suggest(name: str, available: Iterable[str]) -> Optional[str]: +def _suggest(name: str, available: Iterable[str]) -> str | None: """The registered tool ``name`` was probably meant to be, if there is one. Only one mistake is guessed at: a name given without its namespace, or under @@ -69,21 +70,20 @@ def _suggest(name: str, available: Iterable[str]) -> Optional[str]: accepts, and the unqualified form is ambiguous the moment two servers export the same verb -- which is why a suffix shared by several tools yields nothing. """ - wanted = name.rsplit('---', 1)[-1] + wanted = name.rsplit("---", 1)[-1] if not wanted: return None - matches = [n for n in available if n != name and n.rsplit('---', 1)[-1] == wanted] + matches = [n for n in available if n != name and n.rsplit("---", 1)[-1] == wanted] return matches[0] if len(matches) == 1 else None class ToolManager: - def __init__( self, - tools: Optional[Union[Dict[str, Tool], Iterable[Tool]]] = None, + tools: dict[str, Tool] | Iterable[Tool] | None = None, ): if tools is None: - self._tools: Dict[str, Tool] = {} + self._tools: dict[str, Tool] = {} return if isinstance(tools, dict): self._tools = dict(tools) @@ -91,61 +91,59 @@ def __init__( if isinstance(tools, (list, tuple)): self._tools = {} for t in tools: - info = t.tool_info() if hasattr(t, 'tool_info') else None + info = t.tool_info() if hasattr(t, "tool_info") else None name = _extract_name(info) if not name: - raise ValueError(f'tool {type(t).__name__} must expose a non-empty ' - f'tool_info()["function"]["name"]') + raise ValueError(f'tool {type(t).__name__} must expose a non-empty tool_info()["function"]["name"]') self._tools[name] = t return - raise TypeError(f'ToolManager expects dict | Iterable[Tool] | None; ' - f'got {type(tools).__name__}') + raise TypeError(f"ToolManager expects dict | Iterable[Tool] | None; got {type(tools).__name__}") def register(self, tool: Tool): - info = tool.tool_info() if hasattr(tool, 'tool_info') else None + info = tool.tool_info() if hasattr(tool, "tool_info") else None name = _extract_name(info) if not name: - raise ValueError(f'tool {type(tool).__name__} must expose a non-empty ' - f'tool_info()["function"]["name"]') + raise ValueError(f'tool {type(tool).__name__} must expose a non-empty tool_info()["function"]["name"]') self._tools[name] = tool - def unregister(self, name: str) -> Optional[Tool]: + def unregister(self, name: str) -> Tool | None: return self._tools.pop(name, None) - def names(self) -> List[str]: + def names(self) -> list[str]: return list(self._tools) - def copy(self) -> 'ToolManager': + def copy(self) -> "ToolManager": return ToolManager(dict(self._tools)) - def tool_infos(self) -> List[ToolInfo]: + def tool_infos(self) -> list[ToolInfo]: return [t.tool_info() for t in self._tools.values()] - def __call__(self, tool_call: Union[ToolCall, Dict[str, Any]]) -> str: + def __call__(self, tool_call: ToolCall | dict[str, Any]) -> str: name, args, err = _unpack_tool_call(tool_call) if err: return err if (tool := self._tools.get(name)) is None: - available = ', '.join(sorted(self._tools)) or '(none)' - hint = '' + available = ", ".join(sorted(self._tools)) or "(none)" + hint = "" if (suggestion := _suggest(name, self._tools)) is not None: - hint = (f' Did you mean {suggestion!r}? Tool names must be given in ' - f'full, including the part before "---".') - return f'Error: unknown tool {name!r}.{hint} Available: {available}.' + hint = ( + f' Did you mean {suggestion!r}? Tool names must be given in full, including the part before "---".' + ) + return f"Error: unknown tool {name!r}.{hint} Available: {available}." try: return str(tool(name, args)) except Exception as e: # noqa - return f'Error: tool {name!r} raised {type(e).__name__}: {e}' + return f"Error: tool {name!r} raised {type(e).__name__}: {e}" finally: - env = getattr(tool, '_env', None) + env = getattr(tool, "_env", None) if env is not None: env.ensure_idle() def call_many( self, - tool_calls: Iterable[Union[ToolCall, Dict[str, Any]]], - max_workers: Optional[int] = None, - ) -> List[str]: + tool_calls: Iterable[ToolCall | dict[str, Any]], + max_workers: int | None = None, + ) -> list[str]: """Execute many tool calls, preserving input order. ``tool_calls`` are the OpenAI-shaped dicts produced by @@ -178,8 +176,8 @@ def call_many( unpacked = [_unpack_tool_call(tc) for tc in calls] env = self._shared_env() if env is not None: - out: List[Optional[str]] = [None] * len(calls) - batched: List[Tuple[int, str, Dict[str, Any]]] = [] + out: list[str | None] = [None] * len(calls) + batched: list[tuple[int, str, dict[str, Any]]] = [] for i, (name, args, err) in enumerate(unpacked): if err is None and name in self._tools: batched.append((i, name, args)) @@ -194,28 +192,30 @@ def call_many( # thread pool instead -- which is what this used to do, silently # and without even a log line -- sends the turn along the path the # docstring above exists to keep it off. - logger.warning(f'{type(env).__name__}.step_batch raised ' - f'{type(e).__name__}: {e}') - failure = f'Error: tool batch did not run: {type(e).__name__}: {e}' + logger.warning(f"{type(env).__name__}.step_batch raised {type(e).__name__}: {e}") + failure = f"Error: tool batch did not run: {type(e).__name__}: {e}" for i, _name, _args in batched: out[i] = failure - return ['' if x is None else x for x in out] + return ["" if x is None else x for x in out] finally: env.ensure_idle() if len(results) != len(batched): # Same reasoning: a short result list means the calls did not all # run, and pairing them up by position would report one call's # result under another's name. - logger.warning(f'{type(env).__name__}.step_batch returned ' - f'{len(results)} results for {len(batched)} calls') - failure = (f'Error: tool batch did not run: the environment returned ' - f'{len(results)} results for {len(batched)} calls.') + logger.warning( + f"{type(env).__name__}.step_batch returned {len(results)} results for {len(batched)} calls" + ) + failure = ( + f"Error: tool batch did not run: the environment returned " + f"{len(results)} results for {len(batched)} calls." + ) for i, _name, _args in batched: out[i] = failure - return ['' if x is None else x for x in out] - for (i, _name, _args), r in zip(batched, results): - out[i] = r.observation if hasattr(r, 'observation') else str(r) - return ['' if x is None else x for x in out] + return ["" if x is None else x for x in out] + for (i, _name, _args), r in zip(batched, results, strict=False): + out[i] = r.observation if hasattr(r, "observation") else str(r) + return ["" if x is None else x for x in out] workers = max_workers or min(32, len(calls)) out = [None] * len(calls) @@ -223,22 +223,30 @@ def call_many( futs = {pool.submit(self, tc): i for i, tc in enumerate(calls)} for fut in as_completed(futs): out[futs[fut]] = fut.result() - return ['' if x is None else x for x in out] + return ["" if x is None else x for x in out] - def ensure_idle(self): - """所有实际工具线程退出后,还须确认它们使用的远程环境已结束执行。""" + def _iter_bound_envs(self): seen = set() for tool in self._tools.values(): - env = getattr(tool, '_env', None) + env = getattr(tool, "_env", None) if env is not None and id(env) not in seen: seen.add(id(env)) - env.ensure_idle() + yield env + + def bound_envs(self): + """Return the distinct environments borrowed by the registered tools.""" + return tuple(self._iter_bound_envs()) + + def ensure_idle(self): + """Confirm every distinct borrowed environment has no known active execution.""" + for env in self._iter_bound_envs(): + env.ensure_idle() def _shared_env(self): """Return the Env shared by every registered EnvTool, else None.""" env = None for tool in self._tools.values(): - wrapped = getattr(tool, '_env', None) + wrapped = getattr(tool, "_env", None) if wrapped is None: return None if env is None: diff --git a/src/twinkle_client/rollout/_http_sampler.py b/src/twinkle_client/rollout/_http_sampler.py new file mode 100644 index 000000000..02e62ed0e --- /dev/null +++ b/src/twinkle_client/rollout/_http_sampler.py @@ -0,0 +1,68 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Adapt the existing HTTP sampler to the core multi-turn callback contract.""" +from __future__ import annotations + +from typing import Any + +from twinkle.data_format.sampling import SampledSequence + + +class _HTTPSamplerAdapter: + """Capture one call's policy and translate one wire response per turn.""" + + def __init__(self, sampler, *, adapter_name: str = '', adapter_uri: str | None = None): + self._sampler = sampler + self._adapter_name = adapter_name + self._adapter_uri = adapter_uri + + def __call__( + self, + sampler, + api, + sampling_params, + *, + input_feature, + adapter_kwargs, + rollout_state=None, + trajectory_index=0, + **kwargs, + ) -> SampledSequence: + del sampler, api, kwargs + self._validate_state(rollout_state, input_feature, trajectory_index) + adapter_uri = adapter_kwargs.get('adapter_path', self._adapter_uri) + responses = self._sampler.sample( + [input_feature], + sampling_params=sampling_params, + adapter_name=self._adapter_name, + adapter_uri=adapter_uri, + num_samples=1, + ) + if not isinstance(responses, list) or len(responses) != 1: + count = len(responses) if isinstance(responses, list) else type(responses).__name__ + raise RuntimeError(f'sampler returned {count} responses for trajectory {trajectory_index}; expected one') + response = responses[0] + if len(response.sequences) != 1: + raise RuntimeError( + f'sampler returned {len(response.sequences)} sequences for trajectory {trajectory_index}; expected one' + ) + sequence = response.sequences[0] + return SampledSequence( + stop_reason=sequence.stop_reason, + tokens=list(sequence.tokens), + logprobs=sequence.logprobs, + decoded=sequence.decoded, + new_input_feature=sequence.new_input_feature, + ) + + @staticmethod + def _validate_state(state, feature: dict[str, Any], trajectory_index: int) -> None: + feature_ids = list(feature.get('input_ids') or []) + if state is None or state.media: + raise NotImplementedError( + f'trajectory {trajectory_index} requires server-side RolloutState state_id for HTTP multi-turn media' + ) + if feature_ids != list(state.engine_input_ids): + raise NotImplementedError( + f'trajectory {trajectory_index} requires server-side RolloutState state_id because feature and engine ' + 'ids differ' + ) diff --git a/src/twinkle_client/rollout/multi_turn.py b/src/twinkle_client/rollout/multi_turn.py index 7d935d964..9f2bf344e 100644 --- a/src/twinkle_client/rollout/multi_turn.py +++ b/src/twinkle_client/rollout/multi_turn.py @@ -1,19 +1,37 @@ # Copyright (c) ModelScope Contributors. All rights reserved. """Client-side trainable multi-turn rollout over HTTP.""" from __future__ import annotations -import dataclasses +import asyncio +import threading +from concurrent.futures import Future, ThreadPoolExecutor +from dataclasses import dataclass, replace 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.ledger import TurnLedger +from twinkle_agentic.rollout.base import _RolloutExecution +from twinkle_agentic.rollout.multi_turn import MultiTurnRollout from twinkle_agentic.tools.tool_manager import ToolManager from twinkle_client.sampler import vLLMSampler +from ._http_sampler import _HTTPSamplerAdapter + +_CALL_OPTIONS = { + 'sampling_params', 'tool_manager', 'adapter_name', 'adapter_uri', 'followup_fn', 'harness', 'global_step' +} + + +@dataclass +class _Call: + execution: _RolloutExecution + env_ids: tuple[int, ...] + harness_ids: tuple[int, ...] + executor: ThreadPoolExecutor + future: Future | None = None class ClientMultiTurnRollout: - """Drive safe text-only multi-turn sampling through ``/twinkle/sample``.""" + """Run the shared multi-turn core with per-call HTTP sampling adapters.""" def __init__( self, @@ -23,22 +41,158 @@ def __init__( sampling_params: Optional[SamplingParams] = None, max_turns: int = 6, max_trajectory_tokens: Optional[int] = None, + *, + concurrency: Optional[int] = None, + adapter_name: str = '', + adapter_uri: Optional[str] = None, + stop_after_stuck_turns: int = 0, + max_malformed_retries: int = 2, + followup_fn=None, + harness=None, + tracer=None, ): 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 {max_trajectory_tokens}') self.sampler = sampler self.template = template self.tool_manager = tool_manager - self.sampling_params = sampling_params or SamplingParams() + self.sampling_params = self._normalize_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('ClientMultiTurnRollout supports num_samples=1 only, ' - f'got {self.sampling_params.num_samples}') + self.concurrency = concurrency + self.adapter_name = adapter_name + self.adapter_uri = adapter_uri + self.followup_fn = followup_fn + self.harness = harness + self._core = MultiTurnRollout( + sampler=None, + template=template, + tool_manager=tool_manager, + sampling_params=self.sampling_params, + max_turns=max_turns, + max_trajectory_tokens=max_trajectory_tokens, + concurrency=concurrency, + response_callback=_HTTPSamplerAdapter( + sampler, adapter_name=adapter_name, adapter_uri=adapter_uri), + stop_after_stuck_turns=stop_after_stuck_turns, + max_malformed_retries=max_malformed_retries, + followup_fn=followup_fn, + harness=harness, + tracer=tracer, + ) + if hasattr(self._core, '_actors'): + raise RuntimeError('ClientMultiTurnRollout requires the in-process core, not a Ray actor proxy') + self._lock = threading.Lock() + self._active: dict[int, _Call] = {} + self._env_borrows: dict[int, Any] = {} + self._harness_borrows: dict[int, Any] = {} + + @staticmethod + def _normalize_sampling_params(value) -> SamplingParams: + if isinstance(value, dict): + return SamplingParams.from_dict(dict(value)) + if isinstance(value, SamplingParams): + return replace(value) + raise TypeError('sampling_params must be SamplingParams or a dict') + + @staticmethod + def _as_batch(value, count: int, name: str) -> list[Any]: + if isinstance(value, list): + if len(value) != count: + raise ValueError(f'{name} list length ({len(value)}) does not match trajectories ({count})') + return list(value) + return [value] * count + + @staticmethod + def _duplicate_ids(values) -> tuple[int, ...]: + seen: set[int] = set() + duplicates: list[int] = [] + for value in values: + if value is None: + continue + identity = id(value) + if identity in seen and identity not in duplicates: + duplicates.append(identity) + seen.add(identity) + return tuple(duplicates) + + def _borrowed_identities(self, managers, harnesses) -> tuple[tuple[int, ...], tuple[int, ...]]: + envs = [] + for manager in managers: + bound_envs = getattr(manager, 'bound_envs', None) + if callable(bound_envs): + envs.extend(bound_envs()) + bare_harnesses = [harness for harness in harnesses if harness is not None and not hasattr(harness, 'lease')] + duplicate_envs = self._duplicate_ids(envs) + duplicate_harnesses = self._duplicate_ids(bare_harnesses) + if duplicate_envs: + raise ValueError('the same Env cannot be borrowed by multiple episodes; provide an independent Env per ' + 'trajectory and concurrent call') + if duplicate_harnesses: + raise ValueError('the same bare harness cannot be borrowed by multiple episodes; use independent harnesses ' + 'or HarnessLeases') + env_ids = tuple(dict.fromkeys(id(env) for env in envs)) + harness_ids = tuple(dict.fromkeys(id(harness) for harness in bare_harnesses)) + return env_ids, harness_ids + + def _start(self, trajectories: List[Trajectory], kwargs: Dict[str, Any]) -> _Call: + unknown = set(kwargs) - _CALL_OPTIONS + if unknown: + raise ValueError(f'unknown ClientMultiTurnRollout call options: {sorted(unknown)}') + count = len(trajectories) + sampling_params = self._normalize_sampling_params(kwargs.get('sampling_params', self.sampling_params)) + manager_arg = kwargs.get('tool_manager', self.tool_manager) + harness_arg = kwargs.get('harness', self.harness) + managers = self._as_batch(manager_arg, count, 'tool_manager') + harnesses = self._as_batch(harness_arg, count, 'harness') + env_ids, harness_ids = self._borrowed_identities(managers, harnesses) + adapter_name = kwargs.get('adapter_name', self.adapter_name) + adapter_uri = kwargs.get('adapter_uri', self.adapter_uri) + execution = _RolloutExecution() + adapter = _HTTPSamplerAdapter(self.sampler, adapter_name=adapter_name, adapter_uri=adapter_uri) + core_kwargs = { + 'sampling_params': sampling_params, + 'response_callback': adapter, + 'tool_manager': manager_arg, + 'adapter_path': adapter_uri, + 'followup_fn': kwargs.get('followup_fn', self.followup_fn), + 'harness': harness_arg, + 'global_step': kwargs.get('global_step'), + '_execution': execution, + } + executor = ThreadPoolExecutor(max_workers=1) + call = _Call(execution, env_ids, harness_ids, executor) + with self._lock: + if any(identity in self._env_borrows for identity in env_ids): + executor.shutdown(wait=False) + raise ValueError('an Env is still borrowed by another active call on this rollout') + if any(identity in self._harness_borrows for identity in harness_ids): + executor.shutdown(wait=False) + raise ValueError('a bare harness is still borrowed by another active call on this rollout') + self._active[id(execution)] = call + self._env_borrows.update({identity: env for identity, env in zip(env_ids, env_ids)}) + self._harness_borrows.update({identity: harness for identity, harness in zip(harness_ids, harness_ids)}) + try: + call.future = executor.submit(self._core, trajectories, **core_kwargs) + except BaseException: + self._active.pop(id(execution), None) + for identity in env_ids: + self._env_borrows.pop(identity, None) + for identity in harness_ids: + self._harness_borrows.pop(identity, None) + executor.shutdown(wait=False) + raise + call.future.add_done_callback(lambda _future: self._finish(call)) + return call + + def _finish(self, call: _Call) -> None: + call.executor.shutdown(wait=False) + with self._lock: + self._active.pop(id(call.execution), None) + for identity in call.env_ids: + self._env_borrows.pop(identity, None) + for identity in call.harness_ids: + self._harness_borrows.pop(identity, None) def __call__(self, trajectories: List[Trajectory], **kwargs) -> List[Trajectory]: if isinstance(trajectories, dict): @@ -46,126 +200,64 @@ def __call__(self, trajectories: List[Trajectory], **kwargs) -> List[Trajectory] trajectories = list(trajectories) if not trajectories: 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), len(trajectories)) - ledgers = [ - self._open_ledger(trajectory, tool_managers[index], index) - for index, trajectory in enumerate(trajectories) - ] - stop_reasons: List[Optional[str]] = [None] * len(trajectories) - truncated = [False] * len(trajectories) - done = [False] * len(trajectories) - - for _ in range(self.max_turns): - active = [index for index, finished in enumerate(done) if not finished] - if not active: - break - responses = self.sampler.sample( - [ledgers[index].input_feature for index in active], sampling_params=sampling_params) - if len(responses) != len(active): - raise RuntimeError(f'sampler returned {len(responses)} responses for {len(active)} trajectories') - pending: List[tuple[int, List[Dict[str, Any]]]] = [] - for batch_index, trajectory_index in enumerate(active): - response = responses[batch_index] - if len(response.sequences) != 1: - raise RuntimeError('ClientMultiTurnRollout requires exactly one sequence per trajectory') - sequence = response.sequences[0] - ledger = ledgers[trajectory_index] + call = self._start(trajectories, kwargs) + try: + return call.future.result() + except (KeyboardInterrupt, SystemExit) as control_error: + call.execution.cancel() + while True: try: - ledger.record(sequence) - except RuntimeError as error: - raise RuntimeError(f'batch index {batch_index} (trajectory {trajectory_index}): {error}') from error - stop_reasons[trajectory_index] = sequence.stop_reason - if sequence.stop_reason == 'length': - truncated[trajectory_index] = True - done[trajectory_index] = True - continue - if ledger.full(): - truncated[trajectory_index] = True - done[trajectory_index] = True + call.future.result() + except (KeyboardInterrupt, SystemExit) as drain_error: + if call.future.done(): + # A completed Future may contain a control exception raised by the worker itself. + cause = drain_error if drain_error is not control_error else None + raise control_error from cause + # Repeated interrupts must not abandon a Future that is still running. continue - message = ledger.messages[-1] if ledger.messages else {} - tool_calls = list(message.get('tool_calls') or ledger.parse_tool_call(sequence.decoded or '')) - if not tool_calls: - done[trajectory_index] = True - continue - if ledger.turns >= self.max_turns: - truncated[trajectory_index] = True - stop_reasons[trajectory_index] = 'max_turns' - done[trajectory_index] = True - continue - manager = tool_managers[trajectory_index] - if manager is None: - raise ValueError(f'trajectory {trajectory_index} produced tool_calls but no tool_manager was provided') - observations = manager.call_many(tool_calls) - pending.append((trajectory_index, self._tool_messages(tool_calls, observations))) - for trajectory_index, messages in pending: - ledger = ledgers[trajectory_index] - if not ledger.observe(messages): - truncated[trajectory_index] = True - done[trajectory_index] = True - continue - self._require_resubmittable(ledger, trajectory_index) - - return [ - ledger.merge( - trajectory, - turns=ledger.turns, - stop_reason=stop_reasons[index], - truncated=truncated[index], - ) - for index, (trajectory, ledger) in enumerate(zip(trajectories, ledgers)) - ] - - def _open_ledger(self, trajectory: Trajectory, manager: Optional[ToolManager], index: int) -> TurnLedger: - ledger = TurnLedger( - self.template, - label=f'trajectory {index}', - max_tokens=self.max_trajectory_tokens, - backend='vllm', - ) - tools = manager.tool_infos() if manager is not None else list(trajectory.get('tools') or []) - ledger.open(trajectory, tools=tools) - self._require_resubmittable(ledger, index) - return ledger - - def _require_resubmittable(self, ledger: TurnLedger, index: int) -> None: - state = ledger.state - feature_ids = list(ledger.input_feature.get('input_ids') or []) - if state is None or state.media: - raise NotImplementedError( - f'trajectory {index} requires server-side RolloutState state_id for HTTP multi-turn media') - if feature_ids != list(state.engine_input_ids): - raise NotImplementedError( - f'trajectory {index} requires server-side RolloutState state_id because feature and engine ids differ') + except BaseException as drain_error: + raise control_error from drain_error + raise control_error from None - @staticmethod - def _tool_messages(tool_calls: List[Dict[str, Any]], observations: List[str]) -> List[Dict[str, Any]]: - messages = [] - for call, observation in zip(tool_calls, observations): - function = call.get('function') if isinstance(call.get('function'), dict) else {} - message: Dict[str, Any] = {'role': 'tool', 'content': str(observation)} - if call.get('id'): - message['tool_call_id'] = call['id'] - if function.get('name'): - message['name'] = function['name'] - messages.append(message) - return messages - - @staticmethod - def _as_sampling_params_dict(sampling_params) -> Optional[Dict[str, Any]]: - if sampling_params is None: - return None - if isinstance(sampling_params, dict): - return dict(sampling_params) - if dataclasses.is_dataclass(sampling_params): - return dataclasses.asdict(sampling_params) - return sampling_params + async def acall(self, trajectories: List[Trajectory], **kwargs) -> List[Trajectory]: + if isinstance(trajectories, dict): + raise TypeError('ClientMultiTurnRollout expects List[Trajectory]') + trajectories = list(trajectories) + if not trajectories: + return [] + call = self._start(trajectories, kwargs) + wrapped = asyncio.wrap_future(call.future) + task = asyncio.current_task() + cancel_error: asyncio.CancelledError | None = None + while True: + try: + result = await asyncio.shield(wrapped) + if cancel_error is not None: + raise cancel_error + return result + except asyncio.CancelledError as error: + externally_cancelled = task is not None and task.cancelling() > 0 + if cancel_error is None and externally_cancelled: + cancel_error = error + if wrapped.done(): + if cancel_error is not None: + try: + wrapped.result() + except BaseException as drain_error: + raise cancel_error from drain_error + raise cancel_error from None + raise + if cancel_error is None: + cancel_error = error + call.execution.cancel() + except BaseException as error: + if cancel_error is not None: + raise cancel_error from error + raise - @staticmethod - def _resolve_tool_managers(arg, count: int) -> List[Optional[ToolManager]]: - if isinstance(arg, list): - if len(arg) != count: - raise ValueError(f'tool_manager list length ({len(arg)}) does not match trajectories ({count})') - return list(arg) - return [arg] * count + def cancel(self) -> None: + """Stop subsequent operations in calls active at this instant.""" + with self._lock: + executions = [call.execution for call in self._active.values()] + for execution in executions: + execution.cancel() diff --git a/tests/server/integration/agentic_rl/__init__.py b/tests/server/integration/agentic_rl/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/server/integration/agentic_rl/openenv_adapter.py b/tests/server/integration/agentic_rl/openenv_adapter.py new file mode 100644 index 000000000..55394b62b --- /dev/null +++ b/tests/server/integration/agentic_rl/openenv_adapter.py @@ -0,0 +1,229 @@ +"""OpenEnv functional acceptance using session-local virtual files and remote AST execution. + +This is not an OS filesystem or microVM and does not certify AgentENV isolation. +""" + +from __future__ import annotations +import ast +import copy +import posixpath +import shlex +import threading +import uuid + +import json + +from twinkle_agentic.envs import OpenEnvClient +from twinkle_agentic.envs.base import DEFAULT_TOOLS, Env, StepResult + + +class OpenEnvTaskError(RuntimeError): + """Unconfirmed communication, execution, or verifier state must not become a low reward.""" + + +class OpenEnvTaskEnv(Env): + def __init__(self, task, *, base_url: str, env_name: str = "coding_env", client=None): + self.task = task + self._env = client or OpenEnvClient(env_name=env_name, base_url=base_url, message_timeout_s=120) + self._lock = threading.RLock() + self._uncertain: OpenEnvTaskError | None = None + self._episode_id = None + self.closed = False + + @property + def sandbox_id(self): + # Use the server's episode identity rather than fabricating an isolation receipt. + return self._episode_id + + @property + def workspace(self): + return "/workspace" + + def tools(self): + tools = copy.deepcopy(DEFAULT_TOOLS) + for tool in tools: + name = tool["function"]["name"] + if name == "run_command": + tool["function"]["description"] = ( + "Run a Python script in the virtual workspace, e.g. python smoke_test.py. " + "Only python/python3 script commands and python -c are supported; no shell or OS access." + ) + else: + tool["function"]["description"] += " This is a virtual file inside the OpenEnv session." + return tools + + def _fail(self, stage, error=None): + if self._uncertain is None: + detail = f" ({type(error).__name__})" if error is not None else "" + self._uncertain = OpenEnvTaskError(f"OpenEnv {stage} unconfirmed, session={self.sandbox_id}{detail}") + raise self._uncertain from error + + def _execute(self, source, stage): + with self._lock: + self.ensure_idle() + if self.closed: + raise OpenEnvTaskError("OpenEnv session is closed") + try: + result = self._env.execute({"code": source}) + observation = result.observation + if type(observation.exit_code) is not int: + self._fail(stage) + # The pinned smolagents timeout wrapper exits its executor context before + # returning, so a received observation is a confirmed candidate failure. + # Transport exceptions below remain unconfirmed and prevent session reuse. + return observation + except OpenEnvTaskError: + raise + except Exception as error: + # An SDK timeout does not prove the remote AST execution stopped. + self._fail(stage, error) + + def reset(self, trajectory=None): + with self._lock: + self.ensure_idle() + try: + self._env.reset() + self._episode_id = self._env.client.state().episode_id + except Exception as error: + self._fail("reset", error) + if not self.sandbox_id: + self._fail("session_identity") + result = self._execute('import json\n__twinkle_files = {}\nprint("workspace ready")', "setup") + if result.exit_code: + self._fail("setup") + return StepResult(observation="Virtual workspace ready", done=False) + + def ensure_idle(self): + # The lock drains synchronous calls; communication failures remain unconfirmed. + with self._lock: + if self._uncertain is not None: + raise self._uncertain + + def close(self): + with self._lock: + self.ensure_idle() + if not self.closed: + # Bypass best-effort OpenEnvClient.close so close failures remain visible. + try: + self._env.client.close() + except Exception as error: + self._fail("close", error) + self.closed = True + + def _path(self, value): + if not isinstance(value, str) or not value: + raise ValueError("path is required") + path = posixpath.normpath(posixpath.join(self.workspace, value)) + if not path.startswith(self.workspace + "/"): + raise ValueError("files must stay inside /workspace") + return path[len(self.workspace) + 1 :] + + def _value(self, expression, stage): + marker = "TWINKLE_VALUE_" + uuid.uuid4().hex + "=" + result = self._execute(f"import json\nprint({marker!r} + json.dumps({expression}))", stage) + if result.exit_code: + raise ValueError("virtual file is missing or cannot be read") + rows = [row[len(marker) :] for row in result.stdout.splitlines() if row.startswith(marker)] + try: + if len(rows) != 1: + raise ValueError("missing result marker") + return json.loads(rows[0]) + except (ValueError, TypeError) as error: + self._fail(stage, error) + + def read_file(self, path): + value = self._value(f"__twinkle_files[{self._path(path)!r}]", "read_file") + if not isinstance(value, str): + self._fail("read_file_result") + return value + + def _run_command(self, command): + words = shlex.split(command) + if words[:3] == ["cd", self.workspace, "&&"]: + words = words[3:] + if len(words) == 3 and words[0] in {"python", "python3"} and words[1] == "-c": + return self._execute(words[2], "run_command") + if len(words) != 2 or words[0] not in {"python", "python3"}: + raise ValueError("use python smoke_test.py or python -c; shell commands are not available") + path = self._path(words[1]) + source = self.read_file(path) + if path == "smoke_test.py": + # The AST interpreter has no file imports; execute current source with the smoke tests. + module = self.task.filename.removesuffix(".py") + source = "\n".join(line for line in source.splitlines() if not line.startswith(f"from {module} import ")) + source = self.read_file(self.task.filename) + "\n" + source + return self._execute(source + '\nprint("script completed")', "run_command") + + def step(self, tool_name, arguments=None): + arguments = arguments or {} + try: + if tool_name == "read_file": + return StepResult(observation=self.read_file(arguments.get("path")), done=False) + if tool_name == "write_file": + path = self._path(arguments.get("path")) + content = arguments.get("content", "") + if not isinstance(content, str): + content = json.dumps(content, ensure_ascii=False) + result = self._execute(f'__twinkle_files[{path!r}] = {content!r}\nprint("file written")', "write_file") + if result.exit_code: + self._fail("write_file") + elif tool_name == "run_command": + result = self._run_command(arguments.get("command", "")) + else: + raise ValueError(f"unsupported task tool {tool_name!r}") + stderr = result.stderr + if result.exit_code and stderr: + # Keep the confirmed exception, not interpreter host paths that mislead virtual-file tools. + stderr = stderr.strip().splitlines()[-1] + output = (result.stdout + "\n" + stderr).strip() + if result.exit_code: + output = f"[exit code: {result.exit_code}]\n{output}" + return StepResult(observation=output[-4096:], done=False) + except ValueError as error: + # Deterministic refusals and confirmed missing files are recoverable model errors. + return StepResult(observation=f"Error: {error}", done=False) + + def verify(self): + if not self.task.hidden_cases: + self._fail("verify_cases") + source = self.read_file(self.task.filename) + function = self.task.function + if not function.isidentifier(): + self._fail("function_name") + try: + tree = ast.parse(source) + except SyntaxError: + return 0.0 + if not any(isinstance(node, ast.FunctionDef) and node.name == function for node in tree.body): + return 0.0 + fresh_name = "twinkle_candidate_" + uuid.uuid4().hex + + class RenameCandidate(ast.NodeTransformer): + def visit_FunctionDef(self, node): + if node.name == function: + node.name = fresh_name + return self.generic_visit(node) + + def visit_Name(self, node): + if node.id == function: + node.id = fresh_name + return node + + # Use a fresh callable instead of masking a name in the AST executor's persistent state. + # Parsing/renaming is local; execution of the candidate still happens only on OpenEnv. + source = ast.unparse(RenameCandidate().visit(tree)) + loaded = self._execute(source, "verify_candidate") + if loaded.exit_code: + return 0.0 + passed = 0 + for values, expected in self.task.hidden_cases: + marker = "TWINKLE_CASE_" + uuid.uuid4().hex + "=" + result = self._execute(f"print({marker!r} + str({fresh_name}({values!r}) == {expected!r}))", "verify_case") + if result.exit_code: + # Candidate exceptions are completed failures; _execute raises transport errors. + continue + rows = [row[len(marker) :] for row in result.stdout.splitlines() if row.startswith(marker)] + if len(rows) != 1 or rows[0] not in {"True", "False"}: + self._fail("verify_result") + passed += rows[0] == "True" + return passed / len(self.task.hidden_cases) diff --git a/tests/server/integration/agentic_rl/test_openenv_adapter.py b/tests/server/integration/agentic_rl/test_openenv_adapter.py new file mode 100644 index 000000000..e3d3b1a38 --- /dev/null +++ b/tests/server/integration/agentic_rl/test_openenv_adapter.py @@ -0,0 +1,187 @@ +"""Deterministic adapter tests; real OpenEnv coverage is recorded separately.""" + +from __future__ import annotations +import ast +import threading +from types import SimpleNamespace + +import json +import pytest + +from .openenv_adapter import OpenEnvTaskEnv, OpenEnvTaskError + + +class Client: + def __init__(self): + self.client = self + self.files = {} + self.calls = [] + self.closed = False + self.failure = None + self.malformed = False + + def reset(self): + self.files.clear() + + def state(self): + return SimpleNamespace(episode_id="mock-server-episode") + + def close(self): + self.closed = True + + def execute(self, action): + source = action["code"] + self.calls.append(source) + if self.failure is not None: + raise self.failure + output, exit_code = "", 0 + tree = ast.parse(source) + for node in tree.body: + if isinstance(node, ast.Assign) and isinstance(node.targets[0], ast.Subscript): + key = ast.literal_eval(node.targets[0].slice) + self.files[key] = ast.literal_eval(node.value) + if "TWINKLE_VALUE_" in source or "TWINKLE_CASE_" in source: + call = tree.body[-1].value + marker = ast.literal_eval(call.args[0].left) + if "TWINKLE_VALUE_" in source: + key = ast.literal_eval(call.args[0].right.args[0].slice) + if key not in self.files: + exit_code = 1 + else: + output = marker + json.dumps(self.files[key]) + else: + output = marker + ("invalid" if self.malformed else "True") + return SimpleNamespace(observation=SimpleNamespace(exit_code=exit_code, stdout=output, stderr="")) + + +def task(): + return SimpleNamespace(filename="solution.py", function="solve", hidden_cases=(([1], 1), ([2], 2))) + + +def environment(): + client = Client() + env = OpenEnvTaskEnv(task(), base_url="http://not-used", client=client) + env.reset() + return env, client + + +def test_virtual_files_use_the_remote_session_and_server_identity(): + env, client = environment() + env.step("write_file", {"path": "solution.py", "content": "def solve(values):\n return values[0]"}) + assert env.read_file("/workspace/solution.py") == client.files["solution.py"] + assert env.sandbox_id == "mock-server-episode" + assert any("__twinkle_files" in source for source in client.calls) + env.close() + assert env.closed and client.closed + + +def test_local_refusals_do_not_submit_file_or_shell_operations(): + env, client = environment() + before = len(client.calls) + assert env.step("write_file", {"path": "../outside", "content": "x"}).observation.startswith("Error:") + assert env.step("run_command", {"command": "echo unsafe"}).observation.startswith("Error:") + assert len(client.calls) == before + env.ensure_idle() + env.close() + + +def test_transport_timeout_preserves_unknown_state_and_prevents_reuse(): + env, client = environment() + client.failure = TimeoutError("unknown execution") + with pytest.raises(OpenEnvTaskError): + env.step("write_file", {"path": "solution.py", "content": "x"}) + client.failure = None + count = len(client.calls) + with pytest.raises(OpenEnvTaskError): + env.ensure_idle() + with pytest.raises(OpenEnvTaskError): + env.step("read_file", {"path": "solution.py"}) + with pytest.raises(OpenEnvTaskError): + env.close() + assert len(client.calls) == count and not client.closed + + +def test_confirmed_command_failure_keeps_exception_without_host_traceback(): + env, client = environment() + stderr = 'Traceback (most recent call last):\n File "/host/deps/interpreter.py", line 1\nAssertionError: boundary case failed' + observation = SimpleNamespace(exit_code=1, stdout='before failure', stderr=stderr) + client.execute = lambda action: SimpleNamespace(observation=observation) + result = env.step('run_command', {'command': 'python -c "assert False"'}) + assert result.observation.startswith('[exit code: 1]') + assert 'AssertionError: boundary case failed' in result.observation + assert 'before failure' in result.observation + assert '/host/deps' not in result.observation + assert observation.stderr == stderr + env.ensure_idle() + env.close() + + +def test_remote_interpreter_timeout_observation_is_a_confirmed_candidate_failure(): + env, client = environment() + client.execute = lambda action: SimpleNamespace( + observation=SimpleNamespace(exit_code=1, stdout="", stderr="TimeoutError: evaluation timed out") + ) + result = env.step("run_command", {"command": 'python -c "while True: pass"'}) + assert result.observation.startswith('[exit code: 1]') + env.ensure_idle() + env.close() + + +def test_verifier_uses_fresh_callable_names_without_poisoning_smoke_state(): + env, client = environment() + env.step("write_file", {"path": "solution.py", "content": "def solve(values):\n return values[0]"}) + assert env.verify() == 1.0 + assert env.verify() == 1.0 + definitions = [source for source in client.calls if source.startswith("def twinkle_candidate_")] + assert len(definitions) == len(set(definitions)) == 2 + assert not any("solve = None" in source for source in client.calls) + env.step("write_file", {"path": "solution.py", "content": "unrelated = 1"}) + assert env.verify() == 0.0 + env.close() + + +def test_invalid_verifier_marker_is_not_a_low_reward(): + env, client = environment() + env.step("write_file", {"path": "solution.py", "content": "def solve(values):\n return values[0]"}) + client.malformed = True + with pytest.raises(OpenEnvTaskError, match="verify_result"): + env.verify() + + +def test_close_failure_is_visible_and_prevents_session_reuse(): + env, client = environment() + + def failed_close(): + raise ConnectionError("close acknowledgement lost") + + client.close = failed_close + with pytest.raises(OpenEnvTaskError, match="close"): + env.close() + with pytest.raises(OpenEnvTaskError): + env.ensure_idle() + assert not env.closed + + +def test_close_waits_for_the_actual_blocking_operation(): + env, client = environment() + started, release = threading.Event(), threading.Event() + original = client.execute + + def blocked(action): + started.set() + assert release.wait(5) + return original(action) + + client.execute = blocked + operation = threading.Thread(target=lambda: env.step("write_file", {"path": "solution.py", "content": "x"})) + operation.start() + assert started.wait(5) + closing = threading.Thread(target=env.close) + closing.start() + closing.join(0.2) + assert closing.is_alive() and not client.closed + release.set() + operation.join(5) + closing.join(5) + assert not operation.is_alive() and not closing.is_alive() + assert env.closed and client.closed diff --git a/tests/server/integration/agentic_rl/test_validation.py b/tests/server/integration/agentic_rl/test_validation.py new file mode 100644 index 000000000..286b13061 --- /dev/null +++ b/tests/server/integration/agentic_rl/test_validation.py @@ -0,0 +1,590 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from __future__ import annotations + +import asyncio +import json +import os +import subprocess +import sys +import threading +import time +from dataclasses import dataclass, replace +from types import SimpleNamespace + +import pytest + +from twinkle.data_format import SamplingParams + +from .validation import ( + _EVALUATION_TASKS, + ValidationConfig, + _CancelRecorder, + _cancel_tool_manager, + _code_task, + _evaluate, + _load_recipe, + _probe_logps, + _timed_tool_manager, + _Timeline, + _wait_for_stage, + cancel_evidence_gate, + learning_gate, + parameters_changed, + pipeline_evidence_gate, + parse_args, + result_payload, + run_cancel_validation, + run_validation, + slow_tool_isolated, + training_overlaps_rollout, + write_result, +) + + +def _config(tmp_path): + return ValidationConfig( + backend='transformers', + phase='learning', + server_url='http://127.0.0.1:8000', + model_id='ms://Qwen/Qwen3.5-4B', + template_model_id='/models/Qwen3.5-4B', + agentenv_template='python-test', + result_dir=str(tmp_path), + ) + + +def test_learning_gate_requires_updates_and_independent_improvement(): + assert learning_gate(updates=2, baseline=0.25, final=0.5, minimum_improvement=0.0) + assert not learning_gate(updates=1, baseline=0.25, final=0.5, minimum_improvement=0.0) + assert not learning_gate(updates=2, baseline=0.25, final=0.25, minimum_improvement=0.0) + assert not learning_gate(updates=2, baseline=0.25, final=0.3, minimum_improvement=0.1) + + +def test_probe_logps_uses_the_full_rollout_group_for_dp(): + recipe = _load_recipe() + calls = {} + + def forward_only(**kwargs): + calls.update(kwargs) + return SimpleNamespace(result=[-0.25, -0.5]) + + model = SimpleNamespace(forward_only=forward_only) + trajectories = [{'input_ids': [1, 2]}, {'input_ids': [3, 4]}] + result = asyncio.run(_probe_logps(recipe, model, trajectories, asyncio.Event())) + assert calls['inputs'] is trajectories + assert len(calls['inputs']) == 2 + assert result == [-0.25, -0.5] + + +def test_fake_and_real_result_modes_are_explicit(tmp_path): + config = _config(tmp_path) + fake = result_payload(config, mode='fake', status='completed', pipeline_gate=True) + real = result_payload(config, status='blocked', pipeline_gate=False) + assert fake['mode'] == 'fake' + assert real['mode'] == 'real' + assert fake['status'] != real['status'] + + +def test_result_file_contains_no_server_url_or_credentials(tmp_path): + path = write_result(_config(tmp_path), status='blocked', pipeline_gate=False, learning_gate=None) + payload = json.loads(path.read_text(encoding='utf-8')) + assert payload['mode'] == 'real' + assert 'server_url' not in payload['config'] + assert 'api_key' not in payload + assert 'TWINKLE_SERVER_TOKEN' not in payload + + +def _pipeline_evidence(): + return { + 'two_updates': True, + 'parameter_change': True, + 'multi_turn_feedback': True, + 'token_alignment': True, + 'distinct_snapshots': True, + 'new_policy_requests': True, + 'sync_async_equivalent': True, + 'slow_tool_isolated': True, + 'training_overlap': True, + 'environment_isolation': True, + 'cleanup_confirmed': True, + } + + +def test_pipeline_gate_fails_when_any_required_evidence_is_missing(): + evidence = _pipeline_evidence() + assert pipeline_evidence_gate(evidence) + for key in evidence: + incomplete = dict(evidence) + incomplete[key] = False + assert not pipeline_evidence_gate(incomplete), key + + +def test_cancel_gate_rejects_early_env_cancel_new_work_and_cleanup_failure(): + assert not cancel_evidence_gate({ + 'stage_reached': False, + 'cancelled': True, + 'post_cancel_operations': [], + 'cleanup_issues': [], + }) + assert not cancel_evidence_gate({ + 'stage_reached': True, + 'cancelled': True, + 'post_cancel_operations': ['sample'], + 'cleanup_issues': [], + }) + assert not cancel_evidence_gate({ + 'stage_reached': True, + 'cancelled': True, + 'post_cancel_operations': [], + 'cleanup_issues': ['unknown execution'], + }) + + +def test_learning_without_frozen_threshold_is_blocked(tmp_path): + with pytest.raises(ValueError, match='frozen'): + asyncio.run(run_validation(_config(tmp_path))) + + +def test_cli_template_is_set_before_recipe_import(tmp_path, monkeypatch): + config = replace(_config(tmp_path), min_success_improvement=0.1) + monkeypatch.setenv('AGENTENV_TEMPLATE', 'stale-template') + seen = [] + + def stop_at_import(): + seen.append(os.environ['AGENTENV_TEMPLATE']) + raise RuntimeError('stop before provider setup') + + monkeypatch.setattr('tests.server.integration.agentic_rl.validation._load_recipe', stop_at_import) + with pytest.raises(RuntimeError, match='stop before'): + asyncio.run(run_validation(config)) + with pytest.raises(RuntimeError, match='stop before'): + asyncio.run(run_cancel_validation(config)) + assert seen == [config.agentenv_template, config.agentenv_template] + + +def test_evaluation_binds_held_out_tasks_and_distinct_attempt_seeds(tmp_path): + config = _config(tmp_path) + created = [] + seeds = [] + + @dataclass(frozen=True) + class CodeTask: + filename: str + function: str + solution: str + smoke_test: str + prompt: str + hidden_cases: tuple + + class Recipe: + NUM_GENERATIONS = 4 + ADAPTER_NAME = 'adapter' + + @staticmethod + def _ensure_running(stop_event): + return None + + @staticmethod + async def _await_operation(operation, *args, stop_event, on_result=None, **kwargs): + result = operation(*args, **kwargs) + if on_result is not None: + on_result(result) + return result + + @staticmethod + def _create_attempt(task, on_created): + created.append(task) + env = SimpleNamespace(verify=lambda: 1.0) + on_created(env) + return SimpleNamespace(env=env, manager=object(), trajectory={'messages': []}) + + @staticmethod + def _validate_training_rows(rows): + assert len(rows) == 4 + + class Rollout: + sampling_params = SamplingParams() + + async def acall(self, trajectories, *, sampling_params, **kwargs): + assert len(trajectories) == 1 + seeds.append(sampling_params.seed) + return [{}] + + Recipe.CodeTask = CodeTask + rewards, evaluations = asyncio.run(_evaluate(Recipe(), Rollout(), 'policy-0', [], asyncio.Event(), config)) + assert len(rewards) == 16 + assert len(seeds) == len(set(seeds)) == 16 + assert {task.filename for task in created} == {'count_even.py', 'last_negative.py'} + assert all(task.function == task.filename[:-3] for task in created) + assert {item['task_id'] for item in evaluations} == {'count-even', 'last-negative'} + + +_CORRECT_SOLUTIONS = { + 'count-even': 'def count_even(values):\n return sum(value % 2 == 0 for value in values)\n', + 'last-negative': ('def last_negative(values):\n negatives = [value for value in values if value < 0]\n' + ' return negatives[-1] if negatives else None\n'), +} + + +def _verify_locally(recipe, task, source, tmp_path): + env = object.__new__(recipe._TaskAgentEnv) + env.task = task + env._workspace = '/workspace' + uploads = {} + env._sandbox = SimpleNamespace( + sandbox_id='sandbox-eval', files=SimpleNamespace(write=uploads.__setitem__), set_timeout=lambda timeout: None) + env._task_infrastructure_error = None + env._execution_lock = threading.RLock() + env._active_commands = 0 + env._uncertain_execution = None + env._command_timeout = 30 + env._refresh_timeout = False + (tmp_path / task.filename).write_text(source, encoding='utf-8') + + def execute(command, cwd, timeout): + script = tmp_path / 'verify.py' + script.write_text(uploads.popitem()[1], encoding='utf-8') + result = subprocess.run([sys.executable, str(script)], cwd=tmp_path, capture_output=True, text=True) + return result.returncode, result.stdout + result.stderr + + env._execute = execute + return env.verify() + + +@pytest.mark.parametrize('item', _EVALUATION_TASKS, ids=[item['id'] for item in _EVALUATION_TASKS]) +def test_held_out_verifier_scores_its_own_file_function_and_cases(item, tmp_path): + recipe = _load_recipe() + task = _code_task(recipe, item) + assert task.filename != recipe.TASK.filename + assert _verify_locally(recipe, task, _CORRECT_SOLUTIONS[item['id']], tmp_path) == 1.0 + assert _verify_locally(recipe, task, task.solution, tmp_path) < 1.0 + + +def test_parameter_change_requires_aligned_finite_probe_shift(): + assert parameters_changed([-1.0, -2.0], [-1.0, -2.1]) + assert not parameters_changed([-1.0, -2.0], [-1.0, -2.0]) + assert not parameters_changed([], []) + assert not parameters_changed(None, [-1.0]) + assert not parameters_changed([-1.0], [-1.0, -2.0]) + assert not parameters_changed([-1.0], [float('nan')]) + + +def test_timeline_evidence_requires_peer_progress_and_training_overlap(): + slow = {'kind': 'tool', 'start': 1.0, 'end': 5.0, 'slow': True} + inside = {'kind': 'sample', 'start': 2.0, 'end': 3.0} + straddling = {'kind': 'sample', 'start': 4.0, 'end': 6.0} + assert slow_tool_isolated([slow, inside]) + assert not slow_tool_isolated([slow, straddling]) + assert not slow_tool_isolated([dict(slow, slow=False), inside]) + + train = {'kind': 'train', 'start': 10.0, 'end': 12.0} + assert training_overlaps_rollout([train, {'kind': 'tool', 'start': 11.0, 'end': 13.0}]) + assert not training_overlaps_rollout([train, {'kind': 'sample', 'start': 12.0, 'end': 13.0}]) + + +def test_timed_tool_manager_slows_only_the_first_batch(): + class Base: + def call_many(self, tool_calls): + return ['ok'] * len(tool_calls) + + timeline = _Timeline() + manager = _timed_tool_manager(Base, timeline, 0.05)() + assert manager.call_many([{}]) == ['ok'] + assert manager.call_many([{}, {}]) == ['ok', 'ok'] + assert [span['slow'] for span in timeline.spans] == [True, False] + assert timeline.spans[0]['end'] - timeline.spans[0]['start'] >= 0.05 + + +def test_cancel_tool_stage_ignores_setup_writes_and_counts_batches_after_cancel(): + class Env: + steps = 0 + + def step(self, tool_name, arguments=None): + type(self).steps += 1 + + class Base: + def call_many(self, tool_calls): + return ['ok'] * len(tool_calls) + + recorder = _CancelRecorder('tool') + recorder.release.set() + manager = _cancel_tool_manager(Base, recorder)() + Env().step('write_file', {'path': 'solution.py'}) + assert recorder.operations == [] + manager.call_many([{}, {}]) + assert recorder.started.is_set() + recorder.mark_cancelled() + manager.call_many([{}]) + assert recorder.operations == ['tool', 'tool'] + assert recorder.post_cancel_operations == ['tool'] + + +def test_wait_for_stage_returns_without_helper_threads(): + recorder = _CancelRecorder('backward') + + async def finished_early(): + return None + + async def never_reaches_stage(): + await asyncio.sleep(10) + + async def run(): + early = asyncio.create_task(finished_early()) + await asyncio.sleep(0) + assert not await _wait_for_stage(recorder, early, timeout=5) + slow = asyncio.create_task(never_reaches_stage()) + started = time.monotonic() + assert not await _wait_for_stage(recorder, slow, timeout=0.2) + assert time.monotonic() - started < 2 + slow.cancel() + + threads = threading.active_count() + asyncio.run(run()) + assert threading.active_count() <= threads + + +@pytest.mark.parametrize('stage', ['sample', 'tool', 'verify', 'backward', 'step', 'save', 'unreached']) +def test_cancel_runner_drains_real_recipe_call_sites_without_a_provider(stage, tmp_path, monkeypatch): + from . import validation + + recipe = _load_recipe() + created = [] + + class Env: + def __init__(self, task, **kwargs): + self.task = task + self.sandbox_id = f'fake-{len(created)}' + self.reward = len(created) % 2 if stage != 'unreached' else 0.0 + self.closed = False + created.append(self) + + def reset(self): + return None + + def step(self, *args, **kwargs): + return SimpleNamespace(info={}) + + def verify(self): + return float(self.reward) + + def ensure_idle(self): + return None + + def close(self): + self.closed = True + + class Manager: + def __init__(self, tools): + self.env = tools[0].env + + def tool_infos(self): + return [] + + def call_many(self, calls): + for _ in calls: + self.env.step('write_file', {'path': 'solution.py'}) + return ['ok'] * len(calls) + + class Model: + def forward_backward(self, **kwargs): + return None + + def clip_grad_and_step(self, **kwargs): + return None + + def calculate_metric(self, **kwargs): + return {'result': {}} + + def save(self, name, **kwargs): + return SimpleNamespace(twinkle_path=f'twinkle://{name}') + + sampler = SimpleNamespace(sample=lambda: None) + client = SimpleNamespace(closed=False) + client.close = lambda: setattr(client, 'closed', True) + + class Rollout: + async def acall(self, trajectories, tool_manager, **kwargs): + stop_event = kwargs_stop_event + await recipe._await_operation(sampler.sample, stop_event=stop_event) + for manager in tool_manager: + recipe._ensure_running(stop_event) + await recipe._await_operation(manager.call_many, [{}, {}], stop_event=stop_event) + return [{ + 'response_token_ids': [1], 'rollout_logprobs': [-0.1], 'truncated': False, 'stop_reason': 'stop' + } for _ in trajectories] + + kwargs_stop_event = asyncio.Event() + + async def build(*args): + return [client], Model(), sampler, Rollout(), kwargs_stop_event + + monkeypatch.setattr(recipe, '_TaskAgentEnv', Env) + monkeypatch.setattr(recipe, 'EnvTool', SimpleNamespace( + from_schemas=lambda env, schemas: [SimpleNamespace(env=env) for _ in schemas])) + monkeypatch.setattr(recipe, 'ToolManager', Manager) + monkeypatch.setattr(validation, '_load_recipe', lambda: recipe) + monkeypatch.setattr(validation, '_build_runtime', build) + monkeypatch.setattr(validation, '_STAGE_TIMEOUT', 2.0) + config = replace(_config(tmp_path), phase='cancel', cancel_stage='backward' if stage == 'unreached' else stage) + result = asyncio.run(run_cancel_validation(config)) + assert result['pipeline_gate'] is (stage != 'unreached') + assert result['post_cancel_operations'] == [] + assert result['stage_reached'] is (stage != 'unreached') + assert len(created) == 4 and all(env.closed for env in created) + assert client.closed + + +def test_openenv_attempt_keeps_virtual_workspace_tool_schemas(tmp_path, monkeypatch): + from . import validation + from . import openenv_adapter + + recipe = _load_recipe() + + class Env: + def __init__(self, task, **kwargs): + self.task = task + + def reset(self): + return None + + def ensure_idle(self): + return None + + def step(self, *args, **kwargs): + return None + + def tools(self): + return recipe.DEFAULT_TOOLS + + monkeypatch.setattr(validation, '_load_recipe', lambda: recipe) + monkeypatch.setattr(openenv_adapter, 'OpenEnvTaskEnv', Env) + config = replace(_config(tmp_path), environment_backend='openenv', openenv_url='http://not-used') + prepared = validation._prepare_recipe(config) + attempt = prepared._create_attempt(prepared.TASK, lambda env: None) + functions = {info['function']['name']: info['function'] for info in attempt.trajectory['tools']} + assert functions['write_file']['parameters']['required'] == ['path', 'content'] + assert functions['run_command']['parameters']['required'] == ['command'] + assert functions['read_file']['parameters']['properties']['path']['type'] == 'string' + prompt = attempt.trajectory['messages'][0]['content'] + assert '/workspace' in prompt and 'python smoke_test.py' in prompt + assert 'workspace modules' in prompt and 'without repeated reruns' in prompt + + +@pytest.mark.parametrize('invalid', [None, {'truncated': True}, {'stop_reason': 'abort'}, {'rollout_logprobs': []}]) +def test_sync_async_parity_requires_valid_complete_rollouts(tmp_path, monkeypatch, invalid): + from . import validation + recipe = _load_recipe() + created = [] + + def create(task, on_created): + env = SimpleNamespace(sandbox_id=f'session-{len(created)}') + created.append(env) + on_created(env) + return SimpleNamespace(env=env, manager=None, trajectory={}) + + monkeypatch.setattr(recipe, '_create_attempt', create) + row = {'response_token_ids': [1], 'rollout_logprobs': [-0.1], 'stop_reason': 'stop', 'truncated': False} + if invalid: + row.update(invalid) + + class Rollout: + sampling_params = validation.SamplingParams(max_tokens=32) + + def __call__(self, *args, **kwargs): + return [dict(row)] + + async def acall(self, *args, **kwargs): + return [dict(row)] + + operation = validation._sync_async_evidence(recipe, Rollout(), 'policy', [], asyncio.Event(), _config(tmp_path)) + if invalid: + with pytest.raises(RuntimeError): + asyncio.run(operation) + else: + result = asyncio.run(operation) + assert result['sync_async_equivalent'] and result['environment_isolation'] + + +def test_openenv_cli_and_results_cannot_claim_agentenv_security(tmp_path): + config = parse_args([ + '--backend', 'transformers', '--phase', 'pipeline', '--result-dir', str(tmp_path), + '--environment-backend', 'openenv', '--openenv-url', 'http://127.0.0.1:18762', + ]) + assert config.environment_backend == 'openenv' + assert config.gradient_accumulation_steps == 2 + payload = result_payload(config, pipeline_gate=True) + assert payload['acceptance_scope'] == 'openenv_functional' + assert payload['agentenv_security_gate'] is None + assert 'openenv_url' not in payload['config'] + + +def test_cancel_result_paths_keep_every_stage(tmp_path): + config = replace(_config(tmp_path), phase='cancel') + sample = write_result(replace(config, cancel_stage='sample'), pipeline_gate=False) + tool = write_result(replace(config, cancel_stage='tool'), pipeline_gate=False) + assert sample != tool and sample.exists() and tool.exists() + + +@pytest.mark.parametrize('changed', [True, False]) +def test_update_gate_counts_only_complete_ga_cycles_with_probe_changes(changed, tmp_path, monkeypatch): + from . import validation + + recipe = _load_recipe() + calls = [] + + class Model: + micro_steps = 0 + + def forward_backward(self, **kwargs): + calls.append('backward') + + def clip_grad_and_step(self, **kwargs): + self.micro_steps += 1 + calls.append('step') + + def calculate_metric(self, **kwargs): + return {'result': {'grad_norm': 1.0}} + + def forward_only(self, **kwargs): + assert kwargs['compute_loss'] is False + completed = self.micro_steps // 2 if changed else 0 + return {'result': {'logps': [[-1.0 - completed * 0.1]]}} + + def save(self, name, **kwargs): + calls.append(name) + return SimpleNamespace(twinkle_path='twinkle://' + name) + + model = Model() + client = SimpleNamespace(close=lambda: None) + sampler = SimpleNamespace(sample=lambda *args, **kwargs: []) + rollout = SimpleNamespace(sampler=sampler, _core=SimpleNamespace(response_callback=SimpleNamespace())) + + async def build(*args): + return [client], model, sampler, rollout, asyncio.Event() + + async def groups(*args): + return [{'response_token_ids': [1], 'rollout_logprobs': [-0.1]}] * 4, [1.0, 0.0, 1.0, 0.0] + + async def sync_async(*args): + return {'sync_async_equivalent': False, 'environment_isolation': False} + + monkeypatch.setattr(recipe, '_rollout_group', groups) + monkeypatch.setattr(validation, '_load_recipe', lambda: recipe) + monkeypatch.setattr(validation, '_build_runtime', build) + monkeypatch.setattr(validation, '_sync_async_evidence', sync_async) + config = replace(_config(tmp_path), phase='pipeline', max_groups=2) + result = asyncio.run(run_validation(config)) + assert calls.count('step') == 4 + assert result['updates'] == (2 if changed else 0) + assert result['evidence']['two_updates'] is changed + assert [call for call in calls if call.startswith('policy')] == ( + ['policy-0', 'policy-1', 'policy-2'] if changed else ['policy-0']) + + +@pytest.mark.skipif( + not os.environ.get('AGENTENV_TEMPLATE') or not os.environ.get('TWINKLE_SERVER_URL'), + reason='real AgentENV and Twinkle Server are not configured', +) +def test_real_agentenv_configuration_is_present(): + assert os.environ['AGENTENV_TEMPLATE'] + assert os.environ['TWINKLE_SERVER_URL'].startswith(('http://', 'https://')) diff --git a/tests/server/integration/agentic_rl/validation.py b/tests/server/integration/agentic_rl/validation.py new file mode 100644 index 000000000..883023a7a --- /dev/null +++ b/tests/server/integration/agentic_rl/validation.py @@ -0,0 +1,780 @@ +"""Real multi-turn GRPO acceptance with AgentENV or explicitly scoped OpenEnv functional coverage.""" +from __future__ import annotations + +import argparse +import asyncio +import contextlib +import importlib.util +import json +import math +import os +import sys +import threading +import time +from dataclasses import asdict, dataclass, replace +from pathlib import Path + +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_client.rollout import ClientMultiTurnRollout + +_EVALUATION_TASKS = [ + { + 'id': 'count-even', + 'filename': 'count_even.py', + 'function': 'count_even', + 'solution': 'def count_even(values):\n return sum(value % 2 == 1 for value in values)\n', + 'smoke_test': 'from count_even import count_even\nassert count_even([1, 2, 4, 5]) == 2\n', + 'prompt': 'The workspace has count_even.py and smoke_test.py. Run the test, fix the odd/even boundary bug, and rerun it.', + 'hidden_cases': (([], 0), ([2, 4], 2), ([1, 3], 0), ([-2, 0, 7], 2)), + }, + { + 'id': 'last-negative', + 'filename': 'last_negative.py', + 'function': 'last_negative', + 'solution': 'def last_negative(values):\n negatives = [value for value in values if value < 0]\n return negatives[0] if negatives else None\n', + 'smoke_test': 'from last_negative import last_negative\nassert last_negative([-1, -2, 3]) == -2\n', + 'prompt': 'The workspace has last_negative.py and smoke_test.py. Run the test, return the last negative correctly, and rerun it.', + 'hidden_cases': (([-1, -2, 3], -2), ([1, 2], None), ([-4], -4), ([0, -5, -6], -6)), + }, +] +_STAGE_TIMEOUT = 300.0 + + +@dataclass(frozen=True) +class ValidationConfig: + backend: str + phase: str + server_url: str + model_id: str + template_model_id: str + agentenv_template: str + result_dir: str + max_groups: int = 4 + active_groups: int = 2 + seed: int = 7 + eval_repeats: int = 2 + min_success_improvement: float | None = None + cancel_stage: str = 'sample' + slow_tool_seconds: float = 20.0 + environment_backend: str = 'agentenv' + openenv_url: str | None = None + openenv_name: str = 'coding_env' + gradient_accumulation_steps: int = 2 + + +def _load_recipe(): + path = Path(__file__).parents[4] / 'cookbook/client/async_rl/multi_turn_grpo.py' + spec = importlib.util.spec_from_file_location('agentic_rl_validation_recipe', path) + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module + + +def _prepare_recipe(config): + if config.environment_backend == 'agentenv': + os.environ['AGENTENV_TEMPLATE'] = config.agentenv_template + recipe = _load_recipe() + if config.environment_backend == 'openenv': + from tests.server.integration.agentic_rl.openenv_adapter import OpenEnvTaskEnv + + def create(task, on_created): + env = OpenEnvTaskEnv(task, base_url=config.openenv_url, env_name=config.openenv_name) + on_created(env) + env.reset() + for path, content in [(task.filename, task.solution), ('smoke_test.py', task.smoke_test)]: + env.step('write_file', {'path': path, 'content': content}) + env.ensure_idle() + manager = recipe.ToolManager(recipe.EnvTool.from_schemas(env, env.tools())) + return recipe._Attempt(env, manager, { + 'messages': [{'role': 'user', 'content': task.prompt + + '\nAll task files are virtual files under /workspace. Use read_file/write_file with ' + 'relative file names or /workspace paths. Run tests with python smoke_test.py, ' + 'which loads the current task source automatically. Shell commands, OS imports, ' + 'and importing workspace modules with python -c are not available. ' + 'After the test passes, give a short final answer without repeated reruns.'}], + 'tools': manager.tool_infos(), + }) + + recipe._TaskAgentEnv = OpenEnvTaskEnv + recipe._create_attempt = create + return recipe + + +def learning_gate(*, updates: int, baseline: float, final: float, minimum_improvement: float) -> bool: + return updates >= 2 and final - baseline > minimum_improvement + + +def pipeline_evidence_gate(evidence: dict) -> bool: + required = ( + 'two_updates', 'parameter_change', 'multi_turn_feedback', 'token_alignment', 'distinct_snapshots', + 'new_policy_requests', 'sync_async_equivalent', 'slow_tool_isolated', 'training_overlap', + 'environment_isolation', 'cleanup_confirmed', + ) + return all(evidence.get(key) is True for key in required) + + +def cancel_evidence_gate(evidence: dict) -> bool: + return ( + evidence.get('stage_reached') is True + and evidence.get('cancelled') is True + and not evidence.get('post_cancel_operations') + and not evidence.get('cleanup_issues') + ) + + +def result_payload(config: ValidationConfig, *, mode: str = 'real', **values) -> dict: + return { + 'mode': mode, + 'backend': config.backend, + 'phase': config.phase, + 'environment_backend': config.environment_backend, + 'acceptance_scope': ('openenv_functional' if config.environment_backend == 'openenv' else 'agentenv'), + 'agentenv_security_gate': None, + 'config': { + key: value + for key, value in asdict(config).items() + if key not in {'server_url', 'result_dir', 'openenv_url'} + }, + **values, + } + + +def write_result(config: ValidationConfig, **values) -> Path: + destination = Path(config.result_dir) + destination.mkdir(parents=True, exist_ok=True) + suffix = f'-{config.cancel_stage}' if config.phase == 'cancel' else '' + path = destination / f'{config.backend}-{config.environment_backend}-{config.phase}{suffix}.json' + path.write_text(json.dumps(result_payload(config, **values), indent=2, sort_keys=True) + '\n', encoding='utf-8') + return path + + +def _code_task(recipe, item: dict): + return recipe.CodeTask(**{key: value for key, value in item.items() if key != 'id'}) + + +def _training_evidence(rolled: list[dict]) -> dict[str, bool]: + multi_turn = False + aligned = True + for trajectory in rolled: + messages = trajectory.get('messages') or [] + for index, message in enumerate(messages): + if message.get('role') == 'tool' and any( + later.get('role') == 'assistant' for later in messages[index + 1:]): + multi_turn = True + input_ids = list(trajectory.get('input_ids') or []) + labels = list(trajectory.get('labels') or []) + masks = list(trajectory.get('completion_mask') or []) + response_ids = list(trajectory.get('response_token_ids') or []) + old_logps = list(trajectory.get('rollout_logprobs') or []) + if not (len(input_ids) == len(labels) == len(masks) and response_ids and len(response_ids) == len(old_logps)): + aligned = False + if not all(math.isfinite(value) for value in old_logps): + aligned = False + policy_positions = sum(label != -100 and bool(mask) for label, mask in zip(labels, masks, strict=False)) + if policy_positions != len(old_logps): + aligned = False + return {'multi_turn_feedback': multi_turn, 'token_alignment': aligned} + + +def _numbers(value) -> list[float]: + if isinstance(value, bool): + return [] + if isinstance(value, (int, float)): + return [float(value)] + if isinstance(value, dict): + return _numbers(value.get('logps')) + if isinstance(value, (list, tuple)): + return [number for item in value for number in _numbers(item)] + return [] + + +def parameters_changed(before: list[float] | None, after: list[float] | None, *, tolerance: float = 1e-6) -> bool: + """Same probe, same token positions: any log-prob shift beyond tolerance proves the adapter moved.""" + if not before or not after or len(before) != len(after): + return False + if not all(math.isfinite(value) for value in before + after): + return False + return max(abs(left - right) for left, right in zip(before, after, strict=True)) > tolerance + + +def slow_tool_isolated(spans: list[dict]) -> bool: + """A peer generation must start and finish while the injected slow tool batch is still running.""" + for slow in (span for span in spans if span['kind'] == 'tool' and span.get('slow')): + if any(span['kind'] == 'sample' and span['start'] >= slow['start'] and span['end'] <= slow['end'] + for span in spans): + return True + return False + + +def training_overlaps_rollout(spans: list[dict]) -> bool: + training = [span for span in spans if span['kind'] == 'train'] + rollout = [span for span in spans if span['kind'] in {'sample', 'tool'}] + return any(item['start'] < train['end'] and train['start'] < item['end'] for train in training for item in rollout) + + +class _Timeline: + + def __init__(self): + self._lock = threading.Lock() + self._origin = time.monotonic() + self.spans: list[dict] = [] + + @contextlib.contextmanager + def span(self, kind: str, **labels): + start = time.monotonic() - self._origin + try: + yield + finally: + end = time.monotonic() - self._origin + with self._lock: + self.spans.append({'kind': kind, 'start': start, 'end': end, **labels}) + + +class _RecordingSampler: + + def __init__(self, sampler, timeline: _Timeline): + self._sampler = sampler + self._timeline = timeline + self._lock = threading.Lock() + self.uris = [] + + def __getattr__(self, name): + return getattr(self._sampler, name) + + def sample(self, *args, **kwargs): + with self._lock: + self.uris.append(kwargs.get('adapter_uri')) + with self._timeline.span('sample'): + return self._sampler.sample(*args, **kwargs) + + +class _TimedModel: + + def __init__(self, model, timeline: _Timeline): + self._model = model + self._timeline = timeline + + def __getattr__(self, name): + return getattr(self._model, name) + + def forward_backward(self, *args, **kwargs): + with self._timeline.span('train', operation='backward'): + return self._model.forward_backward(*args, **kwargs) + + def clip_grad_and_step(self, *args, **kwargs): + with self._timeline.span('train', operation='step'): + return self._model.clip_grad_and_step(*args, **kwargs) + + +def _timed_tool_manager(base, timeline: _Timeline, slow_seconds: float): + """Time every rollout tool batch and hold the first one for ``slow_seconds``.""" + lock = threading.Lock() + pending = {'slow': slow_seconds > 0} + + class TimedToolManager(base): + + def call_many(self, tool_calls): + with lock: + slow, pending['slow'] = pending['slow'], False + with timeline.span('tool', manager=id(self), slow=slow): + if slow: + time.sleep(slow_seconds) + return super().call_many(tool_calls) + + return TimedToolManager + + +async def _build_runtime(config, recipe): + stop_event = asyncio.Event() + clients = [] + try: + client = await recipe._await_operation( + init_twinkle_client, + base_url=config.server_url, + api_key=os.environ.get('TWINKLE_SERVER_TOKEN', 'EMPTY_TOKEN'), + stop_event=stop_event, + on_result=clients.append, + ) + model = await recipe._await_operation(client.model, config.model_id, stop_event=stop_event) + sampler = await recipe._await_operation(client.sampler, config.model_id, stop_event=stop_event) + await recipe._await_operation( + model.add_adapter_to_model, + recipe.ADAPTER_NAME, + LoraConfig(target_modules='all-linear', r=8, lora_alpha=32), + gradient_accumulation_steps=config.gradient_accumulation_steps, + stop_event=stop_event, + ) + await recipe._await_operation(model.set_loss, 'GRPOLoss', epsilon=0.2, beta=0.0, stop_event=stop_event) + # Adam is accepted by both Transformers and Megatron training backends. + await recipe._await_operation(model.set_optimizer, 'Adam', lr=2e-5, stop_event=stop_event) + await recipe._await_operation(model.set_processor, 'InputProcessor', stop_event=stop_event) + await recipe._await_operation( + model.set_template, 'Qwen3_5Template', model_id=config.template_model_id, stop_event=stop_event) + await recipe._await_operation( + sampler.set_template, 'Qwen3_5Template', model_id=config.template_model_id, stop_event=stop_event) + template = await recipe._await_operation( + Qwen3_5Template, + model_id=config.template_model_id, + max_length=8192, + enable_thinking=False, + stop_event=stop_event, + ) + template.truncation_strategy = 'delete' + rollout = ClientMultiTurnRollout( + sampler, + template, + sampling_params=SamplingParams(max_tokens=512, temperature=1.0, top_p=0.95, logprobs=1), + max_turns=recipe.MAX_TURNS, + ) + return clients, model, sampler, rollout, stop_event + except BaseException: + await recipe._drain_cleanup(clients, [], set(), stop_event, propagate_cancel=False) + raise + + +async def _initial_policy(recipe, model, stop_event, snapshots): + initial = await recipe._await_operation( + model.save, 'policy-0', save_optimizer=False, stop_event=stop_event, on_result=snapshots.append) + return recipe._checkpoint_path(initial) + + +async def _evaluate_group(recipe, rollout, policy_uri, environments, stop_event, task, seeds): + """Score one held-out task; each attempt gets its own seed so a group is not one sample repeated.""" + attempts = [] + for _ in seeds: + recipe._ensure_running(stop_event) + await recipe._await_operation( + recipe._create_attempt, task, environments.append, stop_event=stop_event, on_result=attempts.append) + results = await asyncio.gather( + *[ + rollout.acall( + [attempt.trajectory], + tool_manager=[attempt.manager], + adapter_name=recipe.ADAPTER_NAME, + adapter_uri=policy_uri, + sampling_params=replace(rollout.sampling_params, seed=seed), + ) for attempt, seed in zip(attempts, seeds, strict=True) + ], + return_exceptions=True, + ) + for result in results: + if isinstance(result, BaseException): + raise result + recipe._ensure_running(stop_event) + recipe._validate_training_rows([result[0] for result in results]) + rewards = [] + for attempt in attempts: + recipe._ensure_running(stop_event) + rewards.append(await recipe._await_operation(attempt.env.verify, stop_event=stop_event)) + return rewards + + +async def _evaluate(recipe, rollout, policy_uri, environments, stop_event, config): + rewards = [] + evaluations = [] + for task_index, item in enumerate(_EVALUATION_TASKS): + task = _code_task(recipe, item) + for repeat in range(config.eval_repeats): + first = config.seed + task_index * 1000 + repeat * recipe.NUM_GENERATIONS + seeds = list(range(first, first + recipe.NUM_GENERATIONS)) + group_rewards = await _evaluate_group(recipe, rollout, policy_uri, environments, stop_event, task, seeds) + rewards.extend(group_rewards) + evaluations.append({'task_id': item['id'], 'repeat': repeat, 'seeds': seeds, 'attempts': len(group_rewards)}) + return rewards, evaluations + + +async def _sync_async_evidence(recipe, rollout, policy_uri, environments, stop_event, config): + task = _code_task(recipe, _EVALUATION_TASKS[0]) + attempts = [] + for _ in range(2): + await recipe._await_operation( + recipe._create_attempt, task, environments.append, stop_event=stop_event, on_result=attempts.append) + params = replace(rollout.sampling_params, seed=config.seed + 50000) + common = { + 'adapter_name': recipe.ADAPTER_NAME, + 'adapter_uri': policy_uri, + 'sampling_params': params, + } + sync = rollout([attempts[0].trajectory], tool_manager=[attempts[0].manager], **common)[0] + async_result = await rollout.acall([attempts[1].trajectory], tool_manager=[attempts[1].manager], **common) + async_result = async_result[0] + recipe._validate_training_rows([sync, async_result]) + keys = ('input_ids', 'labels', 'completion_mask', 'rollout_logprobs', 'messages') + equivalent = all(sync.get(key) == async_result.get(key) for key in keys) + sandbox_ids = {attempt.env.sandbox_id for attempt in attempts} + return { + 'sync_async_equivalent': equivalent, + 'environment_isolation': len(sandbox_ids) == 2 and None not in sandbox_ids, + 'sync_stop_reason': sync.get('stop_reason'), + 'async_stop_reason': async_result.get('stop_reason'), + } + + +async def _probe_logps(recipe, model, trajectories, stop_event) -> list[float]: + response = await recipe._await_operation( + model.forward_only, inputs=trajectories, compute_loss=False, stop_event=stop_event) + result = response['result'] if isinstance(response, dict) else response.result + return _numbers(result) + + +async def run_validation(config: ValidationConfig) -> dict: + if config.phase == 'learning' and config.min_success_improvement is None: + raise ValueError('learning validation requires a frozen --min-success-improvement threshold') + recipe = _prepare_recipe(config) + timeline = _Timeline() + recipe.ToolManager = _timed_tool_manager(recipe.ToolManager, timeline, config.slow_tool_seconds) + clients, model, sampler, rollout, stop_event = await _build_runtime(config, recipe) + recording_sampler = _RecordingSampler(sampler, timeline) + rollout.sampler = recording_sampler + rollout._core.response_callback._sampler = recording_sampler + model = _TimedModel(model, timeline) + environments, snapshots = [], [] + active_tasks = set() + values = {'learning_gate': None} + try: + policy_uri = await _initial_policy(recipe, model, stop_event, snapshots) + initial_uri = policy_uri + baseline_rewards = baseline_evaluations = None + if config.phase == 'learning': + baseline_rewards, baseline_evaluations = await _evaluate( + recipe, rollout, policy_uri, environments, stop_event, config) + advantage_fn = GRPOAdvantage() + metrics = [] + all_rolled = [] + probe = {} + + async def train_groups(): + nonlocal policy_uri + updates = 0 + next_group = 0 + tasks = set() + while next_group < config.max_groups or tasks: + while next_group < config.max_groups and len(tasks) < config.active_groups and not stop_event.is_set(): + task = asyncio.create_task(recipe._rollout_group(rollout, policy_uri, environments, stop_event)) + tasks.add(task) + active_tasks.add(task) + next_group += 1 + if not tasks: + break + done, tasks = await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED) + for task in done: + active_tasks.discard(task) + rolled, rewards = task.result() + all_rolled.extend(rolled) + recipe._ensure_running(stop_event) + if 'trajectories' not in probe: + probe['trajectories'] = rolled + probe['before'] = await _probe_logps(recipe, model, rolled, stop_event) + probe['latest'] = probe['before'] + probe['updates'] = [] + updated = False + # Finish all accumulation micro-steps before counting an optimizer update and saving. + for _ in range(config.gradient_accumulation_steps): + updated = await recipe._train_group( + model, rolled, rewards, advantage_fn, stop_event, on_metrics=metrics.append) + if not updated: + break + if updated: + after = await _probe_logps(recipe, model, probe['trajectories'], stop_event) + changed = parameters_changed(probe['latest'], after) + probe['updates'].append({'changed': changed, 'before': probe['latest'], 'after': after}) + probe['latest'] = after + if not changed: + continue + updates += 1 + saved = await recipe._await_operation( + model.save, + f'policy-{updates}', + save_optimizer=False, + stop_event=stop_event, + on_result=snapshots.append, + ) + policy_uri = recipe._checkpoint_path(saved) + return updates + + updates = await train_groups() + probe_after = await _probe_logps(recipe, model, probe['trajectories'], stop_event) if probe else None + sync_async = await _sync_async_evidence(recipe, rollout, policy_uri, environments, stop_event, config) + trajectory = _training_evidence(all_rolled) + snapshot_uris = {recipe._checkpoint_path(snapshot) for snapshot in snapshots} + spans = list(timeline.spans) + evidence = { + 'two_updates': updates >= 2, + 'parameter_change': parameters_changed(probe.get('before'), probe_after), + **trajectory, + 'distinct_snapshots': len(snapshot_uris) >= 3 and policy_uri != initial_uri, + 'new_policy_requests': policy_uri in recording_sampler.uris, + 'slow_tool_isolated': slow_tool_isolated(spans), + 'training_overlap': training_overlaps_rollout(spans), + **sync_async, + } + values.update({ + 'updates': updates, + 'final_policy_uri': policy_uri, + 'snapshot_uris': sorted(snapshot_uris), + 'policy_requests': recording_sampler.uris, + 'metrics': metrics, + 'parameter_probe': { + 'before': probe.get('before'), 'after': probe_after, 'updates': probe.get('updates', []) + }, + 'timeline': spans, + 'evidence': evidence, + }) + if config.phase == 'learning': + final_rewards, final_evaluations = await _evaluate( + recipe, rollout, policy_uri, environments, stop_event, config) + baseline = sum(reward == 1.0 for reward in baseline_rewards) / len(baseline_rewards) + final = sum(reward == 1.0 for reward in final_rewards) / len(final_rewards) + same_plan = baseline_evaluations == final_evaluations + held_out = all(item['filename'] != recipe.TASK.filename for item in _EVALUATION_TASKS) + values.update({ + 'baseline_rewards': baseline_rewards, + 'final_rewards': final_rewards, + 'baseline_success_rate': baseline, + 'final_success_rate': final, + 'evaluation_plan': final_evaluations, + 'learning_gate': ( + same_plan and held_out and learning_gate( + updates=updates, + baseline=baseline, + final=final, + minimum_improvement=config.min_success_improvement, + )), + }) + except BaseException: + stop_event.set() + raise + finally: + cleanup_issues = await recipe._drain_cleanup(clients, environments, active_tasks, stop_event) + values['cleanup_issues'] = cleanup_issues + evidence = values.get('evidence', {}) + evidence['cleanup_confirmed'] = not cleanup_issues + values['evidence'] = evidence + values['pipeline_gate'] = pipeline_evidence_gate(evidence) + return values + + +class _CancelRecorder: + + def __init__(self, stage): + self.stage = stage + self.started = threading.Event() + self.release = threading.Event() + self.operations = [] + self.post_cancel_operations = [] + self._lock = threading.Lock() + self._cancelled = False + self._triggered = False + + def mark_cancelled(self): + with self._lock: + self._cancelled = True + + def block(self, operation): + with self._lock: + self.operations.append(operation) + if self._cancelled: + self.post_cancel_operations.append(operation) + trigger = operation == self.stage and not self._triggered + if trigger: + self._triggered = True + if trigger: + self.started.set() + if not self.release.wait(timeout=_STAGE_TIMEOUT): + raise TimeoutError(f'cancel validation did not release {operation}') + + +class _CancelModel: + + def __init__(self, model, recorder): + self._model = model + self._recorder = recorder + + def __getattr__(self, name): + return getattr(self._model, name) + + def forward_backward(self, *args, **kwargs): + self._recorder.block('backward') + return self._model.forward_backward(*args, **kwargs) + + def clip_grad_and_step(self, *args, **kwargs): + self._recorder.block('step') + return self._model.clip_grad_and_step(*args, **kwargs) + + def save(self, name, *args, **kwargs): + if name != 'policy-0': + self._recorder.block('save') + return self._model.save(name, *args, **kwargs) + + +def _cancel_tool_manager(base, recorder: _CancelRecorder): + """Record rollout tool batches only; environment setup writes go through ``Env.step`` directly.""" + + class CancelToolManager(base): + + def call_many(self, tool_calls): + recorder.block('tool') + return super().call_many(tool_calls) + + return CancelToolManager + + +async def _wait_for_stage(recorder: _CancelRecorder, task: asyncio.Task, timeout: float) -> bool: + """Wait on the event loop (no helper thread) until the stage blocks, the task ends, or time runs out.""" + deadline = time.monotonic() + timeout + while not recorder.started.is_set(): + if task.done() or time.monotonic() >= deadline: + return False + await asyncio.sleep(0.05) + return True + + +async def run_cancel_validation(config: ValidationConfig) -> dict: + recipe = _prepare_recipe(config) + recorder = _CancelRecorder(config.cancel_stage) + recipe.ToolManager = _cancel_tool_manager(recipe.ToolManager, recorder) + clients, model, sampler, rollout, stop_event = await _build_runtime(config, recipe) + environments, snapshots = [], [] + active_tasks = set() + values = {} + original_sample = sampler.sample + original_verify = recipe._TaskAgentEnv.verify + + def sample(*args, **kwargs): + recorder.block('sample') + return original_sample(*args, **kwargs) + + def verify(self): + recorder.block('verify') + return original_verify(self) + + sampler.sample = sample + recipe._TaskAgentEnv.verify = verify + model = _CancelModel(model, recorder) + try: + policy_uri = await _initial_policy(recipe, model, stop_event, snapshots) + + async def operation(): + rolled, rewards = await recipe._rollout_group(rollout, policy_uri, environments, stop_event) + if config.cancel_stage in {'backward', 'step', 'save'}: + updated = await recipe._train_group(model, rolled, rewards, GRPOAdvantage(), stop_event) + if updated and config.cancel_stage == 'save': + await recipe._await_operation( + model.save, + 'policy-1', + save_optimizer=False, + stop_event=stop_event, + on_result=snapshots.append, + ) + + task = asyncio.create_task(operation()) + active_tasks.add(task) + stage_reached = await _wait_for_stage(recorder, task, _STAGE_TIMEOUT) + recorder.mark_cancelled() + if not task.done(): + task.cancel() + # Deliver cancellation before releasing in-flight work, without a measurement blind spot. + await asyncio.sleep(0) + recorder.release.set() + cancelled = False + try: + await task + except asyncio.CancelledError: + cancelled = True + active_tasks.discard(task) + values = { + 'stage_reached': stage_reached, + 'cancelled': cancelled, + 'post_cancel_operations': list(recorder.post_cancel_operations), + 'operations': list(recorder.operations), + } + finally: + recorder.release.set() + cleanup_issues = await recipe._drain_cleanup(clients, environments, active_tasks, stop_event) + values['cleanup_issues'] = cleanup_issues + values['pipeline_gate'] = cancel_evidence_gate(values) + values['learning_gate'] = None + return values + + +def parse_args(argv=None) -> ValidationConfig: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument('--backend', choices=['transformers', 'megatron'], required=True) + parser.add_argument('--phase', choices=['pipeline', 'learning', 'cancel'], required=True) + parser.add_argument('--server-url', default='http://127.0.0.1:8000') + parser.add_argument('--model-id', default='ms://Qwen/Qwen3.5-4B') + parser.add_argument('--template-model-id', default=None) + parser.add_argument('--environment-backend', choices=['agentenv', 'openenv'], default='agentenv') + parser.add_argument('--agentenv-template', default=os.environ.get('AGENTENV_TEMPLATE')) + parser.add_argument('--openenv-url', default=os.environ.get('OPENENV_BASE_URL')) + parser.add_argument('--openenv-name', default='coding_env') + parser.add_argument('--result-dir', required=True) + parser.add_argument('--max-groups', type=int, default=4) + parser.add_argument('--active-groups', type=int, default=2) + parser.add_argument('--seed', type=int, default=7) + parser.add_argument('--eval-repeats', type=int, default=2) + parser.add_argument('--min-success-improvement', type=float, default=None) + parser.add_argument( + '--cancel-stage', choices=['sample', 'tool', 'verify', 'backward', 'step', 'save'], default='sample') + parser.add_argument('--slow-tool-seconds', type=float, default=20.0) + args = parser.parse_args(argv) + if args.max_groups < 2 or args.active_groups < 1 or args.eval_repeats < 1: + parser.error('max-groups must be >=2 and active-groups/eval-repeats must be positive') + if args.slow_tool_seconds <= 0: + parser.error('slow-tool-seconds must be positive') + if args.environment_backend == 'agentenv' and not args.agentenv_template: + parser.error('AgentENV requires --agentenv-template') + if args.environment_backend == 'openenv' and not args.openenv_url: + parser.error('OpenEnv requires --openenv-url') + return ValidationConfig( + backend=args.backend, + phase=args.phase, + server_url=args.server_url, + model_id=args.model_id, + template_model_id=args.template_model_id or args.model_id, + agentenv_template=args.agentenv_template or '', + result_dir=args.result_dir, + max_groups=args.max_groups, + active_groups=args.active_groups, + seed=args.seed, + eval_repeats=args.eval_repeats, + min_success_improvement=args.min_success_improvement, + cancel_stage=args.cancel_stage, + slow_tool_seconds=args.slow_tool_seconds, + environment_backend=args.environment_backend, + openenv_url=args.openenv_url, + openenv_name=args.openenv_name, + ) + + +def main(argv=None) -> int: + config = parse_args(argv) + try: + if config.phase == 'cancel': + values = asyncio.run(run_cancel_validation(config)) + else: + values = asyncio.run(run_validation(config)) + status = 'completed' if values['pipeline_gate'] and values.get('learning_gate') is not False else 'failed' + path = write_result(config, status=status, **values) + print(path) + return 0 if status == 'completed' else 1 + except BaseException as error: + path = write_result( + config, + status='blocked', + error={'type': type(error).__name__, 'message': str(error)}, + pipeline_gate=False, + learning_gate=False if config.phase == 'learning' else None, + ) + print(path, file=sys.stderr) + return 2 + + +if __name__ == '__main__': + raise SystemExit(main()) diff --git a/tests/server/integration/task_http_validation.py b/tests/server/integration/task_http_validation.py index 7a051db22..3f4468b56 100644 --- a/tests/server/integration/task_http_validation.py +++ b/tests/server/integration/task_http_validation.py @@ -70,11 +70,12 @@ def step(model): kwargs = {'task': args.task, 'gradient_accumulation_steps': 2} if args.backend == 'megatron': kwargs['micro_batch_size'] = 6 - result = model.forward_backward(inputs=copy.deepcopy(batch), **kwargs).result - assert np.isfinite(result['loss']) + model.forward_backward(inputs=copy.deepcopy(batch), **kwargs) # Call after every backward; the backend gates updates at its GA boundary. model.clip_grad_and_step() - return float(result['loss']) + loss = float(model.calculate_metric(is_training=True).result['loss']) + assert np.isfinite(loss) + return loss def save(model, name, samples): return model.save(name, save_optimizer=True, consumed_train_samples=samples) diff --git a/tests/server/model/test_task_head.py b/tests/server/model/test_task_head.py index 681d6e52e..055b2ee95 100644 --- a/tests/server/model/test_task_head.py +++ b/tests/server/model/test_task_head.py @@ -99,3 +99,13 @@ def tokenizer(text, **kwargs): parsed = TypeAdapter(WireInputBatch).validate_python(batch) assert len(parsed) == 12 assert all(type(label) is int for row in parsed for label in row.labels) + + +def test_transformers_backend_does_not_forward_head_keys(monkeypatch): + from twinkle.server.model import app + from twinkle.server.model.backends import transformers_model + + received = {} + monkeypatch.setattr(transformers_model, 'TwinkleCompatTransformersModel', lambda **kwargs: received.update(kwargs)) + app._make_transformers_model({'model_id': 'unused', 'task_type': 'seq_cls', 'num_labels': 3}) + assert received == {'model_id': 'unused'} diff --git a/tests/server/validation/test_wire_schema.py b/tests/server/validation/test_wire_schema.py index 09b18e7d7..ac788cb67 100644 --- a/tests/server/validation/test_wire_schema.py +++ b/tests/server/validation/test_wire_schema.py @@ -6,10 +6,10 @@ is evidence of nothing about validation. """ from __future__ import annotations +from typing import Union, get_args, get_origin import pytest from pydantic import TypeAdapter, ValidationError -from typing import Union, get_args, get_origin from twinkle.data_format.encoding import ENCODED_INPUT_KEYS, is_encoded from twinkle.processor.base import InputProcessor @@ -140,6 +140,22 @@ def test_unknown_json_fields_survive_a_round_trip(): assert exported == payload +def test_complete_multi_turn_trajectory_metadata_survives_a_round_trip(): + payload = { + 'input_ids': [1, 10, 90, 20], + 'labels': [10, -100, 20, -100], + 'completion_mask': [1, 0, 1, 0], + 'messages': [{'role': 'user', 'content': 'question'}, {'role': 'assistant', 'content': 'answer'}], + 'tools': [{'type': 'function', 'function': {'name': 'tool', 'parameters': {}}}], + 'turns': 2, + 'stop_reason': 'stop', + 'tool_stop': 'max_turns', + 'truncated': False, + 'rollout_logprobs': [-0.1, -0.2], + } + assert wire.export(_parse([payload])[0]) == payload + + def test_export_omits_unset_optional_fields(): """Twinkle_Core branches on key *presence*, so ``None`` must not be emitted.""" exported = wire.export(_parse([{'input_ids': [1, 2]}])[0]) diff --git a/tests/transformers/test_optimizer_group_regression.py b/tests/transformers/test_optimizer_group_regression.py index d03d16927..993336caf 100644 --- a/tests/transformers/test_optimizer_group_regression.py +++ b/tests/transformers/test_optimizer_group_regression.py @@ -75,3 +75,56 @@ def __init__(self, module): restored = torch.optim.Adam(late_groups, lr=1e-4) restored.load_state_dict(state) assert len(restored.state) == len(optimizer.state) + + +def _slot_module(): + """Resident slots whose names share a prefix, as with max_loras >= 21.""" + root = torch.nn.Module() + root.layer = torch.nn.Module() + root.layer.lora_A = torch.nn.ModuleDict({name: torch.nn.Linear(2, 2) for name in ('lora_2', 'lora_20')}) + root.layer.lora_B = torch.nn.ParameterDict( + {name: torch.nn.Parameter(torch.ones(2)) for name in ('lora_2', 'lora_21')}) + root.embed = torch.nn.Module() + root.embed.modules_to_save = torch.nn.ModuleDict({name: torch.nn.Linear(2, 2) for name in ('lora_2', 'lora_29')}) + return root + + +_EXPECTED_LORA_2 = { + 'layer.lora_A.lora_2.weight', 'layer.lora_A.lora_2.bias', 'layer.lora_B.lora_2', + 'embed.modules_to_save.lora_2.weight', 'embed.modules_to_save.lora_2.bias', +} + + +def test_transformers_trainable_parameters_do_not_select_prefixed_slots(): + module = _slot_module() + harness = SimpleNamespace(model=module, strategy=SimpleNamespace(unwrap_model=lambda wrapped: wrapped)) + assert set(TransformersModel._get_trainable_parameters(harness, 'lora_2')) == _EXPECTED_LORA_2 + + +def test_prefixed_slot_does_not_change_optimizer_restore_shape(): + sparse = torch.nn.Module() + sparse.layer = torch.nn.Module() + sparse.layer.lora_A = torch.nn.ModuleDict({'lora_2': torch.nn.Linear(2, 2)}) + crowded = _slot_module() + crowded.layer.lora_A['lora_2'].load_state_dict(sparse.layer.lora_A['lora_2'].state_dict()) + + def groups(module): + harness = SimpleNamespace(model=module, strategy=SimpleNamespace(unwrap_model=lambda wrapped: wrapped)) + params = {name: parameter for name, parameter in + TransformersModel._get_trainable_parameters(harness, 'lora_2').items() if '.lora_A.' in name} + harness._get_trainable_parameters = lambda adapter: params + return TransformersModel._create_param_group(harness, 'lora_2', weight_decay=0.1) + + source = torch.optim.Adam(groups(sparse), lr=1e-4) + sparse.layer.lora_A['lora_2'](torch.ones(1, 2)).sum().backward() + source.step() + restored = torch.optim.Adam(groups(crowded), lr=1e-4) + restored.load_state_dict(source.state_dict()) + assert len(restored.state) == len(source.state) == 2 + + +def test_megatron_trainable_parameters_do_not_select_prefixed_slots(): + megatron = pytest.importorskip('twinkle.model.megatron.megatron') + module = _slot_module() + harness = SimpleNamespace(model=[module], strategy=SimpleNamespace(unwrap_model=lambda wrapped: wrapped)) + assert set(megatron.MegatronModel._get_trainable_parameters(harness, 'lora_2')) == _EXPECTED_LORA_2 diff --git a/tests/twinkle_agentic/test_multi_turn_rollout.py b/tests/twinkle_agentic/test_multi_turn_rollout.py index 0c06d0271..23e8e97f6 100644 --- a/tests/twinkle_agentic/test_multi_turn_rollout.py +++ b/tests/twinkle_agentic/test_multi_turn_rollout.py @@ -18,20 +18,23 @@ pipeline semantics, and a fake Sampler that queues scripted responses. """ from __future__ import annotations - import copy -import json -import pytest import re import threading +from concurrent.futures import CancelledError, Future, ThreadPoolExecutor from typing import Any +import json +import pytest + from twinkle.data_format.sampling import SampledSequence, SampleResponse, SamplingParams from twinkle.template.base import Template -from twinkle_agentic.rollout.multi_turn import MultiTurnRollout +from twinkle_agentic.rollout.base import _RolloutExecution +from twinkle_agentic.rollout.multi_turn import MultiTurnRollout, default_response_callback from twinkle_agentic.rollout.trace import TraceWriter from twinkle_agentic.tools.base import Tool from twinkle_agentic.tools.tool_manager import ToolManager +from twinkle_agentic.utils.leases import Leases # ============================================================================= @@ -1481,3 +1484,212 @@ def followup(traj, n_before): assert '<|im_end|>' not in last['content'] # And it was not dispatched: a dispatch appends a tool message. assert [m['role'] for m in out['messages'] if m['role'] == 'tool'] == [] + + +def test_call_sampling_params_dict_is_normalized_without_mutating_input(sampler, template): + params = {'temperature': 0.25, 'max_tokens': 32} + sampler.queue('done') + rollout = MultiTurnRollout(sampler=sampler, template=template) + + rollout([_user_traj()], sampling_params=params) + + assert params == {'temperature': 0.25, 'max_tokens': 32} + assert isinstance(sampler.params_seen[0], SamplingParams) + assert sampler.params_seen[0].temperature == 0.25 + assert sampler.params_seen[0].max_tokens == 32 + assert sampler.params_seen[0].num_samples == 1 + + +def test_transport_state_validation_covers_every_append(sampler, template, tool_manager): + class ValidatingCallback: + def __init__(self): + self.validations = [] + + def __call__(self, *args, **kwargs): + return default_response_callback(*args, **kwargs) + + def _validate_state(self, state, feature, index): + self.validations.append((index, len(feature['input_ids']))) + + callback = ValidatingCallback() + followups = iter(('summarize', None)) + rollout = MultiTurnRollout( + sampler=sampler, + template=template, + tool_manager=tool_manager, + response_callback=callback, + followup_fn=lambda *_: next(followups), + max_turns=3, + ) + sampler.queue(_tool_call_text('search', {'q': 'x'})) + sampler.queue('finished tools') + sampler.queue('summary') + + result = rollout([_user_traj()])[0] + + assert result['followups'] == 1 + assert len(callback.validations) == 6 + + +def test_transport_state_validation_covers_malformed_feedback(sampler, template, tool_manager): + class ValidatingCallback: + def __init__(self): + self.count = 0 + + def __call__(self, *args, **kwargs): + return default_response_callback(*args, **kwargs) + + def _validate_state(self, state, feature, index): + self.count += 1 + + callback = ValidatingCallback() + rollout = MultiTurnRollout( + sampler=sampler, + template=template, + tool_manager=tool_manager, + response_callback=callback, + max_malformed_retries=1, + ) + sampler.queue('{bad json}') + sampler.queue('done') + + rollout([_user_traj()]) + + assert callback.count == 4 + + +def test_cancellation_drains_inflight_generation_without_starting_tools(template, tool_manager): + started = threading.Event() + release = threading.Event() + + class BlockingSampler(FakeSampler): + def sample(self, pifs, sampling_params=None): + started.set() + assert release.wait(timeout=5) + return super().sample(pifs, sampling_params) + + sample._enable_continous_work = True + + sampler = BlockingSampler(template) + sampler.queue(_tool_call_text('search', {'q': 'x'})) + sampler.queue('must not run') + rollout = MultiTurnRollout(sampler=sampler, template=template, tool_manager=tool_manager) + execution = _RolloutExecution() + + with ThreadPoolExecutor(max_workers=1) as pool: + future = pool.submit(rollout, [_user_traj()], _execution=execution) + assert started.wait(timeout=5) + execution.cancel() + release.set() + with pytest.raises(CancelledError): + future.result(timeout=5) + + assert sampler.sample_calls == 1 + + +def test_cancellation_drains_tool_batch_without_next_generation(template): + started = threading.Event() + release = threading.Event() + + class BlockingTools(ToolManager): + def __init__(self): + super().__init__({}) + self.register(EchoTool('search')) + self.calls = 0 + + def call_many(self, tool_calls, max_workers=None): + self.calls += 1 + started.set() + assert release.wait(timeout=5) + return super().call_many(tool_calls, max_workers) + + sampler = FakeSampler(template) + sampler.queue(_tool_call_text('search', {'q': 'x'})) + sampler.queue('must not run') + tools = BlockingTools() + rollout = MultiTurnRollout(sampler=sampler, template=template, tool_manager=tools) + execution = _RolloutExecution() + + with ThreadPoolExecutor(max_workers=1) as pool: + future = pool.submit(rollout, [_user_traj()], _execution=execution) + assert started.wait(timeout=5) + execution.cancel() + release.set() + with pytest.raises(CancelledError): + future.result(timeout=5) + + assert tools.calls == 1 + assert sampler.sample_calls == 1 + + +def test_cancellation_interrupts_known_harness_lease_wait_without_closing_pool(template): + started = threading.Event() + release = threading.Event() + sampled = [] + + class BlockingSampler(FakeSampler): + def sample(self, pifs, sampling_params=None): + key = self._key_of(pifs[0]) + sampled.append(key) + if key == 'A': + started.set() + assert release.wait(timeout=5) + return super().sample(pifs, sampling_params) + + sample._enable_continous_work = True + + class Harness: + def before_generate(self, trajectory): + return trajectory + + sampler = BlockingSampler(template) + sampler.queue_for('A', 'done A') + sampler.queue_for('B', 'done B') + pool = Leases([Harness()]) + rollout = MultiTurnRollout(sampler=sampler, template=template, harness=pool, concurrency=2) + execution = _RolloutExecution() + + with ThreadPoolExecutor(max_workers=1) as executor: + future = executor.submit(rollout, [_user_traj('A'), _user_traj('B')], _execution=execution) + assert started.wait(timeout=5) + execution.cancel() + release.set() + with pytest.raises(CancelledError): + future.result(timeout=5) + + assert sampled == ['A'] + assert pool._closing is False + + +def test_failure_during_future_submission_preserves_original_error(template, monkeypatch): + sentinel = RuntimeError('first submission failed') + + class ImmediateExecutor: + def submit(self, function, *args): + future = Future() + try: + future.set_result(function(*args)) + except BaseException as error: + future.set_exception(error) + return future + + @staticmethod + def shutdown(wait=True): + return None + + class FailingSampler(FakeSampler): + def sample(self, pifs, sampling_params=None): + self.sample_calls += 1 + raise sentinel + + sample._enable_continous_work = True + + monkeypatch.setattr('twinkle_agentic.rollout.base.ThreadPoolExecutor', lambda max_workers: ImmediateExecutor()) + sampler = FailingSampler(template) + rollout = MultiTurnRollout(sampler=sampler, template=template, concurrency=3) + + with pytest.raises(RuntimeError) as error: + rollout([_user_traj('A'), _user_traj('B'), _user_traj('C')]) + + assert error.value is sentinel + assert sampler.sample_calls == 1 diff --git a/tests/twinkle_agentic/test_tools.py b/tests/twinkle_agentic/test_tools.py index cfa5def56..7f2283481 100644 --- a/tests/twinkle_agentic/test_tools.py +++ b/tests/twinkle_agentic/test_tools.py @@ -8,7 +8,6 @@ from twinkle_agentic.tools.base import Tool from twinkle_agentic.tools.tool_manager import ToolManager, _extract_name - # --------------------------------------------------------------------------- # Mock tools # --------------------------------------------------------------------------- @@ -41,6 +40,13 @@ def tool_info(self) -> ToolInfo: return {'function': {'name': 'calculator', 'description': 'Basic calculator.'}} +class BoundTool(MockTool): + + def __init__(self, env, name='bound_tool'): + super().__init__(name) + self._env = env + + class FailingTool(Tool): def __call__(self, tool_name: str, arguments: dict) -> str: @@ -245,3 +251,34 @@ def tool_info(self): return {} # missing function.name with pytest.raises(ValueError, match='non-empty'): ToolManager([BadTool()]) + + def test_bound_envs_empty_and_deduplicated_by_identity(self): + class Env: + def __init__(self): + self.idle_calls = 0 + + def ensure_idle(self): + self.idle_calls += 1 + + env = Env() + manager = ToolManager([BoundTool(env, 'first'), BoundTool(env, 'second'), MockTool('stateless')]) + assert ToolManager().bound_envs() == () + assert manager.bound_envs() == (env, ) + + manager.ensure_idle() + assert env.idle_calls == 1 + + def test_bound_envs_returns_each_distinct_environment(self): + first, second = object(), object() + manager = ToolManager([BoundTool(first, 'first'), BoundTool(second, 'second')]) + assert manager.bound_envs() == (first, second) + + def test_ensure_idle_failure_propagates(self): + class Env: + @staticmethod + def ensure_idle(): + raise RuntimeError('execution unknown') + + manager = ToolManager([BoundTool(Env())]) + with pytest.raises(RuntimeError, match='execution unknown'): + manager.ensure_idle() diff --git a/tests/twinkle_client/test_client_multi_turn_rollout.py b/tests/twinkle_client/test_client_multi_turn_rollout.py index f3ce43a93..8d9f52861 100644 --- a/tests/twinkle_client/test_client_multi_turn_rollout.py +++ b/tests/twinkle_client/test_client_multi_turn_rollout.py @@ -1,12 +1,25 @@ # Copyright (c) ModelScope Contributors. All rights reserved. +import asyncio +import threading +from concurrent.futures import CancelledError, Future, ThreadPoolExecutor from dataclasses import replace import pytest -from twinkle.template.contract import (CompletionEncoding, DecodedResponse, MediaBundle, PromptEncoding, RolloutState, - refresh, splice, splice_state) -from twinkle_client.rollout.multi_turn import ClientMultiTurnRollout +from twinkle.data_format.sampling import SampledSequence from twinkle.protocol.types.sampler import SampledSequenceModel, SampleResponseModel +from twinkle.template.contract import ( + CompletionEncoding, + DecodedResponse, + MediaBundle, + PromptEncoding, + RolloutState, + refresh, + splice, + splice_state, +) +from twinkle_agentic.rollout.multi_turn import MultiTurnRollout +from twinkle_client.rollout.multi_turn import ClientMultiTurnRollout class _Template: @@ -62,9 +75,13 @@ class _Sampler: def __init__(self, scripts=((10, 'stop'), (20, 'stop'))): self.scripts = list(scripts) self.calls = [] + self.policies = [] + self.params = [] - def sample(self, features, sampling_params=None): + def sample(self, features, sampling_params=None, adapter_name='', adapter_uri=None, num_samples=1): self.calls.append(features) + self.policies.append((adapter_name, adapter_uri, num_samples)) + self.params.append(sampling_params) token, stop_reason = self.scripts.pop(0) responses = [] for feature in features: @@ -109,6 +126,36 @@ def test_text_http_multi_turn_preserves_tokens_messages_and_logprobs(): assert result['truncated'] is False +def test_core_and_wire_characterization_share_scripted_trajectory(): + def wire_callback(sampler, api, sampling_params, *, input_feature, adapter_kwargs, **kwargs): + response = sampler.sample([input_feature], sampling_params=sampling_params)[0] + sequence = response.sequences[0] + return SampledSequence( + stop_reason=sequence.stop_reason, + tokens=sequence.tokens, + logprobs=sequence.logprobs, + decoded=sequence.decoded, + new_input_feature=sequence.new_input_feature, + ) + + wire_sampler = _Sampler() + core = MultiTurnRollout( + sampler=None, + template=_Template(), + tool_manager=_ToolManager(), + max_turns=2, + response_callback=wire_callback, + ) + core_result = core([_trajectory()], response_callback=lambda _sampler, _api, params, **kwargs: wire_callback( + wire_sampler, _api, params, **kwargs))[0] + http_result = ClientMultiTurnRollout( + _Sampler(), _Template(), tool_manager=_ToolManager(), max_turns=2)([_trajectory()])[0] + + for key in ('input_ids', 'labels', 'completion_mask', 'messages', 'tools', 'logprobs', 'rollout_logprobs'): + assert core_result[key] == http_result[key] + assert core_result['stop_reason'] == http_result['stop_reason'] == 'stop' + + def test_http_multi_turn_complete_trajectory_golden(): result = ClientMultiTurnRollout( _Sampler(), _Template(), tool_manager=_ToolManager(), max_turns=2)([_trajectory()])[0] @@ -133,6 +180,9 @@ def test_http_multi_turn_complete_trajectory_golden(): 'turns': 2, 'stop_reason': 'stop', 'truncated': False, + 'stuck_stop': False, + 'tool_stop': None, + 'followups': 0, } @@ -167,17 +217,17 @@ def test_http_multi_turn_requires_new_input_feature(): class BrokenSampler: @staticmethod - def sample(features, sampling_params=None): + def sample(features, sampling_params=None, **kwargs): sequence = SampledSequenceModel( stop_reason='stop', tokens=[10], logprobs=[[(10, -0.1)]], decoded='x') return [SampleResponseModel(sequences=[sequence])] - with pytest.raises(RuntimeError, match='batch index 0.*trajectory 0.*new_input_feature'): + with pytest.raises(RuntimeError, match='new_input_feature.*trajectory 0'): ClientMultiTurnRollout(BrokenSampler(), _Template())([_trajectory()]) def test_http_multi_turn_requires_tool_manager_only_when_called(): - with pytest.raises(ValueError, match='trajectory 0.*tool_manager'): + with pytest.raises(ValueError, match='tool_calls.*ToolManager'): ClientMultiTurnRollout(_Sampler(), _Template(), max_turns=2)([_trajectory()]) @@ -190,7 +240,7 @@ def test_http_multi_turn_marks_length_stop_as_truncated(): def test_http_multi_turn_preserves_batch_output_order(): trajectories = [dict(_trajectory(), sample_id=index) for index in range(4)] - results = ClientMultiTurnRollout(_Sampler(((30, 'stop'), )), _Template())(trajectories) + results = ClientMultiTurnRollout(_Sampler(((30, 'stop'), ) * 4), _Template())(trajectories) assert [result['sample_id'] for result in results] == [0, 1, 2, 3] @@ -226,7 +276,8 @@ def call_many(cls, tool_calls): result = ClientMultiTurnRollout( _Sampler(((10, 'stop'), )), _Template(), tool_manager=RecordingTools(), max_turns=1)([_trajectory()])[0] assert result['turns'] == 1 - assert result['stop_reason'] == 'max_turns' + assert result['stop_reason'] == 'stop' + assert result['tool_stop'] == 'max_turns' assert result['truncated'] is True assert RecordingTools.calls == 0 @@ -240,7 +291,7 @@ class NetworkError(Exception): class FailingSampler: @staticmethod - def sample(features, sampling_params=None): + def sample(features, sampling_params=None, **kwargs): raise sentinel with pytest.raises(NetworkError) as error: @@ -253,3 +304,443 @@ def test_http_multi_turn_validates_public_arguments(): ClientMultiTurnRollout(object(), None) with pytest.raises(ValueError, match='max_turns'): ClientMultiTurnRollout(object(), _Template(), max_turns=0) + + +def test_policy_and_sampling_params_are_snapshotted_per_call(): + sampler = _Sampler(((30, 'stop'), (31, 'stop'))) + rollout = ClientMultiTurnRollout( + sampler, _Template(), adapter_name='trained', adapter_uri='twinkle://policy-0') + rollout([_trajectory()]) + override = {'temperature': 0.25} + rollout([_trajectory()], adapter_name='base', adapter_uri=None, sampling_params=override) + + assert sampler.policies == [('trained', 'twinkle://policy-0', 1), ('base', None, 1)] + assert sampler.params[0].temperature == 1.0 + assert sampler.params[1].temperature == 0.25 + assert override == {'temperature': 0.25} + + +def test_invalid_call_options_are_rejected_before_http(): + sampler = _Sampler() + rollout = ClientMultiTurnRollout(sampler, _Template()) + with pytest.raises(ValueError, match='unknown.*adapter_path'): + rollout([_trajectory()], adapter_path='not-a-cs-option') + with pytest.raises(ValueError, match='num_samples=1'): + rollout([_trajectory()], sampling_params={'num_samples': 2}) + with pytest.raises(ValueError, match='tool_manager list length'): + rollout([_trajectory()], tool_manager=[]) + assert sampler.calls == [] + assert rollout([]) == [] + + +def test_sync_and_async_calls_return_the_same_core_trajectory(): + sync_rollout = ClientMultiTurnRollout(_Sampler(), _Template(), tool_manager=_ToolManager(), max_turns=2) + async_rollout = ClientMultiTurnRollout(_Sampler(), _Template(), tool_manager=_ToolManager(), max_turns=2) + + expected = sync_rollout([_trajectory()]) + + async def run(): + return await async_rollout.acall([_trajectory()]) + + assert asyncio.run(run()) == expected + + +def test_async_cancellation_drains_http_and_does_not_poison_later_calls(): + started = threading.Event() + release = threading.Event() + + class BlockingSampler(_Sampler): + def sample(self, *args, **kwargs): + started.set() + assert release.wait(timeout=5) + return super().sample(*args, **kwargs) + + sampler = BlockingSampler(((10, 'stop'), (20, 'stop'))) + rollout = ClientMultiTurnRollout(sampler, _Template()) + + async def run(): + task = asyncio.create_task(rollout.acall([_trajectory()])) + assert await asyncio.to_thread(started.wait, 5) + task.cancel() + task.cancel() + release.set() + with pytest.raises(asyncio.CancelledError): + await task + return await rollout.acall([_trajectory()]) + + result = asyncio.run(run()) + assert result[0]['input_ids'] == [1, 20] + assert len(sampler.calls) == 2 + + +def test_sync_repeated_interrupts_keep_draining_until_core_finishes(): + started = threading.Event() + release = threading.Event() + + class BlockingSampler(_Sampler): + def sample(self, *args, **kwargs): + started.set() + assert release.wait(timeout=5) + return super().sample(*args, **kwargs) + + class InterruptingFuture: + + def __init__(self, future, interrupts): + self._future = future + self._interrupts = interrupts + + def done(self): + return self._future.done() + + def result(self): + if self._interrupts: + self._interrupts -= 1 + assert started.wait(timeout=5) + raise KeyboardInterrupt + return self._future.result() + + sampler = BlockingSampler(((10, 'stop'), (20, 'stop'))) + rollout = ClientMultiTurnRollout(sampler, _Template(), tool_manager=_ToolManager(), max_turns=2) + calls = [] + original_start = rollout._start + + def start(trajectories, kwargs): + call = original_start(trajectories, kwargs) + calls.append(call) + call.future = InterruptingFuture(call.future, interrupts=2) + return call + + rollout._start = start + outcome = {} + + def run(): + try: + rollout([_trajectory()]) + except BaseException as error: + outcome['error'] = error + + worker = threading.Thread(target=run) + worker.start() + assert started.wait(timeout=5) + worker.join(timeout=0.3) + assert worker.is_alive(), 'sync call returned before the in-flight sample finished' + release.set() + worker.join(timeout=5) + + assert not worker.is_alive() + assert isinstance(outcome['error'], KeyboardInterrupt) + assert calls[0].future._future.done() + assert len(sampler.calls) == 1 + + +@pytest.mark.parametrize('control_error', [KeyboardInterrupt(), SystemExit(2)]) +def test_sync_completed_future_control_error_is_not_retried_forever(control_error, monkeypatch): + from types import SimpleNamespace + from twinkle_agentic.rollout.base import _RolloutExecution + + future = Future() + future.set_exception(control_error) + rollout = ClientMultiTurnRollout(_Sampler(), _Template()) + execution = _RolloutExecution() + monkeypatch.setattr(rollout, '_start', lambda *args: SimpleNamespace(future=future, execution=execution)) + with pytest.raises(type(control_error)) as error: + rollout([_trajectory()]) + assert error.value is control_error + assert execution.cancelled + + +def test_duplicate_env_and_bare_harness_are_rejected_before_side_effects(): + env = object() + + class EnvManager(_ToolManager): + @staticmethod + def bound_envs(): + return (env, ) + + class Harness: + calls = 0 + + def before_generate(self, trajectory): + type(self).calls += 1 + return trajectory + + sampler = _Sampler() + harness = Harness() + rollout = ClientMultiTurnRollout(sampler, _Template()) + with pytest.raises(ValueError, match='same Env'): + rollout([_trajectory(), _trajectory()], tool_manager=[EnvManager(), EnvManager()]) + with pytest.raises(ValueError, match='same bare harness'): + rollout([_trajectory(), _trajectory()], harness=[harness, harness]) + assert sampler.calls == [] + assert harness.calls == 0 + + +def test_active_env_borrow_is_released_only_after_runner_finishes(): + started = threading.Event() + release = threading.Event() + + class BlockingSampler(_Sampler): + def sample(self, *args, **kwargs): + started.set() + assert release.wait(timeout=5) + return super().sample(*args, **kwargs) + + class Env: + close_calls = 0 + + def close(self): + self.close_calls += 1 + + class EnvManager(_ToolManager): + def __init__(self, env): + self.env = env + + def bound_envs(self): + return (self.env, ) + + env = Env() + manager = EnvManager(env) + sampler = BlockingSampler(((10, 'stop'), (20, 'stop'), (30, 'stop'))) + rollout = ClientMultiTurnRollout(sampler, _Template(), max_turns=2) + with ThreadPoolExecutor(max_workers=1) as pool: + first = pool.submit(rollout, [_trajectory()], tool_manager=manager) + assert started.wait(timeout=5) + with pytest.raises(ValueError, match='still borrowed'): + rollout([_trajectory()], tool_manager=manager) + release.set() + assert first.result(timeout=5)[0]['turns'] == 2 + result = rollout([_trajectory()], tool_manager=manager) + assert result[0]['turns'] == 1 + assert env.close_calls == 0 + + +def test_active_bare_harness_borrow_covers_concurrent_single_episode_calls(): + started = threading.Event() + release = threading.Event() + + class BlockingSampler(_Sampler): + def sample(self, *args, **kwargs): + started.set() + assert release.wait(timeout=5) + return super().sample(*args, **kwargs) + + class Harness: + def before_generate(self, trajectory): + return trajectory + + @staticmethod + def after_generate(context, decoded, tool_calls): + return {'messages': context['messages'] + [{'role': 'assistant', 'content': decoded}]} + + sampler = BlockingSampler(((30, 'stop'), (40, 'stop'))) + rollout = ClientMultiTurnRollout(sampler, _Template(), harness=Harness()) + with ThreadPoolExecutor(max_workers=1) as pool: + first = pool.submit(rollout, [_trajectory()]) + assert started.wait(timeout=5) + with pytest.raises(ValueError, match='still borrowed'): + rollout([_trajectory()]) + release.set() + assert first.result(timeout=5)[0]['input_ids'] == [1, 30] + + +def test_independent_episodes_do_not_wait_for_a_slow_tool_round(): + started = threading.Event() + release = threading.Event() + fast_finished = threading.Event() + + class KeyedTemplate(_Template): + def encode_prompt(self, trajectory, *, backend='vllm', add_generation_prompt=True): + encoding = super().encode_prompt( + trajectory, backend=backend, add_generation_prompt=add_generation_prompt) + token = 1 if trajectory['messages'][0]['content'] == 'slow' else 2 + state = replace( + encoding.state, + feature={'input_ids': [token]}, + engine_input_ids=(token, ), + ) + return PromptEncoding(state) + + class KeyedSampler: + def __init__(self): + self.calls = {1: 0, 2: 0} + + def sample(self, features, sampling_params=None, **kwargs): + feature = features[0] + key = feature['input_ids'][0] + self.calls[key] += 1 + if key == 1 and self.calls[key] == 1: + started.set() + assert release.wait(timeout=5) + token = 10 if self.calls[key] == 1 else 20 + updated = splice(feature, [token], scored=True, sampled=True) + sequence = SampledSequenceModel( + stop_reason='stop', + tokens=[token], + logprobs=[[(token, -0.1)]], + decoded=f'answer-{token}', + new_input_feature=updated, + ) + if key == 2 and self.calls[key] == 2: + fast_finished.set() + return [SampleResponseModel(sequences=[sequence])] + + rollout = ClientMultiTurnRollout( + KeyedSampler(), KeyedTemplate(), tool_manager=_ToolManager(), max_turns=2, concurrency=2) + trajectories = [ + {'messages': [{'role': 'user', 'content': 'slow'}]}, + {'messages': [{'role': 'user', 'content': 'fast'}]}, + ] + with ThreadPoolExecutor(max_workers=1) as pool: + future = pool.submit(rollout, trajectories) + assert started.wait(timeout=5) + assert fast_finished.wait(timeout=5) + release.set() + results = future.result(timeout=5) + assert [result['input_ids'][0] for result in results] == [1, 2] + assert all(result['turns'] == 2 for result in results) + + +def test_first_episode_error_is_preserved_after_peer_drain(): + peer_started = threading.Event() + failing = threading.Event() + release = threading.Event() + sentinel = RuntimeError('sample failed') + + class KeyedTemplate(_Template): + def encode_prompt(self, trajectory, *, backend='vllm', add_generation_prompt=True): + encoding = super().encode_prompt( + trajectory, backend=backend, add_generation_prompt=add_generation_prompt) + token = 1 if trajectory['messages'][0]['content'] == 'slow' else 2 + state = replace(encoding.state, feature={'input_ids': [token]}, engine_input_ids=(token, )) + return PromptEncoding(state) + + class FailingSampler: + @staticmethod + def sample(features, sampling_params=None, **kwargs): + key = features[0]['input_ids'][0] + if key == 2: + assert peer_started.wait(timeout=5) + failing.set() + raise sentinel + peer_started.set() + assert failing.wait(timeout=5) + assert release.wait(timeout=5) + feature = features[0] + updated = splice(feature, [10], scored=True, sampled=True) + sequence = SampledSequenceModel( + stop_reason='stop', + tokens=[10], + logprobs=[[(10, -0.1)]], + decoded='late', + new_input_feature=updated, + ) + return [SampleResponseModel(sequences=[sequence])] + + rollout = ClientMultiTurnRollout(FailingSampler(), KeyedTemplate(), concurrency=2) + trajectories = [ + {'messages': [{'role': 'user', 'content': 'slow'}]}, + {'messages': [{'role': 'user', 'content': 'fail'}]}, + ] + with ThreadPoolExecutor(max_workers=1) as pool: + future = pool.submit(rollout, trajectories) + assert failing.wait(timeout=5) + assert not future.done() + release.set() + with pytest.raises(RuntimeError) as error: + future.result(timeout=5) + assert error.value is sentinel + + +def test_cancel_during_followup_prevents_another_generation(): + followup_started = threading.Event() + release = threading.Event() + sampler = _Sampler(((30, 'stop'), (40, 'stop'))) + + def followup(trajectory, count): + followup_started.set() + assert release.wait(timeout=5) + return 'continue' + + rollout = ClientMultiTurnRollout(sampler, _Template(), followup_fn=followup) + with ThreadPoolExecutor(max_workers=1) as pool: + future = pool.submit(rollout, [_trajectory()]) + assert followup_started.wait(timeout=5) + rollout.cancel() + release.set() + with pytest.raises(CancelledError): + future.result(timeout=5) + assert len(sampler.calls) == 1 + + +def test_sampler_response_cardinality_and_logprob_alignment_are_validated(): + class EmptySampler: + @staticmethod + def sample(features, sampling_params=None, **kwargs): + return [] + + with pytest.raises(RuntimeError, match='expected one'): + ClientMultiTurnRollout(EmptySampler(), _Template())([_trajectory()]) + + class MisalignedSampler: + @staticmethod + def sample(features, sampling_params=None, **kwargs): + updated = splice(features[0], [10, 11], scored=True, sampled=True) + sequence = SampledSequenceModel( + stop_reason='stop', + tokens=[10, 11], + logprobs=[[(10, -0.1)]], + decoded='misaligned', + new_input_feature=updated, + ) + return [SampleResponseModel(sequences=[sequence])] + + with pytest.raises(RuntimeError, match='logprob'): + ClientMultiTurnRollout(MisalignedSampler(), _Template())([_trajectory()]) + + +def test_concurrent_policy_capture_and_completed_call_cancel_race(): + slow_started = threading.Event() + release = threading.Event() + policies = [] + + class KeyedTemplate(_Template): + def encode_prompt(self, trajectory, *, backend='vllm', add_generation_prompt=True): + encoding = super().encode_prompt( + trajectory, backend=backend, add_generation_prompt=add_generation_prompt) + token = {'slow': 1, 'fast': 2, 'later': 3}[trajectory['messages'][0]['content']] + state = replace(encoding.state, feature={'input_ids': [token]}, engine_input_ids=(token, )) + return PromptEncoding(state) + + class PolicySampler: + @staticmethod + def sample(features, sampling_params=None, adapter_name='', adapter_uri=None, num_samples=1): + policies.append((features[0]['input_ids'][0], adapter_name, adapter_uri)) + if adapter_uri == 'slow': + slow_started.set() + assert release.wait(timeout=5) + updated = splice(features[0], [30], scored=True, sampled=True) + sequence = SampledSequenceModel( + stop_reason='stop', + tokens=[30], + logprobs=[[(30, -0.1)]], + decoded='done', + new_input_feature=updated, + ) + return [SampleResponseModel(sequences=[sequence])] + + def trajectory(name): + return {'messages': [{'role': 'user', 'content': name}]} + + rollout = ClientMultiTurnRollout(PolicySampler(), KeyedTemplate()) + with ThreadPoolExecutor(max_workers=1) as pool: + slow = pool.submit(rollout, [trajectory('slow')], adapter_name='policy', adapter_uri='slow') + assert slow_started.wait(timeout=5) + fast = rollout([trajectory('fast')], adapter_name='policy', adapter_uri='fast') + assert fast[0]['input_ids'] == [2, 30] + rollout.cancel() + release.set() + with pytest.raises(CancelledError): + slow.result(timeout=5) + later = rollout([trajectory('later')], adapter_name='base', adapter_uri=None) + assert later[0]['input_ids'] == [3, 30] + assert policies == [(1, 'policy', 'slow'), (2, 'policy', 'fast'), (3, 'base', None)] diff --git a/tests/twinkle_client/test_import_surface.py b/tests/twinkle_client/test_import_surface.py index fcce13399..20584104f 100644 --- a/tests/twinkle_client/test_import_surface.py +++ b/tests/twinkle_client/test_import_surface.py @@ -1,10 +1,10 @@ # Copyright (c) ModelScope Contributors. All rights reserved. from __future__ import annotations - -import json import subprocess import sys +import json + def _run_import_probe(source: str) -> dict[str, object]: completed = subprocess.run( @@ -57,3 +57,10 @@ def test_import_twinkle_client_does_not_load_heavy_data_dependencies() -> None: "import json, sys, twinkle_client; " "print(json.dumps({name: name in sys.modules for name in ('torch', 'datasets', 'pandas')}))") assert result == {'torch': False, 'datasets': False, 'pandas': False} + + +def test_multi_turn_import_does_not_start_ray_or_load_optional_sandbox_engine() -> None: + result = _run_import_probe( + "import json, sys; from twinkle_client.rollout import ClientMultiTurnRollout; " + "print(json.dumps({name: name in sys.modules for name in ('ray', 'e2b', 'vllm.engine', 'transfer_queue')}))") + assert result == {'ray': False, 'e2b': False, 'vllm.engine': False, 'transfer_queue': False} diff --git a/tests/twinkle_client/test_multi_turn_grpo_recipe.py b/tests/twinkle_client/test_multi_turn_grpo_recipe.py new file mode 100644 index 000000000..f3ef31fc9 --- /dev/null +++ b/tests/twinkle_client/test_multi_turn_grpo_recipe.py @@ -0,0 +1,564 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from __future__ import annotations + +import asyncio +import importlib.util +import json +import subprocess +import sys +import threading +from pathlib import Path +from types import SimpleNamespace + +import pytest + + +@pytest.fixture +def recipe(): + path = Path(__file__).parents[2] / 'cookbook/client/async_rl/multi_turn_grpo.py' + spec = importlib.util.spec_from_file_location('multi_turn_grpo_recipe', path) + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module + + +def _cancel_and_release(task, started, release): + assert started.wait(timeout=5) + task.cancel() + task.cancel() + release.set() + + +def test_protected_wait_registers_result_before_repeated_cancellation(recipe): + started, release = threading.Event(), threading.Event() + + def save(): + started.set() + assert release.wait(timeout=5) + return SimpleNamespace(twinkle_path='twinkle://policy-1') + + async def run(): + stop_event = asyncio.Event() + snapshots = [] + task = asyncio.create_task( + recipe._await_operation(save, stop_event=stop_event, on_result=snapshots.append)) + await asyncio.to_thread(_cancel_and_release, task, started, release) + with pytest.raises(asyncio.CancelledError): + await task + assert stop_event.is_set() + return snapshots + + snapshots = asyncio.run(run()) + assert [item.twinkle_path for item in snapshots] == ['twinkle://policy-1'] + + +@pytest.mark.parametrize('stage', ['backward', 'step']) +def test_train_group_cancellation_stops_later_operations(recipe, stage): + started, release = threading.Event(), threading.Event() + + class Model: + def __init__(self): + self.calls = [] + + def forward_backward(self, **kwargs): + self.calls.append('backward') + if stage == 'backward': + started.set() + assert release.wait(timeout=5) + + def clip_grad_and_step(self, **kwargs): + self.calls.append('step') + if stage == 'step': + started.set() + assert release.wait(timeout=5) + + def calculate_metric(self, **kwargs): + self.calls.append('metric') + return {'result': {}} + + class Advantages: + @staticmethod + def tolist(): + return [1.0, -1.0, 0.5, -0.5] + + model = Model() + rolled = [{ + 'response_token_ids': [10], + 'rollout_logprobs': [-0.1], + 'stop_reason': 'stop', + 'truncated': False, + } for _ in range(4)] + + async def run(): + stop_event = asyncio.Event() + task = asyncio.create_task( + recipe._train_group(model, rolled, [1.0, 0.0, 1.0, 0.0], lambda *args, **kwargs: Advantages(), + stop_event)) + await asyncio.to_thread(_cancel_and_release, task, started, release) + with pytest.raises(asyncio.CancelledError): + await task + + asyncio.run(run()) + assert model.calls == ([stage] if stage == 'backward' else ['backward', 'step']) + + +def test_environment_creation_cancellation_preserves_handle_and_skips_rollout(recipe, monkeypatch): + started, release = threading.Event(), threading.Event() + env = SimpleNamespace(sandbox_id='sandbox-1') + attempt = SimpleNamespace(env=env) + + def create(task, on_created): + on_created(env) + started.set() + assert release.wait(timeout=5) + return attempt + + monkeypatch.setattr(recipe, '_create_attempt', create) + + class Rollout: + calls = 0 + + async def acall(self, *args, **kwargs): + self.calls += 1 + return [] + + rollout = Rollout() + environments = [] + + async def run(): + stop_event = asyncio.Event() + task = asyncio.create_task(recipe._rollout_group(rollout, 'policy-0', environments, stop_event)) + await asyncio.to_thread(_cancel_and_release, task, started, release) + with pytest.raises(asyncio.CancelledError): + await task + + asyncio.run(run()) + assert environments == [env] + assert rollout.calls == 0 + + +def test_verification_cancellation_stops_later_verify_calls(recipe, monkeypatch): + started, release = threading.Event(), threading.Event() + + class Env: + def __init__(self): + self.verify_calls = 0 + + def verify(self): + self.verify_calls += 1 + started.set() + assert release.wait(timeout=5) + return 0.0 + + attempts = [SimpleNamespace( + env=Env(), manager=object(), trajectory={'messages': []}) for _ in range(recipe.NUM_GENERATIONS)] + created = [] + + def create(task, on_created): + attempt = attempts.pop(0) + created.append(attempt) + return attempt + + monkeypatch.setattr(recipe, '_create_attempt', create) + + class Rollout: + async def acall(self, trajectories, **kwargs): + return [{ + 'response_token_ids': [10], + 'rollout_logprobs': [-0.1], + 'stop_reason': 'stop', + 'truncated': False, + } for _ in trajectories] + + async def run(): + stop_event = asyncio.Event() + task = asyncio.create_task(recipe._rollout_group(Rollout(), 'policy-0', [], stop_event)) + await asyncio.to_thread(_cancel_and_release, task, started, release) + with pytest.raises(asyncio.CancelledError): + await task + + asyncio.run(run()) + assert [attempt.env.verify_calls for attempt in created] == [1, 0, 0, 0] + + +def test_cleanup_repeated_cancellation_closes_environment_before_client(recipe): + started, release = threading.Event(), threading.Event() + events = [] + + class Env: + sandbox_id = 'sandbox-cleanup' + + @staticmethod + def ensure_idle(): + return None + + @staticmethod + def close(): + started.set() + assert release.wait(timeout=5) + events.append('env-closed') + + class Client: + @staticmethod + def close(): + events.append('client-closed') + + async def run(): + stop_event = asyncio.Event() + task = asyncio.create_task(recipe._drain_cleanup([Client()], [Env()], set(), stop_event)) + await asyncio.to_thread(_cancel_and_release, task, started, release) + with pytest.raises(asyncio.CancelledError): + await task + + asyncio.run(run()) + assert events == ['env-closed', 'client-closed'] + + +def test_cleanup_cancellation_does_not_replace_primary_error(recipe): + started, release = threading.Event(), threading.Event() + + class Env: + sandbox_id = 'sandbox-primary' + + @staticmethod + def ensure_idle(): + return None + + @staticmethod + def close(): + started.set() + assert release.wait(timeout=5) + + class Client: + closed = False + + def close(self): + self.closed = True + + client = Client() + + async def run(): + stop_event = asyncio.Event() + task = asyncio.create_task( + recipe._drain_cleanup([client], [Env()], set(), stop_event, propagate_cancel=False)) + await asyncio.to_thread(_cancel_and_release, task, started, release) + return await task + + assert asyncio.run(run()) == [] + assert client.closed + + +def test_complete_recipe_train_runs_updates_and_closes_client(recipe, monkeypatch): + calls = [] + + class Model: + def __getattr__(self, name): + if name == 'save': + def save(policy, **kwargs): + calls.append(('save', policy)) + return SimpleNamespace(twinkle_path=f'twinkle://{policy}') + return save + if name == 'forward_backward': + return lambda **kwargs: calls.append(('backward', len(kwargs['inputs']))) + if name == 'clip_grad_and_step': + return lambda **kwargs: calls.append(('step', 1)) + if name == 'calculate_metric': + return lambda **kwargs: {'result': {}} + return lambda *args, **kwargs: None + + class Sampler: + @staticmethod + def set_template(*args, **kwargs): + return None + + class Client: + def __init__(self): + self.model_instance = Model() + self.sampler_instance = Sampler() + self.closed = False + + def model(self, model_id): + return self.model_instance + + def sampler(self, model_id): + return self.sampler_instance + + def close(self): + self.closed = True + + class Advantages: + @staticmethod + def tolist(): + return [1.0, -1.0, 0.5, -0.5] + + client = Client() + monkeypatch.setattr(recipe, 'init_twinkle_client', lambda **kwargs: client) + monkeypatch.setattr(recipe, 'Qwen3_5Template', lambda **kwargs: SimpleNamespace(truncation_strategy='delete')) + monkeypatch.setattr(recipe, 'ClientMultiTurnRollout', lambda *args, **kwargs: object()) + monkeypatch.setattr(recipe, 'GRPOAdvantage', lambda: lambda *args, **kwargs: Advantages()) + monkeypatch.setattr(recipe, 'MAX_GROUPS', 2) + + async def fake_rollout(rollout, uri, environments, stop_event): + rolled = [{ + 'response_token_ids': [10], + 'rollout_logprobs': [-0.1], + 'stop_reason': 'stop', + 'truncated': False, + } for _ in range(recipe.NUM_GENERATIONS)] + return rolled, [1.0, 0.0, 1.0, 0.0] + + monkeypatch.setattr(recipe, '_rollout_group', fake_rollout) + asyncio.run(recipe.train()) + + assert calls.count(('step', 1)) == 2 + assert [call for call in calls if call[0] == 'save'] == [ + ('save', 'policy-0'), ('save', 'policy-1'), ('save', 'policy-2')] + assert client.closed + + +def test_client_connection_cancellation_registers_and_closes_client(recipe, monkeypatch): + started, release = threading.Event(), threading.Event() + + class Client: + closed = False + + def model(self, model_id): + raise AssertionError('no model should be created after cancellation') + + def close(self): + self.closed = True + + client = Client() + + def connect(**kwargs): + started.set() + assert release.wait(timeout=5) + return client + + monkeypatch.setattr(recipe, 'init_twinkle_client', connect) + + async def run(): + task = asyncio.create_task(recipe.train()) + await asyncio.to_thread(_cancel_and_release, task, started, release) + with pytest.raises(asyncio.CancelledError): + await task + + asyncio.run(run()) + assert client.closed + + +def test_template_is_built_off_the_event_loop_from_template_model_id(recipe, monkeypatch): + loop_threads = [] + template_calls = [] + + class Client: + closed = False + + def model(self, model_id): + return SimpleNamespace( + add_adapter_to_model=lambda *args, **kwargs: None, + set_loss=lambda *args, **kwargs: None, + set_optimizer=lambda *args, **kwargs: None, + set_processor=lambda *args, **kwargs: None, + set_template=lambda *args, **kwargs: template_calls.append(('model', kwargs['model_id'])), + ) + + def sampler(self, model_id): + return SimpleNamespace( + set_template=lambda *args, **kwargs: template_calls.append(('sampler', kwargs['model_id']))) + + def close(self): + self.closed = True + + def template(**kwargs): + template_calls.append(('local', kwargs['model_id'])) + loop_threads.append(threading.current_thread() is threading.main_thread()) + raise RuntimeError('stop after template') + + client = Client() + monkeypatch.setattr(recipe, 'TEMPLATE_MODEL_ID', '/models/local-qwen') + monkeypatch.setattr(recipe, 'init_twinkle_client', lambda **kwargs: client) + monkeypatch.setattr(recipe, 'Qwen3_5Template', template) + with pytest.raises(RuntimeError, match='stop after template'): + asyncio.run(recipe.train()) + assert template_calls == [ + ('model', '/models/local-qwen'), ('sampler', '/models/local-qwen'), ('local', '/models/local-qwen')] + assert loop_threads == [False] + assert client.closed + + +class _Files: + + def __init__(self): + self.data = {} + self.write_error = None + self.read_error = None + self.write_calls = 0 + + def write(self, path, content): + self.write_calls += 1 + self.data[path] = content + if self.write_error is not None: + raise self.write_error + + def read(self, path): + if self.read_error is not None: + raise self.read_error + return self.data[path] + + +def _make_env(recipe): + env = object.__new__(recipe._TaskAgentEnv) + env._workspace = '/workspace' + env._sandbox = SimpleNamespace(sandbox_id='sandbox-test', files=_Files(), set_timeout=lambda timeout: None) + env._task_infrastructure_error = None + env._execution_lock = threading.RLock() + env._active_commands = 0 + env._uncertain_execution = None + env._command_timeout = 30 + env._refresh_timeout = False + env._pre_tool_call = None + env._custom_handlers = {} + env._custom_tools = [] + env._include_default_tools = True + env.task = recipe.TASK + return env + + +def test_attempt_advertises_required_tool_arguments(recipe, monkeypatch): + env = _make_env(recipe) + env.reset = lambda: None + monkeypatch.setattr(recipe, '_TaskAgentEnv', lambda *args, **kwargs: env) + created = [] + attempt = recipe._create_attempt(recipe.TASK, created.append) + assert created == [env] + functions = {info['function']['name']: info['function'] for info in attempt.trajectory['tools']} + assert set(functions) == {'read_file', 'write_file', 'run_command'} + assert functions['read_file']['parameters']['required'] == ['path'] + assert functions['write_file']['parameters']['required'] == ['path', 'content'] + assert functions['run_command']['parameters']['required'] == ['command'] + assert all(fn['parameters']['properties'] for fn in functions.values()) + + +def test_file_side_effect_response_loss_is_unknown_and_not_retried(recipe): + env = _make_env(recipe) + env.sandbox.files.write_error = ConnectionError('response lost') + + result = env.step('write_file', {'path': 'solution.py', 'content': 'new content'}) + + assert result.info['error'] + assert env.sandbox.files.write_calls == 1 + assert env.sandbox.files.data['/workspace/solution.py'] == 'new content' + with pytest.raises(recipe._TaskInfrastructureError, match='tool:write_file'): + env.ensure_idle() + + +def test_read_failure_and_local_model_error_have_distinct_health_semantics(recipe): + env = _make_env(recipe) + env.sandbox.files.read_error = ConnectionError('read lost') + result = env.step('read_file', {'path': 'solution.py'}) + assert result.info['error'] + with pytest.raises(recipe._TaskInfrastructureError, match='tool:read_file'): + env.ensure_idle() + + clean = _make_env(recipe) + refusal = clean.step('write_file', {'path': '../outside', 'content': 'x'}) + assert refusal.observation.startswith('Error:') + assert not refusal.info + clean.ensure_idle() + assert clean.sandbox.files.write_calls == 0 + + +def test_verification_upload_and_command_failures_are_not_low_rewards(recipe): + env = _make_env(recipe) + env.sandbox.files.write_error = ConnectionError('upload failed') + with pytest.raises(recipe._TaskInfrastructureError, match='verify_upload'): + env.verify() + + command_env = _make_env(recipe) + command_env._execute = lambda command, cwd, timeout: (1, 'verifier did not start') + with pytest.raises(recipe._TaskInfrastructureError, match='verify_command'): + command_env.verify() + + +def test_verification_result_must_match_call_and_statistics(recipe): + env = _make_env(recipe) + payload = {'call_id': 'wrong-call', 'completed': True, 'passed': 5, 'total': 5} + env._execute = lambda command, cwd, timeout: (0, 'TWINKLE_VERIFY=' + json.dumps(payload)) + with pytest.raises(recipe._TaskInfrastructureError, match='verify_result'): + env.verify() + + +def test_candidate_exceptions_receive_low_reward_only_from_valid_verifier(recipe, tmp_path): + env = _make_env(recipe) + (tmp_path / 'solution.py').write_text( + 'def smallest_positive(values):\n raise ValueError("candidate failed")\n', encoding='utf-8') + + def execute(command, cwd, timeout): + uploaded = next(path for path in env.sandbox.files.data if path.startswith('/tmp/twinkle_scripts/verify-')) + source = env.sandbox.files.data[uploaded] + script = tmp_path / 'verify.py' + script.write_text(source, encoding='utf-8') + result = subprocess.run([sys.executable, str(script)], cwd=tmp_path, capture_output=True, text=True) + return result.returncode, result.stdout + result.stderr + + env._execute = execute + assert env.verify() == 0.0 + + +def _run_verifier_locally(env, tmp_path): + + def execute(command, cwd, timeout): + uploaded = next(path for path in env.sandbox.files.data if path.startswith('/tmp/twinkle_scripts/verify-')) + script = tmp_path / 'verify.py' + script.write_text(env.sandbox.files.data.pop(uploaded), encoding='utf-8') + result = subprocess.run([sys.executable, str(script)], cwd=tmp_path, capture_output=True, text=True) + return result.returncode, result.stdout + result.stderr + + env._execute = execute + return env.verify() + + +def test_verifier_scores_the_attempt_task_file_function_and_cases(recipe, tmp_path): + task = recipe.CodeTask( + filename='count_even.py', + function='count_even', + solution='', + smoke_test='', + prompt='', + hidden_cases=(([], 0), ([2, 4], 2), ([1, 3], 0), ([-2, 0, 7], 2)), + ) + env = _make_env(recipe) + env.task = task + (tmp_path / 'count_even.py').write_text( + 'def count_even(values):\n return sum(value % 2 == 0 for value in values)\n', encoding='utf-8') + assert _run_verifier_locally(env, tmp_path) == 1.0 + + (tmp_path / 'count_even.py').write_text( + 'def count_even(values):\n return sum(value % 2 == 1 for value in values)\n', encoding='utf-8') + assert _run_verifier_locally(env, tmp_path) == 0.25 + + +def test_verification_timeout_is_infrastructure_unknown(recipe): + env = _make_env(recipe) + + class Commands: + @staticmethod + def run(command, cwd, timeout): + raise TimeoutError('provider wait expired') + + env.sandbox.commands = Commands() + with pytest.raises(recipe._TaskInfrastructureError, match='verify_command'): + env.verify() + assert env._uncertain_execution is not None + + +def test_training_row_validation_rejects_missing_or_misaligned_probabilities(recipe): + with pytest.raises(RuntimeError, match='logprobs'): + recipe._validate_training_rows([{ + 'response_token_ids': [10, 11], + 'rollout_logprobs': [-0.1], + 'stop_reason': 'stop', + 'truncated': False, + }])