feat(proxy): introduce datastore abstraction (kubeflow#425)

* feat(proxy): introduce datastore abstraction

Signed-off-by: Alessio Pragliola <[email protected]>

* chore(proxy): apply code suggestions

Co-authored-by: Ramesh Reddy <[email protected]>
Signed-off-by: Alessio Pragliola <[email protected]>

---------

Signed-off-by: Alessio Pragliola <[email protected]>
Co-authored-by: Ramesh Reddy <[email protected]>

Author: Al-Pragliola

cmd/proxy.go

@@ -4,110 +4,98 @@ import (
 	"fmt"
 	"net/http"
 	"sync"
-	"time"
 
 	"github.com/golang/glog"
-	"github.com/kubeflow/model-registry/internal/mlmdtypes"
+
+	"github.com/kubeflow/model-registry/internal/datastore"
 	"github.com/kubeflow/model-registry/internal/proxy"
 	"github.com/kubeflow/model-registry/internal/server/openapi"
-	"github.com/kubeflow/model-registry/pkg/core"
 	"github.com/spf13/cobra"
-	"google.golang.org/grpc"
-	"google.golang.org/grpc/codes"
-	"google.golang.org/grpc/credentials/insecure"
-	"google.golang.org/grpc/status"
 )
 
+type ProxyConfig struct {
+	Datastore datastore.Datastore
+}
+
 const (
-	// mlmdUnavailableMessage is the message returned when the MLMD server is down or unavailable.
-	mlmdUnavailableMessage = "MLMD server is down or unavailable. Please check that the database is reachable and try again later."
-	// maxGRPCRetryAttempts is the maximum number of attempts to retry GRPC requests to the MLMD server.
-	maxGRPCRetryAttempts = 25 // 25 attempts with incremental backoff (1s, 2s, 3s, ..., 25s) it's ~5 minutes
+	// datastoreUnavailableMessage is the message returned when the datastore service is down or unavailable.
+	datastoreUnavailableMessage = "Datastore service is down or unavailable. Please check that the database is reachable and try again later."
 )
 
-// proxyCmd represents the proxy command
-var proxyCmd = &cobra.Command{
-	Use:   "proxy",
-	Short: "Starts the ml-metadata go OpenAPI proxy",
-	Long: `This command launches the ml-metadata go OpenAPI proxy server.
-
-The server connects to a mlmd CPP server. It supports options to customize the
-hostname and port where it listens.'`,
-	RunE: runProxyServer,
-}
+var (
+	proxyCfg = ProxyConfig{
+		Datastore: datastore.Datastore{
+			Type: "mlmd",
+		},
+	}
+
+	// proxyCmd represents the proxy command
+	proxyCmd = &cobra.Command{
+		Use:   "proxy",
+		Short: "Starts the go OpenAPI proxy server to connect to a metadata store",
+		Long: `This command launches the go OpenAPI proxy server.
+
+The server connects to a metadata store, currently only MLMD is supported. It supports options to customize the
+hostname and port where it listens.`,
+		RunE: runProxyServer,
+	}
+)
 
 func runProxyServer(cmd *cobra.Command, args []string) error {
-	var conn *grpc.ClientConn
+	var (
+		ds datastore.Connector
+		wg sync.WaitGroup
+	)
 
 	router := proxy.NewDynamicRouter()
 
 	router.SetRouter(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		http.Error(w, mlmdUnavailableMessage, http.StatusServiceUnavailable)
+		http.Error(w, datastoreUnavailableMessage, http.StatusServiceUnavailable)
 	}))
 
 	errChan := make(chan error, 1)
-	var wg sync.WaitGroup
+
 	wg.Add(2)
+
 	go func() {
 		defer close(errChan)
 		wg.Wait()
 	}()
 
