Skip to content

Commit aa00911

Browse files
authored
common : add download cancellation and temp file cleanup (#21813)
Signed-off-by: Adrien Gallouët <angt@huggingface.co>
1 parent ce8fd4b commit aa00911

2 files changed

Lines changed: 13 additions & 0 deletions

File tree

common/download.cpp

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -258,6 +258,9 @@ static bool common_pull_file(httplib::Client & cli,
258258
if (progress_step >= p.total / 1000 || p.downloaded == p.total) {
259259
if (callback) {
260260
callback->on_update(p);
261+
if (callback->is_cancelled()) {
262+
return false;
263+
}
261264
}
262265
progress_step = 0;
263266
}
@@ -373,6 +376,9 @@ static int common_download_file_single_online(const std::string & url,
373376
}
374377

375378
for (int i = 0; i < max_attempts; ++i) {
379+
if (opts.callback && opts.callback->is_cancelled()) {
380+
break;
381+
}
376382
if (i) {
377383
LOG_WRN("%s: retrying after %d seconds...\n", __func__, delay);
378384
std::this_thread::sleep_for(std::chrono::seconds(delay));
@@ -412,6 +418,12 @@ static int common_download_file_single_online(const std::string & url,
412418
if (opts.callback) {
413419
opts.callback->on_done(p, success);
414420
}
421+
if (opts.callback && opts.callback->is_cancelled() &&
422+
std::filesystem::exists(path_temporary)) {
423+
if (remove(path_temporary.c_str()) != 0) {
424+
LOG_ERR("%s: unable to delete temporary file: %s\n", __func__, path_temporary.c_str());
425+
}
426+
}
415427
if (!success) {
416428
LOG_ERR("%s: download failed after %d attempts\n", __func__, max_attempts);
417429
return -1; // max attempts reached

common/download.h

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ class common_download_callback {
2121
virtual void on_start(const common_download_progress & p) = 0;
2222
virtual void on_update(const common_download_progress & p) = 0;
2323
virtual void on_done(const common_download_progress & p, bool ok) = 0;
24+
virtual bool is_cancelled() const { return false; }
2425
};
2526

2627
struct common_remote_params {

0 commit comments

Comments
 (0)