Skip to main content

How to run batch LLM inference over millions of records

Summary

  • AI Functions run batch LLM inference directly on data in Unity Catalog with SQL, eliminating the need to export or move data.
  • For simple use cases, use SQL AI functions like ai_query to process entire tables or datasets without manual batching.
  • Integrate batch inference into production workflows (Lakeflow Jobs, Lakeflow Pipelines) for scheduled or incremental processing.
  • For advanced use cases, use Ray Data and vLLM to run inference across multiple GPU nodes with custom optimization.
  • Databricks handles automatic scaling, retries, and fault tolerance, so you submit your full dataset in a single request.

How to run batch LLM inference over millions of records

Processing very large volumes of records through an LLM is common but challenging. Exporting data, managing batch splits, handling retries, and orchestrating results is tedious at scale. Databricks lets you run batch LLM inference directly on data in Unity Catalog using SQL, without moving data or building custom orchestration. The system automatically scales to handle large workloads and manages retries and fault tolerance for you.

What is batch LLM inference and when you need it

Batch inference is applying a language model to many records offline, not serving real-time requests. Use cases include extracting categories or sentiment from a table of articles, translating text in bulk, generating summaries for historical data, or enriching a customer dataset with AI-generated features. Batch inference is typically lower-latency-sensitive than real-time serving, which lets you prioritize cost-efficiency and throughput over response speed.

Key approaches for batch LLM inference

When building a batch inference pipeline, consider these options:

Approach Best for Key benefit
SQL AI Functions (ai_query, ai_translate, etc.) General-purpose inference, simple use cases Write a SQL query; no code needed. Data stays in place.
AI Functions in production workflows Scheduled or triggered batch jobs Integrate with Lakeflow Jobs or Lakeflow Pipelines for orchestration.
Ray Data + vLLM Advanced optimization, custom tuning Run multiple vLLM replicas across GPU nodes for parallelism.
Model Serving endpoints Flexible orchestration Set up custom processing logic around provisioned or on-demand endpoints.

How Databricks approaches batch LLM inference

SQL-based AI Functions (simplest path). Use built-in SQL functions to apply LLMs directly to your tables:

SELECT 
  article_text,
  ai_query(
    "system.ai.meta-llama-3-3-70b-instruct",
    "Extract the main category: " || article_text
  ) AS category
FROM articles
WHERE processed = false;

Remove the WHERE clause or set a high limit to process your entire dataset in one query. AI Functions automatically parallelize inference across Databricks cluster resources.
No data movement. Data stays in Delta Lake throughout the inference pipeline. Prompts and inputs stay local; only model calls go to the foundation model. Results are written back to Unity Catalog as Delta tables or views.
Databricks-hosted foundation models. Use models prefixed with databricks- (e.g., databricks-meta-llama-3-1-70b-instruct) for provisionless, fully managed inference. These are optimized for batch workloads and scale automatically.
Production workflow integration. Wrap your AI Functions SQL in a Lakeflow Job, Lakeflow Pipeline, or declarative pipeline for scheduling, retries, and monitoring:

CREATE OR REFRESH MATERIALIZED VIEW article_categories
AS
SELECT
  article_id,
  article_text,
  ai_query(
    "system.ai.meta-llama-3-3-70b-instruct",
    "Categorize this article: " || article_text
  ) AS category
FROM articles_raw;

Define this view in a Lakeflow Pipeline, and it updates automatically on a schedule or when source data changes.
Ray Data and vLLM for advanced use cases. If you need custom optimization, launch a Ray cluster across compute nodes and run vLLM replicas for high parallelism:

  • Start Ray across your Databricks cluster nodes
  • Launch one vLLM replica per GPU node
  • Stream prompts through the replicas in parallel
  • Write results back to Unity Catalog as Parquet or Delta

This approach gives you fine-grained control over batching, concurrency, and model deployment.
Automatic scaling and fault tolerance. Databricks handles infrastructure automatically. Resources scale to your workload, requests are retried on transient failures, and results are written transactionally. You focus on the SQL or Ray logic; Databricks manages the rest.

Use cases for batch LLM inference

Bulk content generation and enrichment. Generate product descriptions, summaries, or tags for an entire catalog of items. Update nightly or on-demand.
Data categorization and tagging. Classify support tickets, categorize articles, or label images at scale without manual effort.
Translation workflows. Translate a large corpus of text into multiple languages in a single batch job.
Feature engineering for ML. Generate embeddings or semantic features for a dataset to feed into downstream ML models.
Sentiment and intent analysis. Analyze customer feedback, survey responses, or social media data in bulk.

FAQs

How do I know if my inference will succeed before I process the entire dataset?

Start with a small limit (e.g., LIMIT 100) in your SQL query to test your prompt and model. Once you verify the results look correct, remove or raise the limit to process the full dataset.

Can I process data outside Databricks with AI Functions?

No. AI Functions work on data in Unity Catalog (Delta tables). External data must be ingested into Databricks first.

What happens if my batch job fails halfway through?

Databricks automatically retries transient failures. If your query fails, review the error, fix the issue (e.g., adjust the prompt), and re-run the SQL. Results already written are preserved.

How do I choose between SQL AI Functions and Ray Data + vLLM?

Start with SQL AI Functions (simplest, fewest dependencies). Move to Ray Data + vLLM if you need custom optimization, model fine-tuning, or multi-GPU parallelism.

The information provided herein is for general informational purposes only and may not reflect the most current product capabilities or configurations.