Implement stage-in for zip staging provider (#3395)

Author: benclifford

parsl/data_provider/zip.py

@@ -42,6 +42,12 @@ class ZipFileStaging(Staging):
     """
 
     def can_stage_out(self, file: File) -> bool:
+        return self.is_zip_url(file)
+
+    def can_stage_in(self, file: File) -> bool:
+        return self.is_zip_url(file)
+
+    def is_zip_url(self, file: File) -> bool:
         logger.debug("archive provider checking File {}".format(repr(file)))
 
         # First check if this is the scheme we care about
@@ -76,6 +82,20 @@ def stage_out(self, dm, executor, file, parent_fut):
         app_fut = stage_out_app(zip_path, inside_path, working_dir, inputs=[file], _parsl_staging_inhibit=True, parent_fut=parent_fut)
         return app_fut
 
+    def stage_in(self, dm, executor, file, parent_fut):
+        assert file.scheme == 'zip'
+
+        zip_path, inside_path = zip_path_split(file.path)
+
+        working_dir = dm.dfk.executors[executor].working_dir
+
+        if working_dir:
+            file.local_path = os.path.join(working_dir, inside_path)
+
+        stage_in_app = _zip_stage_in_app(dm)
+        app_fut = stage_in_app(zip_path, inside_path, working_dir, outputs=[file], _parsl_staging_inhibit=True, parent_fut=parent_fut)
+        return app_fut._outputs[0]
+
 
 def _zip_stage_out(zip_file, inside_path, working_dir, parent_fut=None, inputs=[], _parsl_staging_inhibit=True):
     file = inputs[0]
@@ -93,6 +113,18 @@ def _zip_stage_out_app(dm):
     return parsl.python_app(executors=['_parsl_internal'], data_flow_kernel=dm.dfk)(_zip_stage_out)
 
 
+def _zip_stage_in(zip_file, inside_path, working_dir, *, parent_fut, outputs, _parsl_staging_inhibit=True):
+    with filelock.FileLock(zip_file + ".lock"):
+        with zipfile.ZipFile(zip_file, mode='r') as z:
+            content = z.read(inside_path)
+        with open(outputs[0], "wb") as of:
+            of.write(content)
+
+
+def _zip_stage_in_app(dm):
+    return parsl.python_app(executors=['_parsl_internal'], data_flow_kernel=dm.dfk)(_zip_stage_in)
+
+
 def zip_path_split(path: str) -> Tuple[str, str]:
     """Split zip: path into a zipfile name and a contained-file name.
     """

parsl/tests/test_staging/test_zip_in.py

@@ -0,0 +1,42 @@
+import parsl
+import pytest
+import random
+import zipfile
+
+from parsl.data_provider.files import File
+from parsl.data_provider.zip import ZipAuthorityError, ZipFileStaging
+
+from parsl.providers import LocalProvider
+from parsl.channels import LocalChannel
+from parsl.launchers import SimpleLauncher
+
+from parsl.config import Config
+from parsl.executors import HighThroughputExecutor
+
+from parsl.tests.configs.htex_local import fresh_config as local_config
+
+
+@parsl.python_app
+def count_lines(file):
+    with open(file, "r") as f:
+        return len(f.readlines())
+
+
+@pytest.mark.local
+def test_zip_in(tmpd_cwd):
+    # basic test of zip file stage-in
+    zip_path = tmpd_cwd / "container.zip"
+    file_base = "data.txt"
+    zip_file = File(f"zip:{zip_path / file_base}")
+
+    # create a zip file containing one file with some abitrary number of lines
+    n_lines = random.randint(0, 1000)
+
+    with zipfile.ZipFile(zip_path, mode='w') as z:
+        with z.open(file_base, mode='w') as f:
+            for _ in range(n_lines):
+                f.write(b'someline\n')
+
+    app_future = count_lines(zip_file)
+
+    assert app_future.result() == n_lines

parsl/tests/test_staging/test_zip_to_zip.py

@@ -0,0 +1,44 @@
+import parsl
+import pytest
+import random
+import zipfile
+
+from parsl.data_provider.files import File
+from parsl.data_provider.zip import ZipAuthorityError, ZipFileStaging
+
+from parsl.providers import LocalProvider
+from parsl.channels import LocalChannel
+from parsl.launchers import SimpleLauncher
+
+from parsl.config import Config
+from parsl.executors import HighThroughputExecutor
+
+from parsl.tests.configs.htex_local import fresh_config as local_config
+
+
+@parsl.python_app
+def generate_lines(n: int, *, outputs):
+    with open(outputs[0], "w") as f:
+        for x in range(n):
+            # write numbered lines
+            f.write(str(x) + "\n")
+
+
+@parsl.python_app
+def count_lines(file):
+    with open(file, "r") as f:
+        return len(f.readlines())
+
+
+@pytest.mark.local
+def test_zip_pipeline(tmpd_cwd):
+    # basic test of zip file stage-in
+    zip_path = tmpd_cwd / "container.zip"
+    file_base = "data.txt"
+    zip_file = File(f"zip:{zip_path / file_base}")
+
+    n_lines = random.randint(0, 1000)
+    generate_fut = generate_lines(n_lines, outputs=[zip_file])
+    n_lines_out = count_lines(generate_fut.outputs[0]).result()
+
+    assert n_lines == n_lines_out