diff --git a/acceptance/experimental/air/get-ai-runtime/output.txt b/acceptance/experimental/air/get-ai-runtime/output.txt index c47790eaa31..382719320b6 100644 --- a/acceptance/experimental/air/get-ai-runtime/output.txt +++ b/acceptance/experimental/air/get-ai-runtime/output.txt @@ -1,6 +1,6 @@ === get (text) ->>> [CLI] experimental air get 123 +>>> [CLI] experimental air get run 123 ╭─ Configuration ────────────────────────────────────────────────╮ │ │ @@ -35,7 +35,7 @@ Run URL: [DATABRICKS_URL]/jobs/runs/123?o=[NUMID] MLflow URL: [DATABRICKS_URL]/ml/experiments/exp1/runs/run1 === get (json) ->>> [CLI] experimental air get 123 -o json +>>> [CLI] experimental air get run 123 -o json { "v": 1, "ts": "[TIMESTAMP]", diff --git a/acceptance/experimental/air/get-ai-runtime/script b/acceptance/experimental/air/get-ai-runtime/script index 3f41b089cdd..ec17aef7a03 100644 --- a/acceptance/experimental/air/get-ai-runtime/script +++ b/acceptance/experimental/air/get-ai-runtime/script @@ -5,7 +5,7 @@ $CLI workspace mkdirs "/Workspace/Users/user@example.com/.air/cli_launch/my-exp/ $CLI workspace import "/Workspace/Users/user@example.com/.air/cli_launch/my-exp/my-exp_abc/training_config.yaml" --file training_config.yaml --format AUTO &> LOG.import title "get (text)" -trace $CLI experimental air get 123 +trace $CLI experimental air get run 123 title "get (json)" -trace $CLI experimental air get 123 -o json +trace $CLI experimental air get run 123 -o json diff --git a/acceptance/experimental/air/get/output.txt b/acceptance/experimental/air/get/output.txt index 6e51d7debf9..4bc127d79b4 100644 --- a/acceptance/experimental/air/get/output.txt +++ b/acceptance/experimental/air/get/output.txt @@ -1,6 +1,6 @@ === get (text) ->>> [CLI] experimental air get 123 +>>> [CLI] experimental air get run 123 ╭─ Configuration ────────────────────────────────────────────────╮ │ │ @@ -35,7 +35,7 @@ Run URL: [DATABRICKS_URL]/jobs/runs/123?o=[NUMID] MLflow URL: [DATABRICKS_URL]/ml/experiments/exp1/runs/run1 === get (json) ->>> [CLI] experimental air get 123 -o json +>>> [CLI] experimental air get run 123 -o json { "v": 1, "ts": "[TIMESTAMP]", @@ -52,13 +52,13 @@ MLflow URL: [DATABRICKS_URL]/ml/experiments/exp1/runs/run1 } === invalid run id ->>> [CLI] experimental air get notanumber +>>> [CLI] experimental air get run notanumber Error: invalid JOB_RUN_ID "notanumber": must be a positive integer Exit code: 1 === invalid run id (json) ->>> [CLI] experimental air get notanumber -o json +>>> [CLI] experimental air get run notanumber -o json { "v": 1, "ts": "[TIMESTAMP]", diff --git a/acceptance/experimental/air/get/script b/acceptance/experimental/air/get/script index ee66b4aff04..b775d06a48d 100644 --- a/acceptance/experimental/air/get/script +++ b/acceptance/experimental/air/get/script @@ -1,11 +1,11 @@ title "get (text)" -trace $CLI experimental air get 123 +trace $CLI experimental air get run 123 title "get (json)" -trace $CLI experimental air get 123 -o json +trace $CLI experimental air get run 123 -o json title "invalid run id" -errcode trace $CLI experimental air get notanumber +errcode trace $CLI experimental air get run notanumber title "invalid run id (json)" -errcode trace $CLI experimental air get notanumber -o json +errcode trace $CLI experimental air get run notanumber -o json diff --git a/acceptance/experimental/air/help/output.txt b/acceptance/experimental/air/help/output.txt index ee89e778d6b..9a23a15818a 100644 --- a/acceptance/experimental/air/help/output.txt +++ b/acceptance/experimental/air/help/output.txt @@ -11,8 +11,9 @@ Usage: Available Commands: cancel Cancel one or more runs - get Show status, configuration, and timing details for a specific run - list List your active runs for the current profile (use --all-status for finished runs) + delete Delete a resource (bundle) + get Show details for a specific resource (run or bundle) + list List a collection of resources (runs or bundles) logs Stream or fetch logs for a run register-image Mirror a Docker image into the workspace registry run Submit a training workload from a YAML config @@ -29,17 +30,17 @@ Global Flags: Use "databricks experimental air [command] --help" for more information about a command. === list help ->>> [CLI] experimental air list --help +>>> [CLI] experimental air list runs --help List your active runs for the current profile (use --all-status for finished runs) Usage: - databricks experimental air list [flags] + databricks experimental air list runs [flags] Flags: --all-status Show runs in all states (default: active only) --all-users Show runs from all users --filter stringArray Filter runs, e.g. experiment=foo* (repeatable) - -h, --help help for list + -h, --help help for runs --limit int Maximum number of runs to show (default 20) Global Flags: diff --git a/acceptance/experimental/air/help/script b/acceptance/experimental/air/help/script index 81f3907e4f5..29a80575de6 100644 --- a/acceptance/experimental/air/help/script +++ b/acceptance/experimental/air/help/script @@ -5,4 +5,4 @@ title "help" trace $CLI experimental air --help title "list help" -trace $CLI experimental air list --help +trace $CLI experimental air list runs --help diff --git a/acceptance/experimental/air/list/output.txt b/acceptance/experimental/air/list/output.txt index e27cfb0bf39..3fb7eea2ed4 100644 --- a/acceptance/experimental/air/list/output.txt +++ b/acceptance/experimental/air/list/output.txt @@ -1,11 +1,11 @@ === list (text) ->>> [CLI] experimental air list - Run ID Experiment Status Started Duration MLflow User Accelerators - [NUMID] qwen-train ● SUCCESS [TIMESTAMP] 12s …/runs/run1 [USERNAME] 8x H100 +>>> [CLI] experimental air list runs + Run ID Experiment Status Started Duration MLflow User Accelerators + [NUMID] qwen-train ● RUNNING [TIMESTAMP] [DURATION] …/runs/run1 [USERNAME] 8x H100 === list (json) ->>> [CLI] experimental air list -o json +>>> [CLI] experimental air list runs -o json { "v": 1, "ts": "[TIMESTAMP]", @@ -15,7 +15,7 @@ "run_id": "[NUMID]", "run_name": "qwen-train", "user": "[USERNAME]", - "status": "SUCCESS", + "status": "RUNNING", "started_at": "[TIMESTAMP]", "is_sweep": false } @@ -23,13 +23,14 @@ } } -=== list --all-status (text, via AiTrainingService index) ->>> [CLI] experimental air list --all-status - Run ID Experiment Status Started Duration MLflow User Accelerators - [NUMID] qwen-train ● SUCCESS [TIMESTAMP] 12s …/runs/run1 [USERNAME] 8x H100 +=== list --all-status (text) +>>> [CLI] experimental air list runs --all-status + Run ID Experiment Status Started Duration MLflow User Accelerators + [NUMID] qwen-train ● RUNNING [TIMESTAMP] [DURATION] …/runs/run1 [USERNAME] 8x H100 + [NUMID] llama-train ● SUCCESS [TIMESTAMP] 1m 11s …/runs/run1 [USERNAME] 1x A10 === list --all-status (json) ->>> [CLI] experimental air list --all-status -o json +>>> [CLI] experimental air list runs --all-status -o json { "v": 1, "ts": "[TIMESTAMP]", @@ -39,6 +40,14 @@ "run_id": "[NUMID]", "run_name": "qwen-train", "user": "[USERNAME]", + "status": "RUNNING", + "started_at": "[TIMESTAMP]", + "is_sweep": false + }, + { + "run_id": "[NUMID]", + "run_name": "llama-train", + "user": "[USERNAME]", "status": "SUCCESS", "started_at": "[TIMESTAMP]", "is_sweep": false diff --git a/acceptance/experimental/air/list/script b/acceptance/experimental/air/list/script index df547794c6f..02e862c344c 100644 --- a/acceptance/experimental/air/list/script +++ b/acceptance/experimental/air/list/script @@ -1,11 +1,11 @@ title "list (text)" -trace $CLI experimental air list +trace $CLI experimental air list runs title "list (json)" -trace $CLI experimental air list -o json +trace $CLI experimental air list runs -o json -title "list --all-status (text, via AiTrainingService index)" -trace $CLI experimental air list --all-status +title "list --all-status (text)" +trace $CLI experimental air list runs --all-status title "list --all-status (json)" -trace $CLI experimental air list --all-status -o json +trace $CLI experimental air list runs --all-status -o json diff --git a/acceptance/experimental/air/list/test.toml b/acceptance/experimental/air/list/test.toml index 82f1f829d85..89ab3e07fe8 100644 --- a/acceptance/experimental/air/list/test.toml +++ b/acceptance/experimental/air/list/test.toml @@ -2,9 +2,12 @@ [EnvMatrix] DATABRICKS_BUNDLE_ENGINE = [] -# Disable the on-disk run cache so --all-status output is deterministic across runs. -[Env] -DATABRICKS_CACHE_ENABLED = "false" +# A running run's duration is measured against now (no end_time), so mask the +# hours-minutes-seconds cell to keep the golden deterministic. +[[Repls]] +Old = '\d+h \d+m \d+s' +New = "[DURATION]" +Order = 20 # The SDK occasionally probes host reachability with a HEAD request; stub it so # the test is deterministic. @@ -12,11 +15,13 @@ DATABRICKS_CACHE_ENABLED = "false" Pattern = "HEAD /" Response.Body = '' -# `air list` reads Jobs runs/list directly (expand_tasks), filtering to AIR runs. -# The current runs carry an ai_runtime_task (not modeled by the typed SDK), which -# is why the CLI parses the response raw. Both runs belong to the default user -# (tester@databricks.com, from the built-in scim/v2/Me handler). The second run -# has no AI task and must be filtered out. +# `air list` reads Jobs runs/list (expand_tasks), filtering to AIR runs. It first +# tries the server-side AI-Runtime filter (filter=runType=AI_RUNTIME); this fake +# ignores query params and serves the same body, exercising the parse + client-side +# filtering (AIR task-shape, active-only). All runs belong to the default user +# (tester@databricks.com, from the built-in scim/v2/Me handler). qwen-train is +# RUNNING so it shows in the active default; the terminated run shows only under +# --all-status; the notebook run has no AI task and is filtered out. [[Server]] Pattern = "GET /api/2.2/jobs/runs/list" Response.Body = ''' @@ -27,8 +32,7 @@ Response.Body = ''' "run_name": "qwen-train", "creator_user_name": "tester@databricks.com", "start_time": 1717608759000, - "end_time": 1717608771000, - "state": {"life_cycle_state": "TERMINATED", "result_state": "SUCCESS"}, + "state": {"life_cycle_state": "RUNNING"}, "tasks": [{ "run_id": 334747067049497, "ai_runtime_task": { @@ -37,6 +41,21 @@ Response.Body = ''' } }] }, + { + "run_id": 334747067049400, + "run_name": "llama-train", + "creator_user_name": "tester@databricks.com", + "start_time": 1717608700000, + "end_time": 1717608771000, + "state": {"life_cycle_state": "TERMINATED", "result_state": "SUCCESS"}, + "tasks": [{ + "run_id": 334747067049401, + "ai_runtime_task": { + "experiment": "/Users/tester@databricks.com/llama-train", + "deployments": [{"compute": {"accelerator_type": "GPU_1xA10", "accelerator_count": 1}}] + } + }] + }, { "run_id": 999000999000, "run_name": "not-an-air-run", @@ -55,33 +74,3 @@ Pattern = "GET /api/2.2/jobs/runs/get-output" Response.Body = ''' {"ai_runtime_task_output": {"mlflow_experiment_id": "exp1", "mlflow_run_id": "run1"}} ''' - -# `air list --all-status` scoped to the current user is served by the -# AiTrainingService index: it returns cheap (job_run_id, submit_time) pairs, which -# the CLI orders and then hydrates into full runs via runs/get. -[[Server]] -Pattern = "GET /api/2.0/ai-training/workflows" -Response.Body = ''' -{"training_workflows": [{"job_run_id": "334747067049496", "submit_time": "2024-06-05T17:32:39Z"}]} -''' - -# runs/get hydrates one index id into the same shape as a runs/list element. -[[Server]] -Pattern = "GET /api/2.2/jobs/runs/get" -Response.Body = ''' -{ - "run_id": 334747067049496, - "run_name": "qwen-train", - "creator_user_name": "tester@databricks.com", - "start_time": 1717608759000, - "end_time": 1717608771000, - "state": {"life_cycle_state": "TERMINATED", "result_state": "SUCCESS"}, - "tasks": [{ - "run_id": 334747067049497, - "ai_runtime_task": { - "experiment": "/Users/tester@databricks.com/qwen-train", - "deployments": [{"compute": {"accelerator_type": "GPU_8xH100", "accelerator_count": 8}}] - } - }] -} -''' 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/air.go b/experimental/air/cmd/air.go index fbf40a34b52..3b0987fb8b9 100644 --- a/experimental/air/cmd/air.go +++ b/experimental/air/cmd/air.go @@ -20,6 +20,7 @@ experimental and may change in future versions.`, cmd.AddCommand(newRunCommand()) cmd.AddCommand(newGetCommand()) cmd.AddCommand(newListCommand()) + cmd.AddCommand(newDeleteCommand()) cmd.AddCommand(newLogsCommand()) cmd.AddCommand(newCancelCommand()) cmd.AddCommand(newRegisterImageCommand()) diff --git a/experimental/air/cmd/air_test.go b/experimental/air/cmd/air_test.go index 7efac253a2b..d15a94a578d 100644 --- a/experimental/air/cmd/air_test.go +++ b/experimental/air/cmd/air_test.go @@ -14,7 +14,7 @@ func TestNewRegistersAllSubcommands(t *testing.T) { registered[c.Name()] = true } - want := []string{"run", "get", "list", "logs", "cancel", "register-image"} + want := []string{"run", "get", "list", "delete", "logs", "cancel", "register-image"} for _, name := range want { assert.True(t, registered[name], "subcommand %q is not registered", name) } diff --git a/experimental/air/cmd/aitraining.go b/experimental/air/cmd/aitraining.go deleted file mode 100644 index 735c0d71aaf..00000000000 --- a/experimental/air/cmd/aitraining.go +++ /dev/null @@ -1,110 +0,0 @@ -package aircmd - -import ( - "context" - "encoding/json" - "fmt" - "net/http" - "time" - - "github.com/databricks/databricks-sdk-go" - "github.com/databricks/databricks-sdk-go/client" -) - -// aiTrainingWorkflowsPath is the AiTrainingService index of the caller's own AIR -// runs. It returns cheap (job_run_id, submit_time) pairs, letting `air list` -// order and page without scanning the Jobs runs/list firehose. This is called -// with a raw client.Do because the SDK does not model the AiTrainingService. -const aiTrainingWorkflowsPath = "/api/2.0/ai-training/workflows" - -// workflowRef is one run from the index: its Jobs run id and submission time. -type workflowRef struct { - jobRunID int64 - submitTimeMs int64 -} - -type aiTrainingWorkflow struct { - // job_run_id is a Jobs run id; tolerate it arriving as a JSON number or string. - JobRunID json.Number `json:"job_run_id"` - // submit_time is a proto Timestamp, serialized over HTTP as either an RFC3339 - // string or a {seconds, nanos} object. - SubmitTime json.RawMessage `json:"submit_time"` -} - -type aiTrainingWorkflowsResponse struct { - TrainingWorkflows []aiTrainingWorkflow `json:"training_workflows"` - NextPageToken string `json:"next_page_token"` -} - -// listAiTrainingWorkflows pages the index and returns every workflow ref the -// caller owns. Pagination stops at the end or when a page token repeats, which -// guards against a stuck or cycling cursor without an arbitrary page cap. -func listAiTrainingWorkflows(ctx context.Context, w *databricks.WorkspaceClient, activeOnly bool) ([]workflowRef, error) { - apiClient, err := client.New(w.Config) - if err != nil { - return nil, fmt.Errorf("failed to create API client: %w", err) - } - - var refs []workflowRef - seenTokens := map[string]bool{} - // The index can return the same job_run_id on more than one page; dedupe so - // the newest-`limit` truncation counts unique runs, not repeats. - seenIDs := map[int64]bool{} - var pageToken string - for { - query := map[string]any{} - if activeOnly { - query["active_only"] = true - } - if pageToken != "" { - query["page_token"] = pageToken - } - - var resp aiTrainingWorkflowsResponse - err = apiClient.Do(ctx, http.MethodGet, aiTrainingWorkflowsPath, nil, nil, query, &resp) - if err != nil { - return nil, fmt.Errorf("failed to list training workflows: %w", err) - } - - for _, wf := range resp.TrainingWorkflows { - id, err := wf.JobRunID.Int64() - if err != nil || id == 0 || seenIDs[id] { - continue - } - seenIDs[id] = true - refs = append(refs, workflowRef{jobRunID: id, submitTimeMs: parseSubmitTimeMs(wf.SubmitTime)}) - } - - if resp.NextPageToken == "" || seenTokens[resp.NextPageToken] { - break - } - seenTokens[resp.NextPageToken] = true - pageToken = resp.NextPageToken - } - return refs, nil -} - -// parseSubmitTimeMs converts a proto Timestamp (RFC3339 string or {seconds, nanos} -// object) to epoch milliseconds, or 0 when absent or unparseable (so it sorts last). -func parseSubmitTimeMs(raw json.RawMessage) int64 { - if len(raw) == 0 { - return 0 - } - - var s string - if json.Unmarshal(raw, &s) == nil { - if t, err := time.Parse(time.RFC3339, s); err == nil { - return t.UnixMilli() - } - return 0 - } - - var obj struct { - Seconds int64 `json:"seconds"` - Nanos int64 `json:"nanos"` - } - if json.Unmarshal(raw, &obj) == nil { - return obj.Seconds*1000 + obj.Nanos/1_000_000 - } - return 0 -} diff --git a/experimental/air/cmd/aitraining_test.go b/experimental/air/cmd/aitraining_test.go deleted file mode 100644 index 4c44266d8cb..00000000000 --- a/experimental/air/cmd/aitraining_test.go +++ /dev/null @@ -1,92 +0,0 @@ -package aircmd - -import ( - "encoding/json" - "net/http" - "net/http/httptest" - "testing" - - "github.com/stretchr/testify/assert" - "github.com/stretchr/testify/require" -) - -func TestParseSubmitTimeMs(t *testing.T) { - cases := []struct { - name string - raw string - want int64 - }{ - {"rfc3339", `"2023-11-14T22:13:20Z"`, 1700000000000}, - {"rfc3339 offset", `"2023-11-14T22:13:20+00:00"`, 1700000000000}, - {"seconds and nanos", `{"seconds": 1700000000, "nanos": 500000000}`, 1700000000500}, - {"seconds only", `{"seconds": 1700000000}`, 1700000000000}, - {"empty", ``, 0}, - {"garbage string", `"not-a-time"`, 0}, - } - for _, tc := range cases { - t.Run(tc.name, func(t *testing.T) { - assert.Equal(t, tc.want, parseSubmitTimeMs(json.RawMessage(tc.raw))) - }) - } -} - -// indexServer serves paginated AiTrainingService responses, one body per call, -// tracking whether the index was hit. -func indexServer(t *testing.T, hit *bool, bodies ...string) *httptest.Server { - t.Helper() - call := 0 - srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - if r.URL.Path == aiTrainingWorkflowsPath { - *hit = true - body := bodies[min(call, len(bodies)-1)] - call++ - _, _ = w.Write([]byte(body)) - return - } - _, _ = w.Write([]byte(`{}`)) - })) - t.Cleanup(srv.Close) - return srv -} - -func TestListAiTrainingWorkflowsPaginates(t *testing.T) { - page1 := `{"training_workflows":[{"job_run_id":"1","submit_time":"2023-11-14T22:13:20Z"}],"next_page_token":"tok"}` - page2 := `{"training_workflows":[{"job_run_id":2,"submit_time":{"seconds":1700000100}}]}` - var hit bool - srv := indexServer(t, &hit, page1, page2) - - refs, err := listAiTrainingWorkflows(t.Context(), newTestWorkspaceClient(t, srv.URL), false) - require.NoError(t, err) - require.Len(t, refs, 2) - assert.Equal(t, int64(1), refs[0].jobRunID) - assert.Equal(t, int64(1700000000000), refs[0].submitTimeMs) - assert.Equal(t, int64(2), refs[1].jobRunID) -} - -func TestListAiTrainingWorkflowsStopsOnRepeatedToken(t *testing.T) { - // A cursor that always returns the same token must not loop forever. The - // repeated id is also deduped, so only one ref survives. - page := `{"training_workflows":[{"job_run_id":1,"submit_time":"2023-11-14T22:13:20Z"}],"next_page_token":"tok"}` - var hit bool - srv := indexServer(t, &hit, page) - - refs, err := listAiTrainingWorkflows(t.Context(), newTestWorkspaceClient(t, srv.URL), false) - require.NoError(t, err) - require.Len(t, refs, 1) - assert.Equal(t, int64(1), refs[0].jobRunID) -} - -func TestListAiTrainingWorkflowsDedupesIDs(t *testing.T) { - // The same job_run_id on multiple pages must be counted once, so the - // newest-limit truncation doesn't silently return fewer unique runs. - page1 := `{"training_workflows":[{"job_run_id":1,"submit_time":"2023-11-14T22:13:20Z"},{"job_run_id":2,"submit_time":"2023-11-14T22:13:21Z"}],"next_page_token":"tok"}` - page2 := `{"training_workflows":[{"job_run_id":2,"submit_time":"2023-11-14T22:13:21Z"},{"job_run_id":3,"submit_time":"2023-11-14T22:13:22Z"}]}` - var hit bool - srv := indexServer(t, &hit, page1, page2) - - refs, err := listAiTrainingWorkflows(t.Context(), newTestWorkspaceClient(t, srv.URL), false) - require.NoError(t, err) - require.Len(t, refs, 3) - got := []int64{refs[0].jobRunID, refs[1].jobRunID, refs[2].jobRunID} - assert.ElementsMatch(t, []int64{1, 2, 3}, got) -} diff --git a/experimental/air/cmd/bundle.go b/experimental/air/cmd/bundle.go new file mode 100644 index 00000000000..6adb9d90102 --- /dev/null +++ b/experimental/air/cmd/bundle.go @@ -0,0 +1,411 @@ +package aircmd + +import ( + "context" + "fmt" + "strconv" + "strings" + + "github.com/databricks/cli/cmd/root" + "github.com/databricks/cli/libs/cmdctx" + "github.com/databricks/cli/libs/cmdio" + "github.com/databricks/cli/libs/flags" + "github.com/databricks/databricks-sdk-go" + "github.com/databricks/databricks-sdk-go/service/jobs" + "github.com/spf13/cobra" + "golang.org/x/sync/errgroup" +) + +// A bundle is the persistent Jobs resource `air run` deploys: one job per +// experiment, holding an ai_runtime_task. Runs are executions of a bundle. The +// `bundle` noun manages those durable jobs (list / get / delete), distinct from +// `run`, which manages individual executions. + +// bundleScanConcurrency bounds the parallel runs/list calls when checking each +// bundle for active runs. +const bundleScanConcurrency = 16 + +// airBundle is one AIR job, projected for display. +type airBundle struct { + JobID string `json:"job_id"` + Name string `json:"name"` + User string `json:"user"` + Created *int64 `json:"created_time_ms,omitempty"` +} + +// isAirJob reports whether a job's first task is an ai_runtime_task (or a +// foreach sweep of one), or a legacy gen_ai_compute_task with a training script. +// jobs/list analogue of isAirRun. +func isAirJob(j jobs.BaseJob) bool { + if j.Settings == nil || len(j.Settings.Tasks) == 0 { + return false + } + t := j.Settings.Tasks[0] + if t.ForEachTask != nil { + t = t.ForEachTask.Task + } + return t.AiRuntimeTask != nil || + (t.GenAiComputeTask != nil && t.GenAiComputeTask.TrainingScriptPath != "") +} + +// listAirBundles enumerates AIR jobs via jobs/list, newest-first. userFilter, when +// set, keeps only that creator's bundles. truncated reports the maxListScan cap +// was hit. +func listAirBundles(ctx context.Context, w *databricks.WorkspaceClient, userFilter string) (bundles []airBundle, truncated bool, err error) { + iter := w.Jobs.List(ctx, jobs.ListJobsRequest{ExpandTasks: true, Limit: jobsPageLimit}) + inspected := 0 + for iter.HasNext(ctx) && inspected < maxListScan { + job, err := iter.Next(ctx) + if err != nil { + return nil, false, fmt.Errorf("failed to list jobs: %w", err) + } + inspected++ + if !isAirJob(job) { + continue + } + if userFilter != "" && job.CreatorUserName != userFilter { + continue + } + b := airBundle{ + JobID: strconv.FormatInt(job.JobId, 10), + Name: job.Settings.Name, + User: job.CreatorUserName, + } + if job.CreatedTime > 0 { + created := job.CreatedTime + b.Created = &created + } + bundles = append(bundles, b) + } + return bundles, inspected >= maxListScan && iter.HasNext(ctx), nil +} + +// bundleHasActiveRuns reports whether a bundle (job) has any run that is not yet +// terminal. +func bundleHasActiveRuns(ctx context.Context, w *databricks.WorkspaceClient, jobID int64) (bool, error) { + iter := w.Jobs.ListRuns(ctx, jobs.ListRunsRequest{JobId: jobID, ActiveOnly: true, Limit: 1}) + if !iter.HasNext(ctx) { + return false, nil + } + if _, err := iter.Next(ctx); err != nil { + return false, fmt.Errorf("failed to list runs for job %d: %w", jobID, err) + } + return true, nil +} + +// --- list bundles --- + +type listBundlesData struct { + Bundles []airBundle `json:"bundles"` +} + +func newListBundlesCommand() *cobra.Command { + var allUsers bool + + cmd := &cobra.Command{ + Use: "bundles", + Args: root.NoArgs, + Short: "List the persistent AIR jobs (bundles) `air run` has deployed", + } + cmd.Flags().BoolVar(&allUsers, "all-users", false, "Show bundles from all users") + cmd.PreRunE = root.MustWorkspaceClient + + cmd.RunE = func(cmd *cobra.Command, args []string) error { + ctx := cmd.Context() + w := cmdctx.WorkspaceClient(ctx) + + userFilter, err := selfUserFilter(ctx, w, allUsers) + if err != nil { + return err + } + + bundles, truncated, err := listAirBundles(ctx, w, userFilter) + if err != nil { + return renderError(ctx, cmd, "INTERNAL_ERROR", "TRANSIENT", true, err) + } + if truncated { + cmdio.LogString(ctx, fmt.Sprintf("air list bundles: stopped after inspecting %d jobs; results may be incomplete", maxListScan)) + } + + if root.OutputType(cmd) == flags.OutputJSON { + return renderEnvelope(ctx, listBundlesData{Bundles: bundles}) + } + renderBundlesText(cmd, bundles) + return nil + } + return cmd +} + +// --- get bundle --- + +type getBundleData struct { + JobID string `json:"job_id"` + Name string `json:"name"` + User string `json:"user"` + DashboardURL string `json:"dashboard_url"` + Runs []listRow `json:"runs"` +} + +func newGetBundleCommand() *cobra.Command { + cmd := &cobra.Command{ + Use: "bundle NAME", + Args: root.ExactArgs(1), + Short: "Show a bundle (persistent AIR job) and its recent runs", + } + cmd.PreRunE = root.MustWorkspaceClient + + cmd.RunE = func(cmd *cobra.Command, args []string) error { + ctx := cmd.Context() + w := cmdctx.WorkspaceClient(ctx) + + bundle, err := resolveBundle(ctx, w, args[0]) + if err != nil { + return renderError(ctx, cmd, "NOT_FOUND", "NOT_FOUND", false, err) + } + jobID, _ := strconv.ParseInt(bundle.JobID, 10, 64) + + rows, err := runsForBundle(ctx, w, jobID) + if err != nil { + return renderError(ctx, cmd, "INTERNAL_ERROR", "TRANSIENT", true, err) + } + + data := getBundleData{ + JobID: bundle.JobID, + Name: bundle.Name, + User: bundle.User, + DashboardURL: strings.TrimRight(w.Config.Host, "/") + "/jobs/" + bundle.JobID, + Runs: rows, + } + if root.OutputType(cmd) == flags.OutputJSON { + return renderEnvelope(ctx, data) + } + renderBundleDetail(cmd, data) + return nil + } + return cmd +} + +// --- delete bundle --- + +type deleteBundleData struct { + Deleted []string `json:"deleted"` + All bool `json:"all,omitempty"` + Failed []cancelFailure `json:"failed,omitempty"` +} + +func newDeleteCommand() *cobra.Command { + cmd := &cobra.Command{ + Use: "delete", + Short: "Delete a resource (bundle)", + } + cmd.AddCommand(newDeleteBundleCommand()) + return cmd +} + +func newDeleteBundleCommand() *cobra.Command { + var ( + all bool + yes bool + ) + + cmd := &cobra.Command{ + Use: "bundle [NAME...]", + Short: "Delete persistent AIR jobs (bundles) `air run` deployed", + Long: `Delete one or more bundles (persistent AIR jobs) by name, or all of your +bundles with --all. A bundle with running runs is confirmed before deletion.`, + } + cmd.Flags().BoolVarP(&all, "all", "a", false, "Delete all of your bundles") + cmd.Flags().BoolVarP(&yes, "yes", "y", false, "Skip the confirmation prompt") + + cmd.Args = func(cmd *cobra.Command, args []string) error { + switch { + case all && len(args) > 0: + return &root.InvalidArgsError{Command: cmd, Message: "cannot combine NAME arguments with --all"} + case !all && len(args) == 0: + return &root.InvalidArgsError{Command: cmd, Message: "provide at least one bundle NAME, or use --all"} + } + return nil + } + cmd.PreRunE = root.MustWorkspaceClient + + cmd.RunE = func(cmd *cobra.Command, args []string) error { + ctx := cmd.Context() + w := cmdctx.WorkspaceClient(ctx) + jsonOut := root.OutputType(cmd) == flags.OutputJSON + + targets, err := deleteTargets(ctx, w, args, all) + if err != nil { + return renderError(ctx, cmd, "INTERNAL_ERROR", "TRANSIENT", true, err) + } + if len(targets) == 0 { + if jsonOut { + return renderEnvelope(ctx, deleteBundleData{Deleted: []string{}, All: all}) + } + cmdio.LogString(ctx, "No bundles found.") + return nil + } + + // A bundle with running runs needs explicit confirmation; --all always + // confirms the whole set. Non-interactive callers pass -y. + active, err := bundlesWithActiveRuns(ctx, w, targets) + if err != nil { + return renderError(ctx, cmd, "INTERNAL_ERROR", "TRANSIENT", true, err) + } + if !yes && (all || len(active) > 0) { + confirmed, err := confirmDelete(ctx, targets, active, all) + if err != nil { + return err + } + if !confirmed { + cmdio.LogString(ctx, "Deletion aborted.") + return root.ErrAlreadyPrinted + } + } + + data := deleteBundleData{Deleted: []string{}, All: all} + for _, b := range targets { + jobID, _ := strconv.ParseInt(b.JobID, 10, 64) + if err := w.Jobs.DeleteByJobId(ctx, jobID); err != nil { + data.Failed = append(data.Failed, cancelFailure{RunID: b.Name, Error: err.Error()}) + if !jsonOut { + cmdio.LogString(ctx, fmt.Sprintf("Failed to delete bundle %s: %s", b.Name, err)) + } + continue + } + data.Deleted = append(data.Deleted, b.Name) + if !jsonOut { + cmdio.LogString(ctx, "Deleted bundle "+b.Name) + } + } + + if jsonOut { + if err := renderEnvelope(ctx, data); err != nil { + return err + } + if len(data.Failed) > 0 { + return root.ErrAlreadyPrinted + } + return nil + } + if len(data.Failed) > 0 { + cmdio.LogString(ctx, fmt.Sprintf("%d bundle(s) failed to delete.", len(data.Failed))) + return root.ErrAlreadyPrinted + } + return nil + } + return cmd +} + +// deleteTargets resolves the bundles a delete should act on: every AIR bundle +// for --all, otherwise each named bundle (erroring if any name is unknown). +func deleteTargets(ctx context.Context, w *databricks.WorkspaceClient, names []string, all bool) ([]airBundle, error) { + me, err := currentUserName(ctx, w) + if err != nil { + return nil, err + } + bundles, _, err := listAirBundles(ctx, w, me) + if err != nil { + return nil, err + } + if all { + return bundles, nil + } + byName := make(map[string][]airBundle, len(bundles)) + for _, b := range bundles { + byName[b.Name] = append(byName[b.Name], b) + } + var targets []airBundle + for _, name := range names { + matches := byName[name] + if len(matches) == 0 { + return nil, fmt.Errorf("bundle %q not found", name) + } + targets = append(targets, matches...) + } + return targets, nil +} + +// bundlesWithActiveRuns returns the names of bundles that have a non-terminal run. +func bundlesWithActiveRuns(ctx context.Context, w *databricks.WorkspaceClient, bundles []airBundle) ([]string, error) { + active := make([]bool, len(bundles)) + g, gctx := errgroup.WithContext(ctx) + g.SetLimit(bundleScanConcurrency) + for i := range bundles { + g.Go(func() error { + jobID, _ := strconv.ParseInt(bundles[i].JobID, 10, 64) + has, err := bundleHasActiveRuns(gctx, w, jobID) + if err != nil { + return err + } + active[i] = has + return nil + }) + } + if err := g.Wait(); err != nil { + return nil, err + } + var names []string + for i := range bundles { + if active[i] { + names = append(names, bundles[i].Name) + } + } + return names, nil +} + +// confirmDelete shows what will be deleted, flags any bundles with running runs, +// and prompts for confirmation. +func confirmDelete(ctx context.Context, targets []airBundle, active []string, all bool) (bool, error) { + var sb strings.Builder + fmt.Fprintf(&sb, "About to delete %d bundle(s):\n", len(targets)) + activeSet := make(map[string]bool, len(active)) + for _, n := range active { + activeSet[n] = true + } + for _, b := range targets { + if activeSet[b.Name] { + fmt.Fprintf(&sb, " %s (has running runs)\n", b.Name) + } else { + fmt.Fprintf(&sb, " %s\n", b.Name) + } + } + cmdio.LogString(ctx, strings.TrimRight(sb.String(), "\n")) + + prompt := "\nDelete these bundle(s)?" + if len(active) > 0 { + prompt = fmt.Sprintf("\n%d bundle(s) have running runs. Delete anyway?", len(active)) + } + return cmdio.AskYesOrNo(ctx, prompt) +} + +// resolveBundle finds a single AIR bundle by name for the current user. +func resolveBundle(ctx context.Context, w *databricks.WorkspaceClient, name string) (airBundle, error) { + me, err := currentUserName(ctx, w) + if err != nil { + return airBundle{}, err + } + bundles, _, err := listAirBundles(ctx, w, me) + if err != nil { + return airBundle{}, err + } + for _, b := range bundles { + if b.Name == name { + return b, nil + } + } + return airBundle{}, fmt.Errorf("bundle %q not found", name) +} + +// runsForBundle lists a bundle's runs (a single page) as display rows. +func runsForBundle(ctx context.Context, w *databricks.WorkspaceClient, jobID int64) ([]listRow, error) { + iter := w.Jobs.ListRuns(ctx, jobs.ListRunsRequest{JobId: jobID, ExpandTasks: true, Limit: jobsPageLimit}) + var rows []listRow + for iter.HasNext(ctx) { + base, err := iter.Next(ctx) + if err != nil { + return nil, fmt.Errorf("failed to list runs for job %d: %w", jobID, err) + } + rows = append(rows, buildListRow(baseRunToRun(base))) + } + return rows, nil +} diff --git a/experimental/air/cmd/bundle_format.go b/experimental/air/cmd/bundle_format.go new file mode 100644 index 00000000000..f5c2a75260b --- /dev/null +++ b/experimental/air/cmd/bundle_format.go @@ -0,0 +1,59 @@ +package aircmd + +import ( + "fmt" + "strings" + "text/tabwriter" + "time" + + "github.com/spf13/cobra" +) + +// renderBundlesText prints the bundle list as a table (job id, name, user, +// created), or a friendly line when empty. +func renderBundlesText(cmd *cobra.Command, bundles []airBundle) { + out := cmd.OutOrStdout() + if len(bundles) == 0 { + fmt.Fprintln(out, "No bundles found.") + return + } + tw := tabwriter.NewWriter(out, 0, 0, 2, ' ', 0) + fmt.Fprintln(tw, "Job ID\tName\tUser\tCreated") + for _, b := range bundles { + fmt.Fprintf(tw, "%s\t%s\t%s\t%s\n", b.JobID, b.Name, orNA(b.User), bundleCreated(b)) + } + tw.Flush() +} + +// renderBundleDetail prints a single bundle's metadata plus its recent runs. +func renderBundleDetail(cmd *cobra.Command, data getBundleData) { + out := cmd.OutOrStdout() + fmt.Fprintf(out, "Bundle: %s\n", data.Name) + fmt.Fprintf(out, "Job ID: %s\n", data.JobID) + fmt.Fprintf(out, "User: %s\n", orNA(data.User)) + fmt.Fprintf(out, "Job URL: %s\n", data.DashboardURL) + + if len(data.Runs) == 0 { + fmt.Fprintln(out, "\nNo runs yet.") + return + } + fmt.Fprintln(out, "\nRuns:") + tw := tabwriter.NewWriter(out, 0, 0, 2, ' ', 0) + fmt.Fprintln(tw, "Run ID\tStatus\tStarted\tDuration") + for _, r := range data.Runs { + started := na + if r.StartedAt != nil { + started = *r.StartedAt + } + fmt.Fprintf(tw, "%s\t%s\t%s\t%s\n", r.RunID, r.Status, started, orNA(r.Duration)) + } + tw.Flush() +} + +// bundleCreated formats a bundle's creation time, or "N/A" when unset. +func bundleCreated(b airBundle) string { + if b.Created == nil { + return na + } + return strings.TrimSpace(isoFormat(time.UnixMilli(*b.Created))) +} diff --git a/experimental/air/cmd/bundle_test.go b/experimental/air/cmd/bundle_test.go new file mode 100644 index 00000000000..a4fd36f510a --- /dev/null +++ b/experimental/air/cmd/bundle_test.go @@ -0,0 +1,160 @@ +package aircmd + +import ( + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + + "github.com/databricks/databricks-sdk-go/service/jobs" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestIsAirJob(t *testing.T) { + airJob := jobs.BaseJob{Settings: &jobs.JobSettings{ + Tasks: []jobs.Task{{AiRuntimeTask: &jobs.AiRuntimeTask{Experiment: "exp"}}}, + }} + assert.True(t, isAirJob(airJob)) + + sweep := jobs.BaseJob{Settings: &jobs.JobSettings{ + Tasks: []jobs.Task{{ForEachTask: &jobs.ForEachTask{ + Task: jobs.Task{AiRuntimeTask: &jobs.AiRuntimeTask{Experiment: "exp"}}, + }}}, + }} + assert.True(t, isAirJob(sweep), "foreach-wrapped ai_runtime_task is an AIR job") + + legacy := jobs.BaseJob{Settings: &jobs.JobSettings{ + Tasks: []jobs.Task{{GenAiComputeTask: &jobs.GenAiComputeTask{TrainingScriptPath: "/x"}}}, + }} + assert.True(t, isAirJob(legacy)) + + notebook := jobs.BaseJob{Settings: &jobs.JobSettings{ + Tasks: []jobs.Task{{NotebookTask: &jobs.NotebookTask{NotebookPath: "/x"}}}, + }} + assert.False(t, isAirJob(notebook)) + + assert.False(t, isAirJob(jobs.BaseJob{}), "a job with no settings is not an AIR job") +} + +// jobsListServer serves one jobs/list page and stubs everything else (config +// probe, scim Me). The Me response fixes the current user for scoping tests. +func jobsListServer(t *testing.T, user string, jobList ...jobs.BaseJob) *httptest.Server { + t.Helper() + body, err := json.Marshal(jobs.ListJobsResponse{Jobs: jobList}) + require.NoError(t, err) + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/api/2.2/jobs/list": + _, _ = w.Write(body) + case "/api/2.0/preview/scim/v2/Me": + _, _ = w.Write([]byte(`{"userName":"` + user + `"}`)) + default: + _, _ = w.Write([]byte(`{}`)) + } + })) + t.Cleanup(srv.Close) + return srv +} + +func airBaseJob(jobID int64, name, user string) jobs.BaseJob { + return jobs.BaseJob{ + JobId: jobID, + CreatorUserName: user, + Settings: &jobs.JobSettings{ + Name: name, + Tasks: []jobs.Task{{AiRuntimeTask: &jobs.AiRuntimeTask{Experiment: name}}}, + }, + } +} + +func TestListAirBundlesFiltersNonAirAndUser(t *testing.T) { + jobList := []jobs.BaseJob{ + airBaseJob(1, "exp-a", "me@example.com"), + {JobId: 2, CreatorUserName: "me@example.com", Settings: &jobs.JobSettings{ // not AIR + Name: "nb", Tasks: []jobs.Task{{NotebookTask: &jobs.NotebookTask{NotebookPath: "/x"}}}, + }}, + airBaseJob(3, "exp-b", "other@example.com"), // wrong user + airBaseJob(5, "exp-c", "me@example.com"), + } + srv := jobsListServer(t, "me@example.com", jobList...) + w := newTestWorkspaceClient(t, srv.URL) + + bundles, truncated, err := listAirBundles(t.Context(), w, "me@example.com") + require.NoError(t, err) + assert.False(t, truncated) + require.Len(t, bundles, 2) + assert.Equal(t, "exp-a", bundles[0].Name) + assert.Equal(t, "exp-c", bundles[1].Name) + + // No user filter keeps every AIR job (still drops the non-AIR notebook job). + all, _, err := listAirBundles(t.Context(), w, "") + require.NoError(t, err) + assert.Len(t, all, 3) +} + +func TestResolveBundleNotFound(t *testing.T) { + srv := jobsListServer(t, "me@example.com", airBaseJob(1, "exp-a", "me@example.com")) + w := newTestWorkspaceClient(t, srv.URL) + + _, err := resolveBundle(t.Context(), w, "does-not-exist") + require.Error(t, err) + assert.Contains(t, err.Error(), `bundle "does-not-exist" not found`) + + got, err := resolveBundle(t.Context(), w, "exp-a") + require.NoError(t, err) + assert.Equal(t, "1", got.JobID) +} + +func TestDeleteBundleArgs(t *testing.T) { + tests := []struct { + name string + all bool + args []string + wantErr string + }{ + {name: "one name", args: []string{"exp-a"}}, + {name: "many names", args: []string{"exp-a", "exp-b"}}, + {name: "all", all: true}, + {name: "no input", wantErr: "provide at least one bundle NAME, or use --all"}, + {name: "names with all", all: true, args: []string{"exp-a"}, wantErr: "cannot combine NAME arguments with --all"}, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + cmd := newDeleteBundleCommand() + if tc.all { + require.NoError(t, cmd.Flags().Set("all", "true")) + } + err := cmd.Args(cmd, tc.args) + if tc.wantErr == "" { + assert.NoError(t, err) + return + } + require.Error(t, err) + assert.Contains(t, err.Error(), tc.wantErr) + }) + } +} + +func TestDeleteTargetsResolvesNames(t *testing.T) { + srv := jobsListServer(t, "me@example.com", + airBaseJob(1, "exp-a", "me@example.com"), + airBaseJob(2, "exp-b", "me@example.com"), + ) + w := newTestWorkspaceClient(t, srv.URL) + + // --all returns every bundle. + all, err := deleteTargets(t.Context(), w, nil, true) + require.NoError(t, err) + assert.Len(t, all, 2) + + // Named targets resolve to their bundles; an unknown name errors. + named, err := deleteTargets(t.Context(), w, []string{"exp-b"}, false) + require.NoError(t, err) + require.Len(t, named, 1) + assert.Equal(t, "2", named[0].JobID) + + _, err = deleteTargets(t.Context(), w, []string{"nope"}, false) + require.Error(t, err) + assert.Contains(t, err.Error(), `bundle "nope" not found`) +} 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/get.go b/experimental/air/cmd/get.go index 2f8bd8dd09e..9fb6118b742 100644 --- a/experimental/air/cmd/get.go +++ b/experimental/air/cmd/get.go @@ -89,11 +89,23 @@ func authError(ctx context.Context, cmd *cobra.Command, err error) error { fmt.Errorf("failed to verify authentication: %w", err)) } -// newGetCommand returns the `air get JOB_RUN_ID` command, which shows status, -// configuration, and timing details for a specific run. +// newGetCommand returns the `air get` command group. Resources are addressed as +// nouns: `air get run ` and `air get bundle `. func newGetCommand() *cobra.Command { cmd := &cobra.Command{ - Use: "get JOB_RUN_ID", + Use: "get", + Short: "Show details for a specific resource (run or bundle)", + } + cmd.AddCommand(newGetRunCommand()) + cmd.AddCommand(newGetBundleCommand()) + return cmd +} + +// newGetRunCommand returns the `air get run JOB_RUN_ID` command, which shows +// status, configuration, and timing details for a specific run. +func newGetRunCommand() *cobra.Command { + cmd := &cobra.Command{ + Use: "run JOB_RUN_ID", Args: root.ExactArgs(1), Short: "Show status, configuration, and timing details for a specific run", Annotations: map[string]string{ diff --git a/experimental/air/cmd/get_test.go b/experimental/air/cmd/get_test.go index aa23e58159e..c45fee8f98a 100644 --- a/experimental/air/cmd/get_test.go +++ b/experimental/air/cmd/get_test.go @@ -33,23 +33,30 @@ func renderGet(t *testing.T, data getData) string { return buf.String() } -// TestGetCommandShape locks in that `get` takes the run id directly as -// `air get JOB_RUN_ID` and has no `run` subcommand (it was collapsed back into -// `get`). The acceptance test exercises the happy path end to end. +// TestGetCommandShape locks in the verb-noun structure: `get` is a group with +// `run` and `bundle` subcommands, and `get run` takes exactly one run id. func TestGetCommandShape(t *testing.T) { - cmd := newGetCommand() - assert.Equal(t, "get JOB_RUN_ID", cmd.Use) - assert.Empty(t, cmd.Commands(), "get must not register subcommands") + get := newGetCommand() + assert.Equal(t, "get", get.Use) + names := make(map[string]bool) + for _, c := range get.Commands() { + names[c.Name()] = true + } + assert.True(t, names["run"], "get must register the run subcommand") + assert.True(t, names["bundle"], "get must register the bundle subcommand") + + run := newGetRunCommand() + assert.Equal(t, "run JOB_RUN_ID", run.Use) // ExactArgs(1): exactly one run id is required. - assert.NoError(t, cmd.Args(cmd, []string{"123"})) - assert.Error(t, cmd.Args(cmd, []string{})) - assert.Error(t, cmd.Args(cmd, []string{"1", "2"})) + assert.NoError(t, run.Args(run, []string{"123"})) + assert.Error(t, run.Args(run, []string{})) + assert.Error(t, run.Args(run, []string{"1", "2"})) } func TestGetRunInvalidID(t *testing.T) { m := mocks.NewMockWorkspaceClient(t) ctx := cmdctx.SetWorkspaceClient(cmdio.MockDiscard(t.Context()), m.WorkspaceClient) - cmd := withOutput(newGetCommand(), flags.OutputText) + cmd := withOutput(newGetRunCommand(), flags.OutputText) cmd.SetContext(ctx) err := cmd.RunE(cmd, []string{"abc"}) @@ -78,7 +85,7 @@ func notFoundGetServer(t *testing.T) *httptest.Server { func TestGetRunNotFound(t *testing.T) { srv := notFoundGetServer(t) ctx := cmdctx.SetWorkspaceClient(cmdio.MockDiscard(t.Context()), newTestWorkspaceClient(t, srv.URL)) - cmd := withOutput(newGetCommand(), flags.OutputText) + cmd := withOutput(newGetRunCommand(), flags.OutputText) cmd.SetContext(ctx) err := cmd.RunE(cmd, []string{"5"}) @@ -92,7 +99,7 @@ func TestGetRunAuthFailed(t *testing.T) { // fetched, so GetRun is never reached and nothing is rendered. m.GetMockCurrentUserAPI().EXPECT().Me(mock.Anything, mock.Anything).Return(nil, apierr.ErrPermissionDenied) ctx := cmdctx.SetWorkspaceClient(cmdio.MockDiscard(t.Context()), m.WorkspaceClient) - cmd := withOutput(newGetCommand(), flags.OutputText) + cmd := withOutput(newGetRunCommand(), flags.OutputText) cmd.SetContext(ctx) err := cmd.RunE(cmd, []string{"5"}) @@ -106,7 +113,7 @@ func TestGetRunAuthTransient(t *testing.T) { // error; it surfaces as a retryable internal error instead. m.GetMockCurrentUserAPI().EXPECT().Me(mock.Anything, mock.Anything).Return(nil, errors.New("connection reset")) ctx := cmdctx.SetWorkspaceClient(cmdio.MockDiscard(t.Context()), m.WorkspaceClient) - cmd := withOutput(newGetCommand(), flags.OutputText) + cmd := withOutput(newGetRunCommand(), flags.OutputText) cmd.SetContext(ctx) err := cmd.RunE(cmd, []string{"5"}) @@ -117,7 +124,7 @@ func TestGetRunAuthTransient(t *testing.T) { func TestAuthError(t *testing.T) { ctx := cmdio.MockDiscard(t.Context()) - cmd := withOutput(newGetCommand(), flags.OutputText) + cmd := withOutput(newGetRunCommand(), flags.OutputText) // No configurable credentials maps to the missing-profile hint. noProfile := authError(ctx, cmd, config.ErrCannotConfigureDefault) @@ -144,7 +151,7 @@ func TestGetRunNotFoundJSON(t *testing.T) { srv := notFoundGetServer(t) ctx := cmdctx.SetWorkspaceClient(t.Context(), newTestWorkspaceClient(t, srv.URL)) ctx = cmdio.InContext(ctx, cmdio.NewIO(ctx, flags.OutputJSON, nil, &buf, &buf, "", "")) - cmd := withOutput(newGetCommand(), flags.OutputJSON) + cmd := withOutput(newGetRunCommand(), flags.OutputJSON) cmd.SetContext(ctx) // In JSON mode the not-found error is a structured envelope, not a bare error. diff --git a/experimental/air/cmd/joblist.go b/experimental/air/cmd/joblist.go index dd786d8b496..670a18fae61 100644 --- a/experimental/air/cmd/joblist.go +++ b/experimental/air/cmd/joblist.go @@ -2,19 +2,30 @@ package aircmd import ( "context" - "errors" "fmt" - "net/http" "github.com/databricks/databricks-sdk-go" - "github.com/databricks/databricks-sdk-go/apierr" + "github.com/databricks/databricks-sdk-go/service/iam" "github.com/databricks/databricks-sdk-go/service/jobs" - "golang.org/x/sync/errgroup" ) -// hydrateConcurrency bounds the parallel runs/get calls when hydrating a batch -// of run ids from the AiTrainingService index. -const hydrateConcurrency = 16 +// currentUserName resolves the caller's username via the Me probe. +func currentUserName(ctx context.Context, w *databricks.WorkspaceClient) (string, error) { + me, err := w.CurrentUser.Me(ctx, iam.MeRequest{}) + if err != nil { + return "", fmt.Errorf("failed to resolve current user: %w", err) + } + return me.UserName, nil +} + +// selfUserFilter returns the creator to scope a listing to: empty (all users) +// when allUsers is set, otherwise the current user. +func selfUserFilter(ctx context.Context, w *databricks.WorkspaceClient, allUsers bool) (string, error) { + if allUsers { + return "", nil + } + return currentUserName(ctx, w) +} // firstTask returns the run's first task, unwrapping a foreach sweep to the // iterated task, or nil when the run has no tasks. @@ -53,6 +64,24 @@ func isSweep(r *jobs.Run) bool { return len(r.Tasks) > 0 && r.Tasks[0].ForEachTask != nil } +// isActiveRun reports whether a run is not yet in a terminal lifecycle state, +// matching the server's active_only semantics. Used to filter active runs +// client-side on the indexed path, where the filter query and active_only cannot +// be combined server-side. +func isActiveRun(r *jobs.Run) bool { + if r.State == nil { + return false + } + switch r.State.LifeCycleState { + case jobs.RunLifeCycleStateTerminated, + jobs.RunLifeCycleStateInternalError, + jobs.RunLifeCycleStateSkipped: + return false + default: + return true + } +} + // taskRunID returns the run id of the AIR task, used to fetch its MLflow output. func taskRunID(r *jobs.Run) int64 { if len(r.Tasks) == 0 { @@ -106,20 +135,6 @@ func jobTiming(r *jobs.Run) (startMillis, endMillis int64) { return startMillis, endMillis } -// isTerminal reports whether a run has finished and its details are immutable, -// so its row is safe to cache. -func isTerminal(r *jobs.Run) bool { - if r.State == nil { - return false - } - switch r.State.LifeCycleState { - case jobs.RunLifeCycleStateTerminated, jobs.RunLifeCycleStateInternalError, jobs.RunLifeCycleStateSkipped: - return true - default: - return false - } -} - // baseRunToRun converts a runs/list BaseRun into a Run so the display helpers // operate on a single type. The two share every field the list path reads. func baseRunToRun(b jobs.BaseRun) *jobs.Run { @@ -143,38 +158,3 @@ func fetchJobRun(ctx context.Context, w *databricks.WorkspaceClient, runID int64 } return run, nil } - -// hydrateJobRuns fetches the given run ids concurrently via runs/get, preserving -// input order. runs/get enforces per-run view ACLs, so an id the caller can't -// view (403) or that has been purged (404) is dropped; any other error is -// systemic and fails the whole batch. -func hydrateJobRuns(ctx context.Context, w *databricks.WorkspaceClient, ids []int64) ([]*jobs.Run, error) { - runs := make([]*jobs.Run, len(ids)) - g, gctx := errgroup.WithContext(ctx) - g.SetLimit(hydrateConcurrency) - for i, id := range ids { - g.Go(func() error { - run, err := fetchJobRun(gctx, w, id) - if err != nil { - if apiErr, ok := errors.AsType[*apierr.APIError](err); ok && - (apiErr.StatusCode == http.StatusForbidden || apiErr.StatusCode == http.StatusNotFound) { - return nil // not viewable or purged: drop this id - } - return fmt.Errorf("failed to get run %d: %w", id, err) - } - runs[i] = run - return nil - }) - } - if err := g.Wait(); err != nil { - return nil, err - } - - hydrated := make([]*jobs.Run, 0, len(runs)) - for _, run := range runs { - if run != nil { - hydrated = append(hydrated, run) - } - } - return hydrated, nil -} diff --git a/experimental/air/cmd/list.go b/experimental/air/cmd/list.go index 7d5c703a39c..9909ef6250b 100644 --- a/experimental/air/cmd/list.go +++ b/experimental/air/cmd/list.go @@ -16,16 +16,14 @@ import ( "golang.org/x/sync/errgroup" ) -// maxListScan bounds how many runs `air list` inspects while looking for AIR runs -// that match the filters. runs/list returns runs of every kind, so this caps the -// work on a workspace with a large run history. +// maxListScan caps how many runs `air list` inspects while looking for AIR runs. +// runs/list returns every run type, so this bounds the scan on a workspace with a +// large run history where AIR runs are sparse. const maxListScan = 2000 -// jobsPageLimit is the per-request page size for runs/list; enrichConcurrency -// bounds the parallel MLflow lookups. const ( - jobsPageLimit = 25 - enrichConcurrency = 8 + jobsPageLimit = 25 // runs/list page size + enrichConcurrency = 8 // parallel MLflow lookups ) // listData is the payload printed by `air list`. @@ -70,7 +68,19 @@ type listQuery struct { limit int } +// newListCommand returns the `air list` command group. Resources are addressed +// as nouns: `air list runs` and `air list bundles`. func newListCommand() *cobra.Command { + cmd := &cobra.Command{ + Use: "list", + Short: "List a collection of resources (runs or bundles)", + } + cmd.AddCommand(newListRunsCommand()) + cmd.AddCommand(newListBundlesCommand()) + return cmd +} + +func newListRunsCommand() *cobra.Command { var ( limit int allStatus bool @@ -79,7 +89,7 @@ func newListCommand() *cobra.Command { ) cmd := &cobra.Command{ - Use: "list", + Use: "runs", Args: root.NoArgs, Short: "List your active runs for the current profile (use --all-status for finished runs)", } @@ -145,9 +155,7 @@ func newListCommand() *cobra.Command { return cmd } -// listStrategy is a source of matching runs, pulled in batches. Two implement it: -// jobsScanStrategy pages runs/list; indexStrategy hydrates the AiTrainingService -// index. The fetcher wraps whichever is chosen. +// listStrategy is a source of matching runs, pulled in batches by the fetcher. type listStrategy interface { // next returns up to want more matching runs (already row-built + task id). next(want int) ([]listedRun, error) @@ -178,20 +186,13 @@ func newRunFetcher(ctx context.Context, w *databricks.WorkspaceClient, q listQue } } -// newListStrategy picks the fetch source. The AiTrainingService index serves only -// the caller's own runs, so it's used for an all-status self-scoped list; if the -// index load fails (e.g. endpoint unavailable in this workspace), we fall back to -// the Jobs scan so the command still returns. Everything else — the default -// active list, --all-users, and --all-status for another user — uses the scan. +// newListStrategy prefers the server-side AI-Runtime filter (indexed, fast, +// complete) and self-falls-back to the client-side scan when that filter is +// unsupported (flag off) or the raw client can't be built. func newListStrategy(ctx context.Context, w *databricks.WorkspaceClient, q listQuery) listStrategy { - useIndex := !q.activeOnly && !q.allUsers && (q.userFilter == "" || q.userFilter == q.currentUser) - if !useIndex { - return newJobsScanStrategy(ctx, w, q) - } - idx := newIndexStrategy(ctx, w, q, q.limit) - if err := idx.load(); err != nil { - log.Debugf(ctx, "air list: AiTrainingService index unavailable, falling back to Jobs scan: %v", err) - return newJobsScanStrategy(ctx, w, q) + idx, err := newIndexRunStrategy(ctx, w, q) + if err != nil { + return newRunScanStrategy(ctx, w, q) } return idx } @@ -219,10 +220,13 @@ func (f *runFetcher) next(want int) ([]listRow, error) { return rows, nil } -// jobsScanStrategy pages Jobs runs/list, keeping the AIR runs that match the user -// and filters. It buffers a page's leftover runs so successive next() calls -// resume where the last stopped. -type jobsScanStrategy struct { +// runScanStrategy pages Jobs runs/list, keeping the AIR runs that match the user +// and filters. AIR runs are returned newest-first and cluster near the top, so a +// small `limit` is satisfied in the first page or two; the scan stops early. A +// SUBMIT_RUN server filter would hide DABs runs (they are JOB_RUN), so RunType is +// left unset and isAirRun selects by task shape. maxListScan bounds the worst +// case where AIR runs are sparse among other run types. +type runScanStrategy struct { ctx context.Context w *databricks.WorkspaceClient iter listing.Iterator[jobs.BaseRun] @@ -232,23 +236,21 @@ type jobsScanStrategy struct { scanned int } -func newJobsScanStrategy(ctx context.Context, w *databricks.WorkspaceClient, q listQuery) *jobsScanStrategy { - req := jobs.ListRunsRequest{ - RunType: jobs.RunTypeSubmitRun, - ExpandTasks: true, - Limit: jobsPageLimit, - ActiveOnly: q.activeOnly, - } - return &jobsScanStrategy{ - ctx: ctx, - w: w, - iter: w.Jobs.ListRuns(ctx, req), +func newRunScanStrategy(ctx context.Context, w *databricks.WorkspaceClient, q listQuery) *runScanStrategy { + return &runScanStrategy{ + ctx: ctx, + w: w, + iter: w.Jobs.ListRuns(ctx, jobs.ListRunsRequest{ + ExpandTasks: true, + Limit: jobsPageLimit, + ActiveOnly: q.activeOnly, + }), userFilter: q.userFilter, filters: q.filters, } } -func (s *jobsScanStrategy) next(want int) ([]listedRun, error) { +func (s *runScanStrategy) next(want int) ([]listedRun, error) { var entries []listedRun for len(entries) < want && s.scanned < maxListScan && s.iter.HasNext(s.ctx) { base, err := s.iter.Next(s.ctx) @@ -272,16 +274,16 @@ func (s *jobsScanStrategy) next(want int) ([]listedRun, error) { return entries, nil } -func (s *jobsScanStrategy) done() bool { +func (s *runScanStrategy) done() bool { return s.scanned >= maxListScan || !s.iter.HasNext(s.ctx) } -func (s *jobsScanStrategy) truncated() bool { +func (s *runScanStrategy) truncated() bool { return s.scanned >= maxListScan } -// warnIfTruncated logs when a scan hit its safety cap, so one-shot output signals -// its results may be incomplete. +// warnIfTruncated logs when the scan hit its cap, so one-shot output signals its +// results may be incomplete. func warnIfTruncated(ctx context.Context, f *runFetcher) { if f.strategy.truncated() { log.Warnf(ctx, "air list: stopped after scanning %d runs; results may be incomplete", maxListScan) diff --git a/experimental/air/cmd/list_cache.go b/experimental/air/cmd/list_cache.go deleted file mode 100644 index 6f0b95a97d3..00000000000 --- a/experimental/air/cmd/list_cache.go +++ /dev/null @@ -1,81 +0,0 @@ -package aircmd - -import ( - "context" - "time" - - "github.com/databricks/cli/libs/cache" -) - -// The AiTrainingService index path caches hydrated terminal runs on disk: -// terminal runs are immutable, so once we've paid for runs/get + get-output + -// MLflow we persist the finished row and skip those round-trips next time. The -// TTL matches AICM's ~60-day retention, after which the run drops out of the -// index anyway. -const ( - listCacheComponent = "air-list-runs" - listCacheTTL = 60 * 24 * time.Hour -) - -// listCacheKey fingerprints a cached run. Host isolates workspaces (a Jobs run -// id is unique only within one), matching how libs/cache namespaces entries. -type listCacheKey struct { - Host string `json:"host"` - RunID int64 `json:"run_id"` -} - -// cachedRun is the persisted value: every listRow field (including the -// table-only columns, which listRow tags json:"-" and so wouldn't survive a -// direct marshal), the filter inputs, and the submit time. -type cachedRun struct { - RunID string `json:"run_id"` - RunName string `json:"run_name"` - User string `json:"user"` - Status string `json:"status"` - StartedAt *string `json:"started_at"` - IsSweep bool `json:"is_sweep"` - Experiment string `json:"experiment"` - Duration string `json:"duration"` - MLflowURL string `json:"mlflow_url"` - Accelerators string `json:"accelerators"` - Fields filterFields `json:"filter_fields"` - SubmitTimeMs int64 `json:"submit_time_ms"` -} - -func (c cachedRun) toRow() listRow { - return listRow{ - RunID: c.RunID, RunName: c.RunName, User: c.User, Status: c.Status, - StartedAt: c.StartedAt, IsSweep: c.IsSweep, Experiment: c.Experiment, - Duration: c.Duration, MLflowURL: c.MLflowURL, Accelerators: c.Accelerators, - } -} - -func cachedRunFromRow(r listRow, fields filterFields, submitTimeMs int64) cachedRun { - return cachedRun{ - RunID: r.RunID, RunName: r.RunName, User: r.User, Status: r.Status, - StartedAt: r.StartedAt, IsSweep: r.IsSweep, Experiment: r.Experiment, - Duration: r.Duration, MLflowURL: r.MLflowURL, Accelerators: r.Accelerators, - Fields: fields, SubmitTimeMs: submitTimeMs, - } -} - -// newListCache builds the cache for the index path. It fails open, so a nil -// return (or any cache error) just means every run is hydrated from the API. -func newListCache(ctx context.Context) *cache.Cache { - return cache.NewCache(ctx, listCacheComponent, listCacheTTL, nil) -} - -// cachedRow returns the cached row and its filter fields for a run, or -// (zero, zero, false) on miss. -func cachedRow(ctx context.Context, c *cache.Cache, host string, runID int64) (listRow, filterFields, bool) { - entry, ok := cache.Get[cachedRun](ctx, c, listCacheKey{Host: host, RunID: runID}) - if !ok { - return listRow{}, filterFields{}, false - } - return entry.toRow(), entry.Fields, true -} - -// putRow caches a terminal run's finished row and filter fields under its submit time. -func putRow(ctx context.Context, c *cache.Cache, host string, runID, submitTimeMs int64, row listRow, fields filterFields) { - cache.Put(ctx, c, listCacheKey{Host: host, RunID: runID}, cachedRunFromRow(row, fields, submitTimeMs)) -} diff --git a/experimental/air/cmd/list_cache_test.go b/experimental/air/cmd/list_cache_test.go deleted file mode 100644 index 2ed076a98e2..00000000000 --- a/experimental/air/cmd/list_cache_test.go +++ /dev/null @@ -1,105 +0,0 @@ -package aircmd - -import ( - "testing" - - "github.com/databricks/databricks-sdk-go/service/jobs" - "github.com/stretchr/testify/assert" - "github.com/stretchr/testify/require" -) - -func TestListCacheRoundTrip(t *testing.T) { - t.Setenv("DATABRICKS_CACHE_DIR", t.TempDir()) - ctx := t.Context() - c := newListCache(ctx) - - _, _, ok := cachedRow(ctx, c, "https://host.test", 42) - require.False(t, ok, "miss before write") - - row := listRow{RunID: "42", Experiment: "exp", Status: "SUCCESS"} - fields := filterFields{Experiment: "exp", GPUType: "GPU_1xA10", GPUCount: 1} - putRow(ctx, c, "https://host.test", 42, 1700000000000, row, fields) - - got, gotFields, ok := cachedRow(ctx, c, "https://host.test", 42) - require.True(t, ok, "hit after write") - assert.Equal(t, row, got) - assert.Equal(t, fields, gotFields) - - // Different host is a different key. - _, _, ok = cachedRow(ctx, c, "https://other.test", 42) - assert.False(t, ok) -} - -func TestIndexStrategyServesCachedRowWithoutFetch(t *testing.T) { - t.Setenv("DATABRICKS_CACHE_DIR", t.TempDir()) - - refs := []workflowRef{{jobRunID: 7, submitTimeMs: 1000_000}} - srv, hits := indexAndGetServer(t, refs, map[int64]jobs.Run{7: indexRun(7, 1000_000)}, nil, nil) - host := srv.URL - - // Pre-seed the cache for run 7 so hydration should skip runs/get entirely. - ctx := t.Context() - putRow(ctx, newListCache(ctx), host, 7, 1000_000, listRow{RunID: "7", Status: "SUCCESS"}, filterFields{}) - - f := newRunFetcher(ctx, newTestWorkspaceClient(t, host), listQuery{ - userFilter: "me@example.com", currentUser: "me@example.com", limit: 10, - }) - rows, err := f.next(10) - require.NoError(t, err) - require.Len(t, rows, 1) - assert.Equal(t, "7", rows[0].RunID) - assert.Equal(t, 0, hits.get, "cached run must not hit runs/get") -} - -func TestIndexStrategyFiltersCachedRow(t *testing.T) { - t.Setenv("DATABRICKS_CACHE_DIR", t.TempDir()) - - refs := []workflowRef{{jobRunID: 7, submitTimeMs: 1000_000}} - srv, hits := indexAndGetServer(t, refs, map[int64]jobs.Run{7: indexRun(7, 1000_000)}, nil, nil) - host := srv.URL - - // Cache hit under experiment "bar" must be filtered out by an experiment=foo query. - ctx := t.Context() - putRow(ctx, newListCache(ctx), host, 7, 1000_000, - listRow{RunID: "7", Status: "SUCCESS", Experiment: "bar"}, - filterFields{Experiment: "bar"}) - - f := newRunFetcher(ctx, newTestWorkspaceClient(t, host), listQuery{ - userFilter: "me@example.com", currentUser: "me@example.com", limit: 10, - filters: listFilters{Experiment: "foo"}, - }) - rows, err := f.next(10) - require.NoError(t, err) - assert.Empty(t, rows, "cached row not matching the filter must be dropped") - assert.Equal(t, 0, hits.get, "non-matching cached row must not hit runs/get") -} - -func TestIndexStrategyServesMatchingCachedRow(t *testing.T) { - t.Setenv("DATABRICKS_CACHE_DIR", t.TempDir()) - - refs := []workflowRef{{jobRunID: 7, submitTimeMs: 1000_000}} - srv, hits := indexAndGetServer(t, refs, map[int64]jobs.Run{7: indexRun(7, 1000_000)}, nil, nil) - host := srv.URL - - ctx := t.Context() - putRow(ctx, newListCache(ctx), host, 7, 1000_000, - listRow{RunID: "7", Status: "SUCCESS", Experiment: "foo"}, - filterFields{Experiment: "foo"}) - - f := newRunFetcher(ctx, newTestWorkspaceClient(t, host), listQuery{ - userFilter: "me@example.com", currentUser: "me@example.com", limit: 10, - filters: listFilters{Experiment: "foo"}, - }) - rows, err := f.next(10) - require.NoError(t, err) - require.Len(t, rows, 1) - assert.Equal(t, "7", rows[0].RunID) - assert.Equal(t, 0, hits.get, "matching cached row must not hit runs/get") -} - -func TestIsTerminal(t *testing.T) { - assert.True(t, isTerminal(&jobs.Run{State: &jobs.RunState{LifeCycleState: jobs.RunLifeCycleStateTerminated}})) - assert.True(t, isTerminal(&jobs.Run{State: &jobs.RunState{LifeCycleState: jobs.RunLifeCycleStateInternalError}})) - assert.False(t, isTerminal(&jobs.Run{State: &jobs.RunState{LifeCycleState: jobs.RunLifeCycleStateRunning}})) - assert.False(t, isTerminal(&jobs.Run{State: &jobs.RunState{LifeCycleState: jobs.RunLifeCycleStatePending}})) -} diff --git a/experimental/air/cmd/list_filter.go b/experimental/air/cmd/list_filter.go index bd2bf8dff6e..d34d5180476 100644 --- a/experimental/air/cmd/list_filter.go +++ b/experimental/air/cmd/list_filter.go @@ -12,14 +12,6 @@ import ( // supportedFilterKeys are the keys accepted by `air list --filter KEY=VALUE`. var supportedFilterKeys = []string{"accelerator_type", "experiment", "num_accelerators", "user"} -// hasTaskFilter reports whether any filter is applied to a run's task fields -// (experiment or accelerators), i.e. matched after a run is fetched rather than -// while scanning. The index path uses this to skip its newest-N truncation, so a -// dropped match doesn't shrink the result below --limit. -func (f listFilters) hasTaskFilter() bool { - return f.Experiment != "" || f.AcceleratorType != "" || f.NumAccelerators != nil -} - // listFilters holds the parsed `--filter` values for `air list`. type listFilters struct { // User is an exact creator-email match diff --git a/experimental/air/cmd/list_index.go b/experimental/air/cmd/list_index.go index 243affd664f..873f4e516df 100644 --- a/experimental/air/cmd/list_index.go +++ b/experimental/air/cmd/list_index.go @@ -1,146 +1,172 @@ package aircmd import ( - "cmp" "context" - "slices" + "net/http" + "strings" - "github.com/databricks/cli/libs/cache" + "github.com/databricks/cli/libs/log" "github.com/databricks/databricks-sdk-go" + "github.com/databricks/databricks-sdk-go/client" + "github.com/databricks/databricks-sdk-go/service/jobs" ) -// indexStrategy serves the caller's own runs from the AiTrainingService index: -// it fetches every run id up front (cheap id+timestamp pairs), orders them -// newest-first, keeps the newest `limit`, then hydrates them into full rows in -// want-sized batches via Jobs runs/get. Terminal rows are cached so repeat calls -// skip the network. Unlike the Jobs scan it can't lazy-page (it must sort the -// whole id set first), but it still yields in batches so the table paints early. -type indexStrategy struct { +// airRunTypeFilter is the ListRuns `filter` fragment that narrows to AI Runtime +// runs server-side, using the RunsListVisible index. It is a PUBLIC_UNDOCUMENTED +// query param not modeled by the typed SDK, so it is sent via the raw client. +// Gated by the jobs `enableAirRunTypeFilter` flag; when off, the server rejects +// it with INVALID_PARAMETER_VALUE ("not supported in this workspace") and the +// caller falls back to the client-side scan. +const airRunTypeFilter = "runType=AI_RUNTIME" + +// filterUnsupportedMarker identifies the flag-off rejection so it can be told +// apart from a genuinely malformed filter (both are INVALID_PARAMETER_VALUE). +const filterUnsupportedMarker = "not supported in this workspace" + +// indexRunStrategy lists AIR runs via the server-side AI-Runtime filter (the fast, +// complete, indexed path). On the first request it either succeeds — and pages +// the filtered runs/list — or, when the workspace has the filter flag off, +// transparently falls back to the client-side scan for the rest of the listing. +type indexRunStrategy struct { ctx context.Context w *databricks.WorkspaceClient + api *client.DatabricksClient activeOnly bool + userFilter string filters listFilters - limit int - cache *cache.Cache - ids []int64 // newest-first run ids to hydrate, resolved on first next() - pos int - loaded bool + started bool + pageToken string + drained bool + pending []listedRun // matched rows from a page beyond the last want + fallback *runScanStrategy // set when the filter is unsupported } -func newIndexStrategy(ctx context.Context, w *databricks.WorkspaceClient, q listQuery, limit int) *indexStrategy { - return &indexStrategy{ +func newIndexRunStrategy(ctx context.Context, w *databricks.WorkspaceClient, q listQuery) (*indexRunStrategy, error) { + api, err := client.New(w.Config) + if err != nil { + return nil, err + } + return &indexRunStrategy{ ctx: ctx, w: w, + api: api, activeOnly: q.activeOnly, + userFilter: q.userFilter, filters: q.filters, - limit: limit, - cache: newListCache(ctx), - } + }, nil } -// load fetches and orders the index once. It returns an error only when the -// index endpoint itself fails, letting the caller fall back to the Jobs scan. -func (s *indexStrategy) load() error { - refs, err := listAiTrainingWorkflows(s.ctx, s.w, s.activeOnly) - if err != nil { - return err - } - slices.SortFunc(refs, func(a, b workflowRef) int { return cmp.Compare(b.submitTimeMs, a.submitTimeMs) }) - // Keep only the newest `limit` ids so hydration is bounded — but skip that when - // a task filter is active, since it drops matches post-hydration and we'd - // otherwise return fewer than `limit`. The caller stops pulling at `limit`. - if s.limit > 0 && len(refs) > s.limit && !s.filters.hasTaskFilter() { - refs = refs[:s.limit] - } - s.ids = make([]int64, len(refs)) - for i, r := range refs { - s.ids[i] = r.jobRunID - } - s.loaded = true - return nil +// listRunsResponse is the runs/list shape the raw client unmarshals into. It +// mirrors jobs.ListRunsResponse (the typed SDK is bypassed only to pass the +// undocumented filter query param). +type listRunsResponse struct { + Runs []jobs.BaseRun `json:"runs"` + NextPageToken string `json:"next_page_token"` } -func (s *indexStrategy) next(want int) ([]listedRun, error) { - if !s.loaded { - if err := s.load(); err != nil { - return nil, err - } +func (s *indexRunStrategy) next(want int) ([]listedRun, error) { + // A prior request found the filter unsupported: everything now goes through the + // client-side scan. + if s.fallback != nil { + return s.fallback.next(want) } + // Serve leftovers from a prior over-fetched page first. var entries []listedRun - for len(entries) < want && s.pos < len(s.ids) { - end := min(s.pos+want-len(entries), len(s.ids)) - batch := s.ids[s.pos:end] - s.pos = end + if len(s.pending) > 0 { + n := min(want, len(s.pending)) + entries = append(entries, s.pending[:n]...) + s.pending = s.pending[n:] + } - rows, err := s.hydrate(batch) + for len(entries) < want && !s.drained { + page, err := s.fetchPage() if err != nil { + // The workspace has the filter flag off: switch to the scan for the whole + // listing. fetchPage only errs on the first request (before any run is + // yielded), so nothing consumed is lost. + if isFilterUnsupported(err) { + log.Debugf(s.ctx, "air list: AI-Runtime filter unsupported, falling back to scan: %v", err) + s.fallback = newRunScanStrategy(s.ctx, s.w, listQuery{ + activeOnly: s.activeOnly, + userFilter: s.userFilter, + filters: s.filters, + }) + return s.fallback.next(want) + } return nil, err } - entries = append(entries, rows...) - } - return entries, nil -} - -func (s *indexStrategy) done() bool { - return s.loaded && s.pos >= len(s.ids) -} + if s.pageToken == "" { + s.drained = true + } -// truncated is always false: the index path is bounded by limit, not a scan cap. -func (s *indexStrategy) truncated() bool { return false } - -// hydrate turns a batch of run ids into rows, serving cached terminal rows -// without a network call and fetching the rest via runs/get. Freshly hydrated -// terminal runs are cached. Results keep the input (newest-first) order, then -// the batch is re-sorted by start time since concurrent hydration reorders it. -func (s *indexStrategy) hydrate(ids []int64) ([]listedRun, error) { - host := s.w.Config.Host - - rows := make([]listedRun, 0, len(ids)) - var toFetch []int64 - // The cache key excludes the filter, so a cached row is still run through the - // active filter; a non-matching hit is dropped, not re-fetched. - for _, id := range ids { - if row, fields, ok := cachedRow(s.ctx, s.cache, host, id); ok { - if s.filters.matchesFields(fields) { - rows = append(rows, listedRun{row: row, taskRunID: id}) + for i := range page { + run := baseRunToRun(page[i]) + // The server filter narrows to AIR runs; re-check by task shape as + // insurance, and apply active-only, user-scoping, and task filters + // client-side (active_only can't ride the filter query server-side). + if !isAirRun(run) { + continue + } + if s.activeOnly && !isActiveRun(run) { + continue + } + if s.userFilter != "" && run.CreatorUserName != s.userFilter { + continue + } + if !s.filters.matches(run) { + continue + } + row := listedRun{row: buildListRow(run), taskRunID: taskRunID(run)} + if len(entries) < want { + entries = append(entries, row) + } else { + s.pending = append(s.pending, row) // carry to the next next() call } - continue } - toFetch = append(toFetch, id) } + return entries, nil +} - runs, err := hydrateJobRuns(s.ctx, s.w, toFetch) +// fetchPage issues one filtered runs/list request via the raw client and advances +// the page token. +func (s *indexRunStrategy) fetchPage() ([]jobs.BaseRun, error) { + // active_only cannot be combined with the filter query server-side, so the + // indexed path fetches all statuses and applies active-only filtering in next(). + query := map[string]any{ + "expand_tasks": true, + "limit": jobsPageLimit, + "filter": airRunTypeFilter, + } + if s.pageToken != "" { + query["page_token"] = s.pageToken + } + var resp listRunsResponse + err := s.api.Do(s.ctx, http.MethodGet, "/api/2.2/jobs/runs/list", nil, nil, query, &resp) if err != nil { return nil, err } - for _, run := range runs { - fields := filterFieldsFromRun(run) - if !s.filters.matchesFields(fields) { - continue - } - row := buildListRow(run) - rows = append(rows, listedRun{row: row, taskRunID: taskRunID(run)}) - if isTerminal(run) { - start, _ := jobTiming(run) - putRow(s.ctx, s.cache, host, run.RunId, start, row, fields) - } + s.started = true + s.pageToken = resp.NextPageToken + return resp.Runs, nil +} + +func (s *indexRunStrategy) done() bool { + if s.fallback != nil { + return s.fallback.done() } + return s.drained && len(s.pending) == 0 +} - // Concurrent hydration reorders runs, so re-sort the batch newest-first. The - // ISO start timestamp sorts lexicographically; a missing time ("") sorts last. - slices.SortStableFunc(rows, func(a, b listedRun) int { - return cmp.Compare(rowStartKey(b.row), rowStartKey(a.row)) - }) - return rows, nil +// truncated reports the scan cap only when the fallback is active; the indexed +// path is complete (no cap). +func (s *indexRunStrategy) truncated() bool { + return s.fallback != nil && s.fallback.truncated() } -// rowStartKey returns a row's ISO start timestamp for ordering, or "" when the -// run hasn't started (which sorts last under descending comparison). -func rowStartKey(r listRow) string { - if r.StartedAt == nil { - return "" - } - return *r.StartedAt +// isFilterUnsupported reports whether err is the flag-off rejection of the +// AI-Runtime filter, distinguishing it from a genuinely bad filter value. +func isFilterUnsupported(err error) bool { + return strings.Contains(err.Error(), filterUnsupportedMarker) } diff --git a/experimental/air/cmd/list_index_test.go b/experimental/air/cmd/list_index_test.go index f9f6151aae4..84ec76af04a 100644 --- a/experimental/air/cmd/list_index_test.go +++ b/experimental/air/cmd/list_index_test.go @@ -1,11 +1,8 @@ package aircmd import ( - "encoding/json" - "fmt" "net/http" "net/http/httptest" - "strconv" "testing" "github.com/databricks/databricks-sdk-go/service/jobs" @@ -13,201 +10,107 @@ import ( "github.com/stretchr/testify/require" ) -// indexRun is a terminal AIR run with a start time, for index-path hydration. -func indexRun(id, startMillis int64) jobs.Run { - r := airRun(id, "me@example.com", "GPU_1xH100", 1, "/Users/me@example.com/exp") - r.State = &jobs.RunState{LifeCycleState: jobs.RunLifeCycleStateTerminated, ResultState: jobs.RunResultStateSuccess} - r.Tasks[0].StartTime = startMillis - r.Tasks[0].EndTime = startMillis + 1000 - return r -} - -// indexAndGetServer serves the AiTrainingService index (a single page of the -// given refs) and runs/get for each id, recording hit counts per endpoint. A -// runID in forbidden returns 403; in missing returns 404. -type indexHits struct{ index, get int } - -func indexAndGetServer(t *testing.T, refs []workflowRef, runs map[int64]jobs.Run, forbidden, missing map[int64]bool) (*httptest.Server, *indexHits) { +// filterCapturingServer records the `filter` query param on runs/list and serves +// the given body; it lets a test assert the AI-Runtime filter is sent. +func filterCapturingServer(t *testing.T, gotFilter *string, body string) *httptest.Server { t.Helper() - hits := &indexHits{} - wfs := make([]map[string]any, len(refs)) - for i, r := range refs { - wfs[i] = map[string]any{"job_run_id": strconv.FormatInt(r.jobRunID, 10), "submit_time": map[string]any{"seconds": r.submitTimeMs / 1000}} - } srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - switch r.URL.Path { - case aiTrainingWorkflowsPath: - hits.index++ - _ = json.NewEncoder(w).Encode(map[string]any{"training_workflows": wfs}) - case "/api/2.2/jobs/runs/get": - hits.get++ - id, _ := strconv.ParseInt(r.URL.Query().Get("run_id"), 10, 64) - if forbidden[id] { - w.WriteHeader(http.StatusForbidden) - _, _ = w.Write([]byte(`{"message":"forbidden"}`)) - return - } - if missing[id] { - w.WriteHeader(http.StatusNotFound) - _, _ = w.Write([]byte(`{"message":"not found"}`)) - return - } - run := runs[id] - _ = json.NewEncoder(w).Encode(run) - default: - _, _ = w.Write([]byte(`{}`)) + if r.URL.Path == "/api/2.2/jobs/runs/list" { + *gotFilter = r.URL.Query().Get("filter") + _, _ = w.Write([]byte(body)) + return } + _, _ = w.Write([]byte(`{}`)) })) t.Cleanup(srv.Close) - return srv, hits -} - -func TestIndexStrategyOrdersAndLimits(t *testing.T) { - // Three runs, out of submit-time order; newest two should win, newest-first. - refs := []workflowRef{ - {jobRunID: 1, submitTimeMs: 1000_000}, - {jobRunID: 2, submitTimeMs: 3000_000}, - {jobRunID: 3, submitTimeMs: 2000_000}, - } - runs := map[int64]jobs.Run{ - 1: indexRun(1, 1000_000), - 2: indexRun(2, 3000_000), - 3: indexRun(3, 2000_000), - } - srv, _ := indexAndGetServer(t, refs, runs, nil, nil) - t.Setenv("DATABRICKS_CACHE_ENABLED", "false") - - f := newRunFetcher(t.Context(), newTestWorkspaceClient(t, srv.URL), listQuery{ - userFilter: "me@example.com", currentUser: "me@example.com", limit: 2, - }) - rows, err := f.next(10) - require.NoError(t, err) - require.Len(t, rows, 2) - assert.Equal(t, "2", rows[0].RunID) // submit 3000 - assert.Equal(t, "3", rows[1].RunID) // submit 2000 - assert.True(t, f.exhausted) + return srv } -func TestIndexStrategyOverFetchesWithTaskFilter(t *testing.T) { - // With a task filter and limit 1, the newest run doesn't match; the strategy - // must keep hydrating past `limit` to find the match rather than truncating. - refs := []workflowRef{ - {jobRunID: 1, submitTimeMs: 3000_000}, - {jobRunID: 2, submitTimeMs: 2000_000}, - } - run1 := indexRun(1, 3000_000) - run1.Tasks[0].AiRuntimeTask.Experiment = "/Users/me@example.com/llama" - run2 := indexRun(2, 2000_000) - run2.Tasks[0].AiRuntimeTask.Experiment = "/Users/me@example.com/qwen" - srv, _ := indexAndGetServer(t, refs, map[int64]jobs.Run{1: run1, 2: run2}, nil, nil) - t.Setenv("DATABRICKS_CACHE_ENABLED", "false") +func TestIndexStrategySendsAirFilter(t *testing.T) { + var filter string + run := airBaseRun(9, "me@example.com", "GPU_1xA10", 1, "exp") + srv := filterCapturingServer(t, &filter, runsListBody(t, "", run)) - f := newRunFetcher(t.Context(), newTestWorkspaceClient(t, srv.URL), listQuery{ - userFilter: "me@example.com", currentUser: "me@example.com", limit: 1, - filters: listFilters{Experiment: "qwen"}, - }) - rows, err := f.next(1) + rows, err := newRunFetcher(t.Context(), newTestWorkspaceClient(t, srv.URL), listQuery{ + activeOnly: true, userFilter: "me@example.com", + }).next(10) require.NoError(t, err) + assert.Equal(t, "runType=AI_RUNTIME", filter, "list must request the server-side AI-Runtime filter") require.Len(t, rows, 1) - assert.Equal(t, "2", rows[0].RunID) // found despite being the older, second id -} - -func TestIndexStrategyDropsForbiddenAndMissing(t *testing.T) { - refs := []workflowRef{ - {jobRunID: 1, submitTimeMs: 3000_000}, - {jobRunID: 2, submitTimeMs: 2000_000}, - {jobRunID: 3, submitTimeMs: 1000_000}, - } - runs := map[int64]jobs.Run{1: indexRun(1, 3000_000), 3: indexRun(3, 1000_000)} - srv, _ := indexAndGetServer(t, refs, runs, map[int64]bool{2: true}, nil) - t.Setenv("DATABRICKS_CACHE_ENABLED", "false") - - f := newRunFetcher(t.Context(), newTestWorkspaceClient(t, srv.URL), listQuery{ - userFilter: "me@example.com", currentUser: "me@example.com", limit: 10, - }) - rows, err := f.next(10) - require.NoError(t, err) - require.Len(t, rows, 2) // run 2 (403) dropped - assert.Equal(t, "1", rows[0].RunID) - assert.Equal(t, "3", rows[1].RunID) + assert.Equal(t, "9", rows[0].RunID) } -func TestIndexStrategyPropagatesServerError(t *testing.T) { - refs := []workflowRef{{jobRunID: 1, submitTimeMs: 1000_000}} +// paramCapturingServer records the filter and active_only query params on runs/list. +func paramCapturingServer(t *testing.T, gotFilter, gotActiveOnly *string, body string) *httptest.Server { + t.Helper() srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - switch r.URL.Path { - case aiTrainingWorkflowsPath: - wfs := []map[string]any{{"job_run_id": "1", "submit_time": map[string]any{"seconds": int64(1000)}}} - _ = json.NewEncoder(w).Encode(map[string]any{"training_workflows": wfs}) - case "/api/2.2/jobs/runs/get": - w.WriteHeader(http.StatusInternalServerError) - _, _ = w.Write([]byte(`{"message":"boom"}`)) - default: - _, _ = w.Write([]byte(`{}`)) + if r.URL.Path == "/api/2.2/jobs/runs/list" { + *gotFilter = r.URL.Query().Get("filter") + *gotActiveOnly = r.URL.Query().Get("active_only") + _, _ = w.Write([]byte(body)) + return } + _, _ = w.Write([]byte(`{}`)) })) t.Cleanup(srv.Close) - _ = refs - t.Setenv("DATABRICKS_CACHE_ENABLED", "false") - - f := newRunFetcher(t.Context(), newTestWorkspaceClient(t, srv.URL), listQuery{ - userFilter: "me@example.com", currentUser: "me@example.com", limit: 10, - }) - _, err := f.next(10) - require.Error(t, err) // 500 is systemic, not an ACL drop + return srv } -func TestNewListStrategyGate(t *testing.T) { - // --all-users and other-user filters must NOT touch the index endpoint. - cases := []struct { - name string - q listQuery - wantIndex bool - }{ - {"active default → scan", listQuery{activeOnly: true, userFilter: "me@example.com", currentUser: "me@example.com"}, false}, - {"all-status self → index", listQuery{userFilter: "me@example.com", currentUser: "me@example.com", limit: 5}, true}, - {"all-status all-users → scan", listQuery{allUsers: true}, false}, - {"all-status other user → scan", listQuery{userFilter: "other@example.com", currentUser: "me@example.com"}, false}, - } - for _, tc := range cases { - t.Run(tc.name, func(t *testing.T) { - var indexHit bool - srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - if r.URL.Path == aiTrainingWorkflowsPath { - indexHit = true - } - _, _ = w.Write([]byte(`{}`)) - })) - t.Cleanup(srv.Close) - t.Setenv("DATABRICKS_CACHE_ENABLED", "false") +func TestIndexStrategyActiveOnlyFiltersClientSide(t *testing.T) { + // active_only cannot ride the filter query server-side, so the indexed request + // must NOT send active_only; active filtering happens client-side. + running := airBaseRun(1, "me@example.com", "GPU_1xA10", 1, "exp-a") + terminated := airBaseRun(2, "me@example.com", "GPU_1xA10", 1, "exp-b") + terminated.State = &jobs.RunState{LifeCycleState: jobs.RunLifeCycleStateTerminated, ResultState: jobs.RunResultStateSuccess} - newListStrategy(t.Context(), newTestWorkspaceClient(t, srv.URL), tc.q) - assert.Equal(t, tc.wantIndex, indexHit) - }) - } + var filter, activeOnly string + srv := paramCapturingServer(t, &filter, &activeOnly, runsListBody(t, "", running, terminated)) + + rows, err := newRunFetcher(t.Context(), newTestWorkspaceClient(t, srv.URL), listQuery{ + activeOnly: true, userFilter: "me@example.com", + }).next(10) + require.NoError(t, err) + assert.Equal(t, "runType=AI_RUNTIME", filter) + assert.Empty(t, activeOnly, "active_only must not be sent alongside the filter query") + require.Len(t, rows, 1, "the terminated run is filtered out client-side") + assert.Equal(t, "1", rows[0].RunID) } -func TestNewListStrategyFallsBackWhenIndexFails(t *testing.T) { - // Index 500 must silently fall back to the Jobs scan, not fail the command. +// unsupportedFilterServer rejects the AI-Runtime filter like a flag-off +// workspace (400 INVALID_PARAMETER_VALUE, "not supported in this workspace"), +// so a test can exercise the fallback to the client-side scan. +func unsupportedFilterServer(t *testing.T, scanBody string) *httptest.Server { + t.Helper() srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - if r.URL.Path == aiTrainingWorkflowsPath { - w.WriteHeader(http.StatusInternalServerError) - _, _ = w.Write([]byte(`{"message":"boom"}`)) - return - } if r.URL.Path == "/api/2.2/jobs/runs/list" { - _, _ = fmt.Fprint(w, runsListBody(t, "", airBaseRun(1, "me@example.com", "GPU_1xH100", 1, "exp"))) + if r.URL.Query().Get("filter") != "" { + w.WriteHeader(http.StatusBadRequest) + _, _ = w.Write([]byte(`{"error_code":"INVALID_PARAMETER_VALUE","message":"runType=AI_RUNTIME filter is not supported in this workspace"}`)) + return + } + _, _ = w.Write([]byte(scanBody)) return } _, _ = w.Write([]byte(`{}`)) })) t.Cleanup(srv.Close) - t.Setenv("DATABRICKS_CACHE_ENABLED", "false") + return srv +} - f := newRunFetcher(t.Context(), newTestWorkspaceClient(t, srv.URL), listQuery{ - userFilter: "me@example.com", currentUser: "me@example.com", limit: 10, - }) - rows, err := f.next(10) +func TestIndexStrategyFallsBackWhenUnsupported(t *testing.T) { + // The filtered request is rejected (flag off); the scan (no filter) then serves + // the runs. + scan := runsListBody(t, "", + airBaseRun(1, "me@example.com", "GPU_1xA10", 1, "exp-a"), + airBaseRun(2, "me@example.com", "GPU_1xA10", 1, "exp-b"), + ) + srv := unsupportedFilterServer(t, scan) + + rows, err := newRunFetcher(t.Context(), newTestWorkspaceClient(t, srv.URL), listQuery{ + activeOnly: true, userFilter: "me@example.com", + }).next(10) require.NoError(t, err) - require.Len(t, rows, 1) // served by the Jobs scan fallback + require.Len(t, rows, 2, "falls back to the scan and returns its runs") + assert.Equal(t, "1", rows[0].RunID) + assert.Equal(t, "2", rows[1].RunID) } diff --git a/experimental/air/cmd/list_test.go b/experimental/air/cmd/list_test.go index f70330240e5..a79bc29fda5 100644 --- a/experimental/air/cmd/list_test.go +++ b/experimental/air/cmd/list_test.go @@ -99,7 +99,40 @@ func TestListAirRunsExperimentFilter(t *testing.T) { assert.Equal(t, "1", rows[0].RunID) } +// queryCapturingServer records the run_type query param on runs/list, so a test +// can assert the scan does not narrow to SUBMIT_RUN 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 runs of a persistent DABs job (RunType JOB_RUN), so the 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/list returns newest-first; the scan yields in that order and stops at limit. runs := []jobs.BaseRun{ airBaseRun(1, "me@example.com", "GPU_1xH100", 1, "exp-a"), airBaseRun(2, "me@example.com", "GPU_1xH100", 1, "exp-b"), @@ -115,6 +148,7 @@ func TestListAirRunsLimitTruncates(t *testing.T) { } func TestListAirRunsPaginates(t *testing.T) { + // The scan follows runs/list page tokens across pages. page1 := runsListBody(t, "tok", airBaseRun(1, "me@example.com", "GPU_1xH100", 1, "exp-a")) page2 := runsListBody(t, "", airBaseRun(2, "me@example.com", "GPU_1xH100", 1, "exp-b")) srv := runsServer(t, page1, page2) @@ -215,7 +249,7 @@ func TestBuildListRowSweep(t *testing.T) { func TestListInvalidLimit(t *testing.T) { m := mocks.NewMockWorkspaceClient(t) ctx := cmdctx.SetWorkspaceClient(cmdio.MockDiscard(t.Context()), m.WorkspaceClient) - cmd := newListCommand() + cmd := newListRunsCommand() cmd.SetContext(ctx) require.NoError(t, cmd.Flags().Set("limit", "0")) @@ -227,7 +261,7 @@ func TestListInvalidLimit(t *testing.T) { func TestListInvalidFilter(t *testing.T) { m := mocks.NewMockWorkspaceClient(t) ctx := cmdctx.SetWorkspaceClient(cmdio.MockDiscard(t.Context()), m.WorkspaceClient) - cmd := newListCommand() + cmd := newListRunsCommand() cmd.SetContext(ctx) require.NoError(t, cmd.Flags().Set("filter", "bogus=1")) 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..0b6d5d59567 100644 --- a/experimental/air/cmd/runupload.go +++ b/experimental/air/cmd/runupload.go @@ -69,9 +69,12 @@ func buildArtifacts(cfg *runConfig, configPath string) ([]uploadItem, error) { items := []uploadItem{ {trainingConfigName, configData}, - {commandScriptName, []byte(*cfg.Command)}, + {commandScriptName, []byte(commandScript(*cfg.Command))}, } + // Always upload requirements.yaml; the server launcher aborts if it's absent. + // A file path is uploaded verbatim; otherwise the inline list (empty when + // unset) is synthesized. switch reqPath, ok := cfg.requirementsFile(); { case ok: // Resolve a relative requirements path against the config's directory. @@ -84,14 +87,13 @@ func buildArtifacts(cfg *runConfig, configPath string) ([]uploadItem, error) { } items = append(items, uploadItem{requirementsName, data}) default: - if deps, ok := cfg.inlineDependencies(); ok { - version, _ := cfg.runtimeVersion() - data, err := yaml.Marshal(requirementsDoc{Version: version, Dependencies: deps}) - if err != nil { - return nil, fmt.Errorf("failed to synthesize requirements.yaml: %w", err) - } - items = append(items, uploadItem{requirementsName, data}) + deps, _ := cfg.inlineDependencies() + version, _ := cfg.runtimeVersion() + data, err := yaml.Marshal(requirementsDoc{Version: version, Dependencies: deps}) + if err != nil { + return nil, fmt.Errorf("failed to synthesize requirements.yaml: %w", err) } + items = append(items, uploadItem{requirementsName, data}) } if len(cfg.Parameters) > 0 { @@ -102,8 +104,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 +131,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..1f46dd7ca77 100644 --- a/experimental/air/cmd/runupload_test.go +++ b/experimental/air/cmd/runupload_test.go @@ -52,9 +52,21 @@ func TestBuildArtifacts_CommandAndConfig(t *testing.T) { items, err := buildArtifacts(cfg, path) require.NoError(t, err) - assert.Equal(t, []string{trainingConfigName, commandScriptName}, itemNames(items)) + // A config with no dependencies (environment: {}) still uploads an empty + // requirements.yaml: the server launcher aborts if the file is absent. This is + // the regression guard for the no-deps submit failure. + assert.Equal(t, []string{trainingConfigName, commandScriptName, requirementsName}, 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)) + // The synthesized requirements.yaml carries an empty dependency list. + assert.Equal(t, "dependencies: []\n", string(items[2].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 +96,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...) -}