[Notes] [Git][BuildGrid/buildgrid][finn/bot-refactor] Refactored bots, simplified asyncio

[finnball <gitlab mg gitlab com>]

Title: GitLab

finnball pushed to branch finn/bot-refactor at BuildGrid / buildgrid

Commits:

• 54ce7399
  by finn at 2018-07-30T10:03:30Z

  Refactored bots, simplified asyncio
  

6 changed files:

• app/commands/cmd_bot.py
• buildgrid/bot/bot.py
• buildgrid/bot/bot_interface.py
• + buildgrid/bot/bot_session.py
• buildgrid/server/scheduler.py
• tests/integration/bot_interface.py

Changes:

app/commands/cmd_bot.py

---

...	...	@@ -30,10 +30,12 @@ import os
30	30	import random
31	31	import subprocess
32	32	import tempfile
	33	+import time
33	34
34	35	from pathlib import Path, PurePath
35	36
36	37	from buildgrid.bot import bot
	38	+from buildgrid.bot.bot_session import BotSession, Device, Lease, Worker
37	39	from buildgrid._exceptions import BotError
38	40
39	41	from ..cli import pass_context
...	...	@@ -55,9 +57,18 @@ def cli(context, host, port, number_of_leases, parent, continuous):
55	57
56	58	context.continuous = continuous
57	59	context.channel = grpc.insecure_channel('{}:{}'.format(host, port))
58		- context.number_of_leases = number_of_leases
59	60	context.parent = parent
60	61
	62	+ worker = Worker()
	63	+ worker.add_device(Device().get())
	64	+
	65	+ bot_session = BotSession(parent)
	66	+ bot_session.add_worker(worker.get())
	67	+ for i in range(0, number_of_leases):
	68	+ bot_session.add_lease(Lease().get())
	69	+
	70	+ context.bot_session = bot_session.get()
	71	+
61	72	@cli.command('dummy', short_help='Create a dummy bot session')
62	73	@pass_context
63	74	def dummy(context):
...	...	@@ -65,16 +76,13 @@ def dummy(context):
65	76	Simple dummy client. Creates a session, accepts leases, does fake work and
66	77	updates the server.
67	78	"""
68		-
69		- context.logger.info("Creating a bot session")
70		-
71	79	try:
72		- bot.Bot(work=_work_dummy,
73		- context=context,
74		- channel=context.channel,
75		- parent=context.parent,
76		- number_of_leases=context.number_of_leases,
77		- continuous=context.continuous)
	80	+ b = bot.Bot(bot_session=context.bot_session,
	81	+ channel=context.channel)
	82	+ b.session(context.parent,
	83	+ _work_dummy,
	84	+ context,
	85	+ context.continuous)
78	86
79	87	except KeyboardInterrupt:
80	88	pass
...	...	@@ -88,7 +96,7 @@ def dummy(context):
88	96	@click.option('--port', show_default = True, default=11001)
89	97	@click.option('--remote', show_default = True, default='localhost')
90	98	@pass_context
91		-def _work_buildbox(context, remote, port, server_cert, client_key, client_cert, local_cas, fuse_dir):
	99	+def work_buildbox(context, remote, port, server_cert, client_key, client_cert, local_cas, fuse_dir):
92	100	"""
93	101	Uses BuildBox to run commands.
94	102	"""
...	...	@@ -104,12 +112,14 @@ def _work_buildbox(context, remote, port, server_cert, client_key, client_cert,
104	112	context.fuse_dir = fuse_dir
105	113
106	114	try:
107		- bot.Bot(work=_work_buildbox,
108		- context=context,
109		- channel=context.channel,
110		- parent=context.parent,
111		- number_of_leases=context.number_of_leases,
112		- continuous=context.continuous)
	115	+ b = bot.Bot(work=_work_buildbox,
	116	+ bot_session=context.bot_session,
	117	+ channel=context.channel,
	118	+ parent=context.parent)
	119	+
	120	+ b.session(context.parent,
	121	+ _work_buildbox,
	122	+ context)
113	123
114	124	except KeyboardInterrupt:
115	125	pass

buildgrid/bot/bot.py

---

...	...	@@ -23,160 +23,97 @@ Creates a bot session.
23	23	"""
24	24
25	25	import asyncio
26		-import inspect
	26	+import collections
