From ff96504e15eef71bf618c8dbf92dad0879e2ce9d Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Tue, 15 Sep 2026 23:27:13 +0000 Subject: [PATCH 1/3] Preserve NaN log probabilities in trajectory JSON --- src/art/trajectories/_serialization.py | 2 ++ tests/unit/trajectories/test_tokenize.py | 24 ++++++++++++++++++++++++ 2 files changed, 26 insertions(+) diff --git a/src/art/trajectories/_serialization.py b/src/art/trajectories/_serialization.py index b413f7ffa..1a56ee1db 100644 --- a/src/art/trajectories/_serialization.py +++ b/src/art/trajectories/_serialization.py @@ -31,6 +31,8 @@ def _without_pickle_string_interning(): class _StringInterningModel(BaseModel): """Intern strings once, immediately before this graph is pickled.""" + model_config = pydantic.ConfigDict(ser_json_inf_nan="strings") + # Process-local optimization state: omitting it from Pydantic private state keeps # equality and serialization unchanged, and lets a receiving process prepare the # graph again after local mutation. diff --git a/tests/unit/trajectories/test_tokenize.py b/tests/unit/trajectories/test_tokenize.py index 7dae26ecc..5858aa2bb 100644 --- a/tests/unit/trajectories/test_tokenize.py +++ b/tests/unit/trajectories/test_tokenize.py @@ -4380,6 +4380,30 @@ def test_json_round_trip_preserves_exchange_types() -> None: assert isinstance(restored.exchanges.chat_completions[0].response, ChatCompletion) +@pytest.mark.parametrize("grouped", [False, True]) +def test_json_round_trip_preserves_nan_response_logprobs(grouped: bool) -> None: + exchange = _chat_exchange([1], [2, 3]) + logprobs = exchange.response.choices[0].logprobs + assert logprobs is not None and logprobs.content is not None + logprobs.content[1].logprob = math.nan + trajectory = art.Trajectory( + exchanges=TrajectoryExchanges(chat_completions=[exchange]) + ) + expected = trajectory.tokenize() + value = art.TrajectoryGroup([trajectory]) if grouped else trajectory + for _ in range(2): + value = type(value).model_validate_json(value.model_dump_json()) + restored = ( + value.trajectories[0] if isinstance(value, art.TrajectoryGroup) else value + ) + actual = restored.tokenize() + assert actual.tokens == expected.tokens + assert actual.flags == expected.flags + assert actual.logprobs[1] == expected.logprobs[1] + assert math.isnan(actual.logprobs[2]) + assert actual.flags[2] & tr.TokenFlag.SAMPLED + + def _response_exchange( response_id: str, output_id: int, From 75d37c15a9f9f273fd63ad7558fb62e34bc5aeb9 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Wed, 16 Sep 2026 02:41:19 +0000 Subject: [PATCH 2/3] Preserve NaN source bindings across trajectory serialization --- src/art/trajectories/__init__.py | 6 ++-- src/art/trajectories/_compact.py | 4 +-- src/art/trajectories/_serialization.py | 29 +++++++++++++-- src/art/trajectories/tensors.py | 6 ++-- .../trajectories/test_tokenized_models.py | 36 +++++++++++++++++-- 5 files changed, 69 insertions(+), 12 deletions(-) diff --git a/src/art/trajectories/__init__.py b/src/art/trajectories/__init__.py index f72a42205..730c3d096 100644 --- a/src/art/trajectories/__init__.py +++ b/src/art/trajectories/__init__.py @@ -76,6 +76,7 @@ from ..types import Messages, MessagesAndChoices, Tools from ._serialization import ( _CompactModel, + _equal_with_nan, _rebind_history_sources, _StringInterningModel, _StringPool, @@ -1330,9 +1331,8 @@ def _bind_source_graph(self) -> TokenizedTrajectoryGroup[TokenizedTrajectoryT]: for tokenized, trajectory in zip( self.trajectories, self.trajectory_group.trajectories, strict=True ): - if ( - tokenized.trajectory is not trajectory - and tokenized.trajectory.model_dump() != trajectory.model_dump() + if tokenized.trajectory is not trajectory and not _equal_with_nan( + tokenized.trajectory.model_dump(), trajectory.model_dump() ): raise ValueError("Tokenized trajectory does not match its source group") tokenized.trajectory = trajectory diff --git a/src/art/trajectories/_compact.py b/src/art/trajectories/_compact.py index 3582e95b8..f6524677b 100644 --- a/src/art/trajectories/_compact.py +++ b/src/art/trajectories/_compact.py @@ -20,7 +20,7 @@ TrajectoryGroup, _load_tensors, ) -from ._serialization import _rebind_history_sources +from ._serialization import _equal_with_nan, _rebind_history_sources _FORMAT = "art.trajectories" _VERSION = 1 @@ -405,7 +405,7 @@ def _compact_group_data(data: dict[str, pydantic.JsonValue]) -> None: dict[str, pydantic.JsonValue], dict(_mapping(item, "Tokenized trajectory")), ) - if child.pop("trajectory", None) != source: + if not _equal_with_nan(child.pop("trajectory", None), source): raise ValueError("Tokenized trajectory does not match its source group") registry = _ExchangeDataRegistry.from_trajectory_data(source) if "history" in child: diff --git a/src/art/trajectories/_serialization.py b/src/art/trajectories/_serialization.py index 1a56ee1db..7f2a001c5 100644 --- a/src/art/trajectories/_serialization.py +++ b/src/art/trajectories/_serialization.py @@ -3,6 +3,7 @@ from collections.abc import Callable from contextlib import contextmanager from dataclasses import fields, is_dataclass +import math import threading from typing import Any, Literal, SupportsIndex, cast @@ -234,6 +235,26 @@ def validate_history(value: object) -> object: return model.model_validate(data) +def _equal_with_nan(left: Any, right: Any) -> bool: + """Compare serialized source data, treating unknown log probabilities as equal.""" + if left is right: + return True + if isinstance(left, dict) and isinstance(right, dict): + return left.keys() == right.keys() and all( + _equal_with_nan(value, right[key]) for key, value in left.items() + ) + if isinstance(left, (list, tuple)) and type(left) is type(right): + return len(left) == len(right) and all( + _equal_with_nan(a, b) for a, b in zip(left, right, strict=True) + ) + return left == right or ( + isinstance(left, float) + and isinstance(right, float) + and math.isnan(left) + and math.isnan(right) + ) + + def _rebind_history_sources( history: object, trajectory: object | None = None, @@ -257,7 +278,7 @@ def _rebind_history_sources( MessagesExchange, ) - def exchanges(value: object) -> list[object]: + def exchanges(value: object) -> list[BaseModel]: if not isinstance(value, Trajectory): return [] return [ @@ -301,7 +322,11 @@ def visit(value: object) -> None: matches = [ exchange for exchange in canonical - if type(exchange) is type(item) and exchange == item + if type(exchange) is type(item) + and ( + exchange == item + or _equal_with_nan(exchange.model_dump(), item.model_dump()) + ) ] if matches: object.__setattr__(value, name, matches[0]) diff --git a/src/art/trajectories/tensors.py b/src/art/trajectories/tensors.py index 82d16a483..cde8121af 100644 --- a/src/art/trajectories/tensors.py +++ b/src/art/trajectories/tensors.py @@ -28,6 +28,7 @@ _StringInterningModel, ) from ._serialization import ( + _equal_with_nan, _rebind_history_sources, serialize_history, validate_history, @@ -286,9 +287,8 @@ def bind_source_group(self) -> Self: for tensorized, trajectory in zip( self.trajectories, self.trajectory_group.trajectories, strict=True ): - if ( - tensorized.trajectory is not trajectory - and tensorized.trajectory.model_dump() != trajectory.model_dump() + if tensorized.trajectory is not trajectory and not _equal_with_nan( + tensorized.trajectory.model_dump(), trajectory.model_dump() ): raise ValueError( "Tensorized trajectory does not match its source group" diff --git a/tests/unit/trajectories/test_tokenized_models.py b/tests/unit/trajectories/test_tokenized_models.py index 33ee76917..ed0e16376 100644 --- a/tests/unit/trajectories/test_tokenized_models.py +++ b/tests/unit/trajectories/test_tokenized_models.py @@ -238,7 +238,8 @@ def unexpected_dump(*_: object, **__: object) -> object: assert group.trajectories[0].trajectory is source -def test_public_group_tokenization_nan_json_round_trip() -> None: +@pytest.mark.parametrize("logprob", [-0.25, math.nan]) +def test_public_group_tokenization_nan_json_round_trip(logprob: float) -> None: from datetime import datetime from openai.types.chat import ChatCompletion @@ -271,7 +272,7 @@ def test_public_group_tokenization_nan_json_round_trip() -> None: "content": [ { "token": f"token_id:{token_id}", - "logprob": -0.1 * token_id, + "logprob": logprob, "bytes": [], "top_logprobs": [], } @@ -323,6 +324,37 @@ def test_public_group_tokenization_nan_json_round_trip() -> None: tr.TokenizedTrajectoryGroup[tr.TokenizedMultiHistoryTrajectory].model_validate_json( multi_json ) + for group in (single, multi, single.tensorize(), multi.tensorize()): + for restored in ( + type(group).model_validate_json(group.model_dump_json()), + tr.compact_validate(group.compact_dump(), type=type(group)), + ): + assert restored.model_dump_json() == group.model_dump_json() + child = restored.trajectories[0] + assert child.trajectory is restored.trajectory_group.trajectories[0] + histories = ( + child.histories + if isinstance( + child, + ( + tr.TokenizedMultiHistoryTrajectory, + tr.TensorizedMultiHistoryTrajectory, + ), + ) + else [child] + ) + for history in histories: + assert isinstance(history.history, tr.ChatCompletionsHistory) + for source in history.history.message_sources: + assert source is not None + assert ( + source.exchange + is child.trajectory.exchanges.chat_completions[0] + ) + payload = group.model_dump(mode="json") + payload["trajectory_group"]["trajectories"][0]["reward"] = 123 + with pytest.raises(ValueError, match="does not match its source group"): + type(group).model_validate(payload) def test_tokenized_compact_round_trips_retain_source_references() -> None: From b4ef151f033d870fb16e48eb0d997fd487b9f496 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Wed, 16 Sep 2026 03:09:58 +0000 Subject: [PATCH 3/3] Avoid serializing unrelated exchange candidates during rebinding --- src/art/trajectories/_serialization.py | 6 ++- .../trajectories/test_tokenized_models.py | 38 +++++++++++++++++++ 2 files changed, 43 insertions(+), 1 deletion(-) diff --git a/src/art/trajectories/_serialization.py b/src/art/trajectories/_serialization.py index 7f2a001c5..fde796521 100644 --- a/src/art/trajectories/_serialization.py +++ b/src/art/trajectories/_serialization.py @@ -269,6 +269,7 @@ def _rebind_history_sources( MessagesExchange, ResponsesExchange, Trajectory, + _Exchange, ) exchange_types = ( @@ -278,7 +279,7 @@ def _rebind_history_sources( MessagesExchange, ) - def exchanges(value: object) -> list[BaseModel]: + def exchanges(value: object) -> list[_Exchange]: if not isinstance(value, Trajectory): return [] return [ @@ -319,10 +320,13 @@ def visit(value: object) -> None: if replacement is not item: object.__setattr__(value, name, replacement) continue + # Reject unrelated calls before dumping their full payloads. matches = [ exchange for exchange in canonical if type(exchange) is type(item) + and exchange.start_time == item.start_time + and exchange.end_time == item.end_time and ( exchange == item or _equal_with_nan(exchange.model_dump(), item.model_dump()) diff --git a/tests/unit/trajectories/test_tokenized_models.py b/tests/unit/trajectories/test_tokenized_models.py index ed0e16376..1f19d2644 100644 --- a/tests/unit/trajectories/test_tokenized_models.py +++ b/tests/unit/trajectories/test_tokenized_models.py @@ -238,6 +238,44 @@ def unexpected_dump(*_: object, **__: object) -> object: assert group.trajectories[0].trajectory is source +@pytest.mark.parametrize("logprob", [-0.25, math.nan]) +def test_source_rebinding_does_not_dump_unrelated_exchanges( + logprob: float, monkeypatch: pytest.MonkeyPatch +) -> None: + from test_tokenize import _chat_exchange + + exchanges = [ + _chat_exchange(list(range(1, 2 * i + 2)), [2 * i + 2], offset=i) + for i in range(12) + ] + for exchange in exchanges: + old = exchange.response.choices[0].logprobs + assert old and old.content + old.content[0].logprob = logprob + tokenized = art.Trajectory( + exchanges=tr.TrajectoryExchanges(chat_completions=exchanges) + ).tokenize() + payload = tokenized.model_dump_json() + dumps = 0 + original_dump = tr.ChatCompletionsExchange.model_dump + + def counted_dump(self: tr.ChatCompletionsExchange, *args: Any, **kwargs: Any): + nonlocal dumps + dumps += 1 + return original_dump(self, *args, **kwargs) + + monkeypatch.setattr(tr.ChatCompletionsExchange, "model_dump", counted_dump) + restored = tr.TokenizedTrajectory.model_validate_json(payload) + assert isinstance(restored.history, tr.ChatCompletionsHistory) + sources = [s for s in restored.history.message_sources if s is not None] + assert dumps <= (2 * len(sources) if math.isnan(logprob) else 0) + assert all( + any(s.exchange is e for e in restored.trajectory.exchanges.chat_completions) + for s in sources + ) + assert restored.model_dump_json() == payload + + @pytest.mark.parametrize("logprob", [-0.25, math.nan]) def test_public_group_tokenization_nan_json_round_trip(logprob: float) -> None: from datetime import datetime