Skip to content

Harden pipeline publication, trajectories, and trainer telemetry #809

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
bradhilton merged 8 commits into from
Aug 19, 2026
Merged

Harden pipeline publication, trajectories, and trainer telemetry #809

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
757872a
Pipeline LoRA publication off the trainer queue
bradhilton Aug 18, 2026
d845d28
Make token lineage checks linear
bradhilton Aug 18, 2026
291b387
Keep sampled suffix reconciliation linear
bradhilton Aug 18, 2026
ed3efbb
Restore bounded nonstreaming tau-bench completions
bradhilton Aug 18, 2026
1032485
Attribute trainer compile stalls
bradhilton Aug 19, 2026
9e5fa4b
Fix compile telemetry attribution
bradhilton Aug 19, 2026
57406c7
Isolate TrainerRank publication and harden trajectory grouping
bradhilton Aug 19, 2026
9883da6
Apply CI formatter
bradhilton Aug 19, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
111 changes: 11 additions & 100 deletions src/art/tau_bench/rollout.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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: ")},
Expand All @@ -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"] = (
Expand All @@ -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
Expand Down Expand Up @@ -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,
*,
Expand All @@ -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
Expand Down Expand Up @@ -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

Expand Down
41 changes: 0 additions & 41 deletions src/art/trainer_rank/_checkpoint.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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]:
Expand Down
Loading
Loading