[fix][dataloader] use file descripor instead of file system

Author: xysheng-baidu

paddle/common/flags.cc

@@ -1345,6 +1345,19 @@ PHI_DEFINE_EXPORTED_bool(use_shm_cache,
  false,
  "Use shm cache in mmap_allocator.");
 
+/**
+ * mmap_allocator related FLAG
+ * Name: use_file_descripor
+ * Since Version: 2.6.2
+ * Value Range: bool, default=true
+ * Example:
+ * Note: . If True, mmap_allocator will use file descripor to open shared memory
+ * operation.
+ */
+PHI_DEFINE_EXPORTED_bool(use_file_descripor,
+ true,
+ "Use file descripor in mmap_allocator.");
+
 /**
  * Tensor operants related FLAG
  * Name: tensor_operants_mode

paddle/fluid/memory/allocation/mmap_allocator.cc

@@ -54,11 +54,14 @@ struct CountInfo {
  std::atomic<int> refcount;
 };
 
-void AllocateMemoryMap(
- std::string filename, int flags, size_t size, void **map_ptr_, int *fd_) {
+void AllocateMemoryMap(std::string filename,
+ int *shared_fd,
+ int flags,
+ size_t size,
+ void **map_ptr_) {
  // TODO(@ZHUI): support win32
  int file_flags = 0;
- int fd = -1;
+ int fd = *shared_fd;
  if (flags & MAPPED_SHAREDMEM) {
  file_flags = O_RDWR | O_CREAT;
  } else {
@@ -71,7 +74,7 @@ void AllocateMemoryMap(
  file_flags &= ~O_CREAT;
  }
 
- if (!(flags & MAPPED_FROMFD)) {
+ if (!(flags & MAPPED_FROMFD) && fd == -1) {
  if (flags & MAPPED_SHAREDMEM) {
  fd = shm_open(filename.c_str(), file_flags, (mode_t)0600);
  PADDLE_ENFORCE_NE(
@@ -83,8 +86,6 @@ void AllocateMemoryMap(
  VLOG(6) << "shm_open: " << filename;
  MemoryMapFdSet::Instance().Insert(filename);
  }
- } else {
- fd = -1;
  }
 
  PADDLE_ENFORCE_EQ(ftruncate(fd, size),
@@ -98,32 +99,38 @@ void AllocateMemoryMap(
  *map_ptr_ = mmap(nullptr, size, PROT_READ | PROT_WRITE, MAP_PRIVATE, fd, 0);
  }
 
+ if (flags & MAPPED_UNLINK) {
+ VLOG(6) << "shm_unlink: " << filename;
+ shm_unlink(filename.c_str());
+ }
+
  PADDLE_ENFORCE_NE(*map_ptr_,
  MAP_FAILED,
  platform::errors::Unavailable(
  "Memory map failed when create shared memory."));
-
  if (flags & MAPPED_KEEPFD) {
- *fd_ = fd;
+ *shared_fd = fd;
+ VLOG(6) << "keep fd: " << *shared_fd;
  } else {
  PADDLE_ENFORCE_NE(::close(fd),
  -1,
  platform::errors::Unavailable(
  "Error closing memory mapped file <", filename, ">"));
 
- *fd_ = -1;
+ *shared_fd = -1;
  }
 }
 
 std::shared_ptr<RefcountedMemoryMapAllocation>
 AllocateRefcountedMemoryMapAllocation(std::string filename,
+ int shared_fd,
  int flags,
  size_t size,
  int buffer_id) {
- int fd = -1;
+ int fd = shared_fd;
  void *base_ptr = nullptr;
  if (buffer_id == -1) {
- AllocateMemoryMap(filename, flags, size + mmap_alignment, &base_ptr, &fd);
+ AllocateMemoryMap(filename, &fd, flags, size + mmap_alignment, &base_ptr);
  VLOG(4) << "Create and mmap a new shm: " << filename;
  } else {
  base_ptr = MemoryMapAllocationPool::Instance().GetById(buffer_id).mmap_ptr_;
@@ -132,7 +139,7 @@ AllocateRefcountedMemoryMapAllocation(std::string filename,
  void *aligned_base_ptr =
  static_cast<void *>(static_cast<char *>(base_ptr) + mmap_alignment);
  return std::make_shared<RefcountedMemoryMapAllocation>(
- aligned_base_ptr, size, filename, flags, fd, buffer_id);
+ aligned_base_ptr, size, filename, fd, flags, buffer_id);
 }
 
 RefcountedMemoryMapAllocation::RefcountedMemoryMapAllocation(
@@ -145,11 +152,23 @@ RefcountedMemoryMapAllocation::RefcountedMemoryMapAllocation(
  : MemoryMapAllocation(ptr, size, ipc_name, fd, flags) {
  // must reset base ptr first.
  buffer_id_ = buffer_id;
+ fd_ = fd;
+ flags_ = flags;
  resetBaseptr();
  initializeRefercount();
 }
 
 void MemoryMapAllocation::close() {
+ if (!closed_fd_) {
+ closed_fd_ = true;
+ if (flags_ & MAPPED_KEEPFD) {
+ VLOG(6) << "one close fd: " << fd_;
+ PADDLE_ENFORCE_NE(::close(fd_),
+ -1,
+ platform::errors::Unavailable(
+ "Error closing file descriptor <", fd_, ">"));
+ }
+ }
  if (closed_) {
  return;
  }
@@ -193,6 +212,15 @@ void RefcountedMemoryMapAllocation::close() {
  void *data = map_ptr_;
  CountInfo *info = reinterpret_cast<CountInfo *>(data);
  --info->refcount;
+ if (flags_ & MAPPED_KEEPFD) {
+ closed_fd_ = true;
+ PADDLE_ENFORCE_NE(::close(fd_),
+ -1,
+ platform::errors::Unavailable(
+ "Error closing file descriptor <", fd_, ">"));
+ VLOG(6) << "close fd: " << fd_;
+ }
+
  if (FLAGS_use_shm_cache && buffer_id_ != -1) {
  return;
  } else {
@@ -260,6 +288,7 @@ std::shared_ptr<MemoryMapWriterAllocation> AllocateMemoryMapWriterAllocation(
  const std::string &ipc_name = GetIPCName();
  int flags = O_RDWR | O_CREAT;
  int fd = shm_open(ipc_name.c_str(), flags, 0600);
+
  PADDLE_ENFORCE_NE(fd,
  -1,
  platform::errors::Unavailable(
@@ -283,7 +312,6 @@ std::shared_ptr<MemoryMapReaderAllocation> RebuildMemoryMapReaderAllocation(
  const std::string &ipc_name, size_t size) {
  int flags = O_RDWR | O_CREAT;
  flags &= ~O_CREAT;
-
  int fd = shm_open(ipc_name.c_str(), flags, 0600);
  PADDLE_ENFORCE_NE(fd,
  -1,

paddle/fluid/memory/allocation/mmap_allocator.h

@@ -44,13 +44,17 @@ enum MappedModes {
 
 class MemoryMapAllocation : public Allocation {
  public:
- explicit MemoryMapAllocation(void *ptr, size_t size, std::string ipc_name)
+ explicit MemoryMapAllocation(void *ptr,
+ size_t size,
+ std::string ipc_name,
+ int fd)
  : Allocation(ptr, size, platform::CPUPlace()),
  ipc_name_(std::move(ipc_name)),
+ fd_(fd),
  map_ptr_(ptr),
  map_size_(size) {}
  explicit MemoryMapAllocation(
- void *ptr, size_t size, std::string ipc_name, int flags, int fd)
+ void *ptr, size_t size, std::string ipc_name, int fd, int flags)
  : Allocation(ptr, size, platform::CPUPlace()),
  ipc_name_(std::move(ipc_name)),
  fd_(fd),
@@ -59,6 +63,7 @@ class MemoryMapAllocation : public Allocation {
  map_size_(size) {}
 
  inline const std::string &ipc_name() const { return ipc_name_; }
+ inline const int shared_fd() const { return fd_; }
 
  virtual void close();
 
@@ -71,6 +76,7 @@ class MemoryMapAllocation : public Allocation {
  void *map_ptr_ = nullptr;
  size_t map_size_ = 0;
  bool closed_ = false;
+ bool closed_fd_ = false;
 };
 
 class RefcountedMemoryMapAllocation : public MemoryMapAllocation {
@@ -93,11 +99,15 @@ class RefcountedMemoryMapAllocation : public MemoryMapAllocation {
  void resetBaseptr();
 };
 
-void AllocateMemoryMap(
- std::string filename, int flags, size_t size, void **base_ptr_, int *fd_);
+void AllocateMemoryMap(std::string filename,
+ int *shared_fd,
+ int flags,
+ size_t size,
+ void **base_ptr_);
 
 std::shared_ptr<RefcountedMemoryMapAllocation>
 AllocateRefcountedMemoryMapAllocation(std::string filename,
+ int shared_fd,
  int flags,
  size_t size,
  int buffer_id = -1);
@@ -111,11 +121,13 @@ class MemoryMapWriterAllocation : public Allocation {
  ipc_name_(std::move(ipc_name)) {}
 
  inline const std::string &ipc_name() const { return ipc_name_; }
+ inline const int shared_fd() const { return fd_; }
 
  ~MemoryMapWriterAllocation() override;
 
  private:
  std::string ipc_name_;
+ int fd_ = -1;
 };
 
 class MemoryMapReaderAllocation : public Allocation {
@@ -127,11 +139,13 @@ class MemoryMapReaderAllocation : public Allocation {
  ipc_name_(std::move(ipc_name)) {}
 
  inline const std::string &ipc_name() const { return ipc_name_; }
+ inline const int shared_fd() const { return fd_; }
 
  ~MemoryMapReaderAllocation() override;
 
  private:
  std::string ipc_name_;
+ int fd_ = -1;
 };
 
 std::shared_ptr<MemoryMapWriterAllocation> AllocateMemoryMapWriterAllocation(

paddle/fluid/pybind/tensor.cc

@@ -859,7 +859,7 @@ void BindTensor(pybind11::module &m) { // NOLINT
  )DOC")
 #endif
  .def("_share_filename",
- [](phi::DenseTensor &self) {
+ [](phi::DenseTensor &self, bool use_file_descripor) {
  if (!self.IsInitialized() || self.numel() == 0)
  throw std::runtime_error(
  "Tensor not initialized or numel is 0. could not pass to "
@@ -886,6 +886,10 @@ void BindTensor(pybind11::module &m) { // NOLINT
 
  int flags = memory::allocation::MAPPED_SHAREDMEM |
  memory::allocation::MAPPED_EXCLUSIVE;
+ if (use_file_descripor) {
+ flags = flags | memory::allocation::MAPPED_KEEPFD |
+ memory::allocation::MAPPED_UNLINK;
+ }
  std::string handle = memory::allocation::GetIPCName();
  int find_id = -1;
  if (FLAGS_use_shm_cache) {
@@ -894,9 +898,10 @@ void BindTensor(pybind11::module &m) { // NOLINT
  if (find_id != -1) {
  handle = memory::allocation::MemoryMapAllocationPool::Instance().GetById(find_id).file_name_; // NOLINT
  }
+ int shared_fd = -1;
  auto shared_holder =
  memory::allocation::AllocateRefcountedMemoryMapAllocation(
- handle, flags, data_size, find_id);
+ handle, shared_fd, flags, data_size, find_id);
 
  // copy data & reset holder
  if (platform::is_cuda_pinned_place(holder->place())) {
@@ -914,8 +919,10 @@ void BindTensor(pybind11::module &m) { // NOLINT
  int type_idx = static_cast<int>(self.type());
 
  return py::make_tuple(mmap_allocation->ipc_name(),
+ mmap_allocation->shared_fd(),
  mmap_allocation->size(), type_idx,
- common::vectorize(self.dims()), self.lod());
+ common::vectorize(self.dims()), self.lod(),
+ use_file_descripor);
  },
  R"DOC(
  Serialize CPU lod tensor in shared memory to tuple.
@@ -935,30 +942,37 @@ void BindTensor(pybind11::module &m) { // NOLINT
  )DOC")
  .def("_new_shared_filename",
  [](py::tuple t) { // __setstate__
- if (t.size() != 5)
+ if (t.size() != 7)
  throw std::runtime_error("Invalid Tensor meta info state!");
 
  phi::DenseTensor tensor;
 
  // 2. Rebuild Allocation
  const std::string &ipc_name = t[0].cast<std::string>();
- size_t size = t[1].cast<size_t>();
+ const int shared_fd = t[1].cast<int>();
+ const bool use_file_descripor = t[6].cast<bool>();
+
+ size_t size = t[2].cast<size_t>();
  int flags = memory::allocation::MAPPED_SHAREDMEM |
  memory::allocation::MAPPED_NOCREATE;
+ if (use_file_descripor) {
+ flags = flags | memory::allocation::MAPPED_KEEPFD |
+ memory::allocation::MAPPED_UNLINK;
+ }
  int find_id = -1;
  if (FLAGS_use_shm_cache) {
  find_id = memory::allocation::MemoryMapAllocationPool::Instance().FindFromCache(flags, size, ipc_name, /*check_refcount*/ false); // NOLINT
  }
  auto shared_holder =
  memory::allocation::AllocateRefcountedMemoryMapAllocation(
- ipc_name, flags, size, find_id);
+ ipc_name, shared_fd, flags, size, find_id);
 
  // 3. Rebuild Tensor
  tensor.ResetHolderWithType(
  shared_holder,
- static_cast<phi::DataType>(t[2].cast<int>()));
- tensor.Resize(common::make_ddim(t[3].cast<std::vector<int>>()));
- tensor.set_lod(t[4].cast<framework::LoD>());
+ static_cast<phi::DataType>(t[3].cast<int>()));
+ tensor.Resize(common::make_ddim(t[4].cast<std::vector<int>>()));
+ tensor.set_lod(t[5].cast<framework::LoD>());
 
  return tensor;
  },

python/paddle/base/reader.py

@@ -14,6 +14,7 @@
 
 import logging
 import multiprocessing
+import os
 import queue
 import sys
 import threading
@@ -409,6 +410,7 @@ def from_generator(
  ...
  >>> # doctest: -SKIP
  """
+ os.environ["FLAGS_use_file_descripor"] = "0"
  if in_dygraph_mode():
  return DygraphGeneratorLoader(
  feed_list,