MegEngine/imperative/python/megengine/utils/persistent_cache.py

# -*- coding: utf-8 -*-
# MegEngine is Licensed under the Apache License, Version 2.0 (the "License")
#
# Copyright (c) 2014-2021 Megvii Inc. All rights reserved.
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

import argparse
import getpass
import os
import sys

from ..core._imperative_rt import PersistentCacheManager as _PersistentCacheManager
from ..logger import get_logger
from ..version import __version__, git_version


class PersistentCacheManager(_PersistentCacheManager):
    def __init__(self):
        super().__init__()
        if os.getenv("MGE_FASTRUN_CACHE_TYPE") == "MEMORY":
            get_logger().info("fastrun use in-memory cache")
            self.open_memory()
        else:
            self.open_file()

    def open_memory(self):
        pass

    def open_file(self):
        cache_dir = os.getenv("MGE_FASTRUN_CACHE_DIR")
        try:
            if not cache_dir:
                from ..hub.hub import _get_megengine_home

                cache_dir = os.path.expanduser(
                    os.path.join(_get_megengine_home(), "persistent_cache.bin")
                )
            os.makedirs(cache_dir, exist_ok=True)
            cache_file = os.path.join(cache_dir, "cache")
            with open(cache_file, "a"):
                pass
            assert self.try_open_file(cache_file), "cannot create file"
            get_logger().info("fastrun use in-file cache in {}".format(cache_dir))
        except Exception as exc:
            get_logger().error(
                "failed to create cache file in {} {!r}; fallback to "
                "in-memory cache".format(cache_dir, exc)
            )
            self.open_memory()


_manager = None


def get_manager():
    global _manager
    if _manager is None:
        _manager = PersistentCacheManager()
    return _manager