diff --git a/src/art/megatron/weights/lora_publish.py b/src/art/megatron/weights/lora_publish.py index e9d8b4b08..d136a4492 100644 --- a/src/art/megatron/weights/lora_publish.py +++ b/src/art/megatron/weights/lora_publish.py @@ -1,5 +1,7 @@ from collections.abc import Iterable, Sequence -from typing import Any, NamedTuple +from dataclasses import dataclass +import time +from typing import Any, NamedTuple, TypeVar import torch @@ -18,6 +20,18 @@ from art.megatron.model_support.spec import ExpertPackedLoraGroup, ExpertPackedLoraSlot from art.megatron.training.model_chunks import ModelChunks +_K = TypeVar("_K") + + +@dataclass(frozen=True) +class _LoraPublishSnapshot: + 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 PackedExpertShardMeta(NamedTuple): key: str @@ -382,6 +396,34 @@ def _rank_and_device() -> tuple[int, torch.device]: ) +def _raise_on_distributed_phase_error( + error: BaseException | None, + phase: str, + group: torch.distributed.ProcessGroup | None, +) -> None: + if not _distributed_ready(): + if error is not None: + raise error + return + backend = str(torch.distributed.get_backend(group)).lower() # type: ignore[possibly-missing-attribute] + flag = torch.tensor( + error is not None, + dtype=torch.uint8, + device=( + torch.device("cuda", torch.cuda.current_device()) + if backend.endswith("nccl") + else torch.device("cpu") + ), + ) + torch.distributed.all_reduce( # type: ignore[possibly-missing-attribute] + flag, op=torch.distributed.ReduceOp.MAX, group=group + ) + if flag.item(): + if error is not None: + raise error + raise RuntimeError(f"Another rank failed during {phase}") + + def _metadata_by_owner_dtype( metadata: Sequence[Any], ) -> dict[tuple[int, str], list[Any]]: @@ -582,17 +624,17 @@ def merge_packed_expert_adapter_entries( } -def _stage_published_tensors( - tensors: dict[str, torch.Tensor], +def _stage_tensor_mapping( + tensors: dict[_K, torch.Tensor], stager: _PinnedCpuStager, -) -> dict[str, torch.Tensor]: - grouped: dict[tuple[str, int | None, str], list[tuple[str, torch.Tensor]]] = {} +) -> dict[_K, torch.Tensor]: + grouped: dict[tuple[str, int | None, str], list[tuple[_K, torch.Tensor]]] = {} for key, tensor in tensors.items(): dtype_name = _dtype_name(tensor.dtype) group_key = (tensor.device.type, tensor.device.index, dtype_name) grouped.setdefault(group_key, []).append((key, tensor)) - staged: dict[str, torch.Tensor] = {} + staged: dict[_K, torch.Tensor] = {} for _group_key, group in sorted(grouped.items()): flat = torch.cat( [tensor.detach().contiguous().view(-1) for _key, tensor in sorted(group)] @@ -610,6 +652,13 @@ def _stage_published_tensors( return staged +def _stage_published_tensors( + tensors: dict[str, torch.Tensor], + stager: _PinnedCpuStager, +) -> dict[str, torch.Tensor]: + return _stage_tensor_mapping(tensors, stager) + + def _save_rank0_vllm_lora( *, metadata: list[LoraShardMeta], @@ -667,7 +716,7 @@ def _rank0_vllm_lora_tensors( ) -def build_vllm_lora_tensors_from_model( +def _collect_lora_publish_inputs( *, model: ModelChunks, adapter_dtypes: dict[str, torch.dtype], @@ -676,51 +725,85 @@ def build_vllm_lora_tensors_from_model( rank: int, world_size: int, slot_ref: LoRASlotRef | None = None, -) -> tuple[dict[str, torch.Tensor], dict[str, Any]] | None: - actual_rank, device = _rank_and_device() - if _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}) 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 - packed_expert_groups = tuple(handler.expert_packed_lora_groups()) - planner = LoRAPublishPlanner(model, slot_ref) - local_tensors, local_metadata = 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 = collect_local_packed_expert_entries( - model, - adapter_dtypes, - owner_rank=rank, - packed_expert_groups=packed_expert_groups, - slot_ref=slot_ref, - ) - all_packed_metadata = ( - _global_packed_expert_metadata(planner, adapter_dtypes, packed_expert_groups) - if rank == 0 - else local_packed_metadata + timings: dict[str, float] | None = None, + failure_group: torch.distributed.ProcessGroup | None = None, +) -> _LoraPublishSnapshot | None: + started = time.monotonic() + device: torch.device | None = None + error: BaseException | None = None + try: + actual_rank, device = _rank_and_device() + if _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 + except BaseException as exc: + error = exc + _raise_on_distributed_phase_error( + error, "LoRA publish runtime validation", failure_group ) - if rank == 0: - all_metadata = _global_regular_metadata( - planner, + assert device is not None + if timings is not None: + timings["runtime_validation"] = time.monotonic() - started + started = time.monotonic() + local_tensors: dict[str, torch.Tensor] = {} + local_packed_tensors: dict[str, torch.Tensor] = {} + all_metadata: list[LoraShardMeta] = [] + all_packed_metadata: list[PackedExpertShardMeta] = [] + error = None + try: + packed_expert_groups = tuple(handler.expert_packed_lora_groups()) + planner = LoRAPublishPlanner(model, slot_ref) + local_tensors, local_metadata = 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, + ) = collect_local_packed_expert_entries( + model, adapter_dtypes, - packed_expert_groups if all_packed_metadata else (), + owner_rank=rank, + packed_expert_groups=packed_expert_groups, + slot_ref=slot_ref, + ) + all_packed_metadata = ( + _global_packed_expert_metadata( + planner, adapter_dtypes, packed_expert_groups + ) + if rank == 0 + else local_packed_metadata ) - else: - all_metadata = local_metadata + all_metadata = ( + _global_regular_metadata( + planner, + adapter_dtypes, + packed_expert_groups if all_packed_metadata else (), + ) + if rank == 0 + else local_metadata + ) + except BaseException as exc: + error = exc + _raise_on_distributed_phase_error(error, "LoRA publish planning", failure_group) + if timings is not None: + timings["plan_collect"] = time.monotonic() - started + started = time.monotonic() exchanged_tensors = _exchange_batched_tensors( all_metadata, local_tensors=local_tensors, @@ -733,18 +816,114 @@ def build_vllm_lora_tensors_from_model( rank=rank, device=device, ) + if timings is not None: + timings["exchange"] = time.monotonic() - started if rank != 0: return None - return _rank0_vllm_lora_tensors( + return _LoraPublishSnapshot( metadata=all_metadata, tensors_by_owner_key=exchanged_tensors, packed_expert_metadata=all_packed_metadata, packed_expert_tensors_by_owner_key=exchanged_packed_tensors, handler=handler, + adapter_config=dict(adapter_config), + ) + + +def build_vllm_lora_tensors_from_model( + *, + 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, +) -> tuple[dict[str, torch.Tensor], dict[str, Any]] | None: + snapshot = _collect_lora_publish_inputs( + model=model, + adapter_dtypes=adapter_dtypes, + handler=handler, adapter_config=adapter_config, + rank=rank, + world_size=world_size, + slot_ref=slot_ref, + ) + if snapshot is None: + return None + + return _rank0_vllm_lora_tensors( + metadata=snapshot.metadata, + tensors_by_owner_key=snapshot.tensors_by_owner_key, + packed_expert_metadata=snapshot.packed_expert_metadata, + packed_expert_tensors_by_owner_key=snapshot.packed_expert_tensors_by_owner_key, + handler=snapshot.handler, + adapter_config=snapshot.adapter_config, + ) + + +def _capture_lora_publish_snapshot_from_model( + *, + 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, + failure_group: torch.distributed.ProcessGroup | None = None, +) -> tuple[_LoraPublishSnapshot | None, dict[str, float]]: + timings: dict[str, float] = {} + snapshot = _collect_lora_publish_inputs( + model=model, + adapter_dtypes=adapter_dtypes, + handler=handler, + adapter_config=adapter_config, + rank=rank, + world_size=world_size, + slot_ref=slot_ref, + timings=timings, + failure_group=failure_group, + ) + started = time.monotonic() + if snapshot is not None: + stager = _PinnedCpuStager() + snapshot = _LoraPublishSnapshot( + metadata=snapshot.metadata, + tensors_by_owner_key=_stage_tensor_mapping( + snapshot.tensors_by_owner_key, stager + ), + packed_expert_metadata=snapshot.packed_expert_metadata, + packed_expert_tensors_by_owner_key=_stage_tensor_mapping( + snapshot.packed_expert_tensors_by_owner_key, stager + ), + handler=snapshot.handler, + adapter_config=snapshot.adapter_config, + ) + stager.finish() + timings["d2h"] = time.monotonic() - started + return snapshot, timings + + +def _save_lora_publish_snapshot( + output_dir: str, snapshot: _LoraPublishSnapshot +) -> dict[str, float]: + started = time.monotonic() + vllm_tensors, published_config = _rank0_vllm_lora_tensors( + metadata=snapshot.metadata, + tensors_by_owner_key=snapshot.tensors_by_owner_key, + packed_expert_metadata=snapshot.packed_expert_metadata, + packed_expert_tensors_by_owner_key=snapshot.packed_expert_tensors_by_owner_key, + handler=snapshot.handler, + adapter_config=snapshot.adapter_config, ) + 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 save_vllm_lora_from_model( 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..f6e330ac7 100644 --- a/src/art/trainer_rank/_checkpoint.py +++ b/src/art/trainer_rank/_checkpoint.py @@ -13,6 +13,7 @@ import shutil import struct import threading +import time from typing import TYPE_CHECKING, Literal, NotRequired, TypedDict, cast import uuid @@ -21,6 +22,7 @@ if TYPE_CHECKING: from art.megatron.lora import LoRA, LoraShardMeta, LoRASlotRef + from art.megatron.weights.lora_publish import _LoraPublishSnapshot from art.trainer_rank._impl import ( TrainerRank, _AdapterConfig, @@ -1724,8 +1726,24 @@ def commit() -> None: raise -def export_lora(trainer: TrainerRank, output_dir: str, checkpoint_name: str) -> int: +def prepare_lora_export( + trainer: TrainerRank, + export_id: str, + checkpoint_name: str, + *, + owner_id: str, +) -> tuple[int, dict[str, float]]: + started = time.monotonic() group = _ensure_group(trainer) + snapshots: dict[str, tuple[str, _LoraPublishSnapshot]] = getattr( + trainer, "_prepared_lora_exports", {} + ) + duplicate = ( + RuntimeError(f"LoRA export {export_id!r} is already prepared") + if export_id in snapshots + else None + ) + raise_distributed(duplicate, "validate LoRA export ID", group) slot = None error: BaseException | None = None try: @@ -1745,24 +1763,77 @@ def export_lora(trainer: TrainerRank, output_dir: str, checkpoint_name: str) -> 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 + slot_validation = time.monotonic() - started + from art.megatron.weights.lora_publish import ( + _capture_lora_publish_snapshot_from_model, + ) + snapshot = None + capture_timings: dict[str, float] = {} error = None try: - save_vllm_lora_from_model( + snapshot, capture_timings = _capture_lora_publish_snapshot_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), + failure_group=group, ) except BaseException as exc: error = exc + raise_distributed(error, "prepare LoRA export", group) + if snapshot is not None: + snapshots[export_id] = (owner_id, snapshot) + 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]: + from art.megatron.weights.lora_publish import _save_lora_publish_snapshot + + snapshots: dict[str, tuple[str, _LoraPublishSnapshot]] = getattr( + trainer, "_prepared_lora_exports", {} + ) + try: + owner, snapshot = 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_snapshot(output_dir, snapshot) + + +def abort_lora_export(trainer: TrainerRank, export_id: str, *, owner_id: str) -> None: + snapshots: dict[str, tuple[str, _LoraPublishSnapshot]] = 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: + group = _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) raise_distributed(error, "export LoRA", group) - return slot.revision + return revision def _ensure_groups( diff --git a/src/art/trainer_rank/_impl.py b/src/art/trainer_rank/_impl.py index 84c1a7ea2..92a0fde13 100644 --- a/src/art/trainer_rank/_impl.py +++ b/src/art/trainer_rank/_impl.py @@ -58,6 +58,7 @@ from art.megatron.lora import LoRASlotRef from art.megatron.prefix_tree_state import PrefixTreeAttentionState from art.megatron.train import TrainingRuntime + from art.megatron.weights.lora_publish import _LoraPublishSnapshot from art.trainer_rank._checkpoint import ( CustomOptimizerState, LocalOptimizerState, @@ -813,6 +814,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, _LoraPublishSnapshot]] = {} 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 @@ -1280,6 +1282,36 @@ def 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 _checkpoint + + return _checkpoint.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 _checkpoint + + return _checkpoint.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 _checkpoint + + _checkpoint.abort_lora_export(self, export_id, owner_id=owner_id) + @staticmethod def _checkpoint_source_key(path: str) -> str: return str(Path(path).resolve()) 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/tests/integration/megatron/lora/test_lora_disk_codecs.py b/tests/integration/megatron/lora/test_lora_disk_codecs.py index 92606d6a8..c6aa523f9 100644 --- a/tests/integration/megatron/lora/test_lora_disk_codecs.py +++ b/tests/integration/megatron/lora/test_lora_disk_codecs.py @@ -1623,6 +1623,116 @@ def test_save_vllm_lora_from_model_writes_single_vllm_checkpoint(tmp_path: Path) _assert_tensors_equal(roundtrip, full) +@pytest.mark.parametrize( + ("failure_collective", "phase", "expected_plans"), + ((1, "runtime validation", 0), (2, "planning", 1)), +) +def test_remote_preexchange_failure_prevents_tensor_exchange( + monkeypatch: pytest.MonkeyPatch, + failure_collective: int, + phase: str, + expected_plans: int, +) -> None: + exchanges: list[object] = [] + plans: list[object] = [] + collectives = 0 + failure_group = cast(Any, object()) + + def all_reduce(flag: torch.Tensor, **_kwargs: object) -> None: + nonlocal collectives + collectives += 1 + if collectives == failure_collective: + flag.fill_(1) + + monkeypatch.setattr(lora_publish, "_distributed_ready", lambda: True) + monkeypatch.setattr( + lora_publish, "_rank_and_device", lambda: (0, torch.device("cpu")) + ) + monkeypatch.setattr(torch.distributed, "get_world_size", lambda: 2) + monkeypatch.setattr(torch.distributed, "get_backend", lambda group: "gloo") + monkeypatch.setattr(torch.distributed, "all_reduce", all_reduce) + monkeypatch.setattr( + lora_publish, + "LoRAPublishPlanner", + lambda *_args: plans.append(object()) or plans[-1], + ) + monkeypatch.setattr( + lora_publish, "collect_local_lora_entries", lambda *_args, **_kwargs: ({}, []) + ) + monkeypatch.setattr( + lora_publish, + "collect_local_packed_expert_entries", + lambda *_args, **_kwargs: ({}, []), + ) + monkeypatch.setattr(lora_publish, "_global_regular_metadata", lambda *_args: []) + monkeypatch.setattr( + lora_publish, "_global_packed_expert_metadata", lambda *_args: [] + ) + monkeypatch.setattr( + lora_publish, + "_exchange_batched_tensors", + lambda *_args, **_kwargs: exchanges.append(object()), + ) + handler = SimpleNamespace(expert_packed_lora_groups=lambda: ()) + + with pytest.raises( + RuntimeError, match=f"Another rank failed during LoRA publish {phase}" + ): + lora_publish._collect_lora_publish_inputs( + model=cast(Any, object()), + adapter_dtypes={}, + handler=handler, + adapter_config={}, + rank=0, + world_size=2, + failure_group=failure_group, + ) + + assert exchanges == [] + assert len(plans) == expected_plans + + +def test_device_resolution_failure_uses_cpu_failure_group( + monkeypatch: pytest.MonkeyPatch, +) -> None: + failure_group = cast(Any, object()) + reduced_groups: list[object] = [] + monkeypatch.setattr(lora_publish, "_distributed_ready", lambda: True) + monkeypatch.setattr( + lora_publish, + "_rank_and_device", + lambda: (_ for _ in ()).throw(RuntimeError("device resolution failed")), + ) + monkeypatch.setattr( + torch.cuda, + "current_device", + lambda: (_ for _ in ()).throw(AssertionError("unexpected CUDA lookup")), + ) + monkeypatch.setattr( + torch.distributed, + "get_backend", + lambda group: "gloo" if group is failure_group else "nccl", + ) + monkeypatch.setattr( + torch.distributed, + "all_reduce", + lambda _flag, **kwargs: reduced_groups.append(kwargs["group"]), + ) + + with pytest.raises(RuntimeError, match="device resolution failed"): + lora_publish._collect_lora_publish_inputs( + model=cast(Any, object()), + adapter_dtypes={}, + handler=object(), + adapter_config={}, + rank=0, + world_size=2, + failure_group=failure_group, + ) + + assert reduced_groups == [failure_group] + + def test_trainer_rank_publishes_named_checkpoint_slot_without_mutating_base( tmp_path: Path, ): @@ -1662,6 +1772,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_validation.py b/tests/unit/test_trainer_rank_validation.py index c07d3fba1..a3546501c 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: + snapshot = object() + captures: list[object] = [] + saved: list[tuple[str, object]] = [] + module = ModuleType("art.megatron.weights.lora_publish") + + def capture(**_kwargs: object) -> tuple[object, dict[str, float]]: + captures.append(snapshot) + return snapshot, {"d2h": 1.0} + + setattr(module, "_capture_lora_publish_snapshot_from_model", capture) + setattr( + module, + "_save_lora_publish_snapshot", + lambda output, value: saved.append((output, value)) or {"serialize": 2.0}, + ) + monkeypatch.setitem(sys.modules, module.__name__, module) + 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)) 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: