checkpoint_control.hpp

// FEAT3: Finite Element Analysis Toolbox, Version 3
// Copyright (C) 2010 by Stefan Turek & the FEAT group
// FEAT3 is released under the GNU General Public License version 3,
// see the file 'copyright.txt' in the top level directory for details.
 
#pragma once
 
#include <kernel/base_header.hpp>
#include <kernel/util/dist.hpp>
#include <kernel/util/dist_file_io.hpp>
#include <kernel/util/string.hpp>
#include <kernel/util/binary_stream.hpp>
#include <kernel/util/pack.hpp>
#include <kernel/lafem/container.hpp>
 
#include <cstdio>
#include <cstring>
#include <map>
#include <sys/stat.h>
#include <memory>
 
namespace FEAT
{
  namespace Control
  {
    class Checkpointable
    {
    public:
      virtual std::uint64_t get_checkpoint_size(LAFEM::SerialConfig & config) = 0;
 
      virtual void restore_from_checkpoint_data(std::vector<char> & data) = 0;
 
      virtual std::uint64_t set_checkpoint_data(std::vector<char> & data, LAFEM::SerialConfig & config) = 0;
 
    protected:
      virtual ~Checkpointable() {}
    }; // class Checkpointable
 
    template <typename Object_>
    class CheckpointableWrapper : public Checkpointable
    {
    private:
      Object_ & _object;
 
    public:
      explicit CheckpointableWrapper(Object_ & obj) :
        _object(obj) {}
 
      virtual std::uint64_t get_checkpoint_size(LAFEM::SerialConfig & config) override
      {
        return _object.get_checkpoint_size(config);
      }
 
      virtual void restore_from_checkpoint_data(std::vector<char> & data) override
      {
        _object.restore_from_checkpoint_data(data);
      }
 
      virtual std::uint64_t set_checkpoint_data(std::vector<char> & data, LAFEM::SerialConfig & config) override
      {
        return _object.set_checkpoint_data(data, config);
      }
    }; // class CheckpointableWrapper
 
    class CheckpointControl
    {
    private:
      std::map<String, std::shared_ptr<Checkpointable>> _checkpointable_by_identifier;
      std::vector<char> _input_array;
      std::map<String, std::uint64_t> _offset_by_identifier;
      const Dist::Comm & _comm;
      FEAT::LAFEM::SerialConfig _config;
 
      std::uint64_t _collect_checkpoint_data(std::vector<char> & buffer)
      {
        std::uint64_t checkpoint_size(0u);
        std::uint64_t real_size(0u);
        std::map<String, std::tuple<std::uint64_t, std::uint64_t>> sizes;
 
        for (auto const & it : _checkpointable_by_identifier)
        {
          std::uint64_t identifierlength = (std::uint64_t)it.first.length();
          std::uint64_t datalength = it.second->get_checkpoint_size(_config);
 
          checkpoint_size += identifierlength + datalength + sizeof(std::uint64_t) + sizeof(std::uint64_t);
          real_size += identifierlength + sizeof(std::uint64_t) + sizeof(std::uint64_t);
          sizes[it.first] = std::make_tuple(identifierlength, datalength);
        }
 
        buffer.reserve(checkpoint_size);
 
        for (auto const & it : _checkpointable_by_identifier)
        {
          char * cidentifierlength = reinterpret_cast<char *>(&std::get<0>(sizes[it.first]));
          buffer.insert(std::end(buffer), cidentifierlength, cidentifierlength + sizeof(std::uint64_t));
          buffer.insert(std::end(buffer), it.first.begin(), it.first.end());
          std::uint64_t old_size = buffer.size();
          buffer.insert(std::end(buffer), sizeof(std::uint64_t), 0); //set datalength to an arbitrary value
          std::uint64_t ireal_size = it.second->set_checkpoint_data(buffer, _config);
          char * csize = reinterpret_cast<char *>(&ireal_size);
          for(std::uint64_t i(0) ; i < sizeof(std::uint64_t) ; ++i)  //overwrite the the datalength
          {
            buffer[old_size + i] = csize[i];
          }
          real_size += ireal_size;
        }
 
        buffer.resize(real_size);
        return real_size;
      }
 
      void _restore_checkpoint_data()
      {
        std::uint64_t stringsize;
        std::uint64_t datasize(0);
        std::size_t i = 0;
        std::size_t size = _input_array.size();
 
        // Loop over all objects stored in the checkpoint file
        while (i < size)
        {
          // Get size of identifier string
          ::memcpy(&stringsize, _input_array.data() + i, sizeof(std::uint64_t));
          i += sizeof(std::uint64_t);
 
          // Read the identifier string and put it as key for the calculated offset to the map
          _offset_by_identifier[String(_input_array.data() + i, stringsize)] = i + stringsize;
          i += stringsize;
 
          // Get the size of the Data holding the array, to get the start point of the next checkpointable object in the array
          ::memcpy(&datasize, _input_array.data() + i, sizeof(std::uint64_t));
          i += sizeof(std::uint64_t) + datasize;
        }
      }
 