27	27	import logging
28		-import platform
29		-import queue
30	28	import time
31	29
32		-from google.devtools.remoteworkers.v1test2 import bots_pb2, worker_pb2
33		-
34		-from . import bot_interface
	30	+from . import bot_interface, bot_session
	31	+from .bot_session import BotStatus, LeaseState
35	32	from .._exceptions import BotError
36	33
37		-class Bot(object):
	34	+class Bot:
38	35	"""
39	36	Creates a local BotSession.
40	37	"""
41	38
42		- def __init__(self, work, context, channel, parent, number_of_leases, continuous=True):
43		- if not inspect.iscoroutinefunction(work):
44		- raise BotError("work function must be async")
45		-
46		- print(type(context))
	39	+ UPDATE_PERIOD = 1
47	40
48		- self.interface = bot_interface.BotInterface(channel)
	41	+ def __init__(self, bot_session, channel):
49	42	self.logger = logging.getLogger(__name__)
	43	+ self.interface = bot_interface.BotInterface(channel)
	44	+
	45	+ self._bot_session = bot_session
	46	+ self._work_queue = collections.deque()
	47	+
	48	+ def session(self, parent, work, context, continuous = False):
	49	+ self._bot_session = self.interface.create_bot_session(parent, self._bot_session)
	50	+
	51	+ loop = asyncio.get_event_loop()
	52	+
	53	+ try:
	54	+ tasks = [asyncio.ensure_future(self._update_bot_session()),
	55	+ asyncio.ensure_future(self._worker(work, context))]
	56	+ loop.run_forever()
50	57
51		- self._create_session(parent, number_of_leases)
52		- self._work_queue = queue.Queue(maxsize = number_of_leases)
53		-
54		- while continuous:
55		- ## TODO: Leases independently finish
56		- ## Allow leases to queue finished work independently instead
57		- ## of waiting for all to finish
58		- futures = [self._do_work(work, context, lease) for lease in self._get_work()]
59		- if futures:
60		- loop = asyncio.new_event_loop()
61		- leases_complete, _ = loop.run_until_complete(asyncio.wait(futures))
62		- work_complete = [(lease.result().assignment, lease.result(),) for lease in leases_complete]
63		- self._work_complete(work_complete)
64		- loop.close()
65		- self._update_bot_session()
66		- time.sleep(2)
67		-
68		- @property
69		- def bot_session(self):
70		- ## Read only, shouldn't have to set any of the variables in here
71		- return self._bot_session
72		-
73		- def close_session(self):
74		- self.logger.warning("Session closing not yet implemented")
	58	+ except KeyboardInterrupt:
	59	+ pass
	60	+
	61	+ finally:
	62	+ (task.cancel() for task in tasks)
	63	+ loop.close()
	64	+
	65	+ async def _worker(self, work, context):
	66	+ while True:
	67	+ lease = self._get_work()
	68	+ if lease is not None:
	69	+ lease = await self._do_work(work, context, lease)
	70	+ leases = self._update_from_local(lease)
	71	+ del self._bot_session.leases[:]
	72	+ self._bot_session.leases.extend(leases)
	73	+ else:
	74	+ await asyncio.sleep(1)
75	75
76	76	async def _do_work(self, work, context, lease):
77		- """ Work is done here, work function should be asynchronous
78		- """
79	77	self.logger.info("Work found: {}".format(lease.assignment))
80	78	lease = await work(context=context, lease=lease)
81		- lease.state = bots_pb2.LeaseState.Value('COMPLETED')
	79	+ lease.state = LeaseState.COMPLETED.value
82	80	self.logger.info("Work complete: {}".format(lease.assignment))
83	81	return lease
84	82
85		- def _update_bot_session(self):
86		- """ Should call the server periodically to inform the server the client
87		- has not died.
88		- """
89		- self.logger.debug("Updating bot session")
90		- self._bot_session = self.interface.update_bot_session(self._bot_session)
91		- leases_update = ([self._update_lease(lease) for lease in self._bot_session.leases])
92		- del self._bot_session.leases[:]
93		- self._bot_session.leases.extend(leases_update)
	83	+ async def _update_bot_session(self):
	84	+ while True:
	85	+ """ Calls the server periodically to inform the server the client
	86	+ has not died.
	87	+ """
	88	+ self.logger.debug("Updating bot session")
	89	+ self._bot_session = self.interface.update_bot_session(self._bot_session)
	90	+ leases_update = ([self._update_from_server(lease) for lease in self._bot_session.leases])
	91	+ del self._bot_session.leases[:]
	92	+ self._bot_session.leases.extend(leases_update)
	93	+ await asyncio.sleep(self.UPDATE_PERIOD)
