Usher server build and runtime fix

CMakeLists.txt

@@ -571,9 +571,10 @@ if(NOT CMAKE_SYSTEM_NAME STREQUAL "Darwin")
 endif()
 
 if(USHER_SERVER)
+    target_include_directories(usher-sampled-server PUBLIC taskflow)
     TARGET_COMPILE_OPTIONS(usher_server PRIVATE)
-    TARGET_LINK_LIBRARIES(usher_server PRIVATE stdc++ ${Boost_LIBRARIES} TBB::tbb ${Protobuf_LIBRARIES} ZLIB::ZLIB) # OpenMP::OpenMP_CXX)
-    TARGET_LINK_LIBRARIES(usher-sampled-server PRIVATE stdc++ ${CMAKE_DL_LIBS} ${Boost_LIBRARIES} TBB::tbb ${Protobuf_LIBRARIES} ZLIB::ZLIB  ${MPI_CXX_LIBRARIES} ${MPI_CXX_LINK_FLAGS} ${ISAL_LIB} ) # OpenMP::OpenMP_CXX)
+    TARGET_LINK_LIBRARIES(usher_server PRIVATE stdc++ ${Boost_LIBRARIES} TBB::tbb TBB::tbbmalloc ${Protobuf_LIBRARIES} ZLIB::ZLIB) # OpenMP::OpenMP_CXX)
+    TARGET_LINK_LIBRARIES(usher-sampled-server PRIVATE stdc++ ${CMAKE_DL_LIBS} ${Boost_LIBRARIES} TBB::tbb TBB::tbbmalloc ${Protobuf_LIBRARIES} ZLIB::ZLIB  ${MPI_CXX_LIBRARIES} ${MPI_CXX_LINK_FLAGS} ${ISAL_LIB} ) # OpenMP::OpenMP_CXX)
     install(TARGETS usher matUtils matOptimize ripples usher_server DESTINATION bin)
 else()
     install(TARGETS usher matUtils matOptimize ripples DESTINATION bin)

src/usher-sampled/driver/socket.cpp

@@ -3,6 +3,7 @@
 #include <atomic>
 #include <boost/filesystem/operations.hpp>
 #include <boost/filesystem/path.hpp>
+#include <boost/filesystem.hpp>
 #include <boost/program_options.hpp>
 #include <boost/program_options/value_semantic.hpp>
 #include <cerrno>
@@ -114,7 +115,6 @@ void reload_trees(TreeCollectionPtr &to_replace, const std::vector<std::string>&
         }
     }
     to_replace.reset(next);
-    init.terminate();
     fprintf(stderr, "finish loading the tree\n");
 }
 void refresh_tree(TreeCollectionPtr &to_replace, std::fstream &tree_paths) {
@@ -583,7 +583,6 @@ static void tree_update_watch(int refresh_period, std::mutex& done_mutex,std::co
                     }
                 }
             }
-            init.terminate();
         }
     }
 }

src/usher-sampled/import_vcf.cpp

