client/cas.py: Extend the API and implement fallback (fb71df38) · Commits · BuildGrid / buildgrid

buildgrid/client/cas.py

+148 −55

Original line number	Diff line number	Diff line
		@@ -17,9 +17,29 @@ from contextlib import contextmanager
		import uuid
		import os

		from buildgrid.settings import HASH
		import grpc

		from buildgrid._protos.build.bazel.remote.execution.v2 import remote_execution_pb2, remote_execution_pb2_grpc
		from buildgrid._protos.google.bytestream import bytestream_pb2, bytestream_pb2_grpc
		from buildgrid._protos.google.rpc import code_pb2
		from buildgrid.settings import HASH


		class CallCache:
		"""Per remote grpc.StatusCode.UNIMPLEMENTED call cache."""
		__calls = dict()

		@classmethod
		def mark_unimplemented(cls, channel, name):
		if channel not in cls.__calls:
		cls.__calls[channel] = set()
		cls.__calls[channel].add(name)

		@classmethod
		def unimplemented(cls, channel, name):
		if channel not in cls.__calls:
		return False
		return name in cls.__calls[channel]


		@contextmanager
		@@ -47,6 +67,7 @@ class Uploader:

		FILE_SIZE_THRESHOLD = 1 * 1024 * 1024
		MAX_REQUEST_SIZE = 2 * 1024 * 1024
		MAX_REQUEST_COUNT = 500

		def __init__(self, channel, instance=None, u_uid=None):
		"""Initializes a new :class:`Uploader` instance.
		@@ -68,8 +89,61 @@ class Uploader:
		self.__cas_stub = remote_execution_pb2_grpc.ContentAddressableStorageStub(self.channel)

		self.__requests = dict()
		self.__request_count = 0
		self.__request_size = 0

		# --- Public API ---

		def put_blob(self, blob, digest=None, queue=False):
		"""Stores a blob into the remote CAS server.

		If queuing is allowed (`queue=True`), the upload request may be
		defer. An explicit call to :method:`flush` can force the request to be
		send immediately (along with the rest of the queued batch).

		Args:
		blob (bytes): the blob's data.
		digest (:obj:`Digest`, optional): the blob's digest.
		queue (bool, optional): whether or not the upload request may be
		queued and submitted as part of a batch upload request. Defaults
		to False.

		Returns:
		:obj:`Digest`: the sent blob's digest.
		"""
		if not queue or len(blob) > Uploader.FILE_SIZE_THRESHOLD:
		blob_digest = self._send_blob(blob)
		else:
		blob_digest = self._queue_blob(blob)

		return blob_digest

		def put_message(self, message, digest=None, queue=False):
		"""Stores a message into the remote CAS server.

		If queuing is allowed (`queue=True`), the upload request may be
		defer. An explicit call to :method:`flush` can force the request to be
		send immediately (along with the rest of the queued batch).

		Args:
		message (:obj:`Message`): the message object.
		digest (:obj:`Digest`, optional): the message's digest.
		queue (bool, optional): whether or not the upload request may be
		queued and submitted as part of a batch upload request. Defaults
		to False.

		Returns:
		:obj:`Digest`: the sent message's digest.
		"""
		message_blob = message.SerializeToString()

		if not queue or len(message_blob) > Uploader.FILE_SIZE_THRESHOLD:
		message_digest = self._send_blob(message_blob)
		else:
		message_digest = self._queue_blob(message_blob)

		return message_digest

		def upload_file(self, file_path, queue=True):
		"""Stores a local file into the remote CAS storage.

		@@ -79,7 +153,7 @@ class Uploader:

		Args:
		file_path (str): absolute or relative path to a local file.
		queue (bool, optional): wheter or not the upload request may be
		queue (bool, optional): whether or not the upload request may be
		queued and submitted as part of a batch upload request. Defaults
		to True.

		@@ -96,11 +170,11 @@ class Uploader:
		file_bytes = bytes_steam.read()

		if not queue or len(file_bytes) > Uploader.FILE_SIZE_THRESHOLD:
		blob_digest = self._send_blob(file_bytes)
		file_digest = self._send_blob(file_bytes)
		else:
		blob_digest = self._queue_blob(file_bytes)
		file_digest = self._queue_blob(file_bytes)

		return blob_digest
		return file_digest

		def upload_directory(self, directory, queue=True):
		"""Stores a :obj:`Directory` into the remote CAS storage.
		@@ -126,23 +200,14 @@ class Uploader:
		else:
		return self._queue_blob(directory.SerializeToString())

		def put_message(self, message):
		"""Stores a message into the remote CAS storage.

		Message is send immediately, upload is never be deferred.

		Args:
		message (:obj:`Message`): a protobuf message object.

		Returns:
		:obj:`Digest`: The digest of the message.
		"""
		return self._send_blob(message.SerializeToString())

		def flush(self):
		"""Ensures any queued request gets sent."""
		if self.__requests:
		self._send_batch()
		self._send_blob_batch(self.__requests)

		self.__requests.clear()
		self.__request_count = 0
		self.__request_size = 0

		def close(self):
		"""Closes the underlying connection stubs.
		@@ -156,34 +221,16 @@ class Uploader:
		self.__bytestream_stub = None
		self.__cas_stub = None

		def _queue_blob(self, blob):
		"""Queues a memory block for later batch upload"""
		blob_digest = remote_execution_pb2.Digest()
		blob_digest.hash = HASH(blob).hexdigest()
		blob_digest.size_bytes = len(blob)

		if self.__request_size + len(blob) > Uploader.MAX_REQUEST_SIZE:
		self._send_batch()

		update_request = remote_execution_pb2.BatchUpdateBlobsRequest.Request()
		update_request.digest.CopyFrom(blob_digest)
		update_request.data = blob
		# --- Private API ---

		update_request_size = update_request.ByteSize()
		if self.__request_size + update_request_size > Uploader.MAX_REQUEST_SIZE:
		self._send_batch()

		self.__requests[update_request.digest.hash] = update_request
		self.__request_size += update_request_size

		return blob_digest

		def _send_blob(self, blob):
		def _send_blob(self, blob, digest=None):
		"""Sends a memory block using ByteStream.Write()"""
		blob_digest = remote_execution_pb2.Digest()
		if digest is not None:
		blob_digest.CopyFrom(digest)
		else:
		blob_digest.hash = HASH(blob).hexdigest()
		blob_digest.size_bytes = len(blob)

		if self.instance_name is not None:
		resource_name = '/'.join([self.instance_name, 'uploads', self.u_uid, 'blobs',
		blob_digest.hash, str(blob_digest.size_bytes)])
		@@ -218,18 +265,64 @@ class Uploader:

		return blob_digest

		def _send_batch(self):
		def _queue_blob(self, blob, digest=None):
		"""Queues a memory block for later batch upload"""
		blob_digest = remote_execution_pb2.Digest()
		if digest is not None:
		blob_digest.CopyFrom(digest)
		else:
		blob_digest.hash = HASH(blob).hexdigest()
		blob_digest.size_bytes = len(blob)

		if self.__request_size + blob_digest.size_bytes > Uploader.MAX_REQUEST_SIZE:
		self.flush()
		elif self.__request_count >= Uploader.MAX_REQUEST_COUNT:
		self.flush()

		self.__requests[blob_digest.hash] = (blob, blob_digest)
		self.__request_count += 1
		self.__request_size += blob_digest.size_bytes

		return blob_digest

		def _send_blob_batch(self, batch):
		"""Sends queued data using ContentAddressableStorage.BatchUpdateBlobs()"""
		batch_fetched = False
		written_digests = list()

		# First, try BatchUpdateBlobs(), if not already known not being implemented:
		if not CallCache.unimplemented(self.channel, 'BatchUpdateBlobs'):
		batch_request = remote_execution_pb2.BatchUpdateBlobsRequest()
		batch_request.requests.extend(self.__requests.values())
		if self.instance_name is not None:
		batch_request.instance_name = self.instance_name

		batch_response = self.__cas_stub.BatchUpdateBlobs(batch_request)
		for blob, digest in batch.values():
		request = batch_request.requests.add()
		request.digest.CopyFrom(digest)
		request.data = blob

		try:
		batch_response = self.__cas_stub.BatchUpdateBlobs(batch_request)
		for response in batch_response.responses:
		assert response.digest.hash in self.__requests
		assert response.status.code is 0
		assert response.digest.hash in batch

		self.__requests.clear()
		self.__request_size = 0
		written_digests.append(response.digest)
		if response.status.code != code_pb2.OK:
		response.digest.Clear()

		batch_fetched = True

		except grpc.RpcError as e:
		status_code = e.code()
		if status_code == grpc.StatusCode.UNIMPLEMENTED:
		CallCache.mark_unimplemented(self.channel, 'BatchUpdateBlobs')

		else:
		assert False

		# Fallback to Write() if no BatchUpdateBlobs():
		if not batch_fetched:
		for blob, digest in batch.values():
		written_digests.append(self._send_blob(blob, digest=digest))

		return written_digests