Add training js wrapper (#998)

.github/workflows/wasm-ci.yml

@@ -10,6 +10,13 @@ on:
     branches: [dev]
     paths:
       - 'tools/wasm/**'
+      # The binding links the training tools and the CLI's profile graphs.
+      - 'tools/training/**'
+      - 'tools/io/**'
+      - 'tools/logger/**'
+      - 'cli/utils/profile_graphs.*'
+      - 'custom_parsers/**'
+      - 'cpp/**'
       - 'src/**'
       - 'include/**'
       - '**/CMakeLists.txt'
@@ -18,6 +25,13 @@ on:
     branches: [dev, master, actionsTest]
     paths:
       - 'tools/wasm/**'
+      # The binding links the training tools and the CLI's profile graphs.
+      - 'tools/training/**'
+      - 'tools/io/**'
+      - 'tools/logger/**'
+      - 'cli/utils/profile_graphs.*'
+      - 'custom_parsers/**'
+      - 'cpp/**'
       - 'src/**'
       - 'include/**'
       - '**/CMakeLists.txt'

CMakeLists.txt

@@ -42,6 +42,13 @@ set(OPENZL_SUPPORT_SHARED_LIBRARY "${OPENZL_BUILD_SHARED_LIBS}")
 if(EMSCRIPTEN)
     add_compile_options(-sMEMORY64=1)
     add_link_options(-sMEMORY64=1 -sWASM_BIGINT=1)
+
+    # Training throws across the binding and uses std::thread, so these flags
+    # must apply to every object linked into the WASM artifact.
+    if(OPENZL_BUILD_WASM)
+        add_compile_options(-fexceptions -pthread)
+        add_link_options(-fexceptions -pthread)
+    endif()
 endif()
 
 set(CMAKE_EXPORT_COMPILE_COMMANDS ON)

cli/CMakeLists.txt

@@ -36,7 +36,10 @@ add_dependencies(zli
     commands
     utils
     logger
-    xgboost_external
 )
 
+if(TARGET xgboost_external)
+    add_dependencies(zli xgboost_external)
+endif()
+
 apply_openzl_compile_options_to_target(zli)

custom_parsers/CMakeLists.txt

@@ -44,11 +44,10 @@ target_link_libraries(custom_parsers
     openzl
     openzl_cpp
     csv_parser
-    ml_selector
     parquet_graph
     shared_components
 )
-add_dependencies(custom_parsers openzl openzl_cpp xgboost_external)
+add_dependencies(custom_parsers openzl openzl_cpp)
 apply_openzl_compile_options_to_target(custom_parsers
     csv_parser
     parquet_graph

tools/CMakeLists.txt

@@ -19,22 +19,33 @@ if(OPENZL_BUILD_TOOLS)
   set(OPENZL_BUILD_IO_TOOLS ON)
   set(OPENZL_BUILD_SDDL_TOOLS ON)
   set(OPENZL_BUILD_TRAINING_TOOLS ON)
-  set(OPENZL_BUILD_ML_SELECTOR ON)
+  # xgboost has no Emscripten/WASM build
+  if(NOT EMSCRIPTEN)
+    set(OPENZL_BUILD_ML_SELECTOR ON)
+  endif()
 endif()
 if(OPENZL_BUILD_CLI)
   set(OPENZL_BUILD_ARG_TOOLS ON)
   set(OPENZL_BUILD_IO_TOOLS ON)
   set(OPENZL_BUILD_LOGGER ON)
   set(OPENZL_BUILD_SDDL_TOOLS ON)
   set(OPENZL_BUILD_TRAINING_TOOLS ON)
-  set(OPENZL_BUILD_ML_SELECTOR ON)
+  if(NOT EMSCRIPTEN)
+    set(OPENZL_BUILD_ML_SELECTOR ON)
+  endif()
 endif()
 if(OPENZL_BUILD_PYTHON_EXT)
   set(OPENZL_BUILD_LOGGER ON)
   set(OPENZL_BUILD_IO_TOOLS ON)
   set(OPENZL_BUILD_SDDL_TOOLS ON)
 endif()
 
+if(OPENZL_BUILD_WASM AND EMSCRIPTEN)
+  set(OPENZL_BUILD_IO_TOOLS ON)
+  set(OPENZL_BUILD_LOGGER ON)
+  set(OPENZL_BUILD_TRAINING_TOOLS ON)
+endif()
+
 # Common Utilities
 add_subdirectory(arg)
 add_subdirectory(fileio)

tools/training/BUCK

@@ -12,6 +12,9 @@ zs_cxxlibrary(
     headers = [
         "train.h",
     ],
+    preprocessor_flags = [
+        "-DOPENZL_HAS_ML_SELECTOR_TRAINER=1",
+    ],
     deps = [
         "../../custom_parsers:custom_parsers",
         "../ml_selector:ml_selector_trainer",

tools/training/CMakeLists.txt

@@ -24,22 +24,26 @@ if (OPENZL_BUILD_TRAINING_TOOLS)
   target_include_directories(tools_training PUBLIC ${PROJECT_SOURCE_DIR})
   target_compile_definitions(tools_training PUBLIC ZDICT_STATIC_LINKING_ONLY)
   apply_openzl_compile_options_to_target(tools_training)
-    target_link_libraries(
-        tools_training
-        PUBLIC
-            openzl
-            openzl_cpp
-            tools_io
-            logger
-            ml_selector)
-    add_dependencies(
-        tools_training
-        openzl
-        openzl_cpp
-        tools_io
-        logger
-        ml_selector
-        xgboost_external)
+  target_link_libraries(
+      tools_training
+      PUBLIC
+          openzl
+          openzl_cpp
+          tools_io
+          logger)
+  add_dependencies(
+      tools_training
+      openzl
+      openzl_cpp
+      tools_io
+      logger)
+
+  if (OPENZL_BUILD_ML_SELECTOR)
+    target_compile_definitions(
+        tools_training PRIVATE OPENZL_HAS_ML_SELECTOR_TRAINER=1)
+    target_link_libraries(tools_training PUBLIC ml_selector)
+    add_dependencies(tools_training ml_selector xgboost_external)
+  endif()
 
   if (OPENZL_BUILD_TESTS AND OPENZL_ALLOW_INTROSPECTION)
     file(
@@ -53,7 +57,9 @@ if (OPENZL_BUILD_TRAINING_TOOLS)
         CONFIGURE_DEPENDS
         "${CMAKE_CURRENT_LIST_DIR}/tests/*.h")
     add_executable(test_training ${test_training_srcs} ${test_training_headers})
-    add_dependencies(test_training xgboost_external)
+    if (OPENZL_BUILD_ML_SELECTOR)
+      add_dependencies(test_training xgboost_external)
+    endif()
     target_link_libraries(
         test_training
         PRIVATE

tools/training/train.cpp

@@ -4,7 +4,9 @@
 #include <vector>
 
 #include "tools/logger/Logger.h"
-#include "tools/ml_selector/ml_selector_trainer.h"
+#if defined(OPENZL_HAS_ML_SELECTOR_TRAINER)
+#    include "tools/ml_selector/ml_selector_trainer.h"
+#endif
 #include "tools/training/ace/ace.h"
 #include "tools/training/clustering/clustering_graph_trainer.h"
 #include "tools/training/dict/base_dict_trainer.h"
@@ -35,6 +37,13 @@ std::vector<TrainedCandidate> train(
                 "Compressor format version is not set.");
     }
 
+#if !defined(OPENZL_HAS_ML_SELECTOR_TRAINER)
+    if (graph_mutation::hasTargetGraph(compressor, "zl.ml_selector")) {
+        throw NoTrainableGraphError(
+                "ML selector training is not available in this build.");
+    }
+#endif
+
     // 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)) {
@@ -91,11 +100,13 @@ std::vector<TrainedCandidate> train(
         }
     }
 
+#if defined(OPENZL_HAS_ML_SELECTOR_TRAINER)
     if (graph_mutation::hasTargetGraph(compressor, ML_SELECTOR_GRAPH_NAME)) {
         serializedTrainedCompressors.clear();
         serializedTrainedCompressors.push_back(
                 trainMLSelectorGraph(inputs, compressor, trainParams));
     }
+#endif
 
     // Dict training: for each serialized candidate, deserialize, train
     // dicts, re-serialize with bundleID + dictIDs in CBOR.

tools/wasm/BUCK

@@ -15,6 +15,8 @@ zs_library(
     ],
     deps = [
         "../../cli/utils:profile_graphs",
+        "../../custom_parsers:custom_parsers",
+        "../training:train",
     ],
     exported_deps = [
         "../..:public_headers",

tools/wasm/CMakeLists.txt

@@ -5,23 +5,38 @@
 set(OPENZL_WASM_SOURCE src/openzl_wasm.cpp)
 
 if(NOT EMSCRIPTEN)
-  # openzl_wasm_lib is used only for native testing in non-Emscripten builds
-  add_library(openzl_wasm_lib STATIC ${OPENZL_WASM_SOURCE})
-
-  target_include_directories(openzl_wasm_lib PRIVATE
-    ${CMAKE_CURRENT_SOURCE_DIR}/../..
-    ${CMAKE_CURRENT_SOURCE_DIR}/../../include
-    ${CMAKE_CURRENT_SOURCE_DIR}/../../src
-  )
-
-  target_link_libraries(openzl_wasm_lib PRIVATE openzl openzl_profile_graphs)
-  apply_openzl_compile_options_to_target(openzl_wasm_lib)
-
   # Native test - runs under both `buck test` and, with
   # -DOPENZL_BUILD_TESTS=ON, `ctest -R WasmBinding`. gtest_discover_tests
   # registers the gtest names, not the executable name, so the ctest filter
   # matches the suites.
   if(OPENZL_BUILD_TESTS)
+    add_library(openzl_wasm_lib STATIC ${OPENZL_WASM_SOURCE})
+
+    target_include_directories(openzl_wasm_lib PRIVATE
+      ${CMAKE_CURRENT_SOURCE_DIR}/../..
+      ${CMAKE_CURRENT_SOURCE_DIR}/../../include
+      ${CMAKE_CURRENT_SOURCE_DIR}/../../src
+    )
+
+    target_link_libraries(openzl_wasm_lib
+      PRIVATE
+        openzl
+        openzl_profile_graphs
+        tools_training
+        custom_parsers
+    )
+
+    # Both live in subdirectories that can be switched off, and an unknown name
+    # in the link list above is handed to the linker verbatim, so a disabled
+    # subdirectory would otherwise surface as `cannot find -lcustom_parsers`
+    # at link. Naming them here makes it a configure-time error instead.
+    add_dependencies(openzl_wasm_lib
+      openzl
+      tools_training
+      custom_parsers
+    )
+    apply_openzl_compile_options_to_target(openzl_wasm_lib)
+
     add_executable(test_wasm_binding tests/test_wasm_binding.cpp)
     target_include_directories(test_wasm_binding PRIVATE
       ${CMAKE_CURRENT_SOURCE_DIR}/../..
@@ -57,19 +72,52 @@ target_include_directories(openzl_wasm PRIVATE
   ${CMAKE_CURRENT_SOURCE_DIR}/../../src
 )
 
-target_link_libraries(openzl_wasm PRIVATE openzl openzl_profile_graphs)
+target_link_libraries(openzl_wasm
+  PRIVATE
+    openzl
+    openzl_profile_graphs
+    tools_training
+    custom_parsers
+)
+
+add_dependencies(openzl_wasm
+  openzl
+  tools_training
+  custom_parsers
+)
+
+# Emscripten pre-spawns its worker pool at startup. Creating a thread beyond
+# the pool has to round-trip through the main thread's event loop, which
+# deadlocks when that thread is blocked, so the binding never asks for more
+# than this and the pool is sized up front.
+set(OPENZL_WASM_PTHREAD_POOL_SIZE 4 CACHE STRING
+  "Workers pre-allocated for the WASM trainer thread pool")
 
-# -sMEMORY64=1 / -sWASM_BIGINT=1 are set globally in the top-level file so
-# every object, including the core library, uses the wasm64 ABI.
+# -sMEMORY64=1 / -sWASM_BIGINT=1 / -fexceptions / -pthread are set globally in
+# the top-level file so every object, including the core library and the
+# training tools, uses the wasm64 ABI, can unwind, and can spawn threads.
 target_compile_options(openzl_wasm PRIVATE -O3)
+target_compile_definitions(openzl_wasm PRIVATE
+  OPENZL_WASM_TRAIN_MAX_THREADS=${OPENZL_WASM_PTHREAD_POOL_SIZE})
 target_link_options(openzl_wasm PRIVATE
   -O3                        # optimise the linked wasm
   -sWASM=1                   # emit wasm, not asm.js
   -sALLOW_MEMORY_GROWTH=1    # heap can grow; JS must re-read mod.HEAPU8
   -sINITIAL_MEMORY=32MB      # starting heap size
+  # Growth otherwise stops at the 2GB default, which a large sample plus the
+  # trainers' per-candidate working sets can reach. Raising the ceiling costs
+  # nothing up front: it caps growth rather than reserving.  Only wasm64 can
+  # address beyond 4GB, which the global -sMEMORY64=1 provides.
+  -sMAXIMUM_MEMORY=16GB
+  # Reserve extra stack space for ACE. If a future candidate goes deeper than
+  # 4 MB, training will abort with a stack overflow, and the fix
+  # is to raise these numbers.
+  -sSTACK_SIZE=4MB
+  -sDEFAULT_PTHREAD_STACK_SIZE=4MB
   -sMODULARIZE=1             # export a factory function, not globals
   -sEXPORT_ES6=1             # emit an ES module
   -sEXPORT_NAME=openzlModule # name of the factory; js/wasm_api.js imports this
+  -sPTHREAD_POOL_SIZE=${OPENZL_WASM_PTHREAD_POOL_SIZE}
   # Nothing to list: every entry point, the allocator included, exports itself
   # via EMSCRIPTEN_KEEPALIVE. The empty value overrides the _main default.
   -sEXPORTED_FUNCTIONS=

tools/wasm/README.md

@@ -154,6 +154,17 @@ Node.js 24 supports Memory64 by default. Do not pass
 `--experimental-wasm-memory64` or `--experimental-wasm-table64`. Node 24
 rejects those obsolete flags.
 
+The build also applies `-pthread`, because the compressor trainers fan work out
+over a thread pool. Two consequences:
+
+* **Browsers must serve the page cross-origin isolated.** Threads need
+  `SharedArrayBuffer`, which requires `Cross-Origin-Opener-Policy: same-origin`
+  and `Cross-Origin-Embedder-Policy: require-corp` response headers. Without
+  them the module fails to instantiate. Node needs no equivalent setup.
+* **`-sPTHREAD_POOL_SIZE` bounds the trainers.** Emscripten pre-spawns that many
+  workers at startup. To get more threads, override with
+  `-DOPENZL_WASM_PTHREAD_POOL_SIZE=N` at configure time.
+
 ### Test
 
 The generated files must be next to `wasm_api.js` so its relative imports