Retry on HTTP 50x errors

conda/environments/all_cuda-118_arch-aarch64.yaml

@@ -29,6 +29,7 @@ dependencies:
 - pre-commit
 - pytest
 - pytest-cov
+- pytest-timeout
 - python>=3.10,<3.13
 - rangehttpserver
 - rapids-build-backend>=0.3.0,<0.4.0.dev0

conda/environments/all_cuda-118_arch-x86_64.yaml

@@ -31,6 +31,7 @@ dependencies:
 - pre-commit
 - pytest
 - pytest-cov
+- pytest-timeout
 - python>=3.10,<3.13
 - rangehttpserver
 - rapids-build-backend>=0.3.0,<0.4.0.dev0

conda/environments/all_cuda-128_arch-aarch64.yaml

@@ -29,6 +29,7 @@ dependencies:
 - pre-commit
 - pytest
 - pytest-cov
+- pytest-timeout
 - python>=3.10,<3.13
 - rangehttpserver
 - rapids-build-backend>=0.3.0,<0.4.0.dev0

conda/environments/all_cuda-128_arch-x86_64.yaml

@@ -29,6 +29,7 @@ dependencies:
 - pre-commit
 - pytest
 - pytest-cov
+- pytest-timeout
 - python>=3.10,<3.13
 - rangehttpserver
 - rapids-build-backend>=0.3.0,<0.4.0.dev0

cpp/doxygen/main_page.md

@@ -107,10 +107,30 @@ To improve performance of small IO requests, `.pread()` and `.pwrite()` implemen
 This setting can also be controlled by `defaults::gds_threshold()` and `defaults::gds_threshold_reset()`.
 
 #### Size of the Bounce Buffer (KVIKIO_GDS_THRESHOLD)
-KvikIO might have to use intermediate host buffers (one per thread) when copying between files and device memory. Set the environment variable ``KVIKIO_BOUNCE_BUFFER_SIZE`` to the size (in bytes) of these "bounce" buffers. If not set, the default value is 16777216 (16 MiB).
+KvikIO might have to use intermediate host buffers (one per thread) when copying between files and device memory. Set the environment variable `KVIKIO_BOUNCE_BUFFER_SIZE` to the size (in bytes) of these "bounce" buffers. If not set, the default value is 16777216 (16 MiB).
 
 This setting can also be controlled by `defaults::bounce_buffer_size()` and `defaults::bounce_buffer_size_reset()`.
 
+#### HTTP Retries
+
+The behavior when a remote IO read returns a error can be controlled through the `KVIKIO_HTTP_STATUS_CODES` and `KVIKIO_HTTP_MAX_ATTEMPTS` environment variables.
+`KVIKIO_HTTP_STATUS_CODES` controls the status codes to retry, and `KVIKIO_HTTP_MAX_ATTEMPTS` controls the maximum number of attempts to make before throwing an exception.
+
+When a response with a status code in the list of retryable codes is received, KvikIO will wait for some period of time before retrying the request.
+It will keep retrying until reaching the maximum number of attempts.
+
+By default, KvikIO will retry responses with the following status codes:
+
+- 429
+- 500
+- 502
+- 503
+- 504
+
+KvikIO will, by default, make three attempts per read.
+Note that if you're reading a large file that has been split into multiple reads through the KvikIO's task size setting, then *each* task will be retried up to the maximum number of attempts.
+
+These settings can also be controlled by `defaults::http_max_attempts()`, `defaults::http_max_attempts_reset()`, `defaults::http_status_codes()`, and `defaults::http_status_codes_reset()`.
 
 ## Example
 

cpp/include/kvikio/defaults.hpp

@@ -57,6 +57,19 @@ namespace detail {
  */
 CompatMode parse_compat_mode_str(std::string_view compat_mode_str);
 
+/**
+ * @brief Parse a string of comma-separated string of HTTP status codes.
+ *
+ * @param env_var_name The environment variable holding the string.
+ * Used to report errors.
+ * @param status_codes The comma-separated string of HTTP status
+ * codes. Each code should be a 3-digit integer.
+ *
+ * @return The vector with the parsed, integer HTTP status codes.
+ */
+std::vector<int> parse_http_status_codes(std::string_view env_var_name,
+                                         std::string const& status_codes);
+
 }  // namespace detail
 
 template <typename T>
