GNOME.org
 

[Notes] [Git][BuildGrid/buildgrid][finn/48-cancellation-leases] 24 commits: setup.py: Pin moto dependency and require < 1.3.7

Title: GitLab

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

Commits:

23 changed files:

Changes:

  • buildgrid/_app/bots/dummy.py
    ... ... @@ -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

  • 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/_app/commands/cmd_operation.py
    ... ... @@ -155,6 +155,25 @@ def status(context, operation_name, json):
    155 155 
             click.echo(json_format.MessageToJson(operation))
    156 156
    157 157
    158 
    +@cli.command('cancel', short_help="Cancel an operation.")
    159 
    +@click.argument('operation-name', nargs=1, type=click.STRING, required=True)
    160 
    +@pass_context
    161 
    +def cancel(context, operation_name):
    162 
    +    context.logger.info("Cancelling an operation...")
    163 
    +    stub = operations_pb2_grpc.OperationsStub(context.channel)
    164 
    +
    165 
    +    request = operations_pb2.CancelOperationRequest(name=operation_name)
    166 
    +
    167 
    +    try:
    168 
    +        stub.CancelOperation(request)
    169 
    +    except grpc.RpcError as e:
    170 
    +        status_code = e.code()
    171 
    +        if status_code != grpc.StatusCode.CANCELLED:
    172 
    +            raise
    173 
    +
    174 
    +    context.logger.info("Operation cancelled: [{}]".format(request))
    175 
    +
    176 
    +
    158 177 
     @cli.command('list', short_help="List operations.")
    159 178 
     @click.option('--json', is_flag=True, show_default=True,
    160 179 
                   help="Print operations list in JSON format.")

  • buildgrid/_exceptions.py
    ... ... @@ -52,6 +52,12 @@ class BotError(BgdError):
    52 52 
             super().__init__(message, detail=detail, domain=ErrorDomain.BOT, reason=reason)
    53 53
    54 54
    55 
    +class CancelledError(BgdError):
    56 
    +    """The job was cancelled and any callers should be notified"""
    57 
    +    def __init__(self, message, detail=None, reason=None):
    58 
    +        super().__init__(message, detail=detail, domain=ErrorDomain.SERVER, reason=reason)
    59 
    +
    60 
    +
    55 61 
     class InvalidArgumentError(BgdError):
    56 62 
         """A bad argument was passed, such as a name which doesn't exist."""
    57 63 
         def __init__(self, message, detail=None, reason=None):

  • buildgrid/bot/bot.py
    ... ... @@ -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)

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

  • buildgrid/bot/session.py
    1 
    +# Copyright (C) 2018 Bloomberg LP
    2 
    +#
    3 
    +# Licensed under the Apache License, Version 2.0 (the "License");
    4 
    +# you may not use this file except in compliance with the License.
    5 
    +# You may obtain a copy of the License at
    6 
    +#
    7 
    +#  <http://www.apache.org/licenses/LICENSE-2.0>
    8 
    +#
    9 
    +# Unless required by applicable law or agreed to in writing, software
    10 
    +# distributed under the License is distributed on an "AS IS" BASIS,
    11 
    +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12 
    +# See the License for the specific language governing permissions and
    13 
    +# limitations under the License.
    14 
    +
    15 
    +# 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 pdb
    29 
    +# pdb.set_trace()
    30 
    +
    31 
    +import grpc
    32 
    +
    33 
    +from buildgrid._enums import BotStatus, LeaseState
    34 
    +from buildgrid._protos.google.devtools.remoteworkers.v1test2 import bots_pb2
    35 
    +from buildgrid._protos.google.rpc import code_pb2
    36 
    +from buildgrid._exceptions import BotError
    37 
    +
    38 
    +from buildgrid._exceptions import FailedPreconditionError
    39 
    +
    40 
    +from .tenantmanager import TenantManager
    41 
    +
    42 
    +class BotSession:
    43 
    +    def __init__(self, parent, bots_interface, hardware_interface):
    44 
    +        """ Unique bot ID within the farm used to identify this bot
    45 
    +        Needs to be human readable.
    46 
    +        All prior sessions with bot_id of same ID are invalidated.
    47 
    +        If a bot attempts to update an invalid session, it must be rejected and
    48 
    +        may be put in quarantine.
    49 
    +        """
    50 
    +
    51 
    +        self.logger = logging.getLogger(__name__)
    52 
    +
    53 
    +        self._bots_interface = bots_interface
    54 
    +        self._hardware_interface = hardware_interface
    55 
    +
    56 
    +        self._status = BotStatus.OK.value
    57 
    +        self._tenant_manager = TenantManager()
    58 
    +
    59 
    +        self.__parent = parent
    60 
    +        self.__bot_id = '{}.{}'.format(parent, platform.node())
    61 
    +        self.__name = None
    62 
    +
    63 
    +        # Remove these and add to a worker config in the future
    64 
    +        self._work = None
    65 
    +        self._context = None
    66 
    +
    67 
    +    @property
    68 
    +    def bot_id(self):
    69 
    +        return self.__bot_id
    70 
    +
    71 
    +    def create_bot_session(self, work, context):
    72 
    +        # Drop this when properly adding to the work
    73 
    +        self._work = work
    74 
    +        self._context = context
    75 
    +
    76 
    +        self.logger.debug("Creating bot session")
    77 
    +
    78 
    +        session = self._bots_interface.create_bot_session(self.__parent, self.get_pb2())
    79 
    +        self.__name = session.name
    80 
    +
    81 
    +        self.logger.info("Created bot session with name: [{}]".format(self.__name))
    82 
    +
    83 
    +        for lease in session.leases:
    84 
    +            self._register_lease(lease)
    85 
    +
    86 
    +    def update_bot_session(self):
    87 
    +        self.logger.debug("Updating bot session: [{}]".format(self.__bot_id))
    88 
    +
    89 
    +        session = self._bots_interface.update_bot_session(self.get_pb2())
    90 
    +        server_ids = []
    91 
    +
    92 
    +        for lease in session.leases:
    93 
    +            server_ids.append(lease.id)
    94 
    +
    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)
    105 
    +
    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)
    112 
    +
    113 
    +    def _register_lease(self, lease):
    114 
    +        lease_id = lease.id
    115 
    +        try:
    116 
    +            self._tenant_manager.create_tenancy(lease)
    117 
    +
    118 
    +        except KeyError as e:
    119 
    +            self.logger.debug("Cannot register lease=[{}]. {}".format(lease.id, e))
    120 
    +
    121 
    +        else:
    122 
    +            try:
    123 
    +                self._hardware_interface.configure_hardware(lease.requirements)
    124 
    +
    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)
    128 
    +
    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 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))

  • 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.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 
    +            self._tenants.pop(lease_id)
    60 
    +            self._tasks.pop(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)

  • 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):
    ... ... @@ -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,70 +76,54 @@ 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)
    79 
    +        self._check_assigned_leases(bot_session)
    80 
    +
    81 
    +        for lease in bot_session.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
    82 93
    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 
    -
    94 
    +    def _request_leases(self, bot_session):
    88 95 
             # TODO: Send worker capabilities to the scheduler!
    96 
    +        # Only send one lease at a time currently.
    89 97 
             if not bot_session.leases:
    90 98 
                 leases = self._scheduler.request_job_leases({})
    91 99 
                 if leases:
    100 
    +                for lease in leases:
    101 
    +                    self._assigned_leases[bot_session.name].add(lease.id)
    92 102 
                     bot_session.leases.extend(leases)
    93 103
    94 
    -        self._bot_sessions[name] = bot_session
    95 
    -        return bot_session
    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
    96 112
    97 
    -    def check_states(self, client_lease):
    98 
    -        """ Edge detector for states
    99 
    -        """
    100 
    -        # TODO: Handle cancelled states
    101 113 
             try:
    102 
    -            server_lease = self._scheduler.get_job_lease(client_lease.id)
    114 
    +            if self._scheduler.get_job_lease_cancelled(lease.id):
    115 
    +                lease.state.CopyFrom(LeaseState.CANCELLED.value)
    116 
    +                return lease
    103 117 
             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:
    118 
    +            # Job does not exist, remove from bot.
    119 
    +            return None
    110 120
    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))
    121 
    +        self._scheduler.update_job_lease(lease)
    119 122
    120 
    -        elif server_state == LeaseState.ACTIVE:
    123 
    +        if lease_state == LeaseState.COMPLETED:
    124 
    +            return None
    121 125
    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
    126 
    +        return lease
    146 127
    147 128 
         def _check_bot_ids(self, bot_id, name=None):
    148 129 
             """ Checks the ID and the name of the bot.
    ... ... @@ -164,6 +145,19 @@ class BotsInterface:
    164 145 
                             'Bot id already registered. ID sent: [{}].'
    165 146 
                             'Id registered: [{}] with name: [{}]'.format(bot_id, _bot_id, _name))
    166 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 
    +
    167 161 
         def _close_bot_session(self, name):
    168 162 
             """ Before removing the session, close any leases and
    169 163 
             requeue with high priority.
    ... ... @@ -174,10 +168,9 @@ class BotsInterface:
    174 168 
                 raise InvalidArgumentError("Bot id does not exist: [{}]".format(name))
    175 169
    176 170 
             self.logger.debug("Attempting to close [{}] with name: [{}]".format(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)
    171 
    +        for lease_id in self._assigned_leases[name]:
    172 
    +            self._scheduler.retry_job(lease_id)
    173 
    +        self._assigned_leases.pop(name)
    181 174
    182 175 
             self.logger.debug("Closing bot session: [{}]".format(name))
    183 176 
             self._bot_ids.pop(name)

  • buildgrid/server/execution/instance.py
    ... ... @@ -71,8 +71,11 @@ class ExecutionInstance:
    71 71 
                 raise InvalidArgumentError("Operation name does not exist: [{}]".format(name))
    72 72
    73 73 
         def stream_operation_updates(self, message_queue, operation_name):
    74 
    -        operation = message_queue.get()
    75 
    -        while not operation.done:
    76 
    -            yield operation
    77 
    -            operation = message_queue.get()
    78 
    -        yield operation
    74 
    +        job = message_queue.get()
    75 
    +        while not job.operation.done:
    76 
    +            yield job.operation
    77 
    +            job = message_queue.get()
    78 
    +
    79 
    +        job.check_operation_status()
    80 
    +
    81 
    +        yield job.operation

  • buildgrid/server/execution/service.py
    ... ... @@ -26,7 +26,7 @@ from functools import partial
    26 26
    27 27 
     import grpc
    28 28
    29 
    -from buildgrid._exceptions import FailedPreconditionError, InvalidArgumentError
    29 
    +from buildgrid._exceptions import FailedPreconditionError, InvalidArgumentError, CancelledError
    30 30 
     from buildgrid._protos.build.bazel.remote.execution.v2 import remote_execution_pb2_grpc
    31 31 
     from buildgrid._protos.google.longrunning import operations_pb2
    32 32
    ... ... @@ -76,6 +76,12 @@ class ExecutionService(remote_execution_pb2_grpc.ExecutionServicer):
    76 76 
                 context.set_code(grpc.StatusCode.FAILED_PRECONDITION)
    77 77 
                 yield operations_pb2.Operation()
    78 78
    79 
    +        except CancelledError as e:
    80 
    +            self.logger.error(e)
    81 
    +            context.set_details(str(e))
    82 
    +            context.set_code(grpc.StatusCode.CANCELLED)
    83 
    +            yield operations_pb2.Operation()
    84 
    +
    79 85 
         def WaitExecution(self, request, context):
    80 86 
             try:
    81 87 
                 names = request.name.split("/")
    ... ... @@ -106,6 +112,12 @@ class ExecutionService(remote_execution_pb2_grpc.ExecutionServicer):
    106 112 
                 context.set_code(grpc.StatusCode.INVALID_ARGUMENT)
    107 113 
                 yield operations_pb2.Operation()
    108 114
    115 
    +        except CancelledError as e:
    116 
    +            self.logger.error(e)
    117 
    +            context.set_details(str(e))
    118 
    +            context.set_code(grpc.StatusCode.CANCELLED)
    119 
    +            yield operations_pb2.Operation()
    120 
    +
    109 121 
         def _get_instance(self, name):
    110 122 
             try:
    111 123 
                 return self._instances[name]

  • buildgrid/server/job.py
    ... ... @@ -34,6 +34,7 @@ class Job:
    34 34 
             self._operation = operations_pb2.Operation()
    35 35 
             self._lease = None
    36 36
    37 
    +        self.__lease_cancelled = False
    37 38 
             self.__execute_response = None
    38 39 
             self.__operation_metadata = remote_execution_pb2.ExecuteOperationMetadata()
    39 40 
             self.__queued_timestamp = timestamp_pb2.Timestamp()
    ... ... @@ -42,6 +43,7 @@ class Job:
    42 43
    43 44 
             self.__operation_metadata.action_digest.CopyFrom(action_digest)
    44 45 
             self.__operation_metadata.stage = OperationStage.UNKNOWN.value
    46 
    +        self.__operation_cancelled = False
    45 47
    46 48 
             self._action.CopyFrom(action)
    47 49 
             self._do_not_cache = self._action.do_not_cache
    ... ... @@ -92,6 +94,10 @@ class Job:
    92 94 
             else:
    93 95 
                 return None
    94 96
    97 
    +    @property
    98 
    +    def lease_cancelled(self):
    99 
    +        return self.__lease_cancelled
    100 
    +
    95 101 
         @property
    96 102 
         def n_tries(self):
    97 103 
             return self._n_tries
    ... ... @@ -107,7 +113,7 @@ class Job:
    107 113 
                 queue (queue.Queue): the event queue to register.
    108 114 
             """
    109 115 
             self._operation_update_queues.append(queue)
    110 
    -        queue.put(self._operation)
    116 
    +        queue.put(self)
    111 117
    112 118 
         def unregister_client(self, queue):
    113 119 
             """Unsubscribes to the job's :class:`Operation` stage change events.
    ... ... @@ -212,4 +218,34 @@ class Job:
    212 218 
             self._operation.metadata.Pack(self.__operation_metadata)
    213 219
    214 220 
             for queue in self._operation_update_queues:
    215 
    -            queue.put(self._operation)
    221 
    +            queue.put(self)
    222 
    +
    223 
    +    def check_operation_status(self):
    224 
    +        """Reports errors on unexpected job's :class:Operation state.
    225 
    +
    226 
    +        Raises:
    227 
    +            CancelledError: if the job's :class:Operation was cancelled.
    228 
    +        """
    229 
    +        if self.__operation_cancelled:
    230 
    +            raise CancelledError(self.__execute_response.status.message)
    231 
    +
    232 
    +    def cancel_lease(self):
    233 
    +        self.__lease_cancelled = True
    234 
    +        self._update_lease_state(LeaseState.CANCELLED)
    235 
    +
    236 
    +    def cancel_operation(self):
    237 
    +        """Triggers a job's :class:Operation cancellation.
    238 
    +
    239 
    +        This will also cancel any job's :class:Lease that may have been issued.
    240 
    +        """
    241 
    +        self.__operation_cancelled = True
    242 
    +        if self._lease is not None:
    243 
    +            self.cancel_lease()
    244 
    +
    245 
    +        self.__execute_response = remote_execution_pb2.ExecuteResponse()
    246 
    +        self.__execute_response.status.code = code_pb2.CANCELLED
    247 
    +        self.__execute_response.status.message = "Operation cancelled by client."
    248 
    +
    249 
    +        self.update_operation_stage(OperationStage.COMPLETED)
    250 
    +
    251 
    +        raise CancelledError("Operation cancelled: {}".format(self._name))

  • buildgrid/server/operations/instance.py
    ... ... @@ -64,6 +64,13 @@ class OperationsInstance:
    64 64 
             except KeyError:
    65 65 
                 raise InvalidArgumentError("Operation name does not exist: [{}]".format(name))
    66 66
    67 
    +    def cancel_operation(self, name):
    68 
    +        try:
    69 
    +            self._scheduler.cancel_job_operation(name)
    70 
    +
    71 
    +        except KeyError:
    72 
    +            raise InvalidArgumentError("Operation name does not exist: [{}]".format(name))
    73 
    +
    67 74 
         def register_message_client(self, name, queue):
    68 75 
             try:
    69 76 
                 self._scheduler.register_client(name, queue)
    ... ... @@ -79,12 +86,11 @@ class OperationsInstance:
    79 86 
                 raise InvalidArgumentError("Operation name does not exist: [{}]".format(name))
    80 87
    81 88 
         def stream_operation_updates(self, message_queue, operation_name):
    82 
    -        operation = message_queue.get()
    83 
    -        while not operation.done:
    84 
    -            yield operation
    85 
    -            operation = message_queue.get()
    86 
    -        yield operation
    89 
    +        job = message_queue.get()
    90 
    +        while not job.operation.done:
    91 
    +            yield job.operation
    92 
    +            job = message_queue.get()
    87 93
    88 
    -    def cancel_operation(self, name):
    89 
    -        # TODO: Cancel leases
    90 
    -        raise NotImplementedError("Cancelled operations not supported")
    94 
    +        job.check_operation_status()
    95 
    +
    96 
    +        yield job.operation

  • buildgrid/server/operations/service.py
    ... ... @@ -25,7 +25,7 @@ import grpc
    25 25
    26 26 
     from google.protobuf.empty_pb2 import Empty
    27 27
    28 
    -from buildgrid._exceptions import InvalidArgumentError
    28 
    +from buildgrid._exceptions import CancelledError, InvalidArgumentError
    29 29 
     from buildgrid._protos.google.longrunning import operations_pb2_grpc, operations_pb2
    30 30
    31 31
    ... ... @@ -112,10 +112,10 @@ class OperationsService(operations_pb2_grpc.OperationsServicer):
    112 112 
                 operation_name = self._parse_operation_name(name)
    113 113 
                 instance.cancel_operation(operation_name)
    114 114
    115 
    -        except NotImplementedError as e:
    115 
    +        except CancelledError as e:
    116 116 
                 self.logger.error(e)
    117 117 
                 context.set_details(str(e))
    118 
    -            context.set_code(grpc.StatusCode.UNIMPLEMENTED)
    118 
    +            context.set_code(grpc.StatusCode.CANCELLED)
    119 119
    120 120 
             except InvalidArgumentError as e:
    121 121 
                 self.logger.error(e)

  • buildgrid/server/scheduler.py
    ... ... @@ -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,12 +92,19 @@ 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()
    96 95
    97 
    -        return [lease]
    96 
    +        lease = job.lease
    97 
    +
    98 
    +        if not lease:
    99 
    +            # For now, one lease at a time:
    100 
    +            lease = job.create_lease()
    101 
    +
    102 
    +        if lease:
    103 
    +            return [lease]
    98 104
    99 
    -    def update_job_lease_state(self, job_name, lease_state, lease_status=None, lease_result=None):
    105 
    +        return None
    106 
    +
    107 
    +    def update_job_lease(self, lease):
    100 108 
             """Requests a state transition for a job's current :class:Lease.
    101 109
    102 110 
             Args:
    ... ... @@ -107,7 +115,9 @@ class Scheduler:
    107 115 
                 lease_result (google.protobuf.Any): the lease execution result, only
    108 116 
                     required if `lease_state` is `COMPLETED`.
    109 117 
             """
    110 
    -        job = self.jobs[job_name]
    118 
    +
    119 
    +        job = self.jobs[lease.id]
    120 
    +        lease_state = LeaseState(lease.state)
    111 121
    112 122 
             if lease_state == LeaseState.PENDING:
    113 123 
                 job.update_lease_state(LeaseState.PENDING)
    ... ... @@ -119,7 +129,7 @@ class Scheduler:
    119 129
    120 130 
             elif lease_state == LeaseState.COMPLETED:
    121 131 
                 job.update_lease_state(LeaseState.COMPLETED,
    122 
    -                                   status=lease_status, result=lease_result)
    132 
    +                                   status=lease.status, result=lease.result)
    123 133
    124 134 
                 if self._action_cache is not None and not job.do_not_cache:
    125 135 
                     self._action_cache.update_action_result(job.action_digest, job.action_result)
    ... ... @@ -130,6 +140,20 @@ class Scheduler:
    130 140 
             """Returns the lease associated to job, if any have been emitted yet."""
    131 141 
             return self.jobs[job_name].lease
    132 142
    143 
    +    def get_job_lease_cancelled(self, job_name):
    144 
    +        """Returns true if the lease is cancelled"""
    145 
    +        return self.jobs[job_name].lease_cancelled
    146 
    +
    133 147 
         def get_job_operation(self, job_name):
    134 148 
             """Returns the operation associated to job."""
    135 149 
             return self.jobs[job_name].operation
    150 
    +
    151 
    +    def cancel_job_operation(self, job_name):
    152 
    +        """"Cancels the underlying operation of a given job.
    153 
    +
    154 
    +        This will also cancel any job's lease that may have been issued.
    155 
    +
    156 
    +        Args:
    157 
    +            job_name (str): name of the job holding the operation to cancel.
    158 
    +        """
    159 
    +        self.jobs[job_name].cancel_operation()

  • setup.py
    ... ... @@ -87,7 +87,7 @@ def get_cmdclass():
    87 87
    88 88 
     tests_require = [
    89 89 
         'coverage >= 4.5.0',
    90 
    -    'moto',
    90 
    +    'moto < 1.3.7',
    91 91 
         'pep8',
    92 92 
         'psutil',
    93 93 
         'pytest >= 3.8.0',

  • tests/integration/operations_service.py
    ... ... @@ -24,6 +24,7 @@ import grpc
    24 24 
     from grpc._server import _Context
    25 25 
     import pytest
    26 26
    27 
    +from buildgrid._enums import OperationStage
    27 28 
     from buildgrid._exceptions import InvalidArgumentError
    28 29 
     from buildgrid._protos.build.bazel.remote.execution.v2 import remote_execution_pb2
    29 30 
     from buildgrid._protos.google.longrunning import operations_pb2
    ... ... @@ -236,12 +237,26 @@ def test_delete_operation_fail(instance, context):
    236 237 
         context.set_code.assert_called_once_with(grpc.StatusCode.INVALID_ARGUMENT)
    237 238
    238 239
    239 
    -def test_cancel_operation(instance, context):
    240 
    +def test_cancel_operation(instance, controller, execute_request, context):
    241 
    +    response_execute = controller.execution_instance.execute(execute_request.action_digest,
    242 
    +                                                             execute_request.skip_cache_lookup)
    243 
    +
    240 244 
         request = operations_pb2.CancelOperationRequest()
    241 
    -    request.name = "{}/{}".format(instance_name, "runner")
    245 
    +    request.name = "{}/{}".format(instance_name, response_execute.name)
    246 
    +
    242 247 
         instance.CancelOperation(request, context)
    243 248
    244 
    -    context.set_code.assert_called_once_with(grpc.StatusCode.UNIMPLEMENTED)
    249 
    +    context.set_code.assert_called_once_with(grpc.StatusCode.CANCELLED)
    250 
    +
    251 
    +    request = operations_pb2.ListOperationsRequest(name=instance_name)
    252 
    +    response = instance.ListOperations(request, context)
    253 
    +
    254 
    +    assert len(response.operations) is 1
    255 
    +
    256 
    +    for operation in response.operations:
    257 
    +        operation_metadata = remote_execution_pb2.ExecuteOperationMetadata()
    258 
    +        operation.metadata.Unpack(operation_metadata)
    259 
    +        assert operation_metadata.stage == OperationStage.COMPLETED.value
    245 260
    246 261
    247 262 
     def test_cancel_operation_blank(blank_instance, context):
    ... ... @@ -249,7 +264,7 @@ def test_cancel_operation_blank(blank_instance, context):
    249 264 
         request.name = "runner"
    250 265 
         blank_instance.CancelOperation(request, context)
    251 266
    252 
    -    context.set_code.assert_called_once_with(grpc.StatusCode.UNIMPLEMENTED)
    267 
    +    context.set_code.assert_called_once_with(grpc.StatusCode.INVALID_ARGUMENT)
    253 268
    254 269
    255 270 
     def test_cancel_operation_instance_fail(instance, context):

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

    [Date Prev][Date Next]   [Thread Prev][Thread Next]   [Thread Index] [Date Index] [Author Index]