Coverage for /builds/BuildGrid/buildgrid/buildgrid/server/capabilities/instance.py: 97.44%

78 statements  

« prev     ^ index     » next       coverage.py v7.4.1, created at 2024-06-11 15:37 +0000

1# Copyright (C) 2018 Bloomberg LP 

2# 

3# Licensed under the Apache License, Version 2.0 (the "License"); 

4# you may not use this file except in compliance with the License. 

5# You may obtain a copy of the License at 

6# 

7# <http://www.apache.org/licenses/LICENSE-2.0> 

8# 

9# Unless required by applicable law or agreed to in writing, software 

10# distributed under the License is distributed on an "AS IS" BASIS, 

11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 

12# See the License for the specific language governing permissions and 

13# limitations under the License. 

14from typing import Optional, Union 

15 

16from buildgrid._exceptions import RetriableError 

17from buildgrid._protos.build.bazel.remote.execution.v2.remote_execution_pb2 import DESCRIPTOR as RE_DESCRIPTOR 

18from buildgrid._protos.build.bazel.remote.execution.v2.remote_execution_pb2 import ( 

19 ActionCacheUpdateCapabilities, 

20 CacheCapabilities, 

21 ExecutionCapabilities, 

22 ServerCapabilities, 

23) 

24from buildgrid._protos.build.bazel.semver.semver_pb2 import SemVer 

25from buildgrid.server.actioncache.caches.action_cache_abc import ActionCacheABC 

26from buildgrid.server.actioncache.instance import ActionCache 

27from buildgrid.server.cas.instance import ContentAddressableStorageInstance 

28from buildgrid.server.execution.instance import ExecutionInstance 

29from buildgrid.server.metrics_names import CAPABILITIES_GET_CAPABILITIES_EXCEPTION_COUNT_METRIC_NAME 

30from buildgrid.server.metrics_utils import ExceptionCounter 

31from buildgrid.server.servicer import Instance 

32from buildgrid.settings import HIGH_REAPI_VERSION, LOW_REAPI_VERSION 

33 

34ActionCacheInstance = Union[ActionCache, ActionCacheABC] 

35 

36 

37class CapabilitiesInstance(Instance): 

38 SERVICE_NAME = RE_DESCRIPTOR.services_by_name["Capabilities"].full_name 

39 

40 def __init__( 

41 self, 

42 cas_instance: Optional["ContentAddressableStorageInstance"] = None, 

43 action_cache_instance: Optional["ActionCacheInstance"] = None, 

44 execution_instance: Optional["ExecutionInstance"] = None, 

45 ) -> None: 

46 self.__cas_instance = cas_instance 

47 self.__action_cache_instance = action_cache_instance 

48 self.__execution_instance = execution_instance 

49 

50 self.__high_api_version: Optional[SemVer] = None 

51 self.__low_api_version: Optional[SemVer] = None 

52 

53 def add_cas_instance(self, cas_instance: "ContentAddressableStorageInstance") -> None: 

54 self.__cas_instance = cas_instance 

55 

56 def add_action_cache_instance(self, action_cache_instance: "ActionCacheInstance") -> None: 

57 self.__action_cache_instance = action_cache_instance 

58 

59 def add_execution_instance(self, execution_instance: "ExecutionInstance") -> None: 

60 self.__execution_instance = execution_instance 

61 

62 get_capabilities_ignored_exceptions = (RetriableError,) 

63 

64 @ExceptionCounter( 

65 CAPABILITIES_GET_CAPABILITIES_EXCEPTION_COUNT_METRIC_NAME, 

66 ignored_exceptions=get_capabilities_ignored_exceptions, 

67 ) 

68 def get_capabilities(self) -> ServerCapabilities: 

69 cache_capabilities = self._get_cache_capabilities() 

70 execution_capabilities = self._get_capabilities_execution() 

71 

72 if self.__high_api_version is None: 

