AIR CLI: send dependencies inline, stop uploading requirements.yaml

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

@@ -0,0 +1,6 @@
+# These YAML files' contents are uploaded verbatim (as training_config.yaml and
+# requirements.yaml), so their line endings must stay \n on every OS — a Windows
+# \r would change the recorded payload.
+run.yaml text eol=lf
+run-file.yaml text eol=lf
+reqs.yaml text eol=lf

acceptance/experimental/air/run-submit-deps/databricks.yml

@@ -0,0 +1,2 @@
+bundle:
+  name: air-run-submit-deps

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

@@ -0,0 +1,108 @@
+
+=== submit with inline dependencies
+>>> [CLI] experimental air run -f run.yaml
+Submitted run 555
+View at: [DATABRICKS_URL]/jobs/runs/555
+
+=== 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"}
+
+=== declared deps ride on environments[].spec.dependencies
+>>> 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-smoke",
+    "tasks": [
+      {
+        "ai_runtime_task": {
+          "deployments": [
+            {
+              "command_path": "/Workspace/Users/[USERNAME]/.air/cli_launch/deps-smoke/deps-smoke_[RUN_ID]/command.sh",
+              "compute": {
+                "accelerator_count": 1,
+                "accelerator_type": "GPU_1xH100"
+              }
+            }
+          ],
+          "experiment": "deps-smoke"
+        },
+        "environment_key": "default",
+        "max_retries": 3,
+        "retry_on_timeout": true,
+        "run_if": "ALL_SUCCESS",
+        "task_key": "deps-smoke"
+      }
+    ]
+  }
+}
+
+=== file-form deps: version comes from the requirements file
+>>> [CLI] experimental air run -f run-file.yaml
+Submitted run 555
+View at: [DATABRICKS_URL]/jobs/runs/555
+
+=== file-form deps: the requirements file is not uploaded either
+>>> 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]/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\nenvironment:\n  dependencies: ./reqs.yaml\n"}
+
+=== file-form deps ride on environments[].spec.dependencies
+>>> 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"
+              }
+            }
+          ],
+          "experiment": "deps-file-smoke"
+        },
+        "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/reqs.yaml

@@ -0,0 +1,4 @@
+version: 5
+dependencies:
+  - numpy
+  - torch==2.3.0

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

@@ -0,0 +1,7 @@
+experiment_name: deps-file-smoke
+command: python train.py
+compute:
+  accelerator_type: GPU_1xH100
+  num_accelerators: 1
+environment:
+  dependencies: ./reqs.yaml

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

@@ -0,0 +1,10 @@
+experiment_name: deps-smoke
+command: python train.py
+compute:
+  accelerator_type: GPU_1xH100
+  num_accelerators: 1
+environment:
+  version: 5
+  dependencies:
+    - numpy
+    - torch==2.3.0

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

@@ -0,0 +1,17 @@
+title "submit with inline dependencies"
+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"
+trace print_requests.py //api/2.2/jobs/runs/submit
+
+title "file-form deps: version comes from the requirements file"
+trace $CLI experimental air run -f run-file.yaml
+
+title "file-form deps: the requirements file is not uploaded either"
+trace print_requests.py //api/2.0/workspace-files/import-file --oneline --sort --unique --keep
+
+title "file-form deps ride on environments[].spec.dependencies"
+trace print_requests.py //api/2.2/jobs/runs/submit

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

