From 6d2f7851739d542b1e3c06e4034560c0833a2592 Mon Sep 17 00:00:00 2001
From: Joel Tony <github@jaytau.com>
Date: Sun, 2 Aug 2026 23:16:21 +0530
Subject: [PATCH 1/3] fix(kv-cache): drain finite replay requests

---
 kv_cache_benchmark/kv_cache/benchmark.py  |  26 +++-
 kv_cache_benchmark/tests/test_kv_cache.py | 139 +++++++++++++---------
 2 files changed, 104 insertions(+), 61 deletions(-)

diff --git a/kv_cache_benchmark/kv_cache/benchmark.py b/kv_cache_benchmark/kv_cache/benchmark.py
index bf6ab1d2..87fb06ad 100755
--- a/kv_cache_benchmark/kv_cache/benchmark.py
+++ b/kv_cache_benchmark/kv_cache/benchmark.py
@@ -184,6 +184,7 @@ def __init__(self,
         }
         self.results_lock = threading.Lock()
         self.stop_event: Optional[threading.Event] = None
+        self.producer_done = threading.Event()
         self.rag_ingest_done = threading.Event() if self.enable_rag else None
 
     def _ingest_rag_documents(self, num_docs: int, stop_event: Optional[threading.Event] = None):
@@ -258,6 +259,15 @@ def _burst_trace_iterator(self):
                 logger.error(f"Error reading trace file {filepath}: {e}")
                 sys.exit(1)
 
+    def _mark_producer_done(self, stop_event: threading.Event):
+        """Mark finite input complete and stop once every request has finished."""
+        if stop_event.is_set():
+            return
+        self.producer_done.set()
+        with self.results_lock:
+            if self.results['requests_completed'] >= self.request_counter:
+                stop_event.set()
+
     def _generate_requests_from_trace(self, stop_event: threading.Event):
         """Generates InferenceRequest objects from the streaming trace iterator."""
         speedup = self.trace_speedup
@@ -324,6 +334,7 @@ def _generate_requests_from_trace(self, stop_event: threading.Event):
 
             if rows_in_cycle == 0:
                 logger.warning("BurstGPT trace yielded 0 rows.")
+                self._mark_producer_done(stop_event)
                 break
 
             if cycles_remaining > 0:
@@ -331,8 +342,7 @@ def _generate_requests_from_trace(self, stop_event: threading.Event):
                 if cycles_remaining == 0:
                     logger.info(f"Completed {self.replay_cycles} replay cycle(s). "
                                 f"Trace total_tokens sum: {trace_total_tokens_sum:,}")
-                    if self.stop_event:
-                        self.stop_event.set()
+                    self._mark_producer_done(stop_event)
                     break
 
             prev_timestamp = None
@@ -360,8 +370,7 @@ def _generate_requests_from_dataset(self, stop_event: threading.Event):
                         cycles_remaining -= 1
                         if cycles_remaining == 0:
                             logger.info(f"Completed {self.replay_cycles} ShareGPT replay cycle(s).")
-                            if self.stop_event:
-                                self.stop_event.set()
+                            self._mark_producer_done(stop_event)
                             return
                     conversation_iterator = iter(self.sharegpt_loader.iterate_conversations(shuffle=True))
                     continue
@@ -633,8 +642,12 @@ def process_requests(self, stop_event: threading.Event):
                 self.results['generation_latencies'].append(generation_latency)
 
                 if self.max_requests > 0 and self.results['requests_completed'] >= self.max_requests:
-                    if self.stop_event:
-                        self.stop_event.set()
+                    stop_event.set()
+                elif (
+                    self.producer_done.is_set()
+                    and self.results['requests_completed'] >= self.request_counter
+                ):
+                    stop_event.set()
 
             self.qos_monitor.record_request(request)
 
@@ -1216,6 +1229,7 @@ def run(self) -> Dict:
 
         stop_event = threading.Event()
         self.stop_event = stop_event
+        self.producer_done.clear()
 
         threads = []
         if self.use_dataset:
diff --git a/kv_cache_benchmark/tests/test_kv_cache.py b/kv_cache_benchmark/tests/test_kv_cache.py
index c6d14015..a91e3b7a 100644
--- a/kv_cache_benchmark/tests/test_kv_cache.py
+++ b/kv_cache_benchmark/tests/test_kv_cache.py
@@ -2224,79 +2224,108 @@ def test_streaming_iterator_preserves_order(self, benchmark_with_trace):
             assert timestamps[i] > timestamps[i-1], \
                 f"Timestamp at index {i} ({timestamps[i]}) should be > {timestamps[i-1]}"
 
-    def test_replay_cycles_one_pass(self, trace_dir):
-        """With replay_cycles=1, generator should process all rows once then stop."""
-        import threading
-        model_config = MODEL_CONFIGS['tiny-1b']
-        bench = IntegratedBenchmark(
-            model_config=model_config,
-            num_users=5,
-            gpu_memory_gb=0,
-            cpu_memory_gb=0.01,
-            duration_seconds=60,
-            use_burst_trace=True,
-            burst_trace_path=str(trace_dir),
-            generation_mode=GenerationMode.NONE,
-            trace_speedup=0,
-            replay_cycles=1,
-        )
+    @pytest.mark.parametrize(('cycles', 'expected_requests'), [(1, 8), (2, 16)])
+    def test_finite_replay_drains_queue(
+        self, benchmark_with_trace, monkeypatch, cycles, expected_requests
+    ):
+        """Finite replay should stop only after workers drain every generated request."""
+        bench = benchmark_with_trace
+        bench.replay_cycles = cycles
+        bench.prefix_cache_manager = None
+        bench.prefill_only = True
+        monkeypatch.setattr(bench.cache, 'allocate_cache', lambda *args, **kwargs: (True, 'cpu', 0.0))
 
         stop_event = threading.Event()
         bench.stop_event = stop_event
