finn pushed to branch finn/48-cancellation-leases at BuildGrid / buildgrid
Commits:
- 
70dbb7a1
by Finn at 2018-11-07T17:13:42Z
13 changed files:
- buildgrid/_app/bots/dummy.py
- 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/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
Changes:
| ... | ... | @@ -46,4 +46,7 @@ def work_dummy(context, lease): | 
| 46 | 46 |  | 
| 47 | 47 |      lease.result.Pack(action_result)
 | 
| 48 | 48 |  | 
| 49 | +    # while True:
 | |
| 50 | +    #    pass
 | |
| 51 | + | |
| 49 | 52 |      return lease | 
| ... | ... | @@ -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 |  | 
| ... | ... | @@ -17,11 +17,12 @@ | 
| 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
 | 
| 24 | 24 |  import logging
 | 
| 25 | +import sys
 | |
| 25 | 26 |  | 
| 26 | 27 |  | 
| 27 | 28 |  class Bot:
 | 
| ... | ... | @@ -45,6 +46,7 @@ class Bot: | 
| 45 | 46 |              loop.run_forever()
 | 
| 46 | 47 |          except KeyboardInterrupt:
 | 
| 47 | 48 |              pass
 | 
| 49 | + | |
| 48 | 50 |          finally:
 | 
| 49 | 51 |              task.cancel()
 | 
| 50 | 52 |              loop.close()
 | 
| ... | ... | @@ -54,5 +56,11 @@ class Bot: | 
| 54 | 56 |          Calls the server periodically to inform the server the client has not died.
 | 
| 55 | 57 |          """
 | 
| 56 | 58 |          while True:
 | 
| 57 | -            self._bot_session.update_bot_session()
 | |
| 59 | +            try:
 | |
| 60 | +                self._bot_session.update_bot_session()
 | |
| 61 | + | |
| 62 | +            except Exception as e:
 | |
| 63 | +                self.logger.error("Error: [{}]".format(e))
 | |
| 64 | +                raise
 | |
| 65 | + | |
| 58 | 66 |              await asyncio.sleep(self._update_period) | 
| ... | ... | @@ -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)) | 
| 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() | 
| ... | ... | @@ -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)) | 
| ... | ... | @@ -38,10 +38,17 @@ class BotInterface: | 
| 38 | 38 |      def create_bot_session(self, parent, bot_session):
 | 
| 39 | 39 |          request = bots_pb2.CreateBotSessionRequest(parent=parent,
 | 
| 40 | 40 |                                                     bot_session=bot_session)
 | 
| 41 | -        return self._stub.CreateBotSession(request)
 | |
| 41 | +        try:
 | |
| 42 | +            return self._stub.CreateBotSession(request)
 | |
| 43 | +        except Exception as e:
 | |
| 44 | +            self.logger.error("Error creating bot session: [{}]".format(e))
 | |
| 45 | +            raise
 | |
| 42 | 46 |  | 
| 43 | 47 |      def update_bot_session(self, bot_session, update_mask=None):
 | 
| 44 | 48 |          request = bots_pb2.UpdateBotSessionRequest(name=bot_session.name,
 | 
| 45 | 49 |                                                     bot_session=bot_session,
 | 
| 46 | 50 |                                                     update_mask=update_mask)
 | 
| 47 | -        return self._stub.UpdateBotSession(request) | |
| 51 | +        try:
 | |
| 52 | +            return self._stub.UpdateBotSession(request)
 | |
| 53 | +        except Exception as e:
 | |
| 54 | +            self.logger.error("Error updating bot session: [{}]".format(e)) | 
| ... | ... | @@ -25,18 +25,22 @@ Allows connections | 
| 25 | 25 |  import asyncio
 | 
| 26 | 26 |  import logging
 | 
| 27 | 27 |  import platform
 | 
| 28 | -import uuid
 | |
| 28 | +import pdb
 | |
| 29 | +# pdb.set_trace()
 | |
| 29 | 30 |  | 
| 30 | 31 |  import grpc
 | 
| 31 | 32 |  | 
| 32 | 33 |  from buildgrid._enums import BotStatus, LeaseState
 | 
| 34 | +from buildgrid._protos.google.devtools.remoteworkers.v1test2 import bots_pb2
 | |
| 33 | 35 |  from buildgrid._protos.google.rpc import code_pb2
 | 
| 34 | -from buildgrid._protos.google.devtools.remoteworkers.v1test2 import bots_pb2, worker_pb2
 | |
| 35 | 36 |  from buildgrid._exceptions import BotError
 | 
| 36 | 37 |  | 
| 38 | +from buildgrid._exceptions import FailedPreconditionError
 | |
| 39 | + | |
| 40 | +from .tenantmanager import TenantManager
 | |
| 37 | 41 |  | 
| 38 | 42 |  class BotSession:
 | 
| 39 | -    def __init__(self, parent, interface, worker):
 | |
| 43 | +    def __init__(self, parent, bots_interface, hardware_interface):
 | |
| 40 | 44 |          """ Unique bot ID within the farm used to identify this bot
 | 
