From 3bc6bde6e557a626356603bf7baa127137663e44 Mon Sep 17 00:00:00 2001 From: Sidney Glinton Date: Wed, 6 Aug 2025 11:13:34 -0400 Subject: [PATCH 1/2] feat: add ModelRegistry as importable client Signed-off-by: Sidney Glinton --- python/kubeflow/model_registry/__init__.py | 14 ++++++++++++++ python/pyproject.toml | 3 +++ 2 files changed, 17 insertions(+) create mode 100644 python/kubeflow/model_registry/__init__.py diff --git a/python/kubeflow/model_registry/__init__.py b/python/kubeflow/model_registry/__init__.py new file mode 100644 index 00000000..f6f0c0f2 --- /dev/null +++ b/python/kubeflow/model_registry/__init__.py @@ -0,0 +1,14 @@ +# Re-export ModelRegistry from the standalone model-registry package +# This allows users to import from kubeflow.model_registry while +# maintaining backward compatibility with the standalone package + +try: + from model_registry import ModelRegistry as ModelRegistryClient +except ImportError as e: + raise ImportError( + "model-registry package is required. Install with: pip install model-registry" + ) from e + +__all__ = [ + "ModelRegistryClient", +] diff --git a/python/pyproject.toml b/python/pyproject.toml index 6f0e8aa9..b331e5a3 100644 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -31,6 +31,9 @@ dependencies = [ "kubeflow-trainer-api>=2.0.0", ] [project.optional-dependencies] +model-registry = [ + "model-registry>=0.2.21", +] dev = [ "pytest>=7.0", "pytest-mock>=3.10", From f2e6a9a64a7deb00fee519690d5803b00405767b Mon Sep 17 00:00:00 2001 From: Sidney Glinton Date: Wed, 6 Aug 2025 11:42:01 -0400 Subject: [PATCH 2/2] feat: add tests to validate importing Signed-off-by: Sidney Glinton --- .../model_registry/tests/test_import.py | 102 ++++++++++++++++++ 1 file changed, 102 insertions(+) create mode 100644 python/kubeflow/model_registry/tests/test_import.py diff --git a/python/kubeflow/model_registry/tests/test_import.py b/python/kubeflow/model_registry/tests/test_import.py new file mode 100644 index 00000000..cdaaf288 --- /dev/null +++ b/python/kubeflow/model_registry/tests/test_import.py @@ -0,0 +1,102 @@ +import pytest +import sys + + +@pytest.mark.model_registry +def test_import_model_registry_available(): + """Test that ModelRegistry can be imported when model-registry package is available.""" + try: + from kubeflow.model_registry import ModelRegistry, ModelRegistryClient + from model_registry import ModelRegistry as OriginalModelRegistry + + # Verify they are the same class + assert ModelRegistry is OriginalModelRegistry + assert ModelRegistryClient is OriginalModelRegistry + + except ImportError: + pytest.skip("model-registry package is not installed") + + +@pytest.mark.model_registry +def test_import_model_registry_client_available(): + """Test that ModelRegistryClient is an alias for ModelRegistry.""" + try: + from kubeflow.model_registry import ModelRegistryClient, ModelRegistry + + # Verify ModelRegistryClient is an alias for ModelRegistry + assert ModelRegistryClient is ModelRegistry + + except ImportError: + pytest.skip("model-registry package is not installed") + + +@pytest.mark.model_registry +def test_import_error_when_model_registry_not_installed(monkeypatch): + """Test that a helpful error is raised when model-registry is not available.""" + # Remove the module from sys.modules if it exists + modules_to_remove = ["kubeflow.model_registry", "model_registry", "model_registry._client"] + original_modules = {} + + for module_name in modules_to_remove: + if module_name in sys.modules: + original_modules[module_name] = sys.modules[module_name] + del sys.modules[module_name] + + # Store the original import function to avoid recursion + original_import = __import__ + + def mock_import(name, *args, **kwargs): + if name == "model_registry": + raise ImportError("No module named 'model_registry'") + return original_import(name, *args, **kwargs) + + monkeypatch.setattr("builtins.__import__", mock_import) + + try: + with pytest.raises(ImportError) as exc_info: + import kubeflow.model_registry # noqa: F401 + + # Verify the error message is helpful + assert "model-registry package is required" in str(exc_info.value) + assert "pip install model-registry" in str(exc_info.value) + + finally: + # Restore the original modules + for module_name, module in original_modules.items(): + sys.modules[module_name] = module + + +@pytest.mark.model_registry +def test_all_exports_available(): + """Test that all expected exports are available in the module.""" + try: + import kubeflow.model_registry as mr_module + + # Check that __all__ contains expected items + expected_exports = ["ModelRegistry", "ModelRegistryClient"] + assert hasattr(mr_module, "__all__") + assert set(mr_module.__all__) == set(expected_exports) + + # Check that all items in __all__ are actually available + for export in expected_exports: + assert hasattr(mr_module, export) + + except ImportError: + pytest.skip("model-registry package is not installed") + + +@pytest.mark.model_registry +def test_import_via_star_import(): + """Test that star imports work correctly.""" + try: + # This is a bit tricky to test directly, so we'll check the namespace + import kubeflow.model_registry as mr_module + + # Simulate what 'from kubeflow.model_registry import *' would do + star_imports = {name: getattr(mr_module, name) for name in mr_module.__all__} + + assert "ModelRegistryClient" in star_imports + assert star_imports["ModelRegistry"] is star_imports["ModelRegistryClient"] + + except ImportError: + pytest.skip("model-registry package is not installed")