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

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

Commits:

60fcc9d6

by Raoul Hidalgo Charman at 2019-01-03T14:21:39Z

artifactcache: implemented pull_tree similar to pull

New yield_directory_digests function implemented in CASRemote, which given a
tree digest can be iterated over, to give blob digests.

Tests have been updated to ensure the correct tmpdir is set up in process
wrappers, else invalid cross link errors happen in the CI.

Part of #802

6eb87919

by Raoul Hidalgo Charman at 2019-01-03T14:27:42Z

artifact cache: implement new push methods

Part of #802

Changes:

buildstream/_artifactcache.py

@@ -21,6 +21,7 @@ import multiprocessing
  import os
  import string
  from collections.abc import Mapping
 +import grpc
  from .types import _KeyStrength
  from ._exceptions import ArtifactError, CASError, LoadError, LoadErrorReason
@@ -607,16 +608,37 @@ class ArtifactCache():
          for remote in push_remotes:
              remote.init()
 +            pushed_remote = False
              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)
 +                    pushed_remote = True
++
 +            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
++
 +            if pushed_remote is True:
                  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
@@ -694,15 +716,26 @@ class ArtifactCache():
+     #
      # Args:
      #     project (Project): The current project
 -    #     digest (Digest): The digest of the tree
 +    #     tree_digest (Digest): The digest of the tree
+     #
 -    def pull_tree(self, project, digest):
 +    def pull_tree(self, project, tree_digest):
          for remote in self._remotes[project]:
 -            digest = self.cas.pull_tree(remote, digest)
 +            try:
 +                for blob_digest in remote.yield_tree_digests(tree_digest):
 +                    if self.cas.check_blob(blob_digest):
 +                        continue
 +                    remote.request_blob(blob_digest)
 +                    for blob_file in remote.get_blobs():
 +                        self.cas.add_object(path=blob_file.name, link_directly=True)
 -            if digest:
 -                # no need to pull from additional remotes
 -                return digest
 +                # Get the last batch
 +                for blob_file in remote.get_blobs(request_batch=True):
 +                    self.cas.add_object(path=blob_file.name, link_directly=True)
++
 +            except BlobNotFound:
 +                continue
 +            else:
 +                return tree_digest
          return None
@@ -731,7 +764,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 +849,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, final=True)
 +        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.
+ #
@@ -183,29 +176,6 @@ class CASCache():
          return modified, removed, added
 -    # pull_tree():
 -    #
 -    # Pull a single Tree rather than a ref.
 -    # Does not update local refs.
 -    #
 -    # Args:
 -    #     remote (CASRemote): The remote to pull from
 -    #     digest (Digest): The digest of the tree
 -    #
 -    def pull_tree(self, remote, digest):
 -        try:
 -            remote.init()
+-
 -            digest = self._fetch_tree(remote, digest)
+-
 -            return digest
+-
 -        except grpc.RpcError as e:
 -            if e.code() != grpc.StatusCode.NOT_FOUND:
 -                raise
+-
 -        return None
+-
      # link_ref():
+     #
      # Add an alias for an existing ref.
@@ -219,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()
 -                    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()
 -                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.
@@ -557,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            #
      ################################################
@@ -745,105 +669,3 @@ class CASCache():
          for dirnode in directory.directories:
              yield from self._required_blobs(dirnode.digest)
+-
 -    # _ensure_blob():
 -    #
 -    # Fetch and add blob if it's not already local.
 -    #
 -    # Args:
 -    #     remote (Remote): The remote to use.
 -    #     digest (Digest): Digest object for the blob to fetch.
 -    #
 -    # Returns:
 -    #     (str): The path of the object
 -    #
 -    def _ensure_blob(self, remote, digest):
 -        objpath = self.objpath(digest)
 -        if os.path.exists(objpath):
 -            # already in local repository
 -            return objpath
+-
 -        with tempfile.NamedTemporaryFile(dir=self.tmpdir) as f:
 -            remote._fetch_blob(digest, f)
