Fix download progress reporting

superm1 · superm1 · commit c93edd044ee2 · 2026-04-08T09:56:09.000-05:00
diff --git a/src/cpp/include/lemon/backends/backend_utils.h b/src/cpp/include/lemon/backends/backend_utils.h
@@ -79,7 +79,8 @@ namespace lemon::backends {
         static std::string get_therock_install_dir(const std::string& arch, const std::string& version);
 
         /** Download and install TheRock ROCm tarball for the specified architecture (Linux only) */
-        static void install_therock(const std::string& arch, const std::string& version);
+        static void install_therock(const std::string& arch, const std::string& version,
+                                   DownloadProgressCallback progress_cb = nullptr);
 
         /** Clean up old TheRock versions, keeping only the specified version */
         static void cleanup_old_therock_versions(const std::string& current_version);
diff --git a/src/cpp/server/backend_manager.cpp b/src/cpp/server/backend_manager.cpp
@@ -194,8 +194,8 @@ void install_rocm_stable_runtime_if_needed(const std::string& os,
         http_progress_cb = [&progress_cb, &filename](size_t downloaded, size_t total) -> bool {
             DownloadProgress p;
             p.file = filename;
-            p.file_index = 1;
-            p.total_files = 1;
+            p.file_index = 2;  // ROCm runtime is the second file
+            p.total_files = 2;  // Backend binary + ROCm runtime
             p.bytes_downloaded = downloaded;
             p.bytes_total = total;
             p.percent = total > 0 ? static_cast<int>((downloaded * 100) / total) : 0;
@@ -223,12 +223,12 @@ void install_rocm_stable_runtime_if_needed(const std::string& os,
     if (progress_cb) {
         DownloadProgress p;
         p.file = filename;
-        p.file_index = 1;
-        p.total_files = 1;
+        p.file_index = 2;  // ROCm runtime is the second file
+        p.total_files = 2;  // Backend binary + ROCm runtime
         p.bytes_downloaded = download_result.bytes_downloaded;
         p.bytes_total = download_result.total_bytes;
         p.percent = 100;
-        p.complete = true;
+        p.complete = true;  // Now we can send the completion event
         progress_cb(p);
     }
 }
@@ -249,32 +249,27 @@ void uninstall_rocm_stable_runtime_if_needed(const std::string& os) {
     }
 }
 
-void install_therock_if_needed(const std::string& os, const json& backend_versions) {
+bool will_install_therock(const std::string& os, const json& backend_versions) {
     // TheRock is only needed on Linux for ROCm preview channel.
     if (os != "linux") {
-        return;
+        return false;
     }
 
     // Check if system ROCm is available - if so, don't need TheRock
     if (backends::BackendUtils::is_rocm_installed_system_wide()) {
-        LOG(DEBUG, "BackendManager")
-            << "System ROCm detected, skipping TheRock installation" << std::endl;
-        return;
+        return false;
     }
 
     // Get ROCm architecture
     std::string rocm_arch = SystemInfo::get_rocm_arch();
     if (rocm_arch.empty()) {
-        LOG(DEBUG, "BackendManager")
-            << "No ROCm architecture detected, skipping TheRock installation" << std::endl;
-        return;
+        return false;
     }
 
     // Get TheRock version from backend_versions.json
     if (!backend_versions.contains("therock") || !backend_versions["therock"].contains("version")) {
-        throw std::runtime_error("backend_versions.json is missing 'therock.version'");
+        return false;
     }
-    std::string version = backend_versions["therock"]["version"].get<std::string>();
 
     // Check if this architecture is supported
     if (backend_versions["therock"].contains("architectures") &&
@@ -287,14 +282,24 @@ void install_therock_if_needed(const std::string& os, const json& backend_versio
             }
         }
         if (!arch_supported) {
-            LOG(DEBUG, "BackendManager")
-                << "Architecture " << rocm_arch << " not supported by TheRock" << std::endl;
-            return;
+            return false;
         }
     }
 
+    return true;
+}
+
+void install_therock_if_needed(const std::string& os, const json& backend_versions,
+                              DownloadProgressCallback progress_cb = nullptr) {
+    if (!will_install_therock(os, backend_versions)) {
+        return;
+    }
+
+    std::string rocm_arch = SystemInfo::get_rocm_arch();
+    std::string version = backend_versions["therock"]["version"].get<std::string>();
+
     // Install TheRock for this architecture
-    backends::BackendUtils::install_therock(rocm_arch, version);
+    backends::BackendUtils::install_therock(rocm_arch, version, progress_cb);
 }
 
 } // namespace
@@ -374,15 +379,42 @@ void BackendManager::install_backend(const std::string& recipe, const std::strin
         throw std::runtime_error("[BackendManager] Unknown recipe: " + recipe);
     }
 
+    // Check if we need to install additional runtime components after the main backend
+    bool needs_rocm_stable_runtime = (recipe == "llamacpp" || recipe == "sd-cpp") &&
+                                      resolved_backend == "rocm-stable" &&
+                                      get_current_os() == "linux";
+
+    bool needs_therock = (recipe == "llamacpp" || recipe == "sd-cpp") &&
+                         resolved_backend == "rocm-preview" &&
+                         will_install_therock(get_current_os(), backend_versions_);
+
+    // Wrap the progress callback to adjust file indices if runtime download follows
+    DownloadProgressCallback wrapped_progress_cb;
+    if (progress_cb && (needs_rocm_stable_runtime || needs_therock)) {
+        wrapped_progress_cb = [progress_cb](const DownloadProgress& p) -> bool {
+            DownloadProgress adjusted = p;
+            // Adjust to indicate this is file 1 of 2
+            adjusted.file_index = 1;
+            adjusted.total_files = 2;
+            // Suppress the completion event - we'll send it after the runtime download
+            if (p.complete) {
+                adjusted.complete = false;
+            }
+            return progress_cb(adjusted);
+        };
+    } else {
+        wrapped_progress_cb = progress_cb;
+    }
+
     backends::BackendUtils::install_from_github(
-        *spec, params.version, params.repo, params.filename, resolved_backend, progress_cb);
+        *spec, params.version, params.repo, params.filename, resolved_backend, wrapped_progress_cb);
 
