[Notes] [Git][BuildGrid/buildgrid][master] 19 commits: Scheduler can now cancel leases.

[finn <gitlab mg gitlab com>]

Title: GitLab

finn pushed to branch master at BuildGrid / buildgrid

Commits:

• c917480a
  by Finn at 2018-11-23T11:18:49Z

  Scheduler can now cancel leases.
  

• fd0c1534
  by Finn at 2018-11-23T11:18:49Z

  BotsInterface now handles cancelled leases.
  
  Removed state machine. Will help with #66.
  

• 1e58b538
  by Finn at 2018-11-23T11:18:49Z

  Added HardwareInterface class.
  
  This will describe and configure hardware.
  

• b7d62408
  by Finn at 2018-11-23T11:18:49Z

  Add Device class.
  

• 17db2b8c
  by Finn at 2018-11-23T11:18:49Z

  Add Worker class.
  

• 694c42ab
  by Finn at 2018-11-23T11:18:49Z

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

• 87411369
  by Finn at 2018-11-23T11:41:53Z

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

• b031d8de
  by Finn at 2018-11-23T11:55:31Z

  Added TenantManager class.
  
  This will manage the Tenant objects.
  

• 663f1fd8
  by Finn at 2018-11-23T11:55:31Z

  Modified BotSession to use new TenantManager class.
  

• e8972884
  by Finn at 2018-11-23T15:02:39Z

  buildgrid/bot/bot.py: Close tasks when KeyboardInterupt.
  

• d1cef783
  by Finn at 2018-11-23T15:02:44Z

  Added exception handling to BotInterface.
  
  Also Moved bot_interface.py to interface.py.
  

• b0a1c9c9
  by Finn at 2018-11-23T15:11:01Z

  Moved work function and context argument order.
  
  Also moved the initialisation of these variables to
  the `__init__` of `session.py`.
  

• 1dd8291a
  by Finn at 2018-11-23T15:11:01Z

  Subprocess functions now in separate utils file.
  

• 63e2feaf
  by Finn at 2018-11-23T15:11:01Z

  Updated imports for function `run_in_subprocess`
  

• 24f4117f
  by Finn at 2018-11-23T15:11:01Z

  tests/utils/bots_interface.py: Added new bots_interface mock.
  

• cfdb72d9
  by Finn at 2018-11-23T15:11:01Z

  tests/integration/bot_session.py: Updated to use new subprocess tests.
  
  Also adds cancellation tests.
  

• 6b779d3a
  by Finn at 2018-11-23T15:11:01Z

  tests/integration/bots_service.py: Removed redundant tests.
  
  Mostly due to removing the state machine in the BotsService
  

• f9261257
  by Finn at 2018-11-23T15:11:01Z

  tests/integration/operations_service.py: Removing redundant test.
  

• 96eef9a5
  by Finn at 2018-11-23T15:11:01Z

  Added unittests for hardware interface.
  

27 changed files:

• buildgrid/_app/bots/buildbox.py
• buildgrid/_app/bots/dummy.py
• buildgrid/_app/bots/host.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/job.py
• buildgrid/server/scheduler.py
• tests/cas/test_client.py
• tests/cas/test_storage.py
• tests/integration/bot_session.py
• tests/integration/bots_service.py
• + tests/integration/hardware_interface.py
• tests/integration/operations_service.py
• tests/server_instance.py
• + tests/utils/bots_interface.py
• tests/utils/cas.py
• + tests/utils/utils.py

Changes:

buildgrid/_app/bots/buildbox.py

---

...	...	@@ -25,7 +25,7 @@ from buildgrid.settings import HASH_LENGTH
25	25	from buildgrid.utils import read_file, write_file
26	26
27	27
28		-def work_buildbox(context, lease):
	28	+def work_buildbox(lease, context, event):
29	29	"""Executes a lease for a build action, using buildbox.
30	30	"""
31	31	local_cas_directory = context.local_cas

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/bots/host.py

---

