Skip to main content

Reusable CLI for uploading, submitting, validating, fetching logs, and cleaning Databricks job runs

Project description

databricks-job-runner

Reusable CLI for uploading, submitting, and cleaning Databricks job runs.

Wraps the Databricks Python SDK into a small library that each project configures with a Runner instance. One Runner gives you eight CLI subcommands — upload, submit, validate, logs, clean, catalog, schema, and volume — without writing any Databricks API code in your project.

Installation

uv add databricks-job-runner

Or with pip:

pip install databricks-job-runner

For local development against a checkout:

# pyproject.toml
[tool.uv.sources]
databricks-job-runner = { path = "../databricks-job-runner", editable = true }

Warning — do not list databricks-job-runner as a core dependency.

databricks-job-runner is a local-only CLI tool — it is not published to PyPI. If you add it to your project's [project.dependencies] (core dependencies), any wheel you build from that project will declare it as a requirement. When Databricks serverless (or any remote environment) tries to install your wheel, pip will fail because it cannot resolve databricks-job-runner.

Instead, put it in an optional extras group so it is only installed locally:

[project.optional-dependencies]
cli = ["databricks-job-runner"]

Then install locally with uv sync --extra cli (or pip install -e '.[cli]'). Your submitted scripts (e.g. run_my_package.py) should never import databricks_job_runner — they run on Databricks where it is not available.

Quick start

Create a cli/ package in your project with two files:

cli/__init__.py

from databricks_job_runner import Runner

runner = Runner(
    run_name_prefix="my_project",
    wheel_package="my_package",  # optional
)

cli/__main__.py

from cli import runner
runner.main()

Then run from the project root:

python -m cli upload --all          # upload agent_modules/*.py
python -m cli upload test_hello.py  # upload a single file
python -m cli upload --wheel        # build and upload wheel
python -m cli submit test_hello.py  # submit a job and wait
python -m cli submit test_hello.py --no-wait
python -m cli validate              # list remote workspace contents
python -m cli validate test_hello.py  # verify a specific file is uploaded
python -m cli logs                  # stdout/stderr from the most recent run
python -m cli logs 12345            # stdout/stderr from a specific run
python -m cli clean --yes           # clean workspace + runs
python -m cli clean --runs --yes    # clean only runs

# Unity Catalog management
python -m cli catalog list
python -m cli catalog create my_catalog --comment "Analytics"
python -m cli catalog delete my_catalog --force --yes

python -m cli schema list my_catalog
python -m cli schema create my_catalog.my_schema
python -m cli schema delete my_catalog.my_schema --yes

python -m cli volume list my_catalog.my_schema
python -m cli volume create my_catalog.my_schema.my_vol
python -m cli volume create my_catalog.my_schema.ext_vol --volume-type EXTERNAL --storage-location s3://bucket/path
python -m cli volume delete my_catalog.my_schema.my_vol --yes

Configuration

The runner reads a .env file from the project root. Core keys (all prefixed with DATABRICKS_ for consistency):

Key Default Required Description
DATABRICKS_PROFILE no CLI profile in ~/.databrickscfg. When unset, the SDK's unified auth falls back to env vars (DATABRICKS_HOST/DATABRICKS_TOKEN), Azure CLI, service principals, etc.
DATABRICKS_COMPUTE_MODE cluster no cluster or serverless. Selects the compute backend for submitted jobs.
DATABRICKS_CLUSTER_ID when DATABRICKS_COMPUTE_MODE=cluster All-purpose cluster to run jobs on. Started automatically if terminated.
DATABRICKS_SERVERLESS_ENV_VERSION 3 no Serverless environment version (e.g. 3 for Python 3.12).
DATABRICKS_WORKSPACE_DIR yes Remote workspace path (e.g. /Users/you/my_project)
DATABRICKS_VOLUME_PATH when using upload --wheel UC Volume path for wheel uploads.

Precedence: pre-existing environment variables override .env values, matching 12-factor conventions (CI/CD and shell exports can override the file).

Additional non-core keys are captured in RunnerConfig.extras and automatically passed to submitted jobs as KEY=VALUE parameters. Scripts call inject_params() at startup to load them into os.environ, then use pydantic BaseSettings to read configuration.

