pd

type DecodePodInput ¶

type DecodePodInput struct {
	RunningReqs     float64 // active decode requests on this pod (incl. pending)
	Throughput      float64 // AvgGenerationThroughputToksPerS for the model
	FreeGPUPercent  float64 // 100 - GPUCacheUsagePerc*100, floored at 0.1
	MaxRequestCount float64 // max RunningReqs across the candidate decode pods
	MaxThroughput   float64 // max Throughput across the candidate decode pods
	MaxFreeGPUUsage float64 // max FreeGPUPercent across the candidate decode pods
}

DecodePodInput bundles the per-pod metrics and per-batch maxima needed for one decode scoring pass. Maxima are enforced to be positive by the PD router (typically floored at 1.0) so that normalisation denominators are never zero.

type DecodePolicyName ¶

type DecodePolicyName string

const (
	// DecodePolicyLoadBalancing routes to the pod with the best balance of
	// running-request count, generation throughput, and free GPU headroom.
	DecodePolicyLoadBalancing DecodePolicyName = "load_balancing"

	// DecodePolicyLeastRequest routes to the pod with the fewest active decode
	// requests (including pending requests not yet reflected in metrics).
	DecodePolicyLeastRequest DecodePolicyName = "least_request"
)

type DecodeScorePolicy ¶

type DecodeScorePolicy interface {
	// Name returns the canonical policy identifier used in log lines and metrics.
	Name() DecodePolicyName
	// Describe returns a short human-readable summary for observability (e.g. startup logs).
	Describe() string
	// ScoreDecodePod returns a score for one decode pod; lower is better.
	ScoreDecodePod(routingCtx *types.RoutingContext, pod *v1.Pod, in DecodePodInput) float64
}

DecodeScorePolicy is the stateless scoring strategy for decode pod selection. Implementations are selected via AIBRIX_DECODE_SCORE_POLICY or registered dynamically with RegisterDecodePolicy.

To add a new decode scoring strategy: implement this interface and call RegisterDecodePolicy before NewPDRouter is invoked.

type DecodeScoreRun ¶

type DecodeScoreRun struct {
	PerRoleset   map[string]RolesetDecodePick
	MaxScore     float64
	Err          error
	FallbackUsed bool
	Policy       DecodePolicyName
}

DecodeScoreRun is the outcome of a single call to scoreDecodePods.

• PerRoleset empty with Err == nil: no decode pod produced a usable score (e.g. the input list was empty or every score was NaN after fallback).
• Err non-nil: the pass failed in a way the router should surface to the caller.
• FallbackUsed true: at least one pod's primary-policy score was invalid (NaN) and was replaced by load_balancing for that pod.
• MaxScore is the largest raw score produced in the pass, used by finalPDScore to normalise decode scores before combining them with prefill scores.

type LeastRequestDecodePolicy ¶

type LeastRequestDecodePolicy struct{}

LeastRequestDecodePolicy scores decode pods solely by their running-request count (including pending decode requests not yet visible in metrics). It is the simplest policy and useful when GPU-memory headroom differences between pods are negligible or when throughput metrics are unavailable.

func (LeastRequestDecodePolicy) Describe() string

func (LeastRequestDecodePolicy) Name() DecodePolicyName

func (LeastRequestDecodePolicy) ScoreDecodePod(routingCtx *types.RoutingContext, pod *v1.Pod, in DecodePodInput) float64

type LoadBalancingDecodePolicy ¶

type LoadBalancingDecodePolicy struct{}

LoadBalancingDecodePolicy scores decode pods by combining three normalised metrics: running-request count (higher → worse), generation throughput (lower → worse, expressed as inverse), and free GPU headroom (higher → better):

score = (wRun*normRunningReqs + wThru*(1 - normThroughput)) / normFreeGPU

Weights are set by AIBRIX_DECODE_LB_WEIGHT_RUNNING and AIBRIX_DECODE_LB_WEIGHT_THROUGHPUT (default 1.0 each).

func (LoadBalancingDecodePolicy) Describe() string

func (LoadBalancingDecodePolicy) Name() DecodePolicyName

func (LoadBalancingDecodePolicy) ScoreDecodePod(routingCtx *types.RoutingContext, pod *v1.Pod, in DecodePodInput) float64