| 41 | 45 |          Needs to be human readable.
 | 
| 42 | 46 |          All prior sessions with bot_id of same ID are invalidated.
 | 
| ... | ... | @@ -45,88 +49,82 @@ class BotSession: | 
| 45 | 49 |          """
 | 
| 46 | 50 |  | 
| 47 | 51 |          self.logger = logging.getLogger(__name__)
 | 
| 48 | -        self._context = None
 | |
| 49 | 52 |  | 
| 50 | -        self._worker = worker
 | |
| 51 | -        self._interface = interface
 | |
| 52 | -        self._leases = {}
 | |
| 53 | -        self._parent = parent
 | |
| 53 | +        self._bots_interface = bots_interface
 | |
| 54 | +        self._hardware_interface = hardware_interface
 | |
| 55 | + | |
| 54 | 56 |          self._status = BotStatus.OK.value
 | 
| 57 | +        self._tenant_manager = TenantManager()
 | |
| 55 | 58 |  | 
| 59 | +        self.__parent = parent
 | |
| 56 | 60 |          self.__bot_id = '{}.{}'.format(parent, platform.node())
 | 
| 57 | 61 |          self.__name = None
 | 
| 58 | 62 |  | 
| 63 | +        # Remove these and add to a worker config in the future
 | |
| 64 | +        self._work = None
 | |
| 65 | +        self._context = None
 | |
| 66 | + | |
| 59 | 67 |      @property
 | 
| 60 | 68 |      def bot_id(self):
 | 
| 61 | 69 |          return self.__bot_id
 | 
| 62 | 70 |  | 
| 63 | -    def create_bot_session(self, work, context=None):
 | |
| 64 | -        self.logger.debug("Creating bot session")
 | |
| 71 | +    def create_bot_session(self, work, context):
 | |
| 72 | +        # Drop this when properly adding to the work
 | |
| 65 | 73 |          self._work = work
 | 
| 66 | 74 |          self._context = context
 | 
| 67 | 75 |  | 
| 68 | -        session = self._interface.create_bot_session(self._parent, self.get_pb2())
 | |
| 76 | +        self.logger.debug("Creating bot session")
 | |
| 77 | + | |
| 78 | +        session = self._bots_interface.create_bot_session(self.__parent, self.get_pb2())
 | |
| 69 | 79 |          self.__name = session.name
 | 
| 70 | 80 |  | 
| 71 | 81 |          self.logger.info("Created bot session with name: [{}]".format(self.__name))
 | 
| 72 | 82 |  | 
| 73 | 83 |          for lease in session.leases:
 | 
| 74 | -            self._update_lease_from_server(lease)
 | |
| 84 | +            self._register_lease(lease)
 | |
| 75 | 85 |  | 
| 76 | 86 |      def update_bot_session(self):
 | 
| 77 | 87 |          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]
 | |
| 88 | + | |
| 89 | +        session = self._bots_interface.update_bot_session(self.get_pb2())
 | |
| 90 | +        server_ids = []
 | |
| 82 | 91 |  | 
| 83 | 92 |          for lease in session.leases:
 | 
| 84 | -            self._update_lease_from_server(lease)
 | |
| 93 | +            server_ids.append(lease.id)
 | |
| 85 | 94 |  | 
| 86 | -    def get_pb2(self):
 | |
| 87 | -        leases = list(self._leases.values())
 | |
| 88 | -        if not leases:
 | |
| 89 | -            leases = None
 | |
| 95 | +            lease_state = LeaseState(lease.state)
 | |
| 96 | +            if lease_state == LeaseState.PENDING:
 | |
| 97 | +                self._register_lease(lease)
 | |
| 98 | + | |
| 99 | +            elif lease_state == LeaseState.CANCELLED:
 | |
| 100 | +                self._tenant_manager.cancel_tenancy(lease_id)
 | |
| 101 | + | |
| 102 | +        closed_lease_ids = [x for x in self._tenant_manager.get_lease_ids() if x not in server_ids]
 | |
| 103 | +        for lease_id in closed_lease_ids:
 | |
| 104 | +            self._tenant_manager.remove_tenant(lease_id)
 | |
| 90 | 105 |  | 
| 91 | -        return bots_pb2.BotSession(worker=self._worker.get_pb2(),
 | |
| 106 | +    def get_pb2(self):
 | |
| 107 | +        return bots_pb2.BotSession(worker=self._hardware_interface.get_worker_pb2(),
 | |
| 92 | 108 |                                     status=self._status,
 | 
| 93 | -                                   leases=leases,
 | |
| 109 | +                                   leases=self._tenant_manager.get_leases(),
 | |
| 94 | 110 |                                     bot_id=self.__bot_id,
 | 
| 95 | 111 |                                     name=self.__name)
 | 
| 96 | 112 |  | 
| 97 | -    def lease_completed(self, lease):
 | |
| 98 | -        lease.state = LeaseState.COMPLETED.value
 | |
| 99 | -        self._leases[lease.id] = lease
 | |
| 100 | - | |
| 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))
 | |
| 111 | - | |
| 112 | -    async def create_work(self, lease):
 | |
| 113 | -        self.logger.debug("Work created: [{}]".format(lease.id))
 | |
| 114 | -        loop = asyncio.get_event_loop()
 | |
| 115 | - | |
| 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.debug("Cannot register lease=[{}]. {}".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) | 
| 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 asyncio
 | |
| 23 | +import logging
 | |
| 24 | + | |
| 25 | +from functools import partial
 | |
| 26 | + | |
| 27 | +from buildgrid._protos.google.devtools.remoteworkers.v1test2 import bots_pb2
 | |
| 28 | + | |
| 29 | +from buildgrid._enums import LeaseState
 | |
| 30 | + | |
| 31 | + | |
| 32 | +class Tenant:
 | |
| 33 | + | |
| 34 | +    def __init__(self, lease):
 | |
| 35 | + | |
| 36 | +        if lease.state != LeaseState.PENDING.value:
 | |
| 37 | +            raise ValueError("Lease state not `PENDING`: {}".format(lease.state))
 | |
| 38 | + | |
| 39 | +        self.logger = logging.getLogger(__name__)
 | |
| 40 | +        self.lease_finished = False
 | |
| 41 | + | |
| 42 | +        self._lease = lease
 | |
| 43 | + | |
| 44 | +    @property
 | |
| 45 | +    def lease(self):
 | |
| 46 | +        return self._lease
 | |
| 47 | + | |
| 48 | +    def get_lease_state(self):
 | |
| 49 | +        return LeaseState(self._lease.state)
 | |
| 50 | + | |
| 51 | +    def update_lease_state(self, state):
 | |
| 52 | +        self._lease.state = state.value
 | |
| 53 | + | |
| 54 | +    def update_lease_status(self, status):
 | |
| 55 | +        self._lease.status.CopyFrom(status)
 | |
| 56 | + | |
| 57 | +    async def run_work(self, work, context=None, executor=None):
 | |
| 58 | +        self.logger.debug("Work created: [{}]".format(self._lease.id))
 | |
| 59 | + | |
| 60 | +        # Ensures if anything happens to the lease during work, we still have a copy.
 | |
| 61 | +        lease = bots_pb2.Lease()
 | |
| 62 | +        lease.CopyFrom(self._lease)
 | |
| 63 | + | |
| 64 | +        loop = asyncio.get_event_loop()
 | |
| 65 | + | |
| 66 | +        try:
 | |
| 67 | +            lease = await loop.run_in_executor(executor, partial(work, context, self._lease))
 | |
| 68 | +            self._lease.CopyFrom(lease)
 | |
| 69 | + | |
| 70 | +        except asyncio.CancelledError as e:
 | |
| 71 | +            self.logger.error("Task cancelled: [{}]".format(e))
 | |
| 72 | + | |
| 73 | +        except grpc.RpcError as e:
 | |
| 74 | +            self.logger.error("RPC error thrown: [{}]".format(e))
 | |
| 75 | +            lease.status.CopyFrom(e.code())
 | |
| 76 | + | |
| 77 | +        except BotError as e:
 | |
| 78 | +            self.logger.error("Internal bot error thrown: [{}]".format(e))
 | |
| 79 | +            lease.status.code = code_pb2.INTERNAL
 | |
| 80 | + | |
| 81 | +        except Exception as e:
 | |
| 82 | +            self.logger.error("Exception thrown: [{}]".format(e))
 | |
| 83 | +            lease.status.code = code_pb2.INTERNAL
 | |
| 84 | + | |
| 85 | +        self.logger.debug("Work completed: [{}]".format(lease.id)) | 
| 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.PENDING or state == LeaseState.ACTIVE:
 | |
| 54 | +            self.logger.error("Attempting to remove a lease not finished."
 | |
| 55 | +                              "Bot will not remove lease."
 | |
| 56 | +                              "Lease: [{}]".format(self._tenants[lease_id].lease))
 | |
| 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.values():
 | |
| 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 | +        return self._tenants.keys()
 | |
| 74 | + | |
| 75 | +    def get_lease_state(self, lease_id):
 | |
| 76 | +        return self._tenants[lease_id].get_lease_state()
 | |
| 77 | + | |
| 78 | +    def complete_lease(self, lease_id, status, task=None):
 | |
| 79 | +        if status is not None:
 | |
| 80 | +            self._update_lease_status(lease_id, status)
 | |
| 81 | + | |
| 82 | +        if self._tenants[lease_id].get_lease_state() != LeaseState.CANCELLED:
 | |
| 83 | +            self._update_lease_state(lease_id, LeaseState.COMPLETED)
 | |
| 84 | + | |
| 85 | +    def create_work(self, lease_id, work, context):
 | |
| 86 | +        self._update_lease_state(lease_id, LeaseState.ACTIVE)
 | |
| 87 | +        tenant = self._tenants[lease_id]
 | |
| 88 | +        task = asyncio.ensure_future(tenant.run_work(work, context))
 | |
| 89 | + | |
| 90 | +        task.add_done_callback(partial(self.complete_lease, lease_id, None))
 | |
| 91 | + | |
| 92 | +        self._tasks[lease_id] = task
 | |
| 93 | + | |
| 94 | +    def cancel_tenancy(self, lease_id):
 | |
| 95 | +        self._update_lease_state(LeaseState.CANCELLED)
 | |
| 96 | +        self._tasks[lease_id].cancel()
 | |
| 97 | + | |
| 98 | +    def _update_lease_state(self, lease_id, state):
 | |
| 99 | +        self._tenants[lease_id].update_lease_state(state)
 | |
| 100 | + | |
| 101 | +    def _update_lease_status(self, lease_id, status):
 | |
| 102 | +        self._tenants[lease_id].update_lease_status(status) | 
| ... | ... | @@ -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,15 @@ 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
 | |
| 67 | 65 |          self.logger.info("Created bot session name=[{}] with bot_id=[{}]".format(name, bot_id))
 | 
| 68 | 66 |  | 
| 69 | -        # TODO: Send worker capabilities to the scheduler!
 | |
| 70 | -        leases = self._scheduler.request_job_leases({})
 | |
| 71 | -        if leases:
 | |
| 72 | -            bot_session.leases.extend(leases)
 | |
| 67 | +        # We want to keep a copy of lease ids we have assigned
 | |
| 68 | +        self._assigned_leases[name] = set()
 | |
| 73 | 69 |  | 
| 70 | +        self._request_leases(bot_session)
 | |
| 74 | 71 |          return bot_session
 | 
| 75 | 72 |  | 
| 76 | 73 |      def update_bot_session(self, name, bot_session):
 | 
| ... | ... | @@ -79,39 +76,53 @@ class BotsInterface: | 
| 79 | 76 |          """
 | 
