Scope git-pinned AIR snapshots to root_path

acceptance/experimental/air/config-help/output.txt

@@ -123,6 +123,14 @@ config.code_source.snapshot.git.remote
   Type:     bool or string
   Required: no
 
+=== git-pinned subdirectory snapshots document root_path scoping
+>>> [CLI] experimental air run -h config.code_source.snapshot.root_path
+config.code_source.snapshot.root_path
+  Root of the code source to archive. A git-pinned subdirectory packages only that subtree.
+
+  Type:     string
+  Required: when code_source.snapshot is set
+
 === the config. prefix is optional
 >>> [CLI] experimental air run -h compute.num_accelerators
 config.compute.num_accelerators

acceptance/experimental/air/config-help/script

@@ -27,6 +27,9 @@ trace $CLI experimental air run -h config.environment.docker_image.url
 title "union field reports both accepted shapes"
 trace $CLI experimental air run -h config.code_source.snapshot.git.remote
 
+title "git-pinned subdirectory snapshots document root_path scoping"
+trace $CLI experimental air run -h config.code_source.snapshot.root_path
+
 title "the config. prefix is optional"
 trace $CLI experimental air run -h compute.num_accelerators
 

experimental/air/cmd/runconfig.go

@@ -384,7 +384,7 @@ func (c *codeSourceConfig) validate() error {
 
 // snapshotSourceConfig describes a local directory to tar and upload.
 type snapshotSourceConfig struct {
-	RootPath     string   `yaml:"root_path" help:"Local directory to archive, relative or absolute." required:"when code_source.snapshot is set"`
+	RootPath     string   `yaml:"root_path" help:"Root of the code source to archive. A git-pinned subdirectory packages only that subtree." required:"when code_source.snapshot is set"`
 	RemoteVolume *string  `yaml:"remote_volume" help:"Volume to upload the archive to. Must start with /Volumes/."`
 	Git          *gitRef  `yaml:"git" help:"Pin the snapshot to a specific git revision."`
 	IncludePaths []string `yaml:"include_paths" help:"Restrict the archive to these paths, relative to root_path and without \"..\". Omit to include everything."`

experimental/air/cmd/snapshot_cachekey.go

@@ -15,9 +15,10 @@ import (
 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 {
+// SHA-256 digest of (commitSHA, normalized includePaths, snapshotPackagingVersion,
+// subtreePrefix). Changing any input yields a different entry. An empty subtree
+// prefix keeps repository-root keys byte-identical to prior versions.
+func computeSnapshotCacheKey(commitSHA string, includePaths []string, subtreePrefix string) string {
 	var normalizedPaths string
 	if len(includePaths) > 0 {
 		trimmed := make([]string, len(includePaths))
@@ -29,6 +30,9 @@ func computeSnapshotCacheKey(commitSHA string, includePaths []string) string {
 	}
 
 	keyMaterial := commitSHA + "\n" + normalizedPaths + "\n" + snapshotPackagingVersion
+	if subtreePrefix != "" {
+		keyMaterial += "\n" + subtreePrefix
+	}
 	sum := sha256.Sum256([]byte(keyMaterial))
 	return hex.EncodeToString(sum[:])
 }

experimental/air/cmd/snapshot_cachekey_test.go

@@ -29,7 +29,7 @@ func TestComputeSnapshotCacheKeyGolden(t *testing.T) {
 
 	for _, tc := range cases {
 		t.Run(tc.Name, func(t *testing.T) {
-			assert.Equal(t, tc.CacheKey, computeSnapshotCacheKey(tc.CommitSHA, tc.IncludePaths))
+			assert.Equal(t, tc.CacheKey, computeSnapshotCacheKey(tc.CommitSHA, tc.IncludePaths, ""))
 		})
 	}
 }
@@ -41,21 +41,31 @@ func TestComputeSnapshotCacheKeyProperties(t *testing.T) {
 
 	// 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"}),
+		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{}))
+	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 "}),
+		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"}))
+	assert.NotEqual(t, computeSnapshotCacheKey(sha, []string{"x", "y"}, ""), computeSnapshotCacheKey(sha, []string{"x", "x", "y"}, ""))
+
+	// A subtree gets a distinct key while an empty prefix preserves repository-root
+	// keys. Different subtrees at one commit cannot collide, even if their root
+	// directories have the same base name.
+	rootKey := computeSnapshotCacheKey(sha, nil, "")
+	assert.NotEqual(t, rootKey, computeSnapshotCacheKey(sha, nil, "team_a/src"))
+	assert.NotEqual(t,
+		computeSnapshotCacheKey(sha, nil, "team_a/src"),
+		computeSnapshotCacheKey(sha, nil, "team_b/src"),
+	)
 
 	// The version constant participates: a different version is a different key.
 	assert.NotEqual(t, snapshotPackagingVersion, "")

experimental/air/cmd/snapshot_dabs.go

@@ -68,11 +68,12 @@ func snapshotViaDABsUpload(ctx context.Context, w *databricks.WorkspaceClient, s
 	// than failing an otherwise-valid submission.
 	//
 	// The sidecars are deliberately NOT bundled into the code tarball. The git_archive
-	// tarball is content-addressed and cached by (commit, include_paths), so a second
-	// run at the same commit reuses it; but the sidecars vary per run (git_state's
-	// timestamp, and git_diff captures the working tree at submit time). Folding them
-	// in would force a distinct tarball per run (defeating the cache) or serve a prior
-	// run's stale provenance on a cache hit. They also live in the per-run launch dir,
+	// tarball is content-addressed and cached by (commit, include_paths, root_path
+	// subtree), so a second identical snapshot reuses it; but the sidecars vary per
+	// run (git_state's timestamp, and git_diff captures the working tree at submit
+	// time). Folding them in would force a distinct tarball per run (defeating the
+	// cache) or serve a prior run's stale provenance on a cache hit. They also live in
+	// the per-run launch dir,
 	// not the shared artifact dir, so they don't accumulate. Keep them out of the tar.
 	if uploadProvenanceSidecars && plan.isGitRepo {
 		result.GitStatePath, result.GitDiffPath = uploadSnapshotSidecars(ctx, sidecarStore, sidecarBase, newGitRepo(repoPath), plan)
@@ -127,14 +128,14 @@ func uploadSnapshotSidecars(ctx context.Context, sidecarStore filer.Filer, sidec
 }
 
 // snapshotTarballName is the uploaded filename for the snapshot. It is deterministic
-// for git_archive — <dirName>_<cacheKey>.tar.gz keyed on (commit, include_paths) — so
-// an identical commit reuses the same remote object (see the cache check below). For
-// plain_tar it is timestamped so concurrent submissions of the same directory don't
-// clobber each other's upload (working-tree content isn't pinned to a SHA, so it
-// can't be content-addressed).
+// for git_archive — <dirName>_<cacheKey>.tar.gz keyed on (commit, include_paths,
+// root_path subtree) — so an identical snapshot reuses the same remote object (see
+// the cache check below). For plain_tar it is timestamped so concurrent submissions
+// of the same directory don't clobber each other's upload (working-tree content
+// isn't pinned to a SHA, so it can't be content-addressed).
 func snapshotTarballName(plan snapshotPlan, dirName string) string {
 	if plan.mode == modeGitArchive {
-		key := computeSnapshotCacheKey(plan.commitSHA, plan.includePaths)
+		key := computeSnapshotCacheKey(plan.commitSHA, plan.includePaths, plan.subtreePrefix)
 		return fmt.Sprintf("%s_%s.tar.gz", dirName, key[:16])
 	}
 	return fmt.Sprintf("%s_%s.tar.gz", dirName, time.Now().UTC().Format("20060102_150405"))
@@ -145,7 +146,7 @@ func snapshotTarballName(plan snapshotPlan, dirName string) string {
 func packageSnapshot(ctx context.Context, repoPath string, plan snapshotPlan, tarball string) error {
 	dirName := filepath.Base(repoPath)
 	if plan.mode == modeGitArchive {
-		return createGitArchiveSnapshot(ctx, newGitRepo(repoPath), plan.commitSHA, tarball, dirName, plan.includePaths)
+		return createGitArchiveSnapshot(ctx, newGitRepo(repoPath), plan.commitSHA, tarball, dirName, plan.includePaths, plan.subtreePrefix)
 	}
 	return createPlainTarball(ctx, repoPath, tarball, plan.includePaths, plan.isGitRepo)
 }
@@ -157,8 +158,8 @@ func packageSnapshot(ctx context.Context, repoPath string, plan snapshotPlan, ta
 // tarball goes to that UC Volume; otherwise to the user's repo_snapshots dir.
 //
 // git_archive snapshots are cacheable: the tarball name is content-addressed by
-// (commit, include_paths), so if the identical object is already uploaded we skip
-// packaging and upload entirely and just reuse the remote path.
+// (commit, include_paths, root_path subtree), so if the identical object is already
+// uploaded we skip packaging and upload entirely and just reuse the remote path.
 func uploadSnapshotViaDABs(ctx context.Context, w *databricks.WorkspaceClient, repoPath string, plan snapshotPlan, remoteVolume string) (snapshotResult, error) {
 	// artifactPath is where DABs uploads the tarball; GetFilerForLibraries routes to
 	// a Workspace or Volume filer based on its prefix, then appends /.internal.
@@ -210,9 +211,9 @@ func uploadSnapshotViaDABs(ctx context.Context, w *databricks.WorkspaceClient, r
 		return snapshotResult{}, err
 	}
 
-	// git_archive is cacheable by (commit, include_paths): if the identical tarball is
-	// already uploaded, skip packaging + upload and reuse it. Only the config-path
-	// rewrite (ReplaceWithRemotePath) runs — no bytes move.
+	// git_archive is cacheable by (commit, include_paths, root_path subtree): if the
+	// identical tarball is already uploaded, skip packaging + upload and reuse it.
+	// Only the config-path rewrite (ReplaceWithRemotePath) runs — no bytes move.
 	if plan.mode == modeGitArchive {
 		f, uploadPath, diags := libraries.GetFilerForLibraries(ctx, b)
 		if diags.HasError() {

experimental/air/cmd/snapshot_git.go

@@ -62,6 +62,28 @@ func (g gitRepo) isRepository(ctx context.Context) bool {
 	return strings.TrimSpace(out) == "true"
 }
 
+// repoRelativePrefix returns the path from the repository root to g.path. Git
+// emits a trailing slash for subdirectories and an empty string at the root.
+func (g gitRepo) repoRelativePrefix(ctx context.Context) (string, error) {
+	out, err := g.run(ctx, "rev-parse", "--show-prefix")
+	if err != nil {
+		return "", fmt.Errorf("failed to resolve repository-relative path for %s: %w", g.path, err)
+	}
+	out = strings.TrimSuffix(out, "\n")
+	out = strings.TrimSuffix(out, "\r")
+	return strings.TrimSuffix(out, "/"), nil
+}
+
+// repositoryRoot returns the top-level directory of the work tree containing
+// g.path.
+func (g gitRepo) repositoryRoot(ctx context.Context) (string, error) {
+	out, err := g.run(ctx, "rev-parse", "--show-toplevel")
+	if err != nil {
+		return "", fmt.Errorf("failed to resolve repository root for %s: %w", g.path, err)
+	}
+	return strings.TrimSpace(out), nil
+}
+
 // headSHA returns the current HEAD commit SHA.
 func (g gitRepo) headSHA(ctx context.Context) (string, error) {
 	out, err := g.run(ctx, "rev-parse", "HEAD")
@@ -164,6 +186,23 @@ func (g gitRepo) mergeBaseWithUpstream(ctx context.Context, remoteName string) s
 	return ""
 }
 
+// validateSubtreeExists checks that subtreePrefix is a directory at commitSHA.
+func (g gitRepo) validateSubtreeExists(ctx context.Context, commitSHA, subtreePrefix string) error {
+	if subtreePrefix == "" {
+		return nil
+	}
+
+	treeish := commitSHA + ":" + subtreePrefix
+	out, err := g.run(ctx, "cat-file", "-t", treeish)
+	if err != nil {
+		return fmt.Errorf("root_path %q does not exist at commit %s: %w", subtreePrefix, shortSHA(commitSHA), err)
+	}
+	if strings.TrimSpace(out) != "tree" {
+		return fmt.Errorf("root_path %q is not a directory at commit %s", subtreePrefix, shortSHA(commitSHA))
+	}
+	return nil
+}
+
 // validateIncludePathsExist checks that every include path exists at commitSHA.
 // `git ls-tree` (without -d, so both blobs and trees count) reports an entry when the
 // path exists; empty output means missing.

experimental/air/cmd/snapshot_git_test.go

@@ -65,6 +65,37 @@ func TestGitRepo_IsRepository(t *testing.T) {
 	assert.False(t, newGitRepo(t.TempDir()).isRepository(ctx))
 }
 
+func TestGitRepo_RepositoryLayout(t *testing.T) {
+	ctx := t.Context()
+	repo := newTestRepo(t)
+	writeRepoFile(t, repo, "a/b/train.py", "print()")
+
+	prefix, err := newGitRepo(repo).repoRelativePrefix(ctx)
+	require.NoError(t, err)
+	assert.Empty(t, prefix)
+
+	subdir := filepath.Join(repo, "a", "b")
+	g := newGitRepo(subdir)
+	prefix, err = g.repoRelativePrefix(ctx)
+	require.NoError(t, err)
+	assert.Equal(t, "a/b", prefix)
+
+	root, err := g.repositoryRoot(ctx)
+	require.NoError(t, err)
+	assert.Equal(t, repo, root)
+
+	writeRepoFile(t, repo, " leading-space/train.py", "print()")
+	prefix, err = newGitRepo(filepath.Join(repo, " leading-space")).repoRelativePrefix(ctx)
+	require.NoError(t, err)
+	assert.Equal(t, " leading-space", prefix)
+}
+
+func TestGitRepo_RepoRelativePrefixFailure(t *testing.T) {
+	_, err := newGitRepo(t.TempDir()).repoRelativePrefix(t.Context())
+	require.Error(t, err)
+	assert.Contains(t, err.Error(), "failed to resolve repository-relative path")
+}
+
 func TestGitRepo_HeadSHA(t *testing.T) {
 	ctx := t.Context()
 	repo := newTestRepo(t)
@@ -196,6 +227,22 @@ func TestGitRepo_ValidateIncludePathsExist(t *testing.T) {
 	assert.Contains(t, err.Error(), sha[:8])
 }
 
+func TestGitRepo_ValidateSubtreeExists(t *testing.T) {
+	ctx := t.Context()
+	repo := newTestRepo(t)
+	writeRepoFile(t, repo, "subpkg/train.py", "print()")
+	sha := commitAll(t, repo, "init")
+	g := newGitRepo(filepath.Join(repo, "subpkg"))
+
+	require.NoError(t, g.validateSubtreeExists(ctx, sha, "subpkg"))
+	require.NoError(t, g.validateSubtreeExists(ctx, sha, ""))
+
+	err := g.validateSubtreeExists(ctx, sha, "missing")
+	require.Error(t, err)
+	assert.Contains(t, err.Error(), `root_path "missing" does not exist`)
+	assert.Contains(t, err.Error(), sha[:8])
+}
+
 func TestBuildGitStateSidecar_PlainTarClean(t *testing.T) {
 	ctx := t.Context()
 	repo := newTestRepo(t)

experimental/air/cmd/snapshot_package.go

@@ -17,18 +17,29 @@ import (
 // to /databricks/code_source/<dir> — so the --prefix / `-C parent dir` forms preserve it.
 
 // createGitArchiveSnapshot writes a gzipped tar of commitSHA to outputTarball via
-// `git archive`, with every entry prefixed by directoryName/. When includePaths is
-// set, only those paths are archived.
-func createGitArchiveSnapshot(ctx context.Context, git gitRepo, commitSHA, outputTarball, directoryName string, includePaths []string) error {
+// `git archive`, with every entry prefixed by directoryName/. When subtreePrefix is
+// set, only that repository subtree is archived and includePaths are relative to it.
+func createGitArchiveSnapshot(ctx context.Context, git gitRepo, commitSHA, outputTarball, directoryName string, includePaths []string, subtreePrefix string) error {
+	treeish := commitSHA
+	archiveGit := git
+	if subtreePrefix != "" {
+		repoRoot, err := git.repositoryRoot(ctx)
+		if err != nil {
+			return err
+		}
+		archiveGit = newGitRepo(repoRoot)
+		treeish = commitSHA + ":" + subtreePrefix
+	}
+
 	args := []string{
 		"archive",
 		"--format=tar.gz",
 		"--prefix=" + directoryName + "/",
 		"-o", outputTarball,
-		commitSHA,
+		treeish,
 	}
 	args = append(args, includePaths...)
-	if _, err := git.run(ctx, args...); err != nil {
+	if _, err := archiveGit.run(ctx, args...); err != nil {
 		return fmt.Errorf("failed to create git archive: %w", err)
 	}
 	return nil

experimental/air/cmd/snapshot_package_test.go

@@ -46,7 +46,7 @@ func TestCreateGitArchiveSnapshot(t *testing.T) {
 
 	out := filepath.Join(t.TempDir(), "snap.tar.gz")
 	dirName := filepath.Base(repo)
-	require.NoError(t, createGitArchiveSnapshot(ctx, newGitRepo(repo), sha, out, dirName, nil))
+	require.NoError(t, createGitArchiveSnapshot(ctx, newGitRepo(repo), sha, out, dirName, nil, ""))
 
 	entries := tarballEntries(t, out)
 	// Every real entry is prefixed with the directory name. git archive also emits a
@@ -70,13 +70,58 @@ func TestCreateGitArchiveSnapshot_IncludePaths(t *testing.T) {
 
 	out := filepath.Join(t.TempDir(), "snap.tar.gz")
 	dirName := filepath.Base(repo)
-	require.NoError(t, createGitArchiveSnapshot(ctx, newGitRepo(repo), sha, out, dirName, []string{"src"}))
+	require.NoError(t, createGitArchiveSnapshot(ctx, newGitRepo(repo), sha, out, dirName, []string{"src"}, ""))
 
 	entries := tarballEntries(t, out)
 	assert.Contains(t, entries, dirName+"/src/model.py")
 	assert.NotContains(t, entries, dirName+"/a.txt")
 }
 
+func TestCreateGitArchiveSnapshot_SubdirectoryRootPath(t *testing.T) {
+	ctx := t.Context()
+	repo := newTestRepo(t)
+	writeRepoFile(t, repo, "README.md", "repo root")
+	writeRepoFile(t, repo, "subpkg/train.py", "print()")
+	writeRepoFile(t, repo, "subpkg/nested/util.py", "pass")
+	sha := commitAll(t, repo, "init")
+
+	rootPath := filepath.Join(repo, "subpkg")
+	prefix, err := newGitRepo(rootPath).repoRelativePrefix(ctx)
+	require.NoError(t, err)
+
+	out := filepath.Join(t.TempDir(), "snap.tar.gz")
+	require.NoError(t, createGitArchiveSnapshot(ctx, newGitRepo(rootPath), sha, out, "subpkg", nil, prefix))
+
+	entries := tarballEntries(t, out)
+	assert.Contains(t, entries, "subpkg/train.py")
+	assert.Contains(t, entries, "subpkg/nested/util.py")
+	assert.NotContains(t, entries, "subpkg/README.md")
+	for _, entry := range entries {
+		assert.False(t, strings.HasPrefix(entry, "subpkg/subpkg/"), "entry %q is double nested", entry)
+	}
+}
+
+func TestCreateGitArchiveSnapshot_SubdirectoryRootPathWithIncludePaths(t *testing.T) {
+	ctx := t.Context()
+	repo := newTestRepo(t)
+	writeRepoFile(t, repo, "subpkg/train.py", "print()")
+	writeRepoFile(t, repo, "subpkg/src/model.py", "pass")
+	writeRepoFile(t, repo, "subpkg/configs/train.yaml", "x")
+	sha := commitAll(t, repo, "init")
+
+	rootPath := filepath.Join(repo, "subpkg")
+	prefix, err := newGitRepo(rootPath).repoRelativePrefix(ctx)
+	require.NoError(t, err)
+
+	out := filepath.Join(t.TempDir(), "snap.tar.gz")
+	require.NoError(t, createGitArchiveSnapshot(ctx, newGitRepo(rootPath), sha, out, "subpkg", []string{"src"}, prefix))
+
+	entries := tarballEntries(t, out)
+	assert.Contains(t, entries, "subpkg/src/model.py")
+	assert.NotContains(t, entries, "subpkg/train.py")
+	assert.NotContains(t, entries, "subpkg/configs/train.yaml")
+}
+
 func TestCreatePlainTarball(t *testing.T) {
 	ctx := t.Context()
 	repo := t.TempDir()