Add checkpointing to metagraph build

metagraph/integration_tests/test_build.py

@@ -337,6 +337,78 @@ def test_build_chunks_from_kmc_canonical(self, build):
         self.assertEqual('nodes (k): 802920', params_str[1])
         self.assertEqual('canonical mode: yes', params_str[2])
 
+    @parameterized.expand(['succinct_disk'])
+    def test_build_phase(self, build):
+        representation, tmp_dir = build_params[build]
+
+        construct_command = '{exe} build --phase 1 --mask-dummy --graph {repr} --canonical -k 20 ' \
+                            '--disk-swap {tmp_dir} -o {outfile} {input}'.format(
+            exe=METAGRAPH,
+            repr=representation,
+            tmp_dir=tmp_dir,
+            outfile=self.tempdir.name + '/graph',
+            input=TEST_DATA_DIR + '/transcripts_1000.fa'
+        )
+        res = subprocess.run([construct_command], shell=True)
+        self.assertEqual(res.returncode, 0)
+        self.assertTrue(os.path.isfile(self.tempdir.name + '/graph.checkpoint'))
+
+        construct_command = '{exe} build --mask-dummy --graph {repr} --canonical -k 20 ' \
+                            '--disk-swap {tmp_dir} -o {outfile} {input}'.format(
+            exe=METAGRAPH,
+            repr=representation,
+            tmp_dir=tmp_dir,
+            outfile=self.tempdir.name + '/graph',
+            input=TEST_DATA_DIR + '/transcripts_1000.fa'
+        )
+        res = subprocess.run([construct_command], shell=True)
+        self.assertEqual(res.returncode, 0)
+
+        res = self.__get_stats(self.tempdir.name + '/graph' + graph_file_extension[representation])
+        self.assertEqual(res.returncode, 0)
+        params_str = res.stdout.decode().split('\n')[2:]
+        self.assertEqual('k: 20', params_str[0])
+        self.assertEqual('nodes (k): 1159851', params_str[1])
+        self.assertEqual('canonical mode: yes', params_str[2])
+        self.assertFalse(os.path.isfile(self.tempdir.name + '/graph.checkpoint'))
+
+    # tests that we can build and resume 2 separate graphs on the same machine
+    @parameterized.expand(['succinct_disk'])
+    def test_build_phase_parallel(self, build):
+        representation, tmp_dir = build_params[build]
+
+        for name in ('graph1', 'graph2'):
+            construct_command = '{exe} build --phase 1 --mask-dummy --graph {repr} --canonical -k 20 ' \
+                                '--disk-swap {tmp_dir} -o {outfile} {input}'.format(
+                exe=METAGRAPH,
+                repr=representation,
+                tmp_dir=tmp_dir,
+                outfile=self.tempdir.name + '/' + name,
+                input=TEST_DATA_DIR + ('/transcripts_1000.fa' if name == 'graph1' else '/transcripts_100.fa')
+            )
+            res = subprocess.run([construct_command], shell=True)
+            self.assertEqual(res.returncode, 0)
+            self.assertTrue(os.path.isfile(self.tempdir.name + '/' + name + '.checkpoint'))
+
+        for name in ('graph', 'graph2'):
+            construct_command = '{exe} build --mask-dummy --graph {repr} --canonical -k 20 ' \
+                                '--disk-swap {tmp_dir} -o {outfile} {input}'.format(
+                exe=METAGRAPH,
+                repr=representation,
+                tmp_dir=tmp_dir,
+                outfile=self.tempdir.name + '/' + name,
+                input=TEST_DATA_DIR + ('/transcripts_1000.fa' if name == 'graph1' else '/transcripts_100.fa')
+            )
+            res = subprocess.run([construct_command], shell=True)
+            self.assertEqual(res.returncode, 0)
+
+            res = self.__get_stats(self.tempdir.name + '/graph' + graph_file_extension[representation])
+            self.assertEqual(res.returncode, 0)
+            params_str = res.stdout.decode().split('\n')[2:]
+            self.assertEqual('k: 20', params_str[0])
+            self.assertEqual('nodes (k): ' + ('1159851' if name == 'graph1' else '91584'), params_str[1])
+            self.assertEqual('canonical mode: yes', params_str[2])
+            self.assertFalse(os.path.isfile(self.tempdir.name + '/' + name + '.checkpoint'))
 
 if __name__ == '__main__':
     unittest.main()

metagraph/integration_tests/test_build_weighted.py

