air: find the AI Runtime task anywhere in a multi-task job in air list and air get

.nextchanges/air/multitask-air-tasks.md

@@ -0,0 +1 @@
+* `air list` and `air get` now read the AI Runtime task wherever it sits in a multi-task job (for example after a notebook prep task) instead of the first task; a job with several AI Runtime tasks lists one row per task, `air get` gains `--task-key`, and both commands report `task_key` in JSON output. ([#7632](https://github.com/databricks/cli/pull/7632))

acceptance/air/control-chars/output.txt

@@ -38,6 +38,7 @@ MLflow URL: [DATABRICKS_URL]/ml/experiments/exp1/runs/run1
     "runs": [
       {
         "run_id": "123",
+        "task_key": "train",
         "run_name": "qwen\u001b[2J-train",
         "user": "user\u001b[[email protected]",
         "status": "SUCCESS",

acceptance/air/get-ai-runtime/output.txt

@@ -43,6 +43,7 @@ MLflow URL: [DATABRICKS_URL]/ml/experiments/exp1/runs/run1
   "ts": "[TIMESTAMP]",
   "data": {
     "run_id": "123",
+    "task_key": "train",
     "status": "SUCCESS",
     "started_at": "[TIMESTAMP]",
     "duration_seconds": 12,

acceptance/air/get-multitask/out.test.toml

@@ -0,0 +1 @@
+Cloud = false

acceptance/air/get-multitask/output.txt

@@ -0,0 +1,113 @@
+
+=== get (text): the first AI Runtime task, not the notebook prep task
+
+>>> [CLI] air get 123
+
+╭─ Metadata ─────────────────────────────────────────────────────╮
+│                                                                │
+│  Run ID        123                                             │
+│  Task          train                                           │
+│  Status        ● SUCCESS                                       │
+│  Submitted     2023-11-14 22:13 UTC                            │
+│  Retries       0                                               │
+│  Max Retries   0                                               │
+│  Duration      10s                                             │
+│  Experiment    air-byot-multitask-dag                          │
+│  MLflow Run    matmul-smoke                                    │
+│  User          [email protected]                                │
+│  Accelerators  1x A10                                          │
+│  Environment   5                                               │
+│                                                                │
+╰────────────────────────────────────────────────────────────────╯
+
+Run URL:    [DATABRICKS_URL]/jobs/runs/123?o=[WORKSPACE_ID]
+MLflow URL: [DATABRICKS_URL]/ml/experiments/exp1/runs/run1
+
+=== get --task-key eval (text)
+
+>>> [CLI] air get 123 --task-key eval
+
+╭─ Metadata ─────────────────────────────────────────────────────╮
+│                                                                │
+│  Run ID        123                                             │
+│  Task          eval                                            │
+│  Status        ● SUCCESS                                       │
+│  Submitted     2023-11-14 22:13 UTC                            │
+│  Retries       0                                               │
+│  Max Retries   2                                               │
+│  Duration      8s                                              │
+│  Experiment    eval-exp                                        │
+│  MLflow Run    matmul-smoke                                    │
+│  User          [email protected]                                │
+│  Accelerators  8x H100                                         │
+│  Environment   5                                               │
+│                                                                │
+╰────────────────────────────────────────────────────────────────╯
+
+Run URL:    [DATABRICKS_URL]/jobs/runs/123?o=[WORKSPACE_ID]
+MLflow URL: [DATABRICKS_URL]/ml/experiments/exp1/runs/run1
+
+=== get (json)
+
+>>> [CLI] air get 123 -o json
+{
+  "v": 1,
+  "ts": "[TIMESTAMP]",
+  "data": {
+    "run_id": "123",
+    "task_key": "train",
+    "status": "SUCCESS",
+    "started_at": "[TIMESTAMP]",
+    "duration_seconds": 10,
+    "attempt_number": 0,
+    "experiment_name": "air-byot-multitask-dag",
+    "dashboard_url": "[DATABRICKS_URL]/jobs/runs/123?o=[WORKSPACE_ID]",
+    "mlflow_url": "[DATABRICKS_URL]/ml/experiments/exp1/runs/run1/logs",
+    "est_remaining_seconds": null,
+    "est_percent_complete": null
+  }
+}
+
+=== get --task-key eval (json)
+
+>>> [CLI] air get 123 --task-key eval -o json
+{
+  "v": 1,
+  "ts": "[TIMESTAMP]",
+  "data": {
+    "run_id": "123",
+    "task_key": "eval",
+    "status": "SUCCESS",
+    "started_at": "[TIMESTAMP]",
+    "duration_seconds": 8,
+    "attempt_number": 0,
+    "experiment_name": "eval-exp",
+    "dashboard_url": "[DATABRICKS_URL]/jobs/runs/123?o=[WORKSPACE_ID]",
+    "mlflow_url": "[DATABRICKS_URL]/ml/experiments/exp1/runs/run1/logs",
+    "est_remaining_seconds": null,
+    "est_percent_complete": null
+  }
+}
+
+=== get --task-key prep: not an AI Runtime task
+
+>>> [CLI] air get 123 --task-key prep
+Error: run 123 has no AI Runtime task "prep"; its AI Runtime tasks are: train, eval
+
+Exit code: 1
+
+=== get --task-key prep (json)
+
+>>> [CLI] air get 123 --task-key prep -o json
+{
+  "v": 1,
+  "ts": "[TIMESTAMP]",
+  "error": {
+    "code": "INVALID_ARGS",
+    "kind": "PERMANENT",
+    "message": "run 123 has no AI Runtime task \"prep\"; its AI Runtime tasks are: train, eval",
+    "retryable": false
+  }
+}
+
+Exit code: 1

acceptance/air/get-multitask/script

@@ -0,0 +1,17 @@
+title "get (text): the first AI Runtime task, not the notebook prep task"
+trace $CLI air get 123
+
+title "get --task-key eval (text)"
+trace $CLI air get 123 --task-key eval
+
+title "get (json)"
+trace $CLI air get 123 -o json
+
+title "get --task-key eval (json)"
+trace $CLI air get 123 --task-key eval -o json
+
+title "get --task-key prep: not an AI Runtime task"
+errcode trace $CLI air get 123 --task-key prep
+
+title "get --task-key prep (json)"
+errcode trace $CLI air get 123 --task-key prep -o json

acceptance/air/get-multitask/test.toml

@@ -0,0 +1,84 @@
+# The SDK occasionally probes host reachability with a HEAD request; stub it so
+# the test is deterministic.
+
+[[Server]]
+Pattern = "HEAD /"
+Response.Body = ''
+
+# A multi-task job in the shape of the official DAB example (universe
+# ai-compute/dabs/examples/03_multitask_dag): a notebook prep task first, then the
+# ai_runtime_task that depends on it, plus a second ai_runtime_task. `air get`
+# must read the AI Runtime task wherever it sits, not tasks[0].
+[[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": "[email protected]",
+  "start_time": 1700000000000,
+  "end_time": 1700000012000,
+  "state": {"life_cycle_state": "TERMINATED", "result_state": "SUCCESS"},
+  "environments": [
+    {"environment_key": "default", "spec": {"environment_version": "5"}}
+  ],
+  "tasks": [
+    {
+      "task_key": "prep",
+      "run_id": 455,
+      "attempt_number": 0,
+      "start_time": 1700000000000,
+      "end_time": 1700000002000,
+      "state": {"life_cycle_state": "TERMINATED", "result_state": "SUCCESS"},
+      "notebook_task": {"notebook_path": "/Workspace/Users/[email protected]/prep"}
+    },
+    {
+      "task_key": "train",
+      "run_id": 456,
+      "attempt_number": 0,
+      "max_retries": 0,
+      "start_time": 1700000002000,
+      "end_time": 1700000012000,
+      "state": {"life_cycle_state": "TERMINATED", "result_state": "SUCCESS"},
+      "depends_on": [{"task_key": "prep"}],
+      "environment_key": "default",
+      "ai_runtime_task": {
+        "experiment": "air-byot-multitask-dag",
+        "deployments": [
+          {"compute": {"accelerator_type": "GPU_1xA10", "accelerator_count": 1}}
+        ]
+      }
+    },
+    {
+      "task_key": "eval",
+      "run_id": 457,
+      "attempt_number": 0,
+      "max_retries": 2,
+      "start_time": 1700000012000,
+      "end_time": 1700000020000,
+      "state": {"life_cycle_state": "TERMINATED", "result_state": "SUCCESS"},
+      "depends_on": [{"task_key": "train"}],
+      "environment_key": "default",
+      "ai_runtime_task": {
+        "experiment": "/Users/[email protected]/eval-exp",
+        "deployments": [
+          {"compute": {"accelerator_type": "GPU_8xH100", "accelerator_count": 8}}
+        ]
+      }
+    }
+  ]
+}
+'''
+
+# MLflow identifiers for the deep link, fetched for the selected task's run.
+[[Server]]
+Pattern = "GET /api/2.2/jobs/runs/get-output"
+Response.Body = '''
+{"ai_runtime_task_output": {"mlflow_experiment_id": "exp1", "mlflow_run_id": "run1"}}
+'''
+
+[[Server]]
+Pattern = "GET /api/2.0/mlflow/runs/get"
+Response.Body = '''
+{"run": {"info": {"run_name": "matmul-smoke"}}}
+'''

acceptance/air/get-pending/output.txt

@@ -5,6 +5,7 @@
   "ts": "[TIMESTAMP]",
   "data": {
     "run_id": "123",
+    "task_key": "train",
     "status": "PENDING",
     "started_at": null,
     "duration_seconds": null,

acceptance/air/get/output.txt

@@ -43,6 +43,7 @@ MLflow URL: [DATABRICKS_URL]/ml/experiments/exp1/runs/run1
   "ts": "[TIMESTAMP]",
   "data": {
     "run_id": "123",
+    "task_key": "train",
     "status": "SUCCESS",
     "started_at": "[TIMESTAMP]",
     "duration_seconds": 12,

acceptance/air/help/output.txt

@@ -92,11 +92,17 @@ Show a run's status, configuration, and timing.
 
 JOB_RUN_ID is the job run ID returned by "databricks air run".
 
+A job may run other tasks around its AI Runtime task (for example a notebook
+that prepares data first); the view shows the first AI Runtime task in task
+order. When a job has several AI Runtime tasks, pass --task-key to show another
+one.
+
 Usage:
   databricks air get JOB_RUN_ID [flags]
 
 Flags:
-  -h, --help   help for get
+  -h, --help              help for get
+      --task-key string   Show this AI Runtime task of the run (default: the first in task order)
 
 Global Flags:
       --debug            enable debug logging

acceptance/air/list/output.txt

@@ -2,8 +2,9 @@
 === list (text)
 
 >>> [CLI] air list
-   Run ID           Experiment  Status     Started              Duration  Progress  MLflow                  User                   Accelerators
-   334747067049496  qwen-train  ● SUCCESS  [TIMESTAMP]       12s         -  qwen-train-001          [USERNAME]  8x H100
+   Run ID           Experiment              Status     Started              Duration  Progress  MLflow                  User                   Accelerators
+   334747067049496  qwen-train              ● SUCCESS  [TIMESTAMP]       12s         -  qwen-train-001          [USERNAME]  8x H100
+   334747067049500  air-byot-multitask-dag  ● SUCCESS  [TIMESTAMP]       20s         -  qwen-train-001          [USERNAME]  1x A10
 
 === list (json)
 
@@ -15,6 +16,7 @@
     "runs": [
       {
         "run_id": "334747067049496",
+        "task_key": "",
         "run_name": "qwen-train",
         "user": "[USERNAME]",
         "status": "SUCCESS",
@@ -25,6 +27,20 @@
         "accelerator_type": "GPU_8xH100",
         "num_accelerators": 8,
         "dashboard_url": "[DATABRICKS_URL]/jobs/runs/334747067049496?o=[WORKSPACE_ID]"
+      },
+      {
+        "run_id": "334747067049500",
+        "task_key": "train",
+        "run_name": "air-byot-multitask-dag",
+        "user": "[USERNAME]",
+        "status": "SUCCESS",
+        "started_at": "[TIMESTAMP]",
+        "is_sweep": false,
+        "experiment_name": "air-byot-multitask-dag",
+        "duration_seconds": 20,
+        "accelerator_type": "GPU_1xA10",
+        "num_accelerators": 1,
+        "dashboard_url": "[DATABRICKS_URL]/jobs/runs/334747067049500?o=[WORKSPACE_ID]"
       }
     ]
   }
@@ -46,6 +62,7 @@
     "runs": [
       {
         "run_id": "334747067049496",
+        "task_key": "",
         "run_name": "qwen-train",
         "user": "[USERNAME]",
         "status": "SUCCESS",

acceptance/air/list/test.toml

@@ -10,8 +10,10 @@ 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
-# ([email protected], from the built-in scim/v2/Me handler). The second run
-# has no AI task and must be filtered out.
+# ([email protected], from the built-in scim/v2/Me handler). The second run is
+# the official multi-task DAB shape, a notebook prep task before the
+# ai_runtime_task, and must be listed from its AIR task. The third run has no AI
+# task and must be filtered out.
 [[Server]]
 Pattern = "GET /api/2.2/jobs/runs/list"
 Response.Body = '''
@@ -32,6 +34,34 @@ Response.Body = '''
         }
       }]
     },
+    {
+      "run_id": 334747067049500,
+      "run_name": "air-byot-multitask-dag",
+      "creator_user_name": "[email protected]",
+      "start_time": 1717608700000,
+      "end_time": 1717608730000,
+      "state": {"life_cycle_state": "TERMINATED", "result_state": "SUCCESS"},
+      "tasks": [
+        {
+          "task_key": "prep",
+          "run_id": 334747067049501,
+          "start_time": 1717608700000,
+          "end_time": 1717608705000,
+          "notebook_task": {"notebook_path": "/Workspace/Users/[email protected]/prep"}
+        },
+        {
+          "task_key": "train",
+          "run_id": 334747067049502,
+          "start_time": 1717608710000,
+          "end_time": 1717608730000,
+          "depends_on": [{"task_key": "prep"}],
+          "ai_runtime_task": {
+            "experiment": "air-byot-multitask-dag",
+            "deployments": [{"compute": {"accelerator_type": "GPU_1xA10", "accelerator_count": 1}}]
+          }
+        }
+      ]
+    },
     {
       "run_id": 999000999000,
       "run_name": "not-an-air-run",

cmd/air/cancel.go

@@ -106,6 +106,9 @@ all active runs submitted by the current user.
 					fmt.Errorf("failed to list active runs: %w", err))
 			}
 
+			// A job with several AI runtime tasks lists one row per task; the
+			// run is cancelled (and previewed) once.
+			rows = firstRowPerRun(rows)
 			runIDs = make([]string, 0, len(rows))
 			for i := range rows {
 				if rows[i].RunID != "" {
@@ -215,6 +218,20 @@ func cancelRun(ctx context.Context, w *databricks.WorkspaceClient, rid string, s
 	return err
 }
 
+// firstRowPerRun keeps the first row of each run id, in order.
+func firstRowPerRun(rows []listRow) []listRow {
+	seen := make(map[string]bool, len(rows))
+	unique := make([]listRow, 0, len(rows))
+	for _, row := range rows {
+		if seen[row.RunID] {
+			continue
+		}
+		seen[row.RunID] = true
+		unique = append(unique, row)
+	}
+	return unique
+}
+
 // displayCancelPreview shows the runs that `cancel --all` is about to terminate.
 func displayCancelPreview(ctx context.Context, rows []listRow, host string) {
 	var sb strings.Builder