enable large file upload (#13)

* enable large file upload Signed-off-by: Isabella do Amaral <[email protected]> * bump oras-py to 0.2.21 Signed-off-by: tarilabs <[email protected]> * test chunked uploads Signed-off-by: Isabella do Amaral <[email protected]> --------- Signed-off-by: Isabella do Amaral <[email protected]> Signed-off-by: tarilabs <[email protected]> Co-authored-by: tarilabs <[email protected]>
containers · Oct 3, 2024 · f51e18f · f51e18f
1 parent 0b71d27
commit f51e18f
Show file tree

Hide file tree

Showing 5 changed files with 41 additions and 5 deletions.
diff --git a/.gitignore b/.gitignore
@@ -17,3 +17,4 @@ site
 
 # E2E
 venv
+temp
diff --git a/omlmd/helpers.py b/omlmd/helpers.py
@@ -100,6 +100,7 @@ def push(
                 files=files,
                 manifest_annotations=model_metadata.to_annotations_dict(),
                 manifest_config=manifest_cfg,
+                do_chunked=True,
             )
             self.notify_listeners(PushEvent(target, model_metadata))
             return result

diff --git a/poetry.lock b/poetry.lock
diff --git a/pyproject.toml b/pyproject.toml
@@ -14,7 +14,7 @@ Changelog = "https://github.com/containers/omlmd/releases"
 
 [tool.poetry.dependencies]
 python = "^3.9"
-oras = "^0.1.30"
+oras = "^0.2.21"
 pyyaml = "^6.0.1"
 click = "^8.1.7"
 cloup = "^3.0.5"

diff --git a/tests/test_helpers.py b/tests/test_helpers.py
@@ -1,4 +1,5 @@
 import json
+import subprocess
 import tempfile
 from pathlib import Path
 
@@ -64,6 +65,39 @@ def update(self, _, event: Event) -> None:
     assert e0.metadata == ModelMetadata.from_dict(md)
 
 
+@pytest.mark.e2e
+def test_push_pull_chunked(tmp_path, target):
+    omlmd = Helper()
+
+    md = {
+        "name": "mnist",
+        "description": "Lorem ipsum",
+        "author": "John Doe",
+        "accuracy": 0.987,
+    }
+    here = Path.cwd()
+    temp = here / "temp"
+    base_size = 16 * 1024 * 1024 * 3  # 48MB
+    try:
+        subprocess.run(
+            [
+                "dd",
+                "if=/dev/null",
+                f"of={temp}",
+                "bs=1",
+                "count=0",
+                f"seek={base_size}",
+            ],
+        )
+
+        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
+    finally:
+        temp.unlink()
+
+
 @pytest.mark.e2e
 def test_e2e_push_pull(tmp_path, target):
     omlmd = Helper()