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/get-ai-runtime/test.toml b/acceptance/experimental/air/get-ai-runtime/test.toml index de442aca0ba..07d75ae83c8 100644 --- a/acceptance/experimental/air/get-ai-runtime/test.toml +++ b/acceptance/experimental/air/get-ai-runtime/test.toml @@ -1,6 +1,3 @@ -# This command does not deploy a bundle, so no engine matrix is needed. -EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["direct"] - # On Windows, Git Bash rewrites the leading-/ workspace paths passed to # `workspace mkdirs`/`import` into C:/... paths; disable that conversion. [Env] diff --git a/acceptance/experimental/air/get/test.toml b/acceptance/experimental/air/get/test.toml index e0ebbb2ba35..3545a34616e 100644 --- a/acceptance/experimental/air/get/test.toml +++ b/acceptance/experimental/air/get/test.toml @@ -1,6 +1,3 @@ -# This command does not deploy a bundle, so no engine matrix is needed. -EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["direct"] - # The SDK occasionally probes host reachability with a HEAD request; stub it so # the test is deterministic. [[Server]] diff --git a/acceptance/experimental/air/help/output.txt b/acceptance/experimental/air/help/output.txt index ee89e778d6b..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/test.toml b/acceptance/experimental/air/list/test.toml index 10c2a7600b8..bcaa75fb34a 100644 --- a/acceptance/experimental/air/list/test.toml +++ b/acceptance/experimental/air/list/test.toml @@ -1,6 +1,3 @@ -# This command does not deploy a bundle, so no engine matrix is needed. -EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["direct"] - # Disable the on-disk run cache so --all-status output is deterministic across runs. [Env] DATABRICKS_CACHE_ENABLED = "false" diff --git a/acceptance/experimental/air/unimplemented/out.test.toml b/acceptance/experimental/air/logs-mlflow-fallback/out.test.toml similarity index 100% rename from acceptance/experimental/air/unimplemented/out.test.toml rename to acceptance/experimental/air/logs-mlflow-fallback/out.test.toml diff --git a/acceptance/experimental/air/logs-mlflow-fallback/output.txt b/acceptance/experimental/air/logs-mlflow-fallback/output.txt new file mode 100644 index 00000000000..874ad273314 --- /dev/null +++ b/acceptance/experimental/air/logs-mlflow-fallback/output.txt @@ -0,0 +1,8 @@ + +=== logs falls back to mlflow (no logs) +>>> [CLI] experimental air logs 123 +No logs available for run 123. Run terminated in state SUCCESS + +=== logs falls back to mlflow (json) +>>> [CLI] experimental air logs 123 -o json +{"type":"ERROR","ts":"[TIMESTAMP]","node":0,"line":"No logs available for run 123. Run terminated in state SUCCESS"} diff --git a/acceptance/experimental/air/logs-mlflow-fallback/script b/acceptance/experimental/air/logs-mlflow-fallback/script new file mode 100644 index 00000000000..ad8856a036a --- /dev/null +++ b/acceptance/experimental/air/logs-mlflow-fallback/script @@ -0,0 +1,9 @@ +# Bricklens is gated off (FEATURE_DISABLED), so the command falls back to the +# MLflow log path. With no MLflow run id resolvable, the fallback reports no +# logs and exits non-zero — proving the try/catch routes to MLflow. + +title "logs falls back to mlflow (no logs)" +errcode trace $CLI experimental air logs 123 + +title "logs falls back to mlflow (json)" +errcode trace $CLI experimental air logs 123 -o json diff --git a/acceptance/experimental/air/logs-mlflow-fallback/test.toml b/acceptance/experimental/air/logs-mlflow-fallback/test.toml new file mode 100644 index 00000000000..9b9bcab8e68 --- /dev/null +++ b/acceptance/experimental/air/logs-mlflow-fallback/test.toml @@ -0,0 +1,32 @@ +# The SDK occasionally probes host reachability with a HEAD request; stub it so +# the test is deterministic. +[[Server]] +Pattern = "HEAD /" +Response.Body = '' + +# A completed run. +[[Server]] +Pattern = "GET /api/2.2/jobs/runs/get" +Response.Body = ''' +{ + "run_id": 123, + "start_time": 1700000000000, + "end_time": 1700000012000, + "state": {"life_cycle_state": "TERMINATED", "result_state": "SUCCESS"}, + "tasks": [{"task_key": "train", "run_id": 456, "attempt_number": 0}] +} +''' + +# Bricklens is gated off by the backend SAFE flag, forcing the MLflow fallback. +[[Server]] +Pattern = "GET /api/2.0/ai-training/workflows/by-run-id/123/logs" +Response.StatusCode = 403 +Response.Body = ''' +{"error_code": "FEATURE_DISABLED", "message": "training log streaming is not enabled"} +''' + +# The MLflow fallback has no run output to resolve an MLflow run id from, so it +# reports no logs rather than failing — exercising the fallback wiring end to end. +[[Server]] +Pattern = "GET /api/2.2/jobs/runs/get-output" +Response.Body = '{}' diff --git a/acceptance/experimental/air/logs/out.test.toml b/acceptance/experimental/air/logs/out.test.toml new file mode 100644 index 00000000000..e90b6d5d1ba --- /dev/null +++ b/acceptance/experimental/air/logs/out.test.toml @@ -0,0 +1,3 @@ +Local = true +Cloud = false +EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["direct"] diff --git a/acceptance/experimental/air/logs/output.txt b/acceptance/experimental/air/logs/output.txt new file mode 100644 index 00000000000..84d6202afdf --- /dev/null +++ b/acceptance/experimental/air/logs/output.txt @@ -0,0 +1,72 @@ + +=== logs (text, completed run) +>>> [CLI] experimental air logs 123 +step 1 +step 2 +CUDA out of memory + +=== logs (json) +>>> [CLI] experimental air logs 123 -o json +{"type":"LOG","ts":"[TIMESTAMP]","node":0,"line":"step 1"} +{"type":"LOG","ts":"[TIMESTAMP]","node":0,"line":"step 2"} +{"type":"ALERT","ts":"[TIMESTAMP]","node":0,"line":"CUDA out of memory"} +{"type":"LOG","ts":"[TIMESTAMP]","node":0,"line":"CUDA out of memory"} + +=== logs with --minutes +>>> [CLI] experimental air logs 123 --minutes 30 +step 1 +step 2 +CUDA out of memory + +=== logs with --lines +>>> [CLI] experimental air logs 123 --lines 1 +CUDA out of memory + +=== logs with --lines 0 prints nothing +>>> [CLI] experimental air logs 123 --lines 0 +No logs available for run 123. Run terminated in state SUCCESS + +=== logs from a specific retry +>>> [CLI] experimental air logs 123 --retry 0 +step 1 +step 2 +CUDA out of memory + +=== logs --lines and --minutes are mutually exclusive +>>> [CLI] experimental air logs 123 --lines 100 --minutes 30 +Error: cannot combine --lines with --minutes: --lines tails by line count, --minutes by time window + +Exit code: 1 + +=== logs --lines and --minutes are mutually exclusive (json) +>>> [CLI] experimental air logs 123 --lines 100 --minutes 30 -o json +{ + "v": 1, + "ts": "[TIMESTAMP]", + "error": { + "code": "INVALID_ARGS", + "kind": "PERMANENT", + "message": "cannot combine --lines with --minutes: --lines tails by line count, --minutes by time window", + "retryable": false + } +} + +Exit code: 1 + +=== invalid run id +>>> [CLI] experimental air logs notanumber +Error: invalid JOB_RUN_ID "notanumber": must be a positive integer + +Exit code: 1 + +=== negative node +>>> [CLI] experimental air logs 123 --node -1 +Error: invalid --node -1: must not be negative + +Exit code: 1 + +=== --download-to not implemented +>>> [CLI] experimental air logs 123 --download-to /tmp/out +Error: --download-to is not implemented yet + +Exit code: 1 diff --git a/acceptance/experimental/air/logs/script b/acceptance/experimental/air/logs/script new file mode 100644 index 00000000000..b851a8784f8 --- /dev/null +++ b/acceptance/experimental/air/logs/script @@ -0,0 +1,32 @@ +title "logs (text, completed run)" +trace $CLI experimental air logs 123 + +title "logs (json)" +trace $CLI experimental air logs 123 -o json + +title "logs with --minutes" +trace $CLI experimental air logs 123 --minutes 30 + +title "logs with --lines" +trace $CLI experimental air logs 123 --lines 1 + +title "logs with --lines 0 prints nothing" +trace $CLI experimental air logs 123 --lines 0 + +title "logs from a specific retry" +trace $CLI experimental air logs 123 --retry 0 + +title "logs --lines and --minutes are mutually exclusive" +errcode trace $CLI experimental air logs 123 --lines 100 --minutes 30 + +title "logs --lines and --minutes are mutually exclusive (json)" +errcode trace $CLI experimental air logs 123 --lines 100 --minutes 30 -o json + +title "invalid run id" +errcode trace $CLI experimental air logs notanumber + +title "negative node" +errcode trace $CLI experimental air logs 123 --node -1 + +title "--download-to not implemented" +errcode trace $CLI experimental air logs 123 --download-to /tmp/out diff --git a/acceptance/experimental/air/logs/test.toml b/acceptance/experimental/air/logs/test.toml new file mode 100644 index 00000000000..31e95adc8c1 --- /dev/null +++ b/acceptance/experimental/air/logs/test.toml @@ -0,0 +1,30 @@ +# The SDK occasionally probes host reachability with a HEAD request; stub it so +# the test is deterministic. +[[Server]] +Pattern = "HEAD /" +Response.Body = '' + +# A completed run: Bricklens serves its logs via the tail drain. +[[Server]] +Pattern = "GET /api/2.2/jobs/runs/get" +Response.Body = ''' +{ + "run_id": 123, + "start_time": 1700000000000, + "end_time": 1700000012000, + "state": {"life_cycle_state": "TERMINATED", "result_state": "SUCCESS"}, + "tasks": [{"task_key": "train", "run_id": 456, "attempt_number": 0}] +} +''' + +# Bricklens log records, returned newest-first (as the tail fetch requests); +# printed oldest-first. +[[Server]] +Pattern = "GET /api/2.0/ai-training/workflows/by-run-id/123/logs" +Response.Body = ''' +{"log_records": [ + {"time_unix_nano": 1700000003000000000, "body": "CUDA out of memory", "node_index": 0}, + {"time_unix_nano": 1700000002000000000, "body": "step 2", "node_index": 0}, + {"time_unix_nano": 1700000001000000000, "body": "step 1", "node_index": 0} +]} +''' diff --git a/acceptance/experimental/air/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..0d9562f4bfe --- /dev/null +++ b/acceptance/experimental/air/run-submit-deps/.gitattributes @@ -0,0 +1,6 @@ +# These YAML files' contents are uploaded verbatim (as training_config.yaml and +# requirements.yaml), so their line endings must stay \n on every OS — a Windows +# \r would change the recorded payload. +run.yaml text eol=lf +run-file.yaml text eol=lf +reqs.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..98c7dc1dfa9 --- /dev/null +++ b/acceptance/experimental/air/run-submit-deps/output.txt @@ -0,0 +1,112 @@ + +=== submit with inline dependencies +>>> [CLI] experimental air run -f run.yaml +Submitted run 555 +View 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" + } + ] + } +} + +=== file-form deps: version comes from the requirements file +>>> [CLI] experimental air run -f run-file.yaml +Submitted run 555 +View at: [DATABRICKS_URL]/jobs/runs/555 + +Tip: use --watch to stream logs until the run completes. + +=== file-form deps: the requirements file is not uploaded either +>>> 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-file-smoke/deps-file-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-file-smoke/deps-file-smoke_[RUN_ID]/training_config.yaml", "q": {"overwrite": "true"}, "raw_body": "experiment_name: deps-file-smoke\ncommand: python train.py\ncompute:\n accelerator_type: GPU_1xH100\n num_accelerators: 1\nenvironment:\n dependencies: ./reqs.yaml\n"} + +=== file-form 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-file-smoke", + "tasks": [ + { + "ai_runtime_task": { + "deployments": [ + { + "command_path": "/Workspace/Users/[USERNAME]/.air/cli_launch/deps-file-smoke/deps-file-smoke_[RUN_ID]/command.sh", + "compute": { + "accelerator_count": 1, + "accelerator_type": "GPU_1xH100" + } + } + ], + "experiment": "deps-file-smoke" + }, + "environment_key": "default", + "max_retries": 3, + "retry_on_timeout": true, + "run_if": "ALL_SUCCESS", + "task_key": "deps-file-smoke" + } + ] + } +} diff --git a/acceptance/experimental/air/run-submit-deps/reqs.yaml b/acceptance/experimental/air/run-submit-deps/reqs.yaml new file mode 100644 index 00000000000..da4b89534c2 --- /dev/null +++ b/acceptance/experimental/air/run-submit-deps/reqs.yaml @@ -0,0 +1,4 @@ +version: 5 +dependencies: + - numpy + - torch==2.3.0 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..52035bbae65 --- /dev/null +++ b/acceptance/experimental/air/run-submit-deps/script @@ -0,0 +1,17 @@ +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 "file-form deps: version comes from the requirements file" +trace $CLI experimental air run -f run-file.yaml + +title "file-form deps: the requirements file is not uploaded either" +trace print_requests.py //api/2.0/workspace-files/import-file --oneline --sort --unique --keep + +title "file-form deps ride on environments[].spec.dependencies" +trace print_requests.py //api/2.2/jobs/runs/submit 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..590c2918a25 --- /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.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]' + +[[Repls]] +Old = 'deps-file-smoke_[0-9a-f]{16}' +New = 'deps-file-smoke_[RUN_ID]' diff --git a/acceptance/experimental/air/run-submit/output.txt b/acceptance/experimental/air/run-submit/output.txt index 8d52eed1dab..2d92122c8e6 100644 --- a/acceptance/experimental/air/run-submit/output.txt +++ b/acceptance/experimental/air/run-submit/output.txt @@ -1,9 +1,12 @@ === submit with a git code_source >>> [CLI] experimental air run -f run.yaml +Uploading [SNAPSHOT_TARBALL]... Submitted run 555 View at: [DATABRICKS_URL]/jobs/runs/555 +Tip: use --watch to stream logs until the run completes. + === the ai_runtime_task carries the code_source_path >>> print_requests.py //api/2.2/jobs/runs/submit { @@ -23,7 +26,7 @@ View at: [DATABRICKS_URL]/jobs/runs/555 "tasks": [ { "ai_runtime_task": { - "code_source_path": "/Workspace/Users/[USERNAME]/.air/repo_snapshots/001/[SNAPSHOT_TARBALL]", + "code_source_path": "/Workspace/Users/[USERNAME]/.air/repo_snapshots/.internal/[SNAPSHOT_TARBALL]", "deployments": [ { "command_path": "/Workspace/Users/[USERNAME]/.air/cli_launch/submit-smoke/submit-smoke_[RUN_ID]/command.sh", diff --git a/acceptance/experimental/air/run-submit/test.toml b/acceptance/experimental/air/run-submit/test.toml index a077e95bf99..3dc9ff81b99 100644 --- a/acceptance/experimental/air/run-submit/test.toml +++ b/acceptance/experimental/air/run-submit/test.toml @@ -1,14 +1,11 @@ # A real (non-dry-run) submit that packages a git code_source, uploads the -# tarball + provenance sidecars, and POSTs runs/submit. No bundle deploy, so no -# engine matrix. +# tarball + provenance sidecars, and POSTs runs/submit. RecordRequests = true # run.yaml is generated from run.yaml.tmpl at test time (commit SHA templated in); # it isn't a committed input to diff. Ignore = ["run.yaml"] -EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["direct"] - # The SDK probes host reachability with a HEAD request; stub it for determinism. [[Server]] Pattern = "HEAD /" diff --git a/acceptance/experimental/air/run/output.txt b/acceptance/experimental/air/run/output.txt index a753eabd198..13c9b360ece 100644 --- a/acceptance/experimental/air/run/output.txt +++ b/acceptance/experimental/air/run/output.txt @@ -14,18 +14,29 @@ Dry run: configuration for "smoke-test" is valid; not submitting. } } -=== override not yet supported ->>> [CLI] experimental air run -f valid.yaml --dry-run --override a=b -Error: --override is not yet supported +=== override applies and logs the change +>>> [CLI] experimental air run -f valid.yaml --dry-run --override compute.num_accelerators=2 --override timeout_minutes=45 +Override: changing compute.num_accelerators from 1 to 2 +Override: setting timeout_minutes to 45 +Dry run: configuration for "smoke-test" is valid; not submitting. + +=== override of an unknown field is rejected +>>> [CLI] experimental air run -f valid.yaml --dry-run --override bogus=1 +Error: invalid --override "bogus": "bogus" is not a known field; available fields are: code_source, command, compute, env_variables, environment, experiment_name, idempotency_token, max_retries, mlflow_experiment_directory, mlflow_run_name, parameters, permissions, secrets, timeout_minutes, usage_policy_id, usage_policy_name Exit code: 1 -=== watch not yet supported ->>> [CLI] experimental air run -f valid.yaml --dry-run --watch -Error: --watch is not yet supported +=== override still runs schema validation +>>> [CLI] experimental air run -f valid.yaml --dry-run --override compute.num_accelerators=0 +Override: changing compute.num_accelerators from 1 to 0 +Error: compute.num_accelerators must be positive, got 0 Exit code: 1 +=== watch is ignored with dry-run (nothing is submitted) +>>> [CLI] experimental air run -f valid.yaml --dry-run --watch +Dry run: configuration for "smoke-test" is valid; not submitting. + === code_source config passes validation >>> [CLI] experimental air run -f with-code-source.yaml --dry-run Dry run: configuration for "smoke-test" is valid; not submitting. diff --git a/acceptance/experimental/air/run/script b/acceptance/experimental/air/run/script index 312b2f6fecf..2a5263bceb2 100644 --- a/acceptance/experimental/air/run/script +++ b/acceptance/experimental/air/run/script @@ -4,11 +4,17 @@ trace $CLI experimental air run -f valid.yaml --dry-run title "dry-run (json)" trace $CLI experimental air run -f valid.yaml --dry-run -o json -title "override not yet supported" -errcode trace $CLI experimental air run -f valid.yaml --dry-run --override a=b +title "override applies and logs the change" +trace $CLI experimental air run -f valid.yaml --dry-run --override compute.num_accelerators=2 --override timeout_minutes=45 -title "watch not yet supported" -errcode trace $CLI experimental air run -f valid.yaml --dry-run --watch +title "override of an unknown field is rejected" +errcode trace $CLI experimental air run -f valid.yaml --dry-run --override bogus=1 + +title "override still runs schema validation" +errcode trace $CLI experimental air run -f valid.yaml --dry-run --override compute.num_accelerators=0 + +title "watch is ignored with dry-run (nothing is submitted)" +trace $CLI experimental air run -f valid.yaml --dry-run --watch title "code_source config passes validation" trace $CLI experimental air run -f with-code-source.yaml --dry-run diff --git a/acceptance/experimental/air/run/test.toml b/acceptance/experimental/air/run/test.toml deleted file mode 100644 index c228ad415d2..00000000000 --- a/acceptance/experimental/air/run/test.toml +++ /dev/null @@ -1,3 +0,0 @@ -# `air run --dry-run` validates the config locally and makes no workspace calls, -# so no engine matrix or server stubs are needed. -EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["direct"] diff --git a/acceptance/experimental/air/test.toml b/acceptance/experimental/air/test.toml new file mode 100644 index 00000000000..167f8e1a6bf --- /dev/null +++ b/acceptance/experimental/air/test.toml @@ -0,0 +1,5 @@ +# No air command deploys a bundle, so the engine matrix adds no coverage; pin a +# single engine ("direct") so the air tests run once, not across the direct+terraform +# matrix. (An empty list is rejected: it would run on both runners.) +[EnvMatrix] +DATABRICKS_BUNDLE_ENGINE = ["direct"] diff --git a/acceptance/experimental/air/unimplemented/output.txt b/acceptance/experimental/air/unimplemented/output.txt 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/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..8e46c70d48c 100644 --- a/experimental/air/cmd/format.go +++ b/experimental/air/cmd/format.go @@ -12,6 +12,7 @@ import ( "github.com/databricks/cli/libs/cmdio" "github.com/databricks/databricks-sdk-go/service/jobs" + "github.com/muesli/termenv" "go.yaml.in/yaml/v3" ) @@ -26,12 +27,6 @@ func orNA(s string) string { return s } -// osc8Link wraps label in an OSC 8 terminal hyperlink to url. -// See https://gist.github.com/egmontkob/eb114294efbcd5adb1944c9f3cb5feda -func osc8Link(label, url string) string { - return "\x1b]8;;" + url + "\x1b\\" + label + "\x1b]8;;\x1b\\" -} - // hyperlink renders label as a terminal hyperlink to url when out is a rich // terminal, otherwise it returns label unchanged. This mirrors the Python CLI's // Rich link markup, which drops the URL on non-terminals (so piped or captured @@ -40,7 +35,7 @@ func hyperlink(ctx context.Context, out io.Writer, label, url string) string { if url == "" || !cmdio.SupportsColor(ctx, out) { return label } - return osc8Link(label, url) + return termenv.Hyperlink(url, label) } // reformatYAMLForDisplay re-renders a training-config YAML so multi-line strings diff --git a/experimental/air/cmd/format_test.go b/experimental/air/cmd/format_test.go index 62a6d7ac580..1063e20ca1f 100644 --- a/experimental/air/cmd/format_test.go +++ b/experimental/air/cmd/format_test.go @@ -21,10 +21,6 @@ func TestSubmittedDisplay(t *testing.T) { assert.Equal(t, "2023-11-14 22:13 UTC", submittedDisplay(&jobs.Run{StartTime: 1700000000000})) } -func TestOSC8Link(t *testing.T) { - assert.Equal(t, "\x1b]8;;https://h.test/x\x1b\\label\x1b]8;;\x1b\\", osc8Link("label", "https://h.test/x")) -} - func TestHyperlink(t *testing.T) { // On a non-terminal (no color), the URL is dropped and only the label shows. ctx := cmdio.MockDiscard(t.Context()) diff --git a/experimental/air/cmd/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/logbricklens.go b/experimental/air/cmd/logbricklens.go new file mode 100644 index 00000000000..48ab7286649 --- /dev/null +++ b/experimental/air/cmd/logbricklens.go @@ -0,0 +1,87 @@ +package aircmd + +import ( + "context" + "encoding/json" + "fmt" + "net/http" + "strconv" + + "github.com/databricks/databricks-sdk-go/client" +) + +// bricklensLogsPathFmt is the log endpoint, keyed by Jobs run id. Called with a +// raw client.Do because the SDK does not model the AiTrainingService. +const bricklensLogsPathFmt = "/api/2.0/ai-training/workflows/by-run-id/%d/logs" + +// logRecord is one log line from Bricklens. +type logRecord struct { + // TimeUnixNano may arrive as a JSON number or string. + TimeUnixNano json.Number `json:"time_unix_nano"` + Body string `json:"body"` + NodeIndex int `json:"node_index"` +} + +// nano returns time_unix_nano as an int64, or 0 when absent or unparseable. +func (r logRecord) nano() int64 { + n, err := r.TimeUnixNano.Int64() + if err != nil { + return 0 + } + return n +} + +type bricklensLogsResponse struct { + LogRecords []logRecord `json:"log_records"` + NextPageToken string `json:"next_page_token"` +} + +// bricklensLogsQuery is the request-field surface of the log endpoint. +type bricklensLogsQuery struct { + // fromSeconds and toSeconds bound the query window in Unix epoch seconds. + fromSeconds int64 + toSeconds int64 + pageToken string + pageSize int + // attemptNumber selects a retry attempt (0-indexed); -1 means latest. + attemptNumber int + nodeIndex int + // ascending returns oldest-first. The endpoint defaults to ascending when + // absent, so the tail fetch must send an explicit false for newest-first. + ascending bool +} + +// getBricklensLogs fetches one page of logs. The API client is built once by the +// caller and reused across the poll loop. It returns the raw error so the caller +// can classify it via classifyLogError. +func getBricklensLogs(ctx context.Context, apiClient *client.DatabricksClient, runID int64, q bricklensLogsQuery) (*bricklensLogsResponse, error) { + query := map[string]any{ + // Always sent: the tail path relies on an explicit false for newest-first. + "ascending": strconv.FormatBool(q.ascending), + } + if q.fromSeconds > 0 { + query["from"] = q.fromSeconds + } + if q.toSeconds > 0 { + query["to"] = q.toSeconds + } + if q.pageToken != "" { + query["page_token"] = q.pageToken + } + if q.pageSize > 0 { + query["page_size"] = q.pageSize + } + if q.attemptNumber >= 0 { + query["ref.attempt_number"] = q.attemptNumber + } + if q.nodeIndex >= 0 { + query["filter.node_index"] = q.nodeIndex + } + + var resp bricklensLogsResponse + path := fmt.Sprintf(bricklensLogsPathFmt, runID) + if err := apiClient.Do(ctx, http.MethodGet, path, nil, nil, query, &resp); err != nil { + return nil, err + } + return &resp, nil +} diff --git a/experimental/air/cmd/logbricklens_test.go b/experimental/air/cmd/logbricklens_test.go new file mode 100644 index 00000000000..e12e0f75a65 --- /dev/null +++ b/experimental/air/cmd/logbricklens_test.go @@ -0,0 +1,83 @@ +package aircmd + +import ( + "net/http" + "net/http/httptest" + "net/url" + "testing" + + "github.com/databricks/databricks-sdk-go/client" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestGetBricklensLogsQuerySerialization(t *testing.T) { + var got url.Values + var gotPath string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/oidc/.well-known/oauth-authorization-server" { + gotPath = r.URL.Path + got = r.URL.Query() + } + _, _ = w.Write([]byte(`{"log_records": [], "next_page_token": ""}`)) + })) + t.Cleanup(srv.Close) + + w := newTestWorkspaceClient(t, srv.URL) + apiClient, err := client.New(w.Config) + require.NoError(t, err) + _, err = getBricklensLogs(t.Context(), apiClient, 42, bricklensLogsQuery{ + fromSeconds: 100, + toSeconds: 200, + pageToken: "tok", + pageSize: 500, + attemptNumber: 1, + nodeIndex: 3, + ascending: true, + }) + require.NoError(t, err) + + assert.Equal(t, "/api/2.0/ai-training/workflows/by-run-id/42/logs", gotPath) + assert.Equal(t, "100", got.Get("from")) + assert.Equal(t, "200", got.Get("to")) + assert.Equal(t, "tok", got.Get("page_token")) + assert.Equal(t, "500", got.Get("page_size")) + assert.Equal(t, "1", got.Get("ref.attempt_number")) + assert.Equal(t, "3", got.Get("filter.node_index")) + assert.Equal(t, "true", got.Get("ascending")) +} + +func TestGetBricklensLogsOmitsOptionals(t *testing.T) { + var got url.Values + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/oidc/.well-known/oauth-authorization-server" { + got = r.URL.Query() + } + _, _ = w.Write([]byte(`{"log_records": []}`)) + })) + t.Cleanup(srv.Close) + + w := newTestWorkspaceClient(t, srv.URL) + apiClient, err := client.New(w.Config) + require.NoError(t, err) + // attempt -1 (latest) and node 0 are the default request; from/to/page unset. + _, err = getBricklensLogs(t.Context(), apiClient, 7, bricklensLogsQuery{attemptNumber: -1, nodeIndex: 0}) + require.NoError(t, err) + + assert.False(t, got.Has("from")) + assert.False(t, got.Has("to")) + assert.False(t, got.Has("page_token")) + assert.False(t, got.Has("page_size")) + // -1 attempt means "latest" — the field is omitted so the endpoint defaults. + assert.False(t, got.Has("ref.attempt_number")) + // node 0 is a real filter value and must be sent. + assert.Equal(t, "0", got.Get("filter.node_index")) + // ascending is always sent so the tail path can force newest-first. + assert.Equal(t, "false", got.Get("ascending")) +} + +func TestLogRecordNano(t *testing.T) { + assert.Equal(t, int64(123), logRecord{TimeUnixNano: "123"}.nano()) + assert.Equal(t, int64(0), logRecord{TimeUnixNano: ""}.nano()) + assert.Equal(t, int64(0), logRecord{TimeUnixNano: "notanumber"}.nano()) +} diff --git a/experimental/air/cmd/logdetect.go b/experimental/air/cmd/logdetect.go new file mode 100644 index 00000000000..f804f4d5812 --- /dev/null +++ b/experimental/air/cmd/logdetect.go @@ -0,0 +1,53 @@ +package aircmd + +import "regexp" + +// fatalPatterns match log lines that signal a run-ending failure (OOM, NCCL +// timeouts, CUDA errors, segfaults, etc.). In --json mode a matching line emits +// an ALERT event alongside its LOG event, giving an agent an immediate signal. +var fatalPatterns = []*regexp.Regexp{ + // OOM + regexp.MustCompile(`(?i)CUDA out of memory`), + regexp.MustCompile(`(?i)Out of memory: Kill(ed)? process`), + // Signals + regexp.MustCompile(`(?i)signal\s+(9|SIGKILL|SIGTERM)`), + // NCCL / collective + regexp.MustCompile(`Watchdog caught collective operation timeout`), + regexp.MustCompile(`(?i)NCCL WARN .*(Conn|Net|IB|timeout|unhandled)`), + regexp.MustCompile(`Got async error event`), + regexp.MustCompile(`transport/net_ib\.cc:\d+.*WARN`), + // CUDA + regexp.MustCompile(`(?i)CUDA(?: runtime)? error`), + regexp.MustCompile(`(?i)an illegal memory access was encountered`), + regexp.MustCompile(`CUDA kernel errors might be asynchronously reported`), + // Segfault + regexp.MustCompile(`(?i)segmentation fault`), + // Composer / llmfoundry specific + regexp.MustCompile(`composer\.utils\..*Error`), + regexp.MustCompile(`(?i)composer.*OutOfMemory`), + // torch.distributed + regexp.MustCompile(`torch\.distributed\..*(?:Error|Exception)`), + regexp.MustCompile(`(?i)(?:TCPStore|Store).*timed?\s*out`), + // GPU hardware (Xid from dmesg / driver) + regexp.MustCompile(`Xid.*\b(48|63|64|79|94|95)\b`), + // Streaming dataset + regexp.MustCompile(`streaming\.base\..*(?:Error|Exception)`), + // Bare "Killed" on its own line means OOM-killer or similar + regexp.MustCompile(`^\s*Killed\s*$`), + // MLflow stall — training has stopped logging metrics + regexp.MustCompile(`\[MLflow Logger\]\[Warning\] No new logs have been emitted`), + // The launch script prints this when the user's command fails; [1-9]\d* skips exit code 0. + regexp.MustCompile(`ERROR: Script failed with exit code [1-9]\d* after \d+s`), + // Missing command (exit 127), e.g. a typo. + regexp.MustCompile(`(?i)command not found`), +} + +// matchFatalPattern reports whether a log line matches a fatal-failure pattern. +func matchFatalPattern(line string) bool { + for _, p := range fatalPatterns { + if p.MatchString(line) { + return true + } + } + return false +} diff --git a/experimental/air/cmd/logmlflow.go b/experimental/air/cmd/logmlflow.go new file mode 100644 index 00000000000..a8dd016dbbc --- /dev/null +++ b/experimental/air/cmd/logmlflow.go @@ -0,0 +1,302 @@ +package aircmd + +import ( + "bufio" + "context" + "fmt" + "io" + "net" + "net/http" + "os" + "path" + "regexp" + "slices" + "strconv" + "time" + + "github.com/databricks/cli/libs/log" + "github.com/databricks/databricks-sdk-go" + "github.com/databricks/databricks-sdk-go/client" + "github.com/databricks/databricks-sdk-go/listing" + "github.com/databricks/databricks-sdk-go/service/jobs" + "github.com/databricks/databricks-sdk-go/service/ml" +) + +// chunkFilePattern matches a log chunk file (logs-.chunk.txt); group 1 is +// the chunk index. The sidecar splits stdout into 4MB chunks, index ascending. +var chunkFilePattern = regexp.MustCompile(`^logs-(\d+)\.chunk\.txt$`) + +// oldFormatNodeDir matches a bare per-node log dir (logs/node_). The +// attempt-prefixed layout nests these under logs/attempt_/, so a bare +// logs/node_ only appears in the old layout. +var oldFormatNodeDir = regexp.MustCompile(`^logs/node_\d+$`) + +// artifactDownloadClient fetches pre-signed artifact URLs with connect and +// response-header timeouts, so a stalled storage backend can't hang the command. +// Mirrors the Python CLI's (10s connect, 60s read) bounds. +var artifactDownloadClient = &http.Client{ + Transport: &http.Transport{ + DialContext: (&net.Dialer{Timeout: 10 * time.Second}).DialContext, + ResponseHeaderTimeout: 60 * time.Second, + }, +} + +// mlflowLogFallback prints a run's logs from MLflow artifacts, the fallback when +// Bricklens can't serve them. It resolves the MLflow run id, discovers the +// per-node log directory, lists the chunk files, and walks them newest-first +// until it has the requested tail, then prints oldest-first. +// +// The tail length is --lines, else the default cap. MLflow chunks are not +// time-indexed, so --minutes cannot restrict the window here. +func mlflowLogFallback(ctx context.Context, w *databricks.WorkspaceClient, out io.Writer, req logRequest, status logRunStatus) (bool, error) { + if req.windowMinutes > 0 { + log.Debugf(ctx, "air logs: --minutes is not supported on the MLflow fallback path; showing the default tail") + } + + mlflowRunID, logDir, err := resolveMLflowLogPath(ctx, w, req) + if err != nil { + return false, err + } + if mlflowRunID == "" || logDir == "" { + emitNoLogs(out, req, status) + return status.succeeded(), nil + } + + chunks, err := listLogChunks(ctx, w, mlflowRunID, logDir) + if err != nil { + return false, err + } + if len(chunks) == 0 { + // Nothing listed yet: assume the single chunk 0. + chunks = []logChunk{{index: 0, path: path.Join(logDir, chunkFileName(0))}} + } + + target := req.tailTarget() + if target <= 0 { + return status.succeeded(), nil + } + + lines, err := tailChunks(ctx, w, mlflowRunID, chunks, target) + if err != nil { + return false, err + } + if len(lines) == 0 { + emitNoLogs(out, req, status) + return status.succeeded(), nil + } + + if len(lines) > target { + lines = lines[len(lines)-target:] + } + for _, line := range lines { + emitLogLine(out, req, line) + } + return status.succeeded(), nil +} + +// resolveMLflowLogPath returns the run's MLflow run id and per-node log directory. +func resolveMLflowLogPath(ctx context.Context, w *databricks.WorkspaceClient, req logRequest) (string, string, error) { + run, err := w.Jobs.GetRun(ctx, jobs.GetRunRequest{RunId: req.runID}) + if err != nil { + return "", "", err + } + ids := mlflowIDs(ctx, w, run) + if ids == nil || ids.RunID == "" { + return "", "", nil + } + + // -1 (latest) maps to attempt 0's directory. + attempt := max(req.attempt, 0) + withAttempt, err := discoverAttemptPrefix(ctx, w, ids.RunID, attempt) + if err != nil { + return "", "", err + } + return ids.RunID, constructLogPath(req.node, attempt, withAttempt), nil +} + +// discoverAttemptPrefix probes the logs/ dir once to decide whether the layout is +// attempt-prefixed (logs/attempt_X/node_Y) or old (logs/node_Y). A bare +// logs/node_ means old; a logs/attempt_ entry means prefixed. +// Defaults to old when nothing is listed. +func discoverAttemptPrefix(ctx context.Context, w *databricks.WorkspaceClient, mlflowRunID string, attempt int) (bool, error) { + files, err := listArtifacts(ctx, w, mlflowRunID, "logs") + if err != nil { + // Not fatal: default to the old layout; the chunk listing finds it empty if wrong. + log.Debugf(ctx, "air logs: could not list logs dir for format discovery: %v", err) + return false, nil + } + + attemptDir := fmt.Sprintf("logs/attempt_%d", attempt) + for _, f := range files { + if oldFormatNodeDir.MatchString(f.Path) { + return false, nil + } + if f.Path == attemptDir { + return true, nil + } + } + return false, nil +} + +// constructLogPath builds the per-node log directory for a node and attempt. +func constructLogPath(node, attempt int, withAttempt bool) string { + if withAttempt { + return fmt.Sprintf("logs/attempt_%d/node_%d", attempt, node) + } + return fmt.Sprintf("logs/node_%d", node) +} + +// chunkFileName is the artifact filename for a chunk index. +func chunkFileName(index int) string { + return fmt.Sprintf("logs-%d.chunk.txt", index) +} + +// logChunk is one listed chunk: its index and full artifact path. +type logChunk struct { + index int + path string +} + +// listLogChunks lists the chunk files under a log dir, sorted ascending by index. +func listLogChunks(ctx context.Context, w *databricks.WorkspaceClient, mlflowRunID, logDir string) ([]logChunk, error) { + files, err := listArtifacts(ctx, w, mlflowRunID, logDir) + if err != nil { + return nil, err + } + + var chunks []logChunk + for _, f := range files { + base := path.Base(f.Path) + m := chunkFilePattern.FindStringSubmatch(base) + if m == nil { + continue + } + idx, err := strconv.Atoi(m[1]) + if err != nil { + continue + } + chunks = append(chunks, logChunk{index: idx, path: f.Path}) + } + slices.SortFunc(chunks, func(a, b logChunk) int { return a.index - b.index }) + return chunks, nil +} + +// tailChunks walks chunks newest-first, prepending each chunk's lines, until it +// has `target` lines or runs out. A mid-walk download failure stops the walk +// rather than splice non-adjacent chunks. +func tailChunks(ctx context.Context, w *databricks.WorkspaceClient, mlflowRunID string, chunks []logChunk, target int) ([]string, error) { + var accumulated []string + for _, chunk := range slices.Backward(chunks) { + lines, err := downloadChunkLines(ctx, w, mlflowRunID, chunk.path) + if err != nil { + log.Debugf(ctx, "air logs: failed to download chunk %d; showing only logs after it: %v", chunk.index, err) + break + } + accumulated = append(lines, accumulated...) + if len(accumulated) >= target { + break + } + } + return accumulated, nil +} + +// downloadChunkLines fetches one chunk artifact and returns its lines. +func downloadChunkLines(ctx context.Context, w *databricks.WorkspaceClient, mlflowRunID, artifactPath string) ([]string, error) { + f, err := downloadArtifact(ctx, w, mlflowRunID, artifactPath) + if err != nil { + return nil, err + } + defer os.Remove(f) + + file, err := os.Open(f) + if err != nil { + return nil, err + } + defer file.Close() + + var lines []string + scanner := bufio.NewScanner(file) + scanner.Buffer(make([]byte, 0, 64*1024), 4*1024*1024) + for scanner.Scan() { + lines = append(lines, scanner.Text()) + } + return lines, scanner.Err() +} + +// listArtifacts lists a run's artifacts under a path. +func listArtifacts(ctx context.Context, w *databricks.WorkspaceClient, mlflowRunID, path string) ([]ml.FileInfo, error) { + it := w.Experiments.ListArtifacts(ctx, ml.ListArtifactsRequest{RunId: mlflowRunID, Path: path}) + return listing.ToSlice(ctx, it) +} + +// credentialInfo is one credentials-for-read entry: a pre-signed URL plus any +// backend-required request headers. +type credentialInfo struct { + SignedURI string `json:"signed_uri"` + Headers []struct { + Name string `json:"name"` + Value string `json:"value"` + } `json:"headers"` +} + +type credentialsForReadResponse struct { + CredentialInfos []credentialInfo `json:"credential_infos"` +} + +// downloadArtifact downloads one run artifact to a temp file and returns its +// path. credentials-for-read returns a pre-signed URL, which we stream to disk; +// that endpoint is not modeled by the SDK, so it is called via a raw client.Do. +func downloadArtifact(ctx context.Context, w *databricks.WorkspaceClient, mlflowRunID, artifactPath string) (string, error) { + apiClient, err := client.New(w.Config) + if err != nil { + return "", fmt.Errorf("failed to create API client: %w", err) + } + + var resp credentialsForReadResponse + query := map[string]any{ + "run_id": mlflowRunID, + // path is a repeated field, so pass a slice (serialized as path=...&path=...). + "path": []string{artifactPath}, + } + err = apiClient.Do(ctx, http.MethodGet, "/api/2.0/mlflow/artifacts/credentials-for-read", nil, nil, query, &resp) + if err != nil { + return "", fmt.Errorf("failed to get read credentials for %s: %w", artifactPath, err) + } + if len(resp.CredentialInfos) == 0 || resp.CredentialInfos[0].SignedURI == "" { + return "", fmt.Errorf("no download credentials returned for %s", artifactPath) + } + cred := resp.CredentialInfos[0] + + req, err := http.NewRequestWithContext(ctx, http.MethodGet, cred.SignedURI, nil) + if err != nil { + return "", err + } + // Azure SAS / some GCS URIs require backend-supplied headers; AWS returns none. + for _, h := range cred.Headers { + req.Header.Set(h.Name, h.Value) + } + + httpResp, err := artifactDownloadClient.Do(req) + if err != nil { + return "", err + } + defer httpResp.Body.Close() + if httpResp.StatusCode >= 400 { + return "", fmt.Errorf("artifact download failed: HTTP %d", httpResp.StatusCode) + } + + tmp, err := os.CreateTemp("", "air-log-chunk-*") + if err != nil { + return "", err + } + if _, err := io.Copy(tmp, httpResp.Body); err != nil { + tmp.Close() + os.Remove(tmp.Name()) + return "", err + } + if err := tmp.Close(); err != nil { + os.Remove(tmp.Name()) + return "", err + } + return tmp.Name(), nil +} diff --git a/experimental/air/cmd/logmlflow_test.go b/experimental/air/cmd/logmlflow_test.go new file mode 100644 index 00000000000..f3de80abdb6 --- /dev/null +++ b/experimental/air/cmd/logmlflow_test.go @@ -0,0 +1,114 @@ +package aircmd + +import ( + "bytes" + "net/http" + "net/http/httptest" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestConstructLogPath(t *testing.T) { + assert.Equal(t, "logs/node_0", constructLogPath(0, 0, false)) + assert.Equal(t, "logs/node_3", constructLogPath(3, 2, false)) + assert.Equal(t, "logs/attempt_2/node_3", constructLogPath(3, 2, true)) +} + +func TestChunkFileName(t *testing.T) { + assert.Equal(t, "logs-0.chunk.txt", chunkFileName(0)) + assert.Equal(t, "logs-7.chunk.txt", chunkFileName(7)) +} + +func TestChunkFilePattern(t *testing.T) { + m := chunkFilePattern.FindStringSubmatch("logs-12.chunk.txt") + require.NotNil(t, m) + assert.Equal(t, "12", m[1]) + + assert.Nil(t, chunkFilePattern.FindStringSubmatch("logs-12.chunk.txt.bak")) + assert.Nil(t, chunkFilePattern.FindStringSubmatch("node_0")) +} + +// artifactListServer serves a fixed artifacts/list response for any path. +func artifactListServer(t *testing.T, body string) *httptest.Server { + t.Helper() + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == "/api/2.0/mlflow/artifacts/list" { + _, _ = w.Write([]byte(body)) + return + } + _, _ = w.Write([]byte(`{}`)) + })) + t.Cleanup(srv.Close) + return srv +} + +func TestListLogChunksSortsAndFiltersByIndex(t *testing.T) { + // Out-of-order chunks plus a non-chunk file; result is ascending, chunk-only. + srv := artifactListServer(t, `{"files": [ + {"path": "logs/node_0/logs-2.chunk.txt"}, + {"path": "logs/node_0/other.txt"}, + {"path": "logs/node_0/logs-0.chunk.txt"}, + {"path": "logs/node_0/logs-1.chunk.txt"} + ]}`) + w := newTestWorkspaceClient(t, srv.URL) + + chunks, err := listLogChunks(t.Context(), w, "run1", "logs/node_0") + require.NoError(t, err) + require.Len(t, chunks, 3) + assert.Equal(t, 0, chunks[0].index) + assert.Equal(t, 1, chunks[1].index) + assert.Equal(t, 2, chunks[2].index) + assert.Equal(t, "logs/node_0/logs-0.chunk.txt", chunks[0].path) +} + +func TestDiscoverAttemptPrefix(t *testing.T) { + // Old format: a bare logs/node_N dir means no attempt prefix. + old := artifactListServer(t, `{"files": [{"path": "logs/node_0", "is_dir": true}]}`) + got, err := discoverAttemptPrefix(t.Context(), newTestWorkspaceClient(t, old.URL), "run1", 0) + require.NoError(t, err) + assert.False(t, got) + + // New format: a logs/attempt_N entry and no bare node dir. + newFmt := artifactListServer(t, `{"files": [{"path": "logs/attempt_0", "is_dir": true}]}`) + got, err = discoverAttemptPrefix(t.Context(), newTestWorkspaceClient(t, newFmt.URL), "run1", 0) + require.NoError(t, err) + assert.True(t, got) +} + +// noMLflowServer serves a run with no resolvable MLflow run id, so the fallback +// finds no logs. +func noMLflowServer(t *testing.T) *httptest.Server { + t.Helper() + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/api/2.2/jobs/runs/get": + _, _ = w.Write([]byte(`{"run_id": 5, "state": {"life_cycle_state": "TERMINATED", "result_state": "SUCCESS"}, "tasks": [{"run_id": 456}]}`)) + case "/api/2.2/jobs/runs/get-output": + _, _ = w.Write([]byte(`{}`)) + default: + _, _ = w.Write([]byte(`{}`)) + } + })) + t.Cleanup(srv.Close) + return srv +} + +func TestMLflowFallbackNoLogsReflectsRunOutcome(t *testing.T) { + srv := noMLflowServer(t) + w := newTestWorkspaceClient(t, srv.URL) + + // A SUCCESS run with no logs still reports success (exit 0), matching the + // Bricklens path, rather than failing just because no logs exist. + success, err := mlflowLogFallback(t.Context(), w, &bytes.Buffer{}, + logRequest{runID: 5}, logRunStatus{lifeCycleState: "TERMINATED", resultState: "SUCCESS"}) + require.NoError(t, err) + assert.True(t, success) + + // A FAILED run with no logs reports failure (exit 1). + success, err = mlflowLogFallback(t.Context(), w, &bytes.Buffer{}, + logRequest{runID: 5}, logRunStatus{lifeCycleState: "TERMINATED", resultState: "FAILED"}) + require.NoError(t, err) + assert.False(t, success) +} diff --git a/experimental/air/cmd/logs.go b/experimental/air/cmd/logs.go index c34fb62a7df..1bda80b572b 100644 --- a/experimental/air/cmd/logs.go +++ b/experimental/air/cmd/logs.go @@ -1,7 +1,18 @@ package aircmd import ( + "context" + "errors" + "fmt" + "io" + "strconv" + "github.com/databricks/cli/cmd/root" + "github.com/databricks/cli/libs/cmdctx" + "github.com/databricks/cli/libs/flags" + "github.com/databricks/databricks-sdk-go" + "github.com/databricks/databricks-sdk-go/apierr" + "github.com/databricks/databricks-sdk-go/service/iam" "github.com/spf13/cobra" ) @@ -9,6 +20,7 @@ func newLogsCommand() *cobra.Command { var ( node int lines int + minutes int retry int downloadTo string review bool @@ -19,18 +31,142 @@ func newLogsCommand() *cobra.Command { Args: root.ExactArgs(1), Short: "Stream or fetch logs for a run", Long: `Stream logs from an active run, or fetch logs from a completed run.`, - RunE: func(cmd *cobra.Command, args []string) error { - return notImplemented("logs") - }, } cmd.Flags().IntVar(&node, "node", 0, "Fetch logs from this node") - cmd.Flags().IntVar(&lines, "lines", 10000, "For completed runs, print the last N lines") + cmd.Flags().IntVar(&lines, "lines", 0, "For completed runs, print the last N lines (default 10000)") + cmd.Flags().IntVar(&minutes, "minutes", 0, "Fetch only logs from the last N minutes") cmd.Flags().IntVar(&retry, "retry", -1, "View logs from a specific retry attempt; -1 means latest") cmd.Flags().StringVar(&downloadTo, "download-to", "", "Download all logs to this directory instead of printing") cmd.Flags().BoolVar(&review, "review", false, "Download logs from all nodes and filter for error signatures") - // Hidden in the Python `air` CLI (help=argparse.SUPPRESS); keep it internal here to match. cmd.Flags().MarkHidden("review") + // In -o json mode an auth failure should be a JSON error envelope, not a bare + // error. ErrAlreadyPrinted passes through. + cmd.PreRunE = func(cmd *cobra.Command, args []string) error { + err := root.MustWorkspaceClient(cmd, args) + if err == nil || errors.Is(err, root.ErrAlreadyPrinted) { + return err + } + return authError(cmd.Context(), cmd, err) + } + + cmd.RunE = func(cmd *cobra.Command, args []string) error { + ctx := cmd.Context() + + // --download-to and --review are not yet implemented; reject rather than + // silently ignore. + if downloadTo != "" { + return renderError(ctx, cmd, "INVALID_ARGS", "PERMANENT", false, + errors.New("--download-to is not implemented yet")) + } + if review { + return renderError(ctx, cmd, "INVALID_ARGS", "PERMANENT", false, + errors.New("--review is not implemented yet")) + } + + // --lines (line tail) and --minutes (time window) answer the same question + // two ways, so reject both together rather than silently honoring one. + if lines > 0 && minutes > 0 { + return renderError(ctx, cmd, "INVALID_ARGS", "PERMANENT", false, + errors.New("cannot combine --lines with --minutes: --lines tails by line count, --minutes by time window")) + } + if lines < 0 { + return renderError(ctx, cmd, "INVALID_ARGS", "PERMANENT", false, + fmt.Errorf("invalid --lines %d: must be positive", lines)) + } + if minutes < 0 { + return renderError(ctx, cmd, "INVALID_ARGS", "PERMANENT", false, + fmt.Errorf("invalid --minutes %d: must be positive", minutes)) + } + if node < 0 { + return renderError(ctx, cmd, "INVALID_ARGS", "PERMANENT", false, + fmt.Errorf("invalid --node %d: must not be negative", node)) + } + + runID, err := strconv.ParseInt(args[0], 10, 64) + if err != nil || runID <= 0 { + return renderError(ctx, cmd, "INVALID_ARGS", "PERMANENT", false, + fmt.Errorf("invalid JOB_RUN_ID %q: must be a positive integer", args[0])) + } + + // -1 signals "unset" (use the default cap); an explicit --lines 0 stays 0 + // and prints nothing. + tailLines := -1 + if cmd.Flags().Changed("lines") { + tailLines = lines + } + + return runLogs(ctx, cmd, logRequest{ + runID: runID, + node: node, + attempt: retry, + windowMinutes: minutes, + tailLines: tailLines, + jsonOutput: root.OutputType(cmd) == flags.OutputJSON, + }) + } + return cmd } + +// runLogs resolves the run, validates --retry, and fetches logs. It handles error +// reporting; the backend selection lives in fetchLogs. +func runLogs(ctx context.Context, cmd *cobra.Command, req logRequest) error { + w := cmdctx.WorkspaceClient(ctx) + + // Validate credentials server-side before fetching (MustWorkspaceClient only + // attaches them), so a bad token fails clearly here. + if _, err := w.CurrentUser.Me(ctx, iam.MeRequest{}); err != nil { + return authError(ctx, cmd, err) + } + + status, err := resolveRunStatus(ctx, w, req.runID) + if err != nil { + if errors.Is(err, apierr.ErrResourceDoesNotExist) { + return renderError(ctx, cmd, "NOT_FOUND", "NOT_FOUND", false, + fmt.Errorf("run %d not found: check the run ID and that it is a job run ID", req.runID)) + } + return renderError(ctx, cmd, "INTERNAL_ERROR", "TRANSIENT", true, + fmt.Errorf("failed to get status for run %d: %w", req.runID, err)) + } + + // -1 (default) means latest; reject an attempt past the newest. + if req.attempt >= 0 && req.attempt > status.latestAttempt { + return renderError(ctx, cmd, "INVALID_ARGS", "PERMANENT", false, + fmt.Errorf("invalid retry %d: available retries are 0 to %d", req.attempt, status.latestAttempt)) + } + + // A past retry of an active run has immutable logs: render once, don't follow. + if req.attempt >= 0 && req.attempt < status.latestAttempt && !status.terminal() { + req.staticView = true + } + + out := cmd.OutOrStdout() + success, err := fetchLogs(ctx, w, out, req, status) + if err != nil { + if errors.Is(err, apierr.ErrResourceDoesNotExist) { + return renderError(ctx, cmd, "NOT_FOUND", "NOT_FOUND", false, + fmt.Errorf("run %d not found: check the run ID and that it is a job run ID", req.runID)) + } + return renderError(ctx, cmd, "INTERNAL_ERROR", "TRANSIENT", true, + fmt.Errorf("failed to fetch logs for run %d: %w", req.runID, err)) + } + + // A run that finished unsuccessfully exits non-zero; output was already + // written, so don't reprint via Cobra. + if !success { + return root.ErrAlreadyPrinted + } + return nil +} + +// fetchLogs serves logs from Bricklens, falling back to MLflow when Bricklens +// returns errBricklensFeatureDisabled. +func fetchLogs(ctx context.Context, w *databricks.WorkspaceClient, out io.Writer, req logRequest, status logRunStatus) (bool, error) { + success, err := streamBricklensLogs(ctx, w, out, req, status) + if errors.Is(err, errBricklensFeatureDisabled) { + return mlflowLogFallback(ctx, w, out, req, status) + } + return success, err +} diff --git a/experimental/air/cmd/logs_test.go b/experimental/air/cmd/logs_test.go new file mode 100644 index 00000000000..a9eb24b72f3 --- /dev/null +++ b/experimental/air/cmd/logs_test.go @@ -0,0 +1,252 @@ +package aircmd + +import ( + "bytes" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/databricks/cli/cmd/root" + "github.com/databricks/cli/libs/cmdctx" + "github.com/databricks/cli/libs/cmdio" + "github.com/databricks/cli/libs/flags" + "github.com/databricks/databricks-sdk-go/experimental/mocks" + "github.com/spf13/cobra" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestLogsCommandShape(t *testing.T) { + cmd := newLogsCommand() + assert.Equal(t, "logs JOB_RUN_ID", cmd.Use) + assert.Empty(t, cmd.Commands(), "logs must not register subcommands") + assert.NoError(t, cmd.Args(cmd, []string{"123"})) + assert.Error(t, cmd.Args(cmd, []string{})) + assert.Error(t, cmd.Args(cmd, []string{"1", "2"})) + + // --review is hidden. + review := cmd.Flags().Lookup("review") + require.NotNil(t, review) + assert.True(t, review.Hidden) +} + +// runLogsCmd invokes the logs command's RunE with the given flags against a mock +// (no-HTTP) workspace client. Used for input validation that fails before any +// API call. +func runLogsCmd(t *testing.T, args []string, flagsToSet map[string]string) error { + t.Helper() + m := mocks.NewMockWorkspaceClient(t) + ctx := cmdctx.SetWorkspaceClient(cmdio.MockDiscard(t.Context()), m.WorkspaceClient) + cmd := withOutput(newLogsCommand(), flags.OutputText) + for k, v := range flagsToSet { + require.NoError(t, cmd.Flags().Set(k, v)) + } + cmd.SetContext(ctx) + return cmd.RunE(cmd, args) +} + +func TestLogsFlagValidation(t *testing.T) { + tests := []struct { + name string + args []string + flags map[string]string + wantMsg string + }{ + { + name: "lines and minutes are mutually exclusive", + args: []string{"5"}, + flags: map[string]string{"lines": "100", "minutes": "10"}, + wantMsg: "cannot combine --lines with --minutes", + }, + { + name: "negative lines rejected", + args: []string{"5"}, + flags: map[string]string{"lines": "-1"}, + wantMsg: "invalid --lines", + }, + { + name: "negative minutes rejected", + args: []string{"5"}, + flags: map[string]string{"minutes": "-1"}, + wantMsg: "invalid --minutes", + }, + { + name: "download-to not implemented", + args: []string{"5"}, + flags: map[string]string{"download-to": "/tmp/logs"}, + wantMsg: "--download-to is not implemented yet", + }, + { + name: "review not implemented", + args: []string{"5"}, + flags: map[string]string{"review": "true"}, + wantMsg: "--review is not implemented yet", + }, + { + name: "invalid run id", + args: []string{"abc"}, + wantMsg: "invalid JOB_RUN_ID", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := runLogsCmd(t, tt.args, tt.flags) + require.Error(t, err) + assert.Contains(t, err.Error(), tt.wantMsg) + }) + } +} + +// completedRunLogsServer serves the auth probe, a terminal runs/get, and a +// single page of Bricklens logs (newest-first, as the tail fetch requests). +func completedRunLogsServer(t *testing.T) *httptest.Server { + t.Helper() + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch { + case r.URL.Path == "/api/2.2/jobs/runs/get": + _, _ = w.Write([]byte(`{ + "run_id": 5, + "start_time": 1000, + "end_time": 2000, + "state": {"life_cycle_state": "TERMINATED", "result_state": "SUCCESS"}, + "tasks": [{"attempt_number": 0}] + }`)) + case strings.HasPrefix(r.URL.Path, "/api/2.0/ai-training/workflows/by-run-id/"): + // Newest-first, as drainTail requests; reversed to oldest-first on print. + _, _ = w.Write([]byte(`{"log_records": [ + {"time_unix_nano": 2000000000, "body": "line two", "node_index": 0}, + {"time_unix_nano": 1000000000, "body": "line one", "node_index": 0} + ]}`)) + default: + // Me() probe and SDK config discovery. + _, _ = w.Write([]byte(`{"userName":"u@example.com"}`)) + } + })) + t.Cleanup(srv.Close) + return srv +} + +func TestLogsCompletedRunTail(t *testing.T) { + srv := completedRunLogsServer(t) + var buf bytes.Buffer + ctx := cmdctx.SetWorkspaceClient(cmdio.MockDiscard(t.Context()), newTestWorkspaceClient(t, srv.URL)) + cmd := withOutput(&cobra.Command{}, flags.OutputText) + cmd.SetContext(ctx) + cmd.SetOut(&buf) + + // Drive runLogs directly (bypassing PreRunE auth wiring) with a resolved request. + err := runLogs(ctx, cmd, logRequest{runID: 5, node: 0, attempt: -1, tailLines: -1}) + require.NoError(t, err) + + // Records print oldest-first regardless of the newest-first fetch order. + assert.Equal(t, "line one\nline two\n", buf.String()) +} + +// mlflowFallbackServer serves a terminal run whose Bricklens endpoint is gated +// off (FEATURE_DISABLED), plus the full MLflow artifact path the fallback walks: +// runs/get-output (MLflow ids), artifacts/list (logs dir + chunk file), +// credentials-for-read (pre-signed URL), and the pre-signed chunk bytes itself. +func mlflowFallbackServer(t *testing.T) *httptest.Server { + t.Helper() + var base string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch { + case r.URL.Path == "/api/2.2/jobs/runs/get": + _, _ = w.Write([]byte(`{ + "run_id": 5, + "state": {"life_cycle_state": "TERMINATED", "result_state": "SUCCESS"}, + "tasks": [{"run_id": 456, "attempt_number": 0}] + }`)) + case strings.HasPrefix(r.URL.Path, "/api/2.0/ai-training/workflows/by-run-id/"): + w.WriteHeader(http.StatusForbidden) + _, _ = w.Write([]byte(`{"error_code": "FEATURE_DISABLED", "message": "bricklens logs gated off"}`)) + case r.URL.Path == "/api/2.2/jobs/runs/get-output": + _, _ = w.Write([]byte(`{"ai_runtime_task_output": {"mlflow_experiment_id": "exp1", "mlflow_run_id": "run1"}}`)) + case r.URL.Path == "/api/2.0/mlflow/artifacts/list": + // The logs dir probe (format discovery) and the per-node chunk listing + // both hit this; return the old-format node dir and one chunk file. + if r.URL.Query().Get("path") == "logs" { + _, _ = w.Write([]byte(`{"files": [{"path": "logs/node_0", "is_dir": true}]}`)) + return + } + _, _ = w.Write([]byte(`{"files": [{"path": "logs/node_0/logs-0.chunk.txt", "file_size": 12}]}`)) + case r.URL.Path == "/api/2.0/mlflow/artifacts/credentials-for-read": + _, _ = w.Write([]byte(`{"credential_infos": [{"signed_uri": "` + base + `/presigned"}]}`)) + case r.URL.Path == "/presigned": + _, _ = w.Write([]byte("line one\nline two\n")) + default: + _, _ = w.Write([]byte(`{"userName":"u@example.com"}`)) + } + })) + base = srv.URL + t.Cleanup(srv.Close) + return srv +} + +func TestLogsFallsBackToMLflow(t *testing.T) { + srv := mlflowFallbackServer(t) + var buf bytes.Buffer + ctx := cmdctx.SetWorkspaceClient(cmdio.MockDiscard(t.Context()), newTestWorkspaceClient(t, srv.URL)) + cmd := withOutput(&cobra.Command{}, flags.OutputText) + cmd.SetContext(ctx) + cmd.SetOut(&buf) + + // Bricklens is gated off, so fetchLogs routes to the MLflow fallback, which + // resolves the MLflow run, lists the chunk, downloads it via the pre-signed + // URL, and prints its lines. + err := runLogs(ctx, cmd, logRequest{runID: 5, node: 0, attempt: -1, tailLines: -1}) + require.NoError(t, err) + assert.Equal(t, "line one\nline two\n", buf.String()) +} + +// activeRunPastRetryServer serves a still-RUNNING run with two attempts and a +// single page of Bricklens logs. runs/get always returns RUNNING; a test that +// follows the run would poll forever, so it also asserts the static path never +// loops. +func activeRunPastRetryServer(t *testing.T, getRunHits *int) *httptest.Server { + t.Helper() + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch { + case r.URL.Path == "/api/2.2/jobs/runs/get": + *getRunHits++ + _, _ = w.Write([]byte(`{ + "run_id": 9, + "start_time": 1700000000000, + "state": {"life_cycle_state": "RUNNING"}, + "tasks": [{"attempt_number": 0}, {"attempt_number": 1}] + }`)) + case strings.HasPrefix(r.URL.Path, "/api/2.0/ai-training/workflows/by-run-id/"): + _, _ = w.Write([]byte(`{"log_records": [ + {"time_unix_nano": 1700000001000000000, "body": "retry 0 log", "node_index": 0} + ]}`)) + default: + _, _ = w.Write([]byte(`{"userName":"u@example.com"}`)) + } + })) + t.Cleanup(srv.Close) + return srv +} + +func TestLogsPastRetryOfActiveRunIsStatic(t *testing.T) { + var getRunHits int + srv := activeRunPastRetryServer(t, &getRunHits) + var buf bytes.Buffer + ctx := cmdctx.SetWorkspaceClient(cmdio.MockDiscard(t.Context()), newTestWorkspaceClient(t, srv.URL)) + cmd := withOutput(&cobra.Command{}, flags.OutputText) + cmd.SetContext(ctx) + cmd.SetOut(&buf) + + // --retry 0 on a RUNNING run whose latest attempt is 1: the past attempt's + // logs render once instead of following the run (which would never terminate). + // The run has no SUCCESS result yet, so it still exits non-zero via + // ErrAlreadyPrinted; the logs are printed regardless. + err := runLogs(ctx, cmd, logRequest{runID: 9, node: 0, attempt: 0, tailLines: -1}) + require.ErrorIs(t, err, root.ErrAlreadyPrinted) + assert.Equal(t, "retry 0 log\n", buf.String()) + + // Exactly one runs/get: the initial status resolve in runLogs. The static + // streamer must not re-poll (that is the loop this test guards against). + assert.Equal(t, 1, getRunHits) +} diff --git a/experimental/air/cmd/logstream.go b/experimental/air/cmd/logstream.go new file mode 100644 index 00000000000..e5ab7ba8e85 --- /dev/null +++ b/experimental/air/cmd/logstream.go @@ -0,0 +1,483 @@ +package aircmd + +import ( + "context" + "errors" + "fmt" + "io" + "net/http" + "slices" + "time" + + "github.com/databricks/cli/libs/cmdio" + "github.com/databricks/cli/libs/log" + "github.com/databricks/databricks-sdk-go" + "github.com/databricks/databricks-sdk-go/apierr" + "github.com/databricks/databricks-sdk-go/client" + "github.com/databricks/databricks-sdk-go/service/jobs" +) + +const ( + // maxTransientFailures is how many consecutive Bricklens failures to tolerate + // before falling back to MLflow. + maxTransientFailures = 5 + // defaultCompletedRunTailLines caps a completed run's output when neither + // --lines nor --minutes is set. + defaultCompletedRunTailLines = 10000 + // seenRecordsCap bounds the dedup set, evicting oldest-inserted entries first. + seenRecordsCap = 100000 +) + +// retryCheckInterval is the wait between status/log polls. A var so tests can +// shrink it. +var retryCheckInterval = 3 * time.Second + +// errBricklensFeatureDisabled signals the caller to fall back to MLflow: Bricklens +// is gated off (FEATURE_DISABLED), not deployed (ENDPOINT_NOT_FOUND / 404), or +// persistently failing. The flag is evaluated server-side. +var errBricklensFeatureDisabled = errors.New("bricklens logs unavailable; falling back to mlflow") + +// logRequest describes what to fetch, shared by both backends so they honor the +// same flags. windowMinutes and tailLines are mutually exclusive. +type logRequest struct { + runID int64 + // node is the node index to fetch; node 0 always exists. + node int + // attempt is the retry attempt to read; -1 means latest. + attempt int + // windowMinutes, when > 0, restricts the fetch to the last N minutes. + windowMinutes int + // tailLines caps a completed run's output to the last N lines. Negative means + // --lines was unset (use the default cap); 0 prints nothing. + tailLines int + // staticView renders a one-shot tail instead of following the run. Set for a + // past retry of an active run: that attempt's logs are immutable, so streaming + // would poll forever waiting for the run (not the attempt) to finish. + staticView bool + jsonOutput bool + // onStatusChange, when set, is called on each lifecycle transition while + // following the run (current, previous display states). Used by + // `air run --watch -o json` to emit STATUS events. + onStatusChange func(current, previous string) +} + +// logRunStatus is the subset of a run's state the log path needs, resolved once +// and reused. +type logRunStatus struct { + lifeCycleState string + resultState string + stateMessage string + startTimeMs int64 + endTimeMs int64 + // latestAttempt is the highest attempt_number across the run's tasks. + latestAttempt int +} + +// A run is terminal when its lifecycle state is terminal, or a result state is +// set (result states only appear on terminal runs). +var ( + terminalLifeCycleStates = map[string]bool{"TERMINATED": true, "SKIPPED": true, "INTERNAL_ERROR": true} + terminalResultStates = map[string]bool{"SUCCESS": true, "FAILED": true, "CANCELED": true} +) + +func (s logRunStatus) terminal() bool { + return terminalLifeCycleStates[s.lifeCycleState] || terminalResultStates[s.resultState] +} + +func (s logRunStatus) succeeded() bool { + return s.resultState == "SUCCESS" +} + +// resolveRunStatus fetches a run's state and projects it onto logRunStatus. An +// unknown run id surfaces as apierr.ErrResourceDoesNotExist. +func resolveRunStatus(ctx context.Context, w *databricks.WorkspaceClient, runID int64) (logRunStatus, error) { + run, err := w.Jobs.GetRun(ctx, jobs.GetRunRequest{RunId: runID}) + if err != nil { + return logRunStatus{}, err + } + return projectRunStatus(run), nil +} + +// projectRunStatus extracts logRunStatus from a run. Split out so it can be +// tested without an API client. +func projectRunStatus(run *jobs.Run) logRunStatus { + s := logRunStatus{ + startTimeMs: run.StartTime, + endTimeMs: run.EndTime, + } + if run.State != nil { + s.lifeCycleState = string(run.State.LifeCycleState) + s.resultState = string(run.State.ResultState) + s.stateMessage = run.State.StateMessage + } + for i := range run.Tasks { + s.latestAttempt = max(s.latestAttempt, run.Tasks[i].AttemptNumber) + } + return s +} + +// classifyLogError maps a Bricklens failure to one of: +// - errBricklensFeatureDisabled: fall back to MLflow (gated off, endpoint +// absent, or 404). +// - the original error: a genuine not-found, surfaced as-is. +// - nil: a transient failure the caller should retry. +func classifyLogError(err error) error { + if apiErr, ok := errors.AsType[*apierr.APIError](err); ok { + switch apiErr.ErrorCode { + case "FEATURE_DISABLED", "ENDPOINT_NOT_FOUND": + return errBricklensFeatureDisabled + } + if apiErr.StatusCode == http.StatusNotFound { + return errBricklensFeatureDisabled + } + } + if errors.Is(err, apierr.ErrResourceDoesNotExist) { + return err + } + return nil +} + +// fromSeconds computes the `from` bound. With --minutes set it is now-N*60; +// otherwise the run's start second (0 before the run starts, which the endpoint +// reads as "everything stored"). +func (req logRequest) fromSeconds(status logRunStatus, now time.Time) int64 { + if req.windowMinutes > 0 { + return now.Add(-time.Duration(req.windowMinutes) * time.Minute).Unix() + } + if status.startTimeMs > 0 { + return status.startTimeMs / 1000 + } + return 0 +} + +// toSeconds computes the `to` bound. A terminal run caps at its end second (ceil +// of the millisecond time, so the final partial second is kept); otherwise 0 lets +// the endpoint default to now. +func (req logRequest) toSeconds(status logRunStatus) int64 { + if status.terminal() && status.endTimeMs > 0 { + return (status.endTimeMs + 999) / 1000 + } + return 0 +} + +// streamBricklensLogs fetches and prints a run's logs: a bounded tail for a +// completed run, or a poll-and-drain loop that follows an active run to +// completion. It returns whether the run finished with SUCCESS; +// errBricklensFeatureDisabled means the caller should fall back to MLflow. +func streamBricklensLogs(ctx context.Context, w *databricks.WorkspaceClient, out io.Writer, req logRequest, status logRunStatus) (bool, error) { + // Build the API client once and reuse it for every page fetch in the loop. + apiClient, err := client.New(w.Config) + if err != nil { + return false, fmt.Errorf("failed to create API client: %w", err) + } + st := &bricklensStreamer{ + ctx: ctx, + w: w, + apiClient: apiClient, + out: out, + req: req, + status: status, + seen: newSeenSet(seenRecordsCap), + } + return st.run() +} + +// bricklensStreamer holds the poll-loop state: the from-second cursor, the +// highest emitted timestamp, and the dedup set. +type bricklensStreamer struct { + ctx context.Context + w *databricks.WorkspaceClient + apiClient *client.DatabricksClient + out io.Writer + req logRequest + status logRunStatus + + fromSec int64 + lastNano int64 + firstLogSeen bool + seen *seenSet + // previousState is the last display state reported to onStatusChange. + previousState string + // onFirstLog, when set, is called once just before the first log line is + // emitted — used to stop the "waiting for run to start" spinner before any + // log byte reaches stdout. + onFirstLog func() + // updateSpinner, when set, refreshes the waiting-spinner text each poll. + updateSpinner func(string) +} + +// reportStatusChange fires onStatusChange when the run's display state differs +// from the last reported one. +func (st *bricklensStreamer) reportStatusChange() { + if st.req.onStatusChange == nil { + return + } + current := st.status.displayState() + if current == st.previousState { + return + } + st.req.onStatusChange(current, st.previousState) + st.previousState = current +} + +func (st *bricklensStreamer) run() (bool, error) { + now := time.Now() + st.fromSec = st.req.fromSeconds(st.status, now) + + // A past retry's logs are immutable: render a one-shot tail rather than + // following the still-active run, which would poll forever. + if st.req.staticView { + return st.drainStatic(st.req.toSeconds(st.status)) + } + + // Show a "waiting for run to start" spinner on stderr while the run has not + // yet produced logs, so a provisioning run doesn't look hung. Suppressed in + // --json mode and auto-degraded to nothing on a non-interactive terminal. + // The first emitted log line stops it via onFirstLog (before any stdout write). + if !st.req.jsonOutput { + sp := cmdio.NewSpinner(st.ctx) + defer sp.Close() + st.onFirstLog = sp.Close + st.updateSpinner = sp.Update + } + + firstIteration := true + for { + if !firstIteration { + status, err := resolveRunStatus(st.ctx, st.w, st.req.runID) + if err != nil { + if errors.Is(err, apierr.ErrResourceDoesNotExist) { + return false, err + } + // A cancelled context (Ctrl-C) is not a transient blip: stop + // promptly instead of retrying forever. + if st.ctx.Err() != nil { + return false, st.ctx.Err() + } + // A transient status blip should not abort a live stream. + log.Debugf(st.ctx, "air logs: failed to refresh run status: %v", err) + if err := sleepOrCancel(st.ctx, retryCheckInterval); err != nil { + return false, err + } + continue + } + st.status = status + } + + st.reportStatusChange() + + terminal := st.status.terminal() + toSec := st.req.toSeconds(st.status) + + // While waiting on a still-active run with no logs yet, refresh the spinner. + if !terminal && !st.firstLogSeen && st.updateSpinner != nil { + st.updateSpinner(fmt.Sprintf("Waiting for run to start (node %d)...", st.req.node)) + } + + // A run already terminal on the first iteration renders as a tail (most + // recent N lines). An active run streams everything with dedup, so a run + // that terminates while we watch doesn't re-print the boundary second. + var err error + if firstIteration && terminal { + err = st.drainTail(toSec) + } else { + err = st.drainPages(toSec) + } + if err != nil { + return false, err + } + + if terminal { + if !st.firstLogSeen { + // Stop the spinner before the no-logs line so frames don't smear. + if st.onFirstLog != nil { + st.onFirstLog() + } + st.emitNoLogs() + } + log.Infof(st.ctx, "air logs: run %d finished in state %s", st.req.runID, st.status.displayState()) + return st.status.succeeded(), nil + } + + firstIteration = false + if err := sleepOrCancel(st.ctx, retryCheckInterval); err != nil { + return false, err + } + } +} + +// sleepOrCancel waits for d, or returns early with the context error if the +// context is cancelled (e.g. Ctrl-C) so the poll loop exits promptly. +func sleepOrCancel(ctx context.Context, d time.Duration) error { + timer := time.NewTimer(d) + defer timer.Stop() + select { + case <-ctx.Done(): + return ctx.Err() + case <-timer.C: + return nil + } +} + +// drainStatic renders a single tail pass without following the run. Success +// reflects the run's current result state (empty while active). +func (st *bricklensStreamer) drainStatic(toSec int64) (bool, error) { + if err := st.drainTail(toSec); err != nil { + return false, err + } + if !st.firstLogSeen { + st.emitNoLogs() + } + return st.status.succeeded(), nil +} + +// tailTarget is the number of lines a tail keeps. A negative tailLines means +// --lines was unset, so use the default cap; 0 or more is taken literally (an +// explicit --lines 0 prints nothing). +func (req logRequest) tailTarget() int { + if req.tailLines < 0 { + return defaultCompletedRunTailLines + } + return req.tailLines +} + +// drainTail emits the most-recent `target` records oldest-first. Bricklens +// returns records newest-first, so it pages until it has `target`, keeps the +// newest `target`, and reverses to chronological order. +func (st *bricklensStreamer) drainTail(toSec int64) error { + target := st.req.tailTarget() + if target <= 0 { + return nil + } + + var collected []logRecord + var pageToken string + for len(collected) < target { + resp, err := st.requestPage(pageToken, toSec, target, false) + if err != nil { + return err + } + collected = append(collected, resp.LogRecords...) + pageToken = resp.NextPageToken + if pageToken == "" { + break + } + } + + // Keep the newest `target`, then reverse to print oldest -> newest. + if len(collected) > target { + collected = collected[:target] + } + for _, c := range slices.Backward(collected) { + st.emit(c.Body) + } + return nil +} + +// drainPages exhausts all pages from the current from-second in ascending order, +// deduping against the seen-set so a re-queried boundary second is not +// re-printed, then advances fromSec to the newest record's floor-second. +func (st *bricklensStreamer) drainPages(toSec int64) error { + var pageToken string + for { + resp, err := st.requestPage(pageToken, toSec, 0, true) + if err != nil { + return err + } + + for _, rec := range resp.LogRecords { + nano := rec.nano() + if nano != 0 { + // Skip a record older than the last emitted one to keep output + // monotonic (out of order, or a re-queried boundary record). + if st.lastNano != 0 && nano < st.lastNano { + continue + } + if st.seen.has(nano, rec.Body) { + continue + } + } + st.emit(rec.Body) + if nano != 0 { + st.seen.add(nano, rec.Body) + st.lastNano = max(st.lastNano, nano) + } + } + + pageToken = resp.NextPageToken + if pageToken == "" { + break + } + } + + if st.lastNano != 0 { + st.fromSec = st.lastNano / 1_000_000_000 + } + return nil +} + +// requestPage fetches one page, retrying transient failures up to +// maxTransientFailures before falling back to MLflow. A feature-gated response or +// genuine not-found returns immediately. +func (st *bricklensStreamer) requestPage(pageToken string, toSec int64, pageSize int, ascending bool) (*bricklensLogsResponse, error) { + q := bricklensLogsQuery{ + fromSeconds: st.fromSec, + toSeconds: toSec, + pageToken: pageToken, + pageSize: pageSize, + attemptNumber: st.req.attempt, + nodeIndex: st.req.node, + ascending: ascending, + } + + transientFailures := 0 + for { + resp, err := getBricklensLogs(st.ctx, st.apiClient, st.req.runID, q) + if err == nil { + return resp, nil + } + + switch classified := classifyLogError(err); { + case errors.Is(classified, errBricklensFeatureDisabled): + return nil, errBricklensFeatureDisabled + case classified != nil: + return nil, classified + } + + transientFailures++ + if transientFailures >= maxTransientFailures { + log.Debugf(st.ctx, "air logs: bricklens failed %d times; falling back to mlflow", maxTransientFailures) + return nil, errBricklensFeatureDisabled + } + log.Debugf(st.ctx, "air logs: bricklens transient failure (%d/%d): %v", transientFailures, maxTransientFailures, err) + if err := sleepOrCancel(st.ctx, retryCheckInterval); err != nil { + return nil, err + } + } +} + +// emit writes one log line and latches firstLogSeen so an empty terminal run can +// report "no logs". The first line stops the waiting spinner before any byte +// reaches stdout. +func (st *bricklensStreamer) emit(body string) { + if !st.firstLogSeen && st.onFirstLog != nil { + st.onFirstLog() + } + st.firstLogSeen = true + emitLogLine(st.out, st.req, body) +} + +func (st *bricklensStreamer) emitNoLogs() { + emitNoLogs(st.out, st.req, st.status) +} + +// displayState is the result state, else the lifecycle state, else "UNKNOWN". +func (s logRunStatus) displayState() string { + if s.resultState != "" { + return s.resultState + } + if s.lifeCycleState != "" { + return s.lifeCycleState + } + return "UNKNOWN" +} diff --git a/experimental/air/cmd/logstream_support.go b/experimental/air/cmd/logstream_support.go new file mode 100644 index 00000000000..ac250547e16 --- /dev/null +++ b/experimental/air/cmd/logstream_support.go @@ -0,0 +1,179 @@ +package aircmd + +import ( + "container/list" + "encoding/json" + "fmt" + "io" + "time" +) + +// seenNano keys the dedup set. Distinct lines can share a nano (each rank stamps +// from its own clock), so the body disambiguates them. +type seenNano struct { + nano int64 + body string +} + +// seenSet is an insertion-ordered set bounded to a capacity, evicting the +// oldest-inserted entry first. +type seenSet struct { + cap int + items map[seenNano]*list.Element + order *list.List +} + +func newSeenSet(capacity int) *seenSet { + return &seenSet{ + cap: capacity, + items: make(map[seenNano]*list.Element), + order: list.New(), + } +} + +func (s *seenSet) has(nano int64, body string) bool { + _, ok := s.items[seenNano{nano, body}] + return ok +} + +func (s *seenSet) add(nano int64, body string) { + key := seenNano{nano, body} + if _, ok := s.items[key]; ok { + return + } + s.items[key] = s.order.PushBack(key) + if s.order.Len() > s.cap { + oldest := s.order.Front() + s.order.Remove(oldest) + delete(s.items, oldest.Value.(seenNano)) + } +} + +// logEvent is one JSONL streaming event. +type logEvent struct { + Type string `json:"type"` + TS string `json:"ts"` + Node int `json:"node"` + Line string `json:"line"` +} + +// printLogEvent writes a single JSONL event line for --json streaming output. +func printLogEvent(out io.Writer, eventType string, node int, line string) { + b, err := json.Marshal(logEvent{ + Type: eventType, + TS: time.Now().UTC().Format(time.RFC3339), + Node: node, + Line: line, + }) + if err != nil { + return + } + fmt.Fprintln(out, string(b)) +} + +// submittedEvent is the JSONL event `air run --watch -o json` emits before the +// streamed log events, so a consumer sees the run id immediately. +type submittedEvent struct { + Type string `json:"type"` + TS string `json:"ts"` + RunID string `json:"run_id"` + DashboardURL string `json:"dashboard_url"` +} + +// printSubmittedEvent writes the SUBMITTED JSONL event. +func printSubmittedEvent(out io.Writer, runID, dashboardURL string) { + b, err := json.Marshal(submittedEvent{ + Type: "SUBMITTED", + TS: time.Now().UTC().Format(time.RFC3339), + RunID: runID, + DashboardURL: dashboardURL, + }) + if err != nil { + return + } + fmt.Fprintln(out, string(b)) +} + +// statusEvent is a JSONL event emitted on each lifecycle transition while +// following a run with --watch. +type statusEvent struct { + Type string `json:"type"` + TS string `json:"ts"` + Status string `json:"status"` + Previous string `json:"previous_status,omitempty"` +} + +// printStatusEvent writes a STATUS JSONL event for a lifecycle transition. +func printStatusEvent(out io.Writer, current, previous string) { + b, err := json.Marshal(statusEvent{ + Type: "STATUS", + TS: time.Now().UTC().Format(time.RFC3339), + Status: current, + Previous: previous, + }) + if err != nil { + return + } + fmt.Fprintln(out, string(b)) +} + +// terminalEvent is the closing envelope `air run --watch -o json` emits after +// streaming, carrying the run's terminal status. +type terminalEvent struct { + V int `json:"v"` + TS string `json:"ts"` + Data runResult `json:"data"` +} + +// printTerminalEvent writes the closing terminal-status envelope, matching the +// shape of renderEnvelope(runResult). +func printTerminalEvent(out io.Writer, runID, status, dashboardURL string) { + b, err := json.Marshal(terminalEvent{ + V: envelopeVersion, + TS: time.Now().UTC().Format(time.RFC3339), + Data: runResult{ + Status: status, + RunID: runID, + DashboardURL: dashboardURL, + }, + }) + if err != nil { + return + } + fmt.Fprintln(out, string(b)) +} + +// emitLogLine writes one log line: raw in text mode, or a JSONL LOG event under +// --json. In --json mode a line matching a fatal-failure pattern also emits an +// ALERT event first, giving an agent an immediate actionable signal. +func emitLogLine(out io.Writer, req logRequest, body string) { + if !req.jsonOutput { + fmt.Fprintln(out, body) + return + } + if matchFatalPattern(body) { + printLogEvent(out, "ALERT", req.node, body) + } + printLogEvent(out, "LOG", req.node, body) +} + +// emitNoLogs reports that a run produced no logs. A terminal run gets its +// termination reason; a still-active run is reported as having no logs yet, +// since the MLflow fallback is a one-shot that does not follow it to completion. +// Under --json it is a JSONL ERROR, so a consumer never sees an empty stream. +func emitNoLogs(out io.Writer, req logRequest, status logRunStatus) { + var msg string + if status.terminal() { + msg = fmt.Sprintf("No logs available for run %d. Run terminated in state %s", req.runID, status.displayState()) + } else { + msg = fmt.Sprintf("No logs available yet for run %d, which is still in state %s", req.runID, status.displayState()) + } + if status.stateMessage != "" { + msg = fmt.Sprintf("%s: %s", msg, status.stateMessage) + } + if req.jsonOutput { + printLogEvent(out, "ERROR", req.node, msg) + return + } + fmt.Fprintln(out, msg) +} diff --git a/experimental/air/cmd/logstream_test.go b/experimental/air/cmd/logstream_test.go new file mode 100644 index 00000000000..d8a061687a6 --- /dev/null +++ b/experimental/air/cmd/logstream_test.go @@ -0,0 +1,403 @@ +package aircmd + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/databricks/databricks-sdk-go/apierr" + "github.com/databricks/databricks-sdk-go/client" + "github.com/databricks/databricks-sdk-go/service/jobs" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestClassifyLogError(t *testing.T) { + tests := []struct { + name string + err error + want error // errBricklensFeatureDisabled, the input error, or nil + }{ + { + name: "feature disabled falls back", + err: &apierr.APIError{ErrorCode: "FEATURE_DISABLED", StatusCode: http.StatusForbidden}, + want: errBricklensFeatureDisabled, + }, + { + name: "endpoint not found falls back", + err: &apierr.APIError{ErrorCode: "ENDPOINT_NOT_FOUND", StatusCode: http.StatusNotFound}, + want: errBricklensFeatureDisabled, + }, + { + name: "bare 404 falls back", + err: &apierr.APIError{ErrorCode: "SOMETHING", StatusCode: http.StatusNotFound}, + want: errBricklensFeatureDisabled, + }, + { + name: "genuine resource-does-not-exist surfaces", + err: apierr.ErrResourceDoesNotExist, + want: apierr.ErrResourceDoesNotExist, + }, + { + name: "transient 500 is retried", + err: &apierr.APIError{ErrorCode: "INTERNAL", StatusCode: http.StatusInternalServerError}, + want: nil, + }, + { + name: "plain error is retried", + err: errors.New("connection reset"), + want: nil, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := classifyLogError(tt.err) + switch tt.want { + case errBricklensFeatureDisabled: + assert.ErrorIs(t, got, errBricklensFeatureDisabled) + case nil: + assert.NoError(t, got) + default: + assert.ErrorIs(t, got, tt.want) + } + }) + } +} + +func TestProjectRunStatus(t *testing.T) { + run := &jobs.Run{ + StartTime: 1000, + EndTime: 2000, + State: &jobs.RunState{ + LifeCycleState: jobs.RunLifeCycleStateTerminated, + ResultState: jobs.RunResultStateSuccess, + StateMessage: "done", + }, + Tasks: []jobs.RunTask{ + {AttemptNumber: 0}, + {AttemptNumber: 2}, + {AttemptNumber: 1}, + }, + } + + s := projectRunStatus(run) + assert.Equal(t, "TERMINATED", s.lifeCycleState) + assert.Equal(t, "SUCCESS", s.resultState) + assert.Equal(t, "done", s.stateMessage) + assert.Equal(t, int64(1000), s.startTimeMs) + assert.Equal(t, int64(2000), s.endTimeMs) + assert.Equal(t, 2, s.latestAttempt) + assert.True(t, s.terminal()) + assert.True(t, s.succeeded()) + assert.Equal(t, "SUCCESS", s.displayState()) +} + +func TestLogRunStatusTerminal(t *testing.T) { + tests := []struct { + name string + lifeCycle string + resultState string + wantTerminal bool + }{ + {"running", "RUNNING", "", false}, + {"pending", "PENDING", "", false}, + {"terminated lifecycle", "TERMINATED", "", true}, + {"internal error lifecycle", "INTERNAL_ERROR", "", true}, + {"failed result", "TERMINATING", "FAILED", true}, + {"canceled result", "RUNNING", "CANCELED", true}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + s := logRunStatus{lifeCycleState: tt.lifeCycle, resultState: tt.resultState} + assert.Equal(t, tt.wantTerminal, s.terminal()) + }) + } +} + +func TestLogRequestFromSeconds(t *testing.T) { + now := time.Unix(10_000, 0) + + // --minutes narrows the window to now - N*60. + req := logRequest{windowMinutes: 5} + assert.Equal(t, int64(10_000-300), req.fromSeconds(logRunStatus{startTimeMs: 1_000_000}, now)) + + // No window: from the run's start second. + req = logRequest{} + assert.Equal(t, int64(1000), req.fromSeconds(logRunStatus{startTimeMs: 1_000_000}, now)) + + // No window, run not started: everything stored (0). + assert.Equal(t, int64(0), req.fromSeconds(logRunStatus{}, now)) +} + +func TestLogRequestToSeconds(t *testing.T) { + req := logRequest{} + + // Active run: 0 lets the endpoint default to now. + assert.Equal(t, int64(0), req.toSeconds(logRunStatus{lifeCycleState: "RUNNING"})) + + // Terminal run: ceil of the end millisecond so the final partial second is kept. + terminal := logRunStatus{lifeCycleState: "TERMINATED", resultState: "SUCCESS", endTimeMs: 2001} + assert.Equal(t, int64(3), req.toSeconds(terminal)) +} + +func TestLogRequestTailTarget(t *testing.T) { + // Negative (unset) uses the default cap; explicit values are literal. + assert.Equal(t, defaultCompletedRunTailLines, logRequest{tailLines: -1}.tailTarget()) + assert.Equal(t, 42, logRequest{tailLines: 42}.tailTarget()) + assert.Equal(t, 0, logRequest{tailLines: 0}.tailTarget()) +} + +func TestDrainPagesDedupAndOrdering(t *testing.T) { + // Two pages: page 1 has two ascending records; page 2 repeats the last record + // of page 1 (boundary re-query — must dedup) and includes an older record + // (out of order — must skip), then a genuinely newer one. + var page int + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Query().Get("page_token") == "" { + page = 1 + _, _ = w.Write([]byte(`{"log_records": [ + {"time_unix_nano": 1000, "body": "a", "node_index": 0}, + {"time_unix_nano": 2000, "body": "b", "node_index": 0} + ], "next_page_token": "p2"}`)) + return + } + page = 2 + _, _ = w.Write([]byte(`{"log_records": [ + {"time_unix_nano": 2000, "body": "b", "node_index": 0}, + {"time_unix_nano": 1500, "body": "stale", "node_index": 0}, + {"time_unix_nano": 3000, "body": "c", "node_index": 0} + ]}`)) + })) + t.Cleanup(srv.Close) + + var buf bytes.Buffer + w := newTestWorkspaceClient(t, srv.URL) + apiClient, err := client.New(w.Config) + require.NoError(t, err) + st := &bricklensStreamer{ + ctx: t.Context(), + w: w, + apiClient: apiClient, + out: &buf, + req: logRequest{runID: 1, node: 0, attempt: -1}, + seen: newSeenSet(seenRecordsCap), + } + require.NoError(t, st.drainPages(0)) + require.Equal(t, 2, page) + + // "b" prints once (deduped), "stale" is skipped (older than last emitted), and + // fromSec advances to the newest record's floor-second (3000ns -> 0s here). + assert.Equal(t, "a\nb\nc\n", buf.String()) + assert.Equal(t, int64(3000), st.lastNano) +} + +func TestDisplayState(t *testing.T) { + assert.Equal(t, "SUCCESS", logRunStatus{lifeCycleState: "TERMINATED", resultState: "SUCCESS"}.displayState()) + assert.Equal(t, "RUNNING", logRunStatus{lifeCycleState: "RUNNING"}.displayState()) + assert.Equal(t, "UNKNOWN", logRunStatus{}.displayState()) +} + +func TestEmitLogLineJSON(t *testing.T) { + var buf bytes.Buffer + emitLogLine(&buf, logRequest{node: 2, jsonOutput: true}, "hello") + + var ev logEvent + require.NoError(t, json.Unmarshal(buf.Bytes(), &ev)) + assert.Equal(t, "LOG", ev.Type) + assert.Equal(t, 2, ev.Node) + assert.Equal(t, "hello", ev.Line) + assert.NotEmpty(t, ev.TS) +} + +func TestEmitLogLineText(t *testing.T) { + var buf bytes.Buffer + emitLogLine(&buf, logRequest{node: 0}, "hello") + assert.Equal(t, "hello\n", buf.String()) +} + +func TestEmitLogLineJSONFatalEmitsAlert(t *testing.T) { + var buf bytes.Buffer + emitLogLine(&buf, logRequest{node: 1, jsonOutput: true}, "CUDA out of memory") + + // A fatal line emits an ALERT event before its LOG event. + lines := strings.Split(strings.TrimSpace(buf.String()), "\n") + require.Len(t, lines, 2) + + var alert, logEv logEvent + require.NoError(t, json.Unmarshal([]byte(lines[0]), &alert)) + require.NoError(t, json.Unmarshal([]byte(lines[1]), &logEv)) + assert.Equal(t, "ALERT", alert.Type) + assert.Equal(t, "LOG", logEv.Type) + assert.Equal(t, "CUDA out of memory", alert.Line) + + // Text mode never emits ALERT events. + var text bytes.Buffer + emitLogLine(&text, logRequest{node: 1}, "CUDA out of memory") + assert.Equal(t, "CUDA out of memory\n", text.String()) +} + +func TestMatchFatalPattern(t *testing.T) { + fatal := []string{ + "CUDA out of memory", + "cuda OUT OF memory", + "Watchdog caught collective operation timeout", + "Killed", + "ERROR: Script failed with exit code 1 after 42s", + "bash: foo: command not found", + } + for _, l := range fatal { + assert.True(t, matchFatalPattern(l), l) + } + + notFatal := []string{ + "epoch 3 loss 0.5", + "ERROR: Script failed with exit code 0 after 42s", + "just a normal line", + } + for _, l := range notFatal { + assert.False(t, matchFatalPattern(l), l) + } +} + +func TestEmitNoLogs(t *testing.T) { + tests := []struct { + name string + status logRunStatus + want string + }{ + { + name: "terminal", + status: logRunStatus{lifeCycleState: "TERMINATED", resultState: "FAILED", stateMessage: "boom"}, + want: "No logs available for run 7. Run terminated in state FAILED: boom\n", + }, + { + name: "running", + status: logRunStatus{lifeCycleState: "RUNNING"}, + want: "No logs available yet for run 7, which is still in state RUNNING\n", + }, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + var text bytes.Buffer + emitNoLogs(&text, logRequest{runID: 7}, tc.status) + assert.Equal(t, tc.want, text.String()) + + var jsonBuf bytes.Buffer + emitNoLogs(&jsonBuf, logRequest{runID: 7, node: 1, jsonOutput: true}, tc.status) + var ev logEvent + require.NoError(t, json.Unmarshal(jsonBuf.Bytes(), &ev)) + assert.Equal(t, "ERROR", ev.Type) + assert.Equal(t, 1, ev.Node) + assert.Equal(t, strings.TrimRight(tc.want, "\n"), ev.Line) + }) + } +} + +func TestRequestPageRetriesThenFallsBack(t *testing.T) { + // Shrink the retry wait so the transient-failure loop runs fast. + orig := retryCheckInterval + retryCheckInterval = time.Millisecond + t.Cleanup(func() { retryCheckInterval = orig }) + + var calls int + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if !strings.Contains(r.URL.Path, "/logs") { + // Ignore SDK host/config probes so `calls` counts only log requests. + _, _ = w.Write([]byte(`{}`)) + return + } + calls++ + w.WriteHeader(http.StatusInternalServerError) + _, _ = w.Write([]byte(`{"error_code": "INTERNAL_ERROR", "message": "transient"}`)) + })) + t.Cleanup(srv.Close) + + w := newTestWorkspaceClient(t, srv.URL) + apiClient, err := client.New(w.Config) + require.NoError(t, err) + st := &bricklensStreamer{ + ctx: t.Context(), + w: w, + apiClient: apiClient, + req: logRequest{runID: 1, node: 0, attempt: -1}, + seen: newSeenSet(seenRecordsCap), + } + _, err = st.requestPage("", 0, 0, true) + // Persistent transient failures fall back to MLflow after the retry budget. + require.ErrorIs(t, err, errBricklensFeatureDisabled) + assert.Equal(t, maxTransientFailures, calls) +} + +func TestRequestPageRetriesThenSucceeds(t *testing.T) { + orig := retryCheckInterval + retryCheckInterval = time.Millisecond + t.Cleanup(func() { retryCheckInterval = orig }) + + var calls int + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if !strings.Contains(r.URL.Path, "/logs") { + _, _ = w.Write([]byte(`{}`)) + return + } + calls++ + if calls < 3 { + w.WriteHeader(http.StatusInternalServerError) + _, _ = w.Write([]byte(`{"error_code": "INTERNAL_ERROR", "message": "transient"}`)) + return + } + _, _ = w.Write([]byte(`{"log_records": [{"time_unix_nano": 1, "body": "ok", "node_index": 0}]}`)) + })) + t.Cleanup(srv.Close) + + wc := newTestWorkspaceClient(t, srv.URL) + apiClient, err := client.New(wc.Config) + require.NoError(t, err) + st := &bricklensStreamer{ + ctx: t.Context(), + w: wc, + apiClient: apiClient, + req: logRequest{runID: 1, node: 0, attempt: -1}, + seen: newSeenSet(seenRecordsCap), + } + resp, err := st.requestPage("", 0, 0, true) + require.NoError(t, err) + require.Len(t, resp.LogRecords, 1) + assert.Equal(t, "ok", resp.LogRecords[0].Body) + assert.Equal(t, 3, calls) +} + +func TestSeenSetEviction(t *testing.T) { + s := newSeenSet(2) + s.add(1, "a") + s.add(2, "b") + assert.True(t, s.has(1, "a")) + assert.True(t, s.has(2, "b")) + + // Adding a third evicts the oldest-inserted (1,"a"). + s.add(3, "c") + assert.False(t, s.has(1, "a")) + assert.True(t, s.has(2, "b")) + assert.True(t, s.has(3, "c")) + + // Same (nano, body) shares one entry; distinct body under the same nano does not. + s.add(3, "c") + assert.True(t, s.has(3, "c")) + assert.False(t, s.has(3, "d")) +} + +func TestSleepOrCancel(t *testing.T) { + // Returns nil once the duration elapses. + require.NoError(t, sleepOrCancel(t.Context(), time.Millisecond)) + + // Returns the context error promptly when cancelled. + ctx, cancel := context.WithCancel(t.Context()) + cancel() + require.ErrorIs(t, sleepOrCancel(ctx, time.Hour), context.Canceled) +} diff --git a/experimental/air/cmd/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_test.go b/experimental/air/cmd/render_test.go index 2fe125bcf04..067f5817fb1 100644 --- a/experimental/air/cmd/render_test.go +++ b/experimental/air/cmd/render_test.go @@ -151,7 +151,7 @@ func TestRenderFields(t *testing.T) { experiment: "stream-latency-test", mlflowLabel: "stream-latency-test", mlflowURL: "https://h.test/ml/experiments/E1/runs/R1", - user: "riddhi.bhagwat@databricks.com", + user: "user@example.com", accelerators: "1x A10", environment: "ml-runtime-gpu:1.0", }) diff --git a/experimental/air/cmd/run.go b/experimental/air/cmd/run.go index bd32810e9bc..ea00368679f 100644 --- a/experimental/air/cmd/run.go +++ b/experimental/air/cmd/run.go @@ -1,7 +1,7 @@ package aircmd import ( - "errors" + "context" "fmt" "strconv" @@ -9,6 +9,7 @@ import ( "github.com/databricks/cli/libs/cmdctx" "github.com/databricks/cli/libs/cmdio" "github.com/databricks/cli/libs/flags" + "github.com/databricks/databricks-sdk-go" "github.com/spf13/cobra" ) @@ -57,16 +58,7 @@ The workload is described by a YAML config file (see --file).`, cmd.RunE = func(cmd *cobra.Command, args []string) error { ctx := cmd.Context() - // These flags' pipelines are not ported yet; reject rather than silently - // ignore them. - if len(overrides) > 0 { - return errors.New("--override is not yet supported") - } - if watch { - return errors.New("--watch is not yet supported") - } - - cfg, err := loadRunConfig(file) + cfg, err := loadRunConfigWithOverrides(ctx, file, overrides) if err != nil { return err } @@ -86,13 +78,63 @@ The workload is described by a YAML config file (see --file).`, } runIDStr := strconv.FormatInt(runID, 10) - if root.OutputType(cmd) == flags.OutputText { + jsonOut := root.OutputType(cmd) == flags.OutputJSON + + if !watch { + if !jsonOut { + cmdio.LogString(ctx, "Submitted run "+runIDStr) + cmdio.LogString(ctx, "View at: "+dashboardURL) + cmdio.LogString(ctx, "\nTip: use --watch to stream logs until the run completes.") + return nil + } + return renderEnvelope(ctx, runResult{Status: "SUBMITTED", RunID: runIDStr, DashboardURL: dashboardURL}) + } + + // --watch: stream the submitted run's logs until it reaches a terminal + // state, then exit with the run's outcome. This is the same pipeline as + // `air logs ` (Bricklens with MLflow fallback). + req := logRequest{ + runID: runID, + attempt: -1, + tailLines: -1, + jsonOutput: jsonOut, + } + + if !jsonOut { cmdio.LogString(ctx, "Submitted run "+runIDStr) cmdio.LogString(ctx, "View at: "+dashboardURL) - return nil + cmdio.LogString(ctx, "Monitoring run and streaming logs...") + return runLogs(ctx, cmd, req) + } + + // --json: emit SUBMITTED first (so a consumer sees the run id immediately), + // STATUS events on each lifecycle transition, and a closing terminal-status + // envelope after streaming. Mirrors the Python CLI's --watch JSONL contract. + out := cmd.OutOrStdout() + printSubmittedEvent(out, runIDStr, dashboardURL) + req.onStatusChange = func(current, previous string) { + printStatusEvent(out, current, previous) } - return renderEnvelope(ctx, runResult{Status: "SUBMITTED", RunID: runIDStr, DashboardURL: dashboardURL}) + err = runLogs(ctx, cmd, req) + + // Re-resolve the run for the closing envelope. STATUS events only fire on + // the Bricklens path, so the terminal status must come from the run's + // actual state — correct whether Bricklens or the MLflow fallback served + // the logs. + printTerminalEvent(out, runIDStr, watchTerminalStatus(ctx, w, runID), dashboardURL) + return err } return cmd } + +// watchTerminalStatus resolves a watched run's final display state for the +// closing --watch envelope. The run is terminal once streaming returns; if the +// status can't be re-fetched, "UNKNOWN" is reported rather than guessing. +func watchTerminalStatus(ctx context.Context, w *databricks.WorkspaceClient, runID int64) string { + status, err := resolveRunStatus(ctx, w, runID) + if err != nil { + return "UNKNOWN" + } + return status.displayState() +} diff --git a/experimental/air/cmd/run_watch_test.go b/experimental/air/cmd/run_watch_test.go new file mode 100644 index 00000000000..487b0b9989a --- /dev/null +++ b/experimental/air/cmd/run_watch_test.go @@ -0,0 +1,206 @@ +package aircmd + +import ( + "bytes" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/databricks/cli/cmd/root" + "github.com/databricks/cli/libs/cmdctx" + "github.com/databricks/cli/libs/cmdio" + "github.com/databricks/cli/libs/flags" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// watchServer serves submit, the auth probe, a terminal runs/get, and a single +// page of Bricklens logs — everything `air run --watch` touches after submit. +// resultState is the terminal result the run reports (e.g. SUCCESS, FAILED). +func watchServer(t *testing.T, resultState string) *httptest.Server { + t.Helper() + runGet := `{ + "run_id": 777, + "start_time": 1700000000000, + "end_time": 1700000012000, + "state": {"life_cycle_state": "TERMINATED", "result_state": "` + resultState + `"}, + "tasks": [{"run_id": 778, "attempt_number": 0}] + }` + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch { + case strings.HasSuffix(r.URL.Path, "/jobs/runs/submit"): + _, _ = w.Write([]byte(`{"run_id": 777}`)) + case r.URL.Path == "/api/2.2/jobs/runs/get": + _, _ = w.Write([]byte(runGet)) + case strings.HasPrefix(r.URL.Path, "/api/2.0/ai-training/workflows/by-run-id/"): + _, _ = w.Write([]byte(`{"log_records": [ + {"time_unix_nano": 1700000002000000000, "body": "step 2", "node_index": 0}, + {"time_unix_nano": 1700000001000000000, "body": "step 1", "node_index": 0} + ]}`)) + default: + // Me() probe, workspace-id, SDK config discovery. + _, _ = w.Write([]byte(`{"userName": "u@example.com", "workspace_id": 1}`)) + } + })) + t.Cleanup(srv.Close) + return srv +} + +// watchServerMLflow serves a run whose Bricklens endpoint is gated off +// (FEATURE_DISABLED), forcing the MLflow fallback, plus the MLflow artifact +// chain (get-output, artifacts/list, credentials-for-read, the pre-signed bytes). +// STATUS events never fire on this path, so it guards the closing terminal +// envelope against relying on onStatusChange. +func watchServerMLflow(t *testing.T, resultState string) *httptest.Server { + t.Helper() + var base string + runGet := `{ + "run_id": 777, + "state": {"life_cycle_state": "TERMINATED", "result_state": "` + resultState + `"}, + "tasks": [{"run_id": 778, "attempt_number": 0}] + }` + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch { + case strings.HasSuffix(r.URL.Path, "/jobs/runs/submit"): + _, _ = w.Write([]byte(`{"run_id": 777}`)) + case r.URL.Path == "/api/2.2/jobs/runs/get": + _, _ = w.Write([]byte(runGet)) + case strings.HasPrefix(r.URL.Path, "/api/2.0/ai-training/workflows/by-run-id/"): + w.WriteHeader(http.StatusForbidden) + _, _ = w.Write([]byte(`{"error_code": "FEATURE_DISABLED", "message": "gated off"}`)) + case r.URL.Path == "/api/2.2/jobs/runs/get-output": + _, _ = w.Write([]byte(`{"ai_runtime_task_output": {"mlflow_experiment_id": "exp1", "mlflow_run_id": "run1"}}`)) + case r.URL.Path == "/api/2.0/mlflow/artifacts/list": + if r.URL.Query().Get("path") == "logs" { + _, _ = w.Write([]byte(`{"files": [{"path": "logs/node_0", "is_dir": true}]}`)) + return + } + _, _ = w.Write([]byte(`{"files": [{"path": "logs/node_0/logs-0.chunk.txt", "file_size": 12}]}`)) + case r.URL.Path == "/api/2.0/mlflow/artifacts/credentials-for-read": + _, _ = w.Write([]byte(`{"credential_infos": [{"signed_uri": "` + base + `/presigned"}]}`)) + case r.URL.Path == "/presigned": + _, _ = w.Write([]byte("step 1\nstep 2\n")) + default: + _, _ = w.Write([]byte(`{"userName": "u@example.com", "workspace_id": 1}`)) + } + })) + base = srv.URL + t.Cleanup(srv.Close) + return srv +} + +func runWatchCmd(t *testing.T, out flags.Output, buf *bytes.Buffer, srvURL string) error { + t.Helper() + cfgPath := writeConfigFile(t, "run.yaml", minimalConfig) + cmd := withOutput(newRunCommand(), out) + require.NoError(t, cmd.Flags().Set("file", cfgPath)) + require.NoError(t, cmd.Flags().Set("watch", "true")) + + ctx := cmdio.InContext(t.Context(), cmdio.NewIO(t.Context(), out, nil, buf, buf, "", "")) + ctx = cmdctx.SetWorkspaceClient(ctx, newTestWorkspaceClient(t, srvURL)) + cmd.SetContext(ctx) + cmd.SetOut(buf) + return cmd.RunE(cmd, nil) +} + +func TestRunWatchStreamsLogs(t *testing.T) { + var buf bytes.Buffer + err := runWatchCmd(t, flags.OutputText, &buf, watchServer(t, "SUCCESS").URL) + require.NoError(t, err) + + out := buf.String() + assert.Contains(t, out, "Submitted run 777") + assert.Contains(t, out, "Monitoring run and streaming logs...") + // The submitted run's logs stream through, oldest-first. + assert.Contains(t, out, "step 1\nstep 2") +} + +func TestRunWatchJSONEmitsSubmittedThenLogs(t *testing.T) { + var buf bytes.Buffer + err := runWatchCmd(t, flags.OutputJSON, &buf, watchServer(t, "SUCCESS").URL) + require.NoError(t, err) + + all := buf.String() + lines := strings.Split(strings.TrimSpace(all), "\n") + require.GreaterOrEqual(t, len(lines), 3) + // First event is SUBMITTED with the run id; then STATUS + streamed LOG events. + assert.Contains(t, lines[0], `"type":"SUBMITTED"`) + assert.Contains(t, lines[0], `"run_id":"777"`) + assert.Contains(t, all, `"type":"STATUS"`) + assert.Contains(t, all, `"type":"LOG"`) + assert.Contains(t, all, `"line":"step 1"`) + // The last line is the closing terminal-status envelope carrying SUCCESS. + assert.Contains(t, lines[len(lines)-1], `"status":"SUCCESS"`) + assert.Contains(t, lines[len(lines)-1], `"run_id":"777"`) +} + +func TestRunWatchJSONFailedRunTerminalEnvelope(t *testing.T) { + var buf bytes.Buffer + err := runWatchCmd(t, flags.OutputJSON, &buf, watchServer(t, "FAILED").URL) + // Non-zero exit is surfaced as ErrAlreadyPrinted, but the closing envelope + // still carries the terminal status. + require.ErrorIs(t, err, root.ErrAlreadyPrinted) + lines := strings.Split(strings.TrimSpace(buf.String()), "\n") + assert.Contains(t, lines[len(lines)-1], `"status":"FAILED"`) +} + +func TestRunWatchFailedRunExitsNonZero(t *testing.T) { + var buf bytes.Buffer + // A run that ends FAILED streams its logs but exits non-zero, surfaced as + // ErrAlreadyPrinted (the output was already written). + err := runWatchCmd(t, flags.OutputText, &buf, watchServer(t, "FAILED").URL) + require.ErrorIs(t, err, root.ErrAlreadyPrinted) + assert.Contains(t, buf.String(), "step 1\nstep 2") +} + +func TestRunWatchDryRunSkipsSubmit(t *testing.T) { + // --dry-run takes precedence over --watch: nothing is submitted or streamed. + var got []string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + got = append(got, r.URL.Path) + _, _ = w.Write([]byte(`{}`)) + })) + t.Cleanup(srv.Close) + + cfgPath := writeConfigFile(t, "run.yaml", minimalConfig) + cmd := withOutput(newRunCommand(), flags.OutputText) + require.NoError(t, cmd.Flags().Set("file", cfgPath)) + require.NoError(t, cmd.Flags().Set("watch", "true")) + require.NoError(t, cmd.Flags().Set("dry-run", "true")) + var buf bytes.Buffer + ctx := cmdio.InContext(t.Context(), cmdio.NewIO(t.Context(), flags.OutputText, nil, &buf, &buf, "", "")) + ctx = cmdctx.SetWorkspaceClient(ctx, newTestWorkspaceClient(t, srv.URL)) + cmd.SetContext(ctx) + cmd.SetOut(&buf) + + require.NoError(t, cmd.RunE(cmd, nil)) + assert.Contains(t, buf.String(), "Dry run") + for _, p := range got { + assert.NotContains(t, p, "/jobs/runs/submit", "dry-run must not submit") + assert.NotContains(t, p, "/logs", "dry-run must not stream logs") + } +} + +func TestRunWatchJSONMLflowFallbackTerminalEnvelope(t *testing.T) { + // Regression: through the MLflow fallback the terminal status must come from + // the run's actual state, not the onStatusChange callback (which is + // Bricklens-only), so a SUCCESS run isn't mislabeled FAILED in the envelope. + var buf bytes.Buffer + err := runWatchCmd(t, flags.OutputJSON, &buf, watchServerMLflow(t, "SUCCESS").URL) + require.NoError(t, err) + + lines := strings.Split(strings.TrimSpace(buf.String()), "\n") + assert.Contains(t, lines[0], `"type":"SUBMITTED"`) + // Logs stream through the fallback, and the closing envelope reflects the + // real terminal status. + assert.Contains(t, buf.String(), `"line":"step 1"`) + assert.Contains(t, lines[len(lines)-1], `"status":"SUCCESS"`) +} + +func TestRunWatchFlagRegistered(t *testing.T) { + cmd := newRunCommand() + f := cmd.Flags().Lookup("watch") + require.NotNil(t, f) + assert.Equal(t, "false", f.DefValue) +} diff --git a/experimental/air/cmd/runconfig.go b/experimental/air/cmd/runconfig.go index 09437f50a5b..4cfbf3736d1 100644 --- a/experimental/air/cmd/runconfig.go +++ b/experimental/air/cmd/runconfig.go @@ -31,6 +31,11 @@ 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 { @@ -164,8 +169,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 diff --git a/experimental/air/cmd/runconfig_load.go b/experimental/air/cmd/runconfig_load.go index 81b07d3ca50..7f6ad8b5e17 100644 --- a/experimental/air/cmd/runconfig_load.go +++ b/experimental/air/cmd/runconfig_load.go @@ -1,6 +1,8 @@ package aircmd import ( + "bytes" + "context" "errors" "fmt" "io" @@ -10,10 +12,11 @@ import ( ) // decodeRunConfig reads and decodes the run YAML into the schema. Unknown keys -// are rejected (KnownFields), mirroring the Python schema's extra="forbid". +// are rejected (KnownFields). // -// The `_bases_` composition feature and CLI `--override` handling are not yet -// ported; a config using `_bases_` is currently rejected as an unknown field. +// The `_bases_` composition feature is not yet ported; a config using `_bases_` +// is currently rejected as an unknown field. CLI `--override` handling lives in +// runconfig_override.go and is applied to the parsed map before this decode. func decodeRunConfig(path string) (*runConfig, error) { f, err := os.Open(path) if err != nil { @@ -21,7 +24,13 @@ func decodeRunConfig(path string) (*runConfig, error) { } defer f.Close() - dec := yaml.NewDecoder(f) + return decodeRunConfigReader(f, path) +} + +// decodeRunConfigReader decodes and unknown-key-checks a run YAML from r. path is +// used only for error messages. +func decodeRunConfigReader(r io.Reader, path string) (*runConfig, error) { + dec := yaml.NewDecoder(r) dec.KnownFields(true) var cfg runConfig @@ -50,3 +59,53 @@ func loadRunConfig(path string) (*runConfig, error) { } return cfg, nil } + +// loadRunConfigWithOverrides decodes a run YAML config, applies any +// --override KEY=VALUE entries to the parsed map, then re-decodes (with unknown +// keys rejected) and structurally validates the result. Applying overrides to +// the map — rather than the typed config — lets the single decode+validate +// pipeline enforce path existence, type coercion, and the semantic rules at +// once. ctx is used only to log applied overrides. +func loadRunConfigWithOverrides(ctx context.Context, path string, overrides []string) (*runConfig, error) { + if len(overrides) == 0 { + return loadRunConfig(path) + } + + entries, err := parseOverrides(overrides) + if err != nil { + return nil, err + } + if err := validateOverridePaths(entries); err != nil { + return nil, err + } + + raw, err := os.ReadFile(path) + if err != nil { + return nil, err + } + var m map[string]any + if err := yaml.Unmarshal(raw, &m); err != nil { + return nil, fmt.Errorf("invalid config %s: %w", path, err) + } + if m == nil { + // An empty file decodes to a nil map; start from an empty one so overrides + // can populate it (the re-decode still enforces required fields). + m = map[string]any{} + } + if err := applyOverrides(ctx, m, entries); err != nil { + return nil, err + } + + merged, err := yaml.Marshal(m) + if err != nil { + return nil, err + } + cfg, err := decodeRunConfigReader(bytes.NewReader(merged), path) + if err != nil { + return nil, err + } + if err := validateRunConfig(cfg); err != nil { + return nil, err + } + return cfg, nil +} diff --git a/experimental/air/cmd/runconfig_override.go b/experimental/air/cmd/runconfig_override.go new file mode 100644 index 00000000000..a00a178603b --- /dev/null +++ b/experimental/air/cmd/runconfig_override.go @@ -0,0 +1,162 @@ +package aircmd + +import ( + "context" + "fmt" + "maps" + "reflect" + "slices" + "strings" + + "github.com/databricks/cli/libs/cmdio" + "go.yaml.in/yaml/v3" +) + +// This file implements the `--override KEY=VALUE` flag. Overrides are applied to +// the parsed YAML map (not the typed runConfig) before re-decode, so one pipeline +// covers path existence, type coercion, and the semantic validate() rules. + +// freeFormFields hold free-form maps, so path validation stops at them: any +// sub-path is valid. +var freeFormFields = map[string]bool{ + "parameters": true, + "env_variables": true, + "secrets": true, +} + +// parseOverrides parses --override KEY=VALUE arguments, preserving order. +func parseOverrides(overrides []string) ([]overrideEntry, error) { + entries := make([]overrideEntry, 0, len(overrides)) + for _, item := range overrides { + key, value, found := strings.Cut(item, "=") + if !found { + // --override is repeatable, so a config path meant for -f can be + // swallowed here; point at the real fix. + hint := "" + if strings.HasSuffix(item, ".yaml") || strings.HasSuffix(item, ".yml") { + hint = fmt.Sprintf("; %q looks like a config file — pass it with -f/--file", item) + } + return nil, fmt.Errorf("invalid --override %q: expected KEY=VALUE (e.g. compute.num_accelerators=32)%s", item, hint) + } + key = strings.TrimSpace(key) + if key == "" { + return nil, fmt.Errorf("invalid --override %q: empty key", item) + } + entries = append(entries, overrideEntry{path: key, raw: value}) + } + return entries, nil +} + +// overrideEntry is one parsed --override: its dotted path and the raw RHS string. +type overrideEntry struct { + path string + raw string +} + +// validateOverridePaths checks every dotted path against the runConfig schema +// before mutation, so an error names the exact --override key rather than the +// re-decode's Go-type language. +func validateOverridePaths(entries []overrideEntry) error { + for _, e := range entries { + if err := checkOverridePath(strings.Split(e.path, "."), reflect.TypeFor[runConfig](), e.path); err != nil { + return err + } + } + return nil +} + +// checkOverridePath recursively validates one dotted path against a struct type +// whose fields carry `yaml:` tags. +func checkOverridePath(parts []string, t reflect.Type, fullPath string) error { + field := parts[0] + fields := yamlFields(t) + sub, ok := fields[field] + if !ok { + return fmt.Errorf("invalid --override %q: %q is not a known field; available fields are: %s", + fullPath, field, strings.Join(slices.Sorted(maps.Keys(fields)), ", ")) + } + if len(parts) == 1 { + return nil + } + if freeFormFields[field] { + return nil + } + subStruct := underlyingStruct(sub) + if subStruct == nil { + return fmt.Errorf("invalid --override %q: %q is not a nested object; cannot address sub-field %q", + fullPath, field, strings.Join(parts[1:], ".")) + } + return checkOverridePath(parts[1:], subStruct, fullPath) +} + +// yamlFields maps a struct's yaml tag names to their field types, skipping +// fields without a yaml tag (or tagged "-"). +func yamlFields(t reflect.Type) map[string]reflect.Type { + out := map[string]reflect.Type{} + for f := range t.Fields() { + tag := f.Tag.Get("yaml") + if tag == "" || tag == "-" { + continue + } + name, _, _ := strings.Cut(tag, ",") + if name == "" || name == "-" { + continue + } + out[name] = f.Type + } + return out +} + +// underlyingStruct unwraps pointer/slice indirection and returns the struct type +// a field decodes into, or nil if the field is not a struct (a scalar/map/etc.). +func underlyingStruct(t reflect.Type) reflect.Type { + for t.Kind() == reflect.Pointer || t.Kind() == reflect.Slice { + t = t.Elem() + } + if t.Kind() == reflect.Struct { + return t + } + return nil +} + +// applyOverrides walks each dotted path into the parsed YAML map and sets the +// leaf to the RHS parsed as a YAML scalar. Intermediate maps are auto-created so +// an override can add a field the YAML omits; the later re-decode rejects paths +// absent from the schema. Changes are logged to stderr to keep JSON stdout clean. +func applyOverrides(ctx context.Context, m map[string]any, entries []overrideEntry) error { + for _, e := range entries { + var value any + if err := yaml.Unmarshal([]byte(e.raw), &value); err != nil { + return fmt.Errorf("invalid --override %q: cannot parse value %q: %w", e.path, e.raw, err) + } + + parts := strings.Split(e.path, ".") + current := m + for _, part := range parts[:len(parts)-1] { + next, ok := current[part].(map[string]any) + if !ok { + next = map[string]any{} + current[part] = next + } + current = next + } + + leaf := parts[len(parts)-1] + old, had := current[leaf] + current[leaf] = value + if had { + logOverride(ctx, fmt.Sprintf("Override: changing %s from %v to %v", e.path, old, value)) + } else { + logOverride(ctx, fmt.Sprintf("Override: setting %s to %v", e.path, value)) + } + } + return nil +} + +// logOverride writes to stderr only when a cmdIO is present; cmdio.LogString +// panics without one, as in non-command callers such as unit tests. +func logOverride(ctx context.Context, msg string) { + if cmdio.HasIO(ctx) { + cmdio.LogString(ctx, msg) + } +} diff --git a/experimental/air/cmd/runconfig_override_test.go b/experimental/air/cmd/runconfig_override_test.go new file mode 100644 index 00000000000..99eab7dd881 --- /dev/null +++ b/experimental/air/cmd/runconfig_override_test.go @@ -0,0 +1,168 @@ +package aircmd + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestParseOverrides(t *testing.T) { + tests := []struct { + name string + in []string + want []overrideEntry + wantErr string + }{ + { + name: "key=value pairs preserve order", + in: []string{"compute.num_accelerators=8", "timeout_minutes=45"}, + want: []overrideEntry{ + {path: "compute.num_accelerators", raw: "8"}, + {path: "timeout_minutes", raw: "45"}, + }, + }, + { + name: "value may contain =", + in: []string{"env_variables.EXPR=a=b"}, + want: []overrideEntry{{path: "env_variables.EXPR", raw: "a=b"}}, + }, + { + name: "key is trimmed", + in: []string{" timeout_minutes = 45"}, + want: []overrideEntry{{path: "timeout_minutes", raw: " 45"}}, + }, + { + name: "missing = is rejected", + in: []string{"compute.num_accelerators"}, + wantErr: `expected KEY=VALUE`, + }, + { + name: "a .yaml token hints at -f", + in: []string{"train.yaml"}, + wantErr: `looks like a config file`, + }, + { + name: "empty key is rejected", + in: []string{"=5"}, + wantErr: `empty key`, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got, err := parseOverrides(tt.in) + if tt.wantErr != "" { + require.ErrorContains(t, err, tt.wantErr) + return + } + require.NoError(t, err) + assert.Equal(t, tt.want, got) + }) + } +} + +func TestValidateOverridePaths(t *testing.T) { + tests := []struct { + name string + path string + wantErr string + }{ + {name: "known top-level field", path: "experiment_name"}, + {name: "known nested field", path: "compute.num_accelerators"}, + {name: "free-form sub-path", path: "env_variables.MY_VAR"}, + {name: "deep free-form sub-path", path: "parameters.model.layers"}, + { + name: "unknown top-level field", + path: "bogus", + wantErr: `"bogus" is not a known field`, + }, + { + name: "unknown nested field", + path: "compute.bogus", + wantErr: `"bogus" is not a known field`, + }, + { + name: "sub-field of a scalar", + path: "command.sub", + wantErr: `"command" is not a nested object`, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := validateOverridePaths([]overrideEntry{{path: tt.path, raw: "x"}}) + if tt.wantErr != "" { + require.ErrorContains(t, err, tt.wantErr) + return + } + require.NoError(t, err) + }) + } +} + +// overrideBaseConfig is a valid 8-GPU config the override tests mutate. +const overrideBaseConfig = `experiment_name: smoke +command: python train.py +compute: + accelerator_type: GPU_8xH100 + num_accelerators: 8 +env_variables: + EXISTING: hello +` + +func TestLoadRunConfigWithOverrides(t *testing.T) { + t.Run("no overrides matches loadRunConfig", func(t *testing.T) { + cfg, err := loadRunConfigWithOverrides(t.Context(), writeConfig(t, overrideBaseConfig), nil) + require.NoError(t, err) + assert.Equal(t, 8, cfg.Compute.NumAccelerators) + }) + + t.Run("typed scalar override is coerced", func(t *testing.T) { + cfg, err := loadRunConfigWithOverrides(t.Context(), writeConfig(t, overrideBaseConfig), []string{"compute.num_accelerators=16"}) + require.NoError(t, err) + assert.Equal(t, 16, cfg.Compute.NumAccelerators) + }) + + t.Run("multiple overrides all apply", func(t *testing.T) { + cfg, err := loadRunConfigWithOverrides(t.Context(), writeConfig(t, overrideBaseConfig), []string{"compute.num_accelerators=16", "timeout_minutes=45"}) + require.NoError(t, err) + assert.Equal(t, 16, cfg.Compute.NumAccelerators) + require.NotNil(t, cfg.TimeoutMinutes) + assert.Equal(t, 45, *cfg.TimeoutMinutes) + }) + + t.Run("free-form env var adds a key as a string", func(t *testing.T) { + cfg, err := loadRunConfigWithOverrides(t.Context(), writeConfig(t, overrideBaseConfig), []string{"env_variables.RANK=0"}) + require.NoError(t, err) + // A numeric-looking value stays a string because env_variables is map[string]string. + assert.Equal(t, "0", cfg.EnvVariables["RANK"]) + }) + + t.Run("intermediate maps are auto-created", func(t *testing.T) { + cfg, err := loadRunConfigWithOverrides(t.Context(), writeConfig(t, overrideBaseConfig), []string{"environment.docker_image.url=my/img:1"}) + require.NoError(t, err) + require.NotNil(t, cfg.Environment) + require.NotNil(t, cfg.Environment.DockerImage) + assert.Equal(t, "my/img:1", cfg.Environment.DockerImage.URL) + }) + + t.Run("unknown path errors before mutation", func(t *testing.T) { + _, err := loadRunConfigWithOverrides(t.Context(), writeConfig(t, overrideBaseConfig), []string{"bogus=1"}) + require.ErrorContains(t, err, `"bogus" is not a known field`) + }) + + t.Run("semantic validation runs after override", func(t *testing.T) { + // 3 is a known field with a valid type, so only validate() can reject it. + _, err := loadRunConfigWithOverrides(t.Context(), writeConfig(t, overrideBaseConfig), []string{"compute.num_accelerators=3"}) + require.ErrorContains(t, err, "must be a multiple of 8") + }) + + t.Run("type mismatch is rejected on re-decode", func(t *testing.T) { + _, err := loadRunConfigWithOverrides(t.Context(), writeConfig(t, overrideBaseConfig), []string{"compute.num_accelerators=abc"}) + require.Error(t, err) + }) + + t.Run("malformed override is rejected", func(t *testing.T) { + _, err := loadRunConfigWithOverrides(t.Context(), writeConfig(t, overrideBaseConfig), []string{"compute.num_accelerators"}) + require.ErrorContains(t, err, "expected KEY=VALUE") + }) +} diff --git a/experimental/air/cmd/runconfig_test.go b/experimental/air/cmd/runconfig_test.go index 26b54127265..45cae7cd0ae 100644 --- a/experimental/air/cmd/runconfig_test.go +++ b/experimental/air/cmd/runconfig_test.go @@ -243,7 +243,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) { diff --git a/experimental/air/cmd/runsubmit.go b/experimental/air/cmd/runsubmit.go index 8c0be55260d..3aa2436b827 100644 --- a/experimental/air/cmd/runsubmit.go +++ b/experimental/air/cmd/runsubmit.go @@ -2,9 +2,9 @@ package aircmd import ( "context" - "errors" "fmt" "path" + "path/filepath" "strconv" "strings" @@ -38,8 +38,30 @@ func dlRuntimeImage(ctx context.Context, runtimeVersion string) string { return strings.TrimPrefix(img, "CLIENT-GPU-") } +// environmentDependencies resolves the user's declared dependencies as a flat +// list to carry inline on the serverless environment's spec.dependencies: the +// inline list directly, or the dependencies read from a requirements file +// (resolved against the config's directory). For file-form deps it also returns +// the version declared inside that file, which selects the runtime image since +// top-level environment.version is not allowed there. Returns nil when none are +// declared. +func environmentDependencies(cfg *runConfig, configPath string) (deps []string, fileVersion string, err error) { + if deps, ok := cfg.inlineDependencies(); ok { + return deps, "", nil + } + if reqPath, ok := cfg.requirementsFile(); ok { + if !filepath.IsAbs(reqPath) { + reqPath = filepath.Join(filepath.Dir(configPath), reqPath) + } + return readRequirementsDependencies(reqPath) + } + return nil, "", nil +} + // 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 +69,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 +80,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 +102,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, }}, } } @@ -121,14 +145,31 @@ func submitToken(flag string, cfg *runConfig) (string, error) { // 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. + // 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 { - return 0, "", errors.New("usage_policy_name is not yet supported") + usagePolicyID, err = resolveUsagePolicyIDByName(ctx, w, *cfg.UsagePolicyName) + if err != nil { + return 0, "", err + } } - // Resolve the idempotency token first so a bad key fails before any upload. - token, err := submitToken(idempotencyKey, cfg) + // Resolve dependencies before any upload too, so a bad requirements file fails + // fast without leaving orphaned artifacts in the workspace. + deps, fileVersion, err := environmentDependencies(cfg, configPath) if err != nil { return 0, "", err } @@ -163,20 +204,25 @@ func submitWorkload(ctx context.Context, w *databricks.WorkspaceClient, cfg *run return 0, "", err } - // Package and upload the code snapshot, if any. The resulting paths ride on the - // ai_runtime_task; a run with no code_source leaves them empty. Snapshot is the - // only code_source type; guard against a nil block so snapshotCodeSource never - // dereferences a missing snapshot. + // Package and upload the code snapshot, if any, via DABs' artifact-upload + // plumbing; the remote code_source_path rides the ai_runtime_task. A run with no + // code_source leaves it empty. Snapshot is the only code_source type. var snap snapshotResult if cfg.CodeSource != nil && cfg.CodeSource.Snapshot != nil { - snap, err = snapshotCodeSource(ctx, w, cfg.CodeSource.Snapshot, configPath, base, funcDir) + // Sidecars land in the run's launch dir (funcDir) via fc, next to command.sh. + snap, err = snapshotViaDABsUpload(ctx, w, cfg.CodeSource.Snapshot, configPath, fc, funcDir) if err != nil { return 0, "", err } } - runtimeVersion, _ := cfg.runtimeVersion() - payload := buildSubmitPayload(cfg, path.Join(funcDir, commandScriptName), dlRuntimeImage(ctx, runtimeVersion), snap) + // Top-level environment.version wins; for file-form deps it is disallowed, so + // fall back to the version declared inside the requirements file. + runtimeVersion, ok := cfg.runtimeVersion() + if !ok { + runtimeVersion = fileVersion + } + payload := buildSubmitPayload(cfg, path.Join(funcDir, commandScriptName), 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..76ab3aeb3ac 100644 --- a/experimental/air/cmd/runsubmit_test.go +++ b/experimental/air/cmd/runsubmit_test.go @@ -2,10 +2,15 @@ package aircmd import ( "encoding/json" + "io" + "os" + "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" @@ -37,10 +42,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 +79,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 +94,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 +104,88 @@ 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") + } +} + +// TestEnvironmentDependencies covers how declared deps are resolved to a flat list: +// an inline list (no file version), a requirements file (path resolved against the +// config dir, version read from the file), none, and a missing file. +func TestEnvironmentDependencies(t *testing.T) { + inline := &runConfig{Environment: &environmentConfig{ + Dependencies: dependencies{set: true, isList: true, list: []string{"torch", "numpy"}}, + }} + deps, version, err := environmentDependencies(inline, "run.yaml") + require.NoError(t, err) + assert.Equal(t, []string{"torch", "numpy"}, deps) + assert.Empty(t, version) + + dir := t.TempDir() + require.NoError(t, os.WriteFile(filepath.Join(dir, "reqs.yaml"), []byte("version: \"5\"\ndependencies:\n - pandas\n"), 0o600)) + fromFile := &runConfig{Environment: &environmentConfig{ + Dependencies: dependencies{set: true, isList: false, path: "reqs.yaml"}, + }} + deps, version, err = environmentDependencies(fromFile, filepath.Join(dir, "run.yaml")) + require.NoError(t, err) + assert.Equal(t, []string{"pandas"}, deps) + assert.Equal(t, "5", version) + + deps, _, err = environmentDependencies(&runConfig{}, "run.yaml") + require.NoError(t, err) + assert.Nil(t, deps) + + missing := &runConfig{Environment: &environmentConfig{ + Dependencies: dependencies{set: true, isList: false, path: "nope.yaml"}, + }} + _, _, err = environmentDependencies(missing, filepath.Join(dir, "run.yaml")) + require.ErrorContains(t, err, "failed to read requirements file") +} + +// TestReadRequirementsDependencies covers reading a requirements file's dependency +// list and version, with a missing key yielding an empty list and a -r include +// rejected. +func TestReadRequirementsDependencies(t *testing.T) { + dir := t.TempDir() + + reqPath := filepath.Join(dir, "requirements.yaml") + require.NoError(t, os.WriteFile(reqPath, []byte("version: \"5\"\ndependencies:\n - torch==2.3.0\n - numpy\n"), 0o600)) + deps, version, err := readRequirementsDependencies(reqPath) + require.NoError(t, err) + assert.Equal(t, []string{"torch==2.3.0", "numpy"}, deps) + assert.Equal(t, "5", version) + + emptyPath := filepath.Join(dir, "empty.yaml") + require.NoError(t, os.WriteFile(emptyPath, []byte("version: \"5\"\n"), 0o600)) + deps, _, err = readRequirementsDependencies(emptyPath) + require.NoError(t, err) + assert.Empty(t, deps) + + includePath := filepath.Join(dir, "include.yaml") + require.NoError(t, os.WriteFile(includePath, []byte("dependencies:\n - -r other.txt\n"), 0o600)) + _, _, err = readRequirementsDependencies(includePath) + require.ErrorContains(t, err, "requirements-file include") +} + func TestSubmitToken(t *testing.T) { cfg := &runConfig{IdempotencyToken: new("from-config")} @@ -156,8 +245,41 @@ func TestSubmitWorkload(t *testing.T) { assert.Equal(t, jobs.ComputeSpec{AcceleratorType: jobs.ComputeSpecAcceleratorTypeGpu1xH100, AcceleratorCount: 1}, d.Compute) } -// TestSubmitWorkloadWithCodeSource exercises the snapshot path end to end: a -// git-pinned code_source is packaged, uploaded, and its paths attached to the task. +// TestSubmitWorkloadHonorsOverride proves a --override reaches the actual +// runs/submit payload on a real submit, not just dry-run validation: the config +// pins num_accelerators=1, the override bumps it to 4, and the recorded request +// body must carry 4. +func TestSubmitWorkloadHonorsOverride(t *testing.T) { + server := testserver.New(t) + t.Cleanup(server.Close) + + // Register before AddDefaultHandlers: the router is first-wins, so this must + // claim the route ahead of the default jobs/runs/submit handler. + var got jobs.SubmitRun + server.Handle("POST", "/api/2.2/jobs/runs/submit", func(req testserver.Request) any { + require.NoError(t, json.Unmarshal(req.Body, &got)) + return jobs.SubmitRunResponse{RunId: 777} + }) + testserver.AddDefaultHandlers(server) + w, err := databricks.NewWorkspaceClient(&databricks.Config{Host: server.URL, Token: "token"}) + require.NoError(t, err) + + cfgPath := writeConfigFile(t, "run.yaml", minimalConfig) + cfg, err := loadRunConfigWithOverrides(t.Context(), cfgPath, []string{"compute.num_accelerators=4"}) + require.NoError(t, err) + + _, _, err = submitWorkload(t.Context(), w, cfg, cfgPath, "idem-key") + require.NoError(t, err) + + require.Len(t, got.Tasks, 1) + at := got.Tasks[0].AiRuntimeTask + require.NotNil(t, at) + require.Len(t, at.Deployments, 1) + assert.Equal(t, 4, at.Deployments[0].Compute.AcceleratorCount) +} + +// A working-tree code_source is packaged into a tarball, uploaded via DABs' artifact +// plumbing, and its remote code_source_path attached to the submitted task. func TestSubmitWorkloadWithCodeSource(t *testing.T) { server := testserver.New(t) t.Cleanup(server.Close) @@ -172,6 +294,47 @@ func TestSubmitWorkloadWithCodeSource(t *testing.T) { w, err := databricks.NewWorkspaceClient(&databricks.Config{Host: server.URL, Token: "token"}) require.NoError(t, err) + // A plain working-tree directory: packaging is plain-tar. + repo := filepath.Join(t.TempDir(), "src") + writeRepoFile(t, repo, "train.py", "print()") + + cfg := minimalConfig + ` +code_source: + type: snapshot + snapshot: + root_path: ` + repo + ` +` + cfgPath := writeConfigFile(t, "run.yaml", cfg) + loaded, err := loadRunConfig(cfgPath) + require.NoError(t, err) + + // The DABs upload path logs via cmdio; the real `air run` context carries it. + ctx := cmdio.MockDiscard(t.Context()) + _, _, err = submitWorkload(ctx, w, loaded, cfgPath, "idem") + require.NoError(t, err) + + at := got.Tasks[0].AiRuntimeTask + // The tarball is uploaded to the artifact .internal dir and code_source_path + // rewritten to it. + assert.Contains(t, at.CodeSourcePath, "/.air/repo_snapshots/.internal/") + assert.True(t, strings.HasSuffix(at.CodeSourcePath, ".tar.gz"), at.CodeSourcePath) +} + +// A git-pinned code_source is git-archived at the commit, uploaded via DABs' artifact +// plumbing, and its remote code_source_path attached to the submitted task. +func TestSubmitWorkloadWithGitPinnedCodeSource(t *testing.T) { + server := testserver.New(t) + t.Cleanup(server.Close) + + var got jobs.SubmitRun + server.Handle("POST", "/api/2.2/jobs/runs/submit", func(req testserver.Request) any { + require.NoError(t, json.Unmarshal(req.Body, &got)) + return jobs.SubmitRunResponse{RunId: 555} + }) + testserver.AddDefaultHandlers(server) + w, err := databricks.NewWorkspaceClient(&databricks.Config{Host: server.URL, Token: "token"}) + require.NoError(t, err) + // A git repo committed at HEAD, referenced by commit so packaging is git_archive. repo := newTestRepo(t) writeRepoFile(t, repo, "train.py", "print()") @@ -189,28 +352,306 @@ code_source: loaded, err := loadRunConfig(cfgPath) require.NoError(t, err) - _, _, err = submitWorkload(t.Context(), w, loaded, cfgPath, "idem") + ctx := cmdio.MockDiscard(t.Context()) + _, _, err = submitWorkload(ctx, w, loaded, cfgPath, "idem") + require.NoError(t, err) + + at := got.Tasks[0].AiRuntimeTask + assert.Contains(t, at.CodeSourcePath, "/.air/repo_snapshots/.internal/") + assert.True(t, strings.HasSuffix(at.CodeSourcePath, ".tar.gz"), at.CodeSourcePath) +} + +// testSidecarStore builds a workspace filer + base path standing in for the run's +// launch dir, where snapshotViaDABsUpload writes git provenance sidecars. +func testSidecarStore(t *testing.T, w *databricks.WorkspaceClient) (filer.Filer, string) { + t.Helper() + base := "/Workspace/Users/tester@databricks.com/.air/cli_launch/test" + f, err := filer.NewWorkspaceFilesClient(w, base) + require.NoError(t, err) + return f, base +} + +// A plain-tar (working-tree) snapshot is uploaded under a unique, timestamped name so +// two concurrent submissions of the same root_path don't clobber each other's upload. +func TestSubmitWorkloadPlainTarNameIsUnique(t *testing.T) { + server := testserver.New(t) + t.Cleanup(server.Close) + + server.Handle("POST", "/api/2.2/jobs/runs/submit", func(req testserver.Request) any { + return jobs.SubmitRunResponse{RunId: 555} + }) + testserver.AddDefaultHandlers(server) + w, err := databricks.NewWorkspaceClient(&databricks.Config{Host: server.URL, Token: "token"}) + require.NoError(t, err) + + // A plain working-tree directory named "src": the old code named the tarball + // after the dir alone (src.tar.gz), so any two submissions collided. + repo := filepath.Join(t.TempDir(), "src") + writeRepoFile(t, repo, "train.py", "print()") + + cfg := minimalConfig + ` +code_source: + type: snapshot + snapshot: + root_path: ` + repo + ` +` + cfgPath := writeConfigFile(t, "run.yaml", cfg) + loaded, err := loadRunConfig(cfgPath) + require.NoError(t, err) + + // The uploaded name carries a discriminator (timestamp), not the bare dir name. + ctx := cmdio.MockDiscard(t.Context()) + sidecarStore, sidecarBase := testSidecarStore(t, w) + snap, err := snapshotViaDABsUpload(ctx, w, loaded.CodeSource.Snapshot, cfgPath, sidecarStore, sidecarBase) + require.NoError(t, err) + base := path.Base(snap.CodeSourcePath) + assert.NotEqual(t, "src.tar.gz", base, "plain-tar name must be unique, not the bare dir name") + assert.Regexp(t, `^src_\d{8}_\d{6}\.tar\.gz$`, base) +} + +// A git_archive snapshot is content-addressed by (commit, include_paths): submitting +// the same commit twice reuses the already-uploaded tarball and skips the second +// upload (cache hit), while resolving to the identical remote path. +func TestSubmitWorkloadGitArchiveCaching(t *testing.T) { + server := testserver.New(t) + t.Cleanup(server.Close) + + server.Handle("POST", "/api/2.2/jobs/runs/submit", func(req testserver.Request) any { + return jobs.SubmitRunResponse{RunId: 555} + }) + // Track which snapshot tarballs get uploaded, preserving fake-workspace + // persistence so the second submit's cache-existence Stat sees the first upload. + // Dedupe by path: the DABs uploader mkdirs-and-retries the import on a missing + // parent dir, so one logical upload can hit this route more than once. + uploaded := map[string]bool{} + server.Handle("POST", "/api/2.0/workspace-files/import-file/{path...}", func(req testserver.Request) any { + p := req.Vars["path"] + if strings.Contains(p, "/.air/repo_snapshots/") { + uploaded[p] = true + } + return req.Workspace.WorkspaceFilesImportFile(p, req.Body, req.URL.Query().Get("overwrite") == "true") + }) + testserver.AddDefaultHandlers(server) + w, err := databricks.NewWorkspaceClient(&databricks.Config{Host: server.URL, Token: "token"}) + require.NoError(t, err) + + repo := newTestRepo(t) + writeRepoFile(t, repo, "train.py", "print()") + sha := commitAll(t, repo, "init") + + cfg := minimalConfig + ` +code_source: + type: snapshot + snapshot: + root_path: ` + repo + ` + git: + commit: ` + sha + ` +` + cfgPath := writeConfigFile(t, "run.yaml", cfg) + loaded, err := loadRunConfig(cfgPath) + require.NoError(t, err) + + ctx := cmdio.MockDiscard(t.Context()) + sidecarStore, sidecarBase := testSidecarStore(t, w) + first, err := snapshotViaDABsUpload(ctx, w, loaded.CodeSource.Snapshot, cfgPath, sidecarStore, sidecarBase) + require.NoError(t, err) + second, err := snapshotViaDABsUpload(ctx, w, loaded.CodeSource.Snapshot, cfgPath, sidecarStore, sidecarBase) + require.NoError(t, err) + + // Same pinned commit → identical content-addressed remote path, uploaded once + // (the second submit is a cache hit and moves no bytes). + assert.Equal(t, first.CodeSourcePath, second.CodeSourcePath) + assert.Len(t, uploaded, 1, "git_archive cache hit should skip the second upload") +} + +// A git code_source also uploads git provenance sidecars (git_state.json, and +// git_diff.patch when the tree is dirty) next to the run's launch dir, so the +// submitted commit + working-tree diff are inspectable. +func TestSubmitWorkloadUploadsGitSidecars(t *testing.T) { + server := testserver.New(t) + t.Cleanup(server.Close) + + server.Handle("POST", "/api/2.2/jobs/runs/submit", func(req testserver.Request) any { + return jobs.SubmitRunResponse{RunId: 555} + }) + testserver.AddDefaultHandlers(server) + w, err := databricks.NewWorkspaceClient(&databricks.Config{Host: server.URL, Token: "token"}) + require.NoError(t, err) + + // Commit, then dirty the tree so both git_state.json and git_diff.patch are produced. + repo := newTestRepo(t) + writeRepoFile(t, repo, "train.py", "print()") + commitAll(t, repo, "init") + writeRepoFile(t, repo, "train.py", "print('dirty')") + + cfg := minimalConfig + ` +code_source: + type: snapshot + snapshot: + root_path: ` + repo + ` +` + cfgPath := writeConfigFile(t, "run.yaml", cfg) + loaded, err := loadRunConfig(cfgPath) + require.NoError(t, err) + + ctx := cmdio.MockDiscard(t.Context()) + sidecarStore, sidecarBase := testSidecarStore(t, w) + snap, err := snapshotViaDABsUpload(ctx, w, loaded.CodeSource.Snapshot, cfgPath, sidecarStore, sidecarBase) + require.NoError(t, err) + + // Both sidecars are reported under the launch dir and actually exist there. + assert.Equal(t, path.Join(sidecarBase, gitStateName), snap.GitStatePath) + assert.Equal(t, path.Join(sidecarBase, gitDiffName), snap.GitDiffPath) + + r, err := sidecarStore.Read(ctx, gitStateName) + require.NoError(t, err) + stateBytes, err := io.ReadAll(r) + require.NoError(t, err) + var state map[string]any + require.NoError(t, json.Unmarshal(stateBytes, &state)) + assert.Equal(t, "plain_tar", state["packaging_mode"]) + assert.Equal(t, true, state["dirty"]) + assert.Equal(t, "captured", state["diff_status"]) +} + +// remote_volume uploads the snapshot to a UC Volume: DABs' artifact uploader handles +// /Volumes destinations natively, so code_source_path lands under the Volume path. +func TestSubmitWorkloadWithRemoteVolumeCodeSource(t *testing.T) { + server := testserver.New(t) + t.Cleanup(server.Close) + + var got jobs.SubmitRun + server.Handle("POST", "/api/2.2/jobs/runs/submit", func(req testserver.Request) any { + require.NoError(t, json.Unmarshal(req.Body, &got)) + return jobs.SubmitRunResponse{RunId: 555} + }) + // Stub the UC Volume file write: the fake server's default handler 404s when the + // parent dir is absent (no auto-mkdir), so accept the PUT to exercise the Volume + // upload route. This asserts we route to /api/2.0/fs/files/Volumes/... at all. + server.Handle("PUT", "/api/2.0/fs/files/Volumes/{path...}", func(req testserver.Request) any { + return testserver.Response{StatusCode: 204} + }) + testserver.AddDefaultHandlers(server) + w, err := databricks.NewWorkspaceClient(&databricks.Config{Host: server.URL, Token: "token"}) + require.NoError(t, err) + + repo := filepath.Join(t.TempDir(), "src") + writeRepoFile(t, repo, "train.py", "print()") + + cfg := minimalConfig + ` +code_source: + type: snapshot + snapshot: + root_path: ` + repo + ` + remote_volume: /Volumes/main/default/code +` + cfgPath := writeConfigFile(t, "run.yaml", cfg) + loaded, err := loadRunConfig(cfgPath) + require.NoError(t, err) + + ctx := cmdio.MockDiscard(t.Context()) + _, _, err = submitWorkload(ctx, w, loaded, cfgPath, "idem") require.NoError(t, err) at := got.Tasks[0].AiRuntimeTask - // The tarball path is under the user's repo_snapshots dir. git_state_path / - // git_diff_path are not asserted: the typed jobs.AiRuntimeTask has no such fields - // (see the TEMP note in buildSubmitPayload), so they aren't sent. The git_state - // sidecar file is still uploaded next to the tarball — covered by TestRunSnapshot. - assert.Contains(t, at.CodeSourcePath, "/.air/repo_snapshots/"+filepath.Base(repo)+"/") + assert.Contains(t, at.CodeSourcePath, "/Volumes/main/default/code/.internal/") assert.True(t, strings.HasSuffix(at.CodeSourcePath, ".tar.gz"), at.CodeSourcePath) } 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} + }) + testserver.AddDefaultHandlers(server) + pw, err := databricks.NewWorkspaceClient(&databricks.Config{Host: server.URL, Token: "token"}) + require.NoError(t, err) + + cfg := *base + cfg.UsagePolicyName = new("nope") + _, _, err = submitWorkload(t.Context(), pw, &cfg, cfgPath, "") + 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") + } + }) + + t.Run("bad requirements file fails before any upload", func(t *testing.T) { + server := testserver.New(t) + t.Cleanup(server.Close) + var uploaded bool + server.Handle("POST", "/api/2.0/workspace-files/import-file/{path...}", func(testserver.Request) any { + uploaded = true + return nil + }) + testserver.AddDefaultHandlers(server) + tw, 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.Environment = &environmentConfig{Dependencies: dependencies{set: true, isList: false, path: "missing.yaml"}} + _, _, err = submitWorkload(t.Context(), tw, &cfg, cfgPath, "") + require.ErrorContains(t, err, "failed to read requirements file") + assert.False(t, uploaded, "no artifacts should be uploaded when dependency resolution fails") + }) +} + +// 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"}}} + }) + 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") + 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") + 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..a6e68655d5d 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,40 @@ 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. +// requirementsDoc mirrors the on-disk requirements.yaml format, used to read a +// user-provided requirements file's dependencies for the inline spec.dependencies. type requirementsDoc struct { Version string `yaml:"version,omitempty"` Dependencies []string `yaml:"dependencies"` } +// readRequirementsDependencies reads the dependencies and version out of a +// requirements.yaml file so file-form deps can be carried on the serverless +// environment's inline spec.dependencies and its version can select the runtime +// image. Returns an empty list when the file declares no dependencies. +func readRequirementsDependencies(reqPath string) ([]string, string, error) { + data, err := os.ReadFile(reqPath) + if err != nil { + return nil, "", fmt.Errorf("failed to read requirements file %s: %w", reqPath, err) + } + var doc requirementsDoc + if err := yaml.Unmarshal(data, &doc); err != nil { + return nil, "", fmt.Errorf("failed to parse requirements file %s: %w", reqPath, err) + } + for _, dep := range doc.Dependencies { + if fields := strings.Fields(dep); len(fields) > 0 && (fields[0] == "-r" || fields[0] == "--requirement") { + return nil, "", fmt.Errorf("requirements file dependency %q uses a requirements-file include (-r/--requirement), which is not supported; list the dependencies directly instead", dep) + } + } + return doc.Dependencies, doc.Version, nil +} + // 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 +94,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..e57a8aa2c41 100644 --- a/experimental/air/cmd/runupload_test.go +++ b/experimental/air/cmd/runupload_test.go @@ -57,7 +57,7 @@ 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"), @@ -68,19 +68,10 @@ func TestBuildArtifacts_InlineRequirementsAndParameters(t *testing.T) { 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,7 +94,7 @@ 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) { +func TestBuildArtifacts_RequirementsFileNotUploaded(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)) @@ -112,9 +103,10 @@ func TestBuildArtifacts_RequirementsFile(t *testing.T) { Environment: &environmentConfig{Dependencies: dependencies{set: true, isList: false, path: "reqs.yaml"}}, } + // A declared requirements file is not uploaded; its deps travel on spec.dependencies. items, err := buildArtifacts(cfg, filepath.Join(dir, "run.yaml")) require.NoError(t, err) - assert.Contains(t, itemNames(items), requirementsName) + assert.Equal(t, []string{trainingConfigName, commandScriptName}, itemNames(items)) } func TestBuildArtifacts_OversizeConfigRejected(t *testing.T) { @@ -143,13 +135,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") + }) +}