Change state to desiredState for InferenceService (kubeflow#243)

api/openapi/model-registry.yaml

@@ -1576,7 +1576,7 @@ components:
             runtime:
               description: Model runtime.
               type: string
-            state:
+            desiredState:
               $ref: '#/components/schemas/InferenceServiceState'
     InferenceServiceCreate:
       description: >-

internal/converter/mlmd_converter_util_test.go

@@ -544,13 +544,13 @@ func TestMapMLMDModelArtifactState(t *testing.T) {
 func TestMapRegisteredModelState(t *testing.T) {
 	assertion := setup(t)
 
-	state := MapRegisteredModelState(&proto.Context{Properties: map[string]*proto.Value{
+	state := MapRegisteredModelState(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{}})
+	state = MapRegisteredModelState(map[string]*proto.Value{})
 	assertion.Nil(state)
 
 	state = MapRegisteredModelState(nil)
@@ -560,13 +560,13 @@ func TestMapRegisteredModelState(t *testing.T) {
 func TestMapModelVersionState(t *testing.T) {
 	assertion := setup(t)
 
-	state := MapModelVersionState(&proto.Context{Properties: map[string]*proto.Value{
+	state := MapModelVersionState(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{}})
+	state = MapModelVersionState(map[string]*proto.Value{})
 	assertion.Nil(state)
 
 	state = MapModelVersionState(nil)
@@ -576,16 +576,16 @@ func TestMapModelVersionState(t *testing.T) {
 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)}},
-	}})
+	state := MapInferenceServiceDesiredState(map[string]*proto.Value{
+		"desired_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{}})
+	state = MapInferenceServiceDesiredState(map[string]*proto.Value{})
 	assertion.Nil(state)
 
-	state = MapInferenceServiceState(nil)
+	state = MapInferenceServiceDesiredState(nil)
 	assertion.Nil(state)
 }
 
@@ -614,14 +614,16 @@ func TestMapInferenceServiceProperties(t *testing.T) {
 		Runtime:              of("my-runtime"),
 		RegisteredModelId:    "2",
 		ServingEnvironmentId: "3",
+		DesiredState:         openapi.INFERENCESERVICESTATE_DEPLOYED.Ptr(),
 	})
 	assertion.Nil(err)
-	assertion.Equal(5, len(props))
+	assertion.Equal(6, len(props))
 	assertion.Equal("my custom description", props["description"].GetStringValue())
 	assertion.Equal(int64(1), props["model_version_id"].GetIntValue())
 	assertion.Equal("my-runtime", props["runtime"].GetStringValue())
 	assertion.Equal(int64(2), props["registered_model_id"].GetIntValue())
 	assertion.Equal(int64(3), props["serving_environment_id"].GetIntValue())
+	assertion.Equal("DEPLOYED", props["desired_state"].GetStringValue())
 
 	// serving and model id must be provided and must be a valid numeric id
 	_, err = MapInferenceServiceProperties(&openapi.InferenceService{})

internal/converter/mlmd_openapi_converter.go

@@ -15,12 +15,12 @@ import (
 // goverter:extend MapMLMDCustomProperties
 type MLMDToOpenAPIConverter interface {
 	// goverter:map Properties Description | MapDescription
-	// goverter:map . State | MapRegisteredModelState
+	// goverter:map Properties State | MapRegisteredModelState
 	ConvertRegisteredModel(source *proto.Context) (*openapi.RegisteredModel, error)
 
 	// goverter:map Name | MapNameFromOwned
 	// goverter:map Properties Description | MapDescription
-	// goverter:map . State | MapModelVersionState
+	// goverter:map Properties State | MapModelVersionState
 	// goverter:map Properties Author | MapPropertyAuthor
 	ConvertModelVersion(source *proto.Context) (*openapi.ModelVersion, error)
 
@@ -45,7 +45,7 @@ type MLMDToOpenAPIConverter interface {
 	// goverter:map Properties ModelVersionId | MapPropertyModelVersionId
 	// goverter:map Properties RegisteredModelId | MapPropertyRegisteredModelId
 	// goverter:map Properties ServingEnvironmentId | MapPropertyServingEnvironmentId
-	// goverter:map . State | MapInferenceServiceState
+	// goverter:map Properties DesiredState | MapInferenceServiceDesiredState
 	ConvertInferenceService(source *proto.Context) (*openapi.InferenceService, error)
 
 	// goverter:map Name | MapNameFromOwned

internal/converter/mlmd_openapi_converter_util.go

@@ -89,38 +89,26 @@ 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"]
+func MapRegisteredModelState(properties map[string]*proto.Value) *openapi.RegisteredModelState {
+	state, ok := properties["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"]
+func MapModelVersionState(properties map[string]*proto.Value) *openapi.ModelVersionState {
+	state, ok := properties["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"]
+func MapInferenceServiceDesiredState(properties map[string]*proto.Value) *openapi.InferenceServiceState {
+	state, ok := properties["desired_state"]
 	if !ok {
 		return nil
 	}

internal/converter/openapi_converter.go

@@ -76,7 +76,7 @@ type OpenAPIConverter interface {
 	// Ignore all fields that ARE editable
 	// goverter:default InitInferenceServiceWithUpdate
 	// goverter:autoMap Existing
-	// goverter:ignore Id CreateTimeSinceEpoch LastUpdateTimeSinceEpoch Description ExternalID CustomProperties ModelVersionId Runtime State
+	// goverter:ignore Id CreateTimeSinceEpoch LastUpdateTimeSinceEpoch Description ExternalID CustomProperties ModelVersionId Runtime DesiredState
 	OverrideNotEditableForInferenceService(source OpenapiUpdateWrapper[openapi.InferenceService]) (openapi.InferenceService, error)
 
 	// Ignore all fields that ARE editable

internal/converter/openapi_mlmd_converter_util.go

@@ -336,10 +336,10 @@ func MapInferenceServiceProperties(source *openapi.InferenceService) (map[string
 			}
 		}
 
-		if source.State != nil {
-			props["state"] = &proto.Value{
+		if source.DesiredState != nil {
+			props["desired_state"] = &proto.Value{
 				Value: &proto.Value_StringValue{
-					StringValue: string(*source.State),
+					StringValue: string(*source.DesiredState),
 				},
 			}
 		}

pkg/core/core.go

@@ -107,7 +107,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,
+				"desired_state":          proto.PropertyType_STRING,
 			},
 		},
 	}

pkg/core/core_test.go

@@ -2125,7 +2125,7 @@ func (suite *CoreTestSuite) TestCreateInferenceService() {
 	parentResourceId := suite.registerServingEnvironment(service, nil, nil)
 	registeredModelId := suite.registerModel(service, nil, nil)
 	runtime := "model-server"
-	state := openapi.INFERENCESERVICESTATE_DEPLOYED
+	desiredState := openapi.INFERENCESERVICESTATE_DEPLOYED
 
 	eut := &openapi.InferenceService{
 		Name:                 &entityName,
@@ -2134,7 +2134,7 @@ func (suite *CoreTestSuite) TestCreateInferenceService() {
 		ServingEnvironmentId: parentResourceId,
 		RegisteredModelId:    registeredModelId,
 		Runtime:              &runtime,
-		State:                &state,
+		DesiredState:         &desiredState,
 		CustomProperties: &map[string]openapi.MetadataValue{
 			"custom_string_prop": {
 				MetadataStringValue: &openapi.MetadataStringValue{
@@ -2165,7 +2165,7 @@ func (suite *CoreTestSuite) TestCreateInferenceService() {
 	suite.Equal(customString, byId.Contexts[0].CustomProperties["custom_string_prop"].GetStringValue(), "saved custom_string_prop custom property should match the provided one")
 	suite.Equal(entityDescription, byId.Contexts[0].Properties["description"].GetStringValue(), "saved description should match the provided one")
 	suite.Equal(runtime, byId.Contexts[0].Properties["runtime"].GetStringValue(), "saved runtime should match the provided one")
-	suite.Equal(string(state), byId.Contexts[0].Properties["state"].GetStringValue(), "saved state should match the provided one")
+	suite.Equal(string(desiredState), byId.Contexts[0].Properties["desired_state"].GetStringValue(), "saved state should match the provided one")
 	suite.Equalf(*inferenceServiceTypeName, *byId.Contexts[0].Type, "saved context should be of type of %s", *inferenceServiceTypeName)
 
 	getAllResp, err := suite.mlmdClient.GetContexts(context.Background(), &proto.GetContextsRequest{})
@@ -2359,7 +2359,7 @@ func (suite *CoreTestSuite) TestGetInferenceServiceById() {
 		Description:          &entityDescription,
 		ServingEnvironmentId: parentResourceId,
 		RegisteredModelId:    registeredModelId,
-		State:                &state,
+		DesiredState:         &state,
 		CustomProperties: &map[string]openapi.MetadataValue{
 			"custom_string_prop": {
 				MetadataStringValue: &openapi.MetadataStringValue{
@@ -2389,7 +2389,7 @@ func (suite *CoreTestSuite) TestGetInferenceServiceById() {
 	suite.Equal(*getById.Id, *converter.Int64ToString(ctx.Id), "returned id should match the mlmd context one")
 	suite.Equal(*eut.Name, *getById.Name, "saved name should match the provided one")
 	suite.Equal(*eut.ExternalID, *getById.ExternalID, "saved external id should match the provided one")
-	suite.Equal(*eut.State, *getById.State, "saved state should match the provided one")
+	suite.Equal(*eut.DesiredState, *getById.DesiredState, "saved state should match the provided one")
 	suite.Equal(*(*getById.CustomProperties)["custom_string_prop"].MetadataStringValue.StringValue, customString, "saved custom_string_prop custom property should match the provided one")
 }