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..ccfd8531bdf 100644 --- a/NEXT_CHANGELOG.md +++ b/NEXT_CHANGELOG.md @@ -8,6 +8,8 @@ ### 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. 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 ### 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..ff4a7b5dc1b --- /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 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..33fb18eba49 --- /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 00000000000..d0a9e2ea965 --- /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 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/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..6e4639a532a --- /dev/null +++ b/acceptance/bundle/ai_runtime_task/local_code_source/test.toml @@ -0,0 +1,12 @@ +RecordRequests = true + +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 00000000000..4124bbfb060 --- /dev/null +++ b/bundle/config/mutator/aicode/package_upload.go @@ -0,0 +1,243 @@ +// 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 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 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" + "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 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. 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 +} + +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 + } + + userDir, err := userWorkspaceHome(b) + 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, 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 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 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) + + // 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 { + client, err = filer.NewWorkspaceFilesClient(b.WorkspaceClient(ctx), uploadPath) + if err != nil { + return "", err + } + } + + // 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(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 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. + 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 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) { + 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 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..25192332f81 --- /dev/null +++ b/bundle/config/mutator/aicode/package_upload_test.go @@ -0,0 +1,74 @@ +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/iam" + "github.com/databricks/databricks-sdk-go/service/jobs" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// bundleWithCodeSource builds a bundle rooted at dir whose single AI Runtime task +// 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{ + BundleRootPath: dir, + SyncRootPath: dir, + Config: config.Root{ + Bundle: config.Bundle{Target: "default"}, + Workspace: config.Workspace{ + CurrentUser: &config.User{User: &iam.User{UserName: "me@databricks.com"}}, + }, + 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 +} + +func TestCollectLocalCodeSourcesFindsLocalPath(t *testing.T) { + b := bundleWithCodeSource(t, t.TempDir(), "./src") + sources, diags := collectLocalCodeSources(b) + 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", + } { + 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) + } +} 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 new file mode 100644 index 00000000000..ee545df1d6c --- /dev/null +++ b/bundle/config/mutator/aicode/snapshot_package.go @@ -0,0 +1,120 @@ +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 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) + +// 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 = "._" + +// 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 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 iteration order. + 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 { + // 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 := 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. 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, 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 + // 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", 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 00000000000..feb2e10c1e7 --- /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 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/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 e2159189c78..ae15919de72 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,17 @@ 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(), + // 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(), 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