data_store_conduit.hpp

Go to the documentation of this file.

// Copyright (c) 2014-2023, Lawrence Livermore National Security, LLC.
// Produced at the Lawrence Livermore National Laboratory.
// Written by the LBANN Research Team (B. Van Essen, et al.) listed in
// the CONTRIBUTORS file. <lbann-dev@llnl.gov>
//
// LLNL-CODE-697807.
// All rights reserved.
//
// This file is part of LBANN: Livermore Big Artificial Neural Network
// Toolkit. For details, see http://software.llnl.gov/LBANN or
// https://github.com/LLNL/LBANN.
//
// Licensed under the Apache License, Version 2.0 (the "Licensee"); you
// may not use this file except in compliance with the License.  You may
// obtain a copy of the License at:
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
// implied. See the License for the specific language governing
// permissions and limitations under the license.
//

#ifndef __DATA_STORE_CONDUIT_HPP__
#define __DATA_STORE_CONDUIT_HPP__

#include "lbann_config.hpp"

#include "conduit/conduit_node.hpp"
#include "lbann/base.hpp"
#include "lbann/comm.hpp"
#include "lbann/utils/exception.hpp"
#include <mutex>
#include <unordered_map>
#include <unordered_set>

namespace lbann {

// support for encoding data_id in conduit::Node, used by
// conduit_data_store and associated code
#define LBANN_SAMPLE_ID_PAD 9
#define LBANN_DATA_ID_STR(data_id)                                             \
  lbann::pad(std::to_string(data_id), LBANN_SAMPLE_ID_PAD, '0')

class generic_data_reader;

struct size_t_pair_hash
{
  template <class T1, class T2>
  std::size_t operator()(const std::pair<T1, T2>& pair) const
  {
    return std::hash<T1>()(pair.first) ^ std::hash<T2>()(pair.second);
  }
};

class data_store_conduit
{

public:
  // need to quickly change from unordered_map to map for debugging
  using map_ii_t = std::unordered_map<int, int>;
  using map_is_t = std::unordered_map<int, size_t>;

  // Hash map for tracking the node and hyperslab partition ID
  using map_pssi_t =
    std::unordered_map<std::pair<size_t, size_t>, int, size_t_pair_hash>;

  // not currently used; will be in the future
  using map_ss_t = std::unordered_map<size_t, size_t>;

  data_store_conduit(generic_data_reader* reader);

  data_store_conduit(const data_store_conduit&);

  data_store_conduit(const data_store_conduit&, const std::vector<int>&);

  data_store_conduit& operator=(const data_store_conduit&);

  data_store_conduit* copy() const { return new data_store_conduit(*this); }

  ~data_store_conduit();

  void set_data_reader_ptr(generic_data_reader* reader);

  void set_shuffled_indices(const std::vector<int>* indices);

  size_t get_num_global_indices() const;

  void setup(int mini_batch_size);

  // TODO FIXME
  void check_mem_capacity(lbann_comm* comm,
                          const std::string sample_list_file,
                          size_t stride,
                          size_t offset);

  const conduit::Node& get_conduit_node(int data_id) const;

  void set_conduit_node(int data_id,
                        const conduit::Node& node,
                        bool already_have = false);

  void set_preloaded_conduit_node(int data_id, const conduit::Node& node);

  void spill_preloaded_conduit_node(int data_id, const conduit::Node& node);

  const conduit::Node& get_random_node() const;

  const conduit::Node& get_random_node(const std::string& field) const;

  conduit::Node& get_empty_node(int data_id);

  //=================================================================
  // methods for setting and querying the data store's mode
  //=================================================================
  bool is_preloading() const { return m_preloading; }

  bool is_explicitly_loading() const { return m_explicitly_loading; }

  bool is_fully_loaded() const;

  bool is_local_cache() const { return m_is_local_cache; }

  void set_is_preloading(bool flag);

  void set_is_explicitly_loading(bool flag);

  void set_loading_is_complete();

  void set_is_local_cache(bool flag = true) { m_is_local_cache = flag; }

  void check_query_flags() const;

  //=================================================================
  // END methods for setting and querying the data store's mode
  //=================================================================

  // XX   void { m_owner_maps_were_exchanged = false; }
  void exchange_owner_maps();

  void build_preloaded_owner_map(const std::vector<int>& per_rank_list_sizes);

