Add AIR submission and YAML parity

acceptance/experimental/air/config-help/output.txt

@@ -50,6 +50,7 @@ config
     parameters                   Free-form values passed through to the workload.
     mlflow_run_name              Name for the MLflow run.
     mlflow_experiment_directory  Workspace directory holding the MLflow experiment.
+    mlflow_artifact_location     DBFS location where MLflow artifacts are written.
     permissions                  Who may view or manage the run, as a list of principal plus level grants.
     usage_policy_name            Usage policy to bill the run to, by name.
     usage_policy_id              Usage policy to bill the run to, by id.
@@ -62,11 +63,34 @@ config.compute
   Which accelerators to run on and how many.
 
   Fields:
-    num_accelerators  Total number of GPUs to allocate.
-    accelerator_type  Which accelerator to run on, e.g. GPU_1xA10.
+    num_accelerators         Total number of GPUs to allocate.
+    accelerator_type         Which accelerator to run on, e.g. GPU_1xA10.
+    provisioned_capacity_id  Pre-provisioned AIR capacity reservation id.
 
 Use "-h config.compute.<field>" for details on a field.
 
+=== new submission fields are documented
+>>> [CLI] experimental air run -h config.mlflow_artifact_location
+config.mlflow_artifact_location
+  DBFS location where MLflow artifacts are written. A /Volumes path is normalized to dbfs:/Volumes/... .
+
+  Type:     string
+  Required: no
+
+>>> [CLI] experimental air run -h config.compute.provisioned_capacity_id
+config.compute.provisioned_capacity_id
+  Pre-provisioned AIR capacity reservation id. Must be 1-255 characters.
+
+  Type:     string
+  Required: no
+
+>>> [CLI] experimental air run -h config.environment.dependencies
+config.environment.dependencies
+  Inline list of packages to install. Not allowed alongside docker_image.
+
+  Type:     list of strings
+  Required: no
+
 === leaf field
 >>> [CLI] experimental air run -h config.compute.accelerator_type
 config.compute.accelerator_type
@@ -103,7 +127,7 @@ config.compute.num_accelerators
 >>> [CLI] experimental air run -h config.compute.acclerator_type
 Error: unknown config field "config.compute.acclerator_type"; did you mean "accelerator_type"?
 
-fields under "config.compute" are: accelerator_type, num_accelerators
+fields under "config.compute" are: accelerator_type, num_accelerators, provisioned_capacity_id
 
 === free-form map keys are not schema fields
 >>> [CLI] experimental air run -h config.parameters.learning_rate

acceptance/experimental/air/config-help/script

@@ -12,6 +12,11 @@ trace $CLI experimental air run -h config
 title "nested object lists its fields"
 trace $CLI experimental air run -h config.compute
 
+title "new submission fields are documented"
+trace $CLI experimental air run -h config.mlflow_artifact_location
+trace $CLI experimental air run -h config.compute.provisioned_capacity_id
+trace $CLI experimental air run -h config.environment.dependencies
+
 title "leaf field"
 trace $CLI experimental air run -h config.compute.accelerator_type
 

acceptance/experimental/air/run-submit-deps/output.txt

@@ -10,9 +10,9 @@ Tip: use --watch to stream logs until the run completes.
 === only config + command are uploaded; no requirements.yaml
 >>> print_requests.py //api/2.0/workspace-files/import-file --oneline --sort --unique --keep
 {"method": "POST", "path": "/api/2.0/workspace-files/import-file/Workspace/Users/[USERNAME]/.air/cli_launch/deps-smoke/deps-smoke_[RUN_ID]/command.sh", "q": {"overwrite": "true"}, "raw_body": "python train.py"}
-{"method": "POST", "path": "/api/2.0/workspace-files/import-file/Workspace/Users/[USERNAME]/.air/cli_launch/deps-smoke/deps-smoke_[RUN_ID]/training_config.yaml", "q": {"overwrite": "true"}, "raw_body": "experiment_name: deps-smoke\ncommand: python train.py\ncompute:\n  accelerator_type: GPU_1xH100\n  num_accelerators: 1\nenvironment:\n  version: 5\n  dependencies:\n    - numpy\n    - torch==2.3.0\n"}
+{"method": "POST", "path": "/api/2.0/workspace-files/import-file/Workspace/Users/[USERNAME]/.air/cli_launch/deps-smoke/deps-smoke_[RUN_ID]/training_config.yaml", "q": {"overwrite": "true"}, "raw_body": "experiment_name: deps-smoke\ncommand: python train.py\ncompute:\n  accelerator_type: GPU_1xH100\n  num_accelerators: 1\n  provisioned_capacity_id: capacity-123\nenvironment:\n  version: 5\n  dependencies:\n    - numpy\n    - torch==2.3.0\nmlflow_artifact_location: /Volumes/main/default/air-artifacts\n"}
 