@@ -0,0 +1,23 @@
+# A non-dry-run submit with inline dependencies (no code_source): asserts the deps
+# land on spec.dependencies and that requirements.yaml is not uploaded.
+RecordRequests = true
+
+# The SDK probes host reachability with a HEAD request; stub it for determinism.
+[[Server]]
+Pattern = "HEAD /"
+Response.Body = ''
+
+[[Server]]
+Pattern = "POST /api/2.2/jobs/runs/submit"
+Response.Body = '''
+{"run_id": 555}
+'''
+
+# The per-run launch directory ends in <run>_<16 hex>; the random suffix varies.
+[[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/runsubmit.go

@@ -5,6 +5,7 @@ import (
 	"errors"
 	"fmt"
 	"path"
+	"path/filepath"
 	"strconv"
 	"strings"
 
@@ -38,16 +39,37 @@ func dlRuntimeImage(ctx context.Context, runtimeVersion string) string {
 	return strings.TrimPrefix(img, "CLIENT-GPU-")
 }
 
+// environmentDependencies resolves the user's declared dependencies as a flat
+// list to carry inline on the serverless environment's spec.dependencies: the
+// inline list directly, or the dependencies read from a requirements file
+// (resolved against the config's directory). For file-form deps it also returns
+// the version declared inside that file, which selects the runtime image since
+// top-level environment.version is not allowed there. Returns nil when none are
+// declared.
+func environmentDependencies(cfg *runConfig, configPath string) (deps []string, fileVersion string, err error) {
+	if deps, ok := cfg.inlineDependencies(); ok {
+		return deps, "", nil
+	}
+	if reqPath, ok := cfg.requirementsFile(); ok {
+		if !filepath.IsAbs(reqPath) {
+			reqPath = filepath.Join(filepath.Dir(configPath), reqPath)
+		}
+		return readRequirementsDependencies(reqPath)
+	}
+	return nil, "", nil
+}
+
 // buildSubmitPayload assembles the runs/submit payload. commandPath is the
-// workspace path of the uploaded command.sh; dlImage is the runtime channel.
+// workspace path of the uploaded command.sh; dlImage is the runtime channel;
+// deps is the user's declared dependencies (nil when none are declared).
 //
 // max_retries is always sent (including 0) so the user's YAML value is honored:
 // setting it to 0 explicitly disables retries rather than falling back to the
 // server default. retry_on_timeout is sent only when retries are allowed, and is
 // omitempty so the wire form matches the Python CLI (which never emits a bare
 // "false"). Jobs performs the retries — each attempt is a fresh AI Runtime
 // workload.
-func buildSubmitPayload(cfg *runConfig, commandPath, dlImage string, snap snapshotResult) jobs.SubmitRun {
+func buildSubmitPayload(cfg *runConfig, commandPath, dlImage string, snap snapshotResult, deps []string) jobs.SubmitRun {
 	task := jobs.AiRuntimeTask{
 		Experiment: cfg.ExperimentName,
 		Deployments: []jobs.DeploymentSpec{{
@@ -89,13 +111,21 @@ func buildSubmitPayload(cfg *runConfig, commandPath, dlImage string, snap snapsh
 		ForceSendFields: []string{"MaxRetries"},
 	}
 
+	// Carry the user's declared deps inline on spec.dependencies; the AI Runtime
+	// backend installs them via --deps-config. The SDK marshaler drops nil and empty
+	// slices, so a no-deps run omits the key.
+	envSpec := &compute.Environment{EnvironmentVersion: dlImage}
+	if len(deps) > 0 {
+		envSpec.Dependencies = deps
+	}
+
 	return jobs.SubmitRun{
 		RunName:        cfg.ExperimentName,
 		TimeoutSeconds: cfg.timeoutSeconds(),
 		Tasks:          []jobs.SubmitTask{st},
 		Environments: []jobs.JobEnvironment{{
 			EnvironmentKey: aiRuntimeEnvironmentKey,
-			Spec:           &compute.Environment{EnvironmentVersion: dlImage},
+			Spec:           envSpec,
 		}},
 	}
 }
@@ -133,6 +163,13 @@ func submitWorkload(ctx context.Context, w *databricks.WorkspaceClient, cfg *run
 		return 0, "", err
 	}
 
+	// Resolve dependencies before any upload too, so a bad requirements file fails
+	// fast without leaving orphaned artifacts in the workspace.
+	deps, fileVersion, err := environmentDependencies(cfg, configPath)
+	if err != nil {
+		return 0, "", err
+	}
+
 	experimentDir := ""
 	if cfg.MLflowExperimentDirectory != nil {
 		experimentDir = *cfg.MLflowExperimentDirectory
@@ -175,8 +212,13 @@ func submitWorkload(ctx context.Context, w *databricks.WorkspaceClient, cfg *run
 		}
 	}
 
-	runtimeVersion, _ := cfg.runtimeVersion()
-	payload := buildSubmitPayload(cfg, path.Join(funcDir, commandScriptName), dlRuntimeImage(ctx, runtimeVersion), snap)
+	// Top-level environment.version wins; for file-form deps it is disallowed, so
+	// fall back to the version declared inside the requirements file.
+	runtimeVersion, ok := cfg.runtimeVersion()
+	if !ok {
+		runtimeVersion = fileVersion
+	}
+	payload := buildSubmitPayload(cfg, path.Join(funcDir, commandScriptName), dlRuntimeImage(ctx, runtimeVersion), snap, deps)
 	payload.IdempotencyToken = token
 
 	// Submit returns as soon as the run is created; we don't wait for it to finish.