diff --git a/crates/ml/examples/download_baseline.rs b/crates/ml/examples/download_baseline.rs index fa47ebf51..cffe20a13 100644 --- a/crates/ml/examples/download_baseline.rs +++ b/crates/ml/examples/download_baseline.rs @@ -284,6 +284,9 @@ impl DownloadStats { } } +/// Maximum retries for transient download failures (connection resets, EOF). +const MAX_RETRIES: u32 = 3; + async fn download_quarter( client: &mut HistoricalClient, symbol: &str, @@ -307,28 +310,51 @@ async fn download_quarter( let dt_range = date_range(quarter.start, quarter.end)?; - let params = GetRangeToFileParams::builder() - .dataset(dataset.to_owned()) - .symbols(vec![symbol.to_owned()]) - .schema(schema) - .stype_in(SType::Parent) - .date_time_range(dt_range) - .path(file_path.clone()) - .build(); + let mut last_err = None; + for attempt in 0..=MAX_RETRIES { + if attempt > 0 { + // Remove partial file from previous failed attempt + if file_path.exists() { + let _ = fs::remove_file(&file_path); + } + let delay = std::time::Duration::from_secs(2u64.pow(attempt)); + println!( + " [RETRY] {} attempt {}/{} in {}s", + filename, + attempt + 1, + MAX_RETRIES + 1, + delay.as_secs() + ); + tokio::time::sleep(delay).await; + } - let _decoder = client - .timeseries() - .get_range_to_file(¶ms) - .await - .with_context(|| { - format!( - "Failed to download {} for {}", - quarter.label, symbol - ) - })?; + let params = GetRangeToFileParams::builder() + .dataset(dataset.to_owned()) + .symbols(vec![symbol.to_owned()]) + .schema(schema) + .stype_in(SType::Parent) + .date_time_range(dt_range.clone()) + .path(file_path.clone()) + .build(); - let meta = fs::metadata(&file_path).context("Failed to read downloaded file metadata")?; - Ok(meta.len()) + match client.timeseries().get_range_to_file(¶ms).await { + Ok(_decoder) => { + let meta = + fs::metadata(&file_path).context("Failed to read downloaded file metadata")?; + return Ok(meta.len()); + } + Err(e) => { + last_err = Some(anyhow::anyhow!( + "Failed to download {} for {}: {}", + quarter.label, + symbol, + e + )); + } + } + } + + Err(last_err.unwrap()) } /// Upload a single file to MinIO/S3 via rclone. @@ -598,7 +624,7 @@ async fn main() -> Result<()> { (false, false, size, quarter.label.clone()) } Err(e) => { - println!(" [FAIL] {} -- {}", filename, e); + println!(" [FAIL] {} -- {:?}", filename, e); (false, true, 0, quarter.label.clone()) } }