From 6c1d51c08eb77a57c37973a310c671635e3cac6f Mon Sep 17 00:00:00 2001 From: dimk0 <4886215+dimk0@users.noreply.github.com> Date: Sat, 14 Nov 2026 23:42:16 +0000 Subject: [PATCH 1/4] air: include MLflow run and experiment IDs in `air get` JSON `air get -o json` exposed the run's MLflow run only through mlflow_url, so handing an AIR run to an MLflow-based next step (e.g. a registration job that takes the MLflow run ID) meant parsing the URL. Add mlflow_experiment_id and mlflow_run_id next to it; they are null exactly when mlflow_url is. Co-authored-by: Isaac --- acceptance/air/get-ai-runtime/output.txt | 2 + acceptance/air/get-pending/output.txt | 2 + acceptance/air/get/output.txt | 2 + cmd/air/get.go | 22 ++++++---- cmd/air/get_test.go | 54 ++++++++++++++++++++++++ 5 files changed, 73 insertions(+), 9 deletions(-) diff --git a/acceptance/air/get-ai-runtime/output.txt b/acceptance/air/get-ai-runtime/output.txt index c84683e3f3e..8dbc06f966c 100644 --- a/acceptance/air/get-ai-runtime/output.txt +++ b/acceptance/air/get-ai-runtime/output.txt @@ -50,6 +50,8 @@ MLflow URL: [DATABRICKS_URL]/ml/experiments/exp1/runs/run1 "experiment_name": "my-exp", "dashboard_url": "[DATABRICKS_URL]/jobs/runs/123?o=[WORKSPACE_ID]", "mlflow_url": "[DATABRICKS_URL]/ml/experiments/exp1/runs/run1/logs", + "mlflow_experiment_id": "exp1", + "mlflow_run_id": "run1", "est_remaining_seconds": null, "est_percent_complete": null } diff --git a/acceptance/air/get-pending/output.txt b/acceptance/air/get-pending/output.txt index 3a5f9a72eda..887a911a04a 100644 --- a/acceptance/air/get-pending/output.txt +++ b/acceptance/air/get-pending/output.txt @@ -12,6 +12,8 @@ "experiment_name": "my-exp", "dashboard_url": "[DATABRICKS_URL]/jobs/runs/123?o=[WORKSPACE_ID]", "mlflow_url": null, + "mlflow_experiment_id": null, + "mlflow_run_id": null, "status_message": "Waiting for accelerator compute capacity to become available...", "est_remaining_seconds": null, "est_percent_complete": null diff --git a/acceptance/air/get/output.txt b/acceptance/air/get/output.txt index 9c1fed61721..8d85243c45a 100644 --- a/acceptance/air/get/output.txt +++ b/acceptance/air/get/output.txt @@ -50,6 +50,8 @@ MLflow URL: [DATABRICKS_URL]/ml/experiments/exp1/runs/run1 "experiment_name": "my-exp", "dashboard_url": "[DATABRICKS_URL]/jobs/runs/123?o=[WORKSPACE_ID]", "mlflow_url": "[DATABRICKS_URL]/ml/experiments/exp1/runs/run1/logs", + "mlflow_experiment_id": "exp1", + "mlflow_run_id": "run1", "est_remaining_seconds": null, "est_percent_complete": null } diff --git a/cmd/air/get.go b/cmd/air/get.go index bac778ded51..66991d88689 100644 --- a/cmd/air/get.go +++ b/cmd/air/get.go @@ -23,15 +23,17 @@ import ( // 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"` - TerminationReason *string `json:"termination_reason,omitempty"` + 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"` + MLflowExperimentID *string `json:"mlflow_experiment_id"` + MLflowRunID *string `json:"mlflow_run_id"` + TerminationReason *string `json:"termination_reason,omitempty"` // StatusMessage is the progress message the text view shows as the status, // set only while the run is unfinished. StatusMessage *string `json:"status_message,omitempty"` @@ -186,6 +188,8 @@ JOB_RUN_ID is the job run ID returned by "databricks air run".`, if ids != nil { url := mlflowLogsURL(w.Config.Host, ids) data.MLflowURL = &url + data.MLflowExperimentID = &ids.ExperimentID + data.MLflowRunID = &ids.RunID // A remaining-time estimate only makes sense while the run is still // training; a terminal run is either done or stopped mid-progress. if isRunning(run) { diff --git a/cmd/air/get_test.go b/cmd/air/get_test.go index c6e984c3b44..ece0178ff58 100644 --- a/cmd/air/get_test.go +++ b/cmd/air/get_test.go @@ -247,6 +247,58 @@ func TestGetRunRejectsNonNumericWorkspaceID(t *testing.T) { assert.Equal(t, jsonError{Code: "INVALID_ARGS", Kind: "PERMANENT", Message: `configured workspace_id "not-a-number" is not a numeric workspace ID`}, got.Error) } +func TestGetRunJSONMLflowIDs(t *testing.T) { + tests := []struct { + name string + getOutput string + wantExperimentID *string + wantRunID *string + }{ + { + name: "set when the run has MLflow IDs", + getOutput: `{"ai_runtime_task_output": {"mlflow_experiment_id": "E1", "mlflow_run_id": "R1"}}`, + wantExperimentID: new("E1"), + wantRunID: new("R1"), + }, + { + name: "null when the run has no MLflow IDs", + getOutput: `{}`, + }, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/api/2.2/jobs/runs/get": + _, _ = w.Write([]byte(`{"run_id": 5, "state": {"life_cycle_state": "TERMINATED", "result_state": "SUCCESS"}, "tasks": [{"run_id": 50}]}`)) + case "/api/2.2/jobs/runs/get-output": + _, _ = w.Write([]byte(tc.getOutput)) + default: + // SDK config discovery. + _, _ = w.Write([]byte(`{}`)) + } + })) + t.Cleanup(srv.Close) + w, err := databricks.NewWorkspaceClient(&databricks.Config{Host: srv.URL, Token: "token", WorkspaceID: "67890"}) + require.NoError(t, err) + + out, err := runGetJSON(t, w) + require.NoError(t, err) + var got struct { + Data struct { + MLflowURL *string `json:"mlflow_url"` + MLflowExperimentID *string `json:"mlflow_experiment_id"` + MLflowRunID *string `json:"mlflow_run_id"` + } `json:"data"` + } + require.NoError(t, json.Unmarshal([]byte(out), &got)) + assert.Equal(t, tc.wantExperimentID, got.Data.MLflowExperimentID) + assert.Equal(t, tc.wantRunID, got.Data.MLflowRunID) + assert.Equal(t, tc.wantRunID == nil, got.Data.MLflowURL == nil, "the IDs must be null exactly when mlflow_url is") + }) + } +} + func TestGetTemplateSweep(t *testing.T) { out := renderGet(t, getData{ RunID: "456", @@ -324,6 +376,7 @@ func TestBuildGetDataTerminationReasonJSON(t *testing.T) { assert.JSONEq(t, `{ "run_id":"5","status":"FAILED","started_at":null,"duration_seconds":null, "attempt_number":0,"experiment_name":null,"dashboard_url":"","mlflow_url":null, + "mlflow_experiment_id":null,"mlflow_run_id":null, "termination_reason":"out of memory","est_remaining_seconds":null,"est_percent_complete":null }`, string(body)) @@ -346,6 +399,7 @@ func TestBuildGetDataWaitingForComputeStatus(t *testing.T) { assert.JSONEq(t, `{ "run_id":"8","status":"PENDING","started_at":null,"duration_seconds":null, "attempt_number":0,"experiment_name":null,"dashboard_url":"","mlflow_url":null, + "mlflow_experiment_id":null,"mlflow_run_id":null, "est_remaining_seconds":null,"est_percent_complete":null }`, string(body)) } From e21f48dc684fc26e702d7f759ff032dfdf90fc57 Mon Sep 17 00:00:00 2001 From: dimk0 <4886215+dimk0@users.noreply.github.com> Date: Sat, 14 Nov 2026 23:43:27 +0000 Subject: [PATCH 2/4] Add changelog entry Co-authored-by: Isaac --- .nextchanges/air/7636.md | 1 + 1 file changed, 1 insertion(+) create mode 100644 .nextchanges/air/7636.md diff --git a/.nextchanges/air/7636.md b/.nextchanges/air/7636.md new file mode 100644 index 00000000000..ea34773b921 --- /dev/null +++ b/.nextchanges/air/7636.md @@ -0,0 +1 @@ +* `air get -o json` now includes `mlflow_run_id` and `mlflow_experiment_id`, so a run can be handed to MLflow-based steps without parsing `mlflow_url`. ([#7636](https://github.com/databricks/cli/pull/7636)) From e1bcc8c0642c596fe2a8aacb8334c4ed67f0a08a Mon Sep 17 00:00:00 2001 From: dimk0 <4886215+dimk0@users.noreply.github.com> Date: Sun, 15 Nov 2026 22:16:01 +0000 Subject: [PATCH 3/4] air: document MLflow ID semantics in air get JSON and tighten the test Document the earlier-retry fallback and both meanings of null on mlflow_experiment_id/mlflow_run_id, and add the earlier-retry clause to the changelog. Rework TestGetRunJSONMLflowIDs on getRunServer: tell null from a missing key, assert mlflow_url, and cover an earlier retry, a get-output error and a running run. Co-authored-by: Isaac --- .nextchanges/air/7636.md | 2 +- cmd/air/get.go | 20 ++++--- cmd/air/get_test.go | 111 ++++++++++++++++++++++++++------------- 3 files changed, 87 insertions(+), 46 deletions(-) diff --git a/.nextchanges/air/7636.md b/.nextchanges/air/7636.md index ea34773b921..47b37de5cf3 100644 --- a/.nextchanges/air/7636.md +++ b/.nextchanges/air/7636.md @@ -1 +1 @@ -* `air get -o json` now includes `mlflow_run_id` and `mlflow_experiment_id`, so a run can be handed to MLflow-based steps without parsing `mlflow_url`. ([#7636](https://github.com/databricks/cli/pull/7636)) +* `air get -o json` now includes `mlflow_run_id` and `mlflow_experiment_id`, so a run can be handed to MLflow-based steps without parsing `mlflow_url`. For a finished run whose last retry never started an MLflow run, the IDs come from the newest earlier retry that has one. ([#7636](https://github.com/databricks/cli/pull/7636)) diff --git a/cmd/air/get.go b/cmd/air/get.go index 66991d88689..5a839303f4a 100644 --- a/cmd/air/get.go +++ b/cmd/air/get.go @@ -23,14 +23,18 @@ import ( // 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"` + 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"` + // MLflowExperimentID and MLflowRunID identify the MLflow run MLflowURL links + // to. For a finished run whose last retry never started an MLflow run, they + // come from the newest earlier retry that has one. Like MLflowURL, they are + // null both when there is no MLflow run yet and when a get-output call fails. MLflowExperimentID *string `json:"mlflow_experiment_id"` MLflowRunID *string `json:"mlflow_run_id"` TerminationReason *string `json:"termination_reason,omitempty"` diff --git a/cmd/air/get_test.go b/cmd/air/get_test.go index ece0178ff58..a02e5cba598 100644 --- a/cmd/air/get_test.go +++ b/cmd/air/get_test.go @@ -4,6 +4,7 @@ import ( "bytes" "encoding/json" "errors" + "maps" "net/http" "net/http/httptest" "strings" @@ -178,14 +179,25 @@ func TestGetRunNotFoundJSON(t *testing.T) { 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) } -// getRunServer serves a terminal runs/get and a /Me that reports workspace -// 12345 in the X-Databricks-Org-Id header. *meCalls counts /Me requests. -func getRunServer(t *testing.T, meCalls *atomic.Int32) *httptest.Server { +const succeededRun = `{"run_id": 5, "state": {"life_cycle_state": "TERMINATED", "result_state": "SUCCESS"}}` + +// getRunServer serves run from runs/get, outputs[task run ID] from +// runs/get-output (a 500 for a task run with no entry), and a /Me that reports +// workspace 12345 in the X-Databricks-Org-Id header. *meCalls counts /Me +// requests. +func getRunServer(t *testing.T, meCalls *atomic.Int32, run string, outputs map[string]string) *httptest.Server { t.Helper() srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { switch { case r.URL.Path == "/api/2.2/jobs/runs/get": - _, _ = w.Write([]byte(`{"run_id": 5, "state": {"life_cycle_state": "TERMINATED", "result_state": "SUCCESS"}}`)) + _, _ = w.Write([]byte(run)) + case r.URL.Path == "/api/2.2/jobs/runs/get-output": + output, ok := outputs[r.URL.Query().Get("run_id")] + if !ok { + w.WriteHeader(http.StatusInternalServerError) + output = `{"error_code": "INTERNAL_ERROR", "message": "get-output failed"}` + } + _, _ = w.Write([]byte(output)) case strings.HasSuffix(r.URL.Path, "/scim/v2/Me"): meCalls.Add(1) w.Header().Set("X-Databricks-Org-Id", "12345") @@ -212,7 +224,7 @@ func runGetJSON(t *testing.T, w *databricks.WorkspaceClient) (string, error) { func TestGetRunURLUsesConfiguredWorkspaceID(t *testing.T) { var meCalls atomic.Int32 - srv := getRunServer(t, &meCalls) + srv := getRunServer(t, &meCalls, succeededRun, nil) w, err := databricks.NewWorkspaceClient(&databricks.Config{Host: srv.URL, Token: "token", WorkspaceID: "67890"}) require.NoError(t, err) @@ -226,7 +238,7 @@ func TestGetRunURLLooksUpUnsetWorkspaceID(t *testing.T) { // The SDK also reads the workspace ID from the environment. t.Setenv("DATABRICKS_WORKSPACE_ID", "") var meCalls atomic.Int32 - srv := getRunServer(t, &meCalls) + srv := getRunServer(t, &meCalls, succeededRun, nil) out, err := runGetJSON(t, newTestWorkspaceClient(t, srv.URL)) require.NoError(t, err) @@ -236,7 +248,7 @@ func TestGetRunURLLooksUpUnsetWorkspaceID(t *testing.T) { func TestGetRunRejectsNonNumericWorkspaceID(t *testing.T) { var meCalls atomic.Int32 - srv := getRunServer(t, &meCalls) + srv := getRunServer(t, &meCalls, succeededRun, nil) w, err := databricks.NewWorkspaceClient(&databricks.Config{Host: srv.URL, Token: "token", WorkspaceID: "not-a-number"}) require.NoError(t, err) @@ -248,53 +260,78 @@ func TestGetRunRejectsNonNumericWorkspaceID(t *testing.T) { } func TestGetRunJSONMLflowIDs(t *testing.T) { + const ( + succeededTaskRun = `{"run_id": 5, "state": {"life_cycle_state": "TERMINATED", "result_state": "SUCCESS"}, "tasks": [{"run_id": 50}]}` + outputWithIDs = `{"ai_runtime_task_output": {"mlflow_experiment_id": "E1", "mlflow_run_id": "R1"}}` + ) tests := []struct { - name string - getOutput string - wantExperimentID *string - wantRunID *string + name string + run string + outputs map[string]string + // wantExperimentID and wantRunID are the IDs the JSON names, or empty + // when all three MLflow fields must be null. + wantExperimentID string + wantRunID string }{ { - name: "set when the run has MLflow IDs", - getOutput: `{"ai_runtime_task_output": {"mlflow_experiment_id": "E1", "mlflow_run_id": "R1"}}`, - wantExperimentID: new("E1"), - wantRunID: new("R1"), + name: "set for a finished run", + run: succeededTaskRun, + outputs: map[string]string{"50": outputWithIDs}, + wantExperimentID: "E1", + wantRunID: "R1", + }, + { + name: "set while the run is running", + run: `{"run_id": 5, "state": {"life_cycle_state": "RUNNING"}, "tasks": [{"run_id": 50}]}`, + outputs: map[string]string{"50": outputWithIDs}, + wantExperimentID: "E1", + wantRunID: "R1", + }, + { + name: "from an earlier retry when the last retry never started", + run: `{"run_id": 5, "state": {"life_cycle_state": "TERMINATED", "result_state": "FAILED"}, + "tasks": [{"task_key": "train", "run_id": 50}, {"task_key": "train", "run_id": 51, "attempt_number": 1}]}`, + outputs: map[string]string{ + "50": `{"ai_runtime_task_output": {"mlflow_experiment_id": "E1", "mlflow_run_id": "R0"}}`, + "51": `{}`, + }, + wantExperimentID: "E1", + wantRunID: "R0", + }, + { + name: "null when the run has no MLflow run", + run: succeededTaskRun, + outputs: map[string]string{"50": `{}`}, }, { - name: "null when the run has no MLflow IDs", - getOutput: `{}`, + // No output for task run 50, so its get-output fails. + name: "null when get-output fails", + run: succeededTaskRun, }, } for _, tc := range tests { t.Run(tc.name, func(t *testing.T) { - srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - switch r.URL.Path { - case "/api/2.2/jobs/runs/get": - _, _ = w.Write([]byte(`{"run_id": 5, "state": {"life_cycle_state": "TERMINATED", "result_state": "SUCCESS"}, "tasks": [{"run_id": 50}]}`)) - case "/api/2.2/jobs/runs/get-output": - _, _ = w.Write([]byte(tc.getOutput)) - default: - // SDK config discovery. - _, _ = w.Write([]byte(`{}`)) - } - })) - t.Cleanup(srv.Close) + srv := getRunServer(t, new(atomic.Int32), tc.run, tc.outputs) w, err := databricks.NewWorkspaceClient(&databricks.Config{Host: srv.URL, Token: "token", WorkspaceID: "67890"}) require.NoError(t, err) out, err := runGetJSON(t, w) require.NoError(t, err) + // Decode into a map, which tells a missing key from null. var got struct { - Data struct { - MLflowURL *string `json:"mlflow_url"` - MLflowExperimentID *string `json:"mlflow_experiment_id"` - MLflowRunID *string `json:"mlflow_run_id"` - } `json:"data"` + Data map[string]any `json:"data"` } require.NoError(t, json.Unmarshal([]byte(out), &got)) - assert.Equal(t, tc.wantExperimentID, got.Data.MLflowExperimentID) - assert.Equal(t, tc.wantRunID, got.Data.MLflowRunID) - assert.Equal(t, tc.wantRunID == nil, got.Data.MLflowURL == nil, "the IDs must be null exactly when mlflow_url is") + maps.DeleteFunc(got.Data, func(key string, _ any) bool { return !strings.HasPrefix(key, "mlflow_") }) + want := map[string]any{"mlflow_url": nil, "mlflow_experiment_id": nil, "mlflow_run_id": nil} + if tc.wantRunID != "" { + want = map[string]any{ + "mlflow_url": srv.URL + "/ml/experiments/" + tc.wantExperimentID + "/runs/" + tc.wantRunID + "/logs", + "mlflow_experiment_id": tc.wantExperimentID, + "mlflow_run_id": tc.wantRunID, + } + } + assert.Equal(t, want, got.Data) }) } } From 5a5e05d4fd7d922558e7eec80d3ad0390c940b11 Mon Sep 17 00:00:00 2001 From: dimk0 <4886215+dimk0@users.noreply.github.com> Date: Sun, 15 Nov 2026 22:45:33 +0000 Subject: [PATCH 4/4] air: cover an experiment-only MLflow output in the air get test - Add a case whose output has an experiment ID but no run ID; all three MLflow fields must be null. - Say the IDs are null whenever no MLflow run is found, not only in two listed cases. Co-authored-by: Isaac --- cmd/air/get.go | 2 +- cmd/air/get_test.go | 5 +++++ 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/cmd/air/get.go b/cmd/air/get.go index 5a839303f4a..f500c9af9b5 100644 --- a/cmd/air/get.go +++ b/cmd/air/get.go @@ -34,7 +34,7 @@ type getData struct { // MLflowExperimentID and MLflowRunID identify the MLflow run MLflowURL links // to. For a finished run whose last retry never started an MLflow run, they // come from the newest earlier retry that has one. Like MLflowURL, they are - // null both when there is no MLflow run yet and when a get-output call fails. + // null whenever no MLflow run is found, including when a get-output call fails. MLflowExperimentID *string `json:"mlflow_experiment_id"` MLflowRunID *string `json:"mlflow_run_id"` TerminationReason *string `json:"termination_reason,omitempty"` diff --git a/cmd/air/get_test.go b/cmd/air/get_test.go index a02e5cba598..b1b12ceed12 100644 --- a/cmd/air/get_test.go +++ b/cmd/air/get_test.go @@ -303,6 +303,11 @@ func TestGetRunJSONMLflowIDs(t *testing.T) { run: succeededTaskRun, outputs: map[string]string{"50": `{}`}, }, + { + name: "null when the output has only an experiment ID", + run: succeededTaskRun, + outputs: map[string]string{"50": `{"ai_runtime_task_output": {"mlflow_experiment_id": "E1"}}`}, + }, { // No output for task run 50, so its get-output fails. name: "null when get-output fails",