94	94
95	95	def _get_work(self):
96		- while not self._work_queue.empty():
97		- yield self._work_queue.get()
	96	+ if len(self._work_queue) > 0:
	97	+ return self._work_queue.pop()
98	98
99		- def _work_complete(self, leases_complete):
100		- """ Bot updates itself with any completed work.
101		- """
102		- # Should really improve this...
103		- # Maybe add some call back function sentoff work...
104		- leases_active = list(filter(self._lease_active, self._bot_session.leases))
105		- leases_not_active = [lease for lease in self._bot_session.leases if not self._lease_active(lease)]
106		- del self._bot_session.leases[:]
107		- for lease in leases_active:
108		- for lease_tuple in leases_complete:
109		- if lease.assignment == lease_tuple[0]:
110		- leases_not_active.extend([lease_tuple[1]])
111		- self._bot_session.leases.extend(leases_not_active)
112		-
113		- def _update_lease(self, lease):
	99	+ def _update_from_server(self, lease):
114	100	"""
115	101	State machine for any recieved updates to the leases.
116	102	"""
117		- if self._lease_pending(lease):
118		- lease.state = bots_pb2.LeaseState.Value('ACTIVE')
119		- self._work_queue.put(lease)
	103	+ state = lease.state
	104	+ if state == LeaseState.PENDING.value:
	105	+ lease.state = LeaseState.ACTIVE.value
	106	+ self._work_queue.append(lease)
120	107	return lease
121	108
122	109	else:
123	110	return lease
124	111
125		- def _create_session(self, parent, number_of_leases):
126		- self.logger.debug("Creating bot session")
127		- worker = self._create_worker()
128		-
129		- """ Unique bot ID within the farm used to identify this bot
130		- Needs to be human readable.
131		- All prior sessions with bot_id of same ID are invalidated.
132		- If a bot attempts to update an invalid session, it must be rejected and
133		- may be put in quarantine.
134		- """
135		- bot_id = '{}.{}'.format(parent, platform.node())
136		-
137		- leases = [bots_pb2.Lease() for x in range(number_of_leases)]
138		-
139		- bot_session = bots_pb2.BotSession(worker = worker,
140		- status = bots_pb2.BotStatus.Value('OK'),
141		- leases = leases,
142		- bot_id = bot_id)
143		- self._bot_session = self.interface.create_bot_session(parent, bot_session)
144		- self.logger.info("Name: {}, Id: {}".format(self._bot_session.name,
145		- self._bot_session.bot_id))
146		-
147		- def _create_worker(self):
148		- devices = self._create_devices()
149		-
150		- # Contains a list of devices and the connections between them.
151		- worker = worker_pb2.Worker(devices = devices)
152		-
153		- """ Keys supported:
154		- *pool
155		- """
156		- worker.Property.key = "pool"
157		- worker.Property.value = "all"
158		-
159		- return worker
160		-
161		- def _create_devices(self):
162		- """ Creates devices available to the worker
163		- The first device is know as the Primary Device - the revice which
164		- is running a bit and responsible to actually executing commands.
165		- All other devices are known as Attatched Devices and must be controlled
166		- by the Primary Device.
167		- """
168		-
169		- devices = []
170		-
171		- for i in range(0, 1): # Append one device for now
172		- dev = worker_pb2.Device()
173		-
174		- devices.append(dev)
175		-
176		- return devices
177		-
178		- def _lease_pending(self, lease):
179		- return lease.state == bots_pb2.LeaseState.Value('PENDING')
180		-
181		- def _lease_active(self, lease):
182		- return lease.state == bots_pb2.LeaseState.Value('ACTIVE')
	112	+ def _update_from_local(self, lease):
	113	+ updated_leases = []
	114	+ for _lease in self._bot_session.leases:
	115	+ if lease.assignment == _lease.assignment:
	116	+ updated_leases.append(lease)
	117	+ else:
	118	+ updated_leases.append(lease)
	119	+ return updated_leases

