Add EngineFragment

.github/workflows/pr-test.yml

@@ -162,6 +162,12 @@ jobs:
           cd test/srt
           python3 test_update_weights_from_distributed.py
 
+      - name: Test EngineFragment
+        timeout-minutes: 10
+        run: |
+          cd test/srt
+          python3 test_fragment.py
+
       - name: Test expert parallelism (EP=2)
         timeout-minutes: 10
         run: |

docs/backend/function_calling.ipynb

@@ -426,7 +426,7 @@
     "from sglang.srt.managers.io_struct import Tool, Function\n",
     "\n",
     "llm = sgl.Engine(model_path=\"meta-llama/Meta-Llama-3.1-8B-Instruct\")\n",
-    "tokenizer = llm.tokenizer_manager.tokenizer\n",
+    "tokenizer = llm.orchestrator.tokenizer\n",
     "input_ids = tokenizer.apply_chat_template(\n",
     "    messages, tokenize=True, add_generation_prompt=True, tools=tools\n",
     ")\n",

examples/runtime/engine/offline_batch_inference_torchrun.py

@@ -0,0 +1,80 @@
+import datetime
+import os
+import sys
+
+from sglang.srt.entrypoints.engine_fragment import EngineFragment
+
+
+def run():
+    """
+    Example command:
+    ```
+    torchrun --nproc_per_node=4 offline_batch_inference_torchrun.py
+    ```
+    """
+
+    local_rank = int(os.environ["LOCAL_RANK"])
+    rank = int(os.environ["RANK"])
+    world_size = int(os.environ["WORLD_SIZE"])
+
+    def _log(text):
+        t = datetime.datetime.now().strftime("%H:%M:%S")
+        print(f"[{t}] [rank={rank}] {text}")
+
+    _log(
+        f'start {local_rank=} {rank=} {world_size=} {sys.argv=} {os.environ.get("CUDA_VISIBLE_DEVICES")}'
+    )
+
+    tp_size = world_size
+    tp_rank = rank
+    _log(f"{tp_rank=} {tp_size=}")
+
+    model_name, mem_fraction_static = "meta-llama/Llama-3.2-1B-Instruct", 0.1
+    # model_name, mem_fraction_static = "meta-llama/Llama-3.1-70B-Instruct", 0.9 # test large models
+
+    # TODO remove this in next PR
+    for k in [
+        "GROUP_RANK",
+        "GROUP_WORLD_SIZE",
+        "LOCAL_RANK",
+        "LOCAL_WORLD_SIZE",
+        "MASTER_ADDR",
+        "MASTER_PORT",
+        "OMP_NUM_THREADS",
+        "RANK",
+        "ROLE_NAME",
+        "ROLE_RANK",
+        "ROLE_WORLD_SIZE",
+        "TORCHELASTIC_ERROR_FILE",
+        "TORCHELASTIC_MAX_RESTARTS",
+        "TORCHELASTIC_RESTART_COUNT",
+        "TORCHELASTIC_RUN_ID",
+        "TORCHELASTIC_USE_AGENT_STORE",
+        "TORCH_NCCL_ASYNC_ERROR_HANDLING",
+        "WORLD_SIZE",
+    ]:
+        if k in os.environ:
+            del os.environ[k]
+
+    fragment = EngineFragment(
+        model_path=model_name,
+        mem_fraction_static=mem_fraction_static,
+        tp_size=tp_size,
+        tp_rank=tp_rank,
+        nccl_port=23456,
+        gpu_id=tp_rank,
+    )
+    _log(f"{fragment=}")
+
+    output = fragment.generate(
+        prompt=["1+1=2, 1+2=3, 1+3=4, 1+4=", "9-1=8, 8-1=7, 7-1="],
+        sampling_params=dict(max_new_tokens=16, temperature=0.0),
+    )
+    _log(f"{output=}")
+
+    fragment.shutdown()
+    _log(f"End script")
+
+
+if __name__ == "__main__":
+    run()

python/sglang/bench_one_batch.py

@@ -57,11 +57,11 @@
 import torch.distributed as dist
 
 from sglang.srt.configs.model_config import ModelConfig
-from sglang.srt.entrypoints.engine import _set_envs_and_config
 from sglang.srt.hf_transformers_utils import get_tokenizer
 from sglang.srt.managers.schedule_batch import Req, ScheduleBatch
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
 from sglang.srt.model_executor.model_runner import ModelRunner
+from sglang.srt.orchestration.std.launcher import _set_envs_and_config
 from sglang.srt.sampling.sampling_params import SamplingParams
 from sglang.srt.server_args import PortArgs, ServerArgs
 from sglang.srt.speculative.spec_info import SpeculativeAlgorithm