From a0728db61bac6059b32269e8ed60c58b04c10bcb Mon Sep 17 00:00:00 2001 From: Max Brunsfeld Date: Tue, 20 Jan 2026 11:08:15 -0800 Subject: [PATCH] Add --offset flag to ep cli (#47175) Release Notes: - N/A --------- Co-authored-by: Ben Kunkle --- crates/edit_prediction_cli/src/main.rs | 130 ++++++++++-------- .../edit_prediction_cli/src/split_commit.rs | 4 +- 2 files changed, 77 insertions(+), 57 deletions(-) diff --git a/crates/edit_prediction_cli/src/main.rs b/crates/edit_prediction_cli/src/main.rs index 6d61401a798a8a5c465d65b0762a66d54046a7df..5b8d5ba4f6a96540072dcfed36c5f75cf14828c2 100644 --- a/crates/edit_prediction_cli/src/main.rs +++ b/crates/edit_prediction_cli/src/main.rs @@ -56,6 +56,8 @@ struct EpArgs { max_parallelism: usize, #[clap(long, global = true)] limit: Option, + #[clap(long, global = true)] + offset: Option, /// Filter examples by name #[clap(long, global = true)] name: Option, @@ -87,6 +89,8 @@ pub enum FailedHandling { Keep, /// Exclude failed examples from the main output Skip, + /// Skip writing files + SkipNoFiles, } const INPUTS_HELP: &str = r#" @@ -366,12 +370,6 @@ async fn load_examples( examples.retain(|example| example.spec.repository_url.contains(repo_filter)); } - if let Some(limit) = args.limit { - if examples.len() > limit { - examples.truncate(limit); - } - } - // Skip resume logic for --in-place since input and output are the same file, // which would incorrectly treat all input examples as already processed. if !args.in_place { @@ -380,6 +378,14 @@ async fn load_examples( } } + if let Some(offset) = args.offset { + examples.splice(0..offset, []); + } + + if let Some(limit) = args.limit { + examples.truncate(limit); + } + Progress::global().set_total_examples(examples.len()); Ok(examples) @@ -734,57 +740,71 @@ async fn handle_error( example: &Example, ) { Progress::global().increment_failed(); - let example_name = example.spec.filename(); - let failed_example_path = FAILED_EXAMPLES_DIR.join(format!("{}.json", example_name)); - app_state - .fs - .write( - &failed_example_path, - &serde_json::to_vec_pretty(&example).unwrap(), - ) - .await - .unwrap(); - let err_path = FAILED_EXAMPLES_DIR.join(format!("{}_err.txt", example_name)); - app_state - .fs - .write(&err_path, format!("{error:?}").as_bytes()) - .await - .unwrap(); - - let failed_jsonl_path = RUN_DIR.join("failed.jsonl"); - let mut file = OpenOptions::new() - .create(true) - .append(true) - .open(&failed_jsonl_path) - .expect("Failed to open failed.jsonl"); - writeln!(file, "{}", serde_json::to_string(example).unwrap()) - .expect("Failed to write to failed.jsonl"); - - let cursor_path = example - .repo_name() - .unwrap() - .worktree_path() - .join(&example.spec.cursor_path); - - let msg = format!( - indoc::indoc! {" + + let msg; + if !matches!(args.failed, FailedHandling::SkipNoFiles) { + let example_name = example.spec.filename(); + + let failed_example_path = FAILED_EXAMPLES_DIR.join(format!("{}.json", example_name)); + app_state + .fs + .write( + &failed_example_path, + &serde_json::to_vec_pretty(&example).unwrap(), + ) + .await + .unwrap(); + let err_path = FAILED_EXAMPLES_DIR.join(format!("{}_err.txt", example_name)); + app_state + .fs + .write(&err_path, format!("{error:?}").as_bytes()) + .await + .unwrap(); + + let failed_jsonl_path = RUN_DIR.join("failed.jsonl"); + let mut file = OpenOptions::new() + .create(true) + .append(true) + .open(&failed_jsonl_path) + .expect("Failed to open failed.jsonl"); + writeln!(file, "{}", serde_json::to_string(example).unwrap()) + .expect("Failed to write to failed.jsonl"); + + let cursor_path = example + .repo_name() + .unwrap() + .worktree_path() + .join(&example.spec.cursor_path); + msg = format!( + indoc::indoc! {" + While processing \"{}\": + + \x1b[31m{:?}\x1b[0m + + Example: \x1b[36m{}\x1b[0m + Error file: \x1b[36m{}\x1b[0m + Cursor file: \x1b[36m{}\x1b[0m + Re-run: cargo run -p edit_prediction_cli -- {} \x1b[36m{}\x1b[0m + "}, + example.spec.name, + error, + failed_example_path.display(), + err_path.display(), + cursor_path.display(), + command, + failed_example_path.display(), + ); + } else { + msg = format!( + indoc::indoc! {" While processing \"{}\": - \x1b[31m{:?}\x1b[0m - - Example: \x1b[36m{}\x1b[0m - Error file: \x1b[36m{}\x1b[0m - Cursor file: \x1b[36m{}\x1b[0m - Re-run: cargo run -p edit_prediction_cli -- {} \x1b[36m{}\x1b[0m - "}, - example.spec.name, - error, - failed_example_path.display(), - err_path.display(), - cursor_path.display(), - command, - failed_example_path.display(), - ); + \x1b[31m{:?}\x1b[0m + "}, + example.spec.name, error + ); + } + if args.failfast || failfast_on_single_example { Progress::global().finalize(); panic!("{}", msg); diff --git a/crates/edit_prediction_cli/src/split_commit.rs b/crates/edit_prediction_cli/src/split_commit.rs index a992f1af7183d12477d9e19642c7f27dfb3c4b44..b411be9b8e3a0d5f63472b13ce340727e48b2b3a 100644 --- a/crates/edit_prediction_cli/src/split_commit.rs +++ b/crates/edit_prediction_cli/src/split_commit.rs @@ -178,7 +178,7 @@ pub fn run_split_commit( e ); match failed { - FailedHandling::Skip => { + FailedHandling::Skip | FailedHandling::SkipNoFiles => { eprintln!("{}", err_msg); continue; } @@ -219,7 +219,7 @@ pub fn run_split_commit( e ); match failed { - FailedHandling::Skip => { + FailedHandling::Skip | FailedHandling::SkipNoFiles => { eprintln!("{}", err_msg); continue; }