AIR CLI Integration: Port code_source Snapshotting

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

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

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

@@ -0,0 +1,47 @@
+
+=== submit with a git code_source
+>>> [CLI] experimental air run -f run.yaml
+Submitted run 555
+View at: [DATABRICKS_URL]/jobs/runs/555
+
+=== the ai_runtime_task carries the snapshot + provenance paths
+>>> print_requests.py //api/2.2/jobs/runs/submit
+{
+  "method": "POST",
+  "path": "/api/2.2/jobs/runs/submit",
+  "body": {
+    "run_name": "submit-smoke",
+    "tasks": [
+      {
+        "task_key": "submit-smoke",
+        "run_if": "ALL_SUCCESS",
+        "ai_runtime_task": {
+          "experiment": "submit-smoke",
+          "deployments": [
+            {
+              "command_path": "/Workspace/Users/[USERNAME]/.air/cli_launch/submit-smoke/submit-smoke_[RUN_ID]/command.sh",
+              "compute": {
+                "accelerator_type": "GPU_1xH100",
+                "accelerator_count": 1
+              }
+            }
+          ],
+          "code_source_path": "/Workspace/Users/[USERNAME]/.air/repo_snapshots/[SNAPSHOT_TARBALL]",
+          "git_state_path": "/Workspace/Users/[USERNAME]/.air/cli_launch/submit-smoke/submit-smoke_[RUN_ID]/git_state.json"
+        },
+        "environment_key": "default",
+        "max_retries": 3,
+        "retry_on_timeout": true
+      }
+    ],
+    "environments": [
+      {
+        "environment_key": "default",
+        "spec": {
+          "environment_version": "4"
+        }
+      }
+    ],
+    "idempotency_token": "[UUID]"
+  }
+}

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

@@ -0,0 +1,11 @@
+experiment_name: submit-smoke
+command: python train.py
+compute:
+  accelerator_type: GPU_1xH100
+  num_accelerators: 1
+code_source:
+  type: snapshot
+  snapshot:
+    root_path: .
+    git:
+      commit: COMMIT_SHA

acceptance/experimental/air/run-submit/script

@@ -0,0 +1,18 @@
+# Pinned commit dates keep the resolved HEAD SHA — and thus the snapshot cache key
+# baked into the tarball name — stable across runs.
+export GIT_AUTHOR_DATE="2020-01-01T00:00:00Z"
+export GIT_COMMITTER_DATE="2020-01-01T00:00:00Z"
+git-repo-init
+
+# Pin the config to the committed HEAD. git.commit tolerates a dirty working tree
+# (the pinned revision is what gets archived), which matters here because the
+# acceptance harness writes output.txt into the working dir as the script runs.
+sed "s/COMMIT_SHA/$(git rev-parse HEAD)/" run.yaml.tmpl > run.yaml
+
+title "submit with a git code_source"
+trace $CLI experimental air run -f run.yaml
+
+title "the ai_runtime_task carries the snapshot + provenance paths"
+trace print_requests.py //api/2.2/jobs/runs/submit
+
+rm -fr .git

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

@@ -0,0 +1,35 @@
+# A real (non-dry-run) submit that packages a git code_source, uploads the
+# tarball + provenance sidecars, and POSTs runs/submit. No bundle deploy, so no
+# engine matrix.
+RecordRequests = true
+
+# run.yaml is generated from run.yaml.tmpl at test time (commit SHA templated in);
+# it isn't a committed input to diff.
+Ignore = ["run.yaml"]
+
+[EnvMatrix]
+DATABRICKS_BUNDLE_ENGINE = []
+
+# 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 snapshot tarball is named <dir>_<cachekey[:16]>.tar.gz, where <dir> is the
+# test's temp-dir basename and the cache key derives from the pinned commit SHA.
+# Both are stable given the pinned commit dates in the script, but the temp-dir
+# basename varies per run, so collapse the whole tarball filename to a stable token.
+[[Repls]]
+Old = '[a-zA-Z0-9._-]+_[0-9a-f]{16}\.tar\.gz'
+New = '[SNAPSHOT_TARBALL]'
+
+# The per-run launch directory ends in <run>_<16 hex>; the random suffix varies.
+[[Repls]]
+Old = 'submit-smoke_[0-9a-f]{16}'
+New = 'submit-smoke_[RUN_ID]'

