Add aicode mutator: package local-dir code_source_path at deploy

acceptance/bundle/ai_runtime_task/local_code_source/databricks.yml

@@ -0,0 +1,29 @@
+bundle:
+  name: ai-runtime-test
+
+sync:
+  # Exclude a file from both bundle sync and the code snapshot.
+  exclude:
+    - "**/*.log"
+
+resources:
+  jobs:
+    train:
+      name: "[${bundle.target}] AI Runtime training"
+      tasks:
+        - task_key: train
+          environment_key: default
+          ai_runtime_task:
+            experiment: my-training
+            code_source_path: ./src
+            deployments:
+              - command_path: src/command.sh
+                compute:
+                  accelerator_type: GPU_8xH100
+                  accelerator_count: 8
+      environments:
+        - environment_key: default
+          spec:
+            environment_version: "5"
+            dependencies:
+              - torch>=2.0.0

acceptance/bundle/ai_runtime_task/local_code_source/output.txt

@@ -0,0 +1,29 @@
+
+=== deploy packages the local code source
+
+>>> [CLI] bundle deploy
+Uploading bundle files to /Workspace/Users/[USERNAME]/.bundle/ai-runtime-test/default/files...
+Deploying resources...
+Updating deployment state...
+Deployment complete!
+
+=== the snapshot tarball is uploaded to the user's repo_snapshots directory
+/api/2.0/workspace-files/import-file/Workspace/Users/[USERNAME]/.air/repo_snapshots/src/src_[SNAPSHOT].tar.gz
+
+=== code_source_path points at the uploaded archive (no /Workspace prefix)
+/Users/[USERNAME]/.air/repo_snapshots/src/src_[SNAPSHOT].tar.gz
+
+=== command_path is translated to its absolute synced workspace path
+/Workspace/Users/[USERNAME]/.bundle/ai-runtime-test/default/files/src/command.sh
+
+=== a requirements.yaml is written next to command_path
+/api/2.0/workspace-files/import-file/Workspace/Users/[USERNAME]/.bundle/ai-runtime-test/default/files/src/requirements.yaml
+
+=== re-deploying unchanged code is a cache hit: the tarball is not re-uploaded
+
+>>> [CLI] bundle deploy
+Uploading bundle files to /Workspace/Users/[USERNAME]/.bundle/ai-runtime-test/default/files...
+Deploying resources...
+Updating deployment state...
+Deployment complete!
+tarball uploads on redeploy: 0

acceptance/bundle/ai_runtime_task/local_code_source/script

@@ -0,0 +1,30 @@
+# A local code_source_path is packaged into a content-addressed tarball
+# (<dir>_<hash>.tar.gz) and uploaded to the user's repo_snapshots directory. The
+# command_path is translated to its synced workspace path, and a requirements.yaml
+# derived from the job's serverless environment is written beside it.
+
+title "deploy packages the local code source\n"
+trace $CLI bundle deploy
+
+title "the snapshot tarball is uploaded to the user's repo_snapshots directory\n"
+# The upload appears twice: WSFS import-file 404s until the parent dir exists, then
+# the filer mkdirs and retries. Collapse to the distinct path.
+jq -r 'select(.path | test("/import-file/.*/.air/repo_snapshots/src/src_")) | .path' out.requests.txt | sort -u
+
+title "code_source_path points at the uploaded archive (no /Workspace prefix)\n"
+jq -r 'select(.path == "/api/2.2/jobs/create") | .body.tasks[0].ai_runtime_task.code_source_path' out.requests.txt
+
+title "command_path is translated to its absolute synced workspace path\n"
+jq -r 'select(.path == "/api/2.2/jobs/create") | .body.tasks[0].ai_runtime_task.deployments[0].command_path' out.requests.txt
+
+title "a requirements.yaml is written next to command_path\n"
+jq -r 'select(.path | test("/import-file/.*/files/src/requirements.yaml$")) | .path' out.requests.txt | sort -u
+
+rm out.requests.txt
+
+title "re-deploying unchanged code is a cache hit: the tarball is not re-uploaded\n"
+trace $CLI bundle deploy
+echo -n "tarball uploads on redeploy: "
+jq -r 'select(.path | test("/import-file/.*/.air/repo_snapshots/src/src_")) | .path' out.requests.txt | sort -u | wc -l | tr -d ' '
+
+rm out.requests.txt

