abstract into a dmlc stream.

Author: trivialfis

doc/tutorials/external_memory.rst

@@ -8,6 +8,18 @@ The feature is still experimental and not yet ready for production use.  In this
 we will introduce both methods.  Please note that training on data from external memory is
 not supported by ``exact`` tree method.
 
+.. warning::
+
+   The implementation of external memory uses ``mmap`` and is not tested against errors
+   like disconnected network devices. (`SIGBUS`)
+
+.. note::
+
+   When externel memory is used, the CPU training performance is IO bounded. Meaning, the
+   training speed almost exclusively determined by the disk IO speed. For GPU, please read
+   on and see the gradient-based sampling with external memory. During benchmark, we used
+   a NVME connected to a PCIE slot, the performance is "usable" with ``hist`` on CPU.
+
 *************
 Data Iterator
 *************

rabit/include/rabit/internal/io.h

@@ -60,7 +60,7 @@ struct MemoryFixSizeBuffer : public SeekStream {
     return curr_ptr_ == buffer_size_;
   }
 
- private:
+ protected:
   /*! \brief in memory buffer */
   char *p_buffer_;
   /*! \brief current pointer */

src/common/io.h

@@ -9,10 +9,15 @@
 #define XGBOOST_COMMON_IO_H_
 
 #include <dmlc/io.h>
+#include <fcntl.h>  // for open, O_RDONLY
 #include <rabit/rabit.h>
-#include <string>
+#include <sys/mman.h>  // for mmap, munmap
+#include <unistd.h>    // for close
+#include <xgboost/string_view.h>
+
 #include <cstring>
 #include <fstream>
+#include <string>
 
 #include "common.h"
 
@@ -127,6 +132,48 @@ inline std::string ReadAll(std::string const &path) {
   return content;
 }
 
+/**
+ * \brief Private mmap file, copy-on-write
+ */
+class PrivateMmapStream : public MemoryFixSizeBuffer {
+  std::int32_t fd_;
+  std::string path_;
+
+  void* Open(StringView path, bool read_only, std::size_t offset, std::size_t length) {
+    fd_ = open(path.c_str(), O_RDONLY);
+    CHECK_GE(fd_, 0) << "Failed to open:" << path << ". " << strerror(errno);
+
+    char* ptr{nullptr};
+    int prot{PROT_READ};
+    if (!read_only) {
+      prot |= PROT_WRITE;
+    }
+#if defined(__linux__)
+    ptr = reinterpret_cast<char*>(mmap64(nullptr, length, prot, MAP_PRIVATE, fd_, offset));
+#elif defined(__APPLE__)
+    CHECK_LE(offset, std::numeric_limits<off_t>::max())
+        << "File size has exceeded the limit on macos.";
+    ptr = reinterpret_cast<char*>(mmap(nullptr, length, prot, MAP_PRIVATE, fd_, offset));
+#else
+    // fixme: not yet implemented
+    ptr = reinterpret_cast<char*>(mmap(nullptr, length, prot, MAP_PRIVATE, fd_, offset));
+#endif  // defined(__linux__)
+    CHECK_NE(ptr, MAP_FAILED) << "Failed to map: " << path << ". " << strerror(errno);
+    return ptr;
+  }
+
+ public:
+  explicit PrivateMmapStream(std::string path, bool read_only, std::size_t offset,
+                             std::size_t length)
+      : MemoryFixSizeBuffer{Open(StringView{path}, read_only, offset, length), length},
+        path_{path} {}
+
+  ~PrivateMmapStream() override {
+    CHECK_NE(munmap(p_buffer_, buffer_size_), -1)
+        << "Faled to munmap." << path_ << ". " << strerror(errno);
+    CHECK_NE(close(fd_), -1) << "Faled to close: " << path_ << ". " << strerror(errno);
+  }
+};
 }  // namespace common
 }  // namespace xgboost
 #endif  // XGBOOST_COMMON_IO_H_

src/data/sparse_page_source.h

@@ -123,9 +123,9 @@ class SparsePageSourceImpl : public BatchIteratorImpl<S> {
 
     size_t n_prefetch_batches = std::min(kPreFetch, n_batches_);
     CHECK_GT(n_prefetch_batches, 0) << "total batches:" << n_batches_;
-    size_t fetch_it = count_;
+    std::size_t fetch_it = count_;
 
-    for (size_t i = 0; i < n_prefetch_batches; ++i, ++fetch_it) {
+    for (std::size_t i = 0; i < n_prefetch_batches; ++i, ++fetch_it) {
       fetch_it %= n_batches_;  // ring
       if (ring_->at(fetch_it).valid()) {
         continue;
@@ -143,21 +143,9 @@ class SparsePageSourceImpl : public BatchIteratorImpl<S> {
         std::uint64_t offset = self->cache_info_->offset.at(fetch_it);
         std::uint64_t length = self->cache_info_->bytes.at(fetch_it);
 
-        // mmap
-        auto fd = open(n.c_str(), O_RDONLY);
-        CHECK_GE(fd, 0) << "Failed to open:" << n << ". " << strerror(errno);
-        auto ptr = mmap64(nullptr, length, PROT_READ, MAP_PRIVATE, fd, offset);
-        CHECK_NE(ptr, MAP_FAILED) << "Failed to map: " << n << ". " << strerror(errno);
-
-        // read page
-        auto fi = common::MemoryFixSizeBuffer(ptr, length);
-        CHECK(fmt->Read(page.get(), &fi));
+        auto fi = std::make_unique<common::PrivateMmapStream>(n, true, offset, length);
+        CHECK(fmt->Read(page.get(), fi.get()));
         LOG(INFO) << "Read a page in " << timer.ElapsedSeconds() << " seconds.";
-
-        // cleanup
-        CHECK_NE(munmap(ptr, length), -1) << "Faled to munmap: " << n << ". " << strerror(errno);
-        CHECK_NE(close(fd), -1) << "Faled to close: " << n << ". " << strerror(errno);
-
         return page;
       });
     }