Add save_to_oci_registry ptyhon client method

.github/workflows/python-tests.yml

@@ -165,6 +165,11 @@ jobs:
           pipx install --pip-args=--constraint=${{ github.workspace }}/.github/workflows/constraints.txt nox
           pipx inject --pip-args=--constraint=${{ github.workspace }}/.github/workflows/constraints.txt nox nox-poetry
           nox --version
+      - name: Install Storage Clients
+        # Oras should be available in the GH Action CI by default
+        # See: https://oras.land/docs/installation/#runner-machine-of-azure-devops-and-github-actions
+        run: |
+          sudo apt-get -y install skopeo
       - name: Nox test
         working-directory: clients/python
         run: |
@@ -201,6 +206,8 @@ jobs:
         env:
           IMG: "${{ env.IMG_ORG }}/${{ env.IMG_REPO }}:${{ steps.tags.outputs.tag }}"
         run: ./scripts/deploy_on_kind.sh
+      - name: Deploy OCI Registry using manifests
+        run: ./scripts/deploy_local_kind_registry.sh
       - name: Nox test end-to-end
         working-directory: clients/python
         run: |

clients/python/Makefile

@@ -18,8 +18,12 @@ deploy-latest-mr:
 	cd ../../ && IMG_VERSION=${IMG_VERSION} make image/build && LOCAL=1 ./scripts/deploy_on_kind.sh
 	kubectl port-forward -n kubeflow services/model-registry-service 8080:8080 &
 
+.PHONY: deploy-local-registry
+deploy-local-registry:
+	cd ../../ && ./scripts/deploy_local_kind_registry.sh
+
 .PHONY: test-e2e
-test-e2e: deploy-latest-mr
+test-e2e: deploy-latest-mr deploy-local-registry
 	poetry run pytest --e2e -s -rA
 
 .PHONY: test

clients/python/README.md

