air run: package code_source tarball via DABs artifact upload (#6015)

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

@@ -1,6 +1,7 @@
 
 === submit with a git code_source
 >>> [CLI] experimental air run -f run.yaml
+Uploading [SNAPSHOT_TARBALL]...
 Submitted run 555
 View at: [DATABRICKS_URL]/jobs/runs/555
 
@@ -23,7 +24,7 @@ View at: [DATABRICKS_URL]/jobs/runs/555
     "tasks": [
       {
         "ai_runtime_task": {
-          "code_source_path": "/Workspace/Users/[USERNAME]/.air/repo_snapshots/001/[SNAPSHOT_TARBALL]",
+          "code_source_path": "/Workspace/Users/[USERNAME]/.air/repo_snapshots/.internal/[SNAPSHOT_TARBALL]",
           "deployments": [
             {
               "command_path": "/Workspace/Users/[USERNAME]/.air/cli_launch/submit-smoke/submit-smoke_[RUN_ID]/command.sh",

experimental/air/cmd/runsubmit.go

@@ -58,15 +58,6 @@ func buildSubmitPayload(cfg *runConfig, commandPath, dlImage string, snap snapsh
 			},
 		}},
 		CodeSourcePath: snap.CodeSourcePath,
-		// TEMP: git_state_path / git_diff_path are intentionally NOT sent. The typed
-		// jobs.AiRuntimeTask (and its source proto, ai_runtime_task.proto) has no such
-		// fields, so the typed SDK path cannot carry them. This is safe today because
-		// nothing in the backend consumes those fields — the AI Runtime task proto
-		// never declared them, so even the Python CLI's raw-JSON values were dropped
-		// on deserialization. The git_state.json / git_diff.patch sidecars are still
-		// uploaded next to the tarball (see snapshot.go) for human inspection.
-		// If the backend later adds these fields to the proto, regenerate the SDK and
-		// wire snap.GitStatePath / snap.GitDiffPath back in here.
 	}
 	if cfg.MLflowRunName != nil {
 		task.MlflowRun = *cfg.MLflowRunName
@@ -163,13 +154,13 @@ 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. Snapshot is the
-	// only code_source type; guard against a nil block so snapshotCodeSource never
-	// dereferences a missing snapshot.
+	// Package and upload the code snapshot, if any, via DABs' artifact-upload
+	// plumbing; the remote code_source_path rides the ai_runtime_task. A run with no
+	// code_source leaves it empty. Snapshot is the only code_source type.
 	var snap snapshotResult
 	if cfg.CodeSource != nil && cfg.CodeSource.Snapshot != nil {
-		snap, err = snapshotCodeSource(ctx, w, cfg.CodeSource.Snapshot, configPath, base, funcDir)
+		// Sidecars land in the run's launch dir (funcDir) via fc, next to command.sh.
+		snap, err = snapshotViaDABsUpload(ctx, w, cfg.CodeSource.Snapshot, configPath, fc, funcDir)
 		if err != nil {
 			return 0, "", err
 		}

experimental/air/cmd/runsubmit_test.go

@@ -2,10 +2,14 @@ package aircmd
 
 import (
 	"encoding/json"
+	"io"
+	"path"
 	"path/filepath"
 	"strings"
 	"testing"
 
+	"github.com/databricks/cli/libs/cmdio"
+	"github.com/databricks/cli/libs/filer"
 	"github.com/databricks/cli/libs/testserver"
 	"github.com/databricks/databricks-sdk-go"
 	"github.com/databricks/databricks-sdk-go/service/jobs"
@@ -156,8 +160,8 @@ func TestSubmitWorkload(t *testing.T) {
 	assert.Equal(t, jobs.ComputeSpec{AcceleratorType: jobs.ComputeSpecAcceleratorTypeGpu1xH100, AcceleratorCount: 1}, d.Compute)
 }
 
-// TestSubmitWorkloadWithCodeSource exercises the snapshot path end to end: a
-// git-pinned code_source is packaged, uploaded, and its paths attached to the task.
+// A working-tree code_source is packaged into a tarball, uploaded via DABs' artifact
+// plumbing, and its remote code_source_path attached to the submitted task.
 func TestSubmitWorkloadWithCodeSource(t *testing.T) {
 	server := testserver.New(t)
 	t.Cleanup(server.Close)
@@ -172,6 +176,47 @@ func TestSubmitWorkloadWithCodeSource(t *testing.T) {
 	w, err := databricks.NewWorkspaceClient(&databricks.Config{Host: server.URL, Token: "token"})
 	require.NoError(t, err)
 
+	// A plain working-tree directory: packaging is plain-tar.
+	repo := filepath.Join(t.TempDir(), "src")
+	writeRepoFile(t, repo, "train.py", "print()")
+
+	cfg := minimalConfig + `
+code_source:
+  type: snapshot
+  snapshot:
+    root_path: ` + repo + `
+`
+	cfgPath := writeConfigFile(t, "run.yaml", cfg)
+	loaded, err := loadRunConfig(cfgPath)
+	require.NoError(t, err)
+
+	// The DABs upload path logs via cmdio; the real `air run` context carries it.
+	ctx := cmdio.MockDiscard(t.Context())
+	_, _, err = submitWorkload(ctx, w, loaded, cfgPath, "idem")
+	require.NoError(t, err)
+
+	at := got.Tasks[0].AiRuntimeTask
+	// The tarball is uploaded to the artifact .internal dir and code_source_path
+	// rewritten to it.
+	assert.Contains(t, at.CodeSourcePath, "/.air/repo_snapshots/.internal/")
+	assert.True(t, strings.HasSuffix(at.CodeSourcePath, ".tar.gz"), at.CodeSourcePath)
+}
+
+// A git-pinned code_source is git-archived at the commit, uploaded via DABs' artifact
+// plumbing, and its remote code_source_path attached to the submitted task.
+func TestSubmitWorkloadWithGitPinnedCodeSource(t *testing.T) {
+	server := testserver.New(t)
+	t.Cleanup(server.Close)
+
+	var got jobs.SubmitRun
+	server.Handle("POST", "/api/2.2/jobs/runs/submit", func(req testserver.Request) any {
+		require.NoError(t, json.Unmarshal(req.Body, &got))
+		return jobs.SubmitRunResponse{RunId: 555}
+	})
+	testserver.AddDefaultHandlers(server)
+	w, err := databricks.NewWorkspaceClient(&databricks.Config{Host: server.URL, Token: "token"})
+	require.NoError(t, err)
+
 	// A git repo committed at HEAD, referenced by commit so packaging is git_archive.
 	repo := newTestRepo(t)
 	writeRepoFile(t, repo, "train.py", "print()")
@@ -189,15 +234,209 @@ code_source:
 	loaded, err := loadRunConfig(cfgPath)
 	require.NoError(t, err)
 
-	_, _, err = submitWorkload(t.Context(), w, loaded, cfgPath, "idem")
+	ctx := cmdio.MockDiscard(t.Context())
+	_, _, err = submitWorkload(ctx, w, loaded, cfgPath, "idem")
+	require.NoError(t, err)
+
+	at := got.Tasks[0].AiRuntimeTask
+	assert.Contains(t, at.CodeSourcePath, "/.air/repo_snapshots/.internal/")
+	assert.True(t, strings.HasSuffix(at.CodeSourcePath, ".tar.gz"), at.CodeSourcePath)
+}
+
+// testSidecarStore builds a workspace filer + base path standing in for the run's
+// launch dir, where snapshotViaDABsUpload writes git provenance sidecars.
+func testSidecarStore(t *testing.T, w *databricks.WorkspaceClient) (filer.Filer, string) {
+	t.Helper()
+	base := "/Workspace/Users/tester@databricks.com/.air/cli_launch/test"
+	f, err := filer.NewWorkspaceFilesClient(w, base)
+	require.NoError(t, err)
+	return f, base
+}
+
+// A plain-tar (working-tree) snapshot is uploaded under a unique, timestamped name so
+// two concurrent submissions of the same root_path don't clobber each other's upload.
+func TestSubmitWorkloadPlainTarNameIsUnique(t *testing.T) {
+	server := testserver.New(t)
+	t.Cleanup(server.Close)
+
+	server.Handle("POST", "/api/2.2/jobs/runs/submit", func(req testserver.Request) any {
+		return jobs.SubmitRunResponse{RunId: 555}
+	})
+	testserver.AddDefaultHandlers(server)
+	w, err := databricks.NewWorkspaceClient(&databricks.Config{Host: server.URL, Token: "token"})
+	require.NoError(t, err)
+
+	// A plain working-tree directory named "src": the old code named the tarball
+	// after the dir alone (src.tar.gz), so any two submissions collided.
+	repo := filepath.Join(t.TempDir(), "src")
+	writeRepoFile(t, repo, "train.py", "print()")
+
+	cfg := minimalConfig + `
+code_source:
+  type: snapshot
+  snapshot:
+    root_path: ` + repo + `
+`
+	cfgPath := writeConfigFile(t, "run.yaml", cfg)
+	loaded, err := loadRunConfig(cfgPath)
+	require.NoError(t, err)
+
+	// The uploaded name carries a discriminator (timestamp), not the bare dir name.
+	ctx := cmdio.MockDiscard(t.Context())
+	sidecarStore, sidecarBase := testSidecarStore(t, w)
+	snap, err := snapshotViaDABsUpload(ctx, w, loaded.CodeSource.Snapshot, cfgPath, sidecarStore, sidecarBase)
+	require.NoError(t, err)
+	base := path.Base(snap.CodeSourcePath)
+	assert.NotEqual(t, "src.tar.gz", base, "plain-tar name must be unique, not the bare dir name")
+	assert.Regexp(t, `^src_\d{8}_\d{6}\.tar\.gz$`, base)
+}
+
+// A git_archive snapshot is content-addressed by (commit, include_paths): submitting
+// the same commit twice reuses the already-uploaded tarball and skips the second
+// upload (cache hit), while resolving to the identical remote path.
+func TestSubmitWorkloadGitArchiveCaching(t *testing.T) {
+	server := testserver.New(t)
+	t.Cleanup(server.Close)
+
+	server.Handle("POST", "/api/2.2/jobs/runs/submit", func(req testserver.Request) any {
+		return jobs.SubmitRunResponse{RunId: 555}
+	})
+	// Track which snapshot tarballs get uploaded, preserving fake-workspace
+	// persistence so the second submit's cache-existence Stat sees the first upload.
+	// Dedupe by path: the DABs uploader mkdirs-and-retries the import on a missing
+	// parent dir, so one logical upload can hit this route more than once.
+	uploaded := map[string]bool{}
+	server.Handle("POST", "/api/2.0/workspace-files/import-file/{path...}", func(req testserver.Request) any {
+		p := req.Vars["path"]
+		if strings.Contains(p, "/.air/repo_snapshots/") {
+			uploaded[p] = true
+		}
+		return req.Workspace.WorkspaceFilesImportFile(p, req.Body, req.URL.Query().Get("overwrite") == "true")
+	})
+	testserver.AddDefaultHandlers(server)
+	w, err := databricks.NewWorkspaceClient(&databricks.Config{Host: server.URL, Token: "token"})
+	require.NoError(t, err)
+
+	repo := newTestRepo(t)
+	writeRepoFile(t, repo, "train.py", "print()")
+	sha := commitAll(t, repo, "init")
+
+	cfg := minimalConfig + `
+code_source:
+  type: snapshot
+  snapshot:
+    root_path: ` + repo + `
+    git:
+      commit: ` + sha + `
+`
+	cfgPath := writeConfigFile(t, "run.yaml", cfg)
+	loaded, err := loadRunConfig(cfgPath)
+	require.NoError(t, err)
+
+	ctx := cmdio.MockDiscard(t.Context())
+	sidecarStore, sidecarBase := testSidecarStore(t, w)
+	first, err := snapshotViaDABsUpload(ctx, w, loaded.CodeSource.Snapshot, cfgPath, sidecarStore, sidecarBase)
+	require.NoError(t, err)
+	second, err := snapshotViaDABsUpload(ctx, w, loaded.CodeSource.Snapshot, cfgPath, sidecarStore, sidecarBase)
+	require.NoError(t, err)
+
+	// Same pinned commit → identical content-addressed remote path, uploaded once
+	// (the second submit is a cache hit and moves no bytes).
+	assert.Equal(t, first.CodeSourcePath, second.CodeSourcePath)
+	assert.Len(t, uploaded, 1, "git_archive cache hit should skip the second upload")
+}
+
+// A git code_source also uploads git provenance sidecars (git_state.json, and
+// git_diff.patch when the tree is dirty) next to the run's launch dir, so the
+// submitted commit + working-tree diff are inspectable.
+func TestSubmitWorkloadUploadsGitSidecars(t *testing.T) {
+	server := testserver.New(t)
+	t.Cleanup(server.Close)
+
+	server.Handle("POST", "/api/2.2/jobs/runs/submit", func(req testserver.Request) any {
+		return jobs.SubmitRunResponse{RunId: 555}
+	})
+	testserver.AddDefaultHandlers(server)
+	w, err := databricks.NewWorkspaceClient(&databricks.Config{Host: server.URL, Token: "token"})
+	require.NoError(t, err)
+
+	// Commit, then dirty the tree so both git_state.json and git_diff.patch are produced.
+	repo := newTestRepo(t)
+	writeRepoFile(t, repo, "train.py", "print()")
+	commitAll(t, repo, "init")
+	writeRepoFile(t, repo, "train.py", "print('dirty')")
+
+	cfg := minimalConfig + `
+code_source:
+  type: snapshot
+  snapshot:
+    root_path: ` + repo + `
+`
+	cfgPath := writeConfigFile(t, "run.yaml", cfg)
+	loaded, err := loadRunConfig(cfgPath)
+	require.NoError(t, err)
+
+	ctx := cmdio.MockDiscard(t.Context())
+	sidecarStore, sidecarBase := testSidecarStore(t, w)
+	snap, err := snapshotViaDABsUpload(ctx, w, loaded.CodeSource.Snapshot, cfgPath, sidecarStore, sidecarBase)
+	require.NoError(t, err)
+
+	// Both sidecars are reported under the launch dir and actually exist there.
+	assert.Equal(t, path.Join(sidecarBase, gitStateName), snap.GitStatePath)
+	assert.Equal(t, path.Join(sidecarBase, gitDiffName), snap.GitDiffPath)
+
+	r, err := sidecarStore.Read(ctx, gitStateName)
+	require.NoError(t, err)
+	stateBytes, err := io.ReadAll(r)
+	require.NoError(t, err)
+	var state map[string]any
+	require.NoError(t, json.Unmarshal(stateBytes, &state))
+	assert.Equal(t, "plain_tar", state["packaging_mode"])
+	assert.Equal(t, true, state["dirty"])
+	assert.Equal(t, "captured", state["diff_status"])
+}
+
+// remote_volume uploads the snapshot to a UC Volume: DABs' artifact uploader handles
+// /Volumes destinations natively, so code_source_path lands under the Volume path.
+func TestSubmitWorkloadWithRemoteVolumeCodeSource(t *testing.T) {
+	server := testserver.New(t)
+	t.Cleanup(server.Close)
+
+	var got jobs.SubmitRun
+	server.Handle("POST", "/api/2.2/jobs/runs/submit", func(req testserver.Request) any {
+		require.NoError(t, json.Unmarshal(req.Body, &got))
+		return jobs.SubmitRunResponse{RunId: 555}
+	})
+	// Stub the UC Volume file write: the fake server's default handler 404s when the
+	// parent dir is absent (no auto-mkdir), so accept the PUT to exercise the Volume
+	// upload route. This asserts we route to /api/2.0/fs/files/Volumes/... at all.
+	server.Handle("PUT", "/api/2.0/fs/files/Volumes/{path...}", func(req testserver.Request) any {
+		return testserver.Response{StatusCode: 204}
+	})
+	testserver.AddDefaultHandlers(server)
+	w, err := databricks.NewWorkspaceClient(&databricks.Config{Host: server.URL, Token: "token"})
+	require.NoError(t, err)
+
+	repo := filepath.Join(t.TempDir(), "src")
+	writeRepoFile(t, repo, "train.py", "print()")
+
+	cfg := minimalConfig + `
+code_source:
+  type: snapshot
+  snapshot:
+    root_path: ` + repo + `
+    remote_volume: /Volumes/main/default/code
+`
+	cfgPath := writeConfigFile(t, "run.yaml", cfg)
+	loaded, err := loadRunConfig(cfgPath)
+	require.NoError(t, err)
+
+	ctx := cmdio.MockDiscard(t.Context())
+	_, _, err = submitWorkload(ctx, w, loaded, cfgPath, "idem")
 	require.NoError(t, err)
 
 	at := got.Tasks[0].AiRuntimeTask
-	// The tarball path is under the user's repo_snapshots dir. git_state_path /
-	// git_diff_path are not asserted: the typed jobs.AiRuntimeTask has no such fields
-	// (see the TEMP note in buildSubmitPayload), so they aren't sent. The git_state
-	// sidecar file is still uploaded next to the tarball — covered by TestRunSnapshot.
-	assert.Contains(t, at.CodeSourcePath, "/.air/repo_snapshots/"+filepath.Base(repo)+"/")
+	assert.Contains(t, at.CodeSourcePath, "/Volumes/main/default/code/.internal/")
 	assert.True(t, strings.HasSuffix(at.CodeSourcePath, ".tar.gz"), at.CodeSourcePath)
 }