5.0 KiB
Model Registry
⚠️ NOTE: UNSTABLE API ⚠️ This library is provided as a convenience for Kubeflow Model Registry developers. If you are not actively involved in the development of Model Registry, please prefer the REST API.
Getting Started
Model Registry is a high level Go library client for recording and retrieving metadata associated with ML developer and data scientist workflows that provides a high-level model metadata registry API.
You can use Model Registry to manage convenient type definitions for models, model versions, artifacts, serving environments, inference services, and serve models.
Prerequisites
- Go >= 1.24
Install it using:
go get github.com/kubeflow/model-registry
Assuming that an Model Registry server is running at localhost:8080, you can connect with it using:
import (
"net/http"
"github.com/kubeflow/model-registry/pkg/openapi"
)
cfg := &openapi.Configuration{
HTTPClient: http.DefaultClient,
Servers: openapi.ServerConfigurations{
{
URL: "http://localhost:8080",
},
},
}
client := openapi.NewAPIClient(cfg)
Once the connection is established, you can create a ModelRegistryServiceAPI service:
service := client.ModelRegistryServiceAPI
Example usage
After setting up your Model Registry Service, you can try some of the following:
Creating models
Create a RegisteredModel
modelName := "MODEL_NAME"
modelDescription := "MODEL_DESCRIPTION"
// register a new model
registeredModel, _, err := service.CreateRegisteredModel(ctx).RegisteredModelCreate(openapi.RegisteredModelCreate{
Name: modelName,
Description: &modelDescription,
}).Execute()
if err != nil {
return nil, fmt.Errorf("error registering model: %w", err)
}
Create a ModelVersion for the registered model
versionName := "VERSION_NAME"
versionDescription := "VERSION_DESCRIPTION"
versionScore := 0.83
// register model version
modelVersion, _, err := service.CreateModelVersion(ctx).ModelVersionCreate(openapi.ModelVersionCreate{
Name: versionName,
Description: &versionDescription,
RegisteredModelId: *registeredModel.Id,
CustomProperties: &map[string]openapi.MetadataValue{
"score": {
MetadataDoubleValue: &openapi.MetadataDoubleValue{
DoubleValue: &versionScore,
},
},
},
}).Execute()
if err != nil {
return nil, fmt.Errorf("unable to create model version: %w", err)
}
Create a ModelArtifact for the version
artifactName := "ARTIFACT_NAME"
artifactDescription := "ARTIFACT_DESCRIPTION"
artifactUri := "ARTIFACT_URI"
artifactType := "model-artifact"
// register model artifact
modelArtifact, _, err = service.UpsertModelVersionArtifact(ctx, *modelVersion.Id).Artifact(openapi.Artifact{
ModelArtifact: &openapi.ModelArtifact{
Name: &artifactName,
Description: &artifactDescription,
Uri: &artifactUri,
ArtifactType: &artifactType,
},
}).Execute()
if err != nil {
return nil, fmt.Errorf("unable to create model artifact: %w", err)
}
Updating models
The Model Registry Service provides Upsert* methods for all supported models, meaning that you can use them to
insert a new model, or update it:
artifactName := "ARTIFACT_NAME"
artifactDescription := "ARTIFACT_DESCRIPTION"
artifactUri := "ARTIFACT_URI"
artifactType := "model-artifact"
// register model artifact
modelArtifact, _, err = service.UpsertModelVersionArtifact(ctx, *modelVersion.Id).Artifact(openapi.Artifact{
ModelArtifact: &openapi.ModelArtifact{
Name: &artifactName,
Description: &artifactDescription,
Uri: &artifactUri,
ArtifactType: &artifactType,
},
}).Execute()
if err != nil {
return nil, fmt.Errorf("unable to create model artifact: %w", err)
}
// update model artifact
newDescription := "update it!"
modelArtifact, _, err = service.UpsertModelVersionArtifact(ctx, *modelVersion.Id).Artifact(openapi.Artifact{
ModelArtifact: &openapi.ModelArtifact{
Name: &artifactName,
Description: &newDescription,
Uri: &artifactUri,
ArtifactType: &artifactType,
},
}).Execute()
if err != nil {
return nil, fmt.Errorf("unable to update model artifact: %w", err)
}
Querying models
Get a RegisteredModel by name:
modelName := "MODEL_NAME"
registeredModel, _, err := service.FindRegisteredModel(ctx).Name(modelName).Execute()
if err != nil {
return nil, fmt.Errorf("unable to find registered model: %w", err)
}
Get all ModelVersion associated to a specific registered model
allVersions, _, err := service.FindModelVersion(ctx).ParentResourceId(*regModelFound.Id).Execute()
if err != nil {
return nil, fmt.Errorf("unable to find model version: %w", err)
}
API documentation
Check out the Model Registry Service interface and the core layer implementation for additional details.