-    if ((recipe == "llamacpp" || recipe == "sd-cpp") && resolved_backend == "rocm-stable") {
+    if (needs_rocm_stable_runtime) {
         install_rocm_stable_runtime_if_needed(get_current_os(), *spec, backend_versions_, progress_cb);
     }
 
-    if ((recipe == "llamacpp" || recipe == "sd-cpp") && resolved_backend == "rocm-preview") {
-        install_therock_if_needed(get_current_os(), backend_versions_);
+    if (needs_therock) {
+        install_therock_if_needed(get_current_os(), backend_versions_, progress_cb);
     }
 }
 
diff --git a/src/cpp/server/backends/backend_utils.cpp b/src/cpp/server/backends/backend_utils.cpp
@@ -483,7 +483,8 @@ namespace lemon::backends {
 #endif
     }
 
-    void BackendUtils::install_therock(const std::string& arch, const std::string& version) {
+    void BackendUtils::install_therock(const std::string& arch, const std::string& version,
+                                       DownloadProgressCallback progress_cb) {
 #ifndef __linux__
         throw std::runtime_error("TheRock is only supported on Linux");
 #else
@@ -524,10 +525,28 @@ namespace lemon::backends {
         LOG(DEBUG, "BackendUtils") << "Downloading TheRock from: " << url << std::endl;
         LOG(DEBUG, "BackendUtils") << "Downloading to: " << tarball_path << std::endl;
 
+        // Create progress callback for download
+        utils::ProgressCallback http_progress_cb;
+        if (progress_cb) {
+            http_progress_cb = [&progress_cb, &filename](size_t downloaded, size_t total) -> bool {
+                DownloadProgress p;
+                p.file = filename;
+                p.file_index = 1;
+                p.total_files = 1;
+                p.bytes_downloaded = downloaded;
+                p.bytes_total = total;
+                p.percent = total > 0 ? static_cast<int>((downloaded * 100) / total) : 0;
+                p.complete = false;
+                return progress_cb(p);
+            };
+        } else {
+            http_progress_cb = utils::create_throttled_progress_callback();
+        }
+
         auto download_result = utils::HttpClient::download_file(
             url,
             tarball_path,
-            utils::create_throttled_progress_callback()
+            http_progress_cb
         );
 
         if (!download_result.success) {
@@ -567,6 +586,19 @@ namespace lemon::backends {
         fs::remove(tarball_path);
         cleanup_old_therock_versions(version);
 
+        // Send completion notification
+        if (progress_cb) {
+            DownloadProgress p;
+            p.file = filename;
+            p.file_index = 1;
+            p.total_files = 1;
+            p.bytes_downloaded = download_result.bytes_downloaded;
+            p.bytes_total = download_result.total_bytes;
+            p.percent = 100;
+            p.complete = true;
+            progress_cb(p);
+        }
+
         LOG(INFO, "BackendUtils") << "TheRock installation complete" << std::endl;
 #endif
     }
