Productionize training workloads

Important

Use DABs to define an AI Runtime training workload as code. Keep it in source control, deploy it across environments, schedule it, and compose it with other tasks. This page covers the bring-your-own-training path, where an ai_runtime_task runs your own command against a directory of code on serverless GPU compute.

This is a different task from running a notebook on serverless GPU through a bundle. For the basic notebook-on-GPU bundle example, see Jobs API and Declarative Automation Bundles.

Requirements

A workspace with AI Runtime enabled. See Requirements.
The Databricks CLI (command-line interface) installed and configured to deploy bundles.

Define an AI Runtime task in a bundle

An ai_runtime_task names an experiment, points at your training code with code_source_path, and declares one deployment: the command to run and the GPU to run it on. Add it to a job in your bundle:

resources:
  jobs:
    train:
      tasks:
        - task_key: train
          ai_runtime_task:
            experiment: my-experiment
            code_source_path: ./dist/code.tgz
            deployments:
              - command_path: ./command.sh
                compute:
                  accelerator_type: GPU_1xA10
                  accelerator_count: 1

code_source_path points at your packaged training code, and command_path is the script the task runs. Retries, timeouts, and permissions are set on the task and job the same way as any Azure Databricks job, so your existing bundle practices carry over. For how to package and reference your code, see Ship your training code.

`ai_runtime_task` fields

Field	Type	Description
`experiment`	String	Required. The MLflow experiment name for the run. See Experiment tracking and observability.
`code_source_path`	String	The training code to run: the output file of a packaged `tgz` artifact, or a `/Workspace` or `/Volumes` path to code that is already uploaded. See Ship your training code.
`deployments`	Sequence	Required. A single deployment describing the command and the compute to run it on. Each entry contains `command_path`, `compute`, and an optional `name`.
`deployments[].command_path`	String	Required. The script the task runs on each node.
`deployments[].compute.accelerator_type`	String	Required. The GPU type, for example `GPU_1xA10`, `GPU_1xH100`, or `GPU_8xH100`.
`deployments[].compute.accelerator_count`	Integer	Required. The total number of GPUs across all nodes—a multiple of the per-node count encoded in `accelerator_type`.
`deployments[].name`	String	An optional name for the deployment, used in logs and the UI.
`docker_image_url`	String	An optional custom Docker image to run the command in, instead of the managed environment. See Use custom Docker images.
`mlflow_run`	String	An optional display name for the MLflow run.
`mlflow_experiment_directory`	String	An optional workspace directory under which the experiment is created. Must start with `/Workspace`. Set this when running as a service principal that has no default user directory.
`mlflow_artifact_location`	String	An optional root location for MLflow artifacts, for example a `/Volumes/<catalog>/<schema>/<volume>/…` path. Must match an existing experiment's artifact location or be omitted.

Set retries, timeouts, permissions, and the environment (environment_key) on the task and job—not inside ai_runtime_task. For the full task reference, see AI Runtime task.

Configure the hardware accelerator

Set accelerator_type to the GPU your workload needs, and accelerator_count to the total number of GPUs. The count is a multiple of the number of GPUs per node: 1 for GPU_1xA10 and GPU_1xH100, and 8 for GPU_8xH100. A count larger than the per-node size runs the task across multiple nodes—for example, GPU_8xH100 with accelerator_count: 16 runs on two nodes. For guidance on choosing an accelerator, see Hardware options.

Note

For multi-node runs, AI Runtime runs your command on every node and populates the standard distributed-training environment variables in the task environment—NUM_NODES, WORLD_SIZE, LOCAL_WORLD_SIZE, MASTER_ADDR, and MASTER_PORT. Read them from your command (for example, a torchrun launch); you do not set them in the bundle.

Set the environment and dependencies

Declare an environments block on the job and reference it from the task with environment_key. AI Runtime installs the listed dependencies before your command runs:

resources:
  jobs:
    train:
      tasks:
        - task_key: train
          environment_key: default
          ai_runtime_task:
            # experiment, code_source_path, and deployments as above
      environments:
        - environment_key: default
          spec:
            environment_version: '5'
            dependencies:
              - numpy

For the available environments, see Set up your environment.

Ship your training code

code_source_path tells the task where your training code is. It takes one of two forms:

A packaged tgz artifact — declare an artifact and point code_source_path at its output file. Azure Databricks builds the tarball and uploads it on databricks bundle deploy. This is how you ship code from a local project directory or a committed Git revision.
A workspace or volume path — code that is already uploaded, used as-is.

Packaged artifact

Declare a tgz artifact and point code_source_path at its output file. On databricks bundle deploy, the CLI builds the tarball, uploads it, and the task extracts it and runs your command against it:

artifacts:
  code:
    type: tgz
    path: .
    include: [src]
    files:
      - source: ./dist/code.tgz
resources:
  jobs:
    train:
      tasks:
        - task_key: train
          ai_runtime_task:
            code_source_path: ./dist/code.tgz

Use include to package files from your working tree, or git to snapshot a committed branch or commit.

Workspace or volume path

To use code that is already uploaded, set code_source_path to a /Workspace/… or /Volumes/… path. Azure Databricks uses the path as-is and packages nothing.

