fix(hub): add params to register_model()

kubeflow/hub/__init__.py

@@ -13,5 +13,6 @@
 # limitations under the License.
 
 from kubeflow.hub.api.model_registry_client import ModelRegistryClient
+from kubeflow.hub.types.types import StorageConfig
 
-__all__ = ["ModelRegistryClient"]
+__all__ = ["ModelRegistryClient", "StorageConfig"]

kubeflow/hub/api/model_registry_client.py

@@ -17,6 +17,8 @@
 from collections.abc import Iterator, Mapping
 from typing import TYPE_CHECKING
 
+from kubeflow.hub.types.types import StorageConfig
+
 if TYPE_CHECKING:
     from model_registry.types import (
         ModelArtifact,
@@ -107,17 +109,18 @@ def register_model(
         owner: str | None = None,
         version_description: str | None = None,
         metadata: Mapping[str, SupportedTypes] | None = None,
+        storage_config: StorageConfig | None = None,
     ) -> RegisteredModel:
         """Register a model.
 
         This registers a model in the model registry. The model is not downloaded,
         and has to be stored prior to registration.
 
-        Most models can be registered using their URI, along with optional
-        connection-specific parameters, `storage_key` and `storage_path` or,
-        simply a `service_account_name`. URI builder utilities are recommended
-        when referring to specialized storage; for example `utils.s3_uri_from`
-        helper when using S3 object storage data connections.
+        Most models can be registered using their URI, along with an optional
+        `storage_config` describing how KServe should fetch the model at
+        inference time. URI builder utilities are recommended when referring to
+        specialized storage; for example `utils.s3_uri_from` when using S3
+        object storage data connections.
 
         Args:
             name: Name of the model.
@@ -132,10 +135,15 @@ def register_model(
             owner: Owner of the model. Defaults to the client author.
             version_description: Description of the model version.
             metadata: Additional version metadata.
+            storage_config: Storage credentials for the model artifact. Groups
+                           `storage_key`, `storage_path`, and
+                           `service_account_name` used by KServe's
+                           StorageInitializer. See `StorageConfig` for details.
 
         Returns:
             Registered model.
         """
+        storage = storage_config or StorageConfig()
         return self._registry.register_model(
             name=name,
             uri=uri,
@@ -146,6 +154,9 @@ def register_model(
             owner=owner,
             description=version_description,
             metadata=metadata,
+            storage_key=storage.storage_key,
+            storage_path=storage.storage_path,
+            service_account_name=storage.service_account_name,
         )
 
     def update_model(self, model: RegisteredModel) -> RegisteredModel:

kubeflow/hub/api/model_registry_client_test.py

@@ -209,17 +209,83 @@ def test_init(test_case, monkeypatch):
                 "model_format_version": "1.0",
                 "version": "v1",
             },
+            expected_output={
+                "storage_key": None,
+                "storage_path": None,
+                "service_account_name": None,
+            },
+        ),
+        TestCase(
+            name="register_model forwards full StorageConfig fields",
+            expected_status=SUCCESS,
+            config={
+                "name": "test",
+                "uri": "s3://bucket/model",
+                "version": "v1",
+                "storage_config_kwargs": {
+                    "storage_key": "my-s3-secret",
+                    "storage_path": "models/v1",
+                    "service_account_name": "model-sa",
+                },
+            },
+            expected_output={
+                "storage_key": "my-s3-secret",
+                "storage_path": "models/v1",
+                "service_account_name": "model-sa",
+            },
+        ),
+        TestCase(
+            name="register_model forwards partial StorageConfig (storage_key only)",
+            expected_status=SUCCESS,
+            config={
+                "name": "test",
+                "uri": "s3://bucket/model",
+                "version": "v1",
+                "storage_config_kwargs": {"storage_key": "my-s3-secret"},
+            },
+            expected_output={
+                "storage_key": "my-s3-secret",
+                "storage_path": None,
+                "service_account_name": None,
+            },
+        ),
+        TestCase(
+            name="register_model forwards partial StorageConfig (service_account_name only)",
+            expected_status=SUCCESS,
+            config={
+                "name": "test",
+                "uri": "s3://bucket/model",
+                "version": "v1",
+                "storage_config_kwargs": {"service_account_name": "model-sa"},
+            },
+            expected_output={
+                "storage_key": None,
+                "storage_path": None,
+                "service_account_name": "model-sa",
+            },
         ),
     ],
 )
 def test_register_model(test_case, client, mock_registry):
     """Test register_model delegates to ModelRegistry.register_model."""
 
+    from kubeflow.hub.types.types import StorageConfig
+
+    config = dict(test_case.config)
+    storage_kwargs = config.pop("storage_config_kwargs", None)
+    if storage_kwargs is not None:
+        config["storage_config"] = StorageConfig(**storage_kwargs)
+
     try:
-        client.register_model(**test_case.config)
+        client.register_model(**config)
 
         assert test_case.expected_status == SUCCESS
         assert mock_registry.register_model.called
+        forwarded = mock_registry.register_model.call_args[1]
+        for field, expected in test_case.expected_output.items():
+            assert forwarded[field] == expected, (
+                f"expected {field}={expected!r}, got {forwarded[field]!r}"
+            )
 
     except Exception as e:
         assert test_case.expected_status == FAILED

kubeflow/hub/types/types.py

@@ -0,0 +1,61 @@
+# Copyright 2025 The Kubeflow Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+
+@dataclass
+class StorageConfig:
+    """Storage credentials for a model artifact.
+
+    Groups the storage-related fields that KServe's StorageInitializer uses
+    to pull a model at inference time. Pass an instance to
+    ``ModelRegistryClient.register_model`` via the ``storage_config`` keyword
+    argument.
+
+    All fields default to ``None``; populate only the fields relevant to your
+    storage backend:
+
+    - Secret-based (S3 / MinIO data connections): set ``storage_key`` and
+      optionally ``storage_path``.
+    - IRSA / Workload Identity: set ``service_account_name``.
+
+    Args:
+        storage_key: Name of the Kubernetes Secret containing storage
+            credentials (e.g., an S3 access-key/secret pair for MinIO or AWS).
+        storage_path: Subpath within the storage bucket where the model
+            resides.
+        service_account_name: Kubernetes ServiceAccount name annotated for
+            IRSA or Workload Identity (cloud-native auth without a Secret).
+
+    Example:
+        from kubeflow.hub import ModelRegistryClient, StorageConfig
+
+        client = ModelRegistryClient("https://registry.example.com")
+        client.register_model(
+            "my-model",
+            uri="s3://my-bucket/models/v1",
+            version="1.0",
+            storage_config=StorageConfig(
+                storage_key="my-s3-secret",
+                storage_path="models/v1",
+            ),
+        )
+    """
+
+    storage_key: str | None = None
+    storage_path: str | None = None
+    service_account_name: str | None = None