Add checkpointing to metagraph build

metagraph/integration_tests/test_build.py

@@ -3,7 +3,6 @@
 import subprocess
 from subprocess import PIPE
 from tempfile import TemporaryDirectory
-import glob
 import os
 
 
@@ -44,7 +43,7 @@ def test_simple_all_graphs(self, build):
         construct_command = '{exe} build --mask-dummy --graph {repr} --disk-swap {tmp_dir} -k 20 -o {outfile} {input}'.format(
             exe=METAGRAPH,
             repr=representation,
-            tmp_dir=tmp_dir,
+            tmp_dir='""' if tmp_dir == '""' else self.tempdir.name,
             outfile=self.tempdir.name + '/graph',
             input=TEST_DATA_DIR + '/transcripts_1000.fa'
         )
@@ -66,7 +65,7 @@ def test_simple_bloom_graph(self, build):
         construct_command = '{exe} build --mask-dummy --graph {repr} --disk-swap {tmp_dir} -k 20 -o {outfile} {input}'.format(
             exe=METAGRAPH,
             repr=representation,
-            tmp_dir=tmp_dir,
+            tmp_dir='""' if tmp_dir == '""' else self.tempdir.name,
             outfile=self.tempdir.name + '/graph',
             input=TEST_DATA_DIR + '/transcripts_1000.fa'
         )
@@ -111,7 +110,7 @@ def test_simple_all_graphs_canonical(self, build):
                 --graph {repr} --canonical -k 20 -o {outfile} {input}'.format(
             exe=METAGRAPH,
             repr=representation,
-            tmp_dir=tmp_dir,
+            tmp_dir='""' if tmp_dir == '""' else self.tempdir.name,
             outfile=self.tempdir.name + '/graph',
             input=TEST_DATA_DIR + '/transcripts_1000.fa'
         )
@@ -176,7 +175,7 @@ def test_build_from_kmc(self, build):
         construct_command = '{exe} build --mask-dummy --graph {repr} --disk-swap {tmp_dir} -k 11 -o {outfile} {input}'.format(
             exe=METAGRAPH,
             repr=representation,
-            tmp_dir=tmp_dir,
+            tmp_dir='""' if tmp_dir == '""' else self.tempdir.name,
             outfile=self.tempdir.name + '/graph',
             input=TEST_DATA_DIR + '/transcripts_1000_kmc_counters.kmc_suf'
         )
@@ -198,7 +197,7 @@ def test_build_from_kmc_both(self, build):
         construct_command = '{exe} build --mask-dummy --graph {repr} --disk-swap {tmp_dir} -k 11 -o {outfile} {input}'.format(
             exe=METAGRAPH,
             repr=representation,
-            tmp_dir=tmp_dir,
+            tmp_dir='""' if tmp_dir == '""' else self.tempdir.name,
             outfile=self.tempdir.name + '/graph',
             input=TEST_DATA_DIR + '/transcripts_1000_kmc_counters_both_strands.kmc_suf'
         )
@@ -221,7 +220,7 @@ def test_build_from_kmc_canonical(self, build):
                 --graph {repr} --disk-swap {tmp_dir} --canonical -k 11 -o {outfile} {input}'.format(
             exe=METAGRAPH,
             repr=representation,
-            tmp_dir=tmp_dir,
+            tmp_dir='""' if tmp_dir == '""' else self.tempdir.name,
             outfile=self.tempdir.name + '/graph',
             input=TEST_DATA_DIR + '/transcripts_1000_kmc_counters.kmc_suf'
         )
@@ -244,7 +243,7 @@ def test_build_from_kmc_both_canonical(self, build):
                 --graph {repr} --disk-swap {tmp_dir} --canonical -k 11 -o {outfile} {input}'.format(
             exe=METAGRAPH,
             repr=representation,
-            tmp_dir=tmp_dir,
+            tmp_dir='""' if tmp_dir == '""' else self.tempdir.name,
             outfile=self.tempdir.name + '/graph',
             input=TEST_DATA_DIR + '/transcripts_1000_kmc_counters_both_strands.kmc_suf'
         )
@@ -269,7 +268,7 @@ def test_build_chunks_from_kmc(self, build):
                                 --graph {repr} -k 11 --suffix {suffix} -o {outfile} {input}'.format(
                 exe=METAGRAPH,
                 repr=representation,
-                tmp_dir=tmp_dir,
+                tmp_dir='""' if tmp_dir == '""' else self.tempdir.name,
                 outfile=self.tempdir.name + '/graph',
                 input=TEST_DATA_DIR + '/transcripts_1000_kmc_counters.kmc_suf',
                 suffix=suffix
@@ -308,7 +307,7 @@ def test_build_chunks_from_kmc_canonical(self, build):
                     --suffix {suffix} -o {outfile} {input}'.format(
                 exe=METAGRAPH,
                 repr=representation,
