Add state properties for RegisteredModel, ModelVersion, and Inference…

…Service (kubeflow#185) * Add state properties for RegisteredModel, ModelVersion, and InferenceService. Fixes kubeflow#184 * Add property state to mlmd types * Add openapi to mlmd mapping for states --------- Co-authored-by: Andrea Lamparelli <[email protected]>
dhirajsb · Nov 23, 2023 · e0f2925 · e0f2925
1 parent 3f3be43
commit e0f2925
Show file tree

Hide file tree

Showing 22 changed files with 927 additions and 18 deletions.
diff --git a/api/openapi/model-registry.yaml b/api/openapi/model-registry.yaml
@@ -1017,6 +1017,35 @@ components:
         - ABANDONED
         - REFERENCE
       type: string
+    RegisteredModelState:
+      description: |-
+        - LIVE: A state indicating that the `RegisteredModel` exists
+        - ARCHIVED: A state indicating that the `RegisteredModel` has been archived.
+      default: LIVE
+      enum:
+        - LIVE
+        - ARCHIVED
+      type: string
+    ModelVersionState:
+      description: |-
+        - LIVE: A state indicating that the `ModelVersion` exists
+        - ARCHIVED: A state indicating that the `ModelVersion` has been archived.
+      default: LIVE
+      enum:
+        - LIVE
+        - ARCHIVED
+      type: string
+    InferenceServiceState:
+      description: |-
+        - DEPLOYED: A state indicating that the `InferenceService` should be deployed.
+        - UNDEPLOYED: A state indicating that the `InferenceService` should be un-deployed.
+        The state indicates the desired state of inference service.
+        See the associated `ServeModel` for the actual status of service deployment action.
+      default: DEPLOYED
+      enum:
+        - DEPLOYED
+        - UNDEPLOYED
+      type: string
     ExecutionState:
       description: |-
         The state of the Execution. The state transitions are
@@ -1103,6 +1132,11 @@ components:
       allOf:
         -
           $ref: '#/components/schemas/BaseResourceUpdate'
+        -
+          type: object
+          properties:
+            state:
+              $ref: '#/components/schemas/RegisteredModelState'
     ModelVersion:
       description: Represents a ModelVersion belonging to a RegisteredModel.
       allOf:
@@ -1128,6 +1162,11 @@ components:
       allOf:
         -
           $ref: '#/components/schemas/BaseResourceUpdate'
+        -
+          type: object
+          properties:
+            state:
+              $ref: '#/components/schemas/ModelVersionState'
     BaseArtifactCreate:
       allOf:
         -
@@ -1534,6 +1573,8 @@ components:
             runtime:
               description: Model runtime.
               type: string
+            state:
+              $ref: '#/components/schemas/InferenceServiceState'
     InferenceServiceCreate:
       description: >-
         An `InferenceService` entity in a `ServingEnvironment` represents a deployed `ModelVersion`

diff --git a/internal/converter/generated/mlmd_openapi_converter.gen.go b/internal/converter/generated/mlmd_openapi_converter.gen.go
diff --git a/internal/converter/generated/openapi_converter.gen.go b/internal/converter/generated/openapi_converter.gen.go
diff --git a/internal/converter/mlmd_converter_util_test.go b/internal/converter/mlmd_converter_util_test.go
@@ -538,6 +538,54 @@ func TestMapMLMDModelArtifactState(t *testing.T) {
 	assertion.Nil(artifactState)
 }
 
+func TestMapRegisteredModelState(t *testing.T) {
+	assertion := setup(t)
+
+	state := MapRegisteredModelState(&proto.Context{Properties: map[string]*proto.Value{
+		"state": {Value: &proto.Value_StringValue{StringValue: string(openapi.REGISTEREDMODELSTATE_LIVE)}},
+	}})
+	assertion.NotNil(state)
+	assertion.Equal(openapi.REGISTEREDMODELSTATE_LIVE, *state)
+
+	state = MapRegisteredModelState(&proto.Context{Properties: map[string]*proto.Value{}})
+	assertion.Nil(state)
+
+	state = MapRegisteredModelState(nil)
+	assertion.Nil(state)
+}
+
+func TestMapModelVersionState(t *testing.T) {
+	assertion := setup(t)
+
+	state := MapModelVersionState(&proto.Context{Properties: map[string]*proto.Value{
+		"state": {Value: &proto.Value_StringValue{StringValue: string(openapi.MODELVERSIONSTATE_LIVE)}},
+	}})
+	assertion.NotNil(state)
+	assertion.Equal(openapi.MODELVERSIONSTATE_LIVE, *state)
+
+	state = MapModelVersionState(&proto.Context{Properties: map[string]*proto.Value{}})
+	assertion.Nil(state)
+
+	state = MapModelVersionState(nil)
+	assertion.Nil(state)
+}
+
+func TestMapInferenceServiceState(t *testing.T) {
+	assertion := setup(t)
+
+	state := MapInferenceServiceState(&proto.Context{Properties: map[string]*proto.Value{
+		"state": {Value: &proto.Value_StringValue{StringValue: string(openapi.INFERENCESERVICESTATE_DEPLOYED)}},
+	}})
+	assertion.NotNil(state)
+	assertion.Equal(openapi.INFERENCESERVICESTATE_DEPLOYED, *state)
+
+	state = MapInferenceServiceState(&proto.Context{Properties: map[string]*proto.Value{}})
+	assertion.Nil(state)
+
+	state = MapInferenceServiceState(nil)
+	assertion.Nil(state)
+}
+
 func TestMapServingEnvironmentType(t *testing.T) {
 	assertion := setup(t)
 

diff --git a/internal/converter/mlmd_openapi_converter.go b/internal/converter/mlmd_openapi_converter.go
@@ -15,10 +15,12 @@ import (
 // goverter:extend MapMLMDCustomProperties
 type MLMDToOpenAPIConverter interface {
 	// goverter:map Properties Description | MapDescription
+	// goverter:map . State | MapRegisteredModelState
 	ConvertRegisteredModel(source *proto.Context) (*openapi.RegisteredModel, error)
 
 	// goverter:map Name | MapNameFromOwned
 	// goverter:map Properties Description | MapDescription
+	// goverter:map . State | MapModelVersionState
 	ConvertModelVersion(source *proto.Context) (*openapi.ModelVersion, error)
 
 	// goverter:map Name | MapNameFromOwned
@@ -42,6 +44,7 @@ type MLMDToOpenAPIConverter interface {
 	// goverter:map Properties ModelVersionId | MapPropertyModelVersionId
 	// goverter:map Properties RegisteredModelId | MapPropertyRegisteredModelId
 	// goverter:map Properties ServingEnvironmentId | MapPropertyServingEnvironmentId
+	// goverter:map . State | MapInferenceServiceState
 	ConvertInferenceService(source *proto.Context) (*openapi.InferenceService, error)
 
 	// goverter:map Name | MapNameFromOwned

diff --git a/internal/converter/mlmd_openapi_converter_util.go b/internal/converter/mlmd_openapi_converter_util.go
@@ -84,6 +84,45 @@ func MapArtifactType(source *proto.Artifact) (string, error) {
 	return "", fmt.Errorf("invalid artifact type found: %v", source.Type)
 }
 
+func MapRegisteredModelState(source *proto.Context) *openapi.RegisteredModelState {
+	if source == nil || source.GetProperties() == nil {
+		return nil
+	}
+
+	state, ok := source.GetProperties()["state"]
+	if !ok {
+		return nil
+	}
+	str := state.GetStringValue()
+	return (*openapi.RegisteredModelState)(&str)
+}
+
+func MapModelVersionState(source *proto.Context) *openapi.ModelVersionState {
+	if source == nil || source.GetProperties() == nil {
+		return nil
+	}
+
+	state, ok := source.GetProperties()["state"]
+	if !ok {
+		return nil
+	}
+	str := state.GetStringValue()
+	return (*openapi.ModelVersionState)(&str)
+}
+
+func MapInferenceServiceState(source *proto.Context) *openapi.InferenceServiceState {
+	if source == nil || source.GetProperties() == nil {
+		return nil
+	}
+
+	state, ok := source.GetProperties()["state"]
+	if !ok {
+		return nil
+	}
+	str := state.GetStringValue()
+	return (*openapi.InferenceServiceState)(&str)
+}
+
 func MapMLMDModelArtifactState(source *proto.Artifact_State) *openapi.ArtifactState {
 	if source == nil {
 		return nil

diff --git a/internal/converter/openapi_mlmd_converter_util.go b/internal/converter/openapi_mlmd_converter_util.go
@@ -129,6 +129,14 @@ func MapRegisteredModelProperties(source *openapi.RegisteredModel) (map[string]*
 				},
 			}
 		}
+
+		if source.State != nil {
+			props["state"] = &proto.Value{
+				Value: &proto.Value_StringValue{
+					StringValue: string(*source.State),
+				},
+			}
+		}
 	}
 	return props, nil
 }
@@ -163,6 +171,15 @@ func MapModelVersionProperties(source *OpenAPIModelWrapper[openapi.ModelVersion]
 				StringValue: *(*source.Model).Name,
 			},
 		}
+
+		if (*source.Model).State != nil {
+			props["state"] = &proto.Value{
+				Value: &proto.Value_StringValue{
+					StringValue: string(*(*source.Model).State),
+				},
+			}
+		}
+
 		// TODO: not available for now
 		props["author"] = &proto.Value{
 			Value: &proto.Value_StringValue{
@@ -317,6 +334,14 @@ func MapInferenceServiceProperties(source *openapi.InferenceService) (map[string
 			}
 		}
 
+		if source.State != nil {
+			props["state"] = &proto.Value{
+				Value: &proto.Value_StringValue{
+					StringValue: string(*source.State),
+				},
+			}
+		}
+
 		registeredModelId, err := StringToInt64(&source.RegisteredModelId)
 		if err != nil {
 			return nil, err

diff --git a/pkg/core/core.go b/pkg/core/core.go
@@ -43,6 +43,7 @@ func NewModelRegistryService(cc grpc.ClientConnInterface) (api.ModelRegistryApi,
 			Name: registeredModelTypeName,
 			Properties: map[string]proto.PropertyType{
 				"description": proto.PropertyType_STRING,
+				"state":       proto.PropertyType_STRING,
 			},
 		},
 	}
@@ -56,6 +57,7 @@ func NewModelRegistryService(cc grpc.ClientConnInterface) (api.ModelRegistryApi,
 				"model_name":  proto.PropertyType_STRING,
 				"version":     proto.PropertyType_STRING,
 				"author":      proto.PropertyType_STRING,
+				"state":       proto.PropertyType_STRING,
 			},
 		},
 	}
@@ -96,6 +98,7 @@ func NewModelRegistryService(cc grpc.ClientConnInterface) (api.ModelRegistryApi,
 				// same information tracked using ParentContext association
 				"serving_environment_id": proto.PropertyType_INT,
 				"runtime":                proto.PropertyType_STRING,
+				"state":                  proto.PropertyType_STRING,
 			},
 		},
 	}