acceptance/bundle/ai_runtime_task/local_code_source/src/.gitignore

@@ -0,0 +1 @@
+ignored_by_git.txt

acceptance/bundle/ai_runtime_task/local_code_source/src/command.sh

@@ -0,0 +1,2 @@
+cd $CODE_SOURCE_PATH
+python train.py

acceptance/bundle/ai_runtime_task/local_code_source/src/train.py

@@ -0,0 +1 @@
+print("training")

acceptance/bundle/ai_runtime_task/local_code_source/test.toml

@@ -0,0 +1,17 @@
+RecordRequests = true
+
+# Direct engine only: the terraform engine needs code_source_path in the generated
+# terraform provider schema, absent while the field is temporarily removed from the
+# jobs API/provider (returns next week). Restore ["terraform", "direct"] then.
+EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["direct"]
+
+Ignore = [
+    '.databricks',
+]
+
+# The archive is content-addressed: <dir>_<sha256[:16]>.tar.gz. The hash is stable
+# given the committed inputs, but collapse it to a token so the test does not pin a
+# specific digest.
+[[Repls]]
+Old = 'src_[0-9a-f]{16}\.tar\.gz'
+New = 'src_[SNAPSHOT].tar.gz'

bundle/config/mutator/aicode/package_upload.go

@@ -0,0 +1,214 @@
+// Package aicode packages a local code_source_path directory into a
+// content-addressed tarball at deploy, uploads it to the user's workspace, and
+// rewrites code_source_path to the remote archive. Already-remote values are left
+// untouched. The archive name embeds its SHA-256, so unchanged code skips re-upload.
+package aicode
+
+import (
+	"bytes"
+	"context"
+	"errors"
+	"fmt"
+	"io/fs"
+	"path"
+	"path/filepath"
+	"strings"
+
+	"github.com/databricks/cli/bundle"
+	"github.com/databricks/cli/bundle/deploy/files"
+	"github.com/databricks/cli/bundle/libraries"
+	"github.com/databricks/cli/libs/diag"
+	"github.com/databricks/cli/libs/dyn"
+	"github.com/databricks/cli/libs/filer"
+	"github.com/databricks/cli/libs/fileset"
+	"github.com/databricks/cli/libs/log"
+	libsync "github.com/databricks/cli/libs/sync"
+)
+
+// codeSourcePatterns locate code_source_path on a direct task and under a for_each_task.
+var codeSourcePatterns = []dyn.Pattern{
+	dyn.NewPattern(
+		dyn.Key("resources"), dyn.Key("jobs"), dyn.AnyKey(),
+		dyn.Key("tasks"), dyn.AnyIndex(),
+		dyn.Key("ai_runtime_task"), dyn.Key("code_source_path"),
+	),
+	dyn.NewPattern(
+		dyn.Key("resources"), dyn.Key("jobs"), dyn.AnyKey(),
+		dyn.Key("tasks"), dyn.AnyIndex(),
+		dyn.Key("for_each_task"), dyn.Key("task"),
+		dyn.Key("ai_runtime_task"), dyn.Key("code_source_path"),
+	),
+}
+
+// codeSource is a single local code_source_path occurrence to package.
+type codeSource struct {
+	configPath dyn.Path
+	location   dyn.Location
+	value      string // raw code_source_path string as written in config
+}
+
+func PackageAndUpload() bundle.Mutator {
+	return &packageAndUpload{}
+}
+
+type packageAndUpload struct {
+	client filer.Filer // nil in normal use (a filer is built per code source); set only in tests
+}
+
+func (m *packageAndUpload) Name() string {
+	return "aicode.PackageAndUpload"
+}
+
+func (m *packageAndUpload) Apply(ctx context.Context, b *bundle.Bundle) diag.Diagnostics {
+	sources, diags := collectLocalCodeSources(b.Config.Value())
+	if diags.HasError() {
+		return diags
+	}
+	if len(sources) == 0 {
+		return diags
+	}
+
+	userDir, err := userWorkspaceHome(b)
+	if err != nil {
+		return diags.Extend(diag.FromErr(err))
+	}
+
+	// Upload all sources before rewriting any config, so an upload failure leaves
+	// the config untouched.
+	remotePaths := make(map[string]string, len(sources))
+	for _, cs := range sources {
+		remote, err := m.packageOne(ctx, b, cs, userDir)
+		if err != nil {
+			diags = diags.Extend(diag.FromErr(err))
+			return diags
+		}
+		remotePaths[cs.configPath.String()] = remote
+	}
+
+	err = b.Config.Mutate(func(root dyn.Value) (dyn.Value, error) {
+		for _, cs := range sources {
+			remote := remotePaths[cs.configPath.String()]
+			root, err = dyn.SetByPath(root, cs.configPath, dyn.NewValue(remote, []dyn.Location{cs.location}))
+			if err != nil {
+				return root, fmt.Errorf("failed to update code_source_path %q to %q: %w", cs.value, remote, err)
+			}
+		}
+		return root, nil
+	})
+	if err != nil {
+		diags = diags.Extend(diag.FromErr(err))
+	}
+
+	return diags
+}
+
+// repoSnapshotsSubdir holds code snapshots under the user's home. Deliberately not
+// <artifact_path>/.internal, which artifacts.CleanUp() wipes each deploy (matches the Python air CLI).
+const repoSnapshotsSubdir = ".air/repo_snapshots"
+
+// packageOne tarballs one code source, uploads it to repo_snapshots (skipping the
+// upload when a content-identical archive is already there), and returns its remote path.
+func (m *packageAndUpload) packageOne(ctx context.Context, b *bundle.Bundle, cs codeSource, userDir string) (string, error) {
+	localDir := filepath.Join(b.SyncRootPath, filepath.FromSlash(cs.value))
+	dirName := filepath.Base(localDir)
+
+	// relBase scopes the sync file list to the code dir and re-bases archive entries under it.
+	relBase, err := filepath.Rel(b.SyncRootPath, localDir)
+	if err != nil {
+		return "", fmt.Errorf("code_source_path %q: %w", cs.value, err)
+	}
+	relBase = filepath.ToSlash(relBase)
+
+	files, err := codeSourceFiles(ctx, b, relBase)
+	if err != nil {
+		return "", fmt.Errorf("failed to list files for code_source_path %q: %w", cs.value, err)
+	}
+
+	uploadPath := path.Join(userDir, repoSnapshotsSubdir, dirName)
+	client := m.client
+	if client == nil {
+		client, err = filer.NewWorkspaceFilesClient(b.WorkspaceClient(ctx), uploadPath)
+		if err != nil {
+			return "", err
+		}
+	}
+
+	// Build in memory so the content hash (computed while gzipping) can name the upload.
+	var buf bytes.Buffer
+	sha, err := buildCodeSnapshot(b.SyncRoot, relBase, files, dirName, &buf)
+	if err != nil {
+		return "", fmt.Errorf("failed to package code_source_path %q: %w", cs.value, err)
+	}
+
+	archiveName := fmt.Sprintf("%s_%s.tar.gz", dirName, sha[:16])
+	// The AI Runtime fetcher wants the legacy "/Users/..." form (no "/Workspace"), while
+	// the filer uploads via uploadPath; record the de-prefixed path on the task.
+	remotePath := strings.TrimPrefix(path.Join(uploadPath, archiveName), "/Workspace")
+
+	// A matching name means identical content is already uploaded (the archive is reproducible).
+	if _, err := client.Stat(ctx, archiveName); err == nil {
+		log.Debugf(ctx, "code snapshot already present at %s, skipping upload", remotePath)
+		return remotePath, nil
+	} else if !errors.Is(err, fs.ErrNotExist) {
+		return "", fmt.Errorf("failed to check for existing code snapshot %q: %w", remotePath, err)
+	}
+
+	if err := client.Write(ctx, archiveName, &buf, filer.OverwriteIfExists, filer.CreateParentDirectories); err != nil {
+		return "", fmt.Errorf("failed to upload code snapshot %q: %w", remotePath, err)
+	}
+	return remotePath, nil
+}
+
+// codeSourceFiles lists the files under relBase (relative to the sync root) for the
+// snapshot, filtered like bundle file sync: .gitignore-aware plus sync.include/exclude.
+func codeSourceFiles(ctx context.Context, b *bundle.Bundle, relBase string) ([]fileset.File, error) {
+	opts, err := files.GetSyncOptions(ctx, b)
+	if err != nil {
+		return nil, err
+	}
+	fl, err := libsync.NewFileList(ctx, opts.WorktreeRoot, opts.LocalRoot, []string{relBase}, opts.Include, opts.Exclude)
+	if err != nil {
+		return nil, err
+	}
+	return fl.Files(ctx)
+}
+
+// userWorkspaceHome returns /Workspace/Users/<user>, the root for code snapshots.
+func userWorkspaceHome(b *bundle.Bundle) (string, error) {
+	u := b.Config.Workspace.CurrentUser
+	if u == nil || u.User == nil || u.UserName == "" {
+		return "", errors.New("unable to resolve code snapshot location: current user not set")
+	}
+	return "/Workspace/Users/" + u.UserName, nil
+}
+
+// collectLocalCodeSources returns every AI Runtime task code_source_path in root
+// that points at a local directory. Already-remote values are skipped. It reads the
+// dyn value directly (read-only), so it does not depend on the typed SDK field.
+func collectLocalCodeSources(root dyn.Value) ([]codeSource, diag.Diagnostics) {
+	var sources []codeSource
+	var diags diag.Diagnostics
+
+	for _, pattern := range codeSourcePatterns {
+		_, err := dyn.MapByPattern(root, pattern, func(p dyn.Path, v dyn.Value) (dyn.Value, error) {
+			value, ok := v.AsString()
+			if !ok {
+				return v, fmt.Errorf("expected string, got %s", v.Kind())
+			}
+			if !libraries.IsLocalPath(value) {
+				return v, nil
+			}
+			sources = append(sources, codeSource{
+				configPath: p,
+				location:   v.Location(),
+				value:      value,
+			})
+			return v, nil
+		})
+		if err != nil {
+			diags = diags.Extend(diag.FromErr(err))
+		}
+	}
+
+	return sources, diags
+}

