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 diff --git a/experimental/air/cmd/exportbundle.go b/experimental/air/cmd/exportbundle.go new file mode 100644 index 00000000000..fe068072017 --- /dev/null +++ b/experimental/air/cmd/exportbundle.go @@ -0,0 +1,337 @@ +package aircmd + +import ( + "fmt" + "maps" + "path/filepath" + "strings" + + "go.yaml.in/yaml/v3" +) + +// 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). +// +// 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. +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 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 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. 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 + + // 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 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 + // 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"` + 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"` +} + +type exportedJob struct { + Name string `yaml:"name"` + Tasks []exportedTask `yaml:"tasks"` + Environments []exportedEnvironment `yaml:"environments"` + // 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"` + // 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 { + TaskKey string `yaml:"task_key"` + EnvironmentKey string `yaml:"environment_key"` + // 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. +// 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"` +} + +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. +const databricksAITokenPrefix = "databricks_ai_v" + +// databricksAIBaseEnvironment is the system base_environment id the databricks-ai +// token resolves to. +const databricksAIBaseEnvironment = "workspace-base-environments/" + +// 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{{ + // The deployment name is cosmetic (the wire payload carries none). + Name: "worker", + CommandPath: 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, + } + + // 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 + } + + // 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{ + Jobs: map[string]exportedJob{ + cfg.ExperimentName: { + Name: cfg.ExperimentName, + Tasks: []exportedTask{tsk}, + Environments: []exportedEnvironment{{ + EnvironmentKey: aiRuntimeEnvironmentKey, + Spec: exportBundleEnvSpec(cfg), + }}, + EnvironmentVariables: profiles, + Permissions: exportedPermissions(cfg), + }, + }, + }, + Experimental: &exportedExperimental{ImmutableFolder: true}, + } +} + +// 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 +// Jobs resolves 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)) + maps.Copy(variables, cfg.EnvVariables) + for envVar, secretRef := range cfg.Secrets { + variables[envVar] = "{{secrets/" + secretRef + "}}" + } + return []exportedEnvVarProfile{{ + EnvironmentVariablesKey: aiRuntimeEnvVarsKey, + Variables: variables, + }} +} + +// 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 { + 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..3ded218969f --- /dev/null +++ b/experimental/air/cmd/exportbundle_test.go @@ -0,0 +1,244 @@ +package aircmd + +import ( + "os" + "path/filepath" + "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] + // 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) + + // 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 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", + 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 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 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"), + 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") + assert.Contains(t, out, "immutable_folder: true") + + // 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.Contains(t, body, "Generated by `air export-bundle`") +} diff --git a/experimental/air/cmd/list.go b/experimental/air/cmd/list.go index 7d5c703a39c..b9672aad796 100644 --- a/experimental/air/cmd/list.go +++ b/experimental/air/cmd/list.go @@ -233,8 +233,10 @@ type jobsScanStrategy struct { } func newJobsScanStrategy(ctx context.Context, w *databricks.WorkspaceClient, q listQuery) *jobsScanStrategy { + // 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{ - RunType: jobs.RunTypeSubmitRun, ExpandTasks: true, Limit: jobsPageLimit, ActiveOnly: q.activeOnly, diff --git a/experimental/air/cmd/list_test.go b/experimental/air/cmd/list_test.go index f70330240e5..7bd2bbc8d69 100644 --- a/experimental/air/cmd/list_test.go +++ b/experimental/air/cmd/list_test.go @@ -99,6 +99,38 @@ 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 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", + }).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..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" @@ -18,6 +19,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 so the user + // can see exactly what would be deployed on their behalf. + Bundle string `json:"bundle,omitempty"` } func newRunCommand() *cobra.Command { @@ -27,6 +31,7 @@ func newRunCommand() *cobra.Command { overrides []string dryRun bool idempotencyKey string + permissions []string ) cmd := &cobra.Command{ @@ -41,8 +46,9 @@ 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") + 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. @@ -71,15 +77,38 @@ 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) + 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 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 @@ -96,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 +} diff --git a/experimental/air/cmd/rundabs.go b/experimental/air/cmd/rundabs.go new file mode 100644 index 00000000000..fefffe37eb9 --- /dev/null +++ b/experimental/air/cmd/rundabs.go @@ -0,0 +1,314 @@ +package aircmd + +// air run submits a training workload by converting the run YAML to a Databricks +// 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. +// +// 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. +// +// 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" + "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" +) + +// 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" + +// 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 + // (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 does not map onto deploy+run (two calls, not one tokened submit). + 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(ctx, cfg, configPath) + if err != nil { + return 0, "", err + } + defer cleanup() + + // 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")) + + 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) + } + + // 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 { + 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.") + + 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 deploys: the +// 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 + } + 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 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 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() + return "", func() {}, err + } + } + + if err := os.WriteFile(filepath.Join(bundleRoot, "databricks.yml"), []byte(body), 0o600); err != nil { + cleanup() + return "", func() {}, err + } + + // 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 + } + } + 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 +// here (the bundle resolves it from auth at deploy time). +func bundleTargetsBlock() string { + 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" +} + +// 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 ProcessBundle (load, validate, build, sync, deploy) +// via a synthetic carrier command (see newBundleCarrierCommand). +// +// 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, + 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 +} + +// 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"} + + 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. + if w.Config.Profile != "" { + _ = cmd.Flags().Set("profile", w.Config.Profile) + } + + 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 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 { + 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 wait.RunId, nil +} + +// isRunnableJob filters resources.Lookup to job resources. +func isRunnableJob(ref bundleresources.Reference) bool { + _, ok := ref.Resource.(*resources.Job) + return ok +} diff --git a/experimental/air/cmd/rundabs_test.go b/experimental/air/cmd/rundabs_test.go new file mode 100644 index 00000000000..ea4f7331a8b --- /dev/null +++ b/experimental/air/cmd/rundabs_test.go @@ -0,0 +1,154 @@ +package aircmd + +import ( + "io/fs" + "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) { + configPath := writeConfigFile(t, "train.yaml", minimalConfig) + cfg, err := loadRunConfig(configPath) + require.NoError(t, err) + + root, cleanup, err := writeBundleProject(t.Context(), cfg, configPath) + 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") + + // 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, commandScript("python train.py"), string(script)) + + // cleanup removes the temp root. + cleanup() + _, err = os.Stat(root) + assert.ErrorIs(t, err, fs.ErrNotExist) +} + +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)) + + 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, configPath) + 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: the artifact writes happen after the tree copy. + src := t.TempDir() + require.NoError(t, os.WriteFile(filepath.Join(src, bundleCommandScript), []byte("STALE"), 0o644)) + + 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, 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, commandScript("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. + 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) +} + +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{}})) +} 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") - }) -} diff --git a/experimental/air/cmd/runupload.go b/experimental/air/cmd/runupload.go index fb9ca00b987..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(); { @@ -102,8 +102,15 @@ 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. + // 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 { @@ -122,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 0c87524735d..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) { @@ -84,6 +91,10 @@ func TestBuildArtifacts_InlineRequirementsAndParameters(t *testing.T) { } 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"), 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...) -}