Evaluate RAG using Batch Inference with Ray Data LLM#

In this tutorial, we’ll walk through how to efficiently evaluate large-scale requests using batch inference with Ray Data LLM. This approach is ideal when you need to process many prompts without the overhead of real-time online inference. By following these steps, you’ll learn how to load data, generate embeddings, build retrieval-augmented prompts, and finally process them with an LLM.

In previous tutorial (Notebook #6), we demonstrated how to evaluate Retrieval-Augmented Generation (RAG) using a regular pipeline with an LLM deployed as an online service. However, this approach proved to be slow and inefficient even for 64 requests, making it unscalable for evaluating larger workloads.

Relying on a production LLM API for massive requests evaluation can disrupt production stability. Moreover, deploying a separate evaluation service or API adds extra overhead, and if not properly shut down after testing, it can lead to unnecessary costs.

Here is the architecture diagram for evaluating RAG using batch inference with Ray Data:

https://raw.githubusercontent.com/ray-project/ray/refs/heads/master/doc/source/ray-overview/examples/e2e-rag/images/batch_inference_rag_evaluation.png

How to Decide Between Online vs. Offline Inference for LLM#

Online Inference: Use online LLM inference (e.g., via an Anyscale Endpoint) when you require real-time responses or interactive engagement with the LLM. This approach optimizes for low latency, making it ideal for prompt-based applications where speed is critical.
Offline Inference (Batch Inference): Offline LLM inference is best suited for processing a large number of prompts within a specific time frame when real-time responsiveness is not essential. Batch inference can process requests in minutes to hours, allowing for better resource management and cost efficiency.

Key Benefits of Using Batch Inference with Ray Data LLM#

Scalability: Process large-scale datasets efficiently without the constraints of real-time latency requirements.
Optimized Resource Utilization: Maximize throughput by taking full advantage of available compute resources (e.g., GPUs) and scheduling jobs in parallel.
Cost Efficiency: Leverage cost-effective compute resources and schedule batch jobs during off-peak hours to reduce overall expenditure.
Flexibility and Isolation: Maintain a dedicated evaluation environment that can be tuned and scaled independently from production systems, minimizing the risk of disrupting live services.

Anyscale-Specific Configuration

Note: This tutorial is optimized for the Anyscale platform. When running on open source Ray, additional configuration is required. For example, you’ll need to manually:

Configure your Ray Cluster: Set up your multi-node environment (including head and worker nodes) and manage resource allocation (e.g., autoscaling, GPU/CPU assignments) without the Anyscale automation. See the Ray Cluster Setup documentation for details: https://docs.ray.io/en/latest/cluster/getting-started.html.
Manage Dependencies: Install and manage dependencies on each node since you won’t have Anyscale’s Docker-based dependency management. Refer to the Ray Installation Guide for instructions on installing and updating Ray in your environment: https://docs.ray.io/en/latest/ray-core/handling-dependencies.html.
Set Up Storage: Configure your own distributed or shared storage system (instead of relying on Anyscale’s integrated cluster storage). Check out the Ray Cluster Configuration guide for suggestions on setting up shared storage solutions: https://docs.ray.io/en/latest/train/user-guides/persistent-storage.html.

Prerequisites#

Before you move on to the next steps, please make sure you have all the required prerequisites in place.

Pre-requisite: You must have finished the data ingestion in Chroma DB with CHROMA_PATH = "/mnt/cluster_storage/vector_store" and CHROMA_COLLECTION_NAME = "anyscale_jobs_docs_embeddings". For setup details, please refer to Notebook #2.

Load the Evaluation Data#

The evaluation data is stored in a CSV file (evaluation_data/rag-eval-questions.csv) that contains 64 user queries grouped by category.

These queries cover a range of topics—from technical questions about Anyscale and its relationship with Ray, to casual, ethically sensitive, and non-English requests. This diverse dataset helps assess the system’s performance on a wide variety of inputs.

Feel free to add more categories or questions as needed.

import ray
# Load the CSV file directly into a Ray dataset.
csv_file = "evaluation_data/rag-eval-questions.csv"
ds = ray.data.read_csv(csv_file)
# Display the dataset
print(ds.schema())

Generating Embeddings from User Requests#

Like the previous tutorials, We use the same Sentence Transformer model to convert each user request into an embedding. This allows the later retrieval step to find relevant context for each prompt.

Note: since we only have 64 user requests for evaluation, we only use CPU to handle the embedding generation process, instead of using GPU. We also set the compute=ray.data.ActorPoolStrategy(size=1) which only use one CPU node. If you have a large volume of user requests, then you can consider to use multiple CPU nodes or enabel GPU for acceleration.

from typing import Dict
from sentence_transformers import SentenceTransformer
import torch


class UserRequestEmbedder:
    def __init__(self, model_name: str = "intfloat/multilingual-e5-large-instruct"):
        self.model_name = model_name
        self.model = SentenceTransformer(
            self.model_name,
            device="cuda" if torch.cuda.is_available() else "cpu"
        )

    def __call__(self, batch: Dict) -> Dict:
        # Generate embeddings for the 'user_request' field.
        embeddings = self.model.encode(batch["user_request"], convert_to_numpy=True) 
        batch["embeddings"] = embeddings

        return batch

# Use the Embedder class to process the batch and generate embeddings.
ds = ds.map_batches(UserRequestEmbedder, compute=ray.data.ActorPoolStrategy(size=1), batch_size=64)

Querying the Vector Store and Generating Prompts#

Next, we retrieve context for each user request by querying a vector store (using a tool such as Chroma). This context is then used to build a retrieval-augmented generation (RAG) prompt.

from rag_utils import  ChromaQuerier, render_rag_prompt


CHROMA_PATH = "/mnt/cluster_storage/vector_store"
CHROMA_COLLECTION_NAME = "anyscale_jobs_docs_embeddings"


# Initialize the components for rag.
chroma_querier = ChromaQuerier(CHROMA_PATH, CHROMA_COLLECTION_NAME, score_threshold=0.8)


def generate_prompts_batch(batch: dict) -> dict:
    # Extract the embeddings and other columns from the batch.
    embeddings = batch["embeddings"].tolist()
    
    # Perform a batched query.
    batch_results = chroma_querier.query_batch(embeddings, n_results=5)
    
    # Initialize a list to store the generated prompts.
    prompts = []

    # Iterate over the indices of the batch.
    for i, _ in enumerate(embeddings):
        user_request = batch["user_request"][i]
        context = batch_results[i]
        print(f"Render Prompt for user request: {user_request}")
        
        chat_history = ""
        company = "Anyscale"
        # Generate the prompt for this row.
        prompt = render_rag_prompt(company, user_request, context, chat_history)
        prompts.append(prompt)

    
    # Update the batch with the new 'prompt' column.
    batch["prompt"] = prompts
  
    return batch

# Use map_batches on your Ray dataset.
ds = ds.map_batches(generate_prompts_batch, batch_size=128)

Configuring and Running LLM Inference#

Now, use Ray’s LLM processor to send your generated prompts to the LLM for inference. In this example, we use the Qwen model, but you can adjust the configuration as needed.

max_num_batched_tokens: Limits the total number of tokens processed in a batch.
max_num_seqs: Sets the maximum number of individual sequences that can be processed concurrently.
batch_size: Determines how many rows of data are processed at once.

Note:

If you encounter out of memory issue, please decrease the max_num_seqs and batch_size.

from ray.data.llm import vLLMEngineProcessorConfig, build_processor

# Create the LLM processor configuration
model_source='Qwen/Qwen2.5-32B-Instruct' 

config = vLLMEngineProcessorConfig(
    model_source=model_source,
    accelerator_type='L4',
    engine_kwargs={
        'max_num_batched_tokens': 8192,
        'max_model_len': 8192,
        'max_num_seqs': 128, 
        'tensor_parallel_size': 4,
        'trust_remote_code': True,
    },
    concurrency=1,
    batch_size=128,  
)


# Build the processor using a preprocessor that uses the generated prompt.
processor = build_processor(
    config,
    preprocess=lambda row: dict(
        messages=[
            {"role": "user", "content": row["prompt"]},
        ],
        sampling_params=dict(
            temperature=0,
            max_tokens=1024, # max response tokens is 1024
            detokenize=False,
        ),
    ),
    postprocess=lambda row: dict(
        resp=row["generated_text"],
        **row,  # Return all original columns.
    ),
)

# Process the dataset using the LLM inference processor.
ds = processor(ds)
results = ds.take_all()


# Print the output for each row.
for row in results:
    print(row)

Saving the Batch Inference Results#

Finally, convert the Ray dataset to a pandas DataFrame and save the results to a CSV file for further analysis.

import pandas as pd

# Convert the final Ray dataset to a list of dictionaries.
df_eval = pd.DataFrame(results)

# Define the desired order for the first few columns
desired_order = ['category', 'user_request',  'resp', 'time_taken_llm']

# Create a list of the remaining columns that are not in desired_order
remaining_cols = [col for col in df_eval.columns if col not in desired_order]

# Reorder the DataFrame
df_eval = df_eval[desired_order + remaining_cols]

# Save the DataFrame to a CSV file.
df_eval.to_csv(f'eval_results_batch_inference_qwen32b.csv', index=False)

Visualize the Results#

After running the evaluation, open the resulting CSV file (eval_results_online_inference.csv) to review:

The user request.
The retrieved context from the vector store.
The generated answer from the LLM service.

Note after finishing the pipeline, the build_processor also added a few more fileds such as time_taken_llm, num_input_tokens, num_generated_tokens etc.

df_eval = pd.read_csv('eval_results_batch_inference_qwen32b.csv')

# Display the first 5 rows of the DataFrame.
print(df_eval.head())

Evaluate Results and Imrpove RAG Quality#

You can manually review the evaluation results, marking responses as good or bad, and refine the prompt iteratively to improve performance.

Save the high-quality responses as a golden dataset for future reference. Once you have a substantial golden dataset, you can leverage more advanced LLMs—potentially with reasoning capabilities—to act as an LLM judge, comparing new RAG results against the golden dataset.

Final Notes#

This tutorial has provided a comprehensive overview of setting up batch inference with Ray Data LLM for evaluating retrieval-augmented generation (RAG). By following these steps, you can build scalable, cost-effective, and flexible evaluation workflows for your LLM applications.

Experiment with different models and configurations to further optimize performance and resource utilization. Happy evaluating!