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



Title: GitLab

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

Commits:

14 changed files:

Changes:

  • buildgrid/_app/bots/dummy.py
    ... ... @@ -20,7 +20,7 @@ from buildgrid._protos.build.bazel.remote.execution.v2 import remote_execution_p
    20 20
     from buildgrid.utils import get_hostname
    
    21 21
     
    
    22 22
     
    
    23
    -def work_dummy(context, lease):
    
    23
    +def work_dummy(lease, context, event):
    
    24 24
         """ Just returns lease after some random time
    
    25 25
         """
    
    26 26
         action_result = remote_execution_pb2.ActionResult()
    

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

  • buildgrid/bot/bot.py
    ... ... @@ -22,7 +22,6 @@ Creates a bot session and sends updates to the server.
    22 22
     
    
    23 23
     import asyncio
    
    24 24
     import logging
    
    25
    -import sys
    
    26 25
     
    
    27 26
     
    
    28 27
     class Bot:
    

  • buildgrid/bot/hardware/device.py
    ... ... @@ -20,10 +20,10 @@ Device
    20 20
     A device.
    
    21 21
     """
    
    22 22
     
    
    23
    -
    
    24 23
     import uuid
    
    25 24
     from buildgrid._protos.google.devtools.remoteworkers.v1test2 import worker_pb2
    
    26 25
     
    
    26
    +
    
    27 27
     class Device:
    
    28 28
     
    
    29 29
         def __init__(self, properties=None):
    
    ... ... @@ -33,7 +33,7 @@ class Device:
    33 33
             All other devices are known as Attatched Devices and must be controlled
    
    34 34
             by the Primary Device.
    
    35 35
     
    
    36
    -        properties (list(dict(string : string))) : Properties of device. Keys may
    
    36
    +        properties (dict(key: list())) : Properties of device. Keys may
    
    37 37
             repeated.
    
    38 38
             """
    
    39 39
     
    
    ... ... @@ -42,8 +42,9 @@ class Device:
    42 42
             self.__name = str(uuid.uuid4())
    
    43 43
     
    
    44 44
             if properties:
    
    45
    -            for prop in properties:
    
    46
    -                self._add_property(prop)
    
    45
    +            for k, l in properties.items():
    
    46
    +                for v in l:
    
    47
    +                    self._add_property(k, v)
    
    47 48
     
    
    48 49
         @property
    
    49 50
         def name(self):
    

  • buildgrid/bot/hardware/worker.py
    ... ... @@ -13,6 +13,13 @@
    13 13
     # limitations under the License.
    
    14 14
     
    
    15 15
     
    
    16
    +"""
    
    17
    +Worker
    
    18
    +======
    
    19
    +
    
    20
    +A worker.
    
    21
    +"""
    
    22
    +
    
    16 23
     from buildgrid._protos.google.devtools.remoteworkers.v1test2 import worker_pb2
    
    17 24
     
    
    18 25
     
    
    ... ... @@ -26,13 +33,14 @@ class Worker:
    26 33
             self.__config_keys = ['DockerImage']
    
    27 34
     
    
    28 35
             if properties:
    
    29
    -            for k, v in properties.items():
    
    30
    -                if k in self.__property_keys:
    
    31
    -                    self._add_properties(k, v)
    
    36
    +            for k, l in properties.items():
    
    37
    +                for v in l:
    
    38
    +                    self._add_property(k, v)
    
    32 39
     
    
    33 40
             if configs:
    
    34
    -            for k, v in configs.items():
    
    35
    -                self._add_config(k, v)
    
    41
    +            for k, l in configs.items():
    
    42
    +                for v in l:
    
    43
    +                    self._add_config(k, v)
    
    36 44
     
    
    37 45
         @property
    
    38 46
         def configs(self):
    
    ... ... @@ -54,12 +62,12 @@ class Worker:
    54 62
                     property_message = worker_pb2.Device.Property()
    
    55 63
                     property_message.key = k
    
    56 64
                     property_message.value = prop
    
    57
    -                device.properties.extend([property_message])
    
    65
    +                worker.properties.extend([property_message])
    
    58 66
     
    
    59 67
             for k, v in self._configs.items():
    
    60 68
                 for cfg in v:
    
    61 69
                     config_message = worker_pb2.Worker.Config()
    
    62
    -                config.key = k
    
    70
    +                config_message.key = k
    
    63 71
                     config_message.value = cfg
    
    64 72
                     worker.configs.extend([config_message])
    
    65 73
     
    

  • buildgrid/bot/interface.py
    ... ... @@ -21,6 +21,7 @@ Interface to grpc
    21 21
     """
    
    22 22
     
    
    23 23
     import logging
    
    24
    +import grpc
    
    24 25
     
    
    25 26
     from buildgrid._protos.google.devtools.remoteworkers.v1test2 import bots_pb2, bots_pb2_grpc
    
    26 27
     
    
    ... ... @@ -40,7 +41,8 @@ class BotInterface:
    40 41
                                                        bot_session=bot_session)
    
    41 42
             try:
    
    42 43
                 return self._stub.CreateBotSession(request)
    
    43
    -        except Exception as e:
    
    44
    +
    
    45
    +        except grpc.RpcError as e:
    
    44 46
                 self.logger.error("Error creating bot session: [{}]".format(e))
    
    45 47
     
    
    46 48
         def update_bot_session(self, bot_session, update_mask=None):
    
    ... ... @@ -49,5 +51,6 @@ class BotInterface:
    49 51
                                                        update_mask=update_mask)
    
    50 52
             try:
    
    51 53
                 return self._stub.UpdateBotSession(request)
    
    52
    -        except Exception as e:
    
    54
    +
    
    55
    +        except grpc.RpcError as e:
    
    53 56
                 self.logger.error("Error updating bot session: [{}]".format(e))

  • buildgrid/bot/session.py
    ... ... @@ -12,9 +12,6 @@
    12 12
     # See the License for the specific language governing permissions and
    
    13 13
     # limitations under the License.
    
    14 14
     
    
    15
    -# Disable broad exception catch
    
    16
    -# pylint: disable=broad-except
    
    17
    -
    
    18 15
     
    
    19 16
     """
    
    20 17
     Bot Session
    
    ... ... @@ -22,25 +19,20 @@ Bot Session
    22 19
     
    
    23 20
     Allows connections
    
    24 21
     """
    
    25
    -import asyncio
    
    26 22
     import logging
    
    27 23
     import platform
    
    28
    -import pdb
    
    29
    -# pdb.set_trace()
    
    30
    -
    
    31
    -import grpc
    
    32 24
     
    
    33 25
     from buildgrid._enums import BotStatus, LeaseState
    
    34 26
     from buildgrid._protos.google.devtools.remoteworkers.v1test2 import bots_pb2
    
    35 27
     from buildgrid._protos.google.rpc import code_pb2
    
    36
    -from buildgrid._exceptions import BotError
    
    37 28
     
    
    38 29
     from buildgrid._exceptions import FailedPreconditionError
    
    39 30
     
    
    40 31
     from .tenantmanager import TenantManager
    
    41 32
     
    
    33
    +
    
    42 34
     class BotSession:
    
    43
    -    def __init__(self, parent, bots_interface, hardware_interface):
    
    35
    +    def __init__(self, parent, bots_interface, hardware_interface, work, context=None):
    
    44 36
             """ Unique bot ID within the farm used to identify this bot
    
    45 37
             Needs to be human readable.
    
    46 38
             All prior sessions with bot_id of same ID are invalidated.
    
    ... ... @@ -60,19 +52,14 @@ class BotSession:
    60 52
             self.__bot_id = '{}.{}'.format(parent, platform.node())
    
    61 53
             self.__name = None
    
    62 54
     
    
    63
    -        # Remove these and add to a worker config in the future
    
    64
    -        self._work = None
    
    65
    -        self._context = None
    
    55
    +        self._work = work
    
    56
    +        self._context = context
    
    66 57
     
    
    67 58
         @property
    
    68 59
         def bot_id(self):
    
    69 60
             return self.__bot_id
    
    70 61
     
    
    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
    -
    
    62
    +    def create_bot_session(self):
    
    76 63
             self.logger.debug("Creating bot session")
    
    77 64
     
    
    78 65
             session = self._bots_interface.create_bot_session(self.__parent, self.get_pb2())
    
    ... ... @@ -97,10 +84,11 @@ class BotSession:
    97 84
                     self._register_lease(lease)
    
    98 85
     
    
    99 86
                 elif lease_state == LeaseState.CANCELLED:
    
    100
    -                self._tenant_manager.cancel_tenancy(lease_id)
    
    87
    +                self._tenant_manager.cancel_tenancy(lease.id)
    
    101 88
     
    
    102 89
             closed_lease_ids = [x for x in self._tenant_manager.get_lease_ids() if x not in server_ids]
    
    103 90
             for lease_id in closed_lease_ids:
    
    91
    +            self._tenant_manager.cancel_tenancy(lease_id)
    
    104 92
                 self._tenant_manager.remove_tenant(lease_id)
    
    105 93
     
    
    106 94
         def get_pb2(self):
    

  • buildgrid/bot/tenant.py
    ... ... @@ -12,6 +12,10 @@
    12 12
     # See the License for the specific language governing permissions and
    
    13 13
     # limitations under the License.
    
    14 14
     
    
    15
    +# Disable broad exception catch
    
    16
    +# pylint: disable=broad-except
    
    17
    +
    
    18
    +
    
    15 19
     """
    
    16 20
     Tenant
    
    17 21
     ======
    
    ... ... @@ -21,12 +25,15 @@ Handles leased and runs leased work.
    21 25
     
    
    22 26
     import asyncio
    
    23 27
     import logging
    
    28
    +import threading
    
    24 29
     
    
    25 30
     from functools import partial
    
    26 31
     
    
    27
    -from buildgrid._protos.google.devtools.remoteworkers.v1test2 import bots_pb2
    
    32
    +import grpc
    
    28 33
     
    
    34
    +from buildgrid._protos.google.rpc import code_pb2
    
    29 35
     from buildgrid._enums import LeaseState
    
    36
    +from buildgrid._exceptions import BotError
    
    30 37
     
    
    31 38
     
    
    32 39
     class Tenant:
    
    ... ... @@ -37,14 +44,28 @@ class Tenant:
    37 44
                 raise ValueError("Lease state not `PENDING`: {}".format(lease.state))
    
    38 45
     
    
    39 46
             self.logger = logging.getLogger(__name__)
    
    40
    -        self.lease_finished = False
    
    41 47
     
    
    42 48
             self._lease = lease
    
    43 49
     
    
    50
    +        self.__lease_cancelled = False
    
    51
    +        self.__tenant_completed = False
    
    52
    +
    
    44 53
         @property
    
    45 54
         def lease(self):
    
    46 55
             return self._lease
    
    47 56
     
    
    57
    +    @property
    
    58
    +    def tenant_completed(self):
    
    59
    +        return self.__tenant_completed
    
    60
    +
    
    61
    +    @property
    
    62
    +    def lease_cancelled(self):
    
    63
    +        return self.__lease_cancelled
    
    64
    +
    
    65
    +    def cancel_lease(self):
    
    66
    +        self.__lease_cancelled = True
    
    67
    +        self.update_lease_state(LeaseState.CANCELLED)
    
    68
    +
    
    48 69
         def get_lease_state(self):
    
    49 70
             return LeaseState(self._lease.state)
    
    50 71
     
    
    ... ... @@ -58,17 +79,19 @@ class Tenant:
    58 79
             self.logger.debug("Work created: [{}]".format(self._lease.id))
    
    59 80
     
    
    60 81
             # 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 82
     
    
    64 83
             loop = asyncio.get_event_loop()
    
    65 84
     
    
    66 85
             try:
    
    67
    -            lease = await loop.run_in_executor(executor, partial(work, context, self._lease))
    
    86
    +            event = threading.Event()
    
    87
    +            lease = await loop.run_in_executor(executor, partial(work, self._lease, context, event))
    
    68 88
                 self._lease.CopyFrom(lease)
    
    69 89
     
    
    70
    -        except asyncio.CancelledError as e:
    
    71
    -            self.logger.error("Task cancelled: [{}]".format(e))
    
    90
    +        except asyncio.CancelledError:
    
    91
    +            self.logger.error("Lease cancelled: [{}]".format(self._lease.id))
    
    92
    +            event.set()
    
    93
    +            # Propagate error to task wrapper
    
    94
    +            raise
    
    72 95
     
    
    73 96
             except grpc.RpcError as e:
    
    74 97
                 self.logger.error("RPC error thrown: [{}]".format(e))
    
    ... ... @@ -82,4 +105,7 @@ class Tenant:
    82 105
                 self.logger.error("Exception thrown: [{}]".format(e))
    
    83 106
                 lease.status.code = code_pb2.INTERNAL
    
    84 107
     
    
    108
    +        self.__tenant_completed = True
    
    85 109
             self.logger.debug("Work completed: [{}]".format(lease.id))
    
    110
    +
    
    111
    +        return lease

  • buildgrid/bot/tenantmanager.py
    ... ... @@ -20,17 +20,15 @@ TenantManager
    20 20
     Looks after leases of work.
    
    21 21
     """
    
    22 22
     
    
    23
    -
    
    24 23
     import asyncio
    
    25 24
     import logging
    
    26 25
     from functools import partial
    
    27 26
     
    
    28
    -import grpc
    
    29
    -
    
    30 27
     from buildgrid._enums import LeaseState
    
    31 28
     
    
    32 29
     from .tenant import Tenant
    
    33 30
     
    
    31
    +
    
    34 32
     class TenantManager:
    
    35 33
     
    
    36 34
         def __init__(self):
    
    ... ... @@ -49,15 +47,19 @@ class TenantManager:
    49 47
                 raise KeyError("Lease id already exists: [{}]".format(lease_id))
    
    50 48
     
    
    51 49
         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))
    
    50
    +        if not self._tenants[lease_id].lease_cancelled:
    
    51
    +            self.logger.error("Attempting to remove a lease not cancelled."
    
    52
    +                              "Bot will attempt to cancel lease."
    
    53
    +                              "Lease id=[{}]".format(lease_id))
    
    54
    +            self.cancel_tenancy(lease_id)
    
    57 55
     
    
    58
    -        else:
    
    59
    -            self._tenants.pop(lease_id)
    
    60
    -            self._tasks.pop(lease_id)
    
    56
    +        elif not self._tenants[lease_id].tenant_completed:
    
    57
    +            self.logger.debug("Lease cancelled but tenant not completed."
    
    58
    +                              "Lease=[{}]".format(self._tenants[lease_id].lease))
    
    59
    +
    
    60
    +        self.logger.debug("Removing tenant=[{}]".format(lease_id))
    
    61
    +        self._tenants.pop(lease_id)
    
    62
    +        self._tasks.pop(lease_id)
    
    61 63
     
    
    62 64
         def get_leases(self):
    
    63 65
             leases = []
    
    ... ... @@ -79,8 +81,9 @@ class TenantManager:
    79 81
             if status is not None:
    
    80 82
                 self._update_lease_status(lease_id, status)
    
    81 83
     
    
    82
    -        if self._tenants[lease_id].get_lease_state() != LeaseState.CANCELLED:
    
    83
    -            self._update_lease_state(lease_id, LeaseState.COMPLETED)
    
    84
    +        if task:
    
    85
    +            if not task.cancelled():
    
    86
    +                self._update_lease_state(lease_id, LeaseState.COMPLETED)
    
    84 87
     
    
    85 88
         def create_work(self, lease_id, work, context):
    
    86 89
             self._update_lease_state(lease_id, LeaseState.ACTIVE)
    
    ... ... @@ -92,8 +95,12 @@ class TenantManager:
    92 95
             self._tasks[lease_id] = task
    
    93 96
     
    
    94 97
         def cancel_tenancy(self, lease_id):
    
    95
    -        self._update_lease_state(LeaseState.CANCELLED)
    
    96
    -        self._tasks[lease_id].cancel()
    
    98
    +        if not self._tenants[lease_id].lease_cancelled:
    
    99
    +            self._tenants[lease_id].cancel_lease()
    
    100
    +            self._tasks[lease_id].cancel()
    
    101
    +
    
    102
    +    def tenant_completed(self, lease_id):
    
    103
    +        return self._tenants[lease_id].tenant_completed
    
    97 104
     
    
    98 105
         def _update_lease_state(self, lease_id, state):
    
    99 106
             self._tenants[lease_id].update_lease_state(state)
    

  • buildgrid/server/bots/instance.py
    ... ... @@ -23,7 +23,7 @@ Instance of the Remote Workers interface.
    23 23
     import logging
    
    24 24
     import uuid
    
    25 25
     
    
    26
    -from buildgrid._exceptions import InvalidArgumentError, OutOfSyncError
    
    26
    +from buildgrid._exceptions import InvalidArgumentError
    
    27 27
     
    
    28 28
     from ..job import LeaseState
    
    29 29
     
    
    ... ... @@ -34,7 +34,7 @@ class BotsInterface:
    34 34
             self.logger = logging.getLogger(__name__)
    
    35 35
     
    
    36 36
             self._bot_ids = {}
    
    37
    -        self._bot_sessions = {}
    
    37
    +        self._assigned_leases = {}
    
    38 38
             self._scheduler = scheduler
    
    39 39
     
    
    40 40
         def register_instance_with_server(self, instance_name, server):
    
    ... ... @@ -59,18 +59,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)
    

  • tests/integration/bot_session.py
    ... ... @@ -14,56 +14,36 @@
    14 14
     
    
    15 15
     # pylint: disable=redefined-outer-name
    
    16 16
     
    
    17
    -import uuid
    
    17
    +from unittest import mock
    
    18 18
     
    
    19
    +import grpc
    
    19 20
     import pytest
    
    20 21
     
    
    21
    -from buildgrid.bot import bot_session
    
    22
    +from buildgrid._protos.google.devtools.remoteworkers.v1test2 import bots_pb2, bots_pb2_grpc
    
    23
    +from buildgrid.bot.hardware.worker import Worker
    
    24
    +from buildgrid.bot.hardware.interface import HardwareInterface
    
    25
    +from buildgrid.bot.session import BotSession
    
    26
    +from buildgrid.bot.interface import BotInterface
    
    22 27
     
    
    28
    +from ..utils.bots_interface import Server, serve_bots_interface, run_in_subprocess
    
    23 29
     
    
    24
    -@pytest.mark.parametrize("docker_value", ["True", "False"])
    
    25
    -@pytest.mark.parametrize("os_value", ["nexus7", "nexus8"])
    
    26
    -def test_create_device(docker_value, os_value):
    
    27
    -    properties = {'docker': docker_value, 'os': os_value}
    
    28
    -    device = bot_session.Device(properties)
    
    29 30
     
    
    30
    -    assert uuid.UUID(device.name, version=4)
    
    31
    -    assert properties == device.properties
    
    31
    +INSTANCES = ['', 'instance']
    
    32 32
     
    
    33 33
     
    
    34
    -def test_create_device_key_fail():
    
    35
    -    properties = {'voight': 'kampff'}
    
    34
    +@pytest.mark.parametrize('instance', INSTANCES)
    
    35
    +def test_create_bot_session(instance):
    
    36 36
     
    
    37
    -    with pytest.raises(KeyError):
    
    38
    -        bot_session.Device(properties)
    
    37
    +    def __create_bot_session(queue, remote, instance):
    
    38
    +        interface = BotInterface(grpc.insecure_channel(remote))
    
    39
    +        hardware_interface = HardwareInterface(Worker())
    
    40
    +        session = BotSession(instance, interface, hardware_interface, None)
    
    41
    +        session.create_bot_session()
    
    39 42
     
    
    43
    +        queue.put(session.get_pb2().SerializeToString())
    
    40 44
     
    
    41
    -def test_create_device_value_fail():
    
    42
    -    properties = {'docker': True}
    
    45
    +    with serve_bots_interface([instance]) as server:
    
    46
    +        result = run_in_subprocess(__create_bot_session,
    
    47
    +                                   server.remote, instance)
    
    43 48
     
    
    44
    -    with pytest.raises(ValueError):
    
    45
    -        bot_session.Device(properties)
    
    46
    -
    
    47
    -
    
    48
    -def test_create_worker():
    
    49
    -    properties = {'pool': 'swim'}
    
    50
    -    configs = {'DockerImage': 'Windows'}
    
    51
    -    worker = bot_session.Worker(properties, configs)
    
    52
    -
    
    53
    -    assert properties == worker.properties
    
    54
    -    assert configs == worker.configs
    
    55
    -
    
    56
    -    device = bot_session.Device()
    
    57
    -    worker.add_device(device)
    
    58
    -
    
    59
    -    assert worker._devices[0] == device
    
    60
    -
    
    61
    -
    
    62
    -def test_create_worker_key_fail():
    
    63
    -    properties = {'voight': 'kampff'}
    
    64
    -    configs = {'voight': 'kampff'}
    
    65
    -
    
    66
    -    with pytest.raises(KeyError):
    
    67
    -        bot_session.Worker(properties)
    
    68
    -    with pytest.raises(KeyError):
    
    69
    -        bot_session.Worker(configs)
    49
    +        server.compare_bot_sessions(result)

  • tests/integration/bots_service.py
    ... ... @@ -17,7 +17,6 @@
    17 17
     
    
    18 18
     # pylint: disable=redefined-outer-name
    
    19 19
     
    
    20
    -import copy
    
    21 20
     from unittest import mock
    
    22 21
     
    
    23 22
     import grpc
    
    ... ... @@ -151,126 +150,11 @@ def test_update_leases_with_work(bot_session, context, instance):
    151 150
     
    
    152 151
     
    
    153 152
     def test_update_leases_work_complete(bot_session, context, instance):
    
    154
    -    request = bots_pb2.CreateBotSessionRequest(parent='',
    
    155
    -                                               bot_session=bot_session)
    
    156
    -    # Create bot session
    
    157
    -    # Simulated the severed binding between client and server
    
    158
    -    response = copy.deepcopy(instance.CreateBotSession(request, context))
    
    159
    -
    
    160
    -    # Inject work
    
    161
    -    action_digest = remote_execution_pb2.Digest(hash='gaff')
    
    162
    -    _inject_work(instance._instances[""]._scheduler, action_digest=action_digest)
    
    163
    -
    
    164
    -    request = bots_pb2.UpdateBotSessionRequest(name=response.name,
    
    165
    -                                               bot_session=response)
    
    166
    -    response = copy.deepcopy(instance.UpdateBotSession(request, context))
    
    167
    -
    
    168
    -    assert response.leases[0].state == LeaseState.PENDING.value
    
    169
    -    response.leases[0].state = LeaseState.ACTIVE.value
    
    170
    -
    
    171
    -    request = bots_pb2.UpdateBotSessionRequest(name=response.name,
    
    172
    -                                               bot_session=response)
    
    173
    -
    
    174
    -    response = copy.deepcopy(instance.UpdateBotSession(request, context))
    
    175
    -
    
    176
    -    response.leases[0].state = LeaseState.COMPLETED.value
    
    177
    -    response.leases[0].result.Pack(remote_execution_pb2.ActionResult())
    
    178
    -
    
    179
    -    request = bots_pb2.UpdateBotSessionRequest(name=response.name,
    
    180
    -                                               bot_session=response)
    
    181
    -    response = copy.deepcopy(instance.UpdateBotSession(request, context))
    
    182
    -
    
    183
    -    assert len(response.leases) is 0
    
    153
    +    pass
    
    184 154
     
    
    185 155
     
    
    186 156
     def test_work_rejected_by_bot(bot_session, context, instance):
    
    187
    -    request = bots_pb2.CreateBotSessionRequest(parent='',
    
    188
    -                                               bot_session=bot_session)
    
    189
    -    # Inject work
    
    190
    -    action_digest = remote_execution_pb2.Digest(hash='gaff')
    
    191
    -    _inject_work(instance._instances[""]._scheduler, action_digest=action_digest)
    
    192
    -
    
    193
    -    # Simulated the severed binding between client and server
    
    194
    -    response = copy.deepcopy(instance.CreateBotSession(request, context))
    
    195
    -
    
    196
    -    # Reject work
    
    197
    -    assert response.leases[0].state == LeaseState.PENDING.value
    
    198
    -    response.leases[0].state = LeaseState.COMPLETED.value
    
    199
    -    request = bots_pb2.UpdateBotSessionRequest(name=response.name,
    
    200
    -                                               bot_session=response)
    
    201
    -
    
    202
    -    response = instance.UpdateBotSession(request, context)
    
    203
    -
    
    204
    -    context.set_code.assert_called_once_with(grpc.StatusCode.UNIMPLEMENTED)
    
    205
    -
    
    206
    -
    
    207
    -@pytest.mark.parametrize("state", [LeaseState.LEASE_STATE_UNSPECIFIED, LeaseState.PENDING])
    
    208
    -def test_work_out_of_sync_from_pending(state, bot_session, context, instance):
    
    209
    -    request = bots_pb2.CreateBotSessionRequest(parent='',
    
    210
    -                                               bot_session=bot_session)
    
    211
    -    # Inject work
    
    212
    -    action_digest = remote_execution_pb2.Digest(hash='gaff')
    
    213
    -    _inject_work(instance._instances[""]._scheduler, action_digest=action_digest)
    
    214
    -
    
    215
    -    # Simulated the severed binding between client and server
    
    216
    -    response = copy.deepcopy(instance.CreateBotSession(request, context))
    
    217
    -
    
    218
    -    response.leases[0].state = state.value
    
    219
    -
    
    220
    -    request = bots_pb2.UpdateBotSessionRequest(name=response.name,
    
    221
    -                                               bot_session=response)
    
    222
    -
    
    223
    -    response = instance.UpdateBotSession(request, context)
    
    224
    -
    
    225
    -    context.set_code.assert_called_once_with(grpc.StatusCode.DATA_LOSS)
    
    226
    -
    
    227
    -
    
    228
    -@pytest.mark.parametrize("state", [LeaseState.LEASE_STATE_UNSPECIFIED, LeaseState.PENDING])
    
    229
    -def test_work_out_of_sync_from_active(state, bot_session, context, instance):
    
    230
    -    request = bots_pb2.CreateBotSessionRequest(parent='',
    
    231
    -                                               bot_session=bot_session)
    
    232
    -    # Inject work
    
    233
    -    action_digest = remote_execution_pb2.Digest(hash='gaff')
    
    234
    -    _inject_work(instance._instances[""]._scheduler, action_digest=action_digest)
    
    235
    -
    
    236
    -    # Simulated the severed binding between client and server
    
    237
    -    response = copy.deepcopy(instance.CreateBotSession(request, context))
    
    238
    -
    
    239
    -    response.leases[0].state = LeaseState.ACTIVE.value
    
    240
    -
    
    241
    -    request = copy.deepcopy(bots_pb2.UpdateBotSessionRequest(name=response.name,
    
    242
    -                                                             bot_session=response))
    
    243
    -
    
    244
    -    response = instance.UpdateBotSession(request, context)
    
    245
    -
    
    246
    -    response.leases[0].state = state.value
    
    247
    -
    
    248
    -    request = bots_pb2.UpdateBotSessionRequest(name=response.name,
    
    249
    -                                               bot_session=response)
    
    250
    -
    
    251
    -    response = instance.UpdateBotSession(request, context)
    
    252
    -
    
    253
    -    context.set_code.assert_called_once_with(grpc.StatusCode.DATA_LOSS)
    
    254
    -
    
    255
    -
    
    256
    -def test_work_active_to_active(bot_session, context, instance):
    
    257
    -    request = bots_pb2.CreateBotSessionRequest(parent='',
    
    258
    -                                               bot_session=bot_session)
    
    259
    -    # Inject work
    
    260
    -    action_digest = remote_execution_pb2.Digest(hash='gaff')
    
    261
    -    _inject_work(instance._instances[""]._scheduler, action_digest=action_digest)
    
    262
    -
    
    263
    -    # Simulated the severed binding between client and server
    
    264
    -    response = copy.deepcopy(instance.CreateBotSession(request, context))
    
    265
    -
    
    266
    -    response.leases[0].state = LeaseState.ACTIVE.value
    
    267
    -
    
    268
    -    request = bots_pb2.UpdateBotSessionRequest(name=response.name,
    
    269
    -                                               bot_session=response)
    
    270
    -
    
    271
    -    response = instance.UpdateBotSession(request, context)
    
    272
    -
    
    273
    -    assert response.leases[0].state == LeaseState.ACTIVE.value
    
    157
    +    pass
    
    274 158
     
    
    275 159
     
    
    276 160
     def test_post_bot_event_temp(context, instance):
    

  • tests/integration/operations_service.py
    ... ... @@ -28,10 +28,8 @@ from buildgrid._enums import OperationStage
    28 28
     from buildgrid._exceptions import InvalidArgumentError
    
    29 29
     from buildgrid._protos.build.bazel.remote.execution.v2 import remote_execution_pb2
    
    30 30
     from buildgrid._protos.google.longrunning import operations_pb2
    
    31
    -from buildgrid._protos.google.rpc import status_pb2
    
    32 31
     from buildgrid.server.cas.storage import lru_memory_cache
    
    33 32
     from buildgrid.server.controller import ExecutionController
    
    34
    -from buildgrid.server.job import LeaseState
    
    35 33
     from buildgrid.server.operations import service
    
    36 34
     from buildgrid.server.operations.service import OperationsService
    
    37 35
     from buildgrid.utils import create_digest
    
    ... ... @@ -166,31 +164,6 @@ def test_list_operations_instance_fail(instance, controller, execute_request, co
    166 164
         context.set_code.assert_called_once_with(grpc.StatusCode.INVALID_ARGUMENT)
    
    167 165
     
    
    168 166
     
    
    169
    -def test_list_operations_with_result(instance, controller, execute_request, context):
    
    170
    -    response_execute = controller.execution_instance.execute(execute_request.action_digest,
    
    171
    -                                                             execute_request.skip_cache_lookup)
    
    172
    -
    
    173
    -    action_result = remote_execution_pb2.ActionResult()
    
    174
    -    output_file = remote_execution_pb2.OutputFile(path='unicorn')
    
    175
    -    action_result.output_files.extend([output_file])
    
    176
    -
    
    177
    -    controller.operations_instance._scheduler.jobs[response_execute.name].create_lease()
    
    178
    -    controller.operations_instance._scheduler.update_job_lease_state(response_execute.name,
    
    179
    -                                                                     LeaseState.COMPLETED,
    
    180
    -                                                                     lease_status=status_pb2.Status(),
    
    181
    -                                                                     lease_result=_pack_any(action_result))
    
    182
    -
    
    183
    -    request = operations_pb2.ListOperationsRequest(name=instance_name)
    
    184
    -    response = instance.ListOperations(request, context)
    
    185
    -
    
    186
    -    assert response.operations[0].name.split('/')[-1] == response_execute.name
    
    187
    -
    
    188
    -    execute_response = remote_execution_pb2.ExecuteResponse()
    
    189
    -    response.operations[0].response.Unpack(execute_response)
    
    190
    -
    
    191
    -    assert execute_response.result.output_files == action_result.output_files
    
    192
    -
    
    193
    -
    
    194 167
     def test_list_operations_empty(instance, context):
    
    195 168
         request = operations_pb2.ListOperationsRequest(name=instance_name)
    
    196 169
     
    

  • tests/utils/bots_interface.py
    1
    +# Copyright (C) 2018 Bloomberg LP
    
    2
    +#
    
    3
    +# Licensed under the Apache License, Version 2.0 (the "License");
    
    4
    +# you may not use this file except in compliance with the License.
    
    5
    +# You may obtain a copy of the License at
    
    6
    +#
    
    7
    +#  <http://www.apache.org/licenses/LICENSE-2.0>
    
    8
    +#
    
    9
    +# Unless required by applicable law or agreed to in writing, software
    
    10
    +# distributed under the License is distributed on an "AS IS" BASIS,
    
    11
    +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    
    12
    +# See the License for the specific language governing permissions and
    
    13
    +# limitations under the License.
    
    14
    +
    
    15
    +from concurrent import futures
    
    16
    +from contextlib import contextmanager
    
    17
    +import multiprocessing
    
    18
    +import os
    
    19
    +import signal
    
    20
    +import uuid
    
    21
    +
    
    22
    +import grpc
    
    23
    +import psutil
    
    24
    +import pytest_cov
    
    25
    +
    
    26
    +from buildgrid.server.scheduler import Scheduler
    
    27
    +from buildgrid.server.bots import instance, service
    
    28
    +
    
    29
    +
    
    30
    +@contextmanager
    
    31
    +def serve_bots_interface(instances):
    
    32
    +    server = Server(instances)
    
    33
    +    try:
    
    34
    +        yield server
    
    35
    +    finally:
    
    36
    +        server.quit()
    
    37
    +
    
    38
    +
    
    39
    +def kill_process_tree(pid):
    
    40
    +    proc = psutil.Process(pid)
    
    41
    +    children = proc.children(recursive=True)
    
    42
    +
    
    43
    +    def kill_proc(p):
    
    44
    +        try:
    
    45
    +            p.kill()
    
    46
    +        except psutil.AccessDenied:
    
    47
    +            # Ignore this error, it can happen with
    
    48
    +            # some setuid bwrap processes.
    
    49
    +            pass
    
    50
    +
    
    51
    +    # Bloody Murder
    
    52
    +    for child in children:
    
    53
    +        kill_proc(child)
    
    54
    +    kill_proc(proc)
    
    55
    +
    
    56
    +
    
    57
    +def run_in_subprocess(function, *arguments):
    
    58
    +    queue = multiprocessing.Queue()
    
    59
    +    # Use subprocess to avoid creation of gRPC threads in main process
    
    60
    +    # See https://github.com/grpc/grpc/blob/master/doc/fork_support.md
    
    61
    +    process = multiprocessing.Process(target=function,
    
    62
    +                                      args=(queue, *arguments))
    
    63
    +
    
    64
    +    try:
    
    65
    +        process.start()
    
    66
    +        result = queue.get()
    
    67
    +        process.join()
    
    68
    +
    
    69
    +    except KeyboardInterrupt:
    
    70
    +        kill_process_tree(process.pid)
    
    71
    +        raise
    
    72
    +
    
    73
    +    return result
    
    74
    +
    
    75
    +
    
    76
    +class Server:
    
    77
    +
    
    78
    +    def __init__(self, instances):
    
    79
    +        self.instances = instances
    
    80
    +
    
    81
    +        self.__queue = multiprocessing.Queue()
    
    82
    +        self.__bot_session_queue = multiprocessing.Queue()
    
    83
    +        self.__process = multiprocessing.Process(
    
    84
    +            target=Server.serve,
    
    85
    +            args=(self.__queue, self.instances, self.__bot_session_queue))
    
    86
    +        self.__process.start()
    
    87
    +
    
    88
    +        self.port = self.__queue.get()
    
    89
    +        self.remote = 'localhost:{}'.format(self.port)
    
    90
    +
    
    91
    +    @classmethod
    
    92
    +    def serve(cls, queue, instances, bot_session_queue):
    
    93
    +        pytest_cov.embed.cleanup_on_sigterm()
    
    94
    +
    
    95
    +        # Use max_workers default from Python 3.5+
    
    96
    +        max_workers = (os.cpu_count() or 1) * 5
    
    97
    +        server = grpc.server(futures.ThreadPoolExecutor(max_workers))
    
    98
    +        port = server.add_insecure_port('localhost:0')
    
    99
    +
    
    100
    +        bots_service = service.BotsService(server)
    
    101
    +        for name in instances:
    
    102
    +            bots_interface = BotsInterface(bot_session_queue)
    
    103
    +            bots_service.add_instance(name, bots_interface)
    
    104
    +
    
    105
    +        server.start()
    
    106
    +        queue.put(port)
    
    107
    +        signal.pause()
    
    108
    +
    
    109
    +    def get_bot_session(self):
    
    110
    +        return self.__bot_session_queue.get()
    
    111
    +
    
    112
    +    def compare_bot_sessions(self, bot_session):
    
    113
    +        assert bot_session == self.get_bot_session()
    
    114
    +
    
    115
    +    def quit(self):
    
    116
    +        if self.__process:
    
    117
    +            self.__process.terminate()
    
    118
    +            self.__process.join()
    
    119
    +
    
    120
    +
    
    121
    +class BotsInterface:
    
    122
    +
    
    123
    +    def __init__(self, bot_session_queue):
    
    124
    +        self.__bot_session_queue = bot_session_queue
    
    125
    +
    
    126
    +    def register_instance_with_server(self, instance_name, server):
    
    127
    +        server.add_bots_interface(self, instance_name)
    
    128
    +
    
    129
    +    def create_bot_session(self, parent, bot_session):
    
    130
    +        name = "{}/{}".format(parent, str(uuid.uuid4()))
    
    131
    +        bot_session.name = name
    
    132
    +        self.__bot_session_queue.put(bot_session.SerializeToString())
    
    133
    +        return bot_session



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