Save ACE state inside compressor when flag is passed (#733)

Summary:
Pull Request resolved: #733

Adds a flag `--save-ace-state` to the `zli` that saves the ace state to the serialized compressor produced.

Implements this by preserving the ACE state local parameter after graph replacement is done.

Reviewed By: daniellerozenblit

Differential Revision: D102877314

fbshipit-source-id: f0455f2e21aeba3b034bc4fd5eba79a72702275b

Author: kevinjzhang

cli/args/TrainArgs.h

@@ -108,6 +108,12 @@ class TrainArgs : public GlobalArgs, public ProfileArgs {
                 0,
                 false,
                 "Enables pareto frontier training. This will output a directory containing all compressors in the pareto frontier.");
+        parser.addCommandFlag(
+                cmd(),
+                kSaveAceState,
+                0,
+                false,
+                "Save the ACE state as a local parameter in the trained compressor.");
     }
 
     explicit TrainArgs(const arg::ParsedArgs& parsed)
@@ -184,6 +190,7 @@ class TrainArgs : public GlobalArgs, public ProfileArgs {
                 parsed.cmdHasFlag(cmd(), kNoAceSuccessors);
 
         trainParams.noClustering = parsed.cmdHasFlag(cmd(), kNoClustering);
+        trainParams.saveAceState = parsed.cmdHasFlag(cmd(), kSaveAceState);
         trainParams.compressorGenFunc =
                 custom_parsers::createCompressorFromSerialized;
     }
@@ -226,6 +233,7 @@ class TrainArgs : public GlobalArgs, public ProfileArgs {
     inline static const std::string kMaxFileSizeMb   = "max-file-size-mb";
     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";
 };
 
 } // namespace openzl::cli

cli/tests/BUCK

@@ -440,3 +440,17 @@ custom_unittest(
         ":integration_test_bin",
     ],
 )
+
+custom_unittest(
+    name = "csv_save_ace_state_test",
+    command = [
+        "$(location :integration_test_bin)",
+        "$(location ..:zli)",
+        "CsvSaveAceStateTest.test_train_compress_decompress",
+    ],
+    type = "simple",
+    deps = [
+        "..:zli",
+        ":integration_test_bin",
+    ],
+)

cli/tests/cli_integration_tests.py

@@ -118,6 +118,24 @@ def test_train_compress_decompress(self):
         self.train_compress_decompress()
 
 
