diff --git a/ray-operator/apis/config/v1alpha1/configuration_types.go b/ray-operator/apis/config/v1alpha1/configuration_types.go index 892476c146..d039224e46 100644 --- a/ray-operator/apis/config/v1alpha1/configuration_types.go +++ b/ray-operator/apis/config/v1alpha1/configuration_types.go @@ -1,8 +1,10 @@ package v1alpha1 import ( + "github.com/ray-project/kuberay/ray-operator/controllers/ray/utils" corev1 "k8s.io/api/core/v1" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "sigs.k8s.io/controller-runtime/pkg/manager" ) //+kubebuilder:object:root=true @@ -61,3 +63,11 @@ type Configuration struct { // to inject into every Worker pod. WorkerSidecarContainers []corev1.Container `json:"workerSidecarContainers,omitempty"` } + +func (config Configuration) GetDashboardClient(mgr manager.Manager) func() utils.RayDashboardClientInterface { + return utils.GetRayDashboardClientFunc(mgr, config.UseKubernetesProxy) +} + +func (config Configuration) GetHttpProxyClient(mgr manager.Manager) func() utils.RayHttpProxyClientInterface { + return utils.GetRayHttpProxyClientFunc(mgr, config.UseKubernetesProxy) +} diff --git a/ray-operator/controllers/ray/rayjob_controller.go b/ray-operator/controllers/ray/rayjob_controller.go index 6a9312fb5c..922eb6d810 100644 --- a/ray-operator/controllers/ray/rayjob_controller.go +++ b/ray-operator/controllers/ray/rayjob_controller.go @@ -43,7 +43,8 @@ type RayJobReconciler struct { } // NewRayJobReconciler returns a new reconcile.Reconciler -func NewRayJobReconciler(ctx context.Context, mgr manager.Manager, dashboardClientFunc func() utils.RayDashboardClientInterface) *RayJobReconciler { +func NewRayJobReconciler(ctx context.Context, mgr manager.Manager, provider utils.ClientProvider) *RayJobReconciler { + dashboardClientFunc := provider.GetDashboardClient(mgr) return &RayJobReconciler{ Client: mgr.GetClient(), Scheme: mgr.GetScheme(), diff --git a/ray-operator/controllers/ray/rayservice_controller.go b/ray-operator/controllers/ray/rayservice_controller.go index 5272c7a74e..6acf22c1a4 100644 --- a/ray-operator/controllers/ray/rayservice_controller.go +++ b/ray-operator/controllers/ray/rayservice_controller.go @@ -60,7 +60,9 @@ type RayServiceReconciler struct { } // NewRayServiceReconciler returns a new reconcile.Reconciler -func NewRayServiceReconciler(ctx context.Context, mgr manager.Manager, dashboardClientFunc func() utils.RayDashboardClientInterface, httpProxyClientFunc func() utils.RayHttpProxyClientInterface) *RayServiceReconciler { +func NewRayServiceReconciler(ctx context.Context, mgr manager.Manager, provider utils.ClientProvider) *RayServiceReconciler { + dashboardClientFunc := provider.GetDashboardClient(mgr) + httpProxyClientFunc := provider.GetHttpProxyClient(mgr) return &RayServiceReconciler{ Client: mgr.GetClient(), Scheme: mgr.GetScheme(), diff --git a/ray-operator/controllers/ray/suite_test.go b/ray-operator/controllers/ray/suite_test.go index 05b4d385f7..899ac547ab 100644 --- a/ray-operator/controllers/ray/suite_test.go +++ b/ray-operator/controllers/ray/suite_test.go @@ -20,6 +20,8 @@ import ( "path/filepath" "testing" + "sigs.k8s.io/controller-runtime/pkg/manager" + rayv1 "github.com/ray-project/kuberay/ray-operator/apis/ray/v1" "github.com/ray-project/kuberay/ray-operator/controllers/ray/utils" @@ -51,6 +53,20 @@ var ( fakeRayHttpProxyClient *utils.FakeRayHttpProxyClient ) +type TestClientProvider struct{} + +func (testProvider TestClientProvider) GetDashboardClient(mgr manager.Manager) func() utils.RayDashboardClientInterface { + return func() utils.RayDashboardClientInterface { + return fakeRayDashboardClient + } +} + +func (testProvider TestClientProvider) GetHttpProxyClient(mgr manager.Manager) func() utils.RayHttpProxyClientInterface { + return func() utils.RayHttpProxyClientInterface { + return fakeRayHttpProxyClient + } +} + func TestAPIs(t *testing.T) { RegisterFailHandler(Fail) @@ -107,16 +123,11 @@ var _ = BeforeSuite(func(ctx SpecContext) { err = NewReconciler(ctx, mgr, options).SetupWithManager(mgr, 1) Expect(err).NotTo(HaveOccurred(), "failed to setup RayCluster controller") - err = NewRayServiceReconciler(ctx, mgr, func() utils.RayDashboardClientInterface { - return fakeRayDashboardClient - }, func() utils.RayHttpProxyClientInterface { - return fakeRayHttpProxyClient - }).SetupWithManager(mgr) + testClientProvider := TestClientProvider{} + err = NewRayServiceReconciler(ctx, mgr, testClientProvider).SetupWithManager(mgr) Expect(err).NotTo(HaveOccurred(), "failed to setup RayService controller") - err = NewRayJobReconciler(ctx, mgr, func() utils.RayDashboardClientInterface { - return fakeRayDashboardClient - }).SetupWithManager(mgr) + err = NewRayJobReconciler(ctx, mgr, testClientProvider).SetupWithManager(mgr) Expect(err).NotTo(HaveOccurred(), "failed to setup RayJob controller") go func() { diff --git a/ray-operator/controllers/ray/utils/util.go b/ray-operator/controllers/ray/utils/util.go index f567b4c090..371674d996 100644 --- a/ray-operator/controllers/ray/utils/util.go +++ b/ray-operator/controllers/ray/utils/util.go @@ -13,6 +13,8 @@ import ( "time" "unicode" + "sigs.k8s.io/controller-runtime/pkg/manager" + batchv1 "k8s.io/api/batch/v1" "k8s.io/apimachinery/pkg/util/json" @@ -521,3 +523,8 @@ func EnvVarByName(envName string, envVars []corev1.EnvVar) (corev1.EnvVar, bool) } return corev1.EnvVar{}, false } + +type ClientProvider interface { + GetDashboardClient(mgr manager.Manager) func() RayDashboardClientInterface + GetHttpProxyClient(mgr manager.Manager) func() RayHttpProxyClientInterface +} diff --git a/ray-operator/main.go b/ray-operator/main.go index 25e3ff87cf..75ac821a24 100644 --- a/ray-operator/main.go +++ b/ray-operator/main.go @@ -212,9 +212,9 @@ func main() { ctx := ctrl.SetupSignalHandler() exitOnError(ray.NewReconciler(ctx, mgr, rayClusterOptions).SetupWithManager(mgr, config.ReconcileConcurrency), "unable to create controller", "controller", "RayCluster") - exitOnError(ray.NewRayServiceReconciler(ctx, mgr, utils.GetRayDashboardClientFunc(mgr, config.UseKubernetesProxy), utils.GetRayHttpProxyClientFunc(mgr, config.UseKubernetesProxy)).SetupWithManager(mgr), + exitOnError(ray.NewRayServiceReconciler(ctx, mgr, config).SetupWithManager(mgr), "unable to create controller", "controller", "RayService") - exitOnError(ray.NewRayJobReconciler(ctx, mgr, utils.GetRayDashboardClientFunc(mgr, config.UseKubernetesProxy)).SetupWithManager(mgr), + exitOnError(ray.NewRayJobReconciler(ctx, mgr, config).SetupWithManager(mgr), "unable to create controller", "controller", "RayJob") if os.Getenv("ENABLE_WEBHOOKS") == "true" {