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(