diff --git a/NOTICE b/NOTICE index 922dd1fd79a..dd1e4c0c1b6 100644 --- a/NOTICE +++ b/NOTICE @@ -155,6 +155,10 @@ mattn/go-isatty - https://github.com/mattn/go-isatty Copyright (c) Yasuhiro MATSUMOTO License - https://github.com/mattn/go-isatty/blob/master/LICENSE +muesli/termenv - https://github.com/muesli/termenv +Copyright (c) 2019 Christian Muehlhaeuser +License - https://github.com/muesli/termenv/blob/master/LICENSE + sabhiram/go-gitignore - https://github.com/sabhiram/go-gitignore Copyright (c) 2015 Shaba Abhiram License - https://github.com/sabhiram/go-gitignore/blob/master/LICENSE 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..ad0c25bb7f2 --- /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 (runs/get-output is not modeled by the typed SDK). +[[Server]] +Pattern = "GET /api/2.2/jobs/runs/get-output" +Response.Body = ''' +{"gen_ai_compute_output": {"run_info": {"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..606d34ec970 --- /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 (runs/get-output is not modeled by the typed SDK). +[[Server]] +Pattern = "GET /api/2.2/jobs/runs/get-output" +Response.Body = ''' +{"gen_ai_compute_output": {"run_info": {"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/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..180886290ba --- /dev/null +++ b/acceptance/experimental/air/run/output.txt @@ -0,0 +1,39 @@ + +=== 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 + +=== 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..806bd321e6d --- /dev/null +++ b/acceptance/experimental/air/run/script @@ -0,0 +1,17 @@ +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 "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/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..478f3fbb50a --- /dev/null +++ b/experimental/air/cmd/aitraining.go @@ -0,0 +1,109 @@ +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. +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..c676567003c --- /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, "", + airJobRun(111, "me@example.com", "GPU_1xA10", 1, "/Users/me@example.com/exp-a"), + airJobRun(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, "", + airJobRun(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, "", + airJobRun(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 == jobsRunsListPath { + 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..a6b47c977f7 --- /dev/null +++ b/experimental/air/cmd/get.go @@ -0,0 +1,231 @@ +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) + } + + // Fetch the run once, in both the typed and raw shapes: the typed jobs.Run + // drives the display path, and the raw jobRun preserves the ai_runtime_task + // the typed model drops (used below without a second roundtrip). + run, rawRun, err := fetchRun(ctx, w, 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 { + // The typed SDK drops ai_runtime_task, so read it from the raw run we + // already fetched above. + enrichFromRawRun(rawRun, &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 +} + +// enrichFromRawRun fills the config path, experiment, and accelerators from the +// raw run (the ai_runtime_task the typed model drops). Best-effort: empty fields +// leave the existing "N/A" fallbacks in place. +func enrichFromRawRun(raw *jobRun, data *getData) { + if cmdPath := raw.commandPath(); cmdPath != "" { + data.TrainingConfigPath = path.Join(path.Dir(cmdPath), trainingConfigName) + } + if exp := jobExperiment(raw); exp != "" { + data.ExperimentName = &exp + data.ExperimentDisplay = exp + } + if a := acceleratorLabel(jobCompute(raw)); a != "" { + data.AcceleratorsDisplay = a + } +} + +// 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..88a7bdfc4bd --- /dev/null +++ b/experimental/air/cmd/get_test.go @@ -0,0 +1,289 @@ +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 == jobsRunsGetPath { + 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 TestEnrichFromRawRun(t *testing.T) { + rawRun := func(t *testing.T, body string) *jobRun { + t.Helper() + var r jobRun + require.NoError(t, json.Unmarshal([]byte(body), &r)) + return &r + } + + t.Run("fills config path, experiment, and accelerators", func(t *testing.T) { + raw := rawRun(t, `{"run_id":5,"tasks":[{"ai_runtime_task":{ + "experiment":"/Users/me@example.com/my-exp", + "deployments":[{"command_path":"/Workspace/run/command.sh","compute":{"accelerator_type":"GPU_1xA10","accelerator_count":1}}] + }}]}`) + data := &getData{ExperimentDisplay: na, AcceleratorsDisplay: na} + enrichFromRawRun(raw, 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) { + raw := rawRun(t, `{"run_id":5,"tasks":[{}]}`) + data := &getData{ExperimentDisplay: na, AcceleratorsDisplay: na} + enrichFromRawRun(raw, data) + assert.Empty(t, data.TrainingConfigPath) + assert.Equal(t, na, data.ExperimentDisplay) + assert.Equal(t, na, data.AcceleratorsDisplay) + }) +} + +func TestFetchRun(t *testing.T) { + t.Run("parses the run into both the typed and raw shapes from one call", func(t *testing.T) { + body := `{"run_id":5,"creator_user_name":"me@example.com","tasks":[{"ai_runtime_task":{ + "experiment":"/Users/me@example.com/my-exp", + "deployments":[{"command_path":"/Workspace/run/command.sh","compute":{"accelerator_type":"GPU_1xA10","accelerator_count":1}}] + }}]}` + srv := runGetServer(t, body) + typed, raw, err := fetchRun(t.Context(), newTestWorkspaceClient(t, srv.URL), 5) + require.NoError(t, err) + // Typed jobs.Run carries the common fields; raw jobRun keeps ai_runtime_task. + assert.Equal(t, int64(5), typed.RunId) + assert.Equal(t, "me@example.com", typed.CreatorUserName) + assert.Equal(t, "/Workspace/run/command.sh", raw.commandPath()) + }) + + t.Run("propagates a fetch failure", func(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusInternalServerError) + })) + t.Cleanup(srv.Close) + _, _, err := fetchRun(t.Context(), newTestWorkspaceClient(t, srv.URL), 5) + require.Error(t, err) + }) +} + +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..c6f94e1a69f --- /dev/null +++ b/experimental/air/cmd/joblist.go @@ -0,0 +1,285 @@ +package aircmd + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "net/http" + + "github.com/databricks/databricks-sdk-go" + "github.com/databricks/databricks-sdk-go/apierr" + "github.com/databricks/databricks-sdk-go/client" + "github.com/databricks/databricks-sdk-go/service/jobs" + "golang.org/x/sync/errgroup" +) + +// jobsRunsListPath and jobsRunsGetPath are the Jobs endpoints we call directly +// (rather than via the typed SDK) because the SDK's RunTask omits +// ai_runtime_task, the task type the AI runtime now submits. +const ( + jobsRunsListPath = "/api/2.2/jobs/runs/list" + jobsRunsGetPath = "/api/2.2/jobs/runs/get" +) + +// hydrateConcurrency bounds the parallel runs/get calls when hydrating a batch +// of run ids from the AiTrainingService index. +const hydrateConcurrency = 16 + +type jobsRunsListResponse struct { + Runs []jobRun `json:"runs"` + NextPageToken string `json:"next_page_token"` +} + +type jobRun struct { + RunID int64 `json:"run_id"` + RunName string `json:"run_name"` + CreatorUserName string `json:"creator_user_name"` + StartTime int64 `json:"start_time"` + EndTime int64 `json:"end_time"` + State jobState `json:"state"` + Tasks []jobTask `json:"tasks"` +} + +type jobState struct { + LifeCycleState string `json:"life_cycle_state"` + ResultState string `json:"result_state"` +} + +type jobTask struct { + RunID int64 `json:"run_id"` + StartTime int64 `json:"start_time"` + EndTime int64 `json:"end_time"` + AiRuntimeTask *jobAiRuntimeTask `json:"ai_runtime_task"` + GenAiComputeTask *genAiComputeTask `json:"gen_ai_compute_task"` + ForEachTask *forEachTask `json:"for_each_task"` +} + +type forEachTask struct { + Task jobTask `json:"task"` +} + +// jobAiRuntimeTask is the current AI runtime task shape; deployments[0].compute +// carries the accelerator info. +type jobAiRuntimeTask struct { + Experiment string `json:"experiment"` + Deployments []aiRuntimeDeploy `json:"deployments"` +} + +type aiRuntimeDeploy struct { + Compute airCompute `json:"compute"` + // CommandPath is command.sh; training_config.yaml sits beside it. + CommandPath string `json:"command_path"` +} + +type airCompute struct { + AcceleratorType string `json:"accelerator_type"` + AcceleratorCount int `json:"accelerator_count"` +} + +// genAiComputeTask is the legacy task shape, still recognized for older runs. +type genAiComputeTask struct { + TrainingScriptPath string `json:"training_script_path"` + MlflowExperimentName string `json:"mlflow_experiment_name"` + Compute *genAiCompute `json:"compute"` +} + +type genAiCompute struct { + GpuType string `json:"gpu_type"` + NumGpus int `json:"num_gpus"` +} + +// airTask returns the run's AI runtime / legacy GenAI task, unwrapping a foreach +// sweep when present. +func (r *jobRun) airTask() (*jobAiRuntimeTask, *genAiComputeTask) { + if len(r.Tasks) == 0 { + return nil, nil + } + t := r.Tasks[0] + if t.AiRuntimeTask != nil || t.GenAiComputeTask != nil { + return t.AiRuntimeTask, t.GenAiComputeTask + } + if t.ForEachTask != nil { + return t.ForEachTask.Task.AiRuntimeTask, t.ForEachTask.Task.GenAiComputeTask + } + return nil, nil +} + +// 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 *jobRun) bool { + ai, gen := r.airTask() + return ai != nil || (gen != nil && gen.TrainingScriptPath != "") +} + +// isSweep reports whether the run's first task fans out into iterations. +func isSweep(r *jobRun) 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 *jobRun) int64 { + if len(r.Tasks) == 0 { + return 0 + } + t := r.Tasks[0] + if t.ForEachTask != nil { + return t.ForEachTask.Task.RunID + } + return t.RunID +} + +// jobExperiment returns the run's MLflow experiment name (user-folder prefix +// stripped), or "" when there is none. +func jobExperiment(r *jobRun) string { + ai, gen := r.airTask() + switch { + case ai != nil && ai.Experiment != "": + return stripExperimentUserPrefix(ai.Experiment) + case gen != nil && gen.MlflowExperimentName != "": + return stripExperimentUserPrefix(gen.MlflowExperimentName) + } + return "" +} + +// jobCompute returns the run's accelerator type and count, or ("", 0) when it +// has none. +func jobCompute(r *jobRun) (string, int) { + ai, gen := r.airTask() + switch { + case ai != nil && len(ai.Deployments) > 0: + c := ai.Deployments[0].Compute + return c.AcceleratorType, c.AcceleratorCount + case gen != nil && gen.Compute != nil: + return gen.Compute.GpuType, gen.Compute.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 *jobRun) (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 *jobRun) bool { + switch r.State.LifeCycleState { + case "TERMINATED", "INTERNAL_ERROR", "SKIPPED": + return true + } + return false +} + +// fetchJobRunsPage fetches one page of Jobs runs/list. query carries the request +// params (and page_token across calls). +func fetchJobRunsPage(ctx context.Context, w *databricks.WorkspaceClient, query map[string]any) (*jobsRunsListResponse, error) { + apiClient, err := client.New(w.Config) + if err != nil { + return nil, fmt.Errorf("failed to create API client: %w", err) + } + + var resp jobsRunsListResponse + err = apiClient.Do(ctx, http.MethodGet, jobsRunsListPath, nil, nil, query, &resp) + if err != nil { + return nil, fmt.Errorf("failed to list runs: %w", err) + } + return &resp, nil +} + +// fetchJobRun fetches a single run via runs/get. The response mirrors one +// runs/list element, so it deserializes into jobRun directly; expand_tasks +// pulls the ai_runtime_task the typed SDK omits. +func fetchJobRun(ctx context.Context, w *databricks.WorkspaceClient, runID int64) (*jobRun, error) { + apiClient, err := client.New(w.Config) + if err != nil { + return nil, fmt.Errorf("failed to create API client: %w", err) + } + + var run jobRun + query := map[string]any{"run_id": runID, "expand_tasks": true} + err = apiClient.Do(ctx, http.MethodGet, jobsRunsGetPath, nil, nil, query, &run) + if err != nil { + return nil, fmt.Errorf("failed to get run %d: %w", runID, err) + } + return &run, nil +} + +// fetchRun fetches a run once and returns it in both shapes: the typed SDK +// jobs.Run the display path consumes, and the raw jobRun that preserves the +// ai_runtime_task the typed model drops. The single runs/get body is unmarshalled +// into both, avoiding a second identical roundtrip. +func fetchRun(ctx context.Context, w *databricks.WorkspaceClient, runID int64) (*jobs.Run, *jobRun, error) { + apiClient, err := client.New(w.Config) + if err != nil { + return nil, nil, fmt.Errorf("failed to create API client: %w", err) + } + + var body json.RawMessage + query := map[string]any{"run_id": runID, "expand_tasks": true} + if err := apiClient.Do(ctx, http.MethodGet, jobsRunsGetPath, nil, nil, query, &body); err != nil { + return nil, nil, err + } + + var typed jobs.Run + if err := json.Unmarshal(body, &typed); err != nil { + return nil, nil, fmt.Errorf("failed to parse run %d: %w", runID, err) + } + var raw jobRun + if err := json.Unmarshal(body, &raw); err != nil { + return nil, nil, fmt.Errorf("failed to parse run %d: %w", runID, err) + } + return &typed, &raw, 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) ([]*jobRun, error) { + runs := make([]*jobRun, 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([]*jobRun, 0, len(runs)) + for _, run := range runs { + if run != nil { + hydrated = append(hydrated, run) + } + } + return hydrated, nil +} + +// commandPath returns the run's command.sh path, or "" when it has no deployment. +func (r *jobRun) commandPath() string { + ai, _ := r.airTask() + if ai == nil || len(ai.Deployments) == 0 { + return "" + } + return ai.Deployments[0].CommandPath +} diff --git a/experimental/air/cmd/joblist_test.go b/experimental/air/cmd/joblist_test.go new file mode 100644 index 00000000000..d8f03fc60ba --- /dev/null +++ b/experimental/air/cmd/joblist_test.go @@ -0,0 +1,50 @@ +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 == jobsRunsGetPath { + _, _ = 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, "/Workspace/Users/me/.air/cli_launch/my-exp/my-exp_abc/command.sh", run.commandPath()) +} + +func TestCommandPathEmpty(t *testing.T) { + // No ai_runtime_task deployment: no command path. + assert.Empty(t, (&jobRun{Tasks: []jobTask{{}}}).commandPath()) + assert.Empty(t, (&jobRun{}).commandPath()) +} diff --git a/experimental/air/cmd/list.go b/experimental/air/cmd/list.go new file mode 100644 index 00000000000..d5e5e48bcc5 --- /dev/null +++ b/experimental/air/cmd/list.go @@ -0,0 +1,331 @@ +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/service/iam" + "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 + query map[string]any + userFilter string + filters listFilters + + pending []jobRun // runs from the last page not yet inspected + scanned int + drained bool +} + +func newJobsScanStrategy(ctx context.Context, w *databricks.WorkspaceClient, q listQuery) *jobsScanStrategy { + query := map[string]any{ + "run_type": "SUBMIT_RUN", + "expand_tasks": true, + "limit": jobsPageLimit, + } + if q.activeOnly { + query["active_only"] = true + } + return &jobsScanStrategy{ + ctx: ctx, + w: w, + query: query, + userFilter: q.userFilter, + filters: q.filters, + } +} + +func (s *jobsScanStrategy) next(want int) ([]listedRun, error) { + var entries []listedRun + for len(entries) < want { + if len(s.pending) == 0 { + if s.drained || s.scanned >= maxListScan { + break + } + if err := s.fetchPage(); err != nil { + return nil, err + } + continue + } + + run := &s.pending[0] + s.pending = s.pending[1:] + s.scanned++ + + 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.drained && len(s.pending) == 0) || s.scanned >= maxListScan +} + +func (s *jobsScanStrategy) truncated() bool { + return s.scanned >= maxListScan +} + +// fetchPage loads the next runs/list page into the pending buffer, marking the +// strategy drained once the server reports no further pages. +func (s *jobsScanStrategy) fetchPage() error { + resp, err := fetchJobRunsPage(s.ctx, s.w, s.query) + if err != nil { + return err + } + s.pending = resp.Runs + if resp.NextPageToken == "" { + s.drained = true + } else { + s.query["page_token"] = resp.NextPageToken + } + return nil +} + +// 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..f03dd7bc29f --- /dev/null +++ b/experimental/air/cmd/list_cache_test.go @@ -0,0 +1,104 @@ +package aircmd + +import ( + "testing" + + "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]jobRun{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]jobRun{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]jobRun{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(&jobRun{State: jobState{LifeCycleState: "TERMINATED"}})) + assert.True(t, isTerminal(&jobRun{State: jobState{LifeCycleState: "INTERNAL_ERROR"}})) + assert.False(t, isTerminal(&jobRun{State: jobState{LifeCycleState: "RUNNING"}})) + assert.False(t, isTerminal(&jobRun{State: jobState{LifeCycleState: "PENDING"}})) +} diff --git a/experimental/air/cmd/list_filter.go b/experimental/air/cmd/list_filter.go new file mode 100644 index 00000000000..b1cedcb0b99 --- /dev/null +++ b/experimental/air/cmd/list_filter.go @@ -0,0 +1,106 @@ +package aircmd + +import ( + "fmt" + "path" + "strconv" + "strings" +) + +// 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 *jobRun) 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 *jobRun) 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..f8351d33d8d --- /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 := airJobRun(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..5dcb21bfc62 --- /dev/null +++ b/experimental/air/cmd/list_format.go @@ -0,0 +1,45 @@ +package aircmd + +import ( + "strconv" + "time" +) + +// 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 *jobRun) 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: statusWord(run.State.LifeCycleState, run.State.ResultState), + 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..d6c6cf6958f --- /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..ec8d082e3a0 --- /dev/null +++ b/experimental/air/cmd/list_index_test.go @@ -0,0 +1,212 @@ +package aircmd + +import ( + "encoding/json" + "fmt" + "net/http" + "net/http/httptest" + "strconv" + "testing" + + "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) jobRun { + r := airJobRun(id, "me@example.com", "GPU_1xH100", 1, "/Users/me@example.com/exp") + r.State = jobState{LifeCycleState: "TERMINATED", ResultState: "SUCCESS"} + 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]jobRun, 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 jobsRunsGetPath: + 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]jobRun{ + 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]jobRun{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]jobRun{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 jobsRunsGetPath: + 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 == jobsRunsListPath { + _, _ = fmt.Fprint(w, runsListBody(t, "", airJobRun(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..8d186eb048d --- /dev/null +++ b/experimental/air/cmd/list_test.go @@ -0,0 +1,247 @@ +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/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// airJobRun builds a single-task AIR run (ai_runtime_task), as runs/list returns +// it with expand_tasks. +func airJobRun(id int64, user, accelType string, count int, experiment string) jobRun { + return jobRun{ + RunID: id, + RunName: "run-" + strconv.FormatInt(id, 10), + CreatorUserName: user, + State: jobState{LifeCycleState: "RUNNING"}, + Tasks: []jobTask{{AiRuntimeTask: &jobAiRuntimeTask{ + Experiment: experiment, + Deployments: []aiRuntimeDeploy{{ + Compute: airCompute{AcceleratorType: accelType, AcceleratorCount: count}, + }}, + }}}, + } +} + +// runsListBody marshals one runs/list response page. +func runsListBody(t *testing.T, nextToken string, runs ...jobRun) string { + t.Helper() + b, err := json.Marshal(jobsRunsListResponse{Runs: runs, NextPageToken: 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 == jobsRunsListPath { + 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 := []jobRun{ + airJobRun(1, "me@example.com", "GPU_8xH100", 8, "/Users/me@example.com/exp-a"), + {RunID: 2, CreatorUserName: "me@example.com", Tasks: []jobTask{{}}}, // not an AIR run + airJobRun(3, "other@example.com", "GPU_1xA10", 1, "/Users/other/exp-b"), // wrong user + airJobRun(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 := []jobRun{ + airJobRun(1, "me@example.com", "GPU_1xH100", 1, "/Users/me@example.com/qwen-train"), + airJobRun(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 := []jobRun{ + airJobRun(1, "me@example.com", "GPU_1xH100", 1, "exp-a"), + airJobRun(2, "me@example.com", "GPU_1xH100", 1, "exp-b"), + airJobRun(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", airJobRun(1, "me@example.com", "GPU_1xH100", 1, "exp-a")) + page2 := runsListBody(t, "", airJobRun(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 buffer the rest and hand it back on +// the following call, then report exhaustion — without re-fetching. +func TestRunFetcherResumesAcrossCalls(t *testing.T) { + runs := []jobRun{ + airJobRun(1, "me@example.com", "GPU_1xH100", 1, "exp-a"), + airJobRun(2, "me@example.com", "GPU_1xH100", 1, "exp-b"), + airJobRun(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 buffered 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) +} + +// TestFetchJobRunsParsesAiRuntimeTask pins the raw parse against the real +// runs/get shape, since the typed SDK omits ai_runtime_task. +func TestFetchJobRunsParsesAiRuntimeTask(t *testing.T) { + body := `{"runs":[{ + "run_id": 842552489592352, + "run_name": "my-first-air-run", + "creator_user_name": "me@example.com", + "start_time": 1700000000000, + "end_time": 1700000012000, + "state": {"life_cycle_state": "TERMINATED", "result_state": "SUCCESS"}, + "tasks": [{ + "ai_runtime_task": { + "experiment": "my-first-air-run", + "deployments": [{"compute": {"accelerator_count": 1, "accelerator_type": "GPU_1xA10"}}] + } + }] + }]}` + srv := runsServer(t, body) + + resp, err := fetchJobRunsPage(t.Context(), newTestWorkspaceClient(t, srv.URL), map[string]any{}) + require.NoError(t, err) + require.Len(t, resp.Runs, 1) + run := &resp.Runs[0] + 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 := airJobRun(123, "me@example.com", "GPU_8xH100", 8, "/Users/me@example.com/exp") + run.StartTime = 1700000000000 + run.EndTime = 1700000012000 + run.State = jobState{ResultState: "SUCCESS"} + + 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(&jobRun{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 := jobRun{RunID: 9, Tasks: []jobTask{{ + ForEachTask: &forEachTask{Task: jobTask{AiRuntimeTask: &jobAiRuntimeTask{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..88e405c867a --- /dev/null +++ b/experimental/air/cmd/list_tui.go @@ -0,0 +1,253 @@ +package aircmd + +import ( + "fmt" + "io" + "strings" + + tea "github.com/charmbracelet/bubbletea" + "github.com/charmbracelet/lipgloss" + "github.com/databricks/cli/libs/cmdio" + "github.com/muesli/termenv" + "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() + + color := cmdio.SupportsColor(ctx, out) + r := lipgloss.NewRenderer(out) + if !color { + r.SetColorProfile(termenv.Ascii) + } + + // 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..f1cff1e59f0 --- /dev/null +++ b/experimental/air/cmd/list_tui_test.go @@ -0,0 +1,221 @@ +package aircmd + +import ( + "io" + "strconv" + "testing" + + tea "github.com/charmbracelet/bubbletea" + "github.com/charmbracelet/lipgloss" + "github.com/muesli/termenv" + "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() listModel { + r := lipgloss.NewRenderer(io.Discard) + r.SetColorProfile(termenv.Ascii) + 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() + 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() + // 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 := lipgloss.NewRenderer(io.Discard) + r.SetColorProfile(termenv.Ascii) + 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 := lipgloss.NewRenderer(io.Discard) + r.SetColorProfile(termenv.Ascii) + 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() + 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() + 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 := lipgloss.NewRenderer(io.Discard) + r.SetColorProfile(termenv.Ascii) + + // 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() + _, 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().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 := lipgloss.NewRenderer(io.Discard) + r.SetColorProfile(termenv.Ascii) + 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..02a7956d7bd --- /dev/null +++ b/experimental/air/cmd/mlflow.go @@ -0,0 +1,134 @@ +package aircmd + +import ( + "context" + "fmt" + "net/http" + "strings" + + "github.com/databricks/cli/libs/log" + "github.com/databricks/databricks-sdk-go" + "github.com/databricks/databricks-sdk-go/client" + "github.com/databricks/databricks-sdk-go/service/jobs" +) + +// getRunOutputResponse is the part of the jobs runs/get-output response we use. +// The MLflow identifiers live under ai_runtime_task_output (current) or +// gen_ai_compute_output.run_info (legacy), neither modeled by the typed SDK, so +// we call the endpoint directly and parse just these fields. +type getRunOutputResponse struct { + AiRuntimeTaskOutput *struct { + MlflowExperimentID string `json:"mlflow_experiment_id"` + MlflowRunID string `json:"mlflow_run_id"` + } `json:"ai_runtime_task_output"` + GenAiComputeOutput *struct { + RunInfo *struct { + MlflowExperimentID string `json:"mlflow_experiment_id"` + MlflowRunID string `json:"mlflow_run_id"` + } `json:"run_info"` + } `json:"gen_ai_compute_output"` +} + +// 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 + } + + apiClient, err := client.New(w.Config) + if err != nil { + log.Debugf(ctx, "air: could not build API client for MLflow link: %v", err) + return nil + } + + // Pass run_id through the request arg (the SDK serializes it to the query + // string for GET); passing it via queryParams instead leaves a nil body that + // this endpoint rejects with "expected a map". + var out getRunOutputResponse + err = apiClient.Do(ctx, http.MethodGet, "/api/2.2/jobs/runs/get-output", + nil, nil, map[string]any{"run_id": taskRunID}, &out) + 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} + } + if o := out.GenAiComputeOutput; o != nil && o.RunInfo != nil && + o.RunInfo.MlflowExperimentID != "" && o.RunInfo.MlflowRunID != "" { + return &mlflowIdentifiers{ExperimentID: o.RunInfo.MlflowExperimentID, RunID: o.RunInfo.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 { + apiClient, err := client.New(w.Config) + if err != nil { + log.Debugf(ctx, "air get: could not build API client for MLflow run name: %v", err) + return "" + } + var out struct { + Run struct { + Info struct { + RunName string `json:"run_name"` + } `json:"info"` + } `json:"run"` + } + err = apiClient.Do(ctx, http.MethodGet, "/api/2.0/mlflow/runs/get", + nil, nil, map[string]any{"run_id": mlflowRunID}, &out) + if err != nil { + log.Debugf(ctx, "air get: could not fetch MLflow run name: %v", err) + return "" + } + return out.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..06529d5273f --- /dev/null +++ b/experimental/air/cmd/mlflow_test.go @@ -0,0 +1,148 @@ +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. +// mlflowIDs calls the runs/get-output REST endpoint directly (the field it needs +// is not modeled by the typed SDK), so it must be exercised over HTTP. +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, `{"gen_ai_compute_output":{"run_info":{"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("parses legacy gen_ai_compute_output", func(t *testing.T) { + var hit bool + srv := runOutputServer(t, `{"gen_ai_compute_output":{"run_info":{"mlflow_experiment_id":"E2","mlflow_run_id":"R2"}}}`, &hit) + got := mlflowIDsForTask(ctx, newTestWorkspaceClient(t, srv.URL), 99) + require.NotNil(t, got) + assert.Equal(t, &mlflowIdentifiers{ExperimentID: "E2", RunID: "R2"}, 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..3d1fe9637c5 --- /dev/null +++ b/experimental/air/cmd/render.go @@ -0,0 +1,391 @@ +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) { + colorOn := cmdio.SupportsColor(ctx, out) + renderer := lipgloss.NewRenderer(out) + if !colorOn { + // Ascii emits no SGR codes; combined with the link fallback below this + // gives clean, un-escaped output under --no-color / NO_COLOR / piped stdout. + renderer.SetColorProfile(termenv.Ascii) + } + 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..e950bb22331 --- /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/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() palette { + r := lipgloss.NewRenderer(io.Discard) + r.SetColorProfile(termenv.Ascii) + 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() + // 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() + 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() + 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() + // 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..5a0b6f27e1e --- /dev/null +++ b/experimental/air/cmd/runconfig.go @@ -0,0 +1,466 @@ +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 branch/remote names against command injection. 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) + } + if g.Remote.isString { + if g.Remote.name == "" { + return errors.New("git.remote string cannot be empty; use 'true' to auto-detect") + } + if !gitRefRe.MatchString(g.Remote.name) { + return fmt.Errorf("invalid git.remote name %q: only alphanumeric characters, hyphens, dots, slashes, and underscores are allowed", g.Remote.name) + } + } + + 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") + } + if g.Remote.truthy() && g.Branch == nil { + return errors.New("git.remote requires git.branch (only valid with branch refs)") + } + 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..1dd5ce1ea39 --- /dev/null +++ b/experimental/air/cmd/runconfig_test.go @@ -0,0 +1,419 @@ +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 + remote: origin + 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.True(t, cfg.CodeSource.Snapshot.Git.Remote.isString) + assert.Equal(t, "origin", cfg.CodeSource.Snapshot.Git.Remote.name) + 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", func(t *testing.T) { + cfg, err := loadRunConfig(writeConfig(t, minimalConfig+` +code_source: + type: snapshot + snapshot: + root_path: . + git: + branch: main + remote: true +`)) + require.NoError(t, err) + r := cfg.CodeSource.Snapshot.Git.Remote + assert.False(t, r.isString) + assert.True(t, r.enabled) + assert.True(t, r.truthy()) + }) + + 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")}, ""}, + {"branch with remote ok", gitRef{Branch: str("main"), Remote: gitRemote{set: true, enabled: true}}, ""}, + {"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 without branch", gitRef{Commit: str("abc"), Remote: gitRemote{set: true, isString: true, name: "origin"}}, "requires git.branch"}, + {"bad branch chars", gitRef{Branch: str("bad branch")}, "invalid git.branch"}, + {"empty remote string", gitRef{Branch: str("main"), Remote: gitRemote{set: true, isString: true, name: ""}}, "cannot be empty"}, + } + 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..49f51a8aa66 --- /dev/null +++ b/experimental/air/cmd/runsubmit.go @@ -0,0 +1,258 @@ +package aircmd + +import ( + "context" + "errors" + "fmt" + "net/http" + "path" + "strconv" + "strings" + + "github.com/databricks/cli/libs/auth" + "github.com/databricks/cli/libs/env" + "github.com/databricks/cli/libs/filer" + "github.com/databricks/databricks-sdk-go" + "github.com/databricks/databricks-sdk-go/client" + "github.com/google/uuid" +) + +// jobsRunsSubmitPath is the Jobs one-time-run endpoint. air builds the full +// payload and POSTs it here directly — the native ai_runtime_task is not modeled +// by the typed SDK, and we want no genai-mapi forwarding. +const jobsRunsSubmitPath = "/api/2.2/jobs/runs/submit" + +// 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" + +// aiRuntimeCompute is a deployment's accelerator request. +type aiRuntimeCompute struct { + AcceleratorType string `json:"accelerator_type"` + AcceleratorCount int `json:"accelerator_count"` +} + +// aiRuntimeDeployment is one worker deployment of the run. +type aiRuntimeDeployment struct { + CommandPath string `json:"command_path"` + Compute aiRuntimeCompute `json:"compute"` +} + +// aiRuntimeTask is the native AI Runtime task. It routes straight to the training +// service — no genai-mapi forwarding. The proto is lean: env vars, secrets, +// requirements, and hyperparameters are staged as workspace files co-located with +// command.sh (see runupload.go), not carried inline. +type aiRuntimeTask struct { + Experiment string `json:"experiment"` + Deployments []aiRuntimeDeployment `json:"deployments"` + MlflowRun string `json:"mlflow_run,omitempty"` + MlflowExperimentDirectory string `json:"mlflow_experiment_directory,omitempty"` +} + +// environmentSpec carries the bare runtime channel ("4", "5", ...). +type environmentSpec struct { + EnvironmentVersion string `json:"environment_version"` +} + +// jobEnvironment is the serverless environment a task references for its runtime. +type jobEnvironment struct { + EnvironmentKey string `json:"environment_key"` + Spec environmentSpec `json:"spec"` +} + +// submitTask is the single task air submits: a native ai_runtime_task. +// +// 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. +type submitTask struct { + TaskKey string `json:"task_key"` + RunIf string `json:"run_if"` + AiRuntimeTask aiRuntimeTask `json:"ai_runtime_task"` + EnvironmentKey string `json:"environment_key"` + MaxRetries int `json:"max_retries"` + RetryOnTimeout bool `json:"retry_on_timeout,omitempty"` +} + +// jobsSubmitRun is the Jobs runs/submit payload. +type jobsSubmitRun struct { + RunName string `json:"run_name"` + TimeoutSeconds int `json:"timeout_seconds,omitempty"` + Tasks []submitTask `json:"tasks"` + Environments []jobEnvironment `json:"environments"` + BudgetPolicyID string `json:"budget_policy_id,omitempty"` + IdempotencyToken string `json:"idempotency_token,omitempty"` +} + +// 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. +func buildSubmitPayload(cfg *runConfig, commandPath, dlImage string) jobsSubmitRun { + task := aiRuntimeTask{ + Experiment: cfg.ExperimentName, + Deployments: []aiRuntimeDeployment{{ + CommandPath: commandPath, + Compute: aiRuntimeCompute{ + AcceleratorType: cfg.Compute.AcceleratorType, + AcceleratorCount: cfg.Compute.NumAccelerators, + }, + }}, + } + if cfg.MLflowRunName != nil { + task.MlflowRun = *cfg.MLflowRunName + } + if cfg.MLflowExperimentDirectory != nil { + task.MlflowExperimentDirectory = *cfg.MLflowExperimentDirectory + } + + st := submitTask{ + TaskKey: cfg.ExperimentName, + RunIf: "ALL_SUCCESS", + AiRuntimeTask: task, + EnvironmentKey: aiRuntimeEnvironmentKey, + MaxRetries: cfg.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). + st.RetryOnTimeout = st.MaxRetries > 0 + + return jobsSubmitRun{ + RunName: cfg.ExperimentName, + TimeoutSeconds: cfg.timeoutSeconds(), + Tasks: []submitTask{st}, + Environments: []jobEnvironment{{ + EnvironmentKey: aiRuntimeEnvironmentKey, + Spec: environmentSpec{EnvironmentVersion: dlImage}, + }}, + } +} + +// jobsSubmitClient submits one-time runs through the Jobs API. +type jobsSubmitClient struct { + c *client.DatabricksClient +} + +func newJobsSubmitClient(w *databricks.WorkspaceClient) (*jobsSubmitClient, error) { + c, err := client.New(w.Config) + if err != nil { + return nil, err + } + return &jobsSubmitClient{c: c}, nil +} + +type submitRunResponse struct { + RunID int64 `json:"run_id,omitempty"` +} + +// submit POSTs the payload to runs/submit and returns the new run_id. +func (j *jobsSubmitClient) submit(ctx context.Context, payload jobsSubmitRun) (int64, error) { + var resp submitRunResponse + if err := j.c.Do(ctx, http.MethodPost, jobsRunsSubmitPath, auth.WorkspaceIDHeaders(j.c.Config), nil, payload, &resp); err != nil { + return 0, err + } + return resp.RunID, nil +} + +// 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 and packaging a + // code_source snapshot are not ported yet; reject rather than silently drop. + if cfg.UsagePolicyName != nil { + return 0, "", errors.New("usage_policy_name is not yet supported") + } + if cfg.CodeSource != nil { + return 0, "", errors.New("code_source 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 + } + + runtimeVersion, _ := cfg.runtimeVersion() + payload := buildSubmitPayload(cfg, path.Join(funcDir, commandScriptName), dlRuntimeImage(ctx, runtimeVersion)) + payload.IdempotencyToken = token + + jc, err := newJobsSubmitClient(w) + if err != nil { + return 0, "", err + } + runID, err := jc.submit(ctx, payload) + if err != nil { + return 0, "", err + } + + 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..374b584bc63 --- /dev/null +++ b/experimental/air/cmd/runsubmit_test.go @@ -0,0 +1,194 @@ +package aircmd + +import ( + "encoding/json" + "strings" + "testing" + + "github.com/databricks/cli/libs/testserver" + "github.com/databricks/databricks-sdk-go" + "github.com/databricks/databricks-sdk-go/config" + "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") + + 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) + 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, "ALL_SUCCESS", task.RunIf) + assert.Equal(t, aiRuntimeEnvironmentKey, task.EnvironmentKey) + assert.Equal(t, 2, task.MaxRetries) + assert.True(t, task.RetryOnTimeout) + + at := task.AiRuntimeTask + 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, aiRuntimeCompute{AcceleratorType: "GPU_8xH100", 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").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").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 TestJobsSubmitClient(t *testing.T) { + server := testserver.New(t) + t.Cleanup(server.Close) + + var got jobsSubmitRun + server.Handle("POST", "/api/2.2/jobs/runs/submit", func(req testserver.Request) any { + require.NoError(t, json.Unmarshal(req.Body, &got)) + return submitRunResponse{RunID: 999} + }) + + w := &databricks.WorkspaceClient{Config: &config.Config{Host: server.URL, Token: "token"}} + jc, err := newJobsSubmitClient(w) + require.NoError(t, err) + + runID, err := jc.submit(t.Context(), jobsSubmitRun{RunName: "exp", Tasks: []submitTask{{TaskKey: "exp"}}}) + require.NoError(t, err) + assert.Equal(t, int64(999), runID) + assert.Equal(t, "exp", got.RunName) +} + +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 jobsSubmitRun + server.Handle("POST", "/api/2.2/jobs/runs/submit", func(req testserver.Request) any { + require.NoError(t, json.Unmarshal(req.Body, &got)) + return 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.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, aiRuntimeCompute{AcceleratorType: "GPU_1xH100", AcceleratorCount: 1}, d.Compute) +} + +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") + }) + + t.Run("code_source rejected", func(t *testing.T) { + cfg := *base + cfg.CodeSource = &codeSourceConfig{Type: "snapshot"} + _, _, err := submitWorkload(t.Context(), w, &cfg, cfgPath, "") + require.ErrorContains(t, err, "code_source 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/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/go.mod b/go.mod index 91f788b62f5..a6e3179d04e 100644 --- a/go.mod +++ b/go.mod @@ -24,6 +24,7 @@ require ( github.com/hexops/gotextdiff v1.0.3 // BSD-3-Clause github.com/jackc/pgx/v5 v5.10.0 // MIT github.com/mattn/go-isatty v0.0.22 // MIT + github.com/muesli/termenv v0.16.0 // MIT github.com/palantir/pkg/yamlpatch v1.5.0 // BSD-3-Clause github.com/pkg/browser v0.0.0-20240102092130-5ac0b6a4141c // BSD-2-Clause github.com/quasilyte/go-ruleguard/dsl v0.3.22 // BSD-3-Clause @@ -85,7 +86,6 @@ require ( github.com/mitchellh/hashstructure/v2 v2.0.2 // indirect github.com/muesli/ansi v0.0.0-20230316100256-276c6243b2f6 // indirect github.com/muesli/cancelreader v0.2.2 // indirect - github.com/muesli/termenv v0.16.0 // indirect github.com/pkg/errors v0.9.1 // indirect github.com/pmezard/go-difflib v1.0.0 // indirect github.com/rivo/uniseg v0.4.7 // indirect 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 {