benchmarking/benchmarks/benchmarks.py

#!/usr/bin/env python

##############################################################################
# Copyright 2017-present, Facebook, Inc.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
##############################################################################

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals

import copy
import gc
import hashlib
import json
import os
import shutil
import tempfile

import requests
from utils.custom_logger import getLogger
from utils.utilities import deepMerge, deepReplace


COPY_THRESHOLD = 6442450944  # 6 GB


class BenchmarkCollector(object):
    def __init__(self, framework, model_cache, **kwargs):

        self.args = kwargs.get("args", None)
        if not os.path.isdir(model_cache):
            os.makedirs(model_cache)
        self.model_cache = model_cache
        self.framework = framework

    def collectBenchmarks(self, info, source, user_identifier):
        assert os.path.isfile(source), "Source {} is not a file".format(source)
        with open(source, "r") as f:
            content = json.load(f)

        meta = content["meta"] if "meta" in content else {}
        if "meta" in info:
            deepMerge(meta, info["meta"])
        if hasattr(self.args, "timeout"):
            meta["timeout"] = self.args.timeout
        benchmarks = []

        if "benchmarks" in content:
            path = os.path.abspath(os.path.dirname(source))
            assert "meta" in content, "Meta field is missing in benchmarks"
            for benchmark_file in content["benchmarks"]:
                benchmark_file = os.path.join(path, benchmark_file)
                self._collectOneBenchmark(
                    benchmark_file, meta, benchmarks, info, user_identifier
                )
        else:
            self._collectOneBenchmark(source, meta, benchmarks, info, user_identifier)

        for b in benchmarks:
            self._verifyBenchmark(b, b["path"], True)
        return benchmarks

    def _verifyBenchmark(self, benchmark, filename, is_post):
        self.framework.verifyBenchmarkFile(benchmark, filename, is_post)

    def _collectOneBenchmark(self, source, meta, benchmarks, info, user_identifier):
        assert os.path.isfile(source), "Benchmark {} does not exist".format(source)
        with open(source, "r") as b:
            one_benchmark = json.load(b)

        string_map = json.loads(self.args.string_map) if self.args.string_map else {}
        for name in string_map:
            value = string_map[name]
            deepReplace(one_benchmark, "{" + name + "}", value)

        self._verifyBenchmark(one_benchmark, source, False)

        self._updateFiles(one_benchmark, source, user_identifier)

        # following change should not appear in updated_json file
        if meta:
            deepMerge(one_benchmark["model"], meta)

        self._updateTests(one_benchmark, source)
        # Add fields that should not appear in the saved benchmark file
        # Adding path to benchmark file
        one_benchmark["path"] = os.path.abspath(source)

        # One test per benchmark
        if len(one_benchmark["tests"]) == 1:
            benchmarks.append(one_benchmark)
        else:
            tests = copy.deepcopy(one_benchmark["tests"])
            one_benchmark["tests"] = []
            for test in tests:
                new_benchmark = copy.deepcopy(one_benchmark)
                new_benchmark["tests"].append(test)
                benchmarks.append(new_benchmark)

    # Update all files in the benchmark to absolute path
    # download the files if needed
    def _updateFiles(self, one_benchmark, filename, user_identifier):

        model = one_benchmark["model"]
        model_dir = os.path.join(self.model_cache, model["format"], model["name"])
        if not os.path.isdir(model_dir):
            os.makedirs(model_dir)
        collected_files, collected_tmp_files = self._collectFiles(one_benchmark)
        update_json = False
        for file in collected_files:
            update_json |= self._updateOneFile(file, model_dir, filename)

        if update_json:
            s = json.dumps(one_benchmark, indent=2, sort_keys=True)
            with open(filename, "w") as f:
                f.write(s)
            getLogger().info(
                "Model {} is changed. ".format(model["name"])
                + "Please update the meta json file."
            )

        # update the file field with the absolute path
        # needs to be after the file is updated
        # only update files with md5, which means those files are not
        # temporary files
        for file in collected_files:
            if "md5" in file:
                cached_filename = self._getDestFilename(file, model_dir)
                file["location"] = cached_filename

        tmp_dir = tempfile.mkdtemp(
            prefix="_".join(["aibench", str(user_identifier), ""])
        )
        for tmp_file in collected_tmp_files:
            tmp_file["location"] = tmp_file["location"].replace("{TEMPDIR}", tmp_dir)

    def _collectFiles(self, benchmark):
        files = []
        tmp_files = []
        if "model" in benchmark:
            if "files" in benchmark["model"]:
                self._collectOneGroupFiles(benchmark["model"]["files"], files)
            if "libraries" in benchmark["model"]:
                self._collectOneGroupFiles(benchmark["model"]["libraries"], files)

        for test in benchmark["tests"]:
            if "files" in test:
                self._collectOneGroupFiles(test["files"], files, tmp_files)
            if "input_files" in test:
                self._collectOneGroupFiles(test["input_files"], files)
            if "output_files" in test:
                self._collectOneGroupFiles(test["output_files"], files, tmp_files)
            if "preprocess" in test and "files" in test["preprocess"]:
                self._collectOneGroupFiles(
                    test["preprocess"]["files"], files, tmp_files
                )
            if "postprocess" in test and "files" in test["postprocess"]:
                self._collectOneGroupFiles(
                    test["postprocess"]["files"], files, tmp_files
                )
        return files, tmp_files

    def _collectOneGroupFiles(self, group, files, tmp_files=None):
        if isinstance(group, list):
            for f in group:
                self._collectOneFile(f, files, tmp_files)
        elif isinstance(group, dict):
            for name in group:
                f_or_list = group[name]
                if isinstance(f_or_list, list):
                    for f in f_or_list:
                        self._collectOneFile(f, files, tmp_files)
                else:
                    self._collectOneFile(f_or_list, files, tmp_files)

    def _collectOneFile(self, item, files, tmp_files):
        if "location" in item and "{TEMPDIR}" in item["location"]:
            assert tmp_files is not None, "tmp file can only exist for output"
            tmp_files.append(item)
            return
        assert "filename" in item, "field filename must exist"
        if "location" not in item:
            return
        files.append(item)

    def _updateOneFile(self, field, model_dir, filename):
        cached_filename = self._getDestFilename(field, model_dir)
        if (
            "md5" in field
            and field["md5"] is not None
            and (
                not os.path.isfile(cached_filename)
                or self._calculateMD5(cached_filename, field["md5"], filename)
                != field["md5"]
            )
        ):
            return self._copyFile(field, cached_filename, filename)
        return False

    def _calculateMD5(self, model_name, old_md5, filename):
        if os.stat(filename).st_size >= COPY_THRESHOLD or os.path.islink(model_name):
            if not os.path.isfile(model_name):
                getLogger().info(
                    "Create symlink between {} and {}".format(filename, model_name)
                )
                os.symlink(filename, model_name)
            return old_md5
        getLogger().info("Calculate md5 of {}".format(model_name))
        with open(model_name, "rb") as f:
            file_hash = hashlib.md5()
            for chunk in iter(lambda: f.read(8192), b""):
                file_hash.update(chunk)
        md5 = file_hash.hexdigest()
        del file_hash
        gc.collect()
        return md5

    def _copyFile(self, field, destination_name, source):
        if "location" not in field:
            return False
        location = field["location"]
        if location[0:4] == "http":
            abs_name = destination_name
            getLogger().info("Downloading {}".format(location))
            r = requests.get(location)
            if r.status_code == 200:
                with open(destination_name, "wb") as f:
                    f.write(r.content)
        else:
            abs_name = self._getAbsFilename(field, source, None)
            if os.path.isfile(abs_name):
                if os.stat(abs_name).st_size < COPY_THRESHOLD:
                    shutil.copyfile(abs_name, destination_name)
                else:
                    if not os.path.isfile(destination_name):
                        getLogger().info(
                            "Create symlink between {} and {}".format(
                                abs_name, destination_name
                            )
                        )
                        os.symlink(abs_name, destination_name)
            else:
                import distutils.dir_util

                distutils.dir_util.copy_tree(abs_name, destination_name)
        if os.path.isdir(destination_name) and field["md5"] == "directory":
            return False
        assert os.path.isfile(destination_name), "File {} cannot be retrieved".format(
            destination_name
        )
        # verify the md5 matches the file downloaded
        md5 = self._calculateMD5(destination_name, field["md5"], abs_name)
        if md5 != field["md5"]:
            getLogger().info(
                "Source file {} is changed, ".format(location)
                + " updating MD5. "
                + "Please commit the updated json file."
            )
            field["md5"] = md5
            return True
        return False

    def _getDestFilename(self, field, dir):
        fn = os.path.splitext(field["filename"])
        cached_name = os.path.join(dir, fn[0] + fn[1])
        return cached_name

    def _updateTests(self, one_benchmark, source):
        if one_benchmark["tests"][0]["metric"] == "generic":
            return

        # framework specific updates
        self.framework.rewriteBenchmarkTests(one_benchmark, source)

        # rewrite test fields for compatibility reasons
        for test in one_benchmark["tests"]:
            self._rewriteTestFields(test)

        # Update identifiers, the last update
        self._updateNewTestFields(one_benchmark["tests"], one_benchmark)

    def _rewriteTestFields(self, test):
        if "arguments" in test:
            assert (
                "commands" not in test
            ), "Commands and arguments cannot co-exist in test"
            test["commands"] = ["{program} " + test["arguments"]]
            del test["arguments"]
        if "command" in test:
            assert (
                "commands" not in test
            ), "Commands and command cannot co-exist in test"
            test["commands"] = [test["command"]]
            # do not delete for now
            # del test["command"]

    def _updateNewTestFields(self, tests, one_benchmark):
        idx = 0
        for test in tests:
            identifier = test["identifier"].replace("{ID}", str(idx))
            test["identifier"] = identifier
            idx += 1

    def _getAbsFilename(self, file, source, cache_dir):
        location = file["location"]
        filename = file["filename"]
        if location[0:4] == "http":
            # Need to download, return the destination filename
            return os.path.join(cache_dir, filename)
        elif location[0:2] == "//":
            assert self.args.root_model_dir is not None, (
                "When specifying relative directory, the "
                "--root_model_dir must be specified."
            )
            return self.args.root_model_dir + location[1:]
        elif location[0] != "/":
            abs_dir = os.path.dirname(os.path.abspath(source))
            return os.path.join(abs_dir, location)
        else:
            return location