hibachi 0.1.9

use std::sync::Arc;
use async_trait::async_trait;
use tokio::sync::Mutex;

/// A trait that defines how batches of inference requests are processed by a model.
///
/// The `BatchHandler` trait abstracts the core operations required for batched inference
/// processing across different inference patterns (such as feedforward and autoregressive).
/// Implementations handle request batching, model execution, and output distribution,
/// while managing the state of active requests.
///
/// # Type Parameters
///
/// * `Request` - The type representing a single inference request.
/// * `ModelInput` - The type representing the batched input to the model, must be cloneable.
/// * `ModelOutput` - The type representing the output from the model, must be cloneable.
///
#[async_trait]
pub trait BatchHandler {
    /// The type representing a single inference request.
    ///
    /// This type contains all necessary information to process a request, including
    /// input data and a way to send results back to the caller.
    type Request;

    /// The type representing the batched input to the model.
    ///
    /// This is typically a tensor or collection of tensors that can be fed directly
    /// to the model for inference. Must be cloneable to support operations that need
    /// to retain a copy of the input for later stages.
    type ModelInput: Clone;

    /// The type representing the output from the model.
    ///
    /// This is typically a tensor or collection of tensors produced by the model
    /// that will be processed into final results. Must be cloneable to support
    /// distributing outputs to multiple requests.
    type ModelOutput: Clone;

    /// Constructs a model input batch from a collection of requests.
    ///
    /// This method transforms the raw requests into a format that can be directly
    /// fed to the model for inference. The created batch is stored in the `model_input`
    /// parameter.
    ///
    /// # Parameters
    ///
    /// * `model_input` - A mutable reference to an optional model input. This will be
    ///   populated with the constructed batch.
    /// * `requests` - A slice of requests to be processed in this batch.
    ///
    /// # Implementation Notes
    ///
    /// Implementations should:
    /// * Extract relevant data from each request
    /// * Combine them into an appropriate batch format
    /// * Set `model_input` to `Some(batch)`
    ///
    /// For autoregressive models, this might involve updating only the token positions
    /// that are currently active.
    async fn make_batch_input(&self, model_input: &mut Option<Self::ModelInput>, requests: &[Self::Request]);

    /// Executes the model forward pass on the given input batch.
    ///
    /// This method is responsible for running the actual model inference on the
    /// prepared input batch.
    ///
    /// # Parameters
    ///
    /// * `model_input` - The input batch to process.
    ///
    /// # Returns
    ///
    /// The output generated by the model.
    ///
    /// # Implementation Notes
    ///
    /// Implementations should:
    /// * Process the model input through the inference model
    /// * Return the raw model output without additional processing
    ///
    /// This method should contain only the core inference operation, with pre- and
    /// post-processing handled by the other methods.
    async fn forward(&self, model_input: &Self::ModelInput) -> Self::ModelOutput;

    /// Processes model outputs and updates the state of active requests.
    ///
    /// This method is responsible for:
    /// * Sending results to clients (immediately or incrementally)
    /// * Determining which requests are still active
    /// * Updating the batch and input state for the next iteration
    /// * Updating the count of active requests
    ///
    /// # Parameters
    ///
    /// * `batch` - A mutable reference to the current batch of requests being processed.
    ///   After this method returns, this should only contain requests that are still active.
    /// * `input` - A mutable reference to the current model input. This should be updated
    ///   to reflect the next input state.
    /// * `output` - The output from the most recent model execution.
    /// * `active_count` - A shared counter of active requests across all batches.
    ///
    ///
    /// The execution strategy should guide how this model behaves:
    /// * Feedforward models typically process all requests completely and clear the batch
    /// * Autoregressive models may retain some requests for additional processing in
    ///   subsequent iterations
    async fn handle_outputs(
        &self,
        batch: &mut Vec<Self::Request>,
        input: &mut Option<Self::ModelInput>,
        output: Self::ModelOutput,
        active_count: Arc<Mutex<usize>>,
    );
}