Remote I/O egress controls and retryable transport failures

cpp/include/kvikio/detail/remote_callback.hpp

@@ -45,6 +45,32 @@ struct CallbackContext {
   void reset_for_retry() noexcept;
 };
 
+/**
+ * @brief Copy with non-temporal stores, which skip fetching the destination cache lines.
+ *
+ * It requires x86-64 with AVX2, and falls back to `memcpy` elsewhere.
+ *
+ * @param dst Destination.
+ * @param src Source.
+ * @param nbytes Number of bytes to copy.
+ */
+void copy_nontemporal(std::byte* dst, std::byte const* src, std::size_t nbytes);
+
+/**
+ * @brief Copy received data into host memory, with non-temporal stores if
+ * `KVIKIO_REMOTE_IO_NONTEMPORAL_COPY` is enabled and the range request is at least
+ * `KVIKIO_REMOTE_IO_NONTEMPORAL_COPY_THRESHOLD` bytes.
+ *
+ * @param dst Destination in host memory.
+ * @param src Source in libcurl's buffer.
+ * @param nbytes Number of bytes to copy.
+ * @param range_size Size of the range request the data belongs to.
+ */
+void copy_received_data(std::byte* dst,
+                        std::byte const* src,
+                        std::size_t nbytes,
+                        std::size_t range_size);
+
 /**
  * @brief Callback for `CURLOPT_WRITEFUNCTION` that copies received bytes directly into a host
  * buffer.

cpp/src/detail/http_retry.cpp

@@ -30,9 +30,17 @@ std::size_t HttpRetryPolicy::max_attempts() const noexcept { return _max_attempt
 
 bool HttpRetryPolicy::is_retryable(CURLcode curl_code, long http_code) const
 {
-  // TODO: Currently the timeout is the only libcurl transport error treated as retryable. Need to
-  // revisit and add more candidates.
-  if (curl_code == CURLE_OPERATION_TIMEDOUT) { return true; }
+  // Transport failures that may succeed on a new attempt.
+  switch (curl_code) {
+    case CURLE_OPERATION_TIMEDOUT:
+    case CURLE_COULDNT_RESOLVE_HOST:
+    case CURLE_COULDNT_CONNECT:
+    case CURLE_RECV_ERROR:
+    case CURLE_SEND_ERROR:
+    case CURLE_PARTIAL_FILE:
+    case CURLE_GOT_NOTHING: return true;
+    default: break;
+  }
   return std::find(_retryable_status_codes.begin(),
                    _retryable_status_codes.end(),
                    static_cast<int>(http_code)) != _retryable_status_codes.end();
@@ -68,27 +76,27 @@ RetryOutcome HttpRetryPolicy::evaluate(CURLcode curl_code,
     return {RetryDecision::FATAL, std::chrono::milliseconds{0}, ss.str()};
   }
 
+  auto const reason = [&] {
+    std::stringstream rs;
+    if (curl_code == CURLE_HTTP_RETURNED_ERROR) {
+      rs << "Got HTTP code " << http_code << ".";
+    } else {
+      rs << "Transport error: " << curl_easy_strerror(curl_code) << ".";
+    }
+    return rs.str();
+  }();
+
   if (attempt >= _max_attempts) {
     std::stringstream ss;
     ss << "KvikIO: HTTP request reached maximum number of attempts (" << _max_attempts
-       << "). Reason: ";
-    if (curl_code == CURLE_OPERATION_TIMEDOUT) {
-      ss << "Operation timed out.";
-    } else {
-      ss << "Got HTTP code " << http_code << ".";
-    }
+       << "). Reason: " << reason;
     return {RetryDecision::EXHAUSTED, std::chrono::milliseconds{0}, ss.str()};
   }
 
   auto const delay_ms = backoff_for(attempt);
   std::stringstream ss;
-  if (curl_code == CURLE_OPERATION_TIMEDOUT) {
-    ss << "KvikIO: Timeout error. Retrying after " << delay_ms.count() << "ms (attempt " << attempt
-       << " of " << _max_attempts << ").";
-  } else {
-    ss << "KvikIO: Got HTTP code " << http_code << ". Retrying after " << delay_ms.count()
-       << "ms (attempt " << attempt << " of " << _max_attempts << ").";
-  }
+  ss << "KvikIO: " << reason << " Retrying after " << delay_ms.count() << "ms (attempt " << attempt
+     << " of " << _max_attempts << ").";
   return {RetryDecision::RETRY, delay_ms, ss.str()};
 }
 

cpp/src/detail/remote_callback.cpp

@@ -3,17 +3,142 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
+#include <algorithm>
 #include <cstddef>
+#include <cstdint>
 #include <cstring>
+#include <stdexcept>
 #include <string>
 
 #include <curl/curl.h>
 
+#if defined(__x86_64__)
+#include <immintrin.h>
+#endif
+
+#include <kvikio/defaults.hpp>
 #include <kvikio/detail/nvtx.hpp>
 #include <kvikio/detail/remote_callback.hpp>
+#include <kvikio/error.hpp>
 
 namespace kvikio::detail {
 
+namespace {
+/**
+ * @brief Whether `KVIKIO_REMOTE_IO_DISCARD_DATA` is enabled.
+ *
+ * For benchmark purposes only. Drop received data instead of copying it into host memory, to
+ * benchmark the network path alone. Reads into device memory are not affected. The destination
+ * buffer is left untouched, with no error raised.
+ */
+bool discard_data_enabled()
+{
+  static bool const value = getenv_or("KVIKIO_REMOTE_IO_DISCARD_DATA", false);
+  return value;
+}
+
+/**
+ * @brief Whether `KVIKIO_REMOTE_IO_NONTEMPORAL_COPY` is enabled.
+ */
+bool nontemporal_copy_enabled()
+{
+  static bool const value = getenv_or("KVIKIO_REMOTE_IO_NONTEMPORAL_COPY", true);
+  return value;
+}
+
+/**
+ * @brief The value of `KVIKIO_REMOTE_IO_NONTEMPORAL_COPY_THRESHOLD`.
+ */
+std::size_t nontemporal_copy_threshold()
+{
+  static std::size_t const value = [] {
+    auto const env = getenv_or("KVIKIO_REMOTE_IO_NONTEMPORAL_COPY_THRESHOLD", ssize_t{1024 * 1024});
+    KVIKIO_EXPECT(env >= 0,
+                  "KVIKIO_REMOTE_IO_NONTEMPORAL_COPY_THRESHOLD has to be a non-negative integer",
+                  std::invalid_argument);
+    return static_cast<std::size_t>(env);
+  }();
+  return value;
+}
+
+/**
+ * @brief Whether the CPU supports AVX2. Always false on non-x86-64 targets.
+ */
+bool cpu_supports_avx2()
+{
+#if defined(__x86_64__)
+  // `__builtin_cpu_supports` is an x86 built-in function, used here to check at runtime if the
+  // machine supports AVX2.
+  static bool const value = __builtin_cpu_supports("avx2");
+  return value;
+#else
+  return false;
+#endif
+}
+
+/**
+ * @brief Copy with AVX2 non-temporal stores on x86-64, and `memcpy` elsewhere. On x86-64, call only
+ * when `cpu_supports_avx2()` is true.
+ */
+#if defined(__x86_64__)
+// Compile with AVX2 for this function alone (regardless of whether the x86-64 machine at
+// compile-time supports AVX2 or not), as if by -mavx2. The rest of the library is compiled with
+// baseline x86-64 options.
+[[gnu::target("avx2")]] void copy_nontemporal_impl(std::byte* dst,
+                                                   std::byte const* src,
+                                                   std::size_t nbytes)
+{
+  // Non-temporal stores need a 32-byte aligned destination.
+  constexpr std::size_t alignment = 32;
+  auto const misaligned           = reinterpret_cast<std::uintptr_t>(dst) % alignment;
+  if (misaligned != 0) {
+    auto const head = std::min(nbytes, alignment - misaligned);
+    std::memcpy(dst, src, head);
+    dst += head;
+    src += head;
+    nbytes -= head;
+  }
+  while (nbytes >= alignment) {
+    _mm256_stream_si256(reinterpret_cast<__m256i*>(dst),
+                        _mm256_loadu_si256(reinterpret_cast<__m256i const*>(src)));
+    dst += alignment;
+    src += alignment;
+    nbytes -= alignment;
+  }
+  if (nbytes != 0) { std::memcpy(dst, src, nbytes); }
+  // Non-temporal stores are weakly ordered.
+  _mm_sfence();
+}
+#else
+void copy_nontemporal_impl(std::byte* dst, std::byte const* src, std::size_t nbytes)
+{
+  std::memcpy(dst, src, nbytes);
+}
+#endif
+
+}  // namespace
+
+void copy_nontemporal(std::byte* dst, std::byte const* src, std::size_t nbytes)
+{
+  if (cpu_supports_avx2()) {
+    copy_nontemporal_impl(dst, src, nbytes);
+  } else {
+    std::memcpy(dst, src, nbytes);
+  }
+}
+
+void copy_received_data(std::byte* dst,
+                        std::byte const* src,
+                        std::size_t nbytes,
+                        std::size_t range_size)
+{
+  if (nontemporal_copy_enabled() && range_size >= nontemporal_copy_threshold()) {
+    copy_nontemporal(dst, src, nbytes);
+  } else {
+    std::memcpy(dst, src, nbytes);
+  }
+}
+
 void CallbackContext::reset_for_retry() noexcept
 {
   offset         = 0;
@@ -30,7 +155,12 @@ std::size_t callback_host_memory(char* data, std::size_t size, std::size_t nmemb
     return CURL_WRITEFUNC_ERROR;
   }
   KVIKIO_NVTX_FUNC_RANGE(nbytes);
-  std::memcpy(ctx->buf + ctx->offset, data, nbytes);
+  if (!discard_data_enabled()) {
+    copy_received_data(reinterpret_cast<std::byte*>(ctx->buf + ctx->offset),
+                       reinterpret_cast<std::byte const*>(data),
+                       nbytes,
+                       ctx->size);
+  }
   ctx->offset += nbytes;
   return nbytes;
 }
