From bccf691b221835689450ef601e73fd6ee79c2ac3 Mon Sep 17 00:00:00 2001 From: Cheng Wan <54331508+ch-wan@users.noreply.github.com> Date: Mon, 21 Sep 2026 12:29:50 -0700 Subject: [PATCH] Bringing the parallel runtime up becomes a phase, not a side effect (#40345) --- python/sglang/benchmark/one_batch.py | 21 +- python/sglang/srt/distributed/bootstrap.py | 187 +++++++++++++----- .../sglang/srt/distributed/parallel_state.py | 3 + python/sglang/srt/managers/scheduler.py | 16 +- .../sglang/srt/model_executor/model_runner.py | 37 +--- .../test_mooncake_transfer_engine_init.py | 7 +- test/manual/test_forward_split_prefill.py | 15 +- test/manual/test_vlm_accuracy.py | 18 +- test/registered/unit/test_runtime_context.py | 76 ++++++- 9 files changed, 271 insertions(+), 109 deletions(-) diff --git a/python/sglang/benchmark/one_batch.py b/python/sglang/benchmark/one_batch.py index a5753df6d..83a2e42ab 100644 --- a/python/sglang/benchmark/one_batch.py +++ b/python/sglang/benchmark/one_batch.py @@ -71,13 +71,13 @@ from sglang.srt.arg_groups.overrides import ( ) from sglang.srt.configs.hybrid_arch import mambaish_config from sglang.srt.configs.model_config import ModelConfig +from sglang.srt.distributed import bootstrap from sglang.srt.distributed.parallel_state import ( destroy_distributed_environment, destroy_model_parallel, ) from sglang.srt.entrypoints.engine import _set_envs_and_config from sglang.srt.hardware_backend.mlx.runtime import use_mlx -from sglang.srt.layers.dp_attention import compute_dp_attention_world_info from sglang.srt.layers.moe import initialize_moe_config from sglang.srt.layers.quantization.fp4_utils import initialize_fp4_gemm_config from sglang.srt.layers.quantization.fp8_utils import initialize_fp8_gemm_config @@ -93,6 +93,7 @@ from sglang.srt.model_executor.forward_batch_info import ForwardBatch from sglang.srt.model_executor.model_runner import ModelRunner from sglang.srt.runtime_context import ( SpawnRanks, + get_device, get_model, get_parallel, get_schedule, @@ -316,18 +317,8 @@ def load_model(server_args, port_args, gpu_id, tp_rank): cfg = resolving_view(server_args) suppress_other_loggers() rank_print = print if tp_rank == 0 else lambda *args, **kwargs: None - moe_ep_rank = tp_rank // (cfg.tp_size // cfg.ep_size) model_config = ModelConfig.from_server_args(server_args) - attn_tp_rank, attn_tp_size, attn_dp_rank, attn_dp_size = ( - compute_dp_attention_world_info( - cfg.enable_dp_attention, - tp_rank, - cfg.tp_size, - cfg.dp_size, - cfg.attn_cp_size, - ) - ) runner_kwargs = dict( model_config=model_config, mem_fraction_static=cfg.mem_fraction_static, @@ -336,6 +327,14 @@ def load_model(server_args, port_args, gpu_id, tp_rank): server_args=server_args, ) + # Phase two: this entry has no scheduler to run it. + bootstrap.init_parallel_runtime( + server_args=server_args, + model_config=model_config, + device=get_device().device, + dist_port=port_args.nccl_port, + ) + _use_mlx = use_mlx() if _use_mlx: from sglang.srt.hardware_backend.mlx.model_runner_stub import ( diff --git a/python/sglang/srt/distributed/bootstrap.py b/python/sglang/srt/distributed/bootstrap.py index 889917d4f..0fedb3e0a 100644 --- a/python/sglang/srt/distributed/bootstrap.py +++ b/python/sglang/srt/distributed/bootstrap.py @@ -7,6 +7,9 @@ import torch import torch.distributed as dist from sglang.srt.configs.model_config import ModelConfig +from sglang.srt.distributed.device_communicators.mooncake_transfer_engine import ( + maybe_init_shared_mooncake_transfer_engine, +) from sglang.srt.distributed.gated_launch import maybe_wait_for_gated_launch from sglang.srt.distributed.parallel_state import ( _tag_groups_for_flashinfer_allreduce_only, @@ -56,92 +59,176 @@ _is_cpu_arm64 = is_host_cpu_arm64() _TP_ALL_TO_ALL_WARMUP_BYTES_PER_PEER = 4 << 20 -def init_torch_distributed( +#: Set by `init_parallel`; `destroy_model_parallel` clears it, so a test that +#: tears the groups down can build them again. +_PARALLEL_INITIALISED = False + + +def reset_parallel_initialised() -> None: + """Forget that the groups were built. Paired with tearing them down.""" + global _PARALLEL_INITIALISED + _PARALLEL_INITIALISED = False + + +def _bind_threads_if_cpu(*, device: str) -> "Optional[List[int]]": + """Pin OpenMP threads to this process's NUMA node, on CPU. + + A precondition of the CPU group build, which reads the binding, so it is + done here rather than left for a caller to remember. + """ + if device != "cpu": + return None + from sglang.srt.utils import numa_utils + + parallel = get_parallel() + # With --enable-dp-attention, dp partitions the existing TP group rather + # than spawning additional processes, so dp_size must not be multiplied + # into the process count here (unlike regular DP, where dp_size * tp_size * + # pp_size is the true worker count). + dp_size = 1 if parallel.enable_dp_attention else parallel.dp_size + return numa_utils.init_threads_binding( + numa_index=get_device().gpu_id, + world_size=dp_size * parallel.tp_size * parallel.pp_size, + ) + + +def init_parallel_runtime( *, server_args: ServerArgs, model_config: ModelConfig, device: str, dist_port: int, - is_draft_worker: bool, - local_omp_cpuid: Optional[List[int]], -): +) -> None: + """Phase two of startup: bring the parallel runtime up, once. + + Publish says what the topology is; this makes it exist. Nothing returns, + because the groups are read through the runtime context -- a caller that + wants one asks `get_parallel()`, in this process or any later phase. + + "Runtime" rather than "groups": two things have to be in place before the + groups can be built, and they are done here rather than left for every + entry to remember. The OpenMP/NUMA binding is what the CPU group build + reads, and the shared Mooncake transfer engine is what the Mooncake + process-group backend asks for -- create that one late and a second engine + appears. Both are preconditions of the build, not separate work. + + Runs on the target worker only. A draft worker shares its target's groups, + which is why this is a phase the entry runs rather than something a runner + does on its way up: whether the groups exist must not depend on which + runner happened to be constructed first. + """ + global _PARALLEL_INITIALISED + if _PARALLEL_INITIALISED: + raise RuntimeError( + "init_parallel_runtime() ran twice in this process. The groups are built " + "once, before anything that reads one exists; a second build is " + "either a lost race between two entries or a runner trying to " + "bring up its own. An elastic scale-up joins an existing WORLD " + "through initialize_model_parallel directly and does not come " + "through here." + ) + _PARALLEL_INITIALISED = True + tic = time.perf_counter() - logger.info("Init torch distributed begin.") - parallel = get_parallel() + logger.info("Init parallel begin.") backend = _resolve_backend(device=device) - - before_avail_memory = get_available_gpu_memory(device, get_device().gpu_id) if not get_parallel().enable_p2p_check: monkey_patch_p2p_access_check() dist_init_method = _resolve_dist_init_method(dist_port=dist_port) _set_all_reduce_flags() - if not is_draft_worker: - if device == "cpu": - _init_cpu_threads_env( - tp_size=parallel.tp_size, - tp_rank=parallel.tp_rank, - local_omp_cpuid=local_omp_cpuid, - dist_init_method=dist_init_method, - ) + local_omp_cpuid = _bind_threads_if_cpu(device=device) + # Everything below allocates on the current device -- the NCCL warm-up, the + # mooncake all-reduce buffer -- and without this every rank on a node would + # pick device 0, because the default is not to reindex the visible set. + try: + torch.get_device_module(device).set_device(get_device().gpu_id) + except Exception: + logger.warning( + "Context: device=%s gpu_id=%s CUDA_VISIBLE_DEVICES=%s tp_rank=%s", + device, + get_device().gpu_id, + os.environ.get("CUDA_VISIBLE_DEVICES"), + get_parallel().tp_rank, + ) + raise + maybe_init_shared_mooncake_transfer_engine(gpu_id=get_device().gpu_id) - # Only initialize the distributed environment on the target model worker. - # This builds the groups behind the context's live group-handle reads. - _init_parallel_groups( - backend=backend, + parallel = get_parallel() + if device == "cpu": + _init_cpu_threads_env( + tp_size=parallel.tp_size, + tp_rank=parallel.tp_rank, + local_omp_cpuid=local_omp_cpuid, dist_init_method=dist_init_method, - server_args=server_args, - model_config=model_config, - gpu_id=get_device().gpu_id, ) - # Pre-warm NCCL/RCCL/HCCL to eliminate cold-start latency in first request - # Controlled by --pre-warm-nccl flag (default: enabled on AMD GPUs) - if get_exec().comm.pre_warm_nccl and ( - parallel.tp_size > 1 or parallel.pp_size > 1 or parallel.moe_ep_size > 1 - ): - _prewarm_nccl( - tp_size=parallel.tp_size, - pp_size=parallel.pp_size, - moe_ep_size=parallel.moe_ep_size, - ) + _init_parallel_groups( + backend=backend, + dist_init_method=dist_init_method, + server_args=server_args, + model_config=model_config, + gpu_id=get_device().gpu_id, + ) - # CUDA graph capture enables the PyNCCL communicator for TP LM-head - # all-to-all. Exercise that exact send/recv path before measuring - # pre_model_load_memory so its persistent transport allocations are - # included in later KV-cache sizing instead of appearing during capture. - if ( - device == "cuda" - and get_parallel().enable_tp_lm_head_all_to_all - and parallel.tp_size > 1 - ): - _prewarm_tp_lm_head_all_to_all() + # Pre-warm NCCL/RCCL/HCCL to eliminate cold-start latency in first request + # Controlled by --pre-warm-nccl flag (default: enabled on AMD GPUs) + if get_exec().comm.pre_warm_nccl and ( + parallel.tp_size > 1 or parallel.pp_size > 1 or parallel.moe_ep_size > 1 + ): + _prewarm_nccl( + tp_size=parallel.tp_size, + pp_size=parallel.pp_size, + moe_ep_size=parallel.moe_ep_size, + ) + + # CUDA graph capture enables the PyNCCL communicator for TP LM-head + # all-to-all. Exercise that exact send/recv path before measuring + # pre_model_load_memory so its persistent transport allocations are + # included in later KV-cache sizing instead of appearing during capture. + if ( + device == "cuda" + and parallel.enable_tp_lm_head_all_to_all + and parallel.tp_size > 1 + ): + _prewarm_tp_lm_head_all_to_all() + + logger.info(f"Init parallel ends. elapsed={time.perf_counter() - tic:.2f} s") + + +def measure_pre_model_load_memory(*, device: str, is_draft_worker: bool) -> float: + """Available memory after the groups exist and before the model loads. + + Sized into the KV cache later, so it has to be taken at exactly this point + -- which is why it stays with the runner rather than moving into the + parallel phase. + """ + before_avail_memory = get_available_gpu_memory(device, get_device().gpu_id) maybe_wait_for_gated_launch( host=get_serving().host, port=get_parallel().gated_launch_port ) - # Draft workers reuse the target pool config and may exist on only one PP stage; - # including them in this WORLD reduction would deadlock on absent peers. + # Draft workers reuse the target pool config and may exist on only one PP + # stage; including them in this WORLD reduction would deadlock on absent + # peers. pre_model_load_memory = get_available_gpu_memory( device, get_device().gpu_id, distributed=get_world_group().world_size > 1 and not is_draft_worker, cpu_group=get_world_group().cpu_group, ) - # Check memory for tensor parallelism local_gpu_memory = get_available_gpu_memory(device, get_device().gpu_id) - if parallel.tp_size > 1 and not is_draft_worker: + if get_parallel().tp_size > 1 and not is_draft_worker: _check_tp_memory_balance( pre_model_load_memory=pre_model_load_memory, local_gpu_memory=local_gpu_memory, ) - logger.info( - f"Init torch distributed ends. elapsed={time.perf_counter() - tic:.2f} s, " - f"mem usage={(before_avail_memory - local_gpu_memory):.2f} GB" + f"Memory baseline taken. mem usage=" + f"{(before_avail_memory - local_gpu_memory):.2f} GB" ) return pre_model_load_memory diff --git a/python/sglang/srt/distributed/parallel_state.py b/python/sglang/srt/distributed/parallel_state.py index e9e1e7f97..fcc38511d 100644 --- a/python/sglang/srt/distributed/parallel_state.py +++ b/python/sglang/srt/distributed/parallel_state.py @@ -3240,6 +3240,9 @@ def get_moe_tensor_parallel_rank(): def destroy_model_parallel(): """Set the groups to none and destroy them.""" + from sglang.srt.distributed.bootstrap import reset_parallel_initialised + + reset_parallel_initialised() get_parallel().clear_stamp() dwdp_mgr = get_global_dwdp_manager() if dwdp_mgr is not None: diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index f92b88835..51d3c4633 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -101,6 +101,7 @@ from sglang.srt.disaggregation.utils import ( prepare_abort, unified_memory_disagg_move_gate, ) +from sglang.srt.distributed import bootstrap from sglang.srt.distributed.parallel_state import ( abort_distributed_environment, ) @@ -513,6 +514,18 @@ class Scheduler( # Init model configs self.init_model_config() + # Init ZBAL, switch allocator should before any torch alloc action + self.init_zbal_on_npu() + + # The groups are the first thing that allocates, so this comes after the + # allocator switch above and before anything that reads a group. + bootstrap.init_parallel_runtime( + server_args=server_args, + model_config=self.model_config, + device=get_device().device, + dist_port=self.nccl_port, + ) + # Init metrics stats self.init_metrics_collector(tp_rank, pp_rank, dp_rank) @@ -520,9 +533,6 @@ class Scheduler( self.init_ipc_channels(port_args) self.init_idle_sleeper() - # Init ZBAL, switch allocator should before any torch alloc action - self.init_zbal_on_npu() - # Init PD-multiplexing context if self.enable_pdmux: self.init_pdmux() diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 517ced8a8..aba114416 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -34,9 +34,6 @@ from sglang.srt.configs.model_config import ( from sglang.srt.configs.update_config import adjust_config_with_unaligned_cpu_tp from sglang.srt.debug_utils.dumper import dumper from sglang.srt.distributed import bootstrap -from sglang.srt.distributed.device_communicators.mooncake_transfer_engine import ( - maybe_init_shared_mooncake_transfer_engine, -) from sglang.srt.dllm.config import DllmConfig from sglang.srt.elastic_ep.elastic_ep import ( ElasticEPStateManager, @@ -218,7 +215,6 @@ from sglang.srt.utils import ( get_available_gpu_memory, is_host_cpu_arm64, is_npu, - numa_utils, require_gathered_buffer, reserve_rope_cache_for_long_sequences, set_cuda_arch, @@ -402,10 +398,6 @@ class ModelRunner: is_draft_worker=self.is_draft_worker, ) - # Init OpenMP threads binding for CPU - if self.device == "cpu": - self.init_threads_binding() - # Set float32 matmul precision if get_exec().features.enable_tf32_matmul: torch.set_float32_matmul_precision("high") @@ -422,11 +414,6 @@ class ModelRunner: ) raise - # Initialize MooncakeTransferEngine BEFORE init_torch_distributed so - # that the shared TE can be passed to the Mooncake PG backend (avoids - # creating duplicate TransferEngines). - self.init_shared_mooncake_transfer_engine() - # Get available memory before model loading. # Stored for later use by alloc_memory_pool(). self.init_torch_distributed() @@ -1167,13 +1154,8 @@ class ModelRunner: ) def init_torch_distributed(self): - self.pre_model_load_memory = bootstrap.init_torch_distributed( - server_args=self.server_args, - model_config=self.model_config, - device=self.device, - dist_port=self.dist_port, - is_draft_worker=self.is_draft_worker, - local_omp_cpuid=self.local_omp_cpuid if self.device == "cpu" else None, + self.pre_model_load_memory = bootstrap.measure_pre_model_load_memory( + device=self.device, is_draft_worker=self.is_draft_worker ) # Read once, here: a draft runner is constructed inside the scope that # states its topology and used outside it, so what it holds has to be @@ -1197,9 +1179,6 @@ class ModelRunner: self.moe_ep_size = parallel.moe_ep_size self.dp_rank = parallel.dp_rank - def init_shared_mooncake_transfer_engine(self): - maybe_init_shared_mooncake_transfer_engine(gpu_id=self.gpu_id) - def load_model(self): tic_total = time.perf_counter() before_avail_memory = get_available_gpu_memory(self.device, self.gpu_id) @@ -1593,18 +1572,6 @@ class ModelRunner: phases=("prefill", "draft_prefill"), ) - def init_threads_binding(self): - # With --enable-dp-attention, dp partitions the existing TP group - # rather than spawning additional processes, so dp_size must not be - # multiplied into the process count here (unlike regular DP, where - # dp_size * tp_size * pp_size is the true worker count). - parallel = get_parallel() - dp_size = 1 if parallel.enable_dp_attention else parallel.dp_size - self.local_omp_cpuid = numa_utils.init_threads_binding( - numa_index=self.gpu_id, - world_size=dp_size * parallel.tp_size * parallel.pp_size, - ) - def apply_torch_tp(self): model_parallel.apply_torch_tp( model=self.model, device=self.device, tp_size=self.tp_size diff --git a/test/manual/kv_transfer/test_mooncake_transfer_engine_init.py b/test/manual/kv_transfer/test_mooncake_transfer_engine_init.py index bcffda5a0..c121ceee6 100755 --- a/test/manual/kv_transfer/test_mooncake_transfer_engine_init.py +++ b/test/manual/kv_transfer/test_mooncake_transfer_engine_init.py @@ -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 diff --git a/test/manual/test_forward_split_prefill.py b/test/manual/test_forward_split_prefill.py index 7ec1d3d31..68e99672c 100644 --- a/test/manual/test_forward_split_prefill.py +++ b/test/manual/test_forward_split_prefill.py @@ -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, diff --git a/test/manual/test_vlm_accuracy.py b/test/manual/test_vlm_accuracy.py index 0e084f442..f35ee1d5d 100644 --- a/test/manual/test_vlm_accuracy.py +++ b/test/manual/test_vlm_accuracy.py @@ -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, diff --git a/test/registered/unit/test_runtime_context.py b/test/registered/unit/test_runtime_context.py index 08616231f..6cf582ad9 100644 --- a/test/registered/unit/test_runtime_context.py +++ b/test/registered/unit/test_runtime_context.py @@ -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.