model-registry/docs/mr_go_library.md

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.