| 80 | 77 |          self.logger.debug("Updating bot session name={}".format(name))
 | 
| 81 | 78 |          self._check_bot_ids(bot_session.bot_id, name)
 | 
| 82 | - | |
| 83 | -        leases = filter(None, [self._check_lease_state(lease) for lease in bot_session.leases])
 | |
| 79 | +        self._check_assigned_leases(bot_session)
 | |
| 84 | 80 |  | 
| 85 | 81 |          for lease in bot_session.leases:
 | 
| 86 | -            lease.Clear()
 | |
| 87 | - | |
| 88 | -        bot_session.leases.extend(leases)
 | |
| 82 | +            checked_lease = self._check_lease_state(lease)
 | |
| 83 | +            if not checked_lease:
 | |
| 84 | +                # TODO: Make sure we don't need this
 | |
| 85 | +                try:
 | |
| 86 | +                    self._assigned_leases[name].remove(lease.id)
 | |
| 87 | +                except KeyError:
 | |
| 88 | +                    pass
 | |
| 89 | +                lease.Clear()
 | |
| 90 | + | |
| 91 | +        self._request_leases(bot_session)
 | |
| 92 | +        return bot_session
 | |
| 89 | 93 |  | 
| 94 | +    def _request_leases(self, bot_session):
 | |