@@ -81,6 +94,9 @@ bool getenv_or(std::string_view env_var_name, bool default_val);
 template <>
 CompatMode getenv_or(std::string_view env_var_name, CompatMode default_val);
 
+template <>
+std::vector<int> getenv_or(std::string_view env_var_name, std::vector<int> default_val);
+
 /**
  * @brief Singleton class of default values used throughout KvikIO.
  *
@@ -92,6 +108,8 @@ class defaults {
   std::size_t _task_size;
   std::size_t _gds_threshold;
   std::size_t _bounce_buffer_size;
+  std::size_t _http_max_attempts;
+  std::vector<int> _http_status_codes;
 
   static unsigned int get_num_threads_from_env();
 
@@ -181,7 +199,7 @@ class defaults {
    * always use the same thread pool however it is possible to change number of
    * threads in the pool (see `kvikio::default::thread_pool_nthreads_reset()`).
    *
-   * @return The the default thread pool instance.
+   * @return The default thread pool instance.
    */
   [[nodiscard]] static BS::thread_pool& thread_pool();
 
@@ -258,6 +276,47 @@ class defaults {
    * @param nbytes The bounce buffer size in bytes.
    */
   static void bounce_buffer_size_reset(std::size_t nbytes);
+
+  /**
+   * @brief Get the maximum number of attempts per remote IO read.
+   *
+   * Set the value using `kvikio::default::http_max_attempts_reset()` or by setting
+   * the `KVIKIO_HTTP_MAX_ATTEMPTS` environment variable. If not set, the value is 3.
+   *
+   * @return The maximum number of remote IO reads to attempt before raising an
+   * error.
+   */
+  [[nodiscard]] static std::size_t http_max_attempts();
+
+  /**
+   * @brief Reset the maximum number of attempts per remote IO read.
+   *
+   * @param attempts The maximum number of attempts to try before raising an error.
+   */
+  static void http_max_attempts_reset(std::size_t attempts);
+
+  /**
+   * @brief The list of HTTP status codes to retry.
+   *
+   * Set the value using `kvikio::default::http_status_codes()` or by setting the
+   * `KVIKIO_HTTP_STATUS_CODES` environment variable. If not set, the default value is
+   *
+   * - 429
+   * - 500
+   * - 502
+   * - 503
+   * - 504
+   *
+   * @return The list of HTTP status codes to retry.
+   */
+  [[nodiscard]] static std::vector<int> const& http_status_codes();
+
+  /**
+   * @brief Reset the list of HTTP status codes to retry.
+   *
+   * @param status_codes The HTTP status codes to retry.
+   */
+  static void http_status_codes_reset(std::vector<int> status_codes);
 };
 
 }  // namespace kvikio

cpp/src/defaults.cpp

@@ -17,6 +17,7 @@
 #include <algorithm>
 #include <cstddef>
 #include <cstdlib>
+#include <regex>
 #include <sstream>
 #include <stdexcept>
 #include <string>
@@ -49,6 +50,28 @@ CompatMode parse_compat_mode_str(std::string_view compat_mode_str)
   return res;
 }
 
+std::vector<int> parse_http_status_codes(std::string_view env_var_name,
+                                         std::string const& status_codes)
+{
+  // Ensure `status_codes` consists only of 3-digit integers separated by commas, allowing spaces.
+  std::regex const check_pattern(R"(^\s*\d{3}\s*(\s*,\s*\d{3}\s*)*$)");
+  if (!std::regex_match(status_codes, check_pattern)) {
+    throw std::invalid_argument(std::string{env_var_name} +
+                                ": invalid format, expected comma-separated integers.");
+  }
+
+  // Match every integer in `status_codes`.
+  std::regex const number_pattern(R"(\d+)");
+
+  // For each match, we push_back `std::stoi(match.str())` into `ret`.
+  std::vector<int> ret;
+  std::transform(std::sregex_iterator(status_codes.begin(), status_codes.end(), number_pattern),
+                 std::sregex_iterator(),
+                 std::back_inserter(ret),
+                 [](std::smatch const& match) -> int { return std::stoi(match.str()); });
+  return ret;
+}
+
 }  // namespace detail
 
 template <>
@@ -91,6 +114,17 @@ CompatMode getenv_or(std::string_view env_var_name, CompatMode default_val)
   return detail::parse_compat_mode_str(env_val);
 }
 
