Coverage for /builds/BuildGrid/buildgrid/buildgrid/bot/hardware/worker.py: 100.00%
52 statements
« prev ^ index » next coverage.py v6.4.1, created at 2022-06-22 21:04 +0000
« prev ^ index » next coverage.py v6.4.1, created at 2022-06-22 21:04 +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.
16from buildgrid._protos.google.devtools.remoteworkers.v1test2 import worker_pb2
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.
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 """
28 def __init__(self, properties=None, configs=None):
29 """ Create an instance of a :class:`Worker`.
31 Property keys supported: `pool`.
32 Config keys supported: `DockerImage`.
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.
40 """
42 self._devices = []
43 self.__configs = {}
44 self.__properties = {}
45 self.__property_keys = ['pool']
46 self.__config_keys = ['DockerImage']
48 if properties:
49 for k, l in properties.items():
50 for v in l:
51 self._add_property(k, v)
53 if configs:
54 for k, l in configs.items():
55 for v in l:
56 self._add_config(k, v)
58 @property
59 def configs(self):
60 """Returns configurations supported by :class:`Worker`."""
61 return self.__configs
63 @property
64 def properties(self):
65 """Returns properties supported by :class:`Worker`."""
66 return self.__properties
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)
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)
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])
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])
91 return worker
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)
102 else:
103 raise KeyError(f'Key not supported: [{key}]')
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)
114 else:
115 raise KeyError(f'Key not supported: [{key}]')