From 2c18b12d74c116b8bc8bef3b653013cbb4cfd611 Mon Sep 17 00:00:00 2001 From: vinchenzo-db Date: Fri, 31 Jul 2026 05:43:30 +0000 Subject: [PATCH 01/15] air: add convert-to-dabs (run YAML -> Databricks Asset Bundle) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds `air convert-to-dabs`, which translates an AIR CLI run YAML into a deployable Databricks Asset Bundle so a workload authored for `air run` can be managed and deployed with the standard DABs workflow (validate/deploy/run). The emitted bundle is schema-valid: the ai_runtime_task maps to the SDK jobs.AiRuntimeTask (experiment + deployments[].{command_path,compute} + code_source_path), with framework fields (retries, timeout, budget policy) on the surrounding task and the runtime environment in environments[]. Snapshotting is owned by the deploy-time aicode mutator, not by convert: code_source_path points at a local *directory* staged inside the bundle, and `bundle deploy` (aicode.PackageAndUpload) packages it into a content-addressed tarball and uploads it. convert only lays down the source bytes — copying the working tree (honoring .gitignore) or materializing a pinned git commit into the directory. requirements.yaml is likewise not emitted: aicode.SynthesizeRequirements regenerates it from the environments[] spec, so convert folds the whole dependency set (inline or requirements-file) into that spec instead. env_variables / secrets / parameters have no native ai_runtime_task field, so they ride as env_vars.json / secret_env_vars.json / hyperparameters.yaml sidecars (same as `air run`), and a "Notes:" section tells a migrating user what was transformed or staged out-of-band. This is the top of a 2-PR stack: it builds on the aicode deploy-time packaging mutator so the two compose end-to-end. Co-authored-by: Isaac --- .../air/convert-to-dabs/docker.yaml | 8 + .../air/convert-to-dabs/out.test.toml | 3 + .../air/convert-to-dabs/output.txt | 83 +++ .../experimental/air/convert-to-dabs/script | 20 + .../air/convert-to-dabs/src/train.py | 1 + .../air/convert-to-dabs/test.toml | 3 + .../air/convert-to-dabs/train.yaml | 13 + acceptance/experimental/air/help/output.txt | 13 +- experimental/air/cmd/air.go | 1 + experimental/air/cmd/air_test.go | 2 +- experimental/air/cmd/convert_to_dabs.go | 586 ++++++++++++++++++ experimental/air/cmd/convert_to_dabs_test.go | 410 ++++++++++++ 12 files changed, 1136 insertions(+), 7 deletions(-) create mode 100644 acceptance/experimental/air/convert-to-dabs/docker.yaml create mode 100644 acceptance/experimental/air/convert-to-dabs/out.test.toml create mode 100644 acceptance/experimental/air/convert-to-dabs/output.txt create mode 100644 acceptance/experimental/air/convert-to-dabs/script create mode 100644 acceptance/experimental/air/convert-to-dabs/src/train.py create mode 100644 acceptance/experimental/air/convert-to-dabs/test.toml create mode 100644 acceptance/experimental/air/convert-to-dabs/train.yaml create mode 100644 experimental/air/cmd/convert_to_dabs.go create mode 100644 experimental/air/cmd/convert_to_dabs_test.go diff --git a/acceptance/experimental/air/convert-to-dabs/docker.yaml b/acceptance/experimental/air/convert-to-dabs/docker.yaml new file mode 100644 index 00000000000..848ac77fd9c --- /dev/null +++ b/acceptance/experimental/air/convert-to-dabs/docker.yaml @@ -0,0 +1,8 @@ +experiment_name: docker-test +command: python train.py +compute: + accelerator_type: GPU_1xA10 + num_accelerators: 1 +environment: + docker_image: + url: myregistry/img:tag diff --git a/acceptance/experimental/air/convert-to-dabs/out.test.toml b/acceptance/experimental/air/convert-to-dabs/out.test.toml new file mode 100644 index 00000000000..f784a183258 --- /dev/null +++ b/acceptance/experimental/air/convert-to-dabs/out.test.toml @@ -0,0 +1,3 @@ +Local = true +Cloud = false +EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["terraform", "direct"] diff --git a/acceptance/experimental/air/convert-to-dabs/output.txt b/acceptance/experimental/air/convert-to-dabs/output.txt new file mode 100644 index 00000000000..efcb78c4ca8 --- /dev/null +++ b/acceptance/experimental/air/convert-to-dabs/output.txt @@ -0,0 +1,83 @@ + +=== convert an AIR run YAML into a DABs bundle +>>> [CLI] experimental air convert-to-dabs train.yaml --output-dir generated +Wrote a Databricks Asset Bundle to generated: + databricks.yml + training_config.yaml + command.sh + code_source/ + +Notes: + - code_source was staged into the code_source/ directory; bundle deploy packages and uploads it. Re-run convert-to-dabs to re-stage after code changes. + +To deploy and run this workload as a bundle: + 1. cd generated + 2. databricks bundle validate + 3. databricks bundle deploy + 4. databricks bundle run torchrun-a10-smoke-test + +bundle deploy uploads the code source and launch scripts automatically. + +Unlike `air run` (which submits an ephemeral run), bundle deploy creates a +persistent job that is not garbage-collected. When you are done, remove the +job and its uploaded files with: + databricks bundle destroy + +=== emitted databricks.yml +>>> cat generated/databricks.yml +bundle: + name: torchrun-a10-smoke-test +targets: + dev: + mode: development + default: true +resources: + jobs: + torchrun-a10-smoke-test: + name: torchrun-a10-smoke-test + tasks: + - task_key: torchrun-a10-smoke-test + environment_key: default + ai_runtime_task: + experiment: torchrun-a10-smoke-test + deployments: + - command_path: ./command.sh + compute: + accelerator_type: GPU_1xA10 + accelerator_count: 1 + code_source_path: ./code_source + environments: + - environment_key: default + spec: + environment_version: "5" + dependencies: + - numpy + +=== the generated command.sh carries the run command +>>> cat generated/command.sh +torchrun --nproc_per_node=1 train.py +=== the code source is staged as a directory (packaged at deploy time) +>>> ls generated +code_source +command.sh +databricks.yml +training_config.yaml + +>>> ls generated/code_source +train.py + +=== the emitted bundle validates +>>> [CLI] bundle validate +Name: torchrun-a10-smoke-test +Target: dev +Workspace: + User: [USERNAME] + Path: /Workspace/Users/[USERNAME]/.bundle/torchrun-a10-smoke-test/dev + +Validation OK! + +=== docker_image is not supported yet +>>> [CLI] experimental air convert-to-dabs docker.yaml --output-dir generated-docker +Error: environment.docker_image is not yet supported by convert-to-dabs + +Exit code: 1 diff --git a/acceptance/experimental/air/convert-to-dabs/script b/acceptance/experimental/air/convert-to-dabs/script new file mode 100644 index 00000000000..1a9f9cf07c8 --- /dev/null +++ b/acceptance/experimental/air/convert-to-dabs/script @@ -0,0 +1,20 @@ +title "convert an AIR run YAML into a DABs bundle" +trace $CLI experimental air convert-to-dabs train.yaml --output-dir generated + +title "emitted databricks.yml" +trace cat generated/databricks.yml + +title "the generated command.sh carries the run command" +trace cat generated/command.sh + +title "the code source is staged as a directory (packaged at deploy time)" +trace ls generated && true +trace ls generated/code_source && true + +title "the emitted bundle validates" +cd generated +trace $CLI bundle validate +cd .. + +title "docker_image is not supported yet" +errcode trace $CLI experimental air convert-to-dabs docker.yaml --output-dir generated-docker diff --git a/acceptance/experimental/air/convert-to-dabs/src/train.py b/acceptance/experimental/air/convert-to-dabs/src/train.py new file mode 100644 index 00000000000..c859094afdf --- /dev/null +++ b/acceptance/experimental/air/convert-to-dabs/src/train.py @@ -0,0 +1 @@ +print("train") diff --git a/acceptance/experimental/air/convert-to-dabs/test.toml b/acceptance/experimental/air/convert-to-dabs/test.toml new file mode 100644 index 00000000000..19643c87452 --- /dev/null +++ b/acceptance/experimental/air/convert-to-dabs/test.toml @@ -0,0 +1,3 @@ +# The command writes a bundle under generated/; those are generated artifacts, +# not committed inputs, so exclude them from the repo-diff check. +Ignore = ["generated", "generated-docker"] diff --git a/acceptance/experimental/air/convert-to-dabs/train.yaml b/acceptance/experimental/air/convert-to-dabs/train.yaml new file mode 100644 index 00000000000..bae47bc4dbd --- /dev/null +++ b/acceptance/experimental/air/convert-to-dabs/train.yaml @@ -0,0 +1,13 @@ +experiment_name: torchrun-a10-smoke-test +command: torchrun --nproc_per_node=1 train.py +compute: + accelerator_type: GPU_1xA10 + num_accelerators: 1 +environment: + version: 5 + dependencies: + - numpy +code_source: + type: snapshot + snapshot: + root_path: ./src diff --git a/acceptance/experimental/air/help/output.txt b/acceptance/experimental/air/help/output.txt index ee89e778d6b..5cc28fd3628 100644 --- a/acceptance/experimental/air/help/output.txt +++ b/acceptance/experimental/air/help/output.txt @@ -10,12 +10,13 @@ Usage: databricks experimental air [command] Available Commands: - cancel Cancel one or more runs - get Show status, configuration, and timing details for a specific run - list List your active runs for the current profile (use --all-status for finished runs) - logs Stream or fetch logs for a run - register-image Mirror a Docker image into the workspace registry - run Submit a training workload from a YAML config + cancel Cancel one or more runs + convert-to-dabs Convert an AIR run YAML into a Databricks Asset Bundle + get Show status, configuration, and timing details for a specific run + list List your active runs for the current profile (use --all-status for finished runs) + logs Stream or fetch logs for a run + register-image Mirror a Docker image into the workspace registry + run Submit a training workload from a YAML config Flags: -h, --help help for air diff --git a/experimental/air/cmd/air.go b/experimental/air/cmd/air.go index fbf40a34b52..13b3ee18559 100644 --- a/experimental/air/cmd/air.go +++ b/experimental/air/cmd/air.go @@ -23,6 +23,7 @@ experimental and may change in future versions.`, cmd.AddCommand(newLogsCommand()) cmd.AddCommand(newCancelCommand()) cmd.AddCommand(newRegisterImageCommand()) + cmd.AddCommand(newConvertToDabsCommand()) return cmd } diff --git a/experimental/air/cmd/air_test.go b/experimental/air/cmd/air_test.go index 7efac253a2b..1843acfe900 100644 --- a/experimental/air/cmd/air_test.go +++ b/experimental/air/cmd/air_test.go @@ -14,7 +14,7 @@ func TestNewRegistersAllSubcommands(t *testing.T) { registered[c.Name()] = true } - want := []string{"run", "get", "list", "logs", "cancel", "register-image"} + want := []string{"run", "get", "list", "logs", "cancel", "register-image", "convert-to-dabs"} for _, name := range want { assert.True(t, registered[name], "subcommand %q is not registered", name) } diff --git a/experimental/air/cmd/convert_to_dabs.go b/experimental/air/cmd/convert_to_dabs.go new file mode 100644 index 00000000000..2f13ff99612 --- /dev/null +++ b/experimental/air/cmd/convert_to_dabs.go @@ -0,0 +1,586 @@ +package aircmd + +import ( + "bytes" + "context" + "errors" + "fmt" + "os" + "os/exec" + "path/filepath" + "slices" + "strconv" + "strings" + + "github.com/databricks/cli/cmd/root" + "github.com/databricks/cli/libs/cmdio" + "github.com/databricks/cli/libs/dyn" + "github.com/databricks/cli/libs/dyn/yamlsaver" + "github.com/spf13/cobra" + "go.yaml.in/yaml/v3" +) + +// convert_to_dabs turns an AIR CLI run YAML into a Databricks Asset Bundle so a +// workload authored for `air run` can be deployed and managed as a bundle. +// +// The emitted bundle is schema-valid: `databricks bundle validate` accepts it, +// and `databricks bundle deploy` reproduces the same ai_runtime_task workload the +// CLI would submit. Two properties make that work without any manual upload step: +// +// - The DABs `ai_runtime_task` is the SDK jobs.AiRuntimeTask (strict schema: +// experiment + deployments[].{command_path,compute} + code_source_path). +// Framework concerns (retries, timeout) live on the surrounding task, not in +// ai_runtime_task — so they are emitted as task-level fields. +// - code_source_path points at a *local directory* inside the bundle. The +// deploy-time aicode.PackageAndUpload mutator (bundle/config/mutator/aicode) +// owns the snapshotting: at `bundle deploy` it packages that directory into a +// content-addressed tarball, uploads it, and rewrites code_source_path to the +// remote archive. convert-to-dabs' job is only to *stage the source bytes* as +// that directory — copying the working tree, or materializing a pinned git +// commit into it — never to build the tarball itself. +// +// command.sh (and — since the task proto carries no inline env/secrets/parameters — +// the env_vars.json / secret_env_vars.json / hyperparameters.yaml sidecars the +// server-side launcher reads) are written at the bundle root and uploaded by +// deploy, mirroring the CLI's own launch layout. requirements.yaml is NOT emitted: +// the aicode.SynthesizeRequirements mutator derives it from the job's +// environments[] spec at deploy time, so convert folds the dependency set into that +// spec instead. + +// dabsTargetName is the single default target emitted; a development-mode target +// is the conventional starting point for a generated bundle. +const dabsTargetName = "dev" + +func newConvertToDabsCommand() *cobra.Command { + var outputDir string + + cmd := &cobra.Command{ + Use: "convert-to-dabs ", + Args: root.ExactArgs(1), + Short: "Convert an AIR run YAML into a Databricks Asset Bundle", + Long: `Convert an AIR CLI run YAML config into a Databricks Asset Bundle (DABs). + +The emitted bundle can be deployed with the standard DABs workflow: + + databricks bundle validate + databricks bundle deploy + +bundle deploy uploads the code source and launch scripts for you, so no manual +upload step is required. This command performs a purely local translation and +does not contact the workspace.`, + } + + cmd.Flags().StringVar(&outputDir, "output-dir", "", "Directory to write the bundle into (default: a -bundle folder next to the input YAML). Accepts an absolute or relative path.") + + cmd.RunE = func(cmd *cobra.Command, args []string) error { + ctx := cmd.Context() + yamlPath := args[0] + + cfg, err := loadRunConfig(yamlPath) + if err != nil { + return err + } + + // Default the bundle next to the input YAML (a -bundle subfolder + // so the generated files don't scatter across the user's project dir). An + // explicit --output-dir — absolute or relative — overrides. We never write to + // a temp dir: the bundle is the user's artifact to keep, deploy, and manage. + dir := outputDir + if dir == "" { + dir = filepath.Join(filepath.Dir(yamlPath), cfg.ExperimentName+"-bundle") + } + + written, err := writeBundle(ctx, cfg, yamlPath, dir) + if err != nil { + return err + } + + printConvertNextSteps(ctx, dir, written, bundleResourceKey(cfg.ExperimentName), conversionNotes(cfg)) + return nil + } + + return cmd +} + +// codeSourceDirName is the bundle-local directory the code_source is staged into. +// ai_runtime_task.code_source_path points at it; the deploy-time aicode mutator +// packages the directory into a tarball and uploads it (see the file header). A +// fixed name (rather than the source's basename) keeps the emitted databricks.yml +// deterministic regardless of where the user's code lives. +const codeSourceDirName = "code_source" + +// convertToDabs builds the DABs bundle value and the loose launch artifacts for a +// run config. It reads only what the run path's buildArtifacts reads, so the +// mapping is unit-testable in isolation. Returns the bundle root as a +// map[string]dyn.Value (ready for yamlsaver) and the loose artifacts (command.sh + +// env/secret/param sidecars) to write at the bundle root; the code_source directory +// is materialized separately by writeBundle. +func convertToDabs(ctx context.Context, cfg *runConfig, configPath string) (map[string]dyn.Value, []uploadItem, error) { + // idempotency_token is intentionally not mapped: it dedups a single runs/submit + // call, which has no analogue for a persistent, repeatedly-runnable bundle job. + // + // usage_policy_name resolution is not ported (mirrors the submit path), and + // docker images have no ai_runtime_task representation yet. + if cfg.UsagePolicyName != nil { + return nil, nil, errors.New("usage_policy_name is not yet supported by convert-to-dabs") + } + if cfg.Environment != nil && cfg.Environment.DockerImage != nil { + return nil, nil, errors.New("environment.docker_image is not yet supported by convert-to-dabs") + } + // remote_volume points the code archive at a UC Volume. bundle deploy uploads + // code_source_path to the bundle's artifact path, not an arbitrary Volume, so a + // converted bundle can't honor it — reject rather than silently drop it. + if cfg.CodeSource != nil && cfg.CodeSource.Snapshot != nil && cfg.CodeSource.Snapshot.RemoteVolume != nil { + return nil, nil, errors.New("code_source.snapshot.remote_volume is not supported by convert-to-dabs; bundle deploy manages the artifact upload location") + } + + artifacts, err := buildArtifacts(cfg, configPath) + if err != nil { + return nil, nil, err + } + // Drop requirements.yaml: the deploy-time aicode.SynthesizeRequirements mutator + // regenerates it from the job's environments[] spec (which convert populates with + // the same dependency set), so emitting it here would be redundant and could drift. + artifacts = slices.DeleteFunc(artifacts, func(it uploadItem) bool { + return it.name == requirementsName + }) + + root := buildBundleValue(ctx, cfg, configPath) + return root, artifacts, nil +} + +// nv builds a dyn.Value at "position" n: yamlsaver orders a map's keys by their +// Location line, so assigning ascending n values fixes the emitted key order. +// It routes through dyn.V so nested Go maps/slices are converted recursively, +// then stamps the ordering location. +func nv(v any, n int) dyn.Value { + return dyn.V(v).WithLocations([]dyn.Location{{Line: n}}) +} + +// localBundlePath renders a bundle-relative path with a leading "./" so bundle +// deploy classifies it as a local artifact to upload (see IsLibraryLocal). It is +// built from path.Join (forward slashes) so the emitted YAML is identical across +// operating systems. +func localBundlePath(p string) string { + return "./" + p +} + +// buildBundleValue assembles the bundle root as an ordered map[string]dyn.Value. +// command.sh and the code_source directory are bundle-local (emitted "./"-prefixed) +// so `bundle deploy` uploads/packages them. +func buildBundleValue(ctx context.Context, cfg *runConfig, configPath string) map[string]dyn.Value { + name := cfg.ExperimentName + + // ai_runtime_task: experiment + one deployment (command_path + compute) + + // code_source_path. Only the fields the strict schema allows. + // + // Paths are emitted "./"-prefixed so bundle deploy treats them as LOCAL and + // uploads them: libraries.IsLibraryLocal classifies a bare, extensionless path + // as a PyPI package name (not a local file) and skips it, which would deploy a + // path the backend can't resolve. The "./" prefix forces local classification. + deployment := map[string]dyn.Value{ + "command_path": nv(localBundlePath(commandScriptName), 1), + "compute": nv(map[string]dyn.Value{ + "accelerator_type": nv(cfg.Compute.AcceleratorType, 1), + "accelerator_count": nv(cfg.Compute.NumAccelerators, 2), + }, 2), + } + + aiRuntimeTask := map[string]dyn.Value{ + "experiment": nv(name, 1), + "deployments": nv([]dyn.Value{dyn.V(deployment)}, 2), + } + line := 3 + if cfg.CodeSource != nil && cfg.CodeSource.Snapshot != nil { + // A local directory (not a "./"-prefixed file): the aicode mutator recognizes + // a directory code_source_path, packages it, and rewrites this field at deploy. + aiRuntimeTask["code_source_path"] = nv(localBundlePath(codeSourceDirName), line) + line++ + } + if cfg.MLflowRunName != nil { + aiRuntimeTask["mlflow_run"] = nv(*cfg.MLflowRunName, line) + line++ + } + if cfg.MLflowExperimentDirectory != nil { + aiRuntimeTask["mlflow_experiment_directory"] = nv(*cfg.MLflowExperimentDirectory, line) + } + + // Task wrapper: task_key + framework fields (retries/timeout) + env key + + // the ai_runtime_task. Framework fields live here per the schema, not inside + // ai_runtime_task. + task := map[string]dyn.Value{ + "task_key": nv(name, 1), + "environment_key": nv(aiRuntimeEnvironmentKey, 2), + } + taskLine := 3 + if cfg.MaxRetries != nil { + task["max_retries"] = nv(*cfg.MaxRetries, taskLine) + taskLine++ + } + if cfg.TimeoutMinutes != nil { + task["timeout_seconds"] = nv(cfg.timeoutSeconds(), taskLine) + taskLine++ + } + task["ai_runtime_task"] = nv(aiRuntimeTask, taskLine) + + // environments[]: version + the dependency set. The aicode.SynthesizeRequirements + // mutator regenerates requirements.yaml from this spec at deploy time, so the full + // dependency set (whether authored inline or in a requirements file) must live + // here — convert emits no requirements.yaml of its own. Resolve the version + // through the same path `air run` uses (config, else env override, else the + // default channel) so a config without an explicit version still pins the version + // the workload would have run with — not an empty spec. + envVersion, deps := bundleEnvironmentDeps(ctx, cfg, configPath) + envSpec := map[string]dyn.Value{ + "environment_version": nv(envVersion, 1), + } + if len(deps) > 0 { + depVals := make([]dyn.Value, len(deps)) + for i, d := range deps { + depVals[i] = dyn.V(d) + } + envSpec["dependencies"] = nv(depVals, 2) + } + environment := map[string]dyn.Value{ + "environment_key": nv(aiRuntimeEnvironmentKey, 1), + "spec": nv(envSpec, 2), + } + + job := map[string]dyn.Value{ + "name": nv(name, 1), + "tasks": nv([]dyn.Value{dyn.V(task)}, 2), + "environments": nv([]dyn.Value{dyn.V(environment)}, 3), + } + // usage_policy_id is an already-resolved budget policy id, so it maps directly + // to the job's budget_policy_id. (usage_policy_name needs server-side resolution + // and is rejected in convertToDabs.) + if cfg.UsagePolicyID != nil { + job["budget_policy_id"] = nv(*cfg.UsagePolicyID, 4) + } + if perms := buildPermissionsValue(cfg.Permissions); perms.Kind() != dyn.KindInvalid { + job["permissions"] = nv(perms.MustSequence(), 5) + } + + rootValue := map[string]dyn.Value{ + "bundle": nv(map[string]dyn.Value{ + "name": nv(name, 1), + }, 1), + "targets": nv(map[string]dyn.Value{ + dabsTargetName: nv(map[string]dyn.Value{ + "mode": nv("development", 1), + "default": nv(true, 2), + }, 1), + }, 2), + "resources": nv(map[string]dyn.Value{ + "jobs": nv(map[string]dyn.Value{ + bundleResourceKey(name): nv(job, 1), + }, 1), + }, 3), + } + return rootValue +} + +// bundleEnvironmentDeps resolves the runtime version and the flattened dependency +// list to emit in the bundle's environments[] spec. The aicode mutator synthesizes +// requirements.yaml from that spec at deploy, so the whole set must be here — +// whether the user authored dependencies inline or pointed at a requirements file. +// A requirements file is read and its non-comment, non-blank lines are inlined; the +// version, when the file carries one, wins over the config/default version. Any read +// error is best-effort ignored (writeBundle/buildArtifacts surface real problems); +// convert falls back to inline deps so the spec is never silently wrong. +func bundleEnvironmentDeps(ctx context.Context, cfg *runConfig, configPath string) (version string, deps []string) { + cfgVersion, _ := cfg.runtimeVersion() + version = dlRuntimeImage(ctx, cfgVersion) + + if inline, ok := cfg.inlineDependencies(); ok { + return version, inline + } + + reqPath, ok := cfg.requirementsFile() + if !ok { + return version, nil + } + if !filepath.IsAbs(reqPath) { + reqPath = filepath.Join(filepath.Dir(configPath), reqPath) + } + data, err := os.ReadFile(reqPath) + if err != nil { + return version, nil + } + // The requirements file is the same requirements.yaml shape the run path reads + // (version + dependencies), so parse it as such and inline the dependency lines. + var doc requirementsDoc + if err := yaml.Unmarshal(data, &doc); err != nil { + return version, nil + } + if doc.Version != "" { + version = dlRuntimeImage(ctx, doc.Version) + } + return version, doc.Dependencies +} + +// bundleResourceKey derives a job resource key from the experiment name. The key +// is emitted as an unquoted YAML map key, and DABs' strict loader rejects a key +// that parses as a non-string scalar (a purely numeric name like "12345" -> !!int, +// or "true"/"null"). experiment_name allows exactly [alphanumeric, -, _], so the +// only unsafe keys are those that YAML types as int/float/bool/null; prefix those +// with "job_" to force a string key. The human-facing name/experiment fields keep +// the original value (yamlsaver quotes them as scalar string values). +func bundleResourceKey(name string) string { + switch strings.ToLower(name) { + case "true", "false", "null": + return "job_" + name + } + if _, err := strconv.ParseFloat(name, 64); err == nil { + return "job_" + name + } + return name +} + +// buildPermissionsValue maps run-config permissions to DABs job permissions +// (level → principal). Returns an invalid value when there are none. +func buildPermissionsValue(perms []permission) dyn.Value { + if len(perms) == 0 { + return dyn.InvalidValue + } + out := make([]dyn.Value, 0, len(perms)) + for _, p := range perms { + m := map[string]dyn.Value{"level": nv(p.Level, 1)} + switch { + case p.UserName != nil: + m["user_name"] = nv(*p.UserName, 2) + case p.GroupName != nil: + m["group_name"] = nv(*p.GroupName, 2) + case p.ServicePrincipalName != nil: + m["service_principal_name"] = nv(*p.ServicePrincipalName, 2) + } + out = append(out, dyn.V(m)) + } + return dyn.V(out) +} + +// writeBundle writes the bundle into dir: databricks.yml, the loose launch +// artifacts (command.sh + env/secret/param sidecars), and — when the config has a +// code_source — a code_source/ directory holding the staged source tree. All the +// referenced files are bundle-local; `bundle deploy` uploads the launch artifacts +// and the aicode mutator packages+uploads the code_source directory. It refuses to +// overwrite existing files so a re-run can't silently clobber a bundle the user has +// edited. Returns the relative paths written, for the next-steps message. +func writeBundle(ctx context.Context, cfg *runConfig, configPath, dir string) ([]string, error) { + root, artifacts, err := convertToDabs(ctx, cfg, configPath) + if err != nil { + return nil, err + } + + // Restrict perms: the bundle carries env_vars.json (literal env var values), so + // keep the dir owner-only rather than world-readable. + if err := os.MkdirAll(dir, 0o700); err != nil { + return nil, err + } + + // Refuse to clobber an existing file, with one consistent message. A user + // re-running convert into the same dir gets a clear error rather than a silent + // overwrite of edits they may have made. + checkCollision := func(name string) error { + if _, err := os.Stat(filepath.Join(dir, name)); err == nil { + return fmt.Errorf("%s already exists in %s; use --output-dir or remove it", name, dir) + } + return nil + } + writeFile := func(name string, data []byte) error { + if err := checkCollision(name); err != nil { + return err + } + return os.WriteFile(filepath.Join(dir, name), data, 0o600) + } + + if err := checkCollision("databricks.yml"); err != nil { + return nil, err + } + bundlePath := filepath.Join(dir, "databricks.yml") + // force=true: we've already run the collision check above, so SaveAsYAML's own + // guard (which references a --force flag this command doesn't have) can't fire. + if err := yamlsaver.NewSaver().SaveAsYAML(root, bundlePath, true); err != nil { + return nil, err + } + written := []string{"databricks.yml"} + + // Loose launch artifacts (command.sh + sidecars) at the bundle root. + for _, item := range artifacts { + if err := writeFile(item.name, item.data); err != nil { + return nil, fmt.Errorf("failed to write %s: %w", item.name, err) + } + written = append(written, item.name) + } + + // Code source: stage the resolved root_path into the bundle's code_source/ dir. + // The aicode mutator packages+uploads it at deploy; convert only lays down the + // bytes. A pinned git commit/branch is materialized from the archived commit (not + // the dirty working tree), matching what `air run` would submit; otherwise the + // working tree is copied, honoring .gitignore just like the run path's plain-tar. + if cfg.CodeSource != nil && cfg.CodeSource.Snapshot != nil { + snap := cfg.CodeSource.Snapshot + repoPath, err := resolveRootPath(ctx, snap.RootPath, filepath.Dir(configPath)) + if err != nil { + return nil, err + } + plan, err := resolveSnapshotPlan(ctx, newGitRepo(repoPath), snap.Git, snap.IncludePaths) + if err != nil { + return nil, err + } + if err := checkCollision(codeSourceDirName); err != nil { + return nil, err + } + if err := materializeCodeSource(ctx, repoPath, plan, filepath.Join(dir, codeSourceDirName)); err != nil { + return nil, err + } + written = append(written, codeSourceDirName+"/") + } + + return written, nil +} + +// materializeCodeSource stages the snapshot described by plan into destDir (the +// bundle's code_source/ directory). It packages the source with the existing +// snapshot packagers — git archive for a pinned commit, gitignore-aware plain tar +// for a working tree — into a temporary tarball, then extracts it and moves its +// single top-level directory into destDir. Going through the packagers (rather than +// a raw copy) reuses their exact commit-pin and .gitignore/include-path handling, so +// the staged tree matches what `air run` would submit; the deploy-time aicode mutator +// then re-packages destDir into the uploaded, content-addressed archive. +func materializeCodeSource(ctx context.Context, repoPath string, plan snapshotPlan, destDir string) error { + staging, err := os.MkdirTemp("", "air-convert-code-*") + if err != nil { + return err + } + defer os.RemoveAll(staging) + + // Absolute tarball path: the git-archive packager runs git with `-C repoPath`, so + // a relative -o would resolve against the repo dir. Both packagers accept absolute. + tarball, err := filepath.Abs(filepath.Join(staging, "code_source.tar.gz")) + if err != nil { + return err + } + // git archive for a pinned commit (deterministic, ignores the dirty tree), + // gitignore-aware plain tar for a working tree — the same split the run path uses. + dirName := filepath.Base(repoPath) + if plan.mode == modeGitArchive { + err = createGitArchiveSnapshot(ctx, newGitRepo(repoPath), plan.commitSHA, tarball, dirName, plan.includePaths) + } else { + err = createPlainTarball(ctx, repoPath, tarball, plan.includePaths) + } + if err != nil { + return err + } + + // The archive's single top-level entry is the source directory's basename (both + // packagers preserve it). Extract into a scratch dir, then move that top-level + // directory to the fixed destDir name so the emitted code_source_path is stable. + extractDir := filepath.Join(staging, "extract") + if err := os.MkdirAll(extractDir, 0o700); err != nil { + return err + } + if err := extractTarball(ctx, tarball, extractDir); err != nil { + return err + } + entries, err := os.ReadDir(extractDir) + if err != nil { + return err + } + if len(entries) != 1 || !entries[0].IsDir() { + return fmt.Errorf("unexpected code_source archive layout: expected a single top-level directory, got %d entries", len(entries)) + } + if err := os.Rename(filepath.Join(extractDir, entries[0].Name()), destDir); err != nil { + return err + } + return nil +} + +// extractTarball unpacks a gzipped tarball into destDir via `tar`, mirroring the +// snapshot packagers' reliance on the system tar (so symlink/permission handling is +// identical to what the run path produced when it built the archive). +func extractTarball(ctx context.Context, tarball, destDir string) error { + cmd := exec.CommandContext(ctx, "tar", "-xzf", tarball, "-C", destDir) + 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 extract code_source archive: %w: %s", err, msg) + } + return fmt.Errorf("failed to extract code_source archive: %w", err) + } + return nil +} + +// printConvertNextSteps tells the user what was written and the exact deploy +// sequence, since the value of the command is a one-command deploy afterwards. +// +// It also spells out cleanup. This matters specifically for AIR users: `air run` +// submits an ephemeral runs/submit workload that the platform reaps on its own, +// whereas `bundle deploy` creates a *persistent* job that lingers until explicitly +// destroyed — DABs has no automatic GC. A user migrating from `air run` will not +// expect a durable resource, so we call out `bundle destroy` explicitly. +func printConvertNextSteps(ctx context.Context, dir string, written []string, jobKey string, notes []string) { + cmdio.LogString(ctx, fmt.Sprintf("Wrote a Databricks Asset Bundle to %s:", dir)) + for _, w := range written { + cmdio.LogString(ctx, " "+w) + } + + // Notes surface anything the user should know: fields we transformed or dropped, + // and values they may need to fill in. Migrating users otherwise can't tell what + // silently changed between their run YAML and the bundle. + if len(notes) > 0 { + cmdio.LogString(ctx, "") + cmdio.LogString(ctx, "Notes:") + for _, n := range notes { + cmdio.LogString(ctx, " - "+n) + } + } + + cmdio.LogString(ctx, "") + cmdio.LogString(ctx, "To deploy and run this workload as a bundle:") + cmdio.LogString(ctx, " 1. cd "+dir) + cmdio.LogString(ctx, " 2. databricks bundle validate") + cmdio.LogString(ctx, " 3. databricks bundle deploy") + cmdio.LogString(ctx, " 4. databricks bundle run "+jobKey) + cmdio.LogString(ctx, "") + cmdio.LogString(ctx, "bundle deploy uploads the code source and launch scripts automatically.") + cmdio.LogString(ctx, "") + cmdio.LogString(ctx, "Unlike `air run` (which submits an ephemeral run), bundle deploy creates a") + cmdio.LogString(ctx, "persistent job that is not garbage-collected. When you are done, remove the") + cmdio.LogString(ctx, "job and its uploaded files with:") + cmdio.LogString(ctx, " databricks bundle destroy") +} + +// conversionNotes lists what the conversion transformed, staged out-of-band, or +// could not represent natively — so a user migrating from `air run` can see what +// changed between their run YAML and the emitted bundle, and what they may still +// need to fill in. Best-effort: git resolution errors are ignored here (writeBundle +// surfaces them), so a note is only emitted when the state is unambiguous. +func conversionNotes(cfg *runConfig) []string { + var notes []string + + if cfg.CodeSource != nil && cfg.CodeSource.Snapshot != nil { + snap := cfg.CodeSource.Snapshot + if snap.Git != nil { + notes = append(notes, "code_source.git was pinned in the run YAML; the pinned commit was materialized into "+ + "the code_source/ directory. Re-run convert-to-dabs to re-materialize a different revision.") + } + notes = append(notes, "code_source was staged into the code_source/ directory; bundle deploy packages and uploads it. "+ + "Re-run convert-to-dabs to re-stage after code changes.") + } + + // env vars / secrets have no native ai_runtime_task field yet, so they ride as + // sidecar files the server-side launcher reads (same as `air run`). + if len(cfg.EnvVariables) > 0 { + notes = append(notes, "env_variables were written to env_vars.json (no native bundle field yet); they are uploaded with the code and applied at run time.") + } + if len(cfg.Secrets) > 0 { + notes = append(notes, "secrets were written to secret_env_vars.json (no native bundle field yet); they are resolved at run time.") + } + if len(cfg.Parameters) > 0 { + notes = append(notes, "parameters were written to hyperparameters.yaml; they are not a native bundle field and are passed through to the workload.") + } + + return notes +} diff --git a/experimental/air/cmd/convert_to_dabs_test.go b/experimental/air/cmd/convert_to_dabs_test.go new file mode 100644 index 00000000000..51f49f94cba --- /dev/null +++ b/experimental/air/cmd/convert_to_dabs_test.go @@ -0,0 +1,410 @@ +package aircmd + +import ( + "os" + "path/filepath" + "strings" + "testing" + + "github.com/databricks/cli/libs/dyn" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestConvertToDabsCommandShape(t *testing.T) { + cmd := newConvertToDabsCommand() + assert.Equal(t, "convert-to-dabs ", cmd.Use) + assert.Empty(t, cmd.Commands(), "convert-to-dabs must not register subcommands") + // Exactly one positional (the YAML path). + assert.NoError(t, cmd.Args(cmd, []string{"run.yaml"})) + assert.Error(t, cmd.Args(cmd, []string{})) + assert.Error(t, cmd.Args(cmd, []string{"a", "b"})) +} + +// get is a small helper: read a dotted path out of the emitted bundle root. +func get(t *testing.T, root map[string]dyn.Value, path string) dyn.Value { + t.Helper() + v, err := dyn.GetByPath(dyn.V(root), dyn.MustPathFromString(path)) + require.NoError(t, err, "path %q should exist", path) + return v +} + +func has(root map[string]dyn.Value, path string) bool { + _, err := dyn.GetByPath(dyn.V(root), dyn.MustPathFromString(path)) + return err == nil +} + +// A full config maps onto a schema-shaped bundle: bundle name, job/task keys, the +// ai_runtime_task (experiment + single deployment + code_source_path), framework +// fields on the task wrapper, and the environment spec. +func TestConvertToDabsFullMapping(t *testing.T) { + cfg := minimalConfig + ` +max_retries: 2 +timeout_minutes: 30 +mlflow_run_name: run-42 +code_source: + type: snapshot + snapshot: + root_path: ./src +environment: + version: 5 + dependencies: + - numpy + - torch +` + path := writeConfigFile(t, "run.yaml", cfg) + loaded, err := loadRunConfig(path) + require.NoError(t, err) + + root, artifacts, err := convertToDabs(t.Context(), loaded, path) + require.NoError(t, err) + + name := loaded.ExperimentName + assert.Equal(t, name, get(t, root, "bundle.name").MustString()) + assert.Equal(t, "development", get(t, root, "targets.dev.mode").MustString()) + + jobPath := "resources.jobs." + name + assert.Equal(t, name, get(t, root, jobPath+".name").MustString()) + + task := jobPath + ".tasks[0]" + assert.Equal(t, name, get(t, root, task+".task_key").MustString()) + // Framework fields live on the task wrapper, not in ai_runtime_task. + assert.Equal(t, int64(2), get(t, root, task+".max_retries").MustInt()) + assert.Equal(t, int64(1800), get(t, root, task+".timeout_seconds").MustInt()) + assert.False(t, has(root, task+".ai_runtime_task.max_retries"), "retries must not be inside ai_runtime_task") + + art := task + ".ai_runtime_task" + assert.Equal(t, name, get(t, root, art+".experiment").MustString()) + assert.Equal(t, "run-42", get(t, root, art+".mlflow_run").MustString()) + // code_source_path is a local directory (packaged by the deploy-time aicode + // mutator), not a pre-built tarball. + assert.Equal(t, "./"+codeSourceDirName, get(t, root, art+".code_source_path").MustString()) + + dep := art + ".deployments[0]" + assert.Equal(t, "./"+commandScriptName, get(t, root, dep+".command_path").MustString()) + assert.Equal(t, "GPU_1xH100", get(t, root, dep+".compute.accelerator_type").MustString()) + assert.Equal(t, int64(1), get(t, root, dep+".compute.accelerator_count").MustInt()) + + env := jobPath + ".environments[0]" + assert.Equal(t, "default", get(t, root, env+".environment_key").MustString()) + assert.Equal(t, "5", get(t, root, env+".spec.environment_version").MustString()) + deps := get(t, root, env+".spec.dependencies").MustSequence() + require.Len(t, deps, 2) + assert.Equal(t, "numpy", deps[0].MustString()) + + // command.sh is always an artifact. requirements.yaml is NOT emitted: the + // deploy-time aicode.SynthesizeRequirements mutator regenerates it from the + // environments[] spec (asserted above), so convert must not also write it. + assert.Contains(t, itemNames(artifacts), commandScriptName) + assert.NotContains(t, itemNames(artifacts), requirementsName) +} + +// Optional fields are omitted rather than emitted empty: no code_source means no +// code_source_path; unset retries/timeout means no wrapper fields. +func TestConvertToDabsOmitsUnsetFields(t *testing.T) { + path := writeConfigFile(t, "run.yaml", minimalConfig) + loaded, err := loadRunConfig(path) + require.NoError(t, err) + + root, _, err := convertToDabs(t.Context(), loaded, path) + require.NoError(t, err) + + name := loaded.ExperimentName + task := "resources.jobs." + name + ".tasks[0]" + assert.False(t, has(root, task+".max_retries")) + assert.False(t, has(root, task+".timeout_seconds")) + assert.False(t, has(root, task+".ai_runtime_task.code_source_path")) + assert.False(t, has(root, task+".ai_runtime_task.mlflow_run")) + // The command still needs a home even without code_source. + assert.Equal(t, "./"+commandScriptName, get(t, root, task+".ai_runtime_task.deployments[0].command_path").MustString()) + + // Even with no environment block, the default runtime version is pinned (what + // `air run` would have used) rather than emitting an empty environment spec. + env := "resources.jobs." + name + ".environments[0]" + assert.Equal(t, "4", get(t, root, env+".spec.environment_version").MustString()) + assert.False(t, has(root, env+".spec.dependencies")) +} + +// A DATABRICKS_DL_RUNTIME_IMAGE env override flows through the same resolution +// `air run` uses, so a converted bundle pins the same version. +func TestConvertToDabsRuntimeVersionEnvOverride(t *testing.T) { + t.Setenv(dlRuntimeImageEnv, "CLIENT-GPU-7") + path := writeConfigFile(t, "run.yaml", minimalConfig) + loaded, err := loadRunConfig(path) + require.NoError(t, err) + + root, _, err := convertToDabs(t.Context(), loaded, path) + require.NoError(t, err) + + env := "resources.jobs." + loaded.ExperimentName + ".environments[0]" + assert.Equal(t, "7", get(t, root, env+".spec.environment_version").MustString()) +} + +// remote_volume can't be honored by a converted bundle (bundle deploy owns the +// artifact upload location), so it is rejected rather than silently ignored. +func TestConvertToDabsRejectsRemoteVolume(t *testing.T) { + cfg := minimalConfig + ` +code_source: + type: snapshot + snapshot: + root_path: ./src + remote_volume: /Volumes/main/default/code +` + path := writeConfigFile(t, "run.yaml", cfg) + loaded, err := loadRunConfig(path) + require.NoError(t, err) + _, _, err = convertToDabs(t.Context(), loaded, path) + require.ErrorContains(t, err, "remote_volume is not supported") +} + +// env_variables / secrets / parameters ride as sidecar files (the ai_runtime_task +// proto has no inline fields for them), matching the CLI's own launch layout. +func TestConvertToDabsStagesEnvAndSecretSidecars(t *testing.T) { + cfg := minimalConfig + ` +env_variables: + FOO: bar +secrets: + TOKEN: scope/key +parameters: + lr: 0.1 +` + path := writeConfigFile(t, "run.yaml", cfg) + loaded, err := loadRunConfig(path) + require.NoError(t, err) + + root, artifacts, err := convertToDabs(t.Context(), loaded, path) + require.NoError(t, err) + + names := itemNames(artifacts) + assert.Contains(t, names, envVarsName) + assert.Contains(t, names, secretEnvVarsName) + assert.Contains(t, names, hyperparametersName) + + // They are NOT smuggled into ai_runtime_task (which would fail bundle validate). + name := loaded.ExperimentName + art := "resources.jobs." + name + ".tasks[0].ai_runtime_task" + assert.False(t, has(root, art+".env_variables")) + assert.False(t, has(root, art+".secrets")) +} + +// A working-tree (non-git-pinned) code_source is staged into the bundle's +// code_source/ directory, honoring .gitignore just like the run path's plain-tar. +func TestConvertToDabsWorkingTreeMaterialized(t *testing.T) { + repo := t.TempDir() + writeRepoFile(t, repo, "train.py", "print('x')") + writeRepoFile(t, repo, "notes.log", "scratch") + writeRepoFile(t, repo, ".gitignore", "*.log\n") + + cfg := "experiment_name: wt\ncommand: python train.py\n" + + "compute: {accelerator_type: GPU_1xH100, num_accelerators: 1}\n" + + "code_source:\n type: snapshot\n snapshot:\n root_path: " + repo + "\n" + path := writeConfigFile(t, "run.yaml", cfg) + loaded, err := loadRunConfig(path) + require.NoError(t, err) + + outDir := t.TempDir() + _, err = writeBundle(t.Context(), loaded, path, outDir) + require.NoError(t, err) + + codeDir := filepath.Join(outDir, codeSourceDirName) + assert.FileExists(t, filepath.Join(codeDir, "train.py")) + assert.NoFileExists(t, filepath.Join(codeDir, "notes.log"), "gitignored files must be excluded from the staged code_source") +} + +// A requirements-FILE dependency set (environment.dependencies is a path) is folded +// into the environments[] spec so the deploy-time aicode mutator can regenerate +// requirements.yaml from it. Convert emits no requirements.yaml artifact of its own. +func TestConvertToDabsFoldsRequirementsFileIntoEnvSpec(t *testing.T) { + dir := t.TempDir() + reqPath := filepath.Join(dir, "requirements.yaml") + require.NoError(t, os.WriteFile(reqPath, []byte("version: \"6\"\ndependencies:\n - numpy\n - pandas\n"), 0o600)) + + cfg := "experiment_name: reqfile\ncommand: python train.py\n" + + "compute: {accelerator_type: GPU_1xH100, num_accelerators: 1}\n" + + "environment:\n dependencies: " + reqPath + "\n" + path := writeConfigFile(t, "run.yaml", cfg) + loaded, err := loadRunConfig(path) + require.NoError(t, err) + + root, artifacts, err := convertToDabs(t.Context(), loaded, path) + require.NoError(t, err) + + env := "resources.jobs." + loaded.ExperimentName + ".environments[0]" + assert.Equal(t, "6", get(t, root, env+".spec.environment_version").MustString()) + deps := get(t, root, env+".spec.dependencies").MustSequence() + require.Len(t, deps, 2) + assert.Equal(t, "numpy", deps[0].MustString()) + assert.Equal(t, "pandas", deps[1].MustString()) + + // No requirements.yaml artifact: the mutator regenerates it from the spec. + assert.NotContains(t, itemNames(artifacts), requirementsName) +} + +// conversionNotes surfaces what was transformed/staged so a migrating user knows +// what changed between their run YAML and the bundle. +func TestConvertToDabsConversionNotes(t *testing.T) { + cfg := minimalConfig + ` +env_variables: {FOO: bar} +secrets: {TOKEN: scope/key} +parameters: {lr: 0.1} +code_source: + type: snapshot + snapshot: + root_path: ./src + git: {commit: abc123} +` + path := writeConfigFile(t, "run.yaml", cfg) + loaded, err := loadRunConfig(path) + require.NoError(t, err) + + notes := conversionNotes(loaded) + joined := strings.Join(notes, "\n") + assert.Contains(t, joined, "code_source.git") // git pin flagged + assert.Contains(t, joined, "code_source/") // staged-directory behavior + assert.Contains(t, joined, "env_vars.json") // env vars staged + assert.Contains(t, joined, "secret_env_vars.json") // secrets staged + assert.Contains(t, joined, "hyperparameters.yaml") // parameters staged + + // A minimal config with none of those has no notes. + base := writeConfigFile(t, "min.yaml", minimalConfig) + minCfg, err := loadRunConfig(base) + require.NoError(t, err) + assert.Empty(t, conversionNotes(minCfg)) +} + +// usage_policy_id is a resolved budget policy id and maps to the job's +// budget_policy_id (usage_policy_name, which needs resolution, is rejected). +func TestConvertToDabsMapsUsagePolicyID(t *testing.T) { + cfg := minimalConfig + "usage_policy_id: budget-abc-123\n" + path := writeConfigFile(t, "run.yaml", cfg) + loaded, err := loadRunConfig(path) + require.NoError(t, err) + + root, _, err := convertToDabs(t.Context(), loaded, path) + require.NoError(t, err) + assert.Equal(t, "budget-abc-123", get(t, root, "resources.jobs."+loaded.ExperimentName+".budget_policy_id").MustString()) +} + +func TestConvertToDabsMapsPermissions(t *testing.T) { + cfg := minimalConfig + ` +permissions: + - user_name: alice@example.com + level: CAN_MANAGE + - group_name: eng + level: CAN_VIEW +` + path := writeConfigFile(t, "run.yaml", cfg) + loaded, err := loadRunConfig(path) + require.NoError(t, err) + + root, _, err := convertToDabs(t.Context(), loaded, path) + require.NoError(t, err) + + perms := get(t, root, "resources.jobs."+loaded.ExperimentName+".permissions").MustSequence() + require.Len(t, perms, 2) + assert.Equal(t, "CAN_MANAGE", perms[0].Get("level").MustString()) + assert.Equal(t, "alice@example.com", perms[0].Get("user_name").MustString()) + assert.Equal(t, "eng", perms[1].Get("group_name").MustString()) +} + +// A git-pinned code_source is materialized from the commit, not the dirty working +// tree: writeBundle stages a code_source/ directory holding the committed file only. +func TestConvertToDabsGitPinnedMaterialized(t *testing.T) { + repo := newTestRepo(t) + writeRepoFile(t, repo, "train.py", "print('committed')") + sha := commitAll(t, repo, "init") + // Dirty the tree AFTER the commit; the pinned snapshot must not include this. + writeRepoFile(t, repo, "uncommitted.py", "print('dirty')") + + cfg := "experiment_name: git-pin\ncommand: python train.py\n" + + "compute: {accelerator_type: GPU_1xH100, num_accelerators: 1}\n" + + "code_source:\n type: snapshot\n snapshot:\n root_path: " + repo + "\n git:\n commit: " + sha + "\n" + path := writeConfigFile(t, "run.yaml", cfg) + loaded, err := loadRunConfig(path) + require.NoError(t, err) + + outDir := t.TempDir() + _, err = writeBundle(t.Context(), loaded, path, outDir) + require.NoError(t, err) + + codeDir := filepath.Join(outDir, codeSourceDirName) + assert.FileExists(t, filepath.Join(codeDir, "train.py")) + assert.NoFileExists(t, filepath.Join(codeDir, "uncommitted.py"), "git-pinned snapshot must exclude uncommitted files") +} + +// git archive runs with `git -C repoPath`, so the staging tarball path must be +// absolute — otherwise `-o out/...` resolves against the repo dir and fails. Exercise +// writeBundle from a working dir with a RELATIVE output path against a git-pinned source. +func TestConvertToDabsGitPinnedRelativeOutputDir(t *testing.T) { + repo := newTestRepo(t) + writeRepoFile(t, repo, "train.py", "print('x')") + sha := commitAll(t, repo, "init") + + cfg := "experiment_name: rel-out\ncommand: python train.py\n" + + "compute: {accelerator_type: GPU_1xH100, num_accelerators: 1}\n" + + "code_source:\n type: snapshot\n snapshot:\n root_path: " + repo + "\n git:\n commit: " + sha + "\n" + path := writeConfigFile(t, "run.yaml", cfg) + loaded, err := loadRunConfig(path) + require.NoError(t, err) + + // chdir into a scratch dir and pass a relative --output-dir. + work := t.TempDir() + t.Chdir(work) + _, err = writeBundle(t.Context(), loaded, path, "out") + require.NoError(t, err) + + assert.FileExists(t, filepath.Join(work, "out", codeSourceDirName, "train.py")) +} + +// A numeric or reserved-word experiment_name would be an unquoted YAML map key +// that DABs' strict loader rejects (!!int / !!bool). The job resource key is +// prefixed to stay a string, while name/experiment keep the original value. +func TestConvertToDabsSafeJobKey(t *testing.T) { + cases := map[string]string{ + "12345": "job_12345", + "1.5e3": "job_1.5e3", + "true": "job_true", + "null": "job_null", + } + for name, wantKey := range cases { + assert.Equal(t, wantKey, bundleResourceKey(name), "key for %q", name) + } + // A normal name is used as-is. + assert.Equal(t, "my-run_1", bundleResourceKey("my-run_1")) + + // End to end: a numeric name lands under the prefixed key, but name/experiment + // keep the numeric string value. + cfg := "experiment_name: \"12345\"\ncommand: python t.py\n" + + "compute: {accelerator_type: GPU_1xH100, num_accelerators: 1}\n" + path := writeConfigFile(t, "run.yaml", cfg) + loaded, err := loadRunConfig(path) + require.NoError(t, err) + root, _, err := convertToDabs(t.Context(), loaded, path) + require.NoError(t, err) + assert.Equal(t, "12345", get(t, root, "resources.jobs.job_12345.name").MustString()) + assert.Equal(t, "12345", get(t, root, "resources.jobs.job_12345.tasks[0].ai_runtime_task.experiment").MustString()) +} + +func TestConvertToDabsRejectsUnsupported(t *testing.T) { + t.Run("docker_image", func(t *testing.T) { + cfg := minimalConfig + ` +environment: + docker_image: + url: myregistry/img:tag +` + path := writeConfigFile(t, "run.yaml", cfg) + loaded, err := loadRunConfig(path) + require.NoError(t, err) + _, _, err = convertToDabs(t.Context(), loaded, path) + require.ErrorContains(t, err, "docker_image is not yet supported") + }) + + t.Run("usage_policy_name", func(t *testing.T) { + cfg := minimalConfig + "usage_policy_name: my-policy\n" + path := writeConfigFile(t, "run.yaml", cfg) + loaded, err := loadRunConfig(path) + require.NoError(t, err) + _, _, err = convertToDabs(t.Context(), loaded, path) + require.ErrorContains(t, err, "usage_policy_name is not yet supported") + }) +} From 0b5b700e1c147d5d0adb80cfaa20abdc465f7bc5 Mon Sep 17 00:00:00 2001 From: vinchenzo-db Date: Fri, 31 Jul 2026 06:44:45 +0000 Subject: [PATCH 02/15] air/convert-to-dabs: fix Windows code_source staging (tar drive-letter) The Windows CI job failed the convert-to-dabs acceptance test with "tar (child): Cannot connect to C: resolve failed": the system tar reads the `C:` in an absolute archive path as a remote host:path. Two fixes: - createPlainTarball now passes the archive as a bare basename with cmd.Dir set to the output directory (and an absolute parent), so no `C:\...` path reaches tar's -f argument. Mirrors how git archive is invoked; safe on GNU tar and bsdtar. (This helper is shared with the `air run` snapshot path.) - extractTarball is rewritten in pure Go (archive/tar + compress/gzip) instead of shelling out to `tar -xzf`, eliminating the same drive-letter hazard on the extract side and dropping the external-tar dependency for extraction. It rejects entries that would escape the destination (path traversal, absolute/escaping symlinks) and bounds each file copy to its header size. Adds unit tests for extractTarball (happy path incl. nested dirs + in-tree symlink; traversal + escaping-symlink rejection). Co-authored-by: Isaac --- experimental/air/cmd/convert_to_dabs.go | 102 ++++++++++++++++--- experimental/air/cmd/convert_to_dabs_test.go | 66 ++++++++++++ experimental/air/cmd/snapshot_package.go | 25 ++++- 3 files changed, 176 insertions(+), 17 deletions(-) diff --git a/experimental/air/cmd/convert_to_dabs.go b/experimental/air/cmd/convert_to_dabs.go index 2f13ff99612..f6dfb125584 100644 --- a/experimental/air/cmd/convert_to_dabs.go +++ b/experimental/air/cmd/convert_to_dabs.go @@ -1,12 +1,13 @@ package aircmd import ( - "bytes" + "archive/tar" + "compress/gzip" "context" "errors" "fmt" + "io" "os" - "os/exec" "path/filepath" "slices" "strconv" @@ -480,7 +481,7 @@ func materializeCodeSource(ctx context.Context, repoPath string, plan snapshotPl if err := os.MkdirAll(extractDir, 0o700); err != nil { return err } - if err := extractTarball(ctx, tarball, extractDir); err != nil { + if err := extractTarball(tarball, extractDir); err != nil { return err } entries, err := os.ReadDir(extractDir) @@ -496,18 +497,91 @@ func materializeCodeSource(ctx context.Context, repoPath string, plan snapshotPl return nil } -// extractTarball unpacks a gzipped tarball into destDir via `tar`, mirroring the -// snapshot packagers' reliance on the system tar (so symlink/permission handling is -// identical to what the run path produced when it built the archive). -func extractTarball(ctx context.Context, tarball, destDir string) error { - cmd := exec.CommandContext(ctx, "tar", "-xzf", tarball, "-C", destDir) - 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 extract code_source archive: %w: %s", err, msg) +// extractTarball unpacks a gzipped tarball into destDir in pure Go (no `tar` +// subprocess), so extraction is portable — notably on Windows, where handing a +// `C:\...` path to the system tar makes it read the drive letter as a remote host. +// Entries that would escape destDir (path traversal, absolute or escaping symlinks) +// are rejected; unusual entry types are skipped. +func extractTarball(tarball, destDir string) error { + f, err := os.Open(tarball) + if err != nil { + return err + } + defer f.Close() + + gz, err := gzip.NewReader(f) + if err != nil { + return fmt.Errorf("failed to read code_source archive: %w", err) + } + defer gz.Close() + + // destDir is the boundary every extracted entry must stay within: a malicious or + // malformed archive entry ("../x", an absolute path, a symlink escaping the tree) + // must not write outside it. destAbs is compared against each resolved target. + destAbs, err := filepath.Abs(destDir) + if err != nil { + return err + } + + tr := tar.NewReader(gz) + for { + hdr, err := tr.Next() + if errors.Is(err, io.EOF) { + break + } + if err != nil { + return fmt.Errorf("failed to read code_source archive: %w", err) + } + + // Reject path traversal: the cleaned target must stay under destAbs. + target := filepath.Join(destAbs, filepath.FromSlash(hdr.Name)) + if target != destAbs && !strings.HasPrefix(target, destAbs+string(os.PathSeparator)) { + return fmt.Errorf("code_source archive entry %q escapes the destination directory", hdr.Name) + } + + switch hdr.Typeflag { + case tar.TypeDir: + if err := os.MkdirAll(target, 0o700); err != nil { + return err + } + case tar.TypeReg: + if err := os.MkdirAll(filepath.Dir(target), 0o700); err != nil { + return err + } + out, err := os.OpenFile(target, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, os.FileMode(hdr.Mode)&0o777) + if err != nil { + return err + } + // Bounded copy: cap the number of bytes written to the size the header + // declares, so a corrupt/oversized entry can't write unbounded data. + if _, err := io.CopyN(out, tr, hdr.Size); err != nil && !errors.Is(err, io.EOF) { + out.Close() + return err + } + if err := out.Close(); err != nil { + return err + } + case tar.TypeSymlink: + if err := os.MkdirAll(filepath.Dir(target), 0o700); err != nil { + return err + } + // A symlink whose resolved target escapes destAbs is rejected: it could be + // followed on a later write to reach outside the staged code directory. + linkTarget := hdr.Linkname + resolved := linkTarget + if !filepath.IsAbs(resolved) { + resolved = filepath.Join(filepath.Dir(target), filepath.FromSlash(linkTarget)) + } + if resolved != destAbs && !strings.HasPrefix(resolved, destAbs+string(os.PathSeparator)) { + return fmt.Errorf("code_source archive symlink %q -> %q escapes the destination directory", hdr.Name, linkTarget) + } + if err := os.Symlink(linkTarget, target); err != nil { + return err + } + default: + // Skip other entry types (devices, fifos, etc.): code source is regular + // files, dirs, and symlinks; anything else is not meaningful to a workload. } - return fmt.Errorf("failed to extract code_source archive: %w", err) } return nil } diff --git a/experimental/air/cmd/convert_to_dabs_test.go b/experimental/air/cmd/convert_to_dabs_test.go index 51f49f94cba..9e8e80d235c 100644 --- a/experimental/air/cmd/convert_to_dabs_test.go +++ b/experimental/air/cmd/convert_to_dabs_test.go @@ -1,6 +1,8 @@ package aircmd import ( + "archive/tar" + "compress/gzip" "os" "path/filepath" "strings" @@ -356,6 +358,70 @@ func TestConvertToDabsGitPinnedRelativeOutputDir(t *testing.T) { assert.FileExists(t, filepath.Join(work, "out", codeSourceDirName, "train.py")) } +// writeTarball writes a gzipped tar of the given entries to path. Each entry is +// either a regular file (typeflag defaults to file) or, when linkname != "", a +// symlink. Used to exercise extractTarball directly. +func writeTarball(t *testing.T, path string, entries []tar.Header) { + t.Helper() + f, err := os.Create(path) + require.NoError(t, err) + gz := gzip.NewWriter(f) + tw := tar.NewWriter(gz) + for _, h := range entries { + require.NoError(t, tw.WriteHeader(&h)) + if h.Typeflag == tar.TypeReg { + _, err := tw.Write([]byte("data-" + h.Name)) + require.NoError(t, err) + } + } + require.NoError(t, tw.Close()) + require.NoError(t, gz.Close()) + require.NoError(t, f.Close()) +} + +// extractTarball unpacks a well-formed archive (files, nested dirs, in-tree +// symlink) into destDir. +func TestExtractTarballHappyPath(t *testing.T) { + src := filepath.Join(t.TempDir(), "a.tar.gz") + writeTarball(t, src, []tar.Header{ + {Name: "code/", Typeflag: tar.TypeDir, Mode: 0o755}, + {Name: "code/train.py", Typeflag: tar.TypeReg, Mode: 0o644, Size: int64(len("data-code/train.py"))}, + {Name: "code/pkg/", Typeflag: tar.TypeDir, Mode: 0o755}, + {Name: "code/pkg/util.py", Typeflag: tar.TypeReg, Mode: 0o644, Size: int64(len("data-code/pkg/util.py"))}, + {Name: "code/link.py", Typeflag: tar.TypeSymlink, Linkname: "train.py"}, + }) + + dest := t.TempDir() + require.NoError(t, extractTarball(src, dest)) + assert.FileExists(t, filepath.Join(dest, "code", "train.py")) + assert.FileExists(t, filepath.Join(dest, "code", "pkg", "util.py")) + if info, err := os.Lstat(filepath.Join(dest, "code", "link.py")); assert.NoError(t, err) { + assert.NotZero(t, info.Mode()&os.ModeSymlink, "link.py should be a symlink") + } +} + +// A path-traversal entry or an escaping symlink is rejected rather than written +// outside destDir. +func TestExtractTarballRejectsEscape(t *testing.T) { + t.Run("traversal path", func(t *testing.T) { + src := filepath.Join(t.TempDir(), "evil.tar.gz") + writeTarball(t, src, []tar.Header{ + {Name: "../escape.py", Typeflag: tar.TypeReg, Mode: 0o644, Size: int64(len("data-../escape.py"))}, + }) + err := extractTarball(src, t.TempDir()) + require.ErrorContains(t, err, "escapes the destination directory") + }) + + t.Run("escaping symlink", func(t *testing.T) { + src := filepath.Join(t.TempDir(), "evil-link.tar.gz") + writeTarball(t, src, []tar.Header{ + {Name: "link", Typeflag: tar.TypeSymlink, Linkname: "../../etc/passwd"}, + }) + err := extractTarball(src, t.TempDir()) + require.ErrorContains(t, err, "escapes the destination directory") + }) +} + // A numeric or reserved-word experiment_name would be an unquoted YAML map key // that DABs' strict loader rejects (!!int / !!bool). The job resource key is // prefixed to stay a string, while name/experiment keep the original value. diff --git a/experimental/air/cmd/snapshot_package.go b/experimental/air/cmd/snapshot_package.go index 672366086c9..b81f33a65b1 100644 --- a/experimental/air/cmd/snapshot_package.go +++ b/experimental/air/cmd/snapshot_package.go @@ -42,9 +42,24 @@ func createGitArchiveSnapshot(ctx context.Context, git gitRepo, commitSHA, outpu // 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) + // Absolute so it resolves correctly regardless of tar's working dir (set below). + parent, err := filepath.Abs(filepath.Dir(repoPath)) + if err != nil { + return err + } + + // Pass the archive path relative to its own directory (run tar there), never a + // full path: on Windows an absolute path like `C:\out\x.tar.gz` makes tar read + // the `C:` as a remote host ("Cannot connect to C:"), since tar treats a colon + // in the -f arg as host:path. A bare basename with -C avoids that on GNU tar and + // bsdtar alike. + outDirAbs, err := filepath.Abs(filepath.Dir(outputTarball)) + if err != nil { + return err + } + outName := filepath.Base(outputTarball) - args := []string{"-czf", outputTarball} + args := []string{"-czf", outName} // Exclude macOS AppleDouble files: they sort before the real top-level dir and // hijack a remote `head -1` parse. No-op on Linux. @@ -68,7 +83,9 @@ 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). + // prefix each so entries nest under it (matching git archive --prefix). -C only + // affects the file operands that follow it, not the -f archive path (which + // resolves against tar's working dir, set to outDirAbs below). args = append(args, "-C", parent) if len(includePaths) > 0 { for _, p := range includePaths { @@ -79,6 +96,8 @@ func createPlainTarball(ctx context.Context, repoPath, outputTarball string, inc } cmd := exec.CommandContext(ctx, "tar", args...) + // Run tar in the output directory so the bare -f basename lands there. + cmd.Dir = outDirAbs var stderr bytes.Buffer cmd.Stderr = &stderr if err := cmd.Run(); err != nil { From 69140b096a1a49f6a2d9f14a52dc61e367a8aa41 Mon Sep 17 00:00:00 2001 From: vinchenzo-db Date: Mon, 3 Aug 2026 20:39:54 +0000 Subject: [PATCH 03/15] air/convert-to-dabs: pure YAML conversion, no code packaging MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit convert-to-dabs is now a purely local, syntactic translation. It no longer copies, snapshots, or git-archives the code source — the deploy-time aicode mutator (from the base PR) packages the source in place at `bundle deploy`. - code_source_path is emitted as the source directory relative to the bundle (bundle root defaults to the input YAML's directory, which contains it), rather than a copied ./code_source dir. - Removed the materialize/extract-tarball machinery and the git-pinned-commit path. code_source.snapshot.git is now rejected (deploy packages the working tree); a code_source outside the bundle directory is rejected with guidance. - writeBundle now only writes databricks.yml + command.sh + the env/secret/param sidecars. Dependencies stay folded into the environments[] spec. Tests updated: assert code_source_path points at ./src with no copy, reject git-pin and out-of-bundle sources; the acceptance test converts in place. Co-authored-by: Isaac --- .../air/convert-to-dabs/output.txt | 29 +- .../experimental/air/convert-to-dabs/script | 16 +- .../air/convert-to-dabs/test.toml | 6 +- experimental/air/cmd/convert_to_dabs.go | 353 +++++------------- experimental/air/cmd/convert_to_dabs_test.go | 214 ++++------- 5 files changed, 180 insertions(+), 438 deletions(-) diff --git a/acceptance/experimental/air/convert-to-dabs/output.txt b/acceptance/experimental/air/convert-to-dabs/output.txt index efcb78c4ca8..2b3106c4c0d 100644 --- a/acceptance/experimental/air/convert-to-dabs/output.txt +++ b/acceptance/experimental/air/convert-to-dabs/output.txt @@ -1,17 +1,16 @@ -=== convert an AIR run YAML into a DABs bundle ->>> [CLI] experimental air convert-to-dabs train.yaml --output-dir generated -Wrote a Databricks Asset Bundle to generated: +=== convert an AIR run YAML into a DABs bundle (in place, next to the source) +>>> [CLI] experimental air convert-to-dabs train.yaml +Wrote a Databricks Asset Bundle to .: databricks.yml training_config.yaml command.sh - code_source/ Notes: - - code_source was staged into the code_source/ directory; bundle deploy packages and uploads it. Re-run convert-to-dabs to re-stage after code changes. + - code_source points at your source directory; bundle deploy packages and uploads it from there. To deploy and run this workload as a bundle: - 1. cd generated + 1. cd . 2. databricks bundle validate 3. databricks bundle deploy 4. databricks bundle run torchrun-a10-smoke-test @@ -23,8 +22,8 @@ persistent job that is not garbage-collected. When you are done, remove the job and its uploaded files with: databricks bundle destroy -=== emitted databricks.yml ->>> cat generated/databricks.yml +=== emitted databricks.yml (code_source_path points at ./src; no code is copied) +>>> cat databricks.yml bundle: name: torchrun-a10-smoke-test targets: @@ -45,7 +44,7 @@ resources: compute: accelerator_type: GPU_1xA10 accelerator_count: 1 - code_source_path: ./code_source + code_source_path: ./src environments: - environment_key: default spec: @@ -54,18 +53,8 @@ resources: - numpy === the generated command.sh carries the run command ->>> cat generated/command.sh +>>> cat command.sh torchrun --nproc_per_node=1 train.py -=== the code source is staged as a directory (packaged at deploy time) ->>> ls generated -code_source -command.sh -databricks.yml -training_config.yaml - ->>> ls generated/code_source -train.py - === the emitted bundle validates >>> [CLI] bundle validate Name: torchrun-a10-smoke-test diff --git a/acceptance/experimental/air/convert-to-dabs/script b/acceptance/experimental/air/convert-to-dabs/script index 1a9f9cf07c8..2babd7c0aef 100644 --- a/acceptance/experimental/air/convert-to-dabs/script +++ b/acceptance/experimental/air/convert-to-dabs/script @@ -1,20 +1,14 @@ -title "convert an AIR run YAML into a DABs bundle" -trace $CLI experimental air convert-to-dabs train.yaml --output-dir generated +title "convert an AIR run YAML into a DABs bundle (in place, next to the source)" +trace $CLI experimental air convert-to-dabs train.yaml -title "emitted databricks.yml" -trace cat generated/databricks.yml +title "emitted databricks.yml (code_source_path points at ./src; no code is copied)" +trace cat databricks.yml title "the generated command.sh carries the run command" -trace cat generated/command.sh - -title "the code source is staged as a directory (packaged at deploy time)" -trace ls generated && true -trace ls generated/code_source && true +trace cat command.sh title "the emitted bundle validates" -cd generated trace $CLI bundle validate -cd .. title "docker_image is not supported yet" errcode trace $CLI experimental air convert-to-dabs docker.yaml --output-dir generated-docker diff --git a/acceptance/experimental/air/convert-to-dabs/test.toml b/acceptance/experimental/air/convert-to-dabs/test.toml index 19643c87452..386f44a60f0 100644 --- a/acceptance/experimental/air/convert-to-dabs/test.toml +++ b/acceptance/experimental/air/convert-to-dabs/test.toml @@ -1,3 +1,3 @@ -# The command writes a bundle under generated/; those are generated artifacts, -# not committed inputs, so exclude them from the repo-diff check. -Ignore = ["generated", "generated-docker"] +# convert-to-dabs writes the bundle in place (next to train.yaml + src/). These are +# generated artifacts, not committed inputs, so exclude them from the repo-diff check. +Ignore = ["databricks.yml", "command.sh", "training_config.yaml", "generated-docker"] diff --git a/experimental/air/cmd/convert_to_dabs.go b/experimental/air/cmd/convert_to_dabs.go index f6dfb125584..143e896713d 100644 --- a/experimental/air/cmd/convert_to_dabs.go +++ b/experimental/air/cmd/convert_to_dabs.go @@ -1,12 +1,9 @@ package aircmd import ( - "archive/tar" - "compress/gzip" "context" "errors" "fmt" - "io" "os" "path/filepath" "slices" @@ -24,29 +21,19 @@ import ( // convert_to_dabs turns an AIR CLI run YAML into a Databricks Asset Bundle so a // workload authored for `air run` can be deployed and managed as a bundle. // -// The emitted bundle is schema-valid: `databricks bundle validate` accepts it, -// and `databricks bundle deploy` reproduces the same ai_runtime_task workload the -// CLI would submit. Two properties make that work without any manual upload step: +// It is a purely local, syntactic translation: it maps the run config onto a +// schema-valid ai_runtime_task (the SDK jobs.AiRuntimeTask — experiment + +// deployments[].{command_path,compute} + code_source_path, with framework fields +// like retries/timeout on the surrounding task) and writes command.sh plus the +// env_vars.json / secret_env_vars.json / hyperparameters.yaml sidecars at the +// bundle root. It does NOT package, snapshot, or upload anything. // -// - The DABs `ai_runtime_task` is the SDK jobs.AiRuntimeTask (strict schema: -// experiment + deployments[].{command_path,compute} + code_source_path). -// Framework concerns (retries, timeout) live on the surrounding task, not in -// ai_runtime_task — so they are emitted as task-level fields. -// - code_source_path points at a *local directory* inside the bundle. The -// deploy-time aicode.PackageAndUpload mutator (bundle/config/mutator/aicode) -// owns the snapshotting: at `bundle deploy` it packages that directory into a -// content-addressed tarball, uploads it, and rewrites code_source_path to the -// remote archive. convert-to-dabs' job is only to *stage the source bytes* as -// that directory — copying the working tree, or materializing a pinned git -// commit into it — never to build the tarball itself. -// -// command.sh (and — since the task proto carries no inline env/secrets/parameters — -// the env_vars.json / secret_env_vars.json / hyperparameters.yaml sidecars the -// server-side launcher reads) are written at the bundle root and uploaded by -// deploy, mirroring the CLI's own launch layout. requirements.yaml is NOT emitted: -// the aicode.SynthesizeRequirements mutator derives it from the job's -// environments[] spec at deploy time, so convert folds the dependency set into that -// spec instead. +// code_source_path is emitted as the source directory relative to the bundle (the +// bundle root defaults to the YAML's directory, which contains it). At deploy the +// aicode mutator (bundle/config/mutator/aicode) packages that directory and uploads +// it — so convert never touches the code. Dependencies are folded into the job's +// environments[] spec, which the runtime installs from directly; no requirements.yaml +// is emitted. // dabsTargetName is the single default target emitted; a development-mode target // is the conventional starting point for a generated bundle. @@ -82,13 +69,14 @@ does not contact the workspace.`, return err } - // Default the bundle next to the input YAML (a -bundle subfolder - // so the generated files don't scatter across the user's project dir). An - // explicit --output-dir — absolute or relative — overrides. We never write to - // a temp dir: the bundle is the user's artifact to keep, deploy, and manage. + // Default the bundle to the input YAML's directory. The bundle's sync root + // must contain the code_source so `code_source_path` resolves within it (the + // deploy-time aicode mutator packages the source in place), and root_path is + // resolved relative to the YAML, so the YAML's dir is the natural bundle root. + // An explicit --output-dir overrides. dir := outputDir if dir == "" { - dir = filepath.Join(filepath.Dir(yamlPath), cfg.ExperimentName+"-bundle") + dir = filepath.Dir(yamlPath) } written, err := writeBundle(ctx, cfg, yamlPath, dir) @@ -103,20 +91,13 @@ does not contact the workspace.`, return cmd } -// codeSourceDirName is the bundle-local directory the code_source is staged into. -// ai_runtime_task.code_source_path points at it; the deploy-time aicode mutator -// packages the directory into a tarball and uploads it (see the file header). A -// fixed name (rather than the source's basename) keeps the emitted databricks.yml -// deterministic regardless of where the user's code lives. -const codeSourceDirName = "code_source" - // convertToDabs builds the DABs bundle value and the loose launch artifacts for a // run config. It reads only what the run path's buildArtifacts reads, so the // mapping is unit-testable in isolation. Returns the bundle root as a // map[string]dyn.Value (ready for yamlsaver) and the loose artifacts (command.sh + -// env/secret/param sidecars) to write at the bundle root; the code_source directory -// is materialized separately by writeBundle. -func convertToDabs(ctx context.Context, cfg *runConfig, configPath string) (map[string]dyn.Value, []uploadItem, error) { +// env/secret/param sidecars) to write at the bundle root. It does not touch the +// code_source; the deploy-time aicode mutator packages it in place. +func convertToDabs(ctx context.Context, cfg *runConfig, configPath, bundleDir string) (map[string]dyn.Value, []uploadItem, error) { // idempotency_token is intentionally not mapped: it dedups a single runs/submit // call, which has no analogue for a persistent, repeatedly-runnable bundle job. // @@ -128,28 +109,70 @@ func convertToDabs(ctx context.Context, cfg *runConfig, configPath string) (map[ if cfg.Environment != nil && cfg.Environment.DockerImage != nil { return nil, nil, errors.New("environment.docker_image is not yet supported by convert-to-dabs") } - // remote_volume points the code archive at a UC Volume. bundle deploy uploads - // code_source_path to the bundle's artifact path, not an arbitrary Volume, so a - // converted bundle can't honor it — reject rather than silently drop it. - if cfg.CodeSource != nil && cfg.CodeSource.Snapshot != nil && cfg.CodeSource.Snapshot.RemoteVolume != nil { - return nil, nil, errors.New("code_source.snapshot.remote_volume is not supported by convert-to-dabs; bundle deploy manages the artifact upload location") + if snap := codeSnapshot(cfg); snap != nil { + // remote_volume points the code archive at a UC Volume; bundle deploy uploads + // code_source to the bundle artifact path, not an arbitrary Volume. + if snap.RemoteVolume != nil { + return nil, nil, errors.New("code_source.snapshot.remote_volume is not supported by convert-to-dabs") + } + // git pinning would require materializing the pinned commit, which convert no + // longer does — the deploy-time mutator packages the working tree in place. + if snap.Git != nil { + return nil, nil, errors.New("code_source.snapshot.git is not supported by convert-to-dabs; deploy packages the working tree") + } + } + + codeSourcePath, err := bundleCodeSourcePath(ctx, cfg, configPath, bundleDir) + if err != nil { + return nil, nil, err } artifacts, err := buildArtifacts(cfg, configPath) if err != nil { return nil, nil, err } - // Drop requirements.yaml: the deploy-time aicode.SynthesizeRequirements mutator - // regenerates it from the job's environments[] spec (which convert populates with - // the same dependency set), so emitting it here would be redundant and could drift. + // Drop requirements.yaml: the runtime installs pip deps from the job's + // environments[] spec (which convert populates), so a sidecar would be redundant. artifacts = slices.DeleteFunc(artifacts, func(it uploadItem) bool { return it.name == requirementsName }) - root := buildBundleValue(ctx, cfg, configPath) + root := buildBundleValue(ctx, cfg, configPath, codeSourcePath) return root, artifacts, nil } +// codeSnapshot returns the snapshot code source config, or nil if none. +func codeSnapshot(cfg *runConfig) *snapshotSourceConfig { + if cfg.CodeSource == nil { + return nil + } + return cfg.CodeSource.Snapshot +} + +// bundleCodeSourcePath resolves the code_source directory to a "./"-prefixed path +// relative to the bundle dir, for emission as ai_runtime_task.code_source_path. +// Returns "" when the config has no code_source. The path must be inside the bundle +// (the deploy-time mutator packages it in place and only handles in-bundle dirs). +func bundleCodeSourcePath(ctx context.Context, cfg *runConfig, configPath, bundleDir string) (string, error) { + snap := codeSnapshot(cfg) + if snap == nil { + return "", nil + } + root, err := resolveRootPath(ctx, snap.RootPath, filepath.Dir(configPath)) + if err != nil { + return "", err + } + bundleAbs, err := filepath.Abs(bundleDir) + if err != nil { + return "", err + } + rel, err := filepath.Rel(bundleAbs, root) + if err != nil || !filepath.IsLocal(rel) { + return "", fmt.Errorf("code_source root_path %q is not inside the bundle directory %q; run convert-to-dabs with --output-dir set to an ancestor of the code", snap.RootPath, bundleDir) + } + return localBundlePath(filepath.ToSlash(rel)), nil +} + // nv builds a dyn.Value at "position" n: yamlsaver orders a map's keys by their // Location line, so assigning ascending n values fixes the emitted key order. // It routes through dyn.V so nested Go maps/slices are converted recursively, @@ -167,18 +190,17 @@ func localBundlePath(p string) string { } // buildBundleValue assembles the bundle root as an ordered map[string]dyn.Value. -// command.sh and the code_source directory are bundle-local (emitted "./"-prefixed) -// so `bundle deploy` uploads/packages them. -func buildBundleValue(ctx context.Context, cfg *runConfig, configPath string) map[string]dyn.Value { +// codeSourcePath is the "./"-prefixed code_source dir relative to the bundle (empty +// when the config has no code_source); command.sh is a bundle-local artifact. +func buildBundleValue(ctx context.Context, cfg *runConfig, configPath, codeSourcePath string) map[string]dyn.Value { name := cfg.ExperimentName // ai_runtime_task: experiment + one deployment (command_path + compute) + // code_source_path. Only the fields the strict schema allows. // - // Paths are emitted "./"-prefixed so bundle deploy treats them as LOCAL and - // uploads them: libraries.IsLibraryLocal classifies a bare, extensionless path - // as a PyPI package name (not a local file) and skips it, which would deploy a - // path the backend can't resolve. The "./" prefix forces local classification. + // command_path is "./"-prefixed so bundle deploy treats it as LOCAL and uploads + // it: libraries.IsLibraryLocal classifies a bare, extensionless path as a PyPI + // package name, which would deploy a path the backend can't resolve. deployment := map[string]dyn.Value{ "command_path": nv(localBundlePath(commandScriptName), 1), "compute": nv(map[string]dyn.Value{ @@ -192,10 +214,10 @@ func buildBundleValue(ctx context.Context, cfg *runConfig, configPath string) ma "deployments": nv([]dyn.Value{dyn.V(deployment)}, 2), } line := 3 - if cfg.CodeSource != nil && cfg.CodeSource.Snapshot != nil { - // A local directory (not a "./"-prefixed file): the aicode mutator recognizes - // a directory code_source_path, packages it, and rewrites this field at deploy. - aiRuntimeTask["code_source_path"] = nv(localBundlePath(codeSourceDirName), line) + if codeSourcePath != "" { + // The source dir relative to the bundle; the aicode mutator packages it at + // deploy and rewrites this field to the uploaded workspace path. + aiRuntimeTask["code_source_path"] = nv(codeSourcePath, line) line++ } if cfg.MLflowRunName != nil { @@ -360,15 +382,13 @@ func buildPermissionsValue(perms []permission) dyn.Value { return dyn.V(out) } -// writeBundle writes the bundle into dir: databricks.yml, the loose launch -// artifacts (command.sh + env/secret/param sidecars), and — when the config has a -// code_source — a code_source/ directory holding the staged source tree. All the -// referenced files are bundle-local; `bundle deploy` uploads the launch artifacts -// and the aicode mutator packages+uploads the code_source directory. It refuses to +// writeBundle writes the bundle into dir: databricks.yml plus the loose launch +// artifacts (command.sh + env/secret/param sidecars). It does not touch the code +// source — the deploy-time aicode mutator packages it in place. It refuses to // overwrite existing files so a re-run can't silently clobber a bundle the user has // edited. Returns the relative paths written, for the next-steps message. func writeBundle(ctx context.Context, cfg *runConfig, configPath, dir string) ([]string, error) { - root, artifacts, err := convertToDabs(ctx, cfg, configPath) + root, artifacts, err := convertToDabs(ctx, cfg, configPath, dir) if err != nil { return nil, err } @@ -382,23 +402,17 @@ func writeBundle(ctx context.Context, cfg *runConfig, configPath, dir string) ([ // Refuse to clobber an existing file, with one consistent message. A user // re-running convert into the same dir gets a clear error rather than a silent // overwrite of edits they may have made. - checkCollision := func(name string) error { + writeFile := func(name string, data []byte) error { if _, err := os.Stat(filepath.Join(dir, name)); err == nil { return fmt.Errorf("%s already exists in %s; use --output-dir or remove it", name, dir) } - return nil - } - writeFile := func(name string, data []byte) error { - if err := checkCollision(name); err != nil { - return err - } return os.WriteFile(filepath.Join(dir, name), data, 0o600) } - if err := checkCollision("databricks.yml"); err != nil { - return nil, err - } bundlePath := filepath.Join(dir, "databricks.yml") + if _, err := os.Stat(bundlePath); err == nil { + return nil, fmt.Errorf("databricks.yml already exists in %s; use --output-dir or remove it", dir) + } // force=true: we've already run the collision check above, so SaveAsYAML's own // guard (which references a --force flag this command doesn't have) can't fire. if err := yamlsaver.NewSaver().SaveAsYAML(root, bundlePath, true); err != nil { @@ -414,178 +428,9 @@ func writeBundle(ctx context.Context, cfg *runConfig, configPath, dir string) ([ written = append(written, item.name) } - // Code source: stage the resolved root_path into the bundle's code_source/ dir. - // The aicode mutator packages+uploads it at deploy; convert only lays down the - // bytes. A pinned git commit/branch is materialized from the archived commit (not - // the dirty working tree), matching what `air run` would submit; otherwise the - // working tree is copied, honoring .gitignore just like the run path's plain-tar. - if cfg.CodeSource != nil && cfg.CodeSource.Snapshot != nil { - snap := cfg.CodeSource.Snapshot - repoPath, err := resolveRootPath(ctx, snap.RootPath, filepath.Dir(configPath)) - if err != nil { - return nil, err - } - plan, err := resolveSnapshotPlan(ctx, newGitRepo(repoPath), snap.Git, snap.IncludePaths) - if err != nil { - return nil, err - } - if err := checkCollision(codeSourceDirName); err != nil { - return nil, err - } - if err := materializeCodeSource(ctx, repoPath, plan, filepath.Join(dir, codeSourceDirName)); err != nil { - return nil, err - } - written = append(written, codeSourceDirName+"/") - } - return written, nil } -// materializeCodeSource stages the snapshot described by plan into destDir (the -// bundle's code_source/ directory). It packages the source with the existing -// snapshot packagers — git archive for a pinned commit, gitignore-aware plain tar -// for a working tree — into a temporary tarball, then extracts it and moves its -// single top-level directory into destDir. Going through the packagers (rather than -// a raw copy) reuses their exact commit-pin and .gitignore/include-path handling, so -// the staged tree matches what `air run` would submit; the deploy-time aicode mutator -// then re-packages destDir into the uploaded, content-addressed archive. -func materializeCodeSource(ctx context.Context, repoPath string, plan snapshotPlan, destDir string) error { - staging, err := os.MkdirTemp("", "air-convert-code-*") - if err != nil { - return err - } - defer os.RemoveAll(staging) - - // Absolute tarball path: the git-archive packager runs git with `-C repoPath`, so - // a relative -o would resolve against the repo dir. Both packagers accept absolute. - tarball, err := filepath.Abs(filepath.Join(staging, "code_source.tar.gz")) - if err != nil { - return err - } - // git archive for a pinned commit (deterministic, ignores the dirty tree), - // gitignore-aware plain tar for a working tree — the same split the run path uses. - dirName := filepath.Base(repoPath) - if plan.mode == modeGitArchive { - err = createGitArchiveSnapshot(ctx, newGitRepo(repoPath), plan.commitSHA, tarball, dirName, plan.includePaths) - } else { - err = createPlainTarball(ctx, repoPath, tarball, plan.includePaths) - } - if err != nil { - return err - } - - // The archive's single top-level entry is the source directory's basename (both - // packagers preserve it). Extract into a scratch dir, then move that top-level - // directory to the fixed destDir name so the emitted code_source_path is stable. - extractDir := filepath.Join(staging, "extract") - if err := os.MkdirAll(extractDir, 0o700); err != nil { - return err - } - if err := extractTarball(tarball, extractDir); err != nil { - return err - } - entries, err := os.ReadDir(extractDir) - if err != nil { - return err - } - if len(entries) != 1 || !entries[0].IsDir() { - return fmt.Errorf("unexpected code_source archive layout: expected a single top-level directory, got %d entries", len(entries)) - } - if err := os.Rename(filepath.Join(extractDir, entries[0].Name()), destDir); err != nil { - return err - } - return nil -} - -// extractTarball unpacks a gzipped tarball into destDir in pure Go (no `tar` -// subprocess), so extraction is portable — notably on Windows, where handing a -// `C:\...` path to the system tar makes it read the drive letter as a remote host. -// Entries that would escape destDir (path traversal, absolute or escaping symlinks) -// are rejected; unusual entry types are skipped. -func extractTarball(tarball, destDir string) error { - f, err := os.Open(tarball) - if err != nil { - return err - } - defer f.Close() - - gz, err := gzip.NewReader(f) - if err != nil { - return fmt.Errorf("failed to read code_source archive: %w", err) - } - defer gz.Close() - - // destDir is the boundary every extracted entry must stay within: a malicious or - // malformed archive entry ("../x", an absolute path, a symlink escaping the tree) - // must not write outside it. destAbs is compared against each resolved target. - destAbs, err := filepath.Abs(destDir) - if err != nil { - return err - } - - tr := tar.NewReader(gz) - for { - hdr, err := tr.Next() - if errors.Is(err, io.EOF) { - break - } - if err != nil { - return fmt.Errorf("failed to read code_source archive: %w", err) - } - - // Reject path traversal: the cleaned target must stay under destAbs. - target := filepath.Join(destAbs, filepath.FromSlash(hdr.Name)) - if target != destAbs && !strings.HasPrefix(target, destAbs+string(os.PathSeparator)) { - return fmt.Errorf("code_source archive entry %q escapes the destination directory", hdr.Name) - } - - switch hdr.Typeflag { - case tar.TypeDir: - if err := os.MkdirAll(target, 0o700); err != nil { - return err - } - case tar.TypeReg: - if err := os.MkdirAll(filepath.Dir(target), 0o700); err != nil { - return err - } - out, err := os.OpenFile(target, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, os.FileMode(hdr.Mode)&0o777) - if err != nil { - return err - } - // Bounded copy: cap the number of bytes written to the size the header - // declares, so a corrupt/oversized entry can't write unbounded data. - if _, err := io.CopyN(out, tr, hdr.Size); err != nil && !errors.Is(err, io.EOF) { - out.Close() - return err - } - if err := out.Close(); err != nil { - return err - } - case tar.TypeSymlink: - if err := os.MkdirAll(filepath.Dir(target), 0o700); err != nil { - return err - } - // A symlink whose resolved target escapes destAbs is rejected: it could be - // followed on a later write to reach outside the staged code directory. - linkTarget := hdr.Linkname - resolved := linkTarget - if !filepath.IsAbs(resolved) { - resolved = filepath.Join(filepath.Dir(target), filepath.FromSlash(linkTarget)) - } - if resolved != destAbs && !strings.HasPrefix(resolved, destAbs+string(os.PathSeparator)) { - return fmt.Errorf("code_source archive symlink %q -> %q escapes the destination directory", hdr.Name, linkTarget) - } - if err := os.Symlink(linkTarget, target); err != nil { - return err - } - default: - // Skip other entry types (devices, fifos, etc.): code source is regular - // files, dirs, and symlinks; anything else is not meaningful to a workload. - } - } - return nil -} - // printConvertNextSteps tells the user what was written and the exact deploy // sequence, since the value of the command is a one-command deploy afterwards. // @@ -626,22 +471,14 @@ func printConvertNextSteps(ctx context.Context, dir string, written []string, jo cmdio.LogString(ctx, " databricks bundle destroy") } -// conversionNotes lists what the conversion transformed, staged out-of-band, or -// could not represent natively — so a user migrating from `air run` can see what -// changed between their run YAML and the emitted bundle, and what they may still -// need to fill in. Best-effort: git resolution errors are ignored here (writeBundle -// surfaces them), so a note is only emitted when the state is unambiguous. +// conversionNotes lists what the conversion staged out-of-band or could not +// represent natively, so a user migrating from `air run` sees what changed between +// their run YAML and the emitted bundle. func conversionNotes(cfg *runConfig) []string { var notes []string - if cfg.CodeSource != nil && cfg.CodeSource.Snapshot != nil { - snap := cfg.CodeSource.Snapshot - if snap.Git != nil { - notes = append(notes, "code_source.git was pinned in the run YAML; the pinned commit was materialized into "+ - "the code_source/ directory. Re-run convert-to-dabs to re-materialize a different revision.") - } - notes = append(notes, "code_source was staged into the code_source/ directory; bundle deploy packages and uploads it. "+ - "Re-run convert-to-dabs to re-stage after code changes.") + if codeSnapshot(cfg) != nil { + notes = append(notes, "code_source points at your source directory; bundle deploy packages and uploads it from there.") } // env vars / secrets have no native ai_runtime_task field yet, so they ride as diff --git a/experimental/air/cmd/convert_to_dabs_test.go b/experimental/air/cmd/convert_to_dabs_test.go index 9e8e80d235c..8c9d22217a0 100644 --- a/experimental/air/cmd/convert_to_dabs_test.go +++ b/experimental/air/cmd/convert_to_dabs_test.go @@ -1,8 +1,6 @@ package aircmd import ( - "archive/tar" - "compress/gzip" "os" "path/filepath" "strings" @@ -55,10 +53,11 @@ environment: - torch ` path := writeConfigFile(t, "run.yaml", cfg) + require.NoError(t, os.MkdirAll(filepath.Join(filepath.Dir(path), "src"), 0o700)) loaded, err := loadRunConfig(path) require.NoError(t, err) - root, artifacts, err := convertToDabs(t.Context(), loaded, path) + root, artifacts, err := convertToDabs(t.Context(), loaded, path, filepath.Dir(path)) require.NoError(t, err) name := loaded.ExperimentName @@ -78,9 +77,9 @@ environment: art := task + ".ai_runtime_task" assert.Equal(t, name, get(t, root, art+".experiment").MustString()) assert.Equal(t, "run-42", get(t, root, art+".mlflow_run").MustString()) - // code_source_path is a local directory (packaged by the deploy-time aicode - // mutator), not a pre-built tarball. - assert.Equal(t, "./"+codeSourceDirName, get(t, root, art+".code_source_path").MustString()) + // code_source_path is the source dir relative to the bundle; the deploy-time + // aicode mutator packages it in place. + assert.Equal(t, "./src", get(t, root, art+".code_source_path").MustString()) dep := art + ".deployments[0]" assert.Equal(t, "./"+commandScriptName, get(t, root, dep+".command_path").MustString()) @@ -108,7 +107,7 @@ func TestConvertToDabsOmitsUnsetFields(t *testing.T) { loaded, err := loadRunConfig(path) require.NoError(t, err) - root, _, err := convertToDabs(t.Context(), loaded, path) + root, _, err := convertToDabs(t.Context(), loaded, path, filepath.Dir(path)) require.NoError(t, err) name := loaded.ExperimentName @@ -135,7 +134,7 @@ func TestConvertToDabsRuntimeVersionEnvOverride(t *testing.T) { loaded, err := loadRunConfig(path) require.NoError(t, err) - root, _, err := convertToDabs(t.Context(), loaded, path) + root, _, err := convertToDabs(t.Context(), loaded, path, filepath.Dir(path)) require.NoError(t, err) env := "resources.jobs." + loaded.ExperimentName + ".environments[0]" @@ -155,7 +154,7 @@ code_source: path := writeConfigFile(t, "run.yaml", cfg) loaded, err := loadRunConfig(path) require.NoError(t, err) - _, _, err = convertToDabs(t.Context(), loaded, path) + _, _, err = convertToDabs(t.Context(), loaded, path, filepath.Dir(path)) require.ErrorContains(t, err, "remote_volume is not supported") } @@ -174,7 +173,7 @@ parameters: loaded, err := loadRunConfig(path) require.NoError(t, err) - root, artifacts, err := convertToDabs(t.Context(), loaded, path) + root, artifacts, err := convertToDabs(t.Context(), loaded, path, filepath.Dir(path)) require.NoError(t, err) names := itemNames(artifacts) @@ -189,28 +188,66 @@ parameters: assert.False(t, has(root, art+".secrets")) } -// A working-tree (non-git-pinned) code_source is staged into the bundle's -// code_source/ directory, honoring .gitignore just like the run path's plain-tar. -func TestConvertToDabsWorkingTreeMaterialized(t *testing.T) { - repo := t.TempDir() - writeRepoFile(t, repo, "train.py", "print('x')") - writeRepoFile(t, repo, "notes.log", "scratch") - writeRepoFile(t, repo, ".gitignore", "*.log\n") +// code_source_path is emitted as the source dir relative to the bundle and no code +// is copied — the deploy-time mutator packages it in place. writeBundle produces +// only databricks.yml + launch artifacts, not a code_source copy. +func TestConvertToDabsDoesNotCopyCode(t *testing.T) { + dir := t.TempDir() + require.NoError(t, os.MkdirAll(filepath.Join(dir, "src"), 0o700)) + require.NoError(t, os.WriteFile(filepath.Join(dir, "src", "train.py"), []byte("print()\n"), 0o600)) + + cfg := "experiment_name: wt\ncommand: cd \"$CODE_SOURCE_PATH\" && python train.py\n" + + "compute: {accelerator_type: GPU_1xH100, num_accelerators: 1}\n" + + "code_source:\n type: snapshot\n snapshot:\n root_path: ./src\n" + path := filepath.Join(dir, "run.yaml") + require.NoError(t, os.WriteFile(path, []byte(cfg), 0o600)) + loaded, err := loadRunConfig(path) + require.NoError(t, err) + + written, err := writeBundle(t.Context(), loaded, path, dir) + require.NoError(t, err) - cfg := "experiment_name: wt\ncommand: python train.py\n" + + // code_source_path points at the existing source dir; nothing is copied. + root, _, err := convertToDabs(t.Context(), loaded, path, dir) + require.NoError(t, err) + art := "resources.jobs." + loaded.ExperimentName + ".tasks[0].ai_runtime_task" + assert.Equal(t, "./src", get(t, root, art+".code_source_path").MustString()) + assert.NotContains(t, written, "code_source/") +} + +// A code_source root_path outside the bundle directory is rejected: the mutator only +// packages a directory inside the bundle sync root. +func TestConvertToDabsRejectsCodeOutsideBundle(t *testing.T) { + outside := t.TempDir() + require.NoError(t, os.MkdirAll(filepath.Join(outside, "src"), 0o700)) + + cfg := "experiment_name: outside\ncommand: python train.py\n" + "compute: {accelerator_type: GPU_1xH100, num_accelerators: 1}\n" + - "code_source:\n type: snapshot\n snapshot:\n root_path: " + repo + "\n" + "code_source:\n type: snapshot\n snapshot:\n root_path: " + filepath.Join(outside, "src") + "\n" path := writeConfigFile(t, "run.yaml", cfg) loaded, err := loadRunConfig(path) require.NoError(t, err) - outDir := t.TempDir() - _, err = writeBundle(t.Context(), loaded, path, outDir) + // Bundle dir is the config's temp dir; the source is in a different tree. + _, _, err = convertToDabs(t.Context(), loaded, path, filepath.Dir(path)) + require.ErrorContains(t, err, "not inside the bundle") +} + +// A git-pinned code_source is rejected: convert no longer materializes a commit; +// the deploy-time mutator packages the working tree in place. +func TestConvertToDabsRejectsGitPin(t *testing.T) { + dir := t.TempDir() + require.NoError(t, os.MkdirAll(filepath.Join(dir, "src"), 0o700)) + cfg := "experiment_name: git-pin\ncommand: python train.py\n" + + "compute: {accelerator_type: GPU_1xH100, num_accelerators: 1}\n" + + "code_source:\n type: snapshot\n snapshot:\n root_path: ./src\n git:\n commit: abc123\n" + path := filepath.Join(dir, "run.yaml") + require.NoError(t, os.WriteFile(path, []byte(cfg), 0o600)) + loaded, err := loadRunConfig(path) require.NoError(t, err) - codeDir := filepath.Join(outDir, codeSourceDirName) - assert.FileExists(t, filepath.Join(codeDir, "train.py")) - assert.NoFileExists(t, filepath.Join(codeDir, "notes.log"), "gitignored files must be excluded from the staged code_source") + _, _, err = convertToDabs(t.Context(), loaded, path, dir) + require.ErrorContains(t, err, "git is not supported") } // A requirements-FILE dependency set (environment.dependencies is a path) is folded @@ -228,7 +265,7 @@ func TestConvertToDabsFoldsRequirementsFileIntoEnvSpec(t *testing.T) { loaded, err := loadRunConfig(path) require.NoError(t, err) - root, artifacts, err := convertToDabs(t.Context(), loaded, path) + root, artifacts, err := convertToDabs(t.Context(), loaded, path, filepath.Dir(path)) require.NoError(t, err) env := "resources.jobs." + loaded.ExperimentName + ".environments[0]" @@ -253,7 +290,6 @@ code_source: type: snapshot snapshot: root_path: ./src - git: {commit: abc123} ` path := writeConfigFile(t, "run.yaml", cfg) loaded, err := loadRunConfig(path) @@ -261,8 +297,7 @@ code_source: notes := conversionNotes(loaded) joined := strings.Join(notes, "\n") - assert.Contains(t, joined, "code_source.git") // git pin flagged - assert.Contains(t, joined, "code_source/") // staged-directory behavior + assert.Contains(t, joined, "code_source") // source-dir behavior assert.Contains(t, joined, "env_vars.json") // env vars staged assert.Contains(t, joined, "secret_env_vars.json") // secrets staged assert.Contains(t, joined, "hyperparameters.yaml") // parameters staged @@ -282,7 +317,7 @@ func TestConvertToDabsMapsUsagePolicyID(t *testing.T) { loaded, err := loadRunConfig(path) require.NoError(t, err) - root, _, err := convertToDabs(t.Context(), loaded, path) + root, _, err := convertToDabs(t.Context(), loaded, path, filepath.Dir(path)) require.NoError(t, err) assert.Equal(t, "budget-abc-123", get(t, root, "resources.jobs."+loaded.ExperimentName+".budget_policy_id").MustString()) } @@ -299,7 +334,7 @@ permissions: loaded, err := loadRunConfig(path) require.NoError(t, err) - root, _, err := convertToDabs(t.Context(), loaded, path) + root, _, err := convertToDabs(t.Context(), loaded, path, filepath.Dir(path)) require.NoError(t, err) perms := get(t, root, "resources.jobs."+loaded.ExperimentName+".permissions").MustSequence() @@ -309,119 +344,6 @@ permissions: assert.Equal(t, "eng", perms[1].Get("group_name").MustString()) } -// A git-pinned code_source is materialized from the commit, not the dirty working -// tree: writeBundle stages a code_source/ directory holding the committed file only. -func TestConvertToDabsGitPinnedMaterialized(t *testing.T) { - repo := newTestRepo(t) - writeRepoFile(t, repo, "train.py", "print('committed')") - sha := commitAll(t, repo, "init") - // Dirty the tree AFTER the commit; the pinned snapshot must not include this. - writeRepoFile(t, repo, "uncommitted.py", "print('dirty')") - - cfg := "experiment_name: git-pin\ncommand: python train.py\n" + - "compute: {accelerator_type: GPU_1xH100, num_accelerators: 1}\n" + - "code_source:\n type: snapshot\n snapshot:\n root_path: " + repo + "\n git:\n commit: " + sha + "\n" - path := writeConfigFile(t, "run.yaml", cfg) - loaded, err := loadRunConfig(path) - require.NoError(t, err) - - outDir := t.TempDir() - _, err = writeBundle(t.Context(), loaded, path, outDir) - require.NoError(t, err) - - codeDir := filepath.Join(outDir, codeSourceDirName) - assert.FileExists(t, filepath.Join(codeDir, "train.py")) - assert.NoFileExists(t, filepath.Join(codeDir, "uncommitted.py"), "git-pinned snapshot must exclude uncommitted files") -} - -// git archive runs with `git -C repoPath`, so the staging tarball path must be -// absolute — otherwise `-o out/...` resolves against the repo dir and fails. Exercise -// writeBundle from a working dir with a RELATIVE output path against a git-pinned source. -func TestConvertToDabsGitPinnedRelativeOutputDir(t *testing.T) { - repo := newTestRepo(t) - writeRepoFile(t, repo, "train.py", "print('x')") - sha := commitAll(t, repo, "init") - - cfg := "experiment_name: rel-out\ncommand: python train.py\n" + - "compute: {accelerator_type: GPU_1xH100, num_accelerators: 1}\n" + - "code_source:\n type: snapshot\n snapshot:\n root_path: " + repo + "\n git:\n commit: " + sha + "\n" - path := writeConfigFile(t, "run.yaml", cfg) - loaded, err := loadRunConfig(path) - require.NoError(t, err) - - // chdir into a scratch dir and pass a relative --output-dir. - work := t.TempDir() - t.Chdir(work) - _, err = writeBundle(t.Context(), loaded, path, "out") - require.NoError(t, err) - - assert.FileExists(t, filepath.Join(work, "out", codeSourceDirName, "train.py")) -} - -// writeTarball writes a gzipped tar of the given entries to path. Each entry is -// either a regular file (typeflag defaults to file) or, when linkname != "", a -// symlink. Used to exercise extractTarball directly. -func writeTarball(t *testing.T, path string, entries []tar.Header) { - t.Helper() - f, err := os.Create(path) - require.NoError(t, err) - gz := gzip.NewWriter(f) - tw := tar.NewWriter(gz) - for _, h := range entries { - require.NoError(t, tw.WriteHeader(&h)) - if h.Typeflag == tar.TypeReg { - _, err := tw.Write([]byte("data-" + h.Name)) - require.NoError(t, err) - } - } - require.NoError(t, tw.Close()) - require.NoError(t, gz.Close()) - require.NoError(t, f.Close()) -} - -// extractTarball unpacks a well-formed archive (files, nested dirs, in-tree -// symlink) into destDir. -func TestExtractTarballHappyPath(t *testing.T) { - src := filepath.Join(t.TempDir(), "a.tar.gz") - writeTarball(t, src, []tar.Header{ - {Name: "code/", Typeflag: tar.TypeDir, Mode: 0o755}, - {Name: "code/train.py", Typeflag: tar.TypeReg, Mode: 0o644, Size: int64(len("data-code/train.py"))}, - {Name: "code/pkg/", Typeflag: tar.TypeDir, Mode: 0o755}, - {Name: "code/pkg/util.py", Typeflag: tar.TypeReg, Mode: 0o644, Size: int64(len("data-code/pkg/util.py"))}, - {Name: "code/link.py", Typeflag: tar.TypeSymlink, Linkname: "train.py"}, - }) - - dest := t.TempDir() - require.NoError(t, extractTarball(src, dest)) - assert.FileExists(t, filepath.Join(dest, "code", "train.py")) - assert.FileExists(t, filepath.Join(dest, "code", "pkg", "util.py")) - if info, err := os.Lstat(filepath.Join(dest, "code", "link.py")); assert.NoError(t, err) { - assert.NotZero(t, info.Mode()&os.ModeSymlink, "link.py should be a symlink") - } -} - -// A path-traversal entry or an escaping symlink is rejected rather than written -// outside destDir. -func TestExtractTarballRejectsEscape(t *testing.T) { - t.Run("traversal path", func(t *testing.T) { - src := filepath.Join(t.TempDir(), "evil.tar.gz") - writeTarball(t, src, []tar.Header{ - {Name: "../escape.py", Typeflag: tar.TypeReg, Mode: 0o644, Size: int64(len("data-../escape.py"))}, - }) - err := extractTarball(src, t.TempDir()) - require.ErrorContains(t, err, "escapes the destination directory") - }) - - t.Run("escaping symlink", func(t *testing.T) { - src := filepath.Join(t.TempDir(), "evil-link.tar.gz") - writeTarball(t, src, []tar.Header{ - {Name: "link", Typeflag: tar.TypeSymlink, Linkname: "../../etc/passwd"}, - }) - err := extractTarball(src, t.TempDir()) - require.ErrorContains(t, err, "escapes the destination directory") - }) -} - // A numeric or reserved-word experiment_name would be an unquoted YAML map key // that DABs' strict loader rejects (!!int / !!bool). The job resource key is // prefixed to stay a string, while name/experiment keep the original value. @@ -445,7 +367,7 @@ func TestConvertToDabsSafeJobKey(t *testing.T) { path := writeConfigFile(t, "run.yaml", cfg) loaded, err := loadRunConfig(path) require.NoError(t, err) - root, _, err := convertToDabs(t.Context(), loaded, path) + root, _, err := convertToDabs(t.Context(), loaded, path, filepath.Dir(path)) require.NoError(t, err) assert.Equal(t, "12345", get(t, root, "resources.jobs.job_12345.name").MustString()) assert.Equal(t, "12345", get(t, root, "resources.jobs.job_12345.tasks[0].ai_runtime_task.experiment").MustString()) @@ -461,7 +383,7 @@ environment: path := writeConfigFile(t, "run.yaml", cfg) loaded, err := loadRunConfig(path) require.NoError(t, err) - _, _, err = convertToDabs(t.Context(), loaded, path) + _, _, err = convertToDabs(t.Context(), loaded, path, filepath.Dir(path)) require.ErrorContains(t, err, "docker_image is not yet supported") }) @@ -470,7 +392,7 @@ environment: path := writeConfigFile(t, "run.yaml", cfg) loaded, err := loadRunConfig(path) require.NoError(t, err) - _, _, err = convertToDabs(t.Context(), loaded, path) + _, _, err = convertToDabs(t.Context(), loaded, path, filepath.Dir(path)) require.ErrorContains(t, err, "usage_policy_name is not yet supported") }) } From b424aff6b7a5dcea81f5e2ed0f4416ed23faa0e3 Mon Sep 17 00:00:00 2001 From: vinchenzo-db Date: Mon, 3 Aug 2026 21:04:07 +0000 Subject: [PATCH 04/15] air/convert-to-dabs: clearer git / remote_volume rejection messages git and remote_volume stay unsupported, but the errors now explain why and what to do: for git, check out the revision before converting (deploy packages the working tree as-is); for remote_volume, set workspace.artifact_path (the bundle artifact location is bundle-wide, not per-code-source). Co-authored-by: Isaac --- experimental/air/cmd/convert_to_dabs.go | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/experimental/air/cmd/convert_to_dabs.go b/experimental/air/cmd/convert_to_dabs.go index 143e896713d..3d4d240de28 100644 --- a/experimental/air/cmd/convert_to_dabs.go +++ b/experimental/air/cmd/convert_to_dabs.go @@ -110,15 +110,17 @@ func convertToDabs(ctx context.Context, cfg *runConfig, configPath, bundleDir st return nil, nil, errors.New("environment.docker_image is not yet supported by convert-to-dabs") } if snap := codeSnapshot(cfg); snap != nil { - // remote_volume points the code archive at a UC Volume; bundle deploy uploads - // code_source to the bundle artifact path, not an arbitrary Volume. + // remote_volume points the code archive at a specific UC Volume. The bundle's + // artifact location is set bundle-wide via workspace.artifact_path, not + // per-code-source, so a per-source Volume isn't representable here. if snap.RemoteVolume != nil { - return nil, nil, errors.New("code_source.snapshot.remote_volume is not supported by convert-to-dabs") + return nil, nil, errors.New("code_source.snapshot.remote_volume is not supported by convert-to-dabs; set workspace.artifact_path in the bundle instead") } - // git pinning would require materializing the pinned commit, which convert no - // longer does — the deploy-time mutator packages the working tree in place. + // git pins to a committed revision, but convert packages nothing — the + // deploy-time mutator uploads the working tree as it is on disk. Deploying a + // specific revision therefore isn't supported; check it out before converting. if snap.Git != nil { - return nil, nil, errors.New("code_source.snapshot.git is not supported by convert-to-dabs; deploy packages the working tree") + return nil, nil, errors.New("code_source.snapshot.git is not supported by convert-to-dabs; deploy packages your working tree as-is, so check out the revision you want (git checkout ) before converting") } } From d37fd2dccd94d6ad57c7b2effb79ac1e1262beab Mon Sep 17 00:00:00 2001 From: vinchenzo-db Date: Mon, 3 Aug 2026 22:56:27 +0000 Subject: [PATCH 05/15] air/convert-to-dabs: add --force, and name this CLI in the next steps Re-running a conversion previously failed with no way to overwrite, and the suggested remedy (--output-dir) is a dead end in place: code_source must live inside the bundle dir. Add --force, and point the error at it. The next-steps message said "databricks bundle ...". A CLI without ai_runtime_task support only warns on the unknown field and then deploys a job with no AI task, so print the invoked binary instead and call out the risk. Co-authored-by: Isaac --- .../air/convert-to-dabs/output.txt | 43 ++++++++++-- .../experimental/air/convert-to-dabs/script | 6 ++ experimental/air/cmd/convert_to_dabs.go | 67 ++++++++++++++----- experimental/air/cmd/convert_to_dabs_test.go | 37 +++++++++- 4 files changed, 131 insertions(+), 22 deletions(-) diff --git a/acceptance/experimental/air/convert-to-dabs/output.txt b/acceptance/experimental/air/convert-to-dabs/output.txt index 2b3106c4c0d..45c1b476bd9 100644 --- a/acceptance/experimental/air/convert-to-dabs/output.txt +++ b/acceptance/experimental/air/convert-to-dabs/output.txt @@ -11,16 +11,19 @@ Notes: To deploy and run this workload as a bundle: 1. cd . - 2. databricks bundle validate - 3. databricks bundle deploy - 4. databricks bundle run torchrun-a10-smoke-test + 2. [CLI] bundle validate + 3. [CLI] bundle deploy + 4. [CLI] bundle run torchrun-a10-smoke-test bundle deploy uploads the code source and launch scripts automatically. +Run these with this same CLI: a build without ai_runtime_task support +only warns about the unknown field, then deploys a job with no AI task. + Unlike `air run` (which submits an ephemeral run), bundle deploy creates a persistent job that is not garbage-collected. When you are done, remove the job and its uploaded files with: - databricks bundle destroy + [CLI] bundle destroy === emitted databricks.yml (code_source_path points at ./src; no code is copied) >>> cat databricks.yml @@ -65,6 +68,38 @@ Workspace: Validation OK! +=== re-converting refuses to clobber the generated bundle +>>> [CLI] experimental air convert-to-dabs train.yaml +Error: databricks.yml already exists in .; pass --force to overwrite or remove it + +Exit code: 1 + +=== --force overwrites it +>>> [CLI] experimental air convert-to-dabs train.yaml --force +Wrote a Databricks Asset Bundle to .: + databricks.yml + training_config.yaml + command.sh + +Notes: + - code_source points at your source directory; bundle deploy packages and uploads it from there. + +To deploy and run this workload as a bundle: + 1. cd . + 2. [CLI] bundle validate + 3. [CLI] bundle deploy + 4. [CLI] bundle run torchrun-a10-smoke-test + +bundle deploy uploads the code source and launch scripts automatically. + +Run these with this same CLI: a build without ai_runtime_task support +only warns about the unknown field, then deploys a job with no AI task. + +Unlike `air run` (which submits an ephemeral run), bundle deploy creates a +persistent job that is not garbage-collected. When you are done, remove the +job and its uploaded files with: + [CLI] bundle destroy + === docker_image is not supported yet >>> [CLI] experimental air convert-to-dabs docker.yaml --output-dir generated-docker Error: environment.docker_image is not yet supported by convert-to-dabs diff --git a/acceptance/experimental/air/convert-to-dabs/script b/acceptance/experimental/air/convert-to-dabs/script index 2babd7c0aef..8882963ac91 100644 --- a/acceptance/experimental/air/convert-to-dabs/script +++ b/acceptance/experimental/air/convert-to-dabs/script @@ -10,5 +10,11 @@ trace cat command.sh title "the emitted bundle validates" trace $CLI bundle validate +title "re-converting refuses to clobber the generated bundle" +errcode trace $CLI experimental air convert-to-dabs train.yaml + +title "--force overwrites it" +trace $CLI experimental air convert-to-dabs train.yaml --force + title "docker_image is not supported yet" errcode trace $CLI experimental air convert-to-dabs docker.yaml --output-dir generated-docker diff --git a/experimental/air/cmd/convert_to_dabs.go b/experimental/air/cmd/convert_to_dabs.go index 3d4d240de28..1ae604e4d19 100644 --- a/experimental/air/cmd/convert_to_dabs.go +++ b/experimental/air/cmd/convert_to_dabs.go @@ -40,7 +40,10 @@ import ( const dabsTargetName = "dev" func newConvertToDabsCommand() *cobra.Command { - var outputDir string + var ( + outputDir string + force bool + ) cmd := &cobra.Command{ Use: "convert-to-dabs ", @@ -59,6 +62,7 @@ does not contact the workspace.`, } cmd.Flags().StringVar(&outputDir, "output-dir", "", "Directory to write the bundle into (default: a -bundle folder next to the input YAML). Accepts an absolute or relative path.") + cmd.Flags().BoolVar(&force, "force", false, "Overwrite the generated bundle files if they already exist.") cmd.RunE = func(cmd *cobra.Command, args []string) error { ctx := cmd.Context() @@ -79,7 +83,7 @@ does not contact the workspace.`, dir = filepath.Dir(yamlPath) } - written, err := writeBundle(ctx, cfg, yamlPath, dir) + written, err := writeBundle(ctx, cfg, yamlPath, dir, force) if err != nil { return err } @@ -386,10 +390,11 @@ func buildPermissionsValue(perms []permission) dyn.Value { // writeBundle writes the bundle into dir: databricks.yml plus the loose launch // artifacts (command.sh + env/secret/param sidecars). It does not touch the code -// source — the deploy-time aicode mutator packages it in place. It refuses to -// overwrite existing files so a re-run can't silently clobber a bundle the user has -// edited. Returns the relative paths written, for the next-steps message. -func writeBundle(ctx context.Context, cfg *runConfig, configPath, dir string) ([]string, error) { +// source — the deploy-time aicode mutator packages it in place. Unless force is set +// it refuses to overwrite existing files, so a re-run can't silently clobber a +// bundle the user has edited. Returns the relative paths written, for the +// next-steps message. +func writeBundle(ctx context.Context, cfg *runConfig, configPath, dir string, force bool) ([]string, error) { root, artifacts, err := convertToDabs(ctx, cfg, configPath, dir) if err != nil { return nil, err @@ -403,20 +408,26 @@ func writeBundle(ctx context.Context, cfg *runConfig, configPath, dir string) ([ // Refuse to clobber an existing file, with one consistent message. A user // re-running convert into the same dir gets a clear error rather than a silent - // overwrite of edits they may have made. + // overwrite of edits they may have made. The hint names --force rather than + // --output-dir: the code_source must live inside the bundle dir, so redirecting + // the output usually isn't a usable escape hatch for an in-place conversion. writeFile := func(name string, data []byte) error { - if _, err := os.Stat(filepath.Join(dir, name)); err == nil { - return fmt.Errorf("%s already exists in %s; use --output-dir or remove it", name, dir) + if !force { + if _, err := os.Stat(filepath.Join(dir, name)); err == nil { + return fmt.Errorf("%s already exists in %s; pass --force to overwrite or remove it", name, dir) + } } return os.WriteFile(filepath.Join(dir, name), data, 0o600) } bundlePath := filepath.Join(dir, "databricks.yml") - if _, err := os.Stat(bundlePath); err == nil { - return nil, fmt.Errorf("databricks.yml already exists in %s; use --output-dir or remove it", dir) + if !force { + if _, err := os.Stat(bundlePath); err == nil { + return nil, fmt.Errorf("databricks.yml already exists in %s; pass --force to overwrite or remove it", dir) + } } - // force=true: we've already run the collision check above, so SaveAsYAML's own - // guard (which references a --force flag this command doesn't have) can't fire. + // SaveAsYAML's force arg is passed true unconditionally: the collision check + // above already decided whether overwriting is allowed. if err := yamlsaver.NewSaver().SaveAsYAML(root, bundlePath, true); err != nil { return nil, err } @@ -458,19 +469,41 @@ func printConvertNextSteps(ctx context.Context, dir string, written []string, jo } } + // Name this binary rather than a bare "databricks": ai_runtime_task is only + // understood by a CLI carrying it, and an older one on PATH drops the field with + // just a warning, deploying a job with no AI task at all. + self := cliInvocation() + cmdio.LogString(ctx, "") cmdio.LogString(ctx, "To deploy and run this workload as a bundle:") cmdio.LogString(ctx, " 1. cd "+dir) - cmdio.LogString(ctx, " 2. databricks bundle validate") - cmdio.LogString(ctx, " 3. databricks bundle deploy") - cmdio.LogString(ctx, " 4. databricks bundle run "+jobKey) + cmdio.LogString(ctx, " 2. "+self+" bundle validate") + cmdio.LogString(ctx, " 3. "+self+" bundle deploy") + cmdio.LogString(ctx, " 4. "+self+" bundle run "+jobKey) cmdio.LogString(ctx, "") cmdio.LogString(ctx, "bundle deploy uploads the code source and launch scripts automatically.") cmdio.LogString(ctx, "") + cmdio.LogString(ctx, "Run these with this same CLI: a build without ai_runtime_task support") + cmdio.LogString(ctx, "only warns about the unknown field, then deploys a job with no AI task.") + cmdio.LogString(ctx, "") cmdio.LogString(ctx, "Unlike `air run` (which submits an ephemeral run), bundle deploy creates a") cmdio.LogString(ctx, "persistent job that is not garbage-collected. When you are done, remove the") cmdio.LogString(ctx, "job and its uploaded files with:") - cmdio.LogString(ctx, " databricks bundle destroy") + cmdio.LogString(ctx, " "+self+" bundle destroy") +} + +// cliInvocation is how the user should spell this binary in a follow-up command. +// A path-qualified argv[0] (./dbcli, ../dbcli) is kept as typed so copy-paste works +// from the same cwd; a bare name resolved via PATH is reported as "databricks". +func cliInvocation() string { + arg0 := os.Args[0] + if arg0 == "" { + return "databricks" + } + if arg0 == filepath.Base(arg0) { + return "databricks" + } + return arg0 } // conversionNotes lists what the conversion staged out-of-band or could not diff --git a/experimental/air/cmd/convert_to_dabs_test.go b/experimental/air/cmd/convert_to_dabs_test.go index 8c9d22217a0..a9b8ed337a3 100644 --- a/experimental/air/cmd/convert_to_dabs_test.go +++ b/experimental/air/cmd/convert_to_dabs_test.go @@ -11,6 +11,41 @@ import ( "github.com/stretchr/testify/require" ) +// Re-running a conversion into a dir that already holds a generated bundle is +// refused by default (the user may have edited it) and allowed with --force. +func TestConvertToDabsForceOverwrite(t *testing.T) { + dir := t.TempDir() + require.NoError(t, os.MkdirAll(filepath.Join(dir, "src"), 0o700)) + cfg := "experiment_name: overwrite\ncommand: python train.py\n" + + "compute: {accelerator_type: GPU_1xH100, num_accelerators: 1}\n" + + "code_source:\n type: snapshot\n snapshot:\n root_path: ./src\n" + path := filepath.Join(dir, "run.yaml") + require.NoError(t, os.WriteFile(path, []byte(cfg), 0o600)) + loaded, err := loadRunConfig(path) + require.NoError(t, err) + + _, err = writeBundle(t.Context(), loaded, path, dir, false) + require.NoError(t, err) + + // A hand edit must survive a refused re-run. + edited := []byte("# hand-edited\n") + require.NoError(t, os.WriteFile(filepath.Join(dir, "databricks.yml"), edited, 0o600)) + + _, err = writeBundle(t.Context(), loaded, path, dir, false) + require.ErrorContains(t, err, "pass --force to overwrite") + kept, err := os.ReadFile(filepath.Join(dir, "databricks.yml")) + require.NoError(t, err) + assert.Equal(t, edited, kept) + + // With --force the generated bundle replaces it. + _, err = writeBundle(t.Context(), loaded, path, dir, true) + require.NoError(t, err) + regenerated, err := os.ReadFile(filepath.Join(dir, "databricks.yml")) + require.NoError(t, err) + assert.NotEqual(t, edited, regenerated) + assert.Contains(t, string(regenerated), "ai_runtime_task") +} + func TestConvertToDabsCommandShape(t *testing.T) { cmd := newConvertToDabsCommand() assert.Equal(t, "convert-to-dabs ", cmd.Use) @@ -204,7 +239,7 @@ func TestConvertToDabsDoesNotCopyCode(t *testing.T) { loaded, err := loadRunConfig(path) require.NoError(t, err) - written, err := writeBundle(t.Context(), loaded, path, dir) + written, err := writeBundle(t.Context(), loaded, path, dir, false) require.NoError(t, err) // code_source_path points at the existing source dir; nothing is copied. From c7698c980bb7252f34aa07845cfbdb6e622267c5 Mon Sep 17 00:00:00 2001 From: vinchenzo-db Date: Tue, 4 Aug 2026 01:25:22 +0000 Subject: [PATCH 06/15] air/convert-to-dabs: trim the next-steps output The bundle is written next to the input YAML, so `cd .` was a no-op step; it is now printed only when the bundle lands elsewhere. Drop the notes block and the stale-CLI warning, and point at `bundle summary` for what was deployed. Co-authored-by: Isaac --- .../air/convert-to-dabs/output.txt | 28 +++------ experimental/air/cmd/convert_to_dabs.go | 63 ++++++------------- experimental/air/cmd/convert_to_dabs_test.go | 31 --------- 3 files changed, 26 insertions(+), 96 deletions(-) diff --git a/acceptance/experimental/air/convert-to-dabs/output.txt b/acceptance/experimental/air/convert-to-dabs/output.txt index 45c1b476bd9..1f3ca78959a 100644 --- a/acceptance/experimental/air/convert-to-dabs/output.txt +++ b/acceptance/experimental/air/convert-to-dabs/output.txt @@ -6,19 +6,13 @@ Wrote a Databricks Asset Bundle to .: training_config.yaml command.sh -Notes: - - code_source points at your source directory; bundle deploy packages and uploads it from there. - To deploy and run this workload as a bundle: - 1. cd . - 2. [CLI] bundle validate - 3. [CLI] bundle deploy - 4. [CLI] bundle run torchrun-a10-smoke-test + 1. [CLI] bundle validate + 2. [CLI] bundle deploy + 3. [CLI] bundle run torchrun-a10-smoke-test bundle deploy uploads the code source and launch scripts automatically. - -Run these with this same CLI: a build without ai_runtime_task support -only warns about the unknown field, then deploys a job with no AI task. +To see what it deployed and where: [CLI] bundle summary Unlike `air run` (which submits an ephemeral run), bundle deploy creates a persistent job that is not garbage-collected. When you are done, remove the @@ -81,19 +75,13 @@ Wrote a Databricks Asset Bundle to .: training_config.yaml command.sh -Notes: - - code_source points at your source directory; bundle deploy packages and uploads it from there. - To deploy and run this workload as a bundle: - 1. cd . - 2. [CLI] bundle validate - 3. [CLI] bundle deploy - 4. [CLI] bundle run torchrun-a10-smoke-test + 1. [CLI] bundle validate + 2. [CLI] bundle deploy + 3. [CLI] bundle run torchrun-a10-smoke-test bundle deploy uploads the code source and launch scripts automatically. - -Run these with this same CLI: a build without ai_runtime_task support -only warns about the unknown field, then deploys a job with no AI task. +To see what it deployed and where: [CLI] bundle summary Unlike `air run` (which submits an ephemeral run), bundle deploy creates a persistent job that is not garbage-collected. When you are done, remove the diff --git a/experimental/air/cmd/convert_to_dabs.go b/experimental/air/cmd/convert_to_dabs.go index 1ae604e4d19..a3c86a317a3 100644 --- a/experimental/air/cmd/convert_to_dabs.go +++ b/experimental/air/cmd/convert_to_dabs.go @@ -88,7 +88,7 @@ does not contact the workspace.`, return err } - printConvertNextSteps(ctx, dir, written, bundleResourceKey(cfg.ExperimentName), conversionNotes(cfg)) + printConvertNextSteps(ctx, dir, written, bundleResourceKey(cfg.ExperimentName)) return nil } @@ -452,39 +452,37 @@ func writeBundle(ctx context.Context, cfg *runConfig, configPath, dir string, fo // whereas `bundle deploy` creates a *persistent* job that lingers until explicitly // destroyed — DABs has no automatic GC. A user migrating from `air run` will not // expect a durable resource, so we call out `bundle destroy` explicitly. -func printConvertNextSteps(ctx context.Context, dir string, written []string, jobKey string, notes []string) { +func printConvertNextSteps(ctx context.Context, dir string, written []string, jobKey string) { cmdio.LogString(ctx, fmt.Sprintf("Wrote a Databricks Asset Bundle to %s:", dir)) for _, w := range written { cmdio.LogString(ctx, " "+w) } - // Notes surface anything the user should know: fields we transformed or dropped, - // and values they may need to fill in. Migrating users otherwise can't tell what - // silently changed between their run YAML and the bundle. - if len(notes) > 0 { - cmdio.LogString(ctx, "") - cmdio.LogString(ctx, "Notes:") - for _, n := range notes { - cmdio.LogString(ctx, " - "+n) - } - } - // Name this binary rather than a bare "databricks": ai_runtime_task is only // understood by a CLI carrying it, and an older one on PATH drops the field with // just a warning, deploying a job with no AI task at all. self := cliInvocation() + // The bundle is written next to the input YAML by default, so a cd step is only + // worth printing when the user has to leave the current directory. + var steps []string + if dir != "." { + steps = append(steps, "cd "+dir) + } + steps = append(steps, + self+" bundle validate", + self+" bundle deploy", + self+" bundle run "+jobKey, + ) + cmdio.LogString(ctx, "") cmdio.LogString(ctx, "To deploy and run this workload as a bundle:") - cmdio.LogString(ctx, " 1. cd "+dir) - cmdio.LogString(ctx, " 2. "+self+" bundle validate") - cmdio.LogString(ctx, " 3. "+self+" bundle deploy") - cmdio.LogString(ctx, " 4. "+self+" bundle run "+jobKey) + for i, s := range steps { + cmdio.LogString(ctx, fmt.Sprintf(" %d. %s", i+1, s)) + } cmdio.LogString(ctx, "") cmdio.LogString(ctx, "bundle deploy uploads the code source and launch scripts automatically.") - cmdio.LogString(ctx, "") - cmdio.LogString(ctx, "Run these with this same CLI: a build without ai_runtime_task support") - cmdio.LogString(ctx, "only warns about the unknown field, then deploys a job with no AI task.") + cmdio.LogString(ctx, "To see what it deployed and where: "+self+" bundle summary") cmdio.LogString(ctx, "") cmdio.LogString(ctx, "Unlike `air run` (which submits an ephemeral run), bundle deploy creates a") cmdio.LogString(ctx, "persistent job that is not garbage-collected. When you are done, remove the") @@ -505,28 +503,3 @@ func cliInvocation() string { } return arg0 } - -// conversionNotes lists what the conversion staged out-of-band or could not -// represent natively, so a user migrating from `air run` sees what changed between -// their run YAML and the emitted bundle. -func conversionNotes(cfg *runConfig) []string { - var notes []string - - if codeSnapshot(cfg) != nil { - notes = append(notes, "code_source points at your source directory; bundle deploy packages and uploads it from there.") - } - - // env vars / secrets have no native ai_runtime_task field yet, so they ride as - // sidecar files the server-side launcher reads (same as `air run`). - if len(cfg.EnvVariables) > 0 { - notes = append(notes, "env_variables were written to env_vars.json (no native bundle field yet); they are uploaded with the code and applied at run time.") - } - if len(cfg.Secrets) > 0 { - notes = append(notes, "secrets were written to secret_env_vars.json (no native bundle field yet); they are resolved at run time.") - } - if len(cfg.Parameters) > 0 { - notes = append(notes, "parameters were written to hyperparameters.yaml; they are not a native bundle field and are passed through to the workload.") - } - - return notes -} diff --git a/experimental/air/cmd/convert_to_dabs_test.go b/experimental/air/cmd/convert_to_dabs_test.go index a9b8ed337a3..46c5384fb6e 100644 --- a/experimental/air/cmd/convert_to_dabs_test.go +++ b/experimental/air/cmd/convert_to_dabs_test.go @@ -3,7 +3,6 @@ package aircmd import ( "os" "path/filepath" - "strings" "testing" "github.com/databricks/cli/libs/dyn" @@ -314,36 +313,6 @@ func TestConvertToDabsFoldsRequirementsFileIntoEnvSpec(t *testing.T) { assert.NotContains(t, itemNames(artifacts), requirementsName) } -// conversionNotes surfaces what was transformed/staged so a migrating user knows -// what changed between their run YAML and the bundle. -func TestConvertToDabsConversionNotes(t *testing.T) { - cfg := minimalConfig + ` -env_variables: {FOO: bar} -secrets: {TOKEN: scope/key} -parameters: {lr: 0.1} -code_source: - type: snapshot - snapshot: - root_path: ./src -` - path := writeConfigFile(t, "run.yaml", cfg) - loaded, err := loadRunConfig(path) - require.NoError(t, err) - - notes := conversionNotes(loaded) - joined := strings.Join(notes, "\n") - assert.Contains(t, joined, "code_source") // source-dir behavior - assert.Contains(t, joined, "env_vars.json") // env vars staged - assert.Contains(t, joined, "secret_env_vars.json") // secrets staged - assert.Contains(t, joined, "hyperparameters.yaml") // parameters staged - - // A minimal config with none of those has no notes. - base := writeConfigFile(t, "min.yaml", minimalConfig) - minCfg, err := loadRunConfig(base) - require.NoError(t, err) - assert.Empty(t, conversionNotes(minCfg)) -} - // usage_policy_id is a resolved budget policy id and maps to the job's // budget_policy_id (usage_policy_name, which needs resolution, is rejected). func TestConvertToDabsMapsUsagePolicyID(t *testing.T) { From 2cf4f20a0ae2d015b701847c14073472cb4e4da5 Mon Sep 17 00:00:00 2001 From: vinchenzo-db Date: Wed, 5 Aug 2026 00:22:40 +0000 Subject: [PATCH 07/15] air/convert-to-dabs: reject include_paths, quote non-string job keys MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit include_paths narrows the archive to a subset of root_path, which a bundle can't express per code source — deploy packages the whole directory. Converting was silently dropping it and uploading files the user meant to exclude, so reject it and point at sync.exclude / .gitignore. Emit a job resource key that YAML would type as a non-string scalar ("12345" -> !!int) as a quoted key instead of prefixing it with "job_". yamlsaver already had this logic for scalar *values* (isScalarValueInString); apply it to map keys too, so the resource key keeps the experiment name. Co-authored-by: Isaac --- experimental/air/cmd/convert_to_dabs.go | 26 +++---- experimental/air/cmd/convert_to_dabs_test.go | 77 ++++++++++++++------ libs/dyn/yamlsaver/saver.go | 9 ++- 3 files changed, 74 insertions(+), 38 deletions(-) diff --git a/experimental/air/cmd/convert_to_dabs.go b/experimental/air/cmd/convert_to_dabs.go index a3c86a317a3..7b0f33591d2 100644 --- a/experimental/air/cmd/convert_to_dabs.go +++ b/experimental/air/cmd/convert_to_dabs.go @@ -7,8 +7,6 @@ import ( "os" "path/filepath" "slices" - "strconv" - "strings" "github.com/databricks/cli/cmd/root" "github.com/databricks/cli/libs/cmdio" @@ -126,6 +124,13 @@ func convertToDabs(ctx context.Context, cfg *runConfig, configPath, bundleDir st if snap.Git != nil { return nil, nil, errors.New("code_source.snapshot.git is not supported by convert-to-dabs; deploy packages your working tree as-is, so check out the revision you want (git checkout ) before converting") } + // include_paths narrows the archive to a subset of root_path. The bundle has no + // per-code-source equivalent: deploy packages the whole directory, filtered by + // .gitignore and the bundle-wide sync.include/sync.exclude. Silently dropping it + // would upload files the user meant to leave out. + if len(snap.IncludePaths) > 0 { + return nil, nil, errors.New("code_source.snapshot.include_paths is not supported by convert-to-dabs; deploy packages the whole directory, so narrow it with sync.exclude in the bundle (or a .gitignore) instead") + } } codeSourcePath, err := bundleCodeSourcePath(ctx, cfg, configPath, bundleDir) @@ -348,21 +353,10 @@ func bundleEnvironmentDeps(ctx context.Context, cfg *runConfig, configPath strin return version, doc.Dependencies } -// bundleResourceKey derives a job resource key from the experiment name. The key -// is emitted as an unquoted YAML map key, and DABs' strict loader rejects a key -// that parses as a non-string scalar (a purely numeric name like "12345" -> !!int, -// or "true"/"null"). experiment_name allows exactly [alphanumeric, -, _], so the -// only unsafe keys are those that YAML types as int/float/bool/null; prefix those -// with "job_" to force a string key. The human-facing name/experiment fields keep -// the original value (yamlsaver quotes them as scalar string values). +// bundleResourceKey is the job resource key for an experiment name: the name +// itself. A name that YAML would type as a non-string scalar (a numeric "12345", +// or "true"/"null") is emitted quoted by yamlsaver, so it stays a valid string key. func bundleResourceKey(name string) string { - switch strings.ToLower(name) { - case "true", "false", "null": - return "job_" + name - } - if _, err := strconv.ParseFloat(name, 64); err == nil { - return "job_" + name - } return name } diff --git a/experimental/air/cmd/convert_to_dabs_test.go b/experimental/air/cmd/convert_to_dabs_test.go index 46c5384fb6e..14384c36cd6 100644 --- a/experimental/air/cmd/convert_to_dabs_test.go +++ b/experimental/air/cmd/convert_to_dabs_test.go @@ -8,6 +8,7 @@ import ( "github.com/databricks/cli/libs/dyn" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "go.yaml.in/yaml/v3" ) // Re-running a conversion into a dir that already holds a generated bundle is @@ -348,33 +349,67 @@ permissions: assert.Equal(t, "eng", perms[1].Get("group_name").MustString()) } -// A numeric or reserved-word experiment_name would be an unquoted YAML map key -// that DABs' strict loader rejects (!!int / !!bool). The job resource key is -// prefixed to stay a string, while name/experiment keep the original value. +// The job resource key is the experiment name as-is. A name that YAML would type +// as a non-string scalar ("12345" -> !!int, "true" -> !!bool) must be emitted +// quoted, or the bundle loader rejects the key with "invalid key tag". func TestConvertToDabsSafeJobKey(t *testing.T) { - cases := map[string]string{ - "12345": "job_12345", - "1.5e3": "job_1.5e3", - "true": "job_true", - "null": "job_null", + for _, name := range []string{"12345", "1.5e3", "true", "null", "my-run_1"} { + assert.Equal(t, name, bundleResourceKey(name), "key for %q", name) } - for name, wantKey := range cases { - assert.Equal(t, wantKey, bundleResourceKey(name), "key for %q", name) + + for _, name := range []string{"12345", "true", "my-run_1"} { + t.Run(name, func(t *testing.T) { + dir := t.TempDir() + cfg := "experiment_name: \"" + name + "\"\ncommand: python t.py\n" + + "compute: {accelerator_type: GPU_1xH100, num_accelerators: 1}\n" + path := filepath.Join(dir, "run.yaml") + require.NoError(t, os.WriteFile(path, []byte(cfg), 0o600)) + loaded, err := loadRunConfig(path) + require.NoError(t, err) + + root, _, err := convertToDabs(t.Context(), loaded, path, dir) + require.NoError(t, err) + // The key is the raw name; name/experiment keep the same value. + jobs, err := dyn.GetByPath(dyn.V(root), dyn.MustPathFromString("resources.jobs")) + require.NoError(t, err) + job := jobs.Get(name) + require.True(t, job.IsValid(), "job must be keyed by %q", name) + assert.Equal(t, name, job.Get("name").MustString()) + + // The emitted YAML must load back with the key still a string. + _, err = writeBundle(t.Context(), loaded, path, dir, true) + require.NoError(t, err) + emitted, err := os.ReadFile(filepath.Join(dir, "databricks.yml")) + require.NoError(t, err) + var doc struct { + Resources struct { + Jobs map[string]struct { + Name string `yaml:"name"` + } `yaml:"jobs"` + } `yaml:"resources"` + } + require.NoError(t, yaml.Unmarshal(emitted, &doc), "emitted YAML must parse:\n%s", emitted) + require.Contains(t, doc.Resources.Jobs, name, "job key must load as the string %q:\n%s", name, emitted) + }) } - // A normal name is used as-is. - assert.Equal(t, "my-run_1", bundleResourceKey("my-run_1")) +} - // End to end: a numeric name lands under the prefixed key, but name/experiment - // keep the numeric string value. - cfg := "experiment_name: \"12345\"\ncommand: python t.py\n" + - "compute: {accelerator_type: GPU_1xH100, num_accelerators: 1}\n" - path := writeConfigFile(t, "run.yaml", cfg) +// include_paths narrows the archive to a subset of root_path, which a bundle can't +// express per code source. Rejected rather than silently uploading the whole dir. +func TestConvertToDabsRejectsIncludePaths(t *testing.T) { + dir := t.TempDir() + require.NoError(t, os.MkdirAll(filepath.Join(dir, "src"), 0o700)) + cfg := "experiment_name: inc\ncommand: python train.py\n" + + "compute: {accelerator_type: GPU_1xH100, num_accelerators: 1}\n" + + "code_source:\n type: snapshot\n snapshot:\n root_path: ./src\n" + + " include_paths:\n - keep\n" + path := filepath.Join(dir, "run.yaml") + require.NoError(t, os.WriteFile(path, []byte(cfg), 0o600)) loaded, err := loadRunConfig(path) require.NoError(t, err) - root, _, err := convertToDabs(t.Context(), loaded, path, filepath.Dir(path)) - require.NoError(t, err) - assert.Equal(t, "12345", get(t, root, "resources.jobs.job_12345.name").MustString()) - assert.Equal(t, "12345", get(t, root, "resources.jobs.job_12345.tasks[0].ai_runtime_task.experiment").MustString()) + + _, _, err = convertToDabs(t.Context(), loaded, path, dir) + require.ErrorContains(t, err, "include_paths is not supported") } func TestConvertToDabsRejectsUnsupported(t *testing.T) { diff --git a/libs/dyn/yamlsaver/saver.go b/libs/dyn/yamlsaver/saver.go index 01b345923d2..77d6c11a7f4 100644 --- a/libs/dyn/yamlsaver/saver.go +++ b/libs/dyn/yamlsaver/saver.go @@ -92,7 +92,14 @@ func (s *saver) toYamlNodeWithStyle(v dyn.Value, style yaml.Style) (*yaml.Node, for _, pair := range pairs { pk := pair.Key pv := pair.Value - node := yaml.Node{Kind: yaml.ScalarNode, Value: pk.MustString(), Style: style} + // Quote a key that would otherwise round-trip as a non-string scalar (a + // numeric map key like "12345" loads as !!int, which the bundle loader + // rejects). Values get the same treatment below via isScalarValueInString. + keyStyle := style + if isScalarValueInString(pk) { + keyStyle = yaml.DoubleQuotedStyle + } + node := yaml.Node{Kind: yaml.ScalarNode, Value: pk.MustString(), Style: keyStyle} var nestedNodeStyle yaml.Style if customStyle, ok := s.hasStyle(pk.MustString()); ok { nestedNodeStyle = customStyle From c1c3f186f88002799b8151aa01b688392f991910 Mon Sep 17 00:00:00 2001 From: vinchenzo-db Date: Wed, 5 Aug 2026 19:32:47 +0000 Subject: [PATCH 08/15] air/convert-to-dabs: quote the job key locally, revert yamlsaver Keep the change inside experimental/air/ instead of teaching the shared yamlsaver to quote map keys. quoteJobKey rewrites the emitted job resource key when the experiment name would otherwise load as a non-string scalar ("12345" -> !!int), which the bundle loader rejects. NewSaverWithStyle can't do this: its style map applies to a key's whole subtree, so quoting the job key also quotes every nested key and turns accelerator_count: 1 into a string. Co-authored-by: Isaac --- experimental/air/cmd/convert_to_dabs.go | 50 ++++++++++++++++++++++++- libs/dyn/yamlsaver/saver.go | 9 +---- 2 files changed, 49 insertions(+), 10 deletions(-) diff --git a/experimental/air/cmd/convert_to_dabs.go b/experimental/air/cmd/convert_to_dabs.go index 7b0f33591d2..b0e0f9a3b31 100644 --- a/experimental/air/cmd/convert_to_dabs.go +++ b/experimental/air/cmd/convert_to_dabs.go @@ -1,12 +1,15 @@ package aircmd import ( + "bytes" "context" "errors" "fmt" "os" "path/filepath" "slices" + "strconv" + "strings" "github.com/databricks/cli/cmd/root" "github.com/databricks/cli/libs/cmdio" @@ -354,12 +357,52 @@ func bundleEnvironmentDeps(ctx context.Context, cfg *runConfig, configPath strin } // bundleResourceKey is the job resource key for an experiment name: the name -// itself. A name that YAML would type as a non-string scalar (a numeric "12345", -// or "true"/"null") is emitted quoted by yamlsaver, so it stays a valid string key. +// itself. quoteJobKey quotes it in the emitted YAML when it would otherwise load +// as a non-string scalar. func bundleResourceKey(name string) string { return name } +// quoteJobKey rewrites the emitted job resource key as a quoted YAML key when the +// name would otherwise load as a non-string scalar. yamlsaver emits map keys +// unquoted, and the bundle loader rejects a key that types as something other than +// a string ("12345" -> !!int, "true" -> !!bool) with "invalid key tag". Only the +// job key needs this: every other key convert emits is a fixed schema field name. +func quoteJobKey(bundlePath, key string) error { + if !yamlKeyNeedsQuoting(key) { + return nil + } + data, err := os.ReadFile(bundlePath) + if err != nil { + return err + } + // The key is emitted by buildBundleValue at a known depth under resources.jobs, + // so the indented ":" line is unambiguous. + old := []byte("\n " + key + ":\n") + updated := bytes.Replace(data, old, []byte("\n \""+key+"\":\n"), 1) + if bytes.Equal(data, updated) { + return fmt.Errorf("could not quote job key %q in %s", key, bundlePath) + } + return os.WriteFile(bundlePath, updated, 0o600) +} + +// yamlKeyNeedsQuoting reports whether an unquoted YAML key would load as a +// non-string scalar. experiment_name allows only [alphanumeric, -, _], so the +// cases are numbers and the bool/null words. +func yamlKeyNeedsQuoting(key string) bool { + switch strings.ToLower(key) { + case "true", "false", "null", "yes", "no", "on", "off", "~": + return true + } + if _, err := strconv.ParseFloat(key, 64); err == nil { + return true + } + if _, err := strconv.ParseInt(key, 0, 64); err == nil { + return true + } + return false +} + // buildPermissionsValue maps run-config permissions to DABs job permissions // (level → principal). Returns an invalid value when there are none. func buildPermissionsValue(perms []permission) dyn.Value { @@ -425,6 +468,9 @@ func writeBundle(ctx context.Context, cfg *runConfig, configPath, dir string, fo if err := yamlsaver.NewSaver().SaveAsYAML(root, bundlePath, true); err != nil { return nil, err } + if err := quoteJobKey(bundlePath, bundleResourceKey(cfg.ExperimentName)); err != nil { + return nil, err + } written := []string{"databricks.yml"} // Loose launch artifacts (command.sh + sidecars) at the bundle root. diff --git a/libs/dyn/yamlsaver/saver.go b/libs/dyn/yamlsaver/saver.go index 77d6c11a7f4..01b345923d2 100644 --- a/libs/dyn/yamlsaver/saver.go +++ b/libs/dyn/yamlsaver/saver.go @@ -92,14 +92,7 @@ func (s *saver) toYamlNodeWithStyle(v dyn.Value, style yaml.Style) (*yaml.Node, for _, pair := range pairs { pk := pair.Key pv := pair.Value - // Quote a key that would otherwise round-trip as a non-string scalar (a - // numeric map key like "12345" loads as !!int, which the bundle loader - // rejects). Values get the same treatment below via isScalarValueInString. - keyStyle := style - if isScalarValueInString(pk) { - keyStyle = yaml.DoubleQuotedStyle - } - node := yaml.Node{Kind: yaml.ScalarNode, Value: pk.MustString(), Style: keyStyle} + node := yaml.Node{Kind: yaml.ScalarNode, Value: pk.MustString(), Style: style} var nestedNodeStyle yaml.Style if customStyle, ok := s.hasStyle(pk.MustString()); ok { nestedNodeStyle = customStyle From 9adcf131e150ba58ed82b6615d32cea3e013c079 Mon Sep 17 00:00:00 2001 From: vinchenzo-db Date: Wed, 5 Aug 2026 22:29:42 +0000 Subject: [PATCH 09/15] air/convert-to-dabs: stop emitting training_config.yaml It is a verbatim copy of the input YAML, which bundle sync already uploads, and nothing in the emitted bundle references it. Co-authored-by: Isaac --- acceptance/experimental/air/convert-to-dabs/output.txt | 2 -- acceptance/experimental/air/convert-to-dabs/test.toml | 2 +- experimental/air/cmd/convert_to_dabs.go | 4 +--- experimental/air/cmd/convert_to_dabs_test.go | 8 +++++--- 4 files changed, 7 insertions(+), 9 deletions(-) diff --git a/acceptance/experimental/air/convert-to-dabs/output.txt b/acceptance/experimental/air/convert-to-dabs/output.txt index 1f3ca78959a..5ba2e364eef 100644 --- a/acceptance/experimental/air/convert-to-dabs/output.txt +++ b/acceptance/experimental/air/convert-to-dabs/output.txt @@ -3,7 +3,6 @@ >>> [CLI] experimental air convert-to-dabs train.yaml Wrote a Databricks Asset Bundle to .: databricks.yml - training_config.yaml command.sh To deploy and run this workload as a bundle: @@ -72,7 +71,6 @@ Exit code: 1 >>> [CLI] experimental air convert-to-dabs train.yaml --force Wrote a Databricks Asset Bundle to .: databricks.yml - training_config.yaml command.sh To deploy and run this workload as a bundle: diff --git a/acceptance/experimental/air/convert-to-dabs/test.toml b/acceptance/experimental/air/convert-to-dabs/test.toml index 386f44a60f0..1a43491d6e7 100644 --- a/acceptance/experimental/air/convert-to-dabs/test.toml +++ b/acceptance/experimental/air/convert-to-dabs/test.toml @@ -1,3 +1,3 @@ # convert-to-dabs writes the bundle in place (next to train.yaml + src/). These are # generated artifacts, not committed inputs, so exclude them from the repo-diff check. -Ignore = ["databricks.yml", "command.sh", "training_config.yaml", "generated-docker"] +Ignore = ["databricks.yml", "command.sh", "generated-docker"] diff --git a/experimental/air/cmd/convert_to_dabs.go b/experimental/air/cmd/convert_to_dabs.go index b0e0f9a3b31..76d1f8aa59f 100644 --- a/experimental/air/cmd/convert_to_dabs.go +++ b/experimental/air/cmd/convert_to_dabs.go @@ -145,10 +145,8 @@ func convertToDabs(ctx context.Context, cfg *runConfig, configPath, bundleDir st if err != nil { return nil, nil, err } - // Drop requirements.yaml: the runtime installs pip deps from the job's - // environments[] spec (which convert populates), so a sidecar would be redundant. artifacts = slices.DeleteFunc(artifacts, func(it uploadItem) bool { - return it.name == requirementsName + return it.name == requirementsName || it.name == trainingConfigName }) root := buildBundleValue(ctx, cfg, configPath, codeSourcePath) diff --git a/experimental/air/cmd/convert_to_dabs_test.go b/experimental/air/cmd/convert_to_dabs_test.go index 14384c36cd6..7992d2f1eef 100644 --- a/experimental/air/cmd/convert_to_dabs_test.go +++ b/experimental/air/cmd/convert_to_dabs_test.go @@ -128,11 +128,13 @@ environment: require.Len(t, deps, 2) assert.Equal(t, "numpy", deps[0].MustString()) - // command.sh is always an artifact. requirements.yaml is NOT emitted: the - // deploy-time aicode.SynthesizeRequirements mutator regenerates it from the - // environments[] spec (asserted above), so convert must not also write it. + // command.sh is always an artifact. Two sidecars the run path uploads are not + // emitted: requirements.yaml (deps ride on the environments[] spec asserted + // above) and training_config.yaml (a verbatim copy of the input YAML, which is + // already synced into the bundle and referenced by nothing). assert.Contains(t, itemNames(artifacts), commandScriptName) assert.NotContains(t, itemNames(artifacts), requirementsName) + assert.NotContains(t, itemNames(artifacts), trainingConfigName) } // Optional fields are omitted rather than emitted empty: no code_source means no From d3b3a773e58f9c0f1dad915736cad040daf71152 Mon Sep 17 00:00:00 2001 From: vinchenzo-db Date: Wed, 5 Aug 2026 23:58:05 +0000 Subject: [PATCH 10/15] air/convert-to-dabs: emit launch artifacts under generated_artifacts/ Write command.sh and the env/secret/param sidecars into generated_artifacts/ and emit sync.paths listing only that directory. The code directory is no longer synced as loose files: deploy still packages it into the snapshot tarball, so it was previously uploaded twice, which is costly for a large source tree. The server derives the sidecar paths from command_path's parent, so they move with command.sh. Co-authored-by: Isaac --- .../air/convert-to-dabs/output.txt | 11 +++-- .../experimental/air/convert-to-dabs/script | 2 +- .../air/convert-to-dabs/test.toml | 2 +- experimental/air/cmd/convert_to_dabs.go | 40 ++++++++++++++----- experimental/air/cmd/convert_to_dabs_test.go | 31 +++++++++++++- 5 files changed, 67 insertions(+), 19 deletions(-) diff --git a/acceptance/experimental/air/convert-to-dabs/output.txt b/acceptance/experimental/air/convert-to-dabs/output.txt index 5ba2e364eef..3e4b203beda 100644 --- a/acceptance/experimental/air/convert-to-dabs/output.txt +++ b/acceptance/experimental/air/convert-to-dabs/output.txt @@ -3,7 +3,7 @@ >>> [CLI] experimental air convert-to-dabs train.yaml Wrote a Databricks Asset Bundle to .: databricks.yml - command.sh + generated_artifacts/command.sh To deploy and run this workload as a bundle: 1. [CLI] bundle validate @@ -22,6 +22,9 @@ job and its uploaded files with: >>> cat databricks.yml bundle: name: torchrun-a10-smoke-test +sync: + paths: + - generated_artifacts targets: dev: mode: development @@ -36,7 +39,7 @@ resources: ai_runtime_task: experiment: torchrun-a10-smoke-test deployments: - - command_path: ./command.sh + - command_path: ./generated_artifacts/command.sh compute: accelerator_type: GPU_1xA10 accelerator_count: 1 @@ -49,7 +52,7 @@ resources: - numpy === the generated command.sh carries the run command ->>> cat command.sh +>>> cat generated_artifacts/command.sh torchrun --nproc_per_node=1 train.py === the emitted bundle validates >>> [CLI] bundle validate @@ -71,7 +74,7 @@ Exit code: 1 >>> [CLI] experimental air convert-to-dabs train.yaml --force Wrote a Databricks Asset Bundle to .: databricks.yml - command.sh + generated_artifacts/command.sh To deploy and run this workload as a bundle: 1. [CLI] bundle validate diff --git a/acceptance/experimental/air/convert-to-dabs/script b/acceptance/experimental/air/convert-to-dabs/script index 8882963ac91..77b6e06938a 100644 --- a/acceptance/experimental/air/convert-to-dabs/script +++ b/acceptance/experimental/air/convert-to-dabs/script @@ -5,7 +5,7 @@ title "emitted databricks.yml (code_source_path points at ./src; no code is copi trace cat databricks.yml title "the generated command.sh carries the run command" -trace cat command.sh +trace cat generated_artifacts/command.sh title "the emitted bundle validates" trace $CLI bundle validate diff --git a/acceptance/experimental/air/convert-to-dabs/test.toml b/acceptance/experimental/air/convert-to-dabs/test.toml index 1a43491d6e7..aa531fbd106 100644 --- a/acceptance/experimental/air/convert-to-dabs/test.toml +++ b/acceptance/experimental/air/convert-to-dabs/test.toml @@ -1,3 +1,3 @@ # convert-to-dabs writes the bundle in place (next to train.yaml + src/). These are # generated artifacts, not committed inputs, so exclude them from the repo-diff check. -Ignore = ["databricks.yml", "command.sh", "generated-docker"] +Ignore = ["databricks.yml", "generated_artifacts", "generated-docker"] diff --git a/experimental/air/cmd/convert_to_dabs.go b/experimental/air/cmd/convert_to_dabs.go index 76d1f8aa59f..96d8f9b07db 100644 --- a/experimental/air/cmd/convert_to_dabs.go +++ b/experimental/air/cmd/convert_to_dabs.go @@ -6,6 +6,7 @@ import ( "errors" "fmt" "os" + "path" "path/filepath" "slices" "strconv" @@ -26,8 +27,8 @@ import ( // schema-valid ai_runtime_task (the SDK jobs.AiRuntimeTask — experiment + // deployments[].{command_path,compute} + code_source_path, with framework fields // like retries/timeout on the surrounding task) and writes command.sh plus the -// env_vars.json / secret_env_vars.json / hyperparameters.yaml sidecars at the -// bundle root. It does NOT package, snapshot, or upload anything. +// env_vars.json / secret_env_vars.json / hyperparameters.yaml sidecars into +// generated_artifacts/. It does NOT package, snapshot, or upload anything. // // code_source_path is emitted as the source directory relative to the bundle (the // bundle root defaults to the YAML's directory, which contains it). At deploy the @@ -40,6 +41,12 @@ import ( // is the conventional starting point for a generated bundle. const dabsTargetName = "dev" +// generatedArtifactsDir holds command.sh and the env/secret/param sidecars, kept +// apart from the user's tree so sync.paths can list it without the code directory. +// The server derives the sidecar paths from command_path's parent, so they must +// stay beside command.sh. +const generatedArtifactsDir = "generated_artifacts" + func newConvertToDabsCommand() *cobra.Command { var ( outputDir string @@ -100,7 +107,7 @@ does not contact the workspace.`, // run config. It reads only what the run path's buildArtifacts reads, so the // mapping is unit-testable in isolation. Returns the bundle root as a // map[string]dyn.Value (ready for yamlsaver) and the loose artifacts (command.sh + -// env/secret/param sidecars) to write at the bundle root. It does not touch the +// env/secret/param sidecars) to write under generated_artifacts/. It does not touch the // code_source; the deploy-time aicode mutator packages it in place. func convertToDabs(ctx context.Context, cfg *runConfig, configPath, bundleDir string) (map[string]dyn.Value, []uploadItem, error) { // idempotency_token is intentionally not mapped: it dedups a single runs/submit @@ -214,7 +221,7 @@ func buildBundleValue(ctx context.Context, cfg *runConfig, configPath, codeSourc // it: libraries.IsLibraryLocal classifies a bare, extensionless path as a PyPI // package name, which would deploy a path the backend can't resolve. deployment := map[string]dyn.Value{ - "command_path": nv(localBundlePath(commandScriptName), 1), + "command_path": nv(localBundlePath(path.Join(generatedArtifactsDir, commandScriptName)), 1), "compute": nv(map[string]dyn.Value{ "accelerator_type": nv(cfg.Compute.AcceleratorType, 1), "accelerator_count": nv(cfg.Compute.NumAccelerators, 2), @@ -300,17 +307,23 @@ func buildBundleValue(ctx context.Context, cfg *runConfig, configPath, codeSourc "bundle": nv(map[string]dyn.Value{ "name": nv(name, 1), }, 1), + // sync.paths replaces the default of syncing the whole bundle root. The code + // directory is omitted deliberately: deploy still packages it into the + // snapshot tarball, so syncing it too would upload the tree twice. + "sync": nv(map[string]dyn.Value{ + "paths": nv([]dyn.Value{nv(generatedArtifactsDir, 1)}, 1), + }, 2), "targets": nv(map[string]dyn.Value{ dabsTargetName: nv(map[string]dyn.Value{ "mode": nv("development", 1), "default": nv(true, 2), }, 1), - }, 2), + }, 3), "resources": nv(map[string]dyn.Value{ "jobs": nv(map[string]dyn.Value{ bundleResourceKey(name): nv(job, 1), }, 1), - }, 3), + }, 4), } return rootValue } @@ -446,13 +459,18 @@ func writeBundle(ctx context.Context, cfg *runConfig, configPath, dir string, fo // overwrite of edits they may have made. The hint names --force rather than // --output-dir: the code_source must live inside the bundle dir, so redirecting // the output usually isn't a usable escape hatch for an in-place conversion. + artifactDir := filepath.Join(dir, generatedArtifactsDir) + if err := os.MkdirAll(artifactDir, 0o700); err != nil { + return nil, err + } writeFile := func(name string, data []byte) error { + target := filepath.Join(artifactDir, name) if !force { - if _, err := os.Stat(filepath.Join(dir, name)); err == nil { - return fmt.Errorf("%s already exists in %s; pass --force to overwrite or remove it", name, dir) + if _, err := os.Stat(target); err == nil { + return fmt.Errorf("%s already exists in %s; pass --force to overwrite or remove it", name, artifactDir) } } - return os.WriteFile(filepath.Join(dir, name), data, 0o600) + return os.WriteFile(target, data, 0o600) } bundlePath := filepath.Join(dir, "databricks.yml") @@ -471,12 +489,12 @@ func writeBundle(ctx context.Context, cfg *runConfig, configPath, dir string, fo } written := []string{"databricks.yml"} - // Loose launch artifacts (command.sh + sidecars) at the bundle root. + // Launch artifacts (command.sh + sidecars) under generated_artifacts/. for _, item := range artifacts { if err := writeFile(item.name, item.data); err != nil { return nil, fmt.Errorf("failed to write %s: %w", item.name, err) } - written = append(written, item.name) + written = append(written, path.Join(generatedArtifactsDir, item.name)) } return written, nil diff --git a/experimental/air/cmd/convert_to_dabs_test.go b/experimental/air/cmd/convert_to_dabs_test.go index 7992d2f1eef..dee150269e4 100644 --- a/experimental/air/cmd/convert_to_dabs_test.go +++ b/experimental/air/cmd/convert_to_dabs_test.go @@ -117,7 +117,7 @@ environment: assert.Equal(t, "./src", get(t, root, art+".code_source_path").MustString()) dep := art + ".deployments[0]" - assert.Equal(t, "./"+commandScriptName, get(t, root, dep+".command_path").MustString()) + assert.Equal(t, "./"+generatedArtifactsDir+"/"+commandScriptName, get(t, root, dep+".command_path").MustString()) assert.Equal(t, "GPU_1xH100", get(t, root, dep+".compute.accelerator_type").MustString()) assert.Equal(t, int64(1), get(t, root, dep+".compute.accelerator_count").MustInt()) @@ -139,6 +139,33 @@ environment: // Optional fields are omitted rather than emitted empty: no code_source means no // code_source_path; unset retries/timeout means no wrapper fields. +// sync.paths lists only the generated-artifacts dir. The code directory must be +// absent: deploy packages it into the snapshot tarball, so syncing it as loose +// files too would upload the whole tree a second time. +func TestConvertToDabsSyncPathsExcludesCodeDir(t *testing.T) { + cfg := minimalConfig + ` +code_source: + type: snapshot + snapshot: + root_path: ./src +` + path := writeConfigFile(t, "run.yaml", cfg) + require.NoError(t, os.MkdirAll(filepath.Join(filepath.Dir(path), "src"), 0o700)) + loaded, err := loadRunConfig(path) + require.NoError(t, err) + + root, _, err := convertToDabs(t.Context(), loaded, path, filepath.Dir(path)) + require.NoError(t, err) + + paths := get(t, root, "sync.paths").MustSequence() + require.Len(t, paths, 1) + assert.Equal(t, generatedArtifactsDir, paths[0].MustString()) + + // The task still points at the code dir; only the sync set omits it. + task := "resources.jobs." + loaded.ExperimentName + ".tasks[0].ai_runtime_task" + assert.Equal(t, "./src", get(t, root, task+".code_source_path").MustString()) +} + func TestConvertToDabsOmitsUnsetFields(t *testing.T) { path := writeConfigFile(t, "run.yaml", minimalConfig) loaded, err := loadRunConfig(path) @@ -154,7 +181,7 @@ func TestConvertToDabsOmitsUnsetFields(t *testing.T) { assert.False(t, has(root, task+".ai_runtime_task.code_source_path")) assert.False(t, has(root, task+".ai_runtime_task.mlflow_run")) // The command still needs a home even without code_source. - assert.Equal(t, "./"+commandScriptName, get(t, root, task+".ai_runtime_task.deployments[0].command_path").MustString()) + assert.Equal(t, "./"+generatedArtifactsDir+"/"+commandScriptName, get(t, root, task+".ai_runtime_task.deployments[0].command_path").MustString()) // Even with no environment block, the default runtime version is pinned (what // `air run` would have used) rather than emitting an empty environment spec. From ed62059508b2833d5d85a17775dedf1a508d6656 Mon Sep 17 00:00:00 2001 From: vinchenzo-db Date: Thu, 6 Aug 2026 00:45:55 +0000 Subject: [PATCH 11/15] air/convert-to-dabs: point the out-of-bundle error at sync.paths The old message suggested --output-dir, which cannot help here: root_path is resolved relative to the input YAML, so redirecting the output moves the bundle away from the code rather than enclosing it. sync.paths is the mechanism that actually widens the sync root. Co-authored-by: Isaac --- experimental/air/cmd/convert_to_dabs.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/experimental/air/cmd/convert_to_dabs.go b/experimental/air/cmd/convert_to_dabs.go index 96d8f9b07db..d5e0c399dd5 100644 --- a/experimental/air/cmd/convert_to_dabs.go +++ b/experimental/air/cmd/convert_to_dabs.go @@ -187,7 +187,7 @@ func bundleCodeSourcePath(ctx context.Context, cfg *runConfig, configPath, bundl } rel, err := filepath.Rel(bundleAbs, root) if err != nil || !filepath.IsLocal(rel) { - return "", fmt.Errorf("code_source root_path %q is not inside the bundle directory %q; run convert-to-dabs with --output-dir set to an ancestor of the code", snap.RootPath, bundleDir) + return "", fmt.Errorf("code_source root_path %q is not inside the bundle directory %q; bundle deploy can only upload files under the bundle's sync root. Move the code inside the bundle, or set sync.paths in databricks.yml to a directory that contains both (note that widens what gets synced)", snap.RootPath, bundleDir) } return localBundlePath(filepath.ToSlash(rel)), nil } From cf1ceda0b1ddd8be3fe0b173406af2f9d64aec22 Mon Sep 17 00:00:00 2001 From: vinchenzo-db Date: Thu, 6 Aug 2026 19:02:01 +0000 Subject: [PATCH 12/15] air/convert-to-dabs: always emit max_retries, fix --output-dir help MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit air run fills its own default of 3 when max_retries is unset, so omitting it gave the converted bundle the Jobs default instead — the same YAML retried differently depending on which path launched it. The --output-dir help still described a -bundle folder; the default is the input YAML's directory. Co-authored-by: Isaac --- acceptance/experimental/air/convert-to-dabs/output.txt | 1 + experimental/air/cmd/convert_to_dabs.go | 8 +++----- experimental/air/cmd/convert_to_dabs_test.go | 4 +++- 3 files changed, 7 insertions(+), 6 deletions(-) diff --git a/acceptance/experimental/air/convert-to-dabs/output.txt b/acceptance/experimental/air/convert-to-dabs/output.txt index 3e4b203beda..05e584fc8cc 100644 --- a/acceptance/experimental/air/convert-to-dabs/output.txt +++ b/acceptance/experimental/air/convert-to-dabs/output.txt @@ -36,6 +36,7 @@ resources: tasks: - task_key: torchrun-a10-smoke-test environment_key: default + max_retries: 3 ai_runtime_task: experiment: torchrun-a10-smoke-test deployments: diff --git a/experimental/air/cmd/convert_to_dabs.go b/experimental/air/cmd/convert_to_dabs.go index d5e0c399dd5..799124e9206 100644 --- a/experimental/air/cmd/convert_to_dabs.go +++ b/experimental/air/cmd/convert_to_dabs.go @@ -69,7 +69,7 @@ upload step is required. This command performs a purely local translation and does not contact the workspace.`, } - cmd.Flags().StringVar(&outputDir, "output-dir", "", "Directory to write the bundle into (default: a -bundle folder next to the input YAML). Accepts an absolute or relative path.") + cmd.Flags().StringVar(&outputDir, "output-dir", "", "Directory to write the bundle into (default: the input YAML's directory). Must contain the code_source. Accepts an absolute or relative path.") cmd.Flags().BoolVar(&force, "force", false, "Overwrite the generated bundle files if they already exist.") cmd.RunE = func(cmd *cobra.Command, args []string) error { @@ -255,10 +255,8 @@ func buildBundleValue(ctx context.Context, cfg *runConfig, configPath, codeSourc "environment_key": nv(aiRuntimeEnvironmentKey, 2), } taskLine := 3 - if cfg.MaxRetries != nil { - task["max_retries"] = nv(*cfg.MaxRetries, taskLine) - taskLine++ - } + task["max_retries"] = nv(cfg.maxRetries(), taskLine) + taskLine++ if cfg.TimeoutMinutes != nil { task["timeout_seconds"] = nv(cfg.timeoutSeconds(), taskLine) taskLine++ diff --git a/experimental/air/cmd/convert_to_dabs_test.go b/experimental/air/cmd/convert_to_dabs_test.go index dee150269e4..faa7c30cfdd 100644 --- a/experimental/air/cmd/convert_to_dabs_test.go +++ b/experimental/air/cmd/convert_to_dabs_test.go @@ -176,7 +176,9 @@ func TestConvertToDabsOmitsUnsetFields(t *testing.T) { name := loaded.ExperimentName task := "resources.jobs." + name + ".tasks[0]" - assert.False(t, has(root, task+".max_retries")) + // max_retries is always emitted: `air run` fills its own default when unset, so + // omitting it would silently give the bundle the Jobs default instead. + assert.Equal(t, int64(defaultMaxRetries), get(t, root, task+".max_retries").MustInt()) assert.False(t, has(root, task+".timeout_seconds")) assert.False(t, has(root, task+".ai_runtime_task.code_source_path")) assert.False(t, has(root, task+".ai_runtime_task.mlflow_run")) From 7230cb531e1de0e90157f0859dea41b63ef34f8a Mon Sep 17 00:00:00 2001 From: vinchenzo-db Date: Thu, 6 Aug 2026 20:37:24 +0000 Subject: [PATCH 13/15] air/convert-to-dabs: keep emitting training_config.yaml Reverts dropping it. The Jobs run-output page derives its workspace path from command_path (same directory, fixed filename) to show the config a run used, so omitting the file leaves that link pointing at nothing. The input YAML cannot serve instead: `air run -f` accepts any path, and real configs are named train.yaml, workload.yaml, train_kie.yaml and so on, so the UI has no filename to derive. The duplicate copy at a fixed name is what makes the path resolvable. Co-authored-by: Isaac --- acceptance/experimental/air/convert-to-dabs/output.txt | 2 ++ experimental/air/cmd/convert_to_dabs.go | 6 +++++- experimental/air/cmd/convert_to_dabs_test.go | 9 ++++----- 3 files changed, 11 insertions(+), 6 deletions(-) diff --git a/acceptance/experimental/air/convert-to-dabs/output.txt b/acceptance/experimental/air/convert-to-dabs/output.txt index 05e584fc8cc..8b94bf55a7c 100644 --- a/acceptance/experimental/air/convert-to-dabs/output.txt +++ b/acceptance/experimental/air/convert-to-dabs/output.txt @@ -3,6 +3,7 @@ >>> [CLI] experimental air convert-to-dabs train.yaml Wrote a Databricks Asset Bundle to .: databricks.yml + generated_artifacts/training_config.yaml generated_artifacts/command.sh To deploy and run this workload as a bundle: @@ -75,6 +76,7 @@ Exit code: 1 >>> [CLI] experimental air convert-to-dabs train.yaml --force Wrote a Databricks Asset Bundle to .: databricks.yml + generated_artifacts/training_config.yaml generated_artifacts/command.sh To deploy and run this workload as a bundle: diff --git a/experimental/air/cmd/convert_to_dabs.go b/experimental/air/cmd/convert_to_dabs.go index 799124e9206..48d35868a3e 100644 --- a/experimental/air/cmd/convert_to_dabs.go +++ b/experimental/air/cmd/convert_to_dabs.go @@ -152,8 +152,12 @@ func convertToDabs(ctx context.Context, cfg *runConfig, configPath, bundleDir st if err != nil { return nil, nil, err } + // training_config.yaml is kept even though it duplicates the input YAML: the Jobs + // run-output page derives its path from command_path (same directory, fixed name), + // so dropping it breaks the "see what you ran" link. The input YAML can't serve + // instead — `air run -f` accepts any filename, so the UI has nothing to derive. artifacts = slices.DeleteFunc(artifacts, func(it uploadItem) bool { - return it.name == requirementsName || it.name == trainingConfigName + return it.name == requirementsName }) root := buildBundleValue(ctx, cfg, configPath, codeSourcePath) diff --git a/experimental/air/cmd/convert_to_dabs_test.go b/experimental/air/cmd/convert_to_dabs_test.go index faa7c30cfdd..3cad413f20f 100644 --- a/experimental/air/cmd/convert_to_dabs_test.go +++ b/experimental/air/cmd/convert_to_dabs_test.go @@ -128,13 +128,12 @@ environment: require.Len(t, deps, 2) assert.Equal(t, "numpy", deps[0].MustString()) - // command.sh is always an artifact. Two sidecars the run path uploads are not - // emitted: requirements.yaml (deps ride on the environments[] spec asserted - // above) and training_config.yaml (a verbatim copy of the input YAML, which is - // already synced into the bundle and referenced by nothing). + // command.sh and training_config.yaml are always artifacts; the Jobs run-output + // page derives the latter's path from command_path. requirements.yaml is not + // emitted — deps ride on the environments[] spec asserted above. assert.Contains(t, itemNames(artifacts), commandScriptName) + assert.Contains(t, itemNames(artifacts), trainingConfigName) assert.NotContains(t, itemNames(artifacts), requirementsName) - assert.NotContains(t, itemNames(artifacts), trainingConfigName) } // Optional fields are omitted rather than emitted empty: no code_source means no From 7ed25890dca9ba2b12e2936bfcfc610d5a77ab28 Mon Sep 17 00:00:00 2001 From: vinchenzo-db Date: Thu, 6 Aug 2026 23:18:36 +0000 Subject: [PATCH 14/15] Re-trigger CI Co-authored-by: Isaac From 16219ba8e16481448bb5de774da02b5c5cc7c1cf Mon Sep 17 00:00:00 2001 From: vinchenzo-db Date: Thu, 6 Aug 2026 23:53:32 +0000 Subject: [PATCH 15/15] air/convert-to-dabs: fix build after main merge (requirementsName removed) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The main merge (AIR CLI Migration Pt. 2, #6166) stopped emitting a requirements.yaml sidecar — file-form deps now fold into the environments[] spec — and deleted the requirementsName constant. convert_to_dabs still filtered artifacts by it, breaking the build (undefined: requirementsName), which in turn failed lint, validate-generated, and every acceptance test. Drop the now-dead filter: buildArtifacts no longer produces requirements.yaml, so there is nothing to remove. Also update the usage_policy_id test to a UUID (the merge added UUID validation) and regenerate the engine-matrix snapshot. Co-authored-by: Isaac --- .../experimental/air/convert-to-dabs/out.test.toml | 2 +- experimental/air/cmd/convert_to_dabs.go | 13 +++++-------- experimental/air/cmd/convert_to_dabs_test.go | 8 ++++---- 3 files changed, 10 insertions(+), 13 deletions(-) diff --git a/acceptance/experimental/air/convert-to-dabs/out.test.toml b/acceptance/experimental/air/convert-to-dabs/out.test.toml index f784a183258..e90b6d5d1ba 100644 --- a/acceptance/experimental/air/convert-to-dabs/out.test.toml +++ b/acceptance/experimental/air/convert-to-dabs/out.test.toml @@ -1,3 +1,3 @@ Local = true Cloud = false -EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["terraform", "direct"] +EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["direct"] diff --git a/experimental/air/cmd/convert_to_dabs.go b/experimental/air/cmd/convert_to_dabs.go index 48d35868a3e..f5563136176 100644 --- a/experimental/air/cmd/convert_to_dabs.go +++ b/experimental/air/cmd/convert_to_dabs.go @@ -8,7 +8,6 @@ import ( "os" "path" "path/filepath" - "slices" "strconv" "strings" @@ -148,17 +147,15 @@ func convertToDabs(ctx context.Context, cfg *runConfig, configPath, bundleDir st return nil, nil, err } + // buildArtifacts emits command.sh plus the training_config / hyperparameters / + // env / secret sidecars, all co-located so the Jobs run-output page can derive + // their paths from command_path. It no longer produces a requirements.yaml — + // file-form deps are folded into the environments[] spec — so there is nothing to + // filter out here. artifacts, err := buildArtifacts(cfg, configPath) if err != nil { return nil, nil, err } - // training_config.yaml is kept even though it duplicates the input YAML: the Jobs - // run-output page derives its path from command_path (same directory, fixed name), - // so dropping it breaks the "see what you ran" link. The input YAML can't serve - // instead — `air run -f` accepts any filename, so the UI has nothing to derive. - artifacts = slices.DeleteFunc(artifacts, func(it uploadItem) bool { - return it.name == requirementsName - }) root := buildBundleValue(ctx, cfg, configPath, codeSourcePath) return root, artifacts, nil diff --git a/experimental/air/cmd/convert_to_dabs_test.go b/experimental/air/cmd/convert_to_dabs_test.go index 3cad413f20f..0350f010367 100644 --- a/experimental/air/cmd/convert_to_dabs_test.go +++ b/experimental/air/cmd/convert_to_dabs_test.go @@ -133,7 +133,7 @@ environment: // emitted — deps ride on the environments[] spec asserted above. assert.Contains(t, itemNames(artifacts), commandScriptName) assert.Contains(t, itemNames(artifacts), trainingConfigName) - assert.NotContains(t, itemNames(artifacts), requirementsName) + assert.NotContains(t, itemNames(artifacts), "requirements.yaml") } // Optional fields are omitted rather than emitted empty: no code_source means no @@ -341,20 +341,20 @@ func TestConvertToDabsFoldsRequirementsFileIntoEnvSpec(t *testing.T) { assert.Equal(t, "pandas", deps[1].MustString()) // No requirements.yaml artifact: the mutator regenerates it from the spec. - assert.NotContains(t, itemNames(artifacts), requirementsName) + assert.NotContains(t, itemNames(artifacts), "requirements.yaml") } // usage_policy_id is a resolved budget policy id and maps to the job's // budget_policy_id (usage_policy_name, which needs resolution, is rejected). func TestConvertToDabsMapsUsagePolicyID(t *testing.T) { - cfg := minimalConfig + "usage_policy_id: budget-abc-123\n" + cfg := minimalConfig + "usage_policy_id: 12345678-90ab-cdef-1234-567890abcdef\n" path := writeConfigFile(t, "run.yaml", cfg) loaded, err := loadRunConfig(path) require.NoError(t, err) root, _, err := convertToDabs(t.Context(), loaded, path, filepath.Dir(path)) require.NoError(t, err) - assert.Equal(t, "budget-abc-123", get(t, root, "resources.jobs."+loaded.ExperimentName+".budget_policy_id").MustString()) + assert.Equal(t, "12345678-90ab-cdef-1234-567890abcdef", get(t, root, "resources.jobs."+loaded.ExperimentName+".budget_policy_id").MustString()) } func TestConvertToDabsMapsPermissions(t *testing.T) {