diff --git a/src/art/tau_bench/rollout.py b/src/art/tau_bench/rollout.py index be9f855f8..20c05f9ab 100644 --- a/src/art/tau_bench/rollout.py +++ b/src/art/tau_bench/rollout.py @@ -1,34 +1,18 @@ from __future__ import annotations -import asyncio from collections.abc import Mapping import json -import logging import os import time from typing import Any, cast, overload import httpx -from openai import ( - APIConnectionError, - APIError, - APITimeoutError, - AsyncOpenAI, - AsyncStream, - BadRequestError, - DefaultAsyncHttpxClient, -) +from openai import AsyncOpenAI, BadRequestError, DefaultAsyncHttpxClient from openai.types.chat import ChatCompletionMessageParam -from openai.types.chat.chat_completion import ChatCompletion -from openai.types.chat.chat_completion_chunk import ChatCompletionChunk from openai.types.completion_usage import CompletionUsage from art.costs import get_model_pricing, tokens_to_cost from art.model import Model -from art.openai import ( - IncompleteChatCompletionStreamError, - consume_chat_completion_stream, -) from art.trajectories import Trajectory from .client import Scenario, TauBenchClient, _get_default_client @@ -37,7 +21,8 @@ CONTEXT_TOKEN_LIMIT = 32_768 DEFAULT_MAX_COMPLETION_TOKENS = 4096 _POLICY_CONNECTION_LIMIT = 2048 -_STREAM_RETRY_STATUS_CODES = {408, 409, 429, 500, 502, 503, 504} +_POLICY_MAX_RETRIES = 1 +_POLICY_HTTP_TIMEOUT = httpx.Timeout(connect=30, read=10 * 60, write=30, pool=30) @overload @@ -117,29 +102,6 @@ async def rollout( model=model, base_model=base_model, ) - policy_completion_kwargs = dict(chat_completion_kwargs) - stream_policy = bool( - policy_completion_kwargs.setdefault( - "stream", isinstance(base_url_or_model, str) - ) - ) - if stream_policy: - stream_options = dict( - cast( - Mapping[str, Any], - policy_completion_kwargs.get("stream_options") or {}, - ) - ) - stream_options["include_usage"] = True - policy_completion_kwargs["stream_options"] = stream_options - extra_headers = dict( - cast( - Mapping[str, str], - policy_completion_kwargs.get("extra_headers") or {}, - ) - ) - extra_headers.setdefault("X-ART-Stream-Progress", "sse-comments-v1") - policy_completion_kwargs["extra_headers"] = extra_headers messages: list[ChatCompletionMessageParam] = [ {"role": "system", "content": env.info["policy"]}, {"role": "user", "content": env.observation.removeprefix("user: ")}, @@ -162,14 +124,13 @@ async def rollout( break try: policy_started = time.perf_counter() - chat_completion = await _create_policy_completion( - openai_client, + chat_completion = await openai_client.chat.completions.create( messages=messages, model=model_name, - stream_policy=stream_policy, + stream=False, tool_choice="auto", tools=tools, - **policy_completion_kwargs, + **chat_completion_kwargs, ) policy_latency = time.perf_counter() - policy_started trajectory.metrics["latency/policy"] = ( @@ -179,7 +140,7 @@ async def rollout( trajectory.metrics.get("latency/policy_max", 0.0), policy_latency, ) - except (BadRequestError, APIError) as exc: + except BadRequestError as exc: if _is_max_tokens_error(exc): break raise @@ -263,58 +224,6 @@ async def rollout( return trajectory -async def _create_policy_completion( - openai_client: AsyncOpenAI, - *, - stream_policy: bool, - **kwargs: Any, -) -> ChatCompletion: - attempts = max(1, int(getattr(openai_client, "max_retries", 2)) + 1) - for attempt in range(attempts): - completion = await openai_client.chat.completions.create(**kwargs) - if not stream_policy: - return cast(ChatCompletion, completion) - try: - return await consume_chat_completion_stream( - cast(AsyncStream[ChatCompletionChunk], completion), - require_usage=True, - ) - except Exception as error: - if attempt == attempts - 1 or not _retryable_stream_error(error): - raise - delay = 0.25 * (2**attempt) - logging.warning( - "Retrying streamed policy completion after %s", - type(error).__name__, - ) - await asyncio.sleep(delay) - raise AssertionError("unreachable") - - -def _retryable_stream_error(error: Exception) -> bool: - if isinstance( - error, - ( - IncompleteChatCompletionStreamError, - APIConnectionError, - APITimeoutError, - httpx.TransportError, - json.JSONDecodeError, - ), - ): - return True - if not isinstance(error, APIError): - return False - body = getattr(error, "body", None) - code = body.get("code") if isinstance(body, Mapping) else None - if not isinstance(code, (int, str)): - return False - try: - return int(code) in _STREAM_RETRY_STATUS_CODES - except (TypeError, ValueError): - return False - - def _completion_client_and_model( base_url_or_model: str | Model, *, @@ -336,11 +245,13 @@ def _completion_client_and_model( openai_clients[key] = AsyncOpenAI( api_key=api_key, base_url=base_url_or_model, + max_retries=_POLICY_MAX_RETRIES, http_client=DefaultAsyncHttpxClient( + timeout=_POLICY_HTTP_TIMEOUT, limits=httpx.Limits( max_connections=_POLICY_CONNECTION_LIMIT, max_keepalive_connections=_POLICY_CONNECTION_LIMIT, - ) + ), ), ) return openai_clients[key], model, base_model @@ -378,7 +289,7 @@ def _record_tinker_costs( ) -def _is_max_tokens_error(exc: APIError) -> bool: +def _is_max_tokens_error(exc: BadRequestError) -> bool: message = getattr(exc, "message", str(exc)) return "max_tokens" in message or "max_completion_tokens" in message diff --git a/src/art/trainer_rank/_checkpoint.py b/src/art/trainer_rank/_checkpoint.py index 9114d64ba..999a20b8a 100644 --- a/src/art/trainer_rank/_checkpoint.py +++ b/src/art/trainer_rank/_checkpoint.py @@ -1724,47 +1724,6 @@ def commit() -> None: raise -def export_lora(trainer: TrainerRank, output_dir: str, checkpoint_name: str) -> int: - group = _ensure_group(trainer) - slot = None - error: BaseException | None = None - try: - slot = trainer._checkpoint_slots.get(checkpoint_name) - if slot is None: - raise ValueError(f"Unknown checkpoint: {checkpoint_name!r}") - if slot.config is None: - raise trainer._slot_state_error( - f"Checkpoint {checkpoint_name!r} has no adapter_config" - ) - except BaseException as exc: - error = exc - raise_distributed(error, "validate LoRA export", group) - assert slot is not None and slot.config is not None - identity = (dict(slot.config), slot.revision) - if any(value != identity for value in _gather(identity, group)): - raise trainer._slot_state_error( - f"Checkpoint {checkpoint_name!r} differs across ranks" - ) - from art.megatron.weights.lora_publish import save_vllm_lora_from_model - - error = None - try: - save_vllm_lora_from_model( - model=trainer.runtime.model, - adapter_dtypes={}, - handler=trainer.runtime.model_support_handler, - adapter_config=dict(slot.config), - output_dir=output_dir, - rank=trainer.runtime.rank, - world_size=trainer.runtime.world_size, - slot_ref=trainer._slot_ref(checkpoint_name), - ) - except BaseException as exc: - error = exc - raise_distributed(error, "export LoRA", group) - return slot.revision - - def _ensure_groups( trainer: TrainerRank, ) -> tuple[dist.ProcessGroup | None, dist.ProcessGroup | None]: diff --git a/src/art/trainer_rank/_impl.py b/src/art/trainer_rank/_impl.py index 84c1a7ea2..80a2a37d1 100644 --- a/src/art/trainer_rank/_impl.py +++ b/src/art/trainer_rank/_impl.py @@ -46,6 +46,7 @@ estimate_prefix_tree_packed_tokens, prefix_tree_pack, ) +from art.trainer_rank._telemetry import phase as _telemetry_phase if TYPE_CHECKING: from megatron.core.models.gpt.gpt_model import GPTModel @@ -66,6 +67,7 @@ _FinalizedSave, _PreparedSave, ) + from art.trainer_rank._lora_export import _PreparedLoraExport @dataclass(frozen=True) @@ -813,6 +815,7 @@ def __init__( self._default_slot_ref: LoRASlotRef | None = None self._slot_stack: list[LoRASlotRef] = [] self._checkpoint_slots: dict[str, _CheckpointSlot] = {} + self._prepared_lora_exports: dict[str, tuple[str, _PreparedLoraExport]] = {} self._checkpoint_prefetches: dict[str, asyncio.Task[PreparedCheckpoint]] = {} self._checkpoint_mutation_tail: asyncio.Task[None] | None = None self._checkpoint_process_group: dist.ProcessGroup | None = None @@ -1274,12 +1277,42 @@ def export_lora( output_dir: str, checkpoint_path: str | Literal["active"] = "active", ) -> int: - from . import _checkpoint + from . import _lora_export - return _checkpoint.export_lora( + return _lora_export.export_lora( self, output_dir, self._resolve_checkpoint_name(checkpoint_path) ) + def _prepare_lora_export( + self, + export_id: str, + checkpoint_path: str | Literal["active"] = "active", + *, + owner_id: str, + ) -> tuple[int, dict[str, float]]: + from . import _lora_export + + return _lora_export.prepare_lora_export( + self, + export_id, + self._resolve_checkpoint_name(checkpoint_path), + owner_id=owner_id, + ) + + def _finish_lora_export( + self, export_id: str, output_dir: str, *, owner_id: str + ) -> dict[str, float]: + from . import _lora_export + + return _lora_export.finish_lora_export( + self, export_id, output_dir, owner_id=owner_id + ) + + def _abort_lora_export(self, export_id: str, *, owner_id: str) -> None: + from . import _lora_export + + _lora_export.abort_lora_export(self, export_id, owner_id=owner_id) + @staticmethod def _checkpoint_source_key(path: str) -> str: return str(Path(path).resolve()) @@ -1541,9 +1574,13 @@ def _forward_micro_batches( self._validate_replicated_top_level_count(len(items)) start = 0 while start < len(items): - candidate = self._select_next_micro_batch( - items, start, checkpoint=checkpoint - ) + with _telemetry_phase( + "plan", + {"global_start": start, "global_remaining": len(items) - start}, + ): + candidate = self._select_next_micro_batch( + items, start, checkpoint=checkpoint + ) flat_outputs = iter( self._run_flat_plan_with_memory_tracking( candidate.plan, @@ -1557,23 +1594,30 @@ def _forward_micro_batches( self._last_global_micro_batch_size or 0, candidate.stats_global_count, ) - yield MicroBatch( - inputs=candidate.inputs, - outputs=outputs, - indices=candidate.indices, - stats=MicroBatchStats( - global_start=start, - global_stop=stop, - global_count=candidate.stats_global_count, - local_count=len(candidate.inputs), - packed_tokens=candidate.plan.packed_tokens, - logical_tokens=candidate.plan.logical_tokens, - estimated_required_bytes=candidate.check.estimated_required_bytes, - available_bytes=candidate.check.available_bytes, - rejected_candidates=candidate.rejected_candidates, - cold_start=candidate.cold_start, - ), - ) + with _telemetry_phase( + # This interval is controlled by the caller and normally contains + # loss construction and backward for the yielded microbatch. + "caller", + self._telemetry_signature(candidate.plan), + dedup_signature=self._telemetry_plan_signature(candidate.plan), + ): + yield MicroBatch( + inputs=candidate.inputs, + outputs=outputs, + indices=candidate.indices, + stats=MicroBatchStats( + global_start=start, + global_stop=stop, + global_count=candidate.stats_global_count, + local_count=len(candidate.inputs), + packed_tokens=candidate.plan.packed_tokens, + logical_tokens=candidate.plan.logical_tokens, + estimated_required_bytes=candidate.check.estimated_required_bytes, + available_bytes=candidate.check.available_bytes, + rejected_candidates=candidate.rejected_candidates, + cold_start=candidate.cold_start, + ), + ) start = stop @overload @@ -1689,11 +1733,15 @@ def optim_step( selected_checkpoints = self._selected_dynamic_checkpoints(checkpoints) if on_live_graphs == "error": self._guard_checkpoints_can_step(selected_checkpoints) - return self._dynamic_optim_step( - selected_checkpoints, - params=params, - scale_grads=scale_grads, - ) + with _telemetry_phase( + "optim", + {"checkpoint_count": len(selected_checkpoints)}, + ): + return self._dynamic_optim_step( + selected_checkpoints, + params=params, + scale_grads=scale_grads, + ) def _load_checkpoint_slot( self, @@ -2568,7 +2616,15 @@ def _run_flat_plan_with_memory_tracking( else: baseline = 0 try: - outputs = self._execute_flat_plan(plan) + with _telemetry_phase( + "forward", + self._telemetry_signature(plan), + dedup_signature=self._telemetry_plan_signature(plan), + synchronized=torch.cuda.is_available() and self.device.type == "cuda", + ): + outputs = self._execute_flat_plan(plan) + if torch.cuda.is_available() and self.device.type == "cuda": + torch.cuda.synchronize(self.device) except torch.cuda.OutOfMemoryError as exc: check = self._memory_check(plan) self._raise_memory_error( @@ -2579,11 +2635,35 @@ def _run_flat_plan_with_memory_tracking( ) raise AssertionError("unreachable") from exc if torch.cuda.is_available() and self.device.type == "cuda": - torch.cuda.synchronize(self.device) peak = int(torch.cuda.max_memory_allocated(self.device)) self._update_memory_profile(plan, max(0, peak - baseline)) return outputs + @staticmethod + def _telemetry_plan_signature(plan: _FlatForwardPlan) -> dict[str, object]: + return { + "topology": plan.signature.topology, + "shared_prefix_max_depth": plan.signature.shared_prefix_max_depth, + "slot_group_count": plan.signature.slot_group_count, + "request_mix": plan.signature.request_mix, + "grad_enabled": plan.signature.grad_enabled, + } + + @classmethod + def _telemetry_signature(cls, plan: _FlatForwardPlan) -> dict[str, object]: + return { + **cls._telemetry_plan_signature(plan), + "request_count": plan.request_count, + "packed_tokens": plan.packed_tokens, + "logical_tokens": plan.logical_tokens, + "group_packed_tokens": tuple( + int(group.packed.tokens.numel()) for group in plan.groups + ), + "group_segment_counts": tuple( + len(group.packed.segments) for group in plan.groups + ), + } + def _execute_flat_plan(self, plan: _FlatForwardPlan) -> list[AnyForwardOutput]: outputs = [ ForwardOutput(None, None, None, None) for _ in range(plan.request_count) diff --git a/src/art/trainer_rank/_lora_export.py b/src/art/trainer_rank/_lora_export.py new file mode 100644 index 000000000..b4e4e9695 --- /dev/null +++ b/src/art/trainer_rank/_lora_export.py @@ -0,0 +1,416 @@ +"""Trainer-rank lifecycle for pipelined LoRA publication.""" + +from __future__ import annotations + +from dataclasses import dataclass, replace +import time +from typing import TYPE_CHECKING, Any, TypeVar +import uuid + +import torch + +if TYPE_CHECKING: + from art.megatron.lora import LoraShardMeta, LoRASlotRef + from art.megatron.training.model_chunks import ModelChunks + from art.megatron.weights.lora_publish import PackedExpertShardMeta + from art.trainer_rank._impl import TrainerRank + +_K = TypeVar("_K") + + +@dataclass(frozen=True) +class _PreparedLoraExport: + inputs: _VllmLoraPublishInputs + + +@dataclass(frozen=True) +class _VllmLoraPublishPlan: + rank: int + device: torch.device + metadata: list[LoraShardMeta] + local_tensors: dict[str, torch.Tensor] + packed_expert_metadata: list[PackedExpertShardMeta] + local_packed_expert_tensors: dict[str, torch.Tensor] + handler: Any + adapter_config: dict[str, Any] + + +@dataclass(frozen=True) +class _VllmLoraPublishInputs: + metadata: list[LoraShardMeta] + tensors_by_owner_key: dict[tuple[int, str], torch.Tensor] + packed_expert_metadata: list[PackedExpertShardMeta] + packed_expert_tensors_by_owner_key: dict[tuple[int, str], torch.Tensor] + handler: Any + adapter_config: dict[str, Any] + + +class _PinnedCpuStager: + def __init__(self) -> None: + self._events: list[torch.cuda.Event] = [] + self._stream = torch.cuda.Stream() if torch.cuda.is_available() else None + + def stage(self, tensor: torch.Tensor) -> torch.Tensor: + source = tensor.detach() + if self._stream is None or not source.is_cuda: + return source.cpu() + + source = source.contiguous() + target = torch.empty_like(source, device="cpu", pin_memory=True) + source_stream = torch.cuda.current_stream(source.device) + self._stream.wait_stream(source_stream) + with torch.cuda.stream(self._stream): + target.copy_(source, non_blocking=True) + source.record_stream(self._stream) + event = torch.cuda.Event() + event.record(self._stream) + self._events.append(event) + return target + + def finish(self) -> None: + for event in self._events: + event.synchronize() + self._events.clear() + + +def _stage_tensor_mapping( + tensors: dict[_K, torch.Tensor], + stager: _PinnedCpuStager, +) -> dict[_K, torch.Tensor]: + grouped: dict[ + tuple[str, int | None, torch.dtype], list[tuple[_K, torch.Tensor]] + ] = {} + for key, tensor in tensors.items(): + group_key = (tensor.device.type, tensor.device.index, tensor.dtype) + grouped.setdefault(group_key, []).append((key, tensor)) + + staged: dict[_K, torch.Tensor] = {} + for group in grouped.values(): + ordered = sorted(group, key=lambda item: str(item[0])) + flat = torch.cat( + [tensor.detach().contiguous().view(-1) for _key, tensor in ordered] + ) + staged_flat = stager.stage(flat) + offset = 0 + for key, tensor in ordered: + numel = tensor.numel() + if key in staged: + raise RuntimeError(f"Duplicate staged LoRA tensor: {key}") + staged[key] = staged_flat.narrow(0, offset, numel).view(tensor.shape) + offset += numel + return staged + + +def _validate_vllm_lora_publish_runtime( + rank: int, world_size: int +) -> tuple[int, torch.device]: + from art.megatron.weights import lora_publish + + actual_rank, device = lora_publish._rank_and_device() + if lora_publish._distributed_ready(): + actual_world_size = torch.distributed.get_world_size() # type: ignore[possibly-missing-attribute] + if actual_rank != rank or actual_world_size != world_size: + raise RuntimeError( + "LoRA publisher rank/world-size mismatch: " + f"runtime=({rank}, {world_size}) " + f"distributed=({actual_rank}, {actual_world_size})" + ) + else: + if rank != 0 or world_size != 1: + raise RuntimeError( + "Non-distributed LoRA publish requires rank=0 and world_size=1, " + f"got rank={rank} world_size={world_size}" + ) + rank = 0 + return rank, device + + +def _prepare_vllm_lora_publish( + *, + model: ModelChunks, + adapter_dtypes: dict[str, torch.dtype], + handler: Any, + adapter_config: dict[str, Any], + rank: int, + world_size: int, + slot_ref: LoRASlotRef | None = None, + runtime: tuple[int, torch.device] | None = None, +) -> _VllmLoraPublishPlan: + from art.megatron.lora import LoRAPublishPlanner + from art.megatron.weights import lora_publish + + rank, device = ( + _validate_vllm_lora_publish_runtime(rank, world_size) + if runtime is None + else runtime + ) + packed_expert_groups = tuple(handler.expert_packed_lora_groups()) + planner = LoRAPublishPlanner(model, slot_ref) + local_tensors, local_metadata = lora_publish.collect_local_lora_entries( + model, + adapter_dtypes, + owner_rank=rank, + packed_expert_groups=packed_expert_groups, + slot_ref=slot_ref, + ) + ( + local_packed_tensors, + local_packed_metadata, + ) = lora_publish.collect_local_packed_expert_entries( + model, + adapter_dtypes, + owner_rank=rank, + packed_expert_groups=packed_expert_groups, + slot_ref=slot_ref, + ) + all_packed_metadata = ( + lora_publish._global_packed_expert_metadata( + planner, adapter_dtypes, packed_expert_groups + ) + if rank == 0 + else local_packed_metadata + ) + all_metadata = ( + lora_publish._global_regular_metadata( + planner, + adapter_dtypes, + packed_expert_groups if all_packed_metadata else (), + ) + if rank == 0 + else local_metadata + ) + return _VllmLoraPublishPlan( + rank=rank, + device=device, + metadata=all_metadata, + local_tensors=local_tensors, + packed_expert_metadata=all_packed_metadata, + local_packed_expert_tensors=local_packed_tensors, + handler=handler, + adapter_config=dict(adapter_config), + ) + + +def _exchange_vllm_lora_publish( + plan: _VllmLoraPublishPlan, +) -> _VllmLoraPublishInputs | None: + from art.megatron.weights import lora_publish + + exchanged_tensors = lora_publish._exchange_batched_tensors( + plan.metadata, + local_tensors=plan.local_tensors, + rank=plan.rank, + device=plan.device, + ) + exchanged_packed_tensors = lora_publish._exchange_batched_tensors( + plan.packed_expert_metadata, + local_tensors=plan.local_packed_expert_tensors, + rank=plan.rank, + device=plan.device, + ) + if plan.rank != 0: + return None + return _VllmLoraPublishInputs( + metadata=plan.metadata, + tensors_by_owner_key=exchanged_tensors, + packed_expert_metadata=plan.packed_expert_metadata, + packed_expert_tensors_by_owner_key=exchanged_packed_tensors, + handler=plan.handler, + adapter_config=plan.adapter_config, + ) + + +def _build_vllm_lora_tensors_from_inputs( + inputs: _VllmLoraPublishInputs, +) -> tuple[dict[str, torch.Tensor], dict[str, Any]]: + from art.megatron.weights import lora_publish + + return lora_publish._rank0_vllm_lora_tensors( + metadata=inputs.metadata, + tensors_by_owner_key=inputs.tensors_by_owner_key, + packed_expert_metadata=inputs.packed_expert_metadata, + packed_expert_tensors_by_owner_key=inputs.packed_expert_tensors_by_owner_key, + handler=inputs.handler, + adapter_config=inputs.adapter_config, + ) + + +def _capture_lora_publish_inputs( + trainer: TrainerRank, + checkpoint_name: str, + adapter_config: dict[str, object], + group: torch.distributed.ProcessGroup | None, +) -> tuple[_PreparedLoraExport | None, dict[str, float]]: + from art.trainer_rank import _checkpoint + + timings: dict[str, float] = {} + started = time.monotonic() + runtime = _checkpoint._phase( + lambda: _validate_vllm_lora_publish_runtime( + trainer.runtime.rank, trainer.runtime.world_size + ), + "validate LoRA publish runtime", + group, + ) + timings["runtime_validation"] = time.monotonic() - started + + started = time.monotonic() + plan = _checkpoint._phase( + lambda: _prepare_vllm_lora_publish( + model=trainer.runtime.model, + adapter_dtypes={}, + handler=trainer.runtime.model_support_handler, + adapter_config=adapter_config, + rank=trainer.runtime.rank, + world_size=trainer.runtime.world_size, + slot_ref=trainer._slot_ref(checkpoint_name), + runtime=runtime, + ), + "plan LoRA publish", + group, + ) + timings["plan_collect"] = time.monotonic() - started + + started = time.monotonic() + inputs = _exchange_vllm_lora_publish(plan) + timings["exchange"] = time.monotonic() - started + + started = time.monotonic() + + def stage() -> _PreparedLoraExport | None: + if inputs is not None: + stager = _PinnedCpuStager() + staged = replace( + inputs, + tensors_by_owner_key=_stage_tensor_mapping( + inputs.tensors_by_owner_key, stager + ), + packed_expert_tensors_by_owner_key=_stage_tensor_mapping( + inputs.packed_expert_tensors_by_owner_key, stager + ), + ) + stager.finish() + return _PreparedLoraExport(staged) + return None + + prepared = _checkpoint._phase(stage, "stage LoRA publish tensors", group) + timings["d2h"] = time.monotonic() - started + return prepared, timings + + +def _save_lora_publish_inputs( + output_dir: str, prepared: _PreparedLoraExport +) -> dict[str, float]: + from art.megatron.model_support.lora_disk import save_vllm_lora_tensors + + started = time.monotonic() + vllm_tensors, published_config = _build_vllm_lora_tensors_from_inputs( + prepared.inputs + ) + timings = {"convert": time.monotonic() - started} + started = time.monotonic() + save_vllm_lora_tensors(output_dir, vllm_tensors, published_config) + timings["serialize"] = time.monotonic() - started + return timings + + +def prepare_lora_export( + trainer: TrainerRank, + export_id: str, + checkpoint_name: str, + *, + owner_id: str, +) -> tuple[int, dict[str, float]]: + from art.trainer_rank import _checkpoint + + started = time.monotonic() + group = _checkpoint._ensure_group(trainer) + snapshots: dict[str, tuple[str, _PreparedLoraExport]] = getattr( + trainer, "_prepared_lora_exports", {} + ) + duplicate = ( + RuntimeError(f"LoRA export {export_id!r} is already prepared") + if export_id in snapshots + else None + ) + _checkpoint.raise_distributed(duplicate, "validate LoRA export ID", group) + slot = None + error: BaseException | None = None + try: + slot = trainer._checkpoint_slots.get(checkpoint_name) + if slot is None: + raise ValueError(f"Unknown checkpoint: {checkpoint_name!r}") + if slot.config is None: + raise trainer._slot_state_error( + f"Checkpoint {checkpoint_name!r} has no adapter_config" + ) + except BaseException as exc: + error = exc + _checkpoint.raise_distributed(error, "validate LoRA export", group) + assert slot is not None and slot.config is not None + identity = (dict(slot.config), slot.revision) + if any(value != identity for value in _checkpoint._gather(identity, group)): + raise trainer._slot_state_error( + f"Checkpoint {checkpoint_name!r} differs across ranks" + ) + slot_validation = time.monotonic() - started + + prepared = None + capture_timings: dict[str, float] = {} + error = None + try: + prepared, capture_timings = _capture_lora_publish_inputs( + trainer, checkpoint_name, dict(slot.config), group + ) + except BaseException as exc: + error = exc + _checkpoint.raise_distributed(error, "prepare LoRA export", group) + if prepared is not None: + snapshots[export_id] = (owner_id, prepared) + trainer._prepared_lora_exports = snapshots + return slot.revision, {"slot_validation": slot_validation, **capture_timings} + + +def finish_lora_export( + trainer: TrainerRank, export_id: str, output_dir: str, *, owner_id: str +) -> dict[str, float]: + snapshots: dict[str, tuple[str, _PreparedLoraExport]] = getattr( + trainer, "_prepared_lora_exports", {} + ) + try: + owner, prepared = snapshots[export_id] + except KeyError: + raise ValueError(f"Unknown prepared LoRA export: {export_id!r}") from None + if owner != owner_id: + raise ValueError(f"LoRA export {export_id!r} belongs to another owner") + snapshots.pop(export_id) + return _save_lora_publish_inputs(output_dir, prepared) + + +def abort_lora_export(trainer: TrainerRank, export_id: str, *, owner_id: str) -> None: + snapshots: dict[str, tuple[str, _PreparedLoraExport]] = getattr( + trainer, "_prepared_lora_exports", {} + ) + if (prepared := snapshots.get(export_id)) is not None and prepared[0] == owner_id: + snapshots.pop(export_id) + + +def export_lora(trainer: TrainerRank, output_dir: str, checkpoint_name: str) -> int: + from art.trainer_rank import _checkpoint + + group = _checkpoint._ensure_group(trainer) + export_id = uuid.uuid4().hex + owner_id = uuid.uuid4().hex + revision, _timings = prepare_lora_export( + trainer, export_id, checkpoint_name, owner_id=owner_id + ) + error: BaseException | None = None + try: + if trainer.runtime.rank == 0: + finish_lora_export(trainer, export_id, output_dir, owner_id=owner_id) + except BaseException as exc: + error = exc + finally: + abort_lora_export(trainer, export_id, owner_id=owner_id) + _checkpoint.raise_distributed(error, "export LoRA", group) + return revision diff --git a/src/art/trainer_rank/_telemetry.py b/src/art/trainer_rank/_telemetry.py new file mode 100644 index 000000000..22fc117be --- /dev/null +++ b/src/art/trainer_rank/_telemetry.py @@ -0,0 +1,267 @@ +"""Structured host-phase and torch.compile telemetry for trainer-rank processes.""" + +from __future__ import annotations + +from collections.abc import Iterator, Mapping +from contextlib import contextmanager +from dataclasses import dataclass, field +import json +import logging +import threading +import time +from typing import Any, TypedDict + +import torch + +logger = logging.getLogger("art.trainer_rank.telemetry") + +_install_lock = threading.Lock() +_state_lock = threading.Lock() +_installed = False +_guard_counts: dict[object, int] = {} +_compiled_plan_signatures: set[str] = set() +_active_compile: _Compile | None = None +_thread_state = threading.local() +_SLOW_PHASE_SECONDS = 10.0 + + +@dataclass +class _Phase: + name: str + signature: Mapping[str, object] + signature_key: str + start: float + compiles: list[_CompileRecord] = field(default_factory=list) + closed: bool = False + + +class _CompileRecord(TypedDict): + compile_id: str + trigger: str + seconds: float + guard_failures: tuple[str, ...] + frame_id: int | None + frame_compile_id: int | None + graph_status: str + + +@dataclass(frozen=True) +class _Compile: + start: float + compile_id: str + trigger: str + guard_failures: tuple[str, ...] + phase: _Phase | None + + +def _compile_identity(compile_id: str) -> tuple[int | None, int | None, str]: + """Extract Dynamo's frame and per-frame compilation numbers.""" + + parts = compile_id.removeprefix("!").split("/") + if compile_id.startswith("!") and len(parts) == 1: + return None, None, "compiled_autograd" + try: + frame_id, frame_compile_id = int(parts[-2]), int(parts[-1]) + except (IndexError, ValueError): + return None, None, "unknown" + return ( + frame_id, + frame_compile_id, + "new_graph" if frame_compile_id == 0 else "recompile", + ) + + +def _emit(event: Mapping[str, object], *, warning: bool = False) -> None: + log = logger.warning if warning else logger.info + log("ART_TRAINER_EVENT %s", json.dumps(event, sort_keys=True, default=str)) + + +def _phase_stack() -> list[_Phase]: + stack = getattr(_thread_state, "phases", None) + if stack is None: + stack = [] + _thread_state.phases = stack + return stack + + +def _new_guard_failures() -> tuple[str, ...]: + reasons: list[str] = [] + failures = torch._dynamo.guard_failures + with _state_lock: + for code, values in failures.items(): + start = _guard_counts.get(code, 0) + reasons.extend( + str(getattr(value, "reason", value))[:1000] for value in values[start:] + ) + _guard_counts[code] = len(values) + return tuple(reasons) + + +def _compile_start(args: Any) -> None: + global _active_compile + try: + trigger = getattr(getattr(args, "callback_trigger", None), "name", None) + stack = _phase_stack() + active = _Compile( + start=time.perf_counter(), + compile_id=str(getattr(args, "compile_id", "unknown")), + trigger=str(trigger or getattr(args, "callback_trigger", "unknown")), + guard_failures=_new_guard_failures(), + phase=stack[-1] if stack else None, + ) + with _state_lock: + _active_compile = active + except Exception: + logger.debug("Failed to begin ART compile telemetry", exc_info=True) + + +def _compile_end(_args: Any) -> None: + global _active_compile + try: + with _state_lock: + completed = _active_compile + _active_compile = None + if not isinstance(completed, _Compile): + return + seconds = max(0.0, time.perf_counter() - completed.start) + frame_id, frame_compile_id, graph_status = _compile_identity( + completed.compile_id + ) + record: _CompileRecord = { + "compile_id": completed.compile_id, + "trigger": completed.trigger, + "seconds": seconds, + "guard_failures": completed.guard_failures, + "frame_id": frame_id, + "frame_compile_id": frame_compile_id, + "graph_status": graph_status, + } + closed_phase: _Phase | None = None + repeated_plan = False + unique_signatures = 0 + if completed.phase is not None: + with _state_lock: + if completed.phase.closed: + closed_phase = completed.phase + repeated_plan = ( + closed_phase.signature_key in _compiled_plan_signatures + ) + _compiled_plan_signatures.add(closed_phase.signature_key) + unique_signatures = len(_compiled_plan_signatures) + else: + completed.phase.compiles.append(record) + if closed_phase is None: + return + event: dict[str, object] = { + "event": "compile", + "phase": closed_phase.name if closed_phase is not None else "unscoped", + **record, + } + if closed_phase is not None: + event.update( + { + "signature": closed_phase.signature, + "phase_closed": True, + "plan_signature_status": ("repeated" if repeated_plan else "new"), + "unique_compile_plan_signatures": unique_signatures, + } + ) + _emit(event, warning=True) + except Exception: + logger.debug("Failed to finish ART compile telemetry", exc_info=True) + + +def _install() -> None: + global _installed + with _install_lock: + handler = torch._dynamo.callback_handler + if ( + _compile_start in handler.start_callbacks + and _compile_end in handler.end_callbacks + ): + _installed = True + return + with _state_lock: + _guard_counts.clear() + _guard_counts.update( + (code, len(values)) + for code, values in torch._dynamo.guard_failures.items() + ) + if _compile_start not in handler.start_callbacks: + handler.register_start_callback(_compile_start) + if _compile_end not in handler.end_callbacks: + handler.register_end_callback(_compile_end) + _installed = True + + +@contextmanager +def phase( + name: str, + signature: Mapping[str, object], + *, + dedup_signature: Mapping[str, object] | None = None, + synchronized: bool = False, +) -> Iterator[None]: + """Emit one structured phase event, including compile work observed within it.""" + + _install() + signature_key = json.dumps( + { + "phase": name, + "signature": signature if dedup_signature is None else dedup_signature, + }, + sort_keys=True, + default=str, + ) + record = _Phase(name, signature, signature_key, time.perf_counter()) + stack = _phase_stack() + stack.append(record) + error: BaseException | None = None + try: + yield + except BaseException as exc: + error = exc + raise + finally: + if not stack or stack.pop() is not record: + logger.debug("Trainer telemetry phase stack changed unexpectedly") + with _state_lock: + record.closed = True + compiles = list(record.compiles) + repeated_plan = False + if compiles: + repeated_plan = record.signature_key in _compiled_plan_signatures + _compiled_plan_signatures.add(record.signature_key) + unique_signatures = len(_compiled_plan_signatures) + seconds = max(0.0, time.perf_counter() - record.start) + compile_seconds = sum(item["seconds"] for item in compiles) + graph_statuses = {str(item["graph_status"]) for item in compiles} + _emit( + { + "event": "phase", + "phase": name, + "seconds": seconds, + "synchronized": synchronized, + "signature": signature, + "compile_status": ( + "recompile" + if "recompile" in graph_statuses + else "new_graph" + if "new_graph" in graph_statuses + else "compiled_autograd" + if "compiled_autograd" in graph_statuses + else "unknown" + if compiles + else "none" + ), + "compile_seconds": compile_seconds, + "compiles": compiles, + "plan_signature_status": ( + "repeated" if repeated_plan else "new" if compiles else "none" + ), + "unique_compile_plan_signatures": unique_signatures, + "outcome": "error" if error is not None else "ok", + "error_type": type(error).__name__ if error is not None else None, + }, + warning=bool(compiles) or seconds >= _SLOW_PHASE_SECONDS, + ) diff --git a/src/art/trajectories/__init__.py b/src/art/trajectories/__init__.py index 2da1bcb7c..0925f2bb6 100644 --- a/src/art/trajectories/__init__.py +++ b/src/art/trajectories/__init__.py @@ -1336,14 +1336,22 @@ async def trajectory(coroutine: Coroutine[Any, Any, object]) -> Trajectory: async def trajectory_group( - trajectories: Iterable[Coroutine[Any, Any, Trajectory]], + trajectories: Iterable[Trajectory | BaseException | Awaitable[Trajectory]], *, + exceptions: Iterable[BaseException | PydanticException] = (), + metadata: dict[str, MetadataValue] | None = None, + metrics: dict[str, float | int | bool] | None = None, + logs: list[str] | None = None, return_exceptions: bool = False, ) -> TrajectoryGroup: from ._scope import capture_trajectory_group return await capture_trajectory_group( trajectories, + exceptions=exceptions, + metadata=metadata, + metrics=metrics, + logs=logs, return_exceptions=return_exceptions, ) diff --git a/src/art/trajectories/_history.py b/src/art/trajectories/_history.py index 3d4aa3dac..3d174e7a7 100644 --- a/src/art/trajectories/_history.py +++ b/src/art/trajectories/_history.py @@ -189,7 +189,7 @@ def _ordered_choices(choices: Sequence[_IndexedT], *, protocol: str) -> list[_In def _is_prefix(prefix: Sequence[object], value: Sequence[object]) -> bool: return len(prefix) <= len(value) and all( - left == right for left, right in zip(prefix, value[: len(prefix)], strict=True) + left == right for left, right in zip(prefix, value) ) @@ -310,11 +310,37 @@ def _extend_branches( branches.extend(created) +def _token_prefix_lengths(tokens: Sequence[int]) -> list[int]: + prefix_lengths = [0] * len(tokens) + matched = 0 + for index in range(1, len(tokens)): + while matched and tokens[index] != tokens[matched]: + matched = prefix_lengths[matched - 1] + if tokens[index] == tokens[matched]: + matched += 1 + prefix_lengths[index] = matched + return prefix_lengths + + def _contains_tokens(tokens: Sequence[int], sampled: Sequence[int]) -> bool: - return any( - list(tokens[start : start + len(sampled)]) == list(sampled) - for start in range(len(tokens) - len(sampled) + 1) - ) + if not sampled: + return True + if len(sampled) > len(tokens): + return False + + # Knuth-Morris-Pratt keeps reconciliation linear for long sampled outputs. + # Comparing every slice blocks the event loop on long multi-turn histories. + prefix_lengths = _token_prefix_lengths(sampled) + + matched = 0 + for token in tokens: + while matched and token != sampled[matched]: + matched = prefix_lengths[matched - 1] + if token == sampled[matched]: + matched += 1 + if matched == len(sampled): + return True + return False def _retains_output_suffix( @@ -330,7 +356,21 @@ def _retains_output_suffix( ): return False continuation = later_prompt[len(prompt) :] - return any(_is_prefix(output[start:], continuation) for start in range(len(output))) + if not output or not continuation: + return False + + prefix_lengths = _token_prefix_lengths(continuation) + matched = 0 + for index, token in enumerate(output): + while matched and token != continuation[matched]: + matched = prefix_lengths[matched - 1] + if token == continuation[matched]: + matched += 1 + if matched == len(continuation): + if index + 1 == len(output): + return True + matched = prefix_lengths[matched - 1] + return matched > 0 def _chat_generation_tokens( diff --git a/src/art/trajectories/_scope.py b/src/art/trajectories/_scope.py index 6876dcac9..6dad842e7 100644 --- a/src/art/trajectories/_scope.py +++ b/src/art/trajectories/_scope.py @@ -1,15 +1,14 @@ from __future__ import annotations import asyncio -from collections.abc import Coroutine, Iterable, Iterator +from collections.abc import Awaitable, Coroutine, Iterable, Iterator from contextlib import contextmanager import contextvars from dataclasses import dataclass from types import TracebackType from typing import Any -from . import PydanticException, Trajectory, TrajectoryGroup -from ._compat import exception_model +from . import MetadataValue, PydanticException, Trajectory, TrajectoryGroup @dataclass(frozen=True, slots=True) @@ -83,25 +82,72 @@ async def capture_trajectory(coroutine: Coroutine[Any, Any, object]) -> Trajecto async def capture_trajectory_group( - trajectories: Iterable[Coroutine[Any, Any, Trajectory]], + trajectories: Iterable[Trajectory | BaseException | Awaitable[Trajectory]], *, + exceptions: Iterable[BaseException | PydanticException], + metadata: dict[str, MetadataValue] | None, + metrics: dict[str, float | int | bool] | None, + logs: list[str] | None, return_exceptions: bool, ) -> TrajectoryGroup: with no_capture(): - coroutines = list(trajectories) - for coroutine in coroutines: - _require_raw_coroutine(coroutine) - results = await asyncio.gather( - *coroutines, - return_exceptions=return_exceptions, - ) - if not return_exceptions: - return TrajectoryGroup(results) - completed: list[Trajectory] = [] - exceptions: list[PydanticException] = [] - for result in results: - if isinstance(result, BaseException): - exceptions.append(exception_model(result)) - else: - completed.append(result) - return TrajectoryGroup(completed, exceptions=exceptions) + provided_exceptions = list(exceptions) + for error in provided_exceptions: + if isinstance(error, BaseException) and not isinstance(error, Exception): + raise error + results: list[Trajectory | BaseException | None] = [] + pending: dict[asyncio.Future[Trajectory], list[int]] = {} + by_identity: dict[int, asyncio.Future[Trajectory]] = {} + try: + for item in trajectories: + index = len(results) + if isinstance(item, (Trajectory, BaseException)): + if isinstance(item, BaseException) and not isinstance( + item, Exception + ): + raise item + results.append(item) + continue + results.append(None) + task = by_identity.get(id(item)) + if task is None: + task = asyncio.ensure_future(item) + by_identity[id(item)] = task + pending[task] = [] + pending[task].append(index) + while pending: + done, _ = await asyncio.wait( + pending, return_when=asyncio.FIRST_COMPLETED + ) + for task in done: + indexes = pending.pop(task) + try: + result: Trajectory | BaseException = task.result() + except asyncio.CancelledError: + raise + except Exception as error: + if not return_exceptions: + raise + result = error + if not isinstance(result, (Trajectory, BaseException)): + raise TypeError( + "trajectory_group awaitables must resolve to trajectories" + ) + if isinstance(result, BaseException) and not isinstance( + result, Exception + ): + raise result + for index in indexes: + results[index] = result + except BaseException: + for task in pending: + task.cancel() + await asyncio.gather(*pending, return_exceptions=True) + raise + return TrajectoryGroup( + (result for result in results if result is not None), + exceptions=provided_exceptions, + metadata=metadata, + metrics=metrics, + logs=logs, + ) diff --git a/tests/integration/megatron/lora/test_lora_disk_codecs.py b/tests/integration/megatron/lora/test_lora_disk_codecs.py index 92606d6a8..ef9a28e85 100644 --- a/tests/integration/megatron/lora/test_lora_disk_codecs.py +++ b/tests/integration/megatron/lora/test_lora_disk_codecs.py @@ -37,7 +37,7 @@ merge_sharded_adapter_entries, save_vllm_lora_from_model, ) -from art.trainer_rank import TrainerRank +from art.trainer_rank import TrainerRank, _checkpoint, _lora_export from art.trainer_rank._impl import _AdapterConfig, _CheckpointSlot from art.utils.convert_moe_lora import convert_checkpoint_if_needed @@ -1623,6 +1623,89 @@ def test_save_vllm_lora_from_model_writes_single_vllm_checkpoint(tmp_path: Path) _assert_tensors_equal(roundtrip, full) +@pytest.mark.parametrize( + ("failed_phase", "expected_plans"), + (("validate LoRA publish runtime", 0), ("plan LoRA publish", 1)), +) +def test_remote_preexchange_failure_prevents_tensor_exchange( + monkeypatch: pytest.MonkeyPatch, + failed_phase: str, + expected_plans: int, +) -> None: + exchanges: list[object] = [] + plans: list[object] = [] + failure_group = cast(Any, object()) + + monkeypatch.setattr( + _lora_export, + "_validate_vllm_lora_publish_runtime", + lambda *_args: (0, torch.device("cpu")), + ) + monkeypatch.setattr( + _lora_export, + "_prepare_vllm_lora_publish", + lambda *_args, **_kwargs: plans.append(object()) or plans[-1], + ) + monkeypatch.setattr( + _lora_export, + "_exchange_vllm_lora_publish", + lambda *_args: exchanges.append(object()), + ) + + def synchronize(error: BaseException | None, phase: str, group: object) -> None: + assert group is failure_group + if error is not None: + raise error + if phase == failed_phase: + raise RuntimeError(f"Another rank failed to {phase}") + + monkeypatch.setattr(_checkpoint, "raise_distributed", synchronize) + trainer = SimpleNamespace( + runtime=SimpleNamespace( + rank=0, + world_size=2, + model=object(), + model_support_handler=object(), + ), + _slot_ref=lambda _name: None, + ) + + with pytest.raises(RuntimeError, match=f"Another rank failed to {failed_phase}"): + _lora_export._capture_lora_publish_inputs( + cast(Any, trainer), "student", {}, failure_group + ) + + assert exchanges == [] + assert len(plans) == expected_plans + + +def test_runtime_resolution_failure_uses_trainer_failure_group( + monkeypatch: pytest.MonkeyPatch, +) -> None: + failure_group = cast(Any, object()) + synchronized_groups: list[object] = [] + monkeypatch.setattr( + _lora_export, + "_validate_vllm_lora_publish_runtime", + lambda *_args: (_ for _ in ()).throw(RuntimeError("device resolution failed")), + ) + + def synchronize(error: BaseException | None, _phase: str, group: object) -> None: + synchronized_groups.append(group) + if error is not None: + raise error + + monkeypatch.setattr(_checkpoint, "raise_distributed", synchronize) + trainer = SimpleNamespace(runtime=SimpleNamespace(rank=0, world_size=2)) + + with pytest.raises(RuntimeError, match="device resolution failed"): + _lora_export._capture_lora_publish_inputs( + cast(Any, trainer), "student", {}, failure_group + ) + + assert synchronized_groups == [failure_group] + + def test_trainer_rank_publishes_named_checkpoint_slot_without_mutating_base( tmp_path: Path, ): @@ -1662,6 +1745,59 @@ def test_trainer_rank_publishes_named_checkpoint_slot_without_mutating_base( assert torch.equal(lora.B_T, baseline[1]) +def test_prepared_lora_export_is_immutable_and_abortable(tmp_path: Path): + prefix = "base_model.model.model.layers.0.self_attn.q_proj" + lora = LoRA(prefix, 3, 4, 2, 2, torch.float32, torch.device("cpu")) + adapter = { + f"{prefix}.lora_A.weight": torch.arange(6, dtype=torch.float32).reshape(2, 3), + f"{prefix}.lora_B.weight": torch.arange(8, dtype=torch.float32).reshape(4, 2), + } + trainer = TrainerRank.__new__(TrainerRank) + trainer.runtime = SimpleNamespace( + model=[lora], + model_support_handler=DEFAULT_DENSE_HANDLER, + rank=0, + world_size=1, + ) + trainer._slot_stack = [] + trainer._pending_slot_graphs = {} + trainer._checkpoint_slots = {} + config = _config("Qwen/Qwen3-8B", rank=2, alpha=2) + assert trainer._load_checkpoint_slot("student", adapter, alpha=2) == 1 + trainer._checkpoint_slots["student"] = _CheckpointSlot( + tuple(trainer._iter_slot_parameters(trainer._slot_ref("student"))), + cast(_AdapterConfig, config), + ) + + revision, capture_timings = trainer._prepare_lora_export( + "first", "student", owner_id="owner" + ) + for parameter in trainer._checkpoint_slots["student"].params: + parameter.data.fill_(99) + output_dir = tmp_path / "prepared" + finalize_timings = trainer._finish_lora_export( + "first", str(output_dir), owner_id="owner" + ) + + assert revision == 0 + assert set(capture_timings) == { + "slot_validation", + "runtime_validation", + "plan_collect", + "exchange", + "d2h", + } + assert set(finalize_timings) == {"convert", "serialize"} + _assert_tensors_equal(load_file(output_dir / "adapter_model.safetensors"), adapter) + with pytest.raises(ValueError, match="Unknown prepared LoRA export"): + trainer._finish_lora_export("first", str(output_dir), owner_id="owner") + + trainer._prepare_lora_export("aborted", "student", owner_id="owner") + trainer._abort_lora_export("aborted", owner_id="owner") + with pytest.raises(ValueError, match="Unknown prepared LoRA export"): + trainer._finish_lora_export("aborted", str(output_dir), owner_id="owner") + + @pytest.mark.parametrize( ("handler", "base_model"), ( diff --git a/tests/unit/test_tau_bench_client.py b/tests/unit/test_tau_bench_client.py index 6ccbd919b..c1a452746 100644 --- a/tests/unit/test_tau_bench_client.py +++ b/tests/unit/test_tau_bench_client.py @@ -1,14 +1,13 @@ from __future__ import annotations -from collections.abc import AsyncIterator import importlib import json from types import SimpleNamespace from typing import Any import httpx -from openai import APIError, AsyncOpenAI -from openai.types.chat import ChatCompletion, ChatCompletionChunk +from openai import AsyncOpenAI +from openai.types.chat import ChatCompletion import pytest import art @@ -251,92 +250,9 @@ async def delete_environment(self, env_id: str) -> DeleteEnvironmentResponse: return DeleteEnvironmentResponse(id=env_id, deleted=True) -class FakeStream: - def __init__(self, chunks: list[ChatCompletionChunk]) -> None: - self._chunks = iter(chunks) - - def __aiter__(self) -> AsyncIterator[ChatCompletionChunk]: - return self - - async def __anext__(self) -> ChatCompletionChunk: - try: - return next(self._chunks) - except StopIteration: - raise StopAsyncIteration from None - - async def close(self) -> None: - pass - - -class FailedStream: - def __aiter__(self) -> AsyncIterator[ChatCompletionChunk]: - return self - - async def __anext__(self) -> ChatCompletionChunk: - raise httpx.ReadError( - "connection reset", - request=httpx.Request("POST", "http://model.test/v1/chat/completions"), - ) - - async def close(self) -> None: - pass - - -class MaxTokensErrorStream: - def __aiter__(self) -> AsyncIterator[ChatCompletionChunk]: - return self - - async def __anext__(self) -> ChatCompletionChunk: - raise APIError( - message="max_tokens is too large for this model", - request=httpx.Request("POST", "http://model.test/v1/chat/completions"), - body={"code": 400}, - ) - - async def close(self) -> None: - pass - - class FakeCompletions: async def create(self, **kwargs: Any) -> Any: self.kwargs = kwargs - if kwargs.get("stream"): - return FakeStream( - [ - ChatCompletionChunk.model_validate( - { - "id": "chat-1", - "object": "chat.completion.chunk", - "created": 0, - "model": "default", - "choices": [ - { - "index": 0, - "finish_reason": "stop", - "delta": { - "role": "assistant", - "content": "hello", - }, - } - ], - } - ), - ChatCompletionChunk.model_validate( - { - "id": "chat-1", - "object": "chat.completion.chunk", - "created": 0, - "model": "default", - "choices": [], - "usage": { - "prompt_tokens": 10, - "completion_tokens": 5, - "total_tokens": 15, - }, - } - ), - ] - ) return ChatCompletion.model_validate( { "id": "chat-1", @@ -365,36 +281,6 @@ def __init__(self, **kwargs: Any) -> None: self.chat = SimpleNamespace(completions=FakeCompletions()) -class RetryStreamCompletions(FakeCompletions): - def __init__(self) -> None: - self.calls = 0 - - async def create(self, **kwargs: Any) -> Any: - self.calls += 1 - if self.calls == 1: - return FailedStream() - return await super().create(**kwargs) - - -class RetryStreamAsyncOpenAI: - max_retries = 1 - - def __init__(self, **kwargs: Any) -> None: - self.chat = SimpleNamespace(completions=RetryStreamCompletions()) - - -class MaxTokensStreamCompletions: - async def create(self, **kwargs: Any) -> Any: - return MaxTokensErrorStream() - - -class MaxTokensStreamAsyncOpenAI: - max_retries = 0 - - def __init__(self, **kwargs: Any) -> None: - self.chat = SimpleNamespace(completions=MaxTokensStreamCompletions()) - - @pytest.mark.asyncio async def test_rollout_supports_string_model_args( monkeypatch: pytest.MonkeyPatch, @@ -434,133 +320,17 @@ async def test_rollout_supports_string_model_args( policy_client: Any = rollout_module.openai_clients[ ("http://model.test/v1", "model-key") ] - completions = policy_client.chat.completions - assert completions.kwargs["stream"] is True - assert completions.kwargs["stream_options"] == {"include_usage": True} - assert completions.kwargs["extra_headers"] == { - "X-ART-Stream-Progress": "sse-comments-v1" - } - limits = policy_client.kwargs["http_client"].limits - assert limits.max_connections == 2048 - assert limits.max_keepalive_connections == 2048 - - -@pytest.mark.asyncio -async def test_rollout_retries_stream_body_transport_failure( - monkeypatch: pytest.MonkeyPatch, -) -> None: - rollout_module = importlib.import_module("art.tau_bench.rollout") - rollout_module.openai_clients.clear() - monkeypatch.setattr(rollout_module, "AsyncOpenAI", RetryStreamAsyncOpenAI) - - trajectory = await rollout_module.rollout( - Scenario(domain="banking_knowledge", task=Task(id="task_001")), - "http://model.test/v1", - "model-key", - "default", - client=FakeTauBenchClient(), - max_turns=1, + assert policy_client.chat.completions.kwargs["stream"] is False + assert policy_client.kwargs["max_retries"] == 1 + http_client = policy_client.kwargs["http_client"] + assert http_client.timeout == httpx.Timeout( + connect=30, + read=10 * 60, + write=30, + pool=30, ) - - completions: Any = rollout_module.openai_clients[ - ("http://model.test/v1", "model-key") - ].chat.completions - assert completions.calls == 2 - assert trajectory.reward == 1.0 - - -@pytest.mark.asyncio -async def test_rollout_does_not_capture_stream_attempt_missing_requested_usage() -> ( - None -): - rollout_module = importlib.import_module("art.tau_bench.rollout") - rollout_module.openai_clients.clear() - requests = 0 - - def stream_body(content: str, *, include_usage: bool) -> bytes: - chunks = [ - { - "id": "completion", - "object": "chat.completion.chunk", - "created": 0, - "model": "default", - "choices": [ - { - "index": 0, - "delta": {"role": "assistant", "content": content}, - "finish_reason": None, - } - ], - }, - { - "id": "completion", - "object": "chat.completion.chunk", - "created": 0, - "model": "default", - "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}], - }, - ] - if include_usage: - chunks.append( - { - "id": "completion", - "object": "chat.completion.chunk", - "created": 0, - "model": "default", - "choices": [], - "usage": { - "prompt_tokens": 2, - "completion_tokens": 1, - "total_tokens": 3, - }, - } - ) - return b"".join( - [ - *(f"data: {json.dumps(chunk)}\n\n".encode() for chunk in chunks), - b"data: [DONE]\n\n", - ] - ) - - async def handler(_request: httpx.Request) -> httpx.Response: - nonlocal requests - requests += 1 - return httpx.Response( - 200, - content=stream_body( - "bad" if requests == 1 else "good", - include_usage=requests > 1, - ), - headers={"Content-Type": "text/event-stream"}, - ) - - http_client = httpx.AsyncClient(transport=httpx.MockTransport(handler)) - openai_client = AsyncOpenAI( - api_key="model-key", - base_url="http://model.test/v1", - http_client=http_client, - max_retries=1, - ) - rollout_module.openai_clients[("http://model.test/v1", "model-key")] = openai_client - try: - trajectory = await rollout_module.rollout( - Scenario(domain="banking_knowledge", task=Task(id="task_001")), - "http://model.test/v1", - "model-key", - "default", - client=FakeTauBenchClient(), - max_turns=1, - ) - finally: - await openai_client.close() - await http_client.aclose() - rollout_module.openai_clients.clear() - - assert requests == 2 - assert len(trajectory.exchanges.chat_completions) == 1 - assert trajectory.exchanges.chat_completions[0].response.choices[ - 0 - ].message.content == ("good") + assert http_client.limits.max_connections == 2048 + assert http_client.limits.max_keepalive_connections == 2048 @pytest.mark.asyncio @@ -714,7 +484,6 @@ async def handler(request: httpx.Request) -> httpx.Response: "default", client=ToolTauBenchClient(), max_turns=2, - chat_completion_kwargs={"stream": False}, ) finally: await openai_client.close() @@ -802,35 +571,6 @@ async def test_rollout_stops_on_max_tokens_bad_request( "default", client=client, max_turns=10, - chat_completion_kwargs={"stream": False}, - ) - - assert trajectory.metrics["num_turns"] == 0 - assert client.steps == 0 - assert client.deleted == ["env-1"] - - -@pytest.mark.asyncio -async def test_rollout_stops_on_in_band_max_tokens_stream_error( - monkeypatch: pytest.MonkeyPatch, -) -> None: - rollout_module = importlib.import_module("art.tau_bench.rollout") - rollout_module.openai_clients.clear() - monkeypatch.setattr(rollout_module, "AsyncOpenAI", MaxTokensStreamAsyncOpenAI) - monkeypatch.setattr( - rollout_module, - "DefaultAsyncHttpxClient", - lambda **kwargs: SimpleNamespace(**kwargs), - ) - client = CountingTauBenchClient() - - trajectory = await rollout_module.rollout( - Scenario(domain="banking_knowledge", task=Task(id="task_001")), - "http://model.test/v1", - "model-key", - "default", - client=client, - max_turns=10, ) assert trajectory.metrics["num_turns"] == 0 @@ -884,7 +624,6 @@ async def test_rollout_stops_before_next_turn_exceeds_context( "default", client=client, max_turns=10, - chat_completion_kwargs={"stream": False}, ) assert trajectory.metrics["num_turns"] == 1 diff --git a/tests/unit/test_trainer_rank_telemetry.py b/tests/unit/test_trainer_rank_telemetry.py new file mode 100644 index 000000000..2ffdc8956 --- /dev/null +++ b/tests/unit/test_trainer_rank_telemetry.py @@ -0,0 +1,280 @@ +from __future__ import annotations + +import threading +from types import SimpleNamespace +from typing import cast + +import pytest + +from art.trainer_rank import _telemetry + + +def test_guard_failure_delta_accepts_runtime_string_values( + monkeypatch: pytest.MonkeyPatch, +) -> None: + code = object() + failures = {code: ["plain reason", SimpleNamespace(reason="named reason")]} + monkeypatch.setattr(_telemetry.torch._dynamo, "guard_failures", failures) + _telemetry._guard_counts.clear() + + assert _telemetry._new_guard_failures() == ("plain reason", "named reason") + assert _telemetry._new_guard_failures() == () + + +def test_phase_reports_compile_attribution_and_recompiles( + monkeypatch: pytest.MonkeyPatch, +) -> None: + events: list[tuple[dict[str, object], bool]] = [] + clock = iter((0.0, 1.0, 3.0, 5.0, 6.0, 8.0, 11.0, 12.0)) + monkeypatch.setattr(_telemetry, "_install", lambda: None) + monkeypatch.setattr(_telemetry.time, "perf_counter", lambda: next(clock)) + monkeypatch.setattr( + _telemetry, + "_emit", + lambda event, warning=False: events.append((dict(event), warning)), + ) + monkeypatch.setattr(_telemetry, "_new_guard_failures", lambda: ("size mismatch",)) + _telemetry._compiled_plan_signatures.clear() + first_args = SimpleNamespace( + callback_trigger=SimpleNamespace(name="DYNAMO"), compile_id="0/0" + ) + second_args = SimpleNamespace( + callback_trigger=SimpleNamespace(name="DYNAMO"), compile_id="0/1" + ) + + with _telemetry.phase("forward", {"packed_tokens": 128}, synchronized=True): + _telemetry._compile_start(first_args) + _telemetry._compile_end(first_args) + with _telemetry.phase("forward", {"packed_tokens": 128}, synchronized=True): + _telemetry._compile_start(second_args) + _telemetry._compile_end(second_args) + + first, second = events + assert first[0] == { + "event": "phase", + "phase": "forward", + "seconds": 5.0, + "synchronized": True, + "signature": {"packed_tokens": 128}, + "compile_status": "new_graph", + "compile_seconds": 2.0, + "compiles": [ + { + "compile_id": "0/0", + "trigger": "DYNAMO", + "seconds": 2.0, + "guard_failures": ("size mismatch",), + "frame_id": 0, + "frame_compile_id": 0, + "graph_status": "new_graph", + } + ], + "plan_signature_status": "new", + "unique_compile_plan_signatures": 1, + "outcome": "ok", + "error_type": None, + } + assert first[1] + assert second[0]["compile_status"] == "recompile" + second_compiles = cast(list[dict[str, object]], second[0]["compiles"]) + assert second_compiles[0]["compile_id"] == "0/1" + assert second[0]["plan_signature_status"] == "repeated" + assert second[0]["unique_compile_plan_signatures"] == 1 + assert second[1] + + +def test_phase_deduplicates_on_low_cardinality_plan_signature( + monkeypatch: pytest.MonkeyPatch, +) -> None: + events: list[dict[str, object]] = [] + clock = iter((0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0)) + monkeypatch.setattr(_telemetry, "_install", lambda: None) + monkeypatch.setattr(_telemetry.time, "perf_counter", lambda: next(clock)) + monkeypatch.setattr( + _telemetry, + "_emit", + lambda event, warning=False: events.append(dict(event)), + ) + monkeypatch.setattr(_telemetry, "_new_guard_failures", lambda: ()) + _telemetry._compiled_plan_signatures.clear() + plan_signature = {"topology": (1, 1, 1, 1), "request_mix": ("target",)} + + for compile_id, packed_tokens in (("0/0", 128), ("0/1", 256)): + args = SimpleNamespace( + callback_trigger=SimpleNamespace(name="DYNAMO"), compile_id=compile_id + ) + with _telemetry.phase( + "forward", + {**plan_signature, "packed_tokens": packed_tokens}, + dedup_signature=plan_signature, + ): + _telemetry._compile_start(args) + _telemetry._compile_end(args) + + assert events[0]["signature"] == { + **plan_signature, + "packed_tokens": 128, + } + assert events[0]["plan_signature_status"] == "new" + assert events[1]["signature"] == { + **plan_signature, + "packed_tokens": 256, + } + assert events[1]["plan_signature_status"] == "repeated" + assert events[1]["unique_compile_plan_signatures"] == 1 + + +def test_phase_without_compilation_reports_errors( + monkeypatch: pytest.MonkeyPatch, +) -> None: + events: list[dict[str, object]] = [] + clock = iter((2.0, 5.0)) + monkeypatch.setattr(_telemetry, "_install", lambda: None) + monkeypatch.setattr(_telemetry.time, "perf_counter", lambda: next(clock)) + monkeypatch.setattr( + _telemetry, + "_emit", + lambda event, warning=False: events.append(dict(event)), + ) + _telemetry._compiled_plan_signatures.clear() + + with pytest.raises(ValueError, match="failed"): + with _telemetry.phase("optim", {"checkpoint_count": 1}): + raise ValueError("failed") + + assert events == [ + { + "event": "phase", + "phase": "optim", + "seconds": 3.0, + "synchronized": False, + "signature": {"checkpoint_count": 1}, + "compile_status": "none", + "compile_seconds": 0.0, + "compiles": [], + "plan_signature_status": "none", + "unique_compile_plan_signatures": 0, + "outcome": "error", + "error_type": "ValueError", + } + ] + + +def test_compile_end_on_another_thread_keeps_phase_attribution( + monkeypatch: pytest.MonkeyPatch, +) -> None: + events: list[dict[str, object]] = [] + clock = iter((0.0, 1.0, 3.0, 4.0)) + monkeypatch.setattr(_telemetry, "_install", lambda: None) + monkeypatch.setattr(_telemetry.time, "perf_counter", lambda: next(clock)) + monkeypatch.setattr( + _telemetry, + "_emit", + lambda event, warning=False: events.append(dict(event)), + ) + monkeypatch.setattr(_telemetry, "_new_guard_failures", lambda: ()) + _telemetry._compiled_plan_signatures.clear() + args = SimpleNamespace( + callback_trigger=SimpleNamespace(name="DYNAMO"), compile_id="4/0" + ) + + with _telemetry.phase("forward", {"packed_tokens": 64}): + _telemetry._compile_start(args) + thread = threading.Thread(target=_telemetry._compile_end, args=(args,)) + thread.start() + thread.join() + + assert events[0]["compiles"] == [ + { + "compile_id": "4/0", + "trigger": "DYNAMO", + "seconds": 2.0, + "guard_failures": (), + "frame_id": 4, + "frame_compile_id": 0, + "graph_status": "new_graph", + } + ] + + +def test_compile_finishing_after_phase_close_emits_attributed_event( + monkeypatch: pytest.MonkeyPatch, +) -> None: + events: list[dict[str, object]] = [] + clock = iter((0.0, 1.0, 2.0, 4.0, 5.0, 6.0, 7.0, 9.0)) + monkeypatch.setattr(_telemetry, "_install", lambda: None) + monkeypatch.setattr(_telemetry.time, "perf_counter", lambda: next(clock)) + monkeypatch.setattr( + _telemetry, + "_emit", + lambda event, warning=False: events.append(dict(event)), + ) + monkeypatch.setattr(_telemetry, "_new_guard_failures", lambda: ()) + _telemetry._compiled_plan_signatures.clear() + args = SimpleNamespace( + callback_trigger=SimpleNamespace(name="DYNAMO"), compile_id="5/0" + ) + + with _telemetry.phase("forward", {"packed_tokens": 96}): + _telemetry._compile_start(args) + _telemetry._compile_end(args) + + assert events[0]["compiles"] == [] + assert events[1] == { + "event": "compile", + "phase": "forward", + "signature": {"packed_tokens": 96}, + "phase_closed": True, + "plan_signature_status": "new", + "unique_compile_plan_signatures": 1, + "compile_id": "5/0", + "trigger": "DYNAMO", + "seconds": 3.0, + "guard_failures": (), + "frame_id": 5, + "frame_compile_id": 0, + "graph_status": "new_graph", + } + + repeated_args = SimpleNamespace( + callback_trigger=SimpleNamespace(name="DYNAMO"), compile_id="5/1" + ) + with _telemetry.phase("forward", {"packed_tokens": 96}): + _telemetry._compile_start(repeated_args) + _telemetry._compile_end(repeated_args) + + assert events[3]["plan_signature_status"] == "repeated" + assert events[3]["unique_compile_plan_signatures"] == 1 + + +def test_install_restores_callbacks_after_handler_clear( + monkeypatch: pytest.MonkeyPatch, +) -> None: + class Handler: + def __init__(self) -> None: + self.start_callbacks: list[object] = [] + self.end_callbacks: list[object] = [] + + def register_start_callback(self, callback: object) -> None: + self.start_callbacks.append(callback) + + def register_end_callback(self, callback: object) -> None: + self.end_callbacks.append(callback) + + handler = Handler() + monkeypatch.setattr(_telemetry.torch._dynamo, "callback_handler", handler) + code = object() + failures = {code: ["old failure", "old failure 2"]} + monkeypatch.setattr(_telemetry.torch._dynamo, "guard_failures", failures) + + _telemetry._install() + assert handler.start_callbacks == [_telemetry._compile_start] + assert handler.end_callbacks == [_telemetry._compile_end] + handler.start_callbacks.clear() + handler.end_callbacks.clear() + failures[code] = ["fresh failure"] + _telemetry._install() + assert handler.start_callbacks == [_telemetry._compile_start] + assert handler.end_callbacks == [_telemetry._compile_end] + failures[code].append("next failure") + assert _telemetry._new_guard_failures() == ("next failure",) diff --git a/tests/unit/test_trainer_rank_validation.py b/tests/unit/test_trainer_rank_validation.py index c07d3fba1..bc7eb29a5 100644 --- a/tests/unit/test_trainer_rank_validation.py +++ b/tests/unit/test_trainer_rank_validation.py @@ -1153,6 +1153,63 @@ def test_checkpoint_export_requires_retained_adapter_config() -> None: trainer.export_lora("/unused", "student") +def test_prepared_lora_export_lifecycle_without_megatron( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from art.trainer_rank import _lora_export + + snapshot = object() + captures: list[object] = [] + saved: list[tuple[str, object]] = [] + + def capture(*_args: object, **_kwargs: object) -> tuple[object, dict[str, float]]: + captures.append(snapshot) + return snapshot, {"d2h": 1.0} + + monkeypatch.setattr(_lora_export, "_capture_lora_publish_inputs", capture) + monkeypatch.setattr( + _lora_export, + "_save_lora_publish_inputs", + lambda output, value: saved.append((output, value)) or {"serialize": 2.0}, + ) + trainer = TrainerRank(_runtime()) + trainer._checkpoint_slots["student"] = _CheckpointSlot( + config={ + "base_model_name_or_path": "test", + "r": 1, + "lora_alpha": 1, + "target_modules": [], + }, + revision=7, + ) + + revision, timings = trainer._prepare_lora_export( + "publish", "student", owner_id="owner" + ) + assert revision == 7 + assert timings["d2h"] == 1.0 + assert timings["slot_validation"] >= 0 + with pytest.raises(RuntimeError, match="already prepared"): + trainer._prepare_lora_export("publish", "student", owner_id="other") + trainer._abort_lora_export("publish", owner_id="other") + assert captures == [snapshot] + assert trainer._finish_lora_export("publish", "/output", owner_id="owner") == { + "serialize": 2.0 + } + assert saved == [("/output", snapshot)] + with pytest.raises(ValueError, match="Unknown prepared LoRA export"): + trainer._finish_lora_export("publish", "/output", owner_id="owner") + + trainer._prepare_lora_export("aborted", "student", owner_id="owner") + trainer._abort_lora_export("aborted", owner_id="owner") + with pytest.raises(ValueError, match="Unknown prepared LoRA export"): + trainer._finish_lora_export("aborted", "/output", owner_id="owner") + + assert trainer.export_lora("/legacy", "student") == 7 + assert saved[-1] == ("/legacy", snapshot) + assert not trainer._prepared_lora_exports + + def test_checkpoint_save_rejects_accumulated_gradients() -> None: trainer = TrainerRank(_runtime()) parameter = torch.nn.Parameter(torch.ones(2)) @@ -1776,6 +1833,7 @@ def _checkpoint_load_failure_worker( timeout=timedelta(seconds=15), ) from art.trainer_rank import _checkpoint as checkpoint_module + from art.trainer_rank import _lora_export as lora_export_module originals = ( checkpoint_module._load_adapter, @@ -1814,7 +1872,7 @@ def _checkpoint_load_failure_worker( } ) with pytest.raises((ValueError, RuntimeError), match="Unknown|Another"): - checkpoint_module.export_lora(trainer, "/unused", "student") + lora_export_module.export_lora(trainer, "/unused", "student") completed = torch.tensor(1) dist.all_reduce(completed) assert completed.item() == world_size diff --git a/tests/unit/trajectories/test_capture.py b/tests/unit/trajectories/test_capture.py index 660938568..6775c0d5c 100644 --- a/tests/unit/trajectories/test_capture.py +++ b/tests/unit/trajectories/test_capture.py @@ -1,7 +1,13 @@ from __future__ import annotations import asyncio -from collections.abc import AsyncGenerator, AsyncIterator, Coroutine, Generator +from collections.abc import ( + AsyncGenerator, + AsyncIterator, + Coroutine, + Generator, + Iterable, +) import copy from datetime import datetime, timedelta import gzip @@ -384,6 +390,76 @@ def generated() -> Generator[Coroutine[Any, Any, art.Trajectory], None, None]: assert len(generated_result.trajectories[0].exchanges.chat_completions) == 1 assert not outer.exchanges + calls = 0 + + async def once() -> art.Trajectory: + nonlocal calls + calls += 1 + return art.Trajectory(reward=1) + + shared = once() + resolved = art.Trajectory(reward=2) + mixed = await art.trajectory_group( + [resolved, shared, shared, ValueError("recorded")], + exceptions=[RuntimeError("provided")], + metadata={"source": "test"}, + metrics={"score": 3}, + logs=["log"], + ) + assert calls == 1 + assert len(mixed.trajectories) == 3 + assert mixed.trajectories[0] is resolved + assert mixed.trajectories[1] is mixed.trajectories[2] + assert [error.message for error in mixed.exceptions] == [ + "recorded", + "provided", + ] + assert mixed.metadata == {"source": "test"} + assert mixed.metrics == {"score": 3} + assert mixed.logs == ["log"] + + sibling_started = asyncio.Event() + sibling_stopped = asyncio.Event() + + async def sibling() -> art.Trajectory: + sibling_started.set() + try: + await asyncio.Event().wait() + finally: + sibling_stopped.set() + raise AssertionError("unreachable") + + async def cancelled() -> art.Trajectory: + await sibling_started.wait() + raise asyncio.CancelledError + + with pytest.raises(asyncio.CancelledError): + await art.trajectory_group([sibling(), cancelled()]) + assert sibling_stopped.is_set() + + sibling_started.clear() + sibling_stopped.clear() + scheduled = sibling() + + def failing_iterable() -> Iterable[Coroutine[Any, Any, art.Trajectory]]: + yield scheduled + raise ValueError("iteration failed") + + with pytest.raises(ValueError, match="iteration failed"): + await art.trajectory_group(failing_iterable()) + assert scheduled.cr_frame is None + + task = asyncio.create_task(once()) + scheduled = await art.trajectory_group([task, task]) + assert len(scheduled.trajectories) == 2 + assert scheduled.trajectories[0] is scheduled.trajectories[1] + assert calls == 2 + + future = asyncio.get_running_loop().create_future() + future.set_result(resolved) + from_future = await art.trajectory_group([future]) + assert from_future.trajectories == [resolved] + def test_sync_group_generator_initializes_once( monkeypatch: pytest.MonkeyPatch, diff --git a/tests/unit/trajectories/test_history.py b/tests/unit/trajectories/test_history.py index 6ff417702..6c795b384 100644 --- a/tests/unit/trajectories/test_history.py +++ b/tests/unit/trajectories/test_history.py @@ -299,6 +299,84 @@ def test_chat_projection_scales_with_captured_messages() -> None: assert normalized[2] < normalized[1] * 2, measurements +@pytest.mark.parametrize( + ("tokens", "sampled", "expected"), + [ + ([1, 2, 3], [], True), + ([1, 2], [1, 2, 3], False), + ([1, 2, 3, 4], [2, 3], True), + ([1, 1, 1, 2], [1, 1, 2], True), + ([1, 1, 1, 2], [1, 2, 2], False), + ], +) +def test_contains_tokens(tokens: list[int], sampled: list[int], expected: bool) -> None: + history_module = importlib.import_module("art.trajectories._history") + + assert history_module._contains_tokens(tokens, sampled) is expected + + +@pytest.mark.parametrize( + ("prompt", "output", "later_prompt", "expected"), + [ + ([0], [], [0, 1], False), + ([0], [1], [0], False), + ([0], [7, 1, 2], [0, 1, 2], True), + ([0], [1, 1, 1, 2], [0, 1, 1, 2], True), + ([0], [1, 2, 3], [0, 1, 2], False), + ([0], [9, 1], [0, 1, 2], True), + ([0], [1], [9, 1], False), + ], +) +def test_retains_output_suffix( + prompt: list[int], + output: list[int], + later_prompt: list[int], + expected: bool, +) -> None: + history_module = importlib.import_module("art.trajectories._history") + + assert ( + history_module._retains_output_suffix(prompt, output, later_prompt) is expected + ) + + +def test_contains_tokens_scales_linearly() -> None: + history_module = importlib.import_module("art.trajectories._history") + + class CountingTokens: + def __init__(self, values: list[int]) -> None: + self.values = values + self.accesses = 0 + + def __len__(self) -> int: + return len(self.values) + + def __getitem__(self, index: int | slice) -> int | list[int]: + value = self.values[index] + self.accesses += len(value) if isinstance(value, list) else 1 + return value + + def __iter__(self): + for value in self.values: + self.accesses += 1 + yield value + + tokens = CountingTokens([1] * 10_000 + [2]) + sampled = CountingTokens([1] * 1_000 + [2]) + + assert history_module._contains_tokens(cast(Any, tokens), cast(Any, sampled)) + assert tokens.accesses + sampled.accesses < 10 * (len(tokens) + len(sampled)) + + output = CountingTokens([1] * 10_000 + [2]) + later_prompt = CountingTokens([0, *([1] * 1_000), 2]) + assert history_module._retains_output_suffix( + [0], cast(Any, output), cast(Any, later_prompt) + ) + assert output.accesses + later_prompt.accesses < 10 * ( + len(output) + len(later_prompt) + ) + + def test_divergent_chat_projection_parses_each_exact_generation_once( monkeypatch: pytest.MonkeyPatch, ) -> None: