Skip to main content
Last updated on

Databricks CLI quickstart for AI Runtime

Preview

This feature is in Public Preview.

Use the Databricks CLI to submit and monitor GPU workloads on AI Runtime. Start with a minimal GPU check, then run your own Python code and inspect the run.

Requirements​

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

Bash
databricks air --help

Run a minimal workload​

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

workload.yaml
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:

Bash
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.

Output
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

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

main.py
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:

workload.yaml
experiment_name: python-quickstart
environment:
  version: databricks_ai_v6
compute:
  num_accelerators: 1
  accelerator_type: GPU_1xH100
code_source:
  type: snapshot
  snapshot:
    root_path: .
command: python "$CODE_SOURCE_PATH/main.py"

This configuration uses the following settings:

  • environment.version selects the built-in Databricks AI environment version 6, which includes PyTorch.
  • root_path: . uploads the directory containing workload.yaml. At runtime, $CODE_SOURCE_PATH points to the uploaded directory.

AI Runtime runs main.py on one H100 GPU.

From the air-quickstart directory, submit the workload:

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

The streamed logs include:

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

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

Inspect a run​

Use the Job Run ID from the submission output to inspect the run:

Bash
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:

Bash
databricks air logs <job-run-id>

The CLI provides additional commands to list, inspect, cancel, and retrieve logs from workloads. See the air command reference.

Reference​

Additional resources​

On this page