databricks-job-runner 0.3.0

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

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 five CLI subcommands — upload, submit, validate, logs, and clean — 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, RunnerConfig

def build_params(config: RunnerConfig, script: str) -> list[str]:
    """Turn .env values into CLI args for the submitted script."""
    params: list[str] = []
    if config.extras.get("NEO4J_URI") and config.extras.get("NEO4J_PASSWORD"):
        params += ["--neo4j-uri", config.extras["NEO4J_URI"],
                   "--neo4j-password", config.extras["NEO4J_PASSWORD"]]
    return params

runner = Runner(
    run_name_prefix="my_project",
    build_params=build_params,
    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

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 passed to your build_params callback.

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/[e-mail]/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/[e-mail]/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,
    build_params: BuildParamsFn | None = None,
    project_dir: Path | str | None = None,
    wheel_package: str | None = None,
)

Parameter	Description
run_name_prefix	Prefix for job run names and cleanup filtering
build_params	Callback (config: RunnerConfig) -> list[str] that builds CLI args from typed config
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

BuildParamsFn

type BuildParamsFn = Callable[[RunnerConfig, str], list[str]]

Type alias for the build_params callback. The second argument is the script name being submitted, enabling per-script parameter injection.

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

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)
  |-- 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
  |-- 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) 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.
• build_params callback. Project-specific config stays in the consumer's callback rather than the runner's .env schema. Core DATABRICKS_* keys are typed on RunnerConfig; everything else falls into RunnerConfig.extras for the callback to read.
• 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.