Coverage for /builds/BuildGrid/buildgrid/buildgrid/bot/session.py: 70.59%
85 statements
« prev ^ index » next coverage.py v6.4.1, created at 2022-06-22 21:04 +0000
« prev ^ index » next coverage.py v6.4.1, created at 2022-06-22 21:04 +0000
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.
16"""
17Bot Session
18===========
20Allows connections
21"""
22import asyncio
23import logging
24import os
25import platform
27import grpc
29from buildgrid._enums import BotStatus, LeaseState
30from buildgrid._exceptions import FailedPreconditionError
31from buildgrid._protos.google.devtools.remoteworkers.v1test2 import bots_pb2
32from buildgrid._protos.google.rpc import code_pb2
33from .tenantmanager import TenantManager
36class BotSession:
37 def __init__(self, instance_name, bots_interface, hardware_interface, work,
38 context=None):
39 """ Unique bot ID within the farm used to identify this bot
40 Needs to be human readable.
41 All prior sessions with bot_id of same ID are invalidated.
42 If a bot attempts to update an invalid session, it must be rejected and
43 may be put in quarantine.
44 """
45 self.__logger = logging.getLogger(__name__)
47 self._bots_interface = bots_interface
48 self._hardware_interface = hardware_interface
50 self._status = BotStatus.OK.value
51 self._tenant_manager = TenantManager()
53 self.__instance_name = instance_name
54 self.__bot_id = f'{instance_name}.{platform.node()}.{os.getpid()}'
55 self.__name = None
57 self._work = work
58 self._context = context
60 self.__connected = False
62 @property
63 def bot_id(self):
64 return self.__bot_id
66 @property
67 def connected(self):
68 return self.__connected
70 async def run(self):
71 interval = self._bots_interface.interval
72 executing_interval = self._bots_interface.executing_interval
73 try:
74 while True:
75 if not self.connected:
76 self.create_bot_session()
77 else:
78 self.update_bot_session()
80 if not self.connected:
81 await asyncio.sleep(interval)
82 else:
83 await self._tenant_manager.wait_on_tenants(executing_interval)
84 except asyncio.CancelledError:
85 pass
87 def create_bot_session(self):
88 self.__logger.debug("Creating bot session")
90 session = self._bots_interface.create_bot_session(self.__instance_name, self.get_pb2())
91 if session in list(grpc.StatusCode):
92 self.__connected = False
93 return
94 self.__connected = True
95 self.__name = session.name
97 self.__logger.info(f"Created bot session with name: [{self.__name}]")
99 for lease in session.leases:
100 self._register_lease(lease)
102 def update_bot_session(self):
103 self.__logger.debug(f"Updating bot session: [{self.__bot_id}]")
105 session = self._bots_interface.update_bot_session(self.get_pb2())
106 if session in list(grpc.StatusCode):
107 self.__connected = False
108 return
109 self.__connected = True
110 server_ids = []
112 for lease in session.leases:
113 server_ids.append(lease.id)
115 lease_state = LeaseState(lease.state)
116 if lease_state == LeaseState.PENDING:
117 self._register_lease(lease)
119 elif lease_state == LeaseState.CANCELLED:
120 self._tenant_manager.cancel_tenancy(lease.id)
122 closed_lease_ids = [x for x in self._tenant_manager.get_lease_ids() if x not in server_ids]
124 for lease_id in closed_lease_ids:
125 self._tenant_manager.cancel_tenancy(lease_id)
126 self._tenant_manager.remove_tenant(lease_id)
128 def get_pb2(self):
129 return bots_pb2.BotSession(worker=self._hardware_interface.get_worker_pb2(),
130 status=self._status,
131 leases=self._tenant_manager.get_leases(),
132 bot_id=self.__bot_id,
133 name=self.__name)
135 def _register_lease(self, lease):
136 lease_id = lease.id
137 try:
138 self._tenant_manager.create_tenancy(lease)
140 except KeyError as e:
141 self.__logger.error(e)
143 else:
144 try:
145 self._hardware_interface.configure_hardware(lease.requirements)
147 except FailedPreconditionError as e:
148 self.__logger.error(e)
149 self._tenant_manager.complete_lease(lease_id, status=code_pb2.FailedPreconditionError)
151 else:
152 self._tenant_manager.create_work(lease_id, self._work, self._context)