High-throughput async inference: native aiohttp fast-path, scalable async loop, and request-routing keys

core/requirements.txt

@@ -27,7 +27,7 @@ litellm[caching]==1.83.14
 math-verify[antlr4_9_3]
 mcp
 numpy
-openai
+openai[aiohttp]
 openpyxl>=3.1.0
 pandas>=2.0.0
 pyxlsb>=1.0.10

nemo_skills/inference/generate.py

@@ -15,11 +15,22 @@
 import asyncio
 import json
 import logging
+import os
 import random
 import shutil
 import subprocess
 import sys
 import time
+
+# Prefer uvloop when available to reduce event-loop scheduling overhead.
+# Set NEMO_SKILLS_DISABLE_UVLOOP to use the default asyncio event loop.
+if not os.environ.get("NEMO_SKILLS_DISABLE_UVLOOP"):
+    try:
+        import uvloop
+
+        uvloop.install()
+    except ImportError:
+        pass
 from copy import deepcopy
 from dataclasses import asdict, field, is_dataclass
 from pathlib import Path
@@ -125,6 +136,12 @@ class GenerationTaskConfig:
     # maximum number of concurrent requests to the server for the async loop
     # if sync loop is used, this is the batch size
     max_concurrent_requests: int = 512
+    # Target task-creation rate while initially filling the async worker pool.
+    # Tasks are created in short batches to limit connection bursts. Values <= 0 disable ramping.
+    ramp_rate_per_sec: int = 4000
+    # Number of background coroutines that postprocess and optionally evaluate completed datapoints.
+    # Increase this when per-datapoint evaluation benefits from additional concurrency.
+    postprocess_workers: int = 16
     # chunk the dataset into equal sized parts and index into them
     num_chunks: int | None = None  # if specified, will split the data into chunks and only generate for one chunk
     chunk_id: int | None = None  # if specified, will index the specified chunk only
@@ -409,6 +426,9 @@ def __init__(self, cfg: GenerationTaskConfig):
 
         # output_lock will be initialized when async_loop is called
         self.output_lock = None
+        # Pipeline queues are set by async_loop; None selects inline processing for other callers.
+        self._raw_queue = None
+        self._output_queue = None
 
     def setup_prompt(self):
         if self.cfg.prompt_config is None:
@@ -849,16 +869,17 @@ async def _generate_and_save_datapoint(self, data_point, all_data, fout, pbar):
             output["generation_end_time"] = end_time
             output["generation_time"] = end_time - start_time
 
-        await self.postprocess_single_output(output, data_point)
-
-        # evaluate single-data point if requested and evaluator supports that
-        if self.should_run_evaluation and self.evaluator:
-            output = await self.evaluate_single_datapoint({**data_point, **output})
-
-        # Thread-safe output writing
-        async with self.output_lock:
-            self.dump_outputs([output], [data_point], fout)
-            pbar.update(1)
+        # Async-loop callers enqueue completed generations for postprocessing and batched output.
+        if self._raw_queue is not None:
+            await self._raw_queue.put((output, data_point))
+        else:
+            # Other callers postprocess and write inline.
+            await self.postprocess_single_output(output, data_point)
+            if self.should_run_evaluation and self.evaluator:
+                output = await self.evaluate_single_datapoint({**data_point, **output})
+            async with self.output_lock:
+                self.dump_outputs([output], [data_point], fout)
+                pbar.update(1)
 
     async def async_loop(self, data):
         """Async loop to generate generations using asyncio."""
@@ -881,7 +902,8 @@ async def async_loop(self, data):
 
         pbar = tqdm(total=len(remaining_data_points), desc="Remaining generations")
 
-        with open(self.cfg.output_file + "-async", "at", encoding="utf-8", buffering=1) as fout:
+        # The writer flushes after each batch, so default buffering avoids per-record flushes.
+        with open(self.cfg.output_file + "-async", "at", encoding="utf-8") as fout:
             # Dump prefilled data first
             if len(prefilled_data_points) > 0:
                 for output, data_point in zip(prefilled_outputs, prefilled_data_points):
@@ -893,15 +915,119 @@ async def async_loop(self, data):
                 async with self.output_lock:
                     self.dump_outputs(prefilled_outputs, prefilled_data_points, fout)
 
