From 4ecc40e0336913176970aaf90614189e33562880 Mon Sep 17 00:00:00 2001 From: "ben.hansen" Date: Mon, 13 Jul 2026 04:28:35 +0000 Subject: [PATCH 1/5] test air cli dabs --- .github/OWNERS | 4 + NEXT_CHANGELOG.md | 4 + .../local_code_source/databricks.yml | 24 +++ .../local_code_source/out.test.toml | 3 + .../local_code_source/output.txt | 14 ++ .../ai_runtime_task/local_code_source/script | 9 + .../local_code_source/src/command.sh | 2 + .../local_code_source/src/train.py | 1 + .../local_code_source/test.toml | 11 + acceptance/experimental/air/.gitkeep | 0 .../experimental/air/init/out.test.toml | 3 + acceptance/experimental/air/init/output.txt | 43 ++++ acceptance/experimental/air/init/script | 11 + acceptance/experimental/air/init/test.toml | 8 + .../air/init/train_project/src/train.py | 1 + .../air/init/train_project/train.yaml | 18 ++ .../air/init_errors/docker/train.yaml | 12 ++ .../air/init_errors/legacy/train.yaml | 9 + .../air/init_errors/out.test.toml | 3 + .../experimental/air/init_errors/output.txt | 8 + .../experimental/air/init_errors/script | 5 + .../experimental/air/init_errors/test.toml | 3 + .../config/mutator/aicode/package_upload.go | 194 ++++++++++++++++++ .../mutator/aicode/package_upload_test.go | 137 +++++++++++++ bundle/config/mutator/aicode/tarball.go | 122 +++++++++++ bundle/config/mutator/aicode/tarball_test.go | 108 ++++++++++ bundle/config/mutator/aicode/validate.go | 101 +++++++++ bundle/config/mutator/aicode/validate_test.go | 64 ++++++ bundle/phases/deploy.go | 7 + bundle/phases/initialize.go | 7 + cmd/experimental/experimental.go | 2 + experimental/air/.gitkeep | 0 experimental/air/cmd/air.go | 17 ++ experimental/air/cmd/init.go | 145 +++++++++++++ experimental/air/trainyaml/convert.go | 168 +++++++++++++++ experimental/air/trainyaml/convert_test.go | 134 ++++++++++++ experimental/air/trainyaml/parse.go | 110 ++++++++++ 37 files changed, 1512 insertions(+) create mode 100644 acceptance/bundle/ai_runtime_task/local_code_source/databricks.yml create mode 100644 acceptance/bundle/ai_runtime_task/local_code_source/out.test.toml create mode 100644 acceptance/bundle/ai_runtime_task/local_code_source/output.txt create mode 100644 acceptance/bundle/ai_runtime_task/local_code_source/script create mode 100644 acceptance/bundle/ai_runtime_task/local_code_source/src/command.sh create mode 100644 acceptance/bundle/ai_runtime_task/local_code_source/src/train.py create mode 100644 acceptance/bundle/ai_runtime_task/local_code_source/test.toml delete mode 100644 acceptance/experimental/air/.gitkeep create mode 100644 acceptance/experimental/air/init/out.test.toml create mode 100644 acceptance/experimental/air/init/output.txt create mode 100644 acceptance/experimental/air/init/script create mode 100644 acceptance/experimental/air/init/test.toml create mode 100644 acceptance/experimental/air/init/train_project/src/train.py create mode 100644 acceptance/experimental/air/init/train_project/train.yaml create mode 100644 acceptance/experimental/air/init_errors/docker/train.yaml create mode 100644 acceptance/experimental/air/init_errors/legacy/train.yaml create mode 100644 acceptance/experimental/air/init_errors/out.test.toml create mode 100644 acceptance/experimental/air/init_errors/output.txt create mode 100644 acceptance/experimental/air/init_errors/script create mode 100644 acceptance/experimental/air/init_errors/test.toml create mode 100644 bundle/config/mutator/aicode/package_upload.go create mode 100644 bundle/config/mutator/aicode/package_upload_test.go create mode 100644 bundle/config/mutator/aicode/tarball.go create mode 100644 bundle/config/mutator/aicode/tarball_test.go create mode 100644 bundle/config/mutator/aicode/validate.go create mode 100644 bundle/config/mutator/aicode/validate_test.go delete mode 100644 experimental/air/.gitkeep create mode 100644 experimental/air/cmd/air.go create mode 100644 experimental/air/cmd/init.go create mode 100644 experimental/air/trainyaml/convert.go create mode 100644 experimental/air/trainyaml/convert_test.go create mode 100644 experimental/air/trainyaml/parse.go diff --git a/.github/OWNERS b/.github/OWNERS index 197b96f0980..669fd9f4962 100644 --- a/.github/OWNERS +++ b/.github/OWNERS @@ -73,4 +73,8 @@ # Experimental /experimental/air/ team:ai-training /acceptance/experimental/air/ team:ai-training +# AI Runtime code packaging mutator (overloads ai_runtime_task.code_source_path +# at deploy). Lives in core bundle, co-owned by ai-training who owns the feature. +/bundle/config/mutator/aicode/ team:bundle team:ai-training @lennartkats-db +/acceptance/bundle/ai_runtime_task/ team:bundle team:ai-training @lennartkats-db /experimental/aitools/ team:eng-apps-devex @lennartkats-db diff --git a/NEXT_CHANGELOG.md b/NEXT_CHANGELOG.md index 0f71834ff56..7c170b89fa4 100644 --- a/NEXT_CHANGELOG.md +++ b/NEXT_CHANGELOG.md @@ -6,8 +6,12 @@ ### CLI +* Added `databricks experimental air init`, which scaffolds a bundle from an AI Runtime (AIR) CLI `train.yaml`, generating a job with an `ai_runtime_task` so an ephemeral AIR run can be migrated to a durable, DABs-managed job. + ### Bundles +* An `ai_runtime_task.code_source_path` that points at a local directory is now packaged into a tarball during `bundle deploy`, uploaded to the workspace (or the configured UC volume), and rewritten to the uploaded path, mirroring the AIR CLI's `code_source` behavior. + ### Dependency updates ### API Changes diff --git a/acceptance/bundle/ai_runtime_task/local_code_source/databricks.yml b/acceptance/bundle/ai_runtime_task/local_code_source/databricks.yml new file mode 100644 index 00000000000..c6d37475f73 --- /dev/null +++ b/acceptance/bundle/ai_runtime_task/local_code_source/databricks.yml @@ -0,0 +1,24 @@ +bundle: + name: ai-runtime-test + +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 diff --git a/acceptance/bundle/ai_runtime_task/local_code_source/out.test.toml b/acceptance/bundle/ai_runtime_task/local_code_source/out.test.toml new file mode 100644 index 00000000000..f784a183258 --- /dev/null +++ b/acceptance/bundle/ai_runtime_task/local_code_source/out.test.toml @@ -0,0 +1,3 @@ +Local = true +Cloud = false +EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["terraform", "direct"] diff --git a/acceptance/bundle/ai_runtime_task/local_code_source/output.txt b/acceptance/bundle/ai_runtime_task/local_code_source/output.txt new file mode 100644 index 00000000000..be8c3ec21ab --- /dev/null +++ b/acceptance/bundle/ai_runtime_task/local_code_source/output.txt @@ -0,0 +1,14 @@ + +>>> [CLI] bundle deploy +Uploading bundle files to /Workspace/Users/[USERNAME]/.bundle/ai-runtime-test/default/files... +Deploying resources... +Updating deployment state... +Deployment complete! + +=== Expecting the code snapshot tarball to be uploaded to artifact_path/.internal +>>> jq .path +"/api/2.0/workspace-files/import-file/Workspace/Users/[USERNAME]/.bundle/ai-runtime-test/default/artifacts/.internal/src_[HASH].tar.gz" + +=== Expecting the created job's code_source_path to point at the uploaded archive +>>> jq -s .[] | select(.path=="/api/2.2/jobs/create") | .body.tasks[0].ai_runtime_task.code_source_path out.requests.txt +"/Workspace/Users/[USERNAME]/.bundle/ai-runtime-test/default/artifacts/.internal/src_[HASH].tar.gz" diff --git a/acceptance/bundle/ai_runtime_task/local_code_source/script b/acceptance/bundle/ai_runtime_task/local_code_source/script new file mode 100644 index 00000000000..8701b74e6c9 --- /dev/null +++ b/acceptance/bundle/ai_runtime_task/local_code_source/script @@ -0,0 +1,9 @@ +trace $CLI bundle deploy + +title "Expecting the code snapshot tarball to be uploaded to artifact_path/.internal" +trace jq .path < out.requests.txt | grep import-file | grep 'tar.gz' | sort + +title "Expecting the created job's code_source_path to point at the uploaded archive" +trace jq -s '.[] | select(.path=="/api/2.2/jobs/create") | .body.tasks[0].ai_runtime_task.code_source_path' out.requests.txt + +rm out.requests.txt diff --git a/acceptance/bundle/ai_runtime_task/local_code_source/src/command.sh b/acceptance/bundle/ai_runtime_task/local_code_source/src/command.sh new file mode 100644 index 00000000000..7ce3949767c --- /dev/null +++ b/acceptance/bundle/ai_runtime_task/local_code_source/src/command.sh @@ -0,0 +1,2 @@ +cd $CODE_SOURCE_PATH +python train.py diff --git a/acceptance/bundle/ai_runtime_task/local_code_source/src/train.py b/acceptance/bundle/ai_runtime_task/local_code_source/src/train.py new file mode 100644 index 00000000000..d8b062e5dfe --- /dev/null +++ b/acceptance/bundle/ai_runtime_task/local_code_source/src/train.py @@ -0,0 +1 @@ +print("training") diff --git a/acceptance/bundle/ai_runtime_task/local_code_source/test.toml b/acceptance/bundle/ai_runtime_task/local_code_source/test.toml new file mode 100644 index 00000000000..1c37cecae7d --- /dev/null +++ b/acceptance/bundle/ai_runtime_task/local_code_source/test.toml @@ -0,0 +1,11 @@ +RecordRequests = true + +Ignore = [ + '.databricks', +] + +# The code archive filename embeds a content hash. Normalize it so output is +# stable across machines while still asserting the _ shape. +[[Repls]] +Old = 'src_[0-9a-f]{16}\.tar\.gz' +New = 'src_[HASH].tar.gz' diff --git a/acceptance/experimental/air/.gitkeep b/acceptance/experimental/air/.gitkeep deleted file mode 100644 index e69de29bb2d..00000000000 diff --git a/acceptance/experimental/air/init/out.test.toml b/acceptance/experimental/air/init/out.test.toml new file mode 100644 index 00000000000..d6187dcb046 --- /dev/null +++ b/acceptance/experimental/air/init/out.test.toml @@ -0,0 +1,3 @@ +Local = true +Cloud = false +EnvMatrix.DATABRICKS_BUNDLE_ENGINE = [] diff --git a/acceptance/experimental/air/init/output.txt b/acceptance/experimental/air/init/output.txt new file mode 100644 index 00000000000..f08b4a289a3 --- /dev/null +++ b/acceptance/experimental/air/init/output.txt @@ -0,0 +1,43 @@ + +=== migrate a train.yaml into a bundle +>>> [CLI] experimental air init --from train_project/train.yaml --output-dir generated +Warning: "secrets" has no ai_runtime_task equivalent and was dropped; configure secrets on the job or via {{secrets/scope/key}} references +Bundle written to generated +Deploy it with: databricks bundle deploy + +=== generated databricks.yml +>>> cat generated/databricks.yml +bundle: + name: my_training +include: + - resources/*.yml + +=== generated job resource +>>> cat generated/resources/my_training.job.yml +resources: + jobs: + my_training: + name: my-training + tasks: + - task_key: my-training + ai_runtime_task: + code_source_path: ./src + deployments: + - command_path: src/command.sh + compute: + accelerator_count: 8 + accelerator_type: GPU_8xH100 + experiment: my-training + environment_key: default + max_retries: 2 + environments: + - environment_key: default + spec: + dependencies: + - torch>=2.0.0 + environment_version: "5" + +=== command materialized next to the code source +>>> cat generated/src/command.sh +cd $CODE_SOURCE_PATH +python train.py diff --git a/acceptance/experimental/air/init/script b/acceptance/experimental/air/init/script new file mode 100644 index 00000000000..56c24aca87f --- /dev/null +++ b/acceptance/experimental/air/init/script @@ -0,0 +1,11 @@ +title "migrate a train.yaml into a bundle" +trace $CLI experimental air init --from train_project/train.yaml --output-dir generated + +title "generated databricks.yml" +trace cat generated/databricks.yml + +title "generated job resource" +trace cat generated/resources/my_training.job.yml + +title "command materialized next to the code source" +trace cat generated/src/command.sh diff --git a/acceptance/experimental/air/init/test.toml b/acceptance/experimental/air/init/test.toml new file mode 100644 index 00000000000..6270a5aca70 --- /dev/null +++ b/acceptance/experimental/air/init/test.toml @@ -0,0 +1,8 @@ +# This command generates local files from a train.yaml; no bundle engine or +# workspace is needed. +Ignore = [ + 'generated', +] + +[EnvMatrix] +DATABRICKS_BUNDLE_ENGINE = [] diff --git a/acceptance/experimental/air/init/train_project/src/train.py b/acceptance/experimental/air/init/train_project/src/train.py new file mode 100644 index 00000000000..d8b062e5dfe --- /dev/null +++ b/acceptance/experimental/air/init/train_project/src/train.py @@ -0,0 +1 @@ +print("training") diff --git a/acceptance/experimental/air/init/train_project/train.yaml b/acceptance/experimental/air/init/train_project/train.yaml new file mode 100644 index 00000000000..3e3ca0f5f7c --- /dev/null +++ b/acceptance/experimental/air/init/train_project/train.yaml @@ -0,0 +1,18 @@ +experiment_name: my-training +compute: + num_accelerators: 8 + accelerator_type: GPU_8xH100 +environment: + version: "5" + dependencies: + - torch>=2.0.0 +code_source: + type: snapshot + snapshot: + root_path: ./src +command: | + cd $CODE_SOURCE_PATH + python train.py +max_retries: 2 +secrets: + HF_TOKEN: my_scope/hf_token diff --git a/acceptance/experimental/air/init_errors/docker/train.yaml b/acceptance/experimental/air/init_errors/docker/train.yaml new file mode 100644 index 00000000000..df1e22c2056 --- /dev/null +++ b/acceptance/experimental/air/init_errors/docker/train.yaml @@ -0,0 +1,12 @@ +experiment_name: docker-run +compute: + num_accelerators: 1 + accelerator_type: GPU_1xA10 +environment: + docker_image: + url: org/repo:tag +code_source: + type: snapshot + snapshot: + root_path: ./src +command: python train.py diff --git a/acceptance/experimental/air/init_errors/legacy/train.yaml b/acceptance/experimental/air/init_errors/legacy/train.yaml new file mode 100644 index 00000000000..f4d726a4476 --- /dev/null +++ b/acceptance/experimental/air/init_errors/legacy/train.yaml @@ -0,0 +1,9 @@ +experiment_name: legacy-run +compute: + num_accelerators: 1 + accelerator_type: a10 +code_source: + type: snapshot + snapshot: + root_path: ./src +command: python train.py diff --git a/acceptance/experimental/air/init_errors/out.test.toml b/acceptance/experimental/air/init_errors/out.test.toml new file mode 100644 index 00000000000..d6187dcb046 --- /dev/null +++ b/acceptance/experimental/air/init_errors/out.test.toml @@ -0,0 +1,3 @@ +Local = true +Cloud = false +EnvMatrix.DATABRICKS_BUNDLE_ENGINE = [] diff --git a/acceptance/experimental/air/init_errors/output.txt b/acceptance/experimental/air/init_errors/output.txt new file mode 100644 index 00000000000..9a400dc7563 --- /dev/null +++ b/acceptance/experimental/air/init_errors/output.txt @@ -0,0 +1,8 @@ + +=== legacy accelerator type is rejected +>>> [CLI] experimental air init --from legacy/train.yaml --output-dir legacy-out +Error: compute.accelerator_type "a10" is not supported by ai_runtime_task (expected GPU_1xA10, GPU_1xH100, or GPU_8xH100) + +=== docker image is rejected +>>> [CLI] experimental air init --from docker/train.yaml --output-dir docker-out +Error: environment.docker_image is not supported by ai_runtime_task; use environment.version and environment.dependencies instead diff --git a/acceptance/experimental/air/init_errors/script b/acceptance/experimental/air/init_errors/script new file mode 100644 index 00000000000..43fc02501bf --- /dev/null +++ b/acceptance/experimental/air/init_errors/script @@ -0,0 +1,5 @@ +title "legacy accelerator type is rejected" +musterr trace $CLI experimental air init --from legacy/train.yaml --output-dir legacy-out + +title "docker image is rejected" +musterr trace $CLI experimental air init --from docker/train.yaml --output-dir docker-out diff --git a/acceptance/experimental/air/init_errors/test.toml b/acceptance/experimental/air/init_errors/test.toml new file mode 100644 index 00000000000..afca356e837 --- /dev/null +++ b/acceptance/experimental/air/init_errors/test.toml @@ -0,0 +1,3 @@ +# These invocations must fail; no bundle engine or workspace is needed. +[EnvMatrix] +DATABRICKS_BUNDLE_ENGINE = [] diff --git a/bundle/config/mutator/aicode/package_upload.go b/bundle/config/mutator/aicode/package_upload.go new file mode 100644 index 00000000000..234514763c1 --- /dev/null +++ b/bundle/config/mutator/aicode/package_upload.go @@ -0,0 +1,194 @@ +// Package aicode packages a local code directory referenced by an AI Runtime +// task's code_source_path and uploads it to the workspace during deploy. +// +// The SDK jobs.AiRuntimeTask.code_source_path field expects a workspace or UC +// volume path to an uploaded code archive; its doc comment states that the CLI +// is responsible for packaging the user's local code directory into that +// archive. This mutator implements that contract for DABs: when a user points +// code_source_path at a local directory, it tarballs the directory (gitignore +// aware, .git excluded), uploads the archive next to bundle libraries, and +// rewrites the field to the resulting remote path so the deployed job runs +// against the uploaded code. Values that are already remote are left untouched. +package aicode + +import ( + "bytes" + "context" + "errors" + "fmt" + "io/fs" + "os" + "path" + "path/filepath" + + "github.com/databricks/cli/bundle" + "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/log" + "github.com/databricks/cli/libs/vfs" +) + +// codeSourcePatterns are the config locations of an AI Runtime task's +// code_source_path, both as a direct task and nested 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 is the raw code_source_path string as written in config. + value string +} + +func PackageAndUpload() bundle.Mutator { + return &packageAndUpload{} +} + +type packageAndUpload struct { + // client is the filer used for uploads. It defaults to the libraries filer + // and is only overridden in tests. + client filer.Filer +} + +func (m *packageAndUpload) Name() string { + return "aicode.PackageAndUpload" +} + +func (m *packageAndUpload) Apply(ctx context.Context, b *bundle.Bundle) diag.Diagnostics { + sources, diags := collectLocalCodeSources(b) + if diags.HasError() { + return diags + } + if len(sources) == 0 { + return diags + } + + client, uploadPath, filerDiags := libraries.GetFilerForLibraries(ctx, b) + diags = diags.Extend(filerDiags) + if diags.HasError() { + return diags + } + if m.client == nil { + m.client = client + } + + stagingDir, err := b.LocalStateDir(ctx, "ai_code_source") + if err != nil { + return diags.Extend(diag.FromErr(err)) + } + + // remotePaths maps each config location to the remote archive path it should + // point to after upload. Built outside the Mutate closure so upload failures + // are reported before any config is rewritten. + remotePaths := make(map[string]string, len(sources)) + for _, cs := range sources { + remote, err := m.packageOne(ctx, b, cs, stagingDir, uploadPath) + 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 +} + +// packageOne tarballs the local directory for a single code source, uploads it +// (skipping the upload if a content-identical archive already exists), and +// returns the remote path the config should point to. +func (m *packageAndUpload) packageOne(ctx context.Context, b *bundle.Bundle, cs codeSource, stagingDir, uploadPath string) (string, error) { + localDir := filepath.Join(b.SyncRootPath, filepath.FromSlash(cs.value)) + dirName := filepath.Base(localDir) + + var buf bytes.Buffer + sha, err := buildTarball(ctx, vfs.MustNew(localDir), dirName, &buf) + if err != nil { + return "", fmt.Errorf("failed to package code_source_path %q: %w", cs.value, err) + } + + // The SHA of the (reproducible) archive is embedded in the filename so an + // unchanged code directory resolves to the same remote path across deploys + // and the existence check below turns re-uploads into no-ops. + archiveName := fmt.Sprintf("%s_%s.tar.gz", dirName, sha[:16]) + remotePath := path.Join(uploadPath, archiveName) + + if _, err := m.client.Stat(ctx, archiveName); err == nil { + log.Debugf(ctx, "code snapshot %s already present, 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) + } + + // Stage the archive on disk so a large snapshot is not held in memory during + // upload, matching how libraries are uploaded from a file. + localArchive := filepath.Join(stagingDir, archiveName) + if err := os.WriteFile(localArchive, buf.Bytes(), 0o600); err != nil { + return "", err + } + + if err := libraries.UploadFile(ctx, localArchive, m.client); err != nil { + return "", err + } + + return remotePath, nil +} + +// collectLocalCodeSources returns every AI Runtime task code_source_path that +// points at a local directory. Already-remote values are skipped. +func collectLocalCodeSources(b *bundle.Bundle) ([]codeSource, diag.Diagnostics) { + var sources []codeSource + var diags diag.Diagnostics + + for _, pattern := range codeSourcePatterns { + err := b.Config.Mutate(func(root dyn.Value) (dyn.Value, error) { + return 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 +} diff --git a/bundle/config/mutator/aicode/package_upload_test.go b/bundle/config/mutator/aicode/package_upload_test.go new file mode 100644 index 00000000000..a61b01f2264 --- /dev/null +++ b/bundle/config/mutator/aicode/package_upload_test.go @@ -0,0 +1,137 @@ +package aicode + +import ( + "bytes" + "context" + "io" + "io/fs" + "os" + "path/filepath" + "testing" + + "github.com/databricks/cli/bundle" + "github.com/databricks/cli/bundle/config" + "github.com/databricks/cli/bundle/config/resources" + "github.com/databricks/cli/bundle/internal/bundletest" + "github.com/databricks/cli/libs/dyn" + "github.com/databricks/cli/libs/filer" + "github.com/databricks/cli/libs/vfs" + "github.com/databricks/databricks-sdk-go/service/jobs" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// recordingFiler records every Write and reports files as absent, so an upload +// always happens and its target name is captured. +type recordingFiler struct { + filer.Filer + written map[string][]byte + exists map[string]bool +} + +func newRecordingFiler() *recordingFiler { + return &recordingFiler{written: map[string][]byte{}, exists: map[string]bool{}} +} + +func (f *recordingFiler) Write(ctx context.Context, p string, r io.Reader, _ ...filer.WriteMode) error { + b, err := io.ReadAll(r) + if err != nil { + return err + } + f.written[p] = b + return nil +} + +func (f *recordingFiler) Stat(ctx context.Context, p string) (fs.FileInfo, error) { + if f.exists[p] { + return fakeFileInfo{}, nil + } + return nil, fs.ErrNotExist +} + +type fakeFileInfo struct{ fs.FileInfo } + +func bundleWithCodeSource(t *testing.T, codeSourcePath string) (*bundle.Bundle, string) { + t.Helper() + dir := t.TempDir() + require.NoError(t, os.MkdirAll(filepath.Join(dir, "src"), 0o755)) + require.NoError(t, os.WriteFile(filepath.Join(dir, "src", "train.py"), []byte("print('x')"), 0o644)) + + b := &bundle.Bundle{ + BundleRootPath: dir, + SyncRootPath: dir, + Config: config.Root{ + Bundle: config.Bundle{Target: "default"}, + Workspace: config.Workspace{ + ArtifactPath: "/Workspace/Users/me/.bundle/artifacts", + }, + Resources: config.Resources{ + Jobs: map[string]*resources.Job{ + "train": { + JobSettings: jobs.JobSettings{ + Tasks: []jobs.Task{ + { + TaskKey: "train", + AiRuntimeTask: &jobs.AiRuntimeTask{ + Experiment: "exp", + CodeSourcePath: codeSourcePath, + }, + }, + }, + }, + }, + }, + }, + }, + } + bundletest.SetLocation(b, ".", []dyn.Location{{File: filepath.Join(dir, "databricks.yml")}}) + return b, dir +} + +func TestPackageAndUploadRewritesLocalCodeSource(t *testing.T) { + b, _ := bundleWithCodeSource(t, "src") + f := newRecordingFiler() + + diags := bundle.Apply(t.Context(), b, &packageAndUpload{client: f}) + require.Empty(t, diags) + + rewritten := b.Config.Resources.Jobs["train"].Tasks[0].AiRuntimeTask.CodeSourcePath + assert.Regexp(t, `^/Workspace/Users/me/\.bundle/artifacts/\.internal/src_[0-9a-f]{16}\.tar\.gz$`, rewritten) + + require.Len(t, f.written, 1) + for name := range f.written { + assert.Regexp(t, `^src_[0-9a-f]{16}\.tar\.gz$`, name) + } +} + +func TestPackageAndUploadSkipsRemoteCodeSource(t *testing.T) { + b, _ := bundleWithCodeSource(t, "/Volumes/main/default/code/existing.tar.gz") + f := newRecordingFiler() + + diags := bundle.Apply(t.Context(), b, &packageAndUpload{client: f}) + require.Empty(t, diags) + + assert.Equal(t, "/Volumes/main/default/code/existing.tar.gz", + b.Config.Resources.Jobs["train"].Tasks[0].AiRuntimeTask.CodeSourcePath) + assert.Empty(t, f.written, "remote code_source_path must not trigger an upload") +} + +func TestPackageAndUploadSkipsUploadWhenArchiveExists(t *testing.T) { + b, _ := bundleWithCodeSource(t, "src") + + // Determine the archive name the mutator will compute, then mark it present. + var buf bytes.Buffer + sha, err := buildTarball(t.Context(), vfs.MustNew(filepath.Join(b.SyncRootPath, "src")), "src", &buf) + require.NoError(t, err) + archiveName := "src_" + sha[:16] + ".tar.gz" + + f := newRecordingFiler() + f.exists[archiveName] = true + + diags := bundle.Apply(t.Context(), b, &packageAndUpload{client: f}) + require.Empty(t, diags) + + assert.Empty(t, f.written, "existing archive must not be re-uploaded") + assert.Equal(t, "/Workspace/Users/me/.bundle/artifacts/.internal/"+archiveName, + b.Config.Resources.Jobs["train"].Tasks[0].AiRuntimeTask.CodeSourcePath) +} diff --git a/bundle/config/mutator/aicode/tarball.go b/bundle/config/mutator/aicode/tarball.go new file mode 100644 index 00000000000..d09a9a8afde --- /dev/null +++ b/bundle/config/mutator/aicode/tarball.go @@ -0,0 +1,122 @@ +package aicode + +import ( + "archive/tar" + "compress/gzip" + "context" + "crypto/sha256" + "encoding/hex" + "fmt" + "io" + "path" + "slices" + "strings" + "time" + + "github.com/databricks/cli/libs/fileset" + "github.com/databricks/cli/libs/git" + "github.com/databricks/cli/libs/vfs" +) + +// tarEpoch is a fixed modification time stamped on every tar entry so the +// archive is content-addressed: identical file contents always produce an +// identical archive (and therefore an identical SHA-256) regardless of the +// files' mtimes or when the archive was built. This is what makes the cache +// key in packageAndUpload stable across deploys and keeps acceptance output +// deterministic. +var tarEpoch = time.Date(2000, 1, 1, 0, 0, 0, 0, time.UTC) + +// appleDoublePrefix is the basename prefix of macOS AppleDouble metadata files. +// The AIR CLI excludes these (they sort before the real top-level entry and add +// useless per-file metadata); we match that behavior so archives built on macOS +// and Linux are identical. +const appleDoublePrefix = "._" + +// buildTarball writes a reproducible gzipped tarball of the directory rooted at +// codeRoot to out and returns its SHA-256 hex digest. Every entry is prefixed +// with prefix (the code directory's basename) so the archive expands to +// /... on extraction, matching the AIR CLI's `--prefix=/` layout +// and the runtime's /databricks/code_source/ contract. +// +// The file list is gitignore-aware: it honors the code directory's .gitignore +// and always excludes .git and .databricks (via [git.NewFileSetAtRoot]), the +// same walker that backs bundle file sync. +func buildTarball(ctx context.Context, codeRoot vfs.Path, prefix string, out io.Writer) (string, error) { + fsys, err := git.NewFileSetAtRoot(ctx, codeRoot) + if err != nil { + return "", err + } + + fileList, err := fsys.Files() + if err != nil { + return "", err + } + + // Sort by relative path so the archive byte stream (and thus its hash) does + // not depend on filesystem walk order. + slices.SortFunc(fileList, func(a, b fileset.File) int { + return strings.Compare(a.Relative, b.Relative) + }) + + hash := sha256.New() + gzw := gzip.NewWriter(io.MultiWriter(out, hash)) + tw := tar.NewWriter(gzw) + + for _, f := range fileList { + if err := addFileToTar(tw, codeRoot, f, prefix); err != nil { + return "", err + } + } + + if err := tw.Close(); err != nil { + return "", err + } + if err := gzw.Close(); err != nil { + return "", err + } + + return hex.EncodeToString(hash.Sum(nil)), nil +} + +func addFileToTar(tw *tar.Writer, codeRoot vfs.Path, f fileset.File, prefix string) error { + // fileset.File.Relative is already slash-separated, so path.Base is correct. + if strings.HasPrefix(path.Base(f.Relative), appleDoublePrefix) { + return nil + } + + rc, err := codeRoot.Open(f.Relative) + if err != nil { + return fmt.Errorf("open %s: %w", f.Relative, err) + } + defer rc.Close() + + info, err := rc.Stat() + if err != nil { + return fmt.Errorf("stat %s: %w", f.Relative, err) + } + + // Only regular files are archived. The gitignore-aware walker never yields + // directories, and symlinks inside a code snapshot are out of scope for v1. + if !info.Mode().IsRegular() { + return nil + } + + hdr := &tar.Header{ + Typeflag: tar.TypeReg, + Name: path.Join(prefix, f.Relative), + Size: info.Size(), + // Normalize permissions to a fixed 0644 and zero the mtime so the archive + // is reproducible across machines. Executable bits are intentionally + // dropped; the runtime invokes code via an interpreter, not by executing + // files from the snapshot directly. + Mode: 0o644, + ModTime: tarEpoch, + } + if err := tw.WriteHeader(hdr); err != nil { + return fmt.Errorf("tar header for %s: %w", f.Relative, err) + } + if _, err := io.Copy(tw, rc); err != nil { + return fmt.Errorf("write %s: %w", f.Relative, err) + } + return nil +} diff --git a/bundle/config/mutator/aicode/tarball_test.go b/bundle/config/mutator/aicode/tarball_test.go new file mode 100644 index 00000000000..1ac3c9d57a8 --- /dev/null +++ b/bundle/config/mutator/aicode/tarball_test.go @@ -0,0 +1,108 @@ +package aicode + +import ( + "archive/tar" + "bytes" + "compress/gzip" + "io" + "os" + "path/filepath" + "testing" + + "github.com/databricks/cli/libs/vfs" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// writeCodeDir materializes files (relative path -> content) under a fresh temp +// directory and returns a vfs.Path rooted at it. +func writeCodeDir(t *testing.T, files map[string]string) vfs.Path { + t.Helper() + dir := t.TempDir() + for name, content := range files { + p := filepath.Join(dir, filepath.FromSlash(name)) + require.NoError(t, os.MkdirAll(filepath.Dir(p), 0o755)) + require.NoError(t, os.WriteFile(p, []byte(content), 0o644)) + } + return vfs.MustNew(dir) +} + +// tarEntries reads a gzipped tarball and returns entry name -> content. +func tarEntries(t *testing.T, b []byte) map[string]string { + t.Helper() + gzr, err := gzip.NewReader(bytes.NewReader(b)) + require.NoError(t, err) + tr := tar.NewReader(gzr) + out := map[string]string{} + for { + hdr, err := tr.Next() + if err == io.EOF { + break + } + require.NoError(t, err) + content, err := io.ReadAll(tr) + require.NoError(t, err) + out[hdr.Name] = string(content) + } + return out +} + +func TestBuildTarballPrefixesEntriesAndOmitsGitAndGitignored(t *testing.T) { + root := writeCodeDir(t, map[string]string{ + "train.py": "print('train')", + "pkg/util.py": "x = 1", + ".gitignore": "ignored.txt\n*.log\n", + "ignored.txt": "should not be archived", + "debug.log": "should not be archived", + ".git/config": "[core]", + "._resource_fork": "apple double metadata", + }) + + var buf bytes.Buffer + sha, err := buildTarball(t.Context(), root, "mycode", &buf) + require.NoError(t, err) + require.NotEmpty(t, sha) + + entries := tarEntries(t, buf.Bytes()) + assert.Equal(t, "print('train')", entries["mycode/train.py"]) + assert.Equal(t, "x = 1", entries["mycode/pkg/util.py"]) + // .gitignore itself is a tracked file and is archived. + assert.Contains(t, entries, "mycode/.gitignore") + + assert.NotContains(t, entries, "mycode/ignored.txt", "gitignored file must be excluded") + assert.NotContains(t, entries, "mycode/debug.log", "gitignored glob must be excluded") + assert.NotContains(t, entries, "mycode/.git/config", ".git must never be archived") + assert.NotContains(t, entries, "mycode/._resource_fork", "AppleDouble metadata must be excluded") +} + +func TestBuildTarballIsDeterministic(t *testing.T) { + files := map[string]string{ + "a.py": "aaa", + "sub/b.py": "bbb", + "sub/c.txt": "ccc", + } + root1 := writeCodeDir(t, files) + root2 := writeCodeDir(t, files) + + var buf1, buf2 bytes.Buffer + sha1, err := buildTarball(t.Context(), root1, "code", &buf1) + require.NoError(t, err) + sha2, err := buildTarball(t.Context(), root2, "code", &buf2) + require.NoError(t, err) + + assert.Equal(t, sha1, sha2, "identical content must produce an identical hash") + assert.Equal(t, buf1.Bytes(), buf2.Bytes(), "identical content must produce identical bytes") +} + +func TestBuildTarballHashChangesWithContent(t *testing.T) { + root1 := writeCodeDir(t, map[string]string{"main.py": "v1"}) + root2 := writeCodeDir(t, map[string]string{"main.py": "v2"}) + + var buf1, buf2 bytes.Buffer + sha1, err := buildTarball(t.Context(), root1, "code", &buf1) + require.NoError(t, err) + sha2, err := buildTarball(t.Context(), root2, "code", &buf2) + require.NoError(t, err) + + assert.NotEqual(t, sha1, sha2) +} diff --git a/bundle/config/mutator/aicode/validate.go b/bundle/config/mutator/aicode/validate.go new file mode 100644 index 00000000000..e5d7a0cdb8f --- /dev/null +++ b/bundle/config/mutator/aicode/validate.go @@ -0,0 +1,101 @@ +package aicode + +import ( + "context" + "fmt" + "os" + "path/filepath" + + "github.com/databricks/cli/bundle" + "github.com/databricks/cli/bundle/libraries" + "github.com/databricks/cli/libs/diag" + "github.com/databricks/cli/libs/dyn" +) + +// Validate checks AI Runtime tasks that reference a local code_source_path so +// that misconfigurations surface at `bundle validate` time with an actionable +// message, rather than as an obscure failure mid-deploy. It performs no uploads. +func Validate() bundle.ReadOnlyMutator { + return &validate{} +} + +type validate struct{ bundle.RO } + +func (v *validate) Name() string { + return "aicode.Validate" +} + +func (v *validate) Apply(ctx context.Context, b *bundle.Bundle) diag.Diagnostics { + var diags diag.Diagnostics + + jobsPath := dyn.NewPath(dyn.Key("resources"), dyn.Key("jobs")) + + for name, job := range b.Config.Resources.Jobs { + jobPath := jobsPath.Append(dyn.Key(name)) + + for i, task := range job.Tasks { + if task.AiRuntimeTask == nil { + continue + } + codePath := jobPath.Append(dyn.Key("tasks"), dyn.Index(i), + dyn.Key("ai_runtime_task"), dyn.Key("code_source_path")) + diags = diags.Extend(v.validateTask(b, job.GitSource != nil, task.AiRuntimeTask.CodeSourcePath, codePath)) + } + } + + return diags +} + +func (v *validate) validateTask(b *bundle.Bundle, jobHasGitSource bool, codeSourcePath string, codePath dyn.Path) diag.Diagnostics { + // Only local code_source_path values are packaged at deploy; remote values + // are used as-is and need no validation here. + if codeSourcePath == "" || !libraries.IsLocalPath(codeSourcePath) { + return nil + } + + locations := b.Config.GetLocations(codePath.String()) + + // The deploy engine retrieves task files from git when git_source is set, so + // packaging a local directory would be silently ignored. Reject the combination. + if jobHasGitSource { + return diag.Diagnostics{{ + Severity: diag.Error, + Summary: "ai_runtime_task with a local code_source_path cannot be combined with git_source", + Detail: "Remove git_source or set code_source_path to a workspace or volume path", + Locations: locations, + Paths: []dyn.Path{codePath}, + }} + } + + // Immutable-folder deployments upload a single content-addressed snapshot and + // do not support the per-task code packaging this mutator performs. + if b.IsImmutableFolder() { + return diag.Diagnostics{{ + Severity: diag.Error, + Summary: "ai_runtime_task with a local code_source_path is not supported with experimental.immutable_folder", + Locations: locations, + Paths: []dyn.Path{codePath}, + }} + } + + localDir := filepath.Join(b.SyncRootPath, filepath.FromSlash(codeSourcePath)) + info, err := os.Stat(localDir) + if err != nil { + return diag.Diagnostics{{ + Severity: diag.Error, + Summary: fmt.Sprintf("code_source_path %q not found", codeSourcePath), + Locations: locations, + Paths: []dyn.Path{codePath}, + }} + } + if !info.IsDir() { + return diag.Diagnostics{{ + Severity: diag.Error, + Summary: fmt.Sprintf("code_source_path %q must be a directory", codeSourcePath), + Locations: locations, + Paths: []dyn.Path{codePath}, + }} + } + + return nil +} diff --git a/bundle/config/mutator/aicode/validate_test.go b/bundle/config/mutator/aicode/validate_test.go new file mode 100644 index 00000000000..cd3548372e1 --- /dev/null +++ b/bundle/config/mutator/aicode/validate_test.go @@ -0,0 +1,64 @@ +package aicode + +import ( + "path/filepath" + "testing" + + "github.com/databricks/cli/bundle" + "github.com/databricks/cli/bundle/config" + "github.com/databricks/cli/bundle/config/resources" + "github.com/databricks/cli/bundle/internal/bundletest" + "github.com/databricks/cli/libs/dyn" + "github.com/databricks/databricks-sdk-go/service/jobs" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func bundleForValidate(t *testing.T, codeSourcePath string, gitSource *jobs.GitSource) *bundle.Bundle { + t.Helper() + dir := t.TempDir() + b := &bundle.Bundle{ + BundleRootPath: dir, + SyncRootPath: dir, + Config: config.Root{ + Bundle: config.Bundle{Target: "default"}, + Resources: config.Resources{ + Jobs: map[string]*resources.Job{ + "train": { + JobSettings: jobs.JobSettings{ + GitSource: gitSource, + Tasks: []jobs.Task{ + { + TaskKey: "train", + AiRuntimeTask: &jobs.AiRuntimeTask{CodeSourcePath: codeSourcePath}, + }, + }, + }, + }, + }, + }, + }, + } + bundletest.SetLocation(b, ".", []dyn.Location{{File: filepath.Join(dir, "databricks.yml")}}) + return b +} + +func TestValidateMissingCodeSourceDir(t *testing.T) { + b := bundleForValidate(t, "does-not-exist", nil) + diags := Validate().Apply(t.Context(), b) + require.Len(t, diags, 1) + assert.Equal(t, `code_source_path "does-not-exist" not found`, diags[0].Summary) +} + +func TestValidateGitSourceConflict(t *testing.T) { + b := bundleForValidate(t, "src", &jobs.GitSource{GitUrl: "https://example.invalid/repo"}) + diags := Validate().Apply(t.Context(), b) + require.Len(t, diags, 1) + assert.Contains(t, diags[0].Summary, "cannot be combined with git_source") +} + +func TestValidateRemoteCodeSourceIsSkipped(t *testing.T) { + b := bundleForValidate(t, "/Volumes/main/default/code/x.tar.gz", nil) + diags := Validate().Apply(t.Context(), b) + assert.Empty(t, diags) +} diff --git a/bundle/phases/deploy.go b/bundle/phases/deploy.go index e2159189c78..906cfca4f90 100644 --- a/bundle/phases/deploy.go +++ b/bundle/phases/deploy.go @@ -10,6 +10,7 @@ import ( "github.com/databricks/cli/bundle/config" "github.com/databricks/cli/bundle/config/engine" "github.com/databricks/cli/bundle/config/mutator" + "github.com/databricks/cli/bundle/config/mutator/aicode" "github.com/databricks/cli/bundle/deploy" "github.com/databricks/cli/bundle/deploy/files" "github.com/databricks/cli/bundle/deploy/lock" @@ -196,6 +197,12 @@ func Deploy(ctx context.Context, b *bundle.Bundle, outputHandler sync.OutputHand } bundle.ApplySeqContext(ctx, b, + // Package and upload local code referenced by AI Runtime tasks, rewriting + // each ai_runtime_task.code_source_path to its uploaded workspace/volume + // location before the plan is computed so the created job points at the + // remote archive. Runs after file/library upload (artifact_path resolved) + // and before RunPlan for both the direct and terraform engines. + aicode.PackageAndUpload(), deploy.StateUpdate(), deploy.StatePush(), permissions.ApplyWorkspaceRootPermissions(), diff --git a/bundle/phases/initialize.go b/bundle/phases/initialize.go index bfa2af4124b..f1554a4b7a6 100644 --- a/bundle/phases/initialize.go +++ b/bundle/phases/initialize.go @@ -10,6 +10,7 @@ import ( "github.com/databricks/cli/bundle/artifacts" "github.com/databricks/cli/bundle/config" "github.com/databricks/cli/bundle/config/mutator" + "github.com/databricks/cli/bundle/config/mutator/aicode" pythonmutator "github.com/databricks/cli/bundle/config/mutator/python" "github.com/databricks/cli/bundle/config/validate" "github.com/databricks/cli/bundle/deploy/metadata" @@ -189,6 +190,12 @@ func Initialize(ctx context.Context, b *bundle.Bundle) { mutator.TranslatePaths(), + // Reads (typed): resources.jobs.*.tasks[*].ai_runtime_task.code_source_path, job git_source + // Validates that AI Runtime tasks referencing a local code_source_path point at an existing + // directory and are not combined with git_source or immutable-folder deployments, so these + // misconfigurations are caught at validate time rather than mid-deploy. + aicode.Validate(), + // Reads (typed): b.Config.Experimental.PythonWheelWrapper, b.Config.Presets.SourceLinkedDeployment (checks Python wheel wrapper and deployment mode settings) // Reads (dynamic): resources.jobs.*.tasks (checks for tasks with local libraries and incompatible DBR versions) // Provides warnings when Python wheel tasks are used with DBR < 13.3 or when wheel wrapper is incompatible with source-linked deployment diff --git a/cmd/experimental/experimental.go b/cmd/experimental/experimental.go index 8d9827c5c94..d87c893abc5 100644 --- a/cmd/experimental/experimental.go +++ b/cmd/experimental/experimental.go @@ -1,6 +1,7 @@ package experimental import ( + aircmd "github.com/databricks/cli/experimental/air/cmd" aitoolscmd "github.com/databricks/cli/experimental/aitools/cmd" geniecmd "github.com/databricks/cli/experimental/genie/cmd" postgrescmd "github.com/databricks/cli/experimental/postgres/cmd" @@ -22,6 +23,7 @@ These commands provide early access to new features that are still under development. They may change or be removed in future versions without notice.`, } + cmd.AddCommand(aircmd.New()) cmd.AddCommand(aitoolscmd.NewAitoolsCmd()) cmd.AddCommand(geniecmd.NewGenieCmd()) cmd.AddCommand(postgrescmd.New()) diff --git a/experimental/air/.gitkeep b/experimental/air/.gitkeep deleted file mode 100644 index e69de29bb2d..00000000000 diff --git a/experimental/air/cmd/air.go b/experimental/air/cmd/air.go new file mode 100644 index 00000000000..d99c5c660b1 --- /dev/null +++ b/experimental/air/cmd/air.go @@ -0,0 +1,17 @@ +package aircmd + +import "github.com/spf13/cobra" + +// New creates the parent "air" command group for migrating AIR CLI workloads +// onto Databricks Asset Bundles. +func New() *cobra.Command { + cmd := &cobra.Command{ + Use: "air", + Short: "Migrate AI Runtime (AIR) CLI workloads to Databricks Asset Bundles", + Hidden: true, + } + + cmd.AddCommand(newInitCommand()) + + return cmd +} diff --git a/experimental/air/cmd/init.go b/experimental/air/cmd/init.go new file mode 100644 index 00000000000..93a7ec74893 --- /dev/null +++ b/experimental/air/cmd/init.go @@ -0,0 +1,145 @@ +package aircmd + +import ( + "fmt" + "os" + "path/filepath" + + "github.com/databricks/cli/bundle/generate" + "github.com/databricks/cli/experimental/air/trainyaml" + "github.com/databricks/cli/libs/cmdio" + "github.com/databricks/cli/libs/dyn" + "github.com/databricks/cli/libs/dyn/yamlsaver" + "github.com/databricks/cli/libs/textutil" + "github.com/databricks/databricks-sdk-go/service/jobs" + "github.com/spf13/cobra" +) + +func newInitCommand() *cobra.Command { + var fromPath string + var outputDir string + var jobKey string + var force bool + + cmd := &cobra.Command{ + Use: "init", + Short: "Scaffold a Databricks Asset Bundle from an AIR CLI train.yaml", + Long: `Scaffold a Databricks Asset Bundle from an AIR CLI train.yaml. + +This reads an existing AIR CLI train.yaml and generates a bundle that deploys the +same workload as a durable job using an AI Runtime task. The generated bundle +contains: + + - databricks.yml the bundle definition and a dev target + - resources/.job.yml the job with an ai_runtime_task + - /command.sh the command materialized as a script + +Deploy the generated bundle with: + databricks bundle deploy`, + } + + cmd.Flags().StringVarP(&fromPath, "from", "f", "", "Path to the AIR CLI train.yaml to migrate") + cmd.MarkFlagRequired("from") + cmd.Flags().StringVar(&outputDir, "output-dir", ".", "Directory to write the generated bundle into") + cmd.Flags().StringVar(&jobKey, "key", "", "Resource key for the generated job (defaults to the experiment name)") + cmd.Flags().BoolVar(&force, "force", false, "Overwrite existing files in the output directory") + + cmd.RunE = func(cmd *cobra.Command, args []string) error { + ctx := cmd.Context() + + cfg, err := trainyaml.Parse(fromPath) + if err != nil { + return err + } + + res, err := trainyaml.Convert(cfg) + if err != nil { + return err + } + + for _, w := range res.Warnings { + cmdio.LogString(ctx, "Warning: "+w) + } + + if jobKey == "" { + jobKey = textutil.NormalizeString(res.Job.Name) + } + + if err := writeBundle(outputDir, jobKey, res, force); err != nil { + return err + } + + cmdio.LogString(ctx, "Bundle written to "+filepath.ToSlash(outputDir)) + cmdio.LogString(ctx, "Deploy it with: databricks bundle deploy") + return nil + } + + return cmd +} + +// writeBundle materializes the generated bundle files under outputDir. +func writeBundle(outputDir, jobKey string, res *trainyaml.Result, force bool) error { + jobValue, err := generate.ConvertJobToValue(&jobs.Job{Settings: &res.Job.JobSettings}) + if err != nil { + return err + } + + saver := yamlsaver.NewSaver() + + // yamlsaver orders map keys by their location line, so assign increasing + // lines to fix the top-level key order (bundle, include, workspace). + // Without this the order is nondeterministic (all keys default to line 0). + rootValue := map[string]dyn.Value{ + "bundle": dyn.NewValue(map[string]dyn.Value{ + "name": dyn.V(jobKey), + }, []dyn.Location{{Line: 0}}), + "include": dyn.NewValue([]dyn.Value{dyn.V("resources/*.yml")}, []dyn.Location{{Line: 1}}), + } + // The snapshot's remote volume becomes the artifact path, so the code archive + // is uploaded to the volume the user already configured for AIR. + if res.ArtifactPath != "" { + rootValue["workspace"] = dyn.NewValue(map[string]dyn.Value{ + "artifact_path": dyn.V(res.ArtifactPath), + }, []dyn.Location{{Line: 2}}) + } + + resourcesDir := filepath.Join(outputDir, "resources") + if err := os.MkdirAll(resourcesDir, 0o755); err != nil { + return err + } + + if err := saver.SaveAsYAML(rootValue, filepath.Join(outputDir, "databricks.yml"), force); err != nil { + return err + } + + jobFile := map[string]dyn.Value{ + "resources": dyn.V(map[string]dyn.Value{ + "jobs": dyn.V(map[string]dyn.Value{ + jobKey: jobValue, + }), + }), + } + if err := saver.SaveAsYAML(jobFile, filepath.Join(resourcesDir, jobKey+".job.yml"), force); err != nil { + return err + } + + // Materialize the command next to the code it runs against, so command_path + // (which is relative to the extracted code source) resolves at runtime. + codeDir := filepath.Join(outputDir, filepath.FromSlash(codeSourceRootPath(res))) + if err := os.MkdirAll(codeDir, 0o755); err != nil { + return err + } + commandPath := filepath.Join(codeDir, "command.sh") + if !force { + if _, err := os.Stat(commandPath); err == nil { + return fmt.Errorf("%s already exists, use --force to overwrite", filepath.ToSlash(commandPath)) + } + } + return os.WriteFile(commandPath, []byte(res.CommandScript), 0o755) +} + +// codeSourceRootPath returns the local code_source_path of the generated job's +// AI Runtime task, i.e. the directory the command script is written into. +func codeSourceRootPath(res *trainyaml.Result) string { + return res.Job.Tasks[0].AiRuntimeTask.CodeSourcePath +} diff --git a/experimental/air/trainyaml/convert.go b/experimental/air/trainyaml/convert.go new file mode 100644 index 00000000000..11dd926be6f --- /dev/null +++ b/experimental/air/trainyaml/convert.go @@ -0,0 +1,168 @@ +package trainyaml + +import ( + "errors" + "fmt" + "path" + + "github.com/databricks/cli/bundle/config/resources" + "github.com/databricks/databricks-sdk-go/service/compute" + "github.com/databricks/databricks-sdk-go/service/jobs" +) + +// environmentKey is the job environment key the generated AI Runtime task +// references. AI Runtime tasks require a serverless environment; we materialize +// one visibly in the generated config so the user can edit it afterwards. +const environmentKey = "default" + +// commandScriptName is the file the inline train.yaml command is written to. +// The AI Runtime task's command_path points at it, relative to the code source. +const commandScriptName = "command.sh" + +// Result holds everything Convert produces from a train.yaml: the DABs job +// resource, the shell script the command was materialized into, the artifact +// path (set only when the snapshot targets a UC volume), and any warnings about +// train.yaml fields that could not be represented. +type Result struct { + Job *resources.Job + CommandScript string + ArtifactPath string + Warnings []string +} + +// Convert maps a parsed train.yaml Config to a DABs job resource. It returns an +// error for inputs that cannot be represented as an AI Runtime task (legacy +// accelerators, pools, docker images) and collects warnings for fields that are +// dropped because they have no equivalent. +func Convert(cfg *Config) (*Result, error) { + if cfg.ExperimentName == "" { + return nil, errors.New("experiment_name is required") + } + if cfg.Command == "" { + return nil, errors.New("command is required") + } + if cfg.CodeSource == nil || cfg.CodeSource.Snapshot == nil || cfg.CodeSource.Snapshot.RootPath == "" { + return nil, errors.New("code_source.snapshot.root_path is required to migrate to an AI Runtime task") + } + + accelerator, err := convertAcceleratorType(cfg.Compute.AcceleratorType) + if err != nil { + return nil, err + } + if cfg.Compute.NodePoolID != "" || cfg.Compute.PoolName != "" { + return nil, errors.New("compute pools (node_pool_id/pool_name) are not supported by ai_runtime_task; remove them to use serverless AI Runtime compute") + } + + env, err := convertEnvironment(cfg.Environment) + if err != nil { + return nil, err + } + + res := &Result{} + + task := jobs.Task{ + TaskKey: cfg.ExperimentName, + EnvironmentKey: environmentKey, + AiRuntimeTask: &jobs.AiRuntimeTask{ + Experiment: cfg.ExperimentName, + MlflowRun: cfg.MlflowRunName, + MlflowExperimentDirectory: cfg.MlflowExperimentDirectory, + CodeSourcePath: cfg.CodeSource.Snapshot.RootPath, + Deployments: []jobs.DeploymentSpec{ + { + // command_path is relative to the extracted code source root. + CommandPath: path.Join(path.Base(cfg.CodeSource.Snapshot.RootPath), commandScriptName), + Compute: jobs.ComputeSpec{ + AcceleratorType: accelerator, + AcceleratorCount: cfg.Compute.NumAccelerators, + }, + }, + }, + }, + } + if cfg.MaxRetries != nil { + task.MaxRetries = *cfg.MaxRetries + } + if cfg.TimeoutMinutes != nil { + task.TimeoutSeconds = *cfg.TimeoutMinutes * 60 + } + + res.Job = &resources.Job{ + JobSettings: jobs.JobSettings{ + Name: cfg.ExperimentName, + Tasks: []jobs.Task{task}, + Environments: []jobs.JobEnvironment{{EnvironmentKey: environmentKey, Spec: env}}, + }, + } + + res.CommandScript = cfg.Command + if cfg.CodeSource.Snapshot.RemoteVolume != "" { + res.ArtifactPath = cfg.CodeSource.Snapshot.RemoteVolume + } + + res.Warnings = append(res.Warnings, unmappedWarnings(cfg)...) + return res, nil +} + +func convertAcceleratorType(v string) (jobs.ComputeSpecAcceleratorType, error) { + if v == "" { + return "", errors.New("compute.accelerator_type is required") + } + var t jobs.ComputeSpecAcceleratorType + if err := t.Set(v); err != nil { + return "", fmt.Errorf("compute.accelerator_type %q is not supported by ai_runtime_task (expected GPU_1xA10, GPU_1xH100, or GPU_8xH100)", v) + } + return t, nil +} + +func convertEnvironment(env *Environment) (*compute.Environment, error) { + if env == nil { + return &compute.Environment{}, nil + } + if env.DockerImage != nil { + return nil, errors.New("environment.docker_image is not supported by ai_runtime_task; use environment.version and environment.dependencies instead") + } + if env.Dependencies.Path != "" { + return nil, errors.New("environment.dependencies as a requirements file path is not supported; provide an inline list of dependencies instead") + } + + return &compute.Environment{ + EnvironmentVersion: env.Version, + Dependencies: env.Dependencies.List, + }, nil +} + +// unmappedWarnings reports train.yaml fields that have no ai_runtime_task +// equivalent and are therefore dropped from the generated config. +func unmappedWarnings(cfg *Config) []string { + var warnings []string + warn := func(field, hint string) { + warnings = append(warnings, fmt.Sprintf("%q has no ai_runtime_task equivalent and was dropped%s", field, hint)) + } + + if len(cfg.Secrets) > 0 { + warn("secrets", "; configure secrets on the job or via {{secrets/scope/key}} references") + } + if len(cfg.EnvVariables) > 0 { + warn("env_variables", "; set them inside command.sh or the serverless environment") + } + if len(cfg.Parameters) > 0 { + warn("parameters", "") + } + if cfg.Priority != nil { + warn("priority", "") + } + if cfg.UsagePolicyName != "" { + warn("usage_policy_name", "") + } + if cfg.UsagePolicyID != "" { + warn("usage_policy_id", "") + } + if cfg.IdempotencyToken != "" { + warn("idempotency_token", "") + } + if cfg.CodeSource != nil && cfg.CodeSource.Snapshot != nil && len(cfg.CodeSource.Snapshot.IncludePaths) > 0 { + warn("code_source.snapshot.include_paths", "; the entire code_source directory is packaged") + } + return warnings +} diff --git a/experimental/air/trainyaml/convert_test.go b/experimental/air/trainyaml/convert_test.go new file mode 100644 index 00000000000..ea44722771e --- /dev/null +++ b/experimental/air/trainyaml/convert_test.go @@ -0,0 +1,134 @@ +package trainyaml + +import ( + "os" + "path/filepath" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func writeTrainYAML(t *testing.T, content string) string { + t.Helper() + dir := t.TempDir() + p := filepath.Join(dir, "train.yaml") + require.NoError(t, os.WriteFile(p, []byte(content), 0o644)) + return p +} + +func TestParseAndConvertSnapshot(t *testing.T) { + p := writeTrainYAML(t, ` +experiment_name: my-training +compute: + num_accelerators: 8 + accelerator_type: GPU_8xH100 +environment: + version: "5" + dependencies: + - torch>=2.0.0 +code_source: + type: snapshot + snapshot: + root_path: ./src +command: | + cd $CODE_SOURCE_PATH + python train.py +max_retries: 2 +timeout_minutes: 30 +`) + + cfg, err := Parse(p) + require.NoError(t, err) + + res, err := Convert(cfg) + require.NoError(t, err) + require.NotNil(t, res.Job) + + task := res.Job.Tasks[0] + assert.Equal(t, "my-training", task.TaskKey) + assert.Equal(t, "default", task.EnvironmentKey) + assert.Equal(t, 2, task.MaxRetries) + assert.Equal(t, 1800, task.TimeoutSeconds) + + air := task.AiRuntimeTask + require.NotNil(t, air) + assert.Equal(t, "my-training", air.Experiment) + assert.Equal(t, "./src", air.CodeSourcePath) + require.Len(t, air.Deployments, 1) + assert.Equal(t, "src/command.sh", air.Deployments[0].CommandPath) + assert.EqualValues(t, "GPU_8xH100", air.Deployments[0].Compute.AcceleratorType) + assert.Equal(t, 8, air.Deployments[0].Compute.AcceleratorCount) + + require.Len(t, res.Job.Environments, 1) + assert.Equal(t, "default", res.Job.Environments[0].EnvironmentKey) + assert.Equal(t, "5", res.Job.Environments[0].Spec.EnvironmentVersion) + assert.Equal(t, []string{"torch>=2.0.0"}, res.Job.Environments[0].Spec.Dependencies) + + assert.Contains(t, res.CommandScript, "python train.py") + assert.Empty(t, res.ArtifactPath) +} + +func TestConvertRemoteVolumeSetsArtifactPath(t *testing.T) { + cfg := &Config{ + ExperimentName: "exp", + Command: "echo hi", + Compute: Compute{NumAccelerators: 1, AcceleratorType: "GPU_1xA10"}, + CodeSource: &CodeSource{ + Type: "snapshot", + Snapshot: &Snapshot{RootPath: "code", RemoteVolume: "/Volumes/main/default/code"}, + }, + } + res, err := Convert(cfg) + require.NoError(t, err) + assert.Equal(t, "/Volumes/main/default/code", res.ArtifactPath) +} + +func TestConvertRejectsLegacyAccelerator(t *testing.T) { + cfg := &Config{ + ExperimentName: "exp", + Command: "echo hi", + Compute: Compute{NumAccelerators: 1, AcceleratorType: "a10"}, + CodeSource: &CodeSource{Type: "snapshot", Snapshot: &Snapshot{RootPath: "code"}}, + } + _, err := Convert(cfg) + require.Error(t, err) + assert.Contains(t, err.Error(), "not supported by ai_runtime_task") +} + +func TestConvertRejectsDockerImage(t *testing.T) { + cfg := &Config{ + ExperimentName: "exp", + Command: "echo hi", + Compute: Compute{NumAccelerators: 1, AcceleratorType: "GPU_1xA10"}, + Environment: &Environment{DockerImage: &DockerImage{URL: "org/repo:tag"}}, + CodeSource: &CodeSource{Type: "snapshot", Snapshot: &Snapshot{RootPath: "code"}}, + } + _, err := Convert(cfg) + require.Error(t, err) + assert.Contains(t, err.Error(), "docker_image is not supported") +} + +func TestConvertWarnsOnUnmappedFields(t *testing.T) { + cfg := &Config{ + ExperimentName: "exp", + Command: "echo hi", + Compute: Compute{NumAccelerators: 1, AcceleratorType: "GPU_1xA10"}, + Secrets: map[string]string{"HF_TOKEN": "scope/key"}, + CodeSource: &CodeSource{ + Type: "snapshot", + Snapshot: &Snapshot{RootPath: "code", IncludePaths: []string{"a", "b"}}, + }, + } + res, err := Convert(cfg) + require.NoError(t, err) + require.Len(t, res.Warnings, 2) + assert.Contains(t, res.Warnings[0], "secrets") + assert.Contains(t, res.Warnings[1], "include_paths") +} + +func TestParseRejectsUnknownField(t *testing.T) { + p := writeTrainYAML(t, "experiment_name: x\nbogus_field: 1\n") + _, err := Parse(p) + require.Error(t, err) +} diff --git a/experimental/air/trainyaml/parse.go b/experimental/air/trainyaml/parse.go new file mode 100644 index 00000000000..dc29ba264d6 --- /dev/null +++ b/experimental/air/trainyaml/parse.go @@ -0,0 +1,110 @@ +// Package trainyaml parses the AIR CLI's train.yaml launch config and converts +// it into a DABs job resource that uses an AI Runtime task. It exists so users +// who previously launched ephemeral runs with the AIR CLI can migrate onto a +// durable, DABs-managed job while keeping the familiar code_source experience. +package trainyaml + +import ( + "bytes" + "fmt" + "os" + + "go.yaml.in/yaml/v3" +) + +// Config is the subset of the AIR CLI train.yaml schema this migration +// understands. Fields that have no DABs equivalent are still parsed so we can +// warn about them rather than silently drop them (see Convert). +// +// Reference: universe/ai-compute/cli/cli/sdk/config.py (RunConfig and friends). +type Config struct { + ExperimentName string `yaml:"experiment_name"` + Compute Compute `yaml:"compute"` + Command string `yaml:"command"` + + Environment *Environment `yaml:"environment"` + CodeSource *CodeSource `yaml:"code_source"` + + EnvVariables map[string]string `yaml:"env_variables"` + Secrets map[string]string `yaml:"secrets"` + + MaxRetries *int `yaml:"max_retries"` + TimeoutMinutes *int `yaml:"timeout_minutes"` + MlflowRunName string `yaml:"mlflow_run_name"` + MlflowExperimentDirectory string `yaml:"mlflow_experiment_directory"` + + // Fields with no AI Runtime task equivalent. Retained so Convert can warn. + Parameters map[string]any `yaml:"parameters"` + Priority *int `yaml:"priority"` + UsagePolicyName string `yaml:"usage_policy_name"` + UsagePolicyID string `yaml:"usage_policy_id"` + IdempotencyToken string `yaml:"idempotency_token"` +} + +type Compute struct { + NumAccelerators int `yaml:"num_accelerators"` + AcceleratorType string `yaml:"accelerator_type"` + NodePoolID string `yaml:"node_pool_id"` + PoolName string `yaml:"pool_name"` +} + +// Environment mirrors the train.yaml environment block. Dependencies is +// polymorphic in the AIR schema (a requirements.yaml path or an inline list); +// we only support the inline list on the DABs serverless environment and report +// the string form as unsupported. +type Environment struct { + Dependencies dependencies `yaml:"dependencies"` + Version string `yaml:"version"` + DockerImage *DockerImage `yaml:"docker_image"` +} + +type DockerImage struct { + URL string `yaml:"url"` +} + +type CodeSource struct { + Type string `yaml:"type"` + Snapshot *Snapshot `yaml:"snapshot"` +} + +type Snapshot struct { + RootPath string `yaml:"root_path"` + RemoteVolume string `yaml:"remote_volume"` + IncludePaths []string `yaml:"include_paths"` +} + +// dependencies captures the polymorphic environment.dependencies field. Exactly +// one of List / Path is set after unmarshaling. +type dependencies struct { + List []string + Path string +} + +func (d *dependencies) UnmarshalYAML(node *yaml.Node) error { + switch node.Kind { + case yaml.ScalarNode: + return node.Decode(&d.Path) + case yaml.SequenceNode: + return node.Decode(&d.List) + default: + return fmt.Errorf("environment.dependencies must be a string or a list, got %v", node.Kind) + } +} + +// Parse reads and decodes a train.yaml file. Unknown fields are rejected so a +// typo or an unsupported field surfaces to the user instead of being ignored. +func Parse(path string) (*Config, error) { + raw, err := os.ReadFile(path) + if err != nil { + return nil, err + } + + dec := yaml.NewDecoder(bytes.NewReader(raw)) + dec.KnownFields(true) + + var cfg Config + if err := dec.Decode(&cfg); err != nil { + return nil, fmt.Errorf("failed to parse %s: %w", path, err) + } + return &cfg, nil +} From c60b4dc3ab05d0790ee4041f9e4c6bddc6b10c7c Mon Sep 17 00:00:00 2001 From: "ben.hansen" Date: Tue, 14 Jul 2026 02:53:45 +0000 Subject: [PATCH 2/5] add snapshot for ai runtime task --- NEXT_CHANGELOG.md | 4 +- .../databricks.yml | 0 .../out.test.toml | 0 .../ai_runtime_task/git_archive/output.txt | 23 +++ .../bundle/ai_runtime_task/git_archive/script | 28 +++ .../src/command.sh | 0 .../src/train.py | 0 .../ai_runtime_task/git_archive/test.toml | 13 ++ .../local_code_source/output.txt | 14 -- .../ai_runtime_task/local_code_source/script | 9 - .../local_code_source/test.toml | 11 -- .../ai_runtime_task/plain_tar/databricks.yml | 24 +++ .../ai_runtime_task/plain_tar/out.test.toml | 3 + .../ai_runtime_task/plain_tar/output.txt | 14 ++ .../bundle/ai_runtime_task/plain_tar/script | 13 ++ .../ai_runtime_task/plain_tar/src/command.sh | 2 + .../ai_runtime_task/plain_tar}/src/train.py | 0 .../ai_runtime_task/plain_tar/test.toml | 11 ++ acceptance/experimental/air/.gitkeep | 0 .../experimental/air/init/out.test.toml | 3 - acceptance/experimental/air/init/output.txt | 43 ----- acceptance/experimental/air/init/script | 11 -- acceptance/experimental/air/init/test.toml | 8 - .../air/init/train_project/train.yaml | 18 -- .../air/init_errors/docker/train.yaml | 12 -- .../air/init_errors/legacy/train.yaml | 9 - .../air/init_errors/out.test.toml | 3 - .../experimental/air/init_errors/output.txt | 8 - .../experimental/air/init_errors/script | 5 - .../experimental/air/init_errors/test.toml | 3 - .../config/mutator/aicode/package_upload.go | 144 +++++++++++---- .../mutator/aicode/package_upload_test.go | 101 +++++++---- .../mutator/aicode/snapshot_cachekey.go | 32 ++++ .../mutator/aicode/snapshot_cachekey_test.go | 33 ++++ bundle/config/mutator/aicode/snapshot_git.go | 77 ++++++++ .../config/mutator/aicode/snapshot_package.go | 134 ++++++++++++++ .../mutator/aicode/snapshot_package_test.go | 83 +++++++++ .../config/mutator/aicode/snapshot_resolve.go | 56 ++++++ bundle/config/mutator/aicode/snapshot_test.go | 72 ++++++++ bundle/config/mutator/aicode/tarball.go | 122 ------------- bundle/config/mutator/aicode/tarball_test.go | 108 ----------- .../mutator/aicode/testdata/cache_keys.json | 77 ++++++++ cmd/experimental/experimental.go | 2 - experimental/air/.gitkeep | 0 experimental/air/cmd/air.go | 17 -- experimental/air/cmd/init.go | 145 --------------- experimental/air/trainyaml/convert.go | 168 ------------------ experimental/air/trainyaml/convert_test.go | 134 -------------- experimental/air/trainyaml/parse.go | 110 ------------ 49 files changed, 868 insertions(+), 1039 deletions(-) rename acceptance/bundle/ai_runtime_task/{local_code_source => git_archive}/databricks.yml (100%) rename acceptance/bundle/ai_runtime_task/{local_code_source => git_archive}/out.test.toml (100%) create mode 100644 acceptance/bundle/ai_runtime_task/git_archive/output.txt create mode 100644 acceptance/bundle/ai_runtime_task/git_archive/script rename acceptance/bundle/ai_runtime_task/{local_code_source => git_archive}/src/command.sh (100%) rename acceptance/bundle/ai_runtime_task/{local_code_source => git_archive}/src/train.py (100%) create mode 100644 acceptance/bundle/ai_runtime_task/git_archive/test.toml delete mode 100644 acceptance/bundle/ai_runtime_task/local_code_source/output.txt delete mode 100644 acceptance/bundle/ai_runtime_task/local_code_source/script delete mode 100644 acceptance/bundle/ai_runtime_task/local_code_source/test.toml create mode 100644 acceptance/bundle/ai_runtime_task/plain_tar/databricks.yml create mode 100644 acceptance/bundle/ai_runtime_task/plain_tar/out.test.toml create mode 100644 acceptance/bundle/ai_runtime_task/plain_tar/output.txt create mode 100644 acceptance/bundle/ai_runtime_task/plain_tar/script create mode 100644 acceptance/bundle/ai_runtime_task/plain_tar/src/command.sh rename acceptance/{experimental/air/init/train_project => bundle/ai_runtime_task/plain_tar}/src/train.py (100%) create mode 100644 acceptance/bundle/ai_runtime_task/plain_tar/test.toml create mode 100644 acceptance/experimental/air/.gitkeep delete mode 100644 acceptance/experimental/air/init/out.test.toml delete mode 100644 acceptance/experimental/air/init/output.txt delete mode 100644 acceptance/experimental/air/init/script delete mode 100644 acceptance/experimental/air/init/test.toml delete mode 100644 acceptance/experimental/air/init/train_project/train.yaml delete mode 100644 acceptance/experimental/air/init_errors/docker/train.yaml delete mode 100644 acceptance/experimental/air/init_errors/legacy/train.yaml delete mode 100644 acceptance/experimental/air/init_errors/out.test.toml delete mode 100644 acceptance/experimental/air/init_errors/output.txt delete mode 100644 acceptance/experimental/air/init_errors/script delete mode 100644 acceptance/experimental/air/init_errors/test.toml create mode 100644 bundle/config/mutator/aicode/snapshot_cachekey.go create mode 100644 bundle/config/mutator/aicode/snapshot_cachekey_test.go create mode 100644 bundle/config/mutator/aicode/snapshot_git.go create mode 100644 bundle/config/mutator/aicode/snapshot_package.go create mode 100644 bundle/config/mutator/aicode/snapshot_package_test.go create mode 100644 bundle/config/mutator/aicode/snapshot_resolve.go create mode 100644 bundle/config/mutator/aicode/snapshot_test.go delete mode 100644 bundle/config/mutator/aicode/tarball.go delete mode 100644 bundle/config/mutator/aicode/tarball_test.go create mode 100644 bundle/config/mutator/aicode/testdata/cache_keys.json create mode 100644 experimental/air/.gitkeep delete mode 100644 experimental/air/cmd/air.go delete mode 100644 experimental/air/cmd/init.go delete mode 100644 experimental/air/trainyaml/convert.go delete mode 100644 experimental/air/trainyaml/convert_test.go delete mode 100644 experimental/air/trainyaml/parse.go diff --git a/NEXT_CHANGELOG.md b/NEXT_CHANGELOG.md index 7c170b89fa4..51b3ba8b214 100644 --- a/NEXT_CHANGELOG.md +++ b/NEXT_CHANGELOG.md @@ -6,11 +6,9 @@ ### CLI -* Added `databricks experimental air init`, which scaffolds a bundle from an AI Runtime (AIR) CLI `train.yaml`, generating a job with an `ai_runtime_task` so an ephemeral AIR run can be migrated to a durable, DABs-managed job. - ### Bundles -* An `ai_runtime_task.code_source_path` that points at a local directory is now packaged into a tarball during `bundle deploy`, uploaded to the workspace (or the configured UC volume), and rewritten to the uploaded path, mirroring the AIR CLI's `code_source` behavior. +* An `ai_runtime_task.code_source_path` that points at a local directory is now packaged into a tarball during `bundle deploy` (git archive of a clean HEAD, otherwise a plain tar of the working tree), uploaded to the user's workspace code snapshot cache, and rewritten to the uploaded path, mirroring the AIR CLI's `code_source` behavior. ### Dependency updates diff --git a/acceptance/bundle/ai_runtime_task/local_code_source/databricks.yml b/acceptance/bundle/ai_runtime_task/git_archive/databricks.yml similarity index 100% rename from acceptance/bundle/ai_runtime_task/local_code_source/databricks.yml rename to acceptance/bundle/ai_runtime_task/git_archive/databricks.yml diff --git a/acceptance/bundle/ai_runtime_task/local_code_source/out.test.toml b/acceptance/bundle/ai_runtime_task/git_archive/out.test.toml similarity index 100% rename from acceptance/bundle/ai_runtime_task/local_code_source/out.test.toml rename to acceptance/bundle/ai_runtime_task/git_archive/out.test.toml diff --git a/acceptance/bundle/ai_runtime_task/git_archive/output.txt b/acceptance/bundle/ai_runtime_task/git_archive/output.txt new file mode 100644 index 00000000000..0465656efcb --- /dev/null +++ b/acceptance/bundle/ai_runtime_task/git_archive/output.txt @@ -0,0 +1,23 @@ + +=== deploy packages the code source and rewrites code_source_path + +>>> [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 cache +/api/2.0/workspace-files/import-file/Workspace/Users/[USERNAME]/.air/repo_snapshots/src/src_[SNAPSHOT].tar.gz + +=== the created job's code_source_path points at the uploaded archive +/Workspace/Users/[USERNAME]/.air/repo_snapshots/src/src_[SNAPSHOT].tar.gz + +=== re-deploying the same commit 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 diff --git a/acceptance/bundle/ai_runtime_task/git_archive/script b/acceptance/bundle/ai_runtime_task/git_archive/script new file mode 100644 index 00000000000..d0715e02a43 --- /dev/null +++ b/acceptance/bundle/ai_runtime_task/git_archive/script @@ -0,0 +1,28 @@ +# A committed, clean git repo makes the snapshot use git_archive (cacheable by the +# commit SHA). Pinned commit dates keep the resolved HEAD SHA — and thus the 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 +git add src +git commit -qm 'Add code source' + +title "deploy packages the code source and rewrites code_source_path\n" +trace $CLI bundle deploy + +title "the snapshot tarball is uploaded to the user's repo_snapshots cache\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 "the created job's code_source_path points at the uploaded archive\n" +jq -r 'select(.path == "/api/2.2/jobs/create") | .body.tasks[0].ai_runtime_task.code_source_path' out.requests.txt + +rm out.requests.txt + +title "re-deploying the same commit 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 | wc -l | tr -d ' ' + +rm out.requests.txt diff --git a/acceptance/bundle/ai_runtime_task/local_code_source/src/command.sh b/acceptance/bundle/ai_runtime_task/git_archive/src/command.sh similarity index 100% rename from acceptance/bundle/ai_runtime_task/local_code_source/src/command.sh rename to acceptance/bundle/ai_runtime_task/git_archive/src/command.sh diff --git a/acceptance/bundle/ai_runtime_task/local_code_source/src/train.py b/acceptance/bundle/ai_runtime_task/git_archive/src/train.py similarity index 100% rename from acceptance/bundle/ai_runtime_task/local_code_source/src/train.py rename to acceptance/bundle/ai_runtime_task/git_archive/src/train.py diff --git a/acceptance/bundle/ai_runtime_task/git_archive/test.toml b/acceptance/bundle/ai_runtime_task/git_archive/test.toml new file mode 100644 index 00000000000..5087d92c7ec --- /dev/null +++ b/acceptance/bundle/ai_runtime_task/git_archive/test.toml @@ -0,0 +1,13 @@ +RecordRequests = true + +Ignore = [ + '.databricks', + '.git', +] + +# git_archive names the tarball _.tar.gz, where the cache key +# derives from the committed HEAD SHA. Collapse it to a stable token so output does +# not depend on the commit SHA. +[[Repls]] +Old = 'src_[0-9a-f]{16}\.tar\.gz' +New = 'src_[SNAPSHOT].tar.gz' diff --git a/acceptance/bundle/ai_runtime_task/local_code_source/output.txt b/acceptance/bundle/ai_runtime_task/local_code_source/output.txt deleted file mode 100644 index be8c3ec21ab..00000000000 --- a/acceptance/bundle/ai_runtime_task/local_code_source/output.txt +++ /dev/null @@ -1,14 +0,0 @@ - ->>> [CLI] bundle deploy -Uploading bundle files to /Workspace/Users/[USERNAME]/.bundle/ai-runtime-test/default/files... -Deploying resources... -Updating deployment state... -Deployment complete! - -=== Expecting the code snapshot tarball to be uploaded to artifact_path/.internal ->>> jq .path -"/api/2.0/workspace-files/import-file/Workspace/Users/[USERNAME]/.bundle/ai-runtime-test/default/artifacts/.internal/src_[HASH].tar.gz" - -=== Expecting the created job's code_source_path to point at the uploaded archive ->>> jq -s .[] | select(.path=="/api/2.2/jobs/create") | .body.tasks[0].ai_runtime_task.code_source_path out.requests.txt -"/Workspace/Users/[USERNAME]/.bundle/ai-runtime-test/default/artifacts/.internal/src_[HASH].tar.gz" diff --git a/acceptance/bundle/ai_runtime_task/local_code_source/script b/acceptance/bundle/ai_runtime_task/local_code_source/script deleted file mode 100644 index 8701b74e6c9..00000000000 --- a/acceptance/bundle/ai_runtime_task/local_code_source/script +++ /dev/null @@ -1,9 +0,0 @@ -trace $CLI bundle deploy - -title "Expecting the code snapshot tarball to be uploaded to artifact_path/.internal" -trace jq .path < out.requests.txt | grep import-file | grep 'tar.gz' | sort - -title "Expecting the created job's code_source_path to point at the uploaded archive" -trace jq -s '.[] | select(.path=="/api/2.2/jobs/create") | .body.tasks[0].ai_runtime_task.code_source_path' out.requests.txt - -rm out.requests.txt diff --git a/acceptance/bundle/ai_runtime_task/local_code_source/test.toml b/acceptance/bundle/ai_runtime_task/local_code_source/test.toml deleted file mode 100644 index 1c37cecae7d..00000000000 --- a/acceptance/bundle/ai_runtime_task/local_code_source/test.toml +++ /dev/null @@ -1,11 +0,0 @@ -RecordRequests = true - -Ignore = [ - '.databricks', -] - -# The code archive filename embeds a content hash. Normalize it so output is -# stable across machines while still asserting the _ shape. -[[Repls]] -Old = 'src_[0-9a-f]{16}\.tar\.gz' -New = 'src_[HASH].tar.gz' diff --git a/acceptance/bundle/ai_runtime_task/plain_tar/databricks.yml b/acceptance/bundle/ai_runtime_task/plain_tar/databricks.yml new file mode 100644 index 00000000000..c6d37475f73 --- /dev/null +++ b/acceptance/bundle/ai_runtime_task/plain_tar/databricks.yml @@ -0,0 +1,24 @@ +bundle: + name: ai-runtime-test + +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 diff --git a/acceptance/bundle/ai_runtime_task/plain_tar/out.test.toml b/acceptance/bundle/ai_runtime_task/plain_tar/out.test.toml new file mode 100644 index 00000000000..f784a183258 --- /dev/null +++ b/acceptance/bundle/ai_runtime_task/plain_tar/out.test.toml @@ -0,0 +1,3 @@ +Local = true +Cloud = false +EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["terraform", "direct"] diff --git a/acceptance/bundle/ai_runtime_task/plain_tar/output.txt b/acceptance/bundle/ai_runtime_task/plain_tar/output.txt new file mode 100644 index 00000000000..79d4ab65e54 --- /dev/null +++ b/acceptance/bundle/ai_runtime_task/plain_tar/output.txt @@ -0,0 +1,14 @@ + +=== deploy packages the working tree as a plain tarball + +>>> [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 cache +/api/2.0/workspace-files/import-file/Workspace/Users/[USERNAME]/.air/repo_snapshots/src/src_[SNAPSHOT].tar.gz + +=== the created job's code_source_path points at the uploaded archive +/Workspace/Users/[USERNAME]/.air/repo_snapshots/src/src_[SNAPSHOT].tar.gz diff --git a/acceptance/bundle/ai_runtime_task/plain_tar/script b/acceptance/bundle/ai_runtime_task/plain_tar/script new file mode 100644 index 00000000000..9978eb9ba22 --- /dev/null +++ b/acceptance/bundle/ai_runtime_task/plain_tar/script @@ -0,0 +1,13 @@ +# The bundle dir is not a git repo, so the snapshot falls back to plain_tar of the +# working tree (no cache; timestamp-named archive). + +title "deploy packages the working tree as a plain tarball\n" +trace $CLI bundle deploy + +title "the snapshot tarball is uploaded to the user's repo_snapshots cache\n" +jq -r 'select(.path | test("/import-file/.*/.air/repo_snapshots/src/src_")) | .path' out.requests.txt | sort -u + +title "the created job's code_source_path points at the uploaded archive\n" +jq -r 'select(.path == "/api/2.2/jobs/create") | .body.tasks[0].ai_runtime_task.code_source_path' out.requests.txt + +rm out.requests.txt diff --git a/acceptance/bundle/ai_runtime_task/plain_tar/src/command.sh b/acceptance/bundle/ai_runtime_task/plain_tar/src/command.sh new file mode 100644 index 00000000000..7ce3949767c --- /dev/null +++ b/acceptance/bundle/ai_runtime_task/plain_tar/src/command.sh @@ -0,0 +1,2 @@ +cd $CODE_SOURCE_PATH +python train.py diff --git a/acceptance/experimental/air/init/train_project/src/train.py b/acceptance/bundle/ai_runtime_task/plain_tar/src/train.py similarity index 100% rename from acceptance/experimental/air/init/train_project/src/train.py rename to acceptance/bundle/ai_runtime_task/plain_tar/src/train.py diff --git a/acceptance/bundle/ai_runtime_task/plain_tar/test.toml b/acceptance/bundle/ai_runtime_task/plain_tar/test.toml new file mode 100644 index 00000000000..8ef47fe2dc3 --- /dev/null +++ b/acceptance/bundle/ai_runtime_task/plain_tar/test.toml @@ -0,0 +1,11 @@ +RecordRequests = true + +Ignore = [ + '.databricks', +] + +# A non-git code_source uses plain_tar, whose archive is timestamp-named +# (_YYYYMMDD_HHMMSS.tar.gz). Collapse it to a stable token. +[[Repls]] +Old = 'src_\d{8}_\d{6}\.tar\.gz' +New = 'src_[SNAPSHOT].tar.gz' diff --git a/acceptance/experimental/air/.gitkeep b/acceptance/experimental/air/.gitkeep new file mode 100644 index 00000000000..e69de29bb2d diff --git a/acceptance/experimental/air/init/out.test.toml b/acceptance/experimental/air/init/out.test.toml deleted file mode 100644 index d6187dcb046..00000000000 --- a/acceptance/experimental/air/init/out.test.toml +++ /dev/null @@ -1,3 +0,0 @@ -Local = true -Cloud = false -EnvMatrix.DATABRICKS_BUNDLE_ENGINE = [] diff --git a/acceptance/experimental/air/init/output.txt b/acceptance/experimental/air/init/output.txt deleted file mode 100644 index f08b4a289a3..00000000000 --- a/acceptance/experimental/air/init/output.txt +++ /dev/null @@ -1,43 +0,0 @@ - -=== migrate a train.yaml into a bundle ->>> [CLI] experimental air init --from train_project/train.yaml --output-dir generated -Warning: "secrets" has no ai_runtime_task equivalent and was dropped; configure secrets on the job or via {{secrets/scope/key}} references -Bundle written to generated -Deploy it with: databricks bundle deploy - -=== generated databricks.yml ->>> cat generated/databricks.yml -bundle: - name: my_training -include: - - resources/*.yml - -=== generated job resource ->>> cat generated/resources/my_training.job.yml -resources: - jobs: - my_training: - name: my-training - tasks: - - task_key: my-training - ai_runtime_task: - code_source_path: ./src - deployments: - - command_path: src/command.sh - compute: - accelerator_count: 8 - accelerator_type: GPU_8xH100 - experiment: my-training - environment_key: default - max_retries: 2 - environments: - - environment_key: default - spec: - dependencies: - - torch>=2.0.0 - environment_version: "5" - -=== command materialized next to the code source ->>> cat generated/src/command.sh -cd $CODE_SOURCE_PATH -python train.py diff --git a/acceptance/experimental/air/init/script b/acceptance/experimental/air/init/script deleted file mode 100644 index 56c24aca87f..00000000000 --- a/acceptance/experimental/air/init/script +++ /dev/null @@ -1,11 +0,0 @@ -title "migrate a train.yaml into a bundle" -trace $CLI experimental air init --from train_project/train.yaml --output-dir generated - -title "generated databricks.yml" -trace cat generated/databricks.yml - -title "generated job resource" -trace cat generated/resources/my_training.job.yml - -title "command materialized next to the code source" -trace cat generated/src/command.sh diff --git a/acceptance/experimental/air/init/test.toml b/acceptance/experimental/air/init/test.toml deleted file mode 100644 index 6270a5aca70..00000000000 --- a/acceptance/experimental/air/init/test.toml +++ /dev/null @@ -1,8 +0,0 @@ -# This command generates local files from a train.yaml; no bundle engine or -# workspace is needed. -Ignore = [ - 'generated', -] - -[EnvMatrix] -DATABRICKS_BUNDLE_ENGINE = [] diff --git a/acceptance/experimental/air/init/train_project/train.yaml b/acceptance/experimental/air/init/train_project/train.yaml deleted file mode 100644 index 3e3ca0f5f7c..00000000000 --- a/acceptance/experimental/air/init/train_project/train.yaml +++ /dev/null @@ -1,18 +0,0 @@ -experiment_name: my-training -compute: - num_accelerators: 8 - accelerator_type: GPU_8xH100 -environment: - version: "5" - dependencies: - - torch>=2.0.0 -code_source: - type: snapshot - snapshot: - root_path: ./src -command: | - cd $CODE_SOURCE_PATH - python train.py -max_retries: 2 -secrets: - HF_TOKEN: my_scope/hf_token diff --git a/acceptance/experimental/air/init_errors/docker/train.yaml b/acceptance/experimental/air/init_errors/docker/train.yaml deleted file mode 100644 index df1e22c2056..00000000000 --- a/acceptance/experimental/air/init_errors/docker/train.yaml +++ /dev/null @@ -1,12 +0,0 @@ -experiment_name: docker-run -compute: - num_accelerators: 1 - accelerator_type: GPU_1xA10 -environment: - docker_image: - url: org/repo:tag -code_source: - type: snapshot - snapshot: - root_path: ./src -command: python train.py diff --git a/acceptance/experimental/air/init_errors/legacy/train.yaml b/acceptance/experimental/air/init_errors/legacy/train.yaml deleted file mode 100644 index f4d726a4476..00000000000 --- a/acceptance/experimental/air/init_errors/legacy/train.yaml +++ /dev/null @@ -1,9 +0,0 @@ -experiment_name: legacy-run -compute: - num_accelerators: 1 - accelerator_type: a10 -code_source: - type: snapshot - snapshot: - root_path: ./src -command: python train.py diff --git a/acceptance/experimental/air/init_errors/out.test.toml b/acceptance/experimental/air/init_errors/out.test.toml deleted file mode 100644 index d6187dcb046..00000000000 --- a/acceptance/experimental/air/init_errors/out.test.toml +++ /dev/null @@ -1,3 +0,0 @@ -Local = true -Cloud = false -EnvMatrix.DATABRICKS_BUNDLE_ENGINE = [] diff --git a/acceptance/experimental/air/init_errors/output.txt b/acceptance/experimental/air/init_errors/output.txt deleted file mode 100644 index 9a400dc7563..00000000000 --- a/acceptance/experimental/air/init_errors/output.txt +++ /dev/null @@ -1,8 +0,0 @@ - -=== legacy accelerator type is rejected ->>> [CLI] experimental air init --from legacy/train.yaml --output-dir legacy-out -Error: compute.accelerator_type "a10" is not supported by ai_runtime_task (expected GPU_1xA10, GPU_1xH100, or GPU_8xH100) - -=== docker image is rejected ->>> [CLI] experimental air init --from docker/train.yaml --output-dir docker-out -Error: environment.docker_image is not supported by ai_runtime_task; use environment.version and environment.dependencies instead diff --git a/acceptance/experimental/air/init_errors/script b/acceptance/experimental/air/init_errors/script deleted file mode 100644 index 43fc02501bf..00000000000 --- a/acceptance/experimental/air/init_errors/script +++ /dev/null @@ -1,5 +0,0 @@ -title "legacy accelerator type is rejected" -musterr trace $CLI experimental air init --from legacy/train.yaml --output-dir legacy-out - -title "docker image is rejected" -musterr trace $CLI experimental air init --from docker/train.yaml --output-dir docker-out diff --git a/acceptance/experimental/air/init_errors/test.toml b/acceptance/experimental/air/init_errors/test.toml deleted file mode 100644 index afca356e837..00000000000 --- a/acceptance/experimental/air/init_errors/test.toml +++ /dev/null @@ -1,3 +0,0 @@ -# These invocations must fail; no bundle engine or workspace is needed. -[EnvMatrix] -DATABRICKS_BUNDLE_ENGINE = [] diff --git a/bundle/config/mutator/aicode/package_upload.go b/bundle/config/mutator/aicode/package_upload.go index 234514763c1..662c2e1c004 100644 --- a/bundle/config/mutator/aicode/package_upload.go +++ b/bundle/config/mutator/aicode/package_upload.go @@ -5,14 +5,18 @@ // volume path to an uploaded code archive; its doc comment states that the CLI // is responsible for packaging the user's local code directory into that // archive. This mutator implements that contract for DABs: when a user points -// code_source_path at a local directory, it tarballs the directory (gitignore -// aware, .git excluded), uploads the archive next to bundle libraries, and -// rewrites the field to the resulting remote path so the deployed job runs -// against the uploaded code. Values that are already remote are left untouched. +// code_source_path at a local directory, it snapshots the directory (git archive +// of HEAD when the tree is a clean git repo, otherwise a plain tar; .git and +// gitignored files excluded), uploads the archive to the user's workspace code +// snapshot cache, and rewrites the field to the resulting remote path so the +// deployed job runs against the uploaded code. Values that are already remote are +// left untouched. +// +// The snapshot approach mirrors PR #5897, which ported it from the Python air CLI +// (see snapshot_resolve.go / snapshot_package.go / snapshot_cachekey.go). package aicode import ( - "bytes" "context" "errors" "fmt" @@ -20,6 +24,7 @@ import ( "os" "path" "path/filepath" + "time" "github.com/databricks/cli/bundle" "github.com/databricks/cli/bundle/libraries" @@ -27,7 +32,6 @@ import ( "github.com/databricks/cli/libs/dyn" "github.com/databricks/cli/libs/filer" "github.com/databricks/cli/libs/log" - "github.com/databricks/cli/libs/vfs" ) // codeSourcePatterns are the config locations of an AI Runtime task's @@ -59,8 +63,9 @@ func PackageAndUpload() bundle.Mutator { } type packageAndUpload struct { - // client is the filer used for uploads. It defaults to the libraries filer - // and is only overridden in tests. + // client is the filer used for uploads. When nil (the normal case) a filer + // rooted at the code snapshot cache is built per code source. It is only set + // in tests, to inject a recording filer. client filer.Filer } @@ -77,13 +82,9 @@ func (m *packageAndUpload) Apply(ctx context.Context, b *bundle.Bundle) diag.Dia return diags } - client, uploadPath, filerDiags := libraries.GetFilerForLibraries(ctx, b) - diags = diags.Extend(filerDiags) - if diags.HasError() { - return diags - } - if m.client == nil { - m.client = client + userDir, err := userWorkspaceHome(b) + if err != nil { + return diags.Extend(diag.FromErr(err)) } stagingDir, err := b.LocalStateDir(ctx, "ai_code_source") @@ -96,7 +97,7 @@ func (m *packageAndUpload) Apply(ctx context.Context, b *bundle.Bundle) diag.Dia // are reported before any config is rewritten. remotePaths := make(map[string]string, len(sources)) for _, cs := range sources { - remote, err := m.packageOne(ctx, b, cs, stagingDir, uploadPath) + remote, err := m.packageOne(ctx, b, cs, stagingDir, userDir) if err != nil { diags = diags.Extend(diag.FromErr(err)) return diags @@ -121,44 +122,109 @@ func (m *packageAndUpload) Apply(ctx context.Context, b *bundle.Bundle) diag.Dia return diags } -// packageOne tarballs the local directory for a single code source, uploads it -// (skipping the upload if a content-identical archive already exists), and -// returns the remote path the config should point to. -func (m *packageAndUpload) packageOne(ctx context.Context, b *bundle.Bundle, cs codeSource, stagingDir, uploadPath string) (string, error) { +// repoSnapshotsSubdir is the per-user workspace location for cached code +// snapshots, under the user's home. It matches the Python air CLI (and PR #5897) +// so a snapshot is cached in a stable place that bundle deploy does not wipe. +// +// Deliberately NOT /.internal: artifacts.CleanUp() deletes that +// directory at the start of every deploy, which would defeat the git-archive +// snapshot cache. Keeping snapshots under ~/.air/repo_snapshots lets an unchanged +// commit skip re-upload across deploys. +const repoSnapshotsSubdir = ".air/repo_snapshots" + +// packageOne snapshots the local directory for a single code source, uploads the +// archive to the user's repo_snapshots dir (skipping the upload when an identical +// git-archive snapshot already exists), and returns the remote path the config +// should point to. +func (m *packageAndUpload) packageOne(ctx context.Context, b *bundle.Bundle, cs codeSource, stagingDir, userDir string) (string, error) { localDir := filepath.Join(b.SyncRootPath, filepath.FromSlash(cs.value)) dirName := filepath.Base(localDir) - var buf bytes.Buffer - sha, err := buildTarball(ctx, vfs.MustNew(localDir), dirName, &buf) + uploadPath := path.Join(userDir, repoSnapshotsSubdir, dirName) + client := m.client + if client == nil { + var err error + client, err = filer.NewWorkspaceFilesClient(b.WorkspaceClient(ctx), uploadPath) + if err != nil { + return "", err + } + } + + git := newGitRepo(localDir) + plan, err := resolveSnapshotPlan(ctx, git) if err != nil { - return "", fmt.Errorf("failed to package code_source_path %q: %w", cs.value, err) + return "", fmt.Errorf("failed to resolve code_source_path %q: %w", cs.value, err) } - // The SHA of the (reproducible) archive is embedded in the filename so an - // unchanged code directory resolves to the same remote path across deploys - // and the existence check below turns re-uploads into no-ops. - archiveName := fmt.Sprintf("%s_%s.tar.gz", dirName, sha[:16]) - remotePath := path.Join(uploadPath, archiveName) + // git_archive is cacheable by (commit, include_paths): a hit means the identical + // tarball is already uploaded, so packaging and upload are skipped entirely. + if plan.mode == modeGitArchive { + cacheKey := computeSnapshotCacheKey(plan.commitSHA, plan.includePaths) + archiveName := fmt.Sprintf("%s_%s.tar.gz", dirName, cacheKey[:16]) + remotePath := path.Join(uploadPath, archiveName) + + exists, err := archiveExists(ctx, client, archiveName) + if err != nil { + return "", err + } + if exists { + log.Debugf(ctx, "code snapshot for %s already present at %s, skipping upload", shortSHA(plan.commitSHA), remotePath) + return remotePath, nil + } - if _, err := m.client.Stat(ctx, archiveName); err == nil { - log.Debugf(ctx, "code snapshot %s already present, skipping upload", remotePath) + if err := packageAndUploadArchive(ctx, client, stagingDir, archiveName, func(out string) error { + return createGitArchiveSnapshot(ctx, git, plan.commitSHA, out, dirName, plan.includePaths) + }); err != nil { + return "", err + } return remotePath, nil - } else if !errors.Is(err, fs.ErrNotExist) { - return "", fmt.Errorf("failed to check for existing code snapshot %q: %w", remotePath, err) } - // Stage the archive on disk so a large snapshot is not held in memory during - // upload, matching how libraries are uploaded from a file. - localArchive := filepath.Join(stagingDir, archiveName) - if err := os.WriteFile(localArchive, buf.Bytes(), 0o600); err != nil { + // plain_tar is not cacheable (working-tree content isn't pinned to a SHA), so it + // is timestamp-named to avoid clobbering a concurrent deploy. + archiveName := fmt.Sprintf("%s_%s.tar.gz", dirName, time.Now().UTC().Format("20060102_150405")) + remotePath := path.Join(uploadPath, archiveName) + if err := packageAndUploadArchive(ctx, client, stagingDir, archiveName, func(out string) error { + return createPlainTarball(ctx, localDir, out, plan.includePaths) + }); err != nil { return "", err } + return remotePath, nil +} - if err := libraries.UploadFile(ctx, localArchive, m.client); err != nil { - return "", err +// userWorkspaceHome returns the current user's workspace home directory +// (/Workspace/Users/), the root under which code snapshots are cached. +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 +} + +// archiveExists reports whether archiveName already exists in the store, treating +// fs.ErrNotExist as "no". +func archiveExists(ctx context.Context, client filer.Filer, archiveName string) (bool, error) { + _, err := client.Stat(ctx, archiveName) + if err == nil { + return true, nil } + if errors.Is(err, fs.ErrNotExist) { + return false, nil + } + return false, fmt.Errorf("failed to check for existing code snapshot %q: %w", archiveName, err) +} - return remotePath, nil +// packageAndUploadArchive writes the tarball via pkg into a temp file under +// stagingDir, then uploads it as archiveName. The temp file is always removed. +func packageAndUploadArchive(ctx context.Context, client filer.Filer, stagingDir, archiveName string, pkg func(outputPath string) error) error { + localArchive := filepath.Join(stagingDir, archiveName) + defer os.Remove(localArchive) + + if err := pkg(localArchive); err != nil { + return err + } + return libraries.UploadFile(ctx, localArchive, client) } // collectLocalCodeSources returns every AI Runtime task code_source_path that diff --git a/bundle/config/mutator/aicode/package_upload_test.go b/bundle/config/mutator/aicode/package_upload_test.go index a61b01f2264..6b68d98f507 100644 --- a/bundle/config/mutator/aicode/package_upload_test.go +++ b/bundle/config/mutator/aicode/package_upload_test.go @@ -1,7 +1,6 @@ package aicode import ( - "bytes" "context" "io" "io/fs" @@ -15,14 +14,18 @@ import ( "github.com/databricks/cli/bundle/internal/bundletest" "github.com/databricks/cli/libs/dyn" "github.com/databricks/cli/libs/filer" - "github.com/databricks/cli/libs/vfs" + "github.com/databricks/databricks-sdk-go/service/iam" "github.com/databricks/databricks-sdk-go/service/jobs" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) -// recordingFiler records every Write and reports files as absent, so an upload -// always happens and its target name is captured. +// codeSnapshotDir is the workspace location code snapshots are uploaded to in +// tests, derived from the fixture user below. +const codeSnapshotDir = "/Workspace/Users/me@databricks.com/.air/repo_snapshots/src" + +// recordingFiler records every Write and reports files as present per exists, so a +// test can both capture uploads and simulate a cache hit. type recordingFiler struct { filer.Filer written map[string][]byte @@ -51,19 +54,18 @@ func (f *recordingFiler) Stat(ctx context.Context, p string) (fs.FileInfo, error type fakeFileInfo struct{ fs.FileInfo } -func bundleWithCodeSource(t *testing.T, codeSourcePath string) (*bundle.Bundle, string) { +// bundleWithCodeSource builds a bundle rooted at dir whose single AI Runtime task +// points at codeSourcePath. The caller populates dir (and optionally inits a git +// repo) before applying the mutator. +func bundleWithCodeSource(t *testing.T, dir, codeSourcePath string) *bundle.Bundle { t.Helper() - dir := t.TempDir() - require.NoError(t, os.MkdirAll(filepath.Join(dir, "src"), 0o755)) - require.NoError(t, os.WriteFile(filepath.Join(dir, "src", "train.py"), []byte("print('x')"), 0o644)) - b := &bundle.Bundle{ BundleRootPath: dir, SyncRootPath: dir, Config: config.Root{ Bundle: config.Bundle{Target: "default"}, Workspace: config.Workspace{ - ArtifactPath: "/Workspace/Users/me/.bundle/artifacts", + CurrentUser: &config.User{User: &iam.User{UserName: "me@databricks.com"}}, }, Resources: config.Resources{ Jobs: map[string]*resources.Job{ @@ -71,11 +73,8 @@ func bundleWithCodeSource(t *testing.T, codeSourcePath string) (*bundle.Bundle, JobSettings: jobs.JobSettings{ Tasks: []jobs.Task{ { - TaskKey: "train", - AiRuntimeTask: &jobs.AiRuntimeTask{ - Experiment: "exp", - CodeSourcePath: codeSourcePath, - }, + TaskKey: "train", + AiRuntimeTask: &jobs.AiRuntimeTask{Experiment: "exp", CodeSourcePath: codeSourcePath}, }, }, }, @@ -85,45 +84,62 @@ func bundleWithCodeSource(t *testing.T, codeSourcePath string) (*bundle.Bundle, }, } bundletest.SetLocation(b, ".", []dyn.Location{{File: filepath.Join(dir, "databricks.yml")}}) - return b, dir + return b +} + +// bundleWithPlainSrc creates a non-git src/ directory with one file and returns +// the bundle (so packaging falls back to plain_tar). +func bundleWithPlainSrc(t *testing.T, codeSourcePath string) *bundle.Bundle { + t.Helper() + dir := t.TempDir() + require.NoError(t, os.MkdirAll(filepath.Join(dir, "src"), 0o755)) + require.NoError(t, os.WriteFile(filepath.Join(dir, "src", "train.py"), []byte("print('x')"), 0o644)) + return bundleWithCodeSource(t, dir, codeSourcePath) } -func TestPackageAndUploadRewritesLocalCodeSource(t *testing.T) { - b, _ := bundleWithCodeSource(t, "src") +func codeSourcePath(b *bundle.Bundle) string { + return b.Config.Resources.Jobs["train"].Tasks[0].AiRuntimeTask.CodeSourcePath +} + +func TestPackageAndUploadPlainTarRewritesLocalCodeSource(t *testing.T) { + // Non-git src/ dir → plain_tar → timestamp-named archive. + b := bundleWithPlainSrc(t, "src") f := newRecordingFiler() diags := bundle.Apply(t.Context(), b, &packageAndUpload{client: f}) require.Empty(t, diags) - rewritten := b.Config.Resources.Jobs["train"].Tasks[0].AiRuntimeTask.CodeSourcePath - assert.Regexp(t, `^/Workspace/Users/me/\.bundle/artifacts/\.internal/src_[0-9a-f]{16}\.tar\.gz$`, rewritten) + assert.Regexp(t, `^`+codeSnapshotDir+`/src_\d{8}_\d{6}\.tar\.gz$`, codeSourcePath(b)) require.Len(t, f.written, 1) for name := range f.written { - assert.Regexp(t, `^src_[0-9a-f]{16}\.tar\.gz$`, name) + assert.Regexp(t, `^src_\d{8}_\d{6}\.tar\.gz$`, name) } } func TestPackageAndUploadSkipsRemoteCodeSource(t *testing.T) { - b, _ := bundleWithCodeSource(t, "/Volumes/main/default/code/existing.tar.gz") + b := bundleWithPlainSrc(t, "/Volumes/main/default/code/existing.tar.gz") f := newRecordingFiler() diags := bundle.Apply(t.Context(), b, &packageAndUpload{client: f}) require.Empty(t, diags) - assert.Equal(t, "/Volumes/main/default/code/existing.tar.gz", - b.Config.Resources.Jobs["train"].Tasks[0].AiRuntimeTask.CodeSourcePath) + assert.Equal(t, "/Volumes/main/default/code/existing.tar.gz", codeSourcePath(b)) assert.Empty(t, f.written, "remote code_source_path must not trigger an upload") } -func TestPackageAndUploadSkipsUploadWhenArchiveExists(t *testing.T) { - b, _ := bundleWithCodeSource(t, "src") +func TestPackageAndUploadGitArchiveCacheHitSkipsUpload(t *testing.T) { + // A clean git repo → git_archive → cache-key-named archive; when it already + // exists in the store, upload is skipped and the path is still rewritten. + dir := t.TempDir() + require.NoError(t, os.MkdirAll(filepath.Join(dir, "src"), 0o755)) + require.NoError(t, os.WriteFile(filepath.Join(dir, "src", "train.py"), []byte("print('x')"), 0o644)) + commit := initGitRepo(t, dir) + + b := bundleWithCodeSource(t, dir, "src") - // Determine the archive name the mutator will compute, then mark it present. - var buf bytes.Buffer - sha, err := buildTarball(t.Context(), vfs.MustNew(filepath.Join(b.SyncRootPath, "src")), "src", &buf) - require.NoError(t, err) - archiveName := "src_" + sha[:16] + ".tar.gz" + cacheKey := computeSnapshotCacheKey(commit, nil) + archiveName := "src_" + cacheKey[:16] + ".tar.gz" f := newRecordingFiler() f.exists[archiveName] = true @@ -131,7 +147,24 @@ func TestPackageAndUploadSkipsUploadWhenArchiveExists(t *testing.T) { diags := bundle.Apply(t.Context(), b, &packageAndUpload{client: f}) require.Empty(t, diags) - assert.Empty(t, f.written, "existing archive must not be re-uploaded") - assert.Equal(t, "/Workspace/Users/me/.bundle/artifacts/.internal/"+archiveName, - b.Config.Resources.Jobs["train"].Tasks[0].AiRuntimeTask.CodeSourcePath) + assert.Empty(t, f.written, "existing git-archive snapshot must not be re-uploaded") + assert.Equal(t, codeSnapshotDir+"/"+archiveName, codeSourcePath(b)) +} + +func TestPackageAndUploadGitArchiveUploadsWhenAbsent(t *testing.T) { + dir := t.TempDir() + require.NoError(t, os.MkdirAll(filepath.Join(dir, "src"), 0o755)) + require.NoError(t, os.WriteFile(filepath.Join(dir, "src", "train.py"), []byte("print('x')"), 0o644)) + commit := initGitRepo(t, dir) + + b := bundleWithCodeSource(t, dir, "src") + cacheKey := computeSnapshotCacheKey(commit, nil) + archiveName := "src_" + cacheKey[:16] + ".tar.gz" + + f := newRecordingFiler() + diags := bundle.Apply(t.Context(), b, &packageAndUpload{client: f}) + require.Empty(t, diags) + + require.Contains(t, f.written, archiveName, "git-archive snapshot must be uploaded when not cached") + assert.Equal(t, codeSnapshotDir+"/"+archiveName, codeSourcePath(b)) } diff --git a/bundle/config/mutator/aicode/snapshot_cachekey.go b/bundle/config/mutator/aicode/snapshot_cachekey.go new file mode 100644 index 00000000000..7512661beee --- /dev/null +++ b/bundle/config/mutator/aicode/snapshot_cachekey.go @@ -0,0 +1,32 @@ +package aicode + +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. Ported from the Python CLI's +// cli/utils/snapshot.py so cache keys are byte-identical across the two CLIs. +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[:]) +} diff --git a/bundle/config/mutator/aicode/snapshot_cachekey_test.go b/bundle/config/mutator/aicode/snapshot_cachekey_test.go new file mode 100644 index 00000000000..e89ffc01a6f --- /dev/null +++ b/bundle/config/mutator/aicode/snapshot_cachekey_test.go @@ -0,0 +1,33 @@ +package aicode + +import ( + "encoding/json" + "os" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// TestComputeSnapshotCacheKeyGolden checks parity with the Python air CLI: the +// keys in testdata/cache_keys.json are produced by the Python implementation, so +// matching them proves the Go port hashes identical material. +func TestComputeSnapshotCacheKeyGolden(t *testing.T) { + data, err := os.ReadFile("testdata/cache_keys.json") + require.NoError(t, err) + + var cases []struct { + Name string `json:"name"` + CommitSHA string `json:"commit_sha"` + IncludePaths []string `json:"include_paths"` + CacheKey string `json:"cache_key"` + } + 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)) + }) + } +} diff --git a/bundle/config/mutator/aicode/snapshot_git.go b/bundle/config/mutator/aicode/snapshot_git.go new file mode 100644 index 00000000000..94f370dfb67 --- /dev/null +++ b/bundle/config/mutator/aicode/snapshot_git.go @@ -0,0 +1,77 @@ +package aicode + +import ( + "bytes" + "context" + "fmt" + "os/exec" + "strings" +) + +// Local, no-network git introspection for the code_source snapshot, ported from +// the Python CLI's cli/utils/git_state.py and PR #5897's snapshot_git.go. Only the +// subset the DABs resolve path needs is kept: is-repo, HEAD SHA, and dirty check. +// The remote-fetch helpers and the git_state.json/git_diff.patch provenance +// sidecars are intentionally not ported — the SDK ai_runtime_task has nowhere to +// attach sidecar paths, and the snapshot archives the local copy only. + +// gitRepo runs git subcommands scoped to one repository via `git -C`. Arguments +// are passed as a slice, never a shell string, so path values can't inject. +type gitRepo struct { + path string +} + +func newGitRepo(path string) gitRepo { + return gitRepo{path: path} +} + +// run executes `git -C ` and returns trimmed stdout; a non-zero +// exit wraps stderr. +func (g gitRepo) run(ctx context.Context, args ...string) (string, error) { + full := append([]string{"-C", g.path}, args...) + cmd := exec.CommandContext(ctx, "git", full...) + + var stdout, stderr bytes.Buffer + cmd.Stdout = &stdout + cmd.Stderr = &stderr + if err := cmd.Run(); err != nil { + msg := strings.TrimSpace(stderr.String()) + if msg == "" { + return "", fmt.Errorf("git %s: %w", strings.Join(args, " "), err) + } + return "", fmt.Errorf("git %s: %w: %s", strings.Join(args, " "), err, msg) + } + return strings.TrimSpace(stdout.String()), nil +} + +// isRepository reports whether the path is inside a git work tree. Using +// `rev-parse --is-inside-work-tree` (not a .git lookup) means a subdirectory of a +// repo counts — the common case when code_source_path is a subfolder of a monorepo. +func (g gitRepo) isRepository(ctx context.Context) bool { + out, err := g.run(ctx, "rev-parse", "--is-inside-work-tree") + if err != nil { + return false + } + return out == "true" +} + +// headSHA returns the current HEAD commit SHA. +func (g gitRepo) headSHA(ctx context.Context) (string, error) { + return g.run(ctx, "rev-parse", "HEAD") +} + +// hasUncommittedChanges reports whether there are staged or unstaged changes under +// the repo subtree. The `-- .` pathspec scopes the check so a subfolder snapshot +// considers only changes that could land in it. +func (g gitRepo) hasUncommittedChanges(ctx context.Context) (bool, error) { + out, err := g.run(ctx, "status", "--porcelain", "--", ".") + if err != nil { + return false, err + } + return out != "", nil +} + +// shortSHA abbreviates a commit SHA to 8 chars for log/error messages. +func shortSHA(sha string) string { + return sha[:min(len(sha), 8)] +} diff --git a/bundle/config/mutator/aicode/snapshot_package.go b/bundle/config/mutator/aicode/snapshot_package.go new file mode 100644 index 00000000000..69e7fa18165 --- /dev/null +++ b/bundle/config/mutator/aicode/snapshot_package.go @@ -0,0 +1,134 @@ +package aicode + +import ( + "bytes" + "context" + "fmt" + "os" + "os/exec" + "path/filepath" + "strings" +) + +// Tar builders ported from PR #5897 (in turn from cli/utils/snapshot.py). Both +// shell out (git archive / tar) for parity and to reuse git's/tar's symlink, +// gitignore, and AppleDouble handling. The tarball's top-level dir name is +// load-bearing — the remote entry_script extracts to /databricks/code_source/ — +// 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 { + // Single git invocation writes the gzipped tar with the desired prefix; no + // extract/repack. `git -C ` scopes the archive to the code directory's + // subtree even when it is a subfolder of a larger repo. + args := []string{ + "archive", + "--format=tar.gz", + "--prefix=" + directoryName + "/", + "-o", outputTarball, + commitSHA, + } + args = append(args, includePaths...) + if _, err := git.run(ctx, args...); err != nil { + return fmt.Errorf("failed to create git archive: %w", err) + } + return nil +} + +// createPlainTarball writes a gzipped tar of repoPath's working tree to +// outputTarball via `tar`. The archive preserves repoPath's directory name as the +// top-level entry. When includePaths is set, only those paths (nested under the +// directory name) are archived. .git and macOS AppleDouble files are always +// excluded; a .gitignore at repoPath is honored. +func createPlainTarball(ctx context.Context, repoPath, outputTarball string, includePaths []string) error { + dirName := filepath.Base(repoPath) + parent := filepath.Dir(repoPath) + + args := []string{"-czf", outputTarball} + + // Exclude macOS AppleDouble files: they sort before the real top-level dir and + // hijack a remote `head -1` parse. No-op on Linux. + args = append(args, "--exclude=._*") + + // Never ship .git — it is large and unused on the remote. + args = append(args, "--exclude=.git") + + // Honor .gitignore if present. + gitignorePath := filepath.Join(repoPath, ".gitignore") + if patterns, err := parseGitignore(gitignorePath); err == nil { + for _, p := range patterns { + if strings.Contains(p, "/") { + // Anchor path-relative patterns to the archive root so they don't + // match identically-named paths in subdirectories. + args = append(args, "--exclude="+dirName+"/"+strings.TrimPrefix(p, "/")) + } else { + args = append(args, "--exclude="+p) + } + } + } + + // Archive from the parent so the directory name is preserved; with include_paths, + // prefix each so entries nest under it (matching git archive --prefix). + args = append(args, "-C", parent) + if len(includePaths) > 0 { + for _, p := range includePaths { + args = append(args, dirName+"/"+p) + } + } else { + args = append(args, dirName) + } + + cmd := exec.CommandContext(ctx, "tar", args...) + var stderr bytes.Buffer + cmd.Stderr = &stderr + if err := cmd.Run(); err != nil { + if msg := strings.TrimSpace(stderr.String()); msg != "" { + return fmt.Errorf("failed to create plain tarball: %w: %s", err, msg) + } + return fmt.Errorf("failed to create plain tarball: %w", err) + } + return nil +} + +// parseGitignore reads a .gitignore and returns tar --exclude patterns. It mirrors +// the Python CLI's lossy normalization so plain-tar snapshots exclude the same set: +// +// - comments (#…) and blank lines are skipped; +// - negation patterns (!…) are unsupported by tar --exclude and skipped; +// - a trailing "/" (directory marker) is stripped; +// - "**" is not a path-separator-agnostic wildcard in tar, so "**/foo" → "foo" +// and "foo/**" → "foo"; a mid-path "**" has no tar equivalent and is skipped. +// +// A missing file returns (nil, error); callers treat any error as "no patterns". +func parseGitignore(path string) ([]string, error) { + data, err := os.ReadFile(path) + if err != nil { + return nil, err + } + + var patterns []string + for raw := range strings.SplitSeq(string(data), "\n") { + line := strings.TrimRight(raw, " \t\r") + if line == "" || strings.HasPrefix(line, "#") { + continue + } + if strings.HasPrefix(line, "!") { + continue + } + line = strings.TrimRight(line, "/") + if strings.Contains(line, "**") { + switch { + case strings.HasPrefix(line, "**/"): + line = line[len("**/"):] + case strings.HasSuffix(line, "/**"): + line = line[:len(line)-len("/**")] + default: + continue + } + } + patterns = append(patterns, line) + } + return patterns, nil +} diff --git a/bundle/config/mutator/aicode/snapshot_package_test.go b/bundle/config/mutator/aicode/snapshot_package_test.go new file mode 100644 index 00000000000..56252a2e40e --- /dev/null +++ b/bundle/config/mutator/aicode/snapshot_package_test.go @@ -0,0 +1,83 @@ +package aicode + +import ( + "archive/tar" + "compress/gzip" + "io" + "os" + "path/filepath" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// tarEntryNames returns the sorted set of entry names in a gzipped tarball file. +func tarEntryNames(t *testing.T, tarballPath string) map[string]bool { + t.Helper() + f, err := os.Open(tarballPath) + require.NoError(t, err) + defer f.Close() + gzr, err := gzip.NewReader(f) + require.NoError(t, err) + tr := tar.NewReader(gzr) + names := map[string]bool{} + for { + hdr, err := tr.Next() + if err == io.EOF { + break + } + require.NoError(t, err) + names[hdr.Name] = true + } + return names +} + +func TestCreatePlainTarballExcludesGitAndGitignored(t *testing.T) { + dir := t.TempDir() + code := filepath.Join(dir, "mycode") + require.NoError(t, os.MkdirAll(filepath.Join(code, "pkg"), 0o755)) + require.NoError(t, os.MkdirAll(filepath.Join(code, ".git"), 0o755)) + write := func(rel, content string) { + require.NoError(t, os.WriteFile(filepath.Join(code, rel), []byte(content), 0o644)) + } + write("train.py", "print('train')") + write("pkg/util.py", "x = 1") + write(".gitignore", "ignored.txt\n*.log\n") + write("ignored.txt", "nope") + write("debug.log", "nope") + write(".git/config", "[core]") + write("._resource_fork", "apple double") + + out := filepath.Join(dir, "out.tar.gz") + require.NoError(t, createPlainTarball(t.Context(), code, out, nil)) + + names := tarEntryNames(t, out) + assert.True(t, names["mycode/train.py"]) + assert.True(t, names["mycode/pkg/util.py"]) + assert.True(t, names["mycode/.gitignore"]) + assert.False(t, names["mycode/ignored.txt"], "gitignored file must be excluded") + assert.False(t, names["mycode/debug.log"], "gitignored glob must be excluded") + assert.False(t, names["mycode/.git/config"], ".git must never be archived") + assert.False(t, names["mycode/._resource_fork"], "AppleDouble metadata must be excluded") +} + +func TestCreateGitArchiveSnapshotPrefixesAndScopesToSubdir(t *testing.T) { + // Repo with two top-level dirs; archiving the "src" subdir must include only + // src, prefixed by its basename (the runtime extracts to code_source/). + repo := t.TempDir() + require.NoError(t, os.MkdirAll(filepath.Join(repo, "src"), 0o755)) + require.NoError(t, os.MkdirAll(filepath.Join(repo, "other"), 0o755)) + require.NoError(t, os.WriteFile(filepath.Join(repo, "src", "a.py"), []byte("a"), 0o644)) + require.NoError(t, os.WriteFile(filepath.Join(repo, "other", "b.py"), []byte("b"), 0o644)) + commit := initGitRepo(t, repo) + + src := filepath.Join(repo, "src") + out := filepath.Join(t.TempDir(), "out.tar.gz") + require.NoError(t, createGitArchiveSnapshot(t.Context(), newGitRepo(src), commit, out, "src", nil)) + + names := tarEntryNames(t, out) + assert.True(t, names["src/a.py"]) + assert.False(t, names["other/b.py"], "sibling dir must not be in a src-scoped archive") + assert.False(t, names["src/../other/b.py"]) +} diff --git a/bundle/config/mutator/aicode/snapshot_resolve.go b/bundle/config/mutator/aicode/snapshot_resolve.go new file mode 100644 index 00000000000..b58934f964c --- /dev/null +++ b/bundle/config/mutator/aicode/snapshot_resolve.go @@ -0,0 +1,56 @@ +package aicode + +import "context" + +// snapshotMode is how the snapshot tarball is produced. +type snapshotMode int + +const ( + // modeGitArchive packages a pinned commit via `git archive`. The commit is + // deterministic, so the tarball is cacheable by (commit, include_paths). + modeGitArchive snapshotMode = iota + // modePlainTar packages the working tree (including uncommitted changes) via + // `tar`. Not cacheable — working-tree content isn't pinned to a SHA. + modePlainTar +) + +// snapshotPlan is the outcome of resolving how to package a snapshot: the mode and +// the commit SHA to archive (git_archive only; empty for plain_tar). includePaths is +// carried through to packaging and the cache key. +// +// The plan mirrors PR #5897's richer type so this package can later grow git-ref and +// include_paths selection once the ai_runtime_task schema can carry them; for now the +// DABs code_source_path is a plain directory, so the decision reduces to +// "committed git repo → archive HEAD, otherwise tar the working tree". +type snapshotPlan struct { + mode snapshotMode + commitSHA string + includePaths []string +} + +// resolveSnapshotPlan decides how to package a local code_source directory: +// - a git work tree with no uncommitted changes → git_archive of HEAD (cacheable); +// - otherwise (non-git dir, or a dirty tree) → plain_tar of the working tree. +// +// Archiving HEAD when clean lets unchanged code hit the upload cache across deploys; +// falling back to plain_tar when dirty ensures local edits are actually deployed +// rather than silently dropped. +func resolveSnapshotPlan(ctx context.Context, git gitRepo) (snapshotPlan, error) { + if !git.isRepository(ctx) { + return snapshotPlan{mode: modePlainTar}, nil + } + + dirty, err := git.hasUncommittedChanges(ctx) + if err != nil { + return snapshotPlan{}, err + } + if dirty { + return snapshotPlan{mode: modePlainTar}, nil + } + + sha, err := git.headSHA(ctx) + if err != nil { + return snapshotPlan{}, err + } + return snapshotPlan{mode: modeGitArchive, commitSHA: sha}, nil +} diff --git a/bundle/config/mutator/aicode/snapshot_test.go b/bundle/config/mutator/aicode/snapshot_test.go new file mode 100644 index 00000000000..4b6f0fc9135 --- /dev/null +++ b/bundle/config/mutator/aicode/snapshot_test.go @@ -0,0 +1,72 @@ +package aicode + +import ( + "os" + "os/exec" + "path/filepath" + "testing" + + "github.com/stretchr/testify/require" +) + +// initGitRepo initializes a deterministic git repo at dir, commits everything in +// it, and returns the HEAD commit SHA. Author/committer dates are pinned so the +// resulting SHA (and thus any git-archive cache key) is stable across runs. +func initGitRepo(t *testing.T, dir string) string { + t.Helper() + run := func(env []string, args ...string) string { + cmd := exec.Command("git", append([]string{"-C", dir}, args...)...) + cmd.Env = append(os.Environ(), env...) + out, err := cmd.CombinedOutput() + require.NoErrorf(t, err, "git %v: %s", args, out) + return string(out) + } + dates := []string{ + "GIT_AUTHOR_DATE=2020-01-01T00:00:00Z", + "GIT_COMMITTER_DATE=2020-01-01T00:00:00Z", + "GIT_AUTHOR_NAME=Tester", "GIT_AUTHOR_EMAIL=tester@databricks.com", + "GIT_COMMITTER_NAME=Tester", "GIT_COMMITTER_EMAIL=tester@databricks.com", + } + run(nil, "init", "-qb", "main") + run(nil, "config", "core.hooksPath", "no-hooks") + run(nil, "add", "-A") + run(dates, "commit", "-qm", "init") + + cmd := exec.Command("git", "-C", dir, "rev-parse", "HEAD") + out, err := cmd.Output() + require.NoError(t, err) + return string(out[:40]) +} + +func TestResolveSnapshotPlanCleanGitRepoUsesGitArchive(t *testing.T) { + dir := t.TempDir() + require.NoError(t, os.WriteFile(filepath.Join(dir, "a.py"), []byte("x"), 0o644)) + commit := initGitRepo(t, dir) + + plan, err := resolveSnapshotPlan(t.Context(), newGitRepo(dir)) + require.NoError(t, err) + require.Equal(t, modeGitArchive, plan.mode) + require.Equal(t, commit, plan.commitSHA) +} + +func TestResolveSnapshotPlanDirtyGitRepoUsesPlainTar(t *testing.T) { + dir := t.TempDir() + require.NoError(t, os.WriteFile(filepath.Join(dir, "a.py"), []byte("x"), 0o644)) + initGitRepo(t, dir) + // Introduce an uncommitted change. + require.NoError(t, os.WriteFile(filepath.Join(dir, "a.py"), []byte("changed"), 0o644)) + + plan, err := resolveSnapshotPlan(t.Context(), newGitRepo(dir)) + require.NoError(t, err) + require.Equal(t, modePlainTar, plan.mode) + require.Empty(t, plan.commitSHA) +} + +func TestResolveSnapshotPlanNonGitUsesPlainTar(t *testing.T) { + dir := t.TempDir() + require.NoError(t, os.WriteFile(filepath.Join(dir, "a.py"), []byte("x"), 0o644)) + + plan, err := resolveSnapshotPlan(t.Context(), newGitRepo(dir)) + require.NoError(t, err) + require.Equal(t, modePlainTar, plan.mode) +} diff --git a/bundle/config/mutator/aicode/tarball.go b/bundle/config/mutator/aicode/tarball.go deleted file mode 100644 index d09a9a8afde..00000000000 --- a/bundle/config/mutator/aicode/tarball.go +++ /dev/null @@ -1,122 +0,0 @@ -package aicode - -import ( - "archive/tar" - "compress/gzip" - "context" - "crypto/sha256" - "encoding/hex" - "fmt" - "io" - "path" - "slices" - "strings" - "time" - - "github.com/databricks/cli/libs/fileset" - "github.com/databricks/cli/libs/git" - "github.com/databricks/cli/libs/vfs" -) - -// tarEpoch is a fixed modification time stamped on every tar entry so the -// archive is content-addressed: identical file contents always produce an -// identical archive (and therefore an identical SHA-256) regardless of the -// files' mtimes or when the archive was built. This is what makes the cache -// key in packageAndUpload stable across deploys and keeps acceptance output -// deterministic. -var tarEpoch = time.Date(2000, 1, 1, 0, 0, 0, 0, time.UTC) - -// appleDoublePrefix is the basename prefix of macOS AppleDouble metadata files. -// The AIR CLI excludes these (they sort before the real top-level entry and add -// useless per-file metadata); we match that behavior so archives built on macOS -// and Linux are identical. -const appleDoublePrefix = "._" - -// buildTarball writes a reproducible gzipped tarball of the directory rooted at -// codeRoot to out and returns its SHA-256 hex digest. Every entry is prefixed -// with prefix (the code directory's basename) so the archive expands to -// /... on extraction, matching the AIR CLI's `--prefix=/` layout -// and the runtime's /databricks/code_source/ contract. -// -// The file list is gitignore-aware: it honors the code directory's .gitignore -// and always excludes .git and .databricks (via [git.NewFileSetAtRoot]), the -// same walker that backs bundle file sync. -func buildTarball(ctx context.Context, codeRoot vfs.Path, prefix string, out io.Writer) (string, error) { - fsys, err := git.NewFileSetAtRoot(ctx, codeRoot) - if err != nil { - return "", err - } - - fileList, err := fsys.Files() - if err != nil { - return "", err - } - - // Sort by relative path so the archive byte stream (and thus its hash) does - // not depend on filesystem walk order. - slices.SortFunc(fileList, func(a, b fileset.File) int { - return strings.Compare(a.Relative, b.Relative) - }) - - hash := sha256.New() - gzw := gzip.NewWriter(io.MultiWriter(out, hash)) - tw := tar.NewWriter(gzw) - - for _, f := range fileList { - if err := addFileToTar(tw, codeRoot, f, prefix); err != nil { - return "", err - } - } - - if err := tw.Close(); err != nil { - return "", err - } - if err := gzw.Close(); err != nil { - return "", err - } - - return hex.EncodeToString(hash.Sum(nil)), nil -} - -func addFileToTar(tw *tar.Writer, codeRoot vfs.Path, f fileset.File, prefix string) error { - // fileset.File.Relative is already slash-separated, so path.Base is correct. - if strings.HasPrefix(path.Base(f.Relative), appleDoublePrefix) { - return nil - } - - rc, err := codeRoot.Open(f.Relative) - if err != nil { - return fmt.Errorf("open %s: %w", f.Relative, err) - } - defer rc.Close() - - info, err := rc.Stat() - if err != nil { - return fmt.Errorf("stat %s: %w", f.Relative, err) - } - - // Only regular files are archived. The gitignore-aware walker never yields - // directories, and symlinks inside a code snapshot are out of scope for v1. - if !info.Mode().IsRegular() { - return nil - } - - hdr := &tar.Header{ - Typeflag: tar.TypeReg, - Name: path.Join(prefix, f.Relative), - Size: info.Size(), - // Normalize permissions to a fixed 0644 and zero the mtime so the archive - // is reproducible across machines. Executable bits are intentionally - // dropped; the runtime invokes code via an interpreter, not by executing - // files from the snapshot directly. - Mode: 0o644, - ModTime: tarEpoch, - } - if err := tw.WriteHeader(hdr); err != nil { - return fmt.Errorf("tar header for %s: %w", f.Relative, err) - } - if _, err := io.Copy(tw, rc); err != nil { - return fmt.Errorf("write %s: %w", f.Relative, err) - } - return nil -} diff --git a/bundle/config/mutator/aicode/tarball_test.go b/bundle/config/mutator/aicode/tarball_test.go deleted file mode 100644 index 1ac3c9d57a8..00000000000 --- a/bundle/config/mutator/aicode/tarball_test.go +++ /dev/null @@ -1,108 +0,0 @@ -package aicode - -import ( - "archive/tar" - "bytes" - "compress/gzip" - "io" - "os" - "path/filepath" - "testing" - - "github.com/databricks/cli/libs/vfs" - "github.com/stretchr/testify/assert" - "github.com/stretchr/testify/require" -) - -// writeCodeDir materializes files (relative path -> content) under a fresh temp -// directory and returns a vfs.Path rooted at it. -func writeCodeDir(t *testing.T, files map[string]string) vfs.Path { - t.Helper() - dir := t.TempDir() - for name, content := range files { - p := filepath.Join(dir, filepath.FromSlash(name)) - require.NoError(t, os.MkdirAll(filepath.Dir(p), 0o755)) - require.NoError(t, os.WriteFile(p, []byte(content), 0o644)) - } - return vfs.MustNew(dir) -} - -// tarEntries reads a gzipped tarball and returns entry name -> content. -func tarEntries(t *testing.T, b []byte) map[string]string { - t.Helper() - gzr, err := gzip.NewReader(bytes.NewReader(b)) - require.NoError(t, err) - tr := tar.NewReader(gzr) - out := map[string]string{} - for { - hdr, err := tr.Next() - if err == io.EOF { - break - } - require.NoError(t, err) - content, err := io.ReadAll(tr) - require.NoError(t, err) - out[hdr.Name] = string(content) - } - return out -} - -func TestBuildTarballPrefixesEntriesAndOmitsGitAndGitignored(t *testing.T) { - root := writeCodeDir(t, map[string]string{ - "train.py": "print('train')", - "pkg/util.py": "x = 1", - ".gitignore": "ignored.txt\n*.log\n", - "ignored.txt": "should not be archived", - "debug.log": "should not be archived", - ".git/config": "[core]", - "._resource_fork": "apple double metadata", - }) - - var buf bytes.Buffer - sha, err := buildTarball(t.Context(), root, "mycode", &buf) - require.NoError(t, err) - require.NotEmpty(t, sha) - - entries := tarEntries(t, buf.Bytes()) - assert.Equal(t, "print('train')", entries["mycode/train.py"]) - assert.Equal(t, "x = 1", entries["mycode/pkg/util.py"]) - // .gitignore itself is a tracked file and is archived. - assert.Contains(t, entries, "mycode/.gitignore") - - assert.NotContains(t, entries, "mycode/ignored.txt", "gitignored file must be excluded") - assert.NotContains(t, entries, "mycode/debug.log", "gitignored glob must be excluded") - assert.NotContains(t, entries, "mycode/.git/config", ".git must never be archived") - assert.NotContains(t, entries, "mycode/._resource_fork", "AppleDouble metadata must be excluded") -} - -func TestBuildTarballIsDeterministic(t *testing.T) { - files := map[string]string{ - "a.py": "aaa", - "sub/b.py": "bbb", - "sub/c.txt": "ccc", - } - root1 := writeCodeDir(t, files) - root2 := writeCodeDir(t, files) - - var buf1, buf2 bytes.Buffer - sha1, err := buildTarball(t.Context(), root1, "code", &buf1) - require.NoError(t, err) - sha2, err := buildTarball(t.Context(), root2, "code", &buf2) - require.NoError(t, err) - - assert.Equal(t, sha1, sha2, "identical content must produce an identical hash") - assert.Equal(t, buf1.Bytes(), buf2.Bytes(), "identical content must produce identical bytes") -} - -func TestBuildTarballHashChangesWithContent(t *testing.T) { - root1 := writeCodeDir(t, map[string]string{"main.py": "v1"}) - root2 := writeCodeDir(t, map[string]string{"main.py": "v2"}) - - var buf1, buf2 bytes.Buffer - sha1, err := buildTarball(t.Context(), root1, "code", &buf1) - require.NoError(t, err) - sha2, err := buildTarball(t.Context(), root2, "code", &buf2) - require.NoError(t, err) - - assert.NotEqual(t, sha1, sha2) -} diff --git a/bundle/config/mutator/aicode/testdata/cache_keys.json b/bundle/config/mutator/aicode/testdata/cache_keys.json new file mode 100644 index 00000000000..06673499109 --- /dev/null +++ b/bundle/config/mutator/aicode/testdata/cache_keys.json @@ -0,0 +1,77 @@ +[ + { + "name": "commit_no_paths", + "commit_sha": "a3492b801c0ffee00000000000000000000dead", + "include_paths": null, + "cache_key": "8d7bc445ac83dfe432a353718ae8b1ae40eb00c860352662e6df96b7fdf862a5" + }, + { + "name": "commit_empty_paths_none", + "commit_sha": "0000000000000000000000000000000000000000", + "include_paths": null, + "cache_key": "a07e0fa4d38f3c2d08c3ff6bbff0158cea01796aaaf5368b0d0e247cebd27c80" + }, + { + "name": "single_path", + "commit_sha": "a3492b801c0ffee00000000000000000000dead", + "include_paths": [ + "research" + ], + "cache_key": "97bd562ad591af23b9a05651d88c01abbbcc24632e9a8f2534230a61492b9e0a" + }, + { + "name": "multi_path_sorted", + "commit_sha": "a3492b801c0ffee00000000000000000000dead", + "include_paths": [ + "a", + "b", + "c" + ], + "cache_key": "17d021577e3615e9a603a2f701f918d74ad16218efc4b2950ee20f6a18b1a174" + }, + { + "name": "multi_path_unsorted", + "commit_sha": "a3492b801c0ffee00000000000000000000dead", + "include_paths": [ + "c", + "a", + "b" + ], + "cache_key": "17d021577e3615e9a603a2f701f918d74ad16218efc4b2950ee20f6a18b1a174" + }, + { + "name": "path_whitespace", + "commit_sha": "a3492b801c0ffee00000000000000000000dead", + "include_paths": [ + " research ", + " data " + ], + "cache_key": "c6a6fa1d866ac90a5f371610ee08efa8aeb061a43b8a4acae30765db21f4155b" + }, + { + "name": "nested_paths", + "commit_sha": "deadbeefdeadbeefdeadbeefdeadbeefdeadbeef", + "include_paths": [ + "src/models", + "src/data", + "README.md" + ], + "cache_key": "2b18452b1019186a8e5b34af540a0545f96f24bc61771f58fa71fc218cb0275c" + }, + { + "name": "short_sha", + "commit_sha": "a3492b8", + "include_paths": null, + "cache_key": "16be4d7d83ecd01b8d3d274dafb07b3bccc72cf7f7299578e634cd01e1825345" + }, + { + "name": "dup_paths", + "commit_sha": "a3492b801c0ffee00000000000000000000dead", + "include_paths": [ + "x", + "x", + "y" + ], + "cache_key": "407a5d8c3a0e9438d5abf653c4d3a43928c300b5630e2f83f435661b9fbd599e" + } +] diff --git a/cmd/experimental/experimental.go b/cmd/experimental/experimental.go index d87c893abc5..8d9827c5c94 100644 --- a/cmd/experimental/experimental.go +++ b/cmd/experimental/experimental.go @@ -1,7 +1,6 @@ package experimental import ( - aircmd "github.com/databricks/cli/experimental/air/cmd" aitoolscmd "github.com/databricks/cli/experimental/aitools/cmd" geniecmd "github.com/databricks/cli/experimental/genie/cmd" postgrescmd "github.com/databricks/cli/experimental/postgres/cmd" @@ -23,7 +22,6 @@ These commands provide early access to new features that are still under development. They may change or be removed in future versions without notice.`, } - cmd.AddCommand(aircmd.New()) cmd.AddCommand(aitoolscmd.NewAitoolsCmd()) cmd.AddCommand(geniecmd.NewGenieCmd()) cmd.AddCommand(postgrescmd.New()) diff --git a/experimental/air/.gitkeep b/experimental/air/.gitkeep new file mode 100644 index 00000000000..e69de29bb2d diff --git a/experimental/air/cmd/air.go b/experimental/air/cmd/air.go deleted file mode 100644 index d99c5c660b1..00000000000 --- a/experimental/air/cmd/air.go +++ /dev/null @@ -1,17 +0,0 @@ -package aircmd - -import "github.com/spf13/cobra" - -// New creates the parent "air" command group for migrating AIR CLI workloads -// onto Databricks Asset Bundles. -func New() *cobra.Command { - cmd := &cobra.Command{ - Use: "air", - Short: "Migrate AI Runtime (AIR) CLI workloads to Databricks Asset Bundles", - Hidden: true, - } - - cmd.AddCommand(newInitCommand()) - - return cmd -} diff --git a/experimental/air/cmd/init.go b/experimental/air/cmd/init.go deleted file mode 100644 index 93a7ec74893..00000000000 --- a/experimental/air/cmd/init.go +++ /dev/null @@ -1,145 +0,0 @@ -package aircmd - -import ( - "fmt" - "os" - "path/filepath" - - "github.com/databricks/cli/bundle/generate" - "github.com/databricks/cli/experimental/air/trainyaml" - "github.com/databricks/cli/libs/cmdio" - "github.com/databricks/cli/libs/dyn" - "github.com/databricks/cli/libs/dyn/yamlsaver" - "github.com/databricks/cli/libs/textutil" - "github.com/databricks/databricks-sdk-go/service/jobs" - "github.com/spf13/cobra" -) - -func newInitCommand() *cobra.Command { - var fromPath string - var outputDir string - var jobKey string - var force bool - - cmd := &cobra.Command{ - Use: "init", - Short: "Scaffold a Databricks Asset Bundle from an AIR CLI train.yaml", - Long: `Scaffold a Databricks Asset Bundle from an AIR CLI train.yaml. - -This reads an existing AIR CLI train.yaml and generates a bundle that deploys the -same workload as a durable job using an AI Runtime task. The generated bundle -contains: - - - databricks.yml the bundle definition and a dev target - - resources/.job.yml the job with an ai_runtime_task - - /command.sh the command materialized as a script - -Deploy the generated bundle with: - databricks bundle deploy`, - } - - cmd.Flags().StringVarP(&fromPath, "from", "f", "", "Path to the AIR CLI train.yaml to migrate") - cmd.MarkFlagRequired("from") - cmd.Flags().StringVar(&outputDir, "output-dir", ".", "Directory to write the generated bundle into") - cmd.Flags().StringVar(&jobKey, "key", "", "Resource key for the generated job (defaults to the experiment name)") - cmd.Flags().BoolVar(&force, "force", false, "Overwrite existing files in the output directory") - - cmd.RunE = func(cmd *cobra.Command, args []string) error { - ctx := cmd.Context() - - cfg, err := trainyaml.Parse(fromPath) - if err != nil { - return err - } - - res, err := trainyaml.Convert(cfg) - if err != nil { - return err - } - - for _, w := range res.Warnings { - cmdio.LogString(ctx, "Warning: "+w) - } - - if jobKey == "" { - jobKey = textutil.NormalizeString(res.Job.Name) - } - - if err := writeBundle(outputDir, jobKey, res, force); err != nil { - return err - } - - cmdio.LogString(ctx, "Bundle written to "+filepath.ToSlash(outputDir)) - cmdio.LogString(ctx, "Deploy it with: databricks bundle deploy") - return nil - } - - return cmd -} - -// writeBundle materializes the generated bundle files under outputDir. -func writeBundle(outputDir, jobKey string, res *trainyaml.Result, force bool) error { - jobValue, err := generate.ConvertJobToValue(&jobs.Job{Settings: &res.Job.JobSettings}) - if err != nil { - return err - } - - saver := yamlsaver.NewSaver() - - // yamlsaver orders map keys by their location line, so assign increasing - // lines to fix the top-level key order (bundle, include, workspace). - // Without this the order is nondeterministic (all keys default to line 0). - rootValue := map[string]dyn.Value{ - "bundle": dyn.NewValue(map[string]dyn.Value{ - "name": dyn.V(jobKey), - }, []dyn.Location{{Line: 0}}), - "include": dyn.NewValue([]dyn.Value{dyn.V("resources/*.yml")}, []dyn.Location{{Line: 1}}), - } - // The snapshot's remote volume becomes the artifact path, so the code archive - // is uploaded to the volume the user already configured for AIR. - if res.ArtifactPath != "" { - rootValue["workspace"] = dyn.NewValue(map[string]dyn.Value{ - "artifact_path": dyn.V(res.ArtifactPath), - }, []dyn.Location{{Line: 2}}) - } - - resourcesDir := filepath.Join(outputDir, "resources") - if err := os.MkdirAll(resourcesDir, 0o755); err != nil { - return err - } - - if err := saver.SaveAsYAML(rootValue, filepath.Join(outputDir, "databricks.yml"), force); err != nil { - return err - } - - jobFile := map[string]dyn.Value{ - "resources": dyn.V(map[string]dyn.Value{ - "jobs": dyn.V(map[string]dyn.Value{ - jobKey: jobValue, - }), - }), - } - if err := saver.SaveAsYAML(jobFile, filepath.Join(resourcesDir, jobKey+".job.yml"), force); err != nil { - return err - } - - // Materialize the command next to the code it runs against, so command_path - // (which is relative to the extracted code source) resolves at runtime. - codeDir := filepath.Join(outputDir, filepath.FromSlash(codeSourceRootPath(res))) - if err := os.MkdirAll(codeDir, 0o755); err != nil { - return err - } - commandPath := filepath.Join(codeDir, "command.sh") - if !force { - if _, err := os.Stat(commandPath); err == nil { - return fmt.Errorf("%s already exists, use --force to overwrite", filepath.ToSlash(commandPath)) - } - } - return os.WriteFile(commandPath, []byte(res.CommandScript), 0o755) -} - -// codeSourceRootPath returns the local code_source_path of the generated job's -// AI Runtime task, i.e. the directory the command script is written into. -func codeSourceRootPath(res *trainyaml.Result) string { - return res.Job.Tasks[0].AiRuntimeTask.CodeSourcePath -} diff --git a/experimental/air/trainyaml/convert.go b/experimental/air/trainyaml/convert.go deleted file mode 100644 index 11dd926be6f..00000000000 --- a/experimental/air/trainyaml/convert.go +++ /dev/null @@ -1,168 +0,0 @@ -package trainyaml - -import ( - "errors" - "fmt" - "path" - - "github.com/databricks/cli/bundle/config/resources" - "github.com/databricks/databricks-sdk-go/service/compute" - "github.com/databricks/databricks-sdk-go/service/jobs" -) - -// environmentKey is the job environment key the generated AI Runtime task -// references. AI Runtime tasks require a serverless environment; we materialize -// one visibly in the generated config so the user can edit it afterwards. -const environmentKey = "default" - -// commandScriptName is the file the inline train.yaml command is written to. -// The AI Runtime task's command_path points at it, relative to the code source. -const commandScriptName = "command.sh" - -// Result holds everything Convert produces from a train.yaml: the DABs job -// resource, the shell script the command was materialized into, the artifact -// path (set only when the snapshot targets a UC volume), and any warnings about -// train.yaml fields that could not be represented. -type Result struct { - Job *resources.Job - CommandScript string - ArtifactPath string - Warnings []string -} - -// Convert maps a parsed train.yaml Config to a DABs job resource. It returns an -// error for inputs that cannot be represented as an AI Runtime task (legacy -// accelerators, pools, docker images) and collects warnings for fields that are -// dropped because they have no equivalent. -func Convert(cfg *Config) (*Result, error) { - if cfg.ExperimentName == "" { - return nil, errors.New("experiment_name is required") - } - if cfg.Command == "" { - return nil, errors.New("command is required") - } - if cfg.CodeSource == nil || cfg.CodeSource.Snapshot == nil || cfg.CodeSource.Snapshot.RootPath == "" { - return nil, errors.New("code_source.snapshot.root_path is required to migrate to an AI Runtime task") - } - - accelerator, err := convertAcceleratorType(cfg.Compute.AcceleratorType) - if err != nil { - return nil, err - } - if cfg.Compute.NodePoolID != "" || cfg.Compute.PoolName != "" { - return nil, errors.New("compute pools (node_pool_id/pool_name) are not supported by ai_runtime_task; remove them to use serverless AI Runtime compute") - } - - env, err := convertEnvironment(cfg.Environment) - if err != nil { - return nil, err - } - - res := &Result{} - - task := jobs.Task{ - TaskKey: cfg.ExperimentName, - EnvironmentKey: environmentKey, - AiRuntimeTask: &jobs.AiRuntimeTask{ - Experiment: cfg.ExperimentName, - MlflowRun: cfg.MlflowRunName, - MlflowExperimentDirectory: cfg.MlflowExperimentDirectory, - CodeSourcePath: cfg.CodeSource.Snapshot.RootPath, - Deployments: []jobs.DeploymentSpec{ - { - // command_path is relative to the extracted code source root. - CommandPath: path.Join(path.Base(cfg.CodeSource.Snapshot.RootPath), commandScriptName), - Compute: jobs.ComputeSpec{ - AcceleratorType: accelerator, - AcceleratorCount: cfg.Compute.NumAccelerators, - }, - }, - }, - }, - } - if cfg.MaxRetries != nil { - task.MaxRetries = *cfg.MaxRetries - } - if cfg.TimeoutMinutes != nil { - task.TimeoutSeconds = *cfg.TimeoutMinutes * 60 - } - - res.Job = &resources.Job{ - JobSettings: jobs.JobSettings{ - Name: cfg.ExperimentName, - Tasks: []jobs.Task{task}, - Environments: []jobs.JobEnvironment{{EnvironmentKey: environmentKey, Spec: env}}, - }, - } - - res.CommandScript = cfg.Command - if cfg.CodeSource.Snapshot.RemoteVolume != "" { - res.ArtifactPath = cfg.CodeSource.Snapshot.RemoteVolume - } - - res.Warnings = append(res.Warnings, unmappedWarnings(cfg)...) - return res, nil -} - -func convertAcceleratorType(v string) (jobs.ComputeSpecAcceleratorType, error) { - if v == "" { - return "", errors.New("compute.accelerator_type is required") - } - var t jobs.ComputeSpecAcceleratorType - if err := t.Set(v); err != nil { - return "", fmt.Errorf("compute.accelerator_type %q is not supported by ai_runtime_task (expected GPU_1xA10, GPU_1xH100, or GPU_8xH100)", v) - } - return t, nil -} - -func convertEnvironment(env *Environment) (*compute.Environment, error) { - if env == nil { - return &compute.Environment{}, nil - } - if env.DockerImage != nil { - return nil, errors.New("environment.docker_image is not supported by ai_runtime_task; use environment.version and environment.dependencies instead") - } - if env.Dependencies.Path != "" { - return nil, errors.New("environment.dependencies as a requirements file path is not supported; provide an inline list of dependencies instead") - } - - return &compute.Environment{ - EnvironmentVersion: env.Version, - Dependencies: env.Dependencies.List, - }, nil -} - -// unmappedWarnings reports train.yaml fields that have no ai_runtime_task -// equivalent and are therefore dropped from the generated config. -func unmappedWarnings(cfg *Config) []string { - var warnings []string - warn := func(field, hint string) { - warnings = append(warnings, fmt.Sprintf("%q has no ai_runtime_task equivalent and was dropped%s", field, hint)) - } - - if len(cfg.Secrets) > 0 { - warn("secrets", "; configure secrets on the job or via {{secrets/scope/key}} references") - } - if len(cfg.EnvVariables) > 0 { - warn("env_variables", "; set them inside command.sh or the serverless environment") - } - if len(cfg.Parameters) > 0 { - warn("parameters", "") - } - if cfg.Priority != nil { - warn("priority", "") - } - if cfg.UsagePolicyName != "" { - warn("usage_policy_name", "") - } - if cfg.UsagePolicyID != "" { - warn("usage_policy_id", "") - } - if cfg.IdempotencyToken != "" { - warn("idempotency_token", "") - } - if cfg.CodeSource != nil && cfg.CodeSource.Snapshot != nil && len(cfg.CodeSource.Snapshot.IncludePaths) > 0 { - warn("code_source.snapshot.include_paths", "; the entire code_source directory is packaged") - } - return warnings -} diff --git a/experimental/air/trainyaml/convert_test.go b/experimental/air/trainyaml/convert_test.go deleted file mode 100644 index ea44722771e..00000000000 --- a/experimental/air/trainyaml/convert_test.go +++ /dev/null @@ -1,134 +0,0 @@ -package trainyaml - -import ( - "os" - "path/filepath" - "testing" - - "github.com/stretchr/testify/assert" - "github.com/stretchr/testify/require" -) - -func writeTrainYAML(t *testing.T, content string) string { - t.Helper() - dir := t.TempDir() - p := filepath.Join(dir, "train.yaml") - require.NoError(t, os.WriteFile(p, []byte(content), 0o644)) - return p -} - -func TestParseAndConvertSnapshot(t *testing.T) { - p := writeTrainYAML(t, ` -experiment_name: my-training -compute: - num_accelerators: 8 - accelerator_type: GPU_8xH100 -environment: - version: "5" - dependencies: - - torch>=2.0.0 -code_source: - type: snapshot - snapshot: - root_path: ./src -command: | - cd $CODE_SOURCE_PATH - python train.py -max_retries: 2 -timeout_minutes: 30 -`) - - cfg, err := Parse(p) - require.NoError(t, err) - - res, err := Convert(cfg) - require.NoError(t, err) - require.NotNil(t, res.Job) - - task := res.Job.Tasks[0] - assert.Equal(t, "my-training", task.TaskKey) - assert.Equal(t, "default", task.EnvironmentKey) - assert.Equal(t, 2, task.MaxRetries) - assert.Equal(t, 1800, task.TimeoutSeconds) - - air := task.AiRuntimeTask - require.NotNil(t, air) - assert.Equal(t, "my-training", air.Experiment) - assert.Equal(t, "./src", air.CodeSourcePath) - require.Len(t, air.Deployments, 1) - assert.Equal(t, "src/command.sh", air.Deployments[0].CommandPath) - assert.EqualValues(t, "GPU_8xH100", air.Deployments[0].Compute.AcceleratorType) - assert.Equal(t, 8, air.Deployments[0].Compute.AcceleratorCount) - - require.Len(t, res.Job.Environments, 1) - assert.Equal(t, "default", res.Job.Environments[0].EnvironmentKey) - assert.Equal(t, "5", res.Job.Environments[0].Spec.EnvironmentVersion) - assert.Equal(t, []string{"torch>=2.0.0"}, res.Job.Environments[0].Spec.Dependencies) - - assert.Contains(t, res.CommandScript, "python train.py") - assert.Empty(t, res.ArtifactPath) -} - -func TestConvertRemoteVolumeSetsArtifactPath(t *testing.T) { - cfg := &Config{ - ExperimentName: "exp", - Command: "echo hi", - Compute: Compute{NumAccelerators: 1, AcceleratorType: "GPU_1xA10"}, - CodeSource: &CodeSource{ - Type: "snapshot", - Snapshot: &Snapshot{RootPath: "code", RemoteVolume: "/Volumes/main/default/code"}, - }, - } - res, err := Convert(cfg) - require.NoError(t, err) - assert.Equal(t, "/Volumes/main/default/code", res.ArtifactPath) -} - -func TestConvertRejectsLegacyAccelerator(t *testing.T) { - cfg := &Config{ - ExperimentName: "exp", - Command: "echo hi", - Compute: Compute{NumAccelerators: 1, AcceleratorType: "a10"}, - CodeSource: &CodeSource{Type: "snapshot", Snapshot: &Snapshot{RootPath: "code"}}, - } - _, err := Convert(cfg) - require.Error(t, err) - assert.Contains(t, err.Error(), "not supported by ai_runtime_task") -} - -func TestConvertRejectsDockerImage(t *testing.T) { - cfg := &Config{ - ExperimentName: "exp", - Command: "echo hi", - Compute: Compute{NumAccelerators: 1, AcceleratorType: "GPU_1xA10"}, - Environment: &Environment{DockerImage: &DockerImage{URL: "org/repo:tag"}}, - CodeSource: &CodeSource{Type: "snapshot", Snapshot: &Snapshot{RootPath: "code"}}, - } - _, err := Convert(cfg) - require.Error(t, err) - assert.Contains(t, err.Error(), "docker_image is not supported") -} - -func TestConvertWarnsOnUnmappedFields(t *testing.T) { - cfg := &Config{ - ExperimentName: "exp", - Command: "echo hi", - Compute: Compute{NumAccelerators: 1, AcceleratorType: "GPU_1xA10"}, - Secrets: map[string]string{"HF_TOKEN": "scope/key"}, - CodeSource: &CodeSource{ - Type: "snapshot", - Snapshot: &Snapshot{RootPath: "code", IncludePaths: []string{"a", "b"}}, - }, - } - res, err := Convert(cfg) - require.NoError(t, err) - require.Len(t, res.Warnings, 2) - assert.Contains(t, res.Warnings[0], "secrets") - assert.Contains(t, res.Warnings[1], "include_paths") -} - -func TestParseRejectsUnknownField(t *testing.T) { - p := writeTrainYAML(t, "experiment_name: x\nbogus_field: 1\n") - _, err := Parse(p) - require.Error(t, err) -} diff --git a/experimental/air/trainyaml/parse.go b/experimental/air/trainyaml/parse.go deleted file mode 100644 index dc29ba264d6..00000000000 --- a/experimental/air/trainyaml/parse.go +++ /dev/null @@ -1,110 +0,0 @@ -// Package trainyaml parses the AIR CLI's train.yaml launch config and converts -// it into a DABs job resource that uses an AI Runtime task. It exists so users -// who previously launched ephemeral runs with the AIR CLI can migrate onto a -// durable, DABs-managed job while keeping the familiar code_source experience. -package trainyaml - -import ( - "bytes" - "fmt" - "os" - - "go.yaml.in/yaml/v3" -) - -// Config is the subset of the AIR CLI train.yaml schema this migration -// understands. Fields that have no DABs equivalent are still parsed so we can -// warn about them rather than silently drop them (see Convert). -// -// Reference: universe/ai-compute/cli/cli/sdk/config.py (RunConfig and friends). -type Config struct { - ExperimentName string `yaml:"experiment_name"` - Compute Compute `yaml:"compute"` - Command string `yaml:"command"` - - Environment *Environment `yaml:"environment"` - CodeSource *CodeSource `yaml:"code_source"` - - EnvVariables map[string]string `yaml:"env_variables"` - Secrets map[string]string `yaml:"secrets"` - - MaxRetries *int `yaml:"max_retries"` - TimeoutMinutes *int `yaml:"timeout_minutes"` - MlflowRunName string `yaml:"mlflow_run_name"` - MlflowExperimentDirectory string `yaml:"mlflow_experiment_directory"` - - // Fields with no AI Runtime task equivalent. Retained so Convert can warn. - Parameters map[string]any `yaml:"parameters"` - Priority *int `yaml:"priority"` - UsagePolicyName string `yaml:"usage_policy_name"` - UsagePolicyID string `yaml:"usage_policy_id"` - IdempotencyToken string `yaml:"idempotency_token"` -} - -type Compute struct { - NumAccelerators int `yaml:"num_accelerators"` - AcceleratorType string `yaml:"accelerator_type"` - NodePoolID string `yaml:"node_pool_id"` - PoolName string `yaml:"pool_name"` -} - -// Environment mirrors the train.yaml environment block. Dependencies is -// polymorphic in the AIR schema (a requirements.yaml path or an inline list); -// we only support the inline list on the DABs serverless environment and report -// the string form as unsupported. -type Environment struct { - Dependencies dependencies `yaml:"dependencies"` - Version string `yaml:"version"` - DockerImage *DockerImage `yaml:"docker_image"` -} - -type DockerImage struct { - URL string `yaml:"url"` -} - -type CodeSource struct { - Type string `yaml:"type"` - Snapshot *Snapshot `yaml:"snapshot"` -} - -type Snapshot struct { - RootPath string `yaml:"root_path"` - RemoteVolume string `yaml:"remote_volume"` - IncludePaths []string `yaml:"include_paths"` -} - -// dependencies captures the polymorphic environment.dependencies field. Exactly -// one of List / Path is set after unmarshaling. -type dependencies struct { - List []string - Path string -} - -func (d *dependencies) UnmarshalYAML(node *yaml.Node) error { - switch node.Kind { - case yaml.ScalarNode: - return node.Decode(&d.Path) - case yaml.SequenceNode: - return node.Decode(&d.List) - default: - return fmt.Errorf("environment.dependencies must be a string or a list, got %v", node.Kind) - } -} - -// Parse reads and decodes a train.yaml file. Unknown fields are rejected so a -// typo or an unsupported field surfaces to the user instead of being ignored. -func Parse(path string) (*Config, error) { - raw, err := os.ReadFile(path) - if err != nil { - return nil, err - } - - dec := yaml.NewDecoder(bytes.NewReader(raw)) - dec.KnownFields(true) - - var cfg Config - if err := dec.Decode(&cfg); err != nil { - return nil, fmt.Errorf("failed to parse %s: %w", path, err) - } - return &cfg, nil -} From 46b350f1046859501aff9603643df2300b29a3e7 Mon Sep 17 00:00:00 2001 From: "ben.hansen" Date: Tue, 14 Jul 2026 03:27:16 +0000 Subject: [PATCH 3/5] Simplify code_source snapshotting to plain tar Drop the git-archive path (HEAD-only archive + (commit, include_paths) cache) since the SDK ai_runtime_task exposes only a plain code_source_path string with no way to select a git ref, so the cache added complexity for little value. code_source_path now always packages the local directory as a plain tar (.git and gitignored files excluded) and uploads it to the user's ~/.air/repo_snapshots directory. Co-authored-by: Isaac --- NEXT_CHANGELOG.md | 2 +- .../ai_runtime_task/git_archive/output.txt | 23 ---- .../bundle/ai_runtime_task/git_archive/script | 28 ----- .../ai_runtime_task/git_archive/test.toml | 13 --- .../databricks.yml | 0 .../out.test.toml | 0 .../output.txt | 4 +- .../{plain_tar => local_code_source}/script | 10 +- .../src/command.sh | 0 .../src/train.py | 0 .../test.toml | 0 .../ai_runtime_task/plain_tar/databricks.yml | 24 ---- .../ai_runtime_task/plain_tar/out.test.toml | 3 - .../ai_runtime_task/plain_tar/src/command.sh | 2 - .../ai_runtime_task/plain_tar/src/train.py | 1 - .../config/mutator/aicode/package_upload.go | 107 ++++-------------- .../mutator/aicode/package_upload_test.go | 67 ++--------- .../mutator/aicode/snapshot_cachekey.go | 32 ------ .../mutator/aicode/snapshot_cachekey_test.go | 33 ------ bundle/config/mutator/aicode/snapshot_git.go | 77 ------------- .../config/mutator/aicode/snapshot_package.go | 49 ++------ .../mutator/aicode/snapshot_package_test.go | 28 +---- .../config/mutator/aicode/snapshot_resolve.go | 56 --------- bundle/config/mutator/aicode/snapshot_test.go | 72 ------------ .../mutator/aicode/testdata/cache_keys.json | 77 ------------- 25 files changed, 54 insertions(+), 654 deletions(-) delete mode 100644 acceptance/bundle/ai_runtime_task/git_archive/output.txt delete mode 100644 acceptance/bundle/ai_runtime_task/git_archive/script delete mode 100644 acceptance/bundle/ai_runtime_task/git_archive/test.toml rename acceptance/bundle/ai_runtime_task/{git_archive => local_code_source}/databricks.yml (100%) rename acceptance/bundle/ai_runtime_task/{git_archive => local_code_source}/out.test.toml (100%) rename acceptance/bundle/ai_runtime_task/{plain_tar => local_code_source}/output.txt (85%) rename acceptance/bundle/ai_runtime_task/{plain_tar => local_code_source}/script (52%) rename acceptance/bundle/ai_runtime_task/{git_archive => local_code_source}/src/command.sh (100%) rename acceptance/bundle/ai_runtime_task/{git_archive => local_code_source}/src/train.py (100%) rename acceptance/bundle/ai_runtime_task/{plain_tar => local_code_source}/test.toml (100%) delete mode 100644 acceptance/bundle/ai_runtime_task/plain_tar/databricks.yml delete mode 100644 acceptance/bundle/ai_runtime_task/plain_tar/out.test.toml delete mode 100644 acceptance/bundle/ai_runtime_task/plain_tar/src/command.sh delete mode 100644 acceptance/bundle/ai_runtime_task/plain_tar/src/train.py delete mode 100644 bundle/config/mutator/aicode/snapshot_cachekey.go delete mode 100644 bundle/config/mutator/aicode/snapshot_cachekey_test.go delete mode 100644 bundle/config/mutator/aicode/snapshot_git.go delete mode 100644 bundle/config/mutator/aicode/snapshot_resolve.go delete mode 100644 bundle/config/mutator/aicode/snapshot_test.go delete mode 100644 bundle/config/mutator/aicode/testdata/cache_keys.json diff --git a/NEXT_CHANGELOG.md b/NEXT_CHANGELOG.md index 51b3ba8b214..0da6ef56c59 100644 --- a/NEXT_CHANGELOG.md +++ b/NEXT_CHANGELOG.md @@ -8,7 +8,7 @@ ### Bundles -* An `ai_runtime_task.code_source_path` that points at a local directory is now packaged into a tarball during `bundle deploy` (git archive of a clean HEAD, otherwise a plain tar of the working tree), uploaded to the user's workspace code snapshot cache, and rewritten to the uploaded path, mirroring the AIR CLI's `code_source` behavior. +* An `ai_runtime_task.code_source_path` that points at a local directory is now packaged into a tarball during `bundle deploy`, uploaded to the user's workspace code snapshot directory, and rewritten to the uploaded path, mirroring the AIR CLI's `code_source` behavior. ### Dependency updates diff --git a/acceptance/bundle/ai_runtime_task/git_archive/output.txt b/acceptance/bundle/ai_runtime_task/git_archive/output.txt deleted file mode 100644 index 0465656efcb..00000000000 --- a/acceptance/bundle/ai_runtime_task/git_archive/output.txt +++ /dev/null @@ -1,23 +0,0 @@ - -=== deploy packages the code source and rewrites code_source_path - ->>> [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 cache -/api/2.0/workspace-files/import-file/Workspace/Users/[USERNAME]/.air/repo_snapshots/src/src_[SNAPSHOT].tar.gz - -=== the created job's code_source_path points at the uploaded archive -/Workspace/Users/[USERNAME]/.air/repo_snapshots/src/src_[SNAPSHOT].tar.gz - -=== re-deploying the same commit 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 diff --git a/acceptance/bundle/ai_runtime_task/git_archive/script b/acceptance/bundle/ai_runtime_task/git_archive/script deleted file mode 100644 index d0715e02a43..00000000000 --- a/acceptance/bundle/ai_runtime_task/git_archive/script +++ /dev/null @@ -1,28 +0,0 @@ -# A committed, clean git repo makes the snapshot use git_archive (cacheable by the -# commit SHA). Pinned commit dates keep the resolved HEAD SHA — and thus the 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 -git add src -git commit -qm 'Add code source' - -title "deploy packages the code source and rewrites code_source_path\n" -trace $CLI bundle deploy - -title "the snapshot tarball is uploaded to the user's repo_snapshots cache\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 "the created job's code_source_path points at the uploaded archive\n" -jq -r 'select(.path == "/api/2.2/jobs/create") | .body.tasks[0].ai_runtime_task.code_source_path' out.requests.txt - -rm out.requests.txt - -title "re-deploying the same commit 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 | wc -l | tr -d ' ' - -rm out.requests.txt diff --git a/acceptance/bundle/ai_runtime_task/git_archive/test.toml b/acceptance/bundle/ai_runtime_task/git_archive/test.toml deleted file mode 100644 index 5087d92c7ec..00000000000 --- a/acceptance/bundle/ai_runtime_task/git_archive/test.toml +++ /dev/null @@ -1,13 +0,0 @@ -RecordRequests = true - -Ignore = [ - '.databricks', - '.git', -] - -# git_archive names the tarball _.tar.gz, where the cache key -# derives from the committed HEAD SHA. Collapse it to a stable token so output does -# not depend on the commit SHA. -[[Repls]] -Old = 'src_[0-9a-f]{16}\.tar\.gz' -New = 'src_[SNAPSHOT].tar.gz' diff --git a/acceptance/bundle/ai_runtime_task/git_archive/databricks.yml b/acceptance/bundle/ai_runtime_task/local_code_source/databricks.yml similarity index 100% rename from acceptance/bundle/ai_runtime_task/git_archive/databricks.yml rename to acceptance/bundle/ai_runtime_task/local_code_source/databricks.yml diff --git a/acceptance/bundle/ai_runtime_task/git_archive/out.test.toml b/acceptance/bundle/ai_runtime_task/local_code_source/out.test.toml similarity index 100% rename from acceptance/bundle/ai_runtime_task/git_archive/out.test.toml rename to acceptance/bundle/ai_runtime_task/local_code_source/out.test.toml diff --git a/acceptance/bundle/ai_runtime_task/plain_tar/output.txt b/acceptance/bundle/ai_runtime_task/local_code_source/output.txt similarity index 85% rename from acceptance/bundle/ai_runtime_task/plain_tar/output.txt rename to acceptance/bundle/ai_runtime_task/local_code_source/output.txt index 79d4ab65e54..1472320d49c 100644 --- a/acceptance/bundle/ai_runtime_task/plain_tar/output.txt +++ b/acceptance/bundle/ai_runtime_task/local_code_source/output.txt @@ -1,5 +1,5 @@ -=== deploy packages the working tree as a plain tarball +=== deploy packages the local code source and rewrites code_source_path >>> [CLI] bundle deploy Uploading bundle files to /Workspace/Users/[USERNAME]/.bundle/ai-runtime-test/default/files... @@ -7,7 +7,7 @@ Deploying resources... Updating deployment state... Deployment complete! -=== the snapshot tarball is uploaded to the user's repo_snapshots cache +=== 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 === the created job's code_source_path points at the uploaded archive diff --git a/acceptance/bundle/ai_runtime_task/plain_tar/script b/acceptance/bundle/ai_runtime_task/local_code_source/script similarity index 52% rename from acceptance/bundle/ai_runtime_task/plain_tar/script rename to acceptance/bundle/ai_runtime_task/local_code_source/script index 9978eb9ba22..5d729a8e472 100644 --- a/acceptance/bundle/ai_runtime_task/plain_tar/script +++ b/acceptance/bundle/ai_runtime_task/local_code_source/script @@ -1,10 +1,12 @@ -# The bundle dir is not a git repo, so the snapshot falls back to plain_tar of the -# working tree (no cache; timestamp-named archive). +# A local code_source_path is packaged into a tarball and uploaded to the user's +# repo_snapshots directory; the archive is timestamp-named. -title "deploy packages the working tree as a plain tarball\n" +title "deploy packages the local code source and rewrites code_source_path\n" trace $CLI bundle deploy -title "the snapshot tarball is uploaded to the user's repo_snapshots cache\n" +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 "the created job's code_source_path points at the uploaded archive\n" diff --git a/acceptance/bundle/ai_runtime_task/git_archive/src/command.sh b/acceptance/bundle/ai_runtime_task/local_code_source/src/command.sh similarity index 100% rename from acceptance/bundle/ai_runtime_task/git_archive/src/command.sh rename to acceptance/bundle/ai_runtime_task/local_code_source/src/command.sh diff --git a/acceptance/bundle/ai_runtime_task/git_archive/src/train.py b/acceptance/bundle/ai_runtime_task/local_code_source/src/train.py similarity index 100% rename from acceptance/bundle/ai_runtime_task/git_archive/src/train.py rename to acceptance/bundle/ai_runtime_task/local_code_source/src/train.py diff --git a/acceptance/bundle/ai_runtime_task/plain_tar/test.toml b/acceptance/bundle/ai_runtime_task/local_code_source/test.toml similarity index 100% rename from acceptance/bundle/ai_runtime_task/plain_tar/test.toml rename to acceptance/bundle/ai_runtime_task/local_code_source/test.toml diff --git a/acceptance/bundle/ai_runtime_task/plain_tar/databricks.yml b/acceptance/bundle/ai_runtime_task/plain_tar/databricks.yml deleted file mode 100644 index c6d37475f73..00000000000 --- a/acceptance/bundle/ai_runtime_task/plain_tar/databricks.yml +++ /dev/null @@ -1,24 +0,0 @@ -bundle: - name: ai-runtime-test - -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 diff --git a/acceptance/bundle/ai_runtime_task/plain_tar/out.test.toml b/acceptance/bundle/ai_runtime_task/plain_tar/out.test.toml deleted file mode 100644 index f784a183258..00000000000 --- a/acceptance/bundle/ai_runtime_task/plain_tar/out.test.toml +++ /dev/null @@ -1,3 +0,0 @@ -Local = true -Cloud = false -EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["terraform", "direct"] diff --git a/acceptance/bundle/ai_runtime_task/plain_tar/src/command.sh b/acceptance/bundle/ai_runtime_task/plain_tar/src/command.sh deleted file mode 100644 index 7ce3949767c..00000000000 --- a/acceptance/bundle/ai_runtime_task/plain_tar/src/command.sh +++ /dev/null @@ -1,2 +0,0 @@ -cd $CODE_SOURCE_PATH -python train.py diff --git a/acceptance/bundle/ai_runtime_task/plain_tar/src/train.py b/acceptance/bundle/ai_runtime_task/plain_tar/src/train.py deleted file mode 100644 index d8b062e5dfe..00000000000 --- a/acceptance/bundle/ai_runtime_task/plain_tar/src/train.py +++ /dev/null @@ -1 +0,0 @@ -print("training") diff --git a/bundle/config/mutator/aicode/package_upload.go b/bundle/config/mutator/aicode/package_upload.go index 662c2e1c004..e54a1db70d8 100644 --- a/bundle/config/mutator/aicode/package_upload.go +++ b/bundle/config/mutator/aicode/package_upload.go @@ -5,22 +5,19 @@ // volume path to an uploaded code archive; its doc comment states that the CLI // is responsible for packaging the user's local code directory into that // archive. This mutator implements that contract for DABs: when a user points -// code_source_path at a local directory, it snapshots the directory (git archive -// of HEAD when the tree is a clean git repo, otherwise a plain tar; .git and -// gitignored files excluded), uploads the archive to the user's workspace code -// snapshot cache, and rewrites the field to the resulting remote path so the -// deployed job runs against the uploaded code. Values that are already remote are -// left untouched. +// code_source_path at a local directory, it packages the directory into a plain +// tarball (.git and gitignored files excluded), uploads the archive to the user's +// workspace code snapshot directory, and rewrites the field to the resulting +// remote path so the deployed job runs against the uploaded code. Values that are +// already remote are left untouched. // -// The snapshot approach mirrors PR #5897, which ported it from the Python air CLI -// (see snapshot_resolve.go / snapshot_package.go / snapshot_cachekey.go). +// The tarball builder is ported from PR #5897 (see snapshot_package.go). package aicode import ( "context" "errors" "fmt" - "io/fs" "os" "path" "path/filepath" @@ -31,7 +28,6 @@ import ( "github.com/databricks/cli/libs/diag" "github.com/databricks/cli/libs/dyn" "github.com/databricks/cli/libs/filer" - "github.com/databricks/cli/libs/log" ) // codeSourcePatterns are the config locations of an AI Runtime task's @@ -122,20 +118,16 @@ func (m *packageAndUpload) Apply(ctx context.Context, b *bundle.Bundle) diag.Dia return diags } -// repoSnapshotsSubdir is the per-user workspace location for cached code -// snapshots, under the user's home. It matches the Python air CLI (and PR #5897) -// so a snapshot is cached in a stable place that bundle deploy does not wipe. -// -// Deliberately NOT /.internal: artifacts.CleanUp() deletes that -// directory at the start of every deploy, which would defeat the git-archive -// snapshot cache. Keeping snapshots under ~/.air/repo_snapshots lets an unchanged -// commit skip re-upload across deploys. +// repoSnapshotsSubdir is the per-user workspace location for code snapshots, +// under the user's home. It matches the Python air CLI (and PR #5897) and is +// deliberately NOT /.internal, which artifacts.CleanUp() deletes at +// the start of every deploy. const repoSnapshotsSubdir = ".air/repo_snapshots" -// packageOne snapshots the local directory for a single code source, uploads the -// archive to the user's repo_snapshots dir (skipping the upload when an identical -// git-archive snapshot already exists), and returns the remote path the config -// should point to. +// packageOne packages the local directory for a single code source into a plain +// tarball, uploads it to the user's repo_snapshots dir, and returns the remote +// path the config should point to. The archive is timestamp-named so concurrent +// deploys do not clobber each other. func (m *packageAndUpload) packageOne(ctx context.Context, b *bundle.Bundle, cs codeSource, stagingDir, userDir string) (string, error) { localDir := filepath.Join(b.SyncRootPath, filepath.FromSlash(cs.value)) dirName := filepath.Base(localDir) @@ -150,50 +142,22 @@ func (m *packageAndUpload) packageOne(ctx context.Context, b *bundle.Bundle, cs } } - git := newGitRepo(localDir) - plan, err := resolveSnapshotPlan(ctx, git) - if err != nil { - return "", fmt.Errorf("failed to resolve code_source_path %q: %w", cs.value, err) - } - - // git_archive is cacheable by (commit, include_paths): a hit means the identical - // tarball is already uploaded, so packaging and upload are skipped entirely. - if plan.mode == modeGitArchive { - cacheKey := computeSnapshotCacheKey(plan.commitSHA, plan.includePaths) - archiveName := fmt.Sprintf("%s_%s.tar.gz", dirName, cacheKey[:16]) - remotePath := path.Join(uploadPath, archiveName) - - exists, err := archiveExists(ctx, client, archiveName) - if err != nil { - return "", err - } - if exists { - log.Debugf(ctx, "code snapshot for %s already present at %s, skipping upload", shortSHA(plan.commitSHA), remotePath) - return remotePath, nil - } + archiveName := fmt.Sprintf("%s_%s.tar.gz", dirName, time.Now().UTC().Format("20060102_150405")) + localArchive := filepath.Join(stagingDir, archiveName) + defer os.Remove(localArchive) - if err := packageAndUploadArchive(ctx, client, stagingDir, archiveName, func(out string) error { - return createGitArchiveSnapshot(ctx, git, plan.commitSHA, out, dirName, plan.includePaths) - }); err != nil { - return "", err - } - return remotePath, nil + if err := createPlainTarball(ctx, localDir, localArchive); err != nil { + return "", fmt.Errorf("failed to package code_source_path %q: %w", cs.value, err) } - - // plain_tar is not cacheable (working-tree content isn't pinned to a SHA), so it - // is timestamp-named to avoid clobbering a concurrent deploy. - archiveName := fmt.Sprintf("%s_%s.tar.gz", dirName, time.Now().UTC().Format("20060102_150405")) - remotePath := path.Join(uploadPath, archiveName) - if err := packageAndUploadArchive(ctx, client, stagingDir, archiveName, func(out string) error { - return createPlainTarball(ctx, localDir, out, plan.includePaths) - }); err != nil { + if err := libraries.UploadFile(ctx, localArchive, client); err != nil { return "", err } - return remotePath, nil + + return path.Join(uploadPath, archiveName), nil } // userWorkspaceHome returns the current user's workspace home directory -// (/Workspace/Users/), the root under which code snapshots are cached. +// (/Workspace/Users/), the root under which code snapshots are stored. func userWorkspaceHome(b *bundle.Bundle) (string, error) { u := b.Config.Workspace.CurrentUser if u == nil || u.User == nil || u.UserName == "" { @@ -202,31 +166,6 @@ func userWorkspaceHome(b *bundle.Bundle) (string, error) { return "/Workspace/Users/" + u.UserName, nil } -// archiveExists reports whether archiveName already exists in the store, treating -// fs.ErrNotExist as "no". -func archiveExists(ctx context.Context, client filer.Filer, archiveName string) (bool, error) { - _, err := client.Stat(ctx, archiveName) - if err == nil { - return true, nil - } - if errors.Is(err, fs.ErrNotExist) { - return false, nil - } - return false, fmt.Errorf("failed to check for existing code snapshot %q: %w", archiveName, err) -} - -// packageAndUploadArchive writes the tarball via pkg into a temp file under -// stagingDir, then uploads it as archiveName. The temp file is always removed. -func packageAndUploadArchive(ctx context.Context, client filer.Filer, stagingDir, archiveName string, pkg func(outputPath string) error) error { - localArchive := filepath.Join(stagingDir, archiveName) - defer os.Remove(localArchive) - - if err := pkg(localArchive); err != nil { - return err - } - return libraries.UploadFile(ctx, localArchive, client) -} - // collectLocalCodeSources returns every AI Runtime task code_source_path that // points at a local directory. Already-remote values are skipped. func collectLocalCodeSources(b *bundle.Bundle) ([]codeSource, diag.Diagnostics) { diff --git a/bundle/config/mutator/aicode/package_upload_test.go b/bundle/config/mutator/aicode/package_upload_test.go index 6b68d98f507..1e6352513d7 100644 --- a/bundle/config/mutator/aicode/package_upload_test.go +++ b/bundle/config/mutator/aicode/package_upload_test.go @@ -24,16 +24,14 @@ import ( // tests, derived from the fixture user below. const codeSnapshotDir = "/Workspace/Users/me@databricks.com/.air/repo_snapshots/src" -// recordingFiler records every Write and reports files as present per exists, so a -// test can both capture uploads and simulate a cache hit. +// recordingFiler records every Write, so a test can capture the uploaded archive. type recordingFiler struct { filer.Filer written map[string][]byte - exists map[string]bool } func newRecordingFiler() *recordingFiler { - return &recordingFiler{written: map[string][]byte{}, exists: map[string]bool{}} + return &recordingFiler{written: map[string][]byte{}} } func (f *recordingFiler) Write(ctx context.Context, p string, r io.Reader, _ ...filer.WriteMode) error { @@ -46,17 +44,11 @@ func (f *recordingFiler) Write(ctx context.Context, p string, r io.Reader, _ ... } func (f *recordingFiler) Stat(ctx context.Context, p string) (fs.FileInfo, error) { - if f.exists[p] { - return fakeFileInfo{}, nil - } return nil, fs.ErrNotExist } -type fakeFileInfo struct{ fs.FileInfo } - // bundleWithCodeSource builds a bundle rooted at dir whose single AI Runtime task -// points at codeSourcePath. The caller populates dir (and optionally inits a git -// repo) before applying the mutator. +// points at codeSourcePath. The caller populates dir before applying the mutator. func bundleWithCodeSource(t *testing.T, dir, codeSourcePath string) *bundle.Bundle { t.Helper() b := &bundle.Bundle{ @@ -87,9 +79,8 @@ func bundleWithCodeSource(t *testing.T, dir, codeSourcePath string) *bundle.Bund return b } -// bundleWithPlainSrc creates a non-git src/ directory with one file and returns -// the bundle (so packaging falls back to plain_tar). -func bundleWithPlainSrc(t *testing.T, codeSourcePath string) *bundle.Bundle { +// bundleWithSrc creates a src/ directory with one file and returns the bundle. +func bundleWithSrc(t *testing.T, codeSourcePath string) *bundle.Bundle { t.Helper() dir := t.TempDir() require.NoError(t, os.MkdirAll(filepath.Join(dir, "src"), 0o755)) @@ -101,9 +92,8 @@ func codeSourcePath(b *bundle.Bundle) string { return b.Config.Resources.Jobs["train"].Tasks[0].AiRuntimeTask.CodeSourcePath } -func TestPackageAndUploadPlainTarRewritesLocalCodeSource(t *testing.T) { - // Non-git src/ dir → plain_tar → timestamp-named archive. - b := bundleWithPlainSrc(t, "src") +func TestPackageAndUploadRewritesLocalCodeSource(t *testing.T) { + b := bundleWithSrc(t, "src") f := newRecordingFiler() diags := bundle.Apply(t.Context(), b, &packageAndUpload{client: f}) @@ -118,7 +108,7 @@ func TestPackageAndUploadPlainTarRewritesLocalCodeSource(t *testing.T) { } func TestPackageAndUploadSkipsRemoteCodeSource(t *testing.T) { - b := bundleWithPlainSrc(t, "/Volumes/main/default/code/existing.tar.gz") + b := bundleWithSrc(t, "/Volumes/main/default/code/existing.tar.gz") f := newRecordingFiler() diags := bundle.Apply(t.Context(), b, &packageAndUpload{client: f}) @@ -127,44 +117,3 @@ func TestPackageAndUploadSkipsRemoteCodeSource(t *testing.T) { assert.Equal(t, "/Volumes/main/default/code/existing.tar.gz", codeSourcePath(b)) assert.Empty(t, f.written, "remote code_source_path must not trigger an upload") } - -func TestPackageAndUploadGitArchiveCacheHitSkipsUpload(t *testing.T) { - // A clean git repo → git_archive → cache-key-named archive; when it already - // exists in the store, upload is skipped and the path is still rewritten. - dir := t.TempDir() - require.NoError(t, os.MkdirAll(filepath.Join(dir, "src"), 0o755)) - require.NoError(t, os.WriteFile(filepath.Join(dir, "src", "train.py"), []byte("print('x')"), 0o644)) - commit := initGitRepo(t, dir) - - b := bundleWithCodeSource(t, dir, "src") - - cacheKey := computeSnapshotCacheKey(commit, nil) - archiveName := "src_" + cacheKey[:16] + ".tar.gz" - - f := newRecordingFiler() - f.exists[archiveName] = true - - diags := bundle.Apply(t.Context(), b, &packageAndUpload{client: f}) - require.Empty(t, diags) - - assert.Empty(t, f.written, "existing git-archive snapshot must not be re-uploaded") - assert.Equal(t, codeSnapshotDir+"/"+archiveName, codeSourcePath(b)) -} - -func TestPackageAndUploadGitArchiveUploadsWhenAbsent(t *testing.T) { - dir := t.TempDir() - require.NoError(t, os.MkdirAll(filepath.Join(dir, "src"), 0o755)) - require.NoError(t, os.WriteFile(filepath.Join(dir, "src", "train.py"), []byte("print('x')"), 0o644)) - commit := initGitRepo(t, dir) - - b := bundleWithCodeSource(t, dir, "src") - cacheKey := computeSnapshotCacheKey(commit, nil) - archiveName := "src_" + cacheKey[:16] + ".tar.gz" - - f := newRecordingFiler() - diags := bundle.Apply(t.Context(), b, &packageAndUpload{client: f}) - require.Empty(t, diags) - - require.Contains(t, f.written, archiveName, "git-archive snapshot must be uploaded when not cached") - assert.Equal(t, codeSnapshotDir+"/"+archiveName, codeSourcePath(b)) -} diff --git a/bundle/config/mutator/aicode/snapshot_cachekey.go b/bundle/config/mutator/aicode/snapshot_cachekey.go deleted file mode 100644 index 7512661beee..00000000000 --- a/bundle/config/mutator/aicode/snapshot_cachekey.go +++ /dev/null @@ -1,32 +0,0 @@ -package aicode - -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. Ported from the Python CLI's -// cli/utils/snapshot.py so cache keys are byte-identical across the two CLIs. -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[:]) -} diff --git a/bundle/config/mutator/aicode/snapshot_cachekey_test.go b/bundle/config/mutator/aicode/snapshot_cachekey_test.go deleted file mode 100644 index e89ffc01a6f..00000000000 --- a/bundle/config/mutator/aicode/snapshot_cachekey_test.go +++ /dev/null @@ -1,33 +0,0 @@ -package aicode - -import ( - "encoding/json" - "os" - "testing" - - "github.com/stretchr/testify/assert" - "github.com/stretchr/testify/require" -) - -// TestComputeSnapshotCacheKeyGolden checks parity with the Python air CLI: the -// keys in testdata/cache_keys.json are produced by the Python implementation, so -// matching them proves the Go port hashes identical material. -func TestComputeSnapshotCacheKeyGolden(t *testing.T) { - data, err := os.ReadFile("testdata/cache_keys.json") - require.NoError(t, err) - - var cases []struct { - Name string `json:"name"` - CommitSHA string `json:"commit_sha"` - IncludePaths []string `json:"include_paths"` - CacheKey string `json:"cache_key"` - } - 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)) - }) - } -} diff --git a/bundle/config/mutator/aicode/snapshot_git.go b/bundle/config/mutator/aicode/snapshot_git.go deleted file mode 100644 index 94f370dfb67..00000000000 --- a/bundle/config/mutator/aicode/snapshot_git.go +++ /dev/null @@ -1,77 +0,0 @@ -package aicode - -import ( - "bytes" - "context" - "fmt" - "os/exec" - "strings" -) - -// Local, no-network git introspection for the code_source snapshot, ported from -// the Python CLI's cli/utils/git_state.py and PR #5897's snapshot_git.go. Only the -// subset the DABs resolve path needs is kept: is-repo, HEAD SHA, and dirty check. -// The remote-fetch helpers and the git_state.json/git_diff.patch provenance -// sidecars are intentionally not ported — the SDK ai_runtime_task has nowhere to -// attach sidecar paths, and the snapshot archives the local copy only. - -// gitRepo runs git subcommands scoped to one repository via `git -C`. Arguments -// are passed as a slice, never a shell string, so path values can't inject. -type gitRepo struct { - path string -} - -func newGitRepo(path string) gitRepo { - return gitRepo{path: path} -} - -// run executes `git -C ` and returns trimmed stdout; a non-zero -// exit wraps stderr. -func (g gitRepo) run(ctx context.Context, args ...string) (string, error) { - full := append([]string{"-C", g.path}, args...) - cmd := exec.CommandContext(ctx, "git", full...) - - var stdout, stderr bytes.Buffer - cmd.Stdout = &stdout - cmd.Stderr = &stderr - if err := cmd.Run(); err != nil { - msg := strings.TrimSpace(stderr.String()) - if msg == "" { - return "", fmt.Errorf("git %s: %w", strings.Join(args, " "), err) - } - return "", fmt.Errorf("git %s: %w: %s", strings.Join(args, " "), err, msg) - } - return strings.TrimSpace(stdout.String()), nil -} - -// isRepository reports whether the path is inside a git work tree. Using -// `rev-parse --is-inside-work-tree` (not a .git lookup) means a subdirectory of a -// repo counts — the common case when code_source_path is a subfolder of a monorepo. -func (g gitRepo) isRepository(ctx context.Context) bool { - out, err := g.run(ctx, "rev-parse", "--is-inside-work-tree") - if err != nil { - return false - } - return out == "true" -} - -// headSHA returns the current HEAD commit SHA. -func (g gitRepo) headSHA(ctx context.Context) (string, error) { - return g.run(ctx, "rev-parse", "HEAD") -} - -// hasUncommittedChanges reports whether there are staged or unstaged changes under -// the repo subtree. The `-- .` pathspec scopes the check so a subfolder snapshot -// considers only changes that could land in it. -func (g gitRepo) hasUncommittedChanges(ctx context.Context) (bool, error) { - out, err := g.run(ctx, "status", "--porcelain", "--", ".") - if err != nil { - return false, err - } - return out != "", nil -} - -// shortSHA abbreviates a commit SHA to 8 chars for log/error messages. -func shortSHA(sha string) string { - return sha[:min(len(sha), 8)] -} diff --git a/bundle/config/mutator/aicode/snapshot_package.go b/bundle/config/mutator/aicode/snapshot_package.go index 69e7fa18165..28b691a07b7 100644 --- a/bundle/config/mutator/aicode/snapshot_package.go +++ b/bundle/config/mutator/aicode/snapshot_package.go @@ -10,39 +10,16 @@ import ( "strings" ) -// Tar builders ported from PR #5897 (in turn from cli/utils/snapshot.py). Both -// shell out (git archive / tar) for parity and to reuse git's/tar's symlink, -// gitignore, and AppleDouble handling. The tarball's top-level dir name is -// load-bearing — the remote entry_script extracts to /databricks/code_source/ — -// 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 { - // Single git invocation writes the gzipped tar with the desired prefix; no - // extract/repack. `git -C ` scopes the archive to the code directory's - // subtree even when it is a subfolder of a larger repo. - args := []string{ - "archive", - "--format=tar.gz", - "--prefix=" + directoryName + "/", - "-o", outputTarball, - commitSHA, - } - args = append(args, includePaths...) - if _, err := git.run(ctx, args...); err != nil { - return fmt.Errorf("failed to create git archive: %w", err) - } - return nil -} +// Tarball builder ported from PR #5897 (in turn from cli/utils/snapshot.py). It +// shells out to `tar` to reuse tar's symlink, gitignore, and AppleDouble handling. +// The tarball's top-level dir name is load-bearing — the remote entry_script +// extracts to /databricks/code_source/ — so the `-C parent dir` form preserves it. // createPlainTarball writes a gzipped tar of repoPath's working tree to // outputTarball via `tar`. The archive preserves repoPath's directory name as the -// top-level entry. When includePaths is set, only those paths (nested under the -// directory name) are archived. .git and macOS AppleDouble files are always -// excluded; a .gitignore at repoPath is honored. -func createPlainTarball(ctx context.Context, repoPath, outputTarball string, includePaths []string) error { +// top-level entry. .git and macOS AppleDouble files are always excluded; a +// .gitignore at repoPath is honored. +func createPlainTarball(ctx context.Context, repoPath, outputTarball string) error { dirName := filepath.Base(repoPath) parent := filepath.Dir(repoPath) @@ -69,16 +46,8 @@ func createPlainTarball(ctx context.Context, repoPath, outputTarball string, inc } } - // Archive from the parent so the directory name is preserved; with include_paths, - // prefix each so entries nest under it (matching git archive --prefix). - args = append(args, "-C", parent) - if len(includePaths) > 0 { - for _, p := range includePaths { - args = append(args, dirName+"/"+p) - } - } else { - args = append(args, dirName) - } + // Archive from the parent so the directory name is preserved as the top-level entry. + args = append(args, "-C", parent, dirName) cmd := exec.CommandContext(ctx, "tar", args...) var stderr bytes.Buffer diff --git a/bundle/config/mutator/aicode/snapshot_package_test.go b/bundle/config/mutator/aicode/snapshot_package_test.go index 56252a2e40e..b78d24817cf 100644 --- a/bundle/config/mutator/aicode/snapshot_package_test.go +++ b/bundle/config/mutator/aicode/snapshot_package_test.go @@ -12,7 +12,7 @@ import ( "github.com/stretchr/testify/require" ) -// tarEntryNames returns the sorted set of entry names in a gzipped tarball file. +// tarEntryNames returns the set of entry names in a gzipped tarball file. func tarEntryNames(t *testing.T, tarballPath string) map[string]bool { t.Helper() f, err := os.Open(tarballPath) @@ -33,7 +33,7 @@ func tarEntryNames(t *testing.T, tarballPath string) map[string]bool { return names } -func TestCreatePlainTarballExcludesGitAndGitignored(t *testing.T) { +func TestCreatePlainTarballPrefixesExcludesGitAndGitignored(t *testing.T) { dir := t.TempDir() code := filepath.Join(dir, "mycode") require.NoError(t, os.MkdirAll(filepath.Join(code, "pkg"), 0o755)) @@ -50,9 +50,11 @@ func TestCreatePlainTarballExcludesGitAndGitignored(t *testing.T) { write("._resource_fork", "apple double") out := filepath.Join(dir, "out.tar.gz") - require.NoError(t, createPlainTarball(t.Context(), code, out, nil)) + require.NoError(t, createPlainTarball(t.Context(), code, out)) names := tarEntryNames(t, out) + // Entries are prefixed with the code dir basename (the runtime extracts to + // /databricks/code_source/). assert.True(t, names["mycode/train.py"]) assert.True(t, names["mycode/pkg/util.py"]) assert.True(t, names["mycode/.gitignore"]) @@ -61,23 +63,3 @@ func TestCreatePlainTarballExcludesGitAndGitignored(t *testing.T) { assert.False(t, names["mycode/.git/config"], ".git must never be archived") assert.False(t, names["mycode/._resource_fork"], "AppleDouble metadata must be excluded") } - -func TestCreateGitArchiveSnapshotPrefixesAndScopesToSubdir(t *testing.T) { - // Repo with two top-level dirs; archiving the "src" subdir must include only - // src, prefixed by its basename (the runtime extracts to code_source/). - repo := t.TempDir() - require.NoError(t, os.MkdirAll(filepath.Join(repo, "src"), 0o755)) - require.NoError(t, os.MkdirAll(filepath.Join(repo, "other"), 0o755)) - require.NoError(t, os.WriteFile(filepath.Join(repo, "src", "a.py"), []byte("a"), 0o644)) - require.NoError(t, os.WriteFile(filepath.Join(repo, "other", "b.py"), []byte("b"), 0o644)) - commit := initGitRepo(t, repo) - - src := filepath.Join(repo, "src") - out := filepath.Join(t.TempDir(), "out.tar.gz") - require.NoError(t, createGitArchiveSnapshot(t.Context(), newGitRepo(src), commit, out, "src", nil)) - - names := tarEntryNames(t, out) - assert.True(t, names["src/a.py"]) - assert.False(t, names["other/b.py"], "sibling dir must not be in a src-scoped archive") - assert.False(t, names["src/../other/b.py"]) -} diff --git a/bundle/config/mutator/aicode/snapshot_resolve.go b/bundle/config/mutator/aicode/snapshot_resolve.go deleted file mode 100644 index b58934f964c..00000000000 --- a/bundle/config/mutator/aicode/snapshot_resolve.go +++ /dev/null @@ -1,56 +0,0 @@ -package aicode - -import "context" - -// snapshotMode is how the snapshot tarball is produced. -type snapshotMode int - -const ( - // modeGitArchive packages a pinned commit via `git archive`. The commit is - // deterministic, so the tarball is cacheable by (commit, include_paths). - modeGitArchive snapshotMode = iota - // modePlainTar packages the working tree (including uncommitted changes) via - // `tar`. Not cacheable — working-tree content isn't pinned to a SHA. - modePlainTar -) - -// snapshotPlan is the outcome of resolving how to package a snapshot: the mode and -// the commit SHA to archive (git_archive only; empty for plain_tar). includePaths is -// carried through to packaging and the cache key. -// -// The plan mirrors PR #5897's richer type so this package can later grow git-ref and -// include_paths selection once the ai_runtime_task schema can carry them; for now the -// DABs code_source_path is a plain directory, so the decision reduces to -// "committed git repo → archive HEAD, otherwise tar the working tree". -type snapshotPlan struct { - mode snapshotMode - commitSHA string - includePaths []string -} - -// resolveSnapshotPlan decides how to package a local code_source directory: -// - a git work tree with no uncommitted changes → git_archive of HEAD (cacheable); -// - otherwise (non-git dir, or a dirty tree) → plain_tar of the working tree. -// -// Archiving HEAD when clean lets unchanged code hit the upload cache across deploys; -// falling back to plain_tar when dirty ensures local edits are actually deployed -// rather than silently dropped. -func resolveSnapshotPlan(ctx context.Context, git gitRepo) (snapshotPlan, error) { - if !git.isRepository(ctx) { - return snapshotPlan{mode: modePlainTar}, nil - } - - dirty, err := git.hasUncommittedChanges(ctx) - if err != nil { - return snapshotPlan{}, err - } - if dirty { - return snapshotPlan{mode: modePlainTar}, nil - } - - sha, err := git.headSHA(ctx) - if err != nil { - return snapshotPlan{}, err - } - return snapshotPlan{mode: modeGitArchive, commitSHA: sha}, nil -} diff --git a/bundle/config/mutator/aicode/snapshot_test.go b/bundle/config/mutator/aicode/snapshot_test.go deleted file mode 100644 index 4b6f0fc9135..00000000000 --- a/bundle/config/mutator/aicode/snapshot_test.go +++ /dev/null @@ -1,72 +0,0 @@ -package aicode - -import ( - "os" - "os/exec" - "path/filepath" - "testing" - - "github.com/stretchr/testify/require" -) - -// initGitRepo initializes a deterministic git repo at dir, commits everything in -// it, and returns the HEAD commit SHA. Author/committer dates are pinned so the -// resulting SHA (and thus any git-archive cache key) is stable across runs. -func initGitRepo(t *testing.T, dir string) string { - t.Helper() - run := func(env []string, args ...string) string { - cmd := exec.Command("git", append([]string{"-C", dir}, args...)...) - cmd.Env = append(os.Environ(), env...) - out, err := cmd.CombinedOutput() - require.NoErrorf(t, err, "git %v: %s", args, out) - return string(out) - } - dates := []string{ - "GIT_AUTHOR_DATE=2020-01-01T00:00:00Z", - "GIT_COMMITTER_DATE=2020-01-01T00:00:00Z", - "GIT_AUTHOR_NAME=Tester", "GIT_AUTHOR_EMAIL=tester@databricks.com", - "GIT_COMMITTER_NAME=Tester", "GIT_COMMITTER_EMAIL=tester@databricks.com", - } - run(nil, "init", "-qb", "main") - run(nil, "config", "core.hooksPath", "no-hooks") - run(nil, "add", "-A") - run(dates, "commit", "-qm", "init") - - cmd := exec.Command("git", "-C", dir, "rev-parse", "HEAD") - out, err := cmd.Output() - require.NoError(t, err) - return string(out[:40]) -} - -func TestResolveSnapshotPlanCleanGitRepoUsesGitArchive(t *testing.T) { - dir := t.TempDir() - require.NoError(t, os.WriteFile(filepath.Join(dir, "a.py"), []byte("x"), 0o644)) - commit := initGitRepo(t, dir) - - plan, err := resolveSnapshotPlan(t.Context(), newGitRepo(dir)) - require.NoError(t, err) - require.Equal(t, modeGitArchive, plan.mode) - require.Equal(t, commit, plan.commitSHA) -} - -func TestResolveSnapshotPlanDirtyGitRepoUsesPlainTar(t *testing.T) { - dir := t.TempDir() - require.NoError(t, os.WriteFile(filepath.Join(dir, "a.py"), []byte("x"), 0o644)) - initGitRepo(t, dir) - // Introduce an uncommitted change. - require.NoError(t, os.WriteFile(filepath.Join(dir, "a.py"), []byte("changed"), 0o644)) - - plan, err := resolveSnapshotPlan(t.Context(), newGitRepo(dir)) - require.NoError(t, err) - require.Equal(t, modePlainTar, plan.mode) - require.Empty(t, plan.commitSHA) -} - -func TestResolveSnapshotPlanNonGitUsesPlainTar(t *testing.T) { - dir := t.TempDir() - require.NoError(t, os.WriteFile(filepath.Join(dir, "a.py"), []byte("x"), 0o644)) - - plan, err := resolveSnapshotPlan(t.Context(), newGitRepo(dir)) - require.NoError(t, err) - require.Equal(t, modePlainTar, plan.mode) -} diff --git a/bundle/config/mutator/aicode/testdata/cache_keys.json b/bundle/config/mutator/aicode/testdata/cache_keys.json deleted file mode 100644 index 06673499109..00000000000 --- a/bundle/config/mutator/aicode/testdata/cache_keys.json +++ /dev/null @@ -1,77 +0,0 @@ -[ - { - "name": "commit_no_paths", - "commit_sha": "a3492b801c0ffee00000000000000000000dead", - "include_paths": null, - "cache_key": "8d7bc445ac83dfe432a353718ae8b1ae40eb00c860352662e6df96b7fdf862a5" - }, - { - "name": "commit_empty_paths_none", - "commit_sha": "0000000000000000000000000000000000000000", - "include_paths": null, - "cache_key": "a07e0fa4d38f3c2d08c3ff6bbff0158cea01796aaaf5368b0d0e247cebd27c80" - }, - { - "name": "single_path", - "commit_sha": "a3492b801c0ffee00000000000000000000dead", - "include_paths": [ - "research" - ], - "cache_key": "97bd562ad591af23b9a05651d88c01abbbcc24632e9a8f2534230a61492b9e0a" - }, - { - "name": "multi_path_sorted", - "commit_sha": "a3492b801c0ffee00000000000000000000dead", - "include_paths": [ - "a", - "b", - "c" - ], - "cache_key": "17d021577e3615e9a603a2f701f918d74ad16218efc4b2950ee20f6a18b1a174" - }, - { - "name": "multi_path_unsorted", - "commit_sha": "a3492b801c0ffee00000000000000000000dead", - "include_paths": [ - "c", - "a", - "b" - ], - "cache_key": "17d021577e3615e9a603a2f701f918d74ad16218efc4b2950ee20f6a18b1a174" - }, - { - "name": "path_whitespace", - "commit_sha": "a3492b801c0ffee00000000000000000000dead", - "include_paths": [ - " research ", - " data " - ], - "cache_key": "c6a6fa1d866ac90a5f371610ee08efa8aeb061a43b8a4acae30765db21f4155b" - }, - { - "name": "nested_paths", - "commit_sha": "deadbeefdeadbeefdeadbeefdeadbeefdeadbeef", - "include_paths": [ - "src/models", - "src/data", - "README.md" - ], - "cache_key": "2b18452b1019186a8e5b34af540a0545f96f24bc61771f58fa71fc218cb0275c" - }, - { - "name": "short_sha", - "commit_sha": "a3492b8", - "include_paths": null, - "cache_key": "16be4d7d83ecd01b8d3d274dafb07b3bccc72cf7f7299578e634cd01e1825345" - }, - { - "name": "dup_paths", - "commit_sha": "a3492b801c0ffee00000000000000000000dead", - "include_paths": [ - "x", - "x", - "y" - ], - "cache_key": "407a5d8c3a0e9438d5abf653c4d3a43928c300b5630e2f83f435661b9fbd599e" - } -] From dc68a49c4a88a0964e44699410e8a4f4a5f684c9 Mon Sep 17 00:00:00 2001 From: "ben.hansen" Date: Tue, 14 Jul 2026 04:17:03 +0000 Subject: [PATCH 4/5] Make code_source snapshot reproducible and gitignore-aware Replace the shell-out `tar -czf` (which read only the top-level .gitignore via a lossy --exclude translation) with an in-process archive/tar+gzip builder over libs/git.NewFileSet. This honors nested/parent/global gitignore the same way bundle file sync does, so a large tree only packages the tracked payload instead of venvs, caches, and build outputs. Entries use a fixed mtime and are sorted, making the archive reproducible; it is named by its SHA-256 (_.tar.gz) and the upload is skipped when a content-identical archive already exists. Unchanged code no longer re-uploads or accumulates snapshots. The reproducibility technique mirrors bundle/deploy/snapshot/path.go. Co-authored-by: Isaac --- NEXT_CHANGELOG.md | 2 +- .../local_code_source/output.txt | 9 + .../ai_runtime_task/local_code_source/script | 11 +- .../local_code_source/test.toml | 7 +- .../config/mutator/aicode/package_upload.go | 65 ++++--- .../mutator/aicode/package_upload_test.go | 39 +++- .../config/mutator/aicode/snapshot_package.go | 174 ++++++++++-------- .../mutator/aicode/snapshot_package_test.go | 135 ++++++++++---- 8 files changed, 289 insertions(+), 153 deletions(-) diff --git a/NEXT_CHANGELOG.md b/NEXT_CHANGELOG.md index 0da6ef56c59..a618f800931 100644 --- a/NEXT_CHANGELOG.md +++ b/NEXT_CHANGELOG.md @@ -8,7 +8,7 @@ ### Bundles -* An `ai_runtime_task.code_source_path` that points at a local directory is now packaged into a tarball during `bundle deploy`, uploaded to the user's workspace code snapshot directory, and rewritten to the uploaded path, mirroring the AIR CLI's `code_source` behavior. +* An `ai_runtime_task.code_source_path` that points at a local directory is now packaged into a gitignore-aware tarball during `bundle deploy`, uploaded to the user's workspace code snapshot directory, and rewritten to the uploaded path, mirroring the AIR CLI's `code_source` behavior. The archive is content-addressed, so re-deploying unchanged code skips the re-upload. ### Dependency updates diff --git a/acceptance/bundle/ai_runtime_task/local_code_source/output.txt b/acceptance/bundle/ai_runtime_task/local_code_source/output.txt index 1472320d49c..f796067f776 100644 --- a/acceptance/bundle/ai_runtime_task/local_code_source/output.txt +++ b/acceptance/bundle/ai_runtime_task/local_code_source/output.txt @@ -12,3 +12,12 @@ Deployment complete! === the created job's code_source_path points at the uploaded archive /Workspace/Users/[USERNAME]/.air/repo_snapshots/src/src_[SNAPSHOT].tar.gz + +=== 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 diff --git a/acceptance/bundle/ai_runtime_task/local_code_source/script b/acceptance/bundle/ai_runtime_task/local_code_source/script index 5d729a8e472..4dd3edf4a3e 100644 --- a/acceptance/bundle/ai_runtime_task/local_code_source/script +++ b/acceptance/bundle/ai_runtime_task/local_code_source/script @@ -1,5 +1,5 @@ -# A local code_source_path is packaged into a tarball and uploaded to the user's -# repo_snapshots directory; the archive is timestamp-named. +# A local code_source_path is packaged into a content-addressed tarball +# (_.tar.gz) and uploaded to the user's repo_snapshots directory. title "deploy packages the local code source and rewrites code_source_path\n" trace $CLI bundle deploy @@ -13,3 +13,10 @@ title "the created job's code_source_path points at the uploaded archive\n" jq -r 'select(.path == "/api/2.2/jobs/create") | .body.tasks[0].ai_runtime_task.code_source_path' out.requests.txt 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 diff --git a/acceptance/bundle/ai_runtime_task/local_code_source/test.toml b/acceptance/bundle/ai_runtime_task/local_code_source/test.toml index 8ef47fe2dc3..6e4639a532a 100644 --- a/acceptance/bundle/ai_runtime_task/local_code_source/test.toml +++ b/acceptance/bundle/ai_runtime_task/local_code_source/test.toml @@ -4,8 +4,9 @@ Ignore = [ '.databricks', ] -# A non-git code_source uses plain_tar, whose archive is timestamp-named -# (_YYYYMMDD_HHMMSS.tar.gz). Collapse it to a stable token. +# The archive is content-addressed: _.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_\d{8}_\d{6}\.tar\.gz' +Old = 'src_[0-9a-f]{16}\.tar\.gz' New = 'src_[SNAPSHOT].tar.gz' diff --git a/bundle/config/mutator/aicode/package_upload.go b/bundle/config/mutator/aicode/package_upload.go index e54a1db70d8..04f2adf78ee 100644 --- a/bundle/config/mutator/aicode/package_upload.go +++ b/bundle/config/mutator/aicode/package_upload.go @@ -5,29 +5,33 @@ // volume path to an uploaded code archive; its doc comment states that the CLI // is responsible for packaging the user's local code directory into that // archive. This mutator implements that contract for DABs: when a user points -// code_source_path at a local directory, it packages the directory into a plain -// tarball (.git and gitignored files excluded), uploads the archive to the user's -// workspace code snapshot directory, and rewrites the field to the resulting -// remote path so the deployed job runs against the uploaded code. Values that are -// already remote are left untouched. +// code_source_path at a local directory, it packages the directory into a +// reproducible tarball (.git and gitignored files excluded), uploads the archive +// to the user's workspace code snapshot directory, and rewrites the field to the +// resulting remote path so the deployed job runs against the uploaded code. Values +// that are already remote are left untouched. // -// The tarball builder is ported from PR #5897 (see snapshot_package.go). +// The archive is content-addressed: its name embeds the SHA-256 of the +// (reproducible) tarball, so an unchanged code directory resolves to the same +// remote path across deploys and re-uploads are skipped (see snapshot_package.go). package aicode import ( + "bytes" "context" "errors" "fmt" - "os" + "io/fs" "path" "path/filepath" - "time" "github.com/databricks/cli/bundle" "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/log" + "github.com/databricks/cli/libs/vfs" ) // codeSourcePatterns are the config locations of an AI Runtime task's @@ -83,17 +87,12 @@ func (m *packageAndUpload) Apply(ctx context.Context, b *bundle.Bundle) diag.Dia return diags.Extend(diag.FromErr(err)) } - stagingDir, err := b.LocalStateDir(ctx, "ai_code_source") - if err != nil { - return diags.Extend(diag.FromErr(err)) - } - // remotePaths maps each config location to the remote archive path it should // point to after upload. Built outside the Mutate closure so upload failures // are reported before any config is rewritten. remotePaths := make(map[string]string, len(sources)) for _, cs := range sources { - remote, err := m.packageOne(ctx, b, cs, stagingDir, userDir) + remote, err := m.packageOne(ctx, b, cs, userDir) if err != nil { diags = diags.Extend(diag.FromErr(err)) return diags @@ -124,11 +123,11 @@ func (m *packageAndUpload) Apply(ctx context.Context, b *bundle.Bundle) diag.Dia // the start of every deploy. const repoSnapshotsSubdir = ".air/repo_snapshots" -// packageOne packages the local directory for a single code source into a plain -// tarball, uploads it to the user's repo_snapshots dir, and returns the remote -// path the config should point to. The archive is timestamp-named so concurrent -// deploys do not clobber each other. -func (m *packageAndUpload) packageOne(ctx context.Context, b *bundle.Bundle, cs codeSource, stagingDir, userDir string) (string, error) { +// packageOne packages the local directory for a single code source into a +// reproducible tarball, uploads it to the user's repo_snapshots dir (skipping the +// upload when a content-identical archive already exists there), and returns the +// remote path the config should point to. +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) @@ -142,18 +141,30 @@ func (m *packageAndUpload) packageOne(ctx context.Context, b *bundle.Bundle, cs } } - archiveName := fmt.Sprintf("%s_%s.tar.gz", dirName, time.Now().UTC().Format("20060102_150405")) - localArchive := filepath.Join(stagingDir, archiveName) - defer os.Remove(localArchive) - - if err := createPlainTarball(ctx, localDir, localArchive); err != nil { + // Build the archive in memory so its content hash can name the upload; the hash + // is computed while gzipping, so this adds no extra pass over the files. + var buf bytes.Buffer + sha, err := buildCodeSnapshot(ctx, vfs.MustNew(localDir), dirName, &buf) + if err != nil { return "", fmt.Errorf("failed to package code_source_path %q: %w", cs.value, err) } - if err := libraries.UploadFile(ctx, localArchive, client); err != nil { - return "", err + + archiveName := fmt.Sprintf("%s_%s.tar.gz", dirName, sha[:16]) + remotePath := path.Join(uploadPath, archiveName) + + // The archive is reproducible, so a matching name means identical content is + // already uploaded: skip the upload and just point the config at it. + 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) } - return path.Join(uploadPath, archiveName), nil + 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 } // userWorkspaceHome returns the current user's workspace home directory diff --git a/bundle/config/mutator/aicode/package_upload_test.go b/bundle/config/mutator/aicode/package_upload_test.go index 1e6352513d7..cc998a0fcb7 100644 --- a/bundle/config/mutator/aicode/package_upload_test.go +++ b/bundle/config/mutator/aicode/package_upload_test.go @@ -24,14 +24,16 @@ import ( // tests, derived from the fixture user below. const codeSnapshotDir = "/Workspace/Users/me@databricks.com/.air/repo_snapshots/src" -// recordingFiler records every Write, so a test can capture the uploaded archive. +// recordingFiler records every Write and reports the names in exists as already +// present, so a test can capture uploads and simulate a content cache hit. type recordingFiler struct { filer.Filer written map[string][]byte + exists map[string]bool } func newRecordingFiler() *recordingFiler { - return &recordingFiler{written: map[string][]byte{}} + return &recordingFiler{written: map[string][]byte{}, exists: map[string]bool{}} } func (f *recordingFiler) Write(ctx context.Context, p string, r io.Reader, _ ...filer.WriteMode) error { @@ -44,9 +46,14 @@ func (f *recordingFiler) Write(ctx context.Context, p string, r io.Reader, _ ... } func (f *recordingFiler) Stat(ctx context.Context, p string) (fs.FileInfo, error) { + if f.exists[p] { + return fakeFileInfo{}, nil + } return nil, fs.ErrNotExist } +type fakeFileInfo struct{ fs.FileInfo } + // bundleWithCodeSource builds a bundle rooted at dir whose single AI Runtime task // points at codeSourcePath. The caller populates dir before applying the mutator. func bundleWithCodeSource(t *testing.T, dir, codeSourcePath string) *bundle.Bundle { @@ -99,11 +106,12 @@ func TestPackageAndUploadRewritesLocalCodeSource(t *testing.T) { diags := bundle.Apply(t.Context(), b, &packageAndUpload{client: f}) require.Empty(t, diags) - assert.Regexp(t, `^`+codeSnapshotDir+`/src_\d{8}_\d{6}\.tar\.gz$`, codeSourcePath(b)) + // The archive is named by the content hash of the reproducible tarball. + assert.Regexp(t, `^`+codeSnapshotDir+`/src_[0-9a-f]{16}\.tar\.gz$`, codeSourcePath(b)) require.Len(t, f.written, 1) for name := range f.written { - assert.Regexp(t, `^src_\d{8}_\d{6}\.tar\.gz$`, name) + assert.Regexp(t, `^src_[0-9a-f]{16}\.tar\.gz$`, name) } } @@ -117,3 +125,26 @@ func TestPackageAndUploadSkipsRemoteCodeSource(t *testing.T) { assert.Equal(t, "/Volumes/main/default/code/existing.tar.gz", codeSourcePath(b)) assert.Empty(t, f.written, "remote code_source_path must not trigger an upload") } + +func TestPackageAndUploadSkipsUploadWhenArchiveExists(t *testing.T) { + b := bundleWithSrc(t, "src") + + // First deploy uploads and records the archive name. + f1 := newRecordingFiler() + require.Empty(t, bundle.Apply(t.Context(), b, &packageAndUpload{client: f1})) + require.Len(t, f1.written, 1) + var archiveName string + for name := range f1.written { + archiveName = name + } + + // Second deploy against a store that already has that archive: no re-upload, + // same rewritten path (the content hash is stable). + b2 := bundleWithCodeSource(t, b.BundleRootPath, "src") + f2 := newRecordingFiler() + f2.exists[archiveName] = true + require.Empty(t, bundle.Apply(t.Context(), b2, &packageAndUpload{client: f2})) + + assert.Empty(t, f2.written, "content-identical archive must not be re-uploaded") + assert.Equal(t, codeSnapshotDir+"/"+archiveName, codeSourcePath(b2)) +} diff --git a/bundle/config/mutator/aicode/snapshot_package.go b/bundle/config/mutator/aicode/snapshot_package.go index 28b691a07b7..1afb7435568 100644 --- a/bundle/config/mutator/aicode/snapshot_package.go +++ b/bundle/config/mutator/aicode/snapshot_package.go @@ -1,103 +1,119 @@ package aicode import ( - "bytes" + "archive/tar" + "compress/gzip" "context" + "crypto/sha256" + "encoding/hex" "fmt" - "os" - "os/exec" - "path/filepath" + "io" + "path" + "slices" "strings" + "time" + + "github.com/databricks/cli/libs/fileset" + "github.com/databricks/cli/libs/git" + "github.com/databricks/cli/libs/vfs" ) -// Tarball builder ported from PR #5897 (in turn from cli/utils/snapshot.py). It -// shells out to `tar` to reuse tar's symlink, gitignore, and AppleDouble handling. -// The tarball's top-level dir name is load-bearing — the remote entry_script -// extracts to /databricks/code_source/ — so the `-C parent dir` form preserves it. +// tarEpoch is a fixed modification time stamped on every tar entry so the archive +// is content-addressed: identical file contents always produce identical bytes +// (and therefore an identical SHA-256), regardless of file mtimes or when the +// archive was built. This is what lets an unchanged code directory resolve to the +// same uploaded filename across deploys and skip re-upload. The technique mirrors +// bundle/deploy/snapshot/path.go (which does the same for the immutable-folder zip). +var tarEpoch = time.Date(2000, 1, 1, 0, 0, 0, 0, time.UTC) -// createPlainTarball writes a gzipped tar of repoPath's working tree to -// outputTarball via `tar`. The archive preserves repoPath's directory name as the -// top-level entry. .git and macOS AppleDouble files are always excluded; a -// .gitignore at repoPath is honored. -func createPlainTarball(ctx context.Context, repoPath, outputTarball string) error { - dirName := filepath.Base(repoPath) - parent := filepath.Dir(repoPath) +// appleDoublePrefix is the basename prefix of macOS AppleDouble metadata files. +// The AIR CLI excludes these; we match it so archives are identical on macOS and Linux. +const appleDoublePrefix = "._" - args := []string{"-czf", outputTarball} +// buildCodeSnapshot writes a reproducible gzipped tarball of the directory at +// codeDir to out and returns its SHA-256 hex digest. Every entry is prefixed with +// prefix (the code directory's basename) so the archive expands to /... — +// matching the runtime's /databricks/code_source/ extraction contract. +// +// The file list is gitignore-aware via [git.NewFileSetAtRoot]: it honors the code +// directory's .gitignore (including nested .gitignore files) and always excludes +// .git and .databricks — the same walker that backs bundle file sync. This is why +// packaging a large tree only archives the tracked payload rather than venvs, +// caches, and build outputs. +func buildCodeSnapshot(ctx context.Context, codeDir vfs.Path, prefix string, out io.Writer) (string, error) { + fsys, err := git.NewFileSetAtRoot(ctx, codeDir) + if err != nil { + return "", err + } + + files, err := fsys.Files() + if err != nil { + return "", err + } - // Exclude macOS AppleDouble files: they sort before the real top-level dir and - // hijack a remote `head -1` parse. No-op on Linux. - args = append(args, "--exclude=._*") + // Sort by relative path so the archive byte stream (and thus its hash) does not + // depend on filesystem walk order. + slices.SortFunc(files, func(a, b fileset.File) int { + return strings.Compare(a.Relative, b.Relative) + }) - // Never ship .git — it is large and unused on the remote. - args = append(args, "--exclude=.git") + hash := sha256.New() + gzw := gzip.NewWriter(io.MultiWriter(out, hash)) + tw := tar.NewWriter(gzw) - // Honor .gitignore if present. - gitignorePath := filepath.Join(repoPath, ".gitignore") - if patterns, err := parseGitignore(gitignorePath); err == nil { - for _, p := range patterns { - if strings.Contains(p, "/") { - // Anchor path-relative patterns to the archive root so they don't - // match identically-named paths in subdirectories. - args = append(args, "--exclude="+dirName+"/"+strings.TrimPrefix(p, "/")) - } else { - args = append(args, "--exclude="+p) - } + for _, f := range files { + if err := addFileToArchive(tw, codeDir, f, prefix); err != nil { + return "", err } } - // Archive from the parent so the directory name is preserved as the top-level entry. - args = append(args, "-C", parent, dirName) - - cmd := exec.CommandContext(ctx, "tar", args...) - var stderr bytes.Buffer - cmd.Stderr = &stderr - if err := cmd.Run(); err != nil { - if msg := strings.TrimSpace(stderr.String()); msg != "" { - return fmt.Errorf("failed to create plain tarball: %w: %s", err, msg) - } - return fmt.Errorf("failed to create plain tarball: %w", err) + if err := tw.Close(); err != nil { + return "", err } - return nil + if err := gzw.Close(); err != nil { + return "", err + } + return hex.EncodeToString(hash.Sum(nil)), nil } -// parseGitignore reads a .gitignore and returns tar --exclude patterns. It mirrors -// the Python CLI's lossy normalization so plain-tar snapshots exclude the same set: -// -// - comments (#…) and blank lines are skipped; -// - negation patterns (!…) are unsupported by tar --exclude and skipped; -// - a trailing "/" (directory marker) is stripped; -// - "**" is not a path-separator-agnostic wildcard in tar, so "**/foo" → "foo" -// and "foo/**" → "foo"; a mid-path "**" has no tar equivalent and is skipped. -// -// A missing file returns (nil, error); callers treat any error as "no patterns". -func parseGitignore(path string) ([]string, error) { - data, err := os.ReadFile(path) +func addFileToArchive(tw *tar.Writer, codeDir vfs.Path, f fileset.File, prefix string) error { + // fileset.File.Relative is slash-separated, so path.Base is correct here. + if strings.HasPrefix(path.Base(f.Relative), appleDoublePrefix) { + return nil + } + + rc, err := codeDir.Open(f.Relative) if err != nil { - return nil, err + return fmt.Errorf("open %s: %w", f.Relative, err) } + defer rc.Close() - var patterns []string - for raw := range strings.SplitSeq(string(data), "\n") { - line := strings.TrimRight(raw, " \t\r") - if line == "" || strings.HasPrefix(line, "#") { - continue - } - if strings.HasPrefix(line, "!") { - continue - } - line = strings.TrimRight(line, "/") - if strings.Contains(line, "**") { - switch { - case strings.HasPrefix(line, "**/"): - line = line[len("**/"):] - case strings.HasSuffix(line, "/**"): - line = line[:len(line)-len("/**")] - default: - continue - } - } - patterns = append(patterns, line) + info, err := rc.Stat() + if err != nil { + return fmt.Errorf("stat %s: %w", f.Relative, err) } - return patterns, nil + + // Only regular files are archived. The gitignore-aware walker never yields + // directories, and symlinks inside a code snapshot are out of scope. + if !info.Mode().IsRegular() { + return nil + } + + hdr := &tar.Header{ + Typeflag: tar.TypeReg, + Name: path.Join(prefix, f.Relative), + Size: info.Size(), + // Normalize permissions and zero the mtime so the archive is reproducible + // across machines. The runtime invokes code via an interpreter, not by + // executing files from the snapshot, so execute bits are not preserved. + Mode: 0o644, + ModTime: tarEpoch, + } + if err := tw.WriteHeader(hdr); err != nil { + return fmt.Errorf("tar header for %s: %w", f.Relative, err) + } + if _, err := io.Copy(tw, rc); err != nil { + return fmt.Errorf("write %s: %w", f.Relative, err) + } + return nil } diff --git a/bundle/config/mutator/aicode/snapshot_package_test.go b/bundle/config/mutator/aicode/snapshot_package_test.go index b78d24817cf..4d2f86d496f 100644 --- a/bundle/config/mutator/aicode/snapshot_package_test.go +++ b/bundle/config/mutator/aicode/snapshot_package_test.go @@ -2,64 +2,125 @@ package aicode import ( "archive/tar" + "bytes" "compress/gzip" "io" "os" "path/filepath" "testing" + "github.com/databricks/cli/libs/vfs" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) -// tarEntryNames returns the set of entry names in a gzipped tarball file. -func tarEntryNames(t *testing.T, tarballPath string) map[string]bool { +// writeTree materializes files (relative slash path -> content) under a fresh +// temp dir and returns a vfs.Path rooted at it. +func writeTree(t *testing.T, files map[string]string) vfs.Path { t.Helper() - f, err := os.Open(tarballPath) - require.NoError(t, err) - defer f.Close() - gzr, err := gzip.NewReader(f) + dir := t.TempDir() + for name, content := range files { + p := filepath.Join(dir, filepath.FromSlash(name)) + require.NoError(t, os.MkdirAll(filepath.Dir(p), 0o755)) + require.NoError(t, os.WriteFile(p, []byte(content), 0o644)) + } + return vfs.MustNew(dir) +} + +// tarEntries reads a gzipped tarball and returns entry name -> content. +func tarEntries(t *testing.T, b []byte) map[string]string { + t.Helper() + gzr, err := gzip.NewReader(bytes.NewReader(b)) require.NoError(t, err) tr := tar.NewReader(gzr) - names := map[string]bool{} + out := map[string]string{} for { hdr, err := tr.Next() if err == io.EOF { break } require.NoError(t, err) - names[hdr.Name] = true + content, err := io.ReadAll(tr) + require.NoError(t, err) + out[hdr.Name] = string(content) } - return names + return out } -func TestCreatePlainTarballPrefixesExcludesGitAndGitignored(t *testing.T) { - dir := t.TempDir() - code := filepath.Join(dir, "mycode") - require.NoError(t, os.MkdirAll(filepath.Join(code, "pkg"), 0o755)) - require.NoError(t, os.MkdirAll(filepath.Join(code, ".git"), 0o755)) - write := func(rel, content string) { - require.NoError(t, os.WriteFile(filepath.Join(code, rel), []byte(content), 0o644)) - } - write("train.py", "print('train')") - write("pkg/util.py", "x = 1") - write(".gitignore", "ignored.txt\n*.log\n") - write("ignored.txt", "nope") - write("debug.log", "nope") - write(".git/config", "[core]") - write("._resource_fork", "apple double") - - out := filepath.Join(dir, "out.tar.gz") - require.NoError(t, createPlainTarball(t.Context(), code, out)) - - names := tarEntryNames(t, out) - // Entries are prefixed with the code dir basename (the runtime extracts to +func TestBuildCodeSnapshotPrefixesAndExcludes(t *testing.T) { + root := writeTree(t, map[string]string{ + "train.py": "print('train')", + "pkg/util.py": "x = 1", + ".gitignore": "secrets.txt\n*.log\n", + "secrets.txt": "nope", + "debug.log": "nope", + "pkg/nested.log": "nope", + ".git/config": "[core]", + "._resource_fork": "apple double", + }) + + var buf bytes.Buffer + sha, err := buildCodeSnapshot(t.Context(), root, "mycode", &buf) + require.NoError(t, err) + require.NotEmpty(t, sha) + + entries := tarEntries(t, buf.Bytes()) + // Entries are prefixed with the code dir basename (runtime extracts to // /databricks/code_source/). - assert.True(t, names["mycode/train.py"]) - assert.True(t, names["mycode/pkg/util.py"]) - assert.True(t, names["mycode/.gitignore"]) - assert.False(t, names["mycode/ignored.txt"], "gitignored file must be excluded") - assert.False(t, names["mycode/debug.log"], "gitignored glob must be excluded") - assert.False(t, names["mycode/.git/config"], ".git must never be archived") - assert.False(t, names["mycode/._resource_fork"], "AppleDouble metadata must be excluded") + assert.Equal(t, "print('train')", entries["mycode/train.py"]) + assert.Equal(t, "x = 1", entries["mycode/pkg/util.py"]) + assert.Contains(t, entries, "mycode/.gitignore") + + assert.NotContains(t, entries, "mycode/secrets.txt", "gitignored file must be excluded") + assert.NotContains(t, entries, "mycode/debug.log", "gitignored glob must be excluded") + assert.NotContains(t, entries, "mycode/pkg/nested.log", "gitignored glob must be excluded in subdirs") + assert.NotContains(t, entries, "mycode/.git/config", ".git must never be archived") + assert.NotContains(t, entries, "mycode/._resource_fork", "AppleDouble metadata must be excluded") +} + +func TestBuildCodeSnapshotHonorsNestedGitignore(t *testing.T) { + // A .gitignore in a subdirectory applies within that subtree — this is the + // case the previous shell-tar implementation missed. + root := writeTree(t, map[string]string{ + "a.py": "a", + "pkg/keep.py": "keep", + "pkg/.gitignore": "drop.py\n", + "pkg/drop.py": "nope", + }) + + var buf bytes.Buffer + _, err := buildCodeSnapshot(t.Context(), root, "code", &buf) + require.NoError(t, err) + + entries := tarEntries(t, buf.Bytes()) + assert.Contains(t, entries, "code/pkg/keep.py") + assert.NotContains(t, entries, "code/pkg/drop.py", "nested .gitignore must be honored") +} + +func TestBuildCodeSnapshotIsReproducible(t *testing.T) { + files := map[string]string{"a.py": "aaa", "sub/b.py": "bbb"} + root1 := writeTree(t, files) + root2 := writeTree(t, files) + + var buf1, buf2 bytes.Buffer + sha1, err := buildCodeSnapshot(t.Context(), root1, "code", &buf1) + require.NoError(t, err) + sha2, err := buildCodeSnapshot(t.Context(), root2, "code", &buf2) + require.NoError(t, err) + + assert.Equal(t, sha1, sha2, "identical content must produce an identical hash") + assert.Equal(t, buf1.Bytes(), buf2.Bytes(), "identical content must produce identical bytes") +} + +func TestBuildCodeSnapshotHashChangesWithContent(t *testing.T) { + root1 := writeTree(t, map[string]string{"main.py": "v1"}) + root2 := writeTree(t, map[string]string{"main.py": "v2"}) + + var buf1, buf2 bytes.Buffer + sha1, err := buildCodeSnapshot(t.Context(), root1, "code", &buf1) + require.NoError(t, err) + sha2, err := buildCodeSnapshot(t.Context(), root2, "code", &buf2) + require.NoError(t, err) + + assert.NotEqual(t, sha1, sha2) } From f908b0e2c7c42adfbc97b2b6c6714bce053e8962 Mon Sep 17 00:00:00 2001 From: "ben.hansen" Date: Tue, 14 Jul 2026 23:45:06 +0000 Subject: [PATCH 5/5] Translate command_path, synthesize requirements.yaml, honor sync globs Make a plain bundle deploy produce a runnable AI Runtime job end to end: - Translate ai_runtime_task.deployments[].command_path to its absolute synced workspace path (via TranslatePaths), matching notebook_path/python_file. The backend rejects a relative command_path. - Record code_source_path in the legacy /Users/... form (no /Workspace prefix) that the AI Runtime snapshot fetcher expects, while still uploading via the /Workspace filer. - Synthesize a requirements.yaml next to command_path from the job's serverless environments[] spec. The entry script reads it from command_path's directory to set up the workload environment; without it the run fails during setup. - Build the code snapshot from the bundle sync file list so it honors .gitignore plus the top-level sync.include/sync.exclude globs, not just a single top-level .gitignore. Verified end to end on a live workspace: deploy of a plain bundle (no hand-authored files) yields a job whose run terminates SUCCESS; sync.exclude and .gitignore entries are absent from the uploaded snapshot. Co-authored-by: Isaac --- NEXT_CHANGELOG.md | 2 +- .../local_code_source/databricks.yml | 5 + .../local_code_source/output.txt | 12 +- .../ai_runtime_task/local_code_source/script | 14 +- .../local_code_source/src/.gitignore | 1 + .../config/mutator/aicode/package_upload.go | 41 ++++- .../mutator/aicode/package_upload_test.go | 118 +++----------- bundle/config/mutator/aicode/requirements.go | 144 ++++++++++++++++++ .../mutator/aicode/requirements_test.go | 36 +++++ .../config/mutator/aicode/snapshot_package.go | 67 ++++---- .../mutator/aicode/snapshot_package_test.go | 67 ++++---- .../config/mutator/paths/job_paths_visitor.go | 5 + bundle/phases/deploy.go | 5 + 13 files changed, 340 insertions(+), 177 deletions(-) create mode 100644 acceptance/bundle/ai_runtime_task/local_code_source/src/.gitignore create mode 100644 bundle/config/mutator/aicode/requirements.go create mode 100644 bundle/config/mutator/aicode/requirements_test.go diff --git a/NEXT_CHANGELOG.md b/NEXT_CHANGELOG.md index a618f800931..ccfd8531bdf 100644 --- a/NEXT_CHANGELOG.md +++ b/NEXT_CHANGELOG.md @@ -8,7 +8,7 @@ ### Bundles -* An `ai_runtime_task.code_source_path` that points at a local directory is now packaged into a gitignore-aware tarball during `bundle deploy`, uploaded to the user's workspace code snapshot directory, and rewritten to the uploaded path, mirroring the AIR CLI's `code_source` behavior. The archive is content-addressed, so re-deploying unchanged code skips the re-upload. +* An `ai_runtime_task.code_source_path` that points at a local directory is now packaged into a tarball during `bundle deploy`, uploaded to the user's workspace code snapshot directory, and rewritten to the uploaded path, mirroring the AIR CLI's `code_source` behavior. The snapshot honors `.gitignore` and the top-level `sync.include`/`sync.exclude` globs, and is content-addressed so re-deploying unchanged code skips the re-upload. The task's `deployments[].command_path` is translated to its workspace path, and a `requirements.yaml` derived from the job's serverless `environments` is written alongside it so the workload's environment is set up correctly. ### Dependency updates diff --git a/acceptance/bundle/ai_runtime_task/local_code_source/databricks.yml b/acceptance/bundle/ai_runtime_task/local_code_source/databricks.yml index c6d37475f73..ff4a7b5dc1b 100644 --- a/acceptance/bundle/ai_runtime_task/local_code_source/databricks.yml +++ b/acceptance/bundle/ai_runtime_task/local_code_source/databricks.yml @@ -1,6 +1,11 @@ bundle: name: ai-runtime-test +sync: + # Exclude a file from both bundle sync and the code snapshot. + exclude: + - "**/*.log" + resources: jobs: train: diff --git a/acceptance/bundle/ai_runtime_task/local_code_source/output.txt b/acceptance/bundle/ai_runtime_task/local_code_source/output.txt index f796067f776..33fb18eba49 100644 --- a/acceptance/bundle/ai_runtime_task/local_code_source/output.txt +++ b/acceptance/bundle/ai_runtime_task/local_code_source/output.txt @@ -1,5 +1,5 @@ -=== deploy packages the local code source and rewrites code_source_path +=== deploy packages the local code source >>> [CLI] bundle deploy Uploading bundle files to /Workspace/Users/[USERNAME]/.bundle/ai-runtime-test/default/files... @@ -10,8 +10,14 @@ 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 -=== the created job's code_source_path points at the uploaded archive -/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 diff --git a/acceptance/bundle/ai_runtime_task/local_code_source/script b/acceptance/bundle/ai_runtime_task/local_code_source/script index 4dd3edf4a3e..d0a9e2ea965 100644 --- a/acceptance/bundle/ai_runtime_task/local_code_source/script +++ b/acceptance/bundle/ai_runtime_task/local_code_source/script @@ -1,7 +1,9 @@ # A local code_source_path is packaged into a content-addressed tarball -# (_.tar.gz) and uploaded to the user's repo_snapshots directory. +# (_.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 and rewrites code_source_path\n" +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" @@ -9,9 +11,15 @@ title "the snapshot tarball is uploaded to the user's repo_snapshots directory\n # 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 "the created job's code_source_path points at the uploaded archive\n" +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" diff --git a/acceptance/bundle/ai_runtime_task/local_code_source/src/.gitignore b/acceptance/bundle/ai_runtime_task/local_code_source/src/.gitignore new file mode 100644 index 00000000000..7a79c3b103e --- /dev/null +++ b/acceptance/bundle/ai_runtime_task/local_code_source/src/.gitignore @@ -0,0 +1 @@ +ignored_by_git.txt diff --git a/bundle/config/mutator/aicode/package_upload.go b/bundle/config/mutator/aicode/package_upload.go index 04f2adf78ee..4124bbfb060 100644 --- a/bundle/config/mutator/aicode/package_upload.go +++ b/bundle/config/mutator/aicode/package_upload.go @@ -24,14 +24,17 @@ import ( "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" - "github.com/databricks/cli/libs/vfs" + libsync "github.com/databricks/cli/libs/sync" ) // codeSourcePatterns are the config locations of an AI Runtime task's @@ -131,10 +134,22 @@ func (m *packageAndUpload) packageOne(ctx context.Context, b *bundle.Bundle, cs localDir := filepath.Join(b.SyncRootPath, filepath.FromSlash(cs.value)) dirName := filepath.Base(localDir) + // relBase is the code directory relative to the sync root, used both to scope the + // sync file list to this directory and to re-base archive entry names 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 { - var err error client, err = filer.NewWorkspaceFilesClient(b.WorkspaceClient(ctx), uploadPath) if err != nil { return "", err @@ -144,13 +159,17 @@ func (m *packageAndUpload) packageOne(ctx context.Context, b *bundle.Bundle, cs // Build the archive in memory so its content hash can name the upload; the hash // is computed while gzipping, so this adds no extra pass over the files. var buf bytes.Buffer - sha, err := buildCodeSnapshot(ctx, vfs.MustNew(localDir), dirName, &buf) + 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]) - remotePath := path.Join(uploadPath, archiveName) + // The AI Runtime snapshot fetcher expects code_source_path in the legacy + // "/Users/..." form (no "/Workspace" prefix), matching the Python air CLI. The + // filer needs the "/Workspace/Users/..." form to upload, so upload to uploadPath + // but record the de-prefixed path on the task. + remotePath := strings.TrimPrefix(path.Join(uploadPath, archiveName), "/Workspace") // The archive is reproducible, so a matching name means identical content is // already uploaded: skip the upload and just point the config at it. @@ -167,6 +186,20 @@ func (m *packageAndUpload) packageOne(ctx context.Context, b *bundle.Bundle, cs return remotePath, nil } +// codeSourceFiles returns the files under the code directory (relBase, relative to +// the sync root) that should go into the snapshot. It reuses the bundle's sync +// options so the file list is filtered exactly like bundle file sync: .gitignore +// aware, plus the top-level sync.include/exclude globs. Scoping Paths to relBase +// restricts the walk (and the returned relative paths) to the code directory. +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 + } + opts.Paths = []string{relBase} + return libsync.GetFileList(ctx, *opts) +} + // userWorkspaceHome returns the current user's workspace home directory // (/Workspace/Users/), the root under which code snapshots are stored. func userWorkspaceHome(b *bundle.Bundle) (string, error) { diff --git a/bundle/config/mutator/aicode/package_upload_test.go b/bundle/config/mutator/aicode/package_upload_test.go index cc998a0fcb7..25192332f81 100644 --- a/bundle/config/mutator/aicode/package_upload_test.go +++ b/bundle/config/mutator/aicode/package_upload_test.go @@ -1,10 +1,6 @@ package aicode import ( - "context" - "io" - "io/fs" - "os" "path/filepath" "testing" @@ -13,49 +9,20 @@ import ( "github.com/databricks/cli/bundle/config/resources" "github.com/databricks/cli/bundle/internal/bundletest" "github.com/databricks/cli/libs/dyn" - "github.com/databricks/cli/libs/filer" "github.com/databricks/databricks-sdk-go/service/iam" "github.com/databricks/databricks-sdk-go/service/jobs" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) -// codeSnapshotDir is the workspace location code snapshots are uploaded to in -// tests, derived from the fixture user below. -const codeSnapshotDir = "/Workspace/Users/me@databricks.com/.air/repo_snapshots/src" - -// recordingFiler records every Write and reports the names in exists as already -// present, so a test can capture uploads and simulate a content cache hit. -type recordingFiler struct { - filer.Filer - written map[string][]byte - exists map[string]bool -} - -func newRecordingFiler() *recordingFiler { - return &recordingFiler{written: map[string][]byte{}, exists: map[string]bool{}} -} - -func (f *recordingFiler) Write(ctx context.Context, p string, r io.Reader, _ ...filer.WriteMode) error { - b, err := io.ReadAll(r) - if err != nil { - return err - } - f.written[p] = b - return nil -} - -func (f *recordingFiler) Stat(ctx context.Context, p string) (fs.FileInfo, error) { - if f.exists[p] { - return fakeFileInfo{}, nil - } - return nil, fs.ErrNotExist -} - -type fakeFileInfo struct{ fs.FileInfo } - // bundleWithCodeSource builds a bundle rooted at dir whose single AI Runtime task -// points at codeSourcePath. The caller populates dir before applying the mutator. +// points at codeSourcePath. +// +// 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 bundleWithCodeSource(t *testing.T, dir, codeSourcePath string) *bundle.Bundle { t.Helper() b := &bundle.Bundle{ @@ -86,65 +53,22 @@ func bundleWithCodeSource(t *testing.T, dir, codeSourcePath string) *bundle.Bund return b } -// bundleWithSrc creates a src/ directory with one file and returns the bundle. -func bundleWithSrc(t *testing.T, codeSourcePath string) *bundle.Bundle { - t.Helper() - dir := t.TempDir() - require.NoError(t, os.MkdirAll(filepath.Join(dir, "src"), 0o755)) - require.NoError(t, os.WriteFile(filepath.Join(dir, "src", "train.py"), []byte("print('x')"), 0o644)) - return bundleWithCodeSource(t, dir, codeSourcePath) -} - -func codeSourcePath(b *bundle.Bundle) string { - return b.Config.Resources.Jobs["train"].Tasks[0].AiRuntimeTask.CodeSourcePath -} - -func TestPackageAndUploadRewritesLocalCodeSource(t *testing.T) { - b := bundleWithSrc(t, "src") - f := newRecordingFiler() - - diags := bundle.Apply(t.Context(), b, &packageAndUpload{client: f}) - require.Empty(t, diags) - - // The archive is named by the content hash of the reproducible tarball. - assert.Regexp(t, `^`+codeSnapshotDir+`/src_[0-9a-f]{16}\.tar\.gz$`, codeSourcePath(b)) - - require.Len(t, f.written, 1) - for name := range f.written { - assert.Regexp(t, `^src_[0-9a-f]{16}\.tar\.gz$`, name) - } -} - -func TestPackageAndUploadSkipsRemoteCodeSource(t *testing.T) { - b := bundleWithSrc(t, "/Volumes/main/default/code/existing.tar.gz") - f := newRecordingFiler() - - diags := bundle.Apply(t.Context(), b, &packageAndUpload{client: f}) +func TestCollectLocalCodeSourcesFindsLocalPath(t *testing.T) { + b := bundleWithCodeSource(t, t.TempDir(), "./src") + sources, diags := collectLocalCodeSources(b) require.Empty(t, diags) - - assert.Equal(t, "/Volumes/main/default/code/existing.tar.gz", codeSourcePath(b)) - assert.Empty(t, f.written, "remote code_source_path must not trigger an upload") + require.Len(t, sources, 1) + assert.Equal(t, "./src", sources[0].value) } -func TestPackageAndUploadSkipsUploadWhenArchiveExists(t *testing.T) { - b := bundleWithSrc(t, "src") - - // First deploy uploads and records the archive name. - f1 := newRecordingFiler() - require.Empty(t, bundle.Apply(t.Context(), b, &packageAndUpload{client: f1})) - require.Len(t, f1.written, 1) - var archiveName string - for name := range f1.written { - archiveName = name +func TestCollectLocalCodeSourcesSkipsRemotePaths(t *testing.T) { + for _, remote := range []string{ + "/Workspace/Users/me/code.tar.gz", + "/Volumes/main/default/code/existing.tar.gz", + } { + b := bundleWithCodeSource(t, t.TempDir(), remote) + sources, diags := collectLocalCodeSources(b) + require.Empty(t, diags) + assert.Empty(t, sources, "remote code_source_path %q must not be collected", remote) } - - // Second deploy against a store that already has that archive: no re-upload, - // same rewritten path (the content hash is stable). - b2 := bundleWithCodeSource(t, b.BundleRootPath, "src") - f2 := newRecordingFiler() - f2.exists[archiveName] = true - require.Empty(t, bundle.Apply(t.Context(), b2, &packageAndUpload{client: f2})) - - assert.Empty(t, f2.written, "content-identical archive must not be re-uploaded") - assert.Equal(t, codeSnapshotDir+"/"+archiveName, codeSourcePath(b2)) } diff --git a/bundle/config/mutator/aicode/requirements.go b/bundle/config/mutator/aicode/requirements.go new file mode 100644 index 00000000000..df2eee3cbfa --- /dev/null +++ b/bundle/config/mutator/aicode/requirements.go @@ -0,0 +1,144 @@ +package aicode + +import ( + "bytes" + "context" + "fmt" + "path" + + "github.com/databricks/cli/bundle" + "github.com/databricks/cli/libs/diag" + "github.com/databricks/cli/libs/filer" + "github.com/databricks/cli/libs/log" + "github.com/databricks/databricks-sdk-go/service/compute" + "github.com/databricks/databricks-sdk-go/service/jobs" + "go.yaml.in/yaml/v3" +) + +// requirementsFileName is the file the AI Runtime entry script reads from the +// command_path directory to install the workload's pip dependencies. The server +// derives its path from command_path, so it must sit next to command.sh. +const requirementsFileName = "requirements.yaml" + +// requirementsSpec is the requirements.yaml shape the AI Runtime consumes: a +// client image version and a list of pip requirement lines. It mirrors what the +// Python air CLI synthesizes. +type requirementsSpec struct { + Version string `yaml:"version,omitempty"` + Dependencies []string `yaml:"dependencies,omitempty"` +} + +// SynthesizeRequirements uploads a requirements.yaml next to each AI Runtime task's +// command_path, derived from the job-level serverless environment referenced by the +// task's environment_key. The AI Runtime entry script reads this file (resolved +// from command_path's directory) to set up the workload environment; without it the +// run fails during setup. It runs at deploy time, after command_path has been +// translated to its absolute workspace path. +func SynthesizeRequirements() bundle.Mutator { + return &synthesizeRequirements{} +} + +type synthesizeRequirements struct { + // client is the filer used for uploads, keyed by the command_path directory. + // When nil (normal case) a workspace filer is built per directory; only set in + // tests to inject a recording filer. + client filer.Filer +} + +func (m *synthesizeRequirements) Name() string { + return "aicode.SynthesizeRequirements" +} + +func (m *synthesizeRequirements) Apply(ctx context.Context, b *bundle.Bundle) diag.Diagnostics { + var diags diag.Diagnostics + + for name, job := range b.Config.Resources.Jobs { + envs := environmentsByKey(job.Environments) + for i := range job.Tasks { + task := &job.Tasks[i] + if task.AiRuntimeTask == nil { + continue + } + if err := m.synthesizeForTask(ctx, b, name, task, envs); err != nil { + diags = diags.Extend(diag.FromErr(err)) + } + } + } + + return diags +} + +// synthesizeForTask uploads requirements.yaml next to task's command_path. It is a +// no-op when the task has no deployment command_path or no matching environment +// spec (the run can still supply deps another way, so this is not an error). +func (m *synthesizeRequirements) synthesizeForTask(ctx context.Context, b *bundle.Bundle, jobName string, task *jobs.Task, envs map[string]*compute.Environment) error { + if len(task.AiRuntimeTask.Deployments) == 0 { + return nil + } + commandPath := task.AiRuntimeTask.Deployments[0].CommandPath + if commandPath == "" { + return nil + } + + env := envs[task.EnvironmentKey] + if env == nil { + return nil + } + + content, err := renderRequirements(env) + if err != nil { + return err + } + + // command_path is an absolute workspace path (translated during initialize); the + // entry script reads requirements.yaml from its directory. + dir := path.Dir(commandPath) + client := m.client + if client == nil { + client, err = filer.NewWorkspaceFilesClient(b.WorkspaceClient(ctx), dir) + if err != nil { + return err + } + } + + log.Debugf(ctx, "writing %s for job %s task %s to %s", requirementsFileName, jobName, task.TaskKey, dir) + if err := client.Write(ctx, requirementsFileName, bytes.NewReader(content), filer.OverwriteIfExists, filer.CreateParentDirectories); err != nil { + return fmt.Errorf("failed to upload %s next to command_path %q: %w", requirementsFileName, commandPath, err) + } + return nil +} + +// renderRequirements builds the requirements.yaml content from a serverless +// environment spec: version from environment_version (or the legacy client field), +// dependencies from its pip dependency list. +func renderRequirements(env *compute.Environment) ([]byte, error) { + version := env.EnvironmentVersion + if version == "" { + version = env.Client + } + spec := requirementsSpec{ + Version: version, + Dependencies: env.Dependencies, + } + var buf bytes.Buffer + enc := yaml.NewEncoder(&buf) + enc.SetIndent(2) + if err := enc.Encode(spec); err != nil { + return nil, err + } + if err := enc.Close(); err != nil { + return nil, err + } + return buf.Bytes(), nil +} + +// environmentsByKey indexes a job's environments by their key for task lookup. +func environmentsByKey(envs []jobs.JobEnvironment) map[string]*compute.Environment { + out := make(map[string]*compute.Environment, len(envs)) + for i := range envs { + if envs[i].Spec != nil { + out[envs[i].EnvironmentKey] = envs[i].Spec + } + } + return out +} diff --git a/bundle/config/mutator/aicode/requirements_test.go b/bundle/config/mutator/aicode/requirements_test.go new file mode 100644 index 00000000000..a8044e3e35c --- /dev/null +++ b/bundle/config/mutator/aicode/requirements_test.go @@ -0,0 +1,36 @@ +package aicode + +import ( + "testing" + + "github.com/databricks/databricks-sdk-go/service/compute" + "github.com/databricks/databricks-sdk-go/service/jobs" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestRenderRequirementsFromEnvironmentVersion(t *testing.T) { + out, err := renderRequirements(&compute.Environment{ + EnvironmentVersion: "5", + Dependencies: []string{"numpy", "torch>=2.0.0"}, + }) + require.NoError(t, err) + assert.Equal(t, "version: \"5\"\ndependencies:\n - numpy\n - torch>=2.0.0\n", string(out)) +} + +func TestRenderRequirementsFallsBackToClient(t *testing.T) { + out, err := renderRequirements(&compute.Environment{Client: "5"}) + require.NoError(t, err) + assert.Equal(t, "version: \"5\"\n", string(out)) +} + +func TestEnvironmentsByKey(t *testing.T) { + envs := []jobs.JobEnvironment{ + {EnvironmentKey: "default", Spec: &compute.Environment{EnvironmentVersion: "5"}}, + {EnvironmentKey: "nospec"}, + } + got := environmentsByKey(envs) + require.Contains(t, got, "default") + assert.Equal(t, "5", got["default"].EnvironmentVersion) + assert.NotContains(t, got, "nospec", "environments without a spec are skipped") +} diff --git a/bundle/config/mutator/aicode/snapshot_package.go b/bundle/config/mutator/aicode/snapshot_package.go index 1afb7435568..ee545df1d6c 100644 --- a/bundle/config/mutator/aicode/snapshot_package.go +++ b/bundle/config/mutator/aicode/snapshot_package.go @@ -3,7 +3,6 @@ package aicode import ( "archive/tar" "compress/gzip" - "context" "crypto/sha256" "encoding/hex" "fmt" @@ -14,7 +13,6 @@ import ( "time" "github.com/databricks/cli/libs/fileset" - "github.com/databricks/cli/libs/git" "github.com/databricks/cli/libs/vfs" ) @@ -30,29 +28,20 @@ var tarEpoch = time.Date(2000, 1, 1, 0, 0, 0, 0, time.UTC) // The AIR CLI excludes these; we match it so archives are identical on macOS and Linux. const appleDoublePrefix = "._" -// buildCodeSnapshot writes a reproducible gzipped tarball of the directory at -// codeDir to out and returns its SHA-256 hex digest. Every entry is prefixed with -// prefix (the code directory's basename) so the archive expands to /... — -// matching the runtime's /databricks/code_source/ extraction contract. +// buildCodeSnapshot writes a reproducible gzipped tarball of the given files to out +// and returns its SHA-256 hex digest. syncRoot is the root the files' Relative paths +// are against (the bundle sync root); relBase is the code directory relative to that +// root; prefix is the archive's top-level directory name. Each file at +// "/" is written to the archive as "/", so the archive +// expands to /... — matching the runtime's /databricks/code_source/ +// extraction contract. // -// The file list is gitignore-aware via [git.NewFileSetAtRoot]: it honors the code -// directory's .gitignore (including nested .gitignore files) and always excludes -// .git and .databricks — the same walker that backs bundle file sync. This is why -// packaging a large tree only archives the tracked payload rather than venvs, -// caches, and build outputs. -func buildCodeSnapshot(ctx context.Context, codeDir vfs.Path, prefix string, out io.Writer) (string, error) { - fsys, err := git.NewFileSetAtRoot(ctx, codeDir) - if err != nil { - return "", err - } - - files, err := fsys.Files() - if err != nil { - return "", err - } - +// The file list is produced by the bundle's sync walker, so it honors .gitignore +// (including nested files) and the top-level sync.include/exclude globs — the same +// filtering as bundle file sync. +func buildCodeSnapshot(syncRoot vfs.Path, relBase string, files []fileset.File, prefix string, out io.Writer) (string, error) { // Sort by relative path so the archive byte stream (and thus its hash) does not - // depend on filesystem walk order. + // depend on iteration order. slices.SortFunc(files, func(a, b fileset.File) int { return strings.Compare(a.Relative, b.Relative) }) @@ -62,7 +51,7 @@ func buildCodeSnapshot(ctx context.Context, codeDir vfs.Path, prefix string, out tw := tar.NewWriter(gzw) for _, f := range files { - if err := addFileToArchive(tw, codeDir, f, prefix); err != nil { + if err := addFileToArchive(tw, syncRoot, relBase, f, prefix); err != nil { return "", err } } @@ -76,13 +65,25 @@ func buildCodeSnapshot(ctx context.Context, codeDir vfs.Path, prefix string, out return hex.EncodeToString(hash.Sum(nil)), nil } -func addFileToArchive(tw *tar.Writer, codeDir vfs.Path, f fileset.File, prefix string) error { - // fileset.File.Relative is slash-separated, so path.Base is correct here. - if strings.HasPrefix(path.Base(f.Relative), appleDoublePrefix) { +func addFileToArchive(tw *tar.Writer, syncRoot vfs.Path, relBase string, f fileset.File, prefix string) error { + // f.Relative is relative to syncRoot and slash-separated. Re-base it under the + // code directory so the entry nests under the archive prefix. + rel := f.Relative + if relBase != "." { + trimmed, ok := strings.CutPrefix(rel, relBase+"/") + if !ok { + // Not under the code dir; the sync file list is scoped to it, so this + // should not happen, but skip defensively rather than mis-place a file. + return nil + } + rel = trimmed + } + + if strings.HasPrefix(path.Base(rel), appleDoublePrefix) { return nil } - rc, err := codeDir.Open(f.Relative) + rc, err := syncRoot.Open(f.Relative) if err != nil { return fmt.Errorf("open %s: %w", f.Relative, err) } @@ -93,15 +94,15 @@ func addFileToArchive(tw *tar.Writer, codeDir vfs.Path, f fileset.File, prefix s return fmt.Errorf("stat %s: %w", f.Relative, err) } - // Only regular files are archived. The gitignore-aware walker never yields - // directories, and symlinks inside a code snapshot are out of scope. + // Only regular files are archived. The walker never yields directories, and + // symlinks inside a code snapshot are out of scope. if !info.Mode().IsRegular() { return nil } hdr := &tar.Header{ Typeflag: tar.TypeReg, - Name: path.Join(prefix, f.Relative), + Name: path.Join(prefix, rel), Size: info.Size(), // Normalize permissions and zero the mtime so the archive is reproducible // across machines. The runtime invokes code via an interpreter, not by @@ -110,10 +111,10 @@ func addFileToArchive(tw *tar.Writer, codeDir vfs.Path, f fileset.File, prefix s ModTime: tarEpoch, } if err := tw.WriteHeader(hdr); err != nil { - return fmt.Errorf("tar header for %s: %w", f.Relative, err) + return fmt.Errorf("tar header for %s: %w", rel, err) } if _, err := io.Copy(tw, rc); err != nil { - return fmt.Errorf("write %s: %w", f.Relative, err) + return fmt.Errorf("write %s: %w", rel, err) } return nil } diff --git a/bundle/config/mutator/aicode/snapshot_package_test.go b/bundle/config/mutator/aicode/snapshot_package_test.go index 4d2f86d496f..feb2e10c1e7 100644 --- a/bundle/config/mutator/aicode/snapshot_package_test.go +++ b/bundle/config/mutator/aicode/snapshot_package_test.go @@ -9,14 +9,15 @@ import ( "path/filepath" "testing" + "github.com/databricks/cli/libs/fileset" "github.com/databricks/cli/libs/vfs" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) // writeTree materializes files (relative slash path -> content) under a fresh -// temp dir and returns a vfs.Path rooted at it. -func writeTree(t *testing.T, files map[string]string) vfs.Path { +// temp dir and returns a vfs.Path rooted at it plus the fileset for its contents. +func writeTree(t *testing.T, files map[string]string) (vfs.Path, []fileset.File) { t.Helper() dir := t.TempDir() for name, content := range files { @@ -24,7 +25,10 @@ func writeTree(t *testing.T, files map[string]string) vfs.Path { require.NoError(t, os.MkdirAll(filepath.Dir(p), 0o755)) require.NoError(t, os.WriteFile(p, []byte(content), 0o644)) } - return vfs.MustNew(dir) + root := vfs.MustNew(dir) + fs, err := fileset.New(root).Files() + require.NoError(t, err) + return root, fs } // tarEntries reads a gzipped tarball and returns entry name -> content. @@ -47,20 +51,15 @@ func tarEntries(t *testing.T, b []byte) map[string]string { return out } -func TestBuildCodeSnapshotPrefixesAndExcludes(t *testing.T) { - root := writeTree(t, map[string]string{ +func TestBuildCodeSnapshotPrefixesEntries(t *testing.T) { + root, files := writeTree(t, map[string]string{ "train.py": "print('train')", "pkg/util.py": "x = 1", - ".gitignore": "secrets.txt\n*.log\n", - "secrets.txt": "nope", - "debug.log": "nope", - "pkg/nested.log": "nope", - ".git/config": "[core]", "._resource_fork": "apple double", }) var buf bytes.Buffer - sha, err := buildCodeSnapshot(t.Context(), root, "mycode", &buf) + sha, err := buildCodeSnapshot(root, ".", files, "mycode", &buf) require.NoError(t, err) require.NotEmpty(t, sha) @@ -69,43 +68,39 @@ func TestBuildCodeSnapshotPrefixesAndExcludes(t *testing.T) { // /databricks/code_source/). assert.Equal(t, "print('train')", entries["mycode/train.py"]) assert.Equal(t, "x = 1", entries["mycode/pkg/util.py"]) - assert.Contains(t, entries, "mycode/.gitignore") - - assert.NotContains(t, entries, "mycode/secrets.txt", "gitignored file must be excluded") - assert.NotContains(t, entries, "mycode/debug.log", "gitignored glob must be excluded") - assert.NotContains(t, entries, "mycode/pkg/nested.log", "gitignored glob must be excluded in subdirs") - assert.NotContains(t, entries, "mycode/.git/config", ".git must never be archived") assert.NotContains(t, entries, "mycode/._resource_fork", "AppleDouble metadata must be excluded") } -func TestBuildCodeSnapshotHonorsNestedGitignore(t *testing.T) { - // A .gitignore in a subdirectory applies within that subtree — this is the - // case the previous shell-tar implementation missed. - root := writeTree(t, map[string]string{ - "a.py": "a", - "pkg/keep.py": "keep", - "pkg/.gitignore": "drop.py\n", - "pkg/drop.py": "nope", +func TestBuildCodeSnapshotRebasesUnderRelBase(t *testing.T) { + // Files listed relative to a sync root; only the "src" subtree is packaged and + // re-based so entries nest under the prefix (not the intermediate "src/"). + root, files := writeTree(t, map[string]string{ + "src/train.py": "t", + "src/pkg/util.py": "u", + "other/ignore.py": "o", }) var buf bytes.Buffer - _, err := buildCodeSnapshot(t.Context(), root, "code", &buf) + _, err := buildCodeSnapshot(root, "src", files, "src", &buf) require.NoError(t, err) entries := tarEntries(t, buf.Bytes()) - assert.Contains(t, entries, "code/pkg/keep.py") - assert.NotContains(t, entries, "code/pkg/drop.py", "nested .gitignore must be honored") + assert.Contains(t, entries, "src/train.py") + assert.Contains(t, entries, "src/pkg/util.py") + // A file outside relBase is not under "src/", so it is skipped. + assert.NotContains(t, entries, "src/other/ignore.py") + assert.NotContains(t, entries, "other/ignore.py") } func TestBuildCodeSnapshotIsReproducible(t *testing.T) { files := map[string]string{"a.py": "aaa", "sub/b.py": "bbb"} - root1 := writeTree(t, files) - root2 := writeTree(t, files) + root1, fs1 := writeTree(t, files) + root2, fs2 := writeTree(t, files) var buf1, buf2 bytes.Buffer - sha1, err := buildCodeSnapshot(t.Context(), root1, "code", &buf1) + sha1, err := buildCodeSnapshot(root1, ".", fs1, "code", &buf1) require.NoError(t, err) - sha2, err := buildCodeSnapshot(t.Context(), root2, "code", &buf2) + sha2, err := buildCodeSnapshot(root2, ".", fs2, "code", &buf2) require.NoError(t, err) assert.Equal(t, sha1, sha2, "identical content must produce an identical hash") @@ -113,13 +108,13 @@ func TestBuildCodeSnapshotIsReproducible(t *testing.T) { } func TestBuildCodeSnapshotHashChangesWithContent(t *testing.T) { - root1 := writeTree(t, map[string]string{"main.py": "v1"}) - root2 := writeTree(t, map[string]string{"main.py": "v2"}) + root1, fs1 := writeTree(t, map[string]string{"main.py": "v1"}) + root2, fs2 := writeTree(t, map[string]string{"main.py": "v2"}) var buf1, buf2 bytes.Buffer - sha1, err := buildCodeSnapshot(t.Context(), root1, "code", &buf1) + sha1, err := buildCodeSnapshot(root1, ".", fs1, "code", &buf1) require.NoError(t, err) - sha2, err := buildCodeSnapshot(t.Context(), root2, "code", &buf2) + sha2, err := buildCodeSnapshot(root2, ".", fs2, "code", &buf2) require.NoError(t, err) assert.NotEqual(t, sha1, sha2) diff --git a/bundle/config/mutator/paths/job_paths_visitor.go b/bundle/config/mutator/paths/job_paths_visitor.go index 49011e4b1e7..699bad1b9ff 100644 --- a/bundle/config/mutator/paths/job_paths_visitor.go +++ b/bundle/config/mutator/paths/job_paths_visitor.go @@ -46,6 +46,11 @@ func jobTaskRewritePatterns(base dyn.Pattern) []jobRewritePattern { TranslateModeFile, noSkipRewrite, }, + { + base.Append(dyn.Key("ai_runtime_task"), dyn.Key("deployments"), dyn.AnyIndex(), dyn.Key("command_path")), + TranslateModeFile, + noSkipRewrite, + }, } } diff --git a/bundle/phases/deploy.go b/bundle/phases/deploy.go index 906cfca4f90..ae15919de72 100644 --- a/bundle/phases/deploy.go +++ b/bundle/phases/deploy.go @@ -203,6 +203,11 @@ func Deploy(ctx context.Context, b *bundle.Bundle, outputHandler sync.OutputHand // remote archive. Runs after file/library upload (artifact_path resolved) // and before RunPlan for both the direct and terraform engines. aicode.PackageAndUpload(), + // Write a requirements.yaml next to each AI Runtime task's (already + // translated) command_path, derived from the job's serverless environment. + // The AI Runtime entry script reads it from command_path's directory to set + // up the workload environment. + aicode.SynthesizeRequirements(), deploy.StateUpdate(), deploy.StatePush(), permissions.ApplyWorkspaceRootPermissions(),