-=== declared deps ride on environments[].spec.dependencies
+=== inline deps, artifact location, and capacity id reach the submit payload
 >>> print_requests.py //api/2.2/jobs/runs/submit
 {
   "method": "POST",
@@ -40,11 +40,13 @@ Tip: use --watch to stream logs until the run completes.
               "command_path": "/Workspace/Users/[USERNAME]/.air/cli_launch/deps-smoke/deps-smoke_[RUN_ID]/command.sh",
               "compute": {
                 "accelerator_count": 1,
-                "accelerator_type": "GPU_1xH100"
+                "accelerator_type": "GPU_1xH100",
+                "provisioned_capacity_id": "capacity-123"
               }
             }
           ],
-          "experiment": "deps-smoke"
+          "experiment": "deps-smoke",
+          "mlflow_artifact_location": "dbfs:/Volumes/main/default/air-artifacts"
         },
         "environment_key": "default",
         "max_retries": 3,

acceptance/experimental/air/run-submit-deps/run.yaml

@@ -3,8 +3,10 @@ command: python train.py
 compute:
   accelerator_type: GPU_1xH100
   num_accelerators: 1
+  provisioned_capacity_id: capacity-123
 environment:
   version: 5
   dependencies:
     - numpy
     - torch==2.3.0
+mlflow_artifact_location: /Volumes/main/default/air-artifacts

acceptance/experimental/air/run-submit-deps/script

@@ -4,7 +4,7 @@ trace $CLI experimental air run -f run.yaml
 title "only config + command are uploaded; no requirements.yaml"
 trace print_requests.py //api/2.0/workspace-files/import-file --oneline --sort --unique --keep
 
-title "declared deps ride on environments[].spec.dependencies"
+title "inline deps, artifact location, and capacity id reach the submit payload"
 trace print_requests.py //api/2.2/jobs/runs/submit
 
 title "a requirements.yaml file path is rejected; deps must be inline"

acceptance/experimental/air/run-submit-deps/test.toml