@@ -373,5 +373,41 @@ def test_kmer_count_width_large(self, build, k_width_result):
         self.assertEqual('avg weight: {}'.format(avg_count_expected), params_str[4])
 
 
+    def test_build_phase(self):
+        construct_command = '{exe} build --phase 1 --mask-dummy -k 20 --count-kmers --disk-swap {tmp_dir} ' \
+                            '--count-width 16 -o {outfile} {input}'.format(
+            exe=METAGRAPH,
+            tmp_dir=self.tempdir.name,
+            outfile=self.tempdir.name + '/graph',
+            input=TEST_DATA_DIR + '/transcripts_1000.fa'
+        )
+        res = subprocess.run([construct_command], shell=True)
+        self.assertEqual(res.returncode, 0)
+        self.assertTrue(os.path.isfile(self.tempdir.name + '/graph.checkpoint'))
+
+        construct_command = '{exe} build --mask-dummy -k 20 --count-kmers --disk-swap {tmp_dir} --count-width 16 ' \
+                            '-o {outfile} {input}'.format(
+            exe=METAGRAPH,
+            tmp_dir=self.tempdir.name,
+            outfile=self.tempdir.name + '/graph',
+            input=TEST_DATA_DIR + '/transcripts_1000.fa'
+        )
+        res = subprocess.run([construct_command], shell=True)
+        self.assertEqual(res.returncode, 0)
+
+        stats_command = '{exe} stats {graph}'.format(
+            exe=METAGRAPH,
+            graph=self.tempdir.name + '/graph.dbg',
+        )
+        res = subprocess.run(stats_command.split(), stdout=PIPE)
+        self.assertEqual(res.returncode, 0)
+        params_str = res.stdout.decode().split('\n')[2:]
+        self.assertEqual('k: 20', params_str[0])
+        self.assertEqual('nodes (k): 591997', params_str[1])
+        self.assertEqual('canonical mode: no', params_str[2])
+        self.assertEqual('nnz weights: 591997', params_str[3])
+        self.assertEqual('avg weight: 2.48587', params_str[4])
+
+
 if __name__ == '__main__':
     unittest.main()

metagraph/src/cli/build.cpp

@@ -11,6 +11,7 @@
 #include "graph/representation/bitmap/dbg_bitmap_construct.hpp"
 #include "graph/representation/succinct/dbg_succinct.hpp"
 #include "graph/representation/succinct/boss_construct.hpp"
+#include "graph/representation/succinct/build_checkpoint.hpp"
 #include "graph/graph_extensions/node_weights.hpp"
 #include "config/config.hpp"
 #include "parse_sequences.hpp"
@@ -105,6 +106,9 @@ int build_graph(Config *config) {
                 logger->info("k-mer suffix: '{}'", suffix);
             }
 
