[Feature] Support --tokenizer-worker-num > 1 in the offline Engine API

python/sglang/srt/entrypoints/engine.py

@@ -21,6 +21,7 @@
 
 import asyncio
 import atexit
+import copy
 import dataclasses
 import gc
 import logging
@@ -31,6 +32,7 @@
 import tempfile
 import threading
 import time
+import types
 from typing import (
     Any,
     AsyncIterator,
@@ -94,7 +96,11 @@
 )
 from sglang.srt.managers.multi_tokenizer_mixin import (
     MultiTokenizerRouter,
+    get_main_process_id,
+    get_tokenizer_worker_class,
+    read_from_shared_memory,
     run_multi_detokenizer_router_process,
+    write_data_for_multi_tokenizer,
 )
 from sglang.srt.managers.scheduler import run_scheduler_process
 from sglang.srt.managers.tokenizer_manager import TokenizerManager
@@ -243,6 +249,11 @@ class Engine(EngineScoreMixin, EngineBase):
     # Backend-specific launch handle: the Ray engine schedules its actors onto a
     # placement group. Not config — a live cluster object.
     _placement_group = None
+    # Set by attach_tokenizer_worker(): this instance only owns a TokenizerWorker,
+    # the engine subprocesses belong to the parent Engine.
+    _attached = False
+    tokenizer_router: Optional[MultiTokenizerRouter] = None
+    _multi_tokenizer_shm = None
 
     def __init__(self, **kwargs):
         """
@@ -316,6 +327,25 @@ def __init__(self, **kwargs):
             tokenizer_manager._subprocess_watchdog = subprocess_watchdog
         self.port_args = port_args
 
+        if isinstance(tokenizer_manager, MultiTokenizerRouter):
+            # --tokenizer-worker-num > 1: the router only bridges TokenizerWorkers to
+            # the scheduler and has no generate_request. Publish the launch data the
+            # workers need (same shared-memory contract as `sglang serve`) and give
+            # this process its own worker so Engine.generate() keeps working here;
+            # more processes join with Engine.attach_tokenizer_worker().
+            self.tokenizer_router = tokenizer_manager
+            scheduler_info = {
+                **scheduler_init_result.scheduler_infos[0],
+                "startup_time": tokenizer_manager.startup_time,
+            }
+            self._multi_tokenizer_shm = write_data_for_multi_tokenizer(
+                port_args, server_args, scheduler_info
+            )
+            self.tokenizer_manager, self.template_manager = self._init_tokenizer_worker(
+                server_args, port_args, scheduler_info
+            )
+            self.tokenizer_manager._subprocess_watchdog = subprocess_watchdog
+
         # Initialize ZMQ sockets
         context = zmq.Context(2)
         if get_parallel().node_rank == 0:
@@ -339,11 +369,7 @@ def __init__(self, **kwargs):
                 thread_label = "Decode Tokenizer"
             trace_set_thread_info(thread_label)
 
-        try:
-            self.loop = asyncio.get_running_loop()
-        except RuntimeError:
-            self.loop = asyncio.new_event_loop()
-            asyncio.set_event_loop(self.loop)
+        self.loop = self._ensure_event_loop()
 
     def get_all_child_pids(self) -> List[int]:
         """Returns a list of all child process PIDs."""
@@ -1270,10 +1296,91 @@ def _launch_subprocesses(
             weight_cache_daemon_procs,
         )
 
+    @staticmethod
+    def _init_tokenizer_worker(
+        server_args: ServerArgs, port_args: PortArgs, scheduler_info: Dict[str, Any]
+    ) -> Tuple[TokenizerManager, TemplateManager]:
+        # TokenizerWorker.__init__ registers with the router over a zmq.asyncio
+        # socket, which needs the loop the Engine will later run to be current.
+        Engine._ensure_event_loop()
+        port_args = copy.copy(port_args)
+        port_args.tokenizer_ipc_name = (
+            f"ipc://{tempfile.NamedTemporaryFile(delete=False).name}"
+        )
+        tokenizer_manager, template_manager = init_tokenizer_manager(
+            server_args,
+            port_args,
+            TokenizerManagerClass=get_tokenizer_worker_class(server_args),
+        )
+        tokenizer_manager.max_req_input_len = scheduler_info["max_req_input_len"]
+        tokenizer_manager.set_startup_time(scheduler_info["startup_time"])
+        return tokenizer_manager, template_manager
+
+    @staticmethod
+    def _ensure_event_loop() -> asyncio.AbstractEventLoop:
+        try:
+            return asyncio.get_running_loop()
+        except RuntimeError:
+            pass
+        try:
+            return asyncio.get_event_loop_policy().get_event_loop()
+        except RuntimeError:
+            loop = asyncio.new_event_loop()
+            asyncio.set_event_loop(loop)
+            return loop
+
+    @classmethod
+    def attach_tokenizer_worker(cls, parent_pid: Optional[int] = None) -> Engine:
+        """Join an Engine launched with ``tokenizer_worker_num > 1`` from another
+        process on the same host.
+
+        The returned Engine has no subprocesses of its own: it holds a
+        ``TokenizerWorker`` registered with the parent's ``MultiTokenizerRouter``, so
+        ``generate``/``async_generate`` run on this process while prefill and decode
+        stay on the parent's schedulers. Use it to spread the per-request Python
+        work of one engine over several processes (what ``sglang serve`` does with
+        uvicorn workers). ``parent_pid`` defaults to the multiprocessing parent, so
+        children spawned by the parent can omit it. Weight updates and other RPCs
+        must go through the parent Engine.
+        """
+        parent_pid = get_main_process_id() if parent_pid is None else parent_pid
+        port_args, server_args, scheduler_info = read_from_shared_memory(
+            f"multi_tokenizer_args_{parent_pid}"
+        )
+        publish(server_args, role="tokenizer")
+        self = cls.__new__(cls)
+        self._attached = True
+        self.server_args = server_args
+        self.port_args = port_args
+        self._scheduler_init_result = types.SimpleNamespace(
+            scheduler_infos=[scheduler_info], all_child_pids=[]
+        )
+        self._weight_cache_daemon_procs = []
+        self.send_to_rpc = None
+        self.tokenizer_manager, self.template_manager = self._init_tokenizer_worker(
+            server_args, port_args, scheduler_info
+        )
+        self.loop = self._ensure_event_loop()
+        return self
+
     def shutdown(self):
         """Shutdown the engine; block until the scheduler subprocess releases
         its GPU context so the caller can immediately reallocate on the same
         device."""
+        if self._attached:
+            if isinstance(self.tokenizer_manager, TokenizerManager):
+                mm_processor = getattr(self.tokenizer_manager, "mm_processor", None)
+                if mm_processor is not None:
+                    mm_processor.shutdown()
+                self.tokenizer_manager.cuda_vmm_feature_transport.shutdown()
+            return
+        shm = getattr(self, "_multi_tokenizer_shm", None)
+        if shm is not None:
+            try:
+                shm.unlink()
+            except FileNotFoundError:
+                pass
+            self._multi_tokenizer_shm = None
         try:
             if (
                 self.tokenizer_manager is not None

test/registered/e2e/entrypoints/test_engine_multi_tokenizer.py

@@ -0,0 +1,67 @@
+"""Engine API with --tokenizer-worker-num > 1 (issue #15157).
+
+The parent Engine keeps generate() working through an in-process TokenizerWorker,
+and other processes join the same engine with Engine.attach_tokenizer_worker().
+"""
+
+import multiprocessing as mp
+import os
+import unittest
+
+import sglang as sgl
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.test_utils import DEFAULT_SMALL_MODEL_NAME_FOR_TEST, CustomTestCase
+
+register_cuda_ci(est_time=180, stage="base-b", runner_config="1-gpu-large")
+
+PROMPTS = ["The capital of France is", "1 + 1 ="]
+SAMPLING = {"temperature": 0, "max_new_tokens": 8}
+
+
+def _attached_generate(parent_pid, prompts, queue):
+    engine = sgl.Engine.attach_tokenizer_worker(parent_pid)
+    try:
+        outs = engine.generate(prompts, SAMPLING)
+        queue.put([(o["text"], o["meta_info"]["completion_tokens"]) for o in outs])
+    finally:
+        engine.shutdown()
+
+
+class TestEngineMultiTokenizer(CustomTestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.engine = sgl.Engine(
+            model_path=DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
+            random_seed=42,
+            tokenizer_worker_num=2,
+        )
+
+    @classmethod
+    def tearDownClass(cls):
+        cls.engine.shutdown()
+
+    def test_1_parent_generate_uses_worker(self):
+        self.assertIsNotNone(self.engine.tokenizer_router)
+        outs = self.engine.generate(PROMPTS, SAMPLING)
+        self.assertEqual(len(outs), len(PROMPTS))
+        self.assertTrue(all(o["meta_info"]["completion_tokens"] > 0 for o in outs))
+
+    def test_2_attached_process_matches_parent(self):
+        expected = [
+            (o["text"], o["meta_info"]["completion_tokens"])
+            for o in self.engine.generate(PROMPTS, SAMPLING)
+        ]
+        ctx = mp.get_context("spawn")
+        queue = ctx.Queue()
+        proc = ctx.Process(
+            target=_attached_generate, args=(os.getpid(), PROMPTS, queue)
+        )
+        proc.start()
+        got = queue.get(timeout=600)
+        proc.join(timeout=60)
+        self.assertEqual(proc.exitcode, 0)
+        self.assertEqual(got, expected)
+
+
+if __name__ == "__main__":
+    unittest.main()