From a1a44a8dd9d195de16889b962fb7999defce0053 Mon Sep 17 00:00:00 2001 From: kevz8 Date: Wed, 12 Aug 2026 10:17:55 -0700 Subject: [PATCH 1/4] 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. 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 From c71df8bce876381d1c1a18d19afae2f541d8605e Mon Sep 17 00:00:00 2001 From: kevz8 Date: Wed, 12 Aug 2026 11:08:39 -0700 Subject: [PATCH 2/4] Phase 2: thread format version through clustering Summary: Threads the resolved target format version through the clustering trainer: - clusterSuccessors() now filters the clustering (concat) codecs by ZL_Compressor_Node_getMinVersion against the target version and throws FormatVersionUnsupportedError if none survive (so the orchestrator can fall clustering back to zstd), and builds its round-trip CCtx at the target version. - train_cluster() resolves the version from TrainParams and passes it through Trainer::getTrainedClusteringConfig into CompressionUtils, whose per-sample benchmark CCtx now targets that version instead of ZL_MAX_FORMAT_VERSION. FormatVersionUnsupportedError moved to train_params.h (train_common) so it is visible to every trainer and the orchestrator without a BUCK dependency cycle. Differential Revision: D113974921 --- cli/args/TrainArgs.h | 23 +++++++++--- tools/training/BUCK | 1 + .../clustering/clustering_graph_trainer.cpp | 29 +++++++++++++-- tools/training/tests/test_clustering.cpp | 36 +++++++++++++++++++ tools/training/tests/test_train.cpp | 21 +++++++++++ tools/training/train.cpp | 16 +++++++++ tools/training/train.h | 8 +---- tools/training/train_exceptions.h | 25 +++++++++++++ tools/training/train_params.h | 1 + tools/training/utils/BUCK | 1 + 10 files changed, 146 insertions(+), 15 deletions(-) create mode 100644 tools/training/train_exceptions.h diff --git a/cli/args/TrainArgs.h b/cli/args/TrainArgs.h index 247ec3aa1..e38416fe1 100644 --- a/cli/args/TrainArgs.h +++ b/cli/args/TrainArgs.h @@ -126,6 +126,13 @@ class TrainArgs : public GlobalArgs, public ProfileArgs { 0, false, "Save the ACE state as a local parameter in the trained compressor."); + parser.addCommandFlag( + cmd(), + kFormatVersion, + 0, + true, + "Target format version for training. If not provided, defaults " + "to the maximum supported format version."); } explicit TrainArgs(const arg::ParsedArgs& parsed) @@ -134,7 +141,14 @@ class TrainArgs : public GlobalArgs, public ProfileArgs { // Create the compressor setCompressor(createCompressorFromArgs( *this, parsed.cmdFlag(cmd(), kCompressor))); - applyDefaultFormatVersion(); + auto formatVersion = parsed.cmdFlag(cmd(), kFormatVersion); + if (formatVersion) { + compressor()->setParameter( + CParam::FormatVersion, + util::checkedstoi(formatVersion.value())); + } else { + applyDefaultFormatVersion(); + } auto outputPath = parsed.cmdFlag(cmd(), kOutput); if (outputPath) { checkOutput(outputPath.value(), parsed.cmdHasFlag(cmd(), kForce)); @@ -254,10 +268,8 @@ class TrainArgs : public GlobalArgs, public ProfileArgs { // 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); - } + compressor()->setParameter( + CParam::FormatVersion, ZL_MAX_FORMAT_VERSION); } inline static const std::string kSampleDir = "sample-dir"; @@ -279,6 +291,7 @@ class TrainArgs : public GlobalArgs, public ProfileArgs { inline static const std::string kMaxTotalSizeMb = "max-total-size-mb"; inline static const std::string kParetoFrontier = "pareto-frontier"; inline static const std::string kSaveAceState = "save-ace-state"; + inline static const std::string kFormatVersion = "format-version"; }; } // namespace openzl::cli diff --git a/tools/training/BUCK b/tools/training/BUCK index 597f2ff54..e3c67b795 100644 --- a/tools/training/BUCK +++ b/tools/training/BUCK @@ -34,6 +34,7 @@ zs_cxxlibrary( "trained_candidate.cpp", ], headers = [ + "train_exceptions.h", "train_params.h", "trained_candidate.h", ], diff --git a/tools/training/clustering/clustering_graph_trainer.cpp b/tools/training/clustering/clustering_graph_trainer.cpp index a80b1435d..4fdc68ac8 100644 --- a/tools/training/clustering/clustering_graph_trainer.cpp +++ b/tools/training/clustering/clustering_graph_trainer.cpp @@ -15,6 +15,7 @@ #include "tools/training/clustering/train_api.h" #include "tools/training/graph_mutation/graph_mutation_utils.h" #include "tools/training/sample_collection/training_sample_collector.h" +#include "tools/training/train_exceptions.h" #include "tools/training/utils/serialized_compressor_internal.h" #include "tools/training/utils/utils.h" @@ -27,6 +28,10 @@ const std::string CLUSTERING_GRAPH_NAME = "zl.cluster"; namespace { +// Minimum format version at which all concat codecs required by clustering are +// available (CONCAT_SERIAL=16, _NUMERIC/_STRUCT=17, _STRING=18). +constexpr int kMinClusteringFormatVersion = 18; + /** * Add a new parameterized version of the clustering graph to the compressor * which has ACE successors instead of the original successors. @@ -69,15 +74,33 @@ ZL_GraphID clusterSuccessors( const TrainParams& trainParams, GraphID clusteringGraphUniqueIDUntrained) { + const auto formatVersion = compressor.getParameter(CParam::FormatVersion); + // Clustering relies on concat codecs; if any required clustering codec is + // missing at this format version, training is unsupported. + if (formatVersion < kMinClusteringFormatVersion) { + throw FormatVersionUnsupportedError( + "clustering requires concat codecs unavailable at format version " + + std::to_string(formatVersion)); + } + auto cctx = refCCtxForTraining(compressor); const std::string clusteringGraphUniqueNameUntrained = ZL_Compressor_Graph_getName( compressor.get(), clusteringGraphUniqueIDUntrained); - // Get successors for training - const auto successorsVec = - getCustomGraphs(compressor, clusteringGraphUniqueIDUntrained); + // Get successors for training, dropping any that cannot compress at the + // target format version. + const auto successorsVec = filterGraphsByFormatVersion( + compressor, + getCustomGraphs(compressor, clusteringGraphUniqueIDUntrained), + inputs); + + if (successorsVec.size() == 0) { + throw FormatVersionUnsupportedError( + "No compatible successors chosen in clustering graph at format version " + + std::to_string(formatVersion)); + } // Get clustering codecs for training std::vector clusteringCodecs = diff --git a/tools/training/tests/test_clustering.cpp b/tools/training/tests/test_clustering.cpp index 732a4a151..9326baa6b 100644 --- a/tools/training/tests/test_clustering.cpp +++ b/tools/training/tests/test_clustering.cpp @@ -5,15 +5,18 @@ #include #include "openzl/codecs/zl_clustering.h" +#include "openzl/codecs/zl_lz.h" #include "openzl/common/a1cbor_helpers.h" #include "openzl/common/allocation.h" #include "openzl/cpp/Input.hpp" #include "openzl/zl_reflection.h" #include "src/openzl/compress/graphs/generic_clustering_graph.h" #include "tests/datagen/DataGen.h" +#include "tools/training/clustering/clustering_graph_trainer.h" #include "tools/training/clustering/train_api.h" #include "tools/training/clustering/utils.h" #include "tools/training/train.h" +#include "tools/training/train_exceptions.h" #include "tools/training/utils/utils.h" namespace openzl::tests { @@ -89,6 +92,7 @@ class TestTraining : public testing::Test { { // Register the base clustering graph as the starting graph compressor_.selectStartingGraph(ZL_GRAPH_CLUSTERING); + compressor_.setParameter(CParam::FormatVersion, ZL_MAX_FORMAT_VERSION); successors_.push_back(ZL_GRAPH_STORE); successors_.push_back(ZL_GRAPH_FIELD_LZ); @@ -388,5 +392,37 @@ TEST_F(TestTraining, TestTrainingDifferentTypeSameTag) EXPECT_TRUE(!ZL_RES_isError(r)); } +TEST(ClusteringTrainerFormatVersionTest, ThrowsWhenNoSuccessorSupportsVersion) +{ + // ZL_GRAPH_LZ is unavailable before format version 24, so targeting an + // earlier version (still above the clustering minimum) filters out every + // successor and leaves the clustering trainer with nothing to train. + constexpr uint32_t kLzFormatVersion = 24; + constexpr uint32_t kFormatVersionBeforeLz = kLzFormatVersion - 1; + + Compressor compressor; + compressor.setParameter(CParam::FormatVersion, kFormatVersionBeforeLz); + std::vector successors = { ZL_GRAPH_LZ }; + ZL_ClusteringConfig config = { .nbClusters = 0, .nbTypeDefaults = 0 }; + auto clusteringGraph = ZL_Clustering_registerGraph( + compressor.get(), &config, successors.data(), successors.size()); + compressor.selectStartingGraph(clusteringGraph); + + const training::TrainParams trainParams = { + .clusteringTrainer = training::ClusteringTrainer::Greedy, + }; + + std::vector data(1024); + for (size_t i = 0; i < data.size(); ++i) { + data[i] = static_cast(i); + } + training::MultiInput input; + input.add(Input::refSerial(data.data(), data.size())); + const std::vector inputs = { std::move(input) }; + EXPECT_THROW( + training::trainClusteringGraph(inputs, compressor, trainParams), + training::FormatVersionUnsupportedError); +} + } // namespace } // namespace openzl::tests diff --git a/tools/training/tests/test_train.cpp b/tools/training/tests/test_train.cpp index b0133e22b..98a8370f5 100644 --- a/tools/training/tests/test_train.cpp +++ b/tools/training/tests/test_train.cpp @@ -15,6 +15,7 @@ TEST(TrainTest, ThrowsTypedErrorWithoutTrainableGraph) { const std::vector inputs; Compressor compressor; + compressor.setParameter(CParam::FormatVersion, ZL_MAX_FORMAT_VERSION); compressor.selectStartingGraph(ZL_GRAPH_STORE); const training::TrainParams trainParams = { .compressorGenFunc = @@ -30,5 +31,25 @@ TEST(TrainTest, ThrowsTypedErrorWithoutTrainableGraph) training::NoTrainableGraphError); } +TEST(TrainTest, ThrowsWhenCompressorFormatVersionIsNotSet) +{ + const std::vector inputs; + Compressor compressor; + compressor.selectStartingGraph(ZL_GRAPH_STORE); + const training::TrainParams trainParams = { + .compressorGenFunc = + [](poly::string_view, poly::string_view) { + return std::make_unique(); + }, + }; + + try { + training::train(inputs, compressor, trainParams); + FAIL() << "Expected unset compressor format version to throw"; + } catch (const training::FormatVersionUnsupportedError& e) { + EXPECT_EQ(e.msg(), "Compressor format version is not set."); + } +} + } // namespace } // namespace openzl::tests diff --git a/tools/training/train.cpp b/tools/training/train.cpp index 9fdc2ae86..4e3c3346d 100644 --- a/tools/training/train.cpp +++ b/tools/training/train.cpp @@ -28,6 +28,22 @@ std::vector train( throw Exception("Compressor generator function is not set."); } + const auto formatVersion = compressor.getParameter(CParam::FormatVersion); + if (formatVersion == 0) { + throw FormatVersionUnsupportedError( + "Compressor format version is not set."); + } + + // Try compressing with the base graph to train. This is not exhaustive + // because function graphs may select different nodes for other inputs. + if (!compressorIsFormatCompatible(compressor, inputs)) { + throw FormatVersionUnsupportedError( + "Base graph failed to compress at format version " + + std::to_string(formatVersion) + + "; the format version is unsupported for the graph " + "getting trained."); + } + if (graph_mutation::hasTargetGraph(compressor, CLUSTERING_GRAPH_NAME)) { serializedTrainedCompressors.clear(); serializedTrainedCompressors.push_back( diff --git a/tools/training/train.h b/tools/training/train.h index 5bd4bc06d..be01f7a4a 100644 --- a/tools/training/train.h +++ b/tools/training/train.h @@ -3,19 +3,13 @@ #pragma once #include "openzl/cpp/Compressor.hpp" -#include "openzl/cpp/Exception.hpp" +#include "tools/training/train_exceptions.h" #include "tools/training/train_params.h" #include "tools/training/trained_candidate.h" #include "tools/training/utils/utils.h" namespace openzl::training { -/** Thrown when a compressor has no graph that can be trained. */ -class NoTrainableGraphError : public Exception { - public: - using Exception::Exception; -}; - /** * This function trains compressor graphs (clustering and/or ACE graphs). * It takes in a vector of buffer data, processes it through the training diff --git a/tools/training/train_exceptions.h b/tools/training/train_exceptions.h new file mode 100644 index 000000000..350bb7861 --- /dev/null +++ b/tools/training/train_exceptions.h @@ -0,0 +1,25 @@ +// Copyright (c) Meta Platforms, Inc. and affiliates. + +#pragma once + +#include "openzl/cpp/Exception.hpp" + +namespace openzl::training { + +/** Thrown when a compressor has no graph that can be trained. */ +class NoTrainableGraphError : public Exception { + public: + using Exception::Exception; +}; + +/** + * Thrown by a trainer when it cannot produce a graph that supports the target + * format version (e.g. all of its required codecs are below the target + * version). The orchestrator catches this to fall the trainer back to zstd. + */ +class FormatVersionUnsupportedError : public Exception { + public: + using Exception::Exception; +}; + +} // namespace openzl::training diff --git a/tools/training/train_params.h b/tools/training/train_params.h index f8264f0f0..dba63dd80 100644 --- a/tools/training/train_params.h +++ b/tools/training/train_params.h @@ -4,6 +4,7 @@ #include #include #include "openzl/cpp/Compressor.hpp" +#include "openzl/cpp/Exception.hpp" #include "openzl/cpp/poly/Optional.hpp" namespace openzl::training { diff --git a/tools/training/utils/BUCK b/tools/training/utils/BUCK index 97f36c432..d602df215 100644 --- a/tools/training/utils/BUCK +++ b/tools/training/utils/BUCK @@ -17,6 +17,7 @@ zs_cxxlibrary( "utils.h", ], exported_deps = [ + "..:train_common", "../..:io", "../..:logger", "../../..:common", From 8c3769fa6e2f2d9ea57a9597b0c1f453db436d8a Mon Sep 17 00:00:00 2001 From: kevz8 Date: Wed, 12 Aug 2026 11:18:06 -0700 Subject: [PATCH 3/4] Add format versioning to ACE training Summary: Adds format versioning to ACE training by filtering the list of graphs per format version. - Since it is useful to store both the format version and the per format version/ input type list of compatible graphs, refactor the stateless API to a ACECompressorBuilder class. This contains a cache and versions. - Additionally fix a minor bug with graph replacement Differential Revision: D113974924 --- src/openzl/compress/graphmgr.c | 6 +- .../unittest/compress/CompressorUnitTest.cpp | 11 ++++ tools/training/ace/ace.cpp | 14 +++- tools/training/ace/ace_combination.cpp | 26 ++++++-- tools/training/ace/ace_compressor.cpp | 8 ++- tools/training/ace/ace_compressor.h | 8 ++- tools/training/ace/ace_compressors.cpp | 60 +++++++++++++---- tools/training/ace/ace_compressors.h | 15 ++++- tools/training/ace/ace_crossover.h | 7 +- tools/training/ace/ace_mutate.h | 17 ++++- .../ace/automated_compressor_explorer.cpp | 17 +++-- .../ace/automated_compressor_explorer.h | 32 ++++++++-- tools/training/tests/test_ace.cpp | 64 ++++++++++++++++--- 13 files changed, 224 insertions(+), 61 deletions(-) diff --git a/src/openzl/compress/graphmgr.c b/src/openzl/compress/graphmgr.c index b0113de64..98e6a264c 100644 --- a/src/openzl/compress/graphmgr.c +++ b/src/openzl/compress/graphmgr.c @@ -709,9 +709,9 @@ static ZL_Report GM_checkInputTypesAreCompatible( graph_invalid, "Graphs have different number of inputs"); for (size_t i = 0; i < meta0.nbInputs; ++i) { - ZL_ERR_IF_EQ( - meta0.inputTypeMasks[i] & meta1.inputTypeMasks[i], - 0, + ZL_ERR_IF_NOT( + ICONV_isCompatible( + meta0.inputTypeMasks[i], meta1.inputTypeMasks[i]), graph_invalid, "Input %zu types are not compatible", i); diff --git a/tests/unittest/compress/CompressorUnitTest.cpp b/tests/unittest/compress/CompressorUnitTest.cpp index 877c5a511..f800f3d34 100644 --- a/tests/unittest/compress/CompressorUnitTest.cpp +++ b/tests/unittest/compress/CompressorUnitTest.cpp @@ -1718,6 +1718,17 @@ TEST_F(CompressorTest, OverrideBaseGraph) ZL_Compressor_getGraphType(compressor_.get(), paramGraph)); } +TEST_F(CompressorTest, OverrideBaseGraphAllowsImplicitInputConversion) +{ + auto paramGraph = makeParameterizedGraph(); + + EXPECT_ZS_VALID(ZL_Compressor_overrideBaseGraph( + compressor_.get(), paramGraph, ZL_GRAPH_ZSTD)); + EXPECT_EQ( + ZL_Compressor_Graph_getBaseGraphID(compressor_.get(), paramGraph), + ZL_GRAPH_ZSTD); +} + TEST_F(CompressorTest, OverrideBaseGraphRejectsStandardGraph) { EXPECT_ZS_ERROR(ZL_Compressor_overrideBaseGraph( diff --git a/tools/training/ace/ace.cpp b/tools/training/ace/ace.cpp index 75bd9b86e..80b284b87 100644 --- a/tools/training/ace/ace.cpp +++ b/tools/training/ace/ace.cpp @@ -32,7 +32,8 @@ std::string trainBackend( std::vector& samples, const TrainParams& trainParams, size_t graphIdx, - size_t numGraphs) + size_t numGraphs, + uint32_t formatVersion) { if (samples.empty()) { throw Exception( @@ -55,7 +56,8 @@ std::string trainBackend( ? trainParams.threads.value() : std::thread::hardware_concurrency() / 2, }; - params.maxTime = maxTime; + params.maxTime = maxTime; + params.formatVersion = formatVersion; AutomatedCompressorExplorer ace(flattened, params); for (;;) { Logger::logProgress( @@ -131,8 +133,14 @@ std::vector ACETrainer::train( "): no training samples"); continue; } + const auto formatVersion = static_cast( + compressor.getParameter(CParam::FormatVersion)); auto aceState = trainBackend( - samples[backendGraph], trainParams, graphIdx, numGraphs); + samples[backendGraph], + trainParams, + graphIdx, + numGraphs, + formatVersion); auto localParams = LocalParams(); localParams.addCopyParam( AutomatedCompressorExplorer::kAceStateParamId, diff --git a/tools/training/ace/ace_combination.cpp b/tools/training/ace/ace_combination.cpp index 534d766ed..ecec49f83 100644 --- a/tools/training/ace/ace_combination.cpp +++ b/tools/training/ace/ace_combination.cpp @@ -173,8 +173,12 @@ std::vector> benchmarkAce( auto inputs = ace.inputs(); std::vector> result; for (auto&& [candidate, _] : solutions) { - auto benchmark = *candidate.benchmark(inputs); - result.emplace_back(std::move(candidate), std::move(benchmark)); + auto benchmark = candidate.benchmark(inputs, ace.formatVersion()); + // Do not add nullopt to list of benchmark results + if (!benchmark.has_value()) { + continue; + } + result.emplace_back(std::move(candidate), std::move(*benchmark)); if (!trainParams.paretoFrontier) { break; } @@ -183,8 +187,13 @@ std::vector> benchmarkAce( Logger::log( WARNINGS, "No solution found that meets speed constraints: Falling back to store"); - auto store = buildStoreCompressor(); - return { { store, *store.benchmark(inputs) } }; + auto store = buildStoreCompressor(); + auto storeBenchmark = store.benchmark(inputs, ace.formatVersion()); + if (!storeBenchmark.has_value()) { + throw Exception( + "Store compressor failed to compress at the target format version"); + } + return { { store, std::move(*storeBenchmark) } }; } // Register the new graph on the compressor and return the new graph ID @@ -290,7 +299,9 @@ std::vector getCombinedCompressors( *trainedSerializedCompressor, "")); }; - auto compressor = makeCompressor(); + auto compressor = makeCompressor(); + const auto formatVersion = static_cast( + compressor.getParameter(CParam::FormatVersion)); auto cctx = refCCtxForTraining(compressor); auto serialized = compressor.serialize(); @@ -341,7 +352,10 @@ std::vector getCombinedCompressors( } auto aceState = std::string( (const char*)copyParam->paramPtr, copyParam->paramSize); - AutomatedCompressorExplorer ace(flattened, aceState); + AutomatedCompressorExplorer::Parameters aceParams; + // ACE snapshots contain the population, but not the run configuration. + aceParams.formatVersion = formatVersion; + AutomatedCompressorExplorer ace(flattened, aceState, aceParams); auto benchmarks = benchmarkAce(ace, trainParams); allCandidates.emplace(backendGraph, std::move(benchmarks)); } diff --git a/tools/training/ace/ace_compressor.cpp b/tools/training/ace/ace_compressor.cpp index e6453a90b..24f9c9d4b 100644 --- a/tools/training/ace/ace_compressor.cpp +++ b/tools/training/ace/ace_compressor.cpp @@ -692,7 +692,6 @@ poly::optional benchmark( auto cStart = std::chrono::steady_clock::now(); try { cctx.refCompressor(compressor); - cctx.setParameter(CParam::FormatVersion, ZL_MAX_FORMAT_VERSION); compressed = cctx.compress(input); } catch (const Exception&) { return poly::nullopt; @@ -744,11 +743,16 @@ poly::optional benchmark( } poly::optional ACECompressor::benchmark( - poly::span inputs) const + poly::span inputs, + uint32_t formatVersion) const { Compressor compressor; // TODO(terrelln): Allow parameterization compressor.selectStartingGraph(build(compressor)); + // Format version is carried on the compressor, however benchmark builds the + // compressor from an empty compressor. So the format version must be set + // here. + compressor.setParameter(CParam::FormatVersion, formatVersion); return openzl::training::benchmark(compressor, inputs); } } // namespace training diff --git a/tools/training/ace/ace_compressor.h b/tools/training/ace/ace_compressor.h index 07705cbc1..f0b410a54 100644 --- a/tools/training/ace/ace_compressor.h +++ b/tools/training/ace/ace_compressor.h @@ -18,6 +18,8 @@ struct ACENode { poly::optional params; Type inputType; std::vector outputTypes; + /// Minimum format version this node requires; 0 means unconstrained. + unsigned minFormatVersion{ 0 }; }; struct ACEGraph { @@ -271,9 +273,11 @@ class ACECompressor { } /// @returns The benchmark result of the compressor on the @p inputs or - /// poly::nullopt if the compressor fails to compress. + /// poly::nullopt if the compressor fails to compress (including when it + /// requires a newer format version than @p formatVersion). poly::optional benchmark( - poly::span inputs) const; + poly::span inputs, + uint32_t formatVersion) const; private: uint64_t computeHash() const; diff --git a/tools/training/ace/ace_compressors.cpp b/tools/training/ace/ace_compressors.cpp index a887079b3..e07d88b75 100644 --- a/tools/training/ace/ace_compressors.cpp +++ b/tools/training/ace/ace_compressors.cpp @@ -45,11 +45,15 @@ ACENode buildNode(const NodeT& node) for (const auto& meta : NodeT::metadata.variableOutputs) { outputTypes.push_back(meta.type); } + Compressor compressor; + const unsigned minFormatVersion = + ZL_Compressor_Node_getMinVersion(compressor.get(), NodeT::node); return ACENode{ - .name = getName(NodeT::node), - .params = node.parameters(), - .inputType = NodeT::metadata.inputs[0].type, - .outputTypes = std::move(outputTypes), + .name = getName(NodeT::node), + .params = node.parameters(), + .inputType = NodeT::metadata.inputs[0].type, + .outputTypes = std::move(outputTypes), + .minFormatVersion = minFormatVersion, }; } @@ -308,7 +312,12 @@ poly::span getAllGraphs() return *graphs; } -poly::span getNodesComptabileWith(Type inputType) +// Returns the type-compatible nodes further restricted to those usable at +// @p formatVersion (nodes with minFormatVersion 0 are treated as +// unconstrained). +std::vector getNodesComptabileWith( + Type inputType, + uint32_t formatVersion) { static auto nodes = new std::unordered_map>([] { std::unordered_map> m; @@ -322,7 +331,13 @@ poly::span getNodesComptabileWith(Type inputType) } return m; }()); - return nodes->at(inputType); + std::vector compatible; + for (const auto& n : nodes->at(inputType)) { + if (n.minFormatVersion == 0 || n.minFormatVersion <= formatVersion) { + compatible.push_back(n); + } + } + return compatible; } poly::span getGraphsComptabileWith(Type inputType) @@ -367,26 +382,43 @@ ACECompressor buildRandomGraphCompressor(std::mt19937_64& rng, Type inputType) randomChoice(rng, getGraphsComptabileWith(inputType))); } -ACECompressor -buildRandomNodeCompressor(std::mt19937_64& rng, Type inputType, size_t maxDepth) +ACECompressor buildRandomNodeCompressor( + std::mt19937_64& rng, + Type inputType, + uint32_t formatVersion, + size_t maxDepth) { if (maxDepth == 0) { return buildRandomGraphCompressor(rng, inputType); } - auto node = randomChoice(rng, getNodesComptabileWith(inputType)); + const auto compatible = getNodesComptabileWith(inputType, formatVersion); + if (compatible.empty()) { + // No node is available at the target format version; fall back to a + // single graph. + return buildRandomGraphCompressor(rng, inputType); + } + auto node = randomChoice( + rng, + poly::span(compatible.data(), compatible.size())); assert(isCompatible(node.inputType, inputType)); std::vector> successors; successors.reserve(node.outputTypes.size()); for (size_t i = 0; i < node.outputTypes.size(); ++i) { successors.push_back( std::make_unique(buildRandomCompressor( - rng, node.outputTypes[i], maxDepth - 1))); + rng, + node.outputTypes[i], + formatVersion, + maxDepth - 1))); } return ACENodeCompressor(std::move(node), std::move(successors)); } -ACECompressor -buildRandomCompressor(std::mt19937_64& rng, Type inputType, size_t maxDepth) +ACECompressor buildRandomCompressor( + std::mt19937_64& rng, + Type inputType, + uint32_t formatVersion, + size_t maxDepth) { std::bernoulli_distribution dist(0.5); if (dist(rng)) { @@ -394,8 +426,8 @@ buildRandomCompressor(std::mt19937_64& rng, Type inputType, size_t maxDepth) assert(compressor.acceptsInputType(inputType)); return compressor; } else { - auto compressor = - buildRandomNodeCompressor(rng, inputType, maxDepth - 1); + auto compressor = buildRandomNodeCompressor( + rng, inputType, formatVersion, maxDepth - 1); assert(compressor.acceptsInputType(inputType)); return compressor; } diff --git a/tools/training/ace/ace_compressors.h b/tools/training/ace/ace_compressors.h index 5b154114b..189710ccb 100644 --- a/tools/training/ace/ace_compressors.h +++ b/tools/training/ace/ace_compressors.h @@ -2,7 +2,9 @@ #pragma once +#include #include +#include #include "openzl/cpp/poly/Span.hpp" #include "tools/training/ace/ace_compressor.h" @@ -20,8 +22,11 @@ poly::span getAllNodes(); poly::span getAllGraphs(); /// @returns the subset of `getAllNodes()` that are compatible with the given -/// @p inputType -poly::span getNodesComptabileWith(Type inputType); +/// @p inputType and usable at @p formatVersion (nodes whose minFormatVersion is +/// 0 are treated as unconstrained). +std::vector getNodesComptabileWith( + Type inputType, + uint32_t formatVersion); /// @returns the subset of `getAllGraphs()` that are compatible with the given /// @p inputType poly::span getGraphsComptabileWith(Type inputType); @@ -37,16 +42,20 @@ poly::span getPrebuiltCompressors(Type inputType); ACECompressor buildRandomGraphCompressor(std::mt19937_64& rng, Type inputType); /// @returns A random compressor that is compatible with the given @p inputType -/// that is a single node followed by ACECompressor successors. +/// that is a single node followed by ACECompressor successors. Only nodes +/// usable at @p formatVersion are considered. ACECompressor buildRandomNodeCompressor( std::mt19937_64& rng, Type inputType, + uint32_t formatVersion, size_t maxDepth = kDefaultMaxDepth); /// @returns A random compressor that is compatible with the given @p inputType. +/// Only nodes usable at @p formatVersion are considered. ACECompressor buildRandomCompressor( std::mt19937_64& rng, Type inputType, + uint32_t formatVersion, size_t maxDepth = kDefaultMaxDepth); ACECompressor buildStoreCompressor(); diff --git a/tools/training/ace/ace_crossover.h b/tools/training/ace/ace_crossover.h index 08ff7aa35..b5516cea1 100644 --- a/tools/training/ace/ace_crossover.h +++ b/tools/training/ace/ace_crossover.h @@ -13,8 +13,8 @@ namespace training { /// a comination of their traits. class ACECrossover { public: - ACECrossover(std::mt19937_64& rng, Type inputType) - : rng_(rng), inputType_(inputType) + ACECrossover(std::mt19937_64& rng, Type inputType, uint32_t formatVersion) + : rng_(rng), inputType_(inputType), formatVersion_(formatVersion) { } @@ -44,7 +44,7 @@ class ACECrossover { return std::move(*child); } } - return ACEMutate(rng_, inputType_)(recipient); + return ACEMutate(rng_, inputType_, formatVersion_)(recipient); } ACECompressor getRandomComponent(const ACECompressor& donor) @@ -88,6 +88,7 @@ class ACECrossover { std::mt19937_64& rng_; Type inputType_; + uint32_t formatVersion_; }; } // namespace training diff --git a/tools/training/ace/ace_mutate.h b/tools/training/ace/ace_mutate.h index 05fe205c4..90bac9cd9 100644 --- a/tools/training/ace/ace_mutate.h +++ b/tools/training/ace/ace_mutate.h @@ -16,8 +16,12 @@ class ACEMutate { ACEMutate( std::mt19937_64& rng, Type inputType, + uint32_t formatVersion, size_t maxDepth = kDefaultMaxDepth) - : rng_(rng), inputType_(inputType), maxDepth_(maxDepth) + : rng_(rng), + inputType_(inputType), + formatVersion_(formatVersion), + maxDepth_(maxDepth) { } @@ -68,7 +72,8 @@ class ACEMutate { if (depth > maxDepth_) { return buildRandomGraphCompressor(rng_, inputType); } else { - return buildRandomCompressor(rng_, inputType, maxDepth_ - depth); + return buildRandomCompressor( + rng_, inputType, formatVersion_, maxDepth_ - depth); } } @@ -107,7 +112,12 @@ class ACEMutate { return randomSimpleCompressor(inputType); } ACEReservoirSampler sampler(rng_); - for (const auto& node : getNodesComptabileWith(inputType)) { + // Bind to a named local: getNodesComptabileWith returns a vector by + // value, and the sampler stores pointers into it that are dereferenced + // after the loop. + const auto compatibleNodes = + getNodesComptabileWith(inputType, formatVersion_); + for (const auto& node : compatibleNodes) { if (node.outputTypes.size() == 1 && compressor.acceptsInputType(node.outputTypes[0])) { sampler.update(node); @@ -122,6 +132,7 @@ class ACEMutate { std::mt19937_64& rng_; Type inputType_; + uint32_t formatVersion_; size_t maxDepth_; }; diff --git a/tools/training/ace/automated_compressor_explorer.cpp b/tools/training/ace/automated_compressor_explorer.cpp index efa6d43ef..0dd51c1f5 100644 --- a/tools/training/ace/automated_compressor_explorer.cpp +++ b/tools/training/ace/automated_compressor_explorer.cpp @@ -21,7 +21,8 @@ std::vector AutomatedCompressorExplorer::initialPopulation() // Use populationSize() random compressors for (size_t i = 0; i < populationSize(); ++i) { - population.push_back(buildRandomCompressor(rng(), inputType())); + population.push_back( + buildRandomCompressor(rng(), inputType(), formatVersion_)); } return population; } @@ -43,9 +44,10 @@ void adjustResults(const ACECompressor& gene, std::vector& results) /* static */ std::vector AutomatedCompressorExplorer::computeFitness( const ACECompressor& gene, - poly::span inputs) + poly::span inputs, + uint32_t formatVersion) { - auto result = gene.benchmark(inputs); + auto result = gene.benchmark(inputs, formatVersion); std::vector fitness(3, std::numeric_limits::infinity()); if (result.has_value()) { fitness[0] = result->compressedSize; @@ -59,7 +61,7 @@ void adjustResults(const ACECompressor& gene, std::vector& results) std::vector AutomatedCompressorExplorer::computeFitness( const ACECompressor& gene) { - return computeFitness(gene, inputs_); + return computeFitness(gene, inputs_, formatVersion_); } std::vector> AutomatedCompressorExplorer::computeFitness( @@ -78,9 +80,10 @@ std::vector> AutomatedCompressorExplorer::computeFitness( continue; } } - futures.emplace_back(threadPool_.run([&inputs = inputs_, &gene] { - return computeFitness(gene, inputs); - })); + futures.emplace_back(threadPool_.run( + [&inputs = inputs_, &gene, formatVersion = formatVersion_] { + return computeFitness(gene, inputs, formatVersion); + })); } std::vector> results; diff --git a/tools/training/ace/automated_compressor_explorer.h b/tools/training/ace/automated_compressor_explorer.h index ff8f0c3dd..2304c0c0f 100644 --- a/tools/training/ace/automated_compressor_explorer.h +++ b/tools/training/ace/automated_compressor_explorer.h @@ -54,6 +54,9 @@ class AutomatedCompressorExplorer : public GeneticAlgorithm { struct Parameters : public Base::Parameters { size_t numThreads{ std::thread::hardware_concurrency() / 2 }; + /// Target format version. Must be set explicitly; a value of 0 is + /// rejected at construction. + uint32_t formatVersion{ 0 }; }; /** @@ -78,10 +81,14 @@ class AutomatedCompressorExplorer : public GeneticAlgorithm { const Parameters& params) : Base(params), inputs_(std::move(inputs)), + formatVersion_(params.formatVersion), threadPool_(params.numThreads), - crossover_(rng(), inputType()), - mutate_(rng(), inputType()) + crossover_(rng(), inputType(), formatVersion_), + mutate_(rng(), inputType(), formatVersion_) { + if (formatVersion_ == 0) { + throw Exception("Format version must be set in the parameters"); + } for (auto const& input : inputs_) { if (input.type() != inputType()) { throw Exception("All inputs must have the same type"); @@ -89,10 +96,18 @@ class AutomatedCompressorExplorer : public GeneticAlgorithm { } } - explicit AutomatedCompressorExplorer( + AutomatedCompressorExplorer( poly::span inputs, poly::string_view snapshot) - : AutomatedCompressorExplorer(inputs, Parameters{}) + : AutomatedCompressorExplorer(inputs, snapshot, Parameters{}) + { + } + + AutomatedCompressorExplorer( + poly::span inputs, + poly::string_view snapshot, + const Parameters& params) + : AutomatedCompressorExplorer(inputs, params) { loadPopulation(snapshot); } @@ -110,6 +125,11 @@ class AutomatedCompressorExplorer : public GeneticAlgorithm { return inputs_; } + uint32_t formatVersion() const + { + return formatVersion_; + } + /** * Saves the current population to a string. */ @@ -152,9 +172,11 @@ class AutomatedCompressorExplorer : public GeneticAlgorithm { private: static std::vector computeFitness( const ACECompressor& compressor, - poly::span inputs); + poly::span inputs, + uint32_t formatVersion); poly::span inputs_; + uint32_t formatVersion_; ThreadPool threadPool_; ACECrossover crossover_; ACEMutate mutate_; diff --git a/tools/training/tests/test_ace.cpp b/tools/training/tests/test_ace.cpp index ea7a63404..22da05388 100644 --- a/tools/training/tests/test_ace.cpp +++ b/tools/training/tests/test_ace.cpp @@ -1,5 +1,6 @@ // Copyright (c) Meta Platforms, Inc. and affiliates. +#include #include #include #include @@ -53,6 +54,7 @@ class ACETest : public testing::Test { params.numThreads = 4; params.populationSize = 50; params.maxGenerations = 100; + params.formatVersion = ZL_MAX_FORMAT_VERSION; } ACECompressor runOnInput(poly::span input) @@ -75,6 +77,44 @@ class ACETest : public testing::Test { std::unique_ptr ace; }; +TEST_F(ACETest, FormatVersionMustBeSet) +{ + auto data = tripleDeltaData(); + std::vector inputs; + inputs.push_back( + Input::refSerial(data.data(), data.size() * sizeof(data[0]))); + + const AutomatedCompressorExplorer::Parameters defaultParams; + EXPECT_EQ(defaultParams.formatVersion, 0); + + EXPECT_THROW( + AutomatedCompressorExplorer defaultAce(inputs, defaultParams), + Exception); +} + +TEST_F(ACETest, CompressesAtAllFormatVersions) +{ + // A compressor trained for a specific, non-default format version must + // actually compress at that version. + + for (uint32_t version = ZL_MIN_FORMAT_VERSION; + version < ZL_MAX_FORMAT_VERSION; + version++) { + params.formatVersion = version; + + auto data = tripleDeltaData(); + auto input = + Input::refSerial(data.data(), data.size() * sizeof(data[0])); + auto solution = runOnInput(input); + + EXPECT_EQ(ace->formatVersion(), version); + + auto result = solution.benchmark(ace->inputs(), version); + ASSERT_TRUE(result.has_value()); + ASSERT_GT(result->compressedSize, 0u); + } +} + TEST_F(ACETest, ACEReservoirSampler) { std::mt19937_64 rng(0xdeadbeef); @@ -108,15 +148,17 @@ TEST_F(ACETest, SerializeDeserialize) std::mt19937_64 rng(0xdeadbeef); for (auto type : { Type::Serial, Type::Struct, Type::Numeric, Type::String }) { - ACEMutate mutator(rng, type); + ACEMutate mutator(rng, type, ZL_MAX_FORMAT_VERSION); for (const auto& compressor : getPrebuiltCompressors(type)) { testRoundTrip(compressor); auto mutated = mutator(compressor); testRoundTrip(mutated); } testRoundTrip(buildRandomGraphCompressor(rng, type)); - testRoundTrip(buildRandomNodeCompressor(rng, type)); - testRoundTrip(buildRandomCompressor(rng, type)); + testRoundTrip(buildRandomNodeCompressor( + rng, type, ZL_MAX_FORMAT_VERSION, kDefaultMaxDepth)); + testRoundTrip(buildRandomCompressor( + rng, type, ZL_MAX_FORMAT_VERSION, kDefaultMaxDepth)); } } @@ -125,7 +167,7 @@ TEST_F(ACETest, TripleDeltaNumeric) auto data = tripleDeltaData(); auto input = Input::refNumeric(poly::span(data)); auto solution = runOnInput(input); - auto result = solution.benchmark(ace->inputs()); + auto result = solution.benchmark(ace->inputs(), ZL_MAX_FORMAT_VERSION); ASSERT_TRUE(result.has_value()); ASSERT_LE(result->compressedSize, 90); } @@ -135,7 +177,7 @@ TEST_F(ACETest, TripleDeltaSerial) auto data = tripleDeltaData(); auto input = Input::refSerial(data.data(), data.size() * sizeof(data[0])); auto solution = runOnInput(input); - auto result = solution.benchmark(ace->inputs()); + auto result = solution.benchmark(ace->inputs(), ZL_MAX_FORMAT_VERSION); ASSERT_TRUE(result.has_value()); ASSERT_LE(result->compressedSize, 90); } @@ -145,7 +187,7 @@ TEST_F(ACETest, TripleDeltaStruct) auto data = tripleDeltaData(); auto input = Input::refStruct(poly::span(data)); auto solution = runOnInput(input); - auto result = solution.benchmark(ace->inputs()); + auto result = solution.benchmark(ace->inputs(), ZL_MAX_FORMAT_VERSION); ASSERT_TRUE(result.has_value()); ASSERT_LE(result->compressedSize, 90); } @@ -155,7 +197,7 @@ TEST_F(ACETest, TripleDeltaString) auto [content, lengths] = tripleDeltaStringData(); auto input = Input::refString(content, lengths); auto solution = runOnInput(input); - auto result = solution.benchmark(ace->inputs()); + auto result = solution.benchmark(ace->inputs(), ZL_MAX_FORMAT_VERSION); ASSERT_TRUE(result.has_value()); ASSERT_LE(result->compressedSize, 110); } @@ -165,7 +207,7 @@ TEST_F(ACETest, savePopulation) auto data = tripleDeltaData(); auto input = Input::refSerial(data.data(), data.size() * sizeof(data[0])); auto solution = runOnInput(input); - auto result = solution.benchmark(ace->inputs()); + auto result = solution.benchmark(ace->inputs(), ZL_MAX_FORMAT_VERSION); ASSERT_TRUE(result.has_value()); ASSERT_LE(result->compressedSize, 90); auto snapshot = ace->savePopulation(); @@ -178,7 +220,8 @@ TEST_F(ACETest, savePopulation) // Initial population doesn't have a good solution { auto solution2 = ace2.solution()[0].first; - auto result2 = solution2.benchmark(ace->inputs()); + auto result2 = + solution2.benchmark(ace->inputs(), ZL_MAX_FORMAT_VERSION); ASSERT_TRUE(result2.has_value()); ASSERT_GT(result2->compressedSize, 90); ASSERT_NE(solution, solution2); @@ -188,7 +231,8 @@ TEST_F(ACETest, savePopulation) ace2.loadPopulation(snapshot); { auto solution2 = ace2.solution()[0].first; - auto result2 = solution2.benchmark(ace->inputs()); + auto result2 = + solution2.benchmark(ace->inputs(), ZL_MAX_FORMAT_VERSION); ASSERT_TRUE(result2.has_value()); ASSERT_LE(result2->compressedSize, 90); } From 9f99c5561cc81db3d84814e87a7ec468379ee6d9 Mon Sep 17 00:00:00 2001 From: Kevin Zhang Date: Wed, 12 Aug 2026 12:38:47 -0700 Subject: [PATCH 4/4] Add format versioning to ML-selector trainer (#942) Summary: Pull Request resolved: https://github.com/facebook/openzl/pull/942 Add targeting per format version to ML selector. The successor are obtained from the ML Selector graph, so when the format version is inadequate to support them, the graph must be rewritten to exclude the unsupported graphs. Implements this functionality. Reviewed By: terrelln Differential Revision: D113974920 --- tools/ml_selector/ml_selector_trainer.cpp | 58 ++++-- tools/ml_selector/tests/BUCK | 1 + .../tests/test_mlSelectorTrainer.cpp | 182 ++++++++++++++++++ 3 files changed, 229 insertions(+), 12 deletions(-) diff --git a/tools/ml_selector/ml_selector_trainer.cpp b/tools/ml_selector/ml_selector_trainer.cpp index fcbc29841..07517f4c1 100644 --- a/tools/ml_selector/ml_selector_trainer.cpp +++ b/tools/ml_selector/ml_selector_trainer.cpp @@ -14,7 +14,9 @@ #include "tools/logger/Logger.h" #include "tools/training/graph_mutation/graph_mutation_utils.h" #include "tools/training/sample_collection/training_sample_collector.h" +#include "tools/training/train_exceptions.h" #include "tools/training/utils/serialized_compressor_internal.h" +#include "tools/training/utils/utils.h" // Suppress warnings for XGBoost headers #pragma GCC diagnostic push @@ -515,15 +517,22 @@ static GBTPredictorWrapper trainXGBoostModel( static void updateCompressor( Compressor& compressor, ZL_MLSelectorConfig& config, - std::string& mlSelectorGraphName) + std::string& mlSelectorGraphName, + const std::vector& successorGraphs) { - Arena* arena = ALLOC_HeapArena_create(); - A1C_Arena a1cArena = A1C_Arena_wrap(arena); + const auto arena = detail::NonNullUniqueCPtr( + ALLOC_HeapArena_create(), ALLOC_Arena_freeArena); + A1C_Arena a1cArena = A1C_Arena_wrap(arena.get()); ZL_SerializedMLConfig serializedConfig = unwrap( MLSelector_serializeMLSelectorConfig(nullptr, &config, &a1cArena)); + // Override the successor list alongside the trained model so the graph + // offers exactly the successors the model was trained over, keeping the + // model's dense class indices aligned with the graph at inference time. ZL_GraphParameters newParams = { + .customGraphs = successorGraphs.data(), + .nbCustomGraphs = successorGraphs.size(), .mparam = { .content = serializedConfig.data, .size = serializedConfig.size, @@ -537,8 +546,6 @@ static void updateCompressor( auto result = ZL_Compressor_overrideGraphParams( compressor.get(), existingMlSelectorGraphId, &newParams); - ALLOC_Arena_freeArena(arena); - if (ZL_isError(result)) { throw std::runtime_error("Error overriding graph params"); } @@ -550,6 +557,7 @@ SerializedCompressorInternal trainMLSelectorGraph( const TrainParams& trainParams) { (void)trainParams; + const auto formatVersion = compressor.getParameter(CParam::FormatVersion); // Find the ML selector graph by prefix auto mlSelectorGraphs = graph_mutation::findAllGraphsWithPrefix( @@ -585,9 +593,20 @@ SerializedCompressorInternal trainMLSelectorGraph( auto cctx = refCCtxForTraining(compressor); - // Collect inputs for mlSelector graph + // Collect the input streams that reach the selector using a CCtx + // pinned to the maximum format version. Input collection only + // records what flows INTO the selector, but it compresses the whole + // (still untrained) graph, and the untrained selector routes every + // input to its first successor. Collecting at the target version + // would fail here if that successor cannot encode at the target + // version, before we get the chance to filter it out below. The + // successors are still trained/benchmarked at the target version + // via `cctx`, which carries the compressor's format version. + auto collectionCctx = refCCtxForTraining(compressor); + collectionCctx.setParameter( + CParam::FormatVersion, ZL_MAX_FORMAT_VERSION); auto mlSelectorInputs = collectInputStreamsForGraph( - inputs, mlSelectorGraphName, cctx); + inputs, mlSelectorGraphName, collectionCctx); if (mlSelectorInputs.empty()) { continue; @@ -602,19 +621,30 @@ SerializedCompressorInternal trainMLSelectorGraph( successorGraphs.push_back(successors.graphids[i]); } + // Filter out only successors supported by specified format version. + // Note that since the custom graphs are filtered, the set of custom + // graphs must be modified later to be the filtered list. + const auto supportedSuccessors = filterGraphsByFormatVersion( + compressor, successorGraphs, mlSelectorInputs); + if (supportedSuccessors.empty()) { + throw FormatVersionUnsupportedError( + "no ML selector successor supports format version " + + std::to_string(formatVersion)); + } + ProcessedMLTrainingSamples trainingSample = extractMLFeatures( - mlSelectorInputs, compressor, cctx, successorGraphs); + mlSelectorInputs, compressor, cctx, supportedSuccessors); TestTrainData splitData = trainTestSplit( trainingSample.features, trainingSample.numericLabels); GBTPredictorWrapper gbtPred = - trainXGBoostModel(splitData, successors.nbGraphIDs); + trainXGBoostModel(splitData, supportedSuccessors.size()); GBTModel coreModel = { .predictor = gbtPred.core_predictor_.get(), .featureGenerator = FeatureGen_integer, - .nbSuccessors = successors.nbGraphIDs, + .nbSuccessors = supportedSuccessors.size(), .nbFeatures = trainingSample.featurePtrNames.size(), .featureLabels = trainingSample.featurePtrNames.data(), }; @@ -622,8 +652,12 @@ SerializedCompressorInternal trainMLSelectorGraph( ZL_MLSelectorConfig config = { .model = ZL_GBT, .runtimeConfig = &coreModel }; - // Update compressor with new trained config - updateCompressor(compressor, config, mlSelectorGraphName); + // Update compressor with the trained config and filtered successors + updateCompressor( + compressor, + config, + mlSelectorGraphName, + supportedSuccessors); trainedAnyThisPass = true; trainedMlSelectors.insert(mlSelectorGraphName); } diff --git a/tools/ml_selector/tests/BUCK b/tools/ml_selector/tests/BUCK index 68744d44a..cd35a5aae 100644 --- a/tools/ml_selector/tests/BUCK +++ b/tools/ml_selector/tests/BUCK @@ -14,6 +14,7 @@ zs_unittest( "../../../tests:utils", "../../../tests/datagen:datagen", "../../training:train", + "../../training/graph_mutation:graph_mutation", "fbsource//third-party/googletest:gtest", ], ) diff --git a/tools/ml_selector/tests/test_mlSelectorTrainer.cpp b/tools/ml_selector/tests/test_mlSelectorTrainer.cpp index 1ca678865..0e01967e1 100644 --- a/tools/ml_selector/tests/test_mlSelectorTrainer.cpp +++ b/tools/ml_selector/tests/test_mlSelectorTrainer.cpp @@ -1,15 +1,20 @@ // Copyright (c) Meta Platforms, Inc. and affiliates. #include +#include "openzl/codecs/zl_conversion.h" +#include "openzl/codecs/zl_lz.h" #include "openzl/codecs/zl_mlselector.h" #include "openzl/cpp/CCtx.hpp" #include "openzl/cpp/Compressor.hpp" #include "openzl/cpp/DCtx.hpp" +#include "openzl/zl_reflection.h" #include "tests/datagen/DataGen.h" #include "tests/ml_selector_utils.h" #include "tools/ml_selector/ml_features.h" #include "tools/ml_selector/ml_selector_trainer.h" +#include "tools/training/graph_mutation/graph_mutation_utils.h" #include "tools/training/train.h" +#include "tools/training/train_exceptions.h" #include "tools/training/train_params.h" #include "tools/training/utils/serialized_compressor_internal.h" #include "tools/training/utils/utils.h" @@ -94,6 +99,44 @@ class TestMLSelectorTrainer : public testing::Test { return successorGraphs; } + // Mirrors setUpCompressor() but appends a version-gated successor and + // reports the ml selector graph's name so its trained successor list can be + // inspected after training. + std::vector setUpCompressorWithGatedSuccessor( + Compressor& compressor) + { + std::vector successorGraphs; + // Register LZ Graph as an additional potential successor - which + // requires v24 format version + successorGraphs.push_back(ZL_Compressor_registerStaticGraph_fromNode1o( + compressor.get(), + ZL_NODE_CONVERT_NUM_TO_SERIAL_LE, + ZL_GRAPH_LZ)); + // Add normal successors afterwards to ensure removing intermediate + // successor is not problematic + for (auto& successor : registerSuccessors(compressor, true)) { + successorGraphs.push_back(successor); + } + + auto mlSelectorGraphId = ZL_Compressor_buildUntrainedMLSelector( + compressor.get(), + successorGraphs.data(), + successorGraphs.size()); + EXPECT_TRUE(!ZL_RES_isError(mlSelectorGraphId)); + ZL_GraphID mlSelectorGid = ZL_RES_value(mlSelectorGraphId); + + ZL_GraphID staticGraph = ZL_Compressor_registerStaticGraph_fromNode1o( + compressor.get(), + ZL_NODE_CONVERT_SERIAL_TO_NUM_LE64, + mlSelectorGid); + ZL_GraphParameters const wrapperDesc = {}; + auto sgid = ZL_Compressor_parameterizeGraph( + compressor.get(), staticGraph, &wrapperDesc); + EXPECT_TRUE(!ZL_RES_isError(sgid)); + compressor.selectStartingGraph(ZL_RES_value(sgid)); + return successorGraphs; + } + std::vector> generateTestData() { std::vector> data; @@ -297,6 +340,44 @@ class TestMLSelectorTrainer : public testing::Test { return mlCompressor; } + // Number of successors the (trained) ML selector graph currently offers. + // Uses the same prefix-based lookup as the trainer so it works on a + // deserialized compressor regardless of the graph's auto-generated name. + size_t mlSelectorSuccessorCount(const Compressor& compressor) + { + auto graphs = training::graph_mutation::findAllGraphsWithPrefix( + compressor, training::ML_SELECTOR_GRAPH_NAME); + EXPECT_EQ(graphs.size(), 1u); + if (graphs.empty()) { + return 0; + } + return training::graph_mutation::getCustomGraphs( + compressor, graphs.front()) + .size(); + } + + // Attempts to compress input at the given format version, returning whether + // compression succeeded. + bool compressesAtVersion( + Compressor& compressor, + const std::vector& input, + uint32_t formatVersion) + { + compressor.setParameter(CParam::FormatVersion, formatVersion); + CCtx cctx; + cctx.refCompressor(compressor); + auto sInput = + Input::refSerial(input.data(), input.size() * sizeof(uint64_t)); + std::string dst( + ZL_compressBound(input.size() * sizeof(uint64_t)), '\0'); + try { + cctx.compressOne(dst, sInput); + } catch (const openzl::Exception&) { + return false; + } + return true; + } + protected: Compressor compressor_; Compressor trainedCompressor_; @@ -412,6 +493,107 @@ TEST_F(TestMLSelectorTrainer, TrainRoundTrip) testRoundTrip(testData_.front(), mlCompressor); } +// When a successor cannot encode at the target format version, the selector +// should be trained on the successors that survive filtering rather than +// failing the whole selector. The trained selector must then offer exactly the +// surviving successors so its class indices stay aligned at inference. +TEST_F(TestMLSelectorTrainer, + TrainsAndCompressesOnSuccessorsSupportedAtFormatVersion) +{ + // Below ZL_GRAPH_LZ's floor (24), so the LZ-backed successor is dropped. + constexpr uint32_t kReducedFormatVersion = 23; + + auto successors = setUpCompressorWithGatedSuccessor(trainedCompressor_); + + trainedCompressor_.setParameter( + CParam::FormatVersion, kReducedFormatVersion); + + // Trains on compressor setup to target format version. The v24-only LZ + // successor is filtered out and the selector is trained on the survivors + // instead of failing. + auto serializedWithFormatTarget = openzl::training::trainMLSelectorGraph( + multiInputs_, trainedCompressor_, trainParams_); + Compressor filtered = deserializeCompressor(*serializedWithFormatTarget); + + // The trained selector offers exactly the surviving successors (LZ dropped) + // so its class indices stay aligned with the graph at inference. + EXPECT_EQ(mlSelectorSuccessorCount(filtered), successors.size() - 1); + + // The trained compressor compresses successfully at the target format + // version, since no surviving successor requires a newer version. + for (const auto& input : testData_) { + EXPECT_TRUE( + compressesAtVersion(filtered, input, kReducedFormatVersion)); + } + + // Training at the (higher) max format version keeps the LZ successor. + Compressor fullCompressor; + fullCompressor.setParameter(CParam::FormatVersion, ZL_MAX_FORMAT_VERSION); + auto fullSuccessors = setUpCompressorWithGatedSuccessor(fullCompressor); + auto serializedWithoutFormatTarget = openzl::training::trainMLSelectorGraph( + multiInputs_, fullCompressor, trainParams_); + Compressor full = deserializeCompressor(*serializedWithoutFormatTarget); + EXPECT_EQ(mlSelectorSuccessorCount(full), fullSuccessors.size()); + + // The retained LZ successor requires v24, so a graph that reaches it does + // not compress at the reduced version, but does at the max version. + auto lzGraph = fullCompressor.buildStaticGraph( + ZL_NODE_CONVERT_SERIAL_TO_NUM_LE64, { fullSuccessors.front() }); + fullCompressor.selectStartingGraph(lzGraph); + for (const auto& input : testData_) { + EXPECT_FALSE(compressesAtVersion( + fullCompressor, input, kReducedFormatVersion)); + EXPECT_TRUE(compressesAtVersion( + fullCompressor, input, ZL_MAX_FORMAT_VERSION)); + } +} + +// When no successor can encode at the target format version, the selector +// cannot be trained on any survivor and the trainer throws so the orchestrator +// can fall the selector back to zstd. +TEST_F(TestMLSelectorTrainer, + TrainerThrowsWhenNoSupportedFormatVersionSuccessorsExist) +{ + constexpr uint32_t kReducedFormatVersion = 23; + + // Build a selector whose only successors run ZL_GRAPH_LZ (v24+), so nothing + // survives filtering at v23. Two successors are the minimum the untrained + // selector accepts (a single-forest model requires exactly two successors). + std::vector gatedSuccessors = { + ZL_Compressor_registerStaticGraph_fromNode1o( + trainedCompressor_.get(), + ZL_NODE_CONVERT_NUM_TO_SERIAL_LE, + ZL_GRAPH_LZ), + ZL_Compressor_registerStaticGraph_fromNode1o( + trainedCompressor_.get(), + ZL_NODE_CONVERT_NUM_TO_SERIAL_LE, + ZL_GRAPH_LZ), + }; + auto mlSelectorGraphId = ZL_Compressor_buildUntrainedMLSelector( + trainedCompressor_.get(), + gatedSuccessors.data(), + gatedSuccessors.size()); + ASSERT_FALSE(ZL_RES_isError(mlSelectorGraphId)); + + ZL_GraphID staticGraph = ZL_Compressor_registerStaticGraph_fromNode1o( + trainedCompressor_.get(), + ZL_NODE_CONVERT_SERIAL_TO_NUM_LE64, + ZL_RES_value(mlSelectorGraphId)); + ZL_GraphParameters const wrapperDesc = {}; + auto sgid = ZL_Compressor_parameterizeGraph( + trainedCompressor_.get(), staticGraph, &wrapperDesc); + ASSERT_FALSE(ZL_RES_isError(sgid)); + trainedCompressor_.selectStartingGraph(ZL_RES_value(sgid)); + + trainedCompressor_.setParameter( + CParam::FormatVersion, kReducedFormatVersion); + + EXPECT_THROW( + openzl::training::trainMLSelectorGraph( + multiInputs_, trainedCompressor_, trainParams_), + openzl::training::FormatVersionUnsupportedError); +} + TEST_F(TestMLSelectorTrainer, TestFeatureExtraction) { // Use refNumeric here since extactMLFeatures expects numeric data