The tgz artifact fields:

Field	Description
`type`	`tgz` builds a gzipped tarball from source files, instead of running a `build` command.
`path`	The base directory to package. `include` paths and the archive's entry names are relative to it.
`include`	A list of subpaths of `path` to package. Omit to package all of `path`. Honors `.gitignore`; the bundle-wide `sync.include` and `sync.exclude` do not apply. Alternative to a `build` command.
`git`	Snapshot a committed Git ref instead of the working tree. Set `git.branch` or `git.commit` (`commit` wins when both are set). Alternative to a `build` command.
`files[].source`	The path of the built tarball. Point `code_source_path` at this.

Note

AI Runtime extracts your code to a directory and exposes it as the CODE_SOURCE_PATH environment variable. Reference it from your command so relative paths resolve, for example cd "$CODE_SOURCE_PATH" before running your script.

Complete example

This example trains on a single A10 GPU from a local project, with no prior CLI setup beyond installing and configuring the Azure Databricks CLI. The project has three files:

my-training/
├── databricks.yml
├── command.sh
└── src/
    └── train.py

command.sh is the entry point named by command_path. It changes into the extracted code directory and runs the training script:

#!/usr/bin/env bash
set -euo pipefail
cd "$CODE_SOURCE_PATH"
python train.py

databricks.yml names the bundle, packages src/ as a tgz artifact, runs it as an ai_runtime_task, installs numpy in the task environment, and defines development and production targets:

bundle:
  name: my-training

artifacts:
  code:
    type: tgz
    path: .
    include: [src]
    files:
      - source: ./dist/code.tgz

resources:
  jobs:
    train:
      name: my-training
      tasks:
        - task_key: train
          environment_key: default
          ai_runtime_task:
            experiment: /Users/me@example.com/my-training
            code_source_path: ./dist/code.tgz
            deployments:
              - command_path: ./command.sh
                compute:
                  accelerator_type: GPU_1xA10
                  accelerator_count: 1
      environments:
        - environment_key: default
          spec:
            environment_version: '5'
            dependencies:
              - numpy

targets:
  dev:
    mode: development
    default: true
  prod:
    mode: production

Deploy the bundle and run the job. databricks bundle deploy builds and uploads the tgz artifact and creates the job; databricks bundle run starts it:

databricks bundle deploy --target dev
databricks bundle run train --target dev

Build multi-task workflows

An ai_runtime_task is a Azure Databricks job task, so it composes with the rest of a job. You can run a preparation step before training, combine GPU and CPU tasks in one job, and use different accelerators per task.

Order tasks with depends_on

Use depends_on to run tasks in sequence. The following pipeline runs a preparation notebook, then a GPU training task that starts only after the preparation task succeeds:

resources:
  jobs:
    train_pipeline:
      tasks:
        - task_key: prep
          notebook_task:
            notebook_path: ./prep.py
        - task_key: train
          depends_on:
            - task_key: prep
          ai_runtime_task:
            experiment: my-experiment
            code_source_path: ./dist/code.tgz
            deployments:
              - command_path: ./command.sh
                compute:
                  accelerator_type: GPU_1xA10
                  accelerator_count: 1

Combine GPU and CPU tasks

In the pipeline above, only the training step needs a GPU. Keeping non-GPU work such as data preparation in separate tasks keeps GPU time focused on training.

Note

An ai_runtime_task does not support Azure Databricks job task values ({{tasks.<task_key>.values.<name>}} or dbutils.jobs.taskValues). To pass data between steps, write it to a shared location that both tasks can read, such as a Unity Catalog volume or a workspace file, and reference that path from each task.

Schedule the workload

Add a schedule to the job to run it on a cadence. Ship the schedule paused so that deploying the bundle does not start runs on its own, then unpause it when you are ready:

resources:
  jobs:
    train_pipeline:
      schedule:
        quartz_cron_expression: '0 0 9 * * ?'
        timezone_id: UTC
        pause_status: PAUSED

Promote from development to production

Promotion from development to production is a standard bundle feature that the AI Runtime task inherits unchanged.

Bundle targets and modes

Define a production target with mode: production alongside your development target. The target controls where the bundle deploys and how its resources are named:

targets:
  dev:
    mode: development
    default: true
  prod:
    mode: production

Deploy and run

Deploy and run the bundle against a target with the standard bundle commands:

databricks bundle deploy --target dev
databricks bundle run train_pipeline --target dev

Next steps

Run and manage AI Runtime workloads from the command line with the AI Runtime CLI.
Track training runs and manage checkpoints. See Experiment tracking and observability.

Feedback

Was this page helpful?

Last updated on 2026-09-10

Productionize training workloads

Requirements

Define an AI Runtime task in a bundle

ai_runtime_task fields

Configure the hardware accelerator

Set the environment and dependencies

Ship your training code

Packaged artifact

Workspace or volume path

Complete example

Build multi-task workflows

Order tasks with depends_on

Combine GPU and CPU tasks

Schedule the workload

Promote from development to production

Bundle targets and modes

Deploy and run

Next steps

Feedback

Additional resources

`ai_runtime_task` fields