@@ -498,7 +498,6 @@ static void process(infile_t &fd, std::vector<Sample_Muts> &sample_mutations,
     read_size = first_approx_size * single_line_size;
     alloc_size = (first_approx_size + 2) * single_line_size;
     tbb::concurrent_bounded_queue<std::pair<char *, uint8_t *>> queue;
-    queue.set_capacity(10);
     tbb::flow::input_node<line_start_later> line(input_graph, line_align(queue));
     tbb::flow::make_edge(line, parser);
     line.activate();

src/usher-sampled/main_mapper.cpp

@@ -7,7 +7,7 @@
 #include <cstdio>
 #include <signal.h>
 #include <tbb/parallel_for.h>
-#include <tbb/task_group.h>
+#include <taskflow/taskflow.hpp>
 #include <unordered_map>
 #include <utility>
 #include <vector>
@@ -378,7 +378,7 @@ struct Main_Tree_Searcher {
 #ifdef DETAILED_MERGER_CHECK
     Mutation_Set &sample_mutations;
 #endif
-    Main_Tree_Searcher(int curr_lower_bound,MAT::Node *node,
+    Main_Tree_Searcher(int curr_lower_bound, const MAT::Node *node,
                        Output<Main_Tree_Target> &output
 #ifdef DETAILED_MERGER_CHECK
                        ,
@@ -393,8 +393,7 @@ struct Main_Tree_Searcher {
 #endif
     {
     }
-    void execute() {
-        tbb::task_group tg;
+    void operator()(tf::Subflow& sf) {
 #ifndef BOUND_CHECK
         if(curr_lower_bound>output.best_par_score) {
             return;
@@ -404,8 +403,10 @@ struct Main_Tree_Searcher {
             search_serial(node, this_muts, output);
             return;
         }
-        std::vector<Main_Tree_Searcher> children_tasks;
-        children_tasks.reserve(node->children.size() + 1);
+        auto* output_ptr = &output;
+#ifdef DETAILED_MERGER_CHECK
+        auto* samples_ptr = &sample_mutations;
+#endif
         Main_Tree_Target target;
         for (const auto child : node->children) {
             target.target_node = child;
@@ -434,14 +435,20 @@ struct Main_Tree_Searcher {
 #ifndef BOUND_CHECK
                 if (lower_bound <= output.best_par_score) {
 #endif
-                    children_tasks.emplace_back(lower_bound, child, output
+                    sf.emplace([=, muts = std::move(descendant_mutations)
 #ifdef DETAILED_MERGER_CHECK
                                            ,
                                            sample_mutations
 #endif
-                                          );
-                    children_tasks.back().this_muts =
-                        std::move(descendant_mutations);
+                                ](tf::Subflow& child_sf) mutable {
+                        Main_Tree_Searcher child_searcher(lower_bound, child, *output_ptr
+#ifdef DETAILED_MERGER_CHECK
+                                                          , *samples_ptr
+#endif
+                        );
+                    child_searcher.this_muts = std::move(muts);
+                    child_searcher(child_sf);
+                    });
 #ifndef BOUND_CHECK
                 }
 #endif
@@ -452,12 +459,6 @@ struct Main_Tree_Searcher {
             assert(parsimony_score>=curr_lower_bound);
             register_target(target, parsimony_score,output);
         }
-        for (auto& child : children_tasks) {
-            tg.run([&child]{
-                child.execute();
-            });
-        }
-        tg.wait();
     }
 };
 
@@ -487,21 +488,25 @@ place_main_tree(const std::vector<To_Place_Sample_Mutation> &mutations,
         }
     }
     output.targets.push_back(target);
-    std::vector<To_Place_Sample_Mutation> initial_muts = mutations;
-    initial_muts.push_back(temp);
 
-    Main_Tree_Searcher main_tree_task_root{0,main_tree.root,
+    tf::Executor executor;
+    tf::Taskflow taskflow;
+
+    taskflow.emplace([&](tf::Subflow& sf) {
+        Main_Tree_Searcher root_searcher(0,main_tree.root,
                        output
 #ifdef DETAILED_MERGER_CHECK
                        ,
                        sample_mutations
 #endif
-    };
-    main_tree_task_root.this_muts = mutations;
-    main_tree_task_root.this_muts.push_back(temp);
-    main_tree_task_root.execute();
+    );
+    root_searcher.this_muts = mutations;
+    root_searcher.this_muts.push_back(temp);
+    root_searcher(sf);
+    });
+
+    executor.run(taskflow).wait();
 
     assert(!output.targets.empty());
-
     return std::make_tuple(std::move(output.targets), output.best_par_score);
 }

src/usher-sampled/sampler.cpp

@@ -6,6 +6,7 @@
 #include <unordered_map>
 #include <vector>
 #include <taskflow/taskflow.hpp>
+#include <memory>
 
 namespace MAT = Mutation_Annotated_Tree;
 
@@ -51,53 +52,48 @@ struct Assign_Descendant_Possible_Muts_Cont {
     }
 };
 
-struct Assign_Descendant_Possible_Muts {
-    tf::Taskflow& taskflow;
-    MAT::Node *root;
-    std::unordered_map<int, uint8_t> &output;
-    Assign_Descendant_Possible_Muts(tf::Taskflow& t, MAT::Node *root,
-                                    std::unordered_map<int, uint8_t> &output)
-        : taskflow{t}, root(root), output(output) {}
-    void execute() const {
-        Assign_Descendant_Possible_Muts_Cont cont(output, root->children.size(),
-                                                  root);
-        std::vector<Assign_Descendant_Possible_Muts> children_tasks;
-        children_tasks.reserve(root->children.size());
-        for (size_t idx = 0; idx < root->children.size(); idx++) {
-            auto this_child = root->children[idx];
-            if (this_child->children.empty()) {
-                cont.children_out[idx].reserve(this_child->mutations.size());
-                for (auto &mut : this_child->mutations) {
-                    mut.set_descendant_mut(mut.get_mut_one_hot());
-                    cont.children_out[idx].emplace(mut.get_position(),
-                                                   mut.get_mut_one_hot());
-                }
-            } else {
-                children_tasks.emplace_back(taskflow,
-                                            this_child, cont.children_out[idx]);
+void build_task_graph(MAT::Node* node, std::unordered_map<int, uint8_t>& output, tf::Subflow& subflow) {
+    // 1. Allocate Context on the HEAP (Safety Fix)
+    auto cont_ptr = std::make_shared<Assign_Descendant_Possible_Muts_Cont>(output, node->children.size(), node);
+
+    // 2. Create Continuation Task
+    tf::Task continuation = subflow.emplace([cont_ptr]() {
+        cont_ptr->execute();
+    });
+
+    // 3. Spawn Children
+    for (size_t i = 0; i < node->children.size(); ++i) {
+        auto child = node->children[i];
+
+        if (child->children.empty()) {
+            // Leaf optimization (Serial execution)
+            cont_ptr->children_out[i].reserve(child->mutations.size());
+            for (auto &mut : child->mutations) {
+                mut.set_descendant_mut(mut.get_mut_one_hot());
+                cont_ptr->children_out[i].emplace(mut.get_position(), mut.get_mut_one_hot());
             }
-        }
-        if (children_tasks.empty()) {
-            cont.execute();
         } else {
-          taskflow.emplace([&children_tasks](tf::Subflow& subflow){
-            for (auto&& child_task : std::move(children_tasks)) {
-                subflow.emplace([child_task = std::move(child_task)] {
-                    child_task.execute();
-                });
-            }
+            // Recursion: Create a subflow for the child
+            // Capture 'cont_ptr' and 'i' by value to be safe
+            tf::Task child_task = subflow.emplace([child, cont_ptr, i](tf::Subflow& sf) {
+                build_task_graph(child, cont_ptr->children_out[i], sf);
+            });
 
-          });
+            // The continuation (execute) must wait for the child to finish
+            child_task.precede(continuation);
         }
     }
-};
+}
 
+// Entry Point
 void assign_descendant_muts(MAT::Tree &in) {
     std::unordered_map<int, uint8_t> ignore;
     tf::Executor executor;
     tf::Taskflow taskflow;
-    taskflow.emplace([&] {
-      Assign_Descendant_Possible_Muts(taskflow, in.root, ignore).execute();
+
+    taskflow.emplace([&](tf::Subflow& subflow) {
+        build_task_graph(in.root, ignore, subflow);
     });
+
     executor.run(taskflow).wait(); 
-}
+}

src/usher_server.cpp

@@ -30,7 +30,7 @@ int main(int argc, char** argv) {
     //Variables to load command-line options using Boost program_options
     std::string arg_dirname;
     std::string MAT_list_filename;
-    uint32_t num_cores = tbb::task_scheduler_init::default_num_threads();
+    uint32_t num_cores = tbb::global_control::active_value(tbb::global_control::max_allowed_parallelism);
     po::options_description desc{"Options"};
     uint32_t sleep_length;
     uint32_t termination_character;
@@ -81,7 +81,7 @@ int main(int argc, char** argv) {
     // timer object to be used to measure runtimes of individual stages
     Timer timer;
     fprintf(stderr, "Initializing %u worker threads.\n\n", num_threads);
-    tbb::task_scheduler_init init(num_threads);
+    tbb::global_control global_limit(tbb::global_control::max_allowed_parallelism, num_threads);
 
 
     //MAT that is used in the iteration