• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1#!/usr/bin/env python
2# -*- coding: utf-8 -*-
3# Copyright (c) 2021 Huawei Device Co., Ltd.
4# Licensed under the Apache License, Version 2.0 (the "License");
5# you may not use this file except in compliance with the License.
6# You may obtain a copy of the License at
7#
8#     http://www.apache.org/licenses/LICENSE-2.0
9#
10# Unless required by applicable law or agreed to in writing, software
11# distributed under the License is distributed on an "AS IS" BASIS,
12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13# See the License for the specific language governing permissions and
14# limitations under the License.
15
16import shutil
17import os
18import hashlib
19import json
20import http.client as client
21from . import build_utils
22
23
24class Storage():
25    def __init__(self):
26        pass
27
28    @classmethod
29    def retrieve_object(cls, cache_artifact, obj):
30        possible_dir_cache_artifact = '{}.directory'.format(cache_artifact)
31
32        if os.path.exists(cache_artifact):
33            os.makedirs(os.path.dirname(obj), exist_ok=True)
34            shutil.copyfile(cache_artifact, obj)
35            os.utime(cache_artifact)
36            if pycache_debug_enable:
37                print('Retrieve {} from cache'.format(obj))
38        elif os.path.exists(possible_dir_cache_artifact):
39            # Extract zip archive if it's cache artifact for directory.
40            os.makedirs(obj, exist_ok=True)
41            build_utils.extract_all(possible_dir_cache_artifact,
42                                   obj,
43                                   no_clobber=False)
44            os.utime(possible_dir_cache_artifact)
45            if pycache_debug_enable:
46                print('Extract {} from cache'.format(obj))
47        else:
48            if pycache_debug_enable:
49                print('Failed to retrieve {} from cache'.format(obj))
50            return 0
51        return 1
52
53    @classmethod
54    def add_object(cls, cache_artifact, obj):
55        cache_dir = os.path.dirname(cache_artifact)
56        os.makedirs(cache_dir, exist_ok=True)
57
58        if not os.path.exists(obj):
59            return
60        # If path is directory, store an zip archive.
61        if os.path.isdir(obj):
62            dir_cache_artifact = '{}.directory'.format(cache_artifact)
63            build_utils.zip_dir(dir_cache_artifact, obj)
64            if pycache_debug_enable:
65                print("archive {} to {}".format(obj, dir_cache_artifact))
66        else:
67            shutil.copyfile(obj, cache_artifact)
68            if pycache_debug_enable:
69                print("copying {} to {}".format(obj, cache_artifact))
70
71
72class PyCache():
73    def __init__(self, cache_dir=None):
74        cache_dir = os.environ.get('PYCACHE_DIR')
75        if cache_dir:
76            self.pycache_dir = cache_dir
77        else:
78            raise Exception('Error: failed to get PYCACHE_DIR')
79        self.storage = Storage()
80
81    def retrieve(self, output_paths, prefix=''):
82        for path in output_paths:
83            _, cache_artifact = self.descend_directory('{}{}'.format(
84                prefix, path))
85            result = self.storage.retrieve_object(cache_artifact, path)
86            if not result:
87                return result
88
89        try:
90            self.report_cache_stat('cache_hit')
91        except:  # noqa: E722 pylint: disable=bare-except
92            pass
93        return 1
94
95    def save(self, output_paths, prefix=''):
96        for path in output_paths:
97            _, cache_artifact = self.descend_directory('{}{}'.format(
98                prefix, path))
99            self.storage.add_object(cache_artifact, path)
100
101    def report_cache_stat(self, hit_or_miss):
102        pyd_server, pyd_port = self.get_pyd()
103        conn = client.HTTPConnection(pyd_server, pyd_port)
104        conn.request(hit_or_miss, '/')
105        conn.close()
106
107    def get_pyd(self):
108        daemon_config_file = '{}/.config'.format(self.pycache_dir)
109        if not os.path.exists(daemon_config_file):
110            raise Exception('Warning: no pycache daemon process exists.')
111        with open(daemon_config_file, 'r') as jsonfile:
112            data = json.load(jsonfile)
113            return data.get('host'), data.get('port')
114
115    @classmethod
116    def cache_key(cls, path):
117        sha256 = hashlib.sha256()
118        sha256.update(path.encode())
119        return sha256.hexdigest()
120
121    def descend_directory(self, path):
122        digest = self.cache_key(path)
123        cache_dir = os.path.join(self.pycache_dir, digest[:2])
124        return cache_dir, os.path.join(cache_dir, digest[2:])
125
126    # Manifest file to record inputs/outputs/commands.
127    def get_manifest_path(self, path):
128        manifest_dir, manifest_file = self.descend_directory(path)
129        os.makedirs(manifest_dir, exist_ok=True)
130        return manifest_file
131
132
133pycache_enabled = (os.environ.get('PYCACHE_DIR') is not None)
134pycache_debug_enable = int(os.environ.get('PRINT_BUILD_EXPLANATIONS', 0))
135if pycache_enabled:
136    pycache = PyCache()
137else:
138    pycache = None  # pylint: disable=invalid-name
139