Skip to content

Commit 22e86f9

Browse files
committed
Add resumable downloads for llama-server model loading
- Implement resumable downloads in common_download_file_single function - Add detection of partial download files (.downloadInProgress) - Check server support for HTTP Range requests via Accept-Ranges header - Implement HTTP Range request with "bytes=<start>-" header - Open files in append mode when resuming vs create mode for new downloads Signed-off-by: Eric Curtin <[email protected]>
1 parent f161463 commit 22e86f9

File tree

1 file changed

+63
-19
lines changed

1 file changed

+63
-19
lines changed

common/arg.cpp

Lines changed: 63 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -57,12 +57,18 @@ static std::string read_file(const std::string & fname) {
5757
}
5858

5959
static void write_file(const std::string & fname, const std::string & content) {
60-
std::ofstream file(fname);
60+
const std::string fname_tmp = fname + ".tmp";
61+
std::ofstream file(fname_tmp);
6162
if (!file) {
6263
throw std::runtime_error(string_format("error: failed to open file '%s'\n", fname.c_str()));
6364
}
6465
file << content;
6566
file.close();
67+
68+
// Makes write atomic
69+
if (rename(fname_tmp.c_str(), fname.c_str()) != 0) {
70+
LOG_ERR("%s: unable to rename file: %s to %s\n", __func__, fname_tmp.c_str(), fname.c_str());
71+
}
6672
}
6773

