-
Notifications
You must be signed in to change notification settings - Fork 368
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Add batch-scheduler option, deprecate enable-batch-scheduler option (#…
- Loading branch information
Showing
10 changed files
with
377 additions
and
81 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,34 @@ | ||
| package v1alpha1 | ||
|
|
||
| import ( | ||
| "fmt" | ||
|
|
||
| "github.com/go-logr/logr" | ||
|
|
||
| "github.com/ray-project/kuberay/ray-operator/controllers/ray/batchscheduler/volcano" | ||
| "github.com/ray-project/kuberay/ray-operator/controllers/ray/batchscheduler/yunikorn" | ||
| ) | ||
|
|
||
| func ValidateBatchSchedulerConfig(logger logr.Logger, config Configuration) error { | ||
| if config.EnableBatchScheduler { | ||
| logger.Info("Feature flag enable-batch-scheduler is deprecated and will not be supported soon. " + | ||
| "Use batch-scheduler instead. ") | ||
| return nil | ||
| } | ||
|
|
||
| if len(config.BatchScheduler) > 0 { | ||
| // default option, no-opt. | ||
| if config.BatchScheduler == "default" { | ||
| return nil | ||
| } | ||
|
|
||
| // if a customized scheduler is configured, check it is supported | ||
| if config.BatchScheduler == volcano.GetPluginName() || config.BatchScheduler == yunikorn.GetPluginName() { | ||
| logger.Info("Feature flag batch-scheduler is enabled", | ||
| "scheduler name", config.BatchScheduler) | ||
| } else { | ||
| return fmt.Errorf("scheduler is not supported, name=%s", config.BatchScheduler) | ||
| } | ||
| } | ||
| return nil | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,84 @@ | ||
| package v1alpha1 | ||
|
|
||
| import ( | ||
| "testing" | ||
|
|
||
| "github.com/go-logr/logr" | ||
| "github.com/go-logr/logr/testr" | ||
|
|
||
| "github.com/ray-project/kuberay/ray-operator/controllers/ray/batchscheduler/volcano" | ||
| "github.com/ray-project/kuberay/ray-operator/controllers/ray/batchscheduler/yunikorn" | ||
| ) | ||
|
|
||
| func TestValidateBatchSchedulerConfig(t *testing.T) { | ||
| type args struct { | ||
| logger logr.Logger | ||
| config Configuration | ||
| } | ||
| tests := []struct { | ||
| name string | ||
| args args | ||
| wantErr bool | ||
| }{ | ||
| { | ||
| name: "legacy option, enable-batch-scheduler=false", | ||
| args: args{ | ||
| logger: testr.New(t), | ||
| config: Configuration{ | ||
| EnableBatchScheduler: false, | ||
| }, | ||
| }, | ||
| wantErr: false, | ||
| }, | ||
| { | ||
| name: "legacy option, enable-batch-scheduler=true", | ||
| args: args{ | ||
| logger: testr.New(t), | ||
| config: Configuration{ | ||
| EnableBatchScheduler: true, | ||
| }, | ||
| }, | ||
| wantErr: false, | ||
| }, | ||
| { | ||
| name: "valid option, batch-scheduler=yunikorn", | ||
| args: args{ | ||
| logger: testr.New(t), | ||
| config: Configuration{ | ||
| BatchScheduler: yunikorn.GetPluginName(), | ||
| }, | ||
| }, | ||
| wantErr: false, | ||
| }, | ||
| { | ||
| name: "valid option, batch-scheduler=volcano", | ||
| args: args{ | ||
| logger: testr.New(t), | ||
| config: Configuration{ | ||
| BatchScheduler: volcano.GetPluginName(), | ||
| }, | ||
| }, | ||
| wantErr: false, | ||
| }, | ||
| { | ||
| name: "invalid option, invalid scheduler name", | ||
| args: args{ | ||
| logger: testr.New(t), | ||
| config: Configuration{ | ||
| EnableBatchScheduler: false, | ||
| BatchScheduler: "unknown-scheduler-name", | ||
| }, | ||
| }, | ||
| wantErr: true, | ||
| }, | ||
| } | ||
|
|
||
| for _, tt := range tests { | ||
| t.Run(tt.name, func(t *testing.T) { | ||
| t.Logf(tt.name) | ||
| if err := ValidateBatchSchedulerConfig(tt.args.logger, tt.args.config); (err != nil) != tt.wantErr { | ||
| t.Errorf("ValidateBatchSchedulerConfig() error = %v, wantErr %v", err, tt.wantErr) | ||
| } | ||
| }) | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
121 changes: 60 additions & 61 deletions
121
ray-operator/controllers/ray/batchscheduler/schedulermanager.go
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,93 +1,92 @@ | ||
| package batchscheduler | ||
|
|
||
| import ( | ||
| "fmt" | ||
| "sync" | ||
|
|
||
| "k8s.io/apimachinery/pkg/runtime" | ||
| "k8s.io/client-go/rest" | ||
| "sigs.k8s.io/controller-runtime/pkg/builder" | ||
|
|
||
| rayv1 "github.com/ray-project/kuberay/ray-operator/apis/ray/v1" | ||
| schedulerinterface "github.com/ray-project/kuberay/ray-operator/controllers/ray/batchscheduler/interface" | ||
| configapi "github.com/ray-project/kuberay/ray-operator/apis/config/v1alpha1" | ||
| "github.com/ray-project/kuberay/ray-operator/controllers/ray/batchscheduler/volcano" | ||
| "github.com/ray-project/kuberay/ray-operator/controllers/ray/batchscheduler/yunikorn" | ||
|
|
||
| "k8s.io/client-go/rest" | ||
|
|
||
| rayv1 "github.com/ray-project/kuberay/ray-operator/apis/ray/v1" | ||
| schedulerinterface "github.com/ray-project/kuberay/ray-operator/controllers/ray/batchscheduler/interface" | ||
| "github.com/ray-project/kuberay/ray-operator/controllers/ray/utils" | ||
| ) | ||
|
|
||
| var schedulerContainers = map[string]schedulerinterface.BatchSchedulerFactory{ | ||
| schedulerinterface.GetDefaultPluginName(): &schedulerinterface.DefaultBatchSchedulerFactory{}, | ||
| volcano.GetPluginName(): &volcano.VolcanoBatchSchedulerFactory{}, | ||
| yunikorn.GetPluginName(): &yunikorn.YuniKornSchedulerFactory{}, | ||
| } | ||
|
|
||
| func GetRegisteredNames() []string { | ||
| var pluginNames []string | ||
| for key := range schedulerContainers { | ||
| pluginNames = append(pluginNames, key) | ||
| } | ||
| return pluginNames | ||
| type SchedulerManager struct { | ||
| config *rest.Config | ||
| factory schedulerinterface.BatchSchedulerFactory | ||
| scheduler schedulerinterface.BatchScheduler | ||
| rayConfigs configapi.Configuration | ||
| sync.Mutex | ||
| } | ||
|
|
||
| func ConfigureReconciler(b *builder.Builder) *builder.Builder { | ||
| for _, factory := range schedulerContainers { | ||
| b = factory.ConfigureReconciler(b) | ||
| // NewSchedulerManager maintains a specific scheduler plugin based on config | ||
| func NewSchedulerManager(rayConfigs configapi.Configuration, config *rest.Config) (*SchedulerManager, error) { | ||
| // init the scheduler factory from config | ||
| factory := getSchedulerFactory(rayConfigs) | ||
| scheduler, err := factory.New(config) | ||
| if err != nil { | ||
| return nil, err | ||
| } | ||
| return b | ||
| } | ||
|
|
||
| func AddToScheme(scheme *runtime.Scheme) { | ||
| for _, factory := range schedulerContainers { | ||
| factory.AddToScheme(scheme) | ||
| manager := SchedulerManager{ | ||
| rayConfigs: rayConfigs, | ||
| config: config, | ||
| factory: factory, | ||
| scheduler: scheduler, | ||
| } | ||
| } | ||
|
|
||
| type SchedulerManager struct { | ||
| config *rest.Config | ||
| plugins map[string]schedulerinterface.BatchScheduler | ||
| sync.Mutex | ||
| return &manager, nil | ||
| } | ||
|
|
||
| func NewSchedulerManager(config *rest.Config) *SchedulerManager { | ||
| manager := SchedulerManager{ | ||
| config: config, | ||
| plugins: make(map[string]schedulerinterface.BatchScheduler), | ||
| func getSchedulerFactory(rayConfigs configapi.Configuration) schedulerinterface.BatchSchedulerFactory { | ||
| var factory schedulerinterface.BatchSchedulerFactory | ||
| // init with the default factory | ||
| factory = &schedulerinterface.DefaultBatchSchedulerFactory{} | ||
| // when a batch scheduler name is provided | ||
| if len(rayConfigs.BatchScheduler) > 0 { | ||
| switch rayConfigs.BatchScheduler { | ||
| case volcano.GetPluginName(): | ||
| factory = &volcano.VolcanoBatchSchedulerFactory{} | ||
| case yunikorn.GetPluginName(): | ||
| factory = &yunikorn.YuniKornSchedulerFactory{} | ||
| default: | ||
| factory = &schedulerinterface.DefaultBatchSchedulerFactory{} | ||
| } | ||
| } | ||
| return &manager | ||
| } | ||
|
|
||
| func (batch *SchedulerManager) GetSchedulerForCluster(app *rayv1.RayCluster) (schedulerinterface.BatchScheduler, error) { | ||
| if schedulerName, ok := app.ObjectMeta.Labels[utils.RaySchedulerName]; ok { | ||
| return batch.GetScheduler(schedulerName) | ||
| // legacy option, if this is enabled, register volcano | ||
| // this is for backwards compatibility | ||
| if rayConfigs.EnableBatchScheduler { | ||
| factory = &volcano.VolcanoBatchSchedulerFactory{} | ||
| } | ||
|
|
||
| // no scheduler provided | ||
| return &schedulerinterface.DefaultBatchScheduler{}, nil | ||
| return factory | ||
| } | ||
|
|
||
| func (batch *SchedulerManager) GetScheduler(schedulerName string) (schedulerinterface.BatchScheduler, error) { | ||
| factory, registered := schedulerContainers[schedulerName] | ||
| if !registered { | ||
| return nil, fmt.Errorf("unregistered scheduler plugin %s", schedulerName) | ||
| func (batch *SchedulerManager) GetSchedulerForCluster(app *rayv1.RayCluster) (schedulerinterface.BatchScheduler, error) { | ||
| // for backwards compatibility | ||
| if batch.rayConfigs.EnableBatchScheduler { | ||
| if schedulerName, ok := app.ObjectMeta.Labels[utils.RaySchedulerName]; ok { | ||
| if schedulerName == volcano.GetPluginName() { | ||
| return batch.scheduler, nil | ||
| } | ||
| } | ||
| } | ||
|
|
||
| batch.Lock() | ||
| defer batch.Unlock() | ||
| return batch.scheduler, nil | ||
| } | ||
|
|
||
| plugin, existed := batch.plugins[schedulerName] | ||
| func (batch *SchedulerManager) ConfigureReconciler(b *builder.Builder) *builder.Builder { | ||
| batch.factory.ConfigureReconciler(b) | ||
| return b | ||
| } | ||
|
|
||
| if existed && plugin != nil { | ||
| return plugin, nil | ||
| } | ||
| if existed && plugin == nil { | ||
| return nil, fmt.Errorf( | ||
| "failed to get scheduler plugin %s, previous initialization has failed", schedulerName) | ||
| } | ||
| plugin, err := factory.New(batch.config) | ||
| if err != nil { | ||
| batch.plugins[schedulerName] = nil | ||
| return nil, err | ||
| } | ||
| batch.plugins[schedulerName] = plugin | ||
| return plugin, nil | ||
| func (batch *SchedulerManager) AddToScheme(scheme *runtime.Scheme) { | ||
| batch.factory.AddToScheme(scheme) | ||
| } |
Oops, something went wrong.