Further improve GSD performance.

hoomd/GSDDumpWriter.cc

@@ -54,8 +54,7 @@ GSDDumpWriter::GSDDumpWriter(std::shared_ptr<SystemDefinition> sysdef,
                              std::shared_ptr<ParticleGroup> group,
                              std::string mode,
                              bool truncate)
-    : Analyzer(sysdef, trigger), m_fname(fname), m_mode(mode), m_truncate(truncate),
-      m_is_initialized(false), m_group(group)
+    : Analyzer(sysdef, trigger), m_fname(fname), m_mode(mode), m_truncate(truncate), m_group(group)
     {
     m_exec_conf->msg->notice(5) << "Constructing GSDDumpWriter: " << m_fname << " " << mode << " "
                                 << truncate << endl;
@@ -75,6 +74,8 @@ GSDDumpWriter::GSDDumpWriter(std::shared_ptr<SystemDefinition> sysdef,
     m_particle_dynamic.reset();
     m_particle_dynamic[gsd_flag::position] = true;
     m_particle_dynamic[gsd_flag::orientation] = true;
+
+    initFileIO();
     }
 
 pybind11::tuple GSDDumpWriter::getDynamic()
@@ -201,6 +202,40 @@ void GSDDumpWriter::setDynamic(pybind11::object dynamic)
         }
     }
 
+void GSDDumpWriter::flush()
+    {
+    if (m_exec_conf->isRoot())
+        {
+        int retval = gsd_flush(&m_handle);
+        GSDUtils::checkError(retval, m_fname);
+        }
+    }
+
+void GSDDumpWriter::setMaximumWriteBufferSize(uint64_t size)
+    {
+    if (m_exec_conf->isRoot())
+        {
+        int retval = gsd_set_maximum_write_buffer_size(&m_handle, size);
+        GSDUtils::checkError(retval, m_fname);
+
+        // Scale the index buffer entires to write with the write buffer.
+        retval = gsd_set_index_entries_to_buffer(&m_handle, size / 256);
+        GSDUtils::checkError(retval, m_fname);
+        }
+    }
+
+uint64_t GSDDumpWriter::getMaximumWriteBufferSize()
+    {
+    if (m_exec_conf->isRoot())
+        {
+        return gsd_get_maximum_write_buffer_size(&m_handle);
+        }
+    else
+        {
+        return 0;
+        }
+    }
+
 //! Initializes the output file for writing
 void GSDDumpWriter::initFileIO()
     {
@@ -269,20 +304,13 @@ void GSDDumpWriter::initFileIO()
         bcast(m_nondefault, 0, m_exec_conf->getMPICommunicator());
         }
 #endif
-
-    m_is_initialized = true;
     }
 
 GSDDumpWriter::~GSDDumpWriter()
     {
     m_exec_conf->msg->notice(5) << "Destroying GSDDumpWriter" << endl;
 
-    bool root = true;
-#ifdef ENABLE_MPI
-    root = m_exec_conf->isRoot();
-#endif
-
-    if (root && m_is_initialized)
+    if (m_exec_conf->isRoot())
         {
         m_exec_conf->msg->notice(5) << "GSD: close gsd file " << m_fname << endl;
         gsd_close(&m_handle);
@@ -300,10 +328,6 @@ void GSDDumpWriter::analyze(uint64_t timestep)
     Analyzer::analyze(timestep);
     int retval;
 
-    // open the file if it is not yet opened
-    if (!m_is_initialized)
-        initFileIO();
-
     // truncate the file if requested
     if (m_truncate)
         {
@@ -1502,7 +1526,11 @@ void export_GSDDumpWriter(pybind11::module& m)
                                { return gsd->getGroup()->getFilter(); })
         .def_property("write_diameter",
                       &GSDDumpWriter::getWriteDiameter,
-                      &GSDDumpWriter::setWriteDiameter);
+                      &GSDDumpWriter::setWriteDiameter)
+        .def("flush", &GSDDumpWriter::flush)
+        .def_property("maximum_write_buffer_size",
+                      &GSDDumpWriter::getMaximumWriteBufferSize,
+                      &GSDDumpWriter::setMaximumWriteBufferSize);
     }
 
     } // end namespace detail

hoomd/GSDDumpWriter.h

@@ -124,6 +124,15 @@ class PYBIND11_EXPORT GSDDumpWriter : public Analyzer
         m_write_diameter = write_diameter;
         }
 
+    /// Flush the write buffer
+    void flush();
+
+    /// Set the maximum write buffer size (in bytes)
+    void setMaximumWriteBufferSize(uint64_t size);
+
+    /// Get the maximum write buffer size (in bytes)
+    uint64_t getMaximumWriteBufferSize();
+
     protected:
     /// Flags for dynamic/default bitsets.
     struct gsd_flag
@@ -196,7 +205,6 @@ class PYBIND11_EXPORT GSDDumpWriter : public Analyzer
     std::string m_fname;           //!< The file name we are writing to
     std::string m_mode;            //!< The file open mode
     bool m_truncate = false;       //!< True if we should truncate the file on every analyze()
-    bool m_is_initialized = false; //!< True if the file is open
     bool m_write_topology = false; //!< True if topology should be written
     bool m_write_diameter = false; //!< True if the diameter attribute should be written
     gsd_handle m_handle;           //!< Handle to the file

hoomd/__init__.py

@@ -29,10 +29,18 @@
 
 See Also:
     Tutorial: :doc:`tutorial/00-Introducing-HOOMD-blue/00-index`
+
+.. rubric:: Signal handling
+
+On import, `hoomd` installs a ``SIGTERM`` signal handler that calls `sys.exit`
+so that open gsd files have a chance to flush write buffers
+(`hoomd.write.GSD.flush`) when a user's process is terminated. Use
+`signal.signal` to adjust this behavior as needed.
 """
 import sys
 import pathlib
 import os
+import signal
 
 if ((pathlib.Path(__file__).parent / 'CMakeLists.txt').exists()
         and 'SPHINX' not in os.environ):
@@ -92,3 +100,7 @@ def _hoomd_sys_excepthook(type, value, traceback):
 
 
 sys.excepthook = _hoomd_sys_excepthook
+
+# Install a SIGTERM handler that gracefully exits, allowing open files to flush
+# buffered writes and close.
+signal.signal(signal.SIGTERM, lambda n, f: sys.exit(1))