diff --git a/.nextchanges/air/7636.md b/.nextchanges/air/7636.md new file mode 100644 index 00000000000..47b37de5cf3 --- /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`. 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/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..f500c9af9b5 100644 --- a/cmd/air/get.go +++ b/cmd/air/get.go @@ -23,15 +23,21 @@ 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 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 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"` // 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 +192,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..b1b12ceed12 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) @@ -247,6 +259,88 @@ 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) { + 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 + 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 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 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", + run: succeededTaskRun, + }, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + 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 map[string]any `json:"data"` + } + require.NoError(t, json.Unmarshal([]byte(out), &got)) + 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) + }) + } +} + func TestGetTemplateSweep(t *testing.T) { out := renderGet(t, getData{ RunID: "456", @@ -324,6 +418,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 +441,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)) }