The Lakeflow MCP Server enables building, deploying, and managing massively parallel Python compute jobs on Databricks cloud infrastructure.
Core Capabilities:
Build Python wheels - Compile Python packages into distributable wheels using uv build
Upload packages - Transfer built wheels to Databricks workspace storage
Create Databricks jobs - Define new jobs linked to uploaded wheels with specified package names and entry points
Trigger job runs - Execute jobs with custom command-line arguments for parallel processing of different data shards
Monitor executions - List and track run status for specific jobs
Retrieve run logs - Access detailed logs for debugging and analysis
Manage clusters - Create new Databricks clusters with autoscaling or use existing clusters
Secure environment variables - Pass sensitive data through Databricks Secrets without exposing them in command-line arguments
Flexible interfaces - Interact via MCP server, command-line interface, or programmatic Python API
Use Cases:
Training deep learning models at scale
Processing large datasets from S3 buckets
Running thousands of parallel simulations
Executing data pipelines with custom Python dependencies
Managing compute jobs through AI agents via MCP protocol
Enables the management of jobs on Databricks clusters by building and uploading Python wheel packages, creating job definitions, triggering runs with specific arguments, and monitoring job execution status.
Click on "Deploy Server".
Wait a few minutes for the server to deploy. Once ready, it will show a "Started" state.
In the chat, type @ followed by the MCP server name and your instructions, e.g., "@Lakeflow MCP Server launch 4 copies of this job with arguments 'fi', 'fie', 'fo', and 'fum'"
That's it! The server will respond to your query, and you can continue using it as needed.
Here is a step-by-step guide with screenshots.
This tool is an opinionated way to spawn compute jobs on the cloud. By "compute
job", I mean a massively parallel data processing job like training a deep net,
analyzing a large corpus of text that's sitting in an S3 bucket, or 1000
parallel simulations of something. To let you do these things, this package
asks you to author your code as a Python package and forces you to specify your
package dependencies in a pyproject.toml. It then uploads that package (as a
python wheel) for Databricks to execute it.
This is heavier-weight than Databrick's built-in notebook approach of editing a Python script in their web UI. In return, it lets you capture large package dependencies across repos via git submodules, and import third party packages via uv. It's lighter-weight than most other job submission systems because it doesn't require you to build docker containers. Docker containers take a large snapshot of your system, enough to build a full unix environment. These snapshots are on the order of gigabytes and difficult to upload from a home computer. For most of our work, wheels provide all the containerization we need (a wheel is a few kilobytes).
It has one more opinion: That uv is a good way to capture those Python
dependencies, with a pyproject.toml. We're also exploring
Pants as a way to manage more complex packages. Pants can also export wheels, so nothing in this design prevents us from adoptig Pants.
You can use this tool to build your wheel, upload it to Databricks, spawn copies of it each with different command line arguments, and track your jobs's status. You can also use a Databricks UI to check the state of your jobs. The tool provides several interfaces:
An MCP server so you can have AIs spawn jobs for you.
A CLI you can use from the shell.
A programmatic Python interface you can call from a Python program.
Check if you have access to Databrick by visiting this url. If you get stuck in an infinite loop where Databricks sends you a code that doesn't work, it means you don't have an account. Ask for one in #help-data-platform.
Related MCP server: Databricks MCP Server
This package assumes the package you want to run on the cluster has a
structure like this and it can be run with uv run:
my_project/
├── pyproject.toml
├── src/
└── my_package/
├── __init__.py
└── my_package_py.pyIt also assumes you've added an entry point to your pyproject.toml called
"lakeflow-task". If your package is called my_package, and it has a driver
script called my_package_py.py, and the main function in this script is called
main, you would define the "lakeflow-task" entry point like this:
[project.scripts]
lakeflow-task = "my_package.my_package_py:main"The package lakeflow_demo under this directory gives you a
concrete example of how to set up a package.
To run the package on the cluster, first build the wheel, then upload it, then tell Databricks to run it.
To make it easier to track lineage for your artifacts and your runs, the build
step embeds the current git commit hash into the wheel version (e.g.
0.1.0.devabcdef1234...). This requires all changes in your working tree to
bemust committed before building. Otherwise, the build will fail with an error
asking you to commit or stash.
Create the job from source:
You can use create-job-from-source to build, upload, and create the job.
If you don't pass a --cluster-id, a new cluster is created automatically:
uv run lakeflow.py create-job-from-source \
"my-lakeflow-job" \
"my-package" \
--pyproject-dir-path ~/my_project \
--max-workers 4This returns the job ID, which we'll use in the next step. This doesn't yet
run any jobs. It just starts a cluster that can run them. The
--max-workers argument sets the maximum number of workers for autoscaling
on the new cluster.
To use an existing cluster instead, pass --cluster-id:
uv run lakeflow.py create-job-from-source \
"my-lakeflow-job" \
"my-package" \
--pyproject-dir-path ~/my_project \
--cluster-id 0202-235755-w37hoxe8If the cluster is not running, it will be started automatically.
You can also create a cluster explicitly and reuse it across multiple jobs:
uv run lakeflow.py create-cluster --max-workers 4This returns a cluster ID you can pass to create-job-from-source or
create-job via --cluster-id.
Start the job:
uv run lakeflow.py trigger-run 123456 arg11 arg12
uv run lakeflow.py trigger-run 123456 arg21 arg22
uv run lakeflow.py trigger-run 123456 arg31 arg32This starts three instances of the job with three different sets of
arguments. You can have the arguments refer to different shards of data, and
kick off as many parallel jobs as you want. Your job can retrieve these
arguments through argv. It can retrieve its job id from the environment
variable DATABRICKS_RUN_ID.
You can also pass environment variables to the remote job without leaking secrets (like API keys) through your command line:
uv run lakeflow.py trigger-run 123456 arg1 arg2 \
--secret-env-var MY_SECRET_KEY --secret-env-var MY_OTHER_SECRET_KEYThe tool reads the values from your local environment, uploads them to
Databricks Secrets, and passes --lakeflow-secret-scope <scope> as a
command-line argument to the task. Your task can then retrieve secrets using
the Databricks dbutils API with that scope name.
Monitor the runs:
uv run lakeflow.py list-job-runs 123456This lists the runs for the given job ID.
Get Run Logs:
uv run lakeflow.py get-run-logs 987654321This retrieves the logs for a specific run ID. It takes the run returned by trigger-run.
The above illustrated how to use the CLI. You might find it easier to use the programmatic Python interface to the package instead. See run_lakeflow_demo.py for an example.
You can install this package as an MCP server. To do that, add this to ~/.cursor/mcp.json:
{
"mcpServers": {
"lakeflow": {
"command": "uv",
"args": [
"run",
"--quiet",
"--directory",
"/path/to/lakeflow-mcp",
"python",
"lakeflow.py"
],
"env": {
"DATABRICKS_HOST": "https://hims-machine-learning-staging-workspace.cloud.databricks.com",
"DATABRICKS_TOKEN": "<your token>"
}
},
...
}
}Then you can ask the agent to do things like this:
let's launch 4 copies of this job on lakeflow, and pass them the arguments "fi", "fie", "fo", and "fum" respectively.My objective was to build a job submission system that:
Python first: Could run Python packages with ~100s of Python files, and 3rd party dependencies.
Versioned artifacts: Have versioned source and output.
Workflow orchestration: Could break down its steps into tasks that could be cached, checkpointed, retried, and resumed under a workflow orchestrator.
Native-capable: Could accommodate a small amount of non-Python code code written in Rust, C++, or Dafny.
Small scale: Runs jobs on ~100s of remote workers in parallel, for ~20 engineers simultaneously.
The ideal system would use Prefect as a workflow orchestrator, on top of the existing kubernetes scaffolding we currently use to run staging and prod. There are many workflow orchestrators, but Prefect is the only one that provides all of the workflow functionality listed above. The ideal system would be a Prefect front-end VM, which scales a kubernetes cluster up and down on demand. Rolling this out would have taken some conversations with the devops team, and introducing a new tech stack to the company. The time for this will come soon, but this package is not that.
In the mean time, this package uses a tech stack the Data Engineering team is already using. They already use Databricks to run notebooks, and their expertise with Databricks helped me ramp up quickly on this solution. Databricks notebooks are small python files the DE team edits in the Databricks UI. These scripts are versioned under Git. Databricks does provide a workflow orchestrator, but the team uses Airflow for their bigger jobs. In all, the tech stack the Data Processing team already uses provides 70% of the functionality I was trying to devlop. So I decided to build on top of it instead of building an alternative to it. This package upgrades our existing tech stack to support much larger Python packages via Python wheels (not just notebooks).
You'll notice that this package doesn't provide any workflow orchestration. That's to come. Databricks provides some rudimentary workflow capabilities, which I'll gradually incorporate into this system.
Builds the Python wheel using 'uv build --wheel'.
Args:
target: The path to the directory containing pyproject.toml.
Returns:
The path to the generated wheel file.
| Name | Required | Description | Default |
|---|---|---|---|
| target | No | . |
| Name | Required | Description |
|---|---|---|
| result | Yes |
Does the description disclose side effects, auth requirements, rate limits, or destructive behavior?
With no annotations provided, the description carries full burden for behavioral disclosure. It mentions the command executed ('uv build --wheel') and the return value, but doesn't disclose critical behavioral traits like whether this is a read-only operation, what happens on failure, whether it modifies the filesystem, or any side effects. The description is insufficient for a mutation tool with zero annotation coverage.
Agents need to know what a tool does to the world before calling it. Descriptions should go beyond structured annotations to explain consequences.
Is the description appropriately sized, front-loaded, and free of redundancy?
The description is perfectly structured and concise with zero wasted words. It uses a clear main sentence followed by organized Args and Returns sections. Every sentence earns its place by providing essential information about the tool's purpose, parameters, and return value.
Shorter descriptions cost fewer tokens and are easier for agents to parse. Every sentence should earn its place.
Given the tool's complexity, does the description cover enough for an agent to succeed on first attempt?
Given that this is a build tool with no annotations but with an output schema (which handles return value documentation), the description is minimally adequate. It covers the basic purpose and parameter semantics but lacks important behavioral context about what the tool actually does beyond the command execution. For a tool that likely modifies the filesystem, more disclosure would be beneficial.
Complex tools with many parameters or behaviors need more documentation. Simple tools need less. This dimension scales expectations accordingly.
Does the description clarify parameter syntax, constraints, interactions, or defaults beyond what the schema provides?
The description adds meaningful context about the single parameter beyond what the schema provides. While schema description coverage is 0%, the description clarifies that 'target' is 'The path to the directory containing pyproject.toml', which provides essential semantic understanding not present in the schema's minimal title 'Target'. This compensates well for the schema's lack of documentation.
Input schemas describe structure but not intent. Descriptions should explain non-obvious parameter relationships and valid value ranges.
Does the description clearly state what the tool does and how it differs from similar tools?
The description clearly states the specific action ('Builds the Python wheel') and the implementation method ('using "uv build --wheel"'), distinguishing it from sibling tools like upload_wheel. It provides a precise verb+resource combination that leaves no ambiguity about the tool's function.
Agents choose between tools based on descriptions. A clear purpose with a specific verb and resource helps agents select the right tool.
Does the description explain when to use this tool, when not to, or what alternatives exist?
The description provides no guidance on when to use this tool versus alternatives like upload_wheel or trigger_run. It mentions the target parameter but doesn't explain prerequisites (e.g., needing pyproject.toml present) or when this operation is appropriate versus other build methods.
Agents often have multiple tools that could apply. Explicit usage guidance like "use X instead of Y when Z" prevents misuse.
Creates a Databricks job with the specified wheel and entry point.
Args:
job_name: The name of the job to create.
package_name: The name of the Python package.
remote_wheel_path: The remote path to the uploaded wheel file.
Returns:
The ID of the created job.
| Name | Required | Description | Default |
|---|---|---|---|
| job_name | Yes | ||
| package_name | Yes | ||
| remote_wheel_path | Yes |
| Name | Required | Description |
|---|---|---|
| result | Yes |
Does the description disclose side effects, auth requirements, rate limits, or destructive behavior?
With no annotations provided, the description carries full burden but only states it 'Creates a Databricks job' and returns an ID. It doesn't disclose behavioral traits like required permissions, whether it's idempotent, error handling, or side effects, leaving significant gaps for a mutation tool.
Agents need to know what a tool does to the world before calling it. Descriptions should go beyond structured annotations to explain consequences.
Is the description appropriately sized, front-loaded, and free of redundancy?
The description is appropriately sized with a clear purpose statement followed by structured Arg/Return sections. Every sentence adds value, but the 'Args' and 'Returns' labels are slightly redundant with schema fields, keeping it from a perfect score.
Shorter descriptions cost fewer tokens and are easier for agents to parse. Every sentence should earn its place.
Given the tool's complexity, does the description cover enough for an agent to succeed on first attempt?
Given the tool's complexity (mutation with 3 params), no annotations, and an output schema (implied by 'Returns'), the description is minimally adequate. It covers basic purpose and parameters but lacks behavioral context and usage guidelines, making it incomplete for safe agent operation.
Complex tools with many parameters or behaviors need more documentation. Simple tools need less. This dimension scales expectations accordingly.
Does the description clarify parameter syntax, constraints, interactions, or defaults beyond what the schema provides?
The description adds meaningful semantics beyond the schema: it explains that 'job_name' is for naming the job, 'package_name' refers to a Python package, and 'remote_wheel_path' is a path to an uploaded wheel file. This compensates well for the 0% schema description coverage, though it could detail format constraints.
Input schemas describe structure but not intent. Descriptions should explain non-obvious parameter relationships and valid value ranges.
Does the description clearly state what the tool does and how it differs from similar tools?
The description clearly states the action ('Creates a Databricks job') and specifies the key resources involved ('with the specified wheel and entry point'), making the purpose unambiguous. However, it doesn't explicitly differentiate from sibling tools like 'trigger_run' or 'upload_wheel', which prevents a perfect score.
Agents choose between tools based on descriptions. A clear purpose with a specific verb and resource helps agents select the right tool.
Does the description explain when to use this tool, when not to, or what alternatives exist?
The description provides no guidance on when to use this tool versus alternatives like 'trigger_run' or 'upload_wheel', nor does it mention prerequisites (e.g., needing to upload a wheel first). It implies usage through the action but lacks explicit context or exclusions.
Agents often have multiple tools that could apply. Explicit usage guidance like "use X instead of Y when Z" prevents misuse.
Lists runs for a specific job.
Args:
job_id: The ID of the job to list runs for.
| Name | Required | Description | Default |
|---|---|---|---|
| job_id | Yes |
| Name | Required | Description |
|---|---|---|
| result | Yes |
Does the description disclose side effects, auth requirements, rate limits, or destructive behavior?
With no annotations provided, the description carries the full burden of behavioral disclosure. It states the tool lists runs but doesn't describe key behaviors such as pagination, sorting, filtering (e.g., by status or date), rate limits, authentication requirements, or error handling. This leaves significant gaps in understanding how the tool operates beyond its basic purpose.
Agents need to know what a tool does to the world before calling it. Descriptions should go beyond structured annotations to explain consequences.
Is the description appropriately sized, front-loaded, and free of redundancy?
The description is concise and well-structured, with a clear purpose statement followed by a brief parameter explanation. It avoids unnecessary words and is front-loaded with the main action. However, the 'Args:' section is slightly redundant since the parameter is already implied, and more critical details (like behavioral traits) are omitted, preventing a perfect score.
Shorter descriptions cost fewer tokens and are easier for agents to parse. Every sentence should earn its place.
Given the tool's complexity, does the description cover enough for an agent to succeed on first attempt?
Given the tool's low complexity (one parameter) and the presence of an output schema (which likely describes return values), the description is minimally complete. It covers the basic purpose and parameter but lacks behavioral context, usage guidelines, and error handling information. This makes it adequate for simple use cases but insufficient for robust agent operation.
Complex tools with many parameters or behaviors need more documentation. Simple tools need less. This dimension scales expectations accordingly.
Does the description clarify parameter syntax, constraints, interactions, or defaults beyond what the schema provides?
The description adds minimal semantic context beyond the input schema. It explains that 'job_id' is 'The ID of the job to list runs for,' which clarifies the parameter's role but doesn't provide format details, validation rules, or examples. With 0% schema description coverage and only one parameter, this is adequate but not informative, meeting the baseline for such a simple case.
Input schemas describe structure but not intent. Descriptions should explain non-obvious parameter relationships and valid value ranges.
Does the description clearly state what the tool does and how it differs from similar tools?
The description clearly states the action ('Lists runs') and the target resource ('for a specific job'), making the purpose immediately understandable. It distinguishes this from sibling tools like 'create_job' or 'trigger_run' by focusing on retrieval rather than creation or execution. However, it doesn't specify the scope (e.g., all runs, recent runs, or filtered runs), which prevents a perfect score.
Agents choose between tools based on descriptions. A clear purpose with a specific verb and resource helps agents select the right tool.
Does the description explain when to use this tool, when not to, or what alternatives exist?
The description provides no guidance on when to use this tool versus alternatives. It doesn't mention prerequisites (e.g., needing an existing job ID), exclusions, or comparisons with other tools like 'trigger_run' for initiating runs. The agent must infer usage from the tool name and description alone, which is insufficient for optimal selection.
Agents often have multiple tools that could apply. Explicit usage guidance like "use X instead of Y when Z" prevents misuse.
Triggers a run of the specified job.
Args:
job_id: The ID of the job to run.
job_args: A list of Python parameters to pass to the run.
Returns:
The ID of the triggered run.
| Name | Required | Description | Default |
|---|---|---|---|
| job_id | Yes | ||
| job_args | No |
| Name | Required | Description |
|---|---|---|
| result | Yes |
Does the description disclose side effects, auth requirements, rate limits, or destructive behavior?
No annotations are provided, so the description carries full burden. It states the action ('Triggers a run') and return value, but lacks behavioral details such as permissions required, whether the run is synchronous/asynchronous, error handling, rate limits, or side effects. For a mutation tool with zero annotation coverage, this is a significant gap in transparency.
Agents need to know what a tool does to the world before calling it. Descriptions should go beyond structured annotations to explain consequences.
Is the description appropriately sized, front-loaded, and free of redundancy?
The description is well-structured and front-loaded with the core purpose, followed by clear sections for Args and Returns. Every sentence earns its place by defining parameters and output without redundancy. It's appropriately sized for a tool with two parameters and an output schema.
Shorter descriptions cost fewer tokens and are easier for agents to parse. Every sentence should earn its place.
Given the tool's complexity, does the description cover enough for an agent to succeed on first attempt?
Given a mutation tool with no annotations, 0% schema coverage, but an output schema (which handles return values), the description is partially complete. It covers purpose and parameters adequately but lacks behavioral context (e.g., execution model, errors). The output schema reduces the burden, but more guidance on usage and transparency would improve completeness for this complexity level.
Complex tools with many parameters or behaviors need more documentation. Simple tools need less. This dimension scales expectations accordingly.
Does the description clarify parameter syntax, constraints, interactions, or defaults beyond what the schema provides?
Schema description coverage is 0%, so the description must compensate. It adds meaning by explaining 'job_id' as 'The ID of the job to run' and 'job_args' as 'A list of Python parameters to pass to the run', which clarifies their roles beyond schema types. However, it doesn't detail format constraints (e.g., job_id source, job_args syntax), leaving some ambiguity. Baseline 3 is appropriate as it adds value but not fully comprehensive.
Input schemas describe structure but not intent. Descriptions should explain non-obvious parameter relationships and valid value ranges.
Does the description clearly state what the tool does and how it differs from similar tools?
The description clearly states the tool's purpose with a specific verb ('Triggers') and resource ('run of the specified job'). It distinguishes from siblings like 'create_job' (creation vs. execution) and 'list_job_runs' (listing vs. triggering), though it doesn't explicitly mention these distinctions. The purpose is unambiguous but could be more explicit about sibling differentiation.
Agents choose between tools based on descriptions. A clear purpose with a specific verb and resource helps agents select the right tool.
Does the description explain when to use this tool, when not to, or what alternatives exist?
The description provides no guidance on when to use this tool versus alternatives. It doesn't mention prerequisites (e.g., job must exist), exclusions, or comparisons to siblings like 'list_job_runs' for monitoring runs. Usage is implied by the purpose but lacks explicit context or decision criteria.
Agents often have multiple tools that could apply. Explicit usage guidance like "use X instead of Y when Z" prevents misuse.
Uploads a local wheel file to the Databricks workspace.
Args:
local_path: The local path to the wheel file.
Returns:
The full remote path of the uploaded wheel.
| Name | Required | Description | Default |
|---|---|---|---|
| local_path | Yes |
| Name | Required | Description |
|---|---|---|
| result | Yes |
Does the description disclose side effects, auth requirements, rate limits, or destructive behavior?
With no annotations provided, the description carries the full burden of behavioral disclosure. It mentions the upload action and return value but omits critical details like required permissions, file size limits, error handling, or whether the operation is idempotent. For a mutation tool, this leaves significant gaps in understanding its behavior.
Agents need to know what a tool does to the world before calling it. Descriptions should go beyond structured annotations to explain consequences.
Is the description appropriately sized, front-loaded, and free of redundancy?
The description is appropriately sized and front-loaded, with the core purpose stated first, followed by structured Args and Returns sections. Every sentence earns its place, though the formatting is slightly verbose for a single parameter.
Shorter descriptions cost fewer tokens and are easier for agents to parse. Every sentence should earn its place.
Given the tool's complexity, does the description cover enough for an agent to succeed on first attempt?
Given the tool's complexity (a mutation with no annotations) and the presence of an output schema (which covers return values), the description is partially complete. It explains the upload action and parameter but lacks behavioral context like side effects or error conditions, making it adequate but with clear gaps.
Complex tools with many parameters or behaviors need more documentation. Simple tools need less. This dimension scales expectations accordingly.
Does the description clarify parameter syntax, constraints, interactions, or defaults beyond what the schema provides?
The schema description coverage is 0%, so the description must compensate. It adds meaning by explaining that 'local_path' refers to 'The local path to the wheel file', clarifying the parameter's purpose beyond the schema's basic type. However, it doesn't detail format constraints or examples, leaving some ambiguity.
Input schemas describe structure but not intent. Descriptions should explain non-obvious parameter relationships and valid value ranges.
Does the description clearly state what the tool does and how it differs from similar tools?
The description clearly states the specific action ('Uploads') and resource ('a local wheel file to the Databricks workspace'), distinguishing it from sibling tools like 'build_wheel' (creation) or 'list_job_runs' (querying). It precisely defines what the tool does without ambiguity.
Agents choose between tools based on descriptions. A clear purpose with a specific verb and resource helps agents select the right tool.
Does the description explain when to use this tool, when not to, or what alternatives exist?
The description provides no guidance on when to use this tool versus alternatives like 'build_wheel' or 'create_job', nor does it mention prerequisites such as needing an existing wheel file. It lacks explicit usage context or exclusions, relying solely on the tool's name and basic purpose.
Agents often have multiple tools that could apply. Explicit usage guidance like "use X instead of Y when Z" prevents misuse.
Recent tool additions, removals, and schema changes observed during successful MCP inspections.
v0.1.0build_wheelcreate_joblist_job_runstrigger_runupload_wheelScored across 5 tools
Each tool has a clearly distinct purpose with no overlap: building a wheel, creating a job, listing job runs, triggering a run, and uploading a wheel. The descriptions specify unique actions on different resources (wheel files vs. Databricks jobs), making misselection unlikely.
All tool names follow a consistent verb_noun pattern (e.g., build_wheel, create_job, list_job_runs, trigger_run, upload_wheel). The verbs are clear and descriptive, and there are no deviations in naming style across the set.
With 5 tools, this server is well-scoped for its purpose of managing Databricks job workflows with Python wheels. Each tool earns its place by covering distinct steps in the process, from building and uploading wheels to job creation and execution.
The tool set provides strong coverage for core workflows: building, uploading, job creation, triggering runs, and monitoring runs. Minor gaps exist, such as no tools for updating or deleting jobs, but agents can likely work around this for basic operations.
Give your AI agents the tools to build, manage, and run automation workflows.
Create and manage scheduled, guarded AI agent jobs with built-in quality control and 900+ connectors
Agentic CI operations for build inspection, failure diagnosis, and runner troubleshooting.
Deploy and manage your apps, databases, storage, and scheduled jobs from your AI agent
