From ed5e905c89d299a384a40e81318966e3b0744627 Mon Sep 17 00:00:00 2001 From: hzh0425 Date: Fri, 7 Nov 2025 10:58:44 +0800 Subject: [PATCH] [HotFix]: Add missing SGLANG_EPLB_HEATMAP_COLLECTION_INTERVAL env var (#12776) --- python/sglang/srt/environ.py | 1 + python/sglang/srt/eplb/expert_distribution.py | 9 +++------ 2 files changed, 4 insertions(+), 6 deletions(-) diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 51566ba3e..057e3b4f3 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -215,6 +215,7 @@ class Envs: SGLANG_EXPERT_LOCATION_UPDATER_LOG_METRICS = EnvBool(False) SGLANG_LOG_EXPERT_LOCATION_METADATA = EnvBool(False) SGLANG_EXPERT_DISTRIBUTION_RECORDER_DIR = EnvStr("/tmp") + SGLANG_EPLB_HEATMAP_COLLECTION_INTERVAL = EnvInt(0) # TBO SGLANG_TBO_DEBUG = EnvBool(False) diff --git a/python/sglang/srt/eplb/expert_distribution.py b/python/sglang/srt/eplb/expert_distribution.py index 10b50403d..bfec3802e 100644 --- a/python/sglang/srt/eplb/expert_distribution.py +++ b/python/sglang/srt/eplb/expert_distribution.py @@ -31,7 +31,7 @@ from sglang.srt.environ import envs from sglang.srt.metrics.collector import ExpertDispatchCollector from sglang.srt.model_executor.forward_batch_info import ForwardBatch from sglang.srt.server_args import ServerArgs -from sglang.srt.utils import Withable, is_npu +from sglang.srt.utils import Withable, get_int_env_var, is_npu _is_npu = is_npu() @@ -716,11 +716,8 @@ class _UtilizationRateAccumulatorMixin(_Accumulator): def _collect_metrics_if_needed(self, gpu_physical_count: torch.Tensor): # sglang:eplb_gpu_physical_count metric is disabled if SGLANG_EPLB_HEATMAP_COLLECTION_INTERVAL <= 0 - if ( - envs.SGLANG_EPLB_HEATMAP_COLLECTION_INTERVAL > 0 - and self._collection_counter % envs.SGLANG_EPLB_HEATMAP_COLLECTION_INTERVAL - == 0 - ): + interval = get_int_env_var("SGLANG_EPLB_HEATMAP_COLLECTION_INTERVAL", 0) + if interval > 0 and self._collection_counter % interval == 0: for layer_idx in range(self._expert_location_metadata.num_layers): count_of_layer = ( self._expert_dispatch_collector.eplb_gpu_physical_count.labels(