bundle/config/mutator/aicode/package_upload_test.go

@@ -0,0 +1,55 @@
+package aicode
+
+import (
+	"strings"
+	"testing"
+
+	"github.com/databricks/cli/libs/dyn"
+	"github.com/databricks/cli/libs/dyn/yamlloader"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+// configWithCodeSource parses a minimal bundle config whose single AI Runtime task
+// points at codeSourcePath, as a dyn value. It reads YAML directly so the config
+// retains code_source_path regardless of whether the SDK's typed struct has it.
+//
+// The end-to-end package/upload/rewrite behavior (local dir -> tarball -> upload ->
+// rewritten code_source_path) runs the full mutator pipeline (sync file list,
+// workspace filer) and is covered by acceptance tests under
+// acceptance/bundle/ai_runtime_task. This unit test covers only the pure
+// config-collection seam that does not touch the pipeline.
+func configWithCodeSource(t *testing.T, codeSourcePath string) dyn.Value {
+	t.Helper()
+	yml := `
+resources:
+  jobs:
+    train:
+      tasks:
+        - task_key: train
+          ai_runtime_task:
+            experiment: exp
+            code_source_path: ` + codeSourcePath + `
+`
+	v, err := yamlloader.LoadYAML("test.yml", strings.NewReader(yml))
+	require.NoError(t, err)
+	return v
+}
+
+func TestCollectLocalCodeSourcesFindsLocalPath(t *testing.T) {
+	sources, diags := collectLocalCodeSources(configWithCodeSource(t, "./src"))
+	require.Empty(t, diags)
+	require.Len(t, sources, 1)
+	assert.Equal(t, "./src", sources[0].value)
+}
+
+func TestCollectLocalCodeSourcesSkipsRemotePaths(t *testing.T) {
+	for _, remote := range []string{
+		"/Workspace/Users/me/code.tar.gz",
+		"/Volumes/main/default/code/existing.tar.gz",
+	} {
+		sources, diags := collectLocalCodeSources(configWithCodeSource(t, remote))
+		require.Empty(t, diags)
+		assert.Empty(t, sources, "remote code_source_path %q must not be collected", remote)
+	}
+}