wehub-resource-sync
241 lines
10 KiB
C++
241 lines
10 KiB
C++
//
|
|
// Created by ruoyi.sjd on 2024/12/18.
|
|
// Copyright (c) 2024 Alibaba Group Holding Limited All rights reserved.
|
|
//
|
|
|
|
#include "remote_model_downloader.hpp"
|
|
#include "httplib.h"
|
|
#include <cstdio>
|
|
#include <fstream>
|
|
#include <functional>
|
|
#include "file_utils.hpp"
|
|
|
|
static const std::string HUGGINGFACE_HEADER_X_REPO_COMMIT = "x-repo-commit";
|
|
static const std::string HUGGINGFACE_HEADER_X_LINKED_ETAG = "x-linked-etag";
|
|
static const std::string HUGGINGFACE_HEADER_X_LINKED_SIZE = "x-linked-size";
|
|
|
|
namespace mls {
|
|
|
|
size_t ParseContentLength(const std::string& content_length) {
|
|
if (!content_length.empty()) {
|
|
return std::stoul(content_length);
|
|
}
|
|
return 0;
|
|
}
|
|
|
|
std::string NormalizeETag(const std::string& etag) {
|
|
if (!etag.empty() && etag[0] == '"' && etag[etag.size() - 1] == '"') {
|
|
return etag.substr(1, etag.size() - 2);
|
|
}
|
|
return etag;
|
|
}
|
|
|
|
//from get_hf_file_metadata
|
|
HfFileMetadata RemoteModelDownloader::GetFileMetadata(const std::string& url, std::string& error_info) {
|
|
httplib::SSLClient cli(this->host_, 443);
|
|
httplib::Headers headers;
|
|
HfFileMetadata metadata{};
|
|
headers.emplace("Accept-Encoding", "identity");
|
|
auto res = cli.Head(url, headers);
|
|
if (!res || (res->status != 200 && res->status != 302)) {
|
|
error_info = "GetFileMetadata Failed to fetch metadata status " + std::to_string(res ? res->status : -1);
|
|
return metadata;
|
|
}
|
|
metadata.location = url;
|
|
if (res->status == 302) {
|
|
metadata.location = res->get_header_value("Location");
|
|
}
|
|
std::string linked_etag = res->get_header_value(HUGGINGFACE_HEADER_X_LINKED_ETAG);
|
|
std::string etag = res->get_header_value("ETag");
|
|
metadata.etag = NormalizeETag(!linked_etag.empty() ? linked_etag : etag);
|
|
std::string linked_size = res->get_header_value(HUGGINGFACE_HEADER_X_LINKED_SIZE);
|
|
std::string content_length = res->get_header_value("Content-Length");
|
|
metadata.size = ParseContentLength(!linked_size.empty() ? linked_size : content_length);
|
|
metadata.commit_hash = res->get_header_value(HUGGINGFACE_HEADER_X_REPO_COMMIT);
|
|
return metadata;
|
|
}
|
|
|
|
RemoteModelDownloader::RemoteModelDownloader(std::string host, int max_attempts, int retry_delay_seconds)
|
|
: max_attempts_(max_attempts),
|
|
retry_delay_seconds_(retry_delay_seconds),
|
|
host_(std::move(host)){
|
|
}
|
|
|
|
std::string RemoteModelDownloader::DownloadWithRetries(
|
|
const fs::path& storage_folder,
|
|
const std::string& repo,
|
|
const std::string& revision,
|
|
const std::string& relative_path,
|
|
std::string& error_info,
|
|
int max_retries) {
|
|
int attempt = 0;
|
|
bool success = false;
|
|
while (attempt < max_retries) {
|
|
// Reset error message for each attempt
|
|
error_info.clear();
|
|
auto result = this->DownloadFile(storage_folder, repo, revision, relative_path, error_info);
|
|
if (error_info.empty()) {
|
|
success = true;
|
|
return result;
|
|
} else {
|
|
attempt++;
|
|
fprintf(stderr, "DownloadFile error at file: %s error message: %s, attempt: %d\n",
|
|
relative_path.c_str(), error_info.c_str(), attempt);
|
|
std::this_thread::sleep_for(std::chrono::seconds(1));
|
|
}
|
|
}
|
|
return "";
|
|
}
|
|
|
|
std::string RemoteModelDownloader::DownloadFile(
|
|
const fs::path& storage_folder,
|
|
const std::string& repo,
|
|
const std::string& revision,
|
|
const std::string& relative_path,
|
|
std::string& error_info) {
|
|
std::string url = "https://" + this->host_ + "/";
|
|
url += repo;
|
|
url += "/resolve/" + revision + "/" + relative_path;
|
|
auto metadata = GetFileMetadata(url, error_info);
|
|
if (!error_info.empty()) {
|
|
printf("DownloadFile GetFileMetadata faield\n");
|
|
return "";
|
|
}
|
|
auto repo_folder_name = mls::FileUtils::RepoFolderName(repo, "model");
|
|
fs::path blob_path = storage_folder / "blobs" / metadata.etag;
|
|
std::filesystem::path blob_path_incomplete = storage_folder / "blobs" / (metadata.etag + ".incomplete");
|
|
fs::path pointer_path = mls::FileUtils::GetPointerPath(
|
|
storage_folder, metadata.commit_hash, relative_path);
|
|
fs::create_directories(blob_path.parent_path());
|
|
fs::create_directories(pointer_path.parent_path());
|
|
|
|
if (fs::exists(pointer_path)) {
|
|
return pointer_path.string();
|
|
} else if (fs::exists(blob_path)) {
|
|
std::error_code ec;
|
|
mls::FileUtils::CreateSymlink(blob_path, pointer_path, ec);
|
|
if (ec) {
|
|
fprintf(stderr, "DownloadFile create symlink error for pointer_path: %s", pointer_path.string().c_str());
|
|
error_info = ec.message();
|
|
return "";
|
|
}
|
|
printf("DownloadFile %s already exists just create symlink\n", relative_path.c_str());
|
|
return pointer_path.string();
|
|
}
|
|
|
|
std::mutex lock;
|
|
{
|
|
std::lock_guard guard(lock);
|
|
httplib::Headers headers;
|
|
DownloadToTmpAndMove(blob_path_incomplete, blob_path, metadata.location, headers, metadata.size,
|
|
relative_path, false, error_info);
|
|
if (error_info.empty()) {
|
|
std::error_code ec;
|
|
FileUtils::CreateSymlink(blob_path, pointer_path, ec);
|
|
if (ec) {
|
|
error_info = "create link error" + ec.message();
|
|
}
|
|
}
|
|
}
|
|
return pointer_path.string();
|
|
}
|
|
|
|
void RemoteModelDownloader::DownloadToTmpAndMove(
|
|
const fs::path& incomplete_path,
|
|
const fs::path& destination_path,
|
|
const std::string& url_to_download,
|
|
httplib::Headers& headers,
|
|
size_t expected_size,
|
|
const std::string& file_name,
|
|
bool force_download,
|
|
std::string& error_info) {
|
|
if (fs::exists(destination_path) && !force_download) {
|
|
return;
|
|
}
|
|
if (std::filesystem::exists(incomplete_path) && force_download) {
|
|
std::filesystem::remove(incomplete_path);
|
|
}
|
|
size_t resume_size = std::filesystem::exists(incomplete_path) ? std::filesystem::file_size(incomplete_path) : 0;
|
|
DownloadFileInner(url_to_download, incomplete_path, {}, resume_size, headers, expected_size, file_name, error_info);
|
|
if (error_info.empty()) {
|
|
printf("DownloadFile %s success\n", file_name.c_str());
|
|
MoveWithPermissions(incomplete_path, destination_path, error_info);
|
|
} else {
|
|
printf("DownloadFile %s failed\n", file_name.c_str());
|
|
}
|
|
}
|
|
|
|
void RemoteModelDownloader::DownloadFileInner(
|
|
const std::string& url,
|
|
const std::filesystem::path& temp_file,
|
|
const std::unordered_map<std::string, std::string>& proxies,
|
|
size_t resume_size,
|
|
const httplib::Headers& headers,
|
|
const size_t expected_size,
|
|
const std::string& displayed_filename,
|
|
std::string& error_info
|
|
) {
|
|
auto [host, path] = HfApiClient::ParseUrl(url);
|
|
httplib::SSLClient client(host, 443);
|
|
httplib::Headers request_headers(headers.begin(), headers.end());
|
|
if (resume_size > 0) {
|
|
printf("DownloadFile %s resume size %zu", displayed_filename.c_str(), resume_size);
|
|
request_headers.emplace("Range", "bytes=" + std::to_string(resume_size) + "-");
|
|
}
|
|
std::ofstream output(temp_file, std::ios::binary | std::ios::app);
|
|
DownloadProgress progress;
|
|
progress.downloaded = resume_size;
|
|
auto res = client.Get(path, request_headers,
|
|
[&](const httplib::Response& response) {
|
|
auto content_length_str = response.get_header_value("Content-Length");
|
|
if (!content_length_str.empty()) {
|
|
progress.content_length = std::stoull(content_length_str) + resume_size;
|
|
}
|
|
return true;
|
|
},
|
|
[&](const char* data, size_t data_length) {
|
|
output.write(data, data_length);
|
|
progress.downloaded += data_length;
|
|
if (expected_size > 0) {
|
|
double percentage = (static_cast<double>(progress.downloaded) / progress.content_length) * 100.0;
|
|
printf("\rDownloadFile %s progress: %.2f%%", displayed_filename.c_str(), percentage);
|
|
fflush(stdout);
|
|
}
|
|
return true;
|
|
}
|
|
);
|
|
output.flush();
|
|
output.close();
|
|
if (res) {
|
|
if (res->status >= 200 && res->status < 300 || res->status == 416) {
|
|
progress.success = true;
|
|
printf("\n");
|
|
} else {
|
|
error_info = "HTTP error: " + std::to_string(res->status);
|
|
}
|
|
} else {
|
|
error_info = "Connection error: " + std::string(httplib::to_string(res.error()));
|
|
}
|
|
if (!error_info.empty()) {
|
|
printf("HTTP Get Error: %s\n", error_info.c_str());
|
|
}
|
|
}
|
|
|
|
bool RemoteModelDownloader::CheckDiskSpace(size_t required_size, const std::filesystem::path& path) {
|
|
auto space = std::filesystem::space(path);
|
|
if (space.available < required_size) {
|
|
return false;
|
|
}
|
|
return true;
|
|
}
|
|
|
|
void RemoteModelDownloader::MoveWithPermissions(const std::filesystem::path& src,
|
|
const std::filesystem::path& dest,
|
|
std::string& error_info) {
|
|
if (FileUtils::Move(src, dest, error_info)) {
|
|
std::filesystem::permissions(dest, std::filesystem::perms::owner_all);
|
|
}
|
|
}
|
|
|
|
}
|