[Notes] [Git][BuildStream/buildstream][juerg/cas-batch] _artifactcache/cascache.py: Use BatchReadBlobs

[Jürg Billeter <gitlab mg gitlab com>]

Title: GitLab

Jürg Billeter pushed to branch juerg/cas-batch at BuildStream / buildstream

Commits:

• 9ba2395b
  by Jürg Billeter at 2018-09-24T11:17:45Z

  _artifactcache/cascache.py: Use BatchReadBlobs
  
  This uses BatchReadBlobs instead of individual blob download to speed up
  artifact pulling, if the server supports it.
  
  Fixes #554.
  

1 changed file:

• buildstream/_artifactcache/cascache.py

Changes:

buildstream/_artifactcache/cascache.py

---

...	...	@@ -884,6 +884,55 @@ class CASCache(ArtifactCache):
884	884
885	885	return objpath
886	886
	887	+ def _batch_download_complete(self, batch):
	888	+ for digest, data in batch.send():
	889	+ with tempfile.NamedTemporaryFile(dir=self.tmpdir) as f:
	890	+ f.write(data)
	891	+ f.flush()
	892	+
	893	+ added_digest = self.add_object(path=f.name)
	894	+ assert added_digest.hash == digest.hash
	895	+
	896	+ # Helper function for _fetch_directory().
	897	+ def _fetch_directory_batch(self, remote, batch, fetch_queue, fetch_next_queue):
	898	+ self._batch_download_complete(batch)
	899	+
	900	+ # All previously scheduled directories are now locally available,
	901	+ # move them to the processing queue.
	902	+ fetch_queue.extend(fetch_next_queue)
	903	+ fetch_next_queue.clear()
	904	+ return _CASBatchRead(remote)
	905	+
	906	+ # Helper function for _fetch_directory().
	907	+ def _fetch_directory_node(self, remote, digest, batch, fetch_queue, fetch_next_queue, *, recursive=False):
	908	+ in_local_cache = os.path.exists(self.objpath(digest))
	909	+
	910	+ if in_local_cache:
	911	+ # Skip download, already in local cache.
	912	+ pass
	913	+ elif (digest.size_bytes >= remote.max_batch_total_size_bytes or
	914	+ not remote.batch_read_supported):
	915	+ # Too large for batch request, download in independent request.
	916	+ self._ensure_blob(remote, digest)
	917	+ in_local_cache = True
	918	+ else:
	919	+ if not batch.add(digest):
	920	+ # Not enough space left in batch request.
	921	+ # Complete pending batch first.
	922	+ batch = self._fetch_directory_batch(remote, batch, fetch_queue, fetch_next_queue)
	923	+ batch.add(digest)
	924	+
	925	+ if recursive:
	926	+ if in_local_cache:
	927	+ # Add directory to processing queue.
	928	+ fetch_queue.append(digest)
	929	+ else:
	930	+ # Directory will be available after completing pending batch.
	931	+ # Add directory to deferred processing queue.
	932	+ fetch_next_queue.append(digest)
	933	+
	934	+ return batch
	935	+
887	936	# _fetch_directory():
888	937	#
889	938	# Fetches remote directory and adds it to content addressable store.
...	...	@@ -897,23 +946,32 @@ class CASCache(ArtifactCache):
897	946	# dir_digest (Digest): Digest object for the directory to fetch.
898	947	#
899	948	def _fetch_directory(self, remote, dir_digest):
900		- objpath = self.objpath(dir_digest)
901		- if os.path.exists(objpath):
902		- # already in local cache
903		- return
	949	+ fetch_queue = [dir_digest]
	950	+ fetch_next_queue = []
	951	+ batch = _CASBatchRead(remote)
904	952
905		- objpath = self._ensure_blob(remote, dir_digest)
	953	+ while len(fetch_queue) + len(fetch_next_queue) > 0:
	954	+ if len(fetch_queue) == 0:
	955	+ batch = self._fetch_directory_batch(remote, batch, fetch_queue, fetch_next_queue)
906	956
907		- directory = remote_execution_pb2.Directory()
	957	+ dir_digest = fetch_queue.pop(0)
908	958
909		- with open(objpath, 'rb') as f:
910		- directory.ParseFromString(f.read())
	959	+ objpath = self._ensure_blob(remote, dir_digest)
911	960
912		- for filenode in directory.files:
913		- self._ensure_blob(remote, filenode.digest)
	961	+ directory = remote_execution_pb2.Directory()
	962	+ with open(objpath, 'rb') as f:
	963	+ directory.ParseFromString(f.read())
