Reduce the runtime of benchmarking draining phase

src/inference_endpoint/async_utils/services/metrics_aggregator/__main__.py

@@ -18,6 +18,7 @@
 import argparse
 import asyncio
 import logging
+import os
 import signal
 from collections.abc import Callable
 from contextlib import AbstractContextManager, nullcontext
@@ -111,6 +112,19 @@ def _on_sigterm() -> None:
     return _on_sigterm, pending_tasks
 
 
+def _cap_live_lane_rayon_pool(tokenizer_workers: int) -> None:
+    """Cap the Hugging Face rayon pool the live (mid-run) lane encodes on.
+
+    The live lane runs in this process, and rayon sizes its pool once, at the
+    first encode, so the cap must be in the environment before the tokenizer is
+    used. ``setdefault`` lets an operator-exported ``RAYON_NUM_THREADS`` win; the
+    ``max(1, ...)`` keeps ``--tokenizer-workers 0`` from reading as rayon's "use
+    every CPU". The end-of-run drain is not capped by this: its shard processes
+    size their own pools in ``_init_worker``.
+    """
+    os.environ.setdefault("RAYON_NUM_THREADS", str(max(1, tokenizer_workers)))
+
+
 async def main() -> None:
     parser = argparse.ArgumentParser(
         description="Metrics aggregator service - subscribes to EventRecords and computes real-time metrics"
@@ -257,6 +271,7 @@ async def main() -> None:
     # (coalesces to 'object' not 'AbstractContextManager[BatchTokenizer | None]')
     tokenizer_cm: AbstractContextManager[BatchTokenizer | None]
     if args.tokenizer:
+        _cap_live_lane_rayon_pool(args.tokenizer_workers)
         try:
             tokenizer_cm = BatchTokenizer(
                 args.tokenizer, live_workers=args.tokenizer_workers