Brunobkr/llama.cpp_AlgMor24_github

1#include "hf-cache.h"2 3#include "build-info.h"4#include "common.h"5#include "log.h"6#include "http.h"7 8#define JSON_ASSERT GGML_ASSERT9#include <nlohmann/json.hpp>10 11#include <filesystem>12#include <fstream>13#include <atomic>14#include <string>15#include <string_view>16#include <stdexcept>17 18namespace nl = nlohmann;19 20#if defined(_WIN32)21#define WIN32_LEAN_AND_MEAN22#ifndef NOMINMAX23#define NOMINMAX24#endif25#define HOME_DIR "USERPROFILE"26#include <windows.h>27#else28#define HOME_DIR "HOME"29#include <unistd.h>30#include <pwd.h>31#endif32 33namespace hf_cache {34 35namespace fs = std::filesystem;36 37static fs::path get_cache_directory() {38    static const fs::path cache = []() {39        struct {40            const char * var;41            fs::path path;42        } entries[] = {43            {"LLAMA_CACHE",           fs::path()},44            {"HF_HUB_CACHE",          fs::path()},45            {"HUGGINGFACE_HUB_CACHE", fs::path()},46            {"HF_HOME",               fs::path("hub")},47            {"XDG_CACHE_HOME",        fs::path("huggingface") / "hub"},48            {HOME_DIR,                fs::path(".cache") / "huggingface" / "hub"}49        };50        for (const auto & entry : entries) {51            if (auto * p = std::getenv(entry.var); p && *p) {52                fs::path base(p);53                return entry.path.empty() ? base : base / entry.path;54            }55        }56#ifndef _WIN3257        const struct passwd * pw = getpwuid(getuid());58 59        if (pw && pw->pw_dir && *pw->pw_dir) {60            return fs::path(pw->pw_dir) / ".cache" / "huggingface" / "hub";61        }62#endif63        throw std::runtime_error("Failed to determine HF cache directory");64    }();65 66    return cache;67}68 69static std::string folder_name_to_repo(const std::string & folder) {70    constexpr std::string_view prefix = "models--";71    if (folder.rfind(prefix, 0)) {72        return {};73    }74    std::string result = folder.substr(prefix.length());75    string_replace_all(result, "--", "/");76    return result;77}78 79static std::string repo_to_folder_name(const std::string & repo_id) {80    constexpr std::string_view prefix = "models--";81    std::string result = std::string(prefix) + repo_id;82    string_replace_all(result, "/", "--");83    return result;84}85 86static fs::path get_repo_path(const std::string & repo_id) {87    return get_cache_directory() / repo_to_folder_name(repo_id);88}89 90static bool is_hex_char(const char c) {91    return (c >= 'A' && c <= 'F') ||92           (c >= 'a' && c <= 'f') ||93           (c >= '0' && c <= '9');94}95 96static bool is_hex_string(const std::string & s, size_t expected_len) {97    if (s.length() != expected_len) {98        return false;99    }100    for (const char c : s) {101        if (!is_hex_char(c)) {102            return false;103        }104    }105    return true;106}107 108static bool is_alphanum(const char c) {109    return (c >= 'A' && c <= 'Z') ||110           (c >= 'a' && c <= 'z') ||111           (c >= '0' && c <= '9');112}113 114static bool is_special_char(char c) {115    return c == '/' || c == '.' || c == '-';116}117 118// base chars [A-Za-z0-9_] are always valid119// special chars [/.-] must be surrounded by base chars120// exactly one '/' required121static bool is_valid_repo_id(const std::string & repo_id) {122    if (repo_id.empty() || repo_id.length() > 256) {123        return false;124    }125    int slash = 0;126    bool special = true;127 128    for (const char c : repo_id) {129        if (is_alphanum(c) || c == '_') {130            special = false;131        } else if (is_special_char(c)) {132            if (special) {133                return false;134            }135            slash += (c == '/');136            special = true;137        } else {138            return false;139        }140    }141    return !special && slash == 1;142}143 144static bool is_valid_hf_token(const std::string & token) {145    if (token.length() < 37 || token.length() > 256 ||146        !string_starts_with(token, "hf_")) {147        return false;148    }149    for (size_t i = 3; i < token.length(); ++i) {150        if (!is_alphanum(token[i])) {151            return false;152        }153    }154    return true;155}156 157static bool is_valid_commit(const std::string & hash) {158    return is_hex_string(hash, 40);159}160 161static bool is_valid_oid(const std::string & oid) {162    return is_hex_string(oid, 40) || is_hex_string(oid, 64);163}164 165static bool is_valid_subpath(const fs::path & path, const fs::path & subpath) {166    if (subpath.is_absolute()) {167        return false; // never do a / b with b absolute168    }169    auto b = fs::absolute(path).lexically_normal();170    auto t = (b / subpath).lexically_normal();171    auto [b_end, _] = std::mismatch(b.begin(), b.end(), t.begin(), t.end());172 173    return b_end == b.end();174}175 176static void safe_write_file(const fs::path & path, const std::string & data) {177    fs::path path_tmp = path.string() + ".tmp";178 179    if (path.has_parent_path()) {180        fs::create_directories(path.parent_path());181    }182 183    std::ofstream file(path_tmp);184    file << data;185    file.close();186 187    std::error_code ec;188 189    if (!file.fail()) {190        fs::rename(path_tmp, path, ec);191    }192    if (file.fail() || ec) {193        fs::remove(path_tmp, ec);194        throw std::runtime_error("failed to write file: " + path.string());195    }196}197 198static nl::json api_get(const std::string & url,199                        const std::string & token) {200    auto [cli, parts] = common_http_client(url);201 202    httplib::Headers headers = {203        {"User-Agent", "llama-cpp/" + std::string(llama_build_info())},204        {"Accept", "application/json"}205    };206 207    if (is_valid_hf_token(token)) {208        headers.emplace("Authorization", "Bearer " + token);209    } else if (!token.empty()) {210        LOG_WRN("%s: invalid token, authentication disabled\n", __func__);211    }212 213    if (auto res = cli.Get(parts.path, headers)) {214        auto body = res->body;215 216        if (res->status == 200) {217            return nl::json::parse(res->body);218        }219        try {220            body = nl::json::parse(res->body)["error"].get<std::string>();221        } catch (...) { }222 223        throw std::runtime_error("GET failed (" + std::to_string(res->status) + "): " + body);224    } else {225        throw std::runtime_error("HTTPLIB failed: " + httplib::to_string(res.error()));226    }227}228 229static std::string get_repo_commit(const std::string & repo_id,230                                   const std::string & token) {231    try {232        auto endpoint = common_get_model_endpoint();233        auto json = api_get(endpoint + "api/models/" + repo_id + "/refs", token);234 235        if (!json.is_object() ||236            !json.contains("branches") || !json["branches"].is_array()) {237            LOG_WRN("%s: missing 'branches' for '%s'\n", __func__, repo_id.c_str());238            return {};239        }240 241        fs::path refs_path = get_repo_path(repo_id) / "refs";242        std::string name;243        std::string commit;244 245        for (const auto & branch : json["branches"]) {246            if (!branch.is_object() ||247                !branch.contains("name") || !branch["name"].is_string() ||248                !branch.contains("targetCommit") || !branch["targetCommit"].is_string()) {249                continue;250            }251            std::string _name = branch["name"].get<std::string>();252            std::string _commit = branch["targetCommit"].get<std::string>();253 254            if (!is_valid_subpath(refs_path, _name)) {255                LOG_WRN("%s: skip invalid branch: %s\n", __func__, _name.c_str());256                continue;257            }258            if (!is_valid_commit(_commit)) {259                LOG_WRN("%s: skip invalid commit: %s\n", __func__, _commit.c_str());260                continue;261            }262 263            if (_name == "main") {264                name = _name;265                commit = _commit;266                break;267            }268 269            if (name.empty() || commit.empty()) {270                name = _name;271                commit = _commit;272            }273        }274 275        if (name.empty() || commit.empty()) {276            LOG_WRN("%s: no valid branch for '%s'\n", __func__, repo_id.c_str());277            return {};278        }279 280        safe_write_file(refs_path / name, commit);281        return commit;282 283    } catch (const nl::json::exception & e) {284        LOG_ERR("%s: JSON error: %s\n", __func__, e.what());285    } catch (const std::exception & e) {286        LOG_ERR("%s: error: %s\n", __func__, e.what());287    }288    return {};289}290 291hf_files get_repo_files(const std::string & repo_id,292                        const std::string & token) {293    if (!is_valid_repo_id(repo_id)) {294        LOG_WRN("%s: invalid repository: %s\n", __func__, repo_id.c_str());295        return {};296    }297 298    std::string commit = get_repo_commit(repo_id, token);299    if (commit.empty()) {300        LOG_WRN("%s: failed to resolve commit for %s\n", __func__, repo_id.c_str());301        return {};302    }303 304    fs::path blobs_path = get_repo_path(repo_id) / "blobs";305    fs::path commit_path = get_repo_path(repo_id) / "snapshots" / commit;306 307    hf_files files;308 309    try {310        auto endpoint = common_get_model_endpoint();311        auto json = api_get(endpoint + "api/models/" + repo_id + "/tree/" + commit + "?recursive=true", token);312 313        if (!json.is_array()) {314            LOG_WRN("%s: response is not an array for '%s'\n", __func__, repo_id.c_str());315            return {};316        }317 318        for (const auto & item : json) {319            if (!item.is_object() ||320                !item.contains("type") || !item["type"].is_string() || item["type"] != "file" ||321                !item.contains("path") || !item["path"].is_string()) {322                continue;323            }324 325            hf_file file;326            file.repo_id = repo_id;327            file.path = item["path"].get<std::string>();328 329            if (!is_valid_subpath(commit_path, file.path)) {330                LOG_WRN("%s: skip invalid path: %s\n", __func__, file.path.c_str());331                continue;332            }333 334            if (item.contains("lfs") && item["lfs"].is_object()) {335                if (item["lfs"].contains("oid") && item["lfs"]["oid"].is_string()) {336                    file.oid = item["lfs"]["oid"].get<std::string>();337                }338            } else if (item.contains("oid") && item["oid"].is_string()) {339                file.oid = item["oid"].get<std::string>();340            }341 342            if (!file.oid.empty() && !is_valid_oid(file.oid)) {343                LOG_WRN("%s: skip invalid oid: %s\n", __func__, file.oid.c_str());344                continue;345            }346 347            file.url = endpoint + repo_id + "/resolve/" + commit + "/" + file.path;348 349            fs::path final_path = commit_path / file.path;350            file.final_path = final_path.string();351 352            if (!file.oid.empty() && !fs::exists(final_path)) {353                fs::path local_path = blobs_path / file.oid;354                file.local_path = local_path.string();355            } else {356                file.local_path = file.final_path;357            }358 359            files.push_back(file);360        }361    } catch (const nl::json::exception & e) {362        LOG_ERR("%s: JSON error: %s\n", __func__, e.what());363    } catch (const std::exception & e) {364        LOG_ERR("%s: error: %s\n", __func__, e.what());365    }366    return files;367}368 369static std::string get_cached_ref(const fs::path & repo_path) {370    fs::path refs_path = repo_path / "refs";371    if (!fs::is_directory(refs_path)) {372        return {};373    }374    std::string fallback;375 376    for (const auto & entry : fs::directory_iterator(refs_path)) {377        if (!entry.is_regular_file()) {378            continue;379        }380        std::ifstream f(entry.path());381        std::string commit;382        if (!f || !std::getline(f, commit) || commit.empty()) {383            continue;384        }385        if (!is_valid_commit(commit)) {386            LOG_WRN("%s: skip invalid commit: %s\n", __func__, commit.c_str());387            continue;388        }389        if (entry.path().filename() == "main") {390            return commit;391        }392        if (fallback.empty()) {393            fallback = commit;394        }395    }396    return fallback;397}398 399hf_files get_cached_files(const std::string & repo_id) {400    fs::path cache_dir = get_cache_directory();401    if (!fs::exists(cache_dir)) {402        return {};403    }404 405    if (!repo_id.empty() && !is_valid_repo_id(repo_id)) {406        LOG_WRN("%s: invalid repository: %s\n", __func__, repo_id.c_str());407        return {};408    }409 410    hf_files files;411 412    for (const auto & repo : fs::directory_iterator(cache_dir)) {413        if (!repo.is_directory()) {414            continue;415        }416        fs::path snapshots_path = repo.path() / "snapshots";417 418        if (!fs::exists(snapshots_path)) {419            continue;420        }421        std::string _repo_id = folder_name_to_repo(repo.path().filename().string());422 423        if (!is_valid_repo_id(_repo_id)) {424            continue;425        }426        if (!repo_id.empty() && _repo_id != repo_id) {427            continue;428        }429        std::string commit = get_cached_ref(repo.path());430        fs::path commit_path = snapshots_path / commit;431 432        if (commit.empty() || !fs::is_directory(commit_path)) {433            continue;434        }435        for (const auto & entry : fs::recursive_directory_iterator(commit_path)) {436            if (!entry.is_regular_file() && !entry.is_symlink()) {437                continue;438            }439            fs::path path = entry.path().lexically_relative(commit_path);440 441            if (!path.empty()) {442                hf_file file;443                file.repo_id = _repo_id;444                file.path = path.generic_string();445                file.local_path = entry.path().string();446                file.final_path = file.local_path;447                files.push_back(std::move(file));448            }449        }450    }451 452    return files;453}454 455std::string finalize_file(const hf_file & file) {456    static std::atomic<bool> symlinks_disabled{false};457 458    std::error_code ec;459    fs::path local_path(file.local_path);460    fs::path final_path(file.final_path);461 462    if (local_path == final_path || fs::exists(final_path, ec)) {463        return file.final_path;464    }465 466    if (!fs::exists(local_path, ec)) {467        return file.final_path;468    }469 470    fs::create_directories(final_path.parent_path(), ec);471 472    if (!symlinks_disabled) {473        fs::path target = fs::relative(local_path, final_path.parent_path(), ec);474        if (!ec) {475            fs::create_symlink(target, final_path, ec);476        }477        if (!ec) {478            return file.final_path;479        }480    }481 482    if (!symlinks_disabled.exchange(true)) {483        LOG_WRN("%s: failed to create symlink: %s\n", __func__, ec.message().c_str());484        LOG_WRN("%s: switching to degraded mode\n", __func__);485    }486 487    fs::rename(local_path, final_path, ec);488    if (ec) {489        LOG_WRN("%s: failed to move file to snapshots: %s\n", __func__, ec.message().c_str());490        fs::copy(local_path, final_path, ec);491        if (ec) {492            LOG_ERR("%s: failed to copy file to snapshots: %s\n", __func__, ec.message().c_str());493        }494    }495    return file.final_path;496}497 498bool remove_cached_repo(const std::string & repo_id) {499    if (!is_valid_repo_id(repo_id)) {500        LOG_WRN("%s: invalid repository: %s\n", __func__, repo_id.c_str());501        return false;502    }503    fs::path repo_path = get_repo_path(repo_id);504    std::error_code ec;505    auto removed = fs::remove_all(repo_path, ec);506    if (ec) {507        LOG_ERR("%s: failed to remove repo cache %s: %s\n", __func__, repo_path.string().c_str(), ec.message().c_str());508        return false;509    }510    return removed > 0;511}512 513} // namespace hf_cache514