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 0000000000..ff4a7b5dc1 --- /dev/null +++ b/acceptance/bundle/ai_runtime_task/local_code_source/databricks.yml @@ -0,0 +1,29 @@ +bundle: + name: ai-runtime-test + +sync: + # Exclude a file from both bundle sync and the code snapshot. + exclude: + - "**/*.log" + +resources: + jobs: + train: + name: "[${bundle.target}] AI Runtime training" + tasks: + - task_key: train + environment_key: default + ai_runtime_task: + experiment: my-training + code_source_path: ./src + deployments: + - command_path: src/command.sh + compute: + accelerator_type: GPU_8xH100 + accelerator_count: 8 + environments: + - environment_key: default + spec: + environment_version: "5" + dependencies: + - torch>=2.0.0 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 0000000000..e90b6d5d1b --- /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 = ["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 0000000000..33fb18eba4 --- /dev/null +++ b/acceptance/bundle/ai_runtime_task/local_code_source/output.txt @@ -0,0 +1,29 @@ + +=== deploy packages the local code source + +>>> [CLI] bundle deploy +Uploading bundle files to /Workspace/Users/[USERNAME]/.bundle/ai-runtime-test/default/files... +Deploying resources... +Updating deployment state... +Deployment complete! + +=== the snapshot tarball is uploaded to the user's repo_snapshots directory +/api/2.0/workspace-files/import-file/Workspace/Users/[USERNAME]/.air/repo_snapshots/src/src_[SNAPSHOT].tar.gz + +=== code_source_path points at the uploaded archive (no /Workspace prefix) +/Users/[USERNAME]/.air/repo_snapshots/src/src_[SNAPSHOT].tar.gz + +=== command_path is translated to its absolute synced workspace path +/Workspace/Users/[USERNAME]/.bundle/ai-runtime-test/default/files/src/command.sh + +=== a requirements.yaml is written next to command_path +/api/2.0/workspace-files/import-file/Workspace/Users/[USERNAME]/.bundle/ai-runtime-test/default/files/src/requirements.yaml + +=== re-deploying unchanged code is a cache hit: the tarball is not re-uploaded + +>>> [CLI] bundle deploy +Uploading bundle files to /Workspace/Users/[USERNAME]/.bundle/ai-runtime-test/default/files... +Deploying resources... +Updating deployment state... +Deployment complete! +tarball uploads on redeploy: 0 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 0000000000..d0a9e2ea96 --- /dev/null +++ b/acceptance/bundle/ai_runtime_task/local_code_source/script @@ -0,0 +1,30 @@ +# A local code_source_path is packaged into a content-addressed tarball +# (_.tar.gz) and uploaded to the user's repo_snapshots directory. The +# command_path is translated to its synced workspace path, and a requirements.yaml +# derived from the job's serverless environment is written beside it. + +title "deploy packages the local code source\n" +trace $CLI bundle deploy + +title "the snapshot tarball is uploaded to the user's repo_snapshots directory\n" +# The upload appears twice: WSFS import-file 404s until the parent dir exists, then +# the filer mkdirs and retries. Collapse to the distinct path. +jq -r 'select(.path | test("/import-file/.*/.air/repo_snapshots/src/src_")) | .path' out.requests.txt | sort -u + +title "code_source_path points at the uploaded archive (no /Workspace prefix)\n" +jq -r 'select(.path == "/api/2.2/jobs/create") | .body.tasks[0].ai_runtime_task.code_source_path' out.requests.txt + +title "command_path is translated to its absolute synced workspace path\n" +jq -r 'select(.path == "/api/2.2/jobs/create") | .body.tasks[0].ai_runtime_task.deployments[0].command_path' out.requests.txt + +title "a requirements.yaml is written next to command_path\n" +jq -r 'select(.path | test("/import-file/.*/files/src/requirements.yaml$")) | .path' out.requests.txt | sort -u + +rm out.requests.txt + +title "re-deploying unchanged code is a cache hit: the tarball is not re-uploaded\n" +trace $CLI bundle deploy +echo -n "tarball uploads on redeploy: " +jq -r 'select(.path | test("/import-file/.*/.air/repo_snapshots/src/src_")) | .path' out.requests.txt | sort -u | wc -l | tr -d ' ' + +rm out.requests.txt 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 0000000000..7a79c3b103 --- /dev/null +++ b/acceptance/bundle/ai_runtime_task/local_code_source/src/.gitignore @@ -0,0 +1 @@ +ignored_by_git.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 0000000000..7ce3949767 --- /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 0000000000..d8b062e5df --- /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 0000000000..894a7d77b1 --- /dev/null +++ b/acceptance/bundle/ai_runtime_task/local_code_source/test.toml @@ -0,0 +1,17 @@ +RecordRequests = true + +# Direct engine only: the terraform engine needs code_source_path in the generated +# terraform provider schema, absent while the field is temporarily removed from the +# jobs API/provider (returns next week). Restore ["terraform", "direct"] then. +EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["direct"] + +Ignore = [ + '.databricks', +] + +# The archive is content-addressed: _.tar.gz. The hash is stable +# given the committed inputs, but collapse it to a token so the test does not pin a +# specific digest. +[[Repls]] +Old = 'src_[0-9a-f]{16}\.tar\.gz' +New = 'src_[SNAPSHOT].tar.gz' diff --git a/bundle/config/mutator/aicode/package_upload.go b/bundle/config/mutator/aicode/package_upload.go new file mode 100644 index 0000000000..449b6fd971 --- /dev/null +++ b/bundle/config/mutator/aicode/package_upload.go @@ -0,0 +1,214 @@ +// Package aicode packages a local code_source_path directory into a +// content-addressed tarball at deploy, uploads it to the user's workspace, and +// rewrites code_source_path to the remote archive. Already-remote values are left +// untouched. The archive name embeds its SHA-256, so unchanged code skips re-upload. +package aicode + +import ( + "bytes" + "context" + "errors" + "fmt" + "io/fs" + "path" + "path/filepath" + "strings" + + "github.com/databricks/cli/bundle" + "github.com/databricks/cli/bundle/deploy/files" + "github.com/databricks/cli/bundle/libraries" + "github.com/databricks/cli/libs/diag" + "github.com/databricks/cli/libs/dyn" + "github.com/databricks/cli/libs/filer" + "github.com/databricks/cli/libs/fileset" + "github.com/databricks/cli/libs/log" + libsync "github.com/databricks/cli/libs/sync" +) + +// codeSourcePatterns locate code_source_path on a direct task and under a for_each_task. +var codeSourcePatterns = []dyn.Pattern{ + dyn.NewPattern( + dyn.Key("resources"), dyn.Key("jobs"), dyn.AnyKey(), + dyn.Key("tasks"), dyn.AnyIndex(), + dyn.Key("ai_runtime_task"), dyn.Key("code_source_path"), + ), + dyn.NewPattern( + dyn.Key("resources"), dyn.Key("jobs"), dyn.AnyKey(), + dyn.Key("tasks"), dyn.AnyIndex(), + dyn.Key("for_each_task"), dyn.Key("task"), + dyn.Key("ai_runtime_task"), dyn.Key("code_source_path"), + ), +} + +// codeSource is a single local code_source_path occurrence to package. +type codeSource struct { + configPath dyn.Path + location dyn.Location + value string // raw code_source_path string as written in config +} + +func PackageAndUpload() bundle.Mutator { + return &packageAndUpload{} +} + +type packageAndUpload struct { + client filer.Filer // nil in normal use (a filer is built per code source); set only in tests +} + +func (m *packageAndUpload) Name() string { + return "aicode.PackageAndUpload" +} + +func (m *packageAndUpload) Apply(ctx context.Context, b *bundle.Bundle) diag.Diagnostics { + sources, diags := collectLocalCodeSources(b.Config.Value()) + if diags.HasError() { + return diags + } + if len(sources) == 0 { + return diags + } + + userDir, err := userWorkspaceHome(b) + if err != nil { + return diags.Extend(diag.FromErr(err)) + } + + // Upload all sources before rewriting any config, so an upload failure leaves + // the config untouched. + remotePaths := make(map[string]string, len(sources)) + for _, cs := range sources { + remote, err := m.packageOne(ctx, b, cs, userDir) + if err != nil { + diags = diags.Extend(diag.FromErr(err)) + return diags + } + remotePaths[cs.configPath.String()] = remote + } + + err = b.Config.Mutate(func(root dyn.Value) (dyn.Value, error) { + for _, cs := range sources { + remote := remotePaths[cs.configPath.String()] + root, err = dyn.SetByPath(root, cs.configPath, dyn.NewValue(remote, []dyn.Location{cs.location})) + if err != nil { + return root, fmt.Errorf("failed to update code_source_path %q to %q: %w", cs.value, remote, err) + } + } + return root, nil + }) + if err != nil { + diags = diags.Extend(diag.FromErr(err)) + } + + return diags +} + +// repoSnapshotsSubdir holds code snapshots under the user's home. Deliberately not +// /.internal, which artifacts.CleanUp() wipes each deploy (matches the Python air CLI). +const repoSnapshotsSubdir = ".air/repo_snapshots" + +// packageOne tarballs one code source, uploads it to repo_snapshots (skipping the +// upload when a content-identical archive is already there), and returns its remote path. +func (m *packageAndUpload) packageOne(ctx context.Context, b *bundle.Bundle, cs codeSource, userDir string) (string, error) { + localDir := filepath.Join(b.SyncRootPath, filepath.FromSlash(cs.value)) + dirName := filepath.Base(localDir) + + // relBase scopes the sync file list to the code dir and re-bases archive entries under it. + relBase, err := filepath.Rel(b.SyncRootPath, localDir) + if err != nil { + return "", fmt.Errorf("code_source_path %q: %w", cs.value, err) + } + relBase = filepath.ToSlash(relBase) + + files, err := codeSourceFiles(ctx, b, relBase) + if err != nil { + return "", fmt.Errorf("failed to list files for code_source_path %q: %w", cs.value, err) + } + + uploadPath := path.Join(userDir, repoSnapshotsSubdir, dirName) + client := m.client + if client == nil { + client, err = filer.NewWorkspaceFilesClient(b.WorkspaceClient(ctx), uploadPath) + if err != nil { + return "", err + } + } + + // Build in memory so the content hash (computed while gzipping) can name the upload. + var buf bytes.Buffer + sha, err := buildCodeSnapshot(b.SyncRoot, relBase, files, dirName, &buf) + if err != nil { + return "", fmt.Errorf("failed to package code_source_path %q: %w", cs.value, err) + } + + archiveName := fmt.Sprintf("%s_%s.tar.gz", dirName, sha[:16]) + // The AI Runtime fetcher wants the legacy "/Users/..." form (no "/Workspace"), while + // the filer uploads via uploadPath; record the de-prefixed path on the task. + remotePath := strings.TrimPrefix(path.Join(uploadPath, archiveName), "/Workspace") + + // A matching name means identical content is already uploaded (the archive is reproducible). + if _, err := client.Stat(ctx, archiveName); err == nil { + log.Debugf(ctx, "code snapshot already present at %s, skipping upload", remotePath) + return remotePath, nil + } else if !errors.Is(err, fs.ErrNotExist) { + return "", fmt.Errorf("failed to check for existing code snapshot %q: %w", remotePath, err) + } + + if err := client.Write(ctx, archiveName, &buf, filer.OverwriteIfExists, filer.CreateParentDirectories); err != nil { + return "", fmt.Errorf("failed to upload code snapshot %q: %w", remotePath, err) + } + return remotePath, nil +} + +// codeSourceFiles lists the files under relBase (relative to the sync root) for the +// snapshot, filtered like bundle file sync: .gitignore-aware plus sync.include/exclude. +func codeSourceFiles(ctx context.Context, b *bundle.Bundle, relBase string) ([]fileset.File, error) { + opts, err := files.GetSyncOptions(ctx, b) + if err != nil { + return nil, err + } + fl, err := libsync.NewFileList(ctx, opts.WorktreeRoot, opts.LocalRoot, []string{relBase}, opts.Include, opts.Exclude) + if err != nil { + return nil, err + } + return fl.Files(ctx) +} + +// userWorkspaceHome returns /Workspace/Users/, the root for code snapshots. +func userWorkspaceHome(b *bundle.Bundle) (string, error) { + u := b.Config.Workspace.CurrentUser + if u == nil || u.User == nil || u.UserName == "" { + return "", errors.New("unable to resolve code snapshot location: current user not set") + } + return "/Workspace/Users/" + u.UserName, nil +} + +// collectLocalCodeSources returns every AI Runtime task code_source_path in root +// that points at a local directory. Already-remote values are skipped. It reads the +// dyn value directly (read-only), so it does not depend on the typed SDK field. +func collectLocalCodeSources(root dyn.Value) ([]codeSource, diag.Diagnostics) { + var sources []codeSource + var diags diag.Diagnostics + + for _, pattern := range codeSourcePatterns { + _, err := dyn.MapByPattern(root, pattern, func(p dyn.Path, v dyn.Value) (dyn.Value, error) { + value, ok := v.AsString() + if !ok { + return v, fmt.Errorf("expected string, got %s", v.Kind()) + } + if !libraries.IsLocalPath(value) { + return v, nil + } + sources = append(sources, codeSource{ + configPath: p, + location: v.Location(), + value: value, + }) + return v, nil + }) + if err != nil { + diags = diags.Extend(diag.FromErr(err)) + } + } + + return sources, diags +} 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 0000000000..26fdc5f336 --- /dev/null +++ b/bundle/config/mutator/aicode/package_upload_test.go @@ -0,0 +1,55 @@ +package aicode + +import ( + "strings" + "testing" + + "github.com/databricks/cli/libs/dyn" + "github.com/databricks/cli/libs/dyn/yamlloader" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// configWithCodeSource parses a minimal bundle config whose single AI Runtime task +// points at codeSourcePath, as a dyn value. It reads YAML directly so the config +// retains code_source_path regardless of whether the SDK's typed struct has it. +// +// The end-to-end package/upload/rewrite behavior (local dir -> tarball -> upload -> +// rewritten code_source_path) runs the full mutator pipeline (sync file list, +// workspace filer) and is covered by acceptance tests under +// acceptance/bundle/ai_runtime_task. This unit test covers only the pure +// config-collection seam that does not touch the pipeline. +func configWithCodeSource(t *testing.T, codeSourcePath string) dyn.Value { + t.Helper() + yml := ` +resources: + jobs: + train: + tasks: + - task_key: train + ai_runtime_task: + experiment: exp + code_source_path: ` + codeSourcePath + ` +` + v, err := yamlloader.LoadYAML("test.yml", strings.NewReader(yml)) + require.NoError(t, err) + return v +} + +func TestCollectLocalCodeSourcesFindsLocalPath(t *testing.T) { + sources, diags := collectLocalCodeSources(configWithCodeSource(t, "./src")) + require.Empty(t, diags) + require.Len(t, sources, 1) + assert.Equal(t, "./src", sources[0].value) +} + +func TestCollectLocalCodeSourcesSkipsRemotePaths(t *testing.T) { + for _, remote := range []string{ + "/Workspace/Users/me/code.tar.gz", + "/Volumes/main/default/code/existing.tar.gz", + } { + sources, diags := collectLocalCodeSources(configWithCodeSource(t, remote)) + require.Empty(t, diags) + assert.Empty(t, sources, "remote code_source_path %q must not be collected", remote) + } +} diff --git a/bundle/config/mutator/aicode/requirements.go b/bundle/config/mutator/aicode/requirements.go new file mode 100644 index 0000000000..6e55d3ea28 --- /dev/null +++ b/bundle/config/mutator/aicode/requirements.go @@ -0,0 +1,137 @@ +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 read by the AI Runtime entry script from the command_path +// directory, so it must sit next to command.sh. +const requirementsFileName = "requirements.yaml" + +// requirementsSpec is the requirements.yaml shape the AI Runtime consumes. +type requirementsSpec struct { + Version string `yaml:"version,omitempty"` + Dependencies []string `yaml:"dependencies,omitempty"` +} + +// SynthesizeRequirements uploads requirements.yaml next to each AI Runtime task's +// command_path, derived from the task's serverless environment. Runs after +// command_path has been translated to its absolute workspace path. +func SynthesizeRequirements() bundle.Mutator { + return &synthesizeRequirements{} +} + +type synthesizeRequirements struct { + client filer.Filer // nil in normal use (a filer is built per directory); set only in tests +} + +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; a no-op +// when there's no command_path or matching environment (deps can come from elsewhere). +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 + } + + // Skip when the environment has no dependencies: a deps-less file would clobber a + // requirements.yaml the workload supplies another way (e.g. the air CLI's own). + if len(env.Dependencies) == 0 { + return nil + } + + content, err := renderRequirements(env) + if err != nil { + return err + } + + // The entry script reads requirements.yaml from command_path's 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 requirements.yaml from a serverless environment spec. +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 0000000000..a8044e3e35 --- /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 new file mode 100644 index 0000000000..6c3684ec13 --- /dev/null +++ b/bundle/config/mutator/aicode/snapshot_package.go @@ -0,0 +1,100 @@ +package aicode + +import ( + "archive/tar" + "compress/gzip" + "crypto/sha256" + "encoding/hex" + "fmt" + "io" + "path" + "slices" + "strings" + "time" + + "github.com/databricks/cli/libs/fileset" + "github.com/databricks/cli/libs/vfs" +) + +// tarEpoch is stamped on every entry so identical content yields identical bytes +// (and SHA-256) regardless of mtimes, which is what makes the archive content-addressed. +var tarEpoch = time.Date(2000, 1, 1, 0, 0, 0, 0, time.UTC) + +// appleDoublePrefix marks macOS AppleDouble metadata files, excluded to match the AIR CLI. +const appleDoublePrefix = "._" + +// buildCodeSnapshot writes a reproducible gzipped tarball of files to out and returns +// its SHA-256. Each file at "/" becomes "/", so the archive +// expands to /... matching the runtime's /databricks/code_source/ contract. +func buildCodeSnapshot(syncRoot vfs.Path, relBase string, files []fileset.File, prefix string, out io.Writer) (string, error) { + // Sort by relative path so the byte stream (and hash) is order-independent. + slices.SortFunc(files, 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 files { + if err := addFileToArchive(tw, syncRoot, relBase, 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 addFileToArchive(tw *tar.Writer, syncRoot vfs.Path, relBase string, f fileset.File, prefix string) error { + // Re-base f.Relative (relative to syncRoot) under the code dir so it nests under prefix. + rel := f.Relative + if relBase != "." { + trimmed, ok := strings.CutPrefix(rel, relBase+"/") + if !ok { + return nil // outside the code dir; the file list is scoped to it, so skip defensively + } + rel = trimmed + } + + if strings.HasPrefix(path.Base(rel), appleDoublePrefix) { + return nil + } + + rc, err := syncRoot.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 (symlinks are out of scope). + if !info.Mode().IsRegular() { + return nil + } + + hdr := &tar.Header{ + Typeflag: tar.TypeReg, + Name: path.Join(prefix, rel), + Size: info.Size(), + // Fixed mode + mtime for reproducibility; code runs via an interpreter, so exec bits don't matter. + Mode: 0o644, + ModTime: tarEpoch, + } + if err := tw.WriteHeader(hdr); err != nil { + return fmt.Errorf("tar header for %s: %w", rel, err) + } + if _, err := io.Copy(tw, rc); err != nil { + 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 new file mode 100644 index 0000000000..feb2e10c1e --- /dev/null +++ b/bundle/config/mutator/aicode/snapshot_package_test.go @@ -0,0 +1,121 @@ +package aicode + +import ( + "archive/tar" + "bytes" + "compress/gzip" + "io" + "os" + "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 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 { + 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)) + } + 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. +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 TestBuildCodeSnapshotPrefixesEntries(t *testing.T) { + root, files := writeTree(t, map[string]string{ + "train.py": "print('train')", + "pkg/util.py": "x = 1", + "._resource_fork": "apple double", + }) + + var buf bytes.Buffer + sha, err := buildCodeSnapshot(root, ".", files, "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.Equal(t, "print('train')", entries["mycode/train.py"]) + assert.Equal(t, "x = 1", entries["mycode/pkg/util.py"]) + assert.NotContains(t, entries, "mycode/._resource_fork", "AppleDouble metadata must be excluded") +} + +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(root, "src", files, "src", &buf) + require.NoError(t, err) + + entries := tarEntries(t, buf.Bytes()) + 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, fs1 := writeTree(t, files) + root2, fs2 := writeTree(t, files) + + var buf1, buf2 bytes.Buffer + sha1, err := buildCodeSnapshot(root1, ".", fs1, "code", &buf1) + require.NoError(t, err) + sha2, err := buildCodeSnapshot(root2, ".", fs2, "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, 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(root1, ".", fs1, "code", &buf1) + require.NoError(t, err) + sha2, err := buildCodeSnapshot(root2, ".", fs2, "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 0000000000..72946fa99e --- /dev/null +++ b/bundle/config/mutator/aicode/validate.go @@ -0,0 +1,106 @@ +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 surfaces local code_source_path misconfigurations at `bundle validate` +// time rather than 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 { + return validateCodeSources(b.Config.Value(), b.SyncRootPath, b.IsImmutableFolder(), b.Config.GetLocations) +} + +// validateCodeSources checks every local code_source_path in root. It reads the dyn +// value directly (not the typed SDK field, which is a churny private-preview field), +// so it works regardless of the SDK's typed struct. locations resolves a config path +// to its source locations for diagnostics. +func validateCodeSources(root dyn.Value, syncRootPath string, isImmutableFolder bool, locations func(string) []dyn.Location) diag.Diagnostics { + var diags diag.Diagnostics + + for _, pattern := range codeSourcePatterns { + _, err := dyn.MapByPattern(root, pattern, func(codePath dyn.Path, cv dyn.Value) (dyn.Value, error) { + codeSourcePath, ok := cv.AsString() + if !ok { + return cv, nil + } + diags = diags.Extend(validateTask(root, syncRootPath, isImmutableFolder, locations, codeSourcePath, codePath)) + return cv, nil + }) + if err != nil { + diags = diags.Extend(diag.FromErr(err)) + } + } + + return diags +} + +func validateTask(root dyn.Value, syncRootPath string, isImmutableFolder bool, locations func(string) []dyn.Location, codeSourcePath string, codePath dyn.Path) diag.Diagnostics { + // Only local values are packaged; remote ones are used as-is. + if codeSourcePath == "" || !libraries.IsLocalPath(codeSourcePath) { + return nil + } + + locs := locations(codePath.String()) + + // git_source retrieves task files from git, so a local directory would be ignored. + // The job path is code_source_path's path truncated to resources.jobs.. + jobPath := codePath[:3] + if _, err := dyn.GetByPath(root, jobPath.Append(dyn.Key("git_source"))); err == nil { + 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: locs, + Paths: []dyn.Path{codePath}, + }} + } + + // Immutable-folder uploads one snapshot and doesn't support per-task packaging. + if 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: locs, + Paths: []dyn.Path{codePath}, + }} + } + + localDir := filepath.Join(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: locs, + 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: locs, + 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 0000000000..8a971254e0 --- /dev/null +++ b/bundle/config/mutator/aicode/validate_test.go @@ -0,0 +1,76 @@ +package aicode + +import ( + "os" + "path/filepath" + "strings" + "testing" + + "github.com/databricks/cli/libs/dyn" + "github.com/databricks/cli/libs/dyn/yamlloader" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// configForValidate parses a bundle config with one AI Runtime task pointing at +// codeSourcePath (and optionally a job git_source), as a dyn value. Reading YAML +// directly keeps code_source_path regardless of the SDK's typed struct. +func configForValidate(t *testing.T, codeSourcePath string, withGitSource bool) dyn.Value { + t.Helper() + git := "" + if withGitSource { + git = "\n git_source:\n git_url: https://example.invalid/repo\n git_provider: gitHub" + } + yml := ` +resources: + jobs: + train:` + git + ` + tasks: + - task_key: train + ai_runtime_task: + experiment: exp + code_source_path: ` + codeSourcePath + ` +` + v, err := yamlloader.LoadYAML("test.yml", strings.NewReader(yml)) + require.NoError(t, err) + return v +} + +// noLocations is a stub locations resolver for tests (diagnostics content, not +// locations, is what these assert). +func noLocations(string) []dyn.Location { return nil } + +func TestValidateMissingCodeSourceDir(t *testing.T) { + root := configForValidate(t, "does-not-exist", false) + diags := validateCodeSources(root, t.TempDir(), false, noLocations) + require.Len(t, diags, 1) + assert.Equal(t, `code_source_path "does-not-exist" not found`, diags[0].Summary) +} + +func TestValidateGitSourceConflict(t *testing.T) { + root := configForValidate(t, "src", true) + diags := validateCodeSources(root, t.TempDir(), false, noLocations) + require.Len(t, diags, 1) + assert.Contains(t, diags[0].Summary, "cannot be combined with git_source") +} + +func TestValidateImmutableFolderConflict(t *testing.T) { + root := configForValidate(t, "src", false) + diags := validateCodeSources(root, t.TempDir(), true, noLocations) + require.Len(t, diags, 1) + assert.Contains(t, diags[0].Summary, "experimental.immutable_folder") +} + +func TestValidateAcceptsExistingDir(t *testing.T) { + dir := t.TempDir() + require.NoError(t, os.MkdirAll(filepath.Join(dir, "src"), 0o755)) + root := configForValidate(t, "src", false) + diags := validateCodeSources(root, dir, false, noLocations) + assert.Empty(t, diags) +} + +func TestValidateRemoteCodeSourceIsSkipped(t *testing.T) { + root := configForValidate(t, "/Volumes/main/default/code/x.tar.gz", false) + diags := validateCodeSources(root, t.TempDir(), false, noLocations) + assert.Empty(t, diags) +} diff --git a/bundle/phases/build.go b/bundle/phases/build.go index db376e07e2..a386093bc2 100644 --- a/bundle/phases/build.go +++ b/bundle/phases/build.go @@ -7,6 +7,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" "github.com/databricks/cli/bundle/libraries" "github.com/databricks/cli/bundle/scripts" "github.com/databricks/cli/bundle/trampoline" @@ -27,6 +28,17 @@ func Build(ctx context.Context, b *bundle.Bundle) LibLocationMap { artifacts.Build(), scripts.Execute(config.ScriptPostBuild), + // Package any AI Runtime task code_source_path that points at a local + // directory into a tarball, upload it, and rewrite the field to the remote + // path. Runs before ExpandGlobReferences/ReplaceWithRemotePath below so those + // see an already-remote code_source_path (a directory would otherwise be + // collected as a local "library" and fail to upload as a file). + aicode.PackageAndUpload(), + // Write requirements.yaml next to the (already translated) command_path, + // derived from the job's serverless environment, so the AI Runtime harness + // can set up the workload environment. + aicode.SynthesizeRequirements(), + mutator.ResolveVariableReferencesWithoutResources( "artifacts", ), diff --git a/bundle/phases/initialize.go b/bundle/phases/initialize.go index bfa2af4124..f1554a4b7a 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