AI Runtime quickstart

Databricks recommends starting with the Databricks CLI to run your first AI Runtime workload. The Databricks CLI lets you submit, monitor, and manage serverless GPU workloads from your terminal. If you prefer notebooks, see AI Runtime Notebook quickstart.

Requirements

  • A Azure Databricks workspace with AI Runtime enabled. See Requirements.
  • Databricks CLI version 1.24.0 or above. Earlier versions lack the databricks air commands, the code_source format, or the accelerator type listing that this quickstart uses.
  • Git installed and available on your PATH to run the Python example. The CLI uses Git to package local code.

Run your first workload

  1. Install or update the Databricks CLI, and authenticate it to your workspace. See Install or update the Databricks CLI and Authentication for the Databricks CLI.

    Tip

    To select a workspace for the walkthrough without repeating --profile, list your profiles and set a default:

    databricks auth profiles
databricks auth switch --profile <profile>
databricks auth describe

    Check that Host in the output is the workspace where you want to run the workload. For a shell-session default, use export DATABRICKS_CONFIG_PROFILE=<profile> instead. This environment variable takes precedence over the default set by auth switch.

    To override either default for one command, append --profile <other-profile>. Use the same profile for submission, status, and logs: Job Run IDs belong to the workspace where you submitted the workload. See Configuration profiles for the Databricks CLI.

    Confirm that the CLI version is 1.24.0 or above and that the air commands are available:

    databricks --version
databricks air --help

    If databricks air --help returns unknown command "air" for "databricks", your CLI version is too old. Update the CLI and try again.

  2. Create a minimal workload file. The following workload.yaml 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

    accelerator_type values have the form GPU_<GPUs per node>x<GPU model>, such as GPU_1xA10, GPU_1xH100, or GPU_8xH100. For all values, see Accelerator types. To list the accelerator types available in your workspace, run databricks air run -h config.compute.accelerator_type.

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

Tip

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

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: databricks_ai_v6
compute:
  num_accelerators: 1
  accelerator_type: GPU_1xA10
code_source:
  root_path: .
command: python "$CODE_SOURCE_PATH/main.py"

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

  • environment.version selects the 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.

Note

The CLI doesn't upload files that Git ignores, including files ignored by an enclosing repository. If the workload can't find main.py, run git check-ignore -v . to check whether the directory is ignored. See code_source in the workload YAML reference.

To bring your own dependency stack, replace the environment block with Standard environment version 6 and explicit dependencies:

environment:
  version: 6
  dependencies:
    - torch==2.11.0

CLI workloads install these dependencies on every run. This option adds PyTorch and CUDA wheel download time.

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 A10G

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 and monitor your run

The CLI prints the Job Run ID when you submit the workload. To submit in the background, omit --watch:

databricks air run --file workload.yaml

The command returns after submission while the workload continues running. Use the returned Job Run ID and the same workspace profile to check its status:

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

  • Τελευταία ενημέρωση στις