Parallelize memory query for AZURE_BLOB plugin (#1721)

* Parallelize memory query for azure blob

Similar to OBJ plugin, parallelize HEAD requests when batch memory
queries are provided. This specifically improves initial query
times for integrations that checks for a blob existence (e.g.,
checking for KV cache in LMCache). For the most part the implementation,
matches the OBJ plugin to make it easy to update both plugins
if there are adjustments made to the approach in the future.

Signed-off-by: Kyle Knapp <kyleknapp@microsoft.com>

* Clean up memory query wait logic

Specific changes were:
* Make sure we wait for all instances where the callback started
  within the timeout block and not just within the error block
  when queueing work.
* Consolidated logic to single helper function to make it easier
  to follow and maintain in the future.

Signed-off-by: Kyle Knapp <kyleknapp@microsoft.com>

* Add more robust error handling in blob client

Specifically catch any non-standard exceptions to safeguard against
crashes in the asio thread pool. This is primarily defensive as prior
to it, non-standard exceptions had yet to be encountered.

It also adds logging for exceptions in async put and get to have
Azure exceptions visible in NIXL error logs. Alternative was to turn
on AZURE_LOG_LEVEL=verbose to know what the exceptions were

* Add assertions for null clients in tests

Prevents not constructing a client and trying to use a null client
when building a test blob engine for testing.

Signed-off-by: Kyle Knapp <kyleknapp@microsoft.com>

---------

Signed-off-by: Kyle Knapp <kyleknapp@microsoft.com>
Co-authored-by: Adit Ranadive <aranadive@nvidia.com>
Co-authored-by: ovidiusm <ovidium@nvidia.com>

Author: 3 people

src/plugins/azure_blob/azure_blob_backend.cpp

@@ -27,6 +27,7 @@
 #include <vector>
 #include <chrono>
 #include <algorithm>
+#include <atomic>
 
 namespace {
 
@@ -110,6 +111,43 @@ class nixlAzureBlobMetadata : public nixlBackendMD {
     std::string blobName;
 };
 
+struct QueryMemDescriptorState {
+    std::shared_ptr<std::promise<void>> promise;
+    std::shared_ptr<std::atomic<bool>> completed;
+};
+
+// Wait for all asynchronous memory queries to complete with a per-descriptor
+// wait timeout.
+void
+waitForQueryMemFutures(std::vector<std::future<void>> &futures,
+                       std::vector<QueryMemDescriptorState> &states,
+                       std::vector<nixl_query_resp_t> &resp,
+                       std::atomic<bool> &has_error,
+                       std::chrono::seconds timeout) {
+    for (size_t i = 0; i < futures.size(); ++i) {
+        if (futures[i].wait_for(timeout) == std::future_status::timeout) {
+            if (!states[i].completed->exchange(true)) {
+                // We've timed out and the callback never fired. Mark as
+                // completed and resolve the promise ourselves.
+                NIXL_ERROR << "checkBlobExistsAsync timed out for descriptor " << i;
+                resp[i] = std::nullopt;
+                has_error.store(true, std::memory_order_relaxed);
+                try {
+                    states[i].promise->set_value();
+                }
+                catch (...) {
+                }
+            } else {
+                // Handles case where we timed out but the callback started
+                // at the same time as the timeout and won the exchange.
+                // Wait for it to finish so its writes to resp[i]/has_error
+                // complete before stack-captured references go out of scope.
+                futures[i].wait();
+            }
+        }
+    }
+}
+
 } // namespace
 
 // -----------------------------------------------------------------------------
@@ -172,19 +210,60 @@ nixlAzureBlobEngine::deregisterMem(nixlBackendMD *meta) {
 nixl_status_t
 nixlAzureBlobEngine::queryMem(const nixl_reg_dlist_t &descs,
                               std::vector<nixl_query_resp_t> &resp) const {
-    resp.reserve(descs.descCount());
+    resp.assign(descs.descCount(), std::nullopt);
+
+    std::vector<std::future<void>> futures;
+    futures.reserve(descs.descCount());
+    std::vector<QueryMemDescriptorState> states;
+    states.reserve(descs.descCount());
+    std::atomic<bool> has_error{false};
+
+    constexpr auto kQueryTimeout = std::chrono::seconds(30);
 
     try {
-        for (auto &desc : descs)
-            resp.emplace_back(blobClient_->checkBlobExists(desc.metaInfo) ?
-                                  nixl_query_resp_t{nixl_b_params_t{}} :
-                                  std::nullopt);
+        // Launch async existence checks for each descriptor
+        for (int i = 0; i < descs.descCount(); ++i) {
+            auto &desc = descs[i];
+            auto state = QueryMemDescriptorState{std::make_shared<std::promise<void>>(),
+                                                 std::make_shared<std::atomic<bool>>(false)};
+
+            blobClient_->checkBlobExistsAsync(
+                desc.metaInfo,
+                [&resp, &has_error, i, promise = state.promise, completed = state.completed](
+                    std::optional<bool> exists) {
+                    if (completed->exchange(true)) {
+                        return;
+                    }
+                    if (!exists.has_value()) {
+                        resp[i] = std::nullopt;
+                        has_error.store(true, std::memory_order_relaxed);
+                    } else {
+                        resp[i] = *exists ? nixl_query_resp_t{nixl_b_params_t{}} : std::nullopt;
+                    }
+                    try {
+                        promise->set_value();
+                    }
+                    catch (...) {
+                    }
+                });
+
+            futures.push_back(state.promise->get_future());
+            states.push_back(std::move(state));
+        }
     }
-    catch (const std::runtime_error &e) {
+    catch (const std::exception &e) {
+        waitForQueryMemFutures(futures, states, resp, has_error, kQueryTimeout);
         NIXL_ERROR << "Failed to query memory: " << e.what();
         return NIXL_ERR_BACKEND;
     }
 
+    waitForQueryMemFutures(futures, states, resp, has_error, kQueryTimeout);
+
+    if (has_error.load(std::memory_order_relaxed)) {
+        NIXL_ERROR << "Failed to query memory: one or more existence checks failed";
+        return NIXL_ERR_BACKEND;
+    }
+
     return NIXL_SUCCESS;
 }
 

src/plugins/azure_blob/azure_blob_client.cpp

@@ -164,6 +164,11 @@ azureBlobClient::putBlobAsync(std::string_view blob_name,
             callback(true);
         }
         catch (const std::exception &e) {
+            NIXL_ERROR << "putBlobAsync error: " << e.what();
+            callback(false);
+        }
+        catch (...) {
+            NIXL_ERROR << "putBlobAsync: unknown exception";
             callback(false);
         }
     });
@@ -189,24 +194,40 @@ azureBlobClient::getBlobAsync(std::string_view blob_name,
             callback(true);
         }
         catch (const std::exception &e) {
+            NIXL_ERROR << "getBlobAsync error: " << e.what();
+            callback(false);
+        }
+        catch (...) {
+            NIXL_ERROR << "getBlobAsync: unknown exception";
             callback(false);
         }
     });
 }
 
