prefetch_benchmark.py

import logging
import os
import subprocess
from typing import Dict, Any

from benchmarks.base_benchmark import BaseBenchmark
from omegaconf import DictConfig

from benchmarks.benchmark_config_parser import BenchmarkConfigParser

log = logging.getLogger(__name__)


class PrefetchBenchmark(BaseBenchmark):
    def __init__(self, cfg: DictConfig, metadata: Dict[str, Any]):
        self.config_parser = BenchmarkConfigParser(cfg)
        self.common_config = self.config_parser.get_common_config()
        self.prefetch_config = self.config_parser.get_prefetch_config()

    def setup(self) -> None:
        pass

    def run_benchmark(self) -> None:
        subprocess_args = [
            "cargo",
            "run",
            "--release",
            "--example",
            "prefetch_benchmark",
            self.common_config['s3_bucket'],
        ]

        object_size = self.common_config['object_size_in_gib']
        size_gib = str(object_size)
        app_workers = self.common_config['application_workers']

        # Check if objects are specified or if we have to fall back to objects
        # generated by fio.
        objects = self.common_config['s3_keys']
        if not objects:
            objects = self.config_parser.default_object_keys(app_workers, size_gib)

        if len(objects) >= app_workers:
            for i in range(app_workers):
                subprocess_args.append(objects[i])
        else:
            raise ValueError("Seeing fewer objects than app workers. So cannot proceed with the run.")

        region = self.common_config["region"]
        subprocess_args.extend(["--region", region])

        if (max_throughput := self.common_config['max_throughput_gbps']) is not None:
            subprocess_args.extend(["--maximum-throughput-gbps", str(max_throughput)])

        if (max_memory_target := self.prefetch_config['max_memory_target']) is not None:
            subprocess_args.extend(["--max-memory-target", str(max_memory_target)])

        if (read_part_size := self.common_config['read_part_size']) is not None:
            subprocess_args.extend(["--part-size", str(read_part_size)])

        read_size = self.common_config['read_size']
        subprocess_args.extend(["--read-size", str(read_size)])

        for interface in self.common_config['network_interfaces']:
            subprocess_args.extend(["--bind", interface])

        if (run_time := self.common_config['run_time']) is not None:
            subprocess_args.extend(["--max-duration", str(run_time)])

        subprocess_args.extend(["--output-file", "prefetch-output.json"])
        prefetch_env = {}
        if not self.common_config['download_checksums']:
            prefetch_env["EXPERIMENTAL_MOUNTPOINT_NO_DOWNLOAD_INTEGRITY_VALIDATION"] = "ON"

        if self.prefetch_config['max_read_window_size'] is not None:
            prefetch_env["UNSTABLE_MOUNTPOINT_MAX_PREFETCH_WINDOW_SIZE"] = str(self.prefetch_config['max_read_window_size'])

        subprocess_env = os.environ.copy() | prefetch_env
        log.debug("Subprocess env: %s", subprocess_env)

        log.info("Running prefetch benchmark with args: %s", subprocess_args)
        subprocess.run(subprocess_args, check=True, capture_output=True, text=True, env=subprocess_env)
        log.info("Prefetch benchmark completed successfully.")

    def post_process(self) -> None:
        pass