Merge pull request #74 from kubeflow/main

[pull] main from kubeflow:main

Makefile

@@ -2,7 +2,7 @@
 MKFILE_PATH := $(abspath $(lastword $(MAKEFILE_LIST)))
 PROJECT_PATH := $(patsubst %/,%,$(dir $(MKFILE_PATH)))
 PROJECT_BIN := $(PROJECT_PATH)/bin
-GO := $(PROJECT_BIN)/go1.19
+GO := $(PROJECT_BIN)/go1.21.9
 
 # add tools bin directory
 PATH := $(PROJECT_BIN):$(PATH)
@@ -109,8 +109,8 @@ clean/odh:
 	rm -Rf ./model-registry
 
 bin/go:
-	GOBIN=$(PROJECT_BIN) go install golang.org/dl/go1.19@latest
-	$(PROJECT_BIN)/go1.19 download
+	GOBIN=$(PROJECT_BIN) go install golang.org/dl/go1.21.9@latest
+	$(PROJECT_BIN)/go1.21.9 download
 
 bin/protoc:
 	./scripts/install_protoc.sh

clients/python/src/model_registry/_client.py

@@ -69,11 +69,11 @@ def __init__(
                 server_address, port, user_token
             )
 
-    def _register_model(self, name: str) -> RegisteredModel:
+    def _register_model(self, name: str, **kwargs) -> RegisteredModel:
         if rm := self._api.get_registered_model_by_params(name):
             return rm
 
-        rm = RegisteredModel(name)
+        rm = RegisteredModel(name, **kwargs)
         self._api.upsert_registered_model(rm)
         return rm
 
@@ -109,6 +109,7 @@ def register_model(
         storage_path: str | None = None,
         service_account_name: str | None = None,
         author: str | None = None,
+        owner: str | None = None,
         description: str | None = None,
         metadata: dict[str, ScalarType] | None = None,
     ) -> RegisteredModel:
@@ -132,6 +133,7 @@ def register_model(
             model_format_version: Version of the model format.
             description: Description of the model.
             author: Author of the model. Defaults to the client author.
+            owner: Owner of the model. Defaults to the client author.
             storage_key: Storage key.
             storage_path: Storage path.
             service_account_name: Service account name.
@@ -140,7 +142,7 @@ def register_model(
         Returns:
             Registered model.
         """
-        rm = self._register_model(name)
+        rm = self._register_model(name, owner=owner or self._author)
         mv = self._register_new_version(
             rm,
             version,
@@ -169,6 +171,7 @@ def register_hf_model(
         model_format_name: str,
         model_format_version: str,
         author: str | None = None,
+        owner: str | None = None,
         model_name: str | None = None,
         description: str | None = None,
         git_ref: str = "main",
@@ -187,6 +190,7 @@ def register_hf_model(
             model_format_name: Name of the model format.
             model_format_version: Version of the model format.
             author: Author of the model. Defaults to repo owner.
+            owner: Owner of the model. Defaults to the client author.
             model_name: Name of the model. Defaults to the repo name.
             description: Description of the model.
             git_ref: Git reference to use. Defaults to `main`.
@@ -244,6 +248,7 @@ def register_hf_model(
             model_name or model_info.id,
             source_uri,
             author=author or model_author,
+            owner=owner or self._author,
             version=version,
             model_format_name=model_format_name,
             model_format_version=model_format_version,

clients/python/src/model_registry/types/contexts.py

@@ -135,14 +135,12 @@ class RegisteredModel(BaseContext):
     """
 
     name: str
-    owner: str = None
+    owner: str | None = None
 
     @override
     def map(self, type_id: int) -> Context:
         mlmd_obj = super().map(type_id)
-        props = {
-            "owner": self.owner
-        }
+        props = {"owner": self.owner}
         self._map_props(props, mlmd_obj.properties)
         return mlmd_obj
 

clients/python/tests/conftest.py

@@ -160,6 +160,7 @@ def store_wrapper(plain_wrapper: MLMDStore) -> MLMDStore:
         [
             "description",
             "state",
+            "owner",
         ],
     )
 

go.mod

@@ -1,6 +1,6 @@
 module github.com/kubeflow/model-registry
 
-go 1.19
+go 1.21
 
 require (
 	github.com/go-chi/chi/v5 v5.0.12