+template <>
+std::vector<int> getenv_or(std::string_view env_var_name, std::vector<int> default_val)
+{
+  auto* const env_val = std::getenv(env_var_name.data());
+  if (env_val == nullptr) { return std::move(default_val); }
+  std::string const int_str(env_val);
+  if (int_str.empty()) { return std::move(default_val); }
+
+  return detail::parse_http_status_codes(env_var_name, int_str);
+}
+
 unsigned int defaults::get_num_threads_from_env()
 {
   int const ret = getenv_or("KVIKIO_NTHREADS", 1);
@@ -132,6 +166,23 @@ defaults::defaults()
     }
     _bounce_buffer_size = env;
   }
+  // Determine the default value of `http_max_attempts`
+  {
+    const ssize_t env         = getenv_or("KVIKIO_HTTP_MAX_ATTEMPTS", 3);
+    auto upperbound_exclusive = sizeof(int);
+    if (env <= 0) {
+      throw std::invalid_argument("KVIKIO_HTTP_MAX_ATTEMPTS has to be a positive integer");
+    } else if (env >= upperbound_exclusive) {
+      throw std::invalid_argument("KVIKIO_HTTP_MAX_ATTEMPTS must be less than " +
+                                  std::to_string(upperbound_exclusive));
+    }
+    _http_max_attempts = env;
+  }
+  // Determine the default value of `http_status_codes`
+  {
+    _http_status_codes =
+      getenv_or("KVIKIO_HTTP_STATUS_CODES", std::vector<int>{429, 500, 502, 503, 504});
+  }
 }
 
 defaults* defaults::instance()
@@ -200,4 +251,25 @@ void defaults::bounce_buffer_size_reset(std::size_t nbytes)
   instance()->_bounce_buffer_size = nbytes;
 }
 
+std::size_t defaults::http_max_attempts() { return instance()->_http_max_attempts; }
+
+void defaults::http_max_attempts_reset(std::size_t attempts)
+{
+  auto upperbound_exclusive = sizeof(int);
+  if (attempts == 0) {
+    throw std::invalid_argument("attempts must be a positive integer");
+  } else if (attempts >= upperbound_exclusive) {
+    throw std::invalid_argument("attempts must be less than " +
+                                std::to_string(upperbound_exclusive));
+  }
+  instance()->_http_max_attempts = attempts;
+}
+
+std::vector<int> const& defaults::http_status_codes() { return instance()->_http_status_codes; }
+
+void defaults::http_status_codes_reset(std::vector<int> status_codes)
+{
+  instance()->_http_status_codes = std::move(status_codes);
+}
+
 }  // namespace kvikio

cpp/src/shim/libcurl.cpp

@@ -14,12 +14,15 @@
  * limitations under the License.
  */
 
+#include <chrono>
 #include <cstring>
 #include <functional>
+#include <iostream>
 #include <memory>
 #include <sstream>
 #include <stdexcept>
 #include <string>
+#include <thread>
 #include <vector>
 
 #include <curl/curl.h>
