Add format versioning to ML-selector trainer

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"
@@ -125,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)
@@ -133,6 +141,13 @@ class TrainArgs : public GlobalArgs, public ProfileArgs {
         // Create the compressor
         setCompressor(createCompressorFromArgs(
                 *this, parsed.cmdFlag(cmd(), kCompressor)));
+        auto formatVersion = parsed.cmdFlag(cmd(), kFormatVersion);
+        if (formatVersion) {
+            compressor()->setParameter(
+                    CParam::FormatVersion,
+                    util::checkedstoi(formatVersion.value()));
+        }
+        applyDefaultFormatVersion();
         auto outputPath = parsed.cmdFlag(cmd(), kOutput);
         if (outputPath) {
             checkOutput(outputPath.value(), parsed.cmdHasFlag(cmd(), kForce));
@@ -228,6 +243,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 +262,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";
 
@@ -265,6 +292,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

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);

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(

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<ZL_GraphID>& successorGraphs)
 {
-    Arena* arena       = ALLOC_HeapArena_create();
-    A1C_Arena a1cArena = A1C_Arena_wrap(arena);
+    const auto arena = detail::NonNullUniqueCPtr<Arena>(
+            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,28 +621,43 @@ 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(),
             };
 
             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);
         }

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