air run: git-pin code_source snapshots

experimental/air/cmd/runsubmit_test.go

@@ -199,6 +199,47 @@ code_source:
 	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()")
+	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())
+	_, _, 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)
+}
+
 func TestSubmitWorkloadGuards(t *testing.T) {
 	w := newFakeWorkspaceClient(t)
 	cfgPath := writeConfigFile(t, "run.yaml", minimalConfig)

experimental/air/cmd/snapshot_cachekey.go

@@ -0,0 +1,34 @@
+package aircmd
+
+// This file packages a local code directory into a tarball, uploads it to the
+// workspace (or a Volume), and records git provenance sidecars for cache
+// invalidation — the Go port of the Python CLI's code_source snapshot path.
+
+import (
+	"crypto/sha256"
+	"encoding/hex"
+	"slices"
+	"strings"
+)
+
+// snapshotPackagingVersion is bumped when packaging logic changes in a way that invalidates existing caches
+const snapshotPackagingVersion = "v1"
+
+// computeSnapshotCacheKey returns a stable cache key for a snapshot tarball: the
+// SHA-256 digest of (commitSHA, normalized includePaths, snapshotPackagingVersion).
+// Changing any input yields a different entry.
+func computeSnapshotCacheKey(commitSHA string, includePaths []string) string {
+	var normalizedPaths string
+	if len(includePaths) > 0 {
+		trimmed := make([]string, len(includePaths))
+		for i, p := range includePaths {
+			trimmed[i] = strings.TrimSpace(p)
+		}
+		slices.Sort(trimmed)
+		normalizedPaths = strings.Join(trimmed, "\n")
+	}
+
+	keyMaterial := commitSHA + "\n" + normalizedPaths + "\n" + snapshotPackagingVersion
+	sum := sha256.Sum256([]byte(keyMaterial))
+	return hex.EncodeToString(sum[:])
+}

experimental/air/cmd/snapshot_cachekey_test.go

@@ -0,0 +1,62 @@
+package aircmd
+
+import (
+	"encoding/json"
+	"os"
+	"path/filepath"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+type goldenCase struct {
+	Name         string   `json:"name"`
+	CommitSHA    string   `json:"commit_sha"`
+	IncludePaths []string `json:"include_paths"`
+	CacheKey     string   `json:"cache_key"`
+}
+
+// TestComputeSnapshotCacheKeyGolden asserts byte-for-byte parity with golden
+// fixtures across the local-only matrix (commit + include_paths permutations).
+func TestComputeSnapshotCacheKeyGolden(t *testing.T) {
+	data, err := os.ReadFile(filepath.Join("testdata", "cache_keys.json"))
+	require.NoError(t, err)
+
+	var cases []goldenCase
+	require.NoError(t, json.Unmarshal(data, &cases))
+	require.NotEmpty(t, cases)
+
+	for _, tc := range cases {
+		t.Run(tc.Name, func(t *testing.T) {
+			assert.Equal(t, tc.CacheKey, computeSnapshotCacheKey(tc.CommitSHA, tc.IncludePaths))
+		})
+	}
+}
+
+// TestComputeSnapshotCacheKeyProperties pins the normalization behavior the golden cases
+// encode, so a regression is legible without decoding hashes.
+func TestComputeSnapshotCacheKeyProperties(t *testing.T) {
+	sha := "a3492b801c0ffee00000000000000000000dead"
+
+	// Order-independent: sorting means unsorted input yields the sorted key.
+	assert.Equal(t,
+		computeSnapshotCacheKey(sha, []string{"a", "b", "c"}),
+		computeSnapshotCacheKey(sha, []string{"c", "a", "b"}),
+	)
+
+	// nil and empty include_paths are equivalent (both contribute an empty line).
+	assert.Equal(t, computeSnapshotCacheKey(sha, nil), computeSnapshotCacheKey(sha, []string{}))
+
+	// Paths are trimmed before hashing.
+	assert.Equal(t,
+		computeSnapshotCacheKey(sha, []string{"research", "data"}),
+		computeSnapshotCacheKey(sha, []string{"  research  ", "  data "}),
+	)
+
+	// Duplicates are NOT collapsed — they are sorted and kept, matching Python.
+	assert.NotEqual(t, computeSnapshotCacheKey(sha, []string{"x", "y"}), computeSnapshotCacheKey(sha, []string{"x", "x", "y"}))
+
+	// The version constant participates: a different version is a different key.
+	assert.NotEqual(t, snapshotPackagingVersion, "")
+}

experimental/air/cmd/snapshot_dabs.go

@@ -31,16 +31,13 @@ func snapshotViaDABsUpload(ctx context.Context, w *databricks.WorkspaceClient, s
 	if snap.RemoteVolume != nil {
 		return snapshotResult{}, errors.New("code_source.snapshot.remote_volume is not yet supported")
 	}
-	if snap.Git != nil {
-		return snapshotResult{}, errors.New("code_source.snapshot.git is not yet supported on this path")
-	}
 
 	repoPath, err := resolveRootPath(ctx, snap.RootPath, filepath.Dir(configPath))
 	if err != nil {
 		return snapshotResult{}, err
 	}
 
-	tarball, cleanup, err := buildSnapshotTarball(ctx, repoPath, snap.IncludePaths)
+	tarball, cleanup, err := buildSnapshotTarball(ctx, repoPath, snap.Git, snap.IncludePaths)
 	if err != nil {
 		return snapshotResult{}, err
 	}
@@ -49,10 +46,11 @@ func snapshotViaDABsUpload(ctx context.Context, w *databricks.WorkspaceClient, s
 	return uploadSnapshotViaDABs(ctx, w, tarball)
 }
 
-// buildSnapshotTarball writes a plain-tar snapshot of repoPath to a temp file named
+// buildSnapshotTarball writes a snapshot of repoPath to a temp file named
 // <dirName>.tar.gz (the basename becomes the uploaded filename), returning the path
-// and a cleanup func.
-func buildSnapshotTarball(ctx context.Context, repoPath string, includePaths []string) (string, func(), error) {
+// and a cleanup func. With a git ref it archives the resolved commit (git archive);
+// otherwise it plain-tars the working tree.
+func buildSnapshotTarball(ctx context.Context, repoPath string, ref *gitRef, includePaths []string) (string, func(), error) {
 	noop := func() {}
 	tmp, err := os.MkdirTemp("", "air-snapshot-*")
 	if err != nil {
@@ -61,7 +59,22 @@ func buildSnapshotTarball(ctx context.Context, repoPath string, includePaths []s
 	cleanup := func() { _ = os.RemoveAll(tmp) }
 
 	tarball := filepath.Join(tmp, filepath.Base(repoPath)+".tar.gz")
-	if err := createPlainTarball(ctx, repoPath, tarball, includePaths); err != nil {
+	dirName := filepath.Base(repoPath)
+
+	git := newGitRepo(repoPath)
+	plan, err := resolveSnapshotPlan(ctx, git, ref, includePaths)
+	if err != nil {
+		cleanup()
+		return "", noop, err
+	}
+
+	switch plan.mode {
+	case modeGitArchive:
+		err = createGitArchiveSnapshot(ctx, git, plan.commitSHA, tarball, dirName, includePaths)
+	default:
+		err = createPlainTarball(ctx, repoPath, tarball, includePaths)
+	}
+	if err != nil {
 		cleanup()
 		return "", noop, err
 	}