[Notes] [Git][BuildGrid/buildgrid][finn/48-cancellation-leases] Working commit

[finn <gitlab mg gitlab com>]

Title: GitLab

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

Commits:

• 6f70c9c0
  by Finn at 2018-11-05T17:46:08Z

  Working commit
  

9 changed files:

• buildgrid/_app/commands/cmd_bot.py
• buildgrid/bot/bot.py
• buildgrid/bot/hardware/device.py
• + buildgrid/bot/hardware/interface.py
• buildgrid/bot/hardware/worker.py
• buildgrid/bot/session.py
• + buildgrid/bot/tenant.py
• + buildgrid/bot/tenantmanager.py
• buildgrid/server/bots/instance.py

Changes:

buildgrid/_app/commands/cmd_bot.py

---

...	...	@@ -28,8 +28,11 @@ from urllib.parse import urlparse
28	28	import click
29	29	import grpc
30	30
31		-from buildgrid.bot import bot, bot_interface
32		-from buildgrid.bot.bot_session import BotSession, Device, Worker
	31	+from buildgrid.bot import bot, interface, session
	32	+from buildgrid.bot.hardware.interface import HardwareInterface
	33	+from buildgrid.bot.hardware.device import Device
	34	+from buildgrid.bot.hardware.worker import Worker
	35	+
33	36
34	37	from ..bots import buildbox, dummy, host
35	38	from ..cli import pass_context
...	...	@@ -123,13 +126,14 @@ def cli(context, parent, update_period, remote, client_key, client_cert, server_
123	126	context.logger = logging.getLogger(__name__)
124	127	context.logger.debug("Starting for remote {}".format(context.remote))
125	128
126		- interface = bot_interface.BotInterface(context.channel)
	129	+ bot_interface = interface.BotInterface(context.channel)
127	130
128	131	worker = Worker()
129	132	worker.add_device(Device())
130	133
131		- bot_session = BotSession(parent, interface)
132		- bot_session.add_worker(worker)
	134	+ hardware_interface = HardwareInterface(worker)
	135	+
	136	+ bot_session = session.BotSession(parent, bot_interface, hardware_interface)
133	137
134	138	context.bot_session = bot_session
135	139
...	...	@@ -142,8 +146,7 @@ def run_dummy(context):
142	146	"""
143	147	try:
144	148	b = bot.Bot(context.bot_session, context.update_period)
145		- b.session(dummy.work_dummy,
146		- context)
	149	+ b.session(dummy.work_dummy, context)
147	150	except KeyboardInterrupt:
148	151	pass
149	152

buildgrid/bot/bot.py

---

...	...	@@ -17,7 +17,7 @@
17	17	Bot
18	18	====
19	19
20		-Creates a bot session.
	20	+Creates a bot session and sends updates to the server.
21	21	"""
22	22
23	23	import asyncio
...	...	@@ -45,6 +45,7 @@ class Bot:
45	45	loop.run_forever()
46	46	except KeyboardInterrupt:
47	47	pass
	48	+
48	49	finally:
49	50	task.cancel()
50	51	loop.close()

buildgrid/bot/hardware/device.py

---

...	...	@@ -13,6 +13,17 @@
13	13	# limitations under the License.
14	14
15	15
	16	+"""
	17	+Device
	18	+======
	19	+
	20	+A device.
	21	+"""
	22	+
	23	+
	24	+import uuid
	25	+from buildgrid._protos.google.devtools.remoteworkers.v1test2 import worker_pb2
	26	+
16	27	class Device:
17	28
18	29	def __init__(self, properties=None):
...	...	@@ -21,23 +32,18 @@ class Device:
21	32	is running a bit and responsible to actually executing commands.
22	33	All other devices are known as Attatched Devices and must be controlled
23	34	by the Primary Device.
	35	+
	36	+ properties (list(dict(string : string))) : Properties of device. Keys may
	37	+ repeated.
24	38	"""
25	39
26		- self.__properties = {}
	40	+ self._properties = {}
	41	+ self.__property_keys = ['os', 'has-docker']
27	42	self.__name = str(uuid.uuid4())
28	43
29	44	if properties:
30		- for k, v in properties.items():
31		- if k == 'os':
32		- self.__properties[k] = v
33		-
34		- elif k == 'docker':
35		- if v not in ('True', 'False'):
36		- raise ValueError('Value not supported: [{}]'.format(v))
37		- self.__properties[k] = v
38		-
39		- else:
40		- raise KeyError('Key not supported: [{}]'.format(k))
	45	+ for prop in properties:
	46	+ self._add_property(prop)
41	47
42	48	@property
43	49	def name(self):
...	...	@@ -45,13 +51,25 @@ class Device:
45	51
46	52	@property
47	53	def properties(self):
48		- return self.__properties
	54	+ return self._properties
49	55
50	56	def get_pb2(self):
51		- device = worker_pb2.Device(handle=self._name)
52		- property_message = worker_pb2.Device.Property()
	57	+ device = worker_pb2.Device(handle=self.__name)
53	58	for k, v in self._properties.items():
54		- property_message.key = k
55		- property_message.value = v
56		- device.properties.extend([property_message])
	59	+ for prop in v:
	60	+ property_message = worker_pb2.Device.Property()
	61	+ property_message.key = k
	62	+ property_message.value = prop
	63	+ device.properties.extend([property_message])
57	64	return device
	65	+
	66	+ def _add_property(self, key, value):
	67	+ if key in self.__property_keys:
	68	+ prop = self._properties.get(key)
	69	+ if not prop:
	70	+ self._properties[key] = [value]
	71	+ else:
	72	+ prop[key].append(value)
	73	+
	74	+ else:
	75	+ raise KeyError('Key not supported: [{}]'.format(key))

buildgrid/bot/hardware/interface.py

---

	1	+# Copyright (C) 2018 Bloomberg LP
	2	+#
	3	+# Licensed under the Apache License, Version 2.0 (the "License");
	4	+# you may not use this file except in compliance with the License.
	5	+# You may obtain a copy of the License at
	6	+#
	7	+# <http://www.apache.org/licenses/LICENSE-2.0>
	8	+#
	9	+# Unless required by applicable law or agreed to in writing, software
	10	+# distributed under the License is distributed on an "AS IS" BASIS,
	11	+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	12	+# See the License for the specific language governing permissions and
	13	+# limitations under the License.
	14	+
	15	+
	16	+"""
	17	+HardwareInterface
	18	+=================
	19	+
	20	+Class to configure hardware and check requirements of leases.
	21	+
	22	+In the future this could also be used to request and display
	23	+the status of hardware.
	24	+"""
	25	+
	26	+
	27	+from buildgrid._exceptions import FailedPreconditionError
	28	+
	29	+
	30	+class HardwareInterface:
	31	+
	32	+ def __init__(self, worker):
	33	+ self._worker = worker
	34	+
	35	+ def configure_hardware(self, requirements):
	36	+ """ Can check if the requirements can be met and also
	37	+ in the future, potentially configure the hardware.
	38	+ """
	39	+ worker = self._worker
	40	+
	41	+ for config_requirement in requirements.configs:
	42	+ if config_requirement.key not in worker.configs:
	43	+ raise FailedPreconditionError("Config not supported: [{}]".format(config_requirement))
	44	+
	45	+ def get_worker_pb2(self):
	46	+ return self._worker.get_pb2()

buildgrid/bot/hardware/worker.py

---

...	...	@@ -13,39 +13,34 @@
13	13	# limitations under the License.
14	14
15	15
	16	+from buildgrid._protos.google.devtools.remoteworkers.v1test2 import worker_pb2
	17	+
	18	+
16	19	class Worker:
17	20
18	21	def __init__(self, properties=None, configs=None):
19	22	self._devices = []
20		- self.__configs = {}
21		- self.__properties = {}
	23	+ self._configs = {}
	24	+ self._properties = {}
	25	+ self.__property_keys = ['pool']
	26	+ self.__config_keys = ['DockerImage']
22	27
23	28	if properties:
24	29	for k, v in properties.items():
25		- self.update_config(k, v)
	30	+ if k in self.__property_keys:
	31	+ self._add_properties(k, v)
26	32
27	33	if configs:
28	34	for k, v in configs.items():
29		- if k == 'DockerImage':
30		- self.__configs[k] = v
31		- else:
32		- raise KeyError('Key not supported: [{}]'.format(k))
	35	+ self._add_config(k, v)
33	36
34	37	@property
35	38	def configs(self):
36		- return self.__configs
	39	+ return self._configs
37	40
38	41	@property
39	42	def properties(self):
40		- return self.__properties
41		-
42		- def update_config(self, key, value):
43		- """Update the config of the device.
44		- Only certain keys are allowed."""
45		- if k == 'pool':
46		- self.__properties[k] = v
47		- else:
48		- raise KeyError('Key not supported: [{}]'.format(k))
	43	+ return self._properties
49	44
50	45	def add_device(self, device):
51	46	self._devices.append(device)
...	...	@@ -53,16 +48,41 @@ class Worker:
53	48	def get_pb2(self):
54	49	devices = [device.get_pb2() for device in self._devices]
55	50	worker = worker_pb2.Worker(devices=devices)
56		- property_message = worker_pb2.Worker.Property()
57		- for k, v in self.__properties.items():
58		- property_message.key = k
59		- property_message.value = v
60		- worker.properties.extend([property_message])
61		-
62		- config_message = worker_pb2.Worker.Config()
63		- for k, v in self.__properties.items():
64		- property_message.key = k
65		- property_message.value = v
66		- worker.configs.extend([config_message])
	51	+
	52	+ for k, v in self._properties.items():
	53	+ for prop in v:
	54	+ property_message = worker_pb2.Device.Property()
	55	+ property_message.key = k
	56	+ property_message.value = prop
	57	+ device.properties.extend([property_message])
	58	+
	59	+ for k, v in self._configs.items():
	60	+ for cfg in v:
	61	+ config_message = worker_pb2.Worker.Config()
	62	+ config.key = k
	63	+ config_message.value = cfg
	64	+ worker.configs.extend([config_message])
67	65
68	66	return worker
	67	+
	68	+ def _add_config(self, key, value):
	69	+ if key in self.__config_keys:
	70	+ cfg = self._configs.get(key)
	71	+ if not cfg:
	72	+ self._configs[key] = [value]
	73	+ else:
	74	+ cfg[key].append(value)
	75	+
	76	+ else:
	77	+ raise KeyError('Key not supported: [{}]'.format(key))
	78	+
	79	+ def _add_property(self, key, value):
	80	+ if key in self.__property_keys:
	81	+ prop = self._properties.get(key)
	82	+ if not prop:
	83	+ self._properties[key] = [value]
	84	+ else:
	85	+ prop[key].append(value)
	86	+
	87	+ else:
	88	+ raise KeyError('Key not supported: [{}]'.format(key))

buildgrid/bot/session.py

---

...	...	@@ -25,18 +25,20 @@ Allows connections
25	25	import asyncio
26	26	import logging
27	27	import platform
28		-import uuid
29	28
30	29	import grpc
31	30
32	31	from buildgrid._enums import BotStatus, LeaseState
	32	+from buildgrid._protos.google.devtools.remoteworkers.v1test2 import bots_pb2
33	33	from buildgrid._protos.google.rpc import code_pb2
34		-from buildgrid._protos.google.devtools.remoteworkers.v1test2 import bots_pb2, worker_pb2
35	34	from buildgrid._exceptions import BotError
36	35
	36	+from buildgrid._exceptions import FailedPreconditionError
	37	+
	38	+from .tenantmanager import TenantManager
37	39
38	40	class BotSession:
39		- def __init__(self, parent, interface, worker):
	41	+ def __init__(self, parent, bots_interface, hardware_interface):
40	42	""" Unique bot ID within the farm used to identify this bot
41	43	Needs to be human readable.
42	44	All prior sessions with bot_id of same ID are invalidated.
...	...	@@ -45,88 +47,84 @@ class BotSession:
45	47	"""
46	48
47	49	self.logger = logging.getLogger(__name__)
48		- self._context = None
49	50
50		- self._worker = worker
51		- self._interface = interface
52		- self._leases = {}
53		- self._parent = parent
	51	+ self._bots_interface = bots_interface
	52	+ self._hardware_interface = hardware_interface
	53	+
54	54	self._status = BotStatus.OK.value
	55	+ self._tenant_manager = TenantManager()
55	56
	57	+ self.__parent = parent
56	58	self.__bot_id = '{}.{}'.format(parent, platform.node())
57	59	self.__name = None
58	60
	61	+ # Remove these and add to a worker config in the future
	62	+ self._work = None
	63	+ self._context = None
	64	+
59	65	@property
60	66	def bot_id(self):
61	67	return self.__bot_id
62	68
63		- def create_bot_session(self, work, context=None):
64		- self.logger.debug("Creating bot session")
	69	+ def create_bot_session(self, work, context):
	70	+ # Drop this when properly adding to the work
65	71	self._work = work
66	72	self._context = context
67	73
68		- session = self._interface.create_bot_session(self._parent, self.get_pb2())
	74	+ self.logger.debug("Creating bot session")
	75	+
	76	+ session = self._bots_interface.create_bot_session(self.__parent, self.get_pb2())
69	77	self.__name = session.name
70	78
71	79	self.logger.info("Created bot session with name: [{}]".format(self.__name))
72	80
73	81	for lease in session.leases:
74		- self._update_lease_from_server(lease)
	82	+ self._register_lease(lease)
75	83
76	84	def update_bot_session(self):
77	85	self.logger.debug("Updating bot session: [{}]".format(self.__bot_id))
78		- session = self._interface.update_bot_session(self.get_pb2())
79		- for k, v in list(self._leases.items()):
80		- if v.state == LeaseState.COMPLETED.value:
81		- del self._leases[k]
	86	+
	87	+ session = self._bots_interface.update_bot_session(self.get_pb2())
	88	+
	89	+ server_ids = []
82	90
83	91	for lease in session.leases:
84		- self._update_lease_from_server(lease)
	92	+ server_ids.append(lease.id)
85	93
86		- def get_pb2(self):
87		- leases = list(self._leases.values())
88		- if not leases:
89		- leases = None
	94	+ lease_state = LeaseState(lease.state)
	95	+ if lease_state == lease.status.PENDING:
	96	+ self._register_lease(lease)
90	97
91		- return bots_pb2.BotSession(worker=self._worker.get_pb2(),
92		- status=self._status,
93		- leases=leases,
94		- bot_id=self.__bot_id,
95		- name=self.__name)
	98	+ elif lease_state == lease.status.CANCELLED:
	99	+ self._tenant_manager.cancel_tenancy(lease_id)
96	100
97		- def lease_completed(self, lease):
98		- lease.state = LeaseState.COMPLETED.value
99		- self._leases[lease.id] = lease
	101	+ closed_lease_ids = [x for x in self._tenant_manager.get_lease_ids() if x not in server_ids]
100	102
101		- def _update_lease_from_server(self, lease):
102		- """
103		- State machine for any recieved updates to the leases.
104		- """
105		- # TODO: Compare with previous state of lease
106		- if lease.state == LeaseState.PENDING.value:
107		- lease.state = LeaseState.ACTIVE.value
108		- self._leases[lease.id] = lease
109		- self.update_bot_session()
110		- asyncio.ensure_future(self.create_work(lease))
	103	+ for lease_id in closed_lease_ids:
	104	+ self._tenant_manager.remove_tenant(lease_id)
111	105
112		- async def create_work(self, lease):
113		- self.logger.debug("Work created: [{}]".format(lease.id))
114		- loop = asyncio.get_event_loop()
	106	+ def get_pb2(self):
	107	+ return bots_pb2.BotSession(worker=self._hardware_interface.get_worker_pb2(),
	108	+ status=self._status,
	109	+ leases=self._tenant_manager.get_leases(),
	110	+ bot_id=self.__bot_id,
	111	+ name=self.__name)
115	112
	113	+ def _register_lease(self, lease):
	114	+ lease_id = lease.id
116	115	try:
117		- lease = await loop.run_in_executor(None, self._work, self._context, lease)
	116	+ self._tenant_manager.create_tenancy(lease)
118	117
119		- except grpc.RpcError as e:
120		- self.logger.error("RPC error thrown: [{}]".format(e))
121		- lease.status.CopyFrom(e.code())
	118	+ except KeyError as e:
	119	+ self.logger.error("Cannot register lease: [{}]. Error: {}".format(lease.id, e))
122	120
123		- except BotError as e:
124		- self.logger.error("Internal bot error thrown: [{}]".format(e))
125		- lease.status.code = code_pb2.INTERNAL
	121	+ else:
	122	+ try:
	123	+ self._hardware_interface.configure_hardware(lease.requirements)
126	124
127		- except Exception as e:
128		- self.logger.error("Exception thrown: [{}]".format(e))
129		- lease.status.code = code_pb2.INTERNAL
	125	+ except FailedPreconditionError as e:
	126	+ self.logger.error("Failed precondition: [{}]".format(e))
	127	+ self._tenant_manager.complete_lease(lease_id, status=code_pb2.FailedPreconditionError)
130	128
131		- self.logger.debug("Work complete: [{}]".format(lease.id))
132		- self.lease_completed(lease)
	129	+ else:
	130	+ 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	+"""
	16	+Tenant
	17	+======
	18	+
	19	+Handles leased and runs leased work.
	20	+"""
	21	+
	22	+import logging
	23	+
	24	+from functools import partial
	25	+
	26	+from buildgrid._protos.google.devtools.remoteworkers.v1test2 import bots_pb2
	27	+
	28	+from buildgrid._enums import LeaseState
	29	+
	30	+
	31	+class Tenant:
	32	+
	33	+ def __init__(self, lease):
	34	+
	35	+ if lease.state != LeaseState.PENDING.value:
	36	+ raise ValueError("Lease state not `PENDING`: {}".format(lease.state))
	37	+
	38	+ self.logger = logging.getLogger(__name__)
	39	+ self.lease_finished = False
	40	+
	41	+ self._lease = lease
	42	+
	43	+ @property
	44	+ def lease(self):
	45	+ return self._lease
	46	+
	47	+ def get_lease_state(self):
	48	+ return self._lease.state
	49	+
	50	+ def update_lease_state(self, state):
	51	+ self._lease.state = state.value
	52	+
	53	+ def update_lease_status(self, status):
	54	+ self._lease.status.CopyFrom(status)
	55	+
	56	+ async def run_work(self, work, context=None, executor=None):
	57	+ self.logger.debug("Work created: [{}]".format(self._lease.id))
	58	+
	59	+ # Ensures if anything happens to the lease during work, we still have a copy.
	60	+ lease = bots_pb2.Lease()
	61	+ lease.CopyFrom(self._lease)
	62	+
	63	+ loop = asyncio.get_event_loop()
	64	+
	65	+ try:
	66	+ lease = await loop.run_in_executor(executor, partial(work, context, self._lease))
	67	+ self._lease.CopyFrom(lease)
	68	+
	69	+ except asyncio.CancelledError as e:
	70	+ self.logger.error("Task cancelled: [{}]".format(e))
	71	+
	72	+ except grpc.RpcError as e:
	73	+ self.logger.error("RPC error thrown: [{}]".format(e))
	74	+ lease.status.CopyFrom(e.code())
	75	+
	76	+ except BotError as e:
	77	+ self.logger.error("Internal bot error thrown: [{}]".format(e))
	78	+ lease.status.code = code_pb2.INTERNAL
	79	+
	80	+ except Exception as e:
	81	+ self.logger.error("Exception thrown: [{}]".format(e))
	82	+ lease.status.code = code_pb2.INTERNAL
	83	+
	84	+ finally:
	85	+ self.logger.debug("Work completed: [{}]".format(lease.id))

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	+"""
	17	+TenantManager
	18	+=============
	19	+
	20	+Looks after leases of work.
	21	+"""
	22	+
	23	+
	24	+import asyncio
	25	+import logging
	26	+from functools import partial
	27	+
	28	+import grpc
	29	+
	30	+from buildgrid._enums import LeaseState
	31	+
	32	+from .tenant import Tenant
	33	+
	34	+class TenantManager:
	35	+
	36	+ def __init__(self):
	37	+ self.logger = logging.getLogger(__name__)
	38	+ self._tenants = {}
	39	+ self._tasks = {}
	40	+
	41	+ def create_tenancy(self, lease):
	42	+ lease_id = lease.id
	43	+
	44	+ if lease_id not in self._tenants:
	45	+ tenant = Tenant(lease)
	46	+ self._tenants[lease_id] = tenant
	47	+
	48	+ else:
	49	+ raise KeyError("Lease id already exists: [{}]".format(lease_id))
	50	+
	51	+ def remove_tenant(self, lease_id):
	52	+ state = self.get_lease_state(lease_id)
	53	+ if state != LeaseState.COMPLETED or state != LeaseState.CANCELLED:
	54	+ self.logger.error("Attempting to remove a lease not finished."
	55	+ "Bot will not remove lease."
	56	+ "Lease: [{}]".format(self._tenant_manager.get_lease(lease_id)))
	57	+
	58	+ else:
	59	+ del self._tenants[lease_id]
	60	+ del self._tasks[lease_id]
	61	+
	62	+ def get_leases(self):
	63	+ leases = []
	64	+ for _, tenant in self._tenants.items():
	65	+ leases.append(tenant.lease)
	66	+
	67	+ if not leases:
	68	+ return None
	69	+
	70	+ return leases
	71	+
	72	+ def get_lease_ids(self):
	73	+ lease_ids = []
	74	+ for lease_id in self._tenants.keys():
	75	+ lease_ids.append(lease_id)
	76	+
	77	+ if not lease_ids:
	78	+ return []
	79	+
	80	+ return leases
	81	+
	82	+ def get_lease_state(self, lease_id):
	83	+ return self._tenants[lease_id].get_lease_state()
	84	+
	85	+ def complete_lease(self, lease_id, status=None):
	86	+ if status is not None:
	87	+ self._update_lease_status(lease_id, status)
	88	+
	89	+ if self._tenants[lease_id].get_lease_state() != LeaseState.CANCELLED:
	90	+ self._update_lease_state(lease_id, LeaseState.COMPLETED)
	91	+
	92	+ def create_work(self, lease_id, work, context):
	93	+ self._update_lease_state(lease_id, LeaseState.ACTIVE)
	94	+ tenant = self._tenants[lease_id]
	95	+ task = asyncio.ensure_future(tenant.run_work(work, context))
	96	+
	97	+ task.add_done_callback(partial(self.complete_lease, lease_id))
	98	+
	99	+ def cancel_tenancy(self, lease_id):
	100	+ self._update_lease_state(LeaseState.CANCELLED)
	101	+ self._tasks[lease_id].cancel()
	102	+
	103	+ def _update_lease_state(self, lease_id, state):
	104	+ self._tenants[lease_id].update_lease_state(state)
	105	+
	106	+ def _update_lease_status(self, lease_id, status):
	107	+ self._tenants[lease_id].update_lease_status(status)

buildgrid/server/bots/instance.py

---

...	...	@@ -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):
...	...	@@ -63,7 +63,6 @@ class BotsInterface:
63	63	bot_session.name = name
64	64
65	65	self._bot_ids[name] = bot_id
66		- self._bot_sessions[name] = bot_session
67	66	self.logger.info("Created bot session name=[{}] with bot_id=[{}]".format(name, bot_id))
68	67
69	68	# TODO: Send worker capabilities to the scheduler!
...	...	@@ -79,6 +78,7 @@ class BotsInterface:
79	78	"""
80	79	self.logger.debug("Updating bot session name={}".format(name))
81	80	self._check_bot_ids(bot_session.bot_id, name)
	81	+ self._check_assigned_leases(bot_session)
82	82
83	83	leases = filter(None, [self._check_lease_state(lease) for lease in bot_session.leases])
84	84
...	...	@@ -93,16 +93,18 @@ class BotsInterface:
93	93	if leases:
94	94	bot_session.leases.extend(leases)
95	95
96		- self._bot_sessions[name] = bot_session
97	96	return bot_session
98	97
99	98	def _check_lease_state(self, lease):
100		-
101		- # Check for cancelled lease
102		- if self._scheduler.get_lease_cancelled(lease.id):
	99	+ # TODO: Pop assigned leases
	100	+ try:
	101	+ if self._scheduler.get_job_lease_cancelled(lease.id):
	102	+ lease.state.CopyFrom(LeaseState.CANCELLED.value)
	103	+ return lease
	104	+ except KeyError:
	105	+ # Job does not exist, remove from bot.
103	106	return None
104	107
105		- # If not cancelled, update the status
106	108	self._scheduler.update_job_lease(lease)
107	109
108	110	lease_state = LeaseState(lease.state)
...	...	@@ -134,6 +136,15 @@ class BotsInterface:
134	136	'Bot id already registered. ID sent: [{}].'
135	137	'Id registered: [{}] with name: [{}]'.format(bot_id, _bot_id, _name))
136	138
	139	+ def _check_assigned_leases(self, bot_session):
	140	+ session_leases = bot_session.leases
	141	+ for lease in self._assigned_leases:
	142	+ if lease not in session_leases:
	143	+ self.logger.error("Assigned lease=[{}]"
	144	+ " not found on bot with name=[{}]."
	145	+ " Retrying job".format(lease, bot_session.bot_id))
	146	+ self._scheduler.retry_job(lease.id)
	147	+
137	148	def _close_bot_session(self, name):
138	149	""" Before removing the session, close any leases and
139	150	requeue with high priority.
...	...	@@ -144,7 +155,7 @@ class BotsInterface:
144	155	raise InvalidArgumentError("Bot id does not exist: [{}]".format(name))
145	156
146	157	self.logger.debug("Attempting to close [{}] with name: [{}]".format(bot_id, name))
147		- for lease in self._bot_sessions[name].leases:
	158	+ for lease in self._assigned_leases[name]:
148	159	if lease.state != LeaseState.COMPLETED.value:
149	160	# TODO: Be wary here, may need to handle rejected leases in future
150	161	self._scheduler.retry_job(lease.id)

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