...	...	@@ -23,7 +23,7 @@ from buildgrid._protos.build.bazel.remote.execution.v2 import remote_execution_p
23	23	from buildgrid.utils import get_hostname, output_file_maker, output_directory_maker
24	24
25	25
26		-def work_host_tools(context, lease):
	26	+def work_host_tools(lease, context, event):
27	27	"""Executes a lease for a build action, using host tools.
28	28	"""
29	29	instance_name = context.parent

buildgrid/_app/commands/cmd_bot.py

---

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

buildgrid/bot/bot.py

---

...	...	@@ -13,46 +13,52 @@
13	13	# limitations under the License.
14	14
15	15
16		-"""
17		-Bot
18		-====
19		-
20		-Creates a bot session.
21		-"""
22		-
23	16	import asyncio
24	17	import logging
25	18
26	19
27	20	class Bot:
28		- """
29		- Creates a local BotSession.
30		- """
	21	+ """Creates a local BotSession."""
31	22
32	23	def __init__(self, bot_session, update_period=1):
	24	+ """
	25	+ """
33	26	self.__logger = logging.getLogger(__name__)
34	27
35		- self._bot_session = bot_session
36		- self._update_period = update_period
	28	+ self.__bot_session = bot_session
	29	+ self.__update_period = update_period
37	30
38		- def session(self, work, context):
39		- loop = asyncio.get_event_loop()
	31	+ self.__loop = None
40	32
41		- self._bot_session.create_bot_session(work, context)
	33	+ def session(self):
	34	+ """Will create a session and periodically call the server."""
	35	+
	36	+ self.__loop = asyncio.get_event_loop()
	37	+ self.__bot_session.create_bot_session()
42	38
43	39	try:
44		- task = asyncio.ensure_future(self._update_bot_session())
45		- loop.run_forever()
	40	+ task = asyncio.ensure_future(self.__update_bot_session())
	41	+ self.__loop.run_until_complete(task)
	42	+
46	43	except KeyboardInterrupt:
47	44	pass
48		- finally:
49		- task.cancel()
50		- loop.close()
51	45
52		- async def _update_bot_session(self):
53		- """
54		- Calls the server periodically to inform the server the client has not died.
55		- """
56		- while True:
57		- self._bot_session.update_bot_session()
58		- await asyncio.sleep(self._update_period)
	46	+ self.__kill_everyone()
	47	+ self.__logger.info("Bot shutdown.")
	48	+
	49	+ async def __update_bot_session(self):
	50	+ """Calls the server periodically to inform the server the client has not died."""
	51	+ try:
	52	+ while True:
	53	+ self.__bot_session.update_bot_session()
	54	+ await asyncio.sleep(self.__update_period)
	55	+
	56	+ except asyncio.CancelledError:
	57	+ pass
	58	+
	59	+ def __kill_everyone(self):
	60	+ """Cancels and waits for them to stop."""
	61	+ self.__logger.info("Cancelling remaining tasks...")
	62	+ for task in asyncio.Task.all_tasks():
	63	+ task.cancel()
	64	+ self.__loop.run_until_complete(task)

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

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	+ Args:
	38	+ lease (:class:`Lease`) : A new lease.
	39	+
	40	+ Raises:
	41	+ ValueError: If `lease` is not in a `PENDING` state.
	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 id=[{}]".format(lease_id))
	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,70 +77,54 @@ 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)
	80	+ self._check_assigned_leases(bot_session)
	81	+
	82	+ for lease in bot_session.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
82	94
83		- leases = filter(None, [self.check_states(lease) for lease in bot_session.leases])
84		-
85		- del bot_session.leases[:]
86		- bot_session.leases.extend(leases)
87		-
	95	+ def _request_leases(self, bot_session):
88	96	# TODO: Send worker capabilities to the scheduler!
	97	+ # Only send one lease at a time currently.
89	98	if not bot_session.leases:
90	99	leases = self._scheduler.request_job_leases({})
91	100	if leases:
	101	+ for lease in leases:
	102	+ self._assigned_leases[bot_session.name].add(lease.id)
