Add tests, small fixes

Author: Matthew-Whitlock

src/vt/objgroup/proxy/proxy_objgroup.h

@@ -460,8 +460,10 @@ struct Proxy {
 
     if constexpr(has_user_traits_v<Serializer, CheckpointTrait>){
       vtAssert(old_proxy != no_obj_group, "ObjGroups must be pre-instantiated to be checkpointed or restored");
-      vtAssert(old_proxy == proxy_, "The proxy ID bits of this ObjGroup do not match the ID found in the checkpoint!" \
-                                    " Varying IDs is not yet supported.");
+      vtWarnIf(old_proxy != proxy_,
+               "Checkpointed ObjGroup and deserialize target do not match!");
+      proxy_ = old_proxy;
+
       auto objPtr = get();
 
       bool null = objPtr == nullptr;

src/vt/vrt/collection/manager.h

@@ -352,7 +352,7 @@ struct CollectionManager
    * \param[in] is_collective whether the collection is collective
    * \param[in] is_migratable whether the collection is migratable
    * \param[in] request_match an input proxy which we would like to use, if
-   *            there are no existing conflicts. no_vrt_proxy indicates no 
+   *            there are no existing conflicts. no_vrt_proxy indicates no
    *            request.
    *
    * \return the collection proxy bits
@@ -1562,7 +1562,7 @@ struct CollectionManager
    */
   template <typename ColT, typename IndexT = typename ColT::IndexType>
   bool getDynamicMembership(VirtualProxyType proxy);
-  
+
   /**
    * \brief Get the local indices that are currently on this node
    *
@@ -1635,28 +1635,16 @@ struct CollectionManager
   );
 
   /**
-   * \internal \struct MigrateRequestMsg
+   * \internal \brief Migrate element to restore location from checkpoint
    *
-   * \brief Migrate local element, potentially requested by remote location
-   */
-
-  /**
-   * \brief Migrate a remote proxy element to a node, by messaging that
-   *        node to initiate a migration. Immediately returns a rooted epoch
-   *        containing the request message.
-   */
-  template <typename ColT>
-  EpochType requestMigrateDeferred(
-    VrtElmProxy<ColT, typename ColT::IndexType> proxy_elem, NodeType destination
-  );
-
-  /**
-   * \brief Migrate a remote proxy element to a node, by messaging that
-   *        node to initiate a migration. Returns after migration is complete.
+   * \param[in] node the node
+   * \param[in] idx the element index
+   * \param[in] proxy the collection proxy
    */
   template <typename ColT>
-  void requestMigrate(
-    VrtElmProxy<ColT, typename ColT::IndexType> proxy_elem, NodeType destination
+  static void migrateToRestoreLocation(
+    NodeType node, typename ColT::IndexType idx,
+    CollectionProxyWrapType<ColT> proxy
   );
 
   /**

src/vt/vrt/collection/manager.impl.h

@@ -2249,40 +2249,28 @@ void CollectionManager::checkpointToFile(
 
 namespace detail {
 template <typename ColT>
-inline void MigrateRequestHandler (
-  ColT*, VrtElmProxy<ColT, typename ColT::IndexType> proxy_elm, NodeType dest
+inline void restoreOffHomeElement(
+  ColT*, NodeType node, typename ColT::IndexType idx,
+  CollectionProxy<ColT> proxy
 ) {
-  theCollection()->migrate(proxy_elm, dest);
+  theCollection()->migrate(proxy(idx), node);
 }
 } /* end namespace detail */
 
 template <typename ColT>
-EpochType CollectionManager::requestMigrateDeferred(
-  VrtElmProxy<ColT, typename ColT::IndexType> proxy_elm, NodeType destination
-) {
-  auto ep = theTerm()->makeEpochRooted(
-      "Request element migration", term::UseDS{true}
-  );
-  theMsg()->pushEpoch(ep);
-
-  proxy_elm.template send<detail::MigrateRequestHandler<ColT>>(
-      proxy_elm, destination
-  );
-
-  theMsg()->popEpoch(ep);
-  theTerm()->finishedEpoch(ep);
-  return ep;
-}
-
-template <typename ColT>
-void CollectionManager::requestMigrate(
-  VrtElmProxy<ColT, typename ColT::IndexType> proxy_elem, NodeType destination
+/*static*/ void CollectionManager::migrateToRestoreLocation(
+  NodeType node, typename ColT::IndexType idx,
+  CollectionProxyWrapType<ColT> proxy
 ) {
-   auto ep = requestMigrateDeferred(proxy_elem, destination);
-   vt::runSchedulerThrough(ep);
+  if (proxy(idx).tryGetLocalPtr() != nullptr) {
+    theCollection()->migrate(proxy(idx), node);
+  } else {
+    proxy(idx).template send<detail::restoreOffHomeElement<ColT>>(
+      node, idx, proxy
+    );
+  }
 }
 
-
 template <typename ColT>
 void CollectionManager::restoreFromFileInPlace(
   CollectionProxyWrapType<ColT> proxy,
@@ -2309,14 +2297,23 @@ void CollectionManager::restoreFromFileInPlace(
     metadata_file_name
   );
 
-  //Everyone shuffles any elements not where their data is
   runInEpochCollective([&]{
     for (auto&& elm : directory->elements_) {
       auto idx = elm.idx_;
       auto file_name = elm.file_name_;
 
       if (proxy(idx).tryGetLocalPtr() == nullptr) {
-        requestMigrateDeferred(proxy(idx), theContext()->getNode());
+        auto mapped_node = getMappedNode<ColT>(proxy, idx);
+        vtAssertExpr(mapped_node != uninitialized_destination);
+        auto this_node = theContext()->getNode();
+
+        if (mapped_node != this_node) {
+          theMsg()->send<migrateToRestoreLocation<ColT>>(
+            vt::Node{mapped_node}, this_node, idx, proxy
+          );
+        } else {
+          migrateToRestoreLocation<ColT>(this_node, idx, proxy);
+        }
       }
     }
   });

src/vt/vrt/proxy/collection_elm_proxy.impl.h

@@ -58,7 +58,8 @@ void VrtElmProxy<ColT, IndexT>::serialize(SerT& s) {
   //Only serialize the ColT object if checkpointing.
   if constexpr(checkpoint::has_user_traits_v<SerT, CheckpointTrait>){
     ColT* local_elm_ptr = this->tryGetLocalPtr();
-    vtAssert(local_elm_ptr != nullptr || s.isUnpacking(), "Must serialize/size elements from the node they are at");
+    vtAssert(local_elm_ptr != nullptr || s.isUnpacking(),
+        "Must serialize/size elements from the node they are at");
 
     //Traits for nested serialize/deserialize
     using checkpoint::serializerUserTraits::CopyTraits;
@@ -67,43 +68,73 @@ void VrtElmProxy<ColT, IndexT>::serialize(SerT& s) {
                                 ::template With<CheckpointInternalTrait>
     >;
 
-    //Weird nested serialization to enable asynchronous deserializing w/o changing semantics.
+    //Weird nested serialization to enable asynchronous deserializing w/o
+    //changing semantics.
     if(!(s.isPacking() || s.isUnpacking())){
-      int size = checkpoint::getSize<ColT, CheckpointlessTraits>(*local_elm_ptr);
+      int size = checkpoint::getSize<ColT, CheckpointlessTraits>(
+        *local_elm_ptr
+      );
       s | size;
       //Don't use nullptr to avoid warning
       s.contiguousBytes(&size, 1, size);
     } else if(s.isPacking()){
-      auto serialized_elm = checkpoint::serialize<ColT, CheckpointlessTraits>(*local_elm_ptr);
+      auto serialized_elm = checkpoint::serialize<ColT, CheckpointlessTraits>(
+        *local_elm_ptr
+      );
       int size = serialized_elm->getSize();
       s | size;
       s.contiguousBytes(serialized_elm->getBuffer(), 1, size);
     } else if(s.isUnpacking()){
       int size = 0;
       s | size;
 
-      auto buf = std::make_unique<char[]>(size);
-      s.contiguousBytes(buf.get(), 1, size);
+      auto buf = std::make_shared<std::vector<char>>(size);
+      s.contiguousBytes(buf->data(), 1, size);
 
       if(local_elm_ptr != nullptr){
-        checkpoint::deserializeInPlace<ColT, CheckpointlessTraits>(buf.get(), local_elm_ptr);
+        checkpoint::deserializeInPlace<ColT, CheckpointlessTraits>(
+          buf->data(), local_elm_ptr
+        );
       } else {
-        //The element is somewhere else so we'll need to request a migration to here.
-        auto ep = theCollection()->requestMigrateDeferred(*this, theContext()->getNode());
+        //TODO: Investigate skipping the actual full migration and instead
+        //  simply deleting remote node's copy, building a fresh one here,
+        //  and updating system as if migration happened. Would probably be
+        //  noticeably faster on large systems and w/ large elements.
+        CollectionProxy<ColT, IndexT> col(this->getCollectionProxy());
+        std::shared_ptr<int> listener_id = std::make_shared<int>(-1);
 
-        theTerm()->addActionUnique(ep, std::move([elm_proxy = *this, buffer = std::move(buf)]{
-          auto elm_ptr = elm_proxy.tryGetLocalPtr();
-          assert(elm_ptr != nullptr);
-          checkpoint::deserializeInPlace<ColT, CheckpointlessTraits>(buffer.get(), elm_ptr);
-        }));
+        //Listen for this element to migrate in,
+        //then immediately deserialize
+        using listener::ElementEventEnum;
+        listener::ListenFnType<IndexT> m_listener =
+          [*this, buf, listener_id]
+          (ElementEventEnum event, IndexT idx, NodeType) mutable {
+            if(!(idx == getIndex())) return;
+            if(event != ElementEventEnum::ElementMigratedIn) return;
+
+            auto elm_ptr = this->tryGetLocalPtr();
+            checkpoint::deserializeInPlace<ColT, CheckpointlessTraits>(
+              buf->data(), elm_ptr
+            );
+            theCollection()->unregisterElementListener<ColT>(
+                this->getCollectionProxy(), *listener_id
+            );
+          };
+        *listener_id = theCollection()->registerElementListener<ColT>(
+          this->getCollectionProxy(), m_listener
+        );
+
+        theCollection()->migrateToRestoreLocation(
+          theContext()->getNode(), getIndex(), col
+        );
       }
     }
   }
 }
 
 template <typename ColT, typename IndexT>
 template <typename SerT>
-std::unique_ptr<ColT> 
+std::unique_ptr<ColT>
 VrtElmProxy<ColT, IndexT>::deserializeToElm(SerT& s) {
   //Still have to hit data in the same order.
   ProxyCollectionElmTraits<ColT, IndexT>::serialize(s);

src/vt/vrt/proxy/collection_proxy.impl.h

@@ -45,6 +45,7 @@
 #define INCLUDED_VT_VRT_PROXY_COLLECTION_PROXY_IMPL_H
 
 #include "vt/config.h"
+#include "vt/configs/error/soft_error.h"
 #include "vt/vrt/proxy/collection_proxy.h"
 #include "vt/vrt/proxy/base_elm_proxy.h"
 #include "vt/vrt/collection/proxy_traits/proxy_col_traits.h"
@@ -152,12 +153,9 @@ void CollectionProxy<ColT, IndexT>::serialize(
   //TODO: magistrate's virtualized serialization support may enable checkpointing
   //mapper objects. Pre-registered functions could be doable as well.
 
-  //TODO: chkpt location manager so we don't have to message back and forth so much?
-  //auto lm = theLocMan()->getCollectionLM<IndexType>(proxy.getProxy());
-
   //If unpacking, we may need to make the collection to unpack into.
   if(s.isUnpacking() && oldProxy != proxy){
-    vtAssert(oldProxy == no_vrt_proxy,
+    vtWarnIf(oldProxy != no_vrt_proxy,
              "Checkpointed proxy and deserialize target do not match!");
 
     //The checkpointed proxy doesn't exist, we need to create it