6874
common_arg & common_arg::set_examples(std::initializer_list<enum llama_example> examples) {
@@ -220,12 +226,24 @@ struct curl_slist_ptr {
220226
#define CURL_MAX_RETRY 3
221227
#define CURL_RETRY_DELAY_SECONDS 2
222228

223-
static bool curl_perform_with_retry(const std::string & url, CURL * curl, int max_attempts, int retry_delay_seconds, const char * method_name) {
229+
static bool curl_perform_with_retry(const std::string & url,
230+
CURL * curl,
231+
int max_attempts,
232+
int retry_delay_seconds,
233+
const char * method_name,
234+
const std::string & path_temporary = "") {
224235
int remaining_attempts = max_attempts;
225236

226237
while (remaining_attempts > 0) {
227238
LOG_INF("%s: %s %s (attempt %d of %d)...\n", __func__ , method_name, url.c_str(), max_attempts - remaining_attempts + 1, max_attempts);
228239

240+
if (std::filesystem::exists(path_temporary)) {
241+
const long partial_size = static_cast<long>(std::filesystem::file_size(path_temporary));
242+
LOG_INF("%s: server supports range requests, resuming download from byte %ld\n", __func__, partial_size);
243+
const std::string range_str = std::to_string(partial_size) + "-";
244+
curl_easy_setopt(curl, CURLOPT_RANGE, range_str.c_str());
245+
}
246+
229247
CURLcode res = curl_easy_perform(curl);
230248
if (res == CURLE_OK) {
231249
return true;
@@ -246,15 +264,14 @@ static bool curl_perform_with_retry(const std::string & url, CURL * curl, int ma
246264

247265
// download one single file from remote URL to local path
248266
static bool common_download_file_single(const std::string & url, const std::string & path, const std::string & bearer_token, bool offline) {
249-
// Check if the file already exists locally
250-
auto file_exists = std::filesystem::exists(path);
251-
252267
// If the file exists, check its JSON metadata companion file.
253268
std::string metadata_path = path + ".json";
254269
nlohmann::json metadata; // TODO @ngxson : get rid of this json, use regex instead
255270
std::string etag;
256271
std::string last_modified;
257272

273+
// Check if the file already exists locally
274+
auto file_exists = std::filesystem::exists(path);
258275
if (file_exists) {
259276
if (offline) {
260277
LOG_INF("%s: using cached file (offline mode): %s\n", __func__, path.c_str());
@@ -289,6 +306,7 @@ static bool common_download_file_single(const std::string & url, const std::stri
289306
struct common_load_model_from_url_headers {
290307
std::string etag;
291308
std::string last_modified;
309+
std::string accept_ranges;
292310
};
293311

294312
common_load_model_from_url_headers headers;
@@ -328,7 +346,7 @@ static bool common_download_file_single(const std::string & url, const std::stri
328346
static std::regex header_regex("([^:]+): (.*)\r\n");
329347
static std::regex etag_regex("ETag", std::regex_constants::icase);
330348
static std::regex last_modified_regex("Last-Modified", std::regex_constants::icase);
331-
349+
static std::regex accept_ranges_regex("Accept-Ranges", std::regex_constants::icase);
332350
std::string header(buffer, n_items);
333351
std::smatch match;
334352
if (std::regex_match(header, match, header_regex)) {
@@ -338,6 +356,8 @@ static bool common_download_file_single(const std::string & url, const std::stri
338356
headers->etag = value;
339357
} else if (std::regex_match(key, match, last_modified_regex)) {
340358
headers->last_modified = value;
359+
} else if (std::regex_match(key, match, accept_ranges_regex)) {
360+
headers->accept_ranges = value;
341361
}
342362
}
343363
return n_items;
@@ -366,28 +386,48 @@ static bool common_download_file_single(const std::string & url, const std::stri
366386

367387
// if head_request_ok is false, we don't have the etag or last-modified headers
368388
// we leave should_download as-is, which is true if the file does not exist
389+
bool should_download_from_scratch = false;
369390
if (head_request_ok) {
370391
// check if ETag or Last-Modified headers are different
371392
// if it is, we need to download the file again
372393
if (!etag.empty() && etag != headers.etag) {
373394
LOG_WRN("%s: ETag header is different (%s != %s): triggering a new download\n", __func__, etag.c_str(), headers.etag.c_str());
374395
should_download = true;
396+
should_download_from_scratch = true;
375397
} else if (!last_modified.empty() && last_modified != headers.last_modified) {
376398
LOG_WRN("%s: Last-Modified header is different (%s != %s): triggering a new download\n", __func__, last_modified.c_str(), headers.last_modified.c_str());
377399
should_download = true;
400+
should_download_from_scratch = true;
378401
}
379402
}
380403

404+
const bool accept_ranges_supported = !headers.accept_ranges.empty() && headers.accept_ranges != "none";
381405
if (should_download) {
382-
std::string path_temporary = path + ".downloadInProgress";
383-
if (file_exists) {
406+
if (file_exists && !accept_ranges_supported) { // Resumable downloads not supported, delete and start again.
384407
LOG_WRN("%s: deleting previous downloaded file: %s\n", __func__, path.c_str());
385408
if (remove(path.c_str()) != 0) {
386409
LOG_ERR("%s: unable to delete file: %s\n", __func__, path.c_str());
387410
return false;
388411
}
389412
}
390413

414+
std::string path_temporary = path + ".downloadInProgress";
415+
if (should_download_from_scratch) {
416+
if (std::filesystem::exists(path_temporary)) {
417+
if (remove(path_temporary.c_str()) != 0) {
418+
LOG_ERR("%s: unable to delete file: %s\n", __func__, path_temporary.c_str());
419+
return false;
420+
}
421+
}
422+
423+
if (std::filesystem::exists(path)) {
424+
if (remove(path.c_str()) != 0) {
425+
LOG_ERR("%s: unable to delete file: %s\n", __func__, path.c_str());
426+
return false;
427+
}
428+
}
429+
}
430+
391431
// Set the output file
392432

393433
struct FILE_deleter {
@@ -396,7 +436,8 @@ static bool common_download_file_single(const std::string & url, const std::stri
396436
}
397437
};
398438

399-
std::unique_ptr<FILE, FILE_deleter> outfile(fopen(path_temporary.c_str(), "wb"));
439+
// Always open file in append mode could be resuming
440+
std::unique_ptr<FILE, FILE_deleter> outfile(fopen(path_temporary.c_str(), "ab"));
400441
if (!outfile) {
401442
LOG_ERR("%s: error opening local file for writing: %s\n", __func__, path.c_str());
402443
return false;
@@ -431,7 +472,19 @@ static bool common_download_file_single(const std::string & url, const std::stri
431472
// start the download
432473
LOG_INF("%s: trying to download model from %s to %s (server_etag:%s, server_last_modified:%s)...\n", __func__,
433474
llama_download_hide_password_in_url(url).c_str(), path.c_str(), headers.etag.c_str(), headers.last_modified.c_str());
434-
bool was_perform_successful = curl_perform_with_retry(url, curl.get(), CURL_MAX_RETRY, CURL_RETRY_DELAY_SECONDS, "GET");
475+
476+
// Write the updated JSON metadata file.
477+
metadata.update({
478+
{"url", url},
479+
{"etag", headers.etag},
480+
{"lastModified", headers.last_modified}
481+
});
482+
write_file(metadata_path, metadata.dump(4));
483+
LOG_DBG("%s: file metadata saved: %s\n", __func__, metadata_path.c_str());
484+
485+
const bool was_perform_successful =
486+
curl_perform_with_retry(url, curl.get(), CURL_MAX_RETRY, CURL_RETRY_DELAY_SECONDS, "GET",
487+
headers.accept_ranges.empty() ? "" : path_temporary);
435488
if (!was_perform_successful) {
436489
return false;
437490
}
@@ -446,15 +499,6 @@ static bool common_download_file_single(const std::string & url, const std::stri
446499
// Causes file to be closed explicitly here before we rename it.
447500
outfile.reset();
448501

449-
// Write the updated JSON metadata file.
450-
metadata.update({
451-
{"url", url},
452-
{"etag", headers.etag},
453-
{"lastModified", headers.last_modified}
454-
});
455-
write_file(metadata_path, metadata.dump(4));
456-
LOG_DBG("%s: file metadata saved: %s\n", __func__, metadata_path.c_str());
457-
458502
if (rename(path_temporary.c_str(), path.c_str()) != 0) {
459503
LOG_ERR("%s: unable to rename file: %s to %s\n", __func__, path_temporary.c_str(), path.c_str());
460504
return false;

0 commit comments

Comments
 (0)