Repository: huggingface/transformers
File: src/transformers/generation/continuous_batching/scheduler.py (Python)

Code before:

    @traced
    def schedule_batch(self, token_budget: int, cache_budget: int) -> list[RequestState] | None:
        priority_states: list[RequestState] = []
        second_priority_states: list[RequestState] = []

        for state in self.active_requests.values():
            # XXX: when cache is full, state can stay on `PREFILLING_SPLIT` so we need to take those into account
            if state.status in [RequestStatus.PREFILLING_SPLIT, RequestStatus.SPLIT_PENDING_REMAINDER]:
                priority_states.append(state)
            elif state.status == RequestStatus.DECODING:
                second_priority_states.append(state)

        # Add waiting requests to second priority
        if not self.block_new_requests:
            for req_id in self.waiting_requests_order:
                second_priority_states.append(self.waiting_requests[req_id])

        candidates = priority_states + second_priority_states
        request_ids_to_remove_from_waiting = set()
        scheduled_requests, one_allocation_failed = self._process_candidates(
            candidates,
            token_budget,
            cache_budget,
            request_ids_to_remove_from_waiting,
            safety_margin=None,
        )

        # We remove waiting requests before checking requests were scheduled, because there might have been prefill matches
        self._cleanup_waiting_queue(request_ids_to_remove_from_waiting)

        # If no requests were scheduled and the cache is full, we signal it by returning None
        if not scheduled_requests and one_allocation_failed:
            return None

        return scheduled_requests


SCHEDULER_MAPPING = {
    "fifo": FIFOScheduler,
    "prefill_first": PrefillFirstScheduler,
}


Review comment (suggestion):
If we set safety_margins to 0 as a default value, then we are never outside of the margins: 0 would be better than None here

Code after:

    @traced
    def schedule_batch(self, token_budget: int, cache_budget: int) -> list[RequestState] | None:
        priority_states: list[RequestState] = []
        second_priority_states: list[RequestState] = []

        for state in self.active_requests.values():
            # XXX: when cache is full, state can stay on `PREFILLING_SPLIT` so we need to take those into account
            if state.status in [RequestStatus.PREFILLING_SPLIT, RequestStatus.SPLIT_PENDING_REMAINDER]:
                priority_states.append(state)
            elif state.status == RequestStatus.DECODING:
                second_priority_states.append(state)

        # Add waiting requests to second priority
        if not self.block_new_requests:
            for req_id in self.waiting_requests_order:
                second_priority_states.append(self.waiting_requests[req_id])

        candidates = priority_states + second_priority_states
        request_ids_to_remove_from_waiting = set()
        scheduled_requests, one_allocation_failed = self._process_candidates(
            candidates,
            token_budget,
            cache_budget,
            request_ids_to_remove_from_waiting,
            safety_margin=0.0,
        )

        # We remove waiting requests before checking requests were scheduled, because there might have been prefill matches
        self._cleanup_waiting_queue(request_ids_to_remove_from_waiting)

        # If no requests were scheduled and the cache is full, we signal it by returning None
        if not scheduled_requests and one_allocation_failed:
            return None

        return scheduled_requests


SCHEDULER_MAPPING = {
    "fifo": FIFOScheduler,
    "prefill_first": PrefillFirstScheduler,
}