From ad0baf5b89206a63ffbd109d33184a7405f3d100 Mon Sep 17 00:00:00 2001 From: Matteo Mortari Date: Sun, 11 Aug 2024 13:35:27 +0200 Subject: [PATCH] py: add listener mechanism (#8) * py: add listener mechanism Signed-off-by: tarilabs * add e2e test scenarios Signed-off-by: tarilabs * wire GHA jobs Signed-off-by: tarilabs * linting and minor changes Signed-off-by: tarilabs --------- Signed-off-by: tarilabs --- .github/workflows/e2e.yaml | 63 ++++++++- .gitignore | 2 +- Makefile | 8 +- e2e/deploy_model_registry.sh | 29 ++++ e2e/model-registry/README.md | 1 + .../config/ml-metadata/conn_config.pb | 6 + e2e/model-registry/docker-compose.yaml | 19 +++ omlmd/helpers.py | 19 ++- omlmd/listener.py | 27 ++++ omlmd/model_metadata.py | 13 +- omlmd/provider.py | 26 ++++ pyproject.toml | 1 + tests/conftest.py | 24 ++-- tests/test_e2e_model_registry.py | 126 ++++++++++++++++++ tests/test_helpers.py | 29 +++- tests/test_omlmd.py | 18 +++ 16 files changed, 395 insertions(+), 16 deletions(-) create mode 100755 e2e/deploy_model_registry.sh create mode 100644 e2e/model-registry/README.md create mode 100644 e2e/model-registry/config/ml-metadata/conn_config.pb create mode 100644 e2e/model-registry/docker-compose.yaml create mode 100644 omlmd/listener.py create mode 100644 tests/test_e2e_model_registry.py diff --git a/.github/workflows/e2e.yaml b/.github/workflows/e2e.yaml index 1257514..b1bba89 100644 --- a/.github/workflows/e2e.yaml +++ b/.github/workflows/e2e.yaml @@ -1,4 +1,4 @@ -name: E2E testing +name: E2E on: push: @@ -8,6 +8,7 @@ on: jobs: e2e-distribution-registry: + name: E2E using CNCF Distribution Registry runs-on: ubuntu-latest steps: - name: Checkout repository @@ -32,10 +33,36 @@ jobs: - name: Run E2E tests run: | make test-e2e - - name: Run E2E tests for CLI # only running once using distribution/registry intentionally + e2e-cli: + name: E2E of CLI (using Distribution Registry) + needs: + - e2e-distribution-registry # avoid rely on distribution registry for this other e2e if failed. + runs-on: ubuntu-latest + steps: + - name: Checkout repository + uses: actions/checkout@v3 + - name: Set up Python + uses: actions/setup-python@v4 + with: + python-version: '3.9' + - name: Install Poetry + run: | + pipx install poetry + - name: Install dependencies + run: | + make install + - name: Start Kind Cluster + uses: helm/kind-action@v1 + with: + cluster_name: kind + - name: Start distribution-registry + run: | + ./e2e/deploy_distribution_registry.sh + - name: Run E2E tests for CLI run: | ./e2e/test_cli.sh e2e-zot: + name: E2E using CNCF Zot runs-on: ubuntu-latest steps: - name: Checkout repository @@ -61,6 +88,7 @@ jobs: run: | make test-e2e e2e-quay-lite: + name: E2E using Quay (-lite) runs-on: ubuntu-latest steps: - name: Checkout repository @@ -94,3 +122,34 @@ jobs: - name: Run E2E tests run: | make test-e2e + e2e-kubeflow-model-registry: + name: E2E of Kubeflow Model Registry (using Distribution Registry) + needs: + - e2e-distribution-registry # avoid rely on distribution registry for this other e2e if failed. + runs-on: ubuntu-latest + steps: + - name: Checkout repository + uses: actions/checkout@v3 + - name: Set up Python + uses: actions/setup-python@v4 + with: + python-version: '3.9' + - name: Install Poetry + run: | + pipx install poetry + - name: Install dependencies + run: | + make install + - name: Start Kind Cluster + uses: helm/kind-action@v1 + with: + cluster_name: kind + - name: Start distribution-registry + run: | + ./e2e/deploy_distribution_registry.sh + - name: Deploy Kubeflow Model Registry in KinD cluster + run: | + ./e2e/deploy_model_registry.sh + - name: Run Kubeflow Model Registry E2E tests + run: | + make test-e2e-model-registry diff --git a/.gitignore b/.gitignore index a65b023..8e8d683 100644 --- a/.gitignore +++ b/.gitignore @@ -7,7 +7,7 @@ __pycache__ tmp efi-variable-store *.onnx -mr-integration/config/ml-metadata/metadata.sqlite.db +e2e/model-registry/config/ml-metadata/metadata.sqlite.db # Generated python library wheels dist diff --git a/Makefile b/Makefile index a035c9e..5803fa4 100644 --- a/Makefile +++ b/Makefile @@ -27,10 +27,14 @@ test: test-e2e: poetry run pytest --e2e -s -x -rA +.PHONY: test-e2e-model-registry +test-e2e-model-registry: + poetry run pytest --e2e-model-registry -s -x -rA + .PHONY: lint -lint: +lint: install poetry run ruff check --fix .PHONY: mypy -mypy: +mypy: install poetry run mypy . diff --git a/e2e/deploy_model_registry.sh b/e2e/deploy_model_registry.sh new file mode 100755 index 0000000..12ec0d1 --- /dev/null +++ b/e2e/deploy_model_registry.sh @@ -0,0 +1,29 @@ +#!/bin/bash + +SCRIPT_DIR="$(dirname "$(realpath "$BASH_SOURCE")")" +set -e + +if [[ "$OSTYPE" == "darwin"* ]]; then + # Mac OSX + echo "Error: As Kubeflow Model Registry wraps Google's MLMD and no source image for Mchips, this deployment is not supported. Using a temp docker-compose." + rm "$SCRIPT_DIR/model-registry/config/ml-metadata/metadata.sqlite.db" || true + docker compose -f "$SCRIPT_DIR/model-registry/docker-compose.yaml" up + exit 0 +fi + +kubectl create namespace kubeflow +kubectl apply -k "https://github.com/kubeflow/model-registry/manifests/kustomize/overlays/db?ref=v0.2.4-alpha" + +sleep 1 +kubectl get -n kubeflow deployments + +echo "Waiting for Deployment..." +kubectl wait --for=condition=available -n kubeflow deployment/model-registry-deployment --timeout=1m +kubectl logs -n kubeflow deployment/model-registry-deployment +echo "Deployment looks ready." + +echo "Starting port-forward..." +kubectl port-forward svc/model-registry-service -n kubeflow 8081:8080 & +PID=$! +sleep 2 +echo "I have launched port-forward in background with: $PID." diff --git a/e2e/model-registry/README.md b/e2e/model-registry/README.md new file mode 100644 index 0000000..ec94912 --- /dev/null +++ b/e2e/model-registry/README.md @@ -0,0 +1 @@ +This is only for convenience of local development on Mac diff --git a/e2e/model-registry/config/ml-metadata/conn_config.pb b/e2e/model-registry/config/ml-metadata/conn_config.pb new file mode 100644 index 0000000..774b7aa --- /dev/null +++ b/e2e/model-registry/config/ml-metadata/conn_config.pb @@ -0,0 +1,6 @@ +connection_config { + sqlite { + filename_uri: '/tmp/shared/metadata.sqlite.db' + connection_mode: READWRITE_OPENCREATE + } +} diff --git a/e2e/model-registry/docker-compose.yaml b/e2e/model-registry/docker-compose.yaml new file mode 100644 index 0000000..3ab4c6d --- /dev/null +++ b/e2e/model-registry/docker-compose.yaml @@ -0,0 +1,19 @@ +version: '3' +services: + omlmd-integration-mlmd-server: + image: gcr.io/tfx-oss-public/ml_metadata_store_server:1.14.0 + container_name: omlmd-integration-mlmd-server + ports: + - "9090:8080" + environment: + - METADATA_STORE_SERVER_CONFIG_FILE=/tmp/shared/conn_config.pb + volumes: + - ./config/ml-metadata:/tmp/shared + omlmd-integration-model-registry: + image: docker.io/kubeflow/model-registry:latest + command: ["proxy", "--hostname", "0.0.0.0", "--mlmd-hostname", "omlmd-integration-mlmd-server", "--mlmd-port", "8080"] + container_name: omlmd-integration-model-registry + ports: + - "8081:8080" + depends_on: + - omlmd-integration-mlmd-server diff --git a/omlmd/helpers.py b/omlmd/helpers.py index 3a9d5d6..bb86111 100644 --- a/omlmd/helpers.py +++ b/omlmd/helpers.py @@ -1,5 +1,6 @@ from dataclasses import fields from typing import Optional, List +from omlmd.listener import Event, Listener, PushEvent from omlmd.model_metadata import ModelMetadata from omlmd.provider import OMLMDRegistry import os @@ -21,6 +22,9 @@ def download_file(uri): class Helper: + + _listeners: List[Listener] = [] + def __init__(self, registry: Optional[OMLMDRegistry] = None): if registry is None: self._registry = OMLMDRegistry(insecure=True) # TODO: this is a bit limiting when used from CLI, to be refactored @@ -61,12 +65,14 @@ def push( ] try: # print(target, files, model_metadata.to_annotations_dict()) - return self._registry.push( + result = self._registry.push( target=target, files=files, manifest_annotations=model_metadata.to_annotations_dict(), manifest_config="model_metadata.omlmd.json:application/x-config" ) + self.notify_listeners(PushEvent(target, model_metadata)) + return result finally: os.remove("model_metadata.omlmd.json") os.remove("model_metadata.omlmd.yaml") @@ -95,3 +101,14 @@ def crawl( configs = map(self.get_config, targets) joined = "[" + ", ".join(configs) + "]" return joined + + + def add_listener(self, listener: Listener) -> None: + self._listeners.append(listener) + + def remove_listener(self, listener: Listener) -> None: + self._listeners.remove(listener) + + def notify_listeners(self, event: Event) -> None: + for listener in self._listeners: + listener.update(self, event) diff --git a/omlmd/listener.py b/omlmd/listener.py new file mode 100644 index 0000000..b7d4c63 --- /dev/null +++ b/omlmd/listener.py @@ -0,0 +1,27 @@ +from __future__ import annotations +from abc import ABC, abstractmethod +from typing import Any +from omlmd.model_metadata import ModelMetadata + +class Listener(ABC): + """ + TODO: not yet settled for multi-method or current single update method. + """ + @abstractmethod + def update(self, source: Any, event: Event) -> None: + """ + Receive update event. + """ + pass + + +class Event: + pass + + +class PushEvent(Event): + def __init__(self, target: str, metadata: ModelMetadata): + # TODO: cannot just receive yet the push sha, waiting for: https://github.com/oras-project/oras-py/pull/146 in a release. + self.target = target + self.metadata = metadata + diff --git a/omlmd/model_metadata.py b/omlmd/model_metadata.py index 5735e60..3b1a4bd 100644 --- a/omlmd/model_metadata.py +++ b/omlmd/model_metadata.py @@ -1,4 +1,4 @@ -from dataclasses import dataclass, field, asdict +from dataclasses import dataclass, field, asdict, fields from typing import Optional, Dict, Any import json import yaml @@ -45,6 +45,17 @@ def to_yaml(self) -> str: def from_yaml(yaml_str: str) -> 'ModelMetadata': data = yaml.safe_load(yaml_str) return ModelMetadata(**data) + + @staticmethod + def from_dict(data: Dict[str, Any]) -> 'ModelMetadata': + known_keys = {f.name for f in fields(ModelMetadata)} + known_properties = {key: data.get(key) for key in known_keys if key in data} + custom_properties = {key: value for key, value in data.items() if key not in known_keys} + + return ModelMetadata( + **known_properties, + customProperties=custom_properties + ) def deserialize_mdfile(file): diff --git a/omlmd/provider.py b/omlmd/provider.py index 69cd104..4e696d0 100644 --- a/omlmd/provider.py +++ b/omlmd/provider.py @@ -5,8 +5,11 @@ import oras.provider import oras.utils from oras.decorator import ensure_container +from oras.provider import container_type +import oras.schemas import logging import tempfile +from typing import Optional logger = logging.getLogger(__name__) @@ -63,3 +66,26 @@ def get_config(self, package) -> str: os.rmdir(temp_dir) # print("Temporary directory and its contents have been removed.") raise RuntimeError("Unable to locate config layer") + + @ensure_container + def get_manifest_response( + self, + container: container_type, + allowed_media_type: Optional[list] = None, + refresh_headers: bool = True, + ) -> dict: + """ + like get_manifest but return response, + temporary until https://github.com/oras-project/oras-py/pull/146 in a release. + """ + if not allowed_media_type: + allowed_media_type = [oras.defaults.default_manifest_media_type] + headers = {"Accept": ";".join(allowed_media_type)} + + if not refresh_headers: + headers.update(self.headers) + + get_manifest = f"{self.prefix}://{container.manifest_url()}" # type: ignore + response = self.do_request(get_manifest, "GET", headers=headers) + self._check_200_response(response) + return response \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index e92d2cf..628bd92 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -41,6 +41,7 @@ build-backend = "poetry.core.masonry.api" [tool.pytest.ini_options] markers = [ "e2e: end-to-end testing with localhost:5001", + "e2e_model_registry: end-to-end testing with localhost:5001 and Kubeflow Model Registry" ] [tool.ruff] diff --git a/tests/conftest.py b/tests/conftest.py index 8d7e50c..7075826 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -2,22 +2,30 @@ def pytest_collection_modifyitems(config, items): - if config.getoption("--e2e"): - skip_not_e2e = pytest.mark.skip(reason="skipping non-e2e tests; opt-out of --e2e option to run.") - for item in items: - if "e2e" not in item.keywords: - item.add_marker(skip_not_e2e) - return - skip_e2e = pytest.mark.skip(reason="this is an end-to-end test, requires explicit opt-in --e2e option to run.") for item in items: + skip_e2e = pytest.mark.skip(reason="this is an end-to-end test, requires explicit opt-in --e2e option to run.") + skip_e2e_model_registry = pytest.mark.skip(reason="this is an end-to-end test, requires explicit opt-in --e2e-model-registry option to run.") + skip_not_e2e = pytest.mark.skip(reason="skipping non-e2e tests; opt-out of --e2e -like options to run.") if "e2e" in item.keywords: - item.add_marker(skip_e2e) + if not config.getoption("--e2e"): + item.add_marker(skip_e2e) + continue + elif "e2e_model_registry" in item.keywords: + if not config.getoption("--e2e-model-registry"): + item.add_marker(skip_e2e_model_registry) + continue + + if config.getoption("--e2e") or config.getoption("--e2e-model-registry"): + item.add_marker(skip_not_e2e) def pytest_addoption(parser): parser.addoption( "--e2e", action="store_true", default=False, help="opt-in to run tests marked with e2e" ) + parser.addoption( + "--e2e-model-registry", action="store_true", default=False, help="opt-in to run tests marked with e2e_model_registry" + ) @pytest.fixture diff --git a/tests/test_e2e_model_registry.py b/tests/test_e2e_model_registry.py new file mode 100644 index 0000000..79c1b78 --- /dev/null +++ b/tests/test_e2e_model_registry.py @@ -0,0 +1,126 @@ +from omlmd.helpers import Helper +from omlmd.listener import Event, Listener, PushEvent +import pytest +from pathlib import Path +from model_registry import ModelRegistry +from model_registry.types import RegisteredModel +from urllib.parse import quote +from omlmd.helpers import download_file + + +def from_oci_to_kfmr(model_registry: ModelRegistry, push_event: PushEvent, sha: str) -> RegisteredModel: + rm = model_registry.register_model( + name=push_event.metadata.name, + uri=f"oci-artifact://{push_event.target}", + version=quote(sha), + author=push_event.metadata.author, + description=push_event.metadata.description, + model_format_name=push_event.metadata.model_format_name, + model_format_version=push_event.metadata.model_format_version, + metadata=push_event.metadata.customProperties, + ) + return rm + + +@pytest.mark.e2e_model_registry +def test_e2e_model_registry_scenario1(tmp_path, target): + """ + Given a ML model and some metadata, to OCI registry, and then to KF Model Registry (at once) + """ + model_registry = ModelRegistry("http://localhost", 8081, author="mmortari", is_secure=False) + + class ListenerForModelRegistry(Listener): + sha = None + rm = None + + def update(self, source: Helper, event: Event) -> None: + if isinstance(event, PushEvent): + self.sha = source.registry.get_manifest_response(event.target).headers["Docker-Content-Digest"] + print(self.sha) + self.rm = from_oci_to_kfmr(model_registry, event, self.sha) + + listener = ListenerForModelRegistry() + omlmd = Helper() + omlmd.add_listener(listener) + + # assuming a model ... + model_file = Path(__file__).parent / ".." / "README.md" + # ...with some additional characteristics + accuracy_value = 0.987 + + omlmd.push( + target, + model_file, + name="mnist", + description="Lorem ipsum", + author="John Doe", + accuracy=accuracy_value + ) + + v = quote(listener.sha) + + rm = model_registry.get_registered_model("mnist") + assert rm.id == listener.rm.id + assert rm.name == "mnist" + + mv = model_registry.get_model_version("mnist", v) + assert mv.description == "Lorem ipsum" + assert mv.author == "John Doe" + assert mv.custom_properties == {'accuracy': 0.987} + + ma = model_registry.get_model_artifact("mnist", v) + assert ma.uri == f"oci-artifact://{target}" + + # curl http://localhost:5001/v2/testorgns/ml-model-artifact/manifests/v1 -H "Accept: application/vnd.oci.image.manifest.v1+json" --verbose + # or replace tag with target's as needed. + + +@pytest.mark.e2e_model_registry +def test_e2e_model_registry_scenario2(tmp_path, target): + """ + Given some metadata entry in KF model registry, attempt retrieve pointed ML model file asset, then OCI registry + """ + model_registry = ModelRegistry("http://localhost", 8081, author="mmortari", is_secure=False) + + # assuming a model indexed on KF Model Registry ... + registeredmodel_name = "mnist" + version_name = "v0.1" + _ = model_registry.register_model( + registeredmodel_name, + "https://github.com/tarilabs/demo20231212/raw/main/v1.nb20231206162408/mnist.onnx", + model_format_name="onnx", + model_format_version="1", + version=version_name, + description="lorem ipsum mnist", + metadata={ + "accuracy": 3.14, + "license": "apache-2.0", + } + ) + + lookup_name = "mnist" + lookup_version = "v0.1" + + _ = model_registry.get_registered_model(lookup_name) + model_version = model_registry.get_model_version(lookup_name, lookup_version) + model_artifact = model_registry.get_model_artifact(lookup_name, lookup_version) + + file_from_mr = download_file(model_artifact.uri) + + oci_reference = f"localhost:5001/testorgns/{lookup_name}:{lookup_version}" + # ideally, the oci reference should include the registeredmodel name (above), but for e2e testing we're going to use the designated reference + oci_reference = f"localhost:5001/testorgns/ml-model-artifact:{lookup_version}" + + omlmd = Helper() + omlmd.push( + oci_reference, + file_from_mr, + name=lookup_name, + description=model_version.description, + author=model_version.author, + model_format_name=model_artifact.model_format_name, + model_format_version=model_artifact.model_format_version, + **model_version.custom_properties + ) + # curl http://localhost:5001/v2/testorgns/ml-model-artifact/manifests/v0.1 -H "Accept: application/vnd.oci.image.manifest.v1+json" --verbose + # tag v0.1 is defined in this test scenario. diff --git a/tests/test_helpers.py b/tests/test_helpers.py index a88b838..6350a48 100644 --- a/tests/test_helpers.py +++ b/tests/test_helpers.py @@ -1,7 +1,9 @@ from omlmd.helpers import Helper -from omlmd.model_metadata import deserialize_mdfile +from omlmd.listener import Event, Listener +from omlmd.model_metadata import ModelMetadata, deserialize_mdfile import tempfile import json +from omlmd.provider import OMLMDRegistry import pytest from pathlib import Path @@ -31,6 +33,31 @@ def test_call_push_using_md_from_file(mocker): ) +def test_push_event(mocker): + registry = OMLMDRegistry() + mocker.patch.object(registry, "push", return_value=None) + omlmd = Helper(registry) + + events = [] + class MyListener(Listener): + def update(self, _, event: Event) -> None: + events.append(event) + omlmd.add_listener(MyListener()) + + md = { + "name": "mnist", + "description": "Lorem ipsum", + "author": "John Doe", + "accuracy": .987 + } + omlmd.push("unexistent:8080/testorgns/ml-iris:v1", "README.md", **md) + + assert len(events) == 1 + e0 = events[0] + assert e0.target == "unexistent:8080/testorgns/ml-iris:v1" + assert e0.metadata == ModelMetadata.from_dict(md) + + @pytest.mark.e2e def test_e2e_push_pull(tmp_path, target): omlmd = Helper() diff --git a/tests/test_omlmd.py b/tests/test_omlmd.py index ffe36d9..03f6235 100644 --- a/tests/test_omlmd.py +++ b/tests/test_omlmd.py @@ -42,3 +42,21 @@ def test_deserialize_file_yaml(): f.flush() metadata_from_yaml = deserialize_mdfile(f.name) assert md_dict == metadata_from_yaml + + +def test_from_dict(): + data = { + "name": "mnist", + "description": "Lorem ipsum", + "author": "John Doe", + "accuracy": .987 + } + md = ModelMetadata( + name="mnist", + description="Lorem ipsum", + author="John Doe", + customProperties={ + "accuracy": .987 + } + ) + assert ModelMetadata.from_dict(data) == md