[Notes] [Git][BuildStream/buildstream][raoul/802-refactor-artifactcache]

Raoul Hidalgo Charman pushed to branch raoul/802-refactor-artifactcache at BuildStream / buildstream

Commits:

ac8277bc

by Raoul Hidalgo Charman at 2019-01-09T15:52:36Z

artifactcache: implement new push methods

Similar to the pull methods, this implements a yield_directory_digests methods
that iterates over blobs in the local CAS, with the upload_blob sending blobs
to a remote and batching them where appropriate.

Part of #802

3 changed files:

buildstream/_artifactcache.py
buildstream/_cas/cascache.py
buildstream/_cas/casremote.py

Changes:

buildstream/_artifactcache.py

@@ -607,16 +607,41 @@ class ArtifactCache():
          for remote in push_remotes:
              remote.init()
 +            skipped_remote = True
              display_key = element._get_brief_display_key()
              element.status("Pushing artifact {} -> {}".format(display_key, remote.spec.url))
 -            if self.cas.push(refs, remote):
 -                element.info("Pushed artifact {} -> {}".format(display_key, remote.spec.url))
 +            try:
 +                for ref in refs:
 +                    # Check whether ref is already on the server in which case
 +                    # there is no need to push the ref
 +                    root_digest = self.cas.resolve_ref(ref)
 +                    response = remote.get_reference(ref)
 +                    if (response is not None and
 +                            response.hash == root_digest.hash and
 +                            response.size_bytes == root_digest.size_bytes):
 +                        element.info("Remote ({}) already has {} cached".format(
 +                            remote.spec.url, element._get_brief_display_key()))
 +                        continue
++
 +                    # upload blobs
 +                    self._send_directory(root_digest, remote)
 +                    remote.update_reference(ref, root_digest)
++
 +                    skipped_remote = False
++
 +            except CASError as e:
 +                if str(e.reason) == "StatusCode.RESOURCE_EXHAUSTED":
 +                    element.warn("Failed to push element to {}: Resource exhuasted"
 +                                 .format(remote.spec.url))
 +                    continue
 +                else:
 +                    raise ArtifactError("Failed to push refs {}: {}".format(refs, e),
 +                                        temporary=True) from e
++
 +            if skipped_remote is False:
                  pushed = True
 -            else:
 -                element.info("Remote ({}) already has {} cached".format(
 -                    remote.spec.url, element._get_brief_display_key()
 -                ))
 +                element.info("Pushed artifact {} -> {}".format(display_key, remote.spec.url))
          return pushed
@@ -731,7 +756,7 @@ class ArtifactCache():
              return
          for remote in push_remotes:
 -            self.cas.push_directory(remote, directory)
 +            self._send_directory(directory.ref, remote)
      # push_message():
+     #
@@ -816,6 +841,14 @@ class ArtifactCache():
          with self.context.timed_activity("Initializing remote caches", silent_nested=True):
              self.initialize_remotes(on_failure=remote_failed)
 +    def _send_directory(self, root_digest, remote):
 +        required_blobs = self.cas.yield_directory_digests(root_digest)
 +        missing_blobs = remote.find_missing_blobs(required_blobs)
 +        for blob in missing_blobs.values():
 +            blob_file = self.cas.objpath(blob)
 +            remote.upload_blob(blob, blob_file)
 +        remote.send_update_batch()
++
      # _write_cache_size()
+     #
      # Writes the given size of the artifact to the cache's size file

buildstream/_cas/cascache.py

@@ -18,23 +18,16 @@
  #        Jürg Billeter <juerg billeter codethink co uk>
  import hashlib
 -import itertools
  import os
  import stat
  import tempfile
 -import uuid
  import contextlib
 -import grpc
+-
  from .._protos.build.bazel.remote.execution.v2 import remote_execution_pb2
 -from .._protos.buildstream.v2 import buildstream_pb2
  from .. import utils
  from .._exceptions import CASError
 -from .casremote import _CASBatchUpdate
+-
  # A CASCache manages a CAS repository as specified in the Remote Execution API.
+ #
@@ -196,73 +189,6 @@ class CASCache():
          self.set_ref(newref, tree)
 -    # push():
 -    #
 -    # Push committed refs to remote repository.
 -    #
 -    # Args:
 -    #     refs (list): The refs to push
 -    #     remote (CASRemote): The remote to push to
 -    #
 -    # Returns:
 -    #   (bool): True if any remote was updated, False if no pushes were required
 -    #
 -    # Raises:
 -    #   (CASError): if there was an error
 -    #
 -    def push(self, refs, remote):
 -        skipped_remote = True
 -        try:
 -            for ref in refs:
 -                tree = self.resolve_ref(ref)
+-
 -                # Check whether ref is already on the server in which case
 -                # there is no need to push the ref
 -                try:
 -                    request = buildstream_pb2.GetReferenceRequest(instance_name=remote.spec.instance_name)
 -                    request.key = ref
 -                    response = remote.ref_storage.GetReference(request)
+-
 -                    if response.digest.hash == tree.hash and response.digest.size_bytes == tree.size_bytes:
 -                        # ref is already on the server with the same tree
 -                        continue
