[CI] Restore SMG e2e on 2-gpu-h100 / 4-gpu-h100 runners (#24222)
Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
Kangyan-Zhou
co-authored by
Claude Opus 4.7
Claude Opus 4.7
parent
b939d5410f
commit
2e72a36420
@@ -88,7 +88,6 @@ def pytest_collection_modifyitems(
|
||||
|
||||
from infra import (
|
||||
DEFAULT_MODEL,
|
||||
LOG_SEPARATOR_WIDTH,
|
||||
MODEL_SPECS,
|
||||
PARAM_MODEL,
|
||||
PARAM_SETUP_BACKEND,
|
||||
@@ -96,6 +95,8 @@ def pytest_collection_modifyitems(
|
||||
WorkerType,
|
||||
)
|
||||
|
||||
available_gpus = _count_gpus_without_cuda()
|
||||
|
||||
def track_worker(
|
||||
model_id: str, mode: ConnectionMode, worker_type: WorkerType, count: int
|
||||
) -> None:
|
||||
@@ -214,6 +215,27 @@ def pytest_collection_modifyitems(
|
||||
_max_test_gpu_requirement = test_gpus
|
||||
_max_test_name = item.nodeid
|
||||
|
||||
# Mark over-capacity tests as skipped (including when available_gpus
|
||||
# is 0) so pytest_collection_finish can detect the all-skipped case
|
||||
# and fail loudly instead of passing green with zero tests run.
|
||||
if test_gpus > available_gpus:
|
||||
item.add_marker(
|
||||
pytest.mark.skip(
|
||||
reason=(
|
||||
f"requires {test_gpus} GPUs (model={model_id}, "
|
||||
f"tp={MODEL_SPECS.get(model_id, {}).get('tp', 1)}); "
|
||||
f"only {available_gpus} available on this runner"
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
# Prune workers that can never launch on this runner.
|
||||
for key in list(_worker_counts.keys()):
|
||||
spec = MODEL_SPECS.get(key[0], {})
|
||||
if spec.get("tp", 1) > available_gpus:
|
||||
del _worker_counts[key]
|
||||
_first_seen_order[:] = [k for k in _first_seen_order if k in _worker_counts]
|
||||
|
||||
# Log results
|
||||
if _worker_counts:
|
||||
summary = []
|
||||
@@ -233,37 +255,6 @@ def pytest_collection_modifyitems(
|
||||
else:
|
||||
logger.info("Scanned worker requirements: (none)")
|
||||
|
||||
# TEMPORARY: skip every test that would launch an sglang worker subprocess.
|
||||
#
|
||||
# Workers crash at import inside transformers.integrations.hub_kernels →
|
||||
# kernels.deps with:
|
||||
# StrictDataclassFieldValidationError: Validation error for field
|
||||
# 'import_name': TypeError: Unsupported type for field 'import_name': str | None
|
||||
#
|
||||
# The package combo (kernels==0.13.0 + huggingface_hub==1.12.2 +
|
||||
# transformers==5.6.0) is NOT the bug — it imports cleanly in fresh
|
||||
# python:3.10-slim and lmsysorg/sglang:dev containers. The crash is specific
|
||||
# to the 4-gpu-a10 runner image (most likely a stale/partial huggingface_hub
|
||||
# install where _BASIC_TYPE_VALIDATORS[types.UnionType] registration is
|
||||
# missing). Remove this skip once the runner image is rebuilt.
|
||||
#
|
||||
# Filter on `model_pool` in fixturenames (covers setup_backend, model_client,
|
||||
# model_base_url, backend_router — all transitively depend on model_pool) so
|
||||
# tests without an explicit `@pytest.mark.e2e` marker are also skipped. Then
|
||||
# clear scanned worker requirements so model_pool — if realized for any
|
||||
# non-skipped fixture — starts empty and never spawns a worker.
|
||||
skip_marker = pytest.mark.skip(
|
||||
reason="worker-dependent tests disabled: SMG runner image crash on transformers.integrations.hub_kernels import"
|
||||
)
|
||||
for item in items:
|
||||
if (
|
||||
item.get_closest_marker("e2e") is not None
|
||||
or "model_pool" in item.fixturenames
|
||||
):
|
||||
item.add_marker(skip_marker)
|
||||
_worker_counts.clear()
|
||||
_first_seen_order.clear()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Pool requirements
|
||||
@@ -316,9 +307,22 @@ def get_pool_requirements() -> list["WorkerIdentity"]:
|
||||
def _count_gpus_without_cuda() -> int:
|
||||
"""Count available GPUs without initializing CUDA.
|
||||
|
||||
Uses nvidia-smi to avoid CUDA initialization, which is critical for
|
||||
pytest-parallel compatibility. CUDA cannot be re-initialized after a fork.
|
||||
Must avoid CUDA initialization because pytest_collection_modifyitems
|
||||
runs before pytest-parallel forks workers, and CUDA cannot be
|
||||
re-initialized after fork.
|
||||
|
||||
Honors CUDA_VISIBLE_DEVICES first — container runners commonly expose
|
||||
all host GPUs to the container (e.g. NVIDIA_VISIBLE_DEVICES=all) and
|
||||
gate per-process visibility via CUDA_VISIBLE_DEVICES, so nvidia-smi
|
||||
would over-report. Falls back to nvidia-smi only when the env var is
|
||||
unset, and logs (rather than swallows) any nvidia-smi failure so a
|
||||
misconfigured runner is debuggable from CI logs.
|
||||
"""
|
||||
cvd = os.environ.get("CUDA_VISIBLE_DEVICES")
|
||||
if cvd is not None:
|
||||
# CUDA treats "-1" as "no devices"; don't count it as one.
|
||||
return len([d for d in cvd.split(",") if d.strip() and d.strip() != "-1"])
|
||||
|
||||
import subprocess
|
||||
|
||||
try:
|
||||
@@ -328,11 +332,29 @@ def _count_gpus_without_cuda() -> int:
|
||||
text=True,
|
||||
timeout=10,
|
||||
)
|
||||
if result.returncode == 0:
|
||||
return len([line for line in result.stdout.strip().split("\n") if line])
|
||||
except (subprocess.SubprocessError, FileNotFoundError, OSError):
|
||||
pass
|
||||
return 0
|
||||
except FileNotFoundError:
|
||||
logger.error(
|
||||
"nvidia-smi not found and CUDA_VISIBLE_DEVICES is unset; "
|
||||
"cannot determine GPU count, treating as 0"
|
||||
)
|
||||
return 0
|
||||
except (subprocess.SubprocessError, OSError) as e:
|
||||
logger.error(
|
||||
"nvidia-smi failed (%s); cannot determine GPU count, treating as 0",
|
||||
e,
|
||||
exc_info=True,
|
||||
)
|
||||
return 0
|
||||
|
||||
if result.returncode != 0:
|
||||
logger.error(
|
||||
"nvidia-smi exited with code %d; treating as 0 GPUs. stderr=%r stdout=%r",
|
||||
result.returncode,
|
||||
result.stderr,
|
||||
result.stdout,
|
||||
)
|
||||
return 0
|
||||
return len([line for line in result.stdout.strip().split("\n") if line])
|
||||
|
||||
|
||||
def validate_gpu_requirements() -> tuple[int, int]:
|
||||
@@ -350,9 +372,11 @@ def validate_gpu_requirements() -> tuple[int, int]:
|
||||
|
||||
def pytest_collection_finish(session: pytest.Session) -> None:
|
||||
"""Validate GPU requirements after test collection."""
|
||||
from infra import ENV_SKIP_MODEL_POOL, LOG_SEPARATOR_WIDTH
|
||||
from infra import ENV_SKIP_MODEL_POOL
|
||||
|
||||
if not _worker_counts:
|
||||
# _max_test_gpu_requirement survives pruning; _worker_counts may be
|
||||
# emptied above when no test fits, and we still want the loud-fail.
|
||||
if _max_test_gpu_requirement == 0:
|
||||
return
|
||||
|
||||
if os.environ.get(ENV_SKIP_MODEL_POOL, "").lower() in ("1", "true", "yes"):
|
||||
@@ -361,19 +385,31 @@ def pytest_collection_finish(session: pytest.Session) -> None:
|
||||
max_required, available_gpus = validate_gpu_requirements()
|
||||
|
||||
if max_required > available_gpus:
|
||||
sep = "=" * LOG_SEPARATOR_WIDTH
|
||||
raise pytest.UsageError(
|
||||
f"\n{sep}\n"
|
||||
f"GPU REQUIREMENTS EXCEEDED\n"
|
||||
f"{sep}\n"
|
||||
f"Test '{_max_test_name}' requires {max_required} GPUs\n"
|
||||
f"Available: {available_gpus} GPUs\n"
|
||||
f"\nOptions:\n"
|
||||
f" 1. Run tests that fit: pytest -k 'not {_max_test_name.split('::')[0]}'\n"
|
||||
f" 2. Reduce workers: @pytest.mark.workers(prefill=1, decode=1)\n"
|
||||
f" 3. Skip GPU tests: SKIP_MODEL_POOL=1 pytest\n"
|
||||
f"{sep}"
|
||||
# Tests whose individual GPU need exceeds capacity are already skipped
|
||||
# in pytest_collection_modifyitems. If literally every collected test
|
||||
# was skipped this way, refuse to pass green — that's the runner-
|
||||
# mismatch case that should fail loud (e.g. wrong matrix entry,
|
||||
# nvidia-smi returning 0 on a healthy host).
|
||||
non_skipped = [
|
||||
it
|
||||
for it in session.items
|
||||
if not any(m.name == "skip" for m in it.iter_markers())
|
||||
]
|
||||
if not non_skipped:
|
||||
raise pytest.UsageError(
|
||||
f"Runner has {available_gpus} GPU(s); every collected test "
|
||||
f"requires more (largest: {_max_test_name} needs {max_required}). "
|
||||
f"Zero tests would run — refusing to pass silently."
|
||||
)
|
||||
# Otherwise: surface the gap so it's obvious in logs that this runner
|
||||
# only ran the fitting subset.
|
||||
logger.warning(
|
||||
"Runner has %d GPU(s); skipped tests requiring up to %d (largest: %s)",
|
||||
available_gpus,
|
||||
max_required,
|
||||
_max_test_name,
|
||||
)
|
||||
return
|
||||
|
||||
logger.info(
|
||||
"GPU validation passed: max %d required (by %s), %d available",
|
||||
|
||||
@@ -19,13 +19,24 @@ from .markers import get_marker_kwargs, get_marker_value
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@pytest.fixture(scope="class")
|
||||
@pytest.fixture
|
||||
def setup_backend(request: pytest.FixtureRequest, model_pool: "ModelPool"):
|
||||
"""Class-scoped fixture that launches a router for each test class.
|
||||
"""Function-scoped fixture that launches a router for each test.
|
||||
|
||||
Routers are cheap to start (~1-2s) compared to workers (~30-60s), so we
|
||||
launch a fresh router per test class for isolation while reusing the
|
||||
expensive workers from model_pool.
|
||||
launch a fresh router per test for isolation while reusing the expensive
|
||||
workers from the session-scoped model_pool fixture.
|
||||
|
||||
NOTE: This used to be ``scope="class"`` to amortize router startup across
|
||||
tests in the same class. Class-scoped fixtures don't survive
|
||||
pytest-parallel's ``--tests-per-worker N`` thread dispatch — its fixture-
|
||||
finalize handling for non-function scopes is buggy (the project hasn't
|
||||
had a real release since 2019). The class teardown silently never fired,
|
||||
so model_pool references acquired in setup leaked indefinitely, blocking
|
||||
eviction and deadlocking any subsequent test that needed a different
|
||||
model. Function scope walks the canonical pytest finalize path for
|
||||
every test, so each acquire is paired with a real release and the pool
|
||||
can evict cleanly.
|
||||
|
||||
Backend types:
|
||||
- "http", "grpc": Gets existing worker from model_pool, launches router
|
||||
@@ -140,126 +151,142 @@ def _setup_pd_backend(
|
||||
num_decode = workers_config.get("decode") or 1
|
||||
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)
|
||||
prefills: list = []
|
||||
decodes: list = []
|
||||
gateway = None
|
||||
|
||||
# Calculate how many more we need
|
||||
missing_prefill = max(0, num_prefill - len(existing_prefills))
|
||||
missing_decode = max(0, num_decode - len(existing_decodes))
|
||||
# Single try/finally guarantees release() runs for every acquired
|
||||
# worker, even if Gateway.start() / OpenAI() raise after acquisition.
|
||||
# See _setup_local_backend for the full rationale.
|
||||
try:
|
||||
# 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)
|
||||
|
||||
# Calculate how many more we need
|
||||
missing_prefill = max(0, num_prefill - len(existing_prefills))
|
||||
missing_decode = max(0, num_decode - len(existing_decodes))
|
||||
|
||||
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),
|
||||
len(decodes),
|
||||
)
|
||||
else:
|
||||
# Build WorkerIdentity list for missing workers
|
||||
workers_to_launch: list[WorkerIdentity] = []
|
||||
for i in range(missing_prefill):
|
||||
workers_to_launch.append(
|
||||
WorkerIdentity(
|
||||
model_id,
|
||||
ConnectionMode.HTTP,
|
||||
WorkerType.PREFILL,
|
||||
len(existing_prefills) + i,
|
||||
)
|
||||
)
|
||||
for i in range(missing_decode):
|
||||
workers_to_launch.append(
|
||||
WorkerIdentity(
|
||||
model_id,
|
||||
ConnectionMode.HTTP,
|
||||
WorkerType.DECODE,
|
||||
len(existing_decodes) + i,
|
||||
)
|
||||
)
|
||||
|
||||
logger.info(
|
||||
"Have %d/%d prefill, %d/%d decode. Launching %d more workers",
|
||||
len(existing_prefills),
|
||||
num_prefill,
|
||||
len(existing_decodes),
|
||||
num_decode,
|
||||
len(workers_to_launch),
|
||||
)
|
||||
new_instances = model_pool.launch_workers(
|
||||
workers_to_launch, startup_timeout=300
|
||||
)
|
||||
|
||||
if not new_instances:
|
||||
# Existing workers will be released by the outer finally.
|
||||
prefills = existing_prefills
|
||||
decodes = existing_decodes
|
||||
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
|
||||
|
||||
# All workers in prefills and decodes are now acquired
|
||||
|
||||
if not prefills or not decodes:
|
||||
pytest.fail(
|
||||
f"PD setup incomplete: have {len(prefills)} prefill, "
|
||||
f"{len(decodes)} decode "
|
||||
f"(need {num_prefill} prefill, {num_decode} decode)"
|
||||
)
|
||||
|
||||
model_path = prefills[0].model_path
|
||||
|
||||
gateway = Gateway()
|
||||
gateway.start(
|
||||
prefill_workers=prefills,
|
||||
decode_workers=decodes,
|
||||
policy=gateway_config["policy"],
|
||||
timeout=gateway_config["timeout"],
|
||||
extra_args=gateway_config["extra_args"],
|
||||
)
|
||||
|
||||
client = openai.OpenAI(
|
||||
base_url=f"{gateway.base_url}/v1",
|
||||
api_key="not-used",
|
||||
)
|
||||
|
||||
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",
|
||||
"Setup PD backend: model=%s, %d prefill + %d decode workers, "
|
||||
"gateway=%s, policy=%s",
|
||||
model_id,
|
||||
len(prefills),
|
||||
len(decodes),
|
||||
)
|
||||
else:
|
||||
# Build WorkerIdentity list for missing workers
|
||||
workers_to_launch: list[WorkerIdentity] = []
|
||||
for i in range(missing_prefill):
|
||||
workers_to_launch.append(
|
||||
WorkerIdentity(
|
||||
model_id,
|
||||
ConnectionMode.HTTP,
|
||||
WorkerType.PREFILL,
|
||||
len(existing_prefills) + i,
|
||||
)
|
||||
)
|
||||
for i in range(missing_decode):
|
||||
workers_to_launch.append(
|
||||
WorkerIdentity(
|
||||
model_id,
|
||||
ConnectionMode.HTTP,
|
||||
WorkerType.DECODE,
|
||||
len(existing_decodes) + i,
|
||||
)
|
||||
)
|
||||
|
||||
logger.info(
|
||||
"Have %d/%d prefill, %d/%d decode. Launching %d more workers",
|
||||
len(existing_prefills),
|
||||
num_prefill,
|
||||
len(existing_decodes),
|
||||
num_decode,
|
||||
len(workers_to_launch),
|
||||
)
|
||||
new_instances = model_pool.launch_workers(
|
||||
workers_to_launch, startup_timeout=300
|
||||
gateway.base_url,
|
||||
gateway_config["policy"],
|
||||
)
|
||||
|
||||
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
|
||||
|
||||
# All workers in prefills and decodes are now acquired
|
||||
|
||||
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()
|
||||
gateway.start(
|
||||
prefill_workers=prefills,
|
||||
decode_workers=decodes,
|
||||
policy=gateway_config["policy"],
|
||||
timeout=gateway_config["timeout"],
|
||||
extra_args=gateway_config["extra_args"],
|
||||
)
|
||||
|
||||
client = openai.OpenAI(
|
||||
base_url=f"{gateway.base_url}/v1",
|
||||
api_key="not-used",
|
||||
)
|
||||
|
||||
logger.info(
|
||||
"Setup PD backend: model=%s, %d prefill + %d decode workers, "
|
||||
"gateway=%s, policy=%s",
|
||||
model_id,
|
||||
len(prefills),
|
||||
len(decodes),
|
||||
gateway.base_url,
|
||||
gateway_config["policy"],
|
||||
)
|
||||
|
||||
try:
|
||||
yield "pd", model_path, client, gateway
|
||||
finally:
|
||||
logger.info("Tearing down PD gateway")
|
||||
gateway.shutdown()
|
||||
# Release references to allow eviction
|
||||
if gateway is not None:
|
||||
logger.info("Tearing down PD gateway")
|
||||
try:
|
||||
gateway.shutdown()
|
||||
except Exception:
|
||||
logger.exception("Gateway shutdown failed; continuing teardown")
|
||||
for worker in prefills + decodes:
|
||||
worker.release()
|
||||
try:
|
||||
worker.release()
|
||||
except Exception:
|
||||
logger.exception(
|
||||
"Release failed for %s; continuing teardown", worker.key
|
||||
)
|
||||
|
||||
|
||||
def _setup_local_backend(
|
||||
@@ -277,87 +304,106 @@ def _setup_local_backend(
|
||||
|
||||
num_workers = workers_config.get("count") or 1
|
||||
instances: list = [] # Track instances for reference counting
|
||||
gateway = None
|
||||
|
||||
# Single try/finally guarantees release() runs for every acquired
|
||||
# instance — even when Gateway.start() / OpenAI() / launch_workers()
|
||||
# raise after acquisition. Without this, a failed gateway start in
|
||||
# one test pinned the worker as is_in_use=True forever, so subsequent
|
||||
# tests that needed a different model couldn't evict and deadlocked
|
||||
# in model_pool.get().
|
||||
try:
|
||||
if num_workers > 1:
|
||||
# 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 = [
|
||||
WorkerIdentity(
|
||||
model_id,
|
||||
connection_mode,
|
||||
WorkerType.REGULAR,
|
||||
len(existing_for_mode) + i,
|
||||
)
|
||||
for i in range(missing)
|
||||
]
|
||||
new_instances = model_pool.launch_workers(
|
||||
workers_to_launch, startup_timeout=300
|
||||
try:
|
||||
if num_workers > 1:
|
||||
# get_workers_by_type auto-acquires all returned workers
|
||||
all_existing = model_pool.get_workers_by_type(
|
||||
model_id, WorkerType.REGULAR
|
||||
)
|
||||
# Acquire newly launched instances
|
||||
for inst in new_instances:
|
||||
inst.acquire()
|
||||
instances = existing_for_mode + new_instances
|
||||
existing_for_mode = [
|
||||
w for w in all_existing if w.mode == connection_mode
|
||||
]
|
||||
|
||||
if not instances:
|
||||
pytest.fail(f"Failed to get {num_workers} workers for {model_id}")
|
||||
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
|
||||
except RuntimeError as e:
|
||||
pytest.fail(str(e))
|
||||
# Release workers we won't use (wrong mode)
|
||||
for w in all_existing:
|
||||
if w not in existing_for_mode:
|
||||
w.release()
|
||||
|
||||
# Launch gateway
|
||||
gateway = Gateway()
|
||||
gateway.start(
|
||||
worker_urls=worker_urls,
|
||||
model_path=model_path,
|
||||
policy=gateway_config["policy"],
|
||||
timeout=gateway_config["timeout"],
|
||||
extra_args=gateway_config["extra_args"],
|
||||
)
|
||||
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 = [
|
||||
WorkerIdentity(
|
||||
model_id,
|
||||
connection_mode,
|
||||
WorkerType.REGULAR,
|
||||
len(existing_for_mode) + i,
|
||||
)
|
||||
for i in range(missing)
|
||||
]
|
||||
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
|
||||
|
||||
client = openai.OpenAI(
|
||||
base_url=f"{gateway.base_url}/v1",
|
||||
api_key="not-used",
|
||||
)
|
||||
if not instances:
|
||||
pytest.fail(f"Failed to get {num_workers} workers for {model_id}")
|
||||
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
|
||||
except RuntimeError as e:
|
||||
pytest.fail(str(e))
|
||||
|
||||
logger.info(
|
||||
"Setup %s backend: model=%s, workers=%d, gateway=%s, policy=%s",
|
||||
backend_name,
|
||||
model_id,
|
||||
num_workers,
|
||||
gateway.base_url,
|
||||
gateway_config["policy"],
|
||||
)
|
||||
gateway = Gateway()
|
||||
gateway.start(
|
||||
worker_urls=worker_urls,
|
||||
model_path=model_path,
|
||||
policy=gateway_config["policy"],
|
||||
timeout=gateway_config["timeout"],
|
||||
extra_args=gateway_config["extra_args"],
|
||||
)
|
||||
|
||||
client = openai.OpenAI(
|
||||
base_url=f"{gateway.base_url}/v1",
|
||||
api_key="not-used",
|
||||
)
|
||||
|
||||
logger.info(
|
||||
"Setup %s backend: model=%s, workers=%d, gateway=%s, policy=%s",
|
||||
backend_name,
|
||||
model_id,
|
||||
num_workers,
|
||||
gateway.base_url,
|
||||
gateway_config["policy"],
|
||||
)
|
||||
|
||||
try:
|
||||
yield backend_name, model_path, client, gateway
|
||||
finally:
|
||||
logger.info("Tearing down gateway for %s backend", backend_name)
|
||||
gateway.shutdown()
|
||||
# Release references to allow eviction
|
||||
if gateway is not None:
|
||||
logger.info("Tearing down gateway for %s backend", backend_name)
|
||||
try:
|
||||
gateway.shutdown()
|
||||
except Exception:
|
||||
logger.exception("Gateway shutdown failed; continuing teardown")
|
||||
# Release references to allow eviction. Each release is
|
||||
# independently fault-isolated so one failure can't strand the
|
||||
# rest of the acquired instances.
|
||||
for inst in instances:
|
||||
inst.release()
|
||||
try:
|
||||
inst.release()
|
||||
except Exception:
|
||||
logger.exception("Release failed for %s; continuing teardown", inst.key)
|
||||
|
||||
|
||||
def _setup_cloud_backend(
|
||||
|
||||
Reference in New Issue
Block a user