-            # Create tasks for all remaining data points
-            tasks = []
-            for data_point in remaining_data_points:
-                task = asyncio.create_task(self._generate_and_save_datapoint(data_point, data, fout, pbar))
-                tasks.append(task)
-
-            # Wait for all tasks to complete
-            if tasks:
-                await asyncio.gather(*tasks)
+            # Generation tasks -> raw queue -> postprocessors -> output queue -> batched writer.
+            # Bounded queues provide backpressure; None sentinels stop each processing stage.
+            queue_cap = max(self.cfg.max_concurrent_requests * 2, 1024)
+            self._raw_queue: asyncio.Queue = asyncio.Queue(maxsize=queue_cap)
+            self._output_queue = asyncio.Queue(maxsize=queue_cap)
+            num_processors = max(1, int(self.cfg.postprocess_workers))
+
+            async def _processor():
+                # Each processor receives its own sentinel after all raw outputs are queued.
+                while True:
+                    item = await self._raw_queue.get()
+                    if item is None:
+                        return
+                    output, data_point = item
+                    try:
+                        await self.postprocess_single_output(output, data_point)
+                        if self.should_run_evaluation and self.evaluator:
+                            output = await self.evaluate_single_datapoint({**data_point, **output})
+                    except Exception as e:
+                        pos = data_point.get(self.cfg.async_position_key, "?")
+                        LOG.exception("postprocess/evaluate failed for async_pos=%s", pos)
+                        # Preserve ordering metadata and failure details for the output record.
+                        if self.cfg.async_position_key in data_point:
+                            output[self.cfg.async_position_key] = data_point[self.cfg.async_position_key]
+                        output.setdefault("postprocess_error", repr(e))
+                    await self._output_queue.put((output, data_point))
+
+            def _flush_batch(batch):
+                # Use the overridable hook so subclasses retain custom serialization.
+                self.dump_outputs([o for o, _dp in batch], [dp for _o, dp in batch], fout)
+                pbar.update(len(batch))
+                fout.flush()
+
+            async def _writer():
+                # Block for the first item, then drain and flush all currently queued items as one batch.
+                while True:
+                    item = await self._output_queue.get()
+                    if item is None:
+                        return
+                    batch = [item]
+                    stop = False
+                    # Drain anything that's already enqueued.
+                    while not self._output_queue.empty():
+                        nxt = self._output_queue.get_nowait()
+                        if nxt is None:
+                            stop = True
+                            break
+                        batch.append(nxt)
+                    _flush_batch(batch)
+                    if stop:
+                        return
+
+            processor_tasks = [asyncio.create_task(_processor()) for _ in range(num_processors)]
+            writer_task = asyncio.create_task(_writer())
+
+            # Bound both request concurrency and allocated task state to max_in_flight.
+            # Each task handles its own error so other datapoints continue processing.
+            max_in_flight = self.cfg.max_concurrent_requests
+
+            async def _safe_one(dp):
+                try:
+                    await self._generate_and_save_datapoint(dp, data, fout, pbar)
+                except Exception as e:
+                    pos = dp.get(self.cfg.async_position_key, "?")
+                    LOG.exception("datapoint async_pos=%s crashed", pos)
+                    # Emit one placeholder per failed datapoint to preserve positional alignment.
+                    # Route it through the raw queue so it follows normal postprocessing.
+                    if self._raw_queue is not None:
+                        await self._raw_queue.put(({"generation": "", "generation_error": repr(e)}, dp))
+
+            data_iter = iter(remaining_data_points)
+            in_flight: set = set()
+
+            # Fill the concurrency pool at the configured rate to limit initial connection bursts.
+            ramp_rate = self.cfg.ramp_rate_per_sec
+            if ramp_rate > 0:
+                # Quarter-second batches keep the ramp smooth while filling the pool promptly.
+                ramp_batch = max(50, ramp_rate // 4)
+                ramp_interval = ramp_batch / ramp_rate
+            else:
+                ramp_batch = max_in_flight  # instant prime
+                ramp_interval = 0.0
+
+            primed = 0
+            for dp in data_iter:
+                in_flight.add(asyncio.create_task(_safe_one(dp)))
+                primed += 1
+                if primed >= max_in_flight:
+                    break
+                if ramp_interval > 0 and primed % ramp_batch == 0:
+                    await asyncio.sleep(ramp_interval)
+
+            # Refill one worker slot for each completed task.
+            while in_flight:
+                done, in_flight = await asyncio.wait(
+                    in_flight,
+                    return_when=asyncio.FIRST_COMPLETED,
+                )
+                for _ in done:
+                    try:
+                        dp = next(data_iter)
+                    except StopIteration:
+                        continue
+                    in_flight.add(asyncio.create_task(_safe_one(dp)))
+
+            # Stop processors after all raw outputs are queued, then stop the writer after processors finish.
+            for _ in range(num_processors):
+                await self._raw_queue.put(None)
+            await asyncio.gather(*processor_tasks)
+            await self._output_queue.put(None)
+            await writer_task
+            self._raw_queue = None
+            self._output_queue = None
 
             pbar.close()
 
@@ -912,13 +1038,21 @@ def restore_async_order(self):
         with open(self.cfg.output_file + "-async", "rt", encoding="utf-8") as fin:
             generations = [json.loads(line) for line in fin]
 
-        ordered_generations = [None] * len(generations)
+        # Original positions can be sparse, so size from the maximum and skip missing entries.
+        key = self.cfg.async_position_key
+        positions = [gen[key] for gen in generations if key in gen]
+        ordered_generations = [None] * ((max(positions) + 1) if positions else 0)
         for gen_dict in generations:
-            async_pos = gen_dict.pop(self.cfg.async_position_key)
+            async_pos = gen_dict.pop(key, None)
+            if async_pos is None:
+                LOG.warning("async record missing %s, dropping from reorder", key)
+                continue
             ordered_generations[async_pos] = gen_dict
 
         with open(self.cfg.output_file, "wt", encoding="utf-8") as fout:
             for gen_dict in ordered_generations:
+                if gen_dict is None:
+                    continue
                 fout.write(json.dumps(gen_dict) + "\n")
 
         Path(self.cfg.output_file + "-async").unlink()

nemo_skills/inference/model/azure.py

@@ -19,6 +19,9 @@
 
 class AzureOpenAIModel(OpenAIModel):
     MODEL_PROVIDER = "azure"
+    # Azure requires provider-specific endpoint and authentication handling
+    # from litellm.
+    SUPPORTS_NATIVE_OPENAI = False
 
     def __init__(
         self,