+-
 -                except grpc.RpcError as e:
 -                    if e.code() != grpc.StatusCode.NOT_FOUND:
 -                        # Intentionally re-raise RpcError for outer except block.
 -                        raise
+-
 -                self._send_directory(remote, tree)
+-
 -                request = buildstream_pb2.UpdateReferenceRequest(instance_name=remote.spec.instance_name)
 -                request.keys.append(ref)
 -                request.digest.hash = tree.hash
 -                request.digest.size_bytes = tree.size_bytes
 -                remote.ref_storage.UpdateReference(request)
+-
 -                skipped_remote = False
 -        except grpc.RpcError as e:
 -            if e.code() != grpc.StatusCode.RESOURCE_EXHAUSTED:
 -                raise CASError("Failed to push ref {}: {}".format(refs, e), temporary=True) from e
+-
 -        return not skipped_remote
+-
 -    # push_directory():
 -    #
 -    # Push the given virtual directory to a remote.
 -    #
 -    # Args:
 -    #     remote (CASRemote): The remote to push to
 -    #     directory (Directory): A virtual directory object to push.
 -    #
 -    # Raises:
 -    #     (CASError): if there was an error
 -    #
 -    def push_directory(self, remote, directory):
 -        remote.init()
+-
 -        self._send_directory(remote, directory.ref)
+-
      # objpath():
+     #
      # Return the path of an object based on its digest.
@@ -534,6 +460,27 @@ class CASCache():
          else:
              return None
 +    def yield_directory_digests(self, directory_digest):
 +        # parse directory, and recursively add blobs
 +        d = remote_execution_pb2.Digest()
 +        d.hash = directory_digest.hash
 +        d.size_bytes = directory_digest.size_bytes
 +        yield d
++
 +        directory = remote_execution_pb2.Directory()
++
 +        with open(self.objpath(directory_digest), 'rb') as f:
 +            directory.ParseFromString(f.read())
++
 +        for filenode in directory.files:
 +            d = remote_execution_pb2.Digest()
 +            d.hash = filenode.digest.hash
 +            d.size_bytes = filenode.digest.size_bytes
 +            yield d
++
 +        for dirnode in directory.directories:
 +            yield from self.yield_directory_digests(dirnode.digest)
++
      ################################################
      #             Local Private Methods            #
      ################################################
@@ -722,57 +669,3 @@ class CASCache():
          for dirnode in directory.directories:
              yield from self._required_blobs(dirnode.digest)
+-
 -    def _send_directory(self, remote, digest, u_uid=uuid.uuid4()):
 -        required_blobs = self._required_blobs(digest)
+-
 -        missing_blobs = dict()
 -        # Limit size of FindMissingBlobs request
 -        for required_blobs_group in _grouper(required_blobs, 512):
 -            request = remote_execution_pb2.FindMissingBlobsRequest(instance_name=remote.spec.instance_name)
+-
 -            for required_digest in required_blobs_group:
 -                d = request.blob_digests.add()
 -                d.hash = required_digest.hash
 -                d.size_bytes = required_digest.size_bytes
+-
 -            response = remote.cas.FindMissingBlobs(request)
 -            for missing_digest in response.missing_blob_digests:
 -                d = remote_execution_pb2.Digest()
 -                d.hash = missing_digest.hash
 -                d.size_bytes = missing_digest.size_bytes
 -                missing_blobs[d.hash] = d
+-
 -        # Upload any blobs missing on the server
 -        self._send_blobs(remote, missing_blobs.values(), u_uid)
+-
 -    def _send_blobs(self, remote, digests, u_uid=uuid.uuid4()):
 -        batch = _CASBatchUpdate(remote)
+-
 -        for digest in digests:
 -            with open(self.objpath(digest), 'rb') as f:
 -                assert os.fstat(f.fileno()).st_size == digest.size_bytes
+-
 -                if (digest.size_bytes >= remote.max_batch_total_size_bytes or
 -                        not remote.batch_update_supported):
 -                    # Too large for batch request, upload in independent request.
 -                    remote._send_blob(digest, f, u_uid=u_uid)
 -                else:
 -                    if not batch.add(digest, f):
 -                        # Not enough space left in batch request.
 -                        # Complete pending batch first.
 -                        batch.send()
 -                        batch = _CASBatchUpdate(remote)
 -                        batch.add(digest, f)
+-
 -        # Send final batch
 -        batch.send()
+-
+-
 -def _grouper(iterable, n):
 -    while True:
 -        try:
 -            current = next(iterable)
 -        except StopIteration:
 -            return
 -        yield itertools.chain([current], itertools.islice(iterable, n - 1))

buildstream/_cas/casremote.py

  from collections import namedtuple
  import io
 +import itertools
  import os
  import multiprocessing
  import signal
@@ -288,6 +289,18 @@ class CASRemote():
              else:
                  return None
 +    # update_reference():
 +    #
 +    # Args:
 +    #    ref (str): Reference to update
 +    #    digest (Digest): New digest to update ref with
 +    def update_reference(self, ref, digest):
 +        request = buildstream_pb2.UpdateReferenceRequest()
 +        request.keys.append(ref)
 +        request.digest.hash = digest.hash
 +        request.digest.size_bytes = digest.size_bytes
 +        self.ref_storage.UpdateReference(request)