type PendingDecodeTracker ¶

type PendingDecodeTracker struct {
	// contains filtered or unexported fields
}

PendingDecodeTracker tracks decode pods that have been selected for a request but whose RealtimeNumRequestsRunning metric has not yet been incremented by the metrics scrape. This bridges the gap between decode pod selection (during routing) and the moment the decode pod actually starts processing the request, preventing concurrent requests from all being routed to the same decode pod during the prefill phase when the metric is still stale.

All methods are safe for concurrent use, including nil receivers (the tracker is optional and may be nil when disabled).

func NewPendingDecodeTracker() *PendingDecodeTracker

func (t *PendingDecodeTracker) AddPendingDecode(requestID, podName string)

func (t *PendingDecodeTracker) GetPendingDecodeCount(podName string) float64

func (t *PendingDecodeTracker) RemovePendingDecode(requestID string)

type PrefillRequestTracker ¶

type PrefillRequestTracker struct {
	// contains filtered or unexported fields
}

PrefillRequestTracker tracks the number of active prefill requests per pod. It is used by the prefill scorer to avoid routing new requests to pods that are already heavily loaded, and by the load-imbalance detector to select the least-loaded pod when the spread exceeds AIBRIX_PREFILL_LOAD_IMBALANCE_MIN_SPREAD.

All methods are safe for concurrent use.

func NewPrefillRequestTracker() *PrefillRequestTracker

func (t *PrefillRequestTracker) AddPrefillRequest(requestID, podName string)

func (t *PrefillRequestTracker) GetPrefillRequestCountsForPod(podname string) int

func (t *PrefillRequestTracker) GetPrefillRequestCountsForPods(pods []*v1.Pod) map[string]int32

func (t *PrefillRequestTracker) RemovePrefillRequest(requestID string)

type PrefillScorePolicy ¶

type PrefillScorePolicy interface {
	// Prepare is called once per request. pods and readyPodsMap represent the
	// same candidate set; readyPodsMap is provided for O(1) name lookups.
	// Returns an error only when scoring cannot proceed at all (e.g. tokenization
	// failure); in that case the router falls back to skipping the request.
	Prepare(routingCtx *types.RoutingContext, pods []*v1.Pod, readyPodsMap map[string]struct{}) (PrefillScorer, error)

	// Name returns the policy identifier used in log lines and metrics.
	Name() string
}

PrefillScorePolicy is the stateless factory for per-request PrefillScorers. The policy itself holds only immutable config (e.g. tokenizer and cache index handles). All per-request state is captured inside the PrefillScorer returned by Prepare, so the policy is safe to share across concurrent goroutines.

To add a new prefill scoring strategy: implement this interface and register it in NewPDRouter (pd_disaggregation.go) by handling the new policy name in the AIBRIX_PREFILL_SCORE_POLICY switch statement.

func NewLeastRequestPrefillPolicy() PrefillScorePolicy

func NewPrefixCachePrefillPolicy(tok tokenizer.Tokenizer, prefixCacheIndexer *prefixcacheindexer.PrefixHashTable) PrefillScorePolicy

type PrefillScorer ¶

type PrefillScorer interface {
	// ScorePod returns a score for pod (lower is better). reqCnt is the pod's
	// current running-request count; maxRequestCount is the maximum across all
	// candidate pods and is used for normalization. Implementations may use pod
	// for logging or metadata lookups (e.g. prefix_cache); implementations that
	// score purely by count (e.g. least_request) may ignore it.
	ScorePod(pod *v1.Pod, reqCnt, maxRequestCount float64) float64

	// PrefixHashes returns the token-prefix hashes used to warm the prefix-cache
	// index after a pod is selected. Returns nil when the policy does not use
	// the prefix cache (e.g. least_request).
	PrefixHashes() []uint64
}

PrefillScorer is a request-scoped scorer created by PrefillScorePolicy. Prepare for a single request. Each call to Prepare returns a fresh instance; no state is shared across concurrent requests.

type RolesetDecodePick ¶

type RolesetDecodePick struct {
	Pod   *v1.Pod
	Score float64
}

RolesetDecodePick is the winning decode pod for one roleset after comparing scores within that roleset (lower score is better).