Skip to main content
Last updated on

Productionize training workloads

Preview

This feature is in Public Preview.

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:

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

command_path is the script the task runs. Retries, timeouts, and permissions are set on the task and job the same way as any Databricks job, so your existing bundle practices carry over.

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:

YAML
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 accepts three forms:

  • A local directory — the CLI packages and uploads it on databricks bundle deploy. This is the simplest form and the one the examples above use.
  • An explicit tgz artifact — you declare the artifact yourself and point at its output file. Use this when you need to package only a subset of files, or snapshot a committed Git revision instead of your working tree.
  • A workspace or volume path — code that is already uploaded, used as-is.

Point code_source_path at a directory. On databricks bundle deploy, the CLI packages the directory, uploads it, and the task extracts it and runs your command against it:

YAML
code_source_path: ./src
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.

Build multi-task workflows​

An ai_runtime_task is a 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:

YAML
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: ./src
            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 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:

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

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

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

Next steps​

On this page