Optional xet-backed HF downloads (LLAMA_XET)

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

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)

common/download.cpp

@@ -23,6 +23,10 @@
 
 #include "http.h"
 
+#ifdef LLAMA_USE_XET
+#include "xet_session.hpp"
+#endif
+
 #ifndef __EMSCRIPTEN__
 #ifdef __linux__
 #include <linux/limits.h>
@@ -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<common_download_task> & tasks) {
     std::vector<std::future<int>> 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);
             }
         ));
     }

common/download.h

@@ -66,12 +66,19 @@ struct common_download_task {
     std::string local_path;
     std::function<void()> 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<void()> 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<common_download_task> & tasks);

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<size_t>();
+            }
+
+            if (item.contains("xetHash") && item["xetHash"].is_string()) {
+                std::string xet_hash = item["xetHash"].get<std::string>();
+                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;

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<hf_file>;