def test_peer_metadata():
peer_uid = 'peer2'
tool_name = 'tool-name'
tool_version = '1.0'
class DummyServicerContextWithMetadata(DummyServicerContext):
def invocation_metadata(self):
request_metadata = remote_execution_pb2.RequestMetadata()
request_metadata.tool_details.tool_name = tool_name
request_metadata.tool_details.tool_version = tool_version
request_metadata.action_id = '920ea86d6a445df893d0a39815e8856254392ce40e5957f167af8f16485916fb'
request_metadata.tool_invocation_id = 'cec14b04-075e-4f47-9c24-c5e6ac7f9827'
request_metadata.correlated_invocations_id = '9f962383-25f6-43b3-886d-56d761ac524e'
from collections import namedtuple
metadata_entry = namedtuple('metadata_entry', ('key', 'value'))
return [
metadata_entry(REQUEST_METADATA_HEADER_NAME,
request_metadata.SerializeToString())
]
p = Peer.register_peer(peer_uid,
context=DummyServicerContextWithMetadata())
assert Peer.find_peer(peer_uid) is p
assert p.uid == peer_uid
assert p.tool_name == tool_name
assert p.tool_version == tool_version
def test_peer_deregistering_the_last_instance_deletes_peer():
peer_uid = 'peer5'
p = Peer.register_peer(uid=peer_uid, context=context)
assert Peer.find_peer(peer_uid) is p
Peer.deregister_peer(peer_uid)
assert Peer.find_peer(peer_uid) is None
def test_create_new_peer_object():
p = Peer(uid='peer1')
assert p.uid == 'peer1'
assert p.token is None
assert p.tool_name is None
assert p.tool_name is None
def test_peer_token():
peer_uid = 'peer3'
token = 'abcdef1234'
p = Peer.register_peer(uid=peer_uid, token=token, context=context)
assert p.uid == peer_uid
assert p.token == token
def _rpc_termination_callback(self, peer_uid, instance_name,
operation_name):
self.__logger.debug(
"RPC terminated for peer_uid=[%s], instance_name=[%s], operation_name=[%s]",
peer_uid, instance_name, operation_name)
instance = self._get_instance(instance_name)
instance.unregister_operation_peer(operation_name, peer_uid)
if self._is_instrumented:
if self.__peers[peer_uid] > 1:
self.__peers[peer_uid] -= 1
else:
self.__peers_by_instance[instance_name].remove(peer_uid)
del self.__peers[peer_uid]
Peer.deregister_peer(peer_uid)
def Execute(self, request, context):
"""Handles ExecuteRequest messages.
Args:
request (ExecuteRequest): The incoming RPC request.
context (grpc.ServicerContext): Context for the RPC call.
"""
self.__logger.debug("Execute request from [%s]", context.peer())
instance_name = request.instance_name
message_queue = queue.Queue()
peer_uid = context.peer()
Peer.register_peer(uid=peer_uid, context=context)
try:
instance = self._get_instance(instance_name)
job_name = instance.execute(request.action_digest,
request.skip_cache_lookup)
operation_name = instance.register_job_peer(
job_name, peer_uid, message_queue)
context.add_callback(
partial(self._rpc_termination_callback, peer_uid,
instance_name, operation_name))
if self._is_instrumented:
if peer_uid not in self.__peers:
self.__peers_by_instance[instance_name].add(peer_uid)
self.__peers[peer_uid] = 1
else:
self.__peers[peer_uid] += 1
operation_full_name = "{}/{}".format(instance_name, operation_name)
self.__logger.info("Operation [%s] created for job [%s]",
operation_full_name, job_name)
for operation in instance.stream_operation_updates_while_context_is_active(
message_queue, context):
operation.name = operation_full_name
yield operation
if not context.is_active():
self.__logger.info(
"Peer peer_uid=[%s] was holding up a thread for "
" `stream_operation_updates()` for instance_name=[%s],"
" operation_name=[%s], but the rpc context is not active anymore;"
" releasing thread.", peer_uid, instance_name,
operation_name)
except InvalidArgumentError as e:
self.__logger.error(e)
context.set_details(str(e))
context.set_code(grpc.StatusCode.INVALID_ARGUMENT)
yield operations_pb2.Operation()
except FailedPreconditionError as e:
self.__logger.error(e)
context.set_details(str(e))
context.set_code(grpc.StatusCode.FAILED_PRECONDITION)
yield operations_pb2.Operation()
except CancelledError as e:
self.__logger.info("Operation cancelled [%s]", operation_full_name)
context.set_details(str(e))
context.set_code(grpc.StatusCode.CANCELLED)
yield e.last_response
def test_register_new_peer_in_global_list():
p = Peer.register_peer(uid='peer1', context=context)
assert p.uid == 'peer1'
assert Peer.find_peer('peer1') is p
def test_empty_peer_list():
assert Peer.find_peer('peer0') is None
def test_peer_increment_count():
peer_uid = 'peer6'
Peer.register_peer(uid=peer_uid, context=context) # instance_count == 1
Peer.register_peer(peer_uid, context=context) # instance_count == 2
Peer.register_peer(peer_uid, context=context) # instance_count == 3
Peer.deregister_peer(peer_uid) # instance_count == 2
Peer.deregister_peer(peer_uid) # instance_count == 1
assert Peer.find_peer(peer_uid) is not None
Peer.deregister_peer(peer_uid) # instance_count == 0
assert Peer.find_peer(peer_uid) is None
def test_create_peer_that_exists():
peer_uid = 'peer4'
original_peer = Peer.register_peer(uid=peer_uid, context=context)
new_peer = Peer.register_peer(uid=peer_uid, context=context)
assert new_peer is original_peer