  void set_preloaded_owner_map(const std::unordered_map<int, int>& owner)
  {
    for (auto&& i : owner) {
      m_owner[std::make_pair(i.first, m_offset_in_partition)] = i.second;
    }
  }

  void clear_owner_map();

  void set_owner_map(const std::unordered_map<int, int>& m)
  {
    for (auto&& i : m) {
      m_owner[std::make_pair(i.first, m_offset_in_partition)] = i.second;
    }
  }

  void add_owner(int data_id, int owner)
  {
    m_owner[std::make_pair(data_id, m_offset_in_partition)] = owner;
  }

  void set_finished_building_map() { m_owner_maps_were_exchanged = true; }

  void compact_nodes();

  int get_index_owner(int idx);

  void preload_local_cache();

  void start_exchange_mini_batch_data(size_t current_pos, size_t mb_size);
  void finish_exchange_mini_batch_data();

  void set_node_sizes_vary() { m_node_sizes_vary = true; }

  bool has_conduit_node(int data_id) const;

  std::ofstream* m_debug = nullptr;
  std::ofstream* m_profile = nullptr;

  int get_data_size()
  {
    std::lock_guard<std::mutex> lock(m_mutex);
    return m_data.size();
  }

  void copy_members(const data_store_conduit& rhs);

  void flush_debug_file();

  void flush_profile_file() const;

  void write_checkpoint(std::string dir_name);

  void load_checkpoint(std::string dir_name,
                       generic_data_reader* reader = nullptr);

  void set_profile_msg(std::string);

  bool test_local_cache_imagenet(int n);

  void test_imagenet_node(int sample_id, bool dereference = true);

  size_t get_mem_usage();

private:
  bool m_bcast_sample_size = true;

  // if not null, 'm_other' points from a train to a validation
  // data store; this permits communication which is needed in
  // special cases (e.g, see: data_reader_npz_ras_lipid.cpp)
  data_store_conduit* m_other = nullptr;

  bool m_owner_maps_were_exchanged = false;

  bool m_run_checkpoint_test = false;

  size_t m_my_num_indices = 0;

  bool m_spill = false;

  bool m_is_spilled = false;

  std::ofstream m_metadata;

  std::string m_spill_dir_base;

  int m_cur_spill_dir_integer = -1;

  std::string m_cur_spill_dir;

  std::string m_test_dir;

  int m_num_files_in_cur_spill_dir;

  map_ii_t m_spilled_nodes;

  // Guards m_sample_sizes, m_data, m_image_offsets, m_sample_sizes
  mutable std::mutex m_mutex;
  // Guards m_compact_sample_size
  std::mutex m_mutex_2;

  char* m_mem_seg = 0;
  size_t m_mem_seg_length = 0;
  std::string m_seg_name;

  const std::string m_debug_filename_base = "debug";
  std::string m_debug_filename;

  const std::string m_profile_filename_base = "data_store_profile";
  std::string m_profile_filename;

  bool m_was_loaded_from_file = false;
  const std::string m_cereal_fn = "data_store_cereal";

  const int m_max_files_per_directory = 500;

  //===========================================================
  // timers for profiling exchange_data
  //===========================================================

  // applicable to imagenet; NA for JAG
  double m_exchange_sample_sizes_time = 0;

  // time from beginning of exchange_data_by_sample to wait_all
  double m_start_snd_rcv_time = 0;

  // time for wait_all
  double m_wait_all_time = 0;

  // time to unpack nodes received from other ranks
  double m_rebuild_time = 0;

  // total time for exchange_mini_batch_data
  double m_exchange_time = 0;

  // sanity check:
  //   m_start_snd_rcv_time + m_wait_all_time + m_rebuild_time
  // should be only slightly less than m_exchange_time;
  // Note that, for imagenet, the first call to exchange_data_by_sample
  // involves additional communication for exchanging sample sizes

  //===========================================================
  // END: timers for profiling exchange_data
  //===========================================================

  bool m_is_setup = false;

  bool m_loading_is_complete = false;

  bool m_preloading = false;

  bool m_explicitly_loading = false;

  int m_owner_map_mb_size = 0;

  int m_compacted_sample_size = 0;

  bool m_is_local_cache = false;

  bool m_node_sizes_vary = false;

  bool m_have_sample_sizes = false;

  generic_data_reader* m_reader;

  lbann_comm* m_comm = nullptr;