+-
 -            added_digest = self.add_object(path=f.name, link_directly=True)
 -            assert added_digest.hash == digest.hash
+-
 -        return objpath
+-
 -    def _fetch_tree(self, remote, digest):
 -        # download but do not store the Tree object
 -        with tempfile.NamedTemporaryFile(dir=self.tmpdir) as out:
 -            remote._fetch_blob(digest, out)
+-
 -            tree = remote_execution_pb2.Tree()
+-
 -            with open(out.name, 'rb') as f:
 -                tree.ParseFromString(f.read())
+-
 -            tree.children.extend([tree.root])
 -            for directory in tree.children:
 -                for filenode in directory.files:
 -                    self._ensure_blob(remote, filenode.digest)
+-
 -                # place directory blob only in final location when we've downloaded
 -                # all referenced blobs to avoid dangling references in the repository
 -                dirbuffer = directory.SerializeToString()
 -                dirdigest = self.add_object(buffer=dirbuffer)
 -                assert dirdigest.size_bytes == len(dirbuffer)
+-
 -        return dirdigest
+-
 -    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()
+-
 -            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
@@ -283,19 +284,41 @@ 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)
 +        self._fetch_blob(tree_digest, f)
++
 +        tree = remote_execution_pb2.Tree()
 +        with open(f.name, 'rb') as tmp:
 +            tree.ParseFromString(tmp.read())
++
 +        return tree
++
      # yield_directory_digests():
+     #
      # Iterate over blobs digests starting from a root digest
+     #
      # Args:
 -    #     root_digest (str): The root_digest to get a tree of
 +    #     root_digest (digest): The root_digest to get a tree of
      #     progress (callable): The progress callback, if any
      #     subdir (str): The optional specific subdir to pull
      #     excluded_subdirs (list): The optional list of subdirs to not pull
+     #
      # Returns:
 -    #   (iter): True if pull was successful, False if ref was not available
 -    #
 +    #     (iter digests): recursively iterates over digests contained in root directory
      def yield_directory_digests(self, root_digest, *, progress=None, subdir=None, excluded_subdirs=[]):
          self.init()
@@ -303,6 +326,37 @@ class CASRemote():
          # Fetch artifact, excluded_subdirs determined in pullqueue
          yield from self._yield_directory_digests(root_digest, excluded_subdirs=excluded_subdirs)
 +    # yield_tree_digests():
 +    #
 +    # Fetches a tree file from digests and then iterates over child digests
 +    #
 +    # Args:
 +    #     tree_digest (digest): tree digest
 +    #
 +    # Returns:
 +    #     (iter digests): iterates over digests in tree message
 +    def yield_tree_digests(self, tree_digest):
 +        self.init()
++
 +        # get tree file
 +        f = tempfile.NamedTemporaryFile(dir=self.tmpdir)
 +        self._fetch_blob(tree_digest, f)
 +        tree = remote_execution_pb2.Tree()
 +        tree.ParseFromString(f.read())
++
 +        tree.children.extend([tree.root])
 +        for directory in tree.children:
 +            for filenode in directory.files:
 +                yield filenode.digest
++
 +            # add the directory to downloaded tmp files to be added
 +            f2 = tempfile.NamedTemporaryFile(dir=self.tmpdir)
 +            f2.write(directory.SerializeToString())
 +            self.__tmp_downloads.append(f2)
++
 +        # Add the tree directory to downloads right at the end
 +        self.__tmp_downloads.append(f)
++
      # request_blob():
+     #
      # Request blob, triggering download depending via bytestream or cas
@@ -337,6 +391,60 @@ class CASRemote():
          while self.__tmp_downloads:
              yield self.__tmp_downloads.pop()
 +    # upload_blob():
 +    #
 +    # Push blobs given an iterator over blob files
 +    #
 +    def upload_blob(self, digest, blob_file, u_uid=uuid.uuid4(), final=False):
 +        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)
