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

Code before:

                # Update the request and stop if it is complete
                is_finished = state.update_and_check_completion(token)
                # We mark the completed blocks as such
                self.cache.mark_shareable_blocks_as_complete(state)
                if is_finished:
                    self.metrics.record_request_completion(state.created_time, state.request_id)
                    self.scheduler.finish_request(state.request_id, evict_from_cache=(not self.manual_eviction))
                self._maybe_send_output(state)
            #  Otherwise, the request is still prefilling, but the prefill has been split
            elif state.status == RequestStatus.PREFILLING_SPLIT:
                self.cache.mark_shareable_blocks_as_complete(state)
                state.status = RequestStatus.SPLIT_PENDING_REMAINDER
            else:
                raise ValueError(f"Request {state.request_id} is in an unexpected state: {state.status}")

        # If some requests need to be forked, we do it now
        copy_source, copy_destination = [], []
        while self.scheduler._requests_to_fork:
            # Get the number of children and reset it so it's not forked again
            state = self.scheduler._requests_to_fork.pop()
            num_children = state.num_children
            state.num_children = 0
            # Create the new request
            new_request_ids = [f"{state.request_id}__child#{i}" for i in range(num_children)]
            new_requests = [state.fork(new_request_id) for new_request_id in new_request_ids]
            # Fork the cache
            copy_src, copy_dst = self.cache.fork_request(state.request_id, new_request_ids)
            copy_source.extend(copy_src)
            copy_destination.extend(copy_dst)
            # Add the new requests to the scheduler
            for new_request in new_requests:
                self.scheduler.active_requests[new_request.request_id] = new_request
            # FIXME: if fork cant be done, create a new pending request without forking instead of crashing everything

        # The copy induced by the fork is done in one go (if it's even needed)
        if copy_source:
            self.cache.copy_cache(copy_source, copy_destination)

        if self.cache.get_num_free_blocks() == 0:
            raise ValueError("No more free blocks")

    @traced
    def has_pending_requests(self) -> bool:
        """Check if there are any active or waiting requests."""
        return self.scheduler.has_pending_requests()

    @traced
    def handle_batch_error(self, error):
        """Handle errors during batch processing."""
        failed_reqs = self.requests_in_batch

Review comment (suggestion):
```suggestion
            for new_request_id in new_request_ids:
                self.scheduler.active_requests[new_request.request_id] = state.fork(new_request_id)
```

Code after:

                # Update the request and stop if it is complete
                is_finished = state.update_and_check_completion(token)
                # We mark the completed blocks as such
                self.cache.mark_shareable_blocks_as_complete(state)
                if is_finished:
                    self.metrics.record_request_completion(state.created_time, state.request_id)
                    self.scheduler.finish_request(state.request_id, evict_from_cache=(not self.manual_eviction))
                self._maybe_send_output(state)
            #  Otherwise, the request is still prefilling, but the prefill has been split
            elif state.status == RequestStatus.PREFILLING_SPLIT:
                self.cache.mark_shareable_blocks_as_complete(state)
                state.status = RequestStatus.SPLIT_PENDING_REMAINDER
            else:
                raise ValueError(f"Request {state.request_id} is in an unexpected state: {state.status}")

        # If some requests need to be forked, we do it now
        copy_source, copy_destination = [], []
        while self.scheduler._requests_to_fork:
            # Get the number of children and reset it so it's not forked again
            state = self.scheduler._requests_to_fork.pop()
            num_children = state.num_children
            state.num_children = 0
            # Create the new request and add them to the scheduler
            new_request_ids = [f"{state.request_id}__child#{i}" for i in range(num_children)]
            for new_request_id in new_request_ids:
                self.scheduler.active_requests[new_request_id] = state.fork(new_request_id)
            # Fork the cache
            copy_src, copy_dst = self.cache.fork_request(state.request_id, new_request_ids)
            copy_source.extend(copy_src)
            copy_destination.extend(copy_dst)
            # FIXME: if fork cant be done, create a new pending request without forking instead of crashing everything

        # The copy induced by the fork is done in one go (if it's even needed)
        if copy_source:
            self.cache.copy_cache(copy_source, copy_destination)

        if self.cache.get_num_free_blocks() == 0:
            raise ValueError("No more free blocks")

    @traced
    def has_pending_requests(self) -> bool:
        """Check if there are any active or waiting requests."""
        return self.scheduler.has_pending_requests()

    @traced
    def handle_batch_error(self, error):
        """Handle errors during batch processing."""
        failed_reqs = self.requests_in_batch
        for req in failed_reqs:
            self._handle_request_error(error, req)