From 614d4bfa5348ee7b253a22f6218c51b06bd4c131 Mon Sep 17 00:00:00 2001 From: Joseph Viviano Date: Sat, 22 Aug 2026 12:49:11 -0400 Subject: [PATCH] Reject a malformed BASELINE_SYNC instead of ignoring it MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit _apply_baseline_sync read the payload with .get and skipped the assignment when it came back None. A sync with a missing, null or non-finite value left the manager on its shard-local threshold with no signal, so it filtered against a different cutoff from its siblings — the state the message exists to remove. The test that pinned the old behaviour is replaced by one that covers a missing key, a null value, a non-finite value and a non-dict payload. --- src/gfn/containers/replay_buffer_manager.py | 27 ++++++++++++++++++--- testing/test_replay_buffer.py | 21 ++++++++++++---- 2 files changed, 40 insertions(+), 8 deletions(-) diff --git a/src/gfn/containers/replay_buffer_manager.py b/src/gfn/containers/replay_buffer_manager.py index 32e500f3..53465812 100644 --- a/src/gfn/containers/replay_buffer_manager.py +++ b/src/gfn/containers/replay_buffer_manager.py @@ -142,10 +142,31 @@ def _apply_baseline_sync(self, msg) -> None: Sent by a coordinator aggregating local baselines across buffer managers. The manager doesn't need to know who sent it or why. + + A malformed message is an error, not a no-op: the adopted value + gates what every assigned worker sends, so a manager that quietly + keeps its shard-local threshold filters against a different cutoff + from its siblings — the exact condition this message exists to + remove, and one that leaves no trace in the run. """ - global_baseline = msg.message_data.get("global_baseline_log_reward") - if global_baseline is not None: - self._global_baseline = float(global_baseline) + if not isinstance(msg.message_data, dict): + raise ValueError( + "BASELINE_SYNC message_data must be a dict, got " + f"{type(msg.message_data).__name__}" + ) + value = msg.message_data.get("global_baseline_log_reward") + if value is None: + raise ValueError( + "BASELINE_SYNC has no usable 'global_baseline_log_reward'; " + f"message_data={msg.message_data}" + ) + global_baseline = float(value) + if not math.isfinite(global_baseline): + raise ValueError( + f"BASELINE_SYNC global_baseline_log_reward is {global_baseline}; " + "a non-finite baseline keeps every trajectory or none" + ) + self._global_baseline = global_baseline def _inject_baseline_log_reward( self, score_dict: dict[str, float], incoming diff --git a/testing/test_replay_buffer.py b/testing/test_replay_buffer.py index 45d0b3fe..2fdefd69 100644 --- a/testing/test_replay_buffer.py +++ b/testing/test_replay_buffer.py @@ -790,7 +790,19 @@ def test_manager_baseline_sync_overrides_local(simple_env, trajectories): assert score["baseline_log_reward"] == pytest.approx(-7.5) -def test_manager_baseline_sync_ignores_missing_payload(simple_env, trajectories): +@pytest.mark.parametrize( + "payload", + [ + {}, + {"global_baseline_log_reward": None}, + {"global_baseline_log_reward": float("-inf")}, + None, + ], + ids=["missing", "none", "non_finite", "not_a_dict"], +) +def test_manager_baseline_sync_rejects_malformed_payload(simple_env, payload): + # Keeping the shard-local threshold on a malformed sync is the + # inconsistency this message exists to remove, and it leaves no trace. mgr = ReplayBufferManager( simple_env, rank=0, @@ -798,10 +810,9 @@ def test_manager_baseline_sync_ignores_missing_payload(simple_env, trajectories) capacity=5, store_locally=False, ) - mgr._apply_baseline_sync(Message(MessageType.BASELINE_SYNC, {})) - score: dict = {"score": 0.0} - mgr._inject_baseline_log_reward(score, trajectories) - assert score["baseline_log_reward"] == pytest.approx(0.0) # falls back to EMA. + with pytest.raises(ValueError, match="BASELINE_SYNC"): + mgr._apply_baseline_sync(Message(MessageType.BASELINE_SYNC, payload)) + assert mgr._global_baseline is None def test_manager_multi_manager_baseline_warning(simple_env):