acceptance/experimental/air/run/git-remote.yaml

@@ -0,0 +1,12 @@
+experiment_name: smoke-test
+command: python train.py
+compute:
+  accelerator_type: GPU_1xH100
+  num_accelerators: 1
+code_source:
+  type: snapshot
+  snapshot:
+    root_path: .
+    git:
+      branch: main
+      remote: origin

acceptance/experimental/air/run/output.txt

@@ -26,6 +26,16 @@ Error: --watch is not yet supported
 
 Exit code: 1
 
+=== code_source config passes validation
+>>> [CLI] experimental air run -f with-code-source.yaml --dry-run
+Dry run: configuration for "smoke-test" is valid; not submitting.
+
+=== git.remote is rejected
+>>> [CLI] experimental air run -f git-remote.yaml --dry-run
+Error: git.remote is no longer supported: the snapshot archives your local copy, so a branch resolves to its local HEAD. To deploy a specific committed revision, use git.commit
+
+Exit code: 1
+
 === invalid config is rejected
 >>> [CLI] experimental air run -f invalid.yaml --dry-run
 Error: invalid experiment_name "bad.name": only alphanumeric characters, hyphens (-), and underscores (_) are allowed

acceptance/experimental/air/run/script

@@ -10,6 +10,12 @@ errcode trace $CLI experimental air run -f valid.yaml --dry-run --override a=b
 title "watch not yet supported"
 errcode trace $CLI experimental air run -f valid.yaml --dry-run --watch
 
+title "code_source config passes validation"
+trace $CLI experimental air run -f with-code-source.yaml --dry-run
+
+title "git.remote is rejected"
+errcode trace $CLI experimental air run -f git-remote.yaml --dry-run
+
 title "invalid config is rejected"
 errcode trace $CLI experimental air run -f invalid.yaml --dry-run
 

acceptance/experimental/air/run/with-code-source.yaml

@@ -0,0 +1,13 @@
+experiment_name: smoke-test
+command: python train.py
+compute:
+  accelerator_type: GPU_1xH100
+  num_accelerators: 1
+code_source:
+  type: snapshot
+  snapshot:
+    root_path: .
+    git:
+      branch: main
+    include_paths:
+      - src

experimental/air/cmd/runconfig.go