++
 +    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
 +    #    u_uid (str): uuid4
 +    #
 +    # Returns:
 +    #    (Dict(Digest)): missing blobs
 +    def find_missing_blobs(self, required_blobs, u_uid=uuid.uuid4()):
 +        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            #
      ################################################
@@ -418,6 +526,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():

tests/artifactcache/pull.py

@@ -110,7 +110,7 @@ def test_pull(cli, tmpdir, datafiles):
          # See https://github.com/grpc/grpc/blob/master/doc/fork_support.md for details
          process = multiprocessing.Process(target=_queue_wrapper,
                                            args=(_test_pull, queue, user_config_file, project_dir,
 -                                                artifact_dir, 'target.bst', element_key))
 +                                                artifact_dir, tmpdir, 'target.bst', element_key))
          try:
              # Keep SIGINT blocked in the child process
@@ -127,13 +127,14 @@ def test_pull(cli, tmpdir, datafiles):
          assert cas.contains(element, element_key)
 -def _test_pull(user_config_file, project_dir, artifact_dir,
 +def _test_pull(user_config_file, project_dir, artifact_dir, tmpdir,
                 element_name, element_key, queue):
      # Fake minimal context
      context = Context()
      context.load(config=user_config_file)
      context.artifactdir = artifact_dir
      context.set_message_handler(message_handler)
 +    context.tmpdir = tmpdir
      # Load the project manually
      project = Project(project_dir, context)
@@ -218,7 +219,7 @@ def test_pull_tree(cli, tmpdir, datafiles):
          # See https://github.com/grpc/grpc/blob/master/doc/fork_support.md for details
          process = multiprocessing.Process(target=_queue_wrapper,
                                            args=(_test_push_tree, queue, user_config_file, project_dir,
 -                                                artifact_dir, artifact_digest))
 +                                                artifact_dir, tmpdir, artifact_digest))
          try:
              # Keep SIGINT blocked in the child process
@@ -246,7 +247,7 @@ def test_pull_tree(cli, tmpdir, datafiles):
          # Use subprocess to avoid creation of gRPC threads in main BuildStream process
          process = multiprocessing.Process(target=_queue_wrapper,
                                            args=(_test_pull_tree, queue, user_config_file, project_dir,
 -                                                artifact_dir, tree_digest))
 +                                                artifact_dir, tmpdir, tree_digest))
          try:
              # Keep SIGINT blocked in the child process
@@ -268,12 +269,14 @@ def test_pull_tree(cli, tmpdir, datafiles):
          assert os.path.exists(cas.objpath(directory_digest))
 -def _test_push_tree(user_config_file, project_dir, artifact_dir, artifact_digest, queue):
 +def _test_push_tree(user_config_file, project_dir, artifact_dir, tmpdir,
 +                    artifact_digest, queue):
      # Fake minimal context
      context = Context()
      context.load(config=user_config_file)
      context.artifactdir = artifact_dir
      context.set_message_handler(message_handler)
 +    context.tmpdir
      # Load the project manually
      project = Project(project_dir, context)
@@ -304,12 +307,14 @@ def _test_push_tree(user_config_file, project_dir, artifact_dir, artifact_digest
          queue.put("No remote configured")
 -def _test_pull_tree(user_config_file, project_dir, artifact_dir, artifact_digest, queue):
 +def _test_pull_tree(user_config_file, project_dir, artifact_dir, tmpdir,
 +                    artifact_digest, queue):
      # Fake minimal context
      context = Context()
      context.load(config=user_config_file)
      context.artifactdir = artifact_dir
      context.set_message_handler(message_handler)
 +    context.tmpdir = tmpdir
      # Load the project manually
      project = Project(project_dir, context)

[Notes] [Git][BuildStream/buildstream][raoul/802-refactor-artifactcache] 2 commits: artifactcache: implemented pull_tree similar to pull

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

Commits:

4 changed files:

Changes: