scheduler_profile.go

/*
Copyright 2025 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package scheduling

import (
	"context"
	"fmt"
	"sort"
	"strings"
	"time"

	"go.opentelemetry.io/otel/attribute"
	"go.opentelemetry.io/otel/trace"
	"sigs.k8s.io/controller-runtime/pkg/log"

	errcommon "github.com/llm-d/llm-d-router/pkg/common/error"
	logutil "github.com/llm-d/llm-d-router/pkg/common/observability/logging"
	"github.com/llm-d/llm-d-router/pkg/common/observability/tracing"
	"github.com/llm-d/llm-d-router/pkg/epp/framework/interface/plugin"
	fwksched "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/scheduling"
	"github.com/llm-d/llm-d-router/pkg/epp/metrics"
)

// NewSchedulerProfile creates a new SchedulerProfile object and returns its pointer.
func NewSchedulerProfile() *SchedulerProfile {
	return &SchedulerProfile{
		filters: []fwksched.Filter{},
		scorers: []*WeightedScorer{},
		// picker remains nil since profile doesn't support multiple pickers
	}
}

// SchedulerProfile provides a profile configuration for the scheduler which influence routing decisions.
type SchedulerProfile struct {
	filters []fwksched.Filter
	scorers []*WeightedScorer
	picker  fwksched.Picker
}

// WithFilters sets the given filter plugins as the Filter plugins.
// if the SchedulerProfile has Filter plugins, this call replaces the existing plugins with the given ones.
func (p *SchedulerProfile) WithFilters(filters ...fwksched.Filter) *SchedulerProfile {
	p.filters = filters
	return p
}

// WithScorers sets the given scorer plugins as the Scorer plugins.
// if the SchedulerProfile has Scorer plugins, this call replaces the existing plugins with the given ones.
func (p *SchedulerProfile) WithScorers(scorers ...*WeightedScorer) *SchedulerProfile {
	p.scorers = scorers
	return p
}

// WithPicker sets the given picker plugins as the Picker plugin.
// if the SchedulerProfile has Picker plugin, this call replaces the existing plugin with the given one.
func (p *SchedulerProfile) WithPicker(picker fwksched.Picker) *SchedulerProfile {
	p.picker = picker
	return p
}

// AddPlugins adds the given plugins to all scheduler plugins according to the interfaces each plugin implements.
// A plugin may implement more than one scheduler plugin interface.
// Special Case: In order to add a scorer, one must use the scorer.NewWeightedScorer function in order to provide a weight.
// if a scorer implements more than one interface, supplying a WeightedScorer is sufficient. The function will take the internal
// scorer object and register it to all interfaces it implements.
func (p *SchedulerProfile) AddPlugins(pluginObjects ...plugin.Plugin) error {
	for _, plugin := range pluginObjects {
		if weightedScorer, ok := plugin.(*WeightedScorer); ok {
			p.scorers = append(p.scorers, weightedScorer)
			plugin = weightedScorer.Scorer // if we got WeightedScorer, unwrap the plugin
		} else if scorer, ok := plugin.(fwksched.Scorer); ok { // if we got a Scorer instead of WeightedScorer that's an error.
			return fmt.Errorf("failed to register scorer '%s' without a weight. follow function documentation to register a scorer", scorer.TypedName())
		}
		if filter, ok := plugin.(fwksched.Filter); ok {
			p.filters = append(p.filters, filter)
		}
		if picker, ok := plugin.(fwksched.Picker); ok {
			if p.picker != nil {
				return fmt.Errorf("failed to set '%s' as picker, already have a registered picker plugin '%s'", picker.TypedName(), p.picker.TypedName())
			}
			p.picker = picker
		}
	}
	return nil
}

func (p *SchedulerProfile) String() string {
	filterNames := make([]string, len(p.filters))
	for i, filter := range p.filters {
		filterNames[i] = filter.TypedName().String()
	}
	scorerNames := make([]string, len(p.scorers))
	for i, scorer := range p.scorers {
		scorerNames[i] = fmt.Sprintf("%s: %f", scorer.TypedName(), scorer.Weight())
	}

	return fmt.Sprintf(
		"{Filters: [%s], Scorers: [%s], Picker: %s}",
		strings.Join(filterNames, ", "),
		strings.Join(scorerNames, ", "),
		p.picker.TypedName(),
	)
}

// Run runs a SchedulerProfile. It invokes all the SchedulerProfile plugins for the given request in this
// order - Filters, Scorers, Picker. After completing all, it returns the result.
func (p *SchedulerProfile) Run(ctx context.Context, request *fwksched.InferenceRequest, candidateEndpoints []fwksched.Endpoint) (*fwksched.ProfileRunResult, error) {
	endpoints := p.runFilterPlugins(ctx, request, candidateEndpoints)
	if len(endpoints) == 0 {
		return nil, errcommon.Error{Code: errcommon.Internal, Msg: "no endpoints available for the given request"}
	}
	// if we got here, there is at least one endpoint to score
	weightedScorePerEndpoint := p.runScorerPlugins(ctx, request, endpoints)

	result := p.runPickerPlugin(ctx, request, weightedScorePerEndpoint)

	return result, nil
}

func (p *SchedulerProfile) runFilterPlugins(ctx context.Context, request *fwksched.InferenceRequest, endpoints []fwksched.Endpoint) []fwksched.Endpoint {
	logger := log.FromContext(ctx)
	filteredEndpoints := endpoints
	logger.V(logutil.DEBUG).Info("Before running filter plugins", "endpoints", filteredEndpoints)

	ctx, span := tracing.Tracer(TracerScope).Start(ctx, "filter_endpoints",
		trace.WithSpanKind(trace.SpanKindInternal),
	)
	defer span.End()
	span.SetAttributes(attribute.Int("llm_d.epp.filter.candidate_endpoints", len(endpoints)))
	if request != nil {
		if request.TargetModel != "" {
			span.SetAttributes(attribute.String("gen_ai.request.model", request.TargetModel))
		}
		if request.RequestID != "" {
			span.SetAttributes(attribute.String("gen_ai.request.id", request.RequestID))
		}
	}

	for _, filter := range p.filters {
		logger.V(logutil.VERBOSE).Info("Running filter plugin", "plugin", filter.TypedName())
		before := time.Now()
		filteredEndpoints = filter.Filter(ctx, request, filteredEndpoints)
		metrics.RecordPluginProcessingLatency(filterExtensionPoint, filter.TypedName().Type, filter.TypedName().Name, time.Since(before))
		logger.V(logutil.DEBUG).Info("Completed running filter plugin successfully", "plugin", filter.TypedName(), "endpoints", filteredEndpoints)
		if len(filteredEndpoints) == 0 {
			logger.V(logutil.VERBOSE).Info("Filter eliminated all endpoints", "plugin", filter.TypedName(), "endpointsBefore", len(endpoints))
			break
		}
	}
	span.SetAttributes(attribute.Int("llm_d.epp.filter.filtered_endpoints", len(filteredEndpoints)))
	logger.V(logutil.VERBOSE).Info("Completed running filter plugins", "remainingEndpoints", len(filteredEndpoints))

	return filteredEndpoints
}

func (p *SchedulerProfile) runScorerPlugins(ctx context.Context, request *fwksched.InferenceRequest, endpoints []fwksched.Endpoint) map[fwksched.Endpoint]float64 {
	logger := log.FromContext(ctx)
	logger.V(logutil.DEBUG).Info("Before running scorer plugins", "endpoints", endpoints)

	weightedScorePerEndpoint := make(map[fwksched.Endpoint]float64, len(endpoints))
	for _, endpoint := range endpoints {
		weightedScorePerEndpoint[endpoint] = float64(0) // initialize weighted score per endpoint with 0 value
	}
	// Cache the debug logger and its enabled state once. The per-endpoint Info
	// call below evaluates and boxes its variadic args even when the verbosity
	// gate would suppress output; on a 100-endpoint, 4-scorer fleet that line
	// alone accounted for ~80% of total allocations per Scheduler.Schedule call.
	// Guarding by Enabled() preserves debugging behavior while removing the
	// per-endpoint allocation when DEBUG logging is off (the production default).
	debug := logger.V(logutil.DEBUG)
	debugEnabled := debug.Enabled()

	// Iterate through each scorer in the chain and accumulate the weighted scores.
	for _, scorer := range p.scorers {
		logger.V(logutil.VERBOSE).Info("Running scorer plugin", "plugin", scorer.TypedName())
		before := time.Now()
		scores := scorer.Score(ctx, request, endpoints)
		metrics.RecordPluginProcessingLatency(scorerExtensionPoint, scorer.TypedName().Type, scorer.TypedName().Name, time.Since(before))
		for endpoint, score := range scores { // weight is relative to the sum of weights
			if debugEnabled {
				debug.Info("Calculated score", "plugin", scorer.TypedName(), "endpoint", endpoint.GetMetadata().NamespacedName, "score", score)
			}
			weightedScorePerEndpoint[endpoint] += enforceScoreRange(score) * scorer.Weight()
		}
		debug.Info("Completed running scorer plugin successfully", "plugin", scorer.TypedName())
	}
	logger.V(logutil.VERBOSE).Info("Completed running scorer plugins successfully")

	return weightedScorePerEndpoint
}

func (p *SchedulerProfile) runPickerPlugin(ctx context.Context, request *fwksched.InferenceRequest, weightedScorePerEndpoint map[fwksched.Endpoint]float64) *fwksched.ProfileRunResult {
	logger := log.FromContext(ctx)

	// Allocate the ScoredEndpoint values as a single contiguous backing array
	// and build the picker's pointer slice by indexing into it. Previously each
	// per-endpoint &ScoredEndpoint{...} was a separate heap allocation, which
	// at production fleet sizes (~100 pods) dominated per-request picker cost.
	// Pickers reorder the pointer slice (shuffle/sort) but do not realloc, so
	// pointer aliasing into the backing array is safe.
	n := len(weightedScorePerEndpoint)
	storage := make([]fwksched.ScoredEndpoint, n)
	scoredEndpoints := make([]*fwksched.ScoredEndpoint, n)
	i := 0
	for endpoint, score := range weightedScorePerEndpoint {
		storage[i] = fwksched.ScoredEndpoint{Endpoint: endpoint, Score: score}
		scoredEndpoints[i] = &storage[i]
		i++
	}
	logger.V(logutil.VERBOSE).Info("Running picker plugin", "plugin", p.picker.TypedName())
	logger.V(logutil.DEBUG).Info("Candidate pods for picking", "endpoints-weighted-score", scoredEndpoints)

	ctx, span := tracing.Tracer(TracerScope).Start(ctx, "pick_endpoints",
		trace.WithSpanKind(trace.SpanKindInternal),
	)
	defer span.End()

	span.SetAttributes(attribute.Int("llm_d.epp.picker.candidate_endpoints", len(scoredEndpoints)))
	// The picker almost always returns a single target, so its count carries
	// little signal. The score distribution across the strongest candidates is
	// what explains why an endpoint was chosen, so record the highest-scoring
	// few (names with their weighted scores). Captured before Pick because
	// pickers reorder scoredEndpoints in place.
	if names, scores := topScoredEndpoints(scoredEndpoints, maxTracedEndpointScores); len(names) > 0 {
		span.SetAttributes(
			attribute.StringSlice("llm_d.epp.picker.top_endpoints", names),
			attribute.Float64Slice("llm_d.epp.picker.top_scores", scores),
		)
	}
	if request != nil {
		if request.TargetModel != "" {
			span.SetAttributes(attribute.String("gen_ai.request.model", request.TargetModel))
		}
		if request.RequestID != "" {
			span.SetAttributes(attribute.String("gen_ai.request.id", request.RequestID))
		}
	}

	before := time.Now()
	result := p.picker.Pick(ctx, scoredEndpoints)
	metrics.RecordPluginProcessingLatency(pickerExtensionPoint, p.picker.TypedName().Type, p.picker.TypedName().Name, time.Since(before))
	logger.V(logutil.DEBUG).Info("Completed running picker plugin successfully", "plugin", p.picker.TypedName(), "result", result)

	return result
}

// topScoredEndpoints returns the names and weighted scores of the highest
// scoring candidates, ordered by descending score with the endpoint name as a
// stable tiebreaker and capped at limit. The returned slices are index-aligned.
func topScoredEndpoints(scored []*fwksched.ScoredEndpoint, limit int) ([]string, []float64) {
	ranked := make([]*fwksched.ScoredEndpoint, len(scored))
	copy(ranked, scored)
	sort.Slice(ranked, func(i, j int) bool {
		if ranked[i].Score != ranked[j].Score {
			return ranked[i].Score > ranked[j].Score
		}
		return ranked[i].GetMetadata().NamespacedName.String() <
			ranked[j].GetMetadata().NamespacedName.String()
	})
	if limit < len(ranked) {
		ranked = ranked[:limit]
	}
	names := make([]string, len(ranked))
	scores := make([]float64, len(ranked))
	for i, se := range ranked {
		names[i] = se.GetMetadata().NamespacedName.String()
		scores[i] = se.Score
	}
	return names, scores
}

func enforceScoreRange(score float64) float64 {
	if score < 0 {
		return 0
	}
	if score > 1 {
		return 1
	}
	return score
}