Add AIR submission and YAML parity

.nextchanges/cli/air-submission-parity.md

@@ -0,0 +1 @@
+Restored AIR requirements YAML support and added Docker image, MLflow artifact location, and provisioned capacity submission fields.

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 package list or path to a requirements YAML file. 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

@@ -56,6 +56,62 @@ Tip: use --watch to stream logs until the run completes.
   }
 }
 
-=== a requirements.yaml file path is rejected; deps must be inline
+=== submit with file-backed dependencies
 >>> [CLI] experimental air run -f run-file.yaml
-Error: invalid config run-file.yaml: environment.dependencies must be a list of packages or reference a requirements.txt (see https://docs.databricks.com/aws/en/machine-learning/ai-runtime/cli/yaml-config#reference). A direct file reference is not supported
+Submitting experiment: deps-file-smoke
+Submitted workload with Job Run ID: 555
+View job run at: [DATABRICKS_URL]/jobs/runs/555
+
+Tip: use --watch to stream logs until the run completes.
+
+=== requirements.yaml is uploaded beside the launch config
+>>> 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-file-smoke/deps-file-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-file-smoke/deps-file-smoke_[RUN_ID]/requirements.yaml", "q": {"overwrite": "true"}, "raw_body": "version: 5\ndependencies:\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-file-smoke/deps-file-smoke_[RUN_ID]/training_config.yaml", "q": {"overwrite": "true"}, "raw_body": "experiment_name: deps-file-smoke\ncommand: python train.py\ncompute:\n  accelerator_type: GPU_1xH100\n  num_accelerators: 1\n  provisioned_capacity_id: capacity-123\nenvironment:\n  dependencies: ./reqs.yaml\nmlflow_artifact_location: /Volumes/main/default/air-artifacts\n"}
+
+=== file deps, artifact location, and capacity id reach the submit payload
+>>> print_requests.py //api/2.2/jobs/runs/submit
+{
+  "method": "POST",
+  "path": "/api/2.2/jobs/runs/submit",
+  "body": {
+    "environments": [
+      {
+        "environment_key": "default",
+        "spec": {
+          "dependencies": [
+            "numpy",
+            "torch==2.3.0"
+          ],
+          "environment_version": "5"
+        }
+      }
+    ],
+    "idempotency_token": "[UUID]",
+    "run_name": "deps-file-smoke",
+    "tasks": [
+      {
+        "ai_runtime_task": {
+          "deployments": [
+            {
+              "command_path": "/Workspace/Users/[USERNAME]/.air/cli_launch/deps-file-smoke/deps-file-smoke_[RUN_ID]/command.sh",
+              "compute": {
+                "accelerator_count": 1,
+                "accelerator_type": "GPU_1xH100",
+                "provisioned_capacity_id": "capacity-123"
+              }
+            }
+          ],
+          "experiment": "deps-file-smoke",
+          "mlflow_artifact_location": "dbfs:/Volumes/main/default/air-artifacts"
+        },
+        "environment_key": "default",
+        "max_retries": 3,
+        "retry_on_timeout": true,
+        "run_if": "ALL_SUCCESS",
+        "task_key": "deps-file-smoke"
+      }
+    ]
+  }
+}

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

@@ -3,5 +3,7 @@ command: python train.py
 compute:
   accelerator_type: GPU_1xH100
   num_accelerators: 1
+  provisioned_capacity_id: capacity-123
 environment:
   dependencies: ./reqs.yaml
+mlflow_artifact_location: /Volumes/main/default/air-artifacts

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

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

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]'

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/convert_to_dabs.go

@@ -329,8 +329,8 @@ func buildBundleValue(ctx context.Context, cfg *runConfig, configPath, codeSourc
 // bundleEnvironmentDeps resolves the runtime version and the inline dependency
 // list to emit in the bundle's environments[] spec. The aicode mutator synthesizes
 // requirements.yaml from that spec at deploy, so the whole set must be here.
-// Dependencies are inline-only (a requirements-file path is rejected at config
-// load), so an unset list yields no dependencies.
+// File-backed dependencies are resolved during config load, so an unset list
+// yields no dependencies here.
 func bundleEnvironmentDeps(ctx context.Context, cfg *runConfig) (version string, deps []string) {
 	cfgVersion, _ := cfg.runtimeVersion()
 	version = dlRuntimeImage(ctx, cfgVersion)

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
@@ -229,7 +244,7 @@ func validateSecretRefs(secrets map[string]string) error {
 // environmentConfig is the `environment` block: dependencies and/or a custom
 // docker image.
 type environmentConfig struct {
-	Dependencies dependencies       `yaml:"dependencies" help:"Inline list of packages to install. Not allowed alongside docker_image."`
+	Dependencies dependencies       `yaml:"dependencies" help:"Inline package list or path to a requirements YAML file. Not allowed alongside docker_image."`
 	Version      stringOrInt        `yaml:"version" help:"Client image version to pin. Only valid alongside inline dependencies."`
 	DockerImage  *dockerImageConfig `yaml:"docker_image" help:"Custom image supplying the whole runtime. Not allowed alongside dependencies or version."`
 }
@@ -256,24 +271,43 @@ 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
 }
 
-// dependencies is environment.dependencies: an inline list of packages. A scalar
-// (e.g. a path to a requirements file) is rejected — the list may itself reference
-// a requirements.txt, but dependencies must be given as a list.
+// dependencies is environment.dependencies: either an inline list of packages or
+// a path to a requirements YAML file resolved relative to the run config.
 type dependencies struct {
-	set  bool
-	list []string
+	set          bool
+	list         []string
+	path         string
+	resolvedPath string
+	version      string
 }
 
 func (d *dependencies) UnmarshalYAML(node *yaml.Node) error {
-	if node.Kind != yaml.SequenceNode {
-		return errors.New("environment.dependencies must be a list of packages or reference a requirements.txt (see https://docs.databricks.com/aws/en/machine-learning/ai-runtime/cli/yaml-config#reference). A direct file reference is not supported")
-	}
 	d.set = true
-	return node.Decode(&d.list)
+	switch node.Kind {
+	case yaml.SequenceNode:
+		return node.Decode(&d.list)
+	case yaml.ScalarNode:
+		if err := node.Decode(&d.path); err != nil {
+			return err
+		}
+		if strings.TrimSpace(d.path) == "" {
+			return errors.New("environment.dependencies requirements YAML path cannot be empty")
+		}
+		return nil
+	default:
+		return errors.New("environment.dependencies must be a list of packages or a requirements YAML path")
+	}
 }
 
 // stringOrInt holds a scalar that may be a string or an integer in YAML

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
@@ -53,10 +50,23 @@ func (c *runConfig) inlineDependencies() ([]string, bool) {
 	return c.Environment.Dependencies.list, true
 }
 
+func (c *runConfig) requirementsPath() string {
+	if c.Environment == nil {
+		return ""
+	}
+	return c.Environment.Dependencies.resolvedPath
+}
+
 // 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.Dependencies.version != "" {
+		return c.Environment.Dependencies.version, true
+	}
+	if !c.Environment.Version.set {
 		return "", false
 	}
 	return c.Environment.Version.raw, true