diff --git a/e2e/quay-lite/config.yaml b/e2e/quay-lite/config.yaml index 72fd8cc..de87311 100644 --- a/e2e/quay-lite/config.yaml +++ b/e2e/quay-lite/config.yaml @@ -4,8 +4,8 @@ DATA_MODEL_CACHE_CONFIG: primary: host: quay-redis SUPER_USERS: -- admin -- user1 + - admin + - user1 AUTHENTICATION_TYPE: Database DB_URI: postgresql://quay:quay@quay-postgresql:5432/quay BUILDLOGS_REDIS: @@ -15,15 +15,15 @@ USER_EVENTS_REDIS: host: quay-redis port: 6379 BITTORRENT_FILENAME_PEPPER: 0ee18f90-5b6d-42d2-ab5e-ec9fcd846272 -DATABASE_SECRET_KEY: '30060361640793187613697366923211113205676925445650250274752125083971638376224' +DATABASE_SECRET_KEY: "30060361640793187613697366923211113205676925445650250274752125083971638376224" DEFAULT_TAG_EXPIRATION: 2w DISTRIBUTED_STORAGE_CONFIG: default: - - LocalStorage - - storage_path: /datastorage/registry + - LocalStorage + - storage_path: /datastorage/registry DISTRIBUTED_STORAGE_DEFAULT_LOCATIONS: [] DISTRIBUTED_STORAGE_PREFERENCE: -- default + - default ENTERPRISE_LOGO_URL: /static/img/quay-horizontal-color.svg EXTERNAL_TLS_TERMINATION: true FEATURE_ANONYMOUS_ACCESS: true @@ -59,13 +59,13 @@ REPO_MIRROR_TLS_VERIFY: true SETUP_COMPLETE: true SIGNING_ENGINE: gpg2 TAG_EXPIRATION_OPTIONS: -- 0s -- 1d -- 1w -- 2w -- 4w + - 0s + - 1d + - 1w + - 2w + - 4w TEAM_RESYNC_STALE_TIME: 60m -TESTING: false +TESTING: true USERFILES_LOCATION: default USERFILES_PATH: userfiles/ USE_CDN: false @@ -77,4 +77,4 @@ CORS_ORIGIN: - "http://localhost:9000" FEATURE_UI_V2: True FEATURE_USER_METADATA: True -IGNORE_UNKNOWN_MEDIATYPES: True +IGNORE_UNKNOWN_MEDIATYPES: false diff --git a/e2e/quay-lite/quay-app-config.yaml b/e2e/quay-lite/quay-app-config.yaml index 293a324..e91bb5c 100644 --- a/e2e/quay-lite/quay-app-config.yaml +++ b/e2e/quay-lite/quay-app-config.yaml @@ -4,4 +4,4 @@ metadata: name: quay-app-config type: Opaque data: - config.yaml: REFUQV9NT0RFTF9DQUNIRV9DT05GSUc6CiAgZW5naW5lOiByZWRpcwogIHJlZGlzX2NvbmZpZzoKICAgIHByaW1hcnk6CiAgICAgIGhvc3Q6IHF1YXktcmVkaXMKU1VQRVJfVVNFUlM6Ci0gYWRtaW4KLSB1c2VyMQpBVVRIRU5USUNBVElPTl9UWVBFOiBEYXRhYmFzZQpEQl9VUkk6IHBvc3RncmVzcWw6Ly9xdWF5OnF1YXlAcXVheS1wb3N0Z3Jlc3FsOjU0MzIvcXVheQpCVUlMRExPR1NfUkVESVM6CiAgaG9zdDogcXVheS1yZWRpcwogIHBvcnQ6IDYzNzkKVVNFUl9FVkVOVFNfUkVESVM6CiAgaG9zdDogcXVheS1yZWRpcwogIHBvcnQ6IDYzNzkKQklUVE9SUkVOVF9GSUxFTkFNRV9QRVBQRVI6IDBlZTE4ZjkwLTViNmQtNDJkMi1hYjVlLWVjOWZjZDg0NjI3MgpEQVRBQkFTRV9TRUNSRVRfS0VZOiAnMzAwNjAzNjE2NDA3OTMxODc2MTM2OTczNjY5MjMyMTExMTMyMDU2NzY5MjU0NDU2NTAyNTAyNzQ3NTIxMjUwODM5NzE2MzgzNzYyMjQnCkRFRkFVTFRfVEFHX0VYUElSQVRJT046IDJ3CkRJU1RSSUJVVEVEX1NUT1JBR0VfQ09ORklHOgogIGRlZmF1bHQ6CiAgLSBMb2NhbFN0b3JhZ2UKICAtIHN0b3JhZ2VfcGF0aDogL2RhdGFzdG9yYWdlL3JlZ2lzdHJ5CkRJU1RSSUJVVEVEX1NUT1JBR0VfREVGQVVMVF9MT0NBVElPTlM6IFtdCkRJU1RSSUJVVEVEX1NUT1JBR0VfUFJFRkVSRU5DRToKLSBkZWZhdWx0CkVOVEVSUFJJU0VfTE9HT19VUkw6IC9zdGF0aWMvaW1nL3F1YXktaG9yaXpvbnRhbC1jb2xvci5zdmcKRVhURVJOQUxfVExTX1RFUk1JTkFUSU9OOiB0cnVlCkZFQVRVUkVfQU5PTllNT1VTX0FDQ0VTUzogdHJ1ZQpGRUFUVVJFX0FQUF9SRUdJU1RSWTogZmFsc2UKRkVBVFVSRV9BUFBfU1BFQ0lGSUNfVE9LRU5TOiB0cnVlCkZFQVRVUkVfQlVJTERfU1VQUE9SVDogZmFsc2UKRkVBVFVSRV9DSEFOR0VfVEFHX0VYUElSQVRJT046IHRydWUKRkVBVFVSRV9ESVJFQ1RfTE9HSU46IHRydWUKRkVBVFVSRV9NQUlMSU5HOiBmYWxzZQpGRUFUVVJFX1BBUlRJQUxfVVNFUl9BVVRPQ09NUExFVEU6IHRydWUKRkVBVFVSRV9SRVBPX01JUlJPUjogZmFsc2UKRkVBVFVSRV9SRVFVSVJFX1RFQU1fSU5WSVRFOiB0cnVlCkZFQVRVUkVfUkVTVFJJQ1RFRF9WMV9QVVNIOiBmYWxzZQpGRUFUVVJFX1NFQ1VSSVRZX05PVElGSUNBVElPTlM6IGZhbHNlCkZFQVRVUkVfU0VDVVJJVFlfU0NBTk5FUjogZmFsc2UKRkVBVFVSRV9VU0VSTkFNRV9DT05GSVJNQVRJT046IHRydWUKRkVBVFVSRV9VU0VSX0lOSVRJQUxJWkU6IHRydWUKRkVBVFVSRV9VU0VSX0NSRUFUSU9OOiB0cnVlCkZFQVRVUkVfVVNFUl9MT0dfQUNDRVNTOiB0cnVlCkZFQVRVUkVfUFJPWFlfQ0FDSEU6IHRydWUKR0lUSFVCX0xPR0lOX0NPTkZJRzoge30KR0lUSFVCX1RSSUdHRVJfQ09ORklHOiB7fQpHSVRMQUJfVFJJR0dFUl9LSU5EOiB7fQpMT0dfQVJDSElWRV9MT0NBVElPTjogZGVmYXVsdApNQUlMX0RFRkFVTFRfU0VOREVSOiBhZG1pbkBleGFtcGxlLmNvbQpNQUlMX1BPUlQ6IDU4NwpNQUlMX1VTRV9UTFM6IHRydWUKUFJFRkVSUkVEX1VSTF9TQ0hFTUU6IGh0dHAKUkVHSVNUUllfVElUTEU6IFJlZCBIYXQgUXVheSBMSVRFClJFR0lTVFJZX1RJVExFX1NIT1JUOiBSZWQgSGF0IFF1YXkgTElURQpSRVBPX01JUlJPUl9TRVJWRVJfSE9TVE5BTUU6IG51bGwKUkVQT19NSVJST1JfVExTX1ZFUklGWTogdHJ1ZQpTRVRVUF9DT01QTEVURTogdHJ1ZQpTSUdOSU5HX0VOR0lORTogZ3BnMgpUQUdfRVhQSVJBVElPTl9PUFRJT05TOgotIDBzCi0gMWQKLSAxdwotIDJ3Ci0gNHcKVEVBTV9SRVNZTkNfU1RBTEVfVElNRTogNjBtClRFU1RJTkc6IGZhbHNlClVTRVJGSUxFU19MT0NBVElPTjogZGVmYXVsdApVU0VSRklMRVNfUEFUSDogdXNlcmZpbGVzLwpVU0VfQ0ROOiBmYWxzZQpGRUFUVVJFX1FVT1RBX01BTkFHRU1FTlQ6IFRydWUKU0VSVkVSX0hPU1ROQU1FOiBsb2NhbGhvc3Q6NTAwMQpCUk9XU0VSX0FQSV9DQUxMU19YSFJfT05MWTogRmFsc2UKQ09SU19PUklHSU46CiAgLSAiaHR0cHM6Ly9zdGFnZS5mb28ucmVkaGF0LmNvbToxMzM3IgogIC0gImh0dHA6Ly9sb2NhbGhvc3Q6OTAwMCIKRkVBVFVSRV9VSV9WMjogVHJ1ZQpGRUFUVVJFX1VTRVJfTUVUQURBVEE6IFRydWUKSUdOT1JFX1VOS05PV05fTUVESUFUWVBFUzogVHJ1ZQo= + config.yaml: REFUQV9NT0RFTF9DQUNIRV9DT05GSUc6CiAgZW5naW5lOiByZWRpcwogIHJlZGlzX2NvbmZpZzoKICAgIHByaW1hcnk6CiAgICAgIGhvc3Q6IHF1YXktcmVkaXMKU1VQRVJfVVNFUlM6CiAgLSBhZG1pbgogIC0gdXNlcjEKQVVUSEVOVElDQVRJT05fVFlQRTogRGF0YWJhc2UKREJfVVJJOiBwb3N0Z3Jlc3FsOi8vcXVheTpxdWF5QHF1YXktcG9zdGdyZXNxbDo1NDMyL3F1YXkKQlVJTERMT0dTX1JFRElTOgogIGhvc3Q6IHF1YXktcmVkaXMKICBwb3J0OiA2Mzc5ClVTRVJfRVZFTlRTX1JFRElTOgogIGhvc3Q6IHF1YXktcmVkaXMKICBwb3J0OiA2Mzc5CkJJVFRPUlJFTlRfRklMRU5BTUVfUEVQUEVSOiAwZWUxOGY5MC01YjZkLTQyZDItYWI1ZS1lYzlmY2Q4NDYyNzIKREFUQUJBU0VfU0VDUkVUX0tFWTogIjMwMDYwMzYxNjQwNzkzMTg3NjEzNjk3MzY2OTIzMjExMTEzMjA1Njc2OTI1NDQ1NjUwMjUwMjc0NzUyMTI1MDgzOTcxNjM4Mzc2MjI0IgpERUZBVUxUX1RBR19FWFBJUkFUSU9OOiAydwpESVNUUklCVVRFRF9TVE9SQUdFX0NPTkZJRzoKICBkZWZhdWx0OgogICAgLSBMb2NhbFN0b3JhZ2UKICAgIC0gc3RvcmFnZV9wYXRoOiAvZGF0YXN0b3JhZ2UvcmVnaXN0cnkKRElTVFJJQlVURURfU1RPUkFHRV9ERUZBVUxUX0xPQ0FUSU9OUzogW10KRElTVFJJQlVURURfU1RPUkFHRV9QUkVGRVJFTkNFOgogIC0gZGVmYXVsdApFTlRFUlBSSVNFX0xPR09fVVJMOiAvc3RhdGljL2ltZy9xdWF5LWhvcml6b250YWwtY29sb3Iuc3ZnCkVYVEVSTkFMX1RMU19URVJNSU5BVElPTjogdHJ1ZQpGRUFUVVJFX0FOT05ZTU9VU19BQ0NFU1M6IHRydWUKRkVBVFVSRV9BUFBfUkVHSVNUUlk6IGZhbHNlCkZFQVRVUkVfQVBQX1NQRUNJRklDX1RPS0VOUzogdHJ1ZQpGRUFUVVJFX0JVSUxEX1NVUFBPUlQ6IGZhbHNlCkZFQVRVUkVfQ0hBTkdFX1RBR19FWFBJUkFUSU9OOiB0cnVlCkZFQVRVUkVfRElSRUNUX0xPR0lOOiB0cnVlCkZFQVRVUkVfTUFJTElORzogZmFsc2UKRkVBVFVSRV9QQVJUSUFMX1VTRVJfQVVUT0NPTVBMRVRFOiB0cnVlCkZFQVRVUkVfUkVQT19NSVJST1I6IGZhbHNlCkZFQVRVUkVfUkVRVUlSRV9URUFNX0lOVklURTogdHJ1ZQpGRUFUVVJFX1JFU1RSSUNURURfVjFfUFVTSDogZmFsc2UKRkVBVFVSRV9TRUNVUklUWV9OT1RJRklDQVRJT05TOiBmYWxzZQpGRUFUVVJFX1NFQ1VSSVRZX1NDQU5ORVI6IGZhbHNlCkZFQVRVUkVfVVNFUk5BTUVfQ09ORklSTUFUSU9OOiB0cnVlCkZFQVRVUkVfVVNFUl9JTklUSUFMSVpFOiB0cnVlCkZFQVRVUkVfVVNFUl9DUkVBVElPTjogdHJ1ZQpGRUFUVVJFX1VTRVJfTE9HX0FDQ0VTUzogdHJ1ZQpGRUFUVVJFX1BST1hZX0NBQ0hFOiB0cnVlCkdJVEhVQl9MT0dJTl9DT05GSUc6IHt9CkdJVEhVQl9UUklHR0VSX0NPTkZJRzoge30KR0lUTEFCX1RSSUdHRVJfS0lORDoge30KTE9HX0FSQ0hJVkVfTE9DQVRJT046IGRlZmF1bHQKTUFJTF9ERUZBVUxUX1NFTkRFUjogYWRtaW5AZXhhbXBsZS5jb20KTUFJTF9QT1JUOiA1ODcKTUFJTF9VU0VfVExTOiB0cnVlClBSRUZFUlJFRF9VUkxfU0NIRU1FOiBodHRwClJFR0lTVFJZX1RJVExFOiBSZWQgSGF0IFF1YXkgTElURQpSRUdJU1RSWV9USVRMRV9TSE9SVDogUmVkIEhhdCBRdWF5IExJVEUKUkVQT19NSVJST1JfU0VSVkVSX0hPU1ROQU1FOiBudWxsClJFUE9fTUlSUk9SX1RMU19WRVJJRlk6IHRydWUKU0VUVVBfQ09NUExFVEU6IHRydWUKU0lHTklOR19FTkdJTkU6IGdwZzIKVEFHX0VYUElSQVRJT05fT1BUSU9OUzoKICAtIDBzCiAgLSAxZAogIC0gMXcKICAtIDJ3CiAgLSA0dwpURUFNX1JFU1lOQ19TVEFMRV9USU1FOiA2MG0KVEVTVElORzogdHJ1ZQpVU0VSRklMRVNfTE9DQVRJT046IGRlZmF1bHQKVVNFUkZJTEVTX1BBVEg6IHVzZXJmaWxlcy8KVVNFX0NETjogZmFsc2UKRkVBVFVSRV9RVU9UQV9NQU5BR0VNRU5UOiBUcnVlClNFUlZFUl9IT1NUTkFNRTogbG9jYWxob3N0OjUwMDEKQlJPV1NFUl9BUElfQ0FMTFNfWEhSX09OTFk6IEZhbHNlCkNPUlNfT1JJR0lOOgogIC0gImh0dHBzOi8vc3RhZ2UuZm9vLnJlZGhhdC5jb206MTMzNyIKICAtICJodHRwOi8vbG9jYWxob3N0OjkwMDAiCkZFQVRVUkVfVUlfVjI6IFRydWUKRkVBVFVSRV9VU0VSX01FVEFEQVRBOiBUcnVlCklHTk9SRV9VTktOT1dOX01FRElBVFlQRVM6IGZhbHNlCg== diff --git a/e2e/test_cli.sh b/e2e/test_cli.sh index 919d3bd..b9673f0 100755 --- a/e2e/test_cli.sh +++ b/e2e/test_cli.sh @@ -17,15 +17,15 @@ omlmd push localhost:5001/mmortari/mlartifact:v1 README.md --metadata tests/data DIR="tmp/a" omlmd pull localhost:5001/mmortari/mlartifact:v1 -o "$DIR" --plain-http file_count=$(find "$DIR" -type f | wc -l) -if [ "$file_count" -eq 3 ]; then - echo "Expected 3 files in $DIR, ok." +if [ "$file_count" -eq 2 ]; then + echo "Expected 2 files in $DIR, ok." else - echo "Expected 3 files in $DIR, got $file_count, FAIL." + echo "Expected 2 files in $DIR, got $file_count, FAIL." exit 1 fi DIR="tmp/b" -omlmd pull localhost:5001/mmortari/mlartifact:v1 -o "$DIR" --media-types "application/x-mlmodel" --plain-http +omlmd pull localhost:5001/mmortari/mlartifact:v1 -o "$DIR" --media-types "application/vnd.oci.image.layer.v1.tar" --plain-http file_count=$(find "$DIR" -type f | wc -l) if [ "$file_count" -eq 1 ]; then echo "Expected 1 file in $DIR, ok." diff --git a/omlmd/cli.py b/omlmd/cli.py index 8d66eca..20a7400 100644 --- a/omlmd/cli.py +++ b/omlmd/cli.py @@ -73,6 +73,11 @@ def crawl(plain_http: bool, targets: tuple[str]): required=True, type=click.Path(path_type=Path, exists=True, resolve_path=True), ) +@click.option( + "--as-artifact", + is_flag=True, + help="Push as an artifact (default is as a blob)", +) @cloup.option_group( "Metadata options", cloup.option( @@ -88,6 +93,7 @@ def push( plain_http: bool, target: str, path: Path, + as_artifact: bool, metadata: Path | None, empty_metadata: bool, ): @@ -96,4 +102,6 @@ def push( if empty_metadata: logger.warning(f"Pushing to {target} with empty metadata.") md = deserialize_mdfile(metadata) if metadata else {} - click.echo(Helper.from_default_registry(plain_http).push(target, path, **md)) + click.echo( + Helper.from_default_registry(plain_http).push(target, path, as_artifact, **md) + ) diff --git a/omlmd/constants.py b/omlmd/constants.py index 05db48e..d20dad0 100644 --- a/omlmd/constants.py +++ b/omlmd/constants.py @@ -1,3 +1,7 @@ +from oras.defaults import default_blob_media_type + FILENAME_METADATA_JSON = "model_metadata.omlmd.json" -MIME_APPLICATION_CONFIG = "application/x-config" MIME_APPLICATION_MLMODEL = "application/x-mlmodel" +MIME_APPLICATION_MLMETADATA = "application/x-mlmetadata+json" +MIME_BLOB = default_blob_media_type +MIME_MANIFEST_CONFIG = "application/vnd.oci.image.config.v1+json" diff --git a/omlmd/helpers.py b/omlmd/helpers.py index d3b2bd2..4276942 100644 --- a/omlmd/helpers.py +++ b/omlmd/helpers.py @@ -1,7 +1,10 @@ from __future__ import annotations +import json import logging import os +import platform +import tarfile import urllib.request from collections.abc import Sequence from dataclasses import dataclass, field @@ -10,8 +13,10 @@ from .constants import ( FILENAME_METADATA_JSON, - MIME_APPLICATION_CONFIG, + MIME_APPLICATION_MLMETADATA, MIME_APPLICATION_MLMODEL, + MIME_BLOB, + MIME_MANIFEST_CONFIG, ) from .listener import Event, Listener, PushEvent from .model_metadata import ModelMetadata @@ -20,6 +25,18 @@ logger = logging.getLogger(__name__) +def get_arch() -> str: + mac = platform.machine() + if mac == "x86_64": + return "amd64" + if mac == "arm64": + return "arm64" + if mac == "aarch64": + return "arm64" + msg = f"Unsupported architecture: {platform.machine()}" + raise NotImplementedError(msg) + + def download_file(uri: str): file_name = os.path.basename(uri) urllib.request.urlretrieve(uri, file_name) @@ -41,6 +58,7 @@ def push( self, target: str, path: Path | str, + as_artifact: bool = False, **kwargs, ): owns_meta = True @@ -48,30 +66,69 @@ def push( path = Path(path) meta_path = path.parent / FILENAME_METADATA_JSON - if not kwargs and meta_path.exists(): + if meta_path.exists(): owns_meta = False logger.warning("Reusing intermediate metadata files.") logger.debug(f"{meta_path}") - with open(meta_path, "r") as f: - model_metadata = ModelMetadata.from_json(f.read()) - elif meta_path.exists(): - err = dedent(f""" -OMLMD intermediate metadata files found at '{meta_path}'. -Cannot resolve with conflicting keyword args: {kwargs}. -You can reuse the existing metadata by omitting any keywords. -If that was NOT intended, please REMOVE that file from your environment before re-running. - -Note for advanced users: if merging keys with existing metadata is desired, you should create a Feature Request upstream: https://github.com/containers/omlmd""") - raise RuntimeError(err) + model_metadata = ModelMetadata(**json.loads(meta_path.read_bytes())) + if kwargs and ModelMetadata.from_dict(kwargs) != model_metadata: + err = dedent(f""" + OMLMD intermediate metadata files found at '{meta_path}'. + Cannot resolve with conflicting keyword args: {kwargs}. + You can reuse the existing metadata by omitting any keywords. + If that was NOT intended, please REMOVE that file from your environment before re-running. + + Note for advanced users: if merging keys with existing metadata is desired, you should create a Feature Request upstream: https://github.com/containers/omlmd""") + raise RuntimeError(err) else: model_metadata = ModelMetadata.from_dict(kwargs) - meta_path.write_text(model_metadata.to_json()) + meta_path.write_text(json.dumps(model_metadata.to_dict())) + + owns_model_tar = False + owns_md_tar = False + manifest_path = path.parent / "manifest.json" + model_tar = None + meta_tar = None + if not as_artifact: + manifest_path.write_text( + json.dumps( + { + "architecture": get_arch(), + "os": "linux", + } + ) + ) + config = f"{manifest_path}:{MIME_MANIFEST_CONFIG}" + model_tar = path.parent / f"{path.stem}.tar" + meta_tar = path.parent / f"{meta_path.stem}.tar" + if not model_tar.exists(): + owns_model_tar = True + with tarfile.open(model_tar, "w") as tf: + tf.add(path, arcname=path.name) + if not meta_tar.exists(): + owns_md_tar = True + with tarfile.open(meta_tar, "w:gz") as tf: + tf.add(meta_path, arcname=meta_path.name) + files = [ + f"{model_tar}:{MIME_BLOB}", + f"{meta_tar}:{MIME_BLOB}+gzip", + ] + else: + manifest_path.write_text( + json.dumps( + { + "artifactType": MIME_APPLICATION_MLMODEL, + } + ) + ) + config = f"{manifest_path}:{MIME_APPLICATION_MLMODEL}" + files = [ + f"{path}:{MIME_APPLICATION_MLMODEL}", + f"{meta_path}:{MIME_APPLICATION_MLMETADATA}", + ] + + print("uploading", files, "to", target) - config = f"{meta_path}:{MIME_APPLICATION_CONFIG}" - files = [ - f"{path}:{MIME_APPLICATION_MLMODEL}", - config, - ] try: # print(target, files, model_metadata.to_annotations_dict()) result = self._registry.push( @@ -88,6 +145,12 @@ def push( finally: if owns_meta: meta_path.unlink() + if owns_model_tar: + assert isinstance(model_tar, Path) + model_tar.unlink() + if owns_md_tar: + assert isinstance(meta_tar, Path) + meta_tar.unlink() def pull( self, target: str, outdir: Path | str, media_types: Sequence[str] | None = None diff --git a/tests/test_helpers.py b/tests/test_helpers.py index 00bc193..ca47b1d 100644 --- a/tests/test_helpers.py +++ b/tests/test_helpers.py @@ -1,5 +1,7 @@ +import io import json import subprocess +import tarfile import tempfile import typing as t from hashlib import sha256 @@ -7,13 +9,19 @@ import pytest -from omlmd.constants import MIME_APPLICATION_MLMODEL +from omlmd.constants import MIME_BLOB from omlmd.helpers import Helper from omlmd.listener import Event, Listener from omlmd.model_metadata import ModelMetadata, deserialize_mdfile from omlmd.provider import OMLMDRegistry +def untar(tar: Path, out: Path): + out.write_bytes( + t.cast(io.BufferedReader, tarfile.open(tar, "r:*").extractfile(tar.stem)).read() + ) + + def test_call_push_using_md_from_file(mocker): helper = Helper() mocker.patch.object(helper, "push", return_value=None) @@ -75,7 +83,8 @@ def update(self, source: t.Any, event: Event) -> None: @pytest.mark.e2e def test_push_pull_chunked(tmp_path, target): - omlmd = Helper() + omlmd = Helper(OMLMDRegistry()) + target = "quay.io/idoamara/omlmd_test:v4" md = { "name": "mnist", @@ -100,12 +109,33 @@ def test_push_pull_chunked(tmp_path, target): omlmd.push(target, temp, **md) omlmd.pull(target, tmp_path) - assert len(list(tmp_path.iterdir())) == 3 - assert tmp_path.joinpath(temp.name).stat().st_size == base_size + files = list(tmp_path.iterdir()) + print(files) + assert len(files) == 2 + print(tmp_path) + out = tmp_path.joinpath(temp.name) + untar(out.with_suffix(".tar"), out) + assert temp.stat().st_size == base_size finally: temp.unlink() +@pytest.mark.e2e +def test_e2e_push_pull_as_artifact(tmp_path, target): + omlmd = Helper() + omlmd.push( + target, + Path(__file__).parent / ".." / "README.md", + as_artifact=True, + name="mnist", + description="Lorem ipsum", + author="John Doe", + accuracy=0.987, + ) + omlmd.pull(target, tmp_path) + assert len(list(tmp_path.iterdir())) == 2 + + @pytest.mark.e2e def test_e2e_push_pull(tmp_path, target): omlmd = Helper() @@ -118,7 +148,7 @@ def test_e2e_push_pull(tmp_path, target): accuracy=0.987, ) omlmd.pull(target, tmp_path) - assert len(list(tmp_path.iterdir())) == 3 + assert len(list(tmp_path.iterdir())) == 2 @pytest.mark.e2e @@ -132,7 +162,7 @@ def test_e2e_push_pull_with_filters(tmp_path, target): author="John Doe", accuracy=0.987, ) - omlmd.pull(target, tmp_path, media_types=[MIME_APPLICATION_MLMODEL]) + omlmd.pull(target, tmp_path, media_types=[MIME_BLOB]) assert len(list(tmp_path.iterdir())) == 1 @@ -155,10 +185,11 @@ def test_e2e_push_pull_column(tmp_path, target): omlmd.push(target, temp, **md) omlmd.pull(target, tmp_path) - with open(tmp_path.joinpath(temp.name), "r") as f: - pulled = f.read() - assert pulled == content - pulled_sha = sha256(pulled.encode("utf-8")).hexdigest() - assert pulled_sha == content_sha + out = tmp_path.joinpath(temp.name) + untar(out.with_suffix(".tar"), out) + pulled = out.read_text() + assert pulled == content + pulled_sha = sha256(pulled.encode("utf-8")).hexdigest() + assert pulled_sha == content_sha finally: temp.unlink()