Introduce --reuse-disk-images Polybench flag to support reusing images found on disk

Author: DSouzaM

truffle/mx.truffle/mx_polybench/command.py

@@ -43,12 +43,18 @@
 import shlex
 from argparse import ArgumentParser
 from enum import Enum
-from typing import List, Set, Tuple, NamedTuple
+from typing import List, Set, Tuple, NamedTuple, Dict
 
 import mx
 import mx_benchmark
 import mx_sdk
-from mx_polybench.model import _resolve_all_benchmarks, _get_all_suites, PolybenchBenchmarkSuiteEntry
+from mx_polybench.model import (
+    _resolve_all_benchmarks,
+    _get_all_suites,
+    PolybenchBenchmarkSuiteEntry,
+    _extend_env,
+    PolybenchBenchmarkSuite,
+)
 
 _suite = mx.suite("truffle")
 
@@ -200,6 +206,8 @@ def _run_suite(args):
         base_args.append("--dry-run")
     elif args.dry_run_polybench:
         base_args.append("--dry-run-polybench")
+    elif args.reuse_disk_images:
+        base_args.append("--reuse-disk-images")
 
     def polybench_run_function(argument_list: List[str]) -> None:
         raw_args = base_args + argument_list + override_arguments
@@ -258,7 +266,9 @@ def _run_benchmark_pattern(args):
     run_spec = PolybenchRunSpecification(args.benchmarks, _get_vm_features(args), arguments_spec)
     _validate_jdk(run_spec.is_native())
     mx.logv(f"Performing polybench run: {run_spec}")
-    _run_specification(run_spec, pattern_is_glob=args.pattern_is_glob, dry_run=args.dry_run)
+    _run_specification(
+        run_spec, pattern_is_glob=args.pattern_is_glob, dry_run=args.dry_run, reuse_disk_images=args.reuse_disk_images
+    )
 
 
 def _validate_jdk(is_native: bool) -> mx.JDKConfig:
@@ -341,11 +351,6 @@ class PolybenchRunSpecification(NamedTuple):
     vm_features: Set[VMFeature] = set()
     arguments: PolybenchArgumentsSpecification = PolybenchArgumentsSpecification()
 
-    def append_arguments(self, other: PolybenchArgumentsSpecification) -> "PolybenchRunSpecification":
-        return PolybenchRunSpecification(
-            pattern=self.pattern, vm_features=self.vm_features, arguments=self.arguments.append(other)
-        )
-
     def is_native(self) -> bool:
         return VMFeature.NATIVE in self.vm_features
 
@@ -361,7 +366,12 @@ def jvm_config(self) -> str:
         return "-".join(features) if features else "default"
 
 
