diff --git a/experimental/air/cmd/run.go b/experimental/air/cmd/run.go index 2e78986912..45202a09b8 100644 --- a/experimental/air/cmd/run.go +++ b/experimental/air/cmd/run.go @@ -148,11 +148,19 @@ The path must be a separate argument: cobra reserves -h as a boolean, so if !jsonOut { out := cmd.OutOrStdout() + perNode, err := gpusPerNode(gpuType(cfg.Compute.AcceleratorType)) + if err != nil { + return err + } + monitoringMessage := "Monitoring run and streaming logs..." + if cfg.Compute.NumAccelerators > perNode { + monitoringMessage = fmt.Sprintf("Monitoring run and streaming logs from node 0 of %d...", cfg.Compute.NumAccelerators/perNode) + } // The MLflow links stream in via the logs below, so don't poll here. printSubmitResult(ctx, out, runIDStr, dashboardURL) // Separate the submit summary from the streamed logs. fmt.Fprintln(out) - fmt.Fprintln(out, "Monitoring run and streaming logs...") + fmt.Fprintln(out, monitoringMessage) printLogsDivider(ctx, out) return handleWatchResult(out, w.Config.Profile, runIDStr, runLogs(watchCtx, cmd, req)) } diff --git a/experimental/air/cmd/run_watch_test.go b/experimental/air/cmd/run_watch_test.go index 6034fbdb1d..3cf3005561 100644 --- a/experimental/air/cmd/run_watch_test.go +++ b/experimental/air/cmd/run_watch_test.go @@ -93,9 +93,9 @@ func watchServerMLflow(t *testing.T, resultState string) *httptest.Server { return srv } -func runWatchCmd(t *testing.T, out flags.Output, buf *bytes.Buffer, srvURL string) error { +func runWatchCmdWithConfig(t *testing.T, out flags.Output, buf *bytes.Buffer, srvURL, config string) error { t.Helper() - cfgPath := writeConfigFile(t, "run.yaml", minimalConfig) + cfgPath := writeConfigFile(t, "run.yaml", config) cmd := withOutput(newRunCommand(), out) require.NoError(t, cmd.Flags().Set("file", cfgPath)) require.NoError(t, cmd.Flags().Set("watch", "true")) @@ -107,6 +107,19 @@ func runWatchCmd(t *testing.T, out flags.Output, buf *bytes.Buffer, srvURL strin return cmd.RunE(cmd, nil) } +func runWatchCmd(t *testing.T, out flags.Output, buf *bytes.Buffer, srvURL string) error { + t.Helper() + return runWatchCmdWithConfig(t, out, buf, srvURL, minimalConfig) +} + +const multinodeWatchConfig = ` +experiment_name: my-run +command: python train.py +compute: + accelerator_type: GPU_8xH100 + num_accelerators: 16 +` + func TestRunWatchStreamsLogs(t *testing.T) { var buf bytes.Buffer err := runWatchCmd(t, flags.OutputText, &buf, watchServer(t, "SUCCESS").URL) @@ -116,6 +129,7 @@ func TestRunWatchStreamsLogs(t *testing.T) { assert.Contains(t, out, "Submitted workload with Job Run ID: 777") assert.Contains(t, out, "View job run at: ") assert.Contains(t, out, "Monitoring run and streaming logs...") + assert.NotContains(t, out, "from node 0") // A "Logs" divider separates the submit summary from the streamed logs. assert.Contains(t, out, "Logs") assert.Contains(t, out, "───") @@ -123,9 +137,17 @@ func TestRunWatchStreamsLogs(t *testing.T) { assert.Contains(t, out, "step 1\nstep 2") } +func TestRunWatchMultinodeIdentifiesStreamedNode(t *testing.T) { + var buf bytes.Buffer + err := runWatchCmdWithConfig(t, flags.OutputText, &buf, watchServer(t, "SUCCESS").URL, multinodeWatchConfig) + require.NoError(t, err) + + assert.Contains(t, buf.String(), "Monitoring run and streaming logs from node 0 of 2...") +} + func TestRunWatchJSONEmitsSubmittedThenLogs(t *testing.T) { var buf bytes.Buffer - err := runWatchCmd(t, flags.OutputJSON, &buf, watchServer(t, "SUCCESS").URL) + err := runWatchCmdWithConfig(t, flags.OutputJSON, &buf, watchServer(t, "SUCCESS").URL, multinodeWatchConfig) require.NoError(t, err) all := buf.String() @@ -137,6 +159,7 @@ func TestRunWatchJSONEmitsSubmittedThenLogs(t *testing.T) { assert.Contains(t, all, `"type":"STATUS"`) assert.Contains(t, all, `"type":"LOG"`) assert.Contains(t, all, `"line":"step 1"`) + assert.NotContains(t, all, "Monitoring run and streaming logs") // The last line is the closing terminal-status envelope carrying SUCCESS. assert.Contains(t, lines[len(lines)-1], `"status":"SUCCESS"`) assert.Contains(t, lines[len(lines)-1], `"run_id":"777"`)