fix: clear user sessions between load stages (kubernetes-sigs#459)

## Summary

- `LocalUserSession._instances` was never cleared between load stages,
causing session context to accumulate across stage boundaries and
prompts to grow indefinitely in multi-turn shared-prefix workloads.
- **Non-mp path** (`num_workers=0`): call
`LocalUserSession.clear_instances()` between stages in
`LoadGenerator.run()`.
- **Mp path** (`num_workers>0`): add an `mp.Barrier` to synchronize
workers and main at stage boundaries. Workers finish in-flight requests,
clear sessions, then wait at the barrier. Main waits after
`request_queue.join()`, ensuring all workers have cleaned up before the
next stage begins. Without the barrier, `request_phase` can be cleared
and re-set before workers notice, causing them to miss the stage
boundary entirely.

## Test plan

- [x] Unit tests for `LocalUserSession.clear_instances()` (singleton
reset, context reset)
- [x] Integration test: non-mp LoadGenerator with two stages verifies
stage 1 prompts contain no stage 0 context
- [x] Integration test: mp LoadGenerator (`num_workers=1`) with two
stages verifies stage 1 prompts contain no stage 0 context via mp.Queue
- [x] mypy --strict passes
- [x] All 5 tests pass reliably across 10 consecutive runs

Fixes kubernetes-sigs#447 
Fixes kubernetes-sigs#444

Author: alonh

inference_perf/loadgen/load_generator.py

@@ -19,6 +19,7 @@
 from .load_timer import LoadTimer, ConstantLoadTimer, PoissonLoadTimer, TraceReplayLoadTimer
 from inference_perf.datagen import DataGenerator, SessionGenerator, LazyLoadDataMixin
 from inference_perf.apis import InferenceAPIData
+from inference_perf.apis.user_session import LocalUserSession
 from inference_perf.client.modelserver import ModelServerClient
 from inference_perf.client.modelserver.otel_instrumentation import get_otel_instrumentation
 from inference_perf.circuit_breaker import get_circuit_breaker
@@ -57,7 +58,7 @@
 import time
 import multiprocessing as mp
 from queue import Empty
-from multiprocessing.synchronize import Event as SyncEvent
+from multiprocessing.synchronize import Barrier as SyncBarrier, Event as SyncEvent
 from multiprocessing.sharedctypes import Synchronized
 from concurrent.futures import TimeoutError
 from functools import partial
@@ -100,6 +101,7 @@ def __init__(
         active_requests_counter: "Synchronized[int]",
         shared_max_concurrency: Optional["Synchronized[int]"],
         base_seed: int,
+        stage_barrier: Optional[SyncBarrier] = None,
     ):
         super().__init__(daemon=True)  # kill worker process if main process exit unexpected
         self.id = id
@@ -115,6 +117,7 @@ def __init__(
         self.shared_max_concurrency = shared_max_concurrency
         self.skip = False
         self.base_seed = base_seed
+        self.stage_barrier = stage_barrier
 
     async def loop(self) -> None:
         # The self.shared_max_concurrency is initialized to self.max_concurrency
@@ -246,6 +249,9 @@ async def schedule_client(
             if not self.request_phase.is_set():
                 await gather(*tasks)
                 tasks = []
+                LocalUserSession.clear_instances()
+                if self.stage_barrier:
+                    self.stage_barrier.wait()
                 logger.debug(f"[Worker {self.id}] waiting for next phase")
                 self.request_phase.wait()
 
@@ -902,6 +908,9 @@ async def mp_run(self, client: ModelServerClient) -> None:
         request_phase: SyncEvent = mp.Event()
         stop_signal: SyncEvent = mp.Event()
         cancel_signal: SyncEvent = mp.Event()
+        # Synchronize workers and main at stage boundaries so workers finish
+        # in-flight requests and clear session state before the next stage begins.
+        stage_barrier: SyncBarrier = mp.Barrier(self.num_workers + 1)
         # start workers in the request phase
         request_phase.set()
 
@@ -927,6 +936,7 @@ async def mp_run(self, client: ModelServerClient) -> None:
                     active_requests_counter,
                     shared_max_concurrency,
                     self.base_seed,
+                    stage_barrier,
                 )
             )
             self.workers[-1].start()
@@ -1021,6 +1031,8 @@ async def mp_run(self, client: ModelServerClient) -> None:
                 else:
                     raise Exception(f"Stage {stage_id} has the wrong load type")
 
+                stage_barrier.wait()
+
                 # If we encountered a SIGINT, we can break out of run stages loop
                 if self.interrupt_sig:
                     break
@@ -1109,6 +1121,7 @@ async def run(self, client: ModelServerClient) -> None:
                     concurrency_level=None,
                 )
                 progress.update(overall_task, advance=1)
