#Centipede Distiller: Use `ResourcePool` to prevent OOMs Each shard thread blocks until enough free RAM becomes available from other reading threads before proceeding to read/write the shard. PiperOrigin-RevId: 604460185
diff --git a/centipede/BUILD b/centipede/BUILD index afb39ae..2aa9c97 100644 --- a/centipede/BUILD +++ b/centipede/BUILD
@@ -912,7 +912,10 @@ ":feature", ":feature_set", ":logging", + ":remote_file", + ":resource_pool", ":rusage_profiler", + ":rusage_stats", ":shard_reader", ":thread_pool", ":util", @@ -1344,6 +1347,8 @@ ":distill", ":environment", ":feature", + ":resource_pool", + ":rusage_stats", ":shard_reader", ":test_util", ":util",
diff --git a/centipede/distill.cc b/centipede/distill.cc index 05d6412..4c6a8e3 100644 --- a/centipede/distill.cc +++ b/centipede/distill.cc
@@ -40,7 +40,10 @@ #include "./centipede/feature.h" #include "./centipede/feature_set.h" #include "./centipede/logging.h" +#include "./centipede/remote_file.h" +#include "./centipede/resource_pool.h" #include "./centipede/rusage_profiler.h" +#include "./centipede/rusage_stats.h" #include "./centipede/shard_reader.h" #include "./centipede/thread_pool.h" #include "./centipede/util.h" @@ -70,10 +73,11 @@ using CorpusEltVec = std::vector<CorpusElt>; +inline constexpr perf::MemSize kGB = 1024L * 1024L * 1024L; + // The maximum number of threads reading input shards concurrently. This is // mainly to prevent I/O congestion. -// TODO(ussuri): Bump up significantly when RSS-gated mutexing is in. -inline constexpr size_t kMaxReadingThreads = 1; +inline constexpr size_t kMaxReadingThreads = 100; std::string LogPrefix(const Environment &env) { return absl::StrCat("DISTILL[S.", env.my_shard_index, "]: "); @@ -88,6 +92,22 @@ InputCorpusShardReader(const Environment &env) : workdir_{env}, log_prefix_{LogPrefix(env)} {} + perf::MemSize EstimateRamFootprint(size_t shard_idx) const { + const auto corpus_path = workdir_.CorpusFiles().ShardPath(shard_idx); + const auto features_path = workdir_.FeaturesFiles().ShardPath(shard_idx); + const perf::MemSize corpus_file_size = RemoteFileGetSize(corpus_path); + const perf::MemSize features_file_size = RemoteFileGetSize(features_path); + // Conservative compression factors for the two file types. These have been + // observed empirically for the Riegeli blob format. The legacy format is + // approximately 1:1, but use the stricter Riegeli numbers, as the legacy + // should be considered obsolete. + // TODO(b/322880269): Use the actual in-memory footprint once available. + constexpr double kMaxCorpusCompressionRatio = 5.0; + constexpr double kMaxFeaturesCompressionRatio = 10.0; + return corpus_file_size * kMaxCorpusCompressionRatio + + features_file_size * kMaxFeaturesCompressionRatio; + } + // Reads and returns a single shard's elements. Thread-safe. CorpusEltVec ReadShard(size_t shard_idx) { const auto corpus_path = workdir_.CorpusFiles().ShardPath(shard_idx); @@ -233,7 +253,9 @@ } // namespace void DistillTask(const Environment &env, - const std::vector<size_t> &shard_indices) { + const std::vector<size_t> &shard_indices, + perf::ResourcePool<perf::RUsageMemory> &ram_pool, + int parallelism) { // Read and write the shards in parallel, but gate reading of each on the // availability of free RAM to keep the peak RAM usage under control. const size_t num_shards = shard_indices.size(); @@ -242,9 +264,17 @@ DistilledCorpusShardWriter writer{env, /*append=*/false}; { - ThreadPool threads{kMaxReadingThreads}; + ThreadPool threads{parallelism}; for (size_t shard_idx : shard_indices) { - threads.Schedule([shard_idx, &reader, &writer, &env, num_shards] { + threads.Schedule([shard_idx, &reader, &writer, &env, num_shards, + &ram_pool] { + const auto ram_lease = ram_pool.AcquireLeaseBlocking({ + .id = absl::StrCat("out_", env.my_shard_index, "/in_", shard_idx), + .amount = {.mem_rss = reader.EstimateRamFootprint(shard_idx)}, + .timeout = absl::Minutes(30), + }); + CHECK_OK(ram_lease.status()); + CorpusEltVec shard_elts = reader.ReadShard(shard_idx); // Reverse the order of elements. The intuition is as follows: // * If the shard is the result of fuzzing with Centipede, the inputs @@ -270,6 +300,10 @@ /*timelapse_interval=*/absl::Seconds(VLOG_IS_ON(2) ? 10 : 60), // /*also_log_timelapses=*/VLOG_IS_ON(10)); + // The RAM pool shared between all the threads, here and in `DistillTask`. + constexpr perf::RUsageMemory kRamQuota{.mem_rss = 25 * kGB}; + perf::ResourcePool ram_pool{kRamQuota}; + // Run `env.num_threads` independent distillation threads. std::vector<std::thread> threads(env.num_threads); std::vector<Environment> envs(env.num_threads, env); @@ -285,7 +319,8 @@ std::shuffle(shard_indices.begin(), shard_indices.end(), rng); // Run the thread. threads[thread_idx] = - std::thread(DistillTask, std::ref(envs[thread_idx]), shard_indices); + std::thread(DistillTask, std::ref(envs[thread_idx]), shard_indices, + std::ref(ram_pool), kMaxReadingThreads); } // Join threads. for (size_t thread_idx = 0; thread_idx < env.num_threads; thread_idx++) {
diff --git a/centipede/distill.h b/centipede/distill.h index e0bdadb..604703f 100644 --- a/centipede/distill.h +++ b/centipede/distill.h
@@ -19,6 +19,8 @@ #include <vector> #include "./centipede/environment.h" +#include "./centipede/resource_pool.h" +#include "./centipede/rusage_stats.h" namespace centipede { @@ -26,8 +28,13 @@ // by `shard_indices`, distills inputs from them and writes the result to // `WorkDir{env}.DistilledPath()`. Every task gets its own `env.my_shard_index`, // and so every task creates its own independent distilled corpus file. +// `parallelism` is the maximum number of concurrent reading/writing threads. +// Values > 1 can cause non-determinism in which of the same-coverage inputs +// get selected to be written to the output shard; set to 1 for tests. void DistillTask(const Environment &env, - const std::vector<size_t> &shard_indices); + const std::vector<size_t> &shard_indices, + perf::ResourcePool<perf::RUsageMemory> &ram_pool, + int parallelism = 100); // Runs `env.num_threads` independent distill tasks in separate threads. // Returns EXIT_SUCCESS.
diff --git a/centipede/distill_test.cc b/centipede/distill_test.cc index 328caa6..90f6f4f 100644 --- a/centipede/distill_test.cc +++ b/centipede/distill_test.cc
@@ -29,6 +29,8 @@ #include "./centipede/defs.h" #include "./centipede/environment.h" #include "./centipede/feature.h" +#include "./centipede/resource_pool.h" +#include "./centipede/rusage_stats.h" #include "./centipede/shard_reader.h" #include "./centipede/test_util.h" #include "./centipede/util.h" @@ -37,6 +39,8 @@ namespace centipede { namespace { +using testing::UnorderedElementsAreArray; + struct TestCorpusRecord { ByteArray input; FeatureVec feature_vec; @@ -113,6 +117,11 @@ const WorkDir wd{env}; std::filesystem::create_directories(wd.CoverageDirPath()); + // Do not limit the max RAM. + perf::ResourcePool ram_pool{perf::RUsageMemory::Max()}; + // Turn off parallel writes to ensure deterministic outputs. + constexpr int kParallelism = 1; + // Write the shards. for (size_t shard_index = 0; shard_index < shards.size(); ++shard_index) { for (const auto &record : shards[shard_index]) { @@ -120,7 +129,7 @@ } } // Distill. - DistillTask(env, shard_indices); + DistillTask(env, shard_indices, ram_pool, kParallelism); // Read the result back. return ReadFromDistilled(wd); } @@ -135,45 +144,50 @@ ShardVec shards = { // shard 0; note: distillation iterates the shards backwards. - {{in3, {10}}, {in0, {10, 20}}}, + { + {in3, {10}}, + {in0, {10, 20}}, + }, // shard 1 - {{in1, {20, 30, usr0}}}, + { + {in1, {20, 30, usr0}}, + }, // shard 2 - {{in2, {30, 40, usr1}}}, + { + {in2, {30, 40, usr1}}, + }, }; // Distill these 3 shards in different orders, observe different results. EXPECT_THAT(TestDistill(shards, {0, 1, 2}, test_info_->name(), 0), - testing::ElementsAreArray({ + UnorderedElementsAreArray({ EqualsTestCorpusRecord(in0, FeatureVec{10, 20}), EqualsTestCorpusRecord(in1, FeatureVec{20, 30}), EqualsTestCorpusRecord(in2, FeatureVec{30, 40}), })); - EXPECT_THAT(TestDistill(shards, {2, 0, 1}, test_info_->name(), 0), - testing::ElementsAreArray({ + UnorderedElementsAreArray({ EqualsTestCorpusRecord(in2, FeatureVec{30, 40}), EqualsTestCorpusRecord(in0, FeatureVec{10, 20}), })); EXPECT_THAT(TestDistill(shards, {2, 0, 1}, test_info_->name(), 0x1), - testing::ElementsAreArray({ + UnorderedElementsAreArray({ EqualsTestCorpusRecord(in2, FeatureVec{30, 40}), EqualsTestCorpusRecord(in0, FeatureVec{10, 20}), EqualsTestCorpusRecord(in1, FeatureVec{20, 30, usr0}), })); EXPECT_THAT(TestDistill(shards, {2, 0, 1}, test_info_->name(), 0x2), - testing::ElementsAreArray({ + UnorderedElementsAreArray({ EqualsTestCorpusRecord(in2, FeatureVec{30, 40, usr1}), EqualsTestCorpusRecord(in0, FeatureVec{10, 20}), })); EXPECT_THAT(TestDistill(shards, {2, 0, 1}, test_info_->name(), 0x3), - testing::ElementsAreArray({ + UnorderedElementsAreArray({ EqualsTestCorpusRecord(in2, FeatureVec{30, 40, usr1}), EqualsTestCorpusRecord(in0, FeatureVec{10, 20}), EqualsTestCorpusRecord(in1, FeatureVec{20, 30, usr0}), })); - EXPECT_THAT(TestDistill(shards, {1, 0, 2}, test_info_->name(), 0), - testing::ElementsAreArray({ + UnorderedElementsAreArray({ EqualsTestCorpusRecord(in1, FeatureVec{20, 30}), EqualsTestCorpusRecord(in0, FeatureVec{10, 20}), EqualsTestCorpusRecord(in2, FeatureVec{30, 40}),