[air] Add rank-partitioned container support

[vinchenzo-db]

Changes

Add multi-container support in CLI yaml. Now CLI will support both singleton tasks, as well as multi-container/image multi-node tasks. Included an example yaml.

Top level env vars are merged into each container. Each container now has it's own definition of:

1. name
2. command_path
3. ranks this container is assigned to
4. docker_image (though docker images are no longer supported in favor of UC images)
5. unity_catalog_image_path
6. environment variables

I've also added example docker builds for the RL usecase I'm covering, in case another user wants to reproduce it.

KNOWN GAP:
Jobs page will not show container level workloads because Jobs UI reads from task level command_path to find the training workload, and has no concept of containers yet.

Why

For primarily RL usecase, users may want to combine inference workload with training workload. We want to have a way to support this multi-node multi-image usecase

Tests

Unit tests

E2E test:
Replace placeholder images in example with my own DAR images.

# v.chen at ip-10-90-21-154 in ~/databricks-cli-wt-containers (git:air-containers) [20:40:39]
$ sed \
  -e 's|your_catalog.your_schema.vllm_trl:0.24.0|main.vinchenzo.vllm-trl-e2e:0.24.0|' \
  -e 's|your_catalog.your_schema.trl_vllm:0.24.0|main.vinchenzo.trl-vllm-e2e:0.24.0|' \
  cmd/air/examples/multi-image-example.yaml > /tmp/multi-image-example.yaml

Dry run:

# v.chen at ip-10-90-21-154 in ~/databricks-cli-wt-containers (git:air-containers) [20:40:47]
$ ./databricks air run -f /tmp/multi-image-example.yaml \
  --dry-run -p dbc-04ac0685-8857
Dry run: configuration for "qwen05b_grpo_dar_e2e_5epochs" is valid; not submitting.

Launch run:

# v.chen at ip-10-90-21-154 in ~/databricks-cli-wt-containers (git:air-containers) [20:40:54]
$ ./databricks air run -f /tmp/multi-image-example.yaml -p dbc-04ac0685-8857 
+--------------------------------------------------------------------+
| WARNING: EXPERIMENTAL FEATURE                                      |
|                                                                    |
| This feature is experimental and may change without notice,        |
| including changes that could break existing workflows.             |
| Use it only if you accept these risks.                             |
+--------------------------------------------------------------------+
Submitting experiment: qwen05b_grpo_dar_e2e_5epochs
Submitted workload with Job Run ID: 286153620017076
View job run at: https://dbc-04ac0685-8857.staging.cloud.databricks.com/jobs/runs/286153620017076

Tip: use --watch when submitting a run to stream logs to your terminal.
Stream logs after submission using:
  databricks air logs 286153620017076 -p dbc-04ac0685-8857

Job run: https://dbc-04ac0685-8857.staging.cloud.databricks.com/jobs/runs/286153620017076

Sneak peak at exciting logs:

[eng-dev-ecosystem-bot]

Integration test report

Commit: 61f4615

Run: 36928730878

	Env	✅​pass	🙈​skip	Time
✅​	aws linux	276	17	5:22
✅​	aws windows	278	15	3:05
✅​	azure linux	275	17	5:22
✅​	azure windows	277	15	3:01
✅​	gcp linux	276	17	5:43
✅​	gcp windows	278	15	3:02

Top 6 slowest tests (at least 2 minutes):

duration	env	testname
4:03	gcp linux	TestAccept
3:48	aws linux	TestAccept
3:45	azure linux	TestAccept
3:03	aws windows	TestAccept
2:59	gcp windows	TestAccept
2:46	azure windows	TestAccept

[ben-hansen-db]

A couple of comments to address. Nothing blocking

[ben-hansen-db]

The backend looks for hyperparameters.yaml beside each container’s command, but this code uploads it only to the parent launch directory.

We'll need to upload the shared hyperparameters beside every container command.

[ben-hansen-db]

Do we have validation for "containers requires at least two entries"?

[vinchenzo-db]

no need, with one container it will fan out across all nodes. The reason is because in the future, it makes sense to keep the container format and remove the "singleton" format.