ep cli: More substatus granularity (#46266)

Agus Zubiaga created 4 days ago

Helps narrow down steps that are taking too long

Release Notes:

- N/A

Change summary

crates/edit_prediction_cli/src/format_prompt.rs | 6 +++++-
crates/edit_prediction_cli/src/load_project.rs  | 2 ++
crates/edit_prediction_cli/src/predict.rs       | 7 +++++--
crates/edit_prediction_cli/src/score.rs         | 4 +++-
4 files changed, 15 insertions(+), 4 deletions(-)

Detailed changes

crates/edit_prediction_cli/src/format_prompt.rs

@@ -23,10 +23,11 @@ pub async fn run_format_prompt(
 ) -> Result<()> {
     run_context_retrieval(example, app_state.clone(), cx.clone()).await?;
 
-    let _step_progress = Progress::global().start(Step::FormatPrompt, &example.spec.name);
+    let step_progress = Progress::global().start(Step::FormatPrompt, &example.spec.name);
 
     match prompt_format {
         PromptFormat::Teacher => {
+            step_progress.set_substatus("formatting teacher prompt");
             let prompt = TeacherPrompt::format_prompt(example);
             example.prompt = Some(ExamplePrompt {
                 input: prompt,
@@ -40,8 +41,11 @@ pub async fn run_format_prompt(
             });
         }
         PromptFormat::Zeta2 => {
+            step_progress.set_substatus("loading project");
             run_load_project(example, app_state, cx.clone()).await?;
 
+            step_progress.set_substatus("formatting zeta2 prompt");
+
             let ep_store = cx.update(|cx| {
                 EditPredictionStore::try_global(cx).context("EditPredictionStore not initialized")
             })??;

crates/edit_prediction_cli/src/load_project.rs

@@ -27,8 +27,10 @@ pub async fn run_load_project(
 
     let project = setup_project(example, &app_state, &progress, &mut cx).await?;
 
+    progress.set_substatus("applying edit history");
     let open_buffers = apply_edit_history(example, &project, &mut cx).await?;
 
+    progress.set_substatus("resolving cursor");
     let (buffer, cursor_position) =
         cursor_position(example, &project, &open_buffers, &mut cx).await?;
     let (example_buffer, language_name) = buffer.read_with(&cx, |buffer, _cx| {

crates/edit_prediction_cli/src/predict.rs

@@ -56,12 +56,13 @@ pub async fn run_prediction(
 
     run_load_project(example, app_state.clone(), cx.clone()).await?;
 
-    let _step_progress = Progress::global().start(Step::Predict, &example.spec.name);
+    let step_progress = Progress::global().start(Step::Predict, &example.spec.name);
 
     if matches!(
         provider,
         PredictionProvider::Zeta1 | PredictionProvider::Zeta2
     ) {
+        step_progress.set_substatus("authenticating");
         static AUTHENTICATED: OnceLock<Shared<Task<()>>> = OnceLock::new();
         AUTHENTICATED
             .get_or_init(|| {
@@ -93,6 +94,7 @@ pub async fn run_prediction(
         };
         store.set_edit_prediction_model(model);
     })?;
+    step_progress.set_substatus("configuring model");
     let state = example.state.as_ref().context("state must be set")?;
     let run_dir = RUN_DIR.join(&example.spec.name);
 
@@ -173,6 +175,7 @@ pub async fn run_prediction(
                 provider,
             });
 
+        step_progress.set_substatus("requesting prediction");
         let prediction = ep_store
             .update(&mut cx, |store, cx| {
                 store.request_prediction(
@@ -210,7 +213,7 @@ pub async fn run_prediction(
             } else {
                 ("no prediction", InfoStyle::Warning)
             };
-            _step_progress.set_info(info, style);
+            step_progress.set_info(info, style);
         }
     }
 

crates/edit_prediction_cli/src/score.rs

@@ -26,8 +26,9 @@ pub async fn run_scoring(
     )
     .await?;
 
-    let _progress = Progress::global().start(Step::Score, &example.spec.name);
+    let progress = Progress::global().start(Step::Score, &example.spec.name);
 
+    progress.set_substatus("applying patches");
     let original_text = &example.buffer.as_ref().unwrap().content;
     let expected_texts: Vec<String> = example
         .spec
@@ -39,6 +40,7 @@ pub async fn run_scoring(
         })
         .collect::<Result<Vec<_>, _>>()?;
 
+    progress.set_substatus("computing metrics");
     let mut scores = vec![];
     for prediction in &example.predictions {
         let actual_text = match apply_diff_to_string(&prediction.actual_patch, original_text) {