diff --git a/src/aiperf/timing/strategies/fixed_schedule.py b/src/aiperf/timing/strategies/fixed_schedule.py index a557130555..b1afb0c2f5 100644 --- a/src/aiperf/timing/strategies/fixed_schedule.py +++ b/src/aiperf/timing/strategies/fixed_schedule.py @@ -153,7 +153,10 @@ async def handle_credit_return( # This contains the delay_ms or timestamp_ms for the next turn next_meta = self._conversation_source.get_next_turn_metadata(credit) - turn = TurnToSend.from_previous_credit(credit) + # Pass next_meta so has_forks rides onto the continuation turn: the + # sticky router defers parent-entry eviction until DAG children drain + # (dropping it premature-evicts a fork-bearing parent's later turns). + turn = TurnToSend.from_previous_credit(credit, next_meta) if next_meta.timestamp_ms is not None: self._scheduler.schedule_at_perf_sec( diff --git a/src/aiperf/timing/strategies/user_centric_rate.py b/src/aiperf/timing/strategies/user_centric_rate.py index 360e3770c2..b8aa80287b 100644 --- a/src/aiperf/timing/strategies/user_centric_rate.py +++ b/src/aiperf/timing/strategies/user_centric_rate.py @@ -340,7 +340,11 @@ async def handle_credit_return( raise ValueError( f"User not found for x_correlation_id: {credit.x_correlation_id}" ) - turn = TurnToSend.from_previous_credit(credit) + # Pass next-turn metadata so has_forks rides onto the continuation turn + # (the sticky router defers parent-entry eviction until DAG children + # drain); dropping it premature-evicts a fork-bearing parent's turns. + meta = self._conversation_source.get_next_turn_metadata(credit) + turn = TurnToSend.from_previous_credit(credit, meta) # If the next turn time already passed, the max() will # re-align their schedule to account for the delay. diff --git a/tests/unit/timing/strategies/test_fixed_schedule.py b/tests/unit/timing/strategies/test_fixed_schedule.py index 16436ff8ec..8d0fc55719 100644 --- a/tests/unit/timing/strategies/test_fixed_schedule.py +++ b/tests/unit/timing/strategies/test_fixed_schedule.py @@ -227,6 +227,33 @@ async def test_non_final_turn_schedules_next(self) -> None: await strategy.handle_credit_return(credit) scheduler.schedule_at_perf_sec.assert_called_once() + async def test_continuation_turn_carries_has_forks(self) -> None: + """has_forks from the NEXT turn's metadata must ride onto the + continuation turn (the sticky router defers parent-entry eviction until + DAG children drain). Regression: from_previous_credit(credit) was called + without next_meta, so every continuation turn got has_forks=False.""" + strategy, scheduler, _ = make_strategy([(0, "c1"), (100, "c1")]) + await strategy.setup_phase() + # Mark the next turn (index 1) as fork-bearing. + meta = strategy._conversation_source._metadata_lookup["c1"] + meta.turns[1].has_forks = True + + captured: list = [] + strategy._credit_issuer.issue_credit = ( + lambda turn: captured.append(turn) or True + ) + credit = Credit( + id=1, + phase=CreditPhase.PROFILING, + conversation_id="c1", + x_correlation_id="corr-c1", + turn_index=0, + num_turns=2, + issued_at_ns=1000, + ) + await strategy.handle_credit_return(credit) + assert captured and captured[0].has_forks is True + @pytest.mark.asyncio class TestFixedScheduleTimestampConversion: diff --git a/tests/unit/timing/strategies/test_user_centric_rate.py b/tests/unit/timing/strategies/test_user_centric_rate.py index e0b86d3eb2..1552ed75c6 100644 --- a/tests/unit/timing/strategies/test_user_centric_rate.py +++ b/tests/unit/timing/strategies/test_user_centric_rate.py @@ -5,10 +5,13 @@ import pytest from aiperf.common.enums import CreditPhase -from aiperf.plugin.enums import TimingMode +from aiperf.common.models import ConversationMetadata, DatasetMetadata, TurnMetadata +from aiperf.credit.structs import Credit +from aiperf.plugin.enums import DatasetSamplingStrategy, TimingMode from aiperf.timing.config import CreditPhaseConfig +from aiperf.timing.conversation_source import ConversationSource from aiperf.timing.strategies.user_centric_rate import User, UserCentricStrategy -from tests.unit.timing.conftest import OrchestratorHarness +from tests.unit.timing.conftest import OrchestratorHarness, make_sampler TWO_TURN = [("c1", 2), ("c2", 2), ("c3", 2), ("c4", 2), ("c5", 2)] MULTI_TURN = [("c1", 3), ("c2", 3), ("c3", 3), ("c4", 3)] @@ -143,6 +146,75 @@ async def test_stops_at_session_count(self, create_orchestrator_harness) -> None assert len([c for c in h.sent_credits if c.turn_index == 0]) == 10 +def make_strategy(num_turns: int = 2) -> UserCentricStrategy: + """Build a UserCentricStrategy over a real ConversationSource with one + num_turns-turn conversation and mocked scheduler/issuer/lifecycle.""" + scheduler = MagicMock() + stop_checker = MagicMock() + issuer = MagicMock() + issuer.issue_credit = lambda *a, **k: True + lifecycle = MagicMock() + ds = DatasetMetadata( + conversations=[ + ConversationMetadata( + conversation_id="c1", + turns=[TurnMetadata(delay_ms=None) for _ in range(num_turns)], + ) + ], + sampling_strategy=DatasetSamplingStrategy.SEQUENTIAL, + ) + sampler = make_sampler(["c1"], DatasetSamplingStrategy.SEQUENTIAL) + src = ConversationSource(ds, sampler) + cfg = CreditPhaseConfig( + phase=CreditPhase.PROFILING, + timing_mode=TimingMode.USER_CENTRIC_RATE, + request_rate=10.0, + num_users=1, + total_expected_requests=num_turns, + ) + return UserCentricStrategy( + config=cfg, + conversation_source=src, + scheduler=scheduler, + stop_checker=stop_checker, + credit_issuer=issuer, + lifecycle=lifecycle, + ) + + +@pytest.mark.asyncio +class TestUserCentricCreditReturn: + async def test_continuation_turn_carries_has_forks(self) -> None: + """has_forks from the NEXT turn's metadata must ride onto the + continuation turn (the sticky router defers parent-entry eviction until + DAG children drain). Regression: from_previous_credit(credit) was called + without next-turn metadata, so every continuation turn got + has_forks=False.""" + strategy = make_strategy() + await strategy.setup_phase() + # Mark the next turn (index 1) as fork-bearing. + meta = strategy._conversation_source._metadata_lookup["c1"] + meta.turns[1].has_forks = True + + captured: list = [] + strategy._credit_issuer.issue_credit = ( + lambda turn: captured.append(turn) or True + ) + # setup_phase registered the t=0 replacement user; use its session id. + x_correlation_id = next(iter(strategy._session_to_user)) + credit = Credit( + id=1, + phase=CreditPhase.PROFILING, + conversation_id="c1", + x_correlation_id=x_correlation_id, + turn_index=0, + num_turns=2, + issued_at_ns=1000, + ) + await strategy.handle_credit_return(credit) + assert captured and captured[0].has_forks is True + + @pytest.mark.asyncio class TestRealisticScenarios: async def test_chat_benchmark(self, create_orchestrator_harness) -> None: