Add format versioning to clustering training (#941)

cli/commands/cmd_train.cpp

@@ -84,6 +84,12 @@ int cmdTrain(const TrainArgs& args)
             training::SampleLimiter(maxTotalSize, maxFileSize, numSamples);
     auto filteredInputsPtr = limiter.getFilteredInputsPtr(inputs);
 
+    // Defaults to using ZL_MAX_FORMAT_VERSION for training
+    if (args.compressor()->getParameter(CParam::FormatVersion) == 0) {
+        args.compressor()->setParameter(
+                CParam::FormatVersion, ZL_MAX_FORMAT_VERSION);
+    }
+
     // Benchmark the untrained compressor
     BenchmarkArgs benchmarkArgs(args, args.compressor());
     benchmarkArgs.inputs = training::inputSetToMultiInputs(*filteredInputsPtr);

tools/training/BUCK

@@ -34,6 +34,7 @@ zs_cxxlibrary(
         "trained_candidate.cpp",
     ],
     headers = [
+        "train_exceptions.h",
         "train_params.h",
         "trained_candidate.h",
     ],

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,36 @@ ZL_GraphID clusterSuccessors(
         const TrainParams& trainParams,
         GraphID clusteringGraphUniqueIDUntrained)
 {
-    auto cctx = refCCtxForTraining(compressor);
+    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));
+    }
+
+    CCtx cctx;
+    cctx.setParameter(CParam::StickyParameters, 1);
+    cctx.refCompressor(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,
+            formatVersion);
+
+    if (successorsVec.size() == 0) {
+        throw FormatVersionUnsupportedError(
+                "No compatible successors chosen in clustering graph"
+                + std::to_string(formatVersion));
+    }
 
     // Get clustering codecs for training
     std::vector<ZL_NodeID> clusteringCodecs =

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",

tools/training/tests/test_clustering.cpp

@@ -5,15 +5,18 @@
 #include <thread>
 
 #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<ZL_GraphID> 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<uint8_t> data(1024);
+    for (size_t i = 0; i < data.size(); ++i) {
+        data[i] = static_cast<uint8_t>(i);
+    }
+    training::MultiInput input;
+    input.add(Input::refSerial(data.data(), data.size()));
+    const std::vector<training::MultiInput> inputs = { std::move(input) };
+    EXPECT_THROW(
+            training::trainClusteringGraph(inputs, compressor, trainParams),
+            training::FormatVersionUnsupportedError);
+}
+
 } // namespace
 } // namespace openzl::tests

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

tools/training/tests/test_train.cpp

@@ -15,6 +15,7 @@ TEST(TrainTest, ThrowsTypedErrorWithoutTrainableGraph)
 {
     const std::vector<training::MultiInput> 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<training::MultiInput> inputs;
+    Compressor compressor;
+    compressor.selectStartingGraph(ZL_GRAPH_STORE);
+    const training::TrainParams trainParams = {
+        .compressorGenFunc =
+                [](poly::string_view, poly::string_view) {
+                    return std::make_unique<Compressor>();
+                },
+    };
+
+    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

tools/training/tests/test_utils.cpp

@@ -0,0 +1,146 @@
+// Copyright (c) Meta Platforms, Inc. and affiliates.
+
+#include <gtest/gtest.h>
+
+#include <array>
+#include <vector>
+
+#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<ZL_IDType> graphIds(const std::vector<GraphID>& graphs)
+{
+    std::vector<ZL_IDType> ids;
+    ids.reserve(graphs.size());
+    for (const auto graph : graphs) {
+        ids.push_back(graph.gid);
+    }
+    return ids;
+}
+
+std::vector<MultiInput> serialInputs()
+{
+    static const std::array<uint8_t, 1024> data = [] {
+        std::array<uint8_t, 1024> result{};
+        for (size_t i = 0; i < result.size(); ++i) {
+            result[i] = static_cast<uint8_t>(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, RequiresMinimumFormatVersionForAllGraphs)
+{
+    Compressor compressor;
+    const auto customGraph = compressor.buildStaticGraph(
+            ZL_NODE_CONVERT_STRUCT_TO_SERIAL, { ZL_GRAPH_STORE });
+    compressor.selectStartingGraph(ZL_GRAPH_STORE);
+    const std::vector<GraphID> graphs = { ZL_GRAPH_STORE,
+                                          ZL_GRAPH_ZSTD,
+                                          customGraph };
+
+    EXPECT_THROW(
+            filterGraphsByFormatVersion(
+                    compressor,
+                    graphs,
+                    serialInputs(),
+                    ZL_MIN_FORMAT_VERSION - 1),
+            Exception);
+}
+
+TEST(FilterGraphsByFormatVersionTest, FiltersLZByVersion)
+{
+    Compressor compressor;
+    compressor.setParameter(CParam::FormatVersion, ZL_MAX_FORMAT_VERSION);
+    const auto beforeVersion24 = filterGraphsByFormatVersion(
+            compressor, { ZL_GRAPH_LZ }, serialInputs(), 23);
+    EXPECT_TRUE(beforeVersion24.empty());
+
+    const auto atVersion24 = filterGraphsByFormatVersion(
+            compressor, { ZL_GRAPH_LZ }, serialInputs(), 24);
+    EXPECT_EQ(graphIds(atVersion24), graphIds({ ZL_GRAPH_LZ }));
+    EXPECT_EQ(
+            compressor.getParameter(CParam::FormatVersion),
+            ZL_MAX_FORMAT_VERSION);
+}
+
+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(),
+                    ZL_MIN_FORMAT_VERSION),
+            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<uint32_t, 64> data{};
+    MultiInput input;
+    input.add(Input::refStruct(data.data(), data.size()));
+    const std::vector<MultiInput> inputs = { std::move(input) };
+
+    const auto beforeVersion21 =
+            filterGraphsByFormatVersion(compressor, { graph }, inputs, 20);
+    EXPECT_TRUE(beforeVersion21.empty());
+
+    const auto atVersion21 =
+            filterGraphsByFormatVersion(compressor, { graph }, inputs, 21);
+    EXPECT_EQ(graphIds(atVersion21), graphIds({ graph }));
+}
+
+TEST(FilterGraphsByFormatVersionTest, SupportsMultiInputGraphs)
+{
+    Compressor compressor;
+    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<MultiInput> inputs = { std::move(input) };
+
+    const auto supported = filterGraphsByFormatVersion(
+            compressor, { graph }, inputs, ZL_MAX_FORMAT_VERSION);
+
+    EXPECT_EQ(graphIds(supported), graphIds({ graph }));
+}
+
+} // namespace
+} // namespace openzl::training

tools/training/train.cpp

@@ -28,6 +28,22 @@ std::vector<TrainedCandidate> 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(

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