making plough compatible with postprocessing

Author: andrpie

src/pinefarm/cli/run.py

@@ -190,18 +190,18 @@ def run_dataset(runner):
 
             # collect results in the output pineappl grid
             runner.generate_pineappl()
-
-            table.print_table(
-                table.convolute_grid(
-                    runner.grid, runner.pdf, integrated=isinstance(runner, mg5.Mg5)
-                ),
-                runner.results(),
-                runner.dest,
-            )
-
-            # TODO: annotate_version should be a post-processing step
-            # however at the moment only works in 1-grid cases
-            runner.annotate_versions()
+            if not hasattr(runner, "ps_link"):
+                table.print_table(
+                    table.convolute_grid(
+                        runner.grid, runner.pdf, integrated=isinstance(runner, mg5.Mg5)
+                    ),
+                    runner.results(),
+                    runner.dest,
+                )
+
+                # TODO: annotate_version should be a post-processing step
+                # however at the moment only works in 1-grid cases
+                runner.annotate_versions()
 
         runner.postprocess()
 

src/pinefarm/external/plough.py

@@ -2,7 +2,6 @@
 
 import os
 import shutil
-import subprocess
 import tarfile
 import urllib.request
 
@@ -13,7 +12,6 @@
 from . import interface
 
 PLOUGHSHARE_LINK_FILENAME = "ploughshare_link.txt"
-GRIDS_PROCESSOR = "process_grids.sh"
 GRIDS_TMP = "grids"
 
 
@@ -28,19 +26,25 @@ def __init__(self, pinecard, theorycard, *args, **kwargs):
         self.filename = self.link.rsplit("/")[-1]
         self.dir_name = self.filename.rsplit(".", 1)[0]
         self.tarball = self.dest / self.filename
-        self.processor = self.source / GRIDS_PROCESSOR
-        self.run()
-        self.generate_pineappl()
-        self.timestamp = 0
 
     def run(self):
         """Download and extract the .tgz file."""
         print("Downloading from ploughshare...")
-        self.download_to_dest()
+        try:
+            self.download_to_dest()
+            if self.tarball.exists():
+                print(f"Grids successfully downloaded to {self.tarball}")
+            else:
+                raise FileNotFoundError(
+                    f"{self.tarball} not found but the download didn't seem to fail?"
+                )
+        except Exception as e:
+            raise FileNotFoundError(f"{self.tarball} could not be downloaded!") from e
         print(f"Grids successfully downloaded to {self.tarball}")
         print("Extracting files...")
         self.extract_tarball()
-        print(f"Grids successfully extracted to {self.dir_name}")
+        print(f"Grids successfully extracted to {self.dest}")
+        self.cleanup()
 
     def results(self):
         """Results are collected and compared at the pineappl (script) level."""
@@ -51,30 +55,26 @@ def collect_versions(self):
         return {}
 
     def generate_pineappl(self):
-        """Converts donwloaded grids into pineappl format."""
-        print("Grid conversion started...")
-        # the grids are converted and processed here
-        os.environ["PS_DIR"] = str(self.grids_dir)
-        # note that filename is also dir_name
-        os.environ["FILENAME"] = str(self.dir_name)
-        if os.access(self.processor, os.X_OK):
-            shutil.copy2(self.processor, self.dest)
-            subprocess.run("./process_grids.sh", cwd=self.dest, check=True)
-            (self.dest / "process_grids.sh").unlink()
-        else:
-            raise ValueError(
-                f"Grid conversion file present but not executable: {self.processor}"
-            )
-        self.grids = []
-        for g in self.dest.glob("*.pineappl.lz4"):
-            self.grids.append(g)
+        """Grids are converted in postrun.sh."""
+        return
 
     def download_to_dest(self):
-        """Download the file and move it to the output folder."""
+        """Download the file to the output folder."""
         urllib.request.urlretrieve(self.link, self.dest / self.filename)
 
     def extract_tarball(self):
         """Extract the contents."""
         with tarfile.open(self.tarball, "r:*") as tf:
             tf.extractall(self.dest)
         self.grids_dir = self.dest / self.dir_name / GRIDS_TMP
+        grids_list = sorted(os.listdir(self.grids_dir))
+        for i, grid in enumerate(grids_list):
+            extension = grid.split(".", 2)[2]
+            print(extension)
+            os.rename(self.grids_dir / grid, self.dest / f"grid_{i}.{extension}")
+
+    def cleanup(self):
+        """Delete unnecessary files and create tmp.pineappl.lz4 to allow postprocessing."""
+        shutil.rmtree(self.dest / self.dir_name)
+        self.tarball.unlink()
+        open(self.dest / "tmp.pineappl.lz4", "x")