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/unimplemented/out.test.toml b/acceptance/experimental/air/config-help/out.test.toml similarity index 100% rename from acceptance/experimental/air/unimplemented/out.test.toml rename to acceptance/experimental/air/config-help/out.test.toml diff --git a/acceptance/experimental/air/config-help/output.txt b/acceptance/experimental/air/config-help/output.txt new file mode 100644 index 00000000000..2d7dffc135d --- /dev/null +++ b/acceptance/experimental/air/config-help/output.txt @@ -0,0 +1,110 @@ + +=== command help +>>> [CLI] experimental air run --help +Submit a training workload to Databricks serverless GPU compute. + +The workload is described by a YAML config file (see --file). + +To look up a config field, pass its path to -h: + + databricks experimental air run -h config + databricks experimental air run -h config.compute + databricks experimental air run -h config.compute.accelerator_type + +The path must be a separate argument: cobra reserves -h as a boolean, so +-h=config.compute and -hconfig.compute are not accepted. + +Usage: + databricks experimental air run [flags] + +Flags: + --dry-run Validate the config without submitting + -f, --file string Path to the workload YAML config + -h, --help help for run + --idempotency-key string Return the existing run if this key was already used + --override stringArray Override a YAML field, e.g. compute.num_accelerators=8 (repeatable) + --watch Stream logs until the run completes + +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) + +=== schema overview +>>> [CLI] experimental air run -h config +config + The run YAML schema. Pass a field path for details, e.g. config.compute.accelerator_type. + + Fields: + experiment_name (required) Name of the experiment. + compute (required) Which accelerators to run on and how many. + environment Python dependencies, or a custom Docker image, for the run's runtime. + command (required) Shell command that starts the workload. + env_variables Plain environment variables, as NAME: value. + secrets Environment variables sourced from secrets, as NAME: scope/key. + code_source Local code to upload and make available to the run. + max_retries How many times to retry a failed run. + timeout_minutes Wall-clock limit for the run in minutes. + idempotency_token Reuse token: a repeat submission with the same token returns the existing run instead of starting another. + parameters Free-form values passed through to the workload. + mlflow_run_name Name for the MLflow run. + mlflow_experiment_directory Workspace directory holding the MLflow experiment. + permissions Who may view or manage the run, as a list of principal plus level grants. + usage_policy_name Usage policy to bill the run to, by name. + usage_policy_id Usage policy to bill the run to, by id. + +Use "-h config." for details on a field. + +=== nested object lists its fields +>>> [CLI] experimental air run -h config.compute +config.compute + Which accelerators to run on and how many. + + Fields: + num_accelerators Total number of GPUs to allocate. + accelerator_type Which accelerator to run on, e.g. GPU_1xA10. + +Use "-h config.compute." for details on a field. + +=== leaf field +>>> [CLI] experimental air run -h config.compute.accelerator_type +config.compute.accelerator_type + Which accelerator to run on, e.g. GPU_1xA10. See https://docs.databricks.com/aws/en/machine-learning/ai-runtime/cli/yaml-config#reference for the current list of supported GPU types. Matched case-sensitively. + + Type: string + Required: no + +=== conditionally required leaf +>>> [CLI] experimental air run -h config.environment.docker_image.url +config.environment.docker_image.url + Fully qualified image URL, e.g. myregistry.io/team/train:v3. + + Type: string + Required: when environment.docker_image is set + +=== union field reports both accepted shapes +>>> [CLI] experimental air run -h config.code_source.snapshot.git.remote +config.code_source.snapshot.git.remote + No longer supported: the snapshot archives your local copy. Only false is accepted; use commit to pin a revision. + + Type: bool or string + Required: no + +=== the config. prefix is optional +>>> [CLI] experimental air run -h compute.num_accelerators +config.compute.num_accelerators + Total number of GPUs to allocate. Must be a positive multiple of the accelerator type's per-node GPU count. See https://docs.databricks.com/aws/en/machine-learning/ai-runtime/cli/yaml-config#reference for supported GPU types. + + Type: int + Required: no + +=== unknown field suggests the near match +>>> [CLI] experimental air run -h config.compute.acclerator_type +Error: unknown config field "config.compute.acclerator_type"; did you mean "accelerator_type"? + +fields under "config.compute" are: accelerator_type, num_accelerators + +=== free-form map keys are not schema fields +>>> [CLI] experimental air run -h config.parameters.learning_rate +Error: "config.parameters" holds user-defined keys, so "learning_rate" is not part of the schema; see "config.parameters" instead diff --git a/acceptance/experimental/air/config-help/script b/acceptance/experimental/air/config-help/script new file mode 100644 index 00000000000..d1c25be1100 --- /dev/null +++ b/acceptance/experimental/air/config-help/script @@ -0,0 +1,31 @@ +# Pin the rendered config field help. Any change to a `help:` struct tag on the +# run YAML schema shows up as a diff here. + +# Plain -h must keep documenting the command itself, and advertise the config +# path syntax so the feature is discoverable. +title "command help" +trace $CLI experimental air run --help + +title "schema overview" +trace $CLI experimental air run -h config + +title "nested object lists its fields" +trace $CLI experimental air run -h config.compute + +title "leaf field" +trace $CLI experimental air run -h config.compute.accelerator_type + +title "conditionally required leaf" +trace $CLI experimental air run -h config.environment.docker_image.url + +title "union field reports both accepted shapes" +trace $CLI experimental air run -h config.code_source.snapshot.git.remote + +title "the config. prefix is optional" +trace $CLI experimental air run -h compute.num_accelerators + +title "unknown field suggests the near match" +trace $CLI experimental air run -h config.compute.acclerator_type + +title "free-form map keys are not schema fields" +trace $CLI experimental air run -h config.parameters.learning_rate diff --git a/acceptance/experimental/air/get-ai-runtime/output.txt b/acceptance/experimental/air/get-ai-runtime/output.txt index c47790eaa31..21719cb8d78 100644 --- a/acceptance/experimental/air/get-ai-runtime/output.txt +++ b/acceptance/experimental/air/get-ai-runtime/output.txt @@ -27,7 +27,7 @@ │ MLflow Run my-run │ │ User user@example.com │ │ Accelerators 1x A10 │ -│ Environment N/A │ +│ Environment 4 │ │ │ ╰────────────────────────────────────────────────────────────────╯ diff --git a/acceptance/experimental/air/get-ai-runtime/test.toml b/acceptance/experimental/air/get-ai-runtime/test.toml index de442aca0ba..061238e5e37 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] @@ -12,8 +9,10 @@ MSYS_NO_PATHCONV = "1" Pattern = "HEAD /" Response.Body = '' -# The typed SDK GetRun response: an ai_runtime_task run has no gen_ai_compute_task, -# so the task comes back empty (the SDK has no field for ai_runtime_task). +# The GetRun response for an ai_runtime_task run. It has no gen_ai_compute_task, +# so the Environment cell instead comes from the run's environments[].spec (keyed +# by the task's environment_key), which the typed SDK Run drops — `air get` reads +# it with a raw request, so this route is hit twice (typed GetRun + that request). [[Server]] Pattern = "GET /api/2.2/jobs/runs/get" Response.Body = ''' @@ -24,12 +23,16 @@ Response.Body = ''' "start_time": 1700000000000, "end_time": 1700000012000, "state": {"life_cycle_state": "TERMINATED", "result_state": "SUCCESS"}, + "environments": [ + {"environment_key": "default", "spec": {"environment_version": "4"}} + ], "tasks": [ { "task_key": "train", "run_id": 456, "attempt_number": 0, "max_retries": 3, + "environment_key": "default", "ai_runtime_task": { "experiment": "my-exp", "deployments": [ diff --git a/acceptance/experimental/air/get/output.txt b/acceptance/experimental/air/get/output.txt index 6e51d7debf9..ff5cbd2ab04 100644 --- a/acceptance/experimental/air/get/output.txt +++ b/acceptance/experimental/air/get/output.txt @@ -27,7 +27,7 @@ │ MLflow Run my-run │ │ User user@example.com │ │ Accelerators 1x A10 │ -│ Environment ml-runtime-gpu:1.0 │ +│ Environment N/A │ │ │ ╰────────────────────────────────────────────────────────────────╯ 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..52e50016843 100644 --- a/acceptance/experimental/air/help/output.txt +++ b/acceptance/experimental/air/help/output.txt @@ -47,3 +47,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/output.txt b/acceptance/experimental/air/list/output.txt index e27cfb0bf39..6c1eca7d7f2 100644 --- a/acceptance/experimental/air/list/output.txt +++ b/acceptance/experimental/air/list/output.txt @@ -1,8 +1,8 @@ === list (text) >>> [CLI] experimental air list - Run ID Experiment Status Started Duration MLflow User Accelerators - [NUMID] qwen-train ● SUCCESS [TIMESTAMP] 12s …/runs/run1 [USERNAME] 8x H100 + Run ID Experiment Status Started Duration MLflow User Accelerators + [NUMID] qwen-train ● SUCCESS [TIMESTAMP] 12s qwen-train-001 [USERNAME] 8x H100 === list (json) >>> [CLI] experimental air list -o json @@ -25,8 +25,8 @@ === list --all-status (text, via AiTrainingService index) >>> [CLI] experimental air list --all-status - Run ID Experiment Status Started Duration MLflow User Accelerators - [NUMID] qwen-train ● SUCCESS [TIMESTAMP] 12s …/runs/run1 [USERNAME] 8x H100 + Run ID Experiment Status Started Duration MLflow User Accelerators + [NUMID] qwen-train ● SUCCESS [TIMESTAMP] 12s qwen-train-001 [USERNAME] 8x H100 === list --all-status (json) >>> [CLI] experimental air list --all-status -o json diff --git a/acceptance/experimental/air/list/test.toml b/acceptance/experimental/air/list/test.toml index 10c2a7600b8..9b68f61173b 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" @@ -64,6 +61,19 @@ Response.Body = ''' {"training_workflows": [{"job_run_id": "334747067049496", "submit_time": "2024-06-05T17:32:39Z"}]} ''' +# MLflow run names (for the MLflow column label) are fetched per AIR run (text mode). +[[Server]] +Pattern = "GET /api/2.0/mlflow/runs/get" +Response.Body = ''' +{ + "run": { + "info": { + "run_name": "qwen-train-001" + } + } +} +''' + # runs/get hydrates one index id into the same shape as a runs/list element. [[Server]] Pattern = "GET /api/2.2/jobs/runs/get" diff --git a/acceptance/experimental/air/logs-download/out.test.toml b/acceptance/experimental/air/logs-download/out.test.toml new file mode 100644 index 00000000000..e90b6d5d1ba --- /dev/null +++ b/acceptance/experimental/air/logs-download/out.test.toml @@ -0,0 +1,3 @@ +Local = true +Cloud = false +EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["direct"] diff --git a/acceptance/experimental/air/logs-download/output.txt b/acceptance/experimental/air/logs-download/output.txt new file mode 100644 index 00000000000..d3458758c42 --- /dev/null +++ b/acceptance/experimental/air/logs-download/output.txt @@ -0,0 +1,16 @@ + +=== download-to reports no logs when none are available +>>> [CLI] experimental air logs 123 --download-to dl-logs +No logs available for run 123. Run terminated in state SUCCESS + +=== download-to with an out-of-range node is rejected +>>> [CLI] experimental air logs 123 --download-to dl-logs --node 5 +Error: invalid --node 5: run has 2 node(s), indexed 0 to 1 + +Exit code: 1 + +=== download-to cannot be combined with --lines +>>> [CLI] experimental air logs 123 --download-to dl-logs --lines 50 +Error: --download-to writes complete logs, so it cannot be combined with --lines or --minutes + +Exit code: 1 diff --git a/acceptance/experimental/air/logs-download/script b/acceptance/experimental/air/logs-download/script new file mode 100644 index 00000000000..027b5c89528 --- /dev/null +++ b/acceptance/experimental/air/logs-download/script @@ -0,0 +1,12 @@ +# --download-to resolves the run's node count, then downloads each node's logs. +# This run resolves no MLflow run id, so it reports no logs (the full byte +# download is covered by unit tests, since the pre-signed URL host is dynamic). + +title "download-to reports no logs when none are available" +errcode trace $CLI experimental air logs 123 --download-to dl-logs + +title "download-to with an out-of-range node is rejected" +errcode trace $CLI experimental air logs 123 --download-to dl-logs --node 5 + +title "download-to cannot be combined with --lines" +errcode trace $CLI experimental air logs 123 --download-to dl-logs --lines 50 diff --git a/acceptance/experimental/air/logs-download/test.toml b/acceptance/experimental/air/logs-download/test.toml new file mode 100644 index 00000000000..1fd40b58e2b --- /dev/null +++ b/acceptance/experimental/air/logs-download/test.toml @@ -0,0 +1,40 @@ +# The command creates this download directory; don't treat it as test output. +Ignore = ["dl-logs"] + +# The SDK occasionally probes host reachability with a HEAD request; stub it so +# the test is deterministic. +[[Server]] +Pattern = "HEAD /" +Response.Body = '' + +# A completed 2-node run (GPU_1xA10 x 2 = 2 nodes). +[[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, + "ai_runtime_task": { + "experiment": "dl-exp", + "deployments": [ + {"command_path": "/x/command.sh", "compute": {"accelerator_type": "GPU_1xA10", "accelerator_count": 2}} + ] + } + } + ] +} +''' + +# No MLflow run id resolvable, so the download reports no logs rather than +# attempting a (host-dynamic) pre-signed artifact fetch. The full download path +# is covered by unit tests. +[[Server]] +Pattern = "GET /api/2.2/jobs/runs/get-output" +Response.Body = '{}' 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..9ea8deedc5e --- /dev/null +++ b/acceptance/experimental/air/logs/output.txt @@ -0,0 +1,66 @@ + +=== 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 diff --git a/acceptance/experimental/air/logs/script b/acceptance/experimental/air/logs/script new file mode 100644 index 00000000000..167ee331f01 --- /dev/null +++ b/acceptance/experimental/air/logs/script @@ -0,0 +1,29 @@ +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 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/register-image-no-secret-permission/dockercfg/config.json b/acceptance/experimental/air/register-image-no-secret-permission/dockercfg/config.json new file mode 100644 index 00000000000..a27cbb27880 --- /dev/null +++ b/acceptance/experimental/air/register-image-no-secret-permission/dockercfg/config.json @@ -0,0 +1,7 @@ +{ + "auths": { + "nvcr.io": { + "auth": "Ym9iOnNlY3JldHBhdA==" + } + } +} diff --git a/acceptance/experimental/air/register-image-no-secret-permission/out.test.toml b/acceptance/experimental/air/register-image-no-secret-permission/out.test.toml new file mode 100644 index 00000000000..e90b6d5d1ba --- /dev/null +++ b/acceptance/experimental/air/register-image-no-secret-permission/out.test.toml @@ -0,0 +1,3 @@ +Local = true +Cloud = false +EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["direct"] diff --git a/acceptance/experimental/air/register-image-no-secret-permission/output.txt b/acceptance/experimental/air/register-image-no-secret-permission/output.txt new file mode 100644 index 00000000000..36a98a1fe76 --- /dev/null +++ b/acceptance/experimental/air/register-image-no-secret-permission/output.txt @@ -0,0 +1,29 @@ + +=== private image, but the user cannot create a secret scope to store credentials +>>> [CLI] experimental air register-image nvcr.io/org/img:1.0 +Error: image "nvcr.io/org/img:1.0" requires credentials, and the credentials found in your local Docker config could not be stored: creating secret scope "docker-credentials-[USERNAME]" was denied (user does not have permission to create secret scopes). Ask a workspace admin for permission to create secret scopes + +Host: [DATABRICKS_URL] +Auth type: Personal Access Token (pat) + +Next steps: + - Verify you have the required permissions for this operation + - Check your identity: databricks auth describe + - Consider setting up a profile: databricks auth login --profile + +Exit code: 1 + +=== same failure as a JSON envelope +>>> [CLI] experimental air register-image nvcr.io/org/img:1.0 -o json +{ + "v": 1, + "ts": "[TIMESTAMP]", + "error": { + "code": "REGISTRATION_FAILED", + "kind": "PERMANENT", + "message": "image \"nvcr.io/org/img:1.0\" requires credentials, and the credentials found in your local Docker config could not be stored: creating secret scope \"docker-credentials-[USERNAME]\" was denied (user does not have permission to create secret scopes). Ask a workspace admin for permission to create secret scopes", + "retryable": false + } +} + +Exit code: 1 diff --git a/acceptance/experimental/air/register-image-no-secret-permission/script b/acceptance/experimental/air/register-image-no-secret-permission/script new file mode 100644 index 00000000000..a7cad12c4dd --- /dev/null +++ b/acceptance/experimental/air/register-image-no-secret-permission/script @@ -0,0 +1,9 @@ +# Exported rather than passed through trace: $TESTDIR is a backslash path on +# Windows, so echoing it into the output would not match this golden. +export DOCKER_CONFIG="$TESTDIR/dockercfg" + +title "private image, but the user cannot create a secret scope to store credentials" +errcode trace $CLI experimental air register-image nvcr.io/org/img:1.0 + +title "same failure as a JSON envelope" +errcode trace $CLI experimental air register-image nvcr.io/org/img:1.0 -o json diff --git a/acceptance/experimental/air/register-image-no-secret-permission/test.toml b/acceptance/experimental/air/register-image-no-secret-permission/test.toml new file mode 100644 index 00000000000..01a81f90216 --- /dev/null +++ b/acceptance/experimental/air/register-image-no-secret-permission/test.toml @@ -0,0 +1,40 @@ +# A user who cannot create secret scopes: the Docker credentials found locally +# can't be stored, so registration falls back to anonymous access and the private +# image is rejected. The error must name the storage failure as the cause rather +# than telling the user to `docker login` (which they already did successfully). + +# The SDK occasionally probes host reachability with a HEAD request; stub it so +# the test is deterministic. +[[Server]] +Pattern = "HEAD /" +Response.Body = '' + +[[Server]] +Pattern = "GET /api/2.0/ai-compute-manager/images:get" +Response.StatusCode = 404 +Response.Body = ''' +{"error_code": "NOT_FOUND", "message": "image not registered"} +''' + +# The image is private, so credential discovery proceeds. +[[Server]] +Pattern = "GET /api/2.0/ai-compute-manager/images:checkImageAccess" +Response.Body = ''' +{"publicly_accessible": false} +''' + +# The workspace admin has not granted permission to create secret scopes. +[[Server]] +Pattern = "POST /api/2.0/secrets/scopes/create" +Response.StatusCode = 403 +Response.Body = ''' +{"error_code": "PERMISSION_DENIED", "message": "user does not have permission to create secret scopes"} +''' + +# Without credentials the manager cannot pull the private image. +[[Server]] +Pattern = "POST /api/2.0/ai-compute-manager/images" +Response.StatusCode = 403 +Response.Body = ''' +{"error_code": "PERMISSION_DENIED", "message": "cannot pull nvcr.io/org/img:1.0: unauthorized"} +''' diff --git a/acceptance/experimental/air/register-image/dockercfg/config.json b/acceptance/experimental/air/register-image/dockercfg/config.json new file mode 100644 index 00000000000..a27cbb27880 --- /dev/null +++ b/acceptance/experimental/air/register-image/dockercfg/config.json @@ -0,0 +1,7 @@ +{ + "auths": { + "nvcr.io": { + "auth": "Ym9iOnNlY3JldHBhdA==" + } + } +} diff --git a/acceptance/experimental/air/register-image/out.test.toml b/acceptance/experimental/air/register-image/out.test.toml new file mode 100644 index 00000000000..e90b6d5d1ba --- /dev/null +++ b/acceptance/experimental/air/register-image/out.test.toml @@ -0,0 +1,3 @@ +Local = true +Cloud = false +EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["direct"] diff --git a/acceptance/experimental/air/register-image/output.txt b/acceptance/experimental/air/register-image/output.txt new file mode 100644 index 00000000000..ebb73f8cba8 --- /dev/null +++ b/acceptance/experimental/air/register-image/output.txt @@ -0,0 +1,111 @@ + +=== re-registers and reports an updated digest +>>> [CLI] experimental air register-image my-image:latest +Image registered: f9e8d7c6b5a4f3e2... + +To use this image in your training config: + environment: + docker_image: + url: my-image:latest + +=== re-registers and reports an updated digest (json) +>>> [CLI] experimental air register-image my-image:latest -o json +{ + "v": 1, + "ts": "[TIMESTAMP]", + "data": { + "docker_image_url": "my-image:latest", + "manifest_sha256": "f9e8d7c6b5a4f3e2newnewnewnew", + "status": "AVAILABLE", + "image_updated": true, + "cached": false + } +} + +=== tag-policy latest is accepted for backward compatibility +>>> [CLI] experimental air register-image my-image:latest --tag-policy latest -o json +{ + "v": 1, + "ts": "[TIMESTAMP]", + "data": { + "docker_image_url": "my-image:latest", + "manifest_sha256": "f9e8d7c6b5a4f3e2newnewnewnew", + "status": "AVAILABLE", + "image_updated": true, + "cached": false + } +} + +=== empty image url +>>> [CLI] experimental air register-image +Error: IMAGE_URL cannot be empty + +Exit code: 1 + +=== tag-policy auto is rejected +>>> [CLI] experimental air register-image my-image:latest --tag-policy auto +Error: --tag-policy auto is no longer supported: auto mode was removed and registration now always checks the source registry for the latest digest; omit the flag or use --tag-policy latest + +Exit code: 1 + +=== invalid tag policy (json) +>>> [CLI] experimental air register-image my-image:latest --tag-policy bogus -o json +{ + "v": 1, + "ts": "[TIMESTAMP]", + "error": { + "code": "INVALID_ARGS", + "kind": "PERMANENT", + "message": "invalid image tag policy \"bogus\": the only supported value is latest", + "retryable": false + } +} + +Exit code: 1 + +=== non-positive timeout is rejected +>>> [CLI] experimental air register-image my-image:latest --timeout-minutes 0 +Error: --timeout-minutes must be positive, got 0 + +Exit code: 1 + +=== removed credential flags are unknown +>>> [CLI] experimental air register-image my-image:latest --scope s --key k +Error: unknown flag: --scope + +Usage: + databricks experimental air register-image IMAGE_URL [flags] + +Flags: + -h, --help help for register-image + --timeout-minutes int Timeout to wait for the image to become available (default 60) + +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) + + +Exit code: 1 + +=== private image: credentials are discovered from the local Docker config +>>> [CLI] experimental air register-image nvcr.io/org/img:1.0 +Image registered: f9e8d7c6b5a4f3e2... + +To use this image in your training config: + environment: + docker_image: + url: nvcr.io/org/img:1.0 + +=== the registration POST carries the auto-stored credential reference +>>> print_requests.py //api/2.0/ai-compute-manager/images ^//api/2.0/secrets +{ + "method": "POST", + "path": "/api/2.0/ai-compute-manager/images", + "body": { + "credentials_key": "nvcr.io-bob-local", + "credentials_scope": "docker-credentials-[USERNAME]", + "docker_image_url": "nvcr.io/org/img:1.0" + } +} diff --git a/acceptance/experimental/air/register-image/script b/acceptance/experimental/air/register-image/script new file mode 100644 index 00000000000..d34ac9cc02f --- /dev/null +++ b/acceptance/experimental/air/register-image/script @@ -0,0 +1,37 @@ +title "re-registers and reports an updated digest" +trace $CLI experimental air register-image my-image:latest + +title "re-registers and reports an updated digest (json)" +trace $CLI experimental air register-image my-image:latest -o json + +title "tag-policy latest is accepted for backward compatibility" +trace $CLI experimental air register-image my-image:latest --tag-policy latest -o json + +title "empty image url" +errcode trace $CLI experimental air register-image " " + +title "tag-policy auto is rejected" +errcode trace $CLI experimental air register-image my-image:latest --tag-policy auto + +title "invalid tag policy (json)" +errcode trace $CLI experimental air register-image my-image:latest --tag-policy bogus -o json + +title "non-positive timeout is rejected" +errcode trace $CLI experimental air register-image my-image:latest --timeout-minutes 0 + +title "removed credential flags are unknown" +errcode trace $CLI experimental air register-image my-image:latest --scope s --key k + +# Drain the requests recorded by the cases above so the next assertion sees only +# the private-image registration's requests. +print_requests.py //api/2.0/ai-compute-manager/images &> LOG.drain + +# Exported rather than passed through trace: $TESTDIR is a backslash path on +# Windows, so echoing it into the output would not match this golden. +export DOCKER_CONFIG="$TESTDIR/dockercfg" + +title "private image: credentials are discovered from the local Docker config" +trace $CLI experimental air register-image nvcr.io/org/img:1.0 + +title "the registration POST carries the auto-stored credential reference" +trace print_requests.py //api/2.0/ai-compute-manager/images ^//api/2.0/secrets diff --git a/acceptance/experimental/air/register-image/test.toml b/acceptance/experimental/air/register-image/test.toml new file mode 100644 index 00000000000..497dc9ed81d --- /dev/null +++ b/acceptance/experimental/air/register-image/test.toml @@ -0,0 +1,30 @@ +RecordRequests = true + +# The SDK occasionally probes host reachability with a HEAD request; stub it so +# the test is deterministic. +[[Server]] +Pattern = "HEAD /" +Response.Body = '' + +# The image is not yet registered, so :get returns NOT_FOUND. Registration then +# always re-registers via POST regardless of any prior state. +[[Server]] +Pattern = "GET /api/2.0/ai-compute-manager/images:get" +Response.StatusCode = 404 +Response.Body = ''' +{"error_code": "NOT_FOUND", "message": "image not registered"} +''' + +# POST registers the image and returns its digest. +[[Server]] +Pattern = "POST /api/2.0/ai-compute-manager/images" +Response.Body = ''' +{"image": {"docker_image_url": "docker.io/library/my-image:latest", "state": "AVAILABLE", "manifest_sha256": "f9e8d7c6b5a4f3e2newnewnewnew"}} +''' + +# The private image is not publicly pullable, so credential discovery proceeds. +[[Server]] +Pattern = "GET /api/2.0/ai-compute-manager/images:checkImageAccess" +Response.Body = ''' +{"publicly_accessible": false} +''' diff --git a/acceptance/experimental/air/run-submit-deps/.gitattributes b/acceptance/experimental/air/run-submit-deps/.gitattributes new file mode 100644 index 00000000000..3c92a7ca7ab --- /dev/null +++ b/acceptance/experimental/air/run-submit-deps/.gitattributes @@ -0,0 +1,3 @@ +# run.yaml's contents are uploaded verbatim (as training_config.yaml), so its line +# endings must stay \n on every OS — a Windows \r would change the recorded payload. +run.yaml text eol=lf diff --git a/acceptance/experimental/air/run-submit-deps/databricks.yml b/acceptance/experimental/air/run-submit-deps/databricks.yml new file mode 100644 index 00000000000..a073ec04b4f --- /dev/null +++ b/acceptance/experimental/air/run-submit-deps/databricks.yml @@ -0,0 +1,2 @@ +bundle: + name: air-run-submit-deps diff --git a/acceptance/experimental/air/run-submit-deps/out.test.toml b/acceptance/experimental/air/run-submit-deps/out.test.toml new file mode 100644 index 00000000000..e90b6d5d1ba --- /dev/null +++ b/acceptance/experimental/air/run-submit-deps/out.test.toml @@ -0,0 +1,3 @@ +Local = true +Cloud = false +EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["direct"] diff --git a/acceptance/experimental/air/run-submit-deps/output.txt b/acceptance/experimental/air/run-submit-deps/output.txt new file mode 100644 index 00000000000..3c7280bc5ba --- /dev/null +++ b/acceptance/experimental/air/run-submit-deps/output.txt @@ -0,0 +1,61 @@ + +=== submit with inline dependencies +>>> [CLI] experimental air run -f run.yaml +Submitting experiment: deps-smoke +Submitted workload with Job Run ID: 555 +View job run at: [DATABRICKS_URL]/jobs/runs/555 + +Tip: use --watch to stream logs until the run completes. + +=== only config + command are uploaded; no requirements.yaml +>>> print_requests.py //api/2.0/workspace-files/import-file --oneline --sort --unique --keep +{"method": "POST", "path": "/api/2.0/workspace-files/import-file/Workspace/Users/[USERNAME]/.air/cli_launch/deps-smoke/deps-smoke_[RUN_ID]/command.sh", "q": {"overwrite": "true"}, "raw_body": "python train.py"} +{"method": "POST", "path": "/api/2.0/workspace-files/import-file/Workspace/Users/[USERNAME]/.air/cli_launch/deps-smoke/deps-smoke_[RUN_ID]/training_config.yaml", "q": {"overwrite": "true"}, "raw_body": "experiment_name: deps-smoke\ncommand: python train.py\ncompute:\n accelerator_type: GPU_1xH100\n num_accelerators: 1\nenvironment:\n version: 5\n dependencies:\n - numpy\n - torch==2.3.0\n"} + +=== declared deps ride on environments[].spec.dependencies +>>> print_requests.py //api/2.2/jobs/runs/submit +{ + "method": "POST", + "path": "/api/2.2/jobs/runs/submit", + "body": { + "environments": [ + { + "environment_key": "default", + "spec": { + "dependencies": [ + "numpy", + "torch==2.3.0" + ], + "environment_version": "5" + } + } + ], + "idempotency_token": "[UUID]", + "run_name": "deps-smoke", + "tasks": [ + { + "ai_runtime_task": { + "deployments": [ + { + "command_path": "/Workspace/Users/[USERNAME]/.air/cli_launch/deps-smoke/deps-smoke_[RUN_ID]/command.sh", + "compute": { + "accelerator_count": 1, + "accelerator_type": "GPU_1xH100" + } + } + ], + "experiment": "deps-smoke" + }, + "environment_key": "default", + "max_retries": 3, + "retry_on_timeout": true, + "run_if": "ALL_SUCCESS", + "task_key": "deps-smoke" + } + ] + } +} + +=== a requirements.yaml file path is rejected; deps must be inline +>>> [CLI] experimental air run -f run-file.yaml +Error: invalid config run-file.yaml: environment.dependencies must be a list of packages or reference a requirements.txt (see https://docs.databricks.com/aws/en/machine-learning/ai-runtime/cli/yaml-config#reference). A direct file reference is not supported diff --git a/acceptance/experimental/air/run-submit-deps/run-file.yaml b/acceptance/experimental/air/run-submit-deps/run-file.yaml new file mode 100644 index 00000000000..3e5fe54b7b0 --- /dev/null +++ b/acceptance/experimental/air/run-submit-deps/run-file.yaml @@ -0,0 +1,7 @@ +experiment_name: deps-file-smoke +command: python train.py +compute: + accelerator_type: GPU_1xH100 + num_accelerators: 1 +environment: + dependencies: ./reqs.yaml diff --git a/acceptance/experimental/air/run-submit-deps/run.yaml b/acceptance/experimental/air/run-submit-deps/run.yaml new file mode 100644 index 00000000000..fbb7e8681b5 --- /dev/null +++ b/acceptance/experimental/air/run-submit-deps/run.yaml @@ -0,0 +1,10 @@ +experiment_name: deps-smoke +command: python train.py +compute: + accelerator_type: GPU_1xH100 + num_accelerators: 1 +environment: + version: 5 + dependencies: + - numpy + - torch==2.3.0 diff --git a/acceptance/experimental/air/run-submit-deps/script b/acceptance/experimental/air/run-submit-deps/script new file mode 100644 index 00000000000..12dc884740d --- /dev/null +++ b/acceptance/experimental/air/run-submit-deps/script @@ -0,0 +1,11 @@ +title "submit with inline dependencies" +trace $CLI experimental air run -f run.yaml + +title "only config + command are uploaded; no requirements.yaml" +trace print_requests.py //api/2.0/workspace-files/import-file --oneline --sort --unique --keep + +title "declared deps ride on environments[].spec.dependencies" +trace print_requests.py //api/2.2/jobs/runs/submit + +title "a requirements.yaml file path is rejected; deps must be inline" +musterr trace $CLI experimental air run -f run-file.yaml diff --git a/acceptance/experimental/air/run-submit-deps/test.toml b/acceptance/experimental/air/run-submit-deps/test.toml new file mode 100644 index 00000000000..abe4f7f5ebf --- /dev/null +++ b/acceptance/experimental/air/run-submit-deps/test.toml @@ -0,0 +1,23 @@ +# A non-dry-run submit with inline dependencies (no code_source): asserts the deps +# land on spec.dependencies and that requirements.yaml is not uploaded. +RecordRequests = true + +# The SDK probes host reachability with a HEAD request; stub it for determinism. +[[Server]] +Pattern = "HEAD /" +Response.Body = '' + +[[Server]] +Pattern = "POST /api/2.0/ai-training/config:validate" +Response.Body = '{}' + +[[Server]] +Pattern = "POST /api/2.2/jobs/runs/submit" +Response.Body = ''' +{"run_id": 555} +''' + +# The per-run launch directory ends in _<16 hex>; the random suffix varies. +[[Repls]] +Old = 'deps-smoke_[0-9a-f]{16}' +New = 'deps-smoke_[RUN_ID]' diff --git a/acceptance/experimental/air/run-submit/output.txt b/acceptance/experimental/air/run-submit/output.txt index 8d52eed1dab..bd917135d0e 100644 --- a/acceptance/experimental/air/run-submit/output.txt +++ b/acceptance/experimental/air/run-submit/output.txt @@ -1,8 +1,12 @@ === submit with a git code_source >>> [CLI] experimental air run -f run.yaml -Submitted run 555 -View at: [DATABRICKS_URL]/jobs/runs/555 +Submitting experiment: submit-smoke +Uploading [SNAPSHOT_TARBALL]... +Submitted workload with Job Run ID: 555 +View job run 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 +27,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..2e641379092 100644 --- a/acceptance/experimental/air/run-submit/test.toml +++ b/acceptance/experimental/air/run-submit/test.toml @@ -1,19 +1,20 @@ # 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 /" Response.Body = '' +[[Server]] +Pattern = "POST /api/2.0/ai-training/config:validate" +Response.Body = '{}' + [[Server]] Pattern = "POST /api/2.2/jobs/runs/submit" Response.Body = ''' 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 deleted file mode 100644 index 7db6ef1aec2..00000000000 --- a/acceptance/experimental/air/unimplemented/output.txt +++ /dev/null @@ -1,12 +0,0 @@ - -=== 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 - -Exit code: 1 diff --git a/acceptance/experimental/air/unimplemented/script b/acceptance/experimental/air/unimplemented/script deleted file mode 100644 index 19dc13ffe85..00000000000 --- a/acceptance/experimental/air/unimplemented/script +++ /dev/null @@ -1,7 +0,0 @@ -# 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..85cb242db80 100644 --- a/experimental/air/cmd/air.go +++ b/experimental/air/cmd/air.go @@ -1,8 +1,6 @@ package aircmd import ( - "fmt" - "github.com/spf13/cobra" ) @@ -26,8 +24,3 @@ experimental and may change in future versions.`, return cmd } - -// notImplemented returns the placeholder error used by milestone-0 stubs. -func notImplemented(name string) error { - return fmt.Errorf("`air %s` is not implemented yet", name) -} diff --git a/experimental/air/cmd/compute.go b/experimental/air/cmd/compute.go index 07013c53906..e7b69d470be 100644 --- a/experimental/air/cmd/compute.go +++ b/experimental/air/cmd/compute.go @@ -54,8 +54,8 @@ func gpusPerNode(g gpuType) (int, error) { // computeConfig is the `compute` block of the run YAML: which accelerators to // use and how many. type computeConfig struct { - NumAccelerators int `yaml:"num_accelerators"` - AcceleratorType string `yaml:"accelerator_type"` + NumAccelerators int `yaml:"num_accelerators" help:"Total number of GPUs to allocate. Must be a positive multiple of the accelerator type's per-node GPU count. See https://docs.databricks.com/aws/en/machine-learning/ai-runtime/cli/yaml-config#reference for supported GPU types."` + AcceleratorType string `yaml:"accelerator_type" help:"Which accelerator to run on, e.g. GPU_1xA10. See https://docs.databricks.com/aws/en/machine-learning/ai-runtime/cli/yaml-config#reference for the current list of supported GPU types. Matched case-sensitively."` } // validate checks the compute block against the backend's constraints. diff --git a/experimental/air/cmd/docker_config_creds.go b/experimental/air/cmd/docker_config_creds.go new file mode 100644 index 00000000000..448c477b187 --- /dev/null +++ b/experimental/air/cmd/docker_config_creds.go @@ -0,0 +1,141 @@ +package aircmd + +import ( + "context" + "encoding/base64" + "encoding/json" + "os" + "os/exec" + "path/filepath" + "strings" + "time" + + "github.com/databricks/cli/libs/env" + "github.com/databricks/cli/libs/log" +) + +// Docker Hub is keyed under this exact legacy string in ~/.docker/config.json. +const dockerHubAuthKey = "https://index.docker.io/v1/" + +// credHelperTimeout bounds a docker-credential- invocation so a hung +// helper never delays registration. +const credHelperTimeout = 10 * time.Second + +// dockerConfig is the subset of ~/.docker/config.json we read. +type dockerConfig struct { + Auths map[string]struct{ Auth string } `json:"auths"` + CredsStore string `json:"credsStore"` + CredHelpers map[string]string `json:"credHelpers"` +} + +// dockerConfigPath returns ~/.docker/config.json, honoring DOCKER_CONFIG. +func dockerConfigPath(ctx context.Context) (string, error) { + if override, ok := env.Lookup(ctx, "DOCKER_CONFIG"); ok && override != "" { + return filepath.Join(override, "config.json"), nil + } + home, err := env.UserHomeDir(ctx) + if err != nil { + return "", err + } + return filepath.Join(home, ".docker", "config.json"), nil +} + +// registryKey maps a normalized image URL to the key used in docker config. +// Docker Hub uses the legacy key; everything else is the bare hostname. +func registryKey(imageURL string) string { + host, _, _ := strings.Cut(imageURL, "/") + switch host { + case "docker.io", "index.docker.io", "registry-1.docker.io": + return dockerHubAuthKey + } + return host +} + +// decodeDockerAuth decodes a base64 "username:password" auth field. +func decodeDockerAuth(authB64 string) (user, secret string, ok bool) { + decoded, err := base64.StdEncoding.DecodeString(authB64) + if err != nil { + return "", "", false + } + u, s, found := strings.Cut(string(decoded), ":") + if !found || u == "" || s == "" { + return "", "", false + } + return u, s, true +} + +// invokeCredHelper runs `docker-credential- get` for registry and parses +// its JSON. A missing helper, non-zero exit, or timeout yields ok=false; this is +// never an error (we fall through to the next credential source). +func invokeCredHelper(ctx context.Context, helper, registry string) (user, secret string, ok bool) { + ctx, cancel := context.WithTimeout(ctx, credHelperTimeout) + defer cancel() + + cmd := exec.CommandContext(ctx, "docker-credential-"+helper, "get") + cmd.Stdin = strings.NewReader(registry) + out, err := cmd.Output() + if err != nil { + log.Debugf(ctx, "docker-credential-%s get failed: %v", helper, err) + return "", "", false + } + + var payload struct { + Username string `json:"Username"` + Secret string `json:"Secret"` + } + if err := json.Unmarshal(out, &payload); err != nil || payload.Username == "" || payload.Secret == "" { + return "", "", false + } + return payload.Username, payload.Secret, true +} + +// readDockerCredentials looks up registry credentials for imageURL from the local +// Docker config, mirroring Docker's own resolution order: per-registry helper, +// then the global credsStore, then the inline base64 auth. imageURL must already +// be normalized so the first path segment is the registry host. Returns ok=false +// when no credentials are available; it never errors, since a missing or +// unreadable config just means the caller falls back to the public-image path. +func readDockerCredentials(ctx context.Context, imageURL string) (user, secret string, ok bool) { + path, err := dockerConfigPath(ctx) + if err != nil { + return "", "", false + } + // A missing or unreadable config is not an error: fall through to the + // public-image path. + data, err := os.ReadFile(path) + if err != nil { + return "", "", false + } + + var cfg dockerConfig + if err := json.Unmarshal(data, &cfg); err != nil { + log.Debugf(ctx, "could not parse %s: %v", path, err) + return "", "", false + } + + registry := registryKey(imageURL) + + // Per-registry helper takes precedence over everything else. + if helper := cfg.CredHelpers[registry]; helper != "" { + if u, s, ok := invokeCredHelper(ctx, helper, registry); ok { + return u, s, true + } + } + + // Global credential store. Consult it before the inline auth field, which may + // be stale data left from a pre-credsStore `docker login`. + if cfg.CredsStore != "" { + if u, s, ok := invokeCredHelper(ctx, cfg.CredsStore, registry); ok { + return u, s, true + } + } + + // Inline base64 auth as the final fallback. + if entry, present := cfg.Auths[registry]; present && entry.Auth != "" { + if u, s, ok := decodeDockerAuth(entry.Auth); ok { + return u, s, true + } + } + + return "", "", false +} diff --git a/experimental/air/cmd/docker_config_creds_test.go b/experimental/air/cmd/docker_config_creds_test.go new file mode 100644 index 00000000000..41463662d59 --- /dev/null +++ b/experimental/air/cmd/docker_config_creds_test.go @@ -0,0 +1,104 @@ +package aircmd + +import ( + "context" + "encoding/base64" + "os" + "path/filepath" + "runtime" + "testing" + + "github.com/databricks/cli/libs/env" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// writeDockerConfig writes config.json into a temp dir and returns a context with +// DOCKER_CONFIG pointing at it. +func writeDockerConfig(t *testing.T, body string) context.Context { + t.Helper() + dir := t.TempDir() + require.NoError(t, os.WriteFile(filepath.Join(dir, "config.json"), []byte(body), 0o600)) + return env.Set(t.Context(), "DOCKER_CONFIG", dir) +} + +func b64(t *testing.T, s string) string { + t.Helper() + return base64.StdEncoding.EncodeToString([]byte(s)) +} + +func TestRegistryKey(t *testing.T) { + assert.Equal(t, dockerHubAuthKey, registryKey("docker.io/library/ubuntu:latest")) + assert.Equal(t, dockerHubAuthKey, registryKey("index.docker.io/x")) + assert.Equal(t, "nvcr.io", registryKey("nvcr.io/nvidia/pytorch:24.01")) + assert.Equal(t, "ghcr.io", registryKey("ghcr.io/org/img")) +} + +func TestDecodeDockerAuth(t *testing.T) { + u, s, ok := decodeDockerAuth(base64.StdEncoding.EncodeToString([]byte("alice:pat123"))) + require.True(t, ok) + assert.Equal(t, "alice", u) + assert.Equal(t, "pat123", s) + + _, _, ok = decodeDockerAuth("not-base64!!") + assert.False(t, ok) + _, _, ok = decodeDockerAuth(base64.StdEncoding.EncodeToString([]byte("noseparator"))) + assert.False(t, ok) +} + +func TestReadDockerCredentialsInlineAuth(t *testing.T) { + ctx := writeDockerConfig(t, `{"auths":{"nvcr.io":{"auth":"`+b64(t, "bob:secret")+`"}}}`) + u, s, ok := readDockerCredentials(ctx, "nvcr.io/nvidia/pytorch:24.01") + require.True(t, ok) + assert.Equal(t, "bob", u) + assert.Equal(t, "secret", s) +} + +func TestReadDockerCredentialsDockerHubLegacyKey(t *testing.T) { + ctx := writeDockerConfig(t, `{"auths":{"https://index.docker.io/v1/":{"auth":"`+b64(t, "carol:tok")+`"}}}`) + u, _, ok := readDockerCredentials(ctx, "docker.io/library/ubuntu:latest") + require.True(t, ok) + assert.Equal(t, "carol", u) +} + +func TestReadDockerCredentialsNeedsNormalizedURL(t *testing.T) { + // A bare "ubuntu" has no registry host, so it must be normalized before + // lookup; discoverCredentials normalizes so Docker Hub creds are found. + ctx := writeDockerConfig(t, `{"auths":{"https://index.docker.io/v1/":{"auth":"`+b64(t, "dave:tok")+`"}}}`) + _, _, ok := readDockerCredentials(ctx, "ubuntu") + assert.False(t, ok) + u, _, ok := readDockerCredentials(ctx, normalizeDockerImageURL("ubuntu")) + require.True(t, ok) + assert.Equal(t, "dave", u) +} + +func TestReadDockerCredentialsMissingRegistry(t *testing.T) { + ctx := writeDockerConfig(t, `{"auths":{"nvcr.io":{"auth":"`+b64(t, "bob:secret")+`"}}}`) + _, _, ok := readDockerCredentials(ctx, "ghcr.io/org/img:latest") + assert.False(t, ok) +} + +func TestReadDockerCredentialsNoConfigFile(t *testing.T) { + // DOCKER_CONFIG points at an empty dir with no config.json. + ctx := env.Set(t.Context(), "DOCKER_CONFIG", t.TempDir()) + _, _, ok := readDockerCredentials(ctx, "nvcr.io/img:latest") + assert.False(t, ok) +} + +func TestReadDockerCredentialsCredHelper(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("credential helper stub is a POSIX shell script") + } + // A credHelper for the registry takes precedence over inline auth. The stub + // echoes a fixed Username/Secret payload on the `get` protocol. + binDir := t.TempDir() + stub := filepath.Join(binDir, "docker-credential-airtest") + require.NoError(t, os.WriteFile(stub, []byte("#!/bin/sh\necho '{\"Username\":\"helperuser\",\"Secret\":\"helpersecret\"}'\n"), 0o755)) + t.Setenv("PATH", binDir+string(os.PathListSeparator)+os.Getenv("PATH")) + + ctx := writeDockerConfig(t, `{"credHelpers":{"nvcr.io":"airtest"},"auths":{"nvcr.io":{"auth":"`+b64(t, "inline:ignored")+`"}}}`) + u, s, ok := readDockerCredentials(ctx, "nvcr.io/nvidia/pytorch:24.01") + require.True(t, ok) + assert.Equal(t, "helperuser", u) + assert.Equal(t, "helpersecret", s) +} diff --git a/experimental/air/cmd/format.go b/experimental/air/cmd/format.go index c32184517ba..de7046de3f8 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 @@ -290,19 +285,6 @@ func gpuDisplayName(gpuType string) string { return gpuType } -// environment returns the run's runtime image (the training environment), or an -// empty string if the run has no GenAI-compute task. -func environment(run *jobs.Run) string { - if len(run.Tasks) == 0 { - return "" - } - task := run.Tasks[0].GenAiComputeTask - if task == nil { - return "" - } - return task.DlRuntimeImage -} - // maxRetries returns the configured retry limit for the run's latest task as a // display string: "unlimited" for the backend's -1, otherwise the count. func maxRetries(run *jobs.Run) string { 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/get.go b/experimental/air/cmd/get.go index 2f8bd8dd09e..dcde5473d19 100644 --- a/experimental/air/cmd/get.go +++ b/experimental/air/cmd/get.go @@ -236,7 +236,10 @@ func buildGetData(run *jobs.Run) getData { } data.UserDisplay = orNA(run.CreatorUserName) data.AcceleratorsDisplay = orNA(accelerators(run)) - data.EnvironmentDisplay = orNA(environment(run)) + // EnvironmentDisplay is resolved at render time: the serverless environment + // version needs a raw GetRun read (the typed SDK Run drops it), so it is not + // filled here alongside the fields read straight off the run. + data.EnvironmentDisplay = na data.MaxRetriesDisplay = maxRetries(run) return data } diff --git a/experimental/air/cmd/image_client.go b/experimental/air/cmd/image_client.go new file mode 100644 index 00000000000..79613139e7a --- /dev/null +++ b/experimental/air/cmd/image_client.go @@ -0,0 +1,211 @@ +package aircmd + +import ( + "context" + "errors" + "fmt" + "net/http" + "strings" + "time" + + "github.com/databricks/databricks-sdk-go" + "github.com/databricks/databricks-sdk-go/apierr" + "github.com/databricks/databricks-sdk-go/client" +) + +// imagesAPIPath is the AI Compute Manager image service, called with a raw +// client.Do because the SDK does not model it. The ":get" and ":checkImageAccess" +// verbs are literal path suffixes the backend expects, not query-style actions. +const imagesAPIPath = "/api/2.0/ai-compute-manager/images" + +// imageStatus is the lifecycle state of a registered image, as reported in the +// API response's "state" field. +type imageStatus string + +const ( + imageStatusPending imageStatus = "PENDING" + imageStatusImporting imageStatus = "IMPORTING" + imageStatusAvailable imageStatus = "AVAILABLE" + imageStatusFailed imageStatus = "FAILED" +) + +// errImageUploadFailed marks a terminal FAILED upload, so callers can classify +// it as permanent rather than a retryable transient error. +var errImageUploadFailed = errors.New("image upload failed") + +// errImageWaitTimeout marks the poll giving up before the image became +// AVAILABLE. A later run may find it ready, so callers classify it as transient. +var errImageWaitTimeout = errors.New("image did not become AVAILABLE") + +// imageRegistration is a registered image with its status and metadata. +type imageRegistration struct { + DockerImageURL string `json:"docker_image_url"` + Status imageStatus `json:"-"` + StatusMessage string `json:"status_message"` + ManifestSHA256 string `json:"manifest_sha256"` + // State is the raw wire field; Status is derived from it via normalizeStatus + // so an unknown value degrades to PENDING rather than an invalid enum. + State string `json:"state"` +} + +// normalizeStatus maps the raw "state" field to a known status, defaulting to +// PENDING for absent or unrecognized values (matching the Python client). +func (r *imageRegistration) normalizeStatus() { + switch imageStatus(r.State) { + case imageStatusPending, imageStatusImporting, imageStatusAvailable, imageStatusFailed: + r.Status = imageStatus(r.State) + default: + r.Status = imageStatusPending + } +} + +// normalizeDockerImageURL canonicalizes a container image URL for consistent +// hashing. It prepends docker.io/ only for short-form URLs (no explicit +// registry); a registry is identified by a "." in the host portion of the first +// path component (e.g. docker.io, nvcr.io, registry.gitlab.com). +func normalizeDockerImageURL(imageURL string) string { + url := strings.TrimSpace(imageURL) + parts := strings.Split(url, "/") + + // A registry hostname contains a dot. Check only the host portion (before any + // ":") of the first component so version tags like "ubuntu:22.04" are not + // mistaken for a registry hostname. + if !strings.Contains(strings.Split(parts[0], ":")[0], ".") { + if len(parts) == 1 { + // Bare name (e.g. "ubuntu", "ubuntu:latest") — a Docker Hub official image. + url = "docker.io/library/" + url + } else { + // User/org image (e.g. "pytorch/pytorch:2.0.0") — just add the registry. + url = "docker.io/" + url + } + } + + // When a digest is present it takes precedence per the OCI spec — strip any tag. + if idx := strings.Index(url, "@"); idx != -1 { + repo, digest := url[:idx], url[idx+1:] + lastSlash := strings.LastIndex(repo, "/") + if colon := strings.Index(repo[lastSlash+1:], ":"); colon != -1 { + repo = repo[:lastSlash+1+colon] + } + return repo + "@" + digest + } + + // No digest and no tag on the final component — default to :latest. + last := url[strings.LastIndex(url, "/")+1:] + if !strings.Contains(last, ":") { + url += ":latest" + } + return url +} + +// imageClient calls the AI Compute Manager image service. +type imageClient struct { + api *client.DatabricksClient +} + +// newImageClient builds an imageClient from an authenticated workspace client. +func newImageClient(w *databricks.WorkspaceClient) (*imageClient, error) { + api, err := client.New(w.Config) + if err != nil { + return nil, fmt.Errorf("failed to create API client: %w", err) + } + return &imageClient{api: api}, nil +} + +// do issues one request against the image service, decoding the response into +// out. The image URL is normalized by the caller. +func (c *imageClient) do(ctx context.Context, method, endpoint string, query map[string]any, body, out any) error { + return c.api.Do(ctx, method, imagesAPIPath+endpoint, nil, query, body, out) +} + +// createImage registers a Docker image, optionally with registry credentials +// from a Databricks secret. CreateImage is idempotent: re-registering reconciles +// the stored status against the backing image entity. +func (c *imageClient) createImage(ctx context.Context, dockerImageURL, credentialsScope, credentialsKey string) (*imageRegistration, error) { + body := map[string]any{"docker_image_url": normalizeDockerImageURL(dockerImageURL)} + if credentialsScope != "" && credentialsKey != "" { + body["credentials_scope"] = credentialsScope + body["credentials_key"] = credentialsKey + } + + var resp struct { + Image *imageRegistration `json:"image"` + imageRegistration + } + if err := c.do(ctx, http.MethodPost, "", nil, body, &resp); err != nil { + return nil, fmt.Errorf("failed to register image: %w", err) + } + + // The response may wrap the registration under "image" or inline it. + reg := resp.Image + if reg == nil { + reg = &resp.imageRegistration + } + reg.normalizeStatus() + return reg, nil +} + +// getImage returns the registration for an image, or nil if it is not registered. +func (c *imageClient) getImage(ctx context.Context, dockerImageURL string) (*imageRegistration, error) { + query := map[string]any{"docker_image_url": normalizeDockerImageURL(dockerImageURL)} + var reg imageRegistration + if err := c.do(ctx, http.MethodGet, ":get", query, nil, ®); err != nil { + if errors.Is(err, apierr.ErrNotFound) { + return nil, nil + } + return nil, fmt.Errorf("failed to get image: %w", err) + } + reg.normalizeStatus() + return ®, nil +} + +// checkImageAccess reports whether an image is publicly pullable without +// credentials. It returns nil when the answer can't be determined — e.g. the +// manager region does not expose this RPC — so callers can treat nil as unknown. +func (c *imageClient) checkImageAccess(ctx context.Context, dockerImageURL string) *bool { + query := map[string]any{"docker_image_url": normalizeDockerImageURL(dockerImageURL)} + var resp struct { + PubliclyAccessible *bool `json:"publicly_accessible"` + } + if err := c.do(ctx, http.MethodGet, ":checkImageAccess", query, nil, &resp); err != nil { + return nil + } + return resp.PubliclyAccessible +} + +// waitForImageReady polls getImage until the image is AVAILABLE. Callers should +// call createImage first so the status is reconciled before polling begins. +func (c *imageClient) waitForImageReady(ctx context.Context, dockerImageURL string, timeout, pollInterval time.Duration) (*imageRegistration, error) { + deadline := time.Now().Add(timeout) + for { + reg, err := c.getImage(ctx, dockerImageURL) + if err != nil { + return nil, err + } + if reg == nil { + return nil, fmt.Errorf("image registration not found: %s", dockerImageURL) + } + + switch reg.Status { + case imageStatusAvailable: + return reg, nil + case imageStatusFailed: + msg := reg.StatusMessage + if msg == "" { + msg = "unknown error" + } + return nil, fmt.Errorf("%w: %s", errImageUploadFailed, msg) + case imageStatusPending, imageStatusImporting: + // Still uploading; fall through to sleep and poll again. + } + + if time.Now().After(deadline) { + return nil, fmt.Errorf("%w within %s", errImageWaitTimeout, timeout) + } + select { + case <-ctx.Done(): + return nil, ctx.Err() + case <-time.After(pollInterval): + } + } +} diff --git a/experimental/air/cmd/image_client_test.go b/experimental/air/cmd/image_client_test.go new file mode 100644 index 00000000000..8e12393d89d --- /dev/null +++ b/experimental/air/cmd/image_client_test.go @@ -0,0 +1,157 @@ +package aircmd + +import ( + "net/http" + "net/http/httptest" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestNormalizeDockerImageURL(t *testing.T) { + cases := map[string]string{ + // Bare Docker Hub official images get the library/ namespace. + "ubuntu": "docker.io/library/ubuntu:latest", + "ubuntu:22.04": "docker.io/library/ubuntu:22.04", + " ubuntu ": "docker.io/library/ubuntu:latest", + // User/org Docker Hub images get the registry prefix but no library/. + "pytorch/pytorch:2.0.0": "docker.io/pytorch/pytorch:2.0.0", + "pytorch/pytorch": "docker.io/pytorch/pytorch:latest", + // Explicit registries (host has a dot) are left as-is aside from a default tag. + "nvcr.io/nvidia/pytorch:24.01": "nvcr.io/nvidia/pytorch:24.01", + "registry.gitlab.com/org/repo": "registry.gitlab.com/org/repo:latest", + "docker.io/library/ubuntu:latest": "docker.io/library/ubuntu:latest", + // A digest takes precedence over any tag per the OCI spec. + "ubuntu@sha256:abc": "docker.io/library/ubuntu@sha256:abc", + "pytorch/pytorch:2.0.0@sha256:def": "docker.io/pytorch/pytorch@sha256:def", + "nvcr.io/nvidia/pytorch@sha256:xyz": "nvcr.io/nvidia/pytorch@sha256:xyz", + } + for in, want := range cases { + t.Run(in, func(t *testing.T) { + assert.Equal(t, want, normalizeDockerImageURL(in)) + }) + } +} + +func TestNormalizeStatus(t *testing.T) { + cases := map[string]imageStatus{ + "AVAILABLE": imageStatusAvailable, + "PENDING": imageStatusPending, + "IMPORTING": imageStatusImporting, + "FAILED": imageStatusFailed, + // Absent or unrecognized states degrade to PENDING. + "": imageStatusPending, + "UNKNOWN": imageStatusPending, + } + for state, want := range cases { + t.Run(state, func(t *testing.T) { + reg := imageRegistration{State: state} + reg.normalizeStatus() + assert.Equal(t, want, reg.Status) + }) + } +} + +// newTestImageClient builds an imageClient pointed at srv. +func newTestImageClient(t *testing.T, host string) *imageClient { + t.Helper() + c, err := newImageClient(newTestWorkspaceClient(t, host)) + require.NoError(t, err) + return c +} + +func TestImageClientCreateImageUnwrapsResponse(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == imagesAPIPath && r.Method == http.MethodPost { + // The backend wraps the registration under "image". + _, _ = w.Write([]byte(`{"image":{"state":"PENDING","manifest_sha256":"abc"}}`)) + return + } + _, _ = w.Write([]byte(`{}`)) + })) + t.Cleanup(srv.Close) + + reg, err := newTestImageClient(t, srv.URL).createImage(t.Context(), "ubuntu", "", "") + require.NoError(t, err) + assert.Equal(t, imageStatusPending, reg.Status) + assert.Equal(t, "abc", reg.ManifestSHA256) +} + +func TestImageClientGetImageNotFoundReturnsNil(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusNotFound) + _, _ = w.Write([]byte(`{"error_code":"NOT_FOUND","message":"not registered"}`)) + })) + t.Cleanup(srv.Close) + + reg, err := newTestImageClient(t, srv.URL).getImage(t.Context(), "ubuntu") + require.NoError(t, err) + assert.Nil(t, reg) +} + +func TestImageClientCheckImageAccess(t *testing.T) { + var hit bool + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == imagesAPIPath+":checkImageAccess" { + hit = true + _, _ = w.Write([]byte(`{"publicly_accessible":true}`)) + return + } + _, _ = w.Write([]byte(`{}`)) + })) + t.Cleanup(srv.Close) + + got := newTestImageClient(t, srv.URL).checkImageAccess(t.Context(), "ubuntu") + require.True(t, hit) + require.NotNil(t, got) + assert.True(t, *got) +} + +func TestImageClientCheckImageAccessUnknownOnError(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusInternalServerError) + })) + t.Cleanup(srv.Close) + + assert.Nil(t, newTestImageClient(t, srv.URL).checkImageAccess(t.Context(), "ubuntu")) +} + +func TestImageClientWaitForImageReady(t *testing.T) { + call := 0 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == imagesAPIPath+":get" { + // First poll is still PENDING; second reports AVAILABLE. + call++ + if call == 1 { + _, _ = w.Write([]byte(`{"state":"PENDING"}`)) + return + } + _, _ = w.Write([]byte(`{"state":"AVAILABLE","manifest_sha256":"done"}`)) + return + } + _, _ = w.Write([]byte(`{}`)) + })) + t.Cleanup(srv.Close) + + reg, err := newTestImageClient(t, srv.URL).waitForImageReady(t.Context(), "ubuntu", 5*time.Second, time.Millisecond) + require.NoError(t, err) + assert.Equal(t, imageStatusAvailable, reg.Status) + assert.Equal(t, "done", reg.ManifestSHA256) +} + +func TestImageClientWaitForImageReadyFailed(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == imagesAPIPath+":get" { + _, _ = w.Write([]byte(`{"state":"FAILED","status_message":"boom"}`)) + return + } + _, _ = w.Write([]byte(`{}`)) + })) + t.Cleanup(srv.Close) + + _, err := newTestImageClient(t, srv.URL).waitForImageReady(t.Context(), "ubuntu", 5*time.Second, time.Millisecond) + require.Error(t, err) + assert.Contains(t, err.Error(), "boom") +} diff --git a/experimental/air/cmd/image_credentials.go b/experimental/air/cmd/image_credentials.go new file mode 100644 index 00000000000..c4aa4f93cbf --- /dev/null +++ b/experimental/air/cmd/image_credentials.go @@ -0,0 +1,111 @@ +package aircmd + +import ( + "context" + "encoding/base64" + "errors" + "fmt" + "strings" + + "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/service/iam" + "github.com/databricks/databricks-sdk-go/service/workspace" +) + +// Discovered Docker credentials are stored in a per-Databricks-user secret. The +// scope is per-user (creator-only ACL) so one user's registry PAT is never +// readable by other workspace members. The key is "-" +// plus a suffix, so auto-managed keys are distinguishable, never overwrite a +// hand-curated secret, and don't collide across registries. +const ( + dockerCredsScopePrefix = "docker-credentials" + localManagedKeySuffix = "-local" +) + +// errSecretScopeQuota signals the workspace has hit its secret-scope limit. It +// is surfaced to the user rather than swallowed, since it only resolves by +// freeing a scope, not by retrying or falling back to the public-image path. +var errSecretScopeQuota = errors.New("workspace has reached the maximum number of secret scopes; delete an unused scope and retry") + +// encodeDockerCredentials base64-encodes "username:password", the form the +// registration backend decodes. +func encodeDockerCredentials(username, password string) string { + return base64.StdEncoding.EncodeToString([]byte(username + ":" + password)) +} + +// isScopeQuotaError reports whether err is a secret-scope quota rejection. +func isScopeQuotaError(err error) bool { + apiErr, ok := errors.AsType[*apierr.APIError](err) + return ok && apiErr.ErrorCode == "RESOURCE_LIMIT_EXCEEDED" +} + +// ensureSecretScope creates scope if it does not already exist, using the API +// default ACL (creator-only MANAGE). It must not grant workspace-wide access: a +// Docker-credential scope readable by every member would leak the user's PAT. +func ensureSecretScope(ctx context.Context, w *databricks.WorkspaceClient, scope string) error { + scopes, err := w.Secrets.ListScopesAll(ctx) + if err == nil { + for _, s := range scopes { + if s.Name == scope { + return nil + } + } + } else { + // A user without LIST permission can still create their own scope, so + // treat a list failure as "unknown" and proceed to create. + log.Debugf(ctx, "could not list secret scopes: %v", err) + } + + err = w.Secrets.CreateScope(ctx, workspace.CreateScope{Scope: scope}) + switch { + case err == nil: + return nil + case errors.Is(err, apierr.ErrResourceAlreadyExists): + return nil + case isScopeQuotaError(err): + return errSecretScopeQuota + case errors.Is(err, apierr.ErrPermissionDenied): + return fmt.Errorf("creating secret scope %q was denied (%w). Ask a workspace admin for permission to create secret scopes", scope, err) + default: + return fmt.Errorf("creating secret scope %q failed: %w", scope, err) + } +} + +// storeDockerCredentials stores registry credentials in the per-user secret +// scope and returns the (scope, key) reference for registration. A storage +// failure is returned rather than swallowed: registration continues without +// credentials (a public image still succeeds), but the caller reports this as the +// cause if the registry then rejects anonymous access. The caller resolves the +// credentials so the local Docker config is read only once, and passes +// normalizedImageURL so the key is namespaced by registry host. +func storeDockerCredentials(ctx context.Context, w *databricks.WorkspaceClient, normalizedImageURL, username, password string) (scope, key string, err error) { + me, err := w.CurrentUser.Me(ctx, iam.MeRequest{}) + if err != nil { + return "", "", fmt.Errorf("could not resolve the current Databricks user: %w", err) + } + + // Namespace the key by registry host so the same username on two registries + // (e.g. docker.io and nvcr.io) doesn't collide on one secret. + host, _, _ := strings.Cut(normalizedImageURL, "/") + scope = fmt.Sprintf("%s-%s", dockerCredsScopePrefix, me.UserName) + key = fmt.Sprintf("%s-%s%s", host, username, localManagedKeySuffix) + + if err := ensureSecretScope(ctx, w, scope); err != nil { + return "", "", err + } + + if err := w.Secrets.PutSecret(ctx, workspace.PutSecret{ + Scope: scope, + Key: key, + StringValue: encodeDockerCredentials(username, password), + }); err != nil { + if isScopeQuotaError(err) { + return "", "", errSecretScopeQuota + } + return "", "", fmt.Errorf("storing credentials in secret %s/%s failed: %w", scope, key, err) + } + + return scope, key, nil +} diff --git a/experimental/air/cmd/image_credentials_test.go b/experimental/air/cmd/image_credentials_test.go new file mode 100644 index 00000000000..9a7fd3b5a97 --- /dev/null +++ b/experimental/air/cmd/image_credentials_test.go @@ -0,0 +1,130 @@ +package aircmd + +import ( + "encoding/base64" + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestEncodeDockerCredentials(t *testing.T) { + got := encodeDockerCredentials("alice", "pat") + decoded, err := base64.StdEncoding.DecodeString(got) + require.NoError(t, err) + assert.Equal(t, "alice:pat", string(decoded)) +} + +// credServer records secret puts and lets a test choose the scope list and the +// create-scope failure. me is the current-user name returned to the client. +type credServer struct { + existingScopes []string + createStatus int // 0 → 200 + createCode string // error_code for a failed create; defaults to the quota code + putBodies []string +} + +func (cs *credServer) start(t *testing.T) string { + t.Helper() + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch { + case strings.HasSuffix(r.URL.Path, "/scim/v2/Me"): + _, _ = w.Write([]byte(`{"userName":"user@example.com"}`)) + case r.URL.Path == "/api/2.0/secrets/scopes/list": + var scopes []map[string]string + for _, s := range cs.existingScopes { + scopes = append(scopes, map[string]string{"name": s}) + } + _ = json.NewEncoder(w).Encode(map[string]any{"scopes": scopes}) + case r.URL.Path == "/api/2.0/secrets/scopes/create": + if cs.createStatus != 0 { + code := cs.createCode + if code == "" { + code = "RESOURCE_LIMIT_EXCEEDED" + } + w.WriteHeader(cs.createStatus) + _, _ = w.Write([]byte(`{"error_code":"` + code + `","message":"denied"}`)) + return + } + _, _ = w.Write([]byte(`{}`)) + case r.URL.Path == "/api/2.0/secrets/put": + body, _ := io.ReadAll(r.Body) + cs.putBodies = append(cs.putBodies, string(body)) + _, _ = w.Write([]byte(`{}`)) + default: + _, _ = w.Write([]byte(`{}`)) + } + })) + t.Cleanup(srv.Close) + return srv.URL +} + +func TestStoreDockerCredentialsCreatesScope(t *testing.T) { + cs := &credServer{} + w := newTestWorkspaceClient(t, cs.start(t)) + + scope, key, err := storeDockerCredentials(t.Context(), w, "docker.io/library/ubuntu:latest", "bob", "secret") + require.NoError(t, err) + assert.Equal(t, "docker-credentials-user@example.com", scope) + assert.Equal(t, "docker.io-bob-local", key) + require.Len(t, cs.putBodies, 1) + assert.Contains(t, cs.putBodies[0], base64.StdEncoding.EncodeToString([]byte("bob:secret"))) +} + +// TestStoreDockerCredentialsKeyIsPerRegistry guards against the same username on +// two registries colliding on one secret key. +func TestStoreDockerCredentialsKeyIsPerRegistry(t *testing.T) { + cases := map[string]string{ + "docker.io/library/ubuntu:latest": "docker.io-bob-local", + "nvcr.io/nvidia/pytorch:24.01": "nvcr.io-bob-local", + "ghcr.io/org/img:1.0": "ghcr.io-bob-local", + } + for imageURL, wantKey := range cases { + t.Run(imageURL, func(t *testing.T) { + cs := &credServer{} + w := newTestWorkspaceClient(t, cs.start(t)) + _, key, err := storeDockerCredentials(t.Context(), w, imageURL, "bob", "secret") + require.NoError(t, err) + assert.Equal(t, wantKey, key) + }) + } +} + +func TestStoreDockerCredentialsScopeExists(t *testing.T) { + // When the scope already exists, create is not required; storage still succeeds. + cs := &credServer{existingScopes: []string{"docker-credentials-user@example.com"}} + w := newTestWorkspaceClient(t, cs.start(t)) + + _, _, err := storeDockerCredentials(t.Context(), w, "nvcr.io/org/img:1.0", "bob", "secret") + require.NoError(t, err) + assert.Len(t, cs.putBodies, 1) +} + +// TestStoreDockerCredentialsPermissionDenied covers the workspace where the user +// may not create secret scopes: the failure must surface with admin guidance +// rather than be swallowed into a misleading "run docker login" error. +func TestStoreDockerCredentialsPermissionDenied(t *testing.T) { + cs := &credServer{createStatus: http.StatusForbidden, createCode: "PERMISSION_DENIED"} + w := newTestWorkspaceClient(t, cs.start(t)) + + _, _, err := storeDockerCredentials(t.Context(), w, "nvcr.io/org/img:1.0", "bob", "secret") + require.Error(t, err) + assert.NotErrorIs(t, err, errSecretScopeQuota) + assert.Contains(t, err.Error(), `creating secret scope "docker-credentials-user@example.com" was denied`) + assert.Contains(t, err.Error(), "Ask a workspace admin") + assert.Empty(t, cs.putBodies, "must not attempt to store the secret when the scope could not be created") +} + +func TestStoreDockerCredentialsQuotaError(t *testing.T) { + cs := &credServer{createStatus: http.StatusForbidden} + w := newTestWorkspaceClient(t, cs.start(t)) + + _, _, err := storeDockerCredentials(t.Context(), w, "nvcr.io/org/img:1.0", "bob", "secret") + require.Error(t, err) + assert.ErrorIs(t, err, errSecretScopeQuota) +} diff --git a/experimental/air/cmd/list.go b/experimental/air/cmd/list.go index 7d5c703a39c..c5cc64aafdf 100644 --- a/experimental/air/cmd/list.go +++ b/experimental/air/cmd/list.go @@ -46,10 +46,13 @@ type listRow struct { // Experiment, Duration, MLflowURL and Accelerators are table-only columns, // omitted from JSON to match `air list --json`. - Experiment string `json:"-"` - Duration string `json:"-"` - MLflowURL string `json:"-"` - Accelerators string `json:"-"` + Experiment string `json:"-"` + Duration string `json:"-"` + MLflowURL string `json:"-"` + MLflowLabel string `json:"-"` + RunURL string `json:"-"` + ExperimentURL string `json:"-"` + Accelerators string `json:"-"` } // listedRun pairs a row with its task run id, so the MLflow link can be fetched @@ -68,6 +71,7 @@ type listQuery struct { filters listFilters fetchMLflow bool limit int + workspaceID int64 } func newListCommand() *cobra.Command { @@ -118,6 +122,15 @@ func newListCommand() *cobra.Command { userFilter = currentUser } + // Fetch workspace ID once for dashboard links; proceed with 0 on error. + var workspaceID int64 + wsID, err := w.CurrentWorkspaceID(ctx) + if err != nil { + log.Debugf(ctx, "air list: could not fetch workspace ID for dashboard links: %v", err) + } else { + workspaceID = wsID + } + fetcher := newRunFetcher(ctx, w, listQuery{ activeOnly: !allStatus, allUsers: allUsers, @@ -126,6 +139,7 @@ func newListCommand() *cobra.Command { filters: f, fetchMLflow: root.OutputType(cmd) == flags.OutputText, limit: limit, + workspaceID: workspaceID, }) // JSON prints the newest `limit` runs once. Text renders the table: @@ -165,6 +179,7 @@ type runFetcher struct { w *databricks.WorkspaceClient fetchMLflow bool strategy listStrategy + workspaceID int64 exhausted bool } @@ -175,6 +190,7 @@ func newRunFetcher(ctx context.Context, w *databricks.WorkspaceClient, q listQue w: w, fetchMLflow: q.fetchMLflow, strategy: newListStrategy(ctx, w, q), + workspaceID: q.workspaceID, } } @@ -209,7 +225,7 @@ func (f *runFetcher) next(want int) ([]listRow, error) { // MLflow links appear only in the text table, so the per-run get-output // lookups are skipped for JSON output (which omits the column anyway). if f.fetchMLflow { - setMLflowLinks(f.ctx, f.w, entries) + setMLflowLinks(f.ctx, f.w, f.w.Config.Host, entries) } rows := make([]listRow, len(entries)) @@ -223,11 +239,13 @@ func (f *runFetcher) next(want int) ([]listRow, error) { // and filters. It buffers a page's leftover runs so successive next() calls // resume where the last stopped. type jobsScanStrategy struct { - ctx context.Context - w *databricks.WorkspaceClient - iter listing.Iterator[jobs.BaseRun] - userFilter string - filters listFilters + ctx context.Context + w *databricks.WorkspaceClient + iter listing.Iterator[jobs.BaseRun] + userFilter string + filters listFilters + host string + workspaceID int64 scanned int } @@ -240,11 +258,13 @@ func newJobsScanStrategy(ctx context.Context, w *databricks.WorkspaceClient, q l ActiveOnly: q.activeOnly, } return &jobsScanStrategy{ - ctx: ctx, - w: w, - iter: w.Jobs.ListRuns(ctx, req), - userFilter: q.userFilter, - filters: q.filters, + ctx: ctx, + w: w, + iter: w.Jobs.ListRuns(ctx, req), + userFilter: q.userFilter, + filters: q.filters, + host: w.Config.Host, + workspaceID: q.workspaceID, } } @@ -267,7 +287,7 @@ func (s *jobsScanStrategy) next(want int) ([]listedRun, error) { if !s.filters.matches(run) { continue } - entries = append(entries, listedRun{row: buildListRow(run), taskRunID: taskRunID(run)}) + entries = append(entries, listedRun{row: buildListRow(run, s.host, s.workspaceID), taskRunID: taskRunID(run)}) } return entries, nil } @@ -288,15 +308,18 @@ func warnIfTruncated(ctx context.Context, f *runFetcher) { } } -// setMLflowLinks fills in each row's MLflow link in parallel, best-effort: a row -// whose IDs can't be resolved keeps its "-" placeholder. -func setMLflowLinks(ctx context.Context, w *databricks.WorkspaceClient, entries []listedRun) { +// setMLflowLinks fills in each row's MLflow link, label, and experiment URL in +// parallel, best-effort: a row whose IDs can't be resolved keeps its "-" placeholder. +func setMLflowLinks(ctx context.Context, w *databricks.WorkspaceClient, host string, entries []listedRun) { var g errgroup.Group g.SetLimit(enrichConcurrency) for i := range entries { g.Go(func() error { if ids := mlflowIDsForTask(ctx, w, entries[i].taskRunID); ids != nil { - entries[i].row.MLflowURL = mlflowLogsURL(w.Config.Host, ids) + entries[i].row.MLflowURL = mlflowLogsURL(host, ids) + name := fetchMLflowRunName(ctx, w, ids.RunID) + entries[i].row.MLflowLabel = mlflowRunLabel(name, ids.RunID) + entries[i].row.ExperimentURL = mlflowExperimentURL(host, ids) } return nil }) diff --git a/experimental/air/cmd/list_cache.go b/experimental/air/cmd/list_cache.go index 6f0b95a97d3..1b2d47d57d3 100644 --- a/experimental/air/cmd/list_cache.go +++ b/experimental/air/cmd/list_cache.go @@ -28,25 +28,29 @@ type listCacheKey struct { // table-only columns, which listRow tags json:"-" and so wouldn't survive a // direct marshal), the filter inputs, and the submit time. type cachedRun struct { - RunID string `json:"run_id"` - RunName string `json:"run_name"` - User string `json:"user"` - Status string `json:"status"` - StartedAt *string `json:"started_at"` - IsSweep bool `json:"is_sweep"` - Experiment string `json:"experiment"` - Duration string `json:"duration"` - MLflowURL string `json:"mlflow_url"` - Accelerators string `json:"accelerators"` - Fields filterFields `json:"filter_fields"` - SubmitTimeMs int64 `json:"submit_time_ms"` + RunID string `json:"run_id"` + RunName string `json:"run_name"` + User string `json:"user"` + Status string `json:"status"` + StartedAt *string `json:"started_at"` + IsSweep bool `json:"is_sweep"` + Experiment string `json:"experiment"` + Duration string `json:"duration"` + MLflowURL string `json:"mlflow_url"` + MLflowLabel string `json:"mlflow_label"` + RunURL string `json:"run_url"` + ExperimentURL string `json:"experiment_url"` + Accelerators string `json:"accelerators"` + Fields filterFields `json:"filter_fields"` + SubmitTimeMs int64 `json:"submit_time_ms"` } func (c cachedRun) toRow() listRow { return listRow{ RunID: c.RunID, RunName: c.RunName, User: c.User, Status: c.Status, StartedAt: c.StartedAt, IsSweep: c.IsSweep, Experiment: c.Experiment, - Duration: c.Duration, MLflowURL: c.MLflowURL, Accelerators: c.Accelerators, + Duration: c.Duration, MLflowURL: c.MLflowURL, MLflowLabel: c.MLflowLabel, + RunURL: c.RunURL, ExperimentURL: c.ExperimentURL, Accelerators: c.Accelerators, } } @@ -54,7 +58,8 @@ func cachedRunFromRow(r listRow, fields filterFields, submitTimeMs int64) cached return cachedRun{ RunID: r.RunID, RunName: r.RunName, User: r.User, Status: r.Status, StartedAt: r.StartedAt, IsSweep: r.IsSweep, Experiment: r.Experiment, - Duration: r.Duration, MLflowURL: r.MLflowURL, Accelerators: r.Accelerators, + Duration: r.Duration, MLflowURL: r.MLflowURL, MLflowLabel: r.MLflowLabel, + RunURL: r.RunURL, ExperimentURL: r.ExperimentURL, Accelerators: r.Accelerators, Fields: fields, SubmitTimeMs: submitTimeMs, } } diff --git a/experimental/air/cmd/list_detail.go b/experimental/air/cmd/list_detail.go new file mode 100644 index 00000000000..c0eca25b13f --- /dev/null +++ b/experimental/air/cmd/list_detail.go @@ -0,0 +1,64 @@ +package aircmd + +import ( + "bytes" + "context" + "errors" + "fmt" + + "github.com/databricks/databricks-sdk-go" + "github.com/databricks/databricks-sdk-go/apierr" + "github.com/databricks/databricks-sdk-go/service/jobs" +) + +// runDetailText fetches a run and renders the same styled view as `air get` into +// a string, for the list picker's detail pane. +func runDetailText(ctx context.Context, w *databricks.WorkspaceClient, runID int64) (string, error) { + run, err := w.Jobs.GetRun(ctx, jobs.GetRunRequest{RunId: runID}) + if err != nil { + if errors.Is(err, apierr.ErrResourceDoesNotExist) { + return "", fmt.Errorf("run %d not found", runID) + } + return "", fmt.Errorf("failed to fetch run: %w", err) + } + + // A missing workspace id only drops the ?o= org hint from the dashboard link. + workspaceID, _ := w.CurrentWorkspaceID(ctx) + + data := buildGetData(run) + data.DashboardURL = dashboardURL(w.Config.Host, runID, workspaceID) + ids := mlflowIDs(ctx, w, run) + if ids != nil { + url := mlflowLogsURL(w.Config.Host, ids) + data.MLflowURL = &url + } + + var buf bytes.Buffer + renderRunText(ctx, &buf, w, run, &data, ids) + return buf.String(), nil +} + +// runLogsSnapshot fetches a one-shot tail of a run's logs into a string, for the +// list picker's detail pane. +func runLogsSnapshot(ctx context.Context, w *databricks.WorkspaceClient, runID int64) (string, error) { + status, err := resolveRunStatus(ctx, w, runID) + if err != nil { + if errors.Is(err, apierr.ErrResourceDoesNotExist) { + return "", fmt.Errorf("run %d not found", runID) + } + return "", fmt.Errorf("failed to fetch run status: %w", err) + } + + req := logRequest{ + runID: runID, + attempt: -1, // latest attempt + tailLines: 1000, // one-shot tail, not a live follow + staticView: true, + } + + var buf bytes.Buffer + if _, err := fetchLogs(ctx, w, &buf, req, status); err != nil { + return "", fmt.Errorf("failed to fetch logs: %w", err) + } + return buf.String(), nil +} diff --git a/experimental/air/cmd/list_format.go b/experimental/air/cmd/list_format.go index bfa74728cc3..393b48a768c 100644 --- a/experimental/air/cmd/list_format.go +++ b/experimental/air/cmd/list_format.go @@ -9,7 +9,8 @@ import ( // buildListRow extracts the columns shown for one run. Optional cells fall back // to "-"; MLflowURL starts as "-" and setMLflowLinks fills it in for text output. -func buildListRow(run *jobs.Run) listRow { +// host and workspaceID are used for building dashboard URLs. +func buildListRow(run *jobs.Run, host string, workspaceID int64) listRow { experiment := "-" if e := jobExperiment(run); e != "" { experiment = e @@ -42,6 +43,8 @@ func buildListRow(run *jobs.Run) listRow { Experiment: experiment, Duration: duration, MLflowURL: "-", + MLflowLabel: "-", + RunURL: dashboardURL(host, run.RunId, workspaceID), Accelerators: accel, } } diff --git a/experimental/air/cmd/list_index.go b/experimental/air/cmd/list_index.go index 243affd664f..6e42c1695eb 100644 --- a/experimental/air/cmd/list_index.go +++ b/experimental/air/cmd/list_index.go @@ -16,12 +16,14 @@ import ( // skip the network. Unlike the Jobs scan it can't lazy-page (it must sort the // whole id set first), but it still yields in batches so the table paints early. type indexStrategy struct { - ctx context.Context - w *databricks.WorkspaceClient - activeOnly bool - filters listFilters - limit int - cache *cache.Cache + ctx context.Context + w *databricks.WorkspaceClient + activeOnly bool + filters listFilters + limit int + cache *cache.Cache + host string + workspaceID int64 ids []int64 // newest-first run ids to hydrate, resolved on first next() pos int @@ -30,12 +32,14 @@ type indexStrategy struct { func newIndexStrategy(ctx context.Context, w *databricks.WorkspaceClient, q listQuery, limit int) *indexStrategy { return &indexStrategy{ - ctx: ctx, - w: w, - activeOnly: q.activeOnly, - filters: q.filters, - limit: limit, - cache: newListCache(ctx), + ctx: ctx, + w: w, + activeOnly: q.activeOnly, + filters: q.filters, + limit: limit, + cache: newListCache(ctx), + host: w.Config.Host, + workspaceID: q.workspaceID, } } @@ -120,7 +124,7 @@ func (s *indexStrategy) hydrate(ids []int64) ([]listedRun, error) { if !s.filters.matchesFields(fields) { continue } - row := buildListRow(run) + row := buildListRow(run, s.host, s.workspaceID) rows = append(rows, listedRun{row: row, taskRunID: taskRunID(run)}) if isTerminal(run) { start, _ := jobTiming(run) diff --git a/experimental/air/cmd/list_test.go b/experimental/air/cmd/list_test.go index f70330240e5..058b73a2e7d 100644 --- a/experimental/air/cmd/list_test.go +++ b/experimental/air/cmd/list_test.go @@ -167,7 +167,7 @@ func TestBuildListRowFromRun(t *testing.T) { assert.Equal(t, "GPU_1xA10", gpu) assert.Equal(t, 1, count) - row := buildListRow(&run) + row := buildListRow(&run, "https://example.test", 0) assert.Equal(t, "842552489592352", row.RunID) assert.Equal(t, "SUCCESS", row.Status) assert.Equal(t, "my-first-air-run", row.Experiment) @@ -181,7 +181,7 @@ func TestBuildListRow(t *testing.T) { run.EndTime = 1700000012000 run.State = &jobs.RunState{ResultState: jobs.RunResultStateSuccess} - row := buildListRow(&run) + row := buildListRow(&run, "https://example.test", 0) assert.Equal(t, "123", row.RunID) assert.Equal(t, "me@example.com", row.User) assert.Equal(t, "SUCCESS", row.Status) @@ -195,7 +195,7 @@ func TestBuildListRow(t *testing.T) { func TestBuildListRowDashFallbacks(t *testing.T) { // A run with no task, compute, or start time falls back to dashes and UNKNOWN. - row := buildListRow(&jobs.Run{RunId: 7}) + row := buildListRow(&jobs.Run{RunId: 7}, "https://example.test", 0) assert.Equal(t, "-", row.Experiment) assert.Equal(t, "-", row.Duration) assert.Equal(t, "-", row.Accelerators) @@ -208,8 +208,8 @@ func TestBuildListRowSweep(t *testing.T) { run := jobs.Run{RunId: 9, Tasks: []jobs.RunTask{{ ForEachTask: &jobs.RunForEachTask{Task: jobs.Task{AiRuntimeTask: &jobs.AiRuntimeTask{Experiment: "sweep"}}}, }}} - assert.True(t, buildListRow(&run).IsSweep) - assert.Equal(t, "sweep", buildListRow(&run).Experiment) + assert.True(t, buildListRow(&run, "https://example.test", 0).IsSweep) + assert.Equal(t, "sweep", buildListRow(&run, "https://example.test", 0).Experiment) } func TestListInvalidLimit(t *testing.T) { diff --git a/experimental/air/cmd/list_tui.go b/experimental/air/cmd/list_tui.go index 85fe70774b4..4707fd55fce 100644 --- a/experimental/air/cmd/list_tui.go +++ b/experimental/air/cmd/list_tui.go @@ -3,8 +3,10 @@ package aircmd import ( "fmt" "io" + "strconv" "strings" + "github.com/charmbracelet/bubbles/viewport" tea "github.com/charmbracelet/bubbletea" "github.com/charmbracelet/lipgloss" "github.com/databricks/cli/libs/cmdio" @@ -73,6 +75,19 @@ func staticListTable(r *lipgloss.Renderer, rows []listRow, links bool) string { return b.String() } +// Mode constants for the TUI. +const ( + modeList = 0 + modeDetail = 1 +) + +// detailMsg carries fetched run details or error. +type detailMsg struct { + title string + body string + err error +} + // listModel is the inline, navigable runs table. It lazily pages older runs from // the fetcher as the cursor nears the end of the loaded rows. fetcher is nil for // a fixed, non-paging table (e.g. in tests). @@ -88,15 +103,23 @@ type listModel struct { cursor int offset int // index of the first visible row height int // terminal height, for windowing + + mode int // modeList or modeDetail + viewport viewport.Model // for detail pane + detailTitle string // title of detail pane + detailLoading bool // loading detail + detailContent string // rendered detail content } func newListModel(r *lipgloss.Renderer, f *runFetcher, rows []listRow, links bool) listModel { return listModel{ - rows: rows, - styles: newListStyles(r), - cols: computeListCols(rows), - links: links, - fetcher: f, + rows: rows, + styles: newListStyles(r), + cols: computeListCols(rows), + links: links, + fetcher: f, + mode: modeList, + viewport: viewport.New(0, 0), } } @@ -148,6 +171,8 @@ func (m listModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) { switch msg := msg.(type) { case tea.WindowSizeMsg: m.height = msg.Height + m.viewport.Width = msg.Width + m.viewport.Height = max(msg.Height-3, 1) m.offset = m.clampedOffset() return m.maybeFetch() @@ -167,9 +192,46 @@ func (m listModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) { } return m, nil + case detailMsg: + // The fetch is async: if the user pressed esc back to the list before it + // resolved, drop the late result rather than snapping them into a pane + // they already dismissed. + if m.mode != modeDetail { + return m, nil + } + m.detailLoading = false + if msg.err != nil { + m.detailContent = fmt.Sprintf("Error: %v", msg.err) + } else { + m.detailContent = msg.body + } + m.detailTitle = msg.title + m.viewport.SetContent(m.detailContent) + m.viewport.GotoTop() + return m, nil + case tea.KeyMsg: + // Detail pane key handling. + if m.mode == modeDetail { + switch msg.String() { + case "esc", "q": + m.mode = modeList + return m, nil + case "ctrl+c": + return m, tea.Quit + default: + // Delegate scrolling keys to the viewport. + var cmd tea.Cmd + m.viewport, cmd = m.viewport.Update(msg) + return m, cmd + } + } + + // List pane key handling. switch msg.String() { - case "q", "ctrl+c", "esc": + case "q", "ctrl+c": + return m, tea.Quit + case "esc": return m, tea.Quit case "up", "k": if m.cursor > 0 { @@ -194,6 +256,24 @@ func (m listModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) { return m, openURL(url) } } + case "i": + // Open the run details pane; the fetch fills it in. + if m.fetcher != nil && len(m.rows) > 0 { + m.mode = modeDetail + m.detailLoading = true + m.detailTitle = "Run Details" + runID, _ := parseRunID(m.rows[m.cursor].RunID) + return m, m.fetchRunDetail(runID) + } + case "L", "l": + // Open the logs snapshot pane; the fetch fills it in. + if m.fetcher != nil && len(m.rows) > 0 { + m.mode = modeDetail + m.detailLoading = true + m.detailTitle = "Logs Snapshot" + runID, _ := parseRunID(m.rows[m.cursor].RunID) + return m, m.fetchRunLogs(runID) + } } m.offset = m.clampedOffset() return m.maybeFetch() @@ -212,6 +292,19 @@ func (m listModel) clampedOffset() int { } func (m listModel) View() string { + if m.mode == modeDetail { + faint := m.styles.r.NewStyle().Foreground(colN7) + lines := []string{m.detailTitle} + if m.detailLoading { + lines = append(lines, "Loading…") + } else { + lines = append(lines, m.viewport.View()) + } + footer := faint.Render("↑/↓ scroll · esc back · q quit") + lines = append(lines, footer) + return strings.Join(lines, "\n") + } + if len(m.rows) == 0 { return m.styles.r.NewStyle().Foreground(colN9).Render("No runs found.") + "\n" } @@ -229,7 +322,7 @@ func (m listModel) View() string { // paging state (loading / load failed). func (m listModel) renderHint() string { faint := m.styles.r.NewStyle().Foreground(colN7) - hint := fmt.Sprintf("↑/↓ navigate · ←/→ page · ↵ mlflow · q quit · row %d/%d", m.cursor+1, len(m.rows)) + hint := fmt.Sprintf("↑/↓ navigate · ←/→ page · ↵ mlflow · i info · L logs · q quit · row %d/%d", m.cursor+1, len(m.rows)) switch { case m.loadErr != nil: hint += " (load failed)" @@ -239,6 +332,32 @@ func (m listModel) renderHint() string { return faint.Render(hint) } +// fetchRunDetail returns a tea.Cmd that fetches run details in the background. +func (m listModel) fetchRunDetail(runID int64) tea.Cmd { + ctx := m.fetcher.ctx + w := m.fetcher.w + return func() tea.Msg { + body, err := runDetailText(ctx, w, runID) + if err != nil { + return detailMsg{title: "Run Details", err: err} + } + return detailMsg{title: "Run Details", body: body} + } +} + +// fetchRunLogs returns a tea.Cmd that fetches a log snapshot in the background. +func (m listModel) fetchRunLogs(runID int64) tea.Cmd { + ctx := m.fetcher.ctx + w := m.fetcher.w + return func() tea.Msg { + body, err := runLogsSnapshot(ctx, w, runID) + if err != nil { + return detailMsg{title: "Logs Snapshot", err: err} + } + return detailMsg{title: "Logs Snapshot", body: body} + } +} + // openURL opens a URL in the user's default browser, best-effort. func openURL(url string) tea.Cmd { return func() tea.Msg { @@ -246,3 +365,9 @@ func openURL(url string) tea.Cmd { return nil } } + +// parseRunID parses a run id string to int64. Rows carry a formatted int64, so +// this only fails on an unexpectedly malformed value. +func parseRunID(runIDStr string) (int64, error) { + return strconv.ParseInt(runIDStr, 10, 64) +} diff --git a/experimental/air/cmd/list_tui_render.go b/experimental/air/cmd/list_tui_render.go index 6a81deceb5c..ad70369230c 100644 --- a/experimental/air/cmd/list_tui_render.go +++ b/experimental/air/cmd/list_tui_render.go @@ -22,7 +22,7 @@ const ( colBlue = lipgloss.Color("#6CA8F0") // MLflow link ) -const mlflowColWidth = 18 +const mlflowColWidth = 22 // listStyles renders the runs table. The renderer carries the color profile, so // styles render plain under --no-color / non-tty. @@ -105,10 +105,21 @@ func (s listStyles) renderRow(cols listCols, r listRow, selected, links bool) st gutter = "▸" } + runIDLink := "" + if links { + runIDLink = r.RunURL + } + experimentLink := "" + if links { + experimentLink = r.ExperimentURL + } + + // Underline only cells that actually carry a link, so unlinked text isn't + // styled as clickable. cells := []string{ s.cell(base, gutter, 1, fg(colN7), false, false, ""), - s.cell(base, r.RunID, cols.runID, fg(colRunID), false, false, ""), - s.cell(base, r.Experiment, cols.experiment, fg(colN11), false, false, ""), + s.cell(base, r.RunID, cols.runID, fg(colRunID), false, runIDLink != "", runIDLink), + s.cell(base, r.Experiment, cols.experiment, fg(colN11), false, experimentLink != "", experimentLink), s.cell(base, "● "+r.Status, cols.status, fg(statusColor(r.Status)), false, false, ""), s.cell(base, startedDisplay(r), cols.started, fg(colN9), false, false, ""), s.cell(base, r.Duration, cols.duration, fg(colN9), true, false, ""), @@ -159,7 +170,11 @@ func (s listStyles) mlflowCell(base lipgloss.Style, r listRow, selected, links b if links { link = r.MLflowURL } - return s.cell(base, mlflowDisplay(r.MLflowURL), mlflowColWidth, fg, false, true, link) + label := r.MLflowLabel + if label == "" || label == "-" { + label = "-" + } + return s.cell(base, label, mlflowColWidth, fg, false, true, link) } // statusColor maps an air run status word to its data color. @@ -189,29 +204,6 @@ func startedDisplay(r listRow) string { return s } -// mlflowDisplay shortens an MLflow run URL to a "…/runs/" label; the -// OSC 8 target keeps the full URL. -func mlflowDisplay(url string) string { - id := mlflowRunID(url) - if id == "" { - return truncate(url, mlflowColWidth) - } - if len(id) > 8 { - id = id[:8] + "…" - } - return "…/runs/" + id -} - -// mlflowRunID extracts the run-id path segment from an MLflow URL. -func mlflowRunID(url string) string { - _, after, ok := strings.Cut(url, "/runs/") - if !ok { - return "" - } - id, _, _ := strings.Cut(after, "/") - return id -} - // pad pads (or truncates) s to a visible width of n, right-aligned when right is // set. It measures visible width, so it is safe on styled strings. func pad(s string, n int, right bool) string { diff --git a/experimental/air/cmd/list_tui_test.go b/experimental/air/cmd/list_tui_test.go index ba709bbf35c..8b52417ee23 100644 --- a/experimental/air/cmd/list_tui_test.go +++ b/experimental/air/cmd/list_tui_test.go @@ -6,7 +6,6 @@ import ( "testing" tea "github.com/charmbracelet/bubbletea" - "github.com/charmbracelet/lipgloss" "github.com/databricks/cli/libs/cmdio" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" @@ -16,9 +15,9 @@ import ( // absent MLflow link, and a still-running (no end) row. func testListRows() []listRow { return []listRow{ - {RunID: "1", Experiment: "qwen-train", User: "me@example.com", Status: "SUCCESS", StartedAt: new("2026-06-05T17:32:39.000000+00:00"), Duration: "1m 14s", MLflowURL: "https://h/ml/experiments/E/runs/04c41514fbb0/artifacts/logs/node_0", Accelerators: "8x H100"}, - {RunID: "2", Experiment: "llama-train", User: "me@example.com", Status: "RUNNING", StartedAt: new("2026-06-05T18:43:24.000000+00:00"), Duration: "3m 32s", MLflowURL: "-", Accelerators: "1x A10"}, - {RunID: "3", Experiment: "mixtral", User: "me@example.com", Status: "FAILED", StartedAt: nil, Duration: "-", MLflowURL: "-", Accelerators: "-"}, + {RunID: "1", Experiment: "qwen-train", User: "me@example.com", Status: "SUCCESS", StartedAt: new("2026-06-05T17:32:39.000000+00:00"), Duration: "1m 14s", MLflowURL: "https://h/ml/experiments/E/runs/04c41514fbb0/artifacts/logs/node_0", MLflowLabel: "qwen-run-001", Accelerators: "8x H100"}, + {RunID: "2", Experiment: "llama-train", User: "me@example.com", Status: "RUNNING", StartedAt: new("2026-06-05T18:43:24.000000+00:00"), Duration: "3m 32s", MLflowURL: "-", MLflowLabel: "-", Accelerators: "1x A10"}, + {RunID: "3", Experiment: "mixtral", User: "me@example.com", Status: "FAILED", StartedAt: nil, Duration: "-", MLflowURL: "-", MLflowLabel: "-", Accelerators: "-"}, } } @@ -161,7 +160,7 @@ func TestListModelView(t *testing.T) { for _, want := range []string{ "Run ID", "Experiment", "Status", "Started", "Duration", "MLflow", "User", "Accelerators", "qwen-train", "● SUCCESS", "● RUNNING", "● FAILED", - "…/runs/04c41514…", // shortened MLflow link + "qwen-run-001", // MLflow run label "2026-06-05T17:32:39", // started trimmed to seconds "▸", // selection gutter on the first row "↑/↓ navigate", // hint line @@ -176,7 +175,7 @@ func TestStaticListTable(t *testing.T) { assert.NotContains(t, out, "\x1b") assert.NotContains(t, out, "▸", "static table has no selection") - for _, want := range []string{"Run ID", "1", "qwen-train", "…/runs/04c41514…", "Accelerators"} { + for _, want := range []string{"Run ID", "1", "qwen-train", "qwen-run-001", "Accelerators"} { assert.Contains(t, out, want) } @@ -197,15 +196,86 @@ func TestStartedDisplay(t *testing.T) { assert.Equal(t, "2026-06-05T17:32:39", startedDisplay(listRow{StartedAt: new("2026-06-05T17:32:39.000000+00:00")})) } -func TestMLflowDisplay(t *testing.T) { - assert.Equal(t, "…/runs/04c41514…", mlflowDisplay("https://h/ml/experiments/E/runs/04c41514fbb0/artifacts/logs/node_0")) - assert.Equal(t, "…/runs/run1", mlflowDisplay("https://h/ml/experiments/E/runs/run1/artifacts/logs/node_0")) - assert.LessOrEqual(t, lipgloss.Width(mlflowDisplay("https://h/no-runs/here")), mlflowColWidth) +func TestRenderRowHyperlinks(t *testing.T) { + r, _ := cmdio.NewRenderer(cmdio.MockDiscard(t.Context()), io.Discard) + styles := newListStyles(r) + row := listRow{ + RunID: "1", Experiment: "exp", Status: "SUCCESS", Duration: "-", Accelerators: "-", + RunURL: "https://h/jobs/runs/1?o=2", ExperimentURL: "https://h/ml/experiments/E?o=2", + MLflowURL: "https://h/ml/experiments/E/runs/rid", MLflowLabel: "my-run", + } + cols := computeListCols([]listRow{row}) + + linked := styles.renderRow(cols, row, false, true) + assert.Contains(t, linked, "\x1b]8;;https://h/jobs/runs/1?o=2", "run id links to the dashboard") + assert.Contains(t, linked, "\x1b]8;;https://h/ml/experiments/E?o=2", "experiment links to the experiment page") + + plain := styles.renderRow(cols, row, false, false) + assert.NotContains(t, plain, "\x1b]8;;", "no links when links are disabled") +} + +func TestListModelInfoKeyOpensDetail(t *testing.T) { + r, _ := cmdio.NewRenderer(cmdio.MockDiscard(t.Context()), io.Discard) + f := &runFetcher{ctx: t.Context(), w: newTestWorkspaceClient(t, "https://x.test")} + m := newListModel(r, f, testListRows(), false) + + // `i` opens the detail pane in a loading state and dispatches a fetch (not run here). + next, cmd := m.Update(tea.KeyMsg{Type: tea.KeyRunes, Runes: []rune("i")}) + m = next.(listModel) + assert.Equal(t, modeDetail, m.mode) + assert.True(t, m.detailLoading) + assert.NotNil(t, cmd) +} + +// detailLoadingModel returns a sized model in the detail-loading state, as if +// the user just pressed `i` and the fetch is still in flight. +func detailLoadingModel(t *testing.T) listModel { + t.Helper() + r, _ := cmdio.NewRenderer(cmdio.MockDiscard(t.Context()), io.Discard) + f := &runFetcher{ctx: t.Context(), w: newTestWorkspaceClient(t, "https://x.test")} + m := newListModel(r, f, testListRows(), false) + next, _ := m.Update(tea.WindowSizeMsg{Width: 80, Height: 24}) + m = next.(listModel) + next, _ = m.Update(tea.KeyMsg{Type: tea.KeyRunes, Runes: []rune("i")}) + m = next.(listModel) + require.Equal(t, modeDetail, m.mode) + require.True(t, m.detailLoading) + return m +} + +func TestListModelDetailPaneAndBack(t *testing.T) { + m := detailLoadingModel(t) + + // A resolved detailMsg fills the pane. + next, _ := m.Update(detailMsg{title: "Run Details", body: "hello from the detail pane"}) + m = next.(listModel) + require.Equal(t, modeDetail, m.mode) + assert.False(t, m.detailLoading) + view := m.View() + assert.Contains(t, view, "Run Details") + assert.Contains(t, view, "hello from the detail pane") + assert.Contains(t, view, "esc back") + + // esc returns to the list. + next, _ = m.Update(tea.KeyMsg{Type: tea.KeyEsc}) + m = next.(listModel) + assert.Equal(t, modeList, m.mode) + assert.Contains(t, m.View(), "Run ID") } -func TestMLflowRunID(t *testing.T) { - assert.Equal(t, "abc123", mlflowRunID("https://h/ml/experiments/1/runs/abc123/artifacts")) - assert.Empty(t, mlflowRunID("https://h/no-runs-here")) +func TestListModelDetailLateMsgDropped(t *testing.T) { + m := detailLoadingModel(t) + + // User escapes back to the list before the fetch resolves. + next, _ := m.Update(tea.KeyMsg{Type: tea.KeyEsc}) + m = next.(listModel) + require.Equal(t, modeList, m.mode) + + // The late result must be dropped, not snap the user back into the pane. + next, _ = m.Update(detailMsg{title: "Run Details", body: "late result"}) + m = next.(listModel) + assert.Equal(t, modeList, m.mode) + assert.NotContains(t, m.View(), "late result") } func TestPadAndTruncate(t *testing.T) { 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/logdownload.go b/experimental/air/cmd/logdownload.go new file mode 100644 index 00000000000..e6764f49a6c --- /dev/null +++ b/experimental/air/cmd/logdownload.go @@ -0,0 +1,266 @@ +package aircmd + +import ( + "context" + "errors" + "fmt" + "io" + "os" + "path" + "path/filepath" + "slices" + + "github.com/databricks/cli/libs/cmdio" + "github.com/databricks/cli/libs/log" + "github.com/databricks/databricks-sdk-go" + "github.com/databricks/databricks-sdk-go/service/jobs" + "golang.org/x/sync/errgroup" +) + +// downloadConcurrency caps how many nodes download at once, to avoid hammering +// the artifact store on a wide run. +const downloadConcurrency = 8 + +// errNodeOutOfRange marks --node naming a node the run doesn't have. It is user +// input, so the caller reports it as an invalid argument rather than a failure. +var errNodeOutOfRange = errors.New("invalid --node") + +// resolveNodeCount returns how many nodes a run used. +func resolveNodeCount(run *jobs.Run) (int, error) { + accelType, count := jobCompute(run) + if accelType == "" { + return 0, fmt.Errorf("run %d has no AI runtime compute config", run.RunId) + } + if count <= 0 { + return 0, fmt.Errorf("run %d reports %d accelerators", run.RunId, count) + } + g, err := parseGPUType(accelType) + if err != nil { + return 0, err + } + perNode, err := gpusPerNode(g) + if err != nil { + return 0, err + } + // Accelerators come in whole nodes, so a remainder means we can't map the + // count onto node indices. + if count%perNode != 0 { + return 0, fmt.Errorf("run %d reports %d %s accelerators, which is not a multiple of %d per node", run.RunId, count, accelType, perNode) + } + return count / perNode, nil +} + +// downloadLogs writes each node's logs to /logs/node_.log and +// prints a summary. An explicit --node downloads only that node; otherwise all of +// them. Logs come from MLflow artifacts, since Bricklens only streams. The +// returned bool is the run's outcome, so the exit code matches the streaming path. +func downloadLogs(ctx context.Context, w *databricks.WorkspaceClient, out io.Writer, req logRequest, status logRunStatus) (bool, error) { + run, err := w.Jobs.GetRun(ctx, jobs.GetRunRequest{RunId: req.runID}) + if err != nil { + return false, err + } + + numNodes, err := resolveNodeCount(run) + if err != nil { + return false, err + } + + nodes := make([]int, 0, numNodes) + if req.nodeSet { + if req.node >= numNodes { + return false, fmt.Errorf("%w %d: run has %d node(s), indexed 0 to %d", errNodeOutOfRange, req.node, numNodes, numNodes-1) + } + nodes = append(nodes, req.node) + } else { + for n := range numNodes { + nodes = append(nodes, n) + } + } + + // A run with no logs is reported the same way as on the streaming path, so + // the message and exit code agree between them. + ids := mlflowIDs(ctx, w, run) + if ids == nil || ids.RunID == "" { + emitNoLogs(out, req, status) + return status.downloadOutcome(), nil + } + + dir, err := filepath.Abs(req.downloadTo) + if err != nil { + return false, err + } + // Created up front so a bad --download-to fails with a clear message before + // any download work happens. + if err := os.MkdirAll(dir, 0o755); err != nil { + return false, fmt.Errorf("failed to create %s: %w", dir, err) + } + + nodeLogs, failures, err := downloadAllNodeLogs(ctx, w, ids.RunID, dir, nodes, req.attempt) + if err != nil { + return false, err + } + for _, node := range sortedNodeKeys(failures) { + cmdio.LogString(ctx, fmt.Sprintf("warning: node %d: %s", node, failures[node])) + } + + if len(nodeLogs) == 0 { + // "No logs available" would be a lie when the logs exist but couldn't be + // fetched. The warnings above go to stderr, which a -o json consumer reading + // stdout never sees, so fail instead of reporting an empty run. + if len(failures) > 0 { + return false, fmt.Errorf("failed to download logs from any of %d node(s): %s", + len(nodes), failures[sortedNodeKeys(failures)[0]]) + } + emitNoLogs(out, req, status) + return status.downloadOutcome(), nil + } + + cmdio.LogString(ctx, fmt.Sprintf("Downloaded logs from %d of %d node(s) to %s", len(nodeLogs), len(nodes), dir)) + for _, node := range sortedNodeKeys(nodeLogs) { + // Flag it on the file's own line, not just in the warning above. + suffix := "" + if _, truncated := failures[node]; truncated { + suffix = " (incomplete)" + } + cmdio.LogString(ctx, fmt.Sprintf(" node %d: %s%s", node, nodeLogs[node], suffix)) + } + return status.downloadOutcome(), nil +} + +// downloadAllNodeLogs downloads the nodes' logs in parallel. It returns a +// node->path map for the nodes that had logs and a node->reason map for those +// that failed; a truncated node appears in both. The log-dir layout is run-wide, +// so it is probed once here rather than by every worker. +func downloadAllNodeLogs(ctx context.Context, w *databricks.WorkspaceClient, mlflowRunID, dir string, nodes []int, attempt int) (map[int]string, map[int]string, error) { + // -1 (latest) maps to attempt 0's directory, as on the streaming path. + attemptDir := max(attempt, 0) + withAttempt, err := discoverAttemptPrefix(ctx, w, mlflowRunID, attemptDir) + if err != nil { + return nil, nil, err + } + + paths := make([]string, len(nodes)) + reasons := make([]string, len(nodes)) + g, gctx := errgroup.WithContext(ctx) + g.SetLimit(downloadConcurrency) + for i, node := range nodes { + g.Go(func() error { + path, err := downloadNodeLog(gctx, w, mlflowRunID, node, attemptDir, withAttempt, dir) + switch { + case errors.Is(err, context.Canceled), errors.Is(err, context.DeadlineExceeded): + // Interrupting the command must not look like a node with no logs. + return err + case err != nil: + // One bad node shouldn't abort the rest. A truncated node + // returns a path as well as an error, so keep both. + reasons[i] = err.Error() + paths[i] = path + default: + paths[i] = path + } + return nil + }) + } + if err := g.Wait(); err != nil { + return nil, nil, err + } + + nodeLogs := map[int]string{} + failures := map[int]string{} + for i, node := range nodes { + if paths[i] != "" { + nodeLogs[node] = paths[i] + } + if reasons[i] != "" { + failures[node] = reasons[i] + } + } + return nodeLogs, failures, nil +} + +// downloadNodeLog streams a node's chunks in order into dir/logs/node_.log, +// returning the path, or "" if the node logged nothing. A failed chunk is skipped +// rather than ending the walk, so a partial download returns both a path and an +// error naming the gaps. The bytes are copied verbatim: a download should +// reproduce the log exactly, so it must not round-trip through lines (which would +// rewrite line endings and cap long lines). +func downloadNodeLog(ctx context.Context, w *databricks.WorkspaceClient, mlflowRunID string, node, attempt int, withAttempt bool, dir string) (string, error) { + logDir := constructLogPath(node, attempt, withAttempt) + chunks, err := listLogChunks(ctx, w, mlflowRunID, logDir) + if err != nil { + return "", err + } + if len(chunks) == 0 { + // The listing can lag behind the sidecar, so fall back to chunk 0 as the + // streaming path does. + chunks = []logChunk{{index: 0, path: path.Join(logDir, chunkFileName(0))}} + } + + outPath := filepath.Join(dir, "logs", fmt.Sprintf("node_%d.log", node)) + if err := os.MkdirAll(filepath.Dir(outPath), 0o755); err != nil { + return "", err + } + f, err := os.Create(outPath) + if err != nil { + return "", err + } + defer f.Close() + + // Skip a failed chunk and keep going: the tail usually holds the failure + // signature, so losing it to an early bad chunk is worse than a gap. Cancellation + // still aborts, since every remaining chunk would fail too. + var written int64 + var missing []int + for _, chunk := range chunks { + n, err := copyArtifactTo(ctx, w, mlflowRunID, chunk.path, f) + switch { + case errors.Is(err, context.Canceled), errors.Is(err, context.DeadlineExceeded): + os.Remove(outPath) + return "", err + case err != nil: + log.Debugf(ctx, "air logs: node %d chunk %d failed: %v", node, chunk.index, err) + missing = append(missing, chunk.index) + default: + written += n + } + } + if written == 0 { + os.Remove(outPath) + if len(missing) > 0 { + return "", fmt.Errorf("every chunk failed to download (%d total)", len(missing)) + } + return "", nil + } + if len(missing) > 0 { + return outPath, fmt.Errorf("incomplete: chunk(s) %v failed to download", missing) + } + return outPath, nil +} + +// copyArtifactTo streams one artifact's bytes into dst and returns how many were +// written. +func copyArtifactTo(ctx context.Context, w *databricks.WorkspaceClient, mlflowRunID, artifactPath string, dst io.Writer) (int64, error) { + local, err := downloadArtifact(ctx, w, mlflowRunID, artifactPath) + if err != nil { + return 0, err + } + defer os.Remove(local) + + src, err := os.Open(local) + if err != nil { + return 0, err + } + defer src.Close() + return io.Copy(dst, src) +} + +// sortedNodeKeys returns the map's node ids in ascending order, so the summary +// prints deterministically. +func sortedNodeKeys(m map[int]string) []int { + keys := make([]int, 0, len(m)) + for k := range m { + keys = append(keys, k) + } + slices.Sort(keys) + return keys +} diff --git a/experimental/air/cmd/logdownload_test.go b/experimental/air/cmd/logdownload_test.go new file mode 100644 index 00000000000..b7717980f42 --- /dev/null +++ b/experimental/air/cmd/logdownload_test.go @@ -0,0 +1,554 @@ +package aircmd + +import ( + "bytes" + "context" + "net/http" + "net/http/httptest" + "os" + "path" + "path/filepath" + "strings" + "testing" + + "github.com/databricks/cli/libs/cmdio" + "github.com/databricks/databricks-sdk-go/service/jobs" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// airRunWithCompute builds a run reporting the given accelerator type and count. +func airRunWithCompute(accelType string, count int) *jobs.Run { + return &jobs.Run{ + RunId: 123, + Tasks: []jobs.RunTask{{ + RunId: 456, + AiRuntimeTask: &jobs.AiRuntimeTask{ + Deployments: []jobs.DeploymentSpec{{ + Compute: jobs.ComputeSpec{ + AcceleratorType: jobs.ComputeSpecAcceleratorType(accelType), + AcceleratorCount: count, + }, + }}, + }, + }}, + } +} + +func TestResolveNodeCount(t *testing.T) { + tests := []struct { + accelType string + count int + want int + }{ + {"GPU_1xA10", 2, 2}, + {"GPU_1xH100", 4, 4}, + {"GPU_8xH100", 16, 2}, + } + for _, tt := range tests { + n, err := resolveNodeCount(airRunWithCompute(tt.accelType, tt.count)) + require.NoError(t, err) + assert.Equal(t, tt.want, n) + } + + // A run with no AI runtime compute errors. + _, err := resolveNodeCount(&jobs.Run{RunId: 1}) + require.Error(t, err) + + // A count that isn't a whole number of nodes can't be mapped to node indices, + // so it errors rather than truncating to 0. + _, err = resolveNodeCount(airRunWithCompute("GPU_8xH100", 4)) + require.Error(t, err) + assert.Contains(t, err.Error(), "not a multiple of 8") + + // A zero count is reported as such, not as a missing config. + _, err = resolveNodeCount(airRunWithCompute("GPU_1xA10", 0)) + require.Error(t, err) +} + +// downloadServer serves the MLflow artifact chain: the artifact listing, a +// pre-signed URL pointing back at itself, and the chunk bytes. +func downloadServer(t *testing.T) *httptest.Server { + t.Helper() + var base string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/api/2.0/mlflow/artifacts/list": + if r.URL.Query().Get("path") == "logs" { + _, _ = w.Write([]byte(`{"files": [{"path": "logs/node_0", "is_dir": true}, {"path": "logs/node_1", "is_dir": true}]}`)) + } else { + _, _ = w.Write([]byte(`{"files": [{"path": "` + r.URL.Query().Get("path") + `/logs-0.chunk.txt", "file_size": 12}]}`)) + } + case "/api/2.0/mlflow/artifacts/credentials-for-read": + _, _ = w.Write([]byte(`{"credential_infos": [{"signed_uri": "` + base + `/presigned"}]}`)) + case "/presigned": + _, _ = w.Write([]byte("line one\nline two\n")) + default: + _, _ = w.Write([]byte(`{}`)) + } + })) + base = srv.URL + t.Cleanup(srv.Close) + return srv +} + +func TestDownloadNodeLogWritesConcatenatedChunks(t *testing.T) { + w := newTestWorkspaceClient(t, downloadServer(t).URL) + dir := t.TempDir() + + path, err := downloadNodeLog(t.Context(), w, "run1", 0, 0, false, dir) + require.NoError(t, err) + require.NotEmpty(t, path) + + got, err := os.ReadFile(path) + require.NoError(t, err) + assert.Equal(t, "line one\nline two\n", string(got)) + assert.Equal(t, filepath.Join(dir, "logs", "node_0.log"), path) +} + +func TestDownloadAllNodeLogs(t *testing.T) { + w := newTestWorkspaceClient(t, downloadServer(t).URL) + dir := t.TempDir() + + nodeLogs, failures, err := downloadAllNodeLogs(t.Context(), w, "run1", dir, []int{0, 1}, -1) + require.NoError(t, err) + require.Empty(t, failures) + require.Len(t, nodeLogs, 2) + assert.FileExists(t, nodeLogs[0]) + assert.FileExists(t, nodeLogs[1]) +} + +// fullDownloadServer also serves the run and its output, so downloadLogs can run +// end to end against a 2-node run. +func fullDownloadServer(t *testing.T) *httptest.Server { + t.Helper() + var base string + runGet := `{ + "run_id": 123, + "state": {"life_cycle_state": "TERMINATED", "result_state": "SUCCESS"}, + "tasks": [{"run_id": 456, "ai_runtime_task": {"deployments": [{"compute": {"accelerator_type": "GPU_1xA10", "accelerator_count": 2}}]}}] + }` + 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(runGet)) + case "/api/2.2/jobs/runs/get-output": + _, _ = w.Write([]byte(`{"ai_runtime_task_output": {"mlflow_experiment_id": "exp1", "mlflow_run_id": "run1"}}`)) + case "/api/2.0/mlflow/artifacts/list": + if r.URL.Query().Get("path") == "logs" { + _, _ = w.Write([]byte(`{"files": [{"path": "logs/node_0", "is_dir": true}, {"path": "logs/node_1", "is_dir": true}]}`)) + return + } + _, _ = w.Write([]byte(`{"files": [{"path": "` + r.URL.Query().Get("path") + `/logs-0.chunk.txt", "file_size": 6}]}`)) + case "/api/2.0/mlflow/artifacts/credentials-for-read": + _, _ = w.Write([]byte(`{"credential_infos": [{"signed_uri": "` + base + `/presigned"}]}`)) + case "/presigned": + _, _ = w.Write([]byte("hello\n")) + default: + _, _ = w.Write([]byte(`{"userName": "u@example.com", "workspace_id": 1}`)) + } + })) + base = srv.URL + t.Cleanup(srv.Close) + return srv +} + +func TestDownloadLogsAllNodes(t *testing.T) { + ctx := cmdio.MockDiscard(t.Context()) + w := newTestWorkspaceClient(t, fullDownloadServer(t).URL) + dir := t.TempDir() + + success, err := downloadLogs(ctx, w, &bytes.Buffer{}, logRequest{runID: 123, attempt: -1, downloadTo: dir}, logRunStatus{resultState: "SUCCESS"}) + require.NoError(t, err) + assert.True(t, success) + assert.FileExists(t, filepath.Join(dir, "logs", "node_0.log")) + assert.FileExists(t, filepath.Join(dir, "logs", "node_1.log")) +} + +func TestDownloadLogsSingleNode(t *testing.T) { + ctx := cmdio.MockDiscard(t.Context()) + w := newTestWorkspaceClient(t, fullDownloadServer(t).URL) + dir := t.TempDir() + + success, err := downloadLogs(ctx, w, &bytes.Buffer{}, logRequest{runID: 123, node: 1, nodeSet: true, attempt: -1, downloadTo: dir}, logRunStatus{resultState: "SUCCESS"}) + require.NoError(t, err) + assert.True(t, success) + assert.FileExists(t, filepath.Join(dir, "logs", "node_1.log")) + assert.NoFileExists(t, filepath.Join(dir, "logs", "node_0.log")) +} + +func TestDownloadLogsOutOfRangeNode(t *testing.T) { + ctx := cmdio.MockDiscard(t.Context()) + w := newTestWorkspaceClient(t, fullDownloadServer(t).URL) + dir := t.TempDir() + + _, err := downloadLogs(ctx, w, &bytes.Buffer{}, logRequest{runID: 123, node: 5, nodeSet: true, attempt: -1, downloadTo: dir}, logRunStatus{resultState: "SUCCESS"}) + require.Error(t, err) + assert.Contains(t, err.Error(), "invalid --node 5") +} + +func TestDownloadLogsExplicitNodeZero(t *testing.T) { + ctx := cmdio.MockDiscard(t.Context()) + w := newTestWorkspaceClient(t, fullDownloadServer(t).URL) + dir := t.TempDir() + + // An explicit --node 0 must download only node 0, unlike the default which + // downloads every node. + _, err := downloadLogs(ctx, w, &bytes.Buffer{}, logRequest{runID: 123, node: 0, nodeSet: true, attempt: -1, downloadTo: dir}, logRunStatus{resultState: "SUCCESS"}) + require.NoError(t, err) + assert.FileExists(t, filepath.Join(dir, "logs", "node_0.log")) + assert.NoFileExists(t, filepath.Join(dir, "logs", "node_1.log")) +} + +func TestDownloadLogsOutOfRangeNodeIsInvalidArgs(t *testing.T) { + ctx := cmdio.MockDiscard(t.Context()) + w := newTestWorkspaceClient(t, fullDownloadServer(t).URL) + + // The sentinel lets the caller classify this as bad input rather than a + // transient, retryable failure. + _, err := downloadLogs(ctx, w, &bytes.Buffer{}, logRequest{runID: 123, node: 5, nodeSet: true, attempt: -1, downloadTo: t.TempDir()}, logRunStatus{resultState: "SUCCESS"}) + require.ErrorIs(t, err, errNodeOutOfRange) +} + +// noLogsDownloadServer serves a run with no resolvable MLflow run, so there is +// nothing to download. +func noLogsDownloadServer(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": 123, + "state": {"life_cycle_state": "TERMINATED", "result_state": "SUCCESS"}, + "tasks": [{"run_id": 456, "ai_runtime_task": {"deployments": [{"compute": {"accelerator_type": "GPU_1xA10", "accelerator_count": 2}}]}}] + }`)) + default: + _, _ = w.Write([]byte(`{}`)) + } + })) + t.Cleanup(srv.Close) + return srv +} + +func TestDownloadLogsNoLogsMatchesStreamingPath(t *testing.T) { + ctx := cmdio.MockDiscard(t.Context()) + w := newTestWorkspaceClient(t, noLogsDownloadServer(t).URL) + + // A SUCCESS run with no logs reports it and still succeeds, exactly as the + // streaming path does — no error, no non-zero exit. + var buf bytes.Buffer + success, err := downloadLogs(ctx, w, &buf, logRequest{runID: 123, attempt: -1, downloadTo: t.TempDir()}, logRunStatus{lifeCycleState: "TERMINATED", resultState: "SUCCESS"}) + require.NoError(t, err) + assert.True(t, success) + assert.Contains(t, buf.String(), "No logs available for run 123") +} + +// partialFailureServer fails node 1's chunk listing so one node succeeds and the +// other doesn't. +func partialFailureServer(t *testing.T) *httptest.Server { + t.Helper() + var base string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/api/2.0/mlflow/artifacts/list": + p := r.URL.Query().Get("path") + switch p { + case "logs": + _, _ = w.Write([]byte(`{"files": [{"path": "logs/node_0", "is_dir": true}, {"path": "logs/node_1", "is_dir": true}]}`)) + case "logs/node_1": + w.WriteHeader(http.StatusInternalServerError) + _, _ = w.Write([]byte(`{"error_code": "INTERNAL", "message": "boom"}`)) + default: + _, _ = w.Write([]byte(`{"files": [{"path": "` + p + `/logs-0.chunk.txt"}]}`)) + } + case "/api/2.0/mlflow/artifacts/credentials-for-read": + _, _ = w.Write([]byte(`{"credential_infos": [{"signed_uri": "` + base + `/presigned"}]}`)) + case "/presigned": + _, _ = w.Write([]byte("ok\n")) + default: + _, _ = w.Write([]byte(`{}`)) + } + })) + base = srv.URL + t.Cleanup(srv.Close) + return srv +} + +func TestDownloadAllNodeLogsReportsPartialFailure(t *testing.T) { + w := newTestWorkspaceClient(t, partialFailureServer(t).URL) + dir := t.TempDir() + + // Node 1 fails, but node 0 still downloads and the failure is reported rather + // than silently dropped. + nodeLogs, failures, err := downloadAllNodeLogs(t.Context(), w, "run1", dir, []int{0, 1}, -1) + require.NoError(t, err) + require.Len(t, nodeLogs, 1) + assert.FileExists(t, nodeLogs[0]) + require.Contains(t, failures, 1) + assert.NotEmpty(t, failures[1]) +} + +func TestDownloadAllNodeLogsPropagatesCancellation(t *testing.T) { + w := newTestWorkspaceClient(t, downloadServer(t).URL) + ctx, cancel := context.WithCancel(t.Context()) + cancel() + + // A cancelled command must surface the cancellation, not look like a run with + // no logs. + _, _, err := downloadAllNodeLogs(ctx, w, "run1", t.TempDir(), []int{0, 1}, -1) + require.Error(t, err) + assert.ErrorIs(t, err, context.Canceled) +} + +// attemptPrefixServer serves the attempt-prefixed layout (logs/attempt_N/node_M). +func attemptPrefixServer(t *testing.T) *httptest.Server { + t.Helper() + var base string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/api/2.0/mlflow/artifacts/list": + p := r.URL.Query().Get("path") + if p == "logs" { + _, _ = w.Write([]byte(`{"files": [{"path": "logs/attempt_0", "is_dir": true}]}`)) + return + } + _, _ = w.Write([]byte(`{"files": [{"path": "` + p + `/logs-0.chunk.txt"}]}`)) + case "/api/2.0/mlflow/artifacts/credentials-for-read": + // Echo the path so the test can prove the attempt-prefixed dir was used. + _, _ = w.Write([]byte(`{"credential_infos": [{"signed_uri": "` + base + `/presigned?p=` + r.URL.Query().Get("path") + `"}]}`)) + case "/presigned": + _, _ = w.Write([]byte(r.URL.Query().Get("p") + "\n")) + default: + _, _ = w.Write([]byte(`{}`)) + } + })) + base = srv.URL + t.Cleanup(srv.Close) + return srv +} + +func TestDownloadAllNodeLogsAttemptPrefixedLayout(t *testing.T) { + w := newTestWorkspaceClient(t, attemptPrefixServer(t).URL) + dir := t.TempDir() + + nodeLogs, _, err := downloadAllNodeLogs(t.Context(), w, "run1", dir, []int{0}, -1) + require.NoError(t, err) + require.Len(t, nodeLogs, 1) + + body, err := os.ReadFile(nodeLogs[0]) + require.NoError(t, err) + assert.Contains(t, string(body), "logs/attempt_0/node_0") +} + +// truncatingServer lists two chunks for node 0 but fails the second one's +// credential request, so the node downloads partially. +func truncatingServer(t *testing.T) *httptest.Server { + t.Helper() + var base string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/api/2.0/mlflow/artifacts/list": + p := r.URL.Query().Get("path") + if p == "logs" { + _, _ = w.Write([]byte(`{"files": [{"path": "logs/node_0", "is_dir": true}]}`)) + return + } + _, _ = w.Write([]byte(`{"files": [ + {"path": "` + p + `/logs-0.chunk.txt"}, + {"path": "` + p + `/logs-1.chunk.txt"} + ]}`)) + case "/api/2.0/mlflow/artifacts/credentials-for-read": + if strings.HasSuffix(r.URL.Query().Get("path"), "logs-1.chunk.txt") { + w.WriteHeader(http.StatusInternalServerError) + _, _ = w.Write([]byte(`{"error_code": "INTERNAL", "message": "boom"}`)) + return + } + _, _ = w.Write([]byte(`{"credential_infos": [{"signed_uri": "` + base + `/presigned"}]}`)) + case "/presigned": + _, _ = w.Write([]byte("first chunk\n")) + default: + _, _ = w.Write([]byte(`{}`)) + } + })) + base = srv.URL + t.Cleanup(srv.Close) + return srv +} + +func TestDownloadNodeLogReportsFailedChunk(t *testing.T) { + w := newTestWorkspaceClient(t, truncatingServer(t).URL) + dir := t.TempDir() + + // Chunk 0 succeeded and chunk 1 failed: keep the bytes, but return an error + // so the gap isn't silent. + path, err := downloadNodeLog(t.Context(), w, "run1", 0, 0, false, dir) + require.Error(t, err) + assert.Contains(t, err.Error(), "chunk(s) [1]") + require.NotEmpty(t, path) + + got, err := os.ReadFile(path) + require.NoError(t, err) + assert.Equal(t, "first chunk\n", string(got)) +} + +// middleGapServer serves three chunks for node 0 and fails only the middle one, +// so the walk has to continue past a gap to reach the last chunk. +func middleGapServer(t *testing.T) *httptest.Server { + t.Helper() + var base string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/api/2.0/mlflow/artifacts/list": + p := r.URL.Query().Get("path") + if p == "logs" { + _, _ = w.Write([]byte(`{"files": [{"path": "logs/node_0", "is_dir": true}]}`)) + return + } + _, _ = w.Write([]byte(`{"files": [ + {"path": "` + p + `/logs-0.chunk.txt"}, + {"path": "` + p + `/logs-1.chunk.txt"}, + {"path": "` + p + `/logs-2.chunk.txt"} + ]}`)) + case "/api/2.0/mlflow/artifacts/credentials-for-read": + p := r.URL.Query().Get("path") + if strings.HasSuffix(p, "logs-1.chunk.txt") { + w.WriteHeader(http.StatusInternalServerError) + _, _ = w.Write([]byte(`{"error_code": "INTERNAL", "message": "boom"}`)) + return + } + _, _ = w.Write([]byte(`{"credential_infos": [{"signed_uri": "` + base + `/presigned?p=` + p + `"}]}`)) + case "/presigned": + _, _ = w.Write([]byte(path.Base(r.URL.Query().Get("p")) + "\n")) + default: + _, _ = w.Write([]byte(`{}`)) + } + })) + base = srv.URL + t.Cleanup(srv.Close) + return srv +} + +func TestDownloadNodeLogSkipsGapAndKeepsTail(t *testing.T) { + w := newTestWorkspaceClient(t, middleGapServer(t).URL) + dir := t.TempDir() + + // The tail usually carries the failure signature, so a bad middle chunk must + // not cost us the last one. + p, err := downloadNodeLog(t.Context(), w, "run1", 0, 0, false, dir) + require.Error(t, err) + assert.Contains(t, err.Error(), "chunk(s) [1]") + + got, err := os.ReadFile(p) + require.NoError(t, err) + assert.Equal(t, "logs-0.chunk.txt\nlogs-2.chunk.txt\n", string(got)) +} + +func TestDownloadNodeLogErrorsWhenEveryChunkFails(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/api/2.0/mlflow/artifacts/list": + p := r.URL.Query().Get("path") + if p == "logs" { + _, _ = w.Write([]byte(`{"files": [{"path": "logs/node_0", "is_dir": true}]}`)) + return + } + _, _ = w.Write([]byte(`{"files": [{"path": "` + p + `/logs-0.chunk.txt"}]}`)) + case "/api/2.0/mlflow/artifacts/credentials-for-read": + w.WriteHeader(http.StatusInternalServerError) + _, _ = w.Write([]byte(`{"error_code": "INTERNAL", "message": "boom"}`)) + default: + _, _ = w.Write([]byte(`{}`)) + } + })) + t.Cleanup(srv.Close) + + // Nothing downloaded because of failures, not because the node was silent: + // report it rather than returning an empty "no logs" result. + p, err := downloadNodeLog(t.Context(), newTestWorkspaceClient(t, srv.URL), "run1", 0, 0, false, t.TempDir()) + require.Error(t, err) + assert.Contains(t, err.Error(), "every chunk failed") + assert.Empty(t, p) +} + +func TestDownloadAllNodeLogsKeepsTruncatedNode(t *testing.T) { + w := newTestWorkspaceClient(t, truncatingServer(t).URL) + + // A truncated node lands in both maps, so it is still listed as downloaded + // while being reported as incomplete. + nodeLogs, failures, err := downloadAllNodeLogs(t.Context(), w, "run1", t.TempDir(), []int{0}, -1) + require.NoError(t, err) + require.Contains(t, nodeLogs, 0) + assert.FileExists(t, nodeLogs[0]) + require.Contains(t, failures, 0) + assert.Contains(t, failures[0], "chunk(s) [1]") +} + +func TestDownloadOutcomeTreatsActiveRunAsSuccess(t *testing.T) { + // An active run has no result state yet. Fetching its logs succeeded, so the + // command must exit 0 rather than report the run as failed. + for _, lc := range []string{"RUNNING", "PENDING", "QUEUED", "BLOCKED"} { + assert.True(t, logRunStatus{lifeCycleState: lc}.downloadOutcome(), lc) + } + + // A terminal run still decides the exit code by its outcome. + assert.True(t, logRunStatus{lifeCycleState: "TERMINATED", resultState: "SUCCESS"}.downloadOutcome()) + assert.False(t, logRunStatus{lifeCycleState: "TERMINATED", resultState: "FAILED"}.downloadOutcome()) + assert.False(t, logRunStatus{lifeCycleState: "TERMINATED", resultState: "CANCELED"}.downloadOutcome()) + assert.False(t, logRunStatus{lifeCycleState: "INTERNAL_ERROR"}.downloadOutcome()) +} + +func TestDownloadLogsActiveRunExitsZero(t *testing.T) { + ctx := cmdio.MockDiscard(t.Context()) + w := newTestWorkspaceClient(t, fullDownloadServer(t).URL) + + // Downloading a still-running run's logs is not a failure. + success, err := downloadLogs(ctx, w, &bytes.Buffer{}, + logRequest{runID: 123, attempt: -1, downloadTo: t.TempDir()}, + logRunStatus{lifeCycleState: "RUNNING"}) + require.NoError(t, err) + assert.True(t, success) +} + +// allNodesFailServer serves a FAILED 2-node run whose chunk credentials all 404, +// so every node's download fails. +func allNodesFailServer(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": 123, "state": {"life_cycle_state": "TERMINATED", "result_state": "FAILED"}, + "tasks": [{"run_id": 456, "ai_runtime_task": {"deployments": [{"compute": {"accelerator_type": "GPU_1xA10", "accelerator_count": 2}}]}}]}`)) + case "/api/2.2/jobs/runs/get-output": + _, _ = w.Write([]byte(`{"ai_runtime_task_output": {"mlflow_experiment_id": "exp1", "mlflow_run_id": "run1"}}`)) + case "/api/2.0/mlflow/artifacts/list": + p := r.URL.Query().Get("path") + if p == "logs" { + _, _ = w.Write([]byte(`{"files": [{"path": "logs/node_0", "is_dir": true}, {"path": "logs/node_1", "is_dir": true}]}`)) + return + } + _, _ = w.Write([]byte(`{"files": [{"path": "` + p + `/logs-0.chunk.txt"}]}`)) + case "/api/2.0/mlflow/artifacts/credentials-for-read": + w.WriteHeader(http.StatusNotFound) + _, _ = w.Write([]byte(`{"error_code": "NOT_FOUND", "message": "gone"}`)) + default: + _, _ = w.Write([]byte(`{"userName": "u@example.com"}`)) + } + })) + t.Cleanup(srv.Close) + return srv +} + +func TestDownloadLogsAllNodesFailedIsNotReportedAsNoLogs(t *testing.T) { + ctx := cmdio.MockDiscard(t.Context()) + w := newTestWorkspaceClient(t, allNodesFailServer(t).URL) + var stdout bytes.Buffer + + // The logs exist but couldn't be fetched. Reporting "No logs available" would + // tell a caller the run produced nothing, so this fails instead. + _, err := downloadLogs(ctx, w, &stdout, + logRequest{runID: 123, attempt: -1, downloadTo: t.TempDir(), jsonOutput: true}, + logRunStatus{lifeCycleState: "TERMINATED", resultState: "FAILED"}) + require.Error(t, err) + assert.Contains(t, err.Error(), "failed to download logs from any of 2 node(s)") + assert.NotContains(t, stdout.String(), "No logs available") +} diff --git a/experimental/air/cmd/logmlflow.go b/experimental/air/cmd/logmlflow.go new file mode 100644 index 00000000000..debddc0d618 --- /dev/null +++ b/experimental/air/cmd/logmlflow.go @@ -0,0 +1,304 @@ +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.downloadOutcome(), 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.downloadOutcome(), 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.downloadOutcome(), nil + } + + if len(lines) > target { + lines = lines[len(lines)-target:] + } + for _, line := range lines { + emitLogLine(out, req, line) + } + return status.downloadOutcome(), 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 + // A map query is serialized per value with %v, so a []string becomes the + // literal "[path]". The backend signs that bogus path and still returns 200, + // surfacing only as a 404 on the download. + query := map[string]any{ + "run_id": mlflowRunID, + "path": 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..44837999386 --- /dev/null +++ b/experimental/air/cmd/logmlflow_test.go @@ -0,0 +1,141 @@ +package aircmd + +import ( + "bytes" + "net/http" + "net/http/httptest" + "os" + "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) +} + +func TestDownloadArtifactSendsUnbracketedPath(t *testing.T) { + var gotPath string + var base string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/api/2.0/mlflow/artifacts/credentials-for-read": + gotPath = r.URL.Query().Get("path") + _, _ = w.Write([]byte(`{"credential_infos": [{"signed_uri": "` + base + `/presigned"}]}`)) + case "/presigned": + _, _ = w.Write([]byte("bytes\n")) + default: + _, _ = w.Write([]byte(`{}`)) + } + })) + base = srv.URL + t.Cleanup(srv.Close) + + local, err := downloadArtifact(t.Context(), newTestWorkspaceClient(t, srv.URL), "run1", "logs/node_0/logs-0.chunk.txt") + require.NoError(t, err) + t.Cleanup(func() { os.Remove(local) }) + + // The backend signs whatever path it is given and returns 200 even for a + // bracketed one, so only the download 404s. Assert on the path sent. + assert.Equal(t, "logs/node_0/logs-0.chunk.txt", gotPath) +} diff --git a/experimental/air/cmd/logs.go b/experimental/air/cmd/logs.go index c34fb62a7df..19861b9e957 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,162 @@ 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() + + // --review is not yet implemented; reject rather than silently ignore. + if review { + return renderError(ctx, cmd, "INVALID_ARGS", "PERMANENT", false, + errors.New("--review is not implemented yet")) + } + + // A download always writes the full log, so a tail or time window would be + // silently dropped. + if downloadTo != "" && (cmd.Flags().Changed("lines") || minutes > 0) { + return renderError(ctx, cmd, "INVALID_ARGS", "PERMANENT", false, + errors.New("--download-to writes complete logs, so it cannot be combined with --lines or --minutes")) + } + + // --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, + nodeSet: cmd.Flags().Changed("node"), + attempt: retry, + windowMinutes: minutes, + tailLines: tailLines, + downloadTo: downloadTo, + 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)) + } + + // --download-to writes each node's logs to disk instead of streaming. + if req.downloadTo != "" { + success, err := downloadLogs(ctx, w, cmd.OutOrStdout(), req, status) + if errors.Is(err, errNodeOutOfRange) { + return renderError(ctx, cmd, "INVALID_ARGS", "PERMANENT", false, err) + } + if err != nil { + return renderError(ctx, cmd, "INTERNAL_ERROR", "TRANSIENT", true, + fmt.Errorf("failed to download logs for run %d: %w", req.runID, err)) + } + if !success { + return root.ErrAlreadyPrinted + } + return nil + } + + // 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..9ea9026e01e --- /dev/null +++ b/experimental/air/cmd/logs_test.go @@ -0,0 +1,245 @@ +package aircmd + +import ( + "bytes" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "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: "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 is still active, so there is no failure to report — printing the + // logs succeeded, and the command exits 0. + err := runLogs(ctx, cmd, logRequest{runID: 9, node: 0, attempt: 0, tailLines: -1}) + require.NoError(t, err) + 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..51e7e54a768 --- /dev/null +++ b/experimental/air/cmd/logstream.go @@ -0,0 +1,572 @@ +package aircmd + +import ( + "context" + "errors" + "fmt" + "io" + "net/http" + "slices" + "strings" + "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 + // statusMessageRefreshEveryNPolls throttles the status_message fetch so the + // waiting spinner doesn't issue a get-output on every poll tick. + statusMessageRefreshEveryNPolls = 5 +) + +// statusMessageType tags a client-facing message packed into +// ai_runtime_task_output.status_message as ":"; only STATUS-typed +// messages are surfaced. +const statusMessageType = "STATUS" + +// waitingForComputeStatus is the fallback shown while a PENDING run waits for +// accelerator compute. +const waitingForComputeStatus = "Waiting for accelerator compute capacity to become available..." + +// normalizeStatusMessage returns the payload of a "STATUS:" message, +// normalized for display (trailing "." stripped, "..." suffix added), or "" for +// any other type or an empty/absent message. +func normalizeStatusMessage(raw string) string { + messageType, payload, ok := strings.Cut(raw, ":") + if !ok || !strings.EqualFold(strings.TrimSpace(messageType), statusMessageType) { + return "" + } + payload = strings.TrimSpace(payload) + payload = strings.TrimRight(payload, ".") + payload = strings.TrimSpace(payload) + if payload == "" { + return "" + } + return payload + "..." +} + +// 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), +// persistently failing, or served every request successfully but never returned a +// record for a run whose logs may still be in MLflow. 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 + // nodeSet distinguishes an explicit --node 0 from the default, so a download + // knows whether to fetch one node or all of them. + nodeSet bool + // 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 + // downloadTo, when set, writes logs to that directory instead of stdout. + downloadTo string + // 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" +} + +// downloadOutcome is the exit status for a one-shot fetch, which unlike streaming +// can run against an active run. An active run has no result state yet, and +// not-yet-finished is not a failure, so only a terminal run decides the exit code. +func (s logRunStatus) downloadOutcome() bool { + return !s.terminal() || s.succeeded() +} + +// 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) +} + +// waitingSpinnerText returns the waiting-spinner text: the server-set STATUS +// message if present, else the compute-capacity message for a PENDING run, else +// the default "waiting for run to start". +func (st *bricklensStreamer) waitingSpinnerText() string { + if msg := st.serverStatusMessage(); msg != "" { + return msg + } + if st.status.lifeCycleState == "PENDING" { + return waitingForComputeStatus + } + return fmt.Sprintf("Waiting for run to start (node %d)...", st.req.node) +} + +// serverStatusMessage returns the run's server-set STATUS message (normalized for +// display), or "" if unavailable. The message lives on the latest task run's +// output, re-resolved each call so a retry's new task run is picked up. +// Best-effort: any fetch failure logs at debug and returns "". +func (st *bricklensStreamer) serverStatusMessage() string { + run, err := st.w.Jobs.GetRun(st.ctx, jobs.GetRunRequest{RunId: st.req.runID}) + if err != nil || len(run.Tasks) == 0 { + return "" + } + taskRunID := run.Tasks[len(run.Tasks)-1].RunId + out, err := st.w.Jobs.GetRunOutputByRunId(st.ctx, taskRunID) + if err != nil { + log.Debugf(st.ctx, "air logs: status_message fetch failed for run %d: %v", st.req.runID, err) + return "" + } + if out.AiRuntimeTaskOutput == nil { + return "" + } + return normalizeStatusMessage(out.AiRuntimeTaskOutput.StatusMessage) +} + +// 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 + // Throttled refresh of the waiting-spinner text: statusRefreshCounter gates the + // server status_message fetch to every Nth poll, and lastSpinnerText avoids + // redundant spinner updates. + statusRefreshCounter := 0 + lastSpinnerText := "" + 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 + // with the server-set status (throttled), so a run stuck waiting for compute + // shows why rather than a generic "waiting" message. + if !terminal && !st.firstLogSeen && st.updateSpinner != nil { + if statusRefreshCounter%statusMessageRefreshEveryNPolls == 0 { + if desired := st.waitingSpinnerText(); desired != lastSpinnerText { + st.updateSpinner(desired) + lastSpinnerText = desired + } + } + statusRefreshCounter++ + } + + // 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 { + // A successful but empty Bricklens stream isn't proof the run has no + // logs; they may be in MLflow (as --download-to reads). Fall back + // there, which owns the real no-logs report and the same exit code. + return false, errBricklensFeatureDisabled + } + 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 { + // An empty Bricklens tail doesn't mean the attempt has no logs; fall back to + // MLflow, which holds the immutable per-attempt artifacts. See the terminal + // branch in run. + return false, errBricklensFeatureDisabled + } + return st.status.downloadOutcome(), 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) +} + +// 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..be4a9d0c452 --- /dev/null +++ b/experimental/air/cmd/logstream_test.go @@ -0,0 +1,564 @@ +package aircmd + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "net/http" + "net/http/httptest" + "strconv" + "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 TestNormalizeStatusMessage(t *testing.T) { + tests := []struct { + raw string + want string + }{ + {"STATUS: Waiting for GPU capacity.", "Waiting for GPU capacity..."}, + {"STATUS:Waiting for GPU capacity", "Waiting for GPU capacity..."}, + {"status: provisioning", "provisioning..."}, // type match is case-insensitive + {"STATUS: done...", "done..."}, // trailing dots collapse to one "..." + {"INFO: not a status", ""}, // other type ignored + {"no type prefix", ""}, + {"STATUS:", ""}, // empty payload + {"STATUS: ", ""}, // whitespace-only payload + {"", ""}, + } + for _, tt := range tests { + assert.Equal(t, tt.want, normalizeStatusMessage(tt.raw), "raw=%q", tt.raw) + } +} + +func TestWaitingSpinnerText(t *testing.T) { + // A server that returns the run (with a task) and a STATUS-typed status_message. + newStreamer := func(t *testing.T, statusMessage, lifeCycle string) *bricklensStreamer { + 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": 1, "tasks": [{"run_id": 2}]}`)) + case "/api/2.2/jobs/runs/get-output": + _, _ = w.Write([]byte(`{"ai_runtime_task_output": {"status_message": ` + strconv.Quote(statusMessage) + `}}`)) + default: + _, _ = w.Write([]byte(`{}`)) + } + })) + t.Cleanup(srv.Close) + return &bricklensStreamer{ + ctx: t.Context(), + w: newTestWorkspaceClient(t, srv.URL), + req: logRequest{runID: 1, node: 0}, + status: logRunStatus{lifeCycleState: lifeCycle}, + } + } + + // Server STATUS message wins. + assert.Equal(t, "Waiting for GPU capacity...", + newStreamer(t, "STATUS: Waiting for GPU capacity", "PENDING").waitingSpinnerText()) + + // No status message + PENDING -> compute-capacity fallback. + assert.Equal(t, waitingForComputeStatus, + newStreamer(t, "", "PENDING").waitingSpinnerText()) + + // No status message + non-PENDING -> default "waiting for run to start". + assert.Equal(t, "Waiting for run to start (node 0)...", + newStreamer(t, "", "RUNNING").waitingSpinnerText()) +} + +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) +} + +// emptyLogsServer serves an empty Bricklens log response for any /logs request +// and a stub for everything else (SDK config probes, etc.). +func emptyLogsServer(t *testing.T) *httptest.Server { + t.Helper() + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if strings.HasSuffix(r.URL.Path, "/logs") { + _, _ = w.Write([]byte(`{"log_records": []}`)) + return + } + _, _ = w.Write([]byte(`{}`)) + })) + t.Cleanup(srv.Close) + return srv +} + +func TestStreamBricklensEmptyFallsBackToMLflow(t *testing.T) { + // Bricklens served every request but returned no record. That is not proof the + // run has no logs (they may be in MLflow), so the streamer must hand off via + // errBricklensFeatureDisabled and emit nothing itself, rather than reporting + // "No logs available" (the reported bug: the print path did, --download-to did not). + tests := []struct { + name string + req logRequest + status logRunStatus + }{ + { + name: "terminal run", + req: logRequest{runID: 123, node: 0, attempt: -1, tailLines: -1, jsonOutput: true}, + status: logRunStatus{lifeCycleState: "TERMINATED", resultState: "SUCCESS", endTimeMs: 1700000012000}, + }, + { + name: "static view of a past retry", + req: logRequest{runID: 123, node: 0, attempt: 0, tailLines: -1, staticView: true, jsonOutput: true}, + status: logRunStatus{lifeCycleState: "RUNNING"}, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + var buf bytes.Buffer + w := newTestWorkspaceClient(t, emptyLogsServer(t).URL) + _, err := streamBricklensLogs(t.Context(), w, &buf, tt.req, tt.status) + require.ErrorIs(t, err, errBricklensFeatureDisabled) + assert.Empty(t, buf.String(), "nothing should be emitted before the hand-off") + }) + } +} + +func TestStreamBricklensTerminalWithRecordsDoesNotFallBack(t *testing.T) { + // A terminal run whose Bricklens stream has records prints them and reports the + // run's outcome, without triggering the empty-result fallback. + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if strings.HasSuffix(r.URL.Path, "/logs") { + _, _ = w.Write([]byte(`{"log_records": [{"time_unix_nano": 1700000001000000000, "body": "hello", "node_index": 0}]}`)) + return + } + _, _ = w.Write([]byte(`{}`)) + })) + t.Cleanup(srv.Close) + + var buf bytes.Buffer + w := newTestWorkspaceClient(t, srv.URL) + status := logRunStatus{lifeCycleState: "TERMINATED", resultState: "SUCCESS", endTimeMs: 1700000012000} + ok, err := streamBricklensLogs(t.Context(), w, &buf, logRequest{runID: 123, node: 0, attempt: -1, tailLines: -1, jsonOutput: true}, status) + require.NoError(t, err) + assert.True(t, ok) + assert.Contains(t, buf.String(), `"line":"hello"`) +} + +func TestFetchLogsFallsBackToMLflowWhenBricklensEmpty(t *testing.T) { + // End-to-end repro: a terminal SUCCESS run whose Bricklens stream is empty but + // whose logs are in MLflow. The print path must fall back to MLflow and print + // them, exactly as --download-to already does. + var base string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch { + case strings.HasSuffix(r.URL.Path, "/logs"): + _, _ = w.Write([]byte(`{"log_records": []}`)) + case r.URL.Path == "/api/2.2/jobs/runs/get": + _, _ = w.Write([]byte(`{"run_id": 123, "state": {"life_cycle_state": "TERMINATED", "result_state": "SUCCESS"}, "tasks": [{"run_id": 456}]}`)) + case r.URL.Path == "/api/2.2/jobs/runs/get-output": + _, _ = w.Write([]byte(`{"ai_runtime_task_output": {"mlflow_experiment_id": "E1", "mlflow_run_id": "R1"}}`)) + case r.URL.Path == "/api/2.0/mlflow/artifacts/list": + _, _ = w.Write([]byte(`{"files": [{"path": "logs/node_0"}, {"path": "logs/node_0/logs-0.chunk.txt"}]}`)) + 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 1\nline 2\n")) + default: + _, _ = w.Write([]byte(`{}`)) + } + })) + base = srv.URL + t.Cleanup(srv.Close) + + var buf bytes.Buffer + w := newTestWorkspaceClient(t, srv.URL) + status := logRunStatus{lifeCycleState: "TERMINATED", resultState: "SUCCESS", endTimeMs: 1700000012000} + ok, err := fetchLogs(t.Context(), w, &buf, logRequest{runID: 123, node: 0, attempt: -1, tailLines: -1, jsonOutput: true}, status) + require.NoError(t, err) + assert.True(t, ok) + assert.Contains(t, buf.String(), `"line":"line 1"`) + assert.Contains(t, buf.String(), `"line":"line 2"`) +} + +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/mlflow.go b/experimental/air/cmd/mlflow.go index 070caefbb06..0eef26f80c4 100644 --- a/experimental/air/cmd/mlflow.go +++ b/experimental/air/cmd/mlflow.go @@ -4,6 +4,7 @@ import ( "context" "fmt" "strings" + "time" "github.com/databricks/cli/libs/log" "github.com/databricks/databricks-sdk-go" @@ -11,6 +12,15 @@ import ( "github.com/databricks/databricks-sdk-go/service/ml" ) +// mlflowLinkPollAttempts bounds the best-effort poll for a freshly-submitted +// run's MLflow IDs (see resolveMLflowIDsForRun), kept short so a bare `air run` +// returns promptly when the IDs aren't ready yet. +const mlflowLinkPollAttempts = 3 + +// mlflowLinkPollInterval is the delay between poll attempts. A var, not a const, +// so tests can shrink it and avoid a real sleep. +var mlflowLinkPollInterval = 500 * time.Millisecond + // mlflowIdentifiers are the experiment and run IDs MLflow assigns to a run. type mlflowIdentifiers struct { ExperimentID string @@ -62,6 +72,46 @@ func mlflowRunURL(host string, ids *mlflowIdentifiers) string { strings.TrimRight(host, "/"), ids.ExperimentID, ids.RunID) } +// mlflowExperimentURL links to the MLflow experiment page. Omits the ?o= query +// for consistency with mlflowRunURL and the run-submit dashboard URL. +func mlflowExperimentURL(host string, ids *mlflowIdentifiers) string { + return fmt.Sprintf("%s/ml/experiments/%s", strings.TrimRight(host, "/"), ids.ExperimentID) +} + +// resolveMLflowIDsForRun best-effort resolves a just-submitted run's MLflow IDs, +// polling because they are assigned only once the task run starts. Returns nil +// (treated as "no link", not an error) if they don't appear within the budget or +// the context is cancelled. +func resolveMLflowIDsForRun(ctx context.Context, w *databricks.WorkspaceClient, runID int64) *mlflowIdentifiers { + // The task run id is fixed at submit time, so resolve it once; only the MLflow + // output (runs/get-output) fills in later, so that is all we poll. + run, err := w.Jobs.GetRun(ctx, jobs.GetRunRequest{RunId: runID}) + if err != nil { + log.Debugf(ctx, "air run: could not fetch run %d for MLflow link: %v", runID, err) + return nil + } + if len(run.Tasks) == 0 { + return nil + } + // The MLflow output is attached to the task run, not the parent job run. + taskRunID := run.Tasks[len(run.Tasks)-1].RunId + + for attempt := range mlflowLinkPollAttempts { + if ids := mlflowIDsForTask(ctx, w, taskRunID); ids != nil { + return ids + } + if attempt == mlflowLinkPollAttempts-1 { + break + } + select { + case <-ctx.Done(): + return nil + case <-time.After(mlflowLinkPollInterval): + } + } + return nil +} + // fetchMLflowRunName fetches a run's MLflow run_name via the MLflow REST API, // returning "" if it can't be obtained. Best-effort, like the rest of the MLflow // enrichment. diff --git a/experimental/air/cmd/register_image.go b/experimental/air/cmd/register_image.go index 1d8b45044a7..d4195d0d8e2 100644 --- a/experimental/air/cmd/register_image.go +++ b/experimental/air/cmd/register_image.go @@ -1,33 +1,311 @@ package aircmd import ( + "context" + "errors" + "fmt" + "strings" + "time" + "github.com/databricks/cli/cmd/root" + "github.com/databricks/cli/libs/cmdctx" + "github.com/databricks/cli/libs/flags" + "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/service/iam" "github.com/spf13/cobra" ) +// imagePollInterval is how often waitForImageReady polls for a status change. +const imagePollInterval = 5 * time.Second + +// validateTagPolicy checks the deprecated --tag-policy value. Registration +// always re-checks the source registry, so "latest" and the empty default are +// no-ops. "auto" is rejected rather than silently remapped: it used to reuse a +// cached image, so honoring it as always-re-check would be a hidden change. +func validateTagPolicy(value string) error { + switch strings.ToLower(strings.TrimSpace(value)) { + case "", "latest": + return nil + case "auto": + return errors.New("--tag-policy auto is no longer supported: auto mode was removed and registration now always checks the source registry for the latest digest; omit the flag or use --tag-policy latest") + default: + return fmt.Errorf("invalid image tag policy %q: the only supported value is latest", value) + } +} + +// registerImageResult is the JSON payload for `air register-image`. It mirrors +// the Python CLI's success shape so existing consumers keep working. +type registerImageResult struct { + DockerImageURL string `json:"docker_image_url"` + ManifestSHA256 string `json:"manifest_sha256"` + Status string `json:"status"` + ImageUpdated bool `json:"image_updated"` + Cached bool `json:"cached"` +} + func newRegisterImageCommand() *cobra.Command { var ( - scope string - key string - interactiveAuth bool - tagPolicy string - timeoutMinutes int + tagPolicy string + timeoutMinutes int ) cmd := &cobra.Command{ Use: "register-image IMAGE_URL", Args: root.ExactArgs(1), Short: "Mirror a Docker image into the workspace registry", - RunE: func(cmd *cobra.Command, args []string) error { - return notImplemented("register-image") - }, + Long: `Mirror a Docker image into the workspace registry. + +Credentials for private images are discovered from your local Docker +configuration (run ` + "`docker login`" + ` first); there are no credential flags.`, } - cmd.Flags().StringVar(&scope, "scope", "", "Databricks secret scope holding registry credentials") - cmd.Flags().StringVar(&key, "key", "", "Databricks secret key holding registry credentials") - cmd.Flags().BoolVarP(&interactiveAuth, "interactive-authenticate", "i", false, "Prompt for registry credentials and store them as a secret") - cmd.Flags().StringVar(&tagPolicy, "tag-policy", "auto", "Image resolution policy: auto or latest") + // Registration always re-checks the source registry for the latest digest. + // --tag-policy is kept only for backward compatibility (accepts "latest"). + cmd.Flags().StringVar(&tagPolicy, "tag-policy", "", "Deprecated and ignored; registration always checks the source registry for the latest digest") + _ = cmd.Flags().MarkHidden("tag-policy") cmd.Flags().IntVar(&timeoutMinutes, "timeout-minutes", 60, "Timeout to wait for the image to become available") + // Resolve and authenticate the workspace client up front so an auth failure + // fails fast here, before any image is registered or polled. + 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() + + dockerImageURL := strings.TrimSpace(args[0]) + if dockerImageURL == "" { + return renderError(ctx, cmd, "INVALID_ARGS", "PERMANENT", false, + errors.New("IMAGE_URL cannot be empty")) + } + + if err := validateTagPolicy(tagPolicy); err != nil { + return renderError(ctx, cmd, "INVALID_ARGS", "PERMANENT", false, err) + } + + if timeoutMinutes <= 0 { + return renderError(ctx, cmd, "INVALID_ARGS", "PERMANENT", false, + fmt.Errorf("--timeout-minutes must be positive, got %d", timeoutMinutes)) + } + + w := cmdctx.WorkspaceClient(ctx) + + // Validate authentication against the workspace before registering + // anything: MustWorkspaceClient only attaches credentials, so without this + // a bad credential surfaces as a confusing mid-flow failure. + if _, err := w.CurrentUser.Me(ctx, iam.MeRequest{}); err != nil { + return authError(ctx, cmd, err) + } + + c, err := newImageClient(w) + if err != nil { + return renderError(ctx, cmd, "INTERNAL_ERROR", "TRANSIENT", true, err) + } + + timeout := time.Duration(timeoutMinutes) * time.Minute + + // Discover credentials from the local Docker config and store them in a + // per-user secret for the registration call. If storage fails, registration + // still proceeds without credentials — a public image succeeds — and + // credErr is reported as the cause if the registry rejects anonymous access. + // credErr is not fatal on its own: it is reported by registrationError only + // if the registry then rejects anonymous access. + scope, key, credErr := discoverCredentials(ctx, w, c, dockerImageURL) + if credErr != nil { + log.Debugf(ctx, "could not store local Docker credentials: %v", credErr) + } + + updated, sha, err := registerWithCredentialFallback(ctx, c, dockerImageURL, scope, key, timeout) + if err != nil { + kind, retryable := classifyRegistrationError(err) + return renderError(ctx, cmd, "REGISTRATION_FAILED", kind, retryable, + registrationError(dockerImageURL, err, credErr)) + } + + return renderRegisterResult(ctx, cmd, dockerImageURL, + registerImageResult{ + DockerImageURL: dockerImageURL, + ManifestSHA256: sha, + Status: string(imageStatusAvailable), + ImageUpdated: updated, + Cached: !updated, + }) + } + return cmd } + +// discoverCredentials resolves registry credentials from the local Docker config +// and stores them in a per-user secret, returning the (scope, key) reference for +// registration. It first probes whether the image is public: if so, no +// credentials are stored (avoiding a throwaway secret). Returns empty scope/key +// when the image is public or no local credentials exist, both with a nil error. +// A non-nil error means credentials were found but could not be stored (e.g. the +// user lacks permission to create a secret scope); it is advisory, so the caller +// can still attempt an anonymous registration and report this as the cause if +// that fails. +func discoverCredentials(ctx context.Context, w *databricks.WorkspaceClient, c *imageClient, dockerImageURL string) (scope, key string, err error) { + // readDockerCredentials keys off the registry host, so it needs the normalized + // URL (e.g. bare "ubuntu" resolves to the Docker Hub host). + normalized := normalizeDockerImageURL(dockerImageURL) + + // Resolve local creds once (a cheap file read, or one credential-helper call), + // so the no-`docker login` case pays for nothing and helpers aren't invoked + // twice. + username, password, ok := readDockerCredentials(ctx, normalized) + if !ok { + return "", "", nil + } + + if public := c.checkImageAccess(ctx, dockerImageURL); public != nil && *public { + log.Infof(ctx, "image is publicly accessible; skipping local Docker credentials") + return "", "", nil + } + + scope, key, err = storeDockerCredentials(ctx, w, normalized, username, password) + if err != nil { + return "", "", err + } + log.Infof(ctx, "using Docker credentials from local config (stored as %s/%s)", scope, key) + return scope, key, nil +} + +// isAuthError reports whether err is an authentication or permission failure, +// used to decide whether stale auto-discovered credentials warrant an anonymous +// retry. Unlike the Python CLI's substring match on "401"/"403", this keys off +// the SDK's typed sentinels. +func isAuthError(err error) bool { + return errors.Is(err, apierr.ErrUnauthenticated) || errors.Is(err, apierr.ErrPermissionDenied) +} + +// classifyRegistrationError maps a registration failure to the error envelope's +// kind and retryable flag. Only errors we can positively identify as transient +// (rate limits, server-side blips, a poll timeout) are retryable; auth, +// not-found, bad-request, conflict, a terminal FAILED upload, and any +// unclassified error default to permanent, so a consumer never retries a request +// that can't succeed. +func classifyRegistrationError(err error) (kind string, retryable bool) { + switch { + case errors.Is(err, errImageWaitTimeout), + errors.Is(err, apierr.ErrTooManyRequests), + errors.Is(err, apierr.ErrTemporarilyUnavailable), + errors.Is(err, apierr.ErrInternalError), + errors.Is(err, apierr.ErrDeadlineExceeded): + return "TRANSIENT", true + default: + return "PERMANENT", false + } +} + +// registrationError wraps an auth failure with actionable guidance. When +// credentials were found locally but could not be stored, credErr is the real +// cause — telling the user to `docker login` would be wrong, since they already +// have working credentials. +func registrationError(dockerImageURL string, err, credErr error) error { + if !isAuthError(err) { + return err + } + if credErr != nil { + return fmt.Errorf("image %q requires credentials, and the credentials found in your local Docker config could not be stored: %w", dockerImageURL, credErr) + } + return fmt.Errorf("image %q was not found or requires credentials: run `docker login` for its registry, then retry: %w", dockerImageURL, err) +} + +// registerWithCredentialFallback registers the image and, if the stored +// credentials are rejected as an auth failure, retries once anonymously so a +// public image isn't blocked by stale local creds (e.g. a revoked PAT from an +// old `docker login`). The retry only fires when credentials were supplied. +func registerWithCredentialFallback(ctx context.Context, c *imageClient, dockerImageURL, scope, key string, timeout time.Duration) (updated bool, sha string, err error) { + updated, sha, err = resolveImage(ctx, c, dockerImageURL, scope, key, timeout) + if err != nil && scope != "" && isAuthError(err) { + log.Warnf(ctx, "stored Docker credentials were rejected (%v); retrying without credentials in case the image is public", err) + return resolveImage(ctx, c, dockerImageURL, "", "", timeout) + } + return updated, sha, err +} + +// resolveImage always re-registers the image and waits for it to become +// AVAILABLE, returning whether the stored digest changed and the final digest. +// CreateImage is idempotent. The prior registration is fetched only to detect a +// digest change; its status is not consulted. +func resolveImage(ctx context.Context, c *imageClient, dockerImageURL, scope, key string, timeout time.Duration) (updated bool, sha string, err error) { + existing, err := c.getImage(ctx, dockerImageURL) + if err != nil { + return false, "", err + } + + reg, err := createAndWait(ctx, c, dockerImageURL, scope, key, timeout) + if err != nil { + return false, "", err + } + + newSHA := reg.ManifestSHA256 + // Re-read to pick up a digest populated by the background upload. + if final, err := c.getImage(ctx, dockerImageURL); err == nil && final != nil && final.ManifestSHA256 != "" { + newSHA = final.ManifestSHA256 + } + + // A first-time registration, or a changed digest, counts as updated. + if existing == nil { + return true, newSHA, nil + } + cachedSHA := existing.ManifestSHA256 + if newSHA == "" { + newSHA = cachedSHA + } + return cachedSHA != "" && newSHA != "" && cachedSHA != newSHA, newSHA, nil +} + +// createAndWait registers the image and polls until it becomes AVAILABLE. +func createAndWait(ctx context.Context, c *imageClient, dockerImageURL, scope, key string, timeout time.Duration) (*imageRegistration, error) { + reg, err := c.createImage(ctx, dockerImageURL, scope, key) + if err != nil { + return nil, err + } + if reg.Status == imageStatusAvailable { + return reg, nil + } + return c.waitForImageReady(ctx, dockerImageURL, timeout, imagePollInterval) +} + +// renderRegisterResult prints the result as a JSON envelope or human-readable +// text, matching the Python CLI's output. +func renderRegisterResult(ctx context.Context, cmd *cobra.Command, dockerImageURL string, result registerImageResult) error { + if root.OutputType(cmd) != flags.OutputText { + return renderEnvelope(ctx, result) + } + + out := cmd.OutOrStdout() + sha := "unknown" + if result.ManifestSHA256 != "" { + sha = shortManifestSHA(result.ManifestSHA256) + } + + if result.ImageUpdated { + fmt.Fprintf(out, "Image registered: %s\n", sha) + } else { + fmt.Fprintf(out, "Image already up to date: %s\n", sha) + } + + fmt.Fprintln(out, "\nTo use this image in your training config:") + fmt.Fprintln(out, " environment:") + fmt.Fprintln(out, " docker_image:") + fmt.Fprintf(out, " url: %s\n", dockerImageURL) + return nil +} + +// shortManifestSHA truncates a manifest digest to its first 16 characters for +// display, matching the Python CLI. +func shortManifestSHA(sha string) string { + if len(sha) <= 16 { + return sha + } + return sha[:16] + "..." +} diff --git a/experimental/air/cmd/register_image_test.go b/experimental/air/cmd/register_image_test.go new file mode 100644 index 00000000000..8ad5383940b --- /dev/null +++ b/experimental/air/cmd/register_image_test.go @@ -0,0 +1,185 @@ +package aircmd + +import ( + "context" + "errors" + "fmt" + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/databricks/databricks-sdk-go/apierr" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestClassifyRegistrationError(t *testing.T) { + cases := []struct { + name string + err error + kind string + retryable bool + }{ + {"auth", apierr.ErrUnauthenticated, "PERMANENT", false}, + {"permission", apierr.ErrPermissionDenied, "PERMANENT", false}, + {"not found", apierr.ErrNotFound, "PERMANENT", false}, + {"bad request", apierr.ErrBadRequest, "PERMANENT", false}, + {"conflict", apierr.ErrResourceConflict, "PERMANENT", false}, + {"canceled", context.Canceled, "PERMANENT", false}, + {"upload failed", fmt.Errorf("%w: boom", errImageUploadFailed), "PERMANENT", false}, + {"unknown error", errors.New("something odd"), "PERMANENT", false}, + {"wait timeout", fmt.Errorf("%w within 1m0s", errImageWaitTimeout), "TRANSIENT", true}, + {"rate limited", apierr.ErrTooManyRequests, "TRANSIENT", true}, + {"server error", apierr.ErrInternalError, "TRANSIENT", true}, + {"unavailable", apierr.ErrTemporarilyUnavailable, "TRANSIENT", true}, + {"deadline exceeded", apierr.ErrDeadlineExceeded, "TRANSIENT", true}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + kind, retryable := classifyRegistrationError(tc.err) + assert.Equal(t, tc.kind, kind) + assert.Equal(t, tc.retryable, retryable) + }) + } +} + +func TestValidateTagPolicy(t *testing.T) { + require.NoError(t, validateTagPolicy("")) + require.NoError(t, validateTagPolicy("latest")) + require.NoError(t, validateTagPolicy(" LATEST ")) + + err := validateTagPolicy("auto") + require.Error(t, err) + assert.Contains(t, err.Error(), "no longer supported") + + err = validateTagPolicy("bogus") + require.Error(t, err) + assert.Contains(t, err.Error(), "only supported value is latest") +} + +func TestRegistrationError(t *testing.T) { + authErr := apierr.ErrPermissionDenied + credErr := errors.New(`creating secret scope "docker-credentials-you@example.com" was denied`) + + // Credentials were found but couldn't be stored: report that as the cause, + // not "run docker login" — the user already has working credentials. + err := registrationError("nvcr.io/org/img:1.0", authErr, credErr) + assert.Contains(t, err.Error(), "requires credentials, and the credentials found in your local Docker config could not be stored") + assert.Contains(t, err.Error(), "was denied") + assert.NotContains(t, err.Error(), "run `docker login`") + + // No credential-storage problem: the docker login hint is the right guidance. + err = registrationError("nvcr.io/org/img:1.0", authErr, nil) + assert.Contains(t, err.Error(), "run `docker login`") + + // A non-auth failure passes through untouched. + other := errors.New("boom") + assert.Equal(t, other, registrationError("nvcr.io/org/img:1.0", other, credErr)) +} + +// imageServer serves the image API. Each :get call returns the next body in +// getBodies (repeating the last), where an empty string means 404; POST returns +// postBody. Sequencing the :get bodies lets a test set distinct before/after +// digests for the two :get calls resolveImage makes. +func imageServer(t *testing.T, postBody string, getBodies ...string) string { + t.Helper() + call := 0 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case imagesAPIPath + ":get": + body := getBodies[min(call, len(getBodies)-1)] + call++ + if body == "" { + w.WriteHeader(http.StatusNotFound) + _, _ = w.Write([]byte(`{"error_code":"NOT_FOUND","message":"not registered"}`)) + return + } + _, _ = w.Write([]byte(body)) + case imagesAPIPath: + _, _ = w.Write([]byte(postBody)) + default: + _, _ = w.Write([]byte(`{}`)) + } + })) + t.Cleanup(srv.Close) + return srv.URL +} + +func TestResolveImageFreshRegistration(t *testing.T) { + url := imageServer(t, `{"image":{"state":"AVAILABLE","manifest_sha256":"newsha"}}`, "") + updated, sha, err := resolveImage(t.Context(), newTestImageClient(t, url), "ubuntu", "", "", time.Second) + require.NoError(t, err) + assert.True(t, updated) + assert.Equal(t, "newsha", sha) +} + +func TestResolveImageDigestUnchanged(t *testing.T) { + body := `{"state":"AVAILABLE","manifest_sha256":"samesha"}` + url := imageServer(t, `{"image":`+body+`}`, body) + updated, sha, err := resolveImage(t.Context(), newTestImageClient(t, url), "ubuntu", "", "", time.Second) + require.NoError(t, err) + assert.False(t, updated) + assert.Equal(t, "samesha", sha) +} + +func TestResolveImageDigestChanged(t *testing.T) { + // First :get is the pre-existing (old) digest; the re-read after POST returns + // the new digest, so the image reports updated. + url := imageServer(t, `{"image":{"state":"AVAILABLE","manifest_sha256":"newsha"}}`, + `{"state":"AVAILABLE","manifest_sha256":"oldsha"}`, + `{"state":"AVAILABLE","manifest_sha256":"newsha"}`) + updated, sha, err := resolveImage(t.Context(), newTestImageClient(t, url), "ubuntu", "", "", time.Second) + require.NoError(t, err) + assert.True(t, updated) + assert.Equal(t, "newsha", sha) +} + +// credRejectingImageServer 401s a POST that carries credentials and returns +// AVAILABLE for an anonymous POST, so a test can exercise the stale-credential +// anonymous retry. +func credRejectingImageServer(t *testing.T, credentialedPOSTs *int) string { + t.Helper() + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case imagesAPIPath + ":get": + w.WriteHeader(http.StatusNotFound) + _, _ = w.Write([]byte(`{"error_code":"NOT_FOUND","message":"not registered"}`)) + case imagesAPIPath: + body, _ := io.ReadAll(r.Body) + if strings.Contains(string(body), "credentials_scope") { + *credentialedPOSTs++ + w.WriteHeader(http.StatusForbidden) + _, _ = w.Write([]byte(`{"error_code":"PERMISSION_DENIED","message":"denied"}`)) + return + } + _, _ = w.Write([]byte(`{"image":{"state":"AVAILABLE","manifest_sha256":"pubsha"}}`)) + default: + _, _ = w.Write([]byte(`{}`)) + } + })) + t.Cleanup(srv.Close) + return srv.URL +} + +func TestRegisterWithCredentialFallbackRetriesAnonymously(t *testing.T) { + var credentialedPOSTs int + url := credRejectingImageServer(t, &credentialedPOSTs) + updated, sha, err := registerWithCredentialFallback(t.Context(), newTestImageClient(t, url), "nvcr.io/org/img:1.0", "scope", "key", time.Second) + require.NoError(t, err) + assert.True(t, updated) + assert.Equal(t, "pubsha", sha) + assert.Equal(t, 1, credentialedPOSTs, "should try once with creds, then retry anonymously") +} + +func TestRegisterWithCredentialFallbackNoRetryWithoutCreds(t *testing.T) { + // Without credentials there is nothing stale to fall back from, so an auth + // failure surfaces directly. + var credentialedPOSTs int + url := credRejectingImageServer(t, &credentialedPOSTs) + _, _, err := registerWithCredentialFallback(t.Context(), newTestImageClient(t, url), "nvcr.io/org/img:1.0", "", "", time.Second) + require.NoError(t, err) // anonymous POST succeeds on this server + assert.Equal(t, 0, credentialedPOSTs) +} diff --git a/experimental/air/cmd/render.go b/experimental/air/cmd/render.go index f3fb63db807..5d4cf9bb6d0 100644 --- a/experimental/air/cmd/render.go +++ b/experimental/air/cmd/render.go @@ -4,6 +4,7 @@ import ( "context" "fmt" "io" + "net/http" "strconv" "strings" @@ -11,6 +12,7 @@ import ( "github.com/databricks/cli/libs/cmdio" "github.com/databricks/cli/libs/log" "github.com/databricks/databricks-sdk-go" + "github.com/databricks/databricks-sdk-go/client" "github.com/databricks/databricks-sdk-go/service/jobs" "github.com/muesli/termenv" "go.yaml.in/yaml/v3" @@ -22,6 +24,10 @@ const ( metadataBoxTitle = "Metadata" ) +// jobsRunsGetPath is the Jobs GetRun endpoint, called with a raw request to read +// the run's environments[] block, which the typed SDK Run does not expose. +const jobsRunsGetPath = "/api/2.2/jobs/runs/get" + // minBoxInnerWidth keeps all boxes a uniform, comfortable width; boxHPad and // boxVPad are the horizontal and vertical padding inside each box. const ( @@ -91,6 +97,11 @@ func renderRunText(ctx context.Context, out io.Writer, w *databricks.WorkspaceCl renderer, colorOn := cmdio.NewRenderer(ctx, out) p := newPalette(renderer) + // The serverless environment version lives on the run's environments[].spec, + // which the typed SDK Run drops, so read it with a raw request. Empty (a run + // with no serverless environment) stays "N/A". + data.EnvironmentDisplay = orNA(aiRuntimeEnvironmentVersion(ctx, w, run.RunId)) + view := runView{ runID: data.RunID, dashboardURL: data.DashboardURL, @@ -143,6 +154,40 @@ func genAIComputeTask(run *jobs.Run) *jobs.GenAiComputeTask { return run.Tasks[0].GenAiComputeTask } +// aiRuntimeEnvironmentVersion returns the serverless environment version an +// ai_runtime run used (e.g. "4"), read from the Jobs GetRun response's +// environments[] entry keyed by aiRuntimeEnvironmentKey. The typed SDK Run has +// no environments field, so the value is fetched with a raw request. Best-effort: +// returns "" (logged) on any error or when the run declares no environment. +func aiRuntimeEnvironmentVersion(ctx context.Context, w *databricks.WorkspaceClient, runID int64) string { + apiClient, err := client.New(w.Config) + if err != nil { + log.Warnf(ctx, "air get: could not create client to read environment: %v", err) + return "" + } + var resp struct { + Environments []struct { + EnvironmentKey string `json:"environment_key"` + Spec struct { + EnvironmentVersion string `json:"environment_version"` + } `json:"spec"` + } `json:"environments"` + } + // For a GET the SDK serializes the request value into query parameters, so + // run_id is passed as the request, mirroring the other raw calls in this package. + query := map[string]any{"run_id": runID} + if err := apiClient.Do(ctx, http.MethodGet, jobsRunsGetPath, nil, nil, query, &resp); err != nil { + log.Warnf(ctx, "air get: could not read environment for run %d: %v", runID, err) + return "" + } + for _, e := range resp.Environments { + if e.EnvironmentKey == aiRuntimeEnvironmentKey { + return e.Spec.EnvironmentVersion + } + } + return "" +} + // resolveConfigYAML returns the config box body: from the downloaded config file // when we have its path, else from the legacy task. func resolveConfigYAML(ctx context.Context, w *databricks.WorkspaceClient, run *jobs.Run, data *getData) string { diff --git a/experimental/air/cmd/render_test.go b/experimental/air/cmd/render_test.go index 2fe125bcf04..66e5b09deb7 100644 --- a/experimental/air/cmd/render_test.go +++ b/experimental/air/cmd/render_test.go @@ -151,9 +151,9 @@ 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", + environment: "4", }) // Labels are padded to the longest ("Accelerators"), so values align. @@ -161,7 +161,7 @@ func TestRenderFields(t *testing.T) { assert.Contains(t, out, "Accelerators 1x A10") // Max retries and environment show alongside the other fields. assert.Contains(t, out, "Max Retries 3") - assert.Contains(t, out, "Environment ml-runtime-gpu:1.0") + assert.Contains(t, out, "Environment 4") // The status carries its dot prefix. assert.Contains(t, out, "● SUCCESS") // Off a terminal, links render as the bare label (URLs live in JSON output). diff --git a/experimental/air/cmd/run.go b/experimental/air/cmd/run.go index bd32810e9bc..422d4388215 100644 --- a/experimental/air/cmd/run.go +++ b/experimental/air/cmd/run.go @@ -1,14 +1,18 @@ package aircmd import ( - "errors" + "context" "fmt" + "io" "strconv" + "strings" + "unicode/utf8" "github.com/databricks/cli/cmd/root" "github.com/databricks/cli/libs/cmdctx" "github.com/databricks/cli/libs/cmdio" "github.com/databricks/cli/libs/flags" + "github.com/databricks/databricks-sdk-go" "github.com/spf13/cobra" ) @@ -35,9 +39,36 @@ func newRunCommand() *cobra.Command { Short: "Submit a training workload from a YAML config", Long: `Submit a training workload to Databricks serverless GPU compute. -The workload is described by a YAML config file (see --file).`, +The workload is described by a YAML config file (see --file). + +To look up a config field, pass its path to -h: + + databricks experimental air run -h config + databricks experimental air run -h config.compute + databricks experimental air run -h config.compute.accelerator_type + +The path must be a separate argument: cobra reserves -h as a boolean, so +-h=config.compute and -hconfig.compute are not accepted.`, } + // cobra passes -h's positional args to the help func before Args/required-flag + // validation, so a config path documents a field without needing -f. + cmd.SetHelpFunc(func(c *cobra.Command, args []string) { + fields := c.Flags().Args() + if len(fields) == 0 { + // Parent() is nil for a detached command (unit tests). + if parent := c.Parent(); parent != nil { + parent.HelpFunc()(c, args) + return + } + _ = c.Usage() + return + } + if err := writeConfigFieldHelp(c.OutOrStdout(), fields[0]); err != nil { + c.PrintErrln("Error:", err) + } + }) + cmd.Flags().StringVarP(&file, "file", "f", "", "Path to the workload YAML config") cmd.Flags().BoolVar(&watch, "watch", false, "Stream logs until the run completes") cmd.Flags().StringArrayVar(&overrides, "override", nil, "Override a YAML field, e.g. compute.num_accelerators=8 (repeatable)") @@ -57,16 +88,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 } @@ -79,20 +101,128 @@ The workload is described by a YAML config file (see --file).`, return renderEnvelope(ctx, runResult{Status: "DRY_RUN_OK", DryRun: true}) } + jsonOut := root.OutputType(cmd) == flags.OutputJSON + + // Announce the experiment before uploading; skipped in JSON mode to keep + // stdout a clean envelope stream. + if !jsonOut { + cmdio.LogString(ctx, "Submitting experiment: "+cfg.ExperimentName) + } + w := cmdctx.WorkspaceClient(ctx) - runID, dashboardURL, err := submitWorkload(ctx, w, cfg, file, idempotencyKey) + runID, dashboardURL, err := submitWorkload(ctx, w, cfg, file, idempotencyKey, !jsonOut) if err != nil { return err } runIDStr := strconv.FormatInt(runID, 10) - if root.OutputType(cmd) == flags.OutputText { - cmdio.LogString(ctx, "Submitted run "+runIDStr) - cmdio.LogString(ctx, "View at: "+dashboardURL) - return nil + + if !watch { + if !jsonOut { + out := cmd.OutOrStdout() + printSubmitResult(ctx, out, runIDStr, dashboardURL) + // Append the MLflow links only if they resolve; a bare submit is not + // blocked on them since the confirmation above is already printed. + if ids := resolveMLflowIDsForRun(ctx, w, runID); ids != nil { + printMLflowLinks(ctx, out, w.Config.Host, ids) + } + cmdio.LogString(ctx, "\nTip: use --watch to stream logs until the run completes.") + return nil + } + // PENDING is the submit status, distinct from the --watch JSONL + // SUBMITTED event type below. + return renderEnvelope(ctx, runResult{Status: "PENDING", 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, } - return renderEnvelope(ctx, runResult{Status: "SUBMITTED", RunID: runIDStr, DashboardURL: dashboardURL}) + + if !jsonOut { + out := cmd.OutOrStdout() + // The MLflow links stream in via the logs below, so don't poll here. + printSubmitResult(ctx, out, runIDStr, dashboardURL) + // Separate the submit summary from the streamed logs. + fmt.Fprintln(out) + fmt.Fprintln(out, "Monitoring run and streaming logs...") + printLogsDivider(ctx, out) + 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. + out := cmd.OutOrStdout() + printSubmittedEvent(out, runIDStr, dashboardURL) + req.onStatusChange = func(current, previous string) { + printStatusEvent(out, current, previous) + } + 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 } + +// printSubmitResult writes the green success line and Job Run link. These don't +// depend on the MLflow IDs, so they print before any MLflow poll. The link is +// styled (blue, underlined) and clickable, matching the `air get` view, and +// degrades to plain text on non-rich terminals. +func printSubmitResult(ctx context.Context, out io.Writer, runIDStr, dashboardURL string) { + renderer, colorOn := cmdio.NewRenderer(ctx, out) + p := newPalette(renderer) + + fmt.Fprintln(out, p.green.Render("Submitted workload with Job Run ID: "+runIDStr)) + fmt.Fprintln(out, "View job run at: "+link(colorOn, p.blue, dashboardURL, dashboardURL)) +} + +// printMLflowLinks appends the styled, clickable MLflow run and experiment links +// once their IDs are resolved. +func printMLflowLinks(ctx context.Context, out io.Writer, host string, ids *mlflowIdentifiers) { + renderer, colorOn := cmdio.NewRenderer(ctx, out) + p := newPalette(renderer) + + runURL := mlflowRunURL(host, ids) + expURL := mlflowExperimentURL(host, ids) + fmt.Fprintln(out, "View MLflow run at: "+link(colorOn, p.blue, runURL, runURL)) + fmt.Fprintln(out, "View MLflow experiment at: "+link(colorOn, p.blue, expURL, expURL)) +} + +// logsDividerWidth is the total display width of the --watch logs divider. +const logsDividerWidth = 60 + +// printLogsDivider prints a centered "Logs" rule marking where the streamed +// --watch logs begin, separating them from the submit summary. The dim color is +// dropped on non-rich terminals; the rule characters are always printed. +func printLogsDivider(ctx context.Context, out io.Writer) { + renderer, _ := cmdio.NewRenderer(ctx, out) + p := newPalette(renderer) + + const label = " Logs " + side := max((logsDividerWidth-utf8.RuneCountInString(label))/2, 0) + rule := strings.Repeat("─", side) + label + strings.Repeat("─", side) + fmt.Fprintln(out, p.n7.Render(rule)) +} + +// 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_test.go b/experimental/air/cmd/run_test.go new file mode 100644 index 00000000000..c1efd3a8b34 --- /dev/null +++ b/experimental/air/cmd/run_test.go @@ -0,0 +1,136 @@ +package aircmd + +import ( + "bytes" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "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" +) + +// fastMLflowPoll shrinks the MLflow-ID poll interval so tests that exercise the +// poll's retry/timeout path don't sleep a real half-second per attempt. +func fastMLflowPoll(t *testing.T) { + t.Helper() + orig := mlflowLinkPollInterval + mlflowLinkPollInterval = time.Millisecond + t.Cleanup(func() { mlflowLinkPollInterval = orig }) +} + +// submitServer serves a non-watch `air run` submit: runs/submit returns a run +// id, runs/get returns the run with a task, and runs/get-output returns +// getOutput (pass `{}` for "no MLflow IDs yet"). Everything else — the auth +// probe and the workspace-files upload — gets a permissive stub. +func submitServer(t *testing.T, getOutput string) *httptest.Server { + t.Helper() + runGet := `{"run_id": 555, "tasks": [{"run_id": 556, "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": 555}`)) + case r.URL.Path == "/api/2.2/jobs/runs/get": + _, _ = w.Write([]byte(runGet)) + case r.URL.Path == "/api/2.2/jobs/runs/get-output": + _, _ = w.Write([]byte(getOutput)) + default: + _, _ = w.Write([]byte(`{"userName": "u@example.com", "workspace_id": 1}`)) + } + })) + t.Cleanup(srv.Close) + return srv +} + +func runSubmitCmd(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)) + + 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 TestRunSubmitTextOutput(t *testing.T) { + fastMLflowPoll(t) + var buf bytes.Buffer + // get-output carries no MLflow IDs, so the poll times out and the two MLflow + // links are omitted — only the Job Run link is printed. + err := runSubmitCmd(t, flags.OutputText, &buf, submitServer(t, `{}`).URL) + require.NoError(t, err) + + out := buf.String() + assert.Contains(t, out, "Submitting experiment: my-run") + assert.Contains(t, out, "Submitted workload with Job Run ID: 555") + assert.Contains(t, out, "View job run at: ") + assert.Contains(t, out, "/jobs/runs/555") + assert.Contains(t, out, "Tip: use --watch") + assert.NotContains(t, out, "View MLflow run at:") +} + +func TestRunSubmitTextOutputWithMLflowLinks(t *testing.T) { + var buf bytes.Buffer + srvURL := submitServer(t, `{"ai_runtime_task_output": {"mlflow_experiment_id": "exp1", "mlflow_run_id": "run1"}}`).URL + require.NoError(t, runSubmitCmd(t, flags.OutputText, &buf, srvURL)) + + out := buf.String() + assert.Contains(t, out, "View MLflow run at: ") + assert.Contains(t, out, "/ml/experiments/exp1/runs/run1") + assert.Contains(t, out, "View MLflow experiment at: ") + assert.Contains(t, out, "/ml/experiments/exp1") +} + +func TestRunSubmitMLflowLinksResolveOnRetry(t *testing.T) { + fastMLflowPoll(t) + // get-output is empty on the first poll and carries the IDs on the second, + // exercising the poll's sleep-and-retry path. + var getOutputCalls int + runGet := `{"run_id": 555, "tasks": [{"run_id": 556, "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": 555}`)) + case r.URL.Path == "/api/2.2/jobs/runs/get": + _, _ = w.Write([]byte(runGet)) + case r.URL.Path == "/api/2.2/jobs/runs/get-output": + getOutputCalls++ + if getOutputCalls == 1 { + _, _ = w.Write([]byte(`{}`)) + return + } + _, _ = w.Write([]byte(`{"ai_runtime_task_output": {"mlflow_experiment_id": "exp1", "mlflow_run_id": "run1"}}`)) + default: + _, _ = w.Write([]byte(`{"userName": "u@example.com", "workspace_id": 1}`)) + } + })) + t.Cleanup(srv.Close) + + var buf bytes.Buffer + require.NoError(t, runSubmitCmd(t, flags.OutputText, &buf, srv.URL)) + + out := buf.String() + assert.Contains(t, out, "View MLflow run at: ") + assert.Contains(t, out, "/ml/experiments/exp1/runs/run1") + assert.GreaterOrEqual(t, getOutputCalls, 2, "should have polled get-output at least twice") +} + +func TestRunSubmitJSONStatusPending(t *testing.T) { + var buf bytes.Buffer + err := runSubmitCmd(t, flags.OutputJSON, &buf, submitServer(t, `{}`).URL) + require.NoError(t, err) + + out := buf.String() + assert.Contains(t, out, `"status": "PENDING"`) + assert.Contains(t, out, `"run_id": "555"`) + // JSON stdout stays a clean envelope stream — no human-readable submit lines. + assert.NotContains(t, out, "Submitting experiment") +} diff --git a/experimental/air/cmd/run_watch_test.go b/experimental/air/cmd/run_watch_test.go new file mode 100644 index 00000000000..a0b640fc479 --- /dev/null +++ b/experimental/air/cmd/run_watch_test.go @@ -0,0 +1,210 @@ +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 workload with Job Run ID: 777") + assert.Contains(t, out, "View job run at: ") + assert.Contains(t, out, "Monitoring run and streaming logs...") + // A "Logs" divider separates the submit summary from the streamed logs. + assert.Contains(t, out, "Logs") + assert.Contains(t, out, "───") + // 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.go b/experimental/air/cmd/runconfig.go index 09437f50a5b..10295275d25 100644 --- a/experimental/air/cmd/runconfig.go +++ b/experimental/air/cmd/runconfig.go @@ -3,7 +3,9 @@ package aircmd import ( "errors" "fmt" + "io" "maps" + "reflect" "regexp" "slices" "strings" @@ -31,27 +33,32 @@ var taskKeyRe = regexp.MustCompile(`^[A-Za-z0-9_-]+$`) // exec args). Only safe ref characters are allowed. var gitRefRe = regexp.MustCompile(`^[\w./-]+$`) +// Canonical UUID (8-4-4-4-12 hex). Usage policy ids are server-generated UUIDs, +// so an obviously-wrong value (e.g. a policy name pasted into usage_policy_id) +// is rejected up front with a hint pointing at usage_policy_name. +var uuidRe = regexp.MustCompile(`^[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}$`) + // runConfig is the top-level run YAML schema: experiment_name + compute / // environment / code_source plus the command and run options. type runConfig struct { - ExperimentName string `yaml:"experiment_name"` - Compute *computeConfig `yaml:"compute"` - Environment *environmentConfig `yaml:"environment"` - Command *string `yaml:"command"` - EnvVariables map[string]string `yaml:"env_variables"` - Secrets map[string]string `yaml:"secrets"` - CodeSource *codeSourceConfig `yaml:"code_source"` + ExperimentName string `yaml:"experiment_name" help:"Name of the experiment. Becomes the Jobs API task key: max 100 characters, alphanumerics, hyphens, and underscores only." required:"yes"` + Compute *computeConfig `yaml:"compute" help:"Which accelerators to run on and how many." required:"yes"` + Environment *environmentConfig `yaml:"environment" help:"Python dependencies, or a custom Docker image, for the run's runtime."` + Command *string `yaml:"command" help:"Shell command that starts the workload. Max 1000 lines; move longer logic into a script under code_source." required:"yes"` + EnvVariables map[string]string `yaml:"env_variables" help:"Plain environment variables, as NAME: value. A name here cannot also appear in secrets."` + Secrets map[string]string `yaml:"secrets" help:"Environment variables sourced from secrets, as NAME: scope/key."` + CodeSource *codeSourceConfig `yaml:"code_source" help:"Local code to upload and make available to the run."` // MaxRetries defaults to 3 when unset; default-filling is a normalization // concern handled at launch, so a nil pointer is left as-is here. - MaxRetries *int `yaml:"max_retries"` - TimeoutMinutes *int `yaml:"timeout_minutes"` - IdempotencyToken *string `yaml:"idempotency_token"` - Parameters map[string]any `yaml:"parameters"` - MLflowRunName *string `yaml:"mlflow_run_name"` - MLflowExperimentDirectory *string `yaml:"mlflow_experiment_directory"` - Permissions []permission `yaml:"permissions"` - UsagePolicyName *string `yaml:"usage_policy_name"` - UsagePolicyID *string `yaml:"usage_policy_id"` + MaxRetries *int `yaml:"max_retries" help:"How many times to retry a failed run. Must be >= 0. Defaults to 3 when unset."` + TimeoutMinutes *int `yaml:"timeout_minutes" help:"Wall-clock limit for the run in minutes. Must be >= 1."` + IdempotencyToken *string `yaml:"idempotency_token" help:"Reuse token: a repeat submission with the same token returns the existing run instead of starting another. Max 64 characters."` + Parameters map[string]any `yaml:"parameters" help:"Free-form values passed through to the workload. Any nested structure is allowed."` + MLflowRunName *string `yaml:"mlflow_run_name" help:"Name for the MLflow run. Max 100 characters, alphanumerics, hyphens, and underscores only."` + MLflowExperimentDirectory *string `yaml:"mlflow_experiment_directory" help:"Workspace directory holding the MLflow experiment. Must start with /Workspace."` + Permissions []permission `yaml:"permissions" help:"Who may view or manage the run, as a list of principal plus level grants."` + UsagePolicyName *string `yaml:"usage_policy_name" help:"Usage policy to bill the run to, by name. Max 127 characters. Mutually exclusive with usage_policy_id."` + UsagePolicyID *string `yaml:"usage_policy_id" help:"Usage policy to bill the run to, by id. Mutually exclusive with usage_policy_name."` } // validate runs structural validation over the whole config, returning the first @@ -164,8 +171,14 @@ func (c *runConfig) validate() error { return fmt.Errorf("usage_policy_name must be at most 127 characters, got %d", len(v)) } } - if c.UsagePolicyID != nil && strings.TrimSpace(*c.UsagePolicyID) == "" { - return errors.New("usage_policy_id must not be empty") + if c.UsagePolicyID != nil { + v := strings.TrimSpace(*c.UsagePolicyID) + if v == "" { + return errors.New("usage_policy_id must not be empty") + } + if !uuidRe.MatchString(v) { + return fmt.Errorf("usage_policy_id must be a UUID (for example, '12345678-90ab-cdef-1234-567890abcdef'), got: %s. To assign a policy by name instead, use usage_policy_name", v) + } } return nil @@ -216,9 +229,9 @@ func validateSecretRefs(secrets map[string]string) error { // environmentConfig is the `environment` block: dependencies and/or a custom // docker image. type environmentConfig struct { - Dependencies dependencies `yaml:"dependencies"` - Version stringOrInt `yaml:"version"` - DockerImage *dockerImageConfig `yaml:"docker_image"` + Dependencies dependencies `yaml:"dependencies" help:"Inline list of packages to install. Not allowed alongside docker_image."` + Version stringOrInt `yaml:"version" help:"Client image version to pin. Only valid alongside inline dependencies."` + DockerImage *dockerImageConfig `yaml:"docker_image" help:"Custom image supplying the whole runtime. Not allowed alongside dependencies or version."` } func (e *environmentConfig) validate() error { @@ -238,40 +251,29 @@ func (e *environmentConfig) validate() error { return e.DockerImage.validate() } - // version pins the client image version, which is only meaningful for an - // inline (list) dependency set — a requirements.yaml file carries its own. - if e.Version.set { - if e.Dependencies.set && !e.Dependencies.isList { - return errors.New("'environment.version' is only valid with inline dependencies (a list); when 'dependencies' points to a requirements.yaml file, set the version inside that file") - } - if !e.Dependencies.set { - return errors.New("'environment.version' requires inline 'dependencies' (a list of packages)") - } + // version pins the client image version, which is only meaningful alongside an + // inline dependency set. + if e.Version.set && !e.Dependencies.set { + return errors.New("'environment.version' requires inline 'dependencies' (a list of packages)") } return nil } -// dependencies is environment.dependencies, which is polymorphic: a string is a -// path to a requirements.yaml file; a list is an inline package list. +// dependencies is environment.dependencies: an inline list of packages. A scalar +// (e.g. a path to a requirements file) is rejected — the list may itself reference +// a requirements.txt, but dependencies must be given as a list. type dependencies struct { - set bool - isList bool - path string - list []string + set bool + list []string } func (d *dependencies) UnmarshalYAML(node *yaml.Node) error { - switch node.Kind { - case yaml.ScalarNode: - d.set, d.isList = true, false - return node.Decode(&d.path) - case yaml.SequenceNode: - d.set, d.isList = true, true - return node.Decode(&d.list) - default: - return errors.New("environment.dependencies must be a string path or a list of packages") + if node.Kind != yaml.SequenceNode { + return errors.New("environment.dependencies must be a list of packages or reference a requirements.txt (see https://docs.databricks.com/aws/en/machine-learning/ai-runtime/cli/yaml-config#reference). A direct file reference is not supported") } + d.set = true + return node.Decode(&d.list) } // stringOrInt holds a scalar that may be a string or an integer in YAML @@ -293,7 +295,7 @@ func (s *stringOrInt) UnmarshalYAML(node *yaml.Node) error { // dockerImageConfig is environment.docker_image. type dockerImageConfig struct { - URL string `yaml:"url"` + URL string `yaml:"url" help:"Fully qualified image URL, e.g. myregistry.io/team/train:v3." required:"when environment.docker_image is set"` } func (d *dockerImageConfig) validate() error { @@ -305,8 +307,8 @@ func (d *dockerImageConfig) validate() error { // codeSourceConfig is the `code_source` block. Only the "snapshot" type exists. type codeSourceConfig struct { - Type string `yaml:"type"` - Snapshot *snapshotSourceConfig `yaml:"snapshot"` + Type string `yaml:"type" help:"Kind of code source. Must be \"snapshot\", the only supported type." required:"when code_source is set"` + Snapshot *snapshotSourceConfig `yaml:"snapshot" help:"Which local directory to archive and upload." required:"when code_source.type is \"snapshot\""` } func (c *codeSourceConfig) validate() error { @@ -321,10 +323,10 @@ func (c *codeSourceConfig) validate() error { // snapshotSourceConfig describes a local directory to tar and upload. type snapshotSourceConfig struct { - RootPath string `yaml:"root_path"` - RemoteVolume *string `yaml:"remote_volume"` - Git *gitRef `yaml:"git"` - IncludePaths []string `yaml:"include_paths"` + RootPath string `yaml:"root_path" help:"Local directory to archive, relative or absolute." required:"when code_source.snapshot is set"` + RemoteVolume *string `yaml:"remote_volume" help:"Volume to upload the archive to. Must start with /Volumes/."` + Git *gitRef `yaml:"git" help:"Pin the snapshot to a specific git revision."` + IncludePaths []string `yaml:"include_paths" help:"Restrict the archive to these paths, relative to root_path and without \"..\". Omit to include everything."` } func (s *snapshotSourceConfig) validate() error { @@ -363,9 +365,9 @@ func (s *snapshotSourceConfig) validate() error { // gitRef pins a snapshot to a specific git ref. branch and commit are mutually // exclusive; remote is only meaningful with branch. type gitRef struct { - Branch *string `yaml:"branch"` - Commit *string `yaml:"commit"` - Remote gitRemote `yaml:"remote"` + Branch *string `yaml:"branch" help:"Branch to pin to, resolved to its local HEAD. Mutually exclusive with commit." required:"one of branch or commit"` + Commit *string `yaml:"commit" help:"Commit to pin to. Mutually exclusive with branch." required:"one of branch or commit"` + Remote gitRemote `yaml:"remote" help:"No longer supported: the snapshot archives your local copy. Only false is accepted; use commit to pin a revision."` } func (g *gitRef) validate() error { @@ -423,12 +425,12 @@ func (r *gitRemote) truthy() bool { // permission is a DABs-compatible permission grant: exactly one principal plus a // level. type permission struct { - UserName *string `yaml:"user_name"` - GroupName *string `yaml:"group_name"` - ServicePrincipalName *string `yaml:"service_principal_name"` + UserName *string `yaml:"user_name" help:"Grant to this user, by email. Exactly one principal field per grant." required:"one principal per grant"` + GroupName *string `yaml:"group_name" help:"Grant to this group, by name. Exactly one principal field per grant." required:"one principal per grant"` + ServicePrincipalName *string `yaml:"service_principal_name" help:"Grant to this service principal, by name. Exactly one principal field per grant." required:"one principal per grant"` // Level is a databricks PermissionLevel (e.g. CAN_VIEW, CAN_MANAGE). Enum // membership is validated server-side; here we only require it to be set. - Level string `yaml:"level"` + Level string `yaml:"level" help:"Permission level to grant, e.g. CAN_VIEW or CAN_MANAGE. Validated server-side." required:"when a grant is listed"` } func (p *permission) validate() error { @@ -460,3 +462,301 @@ func (p *permission) validate() error { } return nil } + +// Below: `air run -h config.`, which documents the schema above from its +// yaml/help/required struct tags. + +// configHelpRoot is the optional leading segment of a help path (`config.compute` +// or bare `compute`). +const configHelpRoot = "config" + +// writeConfigFieldHelp resolves a dotted config path and writes its docs. +func writeConfigFieldHelp(w io.Writer, path string) error { + field, err := resolveConfigField(path) + if err != nil { + return err + } + renderConfigField(w, field) + return nil +} + +// freeFormConfigFields hold free-form maps, so path resolution stops at them: +// their keys are chosen by the user, not the schema. +var freeFormConfigFields = map[string]bool{ + "parameters": true, + "env_variables": true, + "secrets": true, +} + +// configTypeNames labels the polymorphic unions, whose YAML shape reflection +// can't see (unexported fields filled by a custom UnmarshalYAML). +var configTypeNames = map[reflect.Type]string{ + reflect.TypeFor[dependencies](): "list of strings", + reflect.TypeFor[stringOrInt](): "string or int", + reflect.TypeFor[gitRemote](): "bool or string", +} + +// configField is one resolved node of the run config schema. +type configField struct { + path string + typeName string + required string + help string + // freeForm marks a user-keyed map (parameters/secrets/env_variables): no + // children, yet any sub-path into it is valid. + freeForm bool + children []configField // nil for a leaf +} + +// configSchema is the single reflection walk over runConfig; both +// `-h config.` and --override path validation resolve against it. +func configSchema() configField { + return configField{ + path: configHelpRoot, + help: "The run YAML schema. Pass a field path for details, e.g. " + configHelpRoot + ".compute.accelerator_type.", + children: describeStruct(reflect.TypeFor[runConfig](), configHelpRoot), + } +} + +// resolveConfigField resolves a dotted YAML path against the run config schema. +// The leading "config." is optional; an empty path describes the whole schema. +func resolveConfigField(path string) (configField, error) { + trimmed := strings.TrimPrefix(strings.TrimPrefix(path, configHelpRoot), ".") + root := configSchema() + if trimmed == "" { + return root, nil + } + + current := root + for i, part := range strings.Split(trimmed, ".") { + if current.freeForm { + // Keys are user-defined, so the map is the most specific node. + return configField{}, fmt.Errorf("%q holds user-defined keys, so %q is not part of the schema; see %q instead", current.path, part, current.path) + } + if len(current.children) == 0 { + return configField{}, fmt.Errorf("%q is not an object, so it has no field %q", current.path, part) + } + child, ok := findConfigChild(current.children, part) + if !ok { + return configField{}, unknownConfigFieldError(current, part, strings.Split(trimmed, ".")[:i+1]) + } + current = child + } + return current, nil +} + +// findConfigChild looks up an immediate child by its YAML name. +func findConfigChild(children []configField, name string) (configField, bool) { + for _, c := range children { + if configLeafName(c.path) == name { + return c, true + } + } + return configField{}, false +} + +// configLeafName returns the last segment of a dotted path. +func configLeafName(path string) string { + _, leaf, found := cutLast(path, ".") + if !found { + return path + } + return leaf +} + +// cutLast splits s around the final instance of sep. +func cutLast(s, sep string) (before, after string, found bool) { + i := strings.LastIndex(s, sep) + if i < 0 { + return s, "", false + } + return s[:i], s[i+len(sep):], true +} + +// unknownConfigFieldError reports an unresolvable segment, naming the valid +// siblings and, when one is close enough, a suggestion. +func unknownConfigFieldError(parent configField, part string, matched []string) error { + names := make([]string, 0, len(parent.children)) + for _, c := range parent.children { + names = append(names, configLeafName(c.path)) + } + slices.Sort(names) + + msg := fmt.Sprintf("unknown config field %q", configHelpRoot+"."+strings.Join(matched, ".")) + if suggestion, ok := closestConfigField(part, names); ok { + msg += fmt.Sprintf("; did you mean %q?", suggestion) + } + return fmt.Errorf("%s\n\nfields under %q are: %s", msg, parent.path, strings.Join(names, ", ")) +} + +// closestConfigField returns the nearest candidate by edit distance, if one is +// close enough to be worth suggesting. +func closestConfigField(name string, candidates []string) (string, bool) { + best, bestDist := "", 0 + for _, c := range candidates { + d := configEditDistance(name, c) + // Allow roughly a third of the name to differ, and always accept a + // single edit so short names still get a suggestion. + limit := max(len(c)/3, 1) + if d <= limit && (best == "" || d < bestDist) { + best, bestDist = c, d + } + } + return best, best != "" +} + +// configEditDistance computes the Levenshtein distance between two strings. +func configEditDistance(a, b string) int { + prev := make([]int, len(b)+1) + curr := make([]int, len(b)+1) + for j := range prev { + prev[j] = j + } + for i := 1; i <= len(a); i++ { + curr[0] = i + for j := 1; j <= len(b); j++ { + cost := 1 + if a[i-1] == b[j-1] { + cost = 0 + } + curr[j] = min(min(curr[j-1]+1, prev[j]+1), prev[j-1]+cost) + } + prev, curr = curr, prev + } + return prev[len(b)] +} + +// describeStruct reads a struct's yaml/help/required tags into configFields, +// recursing into nested objects. Declaration order matches validate()'s errors. +func describeStruct(t reflect.Type, prefix string) []configField { + var out []configField + for f := range t.Fields() { + tag := f.Tag.Get("yaml") + if tag == "" || tag == "-" { + continue + } + name, _, _ := strings.Cut(tag, ",") + if name == "" || name == "-" { + continue + } + + field := configField{ + path: prefix + "." + name, + typeName: configTypeName(f.Type), + required: f.Tag.Get("required"), + help: f.Tag.Get("help"), + freeForm: freeFormConfigFields[name], + } + if nested := underlyingConfigStruct(f.Type); nested != nil && !field.freeForm { + field.children = describeStruct(nested, field.path) + } + out = append(out, field) + } + return out +} + +// configTypeName renders a field's YAML-facing type. +func configTypeName(t reflect.Type) string { + if name, ok := configTypeNames[t]; ok { + return name + } + for t.Kind() == reflect.Pointer { + t = t.Elem() + } + switch t.Kind() { + case reflect.String: + return "string" + case reflect.Int: + return "int" + case reflect.Bool: + return "bool" + case reflect.Slice: + return "list of " + configTypeName(t.Elem()) + case reflect.Map: + return fmt.Sprintf("map of %s to %s", configTypeName(t.Key()), configTypeName(t.Elem())) + case reflect.Struct: + return "object" + case reflect.Interface: + return "any" + default: + return t.Kind().String() + } +} + +// underlyingConfigStruct unwraps pointer/slice indirection and returns the struct +// type a field decodes into, or nil if it is not a struct. The polymorphic unions +// are excluded: they are structs, but their YAML shape is scalar or list. +func underlyingConfigStruct(t reflect.Type) reflect.Type { + if _, ok := configTypeNames[t]; ok { + return nil + } + for t.Kind() == reflect.Pointer || t.Kind() == reflect.Slice { + t = t.Elem() + } + if t.Kind() == reflect.Struct { + return t + } + return nil +} + +// renderConfigField writes a resolved field's documentation. An object lists its +// immediate children; a leaf gets its type, required-ness, and description. +func renderConfigField(w io.Writer, f configField) { + fmt.Fprintf(w, "%s\n", f.path) + if f.help != "" { + fmt.Fprintf(w, " %s\n", f.help) + } + + if len(f.children) == 0 { + fmt.Fprintf(w, "\n Type: %s\n", f.typeName) + required := f.required + if required == "" { + required = "no" + } + fmt.Fprintf(w, " Required: %s\n", required) + return + } + + width := 0 + for _, c := range f.children { + width = max(width, len(configLeafName(c.path))) + } + fmt.Fprintf(w, "\n Fields:\n") + for _, c := range f.children { + fmt.Fprintf(w, " %-*s %s\n", width, configLeafName(c.path), configFieldSummary(c)) + } + fmt.Fprintf(w, "\nUse \"-h %s.\" for details on a field.\n", f.path) +} + +// configFieldSummary is the one-line description used in a field listing: the +// first sentence of the help text, annotated when the field is required. +func configFieldSummary(f configField) string { + summary := firstSentence(f.help) + if f.required == "yes" { + summary = "(required) " + summary + } + return summary +} + +// sentenceAbbreviations end in a period that does not close a sentence, so +// firstSentence must not break on them. +var sentenceAbbreviations = []string{"e.g.", "i.e.", "etc.", "vs.", "cf."} + +// firstSentence returns s up to and including the first sentence-ending period, +// i.e. the first ". " boundary not immediately preceded by a known abbreviation. +// Returns s unchanged when it holds a single sentence. +func firstSentence(s string) string { + for i := 0; i+1 < len(s); i++ { + if s[i] != '.' || s[i+1] != ' ' { + continue + } + candidate := s[:i+1] + if slices.ContainsFunc(sentenceAbbreviations, func(a string) bool { + return strings.HasSuffix(candidate, a) + }) { + continue + } + return candidate + } + return s +} diff --git a/experimental/air/cmd/runconfig_launch.go b/experimental/air/cmd/runconfig_launch.go index 1408b600736..75de05ec1be 100644 --- a/experimental/air/cmd/runconfig_launch.go +++ b/experimental/air/cmd/runconfig_launch.go @@ -36,27 +36,17 @@ func (c *runConfig) dockerImageURL() string { return "" } -// requirementsFile returns the path to a requirements file when -// environment.dependencies is a string, and whether it was set. -func (c *runConfig) requirementsFile() (string, bool) { - if c.Environment == nil || !c.Environment.Dependencies.set || c.Environment.Dependencies.isList { - return "", false - } - return c.Environment.Dependencies.path, true -} - -// inlineDependencies returns the inline package list when -// environment.dependencies is a list, and whether it was set. +// inlineDependencies returns the inline package list from +// environment.dependencies, and whether it was set. func (c *runConfig) inlineDependencies() ([]string, bool) { - if c.Environment == nil || !c.Environment.Dependencies.set || !c.Environment.Dependencies.isList { + if c.Environment == nil || !c.Environment.Dependencies.set { return nil, false } return c.Environment.Dependencies.list, true } // runtimeVersion returns the client image version from environment.version when -// set. For a requirements-file dependency set, the version lives in that file and -// is resolved at launch, not here. +// set. func (c *runConfig) runtimeVersion() (string, bool) { if c.Environment == nil || !c.Environment.Version.set { return "", false diff --git a/experimental/air/cmd/runconfig_launch_test.go b/experimental/air/cmd/runconfig_launch_test.go index 289db91c7de..4b250a83fc2 100644 --- a/experimental/air/cmd/runconfig_launch_test.go +++ b/experimental/air/cmd/runconfig_launch_test.go @@ -39,32 +39,17 @@ func TestRunConfigDockerImageURL(t *testing.T) { func TestRunConfigDependencies(t *testing.T) { t.Run("unset", func(t *testing.T) { c := &runConfig{} - _, ok := c.requirementsFile() - assert.False(t, ok) - _, ok = c.inlineDependencies() - assert.False(t, ok) - }) - - t.Run("file path", func(t *testing.T) { - c := &runConfig{Environment: &environmentConfig{ - Dependencies: dependencies{set: true, isList: false, path: "req.yaml"}, - }} - path, ok := c.requirementsFile() - assert.True(t, ok) - assert.Equal(t, "req.yaml", path) - _, ok = c.inlineDependencies() + _, ok := c.inlineDependencies() assert.False(t, ok) }) t.Run("inline list", func(t *testing.T) { c := &runConfig{Environment: &environmentConfig{ - Dependencies: dependencies{set: true, isList: true, list: []string{"torch", "numpy"}}, + Dependencies: dependencies{set: true, list: []string{"torch", "numpy"}}, }} list, ok := c.inlineDependencies() assert.True(t, ok) assert.Equal(t, []string{"torch", "numpy"}, list) - _, ok = c.requirementsFile() - assert.False(t, ok) }) } 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..e8f8930b5b0 --- /dev/null +++ b/experimental/air/cmd/runconfig_override.go @@ -0,0 +1,128 @@ +package aircmd + +import ( + "context" + "fmt" + "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. + +// 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 { + schema := configSchema() + for _, e := range entries { + if err := checkOverridePath(strings.Split(e.path, "."), schema, e.path); err != nil { + return err + } + } + return nil +} + +// checkOverridePath validates one dotted path against a resolved schema node. +// It shares configSchema()'s reflection walk with `-h config.` but keeps +// the --override error voice, which names the offending flag. +func checkOverridePath(parts []string, node configField, fullPath string) error { + name := parts[0] + child, ok := findConfigChild(node.children, name) + if !ok { + names := make([]string, 0, len(node.children)) + for _, c := range node.children { + names = append(names, configLeafName(c.path)) + } + slices.Sort(names) + return fmt.Errorf("invalid --override %q: %q is not a known field; available fields are: %s", + fullPath, name, strings.Join(names, ", ")) + } + if len(parts) == 1 { + return nil + } + // A free-form map's keys are user-defined, so any sub-path into it is valid. + if child.freeForm { + return nil + } + if len(child.children) == 0 { + return fmt.Errorf("invalid --override %q: %q is not a nested object; cannot address sub-field %q", + fullPath, name, strings.Join(parts[1:], ".")) + } + return checkOverridePath(parts[1:], child, fullPath) +} + +// 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/runconfig_test.go b/experimental/air/cmd/runconfig_test.go index 26b54127265..f99bb344e79 100644 --- a/experimental/air/cmd/runconfig_test.go +++ b/experimental/air/cmd/runconfig_test.go @@ -82,7 +82,7 @@ permissions: require.NoError(t, err) assert.Equal(t, gpuType8xH100, gpuType(cfg.Compute.AcceleratorType)) require.NotNil(t, cfg.Environment) - assert.True(t, cfg.Environment.Dependencies.isList) + assert.True(t, cfg.Environment.Dependencies.set) assert.Equal(t, []string{"torch==2.3.0", "numpy"}, cfg.Environment.Dependencies.list) assert.True(t, cfg.Environment.Version.set) assert.Equal(t, "5", cfg.Environment.Version.raw) @@ -94,18 +94,17 @@ permissions: assert.Len(t, cfg.Permissions, 2) } -// TestLoadRunConfig_PolymorphicFields exercises the str|list, str|int, and -// bool|str unions decoded by custom UnmarshalYAML. +// TestLoadRunConfig_PolymorphicFields exercises the str|int and bool|str unions +// decoded by custom UnmarshalYAML, plus the rejection of the removed +// dependencies string form. func TestLoadRunConfig_PolymorphicFields(t *testing.T) { - t.Run("dependencies as string path", func(t *testing.T) { - cfg, err := loadRunConfig(writeConfig(t, minimalConfig+` + t.Run("dependencies as string path is rejected", func(t *testing.T) { + _, err := loadRunConfig(writeConfig(t, minimalConfig+` environment: dependencies: requirements.yaml `)) - require.NoError(t, err) - assert.True(t, cfg.Environment.Dependencies.set) - assert.False(t, cfg.Environment.Dependencies.isList) - assert.Equal(t, "requirements.yaml", cfg.Environment.Dependencies.path) + require.Error(t, err) + assert.Contains(t, err.Error(), "must be a list of packages") }) t.Run("git remote as bool true is rejected", func(t *testing.T) { @@ -243,7 +242,10 @@ func TestRunConfigValidate_FieldRules(t *testing.T) { c.UsagePolicyID = str("id") }, "mutually exclusive"}, {"empty usage_policy_id", func(c *runConfig) { c.UsagePolicyID = str(" ") }, "usage_policy_id must not be empty"}, - {"usage_policy_id alone is ok", func(c *runConfig) { c.UsagePolicyID = str("policy-uuid") }, ""}, + {"non-uuid usage_policy_id", func(c *runConfig) { c.UsagePolicyID = str("policy-uuid") }, "usage_policy_id must be a UUID"}, + // A name pasted into the id field gets pointed at the right field. + {"policy name in usage_policy_id", func(c *runConfig) { c.UsagePolicyID = str("team-a") }, "use usage_policy_name"}, + {"uuid usage_policy_id alone is ok", func(c *runConfig) { c.UsagePolicyID = str("12345678-90ab-cdef-1234-567890abcdef") }, ""}, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { @@ -275,7 +277,7 @@ func TestEnvironmentConfigValidate(t *testing.T) { "docker image with deps conflicts", environmentConfig{ DockerImage: &dockerImageConfig{URL: "org/repo:tag"}, - Dependencies: dependencies{set: true, isList: true, list: []string{"torch"}}, + Dependencies: dependencies{set: true, list: []string{"torch"}}, }, "not allowed: dependencies", }, @@ -284,14 +286,6 @@ func TestEnvironmentConfigValidate(t *testing.T) { environmentConfig{DockerImage: &dockerImageConfig{URL: " "}}, "docker_image.url cannot be empty", }, - { - "version with file deps", - environmentConfig{ - Version: stringOrInt{set: true, raw: "5"}, - Dependencies: dependencies{set: true, isList: false, path: "req.yaml"}, - }, - "only valid with inline dependencies", - }, { "version without deps", environmentConfig{Version: stringOrInt{set: true, raw: "5"}}, @@ -301,7 +295,7 @@ func TestEnvironmentConfigValidate(t *testing.T) { "version with inline deps ok", environmentConfig{ Version: stringOrInt{set: true, raw: "5"}, - Dependencies: dependencies{set: true, isList: true, list: []string{"torch"}}, + Dependencies: dependencies{set: true, list: []string{"torch"}}, }, "", }, @@ -411,3 +405,247 @@ func TestLoadRunConfig_FileErrors(t *testing.T) { assert.Contains(t, err.Error(), "is empty") }) } + +func TestResolveConfigField(t *testing.T) { + tests := []struct { + name string + path string + wantPath string + wantType string + wantRequired string + }{ + {"leaf", "config.compute.accelerator_type", "config.compute.accelerator_type", "string", ""}, + {"bare path", "compute.accelerator_type", "config.compute.accelerator_type", "string", ""}, + {"top-level required", "config.experiment_name", "config.experiment_name", "string", "yes"}, + {"int leaf", "config.max_retries", "config.max_retries", "int", ""}, + {"conditionally required", "config.environment.docker_image.url", "config.environment.docker_image.url", "string", "when environment.docker_image is set"}, + {"through a slice", "config.permissions.level", "config.permissions.level", "string", "when a grant is listed"}, + {"free-form map", "config.parameters", "config.parameters", "map of string to any", ""}, + {"deeply nested", "config.code_source.snapshot.root_path", "config.code_source.snapshot.root_path", "string", "when code_source.snapshot is set"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + field, err := resolveConfigField(tt.path) + require.NoError(t, err) + assert.Equal(t, tt.wantPath, field.path) + assert.Equal(t, tt.wantType, field.typeName) + assert.Equal(t, tt.wantRequired, field.required) + assert.NotEmpty(t, field.help) + assert.Empty(t, field.children) + }) + } +} + +// The union types are structs of unexported fields, so their type label comes +// from configTypeNames rather than reflection, and they must not be walked into. +func TestResolveConfigField_PolymorphicTypes(t *testing.T) { + tests := []struct{ path, wantType string }{ + {"config.environment.dependencies", "list of strings"}, + {"config.environment.version", "string or int"}, + {"config.code_source.snapshot.git.remote", "bool or string"}, + } + + for _, tt := range tests { + t.Run(tt.path, func(t *testing.T) { + field, err := resolveConfigField(tt.path) + require.NoError(t, err) + assert.Equal(t, tt.wantType, field.typeName) + // A union must not be walked into as an object: its fields are unexported. + assert.Empty(t, field.children) + }) + } +} + +func TestResolveConfigField_Containers(t *testing.T) { + root, err := resolveConfigField("config") + require.NoError(t, err) + assert.Equal(t, "config", root.path) + // Children are listed in declaration order, matching the schema. + assert.Equal(t, "experiment_name", configLeafName(root.children[0].path)) + assert.Equal(t, "compute", configLeafName(root.children[1].path)) + + // An empty path describes the whole schema, so `-h config` and a bare + // prefix agree. + bare, err := resolveConfigField("") + require.NoError(t, err) + assert.Equal(t, root.children, bare.children) + + compute, err := resolveConfigField("config.compute") + require.NoError(t, err) + assert.Equal(t, "object", compute.typeName) + require.Len(t, compute.children, 2) +} + +func TestResolveConfigField_Errors(t *testing.T) { + tests := []struct { + name string + path string + wantParts []string + }{ + { + name: "typo suggests the near match", + path: "config.compute.acclerator_type", + wantParts: []string{`unknown config field "config.compute.acclerator_type"`, `did you mean "accelerator_type"?`, "accelerator_type, num_accelerators"}, + }, + { + name: "unknown top-level field lists siblings", + path: "config.bogus", + wantParts: []string{`unknown config field "config.bogus"`, "experiment_name"}, + }, + { + name: "no suggestion when nothing is close", + path: "config.zzzzzzzzzzzz", + wantParts: []string{`unknown config field "config.zzzzzzzzzzzz"`}, + }, + { + name: "free-form keys are not schema fields", + path: "config.parameters.learning_rate", + wantParts: []string{`"config.parameters" holds user-defined keys`, `"learning_rate" is not part of the schema`}, + }, + { + name: "scalar has no sub-fields", + path: "config.command.foo", + wantParts: []string{`"config.command" is not an object`, `no field "foo"`}, + }, + { + name: "nested typo reports the resolved prefix", + path: "config.code_source.snapshot.rootpath", + wantParts: []string{`unknown config field "config.code_source.snapshot.rootpath"`, `did you mean "root_path"?`}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + _, err := resolveConfigField(tt.path) + require.Error(t, err) + for _, part := range tt.wantParts { + assert.Contains(t, err.Error(), part) + } + }) + } + + // A distant name must not attract a suggestion. + _, err := resolveConfigField("config.zzzzzzzzzzzz") + require.Error(t, err) + assert.NotContains(t, err.Error(), "did you mean") +} + +func TestWriteConfigFieldHelp(t *testing.T) { + var leaf strings.Builder + require.NoError(t, writeConfigFieldHelp(&leaf, "config.compute.num_accelerators")) + assert.Contains(t, leaf.String(), "config.compute.num_accelerators") + assert.Contains(t, leaf.String(), "Type: int") + assert.Contains(t, leaf.String(), "Required: no") + + var container strings.Builder + require.NoError(t, writeConfigFieldHelp(&container, "config.compute")) + assert.Contains(t, container.String(), "Fields:") + assert.Contains(t, container.String(), "num_accelerators") + // A container lists its fields rather than printing a type/required pair. + assert.NotContains(t, container.String(), "Required:") + assert.Contains(t, container.String(), `Use "-h config.compute." for details`) + + // Required top-level fields are flagged in a listing. + var root strings.Builder + require.NoError(t, writeConfigFieldHelp(&root, "config")) + assert.Contains(t, root.String(), "(required) Which accelerators to run on") + + require.Error(t, writeConfigFieldHelp(&strings.Builder{}, "config.nope")) +} + +// Guards against adding a schema field without a help: tag. +func TestConfigFieldsAllDocumented(t *testing.T) { + root, err := resolveConfigField("config") + require.NoError(t, err) + + var walk func(fields []configField) + walk = func(fields []configField) { + for _, f := range fields { + assert.NotEmpty(t, f.help, "%s is missing a help: struct tag", f.path) + assert.NotEmpty(t, f.typeName, "%s has no type name", f.path) + walk(f.children) + } + } + walk(root.children) +} + +func TestFirstSentence(t *testing.T) { + tests := []struct{ in, want string }{ + {"One sentence with no period", "One sentence with no period"}, + {"First. Second.", "First."}, + {"Which accelerator, e.g. GPU_1xA10. See the docs.", "Which accelerator, e.g. GPU_1xA10."}, + {"Use numpy, torch, etc. Then run.", "Use numpy, torch, etc. Then run."}, + {"Prod i.e. production. Details.", "Prod i.e. production."}, + {"Ends with abbreviation etc.", "Ends with abbreviation etc."}, + } + for _, tt := range tests { + assert.Equal(t, tt.want, firstSentence(tt.in)) + } +} + +// Plain -h must fall back to command help. A detached command has no parent to +// inherit a help function from, so this also covers that nil case. +func TestRunCommandHelp_NoConfigPath(t *testing.T) { + var out strings.Builder + cmd := newRunCommand() + cmd.SetOut(&out) + cmd.SetArgs([]string{"-h"}) + require.NoError(t, cmd.Execute()) + assert.Contains(t, out.String(), "Usage:") + assert.NotContains(t, out.String(), "Fields:") +} + +// -h with a config path documents the field instead of the command, and does so +// without the otherwise-required --file. +func TestRunCommandHelp_ConfigPath(t *testing.T) { + var out strings.Builder + cmd := newRunCommand() + cmd.SetOut(&out) + cmd.SetArgs([]string{"-h", "config.compute.accelerator_type"}) + require.NoError(t, cmd.Execute()) + assert.Contains(t, out.String(), "config.compute.accelerator_type") + assert.Contains(t, out.String(), "Type: string") + assert.NotContains(t, out.String(), "Usage:") +} + +// An unresolvable path reports to stderr and still exits cleanly, matching how +// cobra treats the help path. +func TestRunCommandHelp_UnknownConfigPath(t *testing.T) { + var out, errOut strings.Builder + cmd := newRunCommand() + cmd.SetOut(&out) + cmd.SetErr(&errOut) + cmd.SetArgs([]string{"-h", "config.nope"}) + require.NoError(t, cmd.Execute()) + assert.Contains(t, errOut.String(), `unknown config field "config.nope"`) + assert.Empty(t, out.String()) +} + +// Both `-h config.` and --override path validation must resolve against +// the one configSchema() walk, so a field valid for one is valid for the other. +// This guards against the two features drifting apart again. +func TestConfigSchemaSharedByHelpAndOverride(t *testing.T) { + paths := []string{ + "compute.num_accelerators", + "environment.docker_image.url", + "code_source.snapshot.root_path", + "env_variables.MY_VAR", // free-form sub-path + } + for _, p := range paths { + t.Run(p, func(t *testing.T) { + require.NoError(t, validateOverridePaths([]overrideEntry{{path: p, raw: "x"}})) + // The help path accepts the same field (free-form sub-paths resolve to + // the map itself, which is the schema's most specific node). + _, err := resolveConfigField(p) + if !strings.Contains(p, "env_variables") { + require.NoError(t, err) + } + }) + } + + // A field unknown to one is unknown to the other. + require.Error(t, validateOverridePaths([]overrideEntry{{path: "compute.bogus", raw: "x"}})) + _, err := resolveConfigField("compute.bogus") + require.Error(t, err) +} diff --git a/experimental/air/cmd/runsubmit.go b/experimental/air/cmd/runsubmit.go index 8c0be55260d..90e306f1b48 100644 --- a/experimental/air/cmd/runsubmit.go +++ b/experimental/air/cmd/runsubmit.go @@ -2,12 +2,12 @@ package aircmd import ( "context" - "errors" "fmt" "path" "strconv" "strings" + "github.com/databricks/cli/libs/cmdio" "github.com/databricks/cli/libs/env" "github.com/databricks/cli/libs/filer" "github.com/databricks/databricks-sdk-go" @@ -39,7 +39,9 @@ func dlRuntimeImage(ctx context.Context, runtimeVersion string) string { } // buildSubmitPayload assembles the runs/submit payload. commandPath is the -// workspace path of the uploaded command.sh; dlImage is the runtime channel. +// workspace path of the uploaded command.sh; dlImage is the runtime channel; +// usagePolicyID is the already-resolved policy id ("" when the run has none); +// deps is the user's declared dependencies (nil when none are declared). // // max_retries is always sent (including 0) so the user's YAML value is honored: // setting it to 0 explicitly disables retries rather than falling back to the @@ -47,7 +49,7 @@ func dlRuntimeImage(ctx context.Context, runtimeVersion string) string { // omitempty so the wire form matches the Python CLI (which never emits a bare // "false"). Jobs performs the retries — each attempt is a fresh AI Runtime // workload. -func buildSubmitPayload(cfg *runConfig, commandPath, dlImage string, snap snapshotResult) jobs.SubmitRun { +func buildSubmitPayload(cfg *runConfig, commandPath, dlImage, usagePolicyID string, snap snapshotResult, deps []string) jobs.SubmitRun { task := jobs.AiRuntimeTask{ Experiment: cfg.ExperimentName, Deployments: []jobs.DeploymentSpec{{ @@ -58,15 +60,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 @@ -89,13 +82,24 @@ func buildSubmitPayload(cfg *runConfig, commandPath, dlImage string, snap snapsh ForceSendFields: []string{"MaxRetries"}, } + // Carry the user's declared deps inline on spec.dependencies; the AI Runtime + // backend installs them via --deps-config. The SDK marshaler drops nil and empty + // slices, so a no-deps run omits the key. + envSpec := &compute.Environment{EnvironmentVersion: dlImage} + if len(deps) > 0 { + envSpec.Dependencies = deps + } + return jobs.SubmitRun{ - RunName: cfg.ExperimentName, + RunName: cfg.ExperimentName, + // budget_policy_id matches what the Python CLI and `ssh connect` send; + // usage_policy_id is the newer alias for the same thing on SubmitRun. + BudgetPolicyId: usagePolicyID, TimeoutSeconds: cfg.timeoutSeconds(), Tasks: []jobs.SubmitTask{st}, Environments: []jobs.JobEnvironment{{ EnvironmentKey: aiRuntimeEnvironmentKey, - Spec: &compute.Environment{EnvironmentVersion: dlImage}, + Spec: envSpec, }}, } } @@ -117,22 +121,69 @@ func submitToken(flag string, cfg *runConfig) (string, error) { return token, nil } +// withSpinner runs fn, showing an stderr spinner labeled msg when show is true. +// The spinner auto-degrades to nothing on a non-interactive terminal; show is +// false in JSON mode so the stdout envelope stream stays clean. +func withSpinner(ctx context.Context, show bool, msg string, fn func() error) error { + if !show { + return fn() + } + sp := cmdio.NewSpinner(ctx) + sp.Update(msg) + defer sp.Close() + return fn() +} + // submitWorkload runs the submit happy path: ensure the experiment directory, // upload the launch artifacts, assemble the Jobs payload, and submit it. It -// returns the new run_id and its dashboard URL. -func submitWorkload(ctx context.Context, w *databricks.WorkspaceClient, cfg *runConfig, configPath, idempotencyKey string) (int64, string, error) { - // Resolving usage_policy_name to a budget policy id is not ported yet; reject - // rather than silently drop. - if cfg.UsagePolicyName != nil { - return 0, "", errors.New("usage_policy_name is not yet supported") +// returns the new run_id and its dashboard URL. showProgress enables the +// stderr upload/packaging spinners (text mode only). +func submitWorkload(ctx context.Context, w *databricks.WorkspaceClient, cfg *runConfig, configPath, idempotencyKey string, showProgress bool) (int64, string, error) { + // Compute the launch dir and command_path up front — a read-only workspace lookup plus a + // local path build, no writes yet — so the pre-flight validates the real command_path. The + // same path is reused for the upload and submit below, so the validated path is the submitted + // one. + base, err := userWorkspaceDir(ctx, w) + if err != nil { + return 0, "", err } + runName := "" + if cfg.MLflowRunName != nil { + runName = *cfg.MLflowRunName + } + funcDir := cliLaunchDir(base, cfg.ExperimentName, runName) + commandPath := path.Join(funcDir, commandScriptName) - // Resolve the idempotency token first so a bad key fails before any upload. + // Pre-flight the config server-side before any upload, so a bad config fails with the + // backend's field-level errors and no orphaned artifacts. + if err := preflightValidate(ctx, w, cfg, commandPath); err != nil { + return 0, "", err + } + + // Resolve the idempotency token first so a bad key fails before any upload, + // and before the policy lookup below spends a round trip on it. token, err := submitToken(idempotencyKey, cfg) if err != nil { return 0, "", err } + // Resolve the usage policy to its id next, so a bad name fails fast with a + // clear (caller-fixable) message before we upload any artifacts. Validation + // guarantees name and id are mutually exclusive: a literal id is used as-is, a + // name is resolved against the workspace. + usagePolicyID := "" + if cfg.UsagePolicyID != nil { + usagePolicyID = strings.TrimSpace(*cfg.UsagePolicyID) + } + if cfg.UsagePolicyName != nil { + usagePolicyID, err = resolveUsagePolicyIDByName(ctx, w, *cfg.UsagePolicyName) + if err != nil { + return 0, "", err + } + } + + deps, _ := cfg.inlineDependencies() + experimentDir := "" if cfg.MLflowExperimentDirectory != nil { experimentDir = *cfg.MLflowExperimentDirectory @@ -141,16 +192,6 @@ func submitWorkload(ctx context.Context, w *databricks.WorkspaceClient, cfg *run return 0, "", err } - base, err := userWorkspaceDir(ctx, w) - if err != nil { - return 0, "", err - } - runName := "" - if cfg.MLflowRunName != nil { - runName = *cfg.MLflowRunName - } - funcDir := cliLaunchDir(base, cfg.ExperimentName, runName) - fc, err := filer.NewWorkspaceFilesClient(w, funcDir) if err != nil { return 0, "", err @@ -159,24 +200,30 @@ func submitWorkload(ctx context.Context, w *databricks.WorkspaceClient, cfg *run if err != nil { return 0, "", err } - if err := uploadArtifacts(ctx, fc, items); err != nil { + if err := withSpinner(ctx, showProgress, "Uploading yaml configuration files…", func() error { + return uploadArtifacts(ctx, fc, items) + }); err != nil { return 0, "", err } - // Package and upload the code snapshot, if any. The resulting paths ride on the - // ai_runtime_task; a run with no code_source leaves them empty. Snapshot is the - // only code_source type; guard against a nil block so snapshotCodeSource never - // dereferences a missing snapshot. + // 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. + err = withSpinner(ctx, showProgress, "Packaging code snapshot…", func() error { + var e error + snap, e = snapshotViaDABsUpload(ctx, w, cfg.CodeSource.Snapshot, configPath, fc, funcDir) + return e + }) if err != nil { return 0, "", err } } runtimeVersion, _ := cfg.runtimeVersion() - payload := buildSubmitPayload(cfg, path.Join(funcDir, commandScriptName), dlRuntimeImage(ctx, runtimeVersion), snap) + payload := buildSubmitPayload(cfg, commandPath, dlRuntimeImage(ctx, runtimeVersion), usagePolicyID, snap, deps) payload.IdempotencyToken = token // Submit returns as soon as the run is created; we don't wait for it to finish. diff --git a/experimental/air/cmd/runsubmit_test.go b/experimental/air/cmd/runsubmit_test.go index fd5103599df..dd8471bd3f5 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" @@ -13,6 +17,14 @@ import ( "github.com/stretchr/testify/require" ) +// stubValidateConfig registers an OK ValidateConfig response so submitWorkload's +// pre-flight passes. Register before AddDefaultHandlers (the router is first-wins). +func stubValidateConfig(server *testserver.Server) { + server.Handle("POST", "/api/2.0/ai-training/config:validate", func(req testserver.Request) any { + return validateConfigResponse{} + }) +} + func TestDlRuntimeImage(t *testing.T) { ctx := t.Context() // A config runtime version wins and is used bare. @@ -37,10 +49,12 @@ func TestBuildSubmitPayload(t *testing.T) { MLflowExperimentDirectory: new("/Workspace/Users/me/exp"), } - p := buildSubmitPayload(cfg, "/d/command.sh", "5", snapshotResult{}) + p := buildSubmitPayload(cfg, "/d/command.sh", "5", "", snapshotResult{}, nil) assert.Equal(t, "exp", p.RunName) assert.Equal(t, 1800, p.TimeoutSeconds) + // No policy configured: the field stays empty and is omitted from the wire form. + assert.Empty(t, p.BudgetPolicyId) require.Len(t, p.Environments, 1) assert.Equal(t, aiRuntimeEnvironmentKey, p.Environments[0].EnvironmentKey) require.NotNil(t, p.Environments[0].Spec) @@ -72,7 +86,7 @@ func TestBuildSubmitPayloadDefaultRetries(t *testing.T) { Command: new("x"), Compute: &computeConfig{AcceleratorType: "GPU_1xH100", NumAccelerators: 1}, } - task := buildSubmitPayload(cfg, "/d/command.sh", "4", snapshotResult{}).Tasks[0] + task := buildSubmitPayload(cfg, "/d/command.sh", "4", "", snapshotResult{}, nil).Tasks[0] assert.Equal(t, defaultMaxRetries, task.MaxRetries) assert.True(t, task.RetryOnTimeout) } @@ -87,7 +101,7 @@ func TestBuildSubmitPayloadNoRetries(t *testing.T) { Compute: &computeConfig{AcceleratorType: "GPU_1xH100", NumAccelerators: 1}, MaxRetries: new(0), } - task := buildSubmitPayload(cfg, "/d/command.sh", "4", snapshotResult{}).Tasks[0] + task := buildSubmitPayload(cfg, "/d/command.sh", "4", "", snapshotResult{}, nil).Tasks[0] assert.Equal(t, 0, task.MaxRetries) assert.False(t, task.RetryOnTimeout) @@ -97,6 +111,30 @@ func TestBuildSubmitPayloadNoRetries(t *testing.T) { assert.NotContains(t, string(b), "retry_on_timeout") } +// TestBuildSubmitPayloadInlineDependencies covers how deps land on the environment +// spec: a non-empty list is set alongside the runtime channel; empty and nil omit +// the key so the payload is unchanged. +func TestBuildSubmitPayloadInlineDependencies(t *testing.T) { + cfg := &runConfig{ + ExperimentName: "exp", + Command: new("x"), + Compute: &computeConfig{AcceleratorType: "GPU_8xH100", NumAccelerators: 8}, + } + + deps := []string{"torch==2.3.0", "--extra-index-url https://internal/pypi", "numpy"} + spec := buildSubmitPayload(cfg, "/d/command.sh", "5", "", snapshotResult{}, deps).Environments[0].Spec + assert.Equal(t, deps, spec.Dependencies) + assert.Equal(t, "5", spec.EnvironmentVersion) + + // The SDK marshaler drops empty/nil slices, so no "dependencies" key is emitted. + for _, empty := range [][]string{{}, nil} { + spec = buildSubmitPayload(cfg, "/d/command.sh", "5", "", snapshotResult{}, empty).Environments[0].Spec + b, err := json.Marshal(spec) + require.NoError(t, err) + assert.NotContains(t, string(b), "dependencies") + } +} + func TestSubmitToken(t *testing.T) { cfg := &runConfig{IdempotencyToken: new("from-config")} @@ -127,6 +165,7 @@ func TestSubmitWorkload(t *testing.T) { require.NoError(t, json.Unmarshal(req.Body, &got)) return jobs.SubmitRunResponse{RunId: 777} }) + stubValidateConfig(server) testserver.AddDefaultHandlers(server) w, err := databricks.NewWorkspaceClient(&databricks.Config{Host: server.URL, Token: "token"}) @@ -136,7 +175,7 @@ func TestSubmitWorkload(t *testing.T) { cfg, err := loadRunConfig(cfgPath) require.NoError(t, err) - runID, dashboardURL, err := submitWorkload(t.Context(), w, cfg, cfgPath, "idem-key") + runID, dashboardURL, err := submitWorkload(t.Context(), w, cfg, cfgPath, "idem-key", false) require.NoError(t, err) assert.Equal(t, int64(777), runID) assert.Contains(t, dashboardURL, "/jobs/runs/777") @@ -156,8 +195,42 @@ 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} + }) + stubValidateConfig(server) + 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", false) + 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) @@ -168,6 +241,49 @@ func TestSubmitWorkloadWithCodeSource(t *testing.T) { require.NoError(t, json.Unmarshal(req.Body, &got)) return jobs.SubmitRunResponse{RunId: 555} }) + stubValidateConfig(server) + testserver.AddDefaultHandlers(server) + 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", false) + 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} + }) + stubValidateConfig(server) testserver.AddDefaultHandlers(server) w, err := databricks.NewWorkspaceClient(&databricks.Config{Host: server.URL, Token: "token"}) require.NoError(t, err) @@ -189,28 +305,293 @@ 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", false) + 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} + }) + stubValidateConfig(server) + 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") + }) + stubValidateConfig(server) + 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} + }) + stubValidateConfig(server) + 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} + }) + stubValidateConfig(server) + 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", false) 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) } func TestSubmitWorkloadGuards(t *testing.T) { - w := newFakeWorkspaceClient(t) cfgPath := writeConfigFile(t, "run.yaml", minimalConfig) base, err := loadRunConfig(cfgPath) require.NoError(t, err) - t.Run("usage_policy_name rejected", func(t *testing.T) { + t.Run("unresolvable usage_policy_name fails before upload", func(t *testing.T) { + // An empty policy list makes the name unresolvable. Record every path the + // server sees so the "fails before any upload" ordering is asserted, not just + // asserted-by-comment: no import/mkdirs request may be made. + server := testserver.New(t) + t.Cleanup(server.Close) + var paths []string + server.Handle("GET", "/api/2.0/serverless-policies", func(req testserver.Request) any { + paths = append(paths, req.URL.Path) + return usagePoliciesResponse{} + }) + server.Handle("POST", "/api/2.0/workspace/{path...}", func(req testserver.Request) any { + paths = append(paths, req.URL.Path) + return testserver.Response{StatusCode: 200} + }) + stubValidateConfig(server) + testserver.AddDefaultHandlers(server) + pw, err := databricks.NewWorkspaceClient(&databricks.Config{Host: server.URL, Token: "token"}) + require.NoError(t, err) + cfg := *base - cfg.UsagePolicyName = new("p") - _, _, err := submitWorkload(t.Context(), w, &cfg, cfgPath, "") - require.ErrorContains(t, err, "usage_policy_name is not yet supported") + cfg.UsagePolicyName = new("nope") + _, _, err = submitWorkload(t.Context(), pw, &cfg, cfgPath, "", false) + require.ErrorContains(t, err, `no usage policy named "nope"`) + for _, p := range paths { + assert.NotContains(t, p, "/workspace/", "no workspace write may precede policy resolution") + } + }) +} + +// The resolved policy id must reach the submit payload, by literal id and by name. +func TestSubmitWorkloadSendsUsagePolicy(t *testing.T) { + const policyID = "12345678-90ab-cdef-1234-567890abcdef" + + setup := func(t *testing.T) (*databricks.WorkspaceClient, *jobs.SubmitRun) { + server := testserver.New(t) + t.Cleanup(server.Close) + + 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: 1} + }) + server.Handle("GET", "/api/2.0/serverless-policies", func(req testserver.Request) any { + return usagePoliciesResponse{Policies: []usagePolicy{{PolicyID: policyID, PolicyName: "team-a"}}} + }) + stubValidateConfig(server) + testserver.AddDefaultHandlers(server) + w, err := databricks.NewWorkspaceClient(&databricks.Config{Host: server.URL, Token: "token"}) + require.NoError(t, err) + return w, got + } + + t.Run("literal id", func(t *testing.T) { + w, got := setup(t) + cfgPath := writeConfigFile(t, "run.yaml", minimalConfig+"usage_policy_id: "+policyID+"\n") + cfg, err := loadRunConfig(cfgPath) + require.NoError(t, err) + + _, _, err = submitWorkload(cmdio.MockDiscard(t.Context()), w, cfg, cfgPath, "idem", false) + require.NoError(t, err) + assert.Equal(t, policyID, got.BudgetPolicyId) + }) + + t.Run("resolved from name", func(t *testing.T) { + w, got := setup(t) + cfgPath := writeConfigFile(t, "run.yaml", minimalConfig+"usage_policy_name: team-a\n") + cfg, err := loadRunConfig(cfgPath) + require.NoError(t, err) + + _, _, err = submitWorkload(cmdio.MockDiscard(t.Context()), w, cfg, cfgPath, "idem", false) + require.NoError(t, err) + assert.Equal(t, policyID, got.BudgetPolicyId) }) } diff --git a/experimental/air/cmd/runupload.go b/experimental/air/cmd/runupload.go index fb9ca00b987..ca0ffd9c480 100644 --- a/experimental/air/cmd/runupload.go +++ b/experimental/air/cmd/runupload.go @@ -8,7 +8,6 @@ import ( "io" "maps" "os" - "path/filepath" "slices" "strings" @@ -22,7 +21,6 @@ import ( const ( trainingConfigName = "training_config.yaml" commandScriptName = "command.sh" - requirementsName = "requirements.yaml" hyperparametersName = "hyperparameters.yaml" envVarsName = "env_vars.json" secretEnvVarsName = "secret_env_vars.json" @@ -45,16 +43,12 @@ type fileWriter interface { Write(ctx context.Context, name string, reader io.Reader, mode ...filer.WriteMode) error } -// requirementsDoc mirrors the on-disk requirements.yaml format so the worker -// parses synthesized inline dependencies identically to a user-provided file. -type requirementsDoc struct { - Version string `yaml:"version,omitempty"` - Dependencies []string `yaml:"dependencies"` -} - // buildArtifacts assembles the files to upload for a run: the merged config, the -// inline command as a script, requirements (from a file or synthesized from -// inline dependencies), and hyperparameters. configPath is the local YAML path. +// inline command as a script, and hyperparameters. configPath is the local YAML +// path. +// +// Dependencies are not uploaded here; they ride inline on the serverless +// environment's spec.dependencies (see buildSubmitPayload). func buildArtifacts(cfg *runConfig, configPath string) ([]uploadItem, error) { // TODO(DABs): with no _bases_/overrides ported yet, the merged config is the // file as-is; once those land, upload the re-serialized merged YAML instead. @@ -72,28 +66,6 @@ func buildArtifacts(cfg *runConfig, configPath string) ([]uploadItem, error) { {commandScriptName, []byte(*cfg.Command)}, } - switch reqPath, ok := cfg.requirementsFile(); { - case ok: - // Resolve a relative requirements path against the config's directory. - if !filepath.IsAbs(reqPath) { - reqPath = filepath.Join(filepath.Dir(configPath), reqPath) - } - data, err := os.ReadFile(reqPath) - if err != nil { - return nil, fmt.Errorf("failed to read requirements file %s: %w", reqPath, err) - } - items = append(items, uploadItem{requirementsName, data}) - default: - if deps, ok := cfg.inlineDependencies(); ok { - version, _ := cfg.runtimeVersion() - data, err := yaml.Marshal(requirementsDoc{Version: version, Dependencies: deps}) - if err != nil { - return nil, fmt.Errorf("failed to synthesize requirements.yaml: %w", err) - } - items = append(items, uploadItem{requirementsName, data}) - } - } - if len(cfg.Parameters) > 0 { data, err := yaml.Marshal(cfg.Parameters) if err != nil { diff --git a/experimental/air/cmd/runupload_test.go b/experimental/air/cmd/runupload_test.go index 0c87524735d..1b3e04ee00d 100644 --- a/experimental/air/cmd/runupload_test.go +++ b/experimental/air/cmd/runupload_test.go @@ -57,30 +57,21 @@ func TestBuildArtifacts_CommandAndConfig(t *testing.T) { assert.Equal(t, "python train.py", string(items[1].data)) } -func TestBuildArtifacts_InlineRequirementsAndParameters(t *testing.T) { +func TestBuildArtifacts_ParametersButNoRequirements(t *testing.T) { path := writeConfigFile(t, "run.yaml", "x: y\n") cfg := &runConfig{ Command: new("echo hi"), Environment: &environmentConfig{ - Dependencies: dependencies{set: true, isList: true, list: []string{"torch", "numpy"}}, + Dependencies: dependencies{set: true, list: []string{"torch", "numpy"}}, Version: stringOrInt{set: true, raw: "5"}, }, Parameters: map[string]any{"lr": 0.1}, } + // Inline deps are not uploaded, so the artifacts are config, command, and params. items, err := buildArtifacts(cfg, path) require.NoError(t, err) - assert.Equal(t, []string{trainingConfigName, commandScriptName, requirementsName, hyperparametersName}, itemNames(items)) - - var reqIdx int - for i, it := range items { - if it.name == requirementsName { - reqIdx = i - } - } - req := string(items[reqIdx].data) - assert.Contains(t, req, "version: \"5\"") - assert.Contains(t, req, "- torch") + assert.Equal(t, []string{trainingConfigName, commandScriptName, hyperparametersName}, itemNames(items)) } func TestBuildArtifacts_EnvVarsAndSecrets(t *testing.T) { @@ -103,20 +94,6 @@ func TestBuildArtifacts_EnvVarsAndSecrets(t *testing.T) { assert.JSONEq(t, `[{"name":"HF_TOKEN","secret_scope":"myscope","secret_key":"hf"}]`, string(byName[secretEnvVarsName])) } -func TestBuildArtifacts_RequirementsFile(t *testing.T) { - dir := t.TempDir() - require.NoError(t, os.WriteFile(filepath.Join(dir, "run.yaml"), []byte("x: y\n"), 0o600)) - require.NoError(t, os.WriteFile(filepath.Join(dir, "reqs.yaml"), []byte("version: 4\n"), 0o600)) - cfg := &runConfig{ - Command: new("echo hi"), - Environment: &environmentConfig{Dependencies: dependencies{set: true, isList: false, path: "reqs.yaml"}}, - } - - items, err := buildArtifacts(cfg, filepath.Join(dir, "run.yaml")) - require.NoError(t, err) - assert.Contains(t, itemNames(items), requirementsName) -} - func TestBuildArtifacts_OversizeConfigRejected(t *testing.T) { path := writeConfigFile(t, "run.yaml", strings.Repeat("a", maxConfigYAMLBytes+1)) _, err := buildArtifacts(&runConfig{Command: new("x")}, path) @@ -143,13 +120,3 @@ func TestUploadArtifacts_WriteError(t *testing.T) { err := uploadArtifacts(t.Context(), errWriter{}, []uploadItem{{trainingConfigName, []byte("x")}}) require.ErrorContains(t, err, "failed to upload "+trainingConfigName) } - -func TestBuildArtifacts_MissingRequirementsFile(t *testing.T) { - cfgPath := writeConfigFile(t, "run.yaml", "x: y\n") - cfg := &runConfig{ - Command: new("echo hi"), - Environment: &environmentConfig{Dependencies: dependencies{set: true, isList: false, path: "nope.yaml"}}, - } - _, err := buildArtifacts(cfg, cfgPath) - require.ErrorContains(t, err, "failed to read requirements file") -} 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..dd043fdfa50 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", 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 deleted file mode 100644 index e28d7f66730..00000000000 --- a/experimental/air/cmd/stubs_test.go +++ /dev/null @@ -1,27 +0,0 @@ -package aircmd - -import ( - "fmt" - "testing" - - "github.com/spf13/cobra" - "github.com/stretchr/testify/assert" - "github.com/stretchr/testify/require" -) - -// TestStubCommandsReturnNotImplemented asserts each unimplemented subcommand -// 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(), - } - - for name, cmd := range stubs { - t.Run(name, func(t *testing.T) { - require.NotNil(t, cmd.RunE, "command should define RunE") - err := cmd.RunE(cmd, nil) - assert.EqualError(t, err, fmt.Sprintf("`air %s` is not implemented yet", name)) - }) - } -} diff --git a/experimental/air/cmd/usagepolicy.go b/experimental/air/cmd/usagepolicy.go new file mode 100644 index 00000000000..fb3c0e1c550 --- /dev/null +++ b/experimental/air/cmd/usagepolicy.go @@ -0,0 +1,178 @@ +package aircmd + +import ( + "context" + "errors" + "fmt" + "net/http" + "slices" + "strings" + + "github.com/databricks/databricks-sdk-go" + "github.com/databricks/databricks-sdk-go/client" +) + +// serverlessPoliciesPath is the workspace-scoped ListBudgetPolicies endpoint on +// the serverless-policy service. This is called with a raw client.Do because the +// SDK only models the account-scoped /api/2.1/accounts/{id}/budget-policies +// service, which is a different (and unusable here) endpoint. +const serverlessPoliciesPath = "/api/2.0/serverless-policies" + +// maxPolicyPageSize is the server's cap: anything larger is coerced down to it. +// Request the max so the common case (a workspace with a handful of policies) is +// a single round-trip. +const maxPolicyPageSize = 1000 + +// maxPolicySuggestions bounds the candidate names surfaced in a "no exact match" +// error: enough to spot a typo or casing mistake without dumping a huge list. +const maxPolicySuggestions = 10 + +// usagePolicy models only the two fields resolution needs. The API also returns +// display_name, which is deliberately not read: policy_name is the unique key +// (unique among active policies) and the only field filter_by can match on. The +// two are identical for user-created policies and diverge only for the system +// defaults, whose display_name is fixed to "Default Policy". +type usagePolicy struct { + PolicyID string `json:"policy_id"` + PolicyName string `json:"policy_name"` +} + +type usagePoliciesResponse struct { + Policies []usagePolicy `json:"policies"` + NextPageToken string `json:"next_page_token"` +} + +// listUsagePolicies pages the serverless-policy index for policies matching +// policyName, which is sent as filter_by.policy_name: a partial, +// case-insensitive server-side filter. Callers must pass a non-empty name. +// +// The filter key is spelled with its dotted proto path rather than as a nested +// map: the SDK only flattens nesting for struct-typed query values, and would +// format a nested map with %v into a useless "map[...]" literal. +func listUsagePolicies(ctx context.Context, w *databricks.WorkspaceClient, policyName string) ([]usagePolicy, error) { + apiClient, err := client.New(w.Config) + if err != nil { + return nil, fmt.Errorf("failed to create API client: %w", err) + } + + var out []usagePolicy + // The index can return the same policy on more than one page, and a stuck or + // cycling cursor can repeat a whole page; dedupe both so an unambiguous name + // never looks like an ambiguous match downstream. + seenIDs := map[string]bool{} + seenTokens := map[string]bool{} + var pageToken string + for { + query := map[string]any{ + "page_size": maxPolicyPageSize, + "filter_by.policy_name": policyName, + } + if pageToken != "" { + query["page_token"] = pageToken + } + + var resp usagePoliciesResponse + err = apiClient.Do(ctx, http.MethodGet, serverlessPoliciesPath, nil, nil, query, &resp) + if err != nil { + return nil, fmt.Errorf("failed to list usage policies: %w", err) + } + + for _, p := range resp.Policies { + if seenIDs[p.PolicyID] { + continue + } + seenIDs[p.PolicyID] = true + out = append(out, p) + } + + if resp.NextPageToken == "" || seenTokens[resp.NextPageToken] { + return out, nil + } + seenTokens[resp.NextPageToken] = true + pageToken = resp.NextPageToken + } +} + +// resolveUsagePolicyIDByName resolves a usage policy name to its UUID policy id. +// +// The server-side filter is a partial match, so the exact (but case-insensitive) +// match is re-applied locally; policy names are unique among active policies. +// +// name is matched against policy_name, not the policy's display_name. For a +// user-created policy the two are the same, so the distinction only surfaces for +// the system defaults; a user who supplies a display name that isn't a +// policy_name gets the not-found error with the real names as candidates. +func resolveUsagePolicyIDByName(ctx context.Context, w *databricks.WorkspaceClient, name string) (string, error) { + target := strings.TrimSpace(name) + // Guard the contract independently of the YAML validator: an empty filter would + // otherwise list (then reject against) every policy in the workspace. + if target == "" { + return "", errors.New("a usage policy name must be a non-empty string") + } + + policies, err := listUsagePolicies(ctx, w, target) + if err != nil { + return "", err + } + + var matches []usagePolicy + for _, p := range policies { + if strings.EqualFold(strings.TrimSpace(p.PolicyName), target) { + matches = append(matches, p) + } + } + + switch len(matches) { + case 1: + if matches[0].PolicyID == "" { + return "", fmt.Errorf("policy %q has no policy_id in the API response", target) + } + return matches[0].PolicyID, nil + + case 0: + // policies holds the partial-match candidates the server returned for this + // name; surface a few to help the user fix a typo or casing. These are + // policy_name values, which is also what a user who typed a policy's UI + // display name needs to see: the two differ only for the system default + // policies, so listing the real names points them at the right one. + return "", fmt.Errorf("no usage policy named %q was found in this workspace%s", target, suggestionHint(policies)) + + default: + // Multiple exact (case-insensitive) matches should not happen given name + // uniqueness, but guard so we never silently pick the wrong policy. + ids := make([]string, 0, len(matches)) + for _, p := range matches { + ids = append(ids, fmt.Sprintf("%q", p.PolicyID)) + } + return "", fmt.Errorf("multiple usage policies match the name %q (ids: %s); please disambiguate with your workspace admin", + target, strings.Join(ids, ", ")) + } +} + +// suggestionHint renders a deduplicated, sorted "did you mean" clause for the +// candidates the partial filter returned, or "" when there are none. +func suggestionHint(candidates []usagePolicy) string { + names := make([]string, 0, len(candidates)) + for _, p := range candidates { + if p.PolicyName != "" { + names = append(names, p.PolicyName) + } + } + slices.Sort(names) + names = slices.Compact(names) + if len(names) == 0 { + return "" + } + + shown := names + suffix := "" + if len(names) > maxPolicySuggestions { + shown = names[:maxPolicySuggestions] + suffix = ", ..." + } + quoted := make([]string, 0, len(shown)) + for _, n := range shown { + quoted = append(quoted, fmt.Sprintf("%q", n)) + } + return fmt.Sprintf(". Did you mean one of: %s%s?", strings.Join(quoted, ", "), suffix) +} diff --git a/experimental/air/cmd/usagepolicy_test.go b/experimental/air/cmd/usagepolicy_test.go new file mode 100644 index 00000000000..441657bdb83 --- /dev/null +++ b/experimental/air/cmd/usagepolicy_test.go @@ -0,0 +1,227 @@ +package aircmd + +import ( + "net/url" + "strconv" + "testing" + + "github.com/databricks/cli/libs/testserver" + "github.com/databricks/databricks-sdk-go" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// policyServer serves /api/2.0/serverless-policies from the given pages, +// returning one page per request and recording each request's query. +func policyServer(t *testing.T, pages ...usagePoliciesResponse) (*databricks.WorkspaceClient, *[]url.Values) { + server := testserver.New(t) + t.Cleanup(server.Close) + + var queries []url.Values + var n int + server.Handle("GET", "/api/2.0/serverless-policies", func(req testserver.Request) any { + queries = append(queries, req.URL.Query()) + page := pages[min(n, len(pages)-1)] + n++ + return page + }) + testserver.AddDefaultHandlers(server) + w, err := databricks.NewWorkspaceClient(&databricks.Config{Host: server.URL, Token: "token"}) + require.NoError(t, err) + return w, &queries +} + +func TestListUsagePoliciesSendsFilterAndPageSize(t *testing.T) { + w, queries := policyServer(t, usagePoliciesResponse{ + Policies: []usagePolicy{{PolicyID: "id-1", PolicyName: "alpha"}}, + }) + + policies, err := listUsagePolicies(t.Context(), w, "alpha") + require.NoError(t, err) + assert.Equal(t, []usagePolicy{{PolicyID: "id-1", PolicyName: "alpha"}}, policies) + + require.Len(t, *queries, 1) + q := (*queries)[0] + // The filter must arrive under its flattened proto path, not as a nested map. + assert.Equal(t, "alpha", q.Get("filter_by.policy_name")) + assert.Equal(t, strconv.Itoa(maxPolicyPageSize), q.Get("page_size")) +} + +func TestListUsagePoliciesPaginates(t *testing.T) { + w, queries := policyServer(t, + usagePoliciesResponse{Policies: []usagePolicy{{PolicyID: "id-1", PolicyName: "a"}}, NextPageToken: "tok"}, + usagePoliciesResponse{Policies: []usagePolicy{{PolicyID: "id-2", PolicyName: "b"}}}, + ) + + policies, err := listUsagePolicies(t.Context(), w, "a") + require.NoError(t, err) + assert.Equal(t, []usagePolicy{ + {PolicyID: "id-1", PolicyName: "a"}, + {PolicyID: "id-2", PolicyName: "b"}, + }, policies) + + require.Len(t, *queries, 2) + assert.Equal(t, "tok", (*queries)[1].Get("page_token")) +} + +// A page token that repeats itself must not spin forever, and the repeated page +// must not be counted twice: a duplicated policy would otherwise look like an +// ambiguous name to resolveUsagePolicyIDByName. +func TestListUsagePoliciesStopsOnRepeatedToken(t *testing.T) { + w, _ := policyServer(t, usagePoliciesResponse{ + Policies: []usagePolicy{{PolicyID: "id-1", PolicyName: "a"}}, + NextPageToken: "same", + }) + + policies, err := listUsagePolicies(t.Context(), w, "a") + require.NoError(t, err) + assert.Equal(t, []usagePolicy{{PolicyID: "id-1", PolicyName: "a"}}, policies) +} + +// An A->B->A token cycle also terminates, rather than only a self-repeat. +func TestListUsagePoliciesStopsOnTokenCycle(t *testing.T) { + w, _ := policyServer(t, + usagePoliciesResponse{Policies: []usagePolicy{{PolicyID: "id-1", PolicyName: "a"}}, NextPageToken: "b"}, + usagePoliciesResponse{Policies: []usagePolicy{{PolicyID: "id-2", PolicyName: "b"}}, NextPageToken: "a"}, + usagePoliciesResponse{Policies: []usagePolicy{{PolicyID: "id-3", PolicyName: "c"}}, NextPageToken: "b"}, + ) + + policies, err := listUsagePolicies(t.Context(), w, "a") + require.NoError(t, err) + assert.Len(t, policies, 3) +} + +// The same policy arriving on two pages is returned once. +func TestListUsagePoliciesDedupesAcrossPages(t *testing.T) { + w, _ := policyServer(t, + usagePoliciesResponse{Policies: []usagePolicy{{PolicyID: "id-1", PolicyName: "team-a"}}, NextPageToken: "tok"}, + usagePoliciesResponse{Policies: []usagePolicy{{PolicyID: "id-1", PolicyName: "team-a"}}}, + ) + + policies, err := listUsagePolicies(t.Context(), w, "team-a") + require.NoError(t, err) + assert.Equal(t, []usagePolicy{{PolicyID: "id-1", PolicyName: "team-a"}}, policies) + + // A repeat must not read as an ambiguous match. + w2, _ := policyServer(t, + usagePoliciesResponse{Policies: []usagePolicy{{PolicyID: "id-1", PolicyName: "team-a"}}, NextPageToken: "tok"}, + usagePoliciesResponse{Policies: []usagePolicy{{PolicyID: "id-1", PolicyName: "team-a"}}}, + ) + got, err := resolveUsagePolicyIDByName(t.Context(), w2, "team-a") + require.NoError(t, err) + assert.Equal(t, "id-1", got) +} + +func TestResolveUsagePolicyIDByName(t *testing.T) { + const id = "12345678-90ab-cdef-1234-567890abcdef" + + t.Run("exact match", func(t *testing.T) { + w, _ := policyServer(t, usagePoliciesResponse{ + Policies: []usagePolicy{{PolicyID: id, PolicyName: "team-a"}}, + }) + got, err := resolveUsagePolicyIDByName(t.Context(), w, "team-a") + require.NoError(t, err) + assert.Equal(t, id, got) + }) + + // The server filter is partial; only the exact name (case-insensitively) wins. + t.Run("case-insensitive exact match wins over partial", func(t *testing.T) { + w, _ := policyServer(t, usagePoliciesResponse{ + Policies: []usagePolicy{ + {PolicyID: "other", PolicyName: "team-a-staging"}, + {PolicyID: id, PolicyName: "Team-A"}, + }, + }) + got, err := resolveUsagePolicyIDByName(t.Context(), w, "team-a") + require.NoError(t, err) + assert.Equal(t, id, got) + }) + + t.Run("no match suggests candidates", func(t *testing.T) { + w, _ := policyServer(t, usagePoliciesResponse{ + Policies: []usagePolicy{{PolicyID: "x", PolicyName: "team-a-staging"}}, + }) + _, err := resolveUsagePolicyIDByName(t.Context(), w, "team-a") + require.ErrorContains(t, err, `no usage policy named "team-a"`) + require.ErrorContains(t, err, `Did you mean one of: "team-a-staging"?`) + }) + + t.Run("no match and no candidates omits the hint", func(t *testing.T) { + w, _ := policyServer(t, usagePoliciesResponse{}) + _, err := resolveUsagePolicyIDByName(t.Context(), w, "team-a") + require.ErrorContains(t, err, `no usage policy named "team-a"`) + assert.NotContains(t, err.Error(), "Did you mean") + }) + + t.Run("suggestions are capped", func(t *testing.T) { + var policies []usagePolicy + for i := range maxPolicySuggestions + 5 { + // Zero-padded so lexical order matches numeric order. + policies = append(policies, usagePolicy{PolicyID: strconv.Itoa(i), PolicyName: "team-a-" + strconv.Itoa(100+i)}) + } + w, _ := policyServer(t, usagePoliciesResponse{Policies: policies}) + _, err := resolveUsagePolicyIDByName(t.Context(), w, "team-a") + require.ErrorContains(t, err, `"team-a-109", ...?`) + assert.NotContains(t, err.Error(), "team-a-110") + }) + + t.Run("ambiguous match refuses to guess", func(t *testing.T) { + w, _ := policyServer(t, usagePoliciesResponse{ + Policies: []usagePolicy{ + {PolicyID: "id-1", PolicyName: "team-a"}, + {PolicyID: "id-2", PolicyName: "TEAM-A"}, + }, + }) + _, err := resolveUsagePolicyIDByName(t.Context(), w, "team-a") + require.ErrorContains(t, err, "multiple usage policies match") + require.ErrorContains(t, err, `"id-1", "id-2"`) + }) + + t.Run("match without an id is an error", func(t *testing.T) { + w, _ := policyServer(t, usagePoliciesResponse{ + Policies: []usagePolicy{{PolicyName: "team-a"}}, + }) + _, err := resolveUsagePolicyIDByName(t.Context(), w, "team-a") + require.ErrorContains(t, err, "has no policy_id") + }) + + // An empty filter would list every policy in the workspace, so a blank name is + // rejected without a round-trip. + t.Run("blank name is rejected", func(t *testing.T) { + w, queries := policyServer(t, usagePoliciesResponse{}) + _, err := resolveUsagePolicyIDByName(t.Context(), w, " ") + require.ErrorContains(t, err, "must be a non-empty string") + assert.Empty(t, *queries) + }) + + // A failed lookup must surface, never fall through to an empty (= no policy) id. + t.Run("api error surfaces", func(t *testing.T) { + server := testserver.New(t) + t.Cleanup(server.Close) + server.Handle("GET", "/api/2.0/serverless-policies", func(req testserver.Request) any { + return testserver.Response{StatusCode: 403, Body: `{"error_code":"PERMISSION_DENIED","message":"nope"}`} + }) + testserver.AddDefaultHandlers(server) + w, err := databricks.NewWorkspaceClient(&databricks.Config{Host: server.URL, Token: "token"}) + require.NoError(t, err) + + _, err = resolveUsagePolicyIDByName(t.Context(), w, "team-a") + require.ErrorContains(t, err, "failed to list usage policies") + }) +} + +// The --override path re-decodes and re-validates the config, so the policy rules +// must hold there too and not just for fields set in the YAML file. +func TestOverrideUsagePolicyValidation(t *testing.T) { + t.Run("override trips mutual exclusion", func(t *testing.T) { + cfgPath := writeConfigFile(t, "run.yaml", minimalConfig+"usage_policy_id: 12345678-90ab-cdef-1234-567890abcdef\n") + _, err := loadRunConfigWithOverrides(t.Context(), cfgPath, []string{"usage_policy_name=team-a"}) + require.ErrorContains(t, err, "mutually exclusive") + }) + + t.Run("override id is UUID-checked", func(t *testing.T) { + cfgPath := writeConfigFile(t, "run.yaml", minimalConfig) + _, err := loadRunConfigWithOverrides(t.Context(), cfgPath, []string{"usage_policy_id=team-a"}) + require.ErrorContains(t, err, "must be a UUID") + }) +} diff --git a/experimental/air/cmd/validateconfig.go b/experimental/air/cmd/validateconfig.go new file mode 100644 index 00000000000..d65339878f0 --- /dev/null +++ b/experimental/air/cmd/validateconfig.go @@ -0,0 +1,143 @@ +package aircmd + +import ( + "context" + "errors" + "fmt" + "net/http" + "strings" + + "github.com/databricks/databricks-sdk-go" + "github.com/databricks/databricks-sdk-go/apierr" + "github.com/databricks/databricks-sdk-go/client" +) + +// validateConfigPath is AiTrainingService's pre-flight: it checks a training +// config server-side and returns the problems, without submitting. Called with a +// raw client.Do because the SDK does not model AiTrainingService. +const validateConfigPath = "/api/2.0/ai-training/config:validate" + +// configFieldError is one problem the server found, addressed to the config +// field that caused it. Mirrors the proto FieldError. +type configFieldError struct { + Path string `json:"path"` + Message string `json:"message"` + Code string `json:"code"` +} + +type validateConfigResponse struct { + Errors []configFieldError `json:"errors"` +} + +// preflightValidate checks the config against the backend before any upload, so +// a bad config fails fast with the server's own field-level errors. +// +// It fails open: the endpoint is behind a SAFE flag and older workspaces do not +// have it, so a disabled or missing endpoint skips the check and lets submission +// proceed (where the config is validated again, authoritatively). A 5xx is a +// backend problem, not the user's config, so it fails open too. Only a 4xx (the +// server rejected the config) or a populated error list blocks. +func preflightValidate(ctx context.Context, w *databricks.WorkspaceClient, cfg *runConfig, commandPath string) error { + apiClient, err := client.New(w.Config) + if err != nil { + return fmt.Errorf("failed to create API client: %w", err) + } + + var resp validateConfigResponse + err = apiClient.Do(ctx, http.MethodPost, validateConfigPath, nil, nil, validateConfigRequest(cfg, commandPath), &resp) + if err != nil { + if endpointUnavailable(err) || serverError(err) { + return nil + } + return fmt.Errorf("failed to validate config: %w", err) + } + if len(resp.Errors) == 0 { + return nil + } + return errors.New(formatConfigErrors(resp.Errors)) +} + +// validateConfigRequest builds the {task, run_options} body from the user's config. commandPath is +// the workspace path where the command script will be uploaded; the caller computes it before this +// call so the server can validate the real path. `parameters` is intentionally omitted: it is +// free-form nested hyperparameters uploaded as a YAML file at submit, not the proto's string map. +func validateConfigRequest(cfg *runConfig, commandPath string) map[string]any { + compute := map[string]any{} + if cfg.Compute != nil { + compute["accelerator_type"] = cfg.Compute.AcceleratorType + compute["accelerator_count"] = cfg.Compute.NumAccelerators + } + task := map[string]any{ + "experiment": cfg.ExperimentName, + "deployments": []any{map[string]any{"command_path": commandPath, "compute": compute}}, + } + putOpt(task, "mlflow_run", cfg.MLflowRunName) + putOpt(task, "mlflow_experiment_directory", cfg.MLflowExperimentDirectory) + + req := map[string]any{"task": task} + if runOptions := validateConfigRunOptions(cfg); len(runOptions) > 0 { + req["run_options"] = runOptions + } + return req +} + +// validateConfigRunOptions gathers the run-level fields into run_options, +// omitting any the user didn't set. +func validateConfigRunOptions(cfg *runConfig) map[string]any { + runOptions := map[string]any{} + putOpt(runOptions, "max_retries", cfg.MaxRetries) + putOpt(runOptions, "timeout_minutes", cfg.TimeoutMinutes) + putOpt(runOptions, "idempotency_token", cfg.IdempotencyToken) + putOpt(runOptions, "usage_policy_name", cfg.UsagePolicyName) + putOpt(runOptions, "usage_policy_id", cfg.UsagePolicyID) + if len(cfg.EnvVariables) > 0 { + runOptions["env_variables"] = cfg.EnvVariables + } + if len(cfg.Secrets) > 0 { + runOptions["secrets"] = cfg.Secrets + } + return runOptions +} + +// putOpt sets key to the pointer's value only when it is non-nil, so an unset +// config field is left out of the request rather than sent as a zero value. +func putOpt[T any](m map[string]any, key string, value *T) { + if value != nil { + m[key] = *value + } +} + +// endpointUnavailable reports whether the failure means the endpoint isn't there +// to answer — the flag is off, or the workspace predates it — as opposed to the +// config being rejected. Those cases fail open. +func endpointUnavailable(err error) bool { + apiErr, ok := errors.AsType[*apierr.APIError](err) + return ok && (apiErr.ErrorCode == "FEATURE_DISABLED" || + apiErr.StatusCode == http.StatusNotFound || + apiErr.StatusCode == http.StatusNotImplemented) +} + +// serverError reports whether the failure is a 5xx: a backend problem, not the +// user's config. The SDK already retries the transient subset (503, 429, IO +// errors); a 5xx that still surfaces here fails open, since blocking a submit on +// a backend blip isn't actionable and submit re-validates anyway. +func serverError(err error) bool { + apiErr, ok := errors.AsType[*apierr.APIError](err) + return ok && apiErr.StatusCode >= 500 +} + +// formatConfigErrors renders the field errors as one message, one problem per +// line, each pointing at the config field the user wrote. +func formatConfigErrors(fieldErrors []configFieldError) string { + var b strings.Builder + b.WriteString("config validation failed:") + for _, e := range fieldErrors { + b.WriteString("\n ") + if e.Path != "" { + b.WriteString(e.Path) + b.WriteString(": ") + } + b.WriteString(e.Message) + } + return b.String() +} diff --git a/experimental/air/cmd/validateconfig_test.go b/experimental/air/cmd/validateconfig_test.go new file mode 100644 index 00000000000..b6245878f78 --- /dev/null +++ b/experimental/air/cmd/validateconfig_test.go @@ -0,0 +1,127 @@ +package aircmd + +import ( + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func baseRunConfig() *runConfig { + return &runConfig{ + ExperimentName: "llama-fine-tune", + Compute: &computeConfig{NumAccelerators: 16, AcceleratorType: "GPU_8xH100"}, + } +} + +// validateServer serves one ValidateConfig response with the given status and +// body, and records the request body it received. +func validateServer(t *testing.T, status int, body string, gotReq *map[string]any) *httptest.Server { + t.Helper() + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == validateConfigPath { + if gotReq != nil { + _ = json.NewDecoder(r.Body).Decode(gotReq) + } + w.WriteHeader(status) + _, _ = w.Write([]byte(body)) + return + } + _, _ = w.Write([]byte(`{}`)) + })) + t.Cleanup(srv.Close) + return srv +} + +func TestPreflightValidatePasses(t *testing.T) { + srv := validateServer(t, http.StatusOK, `{}`, nil) + err := preflightValidate(t.Context(), newTestWorkspaceClient(t, srv.URL), baseRunConfig(), "/Workspace/Users/me/cmd.sh") + assert.NoError(t, err) +} + +func TestPreflightValidateReportsErrors(t *testing.T) { + body := `{"errors":[ + {"path":"experiment","message":"only letters, digits, hyphens, underscores","code":"DISALLOWED_CHARACTERS"}, + {"path":"deployments[0].compute.accelerator_count","message":"must be a multiple of 8","code":"COUNT_NOT_MULTIPLE"} + ]}` + srv := validateServer(t, http.StatusOK, body, nil) + err := preflightValidate(t.Context(), newTestWorkspaceClient(t, srv.URL), baseRunConfig(), "/Workspace/Users/me/cmd.sh") + require.Error(t, err) + // Every problem is surfaced, each pointing at its config field. + assert.Contains(t, err.Error(), "experiment: only letters") + assert.Contains(t, err.Error(), "deployments[0].compute.accelerator_count: must be a multiple of 8") +} + +func TestPreflightValidateFailsOpenWhenDisabled(t *testing.T) { + // The endpoint is behind a SAFE flag; a disabled endpoint must not block the run. + srv := validateServer(t, http.StatusBadRequest, + `{"error_code":"FEATURE_DISABLED","message":"ValidateConfig is not yet enabled."}`, nil) + err := preflightValidate(t.Context(), newTestWorkspaceClient(t, srv.URL), baseRunConfig(), "/Workspace/Users/me/cmd.sh") + assert.NoError(t, err) +} + +func TestPreflightValidateFailsOpenWhenNotFound(t *testing.T) { + // A workspace that predates the endpoint returns 404; skip and let submit proceed. + srv := validateServer(t, http.StatusNotFound, `{"error_code":"ENDPOINT_NOT_FOUND","message":"not found"}`, nil) + err := preflightValidate(t.Context(), newTestWorkspaceClient(t, srv.URL), baseRunConfig(), "/Workspace/Users/me/cmd.sh") + assert.NoError(t, err) +} + +func TestPreflightValidateFailsOpenOnServerError(t *testing.T) { + // A 5xx is a backend problem, not the user's config; blocking a submit on it + // isn't actionable, and submit re-validates anyway. + srv := validateServer(t, http.StatusInternalServerError, + `{"error_code":"INTERNAL_ERROR","message":"backend blew up"}`, nil) + err := preflightValidate(t.Context(), newTestWorkspaceClient(t, srv.URL), baseRunConfig(), "/Workspace/Users/me/cmd.sh") + assert.NoError(t, err) +} + +func TestPreflightValidateBlocksOnClientError(t *testing.T) { + // A 4xx other than the fail-open cases means the request itself was rejected + // (e.g. the proto hook flagged a missing required field); surface it. + srv := validateServer(t, http.StatusBadRequest, + `{"error_code":"INVALID_PARAMETER_VALUE","message":"command_path is required"}`, nil) + err := preflightValidate(t.Context(), newTestWorkspaceClient(t, srv.URL), baseRunConfig(), "/Workspace/Users/me/cmd.sh") + require.Error(t, err) +} + +func TestValidateConfigRequestShape(t *testing.T) { + var gotReq map[string]any + srv := validateServer(t, http.StatusOK, `{}`, &gotReq) + + cfg := baseRunConfig() + cfg.MaxRetries = new(3) + cfg.EnvVariables = map[string]string{"HF_HOME": "/tmp/hf"} + err := preflightValidate(t.Context(), newTestWorkspaceClient(t, srv.URL), cfg, "/Workspace/Users/me/cmd.sh") + require.NoError(t, err) + + task := gotReq["task"].(map[string]any) + assert.Equal(t, "llama-fine-tune", task["experiment"]) + deployment := task["deployments"].([]any)[0].(map[string]any) + compute := deployment["compute"].(map[string]any) + assert.Equal(t, "GPU_8xH100", compute["accelerator_type"]) + assert.EqualValues(t, 16, compute["accelerator_count"]) + + runOptions := gotReq["run_options"].(map[string]any) + assert.EqualValues(t, 3, runOptions["max_retries"]) + assert.Equal(t, map[string]any{"HF_HOME": "/tmp/hf"}, runOptions["env_variables"]) +} + +func TestValidateConfigRequestOmitsUnsetOptions(t *testing.T) { + // A minimal config carries no run_options and only the fields it set, so the + // server never validates values the user didn't provide. + var gotReq map[string]any + srv := validateServer(t, http.StatusOK, `{}`, &gotReq) + + err := preflightValidate(t.Context(), newTestWorkspaceClient(t, srv.URL), baseRunConfig(), "/Workspace/Users/me/cmd.sh") + require.NoError(t, err) + + _, hasRunOptions := gotReq["run_options"] + assert.False(t, hasRunOptions) + task := gotReq["task"].(map[string]any) + _, hasMlflowRun := task["mlflow_run"] + assert.False(t, hasMlflowRun) +}