Databricks CLI quickstart for AI Runtime

Important

This feature is in Public Preview.

The Databricks CLI for AI Runtime lets you submit, monitor, and manage serverless GPU workloads from your terminal. This quickstart shows how to configure and run an AI Runtime workload, stream logs, and inspect the run.

Requirements

Confirm that the air commands are available in the Databricks CLI:

databricks air --help

Tip

To get help writing your workload config and managing runs, install the databricks-ai-runtime skill for your AI agent. See Use an AI agent.

Run a minimal workload

Create a workload.yaml file that requests one A10 GPU and runs nvidia-smi to print the GPU name:

experiment_name: air-quickstart
compute:
  num_accelerators: 1
  accelerator_type: GPU_1xA10
command: nvidia-smi --query-gpu=name --format=csv,noheader

Submit the workload and stream its logs with databricks air run:

databricks air run --file workload.yaml --watch

The CLI returns the Job Run ID and a link to the run in your workspace. --watch keeps the terminal attached and streams logs until the workload finishes.

Submitting experiment: air-quickstart
Submitted workload with Job Run ID: 783940125630184
View job run at: https://example.cloud.databricks.com/jobs/runs/783940125630184?o=1234567890123456

Monitoring run and streaming logs...
...
Environment setup completed at Thu Oct 1 14:23:45 UTC 2026
NVIDIA A10

Run your own Python code

To upload and run local Python code, create an air-quickstart directory with the following files:

air-quickstart/
├── workload.yaml
└── main.py

Tip

To start from a runnable project instead of creating these files manually, clone the AI Runtime cookbook repository and change to the CLI getting-started directory:

git clone https://github.com/databricks-solutions/ai-runtime-cookbook.git
cd ai-runtime-cookbook/getting-started/air-cli

Use the following minimal Python code to verify that AI Runtime runs the uploaded code on the GPU:

import torch

print("Hello from Python on AI Runtime")
print(f"GPU: {torch.cuda.get_device_name(0)}")

To extend this example, you can replace main.py with your own training, batch inference, or other AI application.

Use the following configuration for workload.yaml:

experiment_name: python-quickstart
environment:
  version: 6
  dependencies:
    - torch==2.11.0
compute:
  num_accelerators: 1
  accelerator_type: GPU_1xH100
code_source:
  root_path: .
command: python "$CODE_SOURCE_PATH/main.py"

This configuration uses the following settings and runs main.py on one H100 GPU:

  • environment.version selects Standard environment version 6.
  • environment.dependencies installs PyTorch 2.11.0 into the selected environment.
  • root_path: . uploads the directory containing workload.yaml. At runtime, $CODE_SOURCE_PATH points to the uploaded directory.

Submit the workload from the directory that contains workload.yaml:

databricks air run --file workload.yaml --watch

The streamed logs include:

...
Hello from Python on AI Runtime
GPU: NVIDIA H100 80GB HBM3

To run complete training, fine-tuning, batch inference, and Ray workloads, clone the AI Runtime Cookbook and run a recipe.

To configure environments, code sources, compute, and other workload settings, see Workload YAML reference.

Inspect a run

The CLI prints the Job Run ID when you submit the workload. Use it to inspect the run:

databricks air get <job-run-id>

databricks air get shows the run status, configuration, and timing. It also links to the Job run and, when available, the MLflow run.

To stream logs after submission or resume an interrupted stream, run:

databricks air logs <job-run-id>

The CLI also provides commands to list and inspect runs, stream logs, and cancel active runs. See the air command reference.

Reference

Additional resources

  • Sist oppdatert