[Misc] Migrate SGLANG_SET_CPU_AFFINITY to envs and refactor model config building (#22730)

This commit is contained in:
Lianmin Zheng
2026-04-13 16:10:31 -07:00
committed by GitHub
parent f4f9e68189
commit 8f9553bccb
3 changed files with 16 additions and 10 deletions
+1 -6
View File
@@ -379,10 +379,6 @@ class Scheduler(
)
)
self.enable_kv_cache_events = bool(
server_args.kv_events_config and self.attn_tp_rank == 0
)
# Init model configs
self.init_model_config()
@@ -3688,12 +3684,11 @@ def run_scheduler_process(
dp_rank = configure_scheduler(
server_args, tp_rank, attn_cp_rank, moe_dp_rank, moe_ep_rank, pp_rank, dp_rank
)
kill_itself_when_parent_died()
parent_process = psutil.Process().parent()
# Set cpu affinity to this gpu process
if get_bool_env_var("SGLANG_SET_CPU_AFFINITY"):
if envs.SGLANG_SET_CPU_AFFINITY.get():
set_gpu_proc_affinity(
server_args.pp_size, server_args.tp_size, server_args.nnodes, gpu_id
)
@@ -362,7 +362,7 @@ class ModelRunner(ModelRunnerKVCacheMixin):
self.dflash_draft_num_layers = None
if self.spec_algorithm.is_eagle3() and not self.is_draft_worker:
# load draft config
draft_model_config = ModelConfig.from_server_args(
draft_model_config = self._build_model_config(
server_args,
model_path=(server_args.speculative_draft_model_path),
model_revision=server_args.speculative_draft_model_revision,
@@ -391,7 +391,7 @@ class ModelRunner(ModelRunnerKVCacheMixin):
)
# Select target layers to capture for building DFlash context features.
draft_model_config = ModelConfig.from_server_args(
draft_model_config = self._build_model_config(
server_args,
model_path=(server_args.speculative_draft_model_path),
model_revision=server_args.speculative_draft_model_revision,
@@ -497,6 +497,16 @@ class ModelRunner(ModelRunnerKVCacheMixin):
self._model_update_group = {}
self._weights_send_group = {}
def _build_model_config(
self, server_args, model_path=None, model_revision=None, is_draft_model=False
):
return ModelConfig.from_server_args(
server_args,
model_path=model_path,
model_revision=model_revision,
is_draft_model=is_draft_model,
)
def init_mindspore_runner(self):
# Init the mindspore runner
# for now, there is only some communication initialization work
@@ -166,14 +166,15 @@ class SchedulerMetricsMixin:
reporter=self.metrics_collector.increment_gpu_overlap_wait_seconds,
)
if self.enable_kv_cache_events:
self.init_kv_events(self.server_args.kv_events_config)
self.init_kv_events(self.server_args.kv_events_config)
self.scheduler_status_logger = SchedulerStatusLogger.maybe_create(
enable_metrics=self.enable_metrics
)
def init_kv_events(self: Scheduler, kv_events_config: Optional[str]):
self.enable_kv_cache_events = bool(kv_events_config and self.attn_tp_rank == 0)
if self.enable_kv_cache_events:
self.kv_event_publisher = EventPublisherFactory.create(
kv_events_config, self.attn_dp_rank