Add multiple include paths option to packagers (#117)

* Multiple include paths

Signed-off-by: Igor Gitman <[email protected]>

* Same for pattern packager

Signed-off-by: Igor Gitman <[email protected]>

* Apply formatting

Signed-off-by: Igor Gitman <[email protected]>

* Add tests

Signed-off-by: Igor Gitman <[email protected]>

---------

Signed-off-by: Igor Gitman <[email protected]>

Author: Kipok

src/nemo_run/core/packaging/git.py

@@ -62,11 +62,11 @@ class GitArchivePackager(Packager):
 
     #: Include extra files in the archive which matches include_pattern
     #: This str will be included in the command as: find {include_pattern} -type f to get the list of extra files to include in the archive
-    include_pattern: str = ""
+    include_pattern: str | list[str] = ""
 
     #: Relative path to use as tar -C option - need to be consistent with include_pattern
     #: If not provided, will use git base path.
-    include_pattern_relative_path: str = ""
+    include_pattern_relative_path: str | list[str] = ""
 
     check_uncommitted_changes: bool = False
     check_untracked_files: bool = False
@@ -125,15 +125,32 @@ def package(self, path: Path, job_dir: str, name: str) -> str:
             ctx.run(git_archive_cmd)
             if self.include_submodules:
                 ctx.run(git_submodule_cmd)
+        if isinstance(self.include_pattern, str):
+            self.include_pattern = [self.include_pattern]
 
-        if self.include_pattern:
-            include_pattern_relative_path = self.include_pattern_relative_path or shlex.quote(
+        if isinstance(self.include_pattern_relative_path, str):
+            self.include_pattern_relative_path = [self.include_pattern_relative_path]
+
+        if len(self.include_pattern) != len(self.include_pattern_relative_path):
+            raise ValueError(
+                "include_pattern and include_pattern_relative_path should have the same length"
+            )
+
+        for include_pattern, include_pattern_relative_path in zip(
+            self.include_pattern, self.include_pattern_relative_path
+        ):
+            if include_pattern == "":
+                continue
+            include_pattern_relative_path = include_pattern_relative_path or shlex.quote(
                 str(git_base_path)
             )
             relative_include_pattern = os.path.relpath(
-                self.include_pattern, include_pattern_relative_path
+                include_pattern, include_pattern_relative_path
+            )
+            include_pattern_cmd = (
+                f"find {relative_include_pattern} -type f | "
+                f"tar -cf {os.path.join(git_base_path, 'additional.tmp')} -T -"
             )
-            include_pattern_cmd = f"find {relative_include_pattern} -type f | tar -cf {os.path.join(git_base_path, 'additional.tmp')} -T -"
             tar_concatenate_cmd = f"cat additional.tmp >> {output_file}.tmp && rm additional.tmp"
 
             with ctx.cd(include_pattern_relative_path):

src/nemo_run/core/packaging/pattern.py

@@ -32,23 +32,47 @@ class PatternPackager(Packager):
     #: This str will be included in the command as:
     #: find {include_pattern} -type f to get the list of extra files to include in the archive
     #: best to use an absolute path here and a proper relative path argument to pass to tar
-    include_pattern: str
+    include_pattern: str | list[str]
 
     #: Relative path to use as tar -C option.
-    relative_path: str
+    relative_path: str | list[str]
 
     def package(self, path: Path, job_dir: str, name: str) -> str:
         output_file = os.path.join(job_dir, f"{name}.tar.gz")
         if os.path.exists(output_file):
             return output_file
 
-        relative_include_pattern = os.path.relpath(self.include_pattern, self.relative_path)
-        cmd = (
-            f"tar -czf {output_file} -C {self.relative_path} -T "
-            f"<(cd {self.relative_path} && find {relative_include_pattern} -type f)"
-        )
+        if isinstance(self.include_pattern, str):
+            self.include_pattern = [self.include_pattern]
+
+        if isinstance(self.relative_path, str):
+            self.relative_path = [self.relative_path]
+
+        if len(self.include_pattern) != len(self.relative_path):
+            raise ValueError("include_pattern and relative_path should have the same length")
+
+        # Create initial empty tar file
         ctx = Context()
-        ctx.run(cmd)
+        ctx.run(f"tar -cf {output_file}.tmp --files-from /dev/null")
+
+        for include_pattern, relative_path in zip(self.include_pattern, self.relative_path):
+            if include_pattern == "":
+                continue
+
+            relative_include_pattern = os.path.relpath(include_pattern, relative_path)
+
+            with ctx.cd(relative_path):
+                # Append files directly to the main tar archive
+                cmd = f"find {relative_include_pattern} -type f -print0 | xargs -0 tar -rf {output_file}.tmp"
+                ctx.run(cmd)
+
+        # Gzip the final result
+        gzip_cmd = f"gzip -c {output_file}.tmp > {output_file}"
+        rm_cmd = f"rm {output_file}.tmp"
+
+        ctx.run(gzip_cmd)
+        ctx.run(rm_cmd)
+
         return output_file
 
 

test/core/packaging/test_git.py

@@ -287,6 +287,39 @@ def test_package_with_include_pattern_rel_path(packager, temp_repo, tmpdir):
         assert not cmp.diff_files
 
 
+@patch("nemo_run.core.packaging.git.Context", MockContext)
+def test_package_with_multi_include_pattern_rel_path(packager, temp_repo, tmpdir):
+    temp_repo = Path(temp_repo)
+    # Create extra files in a separate directory
+    (tmpdir / "extra").mkdir()
+    with open(tmpdir / "extra" / "extra_file1.txt", "w") as f:
+        f.write("Extra file 1")
+    with open(tmpdir / "extra" / "extra_file2.txt", "w") as f:
+        f.write("Extra file 2")
+
+    include_pattern = [str(tmpdir / "extra/extra_file1.txt"), str(tmpdir / "extra/extra_file2.txt")]
+    relative_path = [str(tmpdir), str(tmpdir)]
+
+    packager = GitArchivePackager(
+        include_pattern=include_pattern, include_pattern_relative_path=relative_path
+    )
+    with tempfile.TemporaryDirectory() as job_dir:
+        output_file = packager.package(Path(temp_repo), job_dir, "test_package")
+        assert os.path.exists(output_file)
+        subprocess.check_call(shlex.split(f"mkdir -p {os.path.join(job_dir, 'extracted_output')}"))
+        subprocess.check_call(
+            shlex.split(
+                f"tar -xvzf {output_file} -C {os.path.join(job_dir, 'extracted_output')} --ignore-zeros"
+            ),
+        )
+        cmp = filecmp.dircmp(
+            os.path.join(tmpdir, "extra"),
+            os.path.join(job_dir, "extracted_output", "extra"),
+        )
+        assert cmp.left_list == cmp.right_list
+        assert not cmp.diff_files
+
+
 @patch("nemo_run.core.packaging.git.Context", MockContext)
 def test_package_with_check_uncommitted_changes(packager, temp_repo):
     temp_repo = Path(temp_repo)

test/core/packaging/test_pattern.py

@@ -50,3 +50,33 @@ def test_package_with_include_pattern_rel_path(tmpdir):
         )
         assert cmp.left_list == cmp.right_list
         assert not cmp.diff_files
+
+
+@patch("nemo_run.core.packaging.pattern.Context", MockContext)
+def test_package_with_multi_include_pattern_rel_path(tmpdir):
+    # Create extra files in a separate directory
+    (tmpdir / "extra").mkdir()
+    with open(tmpdir / "extra" / "extra_file1.txt", "w") as f:
+        f.write("Extra file 1")
+    with open(tmpdir / "extra" / "extra_file2.txt", "w") as f:
+        f.write("Extra file 2")
+
+    include_pattern = [str(tmpdir / "extra/extra_file1.txt"), str(tmpdir / "extra/extra_file2.txt")]
+    relative_path = [str(tmpdir), str(tmpdir)]
+
+    packager = PatternPackager(include_pattern=include_pattern, relative_path=relative_path)
+    with tempfile.TemporaryDirectory() as job_dir:
+        output_file = packager.package(Path(tmpdir), job_dir, "test_package")
+        assert os.path.exists(output_file)
+        subprocess.check_call(shlex.split(f"mkdir -p {os.path.join(job_dir, 'extracted_output')}"))
+        subprocess.check_call(
+            shlex.split(
+                f"tar -xvzf {output_file} -C {os.path.join(job_dir, 'extracted_output')} --ignore-zeros"
+            ),
+        )
+        cmp = filecmp.dircmp(
+            os.path.join(tmpdir, "extra"),
+            os.path.join(job_dir, "extracted_output", "extra"),
+        )
+        assert cmp.left_list == cmp.right_list
+        assert not cmp.diff_files