Databricks Guide

Databricks Guide

Install

%pip install "distribatch[openai,vllm,mlflow]"     # pick the extras you need
dbutils.library.restartPython()

Notebook-scoped installs reach the executors. For jobs, add distribatch[...] as a PyPI library of the job cluster. pyspark is preinstalled, so the spark extra is not needed.

Which compute

workload	compute
Model Serving / Foundation Model APIs / external APIs	any cluster, incl. shared and serverless (Spark Connect). No GPUs needed.
self-hosted LLMs and embedders (vLLM)	single-user GPU cluster, Databricks Runtime ML (GPU) 15.4 LTS+
Whisper (faster-whisper / transformers)	GPU cluster for volume. CPU clusters work for small batches.
one GPU node, data fits on the driver	single-node GPU cluster + backend="multiprocessing"

Credentials on executors

Notebooks authenticate the driver automatically, but Python code on executors has no notebook context. Anything that calls Databricks APIs from executors needs credentials in the cluster environment: models.databricks(...) endpoints with backend="spark", or UC models without stage_dir.

Cluster → Advanced options → Spark → Environment variables:

DATABRICKS_HOST=https://<workspace-host>
DATABRICKS_TOKEN={{secrets/<scope>/<key>}}

Use a service principal token in the secret. The same mechanism works for other API keys, e.g. OPENAI_API_KEY={{secrets/llm/openai}} with models.openai(..., api_key_env="OPENAI_API_KEY"). distribatch never ships credentials from the driver to executors.

Local and multiprocessing backends run on the driver, where notebook authentication works without any setup.

Unity Catalog models

Register a Hugging Face model with MLflow once:

import mlflow, transformers
mlflow.set_registry_uri("databricks-uc")
pipe = transformers.pipeline("text-generation", model="Qwen/Qwen2.5-7B-Instruct")
with mlflow.start_run():
    mlflow.transformers.log_model(pipe, name="model", registered_model_name="main.llm.qwen2_5_7b_instruct")
mlflow.MlflowClient().set_registered_model_alias("main.llm.qwen2_5_7b_instruct", "prod", 1)

Use it with stage_dir, so executors need no registry credentials and every node copies the weights from a Volume to its local NVMe:

model = dib.models.uc("main.llm.qwen2_5_7b_instruct", alias="prod", stage_dir="/Volumes/main/llm/stage")

Staged versions are reused across runs (one directory per model version). Delete old ones from the Volume when you no longer need them.

GPU clusters

vLLM, one model per GPU. Run one Spark task per GPU, and distribatch pins each Python worker to the GPU Spark assigned to its task. In the cluster's Spark config:

spark.task.resource.gpu.amount 1

Then use partitions= = total GPUs (or a small multiple), and give vLLM the memory:

job = dib.GenerateJob(model=model, prompt=..., engine_options={"gpu_memory_utilization": 0.9, "max_model_len": 8192})
out = job.run(df, backend="spark", partitions=8, batch_size=256, checkpoint=..., id_col="id")

Large models across the GPUs of a node (tensor parallelism): one task per node owns all its GPUs. Set spark.task.resource.gpu.amount = GPUs per node and engine_options={"tensor_parallel_size": <GPUs per node>}.

Batch size: vLLM schedules internally, so a large batch_size (hundreds) keeps the GPU busy. For the transformers engine, micro_batch_size controls the forward-pass size and batch_size the rows per call.

Single-node clusters: backend="multiprocessing" starts one worker per GPU (or per tensor_parallel_size GPUs) on the driver, with no Spark task configuration needed.

Endpoints and rate limits

Total concurrency against an endpoint = running tasks × max_concurrency. For pay-per-token Foundation Model APIs, start with partitions=4, max_concurrency=8 and raise them while watching for RateLimitError in dib.summarize(out).errors. requests_per_minute caps the request rate of each worker process.

Audio from Volumes

audio = (
    spark.read.format("binaryFile")
    .option("pathGlobFilter", "*.{wav,mp3,flac,m4a}")
    .option("recursiveFileLookup", "true")
    .load("/Volumes/main/raw/calls/")
    .select("path")                       # let workers read files: small shuffles
)
out = asr.run(audio, backend="spark", batch_size=8, checkpoint="/Volumes/main/audio/ckpt", id_col="path")

binaryFile reports Volume paths as dbfs:/Volumes/.... distribatch reads them through /Volumes/.... m4a/aac need an ffmpeg binary on the workers for local engines. Check with %sh which ffmpeg and install it with a cluster init script if it's missing.

Model cache on nodes

Local models are cached on /local_disk0/distribatch/cache, the node's NVMe. The cache lives as long as the cluster. Set DISTRIBATCH_CACHE_DIR in the cluster environment to move it.

Serverless and shared clusters

These run Spark Connect: the spark backend is supported (it uses mapInArrow only). Use them for API-based jobs. GPU engines need GPU clusters.