From 2963bf51fc2c100cae1aef84b4dd242c7548519c Mon Sep 17 00:00:00 2001 From: Kevin Zhang Date: Fri, 7 Aug 2026 14:29:48 -0700 Subject: [PATCH] Add filterGraphsByFormatVersion training utility Summary: Adds a shared training utility that filters candidate graphs by compressing caller-provided MultiInput samples at a target format version. Callers are responsible for supplying inputs that exercise the graph paths whose compatibility must be tested. Also ensures format version is set inside managed compression base graph constructors. Reviewed By: terrelln Differential Revision: D114119859 --- cli/args/TrainArgs.h | 14 ++ tools/training/tests/BUCK | 1 + .../tests/test_clustering_benchmarks.cpp | 1 + tools/training/tests/test_dict_training.cpp | 1 + tools/training/tests/test_utils.cpp | 144 ++++++++++++++++++ tools/training/utils/utils.cpp | 69 ++++++++- tools/training/utils/utils.h | 52 ++++++- 7 files changed, 278 insertions(+), 4 deletions(-) create mode 100644 tools/training/tests/test_utils.cpp diff --git a/cli/args/TrainArgs.h b/cli/args/TrainArgs.h index 063d5c376..247ec3aa1 100644 --- a/cli/args/TrainArgs.h +++ b/cli/args/TrainArgs.h @@ -7,6 +7,7 @@ #include "custom_parsers/dependency_registration.h" #include "openzl/cpp/Compressor.hpp" +#include "openzl/zl_version.h" #include "tools/io/InputSetBuilder.h" #include "tools/io/OutputFile.h" @@ -133,6 +134,7 @@ class TrainArgs : public GlobalArgs, public ProfileArgs { // Create the compressor setCompressor(createCompressorFromArgs( *this, parsed.cmdFlag(cmd(), kCompressor))); + applyDefaultFormatVersion(); auto outputPath = parsed.cmdFlag(cmd(), kOutput); if (outputPath) { checkOutput(outputPath.value(), parsed.cmdHasFlag(cmd(), kForce)); @@ -228,6 +230,7 @@ class TrainArgs : public GlobalArgs, public ProfileArgs { // Inline training (e.g. `compress --train-inline`) produces a // standalone compressor only; dictionary training is opt-in via // --dict-bundle-output. + applyDefaultFormatVersion(); trainParams.dictTraining = false; trainParams.compressorGenFunc = custom_parsers::createCompressorFromSerialized; @@ -246,6 +249,17 @@ class TrainArgs : public GlobalArgs, public ProfileArgs { training::TrainParams trainParams; private: + // The trained (and serialized) compressor must carry a format version so + // that downstream training can target it. Default to the maximum supported + // version when the compressor does not already specify one. + void applyDefaultFormatVersion() + { + if (compressor()->getParameter(CParam::FormatVersion) == 0) { + compressor()->setParameter( + CParam::FormatVersion, ZL_MAX_FORMAT_VERSION); + } + } + inline static const std::string kSampleDir = "sample-dir"; inline static const std::string kCompressor = "compressor"; diff --git a/tools/training/tests/BUCK b/tools/training/tests/BUCK index 69c98cf77..da0b69709 100644 --- a/tools/training/tests/BUCK +++ b/tools/training/tests/BUCK @@ -21,6 +21,7 @@ cpp_unittest( "test_sample_limiter.cpp", "test_thread_pool.cpp", "test_train.cpp", + "test_utils.cpp", ], headers = relative_headers([ "benchmark_files/ppmf_unit_segment.h", diff --git a/tools/training/tests/test_clustering_benchmarks.cpp b/tools/training/tests/test_clustering_benchmarks.cpp index 457bde514..88eb6dae2 100644 --- a/tools/training/tests/test_clustering_benchmarks.cpp +++ b/tools/training/tests/test_clustering_benchmarks.cpp @@ -26,6 +26,7 @@ class TestClusteringBenchmarks : public testing::Test { { // Register the graph to train in the compressor trainingGraphFn(compressor.get()); + compressor.setParameter(CParam::FormatVersion, ZL_MAX_FORMAT_VERSION); // Train the compressor and serialize it auto serialized = training::train(inputs_, compressor, params_); // Compress the data using the trained compressor diff --git a/tools/training/tests/test_dict_training.cpp b/tools/training/tests/test_dict_training.cpp index eb6c91dcd..0813e128c 100644 --- a/tools/training/tests/test_dict_training.cpp +++ b/tools/training/tests/test_dict_training.cpp @@ -107,6 +107,7 @@ TEST(BaseDictTrainer, DuplicateDictsAreDeduped) // Generate a compressor that splits an input into 3, and sends each to a // trainable zstd node Compressor compressor; + compressor.setParameter(CParam::FormatVersion, ZL_MAX_FORMAT_VERSION); { constexpr size_t kNumSegments = 3; constexpr size_t kSegmentSizes[kNumSegments] = { 1024, 1024, 0 }; diff --git a/tools/training/tests/test_utils.cpp b/tools/training/tests/test_utils.cpp new file mode 100644 index 000000000..adfc92735 --- /dev/null +++ b/tools/training/tests/test_utils.cpp @@ -0,0 +1,144 @@ +// Copyright (c) Meta Platforms, Inc. and affiliates. + +#include + +#include +#include + +#include "openzl/codecs/zl_concat.h" +#include "openzl/codecs/zl_conversion.h" +#include "openzl/codecs/zl_lz.h" +#include "openzl/codecs/zl_store.h" +#include "openzl/codecs/zl_zstd.h" +#include "openzl/cpp/Compressor.hpp" +#include "openzl/zl_version.h" +#include "tools/training/utils/utils.h" + +namespace openzl::training { +namespace { + +std::vector graphIds(const std::vector& graphs) +{ + std::vector ids; + ids.reserve(graphs.size()); + for (const auto graph : graphs) { + ids.push_back(graph.gid); + } + return ids; +} + +std::vector serialInputs() +{ + static const std::array data = [] { + std::array result{}; + for (size_t i = 0; i < result.size(); ++i) { + result[i] = static_cast(i); + } + return result; + }(); + MultiInput input; + input.add(Input::refSerial(data.data(), data.size())); + return { std::move(input) }; +} + +TEST(CompressorIsFormatCompatibleTest, UsesCompressorFormatVersion) +{ + Compressor compressor; + compressor.selectStartingGraph(ZL_GRAPH_LZ); + compressor.setParameter(CParam::FormatVersion, 23); + EXPECT_FALSE(compressorIsFormatCompatible(compressor, serialInputs())); + + compressor.setParameter(CParam::FormatVersion, 24); + EXPECT_TRUE(compressorIsFormatCompatible(compressor, serialInputs())); +} + +TEST(FilterGraphsByFormatVersionTest, ThrowsWhenFormatVersionBelowMinimum) +{ + Compressor compressor; + const auto customGraph = compressor.buildStaticGraph( + ZL_NODE_CONVERT_STRUCT_TO_SERIAL, { ZL_GRAPH_STORE }); + compressor.selectStartingGraph(ZL_GRAPH_STORE); + const std::vector graphs = { ZL_GRAPH_STORE, + ZL_GRAPH_ZSTD, + customGraph }; + + // Leaving the format version unset reads back as 0, which is below + // ZL_MIN_FORMAT_VERSION. setParameter rejects any explicit value below the + // minimum, so an unset compressor is the way to exercise the guard. + EXPECT_THROW( + filterGraphsByFormatVersion(compressor, graphs, serialInputs()), + Exception); +} + +TEST(FilterGraphsByFormatVersionTest, FiltersLZByVersion) +{ + Compressor compressor; + compressor.setParameter(CParam::FormatVersion, 23); + const auto beforeVersion24 = filterGraphsByFormatVersion( + compressor, { ZL_GRAPH_LZ }, serialInputs()); + EXPECT_TRUE(beforeVersion24.empty()); + + compressor.setParameter(CParam::FormatVersion, 24); + const auto atVersion24 = filterGraphsByFormatVersion( + compressor, { ZL_GRAPH_LZ }, serialInputs()); + EXPECT_EQ(graphIds(atVersion24), graphIds({ ZL_GRAPH_LZ })); + EXPECT_EQ(compressor.getParameter(CParam::FormatVersion), 24); +} + +TEST(FilterGraphsByFormatVersionTest, RestoresCompressorStateAfterException) +{ + Compressor compressor; + compressor.selectStartingGraph(ZL_GRAPH_STORE); + compressor.setParameter(CParam::FormatVersion, ZL_MAX_FORMAT_VERSION); + + EXPECT_THROW( + filterGraphsByFormatVersion( + compressor, { ZL_GRAPH_ILLEGAL }, serialInputs()), + Exception); + + EXPECT_EQ( + compressor.getParameter(CParam::FormatVersion), + ZL_MAX_FORMAT_VERSION); + EXPECT_EQ(compressor.getStartingGraph(), ZL_GRAPH_STORE); +} + +TEST(FilterGraphsByFormatVersionTest, FiltersCustomGraphByConversionVersion) +{ + Compressor compressor; + const auto graph = compressor.buildStaticGraph( + ZL_NODE_CONVERT_STRUCT_TO_NUM_BE, { ZL_GRAPH_STORE }); + const std::array data{}; + MultiInput input; + input.add(Input::refStruct(data.data(), data.size())); + const std::vector inputs = { std::move(input) }; + + compressor.setParameter(CParam::FormatVersion, 20); + const auto beforeVersion21 = + filterGraphsByFormatVersion(compressor, { graph }, inputs); + EXPECT_TRUE(beforeVersion21.empty()); + + compressor.setParameter(CParam::FormatVersion, 21); + const auto atVersion21 = + filterGraphsByFormatVersion(compressor, { graph }, inputs); + EXPECT_EQ(graphIds(atVersion21), graphIds({ graph })); +} + +TEST(FilterGraphsByFormatVersionTest, SupportsMultiInputGraphs) +{ + Compressor compressor; + compressor.setParameter(CParam::FormatVersion, ZL_MAX_FORMAT_VERSION); + const auto graph = compressor.buildStaticGraph( + ZL_NODE_CONCAT_SERIAL, { ZL_GRAPH_STORE, ZL_GRAPH_STORE }); + MultiInput input; + input.add(Input::refSerial("first", 5)); + input.add(Input::refSerial("second", 6)); + const std::vector inputs = { std::move(input) }; + + const auto supported = + filterGraphsByFormatVersion(compressor, { graph }, inputs); + + EXPECT_EQ(graphIds(supported), graphIds({ graph })); +} + +} // namespace +} // namespace openzl::training diff --git a/tools/training/utils/utils.cpp b/tools/training/utils/utils.cpp index a6ba588ba..5a635a47e 100644 --- a/tools/training/utils/utils.cpp +++ b/tools/training/utils/utils.cpp @@ -3,19 +3,32 @@ #include "tools/training/utils/utils.h" #include "openzl/cpp/CCtx.hpp" #include "openzl/cpp/Compressor.hpp" -#include "tools/io/InputSetStatic.h" +#include "openzl/cpp/Exception.hpp" +#include "openzl/zl_reflection.h" namespace openzl::training { CCtx refCCtxForTraining(const Compressor& compressor) { openzl::CCtx cctx; - cctx.setParameter(openzl::CParam::FormatVersion, ZL_MAX_FORMAT_VERSION); - cctx.setParameter(openzl::CParam::StickyParameters, ZL_MAX_FORMAT_VERSION); + cctx.setParameter(openzl::CParam::StickyParameters, 1); cctx.refCompressor(compressor); return cctx; } +size_t MultiInput::compressBound() const +{ + size_t totalSrcSize = 0; + for (const auto& input : *inputs_) { + totalSrcSize += input.contentSize(); + if (input.type() == Type::String) { + totalSrcSize += input.numElts() * sizeof(*input.stringLens()); + } + } + totalSrcSize += inputs_->size() * 256; + return 2 * ZL_compressBound(totalSrcSize) + 1024; +} + std::vector inputSetToMultiInputs(tools::io::InputSet& inputs) { // Convert the io inputs to MultiInputs @@ -28,4 +41,54 @@ std::vector inputSetToMultiInputs(tools::io::InputSet& inputs) return multiInputs; } +bool compressorIsFormatCompatible( + const Compressor& compressor, + const std::vector& inputs) +{ + CCtx cctx; + cctx.refCompressor(compressor); + for (const auto& input : inputs) { + const size_t outputCapacity = input.compressBound(); + std::string output(outputCapacity, '\0'); + try { + cctx.compress(output, *input); + } catch (const Exception& e) { + // Catch only format version unsupported errors. Otherwise it is + // failing compression on the input but is actually supported format + // version-wise. + if (e.code() == ZL_ErrorCode_formatVersion_unsupported + || e.code() == ZL_ErrorCode_node_versionMismatch) { + return false; + } + } + } + return true; +} + +std::vector filterGraphsByFormatVersion( + Compressor& compressor, + const std::vector& graphs, + const std::vector& inputs) +{ + const auto formatVersion = compressor.getParameter(CParam::FormatVersion); + if (formatVersion < ZL_MIN_FORMAT_VERSION) { + throw Exception("Format version is below ZL_MIN_FORMAT_VERSION"); + } + GraphID originalStartingGraph = ZL_GRAPH_ILLEGAL; + const bool hadStartingGraph = ZL_Compressor_getStartingGraphID( + compressor.get(), &originalStartingGraph); + std::vector supported; + supported.reserve(graphs.size()); + for (const auto graph : graphs) { + compressor.selectStartingGraph(graph); + if (compressorIsFormatCompatible(compressor, inputs)) { + supported.push_back(graph); + } + } + if (hadStartingGraph) { + compressor.selectStartingGraph(originalStartingGraph); + } + return supported; +} + } // namespace openzl::training diff --git a/tools/training/utils/utils.h b/tools/training/utils/utils.h index ea9c7497a..9e31760ea 100644 --- a/tools/training/utils/utils.h +++ b/tools/training/utils/utils.h @@ -11,7 +11,7 @@ namespace openzl::training { /** * @brief Create a CCtx for training the compressor. The cctx is configured * so that if training is called multiple times, the parameters will not be - * reset. + * reset. Targets ZL_MAX_FORMAT_VERSION. */ CCtx refCCtxForTraining(const Compressor& compressor); @@ -42,6 +42,12 @@ class MultiInput { return inputs_.get(); } + /** + * @brief Returns maximum compressed size after compression using these + * inputs. + */ + size_t compressBound() const; + // Adds input while not owning the buffer the input references void add(Input&& input) { @@ -67,4 +73,48 @@ class MultiInput { */ std::vector inputSetToMultiInputs(tools::io::InputSet& inputs); +/** + * @brief Returns whether @p compressor is compatible with its configured + * format version for every sample in @p inputs. + * + * It is the caller's responsibility to configure the compressor's format + * version, select its starting graph, and provide inputs that exercise every + * graph path whose compatibility must be tested. Compression errors unrelated + * to format compatibility are ignored. + */ +bool compressorIsFormatCompatible( + const Compressor& compressor, + const std::vector& inputs); + +/** + * @brief Filter @p graphs down to those able to compress @p inputs at the + * target @p formatVersion. + * + * Each candidate graph is used to compress every sample in @p inputs. A graph + * is filtered out if any compression reports a format-version incompatibility. + * Supported graphs are returned in their original order. + * + * It is the caller's responsibility to provide inputs capable of exercising + * every graph path whose format-version compatibility must be tested. A graph + * that is incompatible with @p formatVersion may be retained if @p inputs do + * not exercise the incompatible path. + * + * @throws Exception if @p formatVersion is less than ZL_MIN_FORMAT_VERSION. + * + * Standard graphs follow the guidelines which are required for this function to + * work. Custom graphs are also required to follow these guidelines. These are + * that graphs must either: + * - Always select the same nodes and may not work on older format versions. + * - Dynamically select which nodes to run, in which case they should be + * format-version aware, meaning they should never execute a codec which + * requires a format version above the library's format version. + * + * If these guidelines are not followed, the function may not correctly filter + * out the graph. + */ +std::vector filterGraphsByFormatVersion( + Compressor& compressor, + const std::vector& graphs, + const std::vector& inputs); + } // namespace openzl::training