char-fork/llama.cpp

git clone https://git.t4t.associates/char-fork/llama.cpp

Xuan-Son Nguyencommon: add json.h abstraction (#27511)d9f918d2d

char/aliceai-t5-moe
15.1 KiB509 linesraw
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_MEAN
18#ifndef NOMINMAX
19#define NOMINMAX
20#endif
21#define HOME_DIR "USERPROFILE"
22#include <windows.h>
23#else
24#define HOME_DIR "HOME"
25#include <unistd.h>
26#include <pwd.h>
27#endif
28
29namespace hf_cache {
30
31namespace fs = std::filesystem;
32
33static fs::path get_cache_directory() {
34    static const fs::path 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;
50            }
51        }
52#ifndef _WIN32
53        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";
57        }
58#endif
59        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 get_cache_directory() / 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 valid
115// special chars [/.-] must be surrounded by base chars
116// exactly one '/' required
117static 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 absolute
164    }
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    fs::path cache_dir = get_cache_directory();
397    if (!fs::exists(cache_dir)) {
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_dir)) {
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_cache