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