Add protocol-aware trajectory capture and tokenization

[bradhilton]

Summary

• replace the HTTPX-only auto-trajectory path with typed exchange capture for Chat Completions, Completions, Responses, and Anthropic Messages
• represent requests as protocol-specific TypedDicts with known native SDK fields plus arbitrary extension fields, and responses as native OpenAI/Anthropic models
• lazily capture successful, complete calls made through HTTPX, aiohttp, and requests; reconstruct streaming calls into final responses without exposing transport events
• add context-local Trajectory / TrajectoryGroup scopes and coroutine helpers while preserving legacy trajectory/group construction behind deprecations
• add strict exchange tokenization that prefers exact service token IDs/logprobs, falls back to protocol-aware rendering, and rejects histories that are branching, unresolved, or multi-model without an explicit selection
• expose singular protocol-native history projections with explicit Chat conversion, exact raw-Completions spans, and RULER support
• feed exchange trajectories into existing local/serverless preprocessing and Tinker-native masked datums while retaining legacy history behavior

Design

• art.trajectories.__init__ is the public interface; transport, protocol, scope, compatibility, serialization, and tokenization details live in private modules
• capture hooks install lazily and idempotently, preserve SDK-visible behavior, deduplicate overlapping transport paths, and are effectively no-ops outside an active trajectory scope
• HTTP failures, malformed responses, interrupted streams, and unconsumed streaming calls are not recorded
• exchange requests remain lossless wire dictionaries: typed common fields improve completion/checking while extra_items=Any preserves provider and Caladan extensions
• exchange.model is a serialized computed property derived from the request first and response second, so stored metadata cannot diverge from the captured wire objects
• tokenization sorts exchanges by timestamps, enforces one append-only RL history and one selected model, prefers exact prompt/output IDs and logprobs, and uses NaN only when fallback tokenization cannot recover a logprob
• training rejects missing assistant-token logprobs unless allow_training_without_logprobs=True; prompt-token NaNs remain expected
• explicit model, base_model, chat template, and template kwargs override inferred values; W&B checkpoint metadata and tokenizer defaults provide best-effort inference

Adversarial hardening

Independent Codex and Claude Fable 5 reviews found no P0, then drove focused fixes:

• fail closed when explicit exact token metadata is malformed instead of filtering it into a plausible shorter sequence
• reconstruct interleaved streamed Chat choices by choice.index
• aggregate Responses token/logprob pairs across all compatible output blocks, falling back as a whole when coverage is partial
• reject raw Completions batch prompts and echo=True until their multi-history/masking semantics are implemented correctly
• reject unsupported Responses items only when they would otherwise be silently omitted from a rendered prompt; authoritative raw tokens still support single reasoning responses
• snapshot nested aiohttp request data, distinguish protocol-complete SSE from abandoned JSON transport streams, and recognize Azure/Gemini//v1-less compatible endpoint paths
• treat independently undecodable token bytes as unalignable logprobs rather than crashing
• derive serialized exchange model identity from request/response data
• initialize synchronous TrajectoryGroup generators exactly once while retaining the deprecated awaitable overload
• preserve exact prompt/output token metadata in Tinker-native responses and train multi-turn exchange histories with aligned masks/logprobs

Compatibility and scope

• legacy histories, synchronous TrajectoryGroup construction, and the awaitable group overload remain supported
• auto_trajectory and capture_auto_trajectory remain as deprecated aliases
• trajectories mixing legacy histories and exchanges are rejected rather than ambiguously merged
• multi-history output, normalized string tables, raw stream-event storage, proxy capture, and remote exchange synchronization remain out of scope
• Anthropic is constrained to >=0.77.0, the first verified version exposing the streaming accumulator imports used by capture

Validation

