added context argument to scheduling profile picker (kubernetes-sigs#926)

* added context argument to scheduling profile picker

since this is a custom logic, getting context as argument can be useful.
specific in my case I want to use contextual logging

Signed-off-by: Nir Rozenbaum <[email protected]>

* Update pkg/epp/scheduling/framework/plugins/profile-picker/all_profiles_picker.go

Co-authored-by: Etai Lev Ran <[email protected]>

---------

Signed-off-by: Nir Rozenbaum <[email protected]>
Co-authored-by: Abdullah Gharaibeh <[email protected]>
Co-authored-by: Etai Lev Ran <[email protected]>

Author: 3 people

pkg/epp/scheduling/framework/plugins.go

@@ -35,7 +35,7 @@ const (
 // and the previously executed SchedluderProfile cycles along with their results.
 type ProfilePicker interface {
 	plugins.Plugin
-	Pick(request *types.LLMRequest, profiles map[string]*SchedulerProfile, executionResults map[string]*types.Result) map[string]*SchedulerProfile
+	Pick(ctx context.Context, request *types.LLMRequest, profiles map[string]*SchedulerProfile, executionResults map[string]*types.Result) map[string]*SchedulerProfile
 }
 
 // Filter defines the interface for filtering a list of pods based on context.

pkg/epp/scheduling/framework/plugins/profile-picker/all_profiles_picker.go

@@ -17,6 +17,8 @@ limitations under the License.
 package profilepicker
 
 import (
+	"context"
+
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/scheduling/framework"
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/scheduling/types"
 )
@@ -39,7 +41,8 @@ func (p *AllProfilesPicker) Name() string {
 
 // Pick selects the SchedulingProfiles to run from the list of candidate profiles, while taking into consideration the request properties and the
 // previously executed cycles along with their results.
-func (p *AllProfilesPicker) Pick(request *types.LLMRequest, profiles map[string]*framework.SchedulerProfile, executionResults map[string]*types.Result) map[string]*framework.SchedulerProfile {
+func (p *AllProfilesPicker) Pick(_ context.Context, request *types.LLMRequest, profiles map[string]*framework.SchedulerProfile,
+	executionResults map[string]*types.Result) map[string]*framework.SchedulerProfile {
 	if len(profiles) == len(executionResults) { // all profiles have been executed already in previous call
 		return map[string]*framework.SchedulerProfile{}
 	}

pkg/epp/scheduling/scheduler.go

@@ -112,7 +112,7 @@ func (s *Scheduler) Schedule(ctx context.Context, request *types.LLMRequest) (ma
 
 	for { // get the next set of profiles to run iteratively based on the request and the previous execution results
 		before := time.Now()
-		profiles := s.profilePicker.Pick(request, s.profiles, profileExecutionResults)
+		profiles := s.profilePicker.Pick(ctx, request, s.profiles, profileExecutionResults)
 		metrics.RecordSchedulerPluginProcessingLatency(framework.ProfilePickerType, s.profilePicker.Name(), time.Since(before))
 		if len(profiles) == 0 { // profile picker didn't pick any profile to run
 			break