diff --git a/acto/engine.py b/acto/engine.py index 845c65af3..e6580d3eb 100644 --- a/acto/engine.py +++ b/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 @@ -293,7 +294,9 @@ def __init__( 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( @@ -880,7 +883,6 @@ def __init__( 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 @@ -1135,18 +1137,8 @@ def run(self) -> list[OracleResults]: 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() diff --git a/acto/post_process/post_diff_test.py b/acto/post_process/post_diff_test.py index a2ff587ca..6a41b28a7 100644 --- a/acto/post_process/post_diff_test.py +++ b/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): @@ -384,7 +384,9 @@ def __init__( 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""" @@ -455,7 +457,9 @@ def __init__( 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""" @@ -701,16 +705,9 @@ def post_process(self, workdir: str, num_workers: int = 1): ) 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(): diff --git a/acto/post_process/simple_crash_test.py b/acto/post_process/simple_crash_test.py index bcc1ae707..81ea99e1b 100644 --- a/acto/post_process/simple_crash_test.py +++ b/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( @@ -272,15 +272,9 @@ def post_process(self, workdir: str, num_workers: int = 1): 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 diff --git a/acto/utils/image_helper.py b/acto/utils/image_helper.py new file mode 100644 index 000000000..196bdd29f --- /dev/null +++ b/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