914	964
915		- for dirnode in directory.directories:
916		- self._fetch_directory(remote, dirnode.digest)
	965	+ for dirnode in directory.directories:
	966	+ batch = self._fetch_directory_node(remote, dirnode.digest, batch,
	967	+ fetch_queue, fetch_next_queue, recursive=True)
	968	+
	969	+ for filenode in directory.files:
	970	+ batch = self._fetch_directory_node(remote, filenode.digest, batch,
	971	+ fetch_queue, fetch_next_queue)
	972	+
	973	+ # Fetch final batch
	974	+ self._fetch_directory_batch(remote, batch, fetch_queue, fetch_next_queue)
917	975
918	976	def _fetch_tree(self, remote, digest):
919	977	# download but do not store the Tree object
...	...	@@ -1040,11 +1098,78 @@ class _CASRemote():
1040	1098
1041	1099	self.bytestream = bytestream_pb2_grpc.ByteStreamStub(self.channel)
1042	1100	self.cas = remote_execution_pb2_grpc.ContentAddressableStorageStub(self.channel)
	1101	+ self.capabilities = remote_execution_pb2_grpc.CapabilitiesStub(self.channel)
1043	1102	self.ref_storage = buildstream_pb2_grpc.ReferenceStorageStub(self.channel)
1044	1103
	1104	+ self.max_batch_total_size_bytes = _MAX_PAYLOAD_BYTES
	1105	+ try:
	1106	+ request = remote_execution_pb2.GetCapabilitiesRequest()
	1107	+ response = self.capabilities.GetCapabilities(request)
	1108	+ server_max_batch_total_size_bytes = response.cache_capabilities.max_batch_total_size_bytes
	1109	+ if 0 < server_max_batch_total_size_bytes < self.max_batch_total_size_bytes:
	1110	+ self.max_batch_total_size_bytes = server_max_batch_total_size_bytes
	1111	+ except grpc.RpcError as e:
	1112	+ # Simply use the defaults for servers that don't implement GetCapabilities()
	1113	+ if e.code() != grpc.StatusCode.UNIMPLEMENTED:
	1114	+ raise
	1115	+
	1116	+ # Check whether the server supports BatchReadBlobs()
	1117	+ self.batch_read_supported = False
	1118	+ try:
	1119	+ request = remote_execution_pb2.BatchReadBlobsRequest()
	1120	+ response = self.cas.BatchReadBlobs(request)
	1121	+ self.batch_read_supported = True
	1122	+ except grpc.RpcError as e:
	1123	+ if e.code() != grpc.StatusCode.UNIMPLEMENTED:
	1124	+ raise
	1125	+
1045	1126	self._initialized = True
1046	1127
1047	1128
	1129	+# Represents a batch of blobs queued for fetching.
	1130	+#
	1131	+class _CASBatchRead():
	1132	+ def __init__(self, remote):
	1133	+ self._remote = remote
	1134	+ self._max_total_size_bytes = remote.max_batch_total_size_bytes
	1135	+ self._request = remote_execution_pb2.BatchReadBlobsRequest()
	1136	+ self._size = 0
	1137	+ self._sent = False
	1138	+
	1139	+ def add(self, digest):
	1140	+ assert not self._sent
	1141	+
	1142	+ new_batch_size = self._size + digest.size_bytes
	1143	+ if new_batch_size > self._max_total_size_bytes:
	1144	+ # Not enough space left in current batch
	1145	+ return False
	1146	+
	1147	+ request_digest = self._request.digests.add()
	1148	+ request_digest.hash = digest.hash
	1149	+ request_digest.size_bytes = digest.size_bytes
	1150	+ self._size = new_batch_size
	1151	+ return True
	1152	+
	1153	+ def send(self):
	1154	+ assert not self._sent
	1155	+ self._sent = True
	1156	+
	1157	+ if len(self._request.digests) == 0:
	1158	+ return
	1159	+
	1160	+ batch_response = self._remote.cas.BatchReadBlobs(self._request)
	1161	+
	1162	+ for response in batch_response.responses:
	1163	+ if response.status.code != grpc.StatusCode.OK.value[0]:
	1164	+ raise ArtifactError("Failed to download blob {}: {}".format(
	1165	+ response.digest.hash, response.status.code))
	1166	+ if response.digest.size_bytes != len(response.data):
	1167	+ raise ArtifactError("Failed to download blob {}: expected {} bytes, received {} bytes".format(
	1168	+ response.digest.hash, response.digest.size_bytes, len(response.data)))
	1169	+
	1170	+ yield (response.digest, response.data)
	1171	+
	1172	+
1048	1173	def _grouper(iterable, n):
1049	1174	while True:
1050	1175	try:

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