| 90 | 95 |          # TODO: Send worker capabilities to the scheduler!
 | 
| 96 | +        # Only send one lease at a time currently.
 | |
| 91 | 97 |          if not bot_session.leases:
 | 
| 92 | 98 |              leases = self._scheduler.request_job_leases({})
 | 
| 93 | 99 |              if leases:
 | 
| 100 | +                for lease in leases:
 | |
| 101 | +                    self._assigned_leases[bot_session.name].add(lease.id)
 | |
| 94 | 102 |                  bot_session.leases.extend(leases)
 | 
| 95 | 103 |  | 
| 96 | -        self._bot_sessions[name] = bot_session
 | |
| 97 | -        return bot_session
 | |
| 98 | - | |
| 99 | 104 |      def _check_lease_state(self, lease):
 | 
| 105 | +        # careful here
 | |
| 106 | +        # should store bot name in scheduler
 | |
| 107 | +        lease_state = LeaseState(lease.state)
 | |
| 108 | + | |
| 109 | +        # Lease has replied with cancelled, remove
 | |
| 110 | +        if lease_state == LeaseState.CANCELLED:
 | |
| 111 | +            return None
 | |
| 100 | 112 |  | 
| 101 | -        # Check for cancelled lease
 | |
| 102 | -        if self._scheduler.get_lease_cancelled(lease.id):
 | |
