Workload YAML reference
This page describes the workload YAML fields accepted by databricks air run. For command syntax and flags, see databricks air run. To submit a workload, see AI Runtime quickstart.
Every workload requires experiment_name, compute, and command.
Example
The following example shows a sample workload configuration:
experiment_name: my-training
compute:
accelerator_type: GPU_1xA10
num_accelerators: 1
environment:
version: 6
dependencies:
- torch
code_source:
root_path: .
command: python $CODE_SOURCE_PATH/train.py
experiment_name
Type: String
Required. The experiment_name key is the MLflow experiment name.
If an experiment with this name does not exist, AI Runtime creates it. Otherwise, AI Runtime adds the new run to the existing experiment.
compute
Type: Map
Required. The compute object defines the accelerator configuration.
Key | Type | Description |
|---|---|---|
| string | Required. One of the accelerator types available in the workspace. Matching is case-sensitive. |
| integer | Required. The total number of GPUs to allocate. Must be a multiple of the GPUs per node for the selected accelerator type. |
Accelerator types
Each accelerator_type value has the form GPU_<GPUs per node>x<GPU model>. For example, GPU_8xH100 is a node with eight H100 GPUs. Use the exact value from the following table, including the GPU_ prefix and capitalization:
| GPU | GPUs per node | Valid |
|---|---|---|---|
| NVIDIA A10, 24 GB | 1 |
|
| NVIDIA H100, 80 GB | 1 |
|
| NVIDIA H100, 80 GB per GPU | 8 | A multiple of 8, such as |
| NVIDIA B300, 288 GB per GPU (Public Preview) | 8 | A multiple of 8, such as |
The number of nodes provisioned is num_accelerators divided by the GPUs per node. For example, accelerator_type: GPU_8xH100 with num_accelerators: 16 provisions two 8xH100 nodes.
Accelerator availability depends on the workspace, cloud, and region. To list the accelerator types available in your workspace, run:
databricks air run -h config.compute.accelerator_type
For help choosing hardware, see Hardware.
command
Type: String
Required. The command key contains a nonempty shell command or script. Paths in the command refer to remote file paths in the workload container. When code_source is set, $CODE_SOURCE_PATH contains the absolute path to the extracted code.
environment
Type: Map
The environment object selects a managed environment with optional dependencies or a custom Unity Catalog image.
Key | Type | Description |
|---|---|---|
| string or integer | A numeric environment version, such as |
| list of strings | Package specifications passed to the serverless environment. Defaults to no additional dependencies. |
| string | A custom image in the format |
For guidance about selecting and extending an environment, see Set up your AI Runtime environment.
Managed environment example
environment:
version: 6
dependencies:
- torch
To use a custom image instead of a managed environment, specify only unity_catalog_image:
Custom image example
environment:
unity_catalog_image: main.ml.training:v1
code_source
Type: Map
The code_source object packages local code for the workload.
Key | Type | Description |
|---|---|---|
| string | Required when |
| string | A Unity Catalog volume location for the uploaded archive. Must start with |
| list of strings | Paths to include relative to |
| object | Pins the code source snapshot to a local Git branch or commit. |
If you omit git, the CLI packages the working tree, including uncommitted and untracked files that are not excluded by Git ignore rules. A relative root_path resolves from the directory that contains the workload YAML file.
git
Type: Map
The git object supports the following keys:
Key | Type | Description |
|---|---|---|
| string | Required when |
| string | Required when |
Example
The following example packages selected paths from the local main branch:
code_source:
root_path: ./src
include_paths:
- training
- configs
git:
branch: main
env_variables
Type: Map of strings
The env_variables key contains plain environment variable names and string values. A name cannot also appear in secrets.
Example
env_variables:
BATCH_SIZE: '32'
LEARNING_RATE: '0.001'
idempotency_token
Type: String
The idempotency_token key identifies repeated submissions. A repeated submission with the same token returns the existing run. The CLI generates a UUID when this key is omitted. The --idempotency-key command flag takes precedence.
max_retries
Type: Integer
The max_retries key specifies the number of retries after the initial attempt. The value must be nonnegative. Default: 3.
mlflow_artifact_location
Type: String
The mlflow_artifact_location key contains a dbfs:/ URI or /Volumes/ path. If you omit this key, a new experiment uses the default DBFS location, and an existing experiment uses its configured artifact location. For an existing experiment, the value must match the experiment's artifact location.
mlflow_experiment_directory
Type: String
The mlflow_experiment_directory key contains the workspace directory for the experiment. The path must start with /Workspace.
mlflow_run_name
Type: String
The mlflow_run_name key contains the run name. If you omit this key, the run name defaults to the value of experiment_name.
parameters
Type: Map
The parameters key contains free-form nested values. The CLI serializes the values to hyperparameters.yaml with the other launch artifacts.
Example
parameters:
learning_rate: 0.001
model:
hidden_size: 1024
permissions
Type: List of maps
The permissions key contains grants applied after submission. If a grant fails, the run remains created and the CLI prints a warning.
Each grant supports the following keys:
Key | Type | Description |
|---|---|---|
| string | Required. A permission level, such as |
| string | The user email. Specify exactly one of |
| string | The group name. Specify exactly one of |
| string | The service principal name. Specify exactly one of |
Example
permissions:
- user_name: trainer@example.com
level: CAN_MANAGE
- group_name: training-team
level: CAN_VIEW
secrets
Type: Map of strings
The secrets key contains environment variable names mapped to secret references in the exact scope/key format. The scope and key must be nonempty. A name cannot also appear in env_variables.
Example
secrets:
HF_TOKEN: my-scope/hf-token
For information about creating and managing secrets, see Secret management.
timeout_minutes
Type: Integer
The timeout_minutes key sets the wall-clock limit for the complete run, including retries. The value must be positive. If omitted, the workload uses the AI Runtime default.
usage_policy_id
Type: String
The usage_policy_id key contains the canonical UUID of an existing serverless usage policy. Mutually exclusive with usage_policy_name.
usage_policy_name
Type: String
The usage_policy_name key contains the name of an existing serverless usage policy. Mutually exclusive with usage_policy_id.
For information about creating usage policies, see Attribute usage with serverless usage policies.