Use an env var SGLANG_SET_CPU_AFFINITY to set cpu affinity; turn it off by default (#2222)

This commit is contained in:
Lianmin Zheng
2024-11-27 02:52:46 -08:00
committed by GitHub
parent 6997e28f6e
commit fb6e04a0c2
6 changed files with 23 additions and 20 deletions

View File

@@ -71,9 +71,10 @@ from sglang.srt.utils import (
broadcast_pyobj,
configure_logger,
crash_on_warnings,
get_bool_env_var,
get_zmq_socket,
gpu_proc_affinity,
kill_parent_process,
set_gpu_proc_affinity,
set_random_seed,
suppress_other_loggers,
)
@@ -82,7 +83,7 @@ from sglang.utils import get_exception_traceback
logger = logging.getLogger(__name__)
# Test retract decode
test_retract = os.getenv("SGLANG_TEST_RETRACT", "false").lower() == "true"
test_retract = get_bool_env_var("SGLANG_TEST_RETRACT")
class Scheduler:
@@ -1405,7 +1406,8 @@ def run_scheduler_process(
pipe_writer,
):
# set cpu affinity to this gpu process
gpu_proc_affinity(server_args.tp_size, server_args.nnodes, gpu_id)
if get_bool_env_var("SGLANG_SET_CPU_AFFINITY"):
set_gpu_proc_affinity(server_args.tp_size, server_args.nnodes, gpu_id)
# [For Router] if env var "DP_RANK" exist, set dp_rank to the value of the env var
if dp_rank is None and "DP_RANK" in os.environ: