diff --git a/CMakeLists.txt b/CMakeLists.txt index 81f23d7e70b..9d894db468c 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -118,6 +118,7 @@ option(LLAMA_TESTS_INSTALL "llama: install tests" ON) # 3rd party libs option(LLAMA_OPENSSL "llama: use openssl to support HTTPS" ON) option(LLAMA_LLGUIDANCE "llama-common: include LLGuidance library for structured output in common utils" OFF) +option(LLAMA_XET "llama-common: use the xet client library to download xet-backed HF files (requires XET_CORE_DIR)" OFF) # Required for relocatable CMake package diff --git a/common/CMakeLists.txt b/common/CMakeLists.txt index 4cf580a056c..e2d1eb05483 100644 --- a/common/CMakeLists.txt +++ b/common/CMakeLists.txt @@ -167,4 +167,13 @@ if (LLAMA_LLGUIDANCE) endif() endif() +if (LLAMA_XET) + if (NOT XET_CORE_DIR) + message(FATAL_ERROR "LLAMA_XET requires XET_CORE_DIR to point to a local xet-core checkout") + endif() + add_subdirectory(${XET_CORE_DIR}/xet_capi ${CMAKE_CURRENT_BINARY_DIR}/xet_capi) + target_compile_definitions(${TARGET} PUBLIC LLAMA_USE_XET) + target_link_libraries(${TARGET} PRIVATE xet_capi) +endif() + target_link_libraries(${TARGET} PUBLIC llama Threads::Threads) diff --git a/common/download.cpp b/common/download.cpp index 6b69a441885..0ed6cba4e11 100644 --- a/common/download.cpp +++ b/common/download.cpp @@ -23,6 +23,10 @@ #include "http.h" +#ifdef LLAMA_USE_XET +#include "xet_session.hpp" +#endif + #ifndef __EMSCRIPTEN__ #ifdef __linux__ #include @@ -744,12 +748,108 @@ common_download_hf_plan common_download_get_hf_plan(const common_params_model & return plan; } +#ifdef LLAMA_USE_XET +// download a xet-backed HF file through the xet C API, returns true on success +static bool common_download_file_xet(const common_download_task & task, + common_download_callback * callback) { + const std::string & path = task.local_path; + const std::string path_temporary = path + ".downloadInProgress"; + + { // silent + std::error_code ec; + std::filesystem::create_directories(std::filesystem::path(path).parent_path(), ec); + } + + LOG_INF("%s: downloading %s via xet (hash: %s)\n", __func__, task.url.c_str(), task.xet_hash.c_str()); + + common_download_progress p; + p.url = task.url; + p.total = task.size; + + if (callback) { + callback->on_start(p); + } + + bool success = false; + bool cancelled = false; + try { + static xet::Session session; + + xet::AuthConfig cfg; + cfg.token_refresh_url(task.xet_token_url); + if (!task.opts.bearer_token.empty()) { + cfg.add_refresh_header("Authorization", "Bearer " + task.opts.bearer_token); + } + + auto group = session.new_file_download_group(cfg); + xet::FileInfo file_info = task.sha256.empty() + ? xet::FileInfo(task.xet_hash.c_str(), task.size) + : xet::FileInfo(task.xet_hash.c_str(), task.size, task.sha256.c_str()); + + auto download = group.download_to_path(file_info, path_temporary.c_str()); + auto op = group.finish_start(); + + while (op.poll() == XetPollState_XetPollPending) { + std::this_thread::sleep_for(std::chrono::milliseconds(100)); + if (callback) { + p.downloaded = group.progress().total_bytes_completed; + callback->on_update(p); + if (!cancelled && callback->is_cancelled()) { + cancelled = true; + group.abort(); + } + } + } + op.wait(); // throws on error, including after abort + + if (!cancelled) { + if (std::rename(path_temporary.c_str(), path.c_str()) != 0) { + LOG_ERR("%s: unable to rename file: %s to %s\n", __func__, path_temporary.c_str(), path.c_str()); + } else { + p.downloaded = p.total; + success = true; + } + } + } catch (const std::exception & e) { + if (!cancelled) { + LOG_WRN("%s: xet download of %s failed: %s (falling back to HTTP)\n", __func__, task.url.c_str(), e.what()); + } + } + + if (callback) { + callback->on_done(p, success); + } + if (!success) { + // remove partial output so the HTTP fallback does not resume from it + std::error_code ec; + std::filesystem::remove(path_temporary, ec); + } + return success; +} +#endif + +static int common_download_file_task(const common_download_task & task) { +#ifdef LLAMA_USE_XET + if (!task.xet_hash.empty() && !task.opts.offline && !std::filesystem::exists(task.local_path)) { + ProgressBar tty_cb; + common_download_callback * callback = task.opts.callback ? task.opts.callback : &tty_cb; + if (common_download_file_xet(task, callback)) { + return 200; + } + if (callback->is_cancelled()) { + return -1; + } + } +#endif + return common_download_file_single(task.url, task.local_path, task.opts, task.is_hf); +} + void common_download_run_tasks(const std::vector & tasks) { std::vector> futures; 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_task(task); } )); } diff --git a/common/download.h b/common/download.h index 816e1c7f58a..8ca374948ab 100644 --- a/common/download.h +++ b/common/download.h @@ -66,12 +66,19 @@ struct common_download_task { std::string local_path; std::function on_done; bool is_hf = false; + // xet download info, empty xet_hash means the file is not stored on xet + std::string xet_hash; + std::string xet_token_url; + std::string sha256; + size_t size = 0; 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), + xet_hash(f.xet_hash), xet_token_url(f.xet_token_url), + sha256(f.oid.size() == 64 ? f.oid : ""), size(f.size) {} }; void common_download_run_tasks(const std::vector & tasks); diff --git a/common/hf-cache.cpp b/common/hf-cache.cpp index f1dacaa4778..d2143c2541c 100644 --- a/common/hf-cache.cpp +++ b/common/hf-cache.cpp @@ -344,6 +344,18 @@ hf_files get_repo_files(const std::string & repo_id, continue; } + if (item.contains("size") && item["size"].is_number_unsigned()) { + file.size = item["size"].get(); + } + + if (item.contains("xetHash") && item["xetHash"].is_string()) { + std::string xet_hash = item["xetHash"].get(); + if (is_hex_string(xet_hash, 64) && file.size > 0) { + file.xet_hash = xet_hash; + file.xet_token_url = endpoint + "api/models/" + repo_id + "/xet-read-token/" + commit; + } + } + file.url = endpoint + repo_id + "/resolve/" + commit + "/" + file.path; fs::path final_path = commit_path / file.path; diff --git a/common/hf-cache.h b/common/hf-cache.h index 42c9c6ce34f..6682d3c36e3 100644 --- a/common/hf-cache.h +++ b/common/hf-cache.h @@ -14,6 +14,10 @@ struct hf_file { std::string final_path; std::string oid; std::string repo_id; + size_t size = 0; + // set when the file is stored on xet (has a xetHash in the tree API response) + std::string xet_hash; + std::string xet_token_url; }; using hf_files = std::vector;