@@ -200,4 +200,12 @@ You can use `make test` to execute `pytest`.
 
 Check out our [recommendations on setting up your docker engine](https://github.com/kubeflow/model-registry/blob/main/CONTRIBUTING.md#docker-engine) on an ARM processor.
 
+### Extras
+
+Depending on your development flow, you need to install extra dependencies:
+
+```
+poetry install -E "olot"
+```
+
 <!-- github-only -->

clients/python/noxfile.py

@@ -58,6 +58,7 @@ def tests(session: Session) -> None:
         "requests",
         "pytest",
         "pytest-asyncio",
+        "olot",
     )
     session.run(
         "pytest",
@@ -77,6 +78,7 @@ def e2e_tests(session: Session) -> None:
         "coverage[toml]",
         "pytest-cov",
         "huggingface-hub",
+        "olot",
     )
     try:
         session.run(

clients/python/pyproject.toml

@@ -27,9 +27,11 @@ nest-asyncio = "^1.6.0"
 eval-type-backport = "^0.2.0"
 
 huggingface-hub = { version = ">=0.20.1,<0.29.0", optional = true }
+olot = { version = "^0.1.2", optional = true }
 
 [tool.poetry.extras]
 hf = ["huggingface-hub"]
+olot = ["olot"]
 
 [tool.poetry.group.docs]
 optional = true

clients/python/src/model_registry/utils.py

@@ -3,11 +3,14 @@
 from __future__ import annotations
 
 import os
+import tempfile
+from pathlib import Path
+from typing import Callable, TypedDict
 
 from typing_extensions import overload
 
 from ._utils import required_args
-from .exceptions import MissingMetadata
+from .exceptions import MissingMetadata, StoreError
 
 
 @overload
@@ -90,3 +93,116 @@ def s3_uri_from(
     # https://alexwlchan.net/2020/s3-keys-are-not-file-paths/ nor do they resolve to valid URls
     # FIXME: is this safe?
     return f"s3://{bucket}/{path}?endpoint={endpoint}&defaultRegion={region}"
+
+
+class BackendDefinition(TypedDict):
+    """Holds the 3 core callables for a backend.
+
+    - is_available() -> bool
+    - pull(base_image: str, dest_dir: Path) -> None
+    - push(local_image_path: Path, oci_ref: str) -> None.
+    """
+
+    available: Callable[[], bool]
+    pull: Callable[[str, Path], None]
+    push: Callable[[Path, str], None]
+
+
+def _get_skopeo_backend() -> BackendDefinition:
+    try:
+        from olot.backend.skopeo import is_skopeo, skopeo_pull, skopeo_push
+    except ImportError as e:
+        msg = "Could not import 'olot.backend.skopeo'. Ensure that 'olot' is installed if you want to use the 'skopeo' backend."
+        raise ImportError(msg) from e
+
+    return {
+        "is_available": is_skopeo,
+        "pull": skopeo_pull,
+        "push": skopeo_push
+    }
+
+def _get_oras_backend() -> BackendDefinition:
+    try:
+        from olot.backend.oras_cp import is_oras, oras_pull, oras_push
+    except ImportError as e:
+        msg = "Could not import 'olot.backend.oras_cp'. Ensure that 'olot' is installed if you want to use the 'oras_cp' backend."
+        raise ImportError(msg) from e
+
+    return {
+        "is_available": is_oras,
+        "pull": oras_pull,
+        "push": oras_push,
+    }
+
+# A dict mapping backend names to their definitions
+BackendDict = dict[str, Callable[[], BackendDefinition]]
+
+DEFAULT_BACKENDS: BackendDict = {
+    "skopeo": _get_skopeo_backend,
+    "oras": _get_oras_backend,
+}
+
+def save_to_oci_registry(
+        base_image: str,
+        oci_ref: str,
+        model_files: list[os.PathLike],
+        dest_dir: str | os.PathLike = None,
+        backend: str = "skopeo",
+        modelcard: os.PathLike | None = None,
+        backend_registry: BackendDict | None = DEFAULT_BACKENDS,
+) -> str:
+    """Appends a list of files to an OCI-based image.
+
+    Args:
+        base_image: The image to append model files to. This image will be downloaded to the location at `dest_dir`
+        dest_dir: The location to save the downloaded and extracted base image to.
+        oci_ref: Destination of where to push the newly layered image to. eg, "quay.io/my-org/my-registry:1.0.0"
+        model_files: List of files to add to the base_image as layers
+        backend: The CLI tool to use to perform the oci image pull/push. One of: "skopeo", "oras"
+        modelcard: Optional, path to the modelcard to additionally include as a layer
+        backend_registry: Optional, a dict of backends available to be used to perform the OCI image download/upload
+    Raises:
+        ValueError: If the chosen backend is not installed on the host
+        ValueError: If the chosen backend is an invalid option
+        StoreError: If `olot` is not installed as a python package
+    Returns:
+        None.
+    """
+    try:
+        from olot.basics import oci_layers_on_top
+    except ImportError as e:
+        msg = """Package `olot` is not installed.
+To save models to OCI compatible storage, start by installing the `olot` package, either directly or as an
+extra (available as `model-registry[olot]`), e.g.:
+```sh
+!pip install --pre model-registry[olot]
+```
+or
+```sh
+!pip install olot
+```
+        """
+        raise StoreError(msg) from e
+
+
+    if backend not in backend_registry:
+        msg = f"'{backend}' is not an available backend to use. Available backends: {backend_registry.keys()}"
+        raise ValueError(msg)
+
+    # Fetching the backend definition can throw an error, but it should bubble up as it has the appropriate messaging
+    backend_def = backend_registry[backend]()
+
+    if not backend_def["is_available"]():
+        msg = f"Backend '{backend}' is selected, but not available on the system. Ensure the dependencies for '{backend}' are installed in your environment."
+        raise ValueError(msg)
+
+    if dest_dir is None:
+        dest_dir = tempfile.mkdtemp()
+    local_image_path = Path(dest_dir)
+    backend_def["pull"](base_image, local_image_path)
+    oci_layers_on_top(local_image_path, model_files, modelcard)
+    backend_def["push"](local_image_path, oci_ref)
+
+    # Return the OCI URI
+
+    return f"oci://{oci_ref}"

clients/python/tests/.gitignore

@@ -0,0 +1 @@
+data/

clients/python/tests/test_utils.py

@@ -1,9 +1,11 @@
 import os
+import pathlib
+from unittest.mock import Mock
 
 import pytest
 
 from model_registry.exceptions import MissingMetadata
-from model_registry.utils import s3_uri_from
+from model_registry.utils import s3_uri_from, save_to_oci_registry
 
 
 def test_s3_uri_builder():
@@ -71,3 +73,102 @@ def test_s3_uri_builder_with_complete_env():
     os.environ["AWS_S3_ENDPOINT"] = "test-endpoint"
     os.environ["AWS_DEFAULT_REGION"] = "test-region"
     assert s3_uri_from("test-path") == s3_uri_from("test-path", "test-bucket")
+
+@pytest.mark.e2e(type="oci")
+def test_save_to_oci_registry_with_skopeo():
+    # TODO: We need a good source registry which is oci-compliant and very small in size
+    base_image = "quay.io/mmortari/hello-world-wait:latest"
+    dest_dir = "tests/data"
+    oci_ref = "localhost:5001/foo/bar:latest"
+
+    # Create a sample file named README.md to be added to the registry
+    pathlib.Path(dest_dir).mkdir(parents=True, exist_ok=True)
+    readme_file_path = os.path.join(dest_dir, "README.md")
+    with open(readme_file_path, "w") as f:
+        f.write("")
+
+    model_files = [readme_file_path]
+    backend = "skopeo"
+
+    save_to_oci_registry(base_image, oci_ref, model_files, dest_dir, backend)
+
+
+# These are trimmed down versions of whats found in the example specs found here: https://github.com/opencontainers/image-spec/blob/main/image-layout.md#oci-layout-file
+index_json_contents = """{
+  "schemaVersion": 2,
+  "mediaType": "application/vnd.oci.image.index.v1+json",
+  "manifests": [],
+  "annotations": {
+    "com.example.index.revision": "r124356"
+  }
+}"""
+oci_layout_contents = """{"imageLayoutVersion": "1.0.0"}"""
+
+def test_save_to_oci_registry_with_custom_backend():
+    is_available_mock = Mock()
+    is_available_mock.return_value = True
+    pull_mock = Mock()
+    push_mock = Mock()
+    def pull_mock_imple(base_image, dest_dir):
+        pathlib.Path(dest_dir).joinpath("oci-layout").write_text(oci_layout_contents)
+        pathlib.Path(dest_dir).joinpath("index.json").write_text(index_json_contents)
+
+    pull_mock.side_effect = pull_mock_imple
+
+
+    backend = "something_custom"
+    backend_registry = {
+        "something_custom": lambda: {
+            "is_available": is_available_mock,
+            "pull": pull_mock,
+            "push": push_mock,
+        }
+    }
+
+    # similar to other test
+    base_image = "quay.io/mmortari/hello-world-wait:latest"
+    dest_dir = "tests/data"
+    oci_ref = "localhost:5001/foo/bar:latest"
+
+    # Create a sample file named README.md to be added to the registry
+    pathlib.Path(dest_dir).mkdir(parents=True, exist_ok=True)
+    readme_file_path = os.path.join(dest_dir, "README.md")
+    with open(readme_file_path, "w") as f:
+        f.write("")
+
+    model_files = [readme_file_path]
+
+    uri = save_to_oci_registry(base_image, oci_ref, model_files, dest_dir, backend, None, backend_registry)
+    # Ensure our mocked backend was called
+    is_available_mock.assert_called_once()
+    pull_mock.assert_called_once()
+    push_mock.assert_called_once()
+    assert uri == f"oci://{oci_ref}"
+
+def test_save_to_oci_registry_with_custom_backend_unavailable():
+    is_available_mock = Mock()
+    is_available_mock.return_value = False # Backend is unavailable, expect an error
+    pull_mock = Mock()
+    push_mock = Mock()
+
+    backend = "something_custom"
+    backend_registry = {
+        "something_custom": lambda: {
+            "is_available": is_available_mock,
+            "pull": pull_mock,
+            "push": push_mock,
+        }
+    }
+
+
+    with pytest.raises(ValueError, match=f"Backend '{backend}' is selected, but not available on the system. Ensure the dependencies for '{backend}' are installed in your environment.") as e:
+        save_to_oci_registry("", "", [], "", backend, backend_registry=backend_registry)
+
+    assert f"Backend '{backend}' is selected, but not available on the system." in str(e.value)
+
+def test_save_to_oci_registry_backend_not_found():
+    backend = "non-existent"
+    with pytest.raises(ValueError, match=f"'{backend}' is not an available backend to use.") as e:
+        save_to_oci_registry("", "", [], "", backend)
+
+    assert f"'{backend}' is not an available backend to use." in str(e.value)

scripts/deploy_local_kind_registry.sh

@@ -0,0 +1,17 @@
+#!/bin/bash
+
+SCRIPT_DIR="$(dirname "$(realpath "$BASH_SOURCE")")"
+set -e
+
+kubectl apply -f "${SCRIPT_DIR}/services/container_registry.yaml"
+
+echo "Waiting for Deployment..."
+kubectl wait --for=condition=available deployment/distribution-registry-test-deployment --timeout=5m
+kubectl logs deployment/distribution-registry-test-deployment
+echo "Deployment looks ready."
+
+echo "Starting port-forward..."
+kubectl port-forward service/distribution-registry-test-service 5001:5001 &
+PID=$!
+sleep 2
+echo "I have launched port-forward in background with: $PID."

scripts/deploy_on_kind.sh

@@ -10,7 +10,7 @@ source ./${DIR}/utils.sh
 
 # modularity to allow re-use this script against a remote k8s cluster
 if [[ -n "$LOCAL" ]]; then
-    CLUSTER_NAME="${CLUSTER_NAME:-kind}"
+    CLUSTER_NAME="${CLUSTER_NAME:-mr-e2e}"
 
     echo 'Creating local Kind cluster and loading image'