diff --git a/cmd/epp/runner/runner.go b/cmd/epp/runner/runner.go index cbd3ea0249..9deafb9b6a 100644 --- a/cmd/epp/runner/runner.go +++ b/cmd/epp/runner/runner.go @@ -143,6 +143,7 @@ func NewRunner() *Runner { type Runner struct { requestControlConfig *requestcontrol.Config schedulerConfig *scheduling.SchedulerConfig + customCollectors []prometheus.Collector } func (r *Runner) WithRequestControlConfig(requestControlConfig *requestcontrol.Config) *Runner { @@ -155,6 +156,11 @@ func (r *Runner) WithSchedulerConfig(schedulerConfig *scheduling.SchedulerConfig return r } +func (r *Runner) WithCustomCollectors(collectors ...prometheus.Collector) *Runner { + r.customCollectors = collectors + return r +} + func (r *Runner) Run(ctx context.Context) error { opts := zap.Options{ Development: true, @@ -205,6 +211,9 @@ func (r *Runner) Run(ctx context.Context) error { // --- Setup Metrics Server --- customCollectors := []prometheus.Collector{collectors.NewInferencePoolMetricsCollector(datastore)} + if r.customCollectors != nil { + customCollectors = append(customCollectors, r.customCollectors...) + } metrics.Register(customCollectors...) metrics.RecordInferenceExtensionInfo(version.CommitSHA, version.BuildRef) // Register metrics handler.