| 113 | +        try:
 | |
| 114 | +            if self._scheduler.get_job_lease_cancelled(lease.id):
 | |
| 115 | +                lease.state.CopyFrom(LeaseState.CANCELLED.value)
 | |
| 116 | +                return lease
 | |
| 117 | +        except KeyError:
 | |
| 118 | +            # Job does not exist, remove from bot.
 | |
| 103 | 119 |              return None
 | 
| 104 | 120 |  | 
| 105 | -        # If not cancelled, update the status
 | |
| 106 | 121 |          self._scheduler.update_job_lease(lease)
 | 
| 107 | 122 |  | 
| 108 | -        lease_state = LeaseState(lease.state)
 | |
| 109 | 123 |          if lease_state == LeaseState.COMPLETED:
 | 
| 110 | 124 |              return None
 | 
| 111 | 125 |  | 
| 112 | -        elif lease_state == LeaseState.CANCELLED:
 | |
| 113 | -            return None
 | |
| 114 | - | |
| 115 | 126 |          return lease
 | 
| 116 | 127 |  | 
| 117 | 128 |      def _check_bot_ids(self, bot_id, name=None):
 | 
| ... | ... | @@ -134,6 +145,19 @@ class BotsInterface: | 
| 134 | 145 |                          'Bot id already registered. ID sent: [{}].'
 | 
| 135 | 146 |                          'Id registered: [{}] with name: [{}]'.format(bot_id, _bot_id, _name))
 | 
| 136 | 147 |  | 
| 148 | +    def _check_assigned_leases(self, bot_session):
 | |
| 149 | +        session_lease_ids = []
 | |
