Add docker protocol support for llama-server model loading

- To pull and run models via: llama-server -d ai/smollm2:135M-Q4_K_M
- Implement resumable downloads in common_download_file_single function
- Add detection of partial download files (.downloadInProgress)
- Check server support for HTTP Range requests via Accept-Ranges header
- Implement HTTP Range request with "bytes=<start>-" header
- Open files in append mode when resuming vs create mode for new downloads
- Maintain backwards compatibility with existing functionality
- Add some validators and sanitizers for Docker Model urls and metadata

Signed-off-by: Eric Curtin <[email protected]>

Author: ericcurtin

common/arg.cpp

@@ -220,12 +220,24 @@ struct curl_slist_ptr {
 #define CURL_MAX_RETRY 3
 #define CURL_RETRY_DELAY_SECONDS 2
 
-static bool curl_perform_with_retry(const std::string & url, CURL * curl, int max_attempts, int retry_delay_seconds, const char * method_name) {
+static bool curl_perform_with_retry(const std::string & url,
+                                    CURL *              curl,
+                                    int                 max_attempts,
+                                    int                 retry_delay_seconds,
+                                    const char *        method_name,
+                                    const std::string & path_temporary = "") {
     int remaining_attempts = max_attempts;
 
     while (remaining_attempts > 0) {
         LOG_INF("%s: %s %s (attempt %d of %d)...\n", __func__ , method_name, url.c_str(), max_attempts - remaining_attempts + 1, max_attempts);
 
+        if (std::filesystem::exists(path_temporary)) {
+            const long partial_size = static_cast<long>(std::filesystem::file_size(path_temporary));
+            LOG_INF("%s: server supports range requests, resuming download from byte %ld\n", __func__, partial_size);
+            const std::string range_str = std::to_string(partial_size) + "-";
+            curl_easy_setopt(curl, CURLOPT_RANGE, range_str.c_str());
+        }
+
         CURLcode res = curl_easy_perform(curl);
         if (res == CURLE_OK) {
             return true;
@@ -246,16 +258,16 @@ static bool curl_perform_with_retry(const std::string & url, CURL * curl, int ma
 
 // download one single file from remote URL to local path
 static bool common_download_file_single(const std::string & url, const std::string & path, const std::string & bearer_token, bool offline) {
-    // Check if the file already exists locally
-    auto file_exists = std::filesystem::exists(path);
-
     // If the file exists, check its JSON metadata companion file.
     std::string metadata_path = path + ".json";
     nlohmann::json metadata; // TODO @ngxson : get rid of this json, use regex instead
     std::string etag;
     std::string last_modified;
 
-    if (file_exists) {
+    // Check if the file already exists locally
+    auto file_exists = std::filesystem::exists(path);
+    auto json_file_exists = std::filesystem::exists(metadata_path);
+    if (json_file_exists) {
         if (offline) {
             LOG_INF("%s: using cached file (offline mode): %s\n", __func__, path.c_str());
             return true; // skip verification/downloading
@@ -289,6 +301,7 @@ static bool common_download_file_single(const std::string & url, const std::stri
     struct common_load_model_from_url_headers {
         std::string etag;
         std::string last_modified;
+        std::string accept_ranges;
     };
 
     common_load_model_from_url_headers headers;
@@ -328,7 +341,7 @@ static bool common_download_file_single(const std::string & url, const std::stri
         static std::regex header_regex("([^:]+): (.*)\r\n");
         static std::regex etag_regex("ETag", std::regex_constants::icase);
         static std::regex last_modified_regex("Last-Modified", std::regex_constants::icase);
-
+        static std::regex accept_ranges_regex("Accept-Ranges", std::regex_constants::icase);
         std::string header(buffer, n_items);
         std::smatch match;
         if (std::regex_match(header, match, header_regex)) {
@@ -338,6 +351,8 @@ static bool common_download_file_single(const std::string & url, const std::stri
                 headers->etag = value;
             } else if (std::regex_match(key, match, last_modified_regex)) {
                 headers->last_modified = value;
+            } else if (std::regex_match(key, match, accept_ranges_regex)) {
+                headers->accept_ranges = value;
             }
         }
         return n_items;
@@ -366,28 +381,48 @@ static bool common_download_file_single(const std::string & url, const std::stri
 
     // if head_request_ok is false, we don't have the etag or last-modified headers
     // we leave should_download as-is, which is true if the file does not exist
+    bool should_download_from_scratch = false;
     if (head_request_ok) {
         // check if ETag or Last-Modified headers are different
         // if it is, we need to download the file again
         if (!etag.empty() && etag != headers.etag) {
             LOG_WRN("%s: ETag header is different (%s != %s): triggering a new download\n", __func__, etag.c_str(), headers.etag.c_str());
             should_download = true;
+            should_download_from_scratch = true;
         } else if (!last_modified.empty() && last_modified != headers.last_modified) {
             LOG_WRN("%s: Last-Modified header is different (%s != %s): triggering a new download\n", __func__, last_modified.c_str(), headers.last_modified.c_str());
             should_download = true;
+            should_download_from_scratch = true;
         }
     }
 
     if (should_download) {
-        std::string path_temporary = path + ".downloadInProgress";
-        if (file_exists) {
+        if (file_exists &&
+            headers.accept_ranges.empty()) {  // Resumable downloads not supported, delete and start again.
             LOG_WRN("%s: deleting previous downloaded file: %s\n", __func__, path.c_str());
             if (remove(path.c_str()) != 0) {
                 LOG_ERR("%s: unable to delete file: %s\n", __func__, path.c_str());
                 return false;
             }
         }
 
+        std::string path_temporary = path + ".downloadInProgress";
+        if (should_download_from_scratch) {
+            if (std::filesystem::exists(path_temporary)) {
+                if (remove(path_temporary.c_str()) != 0) {
+                    LOG_ERR("%s: unable to delete file: %s\n", __func__, path_temporary.c_str());
+                    return false;
+                }
+            }
+
+            if (std::filesystem::exists(path)) {
+                if (remove(path.c_str()) != 0) {
+                    LOG_ERR("%s: unable to delete file: %s\n", __func__, path.c_str());
+                    return false;
+                }
+            }
+        }
+
         // Set the output file
 
         struct FILE_deleter {
@@ -396,7 +431,8 @@ static bool common_download_file_single(const std::string & url, const std::stri
             }
         };
 
-        std::unique_ptr<FILE, FILE_deleter> outfile(fopen(path_temporary.c_str(), "wb"));
+        // Always open file in append mode could be resuming
+        std::unique_ptr<FILE, FILE_deleter> outfile(fopen(path_temporary.c_str(), "ab"));
         if (!outfile) {
             LOG_ERR("%s: error opening local file for writing: %s\n", __func__, path.c_str());
             return false;
@@ -431,7 +467,19 @@ static bool common_download_file_single(const std::string & url, const std::stri
         // start the download
         LOG_INF("%s: trying to download model from %s to %s (server_etag:%s, server_last_modified:%s)...\n", __func__,
             llama_download_hide_password_in_url(url).c_str(), path.c_str(), headers.etag.c_str(), headers.last_modified.c_str());
-        bool was_perform_successful = curl_perform_with_retry(url, curl.get(), CURL_MAX_RETRY, CURL_RETRY_DELAY_SECONDS, "GET");
+
+        // Write the updated JSON metadata file.
+        metadata.update({
+            {"url", url},
+            {"etag", headers.etag},
+            {"lastModified", headers.last_modified}
+        });
+        write_file(metadata_path, metadata.dump(4));
+        LOG_DBG("%s: file metadata saved: %s\n", __func__, metadata_path.c_str());
+
+        const bool was_perform_successful =
+            curl_perform_with_retry(url, curl.get(), CURL_MAX_RETRY, CURL_RETRY_DELAY_SECONDS, "GET",
+                                    headers.accept_ranges.empty() ? "" : path_temporary);
         if (!was_perform_successful) {
             return false;
         }
@@ -446,15 +494,6 @@ static bool common_download_file_single(const std::string & url, const std::stri
         // Causes file to be closed explicitly here before we rename it.
         outfile.reset();
 
-        // Write the updated JSON metadata file.
-        metadata.update({
-            {"url", url},
-            {"etag", headers.etag},
-            {"lastModified", headers.last_modified}
-        });
-        write_file(metadata_path, metadata.dump(4));
-        LOG_DBG("%s: file metadata saved: %s\n", __func__, metadata_path.c_str());
-
         if (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());
             return false;
@@ -745,6 +784,118 @@ std::pair<long, std::vector<char>> common_remote_get_content(const std::string &
 
 #endif // LLAMA_USE_CURL
 
+//
+// Docker registry functions
+//
+
+static std::string common_docker_get_token(const std::string & repo) {
+    std::string url = "https://auth.docker.io/token?service=registry.docker.io&scope=repository:" + repo + ":pull";
+
+    common_remote_params params;
+    auto                 res = common_remote_get_content(url, params);
+
+    if (res.first != 200) {
+        throw std::runtime_error("Failed to get Docker registry token, HTTP code: " + std::to_string(res.first));
+    }
+
+    std::string            response_str(res.second.begin(), res.second.end());
+    nlohmann::ordered_json response = nlohmann::ordered_json::parse(response_str);
+
+    if (!response.contains("token")) {
+        throw std::runtime_error("Docker registry token response missing 'token' field");
+    }
+
+    return response["token"].get<std::string>();
+}
+
+static std::string common_docker_resolve_model(const std::string & docker) {
+    // Parse ai/smollm2:135M-Q4_K_M
+    size_t      colon_pos = docker.find(':');
+    std::string repo, tag;
+    if (colon_pos != std::string::npos) {
+        repo = docker.substr(0, colon_pos);
+        tag  = docker.substr(colon_pos + 1);
+    } else {
+        repo = docker;
+        tag  = "latest";
+    }
+
+    LOG_INF("Downloading Docker Model: %s:%s\n", repo.c_str(), tag.c_str());
+    try {
+        // --- helper: digest validation ---
+        auto validate_oci_digest = [](const std::string & digest) -> std::string {
+            // Expected: algo:hex ; start with sha256 (64 hex chars)
+            // You can extend this map if supporting other algorithms in future.
+            static const std::regex re("^sha256:([a-fA-F0-9]{64})$");
+            std::smatch m;
+            if (!std::regex_match(digest, m, re)) {
+                throw std::runtime_error("Invalid OCI digest format received in manifest: " + digest);
+            }
+            // normalize hex to lowercase
+            std::string normalized = digest;
+            std::transform(normalized.begin()+7, normalized.end(), normalized.begin()+7, [](unsigned char c){
+                return std::tolower(c);
+            });
+            return normalized;
+        };
+
+        std::string token = common_docker_get_token(repo);  // Get authentication token
+
+        // Get manifest
+        std::string          manifest_url = "https://registry-1.docker.io/v2/" + repo + "/manifests/" + tag;
+        common_remote_params manifest_params;
+        manifest_params.headers.push_back("Authorization: Bearer " + token);
+        manifest_params.headers.push_back(
+            "Accept: application/vnd.docker.distribution.manifest.v2+json,application/vnd.oci.image.manifest.v1+json");
+        auto manifest_res = common_remote_get_content(manifest_url, manifest_params);
+        if (manifest_res.first != 200) {
+            throw std::runtime_error("Failed to get Docker manifest, HTTP code: " + std::to_string(manifest_res.first));
+        }
+
+        std::string            manifest_str(manifest_res.second.begin(), manifest_res.second.end());
+        nlohmann::ordered_json manifest = nlohmann::ordered_json::parse(manifest_str);
+        std::string            gguf_digest;  // Find the GGUF layer
+        if (manifest.contains("layers")) {
+            for (const auto & layer : manifest["layers"]) {
+                if (layer.contains("mediaType")) {
+                    std::string media_type = layer["mediaType"].get<std::string>();
+                    if (media_type == "application/vnd.docker.ai.gguf.v3" ||
+                        media_type.find("gguf") != std::string::npos) {
+                        gguf_digest = layer["digest"].get<std::string>();
+                        break;
+                    }
+                }
+            }
+        }
+
+        if (gguf_digest.empty()) {
+            throw std::runtime_error("No GGUF layer found in Docker manifest");
+        }
+
+        // Validate & normalize digest
+        gguf_digest = validate_oci_digest(gguf_digest);
+        LOG_DBG("Using validated digest: %s\n", gguf_digest.c_str());
+
+        // Prepare local filename
+        std::string model_filename = repo;
+        std::replace(model_filename.begin(), model_filename.end(), '/', '_');
+        model_filename += "_" + tag + ".gguf";
+        std::string local_path = fs_get_cache_file(model_filename);
+
+        // Download the blob using common_download_file_single with is_docker=true
+        std::string blob_url = "https://registry-1.docker.io/v2/" + repo + "/blobs/" + gguf_digest;
+        if (!common_download_file_single(blob_url, local_path, token, false)) {
+            throw std::runtime_error("Failed to download Docker Model");
+        }
+
+        LOG_INF("%s: Downloaded Docker Model to: %s\n", __func__, local_path.c_str());
+        return local_path;
+    } catch (const std::exception & e) {
+        LOG_ERR("%s: Docker Model download failed: %s\n", __func__, e.what());
+        throw;
+    }
+}
+
 //
 // utils
 //
@@ -795,7 +946,9 @@ static handle_model_result common_params_handle_model(
     handle_model_result result;
     // handle pre-fill default model path and url based on hf_repo and hf_file
     {
-        if (!model.hf_repo.empty()) {
+        if (!model.docker_repo.empty()) {  // Handle Docker URLs by resolving them to local paths
+            model.path = common_docker_resolve_model(model.docker_repo);
+        } else if (!model.hf_repo.empty()) {
             // short-hand to avoid specifying --hf-file -> default it to --model
             if (model.hf_file.empty()) {
                 if (model.path.empty()) {
@@ -2636,6 +2789,15 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
             params.model.url = value;
         }
     ).set_env("LLAMA_ARG_MODEL_URL"));
+    add_opt(common_arg(
+        { "-d", "-dr", "--docker-repo" }, "<repo>/<model>[:quant]",
+        "Docker Hub model repository; quant is optional, default to latest.\n"
+        "example: ai/smollm2:135M-Q4_K_M\n"
+        "(default: unused)",
+        [](common_params & params, const std::string & value) {
+            params.model.docker_repo = value;
+        }
+    ).set_env("LLAMA_ARG_DOCKER"));
     add_opt(common_arg(
         {"-hf", "-hfr", "--hf-repo"}, "<user>/<model>[:quant]",
         "Hugging Face model repository; quant is optional, case-insensitive, default to Q4_K_M, or falls back to the first file in the repo if Q4_K_M doesn't exist.\n"

common/common.h

@@ -193,10 +193,11 @@ struct common_params_sampling {
 };
 
 struct common_params_model {
-    std::string path    = ""; // model local path                                           // NOLINT
-    std::string url     = ""; // model url to download                                      // NOLINT
-    std::string hf_repo = ""; // HF repo                                                    // NOLINT
-    std::string hf_file = ""; // HF file                                                    // NOLINT
+    std::string path        = ""; // model local path                                       // NOLINT
+    std::string url         = ""; // model url to download                                  // NOLINT
+    std::string hf_repo     = ""; // HF repo                                                // NOLINT
+    std::string hf_file     = ""; // HF file                                                // NOLINT
+    std::string docker_repo = ""; // Docker Model url to download                           // NOLINT
 };
 
 struct common_params_speculative {