Use agent skills with AI Runtime

The databricks-ai-runtime agent skill gives AI coding agents, such as Claude Code, Cursor, and Codex, instructions for working with AI Runtime through the databricks air Databricks CLI commands. With the skill installed, you can ask your agent to write a workload YAML, submit a GPU training run, and check its progress.

Requirements

• The Databricks CLI, authenticated to your workspace. The skill's examples use the flattened code_source format, which requires Databricks CLI version 1.20.0 or above.
• A workspace with AI Runtime enabled. See Requirements.

Install the skill

Run the following command:

Bash

databricks aitools install --skills-only --skills databricks-ai-runtime

Follow the prompts to choose your AI agent and installation scope. For installation options, see aitools command group.

What the skill covers

Capability	What the agent can do
Write workloads	Turn a request into a workload YAML with compute, environment, code source, and the launch command for your framework, such as torchrun, Accelerate, or DeepSpeed.
Configure runs	Set environment variables, secrets, hyperparameters, retries, and permissions.
Submit runs	Validate the config with --dry-run, then submit with databricks air run and an idempotency key.
Check and monitor runs	Get status with databricks air get and databricks air list, stream logs with databricks air logs, and point you to the MLflow run for metrics.
Cancel runs	Cancel one, several, or all of your active runs with databricks air cancel.
Debug failures	Pull logs from a failed run and surface the error.
Use custom Docker images	Build and push a custom image, then reference it in the workload.

Capability	What the agent can do
Write workloads	Turn a request into a workload YAML with compute, environment, code source, and the launch command for your framework, such as torchrun, Accelerate, or DeepSpeed.
Configure runs	Set environment variables, secrets, hyperparameters, retries, and permissions.
Submit runs	Validate the config with --dry-run, then submit with databricks air run and an idempotency key.
Check and monitor runs	Get status with databricks air get and databricks air list, stream logs with databricks air logs, and point you to the MLflow run for metrics.
Cancel runs	Cancel one, several, or all of your active runs with databricks air cancel.
Debug failures	Pull logs from a failed run and surface the error.
Use custom Docker images	Build and push a custom image, then reference it in the workload.

Additional resources

• Use the Databricks CLI with AI Runtime
• Workload YAML reference
• Use Genie Code with AI Runtime