Compute modes

  • Classic cluster (DATABRICKS_COMPUTE_MODE=cluster, the default): jobs submit to an existing all-purpose cluster identified by DATABRICKS_CLUSTER_ID. The runner auto-starts the cluster if it is terminated, and attaches wheels via Library(whl=...).
  • Serverless (DATABRICKS_COMPUTE_MODE=serverless): jobs submit to Databricks serverless compute with a job-level environment spec. No cluster ID needed; wheels attach as Environment.dependencies entries (UC Volume paths are supported directly).

Example .env (classic cluster)

DATABRICKS_PROFILE=my-profile
DATABRICKS_CLUSTER_ID=0123-456789-abcdef
DATABRICKS_WORKSPACE_DIR=/Users/ryan.knight@example.com/my_project
DATABRICKS_VOLUME_PATH=/Volumes/catalog/schema/volume
NEO4J_URI=neo4j+s://abc123.databases.neo4j.io
NEO4J_PASSWORD=secret

Example .env (serverless)

DATABRICKS_PROFILE=my-profile
DATABRICKS_COMPUTE_MODE=serverless
DATABRICKS_SERVERLESS_ENV_VERSION=3
DATABRICKS_WORKSPACE_DIR=/Users/ryan.knight@example.com/my_project
DATABRICKS_VOLUME_PATH=/Volumes/catalog/schema/volume

All DATABRICKS_* keys listed above become typed fields on RunnerConfig; any other keys (like NEO4J_URI above) go into config.extras.

API

Runner

Runner(
    run_name_prefix: str,
    project_dir: Path | str | None = None,
    wheel_package: str | None = None,
)
Parameter Description
run_name_prefix Prefix for job run names and cleanup filtering
project_dir Project root (defaults to cwd()). Must contain .env and agent_modules/
wheel_package Package name for wheel builds. Enables upload --wheel. Wheels upload to <DATABRICKS_VOLUME_PATH>/wheels/

RunnerConfig

Pydantic model holding parsed .env values. Frozen (immutable) after construction.

Field Type Description
databricks_profile str | None CLI profile name, or None for unified-auth fallback
databricks_compute_mode Literal["cluster", "serverless"] Compute backend ("cluster" by default)
databricks_cluster_id str | None Cluster ID (required when databricks_compute_mode == "cluster")
databricks_serverless_env_version str Serverless environment version (default "3")
databricks_workspace_dir str Remote workspace root (required)
databricks_volume_path str | None UC Volume path for wheel uploads
extras dict[str, str] All non-core keys from .env

The env_params() method returns extras (plus DATABRICKS_VOLUME_PATH) as KEY=VALUE strings suitable for job parameter injection.

inject_params

from databricks_job_runner import inject_params
inject_params()

Call at the top of submitted scripts to parse KEY=VALUE parameters from sys.argv into os.environ. This lets scripts use pydantic BaseSettings or os.getenv() to read configuration that the runner injected from .env. Uses setdefault so pre-existing env vars take precedence.

Note: databricks_job_runner is not available on the Databricks cluster. For standalone scripts (not part of a wheel), inline the equivalent logic instead of importing:

import os, sys
for _arg in sys.argv[1:]:
    if "=" in _arg and not _arg.startswith("-"):
        _key, _, _value = _arg.partition("=")
        os.environ.setdefault(_key, _value)

For wheel-based scripts, the wheel's entry point can call inject_params() only if databricks_job_runner is listed as a wheel dependency — but since it is not published to PyPI, the inline approach is preferred.

RunnerError

Raised when a runner operation cannot proceed (missing config, file not found, cluster stopped, job failed). The CLI formats and exits; library callers can catch and handle.

Project layout

The runner expects this layout in your project:

my_project/
  .env
  agent_modules/
    test_hello.py
    run_lab2.py
    ...
  cli/
    __init__.py    # Runner config
    __main__.py    # entry point

Scripts in agent_modules/ are uploaded to {DATABRICKS_WORKSPACE_DIR}/agent_modules/ on Databricks and submitted as Spark Python tasks.

Subcommands

