ceph/src/boost/libs/mpi/src/python/py_environment.cpp

   1 // (C) Copyright 2006 Douglas Gregor <doug.gregor -at- gmail.com>
   2
   3 // Use, modification and distribution is subject to the Boost Software
   4 // License, Version 1.0. (See accompanying file LICENSE_1_0.txt or copy at
   5 // http://www.boost.org/LICENSE_1_0.txt)
   6
   7 //  Authors: Douglas Gregor
   8
   9 /** @file environment.cpp
  10  *
  11  *  This file reflects the Boost.MPI "environment" class into Python
  12  *  methods at module level.
  13  */
  14
  15 #include <locale>
  16 #include <string>
  17 #include <boost/python.hpp>
  18 #include <boost/mpi.hpp>
  19
  20 using namespace boost::python;
  21 using namespace boost::mpi;
  22
  23 namespace boost { namespace mpi { namespace python {
  24
  25 extern const char* environment_init_docstring;
  26 extern const char* environment_finalize_docstring;
  27 extern const char* environment_abort_docstring;
  28 extern const char* environment_initialized_docstring;
  29 extern const char* environment_finalized_docstring;
  30
  31 /**
  32  * The environment used by the Boost.MPI Python module. This will be
  33  * zero-initialized before it is used.
  34  */
  35 static environment* env;
  36
  37 bool mpi_init(list python_argv, bool abort_on_exception)
  38 {
  39   // If MPI is already initialized, do nothing.
  40   if (environment::initialized())
  41     return false;
  42
  43 #if PY_MAJOR_VERSION >= 3
  44   #ifdef BOOST_MPI_HAS_NOARG_INITIALIZATION
  45     env = new environment(abort_on_exception);
  46   #else
  47     #error No argument initialization, supported from MPI 1.2 and up, is needed when using Boost.MPI with Python 3.x
  48   #endif
  49 #else
  50
  51   // Convert Python argv into C-style argc/argv.
  52   int my_argc = extract<int>(python_argv.attr("__len__")());
  53   char** my_argv = new char*[my_argc];
  54   for (int arg = 0; arg < my_argc; ++arg)
  55     my_argv[arg] = strdup(extract<const char*>(python_argv[arg]));
  56
  57   // Initialize MPI
  58   int mpi_argc = my_argc;
  59   char** mpi_argv = my_argv;
  60   env = new environment(mpi_argc, mpi_argv, abort_on_exception);
  61
  62   // If anything changed, convert C-style argc/argv into Python argv
  63   if (mpi_argv != my_argv)
  64     PySys_SetArgv(mpi_argc, mpi_argv);
  65
  66   for (int arg = 0; arg < mpi_argc; ++arg)
  67     free(mpi_argv[arg]);
  68   delete [] mpi_argv;
  69 #endif
  70
  71   return true;
  72 }
  73
  74 void mpi_finalize()
  75 {
  76   if (env) {
  77     delete env;
  78     env = 0;
  79   }
  80 }
  81
  82 void export_environment()
  83 {
  84   using boost::python::arg;
  85
  86   def("init", mpi_init, (arg("argv"), arg("abort_on_exception") = true),
  87       environment_init_docstring);
  88   def("finalize", mpi_finalize, environment_finalize_docstring);
  89
  90   // Setup initialization and finalization code
  91   if (!environment::initialized()) {
  92     // MPI_Init from sys.argv
  93     object sys = object(handle<>(PyImport_ImportModule("sys")));
  94     mpi_init(extract<list>(sys.attr("argv")), true);
  95
  96     // Setup MPI_Finalize call when the program exits
  97     object atexit = object(handle<>(PyImport_ImportModule("atexit")));
  98     object finalize = scope().attr("finalize");
  99     atexit.attr("register")(finalize);
 100   }
 101
 102   def("abort", &environment::abort, arg("errcode"),
 103       environment_abort_docstring);
 104   def("initialized", &environment::initialized,
 105       environment_initialized_docstring);
 106   def("finalized", &environment::finalized,
 107       environment_finalized_docstring);
 108   scope().attr("max_tag") = environment::max_tag();
 109   scope().attr("collectives_tag") = environment::collectives_tag();
 110   scope().attr("processor_name") = environment::processor_name();
 111
 112   if (optional<int> host_rank = environment::host_rank())
 113     scope().attr("host_rank") = *host_rank;
 114   else
 115     scope().attr("host_rank") = object();
 116
 117   if (optional<int> io_rank = environment::io_rank())
 118     scope().attr("io_rank") = *io_rank;
 119   else
 120     scope().attr("io_rank") = object();
 121 }
 122
 123 } } } // end namespace boost::mpi::python