py: add listener mechanism (#8)

* py: add listener mechanism

Signed-off-by: tarilabs <[email protected]>

* add e2e test scenarios

Signed-off-by: tarilabs <[email protected]>

* wire GHA jobs

Signed-off-by: tarilabs <[email protected]>

* linting and minor changes

Signed-off-by: tarilabs <[email protected]>

---------

Signed-off-by: tarilabs <[email protected]>

.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

.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

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 .

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."

e2e/model-registry/README.md

@@ -0,0 +1 @@
+This is only for convenience of local development on Mac

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
+  }
+}

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

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)

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
+

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):

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

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]

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