  bool m_world_master;
  bool m_trainer_master;
  int m_rank_in_trainer;
  int m_rank_in_world = -1; // -1 for debugging
  int m_partition_in_trainer;
  int m_offset_in_partition;

  int m_np_in_trainer;
  int m_num_partitions_in_trainer;

  bool m_mini_batch_data_exchange_started = false;

  mutable map_pssi_t m_owner;

  const std::vector<int>* m_shuffled_indices;

  mutable std::unordered_map<int, conduit::Node> m_data;

  std::unordered_map<int, conduit::Node> m_data_cache;

  std::vector<int> m_recv_data_ids;
  map_ii_t m_recv_sample_sizes;

  std::unordered_map<int, conduit::Node> m_minibatch_data;

  std::vector<conduit::Node> m_send_buffer;
  std::vector<conduit::Node> m_send_buffer_2;
  std::vector<El::mpi::Request<El::byte>> m_send_requests;
  std::vector<El::mpi::Request<El::byte>> m_recv_requests;
  std::vector<conduit::Node> m_recv_buffer;
  std::vector<size_t> m_outgoing_msg_sizes;
  std::vector<size_t> m_incoming_msg_sizes;

  map_is_t m_sample_sizes;

  map_is_t m_image_offsets;

  std::vector<std::unordered_set<int>> m_indices_to_send;

  std::vector<std::unordered_set<int>> m_indices_to_recv;

  //=========================================================================
  // methods follow
  //=========================================================================

  void start_exchange_data_by_sample(size_t current_pos, size_t mb_size);
  void finish_exchange_data_by_sample();

  void setup_data_store_buffers();

  void build_node_for_sending(const conduit::Node& node_in,
                              conduit::Node& node_out);

  void exchange_sample_sizes();

  int build_indices_i_will_send(int current_pos, int mb_size);

  int build_indices_i_will_recv(int current_pos, int mb_size);

  void error_check_compacted_node(const conduit::Node& nd, int data_id);

  void exchange_local_caches();

  void get_image_sizes(map_is_t& sizes, std::vector<std::vector<int>>& indices);

  void allocate_shared_segment(map_is_t& sizes,
                               std::vector<std::vector<int>>& indices);

  void read_files(std::vector<char>& work,
                  map_is_t& sizes,
                  std::vector<int>& indices);

  void compute_image_offsets(map_is_t& image_sizes,
                             std::vector<std::vector<int>>& indices);

  void exchange_images(std::vector<char>& work,
                       map_is_t& image_sizes,
                       std::vector<std::vector<int>>& indices);

  void build_conduit_nodes(map_is_t& sizes);

  void fillin_shared_images(char* images, size_t size, size_t offset);

  void test_checkpoint(const std::string&);

  void print_variables();

  void print_partial_owner_map(int n);

  std::string get_conduit_dir() const;
  std::string get_cereal_fn() const;
  std::string get_metadata_fn() const;

  void make_dir_if_it_doesnt_exist(const std::string& dir);

  void spill_conduit_node(const conduit::Node& node, int data_id);

  void load_spilled_conduit_nodes();

  void setup_spill(std::string dir);

  void save_state();

  void open_informational_files();

  void open_next_conduit_spill_directory();

  void profile_timing();

  void setup_checkpoint_test();

  std::string get_lassen_spill_dir();

  void verify_sample_size();

  //=========================================================================
  // functions and templates for optional profiling and debug files follow
  //=========================================================================

  void PROFILE() const
  {
    if (!m_profile) {
      return;
    }
    (*m_profile) << std::endl;
    flush_profile_file();
  }

  template <typename T, typename... Types>
  void PROFILE(T var1, Types... var2) const
  {
    if (!m_world_master) {
      return;
    }
    if (!m_profile) {
      return;
    }
    (*m_profile) << var1 << " ";
    PROFILE(var2...);
    flush_profile_file();
  }

  void DEBUG_DS()
  {
    if (!m_debug) {
      return;
    }
    (*m_debug) << std::endl;
    flush_debug_file();
  }

  template <typename T, typename... Types>
  void DEBUG_DS(T var1, Types... var2)
  {
    if (!m_debug) {
      return;
    }
    (*m_debug) << var1 << " ";
    DEBUG_DS(var2...);
    flush_debug_file();
  }
};

} // namespace lbann

#endif // __DATA_STORE_JAG_HPP__