-	// Start the connection to the MLMD server in a separate goroutine, so that
+	// Start the connection to the Datastore server in a separate goroutine, so that
 	// we can start the proxy server and start serving requests while we wait
 	// for the connection to be established.
 	go func() {
+		var (
+			err error
+		)
+
 		defer wg.Done()
 
-		mlmdAddr := fmt.Sprintf("%s:%d", proxyCfg.MLMDHostname, proxyCfg.MLMDPort)
-		glog.Infof("connecting to MLMD server %s..", mlmdAddr)
-		var err error
-		conn, err = grpc.NewClient(mlmdAddr, grpc.WithTransportCredentials(insecure.NewCredentials()))
+		ds, err = datastore.NewConnector(proxyCfg.Datastore)
 		if err != nil {
-			errChan <- fmt.Errorf("error dialing connection to mlmd server %s: %w", mlmdAddr, err)
+			errChan <- fmt.Errorf("error creating datastore: %w", err)
 			return
 		}
 
-		mlmdTypeNamesConfig := mlmdtypes.NewMLMDTypeNamesConfigFromDefaults()
-
-		// Backoff and retry GRPC requests to the MLMD server, until the server
-		// becomes available or the maximum number of attempts is reached.
-		for i := 0; i < maxGRPCRetryAttempts; i++ {
-			_, err = mlmdtypes.CreateMLMDTypes(conn, mlmdTypeNamesConfig)
-			if err == nil {
-				break
-			}
-
-			st, ok := status.FromError(err)
-			if !ok || st.Code() != codes.Unavailable {
-				errChan <- fmt.Errorf("error creating MLMD types: %w", err)
-				return
-			}
-
-			time.Sleep(time.Duration(i+1) * time.Second)
-		}
-
-		service, err := core.NewModelRegistryService(conn, mlmdTypeNamesConfig)
+		conn, err := ds.Connect()
 		if err != nil {
-			errChan <- fmt.Errorf("error creating core service: %w", err)
-
+			errChan <- fmt.Errorf("error connecting to datastore: %w", err)
 			return
 		}
 
-		ModelRegistryServiceAPIService := openapi.NewModelRegistryServiceAPIService(service)
+		ModelRegistryServiceAPIService := openapi.NewModelRegistryServiceAPIService(conn)
 		ModelRegistryServiceAPIController := openapi.NewModelRegistryServiceAPIController(ModelRegistryServiceAPIService)
 
 		router.SetRouter(openapi.NewRouter(ModelRegistryServiceAPIController))
-
-		glog.Infof("connected to MLMD server")
 	}()
 
 	// Start the proxy server in a separate goroutine so that we can handle
