[Notes] [Git][BuildGrid/buildgrid][finn/48-cancellation-leases] 10 commits: Added HardwareInterface class.

[finn <gitlab mg gitlab com>]

Title: GitLab

finn pushed to branch finn/48-cancellation-leases at BuildGrid / buildgrid

Commits:

• 6eb61ebc
  by Finn at 2018-11-16T11:35:07Z

  Added HardwareInterface class.
  
  This will describe and configure hardware.
  

• 29192175
  by Finn at 2018-11-16T11:35:09Z

  Add Device class.
  

• 8ff0bc50
  by Finn at 2018-11-16T11:35:09Z

  Add Worker class.
  

• 5b7a0241
  by Finn at 2018-11-16T11:35:09Z

  Remove hardware descriptions from bot_session.
  
  Also move `bot_session` to simply `session`.
  

• a28453e8
  by Finn at 2018-11-16T11:35:09Z

  Added Tenant class.
  
  This will contain the information about a running task
  and lease.
  

• eefaee81
  by Finn at 2018-11-16T11:35:09Z

  Added TenantManager class.
  
  This will manage the Tenant objects.
  

• 0451031f
  by Finn at 2018-11-16T11:35:09Z

  Moved bots_interface.py to interface.py.
  

• a118e3a5
  by Finn at 2018-11-16T11:35:09Z

  Modified BotSession to use new TenantManager class.
  

• 1a05afa9
  by Finn at 2018-11-16T11:35:09Z

  Bot now catches errors thrown when updating session.
  

• 20e99038
  by Finn at 2018-11-16T11:35:09Z

  Working commit
  

19 changed files:

• buildgrid/_app/bots/dummy.py
• buildgrid/_app/commands/cmd_bot.py
• buildgrid/bot/bot.py
• − buildgrid/bot/bot_session.py
• + buildgrid/bot/hardware/__init__.py
• + buildgrid/bot/hardware/device.py
• + buildgrid/bot/hardware/interface.py
• + buildgrid/bot/hardware/worker.py
• buildgrid/bot/bot_interface.py → buildgrid/bot/interface.py
• + buildgrid/bot/session.py
• + buildgrid/bot/tenant.py
• + buildgrid/bot/tenantmanager.py
• buildgrid/server/bots/instance.py
• buildgrid/server/execution/service.py
• tests/integration/bot_session.py
• tests/integration/bots_service.py
• tests/integration/operations_service.py
• + tests/utils/bots_interface.py
• + tests/utils/utils.py

Changes:

buildgrid/_app/bots/dummy.py

---

...	...	@@ -20,7 +20,7 @@ from buildgrid._protos.build.bazel.remote.execution.v2 import remote_execution_p
20	20	from buildgrid.utils import get_hostname
21	21
22	22
23		-def work_dummy(context, lease):
	23	+def work_dummy(lease, context, event):
24	24	""" Just returns lease after some random time
25	25	"""
26	26	action_result = remote_execution_pb2.ActionResult()

buildgrid/_app/commands/cmd_bot.py

---

...	...	@@ -28,8 +28,11 @@ from urllib.parse import urlparse
28	28	import click
29	29	import grpc
30	30
31		-from buildgrid.bot import bot, bot_interface
32		-from buildgrid.bot.bot_session import BotSession, Device, Worker
	31	+from buildgrid.bot import bot, interface, session
	32	+from buildgrid.bot.hardware.interface import HardwareInterface
	33	+from buildgrid.bot.hardware.device import Device
	34	+from buildgrid.bot.hardware.worker import Worker
	35	+
