#Centipede Seeder: More RAM-efficient source element sampling

Avoid a spike in peak RSS & VSize usage and speed up source shards processing.

The previous version std::sampled the elements themselves, inserting copies into the destination.

The new version std::shuffles and resizes a source's elements' indices to achieve the same effect of sampling, then moves selected elements from the source's vector to the destination vector.

Example rusage profile for one large shard shows RSS reduction of 30%, from 36GB to 24GB, and sampling time decrease from 40 seconds to 2.

BEFORE:

```
RSS:

seed_corpus_maker_lib.cc:254 @ 14:49:18.95 [P.0:S.80 Done reading ]     25.22G [#################################-----------------]
seed_corpus_maker_lib.cc:271 @ 14:49:20.07 [P.0:S.81 Done merging ]     25.22G [#################################-----------------]
seed_corpus_maker_lib.cc:154 @ 14:49:24.50 [P.0:S.82 Timelapse    ]     26.22G [###################################---------------]
seed_corpus_maker_lib.cc:154 @ 14:49:34.50 [P.0:S.83 Timelapse    ]     29.22G [#######################################-----------]
seed_corpus_maker_lib.cc:154 @ 14:49:44.50 [P.0:S.84 Timelapse    ]     32.22G [############################################------]
seed_corpus_maker_lib.cc:154 @ 14:49:54.50 [P.0:S.85 Timelapse    ]     35.22G [################################################--]
seed_corpus_maker_lib.cc:307 @ 14:49:59.00 [P.0:S.86 Done sampling]     36.22G [##################################################]
seed_corpus_maker_lib.cc:154 @ 14:50:04.50 [P.0:S.87 Timelapse    ]     36.22G [##################################################]
seed_corpus_maker_lib.cc:154 @ 14:50:14.50 [P.0:S.88 Timelapse    ]     36.22G [##################################################]
seed_corpus_maker_lib.cc:154 @ 14:50:24.50 [P.0:S.89 Timelapse    ]     36.22G [##################################################]
seed_corpus_maker_lib.cc:154 @ 14:50:34.50 [P.0:S.90 Timelapse    ]     36.22G [##################################################]
seed_corpus_maker_lib.cc:154 @ 14:50:44.50 [P.0:S.91 FINAL        ]     36.22G [##################################################]
```

AFTER:

```
RSS:

seed_corpus_maker_lib.cc:256 @ 21:05:28.54 [P.0:S.74 Done reading  ]     24.07G [#################################################-]
seed_corpus_maker_lib.cc:276 @ 21:05:29.61 [P.0:S.75 Done merging  ]     24.26G [#################################################-]
seed_corpus_maker_lib.cc:316 @ 21:05:31.14 [P.0:S.76 Done sampling ]     24.26G [#################################################-]
seed_corpus_maker_lib.cc:324 @ 21:05:31.57 [P.0:S.77 Done appending]     24.26G [#################################################-]
seed_corpus_maker_lib.cc:156 @ 21:05:33.96 [P.0:S.78 Timelapse     ]     24.26G [#################################################-]
seed_corpus_maker_lib.cc:156 @ 21:05:43.96 [P.0:S.79 Timelapse     ]     24.26G [#################################################-]
seed_corpus_maker_lib.cc:156 @ 21:05:53.96 [P.0:S.80 FINAL         ]     24.26G [##################################################]
```

PiperOrigin-RevId: 592309603

centipede/seed_corpus_maker_lib.cc

@@ -29,6 +29,7 @@
 #include <functional>
 #include <iterator>
 #include <memory>
+#include <numeric>
 #include <string>
 #include <string_view>
 #include <utility>
@@ -260,11 +261,14 @@ void SampleSeedCorpusElementsFromSource(    //
       src_elts.emplace_back(std::move(elt));
     }
     shard_elts.clear();
+    shard_elts.shrink_to_fit();
     src_num_features += src_elts_with_features_per_shard[s];
   }
 
   src_elts_per_shard.clear();
+  src_elts_per_shard.shrink_to_fit();
   src_elts_with_features_per_shard.clear();
+  src_elts_with_features_per_shard.shrink_to_fit();
 
   RPROF_SNAPSHOT_AND_LOG("Done merging");
 
@@ -293,16 +297,28 @@ void SampleSeedCorpusElementsFromSource(    //
   if (sample_size < src_elts.size()) {
     LOG(INFO) << "Sampling " << sample_size << " elements out of "
               << src_elts.size();
-    std::sample(  //
-        src_elts.cbegin(), src_elts.cend(), std::back_inserter(elements),
-        sample_size, absl::BitGen{});
   } else {
     LOG(INFO) << "Using all " << src_elts.size() << " elements";
-    // TODO(ussuri): Should we still use std::sample() to randomize the order?
-    elements.insert(elements.end(), src_elts.cbegin(), src_elts.cend());
   }
 
+  // Extract a sample by shuffling the elements' indices and resizing to the
+  // requested sample size. We do this, rather than std::sampling the elements
+  // themselves and associated inserting into `elements`, to avoid a spike in
+  // peak RAM usage.
+  std::vector<size_t> src_sample_idxs(src_elts.size());
+  std::iota(src_sample_idxs.begin(), src_sample_idxs.end(), 0);
+  std::shuffle(src_sample_idxs.begin(), src_sample_idxs.end(), absl::BitGen{});
+  src_sample_idxs.resize(sample_size);
+
   RPROF_SNAPSHOT_AND_LOG("Done sampling");
+
+  // Now move each sampled element from `src_elts` to `elements`.
+  elements.reserve(elements.size() + sample_size);
+  for (size_t idx : src_sample_idxs) {
+    elements.emplace_back(std::move(src_elts[idx]));
+  }
+
+  RPROF_SNAPSHOT_AND_LOG("Done appending");
 }
 
 // TODO(ussuri): Refactor into smaller functions.

centipede/seed_corpus_maker_lib_test.cc

@@ -16,6 +16,7 @@
 
 #include <unistd.h>
 
+#include <cmath>
 #include <cstddef>
 #include <filesystem>  // NOLINT
 #include <string>
@@ -38,7 +39,6 @@ namespace {
 namespace fs = std::filesystem;
 using google::protobuf::TextFormat;
 using testing::IsSubsetOf;
-using testing::UnorderedElementsAreArray;
 
 inline constexpr auto kIdxDigits = WorkDir::kDigitsInShardIndex;
 
@@ -189,19 +189,16 @@ TEST(SeedCorpusMakerLibTest, RoundTripWriteReadWrite) {
       }
     )pb";
 
-    for (const double fraction : {1.0, 0.5}) {
+    for (const double fraction : {1.0, 0.5, 0.2}) {
       const SeedCorpusConfig config = ParseSeedCorpusConfig(
           absl::Substitute(kConfigStr, kRelDir1, kCovBin, fraction));
       InputAndFeaturesVec elements;
       SampleSeedCorpusElementsFromSource(  //
           config.sources(0), kCovBin, kCovHash, elements);
       // NOTE: 1.0 has a precise double representation, so `==` is fine.
-      if (fraction == 1.0) {
-        ASSERT_THAT(elements, UnorderedElementsAreArray(kElements))
-            << VV(fraction);
-      } else {
-        ASSERT_THAT(elements, IsSubsetOf(kElements)) << VV(fraction);
-      }
+      ASSERT_EQ(elements.size(), std::llrint(kElements.size() * fraction))
+          << VV(fraction);
+      ASSERT_THAT(elements, IsSubsetOf(kElements)) << VV(fraction);
     }
   }