Model Sources

Model Sources

source	factory	URI string
Unity Catalog	dib.models.uc(name, version=None, alias=None, stage_dir=None)	uc:cat.schema.model@alias, uc:cat.schema.model/3, models:/cat.schema.model@alias
directory / Volume / DBFS	dib.models.path(path, copy_to_local=False)	/..., ./..., ~/..., dbfs:/..., file:/...
OpenAI-compatible API	dib.models.openai(model, base_url=None, api_key_env=None)	openai:<model>
Databricks Model Serving	dib.models.databricks(endpoint, api_key_env=None)	databricks:<endpoint>
fake	dib.models.fake(response=None, fail_marker=None, crash_marker=None, exit_marker=None, delay_ms=0, structured_output=True)	fake

dib.models.parse(uri) converts a URI into a spec. Every job accepts either.

Specs are plain data: they never contain credentials or weights. Workers resolve them on their side.

Unity Catalog

dib.models.uc("main.llm.qwen2_5_7b_instruct", alias="prod")
dib.models.uc("main.llm.qwen2_5_7b_instruct", version=3)

Exactly one of version or alias. The worker:

1. resolves the alias to a version through MLflow (registry_uri="databricks-uc");
2. downloads the version into the node-local model cache, once per node however many processes ask;
3. finds the Hugging Face model inside the artifact.

Supported layouts:

• MLflow transformers flavor (model/, components/tokenizer, components/feature_extractor);
• models logged without weights (save_pretrained=False): the source repository recorded in MLmodel is loaded from the Hugging Face Hub instead;
• any artifact with one Hugging Face model directory (config.json + weights), e.g. sentence-transformers exports.

Requires distribatch[mlflow]. Executors need registry credentials (DATABRICKS_HOST/DATABRICKS_TOKEN in the cluster environment), or use stage_dir:

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

With stage_dir, the Spark backend downloads the model once on the driver (which has credentials in notebooks) into that directory, and executors copy it from there to their local disk.

Directories, Volumes, DBFS

dib.models.path("/Volumes/main/ml/weights/bge-m3")
dib.models.path("dbfs:/models/whisper-small", copy_to_local=True)

A Hugging Face model directory (config.json + weights), or one of the layouts above. dbfs:/x is read through the /dbfs/x FUSE mount (dbfs:/Volumes/... through /Volumes/...).

copy_to_local=True copies the directory into the node-local cache first (once per node, again when the files change). Recommended for vLLM and big models: FUSE reads from Volumes are much slower than local NVMe.

Node-local model cache

Location: $DISTRIBATCH_CACHE_DIR, else /local_disk0/distribatch/cache when it exists (Databricks), else ~/.cache/distribatch. Entries are created under a file lock and published atomically: concurrent workers wait for one download instead of racing. Converted models (faster-whisper) live there too.

OpenAI-compatible APIs

dib.models.openai("gpt-4o-mini")                                       # api.openai.com, OPENAI_API_KEY
dib.models.openai("gpt-4o-mini", api_key_env="TEAM_OPENAI_KEY")
dib.models.openai("Qwen/Qwen2.5-7B-Instruct", base_url="http://vllm.internal:8000/v1")

• base_url=None: the official API (or OPENAI_BASE_URL), key from OPENAI_API_KEY.
• base_url="http(s)://...": any compatible server (vLLM, TGI, LiteLLM, Ollama…). No key is needed unless api_key_env is set.
• api_key_env: name of the environment variable holding the key on the workers. Keys are never stored in the job.

Databricks Model Serving

dib.models.databricks("databricks-meta-llama-3-3-70b-instruct")   # Foundation Model API
dib.models.databricks("my-custom-endpoint")

Calls <workspace>/serving-endpoints through its OpenAI-compatible API. Host and token come from DATABRICKS_HOST/DATABRICKS_TOKEN or databricks-sdk authentication (the notebook context on the driver), or the token from api_key_env. On Spark, executors need the variables in the cluster environment. See Databricks Guide.

Fake

Deterministic, no weights, no network, for dry runs and tests. It echoes the input (or returns response), embeds into hash vectors, and "transcribes" to transcript of <path>. fail_marker / crash_marker / exit_marker simulate a row error, a failed batch, and a dead worker process.