From 806707cc6fc7cea9b09aca16ce8f3c9796359f87 Mon Sep 17 00:00:00 2001 From: Tyler Gu Date: Mon, 15 Dec 2025 17:20:10 -0600 Subject: [PATCH] Add --images-archive argument to Acto to support loading existing image archives Signed-off-by: Tyler Gu --- acto/__main__.py | 13 +++++++++- acto/engine.py | 23 ++++++++++++++---- acto/post_process/post_diff_test.py | 37 ++++++++++++++++++++--------- chactos/__main__.py | 7 ++++++ chactos/fault_injections.py | 26 +++++++------------- 5 files changed, 72 insertions(+), 34 deletions(-) diff --git a/acto/__main__.py b/acto/__main__.py index d01ddd2fc..7b4d90e3c 100644 --- a/acto/__main__.py +++ b/acto/__main__.py @@ -86,6 +86,12 @@ action="store_true", help="Only generate test cases without executing them", ) +parser.add_argument( + "--images-archive", + dest="images_archive", + type=str, + help="Prebuilt images archive to preload into the cluster", +) parser.add_argument("--checkonly", action="store_true") args = parser.parse_args() @@ -148,6 +154,7 @@ input_model=DeterministicInputModel, apply_testcase_f=apply_testcase_f, focus_fields=config.focus_fields, + images_archive=args.images_archive, ) generation_time = datetime.now() logger.info("Acto initialization finished in %s", generation_time - start_time) @@ -159,7 +166,11 @@ # Post processing post_diff_test_dir = os.path.join(args.workdir_path, "post_diff_test") -p = PostDiffTest(testrun_dir=args.workdir_path, config=config) +p = PostDiffTest( + testrun_dir=args.workdir_path, + config=config, + images_archive=args.images_archive, +) if not args.checkonly: p.post_process(post_diff_test_dir, num_workers=args.num_workers) p.check(post_diff_test_dir, num_workers=args.num_workers) diff --git a/acto/engine.py b/acto/engine.py index e6580d3eb..c534f174e 100644 --- a/acto/engine.py +++ b/acto/engine.py @@ -289,14 +289,20 @@ def __init__( acto_namespace: int, additional_exclude_paths: Optional[list[str]] = None, constraints: Optional[list[XorCondition]] = None, + images_archive: Optional[str] = None, ) -> None: self.context = context self.workdir = workdir self.base_workdir = workdir self.cluster = cluster - self.images_archive = ImageHelper.prepare_image_archive( - self.context["preload_images"], - ) + + if images_archive is None: + self.images_archive = ImageHelper.prepare_image_archive( + self.context["preload_images"], + ) + else: + self.images_archive = images_archive + self.worker_id = worker_id self.sequence_base = sequence_base # trial number to start with self.context_name = cluster.get_context_name( @@ -843,6 +849,7 @@ def __init__( mount: Optional[list] = None, focus_fields: Optional[list] = None, acto_namespace: int = 0, + images_archive: Optional[str] = None, ) -> None: logger = get_thread_logger(with_prefix=False) @@ -888,6 +895,11 @@ def __init__( self.is_reproduce = is_reproduce self.apply_testcase_f = apply_testcase_f self.acto_namespace = acto_namespace + self.images_archive = images_archive + + if images_archive is not None: + # early check if the archive exists + assert os.path.exists(images_archive) self.runner_type = Runner self.checker_type = CheckerSet @@ -1134,7 +1146,9 @@ def run(self) -> list[OracleResults]: logger = get_thread_logger(with_prefix=True) # Build an archive to be preloaded - if len(self.context["preload_images"]) > 0: + if self.images_archive is not None: + print_event("Using existing image archive...") + elif len(self.context["preload_images"]) > 0: logger.info("Creating preload images archive") print_event("Preparing required images...") ImageHelper.prepare_image_archive( @@ -1166,6 +1180,7 @@ def run(self) -> list[OracleResults]: self.acto_namespace, self.operator_config.diff_ignore_fields, constraints=self.operator_config.constraints, + images_archive=self.images_archive, ) runners.append(runner) diff --git a/acto/post_process/post_diff_test.py b/acto/post_process/post_diff_test.py index 6a41b28a7..764e5637b 100644 --- a/acto/post_process/post_diff_test.py +++ b/acto/post_process/post_diff_test.py @@ -6,7 +6,6 @@ import json import logging import multiprocessing -import multiprocessing.queues import os import queue import re @@ -370,6 +369,7 @@ def __init__( cluster: base.KubernetesEngine, worker_id, acto_namespace: int, + images_archive: Optional[str] = None, ): self._context = context self._deploy = deploy @@ -384,9 +384,13 @@ def __init__( os.path.expanduser("~"), ".kube", self._context_name ) self._generation = 0 - self._images_archive = ImageHelper.prepare_image_archive( - self._context["preload_images"], - ) + + if images_archive is None: + self._images_archive = ImageHelper.prepare_image_archive( + self._context["preload_images"], + ) + else: + self._images_archive = images_archive def run_cr(self, cr, trial, gen): """Run a CR and return the snapshot""" @@ -443,6 +447,7 @@ def __init__( cluster: base.KubernetesEngine, worker_id, acto_namespace: int, + images_archive: Optional[str] = None, ): self._workqueue = workqueue self._context = context @@ -457,9 +462,13 @@ def __init__( self._kubeconfig = os.path.join( os.path.expanduser("~"), ".kube", self._context_name ) - self._images_archive = ImageHelper.prepare_image_archive( - self._context["preload_images"], - ) + + if images_archive is None: + self._images_archive = ImageHelper.prepare_image_archive( + self._context["preload_images"], + ) + else: + self._images_archive = images_archive def run(self): """Run the deploy runner""" @@ -636,11 +645,19 @@ def __init__( config: OperatorConfig, ignore_invalid: bool = False, acto_namespace: int = 0, + images_archive: Optional[str] = None, ): self.acto_namespace = acto_namespace super().__init__(testrun_dir, config) logger = get_thread_logger(with_prefix=True) + if images_archive is None: + self.images_archive = ImageHelper.prepare_image_archive( + self.context["preload_images"], + ) + else: + self.images_archive = images_archive + self.all_inputs = [] for trial_name, trial in self.trial_to_steps.items(): for step in trial.steps.values(): @@ -704,10 +721,6 @@ def post_process(self, workdir: str, num_workers: int = 1): version=self.config.kubernetes_version, ) deploy = Deploy(self.config.deploy) - # Build an archive to be preloaded - ImageHelper.prepare_image_archive( - self.context["preload_images"], - ) workqueue: multiprocessing.Queue = multiprocessing.Queue() for unique_input_group in self.unique_inputs.values(): @@ -723,6 +736,7 @@ def post_process(self, workdir: str, num_workers: int = 1): cluster, i, self.acto_namespace, + self.images_archive, ) runners.append(runner) @@ -810,6 +824,7 @@ def check_diff_test_result( cluster=cluster, worker_id=worker_id, acto_namespace=self.acto_namespace, + images_archive=self.images_archive, ) while True: diff --git a/chactos/__main__.py b/chactos/__main__.py index ab97df081..941a8bb42 100644 --- a/chactos/__main__.py +++ b/chactos/__main__.py @@ -49,6 +49,12 @@ default=1, help="Number of concurrent workers to run Chactos with", ) +parser.add_argument( + "--images-archive", + dest="images_archive", + type=str, + help="Prebuilt images archive to preload into the cluster", +) args = parser.parse_args() os.makedirs(args.workdir_path, exist_ok=True) @@ -84,5 +90,6 @@ operator_config=operator_config, fault_injection_config=fi_config, num_workers=args.num_workers, + images_archive=args.images_archive, ) driver.run() diff --git a/chactos/fault_injections.py b/chactos/fault_injections.py index 3359fcccd..e149c70f5 100644 --- a/chactos/fault_injections.py +++ b/chactos/fault_injections.py @@ -1,5 +1,4 @@ import multiprocessing -import multiprocessing.queues import os import queue import subprocess @@ -8,8 +7,9 @@ from typing import Optional, Tuple import kubernetes +from utils.image_helper import ImageHelper -from acto.common import kubernetes_client, print_event +from acto.common import kubernetes_client from acto.deploy import Deploy from acto.kubectl_client.kubectl import KubectlClient from acto.kubernetes_engine import kind @@ -45,6 +45,7 @@ def __init__( operator_config: OperatorConfig, fault_injection_config: FaultInjectionConfig, num_workers: int, + images_archive: Optional[str] = None, ): super().__init__(testrun_dir=testrun_dir, config=operator_config) self._operator_config = operator_config @@ -62,23 +63,12 @@ def __init__( ) # Build an archive to be preloaded - container_tool = os.getenv("IMAGE_TOOL", "docker") - self._images_archive = os.path.join(work_dir, "images.tar") - if len(self.context["preload_images"]) > 0: - print_event("Preparing required images...") - # first make sure images are present locally - for image in self.context["preload_images"]: - subprocess.run( - [container_tool, "pull", image], - stdout=subprocess.DEVNULL, - check=True, - ) - subprocess.run( - [container_tool, "image", "save", "-o", self._images_archive] - + list(self.context["preload_images"]), - stdout=subprocess.DEVNULL, - check=True, + if images_archive is None: + self._images_archive = ImageHelper.prepare_image_archive( + self._context["preload_images"], ) + else: + self._images_archive = images_archive self._deployer = Deploy(operator_config.deploy)