-bool
-azureBlobClient::checkBlobExists(std::string_view blob_name) {
-    auto blobClient = blobContainerClient_->GetBlockBlobClient(std::string(blob_name));
-    Azure::Storage::Blobs::GetBlobPropertiesOptions options;
-    try {
-        blobClient.GetProperties(options);
-    }
-    catch (const Azure::Core::RequestFailedException &e) {
-        if (e.StatusCode == Azure::Core::Http::HttpStatusCode::NotFound) {
-            return false;
-        } else {
-            throw std::runtime_error("Failed to check if blob exists: " + std::string(e.what()));
+void
+azureBlobClient::checkBlobExistsAsync(std::string_view blob_name, check_blob_callback_t callback) {
+    std::string blob_name_str(blob_name);
+    asio::post(*executor_, [this, blob_name_str, callback]() {
+        try {
+            auto blobClient = blobContainerClient_->GetBlockBlobClient(blob_name_str);
+            blobClient.GetProperties();
+            callback(true);
         }
-    }
-    return true;
+        catch (const Azure::Core::RequestFailedException &e) {
+            if (e.StatusCode == Azure::Core::Http::HttpStatusCode::NotFound) {
+                callback(false);
+                return;
+            }
+            NIXL_ERROR << "checkBlobExistsAsync error: " << e.what();
+            callback(std::nullopt);
+        }
+        catch (const std::exception &e) {
+            NIXL_ERROR << "checkBlobExistsAsync error: " << e.what();
+            callback(std::nullopt);
+        }
+        catch (...) {
+            NIXL_ERROR << "checkBlobExistsAsync: unknown exception";
+            callback(std::nullopt);
+        }
+    });
 }

src/plugins/azure_blob/azure_blob_client.h

@@ -21,6 +21,7 @@
 
 #include <functional>
 #include <memory>
+#include <optional>
 #include <string_view>
 #include <cstdint>
 #include <azure/storage/blobs.hpp>
@@ -29,6 +30,7 @@
 
 using put_blob_callback_t = std::function<void(bool success)>;
 using get_blob_callback_t = std::function<void(bool success)>;
+using check_blob_callback_t = std::function<void(std::optional<bool> exists)>;
 
 /**
  * Abstract interface for Azure Blob client operations.
@@ -76,12 +78,15 @@ class iBlobClient {
                  get_blob_callback_t callback) = 0;
 
     /**
-     * Check if the blob exists.
+     * Asynchronously check if the blob exists.
+     * Uses check_blob_callback_t: the callback receives a std::optional<bool>
+     * that is true if the blob exists, false if it does not, or std::nullopt
+     * if the request failed (indicating an error).
      * @param blob_name The blob name
-     * @return true if the blob exists, false otherwise
+     * @param callback Callback function invoked with the existence check result
      */
-    virtual bool
-    checkBlobExists(std::string_view blob_name) = 0;
+    virtual void
+    checkBlobExistsAsync(std::string_view blob_name, check_blob_callback_t callback) = 0;
 };
 
 /**
@@ -114,8 +119,8 @@ class azureBlobClient : public iBlobClient {
                  size_t offset,
                  get_blob_callback_t callback) override;
 
-    bool
-    checkBlobExists(std::string_view blob_name) override;
+    void
+    checkBlobExistsAsync(std::string_view blob_name, check_blob_callback_t callback) override;
 
 private:
     std::shared_ptr<asio::thread_pool> executor_;