diff --git a/acceptance/experimental/air/cancel/out.test.toml b/acceptance/experimental/air/cancel/out.test.toml new file mode 100644 index 00000000000..d6187dcb046 --- /dev/null +++ b/acceptance/experimental/air/cancel/out.test.toml @@ -0,0 +1,3 @@ +Local = true +Cloud = false +EnvMatrix.DATABRICKS_BUNDLE_ENGINE = [] diff --git a/acceptance/experimental/air/cancel/output.txt b/acceptance/experimental/air/cancel/output.txt new file mode 100644 index 00000000000..9fd8a055f13 --- /dev/null +++ b/acceptance/experimental/air/cancel/output.txt @@ -0,0 +1,29 @@ + +=== cancel by id (text) +>>> [CLI] experimental air cancel 123 +Successfully requested cancellation for run 123 + +=== cancel by id (json) +>>> [CLI] experimental air cancel 123 -o json +{ + "v": 1, + "ts": "[TIMESTAMP]", + "data": { + "cancelled": [ + "123" + ] + } +} + +=== cancel multiple ids +>>> [CLI] experimental air cancel 123 456 +Successfully requested cancellation for run 123 +Successfully requested cancellation for run 456 +Successfully requested cancellation for 2 run(s). + +=== cancel --all +>>> [CLI] experimental air cancel --all -y +Searching active runs for [USERNAME] in [DATABRICKS_URL]... +Successfully requested cancellation for run [NUMID] +Successfully requested cancellation for run [NUMID] +Successfully requested cancellation for 2 run(s). diff --git a/acceptance/experimental/air/cancel/script b/acceptance/experimental/air/cancel/script new file mode 100644 index 00000000000..ce04a8977fd --- /dev/null +++ b/acceptance/experimental/air/cancel/script @@ -0,0 +1,11 @@ +title "cancel by id (text)" +trace $CLI experimental air cancel 123 + +title "cancel by id (json)" +trace $CLI experimental air cancel 123 -o json + +title "cancel multiple ids" +trace $CLI experimental air cancel 123 456 + +title "cancel --all" +trace $CLI experimental air cancel --all -y diff --git a/acceptance/experimental/air/cancel/test.toml b/acceptance/experimental/air/cancel/test.toml new file mode 100644 index 00000000000..c73594501d8 --- /dev/null +++ b/acceptance/experimental/air/cancel/test.toml @@ -0,0 +1,53 @@ +# This command does not deploy a bundle, so no engine matrix is needed. +[EnvMatrix] +DATABRICKS_BUNDLE_ENGINE = [] + +# The SDK occasionally probes host reachability with a HEAD request; stub it so +# the test is deterministic. +[[Server]] +Pattern = "HEAD /" +Response.Body = '' + +# CancelRun accepts the request and returns an empty body. +[[Server]] +Pattern = "POST /api/2.2/jobs/runs/cancel" +Response.Body = '{}' + +# Jobs runs/list backs `cancel --all`: two active AIR runs for the current user +# (tester@databricks.com, from the built-in scim/v2/Me handler). +[[Server]] +Pattern = "GET /api/2.2/jobs/runs/list" +Response.Body = ''' +{ + "runs": [ + { + "run_id": 334747067049496, + "run_name": "qwen-train", + "creator_user_name": "tester@databricks.com", + "start_time": 1717608759000, + "state": {"life_cycle_state": "RUNNING"}, + "tasks": [{ + "run_id": 334747067049497, + "ai_runtime_task": { + "experiment": "/Users/tester@databricks.com/qwen-train", + "deployments": [{"compute": {"accelerator_type": "GPU_1xA10", "accelerator_count": 1}}] + } + }] + }, + { + "run_id": 566001814929041, + "run_name": "llama-train", + "creator_user_name": "tester@databricks.com", + "start_time": 1717612404000, + "state": {"life_cycle_state": "RUNNING"}, + "tasks": [{ + "run_id": 566001814929042, + "ai_runtime_task": { + "experiment": "/Users/tester@databricks.com/llama-train", + "deployments": [{"compute": {"accelerator_type": "GPU_1xA10", "accelerator_count": 1}}] + } + }] + } + ] +} +''' diff --git a/acceptance/experimental/air/get-ai-runtime/out.test.toml b/acceptance/experimental/air/get-ai-runtime/out.test.toml new file mode 100644 index 00000000000..d6187dcb046 --- /dev/null +++ b/acceptance/experimental/air/get-ai-runtime/out.test.toml @@ -0,0 +1,3 @@ +Local = true +Cloud = false +EnvMatrix.DATABRICKS_BUNDLE_ENGINE = [] diff --git a/acceptance/experimental/air/get-ai-runtime/output.txt b/acceptance/experimental/air/get-ai-runtime/output.txt new file mode 100644 index 00000000000..c47790eaa31 --- /dev/null +++ b/acceptance/experimental/air/get-ai-runtime/output.txt @@ -0,0 +1,52 @@ + +=== get (text) +>>> [CLI] experimental air get 123 + +╭─ Configuration ────────────────────────────────────────────────╮ +│ │ +│ experiment_name: my-exp │ +│ compute: │ +│ accelerator_type: a10 │ +│ num_accelerators: 1 │ +│ command: |- │ +│ for i in $(seq 1 10); do │ +│ echo "step $i" │ +│ done │ +│ │ +╰────────────────────────────────────────────────────────────────╯ + +╭─ Metadata ─────────────────────────────────────────────────────╮ +│ │ +│ Run ID 123 │ +│ Status ● SUCCESS │ +│ Submitted 2023-11-14 22:13 UTC │ +│ Retries 0 │ +│ Max Retries 3 │ +│ Duration 12s │ +│ Experiment my-exp │ +│ MLflow Run my-run │ +│ User user@example.com │ +│ Accelerators 1x A10 │ +│ Environment N/A │ +│ │ +╰────────────────────────────────────────────────────────────────╯ + +Run URL: [DATABRICKS_URL]/jobs/runs/123?o=[NUMID] +MLflow URL: [DATABRICKS_URL]/ml/experiments/exp1/runs/run1 + +=== get (json) +>>> [CLI] experimental air get 123 -o json +{ + "v": 1, + "ts": "[TIMESTAMP]", + "data": { + "run_id": "123", + "status": "SUCCESS", + "started_at": "[TIMESTAMP]", + "duration_seconds": 12, + "attempt_number": 0, + "experiment_name": "my-exp", + "dashboard_url": "[DATABRICKS_URL]/jobs/runs/123?o=[NUMID]", + "mlflow_url": "[DATABRICKS_URL]/ml/experiments/exp1/runs/run1/artifacts/logs/node_0" + } +} diff --git a/acceptance/experimental/air/get-ai-runtime/script b/acceptance/experimental/air/get-ai-runtime/script new file mode 100644 index 00000000000..3f41b089cdd --- /dev/null +++ b/acceptance/experimental/air/get-ai-runtime/script @@ -0,0 +1,11 @@ +# Seed the run's training_config.yaml next to command.sh so `air get` can +# download and render it in the Configuration box. The import API does not +# create missing parents, so mkdirs the containing folder first. +$CLI workspace mkdirs "/Workspace/Users/user@example.com/.air/cli_launch/my-exp/my-exp_abc" &> LOG.mkdirs +$CLI workspace import "/Workspace/Users/user@example.com/.air/cli_launch/my-exp/my-exp_abc/training_config.yaml" --file training_config.yaml --format AUTO &> LOG.import + +title "get (text)" +trace $CLI experimental air get 123 + +title "get (json)" +trace $CLI experimental air get 123 -o json diff --git a/acceptance/experimental/air/get-ai-runtime/test.toml b/acceptance/experimental/air/get-ai-runtime/test.toml new file mode 100644 index 00000000000..d5949a3023e --- /dev/null +++ b/acceptance/experimental/air/get-ai-runtime/test.toml @@ -0,0 +1,60 @@ +# This command does not deploy a bundle, so no engine matrix is needed. +[EnvMatrix] +DATABRICKS_BUNDLE_ENGINE = [] + +# On Windows, Git Bash rewrites the leading-/ workspace paths passed to +# `workspace mkdirs`/`import` into C:/... paths; disable that conversion. +[Env] +MSYS_NO_PATHCONV = "1" + +# The SDK occasionally probes host reachability with a HEAD request; stub it so +# the test is deterministic. +[[Server]] +Pattern = "HEAD /" +Response.Body = '' + +# The typed SDK GetRun response: an ai_runtime_task run has no gen_ai_compute_task, +# so the task comes back empty (the SDK has no field for ai_runtime_task). +[[Server]] +Pattern = "GET /api/2.2/jobs/runs/get" +Response.Body = ''' +{ + "run_id": 123, + "run_page_url": "https://my-workspace.cloud.databricks.test/jobs/runs/123", + "creator_user_name": "user@example.com", + "start_time": 1700000000000, + "end_time": 1700000012000, + "state": {"life_cycle_state": "TERMINATED", "result_state": "SUCCESS"}, + "tasks": [ + { + "task_key": "train", + "run_id": 456, + "attempt_number": 0, + "max_retries": 3, + "ai_runtime_task": { + "experiment": "my-exp", + "deployments": [ + { + "command_path": "/Workspace/Users/user@example.com/.air/cli_launch/my-exp/my-exp_abc/command.sh", + "compute": {"accelerator_type": "GPU_1xA10", "accelerator_count": 1} + } + ] + } + } + ] +} +''' + +# MLflow identifiers for the deep-link, under ai_runtime_task_output. +[[Server]] +Pattern = "GET /api/2.2/jobs/runs/get-output" +Response.Body = ''' +{"ai_runtime_task_output": {"mlflow_experiment_id": "exp1", "mlflow_run_id": "run1"}} +''' + +# The MLflow Run cell shows the run's name, fetched from the MLflow REST API. +[[Server]] +Pattern = "GET /api/2.0/mlflow/runs/get" +Response.Body = ''' +{"run": {"info": {"run_name": "my-run"}}} +''' diff --git a/acceptance/experimental/air/get-ai-runtime/training_config.yaml b/acceptance/experimental/air/get-ai-runtime/training_config.yaml new file mode 100644 index 00000000000..5f6060ecbbc --- /dev/null +++ b/acceptance/experimental/air/get-ai-runtime/training_config.yaml @@ -0,0 +1,8 @@ +experiment_name: my-exp +compute: + accelerator_type: a10 + num_accelerators: 1 +command: |- + for i in $(seq 1 10); do + echo "step $i" + done diff --git a/acceptance/experimental/air/get/out.test.toml b/acceptance/experimental/air/get/out.test.toml new file mode 100644 index 00000000000..d6187dcb046 --- /dev/null +++ b/acceptance/experimental/air/get/out.test.toml @@ -0,0 +1,3 @@ +Local = true +Cloud = false +EnvMatrix.DATABRICKS_BUNDLE_ENGINE = [] diff --git a/acceptance/experimental/air/get/output.txt b/acceptance/experimental/air/get/output.txt new file mode 100644 index 00000000000..6e51d7debf9 --- /dev/null +++ b/acceptance/experimental/air/get/output.txt @@ -0,0 +1,73 @@ + +=== get (text) +>>> [CLI] experimental air get 123 + +╭─ Configuration ────────────────────────────────────────────────╮ +│ │ +│ experiment_name: my-exp │ +│ compute: │ +│ accelerator_type: a10 │ +│ num_accelerators: 1 │ +│ command: | │ +│ for i in $(seq 1 10); do │ +│ echo "step $i" │ +│ done │ +│ │ +╰────────────────────────────────────────────────────────────────╯ + +╭─ Metadata ─────────────────────────────────────────────────────╮ +│ │ +│ Run ID 123 │ +│ Status ● SUCCESS │ +│ Submitted 2023-11-14 22:13 UTC │ +│ Retries 0 │ +│ Max Retries 3 │ +│ Duration 12s │ +│ Experiment my-exp │ +│ MLflow Run my-run │ +│ User user@example.com │ +│ Accelerators 1x A10 │ +│ Environment ml-runtime-gpu:1.0 │ +│ │ +╰────────────────────────────────────────────────────────────────╯ + +Run URL: [DATABRICKS_URL]/jobs/runs/123?o=[NUMID] +MLflow URL: [DATABRICKS_URL]/ml/experiments/exp1/runs/run1 + +=== get (json) +>>> [CLI] experimental air get 123 -o json +{ + "v": 1, + "ts": "[TIMESTAMP]", + "data": { + "run_id": "123", + "status": "SUCCESS", + "started_at": "[TIMESTAMP]", + "duration_seconds": 12, + "attempt_number": 0, + "experiment_name": "my-exp", + "dashboard_url": "[DATABRICKS_URL]/jobs/runs/123?o=[NUMID]", + "mlflow_url": "[DATABRICKS_URL]/ml/experiments/exp1/runs/run1/artifacts/logs/node_0" + } +} + +=== invalid run id +>>> [CLI] experimental air get notanumber +Error: invalid JOB_RUN_ID "notanumber": must be a positive integer + +Exit code: 1 + +=== invalid run id (json) +>>> [CLI] experimental air get notanumber -o json +{ + "v": 1, + "ts": "[TIMESTAMP]", + "error": { + "code": "INVALID_ARGS", + "kind": "PERMANENT", + "message": "invalid JOB_RUN_ID \"notanumber\": must be a positive integer", + "retryable": false + } +} + +Exit code: 1 diff --git a/acceptance/experimental/air/get/script b/acceptance/experimental/air/get/script new file mode 100644 index 00000000000..ee66b4aff04 --- /dev/null +++ b/acceptance/experimental/air/get/script @@ -0,0 +1,11 @@ +title "get (text)" +trace $CLI experimental air get 123 + +title "get (json)" +trace $CLI experimental air get 123 -o json + +title "invalid run id" +errcode trace $CLI experimental air get notanumber + +title "invalid run id (json)" +errcode trace $CLI experimental air get notanumber -o json diff --git a/acceptance/experimental/air/get/test.toml b/acceptance/experimental/air/get/test.toml new file mode 100644 index 00000000000..3f35ddfe658 --- /dev/null +++ b/acceptance/experimental/air/get/test.toml @@ -0,0 +1,51 @@ +# This command does not deploy a bundle, so no engine matrix is needed. +[EnvMatrix] +DATABRICKS_BUNDLE_ENGINE = [] + +# The SDK occasionally probes host reachability with a HEAD request; stub it so +# the test is deterministic. +[[Server]] +Pattern = "HEAD /" +Response.Body = '' + +# A single GenAI-compute run with an experiment, GPUs, and a creator. +[[Server]] +Pattern = "GET /api/2.2/jobs/runs/get" +Response.Body = ''' +{ + "run_id": 123, + "run_page_url": "https://my-workspace.cloud.databricks.test/jobs/runs/123", + "creator_user_name": "user@example.com", + "start_time": 1700000000000, + "end_time": 1700000012000, + "state": {"life_cycle_state": "TERMINATED", "result_state": "SUCCESS"}, + "tasks": [ + { + "task_key": "train", + "run_id": 456, + "attempt_number": 0, + "max_retries": 3, + "gen_ai_compute_task": { + "mlflow_experiment_name": "/Users/user@example.com/my-exp", + "compute": {"gpu_type": "GPU_1xA10", "num_gpus": 1}, + "dl_runtime_image": "ml-runtime-gpu:1.0", + "yaml_parameters": "experiment_name: my-exp\ncompute:\n accelerator_type: a10\n num_accelerators: 1\ncommand: |\n for i in $(seq 1 10); do\n echo \"step $i\"\n done\n" + } + } + ] +} +''' + +# MLflow identifiers for the deep-link, under ai_runtime_task_output. +[[Server]] +Pattern = "GET /api/2.2/jobs/runs/get-output" +Response.Body = ''' +{"ai_runtime_task_output": {"mlflow_experiment_id": "exp1", "mlflow_run_id": "run1"}} +''' + +# The MLflow Run cell shows the run's name, fetched from the MLflow REST API. +[[Server]] +Pattern = "GET /api/2.0/mlflow/runs/get" +Response.Body = ''' +{"run": {"info": {"run_name": "my-run"}}} +''' diff --git a/acceptance/experimental/air/help/out.test.toml b/acceptance/experimental/air/help/out.test.toml new file mode 100644 index 00000000000..d6187dcb046 --- /dev/null +++ b/acceptance/experimental/air/help/out.test.toml @@ -0,0 +1,3 @@ +Local = true +Cloud = false +EnvMatrix.DATABRICKS_BUNDLE_ENGINE = [] diff --git a/acceptance/experimental/air/help/output.txt b/acceptance/experimental/air/help/output.txt new file mode 100644 index 00000000000..ee89e778d6b --- /dev/null +++ b/acceptance/experimental/air/help/output.txt @@ -0,0 +1,49 @@ + +=== help +>>> [CLI] experimental air --help +Run and manage AI runtime training workloads on Databricks serverless GPU compute. + +This command set is the Go port of the standalone Python "air" CLI. It is +experimental and may change in future versions. + +Usage: + databricks experimental air [command] + +Available Commands: + cancel Cancel one or more runs + get Show status, configuration, and timing details for a specific run + list List your active runs for the current profile (use --all-status for finished runs) + logs Stream or fetch logs for a run + register-image Mirror a Docker image into the workspace registry + run Submit a training workload from a YAML config + +Flags: + -h, --help help for air + +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) + +Use "databricks experimental air [command] --help" for more information about a command. + +=== list help +>>> [CLI] experimental air list --help +List your active runs for the current profile (use --all-status for finished runs) + +Usage: + databricks experimental air list [flags] + +Flags: + --all-status Show runs in all states (default: active only) + --all-users Show runs from all users + --filter stringArray Filter runs, e.g. experiment=foo* (repeatable) + -h, --help help for list + --limit int Maximum number of runs to show (default 20) + +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 new file mode 100644 index 00000000000..81f3907e4f5 --- /dev/null +++ b/acceptance/experimental/air/help/script @@ -0,0 +1,8 @@ +# Pin the command tree so any change to a subcommand or its short description +# shows up as a diff here. + +title "help" +trace $CLI experimental air --help + +title "list help" +trace $CLI experimental air list --help diff --git a/acceptance/experimental/air/help/test.toml b/acceptance/experimental/air/help/test.toml new file mode 100644 index 00000000000..49709b578ef --- /dev/null +++ b/acceptance/experimental/air/help/test.toml @@ -0,0 +1,3 @@ +# --help prints without authenticating, so no server stubs are needed. +[EnvMatrix] +DATABRICKS_BUNDLE_ENGINE = [] diff --git a/acceptance/experimental/air/list/out.test.toml b/acceptance/experimental/air/list/out.test.toml new file mode 100644 index 00000000000..d6187dcb046 --- /dev/null +++ b/acceptance/experimental/air/list/out.test.toml @@ -0,0 +1,3 @@ +Local = true +Cloud = false +EnvMatrix.DATABRICKS_BUNDLE_ENGINE = [] diff --git a/acceptance/experimental/air/list/output.txt b/acceptance/experimental/air/list/output.txt new file mode 100644 index 00000000000..e27cfb0bf39 --- /dev/null +++ b/acceptance/experimental/air/list/output.txt @@ -0,0 +1,48 @@ + +=== list (text) +>>> [CLI] experimental air list + Run ID Experiment Status Started Duration MLflow User Accelerators + [NUMID] qwen-train ● SUCCESS [TIMESTAMP] 12s …/runs/run1 [USERNAME] 8x H100 + +=== list (json) +>>> [CLI] experimental air list -o json +{ + "v": 1, + "ts": "[TIMESTAMP]", + "data": { + "runs": [ + { + "run_id": "[NUMID]", + "run_name": "qwen-train", + "user": "[USERNAME]", + "status": "SUCCESS", + "started_at": "[TIMESTAMP]", + "is_sweep": false + } + ] + } +} + +=== list --all-status (text, via AiTrainingService index) +>>> [CLI] experimental air list --all-status + Run ID Experiment Status Started Duration MLflow User Accelerators + [NUMID] qwen-train ● SUCCESS [TIMESTAMP] 12s …/runs/run1 [USERNAME] 8x H100 + +=== list --all-status (json) +>>> [CLI] experimental air list --all-status -o json +{ + "v": 1, + "ts": "[TIMESTAMP]", + "data": { + "runs": [ + { + "run_id": "[NUMID]", + "run_name": "qwen-train", + "user": "[USERNAME]", + "status": "SUCCESS", + "started_at": "[TIMESTAMP]", + "is_sweep": false + } + ] + } +} diff --git a/acceptance/experimental/air/list/script b/acceptance/experimental/air/list/script new file mode 100644 index 00000000000..df547794c6f --- /dev/null +++ b/acceptance/experimental/air/list/script @@ -0,0 +1,11 @@ +title "list (text)" +trace $CLI experimental air list + +title "list (json)" +trace $CLI experimental air list -o json + +title "list --all-status (text, via AiTrainingService index)" +trace $CLI experimental air list --all-status + +title "list --all-status (json)" +trace $CLI experimental air list --all-status -o json diff --git a/acceptance/experimental/air/list/test.toml b/acceptance/experimental/air/list/test.toml new file mode 100644 index 00000000000..82f1f829d85 --- /dev/null +++ b/acceptance/experimental/air/list/test.toml @@ -0,0 +1,87 @@ +# This command does not deploy a bundle, so no engine matrix is needed. +[EnvMatrix] +DATABRICKS_BUNDLE_ENGINE = [] + +# Disable the on-disk run cache so --all-status output is deterministic across runs. +[Env] +DATABRICKS_CACHE_ENABLED = "false" + +# The SDK occasionally probes host reachability with a HEAD request; stub it so +# the test is deterministic. +[[Server]] +Pattern = "HEAD /" +Response.Body = '' + +# `air list` reads Jobs runs/list directly (expand_tasks), filtering to AIR runs. +# The current runs carry an ai_runtime_task (not modeled by the typed SDK), which +# is why the CLI parses the response raw. Both runs belong to the default user +# (tester@databricks.com, from the built-in scim/v2/Me handler). The second run +# has no AI task and must be filtered out. +[[Server]] +Pattern = "GET /api/2.2/jobs/runs/list" +Response.Body = ''' +{ + "runs": [ + { + "run_id": 334747067049496, + "run_name": "qwen-train", + "creator_user_name": "tester@databricks.com", + "start_time": 1717608759000, + "end_time": 1717608771000, + "state": {"life_cycle_state": "TERMINATED", "result_state": "SUCCESS"}, + "tasks": [{ + "run_id": 334747067049497, + "ai_runtime_task": { + "experiment": "/Users/tester@databricks.com/qwen-train", + "deployments": [{"compute": {"accelerator_type": "GPU_8xH100", "accelerator_count": 8}}] + } + }] + }, + { + "run_id": 999000999000, + "run_name": "not-an-air-run", + "creator_user_name": "tester@databricks.com", + "state": {"life_cycle_state": "RUNNING"}, + "tasks": [{"notebook_task": {"notebook_path": "/x"}}] + } + ] +} +''' + +# MLflow IDs for the deep link, fetched per AIR run (text mode). The current +# task type resolves them under ai_runtime_task_output. +[[Server]] +Pattern = "GET /api/2.2/jobs/runs/get-output" +Response.Body = ''' +{"ai_runtime_task_output": {"mlflow_experiment_id": "exp1", "mlflow_run_id": "run1"}} +''' + +# `air list --all-status` scoped to the current user is served by the +# AiTrainingService index: it returns cheap (job_run_id, submit_time) pairs, which +# the CLI orders and then hydrates into full runs via runs/get. +[[Server]] +Pattern = "GET /api/2.0/ai-training/workflows" +Response.Body = ''' +{"training_workflows": [{"job_run_id": "334747067049496", "submit_time": "2024-06-05T17:32:39Z"}]} +''' + +# runs/get hydrates one index id into the same shape as a runs/list element. +[[Server]] +Pattern = "GET /api/2.2/jobs/runs/get" +Response.Body = ''' +{ + "run_id": 334747067049496, + "run_name": "qwen-train", + "creator_user_name": "tester@databricks.com", + "start_time": 1717608759000, + "end_time": 1717608771000, + "state": {"life_cycle_state": "TERMINATED", "result_state": "SUCCESS"}, + "tasks": [{ + "run_id": 334747067049497, + "ai_runtime_task": { + "experiment": "/Users/tester@databricks.com/qwen-train", + "deployments": [{"compute": {"accelerator_type": "GPU_8xH100", "accelerator_count": 8}}] + } + }] +} +''' diff --git a/acceptance/experimental/air/run-submit/databricks.yml b/acceptance/experimental/air/run-submit/databricks.yml new file mode 100644 index 00000000000..4a1c612d600 --- /dev/null +++ b/acceptance/experimental/air/run-submit/databricks.yml @@ -0,0 +1,2 @@ +bundle: + name: air-run-submit diff --git a/acceptance/experimental/air/run-submit/out.test.toml b/acceptance/experimental/air/run-submit/out.test.toml new file mode 100644 index 00000000000..d6187dcb046 --- /dev/null +++ b/acceptance/experimental/air/run-submit/out.test.toml @@ -0,0 +1,3 @@ +Local = true +Cloud = false +EnvMatrix.DATABRICKS_BUNDLE_ENGINE = [] diff --git a/acceptance/experimental/air/run-submit/output.txt b/acceptance/experimental/air/run-submit/output.txt new file mode 100644 index 00000000000..8d52eed1dab --- /dev/null +++ b/acceptance/experimental/air/run-submit/output.txt @@ -0,0 +1,46 @@ + +=== submit with a git code_source +>>> [CLI] experimental air run -f run.yaml +Submitted run 555 +View at: [DATABRICKS_URL]/jobs/runs/555 + +=== the ai_runtime_task carries the code_source_path +>>> print_requests.py //api/2.2/jobs/runs/submit +{ + "method": "POST", + "path": "/api/2.2/jobs/runs/submit", + "body": { + "environments": [ + { + "environment_key": "default", + "spec": { + "environment_version": "4" + } + } + ], + "idempotency_token": "[UUID]", + "run_name": "submit-smoke", + "tasks": [ + { + "ai_runtime_task": { + "code_source_path": "/Workspace/Users/[USERNAME]/.air/repo_snapshots/001/[SNAPSHOT_TARBALL]", + "deployments": [ + { + "command_path": "/Workspace/Users/[USERNAME]/.air/cli_launch/submit-smoke/submit-smoke_[RUN_ID]/command.sh", + "compute": { + "accelerator_count": 1, + "accelerator_type": "GPU_1xH100" + } + } + ], + "experiment": "submit-smoke" + }, + "environment_key": "default", + "max_retries": 3, + "retry_on_timeout": true, + "run_if": "ALL_SUCCESS", + "task_key": "submit-smoke" + } + ] + } +} diff --git a/acceptance/experimental/air/run-submit/run.yaml.tmpl b/acceptance/experimental/air/run-submit/run.yaml.tmpl new file mode 100644 index 00000000000..3fdbf48eb85 --- /dev/null +++ b/acceptance/experimental/air/run-submit/run.yaml.tmpl @@ -0,0 +1,11 @@ +experiment_name: submit-smoke +command: python train.py +compute: + accelerator_type: GPU_1xH100 + num_accelerators: 1 +code_source: + type: snapshot + snapshot: + root_path: . + git: + commit: COMMIT_SHA diff --git a/acceptance/experimental/air/run-submit/script b/acceptance/experimental/air/run-submit/script new file mode 100644 index 00000000000..c44b50e28a1 --- /dev/null +++ b/acceptance/experimental/air/run-submit/script @@ -0,0 +1,18 @@ +# Pinned commit dates keep the resolved HEAD SHA — and thus the snapshot cache key +# baked into the tarball name — stable across runs. +export GIT_AUTHOR_DATE="2020-01-01T00:00:00Z" +export GIT_COMMITTER_DATE="2020-01-01T00:00:00Z" +git-repo-init + +# Pin the config to the committed HEAD. git.commit tolerates a dirty working tree +# (the pinned revision is what gets archived), which matters here because the +# acceptance harness writes output.txt into the working dir as the script runs. +sed "s/COMMIT_SHA/$(git rev-parse HEAD)/" run.yaml.tmpl > run.yaml + +title "submit with a git code_source" +trace $CLI experimental air run -f run.yaml + +title "the ai_runtime_task carries the code_source_path" +trace print_requests.py //api/2.2/jobs/runs/submit + +rm -fr .git diff --git a/acceptance/experimental/air/run-submit/test.toml b/acceptance/experimental/air/run-submit/test.toml new file mode 100644 index 00000000000..fcdf8fd1242 --- /dev/null +++ b/acceptance/experimental/air/run-submit/test.toml @@ -0,0 +1,35 @@ +# A real (non-dry-run) submit that packages a git code_source, uploads the +# tarball + provenance sidecars, and POSTs runs/submit. No bundle deploy, so no +# engine matrix. +RecordRequests = true + +# run.yaml is generated from run.yaml.tmpl at test time (commit SHA templated in); +# it isn't a committed input to diff. +Ignore = ["run.yaml"] + +[EnvMatrix] +DATABRICKS_BUNDLE_ENGINE = [] + +# The SDK probes host reachability with a HEAD request; stub it for determinism. +[[Server]] +Pattern = "HEAD /" +Response.Body = '' + +[[Server]] +Pattern = "POST /api/2.2/jobs/runs/submit" +Response.Body = ''' +{"run_id": 555} +''' + +# The snapshot tarball is named _.tar.gz, where is the +# test's temp-dir basename and the cache key derives from the pinned commit SHA. +# Both are stable given the pinned commit dates in the script, but the temp-dir +# basename varies per run, so collapse the whole tarball filename to a stable token. +[[Repls]] +Old = '[a-zA-Z0-9._-]+_[0-9a-f]{16}\.tar\.gz' +New = '[SNAPSHOT_TARBALL]' + +# The per-run launch directory ends in _<16 hex>; the random suffix varies. +[[Repls]] +Old = 'submit-smoke_[0-9a-f]{16}' +New = 'submit-smoke_[RUN_ID]' diff --git a/acceptance/experimental/air/run/git-remote.yaml b/acceptance/experimental/air/run/git-remote.yaml new file mode 100644 index 00000000000..0161fe14972 --- /dev/null +++ b/acceptance/experimental/air/run/git-remote.yaml @@ -0,0 +1,12 @@ +experiment_name: smoke-test +command: python train.py +compute: + accelerator_type: GPU_1xH100 + num_accelerators: 1 +code_source: + type: snapshot + snapshot: + root_path: . + git: + branch: main + remote: origin diff --git a/acceptance/experimental/air/run/invalid.yaml b/acceptance/experimental/air/run/invalid.yaml new file mode 100644 index 00000000000..c011fc81b37 --- /dev/null +++ b/acceptance/experimental/air/run/invalid.yaml @@ -0,0 +1,5 @@ +experiment_name: bad.name +command: x +compute: + accelerator_type: GPU_8xH100 + num_accelerators: 3 diff --git a/acceptance/experimental/air/run/out.test.toml b/acceptance/experimental/air/run/out.test.toml new file mode 100644 index 00000000000..d6187dcb046 --- /dev/null +++ b/acceptance/experimental/air/run/out.test.toml @@ -0,0 +1,3 @@ +Local = true +Cloud = false +EnvMatrix.DATABRICKS_BUNDLE_ENGINE = [] diff --git a/acceptance/experimental/air/run/output.txt b/acceptance/experimental/air/run/output.txt new file mode 100644 index 00000000000..a753eabd198 --- /dev/null +++ b/acceptance/experimental/air/run/output.txt @@ -0,0 +1,49 @@ + +=== dry-run (text) +>>> [CLI] experimental air run -f valid.yaml --dry-run +Dry run: configuration for "smoke-test" is valid; not submitting. + +=== dry-run (json) +>>> [CLI] experimental air run -f valid.yaml --dry-run -o json +{ + "v": 1, + "ts": "[TIMESTAMP]", + "data": { + "status": "DRY_RUN_OK", + "dry_run": true + } +} + +=== override not yet supported +>>> [CLI] experimental air run -f valid.yaml --dry-run --override a=b +Error: --override is not yet supported + +Exit code: 1 + +=== watch not yet supported +>>> [CLI] experimental air run -f valid.yaml --dry-run --watch +Error: --watch is not yet supported + +Exit code: 1 + +=== code_source config passes validation +>>> [CLI] experimental air run -f with-code-source.yaml --dry-run +Dry run: configuration for "smoke-test" is valid; not submitting. + +=== git.remote is rejected +>>> [CLI] experimental air run -f git-remote.yaml --dry-run +Error: git.remote is no longer supported: the snapshot archives your local copy, so a branch resolves to its local HEAD. To deploy a specific committed revision, use git.commit + +Exit code: 1 + +=== invalid config is rejected +>>> [CLI] experimental air run -f invalid.yaml --dry-run +Error: invalid experiment_name "bad.name": only alphanumeric characters, hyphens (-), and underscores (_) are allowed + +Exit code: 1 + +=== missing --file +>>> [CLI] experimental air run --dry-run +Error: required flag(s) "file" not set + +Exit code: 1 diff --git a/acceptance/experimental/air/run/script b/acceptance/experimental/air/run/script new file mode 100644 index 00000000000..312b2f6fecf --- /dev/null +++ b/acceptance/experimental/air/run/script @@ -0,0 +1,23 @@ +title "dry-run (text)" +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 "watch not yet supported" +errcode 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 + +title "git.remote is rejected" +errcode trace $CLI experimental air run -f git-remote.yaml --dry-run + +title "invalid config is rejected" +errcode trace $CLI experimental air run -f invalid.yaml --dry-run + +title "missing --file" +errcode trace $CLI experimental air run --dry-run diff --git a/acceptance/experimental/air/run/test.toml b/acceptance/experimental/air/run/test.toml new file mode 100644 index 00000000000..2f971c3ed21 --- /dev/null +++ b/acceptance/experimental/air/run/test.toml @@ -0,0 +1,4 @@ +# `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 = [] diff --git a/acceptance/experimental/air/run/valid.yaml b/acceptance/experimental/air/run/valid.yaml new file mode 100644 index 00000000000..b82a321b051 --- /dev/null +++ b/acceptance/experimental/air/run/valid.yaml @@ -0,0 +1,5 @@ +experiment_name: smoke-test +command: python train.py +compute: + accelerator_type: GPU_1xH100 + num_accelerators: 1 diff --git a/acceptance/experimental/air/run/with-code-source.yaml b/acceptance/experimental/air/run/with-code-source.yaml new file mode 100644 index 00000000000..86a32c138cb --- /dev/null +++ b/acceptance/experimental/air/run/with-code-source.yaml @@ -0,0 +1,13 @@ +experiment_name: smoke-test +command: python train.py +compute: + accelerator_type: GPU_1xH100 + num_accelerators: 1 +code_source: + type: snapshot + snapshot: + root_path: . + git: + branch: main + include_paths: + - src diff --git a/acceptance/experimental/air/unimplemented/out.test.toml b/acceptance/experimental/air/unimplemented/out.test.toml new file mode 100644 index 00000000000..d6187dcb046 --- /dev/null +++ b/acceptance/experimental/air/unimplemented/out.test.toml @@ -0,0 +1,3 @@ +Local = true +Cloud = false +EnvMatrix.DATABRICKS_BUNDLE_ENGINE = [] diff --git a/acceptance/experimental/air/unimplemented/output.txt b/acceptance/experimental/air/unimplemented/output.txt new file mode 100644 index 00000000000..7db6ef1aec2 --- /dev/null +++ b/acceptance/experimental/air/unimplemented/output.txt @@ -0,0 +1,12 @@ + +=== 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 new file mode 100644 index 00000000000..19dc13ffe85 --- /dev/null +++ b/acceptance/experimental/air/unimplemented/script @@ -0,0 +1,7 @@ +# 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 new file mode 100644 index 00000000000..c233c30a86c --- /dev/null +++ b/acceptance/experimental/air/unimplemented/test.toml @@ -0,0 +1,3 @@ +# Stubs fail locally before any API call, so no server stubs needed. +[EnvMatrix] +DATABRICKS_BUNDLE_ENGINE = [] diff --git a/cmd/experimental/experimental.go b/cmd/experimental/experimental.go index 8d9827c5c94..d87c893abc5 100644 --- a/cmd/experimental/experimental.go +++ b/cmd/experimental/experimental.go @@ -1,6 +1,7 @@ package experimental import ( + aircmd "github.com/databricks/cli/experimental/air/cmd" aitoolscmd "github.com/databricks/cli/experimental/aitools/cmd" geniecmd "github.com/databricks/cli/experimental/genie/cmd" postgrescmd "github.com/databricks/cli/experimental/postgres/cmd" @@ -22,6 +23,7 @@ These commands provide early access to new features that are still under development. They may change or be removed in future versions without notice.`, } + cmd.AddCommand(aircmd.New()) cmd.AddCommand(aitoolscmd.NewAitoolsCmd()) cmd.AddCommand(geniecmd.NewGenieCmd()) cmd.AddCommand(postgrescmd.New()) diff --git a/experimental/air/cmd/air.go b/experimental/air/cmd/air.go new file mode 100644 index 00000000000..fbf40a34b52 --- /dev/null +++ b/experimental/air/cmd/air.go @@ -0,0 +1,33 @@ +package aircmd + +import ( + "fmt" + + "github.com/spf13/cobra" +) + +// New returns the root command for the experimental AI runtime CLI. +func New() *cobra.Command { + cmd := &cobra.Command{ + Use: "air", + Short: "Run and manage AI runtime training workloads", + Long: `Run and manage AI runtime training workloads on Databricks serverless GPU compute. + +This command set is the Go port of the standalone Python "air" CLI. It is +experimental and may change in future versions.`, + } + + cmd.AddCommand(newRunCommand()) + cmd.AddCommand(newGetCommand()) + cmd.AddCommand(newListCommand()) + cmd.AddCommand(newLogsCommand()) + cmd.AddCommand(newCancelCommand()) + cmd.AddCommand(newRegisterImageCommand()) + + 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/air_test.go b/experimental/air/cmd/air_test.go new file mode 100644 index 00000000000..7efac253a2b --- /dev/null +++ b/experimental/air/cmd/air_test.go @@ -0,0 +1,22 @@ +package aircmd + +import ( + "testing" + + "github.com/stretchr/testify/assert" +) + +// TestNewRegistersAllSubcommands asserts the `air` command wires up every +// expected subcommand, so none is accidentally dropped from New. +func TestNewRegistersAllSubcommands(t *testing.T) { + registered := make(map[string]bool) + for _, c := range New().Commands() { + registered[c.Name()] = true + } + + want := []string{"run", "get", "list", "logs", "cancel", "register-image"} + for _, name := range want { + assert.True(t, registered[name], "subcommand %q is not registered", name) + } + assert.Len(t, registered, len(want), "unexpected number of subcommands") +} diff --git a/experimental/air/cmd/aitraining.go b/experimental/air/cmd/aitraining.go new file mode 100644 index 00000000000..735c0d71aaf --- /dev/null +++ b/experimental/air/cmd/aitraining.go @@ -0,0 +1,110 @@ +package aircmd + +import ( + "context" + "encoding/json" + "fmt" + "net/http" + "time" + + "github.com/databricks/databricks-sdk-go" + "github.com/databricks/databricks-sdk-go/client" +) + +// aiTrainingWorkflowsPath is the AiTrainingService index of the caller's own AIR +// runs. It returns cheap (job_run_id, submit_time) pairs, letting `air list` +// order and page without scanning the Jobs runs/list firehose. This is called +// with a raw client.Do because the SDK does not model the AiTrainingService. +const aiTrainingWorkflowsPath = "/api/2.0/ai-training/workflows" + +// workflowRef is one run from the index: its Jobs run id and submission time. +type workflowRef struct { + jobRunID int64 + submitTimeMs int64 +} + +type aiTrainingWorkflow struct { + // job_run_id is a Jobs run id; tolerate it arriving as a JSON number or string. + JobRunID json.Number `json:"job_run_id"` + // submit_time is a proto Timestamp, serialized over HTTP as either an RFC3339 + // string or a {seconds, nanos} object. + SubmitTime json.RawMessage `json:"submit_time"` +} + +type aiTrainingWorkflowsResponse struct { + TrainingWorkflows []aiTrainingWorkflow `json:"training_workflows"` + NextPageToken string `json:"next_page_token"` +} + +// listAiTrainingWorkflows pages the index and returns every workflow ref the +// caller owns. Pagination stops at the end or when a page token repeats, which +// guards against a stuck or cycling cursor without an arbitrary page cap. +func listAiTrainingWorkflows(ctx context.Context, w *databricks.WorkspaceClient, activeOnly bool) ([]workflowRef, error) { + apiClient, err := client.New(w.Config) + if err != nil { + return nil, fmt.Errorf("failed to create API client: %w", err) + } + + var refs []workflowRef + seenTokens := map[string]bool{} + // The index can return the same job_run_id on more than one page; dedupe so + // the newest-`limit` truncation counts unique runs, not repeats. + seenIDs := map[int64]bool{} + var pageToken string + for { + query := map[string]any{} + if activeOnly { + query["active_only"] = true + } + if pageToken != "" { + query["page_token"] = pageToken + } + + var resp aiTrainingWorkflowsResponse + err = apiClient.Do(ctx, http.MethodGet, aiTrainingWorkflowsPath, nil, nil, query, &resp) + if err != nil { + return nil, fmt.Errorf("failed to list training workflows: %w", err) + } + + for _, wf := range resp.TrainingWorkflows { + id, err := wf.JobRunID.Int64() + if err != nil || id == 0 || seenIDs[id] { + continue + } + seenIDs[id] = true + refs = append(refs, workflowRef{jobRunID: id, submitTimeMs: parseSubmitTimeMs(wf.SubmitTime)}) + } + + if resp.NextPageToken == "" || seenTokens[resp.NextPageToken] { + break + } + seenTokens[resp.NextPageToken] = true + pageToken = resp.NextPageToken + } + return refs, nil +} + +// parseSubmitTimeMs converts a proto Timestamp (RFC3339 string or {seconds, nanos} +// object) to epoch milliseconds, or 0 when absent or unparseable (so it sorts last). +func parseSubmitTimeMs(raw json.RawMessage) int64 { + if len(raw) == 0 { + return 0 + } + + var s string + if json.Unmarshal(raw, &s) == nil { + if t, err := time.Parse(time.RFC3339, s); err == nil { + return t.UnixMilli() + } + return 0 + } + + var obj struct { + Seconds int64 `json:"seconds"` + Nanos int64 `json:"nanos"` + } + if json.Unmarshal(raw, &obj) == nil { + return obj.Seconds*1000 + obj.Nanos/1_000_000 + } + return 0 +} diff --git a/experimental/air/cmd/aitraining_test.go b/experimental/air/cmd/aitraining_test.go new file mode 100644 index 00000000000..4c44266d8cb --- /dev/null +++ b/experimental/air/cmd/aitraining_test.go @@ -0,0 +1,92 @@ +package aircmd + +import ( + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestParseSubmitTimeMs(t *testing.T) { + cases := []struct { + name string + raw string + want int64 + }{ + {"rfc3339", `"2023-11-14T22:13:20Z"`, 1700000000000}, + {"rfc3339 offset", `"2023-11-14T22:13:20+00:00"`, 1700000000000}, + {"seconds and nanos", `{"seconds": 1700000000, "nanos": 500000000}`, 1700000000500}, + {"seconds only", `{"seconds": 1700000000}`, 1700000000000}, + {"empty", ``, 0}, + {"garbage string", `"not-a-time"`, 0}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + assert.Equal(t, tc.want, parseSubmitTimeMs(json.RawMessage(tc.raw))) + }) + } +} + +// indexServer serves paginated AiTrainingService responses, one body per call, +// tracking whether the index was hit. +func indexServer(t *testing.T, hit *bool, bodies ...string) *httptest.Server { + t.Helper() + call := 0 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == aiTrainingWorkflowsPath { + *hit = true + body := bodies[min(call, len(bodies)-1)] + call++ + _, _ = w.Write([]byte(body)) + return + } + _, _ = w.Write([]byte(`{}`)) + })) + t.Cleanup(srv.Close) + return srv +} + +func TestListAiTrainingWorkflowsPaginates(t *testing.T) { + page1 := `{"training_workflows":[{"job_run_id":"1","submit_time":"2023-11-14T22:13:20Z"}],"next_page_token":"tok"}` + page2 := `{"training_workflows":[{"job_run_id":2,"submit_time":{"seconds":1700000100}}]}` + var hit bool + srv := indexServer(t, &hit, page1, page2) + + refs, err := listAiTrainingWorkflows(t.Context(), newTestWorkspaceClient(t, srv.URL), false) + require.NoError(t, err) + require.Len(t, refs, 2) + assert.Equal(t, int64(1), refs[0].jobRunID) + assert.Equal(t, int64(1700000000000), refs[0].submitTimeMs) + assert.Equal(t, int64(2), refs[1].jobRunID) +} + +func TestListAiTrainingWorkflowsStopsOnRepeatedToken(t *testing.T) { + // A cursor that always returns the same token must not loop forever. The + // repeated id is also deduped, so only one ref survives. + page := `{"training_workflows":[{"job_run_id":1,"submit_time":"2023-11-14T22:13:20Z"}],"next_page_token":"tok"}` + var hit bool + srv := indexServer(t, &hit, page) + + refs, err := listAiTrainingWorkflows(t.Context(), newTestWorkspaceClient(t, srv.URL), false) + require.NoError(t, err) + require.Len(t, refs, 1) + assert.Equal(t, int64(1), refs[0].jobRunID) +} + +func TestListAiTrainingWorkflowsDedupesIDs(t *testing.T) { + // The same job_run_id on multiple pages must be counted once, so the + // newest-limit truncation doesn't silently return fewer unique runs. + page1 := `{"training_workflows":[{"job_run_id":1,"submit_time":"2023-11-14T22:13:20Z"},{"job_run_id":2,"submit_time":"2023-11-14T22:13:21Z"}],"next_page_token":"tok"}` + page2 := `{"training_workflows":[{"job_run_id":2,"submit_time":"2023-11-14T22:13:21Z"},{"job_run_id":3,"submit_time":"2023-11-14T22:13:22Z"}]}` + var hit bool + srv := indexServer(t, &hit, page1, page2) + + refs, err := listAiTrainingWorkflows(t.Context(), newTestWorkspaceClient(t, srv.URL), false) + require.NoError(t, err) + require.Len(t, refs, 3) + got := []int64{refs[0].jobRunID, refs[1].jobRunID, refs[2].jobRunID} + assert.ElementsMatch(t, []int64{1, 2, 3}, got) +} diff --git a/experimental/air/cmd/cancel.go b/experimental/air/cmd/cancel.go new file mode 100644 index 00000000000..519a7a82068 --- /dev/null +++ b/experimental/air/cmd/cancel.go @@ -0,0 +1,222 @@ +package aircmd + +import ( + "context" + "errors" + "fmt" + "net/http" + "strconv" + "strings" + "text/tabwriter" + + "github.com/databricks/cli/cmd/root" + "github.com/databricks/cli/libs/cmdctx" + "github.com/databricks/cli/libs/cmdio" + "github.com/databricks/cli/libs/flags" + "github.com/databricks/databricks-sdk-go" + "github.com/databricks/databricks-sdk-go/apierr" + "github.com/databricks/databricks-sdk-go/service/iam" + "github.com/databricks/databricks-sdk-go/service/jobs" + "github.com/spf13/cobra" +) + +// cancelData is the JSON payload printed by `air cancel`. `all` is set only for +// --all, `workspace` only when --all finds no active runs, and `failed` only +// when a run could not be cancelled. +type cancelData struct { + Cancelled []string `json:"cancelled"` + All bool `json:"all,omitempty"` + Workspace string `json:"workspace,omitempty"` + Failed []cancelFailure `json:"failed,omitempty"` +} + +type cancelFailure struct { + RunID string `json:"run_id"` + Error string `json:"error"` +} + +func newCancelCommand() *cobra.Command { + var ( + all bool + yes bool + ) + + cmd := &cobra.Command{ + Use: "cancel [JOB_RUN_ID...]", + Short: "Cancel one or more runs", + Long: `Cancel one or more runs by ID, or cancel all of your active runs with --all.`, + } + + cmd.Flags().BoolVar(&all, "all", false, "Cancel all of your active runs") + cmd.Flags().BoolVarP(&yes, "yes", "y", false, "Skip the confirmation prompt") + + // Require exactly one of: one or more JOB_RUN_IDs, or --all. Cobra parses flags + // before running this, so `all` reflects the user's input. + cmd.Args = func(cmd *cobra.Command, args []string) error { + switch { + case all && len(args) > 0: + return &root.InvalidArgsError{Command: cmd, Message: "cannot combine JOB_RUN_ID arguments with --all"} + case !all && len(args) == 0: + return &root.InvalidArgsError{Command: cmd, Message: "provide at least one JOB_RUN_ID, or use --all"} + } + return nil + } + + // In -o json mode an auth failure should be a JSON error envelope, not a bare + // error. ErrAlreadyPrinted passes through (already handled upstream). + 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 renderError(cmd.Context(), cmd, "INTERNAL_ERROR", "TRANSIENT", true, err) + } + + cmd.RunE = func(cmd *cobra.Command, args []string) error { + ctx := cmd.Context() + w := cmdctx.WorkspaceClient(ctx) + jsonOut := root.OutputType(cmd) == flags.OutputJSON + + runIDs := args + data := cancelData{Cancelled: []string{}} + + if all { + data.All = true + + me, err := w.CurrentUser.Me(ctx, iam.MeRequest{}) + if err != nil { + return renderError(ctx, cmd, "INTERNAL_ERROR", "TRANSIENT", true, + fmt.Errorf("failed to resolve current user: %w", err)) + } + host := strings.TrimRight(w.Config.Host, "/") + + if !jsonOut { + cmdio.LogString(ctx, fmt.Sprintf("Searching active runs for %s in %s...", me.UserName, host)) + } + + // Fetch every active run (up to the scan bound) so --all cancels all + // of them, not just the first page. + fetcher := newRunFetcher(ctx, w, listQuery{activeOnly: true, userFilter: me.UserName}) + rows, err := fetcher.next(maxListScan) + if err != nil { + return renderError(ctx, cmd, "INTERNAL_ERROR", "TRANSIENT", true, + fmt.Errorf("failed to list active runs: %w", err)) + } + + runIDs = make([]string, 0, len(rows)) + for i := range rows { + if rows[i].RunID != "" { + runIDs = append(runIDs, rows[i].RunID) + } + } + + if len(runIDs) == 0 { + if jsonOut { + data.Workspace = host + return renderEnvelope(ctx, data) + } + cmdio.LogString(ctx, "No active runs found.") + return nil + } + + if !yes { + displayCancelPreview(ctx, rows, host) + confirmed, err := cmdio.AskYesOrNo(ctx, fmt.Sprintf("\nCancel %d run(s) in %s?", len(runIDs), host)) + if err != nil { + return err + } + if !confirmed { + cmdio.LogString(ctx, "Cancellation aborted.") + return root.ErrAlreadyPrinted + } + } + } + + for _, rid := range runIDs { + err := cancelRun(ctx, w, rid) + if err != nil { + data.Failed = append(data.Failed, cancelFailure{RunID: rid, Error: err.Error()}) + if !jsonOut { + if runNotFound(err) { + cmdio.LogString(ctx, fmt.Sprintf("Run %s not found. Please check the run ID and ensure you're using a Job Run ID.", rid)) + } else { + cmdio.LogString(ctx, fmt.Sprintf("Failed to cancel run %s: %s", rid, err)) + } + } + continue + } + data.Cancelled = append(data.Cancelled, rid) + if !jsonOut { + cmdio.LogString(ctx, "Successfully requested cancellation for run "+rid) + } + } + + if jsonOut { + if err := renderEnvelope(ctx, data); err != nil { + return err + } + // Print the envelope, but still exit non-zero on any failure. + if len(data.Failed) > 0 { + return root.ErrAlreadyPrinted + } + return nil + } + + if len(data.Failed) > 0 { + cmdio.LogString(ctx, fmt.Sprintf("%d run(s) failed to cancel.", len(data.Failed))) + return root.ErrAlreadyPrinted + } + if all || len(data.Cancelled) > 1 { + cmdio.LogString(ctx, fmt.Sprintf("Successfully requested cancellation for %d run(s).", len(data.Cancelled))) + } + return nil + } + + return cmd +} + +// runNotFound reports whether err means the run does not exist. The cancel +// endpoint returns 400 INVALID_PARAMETER_VALUE ("Run does not exist") for +// an unknown run, and the SDK only remaps that to ErrResourceDoesNotExist for +// the runs/get path, not cancel — so we also detect the raw code here. +func runNotFound(err error) bool { + if errors.Is(err, apierr.ErrResourceDoesNotExist) { + return true + } + if apiErr, ok := errors.AsType[*apierr.APIError](err); ok { + return apiErr.StatusCode == http.StatusBadRequest && apiErr.ErrorCode == "INVALID_PARAMETER_VALUE" + } + return false +} + +// cancelRun requests cancellation of a single job run. The cancel is async, so +// the returned waiter is ignored. +func cancelRun(ctx context.Context, w *databricks.WorkspaceClient, rid string) error { + runID, err := strconv.ParseInt(rid, 10, 64) + if err != nil || runID <= 0 { + return fmt.Errorf("invalid run ID %q: must be a positive integer", rid) + } + _, err = w.Jobs.CancelRun(ctx, jobs.CancelRun{RunId: runID}) + return err +} + +// displayCancelPreview shows the runs that `cancel --all` is about to terminate. +func displayCancelPreview(ctx context.Context, rows []listRow, host string) { + var sb strings.Builder + fmt.Fprintf(&sb, "\nWorkspace: %s\n", host) + fmt.Fprintf(&sb, "Found %d active run(s) to cancel:\n\n", len(rows)) + + tw := tabwriter.NewWriter(&sb, 0, 0, 2, ' ', 0) + fmt.Fprintln(tw, "Run ID\tExperiment\tStarted") + for i := range rows { + experiment := orNA(rows[i].Experiment) + started := na + if rows[i].StartedAt != nil { + started = *rows[i].StartedAt + } + fmt.Fprintf(tw, "%s\t%s\t%s\n", rows[i].RunID, experiment, started) + } + tw.Flush() + + cmdio.LogString(ctx, strings.TrimRight(sb.String(), "\n")) +} diff --git a/experimental/air/cmd/cancel_test.go b/experimental/air/cmd/cancel_test.go new file mode 100644 index 00000000000..5601a99f84a --- /dev/null +++ b/experimental/air/cmd/cancel_test.go @@ -0,0 +1,289 @@ +package aircmd + +import ( + "bytes" + "encoding/json" + "errors" + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" + "testing/iotest" + + "github.com/databricks/cli/cmd/root" + "github.com/databricks/cli/libs/cmdctx" + "github.com/databricks/cli/libs/cmdio" + "github.com/databricks/cli/libs/flags" + "github.com/databricks/databricks-sdk-go" + "github.com/databricks/databricks-sdk-go/apierr" + "github.com/databricks/databricks-sdk-go/experimental/mocks" + "github.com/databricks/databricks-sdk-go/service/iam" + "github.com/databricks/databricks-sdk-go/service/jobs" + "github.com/spf13/cobra" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/mock" + "github.com/stretchr/testify/require" +) + +// runCancelAll runs `cancel --all` against w with the given output mode and +// stdin, capturing output into buf. +func runCancelAll(t *testing.T, w *databricks.WorkspaceClient, out flags.Output, in io.Reader, buf *bytes.Buffer) error { + t.Helper() + cmd := withOutput(newCancelCommand(), out) + require.NoError(t, cmd.Flags().Set("all", "true")) + ctx := cmdio.InContext(t.Context(), cmdio.NewIO(t.Context(), out, in, buf, buf, "", "")) + cmd.SetContext(cmdctx.SetWorkspaceClient(ctx, w)) + return cmd.RunE(cmd, nil) +} + +// cancelEnvelope decodes the air JSON envelope with the cancel payload. +type cancelEnvelope struct { + V int `json:"v"` + Data cancelData `json:"data"` +} + +// runCancel runs the cancel command against w with the given output mode and +// stdin, capturing stdout/stderr into buf. +func runCancel(t *testing.T, w *databricks.WorkspaceClient, out flags.Output, in string, buf *bytes.Buffer, args ...string) (*cobra.Command, error) { + t.Helper() + ctx := cmdio.InContext(t.Context(), cmdio.NewIO(t.Context(), out, strings.NewReader(in), buf, buf, "", "")) + ctx = cmdctx.SetWorkspaceClient(ctx, w) + cmd := withOutput(newCancelCommand(), out) + cmd.SetContext(ctx) + return cmd, cmd.RunE(cmd, args) +} + +func TestCancelArgs(t *testing.T) { + tests := []struct { + name string + all bool + args []string + wantErr string + }{ + {name: "one id", args: []string{"123"}}, + {name: "many ids", args: []string{"123", "456"}}, + {name: "all", all: true}, + {name: "no input", wantErr: "provide at least one JOB_RUN_ID, or use --all"}, + {name: "ids with all", all: true, args: []string{"123"}, wantErr: "cannot combine JOB_RUN_ID arguments with --all"}, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + cmd := newCancelCommand() + if tc.all { + require.NoError(t, cmd.Flags().Set("all", "true")) + } + err := cmd.Args(cmd, tc.args) + if tc.wantErr == "" { + assert.NoError(t, err) + return + } + require.Error(t, err) + assert.Contains(t, err.Error(), tc.wantErr) + }) + } +} + +func TestCancelRunInvalidID(t *testing.T) { + m := mocks.NewMockWorkspaceClient(t) + for _, id := range []string{"abc", "0", "-1"} { + err := cancelRun(t.Context(), m.WorkspaceClient, id) + require.Error(t, err) + assert.Contains(t, err.Error(), "invalid run ID") + } +} + +func TestCancelByIDSuccess(t *testing.T) { + m := mocks.NewMockWorkspaceClient(t) + m.GetMockJobsAPI().EXPECT().CancelRun(mock.Anything, jobs.CancelRun{RunId: 123}).Return(nil, nil) + m.GetMockJobsAPI().EXPECT().CancelRun(mock.Anything, jobs.CancelRun{RunId: 456}).Return(nil, nil) + + var buf bytes.Buffer + _, err := runCancel(t, m.WorkspaceClient, flags.OutputText, "", &buf, "123", "456") + require.NoError(t, err) + + out := buf.String() + assert.Contains(t, out, "Successfully requested cancellation for run 123") + assert.Contains(t, out, "Successfully requested cancellation for run 456") + // More than one run cancelled prints the count summary. + assert.Contains(t, out, "Successfully requested cancellation for 2 run(s).") +} + +func TestCancelByIDNotFound(t *testing.T) { + m := mocks.NewMockWorkspaceClient(t) + m.GetMockJobsAPI().EXPECT().CancelRun(mock.Anything, jobs.CancelRun{RunId: 5}).Return(nil, apierr.ErrResourceDoesNotExist) + + var buf bytes.Buffer + _, err := runCancel(t, m.WorkspaceClient, flags.OutputText, "", &buf, "5") + require.ErrorIs(t, err, root.ErrAlreadyPrinted) + + out := buf.String() + assert.Contains(t, out, "Run 5 not found") + assert.Contains(t, out, "1 run(s) failed to cancel.") +} + +func TestCancelByIDNotFoundInvalidParam(t *testing.T) { + // The cancel endpoint reports an unknown run as 400 INVALID_PARAMETER_VALUE, + // which the SDK does not remap to ErrResourceDoesNotExist for this path. + m := mocks.NewMockWorkspaceClient(t) + apiErr := &apierr.APIError{StatusCode: http.StatusBadRequest, ErrorCode: "INVALID_PARAMETER_VALUE", Message: "Run 5 does not exist."} + m.GetMockJobsAPI().EXPECT().CancelRun(mock.Anything, jobs.CancelRun{RunId: 5}).Return(nil, apiErr) + + var buf bytes.Buffer + _, err := runCancel(t, m.WorkspaceClient, flags.OutputText, "", &buf, "5") + require.ErrorIs(t, err, root.ErrAlreadyPrinted) + assert.Contains(t, buf.String(), "Run 5 not found") +} + +func TestCancelPartialFailureJSON(t *testing.T) { + m := mocks.NewMockWorkspaceClient(t) + m.GetMockJobsAPI().EXPECT().CancelRun(mock.Anything, jobs.CancelRun{RunId: 123}).Return(nil, nil) + m.GetMockJobsAPI().EXPECT().CancelRun(mock.Anything, jobs.CancelRun{RunId: 5}).Return(nil, apierr.ErrResourceDoesNotExist) + + var buf bytes.Buffer + _, err := runCancel(t, m.WorkspaceClient, flags.OutputJSON, "", &buf, "123", "5") + // The envelope is printed, but a failure still exits non-zero. + require.ErrorIs(t, err, root.ErrAlreadyPrinted) + + var got cancelEnvelope + require.NoError(t, json.Unmarshal(buf.Bytes(), &got)) + assert.Equal(t, []string{"123"}, got.Data.Cancelled) + require.Len(t, got.Data.Failed, 1) + assert.Equal(t, "5", got.Data.Failed[0].RunID) + assert.False(t, got.Data.All) +} + +func TestCancelByIDSuccessJSON(t *testing.T) { + m := mocks.NewMockWorkspaceClient(t) + m.GetMockJobsAPI().EXPECT().CancelRun(mock.Anything, jobs.CancelRun{RunId: 123}).Return(nil, nil) + + var buf bytes.Buffer + _, err := runCancel(t, m.WorkspaceClient, flags.OutputJSON, "", &buf, "123") + require.NoError(t, err) + + var got cancelEnvelope + require.NoError(t, json.Unmarshal(buf.Bytes(), &got)) + assert.Equal(t, []string{"123"}, got.Data.Cancelled) + assert.Empty(t, got.Data.Failed) +} + +func TestCancelByIDGenericFailure(t *testing.T) { + m := mocks.NewMockWorkspaceClient(t) + m.GetMockJobsAPI().EXPECT().CancelRun(mock.Anything, jobs.CancelRun{RunId: 7}).Return(nil, errors.New("boom")) + + var buf bytes.Buffer + _, err := runCancel(t, m.WorkspaceClient, flags.OutputText, "", &buf, "7") + require.ErrorIs(t, err, root.ErrAlreadyPrinted) + assert.Contains(t, buf.String(), "Failed to cancel run 7: boom") +} + +func TestCancelAllNoActiveRuns(t *testing.T) { + w := newTestWorkspaceClient(t, runsServer(t, runsListBody(t, "")).URL) + var buf bytes.Buffer + require.NoError(t, runCancelAll(t, w, flags.OutputText, nil, &buf)) + assert.Contains(t, buf.String(), "No active runs found.") +} + +func TestCancelAllNoActiveRunsJSON(t *testing.T) { + srv := runsServer(t, runsListBody(t, "")) + w := newTestWorkspaceClient(t, srv.URL) + + var buf bytes.Buffer + require.NoError(t, runCancelAll(t, w, flags.OutputJSON, nil, &buf)) + + var got cancelEnvelope + require.NoError(t, json.Unmarshal(buf.Bytes(), &got)) + assert.Empty(t, got.Data.Cancelled) + assert.True(t, got.Data.All) + assert.Equal(t, srv.URL, got.Data.Workspace) +} + +func TestCancelAllConfirmYes(t *testing.T) { + srv := runsServer(t, runsListBody(t, "", + airBaseRun(111, "me@example.com", "GPU_1xA10", 1, "/Users/me@example.com/exp-a"), + airBaseRun(222, "me@example.com", "GPU_1xA10", 1, "/Users/me@example.com/exp-b"), + )) + w := newTestWorkspaceClient(t, srv.URL) + + var buf bytes.Buffer + require.NoError(t, runCancelAll(t, w, flags.OutputText, strings.NewReader("y\n"), &buf)) + out := buf.String() + assert.Contains(t, out, "active run(s) to cancel") + assert.Contains(t, out, "Successfully requested cancellation for run 111") + assert.Contains(t, out, "Successfully requested cancellation for run 222") +} + +func TestCancelAllAbort(t *testing.T) { + srv := runsServer(t, runsListBody(t, "", + airBaseRun(111, "me@example.com", "GPU_1xA10", 1, "/Users/me@example.com/exp-a"), + )) + w := newTestWorkspaceClient(t, srv.URL) + + var buf bytes.Buffer + err := runCancelAll(t, w, flags.OutputText, strings.NewReader("n\n"), &buf) + require.ErrorIs(t, err, root.ErrAlreadyPrinted) + assert.Contains(t, buf.String(), "Cancellation aborted.") +} + +func TestCancelAllConfirmReadError(t *testing.T) { + srv := runsServer(t, runsListBody(t, "", + airBaseRun(111, "me@example.com", "GPU_1xA10", 1, "/Users/me@example.com/exp-a"), + )) + w := newTestWorkspaceClient(t, srv.URL) + + var buf bytes.Buffer + err := runCancelAll(t, w, flags.OutputText, iotest.ErrReader(errors.New("read failed")), &buf) + require.Error(t, err) + assert.Contains(t, err.Error(), "read failed") +} + +func TestCancelAllMeError(t *testing.T) { + m := mocks.NewMockWorkspaceClient(t) + m.GetMockCurrentUserAPI().EXPECT().Me(mock.Anything, iam.MeRequest{}).Return(nil, errors.New("nope")) + + var buf bytes.Buffer + err := runCancelAll(t, m.WorkspaceClient, flags.OutputText, nil, &buf) + require.Error(t, err) + assert.Contains(t, err.Error(), "failed to resolve current user") +} + +func TestCancelAllListError(t *testing.T) { + // Me succeeds (default empty user), but listing active runs fails. + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == "/api/2.2/jobs/runs/list" { + w.WriteHeader(http.StatusInternalServerError) + _, _ = w.Write([]byte(`{"error_code":"INTERNAL","message":"boom"}`)) + return + } + _, _ = w.Write([]byte(`{}`)) + })) + t.Cleanup(srv.Close) + w := newTestWorkspaceClient(t, srv.URL) + + var buf bytes.Buffer + err := runCancelAll(t, w, flags.OutputText, nil, &buf) + require.Error(t, err) + assert.Contains(t, err.Error(), "failed to list active runs") +} + +func TestDisplayCancelPreview(t *testing.T) { + var buf bytes.Buffer + ctx := cmdio.InContext(t.Context(), cmdio.NewIO(t.Context(), flags.OutputText, nil, &buf, &buf, "", "")) + + started := "2026-06-05 17:32 UTC" + rows := []listRow{ + {RunID: "111", Experiment: "exp-a", StartedAt: &started}, + {RunID: "222"}, // no experiment or start time -> N/A + } + displayCancelPreview(ctx, rows, "https://my-workspace.cloud.databricks.test") + + out := buf.String() + assert.Contains(t, out, "Workspace: https://my-workspace.cloud.databricks.test") + assert.Contains(t, out, "Found 2 active run(s) to cancel:") + assert.Contains(t, out, "Run ID") + assert.Contains(t, out, "111") + assert.Contains(t, out, "exp-a") + assert.Contains(t, out, "222") + assert.Contains(t, out, na) +} diff --git a/experimental/air/cmd/compute.go b/experimental/air/cmd/compute.go new file mode 100644 index 00000000000..07013c53906 --- /dev/null +++ b/experimental/air/cmd/compute.go @@ -0,0 +1,81 @@ +package aircmd + +import ( + "fmt" + "strings" +) + +// gpuType is a wire-facing accelerator type submitted to the training service. +// The number in the name is the partition count (e.g. GPU_8xH100 is 8 GPUs). +type gpuType string + +const ( + gpuType1xA10 gpuType = "GPU_1xA10" + gpuType8xH100 gpuType = "GPU_8xH100" + gpuType1xH100 gpuType = "GPU_1xH100" +) + +// gpuTypes lists every valid type. Used for validation error messages. +var gpuTypes = []gpuType{gpuType1xA10, gpuType1xH100, gpuType8xH100} + +func validGPUTypesHint() string { + names := make([]string, len(gpuTypes)) + for i, g := range gpuTypes { + names[i] = string(g) + } + return "valid types are: " + strings.Join(names, ", ") +} + +// parseGPUType resolves a YAML accelerator_type string to a gpuType. The match is +// exact: the server's lookup is case-sensitive. +func parseGPUType(value string) (gpuType, error) { + switch gpuType(value) { + case gpuType1xA10, gpuType8xH100, gpuType1xH100: + return gpuType(value), nil + } + return "", fmt.Errorf("invalid GPU type %q: %s", value, validGPUTypesHint()) +} + +// gpusPerNode returns the per-node GPU count, which is the partition count from +// the name (GPU_1xH100 -> 1, GPU_8xH100 -> 8). num_accelerators must be a +// round multiple of this since accelerators are allocated in whole nodes. +func gpusPerNode(g gpuType) (int, error) { + switch g { + case gpuType1xA10, gpuType1xH100: + return 1, nil + case gpuType8xH100: + return 8, nil + } + // Unreachable: callers resolve g through parseGPUType first, which rejects + // unknown types. Kept as a defensive guard. + return 0, fmt.Errorf("invalid GPU type %q", string(g)) +} + +// computeConfig is the `compute` block of the run YAML: which accelerators to +// use and how many. +type computeConfig struct { + NumAccelerators int `yaml:"num_accelerators"` + AcceleratorType string `yaml:"accelerator_type"` +} + +// validate checks the compute block against the backend's constraints. +func (c computeConfig) validate() error { + g, err := parseGPUType(c.AcceleratorType) + if err != nil { + return fmt.Errorf("compute.accelerator_type: %w", err) + } + + if c.NumAccelerators <= 0 { + return fmt.Errorf("compute.num_accelerators must be positive, got %d", c.NumAccelerators) + } + + perNode, err := gpusPerNode(g) + if err != nil { + return err + } + if c.NumAccelerators%perNode != 0 { + return fmt.Errorf("compute.num_accelerators for %s must be a multiple of %d, got %d", c.AcceleratorType, perNode, c.NumAccelerators) + } + + return nil +} diff --git a/experimental/air/cmd/compute_test.go b/experimental/air/cmd/compute_test.go new file mode 100644 index 00000000000..3464afbe9ea --- /dev/null +++ b/experimental/air/cmd/compute_test.go @@ -0,0 +1,86 @@ +package aircmd + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestParseGPUType(t *testing.T) { + tests := []struct { + in string + want gpuType + }{ + {"GPU_1xA10", gpuType1xA10}, + {"GPU_8xH100", gpuType8xH100}, + {"GPU_1xH100", gpuType1xH100}, + } + for _, tt := range tests { + t.Run(tt.in, func(t *testing.T) { + got, err := parseGPUType(tt.in) + require.NoError(t, err) + assert.Equal(t, tt.want, got) + }) + } +} + +func TestParseGPUTypeInvalid(t *testing.T) { + // Wrong casing is rejected rather than fixed up; legacy types (h100_80gb, a10) + // can no longer be submitted; unknown types are rejected. + for _, in := range []string{"gpu_1xa10", "GPU_1XA10", "GPU_2xH100", "h100_80gb", "a10", "b200", ""} { + t.Run(in, func(t *testing.T) { + _, err := parseGPUType(in) + require.Error(t, err) + assert.Contains(t, err.Error(), "valid types are") + }) + } +} + +func TestGPUsPerNode(t *testing.T) { + tests := []struct { + in gpuType + want int + }{ + {gpuType1xA10, 1}, + {gpuType1xH100, 1}, + {gpuType8xH100, 8}, + } + for _, tt := range tests { + t.Run(string(tt.in), func(t *testing.T) { + got, err := gpusPerNode(tt.in) + require.NoError(t, err) + assert.Equal(t, tt.want, got) + }) + } + + _, err := gpusPerNode(gpuType("nonsense")) + require.Error(t, err) +} + +func TestComputeConfigValidate(t *testing.T) { + tests := []struct { + name string + cfg computeConfig + wantErr string // substring; empty means the config is valid + }{ + {"single node", computeConfig{NumAccelerators: 8, AcceleratorType: "GPU_8xH100"}, ""}, + {"multiple nodes", computeConfig{NumAccelerators: 16, AcceleratorType: "GPU_8xH100"}, ""}, + {"single-gpu partitions", computeConfig{NumAccelerators: 3, AcceleratorType: "GPU_1xH100"}, ""}, + {"unknown type", computeConfig{NumAccelerators: 8, AcceleratorType: "b200"}, "accelerator_type"}, + {"legacy type rejected", computeConfig{NumAccelerators: 8, AcceleratorType: "h100_80gb"}, "accelerator_type"}, + {"non-positive count", computeConfig{NumAccelerators: 0, AcceleratorType: "GPU_1xH100"}, "must be positive"}, + {"count not a multiple", computeConfig{NumAccelerators: 4, AcceleratorType: "GPU_8xH100"}, "multiple of 8"}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := tt.cfg.validate() + if tt.wantErr == "" { + require.NoError(t, err) + return + } + require.Error(t, err) + assert.Contains(t, err.Error(), tt.wantErr) + }) + } +} diff --git a/experimental/air/cmd/format.go b/experimental/air/cmd/format.go new file mode 100644 index 00000000000..c32184517ba --- /dev/null +++ b/experimental/air/cmd/format.go @@ -0,0 +1,317 @@ +package aircmd + +import ( + "bytes" + "context" + "fmt" + "io" + "math" + "strconv" + "strings" + "time" + + "github.com/databricks/cli/libs/cmdio" + "github.com/databricks/databricks-sdk-go/service/jobs" + "go.yaml.in/yaml/v3" +) + +// na is the placeholder shown for an empty text-table cell, matching the Python CLI. +const na = "N/A" + +// orNA returns s, or "N/A" when s is empty, for text-table cells. +func orNA(s string) string { + if s == "" { + return na + } + 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 +// output stays plain text). +func hyperlink(ctx context.Context, out io.Writer, label, url string) string { + if url == "" || !cmdio.SupportsColor(ctx, out) { + return label + } + return osc8Link(label, url) +} + +// reformatYAMLForDisplay re-renders a training-config YAML so multi-line strings +// (notably the `command:` field) appear as `|` block literals instead of the +// quoted "\n"-escaped single line they are stored as, which is unreadable. It +// mirrors Python's _reformat_yaml_for_display (cli_display.py); we skip the +// Rich syntax-highlighted panel and only fix the whitespace. On any parse or +// re-encode failure it returns the original content unchanged. +func reformatYAMLForDisplay(content []byte) string { + var node yaml.Node + if err := yaml.Unmarshal(content, &node); err != nil { + return string(content) + } + forceLiteralBlockStrings(&node) + + var buf bytes.Buffer + enc := yaml.NewEncoder(&buf) + enc.SetIndent(2) + if err := enc.Encode(&node); err != nil { + return string(content) + } + enc.Close() + return buf.String() +} + +// forceLiteralBlockStrings walks a YAML node tree and marks every multi-line +// string scalar for `|` block-literal rendering. The encoder automatically +// falls back to a quoted style when a value can't be represented as a block +// literal (e.g. lines with trailing whitespace), so no explicit guard is needed. +func forceLiteralBlockStrings(node *yaml.Node) { + if node.Kind == yaml.ScalarNode && node.Tag == "!!str" && strings.Contains(node.Value, "\n") { + node.Style = yaml.LiteralStyle + } + for _, child := range node.Content { + forceLiteralBlockStrings(child) + } +} + +// gpuDisplayNames maps the GPU identifiers returned by the backend to the short +// names we show to users. Unknown identifiers are shown unchanged. +var gpuDisplayNames = map[string]string{ + "h100_80gb": "H100", + "a10": "A10", + "GPU_1xA10": "A10", + "GPU_8xH100": "H100", + "GPU_1xH100": "H100", +} + +// runStatus returns the single status word to show for a run. The backend +// reports two values: a lifecycle state (e.g. PENDING, RUNNING) and, once the +// run has finished, a result state (e.g. SUCCESS, FAILED). The result state is +// the more meaningful one, so we prefer it when it is set. +func runStatus(state *jobs.RunState) string { + if state == nil { + return "UNKNOWN" + } + return statusWord(string(state.LifeCycleState), string(state.ResultState)) +} + +// statusWord picks the status word to show from a run's lifecycle and result +// states: the result state is the more meaningful one, so it wins when set. +func statusWord(lifeCycle, result string) string { + if result != "" { + return result + } + if lifeCycle != "" { + return lifeCycle + } + return "UNKNOWN" +} + +// reportedTiming returns the run's start and end times (epoch milliseconds), +// preferring the last task's window over the run-level times so a retried run +// reports its latest attempt. Mirrors Python's _reported_attempt_timing +// (cli_display.py:78-87). +func reportedTiming(run *jobs.Run) (startMillis, endMillis int64) { + startMillis, endMillis = run.StartTime, run.EndTime + if n := len(run.Tasks); n > 0 { + last := run.Tasks[n-1] + if last.StartTime > 0 { + startMillis = last.StartTime + } + if last.EndTime > 0 { + endMillis = last.EndTime + } + } + return startMillis, endMillis +} + +// startedAt returns the run's start time as a Python-isoformat string ("+00:00", +// not "Z"; microseconds only when non-zero, cli_entrypoint.py:1899), or nil if it +// hasn't started. +func startedAt(run *jobs.Run) *string { + startMillis, _ := reportedTiming(run) + if startMillis == 0 { + return nil + } + s := isoFormat(time.UnixMilli(startMillis)) + return &s +} + +// isoFormat renders a time as a Python-style isoformat string in UTC ("+00:00", +// not "Z"; microseconds only when the sub-second part is non-zero), matching +// cli_entrypoint.py:1899. +func isoFormat(t time.Time) string { + t = t.UTC() + layout := "2006-01-02T15:04:05-07:00" + if t.Nanosecond() != 0 { + layout = "2006-01-02T15:04:05.000000-07:00" + } + return t.Format(layout) +} + +// submittedDisplay formats the run's start time for the text table as +// "2006-01-02 15:04 UTC", or "N/A" if it hasn't started. Mirrors Python's +// _format_timestamp (cli_display.py); we render in UTC for stable output rather +// than the local zone Python uses. +func submittedDisplay(run *jobs.Run) string { + startMillis, _ := reportedTiming(run) + if startMillis == 0 { + return na + } + return time.UnixMilli(startMillis).UTC().Format("2006-01-02 15:04 MST") +} + +// durationSeconds returns how long the run has taken, in whole seconds, or nil +// if it has not started. For a finished run this is the elapsed time of the +// reported attempt; for a still-running run it is the time since it started. +func durationSeconds(run *jobs.Run) *int64 { + startMillis, endMillis := reportedTiming(run) + if startMillis == 0 { + return nil + } + if endMillis == 0 { + // Still running: measure against the current time. + endMillis = time.Now().UnixMilli() + } + d := roundMillisToSeconds(endMillis - startMillis) + return &d +} + +// roundMillisToSeconds rounds milliseconds to whole seconds, half to even, to +// match Python's round() (cli_entrypoint.py:1903). +func roundMillisToSeconds(ms int64) int64 { + return int64(math.RoundToEven(float64(ms) / 1000)) +} + +// dashboardURL builds {host}/jobs/runs/{id}?o={workspace_id}, matching Python +// (cli_entrypoint.py:1911). The ?o= workspace id deep-links to the right +// workspace on multi-workspace accounts. +func dashboardURL(host string, runID, workspaceID int64) string { + return fmt.Sprintf("%s/jobs/runs/%d?o=%d", strings.TrimRight(host, "/"), runID, workspaceID) +} + +// formatDuration turns a number of seconds into a compact human string such as +// "1h 2m 3s". Trailing zero units are dropped, but a lone "0s" is kept so the +// result is never empty. +func formatDuration(totalSeconds int64) string { + hours := totalSeconds / 3600 + minutes := (totalSeconds % 3600) / 60 + seconds := totalSeconds % 60 + + var parts []string + if hours > 0 { + parts = append(parts, fmt.Sprintf("%dh", hours)) + } + if minutes > 0 { + parts = append(parts, fmt.Sprintf("%dm", minutes)) + } + if seconds > 0 || len(parts) == 0 { + parts = append(parts, fmt.Sprintf("%ds", seconds)) + } + return strings.Join(parts, " ") +} + +// latestAttemptNumber returns the retry count of the run's most recent task. +// Tasks start at attempt 0, so a value of 0 means the run has not been retried. +func latestAttemptNumber(run *jobs.Run) int { + if len(run.Tasks) == 0 { + return 0 + } + return run.Tasks[len(run.Tasks)-1].AttemptNumber +} + +// experimentName returns the MLflow experiment name for the run, or nil if there +// isn't one. Experiment names are often stored under a user's home folder (e.g. +// "/Users/me@example.com/my-experiment"); we strip that prefix so users see just +// the experiment name they chose. +func experimentName(run *jobs.Run) *string { + if len(run.Tasks) == 0 { + return nil + } + task := run.Tasks[0].GenAiComputeTask + if task == nil || task.MlflowExperimentName == "" { + return nil + } + name := stripExperimentUserPrefix(task.MlflowExperimentName) + return &name +} + +// stripExperimentUserPrefix removes a leading "/Users//" from an +// experiment name, leaving the remainder. Names without that prefix are returned +// unchanged. +func stripExperimentUserPrefix(name string) string { + if !strings.HasPrefix(name, "/Users/") { + return name + } + // Split into ["", "Users", "", ""]; keep "". + parts := strings.SplitN(name, "/", 4) + if len(parts) == 4 { + return parts[3] + } + return name +} + +// accelerators returns a short description of the GPUs the run uses, such as +// "8x H100", or an empty string if the run has no GPU compute attached. +func accelerators(run *jobs.Run) string { + if len(run.Tasks) == 0 { + return "" + } + task := run.Tasks[0].GenAiComputeTask + if task == nil || task.Compute == nil { + return "" + } + return acceleratorLabel(task.Compute.GpuType, task.Compute.NumGpus) +} + +// acceleratorLabel renders a GPU count and type as "8x H100", "" for none, or +// the count alone ("8x") when the type is unrecognized. +func acceleratorLabel(gpuType string, count int) string { + if count == 0 { + return "" + } + if name := gpuDisplayName(gpuType); name != "" { + return fmt.Sprintf("%dx %s", count, name) + } + return fmt.Sprintf("%dx", count) +} + +// gpuDisplayName returns the friendly name for a GPU identifier, falling back to +// the identifier itself when it is not one we recognize. +func gpuDisplayName(gpuType string) string { + if name, ok := gpuDisplayNames[gpuType]; ok { + return name + } + return gpuType +} + +// environment returns the run's runtime image (the training environment), or an +// empty string if the run has no GenAI-compute task. +func environment(run *jobs.Run) string { + if len(run.Tasks) == 0 { + return "" + } + task := run.Tasks[0].GenAiComputeTask + if task == nil { + return "" + } + return task.DlRuntimeImage +} + +// maxRetries returns the configured retry limit for the run's latest task as a +// display string: "unlimited" for the backend's -1, otherwise the count. +func maxRetries(run *jobs.Run) string { + if len(run.Tasks) == 0 { + return "0" + } + n := run.Tasks[len(run.Tasks)-1].MaxRetries + if n < 0 { + return "unlimited" + } + return strconv.Itoa(n) +} diff --git a/experimental/air/cmd/format_test.go b/experimental/air/cmd/format_test.go new file mode 100644 index 00000000000..62a6d7ac580 --- /dev/null +++ b/experimental/air/cmd/format_test.go @@ -0,0 +1,241 @@ +package aircmd + +import ( + "io" + "testing" + + "github.com/databricks/cli/libs/cmdio" + "github.com/databricks/databricks-sdk-go/service/jobs" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestOrNA(t *testing.T) { + assert.Equal(t, "x", orNA("x")) + assert.Equal(t, "N/A", orNA("")) +} + +func TestSubmittedDisplay(t *testing.T) { + assert.Equal(t, "N/A", submittedDisplay(&jobs.Run{})) + // 1700000000000 ms == 2023-11-14 22:13:20 UTC. + 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()) + assert.Equal(t, "label", hyperlink(ctx, io.Discard, "label", "https://h.test/x")) + // An empty URL is always rendered as the bare label. + assert.Equal(t, "label", hyperlink(ctx, io.Discard, "label", "")) +} + +func TestFormatDuration(t *testing.T) { + cases := []struct { + seconds int64 + want string + }{ + {0, "0s"}, + {45, "45s"}, + {60, "1m"}, + {63, "1m 3s"}, + {3600, "1h"}, + {3723, "1h 2m 3s"}, + {7260, "2h 1m"}, + } + for _, c := range cases { + assert.Equal(t, c.want, formatDuration(c.seconds)) + } +} + +func TestStripExperimentUserPrefix(t *testing.T) { + cases := []struct { + name string + want string + }{ + {"/Users/me@example.com/my-experiment", "my-experiment"}, + {"/Users/me@example.com/nested/path", "nested/path"}, + {"my-experiment", "my-experiment"}, + {"/Shared/team-experiment", "/Shared/team-experiment"}, + {"/Users/me@example.com", "/Users/me@example.com"}, + } + for _, c := range cases { + assert.Equal(t, c.want, stripExperimentUserPrefix(c.name)) + } +} + +func TestGpuDisplayName(t *testing.T) { + assert.Equal(t, "H100", gpuDisplayName("h100_80gb")) + assert.Equal(t, "A10", gpuDisplayName("GPU_1xA10")) + assert.Equal(t, "A10", gpuDisplayName("a10")) + assert.Equal(t, "H100", gpuDisplayName("GPU_8xH100")) + assert.Equal(t, "H100", gpuDisplayName("GPU_1xH100")) + // Unknown identifiers pass through unchanged. + assert.Equal(t, "b200", gpuDisplayName("b200")) + assert.Empty(t, gpuDisplayName("")) +} + +func TestRunStatusPrefersResultState(t *testing.T) { + // Result state wins once the run has finished. + assert.Equal(t, "SUCCESS", runStatus(&jobs.RunState{ + LifeCycleState: jobs.RunLifeCycleStateTerminated, + ResultState: jobs.RunResultStateSuccess, + })) + // Before completion only the lifecycle state is set. + assert.Equal(t, "RUNNING", runStatus(&jobs.RunState{ + LifeCycleState: jobs.RunLifeCycleStateRunning, + })) + // Non-nil state with neither field set, and nil state. + assert.Equal(t, "UNKNOWN", runStatus(&jobs.RunState{})) + assert.Equal(t, "UNKNOWN", runStatus(nil)) +} + +func TestReportedTiming(t *testing.T) { + // No tasks: run-level times are used. + start, end := reportedTiming(&jobs.Run{StartTime: 100, EndTime: 200}) + assert.Equal(t, int64(100), start) + assert.Equal(t, int64(200), end) + + // The last task's window is preferred over the run-level window, so a + // retried run reports its most recent attempt. + start, end = reportedTiming(&jobs.Run{ + StartTime: 100, EndTime: 200, + Tasks: []jobs.RunTask{ + {StartTime: 100, EndTime: 150}, + {StartTime: 300, EndTime: 450}, + }, + }) + assert.Equal(t, int64(300), start) + assert.Equal(t, int64(450), end) + + // A task missing a field falls back to the run-level value for that field. + start, end = reportedTiming(&jobs.Run{ + StartTime: 100, EndTime: 200, + Tasks: []jobs.RunTask{{StartTime: 300}}, + }) + assert.Equal(t, int64(300), start) + assert.Equal(t, int64(200), end) +} + +func TestStartedAt(t *testing.T) { + // Not started yet. + assert.Nil(t, startedAt(&jobs.Run{})) + // 1700000000000 ms == 2023-11-14T22:13:20+00:00 (Python isoformat, not "Z"). + got := startedAt(&jobs.Run{StartTime: 1700000000000}) + require.NotNil(t, got) + assert.Equal(t, "2023-11-14T22:13:20+00:00", *got) + // A sub-second start time carries microsecond precision. + got = startedAt(&jobs.Run{StartTime: 1700000000500}) + require.NotNil(t, got) + assert.Equal(t, "2023-11-14T22:13:20.500000+00:00", *got) + // The last attempt's start time is reported. 1700000060000 ms == 22:14:20. + got = startedAt(&jobs.Run{ + StartTime: 1700000000000, + Tasks: []jobs.RunTask{{StartTime: 1700000060000}}, + }) + require.NotNil(t, got) + assert.Equal(t, "2023-11-14T22:14:20+00:00", *got) +} + +func TestDurationSeconds(t *testing.T) { + // Not started yet. + assert.Nil(t, durationSeconds(&jobs.Run{})) + + // Finished run: reported end - start. + d := durationSeconds(&jobs.Run{StartTime: 1700000000000, EndTime: 1700000012000}) + require.NotNil(t, d) + assert.Equal(t, int64(12), *d) + + // Sub-second remainders round to the nearest second, matching Python: an + // 11,500 ms run reports 12s, not 11s. + d = durationSeconds(&jobs.Run{StartTime: 1700000000000, EndTime: 1700000011500}) + require.NotNil(t, d) + assert.Equal(t, int64(12), *d) + + // The last attempt's window drives the duration for a retried run. + d = durationSeconds(&jobs.Run{ + StartTime: 1700000000000, EndTime: 1700000012000, + Tasks: []jobs.RunTask{{StartTime: 1700000000000, EndTime: 1700000005000}}, + }) + require.NotNil(t, d) + assert.Equal(t, int64(5), *d) + + // Still running: measured against the current time, so positive. + d = durationSeconds(&jobs.Run{StartTime: 1700000000000}) + require.NotNil(t, d) + assert.Positive(t, *d) +} + +func TestDashboardURL(t *testing.T) { + // The ?o= workspace id and a trailing-slash-trimmed host, matching Python. + assert.Equal(t, "https://example.test/jobs/runs/123?o=42", dashboardURL("https://example.test/", 123, 42)) +} + +func TestLatestAttemptNumber(t *testing.T) { + assert.Equal(t, 0, latestAttemptNumber(&jobs.Run{})) + run := &jobs.Run{Tasks: []jobs.RunTask{{AttemptNumber: 0}, {AttemptNumber: 2}}} + assert.Equal(t, 2, latestAttemptNumber(run)) +} + +func TestExperimentName(t *testing.T) { + assert.Nil(t, experimentName(&jobs.Run{})) + assert.Nil(t, experimentName(&jobs.Run{Tasks: []jobs.RunTask{{}}})) + assert.Nil(t, experimentName(&jobs.Run{Tasks: []jobs.RunTask{{ + GenAiComputeTask: &jobs.GenAiComputeTask{MlflowExperimentName: ""}, + }}})) + got := experimentName(&jobs.Run{Tasks: []jobs.RunTask{{ + GenAiComputeTask: &jobs.GenAiComputeTask{MlflowExperimentName: "/Users/me@example.com/exp"}, + }}}) + require.NotNil(t, got) + assert.Equal(t, "exp", *got) +} + +func TestReformatYAMLForDisplay(t *testing.T) { + // A multi-line command stored as a quoted "\n"-escaped string is re-rendered + // as a `|` block literal, while key order is preserved. + in := "experiment_name: foo\ncommand: \"set -e\\npython train.py --epochs 3\\n\"\nmax_retries: 2\n" + want := "experiment_name: foo\ncommand: |\n set -e\n python train.py --epochs 3\nmax_retries: 2\n" + assert.Equal(t, want, reformatYAMLForDisplay([]byte(in))) + + // A single-line command is left as a plain scalar. + assert.Equal(t, "command: bash train.sh\n", reformatYAMLForDisplay([]byte("command: bash train.sh\n"))) + + // A multi-line value with trailing whitespace can't be a block literal, so the + // encoder falls back to a quoted style rather than emitting invalid YAML. + got := reformatYAMLForDisplay([]byte("command: \"trailing space \\nsecond\"\n")) + assert.Equal(t, "command: \"trailing space \\nsecond\"\n", got) + + // Unparseable content is returned unchanged. + assert.Equal(t, "\tnot: valid: yaml:", reformatYAMLForDisplay([]byte("\tnot: valid: yaml:"))) +} + +func TestAccelerators(t *testing.T) { + assert.Empty(t, accelerators(&jobs.Run{})) + assert.Empty(t, accelerators(&jobs.Run{Tasks: []jobs.RunTask{{}}})) + assert.Empty(t, accelerators(&jobs.Run{Tasks: []jobs.RunTask{{ + GenAiComputeTask: &jobs.GenAiComputeTask{}, + }}})) + assert.Empty(t, accelerators(&jobs.Run{Tasks: []jobs.RunTask{{ + GenAiComputeTask: &jobs.GenAiComputeTask{Compute: &jobs.ComputeConfig{NumGpus: 0}}, + }}})) + assert.Equal(t, "8x H100", accelerators(&jobs.Run{Tasks: []jobs.RunTask{{ + GenAiComputeTask: &jobs.GenAiComputeTask{Compute: &jobs.ComputeConfig{NumGpus: 8, GpuType: "GPU_8xH100"}}, + }}})) +} + +func TestAcceleratorLabel(t *testing.T) { + assert.Empty(t, acceleratorLabel("GPU_8xH100", 0)) + assert.Equal(t, "8x H100", acceleratorLabel("GPU_8xH100", 8)) + assert.Equal(t, "1x A10", acceleratorLabel("GPU_1xA10", 1)) + // The RPC may report a count without a recognized type. + assert.Equal(t, "8x", acceleratorLabel("", 8)) +} + +func TestStatusWord(t *testing.T) { + assert.Equal(t, "SUCCESS", statusWord("TERMINATED", "SUCCESS")) // result wins + assert.Equal(t, "RUNNING", statusWord("RUNNING", "")) // falls back to lifecycle + assert.Equal(t, "UNKNOWN", statusWord("", "")) +} diff --git a/experimental/air/cmd/get.go b/experimental/air/cmd/get.go new file mode 100644 index 00000000000..2f8bd8dd09e --- /dev/null +++ b/experimental/air/cmd/get.go @@ -0,0 +1,242 @@ +package aircmd + +import ( + "context" + "errors" + "fmt" + "path" + "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/apierr" + "github.com/databricks/databricks-sdk-go/config" + "github.com/databricks/databricks-sdk-go/service/iam" + "github.com/databricks/databricks-sdk-go/service/jobs" + "github.com/spf13/cobra" +) + +// getData is the payload printed by `air get`. The json-tagged fields form +// the machine-readable output; fields tagged `json:"-"` are shown only in the +// human-readable text view. +type getData struct { + RunID string `json:"run_id"` + Status string `json:"status"` + StartedAt *string `json:"started_at"` + DurationSeconds *int64 `json:"duration_seconds"` + AttemptNumber int `json:"attempt_number"` + ExperimentName *string `json:"experiment_name"` + DashboardURL string `json:"dashboard_url"` + MLflowURL *string `json:"mlflow_url"` + + // The fields below are pre-rendered text-view cells, excluded from JSON + // (matching `air get --json`). Each shows "N/A" when its value is + // missing. The styled single-run renderer (render.go) consumes them; the + // Run ID, Status, and MLflow Run cells it draws are styled and hyperlinked + // there rather than stored here. + SubmittedDisplay string `json:"-"` + DurationDisplay string `json:"-"` + ExperimentDisplay string `json:"-"` + UserDisplay string `json:"-"` + AcceleratorsDisplay string `json:"-"` + EnvironmentDisplay string `json:"-"` + MaxRetriesDisplay string `json:"-"` + // TrainingConfigPath is the run's config file, downloaded for the config box. + TrainingConfigPath string `json:"-"` + // Sweep replaces the single-run view for foreach runs. + Sweep *sweepInfo `json:"-"` +} + +// getTemplate is the text-mode layout for a sweep (foreach) run. Single runs are +// drawn by the styled renderer in render.go and never reach this template; it is +// used only when .Data.Sweep is set. It reads from the JSON envelope, so every +// field is reached through ".Data". +const getTemplate = `Sweep Run ID: {{.Data.RunID}} +Status: {{.Data.Status}} +Total: {{.Data.Sweep.Total}} +Completed: {{.Data.Sweep.Completed}} +Succeeded: {{.Data.Sweep.Succeeded}} +Failed: {{.Data.Sweep.Failed}} +Active: {{.Data.Sweep.Active}} +{{- if .Data.Sweep.Tasks}} + +Sweep Tasks: +{{printf " %-24s %-14s %-12s %s" "TASK" "RUN ID" "STATUS" "EXPERIMENT"}} +{{- range .Data.Sweep.Tasks}} +{{printf " %-24s %-14s %-12s %s" .TaskKey .RunID .Status .Experiment}} +{{- end}} +{{- end}} +` + +// errNoProfile is the actionable message shown when no credentials are +// configured: no default profile, no --profile (-p), and no auth environment. +var errNoProfile = errors.New("no default profile is set: pass --profile (-p) or configure a default profile in your .databrickscfg") + +// authError classifies a workspace-client or Me() probe failure. Only genuinely +// auth-shaped errors surface as UNAUTHENTICATED/PERMANENT: missing profile, +// SDK auth wrappers, or an API 401/403. Anything else (network blip, 429, 5xx) +// is transient and reported as INTERNAL_ERROR/TRANSIENT so the caller can retry. +func authError(ctx context.Context, cmd *cobra.Command, err error) error { + if errors.Is(err, config.ErrCannotConfigureDefault) { + return renderError(ctx, cmd, "UNAUTHENTICATED", "PERMANENT", false, errNoProfile) + } + if errors.Is(err, apierr.ErrUnauthenticated) || errors.Is(err, apierr.ErrPermissionDenied) { + return renderError(ctx, cmd, "UNAUTHENTICATED", "PERMANENT", false, + fmt.Errorf("authentication was not successful: %w", err)) + } + return renderError(ctx, cmd, "INTERNAL_ERROR", "TRANSIENT", true, + fmt.Errorf("failed to verify authentication: %w", err)) +} + +// newGetCommand returns the `air get JOB_RUN_ID` command, which shows status, +// configuration, and timing details for a specific run. +func newGetCommand() *cobra.Command { + cmd := &cobra.Command{ + Use: "get JOB_RUN_ID", + Args: root.ExactArgs(1), + Short: "Show status, configuration, and timing details for a specific run", + Annotations: map[string]string{ + "template": getTemplate, + }, + } + + // Resolve and authenticate the workspace client up front so an auth failure + // fails fast here, before any run status or config is fetched or printed. + // ErrAlreadyPrinted passes through (it was handled upstream); other failures + // become an actionable auth error (JSON envelope in -o json mode). + 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() + w := cmdctx.WorkspaceClient(ctx) + + 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])) + } + + // Validate authentication against the workspace before fetching or + // rendering anything. MustWorkspaceClient's Config.Authenticate only + // attaches credentials (e.g. it does not check a PAT server-side), so + // without this a bad credential would surface as a confusing failure + // mid-render instead of a clear "not authenticated" error here. + if _, err := w.CurrentUser.Me(ctx, iam.MeRequest{}); err != nil { + return authError(ctx, cmd, err) + } + + run, err := w.Jobs.GetRun(ctx, jobs.GetRunRequest{RunId: runID}) + if err != nil { + // The backend returns this when the run ID is unknown to the user. + 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", runID)) + } + return renderError(ctx, cmd, "INTERNAL_ERROR", "TRANSIENT", true, + fmt.Errorf("failed to get status for run %d: %w", runID, err)) + } + + workspaceID, err := w.CurrentWorkspaceID(ctx) + if err != nil { + return renderError(ctx, cmd, "INTERNAL_ERROR", "TRANSIENT", true, + fmt.Errorf("failed to get workspace id for run %d: %w", runID, err)) + } + + data := buildGetData(run) + data.DashboardURL = dashboardURL(w.Config.Host, runID, workspaceID) + ids := mlflowIDs(ctx, w, run) + if ids != nil { + url := mlflowLogsURL(w.Config.Host, ids) + data.MLflowURL = &url + } + if task := findForEachTask(run); task != nil { + data.Sweep = buildSweepInfo(ctx, w, task) + } else if genAIComputeTask(run) == nil { + enrichFromAiRuntimeTask(run, &data) + } + + if root.OutputType(cmd) != flags.OutputText { + return renderEnvelope(ctx, data) + } + + out := cmd.OutOrStdout() + if data.Sweep != nil { + // A sweep has no single status, config, or timing, so lead with the + // job run link and render the foreach summary table (getTemplate). + fmt.Fprintf(out, "Job Link: %s\n\n", hyperlink(ctx, out, data.DashboardURL, data.DashboardURL)) + return renderEnvelope(ctx, data) + } + + renderRunText(ctx, out, w, run, &data, ids) + return nil + } + + return cmd +} + +// enrichFromAiRuntimeTask fills the config path, experiment, and accelerators +// from the run's ai_runtime_task. Best-effort: empty fields leave the existing +// "N/A" fallbacks in place. +func enrichFromAiRuntimeTask(run *jobs.Run, data *getData) { + task := aiRuntimeTaskOf(run) + if task == nil { + return + } + if len(task.Deployments) > 0 { + d := task.Deployments[0] + if d.CommandPath != "" { + data.TrainingConfigPath = path.Join(path.Dir(d.CommandPath), trainingConfigName) + } + if a := acceleratorLabel(string(d.Compute.AcceleratorType), d.Compute.AcceleratorCount); a != "" { + data.AcceleratorsDisplay = a + } + } + if task.Experiment != "" { + exp := stripExperimentUserPrefix(task.Experiment) + data.ExperimentName = &exp + data.ExperimentDisplay = exp + } +} + +// aiRuntimeTaskOf returns the run's first ai_runtime_task, or nil. +func aiRuntimeTaskOf(run *jobs.Run) *jobs.AiRuntimeTask { + if len(run.Tasks) == 0 { + return nil + } + return run.Tasks[0].AiRuntimeTask +} + +// buildGetData extracts the fields we display from a run. The text-view cells +// are pre-rendered here with their "N/A" fallbacks; the styled renderer adds the +// hyperlinks and colors once the dashboard and MLflow identifiers are known. +func buildGetData(run *jobs.Run) getData { + data := getData{ + RunID: strconv.FormatInt(run.RunId, 10), + Status: runStatus(run.State), + StartedAt: startedAt(run), + DurationSeconds: durationSeconds(run), + AttemptNumber: latestAttemptNumber(run), + ExperimentName: experimentName(run), + } + data.SubmittedDisplay = submittedDisplay(run) + data.DurationDisplay = na + if data.DurationSeconds != nil { + data.DurationDisplay = formatDuration(*data.DurationSeconds) + } + data.ExperimentDisplay = na + if data.ExperimentName != nil { + data.ExperimentDisplay = *data.ExperimentName + } + data.UserDisplay = orNA(run.CreatorUserName) + data.AcceleratorsDisplay = orNA(accelerators(run)) + data.EnvironmentDisplay = orNA(environment(run)) + data.MaxRetriesDisplay = maxRetries(run) + return data +} diff --git a/experimental/air/cmd/get_test.go b/experimental/air/cmd/get_test.go new file mode 100644 index 00000000000..aa23e58159e --- /dev/null +++ b/experimental/air/cmd/get_test.go @@ -0,0 +1,262 @@ +package aircmd + +import ( + "bytes" + "encoding/json" + "errors" + "net/http" + "net/http/httptest" + "testing" + "text/template" + + "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/apierr" + "github.com/databricks/databricks-sdk-go/config" + "github.com/databricks/databricks-sdk-go/experimental/mocks" + "github.com/databricks/databricks-sdk-go/service/jobs" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/mock" + "github.com/stretchr/testify/require" +) + +// renderGet renders the get template against the JSON envelope, exactly as the +// command does for a sweep run, so the test covers the real template branches. +func renderGet(t *testing.T, data getData) string { + t.Helper() + tmpl, err := template.New("get").Parse(getTemplate) + require.NoError(t, err) + var buf bytes.Buffer + require.NoError(t, tmpl.Execute(&buf, envelope{V: envelopeVersion, Data: data})) + return buf.String() +} + +// TestGetCommandShape locks in that `get` takes the run id directly as +// `air get JOB_RUN_ID` and has no `run` subcommand (it was collapsed back into +// `get`). The acceptance test exercises the happy path end to end. +func TestGetCommandShape(t *testing.T) { + cmd := newGetCommand() + assert.Equal(t, "get JOB_RUN_ID", cmd.Use) + assert.Empty(t, cmd.Commands(), "get must not register subcommands") + // ExactArgs(1): exactly one run id is required. + assert.NoError(t, cmd.Args(cmd, []string{"123"})) + assert.Error(t, cmd.Args(cmd, []string{})) + assert.Error(t, cmd.Args(cmd, []string{"1", "2"})) +} + +func TestGetRunInvalidID(t *testing.T) { + m := mocks.NewMockWorkspaceClient(t) + ctx := cmdctx.SetWorkspaceClient(cmdio.MockDiscard(t.Context()), m.WorkspaceClient) + cmd := withOutput(newGetCommand(), flags.OutputText) + cmd.SetContext(ctx) + + err := cmd.RunE(cmd, []string{"abc"}) + require.Error(t, err) + assert.Contains(t, err.Error(), "invalid JOB_RUN_ID") +} + +// notFoundGetServer serves the auth probe plus a runs/get that reports the run +// as missing (400 INVALID_PARAMETER_VALUE, which the SDK maps to +// ErrResourceDoesNotExist for this path). +func notFoundGetServer(t *testing.T) *httptest.Server { + t.Helper() + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == "/api/2.2/jobs/runs/get" { + w.WriteHeader(http.StatusBadRequest) + _, _ = w.Write([]byte(`{"error_code":"INVALID_PARAMETER_VALUE","message":"Run 5 does not exist."}`)) + return + } + // Me() probe and any other config discovery. + _, _ = w.Write([]byte(`{"userName":"u@example.com"}`)) + })) + t.Cleanup(srv.Close) + return srv +} + +func TestGetRunNotFound(t *testing.T) { + srv := notFoundGetServer(t) + ctx := cmdctx.SetWorkspaceClient(cmdio.MockDiscard(t.Context()), newTestWorkspaceClient(t, srv.URL)) + cmd := withOutput(newGetCommand(), flags.OutputText) + cmd.SetContext(ctx) + + err := cmd.RunE(cmd, []string{"5"}) + require.Error(t, err) + assert.Contains(t, err.Error(), "run 5 not found") +} + +func TestGetRunAuthFailed(t *testing.T) { + m := mocks.NewMockWorkspaceClient(t) + // A genuine auth failure (permission denied) is validated before the run is + // fetched, so GetRun is never reached and nothing is rendered. + m.GetMockCurrentUserAPI().EXPECT().Me(mock.Anything, mock.Anything).Return(nil, apierr.ErrPermissionDenied) + ctx := cmdctx.SetWorkspaceClient(cmdio.MockDiscard(t.Context()), m.WorkspaceClient) + cmd := withOutput(newGetCommand(), flags.OutputText) + cmd.SetContext(ctx) + + err := cmd.RunE(cmd, []string{"5"}) + require.Error(t, err) + assert.Contains(t, err.Error(), "authentication was not successful") +} + +func TestGetRunAuthTransient(t *testing.T) { + m := mocks.NewMockWorkspaceClient(t) + // A transient failure at the auth probe must not be misreported as an auth + // error; it surfaces as a retryable internal error instead. + m.GetMockCurrentUserAPI().EXPECT().Me(mock.Anything, mock.Anything).Return(nil, errors.New("connection reset")) + ctx := cmdctx.SetWorkspaceClient(cmdio.MockDiscard(t.Context()), m.WorkspaceClient) + cmd := withOutput(newGetCommand(), flags.OutputText) + cmd.SetContext(ctx) + + err := cmd.RunE(cmd, []string{"5"}) + require.Error(t, err) + assert.Contains(t, err.Error(), "failed to verify authentication") + assert.NotContains(t, err.Error(), "authentication was not successful") +} + +func TestAuthError(t *testing.T) { + ctx := cmdio.MockDiscard(t.Context()) + cmd := withOutput(newGetCommand(), flags.OutputText) + + // No configurable credentials maps to the missing-profile hint. + noProfile := authError(ctx, cmd, config.ErrCannotConfigureDefault) + require.Error(t, noProfile) + assert.Contains(t, noProfile.Error(), "no default profile is set") + + // A 401 / 403 (via the SDK sentinels) is a real auth failure. + unauth := authError(ctx, cmd, apierr.ErrUnauthenticated) + require.Error(t, unauth) + assert.Contains(t, unauth.Error(), "authentication was not successful") + + denied := authError(ctx, cmd, apierr.ErrPermissionDenied) + require.Error(t, denied) + assert.Contains(t, denied.Error(), "authentication was not successful") + + transient := authError(ctx, cmd, errors.New("connection reset")) + require.Error(t, transient) + assert.Contains(t, transient.Error(), "failed to verify authentication") + assert.Contains(t, transient.Error(), "connection reset") +} + +func TestGetRunNotFoundJSON(t *testing.T) { + var buf bytes.Buffer + srv := notFoundGetServer(t) + ctx := cmdctx.SetWorkspaceClient(t.Context(), newTestWorkspaceClient(t, srv.URL)) + ctx = cmdio.InContext(ctx, cmdio.NewIO(ctx, flags.OutputJSON, nil, &buf, &buf, "", "")) + cmd := withOutput(newGetCommand(), flags.OutputJSON) + cmd.SetContext(ctx) + + // In JSON mode the not-found error is a structured envelope, not a bare error. + err := cmd.RunE(cmd, []string{"5"}) + require.ErrorIs(t, err, root.ErrAlreadyPrinted) + + var got errorEnvelope + require.NoError(t, json.Unmarshal(buf.Bytes(), &got)) + assert.Equal(t, jsonError{Code: "NOT_FOUND", Kind: "NOT_FOUND", Message: "run 5 not found: check the run ID and that it is a job run ID"}, got.Error) +} + +func TestGetTemplateSweep(t *testing.T) { + out := renderGet(t, getData{ + RunID: "456", + Status: "RUNNING", + Sweep: &sweepInfo{ + Total: 4, Completed: 2, Succeeded: 1, Failed: 1, Active: 2, + Tasks: []sweepTask{ + {TaskKey: "iter_0", RunID: "789", Status: "SUCCESS", Experiment: "my-exp"}, + {TaskKey: "iter_1", RunID: "790", Status: "FAILED", Experiment: "my-exp"}, + }, + }, + }) + assert.Contains(t, out, "Sweep Run ID: 456") + assert.Contains(t, out, "Total: 4") + assert.Contains(t, out, "Sweep Tasks:") + assert.Contains(t, out, "iter_0") + assert.Contains(t, out, "iter_1") + assert.Contains(t, out, "FAILED") + assert.Contains(t, out, "my-exp") +} + +func TestGetTemplateSweepNoTasks(t *testing.T) { + // A sweep whose iterations haven't materialized yet: counts show, but the + // task table header is hidden. + out := renderGet(t, getData{ + RunID: "456", + Status: "RUNNING", + Sweep: &sweepInfo{Total: 4, Active: 4}, + }) + assert.Contains(t, out, "Sweep Run ID: 456") + assert.Contains(t, out, "Total: 4") + assert.NotContains(t, out, "Sweep Tasks:") +} + +func TestBuildGetData(t *testing.T) { + run := &jobs.Run{ + RunId: 123, + CreatorUserName: "me@example.com", + StartTime: 1700000000000, + EndTime: 1700000012000, + State: &jobs.RunState{ResultState: jobs.RunResultStateSuccess}, + Tasks: []jobs.RunTask{{ + AttemptNumber: 1, + GenAiComputeTask: &jobs.GenAiComputeTask{ + MlflowExperimentName: "/Users/me@example.com/exp", + Compute: &jobs.ComputeConfig{NumGpus: 8, GpuType: "GPU_8xH100"}, + }, + }}, + } + d := buildGetData(run) + assert.Equal(t, "123", d.RunID) + assert.Equal(t, "SUCCESS", d.Status) + assert.Equal(t, 1, d.AttemptNumber) + assert.Equal(t, "2023-11-14 22:13 UTC", d.SubmittedDisplay) + assert.Equal(t, "me@example.com", d.UserDisplay) + assert.Equal(t, "8x H100", d.AcceleratorsDisplay) + assert.Equal(t, "12s", d.DurationDisplay) + assert.Equal(t, "exp", d.ExperimentDisplay) + require.NotNil(t, d.ExperimentName) + assert.Equal(t, "exp", *d.ExperimentName) + require.NotNil(t, d.DurationSeconds) + assert.Equal(t, int64(12), *d.DurationSeconds) +} + +func TestEnrichFromAiRuntimeTask(t *testing.T) { + t.Run("fills config path, experiment, and accelerators", func(t *testing.T) { + run := &jobs.Run{RunId: 5, Tasks: []jobs.RunTask{{ + AiRuntimeTask: &jobs.AiRuntimeTask{ + Experiment: "/Users/me@example.com/my-exp", + Deployments: []jobs.DeploymentSpec{{ + CommandPath: "/Workspace/run/command.sh", + Compute: jobs.ComputeSpec{AcceleratorType: jobs.ComputeSpecAcceleratorTypeGpu1xA10, AcceleratorCount: 1}, + }}, + }, + }}} + data := &getData{ExperimentDisplay: na, AcceleratorsDisplay: na} + enrichFromAiRuntimeTask(run, data) + assert.Equal(t, "/Workspace/run/training_config.yaml", data.TrainingConfigPath) + assert.Equal(t, "my-exp", data.ExperimentDisplay) + require.NotNil(t, data.ExperimentName) + assert.Equal(t, "my-exp", *data.ExperimentName) + assert.Equal(t, "1x A10", data.AcceleratorsDisplay) + }) + + t.Run("leaves fallbacks when the run has no ai_runtime_task", func(t *testing.T) { + run := &jobs.Run{RunId: 5, Tasks: []jobs.RunTask{{}}} + data := &getData{ExperimentDisplay: na, AcceleratorsDisplay: na} + enrichFromAiRuntimeTask(run, data) + assert.Empty(t, data.TrainingConfigPath) + assert.Equal(t, na, data.ExperimentDisplay) + assert.Equal(t, na, data.AcceleratorsDisplay) + }) +} + +func TestBuildGetDataEmpty(t *testing.T) { + // A run with no tasks, creator, or timing renders every text cell as "N/A". + d := buildGetData(&jobs.Run{RunId: 7}) + assert.Equal(t, "7", d.RunID) + assert.Equal(t, na, d.SubmittedDisplay) + assert.Equal(t, na, d.DurationDisplay) + assert.Equal(t, na, d.ExperimentDisplay) + assert.Equal(t, na, d.UserDisplay) + assert.Equal(t, na, d.AcceleratorsDisplay) +} diff --git a/experimental/air/cmd/joblist.go b/experimental/air/cmd/joblist.go new file mode 100644 index 00000000000..dd786d8b496 --- /dev/null +++ b/experimental/air/cmd/joblist.go @@ -0,0 +1,180 @@ +package aircmd + +import ( + "context" + "errors" + "fmt" + "net/http" + + "github.com/databricks/databricks-sdk-go" + "github.com/databricks/databricks-sdk-go/apierr" + "github.com/databricks/databricks-sdk-go/service/jobs" + "golang.org/x/sync/errgroup" +) + +// hydrateConcurrency bounds the parallel runs/get calls when hydrating a batch +// of run ids from the AiTrainingService index. +const hydrateConcurrency = 16 + +// firstTask returns the run's first task, unwrapping a foreach sweep to the +// iterated task, or nil when the run has no tasks. +func firstTask(r *jobs.Run) *jobs.RunTask { + if len(r.Tasks) == 0 { + return nil + } + t := &r.Tasks[0] + if t.AiRuntimeTask != nil || t.GenAiComputeTask != nil { + return t + } + if t.ForEachTask != nil { + // The foreach inner task is a jobs.Task; expose the AI fields via a + // synthetic RunTask so callers can treat both shapes uniformly. + return &jobs.RunTask{ + AiRuntimeTask: t.ForEachTask.Task.AiRuntimeTask, + GenAiComputeTask: t.ForEachTask.Task.GenAiComputeTask, + } + } + return t +} + +// isAirRun reports whether a run is an AI runtime workload: an ai_runtime_task, +// or a legacy gen_ai_compute_task with a training script. +func isAirRun(r *jobs.Run) bool { + t := firstTask(r) + if t == nil { + return false + } + return t.AiRuntimeTask != nil || + (t.GenAiComputeTask != nil && t.GenAiComputeTask.TrainingScriptPath != "") +} + +// isSweep reports whether the run's first task fans out into iterations. +func isSweep(r *jobs.Run) bool { + return len(r.Tasks) > 0 && r.Tasks[0].ForEachTask != nil +} + +// taskRunID returns the run id of the AIR task, used to fetch its MLflow output. +func taskRunID(r *jobs.Run) int64 { + if len(r.Tasks) == 0 { + return 0 + } + return r.Tasks[0].RunId +} + +// jobExperiment returns the run's MLflow experiment name (user-folder prefix +// stripped), or "" when there is none. +func jobExperiment(r *jobs.Run) string { + t := firstTask(r) + switch { + case t == nil: + return "" + case t.AiRuntimeTask != nil && t.AiRuntimeTask.Experiment != "": + return stripExperimentUserPrefix(t.AiRuntimeTask.Experiment) + case t.GenAiComputeTask != nil && t.GenAiComputeTask.MlflowExperimentName != "": + return stripExperimentUserPrefix(t.GenAiComputeTask.MlflowExperimentName) + } + return "" +} + +// jobCompute returns the run's accelerator type and count, or ("", 0) when it +// has none. +func jobCompute(r *jobs.Run) (string, int) { + t := firstTask(r) + switch { + case t == nil: + return "", 0 + case t.AiRuntimeTask != nil && len(t.AiRuntimeTask.Deployments) > 0: + c := t.AiRuntimeTask.Deployments[0].Compute + return string(c.AcceleratorType), c.AcceleratorCount + case t.GenAiComputeTask != nil && t.GenAiComputeTask.Compute != nil: + c := t.GenAiComputeTask.Compute + return c.GpuType, c.NumGpus + } + return "", 0 +} + +// jobTiming returns the run's start and end times (epoch ms), preferring the +// first task's window so a run reports its task attempt rather than the wrapper. +func jobTiming(r *jobs.Run) (startMillis, endMillis int64) { + startMillis, endMillis = r.StartTime, r.EndTime + if len(r.Tasks) > 0 { + if t := r.Tasks[0]; t.StartTime > 0 { + startMillis = t.StartTime + endMillis = t.EndTime + } + } + return startMillis, endMillis +} + +// isTerminal reports whether a run has finished and its details are immutable, +// so its row is safe to cache. +func isTerminal(r *jobs.Run) bool { + if r.State == nil { + return false + } + switch r.State.LifeCycleState { + case jobs.RunLifeCycleStateTerminated, jobs.RunLifeCycleStateInternalError, jobs.RunLifeCycleStateSkipped: + return true + default: + return false + } +} + +// baseRunToRun converts a runs/list BaseRun into a Run so the display helpers +// operate on a single type. The two share every field the list path reads. +func baseRunToRun(b jobs.BaseRun) *jobs.Run { + return &jobs.Run{ + RunId: b.RunId, + RunName: b.RunName, + CreatorUserName: b.CreatorUserName, + StartTime: b.StartTime, + EndTime: b.EndTime, + State: b.State, + Tasks: b.Tasks, + } +} + +// fetchJobRun fetches a single run via runs/get; expand_tasks is implied by the +// typed SDK, which returns the ai_runtime_task in the response. +func fetchJobRun(ctx context.Context, w *databricks.WorkspaceClient, runID int64) (*jobs.Run, error) { + run, err := w.Jobs.GetRun(ctx, jobs.GetRunRequest{RunId: runID}) + if err != nil { + return nil, fmt.Errorf("failed to get run %d: %w", runID, err) + } + return run, nil +} + +// hydrateJobRuns fetches the given run ids concurrently via runs/get, preserving +// input order. runs/get enforces per-run view ACLs, so an id the caller can't +// view (403) or that has been purged (404) is dropped; any other error is +// systemic and fails the whole batch. +func hydrateJobRuns(ctx context.Context, w *databricks.WorkspaceClient, ids []int64) ([]*jobs.Run, error) { + runs := make([]*jobs.Run, len(ids)) + g, gctx := errgroup.WithContext(ctx) + g.SetLimit(hydrateConcurrency) + for i, id := range ids { + g.Go(func() error { + run, err := fetchJobRun(gctx, w, id) + if err != nil { + if apiErr, ok := errors.AsType[*apierr.APIError](err); ok && + (apiErr.StatusCode == http.StatusForbidden || apiErr.StatusCode == http.StatusNotFound) { + return nil // not viewable or purged: drop this id + } + return fmt.Errorf("failed to get run %d: %w", id, err) + } + runs[i] = run + return nil + }) + } + if err := g.Wait(); err != nil { + return nil, err + } + + hydrated := make([]*jobs.Run, 0, len(runs)) + for _, run := range runs { + if run != nil { + hydrated = append(hydrated, run) + } + } + return hydrated, nil +} diff --git a/experimental/air/cmd/joblist_test.go b/experimental/air/cmd/joblist_test.go new file mode 100644 index 00000000000..11e5efac8e9 --- /dev/null +++ b/experimental/air/cmd/joblist_test.go @@ -0,0 +1,47 @@ +package aircmd + +import ( + "net/http" + "net/http/httptest" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// runGetServer serves one runs/get body and a stub for the SDK config probe. +func runGetServer(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.2/jobs/runs/get" { + _, _ = w.Write([]byte(body)) + return + } + _, _ = w.Write([]byte(`{}`)) + })) + t.Cleanup(srv.Close) + return srv +} + +func TestFetchJobRunParsesAiRuntimeTask(t *testing.T) { + body := `{ + "run_id": 5, + "tasks": [{ + "ai_runtime_task": { + "experiment": "my-exp", + "deployments": [{ + "command_path": "/Workspace/Users/me/.air/cli_launch/my-exp/my-exp_abc/command.sh", + "compute": {"accelerator_count": 1, "accelerator_type": "GPU_1xA10"} + }] + } + }] + }` + srv := runGetServer(t, body) + + run, err := fetchJobRun(t.Context(), newTestWorkspaceClient(t, srv.URL), 5) + require.NoError(t, err) + assert.Equal(t, "my-exp", jobExperiment(run)) + gpuType, count := jobCompute(run) + assert.Equal(t, "GPU_1xA10", gpuType) + assert.Equal(t, 1, count) +} diff --git a/experimental/air/cmd/list.go b/experimental/air/cmd/list.go new file mode 100644 index 00000000000..7d5c703a39c --- /dev/null +++ b/experimental/air/cmd/list.go @@ -0,0 +1,306 @@ +package aircmd + +import ( + "context" + "fmt" + + "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/listing" + "github.com/databricks/databricks-sdk-go/service/iam" + "github.com/databricks/databricks-sdk-go/service/jobs" + "github.com/spf13/cobra" + "golang.org/x/sync/errgroup" +) + +// maxListScan bounds how many runs `air list` inspects while looking for AIR runs +// that match the filters. runs/list returns runs of every kind, so this caps the +// work on a workspace with a large run history. +const maxListScan = 2000 + +// jobsPageLimit is the per-request page size for runs/list; enrichConcurrency +// bounds the parallel MLflow lookups. +const ( + jobsPageLimit = 25 + enrichConcurrency = 8 +) + +// listData is the payload printed by `air list`. +type listData struct { + Rows []listRow `json:"runs"` +} + +// listRow is one run in the list. The json-tagged fields form the +// machine-readable output; fields tagged `json:"-"` are shown only in the +// human-readable table. +type listRow struct { + RunID string `json:"run_id"` + RunName string `json:"run_name"` + User string `json:"user"` + Status string `json:"status"` + StartedAt *string `json:"started_at"` + IsSweep bool `json:"is_sweep"` + + // Experiment, Duration, MLflowURL and Accelerators are table-only columns, + // omitted from JSON to match `air list --json`. + Experiment string `json:"-"` + Duration string `json:"-"` + MLflowURL string `json:"-"` + Accelerators string `json:"-"` +} + +// listedRun pairs a row with its task run id, so the MLflow link can be fetched +// after the run has been filtered in. +type listedRun struct { + row listRow + taskRunID int64 +} + +// listQuery holds the resolved inputs to a runFetcher. +type listQuery struct { + activeOnly bool + allUsers bool + userFilter string + currentUser string + filters listFilters + fetchMLflow bool + limit int +} + +func newListCommand() *cobra.Command { + var ( + limit int + allStatus bool + allUsers bool + filters []string + ) + + cmd := &cobra.Command{ + Use: "list", + Args: root.NoArgs, + Short: "List your active runs for the current profile (use --all-status for finished runs)", + } + + cmd.PreRunE = root.MustWorkspaceClient + + cmd.Flags().IntVar(&limit, "limit", 20, "Maximum number of runs to show") + cmd.Flags().BoolVar(&allStatus, "all-status", false, "Show runs in all states (default: active only)") + cmd.Flags().BoolVar(&allUsers, "all-users", false, "Show runs from all users") + cmd.Flags().StringArrayVar(&filters, "filter", nil, "Filter runs, e.g. experiment=foo* (repeatable)") + + cmd.RunE = func(cmd *cobra.Command, args []string) error { + ctx := cmd.Context() + w := cmdctx.WorkspaceClient(ctx) + + if limit <= 0 { + return fmt.Errorf("invalid --limit %d: must be a positive integer", limit) + } + + f, err := parseListFilters(filters) + if err != nil { + return err + } + + // An explicit user= filter wins; otherwise default to the current user + // unless --all-users is set. runs/list has no creator param, so the + // creator is matched while scanning. + userFilter := f.User + var currentUser string + if userFilter == "" && !allUsers { + me, err := w.CurrentUser.Me(ctx, iam.MeRequest{}) + if err != nil { + return fmt.Errorf("failed to resolve current user: %w", err) + } + currentUser = me.UserName + userFilter = currentUser + } + + fetcher := newRunFetcher(ctx, w, listQuery{ + activeOnly: !allStatus, + allUsers: allUsers, + userFilter: userFilter, + currentUser: currentUser, + filters: f, + fetchMLflow: root.OutputType(cmd) == flags.OutputText, + limit: limit, + }) + + // JSON prints the newest `limit` runs once. Text renders the table: + // navigable in a terminal (paging in older runs on demand), printed once + // when piped. + if root.OutputType(cmd) == flags.OutputJSON { + rows, err := fetcher.next(limit) + if err != nil { + return err + } + warnIfTruncated(ctx, fetcher) + return renderEnvelope(ctx, listData{Rows: rows}) + } + return renderListText(cmd, fetcher, limit) + } + + return cmd +} + +// listStrategy is a source of matching runs, pulled in batches. Two implement it: +// jobsScanStrategy pages runs/list; indexStrategy hydrates the AiTrainingService +// index. The fetcher wraps whichever is chosen. +type listStrategy interface { + // next returns up to want more matching runs (already row-built + task id). + next(want int) ([]listedRun, error) + // done reports whether the source has no more runs to yield. + done() bool + // truncated reports whether a safety cap stopped the scan short of the end. + truncated() bool +} + +// runFetcher yields matching rows in batches, driving both one-shot output and +// the interactive table's lazy paging. It wraps a listStrategy and adds the +// shared tail: MLflow enrichment (text only) and row projection. +type runFetcher struct { + ctx context.Context + w *databricks.WorkspaceClient + fetchMLflow bool + strategy listStrategy + + exhausted bool +} + +func newRunFetcher(ctx context.Context, w *databricks.WorkspaceClient, q listQuery) *runFetcher { + return &runFetcher{ + ctx: ctx, + w: w, + fetchMLflow: q.fetchMLflow, + strategy: newListStrategy(ctx, w, q), + } +} + +// newListStrategy picks the fetch source. The AiTrainingService index serves only +// the caller's own runs, so it's used for an all-status self-scoped list; if the +// index load fails (e.g. endpoint unavailable in this workspace), we fall back to +// the Jobs scan so the command still returns. Everything else — the default +// active list, --all-users, and --all-status for another user — uses the scan. +func newListStrategy(ctx context.Context, w *databricks.WorkspaceClient, q listQuery) listStrategy { + useIndex := !q.activeOnly && !q.allUsers && (q.userFilter == "" || q.userFilter == q.currentUser) + if !useIndex { + return newJobsScanStrategy(ctx, w, q) + } + idx := newIndexStrategy(ctx, w, q, q.limit) + if err := idx.load(); err != nil { + log.Debugf(ctx, "air list: AiTrainingService index unavailable, falling back to Jobs scan: %v", err) + return newJobsScanStrategy(ctx, w, q) + } + return idx +} + +// next pulls the next batch from the strategy, enriches it with MLflow links for +// text output, and projects it to rows. It sets exhausted once the strategy is +// drained so the interactive table knows to stop paging. +func (f *runFetcher) next(want int) ([]listRow, error) { + entries, err := f.strategy.next(want) + if err != nil { + return nil, err + } + f.exhausted = f.strategy.done() + + // MLflow links appear only in the text table, so the per-run get-output + // lookups are skipped for JSON output (which omits the column anyway). + if f.fetchMLflow { + setMLflowLinks(f.ctx, f.w, entries) + } + + rows := make([]listRow, len(entries)) + for i, e := range entries { + rows[i] = e.row + } + return rows, nil +} + +// jobsScanStrategy pages Jobs runs/list, keeping the AIR runs that match the user +// and filters. It buffers a page's leftover runs so successive next() calls +// resume where the last stopped. +type jobsScanStrategy struct { + ctx context.Context + w *databricks.WorkspaceClient + iter listing.Iterator[jobs.BaseRun] + userFilter string + filters listFilters + + scanned int +} + +func newJobsScanStrategy(ctx context.Context, w *databricks.WorkspaceClient, q listQuery) *jobsScanStrategy { + req := jobs.ListRunsRequest{ + RunType: jobs.RunTypeSubmitRun, + ExpandTasks: true, + Limit: jobsPageLimit, + ActiveOnly: q.activeOnly, + } + return &jobsScanStrategy{ + ctx: ctx, + w: w, + iter: w.Jobs.ListRuns(ctx, req), + userFilter: q.userFilter, + filters: q.filters, + } +} + +func (s *jobsScanStrategy) next(want int) ([]listedRun, error) { + var entries []listedRun + for len(entries) < want && s.scanned < maxListScan && s.iter.HasNext(s.ctx) { + base, err := s.iter.Next(s.ctx) + if err != nil { + return nil, fmt.Errorf("failed to list runs: %w", err) + } + s.scanned++ + + run := baseRunToRun(base) + if !isAirRun(run) { + continue + } + if s.userFilter != "" && run.CreatorUserName != s.userFilter { + continue + } + if !s.filters.matches(run) { + continue + } + entries = append(entries, listedRun{row: buildListRow(run), taskRunID: taskRunID(run)}) + } + return entries, nil +} + +func (s *jobsScanStrategy) done() bool { + return s.scanned >= maxListScan || !s.iter.HasNext(s.ctx) +} + +func (s *jobsScanStrategy) truncated() bool { + return s.scanned >= maxListScan +} + +// warnIfTruncated logs when a scan hit its safety cap, so one-shot output signals +// its results may be incomplete. +func warnIfTruncated(ctx context.Context, f *runFetcher) { + if f.strategy.truncated() { + log.Warnf(ctx, "air list: stopped after scanning %d runs; results may be incomplete", maxListScan) + } +} + +// setMLflowLinks fills in each row's MLflow link in parallel, best-effort: a row +// whose IDs can't be resolved keeps its "-" placeholder. +func setMLflowLinks(ctx context.Context, w *databricks.WorkspaceClient, entries []listedRun) { + var g errgroup.Group + g.SetLimit(enrichConcurrency) + for i := range entries { + g.Go(func() error { + if ids := mlflowIDsForTask(ctx, w, entries[i].taskRunID); ids != nil { + entries[i].row.MLflowURL = mlflowLogsURL(w.Config.Host, ids) + } + return nil + }) + } + // mlflowIDsForTask never returns an error (it logs and yields nil), so Wait can't fail. + _ = g.Wait() +} diff --git a/experimental/air/cmd/list_cache.go b/experimental/air/cmd/list_cache.go new file mode 100644 index 00000000000..6f0b95a97d3 --- /dev/null +++ b/experimental/air/cmd/list_cache.go @@ -0,0 +1,81 @@ +package aircmd + +import ( + "context" + "time" + + "github.com/databricks/cli/libs/cache" +) + +// The AiTrainingService index path caches hydrated terminal runs on disk: +// terminal runs are immutable, so once we've paid for runs/get + get-output + +// MLflow we persist the finished row and skip those round-trips next time. The +// TTL matches AICM's ~60-day retention, after which the run drops out of the +// index anyway. +const ( + listCacheComponent = "air-list-runs" + listCacheTTL = 60 * 24 * time.Hour +) + +// listCacheKey fingerprints a cached run. Host isolates workspaces (a Jobs run +// id is unique only within one), matching how libs/cache namespaces entries. +type listCacheKey struct { + Host string `json:"host"` + RunID int64 `json:"run_id"` +} + +// cachedRun is the persisted value: every listRow field (including the +// table-only columns, which listRow tags json:"-" and so wouldn't survive a +// direct marshal), the filter inputs, and the submit time. +type cachedRun struct { + RunID string `json:"run_id"` + RunName string `json:"run_name"` + User string `json:"user"` + Status string `json:"status"` + StartedAt *string `json:"started_at"` + IsSweep bool `json:"is_sweep"` + Experiment string `json:"experiment"` + Duration string `json:"duration"` + MLflowURL string `json:"mlflow_url"` + Accelerators string `json:"accelerators"` + Fields filterFields `json:"filter_fields"` + SubmitTimeMs int64 `json:"submit_time_ms"` +} + +func (c cachedRun) toRow() listRow { + return listRow{ + RunID: c.RunID, RunName: c.RunName, User: c.User, Status: c.Status, + StartedAt: c.StartedAt, IsSweep: c.IsSweep, Experiment: c.Experiment, + Duration: c.Duration, MLflowURL: c.MLflowURL, Accelerators: c.Accelerators, + } +} + +func cachedRunFromRow(r listRow, fields filterFields, submitTimeMs int64) cachedRun { + return cachedRun{ + RunID: r.RunID, RunName: r.RunName, User: r.User, Status: r.Status, + StartedAt: r.StartedAt, IsSweep: r.IsSweep, Experiment: r.Experiment, + Duration: r.Duration, MLflowURL: r.MLflowURL, Accelerators: r.Accelerators, + Fields: fields, SubmitTimeMs: submitTimeMs, + } +} + +// newListCache builds the cache for the index path. It fails open, so a nil +// return (or any cache error) just means every run is hydrated from the API. +func newListCache(ctx context.Context) *cache.Cache { + return cache.NewCache(ctx, listCacheComponent, listCacheTTL, nil) +} + +// cachedRow returns the cached row and its filter fields for a run, or +// (zero, zero, false) on miss. +func cachedRow(ctx context.Context, c *cache.Cache, host string, runID int64) (listRow, filterFields, bool) { + entry, ok := cache.Get[cachedRun](ctx, c, listCacheKey{Host: host, RunID: runID}) + if !ok { + return listRow{}, filterFields{}, false + } + return entry.toRow(), entry.Fields, true +} + +// putRow caches a terminal run's finished row and filter fields under its submit time. +func putRow(ctx context.Context, c *cache.Cache, host string, runID, submitTimeMs int64, row listRow, fields filterFields) { + cache.Put(ctx, c, listCacheKey{Host: host, RunID: runID}, cachedRunFromRow(row, fields, submitTimeMs)) +} diff --git a/experimental/air/cmd/list_cache_test.go b/experimental/air/cmd/list_cache_test.go new file mode 100644 index 00000000000..2ed076a98e2 --- /dev/null +++ b/experimental/air/cmd/list_cache_test.go @@ -0,0 +1,105 @@ +package aircmd + +import ( + "testing" + + "github.com/databricks/databricks-sdk-go/service/jobs" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestListCacheRoundTrip(t *testing.T) { + t.Setenv("DATABRICKS_CACHE_DIR", t.TempDir()) + ctx := t.Context() + c := newListCache(ctx) + + _, _, ok := cachedRow(ctx, c, "https://host.test", 42) + require.False(t, ok, "miss before write") + + row := listRow{RunID: "42", Experiment: "exp", Status: "SUCCESS"} + fields := filterFields{Experiment: "exp", GPUType: "GPU_1xA10", GPUCount: 1} + putRow(ctx, c, "https://host.test", 42, 1700000000000, row, fields) + + got, gotFields, ok := cachedRow(ctx, c, "https://host.test", 42) + require.True(t, ok, "hit after write") + assert.Equal(t, row, got) + assert.Equal(t, fields, gotFields) + + // Different host is a different key. + _, _, ok = cachedRow(ctx, c, "https://other.test", 42) + assert.False(t, ok) +} + +func TestIndexStrategyServesCachedRowWithoutFetch(t *testing.T) { + t.Setenv("DATABRICKS_CACHE_DIR", t.TempDir()) + + refs := []workflowRef{{jobRunID: 7, submitTimeMs: 1000_000}} + srv, hits := indexAndGetServer(t, refs, map[int64]jobs.Run{7: indexRun(7, 1000_000)}, nil, nil) + host := srv.URL + + // Pre-seed the cache for run 7 so hydration should skip runs/get entirely. + ctx := t.Context() + putRow(ctx, newListCache(ctx), host, 7, 1000_000, listRow{RunID: "7", Status: "SUCCESS"}, filterFields{}) + + f := newRunFetcher(ctx, newTestWorkspaceClient(t, host), listQuery{ + userFilter: "me@example.com", currentUser: "me@example.com", limit: 10, + }) + rows, err := f.next(10) + require.NoError(t, err) + require.Len(t, rows, 1) + assert.Equal(t, "7", rows[0].RunID) + assert.Equal(t, 0, hits.get, "cached run must not hit runs/get") +} + +func TestIndexStrategyFiltersCachedRow(t *testing.T) { + t.Setenv("DATABRICKS_CACHE_DIR", t.TempDir()) + + refs := []workflowRef{{jobRunID: 7, submitTimeMs: 1000_000}} + srv, hits := indexAndGetServer(t, refs, map[int64]jobs.Run{7: indexRun(7, 1000_000)}, nil, nil) + host := srv.URL + + // Cache hit under experiment "bar" must be filtered out by an experiment=foo query. + ctx := t.Context() + putRow(ctx, newListCache(ctx), host, 7, 1000_000, + listRow{RunID: "7", Status: "SUCCESS", Experiment: "bar"}, + filterFields{Experiment: "bar"}) + + f := newRunFetcher(ctx, newTestWorkspaceClient(t, host), listQuery{ + userFilter: "me@example.com", currentUser: "me@example.com", limit: 10, + filters: listFilters{Experiment: "foo"}, + }) + rows, err := f.next(10) + require.NoError(t, err) + assert.Empty(t, rows, "cached row not matching the filter must be dropped") + assert.Equal(t, 0, hits.get, "non-matching cached row must not hit runs/get") +} + +func TestIndexStrategyServesMatchingCachedRow(t *testing.T) { + t.Setenv("DATABRICKS_CACHE_DIR", t.TempDir()) + + refs := []workflowRef{{jobRunID: 7, submitTimeMs: 1000_000}} + srv, hits := indexAndGetServer(t, refs, map[int64]jobs.Run{7: indexRun(7, 1000_000)}, nil, nil) + host := srv.URL + + ctx := t.Context() + putRow(ctx, newListCache(ctx), host, 7, 1000_000, + listRow{RunID: "7", Status: "SUCCESS", Experiment: "foo"}, + filterFields{Experiment: "foo"}) + + f := newRunFetcher(ctx, newTestWorkspaceClient(t, host), listQuery{ + userFilter: "me@example.com", currentUser: "me@example.com", limit: 10, + filters: listFilters{Experiment: "foo"}, + }) + rows, err := f.next(10) + require.NoError(t, err) + require.Len(t, rows, 1) + assert.Equal(t, "7", rows[0].RunID) + assert.Equal(t, 0, hits.get, "matching cached row must not hit runs/get") +} + +func TestIsTerminal(t *testing.T) { + assert.True(t, isTerminal(&jobs.Run{State: &jobs.RunState{LifeCycleState: jobs.RunLifeCycleStateTerminated}})) + assert.True(t, isTerminal(&jobs.Run{State: &jobs.RunState{LifeCycleState: jobs.RunLifeCycleStateInternalError}})) + assert.False(t, isTerminal(&jobs.Run{State: &jobs.RunState{LifeCycleState: jobs.RunLifeCycleStateRunning}})) + assert.False(t, isTerminal(&jobs.Run{State: &jobs.RunState{LifeCycleState: jobs.RunLifeCycleStatePending}})) +} diff --git a/experimental/air/cmd/list_filter.go b/experimental/air/cmd/list_filter.go new file mode 100644 index 00000000000..bd2bf8dff6e --- /dev/null +++ b/experimental/air/cmd/list_filter.go @@ -0,0 +1,108 @@ +package aircmd + +import ( + "fmt" + "path" + "strconv" + "strings" + + "github.com/databricks/databricks-sdk-go/service/jobs" +) + +// supportedFilterKeys are the keys accepted by `air list --filter KEY=VALUE`. +var supportedFilterKeys = []string{"accelerator_type", "experiment", "num_accelerators", "user"} + +// hasTaskFilter reports whether any filter is applied to a run's task fields +// (experiment or accelerators), i.e. matched after a run is fetched rather than +// while scanning. The index path uses this to skip its newest-N truncation, so a +// dropped match doesn't shrink the result below --limit. +func (f listFilters) hasTaskFilter() bool { + return f.Experiment != "" || f.AcceleratorType != "" || f.NumAccelerators != nil +} + +// listFilters holds the parsed `--filter` values for `air list`. +type listFilters struct { + // User is an exact creator-email match + User string + // Experiment is a case-insensitive glob + Experiment string + // AcceleratorType is a case-insensitive substring matched against the + // display GPU name (e.g. "H100"). + AcceleratorType string + // NumAccelerators is an exact match against the GPU count. + NumAccelerators *int +} + +func parseListFilters(raw []string) (listFilters, error) { + var f listFilters + for _, item := range raw { + key, value, ok := strings.Cut(item, "=") + if !ok || key == "" { + return listFilters{}, fmt.Errorf("invalid --filter %q: expected KEY=VALUE", item) + } + switch key { + case "user": + f.User = value + case "experiment": + f.Experiment = value + case "accelerator_type": + f.AcceleratorType = value + case "num_accelerators": + n, err := strconv.Atoi(value) + if err != nil || n <= 0 { + return listFilters{}, fmt.Errorf("invalid --filter num_accelerators=%q: must be a positive integer", value) + } + f.NumAccelerators = &n + default: + return listFilters{}, fmt.Errorf("unsupported --filter key %q: supported keys are %s", key, strings.Join(supportedFilterKeys, ", ")) + } + } + return f, nil +} + +// filterFields are the task-filter inputs, cached alongside the row so a cache +// hit can be re-filtered without re-fetching the run. +type filterFields struct { + Experiment string `json:"experiment"` + GPUType string `json:"gpu_type"` + GPUCount int `json:"gpu_count"` +} + +func filterFieldsFromRun(run *jobs.Run) filterFields { + gpuType, count := jobCompute(run) + return filterFields{ + Experiment: jobExperiment(run), + GPUType: gpuType, + GPUCount: count, + } +} + +// matches reports whether a run satisfies the experiment, accelerator-type and +// accelerator-count filters. The user filter is applied separately while +// scanning, since it maps onto the run's creator rather than its task. +func (f listFilters) matches(run *jobs.Run) bool { + return f.matchesFields(filterFieldsFromRun(run)) +} + +// matchesFields is the shared comparator for live runs and cached rows, so the +// two paths can't drift. +func (f listFilters) matchesFields(fields filterFields) bool { + if f.Experiment != "" { + matched, err := path.Match(strings.ToLower(f.Experiment), strings.ToLower(fields.Experiment)) + if err != nil || !matched { + return false + } + } + + if f.AcceleratorType != "" { + display := strings.ToLower(gpuDisplayName(fields.GPUType)) + if !strings.Contains(display, strings.ToLower(f.AcceleratorType)) { + return false + } + } + if f.NumAccelerators != nil && fields.GPUCount != *f.NumAccelerators { + return false + } + + return true +} diff --git a/experimental/air/cmd/list_filter_test.go b/experimental/air/cmd/list_filter_test.go new file mode 100644 index 00000000000..381f5e3a62d --- /dev/null +++ b/experimental/air/cmd/list_filter_test.go @@ -0,0 +1,74 @@ +package aircmd + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestParseListFilters(t *testing.T) { + t.Run("valid", func(t *testing.T) { + f, err := parseListFilters([]string{ + "user=me@example.com", + "experiment=qwen*", + "accelerator_type=H100", + "num_accelerators=8", + }) + require.NoError(t, err) + assert.Equal(t, "me@example.com", f.User) + assert.Equal(t, "qwen*", f.Experiment) + assert.Equal(t, "H100", f.AcceleratorType) + require.NotNil(t, f.NumAccelerators) + assert.Equal(t, 8, *f.NumAccelerators) + }) + + t.Run("unknown key", func(t *testing.T) { + _, err := parseListFilters([]string{"region=us"}) + require.Error(t, err) + assert.Contains(t, err.Error(), "unsupported --filter key") + }) + + t.Run("malformed pair", func(t *testing.T) { + _, err := parseListFilters([]string{"experiment"}) + require.Error(t, err) + assert.Contains(t, err.Error(), "expected KEY=VALUE") + }) + + t.Run("bad num_accelerators", func(t *testing.T) { + _, err := parseListFilters([]string{"num_accelerators=lots"}) + require.Error(t, err) + assert.Contains(t, err.Error(), "num_accelerators") + }) + + t.Run("non-positive num_accelerators", func(t *testing.T) { + _, err := parseListFilters([]string{"num_accelerators=0"}) + require.Error(t, err) + assert.Contains(t, err.Error(), "positive integer") + }) +} + +func TestListFiltersMatches(t *testing.T) { + run := airRun(1, "me@example.com", "GPU_8xH100", 8, "/Users/me@example.com/qwen-eval") + + cases := []struct { + name string + f listFilters + want bool + }{ + {"no filters", listFilters{}, true}, + {"experiment prefix glob", listFilters{Experiment: "qwen*"}, true}, + {"experiment suffix glob", listFilters{Experiment: "*-eval"}, true}, + {"experiment case-insensitive", listFilters{Experiment: "QWEN*"}, true}, + {"experiment no match", listFilters{Experiment: "llama*"}, false}, + {"accelerator type substring", listFilters{AcceleratorType: "h100"}, true}, + {"accelerator type no match", listFilters{AcceleratorType: "a10"}, false}, + {"num accelerators match", listFilters{NumAccelerators: new(8)}, true}, + {"num accelerators no match", listFilters{NumAccelerators: new(4)}, false}, + } + for _, c := range cases { + t.Run(c.name, func(t *testing.T) { + assert.Equal(t, c.want, c.f.matches(&run)) + }) + } +} diff --git a/experimental/air/cmd/list_format.go b/experimental/air/cmd/list_format.go new file mode 100644 index 00000000000..bfa74728cc3 --- /dev/null +++ b/experimental/air/cmd/list_format.go @@ -0,0 +1,47 @@ +package aircmd + +import ( + "strconv" + "time" + + "github.com/databricks/databricks-sdk-go/service/jobs" +) + +// buildListRow extracts the columns shown for one run. Optional cells fall back +// to "-"; MLflowURL starts as "-" and setMLflowLinks fills it in for text output. +func buildListRow(run *jobs.Run) listRow { + experiment := "-" + if e := jobExperiment(run); e != "" { + experiment = e + } + + var startedAt *string + duration := "-" + if start, end := jobTiming(run); start > 0 { + s := isoFormat(time.UnixMilli(start)) + startedAt = &s + if end == 0 { + // Still running: measure against the current time. + end = time.Now().UnixMilli() + } + duration = formatDuration(roundMillisToSeconds(end - start)) + } + + accel := "-" + if a := acceleratorLabel(jobCompute(run)); a != "" { + accel = a + } + + return listRow{ + RunID: strconv.FormatInt(run.RunId, 10), + RunName: run.RunName, + User: run.CreatorUserName, + Status: runStatus(run.State), + StartedAt: startedAt, + IsSweep: isSweep(run), + Experiment: experiment, + Duration: duration, + MLflowURL: "-", + Accelerators: accel, + } +} diff --git a/experimental/air/cmd/list_index.go b/experimental/air/cmd/list_index.go new file mode 100644 index 00000000000..243affd664f --- /dev/null +++ b/experimental/air/cmd/list_index.go @@ -0,0 +1,146 @@ +package aircmd + +import ( + "cmp" + "context" + "slices" + + "github.com/databricks/cli/libs/cache" + "github.com/databricks/databricks-sdk-go" +) + +// indexStrategy serves the caller's own runs from the AiTrainingService index: +// it fetches every run id up front (cheap id+timestamp pairs), orders them +// newest-first, keeps the newest `limit`, then hydrates them into full rows in +// want-sized batches via Jobs runs/get. Terminal rows are cached so repeat calls +// skip the network. Unlike the Jobs scan it can't lazy-page (it must sort the +// whole id set first), but it still yields in batches so the table paints early. +type indexStrategy struct { + ctx context.Context + w *databricks.WorkspaceClient + activeOnly bool + filters listFilters + limit int + cache *cache.Cache + + ids []int64 // newest-first run ids to hydrate, resolved on first next() + pos int + loaded bool +} + +func newIndexStrategy(ctx context.Context, w *databricks.WorkspaceClient, q listQuery, limit int) *indexStrategy { + return &indexStrategy{ + ctx: ctx, + w: w, + activeOnly: q.activeOnly, + filters: q.filters, + limit: limit, + cache: newListCache(ctx), + } +} + +// load fetches and orders the index once. It returns an error only when the +// index endpoint itself fails, letting the caller fall back to the Jobs scan. +func (s *indexStrategy) load() error { + refs, err := listAiTrainingWorkflows(s.ctx, s.w, s.activeOnly) + if err != nil { + return err + } + slices.SortFunc(refs, func(a, b workflowRef) int { return cmp.Compare(b.submitTimeMs, a.submitTimeMs) }) + // Keep only the newest `limit` ids so hydration is bounded — but skip that when + // a task filter is active, since it drops matches post-hydration and we'd + // otherwise return fewer than `limit`. The caller stops pulling at `limit`. + if s.limit > 0 && len(refs) > s.limit && !s.filters.hasTaskFilter() { + refs = refs[:s.limit] + } + s.ids = make([]int64, len(refs)) + for i, r := range refs { + s.ids[i] = r.jobRunID + } + s.loaded = true + return nil +} + +func (s *indexStrategy) next(want int) ([]listedRun, error) { + if !s.loaded { + if err := s.load(); err != nil { + return nil, err + } + } + + var entries []listedRun + for len(entries) < want && s.pos < len(s.ids) { + end := min(s.pos+want-len(entries), len(s.ids)) + batch := s.ids[s.pos:end] + s.pos = end + + rows, err := s.hydrate(batch) + if err != nil { + return nil, err + } + entries = append(entries, rows...) + } + return entries, nil +} + +func (s *indexStrategy) done() bool { + return s.loaded && s.pos >= len(s.ids) +} + +// truncated is always false: the index path is bounded by limit, not a scan cap. +func (s *indexStrategy) truncated() bool { return false } + +// hydrate turns a batch of run ids into rows, serving cached terminal rows +// without a network call and fetching the rest via runs/get. Freshly hydrated +// terminal runs are cached. Results keep the input (newest-first) order, then +// the batch is re-sorted by start time since concurrent hydration reorders it. +func (s *indexStrategy) hydrate(ids []int64) ([]listedRun, error) { + host := s.w.Config.Host + + rows := make([]listedRun, 0, len(ids)) + var toFetch []int64 + // The cache key excludes the filter, so a cached row is still run through the + // active filter; a non-matching hit is dropped, not re-fetched. + for _, id := range ids { + if row, fields, ok := cachedRow(s.ctx, s.cache, host, id); ok { + if s.filters.matchesFields(fields) { + rows = append(rows, listedRun{row: row, taskRunID: id}) + } + continue + } + toFetch = append(toFetch, id) + } + + runs, err := hydrateJobRuns(s.ctx, s.w, toFetch) + if err != nil { + return nil, err + } + for _, run := range runs { + fields := filterFieldsFromRun(run) + if !s.filters.matchesFields(fields) { + continue + } + row := buildListRow(run) + rows = append(rows, listedRun{row: row, taskRunID: taskRunID(run)}) + if isTerminal(run) { + start, _ := jobTiming(run) + putRow(s.ctx, s.cache, host, run.RunId, start, row, fields) + } + } + + // Concurrent hydration reorders runs, so re-sort the batch newest-first. The + // ISO start timestamp sorts lexicographically; a missing time ("") sorts last. + slices.SortStableFunc(rows, func(a, b listedRun) int { + return cmp.Compare(rowStartKey(b.row), rowStartKey(a.row)) + }) + return rows, nil +} + +// rowStartKey returns a row's ISO start timestamp for ordering, or "" when the +// run hasn't started (which sorts last under descending comparison). +func rowStartKey(r listRow) string { + if r.StartedAt == nil { + return "" + } + return *r.StartedAt +} diff --git a/experimental/air/cmd/list_index_test.go b/experimental/air/cmd/list_index_test.go new file mode 100644 index 00000000000..f9f6151aae4 --- /dev/null +++ b/experimental/air/cmd/list_index_test.go @@ -0,0 +1,213 @@ +package aircmd + +import ( + "encoding/json" + "fmt" + "net/http" + "net/http/httptest" + "strconv" + "testing" + + "github.com/databricks/databricks-sdk-go/service/jobs" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// indexRun is a terminal AIR run with a start time, for index-path hydration. +func indexRun(id, startMillis int64) jobs.Run { + r := airRun(id, "me@example.com", "GPU_1xH100", 1, "/Users/me@example.com/exp") + r.State = &jobs.RunState{LifeCycleState: jobs.RunLifeCycleStateTerminated, ResultState: jobs.RunResultStateSuccess} + r.Tasks[0].StartTime = startMillis + r.Tasks[0].EndTime = startMillis + 1000 + return r +} + +// indexAndGetServer serves the AiTrainingService index (a single page of the +// given refs) and runs/get for each id, recording hit counts per endpoint. A +// runID in forbidden returns 403; in missing returns 404. +type indexHits struct{ index, get int } + +func indexAndGetServer(t *testing.T, refs []workflowRef, runs map[int64]jobs.Run, forbidden, missing map[int64]bool) (*httptest.Server, *indexHits) { + t.Helper() + hits := &indexHits{} + wfs := make([]map[string]any, len(refs)) + for i, r := range refs { + wfs[i] = map[string]any{"job_run_id": strconv.FormatInt(r.jobRunID, 10), "submit_time": map[string]any{"seconds": r.submitTimeMs / 1000}} + } + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case aiTrainingWorkflowsPath: + hits.index++ + _ = json.NewEncoder(w).Encode(map[string]any{"training_workflows": wfs}) + case "/api/2.2/jobs/runs/get": + hits.get++ + id, _ := strconv.ParseInt(r.URL.Query().Get("run_id"), 10, 64) + if forbidden[id] { + w.WriteHeader(http.StatusForbidden) + _, _ = w.Write([]byte(`{"message":"forbidden"}`)) + return + } + if missing[id] { + w.WriteHeader(http.StatusNotFound) + _, _ = w.Write([]byte(`{"message":"not found"}`)) + return + } + run := runs[id] + _ = json.NewEncoder(w).Encode(run) + default: + _, _ = w.Write([]byte(`{}`)) + } + })) + t.Cleanup(srv.Close) + return srv, hits +} + +func TestIndexStrategyOrdersAndLimits(t *testing.T) { + // Three runs, out of submit-time order; newest two should win, newest-first. + refs := []workflowRef{ + {jobRunID: 1, submitTimeMs: 1000_000}, + {jobRunID: 2, submitTimeMs: 3000_000}, + {jobRunID: 3, submitTimeMs: 2000_000}, + } + runs := map[int64]jobs.Run{ + 1: indexRun(1, 1000_000), + 2: indexRun(2, 3000_000), + 3: indexRun(3, 2000_000), + } + srv, _ := indexAndGetServer(t, refs, runs, nil, nil) + t.Setenv("DATABRICKS_CACHE_ENABLED", "false") + + f := newRunFetcher(t.Context(), newTestWorkspaceClient(t, srv.URL), listQuery{ + userFilter: "me@example.com", currentUser: "me@example.com", limit: 2, + }) + rows, err := f.next(10) + require.NoError(t, err) + require.Len(t, rows, 2) + assert.Equal(t, "2", rows[0].RunID) // submit 3000 + assert.Equal(t, "3", rows[1].RunID) // submit 2000 + assert.True(t, f.exhausted) +} + +func TestIndexStrategyOverFetchesWithTaskFilter(t *testing.T) { + // With a task filter and limit 1, the newest run doesn't match; the strategy + // must keep hydrating past `limit` to find the match rather than truncating. + refs := []workflowRef{ + {jobRunID: 1, submitTimeMs: 3000_000}, + {jobRunID: 2, submitTimeMs: 2000_000}, + } + run1 := indexRun(1, 3000_000) + run1.Tasks[0].AiRuntimeTask.Experiment = "/Users/me@example.com/llama" + run2 := indexRun(2, 2000_000) + run2.Tasks[0].AiRuntimeTask.Experiment = "/Users/me@example.com/qwen" + srv, _ := indexAndGetServer(t, refs, map[int64]jobs.Run{1: run1, 2: run2}, nil, nil) + t.Setenv("DATABRICKS_CACHE_ENABLED", "false") + + f := newRunFetcher(t.Context(), newTestWorkspaceClient(t, srv.URL), listQuery{ + userFilter: "me@example.com", currentUser: "me@example.com", limit: 1, + filters: listFilters{Experiment: "qwen"}, + }) + rows, err := f.next(1) + require.NoError(t, err) + require.Len(t, rows, 1) + assert.Equal(t, "2", rows[0].RunID) // found despite being the older, second id +} + +func TestIndexStrategyDropsForbiddenAndMissing(t *testing.T) { + refs := []workflowRef{ + {jobRunID: 1, submitTimeMs: 3000_000}, + {jobRunID: 2, submitTimeMs: 2000_000}, + {jobRunID: 3, submitTimeMs: 1000_000}, + } + runs := map[int64]jobs.Run{1: indexRun(1, 3000_000), 3: indexRun(3, 1000_000)} + srv, _ := indexAndGetServer(t, refs, runs, map[int64]bool{2: true}, nil) + t.Setenv("DATABRICKS_CACHE_ENABLED", "false") + + f := newRunFetcher(t.Context(), newTestWorkspaceClient(t, srv.URL), listQuery{ + userFilter: "me@example.com", currentUser: "me@example.com", limit: 10, + }) + rows, err := f.next(10) + require.NoError(t, err) + require.Len(t, rows, 2) // run 2 (403) dropped + assert.Equal(t, "1", rows[0].RunID) + assert.Equal(t, "3", rows[1].RunID) +} + +func TestIndexStrategyPropagatesServerError(t *testing.T) { + refs := []workflowRef{{jobRunID: 1, submitTimeMs: 1000_000}} + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case aiTrainingWorkflowsPath: + wfs := []map[string]any{{"job_run_id": "1", "submit_time": map[string]any{"seconds": int64(1000)}}} + _ = json.NewEncoder(w).Encode(map[string]any{"training_workflows": wfs}) + case "/api/2.2/jobs/runs/get": + w.WriteHeader(http.StatusInternalServerError) + _, _ = w.Write([]byte(`{"message":"boom"}`)) + default: + _, _ = w.Write([]byte(`{}`)) + } + })) + t.Cleanup(srv.Close) + _ = refs + t.Setenv("DATABRICKS_CACHE_ENABLED", "false") + + f := newRunFetcher(t.Context(), newTestWorkspaceClient(t, srv.URL), listQuery{ + userFilter: "me@example.com", currentUser: "me@example.com", limit: 10, + }) + _, err := f.next(10) + require.Error(t, err) // 500 is systemic, not an ACL drop +} + +func TestNewListStrategyGate(t *testing.T) { + // --all-users and other-user filters must NOT touch the index endpoint. + cases := []struct { + name string + q listQuery + wantIndex bool + }{ + {"active default → scan", listQuery{activeOnly: true, userFilter: "me@example.com", currentUser: "me@example.com"}, false}, + {"all-status self → index", listQuery{userFilter: "me@example.com", currentUser: "me@example.com", limit: 5}, true}, + {"all-status all-users → scan", listQuery{allUsers: true}, false}, + {"all-status other user → scan", listQuery{userFilter: "other@example.com", currentUser: "me@example.com"}, false}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + var indexHit bool + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == aiTrainingWorkflowsPath { + indexHit = true + } + _, _ = w.Write([]byte(`{}`)) + })) + t.Cleanup(srv.Close) + t.Setenv("DATABRICKS_CACHE_ENABLED", "false") + + newListStrategy(t.Context(), newTestWorkspaceClient(t, srv.URL), tc.q) + assert.Equal(t, tc.wantIndex, indexHit) + }) + } +} + +func TestNewListStrategyFallsBackWhenIndexFails(t *testing.T) { + // Index 500 must silently fall back to the Jobs scan, not fail the command. + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == aiTrainingWorkflowsPath { + w.WriteHeader(http.StatusInternalServerError) + _, _ = w.Write([]byte(`{"message":"boom"}`)) + return + } + if r.URL.Path == "/api/2.2/jobs/runs/list" { + _, _ = fmt.Fprint(w, runsListBody(t, "", airBaseRun(1, "me@example.com", "GPU_1xH100", 1, "exp"))) + return + } + _, _ = w.Write([]byte(`{}`)) + })) + t.Cleanup(srv.Close) + t.Setenv("DATABRICKS_CACHE_ENABLED", "false") + + f := newRunFetcher(t.Context(), newTestWorkspaceClient(t, srv.URL), listQuery{ + userFilter: "me@example.com", currentUser: "me@example.com", limit: 10, + }) + rows, err := f.next(10) + require.NoError(t, err) + require.Len(t, rows, 1) // served by the Jobs scan fallback +} diff --git a/experimental/air/cmd/list_test.go b/experimental/air/cmd/list_test.go new file mode 100644 index 00000000000..f70330240e5 --- /dev/null +++ b/experimental/air/cmd/list_test.go @@ -0,0 +1,237 @@ +package aircmd + +import ( + "encoding/json" + "net/http" + "net/http/httptest" + "strconv" + "testing" + + "github.com/databricks/cli/libs/cmdctx" + "github.com/databricks/cli/libs/cmdio" + "github.com/databricks/databricks-sdk-go/experimental/mocks" + "github.com/databricks/databricks-sdk-go/service/jobs" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// airBaseRun builds a single-task AIR run (ai_runtime_task), as runs/list +// returns it with expand_tasks. +func airBaseRun(id int64, user, accelType string, count int, experiment string) jobs.BaseRun { + return jobs.BaseRun{ + RunId: id, + RunName: "run-" + strconv.FormatInt(id, 10), + CreatorUserName: user, + State: &jobs.RunState{LifeCycleState: jobs.RunLifeCycleStateRunning}, + Tasks: []jobs.RunTask{{AiRuntimeTask: &jobs.AiRuntimeTask{ + Experiment: experiment, + Deployments: []jobs.DeploymentSpec{{ + Compute: jobs.ComputeSpec{AcceleratorType: jobs.ComputeSpecAcceleratorType(accelType), AcceleratorCount: count}, + }}, + }}}, + } +} + +// airRun is the runs/get equivalent of airBaseRun, for the helpers that operate +// on a *jobs.Run. +func airRun(id int64, user, accelType string, count int, experiment string) jobs.Run { + return *baseRunToRun(airBaseRun(id, user, accelType, count, experiment)) +} + +// runsListBody marshals one runs/list response page. +func runsListBody(t *testing.T, nextToken string, runs ...jobs.BaseRun) string { + t.Helper() + b, err := json.Marshal(jobs.ListRunsResponse{Runs: runs, NextPageToken: nextToken, HasMore: nextToken != ""}) + require.NoError(t, err) + return string(b) +} + +// runsServer serves one runs/list response body per call, repeating the last +// once exhausted, and a stub for any other request (the SDK config probe). +func runsServer(t *testing.T, bodies ...string) *httptest.Server { + t.Helper() + call := 0 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == "/api/2.2/jobs/runs/list" { + body := bodies[min(call, len(bodies)-1)] + call++ + _, _ = w.Write([]byte(body)) + return + } + _, _ = w.Write([]byte(`{}`)) + })) + t.Cleanup(srv.Close) + return srv +} + +func TestListAirRunsFiltersUserAndType(t *testing.T) { + runs := []jobs.BaseRun{ + airBaseRun(1, "me@example.com", "GPU_8xH100", 8, "/Users/me@example.com/exp-a"), + {RunId: 2, CreatorUserName: "me@example.com", Tasks: []jobs.RunTask{{}}}, // not an AIR run + airBaseRun(3, "other@example.com", "GPU_1xA10", 1, "/Users/other/exp-b"), // wrong user + airBaseRun(5, "me@example.com", "GPU_1xH100", 1, "/Users/me@example.com/exp-c"), + } + srv := runsServer(t, runsListBody(t, "", runs...)) + + rows, err := newRunFetcher(t.Context(), newTestWorkspaceClient(t, srv.URL), listQuery{ + activeOnly: true, + userFilter: "me@example.com", + }).next(10) + require.NoError(t, err) + require.Len(t, rows, 2) + assert.Equal(t, "1", rows[0].RunID) + assert.Equal(t, "5", rows[1].RunID) +} + +func TestListAirRunsExperimentFilter(t *testing.T) { + runs := []jobs.BaseRun{ + airBaseRun(1, "me@example.com", "GPU_1xH100", 1, "/Users/me@example.com/qwen-train"), + airBaseRun(2, "me@example.com", "GPU_1xH100", 1, "/Users/me@example.com/llama-train"), + } + srv := runsServer(t, runsListBody(t, "", runs...)) + + rows, err := newRunFetcher(t.Context(), newTestWorkspaceClient(t, srv.URL), listQuery{ + activeOnly: true, + filters: listFilters{Experiment: "qwen*"}, + }).next(10) + require.NoError(t, err) + require.Len(t, rows, 1) + assert.Equal(t, "1", rows[0].RunID) +} + +func TestListAirRunsLimitTruncates(t *testing.T) { + runs := []jobs.BaseRun{ + airBaseRun(1, "me@example.com", "GPU_1xH100", 1, "exp-a"), + airBaseRun(2, "me@example.com", "GPU_1xH100", 1, "exp-b"), + airBaseRun(3, "me@example.com", "GPU_1xH100", 1, "exp-c"), + } + srv := runsServer(t, runsListBody(t, "", runs...)) + + rows, err := newRunFetcher(t.Context(), newTestWorkspaceClient(t, srv.URL), listQuery{activeOnly: true}).next(2) + require.NoError(t, err) + require.Len(t, rows, 2) + assert.Equal(t, "1", rows[0].RunID) + assert.Equal(t, "2", rows[1].RunID) +} + +func TestListAirRunsPaginates(t *testing.T) { + page1 := runsListBody(t, "tok", airBaseRun(1, "me@example.com", "GPU_1xH100", 1, "exp-a")) + page2 := runsListBody(t, "", airBaseRun(2, "me@example.com", "GPU_1xH100", 1, "exp-b")) + srv := runsServer(t, page1, page2) + + rows, err := newRunFetcher(t.Context(), newTestWorkspaceClient(t, srv.URL), listQuery{activeOnly: true}).next(10) + require.NoError(t, err) + require.Len(t, rows, 2) + assert.Equal(t, "1", rows[0].RunID) + assert.Equal(t, "2", rows[1].RunID) +} + +// TestRunFetcherResumesAcrossCalls covers the lazy paging the interactive table +// relies on: a next() that stops mid-page must resume on the following call, then +// report exhaustion. +func TestRunFetcherResumesAcrossCalls(t *testing.T) { + runs := []jobs.BaseRun{ + airBaseRun(1, "me@example.com", "GPU_1xH100", 1, "exp-a"), + airBaseRun(2, "me@example.com", "GPU_1xH100", 1, "exp-b"), + airBaseRun(3, "me@example.com", "GPU_1xH100", 1, "exp-c"), + } + srv := runsServer(t, runsListBody(t, "", runs...)) + f := newRunFetcher(t.Context(), newTestWorkspaceClient(t, srv.URL), listQuery{activeOnly: true}) + + first, err := f.next(2) + require.NoError(t, err) + require.Len(t, first, 2) + assert.Equal(t, "1", first[0].RunID) + assert.Equal(t, "2", first[1].RunID) + + second, err := f.next(2) + require.NoError(t, err) + require.Len(t, second, 1) // only the leftover remains + assert.Equal(t, "3", second[0].RunID) + assert.True(t, f.exhausted) + + third, err := f.next(2) + require.NoError(t, err) + assert.Empty(t, third) +} + +func TestBuildListRowFromRun(t *testing.T) { + run := airRun(842552489592352, "me@example.com", "GPU_1xA10", 1, "my-first-air-run") + run.StartTime = 1700000000000 + run.EndTime = 1700000012000 + run.State = &jobs.RunState{LifeCycleState: jobs.RunLifeCycleStateTerminated, ResultState: jobs.RunResultStateSuccess} + + assert.True(t, isAirRun(&run)) + assert.Equal(t, "my-first-air-run", jobExperiment(&run)) + gpu, count := jobCompute(&run) + assert.Equal(t, "GPU_1xA10", gpu) + assert.Equal(t, 1, count) + + row := buildListRow(&run) + assert.Equal(t, "842552489592352", row.RunID) + assert.Equal(t, "SUCCESS", row.Status) + assert.Equal(t, "my-first-air-run", row.Experiment) + assert.Equal(t, "1x A10", row.Accelerators) + assert.Equal(t, "12s", row.Duration) +} + +func TestBuildListRow(t *testing.T) { + run := airRun(123, "me@example.com", "GPU_8xH100", 8, "/Users/me@example.com/exp") + run.StartTime = 1700000000000 + run.EndTime = 1700000012000 + run.State = &jobs.RunState{ResultState: jobs.RunResultStateSuccess} + + row := buildListRow(&run) + assert.Equal(t, "123", row.RunID) + assert.Equal(t, "me@example.com", row.User) + assert.Equal(t, "SUCCESS", row.Status) + assert.Equal(t, "exp", row.Experiment) + assert.Equal(t, "12s", row.Duration) + assert.Equal(t, "8x H100", row.Accelerators) + assert.Equal(t, "-", row.MLflowURL) + assert.False(t, row.IsSweep) + require.NotNil(t, row.StartedAt) +} + +func TestBuildListRowDashFallbacks(t *testing.T) { + // A run with no task, compute, or start time falls back to dashes and UNKNOWN. + row := buildListRow(&jobs.Run{RunId: 7}) + assert.Equal(t, "-", row.Experiment) + assert.Equal(t, "-", row.Duration) + assert.Equal(t, "-", row.Accelerators) + assert.Equal(t, "-", row.MLflowURL) + assert.Equal(t, "UNKNOWN", row.Status) + assert.Nil(t, row.StartedAt) +} + +func TestBuildListRowSweep(t *testing.T) { + run := jobs.Run{RunId: 9, Tasks: []jobs.RunTask{{ + ForEachTask: &jobs.RunForEachTask{Task: jobs.Task{AiRuntimeTask: &jobs.AiRuntimeTask{Experiment: "sweep"}}}, + }}} + assert.True(t, buildListRow(&run).IsSweep) + assert.Equal(t, "sweep", buildListRow(&run).Experiment) +} + +func TestListInvalidLimit(t *testing.T) { + m := mocks.NewMockWorkspaceClient(t) + ctx := cmdctx.SetWorkspaceClient(cmdio.MockDiscard(t.Context()), m.WorkspaceClient) + cmd := newListCommand() + cmd.SetContext(ctx) + require.NoError(t, cmd.Flags().Set("limit", "0")) + + err := cmd.RunE(cmd, nil) + require.Error(t, err) + assert.Contains(t, err.Error(), "invalid --limit") +} + +func TestListInvalidFilter(t *testing.T) { + m := mocks.NewMockWorkspaceClient(t) + ctx := cmdctx.SetWorkspaceClient(cmdio.MockDiscard(t.Context()), m.WorkspaceClient) + cmd := newListCommand() + cmd.SetContext(ctx) + require.NoError(t, cmd.Flags().Set("filter", "bogus=1")) + + err := cmd.RunE(cmd, nil) + require.Error(t, err) + assert.Contains(t, err.Error(), "unsupported --filter key") +} diff --git a/experimental/air/cmd/list_tui.go b/experimental/air/cmd/list_tui.go new file mode 100644 index 00000000000..85fe70774b4 --- /dev/null +++ b/experimental/air/cmd/list_tui.go @@ -0,0 +1,248 @@ +package aircmd + +import ( + "fmt" + "io" + "strings" + + tea "github.com/charmbracelet/bubbletea" + "github.com/charmbracelet/lipgloss" + "github.com/databricks/cli/libs/cmdio" + "github.com/pkg/browser" + "github.com/spf13/cobra" +) + +// renderListText renders the table for text output: an inline navigable table in +// a terminal (paging in older runs on demand), otherwise printed once. JSON is +// handled by the caller. +func renderListText(cmd *cobra.Command, f *runFetcher, limit int) error { + ctx := cmd.Context() + out := cmd.OutOrStdout() + + r, color := cmdio.NewRenderer(ctx, out) + + // Navigate only with a full color TTY and no explicit --limit (which means + // "just print these N"). Everything else — piped, NO_COLOR, --limit — prints + // once. + interactive := color && + cmdio.IsPagerSupported(ctx) && + !cmd.Flags().Changed("limit") + + if interactive { + first, err := f.next(listPageRows) + if err != nil { + return err + } + if len(first) == 0 { + _, err := io.WriteString(out, "No runs found.\n") + return err + } + _, err = tea.NewProgram( + newListModel(r, f, first, color), + tea.WithContext(ctx), + tea.WithInput(cmd.InOrStdin()), + tea.WithOutput(out), + ).Run() + return err + } + + rows, err := f.next(limit) + if err != nil { + return err + } + warnIfTruncated(ctx, f) + _, err = io.WriteString(out, staticListTable(r, rows, color)) + return err +} + +// staticListTable renders the whole table once, with no selection — used when +// piped or non-interactive. +func staticListTable(r *lipgloss.Renderer, rows []listRow, links bool) string { + if len(rows) == 0 { + return "No runs found.\n" + } + styles := newListStyles(r) + cols := computeListCols(rows) + var b strings.Builder + b.WriteString(styles.renderHeader(cols)) + b.WriteByte('\n') + for _, row := range rows { + b.WriteString(styles.renderRow(cols, row, false, links)) + b.WriteByte('\n') + } + return b.String() +} + +// listModel is the inline, navigable runs table. It lazily pages older runs from +// the fetcher as the cursor nears the end of the loaded rows. fetcher is nil for +// a fixed, non-paging table (e.g. in tests). +type listModel struct { + rows []listRow + styles listStyles + cols listCols + links bool + fetcher *runFetcher + loading bool + loadErr error + + cursor int + offset int // index of the first visible row + height int // terminal height, for windowing +} + +func newListModel(r *lipgloss.Renderer, f *runFetcher, rows []listRow, links bool) listModel { + return listModel{ + rows: rows, + styles: newListStyles(r), + cols: computeListCols(rows), + links: links, + fetcher: f, + } +} + +func (m listModel) Init() tea.Cmd { return nil } + +// moreRowsMsg carries a lazily fetched batch of rows, or the error that ended paging. +type moreRowsMsg struct { + rows []listRow + err error +} + +// fetchCmd pulls the next batch of rows in the background; sets loading so only +// one runs at a time, and returns the updated model with the fetch command. +func (m listModel) fetchCmd() (listModel, tea.Cmd) { + m.loading = true + f := m.fetcher + return m, func() tea.Msg { + rows, err := f.next(listPageRows) + return moreRowsMsg{rows: rows, err: err} + } +} + +// maybeFetch starts a fetch when the cursor nears the end of the loaded rows and +// more runs may still exist. +func (m listModel) maybeFetch() (listModel, tea.Cmd) { + if m.fetcher == nil || m.loading || m.loadErr != nil || m.fetcher.exhausted { + return m, nil + } + if m.cursor < len(m.rows)-m.visibleCount() { + return m, nil + } + return m.fetchCmd() +} + +// listPageRows is the most rows shown per page. +const listPageRows = 20 + +// visibleCount is how many rows a page shows: at most listPageRows, and never +// more than fits below the header and hint. +func (m listModel) visibleCount() int { + n := min(listPageRows, len(m.rows)) + if m.height > 0 { + n = min(n, m.height-3) + } + return max(1, n) +} + +func (m listModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) { + switch msg := msg.(type) { + case tea.WindowSizeMsg: + m.height = msg.Height + m.offset = m.clampedOffset() + return m.maybeFetch() + + case moreRowsMsg: + m.loading = false + if msg.err != nil { + m.loadErr = msg.err + return m, nil + } + m.rows = append(m.rows, msg.rows...) + m.cols = computeListCols(m.rows) + m.offset = m.clampedOffset() + // A page with no matches but more to scan: keep paging so the cursor isn't + // stuck at the end of the loaded rows. + if len(msg.rows) == 0 && !m.fetcher.exhausted { + return m.fetchCmd() + } + return m, nil + + case tea.KeyMsg: + switch msg.String() { + case "q", "ctrl+c", "esc": + return m, tea.Quit + case "up", "k": + if m.cursor > 0 { + m.cursor-- + } + case "down", "j": + if m.cursor < len(m.rows)-1 { + m.cursor++ + } + case "right": + m.cursor = min(m.cursor+m.visibleCount(), len(m.rows)-1) + case "left": + m.cursor = max(m.cursor-m.visibleCount(), 0) + case "home", "g": + m.cursor = 0 + case "end", "G": + m.cursor = len(m.rows) - 1 + case "enter": + // Open the selected run's MLflow page in the browser. + if len(m.rows) > 0 { + if url := m.rows[m.cursor].MLflowURL; url != "" && url != "-" { + return m, openURL(url) + } + } + } + m.offset = m.clampedOffset() + return m.maybeFetch() + } + return m, nil +} + +// clampedOffset returns the scroll offset that keeps the cursor visible. +func (m listModel) clampedOffset() int { + visible := m.visibleCount() + offset := min(m.offset, m.cursor) + if m.cursor >= offset+visible { + offset = m.cursor - visible + 1 + } + return max(offset, 0) +} + +func (m listModel) View() string { + if len(m.rows) == 0 { + return m.styles.r.NewStyle().Foreground(colN9).Render("No runs found.") + "\n" + } + + visible := m.visibleCount() + lines := []string{m.styles.renderHeader(m.cols)} + for i := m.offset; i < m.offset+visible && i < len(m.rows); i++ { + lines = append(lines, m.styles.renderRow(m.cols, m.rows[i], i == m.cursor, m.links)) + } + lines = append(lines, m.renderHint()) + return strings.Join(lines, "\n") + "\n" +} + +// renderHint is the faint one-line key legend, with the cursor position and the +// paging state (loading / load failed). +func (m listModel) renderHint() string { + faint := m.styles.r.NewStyle().Foreground(colN7) + hint := fmt.Sprintf("↑/↓ navigate · ←/→ page · ↵ mlflow · q quit · row %d/%d", m.cursor+1, len(m.rows)) + switch { + case m.loadErr != nil: + hint += " (load failed)" + case m.loading: + hint += " (loading…)" + } + return faint.Render(hint) +} + +// openURL opens a URL in the user's default browser, best-effort. +func openURL(url string) tea.Cmd { + return func() tea.Msg { + _ = browser.OpenURL(url) + return nil + } +} diff --git a/experimental/air/cmd/list_tui_render.go b/experimental/air/cmd/list_tui_render.go new file mode 100644 index 00000000000..6a81deceb5c --- /dev/null +++ b/experimental/air/cmd/list_tui_render.go @@ -0,0 +1,239 @@ +package aircmd + +import ( + "strings" + + "github.com/charmbracelet/lipgloss" + "github.com/muesli/termenv" +) + +// DuBois dark-first palette. Neutrals (n7..n12) carry structure; the state +// colors are reserved for the Status column. +const ( + colN7 = lipgloss.Color("#69696B") // gutter / dim + colN9 = lipgloss.Color("#ABABAE") // secondary values + colN11 = lipgloss.Color("#E0E0E3") // body + colN12 = lipgloss.Color("#F1F1F4") // selected text + colOverlay = lipgloss.Color("#1F1F23") // selected-row fill + colRunID = lipgloss.Color("#B7A8E8") // run id + colGreen = lipgloss.Color("#4CD964") // success + colAmber = lipgloss.Color("#E8B84A") // running / pending + colRed = lipgloss.Color("#EB6B6B") // failed + colBlue = lipgloss.Color("#6CA8F0") // MLflow link +) + +const mlflowColWidth = 18 + +// listStyles renders the runs table. The renderer carries the color profile, so +// styles render plain under --no-color / non-tty. +type listStyles struct { + r *lipgloss.Renderer +} + +func newListStyles(r *lipgloss.Renderer) listStyles { + return listStyles{r: r} +} + +// listCols holds the computed width of each variable-width column. MLflow is +// fixed (a short link) and the gutter is one cell. +type listCols struct { + runID, experiment, status, started, duration, user, accel int +} + +// columnCap bounds the widest free-text columns so one long value can't dominate +// the row. +const columnCap = 36 + +func computeListCols(rows []listRow) listCols { + c := listCols{ + runID: len("Run ID"), + experiment: len("Experiment"), + status: len("Status"), + started: len("Started"), + duration: len("Duration"), + user: len("User"), + accel: len("Accelerators"), + } + for _, r := range rows { + c.runID = max(c.runID, lipgloss.Width(r.RunID)) + c.experiment = min(columnCap, max(c.experiment, lipgloss.Width(r.Experiment))) + c.status = max(c.status, lipgloss.Width("● "+r.Status)) + c.started = max(c.started, lipgloss.Width(startedDisplay(r))) + c.duration = max(c.duration, lipgloss.Width(r.Duration)) + c.user = min(columnCap, max(c.user, lipgloss.Width(r.User))) + c.accel = max(c.accel, lipgloss.Width(r.Accelerators)) + } + return c +} + +// renderHeader renders the muted column-title row. +func (s listStyles) renderHeader(cols listCols) string { + h := func(text string, width int, right bool) string { + return s.r.NewStyle().Foreground(colN9).Render(pad(text, width, right)) + } + cells := []string{ + pad(" ", 1, false), + h("Run ID", cols.runID, false), + h("Experiment", cols.experiment, false), + h("Status", cols.status, false), + h("Started", cols.started, false), + h("Duration", cols.duration, true), + h("MLflow", mlflowColWidth, false), + h("User", cols.user, false), + h("Accelerators", cols.accel, false), + } + // Trim trailing pad on the final column so rows carry no trailing whitespace. + return strings.TrimRight(strings.Join(cells, " "), " ") +} + +// renderRow renders one run. The selected row uses a subtle overlay fill + n12 +// text (not full inversion, not per-state color). +func (s listStyles) renderRow(cols listCols, r listRow, selected, links bool) string { + base := s.r.NewStyle() + if selected { + base = base.Background(colOverlay) + } + fg := func(c lipgloss.Color) lipgloss.Color { + if selected { + return colN12 + } + return c + } + + gutter := " " + if selected { + gutter = "▸" + } + + cells := []string{ + s.cell(base, gutter, 1, fg(colN7), false, false, ""), + s.cell(base, r.RunID, cols.runID, fg(colRunID), false, false, ""), + s.cell(base, r.Experiment, cols.experiment, fg(colN11), false, false, ""), + s.cell(base, "● "+r.Status, cols.status, fg(statusColor(r.Status)), false, false, ""), + s.cell(base, startedDisplay(r), cols.started, fg(colN9), false, false, ""), + s.cell(base, r.Duration, cols.duration, fg(colN9), true, false, ""), + s.mlflowCell(base, r, selected, links), + s.cell(base, r.User, cols.user, fg(colN9), false, false, ""), + s.cell(base, r.Accelerators, cols.accel, fg(colN9), false, false, ""), + } + // Trim the final column's trailing pad so rows carry no trailing whitespace. + return strings.TrimRight(strings.Join(cells, base.Render(" ")), " ") +} + +// cell renders one padded, colored cell. The text is truncated to width, then +// padded with (background-only) spaces so columns align even when the text is +// styled or hyperlinked. +func (s listStyles) cell(base lipgloss.Style, text string, width int, fg lipgloss.Color, right, underline bool, link string) string { + text = truncate(text, width) + style := base.Foreground(fg) + if underline { + style = style.Underline(true) + } + rendered := style.Render(text) + if link != "" { + rendered = termenv.Hyperlink(link, rendered) + } + gap := max(width-lipgloss.Width(text), 0) + padStr := base.Render(strings.Repeat(" ", gap)) + if right { + return padStr + rendered + } + return rendered + padStr +} + +// mlflowCell renders the fixed-width MLflow column: a short, blue, underlined +// OSC 8 hyperlink (when links are enabled), or "-" when the run has no link. +func (s listStyles) mlflowCell(base lipgloss.Style, r listRow, selected, links bool) string { + if r.MLflowURL == "" || r.MLflowURL == "-" { + fg := colN9 + if selected { + fg = colN12 + } + return s.cell(base, "-", mlflowColWidth, fg, false, false, "") + } + fg := colBlue + if selected { + fg = colN12 + } + link := "" + if links { + link = r.MLflowURL + } + return s.cell(base, mlflowDisplay(r.MLflowURL), mlflowColWidth, fg, false, true, link) +} + +// statusColor maps an air run status word to its data color. +func statusColor(status string) lipgloss.Color { + switch status { + case "SUCCESS": + return colGreen + case "RUNNING", "PENDING", "TERMINATING": + return colAmber + case "FAILED": + return colRed + default: // CANCELED / UNKNOWN + return colN7 + } +} + +// startedDisplay trims the row's ISO start timestamp to second precision +// ("2006-01-02T15:04:05"), or "-" when the run hasn't started. +func startedDisplay(r listRow) string { + if r.StartedAt == nil { + return "-" + } + s := *r.StartedAt + if len(s) >= 19 { + return s[:19] + } + return s +} + +// mlflowDisplay shortens an MLflow run URL to a "…/runs/" label; the +// OSC 8 target keeps the full URL. +func mlflowDisplay(url string) string { + id := mlflowRunID(url) + if id == "" { + return truncate(url, mlflowColWidth) + } + if len(id) > 8 { + id = id[:8] + "…" + } + return "…/runs/" + id +} + +// mlflowRunID extracts the run-id path segment from an MLflow URL. +func mlflowRunID(url string) string { + _, after, ok := strings.Cut(url, "/runs/") + if !ok { + return "" + } + id, _, _ := strings.Cut(after, "/") + return id +} + +// pad pads (or truncates) s to a visible width of n, right-aligned when right is +// set. It measures visible width, so it is safe on styled strings. +func pad(s string, n int, right bool) string { + s = truncate(s, n) + gap := max(n-lipgloss.Width(s), 0) + if gap == 0 { + return s + } + fill := strings.Repeat(" ", gap) + if right { + return fill + s + } + return s + fill +} + +// truncate shortens s to a visible width of n, appending "…" on overflow. +func truncate(s string, n int) string { + if lipgloss.Width(s) <= n { + return s + } + if n <= 1 { + return "…" + } + return string([]rune(s)[:n-1]) + "…" +} diff --git a/experimental/air/cmd/list_tui_test.go b/experimental/air/cmd/list_tui_test.go new file mode 100644 index 00000000000..ba709bbf35c --- /dev/null +++ b/experimental/air/cmd/list_tui_test.go @@ -0,0 +1,216 @@ +package aircmd + +import ( + "io" + "strconv" + "testing" + + tea "github.com/charmbracelet/bubbletea" + "github.com/charmbracelet/lipgloss" + "github.com/databricks/cli/libs/cmdio" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// testListRows is a small fixture covering each status color, a present and an +// absent MLflow link, and a still-running (no end) row. +func testListRows() []listRow { + return []listRow{ + {RunID: "1", Experiment: "qwen-train", User: "me@example.com", Status: "SUCCESS", StartedAt: new("2026-06-05T17:32:39.000000+00:00"), Duration: "1m 14s", MLflowURL: "https://h/ml/experiments/E/runs/04c41514fbb0/artifacts/logs/node_0", Accelerators: "8x H100"}, + {RunID: "2", Experiment: "llama-train", User: "me@example.com", Status: "RUNNING", StartedAt: new("2026-06-05T18:43:24.000000+00:00"), Duration: "3m 32s", MLflowURL: "-", Accelerators: "1x A10"}, + {RunID: "3", Experiment: "mixtral", User: "me@example.com", Status: "FAILED", StartedAt: nil, Duration: "-", MLflowURL: "-", Accelerators: "-"}, + } +} + +func testListModel(t *testing.T) listModel { + r, _ := cmdio.NewRenderer(cmdio.MockDiscard(t.Context()), io.Discard) + return newListModel(r, nil, testListRows(), false) +} + +func key(t *testing.T, m listModel, s string) listModel { + t.Helper() + next, _ := m.Update(tea.KeyMsg{Type: tea.KeyRunes, Runes: []rune(s)}) + return next.(listModel) +} + +func TestListModelNavigation(t *testing.T) { + m := testListModel(t) + require.Equal(t, 0, m.cursor) + + m = key(t, m, "j") + assert.Equal(t, 1, m.cursor) + m = key(t, key(t, m, "k"), "k") // clamp at top + assert.Equal(t, 0, m.cursor) + + for range len(m.rows) + 2 { // clamp at bottom + m = key(t, m, "j") + } + assert.Equal(t, len(m.rows)-1, m.cursor) +} + +func TestListModelWindowScrolls(t *testing.T) { + m := testListModel(t) + // Height 5 leaves room for ~2 rows (header + hint reserved). + next, _ := m.Update(tea.WindowSizeMsg{Width: 200, Height: 5}) + m = next.(listModel) + require.Equal(t, 2, m.visibleCount()) + require.Equal(t, 0, m.offset) + + m = key(t, key(t, m, "j"), "j") // move to row index 2, past the window + assert.Equal(t, 2, m.cursor) + assert.Equal(t, 1, m.offset, "window scrolled to keep the cursor visible") +} + +func TestListModelPageCap(t *testing.T) { + rows := make([]listRow, 50) + for i := range rows { + rows[i] = listRow{RunID: strconv.Itoa(i)} + } + r, _ := cmdio.NewRenderer(cmdio.MockDiscard(t.Context()), io.Discard) + m := newListModel(r, nil, rows, false) + + // A tall terminal still shows at most listPageRows per page. + next, _ := m.Update(tea.WindowSizeMsg{Width: 80, Height: 100}) + assert.Equal(t, listPageRows, next.(listModel).visibleCount()) +} + +func TestListModelPaging(t *testing.T) { + rows := make([]listRow, 10) + for i := range rows { + rows[i] = listRow{RunID: strconv.Itoa(i)} + } + r, _ := cmdio.NewRenderer(cmdio.MockDiscard(t.Context()), io.Discard) + m := newListModel(r, nil, rows, false) + + // Height 7 leaves a 4-row window (header + hint reserved). + next, _ := m.Update(tea.WindowSizeMsg{Width: 80, Height: 7}) + m = next.(listModel) + require.Equal(t, 4, m.visibleCount()) + + page := func(k tea.KeyType) { + n, _ := m.Update(tea.KeyMsg{Type: k}) + m = n.(listModel) + } + page(tea.KeyRight) + assert.Equal(t, 4, m.cursor) + page(tea.KeyEnd) + assert.Equal(t, 9, m.cursor) + page(tea.KeyLeft) + assert.Equal(t, 5, m.cursor) + page(tea.KeyHome) + assert.Equal(t, 0, m.cursor) +} + +func TestListModelMoreRows(t *testing.T) { + m := testListModel(t) + m.loading = true + before := len(m.rows) + + next, cmd := m.Update(moreRowsMsg{rows: []listRow{{RunID: "4"}, {RunID: "5"}}}) + m = next.(listModel) + + assert.False(t, m.loading, "loading cleared after a batch arrives") + assert.NoError(t, m.loadErr) + require.Len(t, m.rows, before+2) + assert.Equal(t, "5", m.rows[len(m.rows)-1].RunID, "new rows appended") + assert.Nil(t, cmd) +} + +func TestListModelMoreRowsError(t *testing.T) { + m := testListModel(t) + m.loading = true + before := len(m.rows) + + next, cmd := m.Update(moreRowsMsg{err: io.ErrUnexpectedEOF}) + m = next.(listModel) + + assert.False(t, m.loading) + assert.ErrorIs(t, m.loadErr, io.ErrUnexpectedEOF) + assert.Len(t, m.rows, before, "rows unchanged on error") + assert.Nil(t, cmd) +} + +func TestListModelMoreRowsEmptyKeepsPaging(t *testing.T) { + r, _ := cmdio.NewRenderer(cmdio.MockDiscard(t.Context()), io.Discard) + + // An empty page while more runs remain re-fetches; once exhausted it stops. + m := newListModel(r, &runFetcher{}, testListRows(), false) + m.loading = true + next, cmd := m.Update(moreRowsMsg{}) + m = next.(listModel) + assert.NotNil(t, cmd, "empty page with more to scan keeps paging") + + m.fetcher.exhausted = true + m.loading = true + _, cmd = m.Update(moreRowsMsg{}) + assert.Nil(t, cmd, "empty page stops once the fetcher is exhausted") +} + +func TestListModelQuit(t *testing.T) { + m := testListModel(t) + _, cmd := m.Update(tea.KeyMsg{Type: tea.KeyRunes, Runes: []rune("q")}) + require.NotNil(t, cmd) + assert.Equal(t, tea.QuitMsg{}, cmd()) +} + +func TestListModelView(t *testing.T) { + next, _ := testListModel(t).Update(tea.WindowSizeMsg{Width: 200, Height: 24}) + out := next.(listModel).View() + + assert.NotContains(t, out, "\x1b", "Ascii profile + no links should produce no escapes") + for _, want := range []string{ + "Run ID", "Experiment", "Status", "Started", "Duration", "MLflow", "User", "Accelerators", + "qwen-train", "● SUCCESS", "● RUNNING", "● FAILED", + "…/runs/04c41514…", // shortened MLflow link + "2026-06-05T17:32:39", // started trimmed to seconds + "▸", // selection gutter on the first row + "↑/↓ navigate", // hint line + } { + assert.Contains(t, out, want) + } +} + +func TestStaticListTable(t *testing.T) { + r, _ := cmdio.NewRenderer(cmdio.MockDiscard(t.Context()), io.Discard) + out := staticListTable(r, testListRows(), false) + + assert.NotContains(t, out, "\x1b") + assert.NotContains(t, out, "▸", "static table has no selection") + for _, want := range []string{"Run ID", "1", "qwen-train", "…/runs/04c41514…", "Accelerators"} { + assert.Contains(t, out, want) + } + + assert.Equal(t, "No runs found.\n", staticListTable(r, nil, false)) +} + +func TestStatusColor(t *testing.T) { + assert.Equal(t, colGreen, statusColor("SUCCESS")) + assert.Equal(t, colAmber, statusColor("RUNNING")) + assert.Equal(t, colAmber, statusColor("PENDING")) + assert.Equal(t, colRed, statusColor("FAILED")) + assert.Equal(t, colN7, statusColor("CANCELED")) + assert.Equal(t, colN7, statusColor("UNKNOWN")) +} + +func TestStartedDisplay(t *testing.T) { + assert.Equal(t, "-", startedDisplay(listRow{})) + assert.Equal(t, "2026-06-05T17:32:39", startedDisplay(listRow{StartedAt: new("2026-06-05T17:32:39.000000+00:00")})) +} + +func TestMLflowDisplay(t *testing.T) { + assert.Equal(t, "…/runs/04c41514…", mlflowDisplay("https://h/ml/experiments/E/runs/04c41514fbb0/artifacts/logs/node_0")) + assert.Equal(t, "…/runs/run1", mlflowDisplay("https://h/ml/experiments/E/runs/run1/artifacts/logs/node_0")) + assert.LessOrEqual(t, lipgloss.Width(mlflowDisplay("https://h/no-runs/here")), mlflowColWidth) +} + +func TestMLflowRunID(t *testing.T) { + assert.Equal(t, "abc123", mlflowRunID("https://h/ml/experiments/1/runs/abc123/artifacts")) + assert.Empty(t, mlflowRunID("https://h/no-runs-here")) +} + +func TestPadAndTruncate(t *testing.T) { + assert.Equal(t, "ab ", pad("ab", 5, false)) + assert.Equal(t, " ab", pad("ab", 5, true)) + assert.Equal(t, "abcd…", truncate("abcdefgh", 5)) + assert.Equal(t, "abc", truncate("abc", 5)) +} diff --git a/experimental/air/cmd/logs.go b/experimental/air/cmd/logs.go new file mode 100644 index 00000000000..c34fb62a7df --- /dev/null +++ b/experimental/air/cmd/logs.go @@ -0,0 +1,36 @@ +package aircmd + +import ( + "github.com/databricks/cli/cmd/root" + "github.com/spf13/cobra" +) + +func newLogsCommand() *cobra.Command { + var ( + node int + lines int + retry int + downloadTo string + review bool + ) + + cmd := &cobra.Command{ + Use: "logs JOB_RUN_ID", + 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(&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") + + return cmd +} diff --git a/experimental/air/cmd/mlflow.go b/experimental/air/cmd/mlflow.go new file mode 100644 index 00000000000..070caefbb06 --- /dev/null +++ b/experimental/air/cmd/mlflow.go @@ -0,0 +1,91 @@ +package aircmd + +import ( + "context" + "fmt" + "strings" + + "github.com/databricks/cli/libs/log" + "github.com/databricks/databricks-sdk-go" + "github.com/databricks/databricks-sdk-go/service/jobs" + "github.com/databricks/databricks-sdk-go/service/ml" +) + +// mlflowIdentifiers are the experiment and run IDs MLflow assigns to a run. +type mlflowIdentifiers struct { + ExperimentID string + RunID string +} + +// mlflowIDs fetches the MLflow IDs for a run via its latest task. Returns nil if +// they can't be obtained. +func mlflowIDs(ctx context.Context, w *databricks.WorkspaceClient, run *jobs.Run) *mlflowIdentifiers { + if len(run.Tasks) == 0 { + return nil + } + // The MLflow output is attached to the task run, not the parent job run. + return mlflowIDsForTask(ctx, w, run.Tasks[len(run.Tasks)-1].RunId) +} + +// mlflowIDsForTask fetches a task run's MLflow experiment and run IDs from +// runs/get-output, or nil if they can't be obtained. They drive a convenience +// link, so any failure (endpoint error, run not yet started, no MLflow output) +// is logged and treated as "no link" rather than failing the command. +func mlflowIDsForTask(ctx context.Context, w *databricks.WorkspaceClient, taskRunID int64) *mlflowIdentifiers { + if taskRunID == 0 { + return nil + } + + out, err := w.Jobs.GetRunOutputByRunId(ctx, taskRunID) + if err != nil { + log.Debugf(ctx, "air: could not fetch run output for MLflow link: %v", err) + return nil + } + + if o := out.AiRuntimeTaskOutput; o != nil && o.MlflowExperimentId != "" && o.MlflowRunId != "" { + return &mlflowIdentifiers{ExperimentID: o.MlflowExperimentId, RunID: o.MlflowRunId} + } + return nil +} + +// mlflowLogsURL is the deep link to a run's node-0 logs. It is the value of the +// JSON `mlflow_url` field, matching the Python CLI. +func mlflowLogsURL(host string, ids *mlflowIdentifiers) string { + return fmt.Sprintf("%s/ml/experiments/%s/runs/%s/artifacts/logs/node_0", + strings.TrimRight(host, "/"), ids.ExperimentID, ids.RunID) +} + +// mlflowRunURL links to the MLflow run page; it backs the MLflow Run hyperlink +// in the single-run view. +func mlflowRunURL(host string, ids *mlflowIdentifiers) string { + return fmt.Sprintf("%s/ml/experiments/%s/runs/%s", + strings.TrimRight(host, "/"), ids.ExperimentID, ids.RunID) +} + +// fetchMLflowRunName fetches a run's MLflow run_name via the MLflow REST API, +// returning "" if it can't be obtained. Best-effort, like the rest of the MLflow +// enrichment. +func fetchMLflowRunName(ctx context.Context, w *databricks.WorkspaceClient, mlflowRunID string) string { + resp, err := w.Experiments.GetRun(ctx, ml.GetRunRequest{RunId: mlflowRunID}) + if err != nil { + log.Debugf(ctx, "air get: could not fetch MLflow run name: %v", err) + return "" + } + if resp.Run == nil || resp.Run.Info == nil { + return "" + } + return resp.Run.Info.RunName +} + +// mlflowRunLabel is the text shown for the MLflow Run cell: the run's name, or +// "...{last 8 of run id}" when the name is unknown. Mirrors Python's +// _get_mlflow_run_name (cli_display.py). +func mlflowRunLabel(name, mlflowRunID string) string { + if name != "" { + return name + } + if len(mlflowRunID) > 8 { + return "..." + mlflowRunID[len(mlflowRunID)-8:] + } + return "..." + mlflowRunID +} diff --git a/experimental/air/cmd/mlflow_test.go b/experimental/air/cmd/mlflow_test.go new file mode 100644 index 00000000000..1501a2a59d8 --- /dev/null +++ b/experimental/air/cmd/mlflow_test.go @@ -0,0 +1,139 @@ +package aircmd + +import ( + "net/http" + "net/http/httptest" + "testing" + + "github.com/databricks/databricks-sdk-go" + "github.com/databricks/databricks-sdk-go/service/jobs" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// newTestWorkspaceClient builds a WorkspaceClient pointed at a mock HTTP server, +// so the SDK calls these tests exercise are served locally. +func newTestWorkspaceClient(t *testing.T, host string) *databricks.WorkspaceClient { + t.Helper() + w, err := databricks.NewWorkspaceClient(&databricks.Config{Host: host, Token: "token"}) + require.NoError(t, err) + return w +} + +// runOutputServer serves the given runs/get-output body and a stub for the SDK's +// well-known config discovery request. *hit is set when get-output is called. +func runOutputServer(t *testing.T, body string, hit *bool) *httptest.Server { + t.Helper() + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == "/api/2.2/jobs/runs/get-output" { + *hit = true + _, _ = w.Write([]byte(body)) + return + } + _, _ = w.Write([]byte(`{}`)) + })) + t.Cleanup(srv.Close) + return srv +} + +func TestMLflowIDs(t *testing.T) { + ctx := t.Context() + run := &jobs.Run{Tasks: []jobs.RunTask{{RunId: 99}}} + + t.Run("returns the identifiers on success", func(t *testing.T) { + var hit bool + srv := runOutputServer(t, `{"ai_runtime_task_output":{"mlflow_experiment_id":"E1","mlflow_run_id":"R1"}}`, &hit) + + got := mlflowIDs(ctx, newTestWorkspaceClient(t, srv.URL), run) + require.NotNil(t, got) + assert.True(t, hit, "runs/get-output should have been called") + assert.Equal(t, &mlflowIdentifiers{ExperimentID: "E1", RunID: "R1"}, got) + }) + + t.Run("nil when the run has no MLflow info", func(t *testing.T) { + var hit bool + srv := runOutputServer(t, `{}`, &hit) + assert.Nil(t, mlflowIDs(ctx, newTestWorkspaceClient(t, srv.URL), run)) + }) + + t.Run("nil when the run has no tasks", func(t *testing.T) { + // Returns before any HTTP call, so the host is never contacted. + assert.Nil(t, mlflowIDs(ctx, newTestWorkspaceClient(t, "https://unused.invalid"), &jobs.Run{})) + }) + + t.Run("uses the latest attempt's task run", func(t *testing.T) { + // A retried run must link to the last task, not the stale first attempt. + var gotRunID string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == "/api/2.2/jobs/runs/get-output" { + gotRunID = r.URL.Query().Get("run_id") + } + _, _ = w.Write([]byte(`{}`)) + })) + t.Cleanup(srv.Close) + + retried := &jobs.Run{Tasks: []jobs.RunTask{{RunId: 99}, {RunId: 100}}} + mlflowIDs(ctx, newTestWorkspaceClient(t, srv.URL), retried) + assert.Equal(t, "100", gotRunID) + }) +} + +func TestMLflowIDsForTask(t *testing.T) { + ctx := t.Context() + + t.Run("parses ai_runtime_task_output", func(t *testing.T) { + var hit bool + srv := runOutputServer(t, `{"ai_runtime_task_output":{"mlflow_experiment_id":"E1","mlflow_run_id":"R1"}}`, &hit) + got := mlflowIDsForTask(ctx, newTestWorkspaceClient(t, srv.URL), 99) + require.NotNil(t, got) + assert.True(t, hit) + assert.Equal(t, &mlflowIdentifiers{ExperimentID: "E1", RunID: "R1"}, got) + }) + + t.Run("nil when no task run id", func(t *testing.T) { + // Returns before any HTTP call, so the host is never contacted. + assert.Nil(t, mlflowIDsForTask(ctx, newTestWorkspaceClient(t, "https://unused.invalid"), 0)) + }) +} + +func TestMLflowURLs(t *testing.T) { + ids := &mlflowIdentifiers{ExperimentID: "E1", RunID: "R1"} + // A trailing slash on the host must not produce a double slash in the link. + assert.Equal(t, "https://h.test/ml/experiments/E1/runs/R1/artifacts/logs/node_0", mlflowLogsURL("https://h.test/", ids)) + assert.Equal(t, "https://h.test/ml/experiments/E1/runs/R1", mlflowRunURL("https://h.test", ids)) +} + +func TestMLflowRunLabel(t *testing.T) { + // Uses the run name when it is known. + assert.Equal(t, "sunny-cat-42", mlflowRunLabel("sunny-cat-42", "0123456789abcdef")) + // Falls back to the last 8 characters of a long run id. + assert.Equal(t, "...9abcdef0", mlflowRunLabel("", "0123456789abcdef0")) + // A short run id is shown in full behind the ellipsis. + assert.Equal(t, "...short", mlflowRunLabel("", "short")) +} + +func TestFetchMLflowRunName(t *testing.T) { + ctx := t.Context() + + mlflowServer := func(body string) *httptest.Server { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == "/api/2.0/mlflow/runs/get" { + _, _ = w.Write([]byte(body)) + return + } + _, _ = w.Write([]byte(`{}`)) + })) + t.Cleanup(srv.Close) + return srv + } + + t.Run("returns the run name", func(t *testing.T) { + srv := mlflowServer(`{"run":{"info":{"run_name":"sunny-cat-42"}}}`) + assert.Equal(t, "sunny-cat-42", fetchMLflowRunName(ctx, newTestWorkspaceClient(t, srv.URL), "run1")) + }) + + t.Run("empty when the run cannot be fetched", func(t *testing.T) { + srv := mlflowServer(`{}`) + assert.Empty(t, fetchMLflowRunName(ctx, newTestWorkspaceClient(t, srv.URL), "run1")) + }) +} diff --git a/experimental/air/cmd/output.go b/experimental/air/cmd/output.go new file mode 100644 index 00000000000..c00d870e386 --- /dev/null +++ b/experimental/air/cmd/output.go @@ -0,0 +1,82 @@ +package aircmd + +import ( + "context" + "time" + + "github.com/databricks/cli/cmd/root" + "github.com/databricks/cli/libs/cmdio" + "github.com/databricks/cli/libs/flags" + "github.com/spf13/cobra" +) + +// envelopeVersion is the envelope's format-version marker. The Python `air` CLI +// hardcodes it to 1; it lets consumers detect a future incompatible change to +// the envelope shape. +const envelopeVersion = 1 + +// envelope is the JSON shape that the AI runtime CLI prints: +// +// { "v": 1, "ts": "2024-01-15T14:30:45Z", "data": { ... } } +// +// It mirrors the envelope used by the original Python `air` CLI so existing +// consumers keep working after the port to Go. +type envelope struct { + // V is the envelope format-version marker (always 1). + V int `json:"v"` + // TS is the wall-clock time the response was produced, in RFC 3339 UTC. + // It is an absolute timestamp, not an elapsed duration. + TS string `json:"ts"` + // Data is the command-specific payload. + Data any `json:"data"` +} + +// renderEnvelope wraps data in the JSON envelope and prints it. +// Fields that should appear only in text output are tagged `json:"-"` on the payload struct. +func renderEnvelope(ctx context.Context, data any) error { + return cmdio.Render(ctx, envelope{ + V: envelopeVersion, + TS: time.Now().UTC().Format(time.RFC3339), + Data: data, + }) +} + +// jsonError is the error payload, matching the Python `air` CLI's shape (cli/json_output.py). +type jsonError struct { + Code string `json:"code"` + Kind string `json:"kind"` + Message string `json:"message"` + Retryable bool `json:"retryable"` +} + +// errorEnvelope is what a failed command prints in JSON mode: +// +// { "v": 1, "ts": "...", "error": { "code": ..., "kind": ..., "message": ..., "retryable": ... } } +type errorEnvelope struct { + V int `json:"v"` + TS string `json:"ts"` + Error jsonError `json:"error"` +} + +// renderError prints err as a JSON error envelope when output is JSON, returning +// root.ErrAlreadyPrinted so the command exits non-zero without Cobra reprinting +// it; in text mode it returns err unchanged. code/kind/retryable match the +// Python CLI's call site. +func renderError(ctx context.Context, cmd *cobra.Command, code, kind string, retryable bool, err error) error { + if root.OutputType(cmd) != flags.OutputJSON { + return err + } + if rerr := cmdio.Render(ctx, errorEnvelope{ + V: envelopeVersion, + TS: time.Now().UTC().Format(time.RFC3339), + Error: jsonError{ + Code: code, + Kind: kind, + Message: err.Error(), + Retryable: retryable, + }, + }); rerr != nil { + return rerr + } + return root.ErrAlreadyPrinted +} diff --git a/experimental/air/cmd/output_test.go b/experimental/air/cmd/output_test.go new file mode 100644 index 00000000000..3c35dc60a67 --- /dev/null +++ b/experimental/air/cmd/output_test.go @@ -0,0 +1,52 @@ +package aircmd + +import ( + "bytes" + "encoding/json" + "errors" + "testing" + + "github.com/databricks/cli/cmd/root" + "github.com/databricks/cli/libs/cmdio" + "github.com/databricks/cli/libs/flags" + "github.com/spf13/cobra" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestRenderEnvelope(t *testing.T) { + ctx := cmdio.MockDiscard(t.Context()) + require.NoError(t, renderEnvelope(ctx, getData{RunID: "1", Status: "RUNNING"})) +} + +// withOutput registers the --output flag on cmd and sets it, mirroring how the +// root command wires output mode in production. Subcommand unit tests need it +// because they invoke RunE without going through the root command. +func withOutput(cmd *cobra.Command, output flags.Output) *cobra.Command { + cmd.Flags().Var(&output, "output", "") + return cmd +} + +func TestRenderErrorJSON(t *testing.T) { + var buf bytes.Buffer + ctx := cmdio.InContext(t.Context(), cmdio.NewIO(t.Context(), flags.OutputJSON, nil, &buf, &buf, "", "")) + cmd := withOutput(&cobra.Command{}, flags.OutputJSON) + + err := renderError(ctx, cmd, "NOT_FOUND", "NOT_FOUND", false, errors.New("run 1 not found")) + // JSON mode prints the envelope, so Cobra must stay silent but still exit non-zero. + require.ErrorIs(t, err, root.ErrAlreadyPrinted) + + // The envelope must match the Python air CLI's print_json_error shape exactly. + var got errorEnvelope + require.NoError(t, json.Unmarshal(buf.Bytes(), &got)) + assert.Equal(t, 1, got.V) + assert.NotEmpty(t, got.TS) + assert.Equal(t, jsonError{Code: "NOT_FOUND", Kind: "NOT_FOUND", Message: "run 1 not found"}, got.Error) +} + +func TestRenderErrorText(t *testing.T) { + ctx := cmdio.MockDiscard(t.Context()) + cmd := withOutput(&cobra.Command{}, flags.OutputText) + want := errors.New("run 1 not found") + require.Equal(t, want, renderError(ctx, cmd, "NOT_FOUND", "NOT_FOUND", false, want)) +} diff --git a/experimental/air/cmd/register_image.go b/experimental/air/cmd/register_image.go new file mode 100644 index 00000000000..1d8b45044a7 --- /dev/null +++ b/experimental/air/cmd/register_image.go @@ -0,0 +1,33 @@ +package aircmd + +import ( + "github.com/databricks/cli/cmd/root" + "github.com/spf13/cobra" +) + +func newRegisterImageCommand() *cobra.Command { + var ( + scope string + key string + interactiveAuth bool + 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") + }, + } + + 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") + cmd.Flags().IntVar(&timeoutMinutes, "timeout-minutes", 60, "Timeout to wait for the image to become available") + + return cmd +} diff --git a/experimental/air/cmd/render.go b/experimental/air/cmd/render.go new file mode 100644 index 00000000000..f3fb63db807 --- /dev/null +++ b/experimental/air/cmd/render.go @@ -0,0 +1,388 @@ +package aircmd + +import ( + "context" + "fmt" + "io" + "strconv" + "strings" + + "github.com/charmbracelet/lipgloss" + "github.com/databricks/cli/libs/cmdio" + "github.com/databricks/cli/libs/log" + "github.com/databricks/databricks-sdk-go" + "github.com/databricks/databricks-sdk-go/service/jobs" + "github.com/muesli/termenv" + "go.yaml.in/yaml/v3" +) + +// Box titles, rendered into the top border of each box. +const ( + configBoxTitle = "Configuration" + metadataBoxTitle = "Metadata" +) + +// minBoxInnerWidth keeps all boxes a uniform, comfortable width; boxHPad and +// boxVPad are the horizontal and vertical padding inside each box. +const ( + minBoxInnerWidth = 60 + boxHPad = 2 + boxVPad = 1 +) + +// palette holds the lipgloss styles for the single-run view. Two layers: a +// neutral ramp for chrome and text, and an accent palette for syntax and data. +// All styles come from one renderer, so they honor its color profile (Ascii +// under --no-color / non-TTY, which strips every escape). +type palette struct { + n7 lipgloss.Style // dim: block indicator, empty progress, command-block "|" + n8 lipgloss.Style // muted: field labels + n12 lipgloss.Style // content: config and metadata values, percent + + border lipgloss.Style // box borders and titles + + blue lipgloss.Style // yaml keys and hyperlinks + green lipgloss.Style // success status and progress fill + amber lipgloss.Style // in-progress status + red lipgloss.Style // failed status +} + +func newPalette(r *lipgloss.Renderer) palette { + fg := func(hex string) lipgloss.Style { return r.NewStyle().Foreground(lipgloss.Color(hex)) } + return palette{ + n7: fg("#6E6E70"), + n8: fg("#8C8A86"), + n12: fg("#F9F7F4"), // Oat Light + border: fg("#B7A8E8"), // light purple (box borders and titles) + blue: fg("#8FB3DC"), + green: fg("#74C39A"), + amber: fg("#DCAA5C"), + red: fg("#D9756B"), + } +} + +// runView is the resolved, display-ready data the renderer draws. It is built +// from getData plus the MLflow enrichment, so the renderer itself does no API +// calls or formatting decisions. +type runView struct { + runID string + dashboardURL string + status string + submitted string + retries int + maxRetries string + duration string + experiment string + mlflowLabel string + mlflowURL string + user string + accelerators string + environment string +} + +// renderRunText writes the styled single-run view: a training-config box, a +// completed progress bar, and a field list, separated by blank lines. It is a +// one-shot renderer — it builds the full string and writes it once, with no +// streaming, spinner, or redraw. +func renderRunText(ctx context.Context, out io.Writer, w *databricks.WorkspaceClient, run *jobs.Run, data *getData, ids *mlflowIdentifiers) { + // cmdio.NewRenderer sets the Ascii color profile when color is off, which + // emits no SGR codes; combined with the link fallback below this gives clean, + // un-escaped output under --no-color / NO_COLOR / piped stdout. + renderer, colorOn := cmdio.NewRenderer(ctx, out) + p := newPalette(renderer) + + view := runView{ + runID: data.RunID, + dashboardURL: data.DashboardURL, + status: data.Status, + submitted: data.SubmittedDisplay, + retries: data.AttemptNumber, + maxRetries: data.MaxRetriesDisplay, + duration: data.DurationDisplay, + experiment: data.ExperimentDisplay, + mlflowLabel: na, + user: data.UserDisplay, + accelerators: data.AcceleratorsDisplay, + environment: data.EnvironmentDisplay, + } + + if ids != nil { + view.mlflowLabel = mlflowRunLabel(fetchMLflowRunName(ctx, w, ids.RunID), ids.RunID) + view.mlflowURL = mlflowRunURL(w.Config.Host, ids) + } + + var sections []string + if body := colorizeConfig(p, resolveConfigYAML(ctx, w, run, data)); body != "" { + sections = append(sections, renderBox(p, configBoxTitle, body)) + } + sections = append(sections, renderBox(p, metadataBoxTitle, renderFields(p, colorOn, view))) + + // A single write: a blank line before the first box and after the last, and + // one between each box. + fmt.Fprintf(out, "\n%s\n\n", strings.Join(sections, "\n\n")) + + // Bare-URL footer so the job run / MLflow links remain reachable when + // stdout is not a hyperlink-capable terminal (piped, redirected, NO_COLOR). + // In that case the OSC 8 hyperlinks on the Run ID / MLflow Run cells + // degrade to plain labels and the URLs would otherwise disappear from text + // output, breaking workflows like `air get X > out.txt` or + // `NO_COLOR=1 air get X` that the previous `Job Link:` line supported. + if view.dashboardURL != "" { + fmt.Fprintf(out, "Run URL: %s\n", view.dashboardURL) + } + if view.mlflowURL != "" { + fmt.Fprintf(out, "MLflow URL: %s\n", view.mlflowURL) + } +} + +// genAIComputeTask returns the run's first GenAI-compute task, or nil. +func genAIComputeTask(run *jobs.Run) *jobs.GenAiComputeTask { + if len(run.Tasks) == 0 { + return nil + } + return run.Tasks[0].GenAiComputeTask +} + +// resolveConfigYAML returns the config box body: from the downloaded config file +// when we have its path, else from the legacy task. +func resolveConfigYAML(ctx context.Context, w *databricks.WorkspaceClient, run *jobs.Run, data *getData) string { + if data.TrainingConfigPath != "" { + if raw := downloadConfig(ctx, w, data.TrainingConfigPath); len(raw) > 0 { + return strings.TrimRight(reformatYAMLForDisplay(raw), "\n") + } + } + if task := genAIComputeTask(run); task != nil { + return configYAML(ctx, w, task) + } + return "" +} + +// configYAML returns the run's resolved training config as YAML for the box. The +// full config (including the command/script) lives in the run's parameters, not +// the structured task fields, so we prefer the inline parameters, then the +// parameters file, and only synthesize a minimal config as a last resort. +func configYAML(ctx context.Context, w *databricks.WorkspaceClient, task *jobs.GenAiComputeTask) string { + if task.YamlParameters != "" { + return strings.TrimRight(reformatYAMLForDisplay([]byte(task.YamlParameters)), "\n") + } + if task.YamlParametersFilePath != "" { + if raw := downloadConfig(ctx, w, task.YamlParametersFilePath); len(raw) > 0 { + return strings.TrimRight(reformatYAMLForDisplay(raw), "\n") + } + } + return synthConfigYAML(task) +} + +// downloadConfig fetches the run's training-config file, returning nil on +// failure (logged as a warning). Best-effort, like the rest of the enrichment. +func downloadConfig(ctx context.Context, w *databricks.WorkspaceClient, path string) []byte { + r, err := w.Workspace.Download(ctx, path) + if err != nil { + log.Warnf(ctx, "air get: could not download training config %s: %v", path, err) + return nil + } + defer r.Close() + content, err := io.ReadAll(r) + if err != nil { + log.Warnf(ctx, "air get: could not read training config %s: %v", path, err) + return nil + } + return content +} + +// configBox describes the synthesized config we marshal when the run exposes no +// parameters, in the order the fields are shown. +type configBox struct { + ExperimentName string `yaml:"experiment_name,omitempty"` + Compute *configCompute `yaml:"compute,omitempty"` +} + +type configCompute struct { + AcceleratorType string `yaml:"accelerator_type,omitempty"` + NumAccelerators int `yaml:"num_accelerators,omitempty"` +} + +// synthConfigYAML builds a minimal config from the structured task fields. It +// omits the command, which is only available in the run parameters. +func synthConfigYAML(task *jobs.GenAiComputeTask) string { + cfg := configBox{} + if task.MlflowExperimentName != "" { + cfg.ExperimentName = stripExperimentUserPrefix(task.MlflowExperimentName) + } + if task.Compute != nil && task.Compute.NumGpus > 0 { + cfg.Compute = &configCompute{ + AcceleratorType: task.Compute.GpuType, + NumAccelerators: task.Compute.NumGpus, + } + } + if cfg.ExperimentName == "" && cfg.Compute == nil { + return "" + } + b, err := yaml.Marshal(cfg) + if err != nil { + return "" + } + return strings.TrimRight(reformatYAMLForDisplay(b), "\n") +} + +// colorizeConfig styles a YAML config block line by line. +func colorizeConfig(p palette, body string) string { + if body == "" { + return "" + } + lines := strings.Split(body, "\n") + for i, line := range lines { + lines[i] = colorizeConfigLine(p, line) + } + return strings.Join(lines, "\n") +} + +// colorizeConfigLine colors one YAML line: keys blue, the `|` block indicator +// dim, and every value (and the command body that isn't a `key:` pair) in the +// neutral content color. +func colorizeConfigLine(p palette, line string) string { + indent := line[:len(line)-len(strings.TrimLeft(line, " "))] + trimmed := strings.TrimLeft(line, " ") + + if i := strings.IndexByte(trimmed, ':'); i > 0 && isConfigKey(trimmed[:i]) { + key := trimmed[:i] + value := strings.TrimSpace(trimmed[i+1:]) + styled := indent + p.blue.Render(key+":") + switch value { + case "": + // A mapping parent such as "compute:" has no value of its own. + case "|": + styled += " " + p.n7.Render(value) + default: + styled += " " + p.n12.Render(value) + } + return styled + } + return indent + p.n12.Render(trimmed) +} + +// isConfigKey reports whether s is a bare YAML key (lowercase, digits, and +// underscores). It guards against treating a colon inside a command body as a +// key/value separator. +func isConfigKey(s string) bool { + if s == "" { + return false + } + for _, r := range s { + if r != '_' && (r < 'a' || r > 'z') && (r < '0' || r > '9') { + return false + } + } + return true +} + +// renderBox draws a rounded-border box around body, with title rendered into the +// top border in the border color. body lines are padded to the widest one (or +// minBoxInnerWidth), with boxHPad columns and boxVPad rows of padding inside. +func renderBox(p palette, title, body string) string { + border := lipgloss.RoundedBorder() + lines := strings.Split(body, "\n") + pad := strings.Repeat(" ", boxHPad) + + titleWidth := lipgloss.Width(title) + inner := max(minBoxInnerWidth, titleWidth+2) + for _, line := range lines { + inner = max(inner, lipgloss.Width(line)) + } + + left := p.border.Render(border.Left) + right := p.border.Render(border.Right) + blank := left + strings.Repeat(" ", inner+2*boxHPad) + right + + var b strings.Builder + // Top: ╭─ ──…──╮. The dash count makes the row width match the body, + // accounting for the boxHPad columns on each side. + trailing := inner + 2*boxHPad - titleWidth - 3 + b.WriteString(p.border.Render(border.TopLeft + border.Top)) + b.WriteString(" " + p.border.Render(title) + " ") + b.WriteString(p.border.Render(strings.Repeat(border.Top, trailing) + border.TopRight)) + b.WriteByte('\n') + + for range boxVPad { + b.WriteString(blank + "\n") + } + for _, line := range lines { + fill := strings.Repeat(" ", inner-lipgloss.Width(line)) + b.WriteString(left + pad + line + fill + pad + right) + b.WriteByte('\n') + } + for range boxVPad { + b.WriteString(blank + "\n") + } + + b.WriteString(p.border.Render(border.BottomLeft + strings.Repeat(border.Bottom, inner+2*boxHPad) + border.BottomRight)) + return b.String() +} + +// renderFields draws the two-column summary: muted labels right-padded to the +// longest one, neutral values, a status-colored Status, and blue Run ID / MLflow +// Run hyperlinks. +func renderFields(p palette, colorOn bool, v runView) string { + status := statusStyle(p, v.status).Render("● " + v.status) + rows := []string{ + field(p, "Run ID", link(colorOn, p.blue, v.runID, v.dashboardURL)), + field(p, "Status", status), + field(p, "Submitted", p.n12.Render(v.submitted)), + field(p, "Retries", p.n12.Render(strconv.Itoa(v.retries))), + field(p, "Max Retries", p.n12.Render(v.maxRetries)), + field(p, "Duration", p.n12.Render(v.duration)), + field(p, "Experiment", p.n12.Render(v.experiment)), + field(p, "MLflow Run", link(colorOn, p.blue, v.mlflowLabel, v.mlflowURL)), + field(p, "User", p.n12.Render(v.user)), + field(p, "Accelerators", p.n12.Render(v.accelerators)), + field(p, "Environment", p.n12.Render(v.environment)), + } + return strings.Join(rows, "\n") +} + +// fieldLabelWidth is the width of the longest label ("Accelerators"), so values +// line up in a single column. +const fieldLabelWidth = len("Accelerators") + +func field(p palette, label, value string) string { + return p.n8.Render(label+strings.Repeat(" ", fieldLabelWidth-len(label))) + " " + value +} + +// link renders label as an OSC 8 terminal hyperlink to url in the given style +// (underlined). With color off (or no url) it is just the styled label so the +// box stays aligned; the URLs remain available in JSON output. +func link(colorOn bool, style lipgloss.Style, label, url string) string { + if !colorOn || url == "" { + return style.Render(label) + } + // Wrap the already-styled label in the hyperlink. Passing the OSC 8 escape + // through lipgloss.Render instead corrupts it: lipgloss re-styles each rune + // and splits the "\x1b]8;;" introducer, so the terminal can't parse the + // sequence and prints it literally. + return termenv.Hyperlink(url, style.Underline(true).Render(label)) +} + +// statusStyle maps a run status to its accent color: green for success, red for +// terminal failures, amber for everything still in flight. +func statusStyle(p palette, status string) lipgloss.Style { + switch { + case isSuccessStatus(status): + return p.green + case isFailedStatus(status): + return p.red + default: + return p.amber + } +} + +func isSuccessStatus(status string) bool { + return status == "SUCCESS" +} + +func isFailedStatus(status string) bool { + switch status { + case "FAILED", "TIMEDOUT", "CANCELED", "INTERNAL_ERROR", "UPSTREAM_FAILED", "UPSTREAM_CANCELED": + return true + } + return false +} diff --git a/experimental/air/cmd/render_test.go b/experimental/air/cmd/render_test.go new file mode 100644 index 00000000000..2fe125bcf04 --- /dev/null +++ b/experimental/air/cmd/render_test.go @@ -0,0 +1,189 @@ +package aircmd + +import ( + "io" + "strings" + "testing" + + "github.com/charmbracelet/lipgloss" + "github.com/databricks/cli/libs/cmdio" + "github.com/databricks/databricks-sdk-go/experimental/mocks" + "github.com/databricks/databricks-sdk-go/service/jobs" + "github.com/muesli/termenv" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/mock" +) + +// asciiPalette returns a palette whose styles emit no escape codes, so render +// output is plain text and assertions stay readable. +func asciiPalette(t *testing.T) palette { + r, _ := cmdio.NewRenderer(cmdio.MockDiscard(t.Context()), io.Discard) + return newPalette(r) +} + +func TestConfigYAML(t *testing.T) { + ctx := t.Context() + + t.Run("inline parameters include the command as a block literal", func(t *testing.T) { + task := &jobs.GenAiComputeTask{ + YamlParameters: "experiment_name: my-exp\ncompute:\n accelerator_type: a10\n num_accelerators: 1\ncommand: \"for i in $(seq 1 3); do echo $i; done\\n\"\n", + } + got := configYAML(ctx, mocks.NewMockWorkspaceClient(t).WorkspaceClient, task) + assert.Contains(t, got, "experiment_name: my-exp") + assert.Contains(t, got, "accelerator_type: a10") + assert.Contains(t, got, "command: |") + assert.Contains(t, got, " for i in $(seq 1 3); do echo $i; done") + }) + + t.Run("downloads the parameters file when there are no inline parameters", func(t *testing.T) { + m := mocks.NewMockWorkspaceClient(t) + m.GetMockWorkspaceAPI().EXPECT(). + Download(mock.Anything, "/Workspace/cfg.yaml"). + Return(io.NopCloser(strings.NewReader("experiment_name: from-file\n")), nil) + task := &jobs.GenAiComputeTask{YamlParametersFilePath: "/Workspace/cfg.yaml"} + assert.Equal(t, "experiment_name: from-file", configYAML(ctx, m.WorkspaceClient, task)) + }) + + t.Run("falls back to a synthesized config", func(t *testing.T) { + task := &jobs.GenAiComputeTask{ + MlflowExperimentName: "/Users/me@example.com/exp", + Compute: &jobs.ComputeConfig{GpuType: "a10", NumGpus: 1}, + } + got := configYAML(ctx, mocks.NewMockWorkspaceClient(t).WorkspaceClient, task) + assert.Contains(t, got, "experiment_name: exp") + assert.NotContains(t, got, "command") + }) +} + +func TestResolveConfigYAML(t *testing.T) { + ctx := t.Context() + + t.Run("downloads the training config path for new runs", func(t *testing.T) { + m := mocks.NewMockWorkspaceClient(t) + m.GetMockWorkspaceAPI().EXPECT(). + Download(mock.Anything, "/Workspace/run/training_config.yaml"). + Return(io.NopCloser(strings.NewReader("experiment_name: from-air\n")), nil) + run := &jobs.Run{Tasks: []jobs.RunTask{{}}} + data := &getData{TrainingConfigPath: "/Workspace/run/training_config.yaml"} + assert.Equal(t, "experiment_name: from-air", resolveConfigYAML(ctx, m.WorkspaceClient, run, data)) + }) + + t.Run("falls back to the legacy task when no path is set", func(t *testing.T) { + run := &jobs.Run{Tasks: []jobs.RunTask{{GenAiComputeTask: &jobs.GenAiComputeTask{ + MlflowExperimentName: "/Users/me@example.com/exp", + }}}} + got := resolveConfigYAML(ctx, mocks.NewMockWorkspaceClient(t).WorkspaceClient, run, &getData{}) + assert.Contains(t, got, "experiment_name: exp") + }) + + t.Run("empty when neither source is present", func(t *testing.T) { + run := &jobs.Run{Tasks: []jobs.RunTask{{}}} + assert.Empty(t, resolveConfigYAML(ctx, mocks.NewMockWorkspaceClient(t).WorkspaceClient, run, &getData{})) + }) +} + +func TestSynthConfigYAML(t *testing.T) { + task := &jobs.GenAiComputeTask{ + MlflowExperimentName: "/Users/me@example.com/stream-latency-test", + Compute: &jobs.ComputeConfig{GpuType: "a10", NumGpus: 1}, + } + // The accelerator_type uses the raw GPU type; the command is omitted because + // it lives only in the run parameters. + want := "experiment_name: stream-latency-test\n" + + "compute:\n" + + " accelerator_type: a10\n" + + " num_accelerators: 1" + assert.Equal(t, want, synthConfigYAML(task)) + assert.Empty(t, synthConfigYAML(&jobs.GenAiComputeTask{})) +} + +func TestColorizeConfigLine(t *testing.T) { + p := asciiPalette(t) + // Under the Ascii profile colorization adds no escapes, so each line is + // preserved verbatim (indentation included) regardless of its role. + for _, line := range []string{ + "experiment_name: stream-latency-test", + "compute:", + " accelerator_type: a10", + " num_accelerators: 1", + "command: |", + ` for i in $(seq 1 10); do echo "step $i"; done`, + } { + assert.Equal(t, line, colorizeConfigLine(p, line)) + } +} + +func TestIsConfigKey(t *testing.T) { + assert.True(t, isConfigKey("experiment_name")) + assert.True(t, isConfigKey("num_accelerators")) + assert.False(t, isConfigKey("")) + assert.False(t, isConfigKey("for i in $(seq 1 10); do echo ")) + assert.False(t, isConfigKey("Command")) +} + +func TestRenderBox(t *testing.T) { + p := asciiPalette(t) + out := renderBox(p, configBoxTitle, "experiment_name: stream-latency-test\ncompute:") + lines := strings.Split(out, "\n") + + // Title sits in the top border; corners are rounded; every row is the same width. + assert.Contains(t, lines[0], "╭─ "+configBoxTitle+" ") + assert.True(t, strings.HasSuffix(lines[0], "╮")) + assert.True(t, strings.HasPrefix(lines[len(lines)-1], "╰")) + assert.Contains(t, out, "│ experiment_name: stream-latency-test") + + width := lipgloss.Width(lines[0]) + for _, l := range lines { + assert.Equal(t, width, lipgloss.Width(l)) + } +} + +func TestRenderFields(t *testing.T) { + p := asciiPalette(t) + out := renderFields(p, false, runView{ + runID: "836121283738861", + dashboardURL: "https://h.test/jobs/runs/836121283738861", + status: "SUCCESS", + submitted: "2026-06-03 04:17 UTC", + retries: 0, + maxRetries: "3", + duration: "1m 13s", + experiment: "stream-latency-test", + mlflowLabel: "stream-latency-test", + mlflowURL: "https://h.test/ml/experiments/E1/runs/R1", + user: "riddhi.bhagwat@databricks.com", + accelerators: "1x A10", + environment: "ml-runtime-gpu:1.0", + }) + + // Labels are padded to the longest ("Accelerators"), so values align. + assert.Contains(t, out, "Run ID ") + assert.Contains(t, out, "Accelerators 1x A10") + // Max retries and environment show alongside the other fields. + assert.Contains(t, out, "Max Retries 3") + assert.Contains(t, out, "Environment ml-runtime-gpu:1.0") + // The status carries its dot prefix. + assert.Contains(t, out, "● SUCCESS") + // Off a terminal, links render as the bare label (URLs live in JSON output). + assert.Contains(t, out, "Run ID 836121283738861") + assert.NotContains(t, out, "https://h.test") + // The field list is a tight block: no blank lines. + assert.NotContains(t, out, "\n\n") +} + +func TestLink(t *testing.T) { + p := asciiPalette(t) + // Color off: the bare label, no URL. + assert.Equal(t, "label", link(false, p.blue, "label", "https://h.test")) + assert.Equal(t, "label", link(false, p.blue, "label", "")) + // With color on, the label is wrapped in an OSC 8 hyperlink to the url. + assert.Contains(t, link(true, p.blue, "label", "https://h.test"), termenv.Hyperlink("https://h.test", "label")) +} + +func TestStatusStyleSelectors(t *testing.T) { + assert.True(t, isSuccessStatus("SUCCESS")) + assert.False(t, isSuccessStatus("RUNNING")) + assert.True(t, isFailedStatus("FAILED")) + assert.True(t, isFailedStatus("TIMEDOUT")) + assert.False(t, isFailedStatus("RUNNING")) +} diff --git a/experimental/air/cmd/run.go b/experimental/air/cmd/run.go new file mode 100644 index 00000000000..bd32810e9bc --- /dev/null +++ b/experimental/air/cmd/run.go @@ -0,0 +1,98 @@ +package aircmd + +import ( + "errors" + "fmt" + "strconv" + + "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/spf13/cobra" +) + +// runResult is the JSON payload for `air run`. +type runResult struct { + Status string `json:"status"` + DryRun bool `json:"dry_run,omitempty"` + RunID string `json:"run_id,omitempty"` + DashboardURL string `json:"dashboard_url,omitempty"` +} + +func newRunCommand() *cobra.Command { + var ( + file string + watch bool + overrides []string + dryRun bool + idempotencyKey string + ) + + cmd := &cobra.Command{ + Use: "run", + Args: root.NoArgs, + Short: "Submit a training workload from a YAML config", + Long: `Submit a training workload to Databricks serverless GPU compute. + +The workload is described by a YAML config file (see --file).`, + } + + cmd.Flags().StringVarP(&file, "file", "f", "", "Path to the workload YAML config") + cmd.Flags().BoolVar(&watch, "watch", false, "Stream logs until the run completes") + cmd.Flags().StringArrayVar(&overrides, "override", nil, "Override a YAML field, e.g. compute.num_accelerators=8 (repeatable)") + cmd.Flags().BoolVar(&dryRun, "dry-run", false, "Validate the config without submitting") + cmd.Flags().StringVar(&idempotencyKey, "idempotency-key", "", "Return the existing run if this key was already used") + _ = cmd.MarkFlagRequired("file") + + // --dry-run only validates the config locally, so it needs no workspace. + // Submission requires an authenticated client. + cmd.PreRunE = func(cmd *cobra.Command, args []string) error { + if dryRun { + return nil + } + return root.MustWorkspaceClient(cmd, args) + } + + 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) + if err != nil { + return err + } + + if dryRun { + if root.OutputType(cmd) == flags.OutputText { + cmdio.LogString(ctx, fmt.Sprintf("Dry run: configuration for %q is valid; not submitting.", cfg.ExperimentName)) + return nil + } + return renderEnvelope(ctx, runResult{Status: "DRY_RUN_OK", DryRun: true}) + } + + w := cmdctx.WorkspaceClient(ctx) + runID, dashboardURL, err := submitWorkload(ctx, w, cfg, file, idempotencyKey) + if err != nil { + return err + } + + runIDStr := strconv.FormatInt(runID, 10) + if root.OutputType(cmd) == flags.OutputText { + cmdio.LogString(ctx, "Submitted run "+runIDStr) + cmdio.LogString(ctx, "View at: "+dashboardURL) + return nil + } + return renderEnvelope(ctx, runResult{Status: "SUBMITTED", RunID: runIDStr, DashboardURL: dashboardURL}) + } + + return cmd +} diff --git a/experimental/air/cmd/runconfig.go b/experimental/air/cmd/runconfig.go new file mode 100644 index 00000000000..09437f50a5b --- /dev/null +++ b/experimental/air/cmd/runconfig.go @@ -0,0 +1,462 @@ +package aircmd + +import ( + "errors" + "fmt" + "maps" + "regexp" + "slices" + "strings" + + "go.yaml.in/yaml/v3" +) + +// This file ports the run YAML schema and its structural validation from the +// Python CLI's sdk/config.py. "Structural" means types, required fields, and +// format/cross-field rules that need no workspace access. Online checks (e.g. +// GPU availability) and git/filesystem checks run at launch time and are +// intentionally not ported here. +// +// Divergences from the Python schema: compute.node_pool_id / compute.pool_name +// (see compute.go) and the top-level `priority` field are dropped because AIR +// does not support node-pool placement. priority is a pool-queue-ordering knob, +// so it goes with the pool fields. + +// REGEX_TASK_KEY_CHARS: ASCII alphanumeric, hyphen, underscore only (no periods). +// Explicit ASCII class, not \w: \w matches Unicode letters that the ASCII-only +// Jobs API task_key rejects. +var taskKeyRe = regexp.MustCompile(`^[A-Za-z0-9_-]+$`) + +// gitRefRe guards the branch name against command injection (it flows into git +// exec args). Only safe ref characters are allowed. +var gitRefRe = regexp.MustCompile(`^[\w./-]+$`) + +// runConfig is the top-level run YAML schema: experiment_name + compute / +// environment / code_source plus the command and run options. +type runConfig struct { + ExperimentName string `yaml:"experiment_name"` + Compute *computeConfig `yaml:"compute"` + Environment *environmentConfig `yaml:"environment"` + Command *string `yaml:"command"` + EnvVariables map[string]string `yaml:"env_variables"` + Secrets map[string]string `yaml:"secrets"` + CodeSource *codeSourceConfig `yaml:"code_source"` + // MaxRetries defaults to 3 when unset; default-filling is a normalization + // concern handled at launch, so a nil pointer is left as-is here. + MaxRetries *int `yaml:"max_retries"` + TimeoutMinutes *int `yaml:"timeout_minutes"` + IdempotencyToken *string `yaml:"idempotency_token"` + Parameters map[string]any `yaml:"parameters"` + MLflowRunName *string `yaml:"mlflow_run_name"` + MLflowExperimentDirectory *string `yaml:"mlflow_experiment_directory"` + Permissions []permission `yaml:"permissions"` + UsagePolicyName *string `yaml:"usage_policy_name"` + UsagePolicyID *string `yaml:"usage_policy_id"` +} + +// validate runs structural validation over the whole config, returning the first +// failure. Fields are checked in declaration order to keep error output stable. +func (c *runConfig) validate() error { + if err := validateExperimentName(c.ExperimentName); err != nil { + return err + } + + if c.Compute == nil { + return errors.New("compute: section is required") + } + if err := c.Compute.validate(); err != nil { + return err + } + + if c.Environment != nil { + if err := c.Environment.validate(); err != nil { + return err + } + } + + // command is optional in the type system but required in practice, matching + // the Python validate_script_fields model validator. + if c.Command == nil { + return errors.New("command is required") + } + if err := validateCommand(*c.Command); err != nil { + return err + } + + if err := validateSecretRefs(c.Secrets); err != nil { + return err + } + + // A name can't be both a plain env var and a secret: the precedence would be + // ambiguous and could leak the secret. Sorted for a stable error. + for _, name := range slices.Sorted(maps.Keys(c.EnvVariables)) { + if _, ok := c.Secrets[name]; ok { + return fmt.Errorf("%q is set in both env_variables and secrets; remove it from one", name) + } + } + + if c.CodeSource != nil { + if err := c.CodeSource.validate(); err != nil { + return err + } + } + + if c.MaxRetries != nil && *c.MaxRetries < 0 { + return fmt.Errorf("max_retries must be >= 0, got %d", *c.MaxRetries) + } + + if c.TimeoutMinutes != nil && *c.TimeoutMinutes < 1 { + return fmt.Errorf("timeout_minutes must be >= 1, got %d", *c.TimeoutMinutes) + } + + if c.IdempotencyToken != nil { + v := strings.TrimSpace(*c.IdempotencyToken) + if v == "" { + return errors.New("idempotency_token cannot be empty") + } + if len(v) > 64 { + return errors.New("idempotency_token must be 64 characters or less") + } + } + + if c.MLflowRunName != nil { + v := strings.TrimSpace(*c.MLflowRunName) + if v == "" { + return errors.New("mlflow_run_name cannot be empty") + } + if len(v) > 100 { + return fmt.Errorf("mlflow_run_name must be 100 characters or less (got %d)", len(v)) + } + if !taskKeyRe.MatchString(v) { + return fmt.Errorf("invalid mlflow_run_name %q: only alphanumeric characters, hyphens, and underscores are allowed", v) + } + } + + if c.MLflowExperimentDirectory != nil { + v := strings.TrimSpace(*c.MLflowExperimentDirectory) + if v == "" { + return errors.New("mlflow_experiment_directory cannot be empty") + } + // MLflow experiments live under the workspace tree. + if !strings.HasPrefix(v, "/Workspace") { + return fmt.Errorf("mlflow_experiment_directory must start with '/Workspace', got: %s", v) + } + } + + for i := range c.Permissions { + if err := c.Permissions[i].validate(); err != nil { + return err + } + } + + // A usage policy is given by name or id, never both; the name resolves to an + // id at launch. + if c.UsagePolicyName != nil && c.UsagePolicyID != nil { + return errors.New("usage_policy_name and usage_policy_id are mutually exclusive; set only one") + } + if c.UsagePolicyName != nil { + v := strings.TrimSpace(*c.UsagePolicyName) + if v == "" { + return errors.New("usage_policy_name must not be empty") + } + // 127 matches the server-side max_length on the policy name filter. + if len(v) > 127 { + 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") + } + + return nil +} + +// validateExperimentName enforces the Databricks Jobs API task_key constraints: +// the experiment_name becomes a task key, which caps at 100 characters and allows +// only alphanumerics, hyphens, and underscores. +func validateExperimentName(v string) error { + if v == "" { + return errors.New("experiment_name cannot be empty") + } + if len(v) > 100 { + return fmt.Errorf("experiment_name must be 100 characters or less (got %d); this is the Jobs API task_key length limit", len(v)) + } + if !taskKeyRe.MatchString(v) { + return fmt.Errorf("invalid experiment_name %q: only alphanumeric characters, hyphens (-), and underscores (_) are allowed", v) + } + return nil +} + +// validateCommand enforces command is non-empty and within the line-count cap. +func validateCommand(v string) error { + if strings.TrimSpace(v) == "" { + return errors.New("command cannot be empty") + } + lineCount := strings.Count(v, "\n") + 1 + if lineCount > 1000 { + return fmt.Errorf("command is too long (%d lines); maximum is 1000 lines — move complex logic into a script in your code_source", lineCount) + } + return nil +} + +// validateSecretRefs checks that secret references use the "scope/key" format. +func validateSecretRefs(secrets map[string]string) error { + for varName, ref := range secrets { + parts := strings.Split(ref, "/") + if len(parts) != 2 { + return fmt.Errorf("invalid secret reference %q for variable %q: expected format 'scope/key' (e.g., my_scope/hf_token)", ref, varName) + } + if parts[0] == "" || parts[1] == "" { + return fmt.Errorf("invalid secret reference %q for variable %q: scope and key cannot be empty", ref, varName) + } + } + return nil +} + +// environmentConfig is the `environment` block: dependencies and/or a custom +// docker image. +type environmentConfig struct { + Dependencies dependencies `yaml:"dependencies"` + Version stringOrInt `yaml:"version"` + DockerImage *dockerImageConfig `yaml:"docker_image"` +} + +func (e *environmentConfig) validate() error { + // docker_image is exclusive with dependencies/version: the image already pins + // the full runtime. + if e.DockerImage != nil { + var conflicting []string + if e.Dependencies.set { + conflicting = append(conflicting, "dependencies") + } + if e.Version.set { + conflicting = append(conflicting, "version") + } + if len(conflicting) > 0 { + return fmt.Errorf("when 'docker_image' is specified under 'environment', these fields are not allowed: %s", strings.Join(conflicting, ", ")) + } + return e.DockerImage.validate() + } + + // version pins the client image version, which is only meaningful for an + // inline (list) dependency set — a requirements.yaml file carries its own. + if e.Version.set { + if e.Dependencies.set && !e.Dependencies.isList { + return errors.New("'environment.version' is only valid with inline dependencies (a list); when 'dependencies' points to a requirements.yaml file, set the version inside that file") + } + if !e.Dependencies.set { + return errors.New("'environment.version' requires inline 'dependencies' (a list of packages)") + } + } + + return nil +} + +// dependencies is environment.dependencies, which is polymorphic: a string is a +// path to a requirements.yaml file; a list is an inline package list. +type dependencies struct { + set bool + isList bool + path string + list []string +} + +func (d *dependencies) UnmarshalYAML(node *yaml.Node) error { + switch node.Kind { + case yaml.ScalarNode: + d.set, d.isList = true, false + return node.Decode(&d.path) + case yaml.SequenceNode: + d.set, d.isList = true, true + return node.Decode(&d.list) + default: + return errors.New("environment.dependencies must be a string path or a list of packages") + } +} + +// stringOrInt holds a scalar that may be a string or an integer in YAML +// (environment.version). The raw text is kept; integer-format validation is a +// launch-time concern. +type stringOrInt struct { + set bool + raw string +} + +func (s *stringOrInt) UnmarshalYAML(node *yaml.Node) error { + if node.Kind != yaml.ScalarNode { + return errors.New("environment.version must be a string or integer") + } + s.set = true + s.raw = node.Value + return nil +} + +// dockerImageConfig is environment.docker_image. +type dockerImageConfig struct { + URL string `yaml:"url"` +} + +func (d *dockerImageConfig) validate() error { + if strings.TrimSpace(d.URL) == "" { + return errors.New("docker_image.url cannot be empty") + } + return nil +} + +// codeSourceConfig is the `code_source` block. Only the "snapshot" type exists. +type codeSourceConfig struct { + Type string `yaml:"type"` + Snapshot *snapshotSourceConfig `yaml:"snapshot"` +} + +func (c *codeSourceConfig) validate() error { + if c.Type != "snapshot" { + return fmt.Errorf("code_source.type must be 'snapshot', got %q", c.Type) + } + if c.Snapshot == nil { + return errors.New("code_source.type='snapshot' requires a snapshot configuration") + } + return c.Snapshot.validate() +} + +// snapshotSourceConfig describes a local directory to tar and upload. +type snapshotSourceConfig struct { + RootPath string `yaml:"root_path"` + RemoteVolume *string `yaml:"remote_volume"` + Git *gitRef `yaml:"git"` + IncludePaths []string `yaml:"include_paths"` +} + +func (s *snapshotSourceConfig) validate() error { + if strings.TrimSpace(s.RootPath) == "" { + return errors.New("code_source.snapshot.root_path cannot be empty") + } + + if s.RemoteVolume != nil && !strings.HasPrefix(*s.RemoteVolume, "/Volumes/") { + return errors.New("code_source.snapshot.remote_volume must start with '/Volumes/'") + } + + // A non-nil but empty include_paths is an explicit mistake (omit it instead). + if s.IncludePaths != nil && len(s.IncludePaths) == 0 { + return errors.New("code_source.snapshot.include_paths cannot be an empty list; either omit it or provide paths") + } + for _, p := range s.IncludePaths { + p = strings.TrimSpace(p) + if p == "" { + return errors.New("code_source.snapshot.include_paths entry cannot be empty") + } + if strings.HasPrefix(p, "/") { + return fmt.Errorf("code_source.snapshot.include_paths must be relative paths, got: %s", p) + } + // No parent traversal: snapshots must stay within root_path. + if slices.Contains(strings.Split(p, "/"), "..") { + return fmt.Errorf("code_source.snapshot.include_paths cannot contain '..' traversal, got: %s", p) + } + } + + if s.Git != nil { + return s.Git.validate() + } + return nil +} + +// gitRef pins a snapshot to a specific git ref. branch and commit are mutually +// exclusive; remote is only meaningful with branch. +type gitRef struct { + Branch *string `yaml:"branch"` + Commit *string `yaml:"commit"` + Remote gitRemote `yaml:"remote"` +} + +func (g *gitRef) validate() error { + if g.Branch != nil && !gitRefRe.MatchString(*g.Branch) { + return fmt.Errorf("invalid git.branch format %q: only alphanumeric characters, hyphens, dots, slashes, and underscores are allowed", *g.Branch) + } + + // The remote-fetch path (fetching a branch's remote HEAD) is deprecated: the + // snapshot archives the local copy only. A truthy git.remote (a name or `true`) + // is rejected; `remote: false` is the default (local HEAD) and stays valid. + if g.Remote.truthy() { + return errors.New("git.remote is no longer supported: the snapshot archives your local copy, so a branch resolves to its local HEAD. To deploy a specific committed revision, use git.commit") + } + + if g.Branch == nil && g.Commit == nil { + return errors.New("git: must specify either 'branch' or 'commit'") + } + if g.Branch != nil && g.Commit != nil { + return errors.New("git: 'branch' and 'commit' are mutually exclusive — specify only one") + } + return nil +} + +// gitRemote is git.remote: false (default, use local HEAD), true (auto-detect the +// remote), or a remote name string. +type gitRemote struct { + set bool + isString bool + name string + enabled bool +} + +func (r *gitRemote) UnmarshalYAML(node *yaml.Node) error { + if node.Kind != yaml.ScalarNode { + return errors.New("git.remote must be a boolean or a remote name string") + } + r.set = true + if node.Tag == "!!bool" { + return node.Decode(&r.enabled) + } + r.isString = true + r.name = node.Value + return nil +} + +// truthy reports whether remote requests a remote fetch (mirrors Python's +// truthiness of the bool|str union). +func (r *gitRemote) truthy() bool { + if r.isString { + return r.name != "" + } + return r.enabled +} + +// permission is a DABs-compatible permission grant: exactly one principal plus a +// level. +type permission struct { + UserName *string `yaml:"user_name"` + GroupName *string `yaml:"group_name"` + ServicePrincipalName *string `yaml:"service_principal_name"` + // Level is a databricks PermissionLevel (e.g. CAN_VIEW, CAN_MANAGE). Enum + // membership is validated server-side; here we only require it to be set. + Level string `yaml:"level"` +} + +func (p *permission) validate() error { + principals := map[string]*string{ + "user_name": p.UserName, + "group_name": p.GroupName, + "service_principal_name": p.ServicePrincipalName, + } + var set []string + for name, val := range principals { + if val != nil { + set = append(set, name) + } + } + switch len(set) { + case 0: + return errors.New("permissions: one of 'user_name', 'group_name', or 'service_principal_name' must be specified") + case 1: + name := set[0] + if strings.TrimSpace(*principals[name]) == "" { + return fmt.Errorf("permissions: '%s' cannot be empty", name) + } + default: + return errors.New("permissions: only one of 'user_name', 'group_name', or 'service_principal_name' can be specified") + } + + if strings.TrimSpace(p.Level) == "" { + return errors.New("permissions: 'level' is required") + } + return nil +} diff --git a/experimental/air/cmd/runconfig_launch.go b/experimental/air/cmd/runconfig_launch.go new file mode 100644 index 00000000000..1408b600736 --- /dev/null +++ b/experimental/air/cmd/runconfig_launch.go @@ -0,0 +1,65 @@ +package aircmd + +// This file flattens the validated runConfig schema into the derived values the +// launch path consumes, replacing the Python CLI's _convert_to_run_config step. +// There is no separate internal config type: handle_run reads runConfig directly, +// using these accessors for the values that need computing rather than a plain +// field read. + +const defaultMaxRetries = 3 + +// timeoutSeconds converts timeout_minutes to seconds. Zero means the user set no +// timeout and the backend default applies. +func (c *runConfig) timeoutSeconds() int { + if c.TimeoutMinutes == nil { + return 0 + } + return *c.TimeoutMinutes * 60 +} + +// maxRetries returns the retry count, applying the schema default when unset. +func (c *runConfig) maxRetries() int { + if c.MaxRetries == nil { + return defaultMaxRetries + } + return *c.MaxRetries +} + +// dockerImageURL returns the custom docker image URL, or "" when none is set. +// +// TODO: not wired into submission yet — the native ai_runtime_task carries no +// docker field, and full support needs image registration (pending the DCS work). +func (c *runConfig) dockerImageURL() string { + if c.Environment != nil && c.Environment.DockerImage != nil { + return c.Environment.DockerImage.URL + } + return "" +} + +// requirementsFile returns the path to a requirements file when +// environment.dependencies is a string, and whether it was set. +func (c *runConfig) requirementsFile() (string, bool) { + if c.Environment == nil || !c.Environment.Dependencies.set || c.Environment.Dependencies.isList { + return "", false + } + return c.Environment.Dependencies.path, true +} + +// inlineDependencies returns the inline package list when +// environment.dependencies is a list, and whether it was set. +func (c *runConfig) inlineDependencies() ([]string, bool) { + if c.Environment == nil || !c.Environment.Dependencies.set || !c.Environment.Dependencies.isList { + return nil, false + } + return c.Environment.Dependencies.list, true +} + +// runtimeVersion returns the client image version from environment.version when +// set. For a requirements-file dependency set, the version lives in that file and +// is resolved at launch, not here. +func (c *runConfig) runtimeVersion() (string, bool) { + if c.Environment == nil || !c.Environment.Version.set { + return "", false + } + return c.Environment.Version.raw, true +} diff --git a/experimental/air/cmd/runconfig_launch_test.go b/experimental/air/cmd/runconfig_launch_test.go new file mode 100644 index 00000000000..289db91c7de --- /dev/null +++ b/experimental/air/cmd/runconfig_launch_test.go @@ -0,0 +1,80 @@ +package aircmd + +import ( + "testing" + + "github.com/stretchr/testify/assert" +) + +func TestRunConfigTimeoutSeconds(t *testing.T) { + c := &runConfig{} + assert.Equal(t, 0, c.timeoutSeconds()) + + c.TimeoutMinutes = new(2) + assert.Equal(t, 120, c.timeoutSeconds()) +} + +func TestRunConfigMaxRetries(t *testing.T) { + c := &runConfig{} + assert.Equal(t, defaultMaxRetries, c.maxRetries()) + + c.MaxRetries = new(0) + assert.Equal(t, 0, c.maxRetries()) + + c.MaxRetries = new(7) + assert.Equal(t, 7, c.maxRetries()) +} + +func TestRunConfigDockerImageURL(t *testing.T) { + c := &runConfig{} + assert.Empty(t, c.dockerImageURL()) + + c.Environment = &environmentConfig{} + assert.Empty(t, c.dockerImageURL()) + + c.Environment.DockerImage = &dockerImageConfig{URL: "org/repo:tag"} + assert.Equal(t, "org/repo:tag", c.dockerImageURL()) +} + +func TestRunConfigDependencies(t *testing.T) { + t.Run("unset", func(t *testing.T) { + c := &runConfig{} + _, ok := c.requirementsFile() + assert.False(t, ok) + _, ok = c.inlineDependencies() + assert.False(t, ok) + }) + + t.Run("file path", func(t *testing.T) { + c := &runConfig{Environment: &environmentConfig{ + Dependencies: dependencies{set: true, isList: false, path: "req.yaml"}, + }} + path, ok := c.requirementsFile() + assert.True(t, ok) + assert.Equal(t, "req.yaml", path) + _, ok = c.inlineDependencies() + assert.False(t, ok) + }) + + t.Run("inline list", func(t *testing.T) { + c := &runConfig{Environment: &environmentConfig{ + Dependencies: dependencies{set: true, isList: true, list: []string{"torch", "numpy"}}, + }} + list, ok := c.inlineDependencies() + assert.True(t, ok) + assert.Equal(t, []string{"torch", "numpy"}, list) + _, ok = c.requirementsFile() + assert.False(t, ok) + }) +} + +func TestRunConfigRuntimeVersion(t *testing.T) { + c := &runConfig{} + _, ok := c.runtimeVersion() + assert.False(t, ok) + + c.Environment = &environmentConfig{Version: stringOrInt{set: true, raw: "5"}} + v, ok := c.runtimeVersion() + assert.True(t, ok) + assert.Equal(t, "5", v) +} diff --git a/experimental/air/cmd/runconfig_load.go b/experimental/air/cmd/runconfig_load.go new file mode 100644 index 00000000000..81b07d3ca50 --- /dev/null +++ b/experimental/air/cmd/runconfig_load.go @@ -0,0 +1,52 @@ +package aircmd + +import ( + "errors" + "fmt" + "io" + "os" + + "go.yaml.in/yaml/v3" +) + +// decodeRunConfig reads and decodes the run YAML into the schema. Unknown keys +// are rejected (KnownFields), mirroring the Python schema's extra="forbid". +// +// The `_bases_` composition feature and CLI `--override` handling are not yet +// ported; a config using `_bases_` is currently rejected as an unknown field. +func decodeRunConfig(path string) (*runConfig, error) { + f, err := os.Open(path) + if err != nil { + return nil, err + } + defer f.Close() + + dec := yaml.NewDecoder(f) + dec.KnownFields(true) + + var cfg runConfig + if err := dec.Decode(&cfg); err != nil { + if errors.Is(err, io.EOF) { + return nil, fmt.Errorf("config %s is empty", path) + } + return nil, fmt.Errorf("invalid config %s: %w", path, err) + } + return &cfg, nil +} + +// validateRunConfig runs structural validation over a decoded config. +func validateRunConfig(cfg *runConfig) error { + return cfg.validate() +} + +// loadRunConfig decodes and structurally validates a run YAML config file. +func loadRunConfig(path string) (*runConfig, error) { + cfg, err := decodeRunConfig(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_test.go b/experimental/air/cmd/runconfig_test.go new file mode 100644 index 00000000000..26b54127265 --- /dev/null +++ b/experimental/air/cmd/runconfig_test.go @@ -0,0 +1,413 @@ +package aircmd + +import ( + "os" + "path/filepath" + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// writeConfig writes content to a temp YAML file and returns its path. +func writeConfig(t *testing.T, content string) string { + t.Helper() + dir := t.TempDir() + path := filepath.Join(dir, "config.yaml") + require.NoError(t, os.WriteFile(path, []byte(content), 0o600)) + return path +} + +// minimalConfig is the smallest valid config: the three required pieces. +const minimalConfig = ` +experiment_name: my-run +command: python train.py +compute: + accelerator_type: GPU_1xH100 + num_accelerators: 1 +` + +func TestLoadRunConfig_Minimal(t *testing.T) { + cfg, err := loadRunConfig(writeConfig(t, minimalConfig)) + require.NoError(t, err) + assert.Equal(t, "my-run", cfg.ExperimentName) + require.NotNil(t, cfg.Command) + assert.Equal(t, "python train.py", *cfg.Command) + require.NotNil(t, cfg.Compute) + assert.Equal(t, "GPU_1xH100", cfg.Compute.AcceleratorType) + assert.Equal(t, 1, cfg.Compute.NumAccelerators) +} + +func TestLoadRunConfig_FullFeatured(t *testing.T) { + cfg, err := loadRunConfig(writeConfig(t, ` +experiment_name: full_run +command: | + python train.py + echo done +compute: + accelerator_type: GPU_8xH100 + num_accelerators: 16 +environment: + dependencies: + - torch==2.3.0 + - numpy + version: 5 +env_variables: + FOO: bar +secrets: + HF_TOKEN: my_scope/hf_token +code_source: + type: snapshot + snapshot: + root_path: project_root/src + remote_volume: /Volumes/main/default/code + git: + branch: main + include_paths: + - src + - configs/train.yaml +max_retries: 5 +timeout_minutes: 120 +idempotency_token: abc-123 +mlflow_run_name: full_run_v2 +mlflow_experiment_directory: /Workspace/Users/me/exp +usage_policy_name: my-policy +permissions: + - group_name: users + level: CAN_VIEW + - user_name: alice@example.com + level: CAN_MANAGE +`)) + require.NoError(t, err) + assert.Equal(t, gpuType8xH100, gpuType(cfg.Compute.AcceleratorType)) + require.NotNil(t, cfg.Environment) + assert.True(t, cfg.Environment.Dependencies.isList) + assert.Equal(t, []string{"torch==2.3.0", "numpy"}, cfg.Environment.Dependencies.list) + assert.True(t, cfg.Environment.Version.set) + assert.Equal(t, "5", cfg.Environment.Version.raw) + require.NotNil(t, cfg.CodeSource) + require.NotNil(t, cfg.CodeSource.Snapshot) + require.NotNil(t, cfg.CodeSource.Snapshot.Git) + require.NotNil(t, cfg.CodeSource.Snapshot.Git.Branch) + assert.Equal(t, "main", *cfg.CodeSource.Snapshot.Git.Branch) + assert.Len(t, cfg.Permissions, 2) +} + +// TestLoadRunConfig_PolymorphicFields exercises the str|list, str|int, and +// bool|str unions decoded by custom UnmarshalYAML. +func TestLoadRunConfig_PolymorphicFields(t *testing.T) { + t.Run("dependencies as string path", func(t *testing.T) { + cfg, err := loadRunConfig(writeConfig(t, minimalConfig+` +environment: + dependencies: requirements.yaml +`)) + require.NoError(t, err) + assert.True(t, cfg.Environment.Dependencies.set) + assert.False(t, cfg.Environment.Dependencies.isList) + assert.Equal(t, "requirements.yaml", cfg.Environment.Dependencies.path) + }) + + t.Run("git remote as bool true is rejected", func(t *testing.T) { + _, err := loadRunConfig(writeConfig(t, minimalConfig+` +code_source: + type: snapshot + snapshot: + root_path: . + git: + branch: main + remote: true +`)) + require.Error(t, err) + assert.Contains(t, err.Error(), "git.remote is no longer supported") + }) + + t.Run("git remote defaults to false when unset", func(t *testing.T) { + cfg, err := loadRunConfig(writeConfig(t, minimalConfig+` +code_source: + type: snapshot + snapshot: + root_path: . + git: + commit: deadbeef +`)) + require.NoError(t, err) + assert.False(t, cfg.CodeSource.Snapshot.Git.Remote.truthy()) + }) +} + +func TestLoadRunConfig_UnknownFieldRejected(t *testing.T) { + tests := []struct { + name string + extra string + errFrag string + }{ + {"top-level typo", "extra_field: nope\n", "extra_field"}, + // priority was intentionally dropped from the schema (pool-only concept). + {"dropped priority field", "priority: 100\n", "priority"}, + // _bases_ composition is not yet ported, so it surfaces as unknown. + {"unported _bases_", "_bases_: [base.yaml]\n", "_bases_"}, + {"nested typo", "environment:\n bogus: 1\n", "bogus"}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + _, err := loadRunConfig(writeConfig(t, minimalConfig+tt.extra)) + require.Error(t, err) + assert.Contains(t, err.Error(), tt.errFrag) + }) + } +} + +func TestLoadRunConfig_Errors(t *testing.T) { + tests := []struct { + name string + yaml string + errFrag string + }{ + { + "missing experiment_name", + "command: x\ncompute:\n accelerator_type: GPU_1xH100\n num_accelerators: 1\n", + "experiment_name cannot be empty", + }, + { + "experiment_name bad chars", + "experiment_name: my.run\ncommand: x\ncompute:\n accelerator_type: GPU_1xH100\n num_accelerators: 1\n", + "invalid experiment_name", + }, + { + "missing compute", + "experiment_name: r\ncommand: x\n", + "compute: section is required", + }, + { + "missing command", + "experiment_name: r\ncompute:\n accelerator_type: GPU_1xH100\n num_accelerators: 1\n", + "command is required", + }, + { + "bad gpu type", + "experiment_name: r\ncommand: x\ncompute:\n accelerator_type: a100\n num_accelerators: 1\n", + "invalid GPU type", + }, + { + "num_accelerators not a multiple", + "experiment_name: r\ncommand: x\ncompute:\n accelerator_type: GPU_8xH100\n num_accelerators: 3\n", + "must be a multiple of 8", + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + _, err := loadRunConfig(writeConfig(t, tt.yaml)) + require.Error(t, err) + assert.Contains(t, err.Error(), tt.errFrag) + }) + } +} + +// TestRunConfigValidate_FieldRules unit-tests validation rules directly, away +// from YAML decoding, to keep each rule's failure mode explicit. +func TestRunConfigValidate_FieldRules(t *testing.T) { + str := func(s string) *string { return &s } + intp := func(i int) *int { return &i } + base := func() *runConfig { + return &runConfig{ + ExperimentName: "r", + Command: str("x"), + Compute: &computeConfig{AcceleratorType: "GPU_1xH100", NumAccelerators: 1}, + } + } + + tests := []struct { + name string + mutate func(c *runConfig) + errFrag string + }{ + {"ok baseline", func(c *runConfig) {}, ""}, + {"empty command", func(c *runConfig) { c.Command = str(" ") }, "command cannot be empty"}, + {"negative max_retries", func(c *runConfig) { c.MaxRetries = intp(-1) }, "max_retries must be >= 0"}, + {"zero timeout", func(c *runConfig) { c.TimeoutMinutes = intp(0) }, "timeout_minutes must be >= 1"}, + {"empty idempotency", func(c *runConfig) { c.IdempotencyToken = str(" ") }, "idempotency_token cannot be empty"}, + {"long idempotency", func(c *runConfig) { c.IdempotencyToken = str(string(make([]byte, 65))) }, "64 characters or less"}, + {"bad mlflow_run_name", func(c *runConfig) { c.MLflowRunName = str("bad name") }, "invalid mlflow_run_name"}, + {"bad experiment dir", func(c *runConfig) { c.MLflowExperimentDirectory = str("/Users/me") }, "must start with '/Workspace'"}, + {"empty usage policy", func(c *runConfig) { c.UsagePolicyName = str(" ") }, "usage_policy_name must not be empty"}, + {"bad secret ref", func(c *runConfig) { c.Secrets = map[string]string{"T": "noslash"} }, "expected format 'scope/key'"}, + {"empty secret scope", func(c *runConfig) { c.Secrets = map[string]string{"T": "/key"} }, "scope and key cannot be empty"}, + {"env var and secret collide", func(c *runConfig) { + c.EnvVariables = map[string]string{"TOK": "v"} + c.Secrets = map[string]string{"TOK": "scope/key"} + }, `"TOK" is set in both env_variables and secrets`}, + {"long mlflow_run_name", func(c *runConfig) { c.MLflowRunName = str(strings.Repeat("a", 101)) }, "100 characters or less"}, + {"usage policy name and id", func(c *runConfig) { + c.UsagePolicyName = str("p") + 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") }, ""}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + c := base() + tt.mutate(c) + err := c.validate() + if tt.errFrag == "" { + assert.NoError(t, err) + return + } + require.Error(t, err) + assert.Contains(t, err.Error(), tt.errFrag) + }) + } +} + +func TestEnvironmentConfigValidate(t *testing.T) { + tests := []struct { + name string + env environmentConfig + errFrag string + }{ + { + "docker image alone ok", + environmentConfig{DockerImage: &dockerImageConfig{URL: "org/repo:tag"}}, + "", + }, + { + "docker image with deps conflicts", + environmentConfig{ + DockerImage: &dockerImageConfig{URL: "org/repo:tag"}, + Dependencies: dependencies{set: true, isList: true, list: []string{"torch"}}, + }, + "not allowed: dependencies", + }, + { + "empty docker url", + environmentConfig{DockerImage: &dockerImageConfig{URL: " "}}, + "docker_image.url cannot be empty", + }, + { + "version with file deps", + environmentConfig{ + Version: stringOrInt{set: true, raw: "5"}, + Dependencies: dependencies{set: true, isList: false, path: "req.yaml"}, + }, + "only valid with inline dependencies", + }, + { + "version without deps", + environmentConfig{Version: stringOrInt{set: true, raw: "5"}}, + "requires inline 'dependencies'", + }, + { + "version with inline deps ok", + environmentConfig{ + Version: stringOrInt{set: true, raw: "5"}, + Dependencies: dependencies{set: true, isList: true, list: []string{"torch"}}, + }, + "", + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := tt.env.validate() + if tt.errFrag == "" { + assert.NoError(t, err) + return + } + require.Error(t, err) + assert.Contains(t, err.Error(), tt.errFrag) + }) + } +} + +func TestGitRefValidate(t *testing.T) { + str := func(s string) *string { return &s } + tests := []struct { + name string + ref gitRef + errFrag string + }{ + {"branch only ok", gitRef{Branch: str("main")}, ""}, + {"commit only ok", gitRef{Commit: str("abc123")}, ""}, + {"remote false is ok", gitRef{Branch: str("main"), Remote: gitRemote{set: true, enabled: false}}, ""}, + {"neither branch nor commit", gitRef{}, "must specify either 'branch' or 'commit'"}, + {"both branch and commit", gitRef{Branch: str("main"), Commit: str("abc")}, "mutually exclusive"}, + {"remote true rejected", gitRef{Branch: str("main"), Remote: gitRemote{set: true, enabled: true}}, "git.remote is no longer supported"}, + {"remote name rejected", gitRef{Branch: str("main"), Remote: gitRemote{set: true, isString: true, name: "origin"}}, "git.remote is no longer supported"}, + {"bad branch chars", gitRef{Branch: str("bad branch")}, "invalid git.branch"}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := tt.ref.validate() + if tt.errFrag == "" { + assert.NoError(t, err) + return + } + require.Error(t, err) + assert.Contains(t, err.Error(), tt.errFrag) + }) + } +} + +func TestSnapshotSourceConfigValidate(t *testing.T) { + tests := []struct { + name string + snap snapshotSourceConfig + errFrag string + }{ + {"ok", snapshotSourceConfig{RootPath: "src"}, ""}, + {"empty root_path", snapshotSourceConfig{RootPath: " "}, "root_path cannot be empty"}, + {"bad volume", snapshotSourceConfig{RootPath: "src", RemoteVolume: new("/mnt/x")}, "must start with '/Volumes/'"}, + {"empty include list", snapshotSourceConfig{RootPath: "src", IncludePaths: []string{}}, "cannot be an empty list"}, + {"absolute include", snapshotSourceConfig{RootPath: "src", IncludePaths: []string{"/etc"}}, "must be relative"}, + {"traversal include", snapshotSourceConfig{RootPath: "src", IncludePaths: []string{"../x"}}, "'..' traversal"}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := tt.snap.validate() + if tt.errFrag == "" { + assert.NoError(t, err) + return + } + require.Error(t, err) + assert.Contains(t, err.Error(), tt.errFrag) + }) + } +} + +func TestPermissionValidate(t *testing.T) { + str := func(s string) *string { return &s } + tests := []struct { + name string + perm permission + errFrag string + }{ + {"ok user", permission{UserName: str("alice@example.com"), Level: "CAN_VIEW"}, ""}, + {"no principal", permission{Level: "CAN_VIEW"}, "must be specified"}, + {"two principals", permission{UserName: str("a"), GroupName: str("g"), Level: "CAN_VIEW"}, "only one of"}, + {"empty principal", permission{UserName: str(" "), Level: "CAN_VIEW"}, "cannot be empty"}, + {"missing level", permission{GroupName: str("users")}, "'level' is required"}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := tt.perm.validate() + if tt.errFrag == "" { + assert.NoError(t, err) + return + } + require.Error(t, err) + assert.Contains(t, err.Error(), tt.errFrag) + }) + } +} + +func TestLoadRunConfig_FileErrors(t *testing.T) { + t.Run("missing file", func(t *testing.T) { + _, err := loadRunConfig(filepath.Join(t.TempDir(), "nope.yaml")) + assert.Error(t, err) + }) + t.Run("empty file", func(t *testing.T) { + _, err := loadRunConfig(writeConfig(t, "")) + require.Error(t, err) + assert.Contains(t, err.Error(), "is empty") + }) +} diff --git a/experimental/air/cmd/runlaunch.go b/experimental/air/cmd/runlaunch.go new file mode 100644 index 00000000000..b2a7215e66a --- /dev/null +++ b/experimental/air/cmd/runlaunch.go @@ -0,0 +1,73 @@ +package aircmd + +import ( + "context" + "errors" + "fmt" + "path" + "strings" + + "github.com/databricks/cli/libs/env" + "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" + "github.com/google/uuid" +) + +// userWorkspaceDirEnv overrides the per-user workspace directory; mirrors the +// Python CLI's DATABRICKS_INTERNAL_USER_WORKSPACE_DIR escape hatch. +const userWorkspaceDirEnv = "DATABRICKS_INTERNAL_USER_WORKSPACE_DIR" + +// currentUserEmail returns the authenticated user's email (works for any domain). +func currentUserEmail(ctx context.Context, w *databricks.WorkspaceClient) (string, error) { + me, err := w.CurrentUser.Me(ctx, iam.MeRequest{}) + if err != nil { + return "", fmt.Errorf("failed to resolve current user: %w", err) + } + return me.UserName, nil +} + +// userWorkspaceDir returns the user's workspace home, honoring the env override. +func userWorkspaceDir(ctx context.Context, w *databricks.WorkspaceClient) (string, error) { + if override := env.Get(ctx, userWorkspaceDirEnv); override != "" { + return override, nil + } + email, err := currentUserEmail(ctx, w) + if err != nil { + return "", err + } + return "/Workspace/Users/" + email, nil +} + +// cliLaunchDir returns a unique workspace directory for a run's launch artifacts: +// <base>/.air/cli_launch/<experiment>/<run>_<uuid>. run defaults to experiment. +func cliLaunchDir(base, experiment, run string) string { + if run == "" { + run = experiment + } + unique := strings.ReplaceAll(uuid.NewString(), "-", "")[:16] + return path.Join(base, ".air", "cli_launch", experiment, run+"_"+unique) +} + +// ensureExperimentDirectory creates experimentDir if it is missing, matching the +// CLI's convention for its other artifact directories. Without this, a missing +// parent surfaces only as a server-side INTERNAL_ERROR after the run is wasted. +// An empty dir means the default (/Users/<user>/...), which always exists. +func ensureExperimentDirectory(ctx context.Context, w *databricks.WorkspaceClient, experimentDir string) error { + if experimentDir == "" { + return nil + } + + info, err := w.Workspace.GetStatusByPath(ctx, experimentDir) + if errors.Is(err, apierr.ErrNotFound) { + return w.Workspace.MkdirsByPath(ctx, experimentDir) + } + if err != nil { + return fmt.Errorf("failed to check experiment_directory %q: %w", experimentDir, err) + } + if info.ObjectType != workspace.ObjectTypeDirectory { + return fmt.Errorf("experiment_directory %q is not a directory (object_type=%s)", experimentDir, info.ObjectType) + } + return nil +} diff --git a/experimental/air/cmd/runlaunch_test.go b/experimental/air/cmd/runlaunch_test.go new file mode 100644 index 00000000000..af6f0f70d31 --- /dev/null +++ b/experimental/air/cmd/runlaunch_test.go @@ -0,0 +1,65 @@ +package aircmd + +import ( + "strings" + "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 TestCliLaunchDir(t *testing.T) { + dir := cliLaunchDir("/Workspace/Users/me@example.com", "my-exp", "") + assert.True(t, strings.HasPrefix(dir, "/Workspace/Users/me@example.com/.air/cli_launch/my-exp/my-exp_"), dir) + // run name overrides the leaf; the unique suffix keeps successive dirs distinct. + withRun := cliLaunchDir("/base", "exp", "run1") + assert.True(t, strings.HasPrefix(withRun, "/base/.air/cli_launch/exp/run1_"), withRun) + assert.NotEqual(t, dir, cliLaunchDir("/Workspace/Users/me@example.com", "my-exp", "")) +} + +func newFakeWorkspaceClient(t *testing.T) *databricks.WorkspaceClient { + 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 +} + +func TestUserWorkspaceDir(t *testing.T) { + w := newFakeWorkspaceClient(t) + dir, err := userWorkspaceDir(t.Context(), w) + require.NoError(t, err) + assert.True(t, strings.HasPrefix(dir, "/Workspace/Users/"), dir) + + // The env override wins without an API call. + t.Setenv(userWorkspaceDirEnv, "/Workspace/custom") + dir, err = userWorkspaceDir(t.Context(), w) + require.NoError(t, err) + assert.Equal(t, "/Workspace/custom", dir) +} + +func TestEnsureExperimentDirectory(t *testing.T) { + ctx := t.Context() + w := newFakeWorkspaceClient(t) + + // Empty means default (always exists) — no API call, no error. + require.NoError(t, ensureExperimentDirectory(ctx, w, "")) + + // A missing path is created. + require.NoError(t, ensureExperimentDirectory(ctx, w, "/Workspace/Users/me/exp")) + + // An existing directory is accepted as-is. + require.NoError(t, w.Workspace.MkdirsByPath(ctx, "/Workspace/Users/me/existing")) + require.NoError(t, ensureExperimentDirectory(ctx, w, "/Workspace/Users/me/existing")) + + // A path that exists but is a file is rejected. + fc, err := filer.NewWorkspaceFilesClient(w, "/Workspace/Users/me") + require.NoError(t, err) + require.NoError(t, fc.Write(ctx, "afile", strings.NewReader("x"))) + err = ensureExperimentDirectory(ctx, w, "/Workspace/Users/me/afile") + require.ErrorContains(t, err, "is not a directory") +} diff --git a/experimental/air/cmd/runsubmit.go b/experimental/air/cmd/runsubmit.go new file mode 100644 index 00000000000..8c0be55260d --- /dev/null +++ b/experimental/air/cmd/runsubmit.go @@ -0,0 +1,191 @@ +package aircmd + +import ( + "context" + "errors" + "fmt" + "path" + "strconv" + "strings" + + "github.com/databricks/cli/libs/env" + "github.com/databricks/cli/libs/filer" + "github.com/databricks/databricks-sdk-go" + "github.com/databricks/databricks-sdk-go/service/compute" + "github.com/databricks/databricks-sdk-go/service/jobs" + "github.com/google/uuid" +) + +// dlRuntimeImageEnv overrides the default deep-learning runtime image. +const dlRuntimeImageEnv = "DATABRICKS_DL_RUNTIME_IMAGE" + +const defaultDlRuntimeImage = "CLIENT-GPU-4" + +// aiRuntimeEnvironmentKey ties the task to the serverless environment that +// carries the runtime channel. +const aiRuntimeEnvironmentKey = "default" + +// dlRuntimeImage resolves the bare runtime channel (config version, else env, +// else default), always stripping the CLIENT-GPU- prefix. +func dlRuntimeImage(ctx context.Context, runtimeVersion string) string { + img := runtimeVersion + if img == "" { + img = env.Get(ctx, dlRuntimeImageEnv) + } + if img == "" { + img = defaultDlRuntimeImage + } + return strings.TrimPrefix(img, "CLIENT-GPU-") +} + +// buildSubmitPayload assembles the runs/submit payload. commandPath is the +// workspace path of the uploaded command.sh; dlImage is the runtime channel. +// +// max_retries is always sent (including 0) so the user's YAML value is honored: +// setting it to 0 explicitly disables retries rather than falling back to the +// server default. retry_on_timeout is sent only when retries are allowed, and is +// omitempty so the wire form matches the Python CLI (which never emits a bare +// "false"). Jobs performs the retries — each attempt is a fresh AI Runtime +// workload. +func buildSubmitPayload(cfg *runConfig, commandPath, dlImage string, snap snapshotResult) jobs.SubmitRun { + task := jobs.AiRuntimeTask{ + Experiment: cfg.ExperimentName, + Deployments: []jobs.DeploymentSpec{{ + CommandPath: commandPath, + Compute: jobs.ComputeSpec{ + AcceleratorType: jobs.ComputeSpecAcceleratorType(cfg.Compute.AcceleratorType), + AcceleratorCount: cfg.Compute.NumAccelerators, + }, + }}, + CodeSourcePath: snap.CodeSourcePath, + // TEMP: git_state_path / git_diff_path are intentionally NOT sent. The typed + // jobs.AiRuntimeTask (and its source proto, ai_runtime_task.proto) has no such + // fields, so the typed SDK path cannot carry them. This is safe today because + // nothing in the backend consumes those fields — the AI Runtime task proto + // never declared them, so even the Python CLI's raw-JSON values were dropped + // on deserialization. The git_state.json / git_diff.patch sidecars are still + // uploaded next to the tarball (see snapshot.go) for human inspection. + // If the backend later adds these fields to the proto, regenerate the SDK and + // wire snap.GitStatePath / snap.GitDiffPath back in here. + } + if cfg.MLflowRunName != nil { + task.MlflowRun = *cfg.MLflowRunName + } + if cfg.MLflowExperimentDirectory != nil { + task.MlflowExperimentDirectory = *cfg.MLflowExperimentDirectory + } + + maxRetries := cfg.maxRetries() + st := jobs.SubmitTask{ + TaskKey: cfg.ExperimentName, + RunIf: jobs.RunIfAllSuccess, + AiRuntimeTask: &task, + EnvironmentKey: aiRuntimeEnvironmentKey, + MaxRetries: maxRetries, + // retry_on_timeout only makes sense when retries are allowed; otherwise + // omit it (matches Python's native path, which sets retry_on_timeout only + // under the same > 0 gate). + RetryOnTimeout: maxRetries > 0, + ForceSendFields: []string{"MaxRetries"}, + } + + return jobs.SubmitRun{ + RunName: cfg.ExperimentName, + TimeoutSeconds: cfg.timeoutSeconds(), + Tasks: []jobs.SubmitTask{st}, + Environments: []jobs.JobEnvironment{{ + EnvironmentKey: aiRuntimeEnvironmentKey, + Spec: &compute.Environment{EnvironmentVersion: dlImage}, + }}, + } +} + +// submitToken resolves the idempotency token: the --idempotency-key flag wins, +// then the config's token, else a generated one. Over-long tokens error rather +// than truncate, since truncation could make two distinct tokens collide. +func submitToken(flag string, cfg *runConfig) (string, error) { + token := flag + if token == "" && cfg.IdempotencyToken != nil { + token = *cfg.IdempotencyToken + } + if token == "" { + token = uuid.NewString() + } + if len(token) > 64 { + return "", fmt.Errorf("idempotency token must be 64 characters or less, got %d", len(token)) + } + return token, nil +} + +// submitWorkload runs the submit happy path: ensure the experiment directory, +// upload the launch artifacts, assemble the Jobs payload, and submit it. It +// returns the new run_id and its dashboard URL. +func submitWorkload(ctx context.Context, w *databricks.WorkspaceClient, cfg *runConfig, configPath, idempotencyKey string) (int64, string, error) { + // Resolving usage_policy_name to a budget policy id is not ported yet; reject + // rather than silently drop. + if cfg.UsagePolicyName != nil { + return 0, "", errors.New("usage_policy_name is not yet supported") + } + + // Resolve the idempotency token first so a bad key fails before any upload. + token, err := submitToken(idempotencyKey, cfg) + if err != nil { + return 0, "", err + } + + experimentDir := "" + if cfg.MLflowExperimentDirectory != nil { + experimentDir = *cfg.MLflowExperimentDirectory + } + if err := ensureExperimentDirectory(ctx, w, experimentDir); err != nil { + return 0, "", err + } + + base, err := userWorkspaceDir(ctx, w) + if err != nil { + return 0, "", err + } + runName := "" + if cfg.MLflowRunName != nil { + runName = *cfg.MLflowRunName + } + funcDir := cliLaunchDir(base, cfg.ExperimentName, runName) + + fc, err := filer.NewWorkspaceFilesClient(w, funcDir) + if err != nil { + return 0, "", err + } + items, err := buildArtifacts(cfg, configPath) + if err != nil { + return 0, "", err + } + if err := uploadArtifacts(ctx, fc, items); err != nil { + return 0, "", err + } + + // Package and upload the code snapshot, if any. The resulting paths ride on the + // ai_runtime_task; a run with no code_source leaves them empty. Snapshot is the + // only code_source type; guard against a nil block so snapshotCodeSource never + // dereferences a missing snapshot. + var snap snapshotResult + if cfg.CodeSource != nil && cfg.CodeSource.Snapshot != nil { + snap, err = snapshotCodeSource(ctx, w, cfg.CodeSource.Snapshot, configPath, base, funcDir) + if err != nil { + return 0, "", err + } + } + + runtimeVersion, _ := cfg.runtimeVersion() + payload := buildSubmitPayload(cfg, path.Join(funcDir, commandScriptName), dlRuntimeImage(ctx, runtimeVersion), snap) + payload.IdempotencyToken = token + + // Submit returns as soon as the run is created; we don't wait for it to finish. + wait, err := w.Jobs.Submit(ctx, payload) + if err != nil { + return 0, "", err + } + runID := wait.RunId + + dashboardURL := strings.TrimRight(w.Config.Host, "/") + "/jobs/runs/" + strconv.FormatInt(runID, 10) + return runID, dashboardURL, nil +} diff --git a/experimental/air/cmd/runsubmit_test.go b/experimental/air/cmd/runsubmit_test.go new file mode 100644 index 00000000000..fd5103599df --- /dev/null +++ b/experimental/air/cmd/runsubmit_test.go @@ -0,0 +1,216 @@ +package aircmd + +import ( + "encoding/json" + "path/filepath" + "strings" + "testing" + + "github.com/databricks/cli/libs/testserver" + "github.com/databricks/databricks-sdk-go" + "github.com/databricks/databricks-sdk-go/service/jobs" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestDlRuntimeImage(t *testing.T) { + ctx := t.Context() + // A config runtime version wins and is used bare. + assert.Equal(t, "5", dlRuntimeImage(ctx, "5")) + // The CLIENT-GPU- prefix is always stripped, even from the config version. + assert.Equal(t, "5", dlRuntimeImage(ctx, "CLIENT-GPU-5")) + // Default, with the prefix stripped. + assert.Equal(t, "4", dlRuntimeImage(ctx, "")) + // Env override, prefix stripped. + t.Setenv(dlRuntimeImageEnv, "CLIENT-GPU-7") + assert.Equal(t, "7", dlRuntimeImage(ctx, "")) +} + +func TestBuildSubmitPayload(t *testing.T) { + cfg := &runConfig{ + ExperimentName: "exp", + Command: new("python train.py"), + Compute: &computeConfig{AcceleratorType: "GPU_8xH100", NumAccelerators: 16}, + MaxRetries: new(2), + TimeoutMinutes: new(30), + MLflowRunName: new("run-v2"), + MLflowExperimentDirectory: new("/Workspace/Users/me/exp"), + } + + p := buildSubmitPayload(cfg, "/d/command.sh", "5", snapshotResult{}) + + assert.Equal(t, "exp", p.RunName) + assert.Equal(t, 1800, p.TimeoutSeconds) + require.Len(t, p.Environments, 1) + assert.Equal(t, aiRuntimeEnvironmentKey, p.Environments[0].EnvironmentKey) + require.NotNil(t, p.Environments[0].Spec) + assert.Equal(t, "5", p.Environments[0].Spec.EnvironmentVersion) + + require.Len(t, p.Tasks, 1) + task := p.Tasks[0] + assert.Equal(t, "exp", task.TaskKey) + assert.Equal(t, jobs.RunIfAllSuccess, task.RunIf) + assert.Equal(t, aiRuntimeEnvironmentKey, task.EnvironmentKey) + assert.Equal(t, 2, task.MaxRetries) + assert.True(t, task.RetryOnTimeout) + + at := task.AiRuntimeTask + require.NotNil(t, at) + assert.Equal(t, "exp", at.Experiment) + assert.Equal(t, "run-v2", at.MlflowRun) + assert.Equal(t, "/Workspace/Users/me/exp", at.MlflowExperimentDirectory) + require.Len(t, at.Deployments, 1) + assert.Equal(t, "/d/command.sh", at.Deployments[0].CommandPath) + assert.Equal(t, jobs.ComputeSpec{AcceleratorType: jobs.ComputeSpecAcceleratorTypeGpu8xH100, AcceleratorCount: 16}, at.Deployments[0].Compute) +} + +func TestBuildSubmitPayloadDefaultRetries(t *testing.T) { + // max_retries unset defaults to 3 (matching the Python native path), so both + // retry fields are sent. + cfg := &runConfig{ + ExperimentName: "exp", + Command: new("x"), + Compute: &computeConfig{AcceleratorType: "GPU_1xH100", NumAccelerators: 1}, + } + task := buildSubmitPayload(cfg, "/d/command.sh", "4", snapshotResult{}).Tasks[0] + assert.Equal(t, defaultMaxRetries, task.MaxRetries) + assert.True(t, task.RetryOnTimeout) +} + +func TestBuildSubmitPayloadNoRetries(t *testing.T) { + // max_retries: 0 must be sent explicitly so Jobs honors "no retries" instead + // of applying the server default. retry_on_timeout is omitted when retries + // aren't allowed. + cfg := &runConfig{ + ExperimentName: "exp", + Command: new("x"), + Compute: &computeConfig{AcceleratorType: "GPU_1xH100", NumAccelerators: 1}, + MaxRetries: new(0), + } + task := buildSubmitPayload(cfg, "/d/command.sh", "4", snapshotResult{}).Tasks[0] + assert.Equal(t, 0, task.MaxRetries) + assert.False(t, task.RetryOnTimeout) + + b, err := json.Marshal(task) + require.NoError(t, err) + assert.Contains(t, string(b), `"max_retries":0`) + assert.NotContains(t, string(b), "retry_on_timeout") +} + +func TestSubmitToken(t *testing.T) { + cfg := &runConfig{IdempotencyToken: new("from-config")} + + tok, err := submitToken("from-flag", cfg) // flag wins + require.NoError(t, err) + assert.Equal(t, "from-flag", tok) + + tok, err = submitToken("", cfg) // then config + require.NoError(t, err) + assert.Equal(t, "from-config", tok) + + tok, err = submitToken("", &runConfig{}) // else generated + require.NoError(t, err) + assert.NotEmpty(t, tok) + + // An over-long token errors instead of being truncated. + _, err = submitToken(strings.Repeat("a", 65), cfg) + require.ErrorContains(t, err, "64 characters or less") +} + +func TestSubmitWorkload(t *testing.T) { + server := testserver.New(t) + t.Cleanup(server.Close) + + // Register before AddDefaultHandlers: the router is first-wins, so this must claim the route ahead of the default handler. + var got jobs.SubmitRun + server.Handle("POST", "/api/2.2/jobs/runs/submit", func(req testserver.Request) any { + require.NoError(t, json.Unmarshal(req.Body, &got)) + return jobs.SubmitRunResponse{RunId: 777} + }) + testserver.AddDefaultHandlers(server) + + w, err := databricks.NewWorkspaceClient(&databricks.Config{Host: server.URL, Token: "token"}) + require.NoError(t, err) + + cfgPath := writeConfigFile(t, "run.yaml", minimalConfig) + cfg, err := loadRunConfig(cfgPath) + require.NoError(t, err) + + runID, dashboardURL, err := submitWorkload(t.Context(), w, cfg, cfgPath, "idem-key") + require.NoError(t, err) + assert.Equal(t, int64(777), runID) + assert.Contains(t, dashboardURL, "/jobs/runs/777") + + // The submitted payload is a native ai_runtime_task pointing at the uploaded + // command.sh under the run's launch directory. + assert.Equal(t, "my-run", got.RunName) + assert.Equal(t, "idem-key", got.IdempotencyToken) + require.Len(t, got.Environments, 1) + require.Len(t, got.Tasks, 1) + at := got.Tasks[0].AiRuntimeTask + require.NotNil(t, at) + require.Len(t, at.Deployments, 1) + d := at.Deployments[0] + assert.True(t, strings.HasSuffix(d.CommandPath, "/"+commandScriptName), d.CommandPath) + assert.Contains(t, d.CommandPath, "/.air/cli_launch/") + assert.Equal(t, jobs.ComputeSpec{AcceleratorType: jobs.ComputeSpecAcceleratorTypeGpu1xH100, AcceleratorCount: 1}, d.Compute) +} + +// TestSubmitWorkloadWithCodeSource exercises the snapshot path end to end: a +// git-pinned code_source is packaged, uploaded, and its paths attached to the task. +func TestSubmitWorkloadWithCodeSource(t *testing.T) { + server := testserver.New(t) + t.Cleanup(server.Close) + + // Register before AddDefaultHandlers: the router is first-wins, so this must claim the route ahead of the default handler. + var got jobs.SubmitRun + server.Handle("POST", "/api/2.2/jobs/runs/submit", func(req testserver.Request) any { + require.NoError(t, json.Unmarshal(req.Body, &got)) + return jobs.SubmitRunResponse{RunId: 555} + }) + testserver.AddDefaultHandlers(server) + w, err := databricks.NewWorkspaceClient(&databricks.Config{Host: server.URL, Token: "token"}) + require.NoError(t, err) + + // A git repo committed at HEAD, referenced by commit so packaging is git_archive. + repo := newTestRepo(t) + writeRepoFile(t, repo, "train.py", "print()") + sha := commitAll(t, repo, "init") + + cfg := minimalConfig + ` +code_source: + type: snapshot + snapshot: + root_path: ` + repo + ` + git: + commit: ` + sha + ` +` + cfgPath := writeConfigFile(t, "run.yaml", cfg) + loaded, err := loadRunConfig(cfgPath) + require.NoError(t, err) + + _, _, err = submitWorkload(t.Context(), w, loaded, cfgPath, "idem") + require.NoError(t, err) + + at := got.Tasks[0].AiRuntimeTask + // The tarball path is under the user's repo_snapshots dir. git_state_path / + // git_diff_path are not asserted: the typed jobs.AiRuntimeTask has no such fields + // (see the TEMP note in buildSubmitPayload), so they aren't sent. The git_state + // sidecar file is still uploaded next to the tarball — covered by TestRunSnapshot. + assert.Contains(t, at.CodeSourcePath, "/.air/repo_snapshots/"+filepath.Base(repo)+"/") + assert.True(t, strings.HasSuffix(at.CodeSourcePath, ".tar.gz"), at.CodeSourcePath) +} + +func TestSubmitWorkloadGuards(t *testing.T) { + w := newFakeWorkspaceClient(t) + cfgPath := writeConfigFile(t, "run.yaml", minimalConfig) + base, err := loadRunConfig(cfgPath) + require.NoError(t, err) + + t.Run("usage_policy_name rejected", func(t *testing.T) { + cfg := *base + cfg.UsagePolicyName = new("p") + _, _, err := submitWorkload(t.Context(), w, &cfg, cfgPath, "") + require.ErrorContains(t, err, "usage_policy_name is not yet supported") + }) +} diff --git a/experimental/air/cmd/runupload.go b/experimental/air/cmd/runupload.go new file mode 100644 index 00000000000..fb9ca00b987 --- /dev/null +++ b/experimental/air/cmd/runupload.go @@ -0,0 +1,170 @@ +package aircmd + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "maps" + "os" + "path/filepath" + "slices" + "strings" + + "github.com/databricks/cli/libs/filer" + "go.yaml.in/yaml/v3" +) + +// Launch artifact basenames, uploaded into the run's cli_launch directory. The +// server-side launcher derives requirements.yaml / hyperparameters.yaml from the +// same directory, so these names are part of the contract. +const ( + trainingConfigName = "training_config.yaml" + commandScriptName = "command.sh" + requirementsName = "requirements.yaml" + hyperparametersName = "hyperparameters.yaml" + envVarsName = "env_vars.json" + secretEnvVarsName = "secret_env_vars.json" +) + +// maxConfigYAMLBytes caps training_config.yaml. It is referenced by the Jobs +// payload and rendered on the run page, so an oversized parameters/command block +// is rejected here; full parameters still ship in hyperparameters.yaml. +const maxConfigYAMLBytes = 1024 * 1024 + +// uploadItem is a single artifact to write into the launch directory. +type uploadItem struct { + name string + data []byte +} + +// fileWriter is the subset of filer.Filer the upload path needs; a narrow +// interface keeps buildArtifacts/upload testable without a live workspace. +type fileWriter interface { + Write(ctx context.Context, name string, reader io.Reader, mode ...filer.WriteMode) error +} + +// requirementsDoc mirrors the on-disk requirements.yaml format so the worker +// parses synthesized inline dependencies identically to a user-provided file. +type requirementsDoc struct { + Version string `yaml:"version,omitempty"` + Dependencies []string `yaml:"dependencies"` +} + +// buildArtifacts assembles the files to upload for a run: the merged config, the +// inline command as a script, requirements (from a file or synthesized from +// inline dependencies), and hyperparameters. configPath is the local YAML path. +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. + configData, err := os.ReadFile(configPath) + if err != nil { + return nil, fmt.Errorf("failed to read config %s: %w", configPath, err) + } + if len(configData) > maxConfigYAMLBytes { + return nil, fmt.Errorf("config YAML is %.2f MB, over the %d MB limit; reduce 'parameters' or 'command'", + float64(len(configData))/(1024*1024), maxConfigYAMLBytes/(1024*1024)) + } + + items := []uploadItem{ + {trainingConfigName, configData}, + {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 { + return nil, fmt.Errorf("failed to serialize parameters: %w", err) + } + items = append(items, uploadItem{hyperparametersName, data}) + } + + // The ai_runtime_task proto carries no inline env vars or secrets; stage them + // as JSON files co-located with command.sh for the server-side launcher. + if len(cfg.EnvVariables) > 0 { + data, err := json.Marshal(envVarEntries(cfg.EnvVariables)) + if err != nil { + return nil, fmt.Errorf("failed to serialize env_variables: %w", err) + } + items = append(items, uploadItem{envVarsName, data}) + } + if len(cfg.Secrets) > 0 { + data, err := json.Marshal(secretEnvVarEntries(cfg.Secrets)) + if err != nil { + return nil, fmt.Errorf("failed to serialize secrets: %w", err) + } + items = append(items, uploadItem{secretEnvVarsName, data}) + } + + return items, nil +} + +// envVarEntry is one entry in env_vars.json. +type envVarEntry struct { + Name string `json:"name"` + Value string `json:"value"` +} + +// secretEnvVarEntry is one entry in secret_env_vars.json. The YAML side is +// {ENV_VAR: "scope/key"}; the launcher wants the split form. +type secretEnvVarEntry struct { + Name string `json:"name"` + SecretScope string `json:"secret_scope"` + SecretKey string `json:"secret_key"` +} + +// envVarEntries renders env_variables sorted by name for deterministic output. +func envVarEntries(vars map[string]string) []envVarEntry { + out := make([]envVarEntry, 0, len(vars)) + for _, name := range slices.Sorted(maps.Keys(vars)) { + out = append(out, envVarEntry{Name: name, Value: vars[name]}) + } + return out +} + +// secretEnvVarEntries renders secrets sorted by name for deterministic output. +func secretEnvVarEntries(secrets map[string]string) []secretEnvVarEntry { + out := make([]secretEnvVarEntry, 0, len(secrets)) + for _, name := range slices.Sorted(maps.Keys(secrets)) { + scope, key, _ := strings.Cut(secrets[name], "/") + out = append(out, secretEnvVarEntry{Name: name, SecretScope: scope, SecretKey: key}) + } + return out +} + +// uploadArtifacts writes each artifact into the launch directory, overwriting and +// creating parents as needed. +// +// TODO(DABs): this client-side upload could move onto libs/sync / a bundle deploy +// so the CLI reuses DABs' file-staging machinery instead of writing files itself. +func uploadArtifacts(ctx context.Context, w fileWriter, items []uploadItem) error { + for _, it := range items { + if err := w.Write(ctx, it.name, bytes.NewReader(it.data), filer.OverwriteIfExists, filer.CreateParentDirectories); err != nil { + return fmt.Errorf("failed to upload %s: %w", it.name, err) + } + } + return nil +} diff --git a/experimental/air/cmd/runupload_test.go b/experimental/air/cmd/runupload_test.go new file mode 100644 index 00000000000..0c87524735d --- /dev/null +++ b/experimental/air/cmd/runupload_test.go @@ -0,0 +1,155 @@ +package aircmd + +import ( + "context" + "errors" + "io" + "os" + "path/filepath" + "strings" + "testing" + + "github.com/databricks/cli/libs/filer" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// fakeWriter records artifact writes in place of a workspace filer. +type fakeWriter struct { + written map[string]string +} + +func (f *fakeWriter) Write(ctx context.Context, name string, reader io.Reader, mode ...filer.WriteMode) error { + if f.written == nil { + f.written = map[string]string{} + } + data, err := io.ReadAll(reader) + if err != nil { + return err + } + f.written[name] = string(data) + return nil +} + +func writeConfigFile(t *testing.T, name, content string) string { + t.Helper() + path := filepath.Join(t.TempDir(), name) + require.NoError(t, os.WriteFile(path, []byte(content), 0o600)) + return path +} + +func itemNames(items []uploadItem) []string { + names := make([]string, len(items)) + for i, it := range items { + names[i] = it.name + } + return names +} + +func TestBuildArtifacts_CommandAndConfig(t *testing.T) { + path := writeConfigFile(t, "run.yaml", minimalConfig) + cfg := &runConfig{Command: new("python train.py")} + + items, err := buildArtifacts(cfg, path) + require.NoError(t, err) + assert.Equal(t, []string{trainingConfigName, commandScriptName}, itemNames(items)) + assert.Equal(t, minimalConfig, string(items[0].data)) + assert.Equal(t, "python train.py", string(items[1].data)) +} + +func TestBuildArtifacts_InlineRequirementsAndParameters(t *testing.T) { + path := writeConfigFile(t, "run.yaml", "x: y\n") + cfg := &runConfig{ + Command: new("echo hi"), + Environment: &environmentConfig{ + Dependencies: dependencies{set: true, isList: true, list: []string{"torch", "numpy"}}, + Version: stringOrInt{set: true, raw: "5"}, + }, + Parameters: map[string]any{"lr": 0.1}, + } + + 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") +} + +func TestBuildArtifacts_EnvVarsAndSecrets(t *testing.T) { + path := writeConfigFile(t, "run.yaml", "x: y\n") + cfg := &runConfig{ + Command: new("echo hi"), + EnvVariables: map[string]string{"WANDB": "demo"}, + Secrets: map[string]string{"HF_TOKEN": "myscope/hf"}, + } + + items, err := buildArtifacts(cfg, path) + require.NoError(t, err) + assert.Subset(t, itemNames(items), []string{envVarsName, secretEnvVarsName}) + + byName := map[string][]byte{} + for _, it := range items { + byName[it.name] = it.data + } + assert.JSONEq(t, `[{"name":"WANDB","value":"demo"}]`, string(byName[envVarsName])) + assert.JSONEq(t, `[{"name":"HF_TOKEN","secret_scope":"myscope","secret_key":"hf"}]`, string(byName[secretEnvVarsName])) +} + +func TestBuildArtifacts_RequirementsFile(t *testing.T) { + dir := t.TempDir() + require.NoError(t, os.WriteFile(filepath.Join(dir, "run.yaml"), []byte("x: y\n"), 0o600)) + require.NoError(t, os.WriteFile(filepath.Join(dir, "reqs.yaml"), []byte("version: 4\n"), 0o600)) + cfg := &runConfig{ + Command: new("echo hi"), + Environment: &environmentConfig{Dependencies: dependencies{set: true, isList: false, path: "reqs.yaml"}}, + } + + items, err := buildArtifacts(cfg, filepath.Join(dir, "run.yaml")) + require.NoError(t, err) + assert.Contains(t, itemNames(items), requirementsName) +} + +func TestBuildArtifacts_OversizeConfigRejected(t *testing.T) { + path := writeConfigFile(t, "run.yaml", strings.Repeat("a", maxConfigYAMLBytes+1)) + _, err := buildArtifacts(&runConfig{Command: new("x")}, path) + require.Error(t, err) + assert.Contains(t, err.Error(), "over the 1 MB limit") +} + +func TestUploadArtifacts(t *testing.T) { + w := &fakeWriter{} + items := []uploadItem{{trainingConfigName, []byte("cfg")}, {commandScriptName, []byte("cmd")}} + require.NoError(t, uploadArtifacts(t.Context(), w, items)) + assert.Equal(t, "cfg", w.written[trainingConfigName]) + assert.Equal(t, "cmd", w.written[commandScriptName]) +} + +// errWriter fails every Write, exercising the upload error path. +type errWriter struct{} + +func (errWriter) Write(ctx context.Context, name string, reader io.Reader, mode ...filer.WriteMode) error { + return errors.New("boom") +} + +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 new file mode 100644 index 00000000000..aba67f6109c --- /dev/null +++ b/experimental/air/cmd/snapshot.go @@ -0,0 +1,279 @@ +package aircmd + +import ( + "bytes" + "context" + "errors" + "fmt" + "io/fs" + "os" + "path" + "path/filepath" + "strings" + "time" + + "github.com/databricks/cli/libs/env" + "github.com/databricks/cli/libs/filer" + "github.com/databricks/cli/libs/log" + "github.com/databricks/databricks-sdk-go" +) + +// Snapshot orchestrator: resolve → package+upload → sidecars, uploading via +// libs/filer. The Python CLI did this inline; here it's split into steps. + +// snapshotResult holds the paths wired into the submit payload: the uploaded +// tarball and the optional provenance sidecars (empty when not produced). +type snapshotResult struct { + CodeSourcePath string + GitStatePath string + GitDiffPath string +} + +// repoSnapshotsSubdir is the per-user workspace location for cached tarballs, under +// the user's home. Volume uploads use remote_volume directly instead. +const repoSnapshotsSubdir = ".air/repo_snapshots" + +// snapshotCodeSource packages and uploads the code_source snapshot, returning the +// paths to attach to the ai_runtime_task. userDir is the user's workspace home; +// funcDir is the run's launch directory (where sidecars land). +func snapshotCodeSource(ctx context.Context, w *databricks.WorkspaceClient, snap *snapshotSourceConfig, configPath, userDir, funcDir string) (snapshotResult, error) { + repoPath, err := resolveRootPath(ctx, snap.RootPath, filepath.Dir(configPath)) + if err != nil { + return snapshotResult{}, err + } + + up, err := newSnapshotUploader(w, snap, userDir, funcDir, filepath.Base(repoPath)) + if err != nil { + return snapshotResult{}, err + } + return runSnapshot(ctx, up, repoPath, snap) +} + +// resolveRootPath resolves a snapshot root_path the way the Python normalize layer +// does: expand environment variables and ~, strip a leading "project_root/" (meaning +// "relative to the YAML file"), and resolve the rest against the config's directory. +// It then confirms the path exists and is a directory. +func resolveRootPath(ctx context.Context, rawPath, configDir string) (string, error) { + expanded := os.ExpandEnv(rawPath) + if home, err := env.UserHomeDir(ctx); err == nil { + if expanded == "~" { + expanded = home + } else if rest, ok := strings.CutPrefix(expanded, "~/"); ok { + expanded = filepath.Join(home, rest) + } + } + + var resolved string + switch { + case strings.HasPrefix(expanded, "project_root/"): + resolved = filepath.Join(configDir, strings.TrimPrefix(expanded, "project_root/")) + case filepath.IsAbs(expanded): + resolved = expanded + default: + 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) + } + resolved = abs + + info, err := os.Stat(resolved) + if err != nil { + return "", fmt.Errorf("root_path does not exist: %s", resolved) + } + if !info.IsDir() { + return "", fmt.Errorf("root_path must be a directory: %s", resolved) + } + return resolved, nil +} + +// snapshotUploader splits the snapshot's two destinations: the tarball goes to a +// cache location (the user's repo_snapshots dir or a Volume), sidecars to the run's +// funcDir. tarBase/sidecarBase are the absolute roots, for reporting final paths. +type snapshotUploader struct { + tarStore filer.Filer + sidecarStore filer.Filer + tarBase string + sidecarBase string +} + +// runSnapshot resolves the packaging plan, uploads the tarball, then uploads the +// provenance sidecars. repoPath is the resolved root_path. +func runSnapshot(ctx context.Context, up snapshotUploader, repoPath string, snap *snapshotSourceConfig) (snapshotResult, error) { + git := newGitRepo(repoPath) + plan, err := resolveSnapshotPlan(ctx, git, snap.Git, snap.IncludePaths) + if err != nil { + return snapshotResult{}, err + } + + dirName := filepath.Base(repoPath) + + tarName, err := uploadTarball(ctx, up, git, plan, repoPath, dirName) + if err != nil { + return snapshotResult{}, err + } + + result := snapshotResult{CodeSourcePath: path.Join(up.tarBase, tarName)} + + // Provenance sidecars are best-effort: a git/upload hiccup here must not fail an + // otherwise-valid submission. Non-git roots have no provenance to record. + if plan.isGitRepo { + result.GitStatePath, result.GitDiffPath = uploadSidecars(ctx, up, git, plan) + } + return result, nil +} + +// uploadTarball packages the snapshot and uploads it, returning the tarball's name +// within the tar store. For git_archive it checks the cache first and skips +// packaging+upload on a hit. It writes the tarball to a temp file that is always +// cleaned up. +func uploadTarball(ctx context.Context, up snapshotUploader, git gitRepo, plan snapshotPlan, repoPath, dirName string) (string, error) { + // git_archive is cacheable by (commit, include_paths); a hit means the identical + // tarball is already uploaded, so packaging and upload are skipped entirely. + if plan.mode == modeGitArchive { + cacheKey := computeSnapshotCacheKey(plan.commitSHA, plan.includePaths) + tarName := fmt.Sprintf("%s_%s.tar.gz", dirName, cacheKey[:16]) + if exists, err := fileExists(ctx, up.tarStore, tarName); err != nil { + return "", err + } else if exists { + log.Debugf(ctx, "snapshot cache hit for %s at %s", shortSHA(plan.commitSHA), path.Join(up.tarBase, tarName)) + return tarName, nil + } + if err := packageAndUpload(ctx, up, tarName, func(out string) error { + return createGitArchiveSnapshot(ctx, git, plan.commitSHA, out, dirName, plan.includePaths) + }); err != nil { + return "", err + } + return tarName, nil + } + + // plain_tar is not cacheable (working-tree content isn't pinned to a SHA), so it + // is timestamp-named to avoid clobbering a concurrent submission. + tarName := fmt.Sprintf("%s_%s.tar.gz", dirName, time.Now().UTC().Format("20060102_150405")) + if err := packageAndUpload(ctx, up, tarName, func(out string) error { + return createPlainTarball(ctx, repoPath, out, plan.includePaths) + }); err != nil { + return "", err + } + return tarName, nil +} + +// packageAndUpload writes the tarball via pkg into a temp file, then uploads it to +// tarName in the tar store. The temp file is always removed. +func packageAndUpload(ctx context.Context, up snapshotUploader, tarName string, pkg func(outputPath string) error) error { + tmp, err := os.CreateTemp("", "air-snapshot-*.tar.gz") + if err != nil { + return fmt.Errorf("failed to create temp tarball: %w", err) + } + tmpPath := tmp.Name() + tmp.Close() + defer os.Remove(tmpPath) + + if err := pkg(tmpPath); err != nil { + return err + } + + f, err := os.Open(tmpPath) + if err != nil { + return fmt.Errorf("failed to open tarball: %w", err) + } + defer f.Close() + + if err := up.tarStore.Write(ctx, tarName, f, filer.OverwriteIfExists, filer.CreateParentDirectories); err != nil { + return fmt.Errorf("failed to upload snapshot to %s: %w", path.Join(up.tarBase, tarName), err) + } + return nil +} + +// uploadSidecars builds and uploads the git_state.json and optional git_diff.patch +// provenance sidecars into the run's funcDir. It is best-effort: any failure logs a +// warning and returns whatever paths did upload (possibly none), never an error. +func uploadSidecars(ctx context.Context, up snapshotUploader, git gitRepo, plan snapshotPlan) (statePath, diffPath string) { + mode := packagingModePlainTar + pinnedTip := "" + if plan.mode == modeGitArchive { + mode = packagingModeGitArchive + pinnedTip = plan.commitSHA + } + + sidecar, err := buildGitStateSidecar(ctx, git, mode, pinnedTip, time.Now()) + if err != nil { + log.Warnf(ctx, "skipping git provenance sidecar: %v", err) + return "", "" + } + + // Capture the dirty diff first so its status/path land in git_state.json. + if sidecar.Dirty { + status, diff := captureDirtyDiff(ctx, git, dirtyDiffSizeCapBytes, dirtyDiffTimeout) + sidecar.DiffStatus = status + if status == diffStatusCaptured { + if err := up.sidecarStore.Write(ctx, gitDiffName, bytes.NewReader(diff), filer.OverwriteIfExists, filer.CreateParentDirectories); err != nil { + log.Warnf(ctx, "failed to upload git diff sidecar: %v", err) + sidecar.DiffStatus = diffStatusClean + } else { + diffPath = path.Join(up.sidecarBase, gitDiffName) + sidecar.DiffPath = &diffPath + } + } + } + + data, err := sidecar.marshal() + if err != nil { + log.Warnf(ctx, "failed to encode git state sidecar: %v", err) + return "", diffPath + } + if err := up.sidecarStore.Write(ctx, gitStateName, bytes.NewReader(data), filer.OverwriteIfExists, filer.CreateParentDirectories); err != nil { + log.Warnf(ctx, "failed to upload git state sidecar: %v", err) + return "", diffPath + } + return path.Join(up.sidecarBase, gitStateName), diffPath +} + +// gitStateName and gitDiffName are the sidecar basenames read by the backend. +const ( + gitStateName = "git_state.json" + gitDiffName = "git_diff.patch" +) + +// fileExists reports whether name exists in the store, treating fs.ErrNotExist as +// "no". Any other error propagates. +func fileExists(ctx context.Context, store filer.Filer, name string) (bool, error) { + _, err := store.Stat(ctx, name) + if err == nil { + return true, nil + } + if errors.Is(err, fs.ErrNotExist) { + return false, nil + } + return false, fmt.Errorf("failed to check snapshot cache: %w", err) +} + +// newSnapshotUploader builds the uploader for a submission. The tarball store is a +// Volume (when remote_volume is set) or the user's repo_snapshots workspace dir; +// sidecars always go to the run's funcDir in the workspace. +func newSnapshotUploader(w *databricks.WorkspaceClient, snap *snapshotSourceConfig, userDir, funcDir, dirName string) (snapshotUploader, error) { + sidecarStore, err := filer.NewWorkspaceFilesClient(w, funcDir) + if err != nil { + return snapshotUploader{}, err + } + + if snap.RemoteVolume != nil { + tarBase := strings.TrimRight(*snap.RemoteVolume, "/") + tarStore, err := filer.NewFilesClient(w, tarBase) + if err != nil { + return snapshotUploader{}, err + } + return snapshotUploader{tarStore: tarStore, sidecarStore: sidecarStore, tarBase: tarBase, sidecarBase: funcDir}, nil + } + + tarBase := path.Join(userDir, repoSnapshotsSubdir, dirName) + tarStore, err := filer.NewWorkspaceFilesClient(w, tarBase) + if err != nil { + return snapshotUploader{}, err + } + return snapshotUploader{tarStore: tarStore, sidecarStore: sidecarStore, tarBase: tarBase, sidecarBase: funcDir}, nil +} diff --git a/experimental/air/cmd/snapshot_cachekey.go b/experimental/air/cmd/snapshot_cachekey.go new file mode 100644 index 00000000000..44c58ee903b --- /dev/null +++ b/experimental/air/cmd/snapshot_cachekey.go @@ -0,0 +1,34 @@ +package aircmd + +// This file packages a local code directory into a tarball, uploads it to the +// workspace (or a Volume), and records git provenance sidecars for cache +// invalidation — the Go port of the Python CLI's code_source snapshot path. + +import ( + "crypto/sha256" + "encoding/hex" + "slices" + "strings" +) + +// snapshotPackagingVersion is bumped when packaging logic changes in a way that invalidates existing caches +const snapshotPackagingVersion = "v1" + +// computeSnapshotCacheKey returns a stable cache key for a snapshot tarball: the +// SHA-256 digest of (commitSHA, normalized includePaths, snapshotPackagingVersion). +// Changing any input yields a different entry. +func computeSnapshotCacheKey(commitSHA string, includePaths []string) string { + var normalizedPaths string + if len(includePaths) > 0 { + trimmed := make([]string, len(includePaths)) + for i, p := range includePaths { + trimmed[i] = strings.TrimSpace(p) + } + slices.Sort(trimmed) + normalizedPaths = strings.Join(trimmed, "\n") + } + + keyMaterial := commitSHA + "\n" + normalizedPaths + "\n" + snapshotPackagingVersion + sum := sha256.Sum256([]byte(keyMaterial)) + return hex.EncodeToString(sum[:]) +} diff --git a/experimental/air/cmd/snapshot_cachekey_test.go b/experimental/air/cmd/snapshot_cachekey_test.go new file mode 100644 index 00000000000..5743217c003 --- /dev/null +++ b/experimental/air/cmd/snapshot_cachekey_test.go @@ -0,0 +1,62 @@ +package aircmd + +import ( + "encoding/json" + "os" + "path/filepath" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +type goldenCase struct { + Name string `json:"name"` + CommitSHA string `json:"commit_sha"` + IncludePaths []string `json:"include_paths"` + CacheKey string `json:"cache_key"` +} + +// TestComputeSnapshotCacheKeyGolden asserts byte-for-byte parity with golden +// fixtures across the local-only matrix (commit + include_paths permutations). +func TestComputeSnapshotCacheKeyGolden(t *testing.T) { + data, err := os.ReadFile(filepath.Join("testdata", "cache_keys.json")) + require.NoError(t, err) + + var cases []goldenCase + require.NoError(t, json.Unmarshal(data, &cases)) + require.NotEmpty(t, cases) + + for _, tc := range cases { + t.Run(tc.Name, func(t *testing.T) { + assert.Equal(t, tc.CacheKey, computeSnapshotCacheKey(tc.CommitSHA, tc.IncludePaths)) + }) + } +} + +// TestComputeSnapshotCacheKeyProperties pins the normalization behavior the golden cases +// encode, so a regression is legible without decoding hashes. +func TestComputeSnapshotCacheKeyProperties(t *testing.T) { + sha := "a3492b801c0ffee00000000000000000000dead" + + // Order-independent: sorting means unsorted input yields the sorted key. + assert.Equal(t, + computeSnapshotCacheKey(sha, []string{"a", "b", "c"}), + computeSnapshotCacheKey(sha, []string{"c", "a", "b"}), + ) + + // nil and empty include_paths are equivalent (both contribute an empty line). + assert.Equal(t, computeSnapshotCacheKey(sha, nil), computeSnapshotCacheKey(sha, []string{})) + + // Paths are trimmed before hashing. + assert.Equal(t, + computeSnapshotCacheKey(sha, []string{"research", "data"}), + computeSnapshotCacheKey(sha, []string{" research ", " data "}), + ) + + // Duplicates are NOT collapsed — they are sorted and kept, matching Python. + assert.NotEqual(t, computeSnapshotCacheKey(sha, []string{"x", "y"}), computeSnapshotCacheKey(sha, []string{"x", "x", "y"})) + + // The version constant participates: a different version is a different key. + assert.NotEqual(t, snapshotPackagingVersion, "") +} diff --git a/experimental/air/cmd/snapshot_git.go b/experimental/air/cmd/snapshot_git.go new file mode 100644 index 00000000000..616b3049f74 --- /dev/null +++ b/experimental/air/cmd/snapshot_git.go @@ -0,0 +1,309 @@ +package aircmd + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "os/exec" + "strings" + "time" +) + +// Local, no-network git introspection and the git-state provenance sidecar, ported +// from the Python CLI's cli/utils/git_state.py. The remote-fetch helpers +// (fetch_branch_sha, remote detection, partial clone) are deliberately not ported: +// the snapshot archives the local copy only, so a ref must resolve to a local commit. + +// gitRepo runs git subcommands scoped to one repository via `git -C`. Arguments are +// passed as a slice, never a shell string, so branch/commit values can't inject. +type gitRepo struct { + path string +} + +func newGitRepo(path string) gitRepo { + return gitRepo{path: path} +} + +// run executes `git <args...>` and returns stdout; a non-zero exit wraps stderr. +func (g gitRepo) run(ctx context.Context, args ...string) (string, error) { + out, err := g.runBytes(ctx, args...) + return string(out), err +} + +// runBytes is run returning raw stdout bytes, for the dirty-diff capture which needs +// exact bytes and a size measurement. +func (g gitRepo) runBytes(ctx context.Context, args ...string) ([]byte, error) { + full := append([]string{"-C", g.path}, args...) + cmd := exec.CommandContext(ctx, "git", full...) + + var stdout, stderr bytes.Buffer + cmd.Stdout = &stdout + cmd.Stderr = &stderr + if err := cmd.Run(); err != nil { + msg := strings.TrimSpace(stderr.String()) + if msg == "" { + return nil, fmt.Errorf("git %s: %w", strings.Join(args, " "), err) + } + return nil, fmt.Errorf("git %s: %w: %s", strings.Join(args, " "), err, msg) + } + return stdout.Bytes(), nil +} + +// isRepository reports whether the path is inside a git work tree. Using +// `rev-parse --is-inside-work-tree` (not a .git lookup) means a subdirectory of a +// repo counts — the common case when root_path is a subfolder of a monorepo. +func (g gitRepo) isRepository(ctx context.Context) bool { + out, err := g.run(ctx, "rev-parse", "--is-inside-work-tree") + if err != nil { + return false + } + return strings.TrimSpace(out) == "true" +} + +// headSHA returns the current HEAD commit SHA. +func (g gitRepo) headSHA(ctx context.Context) (string, error) { + out, err := g.run(ctx, "rev-parse", "HEAD") + if err != nil { + return "", err + } + return strings.TrimSpace(out), nil +} + +// hasUncommittedChanges reports whether there are staged or unstaged changes under +// the repo subtree. The `-- .` pathspec scopes the check so a subfolder snapshot +// considers only changes that could land in it. +func (g gitRepo) hasUncommittedChanges(ctx context.Context) (bool, error) { + out, err := g.run(ctx, "status", "--porcelain", "--", ".") + if err != nil { + return false, err + } + return strings.TrimSpace(out) != "", nil +} + +// hasUncommittedChangesInPaths reports whether there are uncommitted changes within +// the include paths (empty includePaths yields false). +// +// The pathspecs limit `git status` to those subtrees (it is O(working tree), slow on +// a large monorepo) and already scope the output to what could land in the snapshot. +// Unlike the Python source we don't re-parse the entries to filter by name: git +// reports a rename as `R <new>\x00<old>`, so a name-based re-filter keys off the old +// path and could miss a rename into an include path. The only caller needs the bool. +func (g gitRepo) hasUncommittedChangesInPaths(ctx context.Context, includePaths []string) (bool, error) { + var pathspecs []string + for _, p := range includePaths { + if s := strings.TrimRight(p, "/"); s != "" { + pathspecs = append(pathspecs, s) + } + } + if len(pathspecs) == 0 { + return false, nil + } + + args := append([]string{"status", "--porcelain", "--"}, pathspecs...) + out, err := g.run(ctx, args...) + if err != nil { + return false, err + } + return strings.TrimSpace(out) != "", nil +} + +// resolveLocalBranchSHA resolves a branch to its local-HEAD commit. No remote is +// contacted; the branch must exist locally. +func (g gitRepo) resolveLocalBranchSHA(ctx context.Context, branch string) (string, error) { + out, err := g.run(ctx, "rev-parse", "refs/heads/"+branch) + if err != nil { + return "", fmt.Errorf("failed to resolve local branch %q; ensure the branch exists locally and root_path is correct: %w", branch, err) + } + return strings.TrimSpace(out), nil +} + +// commitExistsLocally reports whether commitSHA is in the local object store, without +// triggering a promisor/lazy fetch. +func (g gitRepo) commitExistsLocally(ctx context.Context, commitSHA string) bool { + _, err := g.run(ctx, "cat-file", "-e", commitSHA) + return err == nil +} + +// currentBranch returns the branch name, or "" for a detached HEAD or on error. +func (g gitRepo) currentBranch(ctx context.Context) string { + out, err := g.run(ctx, "rev-parse", "--abbrev-ref", "HEAD") + if err != nil { + return "" + } + branch := strings.TrimSpace(out) + if branch == "HEAD" { + return "" + } + return branch +} + +// remoteURL returns the URL of the named remote, or "" if it has none. +func (g gitRepo) remoteURL(ctx context.Context, remoteName string) string { + out, err := g.run(ctx, "remote", "get-url", remoteName) + if err != nil { + return "" + } + return strings.TrimSpace(out) +} + +// mergeBaseWithUpstream resolves the merge-base of HEAD and a likely upstream ref, +// trying <remote>/HEAD, /main, then /master. It reads only local remote-tracking +// refs (no fetch), returning "" if none resolve. +func (g gitRepo) mergeBaseWithUpstream(ctx context.Context, remoteName string) string { + for _, ref := range []string{remoteName + "/HEAD", remoteName + "/main", remoteName + "/master"} { + out, err := g.run(ctx, "merge-base", "HEAD", ref) + if err != nil { + continue + } + if base := strings.TrimSpace(out); base != "" { + return base + } + } + return "" +} + +// validateIncludePathsExist checks that every include path exists at commitSHA. +// `git ls-tree` (without -d, so both blobs and trees count) reports an entry when the +// path exists; empty output means missing. +func (g gitRepo) validateIncludePathsExist(ctx context.Context, commitSHA string, includePaths []string) error { + var missing []string + for _, p := range includePaths { + out, err := g.run(ctx, "ls-tree", commitSHA, p) + if err != nil { + return err + } + if strings.TrimSpace(out) == "" { + missing = append(missing, p) + } + } + if len(missing) > 0 { + return fmt.Errorf("include_paths do not exist at commit %s: %s", shortSHA(commitSHA), strings.Join(missing, ", ")) + } + return nil +} + +// shortSHA abbreviates a commit SHA to 8 chars for log/error messages, tolerating +// user-supplied abbreviations shorter than that. +func shortSHA(sha string) string { + return sha[:min(len(sha), 8)] +} + +// --- git-state provenance sidecar (git_state.json + git_diff.patch) --- +// +// The backend reads git_state.json next to the tarball to tag the MLflow run with +// base/tip/dirty provenance, and logs git_diff.patch when the tree was dirty. +// Producing the sidecar is best-effort: callers warn and continue, never fail submit. + +// snapshotStateSchemaVersion is the git_state.json schema version. Bump only in +// coordination with the backend reader. +const snapshotStateSchemaVersion = 1 + +// defaultRemoteName is the remote consulted for merge-base and repo URL (local refs +// only — the remote-fetch path is gone). +const defaultRemoteName = "origin" + +// dirtyDiffSizeCapBytes caps the git_diff.patch sidecar; a larger diff records +// size_exceeded and is skipped to keep the upload small. +const dirtyDiffSizeCapBytes = 1024 * 1024 + +// dirtyDiffTimeout bounds `git diff HEAD` so provenance never delays submission. +const dirtyDiffTimeout = 5 * time.Second + +// packaging_mode values: how the uploaded tarball was produced. +const ( + packagingModeGitArchive = "git_archive" + packagingModePlainTar = "plain_tar" +) + +// diff_status values recorded in the sidecar. +const ( + diffStatusClean = "clean" + diffStatusCaptured = "captured" + diffStatusSizeExceeded = "size_exceeded" + diffStatusTimeout = "timeout" +) + +// gitStateSidecar is the git_state.json record. Field names and the null-for-absent +// encoding match the Python source, so nullable fields are *string (absent → null). +type gitStateSidecar struct { + SchemaVersion int `json:"schema_version"` + PackagingMode string `json:"packaging_mode"` + BaseCommit *string `json:"base_commit"` + TipCommit *string `json:"tip_commit"` + Branch *string `json:"branch"` + RepoURL *string `json:"repo_url"` + Dirty bool `json:"dirty"` + DiffStatus string `json:"diff_status"` + DiffPath *string `json:"diff_path"` + GeneratedAtUTC string `json:"generated_at_utc"` +} + +// nilIfEmpty maps "" to nil so an absent value serializes as JSON null. +func nilIfEmpty(s string) *string { + if s == "" { + return nil + } + return &s +} + +// buildGitStateSidecar gathers git provenance. pinnedTip overrides the HEAD-derived +// tip for git_archive (the tarball reflects that commit, not HEAD); pass "" for +// plain_tar. Metadata is best-effort — unavailable fields become null. +func buildGitStateSidecar(ctx context.Context, git gitRepo, packagingMode, pinnedTip string, now time.Time) (gitStateSidecar, error) { + tip := pinnedTip + if tip == "" { + head, err := git.headSHA(ctx) + if err != nil { + return gitStateSidecar{}, err + } + tip = head + } + + dirty, err := git.hasUncommittedChanges(ctx) + if err != nil { + return gitStateSidecar{}, err + } + + return gitStateSidecar{ + SchemaVersion: snapshotStateSchemaVersion, + PackagingMode: packagingMode, + BaseCommit: nilIfEmpty(git.mergeBaseWithUpstream(ctx, defaultRemoteName)), + TipCommit: nilIfEmpty(tip), + Branch: nilIfEmpty(git.currentBranch(ctx)), + RepoURL: nilIfEmpty(git.remoteURL(ctx, defaultRemoteName)), + Dirty: dirty, + DiffStatus: diffStatusClean, + DiffPath: nil, + GeneratedAtUTC: now.UTC().Format("2006-01-02T15:04:05.000000") + "Z", + }, nil +} + +// marshal renders the sidecar as indented JSON (matching Python's json.dump indent=2). +func (s gitStateSidecar) marshal() ([]byte, error) { + return json.MarshalIndent(s, "", " ") +} + +// captureDirtyDiff runs `git diff HEAD` over the repo subtree, returning a diff_status +// and the diff bytes (non-nil only when captured): clean (no changes or diff failed), +// captured (under the cap), size_exceeded, or timeout. +func captureDirtyDiff(ctx context.Context, git gitRepo, sizeCapBytes int, timeout time.Duration) (string, []byte) { + diffCtx, cancel := context.WithTimeout(ctx, timeout) + defer cancel() + + out, err := git.runBytes(diffCtx, "diff", "HEAD", "--", ".") + if err != nil { + if errors.Is(diffCtx.Err(), context.DeadlineExceeded) { + return diffStatusTimeout, nil + } + return diffStatusClean, nil + } + if len(out) == 0 { + return diffStatusClean, nil + } + if len(out) > sizeCapBytes { + return diffStatusSizeExceeded, nil + } + return diffStatusCaptured, out +} diff --git a/experimental/air/cmd/snapshot_git_test.go b/experimental/air/cmd/snapshot_git_test.go new file mode 100644 index 00000000000..cf1a821b1d6 --- /dev/null +++ b/experimental/air/cmd/snapshot_git_test.go @@ -0,0 +1,296 @@ +package aircmd + +import ( + "encoding/json" + "os" + "os/exec" + "path/filepath" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// newTestRepo initializes a git repo in a temp dir with a deterministic identity +// and returns its path. Tests build up real commits/branches/dirty states on top, +// mirroring the Python git_state tests (which drive real repos, not a fake). +func newTestRepo(t *testing.T) string { + t.Helper() + dir := t.TempDir() + runGit(t, dir, "init", "-q", "-b", "main") + // Deterministic identity so commits succeed in a bare CI environment. + runGit(t, dir, "config", "user.email", "test@example.test") + runGit(t, dir, "config", "user.name", "Test") + return dir +} + +// runGit runs a git command in dir and fails the test on error. +func runGit(t *testing.T, dir string, args ...string) { + t.Helper() + cmd := exec.Command("git", append([]string{"-C", dir}, args...)...) + out, err := cmd.CombinedOutput() + require.NoError(t, err, "git %v: %s", args, out) +} + +// writeRepoFile writes a file at a repo-relative path, creating parent dirs. +func writeRepoFile(t *testing.T, repo, rel, content string) { + t.Helper() + full := filepath.Join(repo, rel) + require.NoError(t, os.MkdirAll(filepath.Dir(full), 0o755)) + require.NoError(t, os.WriteFile(full, []byte(content), 0o600)) +} + +// commitAll stages everything and commits, returning the new HEAD SHA. +func commitAll(t *testing.T, repo, msg string) string { + t.Helper() + runGit(t, repo, "add", "-A") + runGit(t, repo, "commit", "-q", "-m", msg) + sha, err := newGitRepo(repo).headSHA(t.Context()) + require.NoError(t, err) + return sha +} + +func TestGitRepo_IsRepository(t *testing.T) { + ctx := t.Context() + + repo := newTestRepo(t) + assert.True(t, newGitRepo(repo).isRepository(ctx)) + + // A subdirectory of a repo is still inside the work tree. + writeRepoFile(t, repo, "sub/x.txt", "hi") + assert.True(t, newGitRepo(filepath.Join(repo, "sub")).isRepository(ctx)) + + // A plain temp dir with no repo is not. + assert.False(t, newGitRepo(t.TempDir()).isRepository(ctx)) +} + +func TestGitRepo_HeadSHA(t *testing.T) { + ctx := t.Context() + repo := newTestRepo(t) + writeRepoFile(t, repo, "a.txt", "1") + sha := commitAll(t, repo, "init") + + got, err := newGitRepo(repo).headSHA(ctx) + require.NoError(t, err) + assert.Equal(t, sha, got) + assert.Len(t, got, 40) +} + +func TestGitRepo_HasUncommittedChanges(t *testing.T) { + ctx := t.Context() + repo := newTestRepo(t) + writeRepoFile(t, repo, "a.txt", "1") + commitAll(t, repo, "init") + + // Clean tree. + dirty, err := newGitRepo(repo).hasUncommittedChanges(ctx) + require.NoError(t, err) + assert.False(t, dirty) + + // Unstaged modification. + writeRepoFile(t, repo, "a.txt", "2") + dirty, err = newGitRepo(repo).hasUncommittedChanges(ctx) + require.NoError(t, err) + assert.True(t, dirty) +} + +func TestGitRepo_HasUncommittedChangesInPaths(t *testing.T) { + ctx := t.Context() + + repo := newTestRepo(t) + writeRepoFile(t, repo, "src/model.py", "1") + writeRepoFile(t, repo, "other/x.py", "1") + commitAll(t, repo, "init") + g := newGitRepo(repo) + + // No paths: no changes, and git is never consulted. + dirty, err := g.hasUncommittedChangesInPaths(ctx, nil) + require.NoError(t, err) + assert.False(t, dirty) + + // A change outside the included paths is ignored. + writeRepoFile(t, repo, "other/x.py", "2") + dirty, err = g.hasUncommittedChangesInPaths(ctx, []string{"src", "configs"}) + require.NoError(t, err) + assert.False(t, dirty) + + // A change inside an included path is reported. + writeRepoFile(t, repo, "src/model.py", "2") + dirty, err = g.hasUncommittedChangesInPaths(ctx, []string{"src", "configs"}) + require.NoError(t, err) + assert.True(t, dirty) + + // Trailing slashes on include paths are trimmed for the pathspec. + dirty, err = g.hasUncommittedChangesInPaths(ctx, []string{"other/"}) + require.NoError(t, err) + assert.True(t, dirty) +} + +func TestGitRepo_HasUncommittedChangesInPaths_Rename(t *testing.T) { + ctx := t.Context() + repo := newTestRepo(t) + writeRepoFile(t, repo, "src/old.py", "content") + commitAll(t, repo, "init") + + // A rename within an included path counts as a change, however git classifies + // it (rename vs delete+add); we only assert the boolean. + runGit(t, repo, "mv", "src/old.py", "src/new.py") + dirty, err := newGitRepo(repo).hasUncommittedChangesInPaths(ctx, []string{"src"}) + require.NoError(t, err) + assert.True(t, dirty) +} + +func TestGitRepo_ResolveLocalBranchSHA(t *testing.T) { + ctx := t.Context() + repo := newTestRepo(t) + writeRepoFile(t, repo, "a.txt", "1") + mainSHA := commitAll(t, repo, "init") + + // A second branch at its own commit. + runGit(t, repo, "checkout", "-q", "-b", "feature") + writeRepoFile(t, repo, "b.txt", "2") + featSHA := commitAll(t, repo, "feature work") + g := newGitRepo(repo) + + got, err := g.resolveLocalBranchSHA(ctx, "main") + require.NoError(t, err) + assert.Equal(t, mainSHA, got) + + got, err = g.resolveLocalBranchSHA(ctx, "feature") + require.NoError(t, err) + assert.Equal(t, featSHA, got) + + // A branch that does not exist locally errors (no remote is contacted). + _, err = g.resolveLocalBranchSHA(ctx, "nope") + require.Error(t, err) + assert.Contains(t, err.Error(), "resolve local branch") +} + +func TestGitRepo_CommitExistsLocally(t *testing.T) { + ctx := t.Context() + repo := newTestRepo(t) + writeRepoFile(t, repo, "a.txt", "1") + sha := commitAll(t, repo, "init") + g := newGitRepo(repo) + + assert.True(t, g.commitExistsLocally(ctx, sha)) + assert.False(t, g.commitExistsLocally(ctx, "deadbeefdeadbeefdeadbeefdeadbeefdeadbeef")) +} + +func TestGitRepo_ValidateIncludePathsExist(t *testing.T) { + ctx := t.Context() + repo := newTestRepo(t) + writeRepoFile(t, repo, "src/model.py", "1") + writeRepoFile(t, repo, "configs/train.yaml", "x") + writeRepoFile(t, repo, "train.py", "print()") + sha := commitAll(t, repo, "init") + g := newGitRepo(repo) + + // Both directory and file include_paths are accepted (ls-tree without -d). + require.NoError(t, g.validateIncludePathsExist(ctx, sha, []string{"src", "configs", "train.py"})) + + err := g.validateIncludePathsExist(ctx, sha, []string{"src", "missing"}) + require.Error(t, err) + assert.Contains(t, err.Error(), "missing") + assert.Contains(t, err.Error(), sha[:8]) +} + +func TestBuildGitStateSidecar_PlainTarClean(t *testing.T) { + ctx := t.Context() + repo := newTestRepo(t) + writeRepoFile(t, repo, "a.txt", "1") + head := commitAll(t, repo, "init") + + sc, err := buildGitStateSidecar(ctx, newGitRepo(repo), packagingModePlainTar, "", fixedNow) + require.NoError(t, err) + assert.Equal(t, snapshotStateSchemaVersion, sc.SchemaVersion) + assert.Equal(t, packagingModePlainTar, sc.PackagingMode) + require.NotNil(t, sc.TipCommit) + assert.Equal(t, head, *sc.TipCommit) + assert.False(t, sc.Dirty) + assert.Equal(t, diffStatusClean, sc.DiffStatus) + assert.Nil(t, sc.DiffPath) + // No remote in a bare test repo → base_commit and repo_url are null. + assert.Nil(t, sc.BaseCommit) + assert.Nil(t, sc.RepoURL) + require.NotNil(t, sc.Branch) + assert.Equal(t, "main", *sc.Branch) + assert.Equal(t, "2026-07-10T12:00:00.000000Z", sc.GeneratedAtUTC) +} + +func TestBuildGitStateSidecar_GitArchivePinsTip(t *testing.T) { + ctx := t.Context() + repo := newTestRepo(t) + writeRepoFile(t, repo, "a.txt", "1") + first := commitAll(t, repo, "init") + // Advance HEAD; the pinned tip must win over HEAD. + writeRepoFile(t, repo, "b.txt", "2") + commitAll(t, repo, "second") + + sc, err := buildGitStateSidecar(ctx, newGitRepo(repo), packagingModeGitArchive, first, fixedNow) + require.NoError(t, err) + require.NotNil(t, sc.TipCommit) + assert.Equal(t, first, *sc.TipCommit) + assert.Equal(t, packagingModeGitArchive, sc.PackagingMode) +} + +func TestBuildGitStateSidecar_Dirty(t *testing.T) { + ctx := t.Context() + repo := newTestRepo(t) + writeRepoFile(t, repo, "a.txt", "1") + commitAll(t, repo, "init") + writeRepoFile(t, repo, "a.txt", "2") // uncommitted + + sc, err := buildGitStateSidecar(ctx, newGitRepo(repo), packagingModePlainTar, "", fixedNow) + require.NoError(t, err) + assert.True(t, sc.Dirty) +} + +func TestGitStateSidecar_MarshalNullsAbsentFields(t *testing.T) { + ctx := t.Context() + repo := newTestRepo(t) + writeRepoFile(t, repo, "a.txt", "1") + commitAll(t, repo, "init") + + sc, err := buildGitStateSidecar(ctx, newGitRepo(repo), packagingModePlainTar, "", fixedNow) + require.NoError(t, err) + data, err := sc.marshal() + require.NoError(t, err) + + // Absent fields serialize as JSON null (not "" or omitted), matching Python. + var raw map[string]any + require.NoError(t, json.Unmarshal(data, &raw)) + require.Contains(t, raw, "base_commit") + assert.Nil(t, raw["base_commit"]) + require.Contains(t, raw, "repo_url") + assert.Nil(t, raw["repo_url"]) + require.Contains(t, raw, "diff_path") + assert.Nil(t, raw["diff_path"]) + assert.EqualValues(t, 1, raw["schema_version"]) +} + +func TestCaptureDirtyDiff(t *testing.T) { + ctx := t.Context() + repo := newTestRepo(t) + writeRepoFile(t, repo, "a.txt", "one\n") + commitAll(t, repo, "init") + + // Clean tree → no diff. + status, diff := captureDirtyDiff(ctx, newGitRepo(repo), dirtyDiffSizeCapBytes, dirtyDiffTimeout) + assert.Equal(t, diffStatusClean, status) + assert.Nil(t, diff) + + // Dirty tree → captured, and the diff mentions the changed file. + writeRepoFile(t, repo, "a.txt", "two\n") + status, diff = captureDirtyDiff(ctx, newGitRepo(repo), dirtyDiffSizeCapBytes, dirtyDiffTimeout) + assert.Equal(t, diffStatusCaptured, status) + assert.Contains(t, string(diff), "a.txt") + + // A tiny size cap forces size_exceeded and drops the bytes. + status, diff = captureDirtyDiff(ctx, newGitRepo(repo), 1, dirtyDiffTimeout) + assert.Equal(t, diffStatusSizeExceeded, status) + assert.Nil(t, diff) +} + +var fixedNow = time.Date(2026, 7, 10, 12, 0, 0, 0, time.UTC) diff --git a/experimental/air/cmd/snapshot_package.go b/experimental/air/cmd/snapshot_package.go new file mode 100644 index 00000000000..672366086c9 --- /dev/null +++ b/experimental/air/cmd/snapshot_package.go @@ -0,0 +1,132 @@ +package aircmd + +import ( + "bytes" + "context" + "fmt" + "os" + "os/exec" + "path/filepath" + "strings" +) + +// Tar builders ported from cli/utils/snapshot.py. Both shell out (git archive / tar) +// for parity and to reuse git's/tar's symlink, gitignore, and AppleDouble handling. +// The tarball's top-level dir name is load-bearing — the remote entry_script extracts +// to /databricks/code_source/<dir> — so the --prefix / `-C parent dir` forms preserve it. + +// createGitArchiveSnapshot writes a gzipped tar of commitSHA to outputTarball via +// `git archive`, with every entry prefixed by directoryName/. When includePaths is +// set, only those paths are archived. +func createGitArchiveSnapshot(ctx context.Context, git gitRepo, commitSHA, outputTarball, directoryName string, includePaths []string) error { + // Single git invocation writes the gzipped tar with the desired prefix; no + // extract/repack. Provenance lives in the git_state.json sidecar, not here. + args := []string{ + "archive", + "--format=tar.gz", + "--prefix=" + directoryName + "/", + "-o", outputTarball, + commitSHA, + } + args = append(args, includePaths...) + if _, err := git.run(ctx, args...); err != nil { + return fmt.Errorf("failed to create git archive: %w", err) + } + return nil +} + +// createPlainTarball writes a gzipped tar of repoPath's working tree to +// outputTarball via `tar`. The archive preserves repoPath's directory name as the +// top-level entry. When includePaths is set, only those paths (nested under the +// directory name) are archived. .git and macOS AppleDouble files are always +// excluded; a .gitignore at repoPath is honored. +func createPlainTarball(ctx context.Context, repoPath, outputTarball string, includePaths []string) error { + dirName := filepath.Base(repoPath) + parent := filepath.Dir(repoPath) + + args := []string{"-czf", outputTarball} + + // Exclude macOS AppleDouble files: they sort before the real top-level dir and + // hijack a remote `head -1` parse. No-op on Linux. + args = append(args, "--exclude=._*") + + // Never ship .git — provenance flows via the git_state.json sidecar. + args = append(args, "--exclude=.git") + + // Honor .gitignore if present. + gitignorePath := filepath.Join(repoPath, ".gitignore") + if patterns, err := parseGitignore(gitignorePath); err == nil { + for _, p := range patterns { + if strings.Contains(p, "/") { + // Anchor path-relative patterns to the archive root so they don't + // match identically-named paths in subdirectories. + args = append(args, "--exclude="+dirName+"/"+strings.TrimPrefix(p, "/")) + } else { + args = append(args, "--exclude="+p) + } + } + } + + // Archive from the parent so the directory name is preserved; with include_paths, + // prefix each so entries nest under it (matching git archive --prefix). + args = append(args, "-C", parent) + if len(includePaths) > 0 { + for _, p := range includePaths { + args = append(args, dirName+"/"+p) + } + } else { + args = append(args, dirName) + } + + cmd := exec.CommandContext(ctx, "tar", args...) + var stderr bytes.Buffer + cmd.Stderr = &stderr + if err := cmd.Run(); err != nil { + if msg := strings.TrimSpace(stderr.String()); msg != "" { + return fmt.Errorf("failed to create plain tarball: %w: %s", err, msg) + } + return fmt.Errorf("failed to create plain tarball: %w", err) + } + return nil +} + +// parseGitignore reads a .gitignore and returns tar --exclude patterns. It mirrors +// the Python CLI's lossy normalization so plain-tar snapshots exclude the same set: +// +// - comments (#…) and blank lines are skipped; +// - negation patterns (!…) are unsupported by tar --exclude and skipped; +// - a trailing "/" (directory marker) is stripped; +// - "**" is not a path-separator-agnostic wildcard in tar, so "**/foo" → "foo" +// and "foo/**" → "foo"; a mid-path "**" has no tar equivalent and is skipped. +// +// A missing file returns (nil, error); callers treat any error as "no patterns". +func parseGitignore(path string) ([]string, error) { + data, err := os.ReadFile(path) + if err != nil { + return nil, err + } + + var patterns []string + for raw := range strings.SplitSeq(string(data), "\n") { + line := strings.TrimRight(raw, " \t\r") + if line == "" || strings.HasPrefix(line, "#") { + continue + } + if strings.HasPrefix(line, "!") { + continue + } + line = strings.TrimRight(line, "/") + if strings.Contains(line, "**") { + switch { + case strings.HasPrefix(line, "**/"): + line = line[len("**/"):] + case strings.HasSuffix(line, "/**"): + line = line[:len(line)-len("/**")] + default: + continue + } + } + patterns = append(patterns, line) + } + return patterns, nil +} diff --git a/experimental/air/cmd/snapshot_package_test.go b/experimental/air/cmd/snapshot_package_test.go new file mode 100644 index 00000000000..d895d59b98e --- /dev/null +++ b/experimental/air/cmd/snapshot_package_test.go @@ -0,0 +1,163 @@ +package aircmd + +import ( + "archive/tar" + "compress/gzip" + "os" + "path/filepath" + "slices" + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// tarballEntries returns the sorted list of entry names in a .tar.gz. +func tarballEntries(t *testing.T, path string) []string { + t.Helper() + f, err := os.Open(path) + require.NoError(t, err) + defer f.Close() + + gz, err := gzip.NewReader(f) + require.NoError(t, err) + defer gz.Close() + + var names []string + tr := tar.NewReader(gz) + for { + hdr, err := tr.Next() + if err != nil { + break + } + names = append(names, hdr.Name) + } + slices.Sort(names) + return names +} + +func TestCreateGitArchiveSnapshot(t *testing.T) { + ctx := t.Context() + repo := newTestRepo(t) + writeRepoFile(t, repo, "a.txt", "1") + writeRepoFile(t, repo, "src/model.py", "print()") + sha := commitAll(t, repo, "init") + + out := filepath.Join(t.TempDir(), "snap.tar.gz") + dirName := filepath.Base(repo) + require.NoError(t, createGitArchiveSnapshot(ctx, newGitRepo(repo), sha, out, dirName, nil)) + + entries := tarballEntries(t, out) + // Every real entry is prefixed with the directory name; the tracked files are + // present. git archive also emits a `pax_global_header` pseudo-entry carrying + // the commit SHA — it has no prefix and tar ignores it on extraction. + assert.Contains(t, entries, dirName+"/a.txt") + assert.Contains(t, entries, dirName+"/src/model.py") + for _, e := range entries { + if e == "pax_global_header" { + continue + } + assert.True(t, strings.HasPrefix(e, dirName+"/"), "entry %q lacks prefix", e) + } +} + +func TestCreateGitArchiveSnapshot_IncludePaths(t *testing.T) { + ctx := t.Context() + repo := newTestRepo(t) + writeRepoFile(t, repo, "a.txt", "1") + writeRepoFile(t, repo, "src/model.py", "print()") + sha := commitAll(t, repo, "init") + + out := filepath.Join(t.TempDir(), "snap.tar.gz") + dirName := filepath.Base(repo) + require.NoError(t, createGitArchiveSnapshot(ctx, newGitRepo(repo), sha, out, dirName, []string{"src"})) + + entries := tarballEntries(t, out) + assert.Contains(t, entries, dirName+"/src/model.py") + // a.txt is outside the include path, so it must not appear. + assert.NotContains(t, entries, dirName+"/a.txt") +} + +func TestCreatePlainTarball(t *testing.T) { + ctx := t.Context() + repo := newTestRepo(t) + writeRepoFile(t, repo, "a.txt", "1") + writeRepoFile(t, repo, "src/model.py", "print()") + commitAll(t, repo, "init") + // Uncommitted file must be included in a plain tar. + writeRepoFile(t, repo, "dirty.txt", "wip") + + out := filepath.Join(t.TempDir(), "snap.tar.gz") + require.NoError(t, createPlainTarball(ctx, repo, out, nil)) + + dirName := filepath.Base(repo) + entries := tarballEntries(t, out) + assert.Contains(t, entries, dirName+"/a.txt") + assert.Contains(t, entries, dirName+"/dirty.txt") + // .git is never shipped. + for _, e := range entries { + assert.NotContains(t, e, "/.git/") + } +} + +func TestCreatePlainTarball_HonorsGitignore(t *testing.T) { + ctx := t.Context() + repo := newTestRepo(t) + writeRepoFile(t, repo, "keep.txt", "1") + writeRepoFile(t, repo, "junk.log", "noise") + writeRepoFile(t, repo, ".gitignore", "*.log\n") + + out := filepath.Join(t.TempDir(), "snap.tar.gz") + require.NoError(t, createPlainTarball(ctx, repo, out, nil)) + + dirName := filepath.Base(repo) + entries := tarballEntries(t, out) + assert.Contains(t, entries, dirName+"/keep.txt") + assert.NotContains(t, entries, dirName+"/junk.log") +} + +func TestCreatePlainTarball_IncludePaths(t *testing.T) { + ctx := t.Context() + repo := newTestRepo(t) + writeRepoFile(t, repo, "a.txt", "1") + writeRepoFile(t, repo, "src/model.py", "print()") + + out := filepath.Join(t.TempDir(), "snap.tar.gz") + require.NoError(t, createPlainTarball(ctx, repo, out, []string{"src"})) + + dirName := filepath.Base(repo) + entries := tarballEntries(t, out) + assert.Contains(t, entries, dirName+"/src/model.py") + assert.NotContains(t, entries, dirName+"/a.txt") +} + +func TestParseGitignore(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, ".gitignore") + content := "# comment\n" + + "\n" + + "*.log\n" + + "!keep.log\n" + // negation: skipped + "build/\n" + // trailing slash stripped + "**/node_modules\n" + // **/foo -> foo + "dist/**\n" + // foo/** -> foo + "a/**/b\n" + // mid ** : skipped + "src/config\n" // path-relative kept as-is + require.NoError(t, os.WriteFile(path, []byte(content), 0o600)) + + patterns, err := parseGitignore(path) + require.NoError(t, err) + assert.Equal(t, []string{ + "*.log", + "build", + "node_modules", + "dist", + "src/config", + }, patterns) +} + +func TestParseGitignore_Missing(t *testing.T) { + _, err := parseGitignore(filepath.Join(t.TempDir(), "nope")) + require.Error(t, err) +} diff --git a/experimental/air/cmd/snapshot_resolve.go b/experimental/air/cmd/snapshot_resolve.go new file mode 100644 index 00000000000..1146b641b92 --- /dev/null +++ b/experimental/air/cmd/snapshot_resolve.go @@ -0,0 +1,120 @@ +package aircmd + +import ( + "context" + "errors" + "fmt" +) + +// This file ports the mode/ref resolution from the Python CLI's cli_entrypoint +// snapshot block (the if/elif at lines ~1541–1722), local-only. The remote-fetch +// branches are dropped: a git ref must resolve to a commit already present +// locally (git.remote is rejected at validation — see gitRef.validate). + +// snapshotMode is how the snapshot tarball is produced. +type snapshotMode int + +const ( + // modeGitArchive packages a pinned commit via `git archive`. The commit is + // deterministic, so the tarball is cacheable by (commit, include_paths). + modeGitArchive snapshotMode = iota + // modePlainTar packages the working tree (including uncommitted changes) via + // `tar`. Not cacheable — working-tree content isn't pinned to a SHA. + modePlainTar +) + +// snapshotPlan is the outcome of resolving how to package a snapshot: the mode, +// the commit SHA to archive (git_archive only; empty for plain_tar), and whether +// the working tree under the snapshot root has uncommitted changes. +type snapshotPlan struct { + mode snapshotMode + commitSHA string + hasUncommit bool + isGitRepo bool + includePaths []string +} + +// resolveSnapshotPlan decides how to package the snapshot (local-only): +// - git.commit → pin the SHA (must exist locally) → git_archive. +// - git.branch → the branch's local HEAD SHA → git_archive. +// - no ref / non-git dir → the working tree → plain_tar (no caching). +// +// The dirty check runs at most once (git status is O(working tree)) and is threaded +// into the plan. Dirty + git.branch is an error: the committed HEAD wouldn't include +// the uncommitted changes. +func resolveSnapshotPlan(ctx context.Context, git gitRepo, ref *gitRef, includePaths []string) (snapshotPlan, error) { + plan := snapshotPlan{includePaths: includePaths} + plan.isGitRepo = git.isRepository(ctx) + + // Detect uncommitted changes once. When include_paths is set, only changes + // under those paths can land in the snapshot, so scope the check to them — + // both more correct and cheaper than scanning the whole repo. + if plan.isGitRepo { + var err error + if len(includePaths) > 0 { + plan.hasUncommit, err = git.hasUncommittedChangesInPaths(ctx, includePaths) + } else { + plan.hasUncommit, err = git.hasUncommittedChanges(ctx) + } + if err != nil { + return snapshotPlan{}, err + } + } + + // Non-git directory: plain tar, no ref allowed. gitRef.validate already rejects + // git.* on a non-git dir at load time, but guard here too since this function + // is the single decision point. + if !plan.isGitRepo { + if ref != nil { + return snapshotPlan{}, fmt.Errorf("git.* is set but %s is not a git repository", git.path) + } + plan.mode = modePlainTar + return plan, nil + } + + // git repo, no ref: package the working tree as plain tar (uncommitted changes + // included). Provenance is captured separately via the git_state sidecar. + if ref == nil { + plan.mode = modePlainTar + return plan, nil + } + + switch { + case ref.Commit != nil: + // git.commit pins a committed SHA; local uncommitted changes are irrelevant + // and won't be included. The commit must exist locally — no remote fetch. + commit := *ref.Commit + if !git.commitExistsLocally(ctx, commit) { + return snapshotPlan{}, fmt.Errorf("commit %q does not exist locally; fetch it (e.g. `git fetch`) before submitting — the snapshot archives your local copy and does not fetch from a remote", commit) + } + plan.mode = modeGitArchive + plan.commitSHA = commit + + case ref.Branch != nil: + // git.branch deploys the branch's local HEAD. A dirty tree here is an error: + // the committed HEAD wouldn't include the uncommitted changes. + if plan.hasUncommit { + return snapshotPlan{}, fmt.Errorf("uncommitted changes under %s would not be included: git.branch deploys the committed HEAD of %q. Commit your changes, or use git.commit to pin a specific revision", git.path, *ref.Branch) + } + sha, err := git.resolveLocalBranchSHA(ctx, *ref.Branch) + if err != nil { + return snapshotPlan{}, err + } + plan.mode = modeGitArchive + plan.commitSHA = sha + + default: + // gitRef.validate guarantees exactly one of branch/commit is set. + return snapshotPlan{}, errors.New("git: must specify either 'branch' or 'commit'") + } + + // For git_archive with include_paths, verify each path exists at the resolved + // commit so a typo fails fast rather than producing an empty subtree. + if len(includePaths) > 0 { + if err := git.validateIncludePathsExist(ctx, plan.commitSHA, includePaths); err != nil { + return snapshotPlan{}, err + } + } + + return plan, nil +} diff --git a/experimental/air/cmd/snapshot_resolve_test.go b/experimental/air/cmd/snapshot_resolve_test.go new file mode 100644 index 00000000000..c8c946f8394 --- /dev/null +++ b/experimental/air/cmd/snapshot_resolve_test.go @@ -0,0 +1,114 @@ +package aircmd + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestResolveSnapshotPlan_Commit(t *testing.T) { + ctx := t.Context() + repo := newTestRepo(t) + writeRepoFile(t, repo, "a.txt", "1") + sha := commitAll(t, repo, "init") + + plan, err := resolveSnapshotPlan(ctx, newGitRepo(repo), &gitRef{Commit: new(sha)}, nil) + require.NoError(t, err) + assert.Equal(t, modeGitArchive, plan.mode) + assert.Equal(t, sha, plan.commitSHA) + assert.True(t, plan.isGitRepo) + + // A commit pin is valid even with a dirty tree: local changes are irrelevant. + writeRepoFile(t, repo, "a.txt", "2") + plan, err = resolveSnapshotPlan(ctx, newGitRepo(repo), &gitRef{Commit: new(sha)}, nil) + require.NoError(t, err) + assert.Equal(t, modeGitArchive, plan.mode) + assert.True(t, plan.hasUncommit) +} + +func TestResolveSnapshotPlan_CommitNotLocal(t *testing.T) { + ctx := t.Context() + repo := newTestRepo(t) + writeRepoFile(t, repo, "a.txt", "1") + commitAll(t, repo, "init") + + absent := "deadbeefdeadbeefdeadbeefdeadbeefdeadbeef" + _, err := resolveSnapshotPlan(ctx, newGitRepo(repo), &gitRef{Commit: new(absent)}, nil) + require.Error(t, err) + assert.Contains(t, err.Error(), "does not exist locally") +} + +func TestResolveSnapshotPlan_BranchLocalHead(t *testing.T) { + ctx := t.Context() + repo := newTestRepo(t) + writeRepoFile(t, repo, "a.txt", "1") + mainSHA := commitAll(t, repo, "init") + + plan, err := resolveSnapshotPlan(ctx, newGitRepo(repo), &gitRef{Branch: new("main")}, nil) + require.NoError(t, err) + assert.Equal(t, modeGitArchive, plan.mode) + assert.Equal(t, mainSHA, plan.commitSHA) +} + +func TestResolveSnapshotPlan_BranchDirtyIsError(t *testing.T) { + ctx := t.Context() + repo := newTestRepo(t) + writeRepoFile(t, repo, "a.txt", "1") + commitAll(t, repo, "init") + writeRepoFile(t, repo, "a.txt", "2") // uncommitted + + _, err := resolveSnapshotPlan(ctx, newGitRepo(repo), &gitRef{Branch: new("main")}, nil) + require.Error(t, err) + assert.Contains(t, err.Error(), "uncommitted changes") + assert.Contains(t, err.Error(), "git.commit") +} + +func TestResolveSnapshotPlan_NoRefPlainTar(t *testing.T) { + ctx := t.Context() + repo := newTestRepo(t) + writeRepoFile(t, repo, "a.txt", "1") + commitAll(t, repo, "init") + writeRepoFile(t, repo, "a.txt", "2") // dirty tree is fine for plain tar + + plan, err := resolveSnapshotPlan(ctx, newGitRepo(repo), nil, nil) + require.NoError(t, err) + assert.Equal(t, modePlainTar, plan.mode) + assert.Empty(t, plan.commitSHA) + assert.True(t, plan.isGitRepo) + assert.True(t, plan.hasUncommit) +} + +func TestResolveSnapshotPlan_NonGitDir(t *testing.T) { + ctx := t.Context() + dir := t.TempDir() + + plan, err := resolveSnapshotPlan(ctx, newGitRepo(dir), nil, nil) + require.NoError(t, err) + assert.Equal(t, modePlainTar, plan.mode) + assert.False(t, plan.isGitRepo) + + // A git ref on a non-git directory is an error. + _, err = resolveSnapshotPlan(ctx, newGitRepo(dir), &gitRef{Branch: new("main")}, nil) + require.Error(t, err) + assert.Contains(t, err.Error(), "not a git repository") +} + +func TestResolveSnapshotPlan_IncludePaths(t *testing.T) { + ctx := t.Context() + repo := newTestRepo(t) + writeRepoFile(t, repo, "src/model.py", "1") + writeRepoFile(t, repo, "configs/train.yaml", "x") + sha := commitAll(t, repo, "init") + + // All include paths exist at the commit. + plan, err := resolveSnapshotPlan(ctx, newGitRepo(repo), &gitRef{Commit: new(sha)}, []string{"src", "configs"}) + require.NoError(t, err) + assert.Equal(t, modeGitArchive, plan.mode) + assert.Equal(t, []string{"src", "configs"}, plan.includePaths) + + // A missing include path fails fast. + _, err = resolveSnapshotPlan(ctx, newGitRepo(repo), &gitRef{Commit: new(sha)}, []string{"src", "missing"}) + require.Error(t, err) + assert.Contains(t, err.Error(), "missing") +} diff --git a/experimental/air/cmd/snapshot_test.go b/experimental/air/cmd/snapshot_test.go new file mode 100644 index 00000000000..d94fe005fc9 --- /dev/null +++ b/experimental/air/cmd/snapshot_test.go @@ -0,0 +1,155 @@ +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 ._<key>.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 new file mode 100644 index 00000000000..e28d7f66730 --- /dev/null +++ b/experimental/air/cmd/stubs_test.go @@ -0,0 +1,27 @@ +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/sweep.go b/experimental/air/cmd/sweep.go new file mode 100644 index 00000000000..b346f43f1b6 --- /dev/null +++ b/experimental/air/cmd/sweep.go @@ -0,0 +1,76 @@ +package aircmd + +import ( + "context" + "strconv" + + "github.com/databricks/cli/libs/log" + "github.com/databricks/databricks-sdk-go" + "github.com/databricks/databricks-sdk-go/service/jobs" +) + +// sweepInfo summarizes a "foreach" run, which fans a single config out into many +// iterations (a hyperparameter sweep). It is shown only in text output. +type sweepInfo struct { + Total int + Succeeded int + Failed int + Active int + Completed int + Tasks []sweepTask +} + +// sweepTask is one iteration of a sweep. +type sweepTask struct { + TaskKey string + RunID string + Status string + Experiment string +} + +// findForEachTask returns the run's foreach task if it has one, or nil. A run is +// a sweep when one of its tasks fans out into iterations. +func findForEachTask(run *jobs.Run) *jobs.RunTask { + for i := range run.Tasks { + if run.Tasks[i].ForEachTask != nil { + return &run.Tasks[i] + } + } + return nil +} + +// buildSweepInfo gathers the iteration counts and per-iteration rows for a +// sweep. The counts come from the task we already have; the individual +// iterations require a second lookup. If that lookup fails we still return the +// counts (logging the failure) so the user sees the summary. +func buildSweepInfo(ctx context.Context, w *databricks.WorkspaceClient, task *jobs.RunTask) *sweepInfo { + info := &sweepInfo{} + if task.ForEachTask.Stats != nil && task.ForEachTask.Stats.TaskRunStats != nil { + stats := task.ForEachTask.Stats.TaskRunStats + info.Total = stats.TotalIterations + info.Succeeded = stats.SucceededIterations + info.Failed = stats.FailedIterations + info.Active = stats.ActiveIterations + info.Completed = stats.CompletedIterations + } + + // The iterations are returned as part of a run lookup on the foreach task. + iterated, err := w.Jobs.GetRun(ctx, jobs.GetRunRequest{RunId: task.RunId}) + if err != nil { + log.Debugf(ctx, "air get: could not fetch sweep iterations: %v", err) + return info + } + + for _, it := range iterated.Iterations { + row := sweepTask{ + TaskKey: it.TaskKey, + RunID: strconv.FormatInt(it.RunId, 10), + Status: runStatus(it.State), + } + if it.GenAiComputeTask != nil && it.GenAiComputeTask.MlflowExperimentName != "" { + row.Experiment = stripExperimentUserPrefix(it.GenAiComputeTask.MlflowExperimentName) + } + info.Tasks = append(info.Tasks, row) + } + return info +} diff --git a/experimental/air/cmd/sweep_test.go b/experimental/air/cmd/sweep_test.go new file mode 100644 index 00000000000..10134c0df42 --- /dev/null +++ b/experimental/air/cmd/sweep_test.go @@ -0,0 +1,81 @@ +package aircmd + +import ( + "testing" + + "github.com/databricks/databricks-sdk-go/apierr" + "github.com/databricks/databricks-sdk-go/experimental/mocks" + "github.com/databricks/databricks-sdk-go/service/jobs" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/mock" + "github.com/stretchr/testify/require" +) + +func TestFindForEachTask(t *testing.T) { + // No tasks at all. + assert.Nil(t, findForEachTask(&jobs.Run{})) + + // A task that is not a foreach. + assert.Nil(t, findForEachTask(&jobs.Run{Tasks: []jobs.RunTask{{TaskKey: "a"}}})) + + // The foreach task is found even when it isn't first. + run := &jobs.Run{Tasks: []jobs.RunTask{ + {TaskKey: "a"}, + {TaskKey: "sweep", ForEachTask: &jobs.RunForEachTask{}}, + }} + got := findForEachTask(run) + require.NotNil(t, got) + assert.Equal(t, "sweep", got.TaskKey) +} + +func sweepTaskFixture() *jobs.RunTask { + return &jobs.RunTask{ + RunId: 99, + ForEachTask: &jobs.RunForEachTask{ + Stats: &jobs.ForEachStats{TaskRunStats: &jobs.ForEachTaskTaskRunStats{ + TotalIterations: 4, + SucceededIterations: 1, + FailedIterations: 1, + ActiveIterations: 2, + CompletedIterations: 2, + }}, + }, + } +} + +func TestBuildSweepInfo(t *testing.T) { + ctx := t.Context() + + t.Run("counts and iteration rows", func(t *testing.T) { + m := mocks.NewMockWorkspaceClient(t) + m.GetMockJobsAPI().EXPECT().GetRun(mock.Anything, jobs.GetRunRequest{RunId: 99}).Return( + &jobs.Run{Iterations: []jobs.RunTask{{ + TaskKey: "iter_0", + RunId: 100, + State: &jobs.RunState{ResultState: jobs.RunResultStateSuccess}, + GenAiComputeTask: &jobs.GenAiComputeTask{MlflowExperimentName: "/Users/me@example.com/exp"}, + }}}, nil) + + info := buildSweepInfo(ctx, m.WorkspaceClient, sweepTaskFixture()) + assert.Equal(t, 4, info.Total) + assert.Equal(t, 2, info.Completed) + assert.Equal(t, 1, info.Succeeded) + assert.Equal(t, 1, info.Failed) + assert.Equal(t, 2, info.Active) + require.Len(t, info.Tasks, 1) + assert.Equal(t, "iter_0", info.Tasks[0].TaskKey) + assert.Equal(t, "100", info.Tasks[0].RunID) + assert.Equal(t, "SUCCESS", info.Tasks[0].Status) + assert.Equal(t, "exp", info.Tasks[0].Experiment) + }) + + t.Run("iteration lookup failure still returns counts", func(t *testing.T) { + m := mocks.NewMockWorkspaceClient(t) + m.GetMockJobsAPI().EXPECT().GetRun(mock.Anything, jobs.GetRunRequest{RunId: 99}).Return( + nil, apierr.ErrResourceDoesNotExist) + + info := buildSweepInfo(ctx, m.WorkspaceClient, sweepTaskFixture()) + assert.Equal(t, 4, info.Total) + assert.Empty(t, info.Tasks) + }) +} diff --git a/experimental/air/cmd/testdata/cache_keys.json b/experimental/air/cmd/testdata/cache_keys.json new file mode 100644 index 00000000000..06673499109 --- /dev/null +++ b/experimental/air/cmd/testdata/cache_keys.json @@ -0,0 +1,77 @@ +[ + { + "name": "commit_no_paths", + "commit_sha": "a3492b801c0ffee00000000000000000000dead", + "include_paths": null, + "cache_key": "8d7bc445ac83dfe432a353718ae8b1ae40eb00c860352662e6df96b7fdf862a5" + }, + { + "name": "commit_empty_paths_none", + "commit_sha": "0000000000000000000000000000000000000000", + "include_paths": null, + "cache_key": "a07e0fa4d38f3c2d08c3ff6bbff0158cea01796aaaf5368b0d0e247cebd27c80" + }, + { + "name": "single_path", + "commit_sha": "a3492b801c0ffee00000000000000000000dead", + "include_paths": [ + "research" + ], + "cache_key": "97bd562ad591af23b9a05651d88c01abbbcc24632e9a8f2534230a61492b9e0a" + }, + { + "name": "multi_path_sorted", + "commit_sha": "a3492b801c0ffee00000000000000000000dead", + "include_paths": [ + "a", + "b", + "c" + ], + "cache_key": "17d021577e3615e9a603a2f701f918d74ad16218efc4b2950ee20f6a18b1a174" + }, + { + "name": "multi_path_unsorted", + "commit_sha": "a3492b801c0ffee00000000000000000000dead", + "include_paths": [ + "c", + "a", + "b" + ], + "cache_key": "17d021577e3615e9a603a2f701f918d74ad16218efc4b2950ee20f6a18b1a174" + }, + { + "name": "path_whitespace", + "commit_sha": "a3492b801c0ffee00000000000000000000dead", + "include_paths": [ + " research ", + " data " + ], + "cache_key": "c6a6fa1d866ac90a5f371610ee08efa8aeb061a43b8a4acae30765db21f4155b" + }, + { + "name": "nested_paths", + "commit_sha": "deadbeefdeadbeefdeadbeefdeadbeefdeadbeef", + "include_paths": [ + "src/models", + "src/data", + "README.md" + ], + "cache_key": "2b18452b1019186a8e5b34af540a0545f96f24bc61771f58fa71fc218cb0275c" + }, + { + "name": "short_sha", + "commit_sha": "a3492b8", + "include_paths": null, + "cache_key": "16be4d7d83ecd01b8d3d274dafb07b3bccc72cf7f7299578e634cd01e1825345" + }, + { + "name": "dup_paths", + "commit_sha": "a3492b801c0ffee00000000000000000000dead", + "include_paths": [ + "x", + "x", + "y" + ], + "cache_key": "407a5d8c3a0e9438d5abf653c4d3a43928c300b5630e2f83f435661b9fbd599e" + } +] diff --git a/libs/cmdio/io.go b/libs/cmdio/io.go index e57c90974b4..2b9b395ce13 100644 --- a/libs/cmdio/io.go +++ b/libs/cmdio/io.go @@ -83,6 +83,15 @@ func IsPromptSupported(ctx context.Context) bool { return c.capabilities.SupportsPrompt() } +// IsPagerSupported reports whether stdin, stdout, and stderr are all interactive +// terminals. This is the requirement for a full-screen or navigable output +// program: unlike IsPromptSupported it also checks stdout, so it returns false +// when stdout is piped or redirected. +func IsPagerSupported(ctx context.Context) bool { + c := fromContext(ctx) + return c.capabilities.SupportsPager() +} + // SupportsColor returns true if the given writer supports colored output. // This checks both TTY status and environment variables (NO_COLOR, TERM=dumb). func SupportsColor(ctx context.Context, w io.Writer) bool {