@@ -116,19 +119,55 @@ CURL* CurlHandle::handle() noexcept { return _handle.get(); }
 
 void CurlHandle::perform()
 {
-  // Perform the curl operation and check for errors.
-  CURLcode err = curl_easy_perform(handle());
-  if (err != CURLE_OK) {
-    std::string msg(_errbuf);  // We can do this because we always initialize `_errbuf` as empty.
-    std::stringstream ss;
-    ss << "curl_easy_perform() error near " << _source_file << ":" << _source_line;
-    if (msg.empty()) {
-      ss << "(" << curl_easy_strerror(err) << ")";
+  long http_code          = 0;
+  int attempt_count       = 1;
+  int base_delay          = 500;   // milliseconds
+  int max_delay           = 4000;  // milliseconds
+  auto http_max_attempts  = kvikio::defaults::http_max_attempts();
+  auto& http_status_codes = kvikio::defaults::http_status_codes();
+
+  while (attempt_count <= http_max_attempts) {
+    auto err = curl_easy_perform(handle());
+    curl_easy_getinfo(handle(), CURLINFO_RESPONSE_CODE, &http_code);
+
+    // Check if we should retry based on HTTP status code
+    if (std::find(http_status_codes.begin(), http_status_codes.end(), http_code) !=
+        http_status_codes.end()) {
+      // Retry only if one of the specified status codes is returned
+      // TODO: Parse the Retry-After header, if it exists.
+      // TODO: configurable maximum wait.
+      if (attempt_count == http_max_attempts) {
+        std::stringstream ss;
+        ss << "kvikio http_max_attempts_reached. attempts=" << http_max_attempts
+           << " reason=" << http_code;
+        throw std::runtime_error(ss.str());
+      } else {
+        // backoff and retry again. With a base value of 500, we retry after
+        // 500ms, 1s, 2s, 4s, ...
+        int backoff_delay = base_delay * (1 << attempt_count);
+        // up to a maximum of `max_delay` seconds.
+        int delay = std::min(max_delay, backoff_delay);
+
+        attempt_count++;
+        std::cout << "Retrying. reason=" << http_code << " after=" << delay
+                  << " attempt=" << attempt_count << " http_max_attempts=" << http_max_attempts
+                  << std::endl;
+        std::this_thread::sleep_for(std::chrono::milliseconds(delay));
+      }
+    } else if (err != CURLE_OK) {
+      std::string msg(_errbuf);  // We can do this because we always initialize `_errbuf` as empty.
+      std::stringstream ss;
+      ss << "curl_easy_perform() error near " << _source_file << ":" << _source_line;
+      if (msg.empty()) {
+        ss << "(" << curl_easy_strerror(err) << ")";
+      } else {
+        ss << "(" << msg << ")";
+      }
+      throw std::runtime_error(ss.str());
     } else {
-      ss << "(" << msg << ")";
+      // No retry needed
+      break;
     }
-    throw std::runtime_error(ss.str());
   }
 }
-
 }  // namespace kvikio

cpp/tests/test_defaults.cpp

@@ -51,3 +51,24 @@ TEST(Defaults, parse_compat_mode_str)
     }
   }
 }
+
+TEST(Defaults, parse_http_status_codes)
+{
+  {
+    std::vector<std::string> inputs{
+      "429,500", "429, 500", " 429,500", "429,  500", "429 ,500", "429,500 "};
+    std::vector<int> expected = {429, 500};
+    for (const auto& input : inputs) {
+      EXPECT_EQ(kvikio::detail::parse_http_status_codes("KVIKIO_HTTP_STATUS_CODES", input),
+                expected);
+    }
+  }
+
+  {
+    std::vector<std::string> inputs{"429,", ",429", "a,b", "429,,500", "429,1000"};
+    for (const auto& input : inputs) {
+      EXPECT_THROW(kvikio::detail::parse_http_status_codes("KVIKIO_HTTP_STATUS_CODES", input),
+                   std::invalid_argument);
+    }
+  }
+}

dependencies.yaml

@@ -420,6 +420,7 @@ dependencies:
           - rapids-dask-dependency==25.4.*,>=0.0.0a0
           - pytest
           - pytest-cov
+          - pytest-timeout
           - rangehttpserver
           - boto3>=1.21.21
       - output_types: [requirements, pyproject]

docs/source/runtime_settings.rst

@@ -42,3 +42,12 @@ Size of the Bounce Buffer ``KVIKIO_BOUNCE_BUFFER_SIZE``
 KvikIO might have to use intermediate host buffers (one per thread) when copying between files and device memory. Set the environment variable ``KVIKIO_BOUNCE_BUFFER_SIZE`` to the size (in bytes) of these "bounce" buffers. If not set, the default value is 16777216 (16 MiB).
 
 This setting can also be controlled by :py:func:`kvikio.defaults.bounce_buffer_size`, :py:func:`kvikio.defaults.bounce_buffer_size_reset`, and :py:func:`kvikio.defaults.set_bounce_buffer_size`.
+
+#### HTTP Retries
+-----------------
+
+The behavior when a remote IO read returns a error can be controlled through the `KVIKIO_HTTP_STATUS_CODES` and `KVIKIO_HTTP_MAX_ATTEMPTS` environment variables.
+
+`KVIKIO_HTTP_STATUS_CODES` controls the status codes to retry and can be controlled by :py:func:`kvikio.defaults.http_status_codes`, :py:func:`kvikio.defaults.http_status_codes_reset`, and :py:func:`kvikio.defaults.set_http_status_codes`.
+
+`KVIKIO_HTTP_MAX_ATTEMPTS` controls the maximum number of attempts to make before throwing an exception and can be controlled by :py:func:`kvikio.defaults.http_max_attempts`, :py:func:`kvikio.defaults.http_max_attempts_reset`, and :py:func:`kvikio.defaults.set_http_max_attempts`.