Skip to content

fix(engine): support validated TP x PP topologies #29

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
ShuhaoZhangTony wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

fix(engine): support validated TP x PP topologies #29

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
85 changes: 52 additions & 33 deletions tests/test_systemd_service_scripts.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,9 @@
DEPLOY_HELPERS = REPO_ROOT / "tools" / "lib" / "deploy_common.sh"


def test_runtime_dependency_contract_separates_portable_core_from_full_neuromem() -> None:
def test_runtime_dependency_contract_separates_portable_core_from_full_neuromem() -> (
None
):
project = tomllib.loads((REPO_ROOT / "pyproject.toml").read_text(encoding="utf-8"))[
"project"
]
Expand Down Expand Up @@ -49,18 +51,15 @@ def _make_fake_python(path: Path, *, has_uvicorn: bool) -> Path:
exit_code = "0" if has_uvicorn else "1"
_write_executable(
path,
"#!/usr/bin/env bash\n"
f"exit {exit_code}\n",
f"#!/usr/bin/env bash\nexit {exit_code}\n",
)
return path


def _make_fake_systemctl(path: Path) -> Path:
_write_executable(
path,
"#!/usr/bin/env bash\n"
"printf '%s\n' \"$*\" >>\"$SYSTEMCTL_LOG\"\n"
"exit 0\n",
'#!/usr/bin/env bash\nprintf \'%s\n\' "$*" >>"$SYSTEMCTL_LOG"\nexit 0\n',
)
return path

Expand All @@ -82,10 +81,10 @@ def _run_quickstart_install(
network_log = tmp_path / "unexpected-network.log"
_write_executable(
fake_bin_dir / "git",
'#!/usr/bin/env bash\n'
"#!/usr/bin/env bash\n"
'printf "%s\\n" "$*" >> "$UNEXPECTED_NETWORK_LOG"\nexit 1\n',
)
_write_executable(fake_bin_dir / "curl", '#!/usr/bin/env bash\nprintf 200\n')
_write_executable(fake_bin_dir / "curl", "#!/usr/bin/env bash\nprintf 200\n")
for command in ("npu-smi", "nvidia-smi"):
_write_executable(fake_bin_dir / command, "#!/usr/bin/env bash\nexit 0\n")

Expand Down Expand Up @@ -117,14 +116,16 @@ def _run_quickstart_install(
"FACULTY_TWIN_PARENT_DIR": str(tmp_path / "siblings"),
"DIGITAL_TWIN_RUNTIME_DIR": str(tmp_path / "runtime"),
"VLLM_NVIDIA_CONNECT_HOST": "127.0.0.1",
"PYTHON_BIN": python_bin or str(
_make_fake_python(fake_bin_dir / "python", has_uvicorn=True)
),
"PYTHON_BIN": python_bin
or str(_make_fake_python(fake_bin_dir / "python", has_uvicorn=True)),
}

args = [
"bash", str(checkout / "quickstart.sh"),
"--target", "hosted-web", "--systemd-only",
"bash",
str(checkout / "quickstart.sh"),
"--target",
"hosted-web",
"--systemd-only",
]
if extra_args:
args.extend(extra_args)
Expand Down Expand Up @@ -155,12 +156,8 @@ def test_quickstart_install_renders_service_units(tmp_path: Path) -> None:
)

target_dir = tmp_path / "xdg" / "systemd" / "user"
rendered_app = (target_dir / "sage-mate-app.service").read_text(
encoding="utf-8"
)
rendered_site = (target_dir / "sage-mate-site.service").read_text(
encoding="utf-8"
)
rendered_app = (target_dir / "sage-mate-app.service").read_text(encoding="utf-8")
rendered_site = (target_dir / "sage-mate-site.service").read_text(encoding="utf-8")
rendered_tunnel = (target_dir / "sage-mate-tunnel.service").read_text(
encoding="utf-8"
)
Expand Down Expand Up @@ -188,7 +185,10 @@ def test_release_installer_uses_multi_gpu_tp_for_qwen3_32b_awq() -> None:
assert "default_nvidia_tensor_parallel_size" in script
qwen32_case = script.split("qwen3-32b-awq)", 1)[1].split(";;", 1)[0]
assert 'tp="${tp_override:-1}"' not in qwen32_case
assert 'tp="${tp_override:-$(default_nvidia_tensor_parallel_size "$gpus" "$min_mem")}"' in qwen32_case
assert (
'tp="${tp_override:-$(default_nvidia_tensor_parallel_size "$gpus" "$min_mem")}"'
in qwen32_case
)