-	// errors from both the proxy server and the connection to the MLMD server.
+	// errors from both the proxy server and the connection to the Datastore server.
 	go func() {
 		defer wg.Done()
 
-		glog.Infof("proxy server started at %s:%v", cfg.Hostname, cfg.Port)
+		glog.Infof("Proxy server started at %s:%v", cfg.Hostname, cfg.Port)
 
 		err := http.ListenAndServe(fmt.Sprintf("%s:%d", cfg.Hostname, cfg.Port), router)
 		if err != nil {
@@ -116,35 +104,24 @@ func runProxyServer(cmd *cobra.Command, args []string) error {
 	}()
 
 	defer func() {
-		if conn != nil {
-			glog.Info("closing connection to MLMD server")
-
+		if ds != nil {
 			//nolint:errcheck
-			conn.Close()
+			ds.Teardown()
 		}
 	}()
 
-	// Wait for either the MLMD server connection or the proxy server to return an error
+	// Wait for either the Datastore server connection or the proxy server to return an error
 	// or for both to finish successfully.
 	return <-errChan
 }
-
 func init() {
 	rootCmd.AddCommand(proxyCmd)
 
 	proxyCmd.Flags().StringVarP(&cfg.Hostname, "hostname", "n", cfg.Hostname, "Proxy server listen hostname")
 	proxyCmd.Flags().IntVarP(&cfg.Port, "port", "p", cfg.Port, "Proxy server listen port")
 
-	proxyCmd.Flags().StringVar(&proxyCfg.MLMDHostname, "mlmd-hostname", proxyCfg.MLMDHostname, "MLMD hostname")
-	proxyCmd.Flags().IntVar(&proxyCfg.MLMDPort, "mlmd-port", proxyCfg.MLMDPort, "MLMD port")
-}
-
-type ProxyConfig struct {
-	MLMDHostname string
-	MLMDPort     int
-}
+	proxyCmd.Flags().StringVar(&proxyCfg.Datastore.MLMD.Hostname, "mlmd-hostname", proxyCfg.Datastore.MLMD.Hostname, "MLMD hostname")
+	proxyCmd.Flags().IntVar(&proxyCfg.Datastore.MLMD.Port, "mlmd-port", proxyCfg.Datastore.MLMD.Port, "MLMD port")
 
-var proxyCfg = ProxyConfig{
-	MLMDHostname: "localhost",
-	MLMDPort:     9090,
+	proxyCmd.Flags().StringVar(&proxyCfg.Datastore.Type, "datastore-type", proxyCfg.Datastore.Type, "Datastore type")
 }

cmd/root.go

@@ -90,7 +90,7 @@ func initConfig(cmd *cobra.Command) error {
 
 	// If a config file is found, read it in.
 	if err := viper.ReadInConfig(); err == nil {
-		glog.Info("using config file: ", viper.ConfigFileUsed())
+		glog.Info("Using config file: ", viper.ConfigFileUsed())
 	} else {
 		var configFileNotFoundError viper.ConfigFileNotFoundError
 		ok := errors.As(err, &configFileNotFoundError)

docker-compose-local.yaml

@@ -13,7 +13,7 @@ services:
     build:
       context: .
       dockerfile: Dockerfile
-    command: ["proxy", "--hostname", "0.0.0.0", "--mlmd-hostname", "mlmd-server", "--mlmd-port", "8080"]
+    command: ["proxy", "--hostname", "0.0.0.0", "--mlmd-hostname", "mlmd-server", "--mlmd-port", "8080", "--datastore-type", "mlmd"]
     container_name: model-registry
     ports:
       - "8080:8080"

docker-compose.yaml

@@ -11,7 +11,7 @@ services:
       - ./test/config/ml-metadata:/tmp/shared
   model-registry:
     image: docker.io/kubeflow/model-registry:latest
-    command: ["proxy", "--hostname", "0.0.0.0", "--mlmd-hostname", "mlmd-server", "--mlmd-port", "8080"]
+    command: ["proxy", "--hostname", "0.0.0.0", "--mlmd-hostname", "mlmd-server", "--mlmd-port", "8080", "--datastore-type", "mlmd"]
     container_name: model-registry
     ports:
       - "8080:8080"

internal/datastore/connector.go

@@ -0,0 +1,40 @@
+package datastore
+
+import (
+	"errors"
+	"fmt"
+
+	"github.com/kubeflow/model-registry/pkg/api"
+)
+
+var (
+	ErrCreatingDatastore    = errors.New("error creating datastore")
+	ErrUnsupportedDatastore = errors.New("unsupported datastore type")
+)
+
+type TeardownFunc func() error
+
+type Datastore struct {
+	MLMD MLMDConfig
+	Type string
+}
+
+type Connector interface {
+	Connect() (api.ModelRegistryApi, error)
+	Teardown() error
+}
+
+func NewConnector(ds Datastore) (Connector, error) {
+	switch ds.Type {
+	case "mlmd":
+		if err := ds.MLMD.Validate(); err != nil {
+			return nil, fmt.Errorf("invalid MLMD config: %w", err)
+		}
+
+		mlmd := NewMLMDService(&ds.MLMD)
+
+		return mlmd, nil
+	default:
+		return nil, fmt.Errorf("%w: %s. Supported types: mlmd", ErrUnsupportedDatastore, ds.Type)
+	}
+}

internal/datastore/mlmd.go

@@ -0,0 +1,116 @@
+package datastore
+
+import (
+	"errors"
+	"fmt"
+	"time"
+
+	"github.com/golang/glog"
+	"github.com/kubeflow/model-registry/internal/mlmdtypes"
+	"github.com/kubeflow/model-registry/pkg/api"
+	"github.com/kubeflow/model-registry/pkg/core"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/credentials/insecure"
+	"google.golang.org/grpc/status"
+)
+
+const (
+	// maxGRPCRetryAttempts is the maximum number of attempts to retry GRPC requests to the MLMD service.
+	maxGRPCRetryAttempts = 25 // 25 attempts with incremental backoff (1s, 2s, 3s, ..., 25s) it's ~5 minutes
+)
+
+var (
+	ErrMLMDConnectionStart = errors.New("error dialing connection to mlmd service")
+	ErrMLMDTypeCreation    = errors.New("error creating MLMD types")
+	ErrMLMDCoreCreation    = errors.New("error creating core service")
+	ErrMLMDConnectionClose = errors.New("error closing connection to mlmd service")
+)
+
+type MLMDConfig struct {
+	Hostname string
+	Port     int
+}
+
+func (c *MLMDConfig) Validate() error {
+	if c.Hostname == "" {
+		return fmt.Errorf("hostname is required")
+	}
+
+	if c.Port <= 0 || c.Port > 65535 {
+		return fmt.Errorf("port must be in the range 1-65535")
+	}
+
+	return nil
+}
+
+type MLMDService struct {
+	*MLMDConfig
+	gRPCConnection *grpc.ClientConn
+}
+
+func NewMLMDService(cfg *MLMDConfig) *MLMDService {
+	return &MLMDService{
+		MLMDConfig: cfg,
+	}
+}
+
+func (s *MLMDService) Connect() (api.ModelRegistryApi, error) {
+	uri := fmt.Sprintf("%s:%d", s.Hostname, s.Port)
+
+	glog.Infof("Connecting to MLMD service at %s..", uri)
+
+	conn, err := grpc.NewClient(uri, grpc.WithTransportCredentials(insecure.NewCredentials()))
+	if err != nil {
+		return nil, fmt.Errorf("%w %s: %w", ErrMLMDConnectionStart, uri, err)
+	}
+
+	s.gRPCConnection = conn
+
+	mlmdTypeNamesConfig := mlmdtypes.NewMLMDTypeNamesConfigFromDefaults()
+
+	// Backoff and retry GRPC requests to the MLMD service, until the service
+	// becomes available or the maximum number of attempts is reached.
+	for i := range maxGRPCRetryAttempts {
+		_, err = mlmdtypes.CreateMLMDTypes(conn, mlmdTypeNamesConfig)
+		if err == nil {
+			break
+		}
+
+		st, ok := status.FromError(err)
+		if !ok || st.Code() != codes.Unavailable {
+			return nil, fmt.Errorf("%w: %w", ErrMLMDTypeCreation, err)
+		}
+
+		glog.Warningf("Retrying connection to MLMD service (attempt %d/%d): %v", i+1, maxGRPCRetryAttempts, err)
+
+		time.Sleep(time.Duration(i+1) * time.Second)
+	}
+
+	if err != nil {
+		return nil, fmt.Errorf("%w: %w", ErrMLMDTypeCreation, err)
+	}
+
+	service, err := core.NewModelRegistryService(conn, mlmdTypeNamesConfig)
+	if err != nil {
+		return nil, fmt.Errorf("%w: %w", ErrMLMDCoreCreation, err)
+	}
+
+	glog.Infof("Successfully connected to MLMD service")
+
+	return service, nil
+}
+
+func (s *MLMDService) Teardown() error {
+	glog.Info("Closing connection to MLMD service")
+
+	if s.gRPCConnection == nil {
+		return nil
+	}
+
+	if err := s.gRPCConnection.Close(); err != nil {
+		return fmt.Errorf("%w: %w", ErrMLMDConnectionClose, err)
+	}
+
+	return nil
+}

manifests/kustomize/base/model-registry-configmap.yaml

@@ -9,3 +9,4 @@ data:
   MODEL_REGISTRY_REST_SERVICE_PORT: "8080"
   MODEL_REGISTRY_GRPC_SERVICE_HOST: "model-registry-service"
   MODEL_REGISTRY_GRPC_SERVICE_PORT: "9090"
+  MODEL_REGISTRY_DATA_STORE_TYPE: "mlmd"

manifests/kustomize/base/model-registry-deployment.yaml

@@ -26,6 +26,7 @@ spec:
             - --port=8080
             - --mlmd-hostname=localhost
             - --mlmd-port=9090
+            - --datastore-type=mlmd
           command:
             - /model-registry
             - proxy