Cherry pick the Image loader optimization from the NSDI AE

acto/engine.py

@@ -50,6 +50,7 @@
 from acto.serialization import ActoEncoder, ContextEncoder
 from acto.snapshot import Snapshot
 from acto.utils import delete_operator_pod, process_crd
+from acto.utils.image_helper import ImageHelper
 from acto.utils.preprocess import get_existing_images
 from acto.utils.thread_logger import get_thread_logger, set_thread_logger_prefix
 from ssa.analysis import analyze
@@ -290,13 +291,15 @@
         constraints: Optional[list[XorCondition]] = None,
     ) -> None:
         self.context = context
         self.workdir = workdir
         self.base_workdir = workdir
         self.cluster = cluster
-        self.images_archive = os.path.join(workdir, "images.tar")
+        self.images_archive = ImageHelper.prepare_image_archive(
+            self.context["preload_images"],
+        )
         self.worker_id = worker_id
         self.sequence_base = sequence_base  # trial number to start with
         self.context_name = cluster.get_context_name(
             f"acto-{acto_namespace}-cluster-{worker_id}"
         )
         self.kubeconfig = os.path.join(
@@ -880,7 +883,6 @@
         self.crd_name = operator_config.crd_name
         self.crd_version = operator_config.crd_version
         self.workdir_path = workdir_path
-        self.images_archive = os.path.join(workdir_path, "images.tar")
         self.num_workers = num_workers
         self.dryrun = dryrun
         self.is_reproduce = is_reproduce
@@ -1132,24 +1134,14 @@
         logger = get_thread_logger(with_prefix=True)
 
         # Build an archive to be preloaded
         if len(self.context["preload_images"]) > 0:
             logger.info("Creating preload images archive")
             print_event("Preparing required images...")
-            # first make sure images are present locally
-            for image in self.context["preload_images"]:
-                subprocess.run(
-                    [self.tool, "pull", image],
-                    stdout=subprocess.DEVNULL,
-                    check=True,
-                )
-            subprocess.run(
-                [self.tool, "image", "save", "-o", self.images_archive]
-                + list(self.context["preload_images"]),
-                stdout=subprocess.DEVNULL,
-                check=True,
+            ImageHelper.prepare_image_archive(
+                self.context["preload_images"],
             )
 
         start_time = time.time()
 
         errors: list[OracleResults] = []
         runners: list[TrialRunner] = []

acto/post_process/post_diff_test.py

@@ -10,7 +10,6 @@
 import os
 import queue
 import re
-import subprocess
 import sys
 import threading
 import time
@@ -43,6 +42,7 @@
 from acto.snapshot import Snapshot
 from acto.trial import Step
 from acto.utils import add_acto_label, error_handler, get_thread_logger
+from acto.utils.image_helper import ImageHelper
 
 
 class DiffTestResult(pydantic.BaseModel):
@@ -383,8 +383,10 @@
         self._kubeconfig = os.path.join(
             os.path.expanduser("~"), ".kube", self._context_name
         )
         self._generation = 0
-        self._images_archive = os.path.join(workdir, "images.tar")
+        self._images_archive = ImageHelper.prepare_image_archive(
+            self._context["preload_images"],
+        )
 
     def run_cr(self, cr, trial, gen):
         """Run a CR and return the snapshot"""
@@ -452,10 +454,12 @@
         self._context_name = cluster.get_context_name(
             f"acto-{acto_namespace}-cluster-{worker_id}"
         )
         self._kubeconfig = os.path.join(
             os.path.expanduser("~"), ".kube", self._context_name
         )
-        self._images_archive = os.path.join(workdir, "images.tar")
+        self._images_archive = ImageHelper.prepare_image_archive(
+            self._context["preload_images"],
+        )
 
     def run(self):
         """Run the deploy runner"""
@@ -699,21 +703,14 @@
             num_nodes=self.config.num_nodes,
             version=self.config.kubernetes_version,
         )
         deploy = Deploy(self.config.deploy)
         # Build an archive to be preloaded
-        images_archive = os.path.join(workdir, "images.tar")
-        if len(self.context["preload_images"]) > 0:
-            # first make sure images are present locally
-            for image in self.context["preload_images"]:
-                subprocess.run(["docker", "pull", image], check=True)
-            subprocess.run(
-                ["docker", "image", "save", "-o", images_archive]
-                + list(self.context["preload_images"]),
-                check=True,
-            )
+        ImageHelper.prepare_image_archive(
+            self.context["preload_images"],
+        )
 
         workqueue: multiprocessing.Queue = multiprocessing.Queue()
         for unique_input_group in self.unique_inputs.values():
             workqueue.put(unique_input_group)
 
         runners: list[DeployRunner] = []

acto/post_process/simple_crash_test.py

@@ -8,7 +8,6 @@
 import os
 import queue
 import re
-import subprocess
 import sys
 import threading
 import time
@@ -32,6 +31,7 @@
 from acto.serialization import ActoEncoder
 from acto.trial import Step
 from acto.utils.error_handler import handle_excepthook, thread_excepthook
+from acto.utils.image_helper import ImageHelper
 
 
 def get_crash_config_map(
@@ -269,21 +269,15 @@
             num_nodes=self.config.num_nodes,
             version=self.config.kubernetes_version,
         )
         deploy = Deploy(self.config.deploy)
 
         # Build an archive to be preloaded
-        images_archive = os.path.join(workdir, "images.tar")
-        if len(self.context["preload_images"]) > 0:
-            # first make sure images are present locally
-            for image in self.context["preload_images"]:
-                subprocess.run(["docker", "pull", image])
-            subprocess.run(
-                ["docker", "image", "save", "-o", images_archive]
-                + list(self.context["preload_images"])
-            )
+        ImageHelper.prepare_image_archive(
+            self._context["preload_images"],
+        )
 
         ################## Operation sequence crash test ######################
         num_ops = 0
         workqueue = multiprocessing.Queue()
         for trial_name, trial in self._trial_to_steps.items():
             new_steps = {}

acto/utils/image_helper.py

@@ -0,0 +1,51 @@
+import os
+import subprocess
+
+from filelock import FileLock
+
+
+class ImageHelper:
+    """Helper class for managing Docker images, including pulling and archiving them."""
+
+    image_archive_prefix = os.path.join(os.getcwd(), ".acto_images")
+    image_tool = os.getenv("IMAGE_TOOL", "docker")
+
+    @staticmethod
+    def prepare_image_archive(images: list[str]) -> str:
+        """
+        Prepare an archive of images for testing.
+
+        Args:
+            images (list[str]): List of image file paths to include in the archive.
+
+        Returns:
+            str: Path to the created archive.
+        """
+
+        digest = hash("".join(sorted(images)))
+
+        archive_name = f"{digest}.tar"
+        archive_path = os.path.join(
+            ImageHelper.image_archive_prefix, archive_name
+        )
+
+        lock = FileLock(f"{archive_path}.lock")
+        with lock:
+            if os.path.exists(archive_path):
+                return archive_path
+
+            for image in images:
+                subprocess.run(
+                    [ImageHelper.image_tool, "pull", image],
+                    stdout=subprocess.DEVNULL,
+                    check=True,
+                )
+            os.makedirs(ImageHelper.image_archive_prefix, exist_ok=True)
+            subprocess.run(
+                [ImageHelper.image_tool, "image", "save", "-o", archive_path]
+                + list(images),
+                stdout=subprocess.DEVNULL,
+                check=True,
+            )
+
+        return archive_path