add support to disable catalog or registry service endpoints

Signed-off-by: Dhiraj Bokde <[email protected]>
kubeflow · Jan 22, 2025 · 6f329e9 · 6f329e9
1 parent ea8b1a0
commit 6f329e9
Showing 1 changed file with 53 additions and 29 deletions.
diff --git a/cmd/proxy.go b/cmd/proxy.go
@@ -32,40 +32,55 @@ func runProxyServer(cmd *cobra.Command, args []string) error {
 	ctxTimeout, cancel := context.WithTimeout(context.Background(), time.Second*30)
 	defer cancel()
 
-	mlmdAddr := fmt.Sprintf("%s:%d", proxyCfg.MLMDHostname, proxyCfg.MLMDPort)
-	glog.Infof("connecting to MLMD server %s..", mlmdAddr)
-	conn, err := grpc.DialContext( // nolint:staticcheck
-		ctxTimeout,
-		mlmdAddr,
-		grpc.WithReturnConnectionError(), // nolint:staticcheck
-		grpc.WithBlock(),                 // nolint:staticcheck
-		grpc.WithTransportCredentials(insecure.NewCredentials()),
-	)
-	if err != nil {
-		return fmt.Errorf("error dialing connection to mlmd server %s: %v", mlmdAddr, err)
-	}
-	defer conn.Close()
-	glog.Infof("connected to MLMD server")
+	routers := make([]openapi.Router, 0)
 
-	mlmdTypeNamesConfig := mlmdtypes.NewMLMDTypeNamesConfigFromDefaults()
-	_, err = mlmdtypes.CreateMLMDTypes(conn, mlmdTypeNamesConfig)
-	if err != nil {
-		return fmt.Errorf("error creating MLMD types: %v", err)
-	}
-	service, err := core.NewModelRegistryService(conn, mlmdTypeNamesConfig)
-	if err != nil {
-		return fmt.Errorf("error creating core service: %v", err)
+	disableService := proxyCfg.DisableService
+	if _, ok := map[string]struct{}{"": {}, CatalogService: {}, RegistryService: {}}[disableService]; !ok {
+		return fmt.Errorf("invalid disable-service: %v", disableService)
 	}
+	if disableService != CatalogService {
 
-	// TODO read yaml catalog file and instantiate ModelCatalogAPI implementations
-	ModelCatalogServiceAPIService := openapi.NewModelCatalogServiceAPIService(map[string]openapi.ModelCatalogApi{})
-	ModelCatalogServiceAPIController := openapi.NewModelCatalogServiceAPIController(ModelCatalogServiceAPIService)
+		// TODO read yaml catalog file and instantiate ModelCatalogAPI implementations
+		ModelCatalogServiceAPIService := openapi.NewModelCatalogServiceAPIService(map[string]openapi.ModelCatalogApi{})
+		ModelCatalogServiceAPIController := openapi.NewModelCatalogServiceAPIController(ModelCatalogServiceAPIService)
+		routers = append(routers, ModelCatalogServiceAPIController)
+		glog.Infof("started catalog service")
+	}
+	if disableService != RegistryService {
+
+		mlmdAddr := fmt.Sprintf("%s:%d", proxyCfg.MLMDHostname, proxyCfg.MLMDPort)
+		glog.Infof("connecting to MLMD server %s..", mlmdAddr)
+		conn, err := grpc.DialContext( // nolint:staticcheck
+			ctxTimeout,
+			mlmdAddr,
+			grpc.WithReturnConnectionError(), // nolint:staticcheck
+			grpc.WithBlock(),                 // nolint:staticcheck
+			grpc.WithTransportCredentials(insecure.NewCredentials()),
+		)
+		if err != nil {
+			return fmt.Errorf("error dialing connection to mlmd server %s: %v", mlmdAddr, err)
+		}
+		defer conn.Close()
+		glog.Infof("connected to MLMD server")
+
+		mlmdTypeNamesConfig := mlmdtypes.NewMLMDTypeNamesConfigFromDefaults()
+		_, err = mlmdtypes.CreateMLMDTypes(conn, mlmdTypeNamesConfig)
+		if err != nil {
+			return fmt.Errorf("error creating MLMD types: %v", err)
+		}
+		service, err := core.NewModelRegistryService(conn, mlmdTypeNamesConfig)
+		if err != nil {
+			return fmt.Errorf("error creating core service: %v", err)
+		}
+
+		// TODO make registry API optional to support standalone Catalog deployments
+		ModelRegistryServiceAPIService := openapi.NewModelRegistryServiceAPIService(service)
+		ModelRegistryServiceAPIController := openapi.NewModelRegistryServiceAPIController(ModelRegistryServiceAPIService)
+		routers = append(routers, ModelRegistryServiceAPIController)
 
-	// TODO make registry API optional to support standalone Catalog deployments
-	ModelRegistryServiceAPIService := openapi.NewModelRegistryServiceAPIService(service)
-	ModelRegistryServiceAPIController := openapi.NewModelRegistryServiceAPIController(ModelRegistryServiceAPIService)
+	}
 
-	router := openapi.NewRouter(ModelRegistryServiceAPIController, ModelCatalogServiceAPIController)
+	router := openapi.NewRouter(routers...)
 
 	glog.Fatal(http.ListenAndServe(fmt.Sprintf("%s:%d", cfg.Hostname, cfg.Port), router))
 	return nil
@@ -79,11 +94,20 @@ func init() {
 
 	proxyCmd.Flags().StringVar(&proxyCfg.MLMDHostname, "mlmd-hostname", proxyCfg.MLMDHostname, "MLMD hostname")
 	proxyCmd.Flags().IntVar(&proxyCfg.MLMDPort, "mlmd-port", proxyCfg.MLMDPort, "MLMD port")
+
+	proxyCmd.Flags().StringVar(&proxyCfg.DisableService, "disable-service", proxyCfg.DisableService, "Optional name of service/endpoint to disable, can be either \"catalog\" or \"registry\"")
 }
 
+const (
+	CatalogService  string = "catalog"
+	RegistryService string = "registry"
+)
+
 type ProxyConfig struct {
 	MLMDHostname string
 	MLMDPort     int
+
+	DisableService string
 }
 
 var proxyCfg = ProxyConfig{