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

    Confirm that the air commands are available:

    databricks air --help

  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

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

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