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