92	103	bot_session.leases.extend(leases)
93	104
94		- self._bot_sessions[name] = bot_session
95		- return bot_session
	105	+ def _check_lease_state(self, lease):
	106	+ # careful here
	107	+ # should store bot name in scheduler
	108	+ lease_state = LeaseState(lease.state)
	109	+
	110	+ # Lease has replied with cancelled, remove
	111	+ if lease_state == LeaseState.CANCELLED:
	112	+ return None
96	113
97		- def check_states(self, client_lease):
98		- """ Edge detector for states
99		- """
100		- # TODO: Handle cancelled states
101	114	try:
102		- server_lease = self._scheduler.get_job_lease(client_lease.id)
	115	+ if self._scheduler.get_job_lease_cancelled(lease.id):
	116	+ lease.state.CopyFrom(LeaseState.CANCELLED.value)
	117	+ return lease
103	118	except KeyError:
104		- raise InvalidArgumentError("Lease not found on server: [{}]".format(client_lease))
105		-
106		- server_state = LeaseState(server_lease.state)
107		- client_state = LeaseState(client_lease.state)
108		-
109		- if server_state == LeaseState.PENDING:
	119	+ # Job does not exist, remove from bot.
	120	+ return None
110	121
111		- if client_state == LeaseState.ACTIVE:
112		- self._scheduler.update_job_lease_state(client_lease.id,
113		- LeaseState.ACTIVE)
114		- elif client_state == LeaseState.COMPLETED:
115		- # TODO: Lease was rejected
116		- raise NotImplementedError("'Not Accepted' is unsupported")
117		- else:
118		- raise OutOfSyncError("Server lease: [{}]. Client lease: [{}]".format(server_lease, client_lease))
	122	+ self._scheduler.update_job_lease(lease)
119	123
120		- elif server_state == LeaseState.ACTIVE:
	124	+ if lease_state == LeaseState.COMPLETED:
	125	+ return None
121	126
122		- if client_state == LeaseState.ACTIVE:
123		- pass
124		-
125		- elif client_state == LeaseState.COMPLETED:
126		- self._scheduler.update_job_lease_state(client_lease.id,
127		- LeaseState.COMPLETED,
128		- lease_status=client_lease.status,
129		- lease_result=client_lease.result)
130		- return None
131		-
132		- else:
133		- raise OutOfSyncError("Server lease: [{}]. Client lease: [{}]".format(server_lease, client_lease))
134		-
135		- elif server_state == LeaseState.COMPLETED:
136		- raise OutOfSyncError("Server lease: [{}]. Client lease: [{}]".format(server_lease, client_lease))
137		-
138		- elif server_state == LeaseState.CANCELLED:
139		- raise NotImplementedError("Cancelled states not supported yet")
140		-
141		- else:
142		- # Sould never get here
143		- raise OutOfSyncError("State now allowed: {}".format(server_state))
144		-
145		- return client_lease
	127	+ return lease
146	128
147	129	def _check_bot_ids(self, bot_id, name=None):
148	130	""" Checks the ID and the name of the bot.
...	...	@@ -164,6 +146,19 @@ class BotsInterface:
164	146	'Bot id already registered. ID sent: [{}].'
165	147	'Id registered: [{}] with name: [{}]'.format(bot_id, _bot_id, _name))
166	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	+
167	162	def _close_bot_session(self, name):
168	163	""" Before removing the session, close any leases and
169	164	requeue with high priority.
...	...	@@ -174,10 +169,9 @@ class BotsInterface:
174	169	raise InvalidArgumentError("Bot id does not exist: [{}]".format(name))
175	170
176	171	self.__logger.debug("Attempting to close [%s] with name: [%s]", bot_id, name)
177		- for lease in self._bot_sessions[name].leases:
178		- if lease.state != LeaseState.COMPLETED.value:
179		- # TODO: Be wary here, may need to handle rejected leases in future
180		- 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)
181	175
182	176	self.__logger.debug("Closing bot session: [%s]", name)
183	177	self._bot_ids.pop(name)

buildgrid/server/job.py

---

...	...	@@ -98,6 +98,10 @@ class Job:
98	98	else:
99	99	return None
100	100
	101	+ @property
	102	+ def lease_cancelled(self):
	103	+ return self.__lease_cancelled
	104	+