buildgrid/bot/bot_interface.py

---

...	...	@@ -39,22 +39,12 @@ class BotInterface(object):
39	39	self._stub = bots_pb2_grpc.BotsStub(channel)
40	40
41	41	def create_bot_session(self, parent, bot_session):
42		- try:
43		- request = bots_pb2.CreateBotSessionRequest(parent = parent,
44		- bot_session = bot_session)
45		- return self._stub.CreateBotSession(request)
46		-
47		- except Exception as e:
48		- self.logger.error(e)
49		- raise BotError(e)
	42	+ request = bots_pb2.CreateBotSessionRequest(parent = parent,
	43	+ bot_session = bot_session)
	44	+ return self._stub.CreateBotSession(request)
50	45
51	46	def update_bot_session(self, bot_session, update_mask = None):
52		- try:
53		- request = bots_pb2.UpdateBotSessionRequest(name = bot_session.name,
54		- bot_session = bot_session,
55		- update_mask = update_mask)
56		- return self._stub.UpdateBotSession(request)
57		-
58		- except Exception as e:
59		- self.logger.error(e)
60		- raise BotError(e)
	47	+ request = bots_pb2.UpdateBotSessionRequest(name = bot_session.name,
	48	+ bot_session = bot_session,
	49	+ update_mask = update_mask)
	50	+ return self._stub.UpdateBotSession(request)

buildgrid/bot/bot_session.py

---

	1	+# Copyright (C) 2018 Codethink Limited
	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	+# Authors:
	16	+# Finn Ball <finn ball codethink co uk>
	17	+
	18	+"""
	19	+Bot Session
	20	+====
	21	+
	22	+Maintains a Bot Session
	23	+"""
	24	+import platform
	25	+import uuid
	26	+
	27	+from enum import Enum
	28	+
	29	+from google.devtools.remoteworkers.v1test2 import bots_pb2, worker_pb2
	30	+
	31	+class BotStatus(Enum):
	32	+ BOT_STATUS_UNSPECIFIED = bots_pb2.BotStatus.Value('BOT_STATUS_UNSPECIFIED')
	33	+ OK = bots_pb2.BotStatus.Value('OK')
	34	+ UNHEALTHY = bots_pb2.BotStatus.Value('UNHEALTHY');
	35	+ HOST_REBOOTING = bots_pb2.BotStatus.Value('HOST_REBOOTING')
	36	+ BOT_TERMINATING = bots_pb2.BotStatus.Value('BOT_TERMINATING')
	37	+
	38	+class LeaseState(Enum):
	39	+ LEASE_STATE_UNSPECIFIED = bots_pb2.LeaseState.Value('LEASE_STATE_UNSPECIFIED')
	40	+ PENDING = bots_pb2.LeaseState.Value('PENDING')
	41	+ ACTIVE = bots_pb2.LeaseState.Value('ACTIVE')
	42	+ COMPLETED = bots_pb2.LeaseState.Value('COMPLETED')
	43	+ CANCELLED = bots_pb2.LeaseState.Value('CANCELLED')
	44	+
	45	+class BotSession:
	46	+ def __init__(self, parent):
	47	+ """ Unique bot ID within the farm used to identify this bot
	48	+ Needs to be human readable.
	49	+ All prior sessions with bot_id of same ID are invalidated.
	50	+ If a bot attempts to update an invalid session, it must be rejected and
	51	+ may be put in quarantine.
	52	+ """
	53	+
	54	+ self._bot_id = '{}.{}'.format(parent, platform.node())
	55	+
	56	+ self._worker = None
	57	+ self._leases = []
	58	+ self._status = BotStatus.OK.value
	59	+
	60	+ def add_worker(self, worker):
	61	+ self._worker = worker
	62	+
	63	+ def add_lease(self, lease):
	64	+ self._leases.append(lease)
	65	+
	66	+ def get(self):
	67	+ return bots_pb2.BotSession(worker=self._worker,
	68	+ status=self._status,
	69	+ leases=self._leases,
	70	+ bot_id=self._bot_id)
	71	+
	72	+ @property
	73	+ def bot_id(self):
	74	+ return self._bot_id
	75	+
	76	+class Worker:
	77	+ def __init__(self, key=None, value=None):
	78	+ """ Keys supported:
	79	+ *pool
	80	+ """
	81	+
	82	+ self._devices = []
	83	+ self._key = key
	84	+ self._value = value
	85	+
	86	+ if key=='pool':
	87	+ self._key = 'Pool'
	88	+ self._value = 'all'
	89	+
	90	+ elif key is not None:
	91	+ raise Exception
	92	+
	93	+ def add_device(self, device):
	94	+ self._devices.append(device)
	95	+
	96	+ def get(self):
	97	+ worker = worker_pb2.Worker(devices=self._devices)
	98	+ worker.Property.key = self._key
	99	+ worker.Property.value = self._value
	100	+ return worker
	101	+
	102	+class Lease:
	103	+ def __init__(self):
	104	+ self.state = LeaseState.LEASE_STATE_UNSPECIFIED
	105	+
	106	+ def get(self):
	107	+ return bots_pb2.Lease(state = self.state.value)
	108	+
	109	+class Device:
	110	+ def __init__(self, key=None, value=None):
	111	+ """ Creates devices available to the worker
	112	+ The first device is know as the Primary Device - the revice which
	113	+ is running a bit and responsible to actually executing commands.
	114	+ All other devices are known as Attatched Devices and must be controlled
	115	+ by the Primary Device.
	116	+ """
	117	+ if key == 'docker':
	118	+ if value not in ('True', 'False'):
	119	+ raise Exception
	120	+
	121	+ elif key == 'os':
	122	+ pass
	123	+
	124	+ elif key is not None:
	125	+ raise Exception
	126	+
	127	+ name = str(uuid.uuid4())
	128	+ device = worker_pb2.Device(handle=name)
	129	+ device.Property.key = key
	130	+ device.Property.value = value
	131	+ self._device = device
	132	+
	133	+ def get(self):
	134	+ return self._device