-def _run_specification(spec: PolybenchRunSpecification, pattern_is_glob: bool = True, dry_run: bool = False):
+def _run_specification(
+    spec: PolybenchRunSpecification,
+    pattern_is_glob: bool = True,
+    dry_run: bool = False,
+    reuse_disk_images: bool = False,
+):
     pattern = _parse_mx_benchmark_pattern(spec.pattern, pattern_is_glob)
     mx_benchmark_args = (
         [f"polybench:{pattern}"]
@@ -377,7 +387,15 @@ def _run_specification(spec: PolybenchRunSpecification, pattern_is_glob: bool =
         return
 
     mx.logv(f"Running command: {command_string}")
-    mx_benchmark.benchmark(mx_benchmark_args)
+    with _extend_env(_extra_run_variables(reuse_disk_images)):
+        mx_benchmark.benchmark(mx_benchmark_args)
+
+
+def _extra_run_variables(reuse_disk_images) -> Dict[str, str]:
+    result = {}
+    if reuse_disk_images:
+        result[PolybenchBenchmarkSuite.REUSE_DISK_IMAGES] = "true"
+    return result
 
 
 def _base_mx_command() -> List[str]:
@@ -462,6 +480,16 @@ def run_flag(run_arg):
     parser.add_argument(
         run_flag("--g1gc"), action="store_true", default=False, help="use G1GC (only valid for native runs)"
     )
+    parser.add_argument(
+        run_flag("--reuse-disk-images"),
+        action="store_true",
+        default=False,
+        help=(
+            "reuse existing native images found on disk. Polybench will reuse an image from disk if it was built with the same languages and VM arguments. "
+            "This feature does not detect stale images, does not support PGO, and should only be used for development. "
+            f'This feature can also be enabled by setting the environment variable {PolybenchBenchmarkSuite.REUSE_DISK_IMAGES} to "true".'
+        ),
+    )
     benchmark_pattern_group = parser.add_mutually_exclusive_group()
     benchmark_pattern_group.add_argument(
         run_flag("--suite"),

truffle/mx.truffle/mx_polybench/model.py

@@ -284,6 +284,7 @@ class PolybenchBenchmarkSuite(
         "application-memory-metric": "application-memory",
         "none": None,
     }
+    REUSE_DISK_IMAGES = "POLYBENCH_REUSE_DISK_IMAGES"
 
     def __init__(self):
         super(PolybenchBenchmarkSuite, self).__init__()
@@ -333,14 +334,14 @@ def run(self, benchmarks, bmSuiteArgs) -> DataPoints:
         working_directory = self.workingDirectory(benchmarks, bmSuiteArgs) or os.getcwd()
         resolved_benchmark = self._resolve_current_benchmark(benchmarks)
 
-        mx.log(f'Running polybench benchmark "{resolved_benchmark.name}"".')
+        mx.log(f'Running polybench benchmark "{resolved_benchmark.name}".')
         mx.logv(f"CWD: {working_directory}")
         mx.logv(f"Languages included on the classpath: {resolved_benchmark.suite.languages}")
 
         env_vars = PolybenchBenchmarkSuite._prepare_distributions(working_directory, resolved_benchmark)
-        with _extend_env(env_vars):
+        with _extend_env(env_vars), self._set_image_context(resolved_benchmark, bmSuiteArgs):
             if self._can_use_image_cache(bmSuiteArgs):
-                return self._run_with_image_cache(resolved_benchmark, benchmarks, bmSuiteArgs)
+                return self._run_with_image_cache(benchmarks, bmSuiteArgs)
             else:
                 return self.intercept_run(super(), benchmarks, bmSuiteArgs)
 
@@ -380,13 +381,10 @@ def _prepare_distributions(
     def _can_use_image_cache(self, bm_suite_args) -> bool:
         return self.is_native_mode(bm_suite_args) and "pgo" not in self.jvmConfig(bm_suite_args)
 
-    def _run_with_image_cache(
-        self, resolved_benchmark: ResolvedPolybenchBenchmark, benchmarks: List[str], bm_suite_args: List[str]
-    ) -> DataPoints:
-        with self._set_image_context(resolved_benchmark, bm_suite_args):
-            image_build_datapoints = self._build_cached_image(benchmarks, bm_suite_args)
-            image_run_datapoints = self._run_cached_image(benchmarks, bm_suite_args)
-            return list(image_build_datapoints) + list(image_run_datapoints)
+    def _run_with_image_cache(self, benchmarks: List[str], bm_suite_args: List[str]) -> DataPoints:
+        image_build_datapoints = self._build_cached_image(benchmarks, bm_suite_args)
+        image_run_datapoints = self._run_cached_image(benchmarks, bm_suite_args)
+        return list(image_build_datapoints) + list(image_run_datapoints)
 
     @contextlib.contextmanager
     def _set_image_context(self, resolved_benchmark: ResolvedPolybenchBenchmark, bm_suite_args: List[str]):
@@ -404,12 +402,11 @@ def _set_image_context(self, resolved_benchmark: ResolvedPolybenchBenchmark, bm_
 
     def _base_image_name(self) -> Optional[str]:
         """Overrides the image name used to build/run the image."""
-        if self._current_image:
-            return self._current_image.full_executable_name()
-        return None
+        assert self._current_image, "Image should have been set already"
+        return self._current_image.full_executable_name()
 
     def _build_cached_image(self, benchmarks: List[str], bm_suite_args: List[str]) -> DataPoints:
-        if self._current_image in self._image_cache:
+        if self._image_is_cached(bm_suite_args):
             # already built
             return []
 
@@ -422,14 +419,33 @@ def _build_cached_image(self, benchmarks: List[str], bm_suite_args: List[str]) -
             datapoint["benchmark"] = self._current_image.executable_name()
         return image_build_datapoints
 
+    def _image_is_cached(self, bm_suite_args: List[str]) -> bool:
+        if self._current_image in self._image_cache:
+            return True
+
+        if mx.get_env(PolybenchBenchmarkSuite.REUSE_DISK_IMAGES) in ["true", "True"]:
+            full_image_name = self.get_full_image_name(self.get_base_image_name(), self.jvmConfig(bm_suite_args))
+            image_path = self.get_image_output_dir(
+                self.benchmark_output_dir(self.benchmark_name, self.vmArgs(bm_suite_args)), full_image_name
+            ) / self.get_image_file_name(full_image_name)
+            if os.path.exists(image_path):
+                mx.warn(
+                    f"Existing image at {image_path} will be reused ({PolybenchBenchmarkSuite.REUSE_DISK_IMAGES} is set to true). "
+                    "Reusing disk images is a development feature and it does not detect stale images. Use with caution."
+                )
+                self._image_cache.add(self._current_image)
+                return True
+
+        return False
+
     def _run_cached_image(self, benchmarks: List[str], bm_suite_args: List[str]) -> DataPoints:
         return self.intercept_run(
             super(), benchmarks, self._extend_vm_args(bm_suite_args, ["-Dnative-image.benchmark.stages=run"])
         )
 
     def _extend_vm_args(self, bm_suite_args: List[str], new_vm_args: List[str]) -> List[str]:
-        vmArgs, runArgs = self.vmAndRunArgs(bm_suite_args)
-        return vmArgs + new_vm_args + ["--"] + runArgs
+        vm_args, run_args = self.vmAndRunArgs(bm_suite_args)
+        return vm_args + new_vm_args + ["--"] + run_args
 
     def createCommandLineArgs(self, benchmarks, bmSuiteArgs):
         resolved_benchmark = self._resolve_current_benchmark(benchmarks)
@@ -596,6 +612,7 @@ def parse(self, text) -> Iterable[DataPoint]:
 
 @contextlib.contextmanager
 def _extend_env(env_vars: Dict[str, str]):
+    """Temporarily extends the environment variables for the extent of this context."""
     old_env = dict(os.environ)
     try:
         for k, v in env_vars.items():