-                tmp_dir=tmp_dir,
+                tmp_dir='""' if tmp_dir == '""' else self.tempdir.name,
                 outfile=self.tempdir.name + '/graph',
                 input=TEST_DATA_DIR + '/transcripts_1000_kmc_counters.kmc_suf',
                 suffix=suffix
@@ -337,6 +336,77 @@ 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]
+        #TODO: remove -v
+        construct_command = '{exe} build -v --phase 1 --mask-dummy --graph {repr} --canonical -k 20 ' \
+                            '--disk-swap {tmp_dir} -o {outfile} {input}'.format(
+            exe=METAGRAPH,
+            repr=representation,
+            tmp_dir=self.tempdir.name,
+            outfile=self.tempdir.name + '/graph',
+            input=TEST_DATA_DIR + '/transcripts_1000.fa'
+        )
+        print(f'Executing phase1: {construct_command}')
+        res = subprocess.run([construct_command], shell=True)
+        self.assertEqual(res.returncode, 0)
+        self.assertTrue(os.path.isfile(self.tempdir.name + '/graph.checkpoint'))
+
+        construct_command = construct_command.replace('--phase 1', '--phase 2')
+        print(f'Executing phase2: {construct_command}')
+        res = subprocess.run([construct_command], shell=True)
+        self.assertEqual(res.returncode, 0)
+        self.assertTrue(os.path.isfile(self.tempdir.name + '/graph.checkpoint'))
+
+        construct_command = construct_command.replace('--phase 2', '')
+        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 2 --mask-dummy --graph {repr} --canonical -k 20 ' \
+                                '--disk-swap {tmp_dir} -o {outfile} {input}'.format(
+                exe=METAGRAPH,
+                repr=representation,
+                tmp_dir=self.tempdir.name,
+                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 ('graph1', '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=self.tempdir.name,
+                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 + '/' + name + 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,
@@ -116,12 +120,28 @@ int build_graph(Config *config) {
                                         : kmer::ContainerType::VECTOR_DISK,
                 config->tmp_dir.empty() ? std::filesystem::path(config->outfbase).remove_filename()
                                         : 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)");
+            }
 
+            // need to call build_chunk() even if checkpoint.phase()==1, because the
+            // SortedSetDisk needs to be flushed
             boss::BOSS::Chunk *next_chunk = constructor->build_chunk();
+
+            if (checkpoint.phase() <= 2) { // phase 2 stops after generating dummy k-mers
+                assert(next_chunk == nullptr);
+                logger->info("Phase {} (checkpoint {}) successfully finished.",
+                             checkpoint.phase(),
+                             checkpoint.checkpoint_for_phase(checkpoint.phase()));
+                return 0;
+            }
+
             logger->trace("Graph chunk with {} k-mers was built in {} sec",
                           next_chunk->size() - 1, timer.elapsed());
 
@@ -139,6 +159,7 @@ int build_graph(Config *config) {
             } else {
                 graph_data.reset(next_chunk);
             }
+            checkpoint.remove_checkpoint();
         }
 
         assert(graph_data);

metagraph/src/cli/config/config.cpp

@@ -346,6 +346,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 (!strcmp(argv[i], "--anchors-file")) {
             anchors = get_value(i++);
         } else if (argv[i][0] == '-') {
@@ -541,6 +543,30 @@ Config::Config(int argc, char *argv[]) {
     if (identity == COMPARE && fnames.size() != 2)
         print_usage_and_exit = true;
 
+    if (identity != BUILD && phase != 3) {
+        std::cerr << "Error: Phases are only supported for building. Remove --phase"
+                  << std::endl;
+        print_usage_and_exit = true;
+    }
+
+    if (identity != BUILD && phase > 3) {
+        std::cerr << "Error: Invalid phase value. Can be either 1, 2 or 3." << std::endl;
+        print_usage_and_exit = true;
+    }
+
+    if (phase != 3 && 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 != 3 && 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;
 
@@ -787,6 +813,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=collect kmers, 2=generate kmers, 3=build all) [3]\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

@@ -148,6 +148,8 @@ class Config {
 
     size_t disk_cap_bytes = 20e9; // 20GB default
 
+    uint32_t phase = 3; // 1 = collect kmers, 2 = generate kmers, 3 = 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>
@@ -212,6 +213,10 @@ EliasFanoEncoder<T>::~EliasFanoEncoder() {
 template <typename T>
 void EliasFanoEncoder<T>::add(T value) {
 #ifndef NDEBUG
+    if (value < last_value_) {
+        logger->trace("Ooops {} {}", last_value_ , value);
+        logger->info("Ooops {} {}", last_value_ , value);
+    }
     assert(value >= last_value_);
 #endif
     value -= offset_;

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.