buildgrid/server/scheduler.py

---

...	...	@@ -41,17 +41,19 @@ class Scheduler():
41	41	self.queue.append(job)
42	42
43	43	def retry_job(self, name):
44		- job = self.jobs[name]
	44	+ job = self.jobs.get(name)
45	45
46		- if job.n_tries >= self.MAX_N_TRIES:
47		- # TODO: Decide what to do with these jobs
48		- job.execute_stage = ExecuteStage.COMPLETED
49		- else:
50		- job.execute_stage = ExecuteStage.QUEUED
51		- job.n_tries += 1
	46	+ if job is not None:
	47	+
	48	+ if job.n_tries >= self.MAX_N_TRIES:
	49	+ # TODO: Decide what to do with these jobs
	50	+ job.execute_stage = ExecuteStage.COMPLETED
	51	+ else:
	52	+ job.execute_stage = ExecuteStage.QUEUED
	53	+ job.n_tries += 1
52	54	self.queue.appendleft(job)
53	55
54		- self.jobs[name] = job
	56	+ self.jobs[name] = job
55	57
56	58	def create_job(self):
57	59	if len(self.queue) > 0:

tests/integration/bot_interface.py

---

...	...	@@ -22,6 +22,8 @@ import mock
22	22	import pytest
23	23	import uuid
24	24
	25	+from google.devtools.remoteworkers.v1test2 import bots_pb2, worker_pb2
	26	+
25	27	from buildgrid.bot import bot, bot_interface
26	28
27	29	async def _work_dummy(context, lease):
...	...	@@ -39,16 +41,9 @@ def context():
39	41	# GRPC context
40	42	@pytest.fixture
41	43	def channel():
42		- yield mock.MagicMock(spec = grpc.insecure_channel(''))
	44	+ yield mock.MagicMock(spec = grpc.insecure_channel)
43	45
44		-@pytest.fixture
45		-def instance(channel):
46		- yield bot.Bot(work=_work_dummy,
47		- context=ContextMock(),
48		- channel=channel,
49		- parent='rach',
50		- number_of_leases=1,
51		- continuous=False)
52		-
53		-def test_create_job(instance):
54		- instance.bot_session()
	46	+@mock.patch.object(bot.bot_interface, 'bots_pb2', autospec = True)
	47	+@mock.patch.object(bot.bot_interface, 'bots_pb2_grpc', autospec = True)
	48	+def test_me(mock_pb2, mock_pb2_grpc, channel, context):
	49	+ pass

—
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.