minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[smg][ci] Add thread safety to ModelPool and GPUAllocator (#16674)

Browse Source
This commit is contained in:
Simo Lin
2026-01-07 13:25:41 -08:00
committed by GitHub
parent 0241e0460f
commit 6037267f5b
10 changed files with 534 additions and 195 deletions
Download Patch File Download Diff File Expand all files Collapse all files
sgl-model-gateway/e2e_test/fixtures/setup_backend.py
+49 -37
View File
@@ -130,34 +130,15 @@ def _setup_pd_backend(
import openai
from infra import ConnectionMode, Gateway, WorkerIdentity, WorkerType
# Check PD requirements
try:
import sgl_kernel # noqa: F401
except ImportError:
pytest.skip("sgl_kernel not available, required for PD disaggregation")
try:
import torch
except ImportError:
pytest.skip("torch not available")
if not torch.cuda.is_available():
pytest.skip("CUDA not available")
logger.info("Setting up PD backend for model %s", model_id)
# Get PD configuration from workers marker
num_prefill = workers_config.get("prefill") or 1
num_decode = workers_config.get("decode") or 1
# Check GPU requirements
required_gpus = num_prefill + num_decode
gpu_count = torch.cuda.device_count()
if gpu_count < required_gpus:
pytest.skip(
f"PD tests require {required_gpus} GPUs "
f"({num_prefill} prefill + {num_decode} decode), found {gpu_count}"
)
logger.info("PD config: %d prefill, %d decode workers", num_prefill, num_decode)
# Try to use pre-launched PD workers, or launch additional ones if needed
# get_workers_by_type auto-acquires all returned workers
existing_prefills = model_pool.get_workers_by_type(model_id, WorkerType.PREFILL)
existing_decodes = model_pool.get_workers_by_type(model_id, WorkerType.DECODE)
@@ -168,6 +149,11 @@ def _setup_pd_backend(
if missing_prefill == 0 and missing_decode == 0:
prefills = existing_prefills[:num_prefill]
decodes = existing_decodes[:num_decode]
# Release excess workers we won't use
for w in existing_prefills[num_prefill:]:
w.release()
for w in existing_decodes[num_decode:]:
w.release()
logger.info(
"Using pre-launched PD workers: %d prefill, %d decode",
len(prefills),
@@ -207,17 +193,36 @@ def _setup_pd_backend(
workers_to_launch, startup_timeout=300
)
if not new_instances:
# Release any existing workers we acquired
for w in existing_prefills + existing_decodes:
w.release()
pytest.fail(
f"Failed to launch PD workers: needed {len(workers_to_launch)} workers "
f"but could not allocate GPUs (all in use or timeout)"
)
# Acquire newly launched instances (launch_workers doesn't auto-acquire)
for inst in new_instances:
inst.acquire()
new_prefills = [w for w in new_instances if w.worker_type == WorkerType.PREFILL]
new_decodes = [w for w in new_instances if w.worker_type == WorkerType.DECODE]
prefills = existing_prefills + new_prefills
decodes = existing_decodes + new_decodes
# Acquire references to prevent eviction during test
all_workers = prefills + decodes
for worker in all_workers:
worker.acquire()
# All workers in prefills and decodes are now acquired
model_path = prefills[0].model_path if prefills else None
if not prefills or not decodes:
# This shouldn't happen but guard against it
for w in prefills + decodes:
w.release()
pytest.fail(
f"PD setup incomplete: have {len(prefills)} prefill, {len(decodes)} decode "
f"(need {num_prefill} prefill, {num_decode} decode)"
)
model_path = prefills[0].model_path
# Launch PD gateway
gateway = Gateway()
@@ -250,7 +255,7 @@ def _setup_pd_backend(
logger.info("Tearing down PD gateway")
gateway.shutdown()
# Release references to allow eviction
for worker in all_workers:
for worker in prefills + decodes:
worker.release()
@@ -272,11 +277,20 @@ def _setup_local_backend(
try:
if num_workers > 1:
existing = model_pool.get_workers_by_type(model_id, WorkerType.REGULAR)
existing_for_mode = [w for w in existing if w.mode == connection_mode]
# get_workers_by_type auto-acquires all returned workers
all_existing = model_pool.get_workers_by_type(model_id, WorkerType.REGULAR)
existing_for_mode = [w for w in all_existing if w.mode == connection_mode]
# Release workers we won't use (wrong mode)
for w in all_existing:
if w not in existing_for_mode:
w.release()
if len(existing_for_mode) >= num_workers:
instances = existing_for_mode[:num_workers]
# Release excess workers we won't use
for w in existing_for_mode[num_workers:]:
w.release()
else:
missing = num_workers - len(existing_for_mode)
workers_to_launch = [
@@ -291,6 +305,9 @@ def _setup_local_backend(
new_instances = model_pool.launch_workers(
workers_to_launch, startup_timeout=300
)
# Acquire newly launched instances
for inst in new_instances:
inst.acquire()
instances = existing_for_mode + new_instances
if not instances:
@@ -298,14 +315,11 @@ def _setup_local_backend(
worker_urls = [inst.worker_url for inst in instances]
model_path = instances[0].model_path
else:
# get() auto-acquires the returned instance
instance = model_pool.get(model_id, connection_mode)
instances = [instance]
worker_urls = [instance.worker_url]
model_path = instance.model_path
# Acquire references to prevent eviction during test
for inst in instances:
inst.acquire()
except RuntimeError as e:
pytest.fail(str(e))
@@ -393,15 +407,13 @@ def backend_router(request: pytest.FixtureRequest, model_pool: "ModelPool"):
connection_mode = ConnectionMode(backend_name)
try:
# get() auto-acquires the returned instance
instance = model_pool.get(model_id, connection_mode)
except KeyError:
pytest.skip(f"Model {model_id}:{backend_name} not available in pool")
except RuntimeError as e:
pytest.fail(str(e))
# Acquire reference to prevent eviction during test
instance.acquire()
gateway = Gateway()
gateway.start(
worker_urls=[instance.worker_url],