def test_quickstart_install_only_enables_optional_services_with_flags(
Expand Down Expand Up @@ -284,13 +284,16 @@ def test_app_runtime_auto_installs_enabled_sage_anns_backend() -> None:
assert "version('isage-neuromem')" in script
assert "from sage.neuromem import UnifiedCollection" in script
assert '"$py" -m pip install --quiet -e "$repo_root"' in script
assert "ensure_neuromem_collection_runtime \"$repo_root\" \"$py\"" in script
assert 'ensure_neuromem_collection_runtime "$repo_root" "$py"' in script
assert '"$python_bin" -m pip install --quiet --no-deps "$requirement"' in helpers
assert '"$uv_bin" pip install --python "$python_bin" --no-deps "$requirement"' in helpers
assert "ensure_neuromem_collection_runtime \"$repo_root\" \"$python_bin\"" in quickstart
assert (
'"$uv_bin" pip install --python "$python_bin" --no-deps "$requirement"'
in helpers
)
assert 'ensure_neuromem_collection_runtime "$repo_root" "$python_bin"' in quickstart
assert '"$py" -m pip install --quiet -e "$repo_root[vdb-anns]"' in script
assert 'DIGITAL_TWIN_KNOWLEDGE_BACKEND:-neuromem' in script
assert 'DIGITAL_TWIN_CONVERSATION_MEMORY_INDEX_TYPE:-segment' in script
assert "DIGITAL_TWIN_KNOWLEDGE_BACKEND:-neuromem" in script
assert "DIGITAL_TWIN_CONVERSATION_MEMORY_INDEX_TYPE:-segment" in script
assert '"$conversation_index" == "sage_vdb_ann"' in script
assert '"$conversation_index" == "sagedb_ann"' in script

Expand All @@ -300,8 +303,8 @@ def test_engine_lock_waits_for_ascend_namespace_release_with_a_finite_bound() ->

assert "VLLM_ENGINE_DEVICE_RELEASE_TIMEOUT_SECONDS" in script
assert "VLLM_ENGINE_DEVICE_RELEASE_POLL_SECONDS" in script
assert 'release_deadline=$((SECONDS + release_timeout))' in script
assert "while ! \"$python_bin\" \"$repo_root/tools/select_idle_npus.py\"" in script
assert "release_deadline=$((SECONDS + release_timeout))" in script
assert 'while ! "$python_bin" "$repo_root/tools/select_idle_npus.py"' in script
assert "configured NPU devices did not become idle within" in script


Expand Down Expand Up @@ -391,7 +394,7 @@ def test_engine_lock_clears_stale_kv_cache_contract() -> None:

def test_engine_lock_records_configured_model_provenance() -> None:
script = ENGINE_LOCK_SCRIPT.read_text(encoding="utf-8")
receipt_block = script.split("umask 077", 1)[1].split("> \"$lock_file\"", 1)[0]
receipt_block = script.split("umask 077", 1)[1].split('> "$lock_file"', 1)[0]

for key in (
"VLLM_ENGINE_MODEL_PATH",
Expand All @@ -415,13 +418,26 @@ def test_engine_launcher_preserves_explicit_immutable_wheel_profile() -> None:
"""Empty conda/source-path values must not be replaced by dev defaults."""

script = ENGINE_SCRIPT.read_text(encoding="utf-8")
assert '${VLLM_ENGINE_CONDA_ENV-vllm-hust-dev}' in script
assert "${VLLM_ENGINE_CONDA_ENV-vllm-hust-dev}" in script
assert (
'${VLLM_ENGINE_BASE_PYTHONPATH-/workspace/vllm-hust:/workspace/vllm-ascend-hust}'
"${VLLM_ENGINE_BASE_PYTHONPATH-/workspace/vllm-hust:/workspace/vllm-ascend-hust}"
in script
)
assert '${VLLM_ENGINE_CONDA_ENV:-vllm-hust-dev}' not in script
assert '${VLLM_ENGINE_BASE_PYTHONPATH:-/workspace/vllm-hust:/workspace/vllm-ascend-hust}' not in script
assert "${VLLM_ENGINE_CONDA_ENV:-vllm-hust-dev}" not in script
assert (
"${VLLM_ENGINE_BASE_PYTHONPATH:-/workspace/vllm-hust:/workspace/vllm-ascend-hust}"
not in script
)


def test_engine_launcher_supports_product_validated_tp_pp_topology() -> None:
script = ENGINE_SCRIPT.read_text(encoding="utf-8")

assert 'engine_pp_size="${VLLM_ENGINE_PP_SIZE:-1}"' in script
assert "engine_tp_size * engine_pp_size != device_count" in script
assert 'args.extend(("--pipeline-parallel-size", str(pp_size)))' in script
assert 'export VLLM_ENGINE_PP_SIZE="$engine_pp_size"' in script
assert "VLLM_ENGINE_TP_SIZE=$engine_tp_size must equal" not in script


def test_engine_verifier_checks_runtime_import_origins() -> None:
Expand All @@ -446,5 +462,8 @@ def test_engine_verifier_accepts_immutable_wheels_without_source_roots() -> None
encoding="utf-8"
)
assert '[[ -n "$runtime_pythonpath" ]]' in script
assert '[[ -n "$declared_pythonpath" || "$installed_modules_json" != \'{}\' ]]' in script
assert (
'[[ -n "$declared_pythonpath" || "$installed_modules_json" != \'{}\' ]]'
in script
)
assert "runtime or declared engine PYTHONPATH is empty" not in script
49 changes: 47 additions & 2 deletions tools/run_vllm_engine.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -264,10 +264,53 @@ if [[ -z "${runtime_visible_devices:-}" ]]; then
runtime_device_count="$device_count"
fi
engine_tp_size="${VLLM_ENGINE_TP_SIZE:-$device_count}"
if [[ ! "$engine_tp_size" =~ ^[1-9][0-9]*$ || "$engine_tp_size" -ne "$device_count" ]]; then
echo "ERROR: VLLM_ENGINE_TP_SIZE=$engine_tp_size must equal the configured device count ($device_count)." >&2
engine_pp_size="${VLLM_ENGINE_PP_SIZE:-1}"
if [[ ! "$engine_tp_size" =~ ^[1-9][0-9]*$ ]]; then
echo "ERROR: VLLM_ENGINE_TP_SIZE=$engine_tp_size must be a positive integer." >&2
exit 2
fi
if [[ ! "$engine_pp_size" =~ ^[1-9][0-9]*$ ]]; then
echo "ERROR: VLLM_ENGINE_PP_SIZE=$engine_pp_size must be a positive integer." >&2
exit 2
fi
if (( engine_tp_size * engine_pp_size != device_count )); then
echo "ERROR: TP${engine_tp_size} x PP${engine_pp_size} must equal the configured device count ($device_count)." >&2
exit 2
fi

if (( engine_pp_size > 1 )); then
export VLLM_ENGINE_EXTRA_ARGS_JSON="$($selector_python - "$engine_pp_size" <<'PY'
import json
import os
import sys

pp_size = int(sys.argv[1])
raw = os.environ.get("VLLM_ENGINE_EXTRA_ARGS_JSON", "[]")
try:
args = json.loads(raw)
except json.JSONDecodeError as error:
raise SystemExit(f"ERROR: VLLM_ENGINE_EXTRA_ARGS_JSON is invalid JSON: {error}")
if not isinstance(args, list) or not all(isinstance(item, str) for item in args):
raise SystemExit("ERROR: VLLM_ENGINE_EXTRA_ARGS_JSON must be an array of strings.")

configured = None
for index, item in enumerate(args):
if item == "--pipeline-parallel-size":
if index + 1 >= len(args):
raise SystemExit("ERROR: --pipeline-parallel-size requires a value.")
configured = args[index + 1]
elif item.startswith("--pipeline-parallel-size="):
configured = item.split("=", 1)[1]
if configured is not None and configured != str(pp_size):
raise SystemExit(
"ERROR: VLLM_ENGINE_PP_SIZE conflicts with --pipeline-parallel-size."
)
if configured is None:
args.extend(("--pipeline-parallel-size", str(pp_size)))
print(json.dumps(args, separators=(",", ":")))
PY
)"
fi

