[Misc] Migrate SGLANG_SET_CPU_AFFINITY to envs and refactor model config building (#22730)
This commit is contained in:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user