AIR CLI Integration Batch 1

NOTICE

@@ -155,6 +155,10 @@ mattn/go-isatty - https://github.com/mattn/go-isatty
 Copyright (c) Yasuhiro MATSUMOTO <mattn.jp@gmail.com>
 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

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).

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

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}}]
+        }
+      }]
+    }
+  ]
+}
+'''

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"
+  }
+}

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

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"}}}
+'''

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

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

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

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"}}}
+'''