Use canAddFields option in mlmd type setup (kubeflow#160)

* Use canAddFields option in mlmd type setup * Manage proxy errors * Update pkg/core/core_test.go Co-authored-by: Matteo Mortari <[email protected]> --------- Co-authored-by: Matteo Mortari <[email protected]>
dhirajsb · Nov 19, 2023 · 507e93f · 507e93f
1 parent e9d0541
commit 507e93f
Show file tree

Hide file tree

Showing 3 changed files with 307 additions and 16 deletions.
diff --git a/cmd/proxy.go b/cmd/proxy.go
@@ -3,7 +3,6 @@ package cmd
 import (
 	"context"
 	"fmt"
-	"log"
 	"net/http"
 	"time"
 
@@ -29,13 +28,13 @@ hostname and port where it listens.'`,
 )
 
 func runProxyServer(cmd *cobra.Command, args []string) error {
-	glog.Infof("Proxy server started at %s:%v", cfg.Hostname, cfg.Port)
+	glog.Infof("proxy server started at %s:%v", cfg.Hostname, cfg.Port)
 
 	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)
+	glog.Infof("connecting to MLMD server %s..", mlmdAddr)
 	conn, err := grpc.DialContext(
 		ctxTimeout,
 		mlmdAddr,
@@ -44,24 +43,22 @@ func runProxyServer(cmd *cobra.Command, args []string) error {
 		grpc.WithTransportCredentials(insecure.NewCredentials()),
 	)
 	if err != nil {
-		log.Fatalf("Error dialing connection to mlmd server %s: %v", mlmdAddr, err)
-		return err
+		return fmt.Errorf("error dialing connection to mlmd server %s: %v", mlmdAddr, err)
 	}
 	defer conn.Close()
-	glog.Infof("Connected to MLMD server")
+	glog.Infof("connected to MLMD server")
 
 	service, err := core.NewModelRegistryService(conn)
 	if err != nil {
-		log.Fatalf("Error creating core service: %v", err)
-		return err
+		return fmt.Errorf("error creating core service: %v", err)
 	}
 
 	ModelRegistryServiceAPIService := openapi.NewModelRegistryServiceAPIService(service)
 	ModelRegistryServiceAPIController := openapi.NewModelRegistryServiceAPIController(ModelRegistryServiceAPIService)
 
 	router := openapi.NewRouter(ModelRegistryServiceAPIController)
 
-	log.Fatal(http.ListenAndServe(fmt.Sprintf("%s:%d", cfg.Hostname, cfg.Port), router))
+	glog.Fatal(http.ListenAndServe(fmt.Sprintf("%s:%d", cfg.Hostname, cfg.Port), router))
 	return nil
 }
 

diff --git a/pkg/core/core.go b/pkg/core/core.go
@@ -21,6 +21,7 @@ var (
 	servingEnvironmentTypeName = of(converter.ServingEnvironmentTypeName)
 	inferenceServiceTypeName   = of(converter.InferenceServiceTypeName)
 	serveModelTypeName         = of(converter.ServeModelTypeName)
+	canAddFields               = of(true)
 )
 
 // modelRegistryService is the core library of the model registry
@@ -37,6 +38,7 @@ func NewModelRegistryService(cc grpc.ClientConnInterface) (api.ModelRegistryApi,
 	// Setup the needed Type instances if not existing already
 
 	registeredModelReq := proto.PutContextTypeRequest{
+		CanAddFields: canAddFields,
 		ContextType: &proto.ContextType{
 			Name: registeredModelTypeName,
 			Properties: map[string]proto.PropertyType{
@@ -46,6 +48,7 @@ func NewModelRegistryService(cc grpc.ClientConnInterface) (api.ModelRegistryApi,
 	}
 
 	modelVersionReq := proto.PutContextTypeRequest{
+		CanAddFields: canAddFields,
 		ContextType: &proto.ContextType{
 			Name: modelVersionTypeName,
 			Properties: map[string]proto.PropertyType{
@@ -58,6 +61,7 @@ func NewModelRegistryService(cc grpc.ClientConnInterface) (api.ModelRegistryApi,
 	}
 
 	modelArtifactReq := proto.PutArtifactTypeRequest{
+		CanAddFields: canAddFields,
 		ArtifactType: &proto.ArtifactType{
 			Name: modelArtifactTypeName,
 			Properties: map[string]proto.PropertyType{
@@ -72,6 +76,7 @@ func NewModelRegistryService(cc grpc.ClientConnInterface) (api.ModelRegistryApi,
 	}
 
 	servingEnvironmentReq := proto.PutContextTypeRequest{
+		CanAddFields: canAddFields,
 		ContextType: &proto.ContextType{
 			Name: servingEnvironmentTypeName,
 			Properties: map[string]proto.PropertyType{
@@ -81,6 +86,7 @@ func NewModelRegistryService(cc grpc.ClientConnInterface) (api.ModelRegistryApi,
 	}
 
 	inferenceServiceReq := proto.PutContextTypeRequest{
+		CanAddFields: canAddFields,
 		ContextType: &proto.ContextType{
 			Name: inferenceServiceTypeName,
 			Properties: map[string]proto.PropertyType{
@@ -95,6 +101,7 @@ func NewModelRegistryService(cc grpc.ClientConnInterface) (api.ModelRegistryApi,
 	}
 
 	serveModelReq := proto.PutExecutionTypeRequest{
+		CanAddFields: canAddFields,
 		ExecutionType: &proto.ExecutionType{
 			Name: serveModelTypeName,
 			Properties: map[string]proto.PropertyType{
@@ -106,32 +113,32 @@ func NewModelRegistryService(cc grpc.ClientConnInterface) (api.ModelRegistryApi,
 
 	registeredModelResp, err := client.PutContextType(context.Background(), &registeredModelReq)
 	if err != nil {
-		glog.Fatalf("Error setting up context type %s: %v", *registeredModelTypeName, err)
+		return nil, fmt.Errorf("error setting up context type %s: %v", *registeredModelTypeName, err)
 	}
 
 	modelVersionResp, err := client.PutContextType(context.Background(), &modelVersionReq)
 	if err != nil {
-		glog.Fatalf("Error setting up context type %s: %v", *modelVersionTypeName, err)
+		return nil, fmt.Errorf("error setting up context type %s: %v", *modelVersionTypeName, err)
 	}
 
 	modelArtifactResp, err := client.PutArtifactType(context.Background(), &modelArtifactReq)
 	if err != nil {
-		glog.Fatalf("Error setting up artifact type %s: %v", *modelArtifactTypeName, err)
+		return nil, fmt.Errorf("error setting up artifact type %s: %v", *modelArtifactTypeName, err)
 	}
 
 	servingEnvironmentResp, err := client.PutContextType(context.Background(), &servingEnvironmentReq)
 	if err != nil {
-		glog.Fatalf("Error setting up context type %s: %v", *servingEnvironmentTypeName, err)
+		return nil, fmt.Errorf("error setting up context type %s: %v", *servingEnvironmentTypeName, err)
 	}
 
 	inferenceServiceResp, err := client.PutContextType(context.Background(), &inferenceServiceReq)
 	if err != nil {
-		glog.Fatalf("Error setting up context type %s: %v", *inferenceServiceTypeName, err)
+		return nil, fmt.Errorf("error setting up context type %s: %v", *inferenceServiceTypeName, err)
 	}
 
 	serveModelResp, err := client.PutExecutionType(context.Background(), &serveModelReq)
 	if err != nil {
-		glog.Fatalf("Error setting up execution type %s: %v", *serveModelTypeName, err)
+		return nil, fmt.Errorf("error setting up execution type %s: %v", *serveModelTypeName, err)
 	}
 
 	return &modelRegistryService{