101	105	@property
102	106	def n_tries(self):
103	107	return self._n_tries

buildgrid/server/scheduler.py

---

...	...	@@ -77,7 +77,8 @@ class Scheduler:
77	77	# TODO: Mark these jobs as done
78	78	else:
79	79	job.update_operation_stage(OperationStage.QUEUED)
80		- self.queue.appendleft(job)
	80	+ job.update_lease_state(LeaseState.PENDING)
	81	+ self.queue.append(job)
81	82
82	83	def list_jobs(self):
83	84	return self.jobs.values()
...	...	@@ -94,15 +95,19 @@ class Scheduler:
94	95	return []
95	96
96	97	job = self.queue.popleft()
97		- # For now, one lease at a time:
98		- lease = job.create_lease()
	98	+
	99	+ lease = job.lease
	100	+
	101	+ if not lease:
	102	+ # For now, one lease at a time:
	103	+ lease = job.create_lease()
99	104
100	105	if lease:
101	106	return [lease]
102	107
103	108	return None
104	109
105		- def update_job_lease_state(self, job_name, lease_state, lease_status=None, lease_result=None):
	110	+ def update_job_lease(self, lease):
106	111	"""Requests a state transition for a job's current :class:Lease.
107	112
108	113	Args:
...	...	@@ -113,7 +118,9 @@ class Scheduler:
113	118	lease_result (google.protobuf.Any): the lease execution result, only
114	119	required if `lease_state` is `COMPLETED`.
115	120	"""
116		- job = self.jobs[job_name]
	121	+
	122	+ job = self.jobs[lease.id]
	123	+ lease_state = LeaseState(lease.state)
117	124
118	125	if lease_state == LeaseState.PENDING:
119	126	job.update_lease_state(LeaseState.PENDING)
...	...	@@ -125,7 +132,7 @@ class Scheduler:
125	132
126	133	elif lease_state == LeaseState.COMPLETED:
127	134	job.update_lease_state(LeaseState.COMPLETED,
128		- status=lease_status, result=lease_result)
	135	+ status=lease.status, result=lease.result)
129	136
130	137	if self._action_cache is not None and not job.do_not_cache:
131	138	self._action_cache.update_action_result(job.action_digest, job.action_result)
...	...	@@ -136,6 +143,10 @@ class Scheduler:
136	143	"""Returns the lease associated to job, if any have been emitted yet."""
137	144	return self.jobs[job_name].lease
138	145
	146	+ def get_job_lease_cancelled(self, job_name):
	147	+ """Returns true if the lease is cancelled"""
	148	+ return self.jobs[job_name].lease_cancelled
	149	+
139	150	def get_job_operation(self, job_name):
140	151	"""Returns the operation associated to job."""
141	152	return self.jobs[job_name].operation

tests/cas/test_client.py

---

...	...	@@ -26,7 +26,8 @@ from buildgrid.client.cas import download, upload
26	26	from buildgrid._protos.build.bazel.remote.execution.v2 import remote_execution_pb2
27	27	from buildgrid.utils import create_digest
28	28
29		-from ..utils.cas import serve_cas, run_in_subprocess
	29	+from ..utils.cas import serve_cas
	30	+from ..utils.utils import run_in_subprocess
30	31
31	32
32	33	INTANCES = ['', 'instance']

tests/cas/test_storage.py

---

...	...	@@ -32,7 +32,8 @@ from buildgrid.server.cas.storage.s3 import S3Storage
32	32	from buildgrid.server.cas.storage.with_cache import WithCacheStorage
33	33	from buildgrid.settings import HASH
34	34
35		-from ..utils.cas import serve_cas, run_in_subprocess
	35	+from ..utils.cas import serve_cas
	36	+from ..utils.utils import run_in_subprocess
36	37
37	38
38	39	BLOBS = [(b'abc', b'defg', b'hijk', b'')]

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
...	...	@@ -150,129 +149,6 @@ def test_update_leases_with_work(bot_session, context, instance):
150	149	assert response_action == action_digest
151	150
152	151
153		-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
184		-
185		-
186		-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
274		-
275		-
276	152	def test_post_bot_event_temp(context, instance):
277	153	request = bots_pb2.PostBotEventTempRequest()
278	154	instance.PostBotEventTemp(request, context)