+            bool checkpoint_enabled = !config->tmp_dir.empty() && suffixes.size() == 1;
+            boss::BuildCheckpoint checkpoint(checkpoint_enabled ? config->outfbase : "",
+                                             config->phase);
             auto constructor = boss::IBOSSChunkConstructor::initialize(
                 boss_graph->get_k(),
                 config->canonical,
@@ -115,15 +119,25 @@ int build_graph(Config *config) {
                 config->tmp_dir.empty() ? kmer::ContainerType::VECTOR
                                         : kmer::ContainerType::VECTOR_DISK,
                 config->tmp_dir,
-                config->disk_cap_bytes
+                config->disk_cap_bytes,
+                checkpoint
             );
 
-            push_sequences(files, *config, timer, constructor.get());
+            if (checkpoint.checkpoint() == 0) {
+                push_sequences(files, *config, timer, constructor.get());
+            } else {
+                logger->info("Skipping parsing sequences from input file(s)");
+            }
 
             boss::BOSS::Chunk *next_chunk = constructor->build_chunk();
+
+            if (checkpoint.phase() < 2) { // phase 1 stops after generating dummy k-mers
+                assert(next_chunk == nullptr);
+                return 0;
+            }
+
             logger->trace("Graph chunk with {} k-mers was built in {} sec",
                           next_chunk->size() - 1, timer.elapsed());
-
             if (config->suffix.size()) {
                 logger->info("Serialize the graph chunk for suffix '{}'...", suffix);
                 timer.reset();
@@ -140,6 +154,7 @@ int build_graph(Config *config) {
             } else {
                 graph_data.reset(next_chunk);
             }
+            checkpoint.remove_checkpoint();
         }
 
         assert(graph_data);

metagraph/src/cli/config/config.cpp

@@ -334,6 +334,8 @@ Config::Config(int argc, char *argv[]) {
             tmp_dir = get_value(i++);
         } else if (!strcmp(argv[i], "--disk-cap-gb")) {
             disk_cap_bytes = atoi(get_value(i++)) * 1e9;
+        } else if (!strcmp(argv[i], "--phase")) {
+            phase = atoi(get_value(i++));
         } else if (argv[i][0] == '-') {
             fprintf(stderr, "\nERROR: Unknown option %s\n\n", argv[i]);
             print_usage(argv[0], identity);
@@ -521,6 +523,30 @@ Config::Config(int argc, char *argv[]) {
     if (identity == COMPARE && fnames.size() != 2)
         print_usage_and_exit = true;
 
+    if (identity != BUILD && phase != 2) {
+        std::cerr << "Error: Phases are only supported for building. Remove --phase"
+                  << std::endl;
+        print_usage_and_exit = true;
+    }
+
+    if (identity != BUILD && phase > 2) {
+        std::cerr << "Error: Invalid phase value. Can be either 1 or 2." << std::endl;
+        print_usage_and_exit = true;
+    }
+
+    if (phase != 2 && tmp_dir.empty()) {
+        std::cerr << "Error: Phases are only supported for disk-based building. "
+                     "Please set --disk-swap." << std::endl;
+        print_usage_and_exit = true;
+    }
+
+    if (phase != 2 && suffix_len > 0) {
+        std::cerr << "Error: Phases are not supported for multiple suffixes. Remove "
+                     "--phase or specify each suffix separately using --suffix"
+                  << std::endl;
+        print_usage_and_exit = true;
+    }
+
     if (discovery_fraction < 0 || discovery_fraction > 1)
         print_usage_and_exit = true;
 
@@ -751,6 +777,7 @@ void Config::print_usage(const std::string &prog_name, IdentityType identity) {
             fprintf(stderr, "\t-p --parallel [INT] \tuse multiple threads for computation [1]\n");
             fprintf(stderr, "\t   --disk-swap [STR] \tdirectory to use for temporary files [off]\n");
             fprintf(stderr, "\t   --disk-cap-gb [INT] \tmax temp disk space to use before forcing a merge, in GB [20]\n");
+            fprintf(stderr, "\t   --phase [INT] \tmax where to stop the computation (1=generate kmers, 2=build all) [2]\n");
         } break;
         case CLEAN: {
             fprintf(stderr, "Usage: %s clean -o <outfile-base> [options] GRAPH\n\n", prog_name.c_str());

metagraph/src/cli/config/config.hpp

@@ -142,6 +142,8 @@ class Config {
 
     size_t disk_cap_bytes = 20e9; // 20GB default
 
+    uint32_t phase = 2; // build phase; 1 = generate kmers, 2 = complete build
+
     enum IdentityType {
         NO_IDENTITY = -1,
         BUILD = 1,

metagraph/src/common/elias_fano.cpp

@@ -14,29 +14,30 @@
 namespace mtg {
 namespace common {
 
-void concat(const std::vector<std::string> &files, const std::string &result) {
+std::vector<std::string> concat(const std::vector<std::string> &files, const std::string &result) {
     if (files.empty())
-        return;
+        return {};
+
+    std::vector<std::string> original_files;
 
     std::vector<std::string> suffixes = { "", ".up" };
     if (std::filesystem::exists(files[0] + ".count"))
         suffixes.push_back(".count");
 
     for (const std::string &suffix : suffixes) {
         std::string concat_command = "cat ";
-        for (uint32_t i = 1; i < files.size(); ++i) {
+        for (uint32_t i = 0; i < files.size(); ++i) {
             concat_command += files[i] + suffix + " ";
         }
-        concat_command += " >> " + files[0] + suffix;
-
+        concat_command += " > " + result + suffix;
         if (std::system(concat_command.c_str()))
             throw std::runtime_error("Error while cat-ing files: " + concat_command);
 
-        std::filesystem::rename(files[0] + suffix, result + suffix);
         for (const std::string &f : files) {
-            std::filesystem::remove(f + suffix);
+            original_files.push_back(f + suffix);
         }
     }
+    return original_files;
 }
 
 template <class T, class Enable = void>

metagraph/src/common/elias_fano.hpp

@@ -21,7 +21,7 @@ namespace common {
  * The files store data that is ordered and the values in a file are smaller than the
  * values in the next file.
  */
-void concat(const std::vector<std::string> &files, const std::string &result);
+std::vector<std::string> concat(const std::vector<std::string> &files, const std::string &result);
 
 /**
  * Elias-Fano encoder that streams the encoded result into a file.

metagraph/src/common/sorted_sets/sorted_set_disk_base.cpp

@@ -64,13 +64,15 @@ std::vector<std::string> SortedSetDiskBase<T>::files_to_merge() {
 }
 
 template <typename T>
-void SortedSetDiskBase<T>::clear(const std::filesystem::path &tmp_path) {
+void SortedSetDiskBase<T>::clear(const std::filesystem::path &tmp_path, bool remove_files) {
     std::unique_lock<std::mutex> exclusive_lock(mutex_);
     std::unique_lock<std::shared_timed_mutex> multi_insert_lock(multi_insert_mutex_);
     is_merging_ = false;
-    // remove the files that have not been requested to merge
-    for (const auto &chunk_file : get_file_names()) {
-        std::filesystem::remove(chunk_file);
+    if (remove_files) {
+        // remove the files that have not been requested to merge
+        for (const auto &chunk_file : get_file_names()) {
+            std::filesystem::remove(chunk_file);
+        }
     }
     chunk_count_ = 0;
     l1_chunk_count_ = 0;
@@ -91,7 +93,7 @@ void SortedSetDiskBase<T>::start_merging_async() {
     async_worker_.enqueue([file_names, this]() {
         std::function<void(const T &)> on_new_item
                 = [this](const T &v) { merge_queue_.push(v); };
-        merge_files(file_names, on_new_item);
+        merge_files(file_names, on_new_item, false);
         merge_queue_.shutdown();
     });
 }

metagraph/src/common/sorted_sets/sorted_set_disk_base.hpp

@@ -49,10 +49,8 @@ class SortedSetDiskBase {
                       size_t max_disk_space_bytes);
 
     virtual ~SortedSetDiskBase() {
-        // remove the files that have not been requested to merge
-        for (const auto &chunk_file : get_file_names()) {
-            std::filesystem::remove(chunk_file);
-        }
+        // not cleaning up unmerged chunk_*** files so that the computation can be resumed
+        // if building in phases or in case of a crash
         async_worker_.join(); // make sure the data was processed
     }
 
@@ -74,7 +72,7 @@ class SortedSetDiskBase {
      * sorted set may be expensive when #data_ is large. In these cases, prefer calling
      * #clear and re-using the buffer.
      */
-    void clear(const std::filesystem::path &tmp_path = "/tmp/");
+    void clear(const std::filesystem::path &tmp_path = "/tmp/", bool remove_files = true);
 
   protected:
     /** Advances #it by step or points to #end, whichever comes first. */

metagraph/src/common/utils/file_utils.cpp

@@ -38,7 +38,8 @@ void cleanup_tmp_dir_on_exit() {
 }
 
 std::filesystem::path create_temp_dir(std::filesystem::path path,
-                                      const std::string &name) {
+                                      const std::string &name,
+                                      bool clean_on_exit) {
     if (path.empty())
         path = "./";
 
@@ -48,6 +49,12 @@ std::filesystem::path create_temp_dir(std::filesystem::path path,
         exit(1);
     }
 
+    logger->trace("Created temporary directory {}", tmp_dir_str);
+
+    if (!clean_on_exit) {
+        return tmp_dir_str;
+    }
+
     if (TMP_DIRS.empty()) {
         if (std::signal(SIGINT, cleanup_tmp_dir_on_signal) == SIG_ERR)
             logger->error("Couldn't reset the signal handler for SIGINT");
@@ -57,8 +64,6 @@ std::filesystem::path create_temp_dir(std::filesystem::path path,
             logger->error("Couldn't reset the atexit handler");
     }
 
-    logger->trace("Registered temporary directory {}", tmp_dir_str);
-
     static std::mutex mu;
     std::lock_guard<std::mutex> lock(mu);
 

metagraph/src/common/utils/file_utils.hpp

@@ -14,7 +14,8 @@
 namespace utils {
 
 std::filesystem::path create_temp_dir(std::filesystem::path path,
-                                      const std::string &name = "");
+                                      const std::string &name = "",
+                                      bool clean_on_exit = true);
 
 
 bool check_if_writable(const std::string &filename);