[Notes] [Git][BuildGrid/buildgrid][finn/tutorial-fix] 12 commits: _expection.py: Centralise exception declarations



Title: GitLab

finn pushed to branch finn/tutorial-fix at BuildGrid / buildgrid

Commits:

28 changed files:

Changes:

  • buildgrid/_app/bots/buildbox.py
    ... ... @@ -19,32 +19,34 @@ import tempfile
    19 19
     
    
    20 20
     from google.protobuf import any_pb2
    
    21 21
     
    
    22
    -from buildgrid.settings import HASH_LENGTH
    
    23
    -from buildgrid.client.cas import upload
    
    24
    -from buildgrid._protos.build.bazel.remote.execution.v2 import remote_execution_pb2
    
    25
    -from buildgrid._protos.google.bytestream import bytestream_pb2_grpc
    
    22
    +from buildgrid.client.cas import download, upload
    
    26 23
     from buildgrid._exceptions import BotError
    
    27
    -from buildgrid.utils import read_file, write_file, parse_to_pb2_from_fetch
    
    24
    +from buildgrid._protos.build.bazel.remote.execution.v2 import remote_execution_pb2
    
    25
    +from buildgrid.settings import HASH_LENGTH
    
    26
    +from buildgrid.utils import read_file, write_file
    
    28 27
     
    
    29 28
     
    
    30 29
     def work_buildbox(context, lease):
    
    31 30
         """Executes a lease for a build action, using buildbox.
    
    32 31
         """
    
    33 32
     
    
    34
    -    stub_bytestream = bytestream_pb2_grpc.ByteStreamStub(context.cas_channel)
    
    35 33
         local_cas_directory = context.local_cas
    
    34
    +    # instance_name = context.parent
    
    36 35
         logger = context.logger
    
    37 36
     
    
    38 37
         action_digest = remote_execution_pb2.Digest()
    
    39 38
         lease.payload.Unpack(action_digest)
    
    40 39
     
    
    41
    -    action = parse_to_pb2_from_fetch(remote_execution_pb2.Action(),
    
    42
    -                                     stub_bytestream, action_digest)
    
    40
    +    with download(context.cas_channel) as downloader:
    
    41
    +        action = downloader.get_message(action_digest,
    
    42
    +                                        remote_execution_pb2.Action())
    
    43 43
     
    
    44
    -    command = parse_to_pb2_from_fetch(remote_execution_pb2.Command(),
    
    45
    -                                      stub_bytestream, action.command_digest)
    
    44
    +        assert action.command_digest.hash
    
    46 45
     
    
    47
    -    environment = dict()
    
    46
    +        command = downloader.get_message(action.command_digest,
    
    47
    +                                         remote_execution_pb2.Command())
    
    48
    +
    
    49
    +    environment = {}
    
    48 50
         for variable in command.environment_variables:
    
    49 51
             if variable.name not in ['PWD']:
    
    50 52
                 environment[variable.name] = variable.value
    
    ... ... @@ -116,10 +118,11 @@ def work_buildbox(context, lease):
    116 118
     
    
    117 119
                 # TODO: Have BuildBox helping us creating the Tree instance here
    
    118 120
                 # See https://gitlab.com/BuildStream/buildbox/issues/7 for details
    
    119
    -            output_tree = _cas_tree_maker(stub_bytestream, output_digest)
    
    121
    +            with download(context.cas_channel) as downloader:
    
    122
    +                output_tree = _cas_tree_maker(downloader, output_digest)
    
    120 123
     
    
    121
    -            with upload(context.cas_channel) as cas:
    
    122
    -                output_tree_digest = cas.put_message(output_tree)
    
    124
    +            with upload(context.cas_channel) as uploader:
    
    125
    +                output_tree_digest = uploader.put_message(output_tree)
    
    123 126
     
    
    124 127
                 output_directory = remote_execution_pb2.OutputDirectory()
    
    125 128
                 output_directory.tree_digest.CopyFrom(output_tree_digest)
    
    ... ... @@ -135,24 +138,28 @@ def work_buildbox(context, lease):
    135 138
         return lease
    
    136 139
     
    
    137 140
     
    
    138
    -def _cas_tree_maker(stub_bytestream, directory_digest):
    
    141
    +def _cas_tree_maker(cas, directory_digest):
    
    139 142
         # Generates and stores a Tree for a given Directory. This is very inefficient
    
    140 143
         # and only temporary. See https://gitlab.com/BuildStream/buildbox/issues/7.
    
    141 144
         output_tree = remote_execution_pb2.Tree()
    
    142 145
     
    
    143
    -    def list_directories(parent_directory):
    
    144
    -        directory_list = list()
    
    146
    +    def __cas_tree_maker(cas, parent_directory):
    
    147
    +        digests, directories = [], []
    
    145 148
             for directory_node in parent_directory.directories:
    
    146
    -            directory = parse_to_pb2_from_fetch(remote_execution_pb2.Directory(),
    
    147
    -                                                stub_bytestream, directory_node.digest)
    
    148
    -            directory_list.extend(list_directories(directory))
    
    149
    -            directory_list.append(directory)
    
    149
    +            directories.append(remote_execution_pb2.Directory())
    
    150
    +            digests.append(directory_node.digest)
    
    151
    +
    
    152
    +        cas.get_messages(digests, directories)
    
    153
    +
    
    154
    +        for directory in directories[:]:
    
    155
    +            directories.extend(__cas_tree_maker(cas, directory))
    
    156
    +
    
    157
    +        return directories
    
    150 158
     
    
    151
    -        return directory_list
    
    159
    +    root_directory = cas.get_message(directory_digest,
    
    160
    +                                     remote_execution_pb2.Directory())
    
    152 161
     
    
    153
    -    root_directory = parse_to_pb2_from_fetch(remote_execution_pb2.Directory(),
    
    154
    -                                             stub_bytestream, directory_digest)
    
    155
    -    output_tree.children.extend(list_directories(root_directory))
    
    162
    +    output_tree.children.extend(__cas_tree_maker(cas, root_directory))
    
    156 163
         output_tree.root.CopyFrom(root_directory)
    
    157 164
     
    
    158 165
         return output_tree

  • buildgrid/_app/bots/temp_directory.py
    ... ... @@ -19,10 +19,8 @@ import tempfile
    19 19
     
    
    20 20
     from google.protobuf import any_pb2
    
    21 21
     
    
    22
    -from buildgrid.client.cas import upload
    
    22
    +from buildgrid.client.cas import download, upload
    
    23 23
     from buildgrid._protos.build.bazel.remote.execution.v2 import remote_execution_pb2
    
    24
    -from buildgrid._protos.google.bytestream import bytestream_pb2_grpc
    
    25
    -from buildgrid.utils import write_fetch_directory, parse_to_pb2_from_fetch
    
    26 24
     from buildgrid.utils import output_file_maker, output_directory_maker
    
    27 25
     
    
    28 26
     
    
    ... ... @@ -30,29 +28,30 @@ def work_temp_directory(context, lease):
    30 28
         """Executes a lease for a build action, using host tools.
    
    31 29
         """
    
    32 30
     
    
    33
    -    stub_bytestream = bytestream_pb2_grpc.ByteStreamStub(context.cas_channel)
    
    34 31
         instance_name = context.parent
    
    35 32
         logger = context.logger
    
    36 33
     
    
    37 34
         action_digest = remote_execution_pb2.Digest()
    
    38 35
         lease.payload.Unpack(action_digest)
    
    39 36
     
    
    40
    -    action = parse_to_pb2_from_fetch(remote_execution_pb2.Action(),
    
    41
    -                                     stub_bytestream, action_digest, instance_name)
    
    42
    -
    
    43 37
         with tempfile.TemporaryDirectory() as temp_directory:
    
    44
    -        command = parse_to_pb2_from_fetch(remote_execution_pb2.Command(),
    
    45
    -                                          stub_bytestream, action.command_digest, instance_name)
    
    38
    +        with download(context.cas_channel, instance=instance_name) as downloader:
    
    39
    +            action = downloader.get_message(action_digest,
    
    40
    +                                            remote_execution_pb2.Action())
    
    41
    +
    
    42
    +            assert action.command_digest.hash
    
    43
    +
    
    44
    +            command = downloader.get_message(action.command_digest,
    
    45
    +                                             remote_execution_pb2.Command())
    
    46 46
     
    
    47
    -        write_fetch_directory(temp_directory, stub_bytestream,
    
    48
    -                              action.input_root_digest, instance_name)
    
    47
    +            downloader.download_directory(action.input_root_digest, temp_directory)
    
    49 48
     
    
    50 49
             environment = os.environ.copy()
    
    51 50
             for variable in command.environment_variables:
    
    52 51
                 if variable.name not in ['PATH', 'PWD']:
    
    53 52
                     environment[variable.name] = variable.value
    
    54 53
     
    
    55
    -        command_line = list()
    
    54
    +        command_line = []
    
    56 55
             for argument in command.arguments:
    
    57 56
                 command_line.append(argument.strip())
    
    58 57
     
    

  • buildgrid/_app/commands/cmd_cas.py
    ... ... @@ -28,9 +28,9 @@ from urllib.parse import urlparse
    28 28
     import click
    
    29 29
     import grpc
    
    30 30
     
    
    31
    -from buildgrid.client.cas import upload
    
    31
    +from buildgrid.client.cas import download, upload
    
    32 32
     from buildgrid._protos.build.bazel.remote.execution.v2 import remote_execution_pb2
    
    33
    -from buildgrid.utils import merkle_tree_maker
    
    33
    +from buildgrid.utils import create_digest, merkle_tree_maker, read_file
    
    34 34
     
    
    35 35
     from ..cli import pass_context
    
    36 36
     
    
    ... ... @@ -81,38 +81,46 @@ def upload_dummy(context):
    81 81
             click.echo("Error: Failed pushing empty message.", err=True)
    
    82 82
     
    
    83 83
     
    
    84
    -@cli.command('upload-files', short_help="Upload files to the CAS server.")
    
    85
    -@click.argument('files', nargs=-1, type=click.Path(exists=True, dir_okay=False), required=True)
    
    84
    +@cli.command('upload-file', short_help="Upload files to the CAS server.")
    
    85
    +@click.argument('file_path', nargs=-1, type=click.Path(exists=True, dir_okay=False), required=True)
    
    86
    +@click.option('--verify', is_flag=True, show_default=True,
    
    87
    +              help="Check uploaded files integrity.")
    
    86 88
     @pass_context
    
    87
    -def upload_files(context, files):
    
    89
    +def upload_file(context, file_path, verify):
    
    88 90
         sent_digests, files_map = [], {}
    
    89 91
         with upload(context.channel, instance=context.instance_name) as uploader:
    
    90
    -        for file_path in files:
    
    91
    -            context.logger.debug("Queueing {}".format(file_path))
    
    92
    +        for path in file_path:
    
    93
    +            if not os.path.isabs(path):
    
    94
    +                path = os.path.abspath(path)
    
    95
    +            context.logger.debug("Queueing {}".format(path))
    
    92 96
     
    
    93
    -            file_digest = uploader.upload_file(file_path, queue=True)
    
    97
    +            file_digest = uploader.upload_file(path, queue=True)
    
    94 98
     
    
    95
    -            files_map[file_digest.hash] = file_path
    
    99
    +            files_map[file_digest.hash] = path
    
    96 100
                 sent_digests.append(file_digest)
    
    97 101
     
    
    98 102
         for file_digest in sent_digests:
    
    99
    -        file_path = files_map[file_digest.hash]
    
    100
    -        if os.path.isabs(file_path):
    
    101
    -            file_path = os.path.relpath(file_path)
    
    102
    -        if file_digest.ByteSize():
    
    103
    +        file_path = os.path.relpath(files_map[file_digest.hash])
    
    104
    +        if verify and file_digest.size_bytes != os.stat(file_path).st_size:
    
    105
    +            click.echo('Error: Failed to verify "{}"'.format(file_path), err=True)
    
    106
    +        elif file_digest.ByteSize():
    
    103 107
                 click.echo('Success: Pushed "{}" with digest "{}/{}"'
    
    104 108
                            .format(file_path, file_digest.hash, file_digest.size_bytes))
    
    105 109
             else:
    
    106
    -            click.echo('Error: Failed to push "{}"'.format(file_path), err=True)
    
    110
    +            click.echo('Error: Failed pushing "{}"'.format(file_path), err=True)
    
    107 111
     
    
    108 112
     
    
    109 113
     @cli.command('upload-dir', short_help="Upload a directory to the CAS server.")
    
    110
    -@click.argument('directory', nargs=1, type=click.Path(exists=True, file_okay=False), required=True)
    
    114
    +@click.argument('directory-path', nargs=1, type=click.Path(exists=True, file_okay=False), required=True)
    
    115
    +@click.option('--verify', is_flag=True, show_default=True,
    
    116
    +              help="Check uploaded directory's integrity.")
    
    111 117
     @pass_context
    
    112
    -def upload_dir(context, directory):
    
    118
    +def upload_directory(context, directory_path, verify):
    
    113 119
         sent_digests, nodes_map = [], {}
    
    114 120
         with upload(context.channel, instance=context.instance_name) as uploader:
    
    115
    -        for node, blob, path in merkle_tree_maker(directory):
    
    121
    +        for node, blob, path in merkle_tree_maker(directory_path):
    
    122
    +            if not os.path.isabs(path):
    
    123
    +                path = os.path.abspath(path)
    
    116 124
                 context.logger.debug("Queueing {}".format(path))
    
    117 125
     
    
    118 126
                 node_digest = uploader.put_blob(blob, digest=node.digest, queue=True)
    
    ... ... @@ -122,10 +130,85 @@ def upload_dir(context, directory):
    122 130
     
    
    123 131
         for node_digest in sent_digests:
    
    124 132
             node_path = nodes_map[node_digest.hash]
    
    125
    -        if os.path.isabs(node_path):
    
    126
    -            node_path = os.path.relpath(node_path, start=directory)
    
    127
    -        if node_digest.ByteSize():
    
    133
    +        if not os.path.isabs(directory_path):
    
    134
    +            node_path = os.path.relpath(node_path)
    
    135
    +        if verify and (os.path.isfile(node_path) and
    
    136
    +                       node_digest.size_bytes != os.stat(node_path).st_size):
    
    137
    +            click.echo('Error: Failed to verify "{}"'.format(node_path), err=True)
    
    138
    +        elif node_digest.ByteSize():
    
    128 139
                 click.echo('Success: Pushed "{}" with digest "{}/{}"'
    
    129 140
                            .format(node_path, node_digest.hash, node_digest.size_bytes))
    
    130 141
             else:
    
    131
    -            click.echo('Error: Failed to push "{}"'.format(node_path), err=True)
    142
    +            click.echo('Error: Failed pushing "{}"'.format(node_path), err=True)
    
    143
    +
    
    144
    +
    
    145
    +def _create_digest(digest_string):
    
    146
    +    digest_hash, digest_size = digest_string.split('/')
    
    147
    +
    
    148
    +    digest = remote_execution_pb2.Digest()
    
    149
    +    digest.hash = digest_hash
    
    150
    +    digest.size_bytes = int(digest_size)
    
    151
    +
    
    152
    +    return digest
    
    153
    +
    
    154
    +
    
    155
    +@cli.command('download-file', short_help="Download a file from the CAS server.")
    
    156
    +@click.argument('digest-string', nargs=1, type=click.STRING, required=True)
    
    157
    +@click.argument('file-path', nargs=1, type=click.Path(exists=False), required=True)
    
    158
    +@click.option('--verify', is_flag=True, show_default=True,
    
    159
    +              help="Check downloaded file's integrity.")
    
    160
    +@pass_context
    
    161
    +def download_file(context, digest_string, file_path, verify):
    
    162
    +    if os.path.exists(file_path):
    
    163
    +        click.echo('Error: Invalid value for "file-path": ' +
    
    164
    +                   'Path "{}" already exists.'.format(file_path), err=True)
    
    165
    +        return
    
    166
    +
    
    167
    +    digest = _create_digest(digest_string)
    
    168
    +    with download(context.channel, instance=context.instance_name) as downloader:
    
    169
    +        downloader.download_file(digest, file_path)
    
    170
    +
    
    171
    +    if verify:
    
    172
    +        file_digest = create_digest(read_file(file_path))
    
    173
    +        if file_digest != digest:
    
    174
    +            click.echo('Error: Failed to verify "{}"'.format(file_path), err=True)
    
    175
    +            return
    
    176
    +
    
    177
    +    if os.path.isfile(file_path):
    
    178
    +        click.echo('Success: Pulled "{}" from digest "{}/{}"'
    
    179
    +                   .format(file_path, digest.hash, digest.size_bytes))
    
    180
    +    else:
    
    181
    +        click.echo('Error: Failed pulling "{}"'.format(file_path), err=True)
    
    182
    +
    
    183
    +
    
    184
    +@cli.command('download-dir', short_help="Download a directory from the CAS server.")
    
    185
    +@click.argument('digest-string', nargs=1, type=click.STRING, required=True)
    
    186
    +@click.argument('directory-path', nargs=1, type=click.Path(exists=False), required=True)
    
    187
    +@click.option('--verify', is_flag=True, show_default=True,
    
    188
    +              help="Check downloaded directory's integrity.")
    
    189
    +@pass_context
    
    190
    +def download_directory(context, digest_string, directory_path, verify):
    
    191
    +    if os.path.exists(directory_path):
    
    192
    +        if not os.path.isdir(directory_path) or os.listdir(directory_path):
    
    193
    +            click.echo('Error: Invalid value for "directory-path": ' +
    
    194
    +                       'Path "{}" already exists.'.format(directory_path), err=True)
    
    195
    +            return
    
    196
    +
    
    197
    +    digest = _create_digest(digest_string)
    
    198
    +    with download(context.channel, instance=context.instance_name) as downloader:
    
    199
    +        downloader.download_directory(digest, directory_path)
    
    200
    +
    
    201
    +    if verify:
    
    202
    +        last_directory_node = None
    
    203
    +        for node, _, _ in merkle_tree_maker(directory_path):
    
    204
    +            if node.DESCRIPTOR is remote_execution_pb2.DirectoryNode.DESCRIPTOR:
    
    205
    +                last_directory_node = node
    
    206
    +        if last_directory_node.digest != digest:
    
    207
    +            click.echo('Error: Failed to verify "{}"'.format(directory_path), err=True)
    
    208
    +            return
    
    209
    +
    
    210
    +    if os.path.isdir(directory_path):
    
    211
    +        click.echo('Success: Pulled "{}" from digest "{}/{}"'
    
    212
    +                   .format(directory_path, digest.hash, digest.size_bytes))
    
    213
    +    else:
    
    214
    +        click.echo('Error: Failed pulling "{}"'.format(directory_path), err=True)

  • buildgrid/_app/commands/cmd_execute.py
    ... ... @@ -20,7 +20,6 @@ Execute command
    20 20
     Request work to be executed and monitor status of jobs.
    
    21 21
     """
    
    22 22
     
    
    23
    -import errno
    
    24 23
     import logging
    
    25 24
     import os
    
    26 25
     import stat
    
    ... ... @@ -30,10 +29,9 @@ from urllib.parse import urlparse
    30 29
     import click
    
    31 30
     import grpc
    
    32 31
     
    
    33
    -from buildgrid.client.cas import upload
    
    32
    +from buildgrid.client.cas import download, upload
    
    34 33
     from buildgrid._protos.build.bazel.remote.execution.v2 import remote_execution_pb2, remote_execution_pb2_grpc
    
    35
    -from buildgrid._protos.google.bytestream import bytestream_pb2_grpc
    
    36
    -from buildgrid.utils import create_digest, write_fetch_blob
    
    34
    +from buildgrid.utils import create_digest
    
    37 35
     
    
    38 36
     from ..cli import pass_context
    
    39 37
     
    
    ... ... @@ -154,8 +152,6 @@ def run_command(context, input_root, commands, output_file, output_directory):
    154 152
                                                       skip_cache_lookup=True)
    
    155 153
         response = stub.Execute(request)
    
    156 154
     
    
    157
    -    stub = bytestream_pb2_grpc.ByteStreamStub(context.channel)
    
    158
    -
    
    159 155
         stream = None
    
    160 156
         for stream in response:
    
    161 157
             context.logger.info(stream)
    
    ... ... @@ -163,21 +159,16 @@ def run_command(context, input_root, commands, output_file, output_directory):
    163 159
         execute_response = remote_execution_pb2.ExecuteResponse()
    
    164 160
         stream.response.Unpack(execute_response)
    
    165 161
     
    
    166
    -    for output_file_response in execute_response.result.output_files:
    
    167
    -        path = os.path.join(output_directory, output_file_response.path)
    
    168
    -
    
    169
    -        if not os.path.exists(os.path.dirname(path)):
    
    162
    +    with download(context.channel, instance=context.instance_name) as downloader:
    
    170 163
     
    
    171
    -            try:
    
    172
    -                os.makedirs(os.path.dirname(path))
    
    164
    +        for output_file_response in execute_response.result.output_files:
    
    165
    +            path = os.path.join(output_directory, output_file_response.path)
    
    173 166
     
    
    174
    -            except OSError as exc:
    
    175
    -                if exc.errno != errno.EEXIST:
    
    176
    -                    raise
    
    167
    +            if not os.path.exists(os.path.dirname(path)):
    
    168
    +                os.makedirs(os.path.dirname(path), exist_ok=True)
    
    177 169
     
    
    178
    -        with open(path, 'wb+') as f:
    
    179
    -            write_fetch_blob(f, stub, output_file_response.digest, context.instance_name)
    
    170
    +            downloader.download_file(output_file_response.digest, path)
    
    180 171
     
    
    181
    -        if output_file_response.path in output_executeables:
    
    182
    -            st = os.stat(path)
    
    183
    -            os.chmod(path, st.st_mode | stat.S_IXUSR)
    172
    +            if output_file_response.path in output_executeables:
    
    173
    +                st = os.stat(path)
    
    174
    +                os.chmod(path, st.st_mode | stat.S_IXUSR)

  • buildgrid/_exceptions.py
    ... ... @@ -50,3 +50,36 @@ class ServerError(BgdError):
    50 50
     class BotError(BgdError):
    
    51 51
         def __init__(self, message, detail=None, reason=None):
    
    52 52
             super().__init__(message, detail=detail, domain=ErrorDomain.BOT, reason=reason)
    
    53
    +
    
    54
    +
    
    55
    +class InvalidArgumentError(BgdError):
    
    56
    +    """A bad argument was passed, such as a name which doesn't exist."""
    
    57
    +    def __init__(self, message, detail=None, reason=None):
    
    58
    +        super().__init__(message, detail=detail, domain=ErrorDomain.SERVER, reason=reason)
    
    59
    +
    
    60
    +
    
    61
    +class NotFoundError(BgdError):
    
    62
    +    """Requested resource not found."""
    
    63
    +    def __init__(self, message, detail=None, reason=None):
    
    64
    +        super().__init__(message, detail=detail, domain=ErrorDomain.SERVER, reason=reason)
    
    65
    +
    
    66
    +
    
    67
    +class OutOfSyncError(BgdError):
    
    68
    +    """The worker is out of sync with the server, such as having a differing
    
    69
    +    number of leases."""
    
    70
    +    def __init__(self, message, detail=None, reason=None):
    
    71
    +        super().__init__(message, detail=detail, domain=ErrorDomain.SERVER, reason=reason)
    
    72
    +
    
    73
    +
    
    74
    +class OutOfRangeError(BgdError):
    
    75
    +    """ByteStream service read data out of range."""
    
    76
    +    def __init__(self, message, detail=None, reason=None):
    
    77
    +        super().__init__(message, detail=detail, domain=ErrorDomain.SERVER, reason=reason)
    
    78
    +
    
    79
    +
    
    80
    +class FailedPreconditionError(BgdError):
    
    81
    +    """One or more errors occurred in setting up the action requested, such as
    
    82
    +    a missing input or command or no worker being available. The client may be
    
    83
    +    able to fix the errors and retry."""
    
    84
    +    def __init__(self, message, detail=None, reason=None):
    
    85
    +        super().__init__(message, detail=detail, domain=ErrorDomain.SERVER, reason=reason)

  • buildgrid/bot/bot_session.py
    ... ... @@ -154,15 +154,15 @@ class BotSession:
    154 154
                 lease = await loop.run_in_executor(None, self._work, self._context, lease)
    
    155 155
     
    
    156 156
             except grpc.RpcError as e:
    
    157
    -            self.logger.error("Connection error thrown: [{}]".format(e))
    
    158
    -            lease.status.code = e.code()
    
    157
    +            self.logger.error("RPC error thrown: [{}]".format(e))
    
    158
    +            lease.status.CopyFrom(e.code())
    
    159 159
     
    
    160 160
             except BotError as e:
    
    161 161
                 self.logger.error("Internal bot error thrown: [{}]".format(e))
    
    162 162
                 lease.status.code = code_pb2.INTERNAL
    
    163 163
     
    
    164 164
             except Exception as e:
    
    165
    -            self.logger.error("Connection error thrown: [{}]".format(e))
    
    165
    +            self.logger.error("Exception thrown: [{}]".format(e))
    
    166 166
                 lease.status.code = code_pb2.INTERNAL
    
    167 167
     
    
    168 168
             self.logger.debug("Work complete: [{}]".format(lease.id))
    

  • buildgrid/client/cas.py
    ... ... @@ -19,6 +19,7 @@ import os
    19 19
     
    
    20 20
     import grpc
    
    21 21
     
    
    22
    +from buildgrid._exceptions import NotFoundError
    
    22 23
     from buildgrid._protos.build.bazel.remote.execution.v2 import remote_execution_pb2, remote_execution_pb2_grpc
    
    23 24
     from buildgrid._protos.google.bytestream import bytestream_pb2, bytestream_pb2_grpc
    
    24 25
     from buildgrid._protos.google.rpc import code_pb2
    
    ... ... @@ -26,6 +27,16 @@ from buildgrid.settings import HASH
    26 27
     from buildgrid.utils import merkle_tree_maker
    
    27 28
     
    
    28 29
     
    
    30
    +# Maximum size for a queueable file:
    
    31
    +FILE_SIZE_THRESHOLD = 1 * 1024 * 1024
    
    32
    +
    
    33
    +# Maximum size for a single gRPC request:
    
    34
    +MAX_REQUEST_SIZE = 2 * 1024 * 1024
    
    35
    +
    
    36
    +# Maximum number of elements per gRPC request:
    
    37
    +MAX_REQUEST_COUNT = 500
    
    38
    +
    
    39
    +
    
    29 40
     class _CallCache:
    
    30 41
         """Per remote grpc.StatusCode.UNIMPLEMENTED call cache."""
    
    31 42
         __calls = {}
    
    ... ... @@ -43,6 +54,398 @@ class _CallCache:
    43 54
             return name in cls.__calls[channel]
    
    44 55
     
    
    45 56
     
    
    57
    +@contextmanager
    
    58
    +def download(channel, instance=None, u_uid=None):
    
    59
    +    """Context manager generator for the :class:`Downloader` class."""
    
    60
    +    downloader = Downloader(channel, instance=instance)
    
    61
    +    try:
    
    62
    +        yield downloader
    
    63
    +    finally:
    
    64
    +        downloader.close()
    
    65
    +
    
    66
    +
    
    67
    +class Downloader:
    
    68
    +    """Remote CAS files, directories and messages download helper.
    
    69
    +
    
    70
    +    The :class:`Downloader` class comes with a generator factory function that
    
    71
    +    can be used together with the `with` statement for context management::
    
    72
    +
    
    73
    +        from buildgrid.client.cas import download
    
    74
    +
    
    75
    +        with download(channel, instance='build') as downloader:
    
    76
    +            downloader.get_message(message_digest)
    
    77
    +    """
    
    78
    +
    
    79
    +    def __init__(self, channel, instance=None):
    
    80
    +        """Initializes a new :class:`Downloader` instance.
    
    81
    +
    
    82
    +        Args:
    
    83
    +            channel (grpc.Channel): A gRPC channel to the CAS endpoint.
    
    84
    +            instance (str, optional): the targeted instance's name.
    
    85
    +        """
    
    86
    +        self.channel = channel
    
    87
    +
    
    88
    +        self.instance_name = instance
    
    89
    +
    
    90
    +        self.__bytestream_stub = bytestream_pb2_grpc.ByteStreamStub(self.channel)
    
    91
    +        self.__cas_stub = remote_execution_pb2_grpc.ContentAddressableStorageStub(self.channel)
    
    92
    +
    
    93
    +        self.__file_requests = {}
    
    94
    +        self.__file_request_count = 0
    
    95
    +        self.__file_request_size = 0
    
    96
    +        self.__file_response_size = 0
    
    97
    +
    
    98
    +    # --- Public API ---
    
    99
    +
    
    100
    +    def get_blob(self, digest):
    
    101
    +        """Retrieves a blob from the remote CAS server.
    
    102
    +
    
    103
    +        Args:
    
    104
    +            digest (:obj:`Digest`): the blob's digest to fetch.
    
    105
    +
    
    106
    +        Returns:
    
    107
    +            bytearray: the fetched blob data or None if not found.
    
    108
    +        """
    
    109
    +        try:
    
    110
    +            blob = self._fetch_blob(digest)
    
    111
    +        except NotFoundError:
    
    112
    +            return None
    
    113
    +
    
    114
    +        return blob
    
    115
    +
    
    116
    +    def get_blobs(self, digests):
    
    117
    +        """Retrieves a list of blobs from the remote CAS server.
    
    118
    +
    
    119
    +        Args:
    
    120
    +            digests (list): list of :obj:`Digest`s for the blobs to fetch.
    
    121
    +
    
    122
    +        Returns:
    
    123
    +            list: the fetched blob data list.
    
    124
    +        """
    
    125
    +        return self._fetch_blob_batch(digests)
    
    126
    +
    
    127
    +    def get_message(self, digest, message):
    
    128
    +        """Retrieves a :obj:`Message` from the remote CAS server.
    
    129
    +
    
    130
    +        Args:
    
    131
    +            digest (:obj:`Digest`): the message's digest to fetch.
    
    132
    +            message (:obj:`Message`): an empty message to fill.
    
    133
    +
    
    134
    +        Returns:
    
    135
    +            :obj:`Message`: `message` filled or emptied if not found.
    
    136
    +        """
    
    137
    +        try:
    
    138
    +            message_blob = self._fetch_blob(digest)
    
    139
    +        except NotFoundError:
    
    140
    +            message_blob = None
    
    141
    +
    
    142
    +        if message_blob is not None:
    
    143
    +            message.ParseFromString(message_blob)
    
    144
    +        else:
    
    145
    +            message.Clear()
    
    146
    +
    
    147
    +        return message
    
    148
    +
    
    149
    +    def get_messages(self, digests, messages):
    
    150
    +        """Retrieves a list of :obj:`Message`s from the remote CAS server.
    
    151
    +
    
    152
    +        Note:
    
    153
    +            The `digests` and `messages` list **must** contain the same number
    
    154
    +            of elements.
    
    155
    +
    
    156
    +        Args:
    
    157
    +            digests (list):  list of :obj:`Digest`s for the messages to fetch.
    
    158
    +            messages (list): list of empty :obj:`Message`s to fill.
    
    159
    +
    
    160
    +        Returns:
    
    161
    +            list: the fetched and filled message list.
    
    162
    +        """
    
    163
    +        assert len(digests) == len(messages)
    
    164
    +
    
    165
    +        message_blobs = self._fetch_blob_batch(digests)
    
    166
    +
    
    167
    +        assert len(message_blobs) == len(messages)
    
    168
    +
    
    169
    +        for message, message_blob in zip(messages, message_blobs):
    
    170
    +            message.ParseFromString(message_blob)
    
    171
    +
    
    172
    +        return messages
    
    173
    +
    
    174
    +    def download_file(self, digest, file_path, queue=True):
    
    175
    +        """Retrieves a file from the remote CAS server.
    
    176
    +
    
    177
    +        If queuing is allowed (`queue=True`), the download request **may** be
    
    178
    +        defer. An explicit call to :func:`~flush` can force the request to be
    
    179
    +        send immediately (along with the rest of the queued batch).
    
    180
    +
    
    181
    +        Args:
    
    182
    +            digest (:obj:`Digest`): the file's digest to fetch.
    
    183
    +            file_path (str): absolute or relative path to the local file to write.
    
    184
    +            queue (bool, optional): whether or not the download request may be
    
    185
    +                queued and submitted as part of a batch upload request. Defaults
    
    186
    +                to True.
    
    187
    +
    
    188
    +        Raises:
    
    189
    +            NotFoundError: if `digest` is not present in the remote CAS server.
    
    190
    +            OSError: if `file_path` does not exist or is not readable.
    
    191
    +        """
    
    192
    +        if not os.path.isabs(file_path):
    
    193
    +            file_path = os.path.abspath(file_path)
    
    194
    +
    
    195
    +        if not queue or digest.size_bytes > FILE_SIZE_THRESHOLD:
    
    196
    +            self._fetch_file(digest, file_path)
    
    197
    +        else:
    
    198
    +            self._queue_file(digest, file_path)
    
    199
    +
    
    200
    +    def download_directory(self, digest, directory_path):
    
    201
    +        """Retrieves a :obj:`Directory` from the remote CAS server.
    
    202
    +
    
    203
    +        Args:
    
    204
    +            digest (:obj:`Digest`): the directory's digest to fetch.
    
    205
    +
    
    206
    +        Raises:
    
    207
    +            NotFoundError: if `digest` is not present in the remote CAS server.
    
    208
    +            FileExistsError: if `directory_path` already contains parts of their
    
    209
    +                fetched directory's content.
    
    210
    +        """
    
    211
    +        if not os.path.isabs(directory_path):
    
    212
    +            directory_path = os.path.abspath(directory_path)
    
    213
    +
    
    214
    +        self._fetch_directory(digest, directory_path)
    
    215
    +
    
    216
    +    def flush(self):
    
    217
    +        """Ensures any queued request gets sent."""
    
    218
    +        if self.__file_requests:
    
    219
    +            self._fetch_file_batch(self.__file_requests)
    
    220
    +
    
    221
    +            self.__file_requests.clear()
    
    222
    +            self.__file_request_count = 0
    
    223
    +            self.__file_request_size = 0
    
    224
    +            self.__file_response_size = 0
    
    225
    +
    
    226
    +    def close(self):
    
    227
    +        """Closes the underlying connection stubs.
    
    228
    +
    
    229
    +        Note:
    
    230
    +            This will always send pending requests before closing connections,
    
    231
    +            if any.
    
    232
    +        """
    
    233
    +        self.flush()
    
    234
    +
    
    235
    +        self.__bytestream_stub = None
    
    236
    +        self.__cas_stub = None
    
    237
    +
    
    238
    +    # --- Private API ---
    
    239
    +
    
    240
    +    def _fetch_blob(self, digest):
    
    241
    +        """Fetches a blob using ByteStream.Read()"""
    
    242
    +        read_blob = bytearray()
    
    243
    +
    
    244
    +        if self.instance_name is not None:
    
    245
    +            resource_name = '/'.join([self.instance_name, 'blobs',
    
    246
    +                                      digest.hash, str(digest.size_bytes)])
    
    247
    +        else:
    
    248
    +            resource_name = '/'.join(['blobs', digest.hash, str(digest.size_bytes)])
    
    249
    +
    
    250
    +        read_request = bytestream_pb2.ReadRequest()
    
    251
    +        read_request.resource_name = resource_name
    
    252
    +        read_request.read_offset = 0
    
    253
    +
    
    254
    +        try:
    
    255
    +            # TODO: Handle connection loss/recovery
    
    256
    +            for read_response in self.__bytestream_stub.Read(read_request):
    
    257
    +                read_blob += read_response.data
    
    258
    +
    
    259
    +            assert len(read_blob) == digest.size_bytes
    
    260
    +
    
    261
    +        except grpc.RpcError as e:
    
    262
    +            status_code = e.code()
    
    263
    +            if status_code == grpc.StatusCode.NOT_FOUND:
    
    264
    +                raise NotFoundError("Requested data does not exist on the remote.")
    
    265
    +
    
    266
    +            else:
    
    267
    +                assert False
    
    268
    +
    
    269
    +        return read_blob
    
    270
    +
    
    271
    +    def _fetch_blob_batch(self, digests):
    
    272
    +        """Fetches blobs using ContentAddressableStorage.BatchReadBlobs()"""
    
    273
    +        batch_fetched = False
    
    274
    +        read_blobs = []
    
    275
    +
    
    276
    +        # First, try BatchReadBlobs(), if not already known not being implemented:
    
    277
    +        if not _CallCache.unimplemented(self.channel, 'BatchReadBlobs'):
    
    278
    +            batch_request = remote_execution_pb2.BatchReadBlobsRequest()
    
    279
    +            batch_request.digests.extend(digests)
    
    280
    +            if self.instance_name is not None:
    
    281
    +                batch_request.instance_name = self.instance_name
    
    282
    +
    
    283
    +            try:
    
    284
    +                batch_response = self.__cas_stub.BatchReadBlobs(batch_request)
    
    285
    +                for response in batch_response.responses:
    
    286
    +                    assert response.digest.hash in digests
    
    287
    +
    
    288
    +                    read_blobs.append(response.data)
    
    289
    +
    
    290
    +                    if response.status.code != code_pb2.OK:
    
    291
    +                        assert False
    
    292
    +
    
    293
    +                batch_fetched = True
    
    294
    +
    
    295
    +            except grpc.RpcError as e:
    
    296
    +                status_code = e.code()
    
    297
    +                if status_code == grpc.StatusCode.UNIMPLEMENTED:
    
    298
    +                    _CallCache.mark_unimplemented(self.channel, 'BatchReadBlobs')
    
    299
    +
    
    300
    +                elif status_code == grpc.StatusCode.INVALID_ARGUMENT:
    
    301
    +                    read_blobs.clear()
    
    302
    +                    batch_fetched = False
    
    303
    +
    
    304
    +                else:
    
    305
    +                    assert False
    
    306
    +
    
    307
    +        # Fallback to Read() if no BatchReadBlobs():
    
    308
    +        if not batch_fetched:
    
    309
    +            for digest in digests:
    
    310
    +                read_blobs.append(self._fetch_blob(digest))
    
    311
    +
    
    312
    +        return read_blobs
    
    313
    +
    
    314
    +    def _fetch_file(self, digest, file_path):
    
    315
    +        """Fetches a file using ByteStream.Read()"""
    
    316
    +        if self.instance_name is not None:
    
    317
    +            resource_name = '/'.join([self.instance_name, 'blobs',
    
    318
    +                                      digest.hash, str(digest.size_bytes)])
    
    319
    +        else:
    
    320
    +            resource_name = '/'.join(['blobs', digest.hash, str(digest.size_bytes)])
    
    321
    +
    
    322
    +        read_request = bytestream_pb2.ReadRequest()
    
    323
    +        read_request.resource_name = resource_name
    
    324
    +        read_request.read_offset = 0
    
    325
    +
    
    326
    +        os.makedirs(os.path.dirname(file_path), exist_ok=True)
    
    327
    +
    
    328
    +        with open(file_path, 'wb') as byte_file:
    
    329
    +            # TODO: Handle connection loss/recovery
    
    330
    +            for read_response in self.__bytestream_stub.Read(read_request):
    
    331
    +                byte_file.write(read_response.data)
    
    332
    +
    
    333
    +            assert byte_file.tell() == digest.size_bytes
    
    334
    +
    
    335
    +    def _queue_file(self, digest, file_path):
    
    336
    +        """Queues a file for later batch download"""
    
    337
    +        if self.__file_request_size + digest.ByteSize() > MAX_REQUEST_SIZE:
    
    338
    +            self.flush()
    
    339
    +        elif self.__file_response_size + digest.size_bytes > MAX_REQUEST_SIZE:
    
    340
    +            self.flush()
    
    341
    +        elif self.__file_request_count >= MAX_REQUEST_COUNT:
    
    342
    +            self.flush()
    
    343
    +
    
    344
    +        self.__file_requests[digest.hash] = (digest, file_path)
    
    345
    +        self.__file_request_count += 1
    
    346
    +        self.__file_request_size += digest.ByteSize()
    
    347
    +        self.__file_response_size += digest.size_bytes
    
    348
    +
    
    349
    +    def _fetch_file_batch(self, batch):
    
    350
    +        """Sends queued data using ContentAddressableStorage.BatchReadBlobs()"""
    
    351
    +        batch_digests = [digest for digest, _ in batch.values()]
    
    352
    +        batch_blobs = self._fetch_blob_batch(batch_digests)
    
    353
    +
    
    354
    +        for (_, file_path), file_blob in zip(batch.values(), batch_blobs):
    
    355
    +            os.makedirs(os.path.dirname(file_path), exist_ok=True)
    
    356
    +
    
    357
    +            with open(file_path, 'wb') as byte_file:
    
    358
    +                byte_file.write(file_blob)
    
    359
    +
    
    360
    +    def _fetch_directory(self, digest, directory_path):
    
    361
    +        """Fetches a file using ByteStream.GetTree()"""
    
    362
    +        # Better fail early if the local root path cannot be created:
    
    363
    +        os.makedirs(directory_path, exist_ok=True)
    
    364
    +
    
    365
    +        directories = {}
    
    366
    +        directory_fetched = False
    
    367
    +        # First, try GetTree() if not known to be unimplemented yet:
    
    368
    +        if not _CallCache.unimplemented(self.channel, 'GetTree'):
    
    369
    +            tree_request = remote_execution_pb2.GetTreeRequest()
    
    370
    +            tree_request.root_digest.CopyFrom(digest)
    
    371
    +            tree_request.page_size = MAX_REQUEST_COUNT
    
    372
    +            if self.instance_name is not None:
    
    373
    +                tree_request.instance_name = self.instance_name
    
    374
    +
    
    375
    +            try:
    
    376
    +                for tree_response in self.__cas_stub.GetTree(tree_request):
    
    377
    +                    for directory in tree_response.directories:
    
    378
    +                        directory_blob = directory.SerializeToString()
    
    379
    +                        directory_hash = HASH(directory_blob).hexdigest()
    
    380
    +
    
    381
    +                        directories[directory_hash] = directory
    
    382
    +
    
    383
    +                assert digest.hash in directories
    
    384
    +
    
    385
    +                directory = directories[digest.hash]
    
    386
    +                self._write_directory(digest.hash, directory_path,
    
    387
    +                                      directories=directories, root_barrier=directory_path)
    
    388
    +
    
    389
    +                directory_fetched = True
    
    390
    +
    
    391
    +            except grpc.RpcError as e:
    
    392
    +                status_code = e.code()
    
    393
    +                if status_code == grpc.StatusCode.UNIMPLEMENTED:
    
    394
    +                    _CallCache.mark_unimplemented(self.channel, 'BatchUpdateBlobs')
    
    395
    +
    
    396
    +                elif status_code == grpc.StatusCode.NOT_FOUND:
    
    397
    +                    raise NotFoundError("Requested directory does not exist on the remote.")
    
    398
    +
    
    399
    +                else:
    
    400
    +                    assert False
    
    401
    +
    
    402
    +        # TODO: Try with BatchReadBlobs().
    
    403
    +
    
    404
    +        # Fallback to Read() if no GetTree():
    
    405
    +        if not directory_fetched:
    
    406
    +            directory = remote_execution_pb2.Directory()
    
    407
    +            directory.ParseFromString(self._fetch_blob(digest))
    
    408
    +
    
    409
    +            self._write_directory(directory, directory_path,
    
    410
    +                                  root_barrier=directory_path)
    
    411
    +
    
    412
    +    def _write_directory(self, root_directory, root_path, directories=None, root_barrier=None):
    
    413
    +        """Generates a local directory structure"""
    
    414
    +        for file_node in root_directory.files:
    
    415
    +            file_path = os.path.join(root_path, file_node.name)
    
    416
    +
    
    417
    +            self._queue_file(file_node.digest, file_path)
    
    418
    +
    
    419
    +        for directory_node in root_directory.directories:
    
    420
    +            directory_path = os.path.join(root_path, directory_node.name)
    
    421
    +            if directories and directory_node.digest.hash in directories:
    
    422
    +                directory = directories[directory_node.digest.hash]
    
    423
    +            else:
    
    424
    +                directory = remote_execution_pb2.Directory()
    
    425
    +                directory.ParseFromString(self._fetch_blob(directory_node.digest))
    
    426
    +
    
    427
    +            os.makedirs(directory_path, exist_ok=True)
    
    428
    +
    
    429
    +            self._write_directory(directory, directory_path,
    
    430
    +                                  directories=directories, root_barrier=root_barrier)
    
    431
    +
    
    432
    +        for symlink_node in root_directory.symlinks:
    
    433
    +            symlink_path = os.path.join(root_path, symlink_node.name)
    
    434
    +            if not os.path.isabs(symlink_node.target):
    
    435
    +                target_path = os.path.join(root_path, symlink_node.target)
    
    436
    +            else:
    
    437
    +                target_path = symlink_node.target
    
    438
    +            target_path = os.path.normpath(target_path)
    
    439
    +
    
    440
    +            # Do not create links pointing outside the barrier:
    
    441
    +            if root_barrier is not None:
    
    442
    +                common_path = os.path.commonprefix([root_barrier, target_path])
    
    443
    +                if not common_path.startswith(root_barrier):
    
    444
    +                    continue
    
    445
    +
    
    446
    +            os.symlink(symlink_path, target_path)
    
    447
    +
    
    448
    +
    
    46 449
     @contextmanager
    
    47 450
     def upload(channel, instance=None, u_uid=None):
    
    48 451
         """Context manager generator for the :class:`Uploader` class."""
    
    ... ... @@ -63,16 +466,8 @@ class Uploader:
    63 466
     
    
    64 467
             with upload(channel, instance='build') as uploader:
    
    65 468
                 uploader.upload_file('/path/to/local/file')
    
    66
    -
    
    67
    -    Attributes:
    
    68
    -        FILE_SIZE_THRESHOLD (int): maximum size for a queueable file.
    
    69
    -        MAX_REQUEST_SIZE (int): maximum size for a single gRPC request.
    
    70 469
         """
    
    71 470
     
    
    72
    -    FILE_SIZE_THRESHOLD = 1 * 1024 * 1024
    
    73
    -    MAX_REQUEST_SIZE = 2 * 1024 * 1024
    
    74
    -    MAX_REQUEST_COUNT = 500
    
    75
    -
    
    76 471
         def __init__(self, channel, instance=None, u_uid=None):
    
    77 472
             """Initializes a new :class:`Uploader` instance.
    
    78 473
     
    
    ... ... @@ -115,7 +510,7 @@ class Uploader:
    115 510
             Returns:
    
    116 511
                 :obj:`Digest`: the sent blob's digest.
    
    117 512
             """
    
    118
    -        if not queue or len(blob) > Uploader.FILE_SIZE_THRESHOLD:
    
    513
    +        if not queue or len(blob) > FILE_SIZE_THRESHOLD:
    
    119 514
                 blob_digest = self._send_blob(blob, digest=digest)
    
    120 515
             else:
    
    121 516
                 blob_digest = self._queue_blob(blob, digest=digest)
    
    ... ... @@ -141,7 +536,7 @@ class Uploader:
    141 536
             """
    
    142 537
             message_blob = message.SerializeToString()
    
    143 538
     
    
    144
    -        if not queue or len(message_blob) > Uploader.FILE_SIZE_THRESHOLD:
    
    539
    +        if not queue or len(message_blob) > FILE_SIZE_THRESHOLD:
    
    145 540
                 message_digest = self._send_blob(message_blob, digest=digest)
    
    146 541
             else:
    
    147 542
                 message_digest = self._queue_blob(message_blob, digest=digest)
    
    ... ... @@ -174,7 +569,7 @@ class Uploader:
    174 569
             with open(file_path, 'rb') as bytes_steam:
    
    175 570
                 file_bytes = bytes_steam.read()
    
    176 571
     
    
    177
    -        if not queue or len(file_bytes) > Uploader.FILE_SIZE_THRESHOLD:
    
    572
    +        if not queue or len(file_bytes) > FILE_SIZE_THRESHOLD:
    
    178 573
                 file_digest = self._send_blob(file_bytes)
    
    179 574
             else:
    
    180 575
                 file_digest = self._queue_blob(file_bytes)
    
    ... ... @@ -316,7 +711,7 @@ class Uploader:
    316 711
                 finished = False
    
    317 712
                 remaining = len(content)
    
    318 713
                 while not finished:
    
    319
    -                chunk_size = min(remaining, Uploader.MAX_REQUEST_SIZE)
    
    714
    +                chunk_size = min(remaining, MAX_REQUEST_SIZE)
    
    320 715
                     remaining -= chunk_size
    
    321 716
     
    
    322 717
                     request = bytestream_pb2.WriteRequest()
    
    ... ... @@ -347,9 +742,9 @@ class Uploader:
    347 742
                 blob_digest.hash = HASH(blob).hexdigest()
    
    348 743
                 blob_digest.size_bytes = len(blob)
    
    349 744
     
    
    350
    -        if self.__request_size + blob_digest.size_bytes > Uploader.MAX_REQUEST_SIZE:
    
    745
    +        if self.__request_size + blob_digest.size_bytes > MAX_REQUEST_SIZE:
    
    351 746
                 self.flush()
    
    352
    -        elif self.__request_count >= Uploader.MAX_REQUEST_COUNT:
    
    747
    +        elif self.__request_count >= MAX_REQUEST_COUNT:
    
    353 748
                 self.flush()
    
    354 749
     
    
    355 750
             self.__requests[blob_digest.hash] = (blob, blob_digest)
    

  • buildgrid/server/_exceptions.py deleted
    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 .._exceptions import BgdError, ErrorDomain
    
    17
    -
    
    18
    -
    
    19
    -class InvalidArgumentError(BgdError):
    
    20
    -    """A bad argument was passed, such as a name which doesn't exist.
    
    21
    -    """
    
    22
    -
    
    23
    -    def __init__(self, message, detail=None, reason=None):
    
    24
    -        super().__init__(message, detail=detail, domain=ErrorDomain.SERVER, reason=reason)
    
    25
    -
    
    26
    -
    
    27
    -class NotFoundError(BgdError):
    
    28
    -    """Requested resource not found.
    
    29
    -    """
    
    30
    -
    
    31
    -    def __init__(self, message, detail=None, reason=None):
    
    32
    -        super().__init__(message, detail=detail, domain=ErrorDomain.SERVER, reason=reason)
    
    33
    -
    
    34
    -
    
    35
    -class OutofSyncError(BgdError):
    
    36
    -    """The worker is out of sync with the server, such as having a differing number of leases.
    
    37
    -    """
    
    38
    -
    
    39
    -    def __init__(self, message, detail=None, reason=None):
    
    40
    -        super().__init__(message, detail=detail, domain=ErrorDomain.SERVER, reason=reason)
    
    41
    -
    
    42
    -
    
    43
    -class OutOfRangeError(BgdError):
    
    44
    -    """ ByteStream service read data out of range
    
    45
    -    """
    
    46
    -
    
    47
    -    def __init__(self, message, detail=None, reason=None):
    
    48
    -        super().__init__(message, detail=detail, domain=ErrorDomain.SERVER, reason=reason)
    
    49
    -
    
    50
    -
    
    51
    -class FailedPreconditionError(BgdError):
    
    52
    -    """ One or more errors occurred in setting up the action requested, such as a missing input
    
    53
    -    or command or no worker being available. The client may be able to fix the errors and retry.
    
    54
    -    """
    
    55
    -
    
    56
    -    def __init__(self, message, detail=None, reason=None):
    
    57
    -        super().__init__(message, detail=detail, domain=ErrorDomain.SERVER, reason=reason)

  • buildgrid/server/actioncache/service.py
    ... ... @@ -24,11 +24,10 @@ import logging
    24 24
     
    
    25 25
     import grpc
    
    26 26
     
    
    27
    +from buildgrid._exceptions import InvalidArgumentError, NotFoundError
    
    27 28
     from buildgrid._protos.build.bazel.remote.execution.v2 import remote_execution_pb2
    
    28 29
     from buildgrid._protos.build.bazel.remote.execution.v2 import remote_execution_pb2_grpc
    
    29 30
     
    
    30
    -from .._exceptions import InvalidArgumentError, NotFoundError
    
    31
    -
    
    32 31
     
    
    33 32
     class ActionCacheService(remote_execution_pb2_grpc.ActionCacheServicer):
    
    34 33
     
    

  • buildgrid/server/bots/instance.py
    ... ... @@ -23,7 +23,8 @@ Instance of the Remote Workers interface.
    23 23
     import logging
    
    24 24
     import uuid
    
    25 25
     
    
    26
    -from .._exceptions import InvalidArgumentError, OutofSyncError
    
    26
    +from buildgrid._exceptions import InvalidArgumentError, OutOfSyncError
    
    27
    +
    
    27 28
     from ..job import LeaseState
    
    28 29
     
    
    29 30
     
    
    ... ... @@ -108,7 +109,7 @@ class BotsInterface:
    108 109
                     # TODO: Lease was rejected
    
    109 110
                     raise NotImplementedError("'Not Accepted' is unsupported")
    
    110 111
                 else:
    
    111
    -                raise OutofSyncError("Server lease: [{}]. Client lease: [{}]".format(server_lease, client_lease))
    
    112
    +                raise OutOfSyncError("Server lease: [{}]. Client lease: [{}]".format(server_lease, client_lease))
    
    112 113
     
    
    113 114
             elif server_state == LeaseState.ACTIVE:
    
    114 115
     
    
    ... ... @@ -121,17 +122,17 @@ class BotsInterface:
    121 122
                     return None
    
    122 123
     
    
    123 124
                 else:
    
    124
    -                raise OutofSyncError("Server lease: [{}]. Client lease: [{}]".format(server_lease, client_lease))
    
    125
    +                raise OutOfSyncError("Server lease: [{}]. Client lease: [{}]".format(server_lease, client_lease))
    
    125 126
     
    
    126 127
             elif server_state == LeaseState.COMPLETED:
    
    127
    -            raise OutofSyncError("Server lease: [{}]. Client lease: [{}]".format(server_lease, client_lease))
    
    128
    +            raise OutOfSyncError("Server lease: [{}]. Client lease: [{}]".format(server_lease, client_lease))
    
    128 129
     
    
    129 130
             elif server_state == LeaseState.CANCELLED:
    
    130 131
                 raise NotImplementedError("Cancelled states not supported yet")
    
    131 132
     
    
    132 133
             else:
    
    133 134
                 # Sould never get here
    
    134
    -            raise OutofSyncError("State now allowed: {}".format(server_state))
    
    135
    +            raise OutOfSyncError("State now allowed: {}".format(server_state))
    
    135 136
     
    
    136 137
             return client_lease
    
    137 138
     
    

  • buildgrid/server/bots/service.py
    ... ... @@ -25,11 +25,10 @@ import grpc
    25 25
     
    
    26 26
     from google.protobuf.empty_pb2 import Empty
    
    27 27
     
    
    28
    +from buildgrid._exceptions import InvalidArgumentError, OutOfSyncError
    
    28 29
     from buildgrid._protos.google.devtools.remoteworkers.v1test2 import bots_pb2
    
    29 30
     from buildgrid._protos.google.devtools.remoteworkers.v1test2 import bots_pb2_grpc
    
    30 31
     
    
    31
    -from .._exceptions import InvalidArgumentError, OutofSyncError
    
    32
    -
    
    33 32
     
    
    34 33
     class BotsService(bots_pb2_grpc.BotsServicer):
    
    35 34
     
    
    ... ... @@ -73,7 +72,7 @@ class BotsService(bots_pb2_grpc.BotsServicer):
    73 72
                 context.set_details(str(e))
    
    74 73
                 context.set_code(grpc.StatusCode.INVALID_ARGUMENT)
    
    75 74
     
    
    76
    -        except OutofSyncError as e:
    
    75
    +        except OutOfSyncError as e:
    
    77 76
                 self.logger.error(e)
    
    78 77
                 context.set_details(str(e))
    
    79 78
                 context.set_code(grpc.StatusCode.DATA_LOSS)
    

  • buildgrid/server/cas/instance.py
    ... ... @@ -19,11 +19,10 @@ Storage Instances
    19 19
     Instances of CAS and ByteStream
    
    20 20
     """
    
    21 21
     
    
    22
    +from buildgrid._exceptions import InvalidArgumentError, NotFoundError, OutOfRangeError
    
    22 23
     from buildgrid._protos.google.bytestream import bytestream_pb2
    
    23 24
     from buildgrid._protos.build.bazel.remote.execution.v2 import remote_execution_pb2 as re_pb2
    
    24
    -
    
    25
    -from .._exceptions import InvalidArgumentError, NotFoundError, OutOfRangeError
    
    26
    -from ...settings import HASH
    
    25
    +from buildgrid.settings import HASH
    
    27 26
     
    
    28 27
     
    
    29 28
     class ContentAddressableStorageInstance:
    

  • buildgrid/server/cas/service.py
    ... ... @@ -26,12 +26,11 @@ import logging
    26 26
     
    
    27 27
     import grpc
    
    28 28
     
    
    29
    +from buildgrid._exceptions import InvalidArgumentError, NotFoundError, OutOfRangeError
    
    29 30
     from buildgrid._protos.google.bytestream import bytestream_pb2, bytestream_pb2_grpc
    
    30 31
     from buildgrid._protos.build.bazel.remote.execution.v2 import remote_execution_pb2
    
    31 32
     from buildgrid._protos.build.bazel.remote.execution.v2 import remote_execution_pb2_grpc
    
    32 33
     
    
    33
    -from .._exceptions import InvalidArgumentError, NotFoundError, OutOfRangeError
    
    34
    -
    
    35 34
     
    
    36 35
     class ContentAddressableStorageService(remote_execution_pb2_grpc.ContentAddressableStorageServicer):
    
    37 36
     
    

  • buildgrid/server/cas/storage/remote.py
    ... ... @@ -23,14 +23,10 @@ Forwwards storage requests to a remote storage.
    23 23
     import io
    
    24 24
     import logging
    
    25 25
     
    
    26
    -import grpc
    
    27
    -
    
    28
    -from buildgrid.client.cas import upload
    
    29
    -from buildgrid._protos.google.bytestream import bytestream_pb2_grpc
    
    26
    +from buildgrid.client.cas import download, upload
    
    30 27
     from buildgrid._protos.build.bazel.remote.execution.v2 import remote_execution_pb2, remote_execution_pb2_grpc
    
    31 28
     from buildgrid._protos.google.rpc import code_pb2
    
    32 29
     from buildgrid._protos.google.rpc import status_pb2
    
    33
    -from buildgrid.utils import gen_fetch_blob
    
    34 30
     from buildgrid.settings import HASH
    
    35 31
     
    
    36 32
     from .storage_abc import StorageABC
    
    ... ... @@ -44,7 +40,6 @@ class RemoteStorage(StorageABC):
    44 40
             self.instance_name = instance_name
    
    45 41
             self.channel = channel
    
    46 42
     
    
    47
    -        self._stub_bs = bytestream_pb2_grpc.ByteStreamStub(channel)
    
    48 43
             self._stub_cas = remote_execution_pb2_grpc.ContentAddressableStorageStub(channel)
    
    49 44
     
    
    50 45
         def has_blob(self, digest):
    
    ... ... @@ -53,25 +48,12 @@ class RemoteStorage(StorageABC):
    53 48
             return False
    
    54 49
     
    
    55 50
         def get_blob(self, digest):
    
    56
    -        try:
    
    57
    -            fetched_data = io.BytesIO()
    
    58
    -            length = 0
    
    59
    -
    
    60
    -            for data in gen_fetch_blob(self._stub_bs, digest, self.instance_name):
    
    61
    -                length += fetched_data.write(data)
    
    62
    -
    
    63
    -            assert digest.size_bytes == length
    
    64
    -            fetched_data.seek(0)
    
    65
    -            return fetched_data
    
    66
    -
    
    67
    -        except grpc.RpcError as e:
    
    68
    -            if e.code() == grpc.StatusCode.NOT_FOUND:
    
    69
    -                pass
    
    51
    +        with download(self.channel, instance=self.instance_name) as downloader:
    
    52
    +            blob = downloader.get_blob(digest)
    
    53
    +            if blob is not None:
    
    54
    +                return io.BytesIO(blob)
    
    70 55
                 else:
    
    71
    -                self.logger.error(e.details())
    
    72
    -                raise
    
    73
    -
    
    74
    -        return None
    
    56
    +                return None
    
    75 57
     
    
    76 58
         def begin_write(self, digest):
    
    77 59
             return io.BytesIO()
    

  • buildgrid/server/execution/instance.py
    ... ... @@ -21,10 +21,10 @@ An instance of the Remote Execution Service.
    21 21
     
    
    22 22
     import logging
    
    23 23
     
    
    24
    +from buildgrid._exceptions import FailedPreconditionError, InvalidArgumentError
    
    24 25
     from buildgrid._protos.build.bazel.remote.execution.v2.remote_execution_pb2 import Action
    
    25 26
     
    
    26 27
     from ..job import Job
    
    27
    -from .._exceptions import InvalidArgumentError, FailedPreconditionError
    
    28 28
     
    
    29 29
     
    
    30 30
     class ExecutionInstance:
    

  • buildgrid/server/execution/service.py
    ... ... @@ -26,12 +26,10 @@ from functools import partial
    26 26
     
    
    27 27
     import grpc
    
    28 28
     
    
    29
    +from buildgrid._exceptions import FailedPreconditionError, InvalidArgumentError
    
    29 30
     from buildgrid._protos.build.bazel.remote.execution.v2 import remote_execution_pb2_grpc
    
    30
    -
    
    31 31
     from buildgrid._protos.google.longrunning import operations_pb2
    
    32 32
     
    
    33
    -from .._exceptions import InvalidArgumentError, FailedPreconditionError
    
    34
    -
    
    35 33
     
    
    36 34
     class ExecutionService(remote_execution_pb2_grpc.ExecutionServicer):
    
    37 35
     
    

  • buildgrid/server/operations/instance.py
    ... ... @@ -21,7 +21,7 @@ An instance of the LongRunningOperations Service.
    21 21
     
    
    22 22
     import logging
    
    23 23
     
    
    24
    -from .._exceptions import InvalidArgumentError
    
    24
    +from buildgrid._exceptions import InvalidArgumentError
    
    25 25
     
    
    26 26
     
    
    27 27
     class OperationsInstance:
    

  • buildgrid/server/operations/service.py
    ... ... @@ -25,10 +25,9 @@ import grpc
    25 25
     
    
    26 26
     from google.protobuf.empty_pb2 import Empty
    
    27 27
     
    
    28
    +from buildgrid._exceptions import InvalidArgumentError
    
    28 29
     from buildgrid._protos.google.longrunning import operations_pb2_grpc, operations_pb2
    
    29 30
     
    
    30
    -from .._exceptions import InvalidArgumentError
    
    31
    -
    
    32 31
     
    
    33 32
     class OperationsService(operations_pb2_grpc.OperationsServicer):
    
    34 33
     
    

  • buildgrid/server/referencestorage/service.py
    ... ... @@ -17,11 +17,10 @@ import logging
    17 17
     
    
    18 18
     import grpc
    
    19 19
     
    
    20
    +from buildgrid._exceptions import InvalidArgumentError, NotFoundError
    
    20 21
     from buildgrid._protos.buildstream.v2 import buildstream_pb2
    
    21 22
     from buildgrid._protos.buildstream.v2 import buildstream_pb2_grpc
    
    22 23
     
    
    23
    -from .._exceptions import InvalidArgumentError, NotFoundError
    
    24
    -
    
    25 24
     
    
    26 25
     class ReferenceStorageService(buildstream_pb2_grpc.ReferenceStorageServicer):
    
    27 26
     
    

  • buildgrid/server/referencestorage/storage.py
    ... ... @@ -24,10 +24,9 @@ For a given key, it
    24 24
     
    
    25 25
     import collections
    
    26 26
     
    
    27
    +from buildgrid._exceptions import NotFoundError
    
    27 28
     from buildgrid._protos.build.bazel.remote.execution.v2 import remote_execution_pb2
    
    28 29
     
    
    29
    -from .._exceptions import NotFoundError
    
    30
    -
    
    31 30
     
    
    32 31
     class ReferenceCache:
    
    33 32
     
    

  • buildgrid/server/scheduler.py
    ... ... @@ -25,8 +25,7 @@ from collections import deque
    25 25
     
    
    26 26
     from google.protobuf import any_pb2
    
    27 27
     
    
    28
    -
    
    29
    -from buildgrid.server._exceptions import NotFoundError
    
    28
    +from buildgrid._exceptions import NotFoundError
    
    30 29
     from buildgrid._protos.build.bazel.remote.execution.v2 import remote_execution_pb2
    
    31 30
     from buildgrid._protos.google.longrunning import operations_pb2
    
    32 31
     
    

  • buildgrid/utils.py
    ... ... @@ -18,87 +18,6 @@ import os
    18 18
     
    
    19 19
     from buildgrid.settings import HASH
    
    20 20
     from buildgrid._protos.build.bazel.remote.execution.v2 import remote_execution_pb2
    
    21
    -from buildgrid._protos.google.bytestream import bytestream_pb2
    
    22
    -
    
    23
    -
    
    24
    -def gen_fetch_blob(stub, digest, instance_name=""):
    
    25
    -    """ Generates byte stream from a fetch blob request
    
    26
    -    """
    
    27
    -
    
    28
    -    resource_name = os.path.join(instance_name, 'blobs', digest.hash, str(digest.size_bytes))
    
    29
    -    request = bytestream_pb2.ReadRequest(resource_name=resource_name,
    
    30
    -                                         read_offset=0)
    
    31
    -
    
    32
    -    for response in stub.Read(request):
    
    33
    -        yield response.data
    
    34
    -
    
    35
    -
    
    36
    -def write_fetch_directory(root_directory, stub, digest, instance_name=None):
    
    37
    -    """Locally replicates a directory from CAS.
    
    38
    -
    
    39
    -    Args:
    
    40
    -        root_directory (str): local directory to populate.
    
    41
    -        stub (): gRPC stub for CAS communication.
    
    42
    -        digest (Digest): digest for the directory to fetch from CAS.
    
    43
    -        instance_name (str, optional): farm instance name to query data from.
    
    44
    -    """
    
    45
    -
    
    46
    -    if not os.path.isabs(root_directory):
    
    47
    -        root_directory = os.path.abspath(root_directory)
    
    48
    -    if not os.path.exists(root_directory):
    
    49
    -        os.makedirs(root_directory, exist_ok=True)
    
    50
    -
    
    51
    -    directory = parse_to_pb2_from_fetch(remote_execution_pb2.Directory(),
    
    52
    -                                        stub, digest, instance_name)
    
    53
    -
    
    54
    -    for directory_node in directory.directories:
    
    55
    -        child_path = os.path.join(root_directory, directory_node.name)
    
    56
    -
    
    57
    -        write_fetch_directory(child_path, stub, directory_node.digest, instance_name)
    
    58
    -
    
    59
    -    for file_node in directory.files:
    
    60
    -        child_path = os.path.join(root_directory, file_node.name)
    
    61
    -
    
    62
    -        with open(child_path, 'wb') as child_file:
    
    63
    -            write_fetch_blob(child_file, stub, file_node.digest, instance_name)
    
    64
    -
    
    65
    -    for symlink_node in directory.symlinks:
    
    66
    -        child_path = os.path.join(root_directory, symlink_node.name)
    
    67
    -
    
    68
    -        if os.path.isabs(symlink_node.target):
    
    69
    -            continue  # No out of temp-directory links for now.
    
    70
    -        target_path = os.path.join(root_directory, symlink_node.target)
    
    71
    -
    
    72
    -        os.symlink(child_path, target_path)
    
    73
    -
    
    74
    -
    
    75
    -def write_fetch_blob(target_file, stub, digest, instance_name=None):
    
    76
    -    """Extracts a blob from CAS into a local file.
    
    77
    -
    
    78
    -    Args:
    
    79
    -        target_file (str): local file to write.
    
    80
    -        stub (): gRPC stub for CAS communication.
    
    81
    -        digest (Digest): digest for the blob to fetch from CAS.
    
    82
    -        instance_name (str, optional): farm instance name to query data from.
    
    83
    -    """
    
    84
    -
    
    85
    -    for stream in gen_fetch_blob(stub, digest, instance_name):
    
    86
    -        target_file.write(stream)
    
    87
    -    target_file.flush()
    
    88
    -
    
    89
    -    assert digest.size_bytes == os.fstat(target_file.fileno()).st_size
    
    90
    -
    
    91
    -
    
    92
    -def parse_to_pb2_from_fetch(pb2, stub, digest, instance_name=""):
    
    93
    -    """ Fetches stream and parses it into given pb2
    
    94
    -    """
    
    95
    -
    
    96
    -    stream_bytes = b''
    
    97
    -    for stream in gen_fetch_blob(stub, digest, instance_name):
    
    98
    -        stream_bytes += stream
    
    99
    -
    
    100
    -    pb2.ParseFromString(stream_bytes)
    
    101
    -    return pb2
    
    102 21
     
    
    103 22
     
    
    104 23
     def create_digest(bytes_to_digest):
    

  • docs/source/reference_cli.rst
    ... ... @@ -50,17 +50,31 @@ BuildGrid's Command Line Interface (CLI) reference documentation.
    50 50
     
    
    51 51
     ----
    
    52 52
     
    
    53
    +.. _invoking-bgd-cas-download-dir:
    
    54
    +
    
    55
    +.. click:: buildgrid._app.commands.cmd_cas:download_directory
    
    56
    +   :prog: bgd cas download-dir
    
    57
    +
    
    58
    +----
    
    59
    +
    
    60
    +.. _invoking-bgd-cas-download-file:
    
    61
    +
    
    62
    +.. click:: buildgrid._app.commands.cmd_cas:download_file
    
    63
    +   :prog: bgd cas download-file
    
    64
    +
    
    65
    +----
    
    66
    +
    
    53 67
     .. _invoking-bgd-cas-upload-dir:
    
    54 68
     
    
    55
    -.. click:: buildgrid._app.commands.cmd_cas:upload_dir
    
    69
    +.. click:: buildgrid._app.commands.cmd_cas:upload_directory
    
    56 70
        :prog: bgd cas upload-dir
    
    57 71
     
    
    58 72
     ----
    
    59 73
     
    
    60
    -.. _invoking-bgd-cas-upload-files:
    
    74
    +.. _invoking-bgd-cas-upload-file:
    
    61 75
     
    
    62
    -.. click:: buildgrid._app.commands.cmd_cas:upload_files
    
    63
    -   :prog: bgd cas upload-files
    
    76
    +.. click:: buildgrid._app.commands.cmd_cas:upload_file
    
    77
    +   :prog: bgd cas upload-file
    
    64 78
     
    
    65 79
     ----
    
    66 80
     
    
    ... ... @@ -73,7 +87,7 @@ BuildGrid's Command Line Interface (CLI) reference documentation.
    73 87
     
    
    74 88
     .. _invoking-bgd-execute-command:
    
    75 89
     
    
    76
    -.. click:: buildgrid._app.commands.cmd_execute:command
    
    90
    +.. click:: buildgrid._app.commands.cmd_execute:run_command
    
    77 91
        :prog: bgd execute command
    
    78 92
     
    
    79 93
     ----
    
    ... ... @@ -112,6 +126,7 @@ BuildGrid's Command Line Interface (CLI) reference documentation.
    112 126
        :prog: bgd operation wait
    
    113 127
     
    
    114 128
     ----
    
    129
    +
    
    115 130
     .. _invoking-bgd-server:
    
    116 131
     
    
    117 132
     .. click:: buildgrid._app.commands.cmd_server:cli
    

  • docs/source/using_internal.rst
    ... ... @@ -20,7 +20,7 @@ In one terminal, start a server:
    20 20
     
    
    21 21
     In another terminal, upload an action to CAS:
    
    22 22
     
    
    23
    -.. code-block::sh
    
    23
    +.. code-block:: sh
    
    24 24
     
    
    25 25
        bgd cas upload-dummy
    
    26 26
     
    
    ... ... @@ -34,7 +34,7 @@ The stage should show as ``QUEUED`` as it awaits a bot to pick up the work:
    34 34
     
    
    35 35
     .. code-block:: sh
    
    36 36
     
    
    37
    -   bgd execute list
    
    37
    +   bgd operation list
    
    38 38
     
    
    39 39
     Create a bot session:
    
    40 40
     
    
    ... ... @@ -46,7 +46,7 @@ Show the work as completed:
    46 46
     
    
    47 47
     .. code-block:: sh
    
    48 48
     
    
    49
    -   bgd execute list
    
    49
    +   bgd operation list
    
    50 50
     
    
    51 51
     
    
    52 52
     .. _simple-build:
    

  • tests/action_cache.py
    ... ... @@ -17,10 +17,10 @@
    17 17
     
    
    18 18
     import pytest
    
    19 19
     
    
    20
    +from buildgrid._exceptions import NotFoundError
    
    21
    +from buildgrid._protos.build.bazel.remote.execution.v2 import remote_execution_pb2
    
    20 22
     from buildgrid.server.actioncache.storage import ActionCache
    
    21 23
     from buildgrid.server.cas.storage import lru_memory_cache
    
    22
    -from buildgrid.server._exceptions import NotFoundError
    
    23
    -from buildgrid._protos.build.bazel.remote.execution.v2 import remote_execution_pb2
    
    24 24
     
    
    25 25
     
    
    26 26
     @pytest.fixture
    

  • tests/cas/test_client.py
    ... ... @@ -14,12 +14,15 @@
    14 14
     
    
    15 15
     # pylint: disable=redefined-outer-name
    
    16 16
     
    
    17
    +
    
    18
    +from copy import deepcopy
    
    17 19
     import os
    
    20
    +import tempfile
    
    18 21
     
    
    19 22
     import grpc
    
    20 23
     import pytest
    
    21 24
     
    
    22
    -from buildgrid.client.cas import upload
    
    25
    +from buildgrid.client.cas import download, upload
    
    23 26
     from buildgrid._protos.build.bazel.remote.execution.v2 import remote_execution_pb2
    
    24 27
     from buildgrid.utils import create_digest
    
    25 28
     
    
    ... ... @@ -41,6 +44,8 @@ FILES = [
    41 44
         (os.path.join(DATA_DIR, 'hello.cc'),),
    
    42 45
         (os.path.join(DATA_DIR, 'hello', 'hello.c'),
    
    43 46
          os.path.join(DATA_DIR, 'hello', 'hello.h'))]
    
    47
    +FOLDERS = [
    
    48
    +    (os.path.join(DATA_DIR, 'hello'),)]
    
    44 49
     DIRECTORIES = [
    
    45 50
         (os.path.join(DATA_DIR, 'hello'),),
    
    46 51
         (os.path.join(DATA_DIR, 'hello'), DATA_DIR)]
    
    ... ... @@ -214,3 +219,145 @@ def test_upload_tree(instance, directory_paths):
    214 219
                 directory_digest = create_digest(tree.root.SerializeToString())
    
    215 220
     
    
    216 221
                 assert server.compare_directories(directory_digest, directory_path)
    
    222
    +
    
    223
    +
    
    224
    +@pytest.mark.parametrize('blobs', BLOBS)
    
    225
    +@pytest.mark.parametrize('instance', INTANCES)
    
    226
    +def test_download_blob(instance, blobs):
    
    227
    +    # Actual test function, to be run in a subprocess:
    
    228
    +    def __test_download_blob(queue, remote, instance, digests):
    
    229
    +        # Open a channel to the remote CAS server:
    
    230
    +        channel = grpc.insecure_channel(remote)
    
    231
    +
    
    232
    +        blobs = []
    
    233
    +        with download(channel, instance) as downloader:
    
    234
    +            if len(digests) > 1:
    
    235
    +                blobs.extend(downloader.get_blobs(digests))
    
    236
    +            else:
    
    237
    +                blobs.append(downloader.get_blob(digests[0]))
    
    238
    +
    
    239
    +        queue.put(blobs)
    
    240
    +
    
    241
    +    # Start a minimal CAS server in a subprocess:
    
    242
    +    with serve_cas([instance]) as server:
    
    243
    +        digests = []
    
    244
    +        for blob in blobs:
    
    245
    +            digest = server.store_blob(blob)
    
    246
    +            digests.append(digest)
    
    247
    +
    
    248
    +        blobs = run_in_subprocess(__test_download_blob,
    
    249
    +                                  server.remote, instance, digests)
    
    250
    +
    
    251
    +        for digest, blob in zip(digests, blobs):
    
    252
    +            assert server.compare_blobs(digest, blob)
    
    253
    +
    
    254
    +
    
    255
    +@pytest.mark.parametrize('messages', MESSAGES)
    
    256
    +@pytest.mark.parametrize('instance', INTANCES)
    
    257
    +def test_download_message(instance, messages):
    
    258
    +    # Actual test function, to be run in a subprocess:
    
    259
    +    def __test_download_message(queue, remote, instance, digests, empty_messages):
    
    260
    +        # Open a channel to the remote CAS server:
    
    261
    +        channel = grpc.insecure_channel(remote)
    
    262
    +
    
    263
    +        messages = []
    
    264
    +        with download(channel, instance) as downloader:
    
    265
    +            if len(digests) > 1:
    
    266
    +                messages = downloader.get_messages(digests, empty_messages)
    
    267
    +                messages = list([m.SerializeToString() for m in messages])
    
    268
    +            else:
    
    269
    +                message = downloader.get_message(digests[0], empty_messages[0])
    
    270
    +                messages.append(message.SerializeToString())
    
    271
    +
    
    272
    +        queue.put(messages)
    
    273
    +
    
    274
    +    # Start a minimal CAS server in a subprocess:
    
    275
    +    with serve_cas([instance]) as server:
    
    276
    +        empty_messages, digests = [], []
    
    277
    +        for message in messages:
    
    278
    +            digest = server.store_message(message)
    
    279
    +            digests.append(digest)
    
    280
    +
    
    281
    +            empty_message = deepcopy(message)
    
    282
    +            empty_message.Clear()
    
    283
    +            empty_messages.append(empty_message)
    
    284
    +
    
    285
    +        messages = run_in_subprocess(__test_download_message,
    
    286
    +                                     server.remote, instance, digests, empty_messages)
    
    287
    +
    
    288
    +        for digest, message_blob, message in zip(digests, messages, empty_messages):
    
    289
    +            message.ParseFromString(message_blob)
    
    290
    +
    
    291
    +            assert server.compare_messages(digest, message)
    
    292
    +
    
    293
    +
    
    294
    +@pytest.mark.parametrize('file_paths', FILES)
    
    295
    +@pytest.mark.parametrize('instance', INTANCES)
    
    296
    +def test_download_file(instance, file_paths):
    
    297
    +    # Actual test function, to be run in a subprocess:
    
    298
    +    def __test_download_file(queue, remote, instance, digests, paths):
    
    299
    +        # Open a channel to the remote CAS server:
    
    300
    +        channel = grpc.insecure_channel(remote)
    
    301
    +
    
    302
    +        with download(channel, instance) as downloader:
    
    303
    +            if len(digests) > 1:
    
    304
    +                for digest, path in zip(digests, paths):
    
    305
    +                    downloader.download_file(digest, path, queue=False)
    
    306
    +            else:
    
    307
    +                downloader.download_file(digests[0], paths[0], queue=False)
    
    308
    +
    
    309
    +        queue.put(None)
    
    310
    +
    
    311
    +    # Start a minimal CAS server in a subprocess:
    
    312
    +    with serve_cas([instance]) as server:
    
    313
    +        with tempfile.TemporaryDirectory() as temp_folder:
    
    314
    +            paths, digests = [], []
    
    315
    +            for file_path in file_paths:
    
    316
    +                digest = server.store_file(file_path)
    
    317
    +                digests.append(digest)
    
    318
    +
    
    319
    +                path = os.path.relpath(file_path, start=DATA_DIR)
    
    320
    +                path = os.path.join(temp_folder, path)
    
    321
    +                paths.append(path)
    
    322
    +
    
    323
    +                run_in_subprocess(__test_download_file,
    
    324
    +                                  server.remote, instance, digests, paths)
    
    325
    +
    
    326
    +            for digest, path in zip(digests, paths):
    
    327
    +                assert server.compare_files(digest, path)
    
    328
    +
    
    329
    +
    
    330
    +@pytest.mark.parametrize('folder_paths', FOLDERS)
    
    331
    +@pytest.mark.parametrize('instance', INTANCES)
    
    332
    +def test_download_directory(instance, folder_paths):
    
    333
    +    # Actual test function, to be run in a subprocess:
    
    334
    +    def __test_download_directory(queue, remote, instance, digests, paths):
    
    335
    +        # Open a channel to the remote CAS server:
    
    336
    +        channel = grpc.insecure_channel(remote)
    
    337
    +
    
    338
    +        with download(channel, instance) as downloader:
    
    339
    +            if len(digests) > 1:
    
    340
    +                for digest, path in zip(digests, paths):
    
    341
    +                    downloader.download_directory(digest, path)
    
    342
    +            else:
    
    343
    +                downloader.download_directory(digests[0], paths[0])
    
    344
    +
    
    345
    +        queue.put(None)
    
    346
    +
    
    347
    +    # Start a minimal CAS server in a subprocess:
    
    348
    +    with serve_cas([instance]) as server:
    
    349
    +        with tempfile.TemporaryDirectory() as temp_folder:
    
    350
    +            paths, digests = [], []
    
    351
    +            for folder_path in folder_paths:
    
    352
    +                digest = server.store_folder(folder_path)
    
    353
    +                digests.append(digest)
    
    354
    +
    
    355
    +                path = os.path.relpath(folder_path, start=DATA_DIR)
    
    356
    +                path = os.path.join(temp_folder, path)
    
    357
    +                paths.append(path)
    
    358
    +
    
    359
    +                run_in_subprocess(__test_download_directory,
    
    360
    +                                  server.remote, instance, digests, paths)
    
    361
    +
    
    362
    +            for digest, path in zip(digests, paths):
    
    363
    +                assert server.compare_directories(digest, path)

  • tests/integration/operations_service.py
    ... ... @@ -24,16 +24,15 @@ import grpc
    24 24
     from grpc._server import _Context
    
    25 25
     import pytest
    
    26 26
     
    
    27
    -from buildgrid.utils import create_digest
    
    28
    -from buildgrid.server.controller import ExecutionController
    
    29
    -from buildgrid.server.cas.storage import lru_memory_cache
    
    30
    -from buildgrid.server.operations import service
    
    31
    -from buildgrid.server.operations.service import OperationsService
    
    32
    -from buildgrid.server._exceptions import InvalidArgumentError
    
    33
    -
    
    27
    +from buildgrid._exceptions import InvalidArgumentError
    
    34 28
     from buildgrid._protos.build.bazel.remote.execution.v2 import remote_execution_pb2
    
    35 29
     from buildgrid._protos.google.longrunning import operations_pb2
    
    36 30
     from buildgrid._protos.google.rpc import status_pb2
    
    31
    +from buildgrid.server.cas.storage import lru_memory_cache
    
    32
    +from buildgrid.server.controller import ExecutionController
    
    33
    +from buildgrid.server.operations import service
    
    34
    +from buildgrid.server.operations.service import OperationsService
    
    35
    +from buildgrid.utils import create_digest
    
    37 36
     
    
    38 37
     
    
    39 38
     server = mock.create_autospec(grpc.server)
    

  • tests/utils/cas.py
    ... ... @@ -30,6 +30,7 @@ from buildgrid.server.cas.service import ContentAddressableStorageService
    30 30
     from buildgrid.server.cas.instance import ByteStreamInstance
    
    31 31
     from buildgrid.server.cas.instance import ContentAddressableStorageInstance
    
    32 32
     from buildgrid.server.cas.storage.disk import DiskStorage
    
    33
    +from buildgrid.utils import create_digest, merkle_tree_maker
    
    33 34
     
    
    34 35
     
    
    35 36
     @contextmanager
    
    ... ... @@ -124,6 +125,15 @@ class Server:
    124 125
         def get(self, digest):
    
    125 126
             return self.__storage.get_blob(digest).read()
    
    126 127
     
    
    128
    +    def store_blob(self, blob):
    
    129
    +        digest = create_digest(blob)
    
    130
    +        write_buffer = self.__storage.begin_write(digest)
    
    131
    +        write_buffer.write(blob)
    
    132
    +
    
    133
    +        self.__storage.commit_write(digest, write_buffer)
    
    134
    +
    
    135
    +        return digest
    
    136
    +
    
    127 137
         def compare_blobs(self, digest, blob):
    
    128 138
             if not self.__storage.has_blob(digest):
    
    129 139
                 return False
    
    ... ... @@ -133,6 +143,16 @@ class Server:
    133 143
     
    
    134 144
             return blob == stored_blob
    
    135 145
     
    
    146
    +    def store_message(self, message):
    
    147
    +        message_blob = message.SerializeToString()
    
    148
    +        message_digest = create_digest(message_blob)
    
    149
    +        write_buffer = self.__storage.begin_write(message_digest)
    
    150
    +        write_buffer.write(message_blob)
    
    151
    +
    
    152
    +        self.__storage.commit_write(message_digest, write_buffer)
    
    153
    +
    
    154
    +        return message_digest
    
    155
    +
    
    136 156
         def compare_messages(self, digest, message):
    
    137 157
             if not self.__storage.has_blob(digest):
    
    138 158
                 return False
    
    ... ... @@ -144,6 +164,17 @@ class Server:
    144 164
     
    
    145 165
             return message_blob == stored_blob
    
    146 166
     
    
    167
    +    def store_file(self, file_path):
    
    168
    +        with open(file_path, 'rb') as file_bytes:
    
    169
    +            file_blob = file_bytes.read()
    
    170
    +        file_digest = create_digest(file_blob)
    
    171
    +        write_buffer = self.__storage.begin_write(file_digest)
    
    172
    +        write_buffer.write(file_blob)
    
    173
    +
    
    174
    +        self.__storage.commit_write(file_digest, write_buffer)
    
    175
    +
    
    176
    +        return file_digest
    
    177
    +
    
    147 178
         def compare_files(self, digest, file_path):
    
    148 179
             if not self.__storage.has_blob(digest):
    
    149 180
                 return False
    
    ... ... @@ -156,6 +187,17 @@ class Server:
    156 187
     
    
    157 188
             return file_blob == stored_blob
    
    158 189
     
    
    190
    +    def store_folder(self, folder_path):
    
    191
    +        last_digest = None
    
    192
    +        for node, blob, _ in merkle_tree_maker(folder_path):
    
    193
    +            write_buffer = self.__storage.begin_write(node.digest)
    
    194
    +            write_buffer.write(blob)
    
    195
    +
    
    196
    +            self.__storage.commit_write(node.digest, write_buffer)
    
    197
    +            last_digest = node.digest
    
    198
    +
    
    199
    +        return last_digest
    
    200
    +
    
    159 201
         def compare_directories(self, digest, directory_path):
    
    160 202
             if not self.__storage.has_blob(digest):
    
    161 203
                 return False
    



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