upload

  • upload <file> — Upload a single file from agent_modules/
  • upload --all — Upload all *.py files from agent_modules/
  • upload --wheel — Build a wheel with uv build and upload to the UC Volume (requires wheel_package and DATABRICKS_VOLUME_PATH)

submit

  • submit <script> — Submit a script as a one-time Databricks job and wait for completion. Default: test_hello.py
  • submit <script> --no-wait — Submit without waiting

On classic mode, if the target cluster is not already RUNNING, it is started automatically and the submit waits (up to 20 minutes, the SDK default) for it to reach RUNNING. On serverless, no warm-up step is required. When submitting a script named run_{wheel_package}.py, the runner automatically attaches the wheel — as a Library(whl=...) on classic, or as an Environment.dependencies entry on serverless.

validate

  • validate — List the remote workspace directory and its agent_modules/ subdirectory. On classic, auto-starts the cluster if needed; on serverless, this is a no-op.
  • validate <file> — Also verify that {DATABRICKS_WORKSPACE_DIR}/agent_modules/<file> exists; exits non-zero if not.

logs

  • logs — Print stdout/stderr, error, and trace from the most recent run matching {run_name_prefix}:*
  • logs <run_id> — Print output for a specific parent run ID

Output is fetched via the Jobs API's get_run_output, which returns the tail 5 MB of stdout/stderr captured per task (the API caps output size; truncation is signaled in the output). The runner resolves the parent run to its task-level run IDs automatically, so pass the parent run_id shown at submit time. Databricks auto-expires runs after 60 days.

clean

  • clean — Delete the remote workspace directory and all matching job runs
  • clean --workspace — Delete only the workspace directory
  • clean --runs — Delete only job runs
  • clean --yes — Skip confirmation prompt

catalog

Manage Unity Catalog catalogs.

  • catalog list — List all catalogs visible to the current user
  • catalog get <name> — Show details (owner, comment, created date) for a catalog
  • catalog create <name> [--comment TEXT] — Create a new catalog
  • catalog delete <name> [--yes] [--force] — Delete a catalog. --force cascades to all schemas, tables, and volumes inside it. Prompts for confirmation unless --yes is given.

schema

Manage Unity Catalog schemas. Schema names use dotted notation: catalog.schema.

  • schema list <catalog> — List schemas in a catalog
  • schema get <catalog.schema> — Show details for a schema
  • schema create <catalog.schema> [--comment TEXT] — Create a new schema
  • schema delete <catalog.schema> [--yes] — Delete a schema. Prompts for confirmation unless --yes is given.

volume

Manage Unity Catalog volumes. Volume names use dotted notation: catalog.schema.volume.

  • volume list <catalog.schema> — List volumes in a schema
  • volume get <catalog.schema.volume> — Show details (type, owner, storage location) for a volume
  • volume create <catalog.schema.volume> [--volume-type MANAGED|EXTERNAL] [--storage-location URL] [--comment TEXT] — Create a volume. Defaults to MANAGED. EXTERNAL volumes require --storage-location.
  • volume delete <catalog.schema.volume> [--yes] — Delete a volume. Prompts for confirmation unless --yes is given.

Requirements

  • Python 3.12+
  • Databricks authentication: either a Databricks CLI profile, or env vars (DATABRICKS_HOST/DATABRICKS_TOKEN), or any other unified-auth method
  • Either a Databricks all-purpose cluster (auto-started if terminated) or serverless compute enabled for the workspace
  • uv (for wheel building only)

Architecture

databricks-job-runner is layered into a thin CLI, an orchestrator, and a set of single-purpose action modules. Runner is the only class consuming projects need to touch.

cli.py          argparse + dispatch (flags -> Runner method calls)
  |
runner.py       Runner: holds config, owns the WorkspaceClient,
  |             exposes one method per subcommand
  |
  |-- config.py     RunnerConfig (frozen pydantic) + .env parser
  |-- compute.py    ClassicCluster / Serverless strategies (Protocol)
  |-- inject.py     inject_params() for submitted scripts
  |-- upload.py     workspace file + wheel upload
  |-- submit.py     compute-agnostic job submission
  |-- validate.py   workspace listing + file-existence checks
  |-- logs.py       per-task stdout/stderr retrieval
  |-- catalog.py    Unity Catalog catalog/schema/volume management
  |-- clean.py      workspace + run cleanup
  |-- errors.py     RunnerError

