air: add convert-to-dabs (run YAML -> Databricks Asset Bundle)

acceptance/experimental/air/convert-to-dabs/docker.yaml

@@ -0,0 +1,8 @@
+experiment_name: docker-test
+command: python train.py
+compute:
+  accelerator_type: GPU_1xA10
+  num_accelerators: 1
+environment:
+  docker_image:
+    url: myregistry/img:tag

acceptance/experimental/air/convert-to-dabs/output.txt

@@ -0,0 +1,83 @@
+
+=== convert an AIR run YAML into a DABs bundle
+>>> [CLI] experimental air convert-to-dabs train.yaml --output-dir generated
+Wrote a Databricks Asset Bundle to generated:
+  databricks.yml
+  training_config.yaml
+  command.sh
+  code_source/
+
+Notes:
+  - code_source was staged into the code_source/ directory; bundle deploy packages and uploads it. Re-run convert-to-dabs to re-stage after code changes.
+
+To deploy and run this workload as a bundle:
+  1. cd generated
+  2. databricks bundle validate
+  3. databricks bundle deploy
+  4. databricks bundle run torchrun-a10-smoke-test
+
+bundle deploy uploads the code source and launch scripts automatically.
+
+Unlike `air run` (which submits an ephemeral run), bundle deploy creates a
+persistent job that is not garbage-collected. When you are done, remove the
+job and its uploaded files with:
+  databricks bundle destroy
+
+=== emitted databricks.yml
+>>> cat generated/databricks.yml
+bundle:
+  name: torchrun-a10-smoke-test
+targets:
+  dev:
+    mode: development
+    default: true
+resources:
+  jobs:
+    torchrun-a10-smoke-test:
+      name: torchrun-a10-smoke-test
+      tasks:
+        - task_key: torchrun-a10-smoke-test
+          environment_key: default
+          ai_runtime_task:
+            experiment: torchrun-a10-smoke-test
+            deployments:
+              - command_path: ./command.sh
+                compute:
+                  accelerator_type: GPU_1xA10
+                  accelerator_count: 1
+            code_source_path: ./code_source
+      environments:
+        - environment_key: default
+          spec:
+            environment_version: "5"
+            dependencies:
+              - numpy
+
+=== the generated command.sh carries the run command
+>>> cat generated/command.sh
+torchrun --nproc_per_node=1 train.py
+=== the code source is staged as a directory (packaged at deploy time)
+>>> ls generated
+code_source
+command.sh
+databricks.yml
+training_config.yaml
+
+>>> ls generated/code_source
+train.py
+
+=== the emitted bundle validates
+>>> [CLI] bundle validate
+Name: torchrun-a10-smoke-test
+Target: dev
+Workspace:
+  User: [USERNAME]
+  Path: /Workspace/Users/[USERNAME]/.bundle/torchrun-a10-smoke-test/dev
+
+Validation OK!
+
+=== docker_image is not supported yet
+>>> [CLI] experimental air convert-to-dabs docker.yaml --output-dir generated-docker
+Error: environment.docker_image is not yet supported by convert-to-dabs
+
+Exit code: 1

acceptance/experimental/air/convert-to-dabs/script

@@ -0,0 +1,20 @@
+title "convert an AIR run YAML into a DABs bundle"
+trace $CLI experimental air convert-to-dabs train.yaml --output-dir generated
+
+title "emitted databricks.yml"
+trace cat generated/databricks.yml
+
+title "the generated command.sh carries the run command"
+trace cat generated/command.sh
+
+title "the code source is staged as a directory (packaged at deploy time)"
+trace ls generated && true
+trace ls generated/code_source && true
+
+title "the emitted bundle validates"
+cd generated
+trace $CLI bundle validate
+cd ..
+
+title "docker_image is not supported yet"
+errcode trace $CLI experimental air convert-to-dabs docker.yaml --output-dir generated-docker

acceptance/experimental/air/convert-to-dabs/src/train.py

@@ -0,0 +1 @@
+print("train")

acceptance/experimental/air/convert-to-dabs/test.toml

@@ -0,0 +1,3 @@
+# The command writes a bundle under generated/; those are generated artifacts,
+# not committed inputs, so exclude them from the repo-diff check.
+Ignore = ["generated", "generated-docker"]

acceptance/experimental/air/convert-to-dabs/train.yaml

@@ -0,0 +1,13 @@
+experiment_name: torchrun-a10-smoke-test
+command: torchrun --nproc_per_node=1 train.py
+compute:
+  accelerator_type: GPU_1xA10
+  num_accelerators: 1
+environment:
+  version: 5
+  dependencies:
+    - numpy
+code_source:
+  type: snapshot
+  snapshot:
+    root_path: ./src

acceptance/experimental/air/help/output.txt

@@ -10,12 +10,13 @@ Usage:
   databricks experimental air [command]
 
 Available Commands:
-  cancel         Cancel one or more runs
-  get            Show status, configuration, and timing details for a specific run
-  list           List your active runs for the current profile (use --all-status for finished runs)
-  logs           Stream or fetch logs for a run
-  register-image Mirror a Docker image into the workspace registry
-  run            Submit a training workload from a YAML config
+  cancel          Cancel one or more runs
+  convert-to-dabs Convert an AIR run YAML into a Databricks Asset Bundle
+  get             Show status, configuration, and timing details for a specific run
+  list            List your active runs for the current profile (use --all-status for finished runs)
+  logs            Stream or fetch logs for a run
+  register-image  Mirror a Docker image into the workspace registry
+  run             Submit a training workload from a YAML config
 
 Flags:
   -h, --help   help for air

experimental/air/cmd/air.go

@@ -23,6 +23,7 @@ experimental and may change in future versions.`,
 	cmd.AddCommand(newLogsCommand())
 	cmd.AddCommand(newCancelCommand())
 	cmd.AddCommand(newRegisterImageCommand())
+	cmd.AddCommand(newConvertToDabsCommand())
 
 	return cmd
 }

experimental/air/cmd/air_test.go

@@ -14,7 +14,7 @@ func TestNewRegistersAllSubcommands(t *testing.T) {
 		registered[c.Name()] = true
 	}
 
-	want := []string{"run", "get", "list", "logs", "cancel", "register-image"}
+	want := []string{"run", "get", "list", "logs", "cancel", "register-image", "convert-to-dabs"}
 	for _, name := range want {
 		assert.True(t, registered[name], "subcommand %q is not registered", name)
 	}