+                LocalUserSession.clear_instances()
                 if self.stageInterval and stage_id < len(self.stages) - 1:
                     await sleep(self.stageInterval)
 

tests/apis/test_user_session.py

@@ -0,0 +1,241 @@
+# Copyright 2026 The Kubernetes Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for LocalUserSession lifecycle."""
+
+import multiprocessing as mp
+import pytest
+from collections import defaultdict
+from queue import Empty
+from typing import List, Optional, Tuple
+from unittest.mock import MagicMock
+
+from inference_perf.apis.user_session import LocalUserSession, UserSessionCompletionAPIData
+from inference_perf.apis import InferenceAPIData
+from inference_perf.client.modelserver.base import ModelServerClient, PrometheusMetricMetadata
+from inference_perf.config import (
+    APIConfig,
+    APIType,
+    DataConfig,
+    DataGenType,
+    LoadConfig,
+    LoadType,
+    SharedPrefix,
+    StandardLoadStage,
+)
+from inference_perf.datagen.shared_prefix_datagen import SharedPrefixDataGenerator
+from inference_perf.loadgen.load_generator import LoadGenerator
+
+
+def _mock_tokenizer() -> MagicMock:
+    tok = MagicMock()
+    hf = MagicMock()
+    hf.vocab_size = 1000
+    hf.decode = MagicMock(side_effect=lambda ids, **kw: f"tok_{len(ids)}")
+    hf.batch_decode = MagicMock(side_effect=lambda batch, **kw: [f"tok_{len(ids)}" for ids in batch])
+    tok.get_tokenizer.return_value = hf
+    return tok
+
+
+def _make_datagen(num_groups: int = 1, num_prompts_per_group: int = 1) -> SharedPrefixDataGenerator:
+    api_config = APIConfig(type=APIType.Completion)
+    data_config = DataConfig(
+        type=DataGenType.SharedPrefix,
+        shared_prefix=SharedPrefix(
+            num_groups=num_groups,
+            num_prompts_per_group=num_prompts_per_group,
+            enable_multi_turn_chat=True,
+            system_prompt_len=5,
+            question_len=5,
+            output_len=5,
+            seed=42,
+        ),
+    )
+    return SharedPrefixDataGenerator(api_config, data_config, _mock_tokenizer())
+
+
+class SessionTrackingClient(ModelServerClient):
+    """Minimal client that exercises the UserSession to_payload / update_context
+    lifecycle and records the prompt sent per stage.
+
+    When prompt_queue is set (mp mode), prompts are sent to the queue so the
+    main process can read them.  Otherwise they are stored in-process."""
+
+    def __init__(self, prompt_queue: Optional["mp.Queue[Tuple[int, str]]"] = None) -> None:
+        self.api_config = APIConfig(type=APIType.Completion)
+        self.timeout = None
+        self.prompts_by_stage: dict[int, list[str]] = defaultdict(list)
+        self._prompt_queue = prompt_queue
+
+    async def process_request(
+        self, data: InferenceAPIData, stage_id: int, scheduled_time: float, lora_adapter: Optional[str] = None
+    ) -> None:
+        payload = await data.to_payload("model", 64, False, False)
+        prompt = payload["prompt"]
+        self.prompts_by_stage[stage_id].append(prompt)
+
+        if self._prompt_queue is not None:
+            self._prompt_queue.put((stage_id, prompt))
+
+        if isinstance(data, UserSessionCompletionAPIData):
+            data.user_session.update_context(prompt + f" RESPONSE_STAGE{stage_id}")
+
+    def get_supported_apis(self) -> List[APIType]:
+        return [APIType.Completion]
+
+    def get_prometheus_metric_metadata(self) -> PrometheusMetricMetadata:
+        raise NotImplementedError
+
+
+class TestLocalUserSessionLifecycle:
+    def setup_method(self) -> None:
+        LocalUserSession.clear_instances()
+
+    def teardown_method(self) -> None:
+        LocalUserSession.clear_instances()
+
+    def test_get_instance_returns_same_object(self) -> None:
+        s1 = LocalUserSession.get_instance("sess_a")
+        s2 = LocalUserSession.get_instance("sess_a")
+        assert s1 is s2
+
+    def test_clear_instances_resets_all_sessions(self) -> None:
+        s1 = LocalUserSession.get_instance("sess_a")
+        s1.contexts = "accumulated context"
+        s1._current_round = 3
+
+        s2 = LocalUserSession.get_instance("sess_b")
+        s2.contexts = "other context"
+
+        LocalUserSession.clear_instances()
+
+        new_s1 = LocalUserSession.get_instance("sess_a")
+        new_s2 = LocalUserSession.get_instance("sess_b")
+
+        assert new_s1 is not s1
+        assert new_s2 is not s2
+        assert new_s1.contexts == ""
+        assert new_s1._current_round == 0
+        assert new_s2.contexts == ""
+
+    def test_context_does_not_leak_across_stage_boundary(self) -> None:
+        """
+        Simulates two stages. After clearing between stages, a session
+        obtained via get_instance must have empty context and round 0.
+
+
+        """
+        session = LocalUserSession.get_instance("user_0")
+        session.contexts = "system prompt Q1 A1 Q2 A2"
+        session._current_round = 2
+
+        LocalUserSession.clear_instances()
+
+        session_s1 = LocalUserSession.get_instance("user_0")
+        assert session_s1.contexts == ""
+        assert session_s1._current_round == 0
+
+    @pytest.mark.asyncio
+    async def test_loadgen_does_not_leak_session_context_across_stages(self) -> None:
+        """
+        Run the real LoadGenerator with two stages (num_workers=0) using a
+        client that exercises the full to_payload / update_context lifecycle.
+
+        Stage 0 builds up session context.  Stage 1 must NOT see that context
+        in its prompts — if it does, sessions leaked across the stage boundary.
+
+
+        """
+        datagen = _make_datagen()
+        # High rate ensures multiple requests per stage so sessions accumulate context.
+        # ExceptionGroup may fire due to strict zip in the non-mp path when
+        # floating-point timer values land exactly at the stage boundary.
+        load_config = LoadConfig(
+            type=LoadType.CONSTANT,
+            stages=[
+                StandardLoadStage(rate=10, duration=1),
+                StandardLoadStage(rate=10, duration=1),
+            ],
+            num_workers=0,
+            interval=0,
+        )
+        loadgen = LoadGenerator(datagen, load_config)
+        client = SessionTrackingClient()
+
+        try:
+            await loadgen.run(client)
+        except ExceptionGroup:
+            pass
+
+        assert 0 in client.prompts_by_stage, "Expected prompts in stage 0"
+        if 1 not in client.prompts_by_stage:
+            pytest.skip("Stage 1 did not produce prompts (ExceptionGroup aborted early)")
+
+        for prompt in client.prompts_by_stage[1]:
+            assert "RESPONSE_STAGE0" not in prompt, (
+                f"Stage 1 prompt contains stage 0 response context — sessions "
+                f"were not cleared between stages.\n"
+                f"  stage 1 prompt: {prompt!r}"
+            )
+
+    @pytest.mark.asyncio
+    async def test_loadgen_mp_does_not_leak_session_context_across_stages(self) -> None:
+        """
+        Same as the non-mp test but with num_workers=1 so requests flow
+        through a forked Worker subprocess.
+
+        The client writes (stage_id, prompt) tuples to an mp.Queue that the
+        main process drains after the run.  If any stage-1 prompt contains
+        RESPONSE_STAGE0, sessions leaked across the stage boundary inside
+        the worker process.
+
+
+        """
+        mp.set_start_method("fork", force=True)
+
+        datagen = _make_datagen()
+        load_config = LoadConfig(
+            type=LoadType.CONSTANT,
+            stages=[
+                StandardLoadStage(rate=10, duration=1),
+                StandardLoadStage(rate=10, duration=1),
+            ],
+            num_workers=1,
+            interval=0,
+        )
+
+        prompt_queue: "mp.Queue[Tuple[int, str]]" = mp.Queue()
+        client = SessionTrackingClient(prompt_queue=prompt_queue)
+        loadgen = LoadGenerator(datagen, load_config)
+
+        await loadgen.run(client)
+        await loadgen.stop()
+
+        prompts_by_stage: dict[int, list[str]] = defaultdict(list)
+        while True:
+            try:
+                stage_id, prompt = prompt_queue.get_nowait()
+                prompts_by_stage[stage_id].append(prompt)
+            except Empty:
+                break
+
+        assert 0 in prompts_by_stage, "Expected prompts in stage 0"
+        assert 1 in prompts_by_stage, "Expected prompts in stage 1"
+
+        for prompt in prompts_by_stage[1]:
+            assert "RESPONSE_STAGE0" not in prompt, (
+                f"Stage 1 prompt contains stage 0 response context — sessions "
+                f"were not cleared between stages in worker subprocess.\n"
+                f"  stage 1 prompt: {prompt!r}"
+            )