From 82e70e8ef1e3ec8cf8d51e0a0672d39ed4bbb851 Mon Sep 17 00:00:00 2001 From: Tyler Gu Date: Thu, 31 Jul 2025 15:15:04 -0700 Subject: [PATCH 1/2] Optimize the image preload Signed-off-by: Tyler Gu --- acto/engine.py | 20 ++++-------- acto/post_process/post_diff_test.py | 23 ++++++------- acto/post_process/simple_crash_test.py | 14 +++----- acto/utils/image_helper.py | 45 ++++++++++++++++++++++++++ 4 files changed, 65 insertions(+), 37 deletions(-) create mode 100644 acto/utils/image_helper.py diff --git a/acto/engine.py b/acto/engine.py index 845c65af3e..e6580d3ebb 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 a2ff587ca4..6a41b28a78 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 bcc1ae7073..81ea99e1b6 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 0000000000..2f9ee5b3e6 --- /dev/null +++ b/acto/utils/image_helper.py @@ -0,0 +1,45 @@ +import os +import subprocess + + +class ImageHelper: + 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 + ) + + 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 From e02e667f674b63d29657b2d0fc9270ac9c97318d Mon Sep 17 00:00:00 2001 From: Tyler Gu Date: Thu, 31 Jul 2025 15:49:32 -0700 Subject: [PATCH 2/2] Use filelock to prevent race condition for creating image archives Signed-off-by: Tyler Gu --- acto/utils/image_helper.py | 30 ++++++++++++++++++------------ 1 file changed, 18 insertions(+), 12 deletions(-) diff --git a/acto/utils/image_helper.py b/acto/utils/image_helper.py index 2f9ee5b3e6..196bdd29fb 100644 --- a/acto/utils/image_helper.py +++ b/acto/utils/image_helper.py @@ -1,8 +1,12 @@ 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") @@ -25,21 +29,23 @@ def prepare_image_archive(images: list[str]) -> str: ImageHelper.image_archive_prefix, archive_name ) - if os.path.exists(archive_path): - return archive_path - - for image in images: + 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, "pull", image], + [ImageHelper.image_tool, "image", "save", "-o", archive_path] + + list(images), 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