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:

YAML
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
`accelerator_type`	string	Required. An accelerator type accepted by the CLI and available in the workspace. Matching is case-sensitive.
`num_accelerators`	integer	Required. A positive accelerator count supported for the selected accelerator type.
`pool_id`	string	The ID for provisioned GPU capacity.
`priority_class`	string	The scheduling priority for workloads using provisioned GPU capacity. Accepted values are `BEST_EFFORT`, `NORMAL`, and `CRITICAL`. Requires `pool_id`.

Key	Type	Description
`accelerator_type`	string	Required. An accelerator type accepted by the CLI and available in the workspace. Matching is case-sensitive.
`num_accelerators`	integer	Required. A positive accelerator count supported for the selected accelerator type.
`pool_id`	string	The ID for provisioned GPU capacity.
`priority_class`	string	The scheduling priority for workloads using provisioned GPU capacity. Accepted values are `BEST_EFFORT`, `NORMAL`, and `CRITICAL`. Requires `pool_id`.

Accelerator types and counts depend on the workspace, cloud, and region. For hardware selection and availability, see Hardware options.

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
`version`	string or integer	A numeric environment version, such as `6`, or a Databricks AI environment identifier, such as `databricks_ai_v6`. A `databricks_ai_v` identifier requires version 5 or above. If omitted, the CLI uses the current AI Runtime default.
`dependencies`	list of strings	Package specifications passed to the serverless environment. Defaults to no additional dependencies.
`unity_catalog_image`	string	A custom image in the format `<catalog>.<schema>.<image>:<tag>`. Mutually exclusive with `version` and `dependencies`. For setup instructions, see Use custom Docker images with AI Runtime.

Key	Type	Description
`version`	string or integer	A numeric environment version, such as `6`, or a Databricks AI environment identifier, such as `databricks_ai_v6`. A `databricks_ai_v` identifier requires version 5 or above. If omitted, the CLI uses the current AI Runtime default.
`dependencies`	list of strings	Package specifications passed to the serverless environment. Defaults to no additional dependencies.
`unity_catalog_image`	string	A custom image in the format `<catalog>.<schema>.<image>:<tag>`. Mutually exclusive with `version` and `dependencies`. For setup instructions, see Use custom Docker images with AI Runtime.

For guidance about selecting and extending an environment, see Set up your environment.

Managed environment example

YAML
environment:
  version: 6
  dependencies:
    - torch

To use a custom image instead of a managed environment, specify only unity_catalog_image:

Custom image example

YAML
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
`root_path`	string	Required when `code_source` is set. The local directory to package. The directory must exist when the CLI stages the run.
`remote_volume`	string	A Unity Catalog volume location for the uploaded archive. Must start with `/Volumes/`. Defaults to a workspace location.
`include_paths`	list of strings	Paths to include relative to `root_path`. Entries must be nonempty and relative, and cannot contain `..`. Omit the key to include all selected files. An empty list is not valid.
`git`	object	Pins the code source snapshot to a local Git branch or commit.

Key	Type	Description
`root_path`	string	Required when `code_source` is set. The local directory to package. The directory must exist when the CLI stages the run.
`remote_volume`	string	A Unity Catalog volume location for the uploaded archive. Must start with `/Volumes/`. Defaults to a workspace location.
`include_paths`	list of strings	Paths to include relative to `root_path`. Entries must be nonempty and relative, and cannot contain `..`. Omit the key to include all selected files. An empty list is not valid.
`git`	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
`branch`	string	Required when `commit` is omitted. Packages the local HEAD of the branch. Selected files cannot contain uncommitted changes.
`commit`	string	Required when `branch` is omitted. Packages a commit from the local Git object store. Mutually exclusive with `branch`.

Key	Type	Description
`branch`	string	Required when `commit` is omitted. Packages the local HEAD of the branch. Selected files cannot contain uncommitted changes.
`commit`	string	Required when `branch` is omitted. Packages a commit from the local Git object store. Mutually exclusive with `branch`.

Example

The following example packages selected paths from the local main branch:

YAML
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

YAML
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

YAML
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
`level`	string	Required. A permission level, such as `CAN_VIEW` or `CAN_MANAGE`. The workspace validates the value.
`user_name`	string	The user email. Specify exactly one of `user_name`, `group_name`, or `service_principal_name` for each grant.
`group_name`	string	The group name. Specify exactly one of `user_name`, `group_name`, or `service_principal_name` for each grant.
`service_principal_name`	string	The service principal name. Specify exactly one of `user_name`, `group_name`, or `service_principal_name` for each grant.

Key	Type	Description
`level`	string	Required. A permission level, such as `CAN_VIEW` or `CAN_MANAGE`. The workspace validates the value.
`user_name`	string	The user email. Specify exactly one of `user_name`, `group_name`, or `service_principal_name` for each grant.
`group_name`	string	The group name. Specify exactly one of `user_name`, `group_name`, or `service_principal_name` for each grant.
`service_principal_name`	string	The service principal name. Specify exactly one of `user_name`, `group_name`, or `service_principal_name` for each grant.

Example

YAML
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

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

Example​

experiment_name​

compute​

command​

environment​

Managed environment example​

Custom image example​

code_source​

git​

Example​

env_variables​

Example​

idempotency_token​

max_retries​

mlflow_artifact_location​

mlflow_experiment_directory​

mlflow_run_name​

parameters​

Example​

permissions​

Example​

secrets​

Example​

timeout_minutes​

usage_policy_id​

usage_policy_name​

Additional resources​

Example

experiment_name

compute

command

environment

Managed environment example

Custom image example

code_source

git

Example

env_variables

Example

idempotency_token

max_retries

mlflow_artifact_location

mlflow_experiment_directory

mlflow_run_name

parameters

Example

permissions

Example

secrets

Example

timeout_minutes

usage_policy_id

usage_policy_name

Additional resources