diff --git a/common/CMakeLists.txt b/common/CMakeLists.txt index 2b307c59d32d..8e3c88dac9dc 100644 --- a/common/CMakeLists.txt +++ b/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() diff --git a/common/arg.cpp b/common/arg.cpp index 63e342776d54..6c834b40c651 100644 --- a/common/arg.cpp +++ b/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,9 +586,9 @@ 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); } }); } @@ -588,9 +596,9 @@ void common_models_handler_apply(common_models_handler & handler, common_params 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); } }); } @@ -598,9 +606,9 @@ void common_models_handler_apply(common_models_handler & handler, common_params 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); } }); } @@ -608,9 +616,9 @@ void common_models_handler_apply(common_models_handler & handler, common_params 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,16 +639,16 @@ 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); } }); } @@ -648,9 +656,9 @@ void common_models_handler_apply(common_models_handler & handler, common_params 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); } }); } @@ -658,9 +666,9 @@ void common_models_handler_apply(common_models_handler & handler, common_params 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); } }); } @@ -668,9 +676,9 @@ void common_models_handler_apply(common_models_handler & handler, common_params 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); } }); } @@ -678,7 +686,7 @@ void common_models_handler_apply(common_models_handler & handler, common_params 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 }); } diff --git a/common/download.cpp b/common/download.cpp index 4b28a708c86e..baf1b1038ab6 100644 --- a/common/download.cpp +++ b/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 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 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,9 +372,9 @@ 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; @@ -329,8 +382,6 @@ static int common_download_file_single_online(const std::string & url, 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> common_remote_get_content(const std::string & url, @@ -467,14 +510,16 @@ std::pair> 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 & 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); } )); } diff --git a/common/download.h b/common/download.h index 8c30cfc3eadb..cc479910e399 100644 --- a/common/download.h +++ b/common/download.h @@ -69,12 +69,15 @@ struct common_download_task { std::string local_path; std::function on_done; bool is_hf = false; + std::string lock_path; + std::string alt_path; // HF snapshot path, also treated as cached (the blob may be moved there) common_download_task() = default; common_download_task(hf_cache::hf_file f, const common_download_opts & opts, std::function on_done = nullptr) - : opts(opts), url(f.url), local_path(f.local_path), on_done(on_done), is_hf(true) {} + : opts(opts), url(f.url), local_path(f.local_path), on_done(on_done), is_hf(true), + lock_path(hf_cache::get_lock_path(f.repo_id, f.oid)), alt_path(f.final_path) {} }; void common_download_run_tasks(const std::vector & tasks); @@ -92,10 +95,14 @@ std::string common_download_resolve_path(const std::string & hf_repo_with_tag, c // download single file from url to local path // returns status code or -1 on error // skip_etag: if true, don't read/write .etag files (for HF cache where filename is the hash) +// lock_path: lock file used to serialize downloads of the same file; when empty, ".lock" is used +// alt_path: additional path treated as present in cache (see cache_hit in download.cpp) int common_download_file_single(const std::string & url, const std::string & path, const common_download_opts & opts = {}, - bool skip_etag = false); + bool skip_etag = false, + const std::string & lock_path = "", + const std::string & alt_path = ""); // resolve and download model from Docker registry // return local path to downloaded model file diff --git a/common/file-lock.cpp b/common/file-lock.cpp new file mode 100644 index 000000000000..553a06687f9e --- /dev/null +++ b/common/file-lock.cpp @@ -0,0 +1,179 @@ +#include "file-lock.h" + +#include "log.h" + +#include +#include +#include + +#include +#include + +#if defined(_WIN32) +#define WIN32_LEAN_AND_MEAN +#ifndef NOMINMAX +#define NOMINMAX +#endif +#include +#include +#else +#include +#include +#include +#include +#endif + +common_file_lock::common_file_lock(const std::string & lock_path) { + this->lock_path = lock_path; +} + +#if defined(_WIN32) +// the \\?\ prefix (for paths longer than 255 chars) is only accepted on absolute paths +// with backslash separators: normalize first, then check the length of the absolute form +static std::wstring as_extended_path(const std::string & path) { + // convert from the narrow encoding std::filesystem uses on this system + int wlen = MultiByteToWideChar(CP_ACP, 0, path.c_str(), (int) path.size(), nullptr, 0); + if (wlen <= 0) { + return std::wstring(path.begin(), path.end()); // best effort for exotic encodings + } + std::wstring wide(wlen, L'\0'); + MultiByteToWideChar(CP_ACP, 0, path.c_str(), (int) path.size(), wide.data(), wlen); + + if (wide.rfind(L"\\\\?\\", 0) == 0) { + return wide; + } + + // GetFullPathNameW resolves relative paths and normalizes separators + DWORD size = GetFullPathNameW(wide.c_str(), 0, nullptr, nullptr); + if (size == 0) { + return wide; + } + std::wstring full(size, L'\0'); + DWORD written = GetFullPathNameW(wide.c_str(), size, full.data(), nullptr); + if (written == 0 || written >= size) { + return wide; + } + full.resize(written); + + if (full.size() <= 255) { + return wide; // short enough: open the path exactly as given + } + if (full.rfind(L"\\\\", 0) == 0) { + return L"\\\\?\\UNC\\" + full.substr(2); // \\server\share -> \\?\UNC\server\share + } + return L"\\\\?\\" + full; +} +#endif + +common_file_lock::status common_file_lock::try_lock() { + if (lock_path.empty()) { + return status::acquired; // no lock path, nothing to lock + } +#if defined(_WIN32) + if (fd < 0) { + // opened with a wide path: extended paths cannot go through the narrow CRT + std::wstring lock_path_w = as_extended_path(lock_path); + fd = _wopen(lock_path_w.c_str(), _O_RDWR | _O_CREAT | _O_BINARY | _O_NOINHERIT, _S_IREAD | _S_IWRITE); + if (fd < 0) { + LOG_ERR("%s: cannot open lock file: %s\n", __func__, lock_path.c_str()); + return status::error; + } + } + // lock the first byte, the file itself stays empty + OVERLAPPED ov = {}; + locked = LockFileEx((HANDLE) _get_osfhandle(fd), + LOCKFILE_EXCLUSIVE_LOCK | LOCKFILE_FAIL_IMMEDIATELY, 0, 1, 0, &ov); + if (locked) { + return status::acquired; + } + if (GetLastError() == ERROR_LOCK_VIOLATION) { + return status::busy; + } + LOG_ERR("%s: file locking failed (%d)\n", __func__, (int) GetLastError()); + return status::error; +#else + if (fd < 0) { + int open_flags = O_RDWR | O_CREAT | O_CLOEXEC; +#ifdef O_NOFOLLOW + open_flags |= O_NOFOLLOW; +#endif + // group-writable, so users sharing a cache can wait on the lock + fd = open(lock_path.c_str(), open_flags, 0664); + if (fd < 0) { + LOG_ERR("%s: cannot open lock file: %s\n", __func__, lock_path.c_str()); + return status::error; + } + // open() filters the mode through the umask + (void) fchmod(fd, 0664); + } + + if (flock(fd, LOCK_EX | LOCK_NB) == 0) { + locked = true; + return status::acquired; + } + // EACCES is reported as contention on some filesystems + if (errno == EWOULDBLOCK || errno == EAGAIN || errno == EACCES) { + return status::busy; + } + // flock is not implemented (some NFS/FUSE mounts): report it instead of racing + LOG_ERR("%s: file locking failed (%s)\n", __func__, strerror(errno)); + return status::error; +#endif +} + +bool common_file_lock::acquire(const std::function & keep_waiting) { + // the lock file is opened by try_lock, its parent must exist by then + if (!lock_path.empty()) { + std::error_code ec; + auto lock_dir = std::filesystem::path(lock_path).parent_path(); + if (!lock_dir.empty()) { + std::filesystem::create_directories(lock_dir, ec); + if (ec) { + LOG_ERR("%s: unable to create lock directory '%s': %s\n", __func__, lock_dir.string().c_str(), ec.message().c_str()); + return false; + } + } + } + + while (true) { + switch (try_lock()) { + case status::acquired: + return true; + case status::busy: + if (keep_waiting && !keep_waiting()) { + return false; + } + std::this_thread::sleep_for(std::chrono::milliseconds(100)); + continue; + case status::error: + return false; + } + } +} + +// the lock file is never deleted: another process could then create a fresh +// file at the same path and lock that +void common_file_lock::close() { +#if defined(_WIN32) + if (fd >= 0) { + if (locked) { + OVERLAPPED ov = {}; + UnlockFileEx((HANDLE) _get_osfhandle(fd), 0, 1, 0, &ov); + } + _close(fd); + } +#else + if (fd >= 0) { + if (locked) { + flock(fd, LOCK_UN); + } + ::close(fd); + } +#endif + fd = -1; + locked = false; +} + +common_file_lock::~common_file_lock() { + close(); +} diff --git a/common/file-lock.h b/common/file-lock.h new file mode 100644 index 000000000000..6f9443ff81a0 --- /dev/null +++ b/common/file-lock.h @@ -0,0 +1,36 @@ +#pragma once + +#include +#include + +// Cross-process exclusive lock (flock on POSIX, LockFileEx on Windows). +// Lazy: acquire() opens the lock file, so a read-only cache works if the lock is never acquired. +class common_file_lock { +public: + common_file_lock(const std::string & lock_path); + ~common_file_lock(); + + // take the lock, waiting while another process holds it; + // keep_waiting is polled before each retry, false cancels the wait. + // Also returns false when locking is unavailable. + bool acquire(const std::function & keep_waiting = {}); + + void close(); + + common_file_lock(const common_file_lock &) = delete; + common_file_lock & operator=(const common_file_lock &) = delete; + +private: + enum class status { + acquired, + busy, // another process holds the lock + error, // file locking is unavailable: no exclusion is possible + }; + + status try_lock(); + + std::string lock_path; + + int fd = -1; + bool locked = false; +}; diff --git a/common/hf-cache.cpp b/common/hf-cache.cpp index 12d4fcc80447..8f678cd947f0 100644 --- a/common/hf-cache.cpp +++ b/common/hf-cache.cpp @@ -2,10 +2,12 @@ #include "build-info.h" #include "common.h" +#include "file-lock.h" #include "log.h" #include "http.h" #include "json.h" +#include #include #include #include @@ -13,6 +15,9 @@ #include #include +#include +#include + #if defined(_WIN32) #define WIN32_LEAN_AND_MEAN #ifndef NOMINMAX @@ -20,12 +25,15 @@ #endif #define HOME_DIR "USERPROFILE" #include +#include #else #define HOME_DIR "HOME" #include #include #endif +#include + namespace hf_cache { namespace fs = std::filesystem; @@ -167,6 +175,14 @@ static bool is_valid_oid(const std::string & oid) { return is_hex_string(oid, 40) || is_hex_string(oid, 64); } +std::string get_lock_path(const std::string & repo_id, const std::string & oid) { + if (!is_valid_repo_id(repo_id) || !is_valid_oid(oid)) { + return {}; + } + fs::path lock_path = fs::path(get_cache_path()) / ".locks" / repo_to_folder_name(repo_id) / (oid + ".lock"); + return lock_path.string(); +} + static bool is_valid_subpath(const fs::path & path, const fs::path & subpath) { if (subpath.is_absolute()) { return false; // never do a / b with b absolute @@ -179,24 +195,71 @@ static bool is_valid_subpath(const fs::path & path, const fs::path & subpath) { } static void safe_write_file(const fs::path & path, const std::string & data) { - fs::path path_tmp = path.string() + ".tmp"; + // pids are not unique across containers sharing the cache, so two processes + // can derive the same temp name: create it exclusively and try the next name on collision + static std::atomic counter{0}; if (path.has_parent_path()) { fs::create_directories(path.parent_path()); } - std::ofstream file(path_tmp); - file << data; - file.close(); + const std::string path_str = path.string(); + std::string path_tmp; + int fd = -1; - std::error_code ec; + for (int attempt = 0; attempt < 64; ++attempt) { +#if defined(_WIN32) + const uint64_t id = ((uint64_t) GetCurrentProcessId() << 32) | counter++; +#else + const uint64_t id = ((uint64_t) getpid() << 32) | counter++; +#endif + path_tmp = path_str + "." + std::to_string(id) + ".tmp"; + +#if defined(_WIN32) + fd = _open(path_tmp.c_str(), _O_WRONLY | _O_CREAT | _O_EXCL | _O_BINARY | _O_NOINHERIT, + _S_IREAD | _S_IWRITE); +#else + fd = open(path_tmp.c_str(), O_WRONLY | O_CREAT | O_EXCL | O_CLOEXEC, 0666); +#endif + if (fd >= 0) { + break; + } + if (errno != EEXIST) { + break; // report the failure below, retrying cannot fix it + } + } + + if (fd < 0) { + throw std::runtime_error("failed to write file: " + path_str); + } - if (!file.fail()) { + size_t written = 0; + while (written < data.size()) { +#if defined(_WIN32) + const int n = _write(fd, data.data() + written, (unsigned int) (data.size() - written)); +#else + const ssize_t n = write(fd, data.data() + written, data.size() - written); +#endif + if (n <= 0) { + break; + } + written += (size_t) n; + } + + // delayed write errors surface only at close: the rename must not publish them +#if defined(_WIN32) + const int close_ret = _close(fd); +#else + const int close_ret = close(fd); +#endif + + std::error_code ec; + if (written == data.size() && close_ret == 0) { fs::rename(path_tmp, path, ec); } - if (file.fail() || ec) { + if (written != data.size() || close_ret != 0 || ec) { fs::remove(path_tmp, ec); - throw std::runtime_error("failed to write file: " + path.string()); + throw std::runtime_error("failed to write file: " + path_str); } } @@ -457,7 +520,7 @@ hf_files get_cached_files(const std::string & repo_id) { return files; } -std::string finalize_file(const hf_file & file) { +std::string finalize_file(const hf_file & file, const std::function & keep_waiting) { static std::atomic symlinks_disabled{false}; std::error_code ec; @@ -472,6 +535,24 @@ std::string finalize_file(const hf_file & file) { return file.final_path; } + // the lock file is shared with huggingface_hub: another process may be + // finalizing the same blob + common_file_lock lock(get_lock_path(file.repo_id, file.oid)); + if (!lock.acquire(keep_waiting)) { + // without the lock, another process could move the blob while we check it + throw std::runtime_error("failed to lock cache for finalization: " + file.final_path); + } + + // acquire() only polls keep_waiting while waiting for a busy lock + if (keep_waiting && !keep_waiting()) { + throw std::runtime_error("download cancelled before finalization: " + file.final_path); + } + + // re-check under the lock: the winner may have created the link already + if (fs::exists(final_path, ec)) { + return file.final_path; + } + fs::create_directories(final_path.parent_path(), ec); if (!symlinks_disabled) { diff --git a/common/hf-cache.h b/common/hf-cache.h index 41842db788cb..c03e5718260e 100644 --- a/common/hf-cache.h +++ b/common/hf-cache.h @@ -1,5 +1,6 @@ #pragma once +#include #include #include @@ -26,8 +27,14 @@ hf_files get_repo_files( hf_files get_cached_files(const std::string & repo_id = {}); -// Create snapshot path (link or move/copy) and return it -std::string finalize_file(const hf_file & file); +// Lock file path for a blob, shared with the hf CLI: /.locks//.lock +// Pure path computation: creates nothing. Returns an empty string on invalid repo or oid. +std::string get_lock_path(const std::string & repo_id, const std::string & oid); + +// Create snapshot path (link or move/copy) and return it. +// Throws when the lock cannot be taken or keep_waiting returns false. +// A completed snapshot (final_path present) is returned without the lock, so a read-only cache works. +std::string finalize_file(const hf_file & file, const std::function & keep_waiting = {}); // Remove the entire cached directory for a repo, returns true if removed bool remove_cached_repo(const std::string & repo_id);