perf: warm up vendored kmeans1d JIT compile before spawning workers

Palettization spawns many worker processes that each call
_kmeans1d.cluster(). Since the vendored extension JIT-compiles on first
use per-process, every worker independently raced the same on-disk
build lock: one worker compiled while the rest spin-waited for it
instead of just loading the finished .so, serializing a one-time
multi-minute compile across the critical path.

Add _kmeans1d.warmup() and call it once in the parent process before
each worker pool is created (KMeansPalettizer and the coreai-torch
palettize_utils compress pool), so every worker's load() call finds
the extension already built.

Author: guru-desh

src/coreai_opt/_utils/_kmeans1d/__init__.py

@@ -3,6 +3,6 @@
 # Use of this source code is governed by a BSD-3-Clause license that can
 # be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause
 
-from coreai_opt._utils._kmeans1d.core import Clustered, cluster
+from coreai_opt._utils._kmeans1d.core import Clustered, cluster, warmup
 
-__all__ = ["Clustered", "cluster"]
+__all__ = ["Clustered", "cluster", "warmup"]

src/coreai_opt/_utils/_kmeans1d/core.py

@@ -60,6 +60,16 @@ def _dll():
     return _DLL
 
 
+def warmup() -> None:
+    """Force the JIT-compiled extension to build once in the current process.
+
+    Call this in a parent process before spawning workers that call
+    cluster(); each spawned worker's own first-use compile would otherwise
+    race the same on-disk build lock and block on the others' compile time.
+    """
+    _dll()
+
+
 def cluster(array: Sequence[float], k: int, *, weights: Sequence[float] | None = None) -> Clustered:
     """
     :param array: A sequence of floats

src/coreai_opt/coreai_utils/_utils/palettize_utils.py

@@ -256,6 +256,11 @@ def _grouped_channelwise_compress(
 
     if mode == "KMEANS" and num_kmeans_workers > 1:
         if _compress_pool is None:
+            # Warm up the JIT-compiled kmeans1d extension here first: each
+            # worker's own first-use compile would otherwise race the same
+            # on-disk build lock and block on the others' compile time
+            # instead of just loading the built artifact.
+            _kmeans1d.warmup()
             _compress_pool = Pool(processes=num_kmeans_workers)
             atexit.register(lambda: _compress_pool.terminate())
         lut, indices = zip(

src/coreai_opt/palettization/kmeans/palettizer.py

@@ -16,6 +16,7 @@
 import torch.nn.utils.parametrize as P
 from tqdm import tqdm
 
+from coreai_opt._utils import _kmeans1d
 from coreai_opt._utils.eager_utils import (
     EagerCompressionComponentBuilderMixin as _EagerCompressionComponentBuilderMixin,
 )
@@ -501,7 +502,11 @@ def _calculate_centroids_parallel(self, num_workers: int) -> None:
         )
 
         # spawn (not fork) so workers don't inherit the parent's CUDA context
-        # or other process-global state.
+        # or other process-global state. Warm up the JIT-compiled kmeans1d
+        # extension here first: each worker's own first-use compile would
+        # otherwise race the same on-disk build lock and block on the
+        # others' compile time instead of just loading the built artifact.
+        _kmeans1d.warmup()
         pool_args = [(info.fp_module, info.weight, info.layer_name) for info in fp_info]
         ctx = mp.get_context("spawn")
         with ctx.Pool(processes=effective_workers) as pool:

tests/test_utils/test_kmeans1d.py

@@ -164,6 +164,15 @@ def test_centroids_are_ascending_and_labels_in_range(self):
         assert clusters.shape == array.shape
         assert centroids.dtype == np.float64
 
+    def test_warmup_then_cluster_still_works(self):
+        # warmup() must be safe to call (including redundantly, since _dll()
+        # short-circuits once built) and must not affect subsequent results.
+        # k=1 is unambiguous: the single centroid is just the mean.
+        _kmeans1d.warmup()
+        _kmeans1d.warmup()
+        result = _kmeans1d.cluster([1.0, 5.0, 9.0], 1)
+        np.testing.assert_allclose(result.centroids, [5.0], rtol=_RTOL, atol=_ATOL)
+
     def test_accepts_list_and_preserves_original_order(self):
         # A plain Python list (the doc-script call style) must work, and labels are
         # returned in the original (unsorted) order.