feat(sglang): asynchronous support

Author: guicho271828

mellea/backends/sglang.py

@@ -6,8 +6,10 @@
 from __future__ import annotations
 
 import abc
+import asyncio
 import dataclasses
 import datetime
+import functools
 import inspect
 import json
 import os
@@ -16,7 +18,7 @@
 from typing import TYPE_CHECKING, Any, Optional
 
 import sglang as sgl  # type:ignore
-import torch
+from sglang.utils import async_stream_and_merge  # type:ignore
 from transformers import AutoTokenizer, PreTrainedTokenizerBase
 
 from mellea.backends import BaseModelSubclass
@@ -29,12 +31,14 @@
 )
 from mellea.backends.types import ModelOption
 from mellea.backends.utils import extract_model_tool_requests, to_chat
+from mellea.helpers.async_helpers import send_to_queue
 from mellea.helpers.fancy_logger import FancyLogger
 from mellea.stdlib.base import (
     CBlock,
     Component,
     Context,
     GenerateLog,
+    GenerateType,
     ModelOutputThunk,
     TemplateRepresentation,
 )
@@ -145,7 +149,6 @@ def _generate_from_context_standard(
         # Construct input.
         # If the Context is a ChatHistory then we will pretty-print each content as a message and then use apply_chat_template.
         # Otherwise, we will linearize the context and treat it as a raw input.
-        decoded_result: str | None = None
 
         if ctx.is_chat_context:
             system_prompt = model_options.get(ModelOption.SYSTEM_PROMPT, None)
@@ -182,42 +185,100 @@ def _generate_from_context_standard(
             if format is not None:
                 sampling_params["json_schema"] = json.dumps(format.model_json_schema())
 
-            output: dict[str, Any] = self._model.generate(  # type: ignore
-                input_str, sampling_params=sampling_params
-            )  # type: ignore
+            if model_options.get(ModelOption.STREAM, False):
+                generator = async_stream_and_merge(
+                    self._model,  # type: ignore
+                    input_str,
+                    sampling_params=sampling_params,
+                )  # type: ignore
+            else:
+                generator = self._model.async_generate(  # type: ignore
+                    input_str, sampling_params=sampling_params
+                )  # type: ignore
+
+            output = ModelOutputThunk(None)
+            output._context = ctx.render_for_generation()
+            output._action = action
+            output._model_options = model_options
+
+            output._process = self.processing
+            output._post_process = functools.partial(
+                self.post_processing,
+                conversation=ctx_as_chat,
+                tool_calls=tool_calls,
+                tools=tools,
+                seed=model_options.get(ModelOption.SEED, None),
+            )
+
+            try:
+                # This function should always be called from a running event loop so we don't have to worry about
+                # scheduling the task to a specific event loop here.
+                output._generate = asyncio.create_task(
+                    send_to_queue(generator, output._async_queue)  # type: ignore
+                )
+                output._generate_type = GenerateType.ASYNC
+            except RuntimeError as e:
+                # Most likely cause is running this function without an event loop present.
+                raise e
 
-            decoded_result = output["text"]
+            return output
 
         else:
             raise Exception("Does not yet support non-chat contexts.")
 
-        assert decoded_result is not None
+    async def processing(self, mot: ModelOutputThunk, chunk: str | dict[str, Any]):
+        """Process the returned chunks or the complete response."""
+
+        if isinstance(chunk, str):  # via async_stream_and_merge
+            if mot._underlying_value is None:
+                mot._underlying_value = ""
+            mot._underlying_value += chunk
+        else:
+            mot._underlying_value = chunk["text"]
+
+    async def post_processing(
+        self,
+        mot: ModelOutputThunk,
+        conversation: list[dict],
+        tool_calls: bool,
+        tools: dict[str, Callable],
+        seed,
+    ):
+        """Called when generation is done."""
 
-        result = ModelOutputThunk(value=decoded_result)
+        # The ModelOutputThunk must be computed by this point.
+        assert mot.value is not None
 
-        # Only scan for tools if we are not doing structured decoding and tool calls were provided to the model.
+        # Only scan for tools if we are not doing structured output and tool calls were provided to the model.
         if format is None and tool_calls:
-            result.tool_calls = extract_model_tool_requests(tools, decoded_result)
+            mot.tool_calls = extract_model_tool_requests(tools, mot.value)
 
-        parsed_result = self.formatter.parse(action, result)
-        if generate_logs is not None:
-            assert isinstance(generate_logs, list)
-            generate_log = GenerateLog()
-            generate_log.prompt = ctx_as_chat
-            generate_log.backend = f"sglang::{self.model_id!s}"
-            generate_log.model_options = model_options
-            generate_log.date = datetime.datetime.now()
-            generate_log.model_output = decoded_result
-            generate_log.extra = {
-                "format": format,
-                "tools_available": tools,
-                "tools_called": result.tool_calls,
-                "seed": model_options.get(ModelOption.SEED, None),
-            }
-            generate_log.action = action
-            generate_log.result = parsed_result
-            generate_logs.append(generate_log)
-        return parsed_result
+        assert mot._action is not None, (
+            "ModelOutputThunks should have their action assigned during generation"
+        )
+        assert mot._model_options is not None, (
+            "ModelOutputThunks should have their model_opts assigned during generation"
+        )
+
+        self.formatter.parse(mot._action, mot)
+
+        # Generate the log for this ModelOutputThunk.
+        generate_log = GenerateLog()
+        generate_log.prompt = conversation
+        generate_log.backend = f"sglang::{self.model_id!s}"
+        generate_log.model_options = mot._model_options
+        generate_log.date = datetime.datetime.now()
+        generate_log.model_output = mot.value
+        generate_log.extra = {
+            "format": format,
+            "tools_available": tools,
+            "tools_called": mot.tool_calls,
+            "seed": seed,
+        }
+        generate_log.action = mot._action
+        generate_log.result = mot
+
+        mot._generate_log = generate_log
 
     def _generate_from_raw(
         self,