Init dashboardClientFunc and httpProxyClientFunc by the config arg (#…

…2092)
ray-project · May 8, 2024 · 6823da1 · 6823da1
1 parent 7b3adca
commit 6823da1
Show file tree

Hide file tree

Showing 6 changed files with 43 additions and 12 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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" {