Skip to content

Commit 3f49035

Browse files
committed
common : add minimalist multi-thread progress bar
I intentionally kept the bar simple without specifying part numbers (which ultimately don't matter much) the only thing we care about is tracking progress. Signed-off-by: Adrien Gallouët <angt@huggingface.co>
1 parent 59d8d4e commit 3f49035

File tree

1 file changed

+23
-4
lines changed

1 file changed

+23
-4
lines changed

common/download.cpp

Lines changed: 23 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,8 @@
1212
#include <filesystem>
1313
#include <fstream>
1414
#include <future>
15+
#include <map>
16+
#include <mutex>
1517
#include <regex>
1618
#include <string>
1719
#include <thread>
@@ -486,17 +488,36 @@ static void print_progress(size_t current, size_t total) {
486488
return;
487489
}
488490

491+
static std::mutex mutex;
492+
static std::map<std::thread::id, int> lines;
493+
494+
std::lock_guard<std::mutex> lock(mutex);
495+
std::thread::id id = std::this_thread::get_id();
496+
497+
if (lines.find(id) == lines.end()) {
498+
lines[id] = lines.size();
499+
std::cout << "\n";
500+
}
501+
int lines_up = lines.size() - lines[id];
502+
489503
size_t width = 50;
490504
size_t pct = (100 * current) / total;
491505
size_t pos = (width * current) / total;
492506

493-
std::cout << "["
507+
std::cout << "\033[s";
508+
509+
if (lines_up > 0) {
510+
std::cout << "\033[" << lines_up << "A";
511+
}
512+
std::cout << "\033[2K\r["
494513
<< std::string(pos, '=')
495514
<< (pos < width ? ">" : "")
496515
<< std::string(width - pos, ' ')
497516
<< "] " << std::setw(3) << pct << "% ("
498517
<< current / (1024 * 1024) << " MB / "
499-
<< total / (1024 * 1024) << " MB)\r";
518+
<< total / (1024 * 1024) << " MB) "
519+
<< "\033[u";
520+
500521
std::cout.flush();
501522
}
502523

@@ -552,8 +573,6 @@ static bool common_pull_file(httplib::Client & cli,
552573
nullptr
553574
);
554575

555-
std::cout << "\n";
556-
557576
if (!res) {
558577
LOG_ERR("%s: error during download. Status: %d\n", __func__, res ? res->status : -1);
559578
return false;

0 commit comments

Comments
 (0)