npu_selector="$repo_root/tools/select_idle_npus.py"
if [[ -z "$selector_python" || ! -x "$selector_python" ]]; then
Expand Down Expand Up @@ -306,6 +349,7 @@ export VLLM_ENGINE_CONNECT_PORT="${VLLM_ENGINE_CONNECT_PORT:-$VLLM_ENGINE_PORT}"
export VLLM_PROXY_UPSTREAM_BASE_URL="${VLLM_PROXY_UPSTREAM_BASE_URL:-http://$VLLM_ENGINE_CONNECT_HOST:$VLLM_ENGINE_CONNECT_PORT/v1}"
unset resolved_engine_port
export VLLM_ENGINE_TP_SIZE="$engine_tp_size"
export VLLM_ENGINE_PP_SIZE="$engine_pp_size"
export VLLM_ENGINE_MAX_MODEL_LEN="${VLLM_ENGINE_MAX_MODEL_LEN:-32768}"
export VLLM_ENGINE_MAX_NUM_BATCHED_TOKENS="${VLLM_ENGINE_MAX_NUM_BATCHED_TOKENS:-$VLLM_ENGINE_MAX_MODEL_LEN}"
export VLLM_ENGINE_GPU_MEM_UTIL="${VLLM_ENGINE_GPU_MEM_UTIL:-0.9}"
Expand Down Expand Up @@ -547,6 +591,7 @@ echo "[sage-mate] delegating vLLM-HUST launch to $launcher"
echo "[sage-mate] model served name=${VLLM_ENGINE_SERVED_MODEL_NAME:-unknown}"
echo "[sage-mate] quantization=${VLLM_ENGINE_QUANTIZATION:-none}"
echo "[sage-mate] npu devices=${VLLM_ENGINE_NPU_DEVICES}"
echo "[sage-mate] topology=TP${VLLM_ENGINE_TP_SIZE}xPP${VLLM_ENGINE_PP_SIZE}"
echo "[sage-mate] host=${VLLM_ENGINE_HOST:-0.0.0.0} port=${VLLM_ENGINE_PORT:-8000}"
} >>"$VLLM_ENGINE_CONTAINER_LOG_FILE"
exec "$launcher"
Loading