Method setup

src/data_readers/data_reader_python.cpp:178–332

}

void python_reader::setup(int num_io_threads,
                          observer_ptr<thread_pool> io_thread_pool)
{
  generic_data_reader::setup(num_io_threads, io_thread_pool);

  // Acquire Python GIL
  python::global_interpreter_lock gil;

  // Import modules
  python::object main_module = PyImport_ImportModule("__main__");
  python::object ctypes_module = PyImport_ImportModule("ctypes");
  python::object multiprocessing_module =
    PyImport_ImportModule("multiprocessing");

  // Stop process pool if needed
  if (m_process_pool != nullptr) {
    PyObject_CallMethod(m_process_pool, "terminate", nullptr);
    m_process_pool = nullptr;
  }

  // Allocate shared memory array
  /// @todo Figure out more robust way to get max mini-batch size
  const El::Int sample_size = get_linearized_data_size();
  const El::Int mini_batch_size = get_trainer().get_max_mini_batch_size();
  std::string datatype_typecode;
  switch (sizeof(DataType)) {
  case 4:
    datatype_typecode = "f";
    break;
  case 8:
    datatype_typecode = "d";
    break;
  default:
    LBANN_ERROR("invalid data type for Python data reader "
                "(only float and double are supported)");
  }
  m_shared_memory_array = PyObject_CallMethod(multiprocessing_module,
                                              "RawArray",
                                              "(s, l)",
                                              datatype_typecode.c_str(),
                                              sample_size * mini_batch_size);

  // Get address of shared memory buffer
  python::object shared_memory_ptr =
    PyObject_CallMethod(ctypes_module,
                        "addressof",
                        "(O)",
                        m_shared_memory_array.get());
  m_shared_memory_array_ptr =
    reinterpret_cast<DataType*>(PyLong_AsLong(shared_memory_ptr));

  // Create global variables in Python
  // Note: The static counter makes sure variable names are unique.
  static El::Int instance_id = 0;
  instance_id++;
  const std::string sample_func_name =
    ("_DATA_READER_PYTHON_CPP_sample_function_wrapper" +
     std::to_string(instance_id));