GNOME.org
 

[Notes] [Git][BuildGrid/buildgrid][mablanch/77-cas-uploader] 10 commits: Added ReferenceCache parser.

Title: GitLab

Martin Blanchard pushed to branch mablanch/77-cas-uploader at BuildGrid / buildgrid

Commits:

16 changed files:

Changes:

  • buildgrid/_app/bots/buildbox.py
    ... ... @@ -104,7 +104,7 @@ def work_buildbox(context, lease):
    104 104 
                 output_tree = _cas_tree_maker(stub_bytestream, output_digest)
    105 105
    106 106 
                 with upload(context.cas_channel) as cas:
    107 
    -                output_tree_digest = cas.send_message(output_tree)
    107 
    +                output_tree_digest = cas.put_message(output_tree)
    108 108
    109 109 
                 output_directory = remote_execution_pb2.OutputDirectory()
    110 110 
                 output_directory.tree_digest.CopyFrom(output_tree_digest)

  • buildgrid/_app/commands/cmd_cas.py
    ... ... @@ -27,8 +27,9 @@ from urllib.parse import urlparse
    27 27 
     import click
    28 28 
     import grpc
    29 29
    30 
    -from buildgrid.utils import merkle_maker, create_digest
    30 
    +from buildgrid.client.cas import upload
    31 31 
     from buildgrid._protos.build.bazel.remote.execution.v2 import remote_execution_pb2, remote_execution_pb2_grpc
    32 
    +from buildgrid.utils import merkle_maker
    32 33
    33 34 
     from ..cli import pass_context
    34 35
    ... ... @@ -66,27 +67,31 @@ def cli(context, remote, instance_name, client_key, client_cert, server_cert):
    66 67
    67 68
    68 69 
     @cli.command('upload-files', short_help="Upload files to the CAS server.")
    69 
    -@click.argument('files', nargs=-1, type=click.File('rb'), required=True)
    70 
    +@click.argument('files', nargs=-1, type=click.Path(exists=True, dir_okay=False), required=True)
    70 71 
     @pass_context
    71 72 
     def upload_files(context, files):
    72 
    -    stub = remote_execution_pb2_grpc.ContentAddressableStorageStub(context.channel)
    73 
    +    sent_digests, file_map = list(), dict()
    74 
    +    with upload(context.channel, instance=context.instance_name) as cas:
    75 
    +        for file_path in files:
    76 
    +            context.logger.info("Queueing {}".format(file_path))
    73 77
    74 
    -    requests = []
    75 
    -    for file in files:
    76 
    -        chunk = file.read()
    77 
    -        requests.append(remote_execution_pb2.BatchUpdateBlobsRequest.Request(
    78 
    -            digest=create_digest(chunk), data=chunk))
    78 
    +            file_digest = cas.upload_file(file_path, queue=True)
    79 79
    80 
    -    request = remote_execution_pb2.BatchUpdateBlobsRequest(instance_name=context.instance_name,
    81 
    -                                                           requests=requests)
    80 
    +            assert file_digest.hash and file_digest.size_bytes
    82 81
    83 
    -    context.logger.info("Sending: {}".format(request))
    84 
    -    response = stub.BatchUpdateBlobs(request)
    85 
    -    context.logger.info("Response: {}".format(response))
    82 
    +            file_map[file_digest.hash] = file_path
    83 
    +            sent_digests.append(file_digest)
    84 
    +
    85 
    +    for file_digest in sent_digests:
    86 
    +        file_path = file_map[file_digest.hash]
    87 
    +        if file_digest.ByteSize():
    88 
    +            context.logger.info("{}: {}".format(file_path, file_digest.hash))
    89 
    +        else:
    90 
    +            context.logger.info("{}: FAILED".format(file_path))
    86 91
    87 92
    88 93 
     @cli.command('upload-dir', short_help="Upload a directory to the CAS server.")
    89 
    -@click.argument('directory', nargs=1, type=click.Path(), required=True)
    94 
    +@click.argument('directory', nargs=1, type=click.Path(exists=True, file_okay=False), required=True)
    90 95 
     @pass_context
    91 96 
     def upload_dir(context, directory):
    92 97 
         context.logger.info("Uploading directory to cas")

  • buildgrid/_app/commands/cmd_execute.py
    ... ... @@ -30,9 +30,10 @@ from urllib.parse import urlparse
    30 30 
     import click
    31 31 
     import grpc
    32 32
    33 
    -from buildgrid.utils import merkle_maker, create_digest, write_fetch_blob
    33 
    +from buildgrid.client.cas import upload
    34 34 
     from buildgrid._protos.build.bazel.remote.execution.v2 import remote_execution_pb2, remote_execution_pb2_grpc
    35 35 
     from buildgrid._protos.google.bytestream import bytestream_pb2_grpc
    36 
    +from buildgrid.utils import merkle_maker, write_fetch_blob
    36 37
    37 38 
     from ..cli import pass_context
    38 39
    ... ... @@ -119,46 +120,37 @@ def wait_execution(context, operation_name):
    119 120 
     @click.argument('input-root', nargs=1, type=click.Path(), required=True)
    120 121 
     @click.argument('commands', nargs=-1, type=click.STRING, required=True)
    121 122 
     @pass_context
    122 
    -def command(context, input_root, commands, output_file, output_directory):
    123 
    +def run_command(context, input_root, commands, output_file, output_directory):
    123 124 
         stub = remote_execution_pb2_grpc.ExecutionStub(context.channel)
    124 125
    125 
    -    execute_command = remote_execution_pb2.Command()
    126 
    +    output_executeables = list()
    127 
    +    with upload(context.channel, instance=context.instance_name) as cas:
    128 
    +        command = remote_execution_pb2.Command()
    126 129
    127 
    -    for arg in commands:
    128 
    -        execute_command.arguments.extend([arg])
    130 
    +        for arg in commands:
    131 
    +            command.arguments.extend([arg])
    129 132
    130 
    -    output_executeables = []
    131 
    -    for file, is_executeable in output_file:
    132 
    -        execute_command.output_files.extend([file])
    133 
    -        if is_executeable:
    134 
    -            output_executeables.append(file)
    133 
    +        for file, is_executeable in output_file:
    134 
    +            command.output_files.extend([file])
    135 
    +            if is_executeable:
    136 
    +                output_executeables.append(file)
    135 137
    136 
    -    command_digest = create_digest(execute_command.SerializeToString())
    137 
    -    context.logger.info(command_digest)
    138 
    +        command_digest = cas.put_message(command, queue=True)
    138 139
    139 
    -    # TODO: Check for missing blobs
    140 
    -    digest = None
    141 
    -    for _, digest in merkle_maker(input_root):
    142 
    -        pass
    140 
    +        context.logger.info('Sent command: {}'.format(command_digest))
    143 141
    144 
    -    action = remote_execution_pb2.Action(command_digest=command_digest,
    145 
    -                                         input_root_digest=digest,
    146 
    -                                         do_not_cache=True)
    142 
    +        # TODO: Check for missing blobs
    143 
    +        input_root_digest = None
    144 
    +        for _, input_root_digest in merkle_maker(input_root):
    145 
    +            pass
    147 146
    148 
    -    action_digest = create_digest(action.SerializeToString())
    147 
    +        action = remote_execution_pb2.Action(command_digest=command_digest,
    148 
    +                                             input_root_digest=input_root_digest,
    149 
    +                                             do_not_cache=True)
    149 150
    150 
    -    context.logger.info("Sending execution request...")
    151 
    -
    152 
    -    requests = []
    153 
    -    requests.append(remote_execution_pb2.BatchUpdateBlobsRequest.Request(
    154 
    -        digest=command_digest, data=execute_command.SerializeToString()))
    155 
    -
    156 
    -    requests.append(remote_execution_pb2.BatchUpdateBlobsRequest.Request(
    157 
    -        digest=action_digest, data=action.SerializeToString()))
    151 
    +        action_digest = cas.put_message(action, queue=True)
    158 152
    159 
    -    request = remote_execution_pb2.BatchUpdateBlobsRequest(instance_name=context.instance_name,
    160 
    -                                                           requests=requests)
    161 
    -    remote_execution_pb2_grpc.ContentAddressableStorageStub(context.channel).BatchUpdateBlobs(request)
    153 
    +        context.logger.info("Sent action: {}".format(action_digest))
    162 154
    163 155 
         request = remote_execution_pb2.ExecuteRequest(instance_name=context.instance_name,
    164 156 
                                                       action_digest=action_digest,

  • buildgrid/_app/settings/cas.yml
    ... ... @@ -7,7 +7,7 @@ server:
    7 7 
         tls-client-certs: null
    8 8
    9 9 
     description: |
    10 
    -  Just a CAS.
    10 
    +  Just a CAS with some reference storage.
    11 11
    12 12 
     instances:
    13 13 
       - name: main
    ... ... @@ -24,3 +24,8 @@ instances:
    24 24
    25 25 
           - !bytestream
    26 26 
             storage: *main-storage
    27 
    +
    28 
    +      - !reference-cache
    29 
    +        storage: *main-storage
    30 
    +        max_cached_refs: 256
    31 
    +        allow_updates: true

  • buildgrid/_app/settings/parser.py
    ... ... @@ -23,6 +23,7 @@ import yaml
    23 23
    24 24 
     from buildgrid.server.controller import ExecutionController
    25 25 
     from buildgrid.server.actioncache.storage import ActionCache
    26 
    +from buildgrid.server.referencestorage.storage import ReferenceCache
    26 27 
     from buildgrid.server.cas.instance import ByteStreamInstance, ContentAddressableStorageInstance
    27 28 
     from buildgrid.server.cas.storage.disk import DiskStorage
    28 29 
     from buildgrid.server.cas.storage.lru_memory_cache import LRUMemoryCache
    ... ... @@ -126,10 +127,18 @@ class Action(YamlFactory):
    126 127
    127 128 
         yaml_tag = u'!action-cache'
    128 129
    129 
    -    def __new__(cls, storage, max_cached_refs=0, allow_updates=True):
    130 
    +    def __new__(cls, storage, max_cached_refs, allow_updates=True):
    130 131 
             return ActionCache(storage, max_cached_refs, allow_updates)
    131 132
    132 133
    134 
    +class Reference(YamlFactory):
    135 
    +
    136 
    +    yaml_tag = u'!reference-cache'
    137 
    +
    138 
    +    def __new__(cls, storage, max_cached_refs, allow_updates=True):
    139 
    +        return ReferenceCache(storage, max_cached_refs, allow_updates)
    140 
    +
    141 
    +
    133 142 
     class CAS(YamlFactory):
    134 143
    135 144 
         yaml_tag = u'!cas'
    ... ... @@ -160,9 +169,9 @@ def _parse_size(size):
    160 169
    161 170 
     def get_parser():
    162 171
    163 
    -    yaml.SafeLoader.add_constructor(Execution.yaml_tag, Execution.from_yaml)
    164 172 
         yaml.SafeLoader.add_constructor(Execution.yaml_tag, Execution.from_yaml)
    165 173 
         yaml.SafeLoader.add_constructor(Action.yaml_tag, Action.from_yaml)
    174 
    +    yaml.SafeLoader.add_constructor(Reference.yaml_tag, Reference.from_yaml)
    166 175 
         yaml.SafeLoader.add_constructor(Disk.yaml_tag, Disk.from_yaml)
    167 176 
         yaml.SafeLoader.add_constructor(LRU.yaml_tag, LRU.from_yaml)
    168 177 
         yaml.SafeLoader.add_constructor(S3.yaml_tag, S3.from_yaml)

  • buildgrid/client/cas.py
    ... ... @@ -17,9 +17,29 @@ from contextlib import contextmanager
    17 17 
     import uuid
    18 18 
     import os
    19 19
    20 
    -from buildgrid.settings import HASH
    20 
    +import grpc
    21 
    +
    21 22 
     from buildgrid._protos.build.bazel.remote.execution.v2 import remote_execution_pb2, remote_execution_pb2_grpc
    22 23 
     from buildgrid._protos.google.bytestream import bytestream_pb2, bytestream_pb2_grpc
    24 
    +from buildgrid._protos.google.rpc import code_pb2
    25 
    +from buildgrid.settings import HASH
    26 
    +
    27 
    +
    28 
    +class CallCache:
    29 
    +    """Per remote grpc.StatusCode.UNIMPLEMENTED call cache."""
    30 
    +    __calls = dict()
    31 
    +
    32 
    +    @classmethod
    33 
    +    def mark_unimplemented(cls, channel, name):
    34 
    +        if channel not in cls.__calls:
    35 
    +            cls.__calls[channel] = set()
    36 
    +        cls.__calls[channel].add(name)
    37 
    +
    38 
    +    @classmethod
    39 
    +    def unimplemented(cls, channel, name):
    40 
    +        if channel not in cls.__calls:
    41 
    +            return False
    42 
    +        return name in cls.__calls[channel]
    23 43
    24 44
    25 45 
     @contextmanager
    ... ... @@ -28,7 +48,7 @@ def upload(channel, instance=None, u_uid=None):
    28 48 
         try:
    29 49 
             yield uploader
    30 50 
         finally:
    31 
    -        uploader.flush()
    51 
    +        uploader.close()
    32 52
    33 53
    34 54 
     class Uploader:
    ... ... @@ -47,6 +67,7 @@ class Uploader:
    47 67
    48 68 
         FILE_SIZE_THRESHOLD = 1 * 1024 * 1024
    49 69 
         MAX_REQUEST_SIZE = 2 * 1024 * 1024
    70 
    +    MAX_REQUEST_COUNT = 500
    50 71
    51 72 
         def __init__(self, channel, instance=None, u_uid=None):
    52 73 
             """Initializes a new :class:`Uploader` instance.
    ... ... @@ -68,8 +89,61 @@ class Uploader:
    68 89 
             self.__cas_stub = remote_execution_pb2_grpc.ContentAddressableStorageStub(self.channel)
    69 90
    70 91 
             self.__requests = dict()
    92 
    +        self.__request_count = 0
    71 93 
             self.__request_size = 0
    72 94
    95 
    +    # --- Public API ---
    96 
    +
    97 
    +    def put_blob(self, blob, digest=None, queue=False):
    98 
    +        """Stores a blob into the remote CAS server.
    99 
    +
    100 
    +        If queuing is allowed (`queue=True`), the upload request **may** be
    101 
    +        defer. An explicit call to :method:`flush` can force the request to be
    102 
    +        send immediately (along with the rest of the queued batch).
    103 
    +
    104 
    +        Args:
    105 
    +            blob (bytes): the blob's data.
    106 
    +            digest (:obj:`Digest`, optional): the blob's digest.
    107 
    +            queue (bool, optional): whether or not the upload request may be
    108 
    +                queued and submitted as part of a batch upload request. Defaults
    109 
    +                to False.
    110 
    +
    111 
    +        Returns:
    112 
    +            :obj:`Digest`: the sent blob's digest.
    113 
    +        """
    114 
    +        if not queue or len(blob) > Uploader.FILE_SIZE_THRESHOLD:
    115 
    +            blob_digest = self._send_blob(blob, digest=digest)
    116 
    +        else:
    117 
    +            blob_digest = self._queue_blob(blob, digest=digest)
    118 
    +
    119 
    +        return blob_digest
    120 
    +
    121 
    +    def put_message(self, message, digest=None, queue=False):
    122 
    +        """Stores a message into the remote CAS server.
    123 
    +
    124 
    +        If queuing is allowed (`queue=True`), the upload request **may** be
    125 
    +        defer. An explicit call to :method:`flush` can force the request to be
    126 
    +        send immediately (along with the rest of the queued batch).
    127 
    +
    128 
    +        Args:
    129 
    +            message (:obj:`Message`): the message object.
    130 
    +            digest (:obj:`Digest`, optional): the message's digest.
    131 
    +            queue (bool, optional): whether or not the upload request may be
    132 
    +                queued and submitted as part of a batch upload request. Defaults
    133 
    +                to False.
    134 
    +
    135 
    +        Returns:
    136 
    +            :obj:`Digest`: the sent message's digest.
    137 
    +        """
    138 
    +        message_blob = message.SerializeToString()
    139 
    +
    140 
    +        if not queue or len(message_blob) > Uploader.FILE_SIZE_THRESHOLD:
    141 
    +            message_digest = self._send_blob(message_blob, digest=digest)
    142 
    +        else:
    143 
    +            message_digest = self._queue_blob(message_blob, digest=digest)
    144 
    +
    145 
    +        return message_digest
    146 
    +
    73 147 
         def upload_file(self, file_path, queue=True):
    74 148 
             """Stores a local file into the remote CAS storage.
    75 149
    ... ... @@ -79,7 +153,7 @@ class Uploader:
    79 153
    80 154 
             Args:
    81 155 
                 file_path (str): absolute or relative path to a local file.
    82 
    -            queue (bool, optional): wheter or not the upload request may be
    156 
    +            queue (bool, optional): whether or not the upload request may be
    83 157 
                     queued and submitted as part of a batch upload request. Defaults
    84 158 
                     to True.
    85 159
    ... ... @@ -96,11 +170,11 @@ class Uploader:
    96 170 
                 file_bytes = bytes_steam.read()
    97 171
    98 172 
             if not queue or len(file_bytes) > Uploader.FILE_SIZE_THRESHOLD:
    99 
    -            blob_digest = self._send_blob(file_bytes)
    173 
    +            file_digest = self._send_blob(file_bytes)
    100 174 
             else:
    101 
    -            blob_digest = self._queue_blob(file_bytes)
    175 
    +            file_digest = self._queue_blob(file_bytes)
    102 176
    103 
    -        return blob_digest
    177 
    +        return file_digest
    104 178
    105 179 
         def upload_directory(self, directory, queue=True):
    106 180 
             """Stores a :obj:`Directory` into the remote CAS storage.
    ... ... @@ -126,50 +200,37 @@ class Uploader:
    126 200 
             else:
    127 201 
                 return self._queue_blob(directory.SerializeToString())
    128 202
    129 
    -    def send_message(self, message):
    130 
    -        """Stores a message into the remote CAS storage.
    131 
    -
    132 
    -        Args:
    133 
    -            message (:obj:`Message`): a protobuf message object.
    134 
    -
    135 
    -        Returns:
    136 
    -            :obj:`Digest`: The digest of the message.
    137 
    -        """
    138 
    -        return self._send_blob(message.SerializeToString())
    139 
    -
    140 203 
         def flush(self):
    141 204 
             """Ensures any queued request gets sent."""
    142 205 
             if self.__requests:
    143 
    -            self._send_batch()
    144 
    -
    145 
    -    def _queue_blob(self, blob):
    146 
    -        """Queues a memory block for later batch upload"""
    147 
    -        blob_digest = remote_execution_pb2.Digest()
    148 
    -        blob_digest.hash = HASH(blob).hexdigest()
    149 
    -        blob_digest.size_bytes = len(blob)
    206 
    +            self._send_blob_batch(self.__requests)
    150 207
    151 
    -        if self.__request_size + len(blob) > Uploader.MAX_REQUEST_SIZE:
    152 
    -            self._send_batch()
    208 
    +            self.__requests.clear()
    209 
    +            self.__request_count = 0
    210 
    +            self.__request_size = 0
    153 211
    154 
    -        update_request = remote_execution_pb2.BatchUpdateBlobsRequest.Request()
    155 
    -        update_request.digest.CopyFrom(blob_digest)
    156 
    -        update_request.data = blob
    212 
    +    def close(self):
    213 
    +        """Closes the underlying connection stubs.
    157 214
    158 
    -        update_request_size = update_request.ByteSize()
    159 
    -        if self.__request_size + update_request_size > Uploader.MAX_REQUEST_SIZE:
    160 
    -            self._send_batch()
    215 
    +        Note:
    216 
    +            This will always send pending requests before closing connections,
    217 
    +            if any.
    218 
    +        """
    219 
    +        self.flush()
    161 220
    162 
    -        self.__requests[update_request.digest.hash] = update_request
    163 
    -        self.__request_size += update_request_size
    221 
    +        self.__bytestream_stub = None
    222 
    +        self.__cas_stub = None
    164 223
    165 
    -        return blob_digest
    224 
    +    # --- Private API ---
    166 225
    167 
    -    def _send_blob(self, blob):
    226 
    +    def _send_blob(self, blob, digest=None):
    168 227 
             """Sends a memory block using ByteStream.Write()"""
    169 228 
             blob_digest = remote_execution_pb2.Digest()
    170 
    -        blob_digest.hash = HASH(blob).hexdigest()
    171 
    -        blob_digest.size_bytes = len(blob)
    172 
    -
    229 
    +        if digest is not None:
    230 
    +            blob_digest.CopyFrom(digest)
    231 
    +        else:
    232 
    +            blob_digest.hash = HASH(blob).hexdigest()
    233 
    +            blob_digest.size_bytes = len(blob)
    173 234 
             if self.instance_name is not None:
    174 235 
                 resource_name = '/'.join([self.instance_name, 'uploads', self.u_uid, 'blobs',
    175 236 
                                           blob_digest.hash, str(blob_digest.size_bytes)])
    ... ... @@ -204,18 +265,64 @@ class Uploader:
    204 265
    205 266 
             return blob_digest
    206 267
    207 
    -    def _send_batch(self):
    268 
    +    def _queue_blob(self, blob, digest=None):
    269 
    +        """Queues a memory block for later batch upload"""
    270 
    +        blob_digest = remote_execution_pb2.Digest()
    271 
    +        if digest is not None:
    272 
    +            blob_digest.CopyFrom(digest)
    273 
    +        else:
    274 
    +            blob_digest.hash = HASH(blob).hexdigest()
    275 
    +            blob_digest.size_bytes = len(blob)
    276 
    +
    277 
    +        if self.__request_size + blob_digest.size_bytes > Uploader.MAX_REQUEST_SIZE:
    278 
    +            self.flush()
    279 
    +        elif self.__request_count >= Uploader.MAX_REQUEST_COUNT:
    280 
    +            self.flush()
    281 
    +
    282 
    +        self.__requests[blob_digest.hash] = (blob, blob_digest)
    283 
    +        self.__request_count += 1
    284 
    +        self.__request_size += blob_digest.size_bytes
    285 
    +
    286 
    +        return blob_digest
    287 
    +
    288 
    +    def _send_blob_batch(self, batch):
    208 289 
             """Sends queued data using ContentAddressableStorage.BatchUpdateBlobs()"""
    209 
    -        batch_request = remote_execution_pb2.BatchUpdateBlobsRequest()
    210 
    -        batch_request.requests.extend(self.__requests.values())
    211 
    -        if self.instance_name is not None:
    212 
    -            batch_request.instance_name = self.instance_name
    290 
    +        batch_fetched = False
    291 
    +        written_digests = list()
    213 292
    214 
    -        batch_response = self.__cas_stub.BatchUpdateBlobs(batch_request)
    293 
    +        # First, try BatchUpdateBlobs(), if not already known not being implemented:
    294 
    +        if not CallCache.unimplemented(self.channel, 'BatchUpdateBlobs'):
    295 
    +            batch_request = remote_execution_pb2.BatchUpdateBlobsRequest()
    296 
    +            if self.instance_name is not None:
    297 
    +                batch_request.instance_name = self.instance_name
    215 298
    216 
    -        for response in batch_response.responses:
    217 
    -            assert response.digest.hash in self.__requests
    218 
    -            assert response.status.code is 0
    299 
    +            for blob, digest in batch.values():
    300 
    +                request = batch_request.requests.add()
    301 
    +                request.digest.CopyFrom(digest)
    302 
    +                request.data = blob
    219 303
    220 
    -        self.__requests.clear()
    221 
    -        self.__request_size = 0
    304 
    +            try:
    305 
    +                batch_response = self.__cas_stub.BatchUpdateBlobs(batch_request)
    306 
    +                for response in batch_response.responses:
    307 
    +                    assert response.digest.hash in batch
    308 
    +
    309 
    +                    written_digests.append(response.digest)
    310 
    +                    if response.status.code != code_pb2.OK:
    311 
    +                        response.digest.Clear()
    312 
    +
    313 
    +                batch_fetched = True
    314 
    +
    315 
    +            except grpc.RpcError as e:
    316 
    +                status_code = e.code()
    317 
    +                if status_code == grpc.StatusCode.UNIMPLEMENTED:
    318 
    +                    CallCache.mark_unimplemented(self.channel, 'BatchUpdateBlobs')
    319 
    +
    320 
    +                else:
    321 
    +                    assert False
    322 
    +
    323 
    +        # Fallback to Write() if no BatchUpdateBlobs():
    324 
    +        if not batch_fetched:
    325 
    +            for blob, digest in batch.values():
    326 
    +                written_digests.append(self._send_blob(blob, digest=digest))
    327 
    +
    328 
    +        return written_digests

  • buildgrid/server/cas/storage/remote.py
    ... ... @@ -25,9 +25,12 @@ import logging
    25 25
    26 26 
     import grpc
    27 27
    28 
    -from buildgrid.utils import gen_fetch_blob, gen_write_request_blob
    28 
    +from buildgrid.client.cas import upload
    29 29 
     from buildgrid._protos.google.bytestream import bytestream_pb2_grpc
    30 30 
     from buildgrid._protos.build.bazel.remote.execution.v2 import remote_execution_pb2, remote_execution_pb2_grpc
    31 
    +from buildgrid._protos.google.rpc import code_pb2
    32 
    +from buildgrid._protos.google.rpc import status_pb2
    33 
    +from buildgrid.utils import gen_fetch_blob
    31 34
    32 35 
     from .storage_abc import StorageABC
    33 36
    ... ... @@ -36,7 +39,10 @@ class RemoteStorage(StorageABC):
    36 39
    37 40 
         def __init__(self, channel, instance_name):
    38 41 
             self.logger = logging.getLogger(__name__)
    39 
    -        self._instance_name = instance_name
    42 
    +
    43 
    +        self.instance_name = instance_name
    44 
    +        self.channel = channel
    45 
    +
    40 46 
             self._stub_bs = bytestream_pb2_grpc.ByteStreamStub(channel)
    41 47 
             self._stub_cas = remote_execution_pb2_grpc.ContentAddressableStorageStub(channel)
    42 48
    ... ... @@ -50,7 +56,7 @@ class RemoteStorage(StorageABC):
    50 56 
                 fetched_data = io.BytesIO()
    51 57 
                 length = 0
    52 58
    53 
    -            for data in gen_fetch_blob(self._stub_bs, digest, self._instance_name):
    59 
    +            for data in gen_fetch_blob(self._stub_bs, digest, self.instance_name):
    54 60 
                     length += fetched_data.write(data)
    55 61
    56 62 
                 if length:
    ... ... @@ -71,16 +77,14 @@ class RemoteStorage(StorageABC):
    71 77 
             return None
    72 78
    73 79 
         def begin_write(self, digest):
    74 
    -        return io.BytesIO(digest.SerializeToString())
    80 
    +        return io.BytesIO()
    75 81
    76 82 
         def commit_write(self, digest, write_session):
    77 
    -        write_session.seek(0)
    78 
    -
    79 
    -        for request in gen_write_request_blob(write_session, digest, self._instance_name):
    80 
    -            self._stub_bs.Write(request)
    83 
    +        with upload(self.channel, instance=self.instance_name) as cas:
    84 
    +            cas.put_blob(write_session.getvalue())
    81 85
    82 86 
         def missing_blobs(self, blobs):
    83 
    -        request = remote_execution_pb2.FindMissingBlobsRequest(instance_name=self._instance_name)
    87 
    +        request = remote_execution_pb2.FindMissingBlobsRequest(instance_name=self.instance_name)
    84 88
    85 89 
             for blob in blobs:
    86 90 
                 request_digest = request.blob_digests.add()
    ... ... @@ -92,19 +96,12 @@ class RemoteStorage(StorageABC):
    92 96 
             return [x for x in response.missing_blob_digests]
    93 97
    94 98 
         def bulk_update_blobs(self, blobs):
    95 
    -        request = remote_execution_pb2.BatchUpdateBlobsRequest(instance_name=self._instance_name)
    96 
    -
    97 
    -        for digest, data in blobs:
    98 
    -            reqs = request.requests.add()
    99 
    -            reqs.digest.CopyFrom(digest)
    100 
    -            reqs.data = data
    101 
    -
    102 
    -        response = self._stub_cas.BatchUpdateBlobs(request)
    103 
    -
    104 
    -        responses = response.responses
    99 
    +        sent_digests = list()
    100 
    +        with upload(self.channel, instance=self.instance_name) as cas:
    101 
    +            for digest, blob in blobs:
    102 
    +                sent_digests.append(cas.put_blob(blob, digest=digest, queue=True))
    105 103
    106 
    -        # Check everything was sent back, even if order changed
    107 
    -        assert ([x.digest for x in request.requests].sort(key=lambda x: x.hash)) == \
    108 
    -            ([x.digest for x in responses].sort(key=lambda x: x.hash))
    104 
    +        assert len(sent_digests) == len(blobs)
    109 105
    110 
    -        return [x.status for x in responses]
    106 
    +        return [status_pb2.Status(code=code_pb2.OK) if d.ByteSize() > 0
    107 
    +                else status_pb2.Status(code=code_pb2.UNKNOWN) for d in sent_digests]

  • buildgrid/utils.py
    ... ... @@ -15,7 +15,6 @@
    15 15
    16 16 
     from operator import attrgetter
    17 17 
     import os
    18 
    -import uuid
    19 18
    20 19 
     from buildgrid.settings import HASH
    21 20 
     from buildgrid._protos.build.bazel.remote.execution.v2 import remote_execution_pb2
    ... ... @@ -34,32 +33,6 @@ def gen_fetch_blob(stub, digest, instance_name=""):
    34 33 
             yield response.data
    35 34
    36 35
    37 
    -def gen_write_request_blob(digest_bytes, digest, instance_name=""):
    38 
    -    """ Generates a bytestream write request
    39 
    -    """
    40 
    -    resource_name = os.path.join(instance_name, 'uploads', str(uuid.uuid4()),
    41 
    -                                 'blobs', digest.hash, str(digest.size_bytes))
    42 
    -
    43 
    -    offset = 0
    44 
    -    finished = False
    45 
    -    remaining = digest.size_bytes
    46 
    -
    47 
    -    while not finished:
    48 
    -        chunk_size = min(remaining, 64 * 1024)
    49 
    -        remaining -= chunk_size
    50 
    -        finished = remaining <= 0
    51 
    -
    52 
    -        request = bytestream_pb2.WriteRequest()
    53 
    -        request.resource_name = resource_name
    54 
    -        request.write_offset = offset
    55 
    -        request.data = digest_bytes.read(chunk_size)
    56 
    -        request.finish_write = finished
    57 
    -
    58 
    -        yield request
    59 
    -
    60 
    -        offset += chunk_size
    61 
    -
    62 
    -
    63 36 
     def write_fetch_directory(root_directory, stub, digest, instance_name=None):
    64 37 
         """Locally replicates a directory from CAS.
    65 38
    ... ... @@ -280,8 +253,12 @@ def tree_maker(directory_path, cas=None):
    280 253 
         tree.children.extend(child_directories)
    281 254 
         tree.root.CopyFrom(directory)
    282 255
    256 
    +    # Ensure that we've uploded the tree structure first
    257 
    +    if cas is not None:
    258 
    +        cas.flush()
    259 
    +
    283 260 
         if cas is not None:
    284 
    -        tree_digest = cas.send_message(tree)
    261 
    +        tree_digest = cas.put_message(tree)
    285 262 
         else:
    286 263 
             tree_digest = create_digest(tree.SerializeToString())
    287 264

  • setup.py
    ... ... @@ -89,6 +89,7 @@ tests_require = [
    89 89 
         'coverage == 4.4.0',
    90 90 
         'moto',
    91 91 
         'pep8',
    92 
    +    'psutil',
    92 93 
         'pytest == 3.6.4',
    93 94 
         'pytest-cov >= 2.6.0',
    94 95 
         'pytest-pep8',

  • tests/cas/data/hello.cc
    1 
    +#include <iostream>
    2 
    +
    3 
    +int main()
    4 
    +{
    5 
    +  std::cout << "Hello, World!" << std::endl;
    6 
    +  return 0;
    7 
    +}

  • tests/cas/data/hello/hello.c
    1 
    +#include <stdio.h>
    2 
    +
    3 
    +#include "hello.h"
    4 
    +
    5 
    +int main()
    6 
    +{
    7 
    +  printf("%s\n", HELLO_WORLD);
    8 
    +  return 0;
    9 
    +}

  • tests/cas/data/hello/hello.h
    1 
    +#define HELLO_WORLD "Hello, World!"

  • tests/cas/data/void

  • tests/cas/test_client.py
    1 
    +# Copyright (C) 2018 Bloomberg LP
    2 
    +#
    3 
    +# Licensed under the Apache License, Version 2.0 (the "License");
    4 
    +# you may not use this file except in compliance with the License.
    5 
    +# You may obtain a copy of the License at
    6 
    +#
    7 
    +#  <http://www.apache.org/licenses/LICENSE-2.0>
    8 
    +#
    9 
    +# Unless required by applicable law or agreed to in writing, software
    10 
    +# distributed under the License is distributed on an "AS IS" BASIS,
    11 
    +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12 
    +# See the License for the specific language governing permissions and
    13 
    +# limitations under the License.
    14 
    +
    15 
    +# pylint: disable=redefined-outer-name
    16 
    +
    17 
    +import multiprocessing
    18 
    +import os
    19 
    +
    20 
    +import grpc
    21 
    +import pytest
    22 
    +
    23 
    +from buildgrid.client.cas import upload
    24 
    +from buildgrid._protos.build.bazel.remote.execution.v2 import remote_execution_pb2
    25 
    +
    26 
    +from ..utils.cas import serve_cas, kill_process_tree
    27 
    +
    28 
    +
    29 
    +INTANCES = ['', 'instance']
    30 
    +BLOBS = [(b'',), (b'test-string',), (b'test', b'string')]
    31 
    +MESSAGES = [
    32 
    +    (remote_execution_pb2.Directory(),),
    33 
    +    (remote_execution_pb2.SymlinkNode(name='name', target='target'),),
    34 
    +    (remote_execution_pb2.Action(do_not_cache=True),
    35 
    +     remote_execution_pb2.ActionResult(exit_code=12))
    36 
    +]
    37 
    +DATA_DIR = os.path.join(
    38 
    +    os.path.dirname(os.path.realpath(__file__)), 'data')
    39 
    +FILES = [
    40 
    +    (os.path.join(DATA_DIR, 'void'),),
    41 
    +    (os.path.join(DATA_DIR, 'hello.cc'),),
    42 
    +    (os.path.join(DATA_DIR, 'hello', 'hello.c'),
    43 
    +     os.path.join(DATA_DIR, 'hello', 'hello.h'))]
    44 
    +DIRECTORIES = [
    45 
    +    (remote_execution_pb2.Directory(),),
    46 
    +    (remote_execution_pb2.Directory(
    47 
    +        files=[remote_execution_pb2.FileNode(name='helloc.c'),
    48 
    +               remote_execution_pb2.FileNode(name='helloc.h')]),)]
    49 
    +
    50 
    +
    51 
    +def run_in_subprocess(function, *arguments):
    52 
    +    queue = multiprocessing.Queue()
    53 
    +    # Use subprocess to avoid creation of gRPC threads in main process
    54 
    +    # See https://github.com/grpc/grpc/blob/master/doc/fork_support.md
    55 
    +    process = multiprocessing.Process(target=function,
    56 
    +                                      args=(queue, *arguments))
    57 
    +
    58 
    +    try:
    59 
    +        process.start()
    60 
    +
    61 
    +        result = queue.get()
    62 
    +        process.join()
    63 
    +    except KeyboardInterrupt:
    64 
    +        kill_process_tree(process.pid)
    65 
    +        raise
    66 
    +
    67 
    +    return result
    68 
    +
    69 
    +
    70 
    +@pytest.mark.parametrize('blobs', BLOBS)
    71 
    +@pytest.mark.parametrize('instance', INTANCES)
    72 
    +def test_blob_upload(instance, blobs):
    73 
    +    # Actual test function, to be run in a subprocess:
    74 
    +    def __test_blob_upload(queue, remote, instance, blobs):
    75 
    +        # Open a channel to the remote CAS server:
    76 
    +        channel = grpc.insecure_channel(remote)
    77 
    +
    78 
    +        digests = list()
    79 
    +        with upload(channel, instance) as client:
    80 
    +            if len(blobs) > 1:
    81 
    +                for blob in blobs:
    82 
    +                    digest = client.put_blob(blob, queue=True)
    83 
    +                    digests.append(digest.SerializeToString())
    84 
    +            else:
    85 
    +                digest = client.put_blob(blobs[0], queue=False)
    86 
    +                digests.append(digest.SerializeToString())
    87 
    +
    88 
    +        queue.put(digests)
    89 
    +
    90 
    +    # Start a minimal CAS server in a subprocess:
    91 
    +    with serve_cas([instance]) as server:
    92 
    +        digests = run_in_subprocess(__test_blob_upload,
    93 
    +                                    server.remote, instance, blobs)
    94 
    +
    95 
    +        for blob, digest_blob in zip(blobs, digests):
    96 
    +            digest = remote_execution_pb2.Digest()
    97 
    +            digest.ParseFromString(digest_blob)
    98 
    +
    99 
    +            assert server.has(digest)
    100 
    +            assert server.compare_blobs(digest, blob)
    101 
    +
    102 
    +
    103 
    +@pytest.mark.parametrize('messages', MESSAGES)
    104 
    +@pytest.mark.parametrize('instance', INTANCES)
    105 
    +def test_message_upload(instance, messages):
    106 
    +    # Actual test function, to be run in a subprocess:
    107 
    +    def __test_message_upload(queue, remote, instance, messages):
    108 
    +        # Open a channel to the remote CAS server:
    109 
    +        channel = grpc.insecure_channel(remote)
    110 
    +
    111 
    +        digests = list()
    112 
    +        with upload(channel, instance) as client:
    113 
    +            if len(messages) > 1:
    114 
    +                for message in messages:
    115 
    +                    digest = client.put_message(message, queue=True)
    116 
    +                    digests.append(digest.SerializeToString())
    117 
    +            else:
    118 
    +                digest = client.put_message(messages[0], queue=False)
    119 
    +                digests.append(digest.SerializeToString())
    120 
    +
    121 
    +        queue.put(digests)
    122 
    +
    123 
    +    # Start a minimal CAS server in a subprocess:
    124 
    +    with serve_cas([instance]) as server:
    125 
    +        digests = run_in_subprocess(__test_message_upload,
    126 
    +                                    server.remote, instance, messages)
    127 
    +
    128 
    +        for message, digest_blob in zip(messages, digests):
    129 
    +            digest = remote_execution_pb2.Digest()
    130 
    +            digest.ParseFromString(digest_blob)
    131 
    +
    132 
    +            assert server.has(digest)
    133 
    +            assert server.compare_messages(digest, message)
    134 
    +
    135 
    +
    136 
    +@pytest.mark.parametrize('file_paths', FILES)
    137 
    +@pytest.mark.parametrize('instance', INTANCES)
    138 
    +def test_file_upload(instance, file_paths):
    139 
    +    # Actual test function, to be run in a subprocess:
    140 
    +    def __test_file_upload(queue, remote, instance, file_paths):
    141 
    +        # Open a channel to the remote CAS server:
    142 
    +        channel = grpc.insecure_channel(remote)
    143 
    +
    144 
    +        digests = list()
    145 
    +        with upload(channel, instance) as client:
    146 
    +            if len(file_paths) > 1:
    147 
    +                for file_path in file_paths:
    148 
    +                    digest = client.upload_file(file_path, queue=True)
    149 
    +                    digests.append(digest.SerializeToString())
    150 
    +            else:
    151 
    +                digest = client.upload_file(file_paths[0], queue=False)
    152 
    +                digests.append(digest.SerializeToString())
    153 
    +
    154 
    +        queue.put(digests)
    155 
    +
    156 
    +    # Start a minimal CAS server in a subprocess:
    157 
    +    with serve_cas([instance]) as server:
    158 
    +        digests = run_in_subprocess(__test_file_upload,
    159 
    +                                    server.remote, instance, file_paths)
    160 
    +
    161 
    +        for file_path, digest_blob in zip(file_paths, digests):
    162 
    +            digest = remote_execution_pb2.Digest()
    163 
    +            digest.ParseFromString(digest_blob)
    164 
    +
    165 
    +            assert server.has(digest)
    166 
    +            assert server.compare_files(digest, file_path)
    167 
    +
    168 
    +
    169 
    +@pytest.mark.parametrize('directories', DIRECTORIES)
    170 
    +@pytest.mark.parametrize('instance', INTANCES)
    171 
    +def test_directory_upload(instance, directories):
    172 
    +    # Actual test function, to be run in a subprocess:
    173 
    +    def __test_directory_upload(queue, remote, instance, directories):
    174 
    +        # Open a channel to the remote CAS server:
    175 
    +        channel = grpc.insecure_channel(remote)
    176 
    +
    177 
    +        digests = list()
    178 
    +        with upload(channel, instance) as client:
    179 
    +            if len(directories) > 1:
    180 
    +                for directory in directories:
    181 
    +                    digest = client.upload_directory(directory, queue=True)
    182 
    +                    digests.append(digest.SerializeToString())
    183 
    +            else:
    184 
    +                digest = client.upload_directory(directories[0], queue=False)
    185 
    +                digests.append(digest.SerializeToString())
    186 
    +
    187 
    +        queue.put(digests)
    188 
    +
    189 
    +    # Start a minimal CAS server in a subprocess:
    190 
    +    with serve_cas([instance]) as server:
    191 
    +        digests = run_in_subprocess(__test_directory_upload,
    192 
    +                                    server.remote, instance, directories)
    193 
    +
    194 
    +        for directory, digest_blob in zip(directories, digests):
    195 
    +            digest = remote_execution_pb2.Digest()
    196 
    +            digest.ParseFromString(digest_blob)
    197 
    +
    198 
    +            assert server.has(digest)
    199 
    +            assert server.compare_messages(digest, directory)

  • tests/utils/__init__.py

  • tests/utils/cas.py
    1 
    +# Copyright (C) 2018 Bloomberg LP
    2 
    +#
    3 
    +# Licensed under the Apache License, Version 2.0 (the "License");
    4 
    +# you may not use this file except in compliance with the License.
    5 
    +# You may obtain a copy of the License at
    6 
    +#
    7 
    +#  <http://www.apache.org/licenses/LICENSE-2.0>
    8 
    +#
    9 
    +# Unless required by applicable law or agreed to in writing, software
    10 
    +# distributed under the License is distributed on an "AS IS" BASIS,
    11 
    +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12 
    +# See the License for the specific language governing permissions and
    13 
    +# limitations under the License.
    14 
    +
    15 
    +
    16 
    +from concurrent import futures
    17 
    +from contextlib import contextmanager
    18 
    +import multiprocessing
    19 
    +import os
    20 
    +import signal
    21 
    +import tempfile
    22 
    +
    23 
    +import grpc
    24 
    +import psutil
    25 
    +import pytest_cov
    26 
    +
    27 
    +from buildgrid.server.cas.service import ByteStreamService
    28 
    +from buildgrid.server.cas.service import ContentAddressableStorageService
    29 
    +from buildgrid.server.cas.instance import ByteStreamInstance
    30 
    +from buildgrid.server.cas.instance import ContentAddressableStorageInstance
    31 
    +from buildgrid.server.cas.storage.disk import DiskStorage
    32 
    +
    33 
    +
    34 
    +@contextmanager
    35 
    +def serve_cas(instances):
    36 
    +    server = Server(instances)
    37 
    +    try:
    38 
    +        yield server
    39 
    +    finally:
    40 
    +        server.quit()
    41 
    +
    42 
    +
    43 
    +def kill_process_tree(pid):
    44 
    +    proc = psutil.Process(pid)
    45 
    +    children = proc.children(recursive=True)
    46 
    +
    47 
    +    def kill_proc(p):
    48 
    +        try:
    49 
    +            p.kill()
    50 
    +        except psutil.AccessDenied:
    51 
    +            # Ignore this error, it can happen with
    52 
    +            # some setuid bwrap processes.
    53 
    +            pass
    54 
    +
    55 
    +    # Bloody Murder
    56 
    +    for child in children:
    57 
    +        kill_proc(child)
    58 
    +    kill_proc(proc)
    59 
    +
    60 
    +
    61 
    +class Server:
    62 
    +
    63 
    +    def __init__(self, instances):
    64 
    +
    65 
    +        self.instances = instances
    66 
    +
    67 
    +        self.__storage_path = tempfile.TemporaryDirectory()
    68 
    +        self.__storage = DiskStorage(self.__storage_path.name)
    69 
    +
    70 
    +        self.__queue = multiprocessing.Queue()
    71 
    +        self.__process = multiprocessing.Process(
    72 
    +            target=Server.serve,
    73 
    +            args=(self.__queue, self.instances, self.__storage_path.name))
    74 
    +        self.__process.start()
    75 
    +
    76 
    +        self.port = self.__queue.get()
    77 
    +        self.remote = 'localhost:{}'.format(self.port)
    78 
    +
    79 
    +    @classmethod
    80 
    +    def serve(cls, queue, instances, storage_path):
    81 
    +        pytest_cov.embed.cleanup_on_sigterm()
    82 
    +
    83 
    +        bs_instances, cas_instances = dict(), dict()
    84 
    +        for name in instances:
    85 
    +            storage = DiskStorage(storage_path)
    86 
    +
    87 
    +            bs_instances[name] = ByteStreamInstance(storage)
    88 
    +            cas_instances[name] = ContentAddressableStorageInstance(storage)
    89 
    +
    90 
    +        # Use max_workers default from Python 3.5+
    91 
    +        max_workers = (os.cpu_count() or 1) * 5
    92 
    +        server = grpc.server(futures.ThreadPoolExecutor(max_workers))
    93 
    +        port = server.add_insecure_port('localhost:0')
    94 
    +
    95 
    +        ByteStreamService(server, bs_instances)
    96 
    +        ContentAddressableStorageService(server, cas_instances)
    97 
    +
    98 
    +        server.start()
    99 
    +        queue.put(port)
    100 
    +
    101 
    +        signal.pause()
    102 
    +
    103 
    +    def has(self, digest):
    104 
    +        return self.__storage.has_blob(digest)
    105 
    +
    106 
    +    def compare_blobs(self, digest, blob):
    107 
    +        if not self.__storage.has_blob(digest):
    108 
    +            return False
    109 
    +
    110 
    +        stored_blob = self.__storage.get_blob(digest)
    111 
    +        stored_blob = stored_blob.read()
    112 
    +
    113 
    +        return blob == stored_blob
    114 
    +
    115 
    +    def compare_messages(self, digest, message):
    116 
    +        if not self.__storage.has_blob(digest):
    117 
    +            return False
    118 
    +
    119 
    +        message_blob = message.SerializeToString()
    120 
    +
    121 
    +        stored_blob = self.__storage.get_blob(digest)
    122 
    +        stored_blob = stored_blob.read()
    123 
    +
    124 
    +        return message_blob == stored_blob
    125 
    +
    126 
    +    def compare_files(self, digest, file_path):
    127 
    +        if not self.__storage.has_blob(digest):
    128 
    +            return False
    129 
    +
    130 
    +        with open(file_path, 'rb') as file_bytes:
    131 
    +            file_blob = file_bytes.read()
    132 
    +
    133 
    +        stored_blob = self.__storage.get_blob(digest)
    134 
    +        stored_blob = stored_blob.read()
    135 
    +
    136 
    +        return file_blob == stored_blob
    137 
    +
    138 
    +    def quit(self):
    139 
    +        if self.__process:
    140 
    +            self.__process.terminate()
    141 
    +            self.__process.join()
    142 
    +
    143 
    +        self.__storage_path.cleanup()

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

    [Date Prev][Date Next]   [Thread Prev][Thread Next]   [Thread Index] [Date Index] [Author Index]