Skip to content

Feat: improve downloads of models under unstable connections #549

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
stdevPavelmc wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Feat: improve downloads of models under unstable connections #549

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
92 changes: 78 additions & 14 deletions src/pull/download_model.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -52,6 +52,12 @@ std::string calculate_git_blob_oid(const std::string& file_path) {
// Global variable to track if progress bar was shown
static bool g_progress_bar_shown = false;

/// \brief Context passed to the progress callback
struct ProgressContext {
curl_off_t already_downloaded = 0; ///< Bytes already on disk before this session
std::function<void(double)>* user_cb = nullptr;
};

/// \brief Hide the cursor
void hide_cursor() {
std::cout << "\033[?25l" << std::flush;
Expand Down Expand Up @@ -99,14 +105,20 @@ int progress_callback(void* clientp, double dltotal, double dlnow, double ultota
auto now = Clock::now();
auto elapsed = std::chrono::duration_cast<std::chrono::milliseconds>(now - last_print_time);

double percentage = (dlnow / dltotal) * 100.0;

if (elapsed.count() >= 1000) {
utils::enable_ansi_on_windows_once();

double percentage = (dlnow / dltotal) * 100.0;
double mb_now = dlnow / 1024.0 / 1024.0;
double mb_total = dltotal / 1024.0 / 1024.0;
// Adjust for bytes already on disk from a previous partial download
curl_off_t offset = 0;
if (clientp) {
offset = static_cast<ProgressContext*>(clientp)->already_downloaded;
}
double total_now = dlnow + static_cast<double>(offset);
double total_total = dltotal + static_cast<double>(offset);

double percentage = (total_now / total_total) * 100.0;
double mb_now = total_now / 1024.0 / 1024.0;
double mb_total = total_total / 1024.0 / 1024.0;

std::cout << "\r\033[K"
<< "[FLM] Downloading: " << std::fixed << std::setprecision(1)
Expand Down Expand Up @@ -140,29 +152,49 @@ bool download_file(const std::string& url, const std::string& local_path, bool i
std::filesystem::path path(local_path);
std::filesystem::create_directories(path.parent_path());

FILE* fp = fopen(local_path.c_str(), "wb");
// Check for an existing partial download and resume from it
curl_off_t resume_from = 0;
if (std::filesystem::exists(local_path)) {
resume_from = static_cast<curl_off_t>(std::filesystem::file_size(local_path));
}

FILE* fp = fopen(local_path.c_str(), resume_from > 0 ? "ab" : "wb");
if (!fp) {
std::cerr << "Failed to open file for writing: " << local_path << std::endl;
curl_easy_cleanup(curl);
return false;
}

if (resume_from > 0) {
header_print("FLM", "Resuming from " << (resume_from / 1024.0 / 1024.0) << " MB");
}

// Hide cursor before starting download
hide_cursor();

ProgressContext prog_ctx;
prog_ctx.already_downloaded = resume_from;
prog_ctx.user_cb = &progress_cb;

curl_easy_setopt(curl, CURLOPT_URL, url.c_str());
curl_easy_setopt(curl, CURLOPT_WRITEFUNCTION, write_data_to_file);
curl_easy_setopt(curl, CURLOPT_WRITEDATA, fp);
curl_easy_setopt(curl, CURLOPT_FOLLOWLOCATION, 1L);
curl_easy_setopt(curl, CURLOPT_SSL_VERIFYPEER, 0L);
curl_easy_setopt(curl, CURLOPT_SSL_VERIFYHOST, 0L);
curl_easy_setopt(curl, CURLOPT_USERAGENT, "FastFlowLM/1.0");
curl_easy_setopt(curl, CURLOPT_TIMEOUT, 3600L); // 1 hour timeout
// Abort if transfer stalls below 1 byte/sec for 60 seconds
curl_easy_setopt(curl, CURLOPT_LOW_SPEED_LIMIT, 1L);
curl_easy_setopt(curl, CURLOPT_LOW_SPEED_TIME, 60L);
if (resume_from > 0) {
curl_easy_setopt(curl, CURLOPT_RESUME_FROM_LARGE, resume_from);
}

// Set progress callback if provided
if (progress_cb) {
curl_easy_setopt(curl, CURLOPT_NOPROGRESS, 0L);
curl_easy_setopt(curl, CURLOPT_PROGRESSFUNCTION, progress_callback);
curl_easy_setopt(curl, CURLOPT_PROGRESSDATA, &prog_ctx);
}

CURLcode res = curl_easy_perform(curl);
Expand All @@ -175,7 +207,7 @@ bool download_file(const std::string& url, const std::string& local_path, bool i

if (res != CURLE_OK) {
std::cerr << "CURL error: " << curl_easy_strerror(res) << std::endl;
std::filesystem::remove(local_path); // Remove partial download
// Keep the partial file so the next retry can resume from where we left off
return false;
}

Expand All @@ -187,7 +219,8 @@ bool download_file(const std::string& url, const std::string& local_path, bool i
header_print("FLM", "Checking Hash...");
std::string local_oid = is_lfs ? calculate_file_sha256(local_path) : calculate_git_blob_oid(local_path);
if (local_oid != remote_oid) {
header_print("FLM", "Hash not matched!");
header_print("FLM", "Hash not matched! Removing corrupted file.");
std::filesystem::remove(local_path); // Remove corrupted file so next retry starts fresh
show_cursor(); // Show cursor on error
return false;
}
Expand All @@ -197,17 +230,20 @@ bool download_file(const std::string& url, const std::string& local_path, bool i
}

static bool download_with_retry(const std::string& url, const std::string& local_path, bool is_lfs, std::string remote_oid,
std::function<void(double)> progress_cb, int max_retries = 3) {
std::function<void(double)> progress_cb, int max_retries = 10) {
int attempt = 0;
while (attempt < max_retries) {
if (download_file(url, local_path, is_lfs, remote_oid, progress_cb)) {
return true;
}
header_print("FLM", "Download failed (attempt " << (attempt + 1) << "/" << max_retries << ")");
if(attempt < max_retries - 1)
header_print("FLM", "Retrying...");
header_print("FLM", "Download failed (attempt " << (attempt + 1) << "/" << max_retries << ")");
attempt++;
std::this_thread::sleep_for(std::chrono::seconds(1));
if (attempt < max_retries) {
// Exponential backoff: 1s, 2s, 4s, 8s, ... capped at 30s
int wait_seconds = std::min(1 << (attempt - 1), 30);
header_print("FLM", "Retrying in " << wait_seconds << "s...");
std::this_thread::sleep_for(std::chrono::seconds(wait_seconds));
}
}

return false;
Expand Down Expand Up @@ -245,6 +281,34 @@ std::string download_string(const std::string& url) {
return response;
}

/// \brief Query the remote file size via a HEAD request
/// \param url the URL to query
/// \return the file size in bytes, or -1 on failure
curl_off_t get_remote_file_size(const std::string& url) {
CURL* curl = curl_easy_init();
if (!curl) {
return -1;
}

curl_easy_setopt(curl, CURLOPT_URL, url.c_str());
curl_easy_setopt(curl, CURLOPT_NOBODY, 1L); // HEAD-like: no body
curl_easy_setopt(curl, CURLOPT_FOLLOWLOCATION, 1L);
curl_easy_setopt(curl, CURLOPT_SSL_VERIFYPEER, 0L);
curl_easy_setopt(curl, CURLOPT_SSL_VERIFYHOST, 0L);
curl_easy_setopt(curl, CURLOPT_USERAGENT, "FastFlowLM/1.0");
curl_easy_setopt(curl, CURLOPT_TIMEOUT, 30L);

CURLcode res = curl_easy_perform(curl);

curl_off_t content_length = -1;
if (res == CURLE_OK) {
curl_easy_getinfo(curl, CURLINFO_CONTENT_LENGTH_DOWNLOAD_T, &content_length);
}

curl_easy_cleanup(curl);
return content_length;
}

/// \brief Download multiple files with progress tracking
/// \param downloads the downloads
/// \param progress_cb the progress callback
Expand Down
3 changes: 3 additions & 0 deletions src/pull/download_model.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,9 @@ bool download_file(const std::string& url, const std::string& local_path, bool i
// Download content from URL to a string
std::string download_string(const std::string& url);

// Query the remote file size via a HEAD request. Returns -1 on failure.
curl_off_t get_remote_file_size(const std::string& url);

// Download multiple files with progress tracking
bool download_multiple_files(const nlohmann::json downloads,
std::function<void(size_t, size_t)> progress_cb = nullptr);
Expand Down
59 changes: 38 additions & 21 deletions src/pull/model_downloader.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -88,23 +88,13 @@ bool ModelDownloader::pull_model(const std::string& model_tag, bool force_redown
header_print("FLM", "Model: " + new_model_tag);
header_print("FLM", "Name: " + model_name);

// Check if model is already downloaded
if (!force_redownload && is_model_downloaded(new_model_tag)) {
header_print("FLM", "Model already downloaded. Use --force to re-download.");
return true;
}

// If force, remove the model first
if (force_redownload) {
remove_model(new_model_tag);
}

// Get missing files
// Get missing files (by name only — partial files are handled by build_download_list)
auto missing_files = get_missing_files(new_model_tag);
if (missing_files.empty() && !force_redownload) {
header_print("FLM", "All files already present.");
return true;
}

if (!missing_files.empty()) {
header_print("FLM", "Missing files (" + std::to_string(missing_files.size()) + "):");
Expand All @@ -124,16 +114,17 @@ bool ModelDownloader::pull_model(const std::string& model_tag, bool force_redown
}
}

// Build download list
// Build download list (also detects and resumes partial downloads)
header_print("FLM", "Checking all files for missing or partial downloads...");
auto download_list = build_download_list(new_model_tag);
auto downloads = download_list.first;
float sum_fize_size = download_list.second;
if (downloads.empty()) {
header_print("FLM", "No files to download for model: " + new_model_tag);
header_print("FLM", "All files are complete. Nothing to download.");
return true; // Return true since all files are already present
}

header_print("FLM", "Downloading " + std::to_string(downloads.size()) + " missing files...");
header_print("FLM", "Downloading " + std::to_string(downloads.size()) + " file(s)...");

header_print("FLM", "Files to download (" << std::fixed << std::setprecision(2) << sum_fize_size << " MB): ");
for (const auto& download : downloads) {
Expand Down Expand Up @@ -308,14 +299,16 @@ std::pair<nlohmann::json, float> ModelDownloader::build_download_list(const std:
const auto& file = *it;
std::string local_path = get_model_file_path(model_path, filename);

if (!file_exists(local_path)) {
std::string url;
if (std::string(base_url).find("resolve") != std::string::npos) { // resolve provided , may from a specific branch
url = base_url + "/" + filename + "?download=true";
}
else {
url = base_url + "/resolve/main/" + filename + "?download=true";
// Build the download URL (needed for both new and partial files)
auto build_url = [&]() -> std::string {
if (std::string(base_url).find("resolve") != std::string::npos) {
return base_url + "/" + filename + "?download=true";
}
return base_url + "/resolve/main/" + filename + "?download=true";
};

if (!file_exists(local_path)) {
std::string url = build_url();
bool is_lfs = file.contains("lfs");
std::string oid = is_lfs ? file["lfs"]["oid"] : file["oid"];
float file_size = static_cast<float>(file["size"]) / 1024 / 1024;
Expand All @@ -330,6 +323,30 @@ std::pair<nlohmann::json, float> ModelDownloader::build_download_list(const std:
{"is_lfs", is_lfs},
};
downloads.push_back(entry);
} else {
// File exists — check if it is a partial download
curl_off_t remote_size = static_cast<curl_off_t>(file["size"]);
curl_off_t local_size = static_cast<curl_off_t>(std::filesystem::file_size(local_path));
if (local_size < remote_size) {
std::string url = build_url();
bool is_lfs = file.contains("lfs");
std::string oid = is_lfs ? file["lfs"]["oid"] : file["oid"];
float remaining_size = static_cast<float>(remote_size - local_size) / 1024 / 1024;
sum_file_size += remaining_size;
header_print("FLM", "Partial download detected: " + filename +
" (" + std::to_string(local_size / 1024 / 1024) + " MB / " +
std::to_string(remote_size / 1024 / 1024) + " MB)");

nlohmann::json entry = {
{"file", filename},
{"size", remaining_size},
{"url", url},
{"localpath", local_path},
{"oid", oid},
{"is_lfs", is_lfs},
};
downloads.push_back(entry);
}
}

}
Expand Down
26 changes: 5 additions & 21 deletions src/src/main.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -635,27 +635,11 @@ int main(int argc, char* argv[]) {
// server->stop();
}
else if (parsed_args.command == "pull") {
// Check if the model is already downloaded, if true, the model will not be downloaded
// Check if model is already downloaded
if (!parsed_args.force_redownload && downloader.is_model_downloaded(parsed_args.model_tag)) {
header_print("FLM", "Model is already downloaded.");
// Show missing files if any, this will be used to show the missing files
auto missing_files = downloader.get_missing_files(parsed_args.model_tag);
if (!missing_files.empty()) {
header_print("FLM", "Missing files:");
for (const auto& file : missing_files) {
std::cout << " - " << file << std::endl;
}
} else {
header_print("FLM", "All required files are present.");
}
} else {
// Download the model, this will be used to download the model
bool success = downloader.pull_model(parsed_args.model_tag, parsed_args.force_redownload);
if (!success) {
header_print("ERROR", "Failed to pull model: " + parsed_args.model_tag);
return 1;
}
// Always call pull_model on explicit pull; it will detect missing and partial files
bool success = downloader.pull_model(parsed_args.model_tag, parsed_args.force_redownload);
if (!success) {
header_print("ERROR", "Failed to pull model: " + parsed_args.model_tag);
return 1;
}
}
else if (parsed_args.command == "remove") {
Expand Down