tests/integration/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	+# pylint: disable=redefined-outer-name
	16	+
	17	+
	18	+import pytest
	19	+
	20	+from buildgrid._exceptions import FailedPreconditionError
	21	+from buildgrid.bot.hardware.interface import HardwareInterface
	22	+from buildgrid.bot.hardware.device import Device
	23	+from buildgrid.bot.hardware.worker import Worker
	24	+
	25	+
	26	+CONFIGURATIONS_WORKER = [{'DockerImage': ['Blade']}, {'DockerImage': ['Sheep', 'Snake']}]
	27	+PROPERTIES_WORKER = [{'pool': ['Blade']}, {'pool': ['Sheep', 'Snake']}]
	28	+PROPERTIES_DEVICE = [{'os': ['Blade']}, {'has-docker': ['Sheep', 'Snake']}]
	29	+
	30	+
	31	+def test_create_hardware():
	32	+ worker = Worker()
	33	+ interface = HardwareInterface(worker)
	34	+
	35	+ device0 = Device()
	36	+ worker.add_device(device0)
	37	+ protobuf_worker = interface.get_worker_pb2()
	38	+ assert len(protobuf_worker.devices) == 1
	39	+
	40	+ worker.add_device(Device())
	41	+ protobuf_worker = interface.get_worker_pb2()
	42	+ assert len(protobuf_worker.devices) == 2
	43	+
	44	+ assert protobuf_worker.devices[0].handle != protobuf_worker.devices[1].handle
	45	+ assert device0.name == protobuf_worker.devices[0].handle
	46	+
	47	+
	48	+@pytest.mark.parametrize('config', CONFIGURATIONS_WORKER)
	49	+def test_worker_config(config):
	50	+ worker = Worker(configs=config)
	51	+ protobuf_worker = worker.get_pb2()
	52	+
	53	+ proto_cfg = {}
	54	+ for cfg in protobuf_worker.configs:
	55	+ k = cfg.key
	56	+ v = cfg.value
	57	+
	58	+ proto_cfg_values = proto_cfg.get(k)
	59	+ if not proto_cfg_values:
	60	+ proto_cfg[k] = [v]
	61	+ else:
	62	+ proto_cfg_values.append(v)
	63	+
	64	+ assert config == proto_cfg
	65	+ assert worker.configs == config
	66	+
	67	+
	68	+@pytest.mark.parametrize('prop', PROPERTIES_WORKER)
	69	+def test_worker_property(prop):
	70	+ worker = Worker(properties=prop)
	71	+ protobuf_worker = worker.get_pb2()
	72	+
	73	+ proto_prop = {}
	74	+ for p in protobuf_worker.properties:
	75	+ k = p.key
	76	+ v = p.value
	77	+
	78	+ proto_prop_values = proto_prop.get(k)
	79	+ if not proto_prop_values:
	80	+ proto_prop[k] = [v]
	81	+ else:
	82	+ proto_prop_values.append(v)
	83	+
	84	+ assert prop == proto_prop
	85	+ assert worker.properties == prop
	86	+
	87	+
	88	+@pytest.mark.parametrize('prop', PROPERTIES_DEVICE)
	89	+def test_device_property(prop):
	90	+ device = Device(properties=prop)
	91	+ protobuf_device = device.get_pb2()
	92	+
	93	+ proto_prop = {}
	94	+ for p in protobuf_device.properties:
	95	+ k = p.key
	96	+ v = p.value
	97	+
	98	+ proto_prop_values = proto_prop.get(k)
	99	+ if not proto_prop_values:
	100	+ proto_prop[k] = [v]
	101	+ else:
	102	+ proto_prop_values.append(v)
	103	+
	104	+ assert prop == proto_prop
	105	+ assert device.properties == prop
	106	+
	107	+
	108	+@pytest.mark.parametrize('config', [{'piano': ['Blade']}])
	109	+def test_worker_config_fail(config):
	110	+ with pytest.raises(KeyError):
	111	+ Worker(configs=config)
	112	+
	113	+
	114	+@pytest.mark.parametrize('prop', [{'piano': ['Blade']}])
	115	+def test_worker_property_fail(prop):
	116	+ with pytest.raises(KeyError):
	117	+ Worker(properties=prop)
	118	+
	119	+
	120	+@pytest.mark.parametrize('prop', [{'piano': ['Blade']}])
	121	+def test_device_property_fail(prop):
	122	+ with pytest.raises(KeyError):
	123	+ Device(properties=prop)
	124	+
	125	+
	126	+@pytest.mark.parametrize('requirements', CONFIGURATIONS_WORKER)
	127	+def test_configure_hardware(requirements):
	128	+ hardware_interface = HardwareInterface(Worker(configs=requirements))
	129	+
	130	+ worker_requirements = Worker(configs=requirements)
	131	+
	132	+ hardware_interface.configure_hardware(worker_requirements.get_pb2())
	133	+
	134	+
	135	+@pytest.mark.parametrize('requirements', CONFIGURATIONS_WORKER)
	136	+def test_configure_hardware_fail(requirements):
	137	+ hardware_interface = HardwareInterface(Worker())
	138	+
	139	+ worker_requirements = Worker(configs=requirements)
	140	+
	141	+ with pytest.raises(FailedPreconditionError):
	142	+ hardware_interface.configure_hardware(worker_requirements.get_pb2())

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/server_instance.py

