Bringing the parallel runtime up becomes a phase, not a side effect (#40345)
This commit is contained in:
@@ -45,7 +45,6 @@ def test_mooncake_te_condition(server_args: ServerArgs) -> bool:
|
||||
"""
|
||||
Test the condition logic for using MooncakeTransferEngine.
|
||||
"""
|
||||
from sglang.srt.model_executor.model_runner import ModelRunner
|
||||
|
||||
dummy_runner = SimpleNamespace(server_args=server_args, gpu_id=0)
|
||||
init_called = False
|
||||
@@ -69,7 +68,11 @@ def test_mooncake_te_condition(server_args: ServerArgs) -> bool:
|
||||
return_value="127.0.0.1",
|
||||
),
|
||||
):
|
||||
ModelRunner.init_shared_mooncake_transfer_engine(dummy_runner)
|
||||
from sglang.srt.distributed.device_communicators.mooncake_transfer_engine import (
|
||||
maybe_init_shared_mooncake_transfer_engine,
|
||||
)
|
||||
|
||||
maybe_init_shared_mooncake_transfer_engine(gpu_id=dummy_runner.gpu_id)
|
||||
|
||||
return init_called
|
||||
|
||||
|
||||
@@ -15,6 +15,7 @@ import torch
|
||||
|
||||
from sglang.benchmark.one_batch import TreeCacheNamespace
|
||||
from sglang.srt.configs.model_config import ModelConfig
|
||||
from sglang.srt.distributed import bootstrap
|
||||
from sglang.srt.managers.schedule_batch import Req, ScheduleBatch
|
||||
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
|
||||
from sglang.srt.model_executor.forward_context import (
|
||||
@@ -22,7 +23,7 @@ from sglang.srt.model_executor.forward_context import (
|
||||
set_forward_context,
|
||||
)
|
||||
from sglang.srt.model_executor.model_runner import ModelRunner
|
||||
from sglang.srt.runtime_context import publish
|
||||
from sglang.srt.runtime_context import SpawnRanks, publish
|
||||
from sglang.srt.sampling.sampling_params import SamplingParams
|
||||
from sglang.srt.server_args import PortArgs, ServerArgs
|
||||
from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
|
||||
@@ -56,10 +57,20 @@ class TestForwardSplitPrefill(CustomTestCase):
|
||||
|
||||
cls.port_args = PortArgs.init_new(cls.server_args)
|
||||
|
||||
publish(cls.server_args, role="scheduler")
|
||||
publish(
|
||||
cls.server_args,
|
||||
role="scheduler",
|
||||
ranks=SpawnRanks(world_rank=0, gpu_id=0),
|
||||
)
|
||||
|
||||
# Load model and tokenizer
|
||||
cls.model_config = ModelConfig.from_server_args(cls.server_args)
|
||||
bootstrap.init_parallel_runtime(
|
||||
server_args=cls.server_args,
|
||||
model_config=cls.model_config,
|
||||
device=cls.device,
|
||||
dist_port=cls.port_args.nccl_port,
|
||||
)
|
||||
cls.model_runner = ModelRunner(
|
||||
model_config=cls.model_config,
|
||||
mem_fraction_static=cls.server_args.mem_fraction_static,
|
||||
|
||||
@@ -9,6 +9,7 @@ import torch.nn.functional as F
|
||||
from transformers import AutoModel, AutoProcessor, AutoTokenizer
|
||||
|
||||
from sglang.srt.configs.model_config import ModelConfig
|
||||
from sglang.srt.distributed import bootstrap
|
||||
from sglang.srt.entrypoints.openai.protocol import ChatCompletionRequest
|
||||
from sglang.srt.managers.mm_utils import embed_mm_inputs, init_mm_embedding_cache
|
||||
from sglang.srt.managers.schedule_batch import (
|
||||
@@ -19,7 +20,7 @@ from sglang.srt.managers.schedule_batch import (
|
||||
from sglang.srt.model_executor.model_runner import ModelRunner
|
||||
from sglang.srt.multimodal.processors.base_processor import BaseMultimodalProcessor
|
||||
from sglang.srt.parser.conversation import generate_chat_conv
|
||||
from sglang.srt.runtime_context import publish
|
||||
from sglang.srt.runtime_context import SpawnRanks, get_device, publish
|
||||
from sglang.srt.server_args import ServerArgs
|
||||
from sglang.test.test_utils import download_image_with_retry
|
||||
|
||||
@@ -145,9 +146,20 @@ class VisionLLMLogitsBase(unittest.IsolatedAsyncioTestCase):
|
||||
model_path=self.model_path,
|
||||
disable_cuda_graph=True,
|
||||
)
|
||||
publish(server_args, role="scheduler")
|
||||
publish(
|
||||
server_args,
|
||||
role="scheduler",
|
||||
ranks=SpawnRanks(world_rank=0, gpu_id=0),
|
||||
)
|
||||
model_config = ModelConfig(self.model_path, model_override_args="{}")
|
||||
bootstrap.init_parallel_runtime(
|
||||
server_args=server_args,
|
||||
model_config=model_config,
|
||||
device=get_device().device,
|
||||
dist_port=12435,
|
||||
)
|
||||
self.model_runner = ModelRunner(
|
||||
model_config=ModelConfig(self.model_path, model_override_args="{}"),
|
||||
model_config=model_config,
|
||||
mem_fraction_static=0.8,
|
||||
gpu_id=0,
|
||||
nccl_port=12435,
|
||||
|
||||
@@ -2514,9 +2514,8 @@ class TestAnEntryThatBuildsARunnerHandsOverItsPlacement(CustomTestCase):
|
||||
def test_every_publisher_that_builds_a_runner_passes_a_bundle(self):
|
||||
import ast as _ast
|
||||
|
||||
root = _pathlib.Path(next(iter(_sglang.__path__))).resolve()
|
||||
offenders = []
|
||||
for path in root.rglob("*.py"):
|
||||
for path in _sources():
|
||||
text = path.read_text(encoding="utf-8-sig")
|
||||
if "ModelRunner(" not in text or "publish(" not in text:
|
||||
continue
|
||||
@@ -2535,7 +2534,7 @@ class TestAnEntryThatBuildsARunnerHandsOverItsPlacement(CustomTestCase):
|
||||
and getattr(node.func, "id", None) == "publish"
|
||||
and not any(kw.arg == "ranks" for kw in node.keywords)
|
||||
):
|
||||
offenders.append(f"{path.relative_to(root)}:{node.lineno}")
|
||||
offenders.append(f"{path}:{node.lineno}")
|
||||
self.assertEqual(
|
||||
offenders,
|
||||
[],
|
||||
@@ -2929,6 +2928,77 @@ class TestTheTopologyIdentities(CustomTestCase):
|
||||
self.assertEqual(get_parallel().attn_tp_size, 2)
|
||||
|
||||
|
||||
class TestTheParallelPhase(CustomTestCase):
|
||||
"""Publish says what the topology is; one phase builds it, once.
|
||||
|
||||
Before this, whichever runner was constructed first brought the groups up
|
||||
on its way past, so whether they existed depended on construction order --
|
||||
and a draft runner, which must not build them, went down the same path.
|
||||
"""
|
||||
|
||||
def test_building_twice_is_refused(self):
|
||||
from sglang.srt.distributed import bootstrap
|
||||
|
||||
bootstrap.reset_parallel_initialised()
|
||||
self.addCleanup(bootstrap.reset_parallel_initialised)
|
||||
with (
|
||||
patch.object(bootstrap, "_resolve_backend", return_value="gloo"),
|
||||
patch.object(bootstrap, "_resolve_dist_init_method", return_value="env://"),
|
||||
patch.object(bootstrap, "_set_all_reduce_flags"),
|
||||
patch.object(bootstrap, "_init_parallel_groups"),
|
||||
patch.object(bootstrap, "monkey_patch_p2p_access_check"),
|
||||
patch.object(bootstrap, "_init_cpu_threads_env"),
|
||||
patch.object(bootstrap, "_bind_threads_if_cpu", return_value=None),
|
||||
patch.object(bootstrap, "maybe_init_shared_mooncake_transfer_engine"),
|
||||
):
|
||||
reset_context()
|
||||
self.addCleanup(reset_context)
|
||||
publish(
|
||||
ServerArgs(model_path="dummy"),
|
||||
role="test",
|
||||
ranks=SpawnRanks(world_rank=0),
|
||||
)
|
||||
kwargs = dict(
|
||||
server_args=ServerArgs(model_path="dummy"),
|
||||
model_config=None,
|
||||
device="cpu",
|
||||
dist_port=12345,
|
||||
)
|
||||
bootstrap.init_parallel_runtime(**kwargs)
|
||||
with self.assertRaises(RuntimeError) as caught:
|
||||
bootstrap.init_parallel_runtime(**kwargs)
|
||||
self.assertIn("ran twice", str(caught.exception))
|
||||
|
||||
def test_every_publisher_that_builds_a_runner_runs_the_phase(self):
|
||||
"""The companion to the bundle census: an entry that publishes and then
|
||||
builds a runner has to bring the parallel runtime up
|
||||
itself, because the runner no longer does it on the way past."""
|
||||
import ast as _ast
|
||||
|
||||
offenders = []
|
||||
for path in _sources():
|
||||
text = path.read_text(encoding="utf-8-sig")
|
||||
if "ModelRunner(" not in text or "publish(" not in text:
|
||||
continue
|
||||
tree = _ast.parse(text)
|
||||
builds = any(
|
||||
isinstance(n, _ast.Call)
|
||||
and getattr(n.func, "id", getattr(n.func, "attr", None))
|
||||
== "ModelRunner"
|
||||
for n in _ast.walk(tree)
|
||||
)
|
||||
if not builds:
|
||||
continue
|
||||
if "init_parallel_runtime(" not in text:
|
||||
offenders.append(str(path))
|
||||
self.assertEqual(
|
||||
offenders,
|
||||
[],
|
||||
"these publish and then build a ModelRunner without bringing the "
|
||||
"parallel runtime up first:\n " + "\n ".join(offenders),
|
||||
)
|
||||
|
||||
|
||||
class TestWhoAnswersDuringADraftScope(CustomTestCase):
|
||||
"""A draft worker runs in one process with the target, under a scope.
|
||||
|
||||
|
||||
Reference in New Issue
Block a user