@@ -21,3 +21,7 @@ Response.Body = '''
 [[Repls]]
 Old = 'deps-smoke_[0-9a-f]{16}'
 New = 'deps-smoke_[RUN_ID]'
+
+[[Repls]]
+Old = 'deps-file-smoke_[0-9a-f]{16}'
+New = 'deps-file-smoke_[RUN_ID]'

acceptance/experimental/air/run/output.txt

@@ -22,7 +22,7 @@ Dry run: configuration for "smoke-test" is valid; not submitting.
 
 === override of an unknown field is rejected
 >>> [CLI] experimental air run -f valid.yaml --dry-run --override bogus=1
-Error: invalid --override "bogus": "bogus" is not a known field; available fields are: code_source, command, compute, env_variables, environment, experiment_name, idempotency_token, max_retries, mlflow_experiment_directory, mlflow_run_name, parameters, permissions, secrets, timeout_minutes, usage_policy_id, usage_policy_name
+Error: invalid --override "bogus": "bogus" is not a known field; available fields are: code_source, command, compute, env_variables, environment, experiment_name, idempotency_token, max_retries, mlflow_artifact_location, mlflow_experiment_directory, mlflow_run_name, parameters, permissions, secrets, timeout_minutes, usage_policy_id, usage_policy_name
 
 Exit code: 1
 

experimental/air/cmd/compute.go

@@ -1,6 +1,7 @@
 package aircmd
 
 import (
+	"errors"
 	"fmt"
 	"strings"
 )
@@ -54,12 +55,13 @@ func gpusPerNode(g gpuType) (int, error) {
 // computeConfig is the `compute` block of the run YAML: which accelerators to
 // use and how many.
 type computeConfig struct {
-	NumAccelerators int    `yaml:"num_accelerators" help:"Total number of GPUs to allocate. Must be a positive multiple of the accelerator type's per-node GPU count. See https://docs.databricks.com/aws/en/machine-learning/ai-runtime/cli/yaml-config#reference for supported GPU types."`
-	AcceleratorType string `yaml:"accelerator_type" help:"Which accelerator to run on, e.g. GPU_1xA10. See https://docs.databricks.com/aws/en/machine-learning/ai-runtime/cli/yaml-config#reference for the current list of supported GPU types. Matched case-sensitively."`
+	NumAccelerators       int     `yaml:"num_accelerators" help:"Total number of GPUs to allocate. Must be a positive multiple of the accelerator type's per-node GPU count. See https://docs.databricks.com/aws/en/machine-learning/ai-runtime/cli/yaml-config#reference for supported GPU types."`
+	AcceleratorType       string  `yaml:"accelerator_type" help:"Which accelerator to run on, e.g. GPU_1xA10. See https://docs.databricks.com/aws/en/machine-learning/ai-runtime/cli/yaml-config#reference for the current list of supported GPU types. Matched case-sensitively."`
+	ProvisionedCapacityID *string `yaml:"provisioned_capacity_id" help:"Pre-provisioned AIR capacity reservation id. Must be 1-255 characters."`
 }
 
 // validate checks the compute block against the backend's constraints.
-func (c computeConfig) validate() error {
+func (c *computeConfig) validate() error {
 	g, err := parseGPUType(c.AcceleratorType)
 	if err != nil {
 		return fmt.Errorf("compute.accelerator_type: %w", err)
@@ -77,5 +79,16 @@ func (c computeConfig) validate() error {
 		return fmt.Errorf("compute.num_accelerators for %s must be a multiple of %d, got %d", c.AcceleratorType, perNode, c.NumAccelerators)
 	}
 
+	if c.ProvisionedCapacityID != nil {
+		v := strings.TrimSpace(*c.ProvisionedCapacityID)
+		if v == "" {
+			return errors.New("compute.provisioned_capacity_id cannot be empty")
+		}
+		if len(v) > 255 {
+			return fmt.Errorf("compute.provisioned_capacity_id must be 255 characters or less, got %d", len(v))
+		}
+		*c.ProvisionedCapacityID = v
+	}
+
 	return nil
 }

experimental/air/cmd/compute_test.go

@@ -1,6 +1,7 @@
 package aircmd
 
 import (
+	"strings"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -67,6 +68,9 @@ func TestComputeConfigValidate(t *testing.T) {
 		{"single node", computeConfig{NumAccelerators: 8, AcceleratorType: "GPU_8xH100"}, ""},
 		{"multiple nodes", computeConfig{NumAccelerators: 16, AcceleratorType: "GPU_8xH100"}, ""},
 		{"single-gpu partitions", computeConfig{NumAccelerators: 3, AcceleratorType: "GPU_1xH100"}, ""},
+		{"capacity id", computeConfig{NumAccelerators: 1, AcceleratorType: "GPU_1xH100", ProvisionedCapacityID: new(" capacity ")}, ""},
+		{"empty capacity id", computeConfig{NumAccelerators: 1, AcceleratorType: "GPU_1xH100", ProvisionedCapacityID: new(" ")}, "cannot be empty"},
+		{"long capacity id", computeConfig{NumAccelerators: 1, AcceleratorType: "GPU_1xH100", ProvisionedCapacityID: new(strings.Repeat("a", 256))}, "255 characters or less"},
 		{"unknown type", computeConfig{NumAccelerators: 8, AcceleratorType: "b200"}, "accelerator_type"},
 		{"legacy type rejected", computeConfig{NumAccelerators: 8, AcceleratorType: "h100_80gb"}, "accelerator_type"},
 		{"non-positive count", computeConfig{NumAccelerators: 0, AcceleratorType: "GPU_1xH100"}, "must be positive"},

experimental/air/cmd/runconfig.go

@@ -56,6 +56,7 @@ type runConfig struct {
 	Parameters                map[string]any `yaml:"parameters" help:"Free-form values passed through to the workload. Any nested structure is allowed."`
 	MLflowRunName             *string        `yaml:"mlflow_run_name" help:"Name for the MLflow run. Max 100 characters, alphanumerics, hyphens, and underscores only."`
 	MLflowExperimentDirectory *string        `yaml:"mlflow_experiment_directory" help:"Workspace directory holding the MLflow experiment. Must start with /Workspace."`
+	MLflowArtifactLocation    *string        `yaml:"mlflow_artifact_location" help:"DBFS location where MLflow artifacts are written. A /Volumes path is normalized to dbfs:/Volumes/... ."`
 	Permissions               []permission   `yaml:"permissions" help:"Who may view or manage the run, as a list of principal plus level grants."`
 	UsagePolicyName           *string        `yaml:"usage_policy_name" help:"Usage policy to bill the run to, by name. Max 127 characters. Mutually exclusive with usage_policy_id."`
 	UsagePolicyID             *string        `yaml:"usage_policy_id" help:"Usage policy to bill the run to, by id. Mutually exclusive with usage_policy_name."`
@@ -150,6 +151,20 @@ func (c *runConfig) validate() error {
 		}
 	}
 
+	if c.MLflowArtifactLocation != nil {
+		v := strings.TrimSpace(*c.MLflowArtifactLocation)
+		if v == "" {
+			return errors.New("mlflow_artifact_location cannot be empty")
+		}
+		if strings.HasPrefix(v, "/Volumes/") {
+			v = "dbfs:" + v
+		}
+		if !strings.HasPrefix(v, "dbfs:/") {
+			return fmt.Errorf("mlflow_artifact_location must be a dbfs: URI, got: %s", v)
+		}
+		*c.MLflowArtifactLocation = v
+	}
+
 	for i := range c.Permissions {
 		if err := c.Permissions[i].validate(); err != nil {
 			return err
@@ -256,6 +271,13 @@ func (e *environmentConfig) validate() error {
 	if e.Version.set && !e.Dependencies.set {
 		return errors.New("'environment.version' requires inline 'dependencies' (a list of packages)")
 	}
+	if e.Version.set {
+		version, err := validateRuntimeVersion(e.Version.raw, "environment.version")
+		if err != nil {
+			return err
+		}
+		e.Version.raw = version
+	}
 
 	return nil
 }

experimental/air/cmd/runconfig_launch.go

@@ -26,9 +26,6 @@ func (c *runConfig) maxRetries() int {
 }
 
 // dockerImageURL returns the custom docker image URL, or "" when none is set.
-//
-// TODO: not wired into submission yet — the native ai_runtime_task carries no
-// docker field, and full support needs image registration (pending the DCS work).
 func (c *runConfig) dockerImageURL() string {
 	if c.Environment != nil && c.Environment.DockerImage != nil {
 		return c.Environment.DockerImage.URL
@@ -56,7 +53,10 @@ func (c *runConfig) inlineDependencies() ([]string, bool) {
 // runtimeVersion returns the client image version from environment.version when
 // set.
 func (c *runConfig) runtimeVersion() (string, bool) {
-	if c.Environment == nil || !c.Environment.Version.set {
+	if c.Environment == nil {
+		return "", false
+	}
+	if !c.Environment.Version.set {
 		return "", false
 	}
 	return c.Environment.Version.raw, true

experimental/air/cmd/runconfig_load.go

@@ -7,6 +7,9 @@ import (
 	"fmt"
 	"io"
 	"os"
+	"regexp"
+	"strconv"
+	"strings"
 
 	"go.yaml.in/yaml/v3"
 )
@@ -109,3 +112,30 @@ func loadRunConfigWithOverrides(ctx context.Context, path string, overrides []st
 	}
 	return cfg, nil
 }
+
+var runtimeVersionRe = regexp.MustCompile(`^[0-9]+$`)
+
+const databricksAIPrefix = "databricks_ai_v"
+
+func validateRuntimeVersion(version, source string) (string, error) {
+	normalized := strings.ToLower(version)
+	numeric := normalized
+	usesDatabricksAI := strings.HasPrefix(normalized, databricksAIPrefix)
+	if usesDatabricksAI {
+		numeric = strings.TrimPrefix(normalized, databricksAIPrefix)
+	}
+	if !runtimeVersionRe.MatchString(numeric) {
+		return "", fmt.Errorf("unsupported client image version %q in %s: version must be an integer, optionally prefixed with databricks_ai_v", version, source)
+	}
+	if !usesDatabricksAI {
+		return numeric, nil
+	}
+	major, err := strconv.Atoi(numeric)
+	if err != nil {
+		return "", fmt.Errorf("failed to parse client image version %q in %s: %w", version, source, err)
+	}
+	if major < 5 {
+		return "", fmt.Errorf("databricks_ai_v in %s requires AI Runtime version 5 or higher, got %q", source, version)
+	}
+	return databricksAIPrefix + numeric, nil
+}

experimental/air/cmd/runconfig_test.go

@@ -95,18 +95,8 @@ permissions:
 }
 
 // TestLoadRunConfig_PolymorphicFields exercises the str|int and bool|str unions
-// decoded by custom UnmarshalYAML, plus the rejection of the removed
-// dependencies string form.
+// decoded by custom UnmarshalYAML.
 func TestLoadRunConfig_PolymorphicFields(t *testing.T) {
-	t.Run("dependencies as string path is rejected", func(t *testing.T) {
-		_, err := loadRunConfig(writeConfig(t, minimalConfig+`
-environment:
-  dependencies: requirements.yaml
-`))
-		require.Error(t, err)
-		assert.Contains(t, err.Error(), "must be a list of packages")
-	})
-
 	t.Run("git remote as bool true is rejected", func(t *testing.T) {
 		_, err := loadRunConfig(writeConfig(t, minimalConfig+`
 code_source:
@@ -229,6 +219,9 @@ func TestRunConfigValidate_FieldRules(t *testing.T) {
 		{"long idempotency", func(c *runConfig) { c.IdempotencyToken = str(string(make([]byte, 65))) }, "64 characters or less"},
 		{"bad mlflow_run_name", func(c *runConfig) { c.MLflowRunName = str("bad name") }, "invalid mlflow_run_name"},
 		{"bad experiment dir", func(c *runConfig) { c.MLflowExperimentDirectory = str("/Users/me") }, "must start with '/Workspace'"},
+		{"artifact volume path normalizes", func(c *runConfig) { c.MLflowArtifactLocation = str(" /Volumes/main/default/artifacts ") }, ""},
+		{"empty artifact location", func(c *runConfig) { c.MLflowArtifactLocation = str(" ") }, "mlflow_artifact_location cannot be empty"},
+		{"non-dbfs artifact location", func(c *runConfig) { c.MLflowArtifactLocation = str("s3://bucket/path") }, "must be a dbfs: URI"},
 		{"empty usage policy", func(c *runConfig) { c.UsagePolicyName = str(" ") }, "usage_policy_name must not be empty"},
 		{"bad secret ref", func(c *runConfig) { c.Secrets = map[string]string{"T": "noslash"} }, "expected format 'scope/key'"},
 		{"empty secret scope", func(c *runConfig) { c.Secrets = map[string]string{"T": "/key"} }, "scope and key cannot be empty"},
@@ -254,6 +247,9 @@ func TestRunConfigValidate_FieldRules(t *testing.T) {
 			err := c.validate()
 			if tt.errFrag == "" {
 				assert.NoError(t, err)
+				if tt.name == "artifact volume path normalizes" {
+					assert.Equal(t, "dbfs:/Volumes/main/default/artifacts", *c.MLflowArtifactLocation)
+				}
 				return
 			}
 			require.Error(t, err)
@@ -299,12 +295,31 @@ func TestEnvironmentConfigValidate(t *testing.T) {
 			},
 			"",
 		},
+		{
+			"version prefix normalizes",
+			environmentConfig{
+				Version:      stringOrInt{set: true, raw: "DATABRICKS_AI_V5"},
+				Dependencies: dependencies{set: true, list: []string{"torch"}},
+			},
+			"",
+		},
+		{
+			"old databricks ai version rejected",
+			environmentConfig{
+				Version:      stringOrInt{set: true, raw: "databricks_ai_v4"},
+				Dependencies: dependencies{set: true, list: []string{"torch"}},
+			},
+			"requires AI Runtime version 5 or higher",
+		},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
 			err := tt.env.validate()
 			if tt.errFrag == "" {
 				assert.NoError(t, err)
+				if tt.name == "version prefix normalizes" {
+					assert.Equal(t, "databricks_ai_v5", tt.env.Version.raw)
+				}
 				return
 			}
 			require.Error(t, err)
@@ -474,7 +489,7 @@ func TestResolveConfigField_Containers(t *testing.T) {
 	compute, err := resolveConfigField("config.compute")
 	require.NoError(t, err)
 	assert.Equal(t, "object", compute.typeName)
-	require.Len(t, compute.children, 2)
+	require.Len(t, compute.children, 3)
 }
 
 func TestResolveConfigField_Errors(t *testing.T) {