Coverage for /builds/BuildGrid/buildgrid/buildgrid/bot/hardware/worker.py: 100.00%

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. 

14 

15 

16from buildgrid._protos.google.devtools.remoteworkers.v1test2 import worker_pb2 

17 

18 

19class Worker: 

20 """ Describes a worker, which is a list of one of more devices and the 

21 connextions between them. A device could be a computer, a phone or a GPU. 

22 

23 The first device is known as the Primary Device which is responsible for 

24 executing commands. All other devices are known as Attatched Devices and 

25 must be controlled by the Primary Device. 

26 """ 

27 

28 def __init__(self, properties=None, configs=None): 

29 """ Create an instance of a :class:`Worker`. 

30 

31 Property keys supported: `pool`. 

32 Config keys supported: `DockerImage`. 

33 

34 Args: 

35 properties (dict(string : list(string))) : Properties of worker. Keys may have 

36 multiple values. 

37 configs (dict(string : list(string))) : Configurations of a worker. Keys may have 

38 multiple values. 

39 

40 """ 

41 

42 self._devices = [] 

43 self.__configs = {} 

44 self.__properties = {} 

45 self.__property_keys = ['pool'] 

46 self.__config_keys = ['DockerImage'] 

47 

48 if properties: 

49 for k, l in properties.items(): 

50 for v in l: 

51 self._add_property(k, v) 

52 

53 if configs: 

54 for k, l in configs.items(): 

55 for v in l: 

56 self._add_config(k, v) 

57 

58 @property 

59 def configs(self): 

60 """Returns configurations supported by :class:`Worker`.""" 

61 return self.__configs 

62 

63 @property 

64 def properties(self): 

65 """Returns properties supported by :class:`Worker`.""" 

66 return self.__properties 

67 

68 def add_device(self, device): 

69 """ Adds a class:`Device` to this instance. First device added should be the Primary Device.""" 

70 self._devices.append(device) 

71 

72 def get_pb2(self): 

73 """Returns the protobuffer representation of a :class:`Device`.""" 

74 devices = [device.get_pb2() for device in self._devices] 

75 worker = worker_pb2.Worker(devices=devices) 

76 

77 for k, v in self.__properties.items(): 

78 for prop in v: 

79 property_message = worker_pb2.Worker.Property() 

80 property_message.key = k 

81 property_message.value = prop 

82 worker.properties.extend([property_message]) 

83 

84 for k, v in self.__configs.items(): 

85 for cfg in v: 

86 config_message = worker_pb2.Worker.Config() 

87 config_message.key = k 

88 config_message.value = cfg 

89 worker.configs.extend([config_message]) 

90 

91 return worker 

92 

93 def _add_config(self, key, value): 

94 """Adds a configuration. Will raise KeyError if not supported.""" 

95 if key in self.__config_keys: 

96 cfg = self.__configs.get(key) 

97 if not cfg: 

98 self.__configs[key] = [value] 

99 else: 

100 cfg.append(value) 

101 

102 else: 

103 raise KeyError(f'Key not supported: [{key}]') 

104 

105 def _add_property(self, key, value): 

106 """Adds a property. Will raise KeyError if not supported.""" 

107 if key in self.__property_keys: 

108 prop = self.__properties.get(key) 

109 if not prop: 

110 self.__properties[key] = [value] 

111 else: 

112 prop.append(value) 

113 

114 else: 

115 raise KeyError(f'Key not supported: [{key}]')