[JAX SC] Add the optionality in input pre-processing to pass in a custom scheduler callback.

Useful in cases where the client wishes to have a tighter control over scheduling the pre-processing work.

PiperOrigin-RevId: 840582756

Author: Google-ML-Automation

jax_tpu_embedding/sparsecore/lib/core/BUILD

@@ -54,10 +54,12 @@ cc_library(
     deps = [
         ":all_reduce_interface",
         ":coo_format",
+        ":input_preprocessing_threads",
         ":partitioned_coo_tensors",
         "@com_google_absl//absl/base:core_headers",
         "@com_google_absl//absl/base:nullability",
         "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_absl//absl/functional:any_invocable",
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/strings:string_view",
@@ -140,6 +142,8 @@ cc_library(
         "@com_google_absl//absl/algorithm:container",
         "@com_google_absl//absl/base:core_headers",
         "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_absl//absl/functional:any_invocable",
+        "@com_google_absl//absl/functional:bind_front",
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/status:statusor",

jax_tpu_embedding/sparsecore/lib/core/input_preprocessing.cc

@@ -220,7 +220,7 @@ void ExtractSortAndGroupCooTensorsForTable(
   });
   for (int local_device = 0; local_device < options.local_device_count;
        ++local_device) {
-    PreprocessingThreadPool()->Schedule(
+    options.async_task_scheduler(
         [&, local_device, &state = state, input_batches] {
           state.extracted_coo_tensors_per_device[local_device] =
               internal::ExtractCooTensorsForAllFeaturesPerLocalDevice(
@@ -267,7 +267,7 @@ void CreateMinibatchingBucketsForTable(
   state.stats_per_host.dropped_id_count = 0;
   for (int local_device = 0; local_device < options.local_device_count;
        ++local_device) {
-    PreprocessingThreadPool()->Schedule([&, local_device, &state = state] {
+    options.async_task_scheduler([&, local_device, &state = state] {
       // Note: We create a dummy stats object here because we don't want to
       // overwrite the stats from the first pass, which are authoritative.
       // The only stat we care about from this second pass is the number of

jax_tpu_embedding/sparsecore/lib/core/input_preprocessing_util.h

@@ -20,16 +20,19 @@
 #include <limits>
 #include <optional>
 #include <string>
+#include <utility>
 #include <vector>
 
 #include "absl/base/attributes.h"  // from @com_google_absl
 #include "absl/base/nullability.h"  // from @com_google_absl
 #include "absl/container/flat_hash_map.h"  // from @com_google_absl
+#include "absl/functional/any_invocable.h"  // from @com_google_absl
 #include "absl/strings/string_view.h"  // from @com_google_absl
 #include "absl/types/span.h"  // from @com_google_absl
 #include "Eigen/Core"  // from @eigen_archive
 #include "jax_tpu_embedding/sparsecore/lib/core/all_reduce_interface.h"
 #include "jax_tpu_embedding/sparsecore/lib/core/coo_format.h"
+#include "jax_tpu_embedding/sparsecore/lib/core/input_preprocessing_threads.h"
 #include "jax_tpu_embedding/sparsecore/lib/core/partitioned_coo_tensors.h"
 
 namespace jax_sc_embedding {
@@ -215,6 +218,12 @@ struct PreprocessSparseDenseMatmulInputOptions {
   // Hash function used for creating minibatching buckets.
   CooFormat::HashFn minibatching_bucketing_hash_fn = HighwayHash;
 
+  // Callback to schedule async work.
+  absl::AnyInvocable<void(std::function<void()>) const> async_task_scheduler =
+      [](std::function<void()> callback) {
+        PreprocessingThreadPool()->Schedule(std::move(callback));
+      };
+
   // Returns the total number of SparseCores across all devices and hosts.
   uint32_t GetNumScs() const { return num_sc_per_device * global_device_count; }
 };