---

...	...	@@ -22,7 +22,7 @@ from buildgrid.server.operations.service import OperationsService
22	22	from buildgrid.server.bots.service import BotsService
23	23	from buildgrid.server.referencestorage.service import ReferenceStorageService
24	24
25		-from .utils.cas import run_in_subprocess
	25	+from .utils.utils import run_in_subprocess
26	26
27	27
28	28	config = """

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	+# 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	+class Server:
	42	+
	43	+ def __init__(self, instances):
	44	+ self.instances = instances
	45	+
	46	+ self.__queue = multiprocessing.Queue()
	47	+ # Queue purely for bot session updates
	48	+ self.__bot_session_queue = multiprocessing.Queue()
	49	+ # Queue to send messages to subprocess
	50	+ self.__message_queue = multiprocessing.Queue()
	51	+ self.__process = multiprocessing.Process(
	52	+ target=Server.serve,
	53	+ args=(self.__queue, self.instances,
	54	+ self.__bot_session_queue, self.__message_queue))
	55	+ self.__process.start()
	56	+
	57	+ self.port = self.__queue.get()
	58	+ self.remote = 'localhost:{}'.format(self.port)
	59	+
	60	+ @staticmethod
	61	+ def serve(queue, instances, bot_session_queue, message_queue):
	62	+ pytest_cov.embed.cleanup_on_sigterm()
	63	+
	64	+ # Use max_workers default from Python 3.5+
	65	+ max_workers = (os.cpu_count() or 1) * 5
	66	+ server = grpc.server(futures.ThreadPoolExecutor(max_workers))
	67	+ port = server.add_insecure_port('localhost:0')
	68	+
	69	+ bots_service = service.BotsService(server)
	70	+ for name in instances:
	71	+ bots_interface = BotsInterface(bot_session_queue, message_queue)
	72	+ bots_service.add_instance(name, bots_interface)
	73	+
	74	+ server.start()
	75	+ queue.put(port)
	76	+ signal.pause()
	77	+
	78	+ def get_bot_session(self, timeout=1):
	79	+ bot_session = bots_pb2.BotSession()
	80	+ bot_session.ParseFromString(self.__bot_session_queue.get(timeout=timeout))
	81	+ return bot_session
	82	+
	83	+ # Injects leases
	84	+ def inject_work(self, lease=None, timeout=1):
	85	+ if not lease:
	86	+ lease = bots_pb2.Lease()
	87	+ lease.state = LeaseState.PENDING.value
	88	+
	89	+ lease_string = lease.SerializeToString()
	90	+ self.__message_queue.put(('INJECT_WORK', lease_string))
	91	+
	92	+ # Triggers a cancellation of a lease from server
	93	+ def cancel_lease(self, lease_id):
	94	+ self.__message_queue.put(('CANCEL_LEASE', lease_id))
	95	+
	96	+ def quit(self):
	97	+ if self.__process:
	98	+ self.__process.terminate()
	99	+ self.__process.join()
	100	+
	101	+
	102	+class BotsInterface:
	103	+
	104	+ def __init__(self, bot_session_queue, message_queue):
	105	+ self.__bot_session_queue = bot_session_queue
	106	+ self.__message_queue = message_queue
	107	+
	108	+ def register_instance_with_server(self, instance_name, server):
	109	+ server.add_bots_interface(self, instance_name)
	110	+
	111	+ def create_bot_session(self, parent, bot_session):
	112	+ name = "{}/{}".format(parent, str(uuid.uuid4()))
	113	+ bot_session.name = name
	114	+
	115	+ while not self.__message_queue.empty():
	116	+ message = self.__message_queue.get()
	117	+ if message[0] == 'INJECT_WORK':
	118	+ lease_string = message[1]
	119	+ lease = bots_pb2.Lease()
	120	+ lease.ParseFromString(lease_string)
	121	+ bot_session.leases.extend([lease])
	122	+
	123	+ self.__bot_session_queue.put(bot_session.SerializeToString())
	124	+ return bot_session
	125	+
	126	+ def update_bot_session(self, name, bot_session):
	127	+ for lease in bot_session.leases:
	128	+ state = LeaseState(lease.state)
	129	+ if state == LeaseState.COMPLETED:
	130	+ lease.Clear()
	131	+
	132	+ elif state == LeaseState.CANCELLED:
	133	+ lease.Clear()
	134	+
	135	+ while not self.__message_queue.empty():
	136	+ message = self.__message_queue.get()
	137	+
	138	+ if message[0] == 'INJECT_WORK':
	139	+ lease_string = message[1]
	140	+ lease = bots_pb2.Lease()
	141	+ lease.ParseFromString(lease_string)
	142	+ bot_session.leases.extend([lease])
	143	+
	144	+ elif message[0] == 'CANCEL_LEASE':
	145	+ lease_id = message[1]
	146	+ for lease in bot_session.leases:
	147	+ if lease.id == lease_id:
	148	+ lease.state = LeaseState.CANCELLED.value
	149	+
	150	+ self.__bot_session_queue.put(bot_session.SerializeToString())
	151	+ return bot_session

tests/utils/cas.py

---

...	...	@@ -21,7 +21,6 @@ import signal
21	21	import tempfile
22	22
23	23	import grpc
24		-import psutil
25	24	import pytest_cov
26	25
27	26	from buildgrid._protos.build.bazel.remote.execution.v2 import remote_execution_pb2
...	...	@@ -42,43 +41,6 @@ def serve_cas(instances):
42	41	server.quit()
43	42
44	43
45		-def kill_process_tree(pid):
46		- proc = psutil.Process(pid)
47		- children = proc.children(recursive=True)
48		-
49		- def kill_proc(p):
50		- try:
51		- p.kill()
52		- except psutil.AccessDenied:
53		- # Ignore this error, it can happen with
54		- # some setuid bwrap processes.
55		- pass
56		-
57		- # Bloody Murder
58		- for child in children:
59		- kill_proc(child)
60		- kill_proc(proc)
61		-
62		-
63		-def run_in_subprocess(function, *arguments):
64		- queue = multiprocessing.Queue()
65		- # Use subprocess to avoid creation of gRPC threads in main process
66		- # See https://github.com/grpc/grpc/blob/master/doc/fork_support.md
67		- process = multiprocessing.Process(target=function,
68		- args=(queue, *arguments))
69		-
70		- try:
71		- process.start()
72		-
73		- result = queue.get()
74		- process.join()
75		- except KeyboardInterrupt:
76		- kill_process_tree(process.pid)
77		- raise
78		-
79		- return result
80		-
81		-
82	44	class Server:
83	45
84	46	def __init__(self, instances):

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.