From 8b7047acf74d5060a69b9c924dc983aaa0e5f4cd Mon Sep 17 00:00:00 2001 From: vinchenzo-db Date: Fri, 31 Jul 2026 05:43:30 +0000 Subject: [PATCH 1/2] 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 fd0ab4aad9cfd587813d8dc8ee6f7a556818c1ca Mon Sep 17 00:00:00 2001 From: vinchenzo-db Date: Fri, 31 Jul 2026 06:44:45 +0000 Subject: [PATCH 2/2] 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 {