执行摘要
- 一句话:为指标收集器添加类级DI,支持可插拔后端
- 推荐动作:本 PR 的设计模式(类级 DI + 角色映射 + resolve_collector_class 辅助函数)值得嵌入式框架开发者参考。建议关注 _StatLoggerDIMixin 的声明方式以及 from_cli_args 的适应性修改。
功能与动机
实现 Issue #24467 的第 1 项。现有五个 *MetricsCollector 类直接实例化 prometheus_client.Counter/Gauge/Histogram/Summary,导致嵌入式用户无法在不 fork 的情况下替换后端。参考 vLLM 的 stat_loggers 模式,提供可插拔的指标后端注入点。
实现拆解
- 在 metrics_collector.py 中定义五角色常量 (STAT_LOGGER_ROLE_SCHEDULER 等)、resolve_collector_class 函数和 _StatLoggerDIMixin 基类,为每个 collector 提供四个可覆写的类属性(_counter_cls, _gauge_cls, _histogram_cls, _summary_cls)。
- 五个 *MetricsCollector 类继承 _StatLoggerDIMixin,在 init 中使用 self._counter_cls or _PromCounter 模式,允许子类通过修改类属性替换 prometheus_client 的类。
- 在 ServerArgs 中添加 stat_loggers: Optional[Dict[str, type]] 字段并修改 from_cli_args 跳过非 CLI 字段。
- 在六个 collector 实例化站点(如 SchedulerMetricsCollector.init_new、TokenizerManager.init、HiRadixCache._apply_storage_runtime_config 等)调用 resolve_collector_class 获取正确 collector 类再实例化。
- 新增 test_stat_loggers_di.py(14 个单元测试,覆盖类属性默认值、resolve_collector_class 完整逻辑、DI 实际替换)和 test_metrics.py 中的 TestStatLoggersDI(端到端验证引擎子进程使用自定义 collector)。
关键文件:
python/sglang/srt/observability/metrics_collector.py(模块 可观测性;类别 source;类型 core-logic;符号 resolve_collector_class, _StatLoggerDIMixin, SchedulerMetricsCollector, TokenizerMetricsCollector): 核心变更,添加 DI 基类、角色常量、resolve_collector_class 函数,并修改五个 collector 类继承 _DIMixin,实现可插拔 backend 的根基。
test/registered/unit/observability/test_stat_loggers_di.py(模块 DI测试;类别 test;类型 test-coverage;符号 _StubArgs, TestCollectorClassAttrs, test_scheduler_collector_attrs_default_none, TestResolveCollectorClass): 新增的单元测试文件,覆盖 DI 机制所有逻辑路径,确保类属性默认值、resolve_collector_class 行为正确。
test/registered/observability/test_metrics.py(模块 集成测试;类别 test;类型 test-coverage;符号 _MarkingSchedulerCollector, TestStatLoggersDI, test_engine_custom_scheduler_collector): 端到端测试,验证自定义 SchedulerCollector 通过 Engine 注入后实际生效。
python/sglang/srt/server_args.py(模块 配置层;类别 source;类型 core-logic;符号 stat_loggers, from_cli_args): 新增 stat_loggers 配置字段,并调整 from_cli_args 以跳过非 CLI 字段。
python/sglang/srt/mem_cache/hiradix_cache.py(模块 缓存层;类别 source;类型 dependency-wiring;符号 resolve_collector_class, STAT_LOGGER_ROLE_STORAGE): StorageMetricsCollector 实例化点之一,使用 resolve_collector_class 支持 DI。
python/sglang/srt/mem_cache/hi_mamba_radix_cache.py(模块 Mamba缓存;类别 source;类型 dependency-wiring;符号 resolve_collector_class, STAT_LOGGER_ROLE_STORAGE): 与 hiradix_cache.py 相同,另一个 StorageMetricsCollector 实例化点。
关键符号:resolve_collector_class, _StatLoggerDIMixin, SchedulerMetricsCollector.init, TestResolveCollectorClass.test_returns_subclass_when_role_registered, TestStatLoggersDI.test_engine_custom_scheduler_collector
关键源码片段
python/sglang/srt/observability/metrics_collector.py
核心变更,添加 DI 基类、角色常量、resolve_collector_class 函数,并修改五个 collector 类继承 _DIMixin,实现可插拔 backend 的根基。
# 角色常量,供 ServerArgs.stat_loggers 查找 collector 覆盖
# 嵌入式用户(如 Ray Serve LLM)传入 {'scheduler': MyClass, ...}
STAT_LOGGER_ROLE_SCHEDULER = 'scheduler'
STAT_LOGGER_ROLE_TOKENIZER = 'tokenizer'
STAT_LOGGER_ROLE_STORAGE = 'storage'
STAT_LOGGER_ROLE_RADIX_CACHE = 'radix_cache'
STAT_LOGGER_ROLE_EXPERT_DISPATCH = 'expert_dispatch'
def resolve_collector_class(
server_args: Optional['ServerArgs'], role: str, default_cls: type
) -> type:
'''返回 stat_loggers 中为 role 注册的子类,
若未注册则返回 default_cls。能容忍 server_args=None 和 stat_loggers=None。'''
if server_args is None:
return default_cls
stat_loggers = getattr(server_args, 'stat_loggers', None)
if not stat_loggers:
return default_cls
return stat_loggers.get(role, default_cls)
class _StatLoggerDIMixin:
'''所有 *MetricsCollector 类共享的 DI 覆写钩子。
子类(例如 Ray 后端封装)可以替换这些类属性为
镜像 prometheus_client API 但通过不同后端发出的类。
None 表示保留 prometheus_client 默认实现。'''
_counter_cls = None
_gauge_cls = None
_histogram_cls = None
_summary_cls = None
class SchedulerMetricsCollector(_StatLoggerDIMixin):
def __init__(
self,
labels: Dict[str, str],
enable_lora: bool = False,
enable_hierarchical_cache: bool = False,
enable_streaming_session: bool = False,
server_args: Optional['ServerArgs'] = None,
) -> None:
from prometheus_client import Counter as _PromCounter
from prometheus_client import Gauge as _PromGauge
from prometheus_client import Histogram as _PromHistogram
from prometheus_client import Summary as _PromSummary
# DI:若类属性不为 None,则使用自定义后端;否则回退到 prometheus_client
Counter = self._counter_cls or _PromCounter
Gauge = self._gauge_cls or _PromGauge
Histogram = self._histogram_cls or _PromHistogram
Summary = self._summary_cls or _PromSummary
self.labels = labels
self.enable_lora = enable_lora
# ... 后续使用 Counter、Gauge 等创建指标
评论区精华
Review 中的关键讨论:
1) 审阅者 sufeng-buaa 要求添加端到端测试验证 DI 机制在引擎子进程中实际生效,作者后续补充测试并通过 CI。
2) gemini-code-assist[bot] 建议为 resolve_collector_class 的 server_args 参数添加 Optional[ServerArgs] 类型注解以提升代码维护性。两个讨论均以作者采纳或满足需求结束。
- 建议添加端到端测试验证 DI 机制 (testing): 作者添加了 TestStatLoggersDI 测试并上线,CI 通过后被 approve。
- resolve_collector_class 参数类型注解 (style): PR 合并时此类型注解已存在于最终代码中。
风险与影响
- 风险:默认行为完全不变,兼容性风险极低。自定义 collector 子类需镜像 prometheus_client 的 API,否则可能运行时失败,但这属于使用者责任。多进程环境下 resolve_collector_class 依赖 get_global_server_args() 获取配置,需确保在子进程初始化时全局配置已正确设置,当前实现满足此要求。无性能和安全影响。
- 影响:对普通用户无影响,行为向后兼容。嵌入式用户(如 Ray Serve)获得通过 Python API 注入自定义指标后端的能力。团队需维护新增的 DI 钩子、角色常量和测试,但代码量(+411/-22)可控。
- 风险标记:暂无
关联脉络
参与讨论