common : allow concurrent downloads across processes

common/CMakeLists.txt

@@ -78,6 +78,8 @@ add_library(${TARGET}
     download.h
     fit.cpp
     fit.h
+    file-lock.cpp
+    file-lock.h
     hf-cache.cpp
     hf-cache.h
     http.h
@@ -184,3 +186,8 @@ if (LLAMA_LLGUIDANCE)
 endif()
 
 target_link_libraries(${TARGET} PUBLIC llama Threads::Threads)
+
+if (CMAKE_SYSTEM_NAME MATCHES "AIX")
+    # AIX's flock() function comes from libbsd.a
+    target_link_libraries(${TARGET} PRIVATE -lbsd)
+endif()

common/arg.cpp

@@ -467,6 +467,14 @@ void common_models_handler_apply(common_models_handler & handler, common_params
     auto opts = handler.opts; // copy
     opts.callback = callback;
 
+    // another process may hold the cache lock during finalization:
+    // let the download callback cancel that wait too
+    auto finalize = [callback](const hf_cache::hf_file & f) {
+        return hf_cache::finalize_file(f, [callback]() {
+            return !(callback && callback->is_cancelled());
+        });
+    };
+
     // handle plain "url" if needed
     auto handle_url = [&](common_params_model & model) {
         if (!model.url.empty()) {
@@ -515,7 +523,7 @@ void common_models_handler_apply(common_models_handler & handler, common_params
     }
 
     // handle hf_plan tasks
-    auto add_tasks = [&opts, &tasks](const hf_cache::hf_files  & model_files,
+    auto add_tasks = [&opts, &tasks, &finalize](const hf_cache::hf_files  & model_files,
                                     const hf_cache::hf_file    & primary,
                                     common_params_model        & model) {
         for (size_t i = 0; i < model_files.size(); ++i) {
@@ -524,9 +532,9 @@ void common_models_handler_apply(common_models_handler & handler, common_params
             tasks.emplace_back(model_file, opts, [&, is_primary]() {
                 if (is_primary) {
                     // the primary file is the first split (00001-of), use it as model path
-                    model.path = hf_cache::finalize_file(model_file);
+                    model.path = finalize(model_file);
                 } else {
-                    hf_cache::finalize_file(model_file);
+                    finalize(model_file);
                 }
             });
         }
@@ -578,39 +586,39 @@ void common_models_handler_apply(common_models_handler & handler, common_params
         tasks.emplace_back(plan_spec.mtp, opts, [&]() {
             // only use the discovered MTP head when no draft path is set yet
             if (params.speculative.draft.mparams.path.empty()) {
-                params.speculative.draft.mparams.path = hf_cache::finalize_file(plan_spec.mtp);
+                params.speculative.draft.mparams.path = finalize(plan_spec.mtp);
             } else {
-                hf_cache::finalize_file(plan_spec.mtp);
+                finalize(plan_spec.mtp);
             }
         });
     }
     if (!plan_spec.dflash.local_path.empty() && !had_spec_url) {
         tasks.emplace_back(plan_spec.dflash, opts, [&]() {
             // only use the discovered DFlash sidecar when no draft path is set yet
             if (params.speculative.draft.mparams.path.empty()) {
-                params.speculative.draft.mparams.path = hf_cache::finalize_file(plan_spec.dflash);
+                params.speculative.draft.mparams.path = finalize(plan_spec.dflash);
             } else {
-                hf_cache::finalize_file(plan_spec.dflash);
+                finalize(plan_spec.dflash);
             }
         });
     }
     if (!plan_spec.eagle3.local_path.empty() && !had_spec_url) {
         tasks.emplace_back(plan_spec.eagle3, opts, [&]() {
             // only use the discovered Eagle3 sidecar when no draft path is set yet
             if (params.speculative.draft.mparams.path.empty()) {
-                params.speculative.draft.mparams.path = hf_cache::finalize_file(plan_spec.eagle3);
+                params.speculative.draft.mparams.path = finalize(plan_spec.eagle3);
             } else {
-                hf_cache::finalize_file(plan_spec.eagle3);
+                finalize(plan_spec.eagle3);
             }
         });
     }
     if (!plan_spec.dspark.local_path.empty() && !had_spec_url) {
         tasks.emplace_back(plan_spec.dspark, opts, [&]() {
             // only use the discovered DSpark sidecar when no draft path is set yet
             if (params.speculative.draft.mparams.path.empty()) {
-                params.speculative.draft.mparams.path = hf_cache::finalize_file(plan_spec.dspark);
+                params.speculative.draft.mparams.path = finalize(plan_spec.dspark);
             } else {
-                hf_cache::finalize_file(plan_spec.dspark);
+                finalize(plan_spec.dspark);
             }
         });
     }
@@ -631,54 +639,54 @@ void common_models_handler_apply(common_models_handler & handler, common_params
     }
     if (!plan.mmproj.local_path.empty()) {
         tasks.emplace_back(plan.mmproj, opts, [&]() {
-            params.mmproj.path = hf_cache::finalize_file(plan.mmproj);
+            params.mmproj.path = finalize(plan.mmproj);
         });
     }
     if (!plan.mtp.local_path.empty() && !had_spec_url) {
         tasks.emplace_back(plan.mtp, opts, [&]() {
             // only fall back to the discovered MTP head when no draft was explicitly provided
             if (params.speculative.draft.mparams.empty()) {
-                params.speculative.draft.mparams.path = hf_cache::finalize_file(plan.mtp);
+                params.speculative.draft.mparams.path = finalize(plan.mtp);
             } else {
-                hf_cache::finalize_file(plan.mtp);
+                finalize(plan.mtp);
             }
         });
     }
     if (!plan.dflash.local_path.empty() && !had_spec_url) {
         tasks.emplace_back(plan.dflash, opts, [&]() {
             // only fall back to the discovered DFlash sidecar when no draft was explicitly provided
             if (params.speculative.draft.mparams.empty()) {
-                params.speculative.draft.mparams.path = hf_cache::finalize_file(plan.dflash);
+                params.speculative.draft.mparams.path = finalize(plan.dflash);
             } else {
-                hf_cache::finalize_file(plan.dflash);
+                finalize(plan.dflash);
             }
         });
     }
     if (!plan.eagle3.local_path.empty() && !had_spec_url) {
         tasks.emplace_back(plan.eagle3, opts, [&]() {
             // only fall back to the discovered Eagle3 sidecar when no draft was explicitly provided
             if (params.speculative.draft.mparams.empty()) {
-                params.speculative.draft.mparams.path = hf_cache::finalize_file(plan.eagle3);
+                params.speculative.draft.mparams.path = finalize(plan.eagle3);
             } else {
-                hf_cache::finalize_file(plan.eagle3);
+                finalize(plan.eagle3);
             }
         });
     }
     if (!plan.dspark.local_path.empty() && !had_spec_url) {
         tasks.emplace_back(plan.dspark, opts, [&]() {
             // only fall back to the discovered DSpark sidecar when no draft was explicitly provided
             if (params.speculative.draft.mparams.empty()) {
-                params.speculative.draft.mparams.path = hf_cache::finalize_file(plan.dspark);
+                params.speculative.draft.mparams.path = finalize(plan.dspark);
             } else {
-                hf_cache::finalize_file(plan.dspark);
+                finalize(plan.dspark);
             }
         });
     }
     if (!plan.preset.local_path.empty()) {
         tasks.emplace_back(plan.preset, opts, [&]() {
             // if HF repo is a preset repo, we simply run server in router mode with the preset.ini file
             params.models_preset_hf = params.model.hf_repo; // only for showing a warning
-            params.models_preset    = hf_cache::finalize_file(plan.preset);
+            params.models_preset    = finalize(plan.preset);
             params.model = common_params_model{}; // make sure to clear model, so server starts in router mode
         });
     }

common/download.cpp

@@ -4,6 +4,7 @@
 #include "common.h"
 #include "log.h"
 #include "download.h"
+#include "file-lock.h"
 #include "hf-cache.h"
 #include "json.h"
 
@@ -111,6 +112,18 @@ std::pair<std::string, std::string> common_download_split_repo_tag(const std::st
     return {hf_repo, tag};
 }
 
+// file name from url, without the query string
+static std::string url_get_filename(const std::string & url) {
+    std::string filename = url;
+    if (auto pos = filename.rfind('/'); pos != std::string::npos) {
+        filename = filename.substr(pos + 1);
+    }
+    if (auto pos = filename.find('?'); pos != std::string::npos) {
+        filename = filename.substr(0, pos);
+    }
+    return filename;
+}
+
 class ProgressBar : public common_download_callback {
     static inline std::mutex mutex;
     static inline std::map<const ProgressBar *, int> lines;
@@ -138,14 +151,8 @@ class ProgressBar : public common_download_callback {
     ProgressBar() = default;
 
     void on_start(const common_download_progress & p) override {
-        filename = p.url;
+        filename = url_get_filename(p.url);
 
-        if (auto pos = filename.rfind('/'); pos != std::string::npos) {
-            filename = filename.substr(pos + 1);
-        }
-        if (auto pos = filename.find('?'); pos != std::string::npos) {
-            filename = filename.substr(0, pos);
-        }
         for (size_t i = 0; i < filename.size(); ++i) {
             if ((filename[i] & 0xC0) != 0x80) {
                 if (len++ == 39) {
@@ -282,15 +289,61 @@ static bool common_pull_file(httplib::Client & cli,
 static int common_download_file_single_online(const std::string & url,
                                               const std::string & path,
                                               const common_download_opts & opts,
-                                              bool skip_etag) {
+                                              bool skip_etag,
+                                              const std::string & lock_path,
+                                              const std::string & alt_path) {
     static const int max_attempts        = 3;
     static const int retry_delay_seconds = 2;
 
-    const bool file_exists = std::filesystem::exists(path);
+    const std::string path_temporary = path + ".downloadInProgress";
+    const std::string path_lock = lock_path.empty() ? path + ".lock" : lock_path;
+
+    common_download_progress p;
+    p.url = url;
+
+    if (opts.callback) {
+        opts.callback->on_start(p);
+    }
+
+    // the completion callback may itself take the same lock (the HF cache
+    // finalizes blobs under it), so every exit goes through finish()
+    common_file_lock lock(path_lock);
+    auto finish = [&](int status) {
+        lock.close();
+        if (opts.callback) {
+            opts.callback->on_done(p, is_http_status_ok(status));
+        }
+        return status;
+    };
+
+    // alt_path is the HF snapshot: the blob may have been moved there when symlinks are unavailable
+    auto cache_hit = [&] {
+        return std::filesystem::exists(path) || (!alt_path.empty() && std::filesystem::exists(alt_path));
+    };
+
+    if (skip_etag && cache_hit()) {
+        LOG_DBG("%s: using cached file: %s\n", __func__, path.c_str());
+        return finish(304); // 304 Not Modified - fake cached response
+    }
+
+    // take the lock before touching the cache, so two processes cannot race on the same file
+    bool waited = false;
+    if (!lock.acquire([&] {
+            if (!waited) {
+                LOG_INF("%s: file '%s' is being downloaded by another process, waiting...\n", __func__, url_get_filename(url).c_str());
+                waited = true;
+            }
+            return !(opts.callback && opts.callback->is_cancelled());
+        })) {
+        return finish(-1);
+    }
+
+    // under the lock: another process may have finished while we waited
+    bool file_exists = cache_hit();
 
     if (file_exists && skip_etag) {
         LOG_DBG("%s: using cached file: %s\n", __func__, path.c_str());
-        return 304; // 304 Not Modified - fake cached response
+        return finish(304); // 304 Not Modified - fake cached response
     }
 
     auto [cli, parts] = common_http_client(url);
@@ -319,18 +372,16 @@ static int common_download_file_single_online(const std::string & url,
         LOG_TRC("%s: HEAD failed, status: %d\n", __func__, head ? head->status : -1);
         if (file_exists) {
             LOG_TRC("%s: using cached file (HEAD failed): %s\n", __func__, path.c_str());
-            return 304; // 304 Not Modified - fake cached response
+            return finish(304); // 304 Not Modified - fake cached response
         }
-        return head ? head->status : -1;
+        return finish(head ? head->status : -1);
     }
 
     std::string etag;
     if (head->has_header("ETag")) {
         etag = head->get_header_value("ETag");
     }
 
-    common_download_progress p;
-    p.url = url;
     if (head->has_header("Content-Length")) {
         try {
             p.total = std::stoull(head->get_header_value("Content-Length"));
@@ -347,16 +398,16 @@ static int common_download_file_single_online(const std::string & url,
     if (file_exists) {
         if (etag.empty()) {
             LOG_DBG("%s: using cached file (no server etag): %s\n", __func__, path.c_str());
-            return 304; // 304 Not Modified - fake cached response
+            return finish(304); // 304 Not Modified - fake cached response
         }
         if (!last_etag.empty() && last_etag == etag) {
             LOG_DBG("%s: using cached file (same etag): %s\n", __func__, path.c_str());
-            return 304; // 304 Not Modified - fake cached response
+            return finish(304); // 304 Not Modified - fake cached response
         }
         // pass this point, the file exists but is different from the server version, so we need to redownload it
         if (remove(path.c_str()) != 0) {
             LOG_ERR("%s: unable to delete file: %s\n", __func__, path.c_str());
-            return -1;
+            return finish(-1);
         }
     }
 
@@ -366,13 +417,8 @@ static int common_download_file_single_online(const std::string & url,
     }
 
     bool success = false;
-    const std::string path_temporary = path + ".downloadInProgress";
     int delay = retry_delay_seconds;
 
-    if (opts.callback) {
-        opts.callback->on_start(p);
-    }
-
     for (int i = 0; i < max_attempts; ++i) {
         if (opts.callback && opts.callback->is_cancelled()) {
             break;
@@ -413,21 +459,18 @@ static int common_download_file_single_online(const std::string & url,
         }
     }
 
-    if (opts.callback) {
-        opts.callback->on_done(p, success);
-    }
     if (opts.callback && opts.callback->is_cancelled() &&
         std::filesystem::exists(path_temporary)) {
         if (remove(path_temporary.c_str()) != 0) {
             LOG_ERR("%s: unable to delete temporary file: %s\n", __func__, path_temporary.c_str());
         }
     }
+
     if (!success) {
         LOG_ERR("%s: download failed after %d attempts\n", __func__, max_attempts);
-        return -1; // max attempts reached
     }
 
-    return head->status;
+    return finish(success ? head->status : -1);
 }
 
 std::pair<long, std::vector<char>> common_remote_get_content(const std::string          & url,
@@ -467,14 +510,16 @@ std::pair<long, std::vector<char>> common_remote_get_content(const std::string
 int common_download_file_single(const std::string & url,
                                 const std::string & path,
                                 const common_download_opts & opts,
-                                bool skip_etag) {
+                                bool skip_etag,
+                                const std::string & lock_path,
+                                const std::string & alt_path) {
     if (!opts.offline) {
         ProgressBar tty_cb;
         common_download_opts online_opts = opts;
         if (!online_opts.callback) {
             online_opts.callback = &tty_cb;
         }
-        return common_download_file_single_online(url, path, online_opts, skip_etag);
+        return common_download_file_single_online(url, path, online_opts, skip_etag, lock_path, alt_path);
     }
 
     if (!std::filesystem::exists(path)) {
@@ -808,7 +853,7 @@ void common_download_run_tasks(const std::vector<common_download_task> & tasks)
     for (const auto & task : tasks) {
         futures.push_back(std::async(std::launch::async,
             [&task]() {
-                return common_download_file_single(task.url, task.local_path, task.opts, task.is_hf);
+                return common_download_file_single(task.url, task.local_path, task.opts, task.is_hf, task.lock_path, task.alt_path);
             }
         ));
     }