@@ -27,8 +27,8 @@ import (
 // Jobs API task_key rejects.
 var taskKeyRe = regexp.MustCompile(`^[A-Za-z0-9_-]+$`)
 
-// gitRefRe guards branch/remote names against command injection. Only safe ref
-// characters are allowed.
+// gitRefRe guards the branch name against command injection (it flows into git
+// exec args). Only safe ref characters are allowed.
 var gitRefRe = regexp.MustCompile(`^[\w./-]+$`)
 
 // runConfig is the top-level run YAML schema: experiment_name + compute /
@@ -372,13 +372,12 @@ func (g *gitRef) validate() error {
 	if g.Branch != nil && !gitRefRe.MatchString(*g.Branch) {
 		return fmt.Errorf("invalid git.branch format %q: only alphanumeric characters, hyphens, dots, slashes, and underscores are allowed", *g.Branch)
 	}
-	if g.Remote.isString {
-		if g.Remote.name == "" {
-			return errors.New("git.remote string cannot be empty; use 'true' to auto-detect")
-		}
-		if !gitRefRe.MatchString(g.Remote.name) {
-			return fmt.Errorf("invalid git.remote name %q: only alphanumeric characters, hyphens, dots, slashes, and underscores are allowed", g.Remote.name)
-		}
+
+	// The remote-fetch path (fetching a branch's remote HEAD) is deprecated: the
+	// snapshot archives the local copy only. A truthy git.remote (a name or `true`)
+	// is rejected; `remote: false` is the default (local HEAD) and stays valid.
+	if g.Remote.truthy() {
+		return errors.New("git.remote is no longer supported: the snapshot archives your local copy, so a branch resolves to its local HEAD. To deploy a specific committed revision, use git.commit")
 	}
 
 	if g.Branch == nil && g.Commit == nil {
@@ -387,9 +386,6 @@ func (g *gitRef) validate() error {
 	if g.Branch != nil && g.Commit != nil {
 		return errors.New("git: 'branch' and 'commit' are mutually exclusive — specify only one")
 	}
-	if g.Remote.truthy() && g.Branch == nil {
-		return errors.New("git.remote requires git.branch (only valid with branch refs)")
-	}
 	return nil
 }
 

experimental/air/cmd/runconfig_test.go

@@ -64,7 +64,6 @@ code_source:
     remote_volume: /Volumes/main/default/code
     git:
       branch: main
-      remote: origin
     include_paths:
       - src
       - configs/train.yaml
@@ -92,8 +91,6 @@ permissions:
 	require.NotNil(t, cfg.CodeSource.Snapshot.Git)
 	require.NotNil(t, cfg.CodeSource.Snapshot.Git.Branch)
 	assert.Equal(t, "main", *cfg.CodeSource.Snapshot.Git.Branch)
-	assert.True(t, cfg.CodeSource.Snapshot.Git.Remote.isString)
-	assert.Equal(t, "origin", cfg.CodeSource.Snapshot.Git.Remote.name)
 	assert.Len(t, cfg.Permissions, 2)
 }
 
@@ -111,8 +108,8 @@ environment:
 		assert.Equal(t, "requirements.yaml", cfg.Environment.Dependencies.path)
 	})
 
-	t.Run("git remote as bool true", func(t *testing.T) {
-		cfg, err := loadRunConfig(writeConfig(t, minimalConfig+`
+	t.Run("git remote as bool true is rejected", func(t *testing.T) {
+		_, err := loadRunConfig(writeConfig(t, minimalConfig+`
 code_source:
   type: snapshot
   snapshot:
@@ -121,11 +118,8 @@ code_source:
       branch: main
       remote: true
 `))
-		require.NoError(t, err)
-		r := cfg.CodeSource.Snapshot.Git.Remote
-		assert.False(t, r.isString)
-		assert.True(t, r.enabled)
-		assert.True(t, r.truthy())
+		require.Error(t, err)
+		assert.Contains(t, err.Error(), "git.remote is no longer supported")
 	})
 
 	t.Run("git remote defaults to false when unset", func(t *testing.T) {
@@ -334,12 +328,12 @@ func TestGitRefValidate(t *testing.T) {
 	}{
 		{"branch only ok", gitRef{Branch: str("main")}, ""},
 		{"commit only ok", gitRef{Commit: str("abc123")}, ""},
-		{"branch with remote ok", gitRef{Branch: str("main"), Remote: gitRemote{set: true, enabled: true}}, ""},
+		{"remote false is ok", gitRef{Branch: str("main"), Remote: gitRemote{set: true, enabled: false}}, ""},
 		{"neither branch nor commit", gitRef{}, "must specify either 'branch' or 'commit'"},
 		{"both branch and commit", gitRef{Branch: str("main"), Commit: str("abc")}, "mutually exclusive"},
-		{"remote without branch", gitRef{Commit: str("abc"), Remote: gitRemote{set: true, isString: true, name: "origin"}}, "requires git.branch"},
+		{"remote true rejected", gitRef{Branch: str("main"), Remote: gitRemote{set: true, enabled: true}}, "git.remote is no longer supported"},
+		{"remote name rejected", gitRef{Branch: str("main"), Remote: gitRemote{set: true, isString: true, name: "origin"}}, "git.remote is no longer supported"},
 		{"bad branch chars", gitRef{Branch: str("bad branch")}, "invalid git.branch"},
-		{"empty remote string", gitRef{Branch: str("main"), Remote: gitRemote{set: true, isString: true, name: ""}}, "cannot be empty"},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {

experimental/air/cmd/runsubmit.go

@@ -52,6 +52,13 @@ type aiRuntimeTask struct {
 	Deployments               []aiRuntimeDeployment `json:"deployments"`
 	MlflowRun                 string                `json:"mlflow_run,omitempty"`
 	MlflowExperimentDirectory string                `json:"mlflow_experiment_directory,omitempty"`
+	// CodeSourcePath points at the uploaded snapshot tarball; the remote
+	// entry_script extracts it to /databricks/code_source/<dir>. GitStatePath and
+	// GitDiffPath are the optional provenance sidecars (see snapshot.go). All are
+	// omitempty: a workspace-only run (no code_source) carries none of them.
+	CodeSourcePath string `json:"code_source_path,omitempty"`
+	GitStatePath   string `json:"git_state_path,omitempty"`
+	GitDiffPath    string `json:"git_diff_path,omitempty"`
 }
 
 // environmentSpec carries the bare runtime channel ("4", "5", ...).
@@ -106,8 +113,9 @@ func dlRuntimeImage(ctx context.Context, runtimeVersion string) string {
 }
 
 // buildSubmitPayload assembles the runs/submit payload. commandPath is the
-// workspace path of the uploaded command.sh; dlImage is the runtime channel.
-func buildSubmitPayload(cfg *runConfig, commandPath, dlImage string) jobsSubmitRun {
+// workspace path of the uploaded command.sh; dlImage is the runtime channel; snap
+// carries the code snapshot paths (zero value when the run has no code_source).
+func buildSubmitPayload(cfg *runConfig, commandPath, dlImage string, snap snapshotResult) jobsSubmitRun {
 	task := aiRuntimeTask{
 		Experiment: cfg.ExperimentName,
 		Deployments: []aiRuntimeDeployment{{
@@ -117,6 +125,9 @@ func buildSubmitPayload(cfg *runConfig, commandPath, dlImage string) jobsSubmitR
 				AcceleratorCount: cfg.Compute.NumAccelerators,
 			},
 		}},
+		CodeSourcePath: snap.CodeSourcePath,
+		GitStatePath:   snap.GitStatePath,
+		GitDiffPath:    snap.GitDiffPath,
 	}
 	if cfg.MLflowRunName != nil {
 		task.MlflowRun = *cfg.MLflowRunName
@@ -195,14 +206,11 @@ func submitToken(flag string, cfg *runConfig) (string, error) {
 // upload the launch artifacts, assemble the Jobs payload, and submit it. It
 // returns the new run_id and its dashboard URL.
 func submitWorkload(ctx context.Context, w *databricks.WorkspaceClient, cfg *runConfig, configPath, idempotencyKey string) (int64, string, error) {
-	// Resolving usage_policy_name to a budget policy id and packaging a
-	// code_source snapshot are not ported yet; reject rather than silently drop.
+	// Resolving usage_policy_name to a budget policy id is not ported yet; reject
+	// rather than silently drop.
 	if cfg.UsagePolicyName != nil {
 		return 0, "", errors.New("usage_policy_name is not yet supported")
 	}
-	if cfg.CodeSource != nil {
-		return 0, "", errors.New("code_source is not yet supported")
-	}
 
 	// Resolve the idempotency token first so a bad key fails before any upload.
 	token, err := submitToken(idempotencyKey, cfg)
@@ -240,8 +248,18 @@ func submitWorkload(ctx context.Context, w *databricks.WorkspaceClient, cfg *run
 		return 0, "", err
 	}
 
+	// Package and upload the code snapshot, if any. The resulting paths ride on the
+	// ai_runtime_task; a run with no code_source leaves them empty.
+	var snap snapshotResult
+	if cfg.CodeSource != nil {
+		snap, err = snapshotCodeSource(ctx, w, cfg.CodeSource.Snapshot, configPath, base, funcDir)
+		if err != nil {
+			return 0, "", err
+		}
+	}
+
 	runtimeVersion, _ := cfg.runtimeVersion()
-	payload := buildSubmitPayload(cfg, path.Join(funcDir, commandScriptName), dlRuntimeImage(ctx, runtimeVersion))
+	payload := buildSubmitPayload(cfg, path.Join(funcDir, commandScriptName), dlRuntimeImage(ctx, runtimeVersion), snap)
 	payload.IdempotencyToken = token
 
 	jc, err := newJobsSubmitClient(w)