From 3066eb70c828e9b1e3a801872e3d849a563a64ed Mon Sep 17 00:00:00 2001 From: vinchenzo-db Date: Wed, 15 Jul 2026 19:04:29 +0000 Subject: [PATCH 01/11] [air] Prototype: air run/list/cancel/get via DABs (convert -> deploy -> run) with env vars Wires the experimental `air` CLI to submit and list workloads via Databricks Asset Bundles (persistent Jobs resources) alongside the ephemeral runs/submit path. Gated behind --via-dabs / AIR_VIA_DABS=1; the ephemeral path stays the default and is untouched (additive). - run: `air run --via-dabs` converts train.yaml -> databricks.yml (reusing the exportbundle.go converter), then deploys the bundle and runs it. env vars/secrets ride the common Jobs env-var API (API-293) as a job-level environment_variables profile the ai_runtime_task references by key; secrets emit as {{secrets/scope/key}} refs. Verified against staging that this shape persists on runs/get, so the converter's env-var/secret rejection is lifted. - list: `air list --via-dabs` drops the SUBMIT_RUN server-side restriction so DABs-submitted AIR runs (RunType JOB_RUN) are included; isAirRun still selects by task shape. The AiTrainingService index path already returns the caller's runs regardless of submit verb, so only the jobs-scan fallback needed widening. - cancel / get: work unchanged for DABs runs (both key on a run_id that a bundle run still produces); documented inline. Prototype scope: deploy/run shell out to `databricks bundle` (production would call cmd/bundle/utils.ProcessBundle in-process); run_id parsing and code_source staging into the bundle folder are TODO(prototype). Persistent-job GC is accepted for now. Co-authored-by: Isaac --- experimental/air/cmd/cancel.go | 6 + experimental/air/cmd/exportbundle.go | 315 ++++++++++++++++++++++ experimental/air/cmd/exportbundle_test.go | 139 ++++++++++ experimental/air/cmd/get.go | 3 + experimental/air/cmd/list.go | 15 ++ experimental/air/cmd/list_test.go | 44 +++ experimental/air/cmd/run.go | 32 ++- experimental/air/cmd/rundabs.go | 199 ++++++++++++++ 8 files changed, 752 insertions(+), 1 deletion(-) create mode 100644 experimental/air/cmd/exportbundle.go create mode 100644 experimental/air/cmd/exportbundle_test.go create mode 100644 experimental/air/cmd/rundabs.go diff --git a/experimental/air/cmd/cancel.go b/experimental/air/cmd/cancel.go index 519a7a82068..8f93d0ac928 100644 --- a/experimental/air/cmd/cancel.go +++ b/experimental/air/cmd/cancel.go @@ -191,6 +191,12 @@ func runNotFound(err error) bool { // cancelRun requests cancellation of a single job run. The cancel is async, so // the returned waiter is ignored. +// +// DABs note: this works unchanged for DABs-submitted runs. A `bundle run` produces +// a normal Jobs run with a run_id, and cancelling that run terminates the AIR +// workload the same way — BYOT/AICM dispatch is keyed on task type, not on how the +// run was submitted. (`cancel --all` discovers runs via `air list`, which the +// --via-dabs list mode widens to include DABs JOB_RUNs.) func cancelRun(ctx context.Context, w *databricks.WorkspaceClient, rid string) error { runID, err := strconv.ParseInt(rid, 10, 64) if err != nil || runID <= 0 { diff --git a/experimental/air/cmd/exportbundle.go b/experimental/air/cmd/exportbundle.go new file mode 100644 index 00000000000..f4feb2c899b --- /dev/null +++ b/experimental/air/cmd/exportbundle.go @@ -0,0 +1,315 @@ +package aircmd + +import ( + "fmt" + "path/filepath" + "strings" + + "go.yaml.in/yaml/v3" +) + +// This file implements the train.yaml -> databricks.yml conversion behind +// `air export-bundle`: a one-time converter that turns an AIR run YAML into a +// Databricks Asset Bundle deploying the same workload as a native ai_runtime_task. +// It is the durable counterpart to `air run` (which submits an ephemeral one-time +// run): the emitted bundle is a persistent Jobs resource the user owns and deploys +// with `databricks bundle deploy`. +// +// The field mapping mirrors buildSubmitPayload (runsubmit.go) so the generated +// task is the shape `air run` would submit. Structural validity is already +// guaranteed upstream by runConfig.validate(); this file adds a second, +// converter-specific gate — checkBundleConvertible — that rejects the configs a +// bundle cannot represent faithfully, rather than emitting a databricks.yml that +// deploys but misbehaves. The gate exists because train.yaml and a bundle are not +// 1:1: some run fields have no bundle equivalent, and some assume the AIR run +// harness (e.g. $CODE_SOURCE_PATH) that a bundle does not provide. + +// bundleCommandScript is the entrypoint filename the emitted bundle references and +// that `bundle sync` uploads alongside the user's code. +const bundleCommandScript = "command.sh" + +// codeSourcePathVar is the environment variable the AIR run harness sets to the +// extracted snapshot directory. A bundle delivers code via `bundle sync` and never +// sets it, so a command relying on it would break at runtime — the convertibility +// gate rejects such commands. +const codeSourcePathVar = "$CODE_SOURCE_PATH" + +// workspaceFilePathRef is the bundle variable that resolves to where `bundle +// deploy` syncs this folder; the emitted command_path points under it. +const workspaceFilePathRef = "${workspace.file_path}" + +// aiRuntimeEnvVarsKey is the key linking the task's environment_variables_key to +// the single job-level env-var profile the converter emits. Mirrors +// AI_RUNTIME_ENV_VARS_KEY in the Python CLI (common Jobs env-var API, API-293). +const aiRuntimeEnvVarsKey = "default" + +// checkBundleConvertible reports why a structurally-valid runConfig cannot be +// converted to a faithful bundle, or nil if it can. Every reason names the source +// field and why it can't be represented, so the CLI can reject with an actionable +// message instead of emitting a lossy databricks.yml. +func checkBundleConvertible(cfg *runConfig) error { + var reasons []string + + // env_variables / secrets ARE now representable: they ride the common Jobs + // env-var API (API-293) as a job-level environment_variables profile the task + // references by key. convertToBundle emits that profile, so no rejection here. + // (Verified on staging: the profile persists on runs/get; secrets are emitted + // as {{secrets/scope/key}} refs resolved by Jobs at run time.) + + // docker_image: a custom image must be registered before it can be referenced; + // that registration is not part of a bundle deploy. + if cfg.dockerImageURL() != "" { + reasons = append(reasons, "environment.docker_image: needs image registration, which a bundle deploy does not perform") + } + + // usage_policy_*: the name/id resolves to a budget policy via a workspace + // lookup at submit time; a static bundle has nowhere to run that lookup. + if cfg.UsagePolicyName != nil { + reasons = append(reasons, "usage_policy_name: resolved by a workspace lookup at submit time, not representable statically") + } + if cfg.UsagePolicyID != nil { + reasons = append(reasons, "usage_policy_id: budget policy binding is not represented on the ai_runtime_task bundle path yet") + } + + // code_source with a git ref: bundle sync uploads the working tree; it cannot + // pin to a specific commit or fetch a remote branch (the R1/R3 gap in the + // uploads-vs-DABs design doc). Dropping the pin would silently change what runs. + if cfg.CodeSource != nil && cfg.CodeSource.Snapshot != nil && cfg.CodeSource.Snapshot.Git != nil { + reasons = append(reasons, "code_source.snapshot.git: bundle sync cannot pin to a git commit or fetch a remote branch") + } + + // A command that reads $CODE_SOURCE_PATH assumes the AIR run harness, which a + // bundle does not provide; the synced code lives under ${workspace.file_path}. + if cfg.Command != nil && strings.Contains(*cfg.Command, codeSourcePathVar) { + reasons = append(reasons, fmt.Sprintf( + "command references %s, which only exists on the `air run` path; a bundle syncs code to %s instead", + codeSourcePathVar, workspaceFilePathRef)) + } + + if len(reasons) == 0 { + return nil + } + return fmt.Errorf( + "this train.yaml cannot be converted to a faithful bundle:\n - %s\nrun it with `air run` until these are supported on the bundle path", + strings.Join(reasons, "\n - ")) +} + +// exportedBundle is the minimal databricks.yml shape the converter emits: a bundle +// name plus one job with a single ai_runtime_task. It marshals to YAML, so field +// order here is the emitted key order. +type exportedBundle struct { + Bundle bundleBlock `yaml:"bundle"` + Resources exportedResourcesBlock `yaml:"resources"` +} + +type bundleBlock struct { + Name string `yaml:"name"` +} + +type exportedResourcesBlock struct { + Jobs map[string]exportedJob `yaml:"jobs"` +} + +type exportedJob struct { + Name string `yaml:"name"` + Tasks []exportedTask `yaml:"tasks"` + Environments []exportedEnvironment `yaml:"environments"` + // EnvironmentVariables carries env-var profiles via the common Jobs env-var API + // (API-293). Emitted only when the run declares env_variables/secrets; the task + // references a profile by key (see exportedTask.EnvironmentVariablesKey). + EnvironmentVariables []exportedEnvVarProfile `yaml:"environment_variables,omitempty"` +} + +type exportedTask struct { + TaskKey string `yaml:"task_key"` + EnvironmentKey string `yaml:"environment_key"` + // EnvironmentVariablesKey references a job-level environment_variables profile + // (common Jobs env-var API). Omitted when the run has no env vars/secrets so the + // wire form matches the submit path, which sets it only under the same gate. + EnvironmentVariablesKey string `yaml:"environment_variables_key,omitempty"` + MaxRetries int `yaml:"max_retries"` + TimeoutSeconds int `yaml:"timeout_seconds,omitempty"` + AiRuntimeTask exportedAiRuntimeTask `yaml:"ai_runtime_task"` +} + +// exportedEnvVarProfile is one entry in the job-level environment_variables list +// (common Jobs env-var API). variables holds plain values inline and secrets as +// {{secrets/scope/key}} references, resolved by Jobs at run time — the exact shape +// the ai_runtime_task submit path emits (jobs_api_client.py, API-293). +type exportedEnvVarProfile struct { + EnvironmentVariablesKey string `yaml:"environment_variables_key"` + Variables map[string]string `yaml:"variables"` +} + +type exportedAiRuntimeTask struct { + Experiment string `yaml:"experiment"` + MlflowRun string `yaml:"mlflow_run,omitempty"` + MlflowExperimentDirectory string `yaml:"mlflow_experiment_directory,omitempty"` + Deployments []exportedDeployment `yaml:"deployments"` +} + +type exportedDeployment struct { + Name string `yaml:"name"` + CommandPath string `yaml:"command_path"` + Compute exportedCompute `yaml:"compute"` +} + +type exportedCompute struct { + AcceleratorType string `yaml:"accelerator_type"` + AcceleratorCount int `yaml:"accelerator_count"` +} + +type exportedEnvironment struct { + EnvironmentKey string `yaml:"environment_key"` + Spec exportedEnvSpec `yaml:"spec"` +} + +// exportedEnvSpec carries the serverless runtime selection. The two fields are +// mutually exclusive (the Jobs serverless validator rejects setting both): a bare +// numeric channel uses environment_version; the databricks-ai managed environment +// (torch + ML venv preinstalled) uses base_environment. omitempty on both so only +// the resolved one is emitted. +type exportedEnvSpec struct { + EnvironmentVersion string `yaml:"environment_version,omitempty"` + BaseEnvironment string `yaml:"base_environment,omitempty"` +} + +// databricksAITokenPrefix marks an environment.version that selects the +// databricks-ai managed base environment rather than a bare channel. Mirrors +// _DATABRICKS_AI_TOKEN_PREFIX in the Python CLI's jobs_api_client.py. +const databricksAITokenPrefix = "databricks_ai_v" + +// databricksAIBaseEnvironment is the system base_environment id the databricks-ai +// token resolves to. Mirrors _DATABRICKS_AI_BASE_ENVIRONMENT in the Python CLI. +const databricksAIBaseEnvironment = "workspace-base-environments/" + +// convertToBundle maps a convertible runConfig to the emitted bundle. It assumes +// checkBundleConvertible has already passed. The command_path points at the synced +// command.sh under ${workspace.file_path}; the runtime channel is resolved the same +// way `air run` resolves it (config version, else the default), minus the process +// env lookup so a generated artifact is reproducible rather than host-dependent. +func convertToBundle(cfg *runConfig) *exportedBundle { + task := exportedAiRuntimeTask{ + Experiment: cfg.ExperimentName, + Deployments: []exportedDeployment{{ + // `air run` submits a single unnamed deployment; the bundle names it + // "worker" so the authored YAML reads clearly. The name is cosmetic (the + // submit payload carries none), so this does not change behavior. + Name: "worker", + CommandPath: workspaceFilePathRef + "/" + bundleCommandScript, + Compute: exportedCompute{ + AcceleratorType: cfg.Compute.AcceleratorType, + AcceleratorCount: cfg.Compute.NumAccelerators, + }, + }}, + } + if cfg.MLflowRunName != nil { + task.MlflowRun = *cfg.MLflowRunName + } + if cfg.MLflowExperimentDirectory != nil { + task.MlflowExperimentDirectory = *cfg.MLflowExperimentDirectory + } + + tsk := exportedTask{ + TaskKey: cfg.ExperimentName, + EnvironmentKey: aiRuntimeEnvironmentKey, + MaxRetries: cfg.maxRetries(), + TimeoutSeconds: cfg.timeoutSeconds(), + AiRuntimeTask: task, + } + + // Env vars ride the common Jobs env-var API: one job-level profile, referenced + // by the task's environment_variables_key. Emitted only when there are vars or + // secrets, so a var-free run's wire form matches the submit path exactly. + profiles := envVarProfiles(cfg) + if len(profiles) > 0 { + tsk.EnvironmentVariablesKey = aiRuntimeEnvVarsKey + } + + // Resource key shares the task_key charset, which validateExperimentName + // already guarantees, so the experiment name is safe to use verbatim. + return &exportedBundle{ + Bundle: bundleBlock{Name: cfg.ExperimentName}, + Resources: exportedResourcesBlock{ + Jobs: map[string]exportedJob{ + cfg.ExperimentName: { + Name: cfg.ExperimentName, + Tasks: []exportedTask{tsk}, + Environments: []exportedEnvironment{{ + EnvironmentKey: aiRuntimeEnvironmentKey, + Spec: exportBundleEnvSpec(cfg), + }}, + EnvironmentVariables: profiles, + }, + }, + }, + } +} + +// envVarProfiles builds the job-level env-var profile list from the run's +// env_variables and secrets, or nil when there are none. Plain values are inline; +// each secret (ENV_VAR -> "scope/key") becomes a {{secrets/scope/key}} reference +// Jobs resolves at run time. This mirrors the Python CLI's ai_runtime_task path +// (jobs_api_client.py, API-293) and was verified against staging: the profile +// persists on runs/get and secret refs resolve at run time. +func envVarProfiles(cfg *runConfig) []exportedEnvVarProfile { + if len(cfg.EnvVariables) == 0 && len(cfg.Secrets) == 0 { + return nil + } + variables := make(map[string]string, len(cfg.EnvVariables)+len(cfg.Secrets)) + for k, v := range cfg.EnvVariables { + variables[k] = v + } + for envVar, secretRef := range cfg.Secrets { + variables[envVar] = "{{secrets/" + secretRef + "}}" + } + return []exportedEnvVarProfile{{ + EnvironmentVariablesKey: aiRuntimeEnvVarsKey, + Variables: variables, + }} +} + +// exportBundleEnvSpec resolves the serverless runtime selection for the bundle, +// branching on environment.version the same way the Python CLI's jobs_api_client +// does (jobs_api_client.py:1562-1565): a "databricks_ai_v" token selects the +// managed databricks-ai base_environment (torch + ML venv), while a bare numeric +// channel ("4", "5", ...) uses environment_version. Unlike dlRuntimeImage it does +// not read process env, so the generated bundle is reproducible. A requirements-file +// dependency set carries its version in the file, so this falls back to the default +// channel there (same as the submit path). +// +// TODO(air): the Go `air run` submit path (runsubmit.go) only ever emits +// environment_version and cannot select base_environment, so a `version: +// databricks_ai_v` run lands on a bare GPU channel without torch/mlflow and +// fails at import. This converter hardcodes the correct base_environment branch so +// exported bundles run today; once `air run` forwards env vars/dependencies through +// the new BYOT common Jobs API (env_file/env_var work, ETA EOQ2), the submit path +// and this converter should share one runtime-selection helper instead of +// duplicating the branch. +func exportBundleEnvSpec(cfg *runConfig) exportedEnvSpec { + channel := strings.TrimPrefix(defaultDlRuntimeImage, "CLIENT-GPU-") + if v, ok := cfg.runtimeVersion(); ok { + channel = strings.TrimPrefix(v, "CLIENT-GPU-") + } + if strings.HasPrefix(channel, databricksAITokenPrefix) { + return exportedEnvSpec{BaseEnvironment: databricksAIBaseEnvironment + channel} + } + return exportedEnvSpec{EnvironmentVersion: channel} +} + +// marshalBundle renders the bundle to YAML with a header explaining provenance and +// the steps the user must complete before deploying (add a targets block; the code +// and command.sh are synced from the bundle folder). +func marshalBundle(b *exportedBundle, sourcePath string) ([]byte, error) { + body, err := yaml.Marshal(b) + if err != nil { + return nil, err + } + header := "# Generated by `air export-bundle` from " + filepath.Base(sourcePath) + ".\n" + + "#\n" + + "# Deploys the same workload as a durable Jobs resource. `bundle deploy` syncs this\n" + + "# folder (including " + bundleCommandScript + " and your code) to the workspace; the task's\n" + + "# command_path points at the synced " + bundleCommandScript + ". Before deploying, add a\n" + + "# `targets` block with your workspace host. See the ai-compute DABs examples.\n" + return append([]byte(header), body...), nil +} diff --git a/experimental/air/cmd/exportbundle_test.go b/experimental/air/cmd/exportbundle_test.go new file mode 100644 index 00000000000..e7964d4789d --- /dev/null +++ b/experimental/air/cmd/exportbundle_test.go @@ -0,0 +1,139 @@ +package aircmd + +import ( + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestConvertToBundleBasic(t *testing.T) { + cfg := &runConfig{ + ExperimentName: "exp", + Command: new("python train.py"), + Compute: &computeConfig{AcceleratorType: "GPU_8xH100", NumAccelerators: 16}, + MaxRetries: new(2), + TimeoutMinutes: new(30), + } + + b := convertToBundle(cfg) + + assert.Equal(t, "exp", b.Bundle.Name) + job, ok := b.Resources.Jobs["exp"] + require.True(t, ok, "job keyed by experiment name") + require.Len(t, job.Tasks, 1) + task := job.Tasks[0] + assert.Equal(t, "exp", task.TaskKey) + assert.Equal(t, aiRuntimeEnvironmentKey, task.EnvironmentKey) + assert.Equal(t, 2, task.MaxRetries) + assert.Equal(t, 1800, task.TimeoutSeconds) + + require.Len(t, task.AiRuntimeTask.Deployments, 1) + dep := task.AiRuntimeTask.Deployments[0] + assert.Equal(t, workspaceFilePathRef+"/"+bundleCommandScript, dep.CommandPath) + assert.Equal(t, "GPU_8xH100", dep.Compute.AcceleratorType) + assert.Equal(t, 16, dep.Compute.AcceleratorCount) + + // No env vars -> no profile and no task key, so the wire form matches a + // var-free submit exactly. + assert.Empty(t, job.EnvironmentVariables) + assert.Empty(t, task.EnvironmentVariablesKey) +} + +func TestConvertToBundleEnvVarsAndSecrets(t *testing.T) { + cfg := &runConfig{ + ExperimentName: "exp", + Command: new("python train.py"), + Compute: &computeConfig{AcceleratorType: "GPU_1xA10", NumAccelerators: 1}, + EnvVariables: map[string]string{"FOO": "bar", "LOG_LEVEL": "INFO"}, + Secrets: map[string]string{"TOKEN": "myscope/mykey"}, + } + + b := convertToBundle(cfg) + job := b.Resources.Jobs["exp"] + + // The task references the single profile by key. + assert.Equal(t, aiRuntimeEnvVarsKey, job.Tasks[0].EnvironmentVariablesKey) + + // One profile, keyed to match, carrying plain values inline and the secret as a + // {{secrets/scope/key}} reference (common Jobs env-var API; resolved by Jobs at + // run time). Verified against staging: this shape persists on runs/get. + require.Len(t, job.EnvironmentVariables, 1) + prof := job.EnvironmentVariables[0] + assert.Equal(t, aiRuntimeEnvVarsKey, prof.EnvironmentVariablesKey) + assert.Equal(t, "bar", prof.Variables["FOO"]) + assert.Equal(t, "INFO", prof.Variables["LOG_LEVEL"]) + assert.Equal(t, "{{secrets/myscope/mykey}}", prof.Variables["TOKEN"]) +} + +func TestCheckBundleConvertibleAllowsEnvVars(t *testing.T) { + // env_variables and secrets are now representable via the common Jobs env-var + // API, so the gate must NOT reject them (regression guard for the lift). + cfg := &runConfig{ + ExperimentName: "exp", + Command: new("python train.py"), + Compute: &computeConfig{AcceleratorType: "GPU_1xA10", NumAccelerators: 1}, + EnvVariables: map[string]string{"FOO": "bar"}, + Secrets: map[string]string{"TOKEN": "s/k"}, + } + assert.NoError(t, checkBundleConvertible(cfg)) +} + +func TestCheckBundleConvertibleRejectsCodeSourcePathCommand(t *testing.T) { + // A command that reads $CODE_SOURCE_PATH assumes the air run harness a bundle + // doesn't provide; the gate must still reject it with an actionable message. + cfg := &runConfig{ + ExperimentName: "exp", + Command: new("cd $CODE_SOURCE_PATH && python train.py"), + Compute: &computeConfig{AcceleratorType: "GPU_1xA10", NumAccelerators: 1}, + } + err := checkBundleConvertible(cfg) + require.Error(t, err) + assert.Contains(t, err.Error(), codeSourcePathVar) +} + +func TestRenderBundleIncludesTargetsAndConvertGate(t *testing.T) { + // renderBundle (what --dry-run --via-dabs shows and what the run path deploys) + // must include the converted job AND the appended dev targets block. + cfg := &runConfig{ + ExperimentName: "exp", + Command: new("python train.py"), + Compute: &computeConfig{AcceleratorType: "GPU_1xA10", NumAccelerators: 1}, + EnvVariables: map[string]string{"FOO": "bar"}, + } + out, err := renderBundle(cfg, "train.yaml") + require.NoError(t, err) + assert.Contains(t, out, "ai_runtime_task:") + assert.Contains(t, out, "FOO: bar") + assert.Contains(t, out, "targets:") + assert.Contains(t, out, "mode: development") + + // The convertibility gate still applies: an unconvertible config errors instead + // of rendering a lossy bundle. + bad := &runConfig{ + ExperimentName: "exp", + Command: new("cd $CODE_SOURCE_PATH && python train.py"), + Compute: &computeConfig{AcceleratorType: "GPU_1xA10", NumAccelerators: 1}, + } + _, err = renderBundle(bad, "train.yaml") + require.Error(t, err) +} + +func TestMarshalBundleEnvVarsRoundTrip(t *testing.T) { + // The emitted YAML must carry the env-var profile so `bundle deploy` sends it. + cfg := &runConfig{ + ExperimentName: "exp", + Command: new("x"), + Compute: &computeConfig{AcceleratorType: "GPU_1xA10", NumAccelerators: 1}, + EnvVariables: map[string]string{"FOO": "bar"}, + } + out, err := marshalBundle(convertToBundle(cfg), "train.yaml") + require.NoError(t, err) + body := string(out) + assert.Contains(t, body, "environment_variables_key: default") + assert.Contains(t, body, "environment_variables:") + assert.Contains(t, body, "FOO: bar") + // Header provenance line is present. + assert.True(t, strings.Contains(body, "Generated by `air export-bundle`")) +} diff --git a/experimental/air/cmd/get.go b/experimental/air/cmd/get.go index 2f8bd8dd09e..7f94cfbdbd9 100644 --- a/experimental/air/cmd/get.go +++ b/experimental/air/cmd/get.go @@ -91,6 +91,9 @@ func authError(ctx context.Context, cmd *cobra.Command, err error) error { // newGetCommand returns the `air get JOB_RUN_ID` command, which shows status, // configuration, and timing details for a specific run. +// get works unchanged for DABs-submitted runs: it resolves by run_id via +// Jobs.GetRun, and a `bundle run` produces a normal Jobs run with a run_id. No +// submit-verb-specific handling is needed. func newGetCommand() *cobra.Command { cmd := &cobra.Command{ Use: "get JOB_RUN_ID", diff --git a/experimental/air/cmd/list.go b/experimental/air/cmd/list.go index 7d5c703a39c..b963ce352b0 100644 --- a/experimental/air/cmd/list.go +++ b/experimental/air/cmd/list.go @@ -68,6 +68,9 @@ type listQuery struct { filters listFilters fetchMLflow bool limit int + // includeDABs widens the jobs scan to include DABs-submitted AIR runs (JOB_RUN), + // not just ephemeral SUBMIT_RUNs. Set by --via-dabs / AIR_VIA_DABS. + includeDABs bool } func newListCommand() *cobra.Command { @@ -76,6 +79,7 @@ func newListCommand() *cobra.Command { allStatus bool allUsers bool filters []string + viaDABS bool ) cmd := &cobra.Command{ @@ -90,6 +94,9 @@ func newListCommand() *cobra.Command { cmd.Flags().BoolVar(&allStatus, "all-status", false, "Show runs in all states (default: active only)") cmd.Flags().BoolVar(&allUsers, "all-users", false, "Show runs from all users") cmd.Flags().StringArrayVar(&filters, "filter", nil, "Filter runs, e.g. experiment=foo* (repeatable)") + // [experimental] Also include DABs-submitted AIR runs (persistent JOB_RUNs), not + // just ephemeral runs/submit. Honors AIR_VIA_DABS=1. + cmd.Flags().BoolVar(&viaDABS, "via-dabs", false, "[experimental] Include DABs-submitted (persistent job) AIR runs") cmd.RunE = func(cmd *cobra.Command, args []string) error { ctx := cmd.Context() @@ -126,6 +133,7 @@ func newListCommand() *cobra.Command { filters: f, fetchMLflow: root.OutputType(cmd) == flags.OutputText, limit: limit, + includeDABs: resolveViaDABS(ctx, viaDABS), }) // JSON prints the newest `limit` runs once. Text renders the table: @@ -239,6 +247,13 @@ func newJobsScanStrategy(ctx context.Context, w *databricks.WorkspaceClient, q l Limit: jobsPageLimit, ActiveOnly: q.activeOnly, } + // DABs-submitted AIR runs are runs OF a persistent job (RunType JOB_RUN), not + // SUBMIT_RUN, so the SUBMIT_RUN server filter would hide them. When listing the + // DABs path, drop the run-type restriction and let isAirRun (task-shape) select; + // this unions ephemeral SUBMIT_RUNs and DABs JOB_RUNs during migration. + if q.includeDABs { + req.RunType = "" + } return &jobsScanStrategy{ ctx: ctx, w: w, diff --git a/experimental/air/cmd/list_test.go b/experimental/air/cmd/list_test.go index f70330240e5..7ce35a27ab8 100644 --- a/experimental/air/cmd/list_test.go +++ b/experimental/air/cmd/list_test.go @@ -99,6 +99,50 @@ func TestListAirRunsExperimentFilter(t *testing.T) { assert.Equal(t, "1", rows[0].RunID) } +// queryCapturingServer records the run_type query param on the runs/list request, +// so a test can assert how the jobs scan narrows (or doesn't) server-side. +func queryCapturingServer(t *testing.T, gotRunType *string, body string) *httptest.Server { + t.Helper() + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == "/api/2.2/jobs/runs/list" { + *gotRunType = r.URL.Query().Get("run_type") + _, _ = w.Write([]byte(body)) + return + } + _, _ = w.Write([]byte(`{}`)) + })) + t.Cleanup(srv.Close) + return srv +} + +func TestListDefaultScanRestrictsToSubmitRun(t *testing.T) { + // Without --via-dabs the jobs scan restricts to SUBMIT_RUN server-side (the + // ephemeral runs air submits today). + var runType string + srv := queryCapturingServer(t, &runType, runsListBody(t, "", airBaseRun(1, "me@example.com", "GPU_1xA10", 1, "exp"))) + _, err := newRunFetcher(t.Context(), newTestWorkspaceClient(t, srv.URL), listQuery{ + activeOnly: true, userFilter: "me@example.com", + }).next(10) + require.NoError(t, err) + assert.Equal(t, "SUBMIT_RUN", runType) +} + +func TestListViaDABSIncludesJobRuns(t *testing.T) { + // --via-dabs drops the SUBMIT_RUN restriction so DABs-submitted AIR runs + // (RunType JOB_RUN, i.e. runs of a persistent job) are not filtered out + // server-side; isAirRun still selects by task shape. + var runType string + dabsRun := airBaseRun(9, "me@example.com", "GPU_1xA10", 1, "exp") + srv := queryCapturingServer(t, &runType, runsListBody(t, "", dabsRun)) + rows, err := newRunFetcher(t.Context(), newTestWorkspaceClient(t, srv.URL), listQuery{ + activeOnly: true, userFilter: "me@example.com", includeDABs: true, + }).next(10) + require.NoError(t, err) + assert.Empty(t, runType, "run_type must be unset so JOB_RUN DABs runs are included") + require.Len(t, rows, 1) + assert.Equal(t, "9", rows[0].RunID) +} + func TestListAirRunsLimitTruncates(t *testing.T) { runs := []jobs.BaseRun{ airBaseRun(1, "me@example.com", "GPU_1xH100", 1, "exp-a"), diff --git a/experimental/air/cmd/run.go b/experimental/air/cmd/run.go index bd32810e9bc..2e3b5baba80 100644 --- a/experimental/air/cmd/run.go +++ b/experimental/air/cmd/run.go @@ -18,6 +18,9 @@ type runResult struct { DryRun bool `json:"dry_run,omitempty"` RunID string `json:"run_id,omitempty"` DashboardURL string `json:"dashboard_url,omitempty"` + // Bundle is the generated databricks.yml, included in a --dry-run --via-dabs + // so the user can see exactly what would be deployed on their behalf. + Bundle string `json:"bundle,omitempty"` } func newRunCommand() *cobra.Command { @@ -27,6 +30,7 @@ func newRunCommand() *cobra.Command { overrides []string dryRun bool idempotencyKey string + viaDABS bool ) cmd := &cobra.Command{ @@ -43,6 +47,10 @@ The workload is described by a YAML config file (see --file).`, cmd.Flags().StringArrayVar(&overrides, "override", nil, "Override a YAML field, e.g. compute.num_accelerators=8 (repeatable)") cmd.Flags().BoolVar(&dryRun, "dry-run", false, "Validate the config without submitting") cmd.Flags().StringVar(&idempotencyKey, "idempotency-key", "", "Return the existing run if this key was already used") + // [experimental] Submit as a Databricks Asset Bundle (a persistent job: + // convert -> deploy -> run) instead of an ephemeral run. Default off; the + // ephemeral runs/submit path is unchanged. Also honors AIR_VIA_DABS=1. + cmd.Flags().BoolVar(&viaDABS, "via-dabs", false, "[experimental] Submit via a DABs bundle (persistent job) instead of an ephemeral run") _ = cmd.MarkFlagRequired("file") // --dry-run only validates the config locally, so it needs no workspace. @@ -72,6 +80,20 @@ The workload is described by a YAML config file (see --file).`, } if dryRun { + // On the DABs path, a dry run shows the generated bundle so the user can + // see the artifact we'd deploy on their behalf (transparency); it does not + // deploy. Falls back to the plain "config is valid" message otherwise. + if resolveViaDABS(ctx, viaDABS) { + bundleYAML, err := renderBundle(cfg, file) + if err != nil { + return err + } + if root.OutputType(cmd) == flags.OutputText { + cmdio.LogString(ctx, fmt.Sprintf("Dry run: %q would deploy as this bundle (not deploying):\n\n%s", cfg.ExperimentName, bundleYAML)) + return nil + } + return renderEnvelope(ctx, runResult{Status: "DRY_RUN_OK", DryRun: true, Bundle: bundleYAML}) + } if root.OutputType(cmd) == flags.OutputText { cmdio.LogString(ctx, fmt.Sprintf("Dry run: configuration for %q is valid; not submitting.", cfg.ExperimentName)) return nil @@ -80,7 +102,15 @@ The workload is described by a YAML config file (see --file).`, } w := cmdctx.WorkspaceClient(ctx) - runID, dashboardURL, err := submitWorkload(ctx, w, cfg, file, idempotencyKey) + + // submitWorkload (ephemeral runs/submit) and submitViaDABS (persistent + // bundle job) share a signature, so the DABs gate is a call-site swap; the + // whole downstream (envelope, dashboard URL) is reused unchanged. + submit := submitWorkload + if resolveViaDABS(ctx, viaDABS) { + submit = submitViaDABS + } + runID, dashboardURL, err := submit(ctx, w, cfg, file, idempotencyKey) if err != nil { return err } diff --git a/experimental/air/cmd/rundabs.go b/experimental/air/cmd/rundabs.go new file mode 100644 index 00000000000..976a72a0488 --- /dev/null +++ b/experimental/air/cmd/rundabs.go @@ -0,0 +1,199 @@ +package aircmd + +// PROTOTYPE — Option D: `air run --via-dabs` submits the workload as a Databricks +// Asset Bundle (a persistent Jobs resource: deploy then run) instead of the +// ephemeral runs/submit path in runsubmit.go. +// +// Design (per the "AIR CLI + DABs" design meeting, Jul 14): `air run` becomes a +// high-level caller of [convert -> bundle deploy -> bundle run]. This reuses the +// exportbundle.go converter (train.yaml -> databricks.yml) verbatim, so the +// generated ai_runtime_task — including env vars via the common Jobs env-var API — +// is exactly what `air export-bundle` produces and what `air run` would submit. +// +// It is gated behind --via-dabs (and AIR_VIA_DABS=1); the ephemeral path stays the +// default and untouched. This is additive. +// +// PROTOTYPE SCOPE / shell-out choice: deploy+run are driven by shelling out to the +// `databricks bundle` CLI, which is the literal "wrapper around DABs" the meeting +// described and keeps this diff legible. A production version would instead call +// cmd/bundle/utils.ProcessBundle(...) in-process (see TODO(prototype) below) to +// avoid the child-process dependency. The behavior (persistent job, deploy +// latency, no ephemeral-job GC) is identical either way — the point being to +// exercise the real DABs path end to end. +// +// KNOWN GAP (called out in the design doc): the deployed job is PERSISTENT and is +// NOT swept by the Jobs ephemeral-job GC (JobsSoftDeletion only sweeps +// EPHEMERAL/WORKFLOW types). Every `air run --via-dabs` leaves a job behind. A real +// impl must choose a reuse/cleanup policy (stable per-experiment bundle so re-runs +// update one job, or `bundle destroy` after run). The prototype leaves the job and +// logs the tradeoff. + +import ( + "context" + "fmt" + "os" + "os/exec" + "path/filepath" + + "github.com/databricks/cli/libs/cmdio" + "github.com/databricks/cli/libs/env" + "github.com/databricks/databricks-sdk-go" +) + +// viaDABSEnv opts into the DABs submit path without the flag (flag wins). +const viaDABSEnv = "AIR_VIA_DABS" + +// dabsTarget is the bundle target the prototype deploys/runs against. A real impl +// would let the user pick (or default per-workspace); "dev" gets DABs dev-mode +// (per-user name prefix + isolation), which suits AIR's per-user runs. +const dabsTarget = "dev" + +// submitViaDABS is the Option-D analogue of submitWorkload (runsubmit.go), with the +// identical signature so run.go can pick either at the call site. It converts the +// run config to a bundle, writes the project (databricks.yml + command.sh + code), +// then deploys and runs it. Returns the run_id + dashboard URL. +func submitViaDABS(ctx context.Context, w *databricks.WorkspaceClient, cfg *runConfig, configPath, idempotencyKey string) (int64, string, error) { + // The converter's gate rejects configs a bundle can't represent faithfully + // (e.g. git-pinned code_source, $CODE_SOURCE_PATH commands). Fail fast with its + // actionable message before we write or deploy anything. + if err := checkBundleConvertible(cfg); err != nil { + return 0, "", err + } + + // Idempotency doesn't map onto deploy+run (two calls, not one tokened submit); + // note rather than silently honor it. A real impl dedupes on bundle/job identity. + if idempotencyKey != "" || cfg.IdempotencyToken != nil { + cmdio.LogString(ctx, "note: --idempotency-key is ignored on the DABs path (deploy+run is not a single idempotent call)") + } + + root, cleanup, err := writeBundleProject(cfg, configPath) + if err != nil { + return 0, "", err + } + defer cleanup() + + // Surface the generated artifact's path (transparency): the user can inspect + // what we deploy on their behalf. The dir is temporary (auto-removed on exit) — + // labeled as such so we don't imply a durable, user-managed location. For the + // full contents without deploying, use `air run --dry-run --via-dabs`. + cmdio.LogString(ctx, "Generated bundle (temporary): "+filepath.Join(root, "databricks.yml")) + + // deploy: creates/updates the persistent job. This is the step that adds + // latency vs the ephemeral submit (measured ~2-3s warm, ~12s first-of-session). + cmdio.LogString(ctx, "Deploying bundle (creates a persistent job)...") + if err := runBundle(ctx, w, root, "deploy"); err != nil { + return 0, "", fmt.Errorf("bundle deploy: %w", err) + } + + // run: RunNow on the deployed job. --no-wait returns as soon as the run starts, + // matching the ephemeral path's fire-and-return (we don't block on completion). + cmdio.LogString(ctx, "Triggering run...") + if err := runBundle(ctx, w, root, "run", cfg.ExperimentName, "--no-wait"); err != nil { + return 0, "", fmt.Errorf("bundle run: %w", err) + } + + cmdio.LogString(ctx, "note: this run created a PERSISTENT job (not auto-GC'd like an ephemeral run). Use `air list` to find it.") + + // TODO(prototype): parse run_id + dashboard URL from `databricks bundle run + // --output json`. For now return 0 and the jobs page; the run is live either way. + return 0, workspaceJobsURL(w), nil +} + +// renderBundle produces the exact databricks.yml the run path would deploy: the +// exportbundle.go converter output plus the dev targets block the run path appends. +// It touches no filesystem and does not deploy, so `air run --dry-run --via-dabs` +// can show the user the artifact we'd generate on their behalf (transparency), and +// writeBundleProject reuses it so preview and real run never diverge. +func renderBundle(cfg *runConfig, configPath string) (string, error) { + if err := checkBundleConvertible(cfg); err != nil { + return "", err + } + body, err := marshalBundle(convertToBundle(cfg), configPath) + if err != nil { + return "", err + } + // The converter's emitted YAML has no targets block (export-bundle leaves that + // to the user); the run path needs one so deploy is non-interactive. Append a + // dev target — the host is resolved from the CLI profile at deploy time. + return string(body) + bundleTargetsBlock(), nil +} + +// writeBundleProject renders databricks.yml (via renderBundle) plus the command.sh +// and code the bundle syncs, into a temp bundle root. Returns the root and a +// cleanup func. +// +// TODO(prototype): reuse the ephemeral path's snapshot/upload staging so the exact +// code tree + command.sh land in the bundle folder. This sketch writes command.sh +// from cfg.Command; wiring code_source snapshotting through the inherited +// bundle/config/mutator/aicode mutator is the remaining piece. +func writeBundleProject(cfg *runConfig, configPath string) (string, func(), error) { + body, err := renderBundle(cfg, configPath) + if err != nil { + return "", func() {}, err + } + root, err := os.MkdirTemp("", "air-dabs-*") + if err != nil { + return "", func() {}, err + } + cleanup := func() { _ = os.RemoveAll(root) } + + if err := os.WriteFile(filepath.Join(root, "databricks.yml"), []byte(body), 0o600); err != nil { + cleanup() + return "", func() {}, err + } + + // command.sh: the entrypoint the task's command_path points at (synced by deploy). + if cfg.Command != nil { + if err := os.WriteFile(filepath.Join(root, bundleCommandScript), []byte(*cfg.Command), 0o600); err != nil { + cleanup() + return "", func() {}, err + } + } + return root, cleanup, nil +} + +// bundleTargetsBlock is the minimal dev target appended to the converted YAML so +// `bundle deploy` runs non-interactively. dev mode gives per-user name prefix + +// isolation. The workspace host comes from the active CLI profile, so it's omitted +// here (the bundle resolves it from auth at deploy time). +func bundleTargetsBlock() string { + // The comment supersedes the export-bundle header's "add a targets block" + // note: on the run path we append this dev target for you (host resolved from + // your CLI profile at deploy time), so `air run --via-dabs` needs no manual edit. + return "\n# Appended by `air run --via-dabs` (host resolved from your CLI profile at deploy time):\ntargets:\n " + dabsTarget + ":\n mode: development\n default: true\n" +} + +// runBundle shells out to `databricks bundle -t ` in the bundle +// root. Uses the direct engine (no Terraform / no registry.terraform.io dep — GA +// default on new CLIs, set explicitly for older ones). Auth is inherited from the +// same profile/env the parent air command resolved. +// +// TODO(prototype): replace with an in-process cmd/bundle/utils.ProcessBundle call +// so there's no dependency on a `databricks` binary on PATH. +func runBundle(ctx context.Context, w *databricks.WorkspaceClient, root string, args ...string) error { + full := append([]string{"bundle"}, args...) + full = append(full, "-t", dabsTarget) + cmd := exec.CommandContext(ctx, "databricks", full...) + cmd.Dir = root + cmd.Env = append(os.Environ(), + "DATABRICKS_BUNDLE_ENGINE=direct", + "DATABRICKS_HOST="+w.Config.Host, + ) + cmd.Stdout = os.Stderr // keep our stdout clean for the JSON envelope + cmd.Stderr = os.Stderr + return cmd.Run() +} + +func workspaceJobsURL(w *databricks.WorkspaceClient) string { + host := w.Config.Host + for len(host) > 0 && host[len(host)-1] == '/' { + host = host[:len(host)-1] + } + return host + "/jobs" +} + +// resolveViaDABS decides whether to take the DABs submit path: the --via-dabs flag +// or AIR_VIA_DABS=1. Kept here so run.go's wiring stays a one-liner. +func resolveViaDABS(ctx context.Context, flag bool) bool { + return flag || env.Get(ctx, viaDABSEnv) == "1" +} From a12413efcf94b0dd53223e03878fe7b598e15b7f Mon Sep 17 00:00:00 2001 From: vinchenzo-db Date: Wed, 15 Jul 2026 20:35:12 +0000 Subject: [PATCH 02/11] [air] Make DABs the default (and only) submit path for air run Rewrites `air run` to submit exclusively via a Databricks Asset Bundle (convert train.yaml -> databricks.yml -> deploy -> run), removing the ephemeral runs/submit path and the experimental --via-dabs flag. This honors the AIR-CLI/DABs agreement: DECO requires AIR to integrate with DABs and submit through the public Jobs API, not an internal proxy. Deploy and run are driven IN-PROCESS via the bundle libraries (cmd/bundle/utils.ProcessBundle + Jobs.RunNow), the same entry points the `databricks bundle` commands use and the pattern the pipelines CLI established (cmd/pipelines). There is no child `databricks` process, so the deploy always uses this build's ai_runtime_task-aware bundle schema. - rundabs.go: submitWorkload converts the config to a bundle in a temp root, deploys it with ProcessBundle (rooted via DATABRICKS_BUNDLE_ROOT on the context, direct engine), then RunNow on the resolved job to return the run_id (fire-and-return, no wait). A synthetic carrier command supplies the flags ProcessBundle reads and forwards air's resolved auth profile. - run.go: DABs is the only path; --via-dabs removed. --dry-run always renders the generated bundle for transparency. - list.go: the jobs scan no longer restricts to SUBMIT_RUN, since AIR runs are now runs of a persistent DABs job (JOB_RUN); --via-dabs gate removed. isAirRun still selects by task shape. - runsubmit.go: ephemeral submit path retired, leaving only the runtime- image defaults the converter shares. Verified end-to-end on staging (e2-dogfood): run -> get -> list -> cancel all work against a real in-process deploy+run. go build + go test ./experimental/air/... pass. Co-authored-by: Isaac --- experimental/air/cmd/list.go | 21 +- experimental/air/cmd/list_test.go | 22 +-- experimental/air/cmd/run.go | 44 ++--- experimental/air/cmd/rundabs.go | 264 ++++++++++++++++--------- experimental/air/cmd/runsubmit.go | 197 ++---------------- experimental/air/cmd/runsubmit_test.go | 216 -------------------- 6 files changed, 200 insertions(+), 564 deletions(-) delete mode 100644 experimental/air/cmd/runsubmit_test.go diff --git a/experimental/air/cmd/list.go b/experimental/air/cmd/list.go index b963ce352b0..2a64f4118bb 100644 --- a/experimental/air/cmd/list.go +++ b/experimental/air/cmd/list.go @@ -68,9 +68,6 @@ type listQuery struct { filters listFilters fetchMLflow bool limit int - // includeDABs widens the jobs scan to include DABs-submitted AIR runs (JOB_RUN), - // not just ephemeral SUBMIT_RUNs. Set by --via-dabs / AIR_VIA_DABS. - includeDABs bool } func newListCommand() *cobra.Command { @@ -79,7 +76,6 @@ func newListCommand() *cobra.Command { allStatus bool allUsers bool filters []string - viaDABS bool ) cmd := &cobra.Command{ @@ -94,9 +90,6 @@ func newListCommand() *cobra.Command { cmd.Flags().BoolVar(&allStatus, "all-status", false, "Show runs in all states (default: active only)") cmd.Flags().BoolVar(&allUsers, "all-users", false, "Show runs from all users") cmd.Flags().StringArrayVar(&filters, "filter", nil, "Filter runs, e.g. experiment=foo* (repeatable)") - // [experimental] Also include DABs-submitted AIR runs (persistent JOB_RUNs), not - // just ephemeral runs/submit. Honors AIR_VIA_DABS=1. - cmd.Flags().BoolVar(&viaDABS, "via-dabs", false, "[experimental] Include DABs-submitted (persistent job) AIR runs") cmd.RunE = func(cmd *cobra.Command, args []string) error { ctx := cmd.Context() @@ -133,7 +126,6 @@ func newListCommand() *cobra.Command { filters: f, fetchMLflow: root.OutputType(cmd) == flags.OutputText, limit: limit, - includeDABs: resolveViaDABS(ctx, viaDABS), }) // JSON prints the newest `limit` runs once. Text renders the table: @@ -241,19 +233,16 @@ type jobsScanStrategy struct { } func newJobsScanStrategy(ctx context.Context, w *databricks.WorkspaceClient, q listQuery) *jobsScanStrategy { + // AIR runs are now runs OF a persistent DABs job (RunType JOB_RUN), not the + // ephemeral SUBMIT_RUN the retired runs/submit path produced. A SUBMIT_RUN + // server filter would hide them, so we leave RunType unset and let isAirRun + // (task-shape) select AIR runs. Leaving it unset also still surfaces any + // pre-migration ephemeral SUBMIT_RUNs that remain in the workspace. req := jobs.ListRunsRequest{ - RunType: jobs.RunTypeSubmitRun, ExpandTasks: true, Limit: jobsPageLimit, ActiveOnly: q.activeOnly, } - // DABs-submitted AIR runs are runs OF a persistent job (RunType JOB_RUN), not - // SUBMIT_RUN, so the SUBMIT_RUN server filter would hide them. When listing the - // DABs path, drop the run-type restriction and let isAirRun (task-shape) select; - // this unions ephemeral SUBMIT_RUNs and DABs JOB_RUNs during migration. - if q.includeDABs { - req.RunType = "" - } return &jobsScanStrategy{ ctx: ctx, w: w, diff --git a/experimental/air/cmd/list_test.go b/experimental/air/cmd/list_test.go index 7ce35a27ab8..7bd2bbc8d69 100644 --- a/experimental/air/cmd/list_test.go +++ b/experimental/air/cmd/list_test.go @@ -115,27 +115,15 @@ func queryCapturingServer(t *testing.T, gotRunType *string, body string) *httpte return srv } -func TestListDefaultScanRestrictsToSubmitRun(t *testing.T) { - // Without --via-dabs the jobs scan restricts to SUBMIT_RUN server-side (the - // ephemeral runs air submits today). - var runType string - srv := queryCapturingServer(t, &runType, runsListBody(t, "", airBaseRun(1, "me@example.com", "GPU_1xA10", 1, "exp"))) - _, err := newRunFetcher(t.Context(), newTestWorkspaceClient(t, srv.URL), listQuery{ - activeOnly: true, userFilter: "me@example.com", - }).next(10) - require.NoError(t, err) - assert.Equal(t, "SUBMIT_RUN", runType) -} - -func TestListViaDABSIncludesJobRuns(t *testing.T) { - // --via-dabs drops the SUBMIT_RUN restriction so DABs-submitted AIR runs - // (RunType JOB_RUN, i.e. runs of a persistent job) are not filtered out - // server-side; isAirRun still selects by task shape. +func TestListScanIncludesJobRuns(t *testing.T) { + // AIR runs are now runs of a persistent DABs job (RunType JOB_RUN), so the jobs + // scan must NOT restrict to SUBMIT_RUN server-side — that would hide them. + // run_type is left unset and isAirRun selects AIR runs by task shape. var runType string dabsRun := airBaseRun(9, "me@example.com", "GPU_1xA10", 1, "exp") srv := queryCapturingServer(t, &runType, runsListBody(t, "", dabsRun)) rows, err := newRunFetcher(t.Context(), newTestWorkspaceClient(t, srv.URL), listQuery{ - activeOnly: true, userFilter: "me@example.com", includeDABs: true, + activeOnly: true, userFilter: "me@example.com", }).next(10) require.NoError(t, err) assert.Empty(t, runType, "run_type must be unset so JOB_RUN DABs runs are included") diff --git a/experimental/air/cmd/run.go b/experimental/air/cmd/run.go index 2e3b5baba80..def02a785d9 100644 --- a/experimental/air/cmd/run.go +++ b/experimental/air/cmd/run.go @@ -18,8 +18,8 @@ type runResult struct { DryRun bool `json:"dry_run,omitempty"` RunID string `json:"run_id,omitempty"` DashboardURL string `json:"dashboard_url,omitempty"` - // Bundle is the generated databricks.yml, included in a --dry-run --via-dabs - // so the user can see exactly what would be deployed on their behalf. + // Bundle is the generated databricks.yml, included in a --dry-run so the user + // can see exactly what would be deployed on their behalf. Bundle string `json:"bundle,omitempty"` } @@ -30,7 +30,6 @@ func newRunCommand() *cobra.Command { overrides []string dryRun bool idempotencyKey string - viaDABS bool ) cmd := &cobra.Command{ @@ -45,12 +44,8 @@ The workload is described by a YAML config file (see --file).`, cmd.Flags().StringVarP(&file, "file", "f", "", "Path to the workload YAML config") cmd.Flags().BoolVar(&watch, "watch", false, "Stream logs until the run completes") cmd.Flags().StringArrayVar(&overrides, "override", nil, "Override a YAML field, e.g. compute.num_accelerators=8 (repeatable)") - cmd.Flags().BoolVar(&dryRun, "dry-run", false, "Validate the config without submitting") + cmd.Flags().BoolVar(&dryRun, "dry-run", false, "Validate the config and show the generated bundle without submitting") cmd.Flags().StringVar(&idempotencyKey, "idempotency-key", "", "Return the existing run if this key was already used") - // [experimental] Submit as a Databricks Asset Bundle (a persistent job: - // convert -> deploy -> run) instead of an ephemeral run. Default off; the - // ephemeral runs/submit path is unchanged. Also honors AIR_VIA_DABS=1. - cmd.Flags().BoolVar(&viaDABS, "via-dabs", false, "[experimental] Submit via a DABs bundle (persistent job) instead of an ephemeral run") _ = cmd.MarkFlagRequired("file") // --dry-run only validates the config locally, so it needs no workspace. @@ -80,37 +75,24 @@ The workload is described by a YAML config file (see --file).`, } if dryRun { - // On the DABs path, a dry run shows the generated bundle so the user can - // see the artifact we'd deploy on their behalf (transparency); it does not - // deploy. Falls back to the plain "config is valid" message otherwise. - if resolveViaDABS(ctx, viaDABS) { - bundleYAML, err := renderBundle(cfg, file) - if err != nil { - return err - } - if root.OutputType(cmd) == flags.OutputText { - cmdio.LogString(ctx, fmt.Sprintf("Dry run: %q would deploy as this bundle (not deploying):\n\n%s", cfg.ExperimentName, bundleYAML)) - return nil - } - return renderEnvelope(ctx, runResult{Status: "DRY_RUN_OK", DryRun: true, Bundle: bundleYAML}) + // A dry run shows the generated bundle so the user can see the artifact + // we'd deploy on their behalf (transparency); it does not deploy. + bundleYAML, err := renderBundle(cfg, file) + if err != nil { + return err } if root.OutputType(cmd) == flags.OutputText { - cmdio.LogString(ctx, fmt.Sprintf("Dry run: configuration for %q is valid; not submitting.", cfg.ExperimentName)) + cmdio.LogString(ctx, fmt.Sprintf("Dry run: %q would deploy as this bundle (not deploying):\n\n%s", cfg.ExperimentName, bundleYAML)) return nil } - return renderEnvelope(ctx, runResult{Status: "DRY_RUN_OK", DryRun: true}) + return renderEnvelope(ctx, runResult{Status: "DRY_RUN_OK", DryRun: true, Bundle: bundleYAML}) } w := cmdctx.WorkspaceClient(ctx) - // submitWorkload (ephemeral runs/submit) and submitViaDABS (persistent - // bundle job) share a signature, so the DABs gate is a call-site swap; the - // whole downstream (envelope, dashboard URL) is reused unchanged. - submit := submitWorkload - if resolveViaDABS(ctx, viaDABS) { - submit = submitViaDABS - } - runID, dashboardURL, err := submit(ctx, w, cfg, file, idempotencyKey) + // submitWorkload converts the config to a bundle, deploys it, and triggers a + // run in-process (see rundabs.go). It is the only submit path. + runID, dashboardURL, err := submitWorkload(ctx, w, cfg, file, idempotencyKey) if err != nil { return err } diff --git a/experimental/air/cmd/rundabs.go b/experimental/air/cmd/rundabs.go index 976a72a0488..030dc9c7802 100644 --- a/experimental/air/cmd/rundabs.go +++ b/experimental/air/cmd/rundabs.go @@ -1,108 +1,114 @@ package aircmd -// PROTOTYPE — Option D: `air run --via-dabs` submits the workload as a Databricks -// Asset Bundle (a persistent Jobs resource: deploy then run) instead of the -// ephemeral runs/submit path in runsubmit.go. +// air run submits a training workload by converting the run YAML to a Databricks +// Asset Bundle and driving `bundle deploy` + `bundle run` in-process. This is the +// ONLY submit path: per the AIR-CLI/DABs agreement with Pieter Noordhuis (DECO +// requires AIR to integrate with DABs and submit through the public Jobs API, not +// an internal proxy), `air run` is a high-level wrapper of +// [convert train.yaml -> databricks.yml -> deploy -> run]. // -// Design (per the "AIR CLI + DABs" design meeting, Jul 14): `air run` becomes a -// high-level caller of [convert -> bundle deploy -> bundle run]. This reuses the -// exportbundle.go converter (train.yaml -> databricks.yml) verbatim, so the -// generated ai_runtime_task — including env vars via the common Jobs env-var API — -// is exactly what `air export-bundle` produces and what `air run` would submit. +// The bundle is a persistent Jobs resource, so `air list` finds it via the indexed +// jobs-list path and re-runs anchor to a standing job. It reuses the exportbundle.go +// converter verbatim, so the deployed ai_runtime_task — including env vars via the +// common Jobs env-var API — is exactly what `air export-bundle` emits. // -// It is gated behind --via-dabs (and AIR_VIA_DABS=1); the ephemeral path stays the -// default and untouched. This is additive. +// IN-PROCESS, NOT SHELL-OUT: deploy and run call the bundle libraries directly +// (cmd/bundle/utils.ProcessBundle + bundle/run), the same entry points the +// `databricks bundle` commands use and the pattern the pipelines CLI established +// (cmd/pipelines). There is no child `databricks` process, so the deploy always +// uses this build's ai_runtime_task-aware bundle schema — an older `databricks` on +// PATH would drop the unknown ai_runtime_task field and deploy a task-less job. // -// PROTOTYPE SCOPE / shell-out choice: deploy+run are driven by shelling out to the -// `databricks bundle` CLI, which is the literal "wrapper around DABs" the meeting -// described and keeps this diff legible. A production version would instead call -// cmd/bundle/utils.ProcessBundle(...) in-process (see TODO(prototype) below) to -// avoid the child-process dependency. The behavior (persistent job, deploy -// latency, no ephemeral-job GC) is identical either way — the point being to -// exercise the real DABs path end to end. -// -// KNOWN GAP (called out in the design doc): the deployed job is PERSISTENT and is -// NOT swept by the Jobs ephemeral-job GC (JobsSoftDeletion only sweeps -// EPHEMERAL/WORKFLOW types). Every `air run --via-dabs` leaves a job behind. A real -// impl must choose a reuse/cleanup policy (stable per-experiment bundle so re-runs -// update one job, or `bundle destroy` after run). The prototype leaves the job and -// logs the tradeoff. +// KNOWN TRADEOFF (design doc): the deployed job is PERSISTENT and is NOT swept by +// the Jobs ephemeral-job GC (JobsSoftDeletion only sweeps EPHEMERAL/WORKFLOW +// types). Runs of distinct experiments accumulate distinct jobs, drifting toward +// the per-workspace saved-jobs cap. Job reuse (a stable per-experiment bundle name, +// so re-runs update one job) mitigates it; a cleanup/TTL story is future work. import ( "context" + "errors" "fmt" "os" - "os/exec" "path/filepath" + "strconv" + "strings" + "github.com/databricks/cli/bundle" + "github.com/databricks/cli/bundle/config/resources" + bundleresources "github.com/databricks/cli/bundle/resources" + "github.com/databricks/cli/cmd/bundle/utils" "github.com/databricks/cli/libs/cmdio" "github.com/databricks/cli/libs/env" + "github.com/databricks/cli/libs/flags" + "github.com/databricks/cli/libs/logdiag" "github.com/databricks/databricks-sdk-go" + "github.com/databricks/databricks-sdk-go/service/jobs" + "github.com/spf13/cobra" ) -// viaDABSEnv opts into the DABs submit path without the flag (flag wins). -const viaDABSEnv = "AIR_VIA_DABS" - -// dabsTarget is the bundle target the prototype deploys/runs against. A real impl -// would let the user pick (or default per-workspace); "dev" gets DABs dev-mode -// (per-user name prefix + isolation), which suits AIR's per-user runs. +// dabsTarget is the bundle target air deploys/runs against. `dev` mode gives DABs +// dev-mode (per-user name prefix + isolation), which matches AIR's per-user runs +// and keeps one user's jobs from colliding with another's on the same experiment. const dabsTarget = "dev" -// submitViaDABS is the Option-D analogue of submitWorkload (runsubmit.go), with the -// identical signature so run.go can pick either at the call site. It converts the -// run config to a bundle, writes the project (databricks.yml + command.sh + code), -// then deploys and runs it. Returns the run_id + dashboard URL. -func submitViaDABS(ctx context.Context, w *databricks.WorkspaceClient, cfg *runConfig, configPath, idempotencyKey string) (int64, string, error) { +// submitWorkload converts the run config to a bundle, deploys it (creating/updating +// a persistent job), then triggers a run. It returns the new run_id and its +// dashboard URL. The whole flow is in-process: deploy reuses the bundle deploy +// orchestration (state, sync, engine) via ProcessBundle; run calls Jobs.RunNow on +// the deployed job so the run_id is returned directly (fire-and-return, no wait). +func submitWorkload(ctx context.Context, w *databricks.WorkspaceClient, cfg *runConfig, configPath, idempotencyKey string) (int64, string, error) { // The converter's gate rejects configs a bundle can't represent faithfully - // (e.g. git-pinned code_source, $CODE_SOURCE_PATH commands). Fail fast with its - // actionable message before we write or deploy anything. + // (docker_image, usage_policy, git-pinned code_source, $CODE_SOURCE_PATH + // commands). Fail fast with its actionable message before writing or deploying. if err := checkBundleConvertible(cfg); err != nil { return 0, "", err } - // Idempotency doesn't map onto deploy+run (two calls, not one tokened submit); - // note rather than silently honor it. A real impl dedupes on bundle/job identity. + // Idempotency does not map onto deploy+run (two calls, not one tokened submit). + // Note it rather than silently honoring it; job identity (stable bundle name) + // is the DABs-native dedupe mechanism. if idempotencyKey != "" || cfg.IdempotencyToken != nil { cmdio.LogString(ctx, "note: --idempotency-key is ignored on the DABs path (deploy+run is not a single idempotent call)") } - root, cleanup, err := writeBundleProject(cfg, configPath) + bundleRoot, cleanup, err := writeBundleProject(cfg, configPath) if err != nil { return 0, "", err } defer cleanup() // Surface the generated artifact's path (transparency): the user can inspect - // what we deploy on their behalf. The dir is temporary (auto-removed on exit) — - // labeled as such so we don't imply a durable, user-managed location. For the - // full contents without deploying, use `air run --dry-run --via-dabs`. - cmdio.LogString(ctx, "Generated bundle (temporary): "+filepath.Join(root, "databricks.yml")) + // what we deploy on their behalf. The dir is temporary (auto-removed on exit). + // For the full contents without deploying, use `air run --dry-run`. + cmdio.LogString(ctx, "Generated bundle (temporary): "+filepath.Join(bundleRoot, "databricks.yml")) // deploy: creates/updates the persistent job. This is the step that adds - // latency vs the ephemeral submit (measured ~2-3s warm, ~12s first-of-session). + // latency vs the old ephemeral submit (~2-3s warm, ~12s first-of-session). cmdio.LogString(ctx, "Deploying bundle (creates a persistent job)...") - if err := runBundle(ctx, w, root, "deploy"); err != nil { + b, err := deployBundle(ctx, w, bundleRoot) + if err != nil { return 0, "", fmt.Errorf("bundle deploy: %w", err) } - // run: RunNow on the deployed job. --no-wait returns as soon as the run starts, - // matching the ephemeral path's fire-and-return (we don't block on completion). + // run: RunNow on the just-deployed job, returning as soon as the run is created + // (no wait), matching the old submit path's fire-and-return. cmdio.LogString(ctx, "Triggering run...") - if err := runBundle(ctx, w, root, "run", cfg.ExperimentName, "--no-wait"); err != nil { + runID, err := runDeployedJob(ctx, w, b, cfg.ExperimentName) + if err != nil { return 0, "", fmt.Errorf("bundle run: %w", err) } cmdio.LogString(ctx, "note: this run created a PERSISTENT job (not auto-GC'd like an ephemeral run). Use `air list` to find it.") - // TODO(prototype): parse run_id + dashboard URL from `databricks bundle run - // --output json`. For now return 0 and the jobs page; the run is live either way. - return 0, workspaceJobsURL(w), nil + dashboardURL := strings.TrimRight(w.Config.Host, "/") + "/jobs/runs/" + strconv.FormatInt(runID, 10) + return runID, dashboardURL, nil } -// renderBundle produces the exact databricks.yml the run path would deploy: the +// renderBundle produces the exact databricks.yml the run path deploys: the // exportbundle.go converter output plus the dev targets block the run path appends. -// It touches no filesystem and does not deploy, so `air run --dry-run --via-dabs` -// can show the user the artifact we'd generate on their behalf (transparency), and +// It touches no filesystem and does not deploy, so `air run --dry-run` can show the +// user the artifact we'd generate on their behalf (transparency), and // writeBundleProject reuses it so preview and real run never diverge. func renderBundle(cfg *runConfig, configPath string) (string, error) { if err := checkBundleConvertible(cfg); err != nil { @@ -119,37 +125,36 @@ func renderBundle(cfg *runConfig, configPath string) (string, error) { } // writeBundleProject renders databricks.yml (via renderBundle) plus the command.sh -// and code the bundle syncs, into a temp bundle root. Returns the root and a -// cleanup func. +// the bundle syncs, into a temp bundle root. Returns the root and a cleanup func. // -// TODO(prototype): reuse the ephemeral path's snapshot/upload staging so the exact -// code tree + command.sh land in the bundle folder. This sketch writes command.sh -// from cfg.Command; wiring code_source snapshotting through the inherited -// bundle/config/mutator/aicode mutator is the remaining piece. +// TODO(air): reuse the code-snapshot staging (snapshot.go) so the user's code tree +// lands in the bundle folder for `bundle sync` to upload. This writes command.sh +// from cfg.Command; wiring code_source snapshotting into the synced folder is the +// remaining piece before code_source runs are supported on this path. func writeBundleProject(cfg *runConfig, configPath string) (string, func(), error) { body, err := renderBundle(cfg, configPath) if err != nil { return "", func() {}, err } - root, err := os.MkdirTemp("", "air-dabs-*") + bundleRoot, err := os.MkdirTemp("", "air-dabs-*") if err != nil { return "", func() {}, err } - cleanup := func() { _ = os.RemoveAll(root) } + cleanup := func() { _ = os.RemoveAll(bundleRoot) } - if err := os.WriteFile(filepath.Join(root, "databricks.yml"), []byte(body), 0o600); err != nil { + if err := os.WriteFile(filepath.Join(bundleRoot, "databricks.yml"), []byte(body), 0o600); err != nil { cleanup() return "", func() {}, err } // command.sh: the entrypoint the task's command_path points at (synced by deploy). if cfg.Command != nil { - if err := os.WriteFile(filepath.Join(root, bundleCommandScript), []byte(*cfg.Command), 0o600); err != nil { + if err := os.WriteFile(filepath.Join(bundleRoot, bundleCommandScript), []byte(*cfg.Command), 0o600); err != nil { cleanup() return "", func() {}, err } } - return root, cleanup, nil + return bundleRoot, cleanup, nil } // bundleTargetsBlock is the minimal dev target appended to the converted YAML so @@ -157,43 +162,106 @@ func writeBundleProject(cfg *runConfig, configPath string) (string, func(), erro // isolation. The workspace host comes from the active CLI profile, so it's omitted // here (the bundle resolves it from auth at deploy time). func bundleTargetsBlock() string { - // The comment supersedes the export-bundle header's "add a targets block" - // note: on the run path we append this dev target for you (host resolved from - // your CLI profile at deploy time), so `air run --via-dabs` needs no manual edit. - return "\n# Appended by `air run --via-dabs` (host resolved from your CLI profile at deploy time):\ntargets:\n " + dabsTarget + ":\n mode: development\n default: true\n" + return "\n# Appended by `air run` (host resolved from your CLI profile at deploy time):\ntargets:\n " + dabsTarget + ":\n mode: development\n default: true\n" } -// runBundle shells out to `databricks bundle -t ` in the bundle -// root. Uses the direct engine (no Terraform / no registry.terraform.io dep — GA -// default on new CLIs, set explicitly for older ones). Auth is inherited from the -// same profile/env the parent air command resolved. +// deployBundle deploys the bundle rooted at bundleRoot in-process and returns the +// configured *bundle.Bundle (with resource IDs populated, so the caller can resolve +// the deployed job). It reuses cmd/bundle/utils.ProcessBundle — the exact +// orchestration `databricks bundle deploy` runs (load, validate, build, sync, +// deploy) — via a synthetic carrier command: // -// TODO(prototype): replace with an in-process cmd/bundle/utils.ProcessBundle call -// so there's no dependency on a `databricks` binary on PATH. -func runBundle(ctx context.Context, w *databricks.WorkspaceClient, root string, args ...string) error { - full := append([]string{"bundle"}, args...) - full = append(full, "-t", dabsTarget) - cmd := exec.CommandContext(ctx, "databricks", full...) - cmd.Dir = root - cmd.Env = append(os.Environ(), - "DATABRICKS_BUNDLE_ENGINE=direct", - "DATABRICKS_HOST="+w.Config.Host, - ) - cmd.Stdout = os.Stderr // keep our stdout clean for the JSON envelope - cmd.Stderr = os.Stderr - return cmd.Run() +// - air's own cobra command carries none of the flags ProcessBundle reads +// (--var, --target, --profile, --output), so we build a throwaway command that +// declares them and seed the context (logdiag + auth profile + bundle root). +// - the bundle root is steered to bundleRoot with DATABRICKS_BUNDLE_ROOT on the +// context (env.Set is context-scoped, not a real env var), so MustConfigureBundle +// loads our generated databricks.yml instead of walking cwd. +// - the direct engine (DATABRICKS_BUNDLE_ENGINE=direct) avoids the Terraform +// provider/registry dependency; it is the GA default on new CLIs and set +// explicitly here so the deploy is self-contained. +func deployBundle(ctx context.Context, w *databricks.WorkspaceClient, bundleRoot string) (*bundle.Bundle, error) { + cmd := newBundleCarrierCommand(ctx, w, bundleRoot) + + b, err := utils.ProcessBundle(cmd, utils.ProcessOptions{ + FastValidate: true, + Build: true, + Deploy: true, + // air's command context already has logdiag initialized (root's + // PersistentPreRunE); re-initializing panics ("InitContext twice"). Same as + // the pipelines CLI run path. + SkipInitContext: true, + }) + if err != nil { + return nil, err + } + if b == nil || logdiag.HasError(cmd.Context()) { + return nil, errors.New("bundle deploy failed") + } + return b, nil } -func workspaceJobsURL(w *databricks.WorkspaceClient) string { - host := w.Config.Host - for len(host) > 0 && host[len(host)-1] == '/' { - host = host[:len(host)-1] +// newBundleCarrierCommand builds the synthetic cobra command deployBundle hands to +// ProcessBundle. It declares the flags ProcessBundle/configureBundle read and seeds +// the context so the deploy runs with air's resolved auth against the generated +// bundle root. The command is never added to a tree and never Execute()d — it only +// carries flags + context into the bundle libraries. +func newBundleCarrierCommand(ctx context.Context, w *databricks.WorkspaceClient, bundleRoot string) *cobra.Command { + cmd := &cobra.Command{Use: "air-bundle-deploy"} + + // Flags ProcessBundleRet (and the bundle root loader) read by name. + cmd.Flags().StringSlice("var", nil, "") + cmd.Flags().StringP("target", "t", dabsTarget, "") + cmd.Flags().StringP("profile", "p", "", "") + outputFlag := flags.OutputText + cmd.Flags().Var(&outputFlag, "output", "") + + // Forward air's active profile so the bundle authenticates the same way air + // did, instead of falling back to the default profile. + if w.Config.Profile != "" { + _ = cmd.Flags().Set("profile", w.Config.Profile) + } + + // logdiag is how the bundle libraries report diagnostics; ProcessBundle asserts + // it is set up (we pass SkipInitContext=false, so it initializes, but the auth + // profile + bundle root must be seeded before that on the same context). + ctx = env.Set(ctx, "DATABRICKS_BUNDLE_ROOT", bundleRoot) + ctx = env.Set(ctx, "DATABRICKS_BUNDLE_ENGINE", "direct") + cmd.SetContext(ctx) + return cmd +} + +// runDeployedJob triggers a run of the just-deployed job and returns the new run_id +// without waiting for completion. It resolves the job by its bundle resource key +// (the experiment name) to get the server-assigned job_id that deploy populated, +// then calls Jobs.RunNow. AIR runs take no run parameters, so RunNow needs only the +// job_id. (We call RunNow directly rather than the bundle runner's NoWait path, +// which returns nil instead of the run_id.) +func runDeployedJob(ctx context.Context, w *databricks.WorkspaceClient, b *bundle.Bundle, experimentName string) (int64, error) { + ref, err := bundleresources.Lookup(b, experimentName, isRunnableJob) + if err != nil { + return 0, fmt.Errorf("locating the deployed job: %w", err) + } + job, ok := ref.Resource.(*resources.Job) + if !ok { + return 0, fmt.Errorf("deployed resource %q is not a job", experimentName) + } + jobID, err := strconv.ParseInt(job.ID, 10, 64) + if err != nil { + return 0, fmt.Errorf("deployed job has no valid id (deploy may have failed): %q", job.ID) + } + + wait, err := w.Jobs.RunNow(ctx, jobs.RunNow{JobId: jobID}) + if err != nil { + return 0, err } - return host + "/jobs" + return wait.RunId, nil } -// resolveViaDABS decides whether to take the DABs submit path: the --via-dabs flag -// or AIR_VIA_DABS=1. Kept here so run.go's wiring stays a one-liner. -func resolveViaDABS(ctx context.Context, flag bool) bool { - return flag || env.Get(ctx, viaDABSEnv) == "1" +// isRunnableJob filters resources.Lookup to job resources, so an experiment name +// that also matched another resource type can't be selected. AIR bundles only ever +// contain one job, so this is a safety guard rather than a disambiguator. +func isRunnableJob(ref bundleresources.Reference) bool { + _, ok := ref.Resource.(*resources.Job) + return ok } diff --git a/experimental/air/cmd/runsubmit.go b/experimental/air/cmd/runsubmit.go index 8c0be55260d..3fc298d2caf 100644 --- a/experimental/air/cmd/runsubmit.go +++ b/experimental/air/cmd/runsubmit.go @@ -1,191 +1,16 @@ package aircmd -import ( - "context" - "errors" - "fmt" - "path" - "strconv" - "strings" - - "github.com/databricks/cli/libs/env" - "github.com/databricks/cli/libs/filer" - "github.com/databricks/databricks-sdk-go" - "github.com/databricks/databricks-sdk-go/service/compute" - "github.com/databricks/databricks-sdk-go/service/jobs" - "github.com/google/uuid" -) - -// dlRuntimeImageEnv overrides the default deep-learning runtime image. -const dlRuntimeImageEnv = "DATABRICKS_DL_RUNTIME_IMAGE" - +// This file used to hold the ephemeral runs/submit path (submitWorkload + +// buildSubmitPayload). That path was retired: `air run` now submits through DABs +// (convert -> deploy -> run, see rundabs.go), per the AIR-CLI/DABs agreement. What +// remains here is the runtime-image defaults shared by the DABs converter +// (exportbundle.go) — the one piece of the old submit path the bundle path needs. + +// defaultDlRuntimeImage is the default deep-learning runtime channel. The converter +// strips the CLIENT-GPU- prefix; unlike the retired submit path it does not read a +// process-env override, so a generated bundle is reproducible. const defaultDlRuntimeImage = "CLIENT-GPU-4" -// aiRuntimeEnvironmentKey ties the task to the serverless environment that -// carries the runtime channel. +// aiRuntimeEnvironmentKey ties the task to the serverless environment that carries +// the runtime channel. const aiRuntimeEnvironmentKey = "default" - -// dlRuntimeImage resolves the bare runtime channel (config version, else env, -// else default), always stripping the CLIENT-GPU- prefix. -func dlRuntimeImage(ctx context.Context, runtimeVersion string) string { - img := runtimeVersion - if img == "" { - img = env.Get(ctx, dlRuntimeImageEnv) - } - if img == "" { - img = defaultDlRuntimeImage - } - return strings.TrimPrefix(img, "CLIENT-GPU-") -} - -// buildSubmitPayload assembles the runs/submit payload. commandPath is the -// workspace path of the uploaded command.sh; dlImage is the runtime channel. -// -// max_retries is always sent (including 0) so the user's YAML value is honored: -// setting it to 0 explicitly disables retries rather than falling back to the -// server default. retry_on_timeout is sent only when retries are allowed, and is -// omitempty so the wire form matches the Python CLI (which never emits a bare -// "false"). Jobs performs the retries — each attempt is a fresh AI Runtime -// workload. -func buildSubmitPayload(cfg *runConfig, commandPath, dlImage string, snap snapshotResult) jobs.SubmitRun { - task := jobs.AiRuntimeTask{ - Experiment: cfg.ExperimentName, - Deployments: []jobs.DeploymentSpec{{ - CommandPath: commandPath, - Compute: jobs.ComputeSpec{ - AcceleratorType: jobs.ComputeSpecAcceleratorType(cfg.Compute.AcceleratorType), - AcceleratorCount: cfg.Compute.NumAccelerators, - }, - }}, - CodeSourcePath: snap.CodeSourcePath, - // TEMP: git_state_path / git_diff_path are intentionally NOT sent. The typed - // jobs.AiRuntimeTask (and its source proto, ai_runtime_task.proto) has no such - // fields, so the typed SDK path cannot carry them. This is safe today because - // nothing in the backend consumes those fields — the AI Runtime task proto - // never declared them, so even the Python CLI's raw-JSON values were dropped - // on deserialization. The git_state.json / git_diff.patch sidecars are still - // uploaded next to the tarball (see snapshot.go) for human inspection. - // If the backend later adds these fields to the proto, regenerate the SDK and - // wire snap.GitStatePath / snap.GitDiffPath back in here. - } - if cfg.MLflowRunName != nil { - task.MlflowRun = *cfg.MLflowRunName - } - if cfg.MLflowExperimentDirectory != nil { - task.MlflowExperimentDirectory = *cfg.MLflowExperimentDirectory - } - - maxRetries := cfg.maxRetries() - st := jobs.SubmitTask{ - TaskKey: cfg.ExperimentName, - RunIf: jobs.RunIfAllSuccess, - AiRuntimeTask: &task, - EnvironmentKey: aiRuntimeEnvironmentKey, - MaxRetries: maxRetries, - // retry_on_timeout only makes sense when retries are allowed; otherwise - // omit it (matches Python's native path, which sets retry_on_timeout only - // under the same > 0 gate). - RetryOnTimeout: maxRetries > 0, - ForceSendFields: []string{"MaxRetries"}, - } - - return jobs.SubmitRun{ - RunName: cfg.ExperimentName, - TimeoutSeconds: cfg.timeoutSeconds(), - Tasks: []jobs.SubmitTask{st}, - Environments: []jobs.JobEnvironment{{ - EnvironmentKey: aiRuntimeEnvironmentKey, - Spec: &compute.Environment{EnvironmentVersion: dlImage}, - }}, - } -} - -// submitToken resolves the idempotency token: the --idempotency-key flag wins, -// then the config's token, else a generated one. Over-long tokens error rather -// than truncate, since truncation could make two distinct tokens collide. -func submitToken(flag string, cfg *runConfig) (string, error) { - token := flag - if token == "" && cfg.IdempotencyToken != nil { - token = *cfg.IdempotencyToken - } - if token == "" { - token = uuid.NewString() - } - if len(token) > 64 { - return "", fmt.Errorf("idempotency token must be 64 characters or less, got %d", len(token)) - } - return token, nil -} - -// submitWorkload runs the submit happy path: ensure the experiment directory, -// upload the launch artifacts, assemble the Jobs payload, and submit it. It -// returns the new run_id and its dashboard URL. -func submitWorkload(ctx context.Context, w *databricks.WorkspaceClient, cfg *runConfig, configPath, idempotencyKey string) (int64, string, error) { - // Resolving usage_policy_name to a budget policy id is not ported yet; reject - // rather than silently drop. - if cfg.UsagePolicyName != nil { - return 0, "", errors.New("usage_policy_name is not yet supported") - } - - // Resolve the idempotency token first so a bad key fails before any upload. - token, err := submitToken(idempotencyKey, cfg) - if err != nil { - return 0, "", err - } - - experimentDir := "" - if cfg.MLflowExperimentDirectory != nil { - experimentDir = *cfg.MLflowExperimentDirectory - } - if err := ensureExperimentDirectory(ctx, w, experimentDir); err != nil { - return 0, "", err - } - - base, err := userWorkspaceDir(ctx, w) - if err != nil { - return 0, "", err - } - runName := "" - if cfg.MLflowRunName != nil { - runName = *cfg.MLflowRunName - } - funcDir := cliLaunchDir(base, cfg.ExperimentName, runName) - - fc, err := filer.NewWorkspaceFilesClient(w, funcDir) - if err != nil { - return 0, "", err - } - items, err := buildArtifacts(cfg, configPath) - if err != nil { - return 0, "", err - } - if err := uploadArtifacts(ctx, fc, items); err != nil { - return 0, "", err - } - - // Package and upload the code snapshot, if any. The resulting paths ride on the - // ai_runtime_task; a run with no code_source leaves them empty. Snapshot is the - // only code_source type; guard against a nil block so snapshotCodeSource never - // dereferences a missing snapshot. - var snap snapshotResult - if cfg.CodeSource != nil && cfg.CodeSource.Snapshot != nil { - snap, err = snapshotCodeSource(ctx, w, cfg.CodeSource.Snapshot, configPath, base, funcDir) - if err != nil { - return 0, "", err - } - } - - runtimeVersion, _ := cfg.runtimeVersion() - payload := buildSubmitPayload(cfg, path.Join(funcDir, commandScriptName), dlRuntimeImage(ctx, runtimeVersion), snap) - payload.IdempotencyToken = token - - // Submit returns as soon as the run is created; we don't wait for it to finish. - wait, err := w.Jobs.Submit(ctx, payload) - if err != nil { - return 0, "", err - } - runID := wait.RunId - - dashboardURL := strings.TrimRight(w.Config.Host, "/") + "/jobs/runs/" + strconv.FormatInt(runID, 10) - return runID, dashboardURL, nil -} diff --git a/experimental/air/cmd/runsubmit_test.go b/experimental/air/cmd/runsubmit_test.go deleted file mode 100644 index fd5103599df..00000000000 --- a/experimental/air/cmd/runsubmit_test.go +++ /dev/null @@ -1,216 +0,0 @@ -package aircmd - -import ( - "encoding/json" - "path/filepath" - "strings" - "testing" - - "github.com/databricks/cli/libs/testserver" - "github.com/databricks/databricks-sdk-go" - "github.com/databricks/databricks-sdk-go/service/jobs" - "github.com/stretchr/testify/assert" - "github.com/stretchr/testify/require" -) - -func TestDlRuntimeImage(t *testing.T) { - ctx := t.Context() - // A config runtime version wins and is used bare. - assert.Equal(t, "5", dlRuntimeImage(ctx, "5")) - // The CLIENT-GPU- prefix is always stripped, even from the config version. - assert.Equal(t, "5", dlRuntimeImage(ctx, "CLIENT-GPU-5")) - // Default, with the prefix stripped. - assert.Equal(t, "4", dlRuntimeImage(ctx, "")) - // Env override, prefix stripped. - t.Setenv(dlRuntimeImageEnv, "CLIENT-GPU-7") - assert.Equal(t, "7", dlRuntimeImage(ctx, "")) -} - -func TestBuildSubmitPayload(t *testing.T) { - cfg := &runConfig{ - ExperimentName: "exp", - Command: new("python train.py"), - Compute: &computeConfig{AcceleratorType: "GPU_8xH100", NumAccelerators: 16}, - MaxRetries: new(2), - TimeoutMinutes: new(30), - MLflowRunName: new("run-v2"), - MLflowExperimentDirectory: new("/Workspace/Users/me/exp"), - } - - p := buildSubmitPayload(cfg, "/d/command.sh", "5", snapshotResult{}) - - assert.Equal(t, "exp", p.RunName) - assert.Equal(t, 1800, p.TimeoutSeconds) - require.Len(t, p.Environments, 1) - assert.Equal(t, aiRuntimeEnvironmentKey, p.Environments[0].EnvironmentKey) - require.NotNil(t, p.Environments[0].Spec) - assert.Equal(t, "5", p.Environments[0].Spec.EnvironmentVersion) - - require.Len(t, p.Tasks, 1) - task := p.Tasks[0] - assert.Equal(t, "exp", task.TaskKey) - assert.Equal(t, jobs.RunIfAllSuccess, task.RunIf) - assert.Equal(t, aiRuntimeEnvironmentKey, task.EnvironmentKey) - assert.Equal(t, 2, task.MaxRetries) - assert.True(t, task.RetryOnTimeout) - - at := task.AiRuntimeTask - require.NotNil(t, at) - assert.Equal(t, "exp", at.Experiment) - assert.Equal(t, "run-v2", at.MlflowRun) - assert.Equal(t, "/Workspace/Users/me/exp", at.MlflowExperimentDirectory) - require.Len(t, at.Deployments, 1) - assert.Equal(t, "/d/command.sh", at.Deployments[0].CommandPath) - assert.Equal(t, jobs.ComputeSpec{AcceleratorType: jobs.ComputeSpecAcceleratorTypeGpu8xH100, AcceleratorCount: 16}, at.Deployments[0].Compute) -} - -func TestBuildSubmitPayloadDefaultRetries(t *testing.T) { - // max_retries unset defaults to 3 (matching the Python native path), so both - // retry fields are sent. - cfg := &runConfig{ - ExperimentName: "exp", - Command: new("x"), - Compute: &computeConfig{AcceleratorType: "GPU_1xH100", NumAccelerators: 1}, - } - task := buildSubmitPayload(cfg, "/d/command.sh", "4", snapshotResult{}).Tasks[0] - assert.Equal(t, defaultMaxRetries, task.MaxRetries) - assert.True(t, task.RetryOnTimeout) -} - -func TestBuildSubmitPayloadNoRetries(t *testing.T) { - // max_retries: 0 must be sent explicitly so Jobs honors "no retries" instead - // of applying the server default. retry_on_timeout is omitted when retries - // aren't allowed. - cfg := &runConfig{ - ExperimentName: "exp", - Command: new("x"), - Compute: &computeConfig{AcceleratorType: "GPU_1xH100", NumAccelerators: 1}, - MaxRetries: new(0), - } - task := buildSubmitPayload(cfg, "/d/command.sh", "4", snapshotResult{}).Tasks[0] - assert.Equal(t, 0, task.MaxRetries) - assert.False(t, task.RetryOnTimeout) - - b, err := json.Marshal(task) - require.NoError(t, err) - assert.Contains(t, string(b), `"max_retries":0`) - assert.NotContains(t, string(b), "retry_on_timeout") -} - -func TestSubmitToken(t *testing.T) { - cfg := &runConfig{IdempotencyToken: new("from-config")} - - tok, err := submitToken("from-flag", cfg) // flag wins - require.NoError(t, err) - assert.Equal(t, "from-flag", tok) - - tok, err = submitToken("", cfg) // then config - require.NoError(t, err) - assert.Equal(t, "from-config", tok) - - tok, err = submitToken("", &runConfig{}) // else generated - require.NoError(t, err) - assert.NotEmpty(t, tok) - - // An over-long token errors instead of being truncated. - _, err = submitToken(strings.Repeat("a", 65), cfg) - require.ErrorContains(t, err, "64 characters or less") -} - -func TestSubmitWorkload(t *testing.T) { - server := testserver.New(t) - t.Cleanup(server.Close) - - // Register before AddDefaultHandlers: the router is first-wins, so this must claim the route ahead of the default handler. - var got jobs.SubmitRun - server.Handle("POST", "/api/2.2/jobs/runs/submit", func(req testserver.Request) any { - require.NoError(t, json.Unmarshal(req.Body, &got)) - return jobs.SubmitRunResponse{RunId: 777} - }) - testserver.AddDefaultHandlers(server) - - w, err := databricks.NewWorkspaceClient(&databricks.Config{Host: server.URL, Token: "token"}) - require.NoError(t, err) - - cfgPath := writeConfigFile(t, "run.yaml", minimalConfig) - cfg, err := loadRunConfig(cfgPath) - require.NoError(t, err) - - runID, dashboardURL, err := submitWorkload(t.Context(), w, cfg, cfgPath, "idem-key") - require.NoError(t, err) - assert.Equal(t, int64(777), runID) - assert.Contains(t, dashboardURL, "/jobs/runs/777") - - // The submitted payload is a native ai_runtime_task pointing at the uploaded - // command.sh under the run's launch directory. - assert.Equal(t, "my-run", got.RunName) - assert.Equal(t, "idem-key", got.IdempotencyToken) - require.Len(t, got.Environments, 1) - require.Len(t, got.Tasks, 1) - at := got.Tasks[0].AiRuntimeTask - require.NotNil(t, at) - require.Len(t, at.Deployments, 1) - d := at.Deployments[0] - assert.True(t, strings.HasSuffix(d.CommandPath, "/"+commandScriptName), d.CommandPath) - assert.Contains(t, d.CommandPath, "/.air/cli_launch/") - assert.Equal(t, jobs.ComputeSpec{AcceleratorType: jobs.ComputeSpecAcceleratorTypeGpu1xH100, AcceleratorCount: 1}, d.Compute) -} - -// TestSubmitWorkloadWithCodeSource exercises the snapshot path end to end: a -// git-pinned code_source is packaged, uploaded, and its paths attached to the task. -func TestSubmitWorkloadWithCodeSource(t *testing.T) { - server := testserver.New(t) - t.Cleanup(server.Close) - - // Register before AddDefaultHandlers: the router is first-wins, so this must claim the route ahead of the default handler. - var got jobs.SubmitRun - server.Handle("POST", "/api/2.2/jobs/runs/submit", func(req testserver.Request) any { - require.NoError(t, json.Unmarshal(req.Body, &got)) - return jobs.SubmitRunResponse{RunId: 555} - }) - testserver.AddDefaultHandlers(server) - w, err := databricks.NewWorkspaceClient(&databricks.Config{Host: server.URL, Token: "token"}) - require.NoError(t, err) - - // A git repo committed at HEAD, referenced by commit so packaging is git_archive. - repo := newTestRepo(t) - writeRepoFile(t, repo, "train.py", "print()") - sha := commitAll(t, repo, "init") - - cfg := minimalConfig + ` -code_source: - type: snapshot - snapshot: - root_path: ` + repo + ` - git: - commit: ` + sha + ` -` - cfgPath := writeConfigFile(t, "run.yaml", cfg) - loaded, err := loadRunConfig(cfgPath) - require.NoError(t, err) - - _, _, err = submitWorkload(t.Context(), w, loaded, cfgPath, "idem") - require.NoError(t, err) - - at := got.Tasks[0].AiRuntimeTask - // The tarball path is under the user's repo_snapshots dir. git_state_path / - // git_diff_path are not asserted: the typed jobs.AiRuntimeTask has no such fields - // (see the TEMP note in buildSubmitPayload), so they aren't sent. The git_state - // sidecar file is still uploaded next to the tarball — covered by TestRunSnapshot. - assert.Contains(t, at.CodeSourcePath, "/.air/repo_snapshots/"+filepath.Base(repo)+"/") - assert.True(t, strings.HasSuffix(at.CodeSourcePath, ".tar.gz"), at.CodeSourcePath) -} - -func TestSubmitWorkloadGuards(t *testing.T) { - w := newFakeWorkspaceClient(t) - cfgPath := writeConfigFile(t, "run.yaml", minimalConfig) - base, err := loadRunConfig(cfgPath) - require.NoError(t, err) - - t.Run("usage_policy_name rejected", func(t *testing.T) { - cfg := *base - cfg.UsagePolicyName = new("p") - _, _, err := submitWorkload(t.Context(), w, &cfg, cfgPath, "") - require.ErrorContains(t, err, "usage_policy_name is not yet supported") - }) -} From f01b54082d1057e6dc01fd854806592b81ab4d9a Mon Sep 17 00:00:00 2001 From: vinchenzo-db Date: Wed, 15 Jul 2026 21:14:51 +0000 Subject: [PATCH 03/11] [air] Deliver code_source via DABs immutable_folder; trim comments MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Wires code_source snapshots through the DABs content-addressed snapshot path (experimental.immutable_folder) so `air run` uploads the user's code tree the same zip-and-fingerprint way AIR does today — giving an apples-to-apples benchmark against the mechanism the uploads-vs-DABs scoping doc compares (DECO-27087), not the per-file bundle sync. - exportbundle.go: emit experimental.immutable_folder=true; emit command_path as a bundle-root-relative path so translate_paths rewrites it to the deployed snapshot location (emitting ${workspace.file_path} directly is validated as a local file and fails). The convertibility gate now allows working-tree code_source snapshots but still rejects git-pinned (working-tree upload can't pin a commit) and remote_volume (immutable_folder is Workspace-Files-only) snapshots. - rundabs.go: writeBundleProject copies the code_source working tree into the bundle root (honoring include_paths) so deploy uploads it in the immutable snapshot. - Trim AI-narration comments across the air files (provenance dumps, verification history, migration storytelling) down to what/why. Verified on staging: `air run` with a code_source snapshot deploys via "Uploading immutable bundle snapshot..." and submits a run. (The AI Runtime execution then fails on a backend PrincipalContext/entity-owner error, independent of code delivery and of the CLI path.) Co-authored-by: Isaac --- experimental/air/cmd/cancel.go | 6 - experimental/air/cmd/exportbundle.go | 151 ++++++++-------- experimental/air/cmd/exportbundle_test.go | 69 +++++++- experimental/air/cmd/get.go | 3 - experimental/air/cmd/list.go | 8 +- experimental/air/cmd/run.go | 3 +- experimental/air/cmd/rundabs.go | 203 +++++++++++++--------- 7 files changed, 261 insertions(+), 182 deletions(-) diff --git a/experimental/air/cmd/cancel.go b/experimental/air/cmd/cancel.go index 8f93d0ac928..519a7a82068 100644 --- a/experimental/air/cmd/cancel.go +++ b/experimental/air/cmd/cancel.go @@ -191,12 +191,6 @@ func runNotFound(err error) bool { // cancelRun requests cancellation of a single job run. The cancel is async, so // the returned waiter is ignored. -// -// DABs note: this works unchanged for DABs-submitted runs. A `bundle run` produces -// a normal Jobs run with a run_id, and cancelling that run terminates the AIR -// workload the same way — BYOT/AICM dispatch is keyed on task type, not on how the -// run was submitted. (`cancel --all` discovers runs via `air list`, which the -// --via-dabs list mode widens to include DABs JOB_RUNs.) func cancelRun(ctx context.Context, w *databricks.WorkspaceClient, rid string) error { runID, err := strconv.ParseInt(rid, 10, 64) if err != nil || runID <= 0 { diff --git a/experimental/air/cmd/exportbundle.go b/experimental/air/cmd/exportbundle.go index f4feb2c899b..32898240765 100644 --- a/experimental/air/cmd/exportbundle.go +++ b/experimental/air/cmd/exportbundle.go @@ -8,21 +8,14 @@ import ( "go.yaml.in/yaml/v3" ) -// This file implements the train.yaml -> databricks.yml conversion behind -// `air export-bundle`: a one-time converter that turns an AIR run YAML into a -// Databricks Asset Bundle deploying the same workload as a native ai_runtime_task. -// It is the durable counterpart to `air run` (which submits an ephemeral one-time -// run): the emitted bundle is a persistent Jobs resource the user owns and deploys -// with `databricks bundle deploy`. +// This file converts a train.yaml runConfig to a databricks.yml Asset Bundle +// deploying the same workload as a native ai_runtime_task. It backs both +// `air export-bundle` and the `air run` submit path (rundabs.go). // -// The field mapping mirrors buildSubmitPayload (runsubmit.go) so the generated -// task is the shape `air run` would submit. Structural validity is already -// guaranteed upstream by runConfig.validate(); this file adds a second, -// converter-specific gate — checkBundleConvertible — that rejects the configs a -// bundle cannot represent faithfully, rather than emitting a databricks.yml that -// deploys but misbehaves. The gate exists because train.yaml and a bundle are not -// 1:1: some run fields have no bundle equivalent, and some assume the AIR run -// harness (e.g. $CODE_SOURCE_PATH) that a bundle does not provide. +// checkBundleConvertible rejects configs a bundle cannot represent faithfully: +// train.yaml and a bundle are not 1:1, and some run fields have no bundle +// equivalent or assume the AIR run harness (e.g. $CODE_SOURCE_PATH) a bundle does +// not provide. // bundleCommandScript is the entrypoint filename the emitted bundle references and // that `bundle sync` uploads alongside the user's code. @@ -30,32 +23,25 @@ const bundleCommandScript = "command.sh" // codeSourcePathVar is the environment variable the AIR run harness sets to the // extracted snapshot directory. A bundle delivers code via `bundle sync` and never -// sets it, so a command relying on it would break at runtime — the convertibility -// gate rejects such commands. +// sets it, so the gate rejects commands relying on it. const codeSourcePathVar = "$CODE_SOURCE_PATH" // workspaceFilePathRef is the bundle variable that resolves to where `bundle // deploy` syncs this folder; the emitted command_path points under it. const workspaceFilePathRef = "${workspace.file_path}" -// aiRuntimeEnvVarsKey is the key linking the task's environment_variables_key to -// the single job-level env-var profile the converter emits. Mirrors -// AI_RUNTIME_ENV_VARS_KEY in the Python CLI (common Jobs env-var API, API-293). +// aiRuntimeEnvVarsKey links the task's environment_variables_key to the single +// job-level env-var profile the converter emits. const aiRuntimeEnvVarsKey = "default" // checkBundleConvertible reports why a structurally-valid runConfig cannot be -// converted to a faithful bundle, or nil if it can. Every reason names the source -// field and why it can't be represented, so the CLI can reject with an actionable -// message instead of emitting a lossy databricks.yml. +// converted to a faithful bundle, or nil if it can. Each reason names the source +// field so the CLI can reject with an actionable message rather than emit a lossy +// databricks.yml. env_variables/secrets are representable (see envVarProfiles) and +// are not rejected here. func checkBundleConvertible(cfg *runConfig) error { var reasons []string - // env_variables / secrets ARE now representable: they ride the common Jobs - // env-var API (API-293) as a job-level environment_variables profile the task - // references by key. convertToBundle emits that profile, so no rejection here. - // (Verified on staging: the profile persists on runs/get; secrets are emitted - // as {{secrets/scope/key}} refs resolved by Jobs at run time.) - // docker_image: a custom image must be registered before it can be referenced; // that registration is not part of a bundle deploy. if cfg.dockerImageURL() != "" { @@ -71,11 +57,21 @@ func checkBundleConvertible(cfg *runConfig) error { reasons = append(reasons, "usage_policy_id: budget policy binding is not represented on the ai_runtime_task bundle path yet") } - // code_source with a git ref: bundle sync uploads the working tree; it cannot - // pin to a specific commit or fetch a remote branch (the R1/R3 gap in the - // uploads-vs-DABs design doc). Dropping the pin would silently change what runs. - if cfg.CodeSource != nil && cfg.CodeSource.Snapshot != nil && cfg.CodeSource.Snapshot.Git != nil { - reasons = append(reasons, "code_source.snapshot.git: bundle sync cannot pin to a git commit or fetch a remote branch") + // code_source snapshots are delivered by uploading the working tree as an + // immutable-folder snapshot (see writeBundleProject). Two sub-cases can't be + // represented that way: + if cfg.CodeSource != nil && cfg.CodeSource.Snapshot != nil { + snap := cfg.CodeSource.Snapshot + // git ref: the snapshot uploads the live working tree; it cannot pin to a + // commit or fetch a remote branch. Dropping the pin would change what runs. + if snap.Git != nil { + reasons = append(reasons, "code_source.snapshot.git: the immutable-folder snapshot uploads the working tree and cannot pin a git commit or fetch a remote branch") + } + // remote_volume: the immutable-folder snapshot uploads to Workspace Files + // only, not a UC Volume. + if snap.RemoteVolume != nil { + reasons = append(reasons, "code_source.snapshot.remote_volume: the immutable-folder snapshot uploads to Workspace Files, not a UC Volume") + } } // A command that reads $CODE_SOURCE_PATH assumes the AIR run harness, which a @@ -98,14 +94,23 @@ func checkBundleConvertible(cfg *runConfig) error { // name plus one job with a single ai_runtime_task. It marshals to YAML, so field // order here is the emitted key order. type exportedBundle struct { - Bundle bundleBlock `yaml:"bundle"` - Resources exportedResourcesBlock `yaml:"resources"` + Bundle bundleBlock `yaml:"bundle"` + Resources exportedResourcesBlock `yaml:"resources"` + Experimental *exportedExperimental `yaml:"experimental,omitempty"` } type bundleBlock struct { Name string `yaml:"name"` } +// exportedExperimental carries experimental.immutable_folder. air run uploads the +// synced code as a single content-addressed snapshot (/api/2.0/repos/snapshots) +// rather than per-file — the mechanism that mirrors AIR's own zip-and-fingerprint +// model. Requires the direct deployment engine, which the run path uses. +type exportedExperimental struct { + ImmutableFolder bool `yaml:"immutable_folder"` +} + type exportedResourcesBlock struct { Jobs map[string]exportedJob `yaml:"jobs"` } @@ -114,28 +119,26 @@ type exportedJob struct { Name string `yaml:"name"` Tasks []exportedTask `yaml:"tasks"` Environments []exportedEnvironment `yaml:"environments"` - // EnvironmentVariables carries env-var profiles via the common Jobs env-var API - // (API-293). Emitted only when the run declares env_variables/secrets; the task - // references a profile by key (see exportedTask.EnvironmentVariablesKey). + // EnvironmentVariables holds env-var profiles, referenced by a task's + // EnvironmentVariablesKey. Emitted only when the run declares env_variables or + // secrets. EnvironmentVariables []exportedEnvVarProfile `yaml:"environment_variables,omitempty"` } type exportedTask struct { TaskKey string `yaml:"task_key"` EnvironmentKey string `yaml:"environment_key"` - // EnvironmentVariablesKey references a job-level environment_variables profile - // (common Jobs env-var API). Omitted when the run has no env vars/secrets so the - // wire form matches the submit path, which sets it only under the same gate. + // EnvironmentVariablesKey references a job-level environment_variables profile. + // Omitted when the run has no env vars or secrets. EnvironmentVariablesKey string `yaml:"environment_variables_key,omitempty"` MaxRetries int `yaml:"max_retries"` TimeoutSeconds int `yaml:"timeout_seconds,omitempty"` AiRuntimeTask exportedAiRuntimeTask `yaml:"ai_runtime_task"` } -// exportedEnvVarProfile is one entry in the job-level environment_variables list -// (common Jobs env-var API). variables holds plain values inline and secrets as -// {{secrets/scope/key}} references, resolved by Jobs at run time — the exact shape -// the ai_runtime_task submit path emits (jobs_api_client.py, API-293). +// exportedEnvVarProfile is one entry in the job-level environment_variables list. +// Variables holds plain values inline and secrets as {{secrets/scope/key}} +// references, resolved by Jobs at run time. type exportedEnvVarProfile struct { EnvironmentVariablesKey string `yaml:"environment_variables_key"` Variables map[string]string `yaml:"variables"` @@ -175,28 +178,26 @@ type exportedEnvSpec struct { } // databricksAITokenPrefix marks an environment.version that selects the -// databricks-ai managed base environment rather than a bare channel. Mirrors -// _DATABRICKS_AI_TOKEN_PREFIX in the Python CLI's jobs_api_client.py. +// databricks-ai managed base environment rather than a bare channel. const databricksAITokenPrefix = "databricks_ai_v" // databricksAIBaseEnvironment is the system base_environment id the databricks-ai -// token resolves to. Mirrors _DATABRICKS_AI_BASE_ENVIRONMENT in the Python CLI. +// token resolves to. const databricksAIBaseEnvironment = "workspace-base-environments/" -// convertToBundle maps a convertible runConfig to the emitted bundle. It assumes -// checkBundleConvertible has already passed. The command_path points at the synced -// command.sh under ${workspace.file_path}; the runtime channel is resolved the same -// way `air run` resolves it (config version, else the default), minus the process -// env lookup so a generated artifact is reproducible rather than host-dependent. +// convertToBundle maps a convertible runConfig to the emitted bundle, assuming +// checkBundleConvertible has already passed. command_path is emitted as a path +// relative to the bundle root; the bundle's translate_paths mutator rewrites it to +// the deployed location (for immutable_folder, under the content-addressed +// snapshot). Emitting ${workspace.file_path} directly would instead be validated as +// a local file and fail. func convertToBundle(cfg *runConfig) *exportedBundle { task := exportedAiRuntimeTask{ Experiment: cfg.ExperimentName, Deployments: []exportedDeployment{{ - // `air run` submits a single unnamed deployment; the bundle names it - // "worker" so the authored YAML reads clearly. The name is cosmetic (the - // submit payload carries none), so this does not change behavior. + // The deployment name is cosmetic (the wire payload carries none). Name: "worker", - CommandPath: workspaceFilePathRef + "/" + bundleCommandScript, + CommandPath: bundleCommandScript, Compute: exportedCompute{ AcceleratorType: cfg.Compute.AcceleratorType, AcceleratorCount: cfg.Compute.NumAccelerators, @@ -218,16 +219,15 @@ func convertToBundle(cfg *runConfig) *exportedBundle { AiRuntimeTask: task, } - // Env vars ride the common Jobs env-var API: one job-level profile, referenced - // by the task's environment_variables_key. Emitted only when there are vars or - // secrets, so a var-free run's wire form matches the submit path exactly. + // One job-level env-var profile, referenced by the task's + // environment_variables_key. Emitted only when there are vars or secrets. profiles := envVarProfiles(cfg) if len(profiles) > 0 { tsk.EnvironmentVariablesKey = aiRuntimeEnvVarsKey } - // Resource key shares the task_key charset, which validateExperimentName - // already guarantees, so the experiment name is safe to use verbatim. + // The experiment name is safe as a resource key: validateExperimentName already + // guarantees the task_key charset. return &exportedBundle{ Bundle: bundleBlock{Name: cfg.ExperimentName}, Resources: exportedResourcesBlock{ @@ -243,15 +243,14 @@ func convertToBundle(cfg *runConfig) *exportedBundle { }, }, }, + Experimental: &exportedExperimental{ImmutableFolder: true}, } } // envVarProfiles builds the job-level env-var profile list from the run's // env_variables and secrets, or nil when there are none. Plain values are inline; // each secret (ENV_VAR -> "scope/key") becomes a {{secrets/scope/key}} reference -// Jobs resolves at run time. This mirrors the Python CLI's ai_runtime_task path -// (jobs_api_client.py, API-293) and was verified against staging: the profile -// persists on runs/get and secret refs resolve at run time. +// Jobs resolves at run time. func envVarProfiles(cfg *runConfig) []exportedEnvVarProfile { if len(cfg.EnvVariables) == 0 && len(cfg.Secrets) == 0 { return nil @@ -269,23 +268,11 @@ func envVarProfiles(cfg *runConfig) []exportedEnvVarProfile { }} } -// exportBundleEnvSpec resolves the serverless runtime selection for the bundle, -// branching on environment.version the same way the Python CLI's jobs_api_client -// does (jobs_api_client.py:1562-1565): a "databricks_ai_v" token selects the -// managed databricks-ai base_environment (torch + ML venv), while a bare numeric -// channel ("4", "5", ...) uses environment_version. Unlike dlRuntimeImage it does -// not read process env, so the generated bundle is reproducible. A requirements-file -// dependency set carries its version in the file, so this falls back to the default -// channel there (same as the submit path). -// -// TODO(air): the Go `air run` submit path (runsubmit.go) only ever emits -// environment_version and cannot select base_environment, so a `version: -// databricks_ai_v` run lands on a bare GPU channel without torch/mlflow and -// fails at import. This converter hardcodes the correct base_environment branch so -// exported bundles run today; once `air run` forwards env vars/dependencies through -// the new BYOT common Jobs API (env_file/env_var work, ETA EOQ2), the submit path -// and this converter should share one runtime-selection helper instead of -// duplicating the branch. +// exportBundleEnvSpec resolves the serverless runtime selection: a +// "databricks_ai_v" token selects the managed databricks-ai base_environment +// (torch + ML venv), while a bare numeric channel ("4", "5", ...) uses +// environment_version. It does not read process env, so the generated bundle is +// reproducible. func exportBundleEnvSpec(cfg *runConfig) exportedEnvSpec { channel := strings.TrimPrefix(defaultDlRuntimeImage, "CLIENT-GPU-") if v, ok := cfg.runtimeVersion(); ok { diff --git a/experimental/air/cmd/exportbundle_test.go b/experimental/air/cmd/exportbundle_test.go index e7964d4789d..031d2083123 100644 --- a/experimental/air/cmd/exportbundle_test.go +++ b/experimental/air/cmd/exportbundle_test.go @@ -1,6 +1,8 @@ package aircmd import ( + "os" + "path/filepath" "strings" "testing" @@ -31,7 +33,9 @@ func TestConvertToBundleBasic(t *testing.T) { require.Len(t, task.AiRuntimeTask.Deployments, 1) dep := task.AiRuntimeTask.Deployments[0] - assert.Equal(t, workspaceFilePathRef+"/"+bundleCommandScript, dep.CommandPath) + // command_path is relative to the bundle root; translate_paths rewrites it to + // the deployed location. + assert.Equal(t, bundleCommandScript, dep.CommandPath) assert.Equal(t, "GPU_8xH100", dep.Compute.AcceleratorType) assert.Equal(t, 16, dep.Compute.AcceleratorCount) @@ -93,9 +97,67 @@ func TestCheckBundleConvertibleRejectsCodeSourcePathCommand(t *testing.T) { assert.Contains(t, err.Error(), codeSourcePathVar) } +func TestCheckBundleConvertibleCodeSource(t *testing.T) { + base := func() *runConfig { + return &runConfig{ + ExperimentName: "exp", + Command: new("python train.py"), + Compute: &computeConfig{AcceleratorType: "GPU_1xA10", NumAccelerators: 1}, + } + } + + // A plain working-tree snapshot is convertible: it uploads via immutable folder. + ok := base() + ok.CodeSource = &codeSourceConfig{Type: "snapshot", Snapshot: &snapshotSourceConfig{RootPath: "."}} + assert.NoError(t, checkBundleConvertible(ok)) + + // A git-pinned snapshot can't be represented (working-tree upload only). + git := base() + git.CodeSource = &codeSourceConfig{Type: "snapshot", Snapshot: &snapshotSourceConfig{RootPath: ".", Git: &gitRef{Commit: new("abc123")}}} + err := checkBundleConvertible(git) + require.Error(t, err) + assert.Contains(t, err.Error(), "git") + + // A UC Volume destination isn't supported (Workspace Files only). + vol := base() + vol.CodeSource = &codeSourceConfig{Type: "snapshot", Snapshot: &snapshotSourceConfig{RootPath: ".", RemoteVolume: new("/Volumes/c/s/v")}} + err = checkBundleConvertible(vol) + require.Error(t, err) + assert.Contains(t, err.Error(), "remote_volume") +} + +func TestStageCodeSource(t *testing.T) { + // A working-tree snapshot copies the tree into the bundle root; include_paths + // restricts to the named subpaths. + src := t.TempDir() + require.NoError(t, os.WriteFile(filepath.Join(src, "train.py"), []byte("print()"), 0o644)) + require.NoError(t, os.MkdirAll(filepath.Join(src, "pkg"), 0o755)) + require.NoError(t, os.WriteFile(filepath.Join(src, "pkg", "mod.py"), []byte("x=1"), 0o644)) + require.NoError(t, os.WriteFile(filepath.Join(src, "ignore.txt"), []byte("no"), 0o644)) + + t.Run("whole tree", func(t *testing.T) { + dst := t.TempDir() + snap := &snapshotSourceConfig{RootPath: src} + require.NoError(t, stageCodeSource(t.Context(), snap, "train.yaml", dst)) + assert.FileExists(t, filepath.Join(dst, "train.py")) + assert.FileExists(t, filepath.Join(dst, "pkg", "mod.py")) + assert.FileExists(t, filepath.Join(dst, "ignore.txt")) + }) + + t.Run("include_paths only", func(t *testing.T) { + dst := t.TempDir() + snap := &snapshotSourceConfig{RootPath: src, IncludePaths: []string{"train.py", "pkg"}} + require.NoError(t, stageCodeSource(t.Context(), snap, "train.yaml", dst)) + assert.FileExists(t, filepath.Join(dst, "train.py")) + assert.FileExists(t, filepath.Join(dst, "pkg", "mod.py")) + assert.NoFileExists(t, filepath.Join(dst, "ignore.txt")) + }) +} + func TestRenderBundleIncludesTargetsAndConvertGate(t *testing.T) { - // renderBundle (what --dry-run --via-dabs shows and what the run path deploys) - // must include the converted job AND the appended dev targets block. + // renderBundle (what --dry-run shows and what the run path deploys) must include + // the converted job, the appended dev targets block, and the immutable_folder + // flag that routes deploy through the content-addressed snapshot path. cfg := &runConfig{ ExperimentName: "exp", Command: new("python train.py"), @@ -108,6 +170,7 @@ func TestRenderBundleIncludesTargetsAndConvertGate(t *testing.T) { assert.Contains(t, out, "FOO: bar") assert.Contains(t, out, "targets:") assert.Contains(t, out, "mode: development") + assert.Contains(t, out, "immutable_folder: true") // The convertibility gate still applies: an unconvertible config errors instead // of rendering a lossy bundle. diff --git a/experimental/air/cmd/get.go b/experimental/air/cmd/get.go index 7f94cfbdbd9..2f8bd8dd09e 100644 --- a/experimental/air/cmd/get.go +++ b/experimental/air/cmd/get.go @@ -91,9 +91,6 @@ func authError(ctx context.Context, cmd *cobra.Command, err error) error { // newGetCommand returns the `air get JOB_RUN_ID` command, which shows status, // configuration, and timing details for a specific run. -// get works unchanged for DABs-submitted runs: it resolves by run_id via -// Jobs.GetRun, and a `bundle run` produces a normal Jobs run with a run_id. No -// submit-verb-specific handling is needed. func newGetCommand() *cobra.Command { cmd := &cobra.Command{ Use: "get JOB_RUN_ID", diff --git a/experimental/air/cmd/list.go b/experimental/air/cmd/list.go index 2a64f4118bb..b9672aad796 100644 --- a/experimental/air/cmd/list.go +++ b/experimental/air/cmd/list.go @@ -233,11 +233,9 @@ type jobsScanStrategy struct { } func newJobsScanStrategy(ctx context.Context, w *databricks.WorkspaceClient, q listQuery) *jobsScanStrategy { - // AIR runs are now runs OF a persistent DABs job (RunType JOB_RUN), not the - // ephemeral SUBMIT_RUN the retired runs/submit path produced. A SUBMIT_RUN - // server filter would hide them, so we leave RunType unset and let isAirRun - // (task-shape) select AIR runs. Leaving it unset also still surfaces any - // pre-migration ephemeral SUBMIT_RUNs that remain in the workspace. + // AIR runs are runs OF a persistent DABs job (RunType JOB_RUN), so a SUBMIT_RUN + // server filter would hide them. RunType is left unset and isAirRun (task-shape) + // selects AIR runs; this also still surfaces any older ephemeral SUBMIT_RUNs. req := jobs.ListRunsRequest{ ExpandTasks: true, Limit: jobsPageLimit, diff --git a/experimental/air/cmd/run.go b/experimental/air/cmd/run.go index def02a785d9..e9e1c67c6f2 100644 --- a/experimental/air/cmd/run.go +++ b/experimental/air/cmd/run.go @@ -75,8 +75,7 @@ The workload is described by a YAML config file (see --file).`, } if dryRun { - // A dry run shows the generated bundle so the user can see the artifact - // we'd deploy on their behalf (transparency); it does not deploy. + // Show the generated bundle without deploying. bundleYAML, err := renderBundle(cfg, file) if err != nil { return err diff --git a/experimental/air/cmd/rundabs.go b/experimental/air/cmd/rundabs.go index 030dc9c7802..05798d1a0e1 100644 --- a/experimental/air/cmd/rundabs.go +++ b/experimental/air/cmd/rundabs.go @@ -1,34 +1,24 @@ package aircmd // air run submits a training workload by converting the run YAML to a Databricks -// Asset Bundle and driving `bundle deploy` + `bundle run` in-process. This is the -// ONLY submit path: per the AIR-CLI/DABs agreement with Pieter Noordhuis (DECO -// requires AIR to integrate with DABs and submit through the public Jobs API, not -// an internal proxy), `air run` is a high-level wrapper of -// [convert train.yaml -> databricks.yml -> deploy -> run]. +// Asset Bundle and driving deploy + run in-process: convert train.yaml -> +// databricks.yml -> deploy -> run. It reuses the exportbundle.go converter, so the +// deployed ai_runtime_task is exactly what `air export-bundle` emits. // -// The bundle is a persistent Jobs resource, so `air list` finds it via the indexed -// jobs-list path and re-runs anchor to a standing job. It reuses the exportbundle.go -// converter verbatim, so the deployed ai_runtime_task — including env vars via the -// common Jobs env-var API — is exactly what `air export-bundle` emits. +// Deploy and run call the bundle libraries directly (cmd/bundle/utils.ProcessBundle +// + Jobs.RunNow), not a child `databricks` process. Shelling out would risk an +// older `databricks` on PATH dropping the unknown ai_runtime_task field and +// deploying a task-less job. // -// IN-PROCESS, NOT SHELL-OUT: deploy and run call the bundle libraries directly -// (cmd/bundle/utils.ProcessBundle + bundle/run), the same entry points the -// `databricks bundle` commands use and the pattern the pipelines CLI established -// (cmd/pipelines). There is no child `databricks` process, so the deploy always -// uses this build's ai_runtime_task-aware bundle schema — an older `databricks` on -// PATH would drop the unknown ai_runtime_task field and deploy a task-less job. -// -// KNOWN TRADEOFF (design doc): the deployed job is PERSISTENT and is NOT swept by -// the Jobs ephemeral-job GC (JobsSoftDeletion only sweeps EPHEMERAL/WORKFLOW -// types). Runs of distinct experiments accumulate distinct jobs, drifting toward -// the per-workspace saved-jobs cap. Job reuse (a stable per-experiment bundle name, -// so re-runs update one job) mitigates it; a cleanup/TTL story is future work. +// The deployed job is persistent and is NOT swept by the Jobs ephemeral-job GC, so +// runs of distinct experiments accumulate distinct jobs, drifting toward the +// per-workspace saved-jobs cap. A cleanup/TTL story is future work. import ( "context" "errors" "fmt" + "io" "os" "path/filepath" "strconv" @@ -66,33 +56,27 @@ func submitWorkload(ctx context.Context, w *databricks.WorkspaceClient, cfg *run } // Idempotency does not map onto deploy+run (two calls, not one tokened submit). - // Note it rather than silently honoring it; job identity (stable bundle name) - // is the DABs-native dedupe mechanism. if idempotencyKey != "" || cfg.IdempotencyToken != nil { cmdio.LogString(ctx, "note: --idempotency-key is ignored on the DABs path (deploy+run is not a single idempotent call)") } - bundleRoot, cleanup, err := writeBundleProject(cfg, configPath) + bundleRoot, cleanup, err := writeBundleProject(ctx, cfg, configPath) if err != nil { return 0, "", err } defer cleanup() - // Surface the generated artifact's path (transparency): the user can inspect - // what we deploy on their behalf. The dir is temporary (auto-removed on exit). - // For the full contents without deploying, use `air run --dry-run`. + // Surface the generated artifact's path so the user can inspect what is + // deployed. The dir is temporary; `air run --dry-run` shows the full contents. cmdio.LogString(ctx, "Generated bundle (temporary): "+filepath.Join(bundleRoot, "databricks.yml")) - // deploy: creates/updates the persistent job. This is the step that adds - // latency vs the old ephemeral submit (~2-3s warm, ~12s first-of-session). cmdio.LogString(ctx, "Deploying bundle (creates a persistent job)...") b, err := deployBundle(ctx, w, bundleRoot) if err != nil { return 0, "", fmt.Errorf("bundle deploy: %w", err) } - // run: RunNow on the just-deployed job, returning as soon as the run is created - // (no wait), matching the old submit path's fire-and-return. + // RunNow on the just-deployed job, returning as soon as the run is created. cmdio.LogString(ctx, "Triggering run...") runID, err := runDeployedJob(ctx, w, b, cfg.ExperimentName) if err != nil { @@ -106,10 +90,9 @@ func submitWorkload(ctx context.Context, w *databricks.WorkspaceClient, cfg *run } // renderBundle produces the exact databricks.yml the run path deploys: the -// exportbundle.go converter output plus the dev targets block the run path appends. -// It touches no filesystem and does not deploy, so `air run --dry-run` can show the -// user the artifact we'd generate on their behalf (transparency), and -// writeBundleProject reuses it so preview and real run never diverge. +// converter output plus the dev targets block. It touches no filesystem and does +// not deploy, so `air run --dry-run` can show the artifact; writeBundleProject +// reuses it so preview and real run never diverge. func renderBundle(cfg *runConfig, configPath string) (string, error) { if err := checkBundleConvertible(cfg); err != nil { return "", err @@ -124,14 +107,11 @@ func renderBundle(cfg *runConfig, configPath string) (string, error) { return string(body) + bundleTargetsBlock(), nil } -// writeBundleProject renders databricks.yml (via renderBundle) plus the command.sh -// the bundle syncs, into a temp bundle root. Returns the root and a cleanup func. -// -// TODO(air): reuse the code-snapshot staging (snapshot.go) so the user's code tree -// lands in the bundle folder for `bundle sync` to upload. This writes command.sh -// from cfg.Command; wiring code_source snapshotting into the synced folder is the -// remaining piece before code_source runs are supported on this path. -func writeBundleProject(cfg *runConfig, configPath string) (string, func(), error) { +// writeBundleProject renders databricks.yml (via renderBundle) plus command.sh and, +// for a code_source snapshot, the user's code tree, into a temp bundle root. Deploy +// uploads the whole root as an immutable-folder snapshot. Returns the root and a +// cleanup func. +func writeBundleProject(ctx context.Context, cfg *runConfig, configPath string) (string, func(), error) { body, err := renderBundle(cfg, configPath) if err != nil { return "", func() {}, err @@ -142,12 +122,21 @@ func writeBundleProject(cfg *runConfig, configPath string) (string, func(), erro } cleanup := func() { _ = os.RemoveAll(bundleRoot) } + // Copy the code_source working tree into the bundle root first, so a stray + // command.sh / databricks.yml in the user's tree can't shadow ours below. + if cfg.CodeSource != nil && cfg.CodeSource.Snapshot != nil { + if err := stageCodeSource(ctx, cfg.CodeSource.Snapshot, configPath, bundleRoot); err != nil { + cleanup() + return "", func() {}, err + } + } + if err := os.WriteFile(filepath.Join(bundleRoot, "databricks.yml"), []byte(body), 0o600); err != nil { cleanup() return "", func() {}, err } - // command.sh: the entrypoint the task's command_path points at (synced by deploy). + // command.sh: the entrypoint the task's command_path points at. if cfg.Command != nil { if err := os.WriteFile(filepath.Join(bundleRoot, bundleCommandScript), []byte(*cfg.Command), 0o600); err != nil { cleanup() @@ -157,6 +146,75 @@ func writeBundleProject(cfg *runConfig, configPath string) (string, func(), erro return bundleRoot, cleanup, nil } +// stageCodeSource copies the snapshot's working tree into dest so deploy uploads it +// as part of the immutable-folder snapshot. When include_paths is set, only those +// paths (files or directories, relative to root_path) are copied; otherwise the +// whole tree is copied. checkBundleConvertible has already rejected git and +// remote_volume snapshots, so this only handles the local working tree. +func stageCodeSource(ctx context.Context, snap *snapshotSourceConfig, configPath, dest string) error { + root, err := resolveRootPath(ctx, snap.RootPath, filepath.Dir(configPath)) + if err != nil { + return err + } + if len(snap.IncludePaths) == 0 { + return copyTree(root, dest) + } + for _, rel := range snap.IncludePaths { + if err := copyTree(filepath.Join(root, rel), filepath.Join(dest, rel)); err != nil { + return err + } + } + return nil +} + +// copyTree recursively copies the file or directory at src to dst, preserving the +// relative layout and file modes. Symlinks are skipped. +func copyTree(src, dst string) error { + return filepath.WalkDir(src, func(p string, d os.DirEntry, err error) error { + if err != nil { + return err + } + rel, err := filepath.Rel(src, p) + if err != nil { + return err + } + target := filepath.Join(dst, rel) + info, err := d.Info() + if err != nil { + return err + } + switch { + case d.IsDir(): + return os.MkdirAll(target, 0o755) + case info.Mode()&os.ModeSymlink != 0: + return nil + default: + if err := os.MkdirAll(filepath.Dir(target), 0o755); err != nil { + return err + } + return copyFile(p, target, info.Mode()) + } + }) +} + +// copyFile copies a single regular file from src to dst with the given mode. +func copyFile(src, dst string, mode os.FileMode) error { + in, err := os.Open(src) + if err != nil { + return err + } + defer in.Close() + out, err := os.OpenFile(dst, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, mode) + if err != nil { + return err + } + if _, err := io.Copy(out, in); err != nil { + out.Close() + return err + } + return out.Close() +} + // bundleTargetsBlock is the minimal dev target appended to the converted YAML so // `bundle deploy` runs non-interactively. dev mode gives per-user name prefix + // isolation. The workspace host comes from the active CLI profile, so it's omitted @@ -166,30 +224,19 @@ func bundleTargetsBlock() string { } // deployBundle deploys the bundle rooted at bundleRoot in-process and returns the -// configured *bundle.Bundle (with resource IDs populated, so the caller can resolve -// the deployed job). It reuses cmd/bundle/utils.ProcessBundle — the exact -// orchestration `databricks bundle deploy` runs (load, validate, build, sync, -// deploy) — via a synthetic carrier command: +// configured *bundle.Bundle with resource IDs populated, so the caller can resolve +// the deployed job. It reuses ProcessBundle (load, validate, build, sync, deploy) +// via a synthetic carrier command (see newBundleCarrierCommand). // -// - air's own cobra command carries none of the flags ProcessBundle reads -// (--var, --target, --profile, --output), so we build a throwaway command that -// declares them and seed the context (logdiag + auth profile + bundle root). -// - the bundle root is steered to bundleRoot with DATABRICKS_BUNDLE_ROOT on the -// context (env.Set is context-scoped, not a real env var), so MustConfigureBundle -// loads our generated databricks.yml instead of walking cwd. -// - the direct engine (DATABRICKS_BUNDLE_ENGINE=direct) avoids the Terraform -// provider/registry dependency; it is the GA default on new CLIs and set -// explicitly here so the deploy is self-contained. +// SkipInitContext: air's command context already has logdiag initialized (root's +// PersistentPreRunE); re-initializing panics ("InitContext twice"). func deployBundle(ctx context.Context, w *databricks.WorkspaceClient, bundleRoot string) (*bundle.Bundle, error) { cmd := newBundleCarrierCommand(ctx, w, bundleRoot) b, err := utils.ProcessBundle(cmd, utils.ProcessOptions{ - FastValidate: true, - Build: true, - Deploy: true, - // air's command context already has logdiag initialized (root's - // PersistentPreRunE); re-initializing panics ("InitContext twice"). Same as - // the pipelines CLI run path. + FastValidate: true, + Build: true, + Deploy: true, SkipInitContext: true, }) if err != nil { @@ -201,30 +248,27 @@ func deployBundle(ctx context.Context, w *databricks.WorkspaceClient, bundleRoot return b, nil } -// newBundleCarrierCommand builds the synthetic cobra command deployBundle hands to -// ProcessBundle. It declares the flags ProcessBundle/configureBundle read and seeds -// the context so the deploy runs with air's resolved auth against the generated -// bundle root. The command is never added to a tree and never Execute()d — it only -// carries flags + context into the bundle libraries. +// newBundleCarrierCommand builds a throwaway cobra command for ProcessBundle. air's +// own command carries none of the flags ProcessBundle reads (--var, --target, +// --profile, --output), so this declares them and seeds the context with air's auth +// profile and the bundle root. DATABRICKS_BUNDLE_ROOT (context-scoped via env.Set) +// points the loader at the generated databricks.yml instead of walking cwd; the +// direct engine avoids the Terraform provider/registry dependency. The command is +// never added to a tree or Execute()d. func newBundleCarrierCommand(ctx context.Context, w *databricks.WorkspaceClient, bundleRoot string) *cobra.Command { cmd := &cobra.Command{Use: "air-bundle-deploy"} - // Flags ProcessBundleRet (and the bundle root loader) read by name. cmd.Flags().StringSlice("var", nil, "") cmd.Flags().StringP("target", "t", dabsTarget, "") cmd.Flags().StringP("profile", "p", "", "") outputFlag := flags.OutputText cmd.Flags().Var(&outputFlag, "output", "") - // Forward air's active profile so the bundle authenticates the same way air - // did, instead of falling back to the default profile. + // Forward air's active profile so the bundle authenticates the same way air did. if w.Config.Profile != "" { _ = cmd.Flags().Set("profile", w.Config.Profile) } - // logdiag is how the bundle libraries report diagnostics; ProcessBundle asserts - // it is set up (we pass SkipInitContext=false, so it initializes, but the auth - // profile + bundle root must be seeded before that on the same context). ctx = env.Set(ctx, "DATABRICKS_BUNDLE_ROOT", bundleRoot) ctx = env.Set(ctx, "DATABRICKS_BUNDLE_ENGINE", "direct") cmd.SetContext(ctx) @@ -233,10 +277,9 @@ func newBundleCarrierCommand(ctx context.Context, w *databricks.WorkspaceClient, // runDeployedJob triggers a run of the just-deployed job and returns the new run_id // without waiting for completion. It resolves the job by its bundle resource key -// (the experiment name) to get the server-assigned job_id that deploy populated, -// then calls Jobs.RunNow. AIR runs take no run parameters, so RunNow needs only the -// job_id. (We call RunNow directly rather than the bundle runner's NoWait path, -// which returns nil instead of the run_id.) +// (the experiment name) to get the job_id deploy populated, then calls Jobs.RunNow. +// AIR runs take no run parameters. RunNow is called directly rather than via the +// bundle runner's NoWait path, which returns nil instead of the run_id. func runDeployedJob(ctx context.Context, w *databricks.WorkspaceClient, b *bundle.Bundle, experimentName string) (int64, error) { ref, err := bundleresources.Lookup(b, experimentName, isRunnableJob) if err != nil { @@ -258,9 +301,7 @@ func runDeployedJob(ctx context.Context, w *databricks.WorkspaceClient, b *bundl return wait.RunId, nil } -// isRunnableJob filters resources.Lookup to job resources, so an experiment name -// that also matched another resource type can't be selected. AIR bundles only ever -// contain one job, so this is a safety guard rather than a disambiguator. +// isRunnableJob filters resources.Lookup to job resources. func isRunnableJob(ref bundleresources.Reference) bool { _, ok := ref.Resource.(*resources.Job) return ok From e05d89e3cc5d06bc03bcc5a0e0ac8ea805ddd2ab Mon Sep 17 00:00:00 2001 From: vinchenzo-db Date: Wed, 15 Jul 2026 21:19:12 +0000 Subject: [PATCH 04/11] [air] Add unit tests for the DABs deploy-path helpers Covers the parts of rundabs.go that don't need a live workspace: - writeBundleProject: databricks.yml + command.sh land, immutable_folder is set, code_source tree is staged, a user command.sh doesn't shadow the generated one, and the convertibility gate fails before any temp dir is created. - newBundleCarrierCommand: forwards air's profile, declares the flags ProcessBundle reads, and seeds bundle-root + direct engine on the context. - bundleTargetsBlock and isRunnableJob. deployBundle / runDeployedJob still need a live deploy+RunNow and stay covered by staging validation. Co-authored-by: Isaac --- experimental/air/cmd/rundabs_test.go | 144 +++++++++++++++++++++++++++ 1 file changed, 144 insertions(+) create mode 100644 experimental/air/cmd/rundabs_test.go diff --git a/experimental/air/cmd/rundabs_test.go b/experimental/air/cmd/rundabs_test.go new file mode 100644 index 00000000000..7b564e74169 --- /dev/null +++ b/experimental/air/cmd/rundabs_test.go @@ -0,0 +1,144 @@ +package aircmd + +import ( + "os" + "path/filepath" + "testing" + + "github.com/databricks/cli/bundle/config/resources" + bundleresources "github.com/databricks/cli/bundle/resources" + "github.com/databricks/cli/libs/env" + "github.com/databricks/databricks-sdk-go" + "github.com/databricks/databricks-sdk-go/config" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestWriteBundleProject(t *testing.T) { + cfg := &runConfig{ + ExperimentName: "exp", + Command: new("echo hi"), + Compute: &computeConfig{AcceleratorType: "GPU_1xA10", NumAccelerators: 1}, + } + + root, cleanup, err := writeBundleProject(t.Context(), cfg, "train.yaml") + require.NoError(t, err) + defer cleanup() + + // databricks.yml carries the converted job and the appended dev target. + body, err := os.ReadFile(filepath.Join(root, "databricks.yml")) + require.NoError(t, err) + assert.Contains(t, string(body), "ai_runtime_task:") + assert.Contains(t, string(body), "immutable_folder: true") + assert.Contains(t, string(body), "mode: development") + + // command.sh holds the run's command verbatim. + script, err := os.ReadFile(filepath.Join(root, bundleCommandScript)) + require.NoError(t, err) + assert.Equal(t, "echo hi", string(script)) + + // cleanup removes the temp root. + cleanup() + _, err = os.Stat(root) + assert.True(t, os.IsNotExist(err)) +} + +func TestWriteBundleProjectStagesCodeSource(t *testing.T) { + src := t.TempDir() + require.NoError(t, os.WriteFile(filepath.Join(src, "train.py"), []byte("print()"), 0o644)) + require.NoError(t, os.MkdirAll(filepath.Join(src, "pkg"), 0o755)) + require.NoError(t, os.WriteFile(filepath.Join(src, "pkg", "mod.py"), []byte("x=1"), 0o644)) + + cfg := &runConfig{ + ExperimentName: "exp", + Command: new("python train.py"), + Compute: &computeConfig{AcceleratorType: "GPU_1xA10", NumAccelerators: 1}, + CodeSource: &codeSourceConfig{Type: "snapshot", Snapshot: &snapshotSourceConfig{RootPath: src}}, + } + + root, cleanup, err := writeBundleProject(t.Context(), cfg, "train.yaml") + require.NoError(t, err) + defer cleanup() + + // The user's tree is copied into the bundle root alongside our generated files. + assert.FileExists(t, filepath.Join(root, "train.py")) + assert.FileExists(t, filepath.Join(root, "pkg", "mod.py")) + assert.FileExists(t, filepath.Join(root, "databricks.yml")) + assert.FileExists(t, filepath.Join(root, bundleCommandScript)) +} + +func TestWriteBundleProjectCommandShadowProtection(t *testing.T) { + // A command.sh in the user's tree must not shadow the one we generate from the + // run's command: our write happens after the tree copy. + src := t.TempDir() + require.NoError(t, os.WriteFile(filepath.Join(src, bundleCommandScript), []byte("STALE"), 0o644)) + + cfg := &runConfig{ + ExperimentName: "exp", + Command: new("FRESH"), + Compute: &computeConfig{AcceleratorType: "GPU_1xA10", NumAccelerators: 1}, + CodeSource: &codeSourceConfig{Type: "snapshot", Snapshot: &snapshotSourceConfig{RootPath: src}}, + } + + root, cleanup, err := writeBundleProject(t.Context(), cfg, "train.yaml") + require.NoError(t, err) + defer cleanup() + + script, err := os.ReadFile(filepath.Join(root, bundleCommandScript)) + require.NoError(t, err) + assert.Equal(t, "FRESH", string(script)) +} + +func TestWriteBundleProjectRejectsUnconvertible(t *testing.T) { + // A gate failure (here: a $CODE_SOURCE_PATH command) surfaces before any temp + // directory is created. + cfg := &runConfig{ + ExperimentName: "exp", + Command: new("cd $CODE_SOURCE_PATH && python train.py"), + Compute: &computeConfig{AcceleratorType: "GPU_1xA10", NumAccelerators: 1}, + } + _, _, err := writeBundleProject(t.Context(), cfg, "train.yaml") + require.Error(t, err) +} + +func TestBundleTargetsBlock(t *testing.T) { + block := bundleTargetsBlock() + assert.Contains(t, block, "targets:") + assert.Contains(t, block, dabsTarget+":") + assert.Contains(t, block, "mode: development") + assert.Contains(t, block, "default: true") +} + +func TestNewBundleCarrierCommand(t *testing.T) { + // Construct the client directly (not via NewWorkspaceClient, which would try to + // resolve "myprofile" from ~/.databrickscfg); the carrier only reads Config.Profile. + w := &databricks.WorkspaceClient{Config: &config.Config{Profile: "myprofile"}} + + cmd := newBundleCarrierCommand(t.Context(), w, "/tmp/bundle-root") + + // air's active profile is forwarded so the bundle authenticates the same way. + assert.Equal(t, "myprofile", cmd.Flag("profile").Value.String()) + + // The flags ProcessBundle reads are declared. + for _, name := range []string{"var", "target", "profile", "output"} { + assert.NotNil(t, cmd.Flag(name), "flag %q must be declared", name) + } + + // The bundle root and direct engine are seeded on the command's context. + assert.Equal(t, "/tmp/bundle-root", env.Get(cmd.Context(), "DATABRICKS_BUNDLE_ROOT")) + assert.Equal(t, "direct", env.Get(cmd.Context(), "DATABRICKS_BUNDLE_ENGINE")) +} + +func TestNewBundleCarrierCommandNoProfile(t *testing.T) { + // With no profile on the client, the profile flag is left empty (the bundle + // falls back to its own default resolution). + w := &databricks.WorkspaceClient{Config: &config.Config{}} + + cmd := newBundleCarrierCommand(t.Context(), w, "/tmp/bundle-root") + assert.Empty(t, cmd.Flag("profile").Value.String()) +} + +func TestIsRunnableJob(t *testing.T) { + assert.True(t, isRunnableJob(bundleresources.Reference{Resource: &resources.Job{}})) + assert.False(t, isRunnableJob(bundleresources.Reference{Resource: &resources.Pipeline{}})) +} From f711157d3e289696d1a987497477e55c3a6bd61e Mon Sep 17 00:00:00 2001 From: vinchenzo-db Date: Wed, 15 Jul 2026 21:58:30 +0000 Subject: [PATCH 05/11] [air] Deliver env vars only via the common Jobs env-var API (no sidecar) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Match the Python CLI (#2237609): env vars and secrets ride the job-level environment_variables profile the converter emits, not the legacy env_vars.json / secret_env_vars.json workspace side-channel. Drop that sidecar staging from buildArtifacts and delete the now-dead entry encoders and filename constants. The other launch artifacts buildArtifacts stages (command.sh, training_config.yaml, requirements.yaml, hyperparameters.yaml) are unchanged — those fixed the "requirements.yaml not found" harness error and are unrelated to env vars. Known gap (verified on staging): the vendored databricks-sdk-go jobs.JobSettings/Task has no environment_variables field, so the DABs bundle schema strips it client-side ("unknown field") before the wire. Env vars therefore aren't delivered on the DABs path until the SDK carries the common Jobs env-var API. The Jobs server already accepts the field (gated by allowEnvironmentVariables); this is purely a DABs/SDK-schema lag. Co-authored-by: Isaac --- experimental/air/cmd/rundabs.go | 24 +++++---- experimental/air/cmd/rundabs_test.go | 69 +++++++++++++++----------- experimental/air/cmd/runupload.go | 59 ++-------------------- experimental/air/cmd/runupload_test.go | 15 +++--- 4 files changed, 64 insertions(+), 103 deletions(-) diff --git a/experimental/air/cmd/rundabs.go b/experimental/air/cmd/rundabs.go index 05798d1a0e1..fefffe37eb9 100644 --- a/experimental/air/cmd/rundabs.go +++ b/experimental/air/cmd/rundabs.go @@ -107,23 +107,28 @@ func renderBundle(cfg *runConfig, configPath string) (string, error) { return string(body) + bundleTargetsBlock(), nil } -// writeBundleProject renders databricks.yml (via renderBundle) plus command.sh and, -// for a code_source snapshot, the user's code tree, into a temp bundle root. Deploy -// uploads the whole root as an immutable-folder snapshot. Returns the root and a -// cleanup func. +// writeBundleProject renders databricks.yml plus the launch artifacts the AI Runtime +// harness reads (command.sh, training_config.yaml, requirements.yaml, and — when +// present — hyperparameters.yaml and env-var sidecars) and, for a code_source +// snapshot, the user's code tree, into a temp bundle root. Deploy uploads the whole +// root as an immutable-folder snapshot. Returns the root and a cleanup func. func writeBundleProject(ctx context.Context, cfg *runConfig, configPath string) (string, func(), error) { body, err := renderBundle(cfg, configPath) if err != nil { return "", func() {}, err } + artifacts, err := buildArtifacts(cfg, configPath) + if err != nil { + return "", func() {}, err + } bundleRoot, err := os.MkdirTemp("", "air-dabs-*") if err != nil { return "", func() {}, err } cleanup := func() { _ = os.RemoveAll(bundleRoot) } - // Copy the code_source working tree into the bundle root first, so a stray - // command.sh / databricks.yml in the user's tree can't shadow ours below. + // Copy the code_source working tree first, so a stray command.sh / databricks.yml + // in the user's tree can't shadow the generated files written below. if cfg.CodeSource != nil && cfg.CodeSource.Snapshot != nil { if err := stageCodeSource(ctx, cfg.CodeSource.Snapshot, configPath, bundleRoot); err != nil { cleanup() @@ -136,9 +141,10 @@ func writeBundleProject(ctx context.Context, cfg *runConfig, configPath string) return "", func() {}, err } - // command.sh: the entrypoint the task's command_path points at. - if cfg.Command != nil { - if err := os.WriteFile(filepath.Join(bundleRoot, bundleCommandScript), []byte(*cfg.Command), 0o600); err != nil { + // The launch artifacts the harness expects co-located with command.sh, the same + // set the retired ephemeral path uploaded (buildArtifacts). + for _, it := range artifacts { + if err := os.WriteFile(filepath.Join(bundleRoot, it.name), it.data, 0o600); err != nil { cleanup() return "", func() {}, err } diff --git a/experimental/air/cmd/rundabs_test.go b/experimental/air/cmd/rundabs_test.go index 7b564e74169..b84ba90198c 100644 --- a/experimental/air/cmd/rundabs_test.go +++ b/experimental/air/cmd/rundabs_test.go @@ -15,13 +15,11 @@ import ( ) func TestWriteBundleProject(t *testing.T) { - cfg := &runConfig{ - ExperimentName: "exp", - Command: new("echo hi"), - Compute: &computeConfig{AcceleratorType: "GPU_1xA10", NumAccelerators: 1}, - } + configPath := writeConfigFile(t, "train.yaml", minimalConfig) + cfg, err := loadRunConfig(configPath) + require.NoError(t, err) - root, cleanup, err := writeBundleProject(t.Context(), cfg, "train.yaml") + root, cleanup, err := writeBundleProject(t.Context(), cfg, configPath) require.NoError(t, err) defer cleanup() @@ -32,10 +30,12 @@ func TestWriteBundleProject(t *testing.T) { assert.Contains(t, string(body), "immutable_folder: true") assert.Contains(t, string(body), "mode: development") - // command.sh holds the run's command verbatim. + // The launch artifacts the AI Runtime harness reads are staged next to command.sh. + assert.FileExists(t, filepath.Join(root, bundleCommandScript)) + assert.FileExists(t, filepath.Join(root, trainingConfigName)) script, err := os.ReadFile(filepath.Join(root, bundleCommandScript)) require.NoError(t, err) - assert.Equal(t, "echo hi", string(script)) + assert.Equal(t, "python train.py", string(script)) // cleanup removes the temp root. cleanup() @@ -49,14 +49,16 @@ func TestWriteBundleProjectStagesCodeSource(t *testing.T) { require.NoError(t, os.MkdirAll(filepath.Join(src, "pkg"), 0o755)) require.NoError(t, os.WriteFile(filepath.Join(src, "pkg", "mod.py"), []byte("x=1"), 0o644)) - cfg := &runConfig{ - ExperimentName: "exp", - Command: new("python train.py"), - Compute: &computeConfig{AcceleratorType: "GPU_1xA10", NumAccelerators: 1}, - CodeSource: &codeSourceConfig{Type: "snapshot", Snapshot: &snapshotSourceConfig{RootPath: src}}, - } + configPath := writeConfigFile(t, "train.yaml", minimalConfig+` +code_source: + type: snapshot + snapshot: + root_path: `+src+` +`) + cfg, err := loadRunConfig(configPath) + require.NoError(t, err) - root, cleanup, err := writeBundleProject(t.Context(), cfg, "train.yaml") + root, cleanup, err := writeBundleProject(t.Context(), cfg, configPath) require.NoError(t, err) defer cleanup() @@ -69,35 +71,42 @@ func TestWriteBundleProjectStagesCodeSource(t *testing.T) { func TestWriteBundleProjectCommandShadowProtection(t *testing.T) { // A command.sh in the user's tree must not shadow the one we generate from the - // run's command: our write happens after the tree copy. + // run's command: the artifact writes happen after the tree copy. src := t.TempDir() require.NoError(t, os.WriteFile(filepath.Join(src, bundleCommandScript), []byte("STALE"), 0o644)) - cfg := &runConfig{ - ExperimentName: "exp", - Command: new("FRESH"), - Compute: &computeConfig{AcceleratorType: "GPU_1xA10", NumAccelerators: 1}, - CodeSource: &codeSourceConfig{Type: "snapshot", Snapshot: &snapshotSourceConfig{RootPath: src}}, - } + configPath := writeConfigFile(t, "train.yaml", minimalConfig+` +code_source: + type: snapshot + snapshot: + root_path: `+src+` +`) + cfg, err := loadRunConfig(configPath) + require.NoError(t, err) - root, cleanup, err := writeBundleProject(t.Context(), cfg, "train.yaml") + root, cleanup, err := writeBundleProject(t.Context(), cfg, configPath) require.NoError(t, err) defer cleanup() + // minimalConfig's command is "python train.py", not the STALE tree copy. script, err := os.ReadFile(filepath.Join(root, bundleCommandScript)) require.NoError(t, err) - assert.Equal(t, "FRESH", string(script)) + assert.Equal(t, "python train.py", string(script)) } func TestWriteBundleProjectRejectsUnconvertible(t *testing.T) { // A gate failure (here: a $CODE_SOURCE_PATH command) surfaces before any temp // directory is created. - cfg := &runConfig{ - ExperimentName: "exp", - Command: new("cd $CODE_SOURCE_PATH && python train.py"), - Compute: &computeConfig{AcceleratorType: "GPU_1xA10", NumAccelerators: 1}, - } - _, _, err := writeBundleProject(t.Context(), cfg, "train.yaml") + configPath := writeConfigFile(t, "train.yaml", ` +experiment_name: exp +command: cd $CODE_SOURCE_PATH && python train.py +compute: + accelerator_type: GPU_1xA10 + num_accelerators: 1 +`) + cfg, err := loadRunConfig(configPath) + require.NoError(t, err) + _, _, err = writeBundleProject(t.Context(), cfg, configPath) require.Error(t, err) } diff --git a/experimental/air/cmd/runupload.go b/experimental/air/cmd/runupload.go index fb9ca00b987..714c3d2d3c3 100644 --- a/experimental/air/cmd/runupload.go +++ b/experimental/air/cmd/runupload.go @@ -3,14 +3,10 @@ package aircmd import ( "bytes" "context" - "encoding/json" "fmt" "io" - "maps" "os" "path/filepath" - "slices" - "strings" "github.com/databricks/cli/libs/filer" "go.yaml.in/yaml/v3" @@ -24,8 +20,6 @@ const ( commandScriptName = "command.sh" requirementsName = "requirements.yaml" hyperparametersName = "hyperparameters.yaml" - envVarsName = "env_vars.json" - secretEnvVarsName = "secret_env_vars.json" ) // maxConfigYAMLBytes caps training_config.yaml. It is referenced by the Jobs @@ -102,59 +96,14 @@ func buildArtifacts(cfg *runConfig, configPath string) ([]uploadItem, error) { items = append(items, uploadItem{hyperparametersName, data}) } - // The ai_runtime_task proto carries no inline env vars or secrets; stage them - // as JSON files co-located with command.sh for the server-side launcher. - if len(cfg.EnvVariables) > 0 { - data, err := json.Marshal(envVarEntries(cfg.EnvVariables)) - if err != nil { - return nil, fmt.Errorf("failed to serialize env_variables: %w", err) - } - items = append(items, uploadItem{envVarsName, data}) - } - if len(cfg.Secrets) > 0 { - data, err := json.Marshal(secretEnvVarEntries(cfg.Secrets)) - if err != nil { - return nil, fmt.Errorf("failed to serialize secrets: %w", err) - } - items = append(items, uploadItem{secretEnvVarsName, data}) - } + // Env vars and secrets are NOT staged as sidecar files: they ride the common + // Jobs env-var API (the environment_variables profile the converter emits, see + // exportbundle.go), matching the Python CLI which dropped the env_vars.json / + // secret_env_vars.json side-channel. return items, nil } -// envVarEntry is one entry in env_vars.json. -type envVarEntry struct { - Name string `json:"name"` - Value string `json:"value"` -} - -// secretEnvVarEntry is one entry in secret_env_vars.json. The YAML side is -// {ENV_VAR: "scope/key"}; the launcher wants the split form. -type secretEnvVarEntry struct { - Name string `json:"name"` - SecretScope string `json:"secret_scope"` - SecretKey string `json:"secret_key"` -} - -// envVarEntries renders env_variables sorted by name for deterministic output. -func envVarEntries(vars map[string]string) []envVarEntry { - out := make([]envVarEntry, 0, len(vars)) - for _, name := range slices.Sorted(maps.Keys(vars)) { - out = append(out, envVarEntry{Name: name, Value: vars[name]}) - } - return out -} - -// secretEnvVarEntries renders secrets sorted by name for deterministic output. -func secretEnvVarEntries(secrets map[string]string) []secretEnvVarEntry { - out := make([]secretEnvVarEntry, 0, len(secrets)) - for _, name := range slices.Sorted(maps.Keys(secrets)) { - scope, key, _ := strings.Cut(secrets[name], "/") - out = append(out, secretEnvVarEntry{Name: name, SecretScope: scope, SecretKey: key}) - } - return out -} - // uploadArtifacts writes each artifact into the launch directory, overwriting and // creating parents as needed. // diff --git a/experimental/air/cmd/runupload_test.go b/experimental/air/cmd/runupload_test.go index 0c87524735d..38425fda74b 100644 --- a/experimental/air/cmd/runupload_test.go +++ b/experimental/air/cmd/runupload_test.go @@ -83,7 +83,10 @@ func TestBuildArtifacts_InlineRequirementsAndParameters(t *testing.T) { assert.Contains(t, req, "- torch") } -func TestBuildArtifacts_EnvVarsAndSecrets(t *testing.T) { +func TestBuildArtifacts_EnvVarsNotStagedAsSidecars(t *testing.T) { + // Env vars and secrets ride the common Jobs env-var API (the converter's + // environment_variables profile), not sidecar files — matching the Python CLI, + // which dropped the env_vars.json / secret_env_vars.json side-channel. path := writeConfigFile(t, "run.yaml", "x: y\n") cfg := &runConfig{ Command: new("echo hi"), @@ -93,14 +96,8 @@ func TestBuildArtifacts_EnvVarsAndSecrets(t *testing.T) { items, err := buildArtifacts(cfg, path) require.NoError(t, err) - assert.Subset(t, itemNames(items), []string{envVarsName, secretEnvVarsName}) - - byName := map[string][]byte{} - for _, it := range items { - byName[it.name] = it.data - } - assert.JSONEq(t, `[{"name":"WANDB","value":"demo"}]`, string(byName[envVarsName])) - assert.JSONEq(t, `[{"name":"HF_TOKEN","secret_scope":"myscope","secret_key":"hf"}]`, string(byName[secretEnvVarsName])) + assert.NotContains(t, itemNames(items), "env_vars.json") + assert.NotContains(t, itemNames(items), "secret_env_vars.json") } func TestBuildArtifacts_RequirementsFile(t *testing.T) { From e8780b87fd215627a7a09e098b829d5247e2ccb9 Mon Sep 17 00:00:00 2001 From: vinchenzo-db Date: Wed, 15 Jul 2026 22:35:30 +0000 Subject: [PATCH 06/11] [air] Restore env-var sidecar dual-write alongside the common API Mirror the Python CLI's rollout dual-write (#2237609): env vars / secrets ride BOTH the common Jobs env-var API (the converter's environment_variables profile) and the env_vars.json / secret_env_vars.json sidecars the scheduler falls back to. The prior commit removed the sidecars; restore them. The sidecar fallback matters on two fronts: (1) the allowEnvironmentVariables SAFE flag is off in prod, so Jobs drops the profile field, and (2) on the DABs path the bundle SDK schema strips the profile field client-side regardless of the flag (vendored jobs.JobSettings has no environment_variables). Until the flag ramps AND the SDK carries the field, the sidecar is the only working env-var delivery on the DABs path. Verified on staging: a code-free run with env_variables + secrets deploys env_vars.json / secret_env_vars.json into the immutable snapshot with the correct shapes, next to command.sh. Co-authored-by: Isaac --- experimental/air/cmd/runupload.go | 66 ++++++++++++++++++++++++-- experimental/air/cmd/runupload_test.go | 19 +++++--- 2 files changed, 75 insertions(+), 10 deletions(-) diff --git a/experimental/air/cmd/runupload.go b/experimental/air/cmd/runupload.go index 714c3d2d3c3..31f971b69ad 100644 --- a/experimental/air/cmd/runupload.go +++ b/experimental/air/cmd/runupload.go @@ -3,10 +3,14 @@ package aircmd import ( "bytes" "context" + "encoding/json" "fmt" "io" + "maps" "os" "path/filepath" + "slices" + "strings" "github.com/databricks/cli/libs/filer" "go.yaml.in/yaml/v3" @@ -20,6 +24,8 @@ const ( commandScriptName = "command.sh" requirementsName = "requirements.yaml" hyperparametersName = "hyperparameters.yaml" + envVarsName = "env_vars.json" + secretEnvVarsName = "secret_env_vars.json" ) // maxConfigYAMLBytes caps training_config.yaml. It is referenced by the Jobs @@ -96,14 +102,66 @@ func buildArtifacts(cfg *runConfig, configPath string) ([]uploadItem, error) { items = append(items, uploadItem{hyperparametersName, data}) } - // Env vars and secrets are NOT staged as sidecar files: they ride the common - // Jobs env-var API (the environment_variables profile the converter emits, see - // exportbundle.go), matching the Python CLI which dropped the env_vars.json / - // secret_env_vars.json side-channel. + // Dual-write during the common-env-var-API rollout (mirrors the Python CLI's + // sdk/_submit). Env vars / secrets also ride the Jobs payload + // (JobSettings.environment_variables, see exportbundle.go), but that path is a + // no-op until the allowEnvironmentVariables SAFE flag ramps to prod — and on the + // DABs path the field is additionally stripped by the bundle SDK schema. While + // either is true, Jobs drops the field and the run snapshot has no env vars, so + // we also stage the env_vars.json / secret_env_vars.json files the scheduler + // falls back to. Once the flag is fully ramped (and the SDK carries the field), + // this staging and the scheduler fallback are removed together. + if len(cfg.EnvVariables) > 0 { + data, err := json.Marshal(envVarEntries(cfg.EnvVariables)) + if err != nil { + return nil, fmt.Errorf("failed to serialize env_variables: %w", err) + } + items = append(items, uploadItem{envVarsName, data}) + } + if len(cfg.Secrets) > 0 { + data, err := json.Marshal(secretEnvVarEntries(cfg.Secrets)) + if err != nil { + return nil, fmt.Errorf("failed to serialize secrets: %w", err) + } + items = append(items, uploadItem{secretEnvVarsName, data}) + } return items, nil } +// envVarEntry is one entry in env_vars.json. +type envVarEntry struct { + Name string `json:"name"` + Value string `json:"value"` +} + +// secretEnvVarEntry is one entry in secret_env_vars.json. The YAML side is +// {ENV_VAR: "scope/key"}; the launcher wants the split form. +type secretEnvVarEntry struct { + Name string `json:"name"` + SecretScope string `json:"secret_scope"` + SecretKey string `json:"secret_key"` +} + +// envVarEntries renders env_variables sorted by name for deterministic output. +func envVarEntries(vars map[string]string) []envVarEntry { + out := make([]envVarEntry, 0, len(vars)) + for _, name := range slices.Sorted(maps.Keys(vars)) { + out = append(out, envVarEntry{Name: name, Value: vars[name]}) + } + return out +} + +// secretEnvVarEntries renders secrets sorted by name for deterministic output. +func secretEnvVarEntries(secrets map[string]string) []secretEnvVarEntry { + out := make([]secretEnvVarEntry, 0, len(secrets)) + for _, name := range slices.Sorted(maps.Keys(secrets)) { + scope, key, _ := strings.Cut(secrets[name], "/") + out = append(out, secretEnvVarEntry{Name: name, SecretScope: scope, SecretKey: key}) + } + return out +} + // uploadArtifacts writes each artifact into the launch directory, overwriting and // creating parents as needed. // diff --git a/experimental/air/cmd/runupload_test.go b/experimental/air/cmd/runupload_test.go index 38425fda74b..105ff1e65fe 100644 --- a/experimental/air/cmd/runupload_test.go +++ b/experimental/air/cmd/runupload_test.go @@ -83,10 +83,11 @@ func TestBuildArtifacts_InlineRequirementsAndParameters(t *testing.T) { assert.Contains(t, req, "- torch") } -func TestBuildArtifacts_EnvVarsNotStagedAsSidecars(t *testing.T) { - // Env vars and secrets ride the common Jobs env-var API (the converter's - // environment_variables profile), not sidecar files — matching the Python CLI, - // which dropped the env_vars.json / secret_env_vars.json side-channel. +func TestBuildArtifacts_EnvVarsAndSecrets(t *testing.T) { + // Dual-write: env vars / secrets ride the common Jobs env-var API (converter + // profile) AND are staged as env_vars.json / secret_env_vars.json sidecars, the + // scheduler's fallback while the allowEnvironmentVariables flag is off (and, on + // the DABs path, while the bundle SDK schema strips the profile field). path := writeConfigFile(t, "run.yaml", "x: y\n") cfg := &runConfig{ Command: new("echo hi"), @@ -96,8 +97,14 @@ func TestBuildArtifacts_EnvVarsNotStagedAsSidecars(t *testing.T) { items, err := buildArtifacts(cfg, path) require.NoError(t, err) - assert.NotContains(t, itemNames(items), "env_vars.json") - assert.NotContains(t, itemNames(items), "secret_env_vars.json") + assert.Subset(t, itemNames(items), []string{envVarsName, secretEnvVarsName}) + + byName := map[string][]byte{} + for _, it := range items { + byName[it.name] = it.data + } + assert.JSONEq(t, `[{"name":"WANDB","value":"demo"}]`, string(byName[envVarsName])) + assert.JSONEq(t, `[{"name":"HF_TOKEN","secret_scope":"myscope","secret_key":"hf"}]`, string(byName[secretEnvVarsName])) } func TestBuildArtifacts_RequirementsFile(t *testing.T) { From 4b681589a2804c73f01259252a23012063eb1ecc Mon Sep 17 00:00:00 2001 From: vinchenzo-db Date: Thu, 16 Jul 2026 21:17:30 +0000 Subject: [PATCH 07/11] [air] Fix lint: maps.Copy, testifylint, in converter + test - exportbundle.go: replace the env-var copy loop with maps.Copy (modernize mapsloop). - exportbundle_test.go: assert.Contains instead of assert.True(strings.Contains(...)) (testifylint). Co-authored-by: Isaac --- experimental/air/cmd/exportbundle.go | 5 ++--- experimental/air/cmd/exportbundle_test.go | 3 +-- 2 files changed, 3 insertions(+), 5 deletions(-) diff --git a/experimental/air/cmd/exportbundle.go b/experimental/air/cmd/exportbundle.go index 32898240765..6f18c009433 100644 --- a/experimental/air/cmd/exportbundle.go +++ b/experimental/air/cmd/exportbundle.go @@ -2,6 +2,7 @@ package aircmd import ( "fmt" + "maps" "path/filepath" "strings" @@ -256,9 +257,7 @@ func envVarProfiles(cfg *runConfig) []exportedEnvVarProfile { return nil } variables := make(map[string]string, len(cfg.EnvVariables)+len(cfg.Secrets)) - for k, v := range cfg.EnvVariables { - variables[k] = v - } + maps.Copy(variables, cfg.EnvVariables) for envVar, secretRef := range cfg.Secrets { variables[envVar] = "{{secrets/" + secretRef + "}}" } diff --git a/experimental/air/cmd/exportbundle_test.go b/experimental/air/cmd/exportbundle_test.go index 031d2083123..b2969086511 100644 --- a/experimental/air/cmd/exportbundle_test.go +++ b/experimental/air/cmd/exportbundle_test.go @@ -3,7 +3,6 @@ package aircmd import ( "os" "path/filepath" - "strings" "testing" "github.com/stretchr/testify/assert" @@ -198,5 +197,5 @@ func TestMarshalBundleEnvVarsRoundTrip(t *testing.T) { assert.Contains(t, body, "environment_variables:") assert.Contains(t, body, "FOO: bar") // Header provenance line is present. - assert.True(t, strings.Contains(body, "Generated by `air export-bundle`")) + assert.Contains(t, body, "Generated by `air export-bundle`") } From f66ebd8106c8a078cabdc182e577107e3ff4bb59 Mon Sep 17 00:00:00 2001 From: vinchenzo-db Date: Thu, 16 Jul 2026 21:17:45 +0000 Subject: [PATCH 08/11] [air] Run command.sh from the synced code dir so code_source works MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Root cause (found via driver logs on staging): a DABs-deployed ai_runtime_task failed at execution with `python train.py: No such file or directory`. The AI Runtime harness invokes command.sh from an unrelated runtime work dir, not the snapshot dir where bundle deploy syncs the code, so relative references can't find code_source files. Code-free commands (e.g. `python -c ...`) were unaffected, which is why earlier smoke tests passed. Fix: generated command.sh now prepends `cd "$(dirname "$0")"`, running the user's command from command.sh's own directory (the synced code_source location) — the DABs-native equivalent of the ephemeral path's `cd $CODE_SOURCE_PATH`. Also fixes forbidigo lint in rundabs_test.go (errors.Is/fs.ErrNotExist over os.IsNotExist). Verified end-to-end on staging (dbc-04ac0685-8857): a code_source run that failed 4x now SUCCEEDs — driver log shows the cwd contains train.py and "numpy 2.1.3" printed. Co-authored-by: Isaac --- experimental/air/cmd/rundabs_test.go | 7 ++++--- experimental/air/cmd/runupload.go | 9 ++++++++- experimental/air/cmd/runupload_test.go | 9 ++++++++- 3 files changed, 20 insertions(+), 5 deletions(-) diff --git a/experimental/air/cmd/rundabs_test.go b/experimental/air/cmd/rundabs_test.go index b84ba90198c..ea4f7331a8b 100644 --- a/experimental/air/cmd/rundabs_test.go +++ b/experimental/air/cmd/rundabs_test.go @@ -1,6 +1,7 @@ package aircmd import ( + "io/fs" "os" "path/filepath" "testing" @@ -35,12 +36,12 @@ func TestWriteBundleProject(t *testing.T) { assert.FileExists(t, filepath.Join(root, trainingConfigName)) script, err := os.ReadFile(filepath.Join(root, bundleCommandScript)) require.NoError(t, err) - assert.Equal(t, "python train.py", string(script)) + assert.Equal(t, commandScript("python train.py"), string(script)) // cleanup removes the temp root. cleanup() _, err = os.Stat(root) - assert.True(t, os.IsNotExist(err)) + assert.ErrorIs(t, err, fs.ErrNotExist) } func TestWriteBundleProjectStagesCodeSource(t *testing.T) { @@ -91,7 +92,7 @@ code_source: // minimalConfig's command is "python train.py", not the STALE tree copy. script, err := os.ReadFile(filepath.Join(root, bundleCommandScript)) require.NoError(t, err) - assert.Equal(t, "python train.py", string(script)) + assert.Equal(t, commandScript("python train.py"), string(script)) } func TestWriteBundleProjectRejectsUnconvertible(t *testing.T) { diff --git a/experimental/air/cmd/runupload.go b/experimental/air/cmd/runupload.go index 31f971b69ad..7264104ebfc 100644 --- a/experimental/air/cmd/runupload.go +++ b/experimental/air/cmd/runupload.go @@ -69,7 +69,7 @@ func buildArtifacts(cfg *runConfig, configPath string) ([]uploadItem, error) { items := []uploadItem{ {trainingConfigName, configData}, - {commandScriptName, []byte(*cfg.Command)}, + {commandScriptName, []byte(commandScript(*cfg.Command))}, } switch reqPath, ok := cfg.requirementsFile(); { @@ -129,6 +129,13 @@ func buildArtifacts(cfg *runConfig, configPath string) ([]uploadItem, error) { return items, nil } +// commandScript runs the user's command from command.sh's own directory, where the +// synced code_source files live. The harness invokes command.sh from an unrelated +// work dir, so relative references (e.g. `python train.py`) would otherwise fail. +func commandScript(command string) string { + return `cd "$(dirname "$0")"` + "\n" + command +} + // envVarEntry is one entry in env_vars.json. type envVarEntry struct { Name string `json:"name"` diff --git a/experimental/air/cmd/runupload_test.go b/experimental/air/cmd/runupload_test.go index 105ff1e65fe..3ae96a165ec 100644 --- a/experimental/air/cmd/runupload_test.go +++ b/experimental/air/cmd/runupload_test.go @@ -54,7 +54,14 @@ func TestBuildArtifacts_CommandAndConfig(t *testing.T) { require.NoError(t, err) assert.Equal(t, []string{trainingConfigName, commandScriptName}, itemNames(items)) assert.Equal(t, minimalConfig, string(items[0].data)) - assert.Equal(t, "python train.py", string(items[1].data)) + // command.sh cds to its own dir (the synced code_source location) before the command. + assert.Equal(t, "cd \"$(dirname \"$0\")\"\npython train.py", string(items[1].data)) +} + +func TestCommandScript(t *testing.T) { + // Prepends a cd to the script's own directory so relative refs resolve against + // the synced code_source location. + assert.Equal(t, "cd \"$(dirname \"$0\")\"\npython train.py", commandScript("python train.py")) } func TestBuildArtifacts_InlineRequirementsAndParameters(t *testing.T) { From 920c55290648ac39e07b616d1dd203b8727c554d Mon Sep 17 00:00:00 2001 From: vinchenzo-db Date: Thu, 16 Jul 2026 21:29:03 +0000 Subject: [PATCH 09/11] [air] Remove dead snapshot code-upload machinery MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Retiring the ephemeral runs/submit path orphaned AIR's own tarball/git-archive code-upload subsystem — the DABs path delivers code via bundle deploy (immutable_folder), not a tarball. golangci `unused` flagged the entry points; removing them cascades to the whole subsystem. Delete snapshot_git.go, snapshot_package.go, snapshot_cachekey.go, snapshot_resolve.go and their tests; gut snapshot.go to just resolveRootPath (the one piece still used, by rundabs.go's stageCodeSource). Net -1585 lines. Verified: build, `go test ./experimental/air/...`, and the CI-exact lint (go tool -modfile=tools/go.mod golangci-lint run) all pass with 0 issues. Co-authored-by: Isaac --- experimental/air/cmd/snapshot.go | 236 +------------ experimental/air/cmd/snapshot_cachekey.go | 34 -- .../air/cmd/snapshot_cachekey_test.go | 62 ---- experimental/air/cmd/snapshot_git.go | 309 ------------------ experimental/air/cmd/snapshot_git_test.go | 296 ----------------- experimental/air/cmd/snapshot_package.go | 132 -------- experimental/air/cmd/snapshot_package_test.go | 163 --------- experimental/air/cmd/snapshot_resolve.go | 120 ------- experimental/air/cmd/snapshot_resolve_test.go | 114 ------- experimental/air/cmd/snapshot_test.go | 124 +------ 10 files changed, 5 insertions(+), 1585 deletions(-) delete mode 100644 experimental/air/cmd/snapshot_cachekey.go delete mode 100644 experimental/air/cmd/snapshot_cachekey_test.go delete mode 100644 experimental/air/cmd/snapshot_git.go delete mode 100644 experimental/air/cmd/snapshot_git_test.go delete mode 100644 experimental/air/cmd/snapshot_package.go delete mode 100644 experimental/air/cmd/snapshot_package_test.go delete mode 100644 experimental/air/cmd/snapshot_resolve.go delete mode 100644 experimental/air/cmd/snapshot_resolve_test.go diff --git a/experimental/air/cmd/snapshot.go b/experimental/air/cmd/snapshot.go index aba67f6109c..8e47d95b288 100644 --- a/experimental/air/cmd/snapshot.go +++ b/experimental/air/cmd/snapshot.go @@ -1,58 +1,19 @@ package aircmd import ( - "bytes" "context" - "errors" "fmt" - "io/fs" "os" - "path" "path/filepath" "strings" - "time" "github.com/databricks/cli/libs/env" - "github.com/databricks/cli/libs/filer" - "github.com/databricks/cli/libs/log" - "github.com/databricks/databricks-sdk-go" ) -// Snapshot orchestrator: resolve → package+upload → sidecars, uploading via -// libs/filer. The Python CLI did this inline; here it's split into steps. - -// snapshotResult holds the paths wired into the submit payload: the uploaded -// tarball and the optional provenance sidecars (empty when not produced). -type snapshotResult struct { - CodeSourcePath string - GitStatePath string - GitDiffPath string -} - -// repoSnapshotsSubdir is the per-user workspace location for cached tarballs, under -// the user's home. Volume uploads use remote_volume directly instead. -const repoSnapshotsSubdir = ".air/repo_snapshots" - -// snapshotCodeSource packages and uploads the code_source snapshot, returning the -// paths to attach to the ai_runtime_task. userDir is the user's workspace home; -// funcDir is the run's launch directory (where sidecars land). -func snapshotCodeSource(ctx context.Context, w *databricks.WorkspaceClient, snap *snapshotSourceConfig, configPath, userDir, funcDir string) (snapshotResult, error) { - repoPath, err := resolveRootPath(ctx, snap.RootPath, filepath.Dir(configPath)) - if err != nil { - return snapshotResult{}, err - } - - up, err := newSnapshotUploader(w, snap, userDir, funcDir, filepath.Base(repoPath)) - if err != nil { - return snapshotResult{}, err - } - return runSnapshot(ctx, up, repoPath, snap) -} - -// resolveRootPath resolves a snapshot root_path the way the Python normalize layer -// does: expand environment variables and ~, strip a leading "project_root/" (meaning -// "relative to the YAML file"), and resolve the rest against the config's directory. -// It then confirms the path exists and is a directory. +// resolveRootPath resolves a code_source snapshot root_path: expand environment +// variables and ~, strip a leading "project_root/" (meaning "relative to the YAML +// file"), and resolve the rest against the config's directory. It then confirms the +// path exists and is a directory. func resolveRootPath(ctx context.Context, rawPath, configDir string) (string, error) { expanded := os.ExpandEnv(rawPath) if home, err := env.UserHomeDir(ctx); err == nil { @@ -73,8 +34,6 @@ func resolveRootPath(ctx context.Context, rawPath, configDir string) (string, er resolved = filepath.Join(configDir, expanded) } - // Resolve to an absolute path so the directory name (used for the tarball name - // and archive prefix) is a real basename, not "." or a trailing relative segment. abs, err := filepath.Abs(resolved) if err != nil { return "", fmt.Errorf("failed to resolve root_path %s: %w", resolved, err) @@ -90,190 +49,3 @@ func resolveRootPath(ctx context.Context, rawPath, configDir string) (string, er } return resolved, nil } - -// snapshotUploader splits the snapshot's two destinations: the tarball goes to a -// cache location (the user's repo_snapshots dir or a Volume), sidecars to the run's -// funcDir. tarBase/sidecarBase are the absolute roots, for reporting final paths. -type snapshotUploader struct { - tarStore filer.Filer - sidecarStore filer.Filer - tarBase string - sidecarBase string -} - -// runSnapshot resolves the packaging plan, uploads the tarball, then uploads the -// provenance sidecars. repoPath is the resolved root_path. -func runSnapshot(ctx context.Context, up snapshotUploader, repoPath string, snap *snapshotSourceConfig) (snapshotResult, error) { - git := newGitRepo(repoPath) - plan, err := resolveSnapshotPlan(ctx, git, snap.Git, snap.IncludePaths) - if err != nil { - return snapshotResult{}, err - } - - dirName := filepath.Base(repoPath) - - tarName, err := uploadTarball(ctx, up, git, plan, repoPath, dirName) - if err != nil { - return snapshotResult{}, err - } - - result := snapshotResult{CodeSourcePath: path.Join(up.tarBase, tarName)} - - // Provenance sidecars are best-effort: a git/upload hiccup here must not fail an - // otherwise-valid submission. Non-git roots have no provenance to record. - if plan.isGitRepo { - result.GitStatePath, result.GitDiffPath = uploadSidecars(ctx, up, git, plan) - } - return result, nil -} - -// uploadTarball packages the snapshot and uploads it, returning the tarball's name -// within the tar store. For git_archive it checks the cache first and skips -// packaging+upload on a hit. It writes the tarball to a temp file that is always -// cleaned up. -func uploadTarball(ctx context.Context, up snapshotUploader, git gitRepo, plan snapshotPlan, repoPath, dirName string) (string, error) { - // git_archive is cacheable by (commit, include_paths); a hit means the identical - // tarball is already uploaded, so packaging and upload are skipped entirely. - if plan.mode == modeGitArchive { - cacheKey := computeSnapshotCacheKey(plan.commitSHA, plan.includePaths) - tarName := fmt.Sprintf("%s_%s.tar.gz", dirName, cacheKey[:16]) - if exists, err := fileExists(ctx, up.tarStore, tarName); err != nil { - return "", err - } else if exists { - log.Debugf(ctx, "snapshot cache hit for %s at %s", shortSHA(plan.commitSHA), path.Join(up.tarBase, tarName)) - return tarName, nil - } - if err := packageAndUpload(ctx, up, tarName, func(out string) error { - return createGitArchiveSnapshot(ctx, git, plan.commitSHA, out, dirName, plan.includePaths) - }); err != nil { - return "", err - } - return tarName, nil - } - - // plain_tar is not cacheable (working-tree content isn't pinned to a SHA), so it - // is timestamp-named to avoid clobbering a concurrent submission. - tarName := fmt.Sprintf("%s_%s.tar.gz", dirName, time.Now().UTC().Format("20060102_150405")) - if err := packageAndUpload(ctx, up, tarName, func(out string) error { - return createPlainTarball(ctx, repoPath, out, plan.includePaths) - }); err != nil { - return "", err - } - return tarName, nil -} - -// packageAndUpload writes the tarball via pkg into a temp file, then uploads it to -// tarName in the tar store. The temp file is always removed. -func packageAndUpload(ctx context.Context, up snapshotUploader, tarName string, pkg func(outputPath string) error) error { - tmp, err := os.CreateTemp("", "air-snapshot-*.tar.gz") - if err != nil { - return fmt.Errorf("failed to create temp tarball: %w", err) - } - tmpPath := tmp.Name() - tmp.Close() - defer os.Remove(tmpPath) - - if err := pkg(tmpPath); err != nil { - return err - } - - f, err := os.Open(tmpPath) - if err != nil { - return fmt.Errorf("failed to open tarball: %w", err) - } - defer f.Close() - - if err := up.tarStore.Write(ctx, tarName, f, filer.OverwriteIfExists, filer.CreateParentDirectories); err != nil { - return fmt.Errorf("failed to upload snapshot to %s: %w", path.Join(up.tarBase, tarName), err) - } - return nil -} - -// uploadSidecars builds and uploads the git_state.json and optional git_diff.patch -// provenance sidecars into the run's funcDir. It is best-effort: any failure logs a -// warning and returns whatever paths did upload (possibly none), never an error. -func uploadSidecars(ctx context.Context, up snapshotUploader, git gitRepo, plan snapshotPlan) (statePath, diffPath string) { - mode := packagingModePlainTar - pinnedTip := "" - if plan.mode == modeGitArchive { - mode = packagingModeGitArchive - pinnedTip = plan.commitSHA - } - - sidecar, err := buildGitStateSidecar(ctx, git, mode, pinnedTip, time.Now()) - if err != nil { - log.Warnf(ctx, "skipping git provenance sidecar: %v", err) - return "", "" - } - - // Capture the dirty diff first so its status/path land in git_state.json. - if sidecar.Dirty { - status, diff := captureDirtyDiff(ctx, git, dirtyDiffSizeCapBytes, dirtyDiffTimeout) - sidecar.DiffStatus = status - if status == diffStatusCaptured { - if err := up.sidecarStore.Write(ctx, gitDiffName, bytes.NewReader(diff), filer.OverwriteIfExists, filer.CreateParentDirectories); err != nil { - log.Warnf(ctx, "failed to upload git diff sidecar: %v", err) - sidecar.DiffStatus = diffStatusClean - } else { - diffPath = path.Join(up.sidecarBase, gitDiffName) - sidecar.DiffPath = &diffPath - } - } - } - - data, err := sidecar.marshal() - if err != nil { - log.Warnf(ctx, "failed to encode git state sidecar: %v", err) - return "", diffPath - } - if err := up.sidecarStore.Write(ctx, gitStateName, bytes.NewReader(data), filer.OverwriteIfExists, filer.CreateParentDirectories); err != nil { - log.Warnf(ctx, "failed to upload git state sidecar: %v", err) - return "", diffPath - } - return path.Join(up.sidecarBase, gitStateName), diffPath -} - -// gitStateName and gitDiffName are the sidecar basenames read by the backend. -const ( - gitStateName = "git_state.json" - gitDiffName = "git_diff.patch" -) - -// fileExists reports whether name exists in the store, treating fs.ErrNotExist as -// "no". Any other error propagates. -func fileExists(ctx context.Context, store filer.Filer, name string) (bool, error) { - _, err := store.Stat(ctx, name) - if err == nil { - return true, nil - } - if errors.Is(err, fs.ErrNotExist) { - return false, nil - } - return false, fmt.Errorf("failed to check snapshot cache: %w", err) -} - -// newSnapshotUploader builds the uploader for a submission. The tarball store is a -// Volume (when remote_volume is set) or the user's repo_snapshots workspace dir; -// sidecars always go to the run's funcDir in the workspace. -func newSnapshotUploader(w *databricks.WorkspaceClient, snap *snapshotSourceConfig, userDir, funcDir, dirName string) (snapshotUploader, error) { - sidecarStore, err := filer.NewWorkspaceFilesClient(w, funcDir) - if err != nil { - return snapshotUploader{}, err - } - - if snap.RemoteVolume != nil { - tarBase := strings.TrimRight(*snap.RemoteVolume, "/") - tarStore, err := filer.NewFilesClient(w, tarBase) - if err != nil { - return snapshotUploader{}, err - } - return snapshotUploader{tarStore: tarStore, sidecarStore: sidecarStore, tarBase: tarBase, sidecarBase: funcDir}, nil - } - - tarBase := path.Join(userDir, repoSnapshotsSubdir, dirName) - tarStore, err := filer.NewWorkspaceFilesClient(w, tarBase) - if err != nil { - return snapshotUploader{}, err - } - return snapshotUploader{tarStore: tarStore, sidecarStore: sidecarStore, tarBase: tarBase, sidecarBase: funcDir}, nil -} diff --git a/experimental/air/cmd/snapshot_cachekey.go b/experimental/air/cmd/snapshot_cachekey.go deleted file mode 100644 index 44c58ee903b..00000000000 --- a/experimental/air/cmd/snapshot_cachekey.go +++ /dev/null @@ -1,34 +0,0 @@ -package aircmd - -// This file packages a local code directory into a tarball, uploads it to the -// workspace (or a Volume), and records git provenance sidecars for cache -// invalidation — the Go port of the Python CLI's code_source snapshot path. - -import ( - "crypto/sha256" - "encoding/hex" - "slices" - "strings" -) - -// snapshotPackagingVersion is bumped when packaging logic changes in a way that invalidates existing caches -const snapshotPackagingVersion = "v1" - -// computeSnapshotCacheKey returns a stable cache key for a snapshot tarball: the -// SHA-256 digest of (commitSHA, normalized includePaths, snapshotPackagingVersion). -// Changing any input yields a different entry. -func computeSnapshotCacheKey(commitSHA string, includePaths []string) string { - var normalizedPaths string - if len(includePaths) > 0 { - trimmed := make([]string, len(includePaths)) - for i, p := range includePaths { - trimmed[i] = strings.TrimSpace(p) - } - slices.Sort(trimmed) - normalizedPaths = strings.Join(trimmed, "\n") - } - - keyMaterial := commitSHA + "\n" + normalizedPaths + "\n" + snapshotPackagingVersion - sum := sha256.Sum256([]byte(keyMaterial)) - return hex.EncodeToString(sum[:]) -} diff --git a/experimental/air/cmd/snapshot_cachekey_test.go b/experimental/air/cmd/snapshot_cachekey_test.go deleted file mode 100644 index 5743217c003..00000000000 --- a/experimental/air/cmd/snapshot_cachekey_test.go +++ /dev/null @@ -1,62 +0,0 @@ -package aircmd - -import ( - "encoding/json" - "os" - "path/filepath" - "testing" - - "github.com/stretchr/testify/assert" - "github.com/stretchr/testify/require" -) - -type goldenCase struct { - Name string `json:"name"` - CommitSHA string `json:"commit_sha"` - IncludePaths []string `json:"include_paths"` - CacheKey string `json:"cache_key"` -} - -// TestComputeSnapshotCacheKeyGolden asserts byte-for-byte parity with golden -// fixtures across the local-only matrix (commit + include_paths permutations). -func TestComputeSnapshotCacheKeyGolden(t *testing.T) { - data, err := os.ReadFile(filepath.Join("testdata", "cache_keys.json")) - require.NoError(t, err) - - var cases []goldenCase - require.NoError(t, json.Unmarshal(data, &cases)) - require.NotEmpty(t, cases) - - for _, tc := range cases { - t.Run(tc.Name, func(t *testing.T) { - assert.Equal(t, tc.CacheKey, computeSnapshotCacheKey(tc.CommitSHA, tc.IncludePaths)) - }) - } -} - -// TestComputeSnapshotCacheKeyProperties pins the normalization behavior the golden cases -// encode, so a regression is legible without decoding hashes. -func TestComputeSnapshotCacheKeyProperties(t *testing.T) { - sha := "a3492b801c0ffee00000000000000000000dead" - - // Order-independent: sorting means unsorted input yields the sorted key. - assert.Equal(t, - computeSnapshotCacheKey(sha, []string{"a", "b", "c"}), - computeSnapshotCacheKey(sha, []string{"c", "a", "b"}), - ) - - // nil and empty include_paths are equivalent (both contribute an empty line). - assert.Equal(t, computeSnapshotCacheKey(sha, nil), computeSnapshotCacheKey(sha, []string{})) - - // Paths are trimmed before hashing. - assert.Equal(t, - computeSnapshotCacheKey(sha, []string{"research", "data"}), - computeSnapshotCacheKey(sha, []string{" research ", " data "}), - ) - - // Duplicates are NOT collapsed — they are sorted and kept, matching Python. - assert.NotEqual(t, computeSnapshotCacheKey(sha, []string{"x", "y"}), computeSnapshotCacheKey(sha, []string{"x", "x", "y"})) - - // The version constant participates: a different version is a different key. - assert.NotEqual(t, snapshotPackagingVersion, "") -} diff --git a/experimental/air/cmd/snapshot_git.go b/experimental/air/cmd/snapshot_git.go deleted file mode 100644 index 616b3049f74..00000000000 --- a/experimental/air/cmd/snapshot_git.go +++ /dev/null @@ -1,309 +0,0 @@ -package aircmd - -import ( - "bytes" - "context" - "encoding/json" - "errors" - "fmt" - "os/exec" - "strings" - "time" -) - -// Local, no-network git introspection and the git-state provenance sidecar, ported -// from the Python CLI's cli/utils/git_state.py. The remote-fetch helpers -// (fetch_branch_sha, remote detection, partial clone) are deliberately not ported: -// the snapshot archives the local copy only, so a ref must resolve to a local commit. - -// gitRepo runs git subcommands scoped to one repository via `git -C`. Arguments are -// passed as a slice, never a shell string, so branch/commit values can't inject. -type gitRepo struct { - path string -} - -func newGitRepo(path string) gitRepo { - return gitRepo{path: path} -} - -// run executes `git ` and returns stdout; a non-zero exit wraps stderr. -func (g gitRepo) run(ctx context.Context, args ...string) (string, error) { - out, err := g.runBytes(ctx, args...) - return string(out), err -} - -// runBytes is run returning raw stdout bytes, for the dirty-diff capture which needs -// exact bytes and a size measurement. -func (g gitRepo) runBytes(ctx context.Context, args ...string) ([]byte, error) { - full := append([]string{"-C", g.path}, args...) - cmd := exec.CommandContext(ctx, "git", full...) - - var stdout, stderr bytes.Buffer - cmd.Stdout = &stdout - cmd.Stderr = &stderr - if err := cmd.Run(); err != nil { - msg := strings.TrimSpace(stderr.String()) - if msg == "" { - return nil, fmt.Errorf("git %s: %w", strings.Join(args, " "), err) - } - return nil, fmt.Errorf("git %s: %w: %s", strings.Join(args, " "), err, msg) - } - return stdout.Bytes(), nil -} - -// isRepository reports whether the path is inside a git work tree. Using -// `rev-parse --is-inside-work-tree` (not a .git lookup) means a subdirectory of a -// repo counts — the common case when root_path is a subfolder of a monorepo. -func (g gitRepo) isRepository(ctx context.Context) bool { - out, err := g.run(ctx, "rev-parse", "--is-inside-work-tree") - if err != nil { - return false - } - return strings.TrimSpace(out) == "true" -} - -// headSHA returns the current HEAD commit SHA. -func (g gitRepo) headSHA(ctx context.Context) (string, error) { - out, err := g.run(ctx, "rev-parse", "HEAD") - if err != nil { - return "", err - } - return strings.TrimSpace(out), nil -} - -// hasUncommittedChanges reports whether there are staged or unstaged changes under -// the repo subtree. The `-- .` pathspec scopes the check so a subfolder snapshot -// considers only changes that could land in it. -func (g gitRepo) hasUncommittedChanges(ctx context.Context) (bool, error) { - out, err := g.run(ctx, "status", "--porcelain", "--", ".") - if err != nil { - return false, err - } - return strings.TrimSpace(out) != "", nil -} - -// hasUncommittedChangesInPaths reports whether there are uncommitted changes within -// the include paths (empty includePaths yields false). -// -// The pathspecs limit `git status` to those subtrees (it is O(working tree), slow on -// a large monorepo) and already scope the output to what could land in the snapshot. -// Unlike the Python source we don't re-parse the entries to filter by name: git -// reports a rename as `R \x00`, so a name-based re-filter keys off the old -// path and could miss a rename into an include path. The only caller needs the bool. -func (g gitRepo) hasUncommittedChangesInPaths(ctx context.Context, includePaths []string) (bool, error) { - var pathspecs []string - for _, p := range includePaths { - if s := strings.TrimRight(p, "/"); s != "" { - pathspecs = append(pathspecs, s) - } - } - if len(pathspecs) == 0 { - return false, nil - } - - args := append([]string{"status", "--porcelain", "--"}, pathspecs...) - out, err := g.run(ctx, args...) - if err != nil { - return false, err - } - return strings.TrimSpace(out) != "", nil -} - -// resolveLocalBranchSHA resolves a branch to its local-HEAD commit. No remote is -// contacted; the branch must exist locally. -func (g gitRepo) resolveLocalBranchSHA(ctx context.Context, branch string) (string, error) { - out, err := g.run(ctx, "rev-parse", "refs/heads/"+branch) - if err != nil { - return "", fmt.Errorf("failed to resolve local branch %q; ensure the branch exists locally and root_path is correct: %w", branch, err) - } - return strings.TrimSpace(out), nil -} - -// commitExistsLocally reports whether commitSHA is in the local object store, without -// triggering a promisor/lazy fetch. -func (g gitRepo) commitExistsLocally(ctx context.Context, commitSHA string) bool { - _, err := g.run(ctx, "cat-file", "-e", commitSHA) - return err == nil -} - -// currentBranch returns the branch name, or "" for a detached HEAD or on error. -func (g gitRepo) currentBranch(ctx context.Context) string { - out, err := g.run(ctx, "rev-parse", "--abbrev-ref", "HEAD") - if err != nil { - return "" - } - branch := strings.TrimSpace(out) - if branch == "HEAD" { - return "" - } - return branch -} - -// remoteURL returns the URL of the named remote, or "" if it has none. -func (g gitRepo) remoteURL(ctx context.Context, remoteName string) string { - out, err := g.run(ctx, "remote", "get-url", remoteName) - if err != nil { - return "" - } - return strings.TrimSpace(out) -} - -// mergeBaseWithUpstream resolves the merge-base of HEAD and a likely upstream ref, -// trying /HEAD, /main, then /master. It reads only local remote-tracking -// refs (no fetch), returning "" if none resolve. -func (g gitRepo) mergeBaseWithUpstream(ctx context.Context, remoteName string) string { - for _, ref := range []string{remoteName + "/HEAD", remoteName + "/main", remoteName + "/master"} { - out, err := g.run(ctx, "merge-base", "HEAD", ref) - if err != nil { - continue - } - if base := strings.TrimSpace(out); base != "" { - return base - } - } - return "" -} - -// validateIncludePathsExist checks that every include path exists at commitSHA. -// `git ls-tree` (without -d, so both blobs and trees count) reports an entry when the -// path exists; empty output means missing. -func (g gitRepo) validateIncludePathsExist(ctx context.Context, commitSHA string, includePaths []string) error { - var missing []string - for _, p := range includePaths { - out, err := g.run(ctx, "ls-tree", commitSHA, p) - if err != nil { - return err - } - if strings.TrimSpace(out) == "" { - missing = append(missing, p) - } - } - if len(missing) > 0 { - return fmt.Errorf("include_paths do not exist at commit %s: %s", shortSHA(commitSHA), strings.Join(missing, ", ")) - } - return nil -} - -// shortSHA abbreviates a commit SHA to 8 chars for log/error messages, tolerating -// user-supplied abbreviations shorter than that. -func shortSHA(sha string) string { - return sha[:min(len(sha), 8)] -} - -// --- git-state provenance sidecar (git_state.json + git_diff.patch) --- -// -// The backend reads git_state.json next to the tarball to tag the MLflow run with -// base/tip/dirty provenance, and logs git_diff.patch when the tree was dirty. -// Producing the sidecar is best-effort: callers warn and continue, never fail submit. - -// snapshotStateSchemaVersion is the git_state.json schema version. Bump only in -// coordination with the backend reader. -const snapshotStateSchemaVersion = 1 - -// defaultRemoteName is the remote consulted for merge-base and repo URL (local refs -// only — the remote-fetch path is gone). -const defaultRemoteName = "origin" - -// dirtyDiffSizeCapBytes caps the git_diff.patch sidecar; a larger diff records -// size_exceeded and is skipped to keep the upload small. -const dirtyDiffSizeCapBytes = 1024 * 1024 - -// dirtyDiffTimeout bounds `git diff HEAD` so provenance never delays submission. -const dirtyDiffTimeout = 5 * time.Second - -// packaging_mode values: how the uploaded tarball was produced. -const ( - packagingModeGitArchive = "git_archive" - packagingModePlainTar = "plain_tar" -) - -// diff_status values recorded in the sidecar. -const ( - diffStatusClean = "clean" - diffStatusCaptured = "captured" - diffStatusSizeExceeded = "size_exceeded" - diffStatusTimeout = "timeout" -) - -// gitStateSidecar is the git_state.json record. Field names and the null-for-absent -// encoding match the Python source, so nullable fields are *string (absent → null). -type gitStateSidecar struct { - SchemaVersion int `json:"schema_version"` - PackagingMode string `json:"packaging_mode"` - BaseCommit *string `json:"base_commit"` - TipCommit *string `json:"tip_commit"` - Branch *string `json:"branch"` - RepoURL *string `json:"repo_url"` - Dirty bool `json:"dirty"` - DiffStatus string `json:"diff_status"` - DiffPath *string `json:"diff_path"` - GeneratedAtUTC string `json:"generated_at_utc"` -} - -// nilIfEmpty maps "" to nil so an absent value serializes as JSON null. -func nilIfEmpty(s string) *string { - if s == "" { - return nil - } - return &s -} - -// buildGitStateSidecar gathers git provenance. pinnedTip overrides the HEAD-derived -// tip for git_archive (the tarball reflects that commit, not HEAD); pass "" for -// plain_tar. Metadata is best-effort — unavailable fields become null. -func buildGitStateSidecar(ctx context.Context, git gitRepo, packagingMode, pinnedTip string, now time.Time) (gitStateSidecar, error) { - tip := pinnedTip - if tip == "" { - head, err := git.headSHA(ctx) - if err != nil { - return gitStateSidecar{}, err - } - tip = head - } - - dirty, err := git.hasUncommittedChanges(ctx) - if err != nil { - return gitStateSidecar{}, err - } - - return gitStateSidecar{ - SchemaVersion: snapshotStateSchemaVersion, - PackagingMode: packagingMode, - BaseCommit: nilIfEmpty(git.mergeBaseWithUpstream(ctx, defaultRemoteName)), - TipCommit: nilIfEmpty(tip), - Branch: nilIfEmpty(git.currentBranch(ctx)), - RepoURL: nilIfEmpty(git.remoteURL(ctx, defaultRemoteName)), - Dirty: dirty, - DiffStatus: diffStatusClean, - DiffPath: nil, - GeneratedAtUTC: now.UTC().Format("2006-01-02T15:04:05.000000") + "Z", - }, nil -} - -// marshal renders the sidecar as indented JSON (matching Python's json.dump indent=2). -func (s gitStateSidecar) marshal() ([]byte, error) { - return json.MarshalIndent(s, "", " ") -} - -// captureDirtyDiff runs `git diff HEAD` over the repo subtree, returning a diff_status -// and the diff bytes (non-nil only when captured): clean (no changes or diff failed), -// captured (under the cap), size_exceeded, or timeout. -func captureDirtyDiff(ctx context.Context, git gitRepo, sizeCapBytes int, timeout time.Duration) (string, []byte) { - diffCtx, cancel := context.WithTimeout(ctx, timeout) - defer cancel() - - out, err := git.runBytes(diffCtx, "diff", "HEAD", "--", ".") - if err != nil { - if errors.Is(diffCtx.Err(), context.DeadlineExceeded) { - return diffStatusTimeout, nil - } - return diffStatusClean, nil - } - if len(out) == 0 { - return diffStatusClean, nil - } - if len(out) > sizeCapBytes { - return diffStatusSizeExceeded, nil - } - return diffStatusCaptured, out -} diff --git a/experimental/air/cmd/snapshot_git_test.go b/experimental/air/cmd/snapshot_git_test.go deleted file mode 100644 index cf1a821b1d6..00000000000 --- a/experimental/air/cmd/snapshot_git_test.go +++ /dev/null @@ -1,296 +0,0 @@ -package aircmd - -import ( - "encoding/json" - "os" - "os/exec" - "path/filepath" - "testing" - "time" - - "github.com/stretchr/testify/assert" - "github.com/stretchr/testify/require" -) - -// newTestRepo initializes a git repo in a temp dir with a deterministic identity -// and returns its path. Tests build up real commits/branches/dirty states on top, -// mirroring the Python git_state tests (which drive real repos, not a fake). -func newTestRepo(t *testing.T) string { - t.Helper() - dir := t.TempDir() - runGit(t, dir, "init", "-q", "-b", "main") - // Deterministic identity so commits succeed in a bare CI environment. - runGit(t, dir, "config", "user.email", "test@example.test") - runGit(t, dir, "config", "user.name", "Test") - return dir -} - -// runGit runs a git command in dir and fails the test on error. -func runGit(t *testing.T, dir string, args ...string) { - t.Helper() - cmd := exec.Command("git", append([]string{"-C", dir}, args...)...) - out, err := cmd.CombinedOutput() - require.NoError(t, err, "git %v: %s", args, out) -} - -// writeRepoFile writes a file at a repo-relative path, creating parent dirs. -func writeRepoFile(t *testing.T, repo, rel, content string) { - t.Helper() - full := filepath.Join(repo, rel) - require.NoError(t, os.MkdirAll(filepath.Dir(full), 0o755)) - require.NoError(t, os.WriteFile(full, []byte(content), 0o600)) -} - -// commitAll stages everything and commits, returning the new HEAD SHA. -func commitAll(t *testing.T, repo, msg string) string { - t.Helper() - runGit(t, repo, "add", "-A") - runGit(t, repo, "commit", "-q", "-m", msg) - sha, err := newGitRepo(repo).headSHA(t.Context()) - require.NoError(t, err) - return sha -} - -func TestGitRepo_IsRepository(t *testing.T) { - ctx := t.Context() - - repo := newTestRepo(t) - assert.True(t, newGitRepo(repo).isRepository(ctx)) - - // A subdirectory of a repo is still inside the work tree. - writeRepoFile(t, repo, "sub/x.txt", "hi") - assert.True(t, newGitRepo(filepath.Join(repo, "sub")).isRepository(ctx)) - - // A plain temp dir with no repo is not. - assert.False(t, newGitRepo(t.TempDir()).isRepository(ctx)) -} - -func TestGitRepo_HeadSHA(t *testing.T) { - ctx := t.Context() - repo := newTestRepo(t) - writeRepoFile(t, repo, "a.txt", "1") - sha := commitAll(t, repo, "init") - - got, err := newGitRepo(repo).headSHA(ctx) - require.NoError(t, err) - assert.Equal(t, sha, got) - assert.Len(t, got, 40) -} - -func TestGitRepo_HasUncommittedChanges(t *testing.T) { - ctx := t.Context() - repo := newTestRepo(t) - writeRepoFile(t, repo, "a.txt", "1") - commitAll(t, repo, "init") - - // Clean tree. - dirty, err := newGitRepo(repo).hasUncommittedChanges(ctx) - require.NoError(t, err) - assert.False(t, dirty) - - // Unstaged modification. - writeRepoFile(t, repo, "a.txt", "2") - dirty, err = newGitRepo(repo).hasUncommittedChanges(ctx) - require.NoError(t, err) - assert.True(t, dirty) -} - -func TestGitRepo_HasUncommittedChangesInPaths(t *testing.T) { - ctx := t.Context() - - repo := newTestRepo(t) - writeRepoFile(t, repo, "src/model.py", "1") - writeRepoFile(t, repo, "other/x.py", "1") - commitAll(t, repo, "init") - g := newGitRepo(repo) - - // No paths: no changes, and git is never consulted. - dirty, err := g.hasUncommittedChangesInPaths(ctx, nil) - require.NoError(t, err) - assert.False(t, dirty) - - // A change outside the included paths is ignored. - writeRepoFile(t, repo, "other/x.py", "2") - dirty, err = g.hasUncommittedChangesInPaths(ctx, []string{"src", "configs"}) - require.NoError(t, err) - assert.False(t, dirty) - - // A change inside an included path is reported. - writeRepoFile(t, repo, "src/model.py", "2") - dirty, err = g.hasUncommittedChangesInPaths(ctx, []string{"src", "configs"}) - require.NoError(t, err) - assert.True(t, dirty) - - // Trailing slashes on include paths are trimmed for the pathspec. - dirty, err = g.hasUncommittedChangesInPaths(ctx, []string{"other/"}) - require.NoError(t, err) - assert.True(t, dirty) -} - -func TestGitRepo_HasUncommittedChangesInPaths_Rename(t *testing.T) { - ctx := t.Context() - repo := newTestRepo(t) - writeRepoFile(t, repo, "src/old.py", "content") - commitAll(t, repo, "init") - - // A rename within an included path counts as a change, however git classifies - // it (rename vs delete+add); we only assert the boolean. - runGit(t, repo, "mv", "src/old.py", "src/new.py") - dirty, err := newGitRepo(repo).hasUncommittedChangesInPaths(ctx, []string{"src"}) - require.NoError(t, err) - assert.True(t, dirty) -} - -func TestGitRepo_ResolveLocalBranchSHA(t *testing.T) { - ctx := t.Context() - repo := newTestRepo(t) - writeRepoFile(t, repo, "a.txt", "1") - mainSHA := commitAll(t, repo, "init") - - // A second branch at its own commit. - runGit(t, repo, "checkout", "-q", "-b", "feature") - writeRepoFile(t, repo, "b.txt", "2") - featSHA := commitAll(t, repo, "feature work") - g := newGitRepo(repo) - - got, err := g.resolveLocalBranchSHA(ctx, "main") - require.NoError(t, err) - assert.Equal(t, mainSHA, got) - - got, err = g.resolveLocalBranchSHA(ctx, "feature") - require.NoError(t, err) - assert.Equal(t, featSHA, got) - - // A branch that does not exist locally errors (no remote is contacted). - _, err = g.resolveLocalBranchSHA(ctx, "nope") - require.Error(t, err) - assert.Contains(t, err.Error(), "resolve local branch") -} - -func TestGitRepo_CommitExistsLocally(t *testing.T) { - ctx := t.Context() - repo := newTestRepo(t) - writeRepoFile(t, repo, "a.txt", "1") - sha := commitAll(t, repo, "init") - g := newGitRepo(repo) - - assert.True(t, g.commitExistsLocally(ctx, sha)) - assert.False(t, g.commitExistsLocally(ctx, "deadbeefdeadbeefdeadbeefdeadbeefdeadbeef")) -} - -func TestGitRepo_ValidateIncludePathsExist(t *testing.T) { - ctx := t.Context() - repo := newTestRepo(t) - writeRepoFile(t, repo, "src/model.py", "1") - writeRepoFile(t, repo, "configs/train.yaml", "x") - writeRepoFile(t, repo, "train.py", "print()") - sha := commitAll(t, repo, "init") - g := newGitRepo(repo) - - // Both directory and file include_paths are accepted (ls-tree without -d). - require.NoError(t, g.validateIncludePathsExist(ctx, sha, []string{"src", "configs", "train.py"})) - - err := g.validateIncludePathsExist(ctx, sha, []string{"src", "missing"}) - require.Error(t, err) - assert.Contains(t, err.Error(), "missing") - assert.Contains(t, err.Error(), sha[:8]) -} - -func TestBuildGitStateSidecar_PlainTarClean(t *testing.T) { - ctx := t.Context() - repo := newTestRepo(t) - writeRepoFile(t, repo, "a.txt", "1") - head := commitAll(t, repo, "init") - - sc, err := buildGitStateSidecar(ctx, newGitRepo(repo), packagingModePlainTar, "", fixedNow) - require.NoError(t, err) - assert.Equal(t, snapshotStateSchemaVersion, sc.SchemaVersion) - assert.Equal(t, packagingModePlainTar, sc.PackagingMode) - require.NotNil(t, sc.TipCommit) - assert.Equal(t, head, *sc.TipCommit) - assert.False(t, sc.Dirty) - assert.Equal(t, diffStatusClean, sc.DiffStatus) - assert.Nil(t, sc.DiffPath) - // No remote in a bare test repo → base_commit and repo_url are null. - assert.Nil(t, sc.BaseCommit) - assert.Nil(t, sc.RepoURL) - require.NotNil(t, sc.Branch) - assert.Equal(t, "main", *sc.Branch) - assert.Equal(t, "2026-07-10T12:00:00.000000Z", sc.GeneratedAtUTC) -} - -func TestBuildGitStateSidecar_GitArchivePinsTip(t *testing.T) { - ctx := t.Context() - repo := newTestRepo(t) - writeRepoFile(t, repo, "a.txt", "1") - first := commitAll(t, repo, "init") - // Advance HEAD; the pinned tip must win over HEAD. - writeRepoFile(t, repo, "b.txt", "2") - commitAll(t, repo, "second") - - sc, err := buildGitStateSidecar(ctx, newGitRepo(repo), packagingModeGitArchive, first, fixedNow) - require.NoError(t, err) - require.NotNil(t, sc.TipCommit) - assert.Equal(t, first, *sc.TipCommit) - assert.Equal(t, packagingModeGitArchive, sc.PackagingMode) -} - -func TestBuildGitStateSidecar_Dirty(t *testing.T) { - ctx := t.Context() - repo := newTestRepo(t) - writeRepoFile(t, repo, "a.txt", "1") - commitAll(t, repo, "init") - writeRepoFile(t, repo, "a.txt", "2") // uncommitted - - sc, err := buildGitStateSidecar(ctx, newGitRepo(repo), packagingModePlainTar, "", fixedNow) - require.NoError(t, err) - assert.True(t, sc.Dirty) -} - -func TestGitStateSidecar_MarshalNullsAbsentFields(t *testing.T) { - ctx := t.Context() - repo := newTestRepo(t) - writeRepoFile(t, repo, "a.txt", "1") - commitAll(t, repo, "init") - - sc, err := buildGitStateSidecar(ctx, newGitRepo(repo), packagingModePlainTar, "", fixedNow) - require.NoError(t, err) - data, err := sc.marshal() - require.NoError(t, err) - - // Absent fields serialize as JSON null (not "" or omitted), matching Python. - var raw map[string]any - require.NoError(t, json.Unmarshal(data, &raw)) - require.Contains(t, raw, "base_commit") - assert.Nil(t, raw["base_commit"]) - require.Contains(t, raw, "repo_url") - assert.Nil(t, raw["repo_url"]) - require.Contains(t, raw, "diff_path") - assert.Nil(t, raw["diff_path"]) - assert.EqualValues(t, 1, raw["schema_version"]) -} - -func TestCaptureDirtyDiff(t *testing.T) { - ctx := t.Context() - repo := newTestRepo(t) - writeRepoFile(t, repo, "a.txt", "one\n") - commitAll(t, repo, "init") - - // Clean tree → no diff. - status, diff := captureDirtyDiff(ctx, newGitRepo(repo), dirtyDiffSizeCapBytes, dirtyDiffTimeout) - assert.Equal(t, diffStatusClean, status) - assert.Nil(t, diff) - - // Dirty tree → captured, and the diff mentions the changed file. - writeRepoFile(t, repo, "a.txt", "two\n") - status, diff = captureDirtyDiff(ctx, newGitRepo(repo), dirtyDiffSizeCapBytes, dirtyDiffTimeout) - assert.Equal(t, diffStatusCaptured, status) - assert.Contains(t, string(diff), "a.txt") - - // A tiny size cap forces size_exceeded and drops the bytes. - status, diff = captureDirtyDiff(ctx, newGitRepo(repo), 1, dirtyDiffTimeout) - assert.Equal(t, diffStatusSizeExceeded, status) - assert.Nil(t, diff) -} - -var fixedNow = time.Date(2026, 7, 10, 12, 0, 0, 0, time.UTC) diff --git a/experimental/air/cmd/snapshot_package.go b/experimental/air/cmd/snapshot_package.go deleted file mode 100644 index 672366086c9..00000000000 --- a/experimental/air/cmd/snapshot_package.go +++ /dev/null @@ -1,132 +0,0 @@ -package aircmd - -import ( - "bytes" - "context" - "fmt" - "os" - "os/exec" - "path/filepath" - "strings" -) - -// Tar builders ported from cli/utils/snapshot.py. Both shell out (git archive / tar) -// for parity and to reuse git's/tar's symlink, gitignore, and AppleDouble handling. -// The tarball's top-level dir name is load-bearing — the remote entry_script extracts -// to /databricks/code_source/ — so the --prefix / `-C parent dir` forms preserve it. - -// createGitArchiveSnapshot writes a gzipped tar of commitSHA to outputTarball via -// `git archive`, with every entry prefixed by directoryName/. When includePaths is -// set, only those paths are archived. -func createGitArchiveSnapshot(ctx context.Context, git gitRepo, commitSHA, outputTarball, directoryName string, includePaths []string) error { - // Single git invocation writes the gzipped tar with the desired prefix; no - // extract/repack. Provenance lives in the git_state.json sidecar, not here. - args := []string{ - "archive", - "--format=tar.gz", - "--prefix=" + directoryName + "/", - "-o", outputTarball, - commitSHA, - } - args = append(args, includePaths...) - if _, err := git.run(ctx, args...); err != nil { - return fmt.Errorf("failed to create git archive: %w", err) - } - return nil -} - -// createPlainTarball writes a gzipped tar of repoPath's working tree to -// outputTarball via `tar`. The archive preserves repoPath's directory name as the -// top-level entry. When includePaths is set, only those paths (nested under the -// directory name) are archived. .git and macOS AppleDouble files are always -// excluded; a .gitignore at repoPath is honored. -func createPlainTarball(ctx context.Context, repoPath, outputTarball string, includePaths []string) error { - dirName := filepath.Base(repoPath) - parent := filepath.Dir(repoPath) - - args := []string{"-czf", outputTarball} - - // Exclude macOS AppleDouble files: they sort before the real top-level dir and - // hijack a remote `head -1` parse. No-op on Linux. - args = append(args, "--exclude=._*") - - // Never ship .git — provenance flows via the git_state.json sidecar. - args = append(args, "--exclude=.git") - - // Honor .gitignore if present. - gitignorePath := filepath.Join(repoPath, ".gitignore") - if patterns, err := parseGitignore(gitignorePath); err == nil { - for _, p := range patterns { - if strings.Contains(p, "/") { - // Anchor path-relative patterns to the archive root so they don't - // match identically-named paths in subdirectories. - args = append(args, "--exclude="+dirName+"/"+strings.TrimPrefix(p, "/")) - } else { - args = append(args, "--exclude="+p) - } - } - } - - // Archive from the parent so the directory name is preserved; with include_paths, - // prefix each so entries nest under it (matching git archive --prefix). - args = append(args, "-C", parent) - if len(includePaths) > 0 { - for _, p := range includePaths { - args = append(args, dirName+"/"+p) - } - } else { - args = append(args, dirName) - } - - cmd := exec.CommandContext(ctx, "tar", args...) - var stderr bytes.Buffer - cmd.Stderr = &stderr - if err := cmd.Run(); err != nil { - if msg := strings.TrimSpace(stderr.String()); msg != "" { - return fmt.Errorf("failed to create plain tarball: %w: %s", err, msg) - } - return fmt.Errorf("failed to create plain tarball: %w", err) - } - return nil -} - -// parseGitignore reads a .gitignore and returns tar --exclude patterns. It mirrors -// the Python CLI's lossy normalization so plain-tar snapshots exclude the same set: -// -// - comments (#…) and blank lines are skipped; -// - negation patterns (!…) are unsupported by tar --exclude and skipped; -// - a trailing "/" (directory marker) is stripped; -// - "**" is not a path-separator-agnostic wildcard in tar, so "**/foo" → "foo" -// and "foo/**" → "foo"; a mid-path "**" has no tar equivalent and is skipped. -// -// A missing file returns (nil, error); callers treat any error as "no patterns". -func parseGitignore(path string) ([]string, error) { - data, err := os.ReadFile(path) - if err != nil { - return nil, err - } - - var patterns []string - for raw := range strings.SplitSeq(string(data), "\n") { - line := strings.TrimRight(raw, " \t\r") - if line == "" || strings.HasPrefix(line, "#") { - continue - } - if strings.HasPrefix(line, "!") { - continue - } - line = strings.TrimRight(line, "/") - if strings.Contains(line, "**") { - switch { - case strings.HasPrefix(line, "**/"): - line = line[len("**/"):] - case strings.HasSuffix(line, "/**"): - line = line[:len(line)-len("/**")] - default: - continue - } - } - patterns = append(patterns, line) - } - return patterns, nil -} diff --git a/experimental/air/cmd/snapshot_package_test.go b/experimental/air/cmd/snapshot_package_test.go deleted file mode 100644 index d895d59b98e..00000000000 --- a/experimental/air/cmd/snapshot_package_test.go +++ /dev/null @@ -1,163 +0,0 @@ -package aircmd - -import ( - "archive/tar" - "compress/gzip" - "os" - "path/filepath" - "slices" - "strings" - "testing" - - "github.com/stretchr/testify/assert" - "github.com/stretchr/testify/require" -) - -// tarballEntries returns the sorted list of entry names in a .tar.gz. -func tarballEntries(t *testing.T, path string) []string { - t.Helper() - f, err := os.Open(path) - require.NoError(t, err) - defer f.Close() - - gz, err := gzip.NewReader(f) - require.NoError(t, err) - defer gz.Close() - - var names []string - tr := tar.NewReader(gz) - for { - hdr, err := tr.Next() - if err != nil { - break - } - names = append(names, hdr.Name) - } - slices.Sort(names) - return names -} - -func TestCreateGitArchiveSnapshot(t *testing.T) { - ctx := t.Context() - repo := newTestRepo(t) - writeRepoFile(t, repo, "a.txt", "1") - writeRepoFile(t, repo, "src/model.py", "print()") - sha := commitAll(t, repo, "init") - - out := filepath.Join(t.TempDir(), "snap.tar.gz") - dirName := filepath.Base(repo) - require.NoError(t, createGitArchiveSnapshot(ctx, newGitRepo(repo), sha, out, dirName, nil)) - - entries := tarballEntries(t, out) - // Every real entry is prefixed with the directory name; the tracked files are - // present. git archive also emits a `pax_global_header` pseudo-entry carrying - // the commit SHA — it has no prefix and tar ignores it on extraction. - assert.Contains(t, entries, dirName+"/a.txt") - assert.Contains(t, entries, dirName+"/src/model.py") - for _, e := range entries { - if e == "pax_global_header" { - continue - } - assert.True(t, strings.HasPrefix(e, dirName+"/"), "entry %q lacks prefix", e) - } -} - -func TestCreateGitArchiveSnapshot_IncludePaths(t *testing.T) { - ctx := t.Context() - repo := newTestRepo(t) - writeRepoFile(t, repo, "a.txt", "1") - writeRepoFile(t, repo, "src/model.py", "print()") - sha := commitAll(t, repo, "init") - - out := filepath.Join(t.TempDir(), "snap.tar.gz") - dirName := filepath.Base(repo) - require.NoError(t, createGitArchiveSnapshot(ctx, newGitRepo(repo), sha, out, dirName, []string{"src"})) - - entries := tarballEntries(t, out) - assert.Contains(t, entries, dirName+"/src/model.py") - // a.txt is outside the include path, so it must not appear. - assert.NotContains(t, entries, dirName+"/a.txt") -} - -func TestCreatePlainTarball(t *testing.T) { - ctx := t.Context() - repo := newTestRepo(t) - writeRepoFile(t, repo, "a.txt", "1") - writeRepoFile(t, repo, "src/model.py", "print()") - commitAll(t, repo, "init") - // Uncommitted file must be included in a plain tar. - writeRepoFile(t, repo, "dirty.txt", "wip") - - out := filepath.Join(t.TempDir(), "snap.tar.gz") - require.NoError(t, createPlainTarball(ctx, repo, out, nil)) - - dirName := filepath.Base(repo) - entries := tarballEntries(t, out) - assert.Contains(t, entries, dirName+"/a.txt") - assert.Contains(t, entries, dirName+"/dirty.txt") - // .git is never shipped. - for _, e := range entries { - assert.NotContains(t, e, "/.git/") - } -} - -func TestCreatePlainTarball_HonorsGitignore(t *testing.T) { - ctx := t.Context() - repo := newTestRepo(t) - writeRepoFile(t, repo, "keep.txt", "1") - writeRepoFile(t, repo, "junk.log", "noise") - writeRepoFile(t, repo, ".gitignore", "*.log\n") - - out := filepath.Join(t.TempDir(), "snap.tar.gz") - require.NoError(t, createPlainTarball(ctx, repo, out, nil)) - - dirName := filepath.Base(repo) - entries := tarballEntries(t, out) - assert.Contains(t, entries, dirName+"/keep.txt") - assert.NotContains(t, entries, dirName+"/junk.log") -} - -func TestCreatePlainTarball_IncludePaths(t *testing.T) { - ctx := t.Context() - repo := newTestRepo(t) - writeRepoFile(t, repo, "a.txt", "1") - writeRepoFile(t, repo, "src/model.py", "print()") - - out := filepath.Join(t.TempDir(), "snap.tar.gz") - require.NoError(t, createPlainTarball(ctx, repo, out, []string{"src"})) - - dirName := filepath.Base(repo) - entries := tarballEntries(t, out) - assert.Contains(t, entries, dirName+"/src/model.py") - assert.NotContains(t, entries, dirName+"/a.txt") -} - -func TestParseGitignore(t *testing.T) { - dir := t.TempDir() - path := filepath.Join(dir, ".gitignore") - content := "# comment\n" + - "\n" + - "*.log\n" + - "!keep.log\n" + // negation: skipped - "build/\n" + // trailing slash stripped - "**/node_modules\n" + // **/foo -> foo - "dist/**\n" + // foo/** -> foo - "a/**/b\n" + // mid ** : skipped - "src/config\n" // path-relative kept as-is - require.NoError(t, os.WriteFile(path, []byte(content), 0o600)) - - patterns, err := parseGitignore(path) - require.NoError(t, err) - assert.Equal(t, []string{ - "*.log", - "build", - "node_modules", - "dist", - "src/config", - }, patterns) -} - -func TestParseGitignore_Missing(t *testing.T) { - _, err := parseGitignore(filepath.Join(t.TempDir(), "nope")) - require.Error(t, err) -} diff --git a/experimental/air/cmd/snapshot_resolve.go b/experimental/air/cmd/snapshot_resolve.go deleted file mode 100644 index 1146b641b92..00000000000 --- a/experimental/air/cmd/snapshot_resolve.go +++ /dev/null @@ -1,120 +0,0 @@ -package aircmd - -import ( - "context" - "errors" - "fmt" -) - -// This file ports the mode/ref resolution from the Python CLI's cli_entrypoint -// snapshot block (the if/elif at lines ~1541–1722), local-only. The remote-fetch -// branches are dropped: a git ref must resolve to a commit already present -// locally (git.remote is rejected at validation — see gitRef.validate). - -// snapshotMode is how the snapshot tarball is produced. -type snapshotMode int - -const ( - // modeGitArchive packages a pinned commit via `git archive`. The commit is - // deterministic, so the tarball is cacheable by (commit, include_paths). - modeGitArchive snapshotMode = iota - // modePlainTar packages the working tree (including uncommitted changes) via - // `tar`. Not cacheable — working-tree content isn't pinned to a SHA. - modePlainTar -) - -// snapshotPlan is the outcome of resolving how to package a snapshot: the mode, -// the commit SHA to archive (git_archive only; empty for plain_tar), and whether -// the working tree under the snapshot root has uncommitted changes. -type snapshotPlan struct { - mode snapshotMode - commitSHA string - hasUncommit bool - isGitRepo bool - includePaths []string -} - -// resolveSnapshotPlan decides how to package the snapshot (local-only): -// - git.commit → pin the SHA (must exist locally) → git_archive. -// - git.branch → the branch's local HEAD SHA → git_archive. -// - no ref / non-git dir → the working tree → plain_tar (no caching). -// -// The dirty check runs at most once (git status is O(working tree)) and is threaded -// into the plan. Dirty + git.branch is an error: the committed HEAD wouldn't include -// the uncommitted changes. -func resolveSnapshotPlan(ctx context.Context, git gitRepo, ref *gitRef, includePaths []string) (snapshotPlan, error) { - plan := snapshotPlan{includePaths: includePaths} - plan.isGitRepo = git.isRepository(ctx) - - // Detect uncommitted changes once. When include_paths is set, only changes - // under those paths can land in the snapshot, so scope the check to them — - // both more correct and cheaper than scanning the whole repo. - if plan.isGitRepo { - var err error - if len(includePaths) > 0 { - plan.hasUncommit, err = git.hasUncommittedChangesInPaths(ctx, includePaths) - } else { - plan.hasUncommit, err = git.hasUncommittedChanges(ctx) - } - if err != nil { - return snapshotPlan{}, err - } - } - - // Non-git directory: plain tar, no ref allowed. gitRef.validate already rejects - // git.* on a non-git dir at load time, but guard here too since this function - // is the single decision point. - if !plan.isGitRepo { - if ref != nil { - return snapshotPlan{}, fmt.Errorf("git.* is set but %s is not a git repository", git.path) - } - plan.mode = modePlainTar - return plan, nil - } - - // git repo, no ref: package the working tree as plain tar (uncommitted changes - // included). Provenance is captured separately via the git_state sidecar. - if ref == nil { - plan.mode = modePlainTar - return plan, nil - } - - switch { - case ref.Commit != nil: - // git.commit pins a committed SHA; local uncommitted changes are irrelevant - // and won't be included. The commit must exist locally — no remote fetch. - commit := *ref.Commit - if !git.commitExistsLocally(ctx, commit) { - return snapshotPlan{}, fmt.Errorf("commit %q does not exist locally; fetch it (e.g. `git fetch`) before submitting — the snapshot archives your local copy and does not fetch from a remote", commit) - } - plan.mode = modeGitArchive - plan.commitSHA = commit - - case ref.Branch != nil: - // git.branch deploys the branch's local HEAD. A dirty tree here is an error: - // the committed HEAD wouldn't include the uncommitted changes. - if plan.hasUncommit { - return snapshotPlan{}, fmt.Errorf("uncommitted changes under %s would not be included: git.branch deploys the committed HEAD of %q. Commit your changes, or use git.commit to pin a specific revision", git.path, *ref.Branch) - } - sha, err := git.resolveLocalBranchSHA(ctx, *ref.Branch) - if err != nil { - return snapshotPlan{}, err - } - plan.mode = modeGitArchive - plan.commitSHA = sha - - default: - // gitRef.validate guarantees exactly one of branch/commit is set. - return snapshotPlan{}, errors.New("git: must specify either 'branch' or 'commit'") - } - - // For git_archive with include_paths, verify each path exists at the resolved - // commit so a typo fails fast rather than producing an empty subtree. - if len(includePaths) > 0 { - if err := git.validateIncludePathsExist(ctx, plan.commitSHA, includePaths); err != nil { - return snapshotPlan{}, err - } - } - - return plan, nil -} diff --git a/experimental/air/cmd/snapshot_resolve_test.go b/experimental/air/cmd/snapshot_resolve_test.go deleted file mode 100644 index c8c946f8394..00000000000 --- a/experimental/air/cmd/snapshot_resolve_test.go +++ /dev/null @@ -1,114 +0,0 @@ -package aircmd - -import ( - "testing" - - "github.com/stretchr/testify/assert" - "github.com/stretchr/testify/require" -) - -func TestResolveSnapshotPlan_Commit(t *testing.T) { - ctx := t.Context() - repo := newTestRepo(t) - writeRepoFile(t, repo, "a.txt", "1") - sha := commitAll(t, repo, "init") - - plan, err := resolveSnapshotPlan(ctx, newGitRepo(repo), &gitRef{Commit: new(sha)}, nil) - require.NoError(t, err) - assert.Equal(t, modeGitArchive, plan.mode) - assert.Equal(t, sha, plan.commitSHA) - assert.True(t, plan.isGitRepo) - - // A commit pin is valid even with a dirty tree: local changes are irrelevant. - writeRepoFile(t, repo, "a.txt", "2") - plan, err = resolveSnapshotPlan(ctx, newGitRepo(repo), &gitRef{Commit: new(sha)}, nil) - require.NoError(t, err) - assert.Equal(t, modeGitArchive, plan.mode) - assert.True(t, plan.hasUncommit) -} - -func TestResolveSnapshotPlan_CommitNotLocal(t *testing.T) { - ctx := t.Context() - repo := newTestRepo(t) - writeRepoFile(t, repo, "a.txt", "1") - commitAll(t, repo, "init") - - absent := "deadbeefdeadbeefdeadbeefdeadbeefdeadbeef" - _, err := resolveSnapshotPlan(ctx, newGitRepo(repo), &gitRef{Commit: new(absent)}, nil) - require.Error(t, err) - assert.Contains(t, err.Error(), "does not exist locally") -} - -func TestResolveSnapshotPlan_BranchLocalHead(t *testing.T) { - ctx := t.Context() - repo := newTestRepo(t) - writeRepoFile(t, repo, "a.txt", "1") - mainSHA := commitAll(t, repo, "init") - - plan, err := resolveSnapshotPlan(ctx, newGitRepo(repo), &gitRef{Branch: new("main")}, nil) - require.NoError(t, err) - assert.Equal(t, modeGitArchive, plan.mode) - assert.Equal(t, mainSHA, plan.commitSHA) -} - -func TestResolveSnapshotPlan_BranchDirtyIsError(t *testing.T) { - ctx := t.Context() - repo := newTestRepo(t) - writeRepoFile(t, repo, "a.txt", "1") - commitAll(t, repo, "init") - writeRepoFile(t, repo, "a.txt", "2") // uncommitted - - _, err := resolveSnapshotPlan(ctx, newGitRepo(repo), &gitRef{Branch: new("main")}, nil) - require.Error(t, err) - assert.Contains(t, err.Error(), "uncommitted changes") - assert.Contains(t, err.Error(), "git.commit") -} - -func TestResolveSnapshotPlan_NoRefPlainTar(t *testing.T) { - ctx := t.Context() - repo := newTestRepo(t) - writeRepoFile(t, repo, "a.txt", "1") - commitAll(t, repo, "init") - writeRepoFile(t, repo, "a.txt", "2") // dirty tree is fine for plain tar - - plan, err := resolveSnapshotPlan(ctx, newGitRepo(repo), nil, nil) - require.NoError(t, err) - assert.Equal(t, modePlainTar, plan.mode) - assert.Empty(t, plan.commitSHA) - assert.True(t, plan.isGitRepo) - assert.True(t, plan.hasUncommit) -} - -func TestResolveSnapshotPlan_NonGitDir(t *testing.T) { - ctx := t.Context() - dir := t.TempDir() - - plan, err := resolveSnapshotPlan(ctx, newGitRepo(dir), nil, nil) - require.NoError(t, err) - assert.Equal(t, modePlainTar, plan.mode) - assert.False(t, plan.isGitRepo) - - // A git ref on a non-git directory is an error. - _, err = resolveSnapshotPlan(ctx, newGitRepo(dir), &gitRef{Branch: new("main")}, nil) - require.Error(t, err) - assert.Contains(t, err.Error(), "not a git repository") -} - -func TestResolveSnapshotPlan_IncludePaths(t *testing.T) { - ctx := t.Context() - repo := newTestRepo(t) - writeRepoFile(t, repo, "src/model.py", "1") - writeRepoFile(t, repo, "configs/train.yaml", "x") - sha := commitAll(t, repo, "init") - - // All include paths exist at the commit. - plan, err := resolveSnapshotPlan(ctx, newGitRepo(repo), &gitRef{Commit: new(sha)}, []string{"src", "configs"}) - require.NoError(t, err) - assert.Equal(t, modeGitArchive, plan.mode) - assert.Equal(t, []string{"src", "configs"}, plan.includePaths) - - // A missing include path fails fast. - _, err = resolveSnapshotPlan(ctx, newGitRepo(repo), &gitRef{Commit: new(sha)}, []string{"src", "missing"}) - require.Error(t, err) - assert.Contains(t, err.Error(), "missing") -} diff --git a/experimental/air/cmd/snapshot_test.go b/experimental/air/cmd/snapshot_test.go index d94fe005fc9..60fcf7e5610 100644 --- a/experimental/air/cmd/snapshot_test.go +++ b/experimental/air/cmd/snapshot_test.go @@ -1,16 +1,10 @@ package aircmd import ( - "context" - "io" "os" - "path" "path/filepath" "testing" - "github.com/databricks/cli/libs/filer" - "github.com/databricks/cli/libs/testserver" - "github.com/databricks/databricks-sdk-go" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) @@ -21,8 +15,7 @@ func TestResolveRootPath(t *testing.T) { require.NoError(t, os.MkdirAll(filepath.Join(dir, "proj"), 0o755)) // root_path "." resolves against configDir to an absolute path whose basename is - // the real directory name — not "." (which would name the tarball ._.tar.gz, - // colliding with the AppleDouble exclude pattern the remote strips). + // the real directory name — not ".". got, err := resolveRootPath(ctx, ".", filepath.Join(dir, "proj")) require.NoError(t, err) assert.True(t, filepath.IsAbs(got)) @@ -38,118 +31,3 @@ func TestResolveRootPath(t *testing.T) { _, err = resolveRootPath(ctx, "missing", dir) require.Error(t, err) } - -// newSnapshotTestClient returns a workspace client backed by the in-process fake, -// which models workspace get-status / import-file with real state. -func newSnapshotTestClient(t *testing.T) *databricks.WorkspaceClient { - t.Helper() - server := testserver.New(t) - t.Cleanup(server.Close) - testserver.AddDefaultHandlers(server) - w, err := databricks.NewWorkspaceClient(&databricks.Config{Host: server.URL, Token: "token"}) - require.NoError(t, err) - return w -} - -// testUploader builds a snapshotUploader whose tar store and sidecar store both live -// under distinct workspace roots on the fake server. -func testUploader(t *testing.T, w *databricks.WorkspaceClient, tarBase, sidecarBase string) snapshotUploader { - t.Helper() - tarStore, err := filer.NewWorkspaceFilesClient(w, tarBase) - require.NoError(t, err) - sidecarStore, err := filer.NewWorkspaceFilesClient(w, sidecarBase) - require.NoError(t, err) - return snapshotUploader{tarStore: tarStore, sidecarStore: sidecarStore, tarBase: tarBase, sidecarBase: sidecarBase} -} - -func TestRunSnapshot_GitArchive(t *testing.T) { - ctx := t.Context() - w := newSnapshotTestClient(t) - repo := newTestRepo(t) - writeRepoFile(t, repo, "train.py", "print()") - sha := commitAll(t, repo, "init") - - up := testUploader(t, w, "/Workspace/Users/me/.air/repo_snapshots/repo", "/Workspace/Users/me/.air/cli_launch/exp/run") - res, err := runSnapshot(ctx, up, repo, &snapshotSourceConfig{RootPath: repo, Git: &gitRef{Commit: &sha}}) - require.NoError(t, err) - - // Tarball is cache-key-named under the tar base, prefixed with the repo dir name - // (the temp dir's basename); a clean git repo yields a git_state sidecar, no diff. - cacheKey := computeSnapshotCacheKey(sha, nil) - wantName := filepath.Base(repo) + "_" + cacheKey[:16] + ".tar.gz" - assert.Equal(t, path.Join(up.tarBase, wantName), res.CodeSourcePath) - assert.Equal(t, path.Join(up.sidecarBase, gitStateName), res.GitStatePath) - assert.Empty(t, res.GitDiffPath) -} - -func TestRunSnapshot_CacheHitSkipsUpload(t *testing.T) { - ctx := t.Context() - w := newSnapshotTestClient(t) - repo := newTestRepo(t) - writeRepoFile(t, repo, "train.py", "print()") - sha := commitAll(t, repo, "init") - - up := testUploader(t, w, "/Workspace/Users/me/.air/repo_snapshots/repo", "/Workspace/Users/me/.air/cli_launch/exp/run") - snap := &snapshotSourceConfig{RootPath: repo, Git: &gitRef{Commit: &sha}} - - // First submission uploads the tarball. - res1, err := runSnapshot(ctx, up, repo, snap) - require.NoError(t, err) - - // Count uploads to the tarball path on a fresh uploader: the second run should - // see the cached tarball via Stat and not re-upload it. - writes := &countingFiler{Filer: up.tarStore} - up2 := up - up2.tarStore = writes - res2, err := runSnapshot(ctx, up2, repo, snap) - require.NoError(t, err) - - assert.Equal(t, res1.CodeSourcePath, res2.CodeSourcePath) - assert.Zero(t, writes.writes, "cache hit must not re-upload the tarball") -} - -func TestRunSnapshot_PlainTarDirty(t *testing.T) { - ctx := t.Context() - w := newSnapshotTestClient(t) - repo := newTestRepo(t) - writeRepoFile(t, repo, "train.py", "print()") - commitAll(t, repo, "init") - writeRepoFile(t, repo, "train.py", "print('wip')") // dirty, no git ref - - up := testUploader(t, w, "/Workspace/Users/me/.air/repo_snapshots/repo", "/Workspace/Users/me/.air/cli_launch/exp/run") - res, err := runSnapshot(ctx, up, repo, &snapshotSourceConfig{RootPath: repo}) - require.NoError(t, err) - - // Plain tar is timestamp-named (not cache-key-named); a dirty tree captures both - // the state and the diff sidecar. - assert.Contains(t, res.CodeSourcePath, path.Join(up.tarBase, filepath.Base(repo)+"_")) - assert.Equal(t, path.Join(up.sidecarBase, gitStateName), res.GitStatePath) - assert.Equal(t, path.Join(up.sidecarBase, gitDiffName), res.GitDiffPath) -} - -func TestRunSnapshot_NonGitDir(t *testing.T) { - ctx := t.Context() - w := newSnapshotTestClient(t) - dir := t.TempDir() - writeRepoFile(t, dir, "train.py", "print()") - - up := testUploader(t, w, "/Workspace/Users/me/.air/repo_snapshots/proj", "/Workspace/Users/me/.air/cli_launch/exp/run") - res, err := runSnapshot(ctx, up, dir, &snapshotSourceConfig{RootPath: dir}) - require.NoError(t, err) - - // Non-git dir: plain tar, and no provenance sidecars. - assert.NotEmpty(t, res.CodeSourcePath) - assert.Empty(t, res.GitStatePath) - assert.Empty(t, res.GitDiffPath) -} - -// countingFiler wraps a Filer to count Write calls, for asserting cache-hit skips. -type countingFiler struct { - filer.Filer - writes int -} - -func (c *countingFiler) Write(ctx context.Context, name string, reader io.Reader, mode ...filer.WriteMode) error { - c.writes++ - return c.Filer.Write(ctx, name, reader, mode...) -} From 31281802ae74d6d8250ee5bd2aa8f2dacfa4daf0 Mon Sep 17 00:00:00 2001 From: vinchenzo-db Date: Thu, 16 Jul 2026 21:43:05 +0000 Subject: [PATCH 10/11] [air] Emit job permissions from config + add --permissions flag MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The run config already parsed a permissions block but the converter silently dropped it, so it never took effect. Wire cfg.Permissions through to the emitted bundle's job (level + one of user_name/group_name/service_principal_name), and add a repeatable --permissions flag ("CAN_VIEW=group_name:users") that merges into the config's block; flag and YAML grants are re-validated together. This matters because a DABs-deployed job is stamped edit_mode=UI_LOCKED — its permissions cannot be edited via the UI or permissions API after deploy ("Modify bundle sources and redeploy to edit this job"), so declaring them in the bundle is the only way to grant access on AIR-via-DABs jobs. Verified on staging (dbc-04ac0685-8857): `air run --permissions "CAN_VIEW=group_name:users"` produces a job whose ACL grants the users group CAN_VIEW. Co-authored-by: Isaac --- experimental/air/cmd/exportbundle.go | 36 +++++++++++++++++ experimental/air/cmd/exportbundle_test.go | 43 ++++++++++++++++++++ experimental/air/cmd/run.go | 49 +++++++++++++++++++++++ 3 files changed, 128 insertions(+) diff --git a/experimental/air/cmd/exportbundle.go b/experimental/air/cmd/exportbundle.go index 6f18c009433..fe068072017 100644 --- a/experimental/air/cmd/exportbundle.go +++ b/experimental/air/cmd/exportbundle.go @@ -124,6 +124,18 @@ type exportedJob struct { // EnvironmentVariablesKey. Emitted only when the run declares env_variables or // secrets. EnvironmentVariables []exportedEnvVarProfile `yaml:"environment_variables,omitempty"` + // Permissions are the job's ACL grants, from the run's permissions block. + // Emitted only when the run declares any. + Permissions []exportedPermission `yaml:"permissions,omitempty"` +} + +// exportedPermission is one job ACL grant: a level plus exactly one principal. +// Matches the DABs job permissions shape. +type exportedPermission struct { + Level string `yaml:"level"` + UserName string `yaml:"user_name,omitempty"` + GroupName string `yaml:"group_name,omitempty"` + ServicePrincipalName string `yaml:"service_principal_name,omitempty"` } type exportedTask struct { @@ -241,6 +253,7 @@ func convertToBundle(cfg *runConfig) *exportedBundle { Spec: exportBundleEnvSpec(cfg), }}, EnvironmentVariables: profiles, + Permissions: exportedPermissions(cfg), }, }, }, @@ -248,6 +261,29 @@ func convertToBundle(cfg *runConfig) *exportedBundle { } } +// exportedPermissions maps the run's permissions block to job ACL grants, or nil +// when none are declared. Each entry carries the level and its single principal; +// validation (exactly one principal, non-empty level) already ran in runConfig. +func exportedPermissions(cfg *runConfig) []exportedPermission { + if len(cfg.Permissions) == 0 { + return nil + } + out := make([]exportedPermission, 0, len(cfg.Permissions)) + for _, p := range cfg.Permissions { + e := exportedPermission{Level: p.Level} + switch { + case p.UserName != nil: + e.UserName = *p.UserName + case p.GroupName != nil: + e.GroupName = *p.GroupName + case p.ServicePrincipalName != nil: + e.ServicePrincipalName = *p.ServicePrincipalName + } + out = append(out, e) + } + return out +} + // envVarProfiles builds the job-level env-var profile list from the run's // env_variables and secrets, or nil when there are none. Plain values are inline; // each secret (ENV_VAR -> "scope/key") becomes a {{secrets/scope/key}} reference diff --git a/experimental/air/cmd/exportbundle_test.go b/experimental/air/cmd/exportbundle_test.go index b2969086511..3ded218969f 100644 --- a/experimental/air/cmd/exportbundle_test.go +++ b/experimental/air/cmd/exportbundle_test.go @@ -44,6 +44,49 @@ func TestConvertToBundleBasic(t *testing.T) { assert.Empty(t, task.EnvironmentVariablesKey) } +func TestConvertToBundlePermissions(t *testing.T) { + cfg := &runConfig{ + ExperimentName: "exp", + Command: new("python train.py"), + Compute: &computeConfig{AcceleratorType: "GPU_1xA10", NumAccelerators: 1}, + Permissions: []permission{ + {Level: "CAN_VIEW", GroupName: new("users")}, + {Level: "CAN_MANAGE", UserName: new("a@b.com")}, + }, + } + + job := convertToBundle(cfg).Resources.Jobs["exp"] + require.Len(t, job.Permissions, 2) + assert.Equal(t, exportedPermission{Level: "CAN_VIEW", GroupName: "users"}, job.Permissions[0]) + assert.Equal(t, exportedPermission{Level: "CAN_MANAGE", UserName: "a@b.com"}, job.Permissions[1]) + + // No permissions declared -> field omitted (no empty permissions block). + assert.Empty(t, convertToBundle(&runConfig{ + ExperimentName: "exp", Command: new("x"), + Compute: &computeConfig{AcceleratorType: "GPU_1xA10", NumAccelerators: 1}, + }).Resources.Jobs["exp"].Permissions) +} + +func TestParsePermissions(t *testing.T) { + got, err := parsePermissions([]string{ + "CAN_VIEW=group_name:users", + "CAN_MANAGE=user_name:a@b.com", + "CAN_RUN=service_principal_name:1234-abcd", + }) + require.NoError(t, err) + require.Len(t, got, 3) + assert.Equal(t, "CAN_VIEW", got[0].Level) + assert.Equal(t, "users", *got[0].GroupName) + assert.Equal(t, "a@b.com", *got[1].UserName) + assert.Equal(t, "1234-abcd", *got[2].ServicePrincipalName) + + // Malformed inputs error rather than silently drop. + for _, bad := range []string{"CAN_VIEW", "CAN_VIEW=users", "CAN_VIEW=bogus:x"} { + _, err := parsePermissions([]string{bad}) + require.Error(t, err, "expected error for %q", bad) + } +} + func TestConvertToBundleEnvVarsAndSecrets(t *testing.T) { cfg := &runConfig{ ExperimentName: "exp", diff --git a/experimental/air/cmd/run.go b/experimental/air/cmd/run.go index e9e1c67c6f2..c6d07447aa9 100644 --- a/experimental/air/cmd/run.go +++ b/experimental/air/cmd/run.go @@ -4,6 +4,7 @@ import ( "errors" "fmt" "strconv" + "strings" "github.com/databricks/cli/cmd/root" "github.com/databricks/cli/libs/cmdctx" @@ -30,6 +31,7 @@ func newRunCommand() *cobra.Command { overrides []string dryRun bool idempotencyKey string + permissions []string ) cmd := &cobra.Command{ @@ -46,6 +48,7 @@ The workload is described by a YAML config file (see --file).`, cmd.Flags().StringArrayVar(&overrides, "override", nil, "Override a YAML field, e.g. compute.num_accelerators=8 (repeatable)") cmd.Flags().BoolVar(&dryRun, "dry-run", false, "Validate the config and show the generated bundle without submitting") cmd.Flags().StringVar(&idempotencyKey, "idempotency-key", "", "Return the existing run if this key was already used") + cmd.Flags().StringArrayVar(&permissions, "permissions", nil, `Grant a permission on the job, e.g. "CAN_VIEW=group_name:users" or "CAN_MANAGE=user_name:a@b.com" (repeatable). Merged with the config's permissions block.`) _ = cmd.MarkFlagRequired("file") // --dry-run only validates the config locally, so it needs no workspace. @@ -74,6 +77,21 @@ The workload is described by a YAML config file (see --file).`, return err } + // --permissions flags append to the config's permissions block, then the whole + // set is re-validated so flag and YAML grants are held to the same rules. + if len(permissions) > 0 { + parsed, err := parsePermissions(permissions) + if err != nil { + return err + } + cfg.Permissions = append(cfg.Permissions, parsed...) + for i := range cfg.Permissions { + if err := cfg.Permissions[i].validate(); err != nil { + return err + } + } + } + if dryRun { // Show the generated bundle without deploying. bundleYAML, err := renderBundle(cfg, file) @@ -107,3 +125,34 @@ The workload is described by a YAML config file (see --file).`, return cmd } + +// parsePermissions parses --permissions flag values of the form +// "LEVEL=principal_type:name", e.g. "CAN_VIEW=group_name:users". principal_type is +// one of user_name, group_name, service_principal_name. The result is validated by +// the caller alongside any config permissions. +func parsePermissions(specs []string) ([]permission, error) { + out := make([]permission, 0, len(specs)) + for _, spec := range specs { + level, principal, ok := strings.Cut(spec, "=") + if !ok { + return nil, fmt.Errorf("invalid --permissions %q: expected LEVEL=principal_type:name (e.g. CAN_VIEW=group_name:users)", spec) + } + kind, name, ok := strings.Cut(principal, ":") + if !ok { + return nil, fmt.Errorf("invalid --permissions %q: principal must be principal_type:name (e.g. group_name:users)", spec) + } + p := permission{Level: strings.TrimSpace(level)} + switch strings.TrimSpace(kind) { + case "user_name": + p.UserName = &name + case "group_name": + p.GroupName = &name + case "service_principal_name": + p.ServicePrincipalName = &name + default: + return nil, fmt.Errorf("invalid --permissions %q: principal type %q must be user_name, group_name, or service_principal_name", spec, kind) + } + out = append(out, p) + } + return out, nil +} From 2a7f1ab24a16cfcf46bca2406656ad928d228ffe Mon Sep 17 00:00:00 2001 From: vinchenzo-db Date: Thu, 16 Jul 2026 22:09:42 +0000 Subject: [PATCH 11/11] [air] Update acceptance tests for the DABs run path MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit air run no longer submits via runs/submit — it converts to a bundle and deploys/runs. Bring the acceptance suite in line: - Delete run-submit: it tested the retired ephemeral runs/submit path with a git code_source (POST /jobs/runs/submit + tarball upload). air run now deploys via DABs and rejects git code_source, so the scenario no longer exists. - Update run/output.txt: --dry-run now renders the generated databricks.yml (bundle preview) instead of "configuration is valid". - run/with-code-source.yaml: drop the git block so "code_source passes validation" stays a positive case — the DABs path supports working-tree snapshots but rejects git-pinned ones (covered by git-remote.yaml's rejection). Co-authored-by: Isaac --- .../air/run-submit/databricks.yml | 2 - .../experimental/air/run-submit/out.test.toml | 3 - .../experimental/air/run-submit/output.txt | 46 ---------- .../experimental/air/run-submit/run.yaml.tmpl | 11 --- acceptance/experimental/air/run-submit/script | 18 ---- .../experimental/air/run-submit/test.toml | 35 -------- acceptance/experimental/air/run/output.txt | 83 ++++++++++++++++++- .../air/run/with-code-source.yaml | 2 - 8 files changed, 80 insertions(+), 120 deletions(-) delete mode 100644 acceptance/experimental/air/run-submit/databricks.yml delete mode 100644 acceptance/experimental/air/run-submit/out.test.toml delete mode 100644 acceptance/experimental/air/run-submit/output.txt delete mode 100644 acceptance/experimental/air/run-submit/run.yaml.tmpl delete mode 100644 acceptance/experimental/air/run-submit/script delete mode 100644 acceptance/experimental/air/run-submit/test.toml diff --git a/acceptance/experimental/air/run-submit/databricks.yml b/acceptance/experimental/air/run-submit/databricks.yml deleted file mode 100644 index 4a1c612d600..00000000000 --- a/acceptance/experimental/air/run-submit/databricks.yml +++ /dev/null @@ -1,2 +0,0 @@ -bundle: - name: air-run-submit diff --git a/acceptance/experimental/air/run-submit/out.test.toml b/acceptance/experimental/air/run-submit/out.test.toml deleted file mode 100644 index d6187dcb046..00000000000 --- a/acceptance/experimental/air/run-submit/out.test.toml +++ /dev/null @@ -1,3 +0,0 @@ -Local = true -Cloud = false -EnvMatrix.DATABRICKS_BUNDLE_ENGINE = [] diff --git a/acceptance/experimental/air/run-submit/output.txt b/acceptance/experimental/air/run-submit/output.txt deleted file mode 100644 index 8d52eed1dab..00000000000 --- a/acceptance/experimental/air/run-submit/output.txt +++ /dev/null @@ -1,46 +0,0 @@ - -=== submit with a git code_source ->>> [CLI] experimental air run -f run.yaml -Submitted run 555 -View at: [DATABRICKS_URL]/jobs/runs/555 - -=== the ai_runtime_task carries the code_source_path ->>> print_requests.py //api/2.2/jobs/runs/submit -{ - "method": "POST", - "path": "/api/2.2/jobs/runs/submit", - "body": { - "environments": [ - { - "environment_key": "default", - "spec": { - "environment_version": "4" - } - } - ], - "idempotency_token": "[UUID]", - "run_name": "submit-smoke", - "tasks": [ - { - "ai_runtime_task": { - "code_source_path": "/Workspace/Users/[USERNAME]/.air/repo_snapshots/001/[SNAPSHOT_TARBALL]", - "deployments": [ - { - "command_path": "/Workspace/Users/[USERNAME]/.air/cli_launch/submit-smoke/submit-smoke_[RUN_ID]/command.sh", - "compute": { - "accelerator_count": 1, - "accelerator_type": "GPU_1xH100" - } - } - ], - "experiment": "submit-smoke" - }, - "environment_key": "default", - "max_retries": 3, - "retry_on_timeout": true, - "run_if": "ALL_SUCCESS", - "task_key": "submit-smoke" - } - ] - } -} diff --git a/acceptance/experimental/air/run-submit/run.yaml.tmpl b/acceptance/experimental/air/run-submit/run.yaml.tmpl deleted file mode 100644 index 3fdbf48eb85..00000000000 --- a/acceptance/experimental/air/run-submit/run.yaml.tmpl +++ /dev/null @@ -1,11 +0,0 @@ -experiment_name: submit-smoke -command: python train.py -compute: - accelerator_type: GPU_1xH100 - num_accelerators: 1 -code_source: - type: snapshot - snapshot: - root_path: . - git: - commit: COMMIT_SHA diff --git a/acceptance/experimental/air/run-submit/script b/acceptance/experimental/air/run-submit/script deleted file mode 100644 index c44b50e28a1..00000000000 --- a/acceptance/experimental/air/run-submit/script +++ /dev/null @@ -1,18 +0,0 @@ -# Pinned commit dates keep the resolved HEAD SHA — and thus the snapshot cache key -# baked into the tarball name — stable across runs. -export GIT_AUTHOR_DATE="2020-01-01T00:00:00Z" -export GIT_COMMITTER_DATE="2020-01-01T00:00:00Z" -git-repo-init - -# Pin the config to the committed HEAD. git.commit tolerates a dirty working tree -# (the pinned revision is what gets archived), which matters here because the -# acceptance harness writes output.txt into the working dir as the script runs. -sed "s/COMMIT_SHA/$(git rev-parse HEAD)/" run.yaml.tmpl > run.yaml - -title "submit with a git code_source" -trace $CLI experimental air run -f run.yaml - -title "the ai_runtime_task carries the code_source_path" -trace print_requests.py //api/2.2/jobs/runs/submit - -rm -fr .git diff --git a/acceptance/experimental/air/run-submit/test.toml b/acceptance/experimental/air/run-submit/test.toml deleted file mode 100644 index fcdf8fd1242..00000000000 --- a/acceptance/experimental/air/run-submit/test.toml +++ /dev/null @@ -1,35 +0,0 @@ -# A real (non-dry-run) submit that packages a git code_source, uploads the -# tarball + provenance sidecars, and POSTs runs/submit. No bundle deploy, so no -# engine matrix. -RecordRequests = true - -# run.yaml is generated from run.yaml.tmpl at test time (commit SHA templated in); -# it isn't a committed input to diff. -Ignore = ["run.yaml"] - -[EnvMatrix] -DATABRICKS_BUNDLE_ENGINE = [] - -# The SDK probes host reachability with a HEAD request; stub it for determinism. -[[Server]] -Pattern = "HEAD /" -Response.Body = '' - -[[Server]] -Pattern = "POST /api/2.2/jobs/runs/submit" -Response.Body = ''' -{"run_id": 555} -''' - -# The snapshot tarball is named _.tar.gz, where is the -# test's temp-dir basename and the cache key derives from the pinned commit SHA. -# Both are stable given the pinned commit dates in the script, but the temp-dir -# basename varies per run, so collapse the whole tarball filename to a stable token. -[[Repls]] -Old = '[a-zA-Z0-9._-]+_[0-9a-f]{16}\.tar\.gz' -New = '[SNAPSHOT_TARBALL]' - -# The per-run launch directory ends in _<16 hex>; the random suffix varies. -[[Repls]] -Old = 'submit-smoke_[0-9a-f]{16}' -New = 'submit-smoke_[RUN_ID]' diff --git a/acceptance/experimental/air/run/output.txt b/acceptance/experimental/air/run/output.txt index a753eabd198..edb5705b343 100644 --- a/acceptance/experimental/air/run/output.txt +++ b/acceptance/experimental/air/run/output.txt @@ -1,7 +1,45 @@ === dry-run (text) >>> [CLI] experimental air run -f valid.yaml --dry-run -Dry run: configuration for "smoke-test" is valid; not submitting. +Dry run: "smoke-test" would deploy as this bundle (not deploying): + +# Generated by `air export-bundle` from valid.yaml. +# +# Deploys the same workload as a durable Jobs resource. `bundle deploy` syncs this +# folder (including command.sh and your code) to the workspace; the task's +# command_path points at the synced command.sh. Before deploying, add a +# `targets` block with your workspace host. See the ai-compute DABs examples. +bundle: + name: smoke-test +resources: + jobs: + smoke-test: + name: smoke-test + tasks: + - task_key: smoke-test + environment_key: default + max_retries: 3 + ai_runtime_task: + experiment: smoke-test + deployments: + - name: worker + command_path: command.sh + compute: + accelerator_type: GPU_1xH100 + accelerator_count: 1 + environments: + - environment_key: default + spec: + environment_version: "4" +experimental: + immutable_folder: true + +# Appended by `air run` (host resolved from your CLI profile at deploy time): +targets: + dev: + mode: development + default: true + === dry-run (json) >>> [CLI] experimental air run -f valid.yaml --dry-run -o json @@ -10,7 +48,8 @@ Dry run: configuration for "smoke-test" is valid; not submitting. "ts": "[TIMESTAMP]", "data": { "status": "DRY_RUN_OK", - "dry_run": true + "dry_run": true, + "bundle": "# Generated by `air export-bundle` from valid.yaml.\n#\n# Deploys the same workload as a durable Jobs resource. `bundle deploy` syncs this\n# folder (including command.sh and your code) to the workspace; the task's\n# command_path points at the synced command.sh. Before deploying, add a\n# `targets` block with your workspace host. See the ai-compute DABs examples.\nbundle:\n name: smoke-test\nresources:\n jobs:\n smoke-test:\n name: smoke-test\n tasks:\n - task_key: smoke-test\n environment_key: default\n max_retries: 3\n ai_runtime_task:\n experiment: smoke-test\n deployments:\n - name: worker\n command_path: command.sh\n compute:\n accelerator_type: GPU_1xH100\n accelerator_count: 1\n environments:\n - environment_key: default\n spec:\n environment_version: \"4\"\nexperimental:\n immutable_folder: true\n\n# Appended by `air run` (host resolved from your CLI profile at deploy time):\ntargets:\n dev:\n mode: development\n default: true\n" } } @@ -28,7 +67,45 @@ Exit code: 1 === code_source config passes validation >>> [CLI] experimental air run -f with-code-source.yaml --dry-run -Dry run: configuration for "smoke-test" is valid; not submitting. +Dry run: "smoke-test" would deploy as this bundle (not deploying): + +# Generated by `air export-bundle` from with-code-source.yaml. +# +# Deploys the same workload as a durable Jobs resource. `bundle deploy` syncs this +# folder (including command.sh and your code) to the workspace; the task's +# command_path points at the synced command.sh. Before deploying, add a +# `targets` block with your workspace host. See the ai-compute DABs examples. +bundle: + name: smoke-test +resources: + jobs: + smoke-test: + name: smoke-test + tasks: + - task_key: smoke-test + environment_key: default + max_retries: 3 + ai_runtime_task: + experiment: smoke-test + deployments: + - name: worker + command_path: command.sh + compute: + accelerator_type: GPU_1xH100 + accelerator_count: 1 + environments: + - environment_key: default + spec: + environment_version: "4" +experimental: + immutable_folder: true + +# Appended by `air run` (host resolved from your CLI profile at deploy time): +targets: + dev: + mode: development + default: true + === git.remote is rejected >>> [CLI] experimental air run -f git-remote.yaml --dry-run diff --git a/acceptance/experimental/air/run/with-code-source.yaml b/acceptance/experimental/air/run/with-code-source.yaml index 86a32c138cb..7171ddb0ac5 100644 --- a/acceptance/experimental/air/run/with-code-source.yaml +++ b/acceptance/experimental/air/run/with-code-source.yaml @@ -7,7 +7,5 @@ code_source: type: snapshot snapshot: root_path: . - git: - branch: main include_paths: - src