Skip to content

Commit

Permalink
upload as image
Browse files Browse the repository at this point in the history
Signed-off-by: Isabella do Amaral <[email protected]>
  • Loading branch information
isinyaaa committed Nov 11, 2024
1 parent b773b29 commit 9fc1337
Show file tree
Hide file tree
Showing 7 changed files with 152 additions and 50 deletions.
26 changes: 13 additions & 13 deletions e2e/quay-lite/config.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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
2 changes: 1 addition & 1 deletion e2e/quay-lite/quay-app-config.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -4,4 +4,4 @@ metadata:
name: quay-app-config
type: Opaque
data:
config.yaml: REFUQV9NT0RFTF9DQUNIRV9DT05GSUc6CiAgZW5naW5lOiByZWRpcwogIHJlZGlzX2NvbmZpZzoKICAgIHByaW1hcnk6CiAgICAgIGhvc3Q6IHF1YXktcmVkaXMKU1VQRVJfVVNFUlM6Ci0gYWRtaW4KLSB1c2VyMQpBVVRIRU5USUNBVElPTl9UWVBFOiBEYXRhYmFzZQpEQl9VUkk6IHBvc3RncmVzcWw6Ly9xdWF5OnF1YXlAcXVheS1wb3N0Z3Jlc3FsOjU0MzIvcXVheQpCVUlMRExPR1NfUkVESVM6CiAgaG9zdDogcXVheS1yZWRpcwogIHBvcnQ6IDYzNzkKVVNFUl9FVkVOVFNfUkVESVM6CiAgaG9zdDogcXVheS1yZWRpcwogIHBvcnQ6IDYzNzkKQklUVE9SUkVOVF9GSUxFTkFNRV9QRVBQRVI6IDBlZTE4ZjkwLTViNmQtNDJkMi1hYjVlLWVjOWZjZDg0NjI3MgpEQVRBQkFTRV9TRUNSRVRfS0VZOiAnMzAwNjAzNjE2NDA3OTMxODc2MTM2OTczNjY5MjMyMTExMTMyMDU2NzY5MjU0NDU2NTAyNTAyNzQ3NTIxMjUwODM5NzE2MzgzNzYyMjQnCkRFRkFVTFRfVEFHX0VYUElSQVRJT046IDJ3CkRJU1RSSUJVVEVEX1NUT1JBR0VfQ09ORklHOgogIGRlZmF1bHQ6CiAgLSBMb2NhbFN0b3JhZ2UKICAtIHN0b3JhZ2VfcGF0aDogL2RhdGFzdG9yYWdlL3JlZ2lzdHJ5CkRJU1RSSUJVVEVEX1NUT1JBR0VfREVGQVVMVF9MT0NBVElPTlM6IFtdCkRJU1RSSUJVVEVEX1NUT1JBR0VfUFJFRkVSRU5DRToKLSBkZWZhdWx0CkVOVEVSUFJJU0VfTE9HT19VUkw6IC9zdGF0aWMvaW1nL3F1YXktaG9yaXpvbnRhbC1jb2xvci5zdmcKRVhURVJOQUxfVExTX1RFUk1JTkFUSU9OOiB0cnVlCkZFQVRVUkVfQU5PTllNT1VTX0FDQ0VTUzogdHJ1ZQpGRUFUVVJFX0FQUF9SRUdJU1RSWTogZmFsc2UKRkVBVFVSRV9BUFBfU1BFQ0lGSUNfVE9LRU5TOiB0cnVlCkZFQVRVUkVfQlVJTERfU1VQUE9SVDogZmFsc2UKRkVBVFVSRV9DSEFOR0VfVEFHX0VYUElSQVRJT046IHRydWUKRkVBVFVSRV9ESVJFQ1RfTE9HSU46IHRydWUKRkVBVFVSRV9NQUlMSU5HOiBmYWxzZQpGRUFUVVJFX1BBUlRJQUxfVVNFUl9BVVRPQ09NUExFVEU6IHRydWUKRkVBVFVSRV9SRVBPX01JUlJPUjogZmFsc2UKRkVBVFVSRV9SRVFVSVJFX1RFQU1fSU5WSVRFOiB0cnVlCkZFQVRVUkVfUkVTVFJJQ1RFRF9WMV9QVVNIOiBmYWxzZQpGRUFUVVJFX1NFQ1VSSVRZX05PVElGSUNBVElPTlM6IGZhbHNlCkZFQVRVUkVfU0VDVVJJVFlfU0NBTk5FUjogZmFsc2UKRkVBVFVSRV9VU0VSTkFNRV9DT05GSVJNQVRJT046IHRydWUKRkVBVFVSRV9VU0VSX0lOSVRJQUxJWkU6IHRydWUKRkVBVFVSRV9VU0VSX0NSRUFUSU9OOiB0cnVlCkZFQVRVUkVfVVNFUl9MT0dfQUNDRVNTOiB0cnVlCkZFQVRVUkVfUFJPWFlfQ0FDSEU6IHRydWUKR0lUSFVCX0xPR0lOX0NPTkZJRzoge30KR0lUSFVCX1RSSUdHRVJfQ09ORklHOiB7fQpHSVRMQUJfVFJJR0dFUl9LSU5EOiB7fQpMT0dfQVJDSElWRV9MT0NBVElPTjogZGVmYXVsdApNQUlMX0RFRkFVTFRfU0VOREVSOiBhZG1pbkBleGFtcGxlLmNvbQpNQUlMX1BPUlQ6IDU4NwpNQUlMX1VTRV9UTFM6IHRydWUKUFJFRkVSUkVEX1VSTF9TQ0hFTUU6IGh0dHAKUkVHSVNUUllfVElUTEU6IFJlZCBIYXQgUXVheSBMSVRFClJFR0lTVFJZX1RJVExFX1NIT1JUOiBSZWQgSGF0IFF1YXkgTElURQpSRVBPX01JUlJPUl9TRVJWRVJfSE9TVE5BTUU6IG51bGwKUkVQT19NSVJST1JfVExTX1ZFUklGWTogdHJ1ZQpTRVRVUF9DT01QTEVURTogdHJ1ZQpTSUdOSU5HX0VOR0lORTogZ3BnMgpUQUdfRVhQSVJBVElPTl9PUFRJT05TOgotIDBzCi0gMWQKLSAxdwotIDJ3Ci0gNHcKVEVBTV9SRVNZTkNfU1RBTEVfVElNRTogNjBtClRFU1RJTkc6IGZhbHNlClVTRVJGSUxFU19MT0NBVElPTjogZGVmYXVsdApVU0VSRklMRVNfUEFUSDogdXNlcmZpbGVzLwpVU0VfQ0ROOiBmYWxzZQpGRUFUVVJFX1FVT1RBX01BTkFHRU1FTlQ6IFRydWUKU0VSVkVSX0hPU1ROQU1FOiBsb2NhbGhvc3Q6NTAwMQpCUk9XU0VSX0FQSV9DQUxMU19YSFJfT05MWTogRmFsc2UKQ09SU19PUklHSU46CiAgLSAiaHR0cHM6Ly9zdGFnZS5mb28ucmVkaGF0LmNvbToxMzM3IgogIC0gImh0dHA6Ly9sb2NhbGhvc3Q6OTAwMCIKRkVBVFVSRV9VSV9WMjogVHJ1ZQpGRUFUVVJFX1VTRVJfTUVUQURBVEE6IFRydWUKSUdOT1JFX1VOS05PV05fTUVESUFUWVBFUzogVHJ1ZQo=
config.yaml: REFUQV9NT0RFTF9DQUNIRV9DT05GSUc6CiAgZW5naW5lOiByZWRpcwogIHJlZGlzX2NvbmZpZzoKICAgIHByaW1hcnk6CiAgICAgIGhvc3Q6IHF1YXktcmVkaXMKU1VQRVJfVVNFUlM6CiAgLSBhZG1pbgogIC0gdXNlcjEKQVVUSEVOVElDQVRJT05fVFlQRTogRGF0YWJhc2UKREJfVVJJOiBwb3N0Z3Jlc3FsOi8vcXVheTpxdWF5QHF1YXktcG9zdGdyZXNxbDo1NDMyL3F1YXkKQlVJTERMT0dTX1JFRElTOgogIGhvc3Q6IHF1YXktcmVkaXMKICBwb3J0OiA2Mzc5ClVTRVJfRVZFTlRTX1JFRElTOgogIGhvc3Q6IHF1YXktcmVkaXMKICBwb3J0OiA2Mzc5CkJJVFRPUlJFTlRfRklMRU5BTUVfUEVQUEVSOiAwZWUxOGY5MC01YjZkLTQyZDItYWI1ZS1lYzlmY2Q4NDYyNzIKREFUQUJBU0VfU0VDUkVUX0tFWTogIjMwMDYwMzYxNjQwNzkzMTg3NjEzNjk3MzY2OTIzMjExMTEzMjA1Njc2OTI1NDQ1NjUwMjUwMjc0NzUyMTI1MDgzOTcxNjM4Mzc2MjI0IgpERUZBVUxUX1RBR19FWFBJUkFUSU9OOiAydwpESVNUUklCVVRFRF9TVE9SQUdFX0NPTkZJRzoKICBkZWZhdWx0OgogICAgLSBMb2NhbFN0b3JhZ2UKICAgIC0gc3RvcmFnZV9wYXRoOiAvZGF0YXN0b3JhZ2UvcmVnaXN0cnkKRElTVFJJQlVURURfU1RPUkFHRV9ERUZBVUxUX0xPQ0FUSU9OUzogW10KRElTVFJJQlVURURfU1RPUkFHRV9QUkVGRVJFTkNFOgogIC0gZGVmYXVsdApFTlRFUlBSSVNFX0xPR09fVVJMOiAvc3RhdGljL2ltZy9xdWF5LWhvcml6b250YWwtY29sb3Iuc3ZnCkVYVEVSTkFMX1RMU19URVJNSU5BVElPTjogdHJ1ZQpGRUFUVVJFX0FOT05ZTU9VU19BQ0NFU1M6IHRydWUKRkVBVFVSRV9BUFBfUkVHSVNUUlk6IGZhbHNlCkZFQVRVUkVfQVBQX1NQRUNJRklDX1RPS0VOUzogdHJ1ZQpGRUFUVVJFX0JVSUxEX1NVUFBPUlQ6IGZhbHNlCkZFQVRVUkVfQ0hBTkdFX1RBR19FWFBJUkFUSU9OOiB0cnVlCkZFQVRVUkVfRElSRUNUX0xPR0lOOiB0cnVlCkZFQVRVUkVfTUFJTElORzogZmFsc2UKRkVBVFVSRV9QQVJUSUFMX1VTRVJfQVVUT0NPTVBMRVRFOiB0cnVlCkZFQVRVUkVfUkVQT19NSVJST1I6IGZhbHNlCkZFQVRVUkVfUkVRVUlSRV9URUFNX0lOVklURTogdHJ1ZQpGRUFUVVJFX1JFU1RSSUNURURfVjFfUFVTSDogZmFsc2UKRkVBVFVSRV9TRUNVUklUWV9OT1RJRklDQVRJT05TOiBmYWxzZQpGRUFUVVJFX1NFQ1VSSVRZX1NDQU5ORVI6IGZhbHNlCkZFQVRVUkVfVVNFUk5BTUVfQ09ORklSTUFUSU9OOiB0cnVlCkZFQVRVUkVfVVNFUl9JTklUSUFMSVpFOiB0cnVlCkZFQVRVUkVfVVNFUl9DUkVBVElPTjogdHJ1ZQpGRUFUVVJFX1VTRVJfTE9HX0FDQ0VTUzogdHJ1ZQpGRUFUVVJFX1BST1hZX0NBQ0hFOiB0cnVlCkdJVEhVQl9MT0dJTl9DT05GSUc6IHt9CkdJVEhVQl9UUklHR0VSX0NPTkZJRzoge30KR0lUTEFCX1RSSUdHRVJfS0lORDoge30KTE9HX0FSQ0hJVkVfTE9DQVRJT046IGRlZmF1bHQKTUFJTF9ERUZBVUxUX1NFTkRFUjogYWRtaW5AZXhhbXBsZS5jb20KTUFJTF9QT1JUOiA1ODcKTUFJTF9VU0VfVExTOiB0cnVlClBSRUZFUlJFRF9VUkxfU0NIRU1FOiBodHRwClJFR0lTVFJZX1RJVExFOiBSZWQgSGF0IFF1YXkgTElURQpSRUdJU1RSWV9USVRMRV9TSE9SVDogUmVkIEhhdCBRdWF5IExJVEUKUkVQT19NSVJST1JfU0VSVkVSX0hPU1ROQU1FOiBudWxsClJFUE9fTUlSUk9SX1RMU19WRVJJRlk6IHRydWUKU0VUVVBfQ09NUExFVEU6IHRydWUKU0lHTklOR19FTkdJTkU6IGdwZzIKVEFHX0VYUElSQVRJT05fT1BUSU9OUzoKICAtIDBzCiAgLSAxZAogIC0gMXcKICAtIDJ3CiAgLSA0dwpURUFNX1JFU1lOQ19TVEFMRV9USU1FOiA2MG0KVEVTVElORzogdHJ1ZQpVU0VSRklMRVNfTE9DQVRJT046IGRlZmF1bHQKVVNFUkZJTEVTX1BBVEg6IHVzZXJmaWxlcy8KVVNFX0NETjogZmFsc2UKRkVBVFVSRV9RVU9UQV9NQU5BR0VNRU5UOiBUcnVlClNFUlZFUl9IT1NUTkFNRTogbG9jYWxob3N0OjUwMDEKQlJPV1NFUl9BUElfQ0FMTFNfWEhSX09OTFk6IEZhbHNlCkNPUlNfT1JJR0lOOgogIC0gImh0dHBzOi8vc3RhZ2UuZm9vLnJlZGhhdC5jb206MTMzNyIKICAtICJodHRwOi8vbG9jYWxob3N0OjkwMDAiCkZFQVRVUkVfVUlfVjI6IFRydWUKRkVBVFVSRV9VU0VSX01FVEFEQVRBOiBUcnVlCklHTk9SRV9VTktOT1dOX01FRElBVFlQRVM6IGZhbHNlCg==
8 changes: 4 additions & 4 deletions e2e/test_cli.sh
Original file line number Diff line number Diff line change
Expand Up @@ -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."
Expand Down
10 changes: 9 additions & 1 deletion omlmd/cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand All @@ -88,6 +93,7 @@ def push(
plain_http: bool,
target: str,
path: Path,
as_artifact: bool,
metadata: Path | None,
empty_metadata: bool,
):
Expand All @@ -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)
)
6 changes: 5 additions & 1 deletion omlmd/constants.py
Original file line number Diff line number Diff line change
@@ -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"
100 changes: 80 additions & 20 deletions omlmd/helpers.py
Original file line number Diff line number Diff line change
@@ -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
Expand All @@ -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
Expand All @@ -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)
Expand All @@ -41,39 +58,76 @@ def push(
self,
target: str,
path: Path | str,
as_artifact: bool = False,
**kwargs,
):
owns_meta = True
if isinstance(path, str):
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}",
]

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(
target=target,
files=files,
Expand All @@ -88,6 +142,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
Expand Down
50 changes: 40 additions & 10 deletions tests/test_helpers.py
Original file line number Diff line number Diff line change
@@ -1,19 +1,27 @@
import io
import json
import subprocess
import tarfile
import tempfile
import typing as t
from hashlib import sha256
from pathlib import Path

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)
Expand Down Expand Up @@ -100,12 +108,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()
Expand All @@ -118,7 +147,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
Expand All @@ -132,7 +161,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


Expand All @@ -155,10 +184,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()

0 comments on commit 9fc1337

Please sign in to comment.