• protocol-history, legacy-consumer, and RULER slice: 53 passed
• post-review trajectory/legacy-capture suite: 39 passed, 1 optional Tinker test module skipped without its extra
• isolated Tinker unit/rendering suite: 11 passed
• live Tinker-native exchange campaign: six captured rollouts across three groups, one real training step, and inference from both the new sampler and step 0; passed in 127.23 seconds
• repository pre-commit gates: Ruff, format, scoped ty, and lock sync passed on ty 0.0.59
• isolated Anthropic 0.77.0 import smoke passed
• serverless-training complete/legacy trajectory round trips passed against the current ART source
• earlier ART unit slice: 295 passed, 3 skipped
• disposable Caladan Qwen stack: all four protocols, sync/async, stream/non-stream, base/identity-LoRA, custom templates, multi-turn, reasoning, and tool calls
• OpenAI and Anthropic direct APIs; OpenRouter all four protocols; Together, Mistral, Fireworks, LiteLLM, and raw HTTPX/aiohttp/requests; Groq through its HTTPX transport with a mock response because no Groq credential was available
• W&B evidence: caladan-trajectory-validation/trajectory-capture-0714

Open integration decisions

• Megatron MoE expert replay needs optional routed-expert propagation through exchange tokenization; current Caladan does not enable route return, and vLLM 0.24 only exposes it on non-streaming Chat/Completions
• Parquet/local logging remains legacy-history-oriented; gather completion-token metrics now total native usage across exchange protocols and sum legacy multi-choice tokens
• serverless-training PR coreweave/serverless-training#283 targets staging and preserves exchange payloads, but serverless must pin the merged ART revision before exchange preprocessing is operational

Review focus

• streaming reconstruction and capture finalization under arbitrary chunking, cancellation, and early close
• transport monkeypatch idempotency, context isolation, and cross-transport deduplication
• protocol-specific history stitching and exact-token/logprob precedence
• compatibility serialization and legacy constructor behavior
• exchange-to-training mask/logprob alignment in local/serverless and Tinker-native paths

A separate Caladan change will correct the Qwen vLLM tool-parser mismatch found by strict multi-turn validation. Serverless-training compatibility is tracked in coreweave/serverless-training#283.

[bradhilton]

Cross-repository serverless integration is now complete in coreweave/serverless-training#287, which pins this PR exactly at 41e292fb005edbd58876f913718c1b3c7eb053d8.

Dependency ownership and validation:

• ART's backend and megatron extras remain intentionally mutually exclusive. Serverless builds isolated environments and selects one with its existing model router; there is no Transformers override.
• The backend environment resolves Transformers 5.2.0; the Megatron environment resolves Transformers 5.12.1.
• This PR's direct flashinfer-python==0.6.8.post1 and flashinfer-cubin==0.6.8.post1 requirements are inherited by the downstream lock; serverless carries no FlashInfer override.
• Both downstream environments resolve Torch 2.11.0+cu128 and Torchvision 0.26.0+cu128 on the CUDA 12.8 image.
• ART trajectory tokenization slice: 33 passed.
• Serverless compact exchange preprocessing/packing plus runtime-router slice: 5 passed.
• Two-H200 backend validation: the real serverless UnslothTrainer completed a Qwen2.5-14B optimizer step with finite loss, entropy, and grad norm.
• Two-H200 Megatron validation: Transformer Engine/FlashAttention 4/Quack/FlashInfer imported, and the real serverless MegatronTrainer plus two-rank megatron/train.py completed eight Qwen3-0.6B gradient steps with finite loss and grad norm.

The disposable SkyPilot cluster was torn down. Serverless PR #287 targets staging only; production remains untouched.

[bradhilton]

Final downstream validation: coreweave/serverless-training staging now pins exact ART commit 41e292f via merged PR #287. Staging deploy run 29456887370 succeeded; live backend and Megatron pods both report this direct-url commit. Backend is running Transformers 5.2.0, while Megatron is running 5.12.1 with matched FlashInfer 0.6.8.post1; both use Torch 2.11.0+cu128 / CUDA 12.8. All staging deployments are ready and production was untouched.