33	36
34	37	from ..bots import buildbox, dummy, host
35	38	from ..cli import pass_context
...	...	@@ -123,15 +126,13 @@ def cli(context, parent, update_period, remote, client_key, client_cert, server_
123	126	context.logger = logging.getLogger(__name__)
124	127	context.logger.debug("Starting for remote {}".format(context.remote))
125	128
126		- interface = bot_interface.BotInterface(context.channel)
127		-
	129	+ bot_interface = interface.BotInterface(context.channel)
128	130	worker = Worker()
129	131	worker.add_device(Device())
	132	+ hardware_interface = HardwareInterface(worker)
130	133
131		- bot_session = BotSession(parent, interface)
132		- bot_session.add_worker(worker)
133		-
134		- context.bot_session = bot_session
	134	+ context.bot_interface = bot_interface
	135	+ context.hardware_interface = hardware_interface
135	136
136	137
137	138	@cli.command('dummy', short_help="Run a dummy session simply returning leases.")
...	...	@@ -141,9 +142,10 @@ def run_dummy(context):
141	142	Creates a session, accepts leases, does fake work and updates the server.
142	143	"""
143	144	try:
144		- b = bot.Bot(context.bot_session, context.update_period)
145		- b.session(dummy.work_dummy,
146		- context)
	145	+ bot_session = session.BotSession(context.parent, context.bot_interface, context.hardware_interface,
	146	+ dummy.work_dummy, context)
	147	+ b = bot.Bot(bot_session, context.update_period)
	148	+ b.session()
147	149	except KeyboardInterrupt:
148	150	pass
149	151
...	...	@@ -156,9 +158,10 @@ def run_host_tools(context):
156	158	result back to CAS.
157	159	"""
158	160	try:
159		- b = bot.Bot(context.bot_session, context.update_period)
160		- b.session(host.work_host_tools,
161		- context)
	161	+ bot_session = session.BotSession(context.parent, context.bot_interface, context.hardware_interface,
	162	+ host.work_host_tools, context)
	163	+ b = bot.Bot(bot_session, context.update_period)
	164	+ b.session()
162	165	except KeyboardInterrupt:
163	166	pass
164	167
...	...	@@ -177,8 +180,9 @@ def run_buildbox(context, local_cas, fuse_dir):
177	180	context.fuse_dir = fuse_dir
178	181
179	182	try:
180		- b = bot.Bot(context.bot_session, context.update_period)
181		- b.session(buildbox.work_buildbox,
182		- context)
	183	+ bot_session = session.BotSession(context.parent, context.bot_interface, context.hardware_interface,
	184	+ buildbox.work_buildbox, context)
	185	+ b = bot.Bot(bot_session, context.update_period)
	186	+ b.session()
183	187	except KeyboardInterrupt:
184	188	pass

buildgrid/bot/bot.py

---

...	...	@@ -17,7 +17,7 @@
17	17	Bot
18	18	====
19	19
20		-Creates a bot session.
	20	+Creates a bot session and sends updates to the server.
21	21	"""
22	22
23	23	import asyncio
...	...	@@ -35,16 +35,17 @@ class Bot:
35	35	self._bot_session = bot_session
36	36	self._update_period = update_period
37	37
38		- def session(self, work, context):
	38	+ def session(self):
39	39	loop = asyncio.get_event_loop()
40	40
41		- self._bot_session.create_bot_session(work, context)
	41	+ self._bot_session.create_bot_session()
42	42
43	43	try:
44	44	task = asyncio.ensure_future(self._update_bot_session())
45	45	loop.run_forever()
46	46	except KeyboardInterrupt:
47	47	pass
	48	+
48	49	finally:
49	50	task.cancel()
50	51	loop.close()
...	...	@@ -54,5 +55,11 @@ class Bot:
54	55	Calls the server periodically to inform the server the client has not died.
55	56	"""
56	57	while True:
57		- self._bot_session.update_bot_session()
	58	+ try:
	59	+ self._bot_session.update_bot_session()
	60	+
	61	+ except Exception as e:
	62	+ self.__logger.error(e)
	63	+ raise
	64	+
58	65	await asyncio.sleep(self._update_period)

buildgrid/bot/bot_session.py deleted

---

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		-# Disable broad exception catch
16		-# pylint: disable=broad-except
17		-
18		-
19		-"""
20		-Bot Session
21		-====
22		-
23		-Allows connections
24		-"""
25		-import asyncio
26		-import logging
27		-import platform
28		-import uuid
29		-
30		-import grpc
31		-
32		-from buildgrid._enums import BotStatus, LeaseState
33		-from buildgrid._protos.google.rpc import code_pb2
34		-from buildgrid._protos.google.devtools.remoteworkers.v1test2 import bots_pb2, worker_pb2
35		-from buildgrid._exceptions import BotError
36		-
37		-
38		-class BotSession:
39		- def __init__(self, parent, interface):
40		- """ Unique bot ID within the farm used to identify this bot
41		- Needs to be human readable.
42		- All prior sessions with bot_id of same ID are invalidated.
43		- If a bot attempts to update an invalid session, it must be rejected and
44		- may be put in quarantine.
45		- """
46		- self.__logger = logging.getLogger(__name__)
47		-
48		- self._bot_id = '{}.{}'.format(parent, platform.node())
49		- self._context = None
50		- self._interface = interface
51		- self._leases = {}
52		- self._name = None
53		- self._parent = parent
54		- self._status = BotStatus.OK.value
55		- self._work = None
56		- self._worker = None
57		-
58		- @property
59		- def bot_id(self):
60		- return self._bot_id
61		-
62		- def add_worker(self, worker):
63		- self._worker = worker
64		-
65		- def create_bot_session(self, work, context=None):
66		- self.__logger.debug("Creating bot session")
67		- self._work = work
68		- self._context = context
69		-
70		- session = self._interface.create_bot_session(self._parent, self.get_pb2())
71		- self._name = session.name
72		-
73		- self.__logger.info("Created bot session with name: [%s]", self._name)
74		-
75		- for lease in session.leases:
76		- self._update_lease_from_server(lease)
77		-
78		- def update_bot_session(self):
79		- self.__logger.debug("Updating bot session: [%s]", self._bot_id)
80		- session = self._interface.update_bot_session(self.get_pb2())
81		- for k, v in list(self._leases.items()):
82		- if v.state == LeaseState.COMPLETED.value:
83		- del self._leases[k]
84		-
85		- for lease in session.leases:
86		- self._update_lease_from_server(lease)
87		-
88		- def get_pb2(self):
89		- leases = list(self._leases.values())
90		- if not leases:
91		- leases = None
92		-
93		- return bots_pb2.BotSession(worker=self._worker.get_pb2(),
94		- status=self._status,
95		- leases=leases,
96		- bot_id=self._bot_id,
97		- name=self._name)
98		-
99		- def lease_completed(self, lease):
100		- lease.state = LeaseState.COMPLETED.value
101		- self._leases[lease.id] = lease
102		-
103		- def _update_lease_from_server(self, lease):
104		- """
105		- State machine for any recieved updates to the leases.
106		- """
107		- # TODO: Compare with previous state of lease
108		- if lease.state == LeaseState.PENDING.value:
109		- lease.state = LeaseState.ACTIVE.value
110		- self._leases[lease.id] = lease
111		- self.update_bot_session()
112		- asyncio.ensure_future(self.create_work(lease))
113		-
114		- async def create_work(self, lease):
115		- self.__logger.debug("Work created: [%s]", lease.id)
116		- loop = asyncio.get_event_loop()
117		-
118		- try:
119		- lease = await loop.run_in_executor(None, self._work, self._context, lease)
120		-
121		- except grpc.RpcError as e:
122		- self.__logger.error(e)
123		- lease.status.CopyFrom(e.code())
124		-
125		- except BotError as e:
126		- self.__logger.error(e)
127		- lease.status.code = code_pb2.INTERNAL
128		-
129		- except Exception as e:
130		- self.__logger.error(e)
131		- lease.status.code = code_pb2.INTERNAL
132		-
133		- self.__logger.debug("Work complete: [%s]", lease.id)
134		- self.lease_completed(lease)
135		-
136		-
137		-class Worker:
138		- def __init__(self, properties=None, configs=None):
139		- self.properties = {}
140		- self._configs = {}
141		- self._devices = []
142		-
143		- if properties:
144		- for k, v in properties.items():
145		- if k == 'pool':
146		- self.properties[k] = v
147		- else:
148		- raise KeyError('Key not supported: [{}]'.format(k))
149		-
150		- if configs:
151		- for k, v in configs.items():
152		- if k == 'DockerImage':
153		- self.configs[k] = v
154		- else:
155		- raise KeyError('Key not supported: [{}]'.format(k))
156		-
157		- @property
158		- def configs(self):
159		- return self._configs
160		-
161		- def add_device(self, device):
162		- self._devices.append(device)
163		-
164		- def get_pb2(self):
165		- devices = [device.get_pb2() for device in self._devices]
166		- worker = worker_pb2.Worker(devices=devices)
167		- property_message = worker_pb2.Worker.Property()
168		- for k, v in self.properties.items():
169		- property_message.key = k
170		- property_message.value = v
171		- worker.properties.extend([property_message])
172		-
173		- config_message = worker_pb2.Worker.Config()
174		- for k, v in self.properties.items():
175		- property_message.key = k
176		- property_message.value = v
177		- worker.configs.extend([config_message])
178		-
179		- return worker
180		-
181		-
182		-class Device:
183		- def __init__(self, properties=None):
184		- """ Creates devices available to the worker
185		- The first device is know as the Primary Device - the revice which
186		- is running a bit and responsible to actually executing commands.
187		- All other devices are known as Attatched Devices and must be controlled
188		- by the Primary Device.
189		- """
190		-
191		- self._name = str(uuid.uuid4())
192		- self._properties = {}
193		-
194		- if properties:
195		- for k, v in properties.items():
196		- if k == 'os':
197		- self._properties[k] = v
198		-
199		- elif k == 'docker':
200		- if v not in ('True', 'False'):
201		- raise ValueError('Value not supported: [{}]'.format(v))
202		- self._properties[k] = v
203		-
204		- else:
205		- raise KeyError('Key not supported: [{}]'.format(k))
206		-
207		- @property
208		- def name(self):
209		- return self._name
210		-
211		- @property
212		- def properties(self):
213		- return self._properties
214		-
215		- def get_pb2(self):
216		- device = worker_pb2.Device(handle=self._name)
217		- property_message = worker_pb2.Device.Property()
218		- for k, v in self._properties.items():
219		- property_message.key = k
220		- property_message.value = v
221		- device.properties.extend([property_message])
222		- return device

buildgrid/bot/hardware/__init__.py

---

buildgrid/bot/hardware/device.py

---

	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	+
	16	+import uuid
	17	+from buildgrid._protos.google.devtools.remoteworkers.v1test2 import worker_pb2
	18	+
	19	+
	20	+class Device:
	21	+ """ Describes a device.
	22	+
	23	+ A device available to the :class:`Worker`. The first device is known as
	24	+ the Primary Device which is running a bot and responsible for actually
	25	+ executing commands. All other devices are known as Attatched Devices and
	26	+ must be controlled by the Primary Device.
	27	+ """
	28	+
	29	+ def __init__(self, properties=None):
	30	+ """ Initialises a :class:`Device` instances.
	31	+
	32	+ Property keys supported: `os`, `has-docker`.
	33	+
	34	+ Args:
	35	+ properties (dict(string : list(string))) : Properties of device. Keys may have
	36	+ multiple values.
	37	+ """
	38	+
	39	+ self.__properties = {}
	40	+ self.__property_keys = ['os', 'has-docker']
	41	+ self.__name = str(uuid.uuid4())
	42	+
	43	+ if properties:
	44	+ for k, l in properties.items():
	45	+ for v in l:
	46	+ self._add_property(k, v)
	47	+
	48	+ @property
	49	+ def name(self):
	50	+ """Returns the name of the device which is in the form of a `uuid4`."""
	51	+ return self.__name
	52	+
	53	+ @property
	54	+ def properties(self):
	55	+ """Returns the device properties."""
	56	+ return self.__properties
	57	+
	58	+ def get_pb2(self):
	59	+ """Returns the protobuffer representation of the :class:`Device`."""
	60	+ device = worker_pb2.Device(handle=self.__name)
	61	+ for k, v in self.__properties.items():
	62	+ for prop in v:
	63	+ property_message = worker_pb2.Device.Property()
	64	+ property_message.key = k
	65	+ property_message.value = prop
	66	+ device.properties.extend([property_message])
	67	+ return device
	68	+
	69	+ def _add_property(self, key, value):
	70	+ """Adds a property to the :class:`Device` instance."""
	71	+ if key in self.__property_keys:
	72	+ prop = self.__properties.get(key)
	73	+ if not prop:
	74	+ self.__properties[key] = [value]
	75	+ else:
	76	+ prop[key].append(value)
	77	+
	78	+ else:
	79	+ raise KeyError('Key not supported: [{}]'.format(key))

buildgrid/bot/hardware/interface.py

---

	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	+
	16	+"""
	17	+HardwareInterface
	18	+=================
	19	+
	20	+"""
	21	+
	22	+
	23	+from buildgrid._exceptions import FailedPreconditionError
	24	+
	25	+
	26	+class HardwareInterface:
	27	+ """Class to configure hardware and check requirements of a :class:`Lease`.
	28	+
	29	+ In the future this could also be used to request and display
	30	+ the status of hardware.
	31	+ """
	32	+
	33	+ def __init__(self, worker):
	34	+ """Initialises a new :class:`HardwareInstance`.
	35	+
	36	+ Args:
	37	+ worker (:class:`Worker`)
	38	+ """
	39	+ self._worker = worker
	40	+
	41	+ def configure_hardware(self, requirements):
	42	+ """Configure's hardware and checks to see if the requirements can
	43	+ be met.
	44	+
	45	+ Args:
	46	+ requirements (:class:`Worker`): Requirements for the job.
	47	+ """
	48	+ worker = self._worker
	49	+
	50	+ for config_requirement in requirements.configs:
	51	+ if config_requirement.key not in worker.configs:
	52	+ raise FailedPreconditionError("Requirements can't be met. "
	53	+ "Requirements=[{}]".format(config_requirement))
	54	+
	55	+ def get_worker_pb2(self):
	56	+ """ Returns the protobuffer representation of the :class:`Worker`.
	57	+ """
	58	+ return self._worker.get_pb2()

buildgrid/bot/hardware/worker.py

---

	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	+
	16	+"""
	17	+Worker
	18	+======
	19	+
	20	+A worker.
	21	+"""
	22	+
	23	+from buildgrid._protos.google.devtools.remoteworkers.v1test2 import worker_pb2
	24	+
	25	+
	26	+class Worker:
	27	+ """ Describes a worker, which is a list of one of more devices and the
	28	+ connextions between them. A device could be a computer, a phone or a GPU.
	29	+
	30	+ The first device is known as the Primary Device which is responsible for
	31	+ executing commands. All other devices are known as Attatched Devices and
	32	+ must be controlled by the Primary Device.
	33	+ """
	34	+
	35	+ def __init__(self, properties=None, configs=None):
	36	+ """ Create an instance of a :class:`Worker`.
	37	+
	38	+ Property keys supported: `pool`.
	39	+ Config keys supported: `DockerImage`.
	40	+
	41	+ Args:
	42	+ properties (dict(string : list(string))) : Properties of worker. Keys may have
	43	+ multiple values.
	44	+ configs (dict(string : list(string))) : Configurations of a worker. Keys may have
	45	+ multiple values.
	46	+
	47	+ """
	48	+
	49	+ self._devices = []
	50	+ self.__configs = {}
	51	+ self.__properties = {}
	52	+ self.__property_keys = ['pool']
	53	+ self.__config_keys = ['DockerImage']
	54	+
	55	+ if properties:
	56	+ for k, l in properties.items():
	57	+ for v in l:
	58	+ self._add_property(k, v)
	59	+
	60	+ if configs:
	61	+ for k, l in configs.items():
	62	+ for v in l:
	63	+ self._add_config(k, v)
	64	+
	65	+ @property
	66	+ def configs(self):
	67	+ """Returns configurations supported by :class:`Worker`."""
	68	+ return self.__configs
	69	+
	70	+ @property
	71	+ def properties(self):
	72	+ """Returns properties supported by :class:`Worker`."""
	73	+ return self.__properties
	74	+
	75	+ def add_device(self, device):
	76	+ """ Adds a class:`Device` to this instance. First device added should be the Primary Device."""
	77	+ self._devices.append(device)
	78	+
	79	+ def get_pb2(self):
	80	+ """Returns the protobuffer representation of a :class:`Device`."""
	81	+ devices = [device.get_pb2() for device in self._devices]
	82	+ worker = worker_pb2.Worker(devices=devices)
	83	+
	84	+ for k, v in self.__properties.items():
	85	+ for prop in v:
	86	+ property_message = worker_pb2.Device.Property()
	87	+ property_message.key = k
	88	+ property_message.value = prop
	89	+ worker.properties.extend([property_message])
	90	+
	91	+ for k, v in self.__configs.items():
	92	+ for cfg in v:
	93	+ config_message = worker_pb2.Worker.Config()
	94	+ config_message.key = k
	95	+ config_message.value = cfg
	96	+ worker.configs.extend([config_message])
	97	+
	98	+ return worker
	99	+
	100	+ def _add_config(self, key, value):
	101	+ """Adds a configuration. Will raise KeyError if not supported."""
	102	+ if key in self.__config_keys:
	103	+ cfg = self.__configs.get(key)
	104	+ if not cfg:
	105	+ self.__configs[key] = [value]
	106	+ else:
	107	+ cfg[key].append(value)
	108	+
	109	+ else:
	110	+ raise KeyError('Key not supported: [{}]'.format(key))
	111	+
	112	+ def _add_property(self, key, value):
	113	+ """Adds a property. Will raise KeyError if not supported."""
	114	+ if key in self.__property_keys:
	115	+ prop = self.__properties.get(key)
	116	+ if not prop:
	117	+ self.__properties[key] = [value]
	118	+ else:
	119	+ prop[key].append(value)
	120	+
	121	+ else:
	122	+ raise KeyError('Key not supported: [{}]'.format(key))

buildgrid/bot/bot_interface.py → buildgrid/bot/interface.py

---

...	...	@@ -15,12 +15,13 @@
15	15
16	16	"""
17	17	Bot Interface
18		-====
	18	+=============
19	19
20	20	Interface to grpc
21	21	"""
22	22
23	23	import logging
	24	+import grpc
24	25
25	26	from buildgrid._protos.google.devtools.remoteworkers.v1test2 import bots_pb2, bots_pb2_grpc
26	27
...	...	@@ -38,10 +39,18 @@ class BotInterface:
38	39	def create_bot_session(self, parent, bot_session):
39	40	request = bots_pb2.CreateBotSessionRequest(parent=parent,
40	41	bot_session=bot_session)
41		- return self._stub.CreateBotSession(request)
	42	+ try:
	43	+ return self._stub.CreateBotSession(request)
	44	+
	45	+ except grpc.RpcError as e:
	46	+ self.__logger.error(e)
42	47
43	48	def update_bot_session(self, bot_session, update_mask=None):
44	49	request = bots_pb2.UpdateBotSessionRequest(name=bot_session.name,
45	50	bot_session=bot_session,
46	51	update_mask=update_mask)
47		- return self._stub.UpdateBotSession(request)
	52	+ try:
	53	+ return self._stub.UpdateBotSession(request)
	54	+
	55	+ except grpc.RpcError as e:
	56	+ self.__logger.error(e)

buildgrid/bot/session.py

---

	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	+
	16	+"""
	17	+Bot Session
	18	+===========
	19	+
	20	+Allows connections
	21	+"""
	22	+import logging
	23	+import platform
	24	+
	25	+from buildgrid._enums import BotStatus, LeaseState
	26	+from buildgrid._protos.google.devtools.remoteworkers.v1test2 import bots_pb2
	27	+from buildgrid._protos.google.rpc import code_pb2
	28	+
	29	+from buildgrid._exceptions import FailedPreconditionError
	30	+
	31	+from .tenantmanager import TenantManager
	32	+
	33	+
	34	+class BotSession:
	35	+ def __init__(self, parent, bots_interface, hardware_interface, work, context=None):
	36	+ """ Unique bot ID within the farm used to identify this bot
	37	+ Needs to be human readable.
	38	+ All prior sessions with bot_id of same ID are invalidated.
	39	+ If a bot attempts to update an invalid session, it must be rejected and
	40	+ may be put in quarantine.
	41	+ """
	42	+ self.__logger = logging.getLogger(__name__)
	43	+
	44	+ self._bots_interface = bots_interface
	45	+ self._hardware_interface = hardware_interface
	46	+
	47	+ self._status = BotStatus.OK.value
	48	+ self._tenant_manager = TenantManager()
	49	+
	50	+ self.__parent = parent
	51	+ self.__bot_id = '{}.{}'.format(parent, platform.node())
	52	+ self.__name = None
	53	+
	54	+ self._work = work
	55	+ self._context = context
	56	+
	57	+ @property
	58	+ def bot_id(self):
	59	+ return self.__bot_id
	60	+
	61	+ def create_bot_session(self):
	62	+ self.__logger.debug("Creating bot session")
	63	+
	64	+ session = self._bots_interface.create_bot_session(self.__parent, self.get_pb2())
	65	+ self.__name = session.name
	66	+
	67	+ self.__logger.info("Created bot session with name: [%s]", self.__name)
	68	+
	69	+ for lease in session.leases:
	70	+ self._register_lease(lease)
	71	+
	72	+ def update_bot_session(self):
	73	+ self.__logger.debug("Updating bot session: [%s]", self.__bot_id)
	74	+
	75	+ session = self._bots_interface.update_bot_session(self.get_pb2())
	76	+ server_ids = []
	77	+
	78	+ for lease in session.leases:
	79	+ server_ids.append(lease.id)
	80	+
	81	+ lease_state = LeaseState(lease.state)
	82	+ if lease_state == LeaseState.PENDING:
	83	+ self._register_lease(lease)
	84	+
	85	+ elif lease_state == LeaseState.CANCELLED:
	86	+ self._tenant_manager.cancel_tenancy(lease.id)
	87	+
	88	+ closed_lease_ids = [x for x in self._tenant_manager.get_lease_ids() if x not in server_ids]
	89	+
	90	+ for lease_id in closed_lease_ids:
	91	+ self._tenant_manager.cancel_tenancy(lease_id)
	92	+ self._tenant_manager.remove_tenant(lease_id)
	93	+
	94	+ def get_pb2(self):
	95	+ return bots_pb2.BotSession(worker=self._hardware_interface.get_worker_pb2(),
	96	+ status=self._status,
	97	+ leases=self._tenant_manager.get_leases(),
	98	+ bot_id=self.__bot_id,
	99	+ name=self.__name)
	100	+
	101	+ def _register_lease(self, lease):
	102	+ lease_id = lease.id
	103	+ try:
	104	+ self._tenant_manager.create_tenancy(lease)
	105	+
	106	+ except KeyError as e:
	107	+ self.__logger.error(e)
	108	+
	109	+ else:
	110	+ try:
	111	+ self._hardware_interface.configure_hardware(lease.requirements)
	112	+
	113	+ except FailedPreconditionError as e:
	114	+ self.__logger.error(e)
	115	+ self._tenant_manager.complete_lease(lease_id, status=code_pb2.FailedPreconditionError)
	116	+
	117	+ else:
	118	+ self._tenant_manager.create_work(lease_id, self._work, self._context)

buildgrid/bot/tenant.py

---

	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	+# Disable broad exception catch
	16	+# pylint: disable=broad-except
	17	+
	18	+
	19	+import asyncio
	20	+import logging
	21	+import threading
	22	+
	23	+from functools import partial
	24	+
	25	+import grpc
	26	+
	27	+from buildgrid._protos.google.rpc import code_pb2
	28	+from buildgrid._enums import LeaseState
	29	+from buildgrid._exceptions import BotError
	30	+
	31	+
	32	+class Tenant:
	33	+
	34	+ def __init__(self, lease):
	35	+ """Initialises a new :class:`Tenant`.
	36	+
	37	+ Holds a lease of work and keeps the status of it. Will raise a `ValueError`
	38	+ if `lease` is not in a `PENDING` state.
	39	+
	40	+ Args:
	41	+ lease (:class:`Lease`) : A new lease.
	42	+ """
	43	+
	44	+ if lease.state != LeaseState.PENDING.value:
	45	+ raise ValueError("Lease state not `PENDING`")
	46	+
	47	+ self._lease = lease
	48	+ self.__logger = logging.getLogger(__name__)
	49	+ self.__lease_cancelled = False
	50	+ self.__tenant_completed = False
	51	+
	52	+ @property
	53	+ def lease(self):
	54	+ """Returns the lease"""
	55	+ return self._lease
	56	+
	57	+ @property
	58	+ def tenant_completed(self):
	59	+ """Returns `True` if the work has completed or sucessfully stopped its work."""
	60	+ return self.__tenant_completed
	61	+
	62	+ @property
	63	+ def lease_cancelled(self):
	64	+ """Returns `True` if the lease has been cancelled."""
	65	+ return self.__lease_cancelled
	66	+
	67	+ def cancel_lease(self):
	68	+ """Cancel the lease."""
	69	+ self.__lease_cancelled = True
	70	+ self.update_lease_state(LeaseState.CANCELLED)
	71	+
	72	+ def get_lease_state(self):
	73	+ """Returns the :class:`LeaseState`."""
	74	+ return LeaseState(self._lease.state)
	75	+
	76	+ def update_lease_result(self, result):
	77	+ """Update the lease result.
	78	+
	79	+ Args:
	80	+ result (:class:`Any`) : The result of the lease."""
	81	+ self._lease.result.CopyFrom(result)
	82	+
	83	+ def update_lease_state(self, state):
	84	+ """Update the lease state.
	85	+
	86	+ Args:
	87	+ state (:class:`LeaseState`) : State of the lease.
	88	+ """
	89	+ self._lease.state = state.value
	90	+
	91	+ def update_lease_status(self, status):
	92	+ """Update the lease status.
	93	+
	94	+ Args:
	95	+ status (:class:`Status`) : Status of the lease.
	96	+ """
	97	+ self._lease.status.CopyFrom(status)
	98	+
	99	+ async def run_work(self, work, context=None, executor=None):
	100	+ """Runs work.
	101	+
	102	+ Work is run in an executor.
	103	+
	104	+ Args:
	105	+ work (func) : Work to do.
	106	+ context (context) : Context for work.
	107	+ executor (:class:`ThreadPoolExecutor`) : Thread pool for work to run on.
	108	+ """
	109	+ self.__logger.debug("Work created. Lease_id=[%s]", self._lease.id)
	110	+
	111	+ loop = asyncio.get_event_loop()
	112	+
	113	+ try:
	114	+ event = threading.Event()
	115	+ lease = await loop.run_in_executor(executor, partial(work, self._lease, context, event))
	116	+
	117	+ except asyncio.CancelledError:
	118	+ self.__logger.error("Lease cancelled: lease_id=[%s]", self._lease.id)
	119	+ event.set()
	120	+ self.__tenant_completed = True
	121	+ # Propagate error to task wrapper
	122	+ raise
	123	+
	124	+ except grpc.RpcError as e:
	125	+ self.__logger.error(e)
	126	+ lease.status.CopyFrom(e.code())
	127	+
	128	+ except BotError as e:
	129	+ self.__logger.error(e)
	130	+ lease.status.code = code_pb2.INTERNAL
	131	+
	132	+ except Exception as e:
	133	+ self.__logger.error(e)
	134	+ lease.status.code = code_pb2.INTERNAL
	135	+
	136	+ self.__tenant_completed = True
	137	+ self.__logger.debug("Work completed: lease_id=[%s]", lease.id)
	138	+
	139	+ return lease

buildgrid/bot/tenantmanager.py

---

	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	+
	16	+import asyncio
	17	+import logging
	18	+from functools import partial
	19	+
	20	+from buildgrid._enums import LeaseState
	21	+
	22	+from .tenant import Tenant
	23	+
	24	+
	25	+class TenantManager:
	26	+ """Manages a number of :class:`Tenant`s.
	27	+
	28	+ Creates work to do, monitors and removes leases of work.
	29	+ """
	30	+
	31	+ def __init__(self):
	32	+ """Initialises an instance of the :class:`TenantManager`."""
	33	+ self.logger = logging.getLogger(__name__)
	34	+ self._tenants = {}
	35	+ self._tasks = {}
	36	+
	37	+ def create_tenancy(self, lease):
	38	+ """Create a new :class:`Tenant`.
	39	+
	40	+ Args:
	41	+ lease (:class:Lease) : Lease of work to do.
	42	+ """
	43	+ lease_id = lease.id
	44	+
	45	+ if lease_id not in self._tenants:
	46	+ tenant = Tenant(lease)
	47	+ self._tenants[lease_id] = tenant
	48	+
	49	+ else:
	50	+ raise KeyError("Lease id already exists: [{}]".format(lease_id))
	51	+
	52	+ def remove_tenant(self, lease_id):
	53	+ """Attempts to remove a tenant.
	54	+
	55	+ If the tenant has not been cancelled, it will cancel it. If the tenant has
	56	+ not completed, it will not remove it.
	57	+
	58	+ Args:
	59	+ lease_id (string) : The lease id.
	60	+ """
	61	+ if not self._tenants[lease_id].lease_cancelled:
	62	+ self.logger.error("Attempting to remove a lease not cancelled."
	63	+ "Bot will attempt to cancel lease."
	64	+ "Lease id=[{}]".format(lease_id))
	65	+ self.cancel_tenancy(lease_id)
	66	+
	67	+ elif not self._tenants[lease_id].tenant_completed:
	68	+ self.logger.debug("Lease cancelled but tenant not completed."
	69	+ "Lease=[{}]".format(self._tenants[lease_id].lease))
	70	+
	71	+ else:
	72	+ self.logger.debug("Removing tenant=[{}]".format(lease_id))
	73	+ self._tenants.pop(lease_id)
	74	+ self._tasks.pop(lease_id)
	75	+
	76	+ def get_leases(self):
	77	+ """Returns a list of leases managed by this instance."""
	78	+ leases = []
	79	+ for tenant in self._tenants.values():
	80	+ leases.append(tenant.lease)
	81	+
	82	+ if not leases:
	83	+ return None
	84	+
	85	+ return leases
	86	+
	87	+ def get_lease_ids(self):
	88	+ """Returns a list of lease ids."""
	89	+ return self._tenants.keys()
	90	+
	91	+ def get_lease_state(self, lease_id):
	92	+ """Returns the lease state
	93	+
	94	+ Args:
	95	+ lease_id (string) : The lease id.
	96	+ """
	97	+ return self._tenants[lease_id].get_lease_state()
	98	+
	99	+ def complete_lease(self, lease_id, status, task=None):
	100	+ """Informs the :class:`TenantManager` that the lease has completed.
	101	+
	102	+ If it was not cancelled, it will update with the result returned from
	103	+ the task.
	104	+
	105	+ Args:
	106	+ lease_id (string) : The lease id.
	107	+ status (:class:`Status`) : The final status of the lease.
	108	+ task (asyncio.Task) : The task of work.
	109	+ """
	110	+ if status is not None:
	111	+ self._update_lease_status(lease_id, status)
	112	+
	113	+ if task:
	114	+ if not task.cancelled():
	115	+ self._update_lease_result(lease_id, task.result().result)
	116	+ self._update_lease_state(lease_id, LeaseState.COMPLETED)
	117	+
	118	+ def create_work(self, lease_id, work, context):
	119	+ """Creates work to do.
	120	+
	121	+ Will place work on an asyncio loop with a callback to `complete_lease`.
	122	+
	123	+ Args:
	124	+ lease_id (string) : The lease id.
	125	+ work (func) : Work to do.
	126	+ context (context) : Context for work function.
	127	+ """
	128	+ self._update_lease_state(lease_id, LeaseState.ACTIVE)
	129	+ tenant = self._tenants[lease_id]
	130	+ task = asyncio.ensure_future(tenant.run_work(work, context))
	131	+
	132	+ task.add_done_callback(partial(self.complete_lease, lease_id, None))
	133	+ self._tasks[lease_id] = task
	134	+
	135	+ def cancel_tenancy(self, lease_id):
	136	+ """Cancels tenancy and any work being done.
	137	+
	138	+ Args:
	139	+ lease_id (string) : The lease id.
	140	+ """
	141	+ if not self._tenants[lease_id].lease_cancelled:
	142	+ self._tenants[lease_id].cancel_lease()
	143	+ self._tasks[lease_id].cancel()
	144	+
	145	+ def tenant_completed(self, lease_id):
	146	+ """Returns `True` if the work has been completed.
	147	+
	148	+ Args:
	149	+ lease_id (string) : The lease id.
	150	+ """
	151	+ return self._tenants[lease_id].tenant_completed
	152	+
	153	+ def _update_lease_result(self, lease_id, result):
	154	+ """Updates the lease with the result."""
	155	+ self._tenants[lease_id].update_lease_result(result)
	156	+
	157	+ def _update_lease_state(self, lease_id, state):
	158	+ """Updates the lease state."""
	159	+ self._tenants[lease_id].update_lease_state(state)
	160	+
	161	+ def _update_lease_status(self, lease_id, status):
	162	+ """Updates the lease status."""
	163	+ self._tenants[lease_id].update_lease_status(status)

buildgrid/server/bots/instance.py

---

...	...	@@ -23,7 +23,7 @@ Instance of the Remote Workers interface.
23	23	import logging
24	24	import uuid
25	25
26		-from buildgrid._exceptions import InvalidArgumentError, OutOfSyncError
	26	+from buildgrid._exceptions import InvalidArgumentError
27	27
28	28	from ..job import LeaseState
29	29
...	...	@@ -34,7 +34,7 @@ class BotsInterface:
34	34	self.__logger = logging.getLogger(__name__)
35	35
36	36	self._bot_ids = {}
37		- self._bot_sessions = {}
	37	+ self._assigned_leases = {}
38	38	self._scheduler = scheduler
39	39
40	40	def register_instance_with_server(self, instance_name, server):
...	...	@@ -59,18 +59,16 @@ class BotsInterface:
59	59
60	60	# Bot session name, selected by the server
61	61	name = "{}/{}".format(parent, str(uuid.uuid4()))
62		-
63	62	bot_session.name = name
64	63
65	64	self._bot_ids[name] = bot_id
66		- self._bot_sessions[name] = bot_session
	65	+
67	66	self.__logger.info("Created bot session name=[%s] with bot_id=[%s]", name, bot_id)
68	67
69		- # TODO: Send worker capabilities to the scheduler!
70		- leases = self._scheduler.request_job_leases({})
71		- if leases:
72		- bot_session.leases.extend(leases)
	68	+ # We want to keep a copy of lease ids we have assigned
	69	+ self._assigned_leases[name] = set()
73	70
	71	+ self._request_leases(bot_session)
74	72	return bot_session
75	73
76	74	def update_bot_session(self, name, bot_session):
...	...	@@ -79,39 +77,53 @@ class BotsInterface:
79	77	"""
80	78	self.__logger.debug("Updating bot session name=[%s]", name)
81	79	self._check_bot_ids(bot_session.bot_id, name)
82		-
83		- leases = filter(None, [self._check_lease_state(lease) for lease in bot_session.leases])
	80	+ self._check_assigned_leases(bot_session)
84	81
85	82	for lease in bot_session.leases:
86		- lease.Clear()
87		-
88		- bot_session.leases.extend(leases)
	83	+ checked_lease = self._check_lease_state(lease)
	84	+ if not checked_lease:
	85	+ # TODO: Make sure we don't need this
	86	+ try:
	87	+ self._assigned_leases[name].remove(lease.id)
	88	+ except KeyError:
	89	+ pass
	90	+ lease.Clear()
	91	+
	92	+ self._request_leases(bot_session)
	93	+ return bot_session
89	94
	95	+ def _request_leases(self, bot_session):
90	96	# TODO: Send worker capabilities to the scheduler!
	97	+ # Only send one lease at a time currently.
91	98	if not bot_session.leases:
92	99	leases = self._scheduler.request_job_leases({})
93	100	if leases:
	101	+ for lease in leases:
	102	+ self._assigned_leases[bot_session.name].add(lease.id)
94	103	bot_session.leases.extend(leases)
95	104
96		- self._bot_sessions[name] = bot_session
97		- return bot_session
98		-
99	105	def _check_lease_state(self, lease):
	106	+ # careful here
	107	+ # should store bot name in scheduler
	108	+ lease_state = LeaseState(lease.state)
100	109
101		- # Check for cancelled lease
102		- if self._scheduler.get_lease_cancelled(lease.id):
	110	+ # Lease has replied with cancelled, remove
	111	+ if lease_state == LeaseState.CANCELLED:
	112	+ return None
	113	+
	114	+ try:
	115	+ if self._scheduler.get_job_lease_cancelled(lease.id):
	116	+ lease.state.CopyFrom(LeaseState.CANCELLED.value)
	117	+ return lease
	118	+ except KeyError:
	119	+ # Job does not exist, remove from bot.
103	120	return None
104	121
105		- # If not cancelled, update the status
106	122	self._scheduler.update_job_lease(lease)
107	123
108		- lease_state = LeaseState(lease.state)
109	124	if lease_state == LeaseState.COMPLETED:
110	125	return None
111	126
112		- elif lease_state == LeaseState.CANCELLED:
113		- return None
114		-
115	127	return lease
116	128
117	129	def _check_bot_ids(self, bot_id, name=None):
...	...	@@ -134,6 +146,19 @@ class BotsInterface:
134	146	'Bot id already registered. ID sent: [{}].'
135	147	'Id registered: [{}] with name: [{}]'.format(bot_id, _bot_id, _name))
136	148
	149	+ def _check_assigned_leases(self, bot_session):
	150	+ session_lease_ids = []
	151	+
	152	+ for lease in bot_session.leases:
	153	+ session_lease_ids.append(lease.id)
	154	+
	155	+ for lease_id in self._assigned_leases[bot_session.name]:
	156	+ if lease_id not in session_lease_ids:
	157	+ self.__logger.error("Assigned lease id=[%s],"
	158	+ " not found on bot with name=[%s] and id=[%s]."
	159	+ " Retrying job", lease_id, bot_session.name, bot_session.bot_id)
	160	+ self._scheduler.retry_job(lease_id)
	161	+
137	162	def _close_bot_session(self, name):
138	163	""" Before removing the session, close any leases and
139	164	requeue with high priority.
...	...	@@ -144,10 +169,9 @@ class BotsInterface:
144	169	raise InvalidArgumentError("Bot id does not exist: [{}]".format(name))
145	170
146	171	self.__logger.debug("Attempting to close [%s] with name: [%s]", bot_id, name)
147		- for lease in self._bot_sessions[name].leases:
148		- if lease.state != LeaseState.COMPLETED.value:
149		- # TODO: Be wary here, may need to handle rejected leases in future
150		- self._scheduler.retry_job(lease.id)
	172	+ for lease_id in self._assigned_leases[name]:
	173	+ self._scheduler.retry_job(lease_id)
	174	+ self._assigned_leases.pop(name)
151	175
152	176	self.__logger.debug("Closing bot session: [%s]", name)
153	177	self._bot_ids.pop(name)

buildgrid/server/execution/service.py

---

...	...	@@ -80,7 +80,7 @@ class ExecutionService(remote_execution_pb2_grpc.ExecutionServicer):
80	80	yield operations_pb2.Operation()
81	81
82	82	except CancelledError as e:
83		- self.logger.error(e)
	83	+ self.__logger.error(e)
84	84	context.set_details(str(e))
85	85	context.set_code(grpc.StatusCode.CANCELLED)
86	86	yield operations_pb2.Operation()
...	...	@@ -118,7 +118,7 @@ class ExecutionService(remote_execution_pb2_grpc.ExecutionServicer):
118	118	yield operations_pb2.Operation()
119	119
120	120	except CancelledError as e:
121		- self.logger.error(e)
	121	+ self.__logger.error(e)
122	122	context.set_details(str(e))
123	123	context.set_code(grpc.StatusCode.CANCELLED)
124	124	yield operations_pb2.Operation()

tests/integration/bot_session.py

---

...	...	@@ -14,56 +14,174 @@
14	14
15	15	# pylint: disable=redefined-outer-name
16	16
17		-import uuid
	17	+import asyncio
18	18
	19	+import grpc
19	20	import pytest
20	21
21		-from buildgrid.bot import bot_session
	22	+from buildgrid._enums import LeaseState
	23	+from buildgrid._protos.google.devtools.remoteworkers.v1test2 import bots_pb2
	24	+from buildgrid.bot.hardware.worker import Worker
	25	+from buildgrid.bot.hardware.interface import HardwareInterface
	26	+from buildgrid.bot.session import BotSession
	27	+from buildgrid.bot.interface import BotInterface
22	28
	29	+from ..utils.utils import run_in_subprocess
	30	+from ..utils.bots_interface import serve_bots_interface
23	31
24		-@pytest.mark.parametrize("docker_value", ["True", "False"])
25		-@pytest.mark.parametrize("os_value", ["nexus7", "nexus8"])
26		-def test_create_device(docker_value, os_value):
27		- properties = {'docker': docker_value, 'os': os_value}
28		- device = bot_session.Device(properties)
29	32
30		- assert uuid.UUID(device.name, version=4)
31		- assert properties == device.properties
	33	+INSTANCES = ['', 'instance']
32	34
33	35
34		-def test_create_device_key_fail():
35		- properties = {'voight': 'kampff'}
	36	+# Use subprocess to avoid creation of gRPC threads in main process
	37	+# See https://github.com/grpc/grpc/blob/master/doc/fork_support.md
	38	+# Multiprocessing uses pickle which protobufs don't work with
	39	+# Workaround wrapper to send messages as strings
	40	+class ServerInterface:
36	41
37		- with pytest.raises(KeyError):
38		- bot_session.Device(properties)
	42	+ def __init__(self, remote):
	43	+ self.__remote = remote
39	44
	45	+ def create_bot_session(self, parent, bot_session):
40	46
41		-def test_create_device_value_fail():
42		- properties = {'docker': True}
	47	+ def __create_bot_session(queue, remote, parent, string_bot_session):
	48	+ bot_session = bots_pb2.BotSession()
	49	+ bot_session.ParseFromString(string_bot_session)
43	50
44		- with pytest.raises(ValueError):
45		- bot_session.Device(properties)
	51	+ interface = BotInterface(grpc.insecure_channel(remote))
46	52
	53	+ result = interface.create_bot_session(parent, bot_session)
	54	+ queue.put(result.SerializeToString())
47	55
48		-def test_create_worker():
49		- properties = {'pool': 'swim'}
50		- configs = {'DockerImage': 'Windows'}
51		- worker = bot_session.Worker(properties, configs)
	56	+ string_bot_session = bot_session.SerializeToString()
	57	+ result = run_in_subprocess(__create_bot_session,
	58	+ self.__remote, parent, string_bot_session)
52	59
53		- assert properties == worker.properties
54		- assert configs == worker.configs
	60	+ bot_session = bots_pb2.BotSession()
	61	+ bot_session.ParseFromString(result)
	62	+ return bot_session
55	63
56		- device = bot_session.Device()
57		- worker.add_device(device)
	64	+ def update_bot_session(self, bot_session, update_mask=None):
58	65
59		- assert worker._devices[0] == device
	66	+ def __update_bot_session(queue, remote, string_bot_session, update_mask):
	67	+ bot_session = bots_pb2.BotSession()
	68	+ bot_session.ParseFromString(string_bot_session)
60	69
	70	+ interface = BotInterface(grpc.insecure_channel(remote))
61	71
62		-def test_create_worker_key_fail():
63		- properties = {'voight': 'kampff'}
64		- configs = {'voight': 'kampff'}
	72	+ result = interface.update_bot_session(bot_session, update_mask)
	73	+ queue.put(result.SerializeToString())
65	74
66		- with pytest.raises(KeyError):
67		- bot_session.Worker(properties)
68		- with pytest.raises(KeyError):
69		- bot_session.Worker(configs)
	75	+ string_bot_session = bot_session.SerializeToString()
	76	+ result = run_in_subprocess(__update_bot_session,
	77	+ self.__remote, string_bot_session, update_mask)
	78	+
	79	+ bot_session = bots_pb2.BotSession()
	80	+ bot_session.ParseFromString(result)
	81	+ return bot_session
	82	+
	83	+
	84	+@pytest.mark.parametrize('instance', INSTANCES)
	85	+def test_create_bot_session(instance):
	86	+
	87	+ with serve_bots_interface([instance]) as server:
	88	+ interface = ServerInterface(server.remote)
	89	+ hardware_interface = HardwareInterface(Worker())
	90	+ session = BotSession(instance, interface, hardware_interface, None)
	91	+ session.create_bot_session()
	92	+ assert session.get_pb2() == server.get_bot_session()
	93	+
	94	+
	95	+@pytest.mark.parametrize('instance', INSTANCES)
	96	+def test_update_bot_session(instance):
	97	+
	98	+ with serve_bots_interface([instance]) as server:
	99	+ interface = ServerInterface(server.remote)
	100	+ hardware_interface = HardwareInterface(Worker())
	101	+ session = BotSession(instance, interface, hardware_interface, None)
	102	+ session.create_bot_session()
	103	+ assert session.get_pb2() == server.get_bot_session()
	104	+ session.update_bot_session()
	105	+ assert session.get_pb2() == server.get_bot_session()
	106	+
	107	+
	108	+@pytest.mark.parametrize('instance', INSTANCES)
	109	+def test_create_bot_session_with_work(instance):
	110	+
	111	+ def __work(lease, context, event):
	112	+ return lease
	113	+
	114	+ with serve_bots_interface([instance]) as server:
	115	+ interface = ServerInterface(server.remote)
	116	+ hardware_interface = HardwareInterface(Worker())
	117	+ session = BotSession(instance, interface, hardware_interface, __work)
	118	+ server.inject_work()
	119	+ session.create_bot_session()
	120	+
	121	+ assert len(session.get_pb2().leases) == 1
	122	+
	123	+ loop = asyncio.get_event_loop()
	124	+ for task in asyncio.Task.all_tasks():
	125	+ loop.run_until_complete(task)
	126	+
	127	+ assert session.get_pb2().leases[0].state == LeaseState.COMPLETED.value
	128	+
	129	+
	130	+@pytest.mark.parametrize('instance', INSTANCES)
	131	+def test_update_bot_session_with_work(instance):
	132	+
	133	+ def __work(lease, context, event):
	134	+ return lease
	135	+
	136	+ with serve_bots_interface([instance]) as server:
	137	+ interface = ServerInterface(server.remote)
	138	+ hardware_interface = HardwareInterface(Worker())
	139	+ session = BotSession(instance, interface, hardware_interface, __work)
	140	+ session.create_bot_session()
	141	+ server.inject_work()
	142	+ session.update_bot_session()
	143	+
	144	+ assert len(session.get_pb2().leases) == 1
	145	+
	146	+ loop = asyncio.get_event_loop()
	147	+ for task in asyncio.Task.all_tasks():
	148	+ loop.run_until_complete(task)
	149	+
	150	+ assert session.get_pb2().leases[0].state == LeaseState.COMPLETED.value
	151	+
	152	+
	153	+@pytest.mark.parametrize('instance', INSTANCES)
	154	+def test_cancel_leases(instance):
	155	+
	156	+ def __work(lease, context, cancel_event):
	157	+ # while not cancel_event.is_set():
	158	+
	159	+ return lease
	160	+
	161	+ with serve_bots_interface([instance]) as server:
	162	+ interface = ServerInterface(server.remote)
	163	+ hardware_interface = HardwareInterface(Worker())
	164	+ session = BotSession(instance, interface, hardware_interface, __work)
	165	+
	166	+ lease = bots_pb2.Lease()
	167	+ lease.state = LeaseState.PENDING.value
	168	+ lease.id = 'foo'
	169	+ server.inject_work(lease)
	170	+ session.create_bot_session()
	171	+
	172	+ leases_pb2 = session.get_pb2().leases
	173	+ assert len(leases_pb2) == 1
	174	+ assert leases_pb2[0].state == LeaseState.ACTIVE.value
	175	+
	176	+ server.cancel_lease(leases_pb2[0].id)
	177	+ session.update_bot_session()
	178	+ assert len(session.get_pb2().leases) == 1
	179	+
	180	+ loop = asyncio.get_event_loop()
	181	+ for task in asyncio.Task.all_tasks():
	182	+ try:
	183	+ loop.run_until_complete(task)
	184	+ except asyncio.CancelledError:
	185	+ pass
	186	+
	187	+ assert session.get_pb2().leases[0].state == LeaseState.CANCELLED.value

tests/integration/bots_service.py

---

...	...	@@ -17,7 +17,6 @@
17	17
18	18	# pylint: disable=redefined-outer-name
19	19
20		-import copy
21	20	from unittest import mock
22	21
23	22	import grpc
...	...	@@ -151,126 +150,11 @@ def test_update_leases_with_work(bot_session, context, instance):
151	150
152	151
153	152	def test_update_leases_work_complete(bot_session, context, instance):
154		- request = bots_pb2.CreateBotSessionRequest(parent='',
155		- bot_session=bot_session)
156		- # Create bot session
157		- # Simulated the severed binding between client and server
158		- response = copy.deepcopy(instance.CreateBotSession(request, context))
159		-
160		- # Inject work
161		- action_digest = remote_execution_pb2.Digest(hash='gaff')
162		- _inject_work(instance._instances[""]._scheduler, action_digest=action_digest)
163		-
164		- request = bots_pb2.UpdateBotSessionRequest(name=response.name,
165		- bot_session=response)
166		- response = copy.deepcopy(instance.UpdateBotSession(request, context))
167		-
168		- assert response.leases[0].state == LeaseState.PENDING.value
169		- response.leases[0].state = LeaseState.ACTIVE.value
170		-
171		- request = bots_pb2.UpdateBotSessionRequest(name=response.name,
172		- bot_session=response)
173		-
174		- response = copy.deepcopy(instance.UpdateBotSession(request, context))
175		-
176		- response.leases[0].state = LeaseState.COMPLETED.value
177		- response.leases[0].result.Pack(remote_execution_pb2.ActionResult())
178		-
179		- request = bots_pb2.UpdateBotSessionRequest(name=response.name,
180		- bot_session=response)
181		- response = copy.deepcopy(instance.UpdateBotSession(request, context))
182		-
183		- assert len(response.leases) is 0
	153	+ pass
184	154
185	155
186	156	def test_work_rejected_by_bot(bot_session, context, instance):
187		- request = bots_pb2.CreateBotSessionRequest(parent='',
188		- bot_session=bot_session)
189		- # Inject work
190		- action_digest = remote_execution_pb2.Digest(hash='gaff')
191		- _inject_work(instance._instances[""]._scheduler, action_digest=action_digest)
192		-
193		- # Simulated the severed binding between client and server
194		- response = copy.deepcopy(instance.CreateBotSession(request, context))
195		-
196		- # Reject work
197		- assert response.leases[0].state == LeaseState.PENDING.value
198		- response.leases[0].state = LeaseState.COMPLETED.value
199		- request = bots_pb2.UpdateBotSessionRequest(name=response.name,
200		- bot_session=response)
201		-
202		- response = instance.UpdateBotSession(request, context)
203		-
204		- context.set_code.assert_called_once_with(grpc.StatusCode.UNIMPLEMENTED)
205		-
206		-
207		-@pytest.mark.parametrize("state", [LeaseState.LEASE_STATE_UNSPECIFIED, LeaseState.PENDING])
208		-def test_work_out_of_sync_from_pending(state, bot_session, context, instance):
209		- request = bots_pb2.CreateBotSessionRequest(parent='',
210		- bot_session=bot_session)
211		- # Inject work
212		- action_digest = remote_execution_pb2.Digest(hash='gaff')
213		- _inject_work(instance._instances[""]._scheduler, action_digest=action_digest)
214		-
215		- # Simulated the severed binding between client and server
216		- response = copy.deepcopy(instance.CreateBotSession(request, context))
217		-
218		- response.leases[0].state = state.value
219		-
220		- request = bots_pb2.UpdateBotSessionRequest(name=response.name,
221		- bot_session=response)
222		-
223		- response = instance.UpdateBotSession(request, context)
224		-
225		- context.set_code.assert_called_once_with(grpc.StatusCode.DATA_LOSS)
226		-
227		-
228		-@pytest.mark.parametrize("state", [LeaseState.LEASE_STATE_UNSPECIFIED, LeaseState.PENDING])
229		-def test_work_out_of_sync_from_active(state, bot_session, context, instance):
230		- request = bots_pb2.CreateBotSessionRequest(parent='',
231		- bot_session=bot_session)
232		- # Inject work
233		- action_digest = remote_execution_pb2.Digest(hash='gaff')
234		- _inject_work(instance._instances[""]._scheduler, action_digest=action_digest)
235		-
236		- # Simulated the severed binding between client and server
237		- response = copy.deepcopy(instance.CreateBotSession(request, context))
238		-
239		- response.leases[0].state = LeaseState.ACTIVE.value
240		-
241		- request = copy.deepcopy(bots_pb2.UpdateBotSessionRequest(name=response.name,
242		- bot_session=response))
243		-
244		- response = instance.UpdateBotSession(request, context)
245		-
246		- response.leases[0].state = state.value
247		-
248		- request = bots_pb2.UpdateBotSessionRequest(name=response.name,
249		- bot_session=response)
250		-
251		- response = instance.UpdateBotSession(request, context)
252		-
253		- context.set_code.assert_called_once_with(grpc.StatusCode.DATA_LOSS)
254		-
255		-
256		-def test_work_active_to_active(bot_session, context, instance):
257		- request = bots_pb2.CreateBotSessionRequest(parent='',
258		- bot_session=bot_session)
259		- # Inject work
260		- action_digest = remote_execution_pb2.Digest(hash='gaff')
261		- _inject_work(instance._instances[""]._scheduler, action_digest=action_digest)
262		-
263		- # Simulated the severed binding between client and server
264		- response = copy.deepcopy(instance.CreateBotSession(request, context))
265		-
266		- response.leases[0].state = LeaseState.ACTIVE.value
267		-
268		- request = bots_pb2.UpdateBotSessionRequest(name=response.name,
269		- bot_session=response)
270		-
271		- response = instance.UpdateBotSession(request, context)
272		-
273		- assert response.leases[0].state == LeaseState.ACTIVE.value
	157	+ pass
274	158
275	159
276	160	def test_post_bot_event_temp(context, instance):

tests/integration/operations_service.py

---

...	...	@@ -28,10 +28,8 @@ from buildgrid._enums import OperationStage
28	28	from buildgrid._exceptions import InvalidArgumentError
29	29	from buildgrid._protos.build.bazel.remote.execution.v2 import remote_execution_pb2
30	30	from buildgrid._protos.google.longrunning import operations_pb2
31		-from buildgrid._protos.google.rpc import status_pb2
32	31	from buildgrid.server.cas.storage import lru_memory_cache
33	32	from buildgrid.server.controller import ExecutionController
34		-from buildgrid.server.job import LeaseState
35	33	from buildgrid.server.operations import service
36	34	from buildgrid.server.operations.service import OperationsService
37	35	from buildgrid.utils import create_digest
...	...	@@ -166,31 +164,6 @@ def test_list_operations_instance_fail(instance, controller, execute_request, co
166	164	context.set_code.assert_called_once_with(grpc.StatusCode.INVALID_ARGUMENT)
167	165
168	166
169		-def test_list_operations_with_result(instance, controller, execute_request, context):
170		- response_execute = controller.execution_instance.execute(execute_request.action_digest,
171		- execute_request.skip_cache_lookup)
172		-
173		- action_result = remote_execution_pb2.ActionResult()
174		- output_file = remote_execution_pb2.OutputFile(path='unicorn')
175		- action_result.output_files.extend([output_file])
176		-
177		- controller.operations_instance._scheduler.jobs[response_execute.name].create_lease()
178		- controller.operations_instance._scheduler.update_job_lease_state(response_execute.name,
179		- LeaseState.COMPLETED,
180		- lease_status=status_pb2.Status(),
181		- lease_result=_pack_any(action_result))
182		-
183		- request = operations_pb2.ListOperationsRequest(name=instance_name)
184		- response = instance.ListOperations(request, context)
185		-
186		- assert response.operations[0].name.split('/')[-1] == response_execute.name
187		-
188		- execute_response = remote_execution_pb2.ExecuteResponse()
189		- response.operations[0].response.Unpack(execute_response)
190		-
191		- assert execute_response.result.output_files == action_result.output_files
192		-
193		-
194	167	def test_list_operations_empty(instance, context):
195	168	request = operations_pb2.ListOperationsRequest(name=instance_name)
196	169

tests/utils/bots_interface.py

---

	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	+from concurrent import futures
	16	+from contextlib import contextmanager
	17	+import multiprocessing
	18	+import os
	19	+import signal
	20	+import uuid
	21	+
	22	+import grpc
	23	+import pytest_cov
	24	+
	25	+from buildgrid._enums import LeaseState
	26	+from buildgrid._protos.google.devtools.remoteworkers.v1test2 import bots_pb2
	27	+from buildgrid.server.bots import service
	28	+
	29	+
	30	+@contextmanager
	31	+def serve_bots_interface(instances):
	32	+ server = Server(instances)
	33	+ try:
	34	+ yield server
	35	+ finally:
	36	+ server.quit()
	37	+
	38	+
	39	+class Server:
	40	+
	41	+ def __init__(self, instances):
	42	+ self.instances = instances
	43	+
	44	+ self.__queue = multiprocessing.Queue()
	45	+ self.__bot_session_queue = multiprocessing.Queue()
	46	+ self.__message_queue = multiprocessing.Queue()
	47	+ self.__process = multiprocessing.Process(
	48	+ target=Server.serve,
	49	+ args=(self.__queue, self.instances,
	50	+ self.__bot_session_queue, self.__message_queue))
	51	+ self.__process.start()
	52	+
	53	+ self.port = self.__queue.get()
	54	+ self.remote = 'localhost:{}'.format(self.port)
	55	+
	56	+ @classmethod
	57	+ def serve(cls, queue, instances, bot_session_queue, message_queue):
	58	+ pytest_cov.embed.cleanup_on_sigterm()
	59	+
	60	+ # Use max_workers default from Python 3.5+
	61	+ max_workers = (os.cpu_count() or 1) * 5
	62	+ server = grpc.server(futures.ThreadPoolExecutor(max_workers))
	63	+ port = server.add_insecure_port('localhost:0')
	64	+
	65	+ bots_service = service.BotsService(server)
	66	+ for name in instances:
	67	+ bots_interface = BotsInterface(bot_session_queue, message_queue)
	68	+ bots_service.add_instance(name, bots_interface)
	69	+
	70	+ server.start()
	71	+ queue.put(port)
	72	+ signal.pause()
	73	+
	74	+ def get_bot_session(self, timeout=1):
	75	+ bot_session = bots_pb2.BotSession()
	76	+ bot_session.ParseFromString(self.__bot_session_queue.get(timeout=timeout))
	77	+ return bot_session
	78	+
	79	+ def inject_work(self, lease=None, timeout=1):
	80	+ if not lease:
	81	+ lease = bots_pb2.Lease()
	82	+ lease.state = LeaseState.PENDING.value
	83	+
	84	+ lease_string = lease.SerializeToString()
	85	+ self.__message_queue.put(('INJECT_WORK', lease_string))
	86	+
	87	+ def cancel_lease(self, lease_id):
	88	+ self.__message_queue.put(('CANCEL_LEASE', lease_id))
	89	+
	90	+ def quit(self):
	91	+ if self.__process:
	92	+ self.__process.terminate()
	93	+ self.__process.join()
	94	+
	95	+
	96	+class BotsInterface:
	97	+
	98	+ def __init__(self, bot_session_queue, message_queue):
	99	+ self.__bot_session_queue = bot_session_queue
	100	+ self.__message_queue = message_queue
	101	+
	102	+ def register_instance_with_server(self, instance_name, server):
	103	+ server.add_bots_interface(self, instance_name)
	104	+
	105	+ def create_bot_session(self, parent, bot_session):
	106	+ name = "{}/{}".format(parent, str(uuid.uuid4()))
	107	+ bot_session.name = name
	108	+
	109	+ while not self.__message_queue.empty():
	110	+ message = self.__message_queue.get()
	111	+ if message[0] == 'INJECT_WORK':
	112	+ lease_string = message[1]
	113	+ lease = bots_pb2.Lease()
	114	+ lease.ParseFromString(lease_string)
	115	+ bot_session.leases.extend([lease])
	116	+
	117	+ self.__bot_session_queue.put(bot_session.SerializeToString())
	118	+ return bot_session
	119	+
	120	+ def update_bot_session(self, name, bot_session):
	121	+ for lease in bot_session.leases:
	122	+ state = LeaseState(lease.state)
	123	+ if state == LeaseState.COMPLETED:
	124	+ lease.Clear()
	125	+
	126	+ elif state == LeaseState.CANCELLED:
	127	+ lease.Clear()
	128	+
	129	+ while not self.__message_queue.empty():
	130	+ message = self.__message_queue.get()
	131	+
	132	+ if message[0] == 'INJECT_WORK':
	133	+ lease_string = message[1]
	134	+ lease = bots_pb2.Lease()
	135	+ lease.ParseFromString(lease_string)
	136	+ bot_session.leases.extend([lease])
	137	+
	138	+ elif message[0] == 'CANCEL_LEASE':
	139	+ lease_id = message[1]
	140	+ for lease in bot_session.leases:
	141	+ if lease.id == lease_id:
	142	+ lease.state = LeaseState.CANCELLED.value
	143	+
	144	+ self.__bot_session_queue.put(bot_session.SerializeToString())
	145	+ return bot_session

tests/utils/utils.py

---

	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	+import multiprocessing
	16	+import psutil
	17	+
	18	+
	19	+def kill_process_tree(pid):
	20	+ proc = psutil.Process(pid)
	21	+ children = proc.children(recursive=True)
	22	+
	23	+ def kill_proc(p):
	24	+ try:
	25	+ p.kill()
	26	+ except psutil.AccessDenied:
	27	+ # Ignore this error, it can happen with
	28	+ # some setuid bwrap processes.
	29	+ pass
	30	+
	31	+ # Bloody Murder
	32	+ for child in children:
	33	+ kill_proc(child)
	34	+ kill_proc(proc)
	35	+
	36	+
	37	+def run_in_subprocess(function, *arguments, timeout=1):
	38	+ queue = multiprocessing.Queue()
	39	+ # Use subprocess to avoid creation of gRPC threads in main process
	40	+ # See https://github.com/grpc/grpc/blob/master/doc/fork_support.md
	41	+ process = multiprocessing.Process(target=function,
	42	+ args=(queue, *arguments))
	43	+
	44	+ try:
	45	+ process.start()
	46	+ result = queue.get(timeout=timeout)
	47	+ process.join()
	48	+
	49	+ except KeyboardInterrupt:
	50	+ kill_process_tree(process.pid)
	51	+ raise
	52	+
	53	+ return result

—
View it on GitLab.
You're receiving this email because of your account on gitlab.com. If you'd like to receive fewer emails, you can adjust your notification settings.