add owner to RegisteredModel logical model (#70)

* wip Signed-off-by: Matteo Mortari <[email protected]> * complete with e2e testing Signed-off-by: Matteo Mortari <[email protected]> * linting Signed-off-by: Matteo Mortari <[email protected]> * update docs Signed-off-by: Matteo Mortari <[email protected]> --------- Signed-off-by: Matteo Mortari <[email protected]>
kubeflow · Apr 29, 2024 · ea66c4e · ea66c4e
1 parent 972b7cd
commit ea66c4e
Show file tree

Hide file tree

Showing 19 changed files with 334 additions and 37 deletions.
diff --git a/api/openapi/model-registry.yaml b/api/openapi/model-registry.yaml
@@ -1033,6 +1033,8 @@ components:
         - $ref: "#/components/schemas/BaseResourceUpdate"
         - type: object
           properties:
+            owner:
+              type: string
             state:
               $ref: "#/components/schemas/RegisteredModelState"
     ModelVersion:

diff --git a/clients/python/src/model_registry/types/contexts.py b/clients/python/src/model_registry/types/contexts.py
@@ -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
@@ -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
diff --git a/clients/python/tests/conftest.py b/clients/python/tests/conftest.py
@@ -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]):
@@ -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(

diff --git a/clients/python/tests/types/test_context_mapping.py b/clients/python/tests/types/test_context_mapping.py
@@ -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
 
@@ -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
diff --git a/docs/logical_model.md b/docs/logical_model.md
@@ -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{

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
@@ -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)
 

diff --git a/internal/converter/mlmd_openapi_converter.go b/internal/converter/mlmd_openapi_converter.go
@@ -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)
 

diff --git a/internal/converter/mlmd_openapi_converter_util.go b/internal/converter/mlmd_openapi_converter_util.go
@@ -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")
 }

diff --git a/internal/converter/openapi_converter.go b/internal/converter/openapi_converter.go
@@ -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

diff --git a/internal/converter/openapi_mlmd_converter_util.go b/internal/converter/openapi_mlmd_converter_util.go
@@ -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{

diff --git a/internal/mlmdtypes/mlmdtypes.go b/internal/mlmdtypes/mlmdtypes.go
@@ -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,
 			},
 		},

diff --git a/internal/testutils/test_container_utils.go b/internal/testutils/test_container_utils.go
@@ -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"
@@ -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"
@@ -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,
@@ -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"
+}