++
      def get_tree_blob(self, tree_digest):
          self.init()
          f = tempfile.NamedTemporaryFile(dir=self.tmpdir)
@@ -405,6 +418,68 @@ class CASRemote():
          while self.__tmp_downloads:
              yield self.__tmp_downloads.pop()
 +    # upload_blob():
 +    #
 +    # Push blobs given an iterator over blob files
 +    #
 +    # Args:
 +    #    digest (Digest): digest we want to upload
 +    #    blob_file (str): Name of file location
 +    #    u_uid (str): Used to identify to the bytestream service
 +    #
 +    def upload_blob(self, digest, blob_file, u_uid=uuid.uuid4()):
 +        with open(blob_file, 'rb') as f:
 +            assert os.fstat(f.fileno()).st_size == digest.size_bytes
++
 +            if (digest.size_bytes >= self.max_batch_total_size_bytes or
 +                    not self.batch_update_supported):
 +                # Too large for batch request, upload in independent request.
 +                self._send_blob(digest, f, u_uid=u_uid)
 +            else:
 +                if self.__batch_update.add(digest, f) is False:
 +                    self.__batch_update.send()
 +                    self.__batch_update = _CASBatchUpdate(self)
 +                    self.__batch_update.add(digest, f)
++
 +    # send_update_batch():
 +    #
 +    # Sends anything left in the update batch
 +    #
 +    def send_update_batch(self):
 +        # make sure everything is sent
 +        self.__batch_update.send()
 +        self.__batch_update = _CASBatchUpdate(self)
++
 +    # find_missing_blobs()
 +    #
 +    # Does FindMissingBlobs request to remote
 +    #
 +    # Args:
 +    #    required_blobs ([Digest]): list of blobs required
 +    #
 +    # Returns:
 +    #    (Dict(Digest)): missing blobs
 +    def find_missing_blobs(self, required_blobs):
 +        self.init()
 +        missing_blobs = dict()
 +        # Limit size of FindMissingBlobs request
 +        for required_blobs_group in _grouper(required_blobs, 512):
 +            request = remote_execution_pb2.FindMissingBlobsRequest()
++
 +            for required_digest in required_blobs_group:
 +                d = request.blob_digests.add()
 +                d.hash = required_digest.hash
 +                d.size_bytes = required_digest.size_bytes
++
 +            response = self.cas.FindMissingBlobs(request)
 +            for missing_digest in response.missing_blob_digests:
 +                d = remote_execution_pb2.Digest()
 +                d.hash = missing_digest.hash
 +                d.size_bytes = missing_digest.size_bytes
 +                missing_blobs[d.hash] = d
++
 +        return missing_blobs
++
      ################################################
      #             Local Private Methods            #
      ################################################
@@ -443,7 +518,10 @@ class CASRemote():
                  offset += chunk_size
                  finished = request.finish_write
 -        response = self.bytestream.Write(request_stream(resource_name, stream))
 +        try:
 +            response = self.bytestream.Write(request_stream(resource_name, stream))
 +        except grpc.RpcError as e:
 +            raise CASError("Failed to upload blob: {}".format(e), reason=e.code())
          assert response.committed_size == digest.size_bytes
@@ -457,6 +535,15 @@ class CASRemote():
          self.__batch_read = _CASBatchRead(self)
 +def _grouper(iterable, n):
 +    while True:
 +        try:
 +            current = next(iterable)
 +        except StopIteration:
 +            return
 +        yield itertools.chain([current], itertools.islice(iterable, n - 1))
++
++
  # Represents a batch of blobs queued for fetching.
+ #
  class _CASBatchRead():
@@ -488,7 +575,11 @@ class _CASBatchRead():
          if not self._request.digests:
              return
 -        batch_response = self._remote.cas.BatchReadBlobs(self._request)
 +        try:
 +            batch_response = self._remote.cas.BatchReadBlobs(self._request)
 +        except grpc.RpcError as e:
 +            raise CASError("Failed to read blob batch: {}".format(e),
 +                           reason=e.code()) from e
          for response in batch_response.responses:
              if response.status.code == code_pb2.NOT_FOUND:
@@ -536,7 +627,12 @@ class _CASBatchUpdate():
          if not self._request.requests:
              return
 -        batch_response = self._remote.cas.BatchUpdateBlobs(self._request)
 +        # Want to raise a CASError if
 +        try:
 +            batch_response = self._remote.cas.BatchUpdateBlobs(self._request)
 +        except grpc.RpcError as e:
 +            raise CASError("Failed to upload blob batch: {}".format(e),
 +                            reason=e.code()) from e
          for response in batch_response.responses:
              if response.status.code != code_pb2.OK:

[Notes] [Git][BuildStream/buildstream][raoul/802-refactor-artifactcache] artifactcache: implement new push methods

Raoul Hidalgo Charman pushed to branch raoul/802-refactor-artifactcache at BuildStream / buildstream

Commits:

3 changed files:

Changes: