fixing shutdown

Author: fschlimb

ddptensor/__init__.py

@@ -13,6 +13,7 @@
 # At this point there are no checks of input arguments whatsoever, arguments
 # are simply forwarded as-is.
 
+_bool = bool
 from . import _ddptensor as _cdt
 from ._ddptensor import (
     FLOAT64 as float64,
@@ -26,15 +27,22 @@
     UINT16 as uint16,
     UINT8 as uint8,
     BOOL as bool,
-    init,
+    init as _init,
     fini,
     sync
 )
+
 from .ddptensor import dtensor
 from os import getenv
 from . import array_api as api
 from . import spmd
 
+_ddpt_cw = _bool(int(getenv('DDPT_CW', True)))
+
+def init(cw=None):
+    cw = _ddpt_cw if cw is None else cw
+    _init(cw)
+
 for op in api.api_categories["EWBinOp"]:
     if not op.startswith("__"):
         OP = op.upper()

src/Deferred.cpp

@@ -37,6 +37,11 @@ Deferred::ptr_type Deferred::undefer_next()
     return r;
 }
 
+void Deferred::fini()
+{
+    _deferred.clear();
+}
+
 void process_promises()
 {
     while(true) {
@@ -55,3 +60,4 @@ void sync()
         std::this_thread::sleep_for(std::chrono::milliseconds(1));
     }
 }
+

src/MPIMediator.cpp

@@ -21,8 +21,13 @@ static std::mutex ak_mutex;
 void send_to_workers(const Deferred::ptr_type & dfrd, bool self = false);
 
 MPIMediator::MPIMediator()
-    : _listener(&MPIMediator::listen, this)
+    : _listener(nullptr)
 {
+    MPI_Comm comm = MPI_COMM_WORLD;
+    int sz;
+    MPI_Comm_size(comm, &sz);
+    if(sz > 1)
+        _listener = new std::thread(&MPIMediator::listen, this);
 }
 
 MPIMediator::~MPIMediator()
@@ -36,7 +41,11 @@ MPIMediator::~MPIMediator()
     if(is_cw() && rank == 0) to_workers(nullptr);
     MPI_Barrier(comm);
     if(!is_cw() || rank == 0) send_to_workers(nullptr, true);
-    _listener.join();
+    if(_listener) {
+        _listener->join();
+        delete _listener;
+        _listener = nullptr;
+    }
 }
 
 void MPIMediator::pull(rank_type from, id_type guid, const NDSlice & slice, void * rbuff)

src/MPITransceiver.cpp

@@ -15,6 +15,8 @@ MPITransceiver::MPITransceiver()
             throw std::runtime_error("Your MPI implementation is not MPI_THREAD_MULTIPLE. "
                                      "Please use a thread-safe MPI implementation.");
         }
+    } else {
+        std::cerr << "MPI already initialized\n";
     }
     int nranks, rank;
     MPI_Comm_size(MPI_COMM_WORLD, &nranks);
@@ -23,6 +25,14 @@ MPITransceiver::MPITransceiver()
     _rank = rank;
 };
 
+MPITransceiver::~MPITransceiver()
+{
+    int flag;
+    MPI_Finalized(&flag);
+    if(!flag)
+        MPI_Finalize();
+}
+
 static MPI_Datatype to_mpi(DTypeId T)
 {
     switch(T) {

src/Registry.cpp

@@ -37,4 +37,10 @@ namespace Registry {
         locker _l(_mutex);
         _keeper.erase(id);
     }
+
+    void fini()
+    {
+        locker _l(_mutex);
+        _keeper.clear();
+    }
 }

src/ddptensor.cpp

@@ -53,9 +53,13 @@ bool is_cw()
     return _is_cw && theTransceiver->nranks() > 1;
 }
 
+static bool inited = false;
+static bool finied = false;
+
 // users currently need to call fini to make MPI terminate gracefully
 void fini()
 {
+    if(finied) return;
     delete theMediator;  // stop task is sent in here
     theMediator = nullptr;
     if(pprocessor) {
@@ -65,10 +69,15 @@ void fini()
     }
     delete theTransceiver;
     theTransceiver = nullptr;
+    Deferred::fini();
+    Registry::fini();
+    inited = false;
+    finied = true;
 }
 
 void init(bool cw)
 {
+    if(inited) return;
     if(cw) {
         _is_cw = true;
         if(theTransceiver->rank()) {
@@ -78,6 +87,8 @@ void init(bool cw)
         }
     }
     pprocessor = new std::thread(process_promises);
+    inited = true;
+    finied = false;
 }
 
 // #########################################################################

src/include/ddptensor/Deferred.hpp

@@ -24,6 +24,7 @@ struct Deferred : tensor_i::promise_type
 
     static future_type defer(ptr_type &&, bool);
     static ptr_type undefer_next();
+    static void fini();
 };
 
 template<typename T, typename... Ts>

src/include/ddptensor/MPIMediator.hpp

@@ -7,7 +7,7 @@
 
 class MPIMediator : public Mediator
 {
-    std::thread _listener;
+    std::thread * _listener;
 
 public:
     MPIMediator();

src/include/ddptensor/MPITransceiver.hpp

@@ -8,6 +8,7 @@ class MPITransceiver : public Transceiver
 {
 public:
     MPITransceiver();
+    ~MPITransceiver();
 
     rank_type nranks() const
     {

src/include/ddptensor/Registry.hpp

@@ -17,4 +17,5 @@ namespace Registry {
     extern void put(id_type id, tensor_i::ptr_type ptr);
     tensor_i::ptr_type get(id_type id);
     void del(id_type id);
+    void fini();
 };