73 self.__high_api_version = self._split_semantic_version(HIGH_REAPI_VERSION) 

74 if self.__low_api_version is None: 

75 self.__low_api_version = self._split_semantic_version(LOW_REAPI_VERSION) 

76 

77 server_capabilities = ServerCapabilities() 

78 server_capabilities.cache_capabilities.CopyFrom(cache_capabilities) 

79 server_capabilities.execution_capabilities.CopyFrom(execution_capabilities) 

80 server_capabilities.low_api_version.CopyFrom(self.__low_api_version) 

81 server_capabilities.high_api_version.CopyFrom(self.__high_api_version) 

82 

83 return server_capabilities 

84 

85 # --- Private API --- 

86 

87 def _get_cache_capabilities(self) -> CacheCapabilities: 

88 capabilities = CacheCapabilities() 

89 action_cache_update_capabilities = ActionCacheUpdateCapabilities() 

90 

91 if self.__cas_instance: 

92 capabilities.digest_functions.extend([self.__cas_instance.hash_type()]) 

93 capabilities.max_batch_total_size_bytes = self.__cas_instance.max_batch_total_size_bytes() 

94 capabilities.symlink_absolute_path_strategy = self.__cas_instance.symlink_absolute_path_strategy() 

95 # TODO: execution priority #102 

96 # capabilities.cache_priority_capabilities = 

97 

98 if self.__action_cache_instance: 

99 capabilities.digest_functions.extend([self.__action_cache_instance.hash_type()]) 

100 action_cache_update_capabilities.update_enabled = self.__action_cache_instance.allow_updates 

101 

102 # If an endpoint doesn't have a CAS instance, but does have an Execution 

103 # instance, it needs to report the capabilities of the CAS used by the 

104 # Execution instance. See https://gitlab.com/BuildGrid/buildgrid/-/issues/174 

105 # for more context and discussion of this. 

106 # 

107 # There's also a convenient side-effect to this, it means separated BuildGrid 

108 # services can all be deployed behind a single proxy endpoint without needing 

109 # extra work to make GetCapabilities requests make sense. 

110 if self.__execution_instance and not self.__cas_instance: 

111 remote_cache_capabilities = self.__execution_instance.get_storage_capabilities() 

112 capabilities.CopyFrom(remote_cache_capabilities) 

113 

114 # Similarly, if an endpoint has an Execution instance but no ActionCache 

115 # instance, it needs to report the capabilities of the ActionCache used 

116 # by the Execution instance. 

117 if self.__execution_instance and not self.__action_cache_instance: 

118 ac_hash_type, remote_action_cache_capabilities = self.__execution_instance.get_action_cache_capabilities() 

119 if not capabilities.digest_functions and ac_hash_type is not None: 

120 capabilities.digest_functions.extend([ac_hash_type]) 

121 if remote_action_cache_capabilities is not None: 

122 action_cache_update_capabilities.CopyFrom(remote_action_cache_capabilities) 

123 

124 capabilities.action_cache_update_capabilities.CopyFrom(action_cache_update_capabilities) 

125 return capabilities 

126 

127 def _get_capabilities_execution(self) -> ExecutionCapabilities: 

128 capabilities = ExecutionCapabilities() 

129 if self.__execution_instance: 

130 capabilities.exec_enabled = True 

131 capabilities.digest_function = self.__execution_instance.hash_type() 

132 # TODO: execution priority #102 

133 # capabilities.execution_priority = 

134 

135 else: 

136 capabilities.exec_enabled = False 

137 

138 return capabilities 

139 

140 def _split_semantic_version(self, version_string: str) -> SemVer: 

141 major_version, minor_version, patch_version = version_string.split(".") 

142 

143 semantic_version = SemVer() 

144 semantic_version.major = int(major_version) 

145 semantic_version.minor = int(minor_version) 

146 semantic_version.patch = int(patch_version) 

147 

148 return semantic_version