Layers

  • CLI (cli.py) owns all argparse setup and translates the parsed namespace into method calls on Runner. Formats RunnerError into friendly exit messages. No argparse knowledge lives outside this file.
  • Orchestration (runner.py) exposes the Runner class. RunnerConfig and the WorkspaceClient are built lazily on first access, so importing a project's cli/__init__.py doesn't touch Databricks. Each public method coordinates a single subcommand end-to-end.
  • Action modules (upload.py, submit.py, validate.py, logs.py, clean.py, catalog.py) are plain functions wrapping SDK calls. None know about argparse or Runner, keeping each unit composable and independently testable.
  • Compute strategies (compute.py) implement the Compute protocol. A strategy knows how to (1) validate that its backend is ready, (2) decorate a SubmitTask with backend-specific fields, and (3) produce the top-level environments[] list for jobs.submit. submit_job is compute-agnostic — swapping backends is a strategy change, not a conditional branch.

Design choices

  • Strategy pattern for compute. Compute is a typing.Protocol, so adding a new backend is a new frozen dataclass that matches the shape — no changes to submit_job, Runner, or the CLI. ClassicCluster and Serverless are both frozen dataclasses for value-equality and immutability.
  • Single validation point. Required-key enforcement lives entirely in RunnerConfig.from_env_file, branching on DATABRICKS_COMPUTE_MODE (only DATABRICKS_CLUSTER_ID is required when mode is cluster). Downstream code trusts the config is valid.
  • Automatic parameter injection. All non-core .env keys are passed to submitted jobs as KEY=VALUE parameters via RunnerConfig.env_params(). Scripts call inject_params() at startup to load them into os.environ, then use pydantic BaseSettings to read configuration. No callback or per-project wiring needed.
  • Wheel convention. A submitted script named exactly run_{wheel_package}.py auto-attaches the latest wheel from dist/ — as Library(whl=...) on classic, or an Environment.dependencies entry on serverless. Ties upload --wheel and submit run_xxx.py together without adding a CLI flag.
  • 12-factor .env. Pre-existing env vars override .env values, so CI/CD exports and shell overrides trump the file — matching standard .env semantics.

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

databricks_job_runner-0.3.1.tar.gz (24.8 kB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

databricks_job_runner-0.3.1-py3-none-any.whl (31.6 kB view details)

Uploaded Python 3

File details

Details for the file databricks_job_runner-0.3.1.tar.gz.

File metadata

  • Download URL: databricks_job_runner-0.3.1.tar.gz
  • Upload date:
  • Size: 24.8 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? Yes
  • Uploaded via: twine/6.1.0 CPython/3.13.12

File hashes

Hashes for databricks_job_runner-0.3.1.tar.gz
Algorithm Hash digest
SHA256 cee0629ad4251abebd2ee49da7b5a23310722e4098d1ada0a780866c5e3831b8
MD5 cef37ef1d48e13d21cb891c63ea39c62
BLAKE2b-256 36fb263e73937b174d9b5e2718bae86fb5b5a555312223aaefb0671e9b693d30

See more details on using hashes here.

Provenance

The following attestation bundles were made for databricks_job_runner-0.3.1.tar.gz:

Publisher: publish.yml on neo4j-partners/databricks-job-runner

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

File details

Details for the file databricks_job_runner-0.3.1-py3-none-any.whl.

File metadata

File hashes

Hashes for databricks_job_runner-0.3.1-py3-none-any.whl
Algorithm Hash digest
SHA256 a91af56b81315c808b818bf0df111494a77ea42bcfdd53f1eb28577c261d45d7
MD5 573cd0bd2ac50ba082a93bebf0ae477c
BLAKE2b-256 dbb89615f64d56c6a7c1bb07867d9ac0e3c2714aba5363ae1c60c9af875f0f9a

See more details on using hashes here.

Provenance

The following attestation bundles were made for databricks_job_runner-0.3.1-py3-none-any.whl:

Publisher: publish.yml on neo4j-partners/databricks-job-runner

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Pingdom Monitoring Sentry Error logging StatusPage Status page