| 150 | + | |
| 151 | +        for lease in bot_session.leases:
 | |
| 152 | +            session_lease_ids.append(lease.id)
 | |
| 153 | + | |
| 154 | +        for lease_id in self._assigned_leases[bot_session.name]:
 | |
| 155 | +            if lease_id not in session_lease_ids:
 | |
| 156 | +                self.logger.error("Assigned lease id=[{}],"
 | |
| 157 | +                                  " not found on bot with name=[{}] and id=[{}]."
 | |
| 158 | +                                  " Retrying job".format(lease_id, bot_session.name, bot_session.bot_id))
 | |
| 159 | +                self._scheduler.retry_job(lease_id)
 | |
| 160 | + | |
| 137 | 161 |      def _close_bot_session(self, name):
 | 
| 138 | 162 |          """ Before removing the session, close any leases and
 | 
| 139 | 163 |          requeue with high priority.
 | 
| ... | ... | @@ -144,10 +168,9 @@ class BotsInterface: | 
| 144 | 168 |              raise InvalidArgumentError("Bot id does not exist: [{}]".format(name))
 | 
| 145 | 169 |  | 
| 146 | 170 |          self.logger.debug("Attempting to close [{}] with name: [{}]".format(bot_id, name))
 | 
| 147 | -        for lease in self._bot_sessions[name].leases:
 | |
| 148 | -            if lease.state != LeaseState.COMPLETED.value:
 | |
| 149 | -                # TODO: Be wary here, may need to handle rejected leases in future
 | |
| 150 | -                self._scheduler.retry_job(lease.id)
 | |
| 171 | +        for lease_id in self._assigned_leases[name]:
 | |
| 172 | +            self._scheduler.retry_job(lease_id)
 | |
| 173 | +        self._assigned_leases.pop(name)
 | |
| 151 | 174 |  | 
| 152 | 175 |          self.logger.debug("Closing bot session: [{}]".format(name))
 | 
| 153 | 176 |          self._bot_ids.pop(name)
 | 
| ... | ... | @@ -19,11 +19,9 @@ import uuid | 
| 19 | 19 |  from google.protobuf import timestamp_pb2
 | 
| 20 | 20 |  | 
| 21 | 21 |  from buildgrid._enums import LeaseState, OperationStage
 | 
| 22 | -from buildgrid._exceptions import CancelledError
 | |
| 23 | 22 |  from buildgrid._protos.build.bazel.remote.execution.v2 import remote_execution_pb2
 | 
| 24 | 23 |  from buildgrid._protos.google.devtools.remoteworkers.v1test2 import bots_pb2
 | 
| 25 | 24 |  from buildgrid._protos.google.longrunning import operations_pb2
 | 
| 26 | -from buildgrid._protos.google.rpc import code_pb2
 | |
| 27 | 25 |  | 
| 28 | 26 |  | 
| 29 | 27 |  class Job:
 | 
| ... | ... | @@ -36,18 +34,16 @@ class Job: | 
| 36 | 34 |          self._operation = operations_pb2.Operation()
 | 
| 37 | 35 |          self._lease = None
 | 
| 38 | 36 |  | 
| 37 | +        self.__lease_cancelled = False
 | |
| 39 | 38 |          self.__execute_response = None
 | 
| 40 | 39 |          self.__operation_metadata = remote_execution_pb2.ExecuteOperationMetadata()
 | 
| 41 | - | |
| 42 | 40 |          self.__queued_timestamp = timestamp_pb2.Timestamp()
 | 
| 43 | 41 |          self.__worker_start_timestamp = timestamp_pb2.Timestamp()
 | 
| 44 | 42 |          self.__worker_completed_timestamp = timestamp_pb2.Timestamp()
 | 
| 45 | 43 |  | 
| 46 | -        self.__operation_cancelled = False
 | |
| 47 | -        self.__lease_cancelled = False
 | |
| 48 | - | |
| 49 | 44 |          self.__operation_metadata.action_digest.CopyFrom(action_digest)
 | 
| 50 | 45 |          self.__operation_metadata.stage = OperationStage.UNKNOWN.value
 | 
| 46 | +        self.__operation_cancelled = False
 | |
| 51 | 47 |  | 
| 52 | 48 |          self._action.CopyFrom(action)
 | 
| 53 | 49 |          self._do_not_cache = self._action.do_not_cache
 | 
| ... | ... | @@ -140,9 +136,7 @@ class Job: | 
| 140 | 136 |          Only one :class:`Lease` can be emitted for a given job. This method
 | 
| 141 | 137 |          should only be used once, any furhter calls are ignored.
 | 
| 142 | 138 |          """
 | 