@@ -45,7 +175,10 @@ std::size_t callback_pinned_buffer(char* data, std::size_t size, std::size_t nme
     return CURL_WRITEFUNC_ERROR;
   }
   KVIKIO_NVTX_FUNC_RANGE(nbytes);
-  std::memcpy(static_cast<char*>(ctx->pinned_buffer) + ctx->offset, data, nbytes);
+  copy_received_data(static_cast<std::byte*>(ctx->pinned_buffer) + ctx->offset,
+                     reinterpret_cast<std::byte const*>(data),
+                     nbytes,
+                     ctx->size);
   ctx->offset += nbytes;
   return nbytes;
 }

cpp/src/remote_handle.cpp

@@ -117,8 +117,9 @@ class BounceBufferH2D {
    *
    * @param data The host memory source.
    * @param size Number of bytes to write.
+   * @param range_size Size of the range request the data belongs to.
    */
-  void write(char const* data, std::size_t size)
+  void write(char const* data, std::size_t size, std::size_t range_size)
   {
     KVIKIO_NVTX_FUNC_RANGE();
     if (_host_buffer.size() - _host_offset < size) {  // Not enough space left in the bounce buffer
@@ -131,7 +132,10 @@ class BounceBufferH2D {
       // optimize for this case.
       write_to_device(data, size);
     } else if (size > 0) {
-      std::memcpy(_host_buffer.get(_host_offset), data, size);
+      copy_received_data(static_cast<std::byte*>(_host_buffer.get(_host_offset)),
+                         reinterpret_cast<std::byte const*>(data),
+                         size,
+                         range_size);
       _host_offset += size;
     }
   }
@@ -778,7 +782,7 @@ std::size_t callback_device_memory(char* data, std::size_t size, std::size_t nme
   }
   KVIKIO_NVTX_FUNC_RANGE(nbytes);
 
-  ctx->bounce_buffer->write(data, nbytes);
+  ctx->bounce_buffer->write(data, nbytes, ctx->size);
   ctx->offset += nbytes;
   return nbytes;
 }

cpp/src/shim/libcurl.cpp

@@ -6,6 +6,7 @@
 #include <chrono>
 #include <cstddef>
 #include <cstdint>
+#include <cstdlib>
 #include <cstring>
 #include <functional>
 #include <memory>
@@ -125,6 +126,29 @@ CurlHandle::CurlHandle(LibCurl::UniqueHandlePtr handle,
   auto const verbose = getenv_or("KVIKIO_REMOTE_VERBOSE", false);
   if (verbose) { setopt(CURLOPT_VERBOSE, 1L); }
 
+  // Size in bytes of libcurl's receive buffer, one per transfer. When unset, libcurl's own default
+  // of 16 KiB (CURL_MAX_WRITE_SIZE) is used. The value must be between 1 KiB and
+  // CURL_MAX_READ_SIZE (10 MiB in recent versions of curl).
+  static long const buffer_size = [] {
+    auto const env =
+      getenv_or("KVIKIO_REMOTE_IO_BUFFER_SIZE", static_cast<long>(CURL_MAX_WRITE_SIZE));
+    KVIKIO_EXPECT(env >= 1024 && env <= CURL_MAX_READ_SIZE,
+                  "KVIKIO_REMOTE_IO_BUFFER_SIZE has to be an integer between 1024 and " +
+                    std::to_string(CURL_MAX_READ_SIZE),
+                  std::invalid_argument);
+    return env;
+  }();
+  setopt(CURLOPT_BUFFERSIZE, buffer_size);
+
+  // Bind every connection to one network interface, for hosts with several NICs on one subnet.
+  // The value is passed to libcurl verbatim: `<ip>` binds the source address, `if!<name>` binds
+  // the device, and `ifhost!<name>!<ip>` binds both.
+  static std::string const interface_opt = [] {
+    auto const* env = std::getenv("KVIKIO_REMOTE_IO_INTERFACE");
+    return std::string{env == nullptr ? "" : env};
+  }();
+  if (!interface_opt.empty()) { setopt(CURLOPT_INTERFACE, interface_opt.c_str()); }
+
   detail::set_up_ca_paths(*this);
 }
 

cpp/tests/CMakeLists.txt

@@ -167,6 +167,7 @@ kvikio_add_nsys_plugin_test(
 
 if(KvikIO_REMOTE_SUPPORT)
   kvikio_add_test(NAME HTTP_RETRY_TEST SOURCES test_http_retry.cpp utils/env.cpp)
+  kvikio_add_test(NAME REMOTE_CALLBACK_TEST SOURCES test_remote_callback.cpp)
   kvikio_add_test(NAME REMOTE_HANDLE_TEST SOURCES test_remote_handle.cpp utils/env.cpp)
   kvikio_add_test(NAME HDFS_TEST SOURCES test_hdfs.cpp utils/hdfs_helper.cpp)
   kvikio_add_test(NAME TLS_TEST SOURCES test_tls.cpp utils/env.cpp)