      void _save(const String name)
      {
        std::vector<char> buffer;
        _collect_checkpoint_data(buffer);
        std::vector<char> common(0);
 
        DistFileIO::write_combined(common, buffer, name + ".cp", _comm);
      }
 
      void _load(const String name)
      {
 
        //std::vector<char> temp(0);
        std::vector<char> common(0);
        // read the checkpoint file
        DistFileIO::read_combined(common, _input_array, name + ".cp", _comm);
 
        _restore_checkpoint_data();
      }
 
    public:
      CheckpointControl(const Dist::Comm & comm, const LAFEM::SerialConfig & config = LAFEM::SerialConfig()) :
        _comm(comm),
        _config(config)
      {
        _input_array = std::vector<char>(0);
      }
 
      ~CheckpointControl()
      {
      }
 
      void set_config(FEAT::LAFEM::SerialConfig& conf)
      {
        _config = conf;
      }
      void clear_input()
      {
        _offset_by_identifier.clear();
        _input_array.resize(0);
      }
 
      String get_identifier_list()
      {
        String identifiers;
        for (auto item : _checkpointable_by_identifier)
        {
          if (identifiers.length() > 0)
          {
            identifiers += "\n";
          }
          identifiers += item.first;
        }
        return identifiers;
      }
 
      template <typename OT_>
      void add_object(String identifier, OT_ & object)
      {
        XASSERTM(_checkpointable_by_identifier.count(identifier) == 0, "a object with the given identifier is already registered; choose a unique identifier");
        auto checkpointable_object = std::make_shared<CheckpointableWrapper<OT_>>(object);
        _checkpointable_by_identifier[identifier] = checkpointable_object;
      }
 
      void remove_object(String identifier)
      {
        XASSERTM(_checkpointable_by_identifier.count(identifier) == 1, "a object with the given identifier is not registered at this checkpoint instance");
        _checkpointable_by_identifier.erase(identifier);
      }
 
      template <typename OT_>
      void restore_object(String identifier, OT_ & object, bool add_to_checkpoint_control = true)
      {
        XASSERTM(_input_array.size() > 0u, "no input file has been loaded before");
        XASSERTM(_offset_by_identifier.count(identifier) == 1, "input file doesn't contain data for the given identifier");
 
        auto checkpointable_object = std::make_shared<CheckpointableWrapper<OT_>>(object);
 
        std::uint64_t size(0);
        char * in_data = _input_array.data() + _offset_by_identifier[identifier];
        ::memcpy(&size, in_data, sizeof(std::uint64_t));
        std::vector<char> buffer(in_data + sizeof(std::uint64_t), in_data + sizeof(std::uint64_t) + size);
        checkpointable_object->restore_from_checkpoint_data(buffer);
 
        if (add_to_checkpoint_control)
        {
          add_object(identifier, object);
        }
      }
 
      void save(const String filename)
      {
        size_t pos = filename.rfind('.');
        String extension = filename.substr(pos + 1);
        String name = filename.substr(0, pos);
 
        XASSERTM(name != "", "no complete filename consisting of name.extension given");
 
        if (extension == "cp")
        {
          _save(name);
        }
        else
        {
          XASSERTM(extension != "zcp", "no zlib support, activate zlib or write uncompressed files with .cp extension instead");
          XASSERTM(extension == "zcp", "no valid checkpoint filename choosen");
        }
      }
 
      //Should we ignore the extension and just give in filename without .cp?
      void load(const String filename)
      {
        XASSERTM(_input_array.size() == 0u, "another input file was read before");
 
        size_t pos = filename.rfind('.');
        String extension = filename.substr(pos + 1);
        String name = filename.substr(0, pos);
 
        XASSERTM(name != "", "no complete filename consisting of name.extension given");
 
        if (extension == "cp")
        {
          _load(name);
        }
        else
        {
          XABORTM("no valid checkpoint file choosen");
        }
      }
 
      void save(BinaryStream & bs)
      {
        std::vector<char> buffer;
        std::uint64_t slen = _collect_checkpoint_data(buffer);
 
        bs.write(reinterpret_cast<char *>(&slen), sizeof(slen));
        bs.write(buffer.data(), static_cast<std::streamsize>(buffer.size()));
      }
 
      void load(BinaryStream & bs)
      {
        XASSERTM(_input_array.size() == 0, "another input file was read before");
 
        char * buffer = bs.data();
        std::uint64_t size = *(std::uint64_t *)(buffer);
 
        _input_array.resize(size);
        std::copy(buffer + sizeof(std::uint64_t), buffer + sizeof(std::uint64_t) + size - 1, _input_array.data());
 
        _restore_checkpoint_data();
      }
 
    }; // class Checkpoint
 
  } // namespace Control
} // namespace FEAT