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

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/:

    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:

    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:

    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:

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

Additional resources

  • Last updated on