-
-        # Run generator in a thread
-        gen_thread = threading.Thread(
+        producer = threading.Thread(
             target=bench._generate_requests_from_trace,
             args=(stop_event,),
             daemon=True
         )
-        gen_thread.start()
-        gen_thread.join(timeout=10)
+        producer.start()
+        producer.join(timeout=10)
 
-        # stop_event should have been set by the generator after 1 cycle
-        assert stop_event.is_set(), "stop_event should be set after replay_cycles=1 completes"
+        assert not producer.is_alive()
+        assert bench.producer_done.is_set()
+        assert not stop_event.is_set()
+        assert bench.request_queue.qsize() == expected_requests
 
-        # Queue should have exactly 8 requests (5 + 3)
-        count = 0
-        while not bench.request_queue.empty():
-            bench.request_queue.get_nowait()
-            count += 1
-        assert count == 8, f"Expected 8 requests from 1 cycle, got {count}"
+        worker = threading.Thread(target=bench.process_requests, args=(stop_event,), daemon=True)
+        worker.start()
+        worker.join(timeout=10)
 
-    def test_replay_cycles_two_passes(self, trace_dir):
-        """With replay_cycles=2, generator should process all rows twice."""
-        import threading
-        model_config = MODEL_CONFIGS['tiny-1b']
-        bench = IntegratedBenchmark(
-            model_config=model_config,
-            num_users=5,
-            gpu_memory_gb=0,
-            cpu_memory_gb=0.01,
-            duration_seconds=60,
-            use_burst_trace=True,
-            burst_trace_path=str(trace_dir),
-            generation_mode=GenerationMode.NONE,
-            trace_speedup=0,
-            replay_cycles=2,
-        )
+        assert not worker.is_alive()
+        assert stop_event.is_set()
+        assert bench.results['requests_completed'] == expected_requests
+        assert bench.request_queue.empty()
 
+    def test_hard_stop_does_not_drain_queue(self, benchmark_with_trace):
+        """An explicit hard stop should leave queued requests unprocessed."""
         stop_event = threading.Event()
-        bench.stop_event = stop_event
+        benchmark_with_trace._generate_requests_from_trace(stop_event)
+        queued_requests = benchmark_with_trace.request_queue.qsize()
 
-        gen_thread = threading.Thread(
-            target=bench._generate_requests_from_trace,
+        stop_event.set()
+        worker = threading.Thread(
+            target=benchmark_with_trace.process_requests,
             args=(stop_event,),
             daemon=True
         )
-        gen_thread.start()
-        gen_thread.join(timeout=10)
+        worker.start()
+        worker.join(timeout=2)
 
+        assert not worker.is_alive()
+        assert benchmark_with_trace.results['requests_completed'] == 0
+        assert benchmark_with_trace.request_queue.qsize() == queued_requests == 8
+
+    def test_producer_done_stops_after_workers_already_drained(self, benchmark_with_trace):
+        """Producer completion should close the worker-first race."""
+        benchmark_with_trace.request_counter = 8
+        benchmark_with_trace.results['requests_completed'] = 8
+        stop_event = threading.Event()
+
+        benchmark_with_trace._mark_producer_done(stop_event)
+
+        assert benchmark_with_trace.producer_done.is_set()
         assert stop_event.is_set()
-        count = 0
-        while not bench.request_queue.empty():
-            bench.request_queue.get_nowait()
-            count += 1
-        assert count == 16, f"Expected 16 requests from 2 cycles, got {count}"
+
+    def test_empty_trace_completes_immediately(self, benchmark_with_trace, monkeypatch):
+        """A finite trace with no valid rows should not wait for the duration limit."""
+        monkeypatch.setattr(benchmark_with_trace, '_burst_trace_iterator', lambda: iter(()))
+        stop_event = threading.Event()
+
+        benchmark_with_trace._generate_requests_from_trace(stop_event)
+
+        assert benchmark_with_trace.producer_done.is_set()
+        assert stop_event.is_set()
+
+    def test_sharegpt_finite_replay_marks_producer_done(self, tmp_path):
+        """ShareGPT finite replay should use graceful completion too."""
+        dataset = tmp_path / 'sharegpt.json'
+        dataset.write_text(
+            '[{"id":"conv_1","conversations":['
+            '{"from":"human","value":"first"},{"from":"gpt","value":"reply"},'
+            '{"from":"human","value":"second"},{"from":"gpt","value":"reply"}]}]'
+        )
+        bench = IntegratedBenchmark(
+            model_config=MODEL_CONFIGS['tiny-1b'],
+            num_users=2,
+            gpu_memory_gb=0,
+            cpu_memory_gb=0.01,
+            duration_seconds=5,
+            dataset_path=str(dataset),
+            max_conversations=1,
+            generation_mode=GenerationMode.NONE,
+            replay_cycles=1,
+        )
+        stop_event = threading.Event()
+
+        bench._generate_requests_from_dataset(stop_event)
+
+        assert bench.request_counter == 2
+        assert bench.request_queue.qsize() == 2
+        assert bench.producer_done.is_set()
+        assert not stop_event.is_set()
 
     def test_total_tokens_tracked(self, benchmark_with_trace):
         """Total tokens from trace should be summed correctly."""

From 82ce426e296b01dbc9f6acf602849cf27ae16e27 Mon Sep 17 00:00:00 2001
From: Joel Tony <github@jaytau.com>
Date: Mon, 3 Aug 2026 19:27:37 +0530
Subject: [PATCH 2/3] fix(kv-cache): harden replay shutdown

---
 kv_cache_benchmark/kv_cache/benchmark.py    | 258 ++++++++++++++------
 kv_cache_benchmark/kv_cache/prefix_cache.py |  33 ++-
 kv_cache_benchmark/kv_cache/rag.py          |  19 +-
 kv_cache_benchmark/kv_cache/workload.py     |   5 +-
 kv_cache_benchmark/tests/test_kv_cache.py   | 221 +++++++++++++++--
 5 files changed, 441 insertions(+), 95 deletions(-)

diff --git a/kv_cache_benchmark/kv_cache/benchmark.py b/kv_cache_benchmark/kv_cache/benchmark.py
index 87fb06ad..b208ee81 100755
--- a/kv_cache_benchmark/kv_cache/benchmark.py
+++ b/kv_cache_benchmark/kv_cache/benchmark.py
@@ -183,8 +183,8 @@ def __init__(self,
             'seed': self.seed,
         }
         self.results_lock = threading.Lock()
-        self.stop_event: Optional[threading.Event] = None
         self.producer_done = threading.Event()
+        self._thread_error: Optional[BaseException] = None
         self.rag_ingest_done = threading.Event() if self.enable_rag else None
 
     def _ingest_rag_documents(self, num_docs: int, stop_event: Optional[threading.Event] = None):
@@ -208,7 +208,12 @@ def _ingest_rag_documents(self, num_docs: int, stop_event: Optional[threading.Ev
             if stop_event and stop_event.is_set():
                 break
             doc_tokens = random.randint(token_min, token_max)
-            self.rag_manager.ingest_document(f"doc_{i:04d}", doc_tokens, self.model_config)
+            self.rag_manager.ingest_document(
+                f"doc_{i:04d}",
+                doc_tokens,
+                self.model_config,
+                stop_event=stop_event
+            )
 
         if self.rag_ingest_done:
             self.rag_ingest_done.set()
@@ -268,6 +273,15 @@ def _mark_producer_done(self, stop_event: threading.Event):
             if self.results['requests_completed'] >= self.request_counter:
                 stop_event.set()
 
+    def _next_request_id(self) -> Optional[int]:
+        """Reserve the next request ID without producing beyond max_requests."""
+        with self.counter_lock:
+            if self.max_requests > 0 and self.request_counter >= self.max_requests:
+                return None
+            request_id = self.request_counter
+            self.request_counter += 1
+            return request_id
+
     def _generate_requests_from_trace(self, stop_event: threading.Event):
         """Generates InferenceRequest objects from the streaming trace iterator."""
         speedup = self.trace_speedup
@@ -287,22 +301,15 @@ def _generate_requests_from_trace(self, stop_event: threading.Event):
 
                 if prev_timestamp is not None and speedup > 0:
                     delta = timestamp - prev_timestamp
-                    if delta > 0:
-                        sleep_time = delta / speedup
-                        remaining = sleep_time
-                        while remaining > 0 and not stop_event.is_set():
-                            chunk = min(remaining, 5.0)
-                            time.sleep(chunk)
-                            remaining -= chunk
-                        if stop_event.is_set():
-                            break
+                    if delta > 0 and stop_event.wait(delta / speedup):
+                        break
                 prev_timestamp = timestamp
 
                 trace_total_tokens_sum += total_tokens
 
-                with self.counter_lock:
-                    req_id = self.request_counter
-                    self.request_counter += 1
+                req_id = self._next_request_id()
+                if req_id is None:
+                    return
 
                 rand = random.random()
                 if rand < interactive_prob:
@@ -349,7 +356,17 @@ def _generate_requests_from_trace(self, stop_event: threading.Event):
 
     def _generate_requests_from_dataset(self, stop_event: threading.Event):
         """Generates InferenceRequest objects from the loaded ShareGPT dataset."""
+        if (
+            self.replay_cycles > 0
+            and self.sharegpt_loader
+            and self.sharegpt_loader.load_error
+        ):
+            raise RuntimeError("Failed to load ShareGPT dataset") from self.sharegpt_loader.load_error
         if not self.sharegpt_loader or not self.sharegpt_loader.conversations:
+            if self.replay_cycles > 0:
+                logger.warning("ShareGPT dataset is empty or not loaded.")
+                self._mark_producer_done(stop_event)
+                return
             logger.warning("ShareGPT dataset is empty or not loaded. Falling back to synthetic workload.")
             users = UserSimulator.generate_mixed_users(self.num_users)
             self.generate_requests(users, stop_event)
@@ -379,9 +396,9 @@ def _generate_requests_from_dataset(self, stop_event: threading.Event):
             context_tokens = turn['context_tokens']
             generate_tokens = turn['generation_tokens']
 
-            with self.counter_lock:
-                req_id = self.request_counter
-                self.request_counter += 1
+            req_id = self._next_request_id()
+            if req_id is None:
+                return
 
             interactive_prob = cfg('qos_distribution', 'interactive_probability', default=0.15)
             responsive_threshold = cfg('qos_distribution', 'responsive_threshold', default=0.50)
@@ -419,17 +436,36 @@ def _generate_requests_from_dataset(self, stop_event: threading.Event):
             turn_index += 1
 
             if self.request_rate > 0:
-                time.sleep(1.0 / self.request_rate)
+                stop_event.wait(1.0 / self.request_rate)
+
+    def _run_thread(self, role, target, stop_event: threading.Event, args=()):
+        """Run a background target and hard-stop the benchmark if it fails."""
+        try:
+            target(*args)
+        except (Exception, SystemExit) as error:
+            logger.exception("%s failed", role)
+            with self.results_lock:
+                if self._thread_error is None:
+                    self._thread_error = error
+                stop_event.set()
 
     def generate_requests(self, users: List[UserProfile], stop_event: threading.Event):
         """Generate requests concurrently for each simulated user."""
 
+        producer_threads = []
         if self.enable_rag and self.rag_manager and self.rag_ingest_done:
-            threading.Thread(
-                target=self._ingest_rag_documents,
-                args=(self.rag_num_docs, stop_event),
+            rag_thread = threading.Thread(
+                target=self._run_thread,
+                args=(
+                    "Request producer",
+                    self._ingest_rag_documents,
+                    stop_event,
+                    (self.rag_num_docs, stop_event)
+                ),
                 daemon=True
-            ).start()
+            )
+            rag_thread.start()
+            producer_threads.append(rag_thread)
 
         def enqueue_request(request: InferenceRequest):
             priority_tuple = (-QOS_PROFILES[request.qos_level].priority, time.time())
@@ -439,10 +475,7 @@ def user_worker(user: UserProfile):
             """Simulates an individual user generating traffic."""
             local_conv_id = None
 
-            while not stop_event.is_set():
-                time.sleep(user.think_time * random.uniform(0.8, 1.2))
-                if stop_event.is_set():
-                    break
+            while not stop_event.wait(user.think_time * random.uniform(0.8, 1.2)):
 
                 if self.enable_multi_turn and self.conversation_manager:
                     if local_conv_id and random.random() >= 0.8:
@@ -469,9 +502,9 @@ def user_worker(user: UserProfile):
                 new_context = random.randint(max(1, user.context_length // 4), user.context_length)
                 new_gen = random.randint(max(1, user.generation_length // 4), user.generation_length)
 
-                with self.counter_lock:
-                    req_id = self.request_counter
-                    self.request_counter += 1
+                req_id = self._next_request_id()
+                if req_id is None:
+                    return
 
                 if self.enable_multi_turn and self.conversation_manager and local_conv_id:
                     turn_number, cache_key = self.conversation_manager.add_turn(local_conv_id, new_context, new_gen)
@@ -505,9 +538,9 @@ def user_worker(user: UserProfile):
                     retrieved_chunks = self.rag_manager.retrieve_chunks(doc_id)
                     rag_context_tokens = sum(chunk.token_count for chunk in retrieved_chunks)
 
-                    with self.counter_lock:
-                        rag_req_id = self.request_counter
-                        self.request_counter += 1
+                    rag_req_id = self._next_request_id()
+                    if rag_req_id is None:
+                        return
 
                     rag_request = InferenceRequest(
                         user_id=user.user_id,
@@ -522,12 +555,23 @@ def user_worker(user: UserProfile):
                     )
                     enqueue_request(rag_request)
 
-        for user in users:
-            threading.Thread(target=user_worker, args=(user,), daemon=True).start()
-
-        self.active_users = users
+        try:
+            for user in users:
+                user_thread = threading.Thread(
+                    target=self._run_thread,
+                    args=("Request producer", user_worker, stop_event, (user,)),
+                    daemon=True
+                )
+                user_thread.start()
+                producer_threads.append(user_thread)
 
-        stop_event.wait()
+            self.active_users = users
+            stop_event.wait()
+        finally:
+            with self.results_lock:
+                stop_event.set()
+            for thread in producer_threads:
+                thread.join()
 
     def process_requests(self, stop_event: threading.Event):
         """The main worker loop that processes requests from the queue."""
@@ -539,18 +583,31 @@ def process_requests(self, stop_event: threading.Event):
 
             # Check again after dequeue — don't start expensive I/O after stop
             if stop_event.is_set():
+                self.request_queue.put((priority_tuple, request))
                 break
 
             request.start_time = time.perf_counter()
             storage_latency = 0.0
             cache_type = 'user'
+            multi_turn_hits = 0
+            multi_turn_misses = 0
+            prefill_latency = None
+            decode_latency = None
 
             # 1. Check for a prefix cache hit.
             if self.prefix_cache_manager:
-                prefix_entry, remaining_tokens = self.prefix_cache_manager.check_prefix_cache(request, self.model_config)
+                prefix_entry, remaining_tokens = self.prefix_cache_manager.check_prefix_cache(
+                    request,
+                    self.model_config,
+                    record_stats=False
+                )
+                if stop_event.is_set():
+                    return
                 if prefix_entry:
                     cache_type = 'system' if prefix_entry.prefix_type == PrefixType.SYSTEM_PROMPT else 'common'
                     _, read_lat = self.cache.access_cache(prefix_entry.kv_cache_key, request.phase, cache_type)
+                    if stop_event.is_set():
+                        return
                     storage_latency += read_lat
                     request.context_tokens = remaining_tokens
 
@@ -567,19 +624,23 @@ def process_requests(self, stop_event: threading.Event):
                     )
                     for prev_turn_key in prev_keys:
                         location, read_latency = self.cache.access_cache(prev_turn_key, InferencePhase.DECODE, 'multi_turn')
+                        if stop_event.is_set():
+                            return
                         if location is not None:
                             storage_latency += read_latency
-                            with self.results_lock: self.results['multi_turn_cache_hits'] += 1
+                            multi_turn_hits += 1
                         else:
-                            with self.results_lock: self.results['multi_turn_cache_misses'] += 1
+                            multi_turn_misses += 1
 
                 # 3. Perform the main PREFILL operation (a cache WRITE).
                 if request.phase == InferencePhase.PREFILL or request.phase == InferencePhase.PREFILL_DECODE:
                     success, location, write_latency = self.cache.allocate_cache(
                         request.cache_key, request.context_tokens, InferencePhase.PREFILL
                     )
+                    if stop_event.is_set():
+                        return
                     storage_latency += write_latency
-                    with self.results_lock: self.results['prefill_latencies'].append(write_latency)
+                    prefill_latency = write_latency
 
             # 4. Simulate a RAG operation.
             if self.rag_manager and random.random() < cfg('rag', 'request_probability', default=0.1):
@@ -587,8 +648,12 @@ def process_requests(self, stop_event: threading.Event):
                 if doc_keys:
                     doc_id = random.choice(doc_keys)
                     chunks = self.rag_manager.retrieve_chunks(doc_id)
+                    if stop_event.is_set():
+                        return
                     for chunk in chunks:
                         _, read_lat = self.cache.access_cache(chunk.kv_cache_key, InferencePhase.DECODE)
+                        if stop_event.is_set():
+                            return
                         storage_latency += read_lat
 
             # 5. Perform the DECODE operation (a cache READ).
@@ -603,6 +668,8 @@ def process_requests(self, stop_event: threading.Event):
                         decode_key = request.cache_key
                     
                     location, read_latency = self.cache.access_cache(decode_key, InferencePhase.DECODE, cache_type)
+                    if stop_event.is_set():
+                        return
                     storage_latency += read_latency
                     decode_total_latency = read_latency
 
@@ -614,25 +681,34 @@ def process_requests(self, stop_event: threading.Event):
                                 request.context_tokens,
                                 InferencePhase.PREFILL
                             )
+                            if stop_event.is_set():
+                                return
                             storage_latency += write_latency
                     else:
                         decode_batch_size = cfg('decode', 'batch_size', default=32)
                         num_batched_reads = max(1, (request.generate_tokens + decode_batch_size - 1) // decode_batch_size)
                         for _ in range(num_batched_reads):
                             _, batch_read_latency = self.cache.access_cache(decode_key, InferencePhase.DECODE, cache_type)
+                            if stop_event.is_set():
+                                return
                             storage_latency += batch_read_latency
                             decode_total_latency += batch_read_latency
 
-                    with self.results_lock: self.results['decode_latencies'].append(decode_total_latency)
+                    decode_latency = decode_total_latency
 
             # 6. Simulate token generation time.
             generation_latency = request.generate_tokens * GENERATION_TIMING[self.generation_mode]
-            if generation_latency > 0: time.sleep(generation_latency)
+            if generation_latency > 0 and stop_event.wait(generation_latency):
+                return
 
             request.complete_time = time.perf_counter()
 
             # 7. Record all results.
             with self.results_lock:
+                if stop_event.is_set():
+                    return
+                if self.prefix_cache_manager:
+                    self.prefix_cache_manager.record_prefix_result(prefix_entry, self.model_config)
                 self.results['requests_completed'] += 1
                 self.results['total_tokens_generated'] += request.generate_tokens
                 self.results['total_storage_io_latency'] += storage_latency
@@ -640,6 +716,12 @@ def process_requests(self, stop_event: threading.Event):
                 self.results['end_to_end_latencies'].append(request.total_latency_ms / 1000)
                 self.results['storage_latencies'].append(storage_latency)
                 self.results['generation_latencies'].append(generation_latency)
+                self.results['multi_turn_cache_hits'] += multi_turn_hits
+                self.results['multi_turn_cache_misses'] += multi_turn_misses
+                if prefill_latency is not None:
+                    self.results['prefill_latencies'].append(prefill_latency)
+                if decode_latency is not None:
+                    self.results['decode_latencies'].append(decode_latency)
 
                 if self.max_requests > 0 and self.results['requests_completed'] >= self.max_requests:
                     stop_event.set()
@@ -656,8 +738,7 @@ def monitor_stats(self, stop_event: threading.Event):
         start_time = time.time()
         last_log_time = start_time
 
-        while not stop_event.is_set():
-            time.sleep(self.scale_interval)
+        while not stop_event.wait(self.scale_interval):
             now = time.time()
 
             elapsed = now - start_time
@@ -702,7 +783,8 @@ def monitor_stats(self, stop_event: threading.Event):
                     logger.info(f"Autoscaler {action} -> {self.num_users} users (saturation: {saturation_level:.2f})")
                 elif action == 'stop':
                     logger.info("Autoscaler requested stop after reaching capacity peak.")
-                    stop_event.set()
+                    with self.results_lock:
+                        stop_event.set()
                     log_entry = {
                         'timestamp': datetime.now().isoformat(),
                         'mode': self.autoscaler.mode,
@@ -1228,44 +1310,82 @@ def run(self) -> Dict:
         print("-" * 80)
 
         stop_event = threading.Event()
-        self.stop_event = stop_event
-        self.producer_done.clear()
 
         threads = []
         if self.use_dataset:
-            gen_thread = threading.Thread(target=self._generate_requests_from_dataset, args=(stop_event,), daemon=True)
+            producer = self._generate_requests_from_dataset
+            producer_args = (stop_event,)
         elif self.use_burst_trace:
-            gen_thread = threading.Thread(target=self._generate_requests_from_trace, args=(stop_event,), daemon=True)
+            producer = self._generate_requests_from_trace
+            producer_args = (stop_event,)
         else:
-            gen_thread = threading.Thread(target=self.generate_requests, args=(users, stop_event), daemon=True)
-
-        threads.append(gen_thread)
-        gen_thread.start()
-
-        num_workers = min(self.num_users, 500)
-        for _ in range(num_workers):
-            proc_thread = threading.Thread(target=self.process_requests, args=(stop_event,), daemon=True)
-            threads.append(proc_thread)
-            proc_thread.start()
-
-        if self.enable_autoscaling:
-            mon_thread = threading.Thread(target=self.monitor_stats, args=(stop_event,), daemon=True)
-            threads.append(mon_thread)
-            mon_thread.start()
+            producer = self.generate_requests
+            producer_args = (users, stop_event)
+        gen_thread = threading.Thread(
+            target=self._run_thread,
+            args=("Request producer", producer, stop_event, producer_args),
+            daemon=True
+        )
 
         benchmark_start = time.time()
-        stop_event.wait(timeout=self.duration)
-        actual_duration = time.time() - benchmark_start
+        benchmark_deadline = benchmark_start + self.duration
+        try:
+            gen_thread.start()
+            threads.append(gen_thread)
 
-        stop_event.set()
+            num_workers = min(self.num_users, 500)
+            for _ in range(num_workers):
+                if stop_event.is_set():
+                    break
+                if time.time() >= benchmark_deadline:
+                    with self.results_lock:
+                        stop_event.set()
+                    break
+                proc_thread = threading.Thread(
+                    target=self._run_thread,
+                    args=("Request worker", self.process_requests, stop_event, (stop_event,)),
+                    daemon=True
+                )
+                proc_thread.start()
+                threads.append(proc_thread)
+
+            if self.enable_autoscaling and not stop_event.is_set():
+                mon_thread = threading.Thread(
+                    target=self._run_thread,
+                    args=("Autoscaler monitor", self.monitor_stats, stop_event, (stop_event,)),
+                    daemon=True
+                )
+                mon_thread.start()
+                threads.append(mon_thread)
+        except Exception:
+            with self.results_lock:
+                stop_event.set()
+            for thread in threads:
+                thread.join()
+            if self.enable_latency_tracing:
+                self._stop_latency_tracing()
+            if self.io_tracer is not None:
+                self.io_tracer.close()
+            raise
+
+        remaining_duration = max(0.0, benchmark_deadline - time.time())
+        stop_event.wait(timeout=remaining_duration)
+        with self.results_lock:
+            stop_event.set()
+            actual_duration = time.time() - benchmark_start
         for thread in threads:
-            thread.join(timeout=2.0)
+            thread.join()
 
         # Stop tracing and collect results before stats calculation
         trace_data = None
         if self.enable_latency_tracing:
             trace_data = self._stop_latency_tracing()
 
+        if self._thread_error is not None:
+            if self.io_tracer is not None:
+                self.io_tracer.close()
+            raise RuntimeError("Benchmark thread failed") from self._thread_error
+
         self._calculate_stats(actual_duration)
 
         if trace_data:
diff --git a/kv_cache_benchmark/kv_cache/prefix_cache.py b/kv_cache_benchmark/kv_cache/prefix_cache.py
index 24a2792a..369410f3 100755
--- a/kv_cache_benchmark/kv_cache/prefix_cache.py
+++ b/kv_cache_benchmark/kv_cache/prefix_cache.py
@@ -108,7 +108,27 @@ def __init__(self, cache, max_prefix_entries: int = None):
             'bytes_saved': 0
         }
 
-    def check_prefix_cache(self, request: InferenceRequest, model_config: ModelConfig) -> Tuple[Optional[PrefixCacheEntry], int]:
+    def record_prefix_result(
+        self,
+        prefix_entry: Optional[PrefixCacheEntry],
+        model_config: ModelConfig
+    ):
+        """Record prefix-cache metrics for a completed request."""
+        with self.lock:
+            if prefix_entry:
+                self.stats['prefix_hits'] += 1
+                if prefix_entry.prefix_type == PrefixType.SYSTEM_PROMPT:
+                    self.stats['system_prompt_reuse'] += 1
+                self.stats['bytes_saved'] += prefix_entry.token_count * model_config.kv_cache_size_per_token
+            else:
+                self.stats['prefix_misses'] += 1
+
+    def check_prefix_cache(
+        self,
+        request: InferenceRequest,
+        model_config: ModelConfig,
+        record_stats: bool = True
+    ) -> Tuple[Optional[PrefixCacheEntry], int]:
         """
         Checks if the beginning of a request matches a known, cached prefix.
 
@@ -119,15 +139,12 @@ def check_prefix_cache(self, request: InferenceRequest, model_config: ModelConfi
         prefix_entry = self.prefix_matcher.detect_system_prompt(request.context_tokens)
 
         if prefix_entry:
-            with self.lock:
-                self.stats['prefix_hits'] += 1
-                if prefix_entry.prefix_type == PrefixType.SYSTEM_PROMPT:
-                    self.stats['system_prompt_reuse'] += 1
-                self.stats['bytes_saved'] += prefix_entry.token_count * model_config.kv_cache_size_per_token
+            if record_stats:
+                self.record_prefix_result(prefix_entry, model_config)
 
             remaining_tokens = max(0, request.context_tokens - prefix_entry.token_count)
             return prefix_entry, remaining_tokens
         else:
-            with self.lock:
-                self.stats['prefix_misses'] += 1
+            if record_stats:
+                self.record_prefix_result(None, model_config)
             return None, request.context_tokens
diff --git a/kv_cache_benchmark/kv_cache/rag.py b/kv_cache_benchmark/kv_cache/rag.py
index a3b38f77..24d9ac37 100755
--- a/kv_cache_benchmark/kv_cache/rag.py
+++ b/kv_cache_benchmark/kv_cache/rag.py
@@ -91,7 +91,13 @@ def __init__(self, cache, chunk_size: int = None, top_k_chunks: int = None):
             'chunks_retrieved': 0,
         }
 
-    def ingest_document(self, doc_id: str, total_tokens: int, model_config: ModelConfig):
+    def ingest_document(
+        self,
+        doc_id: str,
+        total_tokens: int,
+        model_config: ModelConfig,
+        stop_event: Optional[threading.Event] = None
+    ):
         """
         Simulates the ingestion of a document.
         Splits it into chunks and stores the KV cache for each chunk.
@@ -114,6 +120,8 @@ def ingest_document(self, doc_id: str, total_tokens: int, model_config: ModelCon
         )
 
         for chunk_idx in range(num_chunks):
+            if stop_event and stop_event.is_set():
+                break
             remaining_tokens = total_tokens - chunk_idx * max_tokens_per_chunk
             chunk_tokens = min(max_tokens_per_chunk, remaining_tokens)
 
@@ -137,6 +145,9 @@ def ingest_document(self, doc_id: str, total_tokens: int, model_config: ModelCon
                 logger.error(f"Error ingesting chunk {chunk.chunk_id}: {exc}")
                 continue
 
+            if stop_event and stop_event.is_set():
+                break
+
             if not success:
                 logger.warning(f"Failed to allocate cache for chunk {chunk.chunk_id}.")
                 continue
@@ -145,7 +156,9 @@ def ingest_document(self, doc_id: str, total_tokens: int, model_config: ModelCon
             chunk.size_bytes = chunk_tokens * model_config.kv_cache_size_per_token
 
             doc.chunks.append(chunk)
-            self.chunk_index[chunk.chunk_id] = chunk
+
+        if stop_event and stop_event.is_set():
+            return None
 
         with self.lock:
             # Evict oldest documents if we've hit the limit
@@ -154,6 +167,8 @@ def ingest_document(self, doc_id: str, total_tokens: int, model_config: ModelCon
                     self._evict_oldest_document_unlocked()
 
             self.documents[doc_id] = doc
+            for chunk in doc.chunks:
+                self.chunk_index[chunk.chunk_id] = chunk
             self.ingestion_order.append(doc_id)
             self.stats['documents_ingested'] += 1
             self.stats['chunks_created'] += len(doc.chunks)
diff --git a/kv_cache_benchmark/kv_cache/workload.py b/kv_cache_benchmark/kv_cache/workload.py
index 0c30b415..275ed7d2 100755
--- a/kv_cache_benchmark/kv_cache/workload.py
+++ b/kv_cache_benchmark/kv_cache/workload.py
@@ -312,6 +312,7 @@ def __init__(self, dataset_path: str, max_conversations: int = 1000, seed: Optio
         self.max_conversations = max_conversations
         self.conversations = []
         self.token_stats = {}
+        self.load_error: Optional[Exception] = None
 
         if seed:
             random.seed(seed)
@@ -322,7 +323,8 @@ def __init__(self, dataset_path: str, max_conversations: int = 1000, seed: Optio
     def _load_dataset(self):
         """Load and process the ShareGPT dataset."""
         if not os.path.exists(self.dataset_path):
-            logger.warning(f"Dataset not found at {self.dataset_path}")
+            self.load_error = FileNotFoundError(f"Dataset not found at {self.dataset_path}")
+            logger.warning(str(self.load_error))
             return
 
         try:
@@ -414,6 +416,7 @@ def _load_dataset(self):
         except Exception as e:
             logger.error(f"Error loading dataset: {e}")
             self.conversations = []
+            self.load_error = e
 
     def get_random_conversation(self) -> Optional[Dict]:
         """Get a random conversation from the dataset."""
diff --git a/kv_cache_benchmark/tests/test_kv_cache.py b/kv_cache_benchmark/tests/test_kv_cache.py
index a91e3b7a..de75ba38 100644
--- a/kv_cache_benchmark/tests/test_kv_cache.py
+++ b/kv_cache_benchmark/tests/test_kv_cache.py
@@ -81,6 +81,9 @@
 QoSSLA = kv_cache.QoSSLA
 YAML_AVAILABLE = kv_cache.YAML_AVAILABLE
 IntegratedBenchmark = kv_cache.IntegratedBenchmark
+PrefixCacheEntry = kv_cache.PrefixCacheEntry
+PrefixType = kv_cache.PrefixType
+RAGDocumentManager = kv_cache.RAGDocumentManager
 
 # Input validation imports
 validate_args = kv_cache.validate_args
@@ -2164,6 +2167,19 @@ def benchmark_with_trace(self, trace_dir):
         )
         return bench
 
+    @staticmethod
+    def _sharegpt_benchmark(dataset_path, duration_seconds=5):
+        return IntegratedBenchmark(
+            model_config=MODEL_CONFIGS['tiny-1b'],
+            num_users=2,
+            gpu_memory_gb=0,
+            cpu_memory_gb=0.01,
+            duration_seconds=duration_seconds,
+            dataset_path=str(dataset_path),
+            generation_mode=GenerationMode.NONE,
+            replay_cycles=1,
+        )
+
     def test_resolve_trace_files_from_directory(self, trace_dir):
         """Passing a directory should resolve all CSVs sorted by name."""
         model_config = MODEL_CONFIGS['tiny-1b']
@@ -2236,7 +2252,6 @@ def test_finite_replay_drains_queue(
         monkeypatch.setattr(bench.cache, 'allocate_cache', lambda *args, **kwargs: (True, 'cpu', 0.0))
 
         stop_event = threading.Event()
-        bench.stop_event = stop_event
         producer = threading.Thread(
             target=bench._generate_requests_from_trace,
             args=(stop_event,),
@@ -2259,13 +2274,19 @@ def test_finite_replay_drains_queue(
         assert bench.results['requests_completed'] == expected_requests
         assert bench.request_queue.empty()
 
-    def test_hard_stop_does_not_drain_queue(self, benchmark_with_trace):
+    def test_hard_stop_does_not_drain_queue(self, benchmark_with_trace, monkeypatch):
         """An explicit hard stop should leave queued requests unprocessed."""
         stop_event = threading.Event()
         benchmark_with_trace._generate_requests_from_trace(stop_event)
         queued_requests = benchmark_with_trace.request_queue.qsize()
+        queue_get = benchmark_with_trace.request_queue.get
 
-        stop_event.set()
+        def get_then_stop(*args, **kwargs):
+            item = queue_get(*args, **kwargs)
+            stop_event.set()
+            return item
+
+        monkeypatch.setattr(benchmark_with_trace.request_queue, 'get', get_then_stop)
         worker = threading.Thread(
             target=benchmark_with_trace.process_requests,
             args=(stop_event,),
@@ -2299,6 +2320,162 @@ def test_empty_trace_completes_immediately(self, benchmark_with_trace, monkeypat
         assert benchmark_with_trace.producer_done.is_set()
         assert stop_event.is_set()
 
+    def test_worker_failure_hard_stops_finite_replay(
+        self, benchmark_with_trace, monkeypatch
+    ):
+        """A failed request must not leave a finite replay waiting for its duration."""
+        bench = benchmark_with_trace
+        bench.prefix_cache_manager = None
+        bench.prefill_only = True
+
+        def fail_allocation(*args, **kwargs):
+            raise RuntimeError("allocation failed")
+
+        monkeypatch.setattr(bench.cache, 'allocate_cache', fail_allocation)
+        started = time.perf_counter()
+        with pytest.raises(RuntimeError, match="Benchmark thread failed") as error:
+            bench.run()
+
+        assert time.perf_counter() - started < 5
+        assert isinstance(error.value.__cause__, RuntimeError)
+        assert bench.results['requests_completed'] == 0
+
+    def test_producer_failure_hard_stops_finite_replay(self, benchmark_with_trace):
+        """A failed producer must not leave the benchmark waiting for its duration."""
+        bench = benchmark_with_trace
+        Path(bench.burst_trace_files[0]).unlink()
+        started = time.perf_counter()
+        with pytest.raises(RuntimeError, match="Benchmark thread failed") as error:
+            bench.run()
+
+        assert time.perf_counter() - started < 5
+        assert isinstance(error.value.__cause__, SystemExit)
+        assert not bench.producer_done.is_set()
+
+    def test_in_flight_request_is_cancelled_after_hard_stop(
+        self, benchmark_with_trace, monkeypatch
+    ):
+        """Hard cutoffs discard an admitted request after its blocking I/O returns."""
+        bench = benchmark_with_trace
+        stop_event = threading.Event()
+        started = threading.Event()
+        release = threading.Event()
+        cache_reads = []
+        bench._generate_requests_from_trace(stop_event)
+        prefix_entry = PrefixCacheEntry(
+            prefix_key='system_test',
+            prefix_type=PrefixType.SYSTEM_PROMPT,
+            text_hash='test',
+            token_count=5,
+            kv_cache_key='kv_system_test'
+        )
+
+        def blocking_allocation(*args, **kwargs):
+            started.set()
+            release.wait(timeout=5)
+            return True, 'cpu', 0.0
+
+        def track_cache_read(*args, **kwargs):
+            cache_reads.append(args)
+            return 'cpu', 0.0
+
+        monkeypatch.setattr(
+            bench.prefix_cache_manager.prefix_matcher,
+            'detect_system_prompt',
+            lambda context_tokens: prefix_entry
+        )
+        monkeypatch.setattr(bench.cache, 'allocate_cache', blocking_allocation)
+        monkeypatch.setattr(bench.cache, 'access_cache', track_cache_read)
+        worker = threading.Thread(target=bench.process_requests, args=(stop_event,), daemon=True)
+        worker.start()
+
+        assert started.wait(timeout=2)
+        with bench.results_lock:
+            stop_event.set()
+        release.set()
+        worker.join(timeout=2)
+
+        assert not worker.is_alive()
+        assert bench.results['requests_completed'] == 0
+        assert not bench.results['prefill_latencies']
+        assert bench.request_queue.qsize() == 7
+        assert len(cache_reads) == 1
+        assert not any(bench.prefix_cache_manager.stats.values())
+
+    def test_hard_stop_interrupts_simulated_generation(
+        self, benchmark_with_trace, monkeypatch
+    ):
+        """A duration-style stop should not wait out a long token-generation sleep."""
+        class TrackingEvent(threading.Event):
+            def __init__(self):
+                super().__init__()
+                self.generation_started = threading.Event()
+
+            def wait(self, timeout=None):
+                if timeout and timeout > 10:
+                    self.generation_started.set()
+                return super().wait(timeout)
+
+        bench = benchmark_with_trace
+        bench.generation_mode = GenerationMode.REALISTIC
+        bench.prefix_cache_manager = None
+        bench.prefill_only = True
+        stop_event = TrackingEvent()
+        bench._generate_requests_from_trace(stop_event)
+        monkeypatch.setattr(bench.cache, 'allocate_cache', lambda *args: (True, 'cpu', 0.0))
+        monkeypatch.setitem(GENERATION_TIMING, GenerationMode.REALISTIC, 60.0)
+        worker = threading.Thread(target=bench.process_requests, args=(stop_event,), daemon=True)
+        worker.start()
+
+        assert stop_event.generation_started.wait(timeout=2)
+        stop_event.set()
+        worker.join(timeout=2)
+
+        assert not worker.is_alive()
+        assert bench.results['requests_completed'] == 0
+        assert bench.request_queue.qsize() == 7
+
+    def test_max_requests_caps_trace_production(self, benchmark_with_trace):
+        """A small completion limit must not allow an entire trace to queue first."""
+        bench = benchmark_with_trace
+        bench.max_requests = 1
+        stop_event = threading.Event()
+
+        bench._generate_requests_from_trace(stop_event)
+
+        assert bench.request_counter == 1
+        assert bench.request_queue.qsize() == 1
+        assert not bench.producer_done.is_set()
+        assert not stop_event.is_set()
+
+    def test_hard_stop_interrupts_rag_ingestion(self):
+        """RAG ingestion should stop after the allocation active at cancellation."""
+        stop_event = threading.Event()
+
+        class StoppingCache:
+            def __init__(self):
+                self.allocations = 0
+
+            def allocate_cache(self, **kwargs):
+                self.allocations += 1
+                stop_event.set()
+                return True, 'cpu', 0.0
+
+        cache = StoppingCache()
+        manager = RAGDocumentManager(cache, chunk_size=1)
+
+        manager.ingest_document(
+            'doc',
+            total_tokens=3,
+            model_config=MODEL_CONFIGS['tiny-1b'],
+            stop_event=stop_event
+        )
+
+        assert cache.allocations == 1
+        assert not manager.documents
+        assert not manager.chunk_index
+        assert not any(manager.stats.values())
+
     def test_sharegpt_finite_replay_marks_producer_done(self, tmp_path):
         """ShareGPT finite replay should use graceful completion too."""
         dataset = tmp_path / 'sharegpt.json'
@@ -2307,17 +2484,7 @@ def test_sharegpt_finite_replay_marks_producer_done(self, tmp_path):
             '{"from":"human","value":"first"},{"from":"gpt","value":"reply"},'
             '{"from":"human","value":"second"},{"from":"gpt","value":"reply"}]}]'
         )
-        bench = IntegratedBenchmark(
-            model_config=MODEL_CONFIGS['tiny-1b'],
-            num_users=2,
-            gpu_memory_gb=0,
-            cpu_memory_gb=0.01,
-            duration_seconds=5,
-            dataset_path=str(dataset),
-            max_conversations=1,
-            generation_mode=GenerationMode.NONE,
-            replay_cycles=1,
-        )
+        bench = self._sharegpt_benchmark(dataset)
         stop_event = threading.Event()
 
         bench._generate_requests_from_dataset(stop_event)
@@ -2327,6 +2494,31 @@ def test_sharegpt_finite_replay_marks_producer_done(self, tmp_path):
         assert bench.producer_done.is_set()
         assert not stop_event.is_set()
 
+    def test_empty_sharegpt_finite_replay_completes(self, tmp_path):
+        """An empty finite dataset should finish instead of starting synthetic traffic."""
+        dataset = tmp_path / 'empty-sharegpt.json'
+        dataset.write_text('[]')
+        bench = self._sharegpt_benchmark(dataset)
+        results = bench.run()
+
+        assert bench.request_counter == 0
+        assert bench.request_queue.empty()
+        assert bench.producer_done.is_set()
+        assert results['requests_completed'] == 0
+
+    def test_missing_sharegpt_finite_replay_fails(self, tmp_path):
+        """A missing finite dataset must not look like a successful empty replay."""
+        bench = self._sharegpt_benchmark(
+            tmp_path / 'missing-sharegpt.json', duration_seconds=30
+        )
+
+        started = time.perf_counter()
+        with pytest.raises(RuntimeError, match="Benchmark thread failed"):
+            bench.run()
+
+        assert time.perf_counter() - started < 5
+        assert isinstance(bench.sharegpt_loader.load_error, FileNotFoundError)
+
     def test_total_tokens_tracked(self, benchmark_with_trace):
         """Total tokens from trace should be summed correctly."""
         rows = list(benchmark_with_trace._burst_trace_iterator())
@@ -2352,7 +2544,6 @@ def test_trace_speedup_zero_no_sleep(self, trace_dir):
         )
 
         stop_event = threading.Event()
-        bench.stop_event = stop_event
 
         start = time.time()
         gen_thread = threading.Thread(

From 595857d9e6285404c10a2e6b9a5ace6711d5ec2a Mon Sep 17 00:00:00 2001
From: Joel Tony <github@jaytau.com>
Date: Tue, 4 Aug 2026 22:26:05 +0530
Subject: [PATCH 3/3] fix(kv-cache): bound hard-stop thread joins

---
 kv_cache_benchmark/kv_cache/benchmark.py  |  4 ++--
 kv_cache_benchmark/tests/test_kv_cache.py | 28 +++++++++++++++++++++++
 2 files changed, 30 insertions(+), 2 deletions(-)

diff --git a/kv_cache_benchmark/kv_cache/benchmark.py b/kv_cache_benchmark/kv_cache/benchmark.py
index b208ee81..9d49e5da 100755
--- a/kv_cache_benchmark/kv_cache/benchmark.py
+++ b/kv_cache_benchmark/kv_cache/benchmark.py
@@ -1361,7 +1361,7 @@ def run(self) -> Dict:
             with self.results_lock:
                 stop_event.set()
             for thread in threads:
-                thread.join()
+                thread.join(timeout=2.0)
             if self.enable_latency_tracing:
                 self._stop_latency_tracing()
             if self.io_tracer is not None:
@@ -1374,7 +1374,7 @@ def run(self) -> Dict:
             stop_event.set()
             actual_duration = time.time() - benchmark_start
         for thread in threads:
-            thread.join()
+            thread.join(timeout=2.0)
 
         # Stop tracing and collect results before stats calculation
         trace_data = None
diff --git a/kv_cache_benchmark/tests/test_kv_cache.py b/kv_cache_benchmark/tests/test_kv_cache.py
index de75ba38..00056444 100644
--- a/kv_cache_benchmark/tests/test_kv_cache.py
+++ b/kv_cache_benchmark/tests/test_kv_cache.py
@@ -2402,6 +2402,34 @@ def track_cache_read(*args, **kwargs):
         assert len(cache_reads) == 1
         assert not any(bench.prefix_cache_manager.stats.values())
 
+    def test_duration_hard_stop_does_not_wait_for_blocked_io(
+        self, benchmark_with_trace, monkeypatch
+    ):
+        """A duration cutoff must return while an in-flight storage call is blocked."""
+        bench = benchmark_with_trace
+        bench.duration = 1
+        bench.num_users = 1
+        bench.prefix_cache_manager = None
+        bench.prefill_only = True
+        started = threading.Event()
+        release = threading.Event()
+
+        def blocking_allocation(*args, **kwargs):
+            started.set()
+            release.wait(timeout=10)
+            return True, 'cpu', 0.0
+
+        monkeypatch.setattr(bench.cache, 'allocate_cache', blocking_allocation)
+        run_started = time.perf_counter()
+        try:
+            bench.run()
+        finally:
+            release.set()
+
+        assert started.is_set()
+        assert time.perf_counter() - run_started < 5
+        assert bench.results['requests_completed'] == 0
+
     def test_hard_stop_interrupts_simulated_generation(
         self, benchmark_with_trace, monkeypatch
     ):