@@ -6,13 +6,17 @@ import (
66 "encoding/json"
77 "errors"
88 "fmt"
9+ "net"
10+ "strconv"
911
1012 "sigs.k8s.io/controller-runtime/pkg/log"
1113 "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/plugins"
1214 "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/scheduling/framework"
1315 "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/scheduling/framework/plugins/multi/prefix"
1416 "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/scheduling/types"
1517 logutil "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/util/logging"
18+
19+ "github.com/llm-d/llm-d-inference-scheduler/pkg/common"
1620)
1721
1822const (
@@ -30,6 +34,7 @@ type pdProfileHandlerParameters struct {
3034 PrefillProfile string `json:"prefillProfile"`
3135 PrefixPluginName string `json:"prefixPluginName"`
3236 HashBlockSize int `json:"hashBlockSize"`
37+ PrimaryPort int `json:"primaryPort"`
3338}
3439
3540// compile-time type assertion
@@ -43,6 +48,7 @@ func PdProfileHandlerFactory(name string, rawParameters json.RawMessage, _ plugi
4348 PrefillProfile : defaultPrefillProfile ,
4449 PrefixPluginName : defaultPrefixPluginName ,
4550 HashBlockSize : prefix .DefaultBlockSize ,
51+ PrimaryPort : 0 ,
4652 }
4753 if rawParameters != nil {
4854 if err := json .Unmarshal (rawParameters , & parameters ); err != nil {
@@ -51,19 +57,24 @@ func PdProfileHandlerFactory(name string, rawParameters json.RawMessage, _ plugi
5157 }
5258
5359 return NewPdProfileHandler (parameters .PrefillProfile , parameters .DecodeProfile , parameters .PrefixPluginName ,
54- parameters .Threshold , parameters .HashBlockSize ).WithName (name ), nil
60+ parameters .Threshold , parameters .HashBlockSize , parameters . PrimaryPort ).WithName (name ), nil
5561}
5662
5763// NewPdProfileHandler initializes a new PdProfileHandler and returns its pointer.
58- func NewPdProfileHandler (prefillProfile string , decodeProfile string , prefixPluginName string , pdThreshold int , hashBlockSize int ) * PdProfileHandler {
59- return & PdProfileHandler {
64+ func NewPdProfileHandler (prefillProfile string , decodeProfile string , prefixPluginName string , pdThreshold int , hashBlockSize int , primaryPort int ) * PdProfileHandler {
65+ result := & PdProfileHandler {
6066 typedName : plugins.TypedName {Type : PdProfileHandlerType },
6167 prefixPluginTypedName : plugins.TypedName {Type : prefix .PrefixCachePluginType , Name : prefixPluginName },
6268 decodeProfile : decodeProfile ,
6369 prefillProfile : prefillProfile ,
6470 pdThreshold : pdThreshold ,
6571 hashBlockSize : hashBlockSize ,
6672 }
73+ if primaryPort != 0 {
74+ result .primaryPort = strconv .Itoa (primaryPort )
75+ }
76+
77+ return result
6778}
6879
6980// PdProfileHandler handles scheduler profiles for PD.
@@ -74,6 +85,7 @@ type PdProfileHandler struct {
7485 prefillProfile string
7586 pdThreshold int
7687 hashBlockSize int
88+ primaryPort string
7789}
7890
7991// TypedName returns the typed name of the plugin.
@@ -143,27 +155,47 @@ func (h *PdProfileHandler) Pick(ctx context.Context, cycleState *types.CycleStat
143155// ProcessResults handles the outcome of the profile runs after the selected profiles ran.
144156// In case of an error in any of the profiles, the matching entry in the profileResults will contain nil, to indicate there was
145157// an error while running the profile.
146- func (h * PdProfileHandler ) ProcessResults (_ context.Context , _ * types.CycleState , _ * types.LLMRequest ,
158+ func (h * PdProfileHandler ) ProcessResults (_ context.Context , _ * types.CycleState , request * types.LLMRequest ,
147159 profileResults map [string ]* types.ProfileRunResult ) (* types.SchedulingResult , error ) {
148- if profileResults [h .decodeProfile ] == nil { // if decode profile failed to run, we should fail
160+ decodeRunResults := profileResults [h .decodeProfile ]
161+ if decodeRunResults == nil { // if decode profile failed to run, we should fail
149162 return nil , errors .New ("failed to find available decode workers" )
150163 }
151164 // otherwise, decode ran successfully
152165
166+ updatedResults := map [string ]* types.ProfileRunResult {}
167+
168+ // Add decode profile to result
169+ if h .primaryPort != "" {
170+ // Data Parallel is active
171+
172+ targetPod := decodeRunResults .TargetPods [0 ].GetPod ()
173+ request .Headers [common .DataParallelPodHeader ] = net .JoinHostPort (targetPod .Address , targetPod .Port )
174+
175+ updatedResult := types.ProfileRunResult {
176+ TargetPods : []types.Pod {},
177+ }
178+
179+ for _ , target := range decodeRunResults .TargetPods {
180+ updatedPodInfo := target .GetPod ().Clone ()
181+ updatedPodInfo .Port = h .primaryPort
182+ targetPod := & types.PodMetrics {Pod : updatedPodInfo , MetricsState : target .GetMetrics ().Clone ()}
183+ updatedResult .TargetPods = append (updatedResult .TargetPods , targetPod )
184+ }
185+ updatedResults [h .decodeProfile ] = & updatedResult
186+ } else {
187+ updatedResults [h .decodeProfile ] = decodeRunResults
188+ }
189+
153190 // if both prefill and decode ran successfully
154191 if prefillRunResult , exists := profileResults [h .prefillProfile ]; exists && prefillRunResult != nil {
155- return & types.SchedulingResult {
156- PrimaryProfileName : h .decodeProfile ,
157- ProfileResults : profileResults ,
158- }, nil
192+ // Add the prefill profile to the results
193+ updatedResults [h .prefillProfile ] = prefillRunResult
159194 }
160195
161- // otherwise, decode ran successfully and prefill failed. filter out prefill from the returned results.
162196 return & types.SchedulingResult {
163197 PrimaryProfileName : h .decodeProfile ,
164- ProfileResults : map [string ]* types.ProfileRunResult {
165- h .decodeProfile : profileResults [h .decodeProfile ], // return decode only
166- },
198+ ProfileResults : updatedResults ,
167199 }, nil
168200}
169201
0 commit comments