huggingface/transformers · #48462
[CB] Add pause mechanism
src/transformers/generation/continuous_batching/continuous_api.py162 + / 37 −
@@ -19,10 +19,12 @@ from abc import abstractmethod from collections.abc import Callable, Generator from contextlib import contextmanager, nullcontext+from datetime import timedelta from time import perf_counter from typing import Any import torch+import torch.distributed as dist from torch import nn from tqdm import tqdm from tqdm.contrib.logging import logging_redirect_tqdm@@ -42,7 +44,7 @@ from .offloading_manager import OffloadingManager from .requests import GenerationOutput, RequestState, RequestStatus, logger from .scheduler import SCHEDULER_MAPPING, FIFOScheduler, Scheduler-from .utils import WorkloadHints, drain_queue+from .utils import ThreadLocalCounter, WorkloadHints, drain_queue """@@ -137,58 +139,77 @@ def fail_and_deliver(self, state: RequestState, error: Exception) -> None: class BackgroundThreadStatus: """Tracks the status of the background thread locally and in its TP group. The status is an int that can only- increase, representing how soon the thread should stop."""+ increase, representing how soon the thread should stop.+ Through this object, threads sharing the model can also pause the generation loop. Any number of threads may ask at+ once, and the loop resumes once they are all done.+ """ + # Stop statuses, in increasing order of urgency. The local status can only ever increase. DONT_STOP = 0 FLUSH_AND_STOP = 1 HARD_STOP = 2 STOPPED = 3 def __init__(self) -> None:- self._local_status_lock = threading.Lock()+ self.fatal_error = None+ self._condition = threading.Condition() self._local_status = self.DONT_STOP self._tp_status = self.DONT_STOP- self.fatal_error = None+ self._pauses_requested = 0+ self._local_pauses_requested = ThreadLocalCounter()+ self._paused = False def clear(self) -> None:- """Clear the local and TP statuses. This method should ONLY be called by the main thread itself BEFORE starting- the background thread."""- self._tp_status = self.DONT_STOP- with self._local_status_lock:+ """Clear the statuses and the pause state. This method should ONLY be called by the main thread itself BEFORE+ starting the background thread."""+ with self._condition:+ self._tp_status = self.DONT_STOP self._local_status = self.DONT_STOP- self.fatal_error = None+ self.fatal_error = None+ self._pauses_requested = 0+ self._local_pauses_requested.value = 0+ self._paused = False+ self._condition.notify_all()++ # ---------------------------------------------- STOP STATUS METHODS --------------------------------------------- # def request_stop(self, status: int, global_rank: int) -> None: """Request the background thread to stop. This does not take effect immediately, only after the TP group has communicated.""" if status not in [self.FLUSH_AND_STOP, self.HARD_STOP]: raise ValueError(f"Invalid stop status {status} from rank {global_rank}")- with self._local_status_lock:+ with self._condition: self._local_status = max(status, self._local_status, self._tp_status) logger.info( f"Rank {global_rank} requested background thread to stop with {status = }. Now {self._local_status = }" ) def record_fatal_error(self, error: Exception) -> None:- """Record a fatal error if none has been recorded yet."""- if self.fatal_error is not None:- logger.error(f"A fatal error was already recorded, ignoring later error: {error}")- else:- self.fatal_error = error+ """Record a fatal error if none has been recorded yet. This is called when the thread crashes and the stop+ status goes to HARD_STOP."""+ with self._condition:+ if self.fatal_error is not None:+ logger.error(f"A fatal error was already recorded, ignoring later error: {error}")+ else:+ self.fatal_error = error+ # If the main thread is waiting for a pause, wake it up since no pause will come (thread just crashed)+ self._condition.notify_all() def mark_as_stopped(self) -> None:- """Mark the background thread as stopped. This should be called by the main thread when the generation loop- finishes."""- with self._local_status_lock:+ """Mark the background thread as stopped. This should be called by the background thread itself when the+ generation loop finishes, on every exit path."""+ with self._condition: self._local_status = self.STOPPED+ # If the main thread is waiting for a pause, wake it up since no pause will come (thread just stopped)+ self._condition.notify_all() def update_with_tp_status(self, tp_status: int) -> None: """Update the local and TP statuses with the new TP status.""" if tp_status < self._tp_status: raise ValueError(f"TP communicated a lower stop status: {tp_status = }, {self._tp_status = }") self._tp_status = tp_status- # We need to use the lock here because main thread might change the local status after the comm- with self._local_status_lock:+ # We need to use the condition's lock here because main thread might change the local status after the comm+ with self._condition: self._local_status = max(self._local_status, tp_status) def can_accept_new_requests(self) -> str | None:@@ -209,6 +230,68 @@ def tp_status(self) -> int: """The status last agreed upon by the TP group through a MAX-reduce operation.""" return self._tp_status + # ------------------------------------------------ PAUSE METHODS ------------------------------------------------ #++ def _pause_predicate(self) -> bool:+ """What a thread waiting for a pause blocks on: the loop paused, or it is never going to."""+ loop_is_done = self._local_status == self.STOPPED or self.fatal_error is not None+ return self._paused or loop_is_done++ def _drop_pause_request(self) -> None:+ """Drop a pause request and notifies threads waiting on the condition. Should be called with the condition held."""+ self._pauses_requested -= 1+ self._local_pauses_requested.value -= 1+ if self._pauses_requested == 0: # the pause ends when all threads have released it, not just the local one+ self._paused = False+ self._condition.notify_all()++ def acquire_pause(self, wake_up_loop: threading.Event) -> None:+ """Called by a thread sharing the model to ask the generation loop to pause. The request is picked up by the+ loop at its next TP all-reduce, see `is_pause_requested`. Any number of threads may ask at the same time. The+ thread then waits until the loop is paused, and raises if the loop is gone before waiting is over."""+ with self._condition:+ # Request the pause+ self._pauses_requested += 1+ self._local_pauses_requested.value += 1+ # This wakes up the loop if it is waiting for new work (new or cancelled requests)+ wake_up_loop.set()++ # Wait for the pause, in a try block so we can handle an error that happens while waiting+ try:+ self._condition.wait_for(self._pause_predicate)+ # If an error happens while waiting, we drop the pause request and re-raise the error+ except BaseException:+ self._drop_pause_request()+ raise++ # If the wait_for ended, two possibilities: the loop paused or it is gone. Latter is fatal.+ if not self._paused:+ self._drop_pause_request()+ raise RuntimeError("The generation loop stopped before it could pause.") from self.fatal_error++ def release_pause(self) -> None:+ """Called by a thread that asked for a pause once it is done with the model. The last one out closes the pause+ window, letting the generation loop resume."""+ with self._condition:+ self._drop_pause_request()++ def is_pause_requested(self, local: bool = False) -> bool:+ """Whether a thread on this rank asked for a pause and has not released it yet. If the local flag is True, only+ check the calling thread's counter."""+ # Local-only check, needs no lock by definition+ if local:+ return self._local_pauses_requested.value > 0+ # Global check, needs to be locked+ with self._condition:+ return self._pauses_requested > 0++ def pause_and_wait(self) -> None:+ """Called by the generation loop to open the pause window and park until the last thread out closes it."""+ with self._condition:+ self._paused = True+ self._condition.notify_all()+ self._condition.wait_for(lambda: not self._paused)+ # Continuous Batch Processor (Internal Logic) class ContinuousBatchProcessor:@@ -338,19 +421,36 @@ def _update_tp_group_state(self) -> bool: payload = (drain_queue(self.input_queue), drain_queue(self.cancel_queue)) else: payload = ([], [])- # And the size of the payload is inferred (always 0 for non-TP drivers)- payload_size = len(payload[0]) + len(payload[1]) - # Cheap 2 ints broadcast of payload size (from rank 0) and requested stop status (all to all)- local_requested_status = self.background_thread_status.local_status- payload_size, tp_status = self.distributed_helper.tp_all_reduce_state(payload_size, local_requested_status)+ # Cheap 3 ints broadcast of payload size (from rank 0), requested stop status and requested pause (all to all)+ payload_size, tp_status, pause_requested = self.distributed_helper.tp_all_reduce_state(+ payload_size=len(payload[0]) + len(payload[1]), # always 0 for non-TP drivers+ stop_status=self.background_thread_status.local_status,+ pause_requested=self.background_thread_status.is_pause_requested(),+ )+ # Update the local stop status with the new one self.background_thread_status.update_with_tp_status(tp_status)- # Exit early if the TP group is hard-stopping if self.background_thread_status.tp_status == BackgroundThreadStatus.HARD_STOP: return True- # Same if there is no payload++ # Pause window: if any thread of the TP group asked for a pause, all ranks wait here until their own threads+ # are done pausing. Thanks to the barrier (one per rank), all ranks leave the pause at the same time.+ if pause_requested:+ # Wait for the work on the compute stream to be done before pausing+ if self.inputs_and_outputs.compute_stream is not None:+ self.inputs_and_outputs.compute_stream.synchronize()+ self.background_thread_status.pause_and_wait()+ if self.distributed_helper.cpu_comm_group is not None:+ timeout = self.cb_config.cpu_group_timeout+ dist.monitored_barrier( # ty: ignore[possibly-missing-attribute]+ group=self.distributed_helper.cpu_comm_group,+ timeout=timedelta(seconds=timeout) if timeout is not None else None,+ wait_all_ranks=True,+ )++ # After the pause window is done, we can still exit early if there is no payload if payload_size == 0: return False # Otherwise, distribute the payload of TP rank 0 to all other TP ranks@@ -600,7 +700,7 @@ def __init__( self.cancel_queue: queue.Queue[str] = queue.Queue() self._request_counter = 0 self._request_lock = threading.Lock()- self._has_new_requests = threading.Event()+ self._wake_up_loop = threading.Event() # Processor-related attributes self.background_thread_status = BackgroundThreadStatus()@@ -729,6 +829,14 @@ def stop( logger.warning(msg) return None + # Stopping and pausing are conflicting operations: a thread inside a pause cannot stop the manager, because that+ # would deadlock (pause waits for the stop to complete, stop hangs because the loop is paused)+ if self.background_thread_status.is_pause_requested(local=True):+ raise RuntimeError(+ "Cannot stop the manager from inside a pause: the generation loop is paused and cannot exit, so "+ "this would wait forever. Leave the `pause` context before calling `stop`."+ )+ # Signal the background thread to stop stop_trigger_time = perf_counter() stop_status = BackgroundThreadStatus.HARD_STOP if hard_stop else BackgroundThreadStatus.FLUSH_AND_STOP@@ -777,6 +885,24 @@ def destroy(self) -> None: self.stop(block=True, keep_for_next_session=False) self.distributed_helper.destroy_cpu_comm_group() + @contextmanager+ def pause(self):+ """A context manager that pauses the generation loop, so the calling thread can use the model, typically to+ update it in place. The thread only enters this context once the loop is paused, and the loop resumes on exit,+ keeping its cache and its in-flight requests: nothing is drained and no request is lost.+ Several threads may hold the pause at the same time, and the loop resumes once the last one leaves.+ If TP is on, all ranks must enter this context, otherwise other ranks will hang forever.+ """+ # Error out if the caller asks for a pause while no generation loop is running+ if not self.is_running():+ raise RuntimeError("Cannot pause generation while no generation loop is running.")++ self.background_thread_status.acquire_pause(self._wake_up_loop)+ try:+ yield+ finally:+ self.background_thread_status.release_pause()+ # ---------------------------- REQUEST SUBMISSION, CANCELLATION AND RETRIEVAL METHODS ---------------------------- # def add_request(@@ -834,7 +960,7 @@ def add_request( # Use block=True with timeout to handle backpressure if queue is full self.input_queue.put(state, block=True, timeout=10)- self._has_new_requests.set()+ self._wake_up_loop.set() return request_id def add_requests(@@ -877,7 +1003,7 @@ def cancel_request(self, request_id: str) -> None: driver processes interact with the manager.""" if self.is_tp_driver: self.cancel_queue.put(request_id)- self._has_new_requests.set()+ self._wake_up_loop.set() # TODO (remi-or) : handle benchmarking properly when updating / fixing the requeue logic # TODO (remi-or) : this NEEDS to get fixed in a future PR -- it's quite wasteful@@ -967,11 +1093,10 @@ def _generation_loop_body(self, batch_processor: ContinuousBatchProcessor, boots return False # Otherwise, we wait for new requests and retry else:- self._has_new_requests.wait(timeout=0.1) # wait for new requests instead of busy-spinning.- self._has_new_requests.clear()+ self._wake_up_loop.wait(timeout=0.1) # wait for new requests instead of busy-spinning.+ self._wake_up_loop.clear() return True - @torch.no_grad() def _run_generation_loop(self) -> None: """Main processing loop running in the background thread.""" batch_processor = None@@ -1086,17 +1211,17 @@ def _create_batch_processor(self) -> ContinuousBatchProcessor: def _handle_critical_error(self, error: Exception, batch_processor: ContinuousBatchProcessor | None) -> None: """Handle critical errors that terminate the generation loop."""- # Request a hard stop+ # Request a hard stop., only on this rank. Other ranks will be notified at the comm self.background_thread_status.request_stop( status=BackgroundThreadStatus.HARD_STOP, global_rank=self.distributed_helper.global_rank )- # Communicate to other processes in the TP group that the group is stopping (they could have not crashed)- # Since the other processes need to reach the collective, it may take a few seconds to complete.- self.distributed_helper.tp_all_reduce_state(0, BackgroundThreadStatus.HARD_STOP) # Fail all remaining requests self._fail_all_remaining_requests(error, batch_processor) # After failing the remaining requests (and so retrieving their partial outputs), record the fatal error self.background_thread_status.record_fatal_error(error)+ # Communicate to other ranks in the TP group that the group is stopping (they could have not crashed)+ # Since the other processes need to reach the collective, it may take a few seconds to complete.+ self.distributed_helper.tp_all_reduce_state(0, BackgroundThreadStatus.HARD_STOP) def _fail_all_remaining_requests(self, error: Exception, batch_processor: ContinuousBatchProcessor | None) -> None: """Fail all remaining requests in the input queue and active requests."""src/transformers/generation/continuous_batching/distributed.py10 + / 6 −
@@ -87,7 +87,7 @@ def __init__(self, device_mesh: DeviceMesh | None, cpu_group_timeout: float | No self.dp_size = self.world_size // self.tp_size # Accumulator to CPU integer comm- self._cpu_int_acc = torch.tensor([0, 0], dtype=torch.int64, device="cpu")+ self._cpu_int_acc = torch.tensor([0, 0, 0], dtype=torch.int64, device="cpu") @staticmethod def check_device_mesh_for_cb(device_mesh: DeviceMesh | None) -> None:@@ -130,15 +130,19 @@ def tp_broadcast_from_rank_0(self, value: torch.Tensor) -> torch.Tensor: dist.broadcast(value, src=self.tp_root_global_rank, async_op=False, group=self.tp_group) return value - def tp_all_reduce_state(self, payload_size: int, stop_status: int) -> tuple[int, int]:- """Broadcasts two information: 1. the size of the payload held by the TP driver (all other rank broadcast 0) and- 2. the requested stop status (all to all). These information are broadcasted through a MAX-reduce operation."""+ def tp_all_reduce_state(+ self, payload_size: int, stop_status: int, pause_requested: int = 0+ ) -> tuple[int, int, int]:+ """Broadcasts three information: 1. the size of the payload held by the TP driver (all other rank broadcast 0),+ 2. the requested stop status (all to all) and 3. whether a thread outside the generation loop is asking the+ loop to pause (all to all). These information are broadcasted through a MAX-reduce operation.""" if self.tp_size > 1: self._cpu_int_acc[0] = payload_size self._cpu_int_acc[1] = stop_status+ self._cpu_int_acc[2] = pause_requested dist.all_reduce(self._cpu_int_acc, op=dist.ReduceOp.MAX, async_op=False, group=self.cpu_comm_group)- payload_size, stop_status = self._cpu_int_acc.tolist()- return payload_size, stop_status+ payload_size, stop_status, pause_requested = self._cpu_int_acc.tolist()+ return payload_size, stop_status, pause_requested def tp_all_reduce_min(self, value: torch.Tensor, on_cpu: bool = False) -> torch.Tensor: """Inside each TP group, all-reduces a tensor with the MIN op. No-op when TP is off. If the tensor is on CPU,src/transformers/generation/continuous_batching/utils.py6 + / 0 −
@@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. import queue+import threading from contextlib import contextmanager from dataclasses import dataclass from math import ceil, log2@@ -217,3 +218,8 @@ def mem_pool_ctx(mem_pool): """A context manager to use a CUDA mem pool.""" with torch.cuda.use_mem_pool(mem_pool): yield+++class ThreadLocalCounter(threading.local):+ def __init__(self) -> None:+ self.value = 0tests/generation/test_continuous_batching.py505 + / 1 −
@@ -16,6 +16,8 @@ import gc import itertools import os+import threading+import time import unittest from types import SimpleNamespace from typing import Any@@ -41,7 +43,11 @@ group_layers_by_attn_type, ) from transformers.generation.continuous_batching.cache_manager import FullAttentionCacheAllocator-from transformers.generation.continuous_batching.continuous_api import OutputRouter+from transformers.generation.continuous_batching.continuous_api import (+ BackgroundThreadStatus,+ ContinuousBatchingManager,+ OutputRouter,+) from transformers.generation.continuous_batching.distributed import DistributedHelper from transformers.generation.continuous_batching.input_outputs import build_attention_mask from transformers.generation.continuous_batching.offloading_manager import OffloadingManager@@ -787,6 +793,234 @@ def test_continuous_batching_config_disables_nccl_graph_mixing(self) -> None: os.environ["WORLD_SIZE"] = original_ws +class ContinuousBatchingPauseTest(unittest.TestCase):+ """Tests for the pause API in states a real single-rank generation loop cannot reach. The generation loop is+ replaced by a plain thread driving `BackgroundThreadStatus` directly, so no model or accelerator is needed. Pausing+ a real generation loop is tested in `ContinuousBatchingWithAcceleratorTest`."""++ @staticmethod+ def _get_minimal_manager(status: BackgroundThreadStatus, thread: threading.Thread | None):+ """Builds a bare manager with just enough states for `pause` and `stop` to work."""+ manager = ContinuousBatchingManager.__new__(ContinuousBatchingManager)+ manager._wake_up_loop = threading.Event()+ manager.background_thread_status = status+ manager._generation_thread = thread+ manager.distributed_helper = SimpleNamespace(cpu_comm_group=None, global_rank=0)+ # Attributes torn down by `stop`+ manager.batch_processor = None+ manager._original_attn_impl = None+ return manager++ def test_pause_on_a_peer_initiated_pause(self) -> None:+ """Under TP the pause request is MAX-reduced across ranks, so a loop can be paused before a local thread asked+ for it. A thread arriving on an already paused loop must enter the pause directly."""+ status = BackgroundThreadStatus()+ resumed = threading.Event()++ def generation_loop() -> None:+ """Simulates a loop that was paused by a thread on another rank."""+ try:+ self.assertFalse(status.is_pause_requested()) # should be false because this rank did not pause+ status.pause_and_wait() # enters the pause on behalf of the other rank+ resumed.set() # signals that the pause ended+ finally:+ status.mark_as_stopped()++ # Setup the loop thread and the manager+ loop_thread = threading.Thread(target=generation_loop)+ loop_thread.start()+ manager = self._get_minimal_manager(status, loop_thread)++ # Enter an already paused loop and check that the loop did not resume yet+ entered = []+ with manager.pause():+ entered.append(True)+ self.assertFalse(resumed.is_set(), "the loop resumed while the pause was held")+ loop_thread.join(timeout=3) # should leave enough for the loop to resume++ self.assertEqual(entered, [True], "the thread did not enter the already paused loop")+ self.assertTrue(resumed.is_set(), "the loop did not resume after the pause was released")++ def test_pause_request_is_dropped_when_the_wait_is_interrupted(self) -> None:+ """If an exception (e.g. KeyboardInterrupt) is raised while waiting for the loop to pause, the pause request+ must be dropped. Otherwise `is_pause_requested` stays true and the loop pauses forever with nobody to release+ it. The exception is injected through the wait predicate, so it is raised inside `wait_for` with the lock held,+ like a real interrupt would be."""+ status = BackgroundThreadStatus()+ loop_resumed = threading.Event()++ def generation_loop() -> None:+ try:+ status.pause_and_wait()+ loop_resumed.set()+ finally:+ status.mark_as_stopped()++ loop_thread = threading.Thread(target=generation_loop)+ loop_thread.start()+ manager = self._get_minimal_manager(status, loop_thread)+ # Wait for the loop to pause, so that the interrupted request is the only one keeping it paused+ deadline = time.monotonic() + 3+ while not status.is_pause_requested() and time.monotonic() < deadline:+ time.sleep(0.01)++ def interrupted_predicate() -> bool:+ raise KeyboardInterrupt("interrupted while waiting for the pause")++ # Replace the pause predicate with a function that raises an exception, so that it is raised while waiting+ status._pause_predicate = interrupted_predicate+ with self.assertRaises(KeyboardInterrupt):+ with manager.pause():+ pass++ self.assertFalse(status.is_pause_requested(), "the interrupted pause request was not withdrawn")+ self.assertFalse(status.is_pause_requested(local=True), "the local pause counter was not decremented")+ loop_thread.join(timeout=3)+ self.assertTrue(loop_resumed.is_set(), "the loop stayed paused after the interrupted request")++ def test_stop_from_inside_a_pause_raises(self) -> None:+ """`stop` waits for the generation loop to exit, and the loop cannot exit while a pause is held. So a thread+ calling `stop` from inside its own pause would wait forever: it must raise instead, whatever `block` and+ `timeout` are."""+ status = BackgroundThreadStatus()+ loop_thread = threading.Thread(target=status.pause_and_wait)+ loop_thread.start()+ manager = self._get_minimal_manager(status, loop_thread)++ try:+ with manager.pause():+ for kwargs in [{}, {"block": False}, {"block": True, "timeout": 5.0}]:+ with self.assertRaises(RuntimeError) as context:+ manager.stop(**kwargs)+ self.assertIn("from inside a pause", str(context.exception), f"wrong error for {kwargs}")+ finally:+ loop_thread.join(timeout=10)+ self.assertFalse(loop_thread.is_alive())++ def test_stop_from_another_thread_while_paused(self) -> None:+ """A thread that does not hold a pause can call `stop` while another thread does: the stop waits for the pause+ to be released and the loop to exit. Only a blocking stop is tested, since a non-blocking stop tears down the+ manager while the loop is still running, which is a pre-existing issue of `stop` unrelated to pausing."""+ status = BackgroundThreadStatus()+ loop_thread = threading.Thread(target=status.pause_and_wait)+ loop_thread.start()+ manager = self._get_minimal_manager(status, loop_thread)++ holding = threading.Event()+ release = threading.Event()++ def holder() -> None:+ with manager.pause():+ holding.set()+ release.wait(timeout=10)++ holder_thread = threading.Thread(target=holder)+ holder_thread.start()+ try:+ self.assertTrue(holding.wait(timeout=10), "the holder did not get its pause")+ # A blocking stop should wait for the holder to release the pause+ stopped = threading.Event()+ stopper = threading.Thread(target=lambda: (manager.stop(), stopped.set()))+ stopper.start()+ self.assertFalse(stopped.wait(timeout=0.5), "stop returned while the pause was still held")+ release.set()+ stopper.join(timeout=10)+ self.assertTrue(stopped.is_set(), "stop did not complete after the pause was released")+ finally:+ release.set()+ holder_thread.join(timeout=10)+ loop_thread.join(timeout=10)+ self.assertFalse(loop_thread.is_alive())++ @staticmethod+ def _pause_in_thread(manager, timeout: float = 10) -> tuple[bool, Exception | None]:+ """Enters and exits `pause` in a separate thread, so a hang shows up as a timeout instead of blocking the test+ suite. Returns whether the thread finished and the exception it raised, if any."""+ raised: list[Exception] = []++ def enter_and_leave() -> None:+ try:+ with manager.pause():+ pass+ except Exception as e:+ raised.append(e)++ thread = threading.Thread(target=enter_and_leave)+ thread.start()+ thread.join(timeout=timeout)+ finished = not thread.is_alive()+ return finished, (raised[0] if raised else None)++ def test_pause_raises_when_the_loop_cannot_pause(self) -> None:+ """When the loop will never pause, `pause` must raise instead of blocking. This covers a loop that died on a+ fatal error, a loop that hard-stopped right after seeing the request, a manager that was never started and a+ manager that was already stopped."""+ # 1. The loop dies with a fatal error while a pause is pending: the error is chained as `__cause__`+ fatal_error = RuntimeError("boom in the forward pass")+ status = BackgroundThreadStatus()+ keep_unwinding = threading.Event()++ def dying_loop() -> None:+ try:+ while not status.is_pause_requested():+ time.sleep(0.001)+ # Same order as `_handle_critical_error`: request the hard stop, then record the error+ status.request_stop(BackgroundThreadStatus.HARD_STOP, 0)+ status.record_fatal_error(fatal_error)+ keep_unwinding.wait(timeout=10) # keep the thread alive while `pause` gives up+ finally:+ status.mark_as_stopped()++ loop_thread = threading.Thread(target=dying_loop)+ loop_thread.start()+ manager = self._get_minimal_manager(status, loop_thread)+ finished, error = self._pause_in_thread(manager)+ self.assertTrue(finished, "pause blocked after the generation loop died")+ self.assertIsInstance(error, RuntimeError)+ self.assertIs(error.__cause__, fatal_error)+ # The manager should be stopped rather than left half-dead+ self.assertEqual(status.local_status, BackgroundThreadStatus.HARD_STOP)+ keep_unwinding.set()+ loop_thread.join(timeout=10)++ # 2. The loop hard-stops with a pause pending. `_update_tp_group_state` returns on HARD_STOP before reaching the+ # pause window, so the loop sees the request but never pauses. There is no fatal error, but `pause` must still+ # return.+ status = BackgroundThreadStatus()++ def hard_stopping_loop() -> None:+ try:+ while not status.is_pause_requested():+ time.sleep(0.001)+ finally:+ status.mark_as_stopped()++ loop_thread = threading.Thread(target=hard_stopping_loop)+ loop_thread.start()+ manager = self._get_minimal_manager(status, loop_thread)+ finished, error = self._pause_in_thread(manager)+ loop_thread.join(timeout=10)+ self.assertTrue(finished, "pause blocked after the generation loop hard-stopped")+ self.assertIsInstance(error, RuntimeError)+ self.assertGreaterEqual(status.local_status, BackgroundThreadStatus.HARD_STOP)++ # 3. A manager that was never started+ manager = self._get_minimal_manager(BackgroundThreadStatus(), None)+ finished, error = self._pause_in_thread(manager)+ self.assertTrue(finished, "pause blocked on a manager that was never started")+ self.assertIsInstance(error, RuntimeError)++ # 4. A manager whose generation thread has already exited+ status = BackgroundThreadStatus()+ loop_thread = threading.Thread(target=status.mark_as_stopped)+ loop_thread.start()+ loop_thread.join(timeout=10)+ manager = self._get_minimal_manager(status, loop_thread)+ finished, error = self._pause_in_thread(manager)+ self.assertTrue(finished, "pause blocked on an already stopped manager")+ self.assertIsInstance(error, RuntimeError)++ @require_torch_accelerator class ContinuousBatchingWithAcceleratorTest(unittest.TestCase): # -----------------------------------------------Parity tests----------------------------------------------- #@@ -1410,6 +1644,134 @@ def _test_block_sharing(self, model_id: str, expected_layer_types: dict[str, int # Verify parity with regular generate self.assertEqual(chunk_no_reuse.generated_tokens, expected_output_tokens[0]) + # ------------------------------------------------- Pause tests ------------------------------------------------- #++ def _started_manager_with_requests(self, max_new_tokens: int = 200):+ """Returns a started manager with a few requests in flight, along with the model and the number of requests."""+ tokenizer, model = get_tokenizer_and_model("TinyLlama/TinyLlama-1.1B-Chat-v1.0", "paged|sdpa", torch_device)+ input_ids = get_generation_inputs(_DEFAULT_USER_MESSAGES, tokenizer, for_continuous_batching=True)+ cb_config = ContinuousBatchingConfig(use_cuda_graph=False, use_async_batching=False)+ manager = model.init_continuous_batching(continuous_batching_config=cb_config)+ manager.logit_processor.clear()+ manager.warmup() # so that no graph capture happens on the loop thread during the test+ manager.start()+ for ids in input_ids:+ manager.add_request(ids, max_new_tokens=max_new_tokens)+ return model, manager, len(input_ids)++ def _wait_until_generating(self, manager, timeout: float = 60) -> None:+ """Blocks until the generation loop has run at least one step. Requests are admitted to the scheduler after the+ pause window in `_update_tp_group_state`, so a pause taken on the first iteration would pause an empty scheduler+ and the step counter checks would be meaningless."""+ # `current_batch` is only set once the loop thread has started+ deadline = time.monotonic() + timeout+ while getattr(manager, "current_batch", 0) == 0 and time.monotonic() < deadline:+ time.sleep(0.01)+ self.assertGreater(getattr(manager, "current_batch", 0), 0, "the generation loop never started generating")++ @staticmethod+ def _collect_finished(manager, expected: int, timeout: float = 120):+ """Drains the manager's outputs until `expected` requests have finished or the timeout expires."""+ outputs = {}+ deadline = time.monotonic() + timeout+ while len(outputs) < expected and time.monotonic() < deadline:+ result = manager.get_result(timeout=1.0)+ if result is not None and result.is_finished():+ outputs[result.request_id] = result+ return outputs++ @with_flush_memory+ def test_pause_during_generation(self) -> None:+ """Pauses the manager several times while requests are in flight, like a trainer sharing the model would. The+ generation loop must not run a step while the pause is held, and all requests must still complete."""+ num_pauses = 5+ model, manager, num_requests = self._started_manager_with_requests()+ try:+ self._wait_until_generating(manager)+ for i in range(num_pauses):+ with manager.pause():+ batches_at_pause = manager.current_batch+ # The step counter check below only means something if the loop still has work to do+ self.assertTrue(+ manager.batch_processor.has_pending_requests(),+ f"generation finished before pause {i}",+ )+ # Stands in for a trainer updating the weights in place+ with torch.no_grad():+ for parameter in model.parameters():+ parameter.add_(0.0)+ # A generation step takes a few ms, so hold the pause long enough for a loop that ignored it to+ # run several steps+ time.sleep(0.3)+ self.assertEqual(manager.current_batch, batches_at_pause, f"the loop ran during pause {i}")+ outputs = self._collect_finished(manager, num_requests)+ finally:+ manager.stop(block=True)++ self.assertEqual(len(outputs), num_requests, "some requests did not complete")+ for request_id, output in outputs.items():+ self.assertGreater(len(output.generated_tokens), 0, f"{request_id} generated no tokens")++ @with_flush_memory+ def test_pause_with_several_holders(self) -> None:+ """Several threads can hold the pause at once, and the loop must stay paused until the last one releases it.+ The holders are released one at a time and the step counter is checked after each release."""+ num_holders = 3+ model, manager, num_requests = self._started_manager_with_requests()+ # The main thread is the extra party: it runs its checks once every holder is inside its pause+ all_inside = threading.Barrier(num_holders + 1, timeout=60)+ may_leave = [threading.Event() for _ in range(num_holders)]+ has_left = [threading.Event() for _ in range(num_holders)]+ holder_errors = []++ def holder(index: int) -> None:+ try:+ with manager.pause():+ all_inside.wait()+ may_leave[index].wait(timeout=60)+ has_left[index].set()+ except Exception as error: # a holder that fails to enter the pause would break the barrier+ holder_errors.append(error)++ threads = [threading.Thread(target=holder, args=(index,)) for index in range(num_holders)]+ try:+ self._wait_until_generating(manager)+ for thread in threads:+ thread.start()+ all_inside.wait()+ batches_at_pause = manager.current_batch+ # The step counter checks below only mean something if the loop still has work to do+ self.assertTrue(+ manager.batch_processor.has_pending_requests(),+ "generation finished before the holders entered their pause",+ )++ for index in range(num_holders):+ may_leave[index].set()+ self.assertTrue(has_left[index].wait(timeout=30), f"holder {index} did not exit its pause")+ # Give a loop that resumed too early time to run a few steps+ time.sleep(0.3)+ if index < num_holders - 1:+ self.assertEqual(+ manager.current_batch,+ batches_at_pause,+ f"the loop resumed with {num_holders - index - 1} holders still inside",+ )++ for thread in threads:+ thread.join(timeout=30)+ outputs = self._collect_finished(manager, num_requests)+ finally:+ for event in may_leave:+ event.set()+ manager.stop(block=True)++ self.assertEqual(holder_errors, [], f"some holders failed to enter the pause: {holder_errors}")+ self.assertGreater(+ manager.current_batch, batches_at_pause, "the loop did not resume after the last holder left"+ )+ self.assertEqual(len(outputs), num_requests, "some requests did not complete")+ def test_prefix_sharing(self) -> None: model_id = "Qwen/Qwen2.5-0.5B-Instruct" num_layer_groups = {"full_attention": 1, "sliding_window": 0}@@ -2065,6 +2427,137 @@ def _tp_cancellation_worker( manager.stop(block=True) +def _tp_pause_generation_worker(+ rank: int,+ model_id: str,+ attn_implementation: str,+ max_new_tokens: int,+ num_pauses: int,+ skew_seconds: float,+ use_async_batching: bool = False,+ use_cuda_graph: bool = False,+) -> None:+ """Loads `model_id` with `DistributedConfig(tp_size=...)` and runs a generation loop on every rank, plus a second+ thread that repeatedly pauses it. Each rank starts pausing at a different time, so the ranks reach the pause window+ at different iterations. Since the pause request is MAX-reduced across ranks, a rank ends up paused before its own+ thread asked for it, which a single-rank test cannot reproduce.++ Inside the pause, every rank issues a collective on the TP group. If a rank resumed generation instead of staying+ paused, the collectives of its forward pass would pair up with this probe on the other ranks, and the probe would+ return a wrong value or hang. Rank 0 owns the assertions."""+ import threading+ import time++ import torch+ import torch.distributed as dist++ from transformers.distributed import DistributedConfig++ tokenizer = AutoTokenizer.from_pretrained(model_id, padding_side="left")+ if not hasattr(tokenizer, "pad_token") and hasattr(tokenizer, "eos_token"):+ tokenizer.pad_token = tokenizer.eos_token++ model = AutoModelForCausalLM.from_pretrained(+ model_id,+ attn_implementation=attn_implementation,+ distributed_config=DistributedConfig(tp_size=int(os.environ["WORLD_SIZE"])),+ dtype=torch.float32,+ ).eval()++ chats = [[{"role": "user", "content": message}] for message in _DEFAULT_USER_MESSAGES]+ tokenized = [tokenizer.apply_chat_template(chat, add_generation_prompt=True) for chat in chats]+ input_ids = [(x if isinstance(x, list) else x["input_ids"]) for x in tokenized]++ # A dedicated group for the test's own synchronization: the manager's groups are used by the generation loop+ # thread, and issuing collectives on them from the main thread at the same time would interleave.+ sync_group = dist.new_group(backend="gloo")++ cb_config = ContinuousBatchingConfig(+ use_cuda_graph=use_cuda_graph, use_async_batching=use_async_batching, cpu_group_timeout=120+ )+ manager = model.init_continuous_batching(continuous_batching_config=cb_config)+ manager.logit_processor.clear()+ # Warm up synchronously so that no CUDA-graph capture happens on the loop thread during the test+ manager.warmup()+ manager.start()++ tp_group = manager.distributed_helper.tp_group+ tp_size = manager.distributed_helper.tp_size+ expected_probe = float(sum(range(1, tp_size + 1)))+ pauses_taken = []+ pause_errors = []++ def pauser() -> None:+ """Stands in for a trainer thread updating the model between generation steps."""+ try:+ time.sleep(skew_seconds * rank) # so that the ranks do not ask for the pause at the same iteration+ for _ in range(num_pauses):+ with manager.pause():+ # A collective on the TP group, like a training step would issue. It only returns the right value+ # if every rank is in the same pause window.+ probe = torch.full((16,), float(rank + 1), device=model.device)+ dist.all_reduce(probe, group=tp_group)+ if probe[0].item() != expected_probe:+ raise AssertionError(+ f"TP collective inside the pause returned {probe[0].item()}, expected {expected_probe}: "+ f"the ranks are not in the same pause window"+ )+ pauses_taken.append(1)+ time.sleep(0.02) # let the generation loop progress between pauses+ except Exception as error: # re-raised on the main thread below+ pause_errors.append(error)++ pause_thread = threading.Thread(target=pauser)+ pause_thread.start()++ try:+ # `add_request` only enqueues on the TP driver and returns None on the other ranks. The requests are then+ # broadcast, so every rank produces outputs for the driver's request ids: collect by count and key on the ids+ # that come back.+ for ids in input_ids:+ manager.add_request(ids, max_new_tokens=max_new_tokens, streaming=False)+ outputs = {}+ deadline = time.time() + 300+ while len(outputs) < len(input_ids) and time.time() < deadline:+ result = manager.get_result(timeout=1.0)+ if result is not None and result.is_finished():+ outputs[result.request_id] = result+ pause_thread.join(timeout=120)+ finally:+ # The pause thread must be done before stopping: `stop` joins the generation thread, which cannot exit while a+ # pause is held. Also, no rank may stop while another rank is still pausing, since the stop goes through the+ # same all-reduce as the pause and the late rank would be left waiting on a loop that no longer exists.+ if not pause_thread.is_alive():+ dist.barrier(group=sync_group)+ manager.stop(block=True)++ assert not pause_thread.is_alive(), f"rank {rank}: the pausing thread did not finish"+ if pause_errors:+ # `is_running()` does not distinguish a dead loop from a stopped one, so chain the fatal error if there is one+ loop_error = manager.background_thread_status.fatal_error+ raise pause_errors[0] from loop_error+ assert len(pauses_taken) == num_pauses, f"rank {rank}: took {len(pauses_taken)} pauses, expected {num_pauses}"+ assert len(outputs) == len(input_ids), f"rank {rank}: got {len(outputs)}/{len(input_ids)} requests back"++ # Pausing must not change the generated tokens nor make the ranks diverge. Sorting by request id gives every rank+ # the same order.+ local_tokens = [outputs[request_id].generated_tokens for request_id in sorted(outputs)]+ gathered_tokens = [None] * tp_size+ dist.all_gather_object(gathered_tokens, local_tokens, group=tp_group)++ if rank != 0:+ return++ for index, tokens in enumerate(local_tokens):+ assert len(tokens) > 0, f"Request {index} got no generated tokens"+ for src_rank, src_tokens in enumerate(gathered_tokens):+ if src_tokens != gathered_tokens[0]:+ raise AssertionError(+ f"Generation diverges across ranks when pausing: rank {src_rank} got {src_tokens}, rank 0 got "+ f"{gathered_tokens[0]}"+ )++ @require_torch_multi_accelerator class ContinuousBatchingTensorParallelTest(unittest.TestCase): """Integration tests for continuous batching with tensor parallelism. Each test spawns a TP-sized process group@@ -2093,6 +2586,17 @@ def test_continuous_batching_tp_fast(self) -> None: that all TP ranks agree on the generated tokens.""" self._run_cb_worker(max_new_tokens=4) + def test_continuous_batching_tp_pause(self) -> None:+ """Test that `pause` keeps the TP ranks in the same pause window even when they request it at different+ iterations, and that pausing repeatedly mid-generation loses no request and does not make the ranks diverge."""+ _init_distributed(tp=self.tp_size, backend="nccl")(_tp_pause_generation_worker)(+ model_id="TinyLlama/TinyLlama-1.1B-Chat-v1.0",+ attn_implementation="paged|sdpa",+ max_new_tokens=20,+ num_pauses=5,+ skew_seconds=0.25,+ )+ @slow def test_continuous_batching_tp_greedy(self) -> None: """Test that continuous batching with `DistributedConfig(tp_size=...)` produces non-empty, reproducible greedy outputs and