diff --git a/.gitignore b/.gitignore index 4b82c6d1521..9a49c7c4dac 100644 --- a/.gitignore +++ b/.gitignore @@ -78,3 +78,4 @@ tools/gofumpt .claude/worktrees/ .worktrees/ .isaac/ +/databricks diff --git a/.nextchanges/bundles/5818.md b/.nextchanges/bundles/5818.md new file mode 100644 index 00000000000..5156b3b2a35 --- /dev/null +++ b/.nextchanges/bundles/5818.md @@ -0,0 +1 @@ +* `bundle validate` now reports a clear error when a `sql_warehouse` is missing a `name` (including whitespace-only names), and a warning when a grant is missing a `principal` ([#5818](https://github.com/databricks/cli/pull/5818)). diff --git a/.nextchanges/bundles/bundle-generate-job-download-script.md b/.nextchanges/bundles/bundle-generate-job-download-script.md new file mode 100644 index 00000000000..14c29808030 --- /dev/null +++ b/.nextchanges/bundles/bundle-generate-job-download-script.md @@ -0,0 +1 @@ +* `bundle generate job` can now download workspace files referenced by `spark_python_task`, rewriting them to a relative path like it already does for notebooks. This is opt-in via the `--download-spark-python-files` flag ([#5799](https://github.com/databricks/cli/pull/5799)). diff --git a/.nextchanges/cli/filer-gcs-recursive-delete.md b/.nextchanges/cli/filer-gcs-recursive-delete.md new file mode 100644 index 00000000000..19f446a1b65 --- /dev/null +++ b/.nextchanges/cli/filer-gcs-recursive-delete.md @@ -0,0 +1 @@ +Fixed `databricks fs rm -r` failing on UC Volumes backed by GCS when a directory becomes empty during recursive deletion ([#5958](https://github.com/databricks/cli/pull/5958)). diff --git a/acceptance/experimental/air/cancel/test.toml b/acceptance/experimental/air/cancel/test.toml index e7e3fca1f68..14a54ccb754 100644 --- a/acceptance/experimental/air/cancel/test.toml +++ b/acceptance/experimental/air/cancel/test.toml @@ -1,6 +1,3 @@ -# This command does not deploy a bundle, so no engine matrix is needed. -EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["direct"] - # The SDK occasionally probes host reachability with a HEAD request; stub it so # the test is deterministic. [[Server]] diff --git a/acceptance/experimental/air/convert-to-dabs/docker.yaml b/acceptance/experimental/air/convert-to-dabs/docker.yaml new file mode 100644 index 00000000000..848ac77fd9c --- /dev/null +++ b/acceptance/experimental/air/convert-to-dabs/docker.yaml @@ -0,0 +1,8 @@ +experiment_name: docker-test +command: python train.py +compute: + accelerator_type: GPU_1xA10 + num_accelerators: 1 +environment: + docker_image: + url: myregistry/img:tag diff --git a/acceptance/experimental/air/convert-to-dabs/out.test.toml b/acceptance/experimental/air/convert-to-dabs/out.test.toml new file mode 100644 index 00000000000..e90b6d5d1ba --- /dev/null +++ b/acceptance/experimental/air/convert-to-dabs/out.test.toml @@ -0,0 +1,3 @@ +Local = true +Cloud = false +EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["direct"] diff --git a/acceptance/experimental/air/convert-to-dabs/output.txt b/acceptance/experimental/air/convert-to-dabs/output.txt new file mode 100644 index 00000000000..f7ab2037271 --- /dev/null +++ b/acceptance/experimental/air/convert-to-dabs/output.txt @@ -0,0 +1,82 @@ + +=== convert an AIR run YAML into a DABs bundle +>>> [CLI] experimental air convert-to-dabs train.yaml --output-dir generated +Wrote a Databricks Asset Bundle to generated: + databricks.yml + training_config.yaml + command.sh + requirements.yaml + code_source.tar.gz + +Notes: + - code_source was packaged into code_source.tar.gz; re-run convert-to-dabs to re-snapshot after code changes. + +To deploy and run this workload as a bundle: + 1. cd generated + 2. databricks bundle validate + 3. databricks bundle deploy + 4. databricks bundle run torchrun-a10-smoke-test + +bundle deploy uploads the code source and launch scripts automatically. + +Unlike `air run` (which submits an ephemeral run), bundle deploy creates a +persistent job that is not garbage-collected. When you are done, remove the +job and its uploaded files with: + databricks bundle destroy + +=== emitted databricks.yml +>>> cat generated/databricks.yml +bundle: + name: torchrun-a10-smoke-test +targets: + dev: + mode: development + default: true +resources: + jobs: + torchrun-a10-smoke-test: + name: torchrun-a10-smoke-test + tasks: + - task_key: torchrun-a10-smoke-test + environment_key: default + ai_runtime_task: + experiment: torchrun-a10-smoke-test + deployments: + - command_path: ./command.sh + compute: + accelerator_type: GPU_1xA10 + accelerator_count: 1 + code_source_path: ./code_source.tar.gz + environments: + - environment_key: default + spec: + environment_version: "5" + dependencies: + - numpy + +=== the generated command.sh carries the run command +>>> cat generated/command.sh +torchrun --nproc_per_node=1 train.py +=== the code source is packaged as a tarball +>>> ls generated +code_source.tar.gz +command.sh +databricks.yml +requirements.yaml +training_config.yaml + +=== the emitted bundle validates +>>> [CLI] bundle validate +Name: torchrun-a10-smoke-test +Target: dev +Workspace: + User: [USERNAME] + Path: /Workspace/Users/[USERNAME]/.bundle/torchrun-a10-smoke-test/dev + +Validation OK! + +=== docker_image is not supported yet +>>> [CLI] experimental air convert-to-dabs docker.yaml --output-dir generated-docker +Error: environment.docker_image is not yet supported by convert-to-dabs + +Exit code: 1 diff --git a/acceptance/experimental/air/convert-to-dabs/script b/acceptance/experimental/air/convert-to-dabs/script new file mode 100644 index 00000000000..4e4bd185dbc --- /dev/null +++ b/acceptance/experimental/air/convert-to-dabs/script @@ -0,0 +1,19 @@ +title "convert an AIR run YAML into a DABs bundle" +trace $CLI experimental air convert-to-dabs train.yaml --output-dir generated + +title "emitted databricks.yml" +trace cat generated/databricks.yml + +title "the generated command.sh carries the run command" +trace cat generated/command.sh + +title "the code source is packaged as a tarball" +trace ls generated && true + +title "the emitted bundle validates" +cd generated +trace $CLI bundle validate +cd .. + +title "docker_image is not supported yet" +errcode trace $CLI experimental air convert-to-dabs docker.yaml --output-dir generated-docker diff --git a/acceptance/experimental/air/convert-to-dabs/src/train.py b/acceptance/experimental/air/convert-to-dabs/src/train.py new file mode 100644 index 00000000000..c859094afdf --- /dev/null +++ b/acceptance/experimental/air/convert-to-dabs/src/train.py @@ -0,0 +1 @@ +print("train") diff --git a/acceptance/experimental/air/convert-to-dabs/test.toml b/acceptance/experimental/air/convert-to-dabs/test.toml new file mode 100644 index 00000000000..19643c87452 --- /dev/null +++ b/acceptance/experimental/air/convert-to-dabs/test.toml @@ -0,0 +1,3 @@ +# The command writes a bundle under generated/; those are generated artifacts, +# not committed inputs, so exclude them from the repo-diff check. +Ignore = ["generated", "generated-docker"] diff --git a/acceptance/experimental/air/convert-to-dabs/train.yaml b/acceptance/experimental/air/convert-to-dabs/train.yaml new file mode 100644 index 00000000000..bae47bc4dbd --- /dev/null +++ b/acceptance/experimental/air/convert-to-dabs/train.yaml @@ -0,0 +1,13 @@ +experiment_name: torchrun-a10-smoke-test +command: torchrun --nproc_per_node=1 train.py +compute: + accelerator_type: GPU_1xA10 + num_accelerators: 1 +environment: + version: 5 + dependencies: + - numpy +code_source: + type: snapshot + snapshot: + root_path: ./src diff --git a/acceptance/experimental/air/get-ai-runtime/test.toml b/acceptance/experimental/air/get-ai-runtime/test.toml index de442aca0ba..07d75ae83c8 100644 --- a/acceptance/experimental/air/get-ai-runtime/test.toml +++ b/acceptance/experimental/air/get-ai-runtime/test.toml @@ -1,6 +1,3 @@ -# This command does not deploy a bundle, so no engine matrix is needed. -EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["direct"] - # On Windows, Git Bash rewrites the leading-/ workspace paths passed to # `workspace mkdirs`/`import` into C:/... paths; disable that conversion. [Env] diff --git a/acceptance/experimental/air/get/test.toml b/acceptance/experimental/air/get/test.toml index e0ebbb2ba35..3545a34616e 100644 --- a/acceptance/experimental/air/get/test.toml +++ b/acceptance/experimental/air/get/test.toml @@ -1,6 +1,3 @@ -# This command does not deploy a bundle, so no engine matrix is needed. -EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["direct"] - # The SDK occasionally probes host reachability with a HEAD request; stub it so # the test is deterministic. [[Server]] diff --git a/acceptance/experimental/air/help/output.txt b/acceptance/experimental/air/help/output.txt index ee89e778d6b..f2ab58d0a38 100644 --- a/acceptance/experimental/air/help/output.txt +++ b/acceptance/experimental/air/help/output.txt @@ -10,12 +10,13 @@ Usage: databricks experimental air [command] Available Commands: - cancel Cancel one or more runs - get Show status, configuration, and timing details for a specific run - list List your active runs for the current profile (use --all-status for finished runs) - logs Stream or fetch logs for a run - register-image Mirror a Docker image into the workspace registry - run Submit a training workload from a YAML config + cancel Cancel one or more runs + convert-to-dabs Convert an AIR run YAML into a Databricks Asset Bundle + get Show status, configuration, and timing details for a specific run + list List your active runs for the current profile (use --all-status for finished runs) + logs Stream or fetch logs for a run + register-image Mirror a Docker image into the workspace registry + run Submit a training workload from a YAML config Flags: -h, --help help for air @@ -47,3 +48,24 @@ Global Flags: -o, --output type output type: text or json (default text) -p, --profile string ~/.databrickscfg profile -t, --target string bundle target to use (if applicable) + +=== logs help +>>> [CLI] experimental air logs --help +Stream logs from an active run, or fetch logs from a completed run. + +Usage: + databricks experimental air logs JOB_RUN_ID [flags] + +Flags: + --download-to string Download all logs to this directory instead of printing + -h, --help help for logs + --lines int For completed runs, print the last N lines (default 10000) + --minutes int Fetch only logs from the last N minutes + --node int Fetch logs from this node + --retry int View logs from a specific retry attempt; -1 means latest (default -1) + +Global Flags: + --debug enable debug logging + -o, --output type output type: text or json (default text) + -p, --profile string ~/.databrickscfg profile + -t, --target string bundle target to use (if applicable) diff --git a/acceptance/experimental/air/help/script b/acceptance/experimental/air/help/script index 81f3907e4f5..91dc12567a6 100644 --- a/acceptance/experimental/air/help/script +++ b/acceptance/experimental/air/help/script @@ -6,3 +6,6 @@ trace $CLI experimental air --help title "list help" trace $CLI experimental air list --help + +title "logs help" +trace $CLI experimental air logs --help diff --git a/acceptance/experimental/air/help/test.toml b/acceptance/experimental/air/help/test.toml deleted file mode 100644 index fa9e389f4aa..00000000000 --- a/acceptance/experimental/air/help/test.toml +++ /dev/null @@ -1,2 +0,0 @@ -# --help prints without authenticating, so no server stubs are needed. -EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["direct"] diff --git a/acceptance/experimental/air/list/test.toml b/acceptance/experimental/air/list/test.toml index 10c2a7600b8..bcaa75fb34a 100644 --- a/acceptance/experimental/air/list/test.toml +++ b/acceptance/experimental/air/list/test.toml @@ -1,6 +1,3 @@ -# This command does not deploy a bundle, so no engine matrix is needed. -EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["direct"] - # Disable the on-disk run cache so --all-status output is deterministic across runs. [Env] DATABRICKS_CACHE_ENABLED = "false" diff --git a/acceptance/experimental/air/logs-mlflow-fallback/out.test.toml b/acceptance/experimental/air/logs-mlflow-fallback/out.test.toml new file mode 100644 index 00000000000..e90b6d5d1ba --- /dev/null +++ b/acceptance/experimental/air/logs-mlflow-fallback/out.test.toml @@ -0,0 +1,3 @@ +Local = true +Cloud = false +EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["direct"] diff --git a/acceptance/experimental/air/logs-mlflow-fallback/output.txt b/acceptance/experimental/air/logs-mlflow-fallback/output.txt new file mode 100644 index 00000000000..874ad273314 --- /dev/null +++ b/acceptance/experimental/air/logs-mlflow-fallback/output.txt @@ -0,0 +1,8 @@ + +=== logs falls back to mlflow (no logs) +>>> [CLI] experimental air logs 123 +No logs available for run 123. Run terminated in state SUCCESS + +=== logs falls back to mlflow (json) +>>> [CLI] experimental air logs 123 -o json +{"type":"ERROR","ts":"[TIMESTAMP]","node":0,"line":"No logs available for run 123. Run terminated in state SUCCESS"} diff --git a/acceptance/experimental/air/logs-mlflow-fallback/script b/acceptance/experimental/air/logs-mlflow-fallback/script new file mode 100644 index 00000000000..ad8856a036a --- /dev/null +++ b/acceptance/experimental/air/logs-mlflow-fallback/script @@ -0,0 +1,9 @@ +# Bricklens is gated off (FEATURE_DISABLED), so the command falls back to the +# MLflow log path. With no MLflow run id resolvable, the fallback reports no +# logs and exits non-zero — proving the try/catch routes to MLflow. + +title "logs falls back to mlflow (no logs)" +errcode trace $CLI experimental air logs 123 + +title "logs falls back to mlflow (json)" +errcode trace $CLI experimental air logs 123 -o json diff --git a/acceptance/experimental/air/logs-mlflow-fallback/test.toml b/acceptance/experimental/air/logs-mlflow-fallback/test.toml new file mode 100644 index 00000000000..9b9bcab8e68 --- /dev/null +++ b/acceptance/experimental/air/logs-mlflow-fallback/test.toml @@ -0,0 +1,32 @@ +# The SDK occasionally probes host reachability with a HEAD request; stub it so +# the test is deterministic. +[[Server]] +Pattern = "HEAD /" +Response.Body = '' + +# A completed run. +[[Server]] +Pattern = "GET /api/2.2/jobs/runs/get" +Response.Body = ''' +{ + "run_id": 123, + "start_time": 1700000000000, + "end_time": 1700000012000, + "state": {"life_cycle_state": "TERMINATED", "result_state": "SUCCESS"}, + "tasks": [{"task_key": "train", "run_id": 456, "attempt_number": 0}] +} +''' + +# Bricklens is gated off by the backend SAFE flag, forcing the MLflow fallback. +[[Server]] +Pattern = "GET /api/2.0/ai-training/workflows/by-run-id/123/logs" +Response.StatusCode = 403 +Response.Body = ''' +{"error_code": "FEATURE_DISABLED", "message": "training log streaming is not enabled"} +''' + +# The MLflow fallback has no run output to resolve an MLflow run id from, so it +# reports no logs rather than failing — exercising the fallback wiring end to end. +[[Server]] +Pattern = "GET /api/2.2/jobs/runs/get-output" +Response.Body = '{}' diff --git a/acceptance/experimental/air/logs/out.test.toml b/acceptance/experimental/air/logs/out.test.toml new file mode 100644 index 00000000000..e90b6d5d1ba --- /dev/null +++ b/acceptance/experimental/air/logs/out.test.toml @@ -0,0 +1,3 @@ +Local = true +Cloud = false +EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["direct"] diff --git a/acceptance/experimental/air/logs/output.txt b/acceptance/experimental/air/logs/output.txt new file mode 100644 index 00000000000..84d6202afdf --- /dev/null +++ b/acceptance/experimental/air/logs/output.txt @@ -0,0 +1,72 @@ + +=== logs (text, completed run) +>>> [CLI] experimental air logs 123 +step 1 +step 2 +CUDA out of memory + +=== logs (json) +>>> [CLI] experimental air logs 123 -o json +{"type":"LOG","ts":"[TIMESTAMP]","node":0,"line":"step 1"} +{"type":"LOG","ts":"[TIMESTAMP]","node":0,"line":"step 2"} +{"type":"ALERT","ts":"[TIMESTAMP]","node":0,"line":"CUDA out of memory"} +{"type":"LOG","ts":"[TIMESTAMP]","node":0,"line":"CUDA out of memory"} + +=== logs with --minutes +>>> [CLI] experimental air logs 123 --minutes 30 +step 1 +step 2 +CUDA out of memory + +=== logs with --lines +>>> [CLI] experimental air logs 123 --lines 1 +CUDA out of memory + +=== logs with --lines 0 prints nothing +>>> [CLI] experimental air logs 123 --lines 0 +No logs available for run 123. Run terminated in state SUCCESS + +=== logs from a specific retry +>>> [CLI] experimental air logs 123 --retry 0 +step 1 +step 2 +CUDA out of memory + +=== logs --lines and --minutes are mutually exclusive +>>> [CLI] experimental air logs 123 --lines 100 --minutes 30 +Error: cannot combine --lines with --minutes: --lines tails by line count, --minutes by time window + +Exit code: 1 + +=== logs --lines and --minutes are mutually exclusive (json) +>>> [CLI] experimental air logs 123 --lines 100 --minutes 30 -o json +{ + "v": 1, + "ts": "[TIMESTAMP]", + "error": { + "code": "INVALID_ARGS", + "kind": "PERMANENT", + "message": "cannot combine --lines with --minutes: --lines tails by line count, --minutes by time window", + "retryable": false + } +} + +Exit code: 1 + +=== invalid run id +>>> [CLI] experimental air logs notanumber +Error: invalid JOB_RUN_ID "notanumber": must be a positive integer + +Exit code: 1 + +=== negative node +>>> [CLI] experimental air logs 123 --node -1 +Error: invalid --node -1: must not be negative + +Exit code: 1 + +=== --download-to not implemented +>>> [CLI] experimental air logs 123 --download-to /tmp/out +Error: --download-to is not implemented yet + +Exit code: 1 diff --git a/acceptance/experimental/air/logs/script b/acceptance/experimental/air/logs/script new file mode 100644 index 00000000000..b851a8784f8 --- /dev/null +++ b/acceptance/experimental/air/logs/script @@ -0,0 +1,32 @@ +title "logs (text, completed run)" +trace $CLI experimental air logs 123 + +title "logs (json)" +trace $CLI experimental air logs 123 -o json + +title "logs with --minutes" +trace $CLI experimental air logs 123 --minutes 30 + +title "logs with --lines" +trace $CLI experimental air logs 123 --lines 1 + +title "logs with --lines 0 prints nothing" +trace $CLI experimental air logs 123 --lines 0 + +title "logs from a specific retry" +trace $CLI experimental air logs 123 --retry 0 + +title "logs --lines and --minutes are mutually exclusive" +errcode trace $CLI experimental air logs 123 --lines 100 --minutes 30 + +title "logs --lines and --minutes are mutually exclusive (json)" +errcode trace $CLI experimental air logs 123 --lines 100 --minutes 30 -o json + +title "invalid run id" +errcode trace $CLI experimental air logs notanumber + +title "negative node" +errcode trace $CLI experimental air logs 123 --node -1 + +title "--download-to not implemented" +errcode trace $CLI experimental air logs 123 --download-to /tmp/out diff --git a/acceptance/experimental/air/logs/test.toml b/acceptance/experimental/air/logs/test.toml new file mode 100644 index 00000000000..31e95adc8c1 --- /dev/null +++ b/acceptance/experimental/air/logs/test.toml @@ -0,0 +1,30 @@ +# The SDK occasionally probes host reachability with a HEAD request; stub it so +# the test is deterministic. +[[Server]] +Pattern = "HEAD /" +Response.Body = '' + +# A completed run: Bricklens serves its logs via the tail drain. +[[Server]] +Pattern = "GET /api/2.2/jobs/runs/get" +Response.Body = ''' +{ + "run_id": 123, + "start_time": 1700000000000, + "end_time": 1700000012000, + "state": {"life_cycle_state": "TERMINATED", "result_state": "SUCCESS"}, + "tasks": [{"task_key": "train", "run_id": 456, "attempt_number": 0}] +} +''' + +# Bricklens log records, returned newest-first (as the tail fetch requests); +# printed oldest-first. +[[Server]] +Pattern = "GET /api/2.0/ai-training/workflows/by-run-id/123/logs" +Response.Body = ''' +{"log_records": [ + {"time_unix_nano": 1700000003000000000, "body": "CUDA out of memory", "node_index": 0}, + {"time_unix_nano": 1700000002000000000, "body": "step 2", "node_index": 0}, + {"time_unix_nano": 1700000001000000000, "body": "step 1", "node_index": 0} +]} +''' diff --git a/acceptance/experimental/air/run-submit/output.txt b/acceptance/experimental/air/run-submit/output.txt index 8d52eed1dab..2d92122c8e6 100644 --- a/acceptance/experimental/air/run-submit/output.txt +++ b/acceptance/experimental/air/run-submit/output.txt @@ -1,9 +1,12 @@ === submit with a git code_source >>> [CLI] experimental air run -f run.yaml +Uploading [SNAPSHOT_TARBALL]... Submitted run 555 View at: [DATABRICKS_URL]/jobs/runs/555 +Tip: use --watch to stream logs until the run completes. + === the ai_runtime_task carries the code_source_path >>> print_requests.py //api/2.2/jobs/runs/submit { @@ -23,7 +26,7 @@ View at: [DATABRICKS_URL]/jobs/runs/555 "tasks": [ { "ai_runtime_task": { - "code_source_path": "/Workspace/Users/[USERNAME]/.air/repo_snapshots/001/[SNAPSHOT_TARBALL]", + "code_source_path": "/Workspace/Users/[USERNAME]/.air/repo_snapshots/.internal/[SNAPSHOT_TARBALL]", "deployments": [ { "command_path": "/Workspace/Users/[USERNAME]/.air/cli_launch/submit-smoke/submit-smoke_[RUN_ID]/command.sh", diff --git a/acceptance/experimental/air/run-submit/test.toml b/acceptance/experimental/air/run-submit/test.toml index a077e95bf99..3dc9ff81b99 100644 --- a/acceptance/experimental/air/run-submit/test.toml +++ b/acceptance/experimental/air/run-submit/test.toml @@ -1,14 +1,11 @@ # 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. +# tarball + provenance sidecars, and POSTs runs/submit. 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 = ["direct"] - # The SDK probes host reachability with a HEAD request; stub it for determinism. [[Server]] Pattern = "HEAD /" diff --git a/acceptance/experimental/air/run/output.txt b/acceptance/experimental/air/run/output.txt index a753eabd198..13c9b360ece 100644 --- a/acceptance/experimental/air/run/output.txt +++ b/acceptance/experimental/air/run/output.txt @@ -14,18 +14,29 @@ Dry run: configuration for "smoke-test" is valid; not submitting. } } -=== override not yet supported ->>> [CLI] experimental air run -f valid.yaml --dry-run --override a=b -Error: --override is not yet supported +=== override applies and logs the change +>>> [CLI] experimental air run -f valid.yaml --dry-run --override compute.num_accelerators=2 --override timeout_minutes=45 +Override: changing compute.num_accelerators from 1 to 2 +Override: setting timeout_minutes to 45 +Dry run: configuration for "smoke-test" is valid; not submitting. + +=== override of an unknown field is rejected +>>> [CLI] experimental air run -f valid.yaml --dry-run --override bogus=1 +Error: invalid --override "bogus": "bogus" is not a known field; available fields are: code_source, command, compute, env_variables, environment, experiment_name, idempotency_token, max_retries, mlflow_experiment_directory, mlflow_run_name, parameters, permissions, secrets, timeout_minutes, usage_policy_id, usage_policy_name Exit code: 1 -=== watch not yet supported ->>> [CLI] experimental air run -f valid.yaml --dry-run --watch -Error: --watch is not yet supported +=== override still runs schema validation +>>> [CLI] experimental air run -f valid.yaml --dry-run --override compute.num_accelerators=0 +Override: changing compute.num_accelerators from 1 to 0 +Error: compute.num_accelerators must be positive, got 0 Exit code: 1 +=== watch is ignored with dry-run (nothing is submitted) +>>> [CLI] experimental air run -f valid.yaml --dry-run --watch +Dry run: configuration for "smoke-test" is valid; not submitting. + === 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. diff --git a/acceptance/experimental/air/run/script b/acceptance/experimental/air/run/script index 312b2f6fecf..2a5263bceb2 100644 --- a/acceptance/experimental/air/run/script +++ b/acceptance/experimental/air/run/script @@ -4,11 +4,17 @@ trace $CLI experimental air run -f valid.yaml --dry-run title "dry-run (json)" trace $CLI experimental air run -f valid.yaml --dry-run -o json -title "override not yet supported" -errcode trace $CLI experimental air run -f valid.yaml --dry-run --override a=b +title "override applies and logs the change" +trace $CLI experimental air run -f valid.yaml --dry-run --override compute.num_accelerators=2 --override timeout_minutes=45 -title "watch not yet supported" -errcode trace $CLI experimental air run -f valid.yaml --dry-run --watch +title "override of an unknown field is rejected" +errcode trace $CLI experimental air run -f valid.yaml --dry-run --override bogus=1 + +title "override still runs schema validation" +errcode trace $CLI experimental air run -f valid.yaml --dry-run --override compute.num_accelerators=0 + +title "watch is ignored with dry-run (nothing is submitted)" +trace $CLI experimental air run -f valid.yaml --dry-run --watch title "code_source config passes validation" trace $CLI experimental air run -f with-code-source.yaml --dry-run diff --git a/acceptance/experimental/air/run/test.toml b/acceptance/experimental/air/run/test.toml deleted file mode 100644 index c228ad415d2..00000000000 --- a/acceptance/experimental/air/run/test.toml +++ /dev/null @@ -1,3 +0,0 @@ -# `air run --dry-run` validates the config locally and makes no workspace calls, -# so no engine matrix or server stubs are needed. -EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["direct"] diff --git a/acceptance/experimental/air/test.toml b/acceptance/experimental/air/test.toml new file mode 100644 index 00000000000..167f8e1a6bf --- /dev/null +++ b/acceptance/experimental/air/test.toml @@ -0,0 +1,5 @@ +# No air command deploys a bundle, so the engine matrix adds no coverage; pin a +# single engine ("direct") so the air tests run once, not across the direct+terraform +# matrix. (An empty list is rejected: it would run on both runners.) +[EnvMatrix] +DATABRICKS_BUNDLE_ENGINE = ["direct"] diff --git a/acceptance/experimental/air/unimplemented/output.txt b/acceptance/experimental/air/unimplemented/output.txt index 7db6ef1aec2..261e0c456a9 100644 --- a/acceptance/experimental/air/unimplemented/output.txt +++ b/acceptance/experimental/air/unimplemented/output.txt @@ -1,10 +1,4 @@ -=== logs ->>> [CLI] experimental air logs 123 -Error: `air logs` is not implemented yet - -Exit code: 1 - === register-image >>> [CLI] experimental air register-image my-image:latest Error: `air register-image` is not implemented yet diff --git a/acceptance/experimental/air/unimplemented/script b/acceptance/experimental/air/unimplemented/script index 19dc13ffe85..ae0366a48c9 100644 --- a/acceptance/experimental/air/unimplemented/script +++ b/acceptance/experimental/air/unimplemented/script @@ -1,7 +1,4 @@ # Each stub must fail with "not implemented"; errcode records the exit code. -title "logs" -errcode trace $CLI experimental air logs 123 - title "register-image" errcode trace $CLI experimental air register-image my-image:latest diff --git a/acceptance/experimental/air/unimplemented/test.toml b/acceptance/experimental/air/unimplemented/test.toml deleted file mode 100644 index 0ff461a4579..00000000000 --- a/acceptance/experimental/air/unimplemented/test.toml +++ /dev/null @@ -1,2 +0,0 @@ -# Stubs fail locally before any API call, so no server stubs needed. -EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["direct"] diff --git a/experimental/air/cmd/air.go b/experimental/air/cmd/air.go index fbf40a34b52..13b3ee18559 100644 --- a/experimental/air/cmd/air.go +++ b/experimental/air/cmd/air.go @@ -23,6 +23,7 @@ experimental and may change in future versions.`, cmd.AddCommand(newLogsCommand()) cmd.AddCommand(newCancelCommand()) cmd.AddCommand(newRegisterImageCommand()) + cmd.AddCommand(newConvertToDabsCommand()) return cmd } diff --git a/experimental/air/cmd/air_test.go b/experimental/air/cmd/air_test.go index 7efac253a2b..1843acfe900 100644 --- a/experimental/air/cmd/air_test.go +++ b/experimental/air/cmd/air_test.go @@ -14,7 +14,7 @@ func TestNewRegistersAllSubcommands(t *testing.T) { registered[c.Name()] = true } - want := []string{"run", "get", "list", "logs", "cancel", "register-image"} + want := []string{"run", "get", "list", "logs", "cancel", "register-image", "convert-to-dabs"} for _, name := range want { assert.True(t, registered[name], "subcommand %q is not registered", name) } diff --git a/experimental/air/cmd/convert_to_dabs.go b/experimental/air/cmd/convert_to_dabs.go new file mode 100644 index 00000000000..dc726a04a38 --- /dev/null +++ b/experimental/air/cmd/convert_to_dabs.go @@ -0,0 +1,459 @@ +package aircmd + +import ( + "context" + "errors" + "fmt" + "os" + "path/filepath" + "strconv" + "strings" + + "github.com/databricks/cli/cmd/root" + "github.com/databricks/cli/libs/cmdio" + "github.com/databricks/cli/libs/dyn" + "github.com/databricks/cli/libs/dyn/yamlsaver" + "github.com/spf13/cobra" +) + +// convert_to_dabs turns an AIR CLI run YAML into a Databricks Asset Bundle so a +// workload authored for `air run` can be deployed and managed as a bundle. +// +// The emitted bundle is schema-valid: `databricks bundle validate` accepts it, +// and `databricks bundle deploy` reproduces the same ai_runtime_task workload the +// CLI would submit. Two properties make that work without any manual upload step: +// +// - The DABs `ai_runtime_task` is the SDK jobs.AiRuntimeTask (strict schema: +// experiment + deployments[].{command_path,compute} + code_source_path). +// Framework concerns (retries, timeout) live on the surrounding task, not in +// ai_runtime_task — so they are emitted as task-level fields. +// - `bundle deploy` uploads a *local* code_source_path / command_path itself +// (collectLocalLibraries → aiRuntimeCodeSourcePattern), so we point those at +// local files and let deploy do the uploading. The user runs one deploy. +// +// The command writes the code launch artifacts (command.sh, and — since the task +// proto carries no inline env/secrets/parameters — the env_vars.json / +// secret_env_vars.json / hyperparameters.yaml sidecars the server-side launcher +// reads) into the code_source directory, exactly mirroring the CLI's own launch +// layout so the deployed workload behaves identically to `air run`. + +// dabsTargetName is the single default target emitted; a development-mode target +// is the conventional starting point for a generated bundle. +const dabsTargetName = "dev" + +func newConvertToDabsCommand() *cobra.Command { + var outputDir string + + cmd := &cobra.Command{ + Use: "convert-to-dabs ", + Args: root.ExactArgs(1), + Short: "Convert an AIR run YAML into a Databricks Asset Bundle", + Long: `Convert an AIR CLI run YAML config into a Databricks Asset Bundle (DABs). + +The emitted bundle can be deployed with the standard DABs workflow: + + databricks bundle validate + databricks bundle deploy + +bundle deploy uploads the code source and launch scripts for you, so no manual +upload step is required. This command performs a purely local translation and +does not contact the workspace.`, + } + + cmd.Flags().StringVar(&outputDir, "output-dir", "", "Directory to write the bundle into (default: a -bundle folder next to the input YAML). Accepts an absolute or relative path.") + + cmd.RunE = func(cmd *cobra.Command, args []string) error { + ctx := cmd.Context() + yamlPath := args[0] + + cfg, err := loadRunConfig(yamlPath) + if err != nil { + return err + } + + // Default the bundle next to the input YAML (a -bundle subfolder + // so the generated files don't scatter across the user's project dir). An + // explicit --output-dir — absolute or relative — overrides. We never write to + // a temp dir: the bundle is the user's artifact to keep, deploy, and manage. + dir := outputDir + if dir == "" { + dir = filepath.Join(filepath.Dir(yamlPath), cfg.ExperimentName+"-bundle") + } + + written, err := writeBundle(ctx, cfg, yamlPath, dir) + if err != nil { + return err + } + + printConvertNextSteps(ctx, dir, written, bundleResourceKey(cfg.ExperimentName), conversionNotes(cfg)) + return nil + } + + return cmd +} + +// codeSourceTarballName is the bundle-local filename of the code archive. The +// ai_runtime_task code_source_path points at it; bundle deploy uploads the single +// file and the backend unpacks it on each node (code_source_path is an archive, +// not a directory — deploy's uploader only handles files, so a bare dir fails). +const codeSourceTarballName = "code_source.tar.gz" + +// convertToDabs builds the DABs bundle value and the loose launch artifacts for a +// run config. It reads only what the run path's buildArtifacts reads, so the +// mapping is unit-testable in isolation. Returns the bundle root as a +// map[string]dyn.Value (ready for yamlsaver) and the loose artifacts (command.sh + +// env/secret/param sidecars) to write at the bundle root; the code_source itself +// is packaged into a tarball separately by writeBundle. +func convertToDabs(ctx context.Context, cfg *runConfig, configPath string) (map[string]dyn.Value, []uploadItem, error) { + // idempotency_token is intentionally not mapped: it dedups a single runs/submit + // call, which has no analogue for a persistent, repeatedly-runnable bundle job. + // + // usage_policy_name resolution is not ported (mirrors the submit path), and + // docker images have no ai_runtime_task representation yet. + if cfg.UsagePolicyName != nil { + return nil, nil, errors.New("usage_policy_name is not yet supported by convert-to-dabs") + } + if cfg.Environment != nil && cfg.Environment.DockerImage != nil { + return nil, nil, errors.New("environment.docker_image is not yet supported by convert-to-dabs") + } + // remote_volume points the code archive at a UC Volume. bundle deploy uploads + // code_source_path to the bundle's artifact path, not an arbitrary Volume, so a + // converted bundle can't honor it — reject rather than silently drop it. + if cfg.CodeSource != nil && cfg.CodeSource.Snapshot != nil && cfg.CodeSource.Snapshot.RemoteVolume != nil { + return nil, nil, errors.New("code_source.snapshot.remote_volume is not supported by convert-to-dabs; bundle deploy manages the artifact upload location") + } + + artifacts, err := buildArtifacts(cfg, configPath) + if err != nil { + return nil, nil, err + } + + root := buildBundleValue(ctx, cfg) + return root, artifacts, nil +} + +// nv builds a dyn.Value at "position" n: yamlsaver orders a map's keys by their +// Location line, so assigning ascending n values fixes the emitted key order. +// It routes through dyn.V so nested Go maps/slices are converted recursively, +// then stamps the ordering location. +func nv(v any, n int) dyn.Value { + return dyn.V(v).WithLocations([]dyn.Location{{Line: n}}) +} + +// localBundlePath renders a bundle-relative path with a leading "./" so bundle +// deploy classifies it as a local artifact to upload (see IsLibraryLocal). It is +// built from path.Join (forward slashes) so the emitted YAML is identical across +// operating systems. +func localBundlePath(p string) string { + return "./" + p +} + +// buildBundleValue assembles the bundle root as an ordered map[string]dyn.Value. +// The command script and the code tarball are bundle-local files (emitted +// "./"-prefixed) so `bundle deploy` uploads them. +func buildBundleValue(ctx context.Context, cfg *runConfig) map[string]dyn.Value { + name := cfg.ExperimentName + + // ai_runtime_task: experiment + one deployment (command_path + compute) + + // code_source_path. Only the fields the strict schema allows. + // + // Paths are emitted "./"-prefixed so bundle deploy treats them as LOCAL and + // uploads them: libraries.IsLibraryLocal classifies a bare, extensionless path + // as a PyPI package name (not a local file) and skips it, which would deploy a + // path the backend can't resolve. The "./" prefix forces local classification. + deployment := map[string]dyn.Value{ + "command_path": nv(localBundlePath(commandScriptName), 1), + "compute": nv(map[string]dyn.Value{ + "accelerator_type": nv(cfg.Compute.AcceleratorType, 1), + "accelerator_count": nv(cfg.Compute.NumAccelerators, 2), + }, 2), + } + + aiRuntimeTask := map[string]dyn.Value{ + "experiment": nv(name, 1), + "deployments": nv([]dyn.Value{dyn.V(deployment)}, 2), + } + line := 3 + if cfg.CodeSource != nil && cfg.CodeSource.Snapshot != nil { + aiRuntimeTask["code_source_path"] = nv(localBundlePath(codeSourceTarballName), line) + line++ + } + if cfg.MLflowRunName != nil { + aiRuntimeTask["mlflow_run"] = nv(*cfg.MLflowRunName, line) + line++ + } + if cfg.MLflowExperimentDirectory != nil { + aiRuntimeTask["mlflow_experiment_directory"] = nv(*cfg.MLflowExperimentDirectory, line) + } + + // Task wrapper: task_key + framework fields (retries/timeout) + env key + + // the ai_runtime_task. Framework fields live here per the schema, not inside + // ai_runtime_task. + task := map[string]dyn.Value{ + "task_key": nv(name, 1), + "environment_key": nv(aiRuntimeEnvironmentKey, 2), + } + taskLine := 3 + if cfg.MaxRetries != nil { + task["max_retries"] = nv(*cfg.MaxRetries, taskLine) + taskLine++ + } + if cfg.TimeoutMinutes != nil { + task["timeout_seconds"] = nv(cfg.timeoutSeconds(), taskLine) + taskLine++ + } + task["ai_runtime_task"] = nv(aiRuntimeTask, taskLine) + + // environments[]: version + inline dependencies (a requirements.yaml file + // ships as a launch artifact instead, so only inline deps are inlined here). + // Resolve the version through the same path `air run` uses (config, else env + // override, else the default channel) so a config without an explicit version + // still pins the version the workload would have run with — not an empty spec. + runtimeVersion, _ := cfg.runtimeVersion() + envSpec := map[string]dyn.Value{ + "environment_version": nv(dlRuntimeImage(ctx, runtimeVersion), 1), + } + if deps, ok := cfg.inlineDependencies(); ok && len(deps) > 0 { + depVals := make([]dyn.Value, len(deps)) + for i, d := range deps { + depVals[i] = dyn.V(d) + } + envSpec["dependencies"] = nv(depVals, 2) + } + environment := map[string]dyn.Value{ + "environment_key": nv(aiRuntimeEnvironmentKey, 1), + "spec": nv(envSpec, 2), + } + + job := map[string]dyn.Value{ + "name": nv(name, 1), + "tasks": nv([]dyn.Value{dyn.V(task)}, 2), + "environments": nv([]dyn.Value{dyn.V(environment)}, 3), + } + // usage_policy_id is an already-resolved budget policy id, so it maps directly + // to the job's budget_policy_id. (usage_policy_name needs server-side resolution + // and is rejected in convertToDabs.) + if cfg.UsagePolicyID != nil { + job["budget_policy_id"] = nv(*cfg.UsagePolicyID, 4) + } + if perms := buildPermissionsValue(cfg.Permissions); perms.Kind() != dyn.KindInvalid { + job["permissions"] = nv(perms.MustSequence(), 5) + } + + rootValue := map[string]dyn.Value{ + "bundle": nv(map[string]dyn.Value{ + "name": nv(name, 1), + }, 1), + "targets": nv(map[string]dyn.Value{ + dabsTargetName: nv(map[string]dyn.Value{ + "mode": nv("development", 1), + "default": nv(true, 2), + }, 1), + }, 2), + "resources": nv(map[string]dyn.Value{ + "jobs": nv(map[string]dyn.Value{ + bundleResourceKey(name): nv(job, 1), + }, 1), + }, 3), + } + return rootValue +} + +// bundleResourceKey derives a job resource key from the experiment name. The key +// is emitted as an unquoted YAML map key, and DABs' strict loader rejects a key +// that parses as a non-string scalar (a purely numeric name like "12345" -> !!int, +// or "true"/"null"). experiment_name allows exactly [alphanumeric, -, _], so the +// only unsafe keys are those that YAML types as int/float/bool/null; prefix those +// with "job_" to force a string key. The human-facing name/experiment fields keep +// the original value (yamlsaver quotes them as scalar string values). +func bundleResourceKey(name string) string { + switch strings.ToLower(name) { + case "true", "false", "null": + return "job_" + name + } + if _, err := strconv.ParseFloat(name, 64); err == nil { + return "job_" + name + } + return name +} + +// buildPermissionsValue maps run-config permissions to DABs job permissions +// (level → principal). Returns an invalid value when there are none. +func buildPermissionsValue(perms []permission) dyn.Value { + if len(perms) == 0 { + return dyn.InvalidValue + } + out := make([]dyn.Value, 0, len(perms)) + for _, p := range perms { + m := map[string]dyn.Value{"level": nv(p.Level, 1)} + switch { + case p.UserName != nil: + m["user_name"] = nv(*p.UserName, 2) + case p.GroupName != nil: + m["group_name"] = nv(*p.GroupName, 2) + case p.ServicePrincipalName != nil: + m["service_principal_name"] = nv(*p.ServicePrincipalName, 2) + } + out = append(out, dyn.V(m)) + } + return dyn.V(out) +} + +// writeBundle writes the bundle into dir: databricks.yml, the loose launch +// artifacts (command.sh + env/secret/param sidecars), and — when the config has a +// code_source — a code_source.tar.gz of the resolved root_path. All the referenced +// files are bundle-local so `bundle deploy` uploads them. It refuses to overwrite +// existing files so a re-run can't silently clobber a bundle the user has edited. +// Returns the relative paths written, for the next-steps message. +func writeBundle(ctx context.Context, cfg *runConfig, configPath, dir string) ([]string, error) { + root, artifacts, err := convertToDabs(ctx, cfg, configPath) + if err != nil { + return nil, err + } + + // Restrict perms: the bundle carries env_vars.json (literal env var values), so + // keep the dir owner-only rather than world-readable. + if err := os.MkdirAll(dir, 0o700); err != nil { + return nil, err + } + + // Refuse to clobber an existing file, with one consistent message. A user + // re-running convert into the same dir gets a clear error rather than a silent + // overwrite of edits they may have made. + checkCollision := func(name string) error { + if _, err := os.Stat(filepath.Join(dir, name)); err == nil { + return fmt.Errorf("%s already exists in %s; use --output-dir or remove it", name, dir) + } + return nil + } + writeFile := func(name string, data []byte) error { + if err := checkCollision(name); err != nil { + return err + } + return os.WriteFile(filepath.Join(dir, name), data, 0o600) + } + + if err := checkCollision("databricks.yml"); err != nil { + return nil, err + } + bundlePath := filepath.Join(dir, "databricks.yml") + // force=true: we've already run the collision check above, so SaveAsYAML's own + // guard (which references a --force flag this command doesn't have) can't fire. + if err := yamlsaver.NewSaver().SaveAsYAML(root, bundlePath, true); err != nil { + return nil, err + } + written := []string{"databricks.yml"} + + // Loose launch artifacts (command.sh + sidecars) at the bundle root. + for _, item := range artifacts { + if err := writeFile(item.name, item.data); err != nil { + return nil, fmt.Errorf("failed to write %s: %w", item.name, err) + } + written = append(written, item.name) + } + + // Code source: package the resolved root_path into a single archive next to the + // bundle. code_source_path points at this tarball; the backend unpacks it. + // Reuse the snapshot packager so a pinned git commit/branch is archived exactly + // as `air run` would (git archive of the commit), not the dirty working tree. + if cfg.CodeSource != nil && cfg.CodeSource.Snapshot != nil { + snap := cfg.CodeSource.Snapshot + repoPath, err := resolveRootPath(ctx, snap.RootPath, filepath.Dir(configPath)) + if err != nil { + return nil, err + } + plan, err := resolveSnapshotPlan(ctx, newGitRepo(repoPath), snap.Git, snap.IncludePaths) + if err != nil { + return nil, err + } + if err := checkCollision(codeSourceTarballName); err != nil { + return nil, err + } + // The output path must be absolute: the git-archive packager runs git with + // `-C repoPath`, so a relative -o would resolve against the repo dir, not the + // bundle dir (`git archive -o out/... ` -> repoPath/out/...). plain-tar uses + // the CWD and would tolerate a relative path, but absolute is correct for both. + tarball, err := filepath.Abs(filepath.Join(dir, codeSourceTarballName)) + if err != nil { + return nil, err + } + if err := packageSnapshot(ctx, repoPath, plan, tarball); err != nil { + return nil, err + } + written = append(written, codeSourceTarballName) + } + + return written, nil +} + +// printConvertNextSteps tells the user what was written and the exact deploy +// sequence, since the value of the command is a one-command deploy afterwards. +// +// It also spells out cleanup. This matters specifically for AIR users: `air run` +// submits an ephemeral runs/submit workload that the platform reaps on its own, +// whereas `bundle deploy` creates a *persistent* job that lingers until explicitly +// destroyed — DABs has no automatic GC. A user migrating from `air run` will not +// expect a durable resource, so we call out `bundle destroy` explicitly. +func printConvertNextSteps(ctx context.Context, dir string, written []string, jobKey string, notes []string) { + cmdio.LogString(ctx, fmt.Sprintf("Wrote a Databricks Asset Bundle to %s:", dir)) + for _, w := range written { + cmdio.LogString(ctx, " "+w) + } + + // Notes surface anything the user should know: fields we transformed or dropped, + // and values they may need to fill in. Migrating users otherwise can't tell what + // silently changed between their run YAML and the bundle. + if len(notes) > 0 { + cmdio.LogString(ctx, "") + cmdio.LogString(ctx, "Notes:") + for _, n := range notes { + cmdio.LogString(ctx, " - "+n) + } + } + + cmdio.LogString(ctx, "") + cmdio.LogString(ctx, "To deploy and run this workload as a bundle:") + cmdio.LogString(ctx, " 1. cd "+dir) + cmdio.LogString(ctx, " 2. databricks bundle validate") + cmdio.LogString(ctx, " 3. databricks bundle deploy") + cmdio.LogString(ctx, " 4. databricks bundle run "+jobKey) + cmdio.LogString(ctx, "") + cmdio.LogString(ctx, "bundle deploy uploads the code source and launch scripts automatically.") + cmdio.LogString(ctx, "") + cmdio.LogString(ctx, "Unlike `air run` (which submits an ephemeral run), bundle deploy creates a") + cmdio.LogString(ctx, "persistent job that is not garbage-collected. When you are done, remove the") + cmdio.LogString(ctx, "job and its uploaded files with:") + cmdio.LogString(ctx, " databricks bundle destroy") +} + +// conversionNotes lists what the conversion transformed, staged out-of-band, or +// could not represent natively — so a user migrating from `air run` can see what +// changed between their run YAML and the emitted bundle, and what they may still +// need to fill in. Best-effort: git resolution errors are ignored here (writeBundle +// surfaces them), so a note is only emitted when the state is unambiguous. +func conversionNotes(cfg *runConfig) []string { + var notes []string + + if cfg.CodeSource != nil && cfg.CodeSource.Snapshot != nil { + snap := cfg.CodeSource.Snapshot + if snap.Git != nil { + notes = append(notes, "code_source.git is pinned in the run YAML, but the bundle packages a "+ + "code_source.tar.gz snapshot at convert time; edit code_source_path if you need a different revision.") + } + notes = append(notes, "code_source was packaged into code_source.tar.gz; re-run convert-to-dabs to re-snapshot after code changes.") + } + + // env vars / secrets have no native ai_runtime_task field yet, so they ride as + // sidecar files the server-side launcher reads (same as `air run`). + if len(cfg.EnvVariables) > 0 { + notes = append(notes, "env_variables were written to env_vars.json (no native bundle field yet); they are uploaded with the code and applied at run time.") + } + if len(cfg.Secrets) > 0 { + notes = append(notes, "secrets were written to secret_env_vars.json (no native bundle field yet); they are resolved at run time.") + } + if len(cfg.Parameters) > 0 { + notes = append(notes, "parameters were written to hyperparameters.yaml; they are not a native bundle field and are passed through to the workload.") + } + + return notes +} diff --git a/experimental/air/cmd/convert_to_dabs_test.go b/experimental/air/cmd/convert_to_dabs_test.go new file mode 100644 index 00000000000..420b275b81d --- /dev/null +++ b/experimental/air/cmd/convert_to_dabs_test.go @@ -0,0 +1,355 @@ +package aircmd + +import ( + "path/filepath" + "strings" + "testing" + + "github.com/databricks/cli/libs/dyn" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestConvertToDabsCommandShape(t *testing.T) { + cmd := newConvertToDabsCommand() + assert.Equal(t, "convert-to-dabs ", cmd.Use) + assert.Empty(t, cmd.Commands(), "convert-to-dabs must not register subcommands") + // Exactly one positional (the YAML path). + assert.NoError(t, cmd.Args(cmd, []string{"run.yaml"})) + assert.Error(t, cmd.Args(cmd, []string{})) + assert.Error(t, cmd.Args(cmd, []string{"a", "b"})) +} + +// get is a small helper: read a dotted path out of the emitted bundle root. +func get(t *testing.T, root map[string]dyn.Value, path string) dyn.Value { + t.Helper() + v, err := dyn.GetByPath(dyn.V(root), dyn.MustPathFromString(path)) + require.NoError(t, err, "path %q should exist", path) + return v +} + +func has(root map[string]dyn.Value, path string) bool { + _, err := dyn.GetByPath(dyn.V(root), dyn.MustPathFromString(path)) + return err == nil +} + +// A full config maps onto a schema-shaped bundle: bundle name, job/task keys, the +// ai_runtime_task (experiment + single deployment + code_source_path), framework +// fields on the task wrapper, and the environment spec. +func TestConvertToDabsFullMapping(t *testing.T) { + cfg := minimalConfig + ` +max_retries: 2 +timeout_minutes: 30 +mlflow_run_name: run-42 +code_source: + type: snapshot + snapshot: + root_path: ./src +environment: + version: 5 + dependencies: + - numpy + - torch +` + path := writeConfigFile(t, "run.yaml", cfg) + loaded, err := loadRunConfig(path) + require.NoError(t, err) + + root, artifacts, err := convertToDabs(t.Context(), loaded, path) + require.NoError(t, err) + + name := loaded.ExperimentName + assert.Equal(t, name, get(t, root, "bundle.name").MustString()) + assert.Equal(t, "development", get(t, root, "targets.dev.mode").MustString()) + + jobPath := "resources.jobs." + name + assert.Equal(t, name, get(t, root, jobPath+".name").MustString()) + + task := jobPath + ".tasks[0]" + assert.Equal(t, name, get(t, root, task+".task_key").MustString()) + // Framework fields live on the task wrapper, not in ai_runtime_task. + assert.Equal(t, int64(2), get(t, root, task+".max_retries").MustInt()) + assert.Equal(t, int64(1800), get(t, root, task+".timeout_seconds").MustInt()) + assert.False(t, has(root, task+".ai_runtime_task.max_retries"), "retries must not be inside ai_runtime_task") + + art := task + ".ai_runtime_task" + assert.Equal(t, name, get(t, root, art+".experiment").MustString()) + assert.Equal(t, "run-42", get(t, root, art+".mlflow_run").MustString()) + assert.Equal(t, "./"+codeSourceTarballName, get(t, root, art+".code_source_path").MustString()) + + dep := art + ".deployments[0]" + assert.Equal(t, "./"+commandScriptName, get(t, root, dep+".command_path").MustString()) + assert.Equal(t, "GPU_1xH100", get(t, root, dep+".compute.accelerator_type").MustString()) + assert.Equal(t, int64(1), get(t, root, dep+".compute.accelerator_count").MustInt()) + + env := jobPath + ".environments[0]" + assert.Equal(t, "default", get(t, root, env+".environment_key").MustString()) + assert.Equal(t, "5", get(t, root, env+".spec.environment_version").MustString()) + deps := get(t, root, env+".spec.dependencies").MustSequence() + require.Len(t, deps, 2) + assert.Equal(t, "numpy", deps[0].MustString()) + + // command.sh is always an artifact; inline deps synthesize requirements.yaml. + assert.Contains(t, itemNames(artifacts), commandScriptName) + assert.Contains(t, itemNames(artifacts), requirementsName) +} + +// Optional fields are omitted rather than emitted empty: no code_source means no +// code_source_path; unset retries/timeout means no wrapper fields. +func TestConvertToDabsOmitsUnsetFields(t *testing.T) { + path := writeConfigFile(t, "run.yaml", minimalConfig) + loaded, err := loadRunConfig(path) + require.NoError(t, err) + + root, _, err := convertToDabs(t.Context(), loaded, path) + require.NoError(t, err) + + name := loaded.ExperimentName + task := "resources.jobs." + name + ".tasks[0]" + assert.False(t, has(root, task+".max_retries")) + assert.False(t, has(root, task+".timeout_seconds")) + assert.False(t, has(root, task+".ai_runtime_task.code_source_path")) + assert.False(t, has(root, task+".ai_runtime_task.mlflow_run")) + // The command still needs a home even without code_source. + assert.Equal(t, "./"+commandScriptName, get(t, root, task+".ai_runtime_task.deployments[0].command_path").MustString()) + + // Even with no environment block, the default runtime version is pinned (what + // `air run` would have used) rather than emitting an empty environment spec. + env := "resources.jobs." + name + ".environments[0]" + assert.Equal(t, "4", get(t, root, env+".spec.environment_version").MustString()) + assert.False(t, has(root, env+".spec.dependencies")) +} + +// A DATABRICKS_DL_RUNTIME_IMAGE env override flows through the same resolution +// `air run` uses, so a converted bundle pins the same version. +func TestConvertToDabsRuntimeVersionEnvOverride(t *testing.T) { + t.Setenv(dlRuntimeImageEnv, "CLIENT-GPU-7") + path := writeConfigFile(t, "run.yaml", minimalConfig) + loaded, err := loadRunConfig(path) + require.NoError(t, err) + + root, _, err := convertToDabs(t.Context(), loaded, path) + require.NoError(t, err) + + env := "resources.jobs." + loaded.ExperimentName + ".environments[0]" + assert.Equal(t, "7", get(t, root, env+".spec.environment_version").MustString()) +} + +// remote_volume can't be honored by a converted bundle (bundle deploy owns the +// artifact upload location), so it is rejected rather than silently ignored. +func TestConvertToDabsRejectsRemoteVolume(t *testing.T) { + cfg := minimalConfig + ` +code_source: + type: snapshot + snapshot: + root_path: ./src + remote_volume: /Volumes/main/default/code +` + path := writeConfigFile(t, "run.yaml", cfg) + loaded, err := loadRunConfig(path) + require.NoError(t, err) + _, _, err = convertToDabs(t.Context(), loaded, path) + require.ErrorContains(t, err, "remote_volume is not supported") +} + +// env_variables / secrets / parameters ride as sidecar files (the ai_runtime_task +// proto has no inline fields for them), matching the CLI's own launch layout. +func TestConvertToDabsStagesEnvAndSecretSidecars(t *testing.T) { + cfg := minimalConfig + ` +env_variables: + FOO: bar +secrets: + TOKEN: scope/key +parameters: + lr: 0.1 +` + path := writeConfigFile(t, "run.yaml", cfg) + loaded, err := loadRunConfig(path) + require.NoError(t, err) + + root, artifacts, err := convertToDabs(t.Context(), loaded, path) + require.NoError(t, err) + + names := itemNames(artifacts) + assert.Contains(t, names, envVarsName) + assert.Contains(t, names, secretEnvVarsName) + assert.Contains(t, names, hyperparametersName) + + // They are NOT smuggled into ai_runtime_task (which would fail bundle validate). + name := loaded.ExperimentName + art := "resources.jobs." + name + ".tasks[0].ai_runtime_task" + assert.False(t, has(root, art+".env_variables")) + assert.False(t, has(root, art+".secrets")) +} + +// conversionNotes surfaces what was transformed/staged so a migrating user knows +// what changed between their run YAML and the bundle. +func TestConvertToDabsConversionNotes(t *testing.T) { + cfg := minimalConfig + ` +env_variables: {FOO: bar} +secrets: {TOKEN: scope/key} +parameters: {lr: 0.1} +code_source: + type: snapshot + snapshot: + root_path: ./src + git: {commit: abc123} +` + path := writeConfigFile(t, "run.yaml", cfg) + loaded, err := loadRunConfig(path) + require.NoError(t, err) + + notes := conversionNotes(loaded) + joined := strings.Join(notes, "\n") + assert.Contains(t, joined, "code_source.git") // git pin flagged + assert.Contains(t, joined, "code_source.tar.gz") // snapshot behavior + assert.Contains(t, joined, "env_vars.json") // env vars staged + assert.Contains(t, joined, "secret_env_vars.json") // secrets staged + assert.Contains(t, joined, "hyperparameters.yaml") // parameters staged + + // A minimal config with none of those has no notes. + base := writeConfigFile(t, "min.yaml", minimalConfig) + minCfg, err := loadRunConfig(base) + require.NoError(t, err) + assert.Empty(t, conversionNotes(minCfg)) +} + +// usage_policy_id is a resolved budget policy id and maps to the job's +// budget_policy_id (usage_policy_name, which needs resolution, is rejected). +func TestConvertToDabsMapsUsagePolicyID(t *testing.T) { + cfg := minimalConfig + "usage_policy_id: budget-abc-123\n" + path := writeConfigFile(t, "run.yaml", cfg) + loaded, err := loadRunConfig(path) + require.NoError(t, err) + + root, _, err := convertToDabs(t.Context(), loaded, path) + require.NoError(t, err) + assert.Equal(t, "budget-abc-123", get(t, root, "resources.jobs."+loaded.ExperimentName+".budget_policy_id").MustString()) +} + +func TestConvertToDabsMapsPermissions(t *testing.T) { + cfg := minimalConfig + ` +permissions: + - user_name: alice@example.com + level: CAN_MANAGE + - group_name: eng + level: CAN_VIEW +` + path := writeConfigFile(t, "run.yaml", cfg) + loaded, err := loadRunConfig(path) + require.NoError(t, err) + + root, _, err := convertToDabs(t.Context(), loaded, path) + require.NoError(t, err) + + perms := get(t, root, "resources.jobs."+loaded.ExperimentName+".permissions").MustSequence() + require.Len(t, perms, 2) + assert.Equal(t, "CAN_MANAGE", perms[0].Get("level").MustString()) + assert.Equal(t, "alice@example.com", perms[0].Get("user_name").MustString()) + assert.Equal(t, "eng", perms[1].Get("group_name").MustString()) +} + +// A git-pinned code_source is archived at the commit, not the dirty working tree: +// writeBundle produces code_source.tar.gz containing the committed file only. +func TestConvertToDabsGitPinnedTarball(t *testing.T) { + repo := newTestRepo(t) + writeRepoFile(t, repo, "train.py", "print('committed')") + sha := commitAll(t, repo, "init") + // Dirty the tree AFTER the commit; the pinned archive must not include this. + writeRepoFile(t, repo, "uncommitted.py", "print('dirty')") + + cfg := "experiment_name: git-pin\ncommand: python train.py\n" + + "compute: {accelerator_type: GPU_1xH100, num_accelerators: 1}\n" + + "code_source:\n type: snapshot\n snapshot:\n root_path: " + repo + "\n git:\n commit: " + sha + "\n" + path := writeConfigFile(t, "run.yaml", cfg) + loaded, err := loadRunConfig(path) + require.NoError(t, err) + + outDir := t.TempDir() + _, err = writeBundle(t.Context(), loaded, path, outDir) + require.NoError(t, err) + + entries := tarballEntries(t, filepath.Join(outDir, codeSourceTarballName)) + dir := filepath.Base(repo) + assert.Contains(t, entries, dir+"/train.py") + assert.NotContains(t, entries, dir+"/uncommitted.py", "git-pinned archive must exclude uncommitted files") +} + +// git archive runs with `git -C repoPath`, so a relative output dir must be +// resolved to absolute before packaging — otherwise `-o out/...` resolves against +// the repo dir and fails. Exercise writeBundle from a working dir with a RELATIVE +// output path against a git-pinned source. +func TestConvertToDabsGitPinnedRelativeOutputDir(t *testing.T) { + repo := newTestRepo(t) + writeRepoFile(t, repo, "train.py", "print('x')") + sha := commitAll(t, repo, "init") + + cfg := "experiment_name: rel-out\ncommand: python train.py\n" + + "compute: {accelerator_type: GPU_1xH100, num_accelerators: 1}\n" + + "code_source:\n type: snapshot\n snapshot:\n root_path: " + repo + "\n git:\n commit: " + sha + "\n" + path := writeConfigFile(t, "run.yaml", cfg) + loaded, err := loadRunConfig(path) + require.NoError(t, err) + + // chdir into a scratch dir and pass a relative --output-dir. + work := t.TempDir() + t.Chdir(work) + _, err = writeBundle(t.Context(), loaded, path, "out") + require.NoError(t, err) + + entries := tarballEntries(t, filepath.Join(work, "out", codeSourceTarballName)) + assert.Contains(t, entries, filepath.Base(repo)+"/train.py") +} + +// A numeric or reserved-word experiment_name would be an unquoted YAML map key +// that DABs' strict loader rejects (!!int / !!bool). The job resource key is +// prefixed to stay a string, while name/experiment keep the original value. +func TestConvertToDabsSafeJobKey(t *testing.T) { + cases := map[string]string{ + "12345": "job_12345", + "1.5e3": "job_1.5e3", + "true": "job_true", + "null": "job_null", + } + for name, wantKey := range cases { + assert.Equal(t, wantKey, bundleResourceKey(name), "key for %q", name) + } + // A normal name is used as-is. + assert.Equal(t, "my-run_1", bundleResourceKey("my-run_1")) + + // End to end: a numeric name lands under the prefixed key, but name/experiment + // keep the numeric string value. + cfg := "experiment_name: \"12345\"\ncommand: python t.py\n" + + "compute: {accelerator_type: GPU_1xH100, num_accelerators: 1}\n" + path := writeConfigFile(t, "run.yaml", cfg) + loaded, err := loadRunConfig(path) + require.NoError(t, err) + root, _, err := convertToDabs(t.Context(), loaded, path) + require.NoError(t, err) + assert.Equal(t, "12345", get(t, root, "resources.jobs.job_12345.name").MustString()) + assert.Equal(t, "12345", get(t, root, "resources.jobs.job_12345.tasks[0].ai_runtime_task.experiment").MustString()) +} + +func TestConvertToDabsRejectsUnsupported(t *testing.T) { + t.Run("docker_image", func(t *testing.T) { + cfg := minimalConfig + ` +environment: + docker_image: + url: myregistry/img:tag +` + path := writeConfigFile(t, "run.yaml", cfg) + loaded, err := loadRunConfig(path) + require.NoError(t, err) + _, _, err = convertToDabs(t.Context(), loaded, path) + require.ErrorContains(t, err, "docker_image is not yet supported") + }) + + t.Run("usage_policy_name", func(t *testing.T) { + cfg := minimalConfig + "usage_policy_name: my-policy\n" + path := writeConfigFile(t, "run.yaml", cfg) + loaded, err := loadRunConfig(path) + require.NoError(t, err) + _, _, err = convertToDabs(t.Context(), loaded, path) + require.ErrorContains(t, err, "usage_policy_name is not yet supported") + }) +} diff --git a/experimental/air/cmd/format.go b/experimental/air/cmd/format.go index c32184517ba..8e46c70d48c 100644 --- a/experimental/air/cmd/format.go +++ b/experimental/air/cmd/format.go @@ -12,6 +12,7 @@ import ( "github.com/databricks/cli/libs/cmdio" "github.com/databricks/databricks-sdk-go/service/jobs" + "github.com/muesli/termenv" "go.yaml.in/yaml/v3" ) @@ -26,12 +27,6 @@ func orNA(s string) string { return s } -// osc8Link wraps label in an OSC 8 terminal hyperlink to url. -// See https://gist.github.com/egmontkob/eb114294efbcd5adb1944c9f3cb5feda -func osc8Link(label, url string) string { - return "\x1b]8;;" + url + "\x1b\\" + label + "\x1b]8;;\x1b\\" -} - // hyperlink renders label as a terminal hyperlink to url when out is a rich // terminal, otherwise it returns label unchanged. This mirrors the Python CLI's // Rich link markup, which drops the URL on non-terminals (so piped or captured @@ -40,7 +35,7 @@ func hyperlink(ctx context.Context, out io.Writer, label, url string) string { if url == "" || !cmdio.SupportsColor(ctx, out) { return label } - return osc8Link(label, url) + return termenv.Hyperlink(url, label) } // reformatYAMLForDisplay re-renders a training-config YAML so multi-line strings diff --git a/experimental/air/cmd/format_test.go b/experimental/air/cmd/format_test.go index 62a6d7ac580..1063e20ca1f 100644 --- a/experimental/air/cmd/format_test.go +++ b/experimental/air/cmd/format_test.go @@ -21,10 +21,6 @@ func TestSubmittedDisplay(t *testing.T) { assert.Equal(t, "2023-11-14 22:13 UTC", submittedDisplay(&jobs.Run{StartTime: 1700000000000})) } -func TestOSC8Link(t *testing.T) { - assert.Equal(t, "\x1b]8;;https://h.test/x\x1b\\label\x1b]8;;\x1b\\", osc8Link("label", "https://h.test/x")) -} - func TestHyperlink(t *testing.T) { // On a non-terminal (no color), the URL is dropped and only the label shows. ctx := cmdio.MockDiscard(t.Context()) diff --git a/experimental/air/cmd/list_tui.go b/experimental/air/cmd/list_tui.go index 85fe70774b4..2596bd4e596 100644 --- a/experimental/air/cmd/list_tui.go +++ b/experimental/air/cmd/list_tui.go @@ -25,7 +25,7 @@ func renderListText(cmd *cobra.Command, f *runFetcher, limit int) error { // "just print these N"). Everything else — piped, NO_COLOR, --limit — prints // once. interactive := color && - cmdio.IsPagerSupported(ctx) && + cmdio.IsPromptSupported(ctx) && !cmd.Flags().Changed("limit") if interactive { diff --git a/experimental/air/cmd/logbricklens.go b/experimental/air/cmd/logbricklens.go new file mode 100644 index 00000000000..48ab7286649 --- /dev/null +++ b/experimental/air/cmd/logbricklens.go @@ -0,0 +1,87 @@ +package aircmd + +import ( + "context" + "encoding/json" + "fmt" + "net/http" + "strconv" + + "github.com/databricks/databricks-sdk-go/client" +) + +// bricklensLogsPathFmt is the log endpoint, keyed by Jobs run id. Called with a +// raw client.Do because the SDK does not model the AiTrainingService. +const bricklensLogsPathFmt = "/api/2.0/ai-training/workflows/by-run-id/%d/logs" + +// logRecord is one log line from Bricklens. +type logRecord struct { + // TimeUnixNano may arrive as a JSON number or string. + TimeUnixNano json.Number `json:"time_unix_nano"` + Body string `json:"body"` + NodeIndex int `json:"node_index"` +} + +// nano returns time_unix_nano as an int64, or 0 when absent or unparseable. +func (r logRecord) nano() int64 { + n, err := r.TimeUnixNano.Int64() + if err != nil { + return 0 + } + return n +} + +type bricklensLogsResponse struct { + LogRecords []logRecord `json:"log_records"` + NextPageToken string `json:"next_page_token"` +} + +// bricklensLogsQuery is the request-field surface of the log endpoint. +type bricklensLogsQuery struct { + // fromSeconds and toSeconds bound the query window in Unix epoch seconds. + fromSeconds int64 + toSeconds int64 + pageToken string + pageSize int + // attemptNumber selects a retry attempt (0-indexed); -1 means latest. + attemptNumber int + nodeIndex int + // ascending returns oldest-first. The endpoint defaults to ascending when + // absent, so the tail fetch must send an explicit false for newest-first. + ascending bool +} + +// getBricklensLogs fetches one page of logs. The API client is built once by the +// caller and reused across the poll loop. It returns the raw error so the caller +// can classify it via classifyLogError. +func getBricklensLogs(ctx context.Context, apiClient *client.DatabricksClient, runID int64, q bricklensLogsQuery) (*bricklensLogsResponse, error) { + query := map[string]any{ + // Always sent: the tail path relies on an explicit false for newest-first. + "ascending": strconv.FormatBool(q.ascending), + } + if q.fromSeconds > 0 { + query["from"] = q.fromSeconds + } + if q.toSeconds > 0 { + query["to"] = q.toSeconds + } + if q.pageToken != "" { + query["page_token"] = q.pageToken + } + if q.pageSize > 0 { + query["page_size"] = q.pageSize + } + if q.attemptNumber >= 0 { + query["ref.attempt_number"] = q.attemptNumber + } + if q.nodeIndex >= 0 { + query["filter.node_index"] = q.nodeIndex + } + + var resp bricklensLogsResponse + path := fmt.Sprintf(bricklensLogsPathFmt, runID) + if err := apiClient.Do(ctx, http.MethodGet, path, nil, nil, query, &resp); err != nil { + return nil, err + } + return &resp, nil +} diff --git a/experimental/air/cmd/logbricklens_test.go b/experimental/air/cmd/logbricklens_test.go new file mode 100644 index 00000000000..e12e0f75a65 --- /dev/null +++ b/experimental/air/cmd/logbricklens_test.go @@ -0,0 +1,83 @@ +package aircmd + +import ( + "net/http" + "net/http/httptest" + "net/url" + "testing" + + "github.com/databricks/databricks-sdk-go/client" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestGetBricklensLogsQuerySerialization(t *testing.T) { + var got url.Values + var gotPath string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/oidc/.well-known/oauth-authorization-server" { + gotPath = r.URL.Path + got = r.URL.Query() + } + _, _ = w.Write([]byte(`{"log_records": [], "next_page_token": ""}`)) + })) + t.Cleanup(srv.Close) + + w := newTestWorkspaceClient(t, srv.URL) + apiClient, err := client.New(w.Config) + require.NoError(t, err) + _, err = getBricklensLogs(t.Context(), apiClient, 42, bricklensLogsQuery{ + fromSeconds: 100, + toSeconds: 200, + pageToken: "tok", + pageSize: 500, + attemptNumber: 1, + nodeIndex: 3, + ascending: true, + }) + require.NoError(t, err) + + assert.Equal(t, "/api/2.0/ai-training/workflows/by-run-id/42/logs", gotPath) + assert.Equal(t, "100", got.Get("from")) + assert.Equal(t, "200", got.Get("to")) + assert.Equal(t, "tok", got.Get("page_token")) + assert.Equal(t, "500", got.Get("page_size")) + assert.Equal(t, "1", got.Get("ref.attempt_number")) + assert.Equal(t, "3", got.Get("filter.node_index")) + assert.Equal(t, "true", got.Get("ascending")) +} + +func TestGetBricklensLogsOmitsOptionals(t *testing.T) { + var got url.Values + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/oidc/.well-known/oauth-authorization-server" { + got = r.URL.Query() + } + _, _ = w.Write([]byte(`{"log_records": []}`)) + })) + t.Cleanup(srv.Close) + + w := newTestWorkspaceClient(t, srv.URL) + apiClient, err := client.New(w.Config) + require.NoError(t, err) + // attempt -1 (latest) and node 0 are the default request; from/to/page unset. + _, err = getBricklensLogs(t.Context(), apiClient, 7, bricklensLogsQuery{attemptNumber: -1, nodeIndex: 0}) + require.NoError(t, err) + + assert.False(t, got.Has("from")) + assert.False(t, got.Has("to")) + assert.False(t, got.Has("page_token")) + assert.False(t, got.Has("page_size")) + // -1 attempt means "latest" — the field is omitted so the endpoint defaults. + assert.False(t, got.Has("ref.attempt_number")) + // node 0 is a real filter value and must be sent. + assert.Equal(t, "0", got.Get("filter.node_index")) + // ascending is always sent so the tail path can force newest-first. + assert.Equal(t, "false", got.Get("ascending")) +} + +func TestLogRecordNano(t *testing.T) { + assert.Equal(t, int64(123), logRecord{TimeUnixNano: "123"}.nano()) + assert.Equal(t, int64(0), logRecord{TimeUnixNano: ""}.nano()) + assert.Equal(t, int64(0), logRecord{TimeUnixNano: "notanumber"}.nano()) +} diff --git a/experimental/air/cmd/logdetect.go b/experimental/air/cmd/logdetect.go new file mode 100644 index 00000000000..f804f4d5812 --- /dev/null +++ b/experimental/air/cmd/logdetect.go @@ -0,0 +1,53 @@ +package aircmd + +import "regexp" + +// fatalPatterns match log lines that signal a run-ending failure (OOM, NCCL +// timeouts, CUDA errors, segfaults, etc.). In --json mode a matching line emits +// an ALERT event alongside its LOG event, giving an agent an immediate signal. +var fatalPatterns = []*regexp.Regexp{ + // OOM + regexp.MustCompile(`(?i)CUDA out of memory`), + regexp.MustCompile(`(?i)Out of memory: Kill(ed)? process`), + // Signals + regexp.MustCompile(`(?i)signal\s+(9|SIGKILL|SIGTERM)`), + // NCCL / collective + regexp.MustCompile(`Watchdog caught collective operation timeout`), + regexp.MustCompile(`(?i)NCCL WARN .*(Conn|Net|IB|timeout|unhandled)`), + regexp.MustCompile(`Got async error event`), + regexp.MustCompile(`transport/net_ib\.cc:\d+.*WARN`), + // CUDA + regexp.MustCompile(`(?i)CUDA(?: runtime)? error`), + regexp.MustCompile(`(?i)an illegal memory access was encountered`), + regexp.MustCompile(`CUDA kernel errors might be asynchronously reported`), + // Segfault + regexp.MustCompile(`(?i)segmentation fault`), + // Composer / llmfoundry specific + regexp.MustCompile(`composer\.utils\..*Error`), + regexp.MustCompile(`(?i)composer.*OutOfMemory`), + // torch.distributed + regexp.MustCompile(`torch\.distributed\..*(?:Error|Exception)`), + regexp.MustCompile(`(?i)(?:TCPStore|Store).*timed?\s*out`), + // GPU hardware (Xid from dmesg / driver) + regexp.MustCompile(`Xid.*\b(48|63|64|79|94|95)\b`), + // Streaming dataset + regexp.MustCompile(`streaming\.base\..*(?:Error|Exception)`), + // Bare "Killed" on its own line means OOM-killer or similar + regexp.MustCompile(`^\s*Killed\s*$`), + // MLflow stall — training has stopped logging metrics + regexp.MustCompile(`\[MLflow Logger\]\[Warning\] No new logs have been emitted`), + // The launch script prints this when the user's command fails; [1-9]\d* skips exit code 0. + regexp.MustCompile(`ERROR: Script failed with exit code [1-9]\d* after \d+s`), + // Missing command (exit 127), e.g. a typo. + regexp.MustCompile(`(?i)command not found`), +} + +// matchFatalPattern reports whether a log line matches a fatal-failure pattern. +func matchFatalPattern(line string) bool { + for _, p := range fatalPatterns { + if p.MatchString(line) { + return true + } + } + return false +} diff --git a/experimental/air/cmd/logmlflow.go b/experimental/air/cmd/logmlflow.go new file mode 100644 index 00000000000..a8dd016dbbc --- /dev/null +++ b/experimental/air/cmd/logmlflow.go @@ -0,0 +1,302 @@ +package aircmd + +import ( + "bufio" + "context" + "fmt" + "io" + "net" + "net/http" + "os" + "path" + "regexp" + "slices" + "strconv" + "time" + + "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/listing" + "github.com/databricks/databricks-sdk-go/service/jobs" + "github.com/databricks/databricks-sdk-go/service/ml" +) + +// chunkFilePattern matches a log chunk file (logs-.chunk.txt); group 1 is +// the chunk index. The sidecar splits stdout into 4MB chunks, index ascending. +var chunkFilePattern = regexp.MustCompile(`^logs-(\d+)\.chunk\.txt$`) + +// oldFormatNodeDir matches a bare per-node log dir (logs/node_). The +// attempt-prefixed layout nests these under logs/attempt_/, so a bare +// logs/node_ only appears in the old layout. +var oldFormatNodeDir = regexp.MustCompile(`^logs/node_\d+$`) + +// artifactDownloadClient fetches pre-signed artifact URLs with connect and +// response-header timeouts, so a stalled storage backend can't hang the command. +// Mirrors the Python CLI's (10s connect, 60s read) bounds. +var artifactDownloadClient = &http.Client{ + Transport: &http.Transport{ + DialContext: (&net.Dialer{Timeout: 10 * time.Second}).DialContext, + ResponseHeaderTimeout: 60 * time.Second, + }, +} + +// mlflowLogFallback prints a run's logs from MLflow artifacts, the fallback when +// Bricklens can't serve them. It resolves the MLflow run id, discovers the +// per-node log directory, lists the chunk files, and walks them newest-first +// until it has the requested tail, then prints oldest-first. +// +// The tail length is --lines, else the default cap. MLflow chunks are not +// time-indexed, so --minutes cannot restrict the window here. +func mlflowLogFallback(ctx context.Context, w *databricks.WorkspaceClient, out io.Writer, req logRequest, status logRunStatus) (bool, error) { + if req.windowMinutes > 0 { + log.Debugf(ctx, "air logs: --minutes is not supported on the MLflow fallback path; showing the default tail") + } + + mlflowRunID, logDir, err := resolveMLflowLogPath(ctx, w, req) + if err != nil { + return false, err + } + if mlflowRunID == "" || logDir == "" { + emitNoLogs(out, req, status) + return status.succeeded(), nil + } + + chunks, err := listLogChunks(ctx, w, mlflowRunID, logDir) + if err != nil { + return false, err + } + if len(chunks) == 0 { + // Nothing listed yet: assume the single chunk 0. + chunks = []logChunk{{index: 0, path: path.Join(logDir, chunkFileName(0))}} + } + + target := req.tailTarget() + if target <= 0 { + return status.succeeded(), nil + } + + lines, err := tailChunks(ctx, w, mlflowRunID, chunks, target) + if err != nil { + return false, err + } + if len(lines) == 0 { + emitNoLogs(out, req, status) + return status.succeeded(), nil + } + + if len(lines) > target { + lines = lines[len(lines)-target:] + } + for _, line := range lines { + emitLogLine(out, req, line) + } + return status.succeeded(), nil +} + +// resolveMLflowLogPath returns the run's MLflow run id and per-node log directory. +func resolveMLflowLogPath(ctx context.Context, w *databricks.WorkspaceClient, req logRequest) (string, string, error) { + run, err := w.Jobs.GetRun(ctx, jobs.GetRunRequest{RunId: req.runID}) + if err != nil { + return "", "", err + } + ids := mlflowIDs(ctx, w, run) + if ids == nil || ids.RunID == "" { + return "", "", nil + } + + // -1 (latest) maps to attempt 0's directory. + attempt := max(req.attempt, 0) + withAttempt, err := discoverAttemptPrefix(ctx, w, ids.RunID, attempt) + if err != nil { + return "", "", err + } + return ids.RunID, constructLogPath(req.node, attempt, withAttempt), nil +} + +// discoverAttemptPrefix probes the logs/ dir once to decide whether the layout is +// attempt-prefixed (logs/attempt_X/node_Y) or old (logs/node_Y). A bare +// logs/node_ means old; a logs/attempt_ entry means prefixed. +// Defaults to old when nothing is listed. +func discoverAttemptPrefix(ctx context.Context, w *databricks.WorkspaceClient, mlflowRunID string, attempt int) (bool, error) { + files, err := listArtifacts(ctx, w, mlflowRunID, "logs") + if err != nil { + // Not fatal: default to the old layout; the chunk listing finds it empty if wrong. + log.Debugf(ctx, "air logs: could not list logs dir for format discovery: %v", err) + return false, nil + } + + attemptDir := fmt.Sprintf("logs/attempt_%d", attempt) + for _, f := range files { + if oldFormatNodeDir.MatchString(f.Path) { + return false, nil + } + if f.Path == attemptDir { + return true, nil + } + } + return false, nil +} + +// constructLogPath builds the per-node log directory for a node and attempt. +func constructLogPath(node, attempt int, withAttempt bool) string { + if withAttempt { + return fmt.Sprintf("logs/attempt_%d/node_%d", attempt, node) + } + return fmt.Sprintf("logs/node_%d", node) +} + +// chunkFileName is the artifact filename for a chunk index. +func chunkFileName(index int) string { + return fmt.Sprintf("logs-%d.chunk.txt", index) +} + +// logChunk is one listed chunk: its index and full artifact path. +type logChunk struct { + index int + path string +} + +// listLogChunks lists the chunk files under a log dir, sorted ascending by index. +func listLogChunks(ctx context.Context, w *databricks.WorkspaceClient, mlflowRunID, logDir string) ([]logChunk, error) { + files, err := listArtifacts(ctx, w, mlflowRunID, logDir) + if err != nil { + return nil, err + } + + var chunks []logChunk + for _, f := range files { + base := path.Base(f.Path) + m := chunkFilePattern.FindStringSubmatch(base) + if m == nil { + continue + } + idx, err := strconv.Atoi(m[1]) + if err != nil { + continue + } + chunks = append(chunks, logChunk{index: idx, path: f.Path}) + } + slices.SortFunc(chunks, func(a, b logChunk) int { return a.index - b.index }) + return chunks, nil +} + +// tailChunks walks chunks newest-first, prepending each chunk's lines, until it +// has `target` lines or runs out. A mid-walk download failure stops the walk +// rather than splice non-adjacent chunks. +func tailChunks(ctx context.Context, w *databricks.WorkspaceClient, mlflowRunID string, chunks []logChunk, target int) ([]string, error) { + var accumulated []string + for _, chunk := range slices.Backward(chunks) { + lines, err := downloadChunkLines(ctx, w, mlflowRunID, chunk.path) + if err != nil { + log.Debugf(ctx, "air logs: failed to download chunk %d; showing only logs after it: %v", chunk.index, err) + break + } + accumulated = append(lines, accumulated...) + if len(accumulated) >= target { + break + } + } + return accumulated, nil +} + +// downloadChunkLines fetches one chunk artifact and returns its lines. +func downloadChunkLines(ctx context.Context, w *databricks.WorkspaceClient, mlflowRunID, artifactPath string) ([]string, error) { + f, err := downloadArtifact(ctx, w, mlflowRunID, artifactPath) + if err != nil { + return nil, err + } + defer os.Remove(f) + + file, err := os.Open(f) + if err != nil { + return nil, err + } + defer file.Close() + + var lines []string + scanner := bufio.NewScanner(file) + scanner.Buffer(make([]byte, 0, 64*1024), 4*1024*1024) + for scanner.Scan() { + lines = append(lines, scanner.Text()) + } + return lines, scanner.Err() +} + +// listArtifacts lists a run's artifacts under a path. +func listArtifacts(ctx context.Context, w *databricks.WorkspaceClient, mlflowRunID, path string) ([]ml.FileInfo, error) { + it := w.Experiments.ListArtifacts(ctx, ml.ListArtifactsRequest{RunId: mlflowRunID, Path: path}) + return listing.ToSlice(ctx, it) +} + +// credentialInfo is one credentials-for-read entry: a pre-signed URL plus any +// backend-required request headers. +type credentialInfo struct { + SignedURI string `json:"signed_uri"` + Headers []struct { + Name string `json:"name"` + Value string `json:"value"` + } `json:"headers"` +} + +type credentialsForReadResponse struct { + CredentialInfos []credentialInfo `json:"credential_infos"` +} + +// downloadArtifact downloads one run artifact to a temp file and returns its +// path. credentials-for-read returns a pre-signed URL, which we stream to disk; +// that endpoint is not modeled by the SDK, so it is called via a raw client.Do. +func downloadArtifact(ctx context.Context, w *databricks.WorkspaceClient, mlflowRunID, artifactPath string) (string, error) { + apiClient, err := client.New(w.Config) + if err != nil { + return "", fmt.Errorf("failed to create API client: %w", err) + } + + var resp credentialsForReadResponse + query := map[string]any{ + "run_id": mlflowRunID, + // path is a repeated field, so pass a slice (serialized as path=...&path=...). + "path": []string{artifactPath}, + } + err = apiClient.Do(ctx, http.MethodGet, "/api/2.0/mlflow/artifacts/credentials-for-read", nil, nil, query, &resp) + if err != nil { + return "", fmt.Errorf("failed to get read credentials for %s: %w", artifactPath, err) + } + if len(resp.CredentialInfos) == 0 || resp.CredentialInfos[0].SignedURI == "" { + return "", fmt.Errorf("no download credentials returned for %s", artifactPath) + } + cred := resp.CredentialInfos[0] + + req, err := http.NewRequestWithContext(ctx, http.MethodGet, cred.SignedURI, nil) + if err != nil { + return "", err + } + // Azure SAS / some GCS URIs require backend-supplied headers; AWS returns none. + for _, h := range cred.Headers { + req.Header.Set(h.Name, h.Value) + } + + httpResp, err := artifactDownloadClient.Do(req) + if err != nil { + return "", err + } + defer httpResp.Body.Close() + if httpResp.StatusCode >= 400 { + return "", fmt.Errorf("artifact download failed: HTTP %d", httpResp.StatusCode) + } + + tmp, err := os.CreateTemp("", "air-log-chunk-*") + if err != nil { + return "", err + } + if _, err := io.Copy(tmp, httpResp.Body); err != nil { + tmp.Close() + os.Remove(tmp.Name()) + return "", err + } + if err := tmp.Close(); err != nil { + os.Remove(tmp.Name()) + return "", err + } + return tmp.Name(), nil +} diff --git a/experimental/air/cmd/logmlflow_test.go b/experimental/air/cmd/logmlflow_test.go new file mode 100644 index 00000000000..f3de80abdb6 --- /dev/null +++ b/experimental/air/cmd/logmlflow_test.go @@ -0,0 +1,114 @@ +package aircmd + +import ( + "bytes" + "net/http" + "net/http/httptest" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestConstructLogPath(t *testing.T) { + assert.Equal(t, "logs/node_0", constructLogPath(0, 0, false)) + assert.Equal(t, "logs/node_3", constructLogPath(3, 2, false)) + assert.Equal(t, "logs/attempt_2/node_3", constructLogPath(3, 2, true)) +} + +func TestChunkFileName(t *testing.T) { + assert.Equal(t, "logs-0.chunk.txt", chunkFileName(0)) + assert.Equal(t, "logs-7.chunk.txt", chunkFileName(7)) +} + +func TestChunkFilePattern(t *testing.T) { + m := chunkFilePattern.FindStringSubmatch("logs-12.chunk.txt") + require.NotNil(t, m) + assert.Equal(t, "12", m[1]) + + assert.Nil(t, chunkFilePattern.FindStringSubmatch("logs-12.chunk.txt.bak")) + assert.Nil(t, chunkFilePattern.FindStringSubmatch("node_0")) +} + +// artifactListServer serves a fixed artifacts/list response for any path. +func artifactListServer(t *testing.T, body string) *httptest.Server { + t.Helper() + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == "/api/2.0/mlflow/artifacts/list" { + _, _ = w.Write([]byte(body)) + return + } + _, _ = w.Write([]byte(`{}`)) + })) + t.Cleanup(srv.Close) + return srv +} + +func TestListLogChunksSortsAndFiltersByIndex(t *testing.T) { + // Out-of-order chunks plus a non-chunk file; result is ascending, chunk-only. + srv := artifactListServer(t, `{"files": [ + {"path": "logs/node_0/logs-2.chunk.txt"}, + {"path": "logs/node_0/other.txt"}, + {"path": "logs/node_0/logs-0.chunk.txt"}, + {"path": "logs/node_0/logs-1.chunk.txt"} + ]}`) + w := newTestWorkspaceClient(t, srv.URL) + + chunks, err := listLogChunks(t.Context(), w, "run1", "logs/node_0") + require.NoError(t, err) + require.Len(t, chunks, 3) + assert.Equal(t, 0, chunks[0].index) + assert.Equal(t, 1, chunks[1].index) + assert.Equal(t, 2, chunks[2].index) + assert.Equal(t, "logs/node_0/logs-0.chunk.txt", chunks[0].path) +} + +func TestDiscoverAttemptPrefix(t *testing.T) { + // Old format: a bare logs/node_N dir means no attempt prefix. + old := artifactListServer(t, `{"files": [{"path": "logs/node_0", "is_dir": true}]}`) + got, err := discoverAttemptPrefix(t.Context(), newTestWorkspaceClient(t, old.URL), "run1", 0) + require.NoError(t, err) + assert.False(t, got) + + // New format: a logs/attempt_N entry and no bare node dir. + newFmt := artifactListServer(t, `{"files": [{"path": "logs/attempt_0", "is_dir": true}]}`) + got, err = discoverAttemptPrefix(t.Context(), newTestWorkspaceClient(t, newFmt.URL), "run1", 0) + require.NoError(t, err) + assert.True(t, got) +} + +// noMLflowServer serves a run with no resolvable MLflow run id, so the fallback +// finds no logs. +func noMLflowServer(t *testing.T) *httptest.Server { + t.Helper() + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/api/2.2/jobs/runs/get": + _, _ = w.Write([]byte(`{"run_id": 5, "state": {"life_cycle_state": "TERMINATED", "result_state": "SUCCESS"}, "tasks": [{"run_id": 456}]}`)) + case "/api/2.2/jobs/runs/get-output": + _, _ = w.Write([]byte(`{}`)) + default: + _, _ = w.Write([]byte(`{}`)) + } + })) + t.Cleanup(srv.Close) + return srv +} + +func TestMLflowFallbackNoLogsReflectsRunOutcome(t *testing.T) { + srv := noMLflowServer(t) + w := newTestWorkspaceClient(t, srv.URL) + + // A SUCCESS run with no logs still reports success (exit 0), matching the + // Bricklens path, rather than failing just because no logs exist. + success, err := mlflowLogFallback(t.Context(), w, &bytes.Buffer{}, + logRequest{runID: 5}, logRunStatus{lifeCycleState: "TERMINATED", resultState: "SUCCESS"}) + require.NoError(t, err) + assert.True(t, success) + + // A FAILED run with no logs reports failure (exit 1). + success, err = mlflowLogFallback(t.Context(), w, &bytes.Buffer{}, + logRequest{runID: 5}, logRunStatus{lifeCycleState: "TERMINATED", resultState: "FAILED"}) + require.NoError(t, err) + assert.False(t, success) +} diff --git a/experimental/air/cmd/logs.go b/experimental/air/cmd/logs.go index c34fb62a7df..1bda80b572b 100644 --- a/experimental/air/cmd/logs.go +++ b/experimental/air/cmd/logs.go @@ -1,7 +1,18 @@ package aircmd import ( + "context" + "errors" + "fmt" + "io" + "strconv" + "github.com/databricks/cli/cmd/root" + "github.com/databricks/cli/libs/cmdctx" + "github.com/databricks/cli/libs/flags" + "github.com/databricks/databricks-sdk-go" + "github.com/databricks/databricks-sdk-go/apierr" + "github.com/databricks/databricks-sdk-go/service/iam" "github.com/spf13/cobra" ) @@ -9,6 +20,7 @@ func newLogsCommand() *cobra.Command { var ( node int lines int + minutes int retry int downloadTo string review bool @@ -19,18 +31,142 @@ func newLogsCommand() *cobra.Command { Args: root.ExactArgs(1), Short: "Stream or fetch logs for a run", Long: `Stream logs from an active run, or fetch logs from a completed run.`, - RunE: func(cmd *cobra.Command, args []string) error { - return notImplemented("logs") - }, } cmd.Flags().IntVar(&node, "node", 0, "Fetch logs from this node") - cmd.Flags().IntVar(&lines, "lines", 10000, "For completed runs, print the last N lines") + cmd.Flags().IntVar(&lines, "lines", 0, "For completed runs, print the last N lines (default 10000)") + cmd.Flags().IntVar(&minutes, "minutes", 0, "Fetch only logs from the last N minutes") cmd.Flags().IntVar(&retry, "retry", -1, "View logs from a specific retry attempt; -1 means latest") cmd.Flags().StringVar(&downloadTo, "download-to", "", "Download all logs to this directory instead of printing") cmd.Flags().BoolVar(&review, "review", false, "Download logs from all nodes and filter for error signatures") - // Hidden in the Python `air` CLI (help=argparse.SUPPRESS); keep it internal here to match. cmd.Flags().MarkHidden("review") + // In -o json mode an auth failure should be a JSON error envelope, not a bare + // error. ErrAlreadyPrinted passes through. + cmd.PreRunE = func(cmd *cobra.Command, args []string) error { + err := root.MustWorkspaceClient(cmd, args) + if err == nil || errors.Is(err, root.ErrAlreadyPrinted) { + return err + } + return authError(cmd.Context(), cmd, err) + } + + cmd.RunE = func(cmd *cobra.Command, args []string) error { + ctx := cmd.Context() + + // --download-to and --review are not yet implemented; reject rather than + // silently ignore. + if downloadTo != "" { + return renderError(ctx, cmd, "INVALID_ARGS", "PERMANENT", false, + errors.New("--download-to is not implemented yet")) + } + if review { + return renderError(ctx, cmd, "INVALID_ARGS", "PERMANENT", false, + errors.New("--review is not implemented yet")) + } + + // --lines (line tail) and --minutes (time window) answer the same question + // two ways, so reject both together rather than silently honoring one. + if lines > 0 && minutes > 0 { + return renderError(ctx, cmd, "INVALID_ARGS", "PERMANENT", false, + errors.New("cannot combine --lines with --minutes: --lines tails by line count, --minutes by time window")) + } + if lines < 0 { + return renderError(ctx, cmd, "INVALID_ARGS", "PERMANENT", false, + fmt.Errorf("invalid --lines %d: must be positive", lines)) + } + if minutes < 0 { + return renderError(ctx, cmd, "INVALID_ARGS", "PERMANENT", false, + fmt.Errorf("invalid --minutes %d: must be positive", minutes)) + } + if node < 0 { + return renderError(ctx, cmd, "INVALID_ARGS", "PERMANENT", false, + fmt.Errorf("invalid --node %d: must not be negative", node)) + } + + runID, err := strconv.ParseInt(args[0], 10, 64) + if err != nil || runID <= 0 { + return renderError(ctx, cmd, "INVALID_ARGS", "PERMANENT", false, + fmt.Errorf("invalid JOB_RUN_ID %q: must be a positive integer", args[0])) + } + + // -1 signals "unset" (use the default cap); an explicit --lines 0 stays 0 + // and prints nothing. + tailLines := -1 + if cmd.Flags().Changed("lines") { + tailLines = lines + } + + return runLogs(ctx, cmd, logRequest{ + runID: runID, + node: node, + attempt: retry, + windowMinutes: minutes, + tailLines: tailLines, + jsonOutput: root.OutputType(cmd) == flags.OutputJSON, + }) + } + return cmd } + +// runLogs resolves the run, validates --retry, and fetches logs. It handles error +// reporting; the backend selection lives in fetchLogs. +func runLogs(ctx context.Context, cmd *cobra.Command, req logRequest) error { + w := cmdctx.WorkspaceClient(ctx) + + // Validate credentials server-side before fetching (MustWorkspaceClient only + // attaches them), so a bad token fails clearly here. + if _, err := w.CurrentUser.Me(ctx, iam.MeRequest{}); err != nil { + return authError(ctx, cmd, err) + } + + status, err := resolveRunStatus(ctx, w, req.runID) + if err != nil { + if errors.Is(err, apierr.ErrResourceDoesNotExist) { + return renderError(ctx, cmd, "NOT_FOUND", "NOT_FOUND", false, + fmt.Errorf("run %d not found: check the run ID and that it is a job run ID", req.runID)) + } + return renderError(ctx, cmd, "INTERNAL_ERROR", "TRANSIENT", true, + fmt.Errorf("failed to get status for run %d: %w", req.runID, err)) + } + + // -1 (default) means latest; reject an attempt past the newest. + if req.attempt >= 0 && req.attempt > status.latestAttempt { + return renderError(ctx, cmd, "INVALID_ARGS", "PERMANENT", false, + fmt.Errorf("invalid retry %d: available retries are 0 to %d", req.attempt, status.latestAttempt)) + } + + // A past retry of an active run has immutable logs: render once, don't follow. + if req.attempt >= 0 && req.attempt < status.latestAttempt && !status.terminal() { + req.staticView = true + } + + out := cmd.OutOrStdout() + success, err := fetchLogs(ctx, w, out, req, status) + if err != nil { + if errors.Is(err, apierr.ErrResourceDoesNotExist) { + return renderError(ctx, cmd, "NOT_FOUND", "NOT_FOUND", false, + fmt.Errorf("run %d not found: check the run ID and that it is a job run ID", req.runID)) + } + return renderError(ctx, cmd, "INTERNAL_ERROR", "TRANSIENT", true, + fmt.Errorf("failed to fetch logs for run %d: %w", req.runID, err)) + } + + // A run that finished unsuccessfully exits non-zero; output was already + // written, so don't reprint via Cobra. + if !success { + return root.ErrAlreadyPrinted + } + return nil +} + +// fetchLogs serves logs from Bricklens, falling back to MLflow when Bricklens +// returns errBricklensFeatureDisabled. +func fetchLogs(ctx context.Context, w *databricks.WorkspaceClient, out io.Writer, req logRequest, status logRunStatus) (bool, error) { + success, err := streamBricklensLogs(ctx, w, out, req, status) + if errors.Is(err, errBricklensFeatureDisabled) { + return mlflowLogFallback(ctx, w, out, req, status) + } + return success, err +} diff --git a/experimental/air/cmd/logs_test.go b/experimental/air/cmd/logs_test.go new file mode 100644 index 00000000000..a9eb24b72f3 --- /dev/null +++ b/experimental/air/cmd/logs_test.go @@ -0,0 +1,252 @@ +package aircmd + +import ( + "bytes" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "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/experimental/mocks" + "github.com/spf13/cobra" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestLogsCommandShape(t *testing.T) { + cmd := newLogsCommand() + assert.Equal(t, "logs JOB_RUN_ID", cmd.Use) + assert.Empty(t, cmd.Commands(), "logs must not register subcommands") + assert.NoError(t, cmd.Args(cmd, []string{"123"})) + assert.Error(t, cmd.Args(cmd, []string{})) + assert.Error(t, cmd.Args(cmd, []string{"1", "2"})) + + // --review is hidden. + review := cmd.Flags().Lookup("review") + require.NotNil(t, review) + assert.True(t, review.Hidden) +} + +// runLogsCmd invokes the logs command's RunE with the given flags against a mock +// (no-HTTP) workspace client. Used for input validation that fails before any +// API call. +func runLogsCmd(t *testing.T, args []string, flagsToSet map[string]string) error { + t.Helper() + m := mocks.NewMockWorkspaceClient(t) + ctx := cmdctx.SetWorkspaceClient(cmdio.MockDiscard(t.Context()), m.WorkspaceClient) + cmd := withOutput(newLogsCommand(), flags.OutputText) + for k, v := range flagsToSet { + require.NoError(t, cmd.Flags().Set(k, v)) + } + cmd.SetContext(ctx) + return cmd.RunE(cmd, args) +} + +func TestLogsFlagValidation(t *testing.T) { + tests := []struct { + name string + args []string + flags map[string]string + wantMsg string + }{ + { + name: "lines and minutes are mutually exclusive", + args: []string{"5"}, + flags: map[string]string{"lines": "100", "minutes": "10"}, + wantMsg: "cannot combine --lines with --minutes", + }, + { + name: "negative lines rejected", + args: []string{"5"}, + flags: map[string]string{"lines": "-1"}, + wantMsg: "invalid --lines", + }, + { + name: "negative minutes rejected", + args: []string{"5"}, + flags: map[string]string{"minutes": "-1"}, + wantMsg: "invalid --minutes", + }, + { + name: "download-to not implemented", + args: []string{"5"}, + flags: map[string]string{"download-to": "/tmp/logs"}, + wantMsg: "--download-to is not implemented yet", + }, + { + name: "review not implemented", + args: []string{"5"}, + flags: map[string]string{"review": "true"}, + wantMsg: "--review is not implemented yet", + }, + { + name: "invalid run id", + args: []string{"abc"}, + wantMsg: "invalid JOB_RUN_ID", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := runLogsCmd(t, tt.args, tt.flags) + require.Error(t, err) + assert.Contains(t, err.Error(), tt.wantMsg) + }) + } +} + +// completedRunLogsServer serves the auth probe, a terminal runs/get, and a +// single page of Bricklens logs (newest-first, as the tail fetch requests). +func completedRunLogsServer(t *testing.T) *httptest.Server { + t.Helper() + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch { + case r.URL.Path == "/api/2.2/jobs/runs/get": + _, _ = w.Write([]byte(`{ + "run_id": 5, + "start_time": 1000, + "end_time": 2000, + "state": {"life_cycle_state": "TERMINATED", "result_state": "SUCCESS"}, + "tasks": [{"attempt_number": 0}] + }`)) + case strings.HasPrefix(r.URL.Path, "/api/2.0/ai-training/workflows/by-run-id/"): + // Newest-first, as drainTail requests; reversed to oldest-first on print. + _, _ = w.Write([]byte(`{"log_records": [ + {"time_unix_nano": 2000000000, "body": "line two", "node_index": 0}, + {"time_unix_nano": 1000000000, "body": "line one", "node_index": 0} + ]}`)) + default: + // Me() probe and SDK config discovery. + _, _ = w.Write([]byte(`{"userName":"u@example.com"}`)) + } + })) + t.Cleanup(srv.Close) + return srv +} + +func TestLogsCompletedRunTail(t *testing.T) { + srv := completedRunLogsServer(t) + var buf bytes.Buffer + ctx := cmdctx.SetWorkspaceClient(cmdio.MockDiscard(t.Context()), newTestWorkspaceClient(t, srv.URL)) + cmd := withOutput(&cobra.Command{}, flags.OutputText) + cmd.SetContext(ctx) + cmd.SetOut(&buf) + + // Drive runLogs directly (bypassing PreRunE auth wiring) with a resolved request. + err := runLogs(ctx, cmd, logRequest{runID: 5, node: 0, attempt: -1, tailLines: -1}) + require.NoError(t, err) + + // Records print oldest-first regardless of the newest-first fetch order. + assert.Equal(t, "line one\nline two\n", buf.String()) +} + +// mlflowFallbackServer serves a terminal run whose Bricklens endpoint is gated +// off (FEATURE_DISABLED), plus the full MLflow artifact path the fallback walks: +// runs/get-output (MLflow ids), artifacts/list (logs dir + chunk file), +// credentials-for-read (pre-signed URL), and the pre-signed chunk bytes itself. +func mlflowFallbackServer(t *testing.T) *httptest.Server { + t.Helper() + var base string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch { + case r.URL.Path == "/api/2.2/jobs/runs/get": + _, _ = w.Write([]byte(`{ + "run_id": 5, + "state": {"life_cycle_state": "TERMINATED", "result_state": "SUCCESS"}, + "tasks": [{"run_id": 456, "attempt_number": 0}] + }`)) + case strings.HasPrefix(r.URL.Path, "/api/2.0/ai-training/workflows/by-run-id/"): + w.WriteHeader(http.StatusForbidden) + _, _ = w.Write([]byte(`{"error_code": "FEATURE_DISABLED", "message": "bricklens logs gated off"}`)) + case r.URL.Path == "/api/2.2/jobs/runs/get-output": + _, _ = w.Write([]byte(`{"ai_runtime_task_output": {"mlflow_experiment_id": "exp1", "mlflow_run_id": "run1"}}`)) + case r.URL.Path == "/api/2.0/mlflow/artifacts/list": + // The logs dir probe (format discovery) and the per-node chunk listing + // both hit this; return the old-format node dir and one chunk file. + if r.URL.Query().Get("path") == "logs" { + _, _ = w.Write([]byte(`{"files": [{"path": "logs/node_0", "is_dir": true}]}`)) + return + } + _, _ = w.Write([]byte(`{"files": [{"path": "logs/node_0/logs-0.chunk.txt", "file_size": 12}]}`)) + case r.URL.Path == "/api/2.0/mlflow/artifacts/credentials-for-read": + _, _ = w.Write([]byte(`{"credential_infos": [{"signed_uri": "` + base + `/presigned"}]}`)) + case r.URL.Path == "/presigned": + _, _ = w.Write([]byte("line one\nline two\n")) + default: + _, _ = w.Write([]byte(`{"userName":"u@example.com"}`)) + } + })) + base = srv.URL + t.Cleanup(srv.Close) + return srv +} + +func TestLogsFallsBackToMLflow(t *testing.T) { + srv := mlflowFallbackServer(t) + var buf bytes.Buffer + ctx := cmdctx.SetWorkspaceClient(cmdio.MockDiscard(t.Context()), newTestWorkspaceClient(t, srv.URL)) + cmd := withOutput(&cobra.Command{}, flags.OutputText) + cmd.SetContext(ctx) + cmd.SetOut(&buf) + + // Bricklens is gated off, so fetchLogs routes to the MLflow fallback, which + // resolves the MLflow run, lists the chunk, downloads it via the pre-signed + // URL, and prints its lines. + err := runLogs(ctx, cmd, logRequest{runID: 5, node: 0, attempt: -1, tailLines: -1}) + require.NoError(t, err) + assert.Equal(t, "line one\nline two\n", buf.String()) +} + +// activeRunPastRetryServer serves a still-RUNNING run with two attempts and a +// single page of Bricklens logs. runs/get always returns RUNNING; a test that +// follows the run would poll forever, so it also asserts the static path never +// loops. +func activeRunPastRetryServer(t *testing.T, getRunHits *int) *httptest.Server { + t.Helper() + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch { + case r.URL.Path == "/api/2.2/jobs/runs/get": + *getRunHits++ + _, _ = w.Write([]byte(`{ + "run_id": 9, + "start_time": 1700000000000, + "state": {"life_cycle_state": "RUNNING"}, + "tasks": [{"attempt_number": 0}, {"attempt_number": 1}] + }`)) + case strings.HasPrefix(r.URL.Path, "/api/2.0/ai-training/workflows/by-run-id/"): + _, _ = w.Write([]byte(`{"log_records": [ + {"time_unix_nano": 1700000001000000000, "body": "retry 0 log", "node_index": 0} + ]}`)) + default: + _, _ = w.Write([]byte(`{"userName":"u@example.com"}`)) + } + })) + t.Cleanup(srv.Close) + return srv +} + +func TestLogsPastRetryOfActiveRunIsStatic(t *testing.T) { + var getRunHits int + srv := activeRunPastRetryServer(t, &getRunHits) + var buf bytes.Buffer + ctx := cmdctx.SetWorkspaceClient(cmdio.MockDiscard(t.Context()), newTestWorkspaceClient(t, srv.URL)) + cmd := withOutput(&cobra.Command{}, flags.OutputText) + cmd.SetContext(ctx) + cmd.SetOut(&buf) + + // --retry 0 on a RUNNING run whose latest attempt is 1: the past attempt's + // logs render once instead of following the run (which would never terminate). + // The run has no SUCCESS result yet, so it still exits non-zero via + // ErrAlreadyPrinted; the logs are printed regardless. + err := runLogs(ctx, cmd, logRequest{runID: 9, node: 0, attempt: 0, tailLines: -1}) + require.ErrorIs(t, err, root.ErrAlreadyPrinted) + assert.Equal(t, "retry 0 log\n", buf.String()) + + // Exactly one runs/get: the initial status resolve in runLogs. The static + // streamer must not re-poll (that is the loop this test guards against). + assert.Equal(t, 1, getRunHits) +} diff --git a/experimental/air/cmd/logstream.go b/experimental/air/cmd/logstream.go new file mode 100644 index 00000000000..e5ab7ba8e85 --- /dev/null +++ b/experimental/air/cmd/logstream.go @@ -0,0 +1,483 @@ +package aircmd + +import ( + "context" + "errors" + "fmt" + "io" + "net/http" + "slices" + "time" + + "github.com/databricks/cli/libs/cmdio" + "github.com/databricks/cli/libs/log" + "github.com/databricks/databricks-sdk-go" + "github.com/databricks/databricks-sdk-go/apierr" + "github.com/databricks/databricks-sdk-go/client" + "github.com/databricks/databricks-sdk-go/service/jobs" +) + +const ( + // maxTransientFailures is how many consecutive Bricklens failures to tolerate + // before falling back to MLflow. + maxTransientFailures = 5 + // defaultCompletedRunTailLines caps a completed run's output when neither + // --lines nor --minutes is set. + defaultCompletedRunTailLines = 10000 + // seenRecordsCap bounds the dedup set, evicting oldest-inserted entries first. + seenRecordsCap = 100000 +) + +// retryCheckInterval is the wait between status/log polls. A var so tests can +// shrink it. +var retryCheckInterval = 3 * time.Second + +// errBricklensFeatureDisabled signals the caller to fall back to MLflow: Bricklens +// is gated off (FEATURE_DISABLED), not deployed (ENDPOINT_NOT_FOUND / 404), or +// persistently failing. The flag is evaluated server-side. +var errBricklensFeatureDisabled = errors.New("bricklens logs unavailable; falling back to mlflow") + +// logRequest describes what to fetch, shared by both backends so they honor the +// same flags. windowMinutes and tailLines are mutually exclusive. +type logRequest struct { + runID int64 + // node is the node index to fetch; node 0 always exists. + node int + // attempt is the retry attempt to read; -1 means latest. + attempt int + // windowMinutes, when > 0, restricts the fetch to the last N minutes. + windowMinutes int + // tailLines caps a completed run's output to the last N lines. Negative means + // --lines was unset (use the default cap); 0 prints nothing. + tailLines int + // staticView renders a one-shot tail instead of following the run. Set for a + // past retry of an active run: that attempt's logs are immutable, so streaming + // would poll forever waiting for the run (not the attempt) to finish. + staticView bool + jsonOutput bool + // onStatusChange, when set, is called on each lifecycle transition while + // following the run (current, previous display states). Used by + // `air run --watch -o json` to emit STATUS events. + onStatusChange func(current, previous string) +} + +// logRunStatus is the subset of a run's state the log path needs, resolved once +// and reused. +type logRunStatus struct { + lifeCycleState string + resultState string + stateMessage string + startTimeMs int64 + endTimeMs int64 + // latestAttempt is the highest attempt_number across the run's tasks. + latestAttempt int +} + +// A run is terminal when its lifecycle state is terminal, or a result state is +// set (result states only appear on terminal runs). +var ( + terminalLifeCycleStates = map[string]bool{"TERMINATED": true, "SKIPPED": true, "INTERNAL_ERROR": true} + terminalResultStates = map[string]bool{"SUCCESS": true, "FAILED": true, "CANCELED": true} +) + +func (s logRunStatus) terminal() bool { + return terminalLifeCycleStates[s.lifeCycleState] || terminalResultStates[s.resultState] +} + +func (s logRunStatus) succeeded() bool { + return s.resultState == "SUCCESS" +} + +// resolveRunStatus fetches a run's state and projects it onto logRunStatus. An +// unknown run id surfaces as apierr.ErrResourceDoesNotExist. +func resolveRunStatus(ctx context.Context, w *databricks.WorkspaceClient, runID int64) (logRunStatus, error) { + run, err := w.Jobs.GetRun(ctx, jobs.GetRunRequest{RunId: runID}) + if err != nil { + return logRunStatus{}, err + } + return projectRunStatus(run), nil +} + +// projectRunStatus extracts logRunStatus from a run. Split out so it can be +// tested without an API client. +func projectRunStatus(run *jobs.Run) logRunStatus { + s := logRunStatus{ + startTimeMs: run.StartTime, + endTimeMs: run.EndTime, + } + if run.State != nil { + s.lifeCycleState = string(run.State.LifeCycleState) + s.resultState = string(run.State.ResultState) + s.stateMessage = run.State.StateMessage + } + for i := range run.Tasks { + s.latestAttempt = max(s.latestAttempt, run.Tasks[i].AttemptNumber) + } + return s +} + +// classifyLogError maps a Bricklens failure to one of: +// - errBricklensFeatureDisabled: fall back to MLflow (gated off, endpoint +// absent, or 404). +// - the original error: a genuine not-found, surfaced as-is. +// - nil: a transient failure the caller should retry. +func classifyLogError(err error) error { + if apiErr, ok := errors.AsType[*apierr.APIError](err); ok { + switch apiErr.ErrorCode { + case "FEATURE_DISABLED", "ENDPOINT_NOT_FOUND": + return errBricklensFeatureDisabled + } + if apiErr.StatusCode == http.StatusNotFound { + return errBricklensFeatureDisabled + } + } + if errors.Is(err, apierr.ErrResourceDoesNotExist) { + return err + } + return nil +} + +// fromSeconds computes the `from` bound. With --minutes set it is now-N*60; +// otherwise the run's start second (0 before the run starts, which the endpoint +// reads as "everything stored"). +func (req logRequest) fromSeconds(status logRunStatus, now time.Time) int64 { + if req.windowMinutes > 0 { + return now.Add(-time.Duration(req.windowMinutes) * time.Minute).Unix() + } + if status.startTimeMs > 0 { + return status.startTimeMs / 1000 + } + return 0 +} + +// toSeconds computes the `to` bound. A terminal run caps at its end second (ceil +// of the millisecond time, so the final partial second is kept); otherwise 0 lets +// the endpoint default to now. +func (req logRequest) toSeconds(status logRunStatus) int64 { + if status.terminal() && status.endTimeMs > 0 { + return (status.endTimeMs + 999) / 1000 + } + return 0 +} + +// streamBricklensLogs fetches and prints a run's logs: a bounded tail for a +// completed run, or a poll-and-drain loop that follows an active run to +// completion. It returns whether the run finished with SUCCESS; +// errBricklensFeatureDisabled means the caller should fall back to MLflow. +func streamBricklensLogs(ctx context.Context, w *databricks.WorkspaceClient, out io.Writer, req logRequest, status logRunStatus) (bool, error) { + // Build the API client once and reuse it for every page fetch in the loop. + apiClient, err := client.New(w.Config) + if err != nil { + return false, fmt.Errorf("failed to create API client: %w", err) + } + st := &bricklensStreamer{ + ctx: ctx, + w: w, + apiClient: apiClient, + out: out, + req: req, + status: status, + seen: newSeenSet(seenRecordsCap), + } + return st.run() +} + +// bricklensStreamer holds the poll-loop state: the from-second cursor, the +// highest emitted timestamp, and the dedup set. +type bricklensStreamer struct { + ctx context.Context + w *databricks.WorkspaceClient + apiClient *client.DatabricksClient + out io.Writer + req logRequest + status logRunStatus + + fromSec int64 + lastNano int64 + firstLogSeen bool + seen *seenSet + // previousState is the last display state reported to onStatusChange. + previousState string + // onFirstLog, when set, is called once just before the first log line is + // emitted — used to stop the "waiting for run to start" spinner before any + // log byte reaches stdout. + onFirstLog func() + // updateSpinner, when set, refreshes the waiting-spinner text each poll. + updateSpinner func(string) +} + +// reportStatusChange fires onStatusChange when the run's display state differs +// from the last reported one. +func (st *bricklensStreamer) reportStatusChange() { + if st.req.onStatusChange == nil { + return + } + current := st.status.displayState() + if current == st.previousState { + return + } + st.req.onStatusChange(current, st.previousState) + st.previousState = current +} + +func (st *bricklensStreamer) run() (bool, error) { + now := time.Now() + st.fromSec = st.req.fromSeconds(st.status, now) + + // A past retry's logs are immutable: render a one-shot tail rather than + // following the still-active run, which would poll forever. + if st.req.staticView { + return st.drainStatic(st.req.toSeconds(st.status)) + } + + // Show a "waiting for run to start" spinner on stderr while the run has not + // yet produced logs, so a provisioning run doesn't look hung. Suppressed in + // --json mode and auto-degraded to nothing on a non-interactive terminal. + // The first emitted log line stops it via onFirstLog (before any stdout write). + if !st.req.jsonOutput { + sp := cmdio.NewSpinner(st.ctx) + defer sp.Close() + st.onFirstLog = sp.Close + st.updateSpinner = sp.Update + } + + firstIteration := true + for { + if !firstIteration { + status, err := resolveRunStatus(st.ctx, st.w, st.req.runID) + if err != nil { + if errors.Is(err, apierr.ErrResourceDoesNotExist) { + return false, err + } + // A cancelled context (Ctrl-C) is not a transient blip: stop + // promptly instead of retrying forever. + if st.ctx.Err() != nil { + return false, st.ctx.Err() + } + // A transient status blip should not abort a live stream. + log.Debugf(st.ctx, "air logs: failed to refresh run status: %v", err) + if err := sleepOrCancel(st.ctx, retryCheckInterval); err != nil { + return false, err + } + continue + } + st.status = status + } + + st.reportStatusChange() + + terminal := st.status.terminal() + toSec := st.req.toSeconds(st.status) + + // While waiting on a still-active run with no logs yet, refresh the spinner. + if !terminal && !st.firstLogSeen && st.updateSpinner != nil { + st.updateSpinner(fmt.Sprintf("Waiting for run to start (node %d)...", st.req.node)) + } + + // A run already terminal on the first iteration renders as a tail (most + // recent N lines). An active run streams everything with dedup, so a run + // that terminates while we watch doesn't re-print the boundary second. + var err error + if firstIteration && terminal { + err = st.drainTail(toSec) + } else { + err = st.drainPages(toSec) + } + if err != nil { + return false, err + } + + if terminal { + if !st.firstLogSeen { + // Stop the spinner before the no-logs line so frames don't smear. + if st.onFirstLog != nil { + st.onFirstLog() + } + st.emitNoLogs() + } + log.Infof(st.ctx, "air logs: run %d finished in state %s", st.req.runID, st.status.displayState()) + return st.status.succeeded(), nil + } + + firstIteration = false + if err := sleepOrCancel(st.ctx, retryCheckInterval); err != nil { + return false, err + } + } +} + +// sleepOrCancel waits for d, or returns early with the context error if the +// context is cancelled (e.g. Ctrl-C) so the poll loop exits promptly. +func sleepOrCancel(ctx context.Context, d time.Duration) error { + timer := time.NewTimer(d) + defer timer.Stop() + select { + case <-ctx.Done(): + return ctx.Err() + case <-timer.C: + return nil + } +} + +// drainStatic renders a single tail pass without following the run. Success +// reflects the run's current result state (empty while active). +func (st *bricklensStreamer) drainStatic(toSec int64) (bool, error) { + if err := st.drainTail(toSec); err != nil { + return false, err + } + if !st.firstLogSeen { + st.emitNoLogs() + } + return st.status.succeeded(), nil +} + +// tailTarget is the number of lines a tail keeps. A negative tailLines means +// --lines was unset, so use the default cap; 0 or more is taken literally (an +// explicit --lines 0 prints nothing). +func (req logRequest) tailTarget() int { + if req.tailLines < 0 { + return defaultCompletedRunTailLines + } + return req.tailLines +} + +// drainTail emits the most-recent `target` records oldest-first. Bricklens +// returns records newest-first, so it pages until it has `target`, keeps the +// newest `target`, and reverses to chronological order. +func (st *bricklensStreamer) drainTail(toSec int64) error { + target := st.req.tailTarget() + if target <= 0 { + return nil + } + + var collected []logRecord + var pageToken string + for len(collected) < target { + resp, err := st.requestPage(pageToken, toSec, target, false) + if err != nil { + return err + } + collected = append(collected, resp.LogRecords...) + pageToken = resp.NextPageToken + if pageToken == "" { + break + } + } + + // Keep the newest `target`, then reverse to print oldest -> newest. + if len(collected) > target { + collected = collected[:target] + } + for _, c := range slices.Backward(collected) { + st.emit(c.Body) + } + return nil +} + +// drainPages exhausts all pages from the current from-second in ascending order, +// deduping against the seen-set so a re-queried boundary second is not +// re-printed, then advances fromSec to the newest record's floor-second. +func (st *bricklensStreamer) drainPages(toSec int64) error { + var pageToken string + for { + resp, err := st.requestPage(pageToken, toSec, 0, true) + if err != nil { + return err + } + + for _, rec := range resp.LogRecords { + nano := rec.nano() + if nano != 0 { + // Skip a record older than the last emitted one to keep output + // monotonic (out of order, or a re-queried boundary record). + if st.lastNano != 0 && nano < st.lastNano { + continue + } + if st.seen.has(nano, rec.Body) { + continue + } + } + st.emit(rec.Body) + if nano != 0 { + st.seen.add(nano, rec.Body) + st.lastNano = max(st.lastNano, nano) + } + } + + pageToken = resp.NextPageToken + if pageToken == "" { + break + } + } + + if st.lastNano != 0 { + st.fromSec = st.lastNano / 1_000_000_000 + } + return nil +} + +// requestPage fetches one page, retrying transient failures up to +// maxTransientFailures before falling back to MLflow. A feature-gated response or +// genuine not-found returns immediately. +func (st *bricklensStreamer) requestPage(pageToken string, toSec int64, pageSize int, ascending bool) (*bricklensLogsResponse, error) { + q := bricklensLogsQuery{ + fromSeconds: st.fromSec, + toSeconds: toSec, + pageToken: pageToken, + pageSize: pageSize, + attemptNumber: st.req.attempt, + nodeIndex: st.req.node, + ascending: ascending, + } + + transientFailures := 0 + for { + resp, err := getBricklensLogs(st.ctx, st.apiClient, st.req.runID, q) + if err == nil { + return resp, nil + } + + switch classified := classifyLogError(err); { + case errors.Is(classified, errBricklensFeatureDisabled): + return nil, errBricklensFeatureDisabled + case classified != nil: + return nil, classified + } + + transientFailures++ + if transientFailures >= maxTransientFailures { + log.Debugf(st.ctx, "air logs: bricklens failed %d times; falling back to mlflow", maxTransientFailures) + return nil, errBricklensFeatureDisabled + } + log.Debugf(st.ctx, "air logs: bricklens transient failure (%d/%d): %v", transientFailures, maxTransientFailures, err) + if err := sleepOrCancel(st.ctx, retryCheckInterval); err != nil { + return nil, err + } + } +} + +// emit writes one log line and latches firstLogSeen so an empty terminal run can +// report "no logs". The first line stops the waiting spinner before any byte +// reaches stdout. +func (st *bricklensStreamer) emit(body string) { + if !st.firstLogSeen && st.onFirstLog != nil { + st.onFirstLog() + } + st.firstLogSeen = true + emitLogLine(st.out, st.req, body) +} + +func (st *bricklensStreamer) emitNoLogs() { + emitNoLogs(st.out, st.req, st.status) +} + +// displayState is the result state, else the lifecycle state, else "UNKNOWN". +func (s logRunStatus) displayState() string { + if s.resultState != "" { + return s.resultState + } + if s.lifeCycleState != "" { + return s.lifeCycleState + } + return "UNKNOWN" +} diff --git a/experimental/air/cmd/logstream_support.go b/experimental/air/cmd/logstream_support.go new file mode 100644 index 00000000000..ac250547e16 --- /dev/null +++ b/experimental/air/cmd/logstream_support.go @@ -0,0 +1,179 @@ +package aircmd + +import ( + "container/list" + "encoding/json" + "fmt" + "io" + "time" +) + +// seenNano keys the dedup set. Distinct lines can share a nano (each rank stamps +// from its own clock), so the body disambiguates them. +type seenNano struct { + nano int64 + body string +} + +// seenSet is an insertion-ordered set bounded to a capacity, evicting the +// oldest-inserted entry first. +type seenSet struct { + cap int + items map[seenNano]*list.Element + order *list.List +} + +func newSeenSet(capacity int) *seenSet { + return &seenSet{ + cap: capacity, + items: make(map[seenNano]*list.Element), + order: list.New(), + } +} + +func (s *seenSet) has(nano int64, body string) bool { + _, ok := s.items[seenNano{nano, body}] + return ok +} + +func (s *seenSet) add(nano int64, body string) { + key := seenNano{nano, body} + if _, ok := s.items[key]; ok { + return + } + s.items[key] = s.order.PushBack(key) + if s.order.Len() > s.cap { + oldest := s.order.Front() + s.order.Remove(oldest) + delete(s.items, oldest.Value.(seenNano)) + } +} + +// logEvent is one JSONL streaming event. +type logEvent struct { + Type string `json:"type"` + TS string `json:"ts"` + Node int `json:"node"` + Line string `json:"line"` +} + +// printLogEvent writes a single JSONL event line for --json streaming output. +func printLogEvent(out io.Writer, eventType string, node int, line string) { + b, err := json.Marshal(logEvent{ + Type: eventType, + TS: time.Now().UTC().Format(time.RFC3339), + Node: node, + Line: line, + }) + if err != nil { + return + } + fmt.Fprintln(out, string(b)) +} + +// submittedEvent is the JSONL event `air run --watch -o json` emits before the +// streamed log events, so a consumer sees the run id immediately. +type submittedEvent struct { + Type string `json:"type"` + TS string `json:"ts"` + RunID string `json:"run_id"` + DashboardURL string `json:"dashboard_url"` +} + +// printSubmittedEvent writes the SUBMITTED JSONL event. +func printSubmittedEvent(out io.Writer, runID, dashboardURL string) { + b, err := json.Marshal(submittedEvent{ + Type: "SUBMITTED", + TS: time.Now().UTC().Format(time.RFC3339), + RunID: runID, + DashboardURL: dashboardURL, + }) + if err != nil { + return + } + fmt.Fprintln(out, string(b)) +} + +// statusEvent is a JSONL event emitted on each lifecycle transition while +// following a run with --watch. +type statusEvent struct { + Type string `json:"type"` + TS string `json:"ts"` + Status string `json:"status"` + Previous string `json:"previous_status,omitempty"` +} + +// printStatusEvent writes a STATUS JSONL event for a lifecycle transition. +func printStatusEvent(out io.Writer, current, previous string) { + b, err := json.Marshal(statusEvent{ + Type: "STATUS", + TS: time.Now().UTC().Format(time.RFC3339), + Status: current, + Previous: previous, + }) + if err != nil { + return + } + fmt.Fprintln(out, string(b)) +} + +// terminalEvent is the closing envelope `air run --watch -o json` emits after +// streaming, carrying the run's terminal status. +type terminalEvent struct { + V int `json:"v"` + TS string `json:"ts"` + Data runResult `json:"data"` +} + +// printTerminalEvent writes the closing terminal-status envelope, matching the +// shape of renderEnvelope(runResult). +func printTerminalEvent(out io.Writer, runID, status, dashboardURL string) { + b, err := json.Marshal(terminalEvent{ + V: envelopeVersion, + TS: time.Now().UTC().Format(time.RFC3339), + Data: runResult{ + Status: status, + RunID: runID, + DashboardURL: dashboardURL, + }, + }) + if err != nil { + return + } + fmt.Fprintln(out, string(b)) +} + +// emitLogLine writes one log line: raw in text mode, or a JSONL LOG event under +// --json. In --json mode a line matching a fatal-failure pattern also emits an +// ALERT event first, giving an agent an immediate actionable signal. +func emitLogLine(out io.Writer, req logRequest, body string) { + if !req.jsonOutput { + fmt.Fprintln(out, body) + return + } + if matchFatalPattern(body) { + printLogEvent(out, "ALERT", req.node, body) + } + printLogEvent(out, "LOG", req.node, body) +} + +// emitNoLogs reports that a run produced no logs. A terminal run gets its +// termination reason; a still-active run is reported as having no logs yet, +// since the MLflow fallback is a one-shot that does not follow it to completion. +// Under --json it is a JSONL ERROR, so a consumer never sees an empty stream. +func emitNoLogs(out io.Writer, req logRequest, status logRunStatus) { + var msg string + if status.terminal() { + msg = fmt.Sprintf("No logs available for run %d. Run terminated in state %s", req.runID, status.displayState()) + } else { + msg = fmt.Sprintf("No logs available yet for run %d, which is still in state %s", req.runID, status.displayState()) + } + if status.stateMessage != "" { + msg = fmt.Sprintf("%s: %s", msg, status.stateMessage) + } + if req.jsonOutput { + printLogEvent(out, "ERROR", req.node, msg) + return + } + fmt.Fprintln(out, msg) +} diff --git a/experimental/air/cmd/logstream_test.go b/experimental/air/cmd/logstream_test.go new file mode 100644 index 00000000000..d8a061687a6 --- /dev/null +++ b/experimental/air/cmd/logstream_test.go @@ -0,0 +1,403 @@ +package aircmd + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/databricks/databricks-sdk-go/apierr" + "github.com/databricks/databricks-sdk-go/client" + "github.com/databricks/databricks-sdk-go/service/jobs" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestClassifyLogError(t *testing.T) { + tests := []struct { + name string + err error + want error // errBricklensFeatureDisabled, the input error, or nil + }{ + { + name: "feature disabled falls back", + err: &apierr.APIError{ErrorCode: "FEATURE_DISABLED", StatusCode: http.StatusForbidden}, + want: errBricklensFeatureDisabled, + }, + { + name: "endpoint not found falls back", + err: &apierr.APIError{ErrorCode: "ENDPOINT_NOT_FOUND", StatusCode: http.StatusNotFound}, + want: errBricklensFeatureDisabled, + }, + { + name: "bare 404 falls back", + err: &apierr.APIError{ErrorCode: "SOMETHING", StatusCode: http.StatusNotFound}, + want: errBricklensFeatureDisabled, + }, + { + name: "genuine resource-does-not-exist surfaces", + err: apierr.ErrResourceDoesNotExist, + want: apierr.ErrResourceDoesNotExist, + }, + { + name: "transient 500 is retried", + err: &apierr.APIError{ErrorCode: "INTERNAL", StatusCode: http.StatusInternalServerError}, + want: nil, + }, + { + name: "plain error is retried", + err: errors.New("connection reset"), + want: nil, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := classifyLogError(tt.err) + switch tt.want { + case errBricklensFeatureDisabled: + assert.ErrorIs(t, got, errBricklensFeatureDisabled) + case nil: + assert.NoError(t, got) + default: + assert.ErrorIs(t, got, tt.want) + } + }) + } +} + +func TestProjectRunStatus(t *testing.T) { + run := &jobs.Run{ + StartTime: 1000, + EndTime: 2000, + State: &jobs.RunState{ + LifeCycleState: jobs.RunLifeCycleStateTerminated, + ResultState: jobs.RunResultStateSuccess, + StateMessage: "done", + }, + Tasks: []jobs.RunTask{ + {AttemptNumber: 0}, + {AttemptNumber: 2}, + {AttemptNumber: 1}, + }, + } + + s := projectRunStatus(run) + assert.Equal(t, "TERMINATED", s.lifeCycleState) + assert.Equal(t, "SUCCESS", s.resultState) + assert.Equal(t, "done", s.stateMessage) + assert.Equal(t, int64(1000), s.startTimeMs) + assert.Equal(t, int64(2000), s.endTimeMs) + assert.Equal(t, 2, s.latestAttempt) + assert.True(t, s.terminal()) + assert.True(t, s.succeeded()) + assert.Equal(t, "SUCCESS", s.displayState()) +} + +func TestLogRunStatusTerminal(t *testing.T) { + tests := []struct { + name string + lifeCycle string + resultState string + wantTerminal bool + }{ + {"running", "RUNNING", "", false}, + {"pending", "PENDING", "", false}, + {"terminated lifecycle", "TERMINATED", "", true}, + {"internal error lifecycle", "INTERNAL_ERROR", "", true}, + {"failed result", "TERMINATING", "FAILED", true}, + {"canceled result", "RUNNING", "CANCELED", true}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + s := logRunStatus{lifeCycleState: tt.lifeCycle, resultState: tt.resultState} + assert.Equal(t, tt.wantTerminal, s.terminal()) + }) + } +} + +func TestLogRequestFromSeconds(t *testing.T) { + now := time.Unix(10_000, 0) + + // --minutes narrows the window to now - N*60. + req := logRequest{windowMinutes: 5} + assert.Equal(t, int64(10_000-300), req.fromSeconds(logRunStatus{startTimeMs: 1_000_000}, now)) + + // No window: from the run's start second. + req = logRequest{} + assert.Equal(t, int64(1000), req.fromSeconds(logRunStatus{startTimeMs: 1_000_000}, now)) + + // No window, run not started: everything stored (0). + assert.Equal(t, int64(0), req.fromSeconds(logRunStatus{}, now)) +} + +func TestLogRequestToSeconds(t *testing.T) { + req := logRequest{} + + // Active run: 0 lets the endpoint default to now. + assert.Equal(t, int64(0), req.toSeconds(logRunStatus{lifeCycleState: "RUNNING"})) + + // Terminal run: ceil of the end millisecond so the final partial second is kept. + terminal := logRunStatus{lifeCycleState: "TERMINATED", resultState: "SUCCESS", endTimeMs: 2001} + assert.Equal(t, int64(3), req.toSeconds(terminal)) +} + +func TestLogRequestTailTarget(t *testing.T) { + // Negative (unset) uses the default cap; explicit values are literal. + assert.Equal(t, defaultCompletedRunTailLines, logRequest{tailLines: -1}.tailTarget()) + assert.Equal(t, 42, logRequest{tailLines: 42}.tailTarget()) + assert.Equal(t, 0, logRequest{tailLines: 0}.tailTarget()) +} + +func TestDrainPagesDedupAndOrdering(t *testing.T) { + // Two pages: page 1 has two ascending records; page 2 repeats the last record + // of page 1 (boundary re-query — must dedup) and includes an older record + // (out of order — must skip), then a genuinely newer one. + var page int + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Query().Get("page_token") == "" { + page = 1 + _, _ = w.Write([]byte(`{"log_records": [ + {"time_unix_nano": 1000, "body": "a", "node_index": 0}, + {"time_unix_nano": 2000, "body": "b", "node_index": 0} + ], "next_page_token": "p2"}`)) + return + } + page = 2 + _, _ = w.Write([]byte(`{"log_records": [ + {"time_unix_nano": 2000, "body": "b", "node_index": 0}, + {"time_unix_nano": 1500, "body": "stale", "node_index": 0}, + {"time_unix_nano": 3000, "body": "c", "node_index": 0} + ]}`)) + })) + t.Cleanup(srv.Close) + + var buf bytes.Buffer + w := newTestWorkspaceClient(t, srv.URL) + apiClient, err := client.New(w.Config) + require.NoError(t, err) + st := &bricklensStreamer{ + ctx: t.Context(), + w: w, + apiClient: apiClient, + out: &buf, + req: logRequest{runID: 1, node: 0, attempt: -1}, + seen: newSeenSet(seenRecordsCap), + } + require.NoError(t, st.drainPages(0)) + require.Equal(t, 2, page) + + // "b" prints once (deduped), "stale" is skipped (older than last emitted), and + // fromSec advances to the newest record's floor-second (3000ns -> 0s here). + assert.Equal(t, "a\nb\nc\n", buf.String()) + assert.Equal(t, int64(3000), st.lastNano) +} + +func TestDisplayState(t *testing.T) { + assert.Equal(t, "SUCCESS", logRunStatus{lifeCycleState: "TERMINATED", resultState: "SUCCESS"}.displayState()) + assert.Equal(t, "RUNNING", logRunStatus{lifeCycleState: "RUNNING"}.displayState()) + assert.Equal(t, "UNKNOWN", logRunStatus{}.displayState()) +} + +func TestEmitLogLineJSON(t *testing.T) { + var buf bytes.Buffer + emitLogLine(&buf, logRequest{node: 2, jsonOutput: true}, "hello") + + var ev logEvent + require.NoError(t, json.Unmarshal(buf.Bytes(), &ev)) + assert.Equal(t, "LOG", ev.Type) + assert.Equal(t, 2, ev.Node) + assert.Equal(t, "hello", ev.Line) + assert.NotEmpty(t, ev.TS) +} + +func TestEmitLogLineText(t *testing.T) { + var buf bytes.Buffer + emitLogLine(&buf, logRequest{node: 0}, "hello") + assert.Equal(t, "hello\n", buf.String()) +} + +func TestEmitLogLineJSONFatalEmitsAlert(t *testing.T) { + var buf bytes.Buffer + emitLogLine(&buf, logRequest{node: 1, jsonOutput: true}, "CUDA out of memory") + + // A fatal line emits an ALERT event before its LOG event. + lines := strings.Split(strings.TrimSpace(buf.String()), "\n") + require.Len(t, lines, 2) + + var alert, logEv logEvent + require.NoError(t, json.Unmarshal([]byte(lines[0]), &alert)) + require.NoError(t, json.Unmarshal([]byte(lines[1]), &logEv)) + assert.Equal(t, "ALERT", alert.Type) + assert.Equal(t, "LOG", logEv.Type) + assert.Equal(t, "CUDA out of memory", alert.Line) + + // Text mode never emits ALERT events. + var text bytes.Buffer + emitLogLine(&text, logRequest{node: 1}, "CUDA out of memory") + assert.Equal(t, "CUDA out of memory\n", text.String()) +} + +func TestMatchFatalPattern(t *testing.T) { + fatal := []string{ + "CUDA out of memory", + "cuda OUT OF memory", + "Watchdog caught collective operation timeout", + "Killed", + "ERROR: Script failed with exit code 1 after 42s", + "bash: foo: command not found", + } + for _, l := range fatal { + assert.True(t, matchFatalPattern(l), l) + } + + notFatal := []string{ + "epoch 3 loss 0.5", + "ERROR: Script failed with exit code 0 after 42s", + "just a normal line", + } + for _, l := range notFatal { + assert.False(t, matchFatalPattern(l), l) + } +} + +func TestEmitNoLogs(t *testing.T) { + tests := []struct { + name string + status logRunStatus + want string + }{ + { + name: "terminal", + status: logRunStatus{lifeCycleState: "TERMINATED", resultState: "FAILED", stateMessage: "boom"}, + want: "No logs available for run 7. Run terminated in state FAILED: boom\n", + }, + { + name: "running", + status: logRunStatus{lifeCycleState: "RUNNING"}, + want: "No logs available yet for run 7, which is still in state RUNNING\n", + }, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + var text bytes.Buffer + emitNoLogs(&text, logRequest{runID: 7}, tc.status) + assert.Equal(t, tc.want, text.String()) + + var jsonBuf bytes.Buffer + emitNoLogs(&jsonBuf, logRequest{runID: 7, node: 1, jsonOutput: true}, tc.status) + var ev logEvent + require.NoError(t, json.Unmarshal(jsonBuf.Bytes(), &ev)) + assert.Equal(t, "ERROR", ev.Type) + assert.Equal(t, 1, ev.Node) + assert.Equal(t, strings.TrimRight(tc.want, "\n"), ev.Line) + }) + } +} + +func TestRequestPageRetriesThenFallsBack(t *testing.T) { + // Shrink the retry wait so the transient-failure loop runs fast. + orig := retryCheckInterval + retryCheckInterval = time.Millisecond + t.Cleanup(func() { retryCheckInterval = orig }) + + var calls int + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if !strings.Contains(r.URL.Path, "/logs") { + // Ignore SDK host/config probes so `calls` counts only log requests. + _, _ = w.Write([]byte(`{}`)) + return + } + calls++ + w.WriteHeader(http.StatusInternalServerError) + _, _ = w.Write([]byte(`{"error_code": "INTERNAL_ERROR", "message": "transient"}`)) + })) + t.Cleanup(srv.Close) + + w := newTestWorkspaceClient(t, srv.URL) + apiClient, err := client.New(w.Config) + require.NoError(t, err) + st := &bricklensStreamer{ + ctx: t.Context(), + w: w, + apiClient: apiClient, + req: logRequest{runID: 1, node: 0, attempt: -1}, + seen: newSeenSet(seenRecordsCap), + } + _, err = st.requestPage("", 0, 0, true) + // Persistent transient failures fall back to MLflow after the retry budget. + require.ErrorIs(t, err, errBricklensFeatureDisabled) + assert.Equal(t, maxTransientFailures, calls) +} + +func TestRequestPageRetriesThenSucceeds(t *testing.T) { + orig := retryCheckInterval + retryCheckInterval = time.Millisecond + t.Cleanup(func() { retryCheckInterval = orig }) + + var calls int + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if !strings.Contains(r.URL.Path, "/logs") { + _, _ = w.Write([]byte(`{}`)) + return + } + calls++ + if calls < 3 { + w.WriteHeader(http.StatusInternalServerError) + _, _ = w.Write([]byte(`{"error_code": "INTERNAL_ERROR", "message": "transient"}`)) + return + } + _, _ = w.Write([]byte(`{"log_records": [{"time_unix_nano": 1, "body": "ok", "node_index": 0}]}`)) + })) + t.Cleanup(srv.Close) + + wc := newTestWorkspaceClient(t, srv.URL) + apiClient, err := client.New(wc.Config) + require.NoError(t, err) + st := &bricklensStreamer{ + ctx: t.Context(), + w: wc, + apiClient: apiClient, + req: logRequest{runID: 1, node: 0, attempt: -1}, + seen: newSeenSet(seenRecordsCap), + } + resp, err := st.requestPage("", 0, 0, true) + require.NoError(t, err) + require.Len(t, resp.LogRecords, 1) + assert.Equal(t, "ok", resp.LogRecords[0].Body) + assert.Equal(t, 3, calls) +} + +func TestSeenSetEviction(t *testing.T) { + s := newSeenSet(2) + s.add(1, "a") + s.add(2, "b") + assert.True(t, s.has(1, "a")) + assert.True(t, s.has(2, "b")) + + // Adding a third evicts the oldest-inserted (1,"a"). + s.add(3, "c") + assert.False(t, s.has(1, "a")) + assert.True(t, s.has(2, "b")) + assert.True(t, s.has(3, "c")) + + // Same (nano, body) shares one entry; distinct body under the same nano does not. + s.add(3, "c") + assert.True(t, s.has(3, "c")) + assert.False(t, s.has(3, "d")) +} + +func TestSleepOrCancel(t *testing.T) { + // Returns nil once the duration elapses. + require.NoError(t, sleepOrCancel(t.Context(), time.Millisecond)) + + // Returns the context error promptly when cancelled. + ctx, cancel := context.WithCancel(t.Context()) + cancel() + require.ErrorIs(t, sleepOrCancel(ctx, time.Hour), context.Canceled) +} diff --git a/experimental/air/cmd/render_test.go b/experimental/air/cmd/render_test.go index 2fe125bcf04..067f5817fb1 100644 --- a/experimental/air/cmd/render_test.go +++ b/experimental/air/cmd/render_test.go @@ -151,7 +151,7 @@ func TestRenderFields(t *testing.T) { experiment: "stream-latency-test", mlflowLabel: "stream-latency-test", mlflowURL: "https://h.test/ml/experiments/E1/runs/R1", - user: "riddhi.bhagwat@databricks.com", + user: "user@example.com", accelerators: "1x A10", environment: "ml-runtime-gpu:1.0", }) diff --git a/experimental/air/cmd/run.go b/experimental/air/cmd/run.go index bd32810e9bc..ea00368679f 100644 --- a/experimental/air/cmd/run.go +++ b/experimental/air/cmd/run.go @@ -1,7 +1,7 @@ package aircmd import ( - "errors" + "context" "fmt" "strconv" @@ -9,6 +9,7 @@ import ( "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/spf13/cobra" ) @@ -57,16 +58,7 @@ The workload is described by a YAML config file (see --file).`, cmd.RunE = func(cmd *cobra.Command, args []string) error { ctx := cmd.Context() - // These flags' pipelines are not ported yet; reject rather than silently - // ignore them. - if len(overrides) > 0 { - return errors.New("--override is not yet supported") - } - if watch { - return errors.New("--watch is not yet supported") - } - - cfg, err := loadRunConfig(file) + cfg, err := loadRunConfigWithOverrides(ctx, file, overrides) if err != nil { return err } @@ -86,13 +78,63 @@ The workload is described by a YAML config file (see --file).`, } runIDStr := strconv.FormatInt(runID, 10) - if root.OutputType(cmd) == flags.OutputText { + jsonOut := root.OutputType(cmd) == flags.OutputJSON + + if !watch { + if !jsonOut { + cmdio.LogString(ctx, "Submitted run "+runIDStr) + cmdio.LogString(ctx, "View at: "+dashboardURL) + cmdio.LogString(ctx, "\nTip: use --watch to stream logs until the run completes.") + return nil + } + return renderEnvelope(ctx, runResult{Status: "SUBMITTED", RunID: runIDStr, DashboardURL: dashboardURL}) + } + + // --watch: stream the submitted run's logs until it reaches a terminal + // state, then exit with the run's outcome. This is the same pipeline as + // `air logs ` (Bricklens with MLflow fallback). + req := logRequest{ + runID: runID, + attempt: -1, + tailLines: -1, + jsonOutput: jsonOut, + } + + if !jsonOut { cmdio.LogString(ctx, "Submitted run "+runIDStr) cmdio.LogString(ctx, "View at: "+dashboardURL) - return nil + cmdio.LogString(ctx, "Monitoring run and streaming logs...") + return runLogs(ctx, cmd, req) + } + + // --json: emit SUBMITTED first (so a consumer sees the run id immediately), + // STATUS events on each lifecycle transition, and a closing terminal-status + // envelope after streaming. Mirrors the Python CLI's --watch JSONL contract. + out := cmd.OutOrStdout() + printSubmittedEvent(out, runIDStr, dashboardURL) + req.onStatusChange = func(current, previous string) { + printStatusEvent(out, current, previous) } - return renderEnvelope(ctx, runResult{Status: "SUBMITTED", RunID: runIDStr, DashboardURL: dashboardURL}) + err = runLogs(ctx, cmd, req) + + // Re-resolve the run for the closing envelope. STATUS events only fire on + // the Bricklens path, so the terminal status must come from the run's + // actual state — correct whether Bricklens or the MLflow fallback served + // the logs. + printTerminalEvent(out, runIDStr, watchTerminalStatus(ctx, w, runID), dashboardURL) + return err } return cmd } + +// watchTerminalStatus resolves a watched run's final display state for the +// closing --watch envelope. The run is terminal once streaming returns; if the +// status can't be re-fetched, "UNKNOWN" is reported rather than guessing. +func watchTerminalStatus(ctx context.Context, w *databricks.WorkspaceClient, runID int64) string { + status, err := resolveRunStatus(ctx, w, runID) + if err != nil { + return "UNKNOWN" + } + return status.displayState() +} diff --git a/experimental/air/cmd/run_watch_test.go b/experimental/air/cmd/run_watch_test.go new file mode 100644 index 00000000000..487b0b9989a --- /dev/null +++ b/experimental/air/cmd/run_watch_test.go @@ -0,0 +1,206 @@ +package aircmd + +import ( + "bytes" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "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/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// watchServer serves submit, the auth probe, a terminal runs/get, and a single +// page of Bricklens logs — everything `air run --watch` touches after submit. +// resultState is the terminal result the run reports (e.g. SUCCESS, FAILED). +func watchServer(t *testing.T, resultState string) *httptest.Server { + t.Helper() + runGet := `{ + "run_id": 777, + "start_time": 1700000000000, + "end_time": 1700000012000, + "state": {"life_cycle_state": "TERMINATED", "result_state": "` + resultState + `"}, + "tasks": [{"run_id": 778, "attempt_number": 0}] + }` + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch { + case strings.HasSuffix(r.URL.Path, "/jobs/runs/submit"): + _, _ = w.Write([]byte(`{"run_id": 777}`)) + case r.URL.Path == "/api/2.2/jobs/runs/get": + _, _ = w.Write([]byte(runGet)) + case strings.HasPrefix(r.URL.Path, "/api/2.0/ai-training/workflows/by-run-id/"): + _, _ = w.Write([]byte(`{"log_records": [ + {"time_unix_nano": 1700000002000000000, "body": "step 2", "node_index": 0}, + {"time_unix_nano": 1700000001000000000, "body": "step 1", "node_index": 0} + ]}`)) + default: + // Me() probe, workspace-id, SDK config discovery. + _, _ = w.Write([]byte(`{"userName": "u@example.com", "workspace_id": 1}`)) + } + })) + t.Cleanup(srv.Close) + return srv +} + +// watchServerMLflow serves a run whose Bricklens endpoint is gated off +// (FEATURE_DISABLED), forcing the MLflow fallback, plus the MLflow artifact +// chain (get-output, artifacts/list, credentials-for-read, the pre-signed bytes). +// STATUS events never fire on this path, so it guards the closing terminal +// envelope against relying on onStatusChange. +func watchServerMLflow(t *testing.T, resultState string) *httptest.Server { + t.Helper() + var base string + runGet := `{ + "run_id": 777, + "state": {"life_cycle_state": "TERMINATED", "result_state": "` + resultState + `"}, + "tasks": [{"run_id": 778, "attempt_number": 0}] + }` + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch { + case strings.HasSuffix(r.URL.Path, "/jobs/runs/submit"): + _, _ = w.Write([]byte(`{"run_id": 777}`)) + case r.URL.Path == "/api/2.2/jobs/runs/get": + _, _ = w.Write([]byte(runGet)) + case strings.HasPrefix(r.URL.Path, "/api/2.0/ai-training/workflows/by-run-id/"): + w.WriteHeader(http.StatusForbidden) + _, _ = w.Write([]byte(`{"error_code": "FEATURE_DISABLED", "message": "gated off"}`)) + case r.URL.Path == "/api/2.2/jobs/runs/get-output": + _, _ = w.Write([]byte(`{"ai_runtime_task_output": {"mlflow_experiment_id": "exp1", "mlflow_run_id": "run1"}}`)) + case r.URL.Path == "/api/2.0/mlflow/artifacts/list": + if r.URL.Query().Get("path") == "logs" { + _, _ = w.Write([]byte(`{"files": [{"path": "logs/node_0", "is_dir": true}]}`)) + return + } + _, _ = w.Write([]byte(`{"files": [{"path": "logs/node_0/logs-0.chunk.txt", "file_size": 12}]}`)) + case r.URL.Path == "/api/2.0/mlflow/artifacts/credentials-for-read": + _, _ = w.Write([]byte(`{"credential_infos": [{"signed_uri": "` + base + `/presigned"}]}`)) + case r.URL.Path == "/presigned": + _, _ = w.Write([]byte("step 1\nstep 2\n")) + default: + _, _ = w.Write([]byte(`{"userName": "u@example.com", "workspace_id": 1}`)) + } + })) + base = srv.URL + t.Cleanup(srv.Close) + return srv +} + +func runWatchCmd(t *testing.T, out flags.Output, buf *bytes.Buffer, srvURL string) error { + t.Helper() + cfgPath := writeConfigFile(t, "run.yaml", minimalConfig) + cmd := withOutput(newRunCommand(), out) + require.NoError(t, cmd.Flags().Set("file", cfgPath)) + require.NoError(t, cmd.Flags().Set("watch", "true")) + + ctx := cmdio.InContext(t.Context(), cmdio.NewIO(t.Context(), out, nil, buf, buf, "", "")) + ctx = cmdctx.SetWorkspaceClient(ctx, newTestWorkspaceClient(t, srvURL)) + cmd.SetContext(ctx) + cmd.SetOut(buf) + return cmd.RunE(cmd, nil) +} + +func TestRunWatchStreamsLogs(t *testing.T) { + var buf bytes.Buffer + err := runWatchCmd(t, flags.OutputText, &buf, watchServer(t, "SUCCESS").URL) + require.NoError(t, err) + + out := buf.String() + assert.Contains(t, out, "Submitted run 777") + assert.Contains(t, out, "Monitoring run and streaming logs...") + // The submitted run's logs stream through, oldest-first. + assert.Contains(t, out, "step 1\nstep 2") +} + +func TestRunWatchJSONEmitsSubmittedThenLogs(t *testing.T) { + var buf bytes.Buffer + err := runWatchCmd(t, flags.OutputJSON, &buf, watchServer(t, "SUCCESS").URL) + require.NoError(t, err) + + all := buf.String() + lines := strings.Split(strings.TrimSpace(all), "\n") + require.GreaterOrEqual(t, len(lines), 3) + // First event is SUBMITTED with the run id; then STATUS + streamed LOG events. + assert.Contains(t, lines[0], `"type":"SUBMITTED"`) + assert.Contains(t, lines[0], `"run_id":"777"`) + assert.Contains(t, all, `"type":"STATUS"`) + assert.Contains(t, all, `"type":"LOG"`) + assert.Contains(t, all, `"line":"step 1"`) + // The last line is the closing terminal-status envelope carrying SUCCESS. + assert.Contains(t, lines[len(lines)-1], `"status":"SUCCESS"`) + assert.Contains(t, lines[len(lines)-1], `"run_id":"777"`) +} + +func TestRunWatchJSONFailedRunTerminalEnvelope(t *testing.T) { + var buf bytes.Buffer + err := runWatchCmd(t, flags.OutputJSON, &buf, watchServer(t, "FAILED").URL) + // Non-zero exit is surfaced as ErrAlreadyPrinted, but the closing envelope + // still carries the terminal status. + require.ErrorIs(t, err, root.ErrAlreadyPrinted) + lines := strings.Split(strings.TrimSpace(buf.String()), "\n") + assert.Contains(t, lines[len(lines)-1], `"status":"FAILED"`) +} + +func TestRunWatchFailedRunExitsNonZero(t *testing.T) { + var buf bytes.Buffer + // A run that ends FAILED streams its logs but exits non-zero, surfaced as + // ErrAlreadyPrinted (the output was already written). + err := runWatchCmd(t, flags.OutputText, &buf, watchServer(t, "FAILED").URL) + require.ErrorIs(t, err, root.ErrAlreadyPrinted) + assert.Contains(t, buf.String(), "step 1\nstep 2") +} + +func TestRunWatchDryRunSkipsSubmit(t *testing.T) { + // --dry-run takes precedence over --watch: nothing is submitted or streamed. + var got []string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + got = append(got, r.URL.Path) + _, _ = w.Write([]byte(`{}`)) + })) + t.Cleanup(srv.Close) + + cfgPath := writeConfigFile(t, "run.yaml", minimalConfig) + cmd := withOutput(newRunCommand(), flags.OutputText) + require.NoError(t, cmd.Flags().Set("file", cfgPath)) + require.NoError(t, cmd.Flags().Set("watch", "true")) + require.NoError(t, cmd.Flags().Set("dry-run", "true")) + var buf bytes.Buffer + ctx := cmdio.InContext(t.Context(), cmdio.NewIO(t.Context(), flags.OutputText, nil, &buf, &buf, "", "")) + ctx = cmdctx.SetWorkspaceClient(ctx, newTestWorkspaceClient(t, srv.URL)) + cmd.SetContext(ctx) + cmd.SetOut(&buf) + + require.NoError(t, cmd.RunE(cmd, nil)) + assert.Contains(t, buf.String(), "Dry run") + for _, p := range got { + assert.NotContains(t, p, "/jobs/runs/submit", "dry-run must not submit") + assert.NotContains(t, p, "/logs", "dry-run must not stream logs") + } +} + +func TestRunWatchJSONMLflowFallbackTerminalEnvelope(t *testing.T) { + // Regression: through the MLflow fallback the terminal status must come from + // the run's actual state, not the onStatusChange callback (which is + // Bricklens-only), so a SUCCESS run isn't mislabeled FAILED in the envelope. + var buf bytes.Buffer + err := runWatchCmd(t, flags.OutputJSON, &buf, watchServerMLflow(t, "SUCCESS").URL) + require.NoError(t, err) + + lines := strings.Split(strings.TrimSpace(buf.String()), "\n") + assert.Contains(t, lines[0], `"type":"SUBMITTED"`) + // Logs stream through the fallback, and the closing envelope reflects the + // real terminal status. + assert.Contains(t, buf.String(), `"line":"step 1"`) + assert.Contains(t, lines[len(lines)-1], `"status":"SUCCESS"`) +} + +func TestRunWatchFlagRegistered(t *testing.T) { + cmd := newRunCommand() + f := cmd.Flags().Lookup("watch") + require.NotNil(t, f) + assert.Equal(t, "false", f.DefValue) +} diff --git a/experimental/air/cmd/runconfig_load.go b/experimental/air/cmd/runconfig_load.go index 81b07d3ca50..7f6ad8b5e17 100644 --- a/experimental/air/cmd/runconfig_load.go +++ b/experimental/air/cmd/runconfig_load.go @@ -1,6 +1,8 @@ package aircmd import ( + "bytes" + "context" "errors" "fmt" "io" @@ -10,10 +12,11 @@ import ( ) // decodeRunConfig reads and decodes the run YAML into the schema. Unknown keys -// are rejected (KnownFields), mirroring the Python schema's extra="forbid". +// are rejected (KnownFields). // -// The `_bases_` composition feature and CLI `--override` handling are not yet -// ported; a config using `_bases_` is currently rejected as an unknown field. +// The `_bases_` composition feature is not yet ported; a config using `_bases_` +// is currently rejected as an unknown field. CLI `--override` handling lives in +// runconfig_override.go and is applied to the parsed map before this decode. func decodeRunConfig(path string) (*runConfig, error) { f, err := os.Open(path) if err != nil { @@ -21,7 +24,13 @@ func decodeRunConfig(path string) (*runConfig, error) { } defer f.Close() - dec := yaml.NewDecoder(f) + return decodeRunConfigReader(f, path) +} + +// decodeRunConfigReader decodes and unknown-key-checks a run YAML from r. path is +// used only for error messages. +func decodeRunConfigReader(r io.Reader, path string) (*runConfig, error) { + dec := yaml.NewDecoder(r) dec.KnownFields(true) var cfg runConfig @@ -50,3 +59,53 @@ func loadRunConfig(path string) (*runConfig, error) { } return cfg, nil } + +// loadRunConfigWithOverrides decodes a run YAML config, applies any +// --override KEY=VALUE entries to the parsed map, then re-decodes (with unknown +// keys rejected) and structurally validates the result. Applying overrides to +// the map — rather than the typed config — lets the single decode+validate +// pipeline enforce path existence, type coercion, and the semantic rules at +// once. ctx is used only to log applied overrides. +func loadRunConfigWithOverrides(ctx context.Context, path string, overrides []string) (*runConfig, error) { + if len(overrides) == 0 { + return loadRunConfig(path) + } + + entries, err := parseOverrides(overrides) + if err != nil { + return nil, err + } + if err := validateOverridePaths(entries); err != nil { + return nil, err + } + + raw, err := os.ReadFile(path) + if err != nil { + return nil, err + } + var m map[string]any + if err := yaml.Unmarshal(raw, &m); err != nil { + return nil, fmt.Errorf("invalid config %s: %w", path, err) + } + if m == nil { + // An empty file decodes to a nil map; start from an empty one so overrides + // can populate it (the re-decode still enforces required fields). + m = map[string]any{} + } + if err := applyOverrides(ctx, m, entries); err != nil { + return nil, err + } + + merged, err := yaml.Marshal(m) + if err != nil { + return nil, err + } + cfg, err := decodeRunConfigReader(bytes.NewReader(merged), path) + if err != nil { + return nil, err + } + if err := validateRunConfig(cfg); err != nil { + return nil, err + } + return cfg, nil +} diff --git a/experimental/air/cmd/runconfig_override.go b/experimental/air/cmd/runconfig_override.go new file mode 100644 index 00000000000..a00a178603b --- /dev/null +++ b/experimental/air/cmd/runconfig_override.go @@ -0,0 +1,162 @@ +package aircmd + +import ( + "context" + "fmt" + "maps" + "reflect" + "slices" + "strings" + + "github.com/databricks/cli/libs/cmdio" + "go.yaml.in/yaml/v3" +) + +// This file implements the `--override KEY=VALUE` flag. Overrides are applied to +// the parsed YAML map (not the typed runConfig) before re-decode, so one pipeline +// covers path existence, type coercion, and the semantic validate() rules. + +// freeFormFields hold free-form maps, so path validation stops at them: any +// sub-path is valid. +var freeFormFields = map[string]bool{ + "parameters": true, + "env_variables": true, + "secrets": true, +} + +// parseOverrides parses --override KEY=VALUE arguments, preserving order. +func parseOverrides(overrides []string) ([]overrideEntry, error) { + entries := make([]overrideEntry, 0, len(overrides)) + for _, item := range overrides { + key, value, found := strings.Cut(item, "=") + if !found { + // --override is repeatable, so a config path meant for -f can be + // swallowed here; point at the real fix. + hint := "" + if strings.HasSuffix(item, ".yaml") || strings.HasSuffix(item, ".yml") { + hint = fmt.Sprintf("; %q looks like a config file — pass it with -f/--file", item) + } + return nil, fmt.Errorf("invalid --override %q: expected KEY=VALUE (e.g. compute.num_accelerators=32)%s", item, hint) + } + key = strings.TrimSpace(key) + if key == "" { + return nil, fmt.Errorf("invalid --override %q: empty key", item) + } + entries = append(entries, overrideEntry{path: key, raw: value}) + } + return entries, nil +} + +// overrideEntry is one parsed --override: its dotted path and the raw RHS string. +type overrideEntry struct { + path string + raw string +} + +// validateOverridePaths checks every dotted path against the runConfig schema +// before mutation, so an error names the exact --override key rather than the +// re-decode's Go-type language. +func validateOverridePaths(entries []overrideEntry) error { + for _, e := range entries { + if err := checkOverridePath(strings.Split(e.path, "."), reflect.TypeFor[runConfig](), e.path); err != nil { + return err + } + } + return nil +} + +// checkOverridePath recursively validates one dotted path against a struct type +// whose fields carry `yaml:` tags. +func checkOverridePath(parts []string, t reflect.Type, fullPath string) error { + field := parts[0] + fields := yamlFields(t) + sub, ok := fields[field] + if !ok { + return fmt.Errorf("invalid --override %q: %q is not a known field; available fields are: %s", + fullPath, field, strings.Join(slices.Sorted(maps.Keys(fields)), ", ")) + } + if len(parts) == 1 { + return nil + } + if freeFormFields[field] { + return nil + } + subStruct := underlyingStruct(sub) + if subStruct == nil { + return fmt.Errorf("invalid --override %q: %q is not a nested object; cannot address sub-field %q", + fullPath, field, strings.Join(parts[1:], ".")) + } + return checkOverridePath(parts[1:], subStruct, fullPath) +} + +// yamlFields maps a struct's yaml tag names to their field types, skipping +// fields without a yaml tag (or tagged "-"). +func yamlFields(t reflect.Type) map[string]reflect.Type { + out := map[string]reflect.Type{} + for f := range t.Fields() { + tag := f.Tag.Get("yaml") + if tag == "" || tag == "-" { + continue + } + name, _, _ := strings.Cut(tag, ",") + if name == "" || name == "-" { + continue + } + out[name] = f.Type + } + return out +} + +// underlyingStruct unwraps pointer/slice indirection and returns the struct type +// a field decodes into, or nil if the field is not a struct (a scalar/map/etc.). +func underlyingStruct(t reflect.Type) reflect.Type { + for t.Kind() == reflect.Pointer || t.Kind() == reflect.Slice { + t = t.Elem() + } + if t.Kind() == reflect.Struct { + return t + } + return nil +} + +// applyOverrides walks each dotted path into the parsed YAML map and sets the +// leaf to the RHS parsed as a YAML scalar. Intermediate maps are auto-created so +// an override can add a field the YAML omits; the later re-decode rejects paths +// absent from the schema. Changes are logged to stderr to keep JSON stdout clean. +func applyOverrides(ctx context.Context, m map[string]any, entries []overrideEntry) error { + for _, e := range entries { + var value any + if err := yaml.Unmarshal([]byte(e.raw), &value); err != nil { + return fmt.Errorf("invalid --override %q: cannot parse value %q: %w", e.path, e.raw, err) + } + + parts := strings.Split(e.path, ".") + current := m + for _, part := range parts[:len(parts)-1] { + next, ok := current[part].(map[string]any) + if !ok { + next = map[string]any{} + current[part] = next + } + current = next + } + + leaf := parts[len(parts)-1] + old, had := current[leaf] + current[leaf] = value + if had { + logOverride(ctx, fmt.Sprintf("Override: changing %s from %v to %v", e.path, old, value)) + } else { + logOverride(ctx, fmt.Sprintf("Override: setting %s to %v", e.path, value)) + } + } + return nil +} + +// logOverride writes to stderr only when a cmdIO is present; cmdio.LogString +// panics without one, as in non-command callers such as unit tests. +func logOverride(ctx context.Context, msg string) { + if cmdio.HasIO(ctx) { + cmdio.LogString(ctx, msg) + } +} diff --git a/experimental/air/cmd/runconfig_override_test.go b/experimental/air/cmd/runconfig_override_test.go new file mode 100644 index 00000000000..99eab7dd881 --- /dev/null +++ b/experimental/air/cmd/runconfig_override_test.go @@ -0,0 +1,168 @@ +package aircmd + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestParseOverrides(t *testing.T) { + tests := []struct { + name string + in []string + want []overrideEntry + wantErr string + }{ + { + name: "key=value pairs preserve order", + in: []string{"compute.num_accelerators=8", "timeout_minutes=45"}, + want: []overrideEntry{ + {path: "compute.num_accelerators", raw: "8"}, + {path: "timeout_minutes", raw: "45"}, + }, + }, + { + name: "value may contain =", + in: []string{"env_variables.EXPR=a=b"}, + want: []overrideEntry{{path: "env_variables.EXPR", raw: "a=b"}}, + }, + { + name: "key is trimmed", + in: []string{" timeout_minutes = 45"}, + want: []overrideEntry{{path: "timeout_minutes", raw: " 45"}}, + }, + { + name: "missing = is rejected", + in: []string{"compute.num_accelerators"}, + wantErr: `expected KEY=VALUE`, + }, + { + name: "a .yaml token hints at -f", + in: []string{"train.yaml"}, + wantErr: `looks like a config file`, + }, + { + name: "empty key is rejected", + in: []string{"=5"}, + wantErr: `empty key`, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got, err := parseOverrides(tt.in) + if tt.wantErr != "" { + require.ErrorContains(t, err, tt.wantErr) + return + } + require.NoError(t, err) + assert.Equal(t, tt.want, got) + }) + } +} + +func TestValidateOverridePaths(t *testing.T) { + tests := []struct { + name string + path string + wantErr string + }{ + {name: "known top-level field", path: "experiment_name"}, + {name: "known nested field", path: "compute.num_accelerators"}, + {name: "free-form sub-path", path: "env_variables.MY_VAR"}, + {name: "deep free-form sub-path", path: "parameters.model.layers"}, + { + name: "unknown top-level field", + path: "bogus", + wantErr: `"bogus" is not a known field`, + }, + { + name: "unknown nested field", + path: "compute.bogus", + wantErr: `"bogus" is not a known field`, + }, + { + name: "sub-field of a scalar", + path: "command.sub", + wantErr: `"command" is not a nested object`, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := validateOverridePaths([]overrideEntry{{path: tt.path, raw: "x"}}) + if tt.wantErr != "" { + require.ErrorContains(t, err, tt.wantErr) + return + } + require.NoError(t, err) + }) + } +} + +// overrideBaseConfig is a valid 8-GPU config the override tests mutate. +const overrideBaseConfig = `experiment_name: smoke +command: python train.py +compute: + accelerator_type: GPU_8xH100 + num_accelerators: 8 +env_variables: + EXISTING: hello +` + +func TestLoadRunConfigWithOverrides(t *testing.T) { + t.Run("no overrides matches loadRunConfig", func(t *testing.T) { + cfg, err := loadRunConfigWithOverrides(t.Context(), writeConfig(t, overrideBaseConfig), nil) + require.NoError(t, err) + assert.Equal(t, 8, cfg.Compute.NumAccelerators) + }) + + t.Run("typed scalar override is coerced", func(t *testing.T) { + cfg, err := loadRunConfigWithOverrides(t.Context(), writeConfig(t, overrideBaseConfig), []string{"compute.num_accelerators=16"}) + require.NoError(t, err) + assert.Equal(t, 16, cfg.Compute.NumAccelerators) + }) + + t.Run("multiple overrides all apply", func(t *testing.T) { + cfg, err := loadRunConfigWithOverrides(t.Context(), writeConfig(t, overrideBaseConfig), []string{"compute.num_accelerators=16", "timeout_minutes=45"}) + require.NoError(t, err) + assert.Equal(t, 16, cfg.Compute.NumAccelerators) + require.NotNil(t, cfg.TimeoutMinutes) + assert.Equal(t, 45, *cfg.TimeoutMinutes) + }) + + t.Run("free-form env var adds a key as a string", func(t *testing.T) { + cfg, err := loadRunConfigWithOverrides(t.Context(), writeConfig(t, overrideBaseConfig), []string{"env_variables.RANK=0"}) + require.NoError(t, err) + // A numeric-looking value stays a string because env_variables is map[string]string. + assert.Equal(t, "0", cfg.EnvVariables["RANK"]) + }) + + t.Run("intermediate maps are auto-created", func(t *testing.T) { + cfg, err := loadRunConfigWithOverrides(t.Context(), writeConfig(t, overrideBaseConfig), []string{"environment.docker_image.url=my/img:1"}) + require.NoError(t, err) + require.NotNil(t, cfg.Environment) + require.NotNil(t, cfg.Environment.DockerImage) + assert.Equal(t, "my/img:1", cfg.Environment.DockerImage.URL) + }) + + t.Run("unknown path errors before mutation", func(t *testing.T) { + _, err := loadRunConfigWithOverrides(t.Context(), writeConfig(t, overrideBaseConfig), []string{"bogus=1"}) + require.ErrorContains(t, err, `"bogus" is not a known field`) + }) + + t.Run("semantic validation runs after override", func(t *testing.T) { + // 3 is a known field with a valid type, so only validate() can reject it. + _, err := loadRunConfigWithOverrides(t.Context(), writeConfig(t, overrideBaseConfig), []string{"compute.num_accelerators=3"}) + require.ErrorContains(t, err, "must be a multiple of 8") + }) + + t.Run("type mismatch is rejected on re-decode", func(t *testing.T) { + _, err := loadRunConfigWithOverrides(t.Context(), writeConfig(t, overrideBaseConfig), []string{"compute.num_accelerators=abc"}) + require.Error(t, err) + }) + + t.Run("malformed override is rejected", func(t *testing.T) { + _, err := loadRunConfigWithOverrides(t.Context(), writeConfig(t, overrideBaseConfig), []string{"compute.num_accelerators"}) + require.ErrorContains(t, err, "expected KEY=VALUE") + }) +} diff --git a/experimental/air/cmd/runsubmit.go b/experimental/air/cmd/runsubmit.go index 8c0be55260d..d29bea427ba 100644 --- a/experimental/air/cmd/runsubmit.go +++ b/experimental/air/cmd/runsubmit.go @@ -58,15 +58,6 @@ func buildSubmitPayload(cfg *runConfig, commandPath, dlImage string, snap snapsh }, }}, 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 @@ -163,13 +154,13 @@ func submitWorkload(ctx context.Context, w *databricks.WorkspaceClient, cfg *run 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. + // Package and upload the code snapshot, if any, via DABs' artifact-upload + // plumbing; the remote code_source_path rides the ai_runtime_task. A run with no + // code_source leaves it empty. Snapshot is the only code_source type. var snap snapshotResult if cfg.CodeSource != nil && cfg.CodeSource.Snapshot != nil { - snap, err = snapshotCodeSource(ctx, w, cfg.CodeSource.Snapshot, configPath, base, funcDir) + // Sidecars land in the run's launch dir (funcDir) via fc, next to command.sh. + snap, err = snapshotViaDABsUpload(ctx, w, cfg.CodeSource.Snapshot, configPath, fc, funcDir) if err != nil { return 0, "", err } diff --git a/experimental/air/cmd/runsubmit_test.go b/experimental/air/cmd/runsubmit_test.go index fd5103599df..6418eed5ffe 100644 --- a/experimental/air/cmd/runsubmit_test.go +++ b/experimental/air/cmd/runsubmit_test.go @@ -2,10 +2,14 @@ package aircmd import ( "encoding/json" + "io" + "path" "path/filepath" "strings" "testing" + "github.com/databricks/cli/libs/cmdio" + "github.com/databricks/cli/libs/filer" "github.com/databricks/cli/libs/testserver" "github.com/databricks/databricks-sdk-go" "github.com/databricks/databricks-sdk-go/service/jobs" @@ -156,8 +160,41 @@ func TestSubmitWorkload(t *testing.T) { 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. +// TestSubmitWorkloadHonorsOverride proves a --override reaches the actual +// runs/submit payload on a real submit, not just dry-run validation: the config +// pins num_accelerators=1, the override bumps it to 4, and the recorded request +// body must carry 4. +func TestSubmitWorkloadHonorsOverride(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 jobs/runs/submit 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 := loadRunConfigWithOverrides(t.Context(), cfgPath, []string{"compute.num_accelerators=4"}) + require.NoError(t, err) + + _, _, err = submitWorkload(t.Context(), w, cfg, cfgPath, "idem-key") + require.NoError(t, err) + + require.Len(t, got.Tasks, 1) + at := got.Tasks[0].AiRuntimeTask + require.NotNil(t, at) + require.Len(t, at.Deployments, 1) + assert.Equal(t, 4, at.Deployments[0].Compute.AcceleratorCount) +} + +// A working-tree code_source is packaged into a tarball, uploaded via DABs' artifact +// plumbing, and its remote code_source_path attached to the submitted task. func TestSubmitWorkloadWithCodeSource(t *testing.T) { server := testserver.New(t) t.Cleanup(server.Close) @@ -172,6 +209,47 @@ func TestSubmitWorkloadWithCodeSource(t *testing.T) { w, err := databricks.NewWorkspaceClient(&databricks.Config{Host: server.URL, Token: "token"}) require.NoError(t, err) + // A plain working-tree directory: packaging is plain-tar. + repo := filepath.Join(t.TempDir(), "src") + writeRepoFile(t, repo, "train.py", "print()") + + cfg := minimalConfig + ` +code_source: + type: snapshot + snapshot: + root_path: ` + repo + ` +` + cfgPath := writeConfigFile(t, "run.yaml", cfg) + loaded, err := loadRunConfig(cfgPath) + require.NoError(t, err) + + // The DABs upload path logs via cmdio; the real `air run` context carries it. + ctx := cmdio.MockDiscard(t.Context()) + _, _, err = submitWorkload(ctx, w, loaded, cfgPath, "idem") + require.NoError(t, err) + + at := got.Tasks[0].AiRuntimeTask + // The tarball is uploaded to the artifact .internal dir and code_source_path + // rewritten to it. + assert.Contains(t, at.CodeSourcePath, "/.air/repo_snapshots/.internal/") + assert.True(t, strings.HasSuffix(at.CodeSourcePath, ".tar.gz"), at.CodeSourcePath) +} + +// A git-pinned code_source is git-archived at the commit, uploaded via DABs' artifact +// plumbing, and its remote code_source_path attached to the submitted task. +func TestSubmitWorkloadWithGitPinnedCodeSource(t *testing.T) { + server := testserver.New(t) + t.Cleanup(server.Close) + + 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()") @@ -189,15 +267,209 @@ code_source: loaded, err := loadRunConfig(cfgPath) require.NoError(t, err) - _, _, err = submitWorkload(t.Context(), w, loaded, cfgPath, "idem") + ctx := cmdio.MockDiscard(t.Context()) + _, _, err = submitWorkload(ctx, w, loaded, cfgPath, "idem") + require.NoError(t, err) + + at := got.Tasks[0].AiRuntimeTask + assert.Contains(t, at.CodeSourcePath, "/.air/repo_snapshots/.internal/") + assert.True(t, strings.HasSuffix(at.CodeSourcePath, ".tar.gz"), at.CodeSourcePath) +} + +// testSidecarStore builds a workspace filer + base path standing in for the run's +// launch dir, where snapshotViaDABsUpload writes git provenance sidecars. +func testSidecarStore(t *testing.T, w *databricks.WorkspaceClient) (filer.Filer, string) { + t.Helper() + base := "/Workspace/Users/tester@databricks.com/.air/cli_launch/test" + f, err := filer.NewWorkspaceFilesClient(w, base) + require.NoError(t, err) + return f, base +} + +// A plain-tar (working-tree) snapshot is uploaded under a unique, timestamped name so +// two concurrent submissions of the same root_path don't clobber each other's upload. +func TestSubmitWorkloadPlainTarNameIsUnique(t *testing.T) { + server := testserver.New(t) + t.Cleanup(server.Close) + + server.Handle("POST", "/api/2.2/jobs/runs/submit", func(req testserver.Request) any { + return jobs.SubmitRunResponse{RunId: 555} + }) + testserver.AddDefaultHandlers(server) + w, err := databricks.NewWorkspaceClient(&databricks.Config{Host: server.URL, Token: "token"}) + require.NoError(t, err) + + // A plain working-tree directory named "src": the old code named the tarball + // after the dir alone (src.tar.gz), so any two submissions collided. + repo := filepath.Join(t.TempDir(), "src") + writeRepoFile(t, repo, "train.py", "print()") + + cfg := minimalConfig + ` +code_source: + type: snapshot + snapshot: + root_path: ` + repo + ` +` + cfgPath := writeConfigFile(t, "run.yaml", cfg) + loaded, err := loadRunConfig(cfgPath) + require.NoError(t, err) + + // The uploaded name carries a discriminator (timestamp), not the bare dir name. + ctx := cmdio.MockDiscard(t.Context()) + sidecarStore, sidecarBase := testSidecarStore(t, w) + snap, err := snapshotViaDABsUpload(ctx, w, loaded.CodeSource.Snapshot, cfgPath, sidecarStore, sidecarBase) + require.NoError(t, err) + base := path.Base(snap.CodeSourcePath) + assert.NotEqual(t, "src.tar.gz", base, "plain-tar name must be unique, not the bare dir name") + assert.Regexp(t, `^src_\d{8}_\d{6}\.tar\.gz$`, base) +} + +// A git_archive snapshot is content-addressed by (commit, include_paths): submitting +// the same commit twice reuses the already-uploaded tarball and skips the second +// upload (cache hit), while resolving to the identical remote path. +func TestSubmitWorkloadGitArchiveCaching(t *testing.T) { + server := testserver.New(t) + t.Cleanup(server.Close) + + server.Handle("POST", "/api/2.2/jobs/runs/submit", func(req testserver.Request) any { + return jobs.SubmitRunResponse{RunId: 555} + }) + // Track which snapshot tarballs get uploaded, preserving fake-workspace + // persistence so the second submit's cache-existence Stat sees the first upload. + // Dedupe by path: the DABs uploader mkdirs-and-retries the import on a missing + // parent dir, so one logical upload can hit this route more than once. + uploaded := map[string]bool{} + server.Handle("POST", "/api/2.0/workspace-files/import-file/{path...}", func(req testserver.Request) any { + p := req.Vars["path"] + if strings.Contains(p, "/.air/repo_snapshots/") { + uploaded[p] = true + } + return req.Workspace.WorkspaceFilesImportFile(p, req.Body, req.URL.Query().Get("overwrite") == "true") + }) + testserver.AddDefaultHandlers(server) + w, err := databricks.NewWorkspaceClient(&databricks.Config{Host: server.URL, Token: "token"}) + require.NoError(t, err) + + 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) + + ctx := cmdio.MockDiscard(t.Context()) + sidecarStore, sidecarBase := testSidecarStore(t, w) + first, err := snapshotViaDABsUpload(ctx, w, loaded.CodeSource.Snapshot, cfgPath, sidecarStore, sidecarBase) + require.NoError(t, err) + second, err := snapshotViaDABsUpload(ctx, w, loaded.CodeSource.Snapshot, cfgPath, sidecarStore, sidecarBase) + require.NoError(t, err) + + // Same pinned commit → identical content-addressed remote path, uploaded once + // (the second submit is a cache hit and moves no bytes). + assert.Equal(t, first.CodeSourcePath, second.CodeSourcePath) + assert.Len(t, uploaded, 1, "git_archive cache hit should skip the second upload") +} + +// A git code_source also uploads git provenance sidecars (git_state.json, and +// git_diff.patch when the tree is dirty) next to the run's launch dir, so the +// submitted commit + working-tree diff are inspectable. +func TestSubmitWorkloadUploadsGitSidecars(t *testing.T) { + server := testserver.New(t) + t.Cleanup(server.Close) + + server.Handle("POST", "/api/2.2/jobs/runs/submit", func(req testserver.Request) any { + return jobs.SubmitRunResponse{RunId: 555} + }) + testserver.AddDefaultHandlers(server) + w, err := databricks.NewWorkspaceClient(&databricks.Config{Host: server.URL, Token: "token"}) + require.NoError(t, err) + + // Commit, then dirty the tree so both git_state.json and git_diff.patch are produced. + repo := newTestRepo(t) + writeRepoFile(t, repo, "train.py", "print()") + commitAll(t, repo, "init") + writeRepoFile(t, repo, "train.py", "print('dirty')") + + cfg := minimalConfig + ` +code_source: + type: snapshot + snapshot: + root_path: ` + repo + ` +` + cfgPath := writeConfigFile(t, "run.yaml", cfg) + loaded, err := loadRunConfig(cfgPath) + require.NoError(t, err) + + ctx := cmdio.MockDiscard(t.Context()) + sidecarStore, sidecarBase := testSidecarStore(t, w) + snap, err := snapshotViaDABsUpload(ctx, w, loaded.CodeSource.Snapshot, cfgPath, sidecarStore, sidecarBase) + require.NoError(t, err) + + // Both sidecars are reported under the launch dir and actually exist there. + assert.Equal(t, path.Join(sidecarBase, gitStateName), snap.GitStatePath) + assert.Equal(t, path.Join(sidecarBase, gitDiffName), snap.GitDiffPath) + + r, err := sidecarStore.Read(ctx, gitStateName) + require.NoError(t, err) + stateBytes, err := io.ReadAll(r) + require.NoError(t, err) + var state map[string]any + require.NoError(t, json.Unmarshal(stateBytes, &state)) + assert.Equal(t, "plain_tar", state["packaging_mode"]) + assert.Equal(t, true, state["dirty"]) + assert.Equal(t, "captured", state["diff_status"]) +} + +// remote_volume uploads the snapshot to a UC Volume: DABs' artifact uploader handles +// /Volumes destinations natively, so code_source_path lands under the Volume path. +func TestSubmitWorkloadWithRemoteVolumeCodeSource(t *testing.T) { + server := testserver.New(t) + t.Cleanup(server.Close) + + 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} + }) + // Stub the UC Volume file write: the fake server's default handler 404s when the + // parent dir is absent (no auto-mkdir), so accept the PUT to exercise the Volume + // upload route. This asserts we route to /api/2.0/fs/files/Volumes/... at all. + server.Handle("PUT", "/api/2.0/fs/files/Volumes/{path...}", func(req testserver.Request) any { + return testserver.Response{StatusCode: 204} + }) + testserver.AddDefaultHandlers(server) + w, err := databricks.NewWorkspaceClient(&databricks.Config{Host: server.URL, Token: "token"}) + require.NoError(t, err) + + repo := filepath.Join(t.TempDir(), "src") + writeRepoFile(t, repo, "train.py", "print()") + + cfg := minimalConfig + ` +code_source: + type: snapshot + snapshot: + root_path: ` + repo + ` + remote_volume: /Volumes/main/default/code +` + cfgPath := writeConfigFile(t, "run.yaml", cfg) + loaded, err := loadRunConfig(cfgPath) + require.NoError(t, err) + + ctx := cmdio.MockDiscard(t.Context()) + _, _, err = submitWorkload(ctx, 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.Contains(t, at.CodeSourcePath, "/Volumes/main/default/code/.internal/") assert.True(t, strings.HasSuffix(at.CodeSourcePath, ".tar.gz"), at.CodeSourcePath) } diff --git a/experimental/air/cmd/snapshot.go b/experimental/air/cmd/snapshot.go index 59041f5986d..603b7e6c0f8 100644 --- a/experimental/air/cmd/snapshot.go +++ b/experimental/air/cmd/snapshot.go @@ -1,58 +1,28 @@ 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). +// snapshotResult holds the code_source_path wired into the submit payload (the +// uploaded code archive's remote path) plus the remote paths of the best-effort git +// provenance sidecars (empty when not a git repo or upload failed). 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(ctx, 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 +43,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 +58,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(ctx context.Context, 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(ctx, 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_dabs.go b/experimental/air/cmd/snapshot_dabs.go new file mode 100644 index 00000000000..9becc5e91a2 --- /dev/null +++ b/experimental/air/cmd/snapshot_dabs.go @@ -0,0 +1,278 @@ +package aircmd + +import ( + "bytes" + "context" + "errors" + "fmt" + "io/fs" + "os" + "path" + "path/filepath" + "time" + + "github.com/databricks/cli/bundle" + "github.com/databricks/cli/bundle/config" + "github.com/databricks/cli/bundle/config/resources" + "github.com/databricks/cli/bundle/libraries" + "github.com/databricks/cli/libs/dyn" + "github.com/databricks/cli/libs/filer" + "github.com/databricks/cli/libs/log" + "github.com/databricks/cli/libs/vfs" + "github.com/databricks/databricks-sdk-go" + "github.com/databricks/databricks-sdk-go/service/jobs" +) + +// snapshotViaDABsUpload packages the code_source into a tarball and uploads it using +// DABs' artifact-upload plumbing (the same path a bundle uses for a file-valued +// code_source_path), returning the remote path to attach to the ai_runtime_task. +// +// The packaging + upload logic is CLI-owned (this file, OWNERS = us); it only reuses +// DABs' libraries.ReplaceWithRemotePath + libraries.Upload as the uploader so we do +// not reimplement workspace/volume upload. A minimal in-memory bundle carries the +// local tarball path as code_source_path; ReplaceWithRemotePath rewrites it to the +// artifact .internal path and Upload pushes the bytes. +func snapshotViaDABsUpload(ctx context.Context, w *databricks.WorkspaceClient, snap *snapshotSourceConfig, configPath string, sidecarStore filer.Filer, sidecarBase string) (snapshotResult, error) { + repoPath, err := resolveRootPath(ctx, snap.RootPath, filepath.Dir(configPath)) + if err != nil { + return snapshotResult{}, err + } + + // Resolve how to package before touching the tarball: git_archive (pinned commit, + // cacheable) vs plain_tar (working tree, not cacheable). + plan, err := resolveSnapshotPlan(ctx, newGitRepo(repoPath), snap.Git, snap.IncludePaths) + if err != nil { + return snapshotResult{}, err + } + + // remote_volume, when set, is a UC Volume path; DABs' artifact uploader handles + // /Volumes destinations natively (GetFilerForLibraries → filerForVolume). + remoteVolume := "" + if snap.RemoteVolume != nil { + remoteVolume = *snap.RemoteVolume + } + result, err := uploadSnapshotViaDABs(ctx, w, repoPath, plan, remoteVolume) + if err != nil { + return snapshotResult{}, err + } + + // Upload git provenance sidecars (git_state.json / git_diff.patch) next to the + // run's launch dir so the submitted commit + working-tree diff are inspectable. + // Best-effort and git-only: any failure logs and leaves the paths empty rather + // than failing an otherwise-valid submission. + // + // The sidecars are deliberately NOT bundled into the code tarball. The git_archive + // tarball is content-addressed and cached by (commit, include_paths), so a second + // run at the same commit reuses it; but the sidecars vary per run (git_state's + // timestamp, and git_diff captures the working tree at submit time). Folding them + // in would force a distinct tarball per run (defeating the cache) or serve a prior + // run's stale provenance on a cache hit. They also live in the per-run launch dir, + // not the shared artifact dir, so they don't accumulate. Keep them out of the tar. + if plan.isGitRepo { + result.GitStatePath, result.GitDiffPath = uploadSnapshotSidecars(ctx, sidecarStore, sidecarBase, newGitRepo(repoPath), plan) + } + return result, nil +} + +// uploadSnapshotSidecars writes the git_state.json provenance record — and, when the +// working tree is dirty, a captured git_diff.patch — into the run's launch dir via +// sidecarStore (rooted at sidecarBase, used only to report absolute paths). It is +// best-effort: every failure logs a warning and yields an empty path, never an error, +// so provenance capture cannot fail a submission. +func uploadSnapshotSidecars(ctx context.Context, sidecarStore filer.Filer, sidecarBase string, 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 := 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(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 := 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(sidecarBase, gitStateName), diffPath +} + +// snapshotTarballName is the uploaded filename for the snapshot. It is deterministic +// for git_archive — _.tar.gz keyed on (commit, include_paths) — so +// an identical commit reuses the same remote object (see the cache check below). For +// plain_tar it is timestamped so concurrent submissions of the same directory don't +// clobber each other's upload (working-tree content isn't pinned to a SHA, so it +// can't be content-addressed). +func snapshotTarballName(plan snapshotPlan, dirName string) string { + if plan.mode == modeGitArchive { + key := computeSnapshotCacheKey(plan.commitSHA, plan.includePaths) + return fmt.Sprintf("%s_%s.tar.gz", dirName, key[:16]) + } + return fmt.Sprintf("%s_%s.tar.gz", dirName, time.Now().UTC().Format("20060102_150405")) +} + +// packageSnapshot writes the snapshot to tarball per the resolved plan: `git archive` +// of the pinned commit for git_archive, else a plain tar of the working tree. +func packageSnapshot(ctx context.Context, repoPath string, plan snapshotPlan, tarball string) error { + dirName := filepath.Base(repoPath) + if plan.mode == modeGitArchive { + return createGitArchiveSnapshot(ctx, newGitRepo(repoPath), plan.commitSHA, tarball, dirName, plan.includePaths) + } + return createPlainTarball(ctx, repoPath, tarball, plan.includePaths) +} + +// uploadSnapshotViaDABs uploads the snapshot through DABs' artifact-upload machinery +// and returns its remote code_source_path. It builds a minimal bundle whose only +// artifact is the tarball (as a file-valued code_source_path), rewrites the field to +// the remote .internal path, and uploads the bytes. When remoteVolume is set the +// tarball goes to that UC Volume; otherwise to the user's repo_snapshots dir. +// +// git_archive snapshots are cacheable: the tarball name is content-addressed by +// (commit, include_paths), so if the identical object is already uploaded we skip +// packaging and upload entirely and just reuse the remote path. +func uploadSnapshotViaDABs(ctx context.Context, w *databricks.WorkspaceClient, repoPath string, plan snapshotPlan, remoteVolume string) (snapshotResult, error) { + // artifactPath is where DABs uploads the tarball; GetFilerForLibraries routes to + // a Workspace or Volume filer based on its prefix, then appends /.internal. + artifactPath := remoteVolume + if artifactPath == "" { + base, err := userWorkspaceDir(ctx, w) + if err != nil { + return snapshotResult{}, err + } + // The user's repo_snapshots dir (not the default bundle artifact_path, which a + // deploy would clean up). + artifactPath = path.Join(base, ".air", "repo_snapshots") + } + + tmp, err := os.MkdirTemp("", "air-snapshot-*") + if err != nil { + return snapshotResult{}, err + } + defer os.RemoveAll(tmp) + + tarName := snapshotTarballName(plan, filepath.Base(repoPath)) + + b := &bundle.Bundle{ + BundleRootPath: tmp, + BundleRoot: vfs.MustNew(tmp), + SyncRootPath: tmp, + SyncRoot: vfs.MustNew(tmp), + Config: config.Root{ + Bundle: config.Bundle{Target: "default"}, + Workspace: config.Workspace{ArtifactPath: artifactPath}, + Resources: config.Resources{ + Jobs: map[string]*resources.Job{ + "air": { + JobSettings: jobs.JobSettings{ + Tasks: []jobs.Task{{ + TaskKey: "air", + // Relative to SyncRootPath (the temp dir); collectLocalLibraries + // joins it back and uploads the file. + AiRuntimeTask: &jobs.AiRuntimeTask{CodeSourcePath: tarName}, + }}, + }, + }, + }, + }, + }, + } + b.SetWorkpaceClient(w) + if err := b.Config.Mutate(func(v dyn.Value) (dyn.Value, error) { return v, nil }); err != nil { + return snapshotResult{}, err + } + + // git_archive is cacheable by (commit, include_paths): if the identical tarball is + // already uploaded, skip packaging + upload and reuse it. Only the config-path + // rewrite (ReplaceWithRemotePath) runs — no bytes move. + if plan.mode == modeGitArchive { + f, uploadPath, diags := libraries.GetFilerForLibraries(ctx, b) + if diags.HasError() { + return snapshotResult{}, diags.Error() + } + exists, err := snapshotExists(ctx, f, tarName) + if err != nil { + return snapshotResult{}, err + } + if exists { + if _, diags := libraries.ReplaceWithRemotePath(ctx, b); diags.HasError() { + return snapshotResult{}, diags.Error() + } + remote, err := readCodeSourcePath(b) + if err != nil { + return snapshotResult{}, err + } + log.Debugf(ctx, "snapshot cache hit for %s at %s", shortSHA(plan.commitSHA), path.Join(uploadPath, tarName)) + return snapshotResult{CodeSourcePath: remote}, nil + } + } + + // Cache miss (or plain_tar): package the tarball locally, then upload the bytes. + if err := packageSnapshot(ctx, repoPath, plan, filepath.Join(tmp, tarName)); err != nil { + return snapshotResult{}, err + } + + libs, diags := libraries.ReplaceWithRemotePath(ctx, b) + if diags.HasError() { + return snapshotResult{}, diags.Error() + } + if diags := bundle.Apply(ctx, b, libraries.Upload(libs)); diags.HasError() { + return snapshotResult{}, diags.Error() + } + + remote, err := readCodeSourcePath(b) + if err != nil { + return snapshotResult{}, err + } + return snapshotResult{CodeSourcePath: remote}, nil +} + +// snapshotExists reports whether name already exists in the artifact store, used to +// short-circuit a cacheable git_archive upload. A not-found is a clean miss (false, +// nil); any other error is surfaced. +func snapshotExists(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) +} + +// readCodeSourcePath returns the (rewritten) code_source_path from the bundle config. +func readCodeSourcePath(b *bundle.Bundle) (string, error) { + v, err := dyn.GetByPath(b.Config.Value(), + dyn.MustPathFromString("resources.jobs.air.tasks[0].ai_runtime_task.code_source_path")) + if err != nil { + return "", fmt.Errorf("code snapshot was not packaged: %w", err) + } + s, ok := v.AsString() + if !ok { + return "", fmt.Errorf("unexpected code_source_path value %v", v.AsAny()) + } + return s, nil +} diff --git a/experimental/air/cmd/snapshot_git.go b/experimental/air/cmd/snapshot_git.go index 616b3049f74..70dd34f4da9 100644 --- a/experimental/air/cmd/snapshot_git.go +++ b/experimental/air/cmd/snapshot_git.go @@ -200,6 +200,13 @@ func shortSHA(sha string) string { // coordination with the backend reader. const snapshotStateSchemaVersion = 1 +// gitStateName and gitDiffName are the git provenance sidecar basenames, uploaded +// next to the code snapshot for human/agent inspection of what was submitted. +const ( + gitStateName = "git_state.json" + gitDiffName = "git_diff.patch" +) + // defaultRemoteName is the remote consulted for merge-base and repo URL (local refs // only — the remote-fetch path is gone). const defaultRemoteName = "origin" diff --git a/experimental/air/cmd/snapshot_package.go b/experimental/air/cmd/snapshot_package.go index 672366086c9..9c1fa2714d9 100644 --- a/experimental/air/cmd/snapshot_package.go +++ b/experimental/air/cmd/snapshot_package.go @@ -19,8 +19,6 @@ import ( // `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", @@ -42,9 +40,24 @@ func createGitArchiveSnapshot(ctx context.Context, git gitRepo, commitSHA, outpu // excluded; a .gitignore at repoPath is honored. func createPlainTarball(ctx context.Context, repoPath, outputTarball string, includePaths []string) error { dirName := filepath.Base(repoPath) - parent := filepath.Dir(repoPath) + // Absolute so it resolves correctly regardless of tar's working dir (set below). + parent, err := filepath.Abs(filepath.Dir(repoPath)) + if err != nil { + return err + } + + // Pass the archive path relative to its own directory (run tar there), never a + // full path: on Windows an absolute path like `C:\out\x.tar.gz` makes tar read + // the `C:` as a remote host ("Cannot connect to C:"), since tar treats a colon + // in the -f arg as host:path. A bare basename with -C avoids that on GNU tar and + // bsdtar alike. + outDirAbs, err := filepath.Abs(filepath.Dir(outputTarball)) + if err != nil { + return err + } + outName := filepath.Base(outputTarball) - args := []string{"-czf", outputTarball} + args := []string{"-czf", outName} // Exclude macOS AppleDouble files: they sort before the real top-level dir and // hijack a remote `head -1` parse. No-op on Linux. @@ -68,7 +81,9 @@ func createPlainTarball(ctx context.Context, repoPath, outputTarball string, inc } // Archive from the parent so the directory name is preserved; with include_paths, - // prefix each so entries nest under it (matching git archive --prefix). + // prefix each so entries nest under it (matching git archive --prefix). -C only + // affects the file operands that follow it, not the -f archive path (which + // resolves against tar's working dir, set to outDir below). args = append(args, "-C", parent) if len(includePaths) > 0 { for _, p := range includePaths { @@ -79,6 +94,8 @@ func createPlainTarball(ctx context.Context, repoPath, outputTarball string, inc } cmd := exec.CommandContext(ctx, "tar", args...) + // Run tar in the output directory so the bare -f basename lands there. + cmd.Dir = outDirAbs var stderr bytes.Buffer cmd.Stderr = &stderr if err := cmd.Run(); err != nil { diff --git a/experimental/air/cmd/snapshot_package_test.go b/experimental/air/cmd/snapshot_package_test.go index d895d59b98e..22561f4e2ce 100644 --- a/experimental/air/cmd/snapshot_package_test.go +++ b/experimental/air/cmd/snapshot_package_test.go @@ -49,9 +49,8 @@ func TestCreateGitArchiveSnapshot(t *testing.T) { 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. + // Every real entry is prefixed with the directory name. git archive also emits a + // `pax_global_header` pseudo-entry (no prefix) that tar ignores on extraction. assert.Contains(t, entries, dirName+"/a.txt") assert.Contains(t, entries, dirName+"/src/model.py") for _, e := range entries { @@ -75,18 +74,17 @@ func TestCreateGitArchiveSnapshot_IncludePaths(t *testing.T) { 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) + repo := t.TempDir() 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") + // A .git dir must never be shipped. + writeRepoFile(t, repo, ".git/config", "x") out := filepath.Join(t.TempDir(), "snap.tar.gz") require.NoError(t, createPlainTarball(ctx, repo, out, nil)) @@ -103,7 +101,7 @@ func TestCreatePlainTarball(t *testing.T) { func TestCreatePlainTarball_HonorsGitignore(t *testing.T) { ctx := t.Context() - repo := newTestRepo(t) + repo := t.TempDir() writeRepoFile(t, repo, "keep.txt", "1") writeRepoFile(t, repo, "junk.log", "noise") writeRepoFile(t, repo, ".gitignore", "*.log\n") @@ -119,7 +117,7 @@ func TestCreatePlainTarball_HonorsGitignore(t *testing.T) { func TestCreatePlainTarball_IncludePaths(t *testing.T) { ctx := t.Context() - repo := newTestRepo(t) + repo := t.TempDir() writeRepoFile(t, repo, "a.txt", "1") writeRepoFile(t, repo, "src/model.py", "print()") diff --git a/experimental/air/cmd/snapshot_test.go b/experimental/air/cmd/snapshot_test.go deleted file mode 100644 index d94fe005fc9..00000000000 --- a/experimental/air/cmd/snapshot_test.go +++ /dev/null @@ -1,155 +0,0 @@ -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" -) - -func TestResolveRootPath(t *testing.T) { - ctx := t.Context() - dir := t.TempDir() - 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). - got, err := resolveRootPath(ctx, ".", filepath.Join(dir, "proj")) - require.NoError(t, err) - assert.True(t, filepath.IsAbs(got)) - assert.Equal(t, "proj", filepath.Base(got)) - - // A relative subpath resolves against configDir and keeps its own basename. - require.NoError(t, os.MkdirAll(filepath.Join(dir, "proj", "sub"), 0o755)) - got, err = resolveRootPath(ctx, "sub", filepath.Join(dir, "proj")) - require.NoError(t, err) - assert.Equal(t, "sub", filepath.Base(got)) - - // A non-existent path errors. - _, 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...) -} diff --git a/experimental/air/cmd/stubs_test.go b/experimental/air/cmd/stubs_test.go index e28d7f66730..4007557ef78 100644 --- a/experimental/air/cmd/stubs_test.go +++ b/experimental/air/cmd/stubs_test.go @@ -13,7 +13,6 @@ import ( // fails with a "not implemented" error. Drop a command here once it lands. func TestStubCommandsReturnNotImplemented(t *testing.T) { stubs := map[string]*cobra.Command{ - "logs": newLogsCommand(), "register-image": newRegisterImageCommand(), } diff --git a/libs/cmdio/io.go b/libs/cmdio/io.go index 2b9b395ce13..e57c90974b4 100644 --- a/libs/cmdio/io.go +++ b/libs/cmdio/io.go @@ -83,15 +83,6 @@ func IsPromptSupported(ctx context.Context) bool { return c.capabilities.SupportsPrompt() } -// IsPagerSupported reports whether stdin, stdout, and stderr are all interactive -// terminals. This is the requirement for a full-screen or navigable output -// program: unlike IsPromptSupported it also checks stdout, so it returns false -// when stdout is piped or redirected. -func IsPagerSupported(ctx context.Context) bool { - c := fromContext(ctx) - return c.capabilities.SupportsPager() -} - // SupportsColor returns true if the given writer supports colored output. // This checks both TTY status and environment variables (NO_COLOR, TERM=dumb). func SupportsColor(ctx context.Context, w io.Writer) bool {