+class CsvSaveAceStateTest(_CsvBaseTest):
+    """
+    Test case for CSV training with --save-ace-state flag.
+
+    Verifies that training with --save-ace-state produces a compressor
+    that correctly compresses and decompresses files.
+    """
+
+    def test_train_compress_decompress(self):
+        execute_train(
+            compressor_info=self.training_compressor_info,
+            uncompressed_dir=input_dir_path(self.input_dir_name),
+            trained_compressor_path=self.compressor_info.compressor_str,
+            extra_args="--save-ace-state",
+        )
+        self.compress_and_decompress_samples()
+
+
 class CsvFullSplitTest(_CsvBaseTest):
     """
     Test case for CSV training and compression using the full-split trainer.

tools/training/ace/ace_combination.cpp

@@ -28,7 +28,8 @@ namespace {
  */
 std::shared_ptr<const std::string_view> runReplacements(
         Compressor& compressor,
-        const std::unordered_map<std::string, ACECompressor>& replacements)
+        const std::unordered_map<std::string, ACECompressor>& replacements,
+        bool saveAceState)
 {
     // Add each graph to the compressor
     std::unordered_map<std::string, ZL_GraphID> newGraphIds;
@@ -40,8 +41,17 @@ std::shared_ptr<const std::string_view> runReplacements(
     // Replace each backend graph with the new GraphID
     for (const auto& [backendGraph, newGraphId] : newGraphIds) {
         auto backendGraphId = compressor.getGraph(backendGraph);
+        auto localParams    = LocalParams(ZL_Compressor_Graph_getLocalParams(
+                compressor.get(), backendGraphId.value()));
         compressor.unwrap(ZL_Compressor_overrideBaseGraph(
                 compressor.get(), backendGraphId.value(), newGraphId));
+        if (saveAceState) {
+            auto gp = ZL_GraphParameters{ .localParams = localParams.get() };
+            compressor.unwrap(
+                    ZL_Compressor_overrideGraphParams(
+                            compressor.get(), backendGraphId.value(), &gp),
+                    "Graph replacement failed");
+        }
     }
 
     auto serialized = compressor.serialize();
@@ -85,15 +95,16 @@ std::shared_ptr<const std::string_view> getSmallestCandidate(
         const std::unordered_map<
                 std::string,
                 std::vector<std::pair<ACECompressor, ACECompressionResult>>>&
-                allCandidates)
+                allCandidates,
+        bool saveAceState)
 {
     auto compressor = makeCompressor();
     std::unordered_map<std::string, ACECompressor> replacements;
     replacements.reserve(allCandidates.size());
     for (const auto& [backendGraph, candidates] : allCandidates) {
         replacements.emplace(backendGraph, candidates[0].first);
     }
-    return runReplacements(compressor, replacements);
+    return runReplacements(compressor, replacements, saveAceState);
 }
 
 /**
@@ -127,7 +138,8 @@ std::shared_ptr<const std::string_view> makeCombinedCompressor(
         const std::unordered_map<
                 std::string,
                 std::vector<std::pair<ACECompressor, ACECompressionResult>>>&
-                allCandidates)
+                allCandidates,
+        bool saveAceState)
 {
     const auto& choices = candidate.choices();
     if (allCandidates.size() != choices.size()) {
@@ -143,7 +155,7 @@ std::shared_ptr<const std::string_view> makeCombinedCompressor(
         replacements.emplace(name, compressor);
     }
     auto compressor = makeCompressor();
-    return runReplacements(compressor, replacements);
+    return runReplacements(compressor, replacements, saveAceState);
 }
 
 std::vector<std::pair<ACECompressor, ACECompressionResult>> benchmarkAce(
@@ -276,6 +288,7 @@ std::vector<std::shared_ptr<const std::string_view>> getCombinedCompressors(
         return std::move(
                 *trainParams.compressorGenFunc(*trainedSerializedCompressor));
     };
+
     auto compressor = makeCompressor();
     auto cctx       = refCCtxForTraining(compressor);
     auto serialized = compressor.serialize();
@@ -333,7 +346,8 @@ std::vector<std::shared_ptr<const std::string_view>> getCombinedCompressors(
     }
 
     if (!trainParams.paretoFrontier) {
-        return { getSmallestCandidate(makeCompressor, allCandidates) };
+        return { getSmallestCandidate(
+                makeCompressor, allCandidates, trainParams.saveAceState) };
     }
     std::vector<std::vector<CandidateSelection>> candidates;
     candidates.reserve(allCandidates.size());
@@ -347,7 +361,10 @@ std::vector<std::shared_ptr<const std::string_view>> getCombinedCompressors(
     paretoOptimalResults.reserve(frontier.size());
     for (auto& candidate : frontier) {
         paretoOptimalResults.push_back(makeCombinedCompressor(
-                candidate, makeCompressor, allCandidates));
+                candidate,
+                makeCompressor,
+                allCandidates,
+                trainParams.saveAceState));
     }
     return paretoOptimalResults;
 }

tools/training/tests/test_ace_combination.cpp

@@ -2,6 +2,10 @@
 
 #include <gtest/gtest.h>
 #include <random>
+#include "custom_parsers/dependency_registration.h"
+#include "openzl/cpp/CCtx.hpp"
+#include "openzl/cpp/codecs/ACE.hpp"
+#include "tools/training/ace/ace.h"
 #include "tools/training/ace/ace_combination.h"
 
 namespace openzl {
@@ -142,6 +146,103 @@ TEST_F(ACECombinationTest, ProducesParetoOptimalCombination)
     EXPECT_TRUE(isPareto(frontier));
 }
 
+TEST_F(ACECombinationTest, NoSaveAceStateProducesSmallerCompressor)
+{
+    // Create sample data: triple delta pattern compresses well with ACE
+    std::vector<uint64_t> data(1000, 1);
+    for (size_t i = 1; i < data.size(); ++i) {
+        data[i] += data[i - 1];
+    }
+    for (size_t i = 1; i < data.size(); ++i) {
+        data[i] += data[i - 1];
+    }
+    for (size_t i = 1; i < data.size(); ++i) {
+        data[i] += data[i - 1];
+    }
+    auto input = Input::refSerial(data.data(), data.size() * sizeof(data[0]));
+    std::vector<Input> inputsVec;
+    inputsVec.push_back(std::move(input));
+    std::vector<training::MultiInput> multiInputs;
+    multiInputs.emplace_back(std::move(inputsVec));
+
+    auto compressorGenFunc = [](poly::string_view serialized) {
+        auto compressor = std::make_unique<Compressor>();
+        compressor->deserialize(serialized);
+        return compressor;
+    };
+
+    // Train with saveAceState = true
+    std::shared_ptr<const std::string_view> resultWithAceState;
+    {
+        Compressor compressor;
+        compressor.selectStartingGraph(graphs::ACE()(compressor));
+        compressor.setParameter(CParam::FormatVersion, ZL_MAX_FORMAT_VERSION);
+        training::TrainParams trainParams = {
+            .compressorGenFunc = compressorGenFunc,
+            .threads           = 1,
+            .saveAceState      = true,
+        };
+        ACETrainer trainer;
+        auto results =
+                trainer.train(multiInputs, compressor.serialize(), trainParams);
+        ASSERT_FALSE(results.empty());
+        resultWithAceState = results[0];
+    }
+
+    // Train with saveAceState = false (default)
+    std::shared_ptr<const std::string_view> resultWithoutAceState;
+    {
+        Compressor compressor;
+        compressor.selectStartingGraph(graphs::ACE()(compressor));
+        compressor.setParameter(CParam::FormatVersion, ZL_MAX_FORMAT_VERSION);
+        training::TrainParams trainParams = {
+            .compressorGenFunc = compressorGenFunc,
+            .threads           = 1,
+            .saveAceState      = false,
+        };
+        ACETrainer trainer;
+        auto results =
+                trainer.train(multiInputs, compressor.serialize(), trainParams);
+        ASSERT_FALSE(results.empty());
+        resultWithoutAceState = results[0];
+    }
+
+    auto sizeWithAceState    = resultWithAceState->size();
+    auto sizeWithoutAceState = resultWithoutAceState->size();
+
+    // Serialized compressor without ACE state should be significantly smaller
+    EXPECT_GT(sizeWithAceState, 0);
+    EXPECT_GT(sizeWithoutAceState, 0);
+    EXPECT_LE(sizeWithoutAceState, sizeWithAceState / 2)
+            << "Serialized compressor without ACE state ("
+            << sizeWithoutAceState
+            << " bytes) should be at most half the size of one with ACE state ("
+            << sizeWithAceState << " bytes)";
+
+    // Compress data with both trained compressors and verify identical output
+    auto compressWithResult =
+            [&](const std::string_view& serializedCompressor) {
+                auto comp = compressorGenFunc(serializedCompressor);
+                CCtx cctx;
+                cctx.setParameter(CParam::FormatVersion, ZL_MAX_FORMAT_VERSION);
+                cctx.refCompressor(*comp);
+                auto inputForCompress = Input::refSerial(
+                        data.data(), data.size() * sizeof(data[0]));
+                return cctx.compressOne(inputForCompress);
+            };
+    auto compressedWith    = compressWithResult(*resultWithAceState);
+    auto compressedWithout = compressWithResult(*resultWithoutAceState);
+
+    // Compressed output should be nearly identical — the small difference
+    // is due to ACE training being non-deterministic across independent runs.
+    // Allow up to 10% tolerance.
+    auto maxSize = std::max(compressedWith.size(), compressedWithout.size());
+    auto minSize = std::min(compressedWith.size(), compressedWithout.size());
+    EXPECT_LE(maxSize - minSize, maxSize / 10)
+            << "Compressed data sizes should be within 10%: "
+            << compressedWith.size() << " vs " << compressedWithout.size();
+}
+
 } // namespace tests
 } // namespace training
 } // namespace openzl

tools/training/train_params.h

@@ -29,6 +29,7 @@ struct TrainParams {
     poly::optional<size_t> maxFileSizeMb;
     poly::optional<size_t> maxTotalSizeMb;
     bool paretoFrontier{ false };
+    bool saveAceState{ false };
 };
 
 } // namespace openzl::training