Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

add owner to RegisteredModel logical model #70

Merged
merged 4 commits into from
Apr 29, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions api/openapi/model-registry.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -1033,6 +1033,8 @@ components:
- $ref: "#/components/schemas/BaseResourceUpdate"
- type: object
properties:
owner:
type: string
state:
$ref: "#/components/schemas/RegisteredModelState"
ModelVersion:
Expand Down
12 changes: 12 additions & 0 deletions clients/python/src/model_registry/types/contexts.py
Original file line number Diff line number Diff line change
Expand Up @@ -129,11 +129,22 @@ class RegisteredModel(BaseContext):

Attributes:
name: Registered model name.
owner: Owner of this Registered Model.
description: Description of the object.
external_id: Customizable ID. Has to be unique among instances of the same type.
"""

name: str
owner: str = None

@override
def map(self, type_id: int) -> Context:
mlmd_obj = super().map(type_id)
props = {
"owner": self.owner
}
self._map_props(props, mlmd_obj.properties)
return mlmd_obj

@classmethod
@override
Expand All @@ -142,4 +153,5 @@ def unmap(cls, mlmd_obj: Context) -> RegisteredModel:
assert isinstance(
py_obj, RegisteredModel
), f"Expected RegisteredModel, got {type(py_obj)}"
py_obj.owner = mlmd_obj.properties["owner"].string_value
return py_obj
24 changes: 23 additions & 1 deletion clients/python/tests/conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -89,7 +89,9 @@ def teardown():

request.addfinalizer(teardown)

return MLMDStore(mlmd_conn)
to_return = MLMDStore(mlmd_conn)
sanity_check_mlmd_connection_to_db(to_return)
return to_return


def set_type_attrs(mlmd_obj: ProtoTypeType, name: str, props: list[str]):
Expand All @@ -99,6 +101,26 @@ def set_type_attrs(mlmd_obj: ProtoTypeType, name: str, props: list[str]):
return mlmd_obj


def sanity_check_mlmd_connection_to_db(overview: MLMDStore):
# sanity check before each test: connect to MLMD directly, and dry-run any of the gRPC (read) operations;
# on newer Podman might delay in recognising volume mount files for sqlite3 db,
# hence in case of error "Cannot connect sqlite3 database: unable to open database file" make some retries.
retry_count = 0
while retry_count < 3:
retry_count += 1
try:
overview._mlmd_store.get_artifact_types()
return
except Exception as e:
if str(e) == "Cannot connect sqlite3 database: unable to open database file":
time.sleep(1)
else:
msg = "Failed to sanity check before each test, another type of error detected."
raise RuntimeError(msg) from e
msg = "Failed to sanity check before each test."
raise RuntimeError(msg)


@pytest.fixture()
def store_wrapper(plain_wrapper: MLMDStore) -> MLMDStore:
ma_type = set_type_attrs(
Expand Down
39 changes: 38 additions & 1 deletion clients/python/tests/types/test_context_mapping.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@

import pytest
from ml_metadata.proto import Context
from model_registry.types import ContextState, ModelVersion
from model_registry.types import ContextState, ModelVersion, RegisteredModel

from .. import Mapped

Expand Down Expand Up @@ -95,3 +95,40 @@ def test_full_model_version_unmapping(full_model_version: Mapped):
assert unmapped_version.author == py_version.author
assert unmapped_version.state == py_version.state
assert unmapped_version.metadata == py_version.metadata


@pytest.fixture()
def minimal_registered_model() -> Mapped:
proto_version = Context(
name="mnist",
type_id=1,
external_id="test_external_id")
proto_version.properties["description"].string_value = "test description"
proto_version.properties["state"].string_value = "LIVE"
proto_version.properties["owner"].string_value = "my owner"

py_regmodel = RegisteredModel(name="mnist",
owner="my owner",
external_id="test_external_id",
description="test description")
return Mapped(proto_version, py_regmodel)


def test_minimal_registered_model_mapping(minimal_registered_model: Mapped):
mapped_version = minimal_registered_model.py.map(1)
proto_version = minimal_registered_model.proto
assert mapped_version.name == proto_version.name
assert mapped_version.type_id == proto_version.type_id
assert mapped_version.external_id == proto_version.external_id
assert mapped_version.properties == proto_version.properties
assert mapped_version.custom_properties == proto_version.custom_properties


def test_minimal_registered_model_unmapping(minimal_registered_model: Mapped):
unmapped_regmodel = RegisteredModel.unmap(minimal_registered_model.proto)
py_regmodel = minimal_registered_model.py
assert unmapped_regmodel.name == py_regmodel.name
assert unmapped_regmodel.owner == py_regmodel.owner
assert unmapped_regmodel.description == py_regmodel.description
assert unmapped_regmodel.external_id == py_regmodel.external_id
assert unmapped_regmodel.state == py_regmodel.state
2 changes: 2 additions & 0 deletions docs/logical_model.md
Original file line number Diff line number Diff line change
Expand Up @@ -442,10 +442,12 @@ This diagram summarizes the relationship between the entities:
classDiagram
class RegisteredModel{
+String name
+String owner
+Map customProperties
}
class ModelVersion{
+String name
+String author
+Map customProperties
}
class Artifact{
Expand Down
1 change: 1 addition & 0 deletions internal/converter/generated/mlmd_openapi_converter.gen.go

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

12 changes: 12 additions & 0 deletions internal/converter/generated/openapi_converter.gen.go

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

14 changes: 14 additions & 0 deletions internal/converter/mlmd_converter_util_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -414,6 +414,20 @@ func TestMapDescription(t *testing.T) {
assertion.Equal("my-description", *extracted)
}

func TestMapOwner(t *testing.T) {
assertion := setup(t)

extracted := MapOwner(map[string]*proto.Value{
"owner": {
Value: &proto.Value_StringValue{
StringValue: "my-owner",
},
},
})

assertion.Equal("my-owner", *extracted)
}

func TestPropertyRuntime(t *testing.T) {
assertion := setup(t)

Expand Down
1 change: 1 addition & 0 deletions internal/converter/mlmd_openapi_converter.go
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@ import (
// goverter:extend MapMLMDCustomProperties
type MLMDToOpenAPIConverter interface {
// goverter:map Properties Description | MapDescription
// goverter:map Properties Owner | MapOwner
// goverter:map Properties State | MapRegisteredModelState
ConvertRegisteredModel(source *proto.Context) (*openapi.RegisteredModel, error)

Expand Down
4 changes: 4 additions & 0 deletions internal/converter/mlmd_openapi_converter_util.go
Original file line number Diff line number Diff line change
Expand Up @@ -211,6 +211,10 @@ func MapDescription(properties map[string]*proto.Value) *string {
return MapStringProperty(properties, "description")
}

func MapOwner(properties map[string]*proto.Value) *string {
return MapStringProperty(properties, "owner")
}

func MapModelArtifactFormatName(properties map[string]*proto.Value) *string {
return MapStringProperty(properties, "model_format_name")
}
Expand Down
2 changes: 1 addition & 1 deletion internal/converter/openapi_converter.go
Original file line number Diff line number Diff line change
Expand Up @@ -52,7 +52,7 @@ type OpenAPIConverter interface {
// Ignore all fields that ARE editable
// goverter:default InitRegisteredModelWithUpdate
// goverter:autoMap Existing
// goverter:ignore Id CreateTimeSinceEpoch LastUpdateTimeSinceEpoch Description ExternalId CustomProperties State
// goverter:ignore Id CreateTimeSinceEpoch LastUpdateTimeSinceEpoch Description ExternalId CustomProperties State Owner
OverrideNotEditableForRegisteredModel(source OpenapiUpdateWrapper[openapi.RegisteredModel]) (openapi.RegisteredModel, error)

// Ignore all fields that ARE editable
Expand Down
8 changes: 8 additions & 0 deletions internal/converter/openapi_mlmd_converter_util.go
Original file line number Diff line number Diff line change
Expand Up @@ -123,6 +123,14 @@ func PrefixWhenOwned(ownerId *string, entityName string) string {
func MapRegisteredModelProperties(source *openapi.RegisteredModel) (map[string]*proto.Value, error) {
props := make(map[string]*proto.Value)
if source != nil {
if source.Owner != nil {
props["owner"] = &proto.Value{
Value: &proto.Value_StringValue{
StringValue: *source.Owner,
},
}
}

if source.Description != nil {
props["description"] = &proto.Value{
Value: &proto.Value_StringValue{
Expand Down
1 change: 1 addition & 0 deletions internal/mlmdtypes/mlmdtypes.go
Original file line number Diff line number Diff line change
Expand Up @@ -44,6 +44,7 @@ func CreateMLMDTypes(cc grpc.ClientConnInterface, nameConfig MLMDTypeNamesConfig
Name: &nameConfig.RegisteredModelTypeName,
Properties: map[string]proto.PropertyType{
"description": proto.PropertyType_STRING,
"owner": proto.PropertyType_STRING,
"state": proto.PropertyType_STRING,
},
},
Expand Down
35 changes: 34 additions & 1 deletion internal/testutils/test_container_utils.go
Original file line number Diff line number Diff line change
@@ -1,10 +1,13 @@
package testutils

import (
"bytes"
"context"
"encoding/json"
"errors"
"fmt"
"os"
"os/exec"
"testing"

"github.com/kubeflow/model-registry/internal/ml_metadata/proto"
Expand All @@ -15,7 +18,7 @@ import (
)

const (
useProvider = testcontainers.ProviderDefault // or explicit to testcontainers.ProviderPodman if needed
defaultProvider = testcontainers.ProviderDefault // or explicit to testcontainers.ProviderPodman if needed
mlmdImage = "gcr.io/tfx-oss-public/ml_metadata_store_server:1.14.0"
sqliteFile = "metadata.sqlite.db"
testConfigFolder = "test/config/ml-metadata"
Expand Down Expand Up @@ -93,6 +96,13 @@ func SetupMLMetadataTestContainer(t *testing.T) (*grpc.ClientConn, proto.Metadat
WaitingFor: wait.ForLog("Server listening on"),
}

useProvider := defaultProvider
if useProvider == testcontainers.ProviderDefault { // user did not override
if tryDetectPodmanRunning() {
t.Log("Podman running detected! Will instruct TestContainers to use Podman explicitly.") // see https://github.com/testcontainers/testcontainers-go/issues/2264#issuecomment-2062092012
useProvider = testcontainers.ProviderPodman
}
}
mlmdgrpc, err := testcontainers.GenericContainer(ctx, testcontainers.GenericContainerRequest{
ProviderType: useProvider,
ContainerRequest: req,
Expand Down Expand Up @@ -137,3 +147,26 @@ func SetupMLMetadataTestContainer(t *testing.T) (*grpc.ClientConn, proto.Metadat
}
}
}

// simple utility function with heuristic to detect if Podman is running
func tryDetectPodmanRunning() bool {
cmd := exec.Command("podman", "machine", "info", "--format", "json")
var out bytes.Buffer
cmd.Stdout = &out
err := cmd.Run()
if err != nil {
return false
}
output := out.Bytes()
type MachineInfo struct {
Host struct {
MachineState string `json:"MachineState"`
} `json:"Host"`
}
var info MachineInfo
err = json.Unmarshal(output, &info)
if err != nil {
return false
}
return info.Host.MachineState == "Running"
}
Loading
Loading