Convert an AI Runtime workload to a bundle

Convert an existing AI Runtime workload YAML into Declarative Automation Bundles to manage it as a persistent job, add a schedule, and compose it with preprocessing tasks. Use databricks air convert-to-dabs to generate the bundle, or use the field mappings on this page to convert it manually.

Requirements

• An existing workload YAML and its training code.
• The latest Databricks CLI. Update an existing installation before converting.
• To deploy and run the generated job, a supported workspace with the AI Runtime preview enabled and the CLI authenticated to it. See Schedule GPU workloads and compose tasks.

Convert a workload automatically

If you already run a workload with databricks air run --file train.yaml, convert its workload YAML with databricks air convert-to-dabs.

1. Start with your existing workload YAML and code. For example, train.yaml can point to a training script in src/:

   YAML

   experiment_name: my-training
   environment:
     dependencies:
       - torch
   compute:
     num_accelerators: 1
     accelerator_type: GPU_1xA10
   code_source:
     type: snapshot
     snapshot:
       root_path: src
   command: python $CODE_SOURCE_PATH/train.py

2. From the directory containing train.yaml, run the conversion:

   Bash

   databricks air convert-to-dabs train.yaml

   The command creates databricks.yml and a generated_artifacts/ directory alongside the YAML. It translates the configuration locally. The bundle uploads the code when you deploy it. For other directory layouts and overwrite options, see Converter paths and generated files.

3. Validate, deploy, and run the generated job:

   Bash

   databricks bundle validate
   databricks bundle deploy
   databricks bundle run my-training --no-wait

The converter creates a job with one GPU task. To add a schedule and preprocessing tasks to databricks.yml, follow Schedule GPU workloads and compose tasks. The deployed job persists between runs. Use databricks bundle destroy to remove it when you are done.

Converter paths and generated files

databricks air convert-to-dabs train.yaml writes the bundle to the input YAML's directory by default. It resolves a relative code_source.snapshot.root_path against the YAML's directory.

For a snapshot code source, the resolved source directory must be strictly inside the bundle output directory. With the default output location, root_path: . resolves to the bundle root and is rejected. Either use a source subdirectory, such as src, or select an output directory that contains the source directory.

For example, if /project/training/train.yaml uses root_path: ., the following command writes the bundle to /project and packages /project/training:

Bash

databricks air convert-to-dabs /project/training/train.yaml --output-dir /project

--output-dir changes where the bundle is written. It does not copy or move the source code. Run the subsequent databricks bundle commands from the output directory.

If generated files already exist, conversion stops. Use --force to overwrite them. This replaces the generated bundle configuration, including any manual edits to that file.

The converter writes the following files:

• databricks.yml: the bundle configuration, including the GPU task and code artifact.
• generated_artifacts/command.sh: the command script.
• generated_artifacts/training_config.yaml: the workload configuration.
• generated_artifacts/hyperparameters.yaml: written when parameters is set.
• generated_artifacts/env_vars.json and generated_artifacts/secret_env_vars.json: written when environment variables or secrets are set.

Keep the generated files together. The bundle uploads them during deployment. When hyperparameters.yaml is present beside command.sh, the runtime sets HYPERPARAMETERS_PATH to that file.

Map workload fields manually

Use the following mappings when manually converting to a bundle from an AI Runtime workload YAML. For automatic conversion, follow Convert a workload automatically. For task field definitions, see the AI Runtime task reference.

Workload YAML field	Bundle setting
experiment_name	ai_runtime_task.experiment
command	Move the command into a shell script and reference it with ai_runtime_task.deployments[].command_path.
compute.accelerator_type	ai_runtime_task.deployments[].compute.accelerator_type
compute.num_accelerators	ai_runtime_task.deployments[].compute.accelerator_count
code_source	Package the code with a tgz artifact and reference it with ai_runtime_task.code_source_path, or use an uploaded workspace or volume path.
environment.dependencies	The job's environments[].spec.dependencies.
environment.version	environments[].spec.environment_version for a Standard environment, or environments[].spec.base_environment for a Databricks AI environment.
environment.docker_image.url	ai_runtime_task.docker_image_url
env_variables, secrets	The job's environment_variables and the task's environment_variables_key.
parameters	A hyperparameters.yaml file alongside the script referenced by command_path. The runtime sets HYPERPARAMETERS_PATH to this file.
max_retries	The task's max_retries.
timeout_minutes	The task's timeout_seconds, multiplying minutes by 60.
mlflow_run_name	ai_runtime_task.mlflow_run
mlflow_experiment_directory	ai_runtime_task.mlflow_experiment_directory
mlflow_artifact_location	ai_runtime_task.mlflow_artifact_location

Workload YAML field	Bundle setting
experiment_name	ai_runtime_task.experiment
command	Move the command into a shell script and reference it with ai_runtime_task.deployments[].command_path.
compute.accelerator_type	ai_runtime_task.deployments[].compute.accelerator_type
compute.num_accelerators	ai_runtime_task.deployments[].compute.accelerator_count
code_source	Package the code with a tgz artifact and reference it with ai_runtime_task.code_source_path, or use an uploaded workspace or volume path.
environment.dependencies	The job's environments[].spec.dependencies.
environment.version	environments[].spec.environment_version for a Standard environment, or environments[].spec.base_environment for a Databricks AI environment.
environment.docker_image.url	ai_runtime_task.docker_image_url
env_variables, secrets	The job's environment_variables and the task's environment_variables_key.
parameters	A hyperparameters.yaml file alongside the script referenced by command_path. The runtime sets HYPERPARAMETERS_PATH to this file.
max_retries	The task's max_retries.
timeout_minutes	The task's timeout_seconds, multiplying minutes by 60.
mlflow_run_name	ai_runtime_task.mlflow_run
mlflow_experiment_directory	ai_runtime_task.mlflow_experiment_directory
mlflow_artifact_location	ai_runtime_task.mlflow_artifact_location

Additional resources

• Schedule GPU workloads and compose tasks
• Configure AI Runtime bundle tasks
• Workload YAML reference