diff --git a/Makefile b/Makefile index c0b6e681..a5f9b6b6 100644 --- a/Makefile +++ b/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 diff --git a/clients/python/src/model_registry/_client.py b/clients/python/src/model_registry/_client.py index f8095684..6c70726a 100644 --- a/clients/python/src/model_registry/_client.py +++ b/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, diff --git a/clients/python/src/model_registry/types/contexts.py b/clients/python/src/model_registry/types/contexts.py index 04fbd061..ecd083c7 100644 --- a/clients/python/src/model_registry/types/contexts.py +++ b/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 diff --git a/clients/python/tests/conftest.py b/clients/python/tests/conftest.py index 9cb4f397..b5ceb1e2 100644 --- a/clients/python/tests/conftest.py +++ b/clients/python/tests/conftest.py @@ -160,6 +160,7 @@ def store_wrapper(plain_wrapper: MLMDStore) -> MLMDStore: [ "description", "state", + "owner", ], ) diff --git a/go.mod b/go.mod index 7ab47dc4..9ae3bcaf 100644 --- a/go.mod +++ b/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