| 143 | -        if self.__operation_cancelled:
 | |
| 144 | -            return None
 | |
| 145 | -        elif self._lease is not None:
 | |
| 139 | +        if self._lease is not None:
 | |
| 146 | 140 |              return None
 | 
| 147 | 141 |  | 
| 148 | 142 |          self._lease = bots_pb2.Lease()
 | 
| ... | ... | @@ -183,7 +177,7 @@ class Job: | 
| 183 | 177 |              action_result = remote_execution_pb2.ActionResult()
 | 
| 184 | 178 |  | 
| 185 | 179 |              # TODO: Make a distinction between build and bot failures!
 | 
| 186 | -            if status.code != code_pb2.OK:
 | |
| 180 | +            if status.code != 0:
 | |
| 187 | 181 |                  self._do_not_cache = True
 | 
| 188 | 182 |  | 
| 189 | 183 |              if result is not None:
 | 
| ... | ... | @@ -200,15 +194,6 @@ class Job: | 
| 200 | 194 |              self.__execute_response.cached_result = False
 | 
| 201 | 195 |              self.__execute_response.status.CopyFrom(status)
 | 
| 202 | 196 |  | 
| 203 | -    def cancel_lease(self):
 | |
| 204 | -        """Triggers a job's :class:Lease cancellation.
 | |
| 205 | - | |
| 206 | -        This will not cancel the job's :class:Operation.
 | |
| 207 | -        """
 | |
| 208 | -        self.__lease_cancelled = True
 | |
| 209 | -        if self._lease is not None:
 | |
| 210 | -            self.update_lease_state(LeaseState.CANCELLED)
 | |
| 211 | - | |
| 212 | 197 |      def update_operation_stage(self, stage):
 | 
| 213 | 198 |          """Operates a stage transition for the job's :class:Operation.
 | 
| 214 | 199 |  | 
| ... | ... | @@ -244,6 +229,10 @@ class Job: | 
| 244 | 229 |          if self.__operation_cancelled:
 | 
| 245 | 230 |              raise CancelledError(self.__execute_response.status.message)
 | 
| 246 | 231 |  | 
| 232 | +    def cancel_lease(self):
 | |
| 233 | +        self.__lease_cancelled = True
 | |
| 234 | +        self._update_lease_state(LeaseState.CANCELLED)
 | |
| 235 | + | |
| 247 | 236 |      def cancel_operation(self):
 | 
| 248 | 237 |          """Triggers a job's :class:Operation cancellation.
 | 
| 249 | 238 |  | 
| ... | ... | @@ -74,7 +74,8 @@ class Scheduler: | 
| 74 | 74 |                  # TODO: Mark these jobs as done
 | 
| 75 | 75 |              else:
 | 
| 76 | 76 |                  job.update_operation_stage(OperationStage.QUEUED)
 | 
| 77 | -                self.queue.appendleft(job)
 | |
| 77 | +                job.update_lease_state(LeaseState.PENDING)
 | |
| 78 | +                self.queue.append(job)
 | |
| 78 | 79 |  | 
| 79 | 80 |      def list_jobs(self):
 | 
| 80 | 81 |          return self.jobs.values()
 | 
| ... | ... | @@ -91,8 +92,12 @@ class Scheduler: | 
| 91 | 92 |              return []
 | 
| 92 | 93 |  | 
| 93 | 94 |          job = self.queue.popleft()
 | 
| 94 | -        # For now, one lease at a time:
 | |
| 95 | -        lease = job.create_lease()
 | |
| 95 | + | |
| 96 | +        lease = job.lease
 | |
| 97 | + | |
| 98 | +        if not lease:
 | |
| 99 | +            # For now, one lease at a time:
 | |
| 100 | +            lease = job.create_lease()
 | |
| 96 | 101 |  | 
| 97 | 102 |          if lease:
 | 
| 98 | 103 |              return [lease]
 | 
