databricks-job-runner 0.3.0__tar.gz

databricks_job_runner-0.3.0/PKG-INFO

@@ -0,0 +1,289 @@
+Metadata-Version: 2.4
+Name: databricks-job-runner
+Version: 0.3.0
+Summary: Reusable CLI for uploading, submitting, validating, fetching logs, and cleaning Databricks job runs
+Author: Ryan Knight
+Author-email: Ryan Knight <ryan.knight@neo4j.com>
+License-Expression: MIT
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Developers
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Typing :: Typed
+Requires-Dist: databricks-sdk
+Requires-Dist: pydantic>=2
+Requires-Python: >=3.12
+Project-URL: Repository, https://github.com/neo4j-partners/databricks-job-runner
+Description-Content-Type: text/markdown
+
+# databricks-job-runner
+
+Reusable CLI for uploading, submitting, and cleaning Databricks job runs.
+
+Wraps the [Databricks Python SDK](https://docs.databricks.com/dev-tools/sdk-python.html) 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
+
+```bash
+uv add databricks-job-runner
+```
+
+Or with pip:
+
+```bash
+pip install databricks-job-runner
+```
+
+For local development against a checkout:
+
+```toml
+# 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:
+>
+> ```toml
+> [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`**
+
+```python
+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`**
+
+```python
+from cli import runner
+runner.main()
+```
+
+Then run from the project root:
+
+```bash
+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/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`
+
+```python
+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`
+
+```python
+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](https://docs.databricks.com/dev-tools/cli/index.html), or env vars (`DATABRICKS_HOST`/`DATABRICKS_TOKEN`), or any other [unified-auth](https://docs.databricks.com/dev-tools/auth/) method
+- Either a Databricks all-purpose cluster (auto-started if terminated) or serverless compute enabled for the workspace
+- [uv](https://docs.astral.sh/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.

databricks_job_runner-0.3.0/README.md

@@ -0,0 +1,270 @@
+# databricks-job-runner
+
+Reusable CLI for uploading, submitting, and cleaning Databricks job runs.
+
+Wraps the [Databricks Python SDK](https://docs.databricks.com/dev-tools/sdk-python.html) 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
+
+```bash
+uv add databricks-job-runner
+```
+
+Or with pip:
+
+```bash
+pip install databricks-job-runner
+```
+
+For local development against a checkout:
+
+```toml
+# 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:
+>
+> ```toml
+> [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`**
+
+```python
+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`**
+
+```python
+from cli import runner
+runner.main()
+```
+
+Then run from the project root:
+
+```bash
+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/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`
+
+```python
+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`
+
+```python
+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](https://docs.databricks.com/dev-tools/cli/index.html), or env vars (`DATABRICKS_HOST`/`DATABRICKS_TOKEN`), or any other [unified-auth](https://docs.databricks.com/dev-tools/auth/) method
+- Either a Databricks all-purpose cluster (auto-started if terminated) or serverless compute enabled for the workspace
+- [uv](https://docs.astral.sh/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.

databricks_job_runner-0.3.0/pyproject.toml

@@ -0,0 +1,32 @@
+[project]
+name = "databricks-job-runner"
+version = "0.3.0"
+description = "Reusable CLI for uploading, submitting, validating, fetching logs, and cleaning Databricks job runs"
+readme = "README.md"
+license = "MIT"
+authors = [
+    { name = "Ryan Knight", email = "ryan.knight@neo4j.com" }
+]
+requires-python = ">=3.12"
+dependencies = [
+    "databricks-sdk",
+    "pydantic>=2",
+]
+classifiers = [
+    "Development Status :: 4 - Beta",
+    "Intended Audience :: Developers",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.12",
+    "Programming Language :: Python :: 3.13",
+    "Typing :: Typed",
+]
+
+[project.urls]
+Repository = "https://github.com/neo4j-partners/databricks-job-runner"
+
+[build-system]
+requires = ["uv_build>=0.9.3,<0.10.0"]
+build-backend = "uv_build"
+
+[tool.uv.workspace]
+members = ["examples/serverless_smoke"]

databricks_job_runner-0.3.0/src/databricks_job_runner/__init__.py

@@ -0,0 +1,16 @@
+"""Reusable CLI for uploading, submitting, validating, fetching logs, and cleaning Databricks job runs."""
+
+from databricks_job_runner.compute import ClassicCluster, Compute, Serverless
+from databricks_job_runner.config import RunnerConfig
+from databricks_job_runner.errors import RunnerError
+from databricks_job_runner.runner import BuildParamsFn, Runner
+
+__all__ = [
+    "BuildParamsFn",
+    "ClassicCluster",
+    "Compute",
+    "Runner",
+    "RunnerConfig",
+    "RunnerError",
+    "Serverless",
+]

databricks_job_runner-0.3.0/src/databricks_job_runner/clean.py

@@ -0,0 +1,40 @@
+"""Clean up remote workspace directories and job runs."""
+
+from __future__ import annotations
+
+from databricks.sdk import WorkspaceClient
+from databricks.sdk.errors import NotFound
+from databricks.sdk.service.jobs import RunType
+
+
+def clean_workspace(ws: WorkspaceClient, workspace_dir: str) -> None:
+    """Delete the remote workspace directory recursively."""
+    print(f"Deleting remote workspace: {workspace_dir}")
+    try:
+        ws.workspace.delete(path=workspace_dir, recursive=True)
+        print("  Done.")
+    except NotFound:
+        print("  Directory does not exist or already deleted.")
+
+
+def clean_runs(ws: WorkspaceClient, run_name_prefix: str) -> None:
+    """Find and delete all one-time job runs whose name starts with *prefix*."""
+    print(f"Finding job runs matching '{run_name_prefix}*'...")
+
+    deleted = 0
+    for run in ws.jobs.list_runs(run_type=RunType.SUBMIT_RUN, expand_tasks=False):
+        run_name = run.run_name or ""
+        run_id = run.run_id
+        if run_id is None or not run_name.startswith(run_name_prefix):
+            continue
+        print(f"  Deleting run {run_id} ({run_name})")
+        try:
+            ws.jobs.delete_run(run_id)
+            deleted += 1
+        except NotFound:
+            print(f"    Run {run_id} already deleted.")
+
+    if deleted:
+        print(f"  Deleted {deleted} run(s).")
+    else:
